-- |
-- Module      : Data.Rhythm.ContinuedFractions
-- Copyright   : (c) Eric Bailey, 2024-2025
--
-- License     : MIT
-- Maintainer  : eric@ericb.me
-- Stability   : experimental
-- 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)

-- | 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, 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 = a -> a
forall a. Integral a => a -> a
isqrt a
n
    n' :: a
n' = a -> a
forall a b. (Integral a, Num b) => a -> b
fromIntegral a
n

isqrt :: (Integral a) => a -> a
isqrt :: forall a. Integral a => a -> a
isqrt = Double -> a
forall b. Integral b => Double -> b
forall a b. (RealFrac a, Integral b) => a -> b
truncate (Double -> a) -> (a -> Double) -> a -> a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Floating a => a -> a
sqrt @Double (Double -> Double) -> (a -> Double) -> a -> Double
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral