-- |
-- Module      : Data.Rhythm.ContinuedFractions
-- Description : Simple continued fractions
-- Copyright   : (c) Eric Bailey, 2024-2025
--
-- License     : MIT
-- Maintainer  : eric@ericb.me
-- Stability   : stable
-- Portability : POSIX
--
-- [Simple continued fractions](https://mathworld.wolfram.com/SimpleContinuedFraction.html)
-- represented by nonempty lists of terms.
--
-- \[
--   \begin{align*}
--     [b_0; b_1, b_2, b_3, \dotsc] &=
--     b_0 + \cfrac{1}{b_1 + \cfrac{1}{b_2 + \cfrac{1}{b_3 + \dotsm}}} \\
--     &= b_0 + \mathop{\vcenter{\Huge\mathcal{K}}}_{n=1}^{\infty} \frac{1}{b_n}
--   \end{align*}
-- \]
module Data.Rhythm.ContinuedFractions
  ( ContinuedFraction (..),
    collapseFraction,
    continuedFractionSqrt,
  )
where

import Control.Monad (guard)
import Data.Bifunctor (first)
import Data.Functor.Base (NonEmptyF (..))
import Data.Functor.Foldable (ana, cata)
import Data.List (intercalate)
import Data.List.NonEmpty (NonEmpty (..))
import Data.Maybe (listToMaybe)
import Data.Tuple.Extra (thd3)
import GHC.Generics (Generic)

-- | A 'ContinuedFraction' is a 'NonEmpty', potentially infinite, list of
-- integer 'terms'.
--
-- >>> ContinuedFraction (1 :| [2,3,4])
-- [1;2,3,4]
newtype ContinuedFraction = ContinuedFraction
  {ContinuedFraction -> NonEmpty Integer
terms :: NonEmpty Integer}
  deriving (ContinuedFraction -> ContinuedFraction -> Bool
(ContinuedFraction -> ContinuedFraction -> Bool)
-> (ContinuedFraction -> ContinuedFraction -> Bool)
-> Eq ContinuedFraction
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: ContinuedFraction -> ContinuedFraction -> Bool
== :: ContinuedFraction -> ContinuedFraction -> Bool
$c/= :: ContinuedFraction -> ContinuedFraction -> Bool
/= :: ContinuedFraction -> ContinuedFraction -> Bool
Eq, (forall x. ContinuedFraction -> Rep ContinuedFraction x)
-> (forall x. Rep ContinuedFraction x -> ContinuedFraction)
-> Generic ContinuedFraction
forall x. Rep ContinuedFraction x -> ContinuedFraction
forall x. ContinuedFraction -> Rep ContinuedFraction x
forall a.
(forall x. a -> Rep a x) -> (forall x. Rep a x -> a) -> Generic a
$cfrom :: forall x. ContinuedFraction -> Rep ContinuedFraction x
from :: forall x. ContinuedFraction -> Rep ContinuedFraction x
$cto :: forall x. Rep ContinuedFraction x -> ContinuedFraction
to :: forall x. Rep ContinuedFraction x -> ContinuedFraction
Generic, Eq ContinuedFraction
Eq ContinuedFraction =>
(ContinuedFraction -> ContinuedFraction -> Ordering)
-> (ContinuedFraction -> ContinuedFraction -> Bool)
-> (ContinuedFraction -> ContinuedFraction -> Bool)
-> (ContinuedFraction -> ContinuedFraction -> Bool)
-> (ContinuedFraction -> ContinuedFraction -> Bool)
-> (ContinuedFraction -> ContinuedFraction -> ContinuedFraction)
-> (ContinuedFraction -> ContinuedFraction -> ContinuedFraction)
-> Ord ContinuedFraction
ContinuedFraction -> ContinuedFraction -> Bool
ContinuedFraction -> ContinuedFraction -> Ordering
ContinuedFraction -> ContinuedFraction -> ContinuedFraction
forall a.
Eq a =>
(a -> a -> Ordering)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> a)
-> (a -> a -> a)
-> Ord a
$ccompare :: ContinuedFraction -> ContinuedFraction -> Ordering
compare :: ContinuedFraction -> ContinuedFraction -> Ordering
$c< :: ContinuedFraction -> ContinuedFraction -> Bool
< :: ContinuedFraction -> ContinuedFraction -> Bool
$c<= :: ContinuedFraction -> ContinuedFraction -> Bool
<= :: ContinuedFraction -> ContinuedFraction -> Bool
$c> :: ContinuedFraction -> ContinuedFraction -> Bool
> :: ContinuedFraction -> ContinuedFraction -> Bool
$c>= :: ContinuedFraction -> ContinuedFraction -> Bool
>= :: ContinuedFraction -> ContinuedFraction -> Bool
$cmax :: ContinuedFraction -> ContinuedFraction -> ContinuedFraction
max :: ContinuedFraction -> ContinuedFraction -> ContinuedFraction
$cmin :: ContinuedFraction -> ContinuedFraction -> ContinuedFraction
min :: ContinuedFraction -> ContinuedFraction -> ContinuedFraction
Ord)

