{-# LANGUAGE FlexibleInstances, MultiParamTypeClasses #-}

module Physics.Quantum.Examples.HarmonicOscillator where

import Physics.Quantum.Algebra

import Numeric.Additive.Class          (Additive((+)), Abelian)
import Numeric.Algebra.Class           ( Multiplicative((*), pow1p)
                                       , Monoidal(..), Semiring, LeftModule(..), RightModule(..)
                                       )

import Numeric.Algebra.Division        (Division(..))
import Numeric.Domain.Euclidean        (Euclidean(..))
import Numeric.Domain.PID              (PID(..))
import Numeric.Domain.UFD              (UFD)
import Numeric.Domain.GCD              (GCDDomain(..))
import Numeric.Domain.Integral         (IntegralDomain(..))
import Numeric.Algebra.Unital          (Unital(..))
import Numeric.Algebra.Unital.UnitNormalForm (UnitNormalForm(..))
import Numeric.Decidable.Zero          (DecidableZero(..))
import Numeric.Semiring.ZeroProduct    (ZeroProductSemiring)
import Numeric.Algebra.Commutative     (Commutative)
import Numeric.Decidable.Units         (DecidableUnits(..))
import Numeric.Ring.Class              (Ring(..))
import Numeric.Decidable.Associates    (DecidableAssociates(..))
import Numeric.Rig.Class               (Rig(..))
import Numeric.Additive.Group          (Group(..))

import GHC.Natural                     (Natural)
import GHC.Float                       (double2Float)

import Data.Array.Repa.Eval            (Elt(touch))

import Data.Complex                    (Complex(..), magnitude, realPart)

import Prelude hiding ((+),(-),(/),(*), (^))
import qualified Prelude as P

import qualified Data.Map as M
import qualified Data.Array.Repa as R
import Data.Array.Repa                              (Z(..), (:.)(..))
import Data.Functor.Identity (runIdentity)

import qualified Data.Eigen.SparseMatrix as ES
import qualified Data.Eigen.Matrix as ESD
import Data.Eigen.Matrix                            (CComplex)
import Foreign.C.Types                              (CDouble)

import Data.Maybe                                   (catMaybes)
import qualified Data.List as L

import Graphics.Gloss
import Graphics.Gloss.Data.ViewPort                 (ViewPort)
import Graphics.Gloss.Raster.Field                  (rgb', makePicture)

import Debug.Trace           (trace)
import Control.Concurrent    (threadDelay)

--------------------
-- Add dimensions to quantities
-- Since we are writing operators as endomorphisms of vector spaces
-- if we have to give dimensiones to operators/eigenvalues,
-- we may need to declare the algebra of dimensions as instance of field too
-- Or we may need to relax the condition of operator to homomorphims of vector spaces
-------------------

------------------------------------------
-- Configure
------------------------------------------

data OscillatorInABoxConfig = OscillatorInABoxConfig { plankConstantConfig :: Double
                                                     , particleMassConfig :: Double
                                                     , oscillatorFrequencyConfig :: Double
                                                     , halfWidthConfig :: Int
                                                     , halfHeightConfig :: Int
                                                     , initialAmplitudeConfig :: Position -> Complex Double
                                                     , fpsConfig :: Int
                                                     -- the simulation will play this times faster than real time
                                                     -- mainly to keep time steps small while keeping fps low
                                                     -- the time step in each simulation will be this divided by fps
                                                     , timeFactorConfig :: TimeStep
                                                     }

type TimeStep = Double

data Position = Position { posX :: Int , posY :: Int }
    deriving (Eq, Ord, Show)

type Variance = Double

gaussianWave :: Int -> Int -> Variance -> (Position -> Complex Double)
gaussianWave xCenter yCenter var = \(Position x y) ->
    ( exp ( (
              - ( ((fromIntegral (x - xCenter))P.^2) P.+ ((fromIntegral (y - yCenter) )P.^2) )
            ) P./ (2 P.* (var P.^ 2))
          )
    ) :+ 0

smallBoxConfig :: OscillatorInABoxConfig
smallBoxConfig = OscillatorInABoxConfig { plankConstantConfig = 6.62607004e-16 -- scaling length dimension by 10^9 so that 1 corresponds to 1 nm
                                        , particleMassConfig = 10e-31 -- roughly mass of an electron
                                        , oscillatorFrequencyConfig = 10000
                                        , halfWidthConfig = 10
                                        , halfHeightConfig = 10
                                          -- tune variance of wave based on mass and frequency (and Plank's constant) to get a nice simulation
                                        , initialAmplitudeConfig = gaussianWave 0 0 3.8845e-6
                                        , fpsConfig = 30
                                        , timeFactorConfig = 1
                                        }

bigBoxConfig :: OscillatorInABoxConfig
bigBoxConfig = OscillatorInABoxConfig { plankConstantConfig = 6.62607004e-16 -- scaling length dimension by 10^9 so that 1 corresponds to 1 nm
                                      , particleMassConfig = 10e-31
                                      , oscillatorFrequencyConfig = 10000
                                      , halfWidthConfig = 100
                                      , halfHeightConfig = 100
                                        -- tune variance of wave based on mass and frequency (and Plank's constant) to get a nice simulation
                                      , initialAmplitudeConfig = gaussianWave 0 0 3.8845e-6
                                      , fpsConfig = 30
                                      , timeFactorConfig = 1
                                      }

---------------------
-- Helper functions to generate lattice and translate between two dimensional position and one dimensional array
-- may not need to modify these
--------------------

positionRangeConfig :: OscillatorInABoxConfig -> [Position]
positionRangeConfig config = [Position x y | x <- [-halfWidth .. halfWidth], y <- [-halfHeight .. halfHeight]]
    where halfWidth = halfWidthConfig config
          halfHeight = halfHeightConfig config

indexRangeConfig :: OscillatorInABoxConfig -> Int
indexRangeConfig config = (2 * halfWidth + 1) * (2 * halfHeight + 1)
    where halfWidth = halfWidthConfig config
          halfHeight = halfHeightConfig config

positionToIndexConfig :: OscillatorInABoxConfig -> Position -> Int
positionToIndexConfig config pos = ((posX pos + halfWidth) * (2 * halfHeight + 1)) + (posY pos + halfHeight)
    where halfWidth = halfWidthConfig config
          halfHeight = halfHeightConfig config

indexToPositionConfig :: OscillatorInABoxConfig -> Int -> Position
indexToPositionConfig config index = Position ( (floor ((fromIntegral index) P./ (fromIntegral $! 2 * halfHeight + 1))) - halfWidth )
                                              ( (index `mod` (2 * halfHeight + 1)) - halfHeight )
    where halfWidth = halfWidthConfig config
          halfHeight = halfHeightConfig config

------------------------------------
-- simulate oscillation using Eigen library's sparse matrices
------------------------------------

sparseMatModelSimulator :: OscillatorInABoxConfig
                        -> ( TimeStep -> SparseMatOperator (Complex Double) (CComplex CDouble) Position
                           , SparseMatKet (Complex Double) (CComplex CDouble) Position
                           , SparseMatKet (Complex Double) (CComplex CDouble) Position -> Picture
                           , OscillatorInABoxConfig
                           )
sparseMatModelSimulator config = (oscillationOperator, initialModel, drawing, config)
    where
        -- shorthands for config params
        positionRange = positionRangeConfig config
        halfWidth = halfWidthConfig config
        halfHeight = halfHeightConfig config
        indexRange = indexRangeConfig config

        positionToIndex = positionToIndexConfig config
        indexToPosition = indexToPositionConfig config

        plankConstant = plankConstantConfig config
        particleMass = particleMassConfig config
        oscillatorFrequency = oscillatorFrequencyConfig config

        initialAmplitude = initialAmplitudeConfig config

        -- position and momentum operators

        posXOp :: SparseMatOperator (Complex Double) (CComplex CDouble) Position
        posXOp = SparseMatOperator $! ES.fromList indexRange indexRange $! map (\pos@(Position x _) -> (positionToIndex pos, positionToIndex pos, fromIntegral x)) positionRange

        posYOp :: SparseMatOperator (Complex Double) (CComplex CDouble) Position
        posYOp = SparseMatOperator $! ES.fromList indexRange indexRange $! map (\pos@(Position _ y) -> (positionToIndex pos, positionToIndex pos, fromIntegral y)) positionRange

        momXOp :: SparseMatOperator (Complex Double) (CComplex CDouble) Position
        -- taking dx as 2
        -- <x'| P_x |a>  =  -ih ( d/dx(a)|_x' ) = - ( a(x'+1) - a(x'-1) ) * ih/2   = - (<x'+1|a> - <x'-1|a>) * ih / 2
        momXOp = SparseMatOperator $! ES.fromList indexRange indexRange $! ( (map (\pos@(Position x y) -> ( positionToIndex pos
                                                                                , positionToIndex (Position (x P.+ 1) y)
                                                                                , 0 :+ (- plankConstant/2)
                                                                                )
                                                                                )
                                                                              -- removing max x from position range to keep operator within lattice
                                                                              -- that would give 0 contribution anyway
                                                                              [Position i j | i <- [(-halfWidth) .. (halfWidth-1)], j <- [-(halfHeight) .. halfHeight]]
                                                                             )
                                                                             ++
                                                                             (map (\pos@(Position x y) -> ( positionToIndex pos
                                                                                , positionToIndex (Position (x P.- 1) y)
                                                                                , 0 :+ (  plankConstant/2)
                                                                                )
                                                                                )
                                                                              -- removing min x from position range to keep operator within lattice
                                                                              -- that would give 0 contribution anyway
                                                                              [Position i j | i <- [-(halfWidth-1) .. (halfWidth)], j <- [-(halfHeight) .. halfHeight]]
                                                                             )
                                                                           )

        momYOp :: SparseMatOperator (Complex Double) (CComplex CDouble) Position
        -- taking dy as 2
        -- <y'| P_y |a>  =  -ih ( d/dy(a)|_y' ) = - ( a(y'+1) - a(y'-1) ) * ih/2   = - (<y'+1|a> - <y'-1|a>) * ih / 2
        momYOp = SparseMatOperator $! ES.fromList indexRange indexRange $! ( (map (\pos@(Position x y) -> ( positionToIndex pos
                                                                                , positionToIndex (Position x (y P.+ 1))
                                                                                , 0 :+ (- plankConstant/2)
                                                                                )
                                                                                )
                                                                              -- removing max y from position range to keep operator within lattice
                                                                              -- that would give 0 contribution anyway
                                                                              [Position i j | i <- [(-halfWidth) .. (halfWidth)], j <- [-(halfHeight) .. (halfHeight-1)]]
                                                                             )
                                                                             ++
                                                                             (map (\pos@(Position x y) -> ( positionToIndex pos
                                                                                , positionToIndex (Position x (y P.- 1))
                                                                                , 0 :+ (  plankConstant/2)
                                                                                )
                                                                                )
                                                                              -- removing min y from position range to keep operator within lattice
                                                                              -- that would give 0 contribution anyway
                                                                              [Position i j | i <- [(-halfWidth) .. (halfWidth)], j <- [-(halfHeight-1) .. halfHeight]]
                                                                             )
                                                                           )

        momSquaredOp = (momXOp `pow1p` 2) + (momYOp `pow1p` 2)
        posSquaredOp = (posXOp `pow1p` 2) + (posYOp `pow1p` 2)

        idOp :: SparseMatOperator (Complex Double) (CComplex CDouble) Position
        idOp = SparseMatOperator $! ES.fromList indexRange indexRange $! map (\pos@(Position _ _) -> (positionToIndex pos, positionToIndex pos, 1)) positionRange

        oscillationOperator :: Double -> SparseMatOperator (Complex Double) (CComplex CDouble) Position
        oscillationOperator dt =
              idOp + ( (0 :+ (0 P.- (dt P./ plankConstant))) .*
                       (
                         ( (((particleMass P./ 2) P.* (oscillatorFrequency P.^ 2)) :+ 0) .* posSquaredOp ) +
                         ( ((1 P./ (2 P.* particleMass)) :+ 0)                           .* momSquaredOp )
                       )
                     )

        initialModel :: SparseMatKet (Complex Double) (CComplex CDouble) Position
        initialModel = normalizeVector $! SparseMatKet $! ES.fromList indexRange 1 $! map ( \pos -> ( positionToIndex pos
                                                                                , 0
                                                                                , initialAmplitude pos
                                                                                )
                                                                                ) positionRange

        -- bitmap drawing
        drawing :: SparseMatKet (Complex Double) (CComplex CDouble) Position -> Picture
        drawing (SparseMatKet mat) =
            makePicture (2*halfWidth) (2*halfHeight)
                        1 1
                        ( \(x,y) -> let index = positionToIndex $! Position (round $! x P.* (P.fromIntegral halfWidth))
                                                                            (round $! y P.* (P.fromIntegral halfHeight))
                                        amp = mat ES.! (index, 0)
                                        den = double2Float $! normSquared amp
                                    in {-trace (show $ (index, amp, den)) $ -} rgb' den 0 0
                        )
            where
                normSquared cn = (magnitude cn) P.^ 2

        {-
        -- drawing with objects
        drawing :: SparseMatKet (Complex Double) (CComplex CDouble) Position -> Picture
        drawing (SparseMatKet mat) = {- trace "Drawing" $ -}
            pictures $ catMaybes $ map (\elem -> case elem of
                                           (index, 0, amp) -> let pos = indexToPosition index
                                                              in Just $! translate (fromIntegral $! posX pos) (fromIntegral $! posY pos) $!
                                                                 color (withAlpha (double2Float $! normSquared amp) red) $!
                                                                 circleSolid 1
                                           _ -> trace "OutOfBounds: " Nothing
                                       ) $! ES.toList $! mat
          where
                normSquared cn = (magnitude cn) P.^ 2
        -}

------------------------------------
-- simulate oscillation using sparse maps
------------------------------------

sparseModelSimulator :: OscillatorInABoxConfig
                     -> ( TimeStep -> SparseOperator (Complex Double) Position
                        , SparseKet (Complex Double) Position
                        , SparseKet (Complex Double) Position -> Picture
                        , OscillatorInABoxConfig
                        )
sparseModelSimulator config = (oscillationOperator, initialModel, drawing, config)
    where
        -- shorthands for config params
        positionRange = positionRangeConfig config
        halfWidth = halfWidthConfig config
        halfHeight = halfHeightConfig config

        plankConstant = plankConstantConfig config
        particleMass = particleMassConfig config
        oscillatorFrequency = oscillatorFrequencyConfig config

        initialAmplitude = initialAmplitudeConfig config

        -- position and momentum operators

        posXOp :: SparseOperator (Complex Double) Position
        posXOp = SparseOperator $! M.fromList $! map (\pos@(Position x _) -> ((pos,pos), ((fromIntegral x) :+ 0))) positionRange

        posYOp :: SparseOperator (Complex Double) Position
        posYOp = SparseOperator $! M.fromList $! map (\pos@(Position _ y) -> ((pos,pos), ((fromIntegral y) :+ 0))) positionRange

        momXOp :: SparseOperator (Complex Double) Position
        -- taking dx as 2
        -- <x'| P_x |a>  =  -ih ( d/dx(a)|_x' ) = - ( a(x'+1) - a(x'-1) ) * ih/2   = - (<x'+1|a> - <x'-1|a>) * ih / 2
        momXOp = SparseOperator $! M.fromList $! ( (map (\pos@(Position x y) -> ( (pos, Position (x P.+ 1) y), 0 :+ (- plankConstant/2) ))
                                                        -- removing max x from position range to keeo operator within lattice
                                                        -- that would give 0 contribution anyway
                                                        [Position i j | i <- [(-halfWidth) .. (halfWidth-1)], j <- [-(halfHeight) .. halfHeight]]
                                                   )
                                                   ++
                                                   (map (\pos@(Position x y) -> ( (pos, Position (x P.- 1) y), 0 :+ (  plankConstant/2) ))
                                                        -- removing min x from position range to keeo operator within lattice
                                                        -- that would give 0 contribution anyway
                                                        [Position i j | i <- [-(halfWidth-1) .. (halfWidth)], j <- [-(halfHeight) .. halfHeight]]
                                                   )
                                                 )

        momYOp :: SparseOperator (Complex Double) Position
        -- taking dy as 2
        -- <y'| P_y |a>  =  -ih ( d/dy(a)|_y' ) = - ( a(y'+1) - a(y'-1) ) * ih/2   = - (<y'+1|a> - <y'-1|a>) * ih / 2
        momYOp = SparseOperator $! M.fromList $! ( (map (\pos@(Position x y) -> ( (pos, Position x (y P.+ 1)), 0 :+ (- plankConstant/2) ))
                                                        -- removing max y from position range to keeo operator within lattice
                                                        -- that would give 0 contribution anyway
                                                        [Position i j | i <- [(-halfWidth) .. (halfWidth)], j <- [-(halfHeight) .. (halfHeight-1)]]
                                                   )
                                                   ++
                                                   (map (\pos@(Position x y) -> ( (pos, Position x (y P.- 1)), 0 :+ (  plankConstant/2) ))
                                                        -- removing min y from position range to keeo operator within lattice
                                                        -- that would give 0 contribution anyway
                                                        [Position i j | i <- [(-halfWidth) .. (halfWidth)], j <- [-(halfHeight-1) .. halfHeight]]
                                                   )

                                                 )

        momSquaredOp = (momXOp `pow1p` 2) + (momYOp `pow1p` 2)
        posSquaredOp = (posXOp `pow1p` 2) + (posYOp `pow1p` 2)

        idOp :: SparseOperator (Complex Double) Position
        idOp = SparseOperator $! M.fromList $! map (\pos@(Position _ _) -> ((pos,pos), 1)) positionRange

        oscillationOperator :: Double -> SparseOperator (Complex Double) Position
        oscillationOperator dt =
              idOp + ( (0 :+ (0 P.- (dt P./ plankConstant))) .*
                       (
                         ( (((particleMass P./ 2) P.* (oscillatorFrequency P.^ 2)) :+ 0) .* posSquaredOp ) +
                         ( ((1 P./ (2 P.* particleMass)) :+ 0)                           .* momSquaredOp )
                       )
                     )

        initialModel :: SparseKet (Complex Double) Position
        initialModel = normalizeVector $! SparseKet $! M.fromList $! map ( \pos -> (pos, initialAmplitude pos) ) positionRange

        -- bitmap drawing
        drawing :: SparseKet (Complex Double) Position -> Picture
        drawing (SparseKet coeffMap) =
            makePicture (2*halfWidth) (2*halfHeight)
                        1 1
                        ( \(x,y) -> let pos = Position (round $! x P.* (P.fromIntegral halfWidth))
                                                       (round $! y P.* (P.fromIntegral halfHeight))
                                        maybeAmp = M.lookup pos coeffMap
                                    in case maybeAmp of
                                           Nothing -> rgb' 0 0 0
                                           Just amp -> let den = double2Float $! normSquared amp
                                                       in {-trace (show $ (index, amp, den)) $ -} rgb' den 0 0
                        )
            where
                normSquared cn = (magnitude cn) P.^ 2

        {-
        -- drawing with objects
        drawing :: SparseKet (Complex Double) Position -> Picture
        drawing (SparseKet coeffMap) = {- trace "Drawing" $ -}
            pictures $ map (\((Position x y), amp) -> translate (fromIntegral x) (fromIntegral y) $!
                                                              color (withAlpha (double2Float $! normSquared amp) red) $!
                                                                  circleSolid 1
                           ) $! M.toList $! coeffMap
            where
                normSquared cn = (magnitude cn) P.^ 2
        -}

