{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeApplications    #-}
module ArrayFire.DataSpec where

import Control.Exception
import Data.Complex
import Data.Word
import Foreign.C.Types
import GHC.Int
import Test.Hspec

import ArrayFire

spec :: Spec
spec =
  describe "Data tests" $ do
    it "Should create constant Array" $ do
      constant @Float [1] 1 `shouldBe` 1
      constant @Double [1] 1 `shouldBe` 1
      constant @Int16 [1] 1 `shouldBe` 1
      constant @Int32 [1] 1 `shouldBe` 1
      constant @Int64 [1] 1 `shouldBe` 1
      constant @Int [1] 1 `shouldBe` 1
      constant @Word16 [1] 1 `shouldBe` 1
      constant @Word32 [1] 1 `shouldBe` 1
      constant @Word64 [1] 1 `shouldBe` 1
      constant @Word [1] 1 `shouldBe` 1
      constant @CBool [1] 1 `shouldBe` 1
      constant @(Complex Double) [1] (1.0 :+ 1.0)
        `shouldBe`
          constant @(Complex Double) [1] (1.0 :+ 1.0)
      constant @(Complex Float) [1] (1.0 :+ 1.0)
        `shouldBe`
          constant @(Complex Float) [1] (1.0 :+ 1.0)


{-# LANGUAGE RankNTypes #-}
import Numeric.ADLens.Lens
import Numeric.ADLens.Basic
import Numeric.ADLens.List
import Numeric.ADLens.HMatrix
import Numeric.ADLens.Accelerate as Acc

import Numeric.LinearAlgebra

import Data.Array.Accelerate as A
import Data.Array.Accelerate.LLVM.Native  as CPU

{-
dostuff = CPU.run $ dotp (use xs) (use ys)

fromList (Z:.10) [0..]
-}

type L1 = Matrix Double
type L2 = Matrix Double
type L3 = Matrix Double

type Input = Vector Double
type Output = Vector Double
type Weights = (L1,(L2,(L3,())))

class TupleSum a where
	tupsum :: a -> a -> a
instance TupleSum () where
	tupsum _ _ = ()
instance (Num a, TupleSum b) => TupleSum (a,b) where
	tupsum (a,x) (b,y) = (a + b, tupsum x y)

-- A dense relu neural network example
swaplayer :: Lens' ((Matrix t, b), Vector t) (b, (Matrix t, Vector t))
swaplayer = first' swap' . assoc'

mmultlayer :: Numeric t => Lens' (b, (Matrix t, Vector t)) (b, Vector t)
mmultlayer = second' mdot'

relulayer :: Lens' (b, Vector Double) (b, Vector Double)
relulayer = second' $ cmap' relu'

uselayer :: Lens' ((Matrix Double, b), Vector Double) (b, Vector Double)
uselayer = swaplayer . mmultlayer . relulayer

runNetwork :: Lens' (Weights, Input) ((), Output)
runNetwork =  uselayer . uselayer . uselayer

main :: IO ()
main = do
   putStrLn "Starting Tests"
   print $ grad (pow' 2) 1
   print $ grad (pow' 4) 1
   print $ grad (map' (pow' 2) . sum') $ [1 .. 5]
   print $ grad (map' (pow' 4) . sum') $ [1 .. 5]
   print $ map (\x -> 4 * x ^ 3 )  [1 .. 5]
   l1 <- randn 3 4
   l2 <- randn 2 3
   l3 <- randn 1 2
   let weights = (l1,(l2,(l3,())))
   print $ view runNetwork (weights, vector [1,2,3,4])
   putStrLn "The neural network gradients"
   print $ set runNetwork (weights, vector [1,2,3,4]) ((), vector [1])


module RayLib where

import Numeric.LinearAlgebra
import qualified Graphics.Image as G
import qualified Debug.Trace as Debug

data Ray = Ray {rayOrigin :: Vector Double,
                rayDirection :: Vector Double}
    deriving (Show, Eq)

eps :: Double
eps = 0.000001

data Isect = Isect {iT :: Double,
                    iNormal :: Vector Double,
                    iObject :: Geometry,
                    iMaterial :: Material}

instance Eq Isect where
    (Isect iT1 iNormal1 _ _) == (Isect iT2 iNormal2 _ _) = (iT1 == iT2) && (iNormal1 == iNormal2)

instance Ord Isect where
    compare (Isect iT1 _ _ _) (Isect iT2 _ _ _) = compare iT1 iT2

instance Show Isect where
    show (Isect t n _ _) = "time: " ++ show t ++ ", normal: " ++ show n

class At a where
    at :: a -> Ray -> Vector Double

instance At Double where
    at t (Ray p d) = p + (fromList [t] :: Vector Double) * d

instance At Isect where
    at (Isect t _ _ _) (Ray p d) = p + (fromList [t] :: Vector Double) * d

type Color = Vector Double

data MatParam = MatParam Color
    deriving (Show, Read, Eq)

data Material = Material {ka :: MatParam,
                          ks :: MatParam,
                          kd :: MatParam,
                          ke :: MatParam,
                          kr :: MatParam,
                          kt :: MatParam,
                          shininess :: Double}
    deriving (Show, Read, Eq)

data Light =    PointLight       {lColor :: Color,
                                  lPosition :: Vector Double,
                                  constantTerm :: Double,
                                  linearTerm :: Double,
                                  quadraticTerm :: Double}
            |   DirectionalLight {lColor :: Color,
                                  lOrientation :: Vector Double}
    deriving (Show, Read, Eq)

type Lights = [Light]

lightColor :: Light -> Color
lightColor (PointLight c _ _ _ _) = c
lightColor (DirectionalLight c _) = c

getDirection :: Light -> Vector Double -> Vector Double
getDirection (PointLight _ pos _ _ _) p = normalize (pos - p)
getDirection (DirectionalLight _ o) _ = -o

distanceAttenuation :: Light -> Vector Double -> Double
distanceAttenuation (PointLight _ pos c b a) p = min 1.0 atten
    where atten = (c * dist * dist + b * dist + a) ** (-1.0)
          dist = norm_2 (p - pos)
distanceAttenuation (DirectionalLight _ _) p = 1.0

shadowAttenuation :: Scene -> Light -> Vector Double -> Double
shadowAttenuation scene (PointLight _ pos _ _ _) p =
    let shadowDir = normalize (pos - p) in
    let ray = Ray p shadowDir in
    let isect = intersect scene ray in
    case isect of
        Nothing -> 1
        Just (Isect t _ _ _) ->
            if (t < 0) || (t > norm_2 (pos - p)) then
                1
            else
                0
shadowAttenuation scene dl@(DirectionalLight _ _) p =
    let shadowDir = normalize (getDirection dl p) in
    let ray = Ray p shadowDir in
    let isect = intersect scene ray in
    case isect of
        Nothing -> 1
        Just (Isect t _ _ _) ->
            if t < 0 then
                1
            else
                0

data Geometry = Sphere {sphereCenter :: Vector Double,
                        sphereRadius :: Double,
                        sphereMaterial :: Material}
    deriving (Show, Read, Eq)

type Objects = [Geometry]

data Scene = Scene {scLights :: Lights,
                    scObjects :: Objects,
                    scAmbient :: Color}
    deriving (Show, Read, Eq)

data Camera = Camera {eye :: Vector Double,
                      viewDir :: Vector Double,
                      upDir :: Vector Double,
                      aspect :: Double,
                      fov :: Double}
    deriving (Show, Read, Eq)

rayThrough :: Camera -> Double -> Double -> Ray
rayThrough (Camera eye viewDir upDir aspect fov) xx yy =
    let z = -viewDir in
    let y = upDir in
    let x = cross y z in
    let m = fromRows [x,y,z] in
    let fov' = fov * pi / 180.0 in
    let normalizedHeight = 2 * tan (fov' / 2.0) in
    let u = (flatten (m <> (fromColumns [(fromList [1,0,0])]))) * (fromList [normalizedHeight * aspect]) in
    let v = (flatten (m <> (fromColumns [(fromList [0,1,0])]))) * (fromList [normalizedHeight]) in
    let look = (flatten (m <> (fromColumns [(fromList [0,0,-1])]))) in
    let x' = fromList [xx - 0.5] in
    let y' = fromList [yy - 0.5] in
    let dir = normalize (look + x' * u + y' * v) in
    Ray eye dir

intersectObj :: Ray -> Geometry -> Maybe Isect
intersectObj r@(Ray rayPos rayDir) s@(Sphere center radius material) =
    let v = center - rayPos in
    let b = v <.> rayDir in
    let discriminant = b * b - v <.> v + radius * radius in
    if discriminant < 0.0 then
        Nothing
    else
        let det = sqrt discriminant in
        let (t1,t2) = (b - det, b + det) in
        if t2 < 0.0 then
            Nothing
        else if t1 > 0.0 then
            pure (Isect t1 (normalize (at t1 r)) s material)
        else
            pure (Isect t2 (normalize (at t2 r)) s material)

intersect :: Scene -> Ray -> Maybe Isect
intersect scene@(Scene _ objects _) r =
    let isects = filter (/= Nothing) . map (intersectObj r) $ objects in
    if length isects == 0 then
        Nothing
    else
        minimum isects

perLightDiffuse :: Light -> Vector Double -> Vector Double -> Vector Double -> Double
perLightDiffuse light isectPoint rayDir normal =
    let lightDir = getDirection light isectPoint in
    let ddot = lightDir <.> (if (rayDir <.> normal > 0) then -normal else normal) in
    let dmax = max ddot 0.0 in
    dmax

perLightSpecular :: Light -> Vector Double -> Vector Double -> Vector Double -> Double
perLightSpecular light isectPoint rayDir normal =
    let lightDir = getDirection light isectPoint in
    let reflect = lightDir - (fromList [2 * (lightDir <.> normal)]) * normal in
    if norm_2 reflect == 0 then
        0
    else
        let reflectNorm = normalize reflect in
        let sdot = rayDir <.> reflect in
        let smax = max sdot 0.0 in
        smax

combineLightTerms :: Scene -> Vector Double -> Vector Double -> Vector Double -> Double -> Light -> Color
combineLightTerms scene isectPoint rayDir normal alpha light =
    let diffuse = perLightDiffuse light isectPoint rayDir normal in
    let specular = perLightSpecular light isectPoint rayDir normal in
    let color = lightColor light in
    let atten = (distanceAttenuation light isectPoint) * (shadowAttenuation scene light isectPoint) in
    color * (fromList [(diffuse + specular) * atten])

shade :: Scene -> Ray -> Isect -> Color
shade scene@(Scene lights objects ambient) r@(Ray p d) i@(Isect t normal _ material) =
    let (Material kaM ksM kdM keM krM ktM shiny) = material in
    let (MatParam ke) = keM in
    let emissive = ke in
    let (MatParam ka) = kaM in
    let amb = ka * ambient in
    let isectPoint = at t r in
    let alpha = shiny in
    emissive + amb + (sum . map (combineLightTerms scene isectPoint d normal alpha) $ lights)

data Screen = Screen {screenWidth :: Int,
                      screenHeight :: Int}
    deriving (Show, Read, Eq)

clamp :: Double -> Double -> Vector Double -> Vector Double
clamp a b = fromList . map (\x -> max a (min b x)) . toList

trace :: Scene -> Camera -> Double -> Double -> Color
trace scene camera x y =
    let ray = rayThrough camera x y in
    let color = traceRay scene ray 0 0 in
    clamp 0 1 color

tracePixel :: Scene -> Screen -> Camera -> Int -> Int -> Color
tracePixel scene (Screen width height) camera i j =
    let x = (fromIntegral i :: Double) / (fromIntegral width :: Double) in
    let y = (fromIntegral j :: Double) / (fromIntegral height :: Double) in
    let color = trace scene camera x y in
    (fromList [255.0]) * color

reflection :: Scene -> Ray -> Isect -> Int -> Color
reflection scene ray@(Ray p d) isect@(Isect t n _ mat) depth =
    let reflectR = d - (fromList [2 * d <.> n]) * n in
    if reflectR == 0 then
        fromList ([0,0,0])
    else
        let isectPointApprox = at (t - eps) ray in
        let reflRay = Ray isectPointApprox (normalize reflectR) in
        let (Material _ _ _ _ (MatParam kr)  _ _) = mat in
        kr * (traceRay scene reflRay (depth - 1) 0)

traceRay :: Scene -> Ray -> Int -> Double -> Color
traceRay scene ray@(Ray p d) depth t =
    let mIsect = intersect scene ray in
    case mIsect of
        Nothing -> fromList[0,0,0]
        (Just isect) -> let colorC = shade scene ray isect in
                        if depth == 0 then
                            colorC
                        else
                            let colorRefl = reflection scene ray isect depth in
                            colorC + colorRefl

type Image = ([[Color]], Int, Int)

traceImage :: Scene -> Screen -> Camera -> Image
traceImage scene screen@(Screen width height) camera =
    let (w, h) = (width - 1, height - 1) in
    let img = [[tracePixel scene screen camera i j | j <- [0..h]] | i <- [0..w]] in
    (img, width, height)

rgbify :: [Double] -> G.Pixel G.RGB Double
rgbify vec = G.PixelRGB (r / 255.0) (g / 255.0) (b / 255.0)
    where r = vec !! 0
          g = vec !! 1
          b = vec !! 2

render :: Image -> G.Image G.VU G.RGB Double
render (image, width, height) = G.fromLists . format $ image
    where format = map (map rgbify) . map (map toList)


{-# LANGUAGE FlexibleContexts #-}

------------------------------------------------
-- |
-- Module    :  Numeric.Matrix.Integral.HNFLLL
-- Copyright :  (c) Jun Yoshida 2019
-- License   :  BSD3
--
-- Compute Hermite normal forms in the LLL-based method.
-- The original "pseudo-code" is found in the paper
--   George Havas, Bohdan S. Majewski & Keith R. Matthews (1998) Extended GCD and Hermite Normal Form Algorithms via Lattice Basis Reduction, Experimental Mathematics, 7:2, 125-136, DOI: 10.1080/10586458.1998.10504362
--
------------------------------------------------

module Numeric.Matrix.Integral.HNFLLL (
  hnfLLLST,
  hnfLLL,
  ) where

import Control.Monad (when, forM_)
import Control.Monad.ST
import Control.Monad.Loops (whileM_)

import Data.STRef

import qualified Data.Maybe as M

import Data.Int

import Numeric.LinearAlgebra ((!))
import qualified Numeric.LinearAlgebra as LA
import Numeric.LinearAlgebra.Devel

-- | Type class which guarantees the type can be embedded in the real number field
class (Num t,Eq t,Ord t) => Realizable t where
  toReal :: t -> Double

instance Realizable Double where
  toReal = id

instance Realizable Int64 where
  toReal = fromIntegral

-- | Type class to qualify a type to be used in LLL-based algorithm
class (Realizable t, LA.Indexable (LA.Vector t) t, LA.Container LA.Vector t) => LLLQualified t

instance LLLQualified Double

instance LLLQualified Int64

-- | Data Type to carry resources in the computation of LLL-based algorithm
data HNFLLLResource s t = HNFLLL {
  matA :: STMatrix s t,
  matU :: STMatrix s t,
  lambda :: STMatrix s Double,
  focusC1 :: STRef s (Maybe Int),
  focusC2 :: STRef s (Maybe Int) }

-- | The function which is called Reduce2 in the original "pseudo-code."
reduceLLL :: (LLLQualified t) => Int -> Int -> HNFLLLResource s t -> ST s ()
reduceLLL k i resRef = do
  ma <- freezeMatrix (matA resRef)
  case M.listToMaybe (LA.find (/=0) (ma!i)) of
    (Just j) -> do
      writeSTRef (focusC1 resRef) (Just j)
      when ( (ma!i!j) < 0 ) $ do
        rowOper (SCAL (-1) (Row i) AllCols) (matA resRef)
        rowOper (SCAL (-1) (Row i) AllCols) (matU resRef)
        rowOper (SCAL (-1) (Row (j+1)) (ColRange 0 j)) (lambda resRef)
        rowOper (SCAL (-1) (FromRow (j+2)) (Col (j+1))) (lambda resRef)
    Nothing -> writeSTRef (focusC1 resRef) Nothing
  case M.listToMaybe (LA.find (/=0) (ma!k)) of
    (Just j) -> writeSTRef (focusC2 resRef) (Just j)
    Nothing -> writeSTRef (focusC2 resRef) Nothing
  c1 <- readSTRef (focusC1 resRef)
  q <- fromIntegral.floor <$> do
    case c1 of
      (Just j) ->
        (/) <$> (toReal <$> readMatrix (matA resRef) k j) <*> (toReal <$> readMatrix (matA resRef) i j)
      Nothing -> do
        lki <- readMatrix (lambda resRef) (k+1) (i+1)
        di <- readMatrix (lambda resRef) (i+1) (i+1)
        if 2 * abs lki > di
          then return $ fromIntegral (round (lki / di))
          else return 0
  when (q /= 0) $ do
    rowOper (AXPY (negate q) i k AllCols) (matA resRef)
    rowOper (AXPY (negate q) i k AllCols) (matU resRef)
    rowOper (AXPY (toReal $ negate q) (i+1) (k+1) AllCols) (lambda resRef)

-- | The function which is called Swap1 in the original "pseudo-code."
swapLLL :: (Num t, LA.Element t) => Int -> HNFLLLResource s t -> ST s ()
swapLLL k resRef = do
  rowOper (SWAP k (k-1) AllCols) (matA resRef)
  rowOper (SWAP k (k-1) AllCols) (matU resRef)
  rowOper (SWAP k (k+1) (ColRange 0 (k-1))) (lambda resRef)
  ml <- freezeMatrix (lambda resRef)
  forM_ [(k+2)..(LA.cols ml-1)] $ \i -> do
    writeMatrix (lambda resRef) i k $
      ((ml!i!k)*(ml!(k+1)!k) + (ml!i!(k+1))*(ml!(k-1)!(k-1)))/(ml!k!k)
    writeMatrix (lambda resRef) i (k+1) $
      ((ml!i!k)*(ml!(k+1)!(k+1)) + (ml!i!(k+1))*(ml!(k+1)!k))/(ml!k!k)
  writeMatrix (lambda resRef) k k $
    ((ml!(k+1)!k)*(ml!(k+1)!k) + (ml!(k+1)!(k+1))*(ml!(k-1)!(k-1)))/(ml!k!k)

-- | Recognize whether a swap should occur or not.
shouldSwapST :: (Num t, LA.Element t) => Int -> HNFLLLResource s t -> ST s Bool
shouldSwapST k resRef = do
  ml <- freezeMatrix $ lambda resRef
  mc1 <- readSTRef $ focusC1 resRef
  mc2 <- readSTRef $ focusC2 resRef
  case (mc1,mc2) of
    (Just c1,Just c2) -> return $ c1 <= c2
    (Just _,Nothing) -> return True
    (Nothing,Just _) -> return False
    (Nothing,Nothing) ->
      return $ (ml!(k+1)!k)*(ml!(k+1)!k) + (ml!(k+1)!(k+1))*(ml!(k-1)!(k-1)) < (ml!k!k)*(ml!k!k)

-- | ST to compute the Hermite normal form for the matrix in stMatA.
-- The first argument is also subject to the operation that is executed on stMatA, so one gets the transformation unimodular matrix by passing the identity.
-- See the source of hnfLLL function for an example of usage.
hnfLLLST :: (LLLQualified t) => STMatrix s t -> STMatrix s t -> ST s ()
hnfLLLST stMatU stMatA = do
  m <- LA.rows <$> freezeMatrix stMatA
  newLambda <- thawMatrix $ LA.ident (m+1)
  newC1 <- newSTRef Nothing
  newC2 <- newSTRef Nothing
  let resRef = HNFLLL stMatA stMatU newLambda newC1 newC2
  kRef <- newSTRef (1 :: Int)
  whileM_ ((<m) <$> readSTRef kRef) $ do
    k <- readSTRef kRef
    reduceLLL k (k-1) resRef
    shouldSwap <- shouldSwapST k resRef
    if shouldSwap
      then swapLLL k resRef >> when (k>1) (writeSTRef kRef (k-1))
      else (forM_ [0..k-2] $ \i -> reduceLLL k i resRef) >> writeSTRef kRef (k+1)
  setMatrix stMatA 0 0 =<< (LA.flipud <$> freezeMatrix stMatA)
  setMatrix stMatU 0 0 =<< (LA.flipud <$> freezeMatrix stMatU)

-- | Execute the LLL-based algorithm on a given matrix.
hnfLLL :: LLLQualified t => LA.Matrix t -> (LA.Matrix t, LA.Matrix t)
hnfLLL mtx = runST $ do
  stMatA <- thawMatrix mtx
  stMatU <- thawMatrix $ LA.ident (LA.rows mtx)
  hnfLLLST stMatU stMatA
  h <- freezeMatrix stMatA
  u <- freezeMatrix stMatU
  return (u,h)


{-# LANGUAGE TypeFamilies, BangPatterns, TypeOperators, FlexibleContexts, FlexibleInstances, ScopedTypeVariables #-}
module Main where

import Data.NeuralNetwork
import Data.NeuralNetwork.Backend.HMatrix
import Numeric.LinearAlgebra
import Numeric.LinearAlgebra.Devel
import Parser

import Data.List(foldl',partition,maximumBy)
import Text.Printf (printf)
import Control.Monad
import Control.Monad.Except
import Text.PrettyPrint.Free hiding (flatten, (<>))
import System.Environment

import Data.Constraint
foo :: Backend b s => b -> s -> Inp (ConvertFromSpec b s) -> Env b (Out (ConvertFromSpec b s))
foo b s i = case witness b s of Dict -> compile b s >>= run . flip forward i

main = do x <- runExceptT $ do
                cnn <- compile ByHmatrix (In2D 28 28 :++ Reshape2DAs1D :++ FullConnect 256 :++ FullConnect 32 :++ FullConnect 10)
                debug cnn -- dotrain cnn >>= dotest
          case x of
            Left  _ -> putStrLn "Error."
            Right _ -> putStrLn "Done."

debug :: (Component n, Inp n ~ Image, Out n ~ Label, MonadIO m, RunInEnv (Run n) m)
      => n -> m ()
debug nn = do
    a0:a1:_ <- liftIO $ getArgs
    let cycle = read a0 :: Int
        rate  = read a1 :: Float
    liftIO $ putStrLn "Load training data."
    dataset <- liftIO $ uncurry zip <$> trainingData
    nn <- iterateM cycle (online rate dataset) nn
    liftIO $ putStrLn "Load test data."
    testset <- liftIO $ (uncurry zip <$> testData)
    flip mapM_ (take 10 testset) $ \(ds,ev) -> do
      pv <- run $ forward nn ds
      liftIO $ putStrLn ("+" ++ showPretty (text (printf "%02d:" $ postprocess pv) <+> pretty pv))
      liftIO $ putStrLn ("*" ++ showPretty (text (printf "%02d:" $ postprocess ev) <+> pretty ev))

dotrain :: (Component n, Inp n ~ Image, Out n ~ Label, MonadIO m, RunInEnv (Run n) m)
       => n -> m n
dotrain nn = do
    liftIO $ putStrLn "Load training data."
    dataset <- liftIO $ uncurry zip <$> trainingData
    liftIO $ putStrLn "Learning."
    iterateM 5 (online 0.0001 dataset) nn

dotest :: (Component n, Inp n ~ Image, Out n ~ Label, MonadIO m, RunInEnv (Run n) m)
       => n -> m ()
dotest !nn = do
    liftIO $ putStrLn "Load test data."
    testset <- liftIO $ (uncurry zip <$> testData)
    liftIO $ putStrLn "Start test"
    result  <- mapM ((postprocess <$>) . run . forward nn . fst) testset
    let expect = map (postprocess . snd) testset
        (co,wr)= partition (uncurry (==)) $ zip result expect
    liftIO $ putStrLn $ printf "correct: %d, wrong: %d" (length co) (length wr)

online :: (Component n, Inp n ~ Image, Out n ~ Label, MonadIO m, RunInEnv (Run n) m)
       => Float -> [(Inp n, Out n)] -> n -> m n
online rate ds !nn = walk ds nn
  where
    walk []     !nn = return nn
    walk (d:ds) !nn = do !nn <- run $ learn outcost' rate nn d
                         walk ds nn
    outcost' a b = return $ zipVectorWith cost' a b

iterateM :: (MonadIO m) => Int -> (a -> m a) -> a -> m a
iterateM n f x = walk 0 x
  where
    walk !i !a | i == n    = return a
               | otherwise = do -- when (i `mod` 10 == 0) $ putStrLn ("Iteration " ++ show i)
                                liftIO $ putStrLn ("Iteration " ++ show i)
                                a <- f a
                                walk (i+1) a

postprocess :: Vector Float -> Int
postprocess = fst . maximumBy cmp . zip [0..] . toList
  where cmp a b = compare (snd a) (snd b)

instance Pretty (Vector Float) where
  pretty vec = encloseSep langle rangle comma $ map (text . printf "%.04f") $ toList vec

showPretty x = displayS (renderPretty 0.4 500 x) ""


{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE InstanceSigs #-}
{-# LANGUAGE Rank2Types #-}
{-# LANGUAGE FlexibleInstances    #-}
{-# LANGUAGE UndecidableInstances #-}
{-# LANGUAGE IncoherentInstances  #-}
{-# LANGUAGE StandaloneDeriving  #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE ForeignFunctionInterface #-}
{-# OPTIONS_GHC -fno-warn-orphans #-}

module Data.Vector.Image.Matrix (
  toMatrix
, fromMatrix
, toMatrixList
, fromMatrixList
, toMatrixD
, fromMatrixD
) where

import Prelude hiding(map)
import qualified Prelude as P
import qualified Data.List as L
import Data.Vector.Storable hiding(forM_,map,replicate,fromList,toList,(!?),(!))
import Data.Maybe
import Numeric.LinearAlgebra.Data hiding (fromList,toList,(!))
import Numeric.LinearAlgebra.HMatrix hiding (fromList,toList,(!))
import qualified Data.Vector.Storable as V
import Foreign
import GHC.Real
import Foreign.C.String
import Foreign.C.Types
import qualified Control.Monad as C

import Data.Vector.Image.Color.RGB
import Data.Vector.Image.Image

toMatrix :: (Storable a)
        => Image a
        -> Matrix a
toMatrix (Image w _h v) = reshape w v

fromMatrix :: (Element a)
        => Matrix a
        -> Image a
fromMatrix mat = Image (cols mat) (rows mat) (flatten mat)

toMatrixList :: (ColorSpace c a,Storable (c a),Storable a,Element a)
              => Image (c a)
              -> [Matrix a]
toMatrixList img = P.map toMatrix $ splitColor img

fromMatrixList :: (ColorSpace c a,Storable (c a),Storable a,Element a)
              => [Matrix a]
              -> Image (c a)
fromMatrixList mats = mergeColor $ P.map fromMatrix mats

toMatrixD :: (ColorSpace c a,Storable (c a),Storable a,Element a,Integral a)
         => Image (c a)
         -> [Matrix Double]
toMatrixD img = P.map toMatrix $ P.map toDouble $ splitColor img

fromMatrixD :: (ColorSpace c a,Storable (c a),Storable a,Element a,Integral a)
           => [Matrix Double]
           -> Image (c a)
fromMatrixD mats = mergeColor $ P.map fromDouble $ P.map fromMatrix mats


-- |
-- Module    : Numeric.SpecFunctions.Extra
-- Copyright : (c) 2009, 2011 Bryan O'Sullivan
-- License   : BSD3
--
-- Maintainer  : bos@serpentine.com
-- Stability   : experimental
-- Portability : portable
--
-- Less common mathematical functions.
module Numeric.SpecFunctions.Extra (
    bd0
  , chooseExact
  , logChooseFast
  , logGammaAS245
  , logGammaCorrection
  ) where

import Numeric.MathFunctions.Constants (m_NaN,m_pos_inf)
import Numeric.SpecFunctions.Internal  (chooseExact,logChooseFast,logGammaCorrection)

-- | Evaluate the deviance term @x log(x/np) + np - x@.
bd0 :: Double                   -- ^ @x@
    -> Double                   -- ^ @np@
    -> Double
bd0 x np
  | isInfinite x || isInfinite np || np == 0 = m_NaN
  | abs x_np >= 0.1*(x+np)                   = x * log (x/np) - x_np
  | otherwise                                = loop 1 (ej0*vv) s0
  where
    x_np = x - np
    v    = x_np / (x+np)
    s0   = x_np * v
    ej0  = 2*x*v
    vv   = v*v
    loop j ej s = case s + ej/(2*j+1) of
                    s' | s' == s   -> s'  -- FIXME: Comparing Doubles for equality!
                       | otherwise -> loop (j+1) (ej*vv) s'

-- | Compute the logarithm of the gamma function Γ(/x/).  Uses
-- Algorithm AS 245 by Macleod.
--
-- Gives an accuracy of 10-12 significant decimal digits, except
-- for small regions around /x/ = 1 and /x/ = 2, where the function
-- goes to zero.  For greater accuracy, use 'logGammaL'.
--
-- Returns ∞ if the input is outside of the range (0 < /x/ ≤ 1e305).
logGammaAS245 :: Double -> Double
-- Adapted from http://people.sc.fsu.edu/~burkardt/f_src/asa245/asa245.html
logGammaAS245 x
    | x <= 0    = m_pos_inf
    -- Handle positive infinity. logGamma overflows before 1e308 so
    -- it's safe
    | x > 1e308 = m_pos_inf
    -- Normal cases
    | x < 1.5   = a + c *
                  ((((r1_4 * b + r1_3) * b + r1_2) * b + r1_1) * b + r1_0) /
                  ((((b + r1_8) * b + r1_7) * b + r1_6) * b + r1_5)
    | x < 4     = (x - 2) *
                  ((((r2_4 * x + r2_3) * x + r2_2) * x + r2_1) * x + r2_0) /
                  ((((x + r2_8) * x + r2_7) * x + r2_6) * x + r2_5)
    | x < 12    = ((((r3_4 * x + r3_3) * x + r3_2) * x + r3_1) * x + r3_0) /
                  ((((x + r3_8) * x + r3_7) * x + r3_6) * x + r3_5)
    | x > 3e6   = k
    | otherwise = k + x1 *
                  ((r4_2 * x2 + r4_1) * x2 + r4_0) /
                  ((x2 + r4_4) * x2 + r4_3)
  where
    (a , b , c)
        | x < 0.5   = (-y , x + 1 , x)
        | otherwise = (0  , x     , x - 1)

    y      = log x
    k      = x * (y-1) - 0.5 * y + alr2pi
    alr2pi = 0.918938533204673

    x1 = 1 / x
    x2 = x1 * x1

    r1_0 =  -2.66685511495;   r1_1 =  -24.4387534237;    r1_2 = -21.9698958928
    r1_3 =  11.1667541262;    r1_4 =    3.13060547623;   r1_5 =   0.607771387771
    r1_6 =  11.9400905721;    r1_7 =   31.4690115749;    r1_8 =  15.2346874070

    r2_0 = -78.3359299449;    r2_1 = -142.046296688;     r2_2 = 137.519416416
    r2_3 =  78.6994924154;    r2_4 =    4.16438922228;   r2_5 =  47.0668766060
    r2_6 = 313.399215894;     r2_7 =  263.505074721;     r2_8 =  43.3400022514

    r3_0 =  -2.12159572323e5; r3_1 =    2.30661510616e5; r3_2 =   2.74647644705e4
    r3_3 =  -4.02621119975e4; r3_4 =   -2.29660729780e3; r3_5 =  -1.16328495004e5
    r3_6 =  -1.46025937511e5; r3_7 =   -2.42357409629e4; r3_8 =  -5.70691009324e2

    r4_0 = 0.279195317918525;  r4_1 = 0.4917317610505968;
    r4_2 = 0.0692910599291889; r4_3 = 3.350343815022304
    r4_4 = 6.012459259764103


module MachineLearning.RandomTest
(
  tests
)

where

import Test.Framework (testGroup)
import Test.Framework.Providers.HUnit
import Test.HUnit
import Test.HUnit.Approx
import Test.HUnit.Plus

import MachineLearning.Random

import Data.List (nub)
import Control.Monad (foldM_)
import qualified Data.Vector as V
import qualified Data.Vector.Storable as SV
import qualified Numeric.LinearAlgebra as LA
import qualified System.Random as Rnd
import qualified Control.Monad.Random as RndM

sampleTest = do
  gen <- Rnd.newStdGen
  foldM_ sampleTestIter gen [1..25]

sampleTestIter gen i =
  let xs = V.fromList [1..100]
      n = 10 + i
      (ys, gen') = sample gen n xs
  in do
    assertEqual "uniqness" (V.length xs) (length . nub $ V.toList xs)
    assertEqual "length" n (V.length ys)
    assertBool "maximum" $ (V.maximum ys) <= (V.maximum xs)
    assertBool "minimum" $ (V.minimum ys) >= (V.minimum xs)
    assertEqual "uniqness of elements" n (length . nub $ V.toList ys)
    return gen'

randomRListTest = mapM_ (randomRListTestIter ((-1000, 1000)::(Int, Int))) [10..30]

randomRListTestIter range@(lo, hi) len = do
  rndList <- RndM.evalRandIO (getRandomRListM len range)
  assertEqual "length" len (length rndList)
  assertBool "minimum" $ lo <= (minimum rndList)
  assertBool "maximum" $ hi >= (maximum rndList)

randomRVectorTest = mapM_ (randomRVectorTestIter ((-2, 2))) [10..30]

randomRVectorTestIter range@(lo, hi) len = do
  rndVector <- RndM.evalRandIO (getRandomRVectorM len range)
  assertEqual "length" len (SV.length rndVector)
  assertBool "minimum" $ lo <= (SV.minimum rndVector)
  assertBool "maximum" $ hi >= (SV.maximum rndVector)

randomRMatrixTest = mapM_ (randomRMatrixTestIter ((-2, 2))) $ zip [10..15] [12..17]

randomRMatrixTestIter range@(lo, hi) (rows, cols) = do
  rndMatrix <- RndM.evalRandIO (getRandomRMatrixM rows cols range)
  assertEqual "rows" rows (LA.rows rndMatrix)
  assertEqual "columns" cols (LA.cols rndMatrix)
  assertBool "minimum" $ lo <= (LA.minElement rndMatrix)
  assertBool "maximum" $ hi >= (LA.maxElement rndMatrix)

tests = [ testCase "sample" sampleTest
        , testCase "getRandomRList" randomRListTest
        , testCase "getRandomRVector" randomRVectorTest
        , testCase "getRandomRMatrix" randomRMatrixTest
        ]


{-# LANGUAGE ExistentialQuantification #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE InstanceSigs #-}
{-# LANGUAGE ScopedTypeVariables #-}

module Commons where

import Control.Applicative (liftA2)
import Control.Monad (foldM)
import Data.Array hiding (range)
import Data.Complex (Complex (..), imagPart, realPart)
import Data.Function ((&))
import Data.List (intercalate, sort)
import qualified Data.Map.Strict as Map
import Data.Set (fromList, toList)
import HashedExpression.Internal
import HashedExpression.Internal.Base
import HashedExpression.Internal.Builder
import HashedExpression.Internal.MonadExpression
import HashedExpression.Interp
import HashedExpression.Modeling.Typed
import HashedExpression.Value
import Test.HUnit
import Test.Hspec
import Test.QuickCheck
import Text.Printf
import Var
import Prelude hiding ((**), (^))

sizeReduceFactor :: Int
sizeReduceFactor = 3

shapeToString :: Shape -> String
shapeToString = intercalate "x" . map show

removeDuplicate :: (Ord a) => [a] -> [a]
removeDuplicate = toList . fromList

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

genDouble :: Gen Double
genDouble = arbitrary `suchThat` inSmallRange
  where
    inSmallRange x = x >= 0 && x <= 10

vectorOfDifferent :: Eq a => Int -> Gen a -> Gen [a]
vectorOfDifferent sz gen = foldM f [] [1 .. sz]
  where
    f acc _ = (: acc) <$> gen `suchThat` (not . flip elem acc)

-------------------------------------------------------------------------------
genDimSelector :: Int -> Gen DimSelector
genDimSelector size = do
  let id = elements [0 .. size - 1]
  let step = elements [1 .. size]
  oneof [Range <$> id <*> id <*> step, At <$> id]

genAtSelector :: Int -> Gen DimSelector
genAtSelector size = do
  let id = elements [0 .. size - 1]
  oneof [At <$> id]

genValMapFor :: ExpressionMap -> Gen ValMap
genValMapFor mp = do
  let genVal :: Shape -> Gen Val
      genVal shape = do
        let dbls = vectorOf (product shape) genDouble
        case shape of
          [] -> VScalar <$> genDouble
          [sz] -> V1D . listArray (0, sz - 1) <$> dbls
          [sz1, sz2] -> V2D . listArray ((0, 0), (sz1 - 1, sz2 - 1)) <$> dbls
          [sz1, sz2, sz3] -> V3D . listArray ((0, 0, 0), (sz1 - 1, sz2 - 1, sz3 - 1)) <$> dbls
  vals <-
    (varNodes mp ++ paramNodes mp)
      & mapM
        ( \(name, shape, _) -> do
            val <- genVal shape
            return (name, val)
        )
  return $ Map.fromList vals

primitiveRUntyped :: Shape -> Gen ExprBuilder
primitiveRUntyped shape = do
  dbl <- genDouble
  name <- elements $ map pure ['a' .. 'z']
  let varName = name ++ shapeToString shape
  let parName = "p" ++ name ++ shapeToString shape
  elements
    [ introduceNode (shape, R, Var varName),
      introduceNode (shape, R, Param parName),
      introduceNode (shape, R, Const dbl)
    ]

primitiveCUntyped :: Shape -> Gen ExprBuilder
primitiveCUntyped shape = do
  re <- primitiveRUntyped shape
  im <- primitiveRUntyped shape
  return $ re +: im

genExpUntyped :: Int -> Shape -> ElementType -> Gen ExprBuilder
genExpUntyped qc shape et
  | qc == 0 = case et of
    R -> primitiveRUntyped shape
    C -> primitiveCUntyped shape
  | otherwise =
    let sub = genExpUntyped (qc `div` sizeReduceFactor) shape et
        subScalarR = genExpUntyped (qc `div` sizeReduceFactor) [] R
        subScalarC = genExpUntyped (qc `div` sizeReduceFactor) [] C
        subR = genExpUntyped (qc `div` sizeReduceFactor) shape R
        subC = genExpUntyped (qc `div` sizeReduceFactor) shape C
        subOf shape et = genExpUntyped (qc `div` sizeReduceFactor) shape et
        fromPiecewise = do
          numBranches <- elements [2, 3]
          branches <- vectorOf numBranches sub
          condition <- subR
          marks <- sort <$> vectorOfDifferent (numBranches - 1) arbitrary
          return $ piecewise marks condition branches
        binary op = liftA2 op sub sub
        unary op = op <$> sub
        commonPossibilities =
          [ fromPiecewise,
            binary (+),
            binary (*),
            liftA2 (*.) subScalarR sub,
            unary (negate),
            unary (^ 2)
          ]
        specificElementTypePosibilities
          | et == R =
            [ xRe <$> subC,
              xIm <$> subC
            ]
          | et == C =
            [ liftA2 (+:) subR subR,
              unary ft,
              unary ift,
              liftA2 (*.) subScalarC sub
            ]
        specificShapePossibilities = case shape of
          [] ->
            let fromProjection = do
                  size <- elements [3 .. 9]
                  exp <- subOf [size] et
                  ds <- genAtSelector size
                  return $ project [ds] exp
             in [ do
                    size <- elements [3 .. 9]
                    x <- subOf [size] et
                    y <- subOf [size] et
                    return $ x <.> y,
                  do
                    size1 <- elements [3 .. 9]
                    size2 <- elements [3 .. 9]
                    x <- subOf [size1, size2] et
                    y <- subOf [size1, size2] et
                    return $ x <.> y,
                  fromProjection
                ]
          [n] ->
            let fromRotate = do
                  amount <- elements [- n .. n]
                  exp <- sub
                  return $ rotate [amount] exp
                fromProjectInject = do
                  exp1 <- sub
                  exp2 <- sub
                  ds <- genDimSelector n
                  return $ inject [ds] (project [ds] exp1) exp2
                fromMatrixMul = do
                  m <- elements [3 .. 9]
                  x <- subOf [n, m] et
                  y <- subOf [m] et
                  return $ x ** y
             in [ fromRotate,
                  fromProjectInject,
                  fromMatrixMul
                ]
          [m, n] ->
            let fromRotate = do
                  amount1 <- elements [- m .. m]
                  amount2 <- elements [- n .. n]
                  exp <- sub
                  return $ rotate [amount1, amount2] exp
                fromProjectInject = do
                  exp1 <- sub
                  exp2 <- sub
                  ds1 <- genDimSelector m
                  ds2 <- genDimSelector n
                  return $ inject [ds1, ds2] (project [ds1, ds2] exp1) exp2
                fromMatrixMul = do
                  p <- elements [3 .. 9]
                  x <- subOf [m, p] et
                  y <- subOf [p, n] et
                  return $ x ** y
                fromTranspose = do
                  x <- subOf [n, m] et
                  return $ transpose x
             in [ fromRotate,
                  fromProjectInject,
                  fromMatrixMul,
                  fromTranspose
                ]
     in oneof $ commonPossibilities ++ specificShapePossibilities ++ specificElementTypePosibilities

genExp :: forall d et. (ToShape d, IsElementType et) => Int -> Gen (TypedExpr d et)
genExp size = TypedExpr <$> genExpUntyped size (toShape @d) (toElementType @et)

-------------------------------------------------------------------------------
instance (ToShape d, IsElementType et) => Arbitrary (TypedExpr d et) where
  arbitrary = sized genExp

-------------------------------------------------------------------------------
data Suite d et
  = Suite (TypedExpr d et) ValMap
  deriving (Show)

instance (ToShape d, IsElementType et) => Arbitrary (Suite d et) where
  arbitrary = do
    exp <- arbitrary
    valMap <- genValMapFor (fst $ asRawExpr exp)
    return $ Suite exp valMap

-------------------------------------------------------------------------------
type SuiteScalarR = Suite Scalar R

type SuiteScalarC = Suite Scalar C

type SuiteOneR = Suite '[Default1D] R

type SuiteOneC = Suite '[Default1D] C

type SuiteTwoR = Suite '[Default2D1, Default2D2] R

type SuiteTwoC = Suite '[Default2D1, Default2D2] C

-------------------------------------------------------------------------------
newtype ArbitraryExpr = ArbitraryExpr {unArbitraryExpr :: RawExpr}
  deriving (Show, Ord, Eq)

instance Arbitrary ArbitraryExpr where
  arbitrary = do
    elementType <- elements [R, C]
    shapeLength <- elements [0 .. 2]
    shape <- vectorOf shapeLength $ elements [1 .. 6]
    ArbitraryExpr . buildExpr <$> sized (\sz -> genExpUntyped sz shape elementType)

data XSuite = XSuite RawExpr ValMap deriving (Show, Eq, Ord)

instance Arbitrary XSuite where
  arbitrary = do
    ArbitraryExpr exp <- arbitrary
    valMap <- genValMapFor (fst exp)
    return $ XSuite exp valMap

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

instance (Ix i, Num a) => Num (Array i a) where
  (+) arr1 arr2 = listArray (bounds arr1) $ zipWith (+) (elems arr1) (elems arr2)
  (*) arr1 arr2 = listArray (bounds arr1) $ zipWith (*) (elems arr1) (elems arr2)

instance (Ix i) => ComplexRealOp (Array i Double) (Array i (Complex Double)) where
  (+:) arr1 arr2 = listArray (bounds arr1) $ zipWith (+:) (elems arr1) (elems arr2)
  xRe arr = listArray (bounds arr) $ map xRe (elems arr)
  xIm arr = listArray (bounds arr) $ map xIm (elems arr)
  conjugate arr = listArray (bounds arr) $ map conjugate (elems arr)

-- | Approximable class
class Show a => Approximable a where
  -- | (~=) checks if the two values are within acceptable numerical error
  (~=) :: a -> a -> Bool

  -- | Prettyprint a string produced by show
  prettifyShow :: a -> String

infix 4 ~=

{-
   Calculating the reletive error for two double-precision input arguments.
   It returns the error calculating by |a-b|/ max(|a|.|b|), where a and b are double-precision numbers, so it returns the relative
   error as double.
-}
relativeError :: Double -> Double -> Double
relativeError a b = abs (a - b) / max (abs a) (abs b)

instance Approximable Double where
  (~=) :: Double -> Double -> Bool
  a ~= b
    | abs (a - b) < 1.0e-5 = True
    | a == b = True
    | otherwise = relativeError a b < 0.01
  prettifyShow a
    | abs a < 1e-10 = "0" --  in prettify, inputs less than a small condition value consider as zero,
    | otherwise = printf "%.2f" a --  otherwise, it shows the exact input.

{-
Instance which belongs to approximable class for Complex Double precision inputs.
It takes 2 Complex double precision inputs and apply ~= operation on them.
calculating the error in real and imaginary parts of the complex input seperately.
Returns true if the error is less than the condition value.
-}
instance Approximable (Complex Double) where
  (~=) :: Complex Double -> Complex Double -> Bool -- Using approximable operation ~= for two Complex Double input and returns Bool
  a ~= b = (realPart a ~= realPart b) && (imagPart a ~= imagPart b) --  check for real and imaginary parts seperately
  prettifyShow a =
    prettifyShow (realPart a) ++ " + " ++ prettifyShow (imagPart a) ++ "i" --  Prettyprint a string produced by show

instance Approximable (Array Int Double) where
  (~=) :: Array Int Double -> Array Int Double -> Bool
  a ~= b = (bounds a == bounds b) && and (zipWith (~=) (elems a) (elems b))
  prettifyShow a =
    "[" ++ (intercalate ", " . map prettifyShow . elems $ a) ++ "]" --  concatenate "," to seperate the elements of array in prettify

instance Approximable (Array Int (Complex Double)) where
  (~=) :: Array Int (Complex Double) -> Array Int (Complex Double) -> Bool
  a ~= b = (bounds a == bounds b) && and (zipWith (~=) (elems a) (elems b))
  prettifyShow a =
    "[" ++ (intercalate ", " . map prettifyShow . elems $ a) ++ "]" --  concatenate "," to seperate the elements of array in prettify

instance Approximable (Array (Int, Int) Double) where
  (~=) :: Array (Int, Int) Double -> Array (Int, Int) Double -> Bool
  a ~= b = (bounds a == bounds b) && and (zipWith (~=) (elems a) (elems b))
  prettifyShow a =
    "[" ++ (intercalate ", " . map prettifyShow . elems $ a) ++ "]" --  concatenate "," to seperate the elements of array in prettify

instance Approximable (Array (Int, Int) (Complex Double)) where
  (~=) ::
    Array (Int, Int) (Complex Double) ->
    Array (Int, Int) (Complex Double) ->
    Bool
  a ~= b = (bounds a == bounds b) && and (zipWith (~=) (elems a) (elems b))
  prettifyShow a =
    "[" ++ (intercalate ", " . map prettifyShow . elems $ a) ++ "]" --  concatenate "," to seperate the elements of array in prettify

instance Approximable (Array (Int, Int, Int) Double) where
  (~=) :: Array (Int, Int, Int) Double -> Array (Int, Int, Int) Double -> Bool
  a ~= b = (bounds a == bounds b) && and (zipWith (~=) (elems a) (elems b))
  prettifyShow a =
    "[" ++ (intercalate ", " . map prettifyShow . elems $ a) ++ "]" --  concatenate "," to seperate the elements of array in prettify

instance Approximable (Array (Int, Int, Int) (Complex Double)) where
  (~=) ::
    Array (Int, Int, Int) (Complex Double) ->
    Array (Int, Int, Int) (Complex Double) ->
    Bool
  a ~= b = (bounds a == bounds b) && and (zipWith (~=) (elems a) (elems b))
  prettifyShow a =
    "[" ++ (intercalate ", " . map prettifyShow . elems $ a) ++ "]" --  concatenate "," to seperate the elements of array in prettify

instance Approximable InterpValue where
  VR x ~= VR y = x ~= y
  V1DR x ~= V1DR y = x ~= y
  V2DR x ~= V2DR y = x ~= y
  V3DR x ~= V3DR y = x ~= y
  VC x ~= VC y = x ~= y
  V1DC x ~= V1DC y = x ~= y
  V2DC x ~= V2DC y = x ~= y
  V3DC x ~= V3DC y = x ~= y
  _ ~= _ = False

  prettifyShow a = case a of
    VR x -> prettifyShow x
    V1DR x -> prettifyShow x
    V2DR x -> prettifyShow x
    V3DR x -> prettifyShow x
    VC x -> prettifyShow x
    V1DC x -> prettifyShow x
    V2DC x -> prettifyShow x
    V3DC x -> prettifyShow x

-------------------------------------------------------------------------------
shouldApprox :: (HasCallStack, Approximable a) => a -> a -> Expectation
shouldApprox x y = assertBool msg (x ~= y)
  where
    msg = "Expected: " ++ prettifyShow y ++ "\nGot: " ++ prettifyShow x

infix 1 `shouldApprox`


module Data.Array.DynamicS.MatMul(matMul) where
import Prelude hiding ((<>))
import GHC.Stack
import Data.Array.DynamicS
import Numeric.LinearAlgebra as N

matMul :: (HasCallStack, N.Numeric a) => Array a -> Array a -> Array a
matMul x y =
  case (shapeL x, shapeL y) of
    ([m, n], [n', o]) | n == n' ->
      let x' = N.reshape n $ toVector x
          y' = N.reshape o $ toVector y
          xy' = x' N.<> y'
          xy = fromVector [m, o] $ N.flatten xy'
      in  xy
    sz -> error $ "matMul: expected two conforming two-dimensional arrays, got shape " ++ show sz


{-# LANGUAGE ConstraintKinds, FlexibleContexts, FlexibleInstances, GADTs   #-}
{-# LANGUAGE KindSignatures, MultiParamTypeClasses                         #-}
{-# LANGUAGE NoMonomorphismRestriction, TypeFamilies, TypeSynonymInstances #-}
{-# OPTIONS_GHC -fno-warn-orphans #-}
module Algebra.Matrix (Matrix(..), delta, companion,
                       gaussReduction, maxNorm, rankWith, det,
                       inverse, inverseWith) where
import           Algebra.Internal
import qualified Algebra.LinkedMatrix as LM
import           Algebra.Prelude.Core hiding (maxNorm, zero)

import           Control.Lens                (both, view, (%~), (&), _3)
import           Control.Monad               (when)
import qualified Data.Matrix                 as DM
import qualified Data.Vector                 as V
import           GHC.Exts                    (Constraint)
import           Numeric.Algebra             (Additive, DecidableZero)
import           Numeric.Algebra             (Monoidal, Multiplicative)
import           Numeric.Algebra             (Unital)
import qualified Numeric.Algebra             as NA
import qualified Numeric.Decidable.Zero      as NA
import qualified Numeric.LinearAlgebra       as LA
import qualified Numeric.LinearAlgebra.Devel as LA
import qualified Prelude                     as P

{-# ANN module "HLint: ignore Redundant bang pattern" #-}

class Matrix mat where
  type Elem mat a :: Constraint
  cmap  :: (Elem mat a, Elem mat b) => (a -> b) -> mat a -> mat b
  empty :: Elem mat b => mat b
  fromLists :: Elem mat a => [[a]] -> mat a
  fromCols  :: Elem mat a => [V.Vector a] -> mat a
  fromCols [] = zero 0 0
  fromCols xs = foldr1 (<||>) $ map colVector xs
  fromRows :: Elem mat a => [V.Vector a] -> mat a
  fromRows [] = zero 0 0
  fromRows xs = foldr1 (<-->) $ map rowVector xs
  toCols :: Elem mat a => mat a -> [V.Vector a]
  toCols m = map (`getCol` m) [1..ncols m]
  toRows :: Elem mat a => mat a -> [V.Vector a]
  toRows m = map (`getRow` m) [1..nrows m]
  ncols :: mat a -> Int
  nrows :: mat a -> Int
  identity :: Elem mat a => Int -> mat a
  diag     :: Elem mat a => V.Vector a -> mat a
  getDiag :: Elem mat a => mat a -> V.Vector a
  trace :: Elem mat a => mat a -> a
  diagProd :: Elem mat a => mat a -> a
  zero :: Elem mat a => Int -> Int -> mat a
  colVector :: Elem mat a => V.Vector a -> mat a
  rowVector :: Elem mat a => V.Vector a -> mat a
  getCol :: Elem mat a => Int -> mat a -> V.Vector a
  getRow :: Elem mat a => Int -> mat a -> V.Vector a
  switchRows :: Elem mat a => Int -> Int -> mat a -> mat a
  scaleRow :: Elem mat a => a -> Int -> mat a -> mat a
  combineRows :: Elem mat a => Int -> a -> Int -> mat a -> mat a
  trans :: Elem mat a => mat a -> mat a
  buildMatrix :: Elem mat a => Int -> Int -> ((Int, Int) -> a) -> mat a
  index :: Elem mat a => Int -> Int -> mat a -> Maybe a
  index i j m = if 1 <= i && i <= nrows m && 1 <= j && j <= ncols m
                then Just $ m ! (i, j)
                else Nothing
  (!) :: Elem mat a => mat a -> (Int, Int) -> a
  (<||>) :: Elem mat a => mat a -> mat a -> mat a
  (<-->) :: Elem mat a => mat a -> mat a -> mat a
  nonZeroRows :: (DecidableZero a, Elem mat a) => mat a -> [Int]
  nonZeroRows = map fst . filter (V.any (not . NA.isZero) . snd) . zip [1..] . toRows
  nonZeroCols :: (DecidableZero a, Elem mat a) => mat a -> [Int]
  nonZeroCols = map fst . filter (V.any (not . NA.isZero) . snd) . zip [1..] . toCols

instance Matrix DM.Matrix where
  type Elem DM.Matrix a = P.Num a
  empty = DM.zero 0 0
  cmap  = fmap
  fromLists = DM.fromLists
  ncols = DM.ncols
  nrows = DM.nrows
  trans = DM.transpose
  identity = DM.identity
  diag v = DM.matrix (V.length v) (V.length v) $ \(i, j) ->
    if i == j then v V.! (i-1) else 0
  getDiag = DM.getDiag
  trace = DM.trace
  diagProd = DM.diagProd
  zero = DM.zero
  colVector = DM.colVector
  rowVector = DM.rowVector
  getCol = DM.getCol
  getRow = DM.getRow
  switchRows = DM.switchRows
  combineRows = DM.combineRows
  scaleRow = DM.scaleRow
  buildMatrix = DM.matrix
  (!) = (DM.!)
  (<||>) = (DM.<|>)
  (<-->) = (DM.<->)

swapIJ :: Eq a => a -> a -> a -> a
swapIJ i j k
  | k == i = j
  | k == j = i
  | otherwise = k

instance Matrix LA.Matrix where
  type Elem LA.Matrix a = (P.Num a, LA.Numeric  a, LA.Element a, LA.Container LA.Vector a)
  empty = LA.fromLists [[]]
  fromLists = LA.fromLists
  cmap = LA.cmap
  ncols = LA.cols
  nrows = LA.rows
  trans = LA.tr
  identity = LA.ident
  fromCols = LA.fromColumns . map (LA.fromList . V.toList)
  diag = LA.diag . LA.fromList . V.toList
  getDiag = V.fromList . LA.toList . LA.takeDiag
  trace = LA.sumElements . LA.takeDiag
  diagProd = LA.prodElements . LA.takeDiag
  zero i j = LA.konst 0 (i, j)
  colVector = LA.asColumn . LA.fromList . V.toList
  rowVector = LA.asRow . LA.fromList . V.toList
  toCols = map (V.fromList . LA.toList) . LA.toColumns
  toRows = map (V.fromList . LA.toList) . LA.toRows
  getCol i = V.fromList . LA.toList . (!! (i - 1)) . LA.toColumns
  getRow i = V.fromList . LA.toList . (!! (i - 1)) . LA.toRows
  switchRows i j m = m LA.? map (swapIJ (i-1) (j-1)) [0.. nrows m - 1]
  combineRows j s i m = LA.mapMatrixWithIndex (\(k,l) v -> if k == j - 1 then s P.* (m ! (i,l+1)) P.+ v else v) m
  buildMatrix w h f = LA.build (w, h) (\i j -> f (toIntLA i+1, toIntLA j+1))
  scaleRow a i = (fst .) $ LA.mutable $ \(k, l) m -> do
    v <- LA.readMatrix m k l
    when (k == i - 1) $
      LA.writeMatrix m k l (a P.* v)
  m ! (i, j) = m `LA.atIndex` (i - 1, j - 1)
  m <||> n = LA.fromColumns $ LA.toColumns m ++ LA.toColumns n
  m <--> n = LA.fromRows $ LA.toRows m ++ LA.toRows n

toIntLA :: LA.Container LA.Matrix e => e -> Int
toIntLA e = fromIntegral $ LA.toZ ((1 LA.>< 1) [e]) `LA.atIndex` (0, 0)

instance Matrix LM.Matrix where
  type Elem LM.Matrix a = (CoeffRing a, Unital a, Monoidal a, Multiplicative a, Additive a)
  cmap = LM.cmap
  (!) m pos = m LM.! (pos & both %~ pred)
  index i j = LM.index (i-1) (j-1)
  empty = LM.empty
  buildMatrix h w f = LM.fromList [((i, j), f (i,j)) | j <- [1..w], i <- [1..h]]
  trans = LM.transpose
  combineRows j s i = LM.combineRows s (i-1) (j-1)
  switchRows i j = LM.switchRows (i-1) (j-1)
  scaleRow a i = LM.scaleRow a (pred i)
  zero = LM.zeroMat
  identity = LM.identity
  diag = LM.diag
  getDiag = LM.getDiag
  diagProd = LM.diagProd
  trace = LM.trace
  getRow = LM.getRow . pred
  getCol = LM.getCol . pred
  nrows = LM.nrows
  ncols = LM.ncols
  fromLists = LM.fromLists
  (<||>) = (LM.<||>)
  (<-->) = (LM.<-->)
  rowVector = LM.rowVector
  colVector = LM.colVector
  nonZeroRows = LM.nonZeroRows
  nonZeroCols = LM.nonZeroCols

delta :: (NA.Monoidal r, NA.Unital r) => Int -> Int -> r
delta i j | i == j = NA.one
          | otherwise = NA.zero

companion :: (KnownNat n, CoeffRing r, Matrix mat,
              Elem mat r, IsMonomialOrder n ord)
          => Ordinal n -> OrderedPolynomial r ord n -> mat r
companion odn poly =
  let deg = totalDegree' poly
      vx  = leadingMonomial (var odn `asTypeOf` poly)
  in buildMatrix deg deg $ \(j, k) ->
  if 1 <= k && k <= deg - 1
  then delta j (k+1)
  else NA.negate $ coeff (NA.pow vx (fromIntegral $ j-1 :: NA.Natural)) poly

-- | @gaussReduction a = (a', p)@ where @a'@ is row echelon form and @p@ is pivoting matrix.
gaussReduction :: (Matrix mat, Elem mat a, Normed a, Eq a, NA.Field a)
               => mat a -> (mat a, mat a)
gaussReduction mat =
  let (a, b, _) = gaussReduction' mat in (a, b)

-- | @gaussReduction a = (a', p)@ where @a'@ is row echelon form and @p@ is pivoting matrix.
gaussReduction' :: (Matrix mat, Elem mat a, Normed a, Eq a, NA.Field a)
               => mat a -> (mat a, mat a, a)
gaussReduction' mat = {-# SCC "gaussRed" #-} go 1 1 mat (identity $ nrows mat) NA.one
  where
    go i j a p acc
      | i > nrows mat || j > ncols mat = (a, p, acc)
      | otherwise =
        let (k, new) = maximumBy (comparing $ norm . snd) [(l, a ! (l, j)) | l <- [i..nrows mat]]
        in if new == NA.zero
           then go i (j + 1) a p NA.zero
           else let prc l a0 p0
                      | l == i = prc (l+1) a0 p0
                      | l > nrows mat = (a0, p0)
                      | otherwise     =
                        let coe = NA.negate (a0 ! (l, j))
                            a'' = combineRows l coe i a0
                            p'' = combineRows l coe i p0
                        in prc (l+1) a'' p''
                    (a', p') = prc 1 (scaleRow (NA.recip new) i $ switchRows i k a)
                                     (scaleRow (NA.recip new) i $ switchRows i k p)
                    offset = if i == k then id else NA.negate
                in go (i+1) (j+1) a' p' (offset $ acc NA.* new)

det :: (Elem mat a, Eq a, NA.Field a, Normed a, Matrix mat)
    => mat a -> a
det = view _3 . gaussReduction'

maxNorm :: (Elem mat a, Normed a, Matrix mat) => mat a -> Norm a
maxNorm = maximum . concat . map (map norm . V.toList) . toRows

rankWith :: (Elem mat r, CoeffRing r, Matrix mat)
         => (mat r -> mat r) -> mat r -> Int
rankWith gauss = length . nonZeroRows . gauss

inverse :: (Elem mat a, Eq a, NA.Field a, Normed a, Matrix mat)
        => mat a -> mat a
inverse = snd . gaussReduction

inverseWith :: (mat a -> (mat a, mat a)) -> mat a -> mat a
inverseWith = (snd .)


{-# LANGUAGE BangPatterns  #-}
{-# LANGUAGE DeriveGeneric #-}

module GraphBuilder (
    shortDistMat,
    shortestDists,
    joinContext,
    EContext(..),
    EShortDists(..),
    theContextsAndDists,
    ) where

import           Control.DeepSeq                   (NFData (..))
import           Control.Monad                     (ap, join, liftM2)
import           Control.Monad.ST
import           Control.Parallel.Strategies
import           Data.Graph.Inductive              hiding (ap)
import qualified Data.Graph.Inductive.PatriciaTree as GP
import           Data.List                         (sortOn)
import qualified Data.Strict.Tuple                 as T hiding (uncurry)
import           Foreign.C                         (CInt)
import           GHC.Generics
import           Numeric.LinearAlgebra
import           Numeric.LinearAlgebra.Devel

data EContext =
       EContext
         { vertex   :: {-# UNPACK #-} !Int
         , dstToCtr :: {-# UNPACK #-} !Double
         , dists    :: {-# UNPACK #-} !(Vector Double)
         , ns       :: {-# UNPACK #-} !(Vector Int)
         }
  deriving (Show, Generic)

instance NFData EContext

data EShortDists =
       EShortDists
         { vertex' :: {-# UNPACK #-} !Int
         , dists'  :: {-# UNPACK #-} !(Vector Double)
         }
  deriving (Show, Generic)

instance NFData EShortDists

shortDistMat :: Int -> [EShortDists] -> Matrix Double
shortDistMat size' !sdists =
  let mat = konst 0 (length', size') :: Matrix Double
      length' = length sdists
  in runST $ do
    m <- unsafeThawMatrix mat
    mapM_ (`setRow` m) $ zip sdists [0 ..]
    unsafeFreezeMatrix m
  where
    setRow (e, pt) mat = setMatrix mat pt 0 . asRow . dists' $ e

joinContext :: EContext -> GP.Gr Double Double -> GP.Gr Double Double
joinContext (EContext pt dstCtr edists vs) graph =
  let zipped = zip (toList edists) (toList vs)
  in (zipped, pt, dstCtr, zipped) & graph

matD2 :: Matrix Double -> Matrix Double
matD2 = join pairwiseD2

theContextsAndDists :: (Int, Int) -- ^ Graph neighborhood size and PCA neighborhood size
                    -> Int
                    -> Matrix Double
                    -> ([EContext], Matrix Double, [Int])
theContextsAndDists (n1, n2) bpt mat =
  (makeContexts . theNearest n1 $ mat', extract' . idxs $ idxs'', idxs''')
  where
    makeContexts = zipWith makeContext [0 :: Int ..] . map unzip
    makeContext a (b, c) = EContext a (mat' `atIndex` (a, bpt)) (fromList b) (fromList c)
    mat' = matD2 mat
    idxs'' = map (toEnum . snd) idxs'
    idxs' = (!! bpt) . theNearest n2 $ mat'
    idxs''' = map snd idxs'
    extract' a = mat ?? (Pos a, All)

theNearest :: Int -> Matrix Double -> [[(Double, Int)]]
theNearest n = filterKNearest . findNearestWithInd
  where
    findNearestWithInd = helper1 . helper2
    filterKNearest = map (take n . filter ((/= 0.0) . fst) . uncurry zip)
    helper1 in' = uncurry zip (map toList (T.fst in'), map (map fromEnum . toList) (T.snd in'))
    helper2 = strUnzip . map (liftM2 (T.:!:) sortVector sortIndex) . toRows

shortestDists :: [Int] -> GP.Gr Double Double -> [EShortDists]
shortestDists inds g = parMap rseq (clean `ap` (sortIndDumpDist . makeList)) inds
  where
    clean x' a = EShortDists x' $ fromList a
    sortIndDumpDist = map T.snd . sortOn T.fst
    makeList = map (Prelude.uncurry (T.:!:)) . join . map (take 1 . unLPath) . flip spTree g

-- UTILITY FUNCTION(S) --

strUnzip :: [T.Pair (Vector Double) (Vector CInt)] -> T.Pair [Vector Double] [Vector CInt]
strUnzip = foldr (\((T.:!:) a b) acc -> ((T.:!:) (a : T.fst acc) (b : T.snd acc))) ((T.:!:) [] [])

-- -- For debugging!
-- remember to import Control.Arrow ((&&&)) when using this function.
-- checkSymRowsCols :: (Container Vector a, Eq a, S.Storable a)
--                     => Matrix a
--                     -> Bool
-- checkSymRowsCols = uncurry (==) . (map sumElements . toColumns &&& map sumElements . toRows)


{-# LANGUAGE ConstraintKinds #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE UndecidableInstances #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE FunctionalDependencies #-}
module VAD(
      vad
    ) where

--  let framed = frameWithWinAndOverlap 256 hann 10 s
--      sp = mapS (fft . U.map (:+ 0)) framed
--

import Prelude hiding(splitAt,(:),foldr1,tail,(++))
import Internal
import Signal
import qualified Data.Vector.Unboxed as U
import Data.Vector.Unboxed((!),Unbox(..))
import Data.List.Stream
import Common
import Windows
import Data.Complex
import Transform
import qualified Trace as F
import Fixed
import GHC.TypeLits
import SpecialInt
import Data.Int

ltseF :: Sample a
      => Int
      -> Signal (U.Vector a)
      -> Signal (U.Vector a)
ltseF n (Signal s) =
    let (before,remaining) = splitAt n s
        _lste b (h:t) =
            let (future,tl) = splitAt n t
                valueMax = U.zipWith max (foldr1 (U.zipWith max) (h:future))  (foldl1' (U.zipWith max) b)
            in
            valueMax:(_lste (tail before ++ [h]) t)
    in
    Signal (_lste before remaining)

lstdD :: Int
      -> U.Vector Double
      -> U.Vector Double
      -> Double
lstdD winSize noiseEnergy lste =
        let d a b | b == 0 = fromDouble 1e12
                  | otherwise = a / b
        in
        10 * log(U.sum (U.zipWith (d) lste noiseEnergy) / fromIntegral winSize) / log 10

theTotal :: (SingI r, SingI s, SingI n)
         => Fixed Int32 n s r
         -> U.Vector (Fixed Int32 n s r)
         -> U.Vector (Fixed Int32 n s r)
         -> Fixed Int40 n s r
theTotal m lste noiseEnergy =
           let  d a b | b == 0 = convert m
                      | otherwise = convert a / convert b
           in
           U.sum (U.zipWith (d) lste noiseEnergy)
{-# INLINE [0] theTotal #-}

lstdF :: (SingI n, SingI s, SingI r)
      => Int
      -> U.Vector (Fixed Int32 n s r)
      -> U.Vector (Fixed Int32 n s r)
      -> (Fixed Int32 n s r)
lstdF winSize noiseEnergy lste = 10 * log(convert $ theTotal maxBound lste noiseEnergy/ fromIntegral winSize) / log 10

getDecisionF :: (SingI n, SingI s, SingI r, SingI (n + n))
             => Int
             -> U.Vector (Fixed Int32 (n + n) s r)
             -> [(U.Vector (Fixed Int32 (n + n) s r), U.Vector (Fixed Int32 (n + n) s r) )]
             -> [Fixed Int16 n s r]
getDecisionF winSize energy ((c,currentE):r) =
    let tv = fromDouble 31.0
        tn = fromDouble 27.0
        l = lstdF winSize energy c
        result | l >= tv = 1 : getDecisionF winSize energy r
               | l <= tn = 0 : getDecisionF winSize currentE r
               | otherwise = 0 : getDecisionF winSize energy r
    in
    result
getDecisionF winSize energy [] = 0:getDecisionF winSize energy []

getDecisionD :: Int
             -> U.Vector Double
             -> [(U.Vector Double, U.Vector Double)]
             -> [Double]
getDecisionD winSize energy ((c,currentE):r) =
    let tv = fromDouble 31.0
        tn = fromDouble 27.0
        l = lstdD winSize energy c
        result | l >= tv = 1 : getDecisionD winSize energy r
               | l <= tn = 0 : getDecisionD winSize currentE r
               | otherwise = 0 : getDecisionD winSize energy r
    in
    result
getDecisionD winSize energy [] = 0:getDecisionD winSize energy []

bandEnergy :: Complex Double -> Double
bandEnergy (x :+ y) = x*x + y*y

bandEnergyF :: (SingI n, SingI s, SingI r, SingI (n + n)) => Complex (Fixed Int16 n s r) -> Fixed Int32 (n + n) s r
bandEnergyF (x :+ y) = amul x x + amul y y

class VAD a where
    vad :: (Sample a, FFT a)
        => Sampled Time a
        -> Sampled Time a

instance (SingI n, SingI s, SingI r, SingI (n + n)) => VAD (Fixed Int16 n s r) where
    vad s =
        let winSize = 256
            overlap = 20
            n = 2
            framed = frameWithWinAndOverlap winSize overlap hann s
            energy = mapS (U.map bandEnergyF . fft . U.map (:+ 0)) (getSignal framed)
            noiseEnergy0 = U.generate winSize (const (fromDouble 0.00001))
            lt = ltseF n energy
        in
        --mapS (lstd winSize noiseEnergy0) lt
        Sampled (period framed) (onSamples (getDecisionF winSize noiseEnergy0) (zipS lt (dropS n energy)))

instance VAD Double where
    vad s =
        let winSize = 256
            overlap = 20
            n = 2
            framed = frameWithWinAndOverlap winSize overlap hann s
            energy = mapS (U.map bandEnergy . fft . U.map (:+ 0)) (getSignal framed)
            noiseEnergy0 = U.generate winSize (const (fromDouble 0.00001))
            lt = ltseF n energy
        in
        --mapS (lstd winSize noiseEnergy0) lt
        Sampled (period framed) (onSamples (getDecisionD winSize noiseEnergy0) (zipS lt (dropS n energy)))


{-# LANGUAGE TemplateHaskell, GeneralizedNewtypeDeriving, ScopedTypeVariables #-}
module CNC.Declarative(module CNC.Figure,
                       mirrorX, mirrorY, rotate, scale, posToExpr, renderPath) where

import CNC.Figure
import CNC.Geometry(RealT, eps)
import CNC.HCode(HCode, frame, g, x, y,
                 Expr, toExpr)

import Data.Complex

-- mirrors around the X axis
mirrorX :: Path -> Path
mirrorX = mapPath (\(x :+ y) -> (x :+ negate y))

-- mirrors around the Y axis
mirrorY :: Path -> Path
mirrorY = mapPath (\(x :+ y) -> (negate x :+ y))

-- rotates around (0,0) point counter-clockwise on a given angle
rotate :: RealT -> Path -> Path
rotate a = mapPath (* exp (0 :+ a))

-- scales the path
scale :: RealT -> Path -> Path
scale k = mapPath (* (k :+ 0))

posToExpr :: Pos -> (Expr RealT, Expr RealT)
posToExpr (x :+ y) = (toExpr x, toExpr y)

-- renders a path in HCode monad
renderPath :: (Expr RealT, Expr RealT) -> Path -> HCode ()
renderPath (x0, y0) path = do
  frame [g 0, x x0, y y0]
  mapPathM_ (\(xi :+ yi) -> frame [g 1, x $ toExpr xi + x0, y $ toExpr yi + y0]) path


module RadialPhase where

import           Control.Monad           as M
import           Control.Monad.Parallel  as MP
import           Data.Array.Repa         as R
import           Data.Binary
import           Data.Complex
import           Data.List               as L
import           Data.Vector.Unboxed     as VU
import           DFT.Plan
import           Filter.Utils
import           FokkerPlanck.MonteCarlo
import           FokkerPlanck.Pinwheel
import           Graphics.Gnuplot.Simple
import           Image.IO
import           System.Directory
import           System.Environment
import           System.FilePath
import           Text.Printf
import           Types
import           Utils.Array

main = do
  args@(numPointStr:thetaSigmaStr:scaleSigmaStr:maxScaleStr:maxScaleStr1:taoStr:numTrailStr:maxTrailStr:theta0FreqsStr:thetaFreqsStr:scale0FreqsStr:scaleFreqsStr:histFilePath:alphaStr:deltaThetaStr:numThreadStr:_) <-
    getArgs
  print args
  let numPoint = read numPointStr :: Int
      thetaSigma = read thetaSigmaStr :: Double
      scaleSigma = read scaleSigmaStr :: Double
      maxScale = read maxScaleStr :: Double
      maxScale1 = read maxScaleStr1 :: Double
      tao = read taoStr :: Double
      numTrail = read numTrailStr :: Int
      maxTrail = read maxTrailStr :: Int
      theta0Freq = read theta0FreqsStr :: Double
      theta0Freqs = [-theta0Freq .. theta0Freq]
      thetaFreq = read thetaFreqsStr :: Double
      thetaFreqs = [-thetaFreq .. thetaFreq]
      scale0Freq = read scale0FreqsStr :: Double
      scaleFreq = read scaleFreqsStr :: Double
      scale0Freqs = [-scale0Freq .. scale0Freq]
      scaleFreqs = [-scaleFreq .. scaleFreq]
      alpha = read alphaStr :: Double
      deltaTheta = read deltaThetaStr :: Double
      numThread = read numThreadStr :: Int
      folderPath = "output/test/RadialPhase"
  createDirectoryIfMissing True folderPath
  removePathForcibly (folderPath </> "Magnitude")
  createDirectoryIfMissing True (folderPath </> "Magnitude")
  flag <- doesFileExist histFilePath
  radialArr <-
    if flag
      then R.map magnitude . getNormalizedHistogramArr <$>
           decodeFile histFilePath
      else solveMonteCarloR2Z2T0S0Radial
             numThread
             numTrail
             maxTrail
             numPoint
             numPoint
             thetaSigma
             scaleSigma
             maxScale
             tao
             theta0Freqs
             thetaFreqs
             scale0Freqs
             scaleFreqs
             histFilePath
             (emptyHistogram
                [ (round . sqrt . fromIntegral $ 2 * (div numPoint 2) ^ 2)
                , L.length scale0Freqs
                , L.length theta0Freqs
                , L.length scaleFreqs
                , L.length thetaFreqs
                ]
                0)
  let arr3d =
        R.slice
          radialArr
          (Z :. (L.length thetaFreqs - 1) :. All :. (L.length theta0Freqs - 1) :.
           All :.
           All)
      (Z :. _ :. _ :. _ :. _ :. r) = extent radialArr
      idx =
        [ (tf, sf, t0f, s0f)
        | tf <- L.zip [0 .. (L.length thetaFreqs) - 1] [-3,-2,-1] --[(-thetaFreq) .. 0]
        , sf <- L.zip [0 .. (L.length scaleFreqs) - 1] [(-scaleFreq) .. 0]
        , t0f <- L.zip [0 .. (L.length theta0Freqs) - 1] theta0Freqs
        , s0f <- L.zip [0 .. (L.length scale0Freqs) - 1] scale0Freqs
        ]
      idxs =
        L.groupBy
          (\((_, tf), (_, sf), (_, t0f), (_, s0f)) ((_, tf'), (_, sf'), (_, t0f'), (_, s0f')) ->
             (tf + t0f, sf + s0f) == (tf' + t0f', sf' + s0f')) .
        L.sortBy
          (\((_, tf), (_, sf), (_, t0f), (_, s0f)) ((_, tf'), (_, sf'), (_, t0f'), (_, s0f')) ->
             compare (tf + t0f, sf + s0f) (tf' + t0f', sf' + s0f')) $
        idx
      -- idxs =
      --   L.groupBy
      --     (\((_, tf), (_, sf), (_, t0f), (_, s0f)) ((_, tf'), (_, sf'), (_, t0f'), (_, s0f')) ->
      --        (abs $ tf + t0f, abs $ tf - t0f) == (abs $ tf' + t0f', abs $ tf' - t0f')) .
      --   L.sortBy
      --     (\((_, tf), (_, sf), (_, t0f), (_, s0f)) ((_, tf'), (_, sf'), (_, t0f'), (_, s0f')) ->
      --        compare (abs $ tf + t0f, abs $ tf - t0f) (abs $ tf' + t0f', abs $ tf' - t0f')) $
      --   idx
  -- MP.mapM_
  --   (\((i, sf), (j, s0f)) ->
  --      plotPathStyle
  --        [ PNG (folderPath </> printf "%d_%d.png" i j)
  --        , Title (printf "sf: %.0f, s0f: %.0f" sf s0f)
  --        -- , XRange (0, log (fromIntegral len))
  --        , YRange (0, 2 * pi)
  --        ]
  --        (defaultStyle
  --           {plotType = LinesPoints, lineSpec = CustomStyle [LineTitle ""]}) .
  --      L.zip (L.map log [1 .. fromIntegral len - 1]) .
  --      L.tail . R.toList . R.map (\x -> (phase x) + pi) . R.slice arr3d $
  --      (Z :. i :. j :. All))
  --   [ (a, b)
  --   | a <- L.zip [0 .. (L.length scaleFreqs) - 1] scaleFreqs
  --   , b <- L.zip [0 .. (L.length scale0Freqs) - 1] scale0Freqs
  --   ]
  -- let maxMag = VU.maximum . toUnboxed . computeS . R.map magnitude $ arr3d
  plotPathsStyle
    [ PNG (folderPath </> "Magnitude" </> "all.png")
    , Title ""
    , XRange (60, 120)
    , YRange (0.01, 0.015)
    ] (L.map
         (\((i, tf), (j, sf), (k, t0f), (l, s0f)) ->
            ( defaultStyle
                { plotType = LinesPoints
                , lineSpec =
                    CustomStyle
                      [ LineTitle
                          (printf "%.0f %.0f %.0f %.0f" tf t0f sf s0f)
                      ]
                }
            , L.drop 5 .
              L.zip ([0,5 .. 5*(fromIntegral r - 1)]) .
              R.toList . R.slice radialArr $
              (Z :. i :. j :. k :. l :. All)))
         idx)
  MP.mapM_
    (\xs@(((i, tf), (j, sf), (k, t0f), (l, s0f)):_) ->
       let t = tf + t0f
           s = sf + s0f
        in plotPathsStyle
             [ PNG (folderPath </> "Magnitude" </> printf "%.0f_%.0f.png" t s)
             , Title (printf "tSum: %.0f sSum: %.0f" t s)
             , YRange (0, 0.045)
             ]
             (L.map
                (\((i, tf), (j, sf), (k, t0f), (l, s0f)) ->
                   ( defaultStyle
                       { plotType = LinesPoints
                       , lineSpec =
                           CustomStyle
                             [ LineTitle
                                 (printf "%.0f %.0f %.0f %.0f" tf t0f sf s0f)
                             ]
                       }
                   , L.drop 5 .
                     L.zip ([0 .. fromIntegral r - 1]) .
                     R.toList . R.slice radialArr $
                     (Z :. i :. j :. k :. l :. All)))
                xs))
    idxs
  -- MP.mapM_
  --   (M.mapM_
  --      (\((i, tf), (j, sf), (k, t0f), (l, s0f)) -> do
  --         createDirectoryIfMissing
  --           True
  --           (folderPath </> "Magnitude" </>
  --            (printf "%.0f_%.0f" (tf + t0f) (sf + s0f)))
  --         plotPathStyle
  --           [ PNG
  --               (folderPath </> "Magnitude" </>
  --                (printf "%.0f_%.0f" (tf + t0f) (sf + s0f)) </>
  --                printf "%.0f_%.0f_%.0f_%.0f.png" tf sf t0f s0f)
  --           , Title ""
  --               -- (printf "tf: %.0f sf: %.0f t0f: %.0f s0f: %.0f" tf sf t0f s0f)
  --           ]
  --           (defaultStyle
  --              {plotType = LinesPoints, lineSpec = CustomStyle [LineTitle ""]}) .
  --           L.zip ([0 .. fromIntegral r - 1]) . R.toList . R.slice radialArr $
  --           (Z :. i :. j :. k :. l :. All)))
  --   idxs
  -- MP.mapM_
  --   (\((i, tf), (j, sf), (k, t0f), (l, s0f)) ->
  --      plotPathStyle
  --        [ PNG
  --            (folderPath </> "Magnitude" </>
  --             printf "%.0f_%.0f_%.0f_%.0f.png" tf sf t0f s0f)
  --        , Title (printf "tf: %.0f sf: %.0f t0f: %.0f s0f: %.0f" tf sf t0f s0f)
  --        -- , XRange (0, log (fromIntegral len))
  --        -- , YRange (0, maxMag)
  --        ]
  --        (defaultStyle
  --           {plotType = LinesPoints, lineSpec = CustomStyle [LineTitle ""]}) .
  --      L.zip ([0 .. fromIntegral r - 1]) .
  --      -- L.tail .
  --      R.toList . R.slice radialArr $
  --      (Z :. i :. j :. k :. l :. All))
  --   [ (tf, sf, t0f, s0f)
  --   | tf <- L.zip [0 .. (L.length thetaFreqs) - 1] thetaFreqs
  --   , sf <- L.zip [0 .. (L.length scaleFreqs) - 1] scaleFreqs
  --   , t0f <- L.zip [0 .. (L.length theta0Freqs) - 1] theta0Freqs
  --   , s0f <- L.zip [0 .. (L.length scale0Freqs) - 1] scale0Freqs
  --   ]
  -- MP.mapM_
  --   (\((i, sf), (j, s0f)) ->
  --      plotPathStyle
  --        [ PNG (folderPath </> "Pinwheel" </> printf "%d_%d.png" i j)
  --        , Title (printf "sf: %.0f, s0f: %.0f" sf s0f)
  --        -- , XRange (0, log (fromIntegral len))
  --        , YRange (0, 2 * pi)
  --        ]
  --        (defaultStyle
  --           {plotType = LinesPoints, lineSpec = CustomStyle [LineTitle ""]}) .
  --      L.zip (L.map log [1 .. fromIntegral len - 1]) .
  --      L.map
  --        (\r ->
  --           pi +
  --           (phase $
  --            exp
  --              (0 :+
  --               (s0f *
  --                (log r * 2 * pi / log maxScale1 + (deltaTheta / 180 * pi)) +
  --                sf * (log r * 2 * pi / log maxScale1 + (deltaTheta / 180 * pi)))))) $
  --      [1 .. fromIntegral len - 1])
  --   [ (a, b)
  --   | a <- L.zip [0 .. (L.length scaleFreqs) - 1] scaleFreqs
  --   , b <- L.zip [0 .. (L.length scale0Freqs) - 1] scale0Freqs
  --   ]


{-|
Module : GradientDescent
|-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE TypeFamilies #-}

module GradientDescent
  ( sigmoid
  , sigmoidPrime
  , matrixSigmoid
  , matrixSigmoidPrime
  , getOutcomes
  , gradientDescent
  ) where

import Numeric.LinearAlgebra
import qualified Data.List.Split as S
import System.Random
import Data.Array.IO
import Control.Monad
import GHC.Float
import Network

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

-- | Sigmoid Function
sigmoid :: Double -> Double
sigmoid x = 1.0 / (1.0 + exp (-x))

-- | Apply Sigmoid Function for every element of matrix
matrixSigmoid :: Matrix Double -> Matrix Double
matrixSigmoid = cmap sigmoid

-- | Sigmoid Prime Function
sigmoidPrime :: Double -> Double
sigmoidPrime x = sigmoid x * (1 - sigmoid x)

-- | Apply Sigmoid Prime Function for every element of matrix
matrixSigmoidPrime :: Matrix Double -> Matrix Double
matrixSigmoidPrime = cmap sigmoidPrime

-- | Cost function
costPrime :: Matrix Double -> Matrix Double  -> Matrix Double
costPrime outputActivations y = outputActivations - y

-- | Feed Forward algorithm
feedForward :: Net -- ^ Neural Network
            -> Matrix Double -- ^ Image Matrix
            -> Matrix Double
feedForward net x = foldl (\z (w,b)-> matrixSigmoid ((w <> z) + b)) x (zip (weights net) (biases net))

-- | Evaluate Neural Network predictions
evaluate :: Net -- ^ Neural Network for evaluation
         -> [(Matrix Double, Int)] -- ^ List of Test Images with their corresponding label
         -> Double -- ^ Returns Count of positive predictions
evaluate net testData = sum $ fmap (\(x,y) -> if x==y then 1.0 else 0.0) outcomes
  where
    testImages = map fst testData
    testLabels = map snd testData
    outcomes = zip (fmap (head . fmap maxIndex . toColumns . feedForward net) testImages) testLabels

-- | Get Neural Network predictions
getOutcomes :: Net -- ^ Neural Network
            -> [(Matrix Double, Int)] -- ^ List of Test Images with their corresponding labels
            -> [([Float], (Int, Int))] -- ^ Returns List of Test Images with their corresponding label and predictions
getOutcomes net testData = zip listImages outcomes
  where
    testImages = map fst testData
    testLabels = map snd testData
    listImages = head . toLists . tr . cmap double2Float <$> testImages
    outcomes = zip (fmap (head . fmap maxIndex . toColumns . feedForward net) testImages) testLabels

-- | Get List of neuron activations
neuronActivation :: [Matrix Double] -- ^ Weights of Neural Network
                 -> [Matrix Double] -- ^ Biases of Neural Network
                 -> Matrix Double   -- ^ Image Matrix
                 -> [Matrix Double]
neuronActivation [] [] x = []
neuronActivation (w:ws) (b:bs) x = z : neuronActivation ws bs (matrixSigmoid z)
  where
    z = (w <> x) + b

-- | Calculate deltas
calDelta :: Net -- ^ Neural Network
          -> [Matrix Double] -- ^ Weights of Neural Network
          -> [Matrix Double] -- ^ List of Neuron Activations
          -> Matrix Double   -- ^ delta
          -> Int -- ^ current layer
          -> [Matrix Double]
calDelta net w z delta l
  | l == numLayers net = [delta]
  | otherwise = calDelta net (tail w) (tail z) d (l + 1) ++ [delta]
    where
      sp = matrixSigmoidPrime (head z)
      d = tr (head w) <> delta * sp

-- | Back Propagation algorithm
backProp :: Net -- ^ Neural Network
         -> Matrix Double -- ^ Image Matrix
         -> Matrix Double -- ^ Image vectorized label
         -> ([Matrix Double],[Matrix Double])
backProp net x y = (nablaW, nablaB)
  where
    zs = neuronActivation (weights net) (biases net) x
    activations = x : (matrixSigmoid <$> zs)
    delta = (costPrime (last activations) y * matrixSigmoidPrime (last zs)) :: Matrix Double
    lDelta = calDelta net (reverse $ weights net) (tail $ reverse zs) delta 2
    nablaB = lDelta
    nablaW = (\ (x, y) -> x <> tr y) <$> zip lDelta activations

-- | Calculate Nabla for Weights and Biases
calNabla :: Net -- ^ Neural Network
         -> [(Matrix Double, Matrix Double)] -- ^ Batch of Images with their corresponding labels
         -> [Matrix Double] -- ^ List of Nabla Matrix's of Weights
         -> [Matrix Double] -- ^ List of Nabla Matrix's of Biases
         -> ([Matrix Double], [Matrix Double])
calNabla net [] nW nB = (nW, nB)
calNabla net (x:xs) nW nB = calNabla net xs nablaW nablaB
  where
    (deltaNablaW, deltaNablaB) = uncurry (backProp net) x
    nablaB = uncurry (+) <$> zip nB deltaNablaB
    nablaW = uncurry (+) <$> zip nW deltaNablaW

-- | Update network's weights and biases for each batch
updateMiniBatch :: Net -- ^ Neural Network
                -> [(Matrix Double, Matrix Double)] -- ^ Batch of Images with their corresponding labels
                -> Double -- ^ Alpha
                -> Net
updateMiniBatch net miniBatch alpha = newNet
  where
    nB = fmap (\x -> konst (0 :: Double) ((fst $ size x) :: Int, (snd $ size x) :: Int) :: Matrix Double) (biases net)
    nW = fmap (\x -> konst (0 :: Double) ((fst $ size x) :: Int, (snd $ size x) :: Int) :: Matrix Double) (weights net)
    (nablaW, nablaB) = calNabla net miniBatch nW nB
    newWeights = (\(x,y) -> x - cmap (\ p -> (alpha / fromIntegral (length miniBatch)) * p) y ) <$> zip (weights net) nablaW
    newBiases = (\(x,y) -> x - cmap (\ p -> (alpha / fromIntegral (length miniBatch)) * p) y ) <$> zip (biases net) nablaB
    newNet = Net (numLayers net) (sizes net) newBiases newWeights

-- | Train Neural Network
gradientDescent :: Net -- ^ Neural Network to train
                -> [(Matrix Double, Matrix Double)] -- ^ List of Training Images with column vectorized labels
                -> Int -- ^ Number of training iterations
                -> Int -- ^ Size of Image batches
                -> Double -- ^ Alpha
                -> [(Matrix Double, Int)] -- ^ List of Test Images with Labels
                -> IO Net
gradientDescent net trainingData epochs miniBatchSize alpha testData = loop net 0
  where
    nTest = fromIntegral (length testData)
    loop currentNet epoch = do
      let evaluation = 100.0 * evaluate currentNet testData / nTest
      let evaluationMsg = "Epoch: " ++ show epoch ++ " = " ++ show evaluation ++ "% Out Of " ++ show nTest ++ " Test Images"
      print evaluationMsg
      if epoch == epochs then return currentNet
      else do
        shuffledTrainingData <- shuffle trainingData
        let miniBatches = S.chunksOf miniBatchSize shuffledTrainingData
        let newNet = foldl (\z x-> updateMiniBatch z x alpha) currentNet miniBatches
        loop newNet (epoch + 1)


{-# LANGUAGE DataKinds, FlexibleContexts #-}

module Lib.Practice.P9 (
    f9
) where

import qualified Numeric.LinearAlgebra as LA
import Numeric.LinearAlgebra.Static
import Lib.Utility
import Control.Monad
import Control.Monad.Trans.Maybe
import Control.Monad.IO.Class
import GHC.TypeLits

f9 = downMaybeTIO f9'
f9' :: MaybeT IO ()
f9' = do
    mnist@(tr@(trX, trT), te@(teX, teT)) <- getMnistData
    net <- liftIO (initTestNet 0.01 :: IO (MnistTwoLayerNet 50 100 784 10))
    showAccuracy te net
    net' <- train10000 mnist net
    --net' <- train100 net train'
    --showAccuracy net'
    return ()

showNetParam (MnistTwoLayerNet af1 relu af2 _) = b1
    where (Layer param2 _ _) = af2
          (Layer param1 _ _) = af1
          ((w2, b2), _, _) = param2
          ((w1, b1), _, _) = param1
train10000 mnist net = foldM (\ net' _ -> train600 mnist net') net [1..16]
train600 (tr, te) net = do
    net' <- foldM (\ net' _ -> train tr net') net [1 .. 600]
    showAccuracy te net'
    liftIO $ print $ showNetParam net'
    return net'

showAccuracy te net = do
    (batchX, batchT) <- miniBatch te
    MaybeT $ fmap Just $ print $ accuracy net batchX batchT
train tr net = do
    (batchX, batchT) <- miniBatch tr
    return $ training' net batchX batchT

training' (MnistTwoLayerNet af1 relu af2 updater) x t = MnistTwoLayerNet af1' relu af2' updater'
    where ff layer x = case layer of (Layer _ ff _) -> ff x
          bf layer state = case layer of (Layer _ _ bf) -> bf state
          (y1, af1s) = ff af1 x
          (y2, relus) = ff relu y1
          (y3, af2s) = ff af2 y2
          swl = softmaxWithLossLayer t
          (y4, swls) = ff swl y3
          resultMaybe = do
              let (dx1, _) = bf swl swls 1
              let (dx2, (p2, _, dwdbMaybe2)) = bf af2 af2s dx1
              dp2 <- dwdbMaybe2
              let (dx3, _) = bf relu relus dx2
              let (_, (p1, _, dwdbMaybe1)) = bf af1 af1s dx3
              dp1 <- dwdbMaybe1
              let ((af1s', af2s'), u') = update updater (p1, p2) (dp1, dp2)
              return (affineLayer af1s', affineLayer af2s', u')
          (af1', af2', updater') = case resultMaybe of
              Nothing -> error "Something wrong"
              Just x -> x

testAffine = y
    where (Layer s f bf) = affineLayer (w, b) :: Layer 2 2 3 (AffineState 2 2 3)
          (y, s') = f x
          (dx, (_, _, Just (dw, db))) = bf s' (matrix [1..6] :: L 2 3)
          x = matrix [0,0,1,1] :: L 2 2
          w = matrix [0,0,0,10,10,10] :: L 2 3
          b = vec3 1 2 3
testRelu = dy
    where (Layer s f bf) = reluLayer
          (y, state) = f (matrix [1.0,-0.5,-2.0,3.0] :: L 2 2)
          (dy, _) = bf state (matrix [2.0,5.5,-2.0,3.0] :: L 2 2)
testSoftmax = softmax (matrix [1010, 1000, 990] :: L 1 3)
testCrossEntropyError = crossEntropyError t y
    where t = matrix [0,0,1,0,0,0,0,0,0,0] :: L 1 10
          y = matrix [0.1,0.05,0.6,0.0,0.05,0.1,0.0,0.1,0.0,0.0] :: L 1 10


{-# LANGUAGE FlexibleContexts, NamedFieldPuns #-}

module School.Unit.MultiNoulli
( multiNoulli ) where

import Conduit ((.|), ConduitM, MonadResource, ZipSource(..),
                getZipSource, mapC, mapM_C, mapMC)
import Control.Monad.Except (MonadError, throwError)
import Numeric.LinearAlgebra (Container, Element, Matrix, Vector, assoc, cols, fromRows,
                              rows, takeColumns, toColumns, toList, toLists)
import School.FileIO.FileHeader (FileHeader)
import School.FileIO.Source (source)
import School.Types.Error (Error)
import School.Types.LiftResult (LiftResult)
import School.Types.Slinky (Slinky(..), slinkyAppend, slinkySingleton)
import School.Unit.CostParams (CostParams(..))
import School.Unit.UnitActivation (UnitActivation(..))
import School.Unit.UnitForward (ForwardStack)
import School.Unit.UnitGradient (UnitGradient(..))
import School.Unit.CostFunction (SetupCost, CostFunction(..), defSetupCost)

toTarget :: (Element a, RealFrac a)
         => Vector a
         -> CostParams
toTarget =
  BatchClassTarget . map round . toList

compute :: (Element a, Fractional a, Num a)
        => UnitActivation a
        -> Slinky CostParams
        -> Either Error a
compute (BatchActivation input)
        (SNode (BatchClassTarget target) _) =
  let factor = (-1) / (fromIntegral . rows $ input)
  in Right
   . (*factor)
   . sum
   . zipWith (flip (!!)) target
   . toLists
   $ input
compute _ _ = Left "MultiNoulli expects batch activation and batch class target"

deriv :: (Container Vector a, Fractional a, Num a)
      => UnitActivation a
      -> Slinky CostParams
      -> Either Error (UnitGradient a)
deriv (BatchActivation input)
      (SNode (BatchClassTarget target) _) =
  let factor = (-1) / (fromIntegral . rows $ input)
      c = cols input
  in Right
   . BatchGradient
   . fromRows
   . map (\idx -> assoc c 0 [(idx, factor)])
   $ target
deriv _ _ = Left "MultiNoulli expects batch activation and batch class target"

prepare :: (Element a, RealFrac a, MonadError Error m)
        => Maybe FilePath
        -> ConduitM (ForwardStack a)
                    (ForwardStack a)
                    m
                    ()
prepare Nothing = mapMC $ \(activations, cParams) ->
  case activations of
    [BatchActivation input] -> do
      let c = cols input
      let activation = takeColumns (c - 1) input
      let newParams = toTarget . last . toColumns $ input
      return ([BatchActivation activation], slinkyAppend newParams cParams)
    _ -> throwError "MultiNoulli setup expects single batch activation"
prepare _ = mapC id

setupError :: (MonadError Error m) => ConduitM () (ForwardStack a) m ()
setupError =
  mapM_C . const $ throwError "Need both filepath and file header"

appendTarget :: CostParams -> Matrix a -> ForwardStack a
appendTarget costParams =
  flip (,) (slinkySingleton costParams) . pure . BatchActivation

setup :: (LiftResult m, MonadError Error m, MonadResource m)
      => Maybe FilePath
      -> Maybe FileHeader
      -> SetupCost a m
setup Nothing Nothing byteSource = defSetupCost byteSource
setup Nothing (Just _) _ = setupError
setup (Just _) Nothing _ = setupError
setup (Just path) (Just header) matrixSource =
  getZipSource $ targetZipSource <*> ZipSource matrixSource
  where targetSource = source header path
                    .| mapC BatchClassTarget
                    .| mapMC (return . appendTarget)
        targetZipSource = ZipSource targetSource

multiNoulli :: ( Container Vector a, LiftResult m, RealFrac a
               , MonadError Error m, MonadResource m )
            => Maybe FilePath
            -> Maybe FileHeader
            -> CostFunction a m
multiNoulli path header = CostFunction { computeCost = compute
                                       , derivCost = deriv
                                       , prepareCost = prepare path
                                       , setupCost = setup path header
                                       }


{-# LANGUAGE TemplateHaskell #-}
{-# LANGUAGE TypeSynonymInstances #-}
{-# LANGUAGE FlexibleInstances #-}

module FEECa.FiniteElementTest where

import Control.Monad    ( liftM )

import FEECa.Internal.Form hiding ( inner )
import FEECa.Internal.Spaces
import qualified FEECa.Internal.MultiIndex  as MI
import qualified FEECa.Internal.Vector      as V
import qualified FEECa.Internal.Simplex     as S

import FEECa.Utility.Combinatorics
import FEECa.Utility.Print
import FEECa.Utility.Utility

import FEECa.FiniteElementSpace
import qualified FEECa.Polynomial                 as P
import qualified FEECa.PolynomialDifferentialForm as DF
import qualified FEECa.Bernstein                  as B

import FEECa.Utility.Test
import FEECa.Internal.SimplexTest

import qualified Test.QuickCheck                 as Q
import           Test.QuickCheck  ( (==>) )

import qualified Numeric.LinearAlgebra.HMatrix    as M

--------------------------------------------------------------------------------
-- Random Finite Element Space
--------------------------------------------------------------------------------

max_n = 3
max_r = 3
max_k = 3

arbitraryPL :: (Int -> Int -> Simplex -> FiniteElementSpace)
            -> Q.Gen FiniteElementSpace
arbitraryPL cpl = do
  [n,r,k] <- mapM Q.choose [(1,max_n), (0,max_r), (0,max_k)]
  liftM (cpl r k) (arbitrarySimplex n)

arbitraryPrLk :: Q.Gen FiniteElementSpace
arbitraryPrLk = arbitraryPL PrLk

arbitraryPrmLk :: Q.Gen FiniteElementSpace
arbitraryPrmLk = arbitraryPL PrmLk

instance Q.Arbitrary FiniteElementSpace where
  arbitrary = arbitraryPrLk

n = 4

--------------------------------------------------------------------------------
-- Whitney Forms
--------------------------------------------------------------------------------

data WhitneyTest = WhitneyTest Simplex
  deriving Show

instance Q.Arbitrary WhitneyTest where
  arbitrary = do k <- Q.choose (1,n)
                 liftM WhitneyTest (arbitrarySubsimplex k n)

prop_whitney_integral :: WhitneyTest -> Bool
prop_whitney_integral (WhitneyTest t) =
    abs (DF.integrate t (whitneyForm t [0..k]))
      `eqNum` (1 / fromInt (factorial k))
  where k = S.topologicalDimension t

--------------------------------------------------------------------------------
-- Psi Forms
--------------------------------------------------------------------------------

data PsiTest = PsiTest Simplex Simplex MI.MultiIndex Vector
  deriving Show

instance Q.Arbitrary PsiTest where
  arbitrary = do
    k <- Q.choose (1,n)
    t <- arbitrarySimplex n
    f <- Q.elements $ S.subsimplices t k
    r <- Q.choose (0,10)
    mi <- liftM (MI.extend n (S.sigma f)) (arbitraryMI (k+1) r)
    liftM (PsiTest t f mi) (arbitraryVector n)

prop_psi :: PsiTest -> Q.Property
prop_psi (PsiTest t f mi v) =
    MI.degree mi > 0 ==> all (eqNum 0.0 . evaluate v) ps
  where ps  = [DF.apply (psi' t f mi i) [l] | i <- is, l <- axs]
        axs = projectionAxes t f mi
        is  = S.sigma f
        k   = S.topologicalDimension f

convexCombination :: Simplex -> MI.MultiIndex -> Vector
convexCombination t mi = sumV (zipWith sclV mi' vs)
  where vs    = S.vertices t
        zero  = zeroV (head vs)
        (l,r) = (MI.toList mi, fromInt $ MI.degree mi)
        mi'   = [ let i = fromInt (l !! s)
                  in i / r  -- XXX: mulInv???...
                | s <- S.sigma t]

projectionAxes :: Simplex -> Simplex -> MI.MultiIndex -> [Vector]
projectionAxes t f mi = map (subV xmi) (S.complement t f)
  where xmi = convexCombination f mi

prop_basis :: FiniteElementSpace -> Q.Property
prop_basis s =
    length bs > 0 ==> length bs == dim s && linearIndependent DF.inner bs
  where bs  = basis s
        -- XXX: why the normalisation?? ... removed for now, seems unnecessary
        --      ==> we use bs instead
        --bs' = map (\x -> fmap (sclV (mulInv (sqrt (DF.inner x x)))) x) bs
        --n   = vspaceDim s

--------------------------------------------------------------------------------
-- PrmLk Form associated to a face.
--------------------------------------------------------------------------------

linearIndependent :: Module v => (v -> v -> Double) -> [v] -> Bool
linearIndependent f bs = M.rank mat == n
  where es  = M.eigenvaluesSH' mat
        -- TODO: changed eigenvaluesSH to eigenvaluesSH' for loading; check!
        mat = M.matrix n [ f omega eta | omega <- bs, eta <- bs ]
        n   = length bs

--------------------------------------------------------------------------------
-- DoFs of PrLk space
--------------------------------------------------------------------------------

prop_dof_basis :: FiniteElementSpace -> Q.Property
prop_dof_basis s = dim s > 0 && degree s > 0 ==> M.rank mat == n
  where mat = M.matrix (n) [d b | d <- dofs s, b <- basis s]
        n   = dim s

prop_ndofs :: FiniteElementSpace -> Q.Property
prop_ndofs s = degree s > 0 ==> length (dofs s) == sum [nDofs s k | k <- [0..n]]
  where n = S.topologicalDimension $ simplex s

return []
testFiniteElement = $quickCheckWithAll

space = PrmLk 3 0 (S.referenceSimplex 3)


{-# LANGUAGE LambdaCase, BangPatterns #-}
module Stock where

import Data.Random.Normal
import System.Random
import Data.Random
import Diff
import Numeric.SpecFunctions
import Numeric.LinearAlgebra
import Util

ndigits :: Int
ndigits = 4

roundToN :: Int -> Double -> Double
roundToN n number = (fromInteger (round (number * (10^n)))) / (10.0^n)

-- constants
r :: Double
r = 0.06

variance :: Double
variance = 0.2

days :: Int
days = 360

yperiod :: Double
yperiod = (fromIntegral days) / 360.0

tdiff :: Int
tdiff = 40

tstep :: Double
tstep = (fromIntegral tdiff) / 360.0

s0 :: Double
s0 = 36.0000

k :: Double
k = 40.000

d1 :: Double
d1 = (log (s0 / k) + (r + variance / 2) * yperiod) / (sqrt (variance * yperiod))

d2 :: Double
d2 = d1 - sqrt (variance * yperiod)

-- c = s_0 * N(D_1) - exp (-r * T) * K * N(D_2)
n :: Double -> Double
n = cdf StdNormal

c :: Double
c = s0 * (n d1) - (exp (-r * yperiod)) * k * (n d2)

p :: Double
p = c + k * exp (-r * yperiod) - s0

--- basis functions
l :: Int -> String -> DiffLang
l n x
  | n >= 0    = (Exp ((Const (-0.5)) :*: (Var x))) :*: (Exp (Var x)) :*: (Const (1/(factorial n))) :*: (derivN n (laguerre n x) x)
  | otherwise = error "undefined basis function"

rzs :: Integer -> [Double]
rzs rand = normals' (0.0,1.0) (mkStdGen (fromIntegral rand))

sRand :: Integer -> Int -> Double
sRand j n =
  let zs = rzs j
  in s zs n

s :: [Double] -> Int -> R
s zs 0 = s0
s zs n
  | n > 0 = roundToN ndigits $ (s zs (n - 1)) * (exp $ (r - variance/2) * tstep + ((sqrt (variance * tstep)) * (zs!!n)))

cf :: [R] -> [R]
cf = map (\stdj -> max (k - stdj) 0)

cv :: R -> R
cv = (* (exp (-r * tstep)))

eval :: DiffLang -> R -> R
eval (Var _) x = x
eval (Const n) _ = n
eval (e1 :+: e2) x = eval e1 x + eval e2 x
eval (e1 :*: e2) x = eval e1 x * eval e2 x
eval (e1 :^: e2) x = eval e1 x ** eval e2 x
eval (Exp e) x     = exp (eval e x)

-- compute coefficients {beta_i} using (n-1)'th day's cash as predicator values,
-- n'th days cv as ys
regColN :: ([Vector R], [Vector R]) -> Int -> Matrix R
regColN (paths, cfs) n
  | length paths > n && n > 1 = regCol n laguerreBasis s t
  | otherwise = error "too big a N"
    where st :: ([R], [R], [R])
          st = unzip3 $ filter (\(a,b,c) -> a /= 0) $ zip3 (toList (cfs!!(n-1))) (toList (paths!!(n-1))) (toList (cfs!!n))
          s :: Vector R
          s = fromList $ sndOf3 st
          t :: Vector R
          t = fromList $ thdOf3 st

laguerreJ :: Int -> R -> R
laguerreJ j = eval (laguerre j "X")

laguerreBasis :: [R -> R]
laguerreBasis = map laguerreJ [0..(days `div` tdiff)]

regCol :: Int -> [R -> R] -> Vector R -> Vector R -> Matrix R -- n x 1 result
regCol n basis s t = linearSolveLS a b
  where a :: Matrix R
        -- a = fromColumns $ map (\j -> (fromList . map (basis!!j) . toList) s) [0..n-1]
        a = fromRows $ map (\x -> (fromList . map (\j -> (basis!!j) x)) [0..n-1]) (toList s)
        b :: Matrix R
        b = (asColumn . fromList . map cv) (toList t)

evalPoly :: R -> [R] -> R
evalPoly x betas = foldl (\acc beta -> acc + beta * x) 0 betas

evalPolyL :: [R] -> [R] -> [R]
evalPolyL xs betas = map (\x -> evalPoly x betas) xs

epsilon :: Double
epsilon = 0.00000001

cfP :: ([Vector R], [Vector R]) -> Int -> [Vector R]
cfP (paths, cfs) n
  | length paths > n && n > 1 = [newCFNP, finalCFN]
  | otherwise = error "undefined at the date!"
  where cfn = map (\p -> if fst p == 0 then 0 else snd p) $ zip oldCFP predCF
        finalCFN = fromList $ map (\(s,t) -> if s <= epsilon then 0.0 else t) (zip newCFN oldCFN)
        newCFN = map (\(s,t) -> if t >= s then t else 0.0) (zip oldCFP cfn)
        newCFNP = fromList $ map (\(s,t) -> if t >= s then 0.0 else s) (zip oldCFP cfn)
        -- newCFNP = fromList $ map (\(s,t) -> if t <= epsilon then 0.0 else s) (zip oldCFP cfn)
        oldCFN = toList $ cfs!!n
        oldCFP = toList $ cfs!!(n-1)
        predCF = map (\x -> (foldl (\acc -> \case { (n,beta) -> beta * (laguerreBasis!!n) x + acc }) 0 betas)) (toList (paths!!(n-1)))
        betas :: [(Int,R)]
        betas = zip [0..n-1] ((toList . head . toColumns . regColN (paths, cfs)) n)

cfGenRStep :: ([Vector R], [Vector R]) -> Int -> [Vector R]
cfGenRStep (pathsC, cfsC) n = take (n - 1) cfsC ++ cfP (pathsC, cfsC) n ++ drop (n + 1) cfsC

cfGenD :: (Matrix R, Matrix R) -> [Matrix R]
cfGenD (paths, cfs) = map (fromColumns) $ foldl (\cfsCC n ->
                               let cfs' = head cfsCC
                               in (take (n - 1) cfs' ++ cfP (pathsC, cfs') n ++ drop (n + 1) cfs') : cfsCC)
                      [toColumns cfs] (reverse [2..(days `div` tdiff)])
  where pathsC = toColumns paths

cfGenR :: (Matrix R, Matrix R) -> [Vector R]
cfGenR (paths, cfs) = foldl (\cfs' n -> take (n - 1) cfs' ++ cfP (pathsC, cfs') n ++ drop (n + 1) cfs') cfsC (reverse [2..(days `div` tdiff)])
  where pathsC = toColumns paths
        cfsC = toColumns cfs

tableaux :: Integer -> Integer -> [[R]]
tableaux rand size =
  map (\r -> map (sRand r) [0..(days `div` tdiff)]) ss
  where ss :: [Integer]
        ss = seeds rand size

seeds :: Integer -> Integer -> [Integer]
seeds rand length = take (fromIntegral length) $ map (\x -> (ceiling (abs x * 100000000000))) (rzs rand)

tableauxM :: Integer -> Integer -> (Matrix R, Matrix R)
tableauxM rand size = (matrix (days `div` tdiff + 1) paths, matrix (days `div` tdiff + 1) cfs)
  where paths = concat (tableaux rand size)
        cfs   = cf paths

cfGen :: (Matrix R, Matrix R) -> Matrix R
cfGen = fromColumns . tail . cfGenR

numNonZ :: Matrix R -> Int
numNonZ = sum . map (length . (filter (/= 0)) . toList) . toColumns

aOptionValue :: Matrix R -> [R]
aOptionValue m = map (\lst -> if null lst then 0.0 else (exp (-r * (fromIntegral (fst (last lst))) * tstep)) * (snd (last lst)))
                 $ map (filter (\case {(n,x) -> x /= 0}) . zip [1..(days `div` tdiff)] . toList) (toRows m)

eOptionValue :: Matrix R -> [R]
eOptionValue = map (* (exp (-r * (fromIntegral (days `div` tdiff + 1)) * tstep))) . cf . toList . last . toColumns

cStarMin :: Integer -> (Matrix R, Matrix R) -> R
cStarMin p (cf,m) = (/dist) . sum $ map (\i -> (vs!!(i-1) - vbar) * (qs!!(i-1) - qbar))  [1..(fromIntegral p)]
  where vs = aOptionValue cf
        qs = eOptionValue m
        vbar = sum vs / (fromIntegral (length vs))
        qbar = sum qs / (fromIntegral (length qs))
        dist = sum $ map (\i -> (qs!!(i-1) - qbar)**2) [1..(fromIntegral p)]

rbar :: Integer -> (Matrix R, Matrix R) -> R -> R
rbar pass (cf,m) c = (/(fromIntegral pass)) . sum $ map (\i -> (vs!!(i-1)) - c * ((qs!!(i-1)) - p)) [1..(fromIntegral pass)]
  where vs = aOptionValue cf
        qs = eOptionValue m
        qbar = sum qs / (fromIntegral (length qs))

output :: Integer -> Integer -> R
output rand size = rbar size (cfs, fst m) (cStarMin size (cfs, fst m))
  where cfs = cfGen m
        m = tableauxM rand size

genxs :: Integer -> [Double]
genxs n = map (\i -> let x = aOptionValue (cfGen (tableauxM i 100))
                     in sum x / (fromIntegral (length x)) ) [100..(100+n)]


module Sim.Sim
  ( stepOde,
  )
where

import Linear hiding (cross)
import Math.Frames
import Math.Spatial
import Math.Vectorize (Vectorize (..))
import Models.Aircraft
import Numeric.GSL.ODE (odeSolve)
import Numeric.LinearAlgebra (Vector (..), fromList, toList, toRows)
import Physics.RigidBody

runOde ::
  RigidBodyConstants ->
  V3T Body Double ->
  V3T Body Double ->
  AircraftState Double ->
  AircraftState Double
runOde constants forces moments state =
  AircraftState $ stepRigidBodyOde constants forces moments (dsRigidBody state)

-- | Steps the simulation forward by one timestep us RK45 ode solver.
--
-- This is pretty ugly because we have to convert back and forth from a list of
-- doubles in order to interact with GSL's ode library.
stepOde ::
  RigidBodyConstants ->
  V3T Body Double ->
  V3T Body Double ->
  AircraftState Double ->
  AircraftState Double
stepOde constants forces moments state = devectorize $ toList solution
  where
    solution :: Vector Double
    (_ : solution : _) =
      toRows $
        odeSolve
          wrapOde
          (vectorize state)
          (fromList [0.0, 0.01])

    wrapOde :: Double -> [Double] -> [Double]
    wrapOde t0 x0 =
      vectorize $
        runOde
          constants
          forces
          moments
          (devectorize x0)