-- |
-- 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 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 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)

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