module Num.ContinuedFraction
    ( -- * Conversion between numbers and continued fractions
      continuedFraction
    , collapseFraction
    -- * Rational approximations using continued fractions
    , approximate
    , convergent
    ) where

import           Control.Applicative   (liftA2)
import           Data.Functor.Foldable (ListF (..), ana)
import           Data.List.NonEmpty    (NonEmpty (..), fromList)
import           Data.Maybe            (fromJust)
import           Data.Ratio            (Ratio, denominator, (%))

-- | Check if it's an integer
--
-- >>> isInteger 3
-- True
-- >>> isInteger 2.5
-- False
isInteger :: (RealFrac a) => a -> Bool
isInteger = idem down
    where
        idem = ((==) <*>)
        down = realToFrac . (floor :: (RealFrac a) => a -> Integer)

-- | This take a number and returns its continued fraction expansion as a list of `Integer`s.
--
-- >>> continuedFraction 2
-- [2]
continuedFraction :: (RealFrac a, Integral b) => a -> [b]
continuedFraction = liftA2 (:) floor (ana coalgebra)
    where coalgebra x
              | isInteger x = Nil
              | otherwise = Cons (floor alpha) alpha
                  where alpha = 1 / (x - realToFrac (floor x :: Int))

-- | This takes a list of integers and returns the corresponding rational number, returning "Nothing" on empty lists.
--
-- >>> collapseFractionM []
-- Nothing
-- >>> collapseFractionM [1,2,2,2]
-- Just (17 % 12)
collapseFractionM :: (Integral a) => [Integer] -> Maybe (Ratio a)
collapseFractionM []     = Nothing
collapseFractionM [x]    = Just $ fromIntegral x % 1
collapseFractionM (x:xs) = fmap ((fromIntegral x % 1 +) . (1 /)) (collapseFractionM xs)

-- | Take a non-empty list of integers and return the corresponding rational number.
--
-- >>> collapseFraction (1 :| [2,2,2])
-- 17 % 12
collapseFraction :: (Integral a, Integral b) => NonEmpty b -> Ratio a
collapseFraction (x:|[]) = fromIntegral x % 1
collapseFraction (x:|xs) = (fromIntegral x % 1) + 1 / collapseFraction (fromList xs)

-- | Find a given convergent of the continued fraction expansion of a number
--
-- >>> convergent 4 $ sqrt 2
-- 17 % 12
convergent :: (RealFrac a, Integral b) => Int -> a -> Ratio b
convergent n x = fromJust . collapseFractionM $ take n (continuedFraction x)

-- | Find the best rational approximation to a number such that the denominator is bounded by a given value.
--
-- >>> approximate pi 100
-- 22 % 7
approximate :: (RealFrac a, Integral b) => a -> b -> Ratio b
approximate x d = last . takeWhile ((<= d) . denominator) $ fmap (flip convergent x) [1..]