-- |
-- Module: Symtegration.Polynomial.Indexed
-- Description: A polynomial representation mapping the power of each term to its coefficient.
-- Copyright: Copyright 2025 Yoo Chung
-- License: Apache-2.0
-- Maintainer: dev@chungyc.org
module Symtegration.Polynomial.Indexed
  ( IndexedPolynomial,
    IndexedSymbolicPolynomial,
    IndexedPolynomialWith,
  )
where

import Control.DeepSeq (NFData)
import Data.IntMap (IntMap)
import Data.IntMap qualified as IntMap
import Data.List (intersperse)
import Data.Maybe (fromMaybe)
import Data.Ratio (denominator, numerator)
import Data.Text (unpack)
import GHC.Generics (Generic)
import Symtegration.Polynomial
import Symtegration.Symbolic
import TextShow

-- | Polynomial representation which maps the power of each term to its coefficient.
-- Exponents are represented with 'Int', while coefficients are represented with 'Rational'.
-- It is an instance of the 'Polynomial' type class.
type IndexedPolynomial = IndexedPolynomialWith Rational

-- | Polynomial representation which maps the power of each term to its coefficient.
-- Exponents are represented with 'Int', while coefficients are represented with 'Expression'.
-- It is an instance of the 'Polynomial' type class.
type IndexedSymbolicPolynomial = IndexedPolynomialWith Expression

-- | Polynomial representation which maps the power of each term to its coefficient.
-- Exponents are represented with 'Int'.  Coefficients have a type as specified by the type parameter.
-- These types are an instance of the 'Polynomial' type class.
type IndexedPolynomialWith a = P Int a

-- | Type with two type parameters so that it can be an instance of 'Polynomial'.
-- The first type parameter is not involved in the data constructor;
-- it is used to set the exponent type for 'Polynomial'.
newtype P a b = P (IntMap b) deriving (P a b -> P a b -> Bool
(P a b -> P a b -> Bool) -> (P a b -> P a b -> Bool) -> Eq (P a b)
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
forall k (a :: k) b. Eq b => P a b -> P a b -> Bool
$c== :: forall k (a :: k) b. Eq b => P a b -> P a b -> Bool
== :: P a b -> P a b -> Bool
$c/= :: forall k (a :: k) b. Eq b => P a b -> P a b -> Bool
/= :: P a b -> P a b -> Bool
Eq, (forall x. P a b -> Rep (P a b) x)
-> (forall x. Rep (P a b) x -> P a b) -> Generic (P a b)
forall x. Rep (P a b) x -> P a b
forall x. P a b -> Rep (P a b) x
forall a.
(forall x. a -> Rep a x) -> (forall x. Rep a x -> a) -> Generic a
forall k (a :: k) b x. Rep (P a b) x -> P a b
forall k (a :: k) b x. P a b -> Rep (P a b) x
$cfrom :: forall k (a :: k) b x. P a b -> Rep (P a b) x
from :: forall x. P a b -> Rep (P a b) x
$cto :: forall k (a :: k) b x. Rep (P a b) x -> P a b
to :: forall x. Rep (P a b) x -> P a b
Generic, P a b -> ()
(P a b -> ()) -> NFData (P a b)
forall a. (a -> ()) -> NFData a
forall k (a :: k) b. NFData b => P a b -> ()
$crnf :: forall k (a :: k) b. NFData b => P a b -> ()
rnf :: P a b -> ()
NFData)

instance Show (P Int Rational) where
  show :: P Int Rational -> String
show = Text -> String
unpack (Text -> String)
-> (P Int Rational -> Text) -> P Int Rational -> String
forall b c a. (b -> c) -> (a -> b) -> a -> c
. P Int Rational -> Text
forall a. TextShow a => a -> Text
showt

instance TextShow (P Int Rational) where
  showb :: P Int Rational -> Builder
showb (P IntMap Rational
m)
    | IntMap Rational -> Bool
forall a. IntMap a -> Bool
IntMap.null IntMap Rational
m = Builder
"0"
    | Bool
otherwise =
        [Builder] -> Builder
forall a. Monoid a => [a] -> a
mconcat ([Builder] -> Builder) -> [Builder] -> Builder
forall a b. (a -> b) -> a -> b
$
          Builder -> [Builder] -> [Builder]
forall a. a -> [a] -> [a]
intersperse Builder
" + " ([Builder] -> [Builder]) -> [Builder] -> [Builder]
forall a b. (a -> b) -> a -> b
$
            ((Int, Rational) -> Builder) -> [(Int, Rational)] -> [Builder]
