{- |
Raw, unnormalized versions of the transforms in @fftw@.

Note that the forwards and backwards transforms of this module are not actually
inverses.  For example, @run idft (run dft v) /= v@ in general.

For more information on the individual transforms, see
<http://www.fftw.org/fftw3_doc/What-FFTW-Really-Computes.html>.
-}
module Numeric.FFT.Vector.Unnormalized(
                    -- * Creating and executing 'Plan's
                    run,
                    plan,
                    execute,
                    -- * Complex-to-complex transforms
                    dft,
                    idft,
                    -- * Real-to-complex transforms
                    dftR2C,
                    dftC2R,
                    -- * Real-to-real transforms
                    -- $dct_size
                    -- ** Discrete cosine transforms
                    dct1,
                    dct2,
                    dct3,
                    dct4,
                    -- ** Discrete sine transforms
                    dst1,
                    dst2,
                    dst3,
                    dst4,
                    ) where

import Numeric.FFT.Vector.Base
import Foreign
import Foreign.C
import Data.Complex

#include <fftw3.h>

-- | Whether the complex fft is forwards or backwards.
type CDirection = CInt

-- | The type of the cosine or sine transform.
type CKind = (#type fftw_r2r_kind)

foreign import ccall unsafe fftw_plan_dft_1d
    :: CInt -> Ptr (Complex Double) -> Ptr (Complex Double) -> CDirection
        -> CFlags -> IO (Ptr CPlan)

foreign import ccall unsafe fftw_plan_dft_r2c_1d
    :: CInt -> Ptr Double -> Ptr (Complex Double) -> CFlags -> IO (Ptr CPlan)

foreign import ccall unsafe fftw_plan_dft_c2r_1d
    :: CInt -> Ptr (Complex Double) -> Ptr Double -> CFlags -> IO (Ptr CPlan)

foreign import ccall unsafe fftw_plan_r2r_1d
    :: CInt -> Ptr Double -> Ptr Double -> CKind -> CFlags -> IO (Ptr CPlan)

dft1D :: CDirection -> Transform (Complex Double) (Complex Double)
dft1D d = Transform {
            inputSize = id,
            outputSize = id,
            creationSizeFromInput = id,
            makePlan = \n a b -> withPlanner . fftw_plan_dft_1d n a b d,
            normalization = const id
            }

-- | A forward discrete Fourier transform.  The output and input sizes are the same (@n@).
--
-- @y_k = sum_(j=0)^(n-1) x_j e^(-2pi i j k/n)@
dft :: Transform (Complex Double) (Complex Double)
dft = dft1D (#const FFTW_FORWARD)

-- | A backward discrete Fourier transform.  The output and input sizes are the same (@n@).
--
-- @y_k = sum_(j=0)^(n-1) x_j e^(2pi i j k/n)@
idft :: Transform (Complex Double) (Complex Double)
idft = dft1D (#const FFTW_BACKWARD)

-- | A forward discrete Fourier transform with real data.  If the input size is @n@,
-- the output size will be @n \`div\` 2 + 1@.
dftR2C :: Transform Double (Complex Double)
dftR2C = Transform {
            inputSize = id,
            outputSize = \n -> n `div` 2 + 1,
            creationSizeFromInput = id,
            makePlan = \n a b -> withPlanner . fftw_plan_dft_r2c_1d n a b,
            normalization = const id
        }

-- | A backward discrete Fourier transform which produces real data.
--
-- This 'Transform' behaves differently than the others:
--
--  - Calling @plan dftC2R n@ creates a 'Plan' whose /output/ size is @n@, and whose
--    /input/ size is @n \`div\` 2 + 1@.
--
--  - If @length v == n@, then @length (run dftC2R v) == 2*(n-1)@.
dftC2R :: Transform (Complex Double) Double
dftC2R = Transform {
            inputSize = \n -> n `div` 2 + 1,
            outputSize = id,
            creationSizeFromInput = \n -> 2 * (n-1),
            makePlan = \n a b -> withPlanner . fftw_plan_dft_c2r_1d n a b,
            normalization = const id
        }

r2rTransform :: CKind -> Transform Double Double
r2rTransform kind = Transform {
                    inputSize = id,
                    outputSize = id,
                    creationSizeFromInput = id,
                    makePlan = \n a b -> withPlanner . fftw_plan_r2r_1d n a b kind,
                    normalization = const id
                }

-- $dct_size
-- The real-even (DCT) and real-odd (DST) transforms.  The input and output sizes
-- are the same (@n@).

-- | A type-1 discrete cosine transform.
--
-- @y_k = x_0 + (-1)^k x_(n-1) + 2 sum_(j=1)^(n-2) x_j cos(pi j k\/(n-1))@
dct1 :: Transform Double Double
dct1 = r2rTransform (#const  FFTW_REDFT00)

-- | A type-2 discrete cosine transform.
--
-- @y_k = 2 sum_(j=0)^(n-1) x_j cos(pi(j+1\/2)k\/n)@
dct2 :: Transform Double Double
dct2 = r2rTransform (#const  FFTW_REDFT10)

-- | A type-3 discrete cosine transform.
--
-- @y_k = x_0 + 2 sum_(j=1)^(n-1) x_j cos(pi j(k+1\/2)\/n)@
dct3 :: Transform Double Double
dct3 = r2rTransform (#const  FFTW_REDFT01)

-- | A type-4 discrete cosine transform.
--
-- @y_k = 2 sum_(j=0)^(n-1) x_j cos(pi(j+1\/2)(k+1\/2)\/n)@
dct4 :: Transform Double Double
dct4 = r2rTransform (#const  FFTW_REDFT11)

-- | A type-1 discrete sine transform.
--
-- @y_k = 2 sum_(j=0)^(n-1) x_j sin(pi(j+1)(k+1)\/(n+1))@
dst1 :: Transform Double Double
dst1 = r2rTransform (#const  FFTW_RODFT00)

-- | A type-2 discrete sine transform.
--
-- @y_k = 2 sum_(j=0)^(n-1) x_j sin(pi(j+1\/2)(k+1)\/n)@
dst2 :: Transform Double Double
dst2 = r2rTransform (#const  FFTW_RODFT10)

-- | A type-3 discrete sine transform.
--
-- @y_k = (-1)^k x_(n-1) + 2 sum_(j=0)^(n-2) x_j sin(pi(j+1)(k+1\/2)/n)@
dst3 :: Transform Double Double
dst3 = r2rTransform (#const  FFTW_RODFT01)

-- | A type-4 discrete sine transform.
--
-- @y_k = sum_(j=0)^(n-1) x_j sin(pi(j+1\/2)(k+1\/2)\/n)@
dst4 :: Transform Double Double
dst4 = r2rTransform (#const FFTW_RODFT11)


{-# LANGUAGE GADTs
           , DataKinds
           , TypeApplications
           , ScopedTypeVariables
           , ExplicitForAll
           , AllowAmbiguousTypes
           , KindSignatures
           , TypeOperators
           , RankNTypes
           , GeneralizedNewtypeDeriving
           , StrictData #-}
{-# OPTIONS_GHC -fplugin GHC.TypeLits.KnownNat.Solver #-}

module Lib where

import Numeric.LinearAlgebra hiding (Vector, Matrix)
import qualified Numeric.LinearAlgebra as LA
import qualified Data.Vector.Storable as VS
import GHC.TypeLits
import Data.Proxy
import System.Random
import Control.Monad
import Numeric.AD

newtype Vector (n :: Nat) = Vector { _vec :: VS.Vector Float }
  deriving (Show, Floating)

instance KnownNat n => Num (Vector n) where
  (Vector v1) + (Vector v2) = Vector $ v1 + v2
  (Vector v1) - (Vector v2) = Vector $ v1 - v2
  (Vector v1) * (Vector v2) = Vector $ v1 * v2
  abs                       = vectorMap abs
  fromInteger               = vecFromList . (:[]) . fromInteger
  negate                    = vectorMap negate
  signum                    = vectorMap signum

instance KnownNat n => Fractional (Vector n) where
  (Vector v1) / (Vector v2) = Vector $ v1 / v2
  fromRational              = vecFromList . (:[]) . fromRational

natToInt :: forall n. KnownNat n => Int
natToInt = fromIntegral . natVal $ Proxy @n

vecFromList :: forall n. KnownNat n => [Float] -> Vector n
vecFromList = Vector . VS.fromList . take (natToInt @n) . cycle

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

vecTail :: forall n. KnownNat n => Vector (n+1) -> Vector n
vecTail (Vector v) = Vector $ VS.tail v

newtype Matrix (n :: Nat) (m :: Nat) = Matrix { _mat :: LA.Matrix Float }
  deriving Show

matrixWidth :: forall n m. (KnownNat n, KnownNat m) => Matrix n m -> Int
matrixWidth _ = natToInt @n

matrixHeight :: forall n m. (KnownNat n, KnownNat m) => Matrix n m -> Int
matrixHeight _ = natToInt @m

matFromList :: forall n m. (KnownNat n, KnownNat m) => [Float] -> Matrix n m
matFromList = Matrix . (h><w) . cycle
  where w = natToInt @n
        h = natToInt @m

randomMatrix :: forall n m. (KnownNat n, KnownNat m) => IO (Matrix n m)
randomMatrix = matFromList <$> replicateM (w*h) (randomRIO (0.001, 0.1))
    where w = natToInt @n
          h = natToInt @m

data Network :: Nat -> Nat -> * where
  Layer :: forall m n. (KnownNat n, KnownNat m) =>
                       Matrix (n+1) m ->
                       (forall a. Floating a => a -> a) ->
                       Network n m
  Composition :: (KnownNat n, KnownNat m, KnownNat k) =>
                 Network m k ->
                 Network n m ->
                 Network n k

conductSignal :: forall n m. (KnownNat n, KnownNat m) =>
                               Network n m ->
                               Vector n ->
                               Vector m
conductSignal (Layer mat f) = vectorMap f . mvMul mat . vectorCons 1
conductSignal (Composition network2 network1) = conductSignal network2 . conductSignal network1

vectorMap :: (Float -> Float) -> Vector n -> Vector n
vectorMap f (Vector v) = Vector $ VS.map f v

mvMul :: Matrix n m -> Vector n -> Vector m
mvMul (Matrix m) (Vector v) = Vector $ m #> v

vectorCons :: KnownNat n => Float -> Vector n -> Vector (n+1)
vectorCons n (Vector v) = Vector $ VS.cons n v

randomLayer :: forall n m. (KnownNat n, KnownNat m) => (forall a. Floating a => a -> a) -> IO (Network n m)
randomLayer f = (`Layer` f) <$> randomMatrix

scaleMatrix :: (KnownNat n, KnownNat m) => Float -> Matrix m n -> Matrix m n
scaleMatrix s (Matrix m) = Matrix $ LA.scale s m

transposeMat :: (KnownNat n, KnownNat m) => Matrix n m -> Matrix m n
transposeMat (Matrix mat) = Matrix $ LA.tr mat

outerProduct :: (KnownNat n, KnownNat m) => Vector n -> Vector m -> Matrix n m
outerProduct (Vector v1) (Vector v2) = Matrix $ v2 `outer` v1

vvMul :: (KnownNat n) => Vector n -> Vector n -> Vector n
vvMul (Vector v1) (Vector v2) = Vector $ v1 * v2

matrixSub :: (KnownNat m, KnownNat n) => Matrix n m -> Matrix n m -> Matrix n m
matrixSub (Matrix m1) (Matrix m2) = Matrix $ m1 - m2

useTrainingExample :: forall m n. (KnownNat n, KnownNat m) =>
                        Vector m ->
                        Network n m ->
                        Vector n ->
                        (Network n m, Vector n)
useTrainingExample diffLoss (Layer weights activation) input = (trainedLayer, inputGradient)
  where
    trainedLayer = Layer newWeights activation
    newWeights :: Matrix (n+1) m
    newWeights = matrixSub weights $ scaleMatrix 0.01 weightGradient
    consedInput = vectorCons 1 input
    preActivationOutput = mvMul weights consedInput
    output = vectorMap activation preActivationOutput
    diffActivation = vectorMap (diff activation) preActivationOutput
    weightGradient = outerProduct consedInput (diffLoss * diffActivation)
    inputGradient  = vecTail $ mvMul (transposeMat weights) (diffLoss * diffActivation)
useTrainingExample diffLoss (Composition network2 network1) input =
    (Composition trainedNet2 trainedNet1, inputGradient)
  where (trainedNet2,outputGrad1)   = useTrainingExample diffLoss network2 input2
        input2                      = conductSignal network1 input
        (trainedNet1,inputGradient) = useTrainingExample outputGrad1 network1 input

showExample :: (KnownNat n, KnownNat m) =>
                (Vector m -> Vector m -> Vector m) ->
                Network n m ->
                (Vector n, Vector m) ->
                Network n m
showExample lossGrad network (input,expected) = fst $ useTrainingExample diffLoss network input
  where diffLoss = lossGrad expected output
        output   = conductSignal network input

sigmoid input = 1/(1+exp(-input))
                -- (1/) . (1+) . exp . negate

epsilon :: Floating a => a
epsilon = 0.0000001

logLoss :: Floating a => [a] -> a
logLoss [expected, output] = negate $ expected * log (output + epsilon) + (1 - expected) * log (1 - output + epsilon)

gradLogLoss expected output = grad logLoss [expected, output] !! 1

-- "mnist/train-set"

getTrainSet :: IO [(Vector 784, Vector 10)]
getTrainSet =
  map (\(a,b) -> (vecFromList a, toOneHot b)) .
  map (read @([Float],Int)) .
  lines <$> readFile  "mnist/train-set"
  where
    toOneHot n = vecFromList $ map (\x -> if x == n then 1 else 0) [0..]

eval network (input,expected) = do
  putStrLn "EXPECTED:"
  print expected
  putStrLn "PREDICTED:"
  print $ conductSignal network input

trainTest k n = do
  trainSet <- take n <$> getTrainSet
  untrainedLayer1 <- randomLayer sigmoid
  untrainedLayer2 <- randomLayer @50 @10 sigmoid
  let untrainedNetwork = Composition untrainedLayer2 untrainedLayer1
      trainedNetwork = foldl (showExample gradLogLoss) untrainedNetwork . take (k*n) $ cycle trainSet
      evalSet = take 10 trainSet
  mapM_ (eval trainedNetwork) evalSet

{-

  diff (loss . activation . mvMul(weights, input))

   => diff loss (activation . mvMul (weights,input)) * diff (activation . mvMul (weights,input))

   => diff loss (output) * diff activation (mvMul (weights,input)) * diff mvMul(weights,input)
-}


module Main where

import Network.Layer
import Network.Network
import Network.Neuron
import Network.Trainer

import Numeric.LinearAlgebra
import System.Random

main :: IO ()
main = do

  g <- newStdGen
  let l   = LayerDefinition sigmoidNeuron 2 connectFully
      l'  = LayerDefinition sigmoidNeuron 2 connectFully
      l'' = LayerDefinition sigmoidNeuron 1 connectFully

  let n = createNetwork normals g [l, l', l'']

  let t = BackpropTrainer 3 quadraticCost quadraticCost'

  let dat = [(fromList [0, 1], fromList [1]),
             (fromList [1, 1], fromList [0]),
             (fromList [1, 0], fromList [1]),
             (fromList [0, 0], fromList [0])]

  let n' = trainNTimes n t online dat 10000

  putStrLn "==> XOR predictions: "
  putStr "(0,0): "
  print $ predict (fromList [0, 0]) n'
  putStr "(1,0): "
  print $ predict (fromList [1, 0]) n'
  putStr "(0,1): "
  print $ predict (fromList [0, 1]) n'
  putStr "(1,1): "
  print $ predict (fromList [1, 1]) n'

  putStrLn "\n==> Trained network: "
  print $ map layerToShowable $ layers n'

{-
  saveNetwork "xor.ann" n'

  putStrLn "==> Network saved and reloaded: "
  n'' <- loadNetwork "xor.ann" [l, l', l'']

  print $ predict (fromList [0, 0]) n''
  print $ predict (fromList [1, 0]) n''
  print $ predict (fromList [0, 1]) n''
  print $ predict (fromList [1, 1]) n''
-}


{-# LANGUAGE DataKinds           #-}
{-# LANGUAGE FlexibleContexts    #-}
{-# LANGUAGE MagicHash           #-}
{-# LANGUAGE PatternSynonyms     #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeFamilies        #-}
{-# LANGUAGE UnboxedTuples       #-}
{-# LANGUAGE ViewPatterns        #-}
module Numeric.Quaternion.Internal
    ( Quaternion (..), Quater(Quater)
    ) where

import Numeric.DataFrame.Type  (KnownBackend)
import Numeric.Matrix.Internal (Matrix)
import Numeric.PrimBytes       (PrimBytes)
import Numeric.Vector.Internal (Vector)
import Text.Read

-- | @(x,y,z,w)@ of a quaternion
pattern Quater :: Quaternion t => t -> t -> t -> t -> Quater t
pattern Quater a b c d <- (unpackQ# -> (# a, b, c, d #))
  where
    Quater = packQ
{-# COMPLETE Quater #-}

-- | Quaternion operations
class ( Floating (Quater t), Floating t, Ord t, PrimBytes t
      , KnownBackend t '[3], KnownBackend t '[4]
      , KnownBackend t '[3,3], KnownBackend t '[4, 4]
      ) => Quaternion t where
    -- | Quaternion data type. The ordering of coordinates is @(x,y,z,w)@,
    --   where @w@ is the argument, and @x y z@ are the components of a 3D vector
    data Quater t
    -- | Set the quaternion in format @(x,y,z,w)@
    packQ :: t -> t -> t -> t -> Quater t
    -- | Get the values of the quaternion in format @(x,y,z,w)@
    unpackQ# :: Quater t -> (# t, t, t, t #)
    -- | Set the quaternion from 3D axis vector and argument
    fromVecNum :: Vector t 3 -> t -> Quater t
    -- | Set the quaternion from 4D vector in format @(x,y,z,w)@
    fromVec4 :: Vector t 4 -> Quater t
    -- | Transform the quaternion to 4D vector in format @(x,y,z,w)@
    toVec4 :: Quater t -> Vector t 4
    -- | Get scalar square of the quaternion.
    --
    --   >>> realToFrac (square q) == q * conjugate q
    square :: Quater t -> t
    -- | Imaginary part of the quaternion (orientation vector)
    im :: Quater t -> Quater t
    -- | Real part of the quaternion
    re :: Quater t -> Quater t
    -- | Imaginary part of the quaternion as a 3D vector
    imVec :: Quater t -> Vector t 3
    -- | Real part of the quaternion as a scalar
    taker :: Quater t -> t
    -- | i-th component
    takei :: Quater t -> t
    -- | j-th component
    takej :: Quater t -> t
    -- | k-th component
    takek :: Quater t -> t
    -- | Conjugate quaternion (negate imaginary part)
    conjugate :: Quater t -> Quater t
    -- | Rotates and scales vector in 3D using quaternion.
    --   Let \( q = c (\cos \frac{\alpha}{2}, v \sin \frac{\alpha}{2}) \)
    --     , \( c > 0 \), \( {|v|}^2 = 1 \);
    --   then the rotation angle is \( \alpha \), and the axis of rotation is \(v\).
    --   Scaling is proportional to \( c^2 \).
    --
    --   >>> rotScale q x == q * x * (conjugate q)
    rotScale :: Quater t -> Vector t 3 -> Vector t 3
    -- | Creates a quaternion @q@ from two vectors @a@ and @b@,
    --   such that @rotScale q a == b@.
    getRotScale :: Vector t 3 -> Vector t 3 -> Quater t
    -- | Creates a rotation versor from an axis vector and an angle in radians.
    --   Result is always a unit quaternion (versor).
    --   If the argument vector is zero, then result is a real unit quaternion.
    axisRotation :: Vector t 3 -> t -> Quater t
    -- | Quaternion rotation angle \( \alpha \)
    --   (where \( q = c (\cos \frac{\alpha}{2},  v \sin \frac{\alpha}{2}) \)
    --     , \( c > 0 \), \( {|v|}^2 = 1 \)).
    --
    --   >>> q /= 0 ==> axisRotation (imVec q) (qArg q) == signum q
    qArg :: Quater t -> t
    -- | Create a quaternion from a rotation matrix.
    --   Note, that rotations of \(q\) and \(-q\) are equivalent, there result of this
    --   function may be ambiguious. Assume the sign of the result to be chosen arbitrarily.
    fromMatrix33 :: Matrix t 3 3 -> Quater t
    -- | Create a quaternion from a homogenious coordinates trasform matrix.
    --   Ignores matrix translation transform.
    --   Note, that rotations of \(q\) and \(-q\) are equivalent, there result of this
    --   function may be ambiguious. Assume the sign of the result to be chosen arbitrarily.
    fromMatrix44 :: Matrix t 4 4 -> Quater t
    -- | Create a rotation matrix from a quaternion.
    --   Note, that rotations of \(q\) and \(-q\) are equivalent, so the following property holds:
    --
    --   >>> toMatrix33 q == toMatrix33 (-q)
    toMatrix33 :: Quater t -> Matrix t 3 3
    -- | Create a homogenious coordinates trasform matrix from a quaternion.
    --   Translation of the output matrix is zero.
    --   Note, that rotations of \(q\) and \(-q\) are equivalent, so the following property holds:
    --
    --   >>> toMatrix44 q == toMatrix44 (-q)
    toMatrix44 :: Quater t -> Matrix t 4 4

instance (Show t, Quaternion t, Ord t, Num t) => Show (Quater t) where
    showsPrec p (Quater x y z w)
        = case finS of
            SEmpty -> showChar '0'
            Simple -> finF
            SParen -> showParen (p > 6) finF
      where
        (finS, finF) = go SEmpty
          [(w, Nothing), (x, Just 'i'), (y, Just 'j'), (z, Just 'k')]
        go :: ShowState -> [(t, Maybe Char)] -> (ShowState, ShowS)
        go s ((v,l):xs)
          | (s0, f0) <- showComponent s v l
          , (s', f') <- go s0 xs
            = (s', f0 . f')
        go s [] = (s, id)
        showLabel Nothing  = id
        showLabel (Just c) = showChar c
        showComponent :: ShowState -> t -> Maybe Char -> (ShowState, ShowS)
        showComponent sState val mLabel = case (sState, compare val 0) of
          (_     , EQ) -> ( sState, id )
          (SEmpty, GT) -> ( Simple, shows val . showLabel mLabel )
          (SEmpty, LT) -> ( SParen, shows val . showLabel mLabel )
          (_     , GT) -> ( SParen
                          , showString " + " . shows val . showLabel mLabel )
          (_     , LT) -> ( SParen
                          , showString " - " . shows (negate val) . showLabel mLabel )

data ShowState = SEmpty | Simple | SParen
    deriving Eq

instance (Read t, Quaternion t, Num t) => Read (Quater t) where
    readPrec     = parens $ readPrec >>= go id 0 0 0 0
      where
        go :: (t -> t) -> t -> t -> t -> t -> t -> ReadPrec (Quater t)
        go f x y z w new =
          let def = pure (Quater x y z (f new))
              withLabel EOF         = def
              withLabel (Ident "i")
                = (lexP >>= proceed (f new) y z w) <++ pure (Quater (f new) y z w)
              withLabel (Ident "j")
                = (lexP >>= proceed x (f new) z w) <++ pure (Quater x (f new) z w)
              withLabel (Ident "k")
                = (lexP >>= proceed x y (f new) w) <++ pure (Quater x y (f new) w)
              withLabel l           = proceed x y z (f new) l
          in (lexP >>= withLabel) <++ def
        proceed :: t -> t -> t -> t -> Lexeme -> ReadPrec (Quater t)
        proceed x y z w (Symbol "+") = readPrec >>= go id x y z w
        proceed x y z w (Symbol "-") = readPrec >>= go negate x y z w
        proceed x y z w EOF          = pure (Quater x y z w)
        proceed _ _ _ _ _            = pfail

    readListPrec = readListPrecDefault
    readList     = readListDefault


{-# LANGUAGE DataKinds             #-}
{-# LANGUAGE FlexibleContexts      #-}
{-# LANGUAGE GADTs                 #-}
{-# LANGUAGE KindSignatures        #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE ScopedTypeVariables   #-}
{-# LANGUAGE TypeOperators         #-}

{-# OPTIONS_GHC -fplugin GHC.TypeLits.KnownNat.Solver #-}
{-# OPTIONS_GHC -fplugin GHC.TypeLits.Normalise #-}

{-|
Module      : Numeric.LinearAlgebra.MatrixFree
Description : Represent matrices as a pair of functions for forward and adjoint.
Copyright   : (c) Ryan Orendorff, 2020
License     : BSD3
Stability   : experimental
-}
module Numeric.LinearAlgebra.MatrixFree
    ( LinearMap(..)
    , (##>)
    , (*#)
    , (+#)
    , (<#>)
    , avgRepeatedFree
    , exactDimsFree
    , eyeFree
    , removeAvg
    , fromL
    , innerDimsFree
    , toL
    , konst
    , diag
    , forward
    , adjoint
    )
where

import           Data.Proxy                     ( Proxy(..) )
import           Data.Type.Equality             ( (:~:)(Refl) )
import           GHC.TypeLits
import           Numeric.LinearAlgebra          ( Transposable(..) )
import           Numeric.LinearAlgebra.Static hiding (konst, diag)
import qualified Numeric.LinearAlgebra.Static  as LS
import qualified Data.Vector.Storable          as V

-- | A linear operator/map between two vector spaces
data LinearMap (m :: Nat) (n :: Nat) where
  LinearMap :: (KnownNat m, KnownNat n) =>
               (R n -> R m) -- ^ Forward Function
            -> (R m -> R n) -- ^ Adjoint FunctionLS
            -> LinearMap m n

-- | Extract the foward function from a `LinearMap`
forward :: LinearMap m n -> (R n -> R m)
forward (LinearMap f _) = f

-- | Extract the adjoint function from a `LinearMap`
adjoint :: LinearMap m n -> (R m -> R n)
adjoint (LinearMap _ a) = a

instance Transposable (LinearMap n m) (LinearMap m n) where
   -- We do not have different functions for conjugate transposes in the
   -- LinearMap constructor (yet?)
    tr  (LinearMap f a) = LinearMap a f
    tr' = tr

---------------------------------------------------------------------
--        Conversions between matrices and matrix free forms       --
---------------------------------------------------------------------

-- | Convert a static matrix into the matrix free data type. This merely
-- wraps the matrix in a function (eta abstraction); it does not determine
-- the function associated with a given matrix.
fromL :: (KnownNat m, KnownNat n) =>
         L m n -- ^ Static matrix to convert to matrix free data type
      -> LinearMap m n
fromL a = LinearMap (a #>) (tr a #>)

-- | Convert a matrix free linear map into a static matrix.
--
-- This methods is very inefficient as it traverses several components
-- linearly several times.
toL :: (KnownNat m, KnownNat n) => LinearMap m n -> L m n
toL (LinearMap f _) =
    matrix $ concatMap (V.toList . extract . f) (toColumns eye)

---------------------------------------------------------------------------
--                           Example LinearMaps                          --
---------------------------------------------------------------------------

-- | The identity matrix in matrix free form.
eyeFree :: (KnownNat n) => LinearMap n n
eyeFree = LinearMap id id

-- | Take the average of a vector.
avg :: forall n. (KnownNat n) => R n -> Double
avg v = v <.> 1 / n' where n' = fromIntegral . natVal $ (Proxy :: Proxy n)

-- | mean(v) * 1, where 1 is the ones vector. of length v.
avgRepeatedFree :: (KnownNat n) => LinearMap n n
avgRepeatedFree = LinearMap f f where f v = LS.konst (avg v)

-- | Remove the average value of a vector from each element of the vector.
removeAvg :: (KnownNat n) => LinearMap n n
removeAvg = eyeFree +# ((-1) *# avgRepeatedFree)

-- | The constant matrix, which sums the vector and sets every element of
-- the resulting vector to that sum.
konst :: (KnownNat m, KnownNat n) => Double -> LinearMap m n
konst k = LinearMap sumAndRepeat sumAndRepeat
  where
    sumAndRepeat :: (KnownNat p, KnownNat q) => R p -> R q
    sumAndRepeat v = LS.konst (v <.> LS.konst k)

-- | Convert a vector into a diagonal matrix.
diag :: (KnownNat n) => R n -> LinearMap n n
diag r = LinearMap f f
  where
    f v = zipWithVector (*) r v

-- TODO: Define fft in matrix free form using FFT and IFFT

---------------------------------------------------------------------
--                Operators for composing LinearMaps               --
---------------------------------------------------------------------

-- | Apply matrix free LinearMap to a vector
(##>) :: LinearMap n m -> R m -> R n
(##>) (LinearMap f _) = f

infixr 8 ##>

-- | Matrix multiply of two `LinearMap`s
(<#>) :: LinearMap m n
      -> LinearMap n p
      -> LinearMap m p
(<#>) (LinearMap f_a a_a) (LinearMap f_b a_b) =
    LinearMap (f_a . f_b) (a_b . a_a)

infixr 8 <#>

-- | Multiply a LinearMap matrix by a constant real number
(*#) :: (KnownNat m, KnownNat n) => Double -> LinearMap m n -> LinearMap m n
(*#) s (LinearMap f a) = LinearMap (\v -> LS.konst s * f v) (\v -> LS.konst s * a v)

infixr 8 *#

-- | Add two `LinearMap`s together if they have the same shape.
(+#) :: (KnownNat m, KnownNat n) =>
        LinearMap m n
     -> LinearMap m n
     -> LinearMap m n
(+#) (LinearMap f_a a_a) (LinearMap f_b a_b) =
    LinearMap (\v -> f_b v + f_a v) (\v -> a_b v + a_a v)

infixr 7 +#

-----------------------------------------------------------------------------
--                    Runtime Proofs on LinearMap shape                    --
-----------------------------------------------------------------------------

-- | Useful for constraining two dependently typed LinearMap matrices to
-- match each other in dimensions when they are unknown at compile-time.
exactDimsFree :: forall m n j k.
                 (KnownNat n, KnownNat m, KnownNat j, KnownNat k) =>
                 LinearMap m n
              -> Maybe (LinearMap j k)
exactDimsFree (LinearMap f b) = do
    Refl <- sameNat (Proxy :: Proxy m) (Proxy :: Proxy j)
    Refl <- sameNat (Proxy :: Proxy n) (Proxy :: Proxy k)
    return $ LinearMap f b

-- | Useful for constraining that two dependently typed LinearMap matrices
-- match in their inner dimensions (enabling matrix multiplication) each
-- other in dimensions when they are unknown at compile-time.
innerDimsFree :: forall m n p q. (KnownNat n, KnownNat p) =>
                 LinearMap m n   -- ^ The left LinearMap in a matrix multiply
              -> LinearMap p q   -- ^ The right LinearMap in a matrix multiply
              -> Maybe (n :~: p) -- ^ Proof that the inner dimensions are equal
innerDimsFree (LinearMap _ _) (LinearMap _ _) = do
    Refl <- sameNat (Proxy :: Proxy n) (Proxy :: Proxy p)
    return (Refl)


{-# LANGUAGE DataKinds             #-}
{-# LANGUAGE FlexibleContexts      #-}
{-# LANGUAGE FlexibleInstances     #-}
{-# LANGUAGE MagicHash             #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE PolyKinds             #-}
{-# LANGUAGE RecordWildCards       #-}
{-# LANGUAGE ScopedTypeVariables   #-}
{-# LANGUAGE StandaloneDeriving    #-}
{-# LANGUAGE TypeApplications      #-}
{-# LANGUAGE TypeFamilies          #-}
{-# LANGUAGE TypeOperators         #-}
{-# LANGUAGE UnboxedTuples         #-}
{-# LANGUAGE UndecidableInstances  #-}
module Numeric.Matrix.Internal
  ( MatrixTranspose (..)
  , SquareMatrix (..)
  , MatrixDeterminant (..)
  , MatrixInverse (..)
  , Matrix
  , HomTransform4 (..)
  , Mat22f, Mat23f, Mat24f
  , Mat32f, Mat33f, Mat34f
  , Mat42f, Mat43f, Mat44f
  , Mat22d, Mat23d, Mat24d
  , Mat32d, Mat33d, Mat34d
  , Mat42d, Mat43d, Mat44d
  , mat22, mat33, mat44
  , (%*)
  ) where

import GHC.Base
import Numeric.DataFrame.Contraction        ((%*))
import Numeric.DataFrame.Internal.PrimArray
import Numeric.DataFrame.Type
import Numeric.Dimensions
import Numeric.Scalar.Internal
import Numeric.Vector.Internal

-- | Alias for DataFrames of rank 2
type Matrix (t :: l) (n :: k) (m :: k) = DataFrame t '[n,m]

class MatrixTranspose t (n :: k) (m :: k) where
    -- | Transpose Mat
    transpose :: Matrix t n m -> Matrix t m n

class SquareMatrix t (n :: Nat) where
    -- | Mat with 1 on diagonal and 0 elsewhere
    eye :: Matrix t n n
    -- | Put the same value on the Mat diagonal, 0 otherwise
    diag :: Scalar t -> Matrix t n n
    -- | Sum of diagonal elements
    trace :: Matrix t n n -> Scalar t

class MatrixDeterminant t (n :: Nat) where
    -- | Determinant of  Mat
    det :: Matrix t n n -> Scalar t

class MatrixInverse t (n :: Nat) where
    -- | Matrix inverse
    inverse :: Matrix t n n -> Matrix t n n

-- | Operations on 4x4 transformation matrices and vectors in homogeneous coordinates.
--   All angles are specified in radians.
--
--   Note: since version 2 of @easytensor@, DataFrames and matrices are row-major.
--         A good SIMD implementation may drastically improve performance
--         of 4D vector-matrix products of the form @v %* m@, but not so much
--         for products of the form @m %* v@ (due to memory layout).
--         Thus, all operations here assume the former form to benefit more from
--         SIMD in future.
class HomTransform4 t where
    -- | Create a translation matrix from a vector.  The 4th coordinate is ignored.
    --
    --   If @p ! 3 == 1@ and @v ! 3 == 0@, then
    --
    --   > p %* translate4 v == p + v
    --
    translate4  :: Vector t 4 -> Matrix t 4 4
    -- | Create a translation matrix from a vector.
    --
    --   If @p ! 3 == 1@, then
    --
    --   > p %* translate3 v == p + toHomVector v
    --
    translate3  :: Vector t 3 -> Matrix t 4 4
    -- | Rotation matrix for a rotation around the X axis, angle is given in radians.
    --   e.g. @p %* rotateX (pi/2)@ rotates point @p@ around @Ox@ by 90 degrees.
    rotateX     :: t -> Matrix t 4 4
    -- | Rotation matrix for a rotation around the Y axis, angle is given in radians.
    --   e.g. @p %* rotateY (pi/2)@ rotates point @p@ around @Oy@ by 90 degrees.
    rotateY     :: t -> Matrix t 4 4
    -- | Rotation matrix for a rotation around the Z axis, angle is given in radians.
    --   e.g. @p %* rotateZ (pi/2)@ rotates point @p@ around @Oz@ by 90 degrees.
    rotateZ     :: t -> Matrix t 4 4
    -- | Rotation matrix for a rotation around an arbitrary normalized vector
    --   e.g. @p %* rotate (pi/2) v@ rotates point @p@ around @v@ by 90 degrees.
    rotate      :: Vector t 3 -> t -> Matrix t 4 4
    -- | Rotation matrix from the Euler angles roll (axis @Z@), yaw (axis @Y'@), and pitch (axis @X''@).
    --   This order is known as Tait-Bryan angles (@Z-Y'-X''@ intrinsic rotations), or nautical angles, or Cardan angles.
    --
    --   > rotateEuler pitch yaw roll == rotateZ roll %* rotateY yaw %* rotateX pitch
    --
    --   https://en.wikipedia.org/wiki/Euler_angles#Conventions_2
    rotateEuler :: t -- ^ pitch (axis @X''@)
                -> t -- ^ yaw (axis @Y'@)
                -> t -- ^ roll (axis @Z@)
                -> Matrix t 4 4
    -- | Create a transform matrix using up direction, camera position and a point to look at.
    --   Just the same as GluLookAt.
    lookAt      :: Vector t 3 -- ^ The up direction, not necessary unit length or perpendicular to the view vector
                -> Vector t 3 -- ^ The viewers position
                -> Vector t 3 -- ^ The point to look at
                -> Matrix t 4 4
    -- | A perspective symmetric projection matrix. Right-handed coordinate system. (@x@ - right, @y@ - top)
    --   http://en.wikibooks.org/wiki/GLSL_Programming/Vertex_Transformations
    perspective :: t -- ^ Near plane clipping distance (always positive)
                -> t -- ^ Far plane clipping distance (always positive)
                -> t -- ^ Field of view of the y axis, in radians
                -> t -- ^ Aspect ratio, i.e. screen's width\/height
                -> Matrix t 4 4
    -- | An orthogonal symmetric projection matrix. Right-handed coordinate system. (@x@ - right, @y@ - top)
    --   http://en.wikibooks.org/wiki/GLSL_Programming/Vertex_Transformations
    orthogonal  :: t -- ^ Near plane clipping distance
                -> t -- ^ Far plane clipping distance
                -> t -- ^ width
                -> t -- ^ height
                -> Matrix t 4 4
    -- | Add one more dimension and set it to 1.
    toHomPoint  :: Vector t 3 -> Vector t 4
    -- | Add one more dimension and set it to 0.
    toHomVector :: Vector t 3 -> Vector t 4
    -- | Transform a homogenous vector or point into a normal 3D vector.
    --   If the last coordinate is not zero, divide the rest by it.
    fromHom     :: Vector t 4 -> Vector t 3

-- Type abbreviations

type Mat22f = Matrix Float 2 2
type Mat32f = Matrix Float 3 2
type Mat42f = Matrix Float 4 2
type Mat23f = Matrix Float 2 3
type Mat33f = Matrix Float 3 3
type Mat43f = Matrix Float 4 3
type Mat24f = Matrix Float 2 4
type Mat34f = Matrix Float 3 4
type Mat44f = Matrix Float 4 4

type Mat22d = Matrix Double 2 2
type Mat32d = Matrix Double 3 2
type Mat42d = Matrix Double 4 2
type Mat23d = Matrix Double 2 3
type Mat33d = Matrix Double 3 3
type Mat43d = Matrix Double 4 3
type Mat24d = Matrix Double 2 4
type Mat34d = Matrix Double 3 4
type Mat44d = Matrix Double 4 4

-- | Compose a 2x2D matrix
mat22 :: PrimBytes (t :: Type)
      => Vector t 2 -> Vector t 2 -> Matrix t 2 2
mat22 = packDF @_ @2 @'[2]
{-# INLINE mat22 #-}

-- | Compose a 3x3D matrix
mat33 :: PrimBytes (t :: Type)
      => Vector t 3 -> Vector t 3 -> Vector t 3 -> Matrix t 3 3
mat33 = packDF @_ @3 @'[3]
{-# INLINE mat33 #-}

-- | Compose a 4x4D matrix
mat44 :: PrimBytes (t :: Type)
      => Vector t 4
      -> Vector t 4
      -> Vector t 4
      -> Vector t 4
      -> Matrix t 4 4
mat44 = packDF @_ @4 @'[4]
{-# INLINE mat44 #-}

instance ( KnownDim n, KnownDim m
         , PrimArray t (Matrix t n m)
         , PrimArray t (Matrix t m n)
         ) => MatrixTranspose t (n :: Nat) (m :: Nat) where
    transpose df = case uniqueOrCumulDims df of
      Left a -> broadcast a
      Right _
         | wm <- dimVal' @m
         , wn <- dimVal' @n
         , m <- case wm of W# w -> word2Int# w
         , n <- case wn of W# w -> word2Int# w
         -> let f ( I# i,  I# j )
                  | isTrue# (i ==# n) = f ( 0, I# (j +# 1#) )
                  | otherwise         = (# ( I# (i +# 1#), I# j )
                                         , ix# (i *# m +# j) df
                                         #)
            in case gen# (CumulDims [wm*wn, wn, 1]) f (0,0) of (# _, r #) -> r

instance MatrixTranspose (t :: Type) (xn :: XNat) (xm :: XNat) where
    transpose (XFrame (df :: DataFrame t ns))
      | ((D :: Dim n) :* (D :: Dim m) :* U) <- dims @ns
      , Dict <- inferPrimElem df
      = XFrame (transpose df :: Matrix t m n)
    transpose _ = error "MatrixTranspose/transpose: impossible argument"

instance (KnownDim n, PrimArray t (Matrix t n n), Num t)
      => SquareMatrix t n where
    eye
      | n <- dimVal' @n
      = let f 0 = (# n    , 1 #)
            f k = (# k - 1, 0 #)
        in case gen# (CumulDims [n*n, n, 1]) f 0 of
            (# _, r #) -> r
    diag se
      | n <- dimVal' @n
      , e <- unScalar se
      = let f 0 = (# n    , e #)
            f k = (# k - 1, 0 #)
        in case gen# (CumulDims [n*n, n, 1]) f 0 of
            (# _, r #) -> r
    trace df
      | I# n <- fromIntegral $ dimVal' @n
      , n1 <- n +# 1#
      = let f 0# = ix# 0# df
            f k  = ix# k  df + f (k -# n1)
        in scalar $ f (n *# n -# 1#)


{-# 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]


{-|
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


{-# 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