{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE TypeFamilies #-}
-- |
-- A continued fraction whose terms may be positive, negative or
-- zero. The methods in @Floating@ are supported, with the exception
-- of @asin@, @acos@ and @atan@.
module Math.ContinuedFraction (
  CF,
  CF'(..),
  cfString,
  cfcf
) where

import Data.Maybe (catMaybes, mapMaybe)
import Data.Ratio

import Math.ContinuedFraction.Interval

newtype CF' a = CF [a]
type CF = CF' Integer

type Hom a = (a, a,
              a, a)

type Bihom a = (a, a, a, a,
                a, a, a, a)

class (Fractional (FractionField a)) => HasFractionField a where
  type FractionField a :: *
  insert :: a -> FractionField a
  frac :: (a, a) -> FractionField a
  extract :: FractionField a -> (a, a)

instance HasFractionField Integer where
  type FractionField Integer = Rational
  insert = fromInteger
  {-# INLINE insert #-}
  frac = uncurry (%)
  {-# INLINE frac #-}
  extract r = (numerator r, denominator r)
  {-# INLINE extract #-}

instance HasFractionField Rational where
  type FractionField Rational = Rational
  insert = id
  {-# INLINE insert #-}
  frac = uncurry (/)
  {-# INLINE frac #-}
  extract r = (numerator r % 1, denominator r % 1)
  {-# INLINE extract #-}

instance HasFractionField CF where
  type FractionField CF = CF
  insert = id
  {-# INLINE insert #-}
  frac = uncurry (/)
  {-# INLINE frac #-}
  extract r = (r, 1)
  {-# INLINE extract #-}

homEmit :: Num a => Hom a -> a -> Hom a
homEmit (n0, n1,
         d0, d1) x = (d0,        d1,
                      n0 - d0*x, n1 - d1*x)

homAbsorb :: Num a => Hom a -> a -> Hom a
homAbsorb (n0, n1,
           d0, d1) x = (n0*x + n1, n0,
                        d0*x + d1, d0)

det :: Num a => Hom a -> a
det (n0, n1,
     d0, d1) = n0 * d1 - n1 * d0

homEval :: (Eq a, Num a, HasFractionField a) => Hom a -> Extended (FractionField a) -> Extended (FractionField a)
homEval (n0, n1,
         d0, d1) (Finite q) | denom /= 0 = Finite $ frac (num, denom)
                            | num == 0 = error "0/0 in homQ"
                            | otherwise = Infinity
  where (qnum, qdenom) = extract q
        num   = n0 * qnum + n1 * qdenom
        denom = d0 * qnum + d1 * qdenom
homEval (n0, _n1,
         d0, _d1) Infinity = Finite $ frac (n0, d0)

constantFor :: (Eq a, Num a, HasFractionField a) => Hom a -> Extended (FractionField a)
constantFor (_, _,
             0, 0) = Infinity
constantFor (0, 0,
             0, _) = Finite 0
constantFor (0, 0,
             _, 0) = Finite 0
constantFor (a, 0,
             b, 0) = Finite $ frac (a, b)
constantFor (_, a,
             _, b) = Finite $ frac (a, b)

boundHom :: (Ord a, Num a, HasFractionField a, Ord (FractionField a)) => Hom a -> Interval (FractionField a) -> Interval (FractionField a)
boundHom h (Interval i s _) | d > 0 = interval i' s'
                            | d < 0 = interval s' i'
                            | otherwise = Interval c c True
  where d = det h
        i' = homEval h i
        s' = homEval h s
        c = constantFor h

primitiveBound :: forall a. (Ord a, Num a, HasFractionField a) => a -> Interval (FractionField a)
primitiveBound n | abs n < 1 = Interval (Finite $ insert bot) (Finite $ insert top) True
  where bot = (-2) :: a
        top = 2 :: a
primitiveBound n = Interval (Finite $ an - 0.5) (Finite $ 0.5 - an) False
  where an = insert $ abs n

-- TODO: just take the rational answer from the hom
nthPrimitiveBounds :: (Ord a, Num a, HasFractionField a, Ord (FractionField a)) =>
                       CF' a -> [Interval (FractionField a)]
nthPrimitiveBounds (CF cf) = zipWith boundHom homs (map primitiveBound cf) ++ repeat (Interval ev ev True)
  where homs = scanl homAbsorb (1,0,0,1) cf
        ev = evaluate (CF cf)

evaluate :: (HasFractionField a, Eq (FractionField a)) => CF' a -> Extended (FractionField a)
evaluate (CF []) = Infinity
evaluate (CF [c]) = Finite $ insert c
evaluate (CF (c:cs)) = case next of
                        (Finite 0) -> Infinity
                        Infinity   -> Finite $ insert c
                        (Finite r) -> Finite $ insert c + recip r
  where next = evaluate (CF cs)

valueToCF :: RealFrac a => a -> CF
valueToCF r = if rest == 0 then
                CF [d]
              else
                let (CF ds)  = valueToCF (recip rest) in CF (d:ds)
  where (d, rest) = properFraction r


existsEmittable :: (RealFrac a, Integral b) => Interval a -> Maybe b
existsEmittable (Interval Infinity    Infinity _)  = Nothing
existsEmittable (Interval Infinity   (Finite _) _) = Nothing
existsEmittable (Interval (Finite _)  Infinity _)  = Nothing
existsEmittable int@(Interval (Finite a) (Finite b) _) = euclideanCheck int a b

euclideanCheck :: (Num a, Ord a, RealFrac a, Integral b) => Interval a -> a -> a -> Maybe b
euclideanCheck int a b
  | not isThin = Nothing
  | 0 `elementOf` int && not subsetZero = Nothing
  | zi /= 0 && zs /= 0 = Just z
  | subsetZero = Just 0
  | otherwise = Nothing
    where zi = round a
          zs = round b
          z  = if abs zs < abs zi then zs else zi
          isThin = abs z > 3 || abs (zi - zs) < 2
          subsetZero = int `subset` Interval (Finite (-2)) (Finite 2) True

hom :: (Ord a, Num a, HasFractionField a, RealFrac (FractionField a)) => Hom a -> CF' a -> CF
hom (_n0, _n1,
     0,   0) _  = CF []
hom (_n0, _n1,
     0,   _d1) (CF []) = CF []
hom (n0, _n1,
     d0, _d1) (CF []) = valueToCF $ frac (n0, d0)
hom h (CF (x:xs)) = case existsEmittable $ boundHom h (primitiveBound x) of
                     Just n ->  CF $ n : rest
                       where (CF rest) = hom (homEmit h (fromInteger n)) (CF (x:xs))
                     Nothing -> hom (homAbsorb h x) (CF xs)

bihomEmit :: Num a => Bihom a -> a -> Bihom a
bihomEmit (n0, n1, n2, n3,
           d0, d1, d2, d3) x = (d0,        d1,        d2,        d3,
                                n0 - d0*x, n1 - d1*x, n2 - d2*x, n3 - d3*x)

bihomAbsorbX :: Num a => Bihom a -> a -> Bihom a
bihomAbsorbX (n0, n1, n2, n3,
              d0, d1, d2, d3) x = (n0*x + n1, n0, n2*x + n3, n2,
                                   d0*x + d1, d0, d2*x + d3, d2)

bihomAbsorbY :: Num a => Bihom a -> a -> Bihom a
bihomAbsorbY (n0, n1, n2, n3,
              d0, d1, d2, d3) y = (n0*y + n2, n1*y + n3, n0, n1,
                                   d0*y + d2, d1*y + d3, d0, d1)

bihomSubstituteX :: (Num a, HasFractionField a) => Bihom a -> Extended (FractionField a) -> Hom a
bihomSubstituteX (n0, n1, n2, n3,
                  d0, d1, d2, d3) (Finite x) = (n0*num + n1*den, n2*num + n3*den,
                                                d0*num + d1*den, d2*num + d3*den)
  where (num, den) = extract x
bihomSubstituteX (n0, _n1, n2, _n3,
                  d0, _d1, d2, _d3) Infinity = (n0, n2,
                                                d0, d2)

bihomSubstituteY :: (Num a, HasFractionField a) => Bihom a -> Extended (FractionField a) -> Hom a
bihomSubstituteY (n0, n1, n2, n3,
                  d0, d1, d2, d3) (Finite y) = (n0*num + n2*den, n1*num + n3*den,
                                                d0*num + d2*den, d1*num + d3*den)
  where (num, den) = extract y
bihomSubstituteY (n0, n1, _n2, _n3,
                  d0, d1, _d2, _d3) Infinity = (n0, n1,
                                                d0, d1)

boundBihomAndSelect :: (Ord a, Num a, HasFractionField a, Eq (FractionField a), Ord (FractionField a)) =>
              Bihom a -> Interval (FractionField a) -> Interval (FractionField a) -> (Interval (FractionField a), Bool)
boundBihomAndSelect bh x@(Interval ix sx _) y@(Interval iy sy _) = (interval, intX `smallerThan` intY)
  where interval = ixy `mergeInterval` iyx `mergeInterval` sxy `mergeInterval` syx
        ixy = boundHom (bihomSubstituteX bh ix) y
        iyx = boundHom (bihomSubstituteY bh iy) x
        sxy = boundHom (bihomSubstituteX bh sx) y
        syx = boundHom (bihomSubstituteY bh sy) x
        intX = if ixy `smallerThan` sxy then sxy else ixy
        intY = if iyx `smallerThan` syx then syx else iyx

bihom :: (Ord a, Num a, HasFractionField a, RealFrac (FractionField a))
         => Bihom a -> CF' a -> CF' a -> CF
bihom bh (CF []) y = hom (bihomSubstituteX bh Infinity) y
bihom bh x (CF []) = hom (bihomSubstituteY bh Infinity) x
bihom bh (CF (x:xs)) (CF (y:ys)) =
  let (bound, which) = boundBihomAndSelect bh (primitiveBound x) (primitiveBound y) in
  case existsEmittable bound of
    Just n -> CF $ n : rest
      where (CF rest) = bihom (bihomEmit bh (fromInteger n)) (CF (x:xs)) (CF (y:ys))
    Nothing -> if which then
                 let bh' = bihomAbsorbX bh x in bihom bh' (CF xs) (CF (y:ys))
               else
                 let bh' = bihomAbsorbY bh y in bihom bh' (CF (x:xs)) (CF ys)

homchain :: [Hom Integer] -> CF
homchain (h:h':hs) = case quotEmit h of
                     Just n ->  CF $ n : rest
                       where (CF rest) = homchain ((homEmit h n):h':hs)
                     Nothing -> homchain ((h `mult` h'):hs)
  where quotEmit (n0, n1,
                  d0, d1) = if d0 /= 0 && d1 /= 0 && n0 `quot` d0 == n1 `quot` d1 then Just $ n0 `quot` d0 else Nothing
        mult (n0, n1,
              d0, d1)
             (n0', n1',
              d0', d1') =(n0*n0' + n1*d0', n0*n1' + n1*d1',
                          d0*n0' + d1*d0', d0*n1' + d1*d1')

instance Num CF where
  (+) = bihom (0, 1, 1, 0,
               0, 0, 0, 1)
  (-) = bihom (0, -1, 1, 0,
               0,  0, 0, 1)
  (*) = bihom (1, 0, 0, 0,
               0, 0, 0, 1)

  fromInteger n = CF [n]

  signum x = case 0 `compare` x of
              EQ -> 0
              LT -> 1
              GT -> -1

  abs x | x < 0     = -x
        | otherwise = x

instance Fractional CF where
  (/) = bihom (0, 0, 1, 0,
               0, 1, 0, 0)

  fromRational = valueToCF

base :: Integer
base = 10

rationalDigits :: Rational -> [Integer]
rationalDigits 0 = []
rationalDigits r = let d = num `quot` den in
                   d : rationalDigits (fromInteger base * (r - fromInteger d))
  where num = numerator r
        den = denominator r

digits :: CF -> [Integer]
digits = go (1, 0, 0, 1)
  where go (0, 0, _, _) _ = []
        go (p, _, q, _) (CF []) = rationalDigits (p % q)
        go h (CF (c:cs)) = case intervalDigit $ boundHom h (primitiveBound c) of
                            Nothing -> let h' = homAbsorb h c in go h' (CF cs)
                            Just d  -> d : go (homEmitDigit h d) (CF (c:cs))
        homEmitDigit (n0, n1,
                      d0, d1) d = (base * (n0 - d0*d), base * (n1 - d1*d),
                                   d0,                 d1)

-- | Produce the (possibly infinite) decimal expansion of a continued
-- fraction
cfString :: CF -> String
cfString (CF []) = "Infinity"
cfString cf | cf < 0 = '-' : cfString (-cf)
cfString cf = case digits cf of
               []     -> "0"
               [i]    -> show i
               (i:is) -> show i ++ "." ++ concatMap show is

instance Show CF where
  show = take 50 . cfString

instance Eq CF where
  a == b = a `compare` b == EQ

instance Ord CF where
  a `compare` b = head $ catMaybes $ zipWith comparePosition (nthPrimitiveBounds a) (nthPrimitiveBounds b)

instance Real CF where
  toRational = error "CF: toRational"

instance RealFrac CF where
  properFraction cf = head $ mapMaybe checkValid $ nthPrimitiveBounds cf
    where checkValid (Interval (Finite a) (Finite b) True) =
            if truncate a == truncate b then
              Just (truncate a, cf - fromInteger (truncate a))
            else
              Nothing
          checkValid _ = Nothing

-- | Convert a continued fraction whose terms are continued fractions
-- into an ordinary continued fraction with integer terms
cfcf :: CF' CF -> CF
cfcf = hom (1, 0, 0, 1)

instance Floating CF where
  pi = homchain ((0,4,1,0) : map go [1..])
    where go n = (2*n-1, n^2,
                  1,     0)

  exp r | r < -1 || r > 1 = (exp (r / 2))^2
  exp r = cfcf (CF $ 1 : concatMap go [0..])
    where go n = [fromInteger (4*n+1) / r,
                  -2,
                  -fromInteger (4*n+3) / r,
                  2]

  log r | r < 0.5 = log (2 * r) - log 2
  log r | r > 2   = log (r / 2) + log 2
  log r = cfcf (CF $ 0 : concatMap go [0..])
    where go n = [fromInteger (2*n+1) / (r-1),
                  fromRational $ 2 % (n+1)]

  tan r | r < -1 || r > 1 = bihom ( 0,1,1,0,
                                   -1,0,0,1) tanhalf tanhalf
    where tanhalf = tan (r / 2)
  tan r = cfcf (CF $ 0 : concatMap go [0..])
    where go n = [fromInteger (4*n+1) / r,
                  -fromInteger (4*n+3) / r]

  sin r = bihom (0,2,0,0,
                 1,0,0,1) tanhalf tanhalf
    where tanhalf = tan (r / 2)

  cos r = bihom (-1,0,0,1,
                  1,0,0,1) tanhalf tanhalf
    where tanhalf = tan (r / 2)

  sinh r = bihom (1,0,0,-1,
                  0,1,1, 0) expr expr
    where expr = exp r

  cosh r = bihom (1,0,0,1,
                  0,1,1,0) expr expr
    where expr = exp r

  tanh r = bihom (1,0,0,-1,
                  1,0,0, 1) expr expr
    where expr = exp r