{-# LANGUAGE KindSignatures, DataKinds  #-}
module Padic where

import Data.Ratio
import Data.List (genericLength)
import GHC.TypeLits

data Padic (n :: Nat) = Null
                      | Padic { unit :: [Int], order :: Int }

-- valuation of the base
modulo :: (KnownNat p, Integral i) => Padic p -> i
modulo = fromIntegral . natVal

-- Constructor for zero value
pZero :: KnownNat p => Padic p
pZero = Padic (repeat 0) 0

-- Smart constructor, adjusts trailing zeros with the order.
mkPadic :: (KnownNat p, Integral i) => [i] -> Int -> Padic p
mkPadic u k = go 0 (fromIntegral <$> u)
  where
    go 17 _ = pZero
    go i (0:u) = go (i+1) u
    go i u = Padic u (k-i)

-- Constructor for p-adic unit
mkUnit :: (KnownNat p, Integral i) => [i] -> Padic p
mkUnit u = mkPadic u 0

-- Zero test (up to 1/p^17)
isZero :: KnownNat p => Padic p -> Bool
isZero (Padic u _) = all (== 0) (take 17 u)
isZero _ = False

-- p-adic norm
pNorm :: KnownNat p => Padic p -> Ratio Int
pNorm Null = undefined
pNorm p = fromIntegral (modulo p) ^^ (- order p)

-- test for an integerness up to p^-17
isInteger :: KnownNat p => Padic p -> Bool
isInteger Null = False
isInteger (Padic s k) = case splitAt k s of
  ([],i) -> length (takeWhile (==0) $ reverse (take 20 i)) > 3
  _ -> False

-- p-adics are shown with 1/p^17 precision
instance KnownNat p => Show (Padic p) where
  show Null = "Null"
  show x@(Padic u k) =
    show (modulo x) ++ "-adic: " ++
    (case si of {[] -> "0"; _ -> si})
    ++ "." ++
    (case f of {[] -> "0"; _ -> sf})
    where
      (f,i) = case compare k 0 of
        LT -> ([], replicate (-k) 0 ++ u)
        EQ -> ([], u)
        GT -> splitAt k (u ++ repeat 0)
      sf = foldMap showD $ reverse $ take 17 f
      si = foldMap showD $ dropWhile (== 0) $ reverse $ take 17 i
      el s = if length s > 16 then "…" else ""
      showD n = [(['0'..'9']++['a'..'z']) !! n]

instance KnownNat p => Eq (Padic p) where
  a == b = isZero (a - b)

instance KnownNat p => Ord (Padic p) where
  compare = error "Ordering is undefined fo p-adics."

instance KnownNat p => Num (Padic p) where
  fromInteger 0 = pZero
  fromInteger n = pAdic (fromInteger n)

  x@(Padic a ka) + Padic b kb = mkPadic s k
    where
      k = ka `max` kb
      s = addMod (modulo x)
          (replicate (k-ka) 0 ++ a)
          (replicate (k-kb) 0 ++ b)
  _ + _ = Null

  x@(Padic a ka) * Padic b kb =
    mkPadic (mulMod (modulo x) a b) (ka + kb)
  _ * _ = Null

  negate x@(Padic u k) =
    case map (\y -> modulo x - 1 - y) u of
      n:ns -> Padic ((n+1):ns) k
      [] -> pZero
  negate _ = Null

  abs p = pAdic (pNorm p)

  signum = undefined

------------------------------------------------------------
-- conversion from rationals to p-adics

instance KnownNat p => Fractional (Padic p) where
  fromRational = pAdic

  recip Null = Null
  recip x@(Padic (u:us) k)
    | isZero x = Null
    | gcd p u /= 1 = Null
    | otherwise = mkPadic res (-k)
    where
      p = modulo x
      res = longDivMod p (1:repeat 0) (u:us)

pAdic :: (Show i, Integral i, KnownNat p)
      => Ratio i -> Padic p
pAdic 0 = pZero
pAdic x = res
  where
    p = modulo res
    (k, q) = getUnit p x
    (n, d) = (numerator q, denominator q)
    res = maybe Null process $ recipMod p d

    process r = mkPadic (series n) k
      where
        series n
          | n == 0 = repeat 0
          | n `mod` p == 0 = 0 : series (n `div` p)
          | otherwise =
              let m = (n * r) `mod` p
              in m : series ((n - m * d) `div` p)

------------------------------------------------------------
-- conversion from p-adics to rationals
-- works for relatively small denominators

instance KnownNat p => Real (Padic p) where
  toRational Null = error "no rational representation!"
  toRational x@(Padic s k) = res
    where
      p = modulo x
      res = case break isInteger $ take 10000 $ iterate (x +) x of
        (_,[]) -> - toRational (- x)
        (d, i:_) -> (fromBase p (unit i) * (p^(- order i))) % (genericLength d + 1)

      fromBase p = foldr (\x r -> r*p + x) 0 .
                   take 20 . map fromIntegral

--------------------------------------------------------------------------------
-- helper functions

-- extracts p-adic unit from a rational number
getUnit :: Integral i => i -> Ratio i -> (Int, Ratio i)
getUnit p x = (genericLength k1 - genericLength k2, c)
  where
    (k1,b:_) = span (\n -> denominator n `mod` p == 0) $
               iterate (* fromIntegral p) x
    (k2,c:_) = span (\n -> numerator n `mod` p == 0) $
               iterate (/ fromIntegral p) b


-- Reciprocal of a number modulo p (extended Euclidean algorithm).
-- For non-prime p returns Nothing non-invertible element of the ring.
recipMod :: Integral i => i -> i -> Maybe i
recipMod p 1 = Just 1
recipMod p a | gcd p a == 1 = Just $ go 0 1 p a
             | otherwise = Nothing
  where
    go t _ _ 0 = t `mod` p
    go t nt r nr =
      let q = r `div` nr
      in go nt (t - q*nt) nr (r - q*nr)

-- Addition of two sequences modulo p
addMod p = go 0
  where
    go 0 [] ys = ys
    go 0 xs [] = xs
    go s [] ys = go 0 [s] ys
    go s xs [] = go 0 xs [s]
    go s (x:xs) (y:ys) =
      let (q, r) = (x + y + s) `divMod` p
      in r : go q xs ys

-- Subtraction of two sequences modulo p
subMod p a (b:bs) = addMod p a $ (p-b) : ((p - 1 -) <$> bs)

-- Multiplication of two sequences modulo p
mulMod p as [b] = mulMod p [b] as
mulMod p as bs = case as of
  [0] -> repeat 0
  [1] -> bs
  [a] -> go 0 bs
    where
      go s [] = [s]
      go s (b:bs) =
        let (q, r) = (a * b + s) `divMod` p
        in r : go q bs
  as -> go bs
    where
      go [] = []
      go (b:bs) =
        let c:cs = mulMod p [b] as
        in c : addMod p (go bs) cs

-- Division of two sequences modulo p
longDivMod p a (b:bs) = case recipMod p b of
  Nothing -> error $
    show b ++ " is not invertible modulo " ++ show p
  Just r -> go a
    where
      go [] = []
      go (0:xs) = 0 : go xs
      go (x:xs) =
        let m = (x*r) `mod` p
            _:zs = subMod p (x:xs) (mulMod p [m] (b:bs))
        in m : go zs