module Gates where

import           Protolude

import           Control.Exception (assert)
import           Data.Complex      (Complex (..), magnitude)
import           Data.List         ((!!))
import qualified Data.Matrix       as M
import qualified Data.Vector       as V
import           System.Random     (randomIO)

import           Instruction       (Operation, Qubit, Qubits)
import           State             (QuantumState)

-- Dimension in qubits by taking the log base 2
dimensionQubits :: Integral a => a -> Int
dimensionQubits = ceiling . logBase 2 . fromIntegral

i :: Operation
i = M.identity 2

x :: Operation
x = M.fromLists [[0, 1], [1, 0]]

h :: Operation
h = M.scaleMatrix (1 / sqrt 2) $ M.fromLists [[1, 1], [1, -1]]

cnot :: Operation
cnot = M.fromLists [[1, 0, 0, 0], [0, 1, 0, 0], [0, 0, 0, 1], [0, 0, 1, 0]]

swap :: Operation
swap = M.fromLists [[1, 0, 0, 0], [0, 0, 1, 0], [0, 1, 0, 0], [0, 0, 0, 1]]

-- Apply an operation to a quantum state by matrix multiplication
applyOperation :: Operation -> QuantumState -> QuantumState
applyOperation g s = M.getCol 1 $ g * M.colVector s

-- Compose operators by doing matrix multiplication
-- NOTE: (*) operator chosen over multStd for possible speed improvements for larger matrices
composeOperations :: Operation -> Operation -> Operation
composeOperations = (*)

-- Sample quantum state, randomly selecting a state, weighted by |phi_i|^2
sample :: MonadIO m => QuantumState -> m Int
sample s = do
  r <- liftIO $ randomIO
  return $ sampleIter r 0
  where
    sampleIter :: Double -> Int -> Int
    sampleIter r t =
      let newR = r - (magnitude (s V.! t)) ** 2
       in if newR < 0
            then t
            else sampleIter newR (t + 1)

-- Collapse the quantum state into the state at the given index
collapse :: QuantumState -> Int -> QuantumState
collapse s i =
  V.generate
    (V.length s)
    (\x ->
       if x == i
         then (1 :+ 0)
         else (0 :+ 0))

-- Observe the quantum state and collapse into that state
observe :: MonadIO m => QuantumState -> m QuantumState
observe s = do
  i <- sample s
  return $ collapse s i

-- Oh, to have dependent types
kroneckerMultiply :: (Num a) => M.Matrix a -> M.Matrix a -> M.Matrix a
kroneckerMultiply a b = M.flatten $ M.mapPos (\_ -> flip M.scaleMatrix b) a

-- TODO: re-eval constraints
kroneckerExpt :: (Num a) => M.Matrix a -> Int -> M.Matrix a
kroneckerExpt a n
  | n < 1 = M.fromLists [[1]]
  | n == 1 = a
  | otherwise = kroneckerMultiply (kroneckerExpt a (n - 1)) a

-- Lift an operation by padding with identity
liftOperation :: Operation -> Qubit -> Int -> Operation
liftOperation a index n =
  assert (index < n) kroneckerMultiply left (kroneckerMultiply a right)
  where
    left = kroneckerExpt i (n - index - (dimensionQubits $ M.nrows a))
    right = kroneckerExpt i index

-- Calculate needed transpositions from given permutation
permutationToTranspositions :: [Int] -> [(Int, Int)]
permutationToTranspositions p = permutationToTranspositionsIter 0
  where
    permutationToTranspositionsIter dest
      | dest == (length p) = []
      | otherwise =
        (case compare src dest of
           LT -> [(src, dest)]
           GT -> [(dest, src)]
           _  -> []) ++
        permutationToTranspositionsIter (dest + 1)
      where
        src = nextSrc (p !! dest) dest
    nextSrc s d
      | s < d = nextSrc (p !! s) d
      | otherwise = s

-- Convert transpositions to list of adjacent transpositions
transpositionToAdjacentTranspositions :: [(Int, Int)] -> [Int]
transpositionToAdjacentTranspositions ts = foldMap expandCons ts
  where
    expandCons (x, y)
      | y - x == 1 = [x]
      | otherwise =
        let trans = [x .. y - 1]
         in trans ++ (reverse (initSafe trans))

-- Apply 1 qubit gate
apply1QGate :: QuantumState -> Operation -> Qubit -> QuantumState
apply1QGate s a q =
  applyOperation (liftOperation a q (dimensionQubits (length s))) s

-- Apply multiple qubit gate
applynQGate :: QuantumState -> Operation -> Qubits -> QuantumState
applynQGate s a qubits = do
  applyOperation upq s
  where
    n :: Int
    n = dimensionQubits (length s)
    transpositionsToOperation :: [Int] -> Operation
    transpositionsToOperation ts =
      foldr composeOperations (M.identity (2 ^ n)) $
      map (\x -> liftOperation Gates.swap x n) ts
    u01 :: Operation
    u01 = liftOperation a 0 n
    fromSpace :: [Int]
    fromSpace =
      (reverse qubits) ++ (filter (\x -> not (elem x qubits)) [0 .. n - 1])
    trans :: [Int]
    trans =
      transpositionToAdjacentTranspositions . permutationToTranspositions $
      fromSpace
    toFrom :: Operation
    toFrom = transpositionsToOperation trans
    fromTo :: Operation
    fromTo = transpositionsToOperation (reverse trans)
    upq :: Operation
    upq = composeOperations toFrom (composeOperations u01 fromTo)

-- Apply gate, picking apply1QGate or applynQGate based on size
applyGate :: QuantumState -> Operation -> [Int] -> QuantumState
applyGate s a qubits =
  assert
    (length qubits == dimensionQubits (M.nrows a))
    (if length qubits == 1
       then apply1QGate s a (qubits !! 0)
       else applynQGate s a qubits)


