{-# 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
    License     :  BSD3

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

--{-|
--    Make a univariate polynomial given by a series of coefficients
--    in the Chebyshev basis. 
---}
--chplMakeUnivariate ::
--    (B.ERRealBase b, DomainBox box varid Int, Ord box) =>
--    varid ->
--    [(Int, b)] {-^ list of pairs: degree of Chebyshev polynomial + coefficient -} ->
--    ERChebPoly box b
--chplMakeUnivariate varName powCoeffPairs =
--    ERChebPoly $ Map.fromList $ map encodePow powCoeffPairs
--    where
--    encodePow (pow, coeff) =
--        (DBox.singleton varName pow, coeff)

chplNormaliseDown, chplNormaliseUp ::
    (B.ERRealBase b, DomainBox box varid Int, Ord box) =>
    ERChebPoly box b -> ERChebPoly box b
chplNormaliseUp = snd . chplNormalise
chplNormaliseDown = fst . chplNormalise

chplNormalise ::
    (B.ERRealBase b, DomainBox box varid Int, Ord box) =>
    ERChebPoly box b -> (ERChebPoly box b, ERChebPoly box b)
chplNormalise (ERChebPoly coeffs) =
    (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 =
            (newCoeffs, newErr)
        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 = newCoeffUp - newCoeffDown

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

{-|
    Convert a polynomial to a string representation,
    using the ordinary x^n basis.
-}
chplShow :: 
    (B.ERRealBase b, DomainBox box varid Int, Ord box) =>
    Bool {-^ show the polynomial also in its native Chebyshev basis -} ->
    ERChebPoly box b ->
    String
chplShow showChebyshevBasis (ERChebPoly coeffs) 
    | showChebyshevBasis = "\n" ++ inChebBasis ++ " = \n" ++ inXBasis
    | otherwise = inXBasis
    where
    inChebBasis = 
        showCoeffs showTermT $ 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 8 False 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.foldWithKey addTerm Map.empty coeffs
    where
    addTerm term coeff prevXCoeffs =
        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 =
    foldl addCombination Map.empty $ 
        allCombinations $ 
            map (mapSnd $ \ deg -> chebyXCoeffsLists !! deg) $ 
                DBox.toList term
    where
    addCombination prevMap (varPowerCoeffTriples) =
        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)