------------------------------------
-- Using repa arrays and matrices
-- Matrix multiplication turns out to be the bottleneck
-- Probably because it's not designed for sparse matrices
-- can work for around halfWidth=2, halfHeight=2
------------------------------------

arrayModelSimulator :: OscillatorInABoxConfig
                    -> ( TimeStep -> LatticeOperator (Complex Double) Position
                       , LatticeKet (Complex Double) Position
                       , LatticeKet (Complex Double) Position -> Picture
                       , OscillatorInABoxConfig
                       )
arrayModelSimulator config = (oscillationOperator, initialModel, drawing, config)
    where
        -- shorthands for config params
        positionRange = positionRangeConfig config
        halfWidth = halfWidthConfig config
        halfHeight = halfHeightConfig config

        indexRange = indexRangeConfig config
        positionToIndex = positionToIndexConfig config
        indexToPosition = indexToPositionConfig config

        plankConstant = plankConstantConfig config
        particleMass = particleMassConfig config
        oscillatorFrequency = oscillatorFrequencyConfig config

        initialAmplitude = initialAmplitudeConfig config

        -- position and momentum operators

        posXOp :: LatticeOperator (Complex Double) Position
        posXOp = LatticeOperator $! R.fromFunction (Z :. indexRange :. indexRange) (\(Z:.i:.j) -> if i==j
                                                                                then (fromIntegral $! posX $! indexToPosition i) :+ 0
                                                                                else 0
                                                                                )

        posYOp :: LatticeOperator (Complex Double) Position
        posYOp = LatticeOperator $! R.fromFunction (Z :. indexRange :. indexRange) (\(Z:.i:.j) -> if i==j
                                                                                then (fromIntegral $! posY $! indexToPosition i) :+ 0
                                                                                else 0
                                                                                )

        momXOp :: LatticeOperator (Complex Double) Position
        -- taking dx as 2
        -- <x'| P_x |a>  =  -ih ( d/dx(a)|_x' ) = - ( a(x'+1) - a(x'-1) ) * ih/2   = - (<x'+1|a> - <x'-1|a>) * ih / 2
        momXOp = LatticeOperator $! R.fromFunction (Z :. indexRange :. indexRange) (\(Z:.i:.j) -> let xo = posX $! indexToPosition j
                                                                                yo = posY $! indexToPosition j
                                                                                xn = posX $! indexToPosition i
                                                                                yn = posY $! indexToPosition i
                                                                                in
                                                                                case yn==yo of
                                                                                False -> 0
                                                                                True
                                                                                |xn == xo P.+ 1 -> 0 :+ (  plankConstant/2)
                                                                                |xn == xo P.- 1 -> 0 :+ (- plankConstant/2)
                                                                                |otherwise -> 0
                                                                                )

        momYOp :: LatticeOperator (Complex Double) Position
        -- taking dy as 2
        -- <y'| P_y |a>  =  -ih ( d/dy(a)|_y' ) = - ( a(y'+1) - a(y'-1) ) * ih/2   = - (<y'+1|a> - <y'-1|a>) * ih / 2
        momYOp = LatticeOperator $! R.fromFunction (Z :. indexRange :. indexRange) (\(Z:.i:.j) -> let xo = posX $! indexToPosition j
                                                                                yo = posY $! indexToPosition j
                                                                                xn = posX $! indexToPosition i
                                                                                yn = posY $! indexToPosition i
                                                                                in
                                                                                case xn==xo of
                                                                                False -> 0
                                                                                True
                                                                                |yn == yo P.+ 1 -> 0 :+ (  plankConstant/2)
                                                                                |yn == yo P.- 1 -> 0 :+ (- plankConstant/2)
                                                                                |otherwise -> 0
                                                                                )

        momSquaredOp = (momXOp `pow1p` 2) + (momYOp `pow1p` 2)
        posSquaredOp = (posXOp `pow1p` 2) + (posYOp `pow1p` 2)

        -- identity operator
        -- can make operator an instance of Unital, but don't know any one of writing the identity matrix in repa independent of dimensions
        idOp :: LatticeOperator (Complex Double) Position
        idOp = LatticeOperator $! R.fromFunction (Z :. indexRange :. indexRange) (\(Z:.i:.j) -> case i==j of
                                                                                False -> 0
                                                                                True -> 1
                                                                                )

        -- simulation operator
        oscillationOperator :: Double -> LatticeOperator (Complex Double) Position
        oscillationOperator dt =
              idOp + ( (0 :+ (0 P.- (dt P./ plankConstant))) .*
                       (
                         ( (((particleMass P./ 2) P.* (oscillatorFrequency P.^ 2)) :+ 0) .* posSquaredOp ) +
                         ( ((1 P./ (2 P.* particleMass)) :+ 0)                           .* momSquaredOp )
                       )
                     )

        -- initial ket
        initialModel :: LatticeKet (Complex Double) Position
        initialModel = LatticeKet $! R.fromFunction (Z :. indexRange) $! \(Z:.i) -> initialAmplitude $! indexToPosition i

        -- bitmap drawing
        drawing :: LatticeKet (Complex Double) Position -> Picture
        drawing (LatticeKet vector) =
            makePicture (2*halfWidth) (2*halfHeight)
                        1 1
                        ( \(x,y) -> let pos = Position (round $! x P.* (P.fromIntegral halfWidth))
                                                       (round $! y P.* (P.fromIntegral halfHeight))
                                        amp = vector R.! ( Z:.(positionToIndex pos) )
                                        den = double2Float $! normSquared amp
                                    in {-trace (show $ (pos, amp, den)) $ -} rgb' den 0 0
                        )
            where
                normSquared cn = (magnitude cn) P.^ 2

        {-
        -- drawing with objects
        drawing :: LatticeKet (Complex Double) Position -> Picture
        drawing (LatticeKet vector) = {- trace "Drawing" $ -}
          pictures $ map (\(Position x y) -> translate (fromIntegral x) (fromIntegral y) $!
                                    color (withAlpha (double2Float $ normSquared $! (vector R.! ( Z:.(positionToIndex (Position x y))))) red) $!
                                    circleSolid 1
                         ) $ positionRange
            where normSquared cn = (magnitude cn) P.^ 2
        -}

