module Quox.Polynomial

import public Quox.Syntax.Subst
import public Quox.Context
import public Quox.Semiring
import Data.Maybe
import Data.SortedMap
import Data.Singleton
import Quox.PrettyValExtra

%default total

%hide Prelude.toList

public export
Monom : Nat -> Type
Monom = Context' Nat

export %inline
mZero : {n : Nat} -> Monom n
mZero = replicate n 0

export %inline
mIsZero : Monom n -> Bool
mIsZero = all (== 0)


private
PolyBody : Type -> Nat -> Type
PolyBody coeff scope = SortedMap (Monom scope) coeff

private %inline
getScope : PolyBody coeff scope -> Maybe (Singleton scope)
getScope p = lengthPrf0 . fst <$> leftMost p

export
data Polynomial : (coeff : Type) -> (scope : Nat) -> Type where
  Const : (k : coeff) -> Polynomial coeff scope
  Many  : (p : PolyBody coeff scope) -> Polynomial coeff scope
%name Polynomial p, q
-- the Const constructor exists so that `pconst` (and by extension, `one`, and
-- the TimesMonoid instance) doesn't need to know the scope size

export %inline
toConst' : List (Monom scope, coeff) -> Maybe coeff
toConst' [(m, k)] = k <$ guard (mIsZero m)
toConst' _        = Nothing

export %inline
toConst : PolyBody coeff scope -> Maybe coeff
toConst = toConst' . toList

export %inline
toTrimmedList : AddMonoid v => SortedMap k v -> List (k, v)
toTrimmedList = List.filter (not . isZero . snd) . toList

export %inline
toList : {scope : Nat} -> AddMonoid coeff =>
         Polynomial coeff scope -> List (Monom scope, coeff)
toList (Const k) = [(mZero, k)]
toList (Many  p) = toTrimmedList p

private %inline
addBody : AddMonoid coeff =>
          PolyBody coeff scope -> PolyBody coeff scope -> PolyBody coeff scope
addBody = mergeWith (+.)

export %inline
fromList : AddMonoid coeff =>
           List (Monom scope, coeff) -> Polynomial coeff scope
fromList xs = case toConst' xs of
  Just k  => Const k
  Nothing => Many $ foldl add1 empty xs
where
  add1 : PolyBody coeff scope -> (Monom scope, coeff) -> PolyBody coeff scope
  add1 p (x, k) = p `addBody` singleton x k


export %inline
(AddMonoid coeff, Eq coeff) => Eq (Polynomial coeff scope) where
  Const x == Const y = x == y
  Const x == Many  q = maybe False (== x) $ toConst q
  Many  p == Const y = maybe False (== y) $ toConst p
  Many  p == Many  q = toTrimmedList p == toTrimmedList q

export %inline
(AddMonoid coeff, Ord coeff) => Ord (Polynomial coeff scope) where
  compare (Const x) (Const y) = compare x y
  compare (Const x) (Many  q) = maybe LT (compare x) $ toConst q
  compare (Many  p) (Const y) = maybe GT (flip compare y) $ toConst p
  compare (Many  p) (Many  q) = compare (toTrimmedList p) (toTrimmedList q)

export %inline
{scope : Nat} -> (AddMonoid coeff, Show coeff) =>
Show (Polynomial coeff scope) where
  showPrec d p = showCon d "fromList" $ showArg (toList p)

export %inline
{scope : Nat} -> (AddMonoid coeff, PrettyVal coeff) =>
PrettyVal (Polynomial coeff scope) where
  prettyVal p = Con "fromList" [prettyVal $ toList p]


export %inline
pconst : a -> Polynomial a n
pconst = Const

export %inline
(.at) : AddMonoid a => Polynomial a n -> Monom n -> a
(Const k).at m = if mIsZero m then k else zero
(Many  p).at m = fromMaybe zero $ lookup m p

export %inline
(.at0) : AddMonoid a => Polynomial a n -> a
(Const k).at0 = k
(Many  p).at0 = fromMaybe zero $ do
  (m, k) <- leftMost p
  guard $ mIsZero m
  pure k


private %inline
addConstMany : AddMonoid coeff =>
               coeff -> PolyBody coeff scope -> Polynomial coeff scope
addConstMany k p = case getScope p of
  Nothing      => Const k
  Just (Val _) => Many $ p `addBody` singleton mZero k

export %inline
AddMonoid a => AddMonoid (Polynomial a n) where
  zero   = Const zero

  isZero (Const k) = isZero k
  isZero (Many  p) = maybe False isZero $ toConst p

  Const j +. Const k = Const $ j +. k
  Const j +. Many  q = addConstMany j q
  Many  p +. Const k = addConstMany k p
  Many  p +. Many  q = Many $ p `addBody` q

export %inline
Semiring a => TimesMonoid (Polynomial a n) where
  one = pconst one

  Const j *. Const k = Const $ j *. k
  Const j *. Many  q = Many $ map (j *.) q
  Many  p *. Const k = Many $ map (*. k) p
  Many  p *. Many  q = fromList $ do
    (xs, i) <- toList p
    (ys, j) <- toList q
    let k = i *. j
    guard $ not $ isZero k
    pure (zipWith (+) xs ys, i *. j)

export %inline
Semiring a => VecSpace a (Polynomial a n) where
  x *: xs = Const x *. xs


private
shiftMonom : Shift from to -> Monom from -> Monom to
shiftMonom SZ      m = m
shiftMonom (SS by) m = shiftMonom by m :< 0

export
CanShift (Polynomial coeff) where
  Const k // _  = Const k
  Many  p // by =
    let xs = mapFst (shiftMonom by) <$> toList p in
    Many $ fromList xs


export
TimesMonoid a => FromVarR (Polynomial a) where
  fromVarR i l {n} =
    let m = tabulateVar n $ \j => if i == j then 1 else 0 in
    Many $ singleton m one


private
substMonom : Semiring a => {from, to : Nat} ->
             Subst (Polynomial a) from to -> Monom from -> a -> Polynomial a to
substMonom th m k = sproduct {init = pconst k} $
  zipWith (\i, p => getR th i noLoc ^. p) (allVars from) m

export
Semiring a => CanSubstSelfR (Polynomial a) where
  Const k //? _  = Const k
  Many  p //? th = ssum $ map (uncurry $ substMonom th) $ toList p