instance Show ContinuedFraction where
  show :: ContinuedFraction -> String
show (ContinuedFraction (Integer
x :| [Integer]
xs)) =
    String
"["
      String -> ShowS
forall a. Semigroup a => a -> a -> a
<> Integer -> String
forall a. Show a => a -> String
show Integer
x
      String -> ShowS
forall a. Semigroup a => a -> a -> a
<> String
";"
      String -> ShowS
forall a. Semigroup a => a -> a -> a
<> String -> [String] -> String
forall a. [a] -> [[a]] -> [a]
intercalate String
"," ((Integer -> String) -> [Integer] -> [String]
forall a b. (a -> b) -> [a] -> [b]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Integer -> String
forall a. Show a => a -> String
show [Integer]
ys)
      String -> ShowS
forall a. Semigroup a => a -> a -> a
<> String -> (Integer -> String) -> Maybe Integer -> String
forall b a. b -> (a -> b) -> Maybe a -> b
maybe String
"" (String -> Integer -> String
forall a b. a -> b -> a
const String
",...") ([Integer] -> Maybe Integer
forall a. [a] -> Maybe a
listToMaybe [Integer]
zs)
      String -> ShowS
forall a. Semigroup a => a -> a -> a
<> String
"]"
    where
      ([Integer]
ys, [Integer]
zs) = Int -> [Integer] -> ([Integer], [Integer])
forall a. Int -> [a] -> ([a], [a])
splitAt Int
11 [Integer]
xs

-- | Evaluate a finite 'ContinuedFraction'.
--
-- >>> collapseFraction (ContinuedFraction (1 :| [2,3,4]))
-- 43 % 30
collapseFraction :: ContinuedFraction -> Rational
collapseFraction :: ContinuedFraction -> Rational
collapseFraction = (Base (NonEmpty Integer) Rational -> Rational)
-> NonEmpty Integer -> Rational
forall t a. Recursive t => (Base t a -> a) -> t -> a
forall a. (Base (NonEmpty Integer) a -> a) -> NonEmpty Integer -> a
cata (NonEmptyF Rational Rational -> Rational
forall {a}. Fractional a => NonEmptyF a a -> a
algebra (NonEmptyF Rational Rational -> Rational)
-> (NonEmptyF Integer Rational -> NonEmptyF Rational Rational)
-> NonEmptyF Integer Rational
-> Rational
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Integer -> Rational)
-> NonEmptyF Integer Rational -> NonEmptyF Rational Rational
forall a b c. (a -> b) -> NonEmptyF a c -> NonEmptyF b c
forall (p :: * -> * -> *) a b c.
Bifunctor p =>
(a -> b) -> p a c -> p b c
first Integer -> Rational
forall a. Real a => a -> Rational
toRational) (NonEmpty Integer -> Rational)
-> (ContinuedFraction -> NonEmpty Integer)
-> ContinuedFraction
-> Rational
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ContinuedFraction -> NonEmpty Integer
terms
  where
    algebra :: NonEmptyF a a -> a