forall a b. (a -> b) -> [a] -> [b]
map (Int, Rational) -> Builder
forall {a} {a}.
(Eq a, Integral a, TextShow a, TextShow a, Num a) =>
(a, Ratio a) -> Builder
showTerm ([(Int, Rational)] -> [Builder]) -> [(Int, Rational)] -> [Builder]
forall a b. (a -> b) -> a -> b
$
              IntMap Rational -> [(Int, Rational)]
forall a. IntMap a -> [(Int, a)]
IntMap.toDescList IntMap Rational
m
    where
      showTerm :: (a, Ratio a) -> Builder
showTerm (a
0, Ratio a
1) = Builder
"1"
      showTerm (a
0, Ratio a
c) = Ratio a -> Builder
forall {a}. (Integral a, TextShow a) => Ratio a -> Builder
showCoefficient Ratio a
c
      showTerm (a
1, Ratio a
c) = Ratio a -> Builder
forall {a}. (Integral a, TextShow a) => Ratio a -> Builder
showCoefficient Ratio a
c Builder -> Builder -> Builder
forall a. Semigroup a => a -> a -> a
<> Builder
"x"
      showTerm (a
e, Ratio a
1) = Builder
"x^" Builder -> Builder -> Builder
forall a. Semigroup a => a -> a -> a
<> a -> Builder
forall a. TextShow a => a -> Builder
showb a
e
      showTerm (a
e, Ratio a
c) = Ratio a -> Builder
forall {a}. (Integral a, TextShow a) => Ratio a -> Builder
showCoefficient Ratio a
c Builder -> Builder -> Builder
forall a. Semigroup a => a -> a -> a
<> Builder
"x^" Builder -> Builder -> Builder
forall a. Semigroup a => a -> a -> a
<> a -> Builder
forall a. TextShow a => a -> Builder
showb a
e
      showCoefficient :: Ratio a -> Builder
showCoefficient Ratio a
r
        | Ratio a
1 <- Ratio a
r = Builder
forall a. Monoid a => a
mempty
        | a
1 <- Ratio a -> a
forall a. Ratio a -> a
denominator Ratio a
r, Ratio a
r Ratio a -> Ratio a -> Bool
forall a. Ord a => a -> a -> Bool
> Ratio a
0 = a -> Builder
forall a. TextShow a => a -> Builder
showb (a -> Builder) -> a -> Builder
forall a b. (a -> b) -> a -> b
$ Ratio a -> a
forall a. Ratio a -> a
numerator Ratio a
r
        | a
1 <- Ratio a -> a
forall a. Ratio a -> a
denominator Ratio a
r, Ratio a
r Ratio a -> Ratio a -> Bool
forall a. Ord a => a -> a -> Bool
< Ratio a
0 = Bool -> Builder -> Builder
showbParen Bool
True (Builder -> Builder) -> Builder -> Builder
forall a b. (a -> b) -> a -> b
$ a -> Builder
forall a. TextShow a => a -> Builder
showb (a -> Builder) -> a -> Builder
forall a b. (a -> b) -> a -> b
$ Ratio a -> a
forall a. Ratio a -> a
numerator Ratio a
r
        | Bool
otherwise = Bool -> Builder -> Builder
showbParen Bool
True (Builder -> Builder) -> Builder -> Builder
forall a b. (a -> b) -> a -> b
$ Ratio a -> Builder
forall a. TextShow a => a -> Builder
showb Ratio a
r

instance (Polynomial p e c, TextShow (p e c)) => Show (IndexedPolynomialWith (p e c)) where
  show :: IndexedPolynomialWith (p e c) -> String
show = Text -> String
unpack (Text -> String)
-> (IndexedPolynomialWith (p e c) -> Text)
-> IndexedPolynomialWith (p e c)
-> String
forall b c a. (b -> c) -> (a -> b) -> a -> c
. IndexedPolynomialWith (p e c) -> Text
forall a. TextShow a => a -> Text
showt

instance (Polynomial p e c, TextShow (p e c)) => TextShow (IndexedPolynomialWith (p e c)) where
  showb :: IndexedPolynomialWith (p e c) -> Builder
showb (P IntMap (p e c)
m)
    | IntMap (p e c) -> Bool
forall a. IntMap a -> Bool
IntMap.null IntMap (p e c)
m = Builder
"0"
    | Bool
otherwise = [(Int, p e c)] -> Builder
forall a. TextShow a => a -> Builder
showb ([(Int, p e c)] -> Builder) -> [(Int, p e c)] -> Builder
forall a b. (a -> b) -> a -> b
$ IntMap (p e c) -> [(Int, p e c)]
forall a. IntMap a -> [(Int, a)]
IntMap.toList IntMap (p e c)
m

instance Show (P Int Expression) where
  show :: P Int Expression -> String