{-
  An attempt to implement a basic case of Cooley-Tukey FFT algorithm.
 -}
{-# LANGUAGE ViewPatterns #-}
module Main
  ( main
  ) where

import Control.DeepSeq
import Control.Monad
import Control.Monad.State.Strict
import Data.Bits
import Data.Complex
import Data.Time.Clock
import Statistics.Sample
import System.Random
import Text.Printf

import qualified Data.Vector as V
import qualified Data.Vector.Mutable as VM

type Cpx = Complex Double

-- e^((c pi i) * (k / N))
expAuxCommon :: Int -> Int -> Int -> Cpx
expAuxCommon c k n = cis theta
  where
    theta = pi * fromIntegral (c * k) / fromIntegral n

-- Computes e^((-2 pi i) * (k / N))
expFft :: Int -> Int -> Cpx
expFft = expAuxCommon (-2)

-- Computes e^((2 pi i) * (k / N))
expIfft :: Int -> Int -> Cpx
expIfft = expAuxCommon 2

splitByParity :: V.Vector Cpx -> (V.Vector Cpx, V.Vector Cpx)
splitByParity vs = (evens, odds)
  where
    l = V.length vs
    oddCount = (l + 1) `quot` 2
    evenCount = l - oddCount
    evens =
      V.fromListN evenCount $ (vs V.!) <$> [0,2..]
    odds =
      V.fromListN oddCount $ (vs V.!) <$> [1,3..]

-- TODO: For now we assume that vector length is a power of two,
-- this can be easily extended to any length, by pretending out-of-range
-- values are 0.
gDitFft :: (Int -> Int -> Cpx) -> V.Vector Cpx -> V.Vector Cpx
gDitFft expF = fix $
  \impl vs ->
    let l = V.length vs
    in if l <= 1
      then vs
      else V.create $ do
        let (impl -> es, impl -> os) = splitByParity vs
            hf = l `quot` 2
        vec <- VM.unsafeNew l
        forM_ [0 .. hf - 1] $ \k -> do
          let a = es V.! k
              b = expF k l * os V.! k
          VM.unsafeWrite vec k (a + b)
          VM.unsafeWrite vec (hf + k) (a - b)
        pure vec

ditFft :: V.Vector Cpx -> V.Vector Cpx
ditFft = gDitFft expFft . rightPadZeros

iditFft :: V.Vector Cpx -> V.Vector Cpx
iditFft (rightPadZeros -> vs) = V.map (/ fromIntegral l) (iditFftAux vs)
  where
    l = V.length vs
    iditFftAux :: V.Vector Cpx -> V.Vector Cpx
    iditFftAux = gDitFft expIfft

{-
  right padding zeros in the end of a matrix to make the length of the vector
  a power of 2.

  special case: returns empty vector when the input is empty.
 -}
rightPadZeros :: V.Vector Cpx -> V.Vector Cpx
rightPadZeros vs
    | l <= 1 = vs
    | isPowerOf2 = vs
    | otherwise =
        let closestPowOf2Gt =
              {-
                Find the closest power of two that is greater or equal to input v.
                only works when v > 1 and v is not a power of two.
               -}
              unsafeShiftL 1 (finiteBitSize l - countLeadingZeros l)
        in vs <> V.fromListN (closestPowOf2Gt - l) (repeat 0)
  where
    l = V.length vs
    isPowerOf2 =
      countLeadingZeros l + countTrailingZeros l + 1 == finiteBitSize l

evaluateOnVector :: V.Vector Cpx -> IO ()
evaluateOnVector vs = do
  let result = iditFft (ditFft vs)
      diffs :: V.Vector Double
      diffs = V.zipWith (\x y -> magnitude (x - y)) vs result
  printf "StdDev: %.9f\n" (stdDev diffs)

type M = StateT StdGen IO

genRandomR :: Random a => (a, a) -> M a
genRandomR = state . randomR

genCpxVector :: M (V.Vector Cpx)
genCpxVector =
  genRandomR (1024, 8192) >>= genCpxVectorOfSize

genCpxVectorOfSize :: Int -> M (V.Vector Cpx)
genCpxVectorOfSize l = do
  let rDouble = genRandomR (-100,100)
  xs <- replicateM l $ (:+) <$> rDouble <*> rDouble
  pure $ V.fromListN l xs

evalRandomVector :: M ()
evalRandomVector = do
  vs <- genCpxVector
  liftIO $ do
    printf "Vector length: %d\n" (V.length vs)
    evaluateOnVector vs

directConvolve :: V.Vector Cpx -> V.Vector Cpx -> V.Vector Cpx
directConvolve xs ys = V.create $ do
  let lX = V.length xs
      lY = V.length ys
      lZ = lX + lY - 1
  vec <- VM.replicate lZ 0
  forM_ [0 .. lX - 1] $ \i -> do
    let x = V.unsafeIndex xs i
    forM_ [0 .. lY - 1] $ \j -> do
      let y = V.unsafeIndex ys j
          k = i + j
      v <- VM.unsafeRead vec k
      VM.unsafeWrite vec k $! v + x * y
  pure vec

fftConvolve :: V.Vector Cpx -> V.Vector Cpx -> V.Vector Cpx
fftConvolve xsPre ysPre = V.fromListN lZ (V.toList (iditFft rs))
  where
    lX = V.length xsPre
    lY = V.length ysPre
    lZ = lX + lY - 1
    -- TOOD: admittedly this is a bit wasteful...
    xs = V.fromListN lZ (V.toList xsPre <> repeat 0)
    ys = V.fromListN lZ (V.toList ysPre <> repeat 0)
    xs' = ditFft xs
    ys' = ditFft ys
    rs = V.zipWith (*) xs' ys'

evalRandomConvolve :: M ()
evalRandomConvolve = do
  xs <- genCpxVector
  ys <- genCpxVector
  let lX = V.length xs
      lY = V.length ys
      lZ = lX + lY - 1
      zsDirect = directConvolve xs ys
      zsFft = fftConvolve xs ys
      diffs :: V.Vector Double
      diffs = V.zipWith (\x y -> magnitude (x - y)) zsDirect zsFft
  liftIO $ do
    putStrLn $ "Vector lengths (xs,ys,xs * ys): " <> show (lX, lY, lZ)
    _ <- forceAndMeasure "Direct" zsDirect
    _ <- forceAndMeasure "FFT" zsFft
    printf "StdDev: %.9f\n" (stdDev diffs)

forceAndMeasure :: NFData a => String -> a -> IO a
forceAndMeasure tag r = do
  tBefore <- getCurrentTime
  tAfter <- r `deepseq` getCurrentTime
  printf "Time (%s): %.6f\n" tag ((realToFrac $ diffUTCTime tAfter tBefore) :: Double)
  pure r

main :: IO ()
main = do
  let cs = (\x -> x*2 :+ (x*2 + 1)) <$> [0..19]
      vs = V.fromList cs
  evaluateOnVector vs
  g <- newStdGen
  evalStateT (evalRandomVector >> evalRandomConvolve) g


{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE TypeOperators    #-}
module Filter.Gaussian where

import           Control.Monad       (when)
import           Data.Array.Repa     as R
import           Data.Complex
import           Data.List           as L (length, reverse, (++))
import           Data.Vector.Generic as VG (convert)
import           DFT.Plan
import           Filter.Utils

data Gaussian2DParams = Gaussian2DParams
  { getGaussianFilterSigma :: !Double
  , getGaussianFilterRows  :: !Int
  , getGaussianFilterCols  :: !Int
  } deriving (Show)

{-# INLINE gaussian2D #-}
gaussian2D
  :: (Floating a)
  => a -> Int -> Int -> a
gaussian2D sd i j =
  1 / ((2 * pi) * sd * sd) * exp (-r / (2 * (sd ^ (2 :: Int))))
  where
    r = fromIntegral (i * i + j * j)

gaussian2DFilter ::
     DFTPlan -> Gaussian2DParams -> IO (DFTPlan, Array U DIM2 (Complex Double))
gaussian2DFilter plan (Gaussian2DParams sigma rows cols) = do
  let filterVec =
        VG.convert .
        toUnboxed .
        computeS .
        R.map (:+ 0) . makeFilter2D . R.fromFunction (Z :. cols :. rows) $ \(Z :. i :. j) ->
          gaussian2D sigma (i - (div cols 2)) (j - (div rows 2))
  lock <- getFFTWLock
  (plan1, filterVecF) <- dft1dGPlan lock plan [cols, rows] [0, 1] filterVec
  (newPlan, _) <- idft1dGPlan lock plan1 [cols, rows] [0, 1] filterVecF
  return (newPlan, fromUnboxed (Z :. cols :. rows) . VG.convert $ filterVecF)

{-# INLINE convolveGaussian2D #-}
convolveGaussian2D ::
     (Source r (Complex Double), Shape sh)
  => DFTPlan
  -> Array U DIM2 (Complex Double)
  -> Array r (sh :. Int :. Int) (Complex Double)
  -> IO (Array U (sh :. Int :. Int) (Complex Double))
convolveGaussian2D plan filterF input = do
  let len = rank . extent $ input
  when
    (len < 2)
    (error $ "convolveGaussian2D: rank = " L.++ show len L.++ " is less than 2.")
  inputF <-
    fmap (fromUnboxed (extent input) . VG.convert) .
    dftExecute
      plan
      (DFTPlanID
         DFT1DG
         (L.reverse . listOfShape . extent $ input)
         [len - 2, len - 1]) .
    VG.convert . toUnboxed . computeS . delay $
    input
  fmap (fromUnboxed (extent input) . VG.convert) .
    dftExecute
      plan
      (DFTPlanID
         IDFT1DG
         (L.reverse . listOfShape . extent $ input)
         [len - 2, len - 1]) .
    VG.convert . toUnboxed . computeS . R.traverse2 inputF filterF const $ \fInput fFilter idx@(sh :. i :. j) ->
    fInput idx * fFilter (Z :. i :. j)


{-|
Module      : kMeans
Description : k-Means
Copyright   : (c) Fabrício Olivetti, 2017
License     : GPL-3
Maintainer  : fabricio.olivetti@gmail.com

k-Means Clustering algorithm
-}

module KMeans where

import Data.List
import Data.Ord
--import qualified Numeric.LinearAlgebra as H

import Vector
import Dados

-- |'parseFile' parses a space separated file
-- to a list of lists of Double
parseFile :: String -> [[Double]]
parseFile file = map parseLine (lines file)
  where
    parseLine l = map toDouble (words l)
    toDouble  w = read w :: Double

firstClusters idxs points = [ points !! i | i <- idxs ]

kmeans it points clusters
  | it == 0               = clusters
  | clusters' == clusters = clusters
  | otherwise             = kmeans (it-1) points clusters'
  where
    clusters' = emStep points clusters

emStep points clusters = maximizationStep $ estimationStep points clusters

estimationStep :: [[Double]] -> [[Double]] -> [[[Double]]]
estimationStep points clusters = map unTuple
                               $ groupByValue
                               $ sortByValue
                               $ zip points
                               $ assign points clusters
  where
    unTuple = map fst

maximizationStep :: [[[Double]]] -> [[Double]]
maximizationStep candidates = map mean candidates
  where
    mean xs = (foldl' (.+.) (head xs) (tail xs)) ./ (fromIntegral $ length xs)

assign :: [[Double]] -> [[Double]] -> [Integer]
assign points clusters = map closestTo points
  where
    closestTo p = argmin $ map (euclid p) clusters
    euclid pi ci = sum $ map (^2) $ zipWith (-) pi ci
    argmin xs = fst $ head $ sortByValue $ zip [0..] xs


{-|

 Module      :  HF.Gauss
 Copyright   :  Copyright (c) David Schlegel
 License     :  BSD
 Maintainer  :  David Schlegel
 Stability   :  experimental
 Portability :  Haskell

Gaussian Integral evaluation plays an important role in quantum chemistry. Here functions will be provided to compute the most important integrals involving gaussian-type orbitals (GTOs).

An unnormalized primitive cartesian Gaussian Orbital has the form

<<centered_gaussian.svg centered_gaussian>>

where /A/ denotes the center of the orbital.

An unnormalized contracted Gaussian Orbital is a linear combination of primitive Gaussian:

<<centered_contracted_gaussian.svg centered_contracted_gaussian>>

"Gauss" uses the datatype contstructors 'PG' (Primitive Gaussian) and 'Ctr' (Contraction), defined in "HF.Data".

Notice: Formulas of Gaussian Integrals are taken from:
	Fundamentals of Molecular Integral Evaluation
	by Justin T. Fermann and Edward F. Valeev

-}

module HF.Gauss (
-- * Basic functions
distance, frac, center, itersum, erf, f, factorial2, general_erf,
-- * Normalization
normCtr, normPG,
-- * Integral Evaluation
s_12, t_12,
-- * Contraction operations
zipContractionWith, constr_matrix

) where

import HF.Data
import Numeric.Container hiding (linspace)
import Numeric.LinearAlgebra.Data hiding (linspace)
import Data.Maybe
import Control.DeepSeq
import Debug.Trace

---------------------
---------------------
--Helper Functions---
---------------------
---------------------

itersum function range
	| range == [] = 0
	| otherwise = (function $ head range) + itersum function (tail range)

-- |Calculates the factorial f(x) = x!
factorial n
    | n < 0     = 0
    | otherwise = product [1..n]

-- |Calculates the Doublefactorial f(x) = x!!
factorial2 :: (Show a, Eq a, Num a) => a -> a
factorial2 0 = 1
factorial2 1 = 1
factorial2 (-1) = 1
factorial2 n = n * factorial2 (n-2)

-- |Calculates the binomialcoefficient n over k
--binom :: Integer -> Integer -> Double
binom n 0 = 1
binom 0 k = 0
binom n k = binom (n-1) (k-1) * n `div` k

-- |Computes the distance between to vectors:
--
-- <<norm.svg norm>>
distance :: Vector Double -> Vector Double -> Double
distance rA rB = norm2(rA - rB)

-- | Calculates effective coefficient:
--
-- >>> frac alpha beta = alpha*beta/(alpha + beta)
frac ::  Double -> Double -> Double
frac alpha beta = alpha*beta/(alpha + beta)

-- | Calculates weighted radius of two gaussians
center :: Double -- ^ alpha
			-> Vector Double -- ^ rA
			-> Double -- ^ beta
			-> Vector Double -- ^ rB
			-> Vector Double
center a rA b rB = ((scalar b * rB) * (scalar a * rA))/scalar (a+b)

-- |error function calculation
erf :: Double -> Double
erf x = 2/sqrt pi * sum [x/(2*n+1) * product [-x^2/k |  k <- [1..n] ] | n <- [0..100]]

--error function calculation
f_0 :: Double -> Double
f_0 = erf

-- | General error function. See Eq. 4.118 Thijssen.
general_erf nu u
	| nu == 0 = sqrt pi * erf (sqrt u) * 1/(2*sqrt u)
	| otherwise = 1/(2*u) * (2*nu - 1) * (general_erf (nu-1) u ) - exp (-u)

--------------------------
---Important Functions----
--------------------------

-- |Calculates normalization factor of contracted Gaussian with arbitrary angular momentum
normCtr :: Ctr -> Double
normCtr contr = (prefactor * summation)**(-1.0/2.0)
	--See also Eq. 2.11
	where
		(l, m, n) = lmncontract contr
		--(l,m,n) = (fromIntegral l_, fromIntegral m_, fromIntegral n_) --This looks quite dirty
		n_sum = lengthcontract contr -1
		a = toList $ coefflist contr
		alp = [alpha prim | prim <- (gaussians contr)]
		mom = fromIntegral $ momentumctr contr
		prefactor = 1.0/(2.0**mom)*pi**(3.0/2.0)* factorial2 (2*l-1) * factorial2 (2*m-1)* factorial2 (2*n-1)
		summation = sum $ concat $ [[(a !! i)*(a !! j)/((alp !! i +alp !! j)**(mom + 3.0/2.0)) | i <-[0..n_sum]]| j <- [0..n_sum]]

-- |Calculates normalization factor of primitive Gaussian with arbitrary angular momentum
normPG :: PG -> Double
normPG (PG lmn alpha pos) = (2*alpha/pi)**(3/4) * (a/b)**(1/2)
	where
		(l, m, n) = (fromIntegral $ lmn !! 0, fromIntegral $ lmn !! 1, fromIntegral $ lmn !! 2)
		a = (8*alpha)**(l+m+n) * factorial l * factorial m * factorial n
		b = factorial (2*l) * factorial (2*m) * factorial (2*n)

-- |Calculates f_k (l1 l2 pa_x pb_x) used in Gaussian Product Theorem
f :: Int -> Int -> Int -> Double -> Double -> Double
--trace ("f" ++ " " ++  show k ++ " " ++ show l1_ ++ " " ++ show l2_ ++ " " ++ show pa_x ++ " " ++ show pb_x) $
f k l1_ l2_ pa_x pb_x = sum $ concat $ [[pa_x**(fromIntegral (l1_-i))  *  pb_x**(fromIntegral(l2_-j)) * fromIntegral ((binom l1_ i) * (binom l2_ j))| i <- [0..l1_], i+j <= k]| j <-[0..l2_]]
--See also Eq. 2.45

-- |Evaluates a given function for two contractions. The operation is symmetric in contraction arguments.
zipContractionWith :: (PG -> PG -> Double) -- ^ Function of two primitive Gaussians
							-> Ctr -- ^ Contraction
							-> Ctr -- ^ Contraction
							-> Double
zipContractionWith zipfunction (Ctr pglist1 coeffs1) (Ctr pglist2 coeffs2) = pr1 * pr2 * value
	where
		coefflist1 = toList coeffs1
		coefflist2 = toList coeffs2
		n1 = length coefflist1 -1
		n2 = length coefflist2 -1
		pr1 = normCtr (Ctr pglist1 coeffs1) --I am not so sure about this
		pr2 = normCtr (Ctr pglist2 coeffs2) --I am not so sure about this
		value = sum $ [(coefflist1 !! i)* (coefflist2 !! j)* (zipfunction (pglist1 !! i) (pglist2 !! j)) | i <- [0..n1], j <- [0..n2]]

-- |Construct a matrix out of a list of Contractions and a function for two primitive gaussians
constr_matrix :: [Ctr] -- ^ List of Contraction
						-> (PG -> PG -> Double) -- ^ Function of two primitive Gaussians
						-> Matrix Double -- ^ Matrix of dimensions (length [Ctr]) x (length [Ctr])
constr_matrix contractionlist function = buildMatrix (length contractionlist) (length contractionlist) (\(i,j) -> zipContractionWith function (contractionlist !! i) (contractionlist !! j) )

-- |Calculates overlap integral of two primitive gaussians

-- | <<s_12.svg s_12>>
s_12 :: PG -> PG -> Double
s_12 (PG lmn1 alpha1 pos1) (PG lmn2 alpha2 pos2) = trace ("s_12:   range " ++ show range  ++ "   f 0 =" ++ show (factor 0) ++ "   f 1 = " ++ show (factor 1)++ "   f 2 =" ++ show (factor 2)) $ prefactor * (factor 0) * (factor 1) * (factor 2) --pref * I_x * I_y * I_z
	where
		(l1, m1, n1) = (lmn1 !! 0, lmn1 !! 1, lmn1 !! 2)
		(l2, m2, n2) = (lmn2 !! 0, lmn2 !! 1, lmn2 !! 2)
		g = alpha1 + alpha2
		prefactor = exp (-alpha1*alpha2 * (distance pos1 pos2) /g)
		p = center alpha1 pos1 alpha2 pos2
		pa = toList $ p  - pos1
		pb = toList $ p  - pos2
		range = [0..(fromIntegral (l1+l2) /2)] :: (Enum a, Fractional a) => [a]
		-- See also Eq. 3.15
		function k i = force $ f (round (2*i)) l1 l2 (pa !! k) (pb !! k) *  (factorial2 (2.0* i -1.0)) * ((2*g)** (- i)) * (pi/g)**(1/2)
		factor k = itersum (function k) range
		--i k = sum $ [f (round (2*i)) l1 l2 (pa !! k) (pb !! k) * factorial2 (2.0* i -1.0) * ((2*g)** (- i)) * (pi/g)**(1/2) | i <- range ]

-- |Calculates kinetic energy integral of two primitive gaussians

-- | <<t_12.svg t_12>>
t_12 :: PG -> PG -> Double
t_12 (PG lmn1 alpha1 pos1) (PG lmn2 alpha2 pos2) =  (i 0) + (i 1) + (i 2) -- I_x + I_y + I_z
	where
		--lifts or lowers l, m or n, indicated by index k for a lmn-list, respectively
		pp1 list k = [if i == k then list !! i + 1 else list !! i | i<-[0..2]]	--lifts
		mm1 list k = [if i == k then list !! i - 1 else list !! i | i<-[0..2]]	--lowers
		--overlaps where parts of l, m or n ar lifted +1 or lowered -1
		m1m1 k = s_12 (PG (mm1 lmn1 k) alpha1 pos1) (PG (mm1 lmn2 k) alpha2 pos2) -- <-1|-1>_k
		p1p1 k = s_12 (PG (pp1 lmn1 k) alpha1 pos1) (PG (pp1 lmn2 k) alpha2 pos2) -- <+1|+1>_k
		p1m1 k = s_12 (PG (pp1 lmn1 k) alpha1 pos1) (PG (mm1 lmn2 k) alpha2 pos2) -- <+1|-1>_k
		m1p1 k = s_12 (PG (mm1 lmn1 k) alpha1 pos1) (PG (pp1 lmn2 k) alpha2 pos2) -- <-1|+1>_k
		--See also Eq.4.13
		i k = 1/2 * (fromIntegral (lmn1 !! k)) * (fromIntegral (lmn2 !! k)) * (m1m1 k)
			+ 2 * alpha1 * alpha2 * (p1p1 k)
			- alpha1 * (fromIntegral (lmn2 !! k)) * (p1m1 k)
			- alpha2 * (fromIntegral (lmn1 !! k)) * (m1p1 k)

{-
-- |Calculates nuclear attraction integral of two primitive gaussians
v_12 :: Double -> Vecotr Double -> PG -> PG -> Double
v_12 z r_c (PG lmn1 alpha1 pos1) (PG lmn2 alpha2 pos2) = z*pi*pre/g * 2* summation
	where
		ax, ay, az = (fromList pos1) !! 0, (fromList pos1) !! 1, (fromList pos1) !! 2
		bx, by, bz = (fromList pos2) !! 0, (fromList pos2) !! 1, (fromList pos2) !! 2
		r_cx, r_cy, r_cz = (fromList r_c) !! 0, (fromList r_c) !! 1, (fromList r_c) !! 2
		pref = exp (-alpha1*alpha2 * (distance pos1 pos2) /g)
		g = alpha1 + alpha2
		p = 1/g * (alpha1*pos1 + alpha2*pos2)
		px, py, pz = (fromList p) !! 0, (fromList p) !! 1, (fromList p) !! 2
		eta = (alpha1 + alpha2)/g
		mu_x = l1 + l2 - 2*(ijk1 + ijk2) - (opq1 + opq2)
		mu_y = l1 + l2 - 2*(j1 + j2) - (p1 + p2)
		mu_z = l1 + l2 - 2*(k1 + k2) - (q1 + q2)
		nu = mu_x + mu_y + mu_z - (u + v + w)
		summation = general_erf nu (g * distance p r_c) * a_x * a_y * a_z

		--ranges:
		ijk1_range = [0..floor $ (fromIntegral l1 /2)] :: (Integral a, Enum a) => [a]
		ijk2_range = [0..floor $ (fromIntegral l2 /2)] :: (Integral a, Enum a) => [a]
		opq1_range ijk1 = [0..floor $ (fromIntegral l1 - 2*ijk1)] :: (Integral a, Enum a) => [a]
		opq2_range ijk2 = [0..floor $ (fromIntegral l2 - 2*ijk2)] :: (Integral a, Enum a) => [a]
		rst_range opq1 opq2 = [0..floor $ fromIntegral (opq1+l2) /2] :: (Integral a, Enum a) => [a]
		u_range mu_x = [0..floor $ (fromIntegral mu_x /2)] :: (Integral a, Enum a) => [a]

		a_x = (-1)**(l1+l2) * (factorial l1) * (factorial l2) *
		(-1)**(opq2 + r) * factorial (opq1 + opq2) * 1/(4**(ijk1+ijk2+r) * factorial ijk1 * factorial ijk2 * factorial opq1 * factorial opq2 * factorial r) *
		alpha1**(opq2-ijk1-r) * alpha2**(opq1-ijk2-r) * (ax-bx)**(opq1+opq2-2*r) * 1/(factorial (l1-2*ijk1-opq1) * factorial (l2-2*ijk2-opq2) * factorial (opq1-2*opq2-2*r)) *
		* [ (-1)**u * factorial mu_x * (px -r_cx)**(mu_x - 2*u) * 1/ ((4)**u * factorial u * factorial (mu_x-2*u) * g**(opq1 + opq2 - r + u)) | u <- [0..1/2 * mu_x]]
-}

--v_12_test ::  Double -> Vector Double -> PG -> PG -> [Double]
{- v_12_test z r_c (PG lmn1 alpha1 pos1) (PG lmn2 alpha2 pos2) = summation 0
	where
		--(ax, ay, az) = ((toList pos1) !! 0, (toList pos1) !! 1, (toList pos1) !! 2)
		--(bx, by, bz) = ((toList pos2) !! 0, (toList pos2) !! 1, (toList pos2) !! 2)
		--(r_cx, r_cy, r_cz) = ((toList r_c) !! 0, (toList r_c) !! 1, (toList r_c) !! 2)
		--lmn1_ = fromList lmn1
		--lmn2_ = fromList lmn2

		pref = exp (-alpha1*alpha2 * (distance pos1 pos2) /g)
		g = alpha1 + alpha2
		p = scale (1/g) (add (scale alpha1 pos1) (scale alpha2 pos2))
		--(px, py, pz) = ((toList p) !! 0, (toList p) !! 1, (toList p) !! 2)

		--f_nu ijk1 ijk2 opq1 opq2 = general_erf (nu ijk1 ijk2 opq1 opq2) (g * distance p r_c)

		eta = (alpha1 + alpha2)/g
		mu ijk1 ijk2 opq1 opq2 s = fromIntegral (atIndex lmn1 s) + fromIntegral (atIndex lmn2 s) - 2*(ijk1 + ijk2) - (opq1 + opq2)
		--nu ijk1 ijk2 opq1 opq2 = (mu ijk1 ijk2 opq1 opq2 s) + (mu ijk1 ijk2 opq1 opq2 s) + (mu ijk1 ijk2 opq1 opq2 s) - (u v q)

		--ranges alle ranges (X,Y,Z) mit einer klappe schlagen, dafür index s:
		ijk1_range s = [0..fromIntegral $ floor $ fromIntegral (lmn1 !! s) /2] :: [Double]
		ijk2_range s = [0..fromIntegral $ floor $ fromIntegral (lmn2 !! s) /2] :: [Double]
		opq1_range ijk1 s = [0..fromIntegral $ floor $ (fromIntegral (lmn1 !! s) - 2*ijk1)] :: [Double]
		opq2_range ijk2 s = [0..fromIntegral $ floor $ (fromIntegral (lmn2 !! s) - 2*ijk2)] :: [Double]
		rst_range opq1 opq2 s = [0..fromIntegral $ floor $ (opq1+opq2) /2] :: [Double]
		uvw_range ijk1 ijk2 opq1 opq2 s = [0..fromIntegral $ floor $ ((mu ijk1 ijk2 opq1 opq2 s) /2)] :: [Double]

		--a_x :: Num a => a -> a -> a -> a -> a -> Double

		--(-1)**(l1+l2) * (factorial fromIntegral (lmn1 !! s)) * (factorial fromIntegral (lmn2 !! s)) *
		function ijk1 ijk2 opq1 opq2 rst s = (-1)**(opq2 + rst) * factorial (opq1 + opq2) * 1/(4**(ijk1+ijk2+rst) * factorial ijk1 * factorial ijk2 * factorial opq1 * factorial opq2 * factorial rst) * alpha1**(opq2-ijk1-rst) * alpha2**(opq1-ijk2-rst) * (atIndex pos1 s - atIndex pos2 s)**(opq1+opq2-2*rst) * 1/(factorial ((lmn1 !! s)-2*ijk1-opq1) * factorial ((lmn2 !! s)-2*ijk2-opq2) * factorial (opq1-2*opq2-2*rst)) * sum $ [ (-1)**uvw * factorial (mu ijk1 ijk2 opq1 opq2 s) * ((atIndex p s) - (atIndex r_c s))**((mu ijk1 ijk2 opq1 opq2 s) - 2*uvw) * 1/ ((4)**uvw * factorial uvw * factorial ((mu ijk1 ijk2 opq1 opq2 s)-2*uvw) * g**(opq1 + opq2 - rst + uvw)) | uvw <- (uvw_range ijk1 ijk2 opq1 opq2 s)]

		summation s = (-1)**(lmn1 !! s + lmn2 !! s) * (factorial fromIntegral (lmn1 !! s)) * (factorial fromIntegral (lmn2 !! s)) * sum $ concat $ concat $ [[[function ijk1 ijk2 opq1 opq2 rst s | rst <- (rst_range opq1 opq2 s)] | opq1 <- (opq1_range ijk1 s), opq2 <- (opq2_range ijk2 s)]| ijk1 <- ijk1_range s, ijk2 <- ijk2_range s]
-}

-----------------
--Deprecated-----
-----------------

-- Here all importand integrals involving gaussian functions will be evaluated.

--Overlap integral for s orbitals
--calculates <1s,alpha, A | 1s, beta, B>
overlaps :: Double -> Double -> Vector Double -> Vector Double -> Double
overlaps alpha beta rA rB = prefactor * exp exponent
	where 	prefactor = (pi/(alpha + beta))**(3/2)
		exponent = - (frac alpha beta * distance rA rB)

--Kinetic integral
--calculates <1s,alpha, A | - Laplace | 1s, beta, B>
kinetic :: Double -> Double -> Vector Double -> Vector Double -> Double
kinetic alpha beta rA rB = prefactor * exp(exponent)
	where 	prefactor = (frac alpha beta)
			 *(6 -4 *(frac alpha beta) * distance rA rB)
			 *(pi/(alpha + beta))**(3/2)
		exponent = - ((frac alpha beta) * distance rA rB)

--Nuclear interaction integral
--calculates <1s,alpha, A | - Z/r_C | 1s, beta, B>
nuclear :: Double -> Double -> Vector Double -> Vector Double -> Vector Double -> Double -> Double
nuclear alpha beta rA rB rC z = pref1 * pref2 * exp(exponent)
	where 	pref1 = -2*pi*z/(alpha + beta)
		pref2 = f_0 arg
		arg   = (distance (center alpha rA beta rB) rC )   * (alpha*beta)
		exponent = - ((frac alpha beta) * (distance rA rB))

--two-electron integral
--important!
--calculates <1s,alpha, A ; 1s, beta, B | 1s,gamma, C ; 1s, delta, D >
twoelectron :: Double -> Double -> Double -> Double -> Vector Double -> Vector Double -> Vector Double -> Vector Double -> Double
twoelectron a b g d rA rB rC rD = pref * exp(exponent) * (f_0 arg)
	where	pref = 			(2*pi**(5/2))/
				( (a+g)*(b+d)*(a+b+g+d)**(1/2)   )
		exponent = 	-a*g/(a+g) * (distance rA rC)
				-b*d/(b+d) * (distance rB rD)
		arg = (a+g)*(b+d)/(a+b+g+d) * ( distance (center a rA g rC) (center b rB d rD) )

--Overlap Matrix S_pq = <p|q>
--overlap p q  = t
--where
--	t = buildMatrix (n*n) (n*n) (\(i,j) -> overlaps alpha beta rA rB )


{-# LANGUAGE FlexibleInstances, FlexibleContexts #-}
module Main where
import Test.Hspec
import Test.QuickCheck
import Numeric.LinearAlgebra
import qualified Data.Vector.Storable as V
import Test.Gen
import Test.Utils

main = hspec $ do
  describe "Corr Single" $ do
    it "implements corr2 correct.0" $ do
      forAll (pair (squared_real_matrices 3) (squared_real_matrices 4)) $
        \(m1, m2) -> ioProperty $ do
          r <- test_corr2 2 m1 m2
          return $ good_corr2 2 m1 m2 `eqShowWhenFail` r
    it "implements corr2 correct.1" $ do
      forAll (pair (squared_real_matrices 9) (squared_real_matrices 9)) $
        \(m1, m2) -> ioProperty $ do
          r <- test_corr2 4 m1 m2
          return $ good_corr2 4 m1 m2 `eq` r
    it "implements corr2 correct.2" $ do
      forAll (pair (squared_real_matrices 5) (squared_real_matrices 28)) $
        \(m1, m2) -> ioProperty $ do
          r <- test_corr2 2 m1 m2
          return $ eqShowWhenFail (good_corr2 2 m1 m2) r
    it "implements corr2 correct.3" $ do
      forAll (pair (choose (2,30)) (pair small_matrices small_matrices)) $
        \(p, (m1, m2)) -> ioProperty $ do
          r <- test_corr2 p m1 m2
          return $ good_corr2 p m1 m2 `eqShowWhenFail` r
  describe "Corr Many" $ do
    it "with 2 kernels" $ do
      forAll (pair (sequence $ replicate 2 $ squared_real_matrices 3) (squared_real_matrices 7)) $
        \(m1s, m2) -> ioProperty $ do
          rs <- test_corr2_arr 2 m1s m2
          return $ conjoin $ zipWith (\m r -> good_corr2 2 m m2 `eqShowWhenFail` r) m1s rs
    it "with 5 kernels" $ do
      forAll (pair (sequence $ replicate 4 $ squared_real_matrices 15) (squared_real_matrices 88)) $
        \(m1s, m2) -> ioProperty $ do
          rs <- test_corr2_arr 2 m1s m2
          ss <- return $ map (\m -> good_corr2 2 m m2) m1s
          return $ conjoin $ zipWith eq rs ss

eqShowWhenFail m1 m2 =
    whenFail (do let va = flatten m1
                 let vb = flatten m2
                 let err x 0 = x
                     err x y = abs ((x - y) / y)
                 let ev = (V.zipWith err va vb)
                     ei = V.maxIndex ev
                 putStrLn $ "Max error ration: " ++ show (ev V.! ei, va V.! ei, vb V.! ei)
                 putStrLn $ show m1
                 putStrLn $ show m2)
        (m1 `eq` m2)


-----------------------------------------------------------------------------
-- |
-- Module      :  Data.Random.Distribution.Static.MultivariateNormal
-- Copyright   :  (c) 2016 FP Complete Corporation
-- License     :  MIT (see LICENSE)
-- Maintainer  :  dominic@steinitz.org
--
-- Sample from the multivariate normal distribution with a given
-- vector-valued \(\mu\) and covariance matrix \(\Sigma\). For
-- example, the chart below shows samples from the bivariate normal
-- distribution. The dimension of the mean \(n\) is statically checked
-- to be compatible with the dimension of the covariance matrix \(n \times n\).
--
-- <<diagrams/src_Data_Random_Distribution_Static_MultivariateNormal_diagMS.svg#diagram=diagMS&height=600&width=500>>
--
-- Example code to generate the chart:
--
-- > {-# LANGUAGE DataKinds #-}
-- >
-- > import qualified Graphics.Rendering.Chart as C
-- > import Graphics.Rendering.Chart.Backend.Diagrams
-- >
-- > import Data.Random.Distribution.Static.MultivariateNormal
-- >
-- > import qualified Data.Random as R
-- > import Data.Random.Source.PureMT
-- > import Control.Monad.State
-- > import Numeric.LinearAlgebra.Static
-- >
-- > nSamples :: Int
-- > nSamples = 10000
-- >
-- > sigma1, sigma2, rho :: Double
-- > sigma1 = 3.0
-- > sigma2 = 1.0
-- > rho = 0.5
-- >
-- > singleSample :: R.RVarT (State PureMT) (R 2)
-- > singleSample = R.sample $ Normal (vector [0.0, 0.0])
-- >                (sym $ matrix [ sigma1, rho * sigma1 * sigma2
-- >                              , rho * sigma1 * sigma2, sigma2])
-- >
-- > multiSamples :: [R 2]
-- > multiSamples = evalState (replicateM nSamples $ R.sample singleSample) (pureMT 3)
-- >
-- > pts = map f multiSamples
-- >   where
-- >     f z = (x, y)
-- >       where
-- >         (x, t) = headTail z
-- >         (y, _) = headTail t
-- >
-- > chartPoint pointVals n = C.toRenderable layout
-- >   where
-- >
-- >     fitted = C.plot_points_values .~ pointVals
-- >               $ C.plot_points_style  . C.point_color .~ opaque red
-- >               $ C.plot_points_title .~ "Sample"
-- >               $ def
-- >
-- >     layout = C.layout_title .~ "Sampling Bivariate Normal (" ++ (show n) ++ " samples)"
-- >            $ C.layout_y_axis . C.laxis_generate .~ C.scaledAxis def (-3,3)
-- >            $ C.layout_x_axis . C.laxis_generate .~ C.scaledAxis def (-3,3)
-- >
-- >            $ C.layout_plots .~ [C.toPlot fitted]
-- >            $ def
-- >
-- > diagMS = do
-- >   denv <- defaultEnv C.vectorAlignmentFns 600 500
-- >   return $ fst $ runBackend denv (C.render (chartPoint pts nSamples) (500, 500))
--
-----------------------------------------------------------------------------

{-# OPTIONS_GHC -Wall                     #-}
{-# OPTIONS_GHC -fno-warn-name-shadowing  #-}
{-# OPTIONS_GHC -fno-warn-type-defaults   #-}
{-# OPTIONS_GHC -fno-warn-unused-do-bind  #-}
{-# OPTIONS_GHC -fno-warn-missing-methods #-}
{-# OPTIONS_GHC -fno-warn-orphans         #-}

{-# LANGUAGE MultiParamTypeClasses        #-}
{-# LANGUAGE TypeFamilies                 #-}
{-# LANGUAGE ScopedTypeVariables          #-}
{-# LANGUAGE DataKinds                    #-}

module Data.Random.Distribution.Static.MultivariateNormal
    ( Normal(..)
    ) where

import           Data.Random hiding ( StdNormal, Normal )
import qualified Data.Random as R
import           Control.Monad.State ( replicateM )
import qualified Numeric.LinearAlgebra.HMatrix as H
import           Numeric.LinearAlgebra.Static as S
                 ( R, vector, extract, Sq, Sym, col,
                   tr, linSolve, uncol, chol, (<.>),
                   ℝ, (<>), diag, (#>), eigensystem
                 )
import          GHC.TypeLits ( KnownNat, natVal )
import          Data.Maybe ( fromJust )

normalMultivariate :: KnownNat n =>
                      R n -> Sym n -> RVarT m (R n)
normalMultivariate mu bigSigma = do
  z <- replicateM (fromIntegral $ natVal mu) (rvarT R.StdNormal)
  return $ mu + bigA #> (vector z)
  where
    (vals, bigU) = eigensystem bigSigma
    lSqrt = diag $ mapVector sqrt vals
    bigA = bigU S.<> lSqrt

mapVector :: KnownNat n => (ℝ -> ℝ) -> R n -> R n
mapVector f = vector . H.toList . H.cmap f . extract

sumVector :: KnownNat n => R n -> ℝ
sumVector x = x <.> 1

data family Normal k :: *

data instance Normal (R n) = Normal (R n) (Sym n)

instance KnownNat n => Distribution Normal (R n) where
  rvar (Normal m s) = normalMultivariate m s

normalLogPdf :: KnownNat n =>
                R n -> Sym n -> R n -> Double
normalLogPdf mu bigSigma x = - sumVector (mapVector log (diagonals dec))
                             - 0.5 * (fromIntegral $ natVal mu) * log (2 * pi)
                             - 0.5 * s
  where
    dec = chol bigSigma
    t = uncol $ fromJust $ linSolve (tr dec) (col $ x - mu)
    u = mapVector (\x -> x * x) t
    s = sumVector u

normalPdf :: KnownNat n =>
             R n -> Sym n -> R n -> Double
normalPdf mu sigma x = exp $ normalLogPdf mu sigma x

diagonals :: KnownNat n => Sq n -> R n
diagonals = vector . H.toList . H.takeDiag . extract

instance KnownNat n => PDF Normal (R n) where
  pdf (Normal m s) = normalPdf m s
  logPdf (Normal m s) = normalLogPdf m s


module NeuralNet.NetSpec (netSpec) where

import Test.Hspec
import Numeric.LinearAlgebra
import NeuralNet.Net
import NeuralNet.Activation
import NeuralNet.Layer
import NeuralNet.Example
import Control.Exception (evaluate)
import Data.List

matrixSize :: Matrix Double -> (Int, Int)
matrixSize m = (rows m, cols m)

allUnique :: Matrix Double -> Bool
allUnique m = length (nub list) == length list
  where list = concat (toLists m)

createdLayer :: Activation -> (Int, Int) -> NeuronLayer -> Bool
createdLayer a s l = activationEq && wSizeEq && bCorrect && allUnique m
  where
    m = layerW l
    bCorrect = layerB l == col (replicate (fst s) 0)
    activationEq = layerActivation l == a
    wSizeEq = matrixSize m == s

netSpec :: SpecWith ()
netSpec =
  describe "NeuralNet.Net" $ do
    describe "initNN" $ do
      it "returns correctly for small definition" $
        let
          nn = initNN [1..] (8, [LayerDefinition ReLU 5, LayerDefinition Tanh 1, LayerDefinition Sigmoid 2])
          layers = nnLayers nn
        in do
          layers!!0 `shouldSatisfy` createdLayer ReLU (5, 8)
          layers!!1 `shouldSatisfy` createdLayer Tanh (1, 5)
          layers!!2 `shouldSatisfy` createdLayer Sigmoid (2, 1)

      it "returns correctly for single layer definition" $
        let
          nn = initNN [1..] (6, [LayerDefinition ReLU 5])
          layers = nnLayers nn
        in
          layers!!0 `shouldSatisfy` createdLayer ReLU (5, 6)

      it "throws exception for no layers" $
        evaluate (initNN [1..] (6, [])) `shouldThrow` anyException

      it "throws exception for no inputs" $
        evaluate (initNN [1..] (0, [LayerDefinition ReLU 5])) `shouldThrow` anyException

    describe "buildNNFromList" $ do
      it "creates log reg correctly" $
        let
          nn = buildNNFromList (5, [LayerDefinition ReLU 1]) [1, 2, 3, 4, 5, 6]
          layers = nnLayers nn
        in do
          length layers `shouldBe` 1
          layerW (layers!!0) `shouldBe` fromLists [[1, 2, 3, 4, 5]]
          layerB (layers!!0) `shouldBe` fromLists [[6]]

      it "creates multi output correctly" $
        let
          nn = buildNNFromList (3, [LayerDefinition ReLU 2]) [1, 2, 3, 4, 5, 6, 7, 8]
          layers = nnLayers nn
        in do
          length layers `shouldBe` 1
          layerW (layers!!0) `shouldBe` fromLists [[1, 2, 3], [4, 5, 6]]
          layerB (layers!!0) `shouldBe` fromLists [[7], [8]]

      it "creates multi layer correctly" $
        let
          nn = buildNNFromList (2, [LayerDefinition ReLU 2, LayerDefinition Tanh 1]) [1, 2, 3, 4, 5, 6, 7, 8, 9]
          layers = nnLayers nn
        in do
          length layers `shouldBe` 2
          layerW (layers!!0) `shouldBe` fromLists [[1, 2], [3, 4]]
          layerB (layers!!0) `shouldBe` fromLists [[5], [6]]
          layerW (layers!!1) `shouldBe` fromLists [[7, 8]]
          layerB (layers!!1) `shouldBe` fromLists [[9]]

      it "throw error for incorrect list size" $
        evaluate (buildNNFromList (5, [LayerDefinition ReLU 1]) [1, 2]) `shouldThrow` anyException

    describe "isExampleSetCompatibleWithNN" $ do
      it "calculates correctly for compatible" $
        let
          nn = buildNNFromList (3, [LayerDefinition ReLU 1]) [1, 2, 3, 4]
          examples = createExampleSet [([1, 2, 3], 0), ([4, 5, 6], 1)]
        in isExampleSetCompatibleWithNN examples nn `shouldBe` True

      it "calculates correctly for incompatible" $
        let
          nn = buildNNFromList (3, [LayerDefinition ReLU 1]) [1, 2, 3, 4]
          examples = createExampleSet [([1, 2], 0), ([4, 5], 1)]
        in isExampleSetCompatibleWithNN examples nn `shouldBe` False

    describe "nnForward" $ do
      -- Not working, can't figure out why. When change to shouldBe [] it does throw an error
--      it "throws errors when incorrect num inputs provided" $
--        let nn = buildNNFromList (3, [LayerDefinition ReLU 1]) [1, 2, 3, 4]
--        in evaluate (nnForward nn [1, 1]) `shouldThrow` anyErrorCall

      it "calculates correctly for logreg" $
        let nn = buildNNFromList (3, [LayerDefinition ReLU 1]) [1, 2, 3, 4]
        in nnForward nn [1, 1, 2] `shouldBe` [13.0]

      it "calculates correctly for multi layer" $
        let
          def = (3, [LayerDefinition Sigmoid 2, LayerDefinition ReLU 1])
          nums = [1, -1, -1
            , 1, 1, -1
            , 1, -1
            , 2, 3,
            0]
          nn = buildNNFromList def nums
        in
          nnForward nn [1, 1, 1] `shouldBe` [2.5]

    describe "nnForwardSet" $ do
      it "calculates correctly for logreg" $
        let
          nn = buildNNFromList (3, [LayerDefinition ReLU 1]) [1, 2, 3, 4]
          examples = createExampleSet [([1, 1, 2], 0)]
          resultAs = map forwardPropA (nnForwardSet nn examples)
        in resultAs `shouldBe` [fromLists [[1], [1], [2]], scalar 13.0]

      it "calculates correctly for multiple examples, multiple output" $
        let
          nn = buildNNFromList (3, [LayerDefinition ReLU 2, LayerDefinition ReLU 1]) [1, 4, 2, 5, 3, 6, 7, 8, 1, 1, 0]
          examples = createExampleSet [([1, 1, 2], 0), ([1, 1, 0], 0)]
          resultAs = map forwardPropA (nnForwardSet nn examples)
        in resultAs `shouldBe` [fromLists [[1, 1], [1, 1], [2, 0]], fromLists [[16, 12], [28, 16]], fromLists [[44, 28]]]


{-# LANGUAGE NamedFieldPuns, TemplateHaskell #-}

module School.Unit.Test.LogSoftMax
( logSoftMaxTest ) where

import Control.Monad.IO.Class (liftIO)
import Numeric.LinearAlgebra ((><), size, toLists)
import School.TestUtils (diffInput, matIndexes, randomMatrix, randomMatrixL)
import School.Types.FloatEq ((~=))
import School.Unit.LogSoftMax
import School.Unit.Unit (Unit(..))
import School.Unit.UnitActivation (UnitActivation(..))
import School.Unit.UnitGradient (UnitGradient(..))
import School.Unit.UnitParams (UnitParams(..))
import School.Utils.LinearAlgebra (compareDoubleMatrix, oneMatrix)
import Test.Tasty (TestTree)
import Test.Tasty.QuickCheck hiding ((><), scale)
import Test.Tasty.TH
import Test.QuickCheck.Monadic (assert, monadicIO, pre)

eps :: Double
eps = 1e-5

prec :: Double
prec = 1e-3

prop_norm_size :: (Positive Int) -> (Positive Int) -> Property
prop_norm_size (Positive bSize) (Positive fSize) = monadicIO $ do
  input <- liftIO $ BatchActivation <$> randomMatrix bSize fSize
  let (BatchActivation result) = apply logSoftMax EmptyParams input
  assert $ and . (map (<= prec)) . concat . toLists $ result

prop_norm_sum :: (Positive Int) -> (Positive Int) -> Property
prop_norm_sum (Positive bSize) (Positive fSize) = monadicIO $ do
  input <- liftIO $ BatchActivation <$> randomMatrix bSize fSize
  let (BatchActivation result) = apply logSoftMax EmptyParams input
  assert $ and
         . (map (~= 1))
         . (map (sum . (map exp)))
         . toLists
         $ result

prop_gradient_dims :: (Positive Int) -> Positive Int -> Property
prop_gradient_dims (Positive bSize) (Positive fSize) = monadicIO $ do
  input <- liftIO $ fmap BatchActivation $ randomMatrix bSize fSize
  gradIn <- liftIO $ fmap BatchGradient $ randomMatrix bSize fSize
  let (BatchGradient gradOut, _) = deriv logSoftMax EmptyParams gradIn input
  assert $ (size gradOut == (bSize, fSize))

prop_numerical_gradient :: (Positive Int) -> Positive Int -> Property
prop_numerical_gradient (Positive bSize) (Positive fSize) = monadicIO $ do
  pre $ bSize < 25 && fSize < 25
  input <- liftIO $ fmap BatchActivation $ randomMatrixL 5 bSize fSize
  let params = EmptyParams
  let idxs = matIndexes bSize fSize
  let num = map (\idx -> diffInput logSoftMax params input eps idx)
                idxs
  let numGrad = (bSize >< fSize) num
  let inGrad = BatchGradient $ oneMatrix bSize fSize
  let (BatchGradient result, _) = deriv logSoftMax params inGrad input
  assert $ compareDoubleMatrix prec result numGrad

logSoftMaxTest :: TestTree
logSoftMaxTest = $(testGroupGenerator)


module Stepik where

import Prelude hiding (Monoid, mempty, mappend, mconcat, Functor, fmap)
import Data.Char
import Data.Function
import Data.List (transpose)
import Data.Functor
-- import Data.Semigroup ((<>))
import Control.Monad (ap, liftM)

import Data.Time.Clock
import Data.Time.Format
-- import System.Locale

main = undefined

{- Create own operator -}
infixl 6 *+*
(*+*) a b = a ^ 2 + b ^ 2

infixl 6 |-|
x |-| y = abs (x - y)

twoDigits2Int :: Char -> Char -> Int
twoDigits2Int x y = if isDigit x && isDigit y then read [x, y] :: Int else 100

doubleFact :: Integer -> Integer
doubleFact 1 = 1
doubleFact 2 = 2
doubleFact n = n * doubleFact(n-2)

{- Negative Factorial -}
fibonacci :: Integer -> Integer
fibonacci 0 = 0
fibonacci 1 = 1
fibonacci n
    | n < 0 = (-1)^(abs n + 1) * fibonacci (abs n)
    | otherwise = fibonacci (n - 1) + fibonacci (n - 2)

{- Tail recursion -}
factorial :: Integer -> Integer
factorial n
    | n >= 0 = helper 1 n
    | otherwise = error "Arg must be >= 0!"

helper :: Integer -> Integer -> Integer
helper acc 0 = acc
helper acc n = helper (acc * n) (n - 1)

{- Negative Fibonacci tail recursion -}
fibonacciTail :: Integer -> Integer
fibonacciTail n
    | n < 0 = (-1)^(abs n + 1) * fibonacciTail (abs n)
    | otherwise = helperFib 0 1 n

helperFib :: Integer -> Integer -> Integer -> Integer
helperFib p c 1 = c
helperFib p c n = helperFib c (c+p) (n-1)

{- Negative Fibonacci tail recursion `let ... in ...`. Task -}
fibonacciTail' :: Integer -> Integer
fibonacciTail' n
    | n < 0 = (-1)^(abs n + 1) * fibonacciTail' (abs n)
    | otherwise = let
        helperFib p c 0 = 0
        helperFib p c 1 = c
        helperFib p c n = helperFib c (c + p) (n - 1)
    in helperFib 0 1 n

fibonacciTail'' :: Integer -> Integer
fibonacciTail'' n
    | n < 0 = (-1)^(abs n + 1) * fibonacciTail'' (abs n)
    | otherwise = helper 0 1 n
    where
        helper p c 0 = 0
        helper p c 1 = c
        helper p c n = helper c (c + p) (n - 1)

{- Task -}
seqA :: Integer -> Integer
seqA 0 = 1
seqA 1 = 2
seqA 2 = 3
seqA n = seqA (n-1) + seqA (n-2) - 2 * seqA (n-3)

seqA' :: Integer -> Integer
seqA' n = helper 3 3 2 n
    where helper curr prev multip n
            | n == 0 = 1
            | n == 1 = 2
            | n == 2 = 3
            | n == 3 = curr
            | n > 3 = helper (curr + prev - 2 * multip) curr prev (n - 1)

seqA'' :: Integer -> Integer
seqA'' n = let
    helper curr prev multip 0 = 1
    helper curr prev multip 1 = 2
    helper curr prev multip 2 = 3
    helper curr prev multip 3 = curr
    helper curr prev multip n = helper (curr + prev - 2 * multip) curr prev (n - 1)
    in helper 3 3 2 n

seqA''' :: Integer -> Integer
seqA''' n = helper 3 3 2 n
    where
        helper curr prev multip 0 = 1
        helper curr prev multip 1 = 2
        helper curr prev multip 2 = 3
        helper curr prev multip 3 = curr
        helper curr prev multip n = helper (curr + prev - 2 * multip) curr prev (n - 1)

{- Task -}
sum'n'count :: Integer -> (Integer, Integer)
sum'n'count x = (s, l)
        where
            xs = show $ abs x
            l = toInteger $ length xs
            s = toInteger $ sum $ map digitToInt xs

{- Task -}
integration :: (Double -> Double) -> Double -> Double -> Double
integration f a b = let
                        xs = [a + j * h | j <- [0..1000]]
                        h = (b - a) / 1000
                    in sum [( ( f (xs !! i) + f (xs !! (i + 1)) ) / 2 ) * ( xs !! (i + 1) - xs !! i ) | i <- [0..999]]

integration' :: (Double -> Double) -> Double -> Double -> Double
integration' f a b = h * ( (f a + f b) / 2 + (sum $ map f xs) )
    where
        n = 1000
        h = (b - a) / n
        xs = [a + h * x | x <- [0..n-1]]

{- Task -}
getSecondFrom :: t1 -> t2 -> t3 -> t2
getSecondFrom x y z = y

{- Task -}
multSecond = g `on` h

g :: (Num a) => a -> a -> a
g x y = x * y

h :: (a, b) -> b
h (x, y) = y

{- Task -}
on3 :: (b -> b -> b -> c) -> (a -> b) -> a -> a -> a -> c
on3 op f x y z = op (f x) (f y) (f z)

{- Task -}
doItYourself = doItYourselfF . doItYourselfG . doItYourselfH

doItYourselfF = logBase 2
doItYourselfG = (^ 3)
doItYourselfH = max 42

{- Task -}
swap' = uncurry (flip (,))

{- Task -}
class Printable a where
    toString :: a -> String

instance Printable Bool where
  toString True = "true"
  toString False = "false"

instance Printable () where
  toString _ = "unit type"

instance (Printable a, Printable b) => Printable (a, b) where
  toString (x,y) = "(" ++ toString x ++ "," ++ toString y ++ ")"

{- Task -}
class KnownToGork a where
    stomp :: a -> a
    doesEnrageGork :: a -> Bool

instance KnownToGork Int where
  stomp _ = 4
  doesEnrageGork = (< 10)

class KnownToMork a where
    stab :: a -> a
    doesEnrageMork :: a -> Bool

instance KnownToMork Int where
  stab _ = 3
  doesEnrageMork = (> 5)

class (KnownToGork a, KnownToMork a) => KnownToGorkAndMork a where
    stompOrStab :: a -> a

instance KnownToGorkAndMork Int where
    stompOrStab x
        | doesEnrageGork x && doesEnrageMork x = stomp $ stab x
        | doesEnrageGork x = stab x
        | doesEnrageMork x = stomp x
        | otherwise = x

{- Task -}
class (Enum a, Bounded a, Eq a) => SafeEnum a where
  ssucc :: a -> a
  ssucc x | maxBound == x = minBound | otherwise = succ x
--  ssucc maxBound = minBound
--  ssucc x = succ x

  spred :: a -> a
  spred x | minBound == x = maxBound | otherwise = pred x
--  spred minBound = maxBound
--  spred x = pred x

instance SafeEnum Bool
instance SafeEnum Int

{- Task -}
avg :: Int -> Int -> Int -> Double
avg x y z = sum(map fromIntegral [x,y,z]) / 3

{- Task -}
--bar x y z = x + y
--foo a b = bar a a (a + b)
--value = foo (3 * 10) (5 - 2)

{-
value ~> foo (3 * 10) (5 - 2) ~> bar (3 * 10) (3 * 10) ((3* 10) + (5 - 2)) ~>
(3 * 10) + (3 * 10) ~> 30 + (3 * 10) ~> 30 + 30 ~> 60
-}

{- Task -}
addTwoElements :: a -> a -> [a] -> [a]
addTwoElements x y lst = x : y : lst

{- Task -}
nTimes:: a -> Int -> [a]
nTimes x 0 = []
nTimes x n = x : nTimes x (n-1)

{- Task -}
oddsOnly :: Integral a => [a] -> [a]
oddsOnly xs = [x | x <- xs, odd x]

oddsOnly' :: Integral a => [a] -> [a]
oddsOnly' [] = []
oddsOnly' (x:xs)
  | odd x = x : oddsOnly' xs
  | otherwise = oddsOnly' xs

reverse' :: Integral a => [a] -> [a]
reverse' l = rev l [] where
  rev [] a = a
  rev (x:xs) a = rev xs (x:a)

{- Task -}
isPalindrome :: Eq a => [a] -> Bool
isPalindrome xs = helper xs True where
  helper _ False = False
  helper [] b = b
  helper [_] b = b
  helper xs b = helper (tail $ init xs) (head xs == last xs)

{- Task -}
sum3 :: Num a => [a] -> [a] -> [a] -> [a]
sum3 as bs cs = map sum $ transpose [as, bs, cs]

{- Task -}
groupElems :: Eq a => [a] -> [[a]]
groupElems xs = let
                  helper [] res = res
                  helper (x:xs) [] = helper xs [[x]]
                  helper (x:xs) res
                    | (last $ lres) == x = helper xs (init res ++ [x:lres])
                    | otherwise = helper xs (res ++ [[x]])
                    where lres = last res
                in helper xs []

filter' :: (a -> Bool) -> [a] -> [a]
filter' p [] = []
filter' p (x:xs)
  | p x = x : filter' p xs
  | otherwise = filter' p xs

{- Task -}
readDigits :: String -> (String, String)
readDigits x = span isDigit x

{- Task -}
filterDisj :: (a -> Bool) -> (a -> Bool) -> [a] -> [a]
filterDisj p1 p2 xs = filter (\x -> p1 x || p2 x) xs

{- Task -}
qsort :: Ord a => [a] -> [a]
qsort xs | length xs < 2 = xs
qsort xs = let
              pivot = head xs
              less = [x | x <- tail xs, x <= pivot]
              greater = [x | x <- tail xs, x > pivot]
           in (qsort less) ++ [pivot] ++ (qsort greater)

qsort' :: Ord a => [a] -> [a]
qsort' xs | length xs < 2 = xs
qsort' xs = let
              pivot = head xs
              less = filter (<= pivot) (tail xs)
              greater = filter (> pivot) (tail xs)
           in (qsort' less) ++ [pivot] ++ (qsort' greater)

qsort'' :: Ord a => [a] -> [a]
qsort'' [] = []
qsort'' (x:xs) = (qsort'' $ filter (<=x) xs) ++ [x]  ++ (qsort'' $ filter (>x) xs)

{- Task -}
squares'n'cubes :: Num a => [a] -> [a]
squares'n'cubes xs = concat $ map (\(x,y) -> [x,y]) (zip (map (^2) xs) (map (^3) xs))

squares'n'cubes' :: Num a => [a] -> [a]
squares'n'cubes' = concat . map (\x -> [x^2,x^3])

squares'n'cubes'' :: Num a => [a] -> [a]
squares'n'cubes'' = concatMap (\x -> [x^2,x^3])

{- Task -}
rotations :: Int -> [a] -> [[a]]
rotations len xs = take len (iterate (\(y:ys) -> ys ++ [y]) xs)

perms :: [a] -> [[a]]
perms []        = [[]]
perms il@(x:xs) = concatMap ((rotations len).(x:)) (perms xs)
                  where len = length il

{- Task -}
delAllUpper :: String -> String
delAllUpper = unwords . filter (any isLower) . words

{- Task -}
max3 :: Ord a => [a] -> [a] -> [a] -> [a]
max3 xs ys zs = map (\(x,y,z) -> max x $ max y z) (zip3 xs ys zs)

max3' :: Ord a => [a] -> [a] -> [a] -> [a]
max3' = zipWith3 (\a b c -> a `max` b `max` c)

max3'' :: Ord a => [a] -> [a] -> [a] -> [a]
max3'' = zipWith3 (\x y z -> maximum [x,y,z])

{- Task -}
fibStream :: [Integer]
fibStream = 0 : 1 : zipWith (+) fibStream (tail fibStream)

{- Task -}
data Odd = Odd Integer
  deriving (Eq, Show)

instance Enum Odd where
  succ (Odd x) = Odd (x + 2)
  pred (Odd x) = Odd (x - 2)

  fromEnum (Odd x) = fromInteger $ x `div` 2
  toEnum x = Odd (toInteger x * 2 + 1)

  enumFrom = iterate succ
  enumFromThen (Odd x) (Odd y) = map Odd [x, y ..]
  enumFromTo (Odd x) (Odd y) = map Odd [x, x + 2 .. y]
  enumFromThenTo (Odd x) (Odd y) (Odd z) = map Odd [x , y .. z]

{- Folding: Remember -}
{- foldr (-) x [2,1,5] ~> (2 - (1 - (5 - x))) -}
{- foldl (-) x [2,1,5] ~> (((x - 5) - 1) - 2) -}

{- Task -}
meanList :: [Double] -> Double
meanList = (\(x, y) -> x / y) . foldr (\x (s, l) -> (x + s, 1 + l)) (0, 0)

{- Task -}
evenOnly :: [a] -> [a]
evenOnly = fst . foldl (\(s, p) x -> if even p then (s ++ [x], p + 1) else (s, p + 1))  ([], 1)

-- Endless list
evenOnly' :: [a] -> [a]
evenOnly' = foldr (\(n,x) xs -> if even n then x:xs else xs) [] . zip [1, 2..]

maximum' :: (Ord a) => [a] -> a
maximum' = foldl1 max

{- Task -}
lastElem :: [a] -> a
lastElem = foldl1 (\_ y -> y)

lastElem' :: [a] -> a
lastElem' = foldl1 seq

{- Task -}
unfoldr :: (b -> Maybe (a, b)) -> b -> [a]
unfoldr f ini = helper (f ini) where
  helper (Just (x, ini')) = x : unfoldr f ini'
  helper Nothing          = []

revRange :: (Char,Char) -> [Char]
revRange (a, b) = unfoldr g b
  where g = \x -> if x >= a then Just (x, pred x) else Nothing

{- Task -}
data Color = Red | Green | Blue

instance Show Color where
  show Red = "Red"
  show Green = "Green"
  show Blue = "Blue"

{- Task -}
change :: (Ord a, Num a) => a -> [[a]]
change 0 = [[]]
change amount = [c:cs | c <- [2,3,7], amount >= c, cs <- change (amount - c) ]

{- Task -}
data Point = Point Double Double

origin :: Point
origin = Point 0.0 0.0

distanceToOrigin :: Point -> Double
distanceToOrigin (Point x y) = sqrt (x ^ 2 + y ^ 2)

distance :: Point -> Point -> Double
distance (Point x1 y1) (Point x2 y2)= sqrt ((x2 - x1) ^ 2 + (y2 - y1) ^ 2)

{- Task -}
data Shape = Circle Double | Rectangle Double Double

area :: Shape -> Double
area (Circle r) = pi * r ^ 2
area (Rectangle a b) = a * b

{- Task -}
data Result' = Fail' Int | Success'

instance Show Result' where
    show (Fail' c) = "Fail: " ++ show c
    show Success' = "Success"

--doSomeWork' :: SomeData -> Result'
--doSomeWork' d = case doSomeWork d of
--                    (_, 0) -> Success'
--                    (_, c) -> Fail' c

{- Task -}
square :: Double -> Shape
square a = Rectangle a a

isSquare :: Shape -> Bool
isSquare (Rectangle a b) = a == b
isSquare _ = False

{- Constructor -}
data Person = Person String String Int
  deriving Show

firstName :: Person -> String
firstName (Person f _ _) = f

lastName :: Person -> String
lastName (Person _ l _) = l

age :: Person -> Int
age (Person _ _ a) = a

data Person' = Person' { firstName' :: String, lastName' :: String, age' :: Int}
  deriving (Show, Eq)

john = Person' "John" "Smith" 33
johnAge= john & age'

{- Task -}
timeToString :: UTCTime -> String
timeToString = formatTime defaultTimeLocale "%a %d %T"

data LogLevel = Error | Warning | Info

data LogEntry = LogEntry {timestamp :: UTCTime, logLevel :: LogLevel, message :: String}

logLevelToString :: LogLevel -> String
logLevelToString x = case x of
                        Error   -> "Error"
                        Warning -> "Warning"
                        _    -> "Info"

logEntryToString :: LogEntry -> String
logEntryToString l = timeToString (timestamp l) ++ ": " ++ logLevelToString (logLevel l)  ++ ": " ++ message l

{- Task -}
updateLastName :: Person' -> Person' -> Person'
updateLastName p1 p2 = p2 {lastName' = lastName' p1}

name :: Person' -> String
name person = firstName' person ++ " " ++ lastName' person

name' :: Person' -> String
name' (Person' fn ln _) = fn ++ " " ++ ln

name'' :: Person' -> String
name'' (Person' {lastName' = ln, firstName' = fn}) = fn ++ " " ++ ln

{- Task -}
abbrFirstName :: Person' -> Person'
abbrFirstName person@(Person' {firstName' = fn})
    | length fn > 2 = person {firstName' = (++ ".") . take 1 $ fn}
    | otherwise = person

{- Task -}
data Coord a = Coord a a
  deriving Show

getCenter :: Double -> Coord Int -> Coord Double
getCenter size (Coord x y) = Coord (center x) (center y)
  where center c = fromIntegral c * size + size / 2

getCell :: Double -> Coord Double -> Coord Int
getCell size (Coord x y)  = Coord (point x) (point y)
  where point c = floor $ c / size

roots :: Double -> Double -> Double -> Either [Char] (Double, Double)
roots a b c
  | discr >= 0 = Right (x1, x2)
  | otherwise = Left "Negative discriminant"
  where
      x1 = helper . negate $ d
      x2 = helper d
      helper x = (-b + x) / (2 * a)
      d = sqrt discr
      discr = b ^ 2 - 4 * a * c

{- Task -}
findDigit :: [Char] -> Maybe Char
findDigit xs
  | c /= "" = Just (head c)
  | otherwise = Nothing
  where
    c = filter isDigit xs

findDigit' :: [Char] -> Maybe Char
findDigit' [] = Nothing
findDigit' (x:xs) | isDigit x = Just x
                 | otherwise = findDigit' xs

{- Task -}
findDigitOrX :: [Char] -> Char
findDigitOrX xs = case findDigit xs of
                    Just x -> x
                    Nothing -> 'X'

{- Task -}
maybeToList :: Maybe a -> [a]
maybeToList (Just x) = [x]
maybeToList Nothing = []

maybeToList' :: Maybe a -> [a]
maybeToList' x = case x of
                  Just x -> [x]
                  _ -> []

listToMaybe :: [a] -> Maybe a
listToMaybe [] = Nothing
listToMaybe (x:[]) = Just x

listToMaybe' :: [a] -> Maybe a
listToMaybe' xs = case xs of
                    (x:_) -> Just x
                    _ -> Nothing

eitherToMaybe :: Either a a -> Maybe a
eitherToMaybe (Left a) = Just a
eitherToMaybe (Right _) = Nothing

{- strict and lazy calculations -}
data CoordLazy a = CoordLazy a a
  deriving Show

data CoordStrict a = CoordStrict !a !a
  deriving Show

getXLazy :: CoordLazy a -> a
getXLazy (CoordLazy x _) = x

getXStrict :: CoordStrict a -> a
getXStrict (CoordStrict x _) = x

-- getXLazy 3 undefined - Success
-- getXStrict 3 undefined - Fail

{-
import Data.Complex
import Data.Ratio

data Complex a = !a :+ !a
data Ratio a = !a :% !a
-}

{- Recursive data types -}

-- data [] a = [] | a : ([] a)
data List a = Nil | Cons a (List a)
  deriving Show

yz = Cons 'y' (Cons 'z' Nil)

{- Task -}
fromList :: List a -> [a]
fromList Nil = []
fromList (Cons x y) = x : fromList y

toList :: [a] -> List a
toList [] = Nil
toList (x:xs) = Cons x (toList xs)

{- Task -}
data Nat = Zero | Suc Nat
  deriving Show

nat1 = Suc (Suc Zero)
nat2 = Suc (Suc (Suc Zero))

fromNat :: Nat -> Integer
fromNat Zero = 0
fromNat (Suc n) = fromNat n + 1

toNat :: Integer -> Nat
toNat 0 = Zero
toNat n = Suc (toNat (n-1))

add :: Nat -> Nat -> Nat
add x y = toNat $ fromNat x + fromNat y

mul :: Nat -> Nat -> Nat
mul x y = toNat $ fromNat x * fromNat y

fac :: Nat -> Nat
fac x
  | n >= 0 = toNat $ helper 1 n
  | otherwise = error "Argumnet must be >= 0!"
  where
    n = fromNat x
    helper acc 0 = acc
    helper acc n = helper (acc * n) (n-1)

{- Task -}
data Tree a = Leaf a | Node (Tree a) (Tree a)
  deriving Show

height :: Tree a -> Int
height (Leaf _) = 0
height (Node n1 n2) = 1 + max (height n1) (height n2)

size :: Tree a -> Int
size (Leaf _)  = 1
size (Node n1 n2) = 1 + size n1 + size n2

{- Task -}
avg' :: Tree Int -> Int
avg' t =
    let (s,c) = go t
    in s `div` c
  where
    go :: Tree Int -> (Int,Int)
    go (Leaf n) = (n, 1)
    go (Node n1 n2) = (sum1 + sum2, count1 + count2)
      where
        (sum1, count1) = go n1
        (sum2, count2) = go n2

type AssociativeList k v = [(k, v)]

lookUp :: Eq k => k -> AssociativeList k v -> Maybe v
lookUp _ [] = Nothing
lookUp key ((x,y):xys)
  | key == x = Just y
  | otherwise = lookUp key xys

newtype IntList = IList [Int]
  deriving Show

newtype Identity a = Identity {runIdentity :: a}
  deriving (Eq, Ord, Show)

-- Monoid --
class Monoid a where
  mempty :: a             -- Neutral element
  mappend :: a -> a -> a  -- Operation that introduced on the elements of the monoid

  mconcat :: [a] -> a
  mconcat = foldr mappend mempty

{- Example with List -}
instance Monoid [a] where
  mempty = []
  mappend = (++)

{- Addition implementation for monoid -}
newtype Sum a = Sum {getSum :: a}
  deriving (Eq, Show, Read, Bounded, Ord)

instance Num a => Monoid (Sum a) where
  mempty = Sum 0
  (Sum x) `mappend` (Sum y) = Sum (x + y)

{- Product implementation for monoid -}
newtype Product a = Product {getProduct :: a}
  deriving (Eq, Show, Read, Bounded, Ord)

instance Num a => Monoid (Product a) where
  mempty = Product 1
  (Product x) `mappend` (Product y) = Product (x * y)

{- Task -}
xor :: Bool -> Bool -> Bool
x `xor` y = (x || y) && (not (x && y))

newtype Xor = Xor { getXor :: Bool }
  deriving (Eq, Show)

instance Monoid Xor where
  mempty = Xor False
  (Xor x) `mappend` (Xor y) = Xor (x `xor` y)

{- Monoid Pairs (tuple) -}
instance (Monoid a, Monoid b) => Monoid (a,b) where
  mempty = (mempty, mempty)
  (x1,x2) `mappend` (y1,y2) = (x1 `mappend` y1, x2 `mappend` y2)

{- Example: ("ABC", Product 1) `mappend` ("CDE", Product 3) -}

{- Monoid Maybe -}
instance Monoid a => Monoid (Maybe a) where
  mempty = Nothing
  Nothing `mappend` m = m
  m `mappend` Nothing = m
  Just m1 `mappend` Just m2 = Just (m1 `mappend` m2)

newtype First a = First { getFirst :: Maybe a }
  deriving (Eq, Ord, Read, Show)

instance Monoid (First a) where
  mempty = First Nothing
  First Nothing `mappend` r = r
  l `mappend` _             = l

{- mconcat $ map First [Nothing, Just 3, Just 5] -}
firstConcat = getFirst . mconcat . map First

-- MONAD --

{- Uncomment to check (Data.Functor) -}
--class Functor f where
--  fmap :: (a -> b) -> (f a -> f b)
--
--instance Functor [] where
--  fmap = map
--
--instance Functor Maybe where
--  fmap _ Nothing = Nothing
--  fmap f (Just a) = Just (f a)

{- Task -}
data Point3D a = Point3D a a a
  deriving Show

instance Functor Point3D where
  fmap f (Point3D x1 x2 x3)= Point3D (f x1) (f x2) (f x3)

{- Task -}
data GeomPrimitive a = Point' (Point3D a) | LineSegment (Point3D a) (Point3D a)

instance Functor GeomPrimitive where
  fmap f (Point' p) = Point' (fmap f p)
  fmap f (LineSegment p1 p2) = LineSegment (fmap f p1) (fmap f p2)

{- Binary trees -}
-- SOOO BEAUTIFULL !!!
data Tree' a = Leaf' a | Branch (Tree' a) a (Tree' a)
  deriving Show

testTree = Branch (Leaf' 2) 3 (Branch (Leaf' 4) 5 (Leaf' 6))

instance Functor Tree' where
  fmap g (Leaf' x)       = Leaf' (g x)
  fmap g (Branch l x r) = Branch (fmap g l) (g x) (fmap g r)

{- fmap == (<$>) -}
{- (^2) <$> testTree -}
{- (^2) `fmap` testTree -}

-- Operator (<$) replace all values in data structure

{- Task -}
data Tree'' a = Leaf'' (Maybe a) | Branch' (Tree'' a) (Maybe a) (Tree'' a)
  deriving Show

instance Functor Tree'' where
  fmap g (Leaf'' x)       = case x of
                                Just x  ->  Leaf'' (Just (g x))
                                _       ->  Leaf'' (Nothing)
  fmap g (Branch' l x r)  = case x of
                                Just x  ->  Branch' (fmap g l) (Just (g x)) (fmap g r)
                                _       ->  Branch' (fmap g l) (Nothing) (fmap g r)

--instance Functor ((,), s) where
--  fmap g (x,y) = (x, g y)

-- (a -> b) -> (,) s a -> (,) s b
-- (a -> b) -> (s,a) -> (s,b)
-- fmap succ (1, 'A')

--instance Functor (Either e) where
--  fmap _ (Left x) = Left x
--  fmap g (Right y) = Right (g y)

-- (a -> b) -> Either e a -> Either e b
-- fmap (+3) $ Right 1  ->  Right 4
-- fmap (+3) $ Left 2   ->  Left 2

{- Task -}
data Entry k1 k2 v = Entry (k1, k2) v
  deriving Show

data Map k1 k2 v = Map [Entry k1 k2 v]
  deriving Show

instance Functor (Entry k1 k2) where
  fmap g (Entry (k1, k2) v) = Entry (k1, k2) (g v)

instance Functor (Map k1 k2) where
  fmap g (Map xs) = Map (map (fmap g) xs)

{- Rules for Functors -}
-- 1. fmap id = id
-- 2. fmap (f . g) = fmap f . fmap g

{-
 -
 - Functions with effects
 -
 - f :: a -> Maybe b       -- can sometimes fail (Just b `or` Nothing)
 - f :: a -> [b]           -- can return many results
 - f :: a -> (Either s) b  -- can sometimes end with a typed exception
 - f :: a -> (s,b)         -- can write to the log
 - f :: a -> ((->) e) b    -- can read from external environments
 - f :: a -> (State s) b   -- work with a named state
 - f :: a -> IO a          -- carry out inputs/outputs (files, console ...)
 -
-}

{- Task -}
data Log a = Log [String] a
  deriving Show

add1Log = toLogger (+1) "added one"
mult2Log = toLogger (*2) "multiplied by 2"

toLogger :: (a -> b) -> String -> (a -> Log b)
toLogger f msg = Log [msg] . f

execLoggers :: a -> (a -> Log b) -> (b -> Log c) -> Log c
execLoggers x f g = Log (msg ++ msg') res'
  where
    Log msg res = f x
    Log msg' res' = g res

{-
 - class Monad m where
 -    return :: a -> m a
 -    (>>=) :: m a -> (a -> m b) -> m b   -- bind
 -    ...
 -
 - infixl 1 >>=
 - return True :: [] Bool       ~> [True]
 - return True :: Maybe Bool    ~> [Just True]
-}

toKleisli :: Monad m => (a -> b) -> (a -> m b)
toKleisli f x = return (f x)
-- toKleisli f = \x -> return (f x)
-- toKleisli f = return .

-- toKleisli cos 0 :: [Double]    ~> [1.0]
-- toKleisli cos 0 :: IO Double   ~> 1.0

{- Task -}
returnLog :: a -> Log a
returnLog x = Log [] x

{- Task -}
bindLog :: Log a -> (a -> Log b) -> Log b
bindLog (Log msg x) f =  Log (msg ++ msg') res
  where Log msg' res = f x

{- Task -}
instance Functor Log where
  fmap = liftM

instance Applicative Log where
  pure = return
  (<*>) = ap

instance Monad Log where
    return = returnLog
    (>>=) = bindLog

execLoggersList :: a -> [a -> Log a] -> Log a
execLoggersList x ys = foldl (>>=) (return x) ys

{-
 - class Monad m where
 -    return :: a -> m a
 -    (>>=) :: m a -> (a -> m b) -> m b   -- bind
 -    (>>) :: m a -> m b -> m b   -- bind
 -    x >> y = x >>= \_ -> y
 -    fail :: String -> m a
 -    fail s = error s
-}

-- Control.Monad
(=<<) :: Monad m => (a -> m b) -> m a -> m b
(=<<) = flip (>>=)

-- Fish
(<=<) :: Monad m => (b -> m c) -> (a -> m b) -> (a -> m c)
f <=< g = \x -> g x >>= f

{- Task -}
-- instance Functor SomeType where
--   fmap f x = x >>= (return . f)

-- `do` notation --

-- do { e1 ; e2 }           == e1 >> e2
-- do { p <- e1 ; e2 }      == e1 >>= \p -> e2
-- do { let v = e1 ; e2 }   == let v = e1 in do e2

-- Monad Identity
instance Monad Identity where
  return x = Identity x
  Identity x >>= k = k x

instance Functor Identity where
  fmap  f (Identity x) = Identity (f x)

instance Applicative Identity where
  pure x = Identity x
  Identity f <*> Identity v = Identity (f v)

wrap'n'succ :: Integer -> Identity Integer
wrap'n'succ x = Identity (succ x)

-- Monad Maybe
type Name = String
type Database = [(Name, Name)]

fathers, mothers :: Database
fathers = [("Bill", "John"),
           ("Ann", "John"),
           ("John", "Piter")]

mothers = [("Bill", "Jane"),
           ("Ann", "Jane"),
           ("John", "Alice"),
           ("Jane", "Dorothy"),
           ("Alice", "Mary")]

getM, getF :: Name -> Maybe Name
getM person = lookup person mothers
getF person = lookup person fathers

-- getF "Bill" >>= getM >>= getM
-- do {f <- getF "Bill"; gm <- getM f; getM gm}

granmas :: Name -> Maybe (Name, Name)
granmas person = do
  m <- getM person
  gmm <- getM m
  f <- getF person
  gmf <- getM f
  return (gmm, gmf)

{- Task -}
data Token = Number Int | Plus | Minus | LeftBrace | RightBrace
  deriving (Eq, Show)

asToken :: String -> Maybe Token
asToken x | all isDigit x = Just (Number (read x :: Int))
asToken x = case x of
                "+"       ->  Just Plus
                "-"       ->  Just Minus
                "("       ->  Just LeftBrace
                ")"       ->  Just RightBrace
                _         ->  Nothing

tokenize :: String -> Maybe [Token]
tokenize input = sequence (map asToken (words input))

-- Monad List

-- return 4 :: [Int] ~> [4]
-- [1,2] >>= (\x -> [x,x,x]) ~> [1,1,1,2,2,2]
-- concat $ map (\x -> [x,x,x]) [1,2] ~> [1,1,1,2,2,2]

list = [(x,y) | x <- [1,2,3], y <- [4,5,6]]

list' = do
  x <- [1,2,3]
  y <- [4,5,6]
  return (x,y)

list'' =
  [1,2,3] >>= (\x ->
  [4,5,6] >>= (\y ->
  return (x,y)))

lst = [(x,y) | x <- [1,2,3], y <- [1,2], x /= y]

{- Task -}
-- nextPositionsN :: Board -> Int -> (Board -> Bool) -> [Board]
-- nextPositionsN b 0 pred | pred b = [b]
-- nextPositionsN _ n _    | n <= 0  = []
-- nextPositionsN b n pred = do
--   p <- nextPositions b
--     nextPositionsN p (n-1) pred

lst' = do
  x <- [1,2,3]
  y <- [1,2]
  True <- return (x /= y)
  return (x,y)

lst'' =
  [1,2,3]         >>= (\x ->
  [1,2]           >>= (\y ->
  return (x /= y) >>= (\b ->
  case b of True -> return (x,y)
            _    -> fail "...")))

lst''' = do
  x <- [1,2,3]
  y <- [1,2]
  if x /= y then "E" else []
  return (x,y)

{- Task -}
pythagoreanTriple :: Int -> [(Int, Int, Int)]
pythagoreanTriple x = do
  a <- [1..x]
  b <- [1..x]
  c <- [1..x]
  if a^2 + b^2 == c^2 && a < b then "E" else []
  return (a,b,c)

-- Monad IO

-- type IO a = IO (RealWorld -> (RealWorld, a))
-- return :: a -> IO a ~> return a = \w -> (w,a)
-- (>>=) :: IO a -> (a -> IO b) -> IO b

-- instance Monad IO where
--  return a = \w -> (w,a)
--  (>>=) m k = \w -> case m w of (w',a) -> k a w'

getLine' :: IO String
getLine' = do
  c <- getChar
  if c == '\n' then
    return []
  else do
    cs <- getLine'
    return (c:cs)

putStr' :: String -> IO ()
putStr' [] = return ()
putStr' (x:xs) = putChar x >> putStr' xs

{-
import Control.Monad

sequence_ :: Monad m => [m a] -> m ()
sequence_ = foldr (>>) (return ())

-}

putStr'' :: String -> IO ()
putStr'' = sequence_ . map putChar

{-
mapM_ :: Monad m => (a -> m b) -> [a] -> m ()
mapM_ f = sequence_ , map f
-}

putStr''' :: String -> IO ()
putStr''' = mapM_ putChar


{- |
This module provides normalized versions of the transforms in @fftw@.

All of the transforms are normalized so that

 - Each transform is unitary, i.e., preserves the inner product and the sum-of-squares norm of its input.

 - Each backwards transform is the inverse of the corresponding forwards transform.

(Both conditions only hold approximately, due to floating point precision.)

For more information on the underlying transforms, see
<http://www.fftw.org/fftw3_doc/What-FFTW-Really-Computes.html>.
--
-- @since 0.2
-}

module Numeric.FFT.Vector.Unitary.Multi
  (
        -- * Creating and executing 'Plan's
        run,
        plan,
        execute,
        -- * Complex-to-complex transforms
        dft,
        idft,
        -- * Real-to-complex transforms
        dftR2C,
        dftC2R,
  ) where

import Control.Exception (assert)
import Control.Monad (forM_)
import Numeric.FFT.Vector.Base
import qualified Numeric.FFT.Vector.Unnormalized.Multi as U
import Data.Complex
import qualified Data.Vector.Storable as VS
import qualified Data.Vector.Storable.Mutable as MS
import Control.Monad.Primitive(RealWorld)

-- | A discrete Fourier transform. The output and input sizes are the same (@n@).
--
-- @y_k = (1\/sqrt n) sum_(j=0)^(n-1) x_j e^(-2pi i j k\/n)@
dft :: TransformND (Complex Double) (Complex Double)
dft = U.dft {normalizationND = \ns -> constMultOutput $ 1 / sqrt (toEnum (VS.product ns))}

-- | An inverse discrete Fourier transform.  The output and input sizes are the same (@n@).
--
-- @y_k = (1\/sqrt n) sum_(j=0)^(n-1) x_j e^(2pi i j k\/n)@
idft :: TransformND (Complex Double) (Complex Double)
idft = U.idft {normalizationND = \ns -> constMultOutput $ 1 / sqrt (toEnum (VS.product ns))}

-- | A forward discrete Fourier transform with real data.  If the input size is @n@,
-- the output size will be @n \`div\` 2 + 1@.
dftR2C :: TransformND Double (Complex Double)
dftR2C = U.dftR2C {normalizationND = \ns -> modifyOutput $
                    complexR2CScaling (sqrt 2) ns (outputSizeND U.dftR2C $ VS.last ns)
        }

-- | A normalized backward discrete Fourier transform which is the left inverse of
-- 'U.dftR2C'.  (Specifically, @run dftC2R . run dftR2C == id@.)
--
-- 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 :: TransformND (Complex Double) Double
dftC2R = U.dftC2R {normalizationND = \ns -> modifyInput $
                    complexR2CScaling (sqrt 0.5) ns (inputSizeND U.dftC2R $ VS.last ns)
        }

complexR2CScaling :: Double -> VS.Vector Int -> Int -> MS.MVector RealWorld (Complex Double) -> IO ()
complexR2CScaling !t !ns !len !a = assert (MS.length a == VS.product (VS.init ns) * len) $ do
    let !s1 = sqrt (1/toEnum (VS.product ns))
    let !s2 = t * s1
    -- Justification for the use of unsafeModify:
    -- The output size is 2n+1; so if n>0 then the output size is >=1;
    -- and if n even then the output size is >=3.
    forM_ [0.. VS.product (VS.init ns) - 1] $ \idx -> do
      unsafeModify a (idx * len) $ scaleByD s1
      if odd (VS.last ns)
        then multC s2 (MS.unsafeSlice (idx * len + 1) (len-1) a)
        else do
            unsafeModify a (idx * len + len - 1) $ scaleByD s1
            multC s2 (MS.unsafeSlice (idx * len + 1) (len-2) a)


{-# LANGUAGE RankNTypes #-}

-- |
--  Author: Benedikt Spies
--
--  Utility functions to generate simple Images
--  Image rows are generated in parallel
module Utils.ImageGenerator (generateRainbowImage, generateMandelbrotImage) where

import Codec.Picture
  ( Image,
    Pixel8,
    PixelRGB8 (PixelRGB8),
  )
import Data.Complex (Complex ((:+)), magnitude)
import Utils.ParallelImageUtils (parGenerateImage)

generateRainbowImage :: Image PixelRGB8
generateRainbowImage = parGenerateImage f 256 256
  where
    f :: Int -> Int -> PixelRGB8
    f x y = PixelRGB8 r g b
      where
        r = fromIntegral $ 255 - x
        g = fromIntegral $ 255 - y
        b = fromIntegral x

generateMandelbrotImage :: Int -> Int -> Image PixelRGB8
generateMandelbrotImage width height = parGenerateImage (generateMandelbrotPixel width height) width height

generateMandelbrotPixel :: Int -> Int -> Int -> Int -> PixelRGB8
generateMandelbrotPixel imageWidth imageHeight x y = color $ mandelbrot (x' / w' * 3 - 2) (y' / h' * 3 - 1.5)
  where
    x' = fromIntegral x
    y' = fromIntegral y
    w' = fromIntegral imageWidth
    h' = fromIntegral imageHeight
    maxIter = 300
    color :: Int -> PixelRGB8
    color iter = PixelRGB8 colorVal colorVal colorVal
      where
        colorVal = floor $ 255 - amt * 255 :: Pixel8
        amt = fromIntegral iter / fromIntegral maxIter :: Float
    mandelbrot :: Float -> Float -> Int
    mandelbrot r i =
      length . takeWhile (\z -> magnitude z <= 2)
        . take maxIter
        $ iterate (\z -> z ^ (2 :: Int) + (r :+ i)) 0


{-# LANGUAGE TypeFamilies #-}

import qualified Data.Foldable as F
import Test.HUnit
import Test.Framework
import Test.Framework.Providers.HUnit

import qualified Linear.V2 as V
import Numeric.AD
import Numeric.LinearAlgebra.HMatrix

import Data.SQP

import Debug.Trace

startState :: State2D
startState = vector [(-1), (-1)]

goalState :: State2D
goalState = vector [1, 1]

internalPathLength :: Int
internalPathLength = 10

costMatrix :: Int -> Matrix Double
costMatrix numInternalStates
  | numInternalStates < 1 =
      error "costMatrix can't be generated with fewer than 1 internal states!"
  | otherwise =
      let mainDiag = replicate (2 * numInternalStates) 2
          aboveDiag =
            map (\ix -> ((ix, ix + 2), (-1))) $
              take (2 * (numInternalStates - 1)) [0..]
          belowDiag =
            map (\ix -> ((ix + 2, ix), (-1))) $
              take (2 * (numInternalStates - 1)) [0..]
      in  diagBlock [diagl mainDiag +
            assoc (2 * numInternalStates, 2 * numInternalStates) 0
              (aboveDiag ++ belowDiag), matrix 1 [0]]

cost :: Vector Double -> Double
cost pathV =
  let (internalStates, _) = vectorToPath pathV
      states = startState : internalStates ++ [goalState]
      squaredDists = zipWith (\x x' -> let dx = x - x' in dot dx dx)
        states $ tail states
  in  0.5 * sum squaredDists

approxCost :: Vector Double -> (Matrix Double, Vector Double, Double)
approxCost _ = ( costMatrix internalPathLength
               , vjoin [(vjoin [negate startState, konst 0.0 (2 * (internalPathLength - 1))] +
                 vjoin [konst 0.0 (2 * (internalPathLength - 1)), negate goalState]), vector [0]]
               , 0.5 * (dot startState startState + dot goalState goalState) )

costTest :: Assertion
costTest =
  let pathV = fromList [(-10), 10]
      (costMat, costVec, costAffine) = approxCost pathV
      testApproxCost = 0.5 * (pathV `dot` (costMat #> pathV)) +
        dot pathV costVec + costAffine
  in  (testApproxCost, cost pathV) @?= (400, 400)

type State2D = Vector Double

pathToVector :: [State2D] -> Double -> Vector Double
pathToVector states stepSize = vjoin $ states ++ [vector [stepSize]]

vectorToPath :: Vector Double -> ([State2D], Double)
vectorToPath pathV = ( takesV (replicate internalPathLength 2) pathV
                     , pathV ! (2 * internalPathLength) )

obstacleRadius :: Double
obstacleRadius = 0.5

obstaclePosition :: Vector Double
obstaclePosition = vector [0.001, 0.001]

signedDistance :: State2D -> (Double, Vector Double)
signedDistance p =
  let dp = p - obstaclePosition
      dpNorm = norm_2 dp
      signedDist = dpNorm - obstacleRadius
  in  (signedDist, dp / scalar dpNorm)

ineqs :: Vector Double -> Vector Double
ineqs pathV =
  let (states, stepSize) = vectorToPath pathV
      collisionIneqs = map (negate . fst . signedDistance) states
      boundIneqs = negate stepSize
  in  fromList $ collisionIneqs ++ [boundIneqs]

approxIneqs :: Vector Double -> (Matrix Double, Vector Double)
approxIneqs pathV =
  let (states, _) = vectorToPath pathV
      (dists, normals) = unzip $ map signedDistance states
      assocList =
        concatMap (\(ix, n) -> [ ((ix, 2 * ix), negate $ n ! 0)
                               , ((ix, 2 * ix + 1), negate $ n ! 1)]) $
          zip [0..] normals
      ineqMat = diagBlock
        [ assoc (internalPathLength, 2 * internalPathLength) 0 assocList
        , matrix 1 [(-1)] ]
      affine =
        fromList $ (zipWith (+) (map negate dists) $ zipWith dot states normals)
          ++ [0]
  in  (ineqMat, affine)

unconcat :: [Int] -> [a] -> [[a]]
unconcat [] [] = []
unconcat _ [] = []
unconcat (n : ns) xs = let (xh, xt) = splitAt n xs
                       in  xh : unconcat ns xt

eqsL :: (Mode a, Scalar a ~ Double) => [a] -> [a]
eqsL pathL =
  let stepSize = last pathL
      internalStatesFlat = take (2 * internalPathLength) pathL
      internalStates = unconcat (repeat 2) internalStatesFlat
      states = (map auto $ toList startState) : internalStates ++ [map auto $ toList goalState]
      dist2 s s' = sum $ map (^ (2 :: Int)) $ zipWith (-) s s'
      squaredDists = zipWith dist2 states $ tail states
      stepSizeSqrd = stepSize * stepSize
  in  map (\d2 -> d2 - stepSizeSqrd) squaredDists

eqs :: Vector Double -> Vector Double
eqs pathV =
  vector $ eqsL $ toList pathV

approxEqs :: Vector Double -> (Matrix Double, Vector Double)
approxEqs pathV =
  let f = fromList $ eqsL $ toList pathV
      f' = fromLists $ jacobian eqsL $ toList pathV
  in  ( f', f - (f' #> pathV) )

planningTest :: Assertion
planningTest =
  let problem = Problem
                { _cost = cost
                , _approxCost = approxCost
                , _trueIneqs = ineqs
                , _approxAffineIneqs = approxIneqs
                , _eqs = eqs
                , _approxEqs = approxEqs
                , _numVariables = 2 * internalPathLength + 1
                , _numIneqs = internalPathLength + 1
                , _numEqs = internalPathLength + 1
                }
      initPathStates =
        map (\t -> startState + scalar (fromIntegral t / fromIntegral (internalPathLength + 1)) *
              (goalState - startState))
          [1..internalPathLength]
      initPathStatesPerturbed = map (+ vector [0.0, 0.0001]) initPathStates
      shortestPathLength = norm_2 $ goalState - startState
      initPathVec = pathToVector initPathStatesPerturbed $
       shortestPathLength / fromIntegral (internalPathLength + 1)
      (xResult, fResult) = optimize problem initPathVec
      (xTrue, fTrue) = ((-50), (-50)) -- whatever
  in  (xResult, fResult) @?= (xTrue, fTrue)

main = defaultMain tests
  where tests = [ -- testCase "cost-test" costTest
                -- , testCase "ineq-test" ineqsTest
                 testCase "planning-test" planningTest ]


-- | Dynamic Programming module.
module DP where

import Data.Map (Map)
import qualified Data.Map.Strict as M
import qualified Data.Set as S
import Numeric.LinearAlgebra.Data
       (Matrix, Vector, (!), (¿), (><), maxIndex)
import Numeric.LinearAlgebra.HMatrix as H

import MDP
import Planning

import Debug.Trace

data MDPM s a = MDPM
    { statesm :: Map s Int
    , actionsm :: Map a Int
    , transitionm :: Map a (Matrix Double)
    , rewardm :: Matrix Double
    , gammam :: Double             -- discount
    , sizeStates :: Int           -- S
    , sizeActions :: Int          -- A
    , actionsr :: Map Int a
    }

prepareMDP :: (Ord s, Ord a) => MDP s a -> MDPM s a
prepareMDP mdp =
    MDPM
    { statesm = statesM
    , actionsm = actionsM
    , transitionm = transitionM
    , rewardm = rewardM
    , gammam = gamma mdp
    , sizeStates = S.size (states mdp)
    , sizeActions = S.size (actions mdp)
    , actionsr = M.fromList (zip [0..] (S.toList (actions mdp)))
    }
  where
    statesM = M.fromList (zip (S.toList (states mdp)) [0 ..])
    actionsM = M.fromList (zip (S.toList (actions mdp)) [0 ..])
    transitionM = transitionMap
    transitionMap =
        M.fromList
            (map
                 (\a ->
                       (a, createMatrix a))
                 (S.toList (actions mdp)))
    createMatrix a =
        H.fromLists
            (map
                 (\s ->
                       map (transition mdp a s) (S.toList (states mdp)))
                 (S.toList (states mdp)))
    rewardM = createRewardMatrix
    createRewardMatrix =
        H.fromLists
            (map
                 (\s ->
                       map
                           (\a ->
                                 reward mdp a s)
                           (S.toList (actions mdp)))
                 (S.toList (states mdp)))

valueIterationM :: (Ord a) => MDPM s a -> Int -> Matrix Double -> Matrix Double
valueIterationM mdpm k v
  | k == 0 = v
  | otherwise = valueIterationM mdpm (k - 1) (calc v)
  where
    calc v' =
        rewardm mdpm +
        gammam mdpm `H.scale`
        fromColumns
            (map
                 (\a ->
                       flatten
                           (transitionm mdpm M.! a H.<>
                            (v' ¿ [actionsm mdpm M.! a])))
                 (M.keys (actionsm mdpm)))

valueIteration :: (Ord a, Ord s) => MDP s a -> Int -> s -> a
valueIteration mdp k s = actionsr mdpm M.! maxIndex (vk ! (statesm mdpm M.! s))
  where
    mdpm = prepareMDP mdp
    v0 = (sizeStates mdpm >< sizeActions mdpm) [0,0 ..]
    vk = valueIterationM mdpm k v0


{- |
Description :  Convenience functions to simulate Gamma rate heterogeneity
Copyright   :  (c) Dominik Schrempf 2017
License     :  GPLv3

Maintainer  :  dominik.schrempf@gmail.com
Stability   :  unstable
Portability :  non-portable (not tested)

Sites evolve at different speeds. One way to model this phenomenon is to use
Gamma rate heterogeneity. This module provides the necessary definitions.

* Changelog

-}

module GammaRate where

import           Numeric.Integration.TanhSinh
import           Statistics.Distribution
import           Statistics.Distribution.Gamma

-- | For a given number of rate categories 'n' and a shape parameter 'alpha'
-- (the rate or scale is set such that the mean is 1.0), return a list of rates
-- that represent the respective categories. Use the mean rate for each
-- category.
getMeans :: Int -> Double -> [Double]
getMeans n alpha = means ++ lastMean
  where gamma = gammaDistr alpha (1.0/alpha)
        quantiles = [ quantile gamma (fromIntegral i / fromIntegral n) | i <- [0..n] ]
        -- Calculate the mean rate. Multiplication with the number of rate
        -- categories 'n' is necessary because in each n-quantile the
        -- probability mass is 1/n.
        meanFunc x = fromIntegral n * x * density gamma x
        -- Only calculate the first (n-1) categories with normal integration.
        means = [ intAToB meanFunc (quantiles !! i) (quantiles !! (i+1)) | i <- [0..n-2] ]
        -- The last category has to be calculated with an improper integration.
        lastMean = [intAToInf meanFunc (quantiles !! (n-1))]

-- | The error of integration.
eps :: Double
eps = 1e-6

-- | The integration method to use
method :: (Double -> Double ) -> Double -> Double -> [Result]
method = parSimpson

-- | Helper function for a normal integral from 'a' to 'b'.
intAToB :: (Double -> Double) -> Double -> Double -> Double
intAToB f a b = result . absolute eps $ method f a b

-- | Helper function for an improper integral from 'a' to infinity.
intAToInf :: (Double -> Double) -> Double -> Double
intAToInf f a = (result . absolute eps $ nonNegative method f) - intAToB f 0 a


{-# LANGUAGE BangPatterns #-}
-- file: Astro/Misc.hs
-- Miscellaneous utilities

module Astro.Misc where

import Data.List
import Debug.Trace
import Numeric.LinearAlgebra
import Numeric (showEFloat)
import Text.Printf

{-
 - General math
 -}

type Polynomial = [Double]

-- Evaluate real polynomial given as list of coefficients [a0,a1,...]
evalPoly :: Polynomial->Double->Double
evalPoly p x = sum $ zipWith (*) p $ map (x^) [0..]

-- Derivative of a (real) polynomial
diffPoly :: Polynomial->Polynomial
diffPoly [] = []
diffPoly (a0:[]) = []
diffPoly (a0:as) = zipWith (*) as [1..]

-- Degree of a polynomial
degPoly :: Polynomial->Int
degPoly [] = 0
degPoly (a:[]) = if a /= 0 then 1 else 0
degPoly p = if l == [] then 0 else last l
	where
	l = Data.List.findIndices (/= 0) p

-- Laguerre polynomial root finder
-- Returns Nothing if fails to converge to a real root within maxIt
rootLaguerre (tol,maxIt) p x0  = result
	where
	result = helper x0 0
	p' = diffPoly p
	p'' = diffPoly p'
	deg = fromIntegral $ degPoly p
	helper x n
		| isNaN x = Nothing
		| abs (evalPoly p x) <= tol = Just x
		| n >= maxIt = Nothing -- `debug` ("maxit x,nx:"++show(x,nx))
		| otherwise = helper nx (n+1)
                {-
                        `debug` (printf ("g:n %g h: %g sqa: %g sq: %g x:"
                                ++"%g y: %g") g h sqarg sq nx y)
                                -}
		where
		g = evalPoly p' x / evalPoly p x
		h = g^2 - evalPoly p'' x / evalPoly p x
		sqarg = (deg-1)*(deg*h-g^2)
		sq = sqrt sqarg
		denom = if g>0 then g+sq else g-sq
		nx = x - deg/denom -- `debug` ("n,denom:"++show(n,denom))
                y = evalPoly p nx

{-
 - Vector Math
 -}

-- Standard vector norms
norm :: Vector Double -> Double
--norm v = pnorm PNorm2 v
-- XXX: For 3-vectors this is quite a bit faster
norm !v = sqrt $ v@>0*v@>0 + v@>1*v@>1 + v@>2*v@>2

norm3D :: Vector Double -> Double
norm3D = sqrt . sqrNorm3D

sqrNorm3D :: Vector Double -> Double
sqrNorm3D v = v@>0*v@>0 + v@>1*v@>1 + v@>2*v@>2

-- Normalize vector
unitV v = scale (recip $ norm v) v

-- 3-vector cross product
cross :: Vector Double->Vector Double->Vector Double
cross !a !b =  3|>[a@>1 * b@>2 - a@>2 * b@>1
		,-a@>0 * b@>2 + a@>2 * b@>0
		,a@>0 * b@>1 - a@>1 * b@>0]

-- Rodrigues' rotation formula
rodr :: Vector Double->Double->Vector Double->Vector Double
rodr z th vec = scale (cos th) vec + scale (sin th) (cross z vec)
		+ scale ((z<.>vec)*(1-cos th)) z

{-
 - List and tuple manipulations
 -}

thrd (_,_,a) = a
frth (_,_,_,a) = a

-- Give all pairings, excluding (p,p) and when (p1,p2) excluding (p2,p1)
-- Courtesy of Jussi Knuuttila
chooseTwo :: [a] -> [(a,a)]
chooseTwo as = concatMap pairs (tails as)
    where pairs [] = []
          pairs (a:as) = zip (repeat a) as

-- Take every nth of list
takenth 0 _ = []
takenth _ [] = []
takenth n list =
	iter list []
	where
	iter [] acc = acc
	iter list acc =
		iter (tail newList) (acc ++ [head newList])
		where newList = drop (n-1) list

-- Debugging and printing functions
debug = flip trace
showFloat f = showEFloat (Just 4) f ""


module STFT (stft, istft) where

import Data.Foldable (toList)
import Data.Complex
import Numeric.FFT
import DSP.Window (hanning)

-- Sublists having heads at positions [1,1+noverlap..] of length framesz
split :: Int -> Int -> [a] -> [[a]]
split framesz hop = reverse . split_aux [] where
  split_aux res l
    | length newSeg < framesz = res
    | otherwise = split_aux (newSeg:res) (drop hop l)
    where  newSeg = take framesz l

hannWin :: [Complex Double] ->  [Complex Double]
hannWin l = zipWith (*) hann l
  where hann = map (:+0) $ toList $ hanning $ length l + 1

-- STFT function, calculate X(u,f)
stft :: Int -> Int -> [Complex Double] -> [[Complex Double]]
stft framesz hop = map dft . map hannWin . split framesz hop

-- Overlap function, the list must be no empty
overlap :: Num a => Int -> [[a]] -> [a]
overlap hop = foldr1 overlap_aux where
  overlap_aux seg signal = seg1 ++ (zipWith (+) seg2 signal1) ++ signal2
    where (seg1, seg2) = splitAt hop seg
          noverlap = length seg2
          (signal1, signal2) = splitAt noverlap signal

-- ISTFT function, calculate x(t)
istft :: Int -> [[Complex Double]] -> [Complex Double]
istft hop = overlap hop . map idft


{-# LANGUAGE FlexibleContexts #-}
-----------------------------------------------------------------------------
-- |
-- Module      :  Numeric.Statistics.Information
-- Copyright   :  (c) A. V. H. McPhail 2010, 2014
-- License     :  BSD3
--
-- Maintainer  :  haskell.vivian.mcphail <at> gmail <dot> com
-- Stability   :  provisional
-- Portability :  portable
--
-- Shannon entropy
--
-----------------------------------------------------------------------------

module Numeric.Statistics.Information (
                                   entropy
                                   , mutual_information
                                  ) where

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

import Numeric.Statistics.PDF

import qualified Numeric.LinearAlgebra as LA
--import Numeric.LinearAlgebra.Data hiding(Vector)

--import qualified Data.Vector as DV
import qualified Data.Vector.Storable as V

import Prelude hiding(map,zip)

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

zeroToOne x
    | x == 0.0  = 1.0
    | otherwise = x

logE = V.map (log . zeroToOne)

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

-- | the entropy \sum p_i l\ln{p_i} of a sequence
entropy :: PDF a Double
        => a                       -- ^ the underlying distribution
        -> LA.Vector Double           -- ^ the sequence
        -> Double                  -- ^ the entropy
entropy p x = let ps = probability p x
              in negate $ (LA.dot ps (logE ps))

-- | the mutual information \sum_x \sum_y p(x,y) \ln{\frac{p(x,y)}{p(x)p(y)}}
mutual_information :: (PDF a Double, PDF b (Double,Double))
                   => b                                          -- ^ the underlying distribution
                   -> a                                          -- ^ the first dimension distribution
                   -> a                                          -- ^ the second dimension distribution
                   -> (LA.Vector Double, LA.Vector Double)             -- ^ the sequence
                   -> Double         -- ^ the mutual information
mutual_information p px py (x,y) = let ps = probability p $ V.zipWith (,) x y
                                       xs = probability px x
                                       ys = probability py y
                                   in (LA.dot ps (logE ps - logE (xs*ys)))

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


module Perceptron where

import Prelude hiding ((<>))
import Numeric.LinearAlgebra
import qualified Data.Monoid
import System.IO.Unsafe

orInputs :: Matrix R
orInputs = (4><2)
  [0, 0,
   0, 1,
   1, 0,
   1, 1]

orTargets :: Matrix R
orTargets = col [0, 1, 1, 1]

orInitialWeights :: Matrix R
orInitialWeights = initialWeights (cols orInputs) 1

-- todo: set the initial weights correctly
initialWeights :: Int -> Int -> Matrix R
initialWeights r c = unsafePerformIO $ rand (r + 1) c

addBiasNodes :: Matrix R -> Matrix R
addBiasNodes m = let rs = rows m
                     bias = (rs><1) $ repeat (-1)
                  in m ||| bias

feed :: Matrix R -> Vector R -> Vector R
feed weights = flatten . fwd weights . asColumn

fwd :: Matrix R -> Matrix R -> Matrix R
fwd weights inputs = let withBias = addBiasNodes inputs
                      in step $ withBias <> weights

update :: Double -> Matrix R -> Matrix R -> Matrix R -> Matrix R
update learningRate weights targets inputs  =
  weights - scale learningRate (tr withBias <> (activations - targets))
    where withBias = addBiasNodes inputs
          activations = fwd weights inputs


------------------------------------------------------------
--
-- ShowExample.hs
-- Code sample accompanying topic 1.3.4 "Type classes"
-- See README.md for details
--
-- Fundamentals of Practical Haskell Programming
-- By Richard Cook
--
------------------------------------------------------------

import Data.Complex
import Data.Ratio

main :: IO ()
main = do
    print 5
    print 5.0
    print "five"
    print (5 :+ 6)
    print (5 % 6)

haskell/original/Fundamentals of Practical Haskell Programming [Video]/Volume_01_Section_03_Topic_04/ShowExample.hs


-----------------------------------------------------------------------------
-- |
-- Module     : Numeric.LinearAlgebra.Vector.Statistics
-- Copyright  : Copyright (c) 2010, Patrick Perry <patperry@gmail.com>
-- License    : BSD3
-- Maintainer : Patrick Perry <patperry@gmail.com>
-- Stability  : experimental
--
-- Basic multivariate statistics.
--

module Numeric.LinearAlgebra.Vector.Statistics (
    -- * Immutable interface
    sum,
    mean,
    weightedSum,
    weightedMean,

    -- * Mutable interface
    addSumTo,
    meanTo,
    addWeightedSumTo,
    weightedMeanTo,

    ) where

import Prelude hiding ( sum )

import Control.Monad( forM_ )
import Control.Monad.ST( ST )

import Numeric.LinearAlgebra.Types

import Numeric.LinearAlgebra.Vector.Base( Vector )
import Numeric.LinearAlgebra.Vector.STBase( RVector, STVector )
import qualified Numeric.LinearAlgebra.Vector.STBase as V

-- | Returns the sum of the vectors.  The first argument gives the dimension
-- of the vectors.
sum :: (BLAS1 e) => Int -> [Vector e] -> Vector e
sum p xs = V.create $ do
    s <- V.new_ p
    V.clear s
    addSumTo s xs
    return s

-- | Returns the mean of the vectors.  The first argument gives the dimension
-- of the vectors.
mean :: (BLAS1 e) => Int -> [Vector e] -> Vector e
mean p xs = V.create $ do
      m <- V.new_ p
      meanTo m xs
      return m

-- | Returns the weighted sum of the vectors.  The first argument gives the
-- dimension of the vectors.
weightedSum :: (BLAS1 e) => Int -> [(e, Vector e)] -> Vector e
weightedSum p wxs = V.create $ do
    s <- V.new_ p
    V.clear s
    addWeightedSumTo s wxs
    return s

-- | Returns the weighted mean of the vectors.  The first argument gives the
-- dimension of the vectors.
weightedMean :: (BLAS1 e)
             => Int -> [(Double, Vector e)] -> Vector e
weightedMean p wxs = V.create $ do
    m <- V.new_ p
    weightedMeanTo m wxs
    return m

-- | Adds the sum of the vectors to the target vector.
addSumTo :: (RVector v, BLAS1 e) => STVector s e -> [v e] -> ST s ()
addSumTo dst = addWeightedSumTo dst . zip (repeat 1)

-- | Sets the target vector to the mean of the vectors.
meanTo :: (RVector v, BLAS1 e)
       => STVector s e -> [v e] -> ST s()
meanTo dst = weightedMeanTo dst . zip (repeat 1)

-- | Adds the weigthed sum of the vectors to the target vector.
addWeightedSumTo :: (RVector v, BLAS1 e)
                 => STVector s e ->  [(e, v e)] -> ST s ()
addWeightedSumTo s wxs = do
    n <- V.getDim s
    err <- V.new n 0
    old_s <- V.new_ n
    diff <- V.new_ n
    val <- V.new_ n

    forM_ wxs $ \(w,x) -> do
        V.unsafeCopyTo old_s s -- old_s := s

        V.unsafeCopyTo val x   -- val := w * x
        V.scaleM_ w val

        V.addTo err err val    -- err := err + val
        V.addTo s s err        -- s := s + err

        V.subTo diff old_s s   -- diff := old_s - s
        V.addTo err diff val   -- err := diff + val

-- | Sets the target vector to the weighted mean of the vectors.
weightedMeanTo :: (RVector v, BLAS1 e)
               =>  STVector s e -> [(Double, v e)] -> ST s ()
weightedMeanTo m wxs = let
    go _ _ [] = return ()
    go diff w_sum ((w,x):wxs') | w == 0    = go diff w_sum wxs'
                               | otherwise = let w_sum' = w_sum + w
                                             in do
                                    V.subTo diff x m
                                    V.addWithScaleM_
                                        (realToFrac $ w/w_sum') diff m
                                    go diff w_sum' wxs'
    in do
        n <- V.getDim m
        diff <- V.new_ n
        V.clear m
        go diff 0 wxs