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

import           Control.Applicative
import           Data.Foldable
import           Data.Functor.Base
import           Data.Functor.Foldable (ListF (..), apo, cata)
import           Data.List             (intersperse)
import           Data.List.NonEmpty    (NonEmpty (..))
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 integers.
--
-- >>> continuedFraction 2
-- [2]
continuedFraction :: (RealFrac a, Integral b) => a -> [b]
continuedFraction = apo coalgebra
    where coalgebra x
              | isInteger x = go $ Left []
              | otherwise   = go $ Right alpha
                  where alpha = 1 / (x - realToFrac (floor x :: Integer))
                        go    = Cons (floor x)

prettyFrac :: (Show a) => NonEmpty a -> String
prettyFrac (x :| xs) = "[" ++ show x ++ "; " ++ fold (intersperse ", " (show <$> xs)) ++ "]"

-- | Print a list as a continued fraction.
--
-- >>> prettyFracM (take 5 $ continuedFraction (sqrt 2))
-- Just "[1; 2, 2, 2, 2]"
prettyFracM :: (Show a) => [a] -> Maybe String
prettyFracM []     = Nothing
prettyFracM (x:xs) = Just (prettyFrac (x :| xs))

-- | 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, Integral b) => [a] -> Maybe (Ratio b)
collapseFractionM []     = Nothing
collapseFractionM (x:xs) = Just $ collapseFraction (x:|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 = cata a
    where a (NonEmptyF x xs) = go xs $ fromIntegral x % 1
          go = maybe id ((+) . (1/))

-- | 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 :: Integral a => [Integer] -> Maybe (Ratio a)) $
    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..]