{-# LANGUAGE FunctionalDependencies #-}
{-# LANGUAGE GADTs                  #-}
{-# LANGUAGE MultiParamTypeClasses  #-}
-- | Minimil linear algebra lib.
module Math.Regression.Simple.LinAlg (
    -- * Operations
    Add (..),
    Eye (..),
    Mult (..),
    Det (..),
    Inv (..),
    -- * Zeros
    zerosLin,
    zerosQuad,
    optimaQuad,
    -- * Two dimensions
    V2 (..),
    M22 (..),
    SM22 (..),
    -- * Three dimensions
    V3 (..),
    M33 (..),
    SM33 (..),
) where

import Control.DeepSeq (NFData (..))
import Data.Complex    (Complex (..))

-------------------------------------------------------------------------------
-- Classes
-------------------------------------------------------------------------------

-- | Addition
class Add a where
    zero :: a
    add  :: a -> a -> a

-- | Identity
class Eye a where
    eye :: a

-- | Multiplication of different things.
class Eye a => Mult a b c | a b -> c where
    mult :: a -> b -> c

-- | Determinant
class Eye a => Det a where
    det :: a -> Double

-- | Inverse
class Det a => Inv a where
    inv :: a -> a

infixl 6 `add`
infixl 7 `mult`

instance Eye Double where
    eye = 1

instance Add Double where
    zero = 0
    add = (+)

instance Mult Double Double Double where
    mult = (*)

instance Det Double where
    det = id

instance Inv Double where
    inv = recip

-------------------------------------------------------------------------------
-- Zeros
-------------------------------------------------------------------------------

-- | Solve linear equation.
--
-- >>> zerosLin (V2 1 2)
-- -2.0
--
zerosLin :: V2 -> Double
zerosLin (V2 a b) = negate (b / a)

-- | Solve quadratic equation.
--
-- >>> zerosQuad (V3 2 0 (-1))
-- Right (-0.7071067811865476,0.7071067811865476)
--
-- >>> zerosQuad (V3 2 0 1)
-- Left ((-0.0) :+ (-0.7071067811865476),(-0.0) :+ 0.7071067811865476)
--
-- Double root is not treated separately:
--
-- >>> zerosQuad (V3 1 0 0)
-- Right (-0.0,0.0)
--
-- >>> zerosQuad (V3 1 (-2) 1)
-- Right (1.0,1.0)
--
zerosQuad :: V3 -> Either (Complex Double, Complex Double) (Double, Double)
zerosQuad (V3 a b c)
    | delta < 0 = Left ((-b/da) :+ (-sqrtNDelta/da), (-b/da) :+ (sqrtNDelta/da))
    | otherwise = Right ((- b - sqrtDelta) / da, (-b + sqrtDelta) / da)
  where
    delta = b*b - 4 * a * c
    sqrtDelta = sqrt delta
    sqrtNDelta = sqrt (- delta)
    da = 2 * a

-- | Find an optima point.
--
-- >>> optimaQuad (V3 1 (-2) 0)
-- 1.0
--
-- compare to
--
-- >>> zerosQuad (V3 1 (-2) 0)
-- Right (0.0,2.0)
--
optimaQuad :: V3 -> Double
optimaQuad (V3 a b _) = zerosLin (V2 (2 * a) b)

-------------------------------------------------------------------------------
-- 2 dimensions
-------------------------------------------------------------------------------

-- | 2d vector. Strict pair of 'Double's.
--
-- Also used to represent linear polynomial: @V2 a b@  \(= a x + b\).
--
data V2 = V2 !Double !Double
  deriving (Eq, Show)

instance NFData V2 where
    rnf V2 {} = ()

instance Add V2 where
    zero = V2 0 0
    add (V2 x y) (V2 x' y') = V2 (x + x') (y + y')
    {-# INLINE zero #-}
    {-# INLINE add #-}

instance Mult Double V2 V2 where
    mult k (V2 x y) = V2 (k * x) (k * y)
    {-# INLINE mult #-}

-- | 2×2 matrix.
data M22 = M22 !Double !Double !Double !Double
  deriving (Eq, Show)

instance NFData M22 where
    rnf M22 {} = ()

instance Add M22 where
    zero = M22 0 0 0 0
    add (M22 a b c d) (M22 a' b' c' d') = M22 (a + a') (b + b') (c + c') (d + d')
    {-# INLINE zero #-}
    {-# INLINE add #-}

instance Eye M22 where
    eye = M22 1 0 0 1
    {-# INLINE eye #-}

instance Det M22 where det = det2
instance Inv M22 where inv = inv2

instance Mult Double M22 M22 where
    mult k (M22 a b c d) = M22 (k * a) (k * b) (k * c) (k * d)
    {-# INLINE mult #-}

instance Mult M22 V2 V2 where
    mult (M22 a b c d) (V2 u v) = V2 (a * u + b * v) (c * u + d * v)
    {-# INLINE mult #-}

-- | >>> M22 1 2 3 4 `mult` eye @M22
-- M22 1.0 2.0 3.0 4.0
instance Mult M22 M22 M22 where
    mult (M22 a b c d) (M22 x y z w) = M22
        (a * x + b * z) (a * y + b * w)
        (c * x + d * z) (c * y + d * w)
    {-# INLINE mult #-}

det2 :: M22 -> Double
det2 (M22 a b c d) = a * d - b * c
{-# INLINE det2 #-}

inv2 :: M22 -> M22
inv2 m@(M22 a b c d) = M22
    (  d / detm) (- b / detm)
    (- c / detm) (  a / detm)
  where
    detm = det2 m
{-# INLINE inv2 #-}

-- | Symmetric 2x2 matrix.
data SM22 = SM22 !Double !Double !Double
  deriving (Eq, Show)

instance NFData SM22 where
    rnf SM22 {} = ()

instance Add SM22 where
    zero = SM22 0 0 0
    add (SM22 a b d) (SM22 a' b' d') = SM22 (a + a') (b + b') (d + d')
    {-# INLINE zero #-}
    {-# INLINE add #-}

instance Eye SM22 where
    eye = SM22 1 0 1
    {-# INLINE eye #-}

instance Det SM22 where det = detS2
instance Inv SM22 where inv = invS2

instance Mult Double SM22 SM22 where
    mult k (SM22 a b d) = SM22 (k * a) (k * b) (k * d)
    {-# INLINE mult #-}

instance Mult SM22 V2 V2 where
    mult (SM22 a b d) (V2 u v) = V2 (a * u + b * v) (b * u + d * v)
    {-# INLINE mult #-}

detS2 :: SM22 -> Double
detS2 (SM22 a b d) = a * d - b * b
{-# INLINE detS2 #-}

invS2 :: SM22 -> SM22
invS2 m@(SM22 a b d) = SM22
    (  d / detm)
    (- b / detm) (  a / detm)
  where
    detm = detS2 m
{-# INLINE invS2 #-}

-------------------------------------------------------------------------------
-- 3 dimensions
-------------------------------------------------------------------------------

-- | 3d vector. Strict triple of 'Double's.
--
-- Also used to represent quadratic polynomial: @V3 a b c@  \(= a x^2 + b x + c\).
data V3 = V3 !Double !Double !Double
  deriving (Eq, Show)

instance NFData V3 where
    rnf V3 {} = ()

instance Add V3 where
    zero = V3 0 0 0
    add (V3 x y z) (V3 x' y' z') = V3 (x + x') (y + y') (z + z')
    {-# INLINE zero #-}
    {-# INLINE add #-}

instance Mult Double V3 V3 where
    mult k (V3 x y z) = V3 (k * x) (k * y) (k * z)
    {-# INLINE mult #-}

-- | 3×3 matrix.
data M33 = M33
    !Double !Double !Double
    !Double !Double !Double
    !Double !Double !Double
  deriving (Eq, Show)

instance NFData M33 where
    rnf M33 {} = ()

instance Add M33 where
    zero = M33 0 0 0 0 0 0 0 0 0

    add (M33 a b c d e f g h i) (M33 a' b' c' d' e' f' g' h' i') = M33
        (a + a') (b + b') (c + c')
        (d + d') (e + e') (f + f')
        (g + g') (h + h') (i + i')
    {-# INLINE zero #-}
    {-# INLINE add #-}

instance Eye M33 where
    eye = M33 1 0 0
              0 1 0
              0 0 1
    {-# INLINE eye #-}

instance Det M33 where det = det3
instance Inv M33 where inv = inv3

instance Mult Double M33 M33 where
    mult k (M33 a b c d e f g h i) = M33
        (k * a) (k * b) (k * c)
        (k * d) (k * e) (k * f)
        (k * g) (k * h) (k * i)
    {-# INLINE mult #-}

instance Mult M33 V3 V3 where
    mult (M33 a b c
           d e f
           g h i) (V3 u v w) = V3
        (a * u + b * v + c * w)
        (d * u + e * v + f * w)
        (g * u + h * v + i * w)
    {-# INLINE mult #-}

-- TODO: instance Mult M33 M33 M33 where

det3 :: M33 -> Double
det3 (M33 a b c
          d e f
          g h i)
    = a * (e*i-f*h) - d * (b*i-c*h) + g * (b*f-c*e)
{-# INLINE det3 #-}

inv3 :: M33 -> M33
inv3 m@(M33 a b c
            d e f
            g h i)
    = M33 a' b' c'
          d' e' f'
          g' h' i'
  where
    a' = cofactor e f h i / detm
    b' = cofactor c b i h / detm
    c' = cofactor b c e f / detm
    d' = cofactor f d i g / detm
    e' = cofactor a c g i / detm
    f' = cofactor c a f d / detm
    g' = cofactor d e g h / detm
    h' = cofactor b a h g / detm
    i' = cofactor a b d e / detm
    cofactor q r s t = det2 (M22 q r s t)
    detm = det3 m
{-# INLINE inv3 #-}

-- | Symmetric 3×3 matrix.
data SM33 = SM33
    !Double
    !Double !Double
    !Double !Double !Double
  deriving (Eq, Show)

instance NFData SM33 where
    rnf SM33 {} = ()

instance Add SM33 where
    zero = SM33 0 0 0 0 0 0

    add (SM33 a d e g h i) (SM33 a' d' e' g' h' i') = SM33
        (a + a')
        (d + d') (e + e')
        (g + g') (h + h') (i + i')
    {-# INLINE zero #-}
    {-# INLINE add #-}

instance Eye SM33 where
    eye = SM33 1
               0 1
               0 0 1
    {-# INLINE eye #-}

instance Det SM33 where det = detS3
instance Inv SM33 where inv = invS3

instance Mult Double SM33 SM33 where
    mult k (SM33 a d e g h i) = SM33
        (k * a)
        (k * d) (k * e)
        (k * g) (k * h) (k * i)
    {-# INLINE mult #-}

instance Mult SM33 V3 V3 where
    mult (SM33 a
               d e
               g h i) (V3 u v w) = V3
        (a * u + d * v + g * w)
        (d * u + e * v + h * w)
        (g * u + h * v + i * w)
    {-# INLINE mult #-}

detS3 :: SM33 -> Double
detS3 (SM33 a
            d e
            g h i)
    = a * (e*i-h*h) - d * (d*i-g*h) + g * (d*h-g*e)
{-# INLINE detS3 #-}

invS3 :: SM33 -> SM33
invS3 m@(SM33 a
              d e
              g h i)
     = SM33 a'
            d' e'
            g' h' i'
  where
    a' = cofactor e h h i / detm
    d' = cofactor h d i g / detm
    e' = cofactor a g g i / detm
    g' = cofactor d e g h / detm
    h' = cofactor d a h g / detm
    i' = cofactor a d d e / detm
    cofactor q r s t = det2 (M22 q r s t)
    detm = detS3 m
{-# INLINE invS3 #-}


module NeuralNetwork
  () where

import           Control.Monad         (mapM)
import qualified Data.Sequence         as Seq
import           Debug.Trace
import           Numeric.LinearAlgebra ((<>), (><))
import qualified Numeric.LinearAlgebra as LA

type MatrixNum = LA.Matrix Double
type Cache = ((MatrixNum, MatrixNum, MatrixNum), MatrixNum)

data Parameter = Parameter
  { weightsParam :: MatrixNum
  , biasesParam  :: MatrixNum
  }

-- Utils
emptyx :: MatrixNum
emptyx = (0><0) []

emptyCache :: Cache
emptyCache = ((emptyx, emptyx, emptyx), emptyx)

-- Activation functions
relu :: MatrixNum -> MatrixNum
relu = LA.cmap (max 0)

sigmoid :: MatrixNum -> MatrixNum
sigmoid = LA.cmap (\x -> (1 / (1 + exp (-1 * x))))

relu' :: MatrixNum -> MatrixNum
relu' = LA.cmap f
  where
    f x | x < 0 = 0
        | otherwise = 1

sigmoid' :: MatrixNum -> MatrixNum
sigmoid' xw = (sigmoid xw) * (1 - sigmoid xw)

-- | Initialize parameters. Randomize weights and zeroes biases
initParams :: [Int] -> IO [Parameter]
initParams layerDims =
  (fmap . fmap) (\(bx, wx) -> Parameter wx bx) ((zip biases) <$> weights)
    where
      weights = mapM genParams layerDims'
      layerDims' = (zip layerDims (tail layerDims))
      genParams (prevDim, currentDim) = LA.randn currentDim prevDim
      biases =
        fmap
        (\(_, currentDim) -> (currentDim><1) (replicate currentDim 0))
        layerDims'

-- | Propagate forward to next layer. Calulates next a
linearForward axPrev wx bx activationFun =
  let zx = linearForward' axPrev wx bx
      ax = activationFun zx
      cache = ((axPrev, wx, bx), zx)
      linearForward' axPrev wx bx = (wx <> axPrev) + bx :: MatrixNum
  in  (ax, cache)

-- | Complete forward propagation. Calculates final ax
propagateForward :: MatrixNum
                 -> [Parameter]
                 -> (MatrixNum, [Cache])
propagateForward xx params  =
  (outputAx, hiddenCache ++ [outputCache])
    where
      params' = take ((length params) - 1) params
      f (aPrev, _) (Parameter wx bx) = linearForward aPrev wx bx relu
      hidden = scanl f (xx, emptyCache) params'
      hiddenCache = drop 1 (fmap snd hidden)
      (lastHiddenAx, _) = last hidden
      Parameter wxLast bxLast = last params
      (outputAx, outputCache) = linearForward lastHiddenAx wxLast bxLast sigmoid

-- | Propagate backward to previous layer
linearBackward :: MatrixNum -> Cache -> (MatrixNum -> MatrixNum)
               -> (MatrixNum, MatrixNum, MatrixNum)
linearBackward dax ((axPrev, wx, bx), zx) af' =
  (daxPrev, dwx, dbx)
    where
      dzx = dax * af' zx
      m = fromIntegral $ (snd . LA.size) axPrev
      dwx = (dzx <> (LA.tr axPrev)) / m
      sums = fmap sum (LA.toLists dzx)
      layerSize = (fst . LA.size) wx
      dbx = ((layerSize><1) sums) / m
      daxPrev = (LA.tr wx) <> dzx

-- | Complete backpropagation. Computes gradients (dax, dwx, dbx)
-- of each layer
propagateBackward :: MatrixNum -> (LA.Vector Double) -> [Cache]
                  -> [(MatrixNum, MatrixNum, MatrixNum)]
propagateBackward axOut yx caches =
  let l = length caches
      yx' = LA.reshape 1 yx
      daxOut = - (yx' / axOut) - ((1 - yx') / (1 - axOut))
      gradOut = linearBackward daxOut (last caches) sigmoid'
      hiddenCaches = reverse $ take (length caches - 1) caches
      grads = scanl f gradOut hiddenCaches
      f (dax, _, _) cache = linearBackward dax cache relu'
  in  reverse grads

-- | Perform gradient descent. Computes new decreased gradients (dax, dwx, dbx)
updateParameters :: [Parameter]
                 -> [(MatrixNum, MatrixNum, MatrixNum)]
                 -> Double
                 -> [Parameter]
updateParameters params grads eta = fmap fun (zip params grads)
   where fun ((Parameter wx bx), (_, dwx, dbx)) =
           Parameter
             (wx - (LA.scalar eta) * dwx)
             (bx - (LA.scalar eta) * dbx)

train :: MatrixNum
      -> LA.Vector Double
      -> [Int]
      -> Double
      -> Int
      -> IO [Parameter]
train xx yx layerDims eta iterations =
  do
    params <- initParams layerDims
    return $ train' params iterations
  where
    train' params 0 = params
    train' params times =
      let (ax, caches) = propagateForward xx params
          grads = propagateBackward ax yx caches
          newParams = updateParameters params grads eta
      in  train' newParams (times - 1)


module Mbrot where
import           Codec.Picture.Types
import           Coloring
import           Control.Arrow       ((***))
import           Data.Complex
import           ImageOutput

-- |Size parameters
size :: (Int,Int)
size = (1400,800)

-- |Floating point size, calculated
sizef :: (Float,Float)
sizef = (fromIntegral *** fromIntegral) size

-- |Maximum iterations
maxIter :: Int
maxIter = 1000

-- |Color step for integer esacpe coloring
mult :: Float
mult = 255 / fromIntegral maxIter

-- |Color step for continuously colored
multg :: Float
multg = 400 / fromIntegral maxIter

-- |Function for rendering pixel in greyscale
drawGreyscale :: Bool    -- ^Continuously colored
              -> Int
              -> Int
              -> Pixel8
drawGreyscale c x y
    |c =  fromIntegral (round(float * mult))
    |otherwise = fromIntegral(int * round mult)
    where int = mbrot maxIter (((x1/x2)-0.5) :+ (y1/y2))
          float = mbrotContinuous maxIter (((x1/x2)-0.5) :+ (y1/y2))
          x1 = 2 * (fromIntegral x - x2)
          y1 = 2 * (fromIntegral y - y2)
          x2 = fst sizef /2
          y2 = snd sizef /2

-- |Deprecated. Use 'drawColorZoom' instead
drawColor :: Int -> Int -> PixelRGB8
drawColor x y =  colorToPixel $ color (fromIntegral (int*400)/1000)
    where int = mbrot maxIter ((x1/x2) :+ (y1/y2))
          x1 = 1.75 * (fromIntegral x - x2)
          y1 = fromIntegral y - y2
          x2 = fst sizef /2
          y2 = snd sizef /2

-- |Function for rendering pixel in color, with zoom level and focus
drawColorZoom :: Float         -- ^ Zoom level
              -> (Float,Float) -- ^ Focus point (x,y)
              -> Int
              -> Int
              -> PixelRGB8
drawColorZoom zoom (fx, fy) x y =  colorToPixel $ color (float*multg)
    where float = mbrotContinuous maxIter (((x1/x2)+fx) :+ ((y1/y2)+fy))
          x1 = 1.75 * (fromIntegral x - x2) / zoom
          y1 = (fromIntegral y - y2) / zoom
          x2 = fst sizef /2
          y2 = snd sizef /2

-- |Function for drawing color gradient. Tests for correct linear interpolation of color
drawTestGradient :: Int -> Int -> PixelRGB8
drawTestGradient x _ = colorToPixel $ color (fromIntegral x *400/1000)

-- |Function for integer escape time of mbrot set of Re(x) + Im(y)
mbrot :: RealFloat x
     => Int           -- ^ Maximum iterations
     -> Complex x     -- ^ Re(x) + Im(y)
     -> Int
mbrot n = mbrot' maxIter (0 :+ 0)
    where
        mbrot' :: RealFloat x => Int -> Complex x -> Complex x -> Int
        mbrot' n a c
            |magnitude newIter < 65535 && n > 0 = mbrot' (n-1) newIter c
            |otherwise = maxIter - n
                where
                    newIter = a*a + c

-- |Function for real floating escape time of mbrot set of Re(x) + Im(y)
mbrotContinuous :: RealFloat x
                => Int           -- ^ Maximum iterations
                -> Complex x     -- ^ Re(x) + Im(y)
                -> x
mbrotContinuous n = mbrot' maxIter (0 :+ 0)
    where
        mbrot' :: RealFloat x => Int -> Complex x -> Complex x -> x
        mbrot' n a c
            |magnitude newIter < 65535 && n > 0 = mbrot' (n-1) newIter c
            |otherwise = normalized
                where
                    newIter = a*a + c
                    normalized =  if n > 0
                                  then
                                    fromIntegral(maxIter - n +1) - logBase 2 ( logBase 2 (magnitude newIter ^ 2) / 2 )
                                    --fromIntegral (maxIter - n +1) - (log (log (magnitude newIter)))/ log (2.0)
                                  else fromIntegral maxIter


import Data.Bool
import Data.Complex (Complex((:+)), magnitude)

mandelbrot
  :: RealFloat a
  => Complex a -> Complex a
mandelbrot a = iterate ((a +) . (^ 2)) 0 !! 50

main :: IO ()
main =
  mapM_
    putStrLn
    [ [ bool ' ' '*' (2 > magnitude (mandelbrot (x :+ y)))
      | x <- [-2,-1.9685 .. 0.5] ]
    | y <- [1,0.95 .. -1] ]


{-|
Module      : Univariate
Description : Univariate polynomials (notably cyclotomics)
Copyright   : (c) Matthew Amy, 2020
Maintainer  : matt.e.amy@gmail.com
Stability   : experimental
Portability : portable
-}

module Feynman.Algebra.Polynomial.Univariate(
  Univariate,
  Cyclotomic,
  var,
  constant,
  (*|),
  primitiveRoot,
  evaluate
  )
where

import Data.List
import Data.Map(Map)
import qualified Data.Map as Map
import Data.Complex
import Data.Maybe(maybe)

import qualified Feynman.Util.Unicode as Unicode
import Feynman.Algebra.Base
import Feynman.Algebra.Polynomial

{-------------------------------
 Univariate polynomials
 -------------------------------}

-- | Univariate polynomials over the ring 'r'
data Univariate r = Univariate { getCoeffs :: !(Map Integer r) } deriving (Eq, Ord)

instance (Eq r, Num r, Show r) => Show (Univariate r) where
  show = showWithName "x"

instance Degree (Univariate r) where
  degree = maybe 0 (fromIntegral . fst) . Map.lookupMax . getCoeffs

-- | Print a univariate polynomial with a particular variable name
showWithName :: (Eq r, Num r, Show r) => String -> Univariate r -> String
showWithName x p
    | p == 0    = "0"
    | otherwise = intercalate " + " $ map showTerm (Map.assocs $ getCoeffs p)
    where showTerm (expt, a)
            | expt == 0  = show a
            | a    == 1  = showExpt expt
            | a    == -1 = "-" ++ showExpt expt
            | otherwise  = show a ++ showExpt expt
          showExpt expt
            | expt == 1  = x
            | otherwise  = Unicode.sup x expt

instance (Eq r, Num r) => Num (Univariate r) where
  (+)           = add
  (*)           = mult
  negate        = Univariate . Map.map negate . getCoeffs
  abs           = id
  signum        = id
  fromInteger 0 = Univariate Map.empty
  fromInteger i = Univariate $ Map.singleton 0 (fromInteger i)

-- | Normalize a univariate polynomial
normalize :: (Eq r, Num r) => Univariate r -> Univariate r
normalize = Univariate . Map.filter (/=0) . getCoeffs

-- | The unique univariate variable
var :: Num r => Univariate r
var = Univariate $ Map.singleton 1 1

-- | Constant polynomial
constant :: (Eq r, Num r) => r -> Univariate r
constant a
  | a == 0    = Univariate $ Map.empty
  | otherwise = Univariate $ Map.singleton 0 a

-- | Multiply by a scalar
(*|) :: (Eq r, Num r) => r -> Univariate r -> Univariate r
(*|) 0 = \_p -> zero
(*|) a = Univariate . Map.map (a*) . getCoeffs

-- | Add two univariate polynomials
add :: (Eq r, Num r) => Univariate r -> Univariate r -> Univariate r
add p = normalize . Univariate . Map.unionWith (+) (getCoeffs p) . getCoeffs

-- | Multiply two univariate polynomials
mult :: (Eq r, Num r) => Univariate r -> Univariate r -> Univariate r
mult p = normalize . Map.foldrWithKey (\expt a -> add (mulTerm expt a)) 0 . getCoeffs
  where mulTerm expt a = Univariate . Map.mapKeysMonotonic (+ expt) . Map.map (* a) $ getCoeffs p

{-------------------------------
 Cyclotomics
 -------------------------------}
-- | Cyclotomic polynomials over the ring 'r'
data Cyclotomic r = Cyc { getOrder :: !Integer, getPoly :: !(Univariate r) } deriving (Eq, Ord)

instance (Eq r, Num r, Show r) => Show (Cyclotomic r) where
  show p = showWithName (Unicode.sub Unicode.zeta (getOrder p)) $ getPoly p

instance Degree (Cyclotomic r) where
  degree = degree . getPoly

instance (Eq r, Num r) => Num (Cyclotomic r) where
  p + q = reduceOrder $ Cyc m (p' + q')
    where (Cyc m p', Cyc _ q') = unifyOrder p q
  p * q = reduceOrder $ Cyc m (p' * q')
    where (Cyc m p', Cyc _ q') = unifyOrder p q
  negate (Cyc m p) = Cyc m $ negate p
  abs    p = p
  signum p = p
  fromInteger i = Cyc 1 (fromInteger i)

-- | Unify the order of two cyclotomics
unifyOrder :: Cyclotomic r -> Cyclotomic r -> (Cyclotomic r, Cyclotomic r)
unifyOrder (Cyc n p) (Cyc m q)
  | n == m    = (Cyc n p, Cyc m q)
  | otherwise = (Cyc r p', Cyc r q')
  where r  = lcm n m
        p' = Univariate . Map.mapKeysMonotonic ((r `div` n) *) $ getCoeffs p
        q' = Univariate . Map.mapKeysMonotonic ((r `div` m) *) $ getCoeffs q

-- | Rewrite the cyclotomic in lowest order
reduceOrder :: (Eq r, Num r) => Cyclotomic r -> Cyclotomic r
reduceOrder (Cyc m c) = Cyc m' c'
  where m' = m `div` d
        c' = Univariate . Map.mapKeysMonotonic (`div` d) $ getCoeffs c
        d  = foldr gcd m . Map.keys $ getCoeffs c

-- | Construct the cyclotomic polynomial \(\zeta_i\)
primitiveRoot :: Num r => Integer -> Cyclotomic r
primitiveRoot i = Cyc i var

-- | Convert to a complex number
evaluate :: (Real r, RealFloat f) => Cyclotomic r -> Complex f
evaluate (Cyc m p) = Map.foldrWithKey f (0.0 :+ 0.0) $ getCoeffs p
  where f root coeff = (mkPolar (realToFrac coeff) (expnt root) +)
        expnt root   = 2.0*pi*(fromInteger root)/(fromInteger m)


import Graphics.Gloss
import Graphics.Gloss.Juicy
import Codec.Picture
import Codec.Picture.RGBA8
import Data.Complex

width = 512
height = width `div` 2

mandelbrotness :: Complex Float -> Complex Float -> Float
mandelbrotness c s = (fromIntegral $ length $ takeWhile (\c -> magnitude c <= 2) $ take maxIter $ z c) / (fromIntegral maxIter) where
  z c = iterate (\x -> x^2 + c) s
  maxIter = 32

image :: Float -> Image PixelRGB8
image time = generateImage (\x y -> let [x',y',w,h] = map fromIntegral [x,y,width,height] in mandelbrotColor ((4*x'/w - 2.5) :+ (2*y'/h - 1)) (0.1*time :+ 0) ) width height where
  mandelbrotColor c s = PixelRGB8 (floor $ 255 * mandelbrotness c s) 120 120

main :: IO ()
main = animate (InWindow "Mandelbrot" (width,height) (0,0)) black (fromImageRGB8 . image)


{-# LANGUAGE GADTs #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE DataKinds, PolyKinds #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE KindSignatures #-}
{-# LANGUAGE FlexibleContexts, FlexibleInstances #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE StandaloneDeriving #-}
{-# LANGUAGE ConstraintKinds #-}
{-# LANGUAGE MultiParamTypeClasses #-}

module TensorHMatrix where

import Nats
import Numeric.LinearAlgebra hiding ((!))
import Data.Vector.Storable
import Data.Default
import Data.Singletons.Prelude

import Data.Vinyl

import Prelude hiding (zipWith)

type Usable a = (Element a, Num a, Numeric a, Num (Vector a), Container Vector a)

type IntegralK (p :: KProxy k) = (SingKind p, Integral (DemoteRep p))
type IntegralN (n :: k) = IntegralK ('KProxy :: KProxy k)

natVal' :: (Num a, IntegralN n) => Sing n -> a
natVal' = fromIntegral . fromSing

data Tensor (dims :: [k]) a where
  Scalar :: a -> Tensor '[] a
  Vector :: Vector a -> Tensor '[n] a
  Matrix :: Matrix a -> Tensor '[n, m] a

deriving instance (Show a, Element a) => Show (Tensor dims a)

instance (Default a) => Default (Tensor '[] a) where
  def = Scalar def

fill1 :: forall a n. (SingI n, IntegralN n, Usable a) => a -> Tensor '[n] a
fill1 a = Vector $ konst a (natVal' (sing::Sing n))

instance (SingI n, IntegralN n, Default a, Usable a) => Default (Tensor '[n] a) where
  def = fill1 def

instance (SingI n, IntegralN n, SingI m, Default a, Usable a) => Default (Tensor '[n, m] a) where
  def = Matrix $ konst def (natVal' (sing::Sing n), natVal' (sing::Sing m))

--type instance IndexOf (Tensor l) = Rec LT l
--type instance IndexOf (Tensor l) = Rec (Const Int) l
--type instance IndexOf (Tensor l) = Vec (Length l) Int

{- Fails backwards fundep :(
instance Indexable (Tensor (n ': l) a) (Tensor l a) where
  Matrix m ! i = Vector (m ! i)
  Vector v ! i = Scalar (v ! i)
-}

select :: Storable a => Int -> Tensor '[n] a -> a
select i (Vector v) = v ! i

-- pretty much all of these instances should be automatically derivable

instance Num a => Num (Tensor '[] a) where
  Scalar a + Scalar b = Scalar (a + b)
  Scalar a - Scalar b = Scalar (a - b)
  Scalar a * Scalar b = Scalar (a * b)
  negate (Scalar a) = Scalar (negate a)
  abs (Scalar a) = Scalar (abs a)
  signum (Scalar a) = Scalar (signum a)
  fromInteger n = Scalar (fromInteger n)

instance (SingI n, IntegralN n, Usable a) => Num (Tensor '[n] a) where
  Vector a + Vector b = Vector (a + b)
  Vector a - Vector b = Vector (a - b)
  Vector a * Vector b = Vector (a * b)
  negate (Vector a) = Vector (negate a)
  abs (Vector a) = Vector (abs a)
  signum (Vector a) = Vector (signum a)
  fromInteger n = Vector $ konst (fromInteger n) (natVal' (sing::Sing n))

instance (SingI n, IntegralN n, SingI m, Usable a) => Num (Tensor '[n, m] a) where
  Matrix a + Matrix b = Matrix (a + b)
  Matrix a - Matrix b = Matrix (a - b)
  Matrix a * Matrix b = Matrix (a * b)
  negate (Matrix a) = Matrix (negate a)
  abs (Matrix a) = Matrix (abs a)
  signum (Matrix a) = Matrix (signum a)
  fromInteger n = Matrix $ konst (fromInteger n) (natVal' (sing::Sing n), natVal' (sing::Sing m))

instance Fractional a => Fractional (Tensor '[] a) where
  Scalar a / Scalar b = Scalar (a / b)
  recip (Scalar a) = Scalar (recip a)
  fromRational r = Scalar (fromRational r)

instance (SingI n, IntegralN n, Numeric a, Fractional a, Fractional (Vector a)) => Fractional (Tensor '[n] a) where
  Vector a / Vector b = Vector (a / b)
  recip (Vector a) = Vector (recip a)
  fromRational r = Vector $ konst (fromRational r) (natVal' (sing::Sing n))

instance Floating a => Floating (Tensor '[] a) where
  pi = Scalar pi
  exp (Scalar a) = Scalar (exp a)
  log (Scalar a) = Scalar (log a)
  sqrt (Scalar a) = Scalar (sqrt a)
  Scalar a ** Scalar b = Scalar (a ** b)
  logBase (Scalar a) (Scalar b) = Scalar (logBase a b)
  sin (Scalar a) = Scalar (sin a)
  cos (Scalar a) = Scalar (cos a)
  tan (Scalar a) = Scalar (tan a)
  asin (Scalar a) = Scalar (asin a)
  acos (Scalar a) = Scalar (acos a)
  atan (Scalar a) = Scalar (atan a)
  sinh (Scalar a) = Scalar (sinh a)
  cosh (Scalar a) = Scalar (cosh a)
  tanh (Scalar a) = Scalar (tanh a)
  asinh (Scalar a) = Scalar (asinh a)
  acosh (Scalar a) = Scalar (acosh a)
  atanh (Scalar a) = Scalar (atanh a)

instance (SingI n, IntegralN n, Usable a, Floating a, Floating (Vector a)) => Floating (Tensor '[n] a) where
  pi = fill1 pi
  exp (Vector a) = Vector (exp a)
  log (Vector a) = Vector (log a)
  sqrt (Vector a) = Vector (sqrt a)
  Vector a ** Vector b = Vector (a ** b)
  logBase (Vector a) (Vector b) = Vector (logBase a b)
  sin (Vector a) = Vector (sin a)
  cos (Vector a) = Vector (cos a)
  tan (Vector a) = Vector (tan a)
  asin (Vector a) = Vector (asin a)
  acos (Vector a) = Vector (acos a)
  atan (Vector a) = Vector (atan a)
  sinh (Vector a) = Vector (sinh a)
  cosh (Vector a) = Vector (cosh a)
  tanh (Vector a) = Vector (tanh a)
  asinh (Vector a) = Vector (asinh a)
  acosh (Vector a) = Vector (acosh a)
  atanh (Vector a) = Vector (atanh a)

tmap :: (Container Vector a, Num a, Element b) => (a -> b) -> Tensor dims a -> Tensor dims b
tmap f (Scalar a) = Scalar $ f a
tmap f (Vector v) = Vector $ cmap f v
tmap f (Matrix m) = Matrix $ cmap f m

tZipWith :: (Storable a, Storable b, Storable c) => (a -> b -> c) -> Tensor '[n] a -> Tensor '[n] b -> Tensor '[n] c
tZipWith f (Vector a) (Vector b) = Vector (zipWith f a b)

transpose :: (Numeric a) => Tensor '[n, m] a -> Tensor '[m, n] a
transpose (Matrix m) = Matrix (tr' m)

asCol :: Storable a => Tensor '[n] a -> Tensor '[n, FromInteger 1] a
asCol (Vector v) = Matrix $ asColumn v

asRow' :: Storable a => Tensor '[n] a -> Tensor '[FromInteger 1, n] a
asRow' (Vector v) = Matrix $ asRow v

dot :: (Storable a, Numeric a) => Tensor '[n] a -> Tensor '[n] a -> a
dot (Vector a) (Vector b) = a <.> b

mv :: Numeric a => Tensor '[n, m] a -> Tensor '[m] a -> Tensor '[n] a
mv (Matrix m) (Vector v) = Vector $ m #> v

mm :: Numeric a => Tensor '[n, m] a -> Tensor '[m, k] a -> Tensor '[n, k] a
mm (Matrix m1) (Matrix m2) = Matrix $ m1 <> m2

gradDot :: Numeric a => Tensor '[n] a -> Tensor '[n] a -> a -> (Tensor '[n] a, Tensor '[n] a)
gradDot (Vector a) (Vector b) g = (Vector $ scale g b, Vector $ scale g a)

gradMV :: Numeric a => Tensor '[n, m] a -> Tensor '[m] a -> Tensor '[n] a -> (Tensor '[n, m] a, Tensor '[m] a)
gradMV m v g = (mm (asCol g) (asRow' v), mv (transpose m) g)

gradMM :: Numeric a => Tensor '[n, m] a -> Tensor '[m, k] a -> Tensor '[n, k] a -> (Tensor '[n, m] a, Tensor '[m, k] a)
gradMM a b g = (mm g (transpose b), mm (transpose a) g)

oneHot :: forall a n. (SingI n, IntegralN n, Usable a) => Int -> Tensor '[n] a
oneHot m = Vector $ (konst 0 (natVal' (sing::Sing n))) // [(m, 1)]

gradSelect :: forall a n. (SingI n, IntegralN n, Usable a) => Int -> Tensor '[n] a -> a -> Tensor '[n] a
gradSelect i _ a = Vector $ (konst 0 (natVal' (sing::Sing n))) // [(i, a)]


{-# LANGUAGE FlexibleContexts #-}
module SentimentAnalysis where

import Numeric.LinearAlgebra as LA
import Data.Bifunctor (Bifunctor(bimap))
import Data.List.Split (splitOn)

hingeLossSingle :: (Ord a, Numeric a) => Vector a -> a -> Vector a -> a -> a
hingeLossSingle x y theta theta0 = max 0 (1 - y * (theta LA.<.> x + theta0))

hingeLossFull :: (Numeric a, Ord a, Num (Vector a), Fractional a) => Matrix a -> Vector a -> Vector a -> a -> a
hingeLossFull x y theta theta0 = loss
    where
        theta0s = size y |> repeat theta0
        ones = size y |> repeat 1.0
        agreement = y * (theta LA.<# LA.tr' x + theta0s)
        losses = LA.cmap (max 0) (ones - agreement)
        loss = sumElements losses / fromIntegral (size losses)

perceptronSingleStepUpdate :: (Ord b, Numeric b, Num (Vector b)) => Vector b -> b -> Vector b -> b -> (Vector b, b)
perceptronSingleStepUpdate x y theta theta0 =
    if y * (theta LA.<.> x + theta0) <= 0
        then (theta + (scalar y * x), theta0 + y)
        else (theta, theta0)

perceptronUpdate :: (Ord t, Numeric t, Num (Vector t)) => [Vector t] -> [t] -> Vector t -> t -> [Int] -> (Vector t, t)
perceptronUpdate _ _ theta theta0 [] = (theta, theta0)
perceptronUpdate xs ys theta theta0 (z:zs) =
    let (theta', theta0') = perceptronSingleStepUpdate (xs !! z) (ys !! z) theta theta0
    in perceptronUpdate xs ys theta' theta0' zs

perceptron :: (Ord t, Numeric t, Num (Vector t)) => Matrix t -> Vector t -> Int -> [Int] -> (Vector t, t)
perceptron x y t indices = iterate step (cols x |> repeat 0, 0) !! t
    where step (theta, theta0) = perceptronUpdate (toRows x) (toList y) theta theta0 indices

averagePerceptronUpdate :: (Ord t, Numeric t, Num (Vector t))
    => [Vector t] -> [t] -> Vector t -> t -> [Int] -> [(Vector t, t)]
averagePerceptronUpdate _ _ theta theta0 [] = []
averagePerceptronUpdate xs ys theta theta0 (z:zs) =
    let (theta', theta0') = perceptronSingleStepUpdate (xs !! z) (ys !! z) theta theta0
    in (theta', theta0') : averagePerceptronUpdate xs ys theta' theta0' zs

averagePerceptronList :: (Ord t, Numeric t, Num (Vector t)) => Matrix t -> Vector t -> [Int] -> [[(Vector t, t)]]
averagePerceptronList x y indices = iterate step [(cols x |> repeat 0, 0)]
    where
        step xs = averagePerceptronUpdate (toRows x) (toList y) theta theta0 indices
            where (theta, theta0) = last xs

averagePerceptron :: (Fractional t, Ord t, Numeric t, Num (Vector t))
                        => Matrix t -> Vector t -> Int -> [Int] -> (Vector t, t)
averagePerceptron x y t indices = (bimap ((/ scalar n) . sum) ((/ n) . sum) . unzip . concat) xss
    where
        n = fromIntegral $ t * rows x
        xss = drop 1 $ take (t + 1) $ averagePerceptronList x y indices

pegasosSingleStepUpdate :: (Ord e, Numeric e, Num (Vector e)) => Vector e -> e -> e -> e -> Vector e -> e -> (Vector e, e)
pegasosSingleStepUpdate x y lambda eta theta theta0 =
    if y * (theta LA.<.> x + theta0) <= 1
        then
            (scalar (1 - eta * lambda) * theta + scalar (eta * y) * x, theta0 + eta * y)
        else
            (scalar (1 - eta * lambda) * theta, theta0)

pegasosUpdate :: (Ord t, Numeric t, Num (Vector t), Floating t)
                    => [Vector t] -> [t] -> t -> t -> Vector t -> t -> [Int] -> (Vector t, t, t)
pegasosUpdate _ _ lambda i theta theta0 [] = (theta, theta0, i)
pegasosUpdate xs ys lambda i theta theta0 (z:zs) =
    let (theta', theta0') = pegasosSingleStepUpdate (xs !! z) (ys !! z) lambda (1 / sqrt i) theta theta0
    in pegasosUpdate xs ys lambda (i + 1) theta' theta0' zs

pegasos :: (Ord t, Numeric t, Floating t, Num (Vector t)) => Matrix t -> Vector t -> Int -> t -> [Int] -> (Vector t, t)
pegasos x y t lambda indices = (theta', theta0')
    where
        (theta', theta0', i') = iterate step (cols x |> repeat 0, 0, 1.0) !! t
        step (theta, theta0, i) = pegasosUpdate (toRows x) (toList y) lambda i theta theta0 indices

main :: IO ()
main = do
    dta <- loadMatrix "toy_data.tsv"
    orderDta <- readFile "200.txt"

    let t = 10
    let lambda = 0.2
    let x = dta ¿ [1, 2]
    let y = flatten $ dta ¿ [0]
    let indices = map read $ splitOn "," orderDta

    let (theta, theta0) = perceptron x y t indices
    putStrLn "Perceptron:"
    print theta
    print theta0

    let (theta, theta0) = averagePerceptron x y t indices
    putStrLn "Average Perceptron:"
    print theta
    print theta0

    let (theta, theta0) = pegasos x y t lambda indices
    putStrLn "Pegasos:"
    print theta
    print theta0


{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE DeriveDataTypeable, DeriveGeneric #-}
-- |
-- Module    : Statistics.Distribution.Hypergeometric
-- Copyright : (c) 2009 Bryan O'Sullivan
-- License   : BSD3
--
-- Maintainer  : bos@serpentine.com
-- Stability   : experimental
-- Portability : portable
--
-- The Hypergeometric distribution.  This is the discrete probability
-- distribution that measures the probability of /k/ successes in /l/
-- trials, without replacement, from a finite population.
--
-- The parameters of the distribution describe /k/ elements chosen
-- from a population of /l/, with /m/ elements of one type, and
-- /l/-/m/ of the other (all are positive integers).

module Statistics.Distribution.Hypergeometric
    (
      HypergeometricDistribution
    -- * Constructors
    , hypergeometric
    , hypergeometricE
    -- ** Accessors
    , hdM
    , hdL
    , hdK
    ) where

import Control.Applicative
import Data.Aeson           (FromJSON(..), ToJSON, Value(..), (.:))
import Data.Binary          (Binary(..))
import Data.Data            (Data, Typeable)
import GHC.Generics         (Generic)
import Numeric.MathFunctions.Constants (m_epsilon,m_neg_inf)
import Numeric.SpecFunctions (choose,logChoose)

import qualified Statistics.Distribution as D
import Statistics.Internal

data HypergeometricDistribution = HD {
      hdM :: {-# UNPACK #-} !Int
    , hdL :: {-# UNPACK #-} !Int
    , hdK :: {-# UNPACK #-} !Int
    } deriving (Eq, Typeable, Data, Generic)

instance Show HypergeometricDistribution where
  showsPrec i (HD m l k) = defaultShow3 "hypergeometric" m l k i
instance Read HypergeometricDistribution where
  readPrec = defaultReadPrecM3 "hypergeometric" hypergeometricE

instance ToJSON HypergeometricDistribution
instance FromJSON HypergeometricDistribution where
  parseJSON (Object v) = do
    m <- v .: "hdM"
    l <- v .: "hdL"
    k <- v .: "hdK"
    maybe (fail $ errMsg m l k) return $ hypergeometricE m l k
  parseJSON _ = empty

instance Binary HypergeometricDistribution where
  put (HD m l k) = put m >> put l >> put k
  get = do
    m <- get
    l <- get
    k <- get
    maybe (fail $ errMsg m l k) return $ hypergeometricE m l k

instance D.Distribution HypergeometricDistribution where
    cumulative = cumulative

instance D.DiscreteDistr HypergeometricDistribution where
    probability    = probability
    logProbability = logProbability

instance D.Mean HypergeometricDistribution where
    mean = mean

instance D.Variance HypergeometricDistribution where
    variance = variance

instance D.MaybeMean HypergeometricDistribution where
    maybeMean = Just . D.mean

instance D.MaybeVariance HypergeometricDistribution where
    maybeStdDev   = Just . D.stdDev
    maybeVariance = Just . D.variance

instance D.Entropy HypergeometricDistribution where
  entropy = directEntropy

instance D.MaybeEntropy HypergeometricDistribution where
  maybeEntropy = Just . D.entropy

variance :: HypergeometricDistribution -> Double
variance (HD m l k) = (k' * ml) * (1 - ml) * (l' - k') / (l' - 1)
  where m' = fromIntegral m
        l' = fromIntegral l
        k' = fromIntegral k
        ml = m' / l'

mean :: HypergeometricDistribution -> Double
mean (HD m l k) = fromIntegral k * fromIntegral m / fromIntegral l

directEntropy :: HypergeometricDistribution -> Double
directEntropy d@(HD m _ _)
  = negate . sum
  $ takeWhile (< negate m_epsilon)
  $ dropWhile (not . (< negate m_epsilon))
    [ let x = probability d n in x * log x | n <- [0..m]]

hypergeometric :: Int               -- ^ /m/
               -> Int               -- ^ /l/
               -> Int               -- ^ /k/
               -> HypergeometricDistribution
hypergeometric m l k
  = maybe (error $ errMsg m l k) id $ hypergeometricE m l k

hypergeometricE :: Int               -- ^ /m/
                -> Int               -- ^ /l/
                -> Int               -- ^ /k/
                -> Maybe HypergeometricDistribution
hypergeometricE m l k
  | not (l > 0)            = Nothing
  | not (m >= 0 && m <= l) = Nothing
  | not (k > 0  && k <= l) = Nothing
  | otherwise              = Just (HD m l k)

errMsg :: Int -> Int -> Int -> String
errMsg m l k
  =  "Statistics.Distribution.Hypergeometric.hypergeometric: "
  ++ "m=" ++ show m
  ++ "l=" ++ show l
  ++ "k=" ++ show k
  ++ " should hold: l>0 & m in [0,l] & k in (0,l]"

-- Naive implementation
probability :: HypergeometricDistribution -> Int -> Double
probability (HD mi li ki) n
  | n < max 0 (mi+ki-li) || n > min mi ki = 0
    -- No overflow
  | li < 1000 = choose mi n * choose (li - mi) (ki - n)
              / choose li ki
  | otherwise = exp $ logChoose mi n
                    + logChoose (li - mi) (ki - n)
                    - logChoose li ki

logProbability :: HypergeometricDistribution -> Int -> Double
logProbability (HD mi li ki) n
  | n < max 0 (mi+ki-li) || n > min mi ki = m_neg_inf
  | otherwise = logChoose mi n
              + logChoose (li - mi) (ki - n)
              - logChoose li ki

cumulative :: HypergeometricDistribution -> Double -> Double
cumulative d@(HD mi li ki) x
  | isNaN x      = error "Statistics.Distribution.Hypergeometric.cumulative: NaN argument"
  | isInfinite x = if x > 0 then 1 else 0
  | n <  minN    = 0
  | n >= maxN    = 1
  | otherwise    = D.sumProbabilities d minN n
  where
    n    = floor x
    minN = max 0 (mi+ki-li)
    maxN = min mi ki


------------------------------------------------
-- |
-- Module    :  Numeric.Matrix.Integral
-- Copyright :  (c) Jun Yoshida 2019
-- License   :  BSD3
--
-- Operations on integer matrices
--
------------------------------------------------

module Numeric.Matrix.Integral (
  -- Hermite normal form
  hermiteNF,
  -- Smith normal form
  smithNF,
  smithRep,
  ) where

import Control.Monad
import Control.Monad.ST (ST, runST)
import Control.Monad.Loops (whileM_)
import Data.STRef

import qualified Numeric.LinearAlgebra as LA

--import Numeric.Matrix.Integral.HNFLLL
--import Numeric.Matrix.Integral.SmithNF
import Numeric.Matrix.Integral.NormalForms

-- Compute the submatrix containing all the non-zero diagonal entries
extractMaxNZDiag :: LA.Matrix LA.Z -> (Int,LA.Matrix LA.Z)
extractMaxNZDiag mx = runST $ do
  dRef <- newSTRef 0
  let p d = d < uncurry min (LA.size mx) && (mx LA.! d LA.! d) /= 0
  whileM_ (p <$> readSTRef dRef) $ modifySTRef' dRef (+1)
  d <- readSTRef dRef
  return (d,LA.subMatrix (0,0) (d,d) mx)

-- D = P <> A <> Q
-- | Return (d,ker,im) where
-- d: the diagonal entries of Smith normal form
-- ker: basis for the kernel
-- im: basis for the image
kerImOf :: LA.Matrix LA.Z -> (LA.Vector LA.Z,[LA.Vector LA.Z],[LA.Vector LA.Z])
kerImOf mt =
  let (ul,h,ur) = smithNF mt
      (ull,_,ulr) = smithNF ul
      (rk,dmt) = extractMaxNZDiag h
      ker' = ur LA.¿ [rk..(LA.cols ur-1)]
      (_,kert) = hermiteNF (LA.tr' ker')
      im' = (ulr LA.<> ull LA.<> h) LA.¿ [0..rk-1]
      (_,imt) = hermiteNF (LA.tr' im')
  in (LA.takeDiag dmt, LA.toRows kert, LA.toRows imt)


{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE DeriveDataTypeable, DeriveGeneric #-}
-- |
-- Module    : Statistics.Distribution.Lognormal
-- Copyright : (c) 2020 Ximin Luo
-- License   : BSD3
--
-- Maintainer  : infinity0@pwned.gg
-- Stability   : experimental
-- Portability : portable
--
-- The Weibull distribution.  This is a continuous probability
-- distribution that describes the occurrence of a single event whose
-- probability changes over time, controlled by the shape parameter.

module Statistics.Distribution.Weibull
    (
      WeibullDistribution
      -- * Constructors
    , weibullDistr
    , weibullDistrErr
    , weibullStandard
    , weibullDistrApproxMeanStddevErr
    ) where

import Control.Applicative
import Data.Aeson            (FromJSON(..), ToJSON, Value(..), (.:))
import Data.Binary           (Binary(..))
import Data.Data             (Data, Typeable)
import GHC.Generics          (Generic)
import Numeric.MathFunctions.Constants (m_eulerMascheroni)
import Numeric.SpecFunctions (expm1, log1p, logGamma)
import qualified Data.Vector.Generic as G

import qualified Statistics.Distribution as D
import qualified Statistics.Sample as S
import Statistics.Internal

-- | The Weibull distribution.
data WeibullDistribution = WD {
      wdShape  :: {-# UNPACK #-} !Double
    , wdLambda :: {-# UNPACK #-} !Double
    } deriving (Eq, Typeable, Data, Generic)

instance Show WeibullDistribution where
  showsPrec i (WD k l) = defaultShow2 "weibullDistr" k l i
instance Read WeibullDistribution where
  readPrec = defaultReadPrecM2 "weibullDistr" $
    (either (const Nothing) Just .) . weibullDistrErr

instance ToJSON WeibullDistribution
instance FromJSON WeibullDistribution where
  parseJSON (Object v) = do
    k <- v .: "wdShape"
    l <- v .: "wdLambda"
    either fail return $ weibullDistrErr k l
  parseJSON _ = empty

instance Binary WeibullDistribution where
  put (WD k l) = put k >> put l
  get = do
    k <- get
    l <- get
    either fail return $ weibullDistrErr k l

instance D.Distribution WeibullDistribution where
  cumulative      = cumulative
  complCumulative = complCumulative

instance D.ContDistr WeibullDistribution where
  logDensity    = logDensity
  quantile      = quantile
  complQuantile = complQuantile

instance D.MaybeMean WeibullDistribution where
  maybeMean = Just . D.mean

instance D.Mean WeibullDistribution where
  mean (WD k l) = l * exp (logGamma (1 + 1 / k))

instance D.MaybeVariance WeibullDistribution where
  maybeStdDev   = Just . D.stdDev
  maybeVariance = Just . D.variance

instance D.Variance WeibullDistribution where
  variance (WD k l) = l * l * (exp (logGamma (1 + 2 * invk)) - q * q)
   where
    invk = 1 / k
    q    = exp (logGamma (1 + invk))

instance D.Entropy WeibullDistribution where
  entropy (WD k l) = m_eulerMascheroni * (1 - 1 / k) + log (l / k) + 1

instance D.MaybeEntropy WeibullDistribution where
  maybeEntropy = Just . D.entropy

instance D.ContGen WeibullDistribution where
  genContVar d = D.genContinuous d

-- | Standard Weibull distribution with scale factor (lambda) 1.
weibullStandard :: Double -> WeibullDistribution
weibullStandard k = weibullDistr k 1.0

-- | Create Weibull distribution from parameters.
--
-- If the shape (first) parameter is @1.0@, the distribution is equivalent to a
-- 'Statistics.Distribution.Exponential.ExponentialDistribution' with parameter
-- @1 / lambda@ the scale (second) parameter.
weibullDistr
  :: Double            -- ^ Shape
  -> Double            -- ^ Lambda (scale)
  -> WeibullDistribution
weibullDistr k l = either error id $ weibullDistrErr k l

-- | Create Weibull distribution from parameters.
--
-- If the shape (first) parameter is @1.0@, the distribution is equivalent to a
-- 'Statistics.Distribution.Exponential.ExponentialDistribution' with parameter
-- @1 / lambda@ the scale (second) parameter.
weibullDistrErr
  :: Double            -- ^ Shape
  -> Double            -- ^ Lambda (scale)
  -> Either String WeibullDistribution
weibullDistrErr k l | k <= 0     = Left $ errMsg k l
                    | l <= 0     = Left $ errMsg k l
                    | otherwise = Right $ WD k l

errMsg :: Double -> Double -> String
errMsg k l =
  "Statistics.Distribution.Weibull.weibullDistr: both shape and lambda must be positive. Got shape "
    ++ show k
    ++ " and lambda "
    ++ show l

-- | Create Weibull distribution from mean and standard deviation.
--
-- The algorithm is from "Methods for Estimating Wind Speed Frequency
-- Distributions", C. G. Justus, W. R. Hargreaves, A. Mikhail, D. Graber, 1977.
-- Given the identity:
--
-- \[
-- (\frac{\sigma}{\mu})^2 = \frac{\Gamma(1+2/k)}{\Gamma(1+1/k)^2} - 1
-- \]
--
-- \(k\) can be approximated by
--
-- \[
-- k \approx (\frac{\sigma}{\mu})^{-1.086}
-- \]
--
-- \(\lambda\) is then calculated straightforwardly via the identity
--
-- \[
-- \lambda = \frac{\mu}{\Gamma(1+1/k)}
-- \]
--
-- Numerically speaking, the approximation for \(k\) is accurate only within a
-- certain range. We arbitrarily pick the range \(0.033 \le \frac{\sigma}{\mu} \le 1.45\)
-- where it is good to ~6%, and will refuse to create a distribution outside of
-- this range. The paper does not cover these details but it is straightforward
-- to check them numerically.
weibullDistrApproxMeanStddevErr
  :: Double            -- ^ Mean
  -> Double            -- ^ Stddev
  -> Either String WeibullDistribution
weibullDistrApproxMeanStddevErr m s = if r > 1.45 || r < 0.033
    then Left msg
    else weibullDistrErr k l
  where r = s / m
        k = (s / m) ** (-1.086)
        l = m / exp (logGamma (1 + 1/k))
        msg = "Statistics.Distribution.Weibull.weibullDistr: stddev-mean ratio "
          ++ "outside approximation accuracy range [0.033, 1.45]. Got "
          ++ "stddev " ++ show s ++ " and mean " ++ show m

-- | Uses an approximation based on the mean and standard deviation in
--   'weibullDistrEstMeanStddevErr', with standard deviation estimated
--   using maximum likelihood method (unbiased estimation).
--
--   Returns @Nothing@ if sample contains less than one element or
--   variance is zero (all elements are equal), or if the estimated mean
--   and standard-deviation lies outside the range for which the
--   approximation is accurate.
instance D.FromSample WeibullDistribution Double where
  fromSample xs
    | G.length xs <= 1 = Nothing
    | v == 0           = Nothing
    | otherwise        = either (const Nothing) Just $
      weibullDistrApproxMeanStddevErr m (sqrt v)
    where
      (m,v) = S.meanVarianceUnb xs

logDensity :: WeibullDistribution -> Double -> Double
logDensity (WD k l) x
  | x < 0     = 0
  | otherwise = log k + (k - 1) * log x - k * log l - (x / l) ** k

cumulative :: WeibullDistribution -> Double -> Double
cumulative (WD k l) x | x < 0     = 0
                      | otherwise = -expm1 (-(x / l) ** k)

complCumulative :: WeibullDistribution -> Double -> Double
complCumulative (WD k l) x | x < 0     = 1
                           | otherwise = exp (-(x / l) ** k)

quantile :: WeibullDistribution -> Double -> Double
quantile (WD k l) p
  | p == 0         = 0
  | p == 1         = inf
  | p > 0 && p < 1 = l * (-log1p (-p)) ** (1 / k)
  | otherwise      =
    error $ "Statistics.Distribution.Weibull.quantile: p must be in [0,1] range. Got: " ++ show p
  where inf = 1 / 0

complQuantile :: WeibullDistribution -> Double -> Double
complQuantile (WD k l) q
  | q == 0         = inf
  | q == 1         = 0
  | q > 0 && q < 1 = l * (-log q) ** (1 / k)
  | otherwise      =
    error $ "Statistics.Distribution.Weibull.complQuantile: q must be in [0,1] range. Got: " ++ show q
  where inf = 1 / 0


{-# LANGUAGE BangPatterns        #-}
{-# LANGUAGE DataKinds           #-}
{-# LANGUAGE GADTs               #-}
{-# LANGUAGE KindSignatures      #-}
{-# LANGUAGE LambdaCase          #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE StandaloneDeriving  #-}
{-# LANGUAGE TypeOperators       #-}

module BPTT where

import Control.Monad
import Control.Monad.Random
import Data.Char
import Data.List
import Data.Maybe
import Data.Singletons
import Data.Singletons.Prelude
import Data.Singletons.TypeLits
import GHC.Show
import Numeric (showIntAtBase)
import qualified Numeric.LinearAlgebra as LA
import Numeric.LinearAlgebra.Static

-- dependently typed weights
data Weights i o = W {
    wBiases :: !(R o)
  , wNodes  :: !(L o i)
  }

-- an simple 3-layer RNN, including one recurrent layer,
-- is not like the Network type discussed in Justin Le's
-- blog series (because not all layers are just stacked
-- one on top of the other.)
-- Instead:
data RNN i h o = MkRNN {
    getU    :: !(Weights i h)
  , getW    :: !(Weights h h)
  , getV    :: !(Weights h o)
  }

-- If one were to leverage a.g. ad the Hackage library for
-- autodifferentiable functions, then functions being used for
-- activation could be provided in a symbolic form amenable to
-- being automatically differentiated (for use in the backprop
-- process.) Instead, each activation function used & its derivative
-- are provided separately.

logistic :: Floating a => a -> a
logistic x = 1 / (1 + exp (-x))

logistic' :: Floating a => a -> a
logistic' x = logix * (1 - logix)
  where
    logix = logistic x

softmax :: KnownNat v => R v -> R v
softmax v =
  let vexp = exp v
  in  vexp / (konst $ norm_1 vexp)

-- The formula for the (x, y)-entry of softmax',
-- which is a matrix, is
--
-- softmax'[x, y] = -softmax[x]    * softmax[y],    x /= y
--                  (1-softmax[x]) * softmax[y],    x == y
--
-- which is the same as adding the vector
-- 'softmax v' diagonally down the negative of
-- said vector's outer product with itself.
-- I would prefer to build this using 'build', but
-- I don't know how right now.
softmax' :: KnownNat v => R v -> L v v
softmax' v =
  let s = softmax v
  in  (diag s) + konst (-1) * (s `outer` s)

-- runLayer is for performing the linear mapping
-- due to a particular layer's set of parameters,
-- on an appropriately sized vector
runLayer :: (KnownNat i, KnownNat o)
         => Weights i o
         -> R i -> R o
runLayer (W wB wN) v_in =
  wB + wN #> v_in

-- Verbose 'run RNN', for exposing all intermediate computations.
-- (R i, R h,  R h, R o, R o)
--  x_t  h_t-1 h_t  z_t  y_t

-- runRNN can be expressed in terms of this, if I locate / find
-- generalized tuple accessors
runRNN_verbose  :: (KnownNat i, KnownNat h, KnownNat o)
                => RNN i h o
                -> R i -> R h
                -> (R i, R h, R h, R o, R o)
runRNN_verbose (MkRNN rnn_U rnn_W rnn_V)
  x_t h_t_1 =
    let h_t = logistic $
                -- recurrent parameters, weighting the prior hidden state
                (runLayer rnn_W h_t_1) +
                -- input parameters, weighting the current input
                (runLayer rnn_U x_t)
        -- output parameters, weighting the current hidden state
        z_t = runLayer rnn_V h_t
        -- the final activation provides the prediction
        y_t = logistic z_t
    in  (x_t, h_t_1, h_t, z_t, y_t)

-- Forward phase of training an RNN.
-- No adjustments to the parameters (weights) are made.
-- Instead, foldl is used to accumulate a list of tuples,
-- each tuple bearing all input, intermediate and output vectors
-- for one of the steps (t := 1, 2, ..., n) in the forward direction.
forward_phase :: (KnownNat i, KnownNat h, KnownNat o)
              => RNN i h o    -- RNN
              -> [(R i, R o)] -- list of training inputs paired with target outputs
              -> R h          -- initial hidden state

                              -- All tuples of resulting vectors, from every step carried out:
                              -- For t ranging from 1 to n, n the number of training samples:
              -> [( R i,      -- x_t
                    R h,      -- h_t-1
                    R h,      -- h_t
                    R o,      -- z_t
                    R o,      -- y_t
                    R o)]     -- tgt_t
forward_phase
  rnn
  pairs
  h_0 =
    snd $ foldl unroll (h_0, []) pairs
    where
      -- unroll :: (R h, [(R i, R h, R h, R o, R o, R o)]) -> (R i, R o)
      --        -> (R h, [(R i, R h, R h, R o, R o, R o)])

      unroll (h_t_1, results) pair@(x_t, tgt_t) =
        --                               ^ tgt_t plays no role in the forward direction,
        --                                 but we pattern-match on it so that we can re-package it
        --                                 up into this step's tuple for the result list
        let result@(_, _, h_t, z_t, y_t) = runRNN_verbose rnn x_t h_t_1
        in  (h_t, results ++ [(x_t, h_t_1, h_t, z_t, y_t, tgt_t)])

-- Backward phase of training an RNN.
-- Backpropagation is required.
-- There's a lot that goes into backprop!
-- (Basically, everything that is used or
-- produced at any point in the forward step.)
backward_phase  :: (KnownNat i, KnownNat h, KnownNat o)
                => RNN i h o   -- an RNN to train

                                -- All tuples of resulting vectors, from every step carried out:
                                -- For t ranging from 1 to n, n the number of training samples:
                -> [( R i,      -- x_t
                      R h,      -- h_t-1
                      R h,      -- h_t
                      R o,      -- z_t
                      R o,      -- y_t
                      R o)]     -- tgt_t

                -- At every step #t of BPTT, there are *two* sources of "delta-h_t" the bump
                -- to h_t required for the sake of trying to bump down the error at that same step:
                -- a) a bump due to what W is at the current time, "delta-h given W", and
                -- b) a bump *propagated backwards* from the next step #(t+1).
                -- But in the base case, this dh_t_dh_t' is chosen as 0 (because it's undefined.)

                -- Note the *symmetry* with forward_phase:
                --    just as we 'primed' forward_phase  with an initial hidden state,    h_0        : 0,
                --    we also    'prime'  backward_phase with an initial back-prop'ed dh, dh_n_dh_n' : 0.
                -> R h        -- dh_n_dh_n'

                -> Double     -- a training rate
                -> RNN i h o  -- the trained RNN
backward_phase rnn intermediates dh_n_dh_n' rate =
  snd $ foldr bptt (dh_n_dh_n', rnn) intermediates
    where
      -- -- A BPTT step (at time t) consumes the following data:
      -- --       x_t, h_t, h_t_1, z_t,  y_t, tgt_t
      -- bptt :: (R i, R h, R h,   R o,  R o, R o)
      -- -- It also "accumulates" by swapping out,
      -- -- based upon time-t data  <dh_t   | dh_t+1>, rnn
      --          -> (R h, RNN i h o)
      -- -- the time-(t-1) data     <dh_t_1 | dh_t>,   rnn';
      -- -- in this way, the rnn' "accumulates" training bumps to the parameters
      --          -> (R h, RNN i h o)

      bptt    (x_t, h_t, h_t_1, z_t, y_t, tgt_t)
              (dh_t_dh_t', MkRNN
                -- We'll pattern match into each Weights term u, w, v for convenience
                -- later in deriving the desired infinitesimal bumps.
                u@(W uB uN)
                w@(W wB wN)
                v@(W vB vN)) =
                -- We are hard-coding the use of squared-error for the error function:
                -- e(y_t) := (y_t - tgt_t) ^ 2
                -- whose derivative with respect to the pre-activation z_t is:

                {-
                  de/dz = de/dy           * dy/dz
                        = (e' at y_t)     * (y' at z_t)
                        = 2 (y_t - tgt_t) * (logistic' z_t)
                  These 2 factors appear (in reverse order) as 'dedz' below;
                  the factor of 2 is discarded & we defer all control of the
                  desired training rate to the scalar input param 'rate'.
                -}
            let dedz          = logistic' z_t * (y_t - tgt_t)
                              -- = softmax' z_t #> (y_t - tgt_t)

                {-
                  Informally:
                  dedz is a vector that reports,
                    for a unit "nudge" to each component z_t(j) ->
                      size / direction of resulting "nudge" for error component e_t(j).
                  The choice of dz_t below corresponds to the following training strategy:
                    a) bump each z-component in the direction that will *decrease* its
                       matching e-component (thus, factor of '-1')
                    b) but re-scaled based on the learning rate desired (thus, 'scale rate')
                -}
                dz_t          = -1 * konst rate * dedz

                {-
                  Informally:
                  Having set out to make these infinitesimal changes to the z_t components,
                  we need to figure out
                    what changes to V,    *given* what h_t is,     and
                    what changes to h_t,  *given* what V is
                  will help us accomplish that.

                  Observe that
                    z  = vN h              + vB
                  therefore (product rule on 1st summand)
                    dz = (dvN h) + (vN dh) + dvB

                  we do *not* try to formally solve for dvN, dh, dvB in here, but rather use these
                  as instructions for how to recombine the vectors dz_t / h_t and matrix V
                  so as to arrive at dvN_h ("bump to v's matrix part, given h") and
                  dh_vN ("bump to h, given v's matrix part").

                  As for v's bias-vector part, vB, we simply nudge it directly by the dz (z-nudge)
                  we worked out above. Why? Look at the infinitesimals equation above - a nudge
                  dvB performed on the bias vector will show up directly in dz.
                -}
                dvB           = dz_t
                {-
                  Instead of trying to solve for dvN, dimensional considerations lead us to this formula
                  for dvN in terms of dz_t, h_t:
                    vN   goes    from h-vectors  to o-vectors; i.e.    o * h -matrix
                    h_t            an h-vector
                    dz_t                         an o-vector
                    so the way to combine these to get a "bumps matrix" compatible with vN is
                    dz_t `outer` h_t   <- "column" o-vector times "row" h-vector
                -}
                dvN_h         = dz_t `outer` h_t
                {-
                  Similar to the above, but for building a "bumps h-vector" out of o * h matrix
                  and o-vector.
                -}
                dh_vN         = tr vN #> dz_t
                {-
                  See the note in the type signature of 'backward_phase' about the 2 sources of
                  training bumps to h_t at every step t.
                  Read dh_t_dh_t' as (bump to h at current time t,
                                      propped backwards from bump to h at future time t+1)
                -}
                dh_t          = dh_vN + dh_t_dh_t'

                -- Having now characterized in what way we are going to bump the
                -- entries in the current hidden state h_t, so as to achieve
                -- a desired reduction in error (i.e. dh_t), we now use it to
                -- determine how we should modify the intermediate preactivation,
                -- calling it q, q_t := (W * h_t_1 + U * x_t) to achieve the same goal:

                -- de / dq = (de / dh) * (dh / dq)  -- chain rule
                --         = dh_t      * (dh / dq)  -- back propagation
                --         = dh_t      * (g' at q)  -- h is the activation (g) of q

                dedq          = dh_t * logistic' q
                                  where q = runLayer w h_t_1 + runLayer u x_t

                -- and work out a desired perturbation to q, based upon scaling dedq
                dq_t          = (-1) * konst rate * dedq

                {-
                  q  = W * h_t_1 + U * x_t
                  dq = (dW * h_t_1) + (W * dh_t_1)
                        + (dU * x_t) + [the term in dx_t -> 0, since x constant]
                -}

                {-
                  Recurrent contribution to time-t hidden state:
                  q_t  = wN h_t_1                   + wB        + ...
                  dq_t = (dwN h_t_1) + (wN dh_t_1)  + dwB       + ...
                -}
                dwB           = dq_t
                -- dq_t and h_t_1 both happen to be h-vectors having the same
                -- dimensionality, so here the driving intuition is that such
                -- "matrix-bumps" take the form [destination] `outer` [origin]
                dwN_h         = dq_t `outer` h_t_1

                -- This is what will be given back from the current fold instance
                -- to the next one
                dh_t_1_dh_t   = tr wN #> dq_t
                {-
                  Input contribution to time-t hidden state:
                  q_t  = ... + uN x_t                 + uB
                  dq_t = ... + (duN x_t) + (uN dx_t)  + duB
                       = ... + (duN x_t) + ZERO       + duB
                                        -- ^ x_t, the training inputs, are constant!
                -}
                duB           = dq_t
                duN_x         = dq_t `outer` x_t

                -- Don't forget to apply all the RNN-training nudges!
                vB'           = vB + dvB
                vN'           = vN + dvN_h
                wB'           = wB + dwB
                wN'           = wN + dwN_h
                uB'           = uB + duB
                uN'           = uN + duN_x

            in  (dh_t_1_dh_t, MkRNN (W uB' uN')
                                    (W wB' wN')
                                    (W vB' vN'))

{-
 Making some simple training data.

 Here's a simple sentential grammar for producing "sentences"
 of ints, that later can go in an overall "document":

  <sentence>        ::= <sentence-start> <sentence-middle> <sentence-finish>
  <sentence-start>  ::= a multiple of 20, 20 * s, for some s : 1 <= s <= 10
  <sentence-middle> ::= a sequence of l integers, for some l : 0 <= l <= 10
                        that are derived from the sentence-start as
                          s - m, (s - m) + n, (s - m + n) - m, (s - m + n - m) + n, ...
                        for some choice of parameters      m, n : 1 <= m, n <= 5
  <sentence-finish> ::= a multiple of 17, 17 * f, for some f : 1 <= f <= 10
-}

-- Generate a "biarithmetic" sequence for sentence-middle, given
-- a base number s, a down size m, an up size n
biarithmetic :: Int -> Int -> Int -> [Int]
biarithmetic    s      m      n   =  concat $
 [ [ s - (j * m) + (k * n) | k <- [j-1 .. j] ]
                           | j <- [1..] ]

-- For use as a multiple of 20
genStartParam :: MonadRandom m => m Int
genStartParam = getRandomR (3, 10)

-- For use as a mid-sequence length
genMiddleLenParam :: MonadRandom m => m Int
genMiddleLenParam = getRandomR (0, 10)

-- For use as the down size and up size of a biarithmetic sequence
genMiddleDownUpParams :: MonadRandom m => m (Int, Int)
genMiddleDownUpParams = do
  m <- getRandomR (1, 5)
  n <- getRandomR (1, 5)
  return (m, n)

-- For use as a multiple of 17
genFinishParam :: MonadRandom m => m Int
genFinishParam = getRandomR (1, 10)

-- For use as a number of "sentences".
genSentenceCount :: MonadRandom m => m Int
genSentenceCount = getRandomR (500, 1000)

type Sentence = [Int]

genSentence :: MonadRandom m => m Sentence
genSentence = do
  -- 1. Get sentence-start
  s       <- genStartParam
  let sent_0 = s * 20
  -- 2. Get sentence-middle length
  l       <- genMiddleLenParam
  -- 3. Get sentence up & down size params
  (m, n)  <- genMiddleDownUpParams
  -- 4. Get sentence-middle
  let sent_mid = take l $ biarithmetic sent_0 m n
  -- 5. Get sentence-finish
  f       <- genFinishParam
  let sent_f = f * 17
  -- 6. Return random sentence
  return (sent_0 : sent_mid ++ [sent_f])

type Document = [Sentence]

genDocument :: MonadRandom m => m Document
genDocument = do
  -- 1. Get number of sentences
  sent_count <- genSentenceCount
  -- 2. For as many sentences as 'sent_count', generate one
  --    random sentence each & then combine the results inside
  --    MonadRandom context
  let rand_sent_stream = repeat genSentence
  sentences <- sequenceA $ take sent_count rand_sent_stream
  return sentences

{-
 Now that we have a way of generating "documents" made up out of
  numerical "sentences", we need to convert this into a form that
  would be characteristically used to train a neural net.
  Convert each integer to a (possibly sparse) vector. How?
  Every integer present in a (d :: Document) should always lie
  within the range (0 <= j <= 256), so we could use a binary representation
  taking up 8 bits (tossing in an extra 2 to get 10, for good measure)
-}

-- The full source module citation is probably forgivable if I only do it this once.
-- Best practice is to 'import qualified' with an alias intToVector :: Int -> R 10
intToVector m =
  let binary_string = reverse $ showIntAtBase 2 intToDigit m "" -- :: [Char]
      binary_padded = binary_string ++
                        (take (10 - length binary_string) $ repeat '0')
      doubles_list  = (fromIntegral . digitToInt) <$> binary_padded
      doubles_vect  = Numeric.LinearAlgebra.Static.fromList doubles_list
  in  doubles_vect

-- The other way round, for casting RNN predictions back into Int's.
-- This is unsafe in case the entries of any y_t end up rounding to
-- numbers other than 0, 1.
-- (Does the hmatrix API support a more direct way of casting from
-- R 10 to the list of rounded numbers?)
vectorToInt :: R 10 -> Int
vectorToInt v =
  let bins = (fmap toInteger) <$> LA.toList
                                $ LA.toZ
                                $ LA.roundVector
                                $ extract v
  in  fromIntegral $
        foldr (\p@(bit,  power)  total ->
                   bit * power + total)
        0
        (zip bins [ 2^j | j <- [0..] ])

zero_vec :: KnownNat n => R n
zero_vec = 0

-- via Justin Le's blog series on neural nets
randomWeights :: (MonadRandom m, KnownNat i, KnownNat o)
              => m (Weights i o)
randomWeights = do
    s1 :: Int <- getRandom
    s2 :: Int <- getRandom
    let wB = randomVector  s1 Uniform * 2 - 1
        wN = uniformSample s2 (-1) 1
    return $ W wB wN

randomRNN :: (MonadRandom m, KnownNat i, KnownNat h, KnownNat o)
          => m (RNN i h o)
randomRNN =
  (pure MkRNN)  <*> randomWeights
                <*> randomWeights
                <*> randomWeights

trainRNN :: (KnownNat i, KnownNat h, KnownNat o)
         => RNN i h o     -- an RNN
         -> [(R i, R o)]  -- the training data: pairs of inputs & training outputs
         -> Double
         -> RNN i h o     -- an RNN with trained parameters
trainRNN rnn train_data rate =
  let intermediates = forward_phase rnn train_data zero_vec
      trained       = backward_phase rnn intermediates zero_vec rate
  in  trained

-- Before re-formatting into a vector form suitable for use
-- in training, we have Document ~ [Sentence] ~ [[Int]];
-- now, turn each Int into a binary-repr vector, R 10
genSampleData :: MonadRandom m => m [[ R 10 ]]
genSampleData = do
  doc <- genDocument
  return $ (intToVector <$>) <$> doc

-- Every sentence (sublist) in the document (list of lists) will
-- furnish one sequence of training data for the RNN to "unfold"
-- over and perform BPTT on. So, BPTT is performed once per sentence
-- in the document overall.
pairEachWithNext :: [x] -> [(x, x)]
pairEachWithNext xs =
  zip (init xs) (tail xs)

-- In the special case of an 'RNN i h o' where o = i,
-- we can pipe back a prediction y_t due to x_t and an
-- initial state h, to furnish the next x_(t+1) for further
-- predictions. We'll use this to observe the way that a
-- trained RNN behaves.
rnnStep :: (KnownNat i, KnownNat h)
        => RNN i h i
        -> (R h, R i)
        -> (R h, R i)
rnnStep rnn (h_t_1, x_t) =
  let full_output = runRNN_verbose rnn x_t h_t_1
      third five@(_, _, x3, _, _)  = x3
      fifth five@(_, _, _,  _, x5) = x5
  in  (third full_output, fifth full_output)

getDocumentIO :: IO Document
getDocumentIO = evalRandIO genDocument

takePredictions  :: (KnownNat i, KnownNat h)
                 => RNN i h i
                 -> (R h, R i)
                 -> Int
                 -> [R i]
takePredictions rnn (h_0, x_1) n =
  let hxs_stream = take (n+1) $ iterate (rnnStep rnn) (h_0, x_1)
  in  snd <$> hxs_stream

takePredsFromN :: KnownNat h
               => RNN 10 h 10
               -> Int      -- word of a sentence
               -> Int      -- number of predictions
               -> [R 10]
takePredsFromN rnn x_1 n =
  takePredictions rnn (zero_vec, intToVector x_1) n

trainRNNOnSentence  :: KnownNat h
                    => RNN 10 h 10
                    -> Sentence
                    -> Double
                    -> RNN 10 h 10
trainRNNOnSentence rnn sent rate =
  let training_data = pairEachWithNext (intToVector <$> sent)
  in  trainRNN rnn training_data rate

-- trainRNNOnDocument has to be structured like a fold
-- so that the trained RNN due to each sentence can be
-- passed to the following sentence (at the same time,
-- I think that hidden state "restarts" with the beginning
-- of a new sentence.)
trainRNNOnDocument  :: KnownNat h
                    => RNN 10 h 10
                    -> Document
                    -> Double
                    -> RNN 10 h 10
trainRNNOnDocument rnn sents rate =
  foldl (\rnn sent -> trainRNNOnSentence rnn sent rate)
        rnn
        sents

observeUntrainedPredictions :: Int    -- initial word
                            -> Int    -- number of predictions
                            -> IO ()  -- action printing info to console
observeUntrainedPredictions n_0 num = do
  -- get an RNN of random parameters
  putStrLn "Getting a random RNN ..."
  (rnn :: RNN 10 300 10) <- evalRandIO randomRNN
  -- get a series of predictions from the untrained RNN
  putStrLn $ "Producing: " ++ show num ++
             " predictions, starting from: " ++ show n_0
  let vs = takePredsFromN rnn n_0 num
  let ns = vectorToInt <$> vs
  putStrLn "The RNN has predicted: "
  print ns

observeTrainedPredictions :: Int    -- initial word
                          -> Int    -- number of predictions
                          -> Double -- learning rate
                          -> IO ()  -- action printing info to console
observeTrainedPredictions n_0 num rate = do
  -- get an RNN of random parameters
  putStrLn "Getting a random RNN ..."
  (rnn :: RNN 10 300 10) <- evalRandIO randomRNN
  -- get a random document
  putStrLn "Getting a random document to train on ..."
  document <- evalRandIO genDocument
  -- train RNN on the document
  putStrLn "Training the RNN on the document ..."
  let rnn' = trainRNNOnDocument rnn document rate
  -- get a series of predictions from the trained RNN
  putStrLn $ "Producing: " ++ show num ++
             " predictions, starting from: " ++ show n_0 ++
             " and using a trained RNN"
  let vs = takePredsFromN rnn' n_0 num
  let ns = vectorToInt <$> vs
  putStrLn "The trained RNN has predicted: "
  print ns

observeUntrainedVsTrained :: Int    -- initial word
                          -> Int    -- number of predictions
                          -> Double -- learning rate
                          -> IO ()  -- action printing info to console
observeUntrainedVsTrained n_0 num rate = do
  -- get an RNN of random parameters
  putStrLn "Getting a random RNN ..."
  (rnn :: RNN 10 300 10) <- evalRandIO randomRNN
  -- get a series of predictions from the untrained RNN
  putStrLn $ "Producing: " ++ show num ++
             " predictions, starting from: " ++ show n_0
  let vs = takePredsFromN rnn n_0 num
  let ns = vectorToInt <$> vs
  putStrLn "The RNN has predicted: "
  print ns
  -- get a random document
  putStrLn "Getting a random document to train on ..."
  document <- evalRandIO genDocument
  -- train RNN on the document
  putStrLn "Training the RNN on the document ..."
  let rnn' = trainRNNOnDocument rnn document rate
  -- get a series of predictions from the trained RNN
  putStrLn $ "Producing: " ++ show num ++
             " predictions, starting from: " ++ show n_0 ++
             " and using a trained RNN"
  let vs' = takePredsFromN rnn' n_0 num
  let ns' = vectorToInt <$> vs'
  putStrLn "The trained RNN has predicted: "
  print ns'


{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE Strict           #-}
module FourierPinwheel.AsteriskGaussian where

import           Data.Array.Repa             as R
import           Data.Complex
import           Data.List                   as L
import           Data.Vector.Generic         as VG
import           Data.Vector.Storable        as VS
import           Data.Vector.Unboxed         as VU
import           DFT.Plan
import           Filter.Utils
import           FourierPinwheel.Array
import           FourierPinwheel.Hypergeo1F1
import           Math.Gamma
import           Pinwheel.FourierSeries2D
import           Utils.Distribution
import           Utils.Parallel

-- The envelope in R2 is r ^ alpha, where -2 < alpha < -0.5
{-# INLINE asteriskGaussian #-}
asteriskGaussian ::
     ( VG.Vector vector (Complex e)
     , RealFloat e
     , NFData (vector (Complex e))
     , Gamma (Complex e)
     , Unbox e
     )
  => Int
  -> Int
  -> e
  -> e
  -> e
  -> e
  -> [vector (Complex e)]
asteriskGaussian numR2Freqs thetaFreq alpha periodR2 periodEnv std =
  parMap
    rdeepseq
    (\angularFreq ->
       let arr =
             R.map
               (* (sqrt (pi / std) *
                   exp (fromIntegral angularFreq ^ 2 / (-4) / std) :+
                   0)) .
             centerHollowArray numR2Freqs $
             analyticalFourierCoefficients1
               numR2Freqs
               1
               angularFreq
               0
               alpha
               periodR2
               periodEnv
        in VG.convert . toUnboxed . computeS $ arr)
    [-thetaFreq .. thetaFreq]

{-# INLINE asteriskGaussianEnvelope #-}
asteriskGaussianEnvelope ::
     (R.Source s (Complex Double))
  => DFTPlan
  -> Int
  -> Int
  -> Double
  -> Double
  -> Double
  -> Double
  -> R.Array s DIM2 (Complex Double)
  -> IO (VS.Vector (Complex Double))
asteriskGaussianEnvelope plan numR2Freqs thetaFreq alpha periodR2 periodEnv stdTheta filter = do
  let planID = DFTPlanID DFT1DG [numR2Freqs, numR2Freqs] [0, 1]
      inversePlanID = DFTPlanID IDFT1DG [numR2Freqs, numR2Freqs] [0, 1]
      vecs =
        asteriskGaussian numR2Freqs thetaFreq alpha periodR2 periodEnv stdTheta
  filterF <-
    dftExecute plan planID . VU.convert . toUnboxed . computeS . makeFilter2D $
    filter
  asteriskGaussianF <- dftExecuteBatchP plan planID vecs
  fmap VS.concat .
    dftExecuteBatchP plan inversePlanID .
    parMap rdeepseq (VS.zipWith (*) filterF) $
    asteriskGaussianF

-- The envelope is r^alpha X Gaussian
asteriskGaussian2 ::
     DFTPlan
  -> Int
  -> Int
  -> Double
  -> Double
  -> Double
  -> Double
  -> Double
  -> IO (VS.Vector (Complex Double))
asteriskGaussian2 plan numR2Freqs thetaFreq alpha periodR2 periodEnv stdR2 stdTheta = do
  let centerFreq = div numR2Freqs 2
      gaussian2D =
        fromFunction (Z :. numR2Freqs :. numR2Freqs) $ \(Z :. i' :. j') ->
          let i = i' - centerFreq
              j = j' - centerFreq
           in exp
                (pi * fromIntegral (i ^ 2 + j ^ 2) /
                 ((-1) * periodR2 ^ 2 * stdR2 ^ 2)) /
              (2 * pi * stdR2 ^ 2) :+
              0
  asteriskGaussianEnvelope
    plan
    numR2Freqs
    thetaFreq
    alpha
    periodR2
    periodEnv
    stdTheta
    gaussian2D

-- The envelope is r^alpha X LowPass
asteriskGaussianLowPass ::
     DFTPlan
  -> Int
  -> Int
  -> Double
  -> Double
  -> Double
  -> Double
  -> Double
  -> IO (VS.Vector (Complex Double))
asteriskGaussianLowPass plan numR2Freqs thetaFreq alpha periodR2 periodEnv stdTheta radius = do
  let lpf =
        fromUnboxed (Z :. numR2Freqs :. numR2Freqs) $
        idealLowPassFilter radius periodR2 numR2Freqs
  asteriskGaussianEnvelope
    plan
    numR2Freqs
    thetaFreq
    alpha
    periodR2
    periodEnv
    stdTheta
    lpf

{-# INLINE asteriskGaussianFull #-}
asteriskGaussianFull ::
     ( VG.Vector vector (Complex e)
     , RealFloat e
     , NFData (vector (Complex e))
     , Gamma (Complex e)
     , Unbox e
     )
  => Int
  -> Int
  -> Int
  -> e
  -> e
  -> e
  -> e
  -> e
  -> [vector (Complex e)]
asteriskGaussianFull numR2Freqs thetaFreq rFreq alpha periodR2 periodEnv stdTheta stdR =
  let zeroVec = VG.replicate (numR2Freqs ^ 2) 0
   in parMap
        rdeepseq
        (\(radialFreq, angularFreq) ->
           let pinwheel =
                 analyticalFourierCoefficients1
                   numR2Freqs
                   1
                   angularFreq
                   radialFreq
                   alpha
                   periodR2
                   periodEnv
               -- arr =
               --   R.map
               --     (* ((-- gaussian1DFreq (fromIntegral angularFreq) stdTheta *
               --          gaussian1DFourierCoefficients
               --            (fromIntegral radialFreq)
               --            (log periodEnv)
               --            stdR
               --          :+
               --          0) -- *
               --         -- cis
               --         --   (2 * pi / log periodEnv * fromIntegral radialFreq *
               --         --    log 0.5)
               --        )) $
               --   -- if radialFreq == 0 && angularFreq == 0
               --   --   then pinwheel
               --   --   else
                   -- centerHollowArray numR2Freqs pinwheel
               arr =
                 R.map
                   (* ((1 / (1 + (2 * pi / log periodEnv * fromIntegral radialFreq )^2)) :+ 0)
                      ) $
                   centerHollowArray numR2Freqs pinwheel
            in if angularFreq == 0
                 then VG.convert . toUnboxed . computeS $ arr
                 else zeroVec)
        [ (radialFreq, angularFreq)
        | radialFreq <-  [-rFreq .. rFreq]
        , angularFreq <- [-thetaFreq .. thetaFreq]
        ]

{-# INLINE gaussianFull #-}
gaussianFull ::
     ( VG.Vector vector (Complex e)
     , RealFloat e
     , NFData (vector (Complex e))
     , Gamma (Complex e)
     , Unbox e
     )
  => Int
  -> Int
  -> Int
  -> e
  -> e
  -> [vector (Complex e)]
gaussianFull numR2Freqs thetaFreq rFreq periodR2 stdR2 =
  let zeroVec = VG.replicate (numR2Freqs ^ 2) 0
      periodEnv = periodR2 ^ 2 / 4
   in parMap
        rdeepseq
        (\(radialFreq, angularFreq) ->
           if angularFreq == 0 -- && radialFreq == 0
             then VG.convert . toUnboxed . computeS $
                  fromFunction (Z :. numR2Freqs :. numR2Freqs) $ \(Z :. i' :. j') ->
                    let i = fromIntegral $ i' - div numR2Freqs 2
                        j = fromIntegral $ j' - div numR2Freqs 2
                     in (gaussian2DFourierCoefficients i j periodR2 stdR2 :+ 0) /
                        (1 :+ (2 * pi / log periodEnv * fromIntegral radialFreq)^2)
             else zeroVec)
        [ (radialFreq, angularFreq)
        | radialFreq <- [-rFreq .. rFreq]
        , angularFreq <- [-thetaFreq .. thetaFreq]
        ]

{-# INLINE asteriskGaussianFullEnvelope #-}
asteriskGaussianFullEnvelope ::
     (R.Source s (Complex Double))
  => DFTPlan
  -> Int
  -> Int
  -> Int
  -> Double
  -> Double
  -> Double
  -> Double
  -> Double
  -> R.Array s DIM2 (Complex Double)
  -> IO (VS.Vector (Complex Double))
asteriskGaussianFullEnvelope plan numR2Freqs thetaFreq rFreq alpha periodR2 periodEnv stdTheta stdR filter = do
  let planID = DFTPlanID DFT1DG [numR2Freqs, numR2Freqs] [0, 1]
      inversePlanID = DFTPlanID IDFT1DG [numR2Freqs, numR2Freqs] [0, 1]
      vecs =
        asteriskGaussianFull
          numR2Freqs
          thetaFreq
          rFreq
          alpha
          periodR2
          periodEnv
          stdTheta
          stdR
      -- vecs =
      --  gaussianFull
      --      numR2Freqs
      --      thetaFreq
      --      rFreq
      --      periodR2
      --      stdR
  filterF <-
    dftExecute plan planID . VU.convert . toUnboxed . computeS . makeFilter2D $
    filter
  asteriskGaussianF <- dftExecuteBatchP plan planID vecs
  fmap VS.concat .
    dftExecuteBatchP plan inversePlanID .
    parMap rdeepseq (VS.zipWith (*) filterF) $
    asteriskGaussianF
  -- return . VS.concat $ vecs

-- The envelope is r^alpha X LowPass
asteriskGaussianFullLowPass ::
     DFTPlan
  -> Int
  -> Int
  -> Int
  -> Double
  -> Double
  -> Double
  -> Double
  -> Double
  -> Double
  -> IO (VS.Vector (Complex Double))
asteriskGaussianFullLowPass plan numR2Freqs thetaFreq rFreq alpha periodR2 periodEnv stdTheta stdR radius = do
  let lpf =
        fromUnboxed (Z :. numR2Freqs :. numR2Freqs) $
        idealLowPassFilter radius periodR2 numR2Freqs
  asteriskGaussianFullEnvelope
    plan
    numR2Freqs
    thetaFreq
    rFreq
    alpha
    periodR2
    periodEnv
    stdTheta
    stdR
    lpf

-- The envelope is r^alpha X LowPass
asteriskGaussian2Full ::
     DFTPlan
  -> Int
  -> Int
  -> Int
  -> Double
  -> Double
  -> Double
  -> Double
  -> Double
  -> Double
  -> IO (VS.Vector (Complex Double))
asteriskGaussian2Full plan numR2Freqs thetaFreq rFreq alpha periodR2 periodEnv stdTheta stdR stdR2 = do
  let centerFreq = div numR2Freqs 2
      gaussian2D =
        fromFunction (Z :. numR2Freqs :. numR2Freqs) $ \(Z :. i' :. j') ->
          let i = fromIntegral $ i' - centerFreq
              j = fromIntegral $ j' - centerFreq
           in gaussian2DFourierCoefficients i j periodR2 stdR2 :+ 0
  asteriskGaussianFullEnvelope
    plan
    numR2Freqs
    thetaFreq
    rFreq
    alpha
    periodR2
    periodEnv
    stdTheta
    stdR
    gaussian2D

asteriskGaussianRFull ::
     DFTPlan
  -> Int
  -> Int
  -> Int
  -> Double
  -> Double
  -> Double
  -> Double
  -> Double
  -> Double
  -> IO (VS.Vector (Complex Double))
asteriskGaussianRFull plan numR2Freqs thetaFreq rFreq alpha periodR2 periodEnv stdTheta stdR stdR2 = do
  let planID = DFTPlanID DFT1DG [numR2Freqs, numR2Freqs] [0, 1]
      inversePlanID = DFTPlanID IDFT1DG [numR2Freqs, numR2Freqs] [0, 1]
      centerFreq = div numR2Freqs 2
      gaussian2D =
        fromFunction (Z :. numR2Freqs :. numR2Freqs) $ \(Z :. i' :. j') ->
          let i = i' - centerFreq
              j = j' - centerFreq
           in exp
                (pi * fromIntegral (i ^ 2 + j ^ 2) /
                 ((-1) * periodR2 ^ 2 * stdR2 ^ 2)) /
              (2 * pi * stdR2 ^ 2) :+
              0
      r =
        centerHollowArray' numR2Freqs $ analyticalFourierCoefficients3 numR2Freqs 1 0 0 alpha periodR2 periodEnv
      vecs =
        asteriskGaussianFull
          numR2Freqs
          thetaFreq
          rFreq
          alpha
          periodR2
          periodEnv
          stdTheta
          stdR
  gaussian2DF <-
    dftExecute plan planID . VU.convert . toUnboxed . computeS . makeFilter2D $
    gaussian2D
  rF <-
    dftExecute plan planID . VU.convert . toUnboxed . computeS . makeFilter2D $
    r
  let filterF = VS.zipWith (\x y -> x + y ^ 2) gaussian2DF rF
  asteriskGaussianF <- dftExecuteBatchP plan planID vecs
  fmap VS.concat .
    dftExecuteBatchP plan inversePlanID .
    parMap rdeepseq (VS.zipWith (*) filterF) $
    asteriskGaussianF


module Hw4
    (
    Vec(..),
    range,
    absolute,
    rd,
    dft,
    fft) where

import Data.Complex
import GHC.Base

range :: Double -> Double -> Double -> [Double]
absolute :: Vec (Complex Double) -> Vec Double
rd :: Int -> Vec Double -> Vec Double
dft :: [Double] -> Vec (Complex Double)
fft :: [Double] -> Vec (Complex Double)

newtype Vec a = Vec {runVec :: [a]}

instance Show a => Show (Vec a) where
    show (Vec a) = "[" ++ show' a ++ "]"

show' :: Show a => [a] -> [Char]
show' [] = ""
show' [x] = show x
show' (x:xs) = show x ++ " " ++ show' xs

instance Functor Vec where
    fmap g x = Vec $ g <$> runVec x

instance Applicative Vec where
    pure x = Vec $ repeat x
    liftA2 f a b =  Vec $ zipWith f (runVec a) (runVec b)

instance Semigroup (Vec a) where
    (<>) (Vec a) (Vec b) = Vec $ a ++ b

instance Monoid (Vec a) where
    mempty = Vec []

instance Num a => Num (Vec a) where
    fromInteger x = Vec $ pure $ fromInteger x
    (+) = liftA2 (+)
    (-) = liftA2 (-)
    (*) = liftA2 (*)
    negate (Vec a) = Vec $ fmap negate a
    abs (Vec a) = Vec $ fmap abs a
    signum (Vec a) = Vec $ fmap signum a

instance Fractional a => Fractional (Vec a) where
    (/) = liftA2 (/)
    fromRational x = Vec $ pure $ fromRational x

instance Floating a => Floating (Vec a) where
    exp (Vec a) = Vec (fmap exp a)
    log (Vec a) = Vec (fmap log a)
    sin (Vec a) = Vec (fmap sin a)
    cos (Vec a) = Vec (fmap cos a)
    asin (Vec a) = Vec (fmap asin a)
    acos (Vec a) = Vec (fmap acos a)
    atan (Vec a) = Vec (fmap atan a)
    sinh (Vec a) = Vec (fmap sinh a)
    cosh (Vec a) = Vec (fmap cosh a)
    asinh (Vec a) = Vec (fmap asinh a)
    acosh (Vec a) = Vec (fmap acosh a)
    atanh (Vec a) = Vec (fmap atanh a)
    pi = Vec $ pure pi

range from to count = fmap (\x -> from + x * step) [0..count-1]
    where step = (to - from)/count

absolute = fmap (\x -> sqrt(realPart x * realPart x + imagPart x * imagPart x))

rd n = fmap (\x -> fromIntegral (round $ c * x) / c)
  where c = 10^n

dft x =
    let n = fromIntegral $ length x
        index = range 0 n n
        xn = x `zip` index

        f k = sum y
                where
                    y = map factor xn
                    factor (xi, j) = let y = 2 * pi * j * k / n
                                        in xi * cos y :+ (-xi * sin y)
    in
        Vec $ fmap f index

fft x
    | n <= 16 = dft x
    | otherwise =
        let (even, odd) = split x
            (e, o) = (fft even, fft odd)
            t = -2 * pi / fromIntegral n

            g (e', o', k) =
                        let
                            y = t * fromIntegral k
                            (fr, fi) = (cos y, sin y)
                            (pr, pi) = (fr * realPart o' - fi * imagPart o', fr * imagPart o' + fi * realPart o')
                        in
                            ((realPart e' + pr) :+ (imagPart e' + pi), (realPart e' - pr) :+ (imagPart e' - pi))

            (lower, upper) = unzip $ fmap g $ zip3 (runVec e) (runVec o) [0..]
        in
            Vec (lower ++ upper)

    where
            n = length x
            split [] = ([], [])
            split [a] = ([a], [])
            split (a:b:c) = (a:x, b:y)
                where (x,y) = split c


{-# LANGUAGE DeriveDataTypeable, DeriveGeneric #-}
-- |
-- Module    : Statistics.Distribution.Geometric
-- Copyright : (c) 2009 Bryan O'Sullivan
-- License   : BSD3
--
-- Maintainer  : bos@serpentine.com
-- Stability   : experimental
-- Portability : portable
--
-- The Geometric distribution. There are two variants of
-- distribution. First is the probability distribution of the number
-- of Bernoulli trials needed to get one success, supported on the set
-- [1,2..] ('GeometricDistribution'). Sometimes it's referred to as
-- the /shifted/ geometric distribution to distinguish from another
-- one.
--
-- Second variant is probability distribution of the number of
-- failures before first success, defined over the set [0,1..]
-- ('GeometricDistribution0').
module Statistics.Distribution.Geometric
    (
      GeometricDistribution
    , GeometricDistribution0
    -- * Constructors
    , geometric
    , geometric0
    -- ** Accessors
    , gdSuccess
    , gdSuccess0
    ) where

import Data.Aeson (FromJSON, ToJSON)
import Control.Applicative ((<$>))
import Control.Monad (liftM)
import Data.Binary (Binary)
import Data.Binary (put, get)
import Data.Data (Data, Typeable)
import GHC.Generics (Generic)
import Numeric.MathFunctions.Constants (m_pos_inf, m_neg_inf)
import qualified Statistics.Distribution as D
import qualified System.Random.MWC.Distributions as MWC

----------------------------------------------------------------
-- Distribution over [1..]

newtype GeometricDistribution = GD {
      gdSuccess :: Double
    } deriving (Eq, Read, Show, Typeable, Data, Generic)

instance FromJSON GeometricDistribution
instance ToJSON GeometricDistribution

instance Binary GeometricDistribution where
    get = GD <$> get
    put (GD x) = put x

instance D.Distribution GeometricDistribution where
    cumulative = cumulative

instance D.DiscreteDistr GeometricDistribution where
    probability (GD s) n
      | n < 1     = 0
      | otherwise = s * (1-s) ** (fromIntegral n - 1)
    logProbability (GD s) n
       | n < 1     = m_neg_inf
       | otherwise = log s + log (1-s) * (fromIntegral n - 1)

instance D.Mean GeometricDistribution where
    mean (GD s) = 1 / s

instance D.Variance GeometricDistribution where
    variance (GD s) = (1 - s) / (s * s)

instance D.MaybeMean GeometricDistribution where
    maybeMean = Just . D.mean

instance D.MaybeVariance GeometricDistribution where
    maybeStdDev   = Just . D.stdDev
    maybeVariance = Just . D.variance

instance D.Entropy GeometricDistribution where
  entropy (GD s)
    | s == 0 = m_pos_inf
    | s == 1 = 0
    | otherwise = negate $ (s * log s + (1-s) * log (1-s)) / s

instance D.MaybeEntropy GeometricDistribution where
  maybeEntropy = Just . D.entropy

instance D.DiscreteGen GeometricDistribution where
  genDiscreteVar (GD s) g = MWC.geometric1 s g

instance D.ContGen GeometricDistribution where
  genContVar d g = fromIntegral `liftM` D.genDiscreteVar d g

-- | Create geometric distribution.
geometric :: Double                -- ^ Success rate
          -> GeometricDistribution
geometric x
  | x >= 0 && x <= 1 = GD x
  | otherwise        =
    error $ "Statistics.Distribution.Geometric.geometric: probability must be in [0,1] range. Got " ++ show x

cumulative :: GeometricDistribution -> Double -> Double
cumulative (GD s) x
  | x < 1        = 0
  | isInfinite x = 1
  | isNaN      x = error "Statistics.Distribution.Geometric.cumulative: NaN input"
  | otherwise    = 1 - (1-s) ^ (floor x :: Int)

----------------------------------------------------------------
-- Distribution over [0..]

newtype GeometricDistribution0 = GD0 {
      gdSuccess0 :: Double
    } deriving (Eq, Read, Show, Typeable, Data, Generic)

instance FromJSON GeometricDistribution0
instance ToJSON GeometricDistribution0

instance Binary GeometricDistribution0 where
    get = GD0 <$> get
    put (GD0 x) = put x

instance D.Distribution GeometricDistribution0 where
    cumulative (GD0 s) x = cumulative (GD s) (x + 1)

instance D.DiscreteDistr GeometricDistribution0 where
    probability    (GD0 s) n = D.probability    (GD s) (n + 1)
    logProbability (GD0 s) n = D.logProbability (GD s) (n + 1)

instance D.Mean GeometricDistribution0 where
    mean (GD0 s) = 1 / s - 1

instance D.Variance GeometricDistribution0 where
    variance (GD0 s) = D.variance (GD s)

instance D.MaybeMean GeometricDistribution0 where
    maybeMean = Just . D.mean

instance D.MaybeVariance GeometricDistribution0 where
    maybeStdDev   = Just . D.stdDev
    maybeVariance = Just . D.variance

instance D.Entropy GeometricDistribution0 where
  entropy (GD0 s) = D.entropy (GD s)

instance D.MaybeEntropy GeometricDistribution0 where
  maybeEntropy = Just . D.entropy

instance D.DiscreteGen GeometricDistribution0 where
  genDiscreteVar (GD0 s) g = MWC.geometric0 s g

instance D.ContGen GeometricDistribution0 where
  genContVar d g = fromIntegral `liftM` D.genDiscreteVar d g

-- | Create geometric distribution.
geometric0 :: Double                -- ^ Success rate
           -> GeometricDistribution0
geometric0 x
  | x >= 0 && x <= 1 = GD0 x
  | otherwise        =
    error $ "Statistics.Distribution.Geometric.geometric: probability must be in [0,1] range. Got " ++ show x


module Q.Options.Black76
  (
    Black76(..)
  , atmf
  , euOption
  , eucall
  , euput
  , dPlus
  , dMinus
  )
  where

import           Numeric.GSL (derivCentral)
import           Q.Options (DF, Delta (Delta), Forward (..), Gamma (Gamma), OptionType (..),
                            Premium (Premium), Strike (..), TimeScaleable (scale), TotalVar (..),
                            Valuation (Valuation, vPremium), Vega (Vega), Vol (..),
                            YearFrac (YearFrac), discount)
import           Q.Options.ImpliedVol.TimeSlice (TimeSlice (..))
import           Statistics.Distribution (cumulative, density)
import           Statistics.Distribution.Normal (standard)
import Q.Types (LogRelStrike)

data Black76 = Black76 {
    b76F   :: Forward
  , b76DF  :: DF
  , b76T   :: YearFrac
  , b76Vol :: Vol
}

-- | At the money forward strike.
atmf :: Black76 -> Strike
atmf Black76{..} = Strike f
  where (Forward f) = b76F

-- | European option valuation with black 76
euOption :: Black76 -> OptionType -> Strike -> Valuation
euOption Black76{..} cp k = Valuation premium delta vega gamma where
  (Forward f) = b76F
  n           = cumulative standard
  (Vol sigmaSqt) = scale b76T b76Vol
  d1          = dPlus  b76F b76Vol k b76T
  d2          = dMinus b76F b76Vol k b76T
  nd1         = n d1
  nd2         = n d2
  callDelta   = b76DF `discount` nd1
  putDelta    = b76DF `discount` (- (n (-d1)))
  vega        = Vega  $ b76DF `discount` density standard d1 * f * sigmaSqt
  gamma       = Gamma $ b76DF `discount` density standard d1 / (f * sigmaSqt)
  premium  = Premium $ case cp of
    Call -> b76DF `discount` (f * nd1 - nd2 * k')
    Put  -> b76DF `discount` (n (-d2) * k' - n (-d1) * f)
    where (Strike k') = k
  delta | cp == Call = Delta callDelta
        | otherwise = Delta putDelta

-- | see 'euOption'
euput :: Black76 -> Strike -> Valuation
euput b76 =  euOption b76 Put

-- | see 'euOption'
eucall :: Black76 -> Strike -> Valuation
eucall b76 = euOption b76 Call

dPlus :: Forward -> Vol -> Strike -> YearFrac -> Double
dPlus (Forward f) (Vol sigma) (Strike k) (YearFrac t) =
  recip (sigma * sqrt t) * (log (f/k) + 0.5 * sigma * sigma * t)

dMinus :: Forward -> Vol -> Strike -> YearFrac -> Double
dMinus (Forward f) (Vol sigma) (Strike k) (YearFrac t) =
  recip (sigma * sqrt t) * (log (f/k) - 0.5 * sigma * sigma * t)

instance TimeSlice Black76 Strike where
  totalVar Black76{..} _ = TotalVar $ vol * vol * t
    where (Vol vol) = b76Vol
          (YearFrac t) = b76T
  dW _ _  = 0
  d2W _ _ = 0

  impliedDensity b76 (Strike k) = let
    dk k' = fst (derivCentral 1e-4 (\k2 -> let (Premium v) = vPremium (eucall b76 (Strike k2)) in v) k')
    in fst (derivCentral 1e-4 dk k)

instance TimeSlice Black76 LogRelStrike  where
  totalVar Black76{..} _ = TotalVar $ vol * vol * t
    where (Vol vol) = b76Vol
          (YearFrac t) = b76T
  dW _ _  = 0
  d2W _ _ = 0

  impliedDensity bs k = impliedDensity f k where
    f :: LogRelStrike  -> TotalVar
    f = totalVar bs


{-# LANGUAGE FlexibleContexts                         #-}
{-# LANGUAGE GADTs                                    #-}
{-# LANGUAGE TypeInType                               #-}
{-# LANGUAGE TypeOperators                            #-}
{-# OPTIONS_GHC -fplugin GHC.TypeLits.KnownNat.Solver #-}
{-# OPTIONS_GHC -fplugin GHC.TypeLits.Normalise       #-}

module AOC.Common.Numeric (
    fft
  , ifft
  , convolve
  , rconvolve
  , zconvolve
  , FFT.FFTWReal
  ) where

import           Data.Complex
import           GHC.TypeNats
import qualified Data.Array.CArray         as CA
import qualified Data.Array.IArray         as IA
import qualified Data.Ix                   as Ix
import qualified Data.Vector.Generic       as VG
import qualified Data.Vector.Generic.Sized as SVG
import qualified Foreign.Storable          as FS
import qualified Math.FFT                  as FFT
import qualified Math.FFT.Base             as FFT

fft :: (FFT.FFTWReal a, VG.Vector v (Complex a))
    => SVG.Vector v n (Complex a)
    -> SVG.Vector v n (Complex a)
fft = SVG.withVectorUnsafe $
        fromCA
      . FFT.dft
      . toCA

ifft
    :: (FFT.FFTWReal a, VG.Vector v (Complex a))
    => SVG.Vector v n (Complex a)
    -> SVG.Vector v n (Complex a)
ifft = SVG.withVectorUnsafe $
        fromCA
      . FFT.idft
      . toCA

fromCA
    :: (FS.Storable a, VG.Vector v (Complex a))
    => CA.CArray Int (Complex a)
    -> v (Complex a)
fromCA v = VG.generate (Ix.rangeSize (IA.bounds v)) (v IA.!)

toCA
    :: (FS.Storable a, VG.Vector v (Complex a))
    => v (Complex a)
    -> CA.CArray Int (Complex a)
toCA v = IA.listArray (0, VG.length v - 1) (VG.toList v)

-- | FFT-based convolution
convolve
    :: ( VG.Vector v (Complex a)
       , KnownNat n, 1 <= n
       , KnownNat m, 1 <= m
       , FFT.FFTWReal a
       )
    => SVG.Vector v n (Complex a)
    -> SVG.Vector v m (Complex a)
    -> SVG.Vector v (n + m - 1) (Complex a)
convolve x y = ifft $ fft x' * fft y'
  where
    x' = x SVG.++ 0
    y' = y SVG.++ 0

-- | FFT-based real-valued convolution
rconvolve
    :: ( VG.Vector v (Complex a)
       , VG.Vector v a
       , KnownNat n, 1 <= n
       , KnownNat m, 1 <= m
       , FFT.FFTWReal a
       )
    => SVG.Vector v n a
    -> SVG.Vector v m a
    -> SVG.Vector v (n + m - 1) a
rconvolve x y = SVG.map realPart $ convolve (SVG.map (:+ 0) x) (SVG.map (:+ 0) y)

-- | FFT-based integral convolution
zconvolve
    :: ( VG.Vector v (Complex Double)
       , VG.Vector v Double
       , VG.Vector v a
       , KnownNat n, 1 <= n
       , KnownNat m, 1 <= m
       , Integral a
       )
    => SVG.Vector v n a
    -> SVG.Vector v m a
    -> SVG.Vector v (n + m - 1) a
zconvolve x y = SVG.map (round @Double) $
    rconvolve (SVG.map fromIntegral x) (SVG.map fromIntegral y)