{-# LANGUAGE CPP #-}
#if __GLASGOW_HASKELL__ > 707
{-# LANGUAGE AllowAmbiguousTypes #-}
#endif
{-# LANGUAGE OverlappingInstances #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE NoMonomorphismRestriction #-}
{-# LANGUAGE ConstraintKinds #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE UndecidableInstances #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE StandaloneDeriving #-}
{-# LANGUAGE ViewPatterns #-}
{-# LANGUAGE QuasiQuotes #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE KindSignatures #-}
{-# LANGUAGE TypeFamilies #-}
module T1 where
import GHC.Exts (Constraint)
import DimMat
import Numeric.Units.Dimensional.Prelude
import Numeric.Units.Dimensional

import Text.PrettyPrint.ANSI.Leijen hiding (dot)

import GHC.TypeLits

import qualified Numeric.LinearAlgebra as H
import qualified Prelude as P

{- Example from http://ctms.engin.umich.edu/CTMS/index.php?example=InvertedPendulum&section=ControlStateSpace -}

-- not really stated on that page, but if you go back a couple of pages in their derivation
-- you can see that the type of u is a 1x1 matrix whose sole element is a force

massOfCart = (0.5 :: Double) *~ kilo gram
massOfPendulum = 0.2 *~ kilo gram
coefficientOfFrictionForCart = 0.1 *~ (newton / meter / second)
lengthToPendulumCenterOfMass = 0.3 *~ meter
massMomentOfInertiaOfPendulum = 0.006 *~ (kilo gram * meter * meter)
g = (9.8::Double) *~ (meter / second ^ pos2)

p = massMomentOfInertiaOfPendulum * (massOfCart + massOfPendulum)
  + (massOfCart*massOfPendulum*lengthToPendulumCenterOfMass*lengthToPendulumCenterOfMass)

a22 = negate (massMomentOfInertiaOfPendulum + massOfPendulum * lengthToPendulumCenterOfMass * lengthToPendulumCenterOfMass)
                  * coefficientOfFrictionForCart / p
a23 = (massOfPendulum * massOfPendulum * g * lengthToPendulumCenterOfMass * lengthToPendulumCenterOfMass) / p
a42 = negate (massOfPendulum * lengthToPendulumCenterOfMass * coefficientOfFrictionForCart) / p
a43 = massOfPendulum * g * lengthToPendulumCenterOfMass*(massOfCart + massOfPendulum)/p

{-

>>> inv aSmall `sub` aInvSmall
2><2     1                      m
s        -8.881784197001252e-16 -2.220446049250313e-16
m^-1 s^2 0.0                    0.0

>>> inv aSmall `multiply` aSmall
2><2   1                     m s^-1
1      1.0                   -6.732461810265988e-15
m^-1 s 5.014435047745458e-19 0.9999999999999999

-}
aSmall = [matD| a22, a23; a42, a43 |]

aInvSmall = scale (_1 / det aSmall)
    [matD| a43, negate a23; negate a42, a22 |]

-- id2 = Proxy :: Proxy (0 * a)

colMat :: D '(a, '[ra, '[]]) e -> D '(a, '[ra, '[]]) e
colMat x = x

{-
charEq lam1 lam2 x1 x2 =
        (det ( aSmall `sub` [matD| lam1, _0; _0, lam2 |]),
         colMat $ scale lam1 x1 `sub` dot aSmall x1,
         colMat $ scale lam2 x2 `sub` dot aSmall x2,
         hconcat x1 x2)
         --
         -- scale lam1 x1 `sub` dot aSmall x1, -- `asTypeOf` (undefined :: D '(s, ['[t],'[ ]]) Double),
         -- scale lam2 x2 `sub` dot aSmall x2) --  `asTypeOf` (undefined :: D '(s, ['[t],'[ ]]) Double),
         -- hconcat x1 x2 `asTypeOf` (undefined :: D '(t, ['[v],'[u]]) Double))
-- still has ambiguity. Do the SVD instead?
-}

{-
>>> cmap (Scale' ((2::Double) *~ second)) aSmall
2><2 1                    m s^-1
1    -0.36363636363636365 5.3454545454545475
m^-1 -0.9090909090909091  62.36363636363637

alternative to scale. Also more fancy examples (say swap metres and seconds
powers)...

-}
newtype Scale' f a = Scale' (Quantity f a)
instance (Num a, a ~ Double,
          x ~ Quantity d a,
          y ~ Quantity d' a,
          Mul f d d')
        => ApplyAB (Scale' f a) x y where
  applyAB (Scale' n) x = n * x

b21 = (massMomentOfInertiaOfPendulum + (massOfPendulum * lengthToPendulumCenterOfMass * lengthToPendulumCenterOfMass)) / p
b41 = massOfPendulum * lengthToPendulumCenterOfMass / p

-- example state value
x = [matD| 1.0 *~ meter;
           0.2 *~ (meter / second);
           _0 :: Dimensionless Double;
           0.1 *~ (second^neg1) |]

dx = scale (_1 / (1 *~ second)) x

-- example control input
u = [matD| (0 :: Double) *~ newton |]

--xDot = (a `multiply` x) `add` (b `multiply` u)

y = [matD| 1 *~ meter; _0 :: Dimensionless Double  |]

-- :t isLTI (1 *~ second) x u y, given that x u and y have monomorphic
-- types properly infers constraints on the a,b,c,d arguments!
isLTI time x u y a b c d =
    (scale (_1 /time) (colMat x) `add` dot a x `add` dot b (colMat u),
     colMat $ y `add` dot c x `add` dot d u)

testIsLTI =
  (\ a b c d -> case isLTI (1 *~ second) x u y a (colMat b) c (colMat d) of
   _ -> do
    print $ vsep
        [text "A = " </> indent 0 (pretty (zeroes `asTypeOf` a)),
         text "B = " </> indent 0 (pretty (zeroes `asTypeOf` b)),
         text "C = " </> indent 0 (pretty (zeroes `asTypeOf` c)),
         text "D = " </> indent 0 (pretty (zeroes `asTypeOf` d))]
    ) undefined undefined undefined undefined

{- | data type encoding units required by
http://en.wikibooks.org/wiki/Control_Systems/State-Space_Equations#State-Space_Equations

To refresh:

> dxs/dt = A xs + B us
> ys = C xs + D us

the units of d/dt are 1/iv
-}
class LiSystem (iv :: *) (xs :: (*, [[*]])) (ys :: (*, [[*]])) (us :: (*, [[*]]))
            (a :: (*, [[*]])) (b :: (*, [[*]])) (c :: (*, [[*]])) (d :: (*, [[*]]))
instance LiSystemCxt dxs iv xs ys us a b c d
    => LiSystem iv xs ys us a b c d

type LiSystemCxt dxs iv xs ys us a b c d =
   (Dot a xs dxs,
    Dot b us dxs,
    Dot c xs ys,
    Dot d us ys,
    DotS iv dxs xs
    )

-- | identity function but constrains types
isLiTuple ::
    (LiSystem iv xs ys us a b c d,
    t ~ (D a e, D b e, D c e, D d e)) =>
    t -> t
isLiTuple x = x

isExample ::
     (LiSystem DTime
        '(DLength , '[[DVelocity, DPlaneAngle, DAngularVelocity]])
        '(DLength, '[ '[DPlaneAngle]])
        '(DForce, '[ '[] ] )
        a b c d, t ~ (D a e, D b e, D c e, D d e)) =>
    t -> t
isExample x = x

pendulum = case (a',b',c',d') of
             sys -> isExample sys -- `const` evaluate sys
         where
           a' = [matD| _0, _1, _0, _0;
                       _0, a22, a23, _0;
                       _0, _0, _0, _1;
                       _0, a42, a43, _0 |]
           b' = [matD| _0;
                       b21;
                       _0;
                       b41 |]
           c' = [matD| _1, _0, _0, _0;
                       _0, _0, _1, _0 |]
           d' = zeroes

{-
-- poles (a,_,_,_) = eigenvalues a

-}
-- causes no problems for AV
-- evaluatePendulum = evaluate pendulum

{-
evaluate ::
    ( -- require all matrices to be at least 1x1
      -- really ought to be part of the MultiplyCxt
     a ~ [_1 ': _2, '[] ': _3],
     b ~ [_4 ': _5, '[] ': _6],
     c ~ [_7 ': _8, '[] ': _9],
     d ~ [_a ': _b, '[] ': _c],
     H.Field e,
     LiSystemCxt dxs iv xs ys us a b c d) =>
    (DimMat a e, DimMat b e, DimMat c e, DimMat d e)
    -> DimMat [xs, '[ '[] ]] e
    -> DimMat [us, '[ '[] ]] e
    -> (DimMat [dxs, '[ '[] ]] e, DimMat [ys, '[ '[] ]] e)
    -}
evaluate (a,b,c,d) x u = case (a `dot` x) `add` (b `dot` u) of
    xDot -> case (c `dot` x) `add` (d `dot` u) of
     y -> (xDot, y)

{-

ctrb = let (a,b,_,_) = pendulum
        in [blockD| b, a `multiply` b, a `multiply` a `multiply` b, a `multiply` a `multiply` a `multiply` b |]
        -}


{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE UndecidableInstances #-}

module Section_4_5.TwoLayer where

import Data.List

import Numeric.LinearAlgebra
import System.Random.MWC
import System.Random.MWC.Distributions (normal)

crossEntropyError :: Vector Double -> Vector Double -> Double
crossEntropyError y t =
  let delta = 1.0e-7
   in -1 * (sumElements $ t * cmap (\x -> log (x + delta)) y)

class Indexes c t where
  indexes :: c t -> [IndexOf c]

instance Container Vector t => Indexes Vector t where
  indexes v = [0 .. size v - 1]

instance (Num t, Container Vector t) => Indexes Matrix t where
  indexes m = let (x, y) = size m in (,) <$> [0..x-1] <*> [0..y-1]

numericalGradient :: (Container c Double, Additive (c Double), Indexes c Double) => (c Double -> Double) -> c Double -> c Double
numericalGradient f x =
  foldr add (assoc (size x) 0.0 []) $
    flip map (indexes x) $ \i ->
      let df1 = f (x `add` assoc (size x) 0.0 [(i,  h)])
          df2 = f (x `add` assoc (size x) 0.0 [(i, -h)])
          dfdx = (df1 - df2) / (2*h)
       in assoc (size x) 0.0 [(i, dfdx)]
    where h = 1.0e-4 -- 0.0001

type ActivationFunction = Vector Double -> Vector Double

-- | シグモイド関数
sigmoid :: ActivationFunction
sigmoid = cmap (\x -> 1 / (1 + exp (-x)))

-- | ソフトマックス関数
softmax :: ActivationFunction
softmax a =
  let c = maxElement a
      exp_a = cmap (\x -> exp (x - c)) a
      sum_exp_a = sumElements exp_a
   in cmap (/sum_exp_a) exp_a

type Layer = ( Matrix Double      -- 重み行列
             , Vector Double      -- バイアス
             , ActivationFunction -- 活性化関数
             )
type Network = [Layer]

-- | 学習前のニューラルネットワークを生成する
initNetwork :: [Int]                -- 各層のユニット数
            -> [ActivationFunction] -- 各層の活性化関数
            -> IO Network
initNetwork inputs as =
  withSystemRandom $ \gen ->
    sequence (zipWith (build gen) inouts as)
  where
    inouts :: [(Int, Int)]
    inouts = zip inputs (tail inputs)
    build :: GenIO -> (Int, Int) -> ActivationFunction -> IO Layer
    build gen (x, y) a = do
      w <- (x >< y) <$> (sequence $ replicate (x * y) (normal 0 0.01 gen))
      b <- vector <$> (sequence $ replicate y (normal 0 0.01 gen))
      pure (w, b, a)

-- | 認識（推論）を行う
predict :: Network -> Vector Double -> Vector Double
predict nw x = foldl (\x (w, b, h) -> h (x <# w + b)) x nw

-- | 損失関数の値を求める
loss :: [Vector Double] -> [Vector Double] -> Network -> Double
loss xs ts nw =
  let batchSize = genericLength xs
      ys = map (predict nw) xs
   in sum (zipWith crossEntropyError ys ts) / batchSize

-- | 認識精度を求める
accuracy :: [Vector Double] -> [Vector Double] -> Network -> Double
accuracy xs ts nw =
  let batchSize = genericLength xs
      ys = map (predict nw) xs
      yis = map maxIndex ys
      tis = map maxIndex ts
      correctCount = genericLength (filter id (zipWith (==) yis tis))
   in correctCount / batchSize

dismantle :: [a] -> [(a, a -> [a])]
dismantle xs = map (\i -> (xs!!i, (\x -> (take i xs) ++ [x] ++ (drop (i+1) xs)))) [0..length xs - 1] -- TODO more efficient implementation

updateNetwork :: ([Vector Double] -> [Vector Double] -> Network -> Double) -> [Vector Double] -> [Vector Double] -> Network -> Network
updateNetwork loss xs ts network = map (uncurry $ updateLayer xs ts) $ dismantle network
  where
    rate :: Double
    rate = 0.1
    updateLayer :: [Vector Double] -> [Vector Double] -> Layer -> (Layer -> Network) -> Layer
    updateLayer xs ts (w, b, a) mkNw =
      let dw = numericalGradient (\w -> loss xs ts (mkNw (w, b, a))) w
          db = numericalGradient (\b -> loss xs ts (mkNw (w, b, a))) b
       in (w - rate `scale` dw, b - rate `scale` db, a)


{-# LANGUAGE DataKinds #-}
{-# LANGUAGE ParallelListComp #-}
{-# LANGUAGE TupleSections #-}

module Main where

import Data.Random
import Data.Random.Distribution.MultiNormal
import Numeric.Kalman

import Control.Monad (replicateM)
--import GHC.TypeLits

import Graphics.Rendering.Chart.Easy
import Graphics.Rendering.Chart.Backend.Cairo

import Numeric.LinearAlgebra.Static hiding (mean)

--import Control.Monad (replicateM)
--import Data.Maybe (fromJust)

import qualified Data.List as L
import qualified Numeric.LinearAlgebra as LA

main :: IO ()
--main = putStrLn "Test suite not yet implemented"

--main = runUKFTest1
--main = runNoisySine1
--main = runGSTest1
main = runGSTest2

---------------------------
--- Gaussian sum filter ---
---------------------------
runGSTest2 :: IO ()
runGSTest2 = do
  putStrLn "Running Gaussian sum test 2"
  let numSamples = 1000
      dt = 0.01
      tVals = take numSamples [0, dt ..]
      truey = [(head tVals, 0), (last tVals, 0)] :: [(Double, Double)]
      truem = [(head tVals, 0), (last tVals, 0)] :: [(Double, Double)]
      trueb = [(head tVals, 0), (last tVals, 0)] :: [(Double, Double)]
      truew = [(head tVals, 1), (last tVals, 1)] :: [(Double, Double)]

  let intErr = 0.04
      slpErr = 0.01
      yErr   = 0.03

  let initEst = [(1/2, MultiNormal (vector [2,-0.5,0.1])  (sym . diag . vector $ [intErr, slpErr, yErr])),
                 (1/2, MultiNormal (vector [-2,0.5,0.1]) (sym . diag . vector $ [intErr, slpErr, yErr]))]
                :: [(Double, MultiNormal (R 3))]

      measMat = row . vector $ [0,0,1]
      measModel = const (const measMat)
      measErr = diag . vector $ [0.0001] :: Sq 1

      sysModel :: Double -> R 3 -> R 3
      sysModel t cur =
        let [b,m,y] = LA.toList (extract cur)
        in vector [b, m, m*t + (2*b^2+b-1)]
      sysErr = sym . diag . vector $ [0,0,0]

  measurements <- replicateM numSamples (sample $ MultiNormal (vector [0]) (sym measErr))

  let samples = [(t, head . LA.toList . extract $ s) | t <- tVals | s <- measurements]
      gsfilter = scanl
                 (\[(_w1, est1), (_w2, est2)] (t, meas) ->
                    let pred1 = runUKFPrediction sysModel (const sysErr) t est1
                        pmu1  = mu pred1
                        pcov1 = unSym . cov $ pred1
                        like1 = sym $ measMat <> pcov1 <> tr measMat + measErr
                        neww1 = pdf (MultiNormal (measModel t pmu1 #> pmu1) like1) meas
                        updt1 = runKFUpdate measModel (const (sym measErr)) t pred1 meas

                        pred2 = runUKFPrediction sysModel (const sysErr) t est2
                        pmu2  = mu pred2
                        pcov2 = unSym . cov $ pred2
                        like2 = sym $ measMat <> pcov2 <> tr measMat + measErr
                        neww2 = pdf (MultiNormal (measModel t pmu2 #> pmu2) like2) meas
                        updt2 = runKFUpdate measModel (const (sym measErr)) t pred2 meas

                        scale = neww1 + neww2

                    in [(neww1 / scale, updt1), (neww2 / scale, updt2)]
                 ) initEst $ zip tVals measurements

  let gsy1 = [(t, last . LA.toList . extract . mu . snd . head $ y) | t <- tVals | y <- gsfilter]
      gsy2 = [(t, last . LA.toList . extract . mu . snd . last $ y) | t <- tVals | y <- gsfilter]
      gsy  = zipWith (\t [(w1, est1), (w2, est2)] ->
                        let p1 = last . LA.toList . extract . mu $ est1
                            p2 = last . LA.toList . extract . mu $ est2
                        in (t, w1 * p1 + w2 * p2))
             tVals gsfilter

  let gsm1 = [(t, (!!1) . LA.toList . extract . mu . snd . head $ y) | t <- tVals | y <- gsfilter]
      gsm2 = [(t, (!!1) . LA.toList . extract . mu . snd . last $ y) | t <- tVals | y <- gsfilter]
      gsm  = zipWith (\t [(w1, est1), (w2, est2)] ->
                        let p1 = (!!1) . LA.toList . extract . mu $ est1
                            p2 = (!!1) . LA.toList . extract . mu $ est2
                        in (t, w1 * p1 + w2 * p2))
             tVals gsfilter

  let gsb1 = [(t, (!!0) . LA.toList . extract . mu . snd . head $ y) | t <- tVals | y <- gsfilter]
      gsb2 = [(t, (!!0) . LA.toList . extract . mu . snd . last $ y) | t <- tVals | y <- gsfilter]
      gsb  = zipWith (\t [(w1, est1), (w2, est2)] ->
                        let p1 = (!!0) . LA.toList . extract . mu $ est1
                            p2 = (!!0) . LA.toList . extract . mu $ est2
                        in (t, w1 * p1 + w2 * p2))
             tVals gsfilter

  let gsw1 = zipWith (\t [(w1,_),_] -> (t,w1)) tVals gsfilter
      gsw2 = zipWith (\t [_,(w2,_)] -> (t,w2)) tVals gsfilter
      gsw  = zipWith (\t [(w1, _), (w2, _)] -> (t,w1 + w2))
             tVals gsfilter

  let [(w1,cur1), (w2, cur2)] = last gsfilter
      final1 = (w1, LA.toList . extract . mu $ cur1)
      final2 = (w2, LA.toList . extract . mu $ cur2)

  print final1
  print final2
  print $ w1 + w2

  toFile def "gs2-y.png" $ do
    layout_title .= "Gaussian sums"
    setColors [opaque black, opaque blue
              ,opaque red, opaque green, opaque orange
              --,opaque hotpink
              ,opaque lightgrey, opaque lightgrey
              ,opaque lightslategrey, opaque lightslategrey
              ]
    plot (line   "True signal"  [truey])
    plot (points "Measurements" samples)
    plot (points "G1 estimates" gsy1)
    plot (points "G2 estimates" gsy2)
    plot (points "GS estimates" gsy)

  toFile def "gs2-m.png" $ do
    layout_title .= "Gaussian sums"
    setColors [opaque black --, opaque blue
              ,opaque red, opaque green, opaque orange
              --,opaque hotpink
              ,opaque lightgrey, opaque lightgrey
              ,opaque lightslategrey, opaque lightslategrey
              ]
    plot (line   "True signal"  [truem])
    plot (points "G1 estimates" gsm1)
    plot (points "G2 estimates" gsm2)
    plot (points "GS estimates" gsm)

  toFile def "gs2-b.png" $ do
    layout_title .= "Gaussian sums"
    setColors [opaque black --, opaque blue
              ,opaque red, opaque green, opaque orange
              --,opaque hotpink
              ,opaque lightgrey, opaque lightgrey
              ,opaque lightslategrey, opaque lightslategrey
              ]
    plot (line   "True signal"  [trueb])
    plot (points "G1 estimates" gsb1)
    plot (points "G2 estimates" gsb2)
    plot (points "GS estimates" gsb)

  toFile def "gs2-w.png" $ do
    layout_title .= "Gaussian sums"
    setColors [opaque black --, opaque blue
              ,opaque red, opaque green, opaque orange
              --,opaque hotpink
              ,opaque lightgrey, opaque lightgrey
              ,opaque lightslategrey, opaque lightslategrey
              ]
    plot (line   "True signal"  [truew])
    plot (points "G1 estimates" gsw1)
    plot (points "G2 estimates" gsw2)
    plot (points "GS estimates" gsw)

runGSTest1 :: IO ()
runGSTest1 = do
  putStrLn "Running Gaussian sum test 1"
  let numSamples = 1000
      dt = 0.01
      tVals = take numSamples [0, dt ..]
      truey = [(head tVals, 0), (last tVals, 0)] :: [(Double, Double)]
      truem = [(head tVals, 0), (last tVals, 0)] :: [(Double, Double)]
      trueb = [(head tVals, 0), (last tVals, 0)] :: [(Double, Double)]
      truew = [(head tVals, 1), (last tVals, 1)] :: [(Double, Double)]

  let intErr = 0.04
      slpErr = 0.01
      yErr   = 0.03

  let initEst = [(1/2, MultiNormal (vector [2,-0.5,0.1])  (sym . diag . vector $ [intErr, slpErr, yErr])),
                 (1/2, MultiNormal (vector [-2,0.5,0.1]) (sym . diag . vector $ [intErr, slpErr, yErr]))]
                :: [(Double, MultiNormal (R 3))]

      measMat = row . vector $ [0,0,1]
      measModel = const (const measMat)
      measErr = diag . vector $ [0.0001] :: Sq 1

      sysModel :: Double -> R 3 -> Sq 3
      sysModel t _cur = matrix [1,0,0, 0,1,0, 1,t,0]
      sysErr = sym . diag . vector $ [0,0,0]

  measurements <- replicateM numSamples (sample $ MultiNormal (vector [0]) (sym measErr))

  let samples = [(t, head . LA.toList . extract $ s) | t <- tVals | s <- measurements]
      gsfilter = scanl
                 (\[(_w1, est1), (_w2, est2)] (t, meas) ->
                    let pred1 = runKFPrediction sysModel (const sysErr) t est1
                        pmu1  = mu pred1
                        pcov1 = unSym . cov $ pred1
                        like1 = sym $ measMat <> pcov1 <> tr measMat + measErr
                        neww1 = pdf (MultiNormal (measModel t pmu1 #> pmu1) (like1)) meas
                        updt1 = runKFUpdate measModel (const (sym measErr)) t pred1 meas

                        pred2 = runKFPrediction sysModel (const sysErr) t est2
                        pmu2  = mu pred2
                        pcov2 = unSym . cov $ pred2
                        like2 = sym $ measMat <> pcov2 <> tr measMat + measErr
                        neww2 = pdf (MultiNormal (measModel t pmu2 #> pmu2) (like2)) meas
                        updt2 = runKFUpdate measModel (const (sym measErr)) t pred2 meas

                        scale = neww1 + neww2

                    in [(neww1 / scale, updt1), (neww2 / scale, updt2)]
                 ) initEst $ zip tVals measurements

  let gsy1 = [(t, last . LA.toList . extract . mu . snd . head $ y) | t <- tVals | y <- gsfilter]
      gsy2 = [(t, last . LA.toList . extract . mu . snd . last $ y) | t <- tVals | y <- gsfilter]
      gsy  = zipWith (\t [(w1, est1), (w2, est2)] ->
                        let p1 = last . LA.toList . extract . mu $ est1
                            p2 = last . LA.toList . extract . mu $ est2
                        in (t, w1 * p1 + w2 * p2))
             tVals gsfilter

  let gsm1 = [(t, (!!1) . LA.toList . extract . mu . snd . head $ y) | t <- tVals | y <- gsfilter]
      gsm2 = [(t, (!!1) . LA.toList . extract . mu . snd . last $ y) | t <- tVals | y <- gsfilter]
      gsm  = zipWith (\t [(w1, est1), (w2, est2)] ->
                        let p1 = (!!1) . LA.toList . extract . mu $ est1
                            p2 = (!!1) . LA.toList . extract . mu $ est2
                        in (t, w1 * p1 + w2 * p2))
             tVals gsfilter

  let gsb1 = [(t, (!!0) . LA.toList . extract . mu . snd . head $ y) | t <- tVals | y <- gsfilter]
      gsb2 = [(t, (!!0) . LA.toList . extract . mu . snd . last $ y) | t <- tVals | y <- gsfilter]
      gsb  = zipWith (\t [(w1, est1), (w2, est2)] ->
                        let p1 = (!!0) . LA.toList . extract . mu $ est1
                            p2 = (!!0) . LA.toList . extract . mu $ est2
                        in (t, w1 * p1 + w2 * p2))
             tVals gsfilter

  let gsw1 = zipWith (\t [(w1,_),_] -> (t,w1)) tVals gsfilter
      gsw2 = zipWith (\t [_,(w2,_)] -> (t,w2)) tVals gsfilter
      gsw  = zipWith (\t [(w1, _), (w2, _)] -> (t,w1 + w2))
             tVals gsfilter

  let [(w1,cur1), (w2, cur2)] = last gsfilter
      final1 = (w1, LA.toList . extract . mu $ cur1)
      final2 = (w2, LA.toList . extract . mu $ cur2)

  print final1
  print final2
  print $ w1 + w2

  toFile def "gs1-y.png" $ do
    layout_title .= "Gaussian sums"
    setColors [opaque black, opaque blue
              ,opaque red, opaque green, opaque orange
              --,opaque hotpink
              ,opaque lightgrey, opaque lightgrey
              ,opaque lightslategrey, opaque lightslategrey
              ]
    plot (line   "True signal"  [truey])
    plot (points "Measurements" samples)
    plot (points "G1 estimates" gsy1)
    plot (points "G2 estimates" gsy2)
    plot (points "GS estimates" gsy)

  toFile def "gs1-m.png" $ do
    layout_title .= "Gaussian sums"
    setColors [opaque black --, opaque blue
              ,opaque red, opaque green, opaque orange
              --,opaque hotpink
              ,opaque lightgrey, opaque lightgrey
              ,opaque lightslategrey, opaque lightslategrey
              ]
    plot (line   "True signal"  [truem])
    plot (points "G1 estimates" gsm1)
    plot (points "G2 estimates" gsm2)
    plot (points "GS estimates" gsm)

  toFile def "gs1-b.png" $ do
    layout_title .= "Gaussian sums"
    setColors [opaque black --, opaque blue
              ,opaque red, opaque green, opaque orange
              --,opaque hotpink
              ,opaque lightgrey, opaque lightgrey
              ,opaque lightslategrey, opaque lightslategrey
              ]
    plot (line   "True signal"  [trueb])
    plot (points "G1 estimates" gsb1)
    plot (points "G2 estimates" gsb2)
    plot (points "GS estimates" gsb)

  toFile def "gs1-w.png" $ do
    layout_title .= "Gaussian sums"
    setColors [opaque black --, opaque blue
              ,opaque red, opaque green, opaque orange
              --,opaque hotpink
              ,opaque lightgrey, opaque lightgrey
              ,opaque lightslategrey, opaque lightslategrey
              ]
    plot (line   "True signal"  [truew])
    plot (points "G1 estimates" gsw1)
    plot (points "G2 estimates" gsw2)
    plot (points "GS estimates" gsw)

---------------------------------------
--- Unscented Kalman filter testing ---
---------------------------------------
runUKFTest1 :: IO ()
runUKFTest1 = do
  -- polar coordinate transform test
  let numSamples = 500
      dt = 0.01
      tVals = take (numSamples + 1) [0,dt..] :: [Double] -- +1 because t = 0 is our initial estimate

      initSys = MultiNormal (vector [0, 1]) (sym . diag . vector $ [4e-4, 2.7e-5] :: Sym 2)
      initEst = MultiNormal (vector [1, pi/2]) (sym . diag . vector $ [0.02^2, (15*pi/180)^2] :: Sym 2)

  let coordTrans :: R 2 -> R 2
      coordTrans polar =
        let [r,t] = LA.toList $ extract polar
        in vector [r * cos t, r * sin t]

  measurements <- map coordTrans <$> replicateM numSamples (sample initEst)

  let samples = map (\meas ->
                       let [x,y] = LA.toList $ extract meas
                       in (x,y)
                    ) measurements

      ukf   = scanl (runUKF (const coordTrans)
                            (const (sym . diag . vector $ [0.02^2, (15*pi/180)^2] :: Sym 2))
                            (const id)
                            (const (sym . diag . vector $ [0,0])) 1)
                    initSys measurements

      ukpts = map (\kf ->
                       let [x,y] = LA.toList . extract $ mu kf
                       in (x,y)
                  ) ukf

  let linMeas :: R 2 -> L 2 2
      linMeas est =
        let [x, _] = LA.toList $ extract est
        in matrix [1,0, 0,x]

  let ekf = scanl (runEKF (const coordTrans)
                          (const linMeas)
                          (const (sym . diag . vector $ [0.02^2, (15*pi/180)^2] :: Sym 2))
                          (const id)
                          (const (const eye))
                          (const (sym . diag . vector $ [0,0])) 1)
                  initSys measurements

      ekpts = map (\kf ->
                       let [x,y] = LA.toList . extract $ mu kf
                       in (x,y)
                    ) ekf

      -- predTest = runUKFPrediction id (sym . diag . vector $ [0,0]) initSys
      -- updtTest = runUKFUpdate coordTrans (sym . diag . vector $ [0.02^2, (15*pi/180)^2] :: Sym 2) initSys (head measurements)

  -- print $ extract . mu $ predTest
  -- print $ extract . unSym . cov $ predTest
  -- putStrLn "\n"
  -- print $ extract . mu $ updtTest
  -- print $ extract . unSym . cov $ updtTest
  -- putStrLn "\n"
  -- print $ last ukpts
  -- print $ last ekpts

  toFile def "noisy-coords.png" $ do
    layout_title .= "Noisy coord transform"
    setColors [opaque black, opaque blue
              ,opaque red, opaque green
              --,opaque hotpink
              ,opaque lightgrey, opaque lightgrey
              ,opaque lightslategrey, opaque lightslategrey
              ]
    plot (points "True signal"  [(0,1)])
    plot (points "Measurements" samples)
    plot (points "UKF estimates" ukpts)
    plot (points "EKF estimates" ekpts)

  toFile def "noisy-params.png" $ do
    layout_title .= "Noisy coord transform"
    setColors [opaque black, opaque black
              ,opaque green, opaque green
              ,opaque red, opaque red
              --,opaque hotpink
              ,opaque lightgrey, opaque lightgrey
              ,opaque lightslategrey, opaque lightslategrey
              ]
    plot (line "True X signal"  [[(head tVals, 0), (last tVals, 0)]])
    plot (line "True Y signal"  [[(head tVals, 1), (last tVals, 1)]])
    plot (points "UKF X estimates" $ zipWith (\t (x,_) -> (t, x)) tVals ukpts)
    plot (points "UKF Y estimates" $ zipWith (\t (_,y) -> (t, y)) tVals ukpts)
    plot (points "EKF X estimates" $ zipWith (\t (x,_) -> (t, x)) tVals ekpts)
    plot (points "EKF Y estimates" $ zipWith (\t (_,y) -> (t, y)) tVals ekpts)

---------------------
--- Noisy sine # 1---
---------------------
-- In this example, we estimate the parameters (amplitude, period, vertical shift, horizontal shift)
-- defining sinusoidal evolution.
-- Figure 3 in my notes

runNoisySine1 :: IO ()
runNoisySine1 = do

  let (a, p, v) = (0.7, 2.2*pi, -0.5) -- (hidden) internal system parameters
      (a0, p0, v0) = (0.8, 2*pi, -0.6)    -- initial estimates
      --(aSig, pSig, vSig) = (0.2, 0.2*pi, 0.001) -- initial uncertainties
      measCov = 0.01

      --- Figure 1 in my notes
      --numSamples = 499
      --dt = 0.002
      --- Figure 2 in my notes
      numSamples = 199
      dt = 0.02
      tVals = take (numSamples + 1) [0,dt..] -- +1 because t = 0 is our initial estimate

      --initSys = vector [a, p, v] -- initial (true) system: (parameter1, parameter2, state)
      --initEst = MultiNormal (vector [a0, p0, v0, a0*p0, v0]) (sym . diag . vector $ [0.04, 0.4, 0.04, 0.001, 0.001] :: Sym 5)
      initEst = MultiNormal (vector [a0, p0, v0, a0*p0, v0])
                (sym . matrix $ [0.04,  0,    0,    1e-5, -2e-3,
                                 0,     0.4,  0,    1e-2, 12e-3,
                                 0,     0,    0.04, 1e-4, 1e-3,
                                 1e-5,  1e-2, 1e-4, 1e-3, 1e-4,
                                 -2e-3, 12e-3, 1e-3, 1e-4, 1e-3] :: Sym 5)

      evolMat :: R 5 -> Sq 5
      evolMat cur =
        let pCur  = (extract cur) LA.! 1
            p2dt  = pCur * pCur * dt
        in matrix [1,0,0,0,0, 0,1,0,0,0, 0,0,1,0,0, 0,0,p2dt,1,-p2dt, 0,0,0,dt,1]

      evolSys :: Double -> R 5 -> R 5
      evolSys t cur =
        let [a,p,v,ap,b] = LA.toList $ extract cur
        in vector [a, p, v, a * p * cos (p * t), a * sin (p * t) + v]

      sysCovariance = sym . diag . vector $ [0,0,0,5e-4,7e-5] :: Sym 5

      measMat = row (vector [0,0,0,0,1]) :: L 1 5
      measTrans = (measMat #>)
      measCovariance = sym . diag . vector $ [measCov] :: Sym 1
      measurementNoise = MultiNormal (vector [0]) measCovariance

      trueSystem = [vector [a, p, v, a * p * cos(p*t), a * sin(p*t) + v] | t <- tVals] :: [R 5]

  --trueSystem   <- simulateSystem numSamples sysModel initSys
  measurements <- mapM (\sys -> do
                            noise <- sample measurementNoise
                            return $ measTrans sys + noise
                       ) $ tail trueSystem

  --mapM_ print $ map (LA.toList . extract) trueSystem
  -- print $ map ((LA.! 0) . extract) measurements

  let truesine = zip tVals $ map ((LA.! 4) . extract) trueSystem
      samples  = zip (tail tVals) $ map ((LA.! 0) . extract) measurements

      kfiltered = scanl (runKF (const (const measMat)) (const measCovariance) (const evolMat) (const sysCovariance) 1) initEst measurements
      ekf = scanl (\sys (t,meas) ->
                     runEKF (\_t sys -> measMat #> sys)
                     (\_t _sys -> measMat)
                     (const measCovariance)
                     evolSys
                     (const evolMat)
                     (const sysCovariance)
                     t sys meas
                  ) initEst $ zip (tail tVals) measurements

      ukf = scanl (\sys (t,meas) ->
                     runUKF (\_t sys -> measMat #> sys)
                     (const measCovariance)
                     evolSys
                     (const sysCovariance)
                     t sys meas
                  ) initEst $ zip (tail tVals) measurements

      kfstates  = zipWith (\t kf -> (t, (extract . mu $ kf) LA.! 4)) tVals kfiltered
      ekfstates = zipWith (\t kf -> (t, (extract . mu $ kf) LA.! 4)) tVals ekf
      ukfstates = zipWith (\t kf -> (t, (extract . mu $ kf) LA.! 4)) tVals ukf

      [kfa,kfp,kfv] = L.transpose $ zipWith (\t kf -> map (t,) . take 3 . LA.toList . extract . mu $ kf) tVals kfiltered
      [ekfa,ekfp,ekfv] = L.transpose $ zipWith (\t kf -> map (t,) . take 3 . LA.toList . extract . mu $ kf) tVals ekf
      [ukfa,ukfp,ukfv] = L.transpose $ zipWith (\t kf -> map (t,) . take 3 . LA.toList . extract . mu $ kf) tVals ukf

      [kfaCovs, kfpCovs, kfvCovs, kfsCovs] = L.transpose $
                                             map (\kf -> let kfcov = extract . unSym $ cov kf
                                                         in [kfcov LA.! 0 LA.! 0, kfcov LA.! 1 LA.! 1,
                                                             kfcov LA.! 2 LA.! 2, kfcov LA.! 4 LA.! 4])
                                             kfiltered

      kfsUpper1 = zipWith (\(t, kfs) kfsCov -> (t, kfs + sqrt kfsCov)) kfstates kfsCovs
      kfsLower1 = zipWith (\(t, kfs) kfsCov -> (t, kfs - sqrt kfsCov)) kfstates kfsCovs

      kfsUpper2 = zipWith (\(t, kfs) kfsCov -> (t, kfs + 2 * sqrt kfsCov)) kfstates kfsCovs
      kfsLower2 = zipWith (\(t, kfs) kfsCov -> (t, kfs - 2 * sqrt kfsCov)) kfstates kfsCovs

      kfaUpper1 = zipWith (\(t, kfs) kfsCov -> (t, kfs + sqrt kfsCov)) kfa kfaCovs
      kfaLower1 = zipWith (\(t, kfs) kfsCov -> (t, kfs - sqrt kfsCov)) kfa kfaCovs

      kfaUpper2 = zipWith (\(t, kfs) kfsCov -> (t, kfs + 2 * sqrt kfsCov)) kfa kfaCovs
      kfaLower2 = zipWith (\(t, kfs) kfsCov -> (t, kfs - 2 * sqrt kfsCov)) kfa kfaCovs

      kfpUpper1 = zipWith (\(t, kfs) kfsCov -> (t, kfs + sqrt kfsCov)) kfp kfpCovs
      kfpLower1 = zipWith (\(t, kfs) kfsCov -> (t, kfs - sqrt kfsCov)) kfp kfpCovs

      kfpUpper2 = zipWith (\(t, kfs) kfsCov -> (t, kfs + 2 * sqrt kfsCov)) kfp kfpCovs
      kfpLower2 = zipWith (\(t, kfs) kfsCov -> (t, kfs - 2 * sqrt kfsCov)) kfp kfpCovs

      kfvUpper1 = zipWith (\(t, kfs) kfsCov -> (t, kfs + sqrt kfsCov)) kfv kfvCovs
      kfvLower1 = zipWith (\(t, kfs) kfsCov -> (t, kfs - sqrt kfsCov)) kfv kfvCovs

      kfvUpper2 = zipWith (\(t, kfs) kfsCov -> (t, kfs + 2 * sqrt kfsCov)) kfv kfvCovs
      kfvLower2 = zipWith (\(t, kfs) kfsCov -> (t, kfs - 2 * sqrt kfsCov)) kfv kfvCovs

      ksmoothed = scanr1 (\pres fut -> runKS (const evolMat) (const sysCovariance) 1 fut pres) kfiltered
      ksstates  = zipWith (\t ks -> (t, (extract . mu $ ks) LA.! 4)) tVals ksmoothed

  --mapM_ print $ map (LA.toList . extract . mu) kfiltered
  --mapM_ print kfa

  toFile def "noisy-sine-points.png" $ do
    layout_title .= "Noisy sine (state estimate)"
    setColors [opaque black, opaque blue
              ,opaque red, opaque purple, opaque hotpink, opaque green
              ,opaque lightgrey, opaque lightgrey
              ,opaque lightslategrey, opaque lightslategrey]
    plot (line   "True signal"  [truesine])
    plot (points "Measurements" samples)
    plot (line   "KF estimated signal" [kfstates])
    plot (line   "EKF estimated signal" [ekfstates])
    plot (line   "UKF estimated signal" [ukfstates])
    plot (line   "KS smoothed signal"  [ksstates])
    plot (line   "Filtered 1 standard deviation" [kfsUpper1])
    plot (line   "Filtered 1 standard deviation" [kfsLower1])
    plot (line   "Filtered 2 standard deviations" [kfsUpper2])
    plot (line   "Filtered 2 standard deviations" [kfsLower2])

  let truea = [(head tVals, a), (last tVals, a)]
      truep = [(head tVals, p), (last tVals, p)]
      truev = [(head tVals, v), (last tVals, v)]

  toFile def "noisy-sine-amplitude.png" $ do
    layout_title .= "Noisy sine (amplitude estimate)"
    setColors [opaque black, opaque red, opaque purple, opaque hotpink
              ,opaque lightgrey, opaque lightgrey
              ,opaque lightslategrey, opaque lightslategrey
              ]
    plot (line   "True signal"  [truea])
    plot (line   "KF estimated signal" [kfa])
    plot (line   "EKF estimated signal" [ekfa])
    plot (line   "UKF estimated signal" [ukfa])
    plot (line   "Filtered 1 standard deviation" [kfaUpper1])
    plot (line   "Filtered 1 standard deviation" [kfaLower1])
    plot (line   "Filtered 2 standard deviations" [kfaUpper2])
    plot (line   "Filtered 2 standard deviations" [kfaLower2])

  toFile def "noisy-sine-period.png" $ do
    layout_title .= "Noisy sine (period estimate)"
    setColors [opaque black, opaque red, opaque purple, opaque hotpink
              ,opaque lightgrey, opaque lightgrey
              ,opaque lightslategrey, opaque lightslategrey
              ]
    plot (line   "True signal"  [truep])
    plot (line   "KF estimated signal" [kfp])
    plot (line   "EKF estimated signal" [ekfp])
    plot (line   "UKF estimated signal" [ukfp])
    plot (line   "Filtered 1 standard deviation" [kfpUpper1])
    plot (line   "Filtered 1 standard deviation" [kfpLower1])
    plot (line   "Filtered 2 standard deviations" [kfpUpper2])
    plot (line   "Filtered 2 standard deviations" [kfpLower2])

  toFile def "noisy-sine-vertical.png" $ do
    layout_title .= "Noisy sine (vertical estimate)"
    setColors [opaque black, opaque red, opaque purple, opaque hotpink
              ,opaque lightgrey, opaque lightgrey
              ,opaque lightslategrey, opaque lightslategrey
              ]
    plot (line   "True signal"  [truev])
    plot (line   "KF estimated signal" [kfv])
    plot (line   "EKF estimated signal" [ekfv])
    plot (line   "UKF estimated signal" [ukfv])
    plot (line   "Filtered 1 standard deviation" [kfvUpper1])
    plot (line   "Filtered 1 standard deviation" [kfvLower1])
    plot (line   "Filtered 2 standard deviations" [kfvUpper2])
    plot (line   "Filtered 2 standard deviations" [kfvLower2])

  putStrLn $ "Kalman processing complete."


--{-# LANGUAGE UndecidableInstances,
--             FlexibleInstances,
--             FlexibleContexts,
--             TypeFamilies,
--             ScopedTypeVariables #-}
-----------------------------------------------------------------------------
-- |
-- Module      :  Numeric.Signal.Noise
-- Copyright   :  (c) Alexander Vivian Hugh McPhail 2010, 2014, 2015
-- License     :  BSD3
--
-- Maintainer  :  haskell.vivian.mcphail <at> gmail <dot> com
-- Stability   :  provisional
-- Portability :  uses Concurrency
--
-- Noise generation functions
--
-----------------------------------------------------------------------------

--             IncoherentInstances,

module Numeric.Signal.Noise (
                             pinkNoise
                            , spatialNoise
                            , powerNoise
                            ) where

-----------------------------------------------------------------------------

--import qualified Numeric.Signal as S

--import Complex

--import qualified Data.Array.IArray as I
--import Data.Ix

--import Data.Word

--import System.IO.Unsafe(unsafePerformIO)

--import qualified Data.List as L

--import Data.Binary

--import Foreign.Storable

import Numeric.LinearAlgebra

import qualified Data.Vector.Generic as GV

import qualified Numeric.GSL.Fourier as F

--import qualified Numeric.GSL.Histogram as H
--import qualified Numeric.GSL.Histogram2D as H2

--import qualified Numeric.Statistics.Information as SI

import Prelude hiding(filter)

--import Control.Monad(replicateM)

-------------------------------------------------------------------

-- | The method is briefly descirbed in Lennon, J.L. "Red-shifts and red
-- herrings in geographical ecology", Ecography, Vol. 23, p101-113 (2000)
--
-- Matlab version Written by Jon Yearsley  1 May 2004
-- j.yearsley@macaulay.ac.uk
--
-- Creates 1/f scale invariant spatial noise
spatialNoise :: Double    -- ^  β: spectral distribution
                         --    0: White noise
                         --   -1: Pink noise
                         --   -2: Brownian noise
             -> Int -> Int -- ^ matrix dimensions
             -> Int       -- ^ random seed
             -> Matrix Double
spatialNoise b r' c' s = let c = fromIntegral c'
                             r = fromIntegral r'
                             pre_x = linspace c' (0::Double,c-1)
                             post_x = linspace c' (c,1)
                             freq_x = GV.map (/c) $ vjoin [pre_x,post_x]
                             u = fromRows (replicate (2*r') freq_x)
                             pre_y = linspace r' (0::Double,r-1)
                             post_y = linspace r' (r,1)
                             freq_y = GV.map (/c) $ vjoin [pre_y,post_y]
                             v = fromColumns (replicate (2*c') freq_y)
                             s_f = cmap (**(b/2)) ((u**2) + (v**2))
                             s_f' = cmap (\x -> if isInfinite x then 0 else x) s_f
                             phi = reshape (2*c') (randomVector s Uniform (4*r'*c'))
                         in subMatrix (1,1) (r',c') $ fst $ fromComplex $ fromRows $ map F.ifft $ toRows $ ((complex $ s_f'**0.5) * (toComplex (cos(2*pi*phi),sin(2*pi*phi))))

-- | 1/f scale invariant noise
pinkNoise ::
    Double    -- ^  β: spectral distribution
              --    0: White noise
              --   -1: Pink noise
              --   -2: Brownian (red) noise
    -> Int       -- ^ samples
    -> Int       -- ^ random seed
    -> Vector Double
pinkNoise b s r = let pre = linspace s (0::Double,fromIntegral (s-1))
                      post = linspace s (fromIntegral s,1)
                      freq = vjoin [pre/(fromIntegral s),post/(fromIntegral s)]
                      s_f = GV.map (**(b/2)) (freq**2)
                      s_f' = GV.map (\x -> if isInfinite x then 0 else x) s_f
                      phi = randomVector r Uniform (2*s)
                  in subVector 0 s $ fst $ fromComplex $ F.ifft ((complex $ s_f'**0.5) * (toComplex (cos(2*pi*phi),sin(2*pi*phi))))

-- | generate noise from a power spectrum
powerNoise :: Vector Double   -- ^ the power spectrum
           -> Int             -- ^ random seed
           -> Vector Double
powerNoise psd r = let ln = GV.length psd
                       freq = vjoin [fromList [0],psd, (fromList . reverse . tail . toList) psd]
                       phi = randomVector r Uniform (2*ln)
                   in (fromIntegral ln) * (subVector 0 (ln-1) $ fst $ fromComplex $ F.ifft ((complex $ freq) * (toComplex (cos(2*pi*phi),sin(2*pi*phi)))))


module Math.IRT.Fisher
    ( FisherInfo(..)
    , fisherInfoObserved
    , fisherInfoExpected
    , l'
    , l''
    ) where

import Statistics.Distribution
import Math.IRT.Internal.Distribution

data FisherInfo = FisherInfo { items :: [Double]
                             , test  :: !Double
                             , sem   :: !Double
                             } deriving (Show)

-- |The observed Fisher Information
--  Applies a min/max prior to ensure a real-valued SEM
fisherInfoObserved :: (ContDistr d, DensityDeriv d) => Double -> [Bool] -> [d] -> FisherInfo
fisherInfoObserved theta resps params =
    let is = zipWith go resps params
        t  = sum is
        s  = sqrt (1 / t)
    in FisherInfo is t s
    where go u x = negate $ l'' u x theta

-- |The expected Fisher Information
fisherInfoExpected :: ContDistr d => Double -> [d] -> FisherInfo
fisherInfoExpected theta params =
    let is = map go params
        t  = sum is
        s  = sqrt (1 / t)
    in FisherInfo is t s
    where go x = let (cdf, ccdf, pdf) = pqDers x theta
                 in (square pdf) / (cdf * ccdf)

-- |The first derivative of `l`, whatever `l` is... Once again, the catIrt documentation disappoints.
-- According to catIrt, "u is the response, and x are the parameters."
-- We only implement the MLE route
l' :: ContDistr d => Bool -> d -> Double -> Double
l' u x theta =
    let (cdf, ccdf, pdf) = pqDers x theta
        denom = cdf * ccdf
        f k   = k * pdf / denom
    in case u of
         True  -> f ccdf
         False -> f cdf

-- |The second derivative of `l` (same one as in `l'`)
l'' :: (ContDistr d, DensityDeriv d) => Bool -> d -> Double -> Double
l'' u x theta =
    let (cdf, ccdf, pdf) = pqDers x theta
        pdf'             = densityDeriv x theta
    in case u of
         True  -> (((-1) / square cdf) * (square pdf))
                  + (1 / cdf * pdf')
         False -> negate $ ((1 / square ccdf) * square pdf)
                           + (1 / ccdf * pdf')

pqDers :: ContDistr d => d -> Double -> (Double, Double, Double)
pqDers x theta = let pComp = cumulative x theta
                     p'    = density x theta
                 in ( pComp
                    , (1 - pComp)
                    , p')

square :: Double -> Double
square = (^^ (2 :: Int))


module Scenes.Resch (getScene, vectorToFloats) where

import Numeric.LinearAlgebra (vector)

import qualified Group as G
import Group (pattern Group, optimizeGroup)
import Helpers (vectorToFloats)
import Layer (mkLayer)
import Line (pattern Line)
import Point (pattern Point, cartesianProduct, fromFloat)
import Rectangle (mkRectangle)
import Scene (Scene, mkSceneWithStyle)
import Shape (mirror, rotate)
import Style

getScene :: IO Scene
getScene = do
    let width      = 3 * (m + 0.5) * l
    let height     = sqrt 3 * (n + 1) * l
    let layerStyle = StyleAttrs { strokeColor=Just "#03161d"
                                , strokeWidth=Just 0.01
                                , fillColor=Nothing }
    pure $ mkSceneWithStyle width height layerStyle [triangles, frame]
  where
    l = 0.5  -- triangle circumradius
    n = 6    -- twice number in height
    m = 6    -- twice number in width

    lp = fromFloat l
    p0 = Point (-1) 0 * lp
    p1 = Point 0 (sqrt 3  / 3) * lp
    p2 = Point ((-1) / 2) (-sqrt 3 / 6) * lp

    -- Create a triangle from our three points.
    baseTriangle = Group [ Line p0 p1, Line p1 p2, Line p0 p2]
    -- Create a group from the base triangle rotated three times.
    rotationalVerticies = map (\r -> 2 * pi * r / 3 ) [0..2]
    t0 = mconcat $ map (\t -> rotate (Point 0 0) t baseTriangle) rotationalVerticies
    -- Mirror this group twice.
    t1 = t0 <> mirror p0 (p1 - p0) t0
    t2 = t1 <> mirror (Point (0.5 * l) 0) (Point 0 1) t1

    -- Clone the above group across a series of positions on our convas.
    xs = vectorToFloats $ (vector [0..5] - 6 / 2 + 0.25) * 3 * 0.5
    ys = vectorToFloats $ (vector [0..5] - 6 / 2 + 0.25) * sqrt 3 * 0.5
    t3 = optimizeGroup (G.translateGroupOverPoints (cartesianProduct xs ys) t2) 0.01
    triangles = G.toLayer "triangles" t3

  -- Create a framing rectangle around the group.
    topLeft = Point ((-1.5 * m - 0.25) * l) (0.5 * sqrt 3 * (n + 0.5) * l)
    bottomRight = Point ((1.5 * m - 0.25) * l) (-0.5 * sqrt 3 * (n + 0.5) * l)
    frame = mkLayer "frame" [mkRectangle topLeft bottomRight]


{-# LANGUAGE FlexibleInstances #-}
-----------------------------------------------------------------------------
-- |
-- Module     : Foreign.BLAS.Level3
-- Copyright  : Copyright (c) 2010, Patrick Perry <patperry@gmail.com>
-- License    : BSD3
-- Maintainer : Patrick Perry <patperry@gmail.com>
-- Stability  : experimental
--
-- Matrix-Matrix operations.
--

module Foreign.BLAS.Level3 (
    BLAS3(..),
    ) where

import Data.Complex
import Foreign( Ptr, Storable, with )

import Foreign.BLAS.Types
import Foreign.BLAS.Level2
import Foreign.BLAS.Double
import Foreign.BLAS.Zomplex

-- | Types with matrix-matrix operations.
class (BLAS2 a) => BLAS3 a where
    gemm  :: Trans -> Trans -> Int -> Int -> Int -> a -> Ptr a -> Int -> Ptr a -> Int -> a -> Ptr a -> Int -> IO ()
    symm  :: Side -> Uplo -> Int -> Int -> a -> Ptr a -> Int -> Ptr a -> Int -> a -> Ptr a -> Int -> IO ()
    hemm  :: Side -> Uplo -> Int -> Int -> a -> Ptr a -> Int -> Ptr a -> Int -> a -> Ptr a -> Int -> IO ()
    trmm  :: Side -> Uplo -> Trans -> Diag -> Int -> Int -> a -> Ptr a -> Int -> Ptr a -> Int -> IO ()
    trsm  :: Side -> Uplo -> Trans -> Diag -> Int -> Int -> a -> Ptr a -> Int -> Ptr a -> Int -> IO ()
    syrk  :: Uplo -> Trans -> Int -> Int -> a -> Ptr a -> Int -> a -> Ptr a -> Int -> IO ()
    syr2k :: Uplo -> Trans -> Int -> Int -> a -> Ptr a -> Int -> Ptr a -> Int -> a -> Ptr a -> Int -> IO ()
    herk  :: Uplo -> Trans -> Int -> Int -> a -> Ptr a -> Int -> a -> Ptr a -> Int -> IO ()
    her2k :: Uplo -> Trans -> Int -> Int -> a -> Ptr a -> Int -> Ptr a -> Int -> a -> Ptr a -> Int -> IO ()

withEnum :: (Enum a, Storable a) => Int -> (Ptr a -> IO b) -> IO b
withEnum = with . toEnum
{-# INLINE withEnum #-}

instance BLAS3 Double where
    gemm transa transb m n k alpha pa lda pb ldb beta pc ldc =
        withTrans transa $ \ptransa ->
        withTrans transb $ \ptransb ->
        withEnum m $ \pm ->
        withEnum n $ \pn ->
        withEnum k $ \pk ->
        with alpha $ \palpha ->
        withEnum lda $ \plda ->
        withEnum ldb $ \pldb ->
        with beta $ \pbeta ->
        withEnum ldc $ \pldc ->
            dgemm ptransa ptransb pm pn pk palpha pa plda pb pldb pbeta pc pldc
    {-# INLINE gemm #-}

    symm side uplo m n alpha pa lda pb ldb beta pc ldc =
        withSide side $ \pside ->
        withUplo uplo $ \puplo ->
        withEnum m $ \pm ->
        withEnum n $ \pn ->
        with alpha $ \palpha ->
        withEnum lda $ \plda ->
        withEnum ldb $ \pldb ->
        with beta $ \pbeta ->
        withEnum ldc $ \pldc ->
            dsymm pside puplo pm pn palpha pa plda pb pldb pbeta pc pldc
    {-# INLINE symm #-}

    hemm = symm
    {-# INLINE hemm #-}

    trmm side uplo transa diag m n alpha pa lda pb ldb =
        withSide side $ \pside ->
        withUplo uplo $ \puplo ->
        withTrans transa $ \ptransa ->
        withDiag diag $ \pdiag ->
        withEnum m $ \pm ->
        withEnum n $ \pn ->
        with alpha $ \palpha ->
        withEnum lda $ \plda ->
        withEnum ldb $ \pldb ->
            dtrmm pside puplo ptransa pdiag pm pn palpha pa plda pb pldb
    {-# INLINE trmm #-}

    trsm side uplo transa diag m n alpha pa lda pb ldb =
        withSide side $ \pside ->
        withUplo uplo $ \puplo ->
        withTrans transa $ \ptransa ->
        withDiag diag $ \pdiag ->
        withEnum m $ \pm ->
        withEnum n $ \pn ->
        with alpha $ \palpha ->
        withEnum lda $ \plda ->
        withEnum ldb $ \pldb ->
            dtrsm pside puplo ptransa pdiag pm pn palpha pa plda pb pldb
    {-# INLINE trsm #-}

    syrk uplo transa n k alpha pa lda beta pc ldc =
        withUplo uplo $ \puplo ->
        withTrans transa $ \ptransa ->
        withEnum n $ \pn ->
        withEnum k $ \pk ->
        with alpha $ \palpha ->
        withEnum lda $ \plda ->
        with beta $ \pbeta ->
        withEnum ldc $ \pldc ->
            dsyrk puplo ptransa pn pk palpha pa plda pbeta pc pldc
    {-# INLINE syrk #-}

    syr2k uplo transa n k alpha pa lda pb ldb beta pc ldc =
        withUplo uplo $ \puplo ->
        withTrans transa $ \ptransa ->
        withEnum n $ \pn ->
        withEnum k $ \pk ->
        with alpha $ \palpha ->
        withEnum lda $ \plda ->
        withEnum ldb $ \pldb ->
        with beta $ \pbeta ->
        withEnum ldc $ \pldc ->
            dsyr2k puplo ptransa pn pk palpha pa plda pb pldb pbeta pc pldc
    {-# INLINE syr2k #-}

    herk = syrk
    {-# INLINE herk #-}

    her2k = syr2k
    {-# INLINE her2k #-}

instance BLAS3 (Complex Double) where
    gemm transa transb m n k alpha pa lda pb ldb beta pc ldc =
        withTrans transa $ \ptransa ->
        withTrans transb $ \ptransb ->
        withEnum m $ \pm ->
        withEnum n $ \pn ->
        withEnum k $ \pk ->
        with alpha $ \palpha ->
        withEnum lda $ \plda ->
        withEnum ldb $ \pldb ->
        with beta $ \pbeta ->
        withEnum ldc $ \pldc ->
            zgemm ptransa ptransb pm pn pk palpha pa plda pb pldb pbeta pc pldc
    {-# INLINE gemm #-}

    symm side uplo m n alpha pa lda pb ldb beta pc ldc =
        withSide side $ \pside ->
        withUplo uplo $ \puplo ->
        withEnum m $ \pm ->
        withEnum n $ \pn ->
        with alpha $ \palpha ->
        withEnum lda $ \plda ->
        withEnum ldb $ \pldb ->
        with beta $ \pbeta ->
        withEnum ldc $ \pldc ->
            zsymm pside puplo pm pn palpha pa plda pb pldb pbeta pc pldc
    {-# INLINE symm #-}

    hemm side uplo m n alpha pa lda pb ldb beta pc ldc =
        withSide side $ \pside ->
        withUplo uplo $ \puplo ->
        withEnum m $ \pm ->
        withEnum n $ \pn ->
        with alpha $ \palpha ->
        withEnum lda $ \plda ->
        withEnum ldb $ \pldb ->
        with beta $ \pbeta ->
        withEnum ldc $ \pldc ->
            zhemm pside puplo pm pn palpha pa plda pb pldb pbeta pc pldc
    {-# INLINE hemm #-}

    trmm side uplo transa diag m n alpha pa lda pb ldb =
        withSide side $ \pside ->
        withUplo uplo $ \puplo ->
        withTrans transa $ \ptransa ->
        withDiag diag $ \pdiag ->
        withEnum m $ \pm ->
        withEnum n $ \pn ->
        with alpha $ \palpha ->
        withEnum lda $ \plda ->
        withEnum ldb $ \pldb ->
            ztrmm pside puplo ptransa pdiag pm pn palpha pa plda pb pldb
    {-# INLINE trmm #-}

    trsm side uplo transa diag m n alpha pa lda pb ldb =
        withSide side $ \pside ->
        withUplo uplo $ \puplo ->
        withTrans transa $ \ptransa ->
        withDiag diag $ \pdiag ->
        withEnum m $ \pm ->
        withEnum n $ \pn ->
        with alpha $ \palpha ->
        withEnum lda $ \plda ->
        withEnum ldb $ \pldb ->
            ztrsm pside puplo ptransa pdiag pm pn palpha pa plda pb pldb
    {-# INLINE trsm #-}

    syrk uplo transa n k alpha pa lda beta pc ldc =
        withUplo uplo $ \puplo ->
        withTrans transa $ \ptransa ->
        withEnum n $ \pn ->
        withEnum k $ \pk ->
        with alpha $ \palpha ->
        withEnum lda $ \plda ->
        with beta $ \pbeta ->
        withEnum ldc $ \pldc ->
            zsyrk puplo ptransa pn pk palpha pa plda pbeta pc pldc
    {-# INLINE syrk #-}

    syr2k uplo transa n k alpha pa lda pb ldb beta pc ldc =
        withUplo uplo $ \puplo ->
        withTrans transa $ \ptransa ->
        withEnum n $ \pn ->
        withEnum k $ \pk ->
        with alpha $ \palpha ->
        withEnum lda $ \plda ->
        withEnum ldb $ \pldb ->
        with beta $ \pbeta ->
        withEnum ldc $ \pldc ->
            zsyr2k puplo ptransa pn pk palpha pa plda pb pldb pbeta pc pldc
    {-# INLINE syr2k #-}

    herk uplo transa n k alpha pa lda beta pc ldc =
        withUplo uplo $ \puplo ->
        withTrans transa $ \ptransa ->
        withEnum n $ \pn ->
        withEnum k $ \pk ->
        with alpha $ \palpha ->
        withEnum lda $ \plda ->
        with beta $ \pbeta ->
        withEnum ldc $ \pldc ->
            zherk puplo ptransa pn pk palpha pa plda pbeta pc pldc
    {-# INLINE herk #-}

    her2k uplo transa n k alpha pa lda pb ldb beta pc ldc =
        withUplo uplo $ \puplo ->
        withTrans transa $ \ptransa ->
        withEnum n $ \pn ->
        withEnum k $ \pk ->
        with alpha $ \palpha ->
        withEnum lda $ \plda ->
        withEnum ldb $ \pldb ->
        with beta $ \pbeta ->
        withEnum ldc $ \pldc ->
            zher2k puplo ptransa pn pk palpha pa plda pb pldb pbeta pc pldc
    {-# INLINE her2k #-}


{-# LANGUAGE OverloadedStrings #-}
{-# 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
    , geometricE
    , geometric0
    , geometric0E
    -- ** Accessors
    , gdSuccess
    , gdSuccess0
    ) where

import Control.Applicative
import Control.Monad       (liftM)
import Data.Aeson          (FromJSON(..), ToJSON, Value(..), (.:))
import Data.Binary         (Binary(..))
import Data.Data           (Data, Typeable)
import GHC.Generics        (Generic)
import Numeric.MathFunctions.Constants (m_pos_inf, m_neg_inf)
import qualified System.Random.MWC.Distributions as MWC

import qualified Statistics.Distribution as D
import Statistics.Internal

----------------------------------------------------------------

-- | Distribution over [1..]
newtype GeometricDistribution = GD {
      gdSuccess :: Double
    } deriving (Eq, Typeable, Data, Generic)

instance Show GeometricDistribution where
  showsPrec i (GD x) = defaultShow1 "geometric" x i
instance Read GeometricDistribution where
  readPrec = defaultReadPrecM1 "geometric" geometricE

instance ToJSON GeometricDistribution
instance FromJSON GeometricDistribution where
  parseJSON (Object v) = do
    x <- v .: "gdSuccess"
    maybe (fail $ errMsg x) return  $ geometricE x
  parseJSON _ = empty

instance Binary GeometricDistribution where
  put (GD x) = put x
  get = do
    x <- get
    maybe (fail $ errMsg x) return  $ geometricE 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

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)

-- | Create geometric distribution.
geometric :: Double                -- ^ Success rate
          -> GeometricDistribution
geometric x = maybe (error $ errMsg x) id $ geometricE x

-- | Create geometric distribution.
geometricE :: Double                -- ^ Success rate
           -> Maybe GeometricDistribution
geometricE x
  | x >= 0 && x <= 1 = Just (GD x)
  | otherwise        = Nothing

errMsg :: Double -> String
errMsg x = "Statistics.Distribution.Geometric.geometric: probability must be in [0,1] range. Got " ++ show x

----------------------------------------------------------------

-- | Distribution over [0..]
newtype GeometricDistribution0 = GD0 {
      gdSuccess0 :: Double
    } deriving (Eq, Typeable, Data, Generic)

instance Show GeometricDistribution0 where
  showsPrec i (GD0 x) = defaultShow1 "geometric0" x i
instance Read GeometricDistribution0 where
  readPrec = defaultReadPrecM1 "geometric0" geometric0E

instance ToJSON GeometricDistribution0
instance FromJSON GeometricDistribution0 where
  parseJSON (Object v) = do
    x <- v .: "gdSuccess0"
    maybe (fail $ errMsg x) return  $ geometric0E x
  parseJSON _ = empty

instance Binary GeometricDistribution0 where
  put (GD0 x) = put x
  get = do
    x <- get
    maybe (fail $ errMsg x) return  $ geometric0E 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 = maybe (error $ errMsg0 x) id $ geometric0E x

-- | Create geometric distribution.
geometric0E :: Double                -- ^ Success rate
            -> Maybe GeometricDistribution0
geometric0E x
  | x >= 0 && x <= 1 = Just (GD0 x)
  | otherwise        = Nothing

errMsg0 :: Double -> String
errMsg0 x = "Statistics.Distribution.Geometric.geometric0: probability must be in [0,1] range. Got " ++ show x


module Lec7 where

-- Функторы
--
-- Функция map была бы логична не только для списков. Как бы она работала для Maybe?
-- f x = x + 1
-- x1 = Just 42
-- x2 = Nothing
--
-- -хочу применить f к x1, x2, точнее, к значению внутри
-- maybeMap fun (Just n) = Just (fun n)
-- maybeMap _ Nothing = Nothing
-- print $ maybeMap f x1  -> Just 43
-- print $ maybeMap f x2  -> Nothing
--
-- Допустим, я хочу уметь так же применять функции ко вторым элементам пары:
-- f x = x + 1
-- x1 = ("hello", 42)
-- x2 = ([1, 2, 3], 42.0)
-- хочу применить f ко второму элементу
-- pairMap fun (x, y) = (x, fun y)
-- print $ pairMap f x1  -> ("hello", 43)
-- print $ pairMap f x2  -> ([1, 2, 3], 43.0)
--
-- Есть и другие аналогичные ситуации. Когда значение (или значения) хранятся в какой-то структуре с дополнительными
-- данными.
-- Список - несколько значений и информация об их порядке
-- Maybe - одно значение или ноль
-- Пара - значение вместе с дополнительной информацией
--
-- Функции тоже можно понимать как такие структуры:
-- fun :: a -> b
-- получается, функция хранит значение типа B, доступ к которым осуществляется по аргументу типа a:
-- Например,
-- f 1 = "один"
-- f 2 = "два"
-- функция "хранит" внутри себя значения "один", "два", получить их можно, подставив соответствующий аргумент в функцию.
-- f x = x + 1
-- x1 = \x -> 2 * x
-- x2 = \x -> 3 * x
-- funMap f x1 - должна получиться функция, которая сопадает с x1, но после применения x1 дополнительно вызывает f
-- funMap fun x = \t -> fun (x t)
--
-- xx1 = funMap f x1
-- xx2 = funMap f x2
-- print $ xx1 42  -> 85
-- print $ xx2 42  -> 127
--
-- На самом деле все эти map функции уже есть. Они называются fmap или <$>, и они делают ровно то же, что делаем мы:
-- print $ fmap (+1) (Just 42)  -> Just 43
-- print $ (+1) <$> Just 42  -> Just 43
-- print $ (+1) <$> [10, 20, 30]  -> [11, 21, 31]
-- print $ (+1) <$> ("abc", 123)  -? ("abc", 124)
--
-- Это работает, потому что в Haskell есть класс типов Functor - функтор. Он определен так:
-- class Functor f where
--   fmap :: (a -> b) -> f a -> f b
--
-- У него есть реализации, например, примерно так выглядит реализация для Maybe.
-- instance Functor Maybe where
--   fmap f (Just x) = Just (f x)
--   fmap _ Nothing = Nothing
--
-- Если вы определяете функтор, вы должны гарантировать выполнение следующего правила. Функция id - это функция, которая
-- ничего не делает: id x = x, её тип a -> a
-- **Правило**:
-- 1) fmap id x == x (или эквивалентно) fmap id == id
-- 2) fmap (g . f) x == fmap g . fmap f
--
-- ## Аппликативные функторы
-- Расширение обычных функторов дополнительными операциями. Представьте опять же Maybe.
-- Было в функторах:
-- 1) Обычная функция
-- 2) Данные внутри структуры
-- 3) Применяем функцию к каждому данному внутри структуры
-- Теперь в аппликативных функторах можно дополнительно:
-- 1) Обычные функция внутри структуры
-- 2) Данные внутри структуры
-- 3) Применяем функции из структуры к данным. Как именно применять - очень сильно зависит от ситуации.
--
-- Пример с Maybe:
-- f = Just (+1) -- увеличение на 1, но внутри Maybe
-- x = Just 42 -- значение внутри Maybe
-- применение:
-- maybeApply (Just fun) (Just x) = Just (fun x)
-- maybeApply _ _ = Nothing
--
-- print $ maybeApply f x  -> Just 43
-- print $ maybeApply Nothing x  -> Nothing
-- print $ maybeApply (Just (*2)) Nothing  -> Nothing
--
-- А для списков?
-- f = [(+1), (*2)]
-- x = [10, 20, 30]
-- что может означать listApply, когда есть список функций и список значений?
-- Применить каждую функцию к каждому элементу, получить список новых значений.
--
-- listApply fs xs = concatMap (\f -> map f xs) fs
-- print $ listApply f x  -> [11, 21, 31, 20, 40, 60]
-- Немного философии, почему для списков это логично делать именно так. Можно считать, что Maybe - это вычисления,
-- в которых иногда отсутствуют результаты. А в списках можно считать, что мы делаем вычисления, в которых может быть
-- несколько результатов.
-- В Haskell есть класс типов Аппликативные функторы:
-- class Functor f => Applicative f where
--   pure :: a -> f a -- вставить значение в функтор
--   (<*>) :: f (a -> b) -> f a -> f b -- та самая функция применения, рассмотренная выше
--   -- и ещё несколько, см. дальше
--
-- pure - это функция, которая помещает значение внутрь функции, "самым естественным способом". Что это значит будет
-- понятно из требований к аппликативному функтору, см. дальше.
-- Для Maybe: pure x = Just x, для списка pure x = [x]
--
-- Допустим, хотим сложить значения в двух Maybe. Или в двух списках сложить попарно значения. Как это будет написано?
-- x1 = Just 110
-- x2 = Just 654
-- надо сложить. Сначала давайте получим Just (+110), потом применим её к Just 654:
-- (+) <*> x1 <*> x2
--
-- Аппликативные функторы должны уметь делать не только <*>, но и функцию pure, которая вводит обычное значение в функтор.
-- 'pure 42' в типе Maybe это 'Just 42' . 'pure 42' в типе списка: [42]
--
-- Правила для функторов:
--   pure id <*> v = v
--   pure f <*> pure x = pure (f x)
--   и др.
--
--
-- #Монады
-- Вычисления с Maybe - это вычисления, в которых может не получиться результата. Если мы делаем последовательность
-- вычислений, в которой в какой-то момент результата не получилось, то и дальше считаем, что результата нет.
-- mSqrt :: Double -> Maybe Double
-- mSqrt x | x < 0 = Nothing
--         | otherwise = Just $ sqrt x
-- [хотим вычислить (sqrt(x) + 1) / 2]
-- eval x = let sx = mSqrt x
--              y = (+1) <$> sx
--            in (/2) <$> y
-- eval 9 -> Just 2.0
-- eval (-9) -> Nothing
--
-- sqrt(x) + sqrt(y) + 1
-- eval2 x y = let sx = mSqrt x
--                 sy = mSqrt y
--                 sq_plus_sq = Just (+) <*> sx <*> sy -- liftA2 (+) sx sy
--                in
--                 (+1) <$> sq_plus_sq
-- eval2 9 16 -> Just 8.0
-- eval2 (-9) 16 -> Nothing
--
-- sqrt(sqrt(x) - 3)
-- eval3 x = let sx = mSqrt x
--               sx_3 = (\t -> t - 3) <$> sx
--               -- ?? не получается сделать ни с <*>, ни с <$>
--               -- есть Just y и есть mSqrt :: Double -> Maybe Double
--               z = case sx_3 of
--                      Just n -> mSqrt n
--                      Nothing -> Nothing
--              in
--               z
-- eval3 100 -> Just ...
-- eval3 (-100) -> Nothing
-- eval3 4 -> Nothing
--
-- Хочется вспомогательную функцию типа <*> и <$>, чтобы она имела такой заголовок
-- bind :: Maybe a -> (a -> Maybe b) -> Maybe b
--
-- Если подумать, то
-- 1) fmap - частный случай этой функции. Сравним
--    bind :: Maybe a -> (a -> Maybe b) -> Maybe b
--    fmap :: Maybe a -> (a -> b) -> Maybe b
--    fmap f t = bind t (\x -> pure $ f x) или fmap f t = bind t (pure . f)
-- 2) <*> - тоже частный случай. Попробуйте её сам выразить через bind
--
-- Реализуем bind для Maybe:
-- maybeBind :: Maybe a -> (a -> Maybe b) -> Maybe b
-- maybeBind Nothing _ = Nothing
-- maybeBind (Just x) f = f x
--
-- снова sqrt(sqrt(x) - 3)
-- eval3 x = let sx = mSqrt(x)
--               sx_3 = maybeBind sx (\t -> Just (t - 3))
--              in
--               maybeBind sx_3 mSqrt
-- eval3 100 -> Just ...
-- eval3 (-100) -> Nothing
-- eval3 4 -> Nothing
--
-- короткий синтаксис
-- eval3 x = do
--             sx <- mSqrt x
--             sx_3 <- sx - 3
--             mSqrt sx_3
-- ответы те же
--
-- Давайте остановимся на одной вспомогательной функции
-- maybeBind :: Maybe a -> (a -> Maybe b) -> Maybe b
-- Функции даётся значение, очередной шаг вычисления, она возвращает новое значение.
-- maybeBind :: Maybe a -> (a -> Maybe b) -> Maybe b
-- maybeBind Nothing _ = Nothing
-- maybeBind (Just x) f = f x
--
-- взятие корня
-- mSqrt :: Double -> Maybe Double
-- mSqrt x | x < 0 = Nothing
--         | otherwise = Just $ sqrt x
-- sqrt(sqrt(x) - 3)
--
-- eval4 :: Double -> Maybe Double
-- maybeBind (Just x) mSqrt - это sqrt(x) внутри maybe
-- eval4 x = maybeBind (mSqrt x) (\a -> mSqrt (a - 3))
--
-- Эта программа похожа на последловательность шагов:
-- 1. mSqrt
-- 2. \a -> mSqrt (a - 3)
--
-- Посчитаем теперь sqrt x + sqrt y + 1
-- Шаги:
-- 1. sqrt x
-- 2. sqrt y
-- 3. sqrt x + sqrt y + 1
--
-- eval5 :: Double -> Double -> Maybe Double
-- eval5 x y = maybeBind (mSqrt x) (\a -> maybeBind (mSqrt y) (\b -> Just (a + b + 1)))
--
-- maybeBind - это функция '>>=', которая существует для Maybe и других типов.
-- eval5 :: Double -> Double -> Maybe Double
-- eval5 x y = (mSqrt x) >>= (\a -> (mSqrt y) >>= (\b -> Just (a + b + 1)))
--
-- для подообной записи есть do-нотация
-- eval5 :: Double -> Double -> Maybe Double
-- eval5 x y = do
--              a <- mSqrt x
--              b <- mSqrt y
--              return (a + b + 1)
-- Получился "императивный" код. return для Maybe это Just.
-- Перепишем теперь eval4
-- eval4 :: Double -> Maybe Double
-- eval4 x = (mSqrt x) >>= (\a -> mSqrt (a - 3))
-- или в do-нотации
-- eval4 x = do
--             a <- mSqrt x
--             mSqrt (a - 3) -- сразу возвращает Maybe
--
-- Комплексное число 1 :+ 2 это 1 + 2i
-- import Data.Complex
-- lSqrt :: Complex Double -> [Complex Double]
-- lSqrt x = [sqrt x, -sqrt x]
--
-- eval5 :: Complex Double -> Complex Double -> [Complex Double]
-- eval5 x y = do
--              a <- lSqrt x
--              b <- lSqrt y
--              return (a + b + 1)
--
-- Монада в хаскеле:
-- class Applicative a => Monad (m a) where
-- (>>=) :: m a -> (a -> m b) -> m b
-- return :: a -> m a -- превращает обычное значение в значение в монаде
-- (>>) :: m a -> m b -> m b
-- (>>) x y = x >>= (\_ -> y)
--
-- Пример монады. Вычисления с журналированием действий.
data Log a =
  Log a [String]
  deriving (Show) -- храним значение и список сообщений

-- Будем пользоваться функциями, которые при вычислениях могут выдать несколько сообщений:
add1 :: Int -> Log Int
add1 x = Log (x + 1) ["added 1"]

--
mul2 :: Int -> Log Int
mul2 x = Log (2 * x) ["multiplied by 2", "that was hard"]

--
-- Скажем, что Log это монада и объясним, как комбинировать вычисления.
instance Functor Log where
  fmap f (Log v msg) = Log (f v) msg

--
instance Applicative Log where
  pure x = Log x []
  Log f newMsg <*> Log v msg = Log (f v) (msg ++ newMsg)

--
instance Monad Log where
-- (>>=) :: Log a -> (a -> Log b) -> Log b
  (Log x messages) >>= f =
    let (Log y newMessages) = f x
      in Log y (messages ++ newMessages)

  -- return :: a -> Log a
  return x = Log x []

--
-- 2 + x + 1 через функции add1 и mul2
eval6 :: Int -> Log Int
eval6 x = do
  a <- mul2 x
  add1 a -- или b <- add1 a; return b
-- eval6 10 -- должен вернуть 21 и журнал ["умножила на 2", "добавила 1"]

--sumIO :: IO ()
--sumIO = do
--  a <- getLine
--  b <- getLine
--  print (read a + read b)
sumIO :: IO Int
sumIO = do
    a <- getLine
    b <- getLine
    return (read a + read b)


{-# LANGUAGE BangPatterns #-}

{- |
Module      : Main
Description : Haskell translation of Neural Networks and Deep Learning code sample
Copyright   : (c) FUJII Satohi, 2016
License     : BSD3
Maintainer  : fujiisat@gmaill.com
Stability   : experimental
Portability : portable

My personal Haskell translation of Neural Networks and Deep Learning
code sample
(https://github.com/mnielsen/neural-networks-and-deep-learning,
Copyright (c) 2012-2015 Michael Nielsen)

-}

module Main where

import           Control.Monad         (replicateM, zipWithM)
import           Data.List             (foldl')
import qualified Data.Vector           as V
import           Numeric.LinearAlgebra (( #> ))
import qualified Numeric.LinearAlgebra as H
import           System.Environment    (getArgs)
import           System.Random         (newStdGen, randomIO)
import           System.Random.Shuffle (shuffle')
import           Text.Printf           (printf)

import           MnistData

type Vec = H.Vector H.R
type Mat = H.Matrix H.R
type Data = (Vec, Vec)
data Network = Network
               { numLayers :: !Int
               , sizes     :: V.Vector Int
               , biases    :: V.Vector Vec
               , weights   :: V.Vector Mat
               } deriving (Eq, Show)

main :: IO ()
main = do
  -- ([30, 10, 3.0], [784, 30, 10])
  ([epo, mbsize, e], networkArgs) <-  splitAt 3  <$> getArgs
  putStrLn $ "Epoch: " ++ epo
  putStrLn $ "Minibatch: " ++ mbsize
  putStrLn $ "Eta: " ++ e
  putStrLn $ "Network: " ++ show networkArgs
  (trainingData, validateDate, testData) <- loadDataWrapper
  network <- networkInit (map read networkArgs)
  !result <- sgd network (V.fromList trainingData) (read epo) (read mbsize) (read e) (Just $ V.fromList testData)
  -- print result
  return ()

networkInit :: [Int] -> IO Network
networkInit sizes = do
  b <- map H.flatten <$> mapM (`H.randn` 1) (tail sizes)
  w <- zipWithM H.randn sizes (tail sizes)
  return Network { numLayers = length sizes
                 , sizes     = V.fromList sizes
                 , biases    = V.fromList b
                 , weights   = V.fromList w
                 }

feedforward :: Network -> Vec -> Vec
feedforward (Network n s biases weights) a =
  V.foldl' (\ !a (w, b) -> sigmoidVec ((H.tr w #> a) + b) ) a $ V.zip weights biases

sgd :: Network
    -> V.Vector Data               -- Training data
    -> Int                         -- Epoch
    -> Int                         -- Minibatch size
    -> Double                      -- Eta
    -> Maybe (V.Vector (Vec, Int)) -- Test data
    -> IO Network
sgd network trainingData 0 miniBatchSize eta testData = return network
sgd network trainingData epohcs miniBatchSize eta testData = do
  miniBatches <- divide miniBatchSize <$> shuffleList trainingData
  let network' = V.foldl' (\ !n b -> updateMiniBatch n b eta) network miniBatches
  case testData of
    Just t  -> putStrLn $ printf "Epoch %d: %d / %d" epohcs (evaluate network' t) (length t)
    Nothing -> putStrLn $ printf "Epoch %d complete" epohcs
  sgd network' trainingData (epohcs - 1) miniBatchSize eta testData

updateMiniBatch :: Network -> V.Vector Data -> Double -> Network
updateMiniBatch network miniBatch eta = network { weights=weights', biases=biases'}
  where
    f :: (V.Vector Vec, V.Vector Mat) -> Data -> (V.Vector Vec, V.Vector Mat)
    f nabla mb = nablaNew
      where
        !deltaNabla = backprop network mb
        !nablaNew = ( V.zipWith (+) (fst nabla) (fst deltaNabla)
                    , V.zipWith (+) (snd nabla) (snd deltaNabla))

    (!nablaB, !nablaW) = V.foldl' f zeros miniBatch

    !zeros = ( V.map (\a -> a - a) $ biases network
             , V.map (\a -> a - a) $ weights network )

    !etaN = eta / fromIntegral (length miniBatch)
    !weights' = V.map (\(w, nw) -> w - H.cmap (etaN *) nw) $ V.zip (weights network) nablaW
    !biases'  = V.map (\(b, nb) -> b - H.cmap (etaN *) nb) $ V.zip (biases network) nablaB

backprop :: Network -> Data -> (V.Vector Vec, V.Vector Mat)
backprop network (x, y) = (V.fromList nablaB, V.fromList nablaW)
  where
    (ac0:ac1:activations, z:zs) = V.foldl' forward ([x], []) $ V.zip (biases network) (weights network)
--    !delta = costDerivative ac0 y * sigmoidPrimeVec z -- quadratic
    !delta = costDerivative ac0 y  -- cross entropy
    !nablaIni = ([delta], [H.tr (delta `H.outer` ac1)], delta)
    (nablaB, nablaW, _) = foldl' backward nablaIni
      $ zip3 zs (V.toList $ V.reverse $ weights network) activations

backward :: ([Vec], [Mat], Vec) -> (Vec, Mat, Vec) -> ([Vec], [Mat], Vec)
backward  (nabB, nabW, del) (z, w, ac) = (delta : nabB, H.tr (delta `H.outer` ac) : nabW, delta)
  where
    !sp = sigmoidPrimeVec z
    !delta = (w #> del) * sp

forward :: ([Vec], [Vec]) -> (Vec, Mat) -> ([Vec], [Vec])
forward (a@(ac:acs), zzs) (b, w) = (activation:a, z:zzs)
  where
    !z = (H.tr w #> ac) + b
    !activation = sigmoidVec z

costDerivative :: Vec -> Vec -> Vec
costDerivative outputActivations y = outputActivations - y

sigmoid :: Double -> Double
sigmoid z = 1.0 / (1.0 + exp (-z))

sigmoidVec :: Vec -> Vec
sigmoidVec = H.cmap sigmoid

sigmoidPrime :: Double -> Double
sigmoidPrime z = sigmoid z * (1 - sigmoid z)

sigmoidPrimeVec :: Vec -> Vec
sigmoidPrimeVec = H.cmap sigmoidPrime

evaluate :: Network -> V.Vector (Vec, Int) -> Int
evaluate network testData = length $ V.filter (uncurry (==)) testResults
  where
    testResults = V.map (\(x, y) -> (H.maxIndex (feedforward network x), y)) testData

divide :: Int -> V.Vector a -> V.Vector (V.Vector a)
divide n xs | V.null xs = V.empty
            | otherwise =  V.cons hd (divide n tl)
  where
    (hd, tl) = V.splitAt n xs

shuffleList :: V.Vector a -> IO (V.Vector a)
shuffleList xs = do
  let len = length xs
  gen <- newStdGen
  return $ V.fromList $ shuffle' (V.toList xs) len gen


-- By Christian Oliveros and Minmin Chen
module PingPong.Player.Stig.Kinematics where

import Control.Lens (view)
import Data.Bool (bool)
import Data.Geometry hiding (Vector, dot, head)
import Data.Maybe
import GHC.Float
import Numeric.LinearAlgebra hiding (Element)
import PingPong.Model
import PingPong.Player.Stig.General
import PingPong.Player.Stig.GeometryHelpers
import Prelude hiding ((<>))

-- Inverse and Forward Kinematics

-- | Calculate the forward kinematic matrix for a joint / link
calFwdMatrix :: Element -> Element -> [[Float]]
calFwdMatrix (Joint _ a) (Link _ l) =
  [ [cos a, - sin a, l * cos a],
    [sin a, cos a, l * sin a],
    [0, 0, 1]
  ]
calFwdMatrix _ _ = [[1, 0, 0], [0, 1, 0], [0, 0, 1]]

data ArmKinematicPart
  = ArmBase
      { footParam :: R
      }
  | ArmJoint
      { translationParam :: R
      }
  | ArmBat
      { translationParam :: R
      }
  deriving (Show)

type ArmKinematic = [ArmKinematicPart]

-- | Goes from rough representation to a more refined one
fromFootArmToArmKinematic :: Float -> Arm -> ArmKinematic
fromFootArmToArmKinematic foot arm = ArmBase (float2Double foot) : processArm 0 arm
  where
    processArm :: R -> Arm -> ArmKinematic
    processArm t [Link _ t2] = [ArmBat $ t + float2Double t2]
    processArm t (Link _ t2 : as) = processArm (t + float2Double t2) as
    processArm t (Joint _ _ : as) = ArmJoint t : processArm 0 as

-- | Goes from rough representation to a more refined one.
-- As well, get the motion vector
getArmKinematicAndMotion :: Float -> Arm -> (ArmKinematic, Motion)
getArmKinematicAndMotion foot arm =
  (fromFootArmToArmKinematic foot simpleArm, getCurrentJoints simpleArm)
  where
    simpleArm = simplifyArm arm

-- | Converts a motion list to a vector of joint coordinates
motionToJointVector :: Motion -> Vector R
motionToJointVector m = fromList $ map float2Double m

-- | Converts vector of joint coordinates to a motion list
jointVectorToMotion :: Vector R -> Motion
jointVectorToMotion v = map double2Float $ toList v

-- | Translation in the X axis matrix
translateXTrans :: R -> Matrix R
translateXTrans t =
  fromLists [[1, 0, t], [0, 1, 0], [0, 0, 1]] :: Matrix R

-- | Translation in the X axis inverse matrix
translateXInvTrans :: R -> Matrix R
translateXInvTrans t =
  fromLists [[1, 0, - t], [0, 1, 0], [0, 0, 1]] :: Matrix R

-- | Rotation matrix
rotateTrans :: R -> Matrix R
rotateTrans a =
  fromLists [[c, - s, 0], [s, c, 0], [0, 0, 1]] :: Matrix R
  where
    c = cos a
    s = sin a

-- | Rotation inverse matrix
rotateInvTrans :: R -> Matrix R
rotateInvTrans a =
  fromLists [[c, s, 0], [- s, c, 0], [0, 0, 1]] :: Matrix R
  where
    c = cos a
    s = sin a

-- | Create an homogeneous 2D point
homogeneousPoint :: R -> R -> Vector R
homogeneousPoint x y = fromList [x, y, 1]

-- | Zero Point in Homogenous Coordinates
homogeneousZero :: Vector R
homogeneousZero = homogeneousPoint 0 0

-- | 2d Point to Homogeneous Coordinates
pointToHomogenousPoint :: Point 2 Float -> Vector R
pointToHomogenousPoint p = homogeneousPoint (float2Double $ view xCoord p) (float2Double $ view yCoord p)

-- | Create an homogeneous 2D vector
homogeneousVector :: R -> R -> Vector R
homogeneousVector x y = fromList [x, y, 0]

-- | X Vector
homogeneousVectorX :: Vector R
homogeneousVectorX = homogeneousVector 1 0

-- | Y Vector
homogeneousVectorY :: Vector R
homogeneousVectorY = homogeneousVector 0 1

-- | Homogeneous 2D Identity Matrix
homogeneousIdent :: Matrix R
homogeneousIdent = ident 3

data ArmKinematicPartMatrix
  = ArmBaseMatrix
      { footMatrix :: Matrix R
      }
  | ArmJointMatrix
      { translationMatrix :: Matrix R,
        jointMatrix :: Matrix R
      }
  | ArmBatMatrix
      { batMatrix :: Matrix R
      }
  deriving (Show)

type ArmKinematicMatrix = [ArmKinematicPartMatrix]

-- | Get Transform from Part Matrix
getTrans :: ArmKinematicPartMatrix -> Matrix R
getTrans (ArmBaseMatrix t) = t
getTrans (ArmJointMatrix t r) = t <> r
getTrans (ArmBatMatrix t) = t

-- | Get Inverse Transform from Part Matrix
getTransInv :: ArmKinematicPartMatrix -> Matrix R
getTransInv (ArmBaseMatrix t) = t
getTransInv (ArmJointMatrix t r) = r <> t
getTransInv (ArmBatMatrix t) = t

-- | Calculate the transforms for forward kinematics
applyForwardKinematicTrans ::
  ArmKinematic -> Vector R -> ArmKinematicMatrix
applyForwardKinematicTrans arm v = toTrans arm motion
  where
    motion = toList v :: [R]

    toTrans :: ArmKinematic -> [R] -> ArmKinematicMatrix
    toTrans (ArmBase t : as) js = (ArmBaseMatrix $ translateXTrans t <> rotateTrans (pi / 2)) : toTrans as js
    toTrans (ArmJoint t : as) (j : js) = ArmJointMatrix (translateXTrans t) (rotateTrans j) : toTrans as js
    toTrans [ArmBat t] [] = [ArmBatMatrix $ translateXTrans t]

-- | Calculate the inverse transforms for forward kinematics
applyForwardKinematicTransInv ::
  ArmKinematic -> Vector R -> ArmKinematicMatrix
applyForwardKinematicTransInv arm v = reverse $ toTrans arm motion
  where
    motion = toList v :: [R]

    toTrans :: ArmKinematic -> [R] -> ArmKinematicMatrix
    toTrans (ArmBase t : as) js = (ArmBaseMatrix $ rotateInvTrans (pi / 2) <> translateXInvTrans t) : toTrans as js
    toTrans (ArmJoint t : as) (j : js) = ArmJointMatrix (translateXInvTrans t) (rotateInvTrans j) : toTrans as js
    toTrans [ArmBat t] [] = [ArmBatMatrix $ translateXInvTrans t]

-- | Apply Forward Kinematic Transformations
applyForwardKinematicMatrixTrans :: ArmKinematicMatrix -> Matrix R
applyForwardKinematicMatrixTrans = foldr ((<>) . getTrans) homogeneousIdent

-- | Apply Forward Kinematic Inverse Transformations (the list must go from end effecto to base)
applyForwardKinematicMatrixTransInv :: ArmKinematicMatrix -> Matrix R
applyForwardKinematicMatrixTransInv = foldr ((<>) . getTransInv) homogeneousIdent

-- | Compress the transforms of a Forward Kinematic Arm to only the joints and the end effector
compressForwardKinematicsJointsTrans :: ArmKinematicMatrix -> ArmKinematicMatrix
compressForwardKinematicsJointsTrans a = compress a homogeneousIdent
  where
    compress :: ArmKinematicMatrix -> Matrix R -> ArmKinematicMatrix
    compress [ArmBatMatrix t] m = [ArmBatMatrix (m <> t)]
    compress (ArmJointMatrix t r : as) m = ArmJointMatrix (m <> t) r : compress as homogeneousIdent
    compress (a : as) m = compress as (m <> getTrans a)

-- | Compress the transforms of a Forward Kinematic Arm to only the joints and the end base
compressForwardKinematicsJointsInvTrans :: ArmKinematicMatrix -> ArmKinematicMatrix
compressForwardKinematicsJointsInvTrans a = compress a homogeneousIdent
  where
    compress :: ArmKinematicMatrix -> Matrix R -> ArmKinematicMatrix
    compress [ArmBaseMatrix t] m = [ArmBaseMatrix (t <> m)]
    compress (ArmJointMatrix t r : as) m = ArmJointMatrix (t <> m) r : compress as homogeneousIdent
    compress (a : as) m = compress as (getTransInv a <> m)

-- | Calculate the 2D Homogeneous Jacobian of an arm [J^T|0]^T .
-- fwdMatTrans : Kinematic Forward Transforms (From base to end effector).
-- fwdMatTransInv : Kinematic Forward Inverse Transforms (From end effector to base).
-- xLocal : Position of Tool on End Effector Local Coordinates
calculateJacobian ::
  ArmKinematicMatrix -> ArmKinematicMatrix -> Vector R -> Matrix R
calculateJacobian fwdMatTrans fwdMatTransInv xLocal = jH
  where
    -- Compressed form guarantees that the links were applied to the joints
    fwdTransCompressed = compressForwardKinematicsJointsTrans fwdMatTrans
    fwdTransInvCompressed = compressForwardKinematicsJointsInvTrans fwdMatTransInv

    -- Derivative of a revolute Joint
    revoluteDeriv = fromLists [[0, -1, 0], [1, 0, 0], [0, 0, 0]] :: Matrix R

    rollingTransInv :: ArmKinematicMatrix -> Matrix R -> [Matrix R]
    rollingTransInv [ArmBaseMatrix t] m = [t <> m]
    rollingTransInv (a : as) m = roll : rollingTransInv as roll
      where
        roll = getTransInv a <> m
    calculate ::
      ArmKinematicMatrix -> Matrix R -> [Matrix R] -> [Vector R]
    calculate [ArmBatMatrix _] _ [_] = [] -- We reach the bat, we no longer need to process
    calculate (j@(ArmJointMatrix _ _) : as) m (endToJ : mIs) =
      (baseToJ <> revoluteDeriv <> endToJ) #> xLocal : calculate as baseToJ mIs
      where
        baseToJ = m <> getTrans j

    -- Jacobian in this form: [J^T|0]^T there is a row of 0s at the end due to 2d homogeneous coordinates
    jH =
      fromColumns $
        calculate fwdTransCompressed homogeneousIdent (rollingTransInv fwdTransInvCompressed homogeneousIdent)

-- | Get the Jacobian from the Homogeneous Jacobian (Note this doesn't work for the inverse of the Jacobian)
getJacobianFromHomogeneousJacobian :: Matrix R -> Matrix R
getJacobianFromHomogeneousJacobian jH = jH ?? (DropLast 1, All)

-- | Newton Raphson Threshold
newtonRaphsonThreshold :: (Num r, Ord r, Fractional r) => r -> r -> r
newtonRaphsonThreshold = threshold 0.0000001

-- | Calculates a Newton-Raphson IK step (if singular it fails). Returns the new Joints and its error agains target
-- a : Arm Description.
-- q : Current Joints.
-- xLocal : Tool in End-Effector Local Coordinates
-- xTargetGlobal : Target in Global Coordinates
newtonRaphsonStep ::
  ArmKinematic ->
  Vector R ->
  Vector R ->
  Vector R ->
  Maybe (Vector R, R)
newtonRaphsonStep a q xLocal xTargetGlobal
  | singular = Nothing
  | otherwise = Just (qn, norm_2 eN)
  where
    -- Forward Transforms
    fwdT = applyForwardKinematicTrans a q
    fwdTI = applyForwardKinematicTransInv a q
    jH = calculateJacobian fwdT fwdTI xLocal

    -- Bat Global Position
    batGlobal = applyForwardKinematicMatrixTrans fwdT #> xLocal

    -- Error
    e = xTargetGlobal - batGlobal

    -- Drop Homogeneous part
    er = subVector 0 2 e
    j = getJacobianFromHomogeneousJacobian jH

    -- Calculate delta q
    dq = - pinv j #> er

    dqNorm = norm_2 dq

    qnUnormalized = q + dq
    -- Re normalizes angles for more accuracy
    qn = fromList $ map normalizeAngle $ toList qnUnormalized

    -- New Bat Position
    batN = applyForwardKinematicMatrixTrans (applyForwardKinematicTrans a qn) #> xLocal

    -- New Error
    eN = xTargetGlobal - batN

    -- If the move was singular
    singular = rank jH < 2 && newtonRaphsonThreshold 0 dqNorm == 0

-- | Maximum Newton Raphson Step
newtonRaphsonMaxStep :: Int
newtonRaphsonMaxStep = 10000

-- | Really Bad Step
newtonRaphsonBadStep :: forall a. Floating a => a
newtonRaphsonBadStep = 2

-- | When to do a random Restart Newton Raphson Step
newtonRaphsonMaxRandomRestartStep :: Int
newtonRaphsonMaxRandomRestartStep = round (fromIntegral newtonRaphsonMaxStep / 10)

-- | NewtonRaphson Loop Iteration
newtonRaphsonIKIter ::
  Int ->
  Int ->
  ArmKinematic ->
  (Vector R, Vector R) ->
  Vector R ->
  (Vector R, R) ->
  (Vector R, R)
newtonRaphsonIKIter i j a (xLocal, xTargetGlobal) q (qBest, eBest)
  | i >= newtonRaphsonMaxStep = (qBest, eBest)
  | newtonRaphsonThreshold 0 eBest == 0 = (qBest, eBest)
  | singular = newtonRaphsonIKIter (i + 1) 0 a (xLocal, xTargetGlobal) qR (qBest, eBest)
  | needReset = newtonRaphsonIKIter (i + 1) 0 a (xLocal, xTargetGlobal) qR (qBest, eBest)
  | eN < eBest = newtonRaphsonIKIter (i + 1) (j + 1) a (xLocal, xTargetGlobal) qN (qN, eN)
  | otherwise = newtonRaphsonIKIter (i + 1) (j + 1) a (xLocal, xTargetGlobal) qN (qBest, eBest)
  where
    -- Random Vector
    qSize = size q
    pseudoInt = round $ fromIntegral i ** (dot q qBest * bool (eBest + 1) (1 / eBest) (eBest < 1 && 0 < eBest))
    qR = pi / 2 * (2 * randomVector pseudoInt Uniform qSize - 1)

    -- Perform the step
    step = newtonRaphsonStep a q xLocal xTargetGlobal

    singular = case step of
      Nothing -> True
      Just _ -> False

    (qN, eN) = fromJust step

    needReset =
      ( eN >= ((eBest / (1 + fromIntegral j)) * newtonRaphsonBadStep)
          || (newtonRaphsonThreshold 0 (norm_2 (qN - q)) == 0)
      )
        && j >= newtonRaphsonMaxRandomRestartStep

-- | Newton Raphson IK Algorithm
newtonRaphsonIK ::
  ArmKinematic ->
  (Vector R, Vector R) ->
  Vector R ->
  (Vector R, R)
newtonRaphsonIK a (xLocal, xTargetGlobal) q = newtonRaphsonIKIter 0 0 a (xLocal, xTargetGlobal) q (q, eNorm)
  where
    -- Forward Transforms
    fwdT = applyForwardKinematicTrans a q

    -- Bat Global Position
    batGlobal = applyForwardKinematicMatrixTrans fwdT #> xLocal

    -- Error
    e = xTargetGlobal - batGlobal
    eNorm = norm_2 e

-- Calculates newton raphson step to approximate acos
newtonRaphsonAcosStep :: Int -> Int -> Double -> (Double, Double) -> Double -> (Double, Double)
newtonRaphsonAcosStep i j q (qB, eB) t
  | i == newtonRaphsonMaxStep = (qB, eB)
  | newtonRaphsonThreshold 0 eB == 0 = (qB, eB)
  | isSingular || needReset = newtonRaphsonAcosStep (i + 1) 0 qR (qB, eB) t
  | eN < eB = newtonRaphsonAcosStep (i + 1) (j + 1) qN (qN, eN) t
  | otherwise = newtonRaphsonAcosStep (i + 1) (j + 1) qN (qB, eB) t
  where
    -- Calculate derivative and function
    f = t - cos q
    fDeriv = sin q

    -- Check if we are in a singularity
    isSingular = newtonRaphsonThreshold 0 fDeriv == 0

    -- Step
    qN = normalizeAngle $ q - f / fDeriv

    -- New Error
    eN = t - cos qN

    -- Random Restarts
    needReset = eN >= eB * newtonRaphsonBadStep && j >= newtonRaphsonMaxRandomRestartStep

    pseudoInt = round $ fromIntegral i ** (q * qB * bool (eB + 1) (1 / eB) (eB < 1 && 0 < eB))
    randNum = atIndex (randomVector pseudoInt Uniform 1) 0
    qR = pi / 2 * (2 * randNum - 1)

-- | Approximate Acos with newton raphson
newtonRaphsonAcos :: Double -> Double -> (Double, Double)
newtonRaphsonAcos q t = newtonRaphsonAcosStep 0 0 qN (qN, eN) t
  where
    qN = normalizeAngle q
    eN = t - cos qN


{-# LANGUAGE ExistentialQuantification,
             DataKinds #-}
module MVDistributions
where

import Prelude hiding ((<>))

import Control.Arrow (first)
import Control.Monad (replicateM, zipWithM)
import Control.Monad.Bayes.Class (MonadSample, normal, shuffle)
import Data.Aeson (ToJSON (toJSON))
import Data.List (permutations, sortBy, maximumBy)
import Data.Ord (comparing, Down (Down))
import qualified Data.Vector as V
import Data.Proxy (Proxy (Proxy))
import GHC.TypeLits (KnownNat, natVal)

import Numeric.LinearAlgebra.Static
import qualified Numeric.LinearAlgebra.Data as LAD
import Numeric.SpecFunctions (logGamma)

import Distributions (Distr, Distr' (..))

instance KnownNat n => ToJSON (R n) where
  toJSON x = toJSON (extract x :: LAD.Vector Double)

instance (KnownNat m, KnownNat n) => ToJSON (L m n) where
  toJSON x = toJSON (LAD.toLists (extract x :: LAD.Matrix Double))

instance KnownNat n => ToJSON (Sym n) where
  toJSON x = toJSON (unSym x)

eyeSym :: KnownNat n => Sym n
eyeSym = sym eye

sampleNormalDist :: forall n m. MonadSample m => KnownNat n => R n -> Sym n -> m (R n)
sampleNormalDist mu sig = do
  xs <- replicateM (fromInteger (natVal (Proxy :: Proxy n))) (normal 0 1)
  let x = vector xs
  return (mu + chol sig #> x)

log2pi :: Double
log2pi = log (2 * pi)

mvNormal_ll0 :: forall n. KnownNat n => Sym n -> R n -> Double
mvNormal_ll0 s x = -1/2 * (x <.> (sinv #> x) + lndets + fromInteger n * log2pi)
  where
  (sinv, (lndets, _)) = invlndet (unSym s)
  sinv' :: Sym n
  sinv' = sym sinv
  n = natVal (Proxy :: Proxy n)

mvNormal_ll :: KnownNat n => R n -> Sym n -> R n -> Double
mvNormal_ll mu s x = mvNormal_ll0 s (x - mu)

mvNormal :: KnownNat n => R n -> Sym n -> Distr (R n)
mvNormal mu s = Distr (sampleNormalDist mu s) (mvNormal_ll mu s)

uniformNSphere :: forall n. KnownNat n => Distr (R n)
uniformNSphere = Distr (do { x <- sampleNormalDist 0 (sym eye); pure (konst (1 / (norm_2 x)) * x) })
  (\_ -> - logNSphereVolume (natVal (Proxy :: Proxy n)))

-- https://en.wikipedia.org/wiki/N-sphere#Closed_forms
logNSphereVolume :: Integer -> Double
logNSphereVolume n = log (fromIntegral n) + fromIntegral n / 2 * log pi
  - logGamma (fromIntegral n / 2 + 1)

conjugate :: (KnownNat m, KnownNat n) => L m n -> Sym n -> Sym m
conjugate f p = sym (f <> unSym p <> tr f)

conjugate' :: (KnownNat m, KnownNat n) => L n m -> Sym n -> Sym m
conjugate' f p = sym (tr f <> unSym p <> f)

kalmanPredict :: KnownNat n => Sq n -> Sym n -> (R n, Sym n) -> (R n, Sym n)
kalmanPredict f q (x, p) = (f #> x, conjugate f p + q)

symEachSide :: KnownNat n => Sym n -> Sym n -> Sym n
symEachSide p x = sym (p' <> unSym x <> p')
  where p' = unSym p

-- The below is not ideal numerically
kalmanUpdate :: forall m n. (KnownNat m, KnownNat n) => R m -> L m n -> Sym m
  -> (R n, Sym n) -> ((R n, Sym n), Double)
kalmanUpdate z h r (x, p) =
  ((x + k #> ytilde, p - symEachSide p (conjugate' h sinv'))
   , loglike)
  where
  ytilde :: R m
  ytilde = z - h #> x
  trh :: L n m
  trh = tr h
  s :: Sym m
  s = r + conjugate h p
  (sinv, (lndets, _)) = invlndet (unSym s)
  sinv' :: Sym m
  sinv' = sym sinv
  k = unSym p <> trh <> sinv
  n = natVal (Proxy :: Proxy n)
  loglike = -1/2 * (ytilde <.> (sinv #> ytilde) + lndets + fromInteger n * log2pi)

shuffleList :: MonadSample m => [a] ->  m [a]
shuffleList xs = V.toList <$> shuffle (V.fromList xs)

printVector :: KnownNat n => R n -> IO ()
printVector x = print (extract x)

scaleSym :: KnownNat n => Double -> Sym n -> Sym n
scaleSym c m = sym $ konst c * unSym m


import Data.Complex
import Data.List
import Data.Maybe

direction :: Char -> Complex Float
direction 'L' = 0 :+ 1
direction 'R' = 0 :+ (-1)

blocks :: String -> Float
blocks str =
  let [(count, _)] = reads str
  in count

parse :: String -> [(Complex Float, Float)]
parse = map (\(x:xs) -> (direction x, blocks xs)) . words

advance :: (Complex Float, Complex Float) -> (Complex Float, Float) -> (Complex Float, Complex Float)
advance (dir, pos) (turn, steps) =
  let dir' = dir * turn
      pos' = pos + dir' * (steps :+ 0)
  in (dir', pos')

part1 :: String -> Int
part1 input =
  let parts = parse input
      (_, (x :+ y)) = foldl advance (0 :+ 1, 0 :+ 0) parts
  in grid x + grid y
  where grid = round . abs

firstTwice :: (Complex Float, Complex Float) -> [(Complex Float, Float)] -> [Complex Float] -> Complex Float
firstTwice (dir, pos) (x:xs) prev =
  let (turn, steps) = x
      (dir':_, poss@(pos':_)) = unzip $ map (advance (dir, pos)) $ map (\s -> (turn, s)) $ reverse [1..steps]
  in fromMaybe (firstTwice (dir', pos') xs (poss ++ prev)) (listToMaybe $ intersect poss prev)

part2 :: String -> Int
part2 input =
  let parts = parse input
      (x :+ y) = firstTwice (0 :+ 1, 0 :+ 0) parts []
  in grid x + grid y
  where grid = round . abs