{-|
Module      : Qubism.QGate
Description : Types and functions for quantum gates
Copyright   : (c) Keith Pearson, 2018
License     : MIT
Maintainer  : keith@qubitrot.org
-}

{-# LANGUAGE DataKinds           #-}
{-# LANGUAGE KindSignatures      #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeOperators       #-}

module Qubism.QGate
  ( QGate
  , (#>)
  , gate
  , ident
  , pauliX
  , pauliY
  , pauliZ
  , hadamard
  , unitary
  , cnot
  , controlled
  , ifBit
  , kronecker
  , onJust
  , onEvery
  , onRange
  ) where

-- For dependent typing
import GHC.TypeLits
import Data.Singletons
import Data.Singletons.TypeLits
import Data.Finite

import           Data.Complex
import           Data.Monoid
import           Control.Monad.Trans.State.Strict
import           Numeric.LinearAlgebra ((><))
import qualified Numeric.LinearAlgebra as LA

import Qubism.Algebra
import Qubism.StateVec
import Qubism.CReg

newtype QGate (n :: Nat) =
  UnsafeMkQGate (LA.Matrix C) -- must be a (2^n) x (2^n) unitary matrix
  deriving (Show)

-- | "Close enough" equality testing
instance Eq (QGate n) where
  (UnsafeMkQGate a) == (UnsafeMkQGate b) =
    LA.norm_2 (a - b) < 0.000001

instance KnownNat n => Semigroup (QGate n) where
  (UnsafeMkQGate a) <> (UnsafeMkQGate b) = UnsafeMkQGate $ a LA.<> b

instance KnownNat n => Monoid (QGate n) where
  mempty = ident

instance KnownNat n => VectorSpace (QGate n) where
  zero = UnsafeMkQGate . (l><l) $ repeat 0 where l = internalLen (sing :: Sing n)
  z                 .: (UnsafeMkQGate b) = UnsafeMkQGate $ LA.scalar z * b
  (UnsafeMkQGate a) +: (UnsafeMkQGate b) = UnsafeMkQGate $ a + b
  neg (UnsafeMkQGate a) = UnsafeMkQGate $ -a

instance KnownNat n => Algebra (QGate n) where
  (*:) = (<>)

internalLen :: (KnownNat n, Num a) => Sing n -> a
internalLen = (2 ^) . fromIntegral . fromSing

-- | Apply a quantum gate to a quantum register. Note that this
-- operator conflicts with the one from Numeric.LinearAlgebra.
infixr 5 #>
(#>) :: QGate n -> StateVec n -> StateVec n
(#>) (UnsafeMkQGate m) (UnsafeMkStateVec v) = UnsafeMkStateVec $ m LA.#> v

-- | Helper function to make state computations nicer
gate :: Monad m => QGate n -> StateT (StateVec n) m ()
gate g = state $ \qr -> ((), g #> qr)

ident :: forall n . KnownNat n => QGate n
ident = UnsafeMkQGate . LA.ident $ internalLen (sing :: Sing n)

-- | Also known as a NOT gate
pauliX :: QGate 1
pauliX = UnsafeMkQGate $
  (2><2) [0 :+ 0, 1 :+ 0,
          1 :+ 0, 0 :+ 0]

pauliY :: QGate 1
pauliY = UnsafeMkQGate $
  (2><2) [0 :+ 0, 0 :+ (-1),
          0 :+ 1, 0 :+ 0]

pauliZ :: QGate 1
pauliZ = UnsafeMkQGate $
  (2><2) [1 :+ 0, 0 :+ 0,
          0 :+ 0, (-1) :+ 0]

hadamard :: QGate 1
hadamard = UnsafeMkQGate $ 1 / sqrt 2 *
  (2><2) [1 :+ 0, 1 :+ 0,
          1 :+ 0, (-1) :+ 0]

-- | An arbitrary element of SU(2),
-- U(theta, phi, lambda)
unitary :: Double -> Double -> Double -> QGate 1
unitary theta phi lambda = UnsafeMkQGate $
  (2><2) [a, b, c, d]
  where a =  cis (phi+lambda/2) * ( cos (theta/2) :+ 0 )
        b = -cis (phi-lambda/2) * ( sin (theta/2) :+ 0 )
        c =  cis (phi-lambda/2) * ( sin (theta/2) :+ 0 )
        d =  cis (phi+lambda/2) * ( cos (theta/2) :+ 0 )

-- | Controlled not between control and target qubits
cnot :: KnownNat n => Finite n -> Finite n -> QGate n
cnot c t = controlled c . onJust t $ pauliX

-- | Transform an arbitrary QGate into a controled one depending on qubit i
controlled :: forall n . KnownNat n => Finite n -> QGate n -> QGate n
controlled finite (UnsafeMkQGate m) =
  UnsafeMkQGate $ (m <> projection) + LA.ident (2^n) - projection
  where
    projection = LA.diag $ (2^n) LA.|> fmap f [0..]
    f j = fromIntegral $ j `quot` 2^(n-i-1) `mod` 2
    n   = fromIntegral $ fromSing (sing :: Sing n)
    i   = fromIntegral $ getFinite finite

-- | If the given Bit is One then apply the gate, otherwise
--  do nothing
ifBit :: KnownNat n => Bit -> QGate n -> QGate n
ifBit b g = if (b == One) then g else ident

-- | The tensor product of a QGate's A and B is a QGate that acts as A on the
-- first n qubits and B on the rest. In the computational basis this is simply
-- the kronecker product of matricies.
kronecker :: QGate n -> QGate m -> QGate (m+n)
kronecker (UnsafeMkQGate a) (UnsafeMkQGate b) =
  UnsafeMkQGate $ LA.kronecker a b

-- | Promote a 1-qubit gate to an n-qubit gate with the original gate acting
-- on qubit i. Other qubits are unaffected.
onJust :: forall n . KnownNat n => Finite n -> QGate 1 -> QGate n
onJust i (UnsafeMkQGate m) = UnsafeMkQGate $  -- it should be possible to do
  pre `LA.kronecker` m `LA.kronecker` post    -- this with just QGate's
  where pre  = LA.ident $ 2^j                 -- kronecker, but I can't get
        post = LA.ident $ 2^(n-j-1)           -- the types to check.
        j    = fromIntegral $ getFinite i
        n    = fromIntegral $ fromSing (sing :: Sing n)

-- | Promote a 1-qubit gate to an n-qubit gate which acts on every qubit
-- identically.
onEvery :: forall n . KnownNat n => QGate 1 -> QGate n
onEvery (UnsafeMkQGate m) = UnsafeMkQGate $ iterate (LA.kronecker m) m !! (n-1)
  where n = fromIntegral $ fromSing (sing :: Sing n)

-- | Promote a 1-qubit gate to an n-qubit gate which acts on a range of
-- qubits, leaving others untouched
onRange :: forall n . KnownNat n => Finite n -> Finite n -> QGate 1 -> QGate n
onRange f l m = mconcat $ map (\i -> onJust i m) [f..l]


{-# LANGUAGE DeriveDataTypeable, DeriveGeneric #-}
-- |
-- Module    : Statistics.Distribution.Binomial
-- Copyright : (c) 2009 Bryan O'Sullivan
-- License   : BSD3
--
-- Maintainer  : bos@serpentine.com
-- Stability   : experimental
-- Portability : portable
--
-- The binomial distribution.  This is the discrete probability
-- distribution of the number of successes in a sequence of /n/
-- independent yes\/no experiments, each of which yields success with
-- probability /p/.

module Statistics.Distribution.Binomial
    (
      BinomialDistribution
    -- * Constructors
    , binomial
    -- * Accessors
    , bdTrials
    , bdProbability
    ) where

import Data.Aeson (FromJSON, ToJSON)
import Data.Binary (Binary)
import Data.Data (Data, Typeable)
import GHC.Generics (Generic)
import qualified Statistics.Distribution as D
import qualified Statistics.Distribution.Poisson.Internal as I
import Numeric.SpecFunctions (choose,incompleteBeta)
import Numeric.MathFunctions.Constants (m_epsilon)
import Data.Binary (put, get)
import Control.Applicative ((<$>), (<*>))

-- | The binomial distribution.
data BinomialDistribution = BD {
      bdTrials      :: {-# UNPACK #-} !Int
    -- ^ Number of trials.
    , bdProbability :: {-# UNPACK #-} !Double
    -- ^ Probability.
    } deriving (Eq, Read, Show, Typeable, Data, Generic)

instance FromJSON BinomialDistribution
instance ToJSON BinomialDistribution

instance Binary BinomialDistribution where
    put (BD x y) = put x >> put y
    get = BD <$> get <*> get

instance D.Distribution BinomialDistribution where
    cumulative = cumulative

instance D.DiscreteDistr BinomialDistribution where
    probability = probability

instance D.Mean BinomialDistribution where
    mean = mean

instance D.Variance BinomialDistribution where
    variance = variance

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

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

instance D.Entropy BinomialDistribution where
  entropy (BD n p)
    | n == 0 = 0
    | n <= 100 = directEntropy (BD n p)
    | otherwise = I.poissonEntropy (fromIntegral n * p)

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

-- This could be slow for big n
probability :: BinomialDistribution -> Int -> Double
probability (BD n p) k
  | k < 0 || k > n = 0
  | n == 0         = 1
  | otherwise      = choose n k * p^k * (1-p)^(n-k)

-- Summation from different sides required to reduce roundoff errors
cumulative :: BinomialDistribution -> Double -> Double
cumulative (BD n p) x
  | isNaN x      = error "Statistics.Distribution.Binomial.cumulative: NaN input"
  | isInfinite x = if x > 0 then 1 else 0
  | k <  0       = 0
  | k >= n       = 1
  | otherwise    = incompleteBeta (fromIntegral (n-k)) (fromIntegral (k+1)) (1 - p)
  where
    k = floor x

mean :: BinomialDistribution -> Double
mean (BD n p) = fromIntegral n * p

variance :: BinomialDistribution -> Double
variance (BD n p) = fromIntegral n * p * (1 - p)

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

-- | Construct binomial distribution. Number of trials must be
--   non-negative and probability must be in [0,1] range
binomial :: Int                 -- ^ Number of trials.
         -> Double              -- ^ Probability.
         -> BinomialDistribution
binomial n p
  | n < 0          =
    error $ msg ++ "number of trials must be non-negative. Got " ++ show n
  | p < 0 || p > 1 =
    error $ msg++"probability must be in [0,1] range. Got " ++ show p
  | otherwise      = BD n p
    where msg = "Statistics.Distribution.Binomial.binomial: "


{-# LANGUAGE ViewPatterns, ScopedTypeVariables, TupleSections #-}

import qualified Text.Printf as TP (printf)
import qualified Data.List as L (intercalate)

import Data.Array.Repa (Z(Z), (:.)((:.)), (!), (-^), (+^), (*^), All(All), Any(Any), extent, size)
import qualified Data.Array.Repa as RP (fromFunction, fromListUnboxed, Array, U, D, DIM1, DIM2, DIM4, map)
import qualified Data.Array.Repa as RP (sumAllS, slice, delay, computeS, toList)
import qualified Data.Array.Repa.Repr.Vector as RP (V, fromListVector)
import Data.Complex (Complex((:+)), realPart, imagPart)
import Data.Bits (bit)

pdim1 :: RP.Array RP.U RP.DIM1 (Complex Double) -> String
pdim1 mat = pt ++ "\n" where
    (Z :. n) = extent mat
    pt = concat [ let (xr :+ xi) = mat ! (Z :. i) in TP.printf "%14.4f +I %8.4f" xr xi | i <- [0 .. n - 1] ]

pdim2 :: RP.Array RP.U RP.DIM2 (Complex Double) -> String
pdim2 mat = L.intercalate "\n" [ pt i | i <- [0 .. m - 1]] ++ "\n\n" where
    (Z :. m :. n) = extent mat
    pt i = concat [ let (xr :+ xi) = mat ! (Z :. i :. j) in TP.printf "%14.4f +I %8.4f" xr xi | j <- [0 .. n - 1] ]

identity :: RP.Array RP.U RP.DIM2 (Complex Double)
identity = RP.fromListUnboxed (Z :. 2 :. 2) [1, 0, 0, 1]

sigmaX :: RP.Array RP.U RP.DIM2 (Complex Double)
sigmaX = RP.fromListUnboxed (Z :. 2 :. 2) [0, 1, 1, 0]

sigmaY :: RP.Array RP.U RP.DIM2 (Complex Double)
sigmaY = RP.fromListUnboxed (Z :. 2 :. 2) [0, 0 :+ (-1), 0 :+ 1, 0]

sigmaZ :: RP.Array RP.U RP.DIM2 (Complex Double)
sigmaZ = RP.fromListUnboxed (Z :. 2 :. 2) [1, 0, 0, -1]

-- | Tensor product.
kron :: RP.Array RP.U RP.DIM2 (Complex Double)
     -> RP.Array RP.U RP.DIM2 (Complex Double)
     -> RP.Array RP.U RP.DIM2 (Complex Double)
kron mata matb = reduceBlock . RP.computeS $ RP.fromFunction (extent mata) iter where
    iter i = RP.computeS $ RP.map (* (mata ! i)) matb :: RP.Array RP.U RP.DIM2 (Complex Double)

-- | Heisenberg Hamiltonian
hamil :: RP.Array RP.U RP.DIM2 (Complex Double)
hamil = RP.computeS $ RP.map (* 0.25) (((kron sigmaX identity) `mmult` (kron identity sigmaX))
    +^ ((kron sigmaY identity) `mmult` (kron identity sigmaY))
    +^ ((kron sigmaZ identity) `mmult` (kron identity sigmaZ)))

-- | S^2 Operator
s2 :: RP.Array RP.U RP.DIM2 (Complex Double)
s2 = RP.computeS $ RP.map (* 0.25) (
    square (RP.computeS $ kron sigmaX identity +^ kron identity sigmaX)
    +^ square (RP.computeS $ kron sigmaY identity +^ kron identity sigmaY)
    +^ square (RP.computeS $ kron sigmaZ identity +^ kron identity sigmaZ))

-- | Square of a matrix
square :: RP.Array RP.U RP.DIM2 (Complex Double) -> RP.Array RP.U RP.DIM2 (Complex Double)
square mat = mmult mat mat

-- | Matrix Matrix Multiplication
mmult :: RP.Array RP.U RP.DIM2 (Complex Double)
      -> RP.Array RP.U RP.DIM2 (Complex Double)
      -> RP.Array RP.U RP.DIM2 (Complex Double)
mmult mata matb = RP.computeS $ RP.fromFunction (Z :. n :. m) iter where
    (Z :. n :. k) = extent mata
    (Z :. kp :. m) = extent matb
    iter :: RP.DIM2 -> Complex Double
    iter (Z :. i :. j) = RP.sumAllS $ RP.slice mata (Any :. i :. All) *^ RP.slice matb (Any :. All :. j)

-- | Matrix Vector Multiplication
mvmult :: RP.Array RP.U RP.DIM2 (Complex Double)
       -> RP.Array RP.U RP.DIM1 (Complex Double)
       -> RP.Array RP.U RP.DIM1 (Complex Double)
mvmult mat vec = RP.computeS $ RP.fromFunction (Z :. n) iter where
    (Z :. n :. m) = extent mat
    iter :: RP.DIM1 -> Complex Double
    iter (Z :. i) = RP.sumAllS $ RP.slice mat (Any :. i :. All) *^ vec

-- | Vector Vector Dot Product
dot :: RP.Array RP.U RP.DIM1 (Complex Double) -> RP.Array RP.U RP.DIM1 (Complex Double) -> Complex Double
dot veca vecb = RP.sumAllS $ veca *^ vecb

-- | Reduce blocked matrix to regular matrix.
reduceBlock :: RP.Array RP.V RP.DIM2 (RP.Array RP.U RP.DIM2 (Complex Double)) -- ^ matrix of blocks
            -> RP.Array RP.U RP.DIM2 (Complex Double)
reduceBlock mat = RP.computeS $ RP.fromFunction (Z :. nl :. nl) iter where
    f :: Int -> RP.Array RP.U RP.DIM2 (Complex Double) -> [(Int, Int)]
    f i block = let (Z :. _ :. n) = extent block in map (i, ) [0..n - 1]
    idx = let l = concat (zipWith f [0..] (RP.toList (RP.slice mat (Any :. (0 :: Int) :. All))))
        in RP.fromListUnboxed (Z :. (length l)) l
    (Z :. nl) = extent idx
    iter (Z :. i :. j) = let (ip, im) = idx ! (Z :. i); (jp, jm) = idx ! (Z :. j)
        in (mat ! (Z :. ip :. jp)) ! (Z :. im :. jm)

-- alpha = 0, beta = 1
-- |babb> = 1011 = 11
stateIndex :: String -> Int
stateIndex x = iter (length x - 1) x where
    iter n (x:xs) = (if n == 0 then 0 else iter (n - 1) xs) + (if x == 'a' then 0 else bit n)

-- | The kth eigenstate represented as n-component vector
eigenState :: Int -> Int -> RP.Array RP.U RP.DIM1 (Complex Double)
eigenState n k = RP.computeS $ RP.fromFunction (Z :. n) (\(Z :. i) -> if k == i then 1 else 0)

state_ab = eigenState (bit 2) $ stateIndex "ab"
state_ba = eigenState (bit 2) $ stateIndex "ba"
state_aa = eigenState (bit 2) $ stateIndex "aa"
state_bb = eigenState (bit 2) $ stateIndex "bb"
singlet = RP.computeS $ RP.map (* (2 ** (-0.5))) (state_ab -^ state_ba)
triplet = RP.computeS $ RP.map (* (2 ** (-0.5))) (state_ab +^ state_ba)

main = do
    putStrLn $ "[Singlet] (Si . Sj) S^2"
    print $ let state = singlet in dot state (hamil `mvmult` state)
    print $ let state = singlet in dot state (s2 `mvmult` state)
    putStrLn $ "[Triplet] (Si . Sj) S^2"
    print $ let state = triplet in dot state (hamil `mvmult` state)
    print $ let state = triplet in dot state (s2 `mvmult` state)
    putStrLn $ "[|aa>] (Si . Sj) S^2"
    print $ let state = state_aa in dot state (hamil `mvmult` state)
    print $ let state = state_aa in dot state (s2 `mvmult` state)


{-# LANGUAGE TemplateHaskell #-}
{-# LANGUAGE TypeApplications, DataKinds, GADTs #-}
{-# LANGUAGE FlexibleInstances, MultiParamTypeClasses, FunctionalDependencies #-}

module Lib
    ( someFunc
    ) where

import Data.Maybe ( fromMaybe )
import Data.Int (Int32)
import Linear.V1
import Linear.V3
import Linear.V4 hiding (vector)
import Linear.Matrix
import Linear.Metric (dot, norm)
import Graphics.UI.GLUT ( Key(..), SpecialKey(..), MouseButton(..), KeyState(..) )
import Vis
import Vis.Vis ( vis )
import Vis.Camera ( makeCamera, setCamera, cameraMotion, cameraKeyboardMouse, Camera )
import Control.Lens
import Control.Applicative
import Control.Monad (forever)
import Numeric.AD
import Numeric.LinearAlgebra ((#>), det, fromLists, fromList, toList, pinv, inv, vector, matrix, Vector, Matrix, scale, norm_2)
--import Debug.Trace

rotM33XYZ :: (Floating f) => f -> f -> f -> M33 f
rotM33XYZ r p y = rotM33Z y !*! rotM33Y p !*! rotM33X r

rotM33X :: (Floating f) => f -> M33 f
rotM33X a = V3 (V3 1 0 0) (V3 0 ca (-sa))  (V3 0 sa ca) where
    ca = cos a
    sa = sin a

rotM33Y :: (Floating f) => f -> M33 f
rotM33Y a = V3 (V3 ca 0 sa) (V3 0 1 0)  (V3 (-sa) 0 ca) where
    ca = cos a
    sa = sin a

rotM33Z :: (Floating f) => f -> M33 f
rotM33Z a = V3 (V3 ca (-sa) 0)  (V3 sa ca 0) (V3 0 0 1) where
    ca = cos a
    sa = sin a

printV4 :: (Show f) => V4 f -> IO ()
printV4 (V4 v1 v2 v3 v4) = print [v1, v2, v3, v4]

printV3 :: (Show f) => V3 f -> IO ()
printV3 (V3 v1 v2 v3) = print [v1, v2, v3]

printM33 :: (Show f) => M33 f -> IO ()
printM33 (V3 v1 v2 v3) = do
    putStrLn ""
    printV3 v1
    printV3 v2
    printV3 v3
    putStrLn ""

printM44 :: (Show f) => M44 f -> IO ()
printM44 (V4 v1 v2 v3 v4) = do
    putStrLn ""
    printV4 v1
    printV4 v2
    printV4 v3
    printV4 v4
    putStrLn ""

rotM44X :: (Floating f) => f -> M44 f
rotM44X a = identity & _m33 .~ rotM33X a

rotM44Z :: (Floating f) => f -> M44 f
rotM44Z a = identity & _m33 .~ rotM33Z a

transM44X :: (Floating f) => f -> M44 f
transM44X a =  identity & translation .~ V3 a 0 0

transM44Z :: (Floating f) => f -> M44 f
transM44Z a =  identity & translation .~ V3 0 0 a

twistX :: (Floating f) => f -> f -> M44 f
twistX d theta = transM44X d !*! rotM44X theta

twistZ :: (Floating f) => f -> f -> M44 f
twistZ d theta = transM44Z d !*! rotM44Z theta

skew :: (Floating f) => V3 f -> M33 f
skew (V3 x y z) = V3 (V3 0 (-z) y) (V3 z 0 (-x)) (V3 (-y) x 0)

invSkew :: (Floating f) => M33 f -> V3 f
invSkew (V3 (V3 _ _ y) (V3 z _ _) (V3 _ x _)) = V3 x y z

rodrigues :: (Eq f, Floating f) => V3 f -> M33 f
rodrigues r = identity !!* ct  !+! rrt !!* (1 - ct) !+! sr !!* st
    where
    theta = norm r
    ct    = cos theta
    st    = sin theta
    r'    = if theta == 0 then V3 1 0 0 else r / V3 theta theta theta
    sr    = skew r'
    mrt   = V1 r'
    mr    = transpose mrt
    rrt   = mr !*! mrt

invRodrigues :: (Floating f) => M33 f -> V3 f
invRodrigues m = V3 x' y' z' * V3 t t t
    where
    V3 x y z = invSkew (m - transpose m) / 2
    st = norm (V3 x y z)
    t  = asin st
    V3 x' y' z' = V3 x y z / V3 t t t

frame2xyzaxayaz :: (Floating f) => M44 f -> [f]
frame2xyzaxayaz m = [x, y, z, ax, ay, az]
    where
    V3 ax ay az = invRodrigues $ m ^. _m33
    V3 x y z = m ^. translation

class Renderable b f  | b -> f where
    renderFC :: Flavour -> Color -> b -> VisObject f
    render   :: b -> VisObject f
    render = renderFC Wireframe blue

-- M44 を render
instance (Floating f) => Renderable (M44 f) f where
    renderFC _ _ m = convfrm m $ axes (0.05, 10)
        where
        axes (size, aspectRatio) = VisObjects [ Arrow (size, aspectRatio) (V3 1 0 0) red
                                              , Arrow (size, aspectRatio) (V3 0 0 1) blue
                                              ]

class FrameConvertable b f | b -> f where
    convfrm :: M44 f -> b -> b

instance (Floating f) => FrameConvertable (VisObject f) f where
    convfrm m v = Trans t $ RotDcm rot v
        where
        rot = transpose $ m ^. _m33 -- NOTE: vis function use row vector??
        t   = m ^. translation

---- for rendering and collision detection
data LinkShape f = OBB (V3 f) (V3 f) (M33 f) -- center size/2 axis-vector
    deriving Show

instance (Floating f) => FrameConvertable (LinkShape f) f where
    convfrm m (OBB c a r) = OBB c' a r'
        where
        c' = (m !* point c) ^. _xyz
        r' = m^._m33 !*! r

instance (Floating f) => Renderable (LinkShape f) f where
    renderFC f color (OBB c b@(V3 bx by bz) m) = Trans c $ RotDcm rot $ Box (2*bx, 2*by, 2*bz) f color
        where
        rot = transpose $ m ^. _m33 -- NOTE: vis function use row vector??

hasIntersection :: (Ord f, Floating f) => LinkShape f -> LinkShape f -> Bool
hasIntersection (OBB c0 a@(V3 a0 a1 a2)  m0) (OBB c1 b@(V3 b0 b1 b2) m1) = not $ easy || i0 || i1 || i2 || i3 || i4 || i5 || i6 || i7 || i8 || i9 || i10 || i11 || i12 || i13 || i14
    where

    -- not intersection test using bounding sphere
    easy = norm a + norm b < norm vD

    -- not intersection test
    i0  = a0 + (b0*c00' + b1*c01' + b2*c02') < abs sA0D
    i1  = a1 + (b0*c10' + b1*c11' + b2*c12') < abs sA1D
    i2  = a2 + (b0*c20' + b1*c21' + b2*c22') < abs sA2D

    i3  = (a0*c00' + a1*c10' + a2*c20') + b0 < abs (vB0 `dot` vD)
    i4  = (a0*c01' + a1*c11' + a2*c21') + b1 < abs (vB1 `dot` vD)
    i5  = (a0*c02' + a1*c12' + a2*c22') + b2 < abs (vB2 `dot` vD)

    i6  = (a1*c20' + a2*c10') + (b1*c02' + b2*c01') < abs (c10*sA2D - c20*sA1D)
    i7  = (a1*c21' + a2*c11') + (b0*c02' + b2*c00') < abs (c11*sA2D - c21*sA1D)
    i8  = (a1*c22' + a2*c12') + (b0*c01' + b1*c00') < abs (c12*sA2D - c22*sA1D)

    i9  = (a0*c20' + a2*c00') + (b1*c12' + b2*c11') < abs (c20*sA0D - c00*sA2D)
    i10 = (a0*c21' + a2*c01') + (b0*c12' + b2*c10') < abs (c21*sA0D - c01*sA2D)
    i11 = (a0*c22' + a2*c02') + (b0*c11' + b1*c10') < abs (c22*sA0D - c02*sA2D)

    i12 = (a0*c10' + a1*c00') + (b1*c22' + b2*c21') < abs (c00*sA1D - c10*sA0D)
    i13 = (a0*c11' + a1*c01') + (b0*c22' + b2*c20') < abs (c01*sA1D - c11*sA0D)
    i14 = (a0*c12' + a1*c02') + (b0*c21' + b1*c20') < abs (c02*sA1D - c12*sA0D)

    vD = c1 - c0
    V3 vA0 vA1  vA2 = transpose m0
    V3 vB0 vB1  vB2 = transpose m1
    V3 (V3 c00 c01 c02) (V3 c10 c11 c12) (V3 c20 c21 c22)  = transpose m0 !*! m1

    ((c00', c01', c02'), (c10', c11', c12'), (c20', c21', c22'))  = ((abs c00, abs c01, abs c02), (abs c10, abs c11, abs c12), (abs c20, abs c21, abs c22))

    sA0D = vA0 `dot` vD
    sA1D = vA1 `dot` vD
    sA2D = vA2 `dot` vD

-- TODO
hasCollisions :: (Ord f, Floating f) => [LinkShape f] -> [Bool]
hasCollisions (x:xs) = any id (hasIntersection x <$> xs) : hasCollisions xs
hasCollisions [] = []

aabb :: (Ord f, Floating f) => (f, f) -> (f, f) -> (f, f) -> LinkShape f
aabb (x0,x1) (y0,y1) (z0,z1) = OBB c b identity
        where
        c = V3 ((x1+x0)/2) ((y1+y0)/2) ((z1+z0)/2)
        b = V3 ((x1-x0)/2) ((y1-y0)/2) ((z1-z0)/2)

type FrameTrans f = f -> M44 f

data DH f = DH {   __a :: f      -- [m]
               ,   __d :: f      -- [m]
               ,   __alpha :: f  -- [rad]
               } deriving Show

makeLenses ''DH

d1 :: (Floating f) => f
d1 = 0.089159

a2 :: (Floating f) => f
a2 = 0.425

a3 :: (Floating f) => f
a3 = 0.39225

d4 :: (Floating f) => f
d4 = 0.10915

d5 :: (Floating f) => f
d5 = 0.09465

d6 :: (Floating f) => f
d6 = 0.0823

ur5dh :: (Floating f) => [DH f]
ur5dh =
    [ DH 0             d1          (pi/2) -- J1
    , DH (-a2)         0           0      -- J2
    , DH (-a3)         0           0      -- J3
    , DH 0             d4          (pi/2) -- J4
    , DH 0             d5          (-pi/2)-- J5
    , DH 0             d6          0      -- J6(flange)
    ]

tcpOffsetZ :: (Floating f) => f
tcpOffsetZ = 0.15

-- how to draw link (from frame{i} to frame{i+1})
ur5link :: (Ord f, Floating f) => [LinkShape f]
ur5link =
    [ aabb (-r/2, r/2) (-d1, r/2) (-r/2, r/2) -- 1,2
    , aabb (-r/2, a2 + r/2) (-r/2, r/2) (delta+r/2, 3*r/2) -- 2,3
    , aabb (-r/2, a3 + r/2) (-r/2, r/2) (-r/2, r/2) -- 3,4
    , aabb (-r/2, r/2) (r/2 + delta-d4, r/2) (-r/2, r/2) -- 4,5
    , aabb (-r/2, r/2) (-r/2, d5 - r/2 - delta) (-r/2, r/2) -- 5,6
    , aabb (-r/2, r/2) (-r/2, r/2) (delta-d6+r, 0) -- 6,frange
    , aabb (-r/2, r/2) (-r/2, r/2) (delta, tcpOffsetZ)  -- tool
    ]
    where
    r = 0.05
    delta = 0.01

-- https://en.wikipedia.org/wiki/Denavit%E2%80%93Hartenberg_parameters
dr2ft :: (Floating f) => DH f -> FrameTrans f
dr2ft dh theta = twistZ (dh ^. _d) theta !*!  twistX (dh^. _a) (dh^. _alpha)

ur5fts :: (Floating f) => [FrameTrans f]
ur5fts = dr2ft <$> ur5dh

-- calc each frame in base-coordinate for each joint-angle
kinematics :: (Floating f) => [FrameTrans f] -> [f] -> [M44 f]
kinematics (e:es) (a:as) = t : ((t !*! ) <$> kinematics es as)
    where
    t = e a
kinematics [] [] = []
kinematics _ _ = error "invalid args"

joint2tcp :: (Floating f) => [FrameTrans f] -> [f] -> [f]
joint2tcp es js = frame2xyzaxayaz $ last ( kinematics es js) !*! transM44Z tcpOffsetZ

poslabel = ["x", "y", "z", "ax", "ay", "az"]

data World = World { _jPos    :: [Double]
                   , _tcpVel  :: [Double]
                   , _trajct   :: [V3 Double]
                   , _camera :: Camera
                   }

makeLenses ''World

world0 = World [0.1, 0.4, 0.4, 0.1, 0.1, 0.1] [0, 0, 0, 0, 0, 0] [] cameraState0
    where
    defaultCamera = Camera0 { phi0 = 60 , theta0 = 20 , rho0 = 7 }
    cameraState0 = makeCamera $ fromMaybe defaultCamera (optInitialCamera defaultOpts)

motionPolicy :: [Double] -> [Double] ->  [Double] -> [Double]
motionPolicy goal pos vel = toList acc
    where
    alpha = 0.001
    c     = 0.01
    beta  = 0.5
    acc :: Vector Double
    acc = alpha * s (vector goal - vector pos) - beta * vector vel
    s :: Vector Double -> Vector Double
    s v = scale (1 / h (norm_2 v)) v
    h :: Double -> Double
    h z = z + c * log (1 + exp (-2) * c * z)

mySimulateIO :: IO ()
mySimulateIO = vis defaultOpts 0.016 world0 simFun drawFun setCameraFun (Just kmCallback) (Just motionCallback) Nothing
  where
    delta = 0.001

    simFun (w,t) = return $ w  &~ do
                        let tcpp = joint2tcp ur5fts $ w ^. jPos
                        let tcpv = w ^. tcpVel
                        let goal = [0.192428507378705,-0.39791168805212945,-0.43347675593345414,0.7684417340393249,-0.40013839074582813,0.3969016751654795]
                        let tcpa = motionPolicy goal tcpp tcpv
                        let tcpv' = zipWith (+) tcpv  tcpa
                        let j    = fromLists (jacobian (joint2tcp $ ur5fts) (w ^. jPos) :: [[Double]])
                        let jvel = if det j > 0.0001 then toList $ pinv j #> fromList tcpv else [0, 0, 0, 0, 0, 0]
                        let jpos' = zipWith (+) (w ^. jPos) jvel
                        let pos  = joint2tcp ur5fts jpos'
                        jPos .= jpos'
                        tcpVel .= tcpv'
                        trajct %= (:) (V3 (pos !! 0) (pos !! 1) (pos !! 2))

    drawFun (w, t) = do
        let jpos = w ^. jPos
        let frames =  identity : kinematics ur5fts jpos
        let vframes = render <$> frames
        let ts     = tail frames ++ [last frames]
        let links  = zipWith convfrm ts ur5link
        let colors  = (\b -> if b then red else blue) <$> hasCollisions links
        let vlinks  = zipWith (renderFC Wireframe) colors links
        let pos = joint2tcp ur5fts jpos
        let info = [Text2d (show (l, j)) (0, 110  - i * 20) Fixed8By13 red | (l, i,j) <- zip3 poslabel [0..] pos]
        let infoj = [Text2d (show ("J", i, j)) (300, 110  - i * 20) Fixed8By13 red | (i,j) <- zip [0..] jpos]
        let points = Points (w ^. trajct) (Just 1.0) yellow
        print pos
        return (VisObjects $ vlinks ++ vframes ++ info ++ infoj ++ [points], Nothing)

    setCameraFun w = setCamera $ w ^. camera

    kmCallback w k0 k1 _ _ = case (k0,  k1) of
                         (Char 'h', Down)            -> w & tcpVel .~ [ delta, 0, 0, 0, 0, 0]
                         (Char 'l', Down)            -> w & tcpVel .~ [-delta, 0, 0, 0, 0, 0]
                         (Char 'j', Down)            -> w & tcpVel .~ [0,  delta, 0, 0, 0, 0]
                         (Char 'k', Down)            -> w & tcpVel .~ [0, -delta, 0, 0, 0, 0]
                         (Char 'n', Down)            -> w & tcpVel .~ [0, 0,  delta, 0, 0, 0]
                         (Char 'p', Down)            -> w & tcpVel .~ [0, 0, -delta, 0, 0, 0]
                         (SpecialKey KeyLeft,  Down) -> w & tcpVel .~ [0, 0, 0,  delta, 0, 0]
                         (SpecialKey KeyRight, Down) -> w & tcpVel .~ [0, 0, 0, -delta, 0, 0]
                         (SpecialKey KeyDown,  Down) -> w & tcpVel .~ [0, 0, 0, 0,  delta, 0]
                         (SpecialKey KeyUp,    Down) -> w & tcpVel .~ [0, 0, 0, 0, -delta, 0]
                         (Char ',', Down)            -> w & tcpVel .~ [0, 0, 0, 0, 0,  delta]
                         (Char '.', Down)            -> w & tcpVel .~ [0, 0, 0, 0, 0, -delta]
                         (Char 'a', Down)            -> world0 & camera .~ (w ^. camera)
                         (Char 's', Down)            -> w & tcpVel .~ [0, 0, 0, 0, 0, 0]
                         (_, _)           -> w & camera %~ (\x -> cameraKeyboardMouse x k0 k1)

    motionCallback w pos = w & camera %~ (`cameraMotion` pos)

someFunc :: IO ()
someFunc = mySimulateIO

main = mySimulateIO


{-# LANGUAGE RankNTypes, BangPatterns, ScopedTypeVariables, FlexibleInstances #-}

module Math.Probably.HamMC where

import Math.Probably.MCMC
import qualified Math.Probably.PDF as PDF
import Math.Probably.RandIO
import Math.Probably.FoldingStats
import Math.Probably.Sampler
import Control.Applicative

import Numeric.LinearAlgebra
import Numeric.AD
import Text.Printf
import System.IO
import Data.Maybe

import Statistics.Test.KolmogorovSmirnov
import Statistics.Test.MannWhitneyU
import qualified Data.Vector.Unboxed as U
import qualified Control.Monad.State.Strict as S

import Debug.Trace
import Data.IORef
import Control.Spoon

data HMCPar = HMCPar { hpXi :: !(Vector Double),
                       hpL :: !Int,
                       hpPi :: !Double,
                       hpEpsilon :: !Double,
                       hpCount :: !Int,
                       hpAccept :: !Int,
                       hpFreezeEps :: Bool }
    deriving Show

hmc1 :: (Vector Double -> (Double,Vector Double))
         -> HMCPar
         -> Sampler HMCPar
hmc1 postGrad (HMCPar current_q l _ epsMean count accept freeze) = do
   let dims = dim current_q
       grad_u = negate . snd . postGrad
       u = negate . fst . postGrad
   current_p <- fmap fromList $ gaussManyUnitD dims
   eps <- uniform (epsMean*0.8) (epsMean*1.2)
   --the initial half momentum step

   let p_half_step = current_p - scale (eps/2) (grad_u current_q)

   let step :: Int -> Vector Double -> Vector Double -> (Vector Double, Vector Double)
       step n p q  -- note that my l is Radford Neal's L+1
        | n == 0 = let qfinal = q + scale eps p -- only q not p update in last loop
                       pfinal = p - scale (eps/2) (grad_u qfinal) -- half momentum step at end
                   in (negate pfinal, -- negate momentum for symmetric proposal
                       qfinal)
        | otherwise =
            let q1 = q + scale eps p
                p1 = p - scale eps (grad_u q1)
            in step (n-1) p1 q1

   let (propose_p, propose_q) = step l p_half_step current_q
       current_U =  u current_q
       current_K =  (current_p `dot` current_p) / 2
       propose_U = u propose_q
       propose_K =  (propose_p `dot` propose_p) / 2
       ratio = exp $ current_U - propose_U + current_K - propose_K
       tr = max 1.0 $ realToFrac count

   u <- unitSample -- 0 to 1
   return $ if  u < ratio
               then HMCPar propose_q l propose_U (if freeze then epsMean else (min 2 $ 1+k_hmc/tr)*epsMean) (count+1) (accept+1) freeze
               else HMCPar current_q l current_U (if freeze then epsMean else (max 0.5 $ 1-k_hmc/tr)**1.8*epsMean) (count+1) (accept) freeze

traceAs s x = trace (s++" = "++show x) x

runHMC postgrad nsam hmc0 = go nsam hmc0 [] where
  go 0 hmcp xs = do io $ putStrLn $ "done with "++show hmcp
                    return (hmcp, xs)
  go n hmcp xs = do hmcnext <- sample $ hmc1 postgrad hmcp
                    if isNaN $ hpPi hmcnext
                       then return (hmcnext, [])
                       else do io $ do putStrLn $ show $ (hpCount hmcnext,
                                                          hpPi hmcnext,
                                                          hpEpsilon hmcnext)
                                       hFlush stdout
                               go (n-1) hmcnext $ hpXi hmcnext : xs

k_hmc = 2


{-|
Module      : ODMatrix.ElementaryDecomposition
Description : This module is responsible for the representation of Elementary operations and the process of elementary decomposition of the path between Origin-Destination matrices.

-}
module ODMatrix.ElementaryDecomposition (
    ElementaryMatrix(..)
  , getChildren
  , applyElementary
  , applyPath
  , applicables
  , elementaryValue
  , opposite
  , isNullElementary
  , ODM
  , Bounds
  ) where

  import Numeric.LinearAlgebra
  --import Data.Tree (Tree(..), unfoldTree)

  import ODMatrix (ODM)

  data ElementaryMatrix =
    NullElementary Int |
    Elementary Int (Int, Int) (Int, Int)
    deriving (Show)

  -- | Size of the elementary matrix
  emSize :: ElementaryMatrix -> Int
  emSize (NullElementary s) = s
  emSize (Elementary s _ _) = s

  -- | Bounds of the capacity of the unit (Lower bound, Upper bound).
  type Bounds = (Double, Double)

  -- Size, visited, s
  type ES = (ElementaryMatrix, ODM)

  applyElementaries :: ODM -> [ElementaryMatrix] -> [ODM]
  applyElementaries s es = map (applyElementary s) es

  -- | Apply a elementary operation to the given ODM.
  applyElementary :: ODM -> ElementaryMatrix -> ODM
  applyElementary a e = a + (toMatrix e)

  -- | Apply a sequence of elementary operations to the given ODM.
  applyPath :: ODM -> [ElementaryMatrix] -> ODM
  applyPath = foldl applyElementary

  -- | Given S, such that A + S = B, this function computes the first elementary permutation matrix E such that S' + E = S.
  getChildren :: ODM                  -- ^ Full path permutation matrix
               -> [ElementaryMatrix]  -- ^ All possible elementary matrices
  getChildren odm = es ++ map opposite es
    where al = toAssocList odm
          s = rows odm
          es = [ Elementary s (r1,c1) (r2,c2) |
                  ((r1,c1),_) <- al,
                  ((r2,c2),_) <- al,
                  r1 > r2, c1 < c2 ]

  opposite :: ElementaryMatrix -> ElementaryMatrix
  opposite (Elementary s (r1,c1) (r2,c2)) = Elementary s (r2,c1) (r1,c2)
  opposite (NullElementary s) = NullElementary s

  -- getChildren s = foldl (\acc p -> acc ++ opposites s p pivots) [] pivots
  --   where -- Filtra solo los potenciales pivotes a los que puede agregarse o quitarse un pasajero.
  --         pivots = [ p | p@((i,j),x) <- toAssocList s, i <= j, x /= 0]

  -- reciprocals :: Pivot -> [Pivot] -> [Pivot]
  -- reciprocals ((i,j),x) pivots =
  --   [ p | p@((h,k),y) <- pivots, h < i, j < k, x/y > 0]

  -- opposites :: ODM -> Pivot -> [Pivot] -> [ElementaryMatrix]
  -- opposites s p@(i,j) pivots =
  --   [ Elementary (rows s) (i,j) (h,k)
  --     | ((h,k),_) <- reciprocals p pivots,
  --       atIndex s (h,j) / x < 0,  -- Opposite signs
  --       atIndex s (i,k) / x < 0 ]

  -- | Filter the elementary matrices that cannot be applied to the odm
  applicables :: Bounds             -- ^ Bounds of the values of the target
              -> ODM      -- ^ Target Origin-Destination Matrix
              -> [ElementaryMatrix] -- ^ List of potential elementary operations
              -> [ElementaryMatrix] -- ^ List of applicable elementary operations
  applicables b s = filter (applicable b s)

  -- | Indicates if an elementary operation is applicable for a given origin destination matrix with the given bounds.
  applicable :: Bounds            -- ^ Bounds of the values of the target
             -> ODM               -- ^ Target Origin-Destination Matrix
             -> ElementaryMatrix  -- ^ Elementary operation to apply
             -> Bool
  applicable _ _ (NullElementary _) = True
  applicable (mn,mx) s (Elementary _ (i,j) (h,k)) =
    p (i,k) > mn && p (h,k) < mx &&
    p (i,j) < mx && p (h,j) > mn
    where p = atIndex s

  -- | Compute the value of a elementary matrix in base of a distance functions between cells.
  elementaryValue :: (Int -> Int -> Double) -- ^ Measure function
                  -> ElementaryMatrix       -- ^ E
                  -> Double                 -- ^ Value of E
  elementaryValue _ (NullElementary _) = 0
  elementaryValue m (Elementary _ (r1,_ ) (r2,_)) = s * 2 * m r1 r2
    where s = fromIntegral . signum $ (r1 - r2)

  -- * Convertions

  -- | Transform the internal representation of a Elementary Matrix in a Matrix Double
  toMatrix :: ElementaryMatrix -> ODM
  toMatrix (NullElementary n) = assoc (n,n) 0 []
  toMatrix (Elementary n idx1@(r1,c1) idx2@(r2,c2)) =
    assoc (n,n) 0 [(idx1,1), (idx2,1), ((r1,c2),-1), ((r2,c1),-1)]

  -- | Transform a matrix to an assoc list.
  -- Recovers only the valid, different to 0, elements of the ODM.
  toAssocList :: ODM -> [((Int, Int), Double)]
  toAssocList s = [ c | c@((i,j),v) <- zip idxs elems, i <= j, v /= 0]
    where n = rows s
          idxs = [ (i,j) | i <- [0..(n-1)], j <- [0..(n-1)] ]
          elems = toList . flatten $ s

  -- | Return True if the elementary matrix is a null elementary matrix.
  isNullElementary :: ElementaryMatrix -> Bool
  isNullElementary (NullElementary s) = True
  isNullElementary _ = False


{-# LANGUAGE DeriveFunctor #-}
{-# LANGUAGE TypeFamilies  #-}

-- | Provides spherical harmonic models of scalar-valued functions.
module Math.SphericalHarmonics
(
  SphericalHarmonicModel
, sphericalHarmonicModel
, scaledSphericalHarmonicModel
, evaluateModel
, evaluateModelCartesian
, evaluateModelGradient
, evaluateModelGradientCartesian
, evaluateModelGradientInLocalTangentPlane
)
where

import Data.Complex
import Data.VectorSpace hiding (magnitude)
import Math.SphericalHarmonics.AssociatedLegendre
import Numeric.AD

-- | Represents a spherical harmonic model of a scalar-valued function.
data SphericalHarmonicModel a = SphericalHarmonicModel [[(a, a)]]
  deriving (Functor)

-- | Creates a spherical harmonic model.
-- Result in an error if the length of the list is not a triangular number.
sphericalHarmonicModel :: (Fractional a) => [[(a, a)]] -- ^ A list of g and h coefficients for the model
                       -> SphericalHarmonicModel a -- ^ The spherical harmonic model
sphericalHarmonicModel cs | valid = SphericalHarmonicModel cs
                          | otherwise = error "The number of coefficients is not a triangular number."
  where
    valid = and $ zipWith (==) (fmap length cs) [1..length cs]

-- | Creates a spherical harmonic model, scaling coefficients for the supplied reference radius.
-- Result in an error if the length of the list is not a triangular number.
scaledSphericalHarmonicModel :: (Fractional a) => a -- ^ The reference radius
                       -> [[(a, a)]] -- ^ A list of g and h coefficients for the model
                       -> SphericalHarmonicModel a -- ^ The spherical harmonic model
scaledSphericalHarmonicModel r cs = sphericalHarmonicModel cs'
  where
    cs' = normalizeReferenceRadius r cs

instance(Fractional a, Eq a) => AdditiveGroup (SphericalHarmonicModel a) where
  zeroV = SphericalHarmonicModel [[(0,0)]]
  negateV = fmap negate
  (SphericalHarmonicModel m1) ^+^ (SphericalHarmonicModel m2) = SphericalHarmonicModel (combineCoefficients m1 m2)
    where
      combineCoefficients [] cs = cs
      combineCoefficients cs [] = cs
      combineCoefficients (c1:cs1) (c2:cs2) = zipWith addPairs c1 c2 : combineCoefficients cs1 cs2
      addPairs (g1, h1) (g2, h2) = (g1 + g2, h1 + h2)

instance (Fractional a, Eq a) => VectorSpace (SphericalHarmonicModel a) where
  type Scalar (SphericalHarmonicModel a) = a
  x *^ m = fmap (* x) m

normalizeReferenceRadius :: (Fractional a) => a -> [[(a, a)]] -> [[(a, a)]]
normalizeReferenceRadius r = zipWith (fmap . mapWholePair . transform) [0 :: Int ..]
  where
    transform n = (* (r ^ (2 + n)))

-- | Computes the scalar value of the spherical harmonic model at a specified spherical position.
evaluateModel :: (RealFloat a, Ord a) => SphericalHarmonicModel a -- ^ Spherical harmonic model
              -> a -- ^ Spherical radius
              -> a -- ^ Spherical colatitude (radian)
              -> a -- ^ Spherical longitude (radian)
              -> a -- ^ Model value
evaluateModel m r colat lon = evaluateModel' m r (cos colat) (cis lon)

-- | Computes the scalar value of the spherical harmonic model at a specified Cartesian position.
evaluateModelCartesian :: (RealFloat a, Ord a) => SphericalHarmonicModel a -- ^ Spherical harmonic model
                       -> a -- ^ X position
                       -> a -- ^ Y position
                       -> a -- ^ Z position
                       -> a -- ^ Model value
evaluateModelCartesian m x y z = evaluateModel' m r cosColat cisLon
  where
    r = sqrt $ (x*x) + (y*y) + (z*z)
    cosColat = z / r
    cisLon = normalize $ mkPolar x y

evaluateModel' :: (RealFloat a, Ord a) => SphericalHarmonicModel a
               -> a -- r
               -> a -- cosColat
               -> Complex a -- cisLon
               -> a
evaluateModel' (SphericalHarmonicModel cs) r cosColat cisLon = sum $ zipWith (*) (iterate (/ r) (recip r)) (zipWith evaluateDegree [0..] cs)
  where
    sines = 1 : iterate (* cisLon) cisLon
    evaluateDegree n cs' = sum $ zipWith3 evaluateOrder (fmap (schmidtSemiNormalizedAssociatedLegendreFunction n) [0..n]) cs' sines
    evaluateOrder p (g, h) cisMLon = ((g * realPart cisMLon) + (h * imagPart cisMLon)) * (p (cosColat))

-- | Computes the gradient of the scalar value of the spherical harmonic model, in spherical coordinates, at a specified location.
evaluateModelGradient :: (RealFloat a, Ord a) => SphericalHarmonicModel a -- ^ Spherical harmonic model
                      -> a -- ^ Spherical radius
                      -> a -- ^ Spherical colatitude (radian)
                      -> a -- ^ Spherical longitude (radian)
                      -> (a, a, a) -- ^ Radial, colatitudinal, and longitudinal components of gradient
evaluateModelGradient model r colat lon = makeTuple . fmap negate $ modelGrad [r, colat, lon]
  where
    modelGrad = grad (\[r', c', l'] -> evaluateModel (fmap auto model) r' c' l')

-- | Computes the gradient of the scalar value of the spherical harmonic model at a specified location, in Cartesian coordinates.
-- The result is expressed in right-handed coordinates centered at the origin of the sphere, with the positive Z-axis piercing the
-- north pole and the positive x-axis piercing the reference meridian.
evaluateModelGradientCartesian :: (RealFloat a, Ord a) => SphericalHarmonicModel a -- ^ Spherical harmonic model
                               -> a -- ^ X position
                               -> a -- ^ Y position
                               -> a -- ^ Z position
                               -> (a, a, a) -- X, Y, and Z components of gradient
evaluateModelGradientCartesian model x y z = makeTuple . fmap negate $ modelGrad [x, y, z]
  where
    modelGrad = grad (\[x', y', z'] -> evaluateModelCartesian (fmap auto model) x' y' z')

-- | Computes the gradient of the scalar value of the spherical harmonic model at a specified location, in Cartesian coordinates.
-- The result is expressed in a reference frame locally tangent to the sphere at the specified location.
evaluateModelGradientInLocalTangentPlane :: (RealFloat a, Ord a) => SphericalHarmonicModel a -- ^ Spherical harmonic model
                                         -> a -- ^ Spherical radius
                                         -> a -- ^ Spherical colatitude (radian)
                                         -> a -- ^ Spherical longitude (radian)
                                         -> (a, a, a) -- ^ East, North, and up components of gradient
evaluateModelGradientInLocalTangentPlane model r colat lon = (e, n, u)
  where
    (r', colat', lon') = evaluateModelGradient model r colat lon
    e = lon' / (r * sin colat)
    n = -colat' / r -- negated because the colatitude increase southward
    u = r'

normalize :: (RealFloat a) => Complex a -> Complex a
normalize r@(x :+ y) | isInfinite m' = 0
                     | otherwise = (x * m') :+ (y * m')
  where
    m' = recip . magnitude $ r

mapWholePair :: (a -> b) -> (a, a) -> (b, b)
mapWholePair f (a, b) = (f a, f b)

makeTuple :: [a] -> (a, a, a)
makeTuple [x, y, z] = (x, y, z)


{-# LANGUAGE BangPatterns     #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE Strict           #-}
{-# LANGUAGE StrictData       #-}
module FourierPinwheel.Harmonics where

import           Data.Array.IArray   as IA
import           Data.Array.Repa     as R
import           Data.Complex
import           Data.List
import           Data.List           as L
import           Data.Vector.Generic as VG
import           Data.Vector.Unboxed as VU
import           Filter.Utils
import           Math.Gamma
import           Utils.Parallel

data FPData vector = FPData
  { getFPDataHarmonics       :: [vector]
  , getFPDataHarmonicsOffset :: [vector]
  , getFPDataCoef            :: [vector]
  , getFPDataCoefHollow      :: [vector]
  }

createHarmonics ::
     ( RealFloat e
     , Unbox e
     , NFData e
     , VG.Vector vector (Complex e)
     , NFData (vector (Complex e))
     , Gamma (Complex e)
     )
  => Int
  -> Int
  -> Int
  -> Int
  -> Int
  -> e
  -> e
  -> e
  -> R.Array U DIM4 (Complex e)
  -> IO (FPData (vector (Complex e)))
createHarmonics numR2Freqs phiFreq rhoFreq thetaFreq rFreq alpha periodR2 periodEnv coefficients = do
  let harmonicVecs =
        parMap
          rdeepseq
          (\radialFreq ->
             createVector
               numR2Freqs
               (2 * phiFreq + 1)
               1
               (pinwheel periodR2 radialConst alpha radialFreq))
          [-rhoFreq .. rhoFreq]
      harmonicVecsOffset =
        parMap
          rdeepseq
          (\radialFreq ->
             createVector
               numR2Freqs
               (2 * thetaFreq + 1)
               (-1)
               (logpolarHarmonics periodR2 radialConst radialFreq)) .
        L.reverse $
        [-rFreq .. rFreq]
  phaseShiftedCoefArr <-
    computeUnboxedP .
    R.zipWith (*) coefficients . fromFunction (extent coefficients) $ \(Z :. r :. theta :. rho :. phi) ->
      phaseShift
        (phi - phiFreq - (theta - thetaFreq))
        (rho - rhoFreq - (r - rFreq))
        radialConst
        alpha
  let phaseShiftedCoefVecs =
        parMap
          rdeepseq
          (\r ->
             VU.convert . toUnboxed . computeS . R.slice phaseShiftedCoefArr $
             (Z :. r :. All :. All :. All))
          [0 .. 2 * rFreq]
  coefHollowArr <- computeUnboxedP $ sumArr *^ coefficients
  let coefHollowVecs =
        parMap
          rdeepseq
          (\r ->
             VU.convert . toUnboxed . computeS . R.slice coefHollowArr $
             (Z :. r :. All :. All :. All))
          [0 .. 2 * rFreq]
  return $
    FPData harmonicVecs harmonicVecsOffset phaseShiftedCoefVecs coefHollowVecs
  where
    radialConst = 2 * pi / log periodEnv
    !sumArr1 =
      computePinwheelArray
        numR2Freqs
        phiFreq
        rhoFreq
        thetaFreq
        rFreq
        alpha
        periodR2
        periodEnv
    sumArr =
      fromListUnboxed (extent coefficients) .
      parMap
        rdeepseq
        (\(r, theta, rho, phi) ->
           let angularFreq = phi - theta
               radialFreq = rho - r
            in sumArr1 IA.! (radialFreq, angularFreq)) $
      [ (r, theta, rho, phi)
      | r <- [-rFreq .. rFreq]
      , theta <- [-thetaFreq .. thetaFreq]
      , rho <- [-rhoFreq .. rhoFreq]
      , phi <- [-phiFreq .. phiFreq]
      ]

{-# INLINE createVector #-}
createVector ::
     (RealFloat e, Unbox e, VG.Vector vector (Complex e))
  => Int
  -> Int
  -> Int
  -> (Int -> e -> e -> Complex e)
  -> vector (Complex e)
createVector numR2Freqs numAngularFreqs sign f =
  let centerR2 = div numR2Freqs 2
      centerAngular = div numAngularFreqs 2
   in VG.convert .
      toUnboxed .
      computeS .
      makeFilter2D .
      fromFunction (Z :. numAngularFreqs :. numR2Freqs :. numR2Freqs) $ \(Z :. k :. i :. j) ->
        let x = fromIntegral (i - centerR2)
            y = fromIntegral (j - centerR2)
            rho = sqrt $ x ^ 2 + y ^ 2
            phi = atan2 y x
         in if rho == 0
              then 0
              else f (sign * (k - centerAngular)) phi rho

{-# INLINE pinwheel #-}
pinwheel :: (RealFloat e) => e -> e -> e -> Int -> Int -> e -> e -> Complex e
pinwheel periodR2 radialConst alpha radialFreq angularFreq phi rho =
  pi * cis (fromIntegral angularFreq * phi) *
  ((periodR2 / (pi * rho) :+ 0) **
   ((2 + alpha) :+ (radialConst * fromIntegral radialFreq)))

{-# INLINE logpolarHarmonics #-}
logpolarHarmonics ::
     (RealFloat e) => e -> e -> Int -> Int -> e -> e -> Complex e
logpolarHarmonics periodR2 radialConst radialFreq angularFreq phi rho =
  cis (fromIntegral angularFreq * phi) *
  ((periodR2 / (pi * rho) :+ 0) **
   (0 :+ (radialConst * fromIntegral radialFreq)))

{-# INLINE phaseShift #-}
phaseShift ::
     (RealFloat e, Gamma (Complex e)) => Int -> Int -> e -> e -> Complex e
phaseShift angularFreq radialFreq radialConst alpha =
  ((0 :+ (-1)) ^ abs angularFreq) *
  (gamma $
   ((2 + fromIntegral (abs angularFreq) + alpha) :+
    (radialConst * fromIntegral radialFreq)) /
   2) /
  (gamma $
   ((fromIntegral (abs angularFreq) - alpha) :+
    (radialConst * fromIntegral (-radialFreq))) /
   2)

{-# INLINE computeIndicesFromRadius #-}
computeIndicesFromRadius :: Double -> [(Int,Int)]
computeIndicesFromRadius radius =
  let rho = round radius
      r2 = rho ^ 2
   in L.filter
        (\(x, y) ->
           let r = fromIntegral $ x ^ 2 + y ^ 2
            in r < r2)
        [(x, y) | x <- [-rho .. rho], y <- [-rho .. rho]]

{-# INLINE computePinwheelArray #-}
computePinwheelArray ::
     ( Num e
     , RealFloat e
     , Gamma (Complex e)
     , Unbox e
     , NFData e
     )
  => Int
  -> Int
  -> Int
  -> Int
  -> Int
  -> e
  -> e
  -> e
  -> IA.Array (Int, Int) (Complex e)
computePinwheelArray numR2Freqs phiFreq rhoFreq thetaFreq rFreq alpha periodR2 periodEnv =
  let maxAngularFreq = phiFreq + thetaFreq
      maxRadialFreq = rhoFreq + rFreq
      radialConst = 2 * pi / log periodEnv
      centerR2 = div numR2Freqs 2
      idxs =
        [ (radialFreq, angularFreq)
        | radialFreq <- [-maxRadialFreq .. maxRadialFreq]
        , angularFreq <- [-maxAngularFreq .. maxAngularFreq]
        ]
      pinwheels =
        parMap
          rdeepseq
          (\(radialFreq, angularFreq) ->
             let arr =
                   fromFunction (Z :. numR2Freqs :. numR2Freqs) $ \(Z :. i :. j) ->
                     let x = fromIntegral $ i - centerR2
                         y = fromIntegral $ j - centerR2
                         rho = sqrt $ x ^ 2 + y ^ 2
                         phi = atan2 y x
                      in if rho == 0
                           then 0
                           else pinwheel
                                  periodR2
                                  radialConst
                                  alpha
                                  radialFreq
                                  angularFreq
                                  phi
                                  rho *
                                phaseShift
                                  angularFreq
                                  radialFreq
                                  radialConst
                                  alpha
              in sumAllS arr / (fromIntegral numR2Freqs^2))
          idxs
   in listArray
        ((-maxRadialFreq, -maxAngularFreq), (maxRadialFreq, maxAngularFreq))
        pinwheels


{-# LANGUAGE TemplateHaskell #-}
module AI.NeuralNetwork
    ( Network
    , layerSizes
    , biases
    , weights
    , sigmoid
    , newNetwork
    , feedforward
    , toFile
    , fromFile ) where

import           Control.Lens
import           Control.Monad
import           Control.Monad.State
import           Data.Random.Normal
import           Numeric.LinearAlgebra
import           System.IO
import           System.Random

data Network = Network { _layerSizes :: [Int]
                       , _biases     :: [Vector Double]
                       , _weights    :: [Matrix Double] }
                       deriving (Eq, Show)

makeLenses ''Network

sigmoid :: Floating a => a -> a
sigmoid z = 1.0 / (1 + exp (negate z))

newNetwork :: Monad m => [Int] -> StateT StdGen m Network
newNetwork sizes = do
    let matrixSizes = zip (init sizes) (tail sizes)

    biasesL  <- mapM sequence
        [[state normal | _ <- [1..y]] | y <- tail sizes]
    weightsL <- mapM sequence $
        [[state normal | _ <- [1..x], x <- [1..y]] | (x, y) <- matrixSizes]

    let randBiases  = map fromList biasesL
    let randWeights = map (\((x, y), l) -> (y><x) l) (zip matrixSizes weightsL)

    return Network { _layerSizes = sizes
                   , _biases = randBiases
                   , _weights = randWeights }

feedforward :: Vector Double -> Network -> Vector Double
feedforward input nn = foldl (\a (b, w) -> cmap sigmoid ((w #> a) + b)) input bw
    where bw = zip (nn^.biases) (nn^.weights)

toFile :: FilePath -> Network -> IO ()
toFile filepath nn =
    let headerStr  = show (nn^.layerSizes) ++ "\n"
        biasesStr  = unlines $ show . toList  <$> (nn^.biases)
        weightsStr = unlines $ show . toLists <$> (nn^.weights)
        nnStr      = headerStr ++ biasesStr ++ weightsStr
     in writeFile filepath nnStr

fromFile :: FilePath -> IO Network
fromFile filepath = withFile filepath ReadMode (\handle -> do
    layerSizes' <- (read :: String -> [Int]) <$> hGetLine handle
    biases'     <- forM (tail layerSizes') $ \_ -> do
        line <- hGetLine handle
        return $ fromList $ (read :: String -> [Double]) line
    weights'    <- forM (tail layerSizes') $ \_ -> do
        line <- hGetLine handle
        return $ fromLists $ (read :: String -> [[Double]]) line
    return $ Network { _layerSizes = layerSizes'
                     , _biases     = biases'
                     , _weights    = weights' })


module Quantum
  ( QRegister,
    QGate,
    C,
    qbitsNum,
    dimensionsNum,
    qbit,
    qregister,
    qbitL,
    qregisterL,
    makeGateC,
    makeGate,
    (|>),
    (|+>),
    measure,
  )
where

import Data.Complex
import Data.List (sort)
import Matrix
import System.Random (randomRIO)
import Utils

type C = Complex Double

data QRegister = QRegister
  { qbitsNum :: Int,
    states :: Matrix C
  }

instance Show QRegister where
  show (QRegister n xm) = show xm

data QGate = QGate
  { dimensionsNum :: Int,
    operator :: Matrix C
  }

instance Show QGate where
  show (QGate n xm) = show xm

class Entanglable a where
  entangle :: a -> a -> a

instance Entanglable QRegister where
  entangle (QRegister n1 xm) (QRegister n2 ym) = QRegister n zm
    where
      n = n1 + n2
      zm = kron xm ym

instance Entanglable QGate where
  entangle (QGate n1 xm) (QGate n2 ym) = QGate n zm
    where
      n = n1 + n2
      zm = kron xm ym

qbitL :: Int -> [C]
qbitL 0 = [1 :+ 0, 0 :+ 0]
qbitL 1 = [0 :+ 0, 1 :+ 0]
qbitL _ = error "Expecting |0> or |1>"

qbit :: Int -> QRegister
qbit = QRegister 1 . vectorV . qbitL

applyGate :: QGate -> QRegister -> QRegister
applyGate (QGate n1 xm) (QRegister n2 ym)
  | n1 /= n2 = error "Dimensions of gate and number of qbits in register disagree"
  | otherwise = QRegister n2 (mult xm ym)

qregister :: [Int] -> QRegister
qregister = foldl1 entangle . map qbit

qregisterL :: [Int] -> [C]
qregisterL = concat . toList . states . qregister

makeGateC :: Matrix C -> QGate
makeGateC xm
  | r /= c = error "Gate matrix must be square"
  | not . powerOfTwo $ r = error "Gate matrix dimensions must be power of two"
  | otherwise = QGate n xm
  where
    r = rowsNum xm
    c = colsNum xm
    n = log2 r
    powerOfTwo n = odd n || (elem n . take (log2 n + 1) . map (2 ^) $ [1 ..])
    log2 = fromIntegral . truncate . logBase 2.0 . fromIntegral

makeGate :: Matrix Double -> QGate
makeGate = makeGateC . fmap (:+ 0.0)

infixl 7 |>

(|>) :: QRegister -> QGate -> QRegister
(|>) = flip applyGate

infixl 8 |+>

(|+>) :: Entanglable a => a -> a -> a
(|+>) = entangle

measure :: QRegister -> IO Int
measure (QRegister n xm) = do
  rnd <- randomRIO (0, 1) :: IO Double
  let state = length . takeWhile (< rnd) $ probs
  return (state - 1)
  where
    probs = sort . (\xs -> 0.0 : xs ++ [1]) . scanl1 (+) . map ((^ 2) . magnitude) . concat . toList $ xm


module Space where

import Data.Set
import qualified Data.Set as S

import Data.Vector
import qualified Data.Vector as V

import Numeric.LinearAlgebra
import qualified Numeric.LinearAlgebra as M

-- INITIALISATION ---------------------------------------------------------------------------------

-- The data Space (short for simplicial complex) consist of a set of simplices with
-- the following property: all the faces of all the simplices in the set must also
-- belong to the set. A simplex is implemented as a Set for efficiency.
-- For convenience, the set of simplices is split into a list according to the
-- dimension. The defining condition is unchecked as, in practice, the helper
-- function allSimplices (below) is always used to iniatialise the data. The output
-- of the helper function is always valid.

data Space t = Space { subsimplices :: [Set (Set t)] } deriving (Show)

-- This helper function takes a set of simplices as input and returns a valid
-- simplicial complex. The input consists of all the highest dimensional simplices
-- which make up the space. The function takes all their boundaries and sorts
-- them by dimension. The lowest dimension is -1 and it is the dimension of the
-- empty set. It is considered a simplex.
allSimplices :: (Ord t) => [[t]] -> [Set (Set t)]
allSimplices [] = error "This space has no faces. For empty space input '[[]]'."
allSimplices faces =

    let associatedSet = S.fromList (Prelude.map (S.fromList) faces)
        powerFaces = S.map (S.powerSet) associatedSet
        allMixedFaces = S.foldl (S.union) S.empty powerFaces
        lengthFunctions = [ \xs -> (S.size xs == n) | n <- [0,1..] ]

    in  sieveAndForget lengthFunctions allMixedFaces

sieveAndForget :: [t -> Bool] -> Set t -> [Set t]
sieveAndForget (f:fs) xs
    | (S.size xs == 0) = []
    | otherwise        = p1 : (sieveAndForget fs p2)
        where (p1,p2) = S.partition (f) xs

-- HOMOLOGY ---------------------------------------------------------------------------------------

-- This function only computes the free rank of the homology groups. It does not give
-- any information about the torsion subgroups. These require knowledge of the
-- Smith decomposition of a matrix, an algorithm which has never been developed in
-- Haskell (over a PID). To compute the free rank, it is enough to know the rank
-- of the boundary operators (through the package Numeric.LinearALgebra)

-- This function returns the homology numbers of a space from 0 up to its dimension.
-- All the other homology numbers must be zero.
homology :: (Ord t) => Space t -> [Int]
homology (Space { subsimplices = xs }) =
    let operators = chainComplex (Space { subsimplices = xs })
        dimensionDomains = Prelude.map (M.cols) operators
        dimensionImages = Prelude.map (M.rank) (Prelude.tail operators)
        dimensionAugmentedImages = 0:dimensionImages
        dimensionShiftedImages = dimensionImages Prelude.++ [0]
        dimensionKernels = Prelude.zipWith (-) dimensionDomains dimensionAugmentedImages

    in  Prelude.zipWith (-) dimensionKernels dimensionShiftedImages

-- Very small variation on homology. It is used to understand if a space is connected.
reducedHomology :: (Ord t) => Space t -> [Int]
reducedHomology space = reduced
    where reduced = (x-1):xs
          (x:xs) = homology space

-- Computes all boundary operators in a chain complex
chainComplex :: (Ord t) => Space t -> [M.Matrix R]
chainComplex (Space { subsimplices = xs }) =

    let listOfSimplices = Prelude.map (S.toList) xs
        domainListOfSimplices = Prelude.tail listOfSimplices
        imageListOfSimplices = Prelude.init listOfSimplices
        domainVectorOfSimplices = Prelude.map (V.fromList) domainListOfSimplices
        imageVectorOfSimplices = Prelude.map (V.fromList) imageListOfSimplices
        domainImage = Prelude.zip domainVectorOfSimplices imageVectorOfSimplices

    in  Prelude.map (boundaryOperator) domainImage

-- Computes boundary operator for all simplices of a fixed dimension.
boundaryOperator :: (Ord t) => (V.Vector (Set t) , V.Vector (Set t))-> M.Matrix R
boundaryOperator (domain , image) = convertToMatrix vectorMatrix
    where vectorMatrix = V.map (f) domain
          f = \x -> (piecewiseBoundaryOperator x image)

-- Computes boundary operator of a single simplex.
piecewiseBoundaryOperator :: (Ord t) => Set t -> V.Vector (Set t) -> V.Vector Int
piecewiseBoundaryOperator element image =
    let f imageFace
              | (what == False) = 0
              | odd indicator   = 1
              | otherwise       = -1

              where (s1 , what , s2) = S.splitMember imageFace facesOfElement
                    -- Shortcut to compute the sign of a face
                    indicator = S.size element + S.size s1
                    facesOfElement = S.map (g) element
                    g = \y -> (S.delete y element)

    in  V.map (f) image

-- MATRIX CONVERSION ------------------------------------------------------------------------------

-- These algorithms are necessary to convert from V.Vector to M.Matrix.
-- Unfortunately these two data lack compatibility. Furthermore
-- M.Vector cannot store Sets, hence they cannot be used from the beginning.
-- Ideally, the Smith Normal Form should be called directly on the V.Vector
-- so to avoid using M.Matrix all together.

convertToMatrix :: V.Vector (V.Vector Int) -> M.Matrix R
convertToMatrix vectorMatrix = M.matrix s doubleFlattenedListMatrix
    where s = V.length vectorMatrix
          doubleFlattenedListMatrix = Prelude.map (fromIntegral) flattenedListMatrix
          flattenedListMatrix = Prelude.foldl (Prelude.++) [] listMatrix
          listMatrix = Prelude.map (V.toList) listVectorMatrix
          listVectorMatrix = V.toList transposedVectorMatrix
          transposedVectorMatrix = transpose vectorMatrix

transpose :: V.Vector (V.Vector Int) -> V.Vector (V.Vector Int)
transpose matrix = V.foldr (g) emptyVector matrix
    where emptyVector = V.replicate s V.empty
          s = V.length ((V.!) matrix 0)
          g = \x y -> V.zipWith (V.cons) x y

-- BASIC FUNCTIONS --------------------------------------------------------------------------------

-- These functions are basic as they do not rely on computing homology.

-- The function returns the dimension of the highest dimensional simplices.
-- 2 is subtracted to account for the simplices of dimension 0 and -1.
dimension :: (Ord t) => Space t -> Int
dimension (Space { subsimplices = xs }) = Prelude.length xs - 2

-- This function returns the number of simplices of a given dimension.
numberOfSimplices :: (Ord t) => Space t -> Int -> Int
numberOfSimplices (Space { subsimplices = xs }) dim
    | dim < -1 || dim > maxDim = 0
    | otherwise                   = S.size simplicesDim

    where maxDim = dimension (Space { subsimplices = xs })
          simplicesDim = xs !! index
          index = dim + 1

eulerCharacteristic :: (Ord t) => Space t -> Int
eulerCharacteristic (Space { subsimplices = xs }) =
    alternatingSum (Prelude.tail lengths)

        where alternatingSum :: (Num t) => [t] -> t
              alternatingSum [] = error "Not defined."
              alternatingSum [x] = x
              alternatingSum (x:xs) = x - alternatingSum xs

              lengths = (Prelude.map (S.size) xs)

-- ADVANCED FUNCTIONS -----------------------------------------------------------------------------

-- These functions all rely on homology.

isConnected :: (Ord t) => Space t -> Bool
isConnected space = (zeroHomology == 1)
    where zeroHomology = Prelude.head (homology space)

-- This function can tell if a space is not homotopic to a topological manifold.
-- Note that, for the reasons above, it only checks Poincaré duality up to torsion.
rationalPoincareDuality :: (Ord t) => Space t -> Bool
rationalPoincareDuality space = (homNumbers == revHomNumbers)
    where revHomNumbers = Prelude.reverse homNumbers
          homNumbers = homology space

-- If this function returns False then the answer is definitely No. If it returns
-- True, then the algorthm is inconcludent. The function first checks the Euler
-- characteristic as this is a less expensice invariant to compute, albeit coarser.
-- If it is the same, homology ranks are computed.
couldItBeContractible :: (Ord t) => Space t -> Bool
couldItBeContractible space = (euler == 1) && (isZero reduced)
    where euler = eulerCharacteristic space
          reduced = reducedHomology space

-- Simple helper function to check if all the entries in a list are zero
isZero :: [Int] -> Bool
isZero [] = error "Empty list."
isZero [0] = True
isZero (x:xs) = (x == 0) && (isZero xs)

-- Again, if the function returns True, no new information is gained. The
-- Euler characteristic is computed first. This algorithm only checks
-- if the homology numbers are the same, however, it is necessary to add
-- enough zeros to the shorter list to equate the lengths.
couldTheyBeHomotopic :: (Ord t) => Space t -> Space t -> Bool
couldTheyBeHomotopic space1 space2 = (euler1 == euler2) && (hom1 == hom2)
    where euler1 = eulerCharacteristic space1
          euler2 = eulerCharacteristic space2
          hom1 = tempHom1 Prelude.++ zeroList1
          hom2 = tempHom2 Prelude.++ zeroList2
          tempHom1 = homology space1
          tempHom2 = homology space2
          zeroList1 = [ 0 | x <- [1..s1]]
          zeroList2 = [ 0 | x <- [1..s2]]
          s1 = maxLength - length1
          s2 = maxLength - length2
          length1 = Prelude.length tempHom1
          length2 = Prelude.length tempHom2
          maxLength = Prelude.maximum [length1, length2]


{-# LANGUAGE FlexibleInstances #-}

module Main where

import Control.Monad
import Graphics.Rendering.Chart
import Data.Colour
import Data.Colour.Names
import Data.Default.Class
import Graphics.Rendering.Chart.Backend.Cairo
import Control.Lens
import Data.Complex

setLinesBlue :: PlotLines a b -> PlotLines a b
setLinesBlue = plot_lines_style  . line_color .~ opaque blue

chart :: Renderable ()
chart = toRenderable layout
  where
    sumf = plot_lines_values .~ [[ (x,(s x)) | x <- [1.3,(1.31)..10]]]
              $ plot_lines_style  . line_color .~ opaque black
              $ plot_lines_title .~ "S(a)"
              $ def

    tln = plot_lines_values .~ [[ (x,(taylor1 2.0 x)) | x <- [1.3,(1.4)..3.5]]]
              $ plot_lines_style  . line_color .~ opaque red
              $ plot_lines_title .~ "First taylor"
              $ def

    tln2 = plot_lines_values .~ [[ (x,(taylor2 2.0 x)) | x <- [1.3,(1.4)..4.0]]]
              $ plot_lines_style  . line_color .~ opaque orange
              $ plot_lines_title .~ "Second taylor"
              $ def

    tln3 = plot_lines_values .~ [[ (x,(taylor3 2.0 x)) | x <- [1.3,(1.4)..3.2]]]
              $ plot_lines_style  . line_color .~ opaque blue
              $ plot_lines_title .~ "Third taylor"
              $ def

    layout = layout_title .~ "Amplitude Modulation"
           $ layout_plots .~ [toPlot sumf, toPlot tln, toPlot tln2, toPlot tln3]
           $ def

renderChart :: IO (PickFn ())
renderChart = renderableToFile def "example1_big.png" chart

approxS :: Int -> Double -> Double
approxS i a = sum $ map (f.fromIntegral) [1..i]
  where f n = n**(-a)

class Abs a where
  absoluteValue :: a -> Double

instance Abs Double where
  absoluteValue = abs

instance Abs (Complex Double) where
  absoluteValue = magnitude

untilConverged :: (Abs a, Floating a, Num a) => Double -> [a] -> a
untilConverged epsilon (x1:xs) = go x1 xs
  where go last (y:ys) =
         if absoluteValue (last - y) < epsilon
           then y
           else go y ys

sumSeries :: (Abs a, Floating a, Num a) => (Int -> a) -> a
sumSeries seriesGen = untilConverged 0.000001 $ tail $ scanl1 (+) $ map seriesGen [1..]

s :: Double -> Double
s a = sumSeries $ \n-> (fromIntegral n)**(-a)

sDerivative :: (Abs a, Floating a, Num a) => Int -> a -> a
sDerivative d a = sign * (sumSeries $ \nInt ->
    let n = fromIntegral nInt in
    (fromIntegral nInt)**(-a) * (realToFrac $ log n)^d)
  where sign = if even d then 1 else (-1)

fact :: Int -> Int
fact 0 = 1
fact n = n * fact (n - 1)

taylorCoeffients :: (Fractional a, Num a) => (Int -> a -> a) -> a -> [a]
taylorCoeffients f a0 = map g [0..]
  where g n = (f n a0) / fromIntegral (fact n)

taylor1 :: Double -> Double -> Double
taylor1 a0 a = s a0 + (a - a0) * sDerivative 1 a0

taylor2 :: Double -> Double -> Double
taylor2 a0 a = s a0 + (a - a0) * sDerivative 1 a0 + 0.5 * (a - a0)^2 * sDerivative 2 a0

taylor3 :: Double -> Double -> Double
taylor3 a0 a = s a0 + (a - a0) * sDerivative 1 a0 + 0.5 * (a - a0)^2 * sDerivative 2 a0 + (1.0/6.0) * (a - a0)^3 * sDerivative 3 a0

main :: IO ()
main = void $ renderChart


module EKF
  ( KF
  , kf_x
  , kf_p
  , kf_t

  , makeFilter
  , predict
  , update
  ) where

import Prelude hiding ((<>))

import App (Sensor (..))

import Numeric.LinearAlgebra

data KF = KF
  { kf_x :: Vector Double
  , kf_p :: Matrix Double
  , kf_t :: Word
  } deriving Show

makeFilter (Laser px py t) =
  KF { kf_x = 4 |> [px,
                    py,
                    0,
                    0]
     , kf_p = (4><4) [1, 0,    0,    0,
                      0, 1,    0,    0,
                      0, 0, 1000,    0,
                      0, 0,    0, 1000]
     , kf_t = t }

makeFilter (Radar rho phi rho' t) =
  KF { kf_x = 4 |> [rho  * cos phi,
                    rho  * sin phi,
                    rho' * cos phi,
                    rho' * sin phi]
     , kf_p =  (4><4) [1, 0,    0,    0,
                       0, 1,    0,    0,
                       0, 0, 1000,    0,
                       0, 0,    0, 1000]
     , kf_t = t }

predict :: Word -> KF -> KF
predict t kf =
  KF { kf_x = x
     , kf_p = p
     , kf_t = t }
  where
    -- Time inside the filter.
    filterTime =
      kf_t kf

    -- Delta-time. Essentially a timestep.
    dt =
      (fromIntegral (t - filterTime)) / 1000000.0
    dt2 = dt ** 2
    dt3 = dt ** 3
    dt4 = dt ** 4

    -- Acceleration noise tuning parameters
    ax = 3
    ay = 3

    noiseAx = ax ** 2
    noiseAy = ay ** 2

    -- State transition matrix
    f = (4><4) [1, 0, dt,  0,
                0, 1,  0, dt,
                0, 0,  1,  0,
                0, 0,  0,  1]

    -- Noise covariance matrix
    q = (4><4) [-- Row 1
                noiseAx * 0.25 * dt4, 0, noiseAx * 0.5 * dt3, 0,
                -- Row 2
                0, noiseAy * 0.25 * dt4, 0, noiseAy * 0.5 * dt3,
                --  Row 3
                noiseAx *0.5 * dt3, 0, noiseAx * dt2, 0,
                -- Row 4
                0, noiseAy * 0.5 * dt3, 0, noiseAy * dt2]

    -- Noise vector
    v  = 4 |> [ax * 0.5 * dt2,
               ay * 0.5 * dt2,
               ax * dt,
               ay * dt]

    -- Estimated step based on constant velocity (CV) motion model
    x = f #> kf_x kf + v
    p = f <> kf_p kf <> (tr f) + q

update (Laser px py t) kf =
  KF { kf_x = x'
     , kf_p = p'
     , kf_t = t }
  where
    kf' = predict t kf

    -- Measurment transition matricies
    h = (2><4) [1, 0, 0, 0,
                0, 1, 0, 0]
    r = (2><2) [0.0225, 0.0000,
                0.0000, 0.0225]
    y = 2 |> [px, py] - (h #> kf_x kf')

    -- Kalman Filter Equations
    i  = ident 4
    ht = tr h
    s  = h <> kf_p kf' <> ht + r
    si = inv s
    k  = kf_p kf' <> ht <> si

    -- Estimation based on kalman filter gain
    x' = kf_x kf' + (k #> y)
    p' = (i - k <> h) <> kf_p kf'

update (Radar rho phi rho' t) kf =
  KF { kf_x = x'
     , kf_p = p'
     , kf_t = t }
  where
    kf' = predict t kf

    -- Calculate the Jacobian matrix values
    px = kf_x kf' ! 0
    py = kf_x kf' ! 1
    vx = kf_x kf' ! 2
    vy = kf_x kf' ! 3

    c1 = px ** 2 + py ** 2
    c2 = sqrt c1
    c3 = c1 * c2

    zRho  = c2;
    zPhi  = (atan2 py px)
    zRho' = (px*vx + py*vy) / zRho

    -- Measurment transition matricies
    h = (3><4) [-- Row 1
                px / c2, py / c2, 0, 0,
                -- Row 2
                -py / c1, px / c1, 0, 0,
                -- Row 3
                py*(vx * py - vy * px)/c3,
                px*(vy * px - vx * py)/c3,
                px/c2,
                py/c2]
    r = (3><3) [0.09, 0.0000, 0.00,
                0.00, 0.0009, 0.00,
                0.00, 0.0000, 0.09]
    z = 3 |> [rho, phi, rho'] - 3 |> [zRho, zPhi, zRho']
    y = 3 |> [z ! 0, atan2 (sin (z ! 1)) (cos (z ! 1)), z ! 2]

    -- Kalman filter equations
    i  = ident 4
    ht = tr h
    s  = h <> kf_p kf' <> ht + r
    si = inv s
    k  = kf_p kf' <> ht <> si

    -- Estimation based on kalman filter gain
    x' = kf_x kf' + (k #> y)
    p' = (i - k <> h) <> kf_p kf'


module Echo where

import Data.Complex
import Linear.Metric
import Linear.V3

echo :: (Floating a) => (a -> Complex a) -> V3 a -> V3 a -> a -> Complex a
echo signal carrier obj t = signal (t - 2 * d / c)
  where
    d = distance carrier obj
    c = 299792458


{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE TypeOperators    #-}
module FokkerPlanck.Interpolation where

import           Data.Array.Repa as R
import           Data.Complex
import           Data.List       as L
import           Filter.Utils
import           Types

{-# INLINE radialCubicInterpolation #-}
radialCubicInterpolation ::
     (R.Source r1 Double, R.Source r2 (Complex Double), Shape sh)
  => R.Array r1 (sh :. Int) Double
  -> Double
  -> R.Array r2 (sh :. Int :. Int) (Complex Double)
  -> R.Array D (sh :. Int :. Int) (Complex Double)
radialCubicInterpolation radialArr scaleFactor inputArr =
  let (_ :. n) = extent radialArr
      (_ :. nx :. ny) = extent inputArr
   in if scaleFactor == 0
        then error "radialCubicInterpolation: scaleFactor = 0"
        else R.traverse2 inputArr radialArr const $ \fInput fRadial idx@(sh :. i :. j) ->
               let r =
                     (sqrt . fromIntegral $
                      (i - center nx) ^ 2 + (j - center ny) ^ 2) /
                     scaleFactor
                in if r < 0 || r > (fromIntegral $ n - 1) -- out of boundary. When r = 0, log r makes no sense, and therefore the information at r = 0 should not be used to interpoalte 0 < r <= 1.
                     then 0
                     else if r <= 1 -- linear interpolation
                            then fInput idx *
                                 ((fRadial (sh :. 1) +
                                   (r - 1) *
                                   (fRadial (sh :. 2) - fRadial (sh :. 1))) :+
                                  0)
                            else if r >= (fromIntegral $ n - 2) -- linear interpolation
                                   then fInput idx *
                                        ((fRadial (sh :. n - 2) +
                                          (r - (fromIntegral $ n - 2)) *
                                          (fRadial (sh :. n - 1) -
                                           fRadial (sh :. n - 2))) :+
                                         0)
                                   else let x0 = floor r -- cubic interpolation
                                            p0 = fRadial (sh :. x0 - 1)
                                            p1 = fRadial (sh :. x0)
                                            p2 = fRadial (sh :. x0 + 1)
                                            p3 = fRadial (sh :. x0 + 2)
                                            x = r - fromIntegral x0
                                         in fInput idx *
                                            (((-0.5 * p0 + 1.5 * p1 - 1.5 * p2 +
                                               0.5 * p3) *
                                              x ^ 3 +
                                              (p0 - 2.5 * p1 + 2 * p2 - 0.5 * p3) *
                                              x ^ 2 +
                                              (-0.5 * p0 + 0.5 * p2) * x +
                                              p1) :+
                                             0)

{-# INLINE cubicInterpolation #-}
cubicInterpolation ::
     (R.Source s Double)
  => R.Array s DIM5 Double
  -> Double
  -> R.Array D DIM4 Double
cubicInterpolation arr r =
  let (Z :. _ :. _ :. _ :. _ :. n) = extent arr
   in R.traverse arr (\(sh :. _) -> sh) $ \f sh@(Z :. tf :. sf :. t0f :. s0f) ->
        if r < 0 || r > (fromIntegral $ n - 1) -- out of boundary. When r = 0, log r makes no sense, and therefore the information at r = 0 should not be used to interpoalte 0 < r <= 1.
          then 0
          else if r <= 1 -- linear interpolation
                 then (f (sh :. 1) + (r - 1) * (f (sh :. 2) - f (sh :. 1)))
                 else if r >= (fromIntegral $ n - 2) -- linear interpolation
                        then (f (sh :. n - 2) +
                              (r - (fromIntegral $ n - 2)) *
                              (f (sh :. n - 1) - f (sh :. n - 2)))
                        else let x0 = floor r -- cubic interpolation
                                 p0 = f (sh :. x0 - 1)
                                 p1 = f (sh :. x0)
                                 p2 = f (sh :. x0 + 1)
                                 p3 = f (sh :. x0 + 2)
                                 x = r - fromIntegral x0
                              in ((-0.5 * p0 + 1.5 * p1 - 1.5 * p2 + 0.5 * p3) *
                                  x ^ 3 +
                                  (p0 - 2.5 * p1 + 2 * p2 - 0.5 * p3) * x ^ 2 +
                                  (-0.5 * p0 + 0.5 * p2) * x +
                                  p1)


{-# LANGUAGE FlexibleContexts #-}

module Learning
    ( ols
    , gls
    , kmeans
    , Centroids
    , Clusters
    , Clustering
    , Initializer
    , average
    , invertMap
    ) where

import Debug.Trace
import Data.Set
import Data.Map.Strict
import Numeric.LinearAlgebra hiding ((!))

--covariance :: (Field n, Ord n, Num (Vector n))
--           => Set (Vector n)
--           -> Matrix n

--covariance set = exxt - uut where
--  exxt = mean $ Prelude.map (\a -> asColumn a <> asRow a) (Data.Set.elems set)
--  uut  = asColumn mu <> asRow mu
--  mu   = average set

-- | Ordinary least squares
ols :: (Field n, Monoid (Matrix n))
    => Matrix n
    -> Vector n
    -> Vector n

ols x y = head $ toColumns $ optimiseMult [l', tr l', tr x, asColumn y] where
  xTx = mTm x
  l   = chol xTx
  l'  = inv l

-- | Generalized least squares
gls :: (Eq n, Num (Vector n), Monoid (Matrix n), Field n)
    => Matrix n
    -> Matrix n
    -> Vector n
    -> Vector n

gls c x y = t where
  x' = tr x
  ic = inv c
  r  = inv (optimiseMult [x', ic, x])
  s  = optimiseMult [r, x, ic]
  t  = s #> y

type Centroids n = Set (Vector n)

type Clusters n = Map (Vector n) (Set (Vector n))

type Clustering n = (Centroids n, Clusters n)

type Initializer m n = Int -> Set (Vector n) -> m (Centroids n)

-- | KMeans clustering algorithm
kmeans :: (Functor m, Container Vector n, Ord n, Num (Vector n), Normed (Vector n), Field n)
       => Initializer m n
       -> Int
       -> Set (Vector n)
       -> m (Clustering n)

-- KMeans runs by first getting an original clustering, then reclustering
-- until a fixed point is reached.
kmeans init k set = (untilEqual fst (recluster k) . (\centroids -> (centroids, reassign centroids set))) <$> init k set

untilEqual :: Eq a => (e -> a) -> (e -> e) -> e -> e
untilEqual g f e = if g e' == g e then e else untilEqual g f e' where
  e' = f e

-- Reclustering means taking each cluster, computing a new centroid based on
-- the contents, and then reassigning each vector to the cluster which has a
-- centroid closest to it.
recluster :: (Container Vector n, Ord n, Fractional n, Normed (Vector n), Num (Vector n), Field n)
          => Int
          -> Clustering n
          -> Clustering n
recluster k (centroids, clusters) = (centroids', reassign centroids' (Data.Set.unions $ Data.Map.Strict.elems clusters)) where
  centroids' = Data.Set.map (\centroid -> average (clusters ! centroid)) centroids

average :: (Container c n, Fractional n, Ord n, Num (c n), Linear n c)
        => Set (c n) -> c n

average set = scale (1 / (fromInteger . toInteger $ Data.Set.size set)) $ Data.Set.foldr (+) one set' where
  (one, set') = Data.Set.deleteFindMin set

-- Reassignment goes like this: given a set of vectors and a set of centroids,
-- return a map from centroid to sets of vectors which have been found to be within this
-- centroid.
reassign :: (Normed (Vector n), Ord n, Eq n, Num (Vector n), Field n)
         => Set (Vector n)
         -> Set (Vector n)
         -> Clusters n

reassign centroids points = invertMap $ fromSet (getClosest centroids) points

-- | Given a map from x to y, returns a map from y to the preimage of y in the
--   original map
invertMap :: (Ord y, Ord x) => Map x y -> Map y (Set x)
invertMap map = st where
  st = fromListWith Data.Set.union [ (e, Data.Set.singleton k)
                                   | k <- ks, e <- es, map ! k == e ]
  es = Data.Map.Strict.elems map
  ks = keys map

getClosest :: (Normed (Vector n), Num (Vector n), Field n, Eq n) => Set (Vector n) -> Vector n -> Vector n
getClosest points point = if Data.Set.null points then error "getClosest" else fst $ Data.Set.foldr
  (\candidate (c, d) -> let d' = norm_2 (candidate - point)
                        in if d > d' then (candidate, d') else (c, d))
  (candidate, norm_2 (candidate - point))
   points' where
  (candidate, points') = Data.Set.deleteFindMin points


{-# LANGUAGE BangPatterns        #-}
{-# LANGUAGE CPP                 #-}
{-# LANGUAGE DataKinds           #-}
{-# LANGUAGE DeriveAnyClass      #-}
{-# LANGUAGE DeriveGeneric       #-}
{-# LANGUAGE FlexibleContexts    #-}
{-# LANGUAGE GADTs               #-}
{-# LANGUAGE PolyKinds           #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeFamilies        #-}
{-# LANGUAGE TypeOperators       #-}
{-|
Module      : Grenade.Core.WeightInitialization
Description : Defines the Weight Initialization methods of Grenade.
Copyright   : (c) Manuel Schneckenreither, 2018
License     : BSD2
Stability   : experimental

This module defines the weight initialization methods.

-}

module Grenade.Core.WeightInitialization
    ( getRandomVector
    , getRandomVectorV
    , getRandomMatrix
    , getRandomList
    , WeightInitMethod (..)
    ) where

import           Control.DeepSeq
import           Control.Monad
import           Control.Monad.Primitive         (PrimBase, PrimState)
import           Data.Proxy
import           Data.Serialize
import           Data.Singletons.TypeLits
import qualified Data.Vector.Storable            as V
import           GHC.Generics                    (Generic)
import           GHC.TypeLits                    hiding (natVal)
import           Numeric.LinearAlgebra.Static
import           System.IO.Unsafe                (unsafePerformIO)
import           System.Random.MWC
import           System.Random.MWC.Distributions

import           Grenade.Types
import           Grenade.Utils.Vector

import           Debug.Trace

-- ^ Weight initialization method.
data WeightInitMethod
  = UniformInit -- ^ W_l,i ~ U(-1/sqrt(n_l),1/sqrt(n_l))                   where n_l is the number of nodes in layer l
  | Xavier      -- ^ W_l,i ~ U(-sqrt (6/n_l+n_{l+1}),sqrt (6/n_l+n_{l+1})) where n_l is the number of nodes in layer l
  | HeEtAl      -- ^ W_l,i ~ N(0,sqrt(2/n_l))                              where n_l is the number of nodes in layer l
  deriving (Show, Eq, Ord, Enum, Bounded, NFData, Serialize, Generic)

type LayerInput = Integer
type LayerOutput = Integer

getRandomList ::
     (PrimBase m)
  => LayerInput
  -> LayerOutput
  -> Int -- ^ Length of list
  -> WeightInitMethod
  -> Gen (PrimState m)
  -> m [RealNum]
getRandomList i o n method gen = do
  let unifRands = replicateM n (uniformR (-1, 1) gen)
      gaussRands = replicateM n (realToFrac <$> standard gen)
  case method of
    UniformInit -> map (1 / sqrt (fromIntegral i) *) <$> unifRands
    Xavier      -> map ((sqrt 6 / sqrt (fromIntegral i + fromIntegral o)) *) <$> unifRands
    HeEtAl      -> map (sqrt (2 / fromIntegral i) *) <$> gaussRands

getRandomVectorV ::
     (PrimBase m)
  => LayerInput
  -> LayerOutput
  -> Int -- ^ Length of list
  -> WeightInitMethod
  -> Gen (PrimState m)
  -> m (V.Vector RealNum)
getRandomVectorV i o n method gen = do
  let mkVec :: [RealNum] -> V.Vector RealNum
      mkVec xs = V.imap (\idx (_ :: RealNum) -> xs' V.! idx) (createVectorUnsafe n)
        where
          xs' = V.fromList xs -- convert first to vector, otherwise this function is quadratic!
  let unifRands = replicateM n (uniformR (-1, 1) gen)
      gaussRands = replicateM n (realToFrac <$> standard gen)
  case method of
    UniformInit -> mkVec . map (1 / sqrt (fromIntegral i) *) <$> unifRands
    Xavier -> mkVec . map ((sqrt 6 / sqrt (fromIntegral i + fromIntegral o)) *) <$> unifRands
    HeEtAl -> mkVec . map (sqrt (2 / fromIntegral i) *) <$> gaussRands

-- test :: IO [Double]
-- test = withSystemRandom $ asGenIO $ \gen -> replicateM 10000 (uniformR (-1, 1) gen)

-- test :: Int -> IO (V.Vector Double)
-- test n = withSystemRandom $ asGenIO $ getRandomVectorV 10 100 n UniformInit

-- | Get a random vector initialized according to the specified method.
getRandomVector ::
     forall m n. (PrimBase m, KnownNat n)
  => Integer
  -> Integer
  -> WeightInitMethod
  -> Gen (PrimState m)
  -> m (R n)
getRandomVector i o method gen = do
  let unifRands = vector <$> replicateM n (uniformR (-1, 1) gen)
      gaussRands = vector <$> replicateM n (realToFrac <$> standard gen)
  case method of
    UniformInit -> ((1 / sqrt (fromIntegral i)) *) <$> unifRands
    Xavier      -> ((sqrt 6 / sqrt (fromIntegral i + fromIntegral o)) *) <$> unifRands
    HeEtAl      -> (sqrt (2 / fromIntegral i) *) <$> gaussRands
  where
    n = fromIntegral $ natVal (Proxy :: Proxy n)

-- | Get a matrix with weights initialized according to the specified method.
getRandomMatrix ::
     forall m r n. (PrimBase m, KnownNat r, KnownNat n, KnownNat (n * r))
  => Integer
  -> Integer
  -> WeightInitMethod
  -> Gen (PrimState m)
  -> m (L r n)
getRandomMatrix i o method gen = do
  let unifRands = matrix <$> replicateM nr (uniformR (-1, 1) gen)
      gaussRands = matrix <$> replicateM nr (realToFrac <$> standard gen)
  case method of
    UniformInit -> ((1 / sqrt (fromIntegral i)) *) <$> unifRands
    Xavier      -> ((sqrt 6 / sqrt (fromIntegral i + fromIntegral o)) *) <$> unifRands
    HeEtAl      -> (sqrt (2 / fromIntegral i) *) <$> gaussRands
  where
    nr = fromIntegral $ natVal (Proxy :: Proxy (n * r))


{-# OPTIONS_GHC -Wall                   #-}
{-# OPTIONS_GHC -fno-warn-type-defaults #-}

{-# LANGUAGE DataKinds #-}

module HaddockExample ( main ) where

import           Numeric.LinearAlgebra.Static
                 ( R, vector, Sym,
                   headTail, matrix, sym
                 )
import           Data.Random.Distribution.MultiNormal
import           Numeric.Kalman
import           Numeric.LinearAlgebra.Static

import Graphics.Rendering.Chart hiding ( Vector )
import Graphics.Rendering.Chart.Backend.Diagrams
import Diagrams.Backend.Cairo.CmdLine
import Diagrams.Prelude hiding ( render, Renderable, sample )
import Diagrams.Backend.CmdLine

import Data.Csv
import System.IO hiding ( hGetContents )
import Data.ByteString.Lazy ( hGetContents )
import qualified Data.Vector as V

deltaT, g :: Double
deltaT = 0.01
g  = 9.81

bigQ :: Sym 2
bigQ = sym $ matrix bigQl

qc1 :: Double
qc1 = 0.01

bigQl :: [Double]
bigQl = [ qc1 * deltaT^3 / 3, qc1 * deltaT^2 / 2,
           qc1 * deltaT^2 / 2,       qc1 * deltaT
         ]

bigR :: Sym 1
bigR  = sym $ matrix [0.1]

observe :: R 2 -> R 1
observe a = vector [sin x] where x = fst $ headTail a

linearizedObserve :: R 2 -> L 1 2
linearizedObserve a = matrix [cos x, 0.0] where x = fst $ headTail a

stateUpdate :: R 2 -> R 2
stateUpdate u =  vector [x1 + x2 * deltaT, x2 - g * (sin x1) * deltaT]
  where
    (x1, w) = headTail u
    (x2, _) = headTail w

linearizedStateUpdate :: R 2 -> Sq 2
linearizedStateUpdate u = matrix [1.0,                    deltaT,
                                  -g * (cos x1) * deltaT,    1.0]
  where
    (x1, _) = headTail u

singleEKF :: MultiNormal (R 2) -> R 1 -> MultiNormal (R 2)
singleEKF = runEKF (const observe) (const linearizedObserve) (const bigR)
             (const stateUpdate) (const linearizedStateUpdate) (const bigQ)
             undefined

singleUKF :: MultiNormal (R 2) -> R 1 -> MultiNormal (R 2)
singleUKF = runUKF (const observe) (const bigR) (const stateUpdate) (const bigQ)
             undefined

initialDist :: MultiNormal (R 2)
initialDist = MultiNormal (vector [1.6, 0.0])
                          (sym $ matrix [0.1, 0.0,
                                         0.0, 0.1])

multiEKF :: [ℝ] -> [MultiNormal (R 2)]
multiEKF obs = scanl singleEKF initialDist (map (vector . pure) obs)

multiUKF :: [ℝ] -> [MultiNormal (R 2)]
multiUKF obs = scanl singleUKF initialDist (map (vector . pure) obs)

denv :: IO (DEnv Double)
denv = defaultEnv vectorAlignmentFns 600 500

diagEstimated :: String ->
                 [(Double, Double)] ->
                 [(Double, Double)] ->
                 [(Double, Double)] ->
                 IO (Diagram Cairo)
diagEstimated t l xs es = do
  env <- denv
  return $ fst $ runBackend env (render (chartEstimated t l xs es) (600, 500))

chartEstimated :: String ->
              [(Double, Double)] ->
              [(Double, Double)] ->
              [(Double, Double)] ->
              Renderable ()
chartEstimated title acts obs ests = toRenderable layout
  where

    actuals = plot_lines_values .~ [acts]
            $ plot_lines_style  . line_color .~ opaque red
            $ plot_lines_title .~ "Actual Trajectory"
            $ plot_lines_style  . line_width .~ 1.0
            $ def

    measurements = plot_points_values .~ obs
                 $ plot_points_style  . point_color .~ opaque blue
                 $ plot_points_title .~ "Measurements"
                 $ def

    estimas = plot_lines_values .~ [ests]
            $ plot_lines_style  . line_color .~ opaque black
            $ plot_lines_title .~ "Inferred Trajectory"
            $ plot_lines_style  . line_width .~ 1.0
            $ def

    layout = layout_title .~ title
           $ layout_plots .~ [toPlot actuals, toPlot measurements, toPlot estimas]
           $ layout_y_axis . laxis_title .~ "Angle / Horizontal Displacement"
           $ layout_y_axis . laxis_override .~ axisGridHide
           $ layout_x_axis . laxis_title .~ "Time"
           $ layout_x_axis . laxis_override .~ axisGridHide
           $ def

displayHeader :: FilePath -> Diagram B -> IO ()
displayHeader fn =
  mainRender ( DiagramOpts (Just 900) (Just 700) fn
             , DiagramLoopOpts False Nothing 0
             )

main :: IO ()
main = do
  h <- openFile "matlabRNGs.csv" ReadMode
  cs <- hGetContents h
  let df = (decode NoHeader cs) :: Either String (V.Vector (Double, Double))
  case df of
    Left _ -> error "Whatever"
    Right generatedSamples -> do
      let xs = take 500 (multiEKF $ V.toList $ V.map fst generatedSamples)
      let mus = map (fst . headTail . mu) xs
      let obs = V.toList $ V.map fst generatedSamples
      let acts = V.toList $ V.map snd generatedSamples
      de1 <- diagEstimated "Fitted Pendulum"
             (zip [0,1..] acts)
             (zip [0,1..] obs)
             (zip [0,1..] mus)
      displayHeader "diagrams/PendulumFittedEkf.png" de1
      let ys = take 500 (multiUKF $ V.toList $ V.map fst generatedSamples)
      let nus = map (fst . headTail . mu) ys
      de2 <- diagEstimated "Fitted Pendulum"
             (zip [0,1..] acts)
             (zip [0,1..] obs)
             (zip [0,1..] nus)
      displayHeader "diagrams/PendulumFittedUkf.png" de2


{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE RecordWildCards #-}
import System.IO
import Control.Monad
import Data.Complex
import Data.List
import Data.ByteString (ByteString)

data Obj2D = Obj2D [String]
data Obj3D = Obj3D [String]

class Render2D a where
  render2d :: a -> [Complex Double]

class Render3D a where
  render3d :: a -> Obj3D

data Gear = Gear {
  gAlpha :: Double
, gTooth :: Double
, gModule :: Double
, gWidth :: Double
} deriving Show

data Bar = Bar {
  bWidth :: Double
} deriving Show

involute :: Double -> Double -> Complex Double
involute r theta = (r*(cos theta + theta * sin theta)) :+ (r*(sin theta - theta * cos theta))

yrev :: Complex Double -> Complex Double
yrev (a:+b) = (a:+(-b))

alpha :: Gear -> Double
alpha Gear{..} = 2 * pi * gAlpha / 360

radius :: Gear -> Double
radius g@Gear{..} = (cos (alpha g)) * radius' g

radius' :: Gear -> Double
radius' Gear{..} = gModule * gTooth / 2

instance Render2D Gear where
  render2d g@Gear{..} =
    let dat :: [Complex Double]
        dat = ((radius g) - 0.5) :+ 0 : map (\i-> involute (radius g) (i*2.0*pi/100.0))  [0..10]
        dat''' :: [Complex Double]
        dat''' = map (\i-> mkPolar (radius g) (i*2.0*pi/100.0))  [0..100]
        gear''  =  concat $
                   flip map [0..(gTooth-1)] $
                   \t -> map (\i -> i * (cis (t*2*pi/gTooth) )) $
                   dat ++
                   map (\t -> t * cis (0.9*pi/gTooth))
                   (reverse (map yrev dat))
    in gear''

instance Render3D Gear where
  render3d g@Gear{..} =
    let gear2d = render2d g
        gear2d' = intercalate "," $ map (\(r:+i) -> "[" ++ show r ++ "," ++ show i ++ "]") gear2d
    in Obj3D ["linear_extrude(height = " ++ show gWidth ++ ") { polygon( points= [",
              gear2d',
              "] );}"]

writePlot :: String -> [Complex Double] -> IO ()
writePlot file dat = do
  withFile file WriteMode $ \h -> do
    forM_ dat $ \(r:+i) -> do
      hPutStrLn h $ show r ++ " " ++ show i

write3DPlot :: String -> Obj3D -> IO ()
write3DPlot file (Obj3D dat) = do
  withFile file WriteMode $ \h -> do
    forM_ dat $ \str -> do
      hPutStr h $ str

main = do
  writePlot "involute.dat" $ render2d $ Gear 20 8 1 6
  write3DPlot "involute.scad" $ render3d $ Gear 20 8 1 6


module Math.FROG.Tools
( mkTrace
, mkSigGate
, shift
, mkInitialGuess
, flatPolar2Complex
, center
) where

import qualified System.Random as R
import qualified Numeric.LinearAlgebra as LA
import qualified Numeric.GSL.Fourier as F
import qualified Foreign

import Math.FROG.Types

mkTrace :: ComplexSignal -> ComplexSignal -> Trace
mkTrace field gate = LA.fromColumns $ map atEachDelayApply [upper,upper-1..lower]
    where
    atEachDelayApply = (^2) . (LA.mapVector LA.magnitude)
                            . (flip shift upper)
                            .  F.fft
                            . (*field)
                            . (shift gate)
    sz = LA.dim field
    upper = (sz `quot` 2)
    lower = (-1) * (sz `quot` 2) + 1

mkSigGate :: Nonlinearity -> ComplexSignal -> ComplexSignal -> (ComplexSignal, ComplexSignal)
mkSigGate PG f g = (id f, LA.mapVector ((^2) . abs) g)
mkSigGate SHG f g = (id f, id g)
mkSigGate SD f g = (LA.mapVector (^2) f, LA.mapVector LA.conjugate g)
mkSigGate THG f g = (LA.mapVector (^2) f, id g)

shift :: ComplexSignal -> Int -> ComplexSignal
shift gate delay = LA.buildVector sz (builder delay)
    where
    sz = LA.dim gate
    builder d k = if any (<0) [(sz-1-d-k), (k+d)] then (getElem (k+d-(signum d)*sz)) else (getElem (k+d))
        where
        getElem = (LA.@>) gate

mkInitialGuess :: Int -> IO (LA.Vector Double)
mkInitialGuess sz = do
    stdgen <- R.getStdGen
    let randMags = R.randomRs (0, 1.0) stdgen
    let randAngs = R.randomRs (-pi, pi) stdgen
    return $ LA.buildVector (2*sz) (\k -> if k < sz then (randMags !! k) else (randAngs !! k))

flatPolar2Complex :: LA.Vector Double -> LA.Vector (LA.Complex Double)
flatPolar2Complex v = LA.buildVector sz (\k -> LA.mkPolar (mags LA.@> k) (angs LA.@> k))
    where
    sz = (LA.dim v) `quot` 2
    mags = LA.subVector 0 sz v
    angs = LA.subVector sz sz v

center :: ComplexSignal -> ComplexSignal
center f = shift f $ getDelay . getMaxInd $ f
    where
    sz = LA.dim f
    getMaxInd = LA.maxIndex . (LA.mapVector LA.magnitude)
    getDelay = (*(-1)) . ((-) (sz `quot` 2))