{-# LANGUAGE FlexibleInstances #-}
-----------------------------------------------------------------------------
-- Copyright 2019, Advise-Me project team. This file is distributed under 
-- the terms of the Apache License 2.0. For more information, see the files
-- "LICENSE.txt" and "NOTICE.txt", which are included in the distribution.
-----------------------------------------------------------------------------
-- |
-- Maintainer  :  bastiaan.heeren@ou.nl
-- Stability   :  provisional
-- Portability :  portable (depends on ghc)
--
-----------------------------------------------------------------------------

module Domain.Math.Data.MultivariatePolynomial where

import Control.Monad
import qualified Data.Map as M
import Data.Maybe
import Domain.Algebra.Field
import Domain.Math.CleanUp
import Domain.Math.Expr
import Domain.Math.Numeric.Views
import Ideas.Common.View
import Ideas.Utils.Prelude
import Ideas.Utils.Uniplate
import Prelude hiding (subtract)
import Recognize.Expr.Normalform (distributeExponent, rewriteSqrt)

multivariatePolynomialView :: View Expr (MultivariatePolynomial Expr)
multivariatePolynomialView = makeView toMultivariatePolynomial fromMultivariatePolynomial

multivariatePolynomialViewWith :: View Expr a -> View Expr (MultivariatePolynomial a)
multivariatePolynomialViewWith v = multivariatePolynomialView >>> traverseView v

toMultivariatePolynomial :: Expr -> Maybe (MultivariatePolynomial Expr)
toMultivariatePolynomial expr =
   case fixpoint (transform (distributeExponent . rewriteSqrt)) expr of
      x :+: y -> (+) <$> toMultivariatePolynomial x <*> toMultivariatePolynomial y
      x :*: y -> (*) <$> toMultivariatePolynomial x <*> toMultivariatePolynomial y
      x :-: y -> (-) <$> toMultivariatePolynomial x <*> toMultivariatePolynomial y
      Negate e -> (0-) <$> toMultivariatePolynomial e
      Nat e -> Just (con $ fromInteger e)
      x :/: y -> join (divide <$> toMultivariatePolynomial x <*> return y) --match rationalView (cleanUpExpr y)) -- SAVE DIV!
      Sqrt _ -> Nothing -- flip expo (1/2) <$> (toMultivariatePolynomial e)
      Number d -> let t = toRational d in Just (con $ fromRational t ) -- (numerator t) % (denominator t))
      Var s -> Just (var s)
      Sym s [x,y] | isPowerSymbol s -> expo <$> toMultivariatePolynomial x <*> match naturalView (cleanUpExpr y)
      Sym _ _ -> Nothing -- Sym s <$> map toMultivariatePolynomial es

fromMultivariatePolynomial :: MultivariatePolynomial Expr -> Expr
fromMultivariatePolynomial (Sum m) | m == M.empty = 0
                                   | otherwise = foldr1 (.+.) [ fromFactors fac .*. r | (fac, r) <- M.toList m ]

fromFactors :: Factors -> Expr
fromFactors (Product m) | m == M.empty = 1
                        | otherwise = foldr1 (.*.) [
                                        if n == 1 then
                                          Var x
                                        else
                                          Var x ** fromInteger n
                                        | (x, n) <- M.toList m ]

multivariatePolynomialToMap :: MultivariatePolynomial a -> M.Map Factors a
multivariatePolynomialToMap (Sum s) = s
-- 2ab+b^2-2ab+a^2
ex :: MultivariatePolynomial Expr
ex = 2 * a * b + expo b 2 - 2 * a * b + expo a 2
 where
   a = var "a"
   b = var "b"

type Var = String

-- invariant: values should not be 0
newtype MultivariatePolynomial a = Sum (M.Map Factors a)
 deriving Eq

instance Show (MultivariatePolynomial Expr) where
  show = show . fromMultivariatePolynomial

-- invariant: values (exponents) must be > 0
newtype Factors = Product (M.Map Var Integer) -- a^3 b^1 c^2
  deriving (Eq, Ord)

instance Show Factors where
  show = show . fromFactors

instance Functor MultivariatePolynomial where
  fmap f (Sum xs) = Sum (f <$> xs)

instance Foldable MultivariatePolynomial where
  foldMap f (Sum xs) = foldMap f xs

instance Traversable MultivariatePolynomial where
   sequenceA (Sum m) = Sum <$> sequenceA m

-- smart constructor that takes care of MultivariatePolynomial's invariant
mkSum :: (Eq a, Field a) => M.Map Factors a -> MultivariatePolynomial a
mkSum = Sum . M.filter (/= zero)

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

instance (Num a, Eq a, Field a) => Num (MultivariatePolynomial a) where
   (+) = plus
   (-) = subtract
   (*) = times
   negate = subtract 0
   fromInteger = con . fromInteger

var :: (Eq a, Field a) => Var -> MultivariatePolynomial a
var x = mkSum (M.singleton (varF x) one)

con :: (Eq a, Field a) => a -> MultivariatePolynomial a
con n = mkSum (M.singleton oneF n)

plus :: (Eq a, Field a) => MultivariatePolynomial a -> MultivariatePolynomial a -> MultivariatePolynomial a
plus (Sum m1) (Sum m2) = mkSum (M.unionWith (|+|) m1 m2)

negateMV ::  (Eq a,Num a, Field a) => MultivariatePolynomial a -> MultivariatePolynomial a
negateMV (Sum s) = Sum (M.map negate s)

subtract :: (Eq a,Num a, Field a) => MultivariatePolynomial a -> MultivariatePolynomial a -> MultivariatePolynomial a
subtract a b = a `plus` negateMV b

times :: (Eq a, Field a) => MultivariatePolynomial a -> MultivariatePolynomial a -> MultivariatePolynomial a
times p (Sum m) = foldr plus (con zero) [ times3 r (times2 fac p) | (fac, r) <- M.toList m ]

times2 :: (Eq a, Field a) => Factors -> MultivariatePolynomial a -> MultivariatePolynomial a
times2 f (Sum m) = mkSum (M.mapKeys (timesF f) m)

times3 :: (Eq a, Field a) => a -> MultivariatePolynomial a -> MultivariatePolynomial a
times3 r (Sum m) = mkSum (M.map (|*| r) m)

divide :: (Eq a, Field a) => MultivariatePolynomial a -> a -> Maybe (MultivariatePolynomial a) -- eigenlijk Maybe: r /= 0, dus safe division
divide p r | r == zero = Nothing
           | otherwise = Just $ times p (con (one |/| r))

unsafeDivide :: (Eq a, Field a) => MultivariatePolynomial a -> a -> MultivariatePolynomial a
unsafeDivide p r = fromMaybe (error "unsafe multivariate polynomial division by zero") (divide p r)

-- we would like to have expo :: MultivariatePolynomial -> Rational -> MultivariatePolynomial
-- but, then we would need to factorize the MultivariatePolynomial: https://en.wikipedia.org/wiki/Factorization_of_polynomials
expo :: (Eq a, Field a) => MultivariatePolynomial a -> Integer -> MultivariatePolynomial a -- kan vast slimmer
expo _ 0 = con one
expo p n = times p (expo p (n-1))

----------------------
varF :: Var -> Factors
varF x = Product (M.singleton x 1)

oneF :: Factors
oneF = Product M.empty

timesF :: Factors -> Factors -> Factors
timesF (Product m1) (Product m2) = Product (M.unionWith (+) m1 m2)