------------------------------------
-- simulate oscillation using index functions
-- doesn't work at the moment
-- no NormalizeVector instance for InfKet
------------------------------------

indexFunctionsSimulator :: OscillatorInABoxConfig
                        -> ( TimeStep -> InfOperator (Complex Double) Position
                           , InfKet (Complex Double) Position
                           , InfKet (Complex Double) Position -> Picture
                           , OscillatorInABoxConfig
                           )
indexFunctionsSimulator config = (oscillationOperator, initialModel, drawing, config)
    where
        -- shorthands for config params
        positionRange = positionRangeConfig config
        halfWidth = halfWidthConfig config
        halfHeight = halfHeightConfig config

        plankConstant = plankConstantConfig config
        particleMass = particleMassConfig config
        oscillatorFrequency = oscillatorFrequencyConfig config

        initialAmplitude = initialAmplitudeConfig config

        -- position and momentum operators
        posXOp :: InfOperator (Complex Double) Position
        posXOp = InfOperator $ \pos@(Position x _) -> M.fromList [(pos, (fromIntegral x) :+ 0) ]

        posYOp :: InfOperator (Complex Double) Position
        posYOp = InfOperator $ \pos@(Position _ y) -> M.fromList [(pos, (fromIntegral y) :+ 0)]

        momXOp :: InfOperator (Complex Double) Position
        -- taking dx as 2
        -- <x'| P_x |a>  =  -ih ( d/dx(a)|_x' ) = - ( a(x'+1) - a(x'-1) ) * ih/2   = - (<x'+1|a> - <x'-1|a>) * ih / 2
        momXOp = InfOperator $ \(Position x y) -> M.fromList [ ( Position (x P.- 1) y, (0 :+ (  plankConstant/2)) ),
                                                               ( Position (x P.+ 1) y, (0 :+ (- plankConstant/2)) )
                                                             ]

        momYOp :: InfOperator (Complex Double) Position
        -- taking dy as 2
        -- <y'| P_y |a>  =  -ih ( d/dy(a)|_y' ) = - ( a(y'+1) - a(y'-1) ) * ih/2   = - (<y'+1|a> - <y'-1|a>) * ih / 2
        momYOp = InfOperator $ \(Position x y) -> M.fromList [ ( Position x (y P.- 1), 0 :+ (  plankConstant/2) ),
                                                               ( Position x (y P.+ 1), 0 :+ (- plankConstant/2) )
                                                             ]

        momSquaredOp = (momXOp `pow1p` 2) + (momYOp `pow1p` 2)
        posSquaredOp = (posXOp `pow1p` 2) + (posYOp `pow1p` 2)

        scOp :: (Complex Double) -> InfOperator (Complex Double) Position
        scOp scalar = InfOperator $! \pos -> M.fromList [ (pos, scalar) ]

        oscillationOperator :: Double -> InfOperator (Complex Double) Position
        oscillationOperator dt = (scOp $! 1 :+ 0) + ((scOp $! 0 :+ (0 P.- (dt P./ plankConstant))) *
                                                     (
                                                       ( momSquaredOp * (scOp $ (1 P./ (2 P.* particleMass)) :+ 0) ) +
                                                       ( posSquaredOp * (scOp $ ((particleMass P./ 2) P.* oscillatorFrequency P.^ 2) :+ 0) )
                                                     )
                                                    )

        initialModel :: InfKet (Complex Double) Position
        initialModel = InfKet $! initialAmplitude

        -- bitmap drawing
        drawing :: InfKet (Complex Double) Position -> Picture
        drawing (InfKet coeffMap) = {- trace "Drawing" $ -}
            makePicture (2*halfWidth) (2*halfHeight)
                        1 1
                        ( \(x,y) -> let pos = Position (round $! x P.* (P.fromIntegral halfWidth))
                                                       (round $! y P.* (P.fromIntegral halfHeight))
                                        amp = coeffMap pos
                                        den = double2Float $! normSquared amp
                                    in {-trace (show $ (pos, amp, den)) $ -} rgb' den 0 0
                        )
            where
                normSquared cn = (magnitude cn) P.^ 2

        {-
        -- drawing with objects
        drawing :: InfKet (Complex Double) Position -> Picture
        drawing (InfKet coeffMap) = {- trace "Drawing" $ -}
            pictures $ map (\(Position x y) -> translate (fromIntegral x) (fromIntegral y) $!
                                               color (withAlpha (double2Float $! normSquared $! coeffMap (Position x y)) red) $!
                                               circleSolid 1
                           ) $ positionRange
            where normSquared cn = (magnitude cn) P.^ 2
        -}

------------------------------------------
-- Simulate
------------------------------------------

simulateOscillation :: (VectorSpace k v, SimpleOperator k v o, NormalizeVector v)
                    => (TimeStep -> o)
                    -> v
                    -> (v -> Picture)
                    -> OscillatorInABoxConfig
                    -> IO ()
simulateOscillation oscillationOperator initialModel drawing config = do
    -- return $! trace "Start Calculating Operator." ()
    -- keeping fps constant to avoid recalculation of operator
    operator <- return $! oscillationOperator $! timeFactor P./ (fromIntegral fps)
    -- return $! trace "Finish Calculating Operator." ()
    simulate (InWindow "Quantum Oscillation Simulator" (2*halfWidth, 2*halfHeight) (halfWidth, halfHeight))
             black
             fps
             initialModel
             drawing
             (simulation operator)

    where
        fps = fpsConfig config
        halfHeight = halfHeightConfig config
        halfWidth = halfWidthConfig config
        timeFactor = timeFactorConfig config

        simulation :: (VectorSpace k v, SimpleOperator k v o, NormalizeVector v)
                   => o -> ViewPort -> Float
                   -> v -> v
        simulation operator _ _ ket = {- trace "Simulating" $! -} normalizeVector $! operator `op` ket

textSimulateOscillation :: (VectorSpace k v, Show v, SimpleOperator k v o, NormalizeVector v)
                        => (TimeStep -> o)
                        -> v
                        -> OscillatorInABoxConfig
                        -> IO ()
textSimulateOscillation oscillationOperator initialModel config = do
    -- return $! trace "Start Calculating Operator." ()
    operator <- return $ oscillationOperator $! timeFactor P./ (fromIntegral fps)
    -- return $! trace "Finish Calculating Operator." ()
    go operator initialModel

    where
        fps = fpsConfig config
        timeFactor = timeFactorConfig config

        go :: ( VectorSpace k v, SimpleOperator k v o
              , Show v, NormalizeVector v
              )
            => o -> v -> IO ()
        go operator ket = do
                print ket
                threadDelay 1000000
                let nextMap = operator `op` ket
                go operator $! normalizeVector nextMap

--------------------------------
--Class Instances
--------------------------------

class NormalizeVector a where
    normalizeVector :: a -> a

instance NormalizeVector (SparseMatKet (Complex Double) (CComplex CDouble) Position) where
    normalizeVector ket@(SparseMatKet mat) = (1 / (norm :+ 0)) *^ ket
        where
            norm :: Double
            -- norm = realPart $! ES.norm mat
            -- normalizing with max of norms instead of sum of squares of norm to make the states more visible in simulation
            norm = ESD.maxCoeff $ ESD.convert normSquared $ ES.toMatrix mat

            normSquared cn = (magnitude cn) P.^ 2

instance NormalizeVector (SparseKet (Complex Double) Position) where
    normalizeVector ket@(SparseKet ketMap) = (1 / (norm :+ 0)) *^ ket
        where
            -- ketNormSquared = M.foldl' (P.+) 0 $! M.map (\coeff -> normSquared coeff) ketMap
            -- normalizing with max of norms instead of sum of squares of norm to make the states more visible in simulation
            maxNormSquared = M.foldl' (\a b -> max a b) 0 $! M.map (\coeff -> normSquared coeff) ketMap
            normSquared cn = (magnitude cn) P.^ 2
            norm :: Double
            norm = sqrt $! maxNormSquared

instance NormalizeVector (LatticeKet (Complex Double) Position) where
    normalizeVector ket@(LatticeKet array) = (1 / (norm :+ 0)) *^ ket
        where
            -- ketNormSquared = R.foldlAllS (P.+) 0 $! R.map (\coeff -> normSquared coeff) array
            -- normalizing with max of norms instead of sum of squares of norm to make the states more visible in simulation
            maxNormSquared = R.foldAllS (\a b -> max a b) 0 $! R.map (\coeff -> normSquared coeff) array
            normSquared cn = (magnitude cn) P.^ 2
            norm :: Double
            norm = sqrt $! maxNormSquared

{-
-- Will have to encode positionRange in Position to be able to traverse index map
instance NormalizeVector (InfKet (Complex Double) Position) where
    normalizeVector ket@(InfKet coeffMap) = (1 / (norm :+ 0)) *^ ket
        where
            -- ketNormSquared = foldl' (P.+) 0 $! map (normSquared .coeffMap) positionRange
            -- normalizing with max of norms instead of sum of squares of norm to make the states more visible in simulation
            maxNormSquared = foldl' (\a b -> max a b) 0 $! map (normSquared .coeffMap) positionRange
            normSquared cn = (magnitude cn) P.^ 2
            norm :: Double
            norm = sqrt $! maxNormSquared
-}

-- Hermitian Inner Product
instance (InnerProductSpace (Complex Double) (SparseKet (Complex Double) Position)) where
    (<.>) (SparseKet coeffA) (SparseKet coeffB) = M.foldl' (+) 0 $! M.intersectionWith (\ampA ampB -> ampA * (complexConjugate ampB)) coeffA coeffB

-------------------
-- For field instance of Complex Double

instance Euclidean (Complex Double) where degree _ = Just 1
instance Division (Complex Double) where (/) = (P./)
instance PID (Complex Double)
instance UFD (Complex Double)
instance GCDDomain (Complex Double)
instance IntegralDomain (Complex Double)
instance Unital (Complex Double) where one = 1 :+ 0
instance DecidableZero (Complex Double) where isZero = ((0 :+ 0) ==)
instance ZeroProductSemiring (Complex Double)
instance Commutative (Complex Double)
instance Multiplicative (Complex Double) where (*) = (P.*)
instance DecidableUnits (Complex Double) where recipUnit a = case a of
                                                (1 :+ 0) -> Just a
                                                (-1 :+ 0) -> Just a
                                                _ -> Nothing