show = Text -> String
unpack (Text -> String)
-> (P Int Expression -> Text) -> P Int Expression -> String
forall b c a. (b -> c) -> (a -> b) -> a -> c
. P Int Expression -> Text
forall a. TextShow a => a -> Text
showt

instance TextShow (P Int Expression) where
  showb :: P Int Expression -> Builder
showb (P IntMap Expression
m)
    | IntMap Expression -> Bool
forall a. IntMap a -> Bool
IntMap.null IntMap Expression
m = Builder
"0"
    | Bool
otherwise = [(Int, Expression)] -> Builder
forall a. TextShow a => a -> Builder
showb ([(Int, Expression)] -> Builder) -> [(Int, Expression)] -> Builder
forall a b. (a -> b) -> a -> b
$ IntMap Expression -> [(Int, Expression)]
forall a. IntMap a -> [(Int, a)]
IntMap.toList IntMap Expression
m

instance (Eq a, Num a) => Num (P Int a) where
  (P IntMap a
p) + :: P Int a -> P Int a -> P Int a
+ (P IntMap a
q) = IntMap a -> P Int a
forall {k} (a :: k) b. IntMap b -> P a b
P (IntMap a -> P Int a) -> IntMap a -> P Int a
forall a b. (a -> b) -> a -> b
$ IntMap a -> IntMap a
forall a. (Eq a, Num a) => IntMap a -> IntMap a
filterNonzero (IntMap a -> IntMap a) -> IntMap a -> IntMap a
forall a b. (a -> b) -> a -> b
$ (a -> a -> a) -> IntMap a -> IntMap a -> IntMap a
forall a. (a -> a -> a) -> IntMap a -> IntMap a -> IntMap a
IntMap.unionWith a -> a -> a
forall a. Num a => a -> a -> a
(+) IntMap a
p IntMap a
q

  (P IntMap a
p) * :: P Int a -> P Int a -> P Int a
* (P IntMap a
q) = IntMap a -> P Int a
forall {k} (a :: k) b. IntMap b -> P a b
P (IntMap a -> P Int a) -> IntMap a -> P Int a
forall a b. (a -> b) -> a -> b
$ IntMap a -> IntMap a
forall a. (Eq a, Num a) => IntMap a -> IntMap a
filterNonzero (IntMap a -> IntMap a) -> IntMap a -> IntMap a
forall a b. (a -> b) -> a -> b
$ (IntMap a -> Int -> a -> IntMap a)
-> IntMap a -> IntMap a -> IntMap a
forall a b. (a -> Int -> b -> a) -> a -> IntMap b -> a
IntMap.foldlWithKey' IntMap a -> Int -> a -> IntMap a
accumulate IntMap a
forall a. IntMap a
IntMap.empty IntMap a
p
    where
      accumulate :: IntMap a -> Int -> a -> IntMap a
accumulate IntMap a
m Int
e a
c = (a -> a -> a) -> IntMap a -> IntMap a -> IntMap a
forall a. (a -> a -> a) -> IntMap a -> IntMap a -> IntMap a
IntMap.unionWith a -> a -> a
forall a. Num a => a -> a -> a
(+) IntMap a
m (IntMap a -> IntMap a) -> IntMap a -> IntMap a
forall a b. (a -> b) -> a -> b
$ Int -> a -> IntMap a
multiplyTerm Int
e a
c
      multiplyTerm :: Int -> a -> IntMap a
multiplyTerm Int
e a
c = (Int -> Int) -> IntMap a -> IntMap a
forall a. (Int -> Int) -> IntMap a -> IntMap a
IntMap.mapKeysMonotonic (Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
e) (IntMap a -> IntMap a) -> IntMap a -> IntMap a
forall a b. (a -> b) -> a -> b
$ (a -> a) -> IntMap a -> IntMap a
forall a b. (a -> b) -> IntMap a -> IntMap b
IntMap.map (a -> a -> a
forall a. Num a => a -> a -> a
* a
c) IntMap a
q

  abs :: P Int a -> P Int a
abs = P Int a -> P Int a
forall a. a -> a
id
  signum :: P Int a -> P Int a
signum P Int a
0 = P Int a
0
  signum P Int a
_ = P Int a
1
  fromInteger :: Integer -> P Int a
fromInteger Integer
0 = IntMap a -> P Int a
forall {k} (a :: k) b. IntMap b -> P a b
P IntMap a
forall a. IntMap a
IntMap.empty
  fromInteger Integer
n = IntMap a -> P Int a
forall {k} (a :: k) b. IntMap b -> P a b
P (IntMap a -> P Int a) -> IntMap a -> P Int a
forall a b. (a -> b) -> a -> b
$ Int -> a -> IntMap a
forall a. Int -> a -> IntMap a
IntMap.singleton Int
0 (a -> IntMap a) -> a -> IntMap a
forall a b. (a -> b) -> a -> b
$ Integer -> a
forall a. Num a => Integer -> a
fromInteger Integer
n
  negate :: P Int a -> P Int a
