module Data.TrueReal where

import Data.Ratio
import Data.Bits (Bits, shiftL, shiftR, (.&.))
import qualified Data.Digits as D

factorial :: Integer -> Integer
factorial n | n < 0 = error "factorial: negative argument"
factorial 0 = 1
factorial n = fac 1 n
  where fac acc 1 = acc
        fac acc n = fac (acc * n) (n-1)

minusN :: Num a => Integer -> a -> a
minusN n k | odd n     = -k
           | otherwise = k


applyN :: Integer -> (a -> a) -> a -> a
applyN n f x = go n x
  where
    go 0 acc = acc
    go n acc = go (n - 1) (f acc)

data R = R Integer D.D

check :: Int -> R -> Maybe (Int, Integer, (Integer, Integer))
check n (R _ d) = D.check n d

align :: Integer -> R -> Maybe R
align e0 (R e d)
  | e0 >= e   = Just $ R e0 $ D.extend (e0 - e) d
  | otherwise = do dd <- D.reduceTo (e - e0) d
                   return $ R e0 dd

reduceN :: Integer -> R -> R
reduceN n (R e d) = R (e - de) dr
  where (de, dr) = D.reduce n d

reduceTo :: Integer -> R -> R
reduceTo e0 x@(R e _) | e0 >= e   = x
                      | otherwise = reduceN (e - e0) x

reduce :: R -> R
reduce = reduceN (-1)

build :: Integer -> Integer -> D.D -> R
build n e d = reduceN n $ R (e + n) d

sgn :: R -> Int
sgn (R _ d) = D.sgn d

debugR :: Integer -> R -> (Integer, D.D)
debugR n (R e d) = (e, D.chop n d)

op1 :: (D.D -> D.D) -> R -> R
op1 f (R e d) = build 1 e $ f d

op2 :: (D.D -> D.D -> D.D) -> R -> R -> R
op2 f (R e1 d1) (R e2 d2) =
  build 1 e $ f de1 de2
  where (e, de1, de2) = D.align e1 e2 d1 d2

instance Num R where
  (+) = op2 D.add
  (-) = op2 D.sub

  R e1 d1 * R e2 d2 =
    build 3 (e1 + e2) $ D.mul d1 d2

  negate = op1 D.neg
  abs = op1 D.abs

  fromInteger n = R e d
    where (e, d) = D.buildInteger n

  signum = fromIntegral . sgn

sqr :: R -> R
sqr (R e d) = build 3 (e `shiftL` 1) $ D.sqr d

shift :: R -> Integer -> R
shift (R e d) n = build 1 (e + de) dd
  where
    (de, dd) = D.shift d n

chop :: Integer -> R -> R
chop e0 (R e d)
  | de <= 0 = 0
  | otherwise = R e $ D.chop de d
  where
    de = e - e0 + 1

minmax :: R -> R -> (R, R)
minmax (R e1 d1) (R e2 d2) =
  (build 1 e mn, build 1 e mx)
  where (e, de1, de2) = D.align e1 e2 d1 d2
        (mn, mx) = D.minmax de1 de2

approx :: RealFloat a => R -> a
approx r
  | floatRadix zr /= 2 = error "Data.TrueReal.approx: can only work with radix 2"
  | ef >=  1000000 = sg * (1 / 0)
  | ef <= -1000000 = sg * 0
  | otherwise = encodeFloat (fst $ D.split (fromIntegral n) d) $ fromIntegral ef
    where
      zr = undefined :: a
      fd = floatDigits zr
      (R e d) = reduce r
      sg = fromIntegral $ D.sgn d
      n = (fd + 2 * D.nBits - 2) `quot` D.nBits
      ef = (e - fromIntegral n + 1) * fromIntegral D.nBits

split :: R -> (Integer, R)
split r@(R e d)
  | e < 0     = (0, r)
  | otherwise = (n, R (-1) f)
    where (n, f) = D.split (e + 1) d

instance Eq R where
  a == b = sgn (a - b) == 0

instance Ord R where
  compare a b = compare (sgn $ a - b) 0

  min a b = fst $ minmax a b

  max a b = snd $ minmax a b

instance Fractional R where
  fromRational r = R e d
    where (e, d) = D.buildRational r

  n / d =
    build 4 (en - ed - 1) $ D.div dn dd
    where
      R en dn = n
      R ed dd = reduce d

infixr 7 *.
(*.) :: Integer -> R -> R
0 *. _ = 0
1 *. r = r
2 *. r = r `shift` 1
n *. r@(R e d) | abs n >= D.mulNmax = fromIntegral n * r
               | otherwise          = build 2 e $ D.mulN n d

