{-# LANGUAGE NoImplicitPrelude #-}
{-
Complex Gaussian bell multiplied with a polynomial.

In order to make this free of @pi@ factors,
we have to choose @recip (sqrt pi)@
as unit for translations and modulations,
for linear factors and in the differentiation.
-}
{-
ToDo:

* In order to avoid the weird @sqrt pi@ factor,
  use a polynomial expression in @pi@.

* sum of multiple bells using Data.Map from exponent polynomial to coefficient polynomial
  use of Algebra object.

* Projective geometry in order to support Dirac impulse.
-}
module MathObj.Gaussian.Polynomial where

import qualified MathObj.Gaussian.Bell as Bell

import qualified MathObj.LaurentPolynomial as LPoly
import qualified MathObj.Polynomial as Poly
import qualified Number.Complex     as Complex

import qualified Algebra.ZeroTestable   as ZeroTestable
import qualified Algebra.Differential   as Differential
import qualified Algebra.Transcendental as Trans
import qualified Algebra.Field          as Field
import qualified Algebra.Real           as Real
import qualified Algebra.Ring           as Ring
import qualified Algebra.Additive       as Additive

import Algebra.Transcendental (pi, )
import Algebra.Ring ((*), )
-- import Algebra.Additive ((+))

import qualified Data.Record.HT as Rec
import qualified Data.List as List
import Data.Function.HT (nest, )
import Data.Eq.HT (equating, )
import Data.List.HT (mapAdjacent, )
import Data.Tuple.HT (forcePair, )

import Test.QuickCheck (Arbitrary, arbitrary, coarbitrary, )
import Control.Monad (liftM2, )

import NumericPrelude
import PreludeBase hiding (reverse, )
-- import Prelude ()


data T a = Cons {bell :: Bell.T a, polynomial :: Poly.T (Complex.T a)}
   deriving (Show)

instance Real.C a => Eq (T a) where
   (==) = equal

{-
The derived Eq is not correct.
We have to combine the amplitude of the bell with the polynomial,
respecting signs and the square root of the bell amplitude.
-}
equal :: Real.C a => T a -> T a -> Bool
equal x y =
   let bx = bell x
       by = bell y
       csign c =
          Complex.real c > 0 ||
          (Complex.real c == 0 && Complex.imag c > 0)
       scaleSqr b =
          map (\c -> (Complex.scale (Bell.amp b) (c^2), csign c)) .
          Poly.coeffs . polynomial
   in  Rec.equal
          (equating Bell.c0 :
           equating Bell.c1 :
           equating Bell.c2 :
           [])
          bx by
       &&
       scaleSqr by x == scaleSqr bx y


instance (Real.C a, Arbitrary a) => Arbitrary (T a) where
   arbitrary =
--      liftM2 Cons arbitrary arbitrary
      liftM2 Cons
         arbitrary
         -- we have to restrict the number of polynomial coefficients,
         -- since with the quadratic time algorithms like fourier and convolve,
         -- in connection with Rational slow down tests too much.
         (fmap (Poly.fromCoeffs . take 5 . Poly.coeffs) arbitrary)
   coarbitrary = undefined



{-# INLINE evaluateSqRt #-}
evaluateSqRt :: (Trans.C a) =>
   T a -> a -> Complex.T a
evaluateSqRt f x =
   Bell.evaluateSqRt (bell f) x *
   Poly.evaluate (polynomial f) (Complex.fromReal $ sqrt pi * x)
{- ToDo: evaluating a complex polynomial for a real argument can be optimized -}


constant :: (Ring.C a) => T a
constant =
   Cons Bell.constant (Poly.const one)

scale :: (Ring.C a) => a -> T a -> T a
scale x f =
   f{polynomial = fmap (Complex.scale x) $ polynomial f}

scaleComplex :: (Ring.C a) => Complex.T a -> T a -> T a
scaleComplex x f =
   f{polynomial = fmap (x*) $ polynomial f}


eigenfunction :: (Field.C a) => Int -> T a
eigenfunction =
   eigenfunctionDifferential

eigenfunction0 :: (Ring.C a) => T a
eigenfunction0 =
   Cons Bell.unit (Poly.fromCoeffs [one])

eigenfunction1 :: (Ring.C a) => T a
eigenfunction1 =
   Cons Bell.unit (Poly.fromCoeffs [zero, one])

eigenfunction2 :: (Field.C a) => T a
eigenfunction2 =
   Cons Bell.unit (Poly.fromCoeffs [-(1/4), zero, one])

eigenfunction3 :: (Field.C a) => T a
eigenfunction3 =
   Cons Bell.unit (Poly.fromCoeffs [zero, -(3/4), zero, one])


eigenfunctionDifferential :: (Field.C a) => Int -> T a
eigenfunctionDifferential n =
   (\f -> f{bell = Bell.unit}) $
   nest n (scale (-1/4) . differentiate) $
   Cons (Bell.Cons one zero zero 2) one

eigenfunctionIterative :: (Field.C a, Real.C a) => Int -> T a
eigenfunctionIterative n =
   fst . head . dropWhile (uncurry (/=)) . mapAdjacent (,) $
   eigenfunctionIteration $
   Cons
      Bell.unit
      (Poly.fromCoeffs $ replicate n zero ++ [one])

eigenfunctionIteration :: (Field.C a) => T a -> [T a]
eigenfunctionIteration =
   iterate (\x ->
      let y = fourier x
          px = polynomial x
          py = polynomial y
          c = last (Poly.coeffs px) / last (Poly.coeffs py)
      in  y{polynomial = fmap (0.5*) (px + fmap (c*) py)})


multiply :: (Ring.C a) =>
   T a -> T a -> T a
multiply f g =
   Cons
      (Bell.multiply (bell f) (bell g))
      (polynomial f * polynomial g)

convolve :: (Field.C a) =>
   T a -> T a -> T a
convolve f g =
   reverse $ fourier $ multiply (fourier f) (fourier g)

{-
We use a Horner like scheme
in order to translate multiplications with @id@
to differentations on the Fourier side.
Quadratic runtime.

fourier (Cons bell (Poly.const a + Poly.shift f))
  = fourier (Cons bell (Poly.const a)) + fourier (Cons bell (Poly.shift f))
  = fourier (Cons bell (Poly.const a)) + C * differentiate (fourier (Cons bell f))
-}
fourier :: (Field.C a) =>
   T a -> T a
fourier f =
   foldr
      (\c p ->
          let q = differentiate p
          in  q{polynomial =
                   Poly.const c +
                   fmap (Complex.scale (1/2) . Complex.quarterLeft) (polynomial q)})
      (Cons (Bell.fourier $ bell f) zero) $
   Poly.coeffs $ polynomial f

{-
Differentiate and divide by @sqrt pi@ in order to stay in a ring.
This way, we do not need to fiddle with pi factors.
-}
differentiate :: (Ring.C a) => T a -> T a
differentiate f =
   f{polynomial =
        Differential.differentiate (polynomial f)
        - Differential.differentiate (Bell.exponentPolynomial (bell f))
           * polynomial f}

{-
snd $ integrate $ differentiate (Cons Bell.unit (Poly.fromCoeffs [7,7,7,7]) :: T Double)
-}
integrate ::
   (Field.C a, ZeroTestable.C a) =>
   T a -> (Complex.T a, T a)
integrate f =
   let fs = Poly.coeffs $ polynomial f
       (ys,~[r]) =
          Poly.divModRev
             {-
             We need the shortening convention of 'zipWith'
             in order to limit the result list,
             we cannot use list instance for (-).
             -}
             (zipWith (-)
                (0 : 0 : diffRev ys)
                (List.reverse fs))
             (List.reverse $ Poly.coeffs $
              Differential.differentiate $
              Bell.exponentPolynomial $ bell f)
   in  forcePair $
       if null fs
         then (zero, f)
         else (r, f{polynomial = Poly.fromCoeffs $ List.reverse ys})

diffRev :: Ring.C a => [a] -> [a]
diffRev xs =
   zipWith (*) xs
      (drop 1 (iterate (subtract 1) (fromIntegral $ length xs)))

translate :: Ring.C a => a -> T a -> T a
translate d =
   translateComplex (Complex.fromReal d)

translateComplex :: Ring.C a => Complex.T a -> T a -> T a
translateComplex d f =
   Cons
      (Bell.translateComplex d $ bell f)
      (Poly.translate d $ polynomial f)

modulate :: Ring.C a => a -> T a -> T a
modulate d f =
   Cons
      (Bell.modulate d $ bell f)
      (polynomial f)

turn :: Ring.C a => a -> T a -> T a
turn d f =
   Cons
      (Bell.turn d $ bell f)
      (polynomial f)

reverse :: Additive.C a => T a -> T a
reverse f =
   Cons
      (Bell.reverse $ bell f)
      (Poly.reverse $ polynomial f)

dilate :: Field.C a => a -> T a -> T a
dilate k f =
   Cons
      (Bell.dilate k $ bell f)
      (Poly.dilate (Complex.fromReal k) $ polynomial f)

shrink :: Ring.C a => a -> T a -> T a
shrink k f =
   Cons
      (Bell.shrink k $ bell f)
      (Poly.shrink (Complex.fromReal k) $ polynomial f)

{-
We could also amplify the polynomial coefficients.
-}
amplify :: Ring.C a => a -> T a -> T a
amplify k f =
   Cons
      (Bell.amplify k $ bell f)
      (polynomial f)


{- |
Approximate a @T a@ using a linear combination of translated @Bell.T a@.
The smaller the unit (e.g. 0.1, 0.01, 0.001)
the better the approximation but the worse the numeric properties.

We cannot put all information into @amp@ of @Bell@,
since @amp@ must be real, but is complex here by construction.
We really need at least signed amplitudes at this place,
since we want to represent differences of Gaussians.
-}
approximateByBells ::
   Field.C a =>
   a -> T a -> [(Complex.T a, Bell.T a)]
approximateByBells unit f =
   let b = bell f
       amps =
          -- approximateByBellsByTranslation
          approximateByBellsAtOnce
             unit
             (Complex.scale (recip (2 * Bell.c2 b)) (Bell.c1 b))
             (recip (2*unit*Bell.c2 b))
             (polynomial f)
   in  zip (LPoly.coeffs amps) $
       map
          (\d -> Bell.translate d b)
          (laurentAbscissas (unit/2) amps)

approximateByBellsAtOnce ::
   Field.C a =>
   a -> Complex.T a -> a -> Poly.T (Complex.T a) -> LPoly.T (Complex.T a)
approximateByBellsAtOnce unit d s p =
   foldr
      (\x amps0 ->
         {-
         Decompose (bell t * (t-d)) = bell t * t - bell t * d
         -}
         let y = fmap (Complex.scale s) amps0
         in  -- \t -> bell t * t
             --    ~   (translate unit bell - translate (-unit) bell) / unit
             LPoly.shift 1 y -
             LPoly.shift (-1) y +
             -- bell t * d
             zipWithAbscissas
                (\t z -> (Complex.fromReal t - d) * z)
                (unit/2) amps0 +
             LPoly.const x)
      (LPoly.fromCoeffs [])
      (Poly.coeffs p)

approximateByBellsByTranslation ::
   Field.C a =>
   a -> Complex.T a -> a -> Poly.T (Complex.T a) -> LPoly.T (Complex.T a)
approximateByBellsByTranslation unit d s p =
   foldr
      (\x amps0 ->
         {-
         Decompose (bell t * (t-d)) = bell t * t - bell t * d
         -}
         let y = fmap (Complex.scale s) amps0
         in  -- \t -> bell t * t
             --    ~   (translate unit bell - translate (-unit) bell) / unit
             LPoly.shift 1 y -
             LPoly.shift (-1) y +
             -- bell t * d
             zipWithAbscissas Complex.scale (unit/2) amps0 +
             LPoly.const x)
      (LPoly.fromCoeffs [])
      (Poly.coeffs $ Poly.translate d p)

zipWithAbscissas ::
   (Ring.C a) =>
   (a -> b -> c) -> a -> LPoly.T b -> LPoly.T c
zipWithAbscissas h unit y =
   LPoly.fromShiftCoeffs (LPoly.expon y) $
   zipWith h
      (laurentAbscissas unit y)
      (LPoly.coeffs y)

laurentAbscissas :: Ring.C a => a -> LPoly.T c -> [a]
laurentAbscissas unit =
   map (\d -> fromIntegral d * unit) .
   iterate (1+) . LPoly.expon


{- No Ring instance for Gaussians
instance (Ring.C a) => Differential.C (T a) where
   differentiate = differentiate
-}

{- laws
differentiate (f*g) =
   (differentiate f) * g + f * (differentiate g)
-}