instance UnitNormalForm (Complex Double)
instance Ring (Complex Double)
instance Monoidal (Complex Double) where zero = 0 :+ 0
instance Semiring (Complex Double)
instance Additive (Complex Double) where (+) = (P.+)
instance DecidableAssociates (Complex Double) where isAssociate = (==)
instance Rig (Complex Double)
instance Group (Complex Double) where (-) = (P.-)
instance LeftModule Natural (Complex Double) where
    n .* m = (fromIntegral $ toInteger n) P.* m
instance RightModule Natural (Complex Double) where
    m *. n = m P.* (fromIntegral n)
instance LeftModule Integer (Complex Double) where
    n .* m = (fromIntegral $ toInteger n) P.* m
instance RightModule Integer (Complex Double) where
    m *. n = m P.* (fromIntegral n)
instance Abelian (Complex Double)

--------------
-- For field instance of Double

instance Euclidean Double where degree _ = Just 1
instance Division Double where (/) = (P./)
instance PID Double
instance UFD Double
instance GCDDomain Double
instance IntegralDomain Double
instance Unital Double where one = 1
instance DecidableZero Double where isZero = (0 ==)
instance ZeroProductSemiring Double
instance Commutative Double
instance Multiplicative Double where (*) = (P.*)
instance DecidableUnits Double where recipUnit a = case a of
                                                       1 -> Just a
                                                       -1 -> Just a
                                                       _ -> Nothing
instance UnitNormalForm Double
instance Ring Double
instance Monoidal Double where zero = 0
instance Semiring Double
instance Additive Double where (+) = (P.+)
instance DecidableAssociates Double where isAssociate = (==)
instance Rig Double
instance Group Double where (-) = (P.-)
instance LeftModule Natural Double where
    n .* m = (fromIntegral $ toInteger n) P.* m
instance RightModule Natural Double where
    m *. n = m P.* (fromIntegral n)
instance LeftModule Integer Double where
    n .* m = (fromIntegral $ toInteger n) P.* m
instance RightModule Integer Double where
    m *. n = m P.* (fromIntegral n)
instance Abelian Double

---------------
-- For use by repa
instance Elt (Complex Double) where
    touch (a :+ b) = touch a >> touch b