infixl 7 /.
(/.) :: R -> Integer -> R
r /. 0 = error "Data.TrueReal.(/.): division by zero"
r /. 1 = r
r /. (-1) = -r
r /. 2 = r `shift` (-1)
R e d /. n = build 1 e $ D.divN d n

infixr 7 *% 
(*%) :: Rational -> R -> R
(*%) q r = (numerator q *. r) /. denominator q

alignAll :: [R] -> (Integer, [D.D])
alignAll [] = (0, [])
alignAll l@(R e _ : _) = (e, map go l)
  where
    go r = 
      case align e r of
        Just (R _ d) -> d
        Nothing      -> error "Data.TrueReal.series: non decreasing"

series :: [R] -> R
series [] = 0
series l =
  build 2 e $ D.series ds
  where
    (e, ds) = alignAll l

limit :: [R] -> R
limit [] = error "limit: empty list"
limit l =
  build 2 e $ D.limit ds
  where
    (e, ds) = alignAll l

indexSeries :: (Integer -> R) -> R
indexSeries f = series $ [f n | n <- [0..]]

iterationLimit :: (R -> R) -> R -> (Integer -> Integer) -> R
iterationLimit f r0 prec =
  limit $ go 0 r0
  where
    go n r = rc : go (n + 1) (f rc)
      where rc = chop (prec n) r

powerList :: R -> [R]
powerList x = 1 : xs
  where
    xs = x : go xs
    go [] = error "powerList: impossible"
    go (y : ys) =  y2 : x * y2 : go ys
      where y2 = sqr y

powerSeries :: [Rational] -> R -> R
powerSeries l x = series $ zipWith (*%) l $ powerList x

evenPowerSeries :: [Rational] -> R -> R
evenPowerSeries l x = series $ zipWith (*%) l $ powerList $ sqr x

oddPowerSeries :: [Rational] -> R -> R
oddPowerSeries l x = x * evenPowerSeries l x

indexPowerSeries :: (Integer -> Rational) -> R -> R
indexPowerSeries f x = powerSeries [f n | n <- [0..]] x

evenIndexPowerSeries :: (Integer -> Rational) -> R -> R
evenIndexPowerSeries f x = evenPowerSeries [f n | n <- [0..]] x

oddIndexPowerSeries :: (Integer -> Rational) -> R -> R
oddIndexPowerSeries f x = oddPowerSeries [f n | n <- [0..]] x

ln2 :: R
ln2 = indexSeries f
  where f 0 = 2 / 3
        f n = fromRational $ 1 % (2 * (16 ^ n))
              * (1 % (2 * n) + 1 % (4 * n + 1) + 1 % (8 * n + 4) + 1 % (16 * n + 12))

roots2 :: [R]
roots2 = go $ sqrt 0.5
  where go x = x : go (sqrt x)

doublePi :: R
doublePi = 2 *. pi

halfPi :: R
halfPi = pi /. 2

agmList :: R -> R -> [(R, R)]
agmList x y = (x, y) : agmList ((x + y) /. 2) (sqrt $ x * y)

agm :: R -> R -> R
agm x y = limit $ map fst $ agmList x y

piAgm :: R
piAgm = (2 *. sqr a) / (1 - series c)
  where
    a = limit $ map fst l
    l = agmList 1 (sqrt 0.5)
    c = go 0 l
      where go _ [] = error "Data.TrueReal.piAgm: impossible"
            go n ((an, bn) : ll) = (sqr an - sqr bn) `shift` n : go (n + 1) ll

logAgm :: R -> R
logAgm x = s *. (pi / agm u v - k *. ln2)
  where
    x0 = reduce x
    (s, ax) | approx x0 > 1 = (1, 1 / x0)
            | otherwise     = (-1, x0)
    t = 1e-4
    ra = approx ax
    k | ra < t    = 0
      | otherwise = round (log (ra / t) / log 2)
    r = ax `shift` (-k)
    rn = powerList r
    rno = map (\x -> r * sqr x) rn
    rn2 = scanl (*) 1 rno
    rnn = zipWith (*) rn2 rn
    u = sqr (2 *. series rn2 - 1)
    v = sqrt r * sqr (2 *. series rnn)

