{-# LANGUAGE NoImplicitPrelude #-}
{-
Complex translated Gaussian bell curve
with amplitude abstracted away.
-}
module MathObj.Gaussian.Bell where

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

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 Number.Complex ((+:), )
import Algebra.Transcendental (pi, )
import Algebra.Ring ((*), (^), )
import Algebra.Additive ((+), )

import Test.QuickCheck (Arbitrary, arbitrary, )
import Control.Monad (liftM4, )

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


data T a = Cons {amp :: a, c0, c1 :: Complex.T a, c2 :: a}
   deriving (Eq, Show)

instance (Real.C a, Arbitrary a) => Arbitrary (T a) where
   arbitrary =
      liftM4
         (\k a b c -> Cons k a b (1 + abs c))
         arbitrary arbitrary arbitrary arbitrary


constant :: Ring.C a => T a
constant = Cons one zero zero zero

{- |
eigenfunction of 'fourier'
-}
unit :: Ring.C a => T a
unit = Cons one zero zero one

{-# INLINE evaluate #-}
evaluate :: (Trans.C a) =>
   T a -> a -> Complex.T a
evaluate f x =
   Complex.scale
     (sqrt (amp f))
     (Complex.exp $ Complex.scale (-pi) $
      c0 f + Complex.scale x (c1 f) + Complex.fromReal (c2 f * x^2))

evaluateSqRt :: (Trans.C a) =>
   T a -> a -> Complex.T a
evaluateSqRt f x0 =
   Complex.scale
     (sqrt (amp f))
     (let x = sqrt pi * x0
      in  Complex.exp $ negate $
          c0 f + Complex.scale x (c1 f) + Complex.fromReal (c2 f * x^2))

exponentPolynomial :: (Additive.C a) =>
   T a -> Poly.T (Complex.T a)
exponentPolynomial f =
   Poly.fromCoeffs [c0 f, c1 f, Complex.fromReal (c2 f)]


variance :: (Trans.C a) =>
   T a -> a
variance f =
   recip $ c2 f * 2*pi

multiply :: (Ring.C a) =>
   T a -> T a -> T a
multiply f g =
   Cons
      (amp f * amp g)
      (c0 f + c0 g) (c1 f + c1 g) (c2 f + c2 g)


{-
let x=Cons (1+:3) (4+:5) (7::Rational); y=Cons (1+:4) (3+:2) (5::Rational)
-}
convolve :: (Field.C a) =>
   T a -> T a -> T a
convolve f g =
   let s = c2 f + c2 g
       {-
       fd = f1/(2*f2)
       gd = g1/(2*g2)
       c = f2*g2/(f2+g2)

       c*(fd+gd) = (f1*g2+f2*g1)/(2*(f2+g2)) = b/2

       c*(fd+gd)^2 - fd^2*f2 - gd^2*g2
         = f2*g2*(fd+gd)^2/(f2 + g2) - (fd^2*f2 + gd^2*g2)
         = (f2*g2*(fd+gd)^2 - (f2+g2)*(fd^2*f2+gd^2*g2)) / (f2 + g2)
         = (2*f2*g2*fd*gd - (fd^2*f2^2+gd^2*g2^2)) / (f2 + g2)
         = (2*f1*g1 - (f1^2+g1^2)) / (4*(f2 + g2))
         = -(f1 - g1)^2/(4*(f2 + g2))
       -}
   in  Cons
          ((amp f * amp g) / (c2 f + c2 g))
          (c0 f + c0 g
              - Complex.scale (recip (4*s)) ((c1 f - c1 g)^2))
          (Complex.scale (c2 g / s) (c1 f) +
           Complex.scale (c2 f / s) (c1 g))
          (c2 f * c2 g / s)
            -- recip $ recip (c2 f) + recip (c2 g)
{-
   Cons
      (c0 f + c0 g) (c1 f + c1 g)
      (recip $ recip (c2 f) + recip (c2 g))
-}

convolveByTranslation :: (Field.C a) =>
   T a -> T a -> T a
convolveByTranslation f0 g0 =
   let fd = Complex.scale (recip (2 * c2 f0)) $ c1 f0
       gd = Complex.scale (recip (2 * c2 g0)) $ c1 g0
       f1 = translateComplex fd f0
       g1 = translateComplex gd g0
   in  translateComplex (negate $ fd + gd) $
       Cons
          ((amp f0 * amp g0) / (c2 f0 + c2 g0))
          (c0 f1 + c0 g1) zero
          (recip $ recip (c2 f1) + recip (c2 g1))

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

fourier :: (Field.C a) =>
   T a -> T a
fourier f =
   let a = c0 f
       b = c1 f
       c = c2 f
       rc = recip c
   in  Cons
          (amp f / c2 f)
          (Complex.scale (rc/4) (-b^2) + a)
          (Complex.scale rc $ Complex.quarterRight b)
          rc

fourierByTranslation :: (Field.C a) =>
   T a -> T a
fourierByTranslation f =
   translateComplex (Complex.scale (1/2) $ Complex.quarterLeft $ c1 f) $
   Cons (amp f / c2 f) (c0 f) zero (recip $ c2 f)

{-
a + b*x + c*x^2
 = c*(a/c + b/c*x + x^2)
 = c*((x-b/(2*c))^2 + (4*a*c+b^2)/(4*c^2))
 = c*(x-b/(2*c))^2 + (4*a*c+b^2)/(4*c)

fourier ->
   x^2/c - i*b/c*x + (4*a*c+b^2)/(4*c)

fourier (x -> exp(-pi*c*(x-t)^2))
 = fourier $ translate t $ shrink (sqrt c) $ x -> exp(-pi*x^2)
 = modulate t $ dilate (sqrt c) $ fourier $ x -> exp(-pi*x^2)
 = modulate t $ dilate (sqrt c) $ x -> exp(-pi*x^2)
 = modulate t $ x -> exp(-pi*x^2/c)
 = x -> exp(-pi*x^2/c) * exp(-2*pi*i*x*t)
 = x -> exp(-pi*(x^2/c - 2*i*x*t))
-}

{-
b*x + c*x^2
 = c*(b/c*x + x^2)
 = c*((x-br/(2*c))^2 + i*x*bi/c - br^2/(4*c^2))
 = c*(x-br/(2*c))^2 + i*x*bi - br^2/(4*c)

fourier ->
   (x+bi/2)^2/c - i*br/c*(x+bi/2) - br^2/(4*c)
 = (1/c) * ((x+bi/2)^2 - i*br*(x+bi/2) + (br/2)^2)
 = (1/c) * (x^2 - i*b*x + -(br/2)^2 + (bi/2)^2 - i*br*bi/2)
 = (1/c) * (x^2 - i*b*x - (br^2-bi^2+2*br*bi*i)^2 /4)
 = (1/c) * (x^2 - i*b*x - b^2 / 4)
 = (1/c) * (x^2 - i*b*x + (i*b/2)^2)
 = (1/c) * (x - i*b/2)^2

Example:
  (x-b)^2 = b^2 - 2*b*x + x^2
    ->  (- i*2*b*x + x^2)


fourier (x -> exp(-pi*(c*(x-t)^2 + 2*i*m*x)))
 = fourier $ modulate m $ translate t $ shrink (sqrt c) $ x -> exp(-pi*x^2)
 = translate (-m) $ modulate t $ dilate (sqrt c) $ fourier $ x -> exp(-pi*x^2)
 = translate (-m) $ modulate t $ dilate (sqrt c) $ x -> exp(-pi*x^2)
 = translate (-m) $ modulate t $ x -> exp(-pi*x^2/c)
 = translate (-m) $ x -> exp(-pi*x^2/c) * exp(-2*pi*i*x*t)
 = x -> exp(-pi*(x+m)^2/c) * exp(-2*pi*i*(x+m)*t)
 = x -> exp(-pi*((x+m)^2/c - 2*i*(x+m)*t))
-}

{-
fourier (Cons a 0 0) =
  Cons a 0 infinity

fourier (Cons 0 0 c) =
  Cons 0 0 (recip c)

fourier (Cons 0 b 1) =
  Cons 0 (i*b) 1
-}

translate :: Ring.C a => a -> T a -> T a
translate d f =
   let a = c0 f
       b = c1 f
       c = c2 f
   in  Cons
          (amp f)
          (Complex.fromReal (c*d^2) - Complex.scale d b + a)
          (Complex.fromReal (-2*c*d) + b)
          c

translateComplex :: Ring.C a => Complex.T a -> T a -> T a
translateComplex d f =
   let a = c0 f
       b = c1 f
       c = c2 f
   in  Cons
          (amp f)
          (Complex.scale c (d^2) - b*d + a)
          (Complex.scale (-2*c) d + b)
          c

modulate :: Ring.C a => a -> T a -> T a
modulate d f =
   Cons
      (amp f)
      (c0 f)
      (c1 f + (zero +: 2*d))
      (c2 f)

turn :: Ring.C a => a -> T a -> T a
turn d f =
   Cons
      (amp f)
      (c0 f + (zero +: 2*d))
      (c1 f)
      (c2 f)

reverse :: Additive.C a => T a -> T a
reverse f =
   f{c1 = negate $ c1 f}


dilate :: Field.C a => a -> T a -> T a
dilate k f =
   Cons
      (amp f)
      (c0 f)
      (Complex.scale (recip k) $ c1 f)
      (c2 f / k^2)

shrink :: Ring.C a => a -> T a -> T a
shrink k f =
   Cons
      (amp f)
      (c0 f)
      (Complex.scale k $ c1 f)
      (k^2 * c2 f)

amplify :: (Ring.C a) => a -> T a -> T a
amplify k f =
   Cons
      (k^2 * amp f)
      (c0 f)
      (c1 f)
      (c2 f)


{- laws
fourier (convolve f g) = fourier f * fourier g

fourier (fourier f) = reverse f
-}