algebra (NonEmptyF a
i Maybe a
Nothing) = a
i
    algebra (NonEmptyF a
i (Just a
d)) = a
i a -> a -> a
forall a. Num a => a -> a -> a
+ a -> a
forall a. Fractional a => a -> a
recip a
d

-- | Calculate the 'ContinuedFraction' representation of the square root of a
-- given number.
--
-- >>> continuedFractionSqrt 7
-- [2;1,1,1,4]
continuedFractionSqrt :: (Integral a) => a -> ContinuedFraction
continuedFractionSqrt :: forall a. Integral a => a -> ContinuedFraction
continuedFractionSqrt a
n
  | a
root a -> a -> a
forall a. Num a => a -> a -> a
* a
root a -> a -> Bool
forall a. Ord a => a -> a -> Bool
< a
n =
      NonEmpty Integer -> ContinuedFraction
ContinuedFraction (NonEmpty Integer -> ContinuedFraction)
-> NonEmpty Integer -> ContinuedFraction
forall a b. (a -> b) -> a -> b
$ a -> Integer
forall a b. (Integral a, Num b) => a -> b
fromIntegral (a -> Integer) -> ((a, a, a) -> a) -> (a, a, a) -> Integer
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a, a, a) -> a
forall a b c. (a, b, c) -> c
thd3 ((a, a, a) -> Integer) -> NonEmpty (a, a, a) -> NonEmpty Integer
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ((a, a, a) -> Base (NonEmpty (a, a, a)) (a, a, a))
-> (a, a, a) -> NonEmpty (a, a, a)
forall t a. Corecursive t => (a -> Base t a) -> a -> t
forall a.
(a -> Base (NonEmpty (a, a, a)) a) -> a -> NonEmpty (a, a, a)
ana (a, a, a) -> NonEmptyF (a, a, a) (a, a, a)
(a, a, a) -> Base (NonEmpty (a, a, a)) (a, a, a)
coalgebra (a
0, a
1, a
root)
  | Bool
otherwise =
      NonEmpty Integer -> ContinuedFraction
ContinuedFraction (a -> Integer
forall a b. (Integral a, Num b) => a -> b
fromIntegral a
root Integer -> [Integer] -> NonEmpty Integer
forall a. a -> [a] -> NonEmpty a
:| [])
  where
    coalgebra :: (a, a, a) -> NonEmptyF (a, a, a) (a, a, a)
coalgebra (a
a, a
b, a
c) =
      (a, a, a) -> Maybe (a, a, a) -> NonEmptyF (a, a, a) (a, a, a)
forall a b. a -> Maybe b -> NonEmptyF a b
NonEmptyF (a
a, a
b, a
c) (Maybe (a, a, a) -> NonEmptyF (a, a, a) (a, a, a))
-> Maybe (a, a, a) -> NonEmptyF (a, a, a) (a, a, a)
forall a b. (a -> b) -> a -> b
$
        do
          Bool -> Maybe ()
forall (f :: * -> *). Alternative f => Bool -> f ()
guard (a
c a -> a -> Bool
forall a. Eq a => a -> a -> Bool
/= a
2 a -> a -> a
forall a. Num a => a -> a -> a
* a
root)
          let a' :: a
a' = a
b a -> a -> a
forall a. Num a => a -> a -> a
* a
c a -> a -> a
forall a. Num a => a -> a -> a
- a
a
          let b' :: a
b' = (a
n a -> a -> a
forall a. Num a => a -> a -> a
- a
a' a -> a -> a
forall a. Num a => a -> a -> a
* a
a') a -> a -> a
forall a. Integral a => a -> a -> a
`div` a
b
          let c' :: a
c' = (a
root a -> a -> a
forall a. Num a => a -> a -> a
+ a
a') a -> a -> a
forall a. Integral a => a -> a -> a
`div` a
b'
          (a, a, a) -> Maybe (a, a, a)
forall a. a -> Maybe a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (a
a', a
b', a
c')
    root :: a
root = Double -> a
forall b. Integral b => Double -> b
forall a b. (RealFrac a, Integral b) => a -> b
truncate (forall a. Floating a => a -> a
sqrt @Double (a -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral a
n))