module SubHask.Category.Polynomial
    where

import Data.List (intersperse,filter,reverse)
import qualified Prelude as P

import SubHask.Internal.Prelude
import SubHask.Category
import SubHask.Algebra
import SubHask.Monad
import SubHask.SubType

-------------------------------------------------------------------------------


-- | The type of polynomials over an arbitrary ring.
--
-- See <https://en.wikipedia.org/wiki/Polynomial__ring wikipedia> for more detail.
type Polynomial a = Polynomial_ a a

-- |
-- FIXME:
-- "Polynomial_" takes two type parameters in order to be compatible with the "Category" hierarchy of classes.
-- But currently, both types must match each other.
-- Can/Should we generalize this to allow polynomials between types?
--
data Polynomial_ a b where
    Polynomial_ :: (ValidLogic a, Ring a, a~b) => {-#UNPACK#-}![a] -> Polynomial_ a b

mkMutable [t| forall a b. Polynomial_ a b |]

instance (Eq r, Show r) => Show (Polynomial_ r r) where
    show (Polynomial_ xs) = concat $ intersperse " + " $ filter (/=[]) $ reverse $ imap go xs
        where
            -- FIXME:
            -- The code below results in prettier output but incurs an "Eq" constraint that confuses ghci
            go :: Int -> r -> String
            go 0 x = when (zero/=x) $ show x
            go 1 x = when (zero/=x) $ when (one/=x) (show x) ++ "x"
            go i x = when (zero/=x) $ when (one/=x) (show x) ++ "x^" ++ show i

            when :: Monoid a => Bool -> a -> a
            when cond x = if cond then x else zero


-------------------------------------------------------------------------------

newtype instance ProofOf Polynomial_ a = ProofOf { unProofOf :: Polynomial_ a a }

mkMutable [t| forall a. ProofOf Polynomial_ a |]

instance Ring a => Semigroup (ProofOf Polynomial_ a) where
    (ProofOf p1)+(ProofOf p2) = ProofOf $ p1+p2

instance (ValidLogic a, Ring a) => Cancellative (ProofOf Polynomial_ a) where
    (ProofOf p1)-(ProofOf p2) = ProofOf $ p1-p2

instance (ValidLogic a, Ring a) => Monoid (ProofOf Polynomial_ a) where
    zero = ProofOf zero

instance (Ring a, Abelian a) => Abelian (ProofOf Polynomial_ a)

instance (ValidLogic a, Ring a) => Group (ProofOf Polynomial_ a) where
    negate (ProofOf p) = ProofOf $ negate p

instance (ValidLogic a, Ring a) => Rg (ProofOf Polynomial_ a) where
    (ProofOf p1)*(ProofOf p2) = ProofOf $ p1*p2

instance (ValidLogic a, Ring a) => Rig (ProofOf Polynomial_ a) where
    one = ProofOf one

instance (ValidLogic a, Ring a) => Ring (ProofOf Polynomial_ a) where
    fromInteger i = ProofOf $ fromInteger i

provePolynomial :: (ValidLogic a, Ring a) => (ProofOf Polynomial_ a -> ProofOf Polynomial_ a) -> Polynomial_ a a
provePolynomial f = unProofOf $ f $ ProofOf $ Polynomial_ [0,1]
---------------------------------------

type instance Scalar (Polynomial_ a b) = Scalar b
type instance Logic (Polynomial_ a b) = Logic b

instance Eq b => Eq_ (Polynomial_ a b) where
    (Polynomial_ xs)==(Polynomial_ ys) = xs==ys

instance Ring r => Semigroup (Polynomial_ r r) where
    (Polynomial_ p1)+(Polynomial_ p2) = Polynomial_ $ sumList (+) p1 p2

instance (ValidLogic r, Ring r) => Monoid (Polynomial_ r r) where
    zero = Polynomial_ []

instance (ValidLogic r, Ring r) => Cancellative (Polynomial_ r r) where
    (Polynomial_ p1)-(Polynomial_ p2) = Polynomial_ $ sumList (-) p1 p2

instance (ValidLogic r, Ring r) => Group (Polynomial_ r r) where
    negate (Polynomial_ p) = Polynomial_ $ P.map negate p

instance (Ring r, Abelian r) => Abelian (Polynomial_ r r)

instance (ValidLogic r, Ring r) => Rg (Polynomial_ r r) where
    (Polynomial_ p1)*(Polynomial_ p2) = Polynomial_ $ P.foldl (sumList (+)) [] $ go p1 zero
        where
            go []     i = []
            go (x:xs) i = (P.replicate i zero ++ P.map (*x) p2):go xs (i+one)

instance (ValidLogic r, Ring r) => Rig (Polynomial_ r r) where
    one = Polynomial_ [one]

instance (ValidLogic r, Ring r) => Ring (Polynomial_ r r) where
    fromInteger i = Polynomial_ [fromInteger i]

type instance Polynomial_ r r >< r = Polynomial_ r r

instance IsScalar r => Module (Polynomial_ r r) where
    (Polynomial_ xs) .*  r               = Polynomial_ $ P.map (*r) xs

instance IsScalar r => FreeModule (Polynomial_ r r) where
    (Polynomial_ xs) .*. (Polynomial_ ys) = Polynomial_ $ P.zipWith (*) xs ys
    ones = Polynomial_ $ P.repeat one

sumList f [] ys = ys
sumList f xs [] = xs
sumList f (x:xs) (y:ys) = f x y:sumList f xs ys

---------------------------------------

instance Category Polynomial_ where
    type ValidCategory Polynomial_ a = (ValidLogic a, Ring a)
    id = Polynomial_ [zero, one]
    (Polynomial_ xs) . p2@(Polynomial_ _) = Polynomial_ (map (\x -> Polynomial_ [x]) xs) $ p2

instance Sup Polynomial_ Hask Hask
instance Sup Hask Polynomial_ Hask

instance Polynomial_ <: Hask where
    embedType_ = Embed2 evalPolynomial_

pow :: Rig r => r -> Int -> r
pow r i = foldl' (*) one $ P.replicate i r

evalPolynomial_ :: Polynomial_ a b -> a -> b
evalPolynomial_ (Polynomial_ xs) r = sum $ imap go xs
    where
        go i x = x*pow r i

-------------------------------------------------------------------------------

-- FIXME:
-- Polynomial_s should use the derivative interface from the Derivative module
--
-- class Category cat => Smooth cat where
--     derivative :: ValidCategory cat a b => cat a b Linear.+> cat a b
--
-- instance Smooth Polynomial_ where
--     derivative = unsafeProveLinear go
--         where
--             go (Polynomial_ xs) =  Polynomial_ $ P.tail $ P.zipWith (*) (inflist zero one) xs
--             inflist xs x = xs : inflist (xs+x) x
--
-- data MonoidT c a b = MonoidT (c a)