instance Floating R where
  pi = indexSeries f
    where f n = shift q (-10 * n - 6)
            where q = fromRational $ minusN n $
                    -32%(4 * n + 1)-1%(4 * n + 3)+256%(10 * n + 1)-64%(10 * n + 3)
                    -4%(10 * n + 5)-4%(10 * n + 7)+1%(10 * n + 9)

  sqrt r | odd e     = build 2 ((e - 1) `shiftR` 1) $ D.sqrt d
         | otherwise = build 2 (e `shiftR` 1) $ D.sqrt $ D.extend 1 d
    where
      R e d = r

  exp x = applyN k sqr (s + sqrt (1 + sqr s)) `shift` n
    where
      (n, _) = split $ x / ln2
      r = reduceTo 0 $ x - n *. ln2
      ra = abs $ approx $ chop (-1) r
      t = 1e-3
      k | ra < t    = 0
        | otherwise = round (log (ra / t) / log 2)
      s = oddIndexPowerSeries (\n -> 1 % factorial (2 * n + 1)) (r `shift` (-k))

  log x
    | D.sgn d <= 0 = error "log: non positive argument"
    | otherwise  =  ((e * fromIntegral D.nBits)%1 + k0) *% ln2 + s
    where
      R e d = reduce x
      x0 = R 0 d
      n0 = 10
      k0 = round (2^n0 * log (approx x0) / log 2)%2^n0
      k = fromRational k0 :: Double
      n = floor k0
      r = go (k0 - fromIntegral n) (x0 `shift` (-n))  roots2
      go f z (rt : rts)
        | f == 0    = z
        | ff >= 1   = go (ff - 1) (z * rt) rts
        | otherwise = go ff z rts
        where ff = 2 * f
      go _ _ [] = error "Data.TrueReal.log: impossible"
      s = oddIndexPowerSeries (\n -> 2%(2 * n + 1)) $ (r - 1) / (r + 1)

  cos x = applyN k (\x -> (2 *. sqr x) - 1) s
    where
      (n, _) = split $ x / doublePi
      r = reduceTo 0 $ x - n *. doublePi
      ra = abs $ approx $ chop (-1) r
      t = 1e-3
      k | ra < t    = 0
        | otherwise = round (log (ra / t) / log 2)
      s = evenIndexPowerSeries (\n -> 1 % minusN n (factorial (2 * n))) (r `shift` (-k))

  atan x = s `shift` (k + 1)
    where
      f z = z / (1 + sqrt (1 + sqr z))
      r = reduceTo 0 $ f x
      ra = abs $ approx $ chop (-1) r
      t = 1e-3
      k | ra < t    = 0
        | otherwise = round (log (ra / t) / log 2)
      s = oddIndexPowerSeries (\n -> 1 % minusN n (2 * n + 1)) $ applyN k f r

  sin x = cos (x - halfPi)
  asin x = atan (x / sqrt (1 - sqr x))
  acos x = halfPi - asin x
  cosh x = let e = exp x in (x + 1 / x) /. 2
  sinh x = let e = exp x in (x - 1 / x) /. 2
  acosh x = log (x + sqrt (sqr x - 1))
  asinh x = log (x + sqrt (sqr x + 1))
  atanh x = log ((1 + x) / (1 - x)) /. 2

baseRepr :: Int -> R -> (Int, Integer, Int, [ Int ])
baseRepr b x = (sgn x0, k, fromIntegral ip + is, fs)
  where
    x0 = reduce x
    ax0 = abs x0
    (k0, R _ d) = split $ log ax0 / log (fromIntegral b)
    k | D.sgn (D.chop 1 d) < 0 = k0 - 1
      | otherwise              = k0
    (ip, fp) = split $ (fromIntegral b)^^(-k) * ax0
    Just (R (-1) ds) = align (-1) fp
    (is, fs) = D.digits b ds

baseStrings :: Int -> R -> (String, String, String, String)
baseStrings b x
  | sg == 0   = ("", "0", "0", "")
  | sg == -1  = ("-", show e, show ip, go fp)
  | otherwise = ("", show e, show ip, go fp)
  where
    (sg, e, ip, fp) = baseRepr b x

    go [] = ""
    go (d : ds)
      | d == -1 = '(' : '-' : '1' : ')' : go ds
      | d >= 0 && d < 10 = toEnum (d + 48) : go ds
      | d >= 10 && d < 36 = toEnum (d + 87) : go ds
      | otherwise = error $ "Data.TrueReal.baseStrings: invalid digit" ++ show d

showBase :: Int -> R -> String
showBase b x@(R e _)
  | b <= 1 || b > 36 = error "Data.TrueReal.showBase: invalid base"
  | Just _ <- align (e - 10) x = "0 ?"
  | ex == "0" = sg ++ ip ++ '.' : fp
  | otherwise = sg ++ show b ++ "^" ++ ex ++ " * " ++ ip ++ '.' : fp
  where
    (sg, ex, ip, fp) = baseStrings b x

instance Show R where
  show x = showBase 10 x