negate (P IntMap a
m) = IntMap a -> P Int a
forall {k} (a :: k) b. IntMap b -> P a b
P (IntMap a -> P Int a) -> IntMap a -> P Int a
forall a b. (a -> b) -> a -> b
$ (a -> a) -> IntMap a -> IntMap a
forall a b. (a -> b) -> IntMap a -> IntMap b
IntMap.map a -> a
forall a. Num a => a -> a
negate IntMap a
m

-- | Get rid of zero coefficients to ensure that zero coefficients do not cause
-- two polynomials represented by an 'IntMap' are not considered different.
filterNonzero :: (Eq a, Num a) => IntMap a -> IntMap a
filterNonzero :: forall a. (Eq a, Num a) => IntMap a -> IntMap a
filterNonzero = (a -> Bool) -> IntMap a -> IntMap a
forall a. (a -> Bool) -> IntMap a -> IntMap a
IntMap.filter (a -> a -> Bool
forall a. Eq a => a -> a -> Bool
/= a
0)

instance (Eq a, Num a) => Polynomial P Int a where
  degree :: P Int a -> Int
degree (P IntMap a
m) = Int -> ((Int, a) -> Int) -> Maybe (Int, a) -> Int
forall b a. b -> (a -> b) -> Maybe a -> b
maybe Int
0 (Int, a) -> Int
forall a b. (a, b) -> a
fst (Maybe (Int, a) -> Int) -> Maybe (Int, a) -> Int
forall a b. (a -> b) -> a -> b
$ IntMap a -> Maybe (Int, a)
forall a. IntMap a -> Maybe (Int, a)
IntMap.lookupMax IntMap a
m
  coefficient :: P Int a -> Int -> a
coefficient (P IntMap a
m) Int
k = a -> Maybe a -> a
forall a. a -> Maybe a -> a
fromMaybe a
0 (Maybe a -> a) -> Maybe a -> a
forall a b. (a -> b) -> a -> b
$ Int -> IntMap a -> Maybe a
forall a. Int -> IntMap a -> Maybe a
IntMap.lookup Int
k IntMap a
m
  leadingCoefficient :: P Int a -> a
leadingCoefficient (P IntMap a
m) = a -> ((Int, a) -> a) -> Maybe (Int, a) -> a
forall b a. b -> (a -> b) -> Maybe a -> b
maybe a
0 (Int, a) -> a
forall a b. (a, b) -> b
snd (Maybe (Int, a) -> a) -> Maybe (Int, a) -> a
forall a b. (a -> b) -> a -> b
$ IntMap a -> Maybe (Int, a)
forall a. IntMap a -> Maybe (Int, a)
IntMap.lookupMax IntMap a
m
  deleteLeadingTerm :: P Int a -> P Int a
deleteLeadingTerm (P IntMap a
m) = IntMap a -> P Int a
forall {k} (a :: k) b. IntMap b -> P a b
P (IntMap a -> P Int a) -> IntMap a -> P Int a
forall a b. (a -> b) -> a -> b
$ IntMap a -> IntMap a
forall a. IntMap a -> IntMap a
IntMap.deleteMax IntMap a
m
  foldTerms :: forall m. Monoid m => (Int -> a -> m) -> P Int a -> m
foldTerms Int -> a -> m
f (P IntMap a
m) = (Int -> a -> m) -> IntMap a -> m
forall m a. Monoid m => (Int -> a -> m) -> IntMap a -> m
IntMap.foldMapWithKey Int -> a -> m
f IntMap a
m
  scale :: a -> P Int a -> P Int a
scale a
0 P Int a
_ = IntMap a -> P Int a
forall {k} (a :: k) b. IntMap b -> P a b
P IntMap a
forall a. IntMap a
IntMap.empty
  scale a
x (P IntMap a
m) = IntMap a -> P Int a
forall {k} (a :: k) b. IntMap b -> P a b
P (IntMap a -> P Int a) -> IntMap a -> P Int a
forall a b. (a -> b) -> a -> b
$ (a -> a) -> IntMap a -> IntMap a
forall a b. (a -> b) -> IntMap a -> IntMap b
IntMap.map (a -> a -> a
forall a. Num a => a -> a -> a
* a
x) IntMap a
m
  power :: Int -> P Int a
power Int
n = IntMap a -> P Int a
forall {k} (a :: k) b. IntMap b -> P a b
P (IntMap a -> P Int a) -> IntMap a -> P Int a
forall a b. (a -> b) -> a -> b
$ Int -> a -> IntMap a
forall a. Int -> a -> IntMap a
IntMap.singleton (Int -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
n) a
1