```{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE UndecidableInstances #-}
{-# LANGUAGE DeriveDataTypeable #-}
{-|
Module      :  Data.Number.ER.RnToRm.UnitDom.ChebyshevBase.Polynom.Basic
Description :  (internal) polynomial datatype and simple functions
Copyright   :  (c) 2007-2008 Michal Konecny

Maintainer  :  mik@konecny.aow.cz
Stability   :  experimental
Portability :  portable

Internal module for "Data.Number.ER.RnToRm.UnitDom.ChebyshevBase.Polynom".

Definition of the polynomial datatype and simple related functions.
-}
module Data.Number.ER.RnToRm.UnitDom.ChebyshevBase.Polynom.Basic
where

import qualified Data.Number.ER.Real.Base as B
import qualified Data.Number.ER.Real.DomainBox as DBox
import Data.Number.ER.Real.DomainBox (VariableID(..), DomainBox, DomainIntBox)
import Data.Number.ER.Misc

import qualified Data.Map as Map

import Data.Typeable
import Data.Generics.Basics
import Data.Binary

errorModule msg = error \$ "Data.Number.ER.RnToRm.UnitDom.ChebyshevBase.Polynom: " ++ msg

{-|
A polynomial represented by its coefficients it the Chebyshev basis.

The polynomials are never to be used outside the domain @[-1,1]^n@.

All operations are rounded in such a way that the resulting polynomial
is a /point-wise upper or lower bound/ of the exact result.
-}
data ERChebPoly box b =
ERChebPoly
--         Map (MultiSet Int) b
{
chplCoeffs :: (Map.Map (TermKey box) b)
}
deriving (Eq, Ord, Typeable, Data)

type TermKey box = box

instance (Ord a, Binary a, Binary b) => Binary (ERChebPoly a b) where
put (ERChebPoly a) = put a
get = get >>= \a -> return (ERChebPoly a)

chplHasNoNaN p@(ERChebPoly coeffs) =
Map.fold (&&) True \$ Map.map coeffOK coeffs
where
coeffOK c =
not \$ B.isERNaN c

chplHasNoNaNOrInfty p@(ERChebPoly coeffs) =
Map.fold (&&) True \$ Map.map coeffOK coeffs
where
coeffOK = isBounded
isBounded c
| B.isERNaN c = False
| B.isPlusInfinity c = False
| B.isPlusInfinity (- c) = False
| otherwise = True

chplCheck prgLocation p
| chplHasNoNaNOrInfty p = p
| otherwise =
unsafePrint (prgLocation ++ " problem with p = \n" ++ show p) p

chplCompareApprox (ERChebPoly coeffs1) (ERChebPoly coeffs2) =
compare coeffs1 coeffs2

chplConstTermKey :: (DomainBox box varid d) => box
chplConstTermKey = DBox.noinfo

chplIsConstTermKey :: (DomainBox box varid d) => box -> Bool
chplIsConstTermKey = DBox.isNoinfo

chplTermOrder :: (DomainBox box varid d, Num d) => box -> d
chplTermOrder termKey = DBox.fold (+) 0 termKey

chplTermArity :: (DomainBox box varid d) => box -> Int
chplTermArity termKey = length \$ DBox.keys termKey

{-|
Inspect all terms of the polynomial and return the
degree of the highest degree term.
-}
chplGetDegree ::
(B.ERRealBase b, DomainBox box varid d, Num d, Ord d) =>
(ERChebPoly box b) ->
d
chplGetDegree (ERChebPoly coeffs) =
foldl max 0 \$ map chplTermOrder \$ Map.keys coeffs

{-|
If the polynomial is constant, return the constant,
otherwise return Nothing.
-}
chplGetConst ::
(B.ERRealBase b, DomainBox box varid d, Num d, Ord d) =>
(ERChebPoly box b) ->
Maybe b
chplGetConst (ERChebPoly coeffs) =
case Map.keys coeffs of
[key] | chplIsConstTermKey key ->
Just \$ head \$ Map.elems coeffs
_ -> Nothing

-- chplGetArity = length . chplGetVars

chplGetVars (ERChebPoly coeffs) =
DBox.keys \$ foldl DBox.union DBox.noinfo \$ Map.keys coeffs

chplGetGranularity (ERChebPoly coeffs) =
foldl max 0 \$ map B.getGranularity \$ Map.elems coeffs

chplSetMinGranularity gran (ERChebPoly coeffs) =
ERChebPoly \$ Map.map (B.setMinGranularity gran) coeffs

chplSetGranularity gran (ERChebPoly coeffs) =
ERChebPoly \$ Map.map (B.setGranularity gran) coeffs

chplConst ::
(B.ERRealBase b, DomainBox box varid Int, Ord box) =>
b ->
ERChebPoly box b
chplConst val =
(ERChebPoly \$ Map.singleton chplConstTermKey val)

{-|
Make a basic "x" polynomial for a given variable number.
-}
chplVar ::
(B.ERRealBase b, DomainBox box varid Int, Ord box) =>
varid ->
ERChebPoly box b
chplVar varName =
ERChebPoly \$ Map.singleton (DBox.singleton varName 1) 1

{-|
Construct an affine polynomial.
-}
chplAffine ::
(B.ERRealBase b, DomainBox box varid Int, Ord box) =>
b ->
Map.Map varid b ->
ERChebPoly box b
chplAffine at0 varCoeffs =
ERChebPoly \$
Map.insert chplConstTermKey at0 \$
Map.mapKeys (\ i -> DBox.singleton i 1) varCoeffs

--chplRemoveZeroTermsDown, chplRemoveZeroTermsUp ::
--    (B.ERRealBase b, DomainBox box varid Int, Ord box) =>
--    ERChebPoly box b -> ERChebPoly box b
--chplRemoveZeroTermsDown = chplNeg . fst . chplRemoveZeroTerms
--chplRemoveZeroTermsUp = snd . chplRemoveZeroTerms

--chplRemoveZeroTerms ::
--    (B.ERRealBase b, DomainBox box varid Int, Ord box) =>
--    ERChebPoly box b -> (ERChebPoly box b, ERChebPoly box b)
--chplRemoveZeroTerms (ERChebPoly coeffs) =
--    (chplNeg \$ ERChebPoly \$ coeffsNo0T0Down,
--     ERChebPoly \$ coeffsNo0T0Up)
--    where
--    coeffsNo0T0Down =
--        Map.insertWith plusDown chplConstTermKey (- err) coeffsNo0T0
--    coeffsNo0T0Up =
--        Map.insertWith plusUp chplConstTermKey err coeffsNo0T0
--    (coeffsNo0T0, err) =
--        foldl addTermNo0T0 (Map.empty, 0) \$ Map.toList coeffs
--    addTermNo0T0 (prevCoeffs, prevErr) (term, coeff)
--        | coeff == 0 =
--            (prevCoeffs, prevErr)
--        | otherwise =
--        where
--        newTerm =
--            DBox.filter (> 0) term
--        newCoeffs =
--            Map.insert newTerm newCoeffUp prevCoeffs
--        newCoeffUp = prevCoeff + coeff
--        newCoeffDown = prevCoeff `plusDown` coeff
--        prevCoeff =
--            Map.findWithDefault 0 newTerm prevCoeffs
--        newErr = prevErr +  newCoeffUp - newCoeffDown

chplRemoveZeroTermsUp ::
(B.ERRealBase b, DomainBox box varid Int, Ord box) =>
ERChebPoly box b -> ERChebPoly box b
chplRemoveZeroTermsUp (ERChebPoly coeffs) =
ERChebPoly coeffsNo0T0Up
where
coeffsNo0T0Up =
Map.insertWith plusUp chplConstTermKey err coeffsNo0T0
(coeffsNo0T0, err) =
foldl addTermNo0T0 (Map.empty, 0) \$ Map.toList coeffs
| coeff == 0 =
(prevCoeffs, prevErr)
| otherwise =
where
newTerm =
DBox.filter (> 0) term
newCoeffs =
Map.insert newTerm newCoeffUp prevCoeffs
newCoeffUp = prevCoeff + coeff
newCoeffDown = prevCoeff `plusDown` coeff
prevCoeff =
Map.findWithDefault 0 newTerm prevCoeffs
newErr = prevErr +  newCoeffUp - newCoeffDown

--chplRemoveLowCoeffsDown, chplRemoveLowCoeffsUp ::
--    (B.ERRealBase b, DomainBox box varid Int, Ord box) =>
--    b -> ERChebPoly box b -> ERChebPoly box b
--chplRemoveLowCoeffsDown maxCoeff = chplNeg . fst . chplRemoveLowCoeffs maxCoeff
--chplRemoveLowCoeffsUp maxCoeff = snd . chplRemoveLowCoeffs maxCoeff

--chplRemoveLowCoeffs ::
--    (B.ERRealBase b, DomainBox box varid Int, Ord box) =>
--    b -> ERChebPoly box b -> (ERChebPoly box b, ERChebPoly box b)
--chplRemoveLowCoeffs maxCoeff (ERChebPoly coeffs) =
--    (chplNeg \$ ERChebPoly \$ coeffsNoLowDown,
--     ERChebPoly \$ coeffsNoLowUp)
--    where
--    coeffsNoLowDown =
--        Map.insertWith plusDown chplConstTermKey (- err) coeffsNoLow
--    coeffsNoLowUp =
--        Map.insertWith plusUp chplConstTermKey err coeffsNoLow
--    err = sum \$ map abs \$ Map.elems coeffsLow
--    (coeffsLow, coeffsNoLow) =
--        Map.partition (\ c -> abs c < maxCoeff) coeffs

chplCountTerms ::
(B.ERRealBase b, DomainBox box varid Int, Ord box) =>
ERChebPoly box b -> Int
chplCountTerms (ERChebPoly coeffs) =
Map.size coeffs

{------------------ Formatting ------------------------}

instance (B.ERRealBase b, DomainBox box varid Int, Ord box) => Show (ERChebPoly box b)
where
--    show = chplShow 8 False True
show = chplShow 8 False False

{-|
Convert a polynomial to a string representation,
using the ordinary x^n basis.
-}
chplShow ::
(B.ERRealBase b, DomainBox box varid Int, Ord box) =>
Int {- ^ number of decimal digits to show -} ->
Bool {-^ whether to show granularity -} ->
Bool {-^ show the polynomial also in its native Chebyshev basis -} ->
ERChebPoly box b ->
String
chplShow digitsToShow showGranularity showChebyshevBasis (ERChebPoly coeffs)
| showChebyshevBasis = "\n" ++ inChebBasis ++ " = \n" ++ inXBasis
| otherwise = inXBasis
where
inChebBasis =
showCoeffs showTermT \$ Map.filter (\c -> c /= 0) \$ coeffs
inXBasis =
showCoeffs showTermX \$ chebToXBasis coeffs
showCoeffs showTerm coeffs =
concatWith " + " \$ map showTerm \$ Map.toAscList coeffs
showTermT (term, coeff)
| chplIsConstTermKey term = show coeff
| otherwise =
show coeff ++ "*" ++ (concat \$ map showT \$ DBox.toList term)
showTermX (term, coeff)
| chplIsConstTermKey term = showC coeff
| otherwise =
showC coeff ++ "*" ++ (concat \$ map showX \$ DBox.toList term)
showT (var, deg) = "T" ++ show deg ++ "(" ++ showVar var ++ ")"
showX (var, deg) = showVar var ++ "^" ++ show deg
showC = B.showDiGrCmp digitsToShow showGranularity False

{-|
conversion of polynomials from Chebyshev basis to the X^n basis

(not exact - suffering from rounding in the coefficient conversions)
-}
chebToXBasis ::
(B.ERRealBase b, DomainBox box varid Int, Ord box) =>
(Map.Map (TermKey box) b) {-^ polynomial in Chebyshev basis -} ->
(Map.Map (TermKey box) b) {-^ approxition of the equivalent polynomial in X^n basis -}
chebToXBasis coeffs =
Map.filter (\c -> c /= 0) \$
where
Map.unionWith (+) prevXCoeffs \$
Map.map (\ c -> coeff * (fromInteger c)) \$
termXterms term

{-|
conversion of one Chebyshev term to the X^n basis
-}
termXterms ::
(DomainBox box varid Int, Ord box) =>
TermKey box
{-^ a Chebyshev term represented by the Chebyshev degrees
for each variable in the term -} ->
Map.Map (TermKey box) Integer
{-^ the polynomial equivalent to the given Chebyshev term
(using integer coefficients) -}
termXterms term =
allCombinations \$
map (mapSnd \$ \ deg -> chebyXCoeffsLists !! deg) \$
DBox.toList term
where
Map.insertWith (+) term coeffProduct prevMap
where
term =
DBox.fromList \$
filter (\(v,p) -> p > 0) \$
map (\(v,(p,_)) -> (v,p)) varPowerCoeffTriples
coeffProduct =
fromInteger \$
product \$
map (\(_,(_,c)) -> c) varPowerCoeffTriples

{-| Chebyshev polynomials expressed as associative lists power -> coeff -}
chebyXCoeffsLists ::
(Num d1, Enum d1, Num d2, Enum d2) =>
[[(d1, d2)]]
chebyXCoeffsLists =
map convertCoeffs chebyXCoeffs
where
convertCoeffs coeffs =
filter ((/= 0) . snd) \$ zip [0,1..] coeffs

{-| Chebyshev polynomials expressed as lists of integer coefficients for powers 0,1,2... -}
chebyXCoeffs ::
(Num d, Enum d) =>
[[d]]
chebyXCoeffs =
aux
[1] -- T_0(x) = 1
[0,1] -- T_1(x) = x
where
aux tnM2 tnM1 =
tnM2 : (aux tnM1 (newTerm tnM2 tnM1))
newTerm tnM2 tnM1 =
zipWith (-) (0 : (map (*2) tnM1)) (tnM2 ++ [0,0..])
-- T_n(x) = 2 * x * T_{n-1}(x) - T_{n-2}(x)

```