{-# 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, coarbitrary, )
import Control.Monad (liftM3, )

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


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

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


constant :: Additive.C a => T a
constant = Cons zero zero zero

{-# INLINE evaluate #-}
evaluate :: (Trans.C a) =>
   T a -> a -> Complex.T a
evaluate f x =
   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 =
   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)]


multiply :: (Additive.C a) =>
   T a -> T a -> T a
multiply f g =
   Cons (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
          (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
          (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
          (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 (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
          (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
          (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
      (c0 f)
      (c1 f + (zero +: 2*d))
      (c2 f)

turn :: Ring.C a => a -> T a -> T a
turn d f =
   Cons
      (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
      (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
      (c0 f)
      (Complex.scale k $ c1 f)
      (k^2 * c2 f)


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

fourier (fourier f) = reverse f
-}