{-# LANGUAGE ScopedTypeVariables #-}

module Math.Probably.IterLap where

import Math.Probably.Sampler
import Math.Probably.FoldingStats
import Math.Probably.MCMC (empiricalMean, empiricalCovariance)
import Math.Probably.PDF (posdefify)
import Debug.Trace

import Data.Maybe

import Numeric.LinearAlgebra hiding (find)

weightedMeanCov :: [(Vector Double, Double)] -> (Vector Double, Matrix Double)
weightedMeanCov pts = (mu, cov) where
  mu = sum $ flip map pts $ \(x,w) -> scale w x
  npts = dim $ fst $ head pts
  sumSqrWs = sum $ flip map pts $ \(x,w) -> w*w
  factor = 1/(1-sumSqrWs)
  xmeans :: [Double]
  xmeans = flip map [0..npts-1] $ \i -> mean $ flip map pts $ \(x,w) -> x@>i
  cov = scale factor $ buildMatrix npts npts $ \(j,k) ->
         sum $ flip map pts $ \(xi,wi) -> wi*(xi@>j - xmeans!!j)*
                                             (xi@>k - xmeans!!k)

mean :: [Double] -> Double
mean xs = sum xs / realToFrac (length xs)

improve :: Int -> (Vector Double -> Double)
        -> (Vector Double, Matrix Double)
        -> Sampler (Vector Double, Matrix Double)
improve n f (mu, cov) = do
  xs <- sequence $ replicate n $ multiNormal mu cov
  let xps = catMaybes $ map (\x-> let fx = f x in if isNaN fx || isInfinite fx then Nothing else Just (x,fx)) xs
      pmin = runStat (before (minFrom 1e80) snd) xps
      psubmin =  map (\(x,p)-> (x, exp $ p - pmin)) xps
      psum = sum $ map snd psubmin
      ws =  map (\(x,p)-> (x,p/psum)) psubmin
      (mu', cov') = weightedMeanCov $ ws
  return $ (mu', posdefify $ scale 2 cov')

iterateM :: Int -> (a -> Sampler a) -> a-> Sampler a
iterateM 0 _ x = return x
iterateM n f x = do
   newx <- f x
   iterateM (n-1) f newx

iterLap :: [Int]
        -> (Vector Double -> Double)
        -> (Vector Double, Matrix Double)
        -> Sampler (Vector Double, Matrix Double)
iterLap [] _ x = return x
iterLap (n:ns) f x = do
   newx <- improve n f x
   iterLap (ns) f newx


{-# LANGUAGE FlexibleContexts #-}

module Geometry where

import Graphics.Gloss as G
import Graphics.Gloss.Data.ViewPort as G
import Graphics.Gloss.Data.ViewState as G
import Data.Array.Accelerate as A
import Data.Array.Accelerate.Data.Complex as A
import qualified Prelude as P
import qualified GHC.Float as P
import qualified Data.Complex as C

type FractalPoint = Complex Double

-- TODO: phantom types

width :: Int
width = 400

height :: Int
height = 400

blWorldByWidthHeight :: (Int, Int) -> Point
blWorldByWidthHeight (w, h) = (-P.fromIntegral w P./ 2.0, -P.fromIntegral h P./ 2.0)

urWorldByWidthHeight :: (Int, Int) -> Point
urWorldByWidthHeight (w, h) = (P.fromIntegral w P./ 2.0, P.fromIntegral h P./ 2.0)

blWorldInit :: Point
blWorldInit = blWorldByWidthHeight (width, height)

urWorldInit :: Point
urWorldInit = urWorldByWidthHeight (width, height)

blComplexInit :: FractalPoint
blComplexInit = (-1) :+ (-1)

urComplexInit :: FractalPoint
urComplexInit = (1) :+ (1)

convertWorldToComplex :: Point -> FractalPoint -> Point -> FractalPoint
convertWorldToComplex w_ur c_ur (x, y) = xComplex C.:+ yComplex
                                where xRel = x P./ P.fst w_ur
                                      yRel = y P./ P.snd w_ur
                                      xComplex = P.float2Double xRel P.* C.realPart c_ur
                                      yComplex = P.float2Double yRel P.* C.imagPart c_ur

makeComplexGrid :: Int -> Int -> FractalPoint -> FractalPoint -> Acc (A.Matrix FractalPoint)
makeComplexGrid w h bl ur = use $ A.fromList (Z:.h :. w) $ do
    imag <- P.reverse $ make1DComplexGrid h (C.imagPart bl) (C.imagPart ur)
    real <- make1DComplexGrid w (C.realPart bl) (C.realPart ur)
    P.return (real :+ imag)
        where
            make1DComplexGrid n l r = (P.+l) . (P./ P.fromIntegral n) . (P.*(r P.- l)) . P.fromIntegral P.<$> [0..n - 1]


{-# LANGUAGE TypeSynonymInstances, FlexibleInstances#-}

module Geometry.Matrix
(
    -- * Functions
    getIndependent
)

where

import Numeric.LinearAlgebra.HMatrix
import qualified Data.Matrix as M
import Data.Function
-- instance Element Rational

-- getIndependent :: [[Rational]] -> [[Rational]] -> Int -> Int -> M.Matrix Rational-> [[Rational]]
-- getIndependent initial accum currRank maxRank slackMat
--     | currRank == maxRank = accum
--     | accum == [] = getIndependent (tail initial) (return $ head initial) 1 maxRank slackMat
--     | otherwise =   if newRank > currRank
--                     then getIndependent (tail initial) (candidate:accum) (currRank+1) maxRank slackMat
--                     else getIndependent (tail initial) accum currRank maxRank slackMat
--         where
--             slack = map (map fromRational) $ M.toLists slackMat
--             candidate = head initial
--             setVectors = map (1:) $ candidate:accum
--             matrix =  fromLists $ zipRev ((map (map fromRational)) setVectors :: [[Double]]) (slack)
--             newRank = rank matrix

-- | Get independet rows/columns of a non square matrix
getIndependent :: [[Rational]] -> [[Rational]] -> Int -> Int -> M.Matrix Rational-> [[Rational]]
getIndependent initial accum currRank maxRank slackMat
    | currRank == maxRank = accum
    | accum == [] = getIndependent (tail initial) (return $ head initial) 1 maxRank slackMat
    | otherwise =   if newRank > currRank
                    then getIndependent (tail initial) (candidate:accum) (currRank+1) maxRank slackMat
                    else getIndependent (tail initial ++ [candidate]) accum currRank maxRank slackMat
        where
            slack = map (map fromRational) $ M.toLists slackMat
            candidate = head initial
            setVectors = map (1:) $ candidate:accum
            matrix =  fromLists $ zipRev ((map (map fromRational)) setVectors :: [[Double]]) (slack)
            newRank = rank matrix

zipRev :: [[Double]] -> [[Double]] -> [[Double]]
zipRev a = (reverse . ((zipWith (++)) `on` reverse) a)


{-|
Module: Numeric.Morpheus.Statistics
Description: Statistics Functions
Copyright: (c) Alexander Ignatyev, 2017
License: BSD-3
Stability: experimental
Portability: POSIX
-}

module Numeric.Morpheus.Statistics
(
  mean
  , stddev_m
  , stddev
  , columnMean
  , rowMean
  , columnStddev_m
  , rowStddev_m
)

where

import Numeric.Morpheus.Utils(morpheusLayout)

import Numeric.LinearAlgebra
import Numeric.LinearAlgebra.Devel
import System.IO.Unsafe(unsafePerformIO)
import Foreign.C.Types
import Foreign.Ptr(Ptr)
import Foreign.Storable

{- morpheus_mean -}
foreign import ccall unsafe "morpheus_mean"
  c_morpheus_mean :: CInt -> Ptr R -> IO R

-- | Calculates mean (average).
mean :: Vector R -> R
mean x = unsafePerformIO $ do
    apply x id c_morpheus_mean

{- morpheus_stddev_m -}
foreign import ccall unsafe "morpheus_stddev"
  c_morpheus_stddev_m :: R -> CInt -> Ptr R -> IO R

-- | Calculates sample standard deviation using given mean value.
stddev_m :: R -> Vector R -> R
stddev_m m x = unsafePerformIO $ do
    apply x id (c_morpheus_stddev_m m)

{- morpheus_stddev -}
foreign import ccall unsafe "morpheus_stddev"
  c_morpheus_stddev :: CInt -> Ptr R -> IO R

-- | Calculates sample standard deviation.
stddev :: Vector R -> R
stddev x = unsafePerformIO $ do
    apply x id c_morpheus_stddev

{- morpheus_column_mean -}
foreign import ccall unsafe "morpheus_column_mean"
  c_morpheus_column_mean :: CInt -> CInt -> CInt -> Ptr R -> Ptr R -> IO ()

call_morpheus_column_mean :: CInt -> CInt -> CInt -> CInt -> Ptr R
                          -> CInt -> Ptr R
                          -> IO ()
call_morpheus_column_mean rows cols xRow xCol mat _ means = do
  let layout = morpheusLayout xCol cols
  c_morpheus_column_mean layout rows cols mat means

-- | Calculates mean (average) for every column.
columnMean :: Matrix R -> Vector R
columnMean m = unsafePerformIO $ do
    v <- createVector (cols m)
    apply m (apply v id) call_morpheus_column_mean
    return v

{- morpheus_row_mean -}
foreign import ccall unsafe "morpheus_row_mean"
  c_morpheus_row_mean :: CInt -> CInt -> CInt -> Ptr R -> Ptr R -> IO ()

call_morpheus_row_mean :: CInt -> CInt -> CInt -> CInt -> Ptr R
                       -> CInt -> Ptr R
                       -> IO ()
call_morpheus_row_mean rows cols xRow xCol mat _ means = do
  let layout = morpheusLayout xCol cols
  c_morpheus_row_mean layout rows cols mat means

-- | Calculates mean (average) for every row.
rowMean :: Matrix R -> Vector R
rowMean m = unsafePerformIO $ do
    v <- createVector (rows m)
    apply m (apply v id) call_morpheus_row_mean
    return v

{- morpheus_column_stddev_m -}
foreign import ccall unsafe "morpheus_column_stddev_m"
  c_morpheus_column_stddev_m :: CInt -> CInt -> CInt -> Ptr R -> Ptr R -> Ptr R -> IO ()

call_morpheus_column_stddev_m :: CInt -> Ptr R
                              -> CInt -> CInt -> CInt -> CInt -> Ptr R
                              -> CInt -> Ptr R
                              -> IO ()
call_morpheus_column_stddev_m _ means rows cols xRow xCol mat _ stddevs = do
  let layout = morpheusLayout xCol cols
  c_morpheus_column_stddev_m layout rows cols means mat stddevs

-- | Calculates sample standard deviation using given mean value for every column.
columnStddev_m :: Vector R -> Matrix R -> Vector R
columnStddev_m means m = unsafePerformIO $ do
    v <- createVector (cols m)
    apply means (apply m (apply v id)) call_morpheus_column_stddev_m
    return v

{- morpheus_row_stddev_m -}
foreign import ccall unsafe "morpheus_row_stddev_m"
  c_morpheus_row_stddev_m :: CInt -> CInt -> CInt -> Ptr R -> Ptr R -> Ptr R -> IO ()

call_morpheus_row_stddev_m :: CInt -> Ptr R
                              -> CInt -> CInt -> CInt -> CInt -> Ptr R
                              -> CInt -> Ptr R
                              -> IO ()
call_morpheus_row_stddev_m _ means rows cols xRow xCol mat _ stddevs = do
  let layout = morpheusLayout xCol cols
  c_morpheus_row_stddev_m layout rows cols means mat stddevs

-- | Calculates sample standard deviation using given mean value for every row.
rowStddev_m :: Vector R -> Matrix R -> Vector R
rowStddev_m means m = unsafePerformIO $ do
    v <- createVector (rows m)
    apply means (apply m (apply v id)) call_morpheus_row_stddev_m
    return v


{-# LANGUAGE FlexibleInstances #-}
module Main where

import Data.Complex

import Debug.Trace

data Nav = N Float
         | S Float
         | E Float
         | W Float
         | L Int
         | R Int
         | F Float
 deriving (Show)

data Change = Dir (Complex Float)
            | Rot (Complex Float)
            | Move (Complex Float)

 deriving (Show)

degToRot :: Int -> Float
degToRot d = fromIntegral (d `div` 90)*pi/2

navToChange :: Nav -> Change
navToChange (N m) = Dir  $ 0 :+ m
navToChange (S m) = Dir  $ 0 :+ (-m)
navToChange (E m) = Dir  $ m :+ 0
navToChange (W m) = Dir  $ (-m) :+ 0
navToChange (L d) = Rot  $ mkPolar 1.0 (degToRot d)
navToChange (R d) = Rot  $ mkPolar 1.0 (-(degToRot d))
navToChange (F m) = Move $ mkPolar m 0

instance Read Nav where
  readsPrec _ ('N':r) = [( N $ read r, "")]
  readsPrec _ ('S':r) = [( S $ read r, "")]
  readsPrec _ ('E':r) = [( E $ read r, "")]
  readsPrec _ ('W':r) = [( W $ read r, "")]
  readsPrec _ ('L':r) = [( L $ read r, "")]
  readsPrec _ ('R':r) = [( R $ read r, "")]
  readsPrec _ ('F':r) = [( F $ read r, "")]

getInput :: FilePath -> IO [Nav]
getInput = fmap (map read . lines) . readFile

dirToComplex :: Nav -> Complex Float
dirToComplex (N n) = 0 :+ n
dirToComplex (S s) = 0 :+ (-s)
dirToComplex (E e) = e :+ 0
dirToComplex (W w) = (-w) :+ 0
dirToComplex _ = undefined

solution :: [Nav] -> Int
solution navs = round $ abs (dx + ix) + abs (dy + iy)
  where (dirs, instr) = (filter isDir navs, filter (not . isDir) navs)
        (dx :+ dy) = sum $ map dirToComplex dirs
        (ix :+ iy) = instrsToComplex 0 (0 :+ 0) instr

        isDir :: Nav -> Bool
        isDir (N _) = True
        isDir (S _) = True
        isDir (E _) = True
        isDir (W _) = True
        isDir _ = False

        instrsToComplex :: Float -> Complex Float -> [Nav] -> Complex Float
        instrsToComplex _ curLoc [] = curLoc
        instrsToComplex curPhase curLoc ((L d):rest)
            = instrsToComplex (curPhase + degToRot d) curLoc rest
        instrsToComplex curPhase curLoc ((R d):rest)
            = instrsToComplex (curPhase - degToRot d) curLoc rest
        instrsToComplex curPhase ship ((F m):rest)
            = instrsToComplex curPhase (ship+rot) rest
            where rot = mkPolar m curPhase

solution2 :: [Nav] -> Int
solution2 navs = round $ abs ix + abs iy
  where (ix :+ iy) = run 0 (10 :+ 1) $ map navToChange navs
        run :: Complex Float -> Complex Float -> [Change] -> Complex Float
        run ship _ [] = ship
        run ship wp (Dir d:rest) = run ship (wp + d) rest
        run ship wp (Rot r:rest) = run ship (wp*r) rest
        run ship wp (Move m:rest)= run (ship + m*wp) wp rest

main :: IO ()
main = do getInput "test-input" >>= print . solution
          getInput "test-input" >>= print . solution2
          getInput "input" >>= print . solution
          getInput "input" >>= print . solution2


module STCR2S1PointSet where

import           Control.Monad           as M
import           Control.Monad.Parallel  as MP
import           Data.Array.Repa         as R
import           Data.Binary             (decodeFile)
import           Data.Complex
import           Data.List               as L
import           Data.Vector.Unboxed     as VU
import           DFT.Plan
import           Filter.Utils
import           FokkerPlanck.MonteCarlo
import           Image.IO
import           STC.CompletionFieldR2S1
import STC.Shape
import           System.Directory
import           System.Environment
import           System.FilePath
import           Text.Printf
import           Types
import           Utils.Array
import FokkerPlanck.Histogram
import FokkerPlanck.GreensFunction

main = do
  args@(numPointStr:numOrientationStr:sigmaStr:taoStr:numTrailStr:maxTrailStr:initOriStr:initSpeedStr:writeFlagStr:numIterationStr:thresholdStr:histFilePath:rStr:deltaStr:stdStr:numThreadStr:_) <-
    getArgs
  print args
  let numPoint = read numPointStr :: Int
      numOrientation = read numOrientationStr :: Int
      sigma = read sigmaStr :: Double
      tao = read taoStr :: Double
      numTrail = read numTrailStr :: Int
      maxTrail = read maxTrailStr :: Int
      initOri = read initOriStr :: Double
      initSpeed = read initSpeedStr :: Double
      writeFlag = read writeFlagStr :: Bool
      numIteration = read numIterationStr :: Int
      threshold = read thresholdStr :: Double
      r = read rStr :: Double
      delta = read deltaStr :: Double
      std = read stdStr :: Double
      numThread = read numThreadStr :: Int
      folderPath = "output/test/STCR2S1PointSet"
  createDirectoryIfMissing True folderPath
  flag <- doesFileExist histFilePath
  -- arrG' <-
  --   if flag
  --     then getNormalizedHistogramArr <$> decodeFile histFilePath
  --     else do
  --       putStrLn "Couldn't find a Green's function data. Start simulation..."
  --       solveMonteCarloR2S1
  --         numThread
  --         numTrail
  --         maxTrail
  --         numPoint
  --         numPoint
  --         numOrientation
  --         sigma
  --         tao
  --         ((sqrt 2) * fromIntegral numPoint)
  --         initSpeed
  --         histFilePath
  --         -- (0, 0, 0, 0, initOri / 180 * pi, initSpeed)
  arrG' <-
    sampleR2S1
      numPoint
      numPoint
      delta
      (2 * pi / fromIntegral numOrientation)
      initSpeed
      sigma
      tao
      [0 .. numOrientation - 1]
  let -- arrG' =
      --   sampleR2S1
      --     numPoint
      --     numPoint
      --     delta
      --     (2 * pi / fromIntegral numOrientation)
      --     initSpeed
      --     sigma
      --     tao
      --     [0 .. numOrientation - 1]
      -- arrG =
      --   computeS . R.traverse arrG' id $ \f idx@(Z :. k :. i :. j) ->
      --     let r' =
      --           sqrt . fromIntegral $
      --           (i - div numPoint 2) ^ 2 + (j - div numPoint 2) ^ 2
      --     in -- if r' > 49 -- || r' <= 2
      --        --   then 0
      --        --   else
      --       f idx
      arrG = computeS . rotate3DR $ arrG'
      -- r = 16
  let numTheta = 8
      deltaTheta = (2 * pi) / numTheta
      xs -- [R2S1RPPoint (-i, 0, 0, 1) | i <- [-24,-18 .. 24]]
       =
        ((L.map
            (\(i, j) -> R2S1RPPoint (round i, round j, 0, 1))
            [ (r * cos (k * deltaTheta) + 0, r * sin (k * deltaTheta) + 0)
            | k <- [0 .. numTheta - 1]
            ]) -- L.++
         -- [R2S1RPPoint (-3, -1, 0, 1)]
         )
      bias = computeBias numPoint numPoint numOrientation xs
      -- initialEigenVec =
      --   computeInitialEigenVec numPoint numPoint numOrientation xs
      circle = Points (0, 0) 2 Circle {circleNum = 8, circleRadius = round r}
      circlePoints = makeShape2D $ circle
      ys = getShape2DIndexListGaussian 5 2 $ circlePoints
      initialEigenVec =
        computeInitialEigenVecGaussian numPoint numPoint numOrientation ys
  print xs
  print circlePoints
  -- powerMethod
  --   emptyPlan
  --   folderPath
  --   arrG
  --   numIteration
  --   writeFlag
  --   ""
  --   threshold
  --   bias
  --   initialEigenVec
  computeContourR2S1Analytic
    folderPath
    arrG
    numOrientation
    numPoint
    numPoint
    sigma
    tao
    initSpeed
    std
    delta
    bias
    circlePoints
  -- computeContourR2S1Tangent
  --   folderPath
  --   arrG
  --   numOrientation
  --   numPoint
  --   numPoint
  --   delta
  --   circlePoints


{-# LANGUAGE FlexibleContexts #-}

module Lib.Utility.MBool where

import Numeric.LinearAlgebra

infix  4  .==., ./=., .<., .<=., .>=., .>.
infixr 3  .&&.
infixr 2  .||.

-- specialized for Int result
cond'
    :: (Element t, Ord t, Container c I, Container c t)
    => c t -> c t -> c I -> c I -> c I -> c I
cond' = cond

a .<.  b = cond' a b 1 0 0
a .<=. b = cond' a b 1 1 0
a .==. b = cond' a b 0 1 0
a ./=. b = cond' a b 1 0 1
a .>=. b = cond' a b 0 1 1
a .>.  b = cond' a b 0 0 1

a .&&. b  = step (a*b)
a .||. b  = step (a+b)
no a      = 1-a
xor a b   = a ./=. b
equiv a b = a .==. b
imp a b   = no a .||. b

taut x = minElement x == 1

minEvery a b = cond a b a a b
maxEvery a b = cond a b b b a


{-# OPTIONS_GHC -fplugin GHC.TypeLits.KnownNat.Solver #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE KindSignatures #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE StandaloneDeriving #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE TypeSynonymInstances #-}
{-# LANGUAGE FlexibleInstances #-}
module AI.Singularity.Data.Vector where
import GHC.TypeLits
import qualified Numeric.LinearAlgebra.Data as LD
import qualified Numeric.LinearAlgebra.Devel as LDev
import qualified Numeric.LinearAlgebra as LA
import Control.Lens
import Data.Proxy
import Control.Monad.ST
import System.Random
import Control.Monad
import Control.Arrow
import qualified Data.Vector.Storable as VS

import AI.Singularity.Utils
import AI.Singularity.Data.Matrix

data GeneralVector :: (* -> *) -> Nat -> * where
  Vec :: vec Float -> GeneralVector vec n
deriving instance Show (vec Float) => Show (GeneralVector vec n)

vec :: Lens' (GeneralVector vec n) (vec Float)
{-# INLINE vec #-}
vec f (Vec v) = Vec <$> f v

type Vector = GeneralVector LD.Vector

type MutVector s = GeneralVector (LDev.STVector s)

thawV :: Vector n -> ST s (MutVector s n)
thawV = Vec <$< LDev.thawVector . view vec

cmap :: (Float -> Float) -> Vector n -> Vector n
cmap = over vec . LD.cmap

foldV :: KnownNat n => (Float -> Float -> Float) -> Float -> Vector n -> Float
foldV f z = VS.foldr f z . view vec

lenV :: forall n. KnownNat n => Vector n -> Int
lenV _ = natToInt @ n

fromListV :: forall n. KnownNat n => [Float] -> Vector n
fromListV = Vec . LD.fromList . take l . cycle
  where l = natToInt @ n

fromValV :: forall n. KnownNat n => Float -> Vector n
fromValV = fromListV . repeat

toListV :: forall n. KnownNat n => Vector n -> [Float]
toListV (Vec v) = VS.toList v

consV :: forall n. (KnownNat n) => Float -> Vector n -> Vector (n+1)
consV d = Vec . VS.cons d . view vec

tailV :: forall n. (KnownNat n) => Vector (n+1) -> Vector n
tailV = Vec . VS.tail . view vec

takeV :: forall n m. (KnownNat n, KnownNat m, m <= n) => Vector n -> Vector m
takeV = Vec . VS.take m . view vec
  where m = natToInt @ n

dropV :: forall n m. (KnownNat n, KnownNat m, m <= n) => Vector n -> Vector (n - m)
dropV = Vec . VS.drop m . view vec
  where m = natToInt @ m

splitV :: forall k l n. (KnownNat k, KnownNat l, KnownNat n, (k+l) ~ n) => Vector n -> (Vector k, Vector l)
splitV = view vec >>> VS.take k &&& VS.drop k >>> Vec *** Vec
  where k = natToInt @ k

concatVec :: forall n m. (KnownNat n, KnownNat m) => Vector n -> Vector m -> Vector (n + m)
concatVec (Vec v1) (Vec v2) = Vec (v1 VS.++ v2)

(<.>) :: Vector n -> Vector n -> Float
(<.>) (Vec a) (Vec b) = (LA.<.>) a b

(#>) :: Matrix m n -> Vector m -> Vector n
(#>) (Mat m) (Vec v) = Vec (m LA.#> v)

outer :: Vector m -> Vector n -> Matrix n m
outer (Vec a) (Vec b) = Mat (a `LA.outer` b)

appendV :: (KnownNat n, KnownNat m) => Vector n -> Vector m -> Vector (n+m)
appendV (Vec v1) (Vec v2) = Vec $ v1 VS.++ v2

instance forall n. KnownNat n => Num (Vector n) where
  (+) (Vec v1) (Vec v2) = Vec (v1 + v2)
  (*) (Vec v1) (Vec v2) = Vec (v1 * v2)
  negate (Vec v1)       = Vec (LD.cmap negate v1)
  fromInteger i         = fromListV [fromIntegral i]
  abs    (Vec v1)       = Vec (LD.cmap abs v1)
  signum (Vec v1)       = Vec (LD.cmap signum v1)

instance forall n. KnownNat n => Fractional (Vector n) where
  (/) (Vec v1) (Vec v2) = Vec (v1 / v2)
  fromRational          = fromValV . fromRational

instance forall n. KnownNat n => Floating (Vector n) where
  pi    = fromValV pi
  exp   = cmap exp
  log   = cmap log
  sin   = cmap sin
  cos   = cmap cos
  sinh  = cmap sinh
  cosh  = cmap cosh
  tanh  = cmap tanh
  asinh = cmap asinh
  acosh = cmap acosh
  atanh = cmap atanh

answer :: Vector n -> Int
answer (Vec v) = fst . VS.ifoldr (\i x b -> if x > snd b then (i,x) else b) (0,0) $ v


--stolen from kmeans!

module Math.Probably.KMedoids
    where
import Data.List (transpose, sort, groupBy, minimumBy, nub)
import Data.Ord (comparing)
--import Numeric.LinearAlgebra
import qualified Data.Map.Strict as Map

import Numeric.LinearAlgebra
import Control.Monad
import Debug.Trace
import Data.Random

type DistMap = ([Int] , (Int -> Int -> Double))

kmedoids :: Int -> Int -> DistMap -> RVar (Map.Map Int [Int])
kmedoids nclust iters dist = do
  inits <- kmedoidInit nclust dist
  return $ kmedoidIter dist inits iters

kmedoidsMany :: Int -> Int -> Int -> DistMap -> RVar [Int]
kmedoidsMany nclust iters tries dist = do
  ress <- forM [1..tries] $ \i -> do
     clusters <- kmedoids nclust iters dist
     let cost = totalCost dist clusters
     return $ trace ("KMed try = "++show i++ " nclus="++show nclust++" cost= "++show cost) (clusters, cost)
  return $ Map.keys $ fst $ minimumBy (comparing snd) ress

kmedoidInit :: Int -> DistMap -> RVar [Int]
kmedoidInit nclust dist = fmap (sort . take nclust) $ shuffle $  fst dist

kmedoidIter :: DistMap -> [Int] -> Int -> Map.Map Int [Int]
kmedoidIter dists meds 0 = assignToClosestMedoid dists meds
kmedoidIter dists meds0 iter = next where
  clusters = assignToClosestMedoid dists meds0
  meds1 = sort $ nub $ map (\(_, pts)-> reassignMedoid dists pts) $ Map.toList clusters
  next = {-trace (show (meds0, totalCost dists clusters))
         $ -} if meds1 == meds0
                 then assignToClosestMedoid dists meds1
                 else kmedoidIter dists meds1 (iter-1)

totalCost :: DistMap -> Map.Map Int [Int] -> Double
totalCost dists  = sum . map ccost . Map.toList where
  ccost (med, pts) = sum $ map (getDist dists med) pts

getDist :: DistMap -> Int -> Int -> Double
getDist (_,f) from to
  | from == to = 0
  | otherwise = f from to

assignToClosestMedoid :: DistMap -> [Int] -> Map.Map Int [Int]
assignToClosestMedoid dists medoids = clusmap where
  assignmap = Map.fromList $ map (\k -> (k, findNearestMedoid dists medoids k)) $ fst dists
  clusmap = Map.fromList $ map (\med -> (med, keysWhereValIs assignmap med)) medoids

keysWhereValIs :: Eq a => Map.Map k a -> a -> [k]
keysWhereValIs mp val = Map.keys $ Map.filter (==val) mp

findNearestMedoid :: DistMap -> [Int] -> Int -> Int
findNearestMedoid dists meds k
  = fst $ minimumBy (comparing snd) $ map (\med-> (med, getDist dists k med)) meds

calcCost :: DistMap -> Int -> [Int] -> Double
calcCost dists med  = sum . map (getDist dists med)

reassignMedoid :: DistMap -> [Int] -> Int
reassignMedoid dists points
  = fst $ minimumBy (comparing snd) $ map (\pt-> (pt, calcCost dists pt points) ) points


{- | Test of learning

In this example, two networks are used : 'simple' which is the reference and 'wrong' which is a wrong start.
The goal is to use test patterns to learn the right 'simple' network from 'wrong'. Only the values are learnt.
The topology of both networks is the same.

@
simple :: (['TDV' Bool],'SBN' 'CPT')
simple = 'runBN' $ do
    a <- 'variable' \"a\" ('t' :: Bool)
    b <- 'variable' \"b\" ('t' :: Bool)
--
    'proba' a '~~' [0.4,0.6]
    'cpt' b [a] '~~' [0.8,0.2,0.2,0.8]
--
    return [a,b]
@

and 'wrong' where the probability for a is wrong.

@
wrong :: (['TDV' Bool],'SBN' 'CPT')
wrong = 'runBN' $ do
    a <- 'variable' \"a\" ('t' :: Bool)
    b <- 'variable' \"b\" ('t' :: Bool)
--
    'proba' a '~~' [0.2,0.8]
    'cpt' b [a] '~~' [0.8,0.2,0.2,0.8]
--
    return [a,b]
@

So, the first thing to do is generate test patterns. We are using the 'discreteAncestralSampler' for this. This function is
generating a sequence of graphs. We are just interested in the values. So, we get the values with 'allVertexValues'.

@
generatePatterns :: IO [[DVI]]
generatePatterns = do
    let (vars\@[a,b],exampleG) = simple
    r <- 'runSampling' 5000 0 ('discreteAncestralSampler' exampleG)
    return (map 'allVertexValues' r)
@

Once we have the data, we can try to learn the network:

@
emTest = do
  samples <- generatePatterns
  let (_,simpleG) = simple
      (_,wrongG) = wrong
  print simpleG
  'printGraphValues' simpleG
  'printGraphValues' wrongG
--
  'printGraphValues' ('learnEM' samples wrongG)
@

First, we display the topology of the graph and the values for the reference graph and the wrong one.
Then, we use the 'learnEM' function to learn a new network from the samples. And, we print the new values to check.

-}
module Bayes.Examples.EMTest(
    -- * Test function
    emTest
  ) where

import Bayes.Sampling
import System.Random.MWC.CondensedTable
import qualified Data.Vector as V
import Bayes.Factor
import Bayes
import Bayes.FactorElimination
import Data.Maybe(fromJust)
import Bayes.BayesianNetwork
import Bayes.Factor.CPT
import Statistics.Sample.Histogram
import Bayes.EM

simple :: ([TDV Bool],SBN CPT)
simple = runBN $ do
    a <- variable "a" (t :: Bool)
    b <- variable "b" (t :: Bool)

    proba a ~~ [0.4,0.6]
    cpt b [a] ~~ [0.8,0.2,0.2,0.8]

    return [a,b]

wrong :: ([TDV Bool],SBN CPT)
wrong = runBN $ do
    a <- variable "a" (t :: Bool)
    b <- variable "b" (t :: Bool)

    proba a ~~ [0.2,0.8]
    cpt b [a] ~~ [0.8,0.2,0.2,0.8]

    return [a,b]

setJust vertex v = Just v

generatePatterns :: IO [[DVI]]
generatePatterns = do
    let (vars@[a,b],exampleG) = simple
    r <- runSampling 5000 0 (discreteAncestralSampler exampleG)
    return (map allVertexValues r)

emTest = do
  samples <- generatePatterns
  let (_,simpleG) = simple
      (_,wrongG) = wrong
  print simpleG
  printGraphValues simpleG
  printGraphValues wrongG

  printGraphValues (learnEM samples wrongG)


module Raytracer.Geometry where

import Numeric.LinearAlgebra.Data (Vector, (|>), toList, (!), fromList, toList)
import Numeric.LinearAlgebra.Data (fromColumns, asColumn, toColumns)
import Data.Maybe.HT (toMaybe)
import Data.Monoid (First(First), Monoid, getFirst, mempty, mappend, mconcat)
import Numeric.LinearAlgebra (luSolve, luPacked, pinv, (<>))
import Control.Applicative ((<*>), pure)
import Control.Monad (join)
import Data.Foldable (find)
import Data.List (sort)
import Data.Function (on)
import Data.Maybe (isJust)
import Data.Bool.HT (if')

(|*) = (*).(3|>).repeat

all_of :: [(a -> Bool)] -> a -> Bool
all_of conditions value = and $ conditions <*> (pure value)

type Pos = Vector Double
type Dir = Vector Double

-- This represents a ray of the form r = ax + b
data Ray = Ray Dir Pos deriving (Show)
calc_ray (Ray a b) n1 = round_to $ (n1 |* a) + b
  where
  round_to = fromList . map (epsilon_round 6) . toList
  epsilon_round n f = (fromInteger $ round $ f * (10^n)) / (10.0^^n)

-- This computes the ray that goes from the first point to the second
rayTo :: Pos -> Pos -> Ray
p1 `rayTo` p2 = Ray (p2 - p1) p1

-- This wraps a collision with a ray and an object
-- The `a` is what the object was holding, and the double is how far down the
-- ray the collision occured (used to find the "first" collision)
data Collision a = Collision Double a deriving (Show)

dist (Collision d _) = d

instance Eq (Collision a) where
  (==) = (==) `on` dist

instance Ord (Collision a) where
  compare = compare `on` dist

class Intersectable a where
	intersection :: Ray -> a b -> Maybe (Collision b)

ray_plane_intersect :: Ray -> Plane a -> (Vector Double)
ray_plane_intersect (Ray r1 r2) (Plane a p1 p2 p3) = head.toColumns.(\m -> (luSolve.luPacked) m (asColumn (r2 - p3))) $ fromColumns [p1, p2, -1 |* r1]

-- This represents a plane of the form p = ax + by + c
data Plane a = Plane a Dir Dir Pos deriving (Show)
calc_plane (Plane _ a b c) n1 n2 = (n1 |* a) + (n2 |* b) + c

instance Intersectable Plane where
	intersection ray plane@(Plane a _ _ _) = find (all_of [not.isInfinite, (>0)] . dist) $ Just $ Collision n a
		where
		n = (ray_plane_intersect ray plane) ! 2

-- Each of these Vectors represent a point in 3-space
data Triangle a = Triangle a Pos Pos Pos deriving (Show)

instance Intersectable Triangle where
	intersection ray (Triangle a p1 p2 p3) = if' (not $ all_of [all (>0), (1>).sum] [n1,n2]) Nothing $ find (all_of [not.isInfinite, (>0)] . dist) $ Just $ Collision n a
		where
		plane = Plane a (p2 - p1) (p3 - p1) p1
		(n1:n2:n:_) = toList $ ray_plane_intersect ray plane

-- We'll call a collection of triangles a mesh...
data Mesh a = Mesh [Triangle a] deriving (Show)

-- Most of this is just proxying the list of triangles.
instance Monoid (Mesh a) where
	mempty = Mesh []
	mappend (Mesh l1) (Mesh l2) = Mesh $ mappend l1 l2

instance Intersectable Mesh where
	-- We'll return, for the mesh, the first one we hit
	intersection ray (Mesh triangles) = join $ find isJust $ sort $ fmap (intersection ray) triangles

-- This function takes a point and two vectors and generates a mesh representing that square
square a v1 v2 p1 = Mesh [Triangle a p1 (p1 + v1) (p1 + v1 + v2), Triangle a (p1 + v1 + v2) (p1 + v2) p1]

-- This function takes a point and three vectors and generates a mesh representing that cube
cube a v1 v2 v3 p1 = mconcat $ fmap (uncurry3 $ square a) [(v1, v2, p1), (v2, v3, p1), (v3, v1, p1), (neg v1, neg v2, p2), (neg v2, neg v3, p2), (neg v3, neg v1, p2)]
	where
	p2 = p1 + v1 + v2 + v3
	neg = ((-1) |*)
	uncurry3 func (a, b, c) = func a b c


module Main where

import HopfieldMat as H
import Numeric.LinearAlgebra
import Codec.Picture
import Data.Either
import Control.Applicative

pixelsAsList :: Image PixelRGB8 -> [Int]
pixelsAsList img@(Image w h _) = go <$> [0..w-1] <*> [0..h-1]
  where go y x = pixelToBinary $ pixelAt img x y
        -- black is 1, everything else is zero
        pixelToBinary (PixelRGB8 r _ _) =
          fromIntegral $ 1 - signum r

listToImage :: Int -> Int -> [Int] -> Image PixelRGB8
listToImage w h d =  generateImage go w h
  where
    -- this is a horrible idea performance wise. yolo
    go x y = let v = d !! (y * w + x)
                 pv = fromIntegral $ 255 * (1 - v)
              in PixelRGB8 pv pv pv

loadPixels :: FilePath -> IO (Vector R)
loadPixels path = do
  img <- readImage path
  let pat = fromRight [] $ pixelsAsList . convertRGB8 <$> img
  let vpat = fromList $ fromIntegral <$> pat
  return vpat

main :: IO ()
main = do
  clean <- loadPixels "data/test.png"
  noisy <- loadPixels "data/test_noisy.png"
  let ws = H.train clean (H.initialWeights (size clean))
  putStrLn . show $ H.energy ws clean
  putStrLn . show $ H.energy ws noisy
  let newImg = H.feed ws noisy
  savePngImage "data/output_test.png" . ImageRGB8
    . listToImage 25 25 $ round <$> toList newImg


{-# LANGUAGE DeriveGeneric             #-}
{-# LANGUAGE ExistentialQuantification #-}
{-# LANGUAGE FlexibleInstances         #-}
{-# LANGUAGE InstanceSigs              #-}
{-# LANGUAGE LambdaCase                #-}
{-# LANGUAGE MultiParamTypeClasses     #-}
{-# LANGUAGE PartialTypeSignatures     #-}
{-# LANGUAGE ScopedTypeVariables       #-}
{-# LANGUAGE TypeApplications          #-}
{-# LANGUAGE TypeFamilies              #-}
{-# LANGUAGE TypeInType                #-}
{-# LANGUAGE UndecidableInstances      #-}

module Learn.Neural.Layer.Recurrent.LSTM (
  LSTM
  ) where

import           Data.Kind
import           Data.Singletons.Prelude
import           GHC.Generics                   (Generic)
import           GHC.Generics.Numeric
import           GHC.TypeLits
import           Learn.Neural.Layer
import           Numeric.BLAS
import           Numeric.Backprop
import           Numeric.Backprop.Iso
import           Statistics.Distribution
import           Statistics.Distribution.Normal
import qualified Generics.SOP                   as SOP

data LSTM :: Type

instance ( BLAS b
         , KnownNat i
         , KnownNat o
         , Floating (b '[o])
         , Floating (b '[o,i])
         , Floating (b '[o,o])
         )
        => Component LSTM b '[i] '[o] where

    data CParam LSTM b '[i] '[o] =
        LP { _lpForgetInp     :: !(b '[o,i])
           , _lpForgetState   :: !(b '[o,o])
           , _lpForgetBias    :: !(b '[o])
           , _lpRememberInp   :: !(b '[o,i])
           , _lpRememberState :: !(b '[o,o])
           , _lpRememberBias  :: !(b '[o])
           , _lpCommitInp     :: !(b '[o,i])
           , _lpCommitState   :: !(b '[o,o])
           , _lpCommitBias    :: !(b '[o])
           , _lpOutInp        :: !(b '[o,i])
           , _lpOutState      :: !(b '[o,o])
           , _lpOutBias       :: !(b '[o])
           }
      deriving (Generic)
    data CState LSTM b '[i] '[o] =
        LS { _lsCellState   :: !(b '[o])
           , _lsHiddenState :: !(b '[o])
           }
      deriving (Generic)
    type CConstr LSTM b '[i] '[o] = (Num (b '[o,i]), Num (b '[o,o]), Num (b '[o]))
    data CConf   LSTM b '[i] '[o] = forall d. ContGen d => LC d

    componentOp :: forall s. (Num (b '[i]), Num (b '[o]), CConstr LSTM b '[i] '[o])
        => OpB s '[b '[i], CParam LSTM b '[i] '[o], CState LSTM b '[i] '[o]] '[b '[o], CState LSTM b '[i] '[o]]
    componentOp = bpOp . withInps $ \(x :< p :< s :< Ø) -> do
        fI :< fS :< fB :< rI :< rS :< rB :< cI :< cS :< cB :< oI :< oS :< oB :< Ø <- gTuple #<~ p
        sC :< sH :< Ø <- gTuple #<~ s
        let forget   = sum [ matVecOp .$ (fI :< x  :< Ø)
                           , matVecOp .$ (fS :< sH :< Ø)
                           , fB
                           ]
            remember = sum [ matVecOp .$ (rI :< x  :< Ø)
                           , matVecOp .$ (rS :< sH :< Ø)
                           , rB
                           ]
            commit   = sum [ matVecOp .$ (cI :< x  :< Ø)
                           , matVecOp .$ (cS :< sH :< Ø)
                           , cB
                           ]
            out      = sum [ matVecOp .$ (oI :< x  :< Ø)
                           , matVecOp .$ (oS :< sH :< Ø)
                           , oB
                           ]
        forget'   <- tmapOp logistic ~$ (forget   :< Ø)
        remember' <- tmapOp logistic ~$ (remember :< Ø)
        commit'   <- tmapOp tanh     ~$ (commit   :< Ø)
        out'      <- tmapOp logistic ~$ (out      :< Ø)
        sC'       <- tmapOp tanh ~$ ((forget' * sC + remember' * commit') :< Ø)
        finalOut  <- bindVar $ out' * sC'
        s' :< Ø   <- isoVar (from (gTuple @(CState LSTM b '[i] '[o])) . tup1)
                            (sC' :< finalOut :< Ø)
        return $ finalOut :< s' :< Ø
      where
        logistic :: Floating a => a -> a
        logistic x = 1 / (1 + exp (-x))

    defConf = LC (normalDistr 0 0.01)
    initParam = \case
      i `SCons` SNil -> \case
        so@(o `SCons` SNil) -> \(LC d) g -> do
          LP <$> genA (o `SCons` (i `SCons` SNil)) (\_ -> realToFrac <$> genContVar d g)
             <*> genA (o `SCons` (o `SCons` SNil)) (\_ -> realToFrac <$> genContVar d g)
             <*> genA so                           (\_ -> realToFrac <$> genContVar d g)
             <*> genA (o `SCons` (i `SCons` SNil)) (\_ -> realToFrac <$> genContVar d g)
             <*> genA (o `SCons` (o `SCons` SNil)) (\_ -> realToFrac <$> genContVar d g)
             <*> genA so                           (\_ -> realToFrac <$> genContVar d g)
             <*> genA (o `SCons` (i `SCons` SNil)) (\_ -> realToFrac <$> genContVar d g)
             <*> genA (o `SCons` (o `SCons` SNil)) (\_ -> realToFrac <$> genContVar d g)
             <*> genA so                           (\_ -> realToFrac <$> genContVar d g)
             <*> genA (o `SCons` (i `SCons` SNil)) (\_ -> realToFrac <$> genContVar d g)
             <*> genA (o `SCons` (o `SCons` SNil)) (\_ -> realToFrac <$> genContVar d g)
             <*> genA so                           (\_ -> realToFrac <$> genContVar d g)
        _ -> error "inaccessible"
      _ -> error "inaccessible"
    initState _ so (LC d) g =
        LS <$> genA so (\_ -> realToFrac <$> genContVar d g)
           <*> genA so (\_ -> realToFrac <$> genContVar d g)

instance ( BLAS b
         , KnownNat i
         , KnownNat o
         , Floating (b '[o])
         , Floating (b '[o,i])
         , Floating (b '[o,o])
         )
      => ComponentLayer 'Recurrent LSTM b '[i] '[o] where
    componentRunMode = RMNotFF

instance SOP.Generic (CState LSTM b '[i] '[o])
instance SOP.Generic (CParam LSTM b '[i] '[o])

instance (Num (b '[o,i]), Num (b '[o,o]), Num (b '[o])) => Num (CParam LSTM b '[i] '[o]) where
    (+)         = genericPlus
    (-)         = genericMinus
    (*)         = genericTimes
    negate      = genericNegate
    abs         = genericAbs
    signum      = genericSignum
    fromInteger = genericFromInteger

instance (Fractional (b '[o,i]), Fractional (b '[o,o]), Fractional (b '[o])) => Fractional (CParam LSTM b '[i] '[o]) where
    (/)          = genericDivide
    recip        = genericRecip
    fromRational = genericFromRational

instance (Floating (b '[o,i]), Floating (b '[o,o]), Floating (b '[o])) => Floating (CParam LSTM b '[i] '[o]) where
    pi      = genericPi
    exp     = genericExp
    (**)    = genericPower
    log     = genericLog
    logBase = genericLogBase
    sin     = genericSin
    cos     = genericCos
    tan     = genericTan
    asin    = genericAsin
    acos    = genericAcos
    atan    = genericAtan
    sinh    = genericSinh
    cosh    = genericCosh
    tanh    = genericTanh
    asinh   = genericAsinh
    acosh   = genericAcosh
    atanh   = genericAtanh

instance (Num (b '[o,i]), Num (b '[o,o]), Num (b '[o])) => Num (CState LSTM b '[i] '[o]) where
    (+)         = genericPlus
    (-)         = genericMinus
    (*)         = genericTimes
    negate      = genericNegate
    abs         = genericAbs
    signum      = genericSignum
    fromInteger = genericFromInteger

instance (Fractional (b '[o,i]), Fractional (b '[o,o]), Fractional (b '[o])) => Fractional (CState LSTM b '[i] '[o]) where
    (/)          = genericDivide
    recip        = genericRecip
    fromRational = genericFromRational

instance (Floating (b '[o,i]), Floating (b '[o,o]), Floating (b '[o])) => Floating (CState LSTM b '[i] '[o]) where
    pi      = genericPi
    exp     = genericExp
    (**)    = genericPower
    log     = genericLog
    logBase = genericLogBase
    sin     = genericSin
    cos     = genericCos
    tan     = genericTan
    asin    = genericAsin
    acos    = genericAcos
    atan    = genericAtan
    sinh    = genericSinh
    cosh    = genericCosh
    tanh    = genericTanh
    asinh   = genericAsinh
    acosh   = genericAcosh
    atanh   = genericAtanh


{-# LANGUAGE BangPatterns   #-}
{-# LANGUAGE NamedFieldPuns #-}
module Network
    ( NeuralNet (..),
      train,
      test,
      initWB,
      loadData,
      loadTestData,
    ) where

import           Codec.Compression.GZip (decompress)
import           Control.Monad
import qualified Data.Array.IO          as A
import qualified Data.ByteString        as BS
import qualified Data.ByteString.Lazy   as BL
import           Data.Int               (Int64)
import           Data.List
import           GHC.Base               (build)
import           Numeric.LinearAlgebra  hiding (build, normalize)
import           Prelude                hiding ((<>))
import           System.Directory       (getCurrentDirectory)
import qualified System.Random          as R
import           Util

type Image        = Vector Double
type Label        = Vector Double
type Bias         = Vector Double
type Weight       = Matrix Double
type Activation   = Matrix Double
type Nabla        = Matrix Double
type ImagesMat    = Matrix Double
type LabelsMat    = Matrix Double

data NeuralNet = NN
    { weights   :: ![Weight]   -- ^ List of all weights
    , biases    :: ![Bias]    -- ^ List of all biases
    , eta       :: !Double    -- ^ Learning rate
    , epochs    :: !Int       -- ^ No. of epochs
    , layers    :: !Int       -- ^ No. of layers
    , layerSize :: !Int       -- ^ Size of hidden layers
    , batchSize :: !Int       -- ^ Size of mini-batch
    , trainData :: ![(Image, Label)]
    , testData  :: ![(ImagesMat, Int)]
    } deriving (Show, Eq)

test :: NeuralNet -> IO ()
test net@NN{epochs, biases, weights, testData} = do
    let (imgs, labels) = unzip testData
    let guesses = [guess img (zip weights $ toMatrix <$> biases) | img <- imgs]
    let c = sum $ (\(a,b) -> if a == b then 1 else 0) <$> zip guesses labels
    putStrLn $ "Epoch #" ++ (show $ 30-epochs) ++ ": " ++ (show $ c/100) ++ "%"
  where
    guess :: ImagesMat -> [(Weight, Matrix Double)] -> Int
    guess i wb = maxIndex . flatten . head $ feedforward i wb

-- | Fully matrix-based approach to backpropagation.
train :: NeuralNet -> IO NeuralNet
train net@NN{weights, batchSize, epochs, trainData} = do
    shuffledData <- shuffle trainData
    let miniBatches = chunkList batchSize shuffledData
    let batches = fmap getMiniBatch miniBatches
    let !(newWs, newBs) = (recurTest batchSize (weights, biases net) batches) :: ([Weight], [Bias])
    let newNet = net { weights = newWs, biases = newBs, epochs = epochs-1 }
    test newNet
    case epochs of
        1 -> return newNet
        _ -> train newNet

recurTest :: Int -> ([Weight], [Bias]) -> [(ImagesMat, LabelsMat)] -> ([Weight], [Bias])
recurTest batchSize !wb [batch] = trainBatch batchSize wb batch
recurTest batchSize !wb (firstBatch:batches) = recurTest batchSize (trainBatch batchSize wb firstBatch) batches

trainBatch :: Int -> ([Weight], [Bias]) -> (ImagesMat, LabelsMat) -> ([Weight], [Bias])
trainBatch batchSize !(ws, bs) (mX, mY) = backprop mY ws bs $ feedforward mX (zip ws $ bTm bs batchSize)

-- | Make biases suitable for full-matrix backprop.
bTm :: [Bias] -> Int -> [Matrix Double]
bTm biases batchSize = fmap (fromColumns . replicate batchSize) biases

-- | Turn mini-batch matrices into one matrix.
getMiniBatch :: [(Image, Label)] -> (ImagesMat, LabelsMat)
getMiniBatch mBatchData = (fromColumns imgs, fromColumns labels)
  where
    (imgs, labels) = unzip mBatchData

feedforward :: ImagesMat -> [(Weight, Matrix Double)] -> [Activation]
feedforward = scanr (\(w, b) a -> sigmoid $ w <> a + b)

backprop :: LabelsMat -> [Weight] -> [Bias] -> [Activation] -> ([Weight], [Bias])
backprop y ws bs (aL:as) = (zipWithSafe (-) ws (zipWithSafe (<>) deltaList (tr <$> as)), zipWithSafe (-) bs $ meanRows <$> deltaList)
    where
        deltaList :: [Matrix Double]
        deltaList = init $ deltas [edgeDelta y aL] y as ws

deltas :: [Matrix Double]
       -> LabelsMat
       -> [Activation]
       -> [Weight]
       -> [Matrix Double]
deltas deltal _ [] [] = deltal
deltas deltal y (a:as) (wlp1:ws) = deltas (deltal++[(tr wlp1 <> last deltal) * sigmoid' a]) y as ws

edgeDelta :: LabelsMat -- ^ Correct output Y for input X
           -> Activation -- ^ Predicted ouput for input X
           -> Matrix Double
edgeDelta y aL = (aL - y) * sigmoid' aL

{-- Arithmetic Functions --}
cost :: Matrix Double -> Matrix Double -> Matrix Double
cost a y = ((a-y)**2) / 2

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

sigmoid' :: Floating a => a -> a
sigmoid' x = x * (1 - x)

normalize :: (Integral a, Floating b) => a -> b
normalize x = (fromIntegral x) / 255

gauss :: IO Double
gauss = do
    u1 <- R.randomRIO (0 :: Double, 1 :: Double)
    u2 <- R.randomRIO (0 :: Double, 1 :: Double)
    return $ (sqrt (-2*log u1)) * (cos (2*pi*u2))

{-- Data Processing Functions --}
initWB :: NeuralNet -> IO NeuralNet
initWB net@NN{batchSize = bs, layerSize = ls, layers = n} = do
    headW <- randn ls 784
    midW  <- replicateM (n-2) (randn ls ls)
    lastW <- randn 10 ls

    initB <- replicateM (n-1) (randn ls 1)
    lastB <- randn 10 1

    let weights = reverse $ headW : midW ++ [lastW]
    let biases = reverse . fmap flatten $ initB ++ [lastB]

    return net { weights = weights
               , biases  = biases }

randomMatrix :: Int -> Int -> IO (Matrix Double)
randomMatrix nrows ncols = do
    gaussList <- replicateM (ncols*nrows) gauss
    return $ (nrows><ncols) gaussList

loadData :: IO [(Vector Double, Vector Double)]
loadData = do
    trainImgs <- getData "train-images-idx3-ubyte.gz"
    trainLabels <- getData "train-labels-idx1-ubyte.gz"
    let !labels = BL.toStrict trainLabels
    let !imgs = BL.toStrict trainImgs
    let !l = [vectorizeLabel $ getLabel n labels | n <- [0..49999]] :: [Label]
    let !i = [getImage n imgs | n <- [0..49999]] :: [Image]
    return $ zip i l

loadTestData :: IO [(Matrix Double, Int)]
loadTestData = do
    testImgs <- getData "t10k-images-idx3-ubyte.gz"
    testLabels <- getData "t10k-labels-idx1-ubyte.gz"
    let !labels = BL.toStrict testLabels
    let !imgs = BL.toStrict testImgs
    let !l = [getLabel n labels | n <- [0..9999]] :: [Int]
    let !i = asColumn <$> [getImage n imgs | n <- [0..9999]] :: [Matrix Double]
    return $ zip i l

getData :: FilePath -> IO BL.ByteString
getData path = do
    currentDir <- getCurrentDirectory
    fileData <- decompress <$> BL.readFile (currentDir ++ "/data/mnist_dataset/" ++ path)
    return fileData

getImage :: Int -> BS.ByteString -> Image
getImage n imgs = fromList [normalize $ BS.index imgs (16 + n*784 + s) | s <- [0..783]]

-- | `getImage` but lazy
getImage' :: Int64 -> BL.ByteString -> Image
getImage' n imgs = fromList [normalize $ BL.index imgs (16 + n*784 + s) | s <- [0..783]]

getLabel :: Num a => Int -> BS.ByteString -> a
getLabel n labels = fromIntegral $ BS.index labels (n+8)

getGuess :: [Activation] -> Int
getGuess = maxIndex . flatten . head

vectorizeLabel :: Int -> Vector Double
vectorizeLabel l = fromList $ x ++ 1 : y
    where (x,y) = splitAt l $ replicate 9 0

{-- HELPER FUNCTIONS --}
toMatrix :: Vector Double -> Matrix Double
toMatrix = reshape 1

-- | Means of rows in matrix as a column vector
meanRows :: Matrix Double -> Vector Double
meanRows = fst . meanCov . tr

-- | zipWith, but throws error if lengths don't match
zipWithSafe :: (a -> b -> c) -> [a] -> [b] -> [c]
zipWithSafe f as bs
    | length as == length bs = zipWith f as bs
    | otherwise = error $ "zipWith: lengths don't match (" ++ (show $ length as) ++ " " ++ (show $ length bs) ++ ")"

-- | Randomly shuffle a list
--   /O(N)/
shuffle :: [a] -> IO [a]
shuffle xs = do
        ar <- newArray n xs
        forM [1..n] $ \i -> do
            j <- R.randomRIO (i,n)
            vi <- A.readArray ar i
            vj <- A.readArray ar j
            A.writeArray ar j vi
            return vj
  where
    n = length xs
    newArray :: Int -> [a] -> IO (A.IOArray Int a)
    newArray n xs =  A.newListArray (1,n) xs