-- |
-- Module      :  HashedExpression.Interp
-- Copyright   :  (c) OCA 2020
-- License     :  MIT (see the LICENSE file)
-- Maintainer  :  anandc@mcmaster.ca
-- Stability   :  provisional
-- Portability :  unportable
--
-- Evaluate expressions. Mainly useful for testings.
module HashedExpression.Interp where

import Data.Array
import Data.Complex hiding (conjugate)
import qualified Data.List as List
import qualified Data.Map as Map
import HashedExpression.Internal.Base hiding ((**))
import HashedExpression.Internal.Node
import HashedExpression.Utils
import HashedExpression.Value

data InterpValue
  = VR Double
  | VC (Complex Double)
  | V1DR (Array Int Double)
  | V1DC (Array Int (Complex Double))
  | V2DR (Array (Int, Int) Double)
  | V2DC (Array (Int, Int) (Complex Double))
  | V3DR (Array (Int, Int, Int) Double)
  | V3DC (Array (Int, Int, Int) (Complex Double))
  deriving (Show, Eq)

eval :: IsExpression e => ValMap -> e -> InterpValue
eval valMap e =
  let (mp, nID) = asRawExpr e
      (shape, et, op) = retrieveNode nID mp
      eval' :: NodeID -> InterpValue
      eval' x = eval valMap (mp, x)
      -------------------------------------------------------------------------------
      -- Partial helper functions
      constructR :: [Double] -> InterpValue
      constructR vs = case (shape, vs) of
        ([], [v]) -> VR v
        ([size], _) -> V1DR $ listArray (0, size - 1) vs
        ([size1, size2], _) -> V2DR $ listArray ((0, 0), (size1 - 1, size2 - 1)) vs
        ([size1, size2, size3], _) -> V3DR $ listArray ((0, 0, 0), (size1 - 1, size2 - 1, size3 - 1)) vs
      extractR :: InterpValue -> [Double]
      extractR val = case val of
        VR v -> [v]
        V1DR arr -> elems arr
        V2DR arr -> elems arr
        V3DR arr -> elems arr
      constructC :: [Complex Double] -> InterpValue
      constructC vs = case (shape, vs) of
        ([], [v]) -> VC v
        ([size], _) -> V1DC $ listArray (0, size - 1) vs
        ([size1, size2], _) -> V2DC $ listArray ((0, 0), (size1 - 1, size2 - 1)) vs
        ([size1, size2, size3], _) -> V3DC $ listArray ((0, 0, 0), (size1 - 1, size2 - 1, size3 - 1)) vs
      extractC :: InterpValue -> [Complex Double]
      extractC val = case val of
        VC v -> [v]
        V1DC arr -> elems arr
        V2DC arr -> elems arr
        V3DC arr -> elems arr
      unaryR op arg = arg |> eval' |> extractR |> map op |> constructR
      unaryC op arg = arg |> eval' |> extractC |> map op |> constructC
      binaryR op arg1 arg2 = zipWith op (extractR $ eval' arg1) (extractR $ eval' arg2) |> constructR
      binaryC op arg1 arg2 = zipWith op (extractC $ eval' arg1) (extractC $ eval' arg2) |> constructC
   in case op of
        Var name -> case (shape, Map.lookup name valMap) of
          ([], Just (VScalar val)) -> VR val
          ([_], Just (V1D arr)) -> V1DR arr
          ([_, _], Just (V2D arr)) -> V2DR arr
          ([_, _, _], Just (V3D arr)) -> V3DR arr
        Param name -> case (shape, Map.lookup name valMap) of
          ([], Just (VScalar val)) -> VR val
          ([_], Just (V1D arr)) -> V1DR arr
          ([_, _], Just (V2D arr)) -> V2DR arr
          ([_, _, _], Just (V3D arr)) -> V3DR arr
        Const v -> constructR $ replicate (product shape) v
        Sum args
          | et == R -> args |> map eval' |> map extractR |> foldl1 (zipWith (+)) |> constructR
          | et == C -> args |> map eval' |> map extractC |> foldl1 (zipWith (+)) |> constructC
        Mul args
          | et == R -> args |> map eval' |> map extractR |> foldl1 (zipWith (*)) |> constructR
          | et == C -> args |> map eval' |> map extractC |> foldl1 (zipWith (*)) |> constructC
        Power x arg
          | et == R -> unaryR (** fromIntegral x) arg
          | et == C -> unaryC (** fromIntegral x) arg
        Neg arg
          | et == R -> unaryR negate arg
          | et == C -> unaryC negate arg
        Scale arg1 arg2 -> case (retrieveElementType arg1 mp, retrieveElementType arg2 mp) of
          (R, R) ->
            let VR v = eval' arg1
             in unaryR (v *) arg2
          (R, C) ->
            let VR v = eval' arg1
             in unaryC ((v :+ 0) *) arg2
          (C, C) ->
            let VC v = eval' arg1
             in unaryC (v *) arg2
        Div arg1 arg2
          | et == R -> binaryR (/) arg1 arg2
          | et == C -> binaryC (/) arg1 arg2
        -------------------------------------------------------------------------------
        Sqrt arg -> unaryR sqrt arg
        Sin arg -> unaryR sin arg
        Cos arg -> unaryR cos arg
        Tan arg -> unaryR tan arg
        Exp arg -> unaryR exp arg
        Log arg -> unaryR log arg
        Sinh arg -> unaryR sinh arg
        Cosh arg -> unaryR cosh arg
        Tanh arg -> unaryR tanh arg
        Asin arg -> unaryR asin arg
        Acos arg -> unaryR acos arg
        Atan arg -> unaryR atan arg
        Asinh arg -> unaryR asinh arg
        Acosh arg -> unaryR acosh arg
        Atanh arg -> unaryR atanh arg
        -------------------------------------------------------------------------------
        RealImag arg1 arg2 ->
          let re = extractR (eval' arg1)
              im = extractR (eval' arg2)
           in constructC $ zipWith (:+) re im
        RealPart arg -> extractC (eval' arg) |> map realPart |> constructR
        ImagPart arg -> extractC (eval' arg) |> map imagPart |> constructR
        Conjugate arg -> unaryC conjugate arg
        -------------------------------------------------------------------------------
        InnerProd arg1 arg2
          | et == R ->
            let x = extractR (eval' arg1)
                y = extractR (eval' arg2)
             in VR $ sum $ zipWith (*) x y
          | et == C ->
            let x = extractC (eval' arg1)
                y = extractC (eval' arg2)
             in VC $ sum $ zipWith (*) x (map conjugate y)
        Piecewise marks conditionArg branchArgs
          | et == R ->
            let condition = extractR (eval' conditionArg)
                branches = map (extractR . eval') branchArgs
             in constructR $ zipWith (chooseBranch marks) condition (List.transpose branches)
          | et == C ->
            let condition = extractR (eval' conditionArg)
                branches = map (extractC . eval') branchArgs
             in constructC $ zipWith (chooseBranch marks) condition (List.transpose branches)
        Rotate rotateAmount arg
          | et == R -> case (rotateAmount, shape) of
            ([amount], [size]) ->
              let V1DR arr = eval' arg
               in V1DR $ rotate1D size amount arr
            ([amount1, amount2], [size1, size2]) ->
              let V2DR arr = eval' arg
               in V2DR $ rotate2D (size1, size2) (amount1, amount2) arr
            ([amount1, amount2, amount3], [size1, size2, size3]) ->
              let V3DR arr = eval' arg
               in V3DR $ rotate3D (size1, size2, size3) (amount1, amount2, amount3) arr
          | et == C -> case (rotateAmount, shape) of
            ([amount], [size]) ->
              let V1DC arr = eval' arg
               in V1DC $ rotate1D size amount arr
            ([amount1, amount2], [size1, size2]) ->
              let V2DC arr = eval' arg
               in V2DC $ rotate2D (size1, size2) (amount1, amount2) arr
            ([amount1, amount2, amount3], [size1, size2, size3]) ->
              let V3DC arr = eval' arg
               in V3DC $ rotate3D (size1, size2, size3) (amount1, amount2, amount3) arr
        FT arg -> case shape of
          [] -> eval' arg
          [size] ->
            let V1DC arr = eval' arg
             in V1DC $ fourierTransform1D FT_FORWARD size arr
          [size1, size2] ->
            let V2DC arr = eval' arg
             in V2DC $ fourierTransform2D FT_FORWARD (size1, size2) arr
          [size1, size2, size3] ->
            let V3DC arr = eval' arg
             in V3DC $ fourierTransform3D FT_FORWARD (size1, size2, size3) arr
        IFT arg -> case shape of
          [] -> eval' arg
          [size] ->
            let V1DC arr = eval' arg
             in V1DC $ fourierTransform1D FT_BACKWARD size arr
          [size1, size2] ->
            let V2DC arr = eval' arg
             in V2DC $ fourierTransform2D FT_BACKWARD (size1, size2) arr
          [size1, size2, size3] ->
            let V3DC arr = eval' arg
             in V3DC $ fourierTransform3D FT_BACKWARD (size1, size2, size3) arr
        Project dss arg -> case shape of
          [] -> case (retrieveShape arg mp, dss) of
            ([size], [At i]) -> case eval' arg of
              V1DR base -> VR $ base ! i
              V1DC base -> VC $ base ! i
            ([size1, size2], [At i, At j]) -> case eval' arg of
              V2DR base -> VR $ base ! (i, j)
              V2DC base -> VC $ base ! (i, j)
            ([size1, size2, size3], [At i, At j, At k]) -> case eval' arg of
              V3DR base -> VR $ base ! (i, j, k)
              V3DC base -> VC $ base ! (i, j, k)
          [size] -> case (retrieveShape arg mp, dss) of
            ([bSize], [ds]) -> case eval' arg of
              V1DR base -> V1DR $ listArray (0, size - 1) [base ! i | i <- mkIndices ds bSize]
              V1DC base -> V1DC $ listArray (0, size - 1) [base ! i | i <- mkIndices ds bSize]
            ([bSize1, bSize2], [ds1, ds2]) -> case eval' arg of
              V2DR base ->
                V1DR $
                  listArray (0, size - 1) $
                    [base ! (i, j) | i <- mkIndices ds1 bSize1, j <- mkIndices ds2 bSize2]
              V2DC base ->
                V1DC $
                  listArray (0, size - 1) $
                    [base ! (i, j) | i <- mkIndices ds1 bSize1, j <- mkIndices ds2 bSize2]
            ([bSize1, bSize2, bSize3], [ds1, ds2, ds3]) -> case eval' arg of
              V3DR base ->
                V1DR $
                  listArray (0, size - 1) $
                    [base ! (i, j, k) | i <- mkIndices ds1 bSize1, j <- mkIndices ds2 bSize2, k <- mkIndices ds3 bSize3]
              V3DC base ->
                V1DC $
                  listArray (0, size - 1) $
                    [base ! (i, j, k) | i <- mkIndices ds1 bSize1, j <- mkIndices ds2 bSize2, k <- mkIndices ds3 bSize3]
          [size1, size2] -> case (retrieveShape arg mp, dss) of
            ([bSize1, bSize2], [ds1, ds2]) -> case eval' arg of
              V2DR base ->
                V2DR $
                  listArray
                    ((0, 0), (size1 - 1, size2 - 1))
                    [base ! (i, j) | i <- mkIndices ds1 bSize1, j <- mkIndices ds2 bSize2]
              V2DC base ->
                V2DC $
                  listArray
                    ((0, 0), (size1 - 1, size2 - 1))
                    [base ! (i, j) | i <- mkIndices ds1 bSize1, j <- mkIndices ds2 bSize2]
            ([bSize1, bSize2, bSize3], [ds1, ds2, ds3]) -> case eval' arg of
              V3DR base ->
                V2DR $
                  listArray
                    ((0, 0), (size1 - 1, size2 - 1))
                    [base ! (i, j, k) | i <- mkIndices ds1 bSize1, j <- mkIndices ds2 bSize2, k <- mkIndices ds3 bSize3]
              V3DC base ->
                V2DC $
                  listArray
                    ((0, 0), (size1 - 1, size2 - 1))
                    [base ! (i, j, k) | i <- mkIndices ds1 bSize1, j <- mkIndices ds2 bSize2, k <- mkIndices ds3 bSize3]
          [size1, size2, size3] -> case (retrieveShape arg mp, dss) of
            ([bSize1, bSize2, bSize3], [ds1, ds2, ds3]) -> case eval' arg of
              V3DR base ->
                V3DR $
                  listArray
                    ((0, 0, 0), (size1 - 1, size2 - 1, size3 - 1))
                    [base ! (i, j, k) | i <- mkIndices ds1 bSize1, j <- mkIndices ds2 bSize2, k <- mkIndices ds3 bSize3]
              V3DC base ->
                V3DC $
                  listArray
                    ((0, 0, 0), (size1 - 1, size2 - 1, size3 - 1))
                    [base ! (i, j, k) | i <- mkIndices ds1 bSize1, j <- mkIndices ds2 bSize2, k <- mkIndices ds3 bSize3]
        Inject dss subArg baseArg
          | et == R ->
            let injectingElements = extractR $ eval' subArg
             in case (eval' baseArg, dss, shape) of
                  (V1DR base, [ds], [size]) ->
                    let indices = mkIndices ds size
                     in V1DR $ base // zip indices injectingElements
                  (V2DR base, [ds1, ds2], [size1, size2]) ->
                    let indices = [(i, j) | i <- mkIndices ds1 size1, j <- mkIndices ds2 size2]
                     in V2DR $ base // zip indices injectingElements
                  (V3DR base, [ds1, ds2, ds3], [size1, size2, size3]) ->
                    let indices = [(i, j, k) | i <- mkIndices ds1 size1, j <- mkIndices ds2 size2, k <- mkIndices ds3 size3]
                     in V3DR $ base // zip indices injectingElements
          | et == C ->
            let injectingElements = extractC $ eval' subArg
             in case (eval' baseArg, dss, shape) of
                  (V1DC base, [ds], [size]) ->
                    let indices = mkIndices ds size
                     in V1DC $ base // zip indices injectingElements
                  (V2DC base, [ds1, ds2], [size1, size2]) ->
                    let indices = [(i, j) | i <- mkIndices ds1 size1, j <- mkIndices ds2 size2]
                     in V2DC $ base // zip indices injectingElements
                  (V3DC base, [ds1, ds2, ds3], [size1, size2, size3]) ->
                    let indices = [(i, j, k) | i <- mkIndices ds1 size1, j <- mkIndices ds2 size2, k <- mkIndices ds3 size3]
                     in V3DC $ base // zip indices injectingElements
        MatMul arg1 arg2
          | et == R ->
            case (shape, retrieveShape arg1 mp, retrieveShape arg2 mp) of
              ([_m], [m, _n], [n]) ->
                let V2DR x = eval' arg1
                    V1DR y = eval' arg2
                 in V1DR $
                      listArray
                        (0, m - 1)
                        [ sum [(x ! (i, j)) * (y ! j) | j <- [0 .. n - 1]]
                          | i <- [0 .. m - 1]
                        ]
              ([_m, _p], [m, _n], [n, p]) ->
                let V2DR x = eval' arg1
                    V2DR y = eval' arg2
                 in V2DR $
                      listArray
                        ((0, 0), (m - 1, p - 1))
                        [ sum [(x ! (i, k)) * (y ! (k, j)) | k <- [0 .. n - 1]]
                          | i <- [0 .. m - 1],
                            j <- [0 .. p - 1]
                        ]
          | et == C ->
            case (shape, retrieveShape arg1 mp, retrieveShape arg2 mp) of
              ([_m], [m, _n], [n]) ->
                let V2DC x = eval' arg1
                    V1DC y = eval' arg2
                 in V1DC $
                      listArray
                        (0, m - 1)
                        [ sum [(x ! (i, j)) * (y ! j) | j <- [0 .. n - 1]]
                          | i <- [0 .. m - 1]
                        ]
              ([_m, _p], [m, _n], [n, p]) ->
                let V2DC x = eval' arg1
                    V2DC y = eval' arg2
                 in V2DC $
                      listArray
                        ((0, 0), (m - 1, p - 1))
                        [ sum [(x ! (i, k)) * (y ! (k, j)) | k <- [0 .. n - 1]]
                          | i <- [0 .. m - 1],
                            j <- [0 .. p - 1]
                        ]
        Transpose arg
          | et == R ->
            case (retrieveShape arg mp, shape) of
              -- ([m], [1, _m]) ->
              --   let V1DR x = eval' arg
              --    in V2DR $ listArray ((0, 0), (1, m - 1)) $ elems x
              ([m, n], [_n, _m]) ->
                let V2DR x = eval' arg
                 in V2DR $ listArray ((0, 0), (n - 1, m - 1)) [x ! (j, i) | i <- [0 .. n - 1], j <- [0 .. m - 1]]
          | et == C ->
            case (retrieveShape arg mp, shape) of
              -- ([m], [1, _m]) ->
              --   let V1DC x = eval' arg
              --    in V2DC $ listArray ((0, 0), (1, m - 1)) $ elems x
              ([m, n], [_n, _m]) ->
                let V2DC x = eval' arg
                 in V2DC $ listArray ((0, 0), (n - 1, m - 1)) [x ! (j, i) | i <- [0 .. n - 1], j <- [0 .. m - 1]]
        Coerce _ arg
          | et == R -> constructR . extractR $ eval' arg
          | et == C -> constructC . extractC $ eval' arg
        node -> error $ show node

zipWithA :: Ix x => (a -> b -> c) -> Array x a -> Array x b -> Array x c
zipWithA f xs ys = listArray (bounds xs) $ zipWith f (elems xs) (elems ys)

-- | Choose branch base on condition value.
-- In Decision tree, there are 2 possible outcomes, Head and Tail.
-- The decision of being Head or Tail is made based on the the condition value.
chooseBranch :: [Double] -> Double -> [a] -> a
chooseBranch marks val branches
  | val < head marks = head branches
  | otherwise =
    snd . last . filter ((val >=) . fst) $ zip marks (tail branches)

-- NOTE: `mod` in Haskell with negative number, e.g, (-5) `mod` 3 = 1

-- | One dimension rotation.
--   The elemnts falling off of the length of the 1D array will appear at the beginning of the array
rotate1D ::
  -- | Size of the input array
  Int ->
  -- | amount of rotation
  Int ->
  -- | Input array
  Array Int a ->
  -- | Rotated array
  Array Int a
rotate1D size amount arr =
  listArray
    (0, size - 1)
    [arr ! ((i - amount) `mod` size) | i <- [0 .. size - 1]]

-- | Two dimension rotation.
--   The elements falling off of the length of the 2D array will appear at the beginning of the each row or column of the array
rotate2D ::
  -- | Size of the 2d input array
  (Int, Int) ->
  -- | amount of rotation for 2d array
  (Int, Int) ->
  -- | Input 2d array
  Array (Int, Int) a ->
  -- | Rotated 2d array
  Array (Int, Int) a
rotate2D (size1, size2) (amount1, amount2) arr =
  listArray
    ((0, 0), (size1 - 1, size2 - 1))
    [ arr ! ((i - amount1) `mod` size1, (j - amount2) `mod` size2)
      | i <- [0 .. size1 - 1],
        j <- [0 .. size2 - 1]
    ]

-- | Three dimension rotation
rotate3D ::
  -- | Size of 3d input array
  (Int, Int, Int) ->
  -- | Amount of Rotation for 3d array
  (Int, Int, Int) ->
  -- | Input 3d array
  Array (Int, Int, Int) a ->
  -- | Rotated 3d array
  Array (Int, Int, Int) a
rotate3D (size1, size2, size3) (amount1, amount2, amount3) arr =
  listArray
    ((0, 0, 0), (size1 - 1, size2 - 1, size3 - 1))
    [ arr
        ! ( (i - amount1) `mod` size1,
            (j - amount2) `mod` size2,
            (k - amount3) `mod` size3
          )
      | i <- [0 .. size1 - 1],
        j <- [0 .. size2 - 1],
        k <- [0 .. size3 - 1]
    ]

data FTMode = FT_FORWARD | FT_BACKWARD deriving (Eq, Ord)

-- | Fourier Transform in 1D.
--  Frequency is just in one dimension.
--  Consider a real-valued function, S(x),
--  that is integrable on an interval of P, which will be the period of the Fourier series.
--  number of cycles is n.
--  length of cycle is P/n, and frequency is n/P.
--  so for input i the frequency is (2*pi*i*n)/P
fourierTransform1D ::
  FTMode -> Int -> Array Int (Complex Double) -> Array Int (Complex Double)
fourierTransform1D mode size arr =
  listArray (0, size - 1) [computeX i | i <- [0 .. size - 1]]
  where
    s = if mode == FT_BACKWARD then fromIntegral size else 1
    computeX i = (sum $ zipWithA (*) arr (fourierBasis i)) / s
    fourierBasis i =
      let frequency n = (2 * pi * fromIntegral (i * n) / fromIntegral size) * (if mode == FT_BACKWARD then -1 else 1)
       in listArray
            (0, size - 1)
            [ cos (frequency n) :+ (- sin (frequency n))
              | n <- [0 .. size - 1]
            ]

-- | Fourier Transform in 2D
--  the frequency should be calculated in 2D
--  Consider a real-valued function, S(x),
--  that is integrable on an interval of P, which will be the period of the Fourier series.
--  numbber of cycles is n.
--  length of cycle is P/n, and frequency is n/P.
--  so for input i the frequency is (2*pi*i*n)/P
--  the frequency should be calculated in both dimensions for i and j
fourierTransform2D ::
  FTMode ->
  (Int, Int) ->
  Array (Int, Int) (Complex Double) ->
  Array (Int, Int) (Complex Double)
fourierTransform2D mode (size1, size2) arr =
  listArray
    ((0, 0), (size1 - 1, size2 - 1))
    [computeX i j | i <- [0 .. size1 - 1], j <- [0 .. size2 - 1]]
  where
    s = if mode == FT_BACKWARD then fromIntegral (size1 * size2) else 1
    computeX i j = (sum $ zipWithA (*) arr (fourierBasis i j)) / s
    fourierBasis i j =
      let frequency m n =
            ( 2 * pi * fromIntegral (i * m) / fromIntegral size1
                + 2 * pi * fromIntegral (j * n) / fromIntegral size2
            )
              * (if mode == FT_BACKWARD then -1 else 1)
       in listArray
            ((0, 0), (size1 - 1, size2 - 1))
            [ cos (frequency m n) :+ (- sin (frequency m n))
              | m <- [0 .. size1 - 1],
                n <- [0 .. size2 - 1]
            ]

-- | Fourier Transform in 3D
--   the frequency should be calculated in 3D
--   Consider a real-valued function, S(x),
--   that is integrable on an interval of P, which will be the period of the Fourier series.
--   numbber of cycles is n.
--   length of cycle is P/n, and frequency is n/P.
--   so for input i the frequency is (2*pi*i*n)/P
--   the frequency should be calculated for all dimensions, i , j , k
fourierTransform3D ::
  FTMode ->
  (Int, Int, Int) ->
  Array (Int, Int, Int) (Complex Double) ->
  Array (Int, Int, Int) (Complex Double)
fourierTransform3D mode (size1, size2, size3) arr =
  listArray
    ((0, 0, 0), (size1 - 1, size2 - 1, size3 - 1))
    [ computeX i j k
      | i <- [0 .. size1 - 1],
        j <- [0 .. size2 - 1],
        k <- [0 .. size3 - 1]
    ]
  where
    s = if mode == FT_BACKWARD then fromIntegral (size1 * size2) else 1
    computeX i j k = (sum $ zipWithA (*) arr (fourierBasis i j k)) / s
    fourierBasis i j k =
      let frequency m n p =
            ( 2 * pi * fromIntegral (i * m) / fromIntegral size1
                + 2 * pi * fromIntegral (j * n) / fromIntegral size2
                + 2 * pi * fromIntegral (k * p) / fromIntegral size3
            )
              * (if mode == FT_BACKWARD then -1 else 1)
       in listArray
            ((0, 0, 0), (size1 - 1, size2 - 1, size3 - 1))
            [ cos (frequency m n p) :+ (- sin (frequency m n p))
              | m <- [0 .. size1 - 1],
                n <- [0 .. size2 - 1],
                p <- [0 .. size3 - 1]
            ]


{-# LANGUAGE CPP                   #-}
{-# LANGUAGE DataKinds             #-}
{-# LANGUAGE DeriveGeneric         #-}
{-# LANGUAGE FlexibleContexts      #-}
{-# LANGUAGE FlexibleInstances     #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE RankNTypes            #-}
{-# LANGUAGE ScopedTypeVariables   #-}
{-# LANGUAGE TypeFamilies          #-}
{-# LANGUAGE UndecidableInstances  #-}
{-# LANGUAGE TypeOperators         #-}
{-# LANGUAGE OverloadedStrings     #-}
{-# LANGUAGE OverloadedLabels      #-}

module Grenade.Dynamic.Layers.FullyConnected
  ( SpecFullyConnected (..)
  , specFullyConnected
  , fullyConnected
  ) where

import           Data.Proxy
import           Data.Reflection                     (reifyNat)
import           Data.Singletons
import           Data.Singletons.Prelude.Num         ((%*))
import           Data.Singletons.TypeLits            hiding (natVal)
import           GHC.TypeLits

import           Numeric.LinearAlgebra.Static        hiding (build, toRows, (&),
                                                      (|||), size)

import           Grenade.Core
import           Grenade.Utils.ListStore
import           Grenade.Dynamic.Internal.Build
import           Grenade.Dynamic.Specification
import           Grenade.Layers.FullyConnected

-------------------- DynamicNetwork instance --------------------

instance (KnownNat i, KnownNat o) => FromDynamicLayer (FullyConnected i o) where
  fromDynamicLayer _ _ _ = SpecNetLayer $ SpecFullyConnected (natVal (Proxy :: Proxy i)) (natVal (Proxy :: Proxy o))

instance ToDynamicLayer SpecFullyConnected where
  toDynamicLayer wInit gen (SpecFullyConnected nrI nrO) =
    reifyNat nrI $ \(pxInp :: (KnownNat i) => Proxy i) ->
      reifyNat nrO $ \(pxOut :: (KnownNat o') => Proxy o') ->
        case singByProxy pxInp %* singByProxy pxOut of
          SNat -> do
            (layer :: FullyConnected i o') <- randomFullyConnected wInit gen
            return $ SpecLayer layer (sing :: Sing ('D1 i)) (sing :: Sing ('D1 o'))

-- | Make a specification of a fully connected layer (see Grenade.Dynamic.Build for a user-interface to specifications).
specFullyConnected :: Integer -> Integer -> SpecNet
specFullyConnected nrI nrO = SpecNetLayer $ SpecFullyConnected nrI nrO

-- | A Fully-connected layer with input dimensions as given in last output layer and output dimensions specified. 1D only!
fullyConnected :: Integer -> BuildM ()
fullyConnected rows = do
  (inRows, _, _) <- buildRequireLastLayerOut Is1D
  buildAddSpec (SpecNetLayer $ SpecFullyConnected inRows rows)
  buildSetLastLayer (rows, 1, 1)

-------------------- GNum instances --------------------

instance (KnownNat i, KnownNat o) => GNum (FullyConnected i o) where
  s |* FullyConnected w store = FullyConnected (s |* w) (s |* store)
  FullyConnected w1 store1 |+ FullyConnected w2 store2 = FullyConnected (w1 |+ w2) (store1 |+ store2)
  gFromRational r = FullyConnected (gFromRational r) mkListStore

instance (KnownNat i, KnownNat o) => GNum (FullyConnected' i o) where
  s |* FullyConnected' b w = FullyConnected' (dvmap (fromRational s *) b) (dmmap (fromRational s *) w)
  FullyConnected' b1 w1 |+ FullyConnected' b2 w2 = FullyConnected' (b1 + b2) (w1 + w2)
  gFromRational r = FullyConnected' (fromRational r) (fromRational r)


{-# LANGUAGE DataKinds #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE TemplateHaskell #-}
-- | a variation on <http://hackage.haskell.org/package/hmatrix-syntax>,
-- which constructs a matrix which has dimensions stored.
module Numeric.LinearAlgebra.Dimensional.DK.QuasiQuotes
(
  vec,
  vecShape,
  mat,
) where

import Language.Haskell.TH

import Language.Haskell.TH as TH
import Language.Haskell.TH.Quote as TH
import Language.Haskell.TH.Syntax as TH
import Language.Haskell.TH.Quote
import Language.Haskell.Meta.Parse as Meta

import Numeric.Units.Dimensional.DK (Quantity, Dimension, DLength, DMass)
import Numeric.LinearAlgebra.Dimensional.DK.Internal (DimMat(..), vecSingleton, vecCons, fromRowVector, vconcat, vconcat', vconcat'')
import Numeric.LinearAlgebra.Dimensional.DK.Shapes

import Data.List.Split

vec = QuasiQuoter {
  quoteExp = parseVectorExp,
  quotePat = error "vec",
  quoteDec = error "vec",
  quoteType = parseVectorType
}

vecShape = QuasiQuoter {
  quoteExp = error "vecShape",
  quotePat = error "vecShape",
  quoteDec = error "vecShape",
  quoteType = parseVectorShapeType
}

mat = QuasiQuoter {
  quoteExp = parseMatrixExp,
  quotePat = error "mat",
  quoteDec = error "mat",
  quoteType = error "mat"
}

parseMatrixExp :: String -> Q Exp
parseMatrixExp s = let rs = fmap parseVectorExp $ splitSemicolonList s
                    in makeMatrixExp rs

parseVectorExp :: String -> Q Exp
parseVectorExp s = let qs = fmap parseQuantityExp $ splitCommaList s
                    in makeVectorExp qs

parseQuantityExp :: String -> Q Exp
parseQuantityExp s = do
                       case (Meta.parseExp s) of
                         Left err -> fail $ show err
                         Right e -> return e

parseVectorType :: String -> Q Type
parseVectorType s = [t| DimMat $(parseVectorShapeType s) |]

parseVectorShapeType :: String -> Q Type
parseVectorShapeType s = let dimTypes = fmap parseDimensionType $ splitCommaList s
                          in makeVectorShapeType dimTypes

parseDimensionType :: String -> Q Type
parseDimensionType s = do
                          case (Meta.parseType s) of
                            Left err -> fail $ show err
                            Right t -> return t

makeMatrixExp :: [Q Exp] -> Q Exp
makeMatrixExp [] = fail "Empty matrices not permitted."
makeMatrixExp [e] = [| fromRowVector $(e) |]
--makeMatrixExp (e:[e2]) = [| vconcat'' $(e) $(e2) |]
makeMatrixExp (e:es) = [| vconcat' $(e) $(makeMatrixExp es) |]

makeVectorExp :: [Q Exp] -> Q Exp
makeVectorExp [] = fail "Empty vectors not permitted."
makeVectorExp [e] = [| vecSingleton $(e) |]
makeVectorExp (e:es) = [| vecCons $(e) $(makeVectorExp es) |]

makeVectorShapeType :: [Q Type] -> Q Type
makeVectorShapeType (d : ds) = [t| 'VectorShape $(d) $(makeTypeLevelList ds) |]
makeVectorShapeType _ = fail "Empty vectors not permitted."

makeTypeLevelList :: [Q Type] -> Q Type
makeTypeLevelList (t : ts) = [t| $(t) ': $(makeTypeLevelList ts) |]
makeTypeLevelList [] = [t| '[] |]

splitCommaList :: String -> [String]
splitCommaList = splitOn ","

splitSemicolonList :: String -> [String]
splitSemicolonList = splitOn ";"


{-# LANGUAGE KindSignatures #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE FunctionalDependencies #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE DeriveGeneric #-}
{-# LANGUAGE DeriveAnyClass #-}
{-# LANGUAGE ViewPatterns #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE UnicodeSyntax #-}
{-# LANGUAGE AllowAmbiguousTypes #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE TemplateHaskell #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE CPP #-}

module NQS.SR
  ( -- * Stochastic Reconfiguration
    --
    -- Globally, the SR algorithm looks more or less like this
    --
    -- @
    --     for n in {0,...,maxIter - 1} do
    --         (F, ∂ψ) <- sample ψ
    --         S <- covariance⁽ⁿ⁾ ∂ψ
    --         δ <- (S + λ⁽ⁿ⁾)⁻¹ F
    --         ψ <- ψ - learningRate * δ
    --     done
    -- @

    -- * AXPY
    -- * Sampling
    -- * Constructing S
    -- * Solving S
    sr
  , IterInfo(..)
  , HasIter(..)
  , HasSampler(..)
  , HasSolver(..)
  , HasMoments(..)
  , HasForceNorm(..)
  ) where

import           Prelude                 hiding ( zipWith
                                                , zipWithM
                                                , map
                                                , mapM
                                                )

import           GHC.Generics                   ( Generic )

import           Debug.Trace
import           Control.Exception              ( assert )
import           Control.Monad.Identity         ( Identity(..) )
import           Control.Monad.ST
import           System.IO.Unsafe               ( unsafePerformIO )
import           Foreign.Storable
import           Foreign.ForeignPtr
import           Data.Vector.Storable           ( Vector )
import qualified Data.Vector.Storable          as V
import           Data.Vector.Storable.Mutable   ( MVector )
import qualified Data.Vector.Storable.Mutable  as MV
import qualified Data.Vector.Unboxed

import           Data.Singletons
import           Data.Complex
import           Data.Semigroup                 ( (<>) )
import           Control.Monad                  ( (>=>) )
import           Control.Monad.Primitive

import           Control.DeepSeq
import           System.CPUTime
import Data.Aeson

import qualified NQS.CG                        as CG
import           NQS.CG                         ( Operator )
import           NQS.Rbm (Rbm(..))
import           NQS.Rbm.Mutable
import           NQS.Internal.BLAS
import           NQS.Internal.LAPACK
import           NQS.Internal.Types
import           NQS.Internal.Hamiltonian
import           NQS.Internal.Sampling
import           NQS.Internal.Rbm (unsafeFreezeRbm)

import           Lens.Micro
import           Lens.Micro.TH
import           Lens.Micro.Extras

import           GHC.Float                      ( int2Float )

import           Data.Aeson
import qualified Data.ByteString.Lazy.Char8     as BS

-- | In Stochastic Reconfiguration we only ever deal with 'Direct' vectors.
type V = MDenseVector 'Direct

-- | A shorter name for dense matrices.
type M orient s a = MDenseMatrix orient s a

data SolverStats = SolverStats
  { _solverStatsIters :: {-# UNPACK #-}!Int
  , _solverStatsErr   :: {-# UNPACK #-}!ℝ
  , _solverStatsTime  :: {-# UNPACK #-}!Double
  }

makeFields ''SolverStats

instance ToJSON SolverStats where
  toJSON stats =
    object [ "iters" .= (stats ^. iters)
           , "error" .= (stats ^. err)
           , "time"  .= (stats ^. time)
           ]
  toEncoding stats =
    pairs $ "iters" .= (stats ^. iters)
         <> "error" .= (stats ^. err)
         <> "time"  .= (stats ^. time)

instance FromJSON SolverStats where
  parseJSON = withObject "SolverStats" $ \v ->
    SolverStats <$> v .: "iters"
                <*> v .: "error"
                <*> v .: "time"

data SamplerStats = SamplerStats
  { _samplerStatsMoments  :: {-# UNPACK #-}!(Vector ℂ)
  , _samplerStatsStdDev   :: !(Maybe ℝ)
  , _samplerStatsDim      :: {-# UNPACK #-}!Int
  , _samplerStatsTime     :: {-# UNPACK #-}!Double
  }

makeFields ''SamplerStats

instance ToJSON SamplerStats where
  toJSON stats =
    object [ "moments" .= (stats ^. moments)
           , "stddev"  .= (stats ^. stdDev)
           , "dim"     .= (stats ^. dim)
           , "time"    .= (stats ^. time)
           ]
  toEncoding stats =
    pairs $ "moments" .= (stats ^. moments)
         <> "stddev"  .= (stats ^. stdDev)
         <> "dim"     .= (stats ^. dim)
         <> "time"    .= (stats ^. time)

instance FromJSON SamplerStats where
  parseJSON = withObject "SamplerStats" $ \v ->
    SamplerStats <$> v .: "moments"
                 <*> v .: "stddev"
                 <*> v .: "dim"
                 <*> v .: "time"

data IterInfo = IterInfo
  { _iterInfoIter      :: {-# UNPACK #-}!Int
  , _iterInfoState     :: Rbm
  , _iterInfoSampler   :: !SamplerStats
  , _iterInfoSolver    :: !SolverStats
  , _iterInfoForceNorm :: {-# UNPACK #-}!ℝ
  , _iterInfoDeltaNorm :: {-# UNPACK #-}!ℝ
  }

makeFields ''IterInfo

instance ToJSON IterInfo where
  toJSON stats =
    object [ "iter"      .= (stats ^. iter)
           , "state"     .= (stats ^. state)
           , "sampler"   .= (stats ^. sampler)
           , "solver"    .= (stats ^. solver)
           , "forceNorm" .= (stats ^. forceNorm)
           , "deltaNorm" .= (stats ^. deltaNorm)
           ]
  toEncoding stats =
    pairs $ "iter"      .= (stats ^. iter)
         <> "state"     .= (stats ^. state)
         <> "sampler"   .= (stats ^. sampler)
         <> "solver"    .= (stats ^. solver)
         <> "forceNorm" .= (stats ^. forceNorm)
         <> "deltaNorm" .= (stats ^. deltaNorm)

instance FromJSON IterInfo where
  parseJSON = withObject "IterInfo" $ \v ->
    IterInfo <$> v .: "iter"
             <*> v .: "state"
             <*> v .: "sampler"
             <*> v .: "solver"
             <*> v .: "forceNorm"
             <*> v .: "deltaNorm"

-- AXPY
-- -----------------------------------------------------------------------------

type AxpyFun v m = ℂ -> V (PrimState m) ℂ -> v (PrimState m) -> m ()

-- | \(Y \leftarrow \alpha X + Y\) for 'Column' matrices.
axpyMatrix
  :: (PrimMonad m, Storable a, AXPY (MDenseVector 'Direct) a)
  => a
  -> MDenseVector 'Direct (PrimState m) a
  -> MDenseMatrix 'Column (PrimState m) a
  -> m ()
axpyMatrix !α !x !y@(MDenseMatrix !rows !cols !stride !buff) =
  assert (rows * cols == x ^. dim) $ loop 0
 where
  getX !i = slice (i * rows) rows x
  getY !i = unsafeColumn i y
  loop !i | i < cols  = axpy α (getX i) (getY i) >> loop (i + 1)
          | otherwise = return ()

-- | Performs @y <- αx + y@ where @x@ is a vector and @y@ -- an RBM.
axpyRbm :: AxpyFun MRbm IO
axpyRbm α x y =
  let n = sizeVisible y
      m = sizeHidden y
  in  assert (x ^. dim == size y) $ do
        withVisible y $ \a -> do
          axpy α (slice 0 n x) a
          -- unsafeFreeze a >>= \a' -> print (a' ^. buffer)
        withHidden y $ \b -> do
          axpy α (slice n m x) b
          -- unsafeFreeze b >>= \b' -> print (b' ^. buffer)
        withWeights y $ \w -> do
          axpyMatrix α (slice (n + m) (n * m) x) w
          -- unsafeFreeze w >>= \w' -> print (w' ^. buffer)

-- Monte-Carlo sampling
-- -----------------------------------------------------------------------------

-- | Sampling
type SampleFun v m
  = v (PrimState m)
 -> V (PrimState m) ℂ
 -> M 'Row (PrimState m) ℂ
 -> m SamplerStats

-- | Monte-Carlo sampling an RBM
sampleRbm :: MCConfig -> Hamiltonian -> SampleFun MRbm IO
sampleRbm config hamiltonian = doSample
 where
  doSample rbm force derivatives = do
    t1                    <- getCPUTime
    moments               <- MV.new 2
    (dimension, variance) <- sampleGradients config
                                             hamiltonian
                                             rbm
                                             moments
                                             force
                                             derivatives
    t2 <- getCPUTime
    V.unsafeFreeze moments >>= \moments' -> return
      (SamplerStats moments'
                    (sqrt <$> variance)
                    dimension
                    (fromIntegral (t2 - t1) * 1.0E-12)
      )

type SolveFun wrapper m
  = wrapper -- ^ S
 -> V (PrimState m) ℂ -- ^ b
 -> V (PrimState m) ℂ -- ^ x
 -> m SolverStats

type MakeSFun wrapper m = Int -> M 'Row (PrimState m) ℂ -> m wrapper

-- | Sparse representation of the @S@ matrix.
data SMatrix s =
  SMatrix !(M 'Row s ℂ) -- ^ Derivatives (O - 〈O〉)
          !(V s ℂ) -- ^ Workspace of size #steps.
          !(Maybe ℂ) -- ^ Regulariser λ

makeS :: PrimMonad m => Maybe (Int -> ℂ) -> MakeSFun (SMatrix (PrimState m)) m
makeS regulariser i gradients = do
  workspace <- newDenseVector (gradients ^. dim ^. _1)
  return $! SMatrix gradients workspace ((\f -> f i) <$> regulariser)

makeS' :: PrimMonad m => Maybe (Int -> ℂ) -> MakeSFun (M 'Row (PrimState m) ℂ) m
makeS' regulariser i gradients = do
  let (steps, params) = gradients ^. dim
  s <- newDenseMatrix params params
  herk MatUpper ConjTranspose (1 / int2Float steps) gradients 0 s
  herk MatLower ConjTranspose (1 / int2Float steps) gradients 0 s
  case regulariser of
    Just f  -> do
      one <- MDenseVector @'Direct params 0 <$> V.unsafeThaw (V.singleton 1)
      axpy (f i) one (MDenseVector @'Direct params (s ^. stride + 1) (s ^. buffer))
    Nothing -> return ()
  return $! s

opS :: PrimMonad m => SMatrix (PrimState m) -> Operator m ℂ
opS (SMatrix o temp λ) x out = do
  let scale = (1 / int2Float (o ^. dim ^. _1)) :+ 0
  gemv NoTranspose 1 o x 0 temp
  case λ of
    (Just λ') -> copy x out >> gemv ConjTranspose scale o temp λ' out
    Nothing   -> gemv ConjTranspose scale o temp 0 out

{-
solveS :: PrimMonad m => CGConfig ℝ -> SolveFun (MDenseMatrix 'Row (PrimState m) ℂ) m
solveS (CGConfig maxIter tol) = doSolve
  where doSolve s b x = do
          let operator = \input output -> gemv NoTranspose 1 s input 0 output
          zero <- MDenseVector @'Direct (b ^. dim) 0 <$> V.unsafeThaw (V.singleton 0)
          copy zero x
          !answer <- CG.cg maxIter tol operator b x
          trace (show answer) (return ())
-}

solveS :: CGConfig ℝ -> SolveFun (SMatrix RealWorld) IO
solveS (CGConfig maxIter tol) = doSolve
 where
  doSolve s b x = do
    t1  <- getCPUTime
    one <- MDenseVector @ 'Direct (b ^. dim) 0 <$> V.unsafeThaw (V.singleton 0)
    copy one x
    (iters, err) <- CG.cg maxIter tol (opS s) b x
    t2           <- getCPUTime
    return $! SolverStats iters err (fromIntegral (t2 - t1) * 1.0E-12)

solveS' :: SolveFun (MDenseMatrix 'Row RealWorld ℂ) IO
solveS' s b x = assert (x ^. stride == 1) $ do
  t1 <- getCPUTime
  copy b x
  cgelsd s (MDenseMatrix (b ^. dim) 1 1 (x ^. buffer)) (-1.0)
  t2 <- getCPUTime
  return $! SolverStats (-1) (-1.0) (fromIntegral (t2 - t1) * 1.0E-12)

type Stepper v m a = v (PrimState m) a -> m ()

numberSteps :: (Int, Int, Int) -> Int
numberSteps (low, high, step) = (high - low - 1) `div` step + 1

newWorkspace :: PrimMonad m => Int -> Int -> m (DenseWorkspace (PrimState m) ℂ)
newWorkspace nParams nSteps =
  DenseWorkspace
    <$> newDenseVector nParams
    <*> newDenseMatrix nSteps nParams
    <*> newDenseVector nParams

sr
  :: SRConfig ℂ
  -> Hamiltonian
  -> MRbm (PrimState IO)
  -> (IterInfo -> IO ())
  -> IO ()
sr config hamiltonian ψ process = newWorkspace nParams nSteps
  >>= \w -> loop w 0
 where
  nParams  = size ψ
  nSteps   = config ^. mc ^. runs * numberSteps (config ^. mc ^. steps)
  doAxpy   = axpyRbm
  doMakeS  = makeS (config ^. regulariser)
  doSolveS = solveS (config ^. cg)
  doSample = sampleRbm (config ^. mc) hamiltonian
  loop !w@(DenseWorkspace f grad δ) !i
    | i >= config ^. maxIter = return ()
    | otherwise = do
      mcStats <- doSample ψ f grad
      fNorm   <- nrm2 f
      cgStats <- doMakeS i grad >>= \s -> doSolveS s f δ
      δNorm   <- nrm2 δ
      unsafeFreezeRbm ψ
        >>= \ψ' ->
              let stats = IterInfo i ψ' mcStats cgStats fNorm δNorm
              in  process stats
      doAxpy (-(config ^. rate) i) δ ψ
      loop w (i + 1)


{-# LANGUAGE FlexibleContexts #-}

module Numeric.HHT.Internal.FFT (
    fft
  , ifft
  ) where

import           Data.Complex
import qualified Data.Array.CArray         as CA
import qualified Data.Array.IArray         as IA
import qualified Data.Ix                   as Ix
import qualified Data.Vector.Generic       as VG
import qualified Data.Vector.Generic.Sized as SVG
import qualified Foreign.Storable          as FS
import qualified Math.FFT                  as FFT
import qualified Math.FFT.Base             as FFT

fft :: (FFT.FFTWReal a, VG.Vector v (Complex a))
    => SVG.Vector v n (Complex a)
    -> SVG.Vector v n (Complex a)
fft = SVG.withVectorUnsafe $
        fromCA
      . FFT.dft
      . toCA

ifft
    :: (FFT.FFTWReal a, VG.Vector v (Complex a))
    => SVG.Vector v n (Complex a)
    -> SVG.Vector v n (Complex a)
ifft = SVG.withVectorUnsafe $
        fromCA
      . FFT.idft
      . toCA

fromCA
    :: (FS.Storable a, VG.Vector v (Complex a))
    => CA.CArray Int (Complex a)
    -> v (Complex a)
fromCA v = VG.generate (Ix.rangeSize (IA.bounds v)) (v IA.!)

toCA
    :: (FS.Storable a, VG.Vector v (Complex a))
    => v (Complex a)
    -> CA.CArray Int (Complex a)
toCA v = IA.listArray (0, VG.length v - 1) (VG.toList v)


{-|
Module: MachineLearning.Utils
Description: Utils
Copyright: (c) Alexander Ignatyev, 2016
License: BSD-3
Stability: experimental
Portability: POSIX

Various helpful utilities.
-}

module MachineLearning.Utils
(
  reduceByRowsV
  , reduceByColumnsV
  , reduceByRows
  , reduceByColumns
  , sumByRows
  , sumByColumns
  , listOfTuplesToList
)

where

import MachineLearning.Types (R, Vector, Matrix)
import qualified Data.Vector.Storable as V
import qualified Numeric.LinearAlgebra as LA

reduceByRowsV :: (Vector -> R) -> Matrix -> Vector
reduceByRowsV f = LA.vector . map f . LA.toRows

reduceByColumnsV :: (Vector -> R) -> Matrix -> Vector
reduceByColumnsV f = LA.vector . map f . LA.toColumns

reduceByRows :: (Vector -> R) -> Matrix -> Matrix
reduceByRows f = LA.asColumn . reduceByRowsV f

reduceByColumns :: (Vector -> R) -> Matrix -> Matrix
reduceByColumns f = LA.asRow . reduceByColumnsV f

sumByColumns :: Matrix -> Matrix
sumByColumns = reduceByColumns V.sum

sumByRows :: Matrix -> Matrix
sumByRows = reduceByRows V.sum

-- | Converts list of tuples into list.
listOfTuplesToList :: [(a, a)] -> [a]
listOfTuplesToList [] = []
listOfTuplesToList ((a, b):xs) = a : b : listOfTuplesToList xs


{-# LANGUAGE FlexibleContexts #-}
module Statistics.Test.Internal (
    rank
  , rankUnsorted
  , splitByTags
  ) where

import Data.Ord
import           Data.Vector.Generic           ((!))
import qualified Data.Vector.Generic         as G
import qualified Data.Vector.Generic.Mutable as M
import Statistics.Function

-- Private data type for unfolding
data Rank v a = Rank {
      rankCnt :: {-# UNPACK #-} !Int        -- Number of ranks to return
    , rankVal :: {-# UNPACK #-} !Double     -- Rank to return
    , rankNum :: {-# UNPACK #-} !Double     -- Current rank
    , rankVec :: v a                        -- Remaining vector
    }

-- | Calculate rank of every element of sample. In case of ties ranks
--   are averaged. Sample should be already sorted in ascending order.
--
--   Rank is index of element in the sample, numeration starts from 1.
--   In case of ties average of ranks of equal elements is assigned
--   to each
--
-- >>> rank (==) (fromList [10,20,30::Int])
-- > fromList [1.0,2.0,3.0]
--
-- >>> rank (==) (fromList [10,10,10,30::Int])
-- > fromList [2.0,2.0,2.0,4.0]
rank :: (G.Vector v a, G.Vector v Double)
     => (a -> a -> Bool)        -- ^ Equivalence relation
     -> v a                     -- ^ Vector to rank
     -> v Double
rank eq vec = G.unfoldr go (Rank 0 (-1) 1 vec)
  where
    go (Rank 0 _ r v)
      | G.null v  = Nothing
      | otherwise =
          case G.length h of
            1 -> Just (r, Rank 0 0 (r+1) rest)
            n -> go Rank { rankCnt = n
                         , rankVal = 0.5 * (r*2 + fromIntegral (n-1))
                         , rankNum = r + fromIntegral n
                         , rankVec = rest
                         }
          where
            (h,rest) = G.span (eq $ G.head v) v
    go (Rank n val r v) = Just (val, Rank (n-1) val r v)
{-# INLINE rank #-}

-- | Compute rank of every element of vector. Unlike rank it doesn't
--   require sample to be sorted.
rankUnsorted :: ( Ord a
                , G.Vector v a
                , G.Vector v Int
                , G.Vector v Double
                , G.Vector v (Int, a)
                )
             => v a
             -> v Double
rankUnsorted xs = G.create $ do
    -- Put ranks into their original positions
    -- NOTE: backpermute will do wrong thing
    vec <- M.new n
    for 0 n $ \i ->
      M.unsafeWrite vec (index ! i) (ranks ! i)
    return vec
  where
    n = G.length xs
    -- Calculate ranks for sorted array
    ranks = rank (==) sorted
    -- Sort vector and retain original indices of elements
    (index, sorted)
      = G.unzip
      $ sortBy (comparing snd)
      $ indexed xs
{-# INLINE rankUnsorted #-}

-- | Split tagged vector
splitByTags :: (G.Vector v a, G.Vector v (Bool,a)) => v (Bool,a) -> (v a, v a)
splitByTags vs = (G.map snd a, G.map snd b)
  where
    (a,b) = G.unstablePartition fst vs
{-# INLINE splitByTags #-}


module HaskovSpec where

import Haskov (fromList,imap,hmatrix,walk,walkFrom,steady,steadyState,statesI)

import System.Random
import Test.Hspec
import qualified Data.Set as Set
import qualified Numeric.LinearAlgebra.Data as Dat
import Data.List (intercalate)

import           Control.Monad (unless,when)

testTransitions =
  [ (("A", "B"), 0.3)
  , (("A", "A"), 0.7)
  , (("B", "C"), 1.0)
  , (("C", "A"), 1.0)
  ]

spec :: Spec
spec =

  describe "haskov" $ do

--    describe "steadyState" $ do
--
--      it "should always contain positive numbers" $ do
--        let
--          transitions = [ (("A", "B"), 1.0), (("B", "C"), 1.0), (("C", "A"), 1.0)]
--          haskov = fromList transitions
--          s = steadyState haskov
--        --mapM (\(s, p) -> s `shouldSatisfy` (> 0))
--        print s
--        True `shouldBe` True

    describe "A haskov walk" $ do

    -- could use quickcheck for better tests

      it "should never do invalid transitions when transitions are ordered" $ do
        let
          transitions = [(("A", "B"), 1.0), (("B", "C"), 1.0), (("C", "A"), 1.0)]
          valid = map fst transitions
          markov = fromList transitions
        gen <- getStdGen
        res <- walk 10 markov gen
        expectOnlyValidTransitions transitions res

      it "should never do invalid transitions when transtions are not ordered" $ do
        let
          transitions = [ (("C", "A"), 1.0), (("A", "B"), 1.0), (("B", "C"), 1.0)]
          valid = map fst transitions
          haskov = fromList transitions
        gen <- getStdGen
        res <- walk 3 haskov gen
        expectOnlyValidTransitions transitions res

    describe "A haskov walkFrom" $ do

      it "starts with the head initial state" $ do
        let
          markov = fromList testTransitions
        gen <- getStdGen
        res <- walkFrom "B" 10 markov gen
        head res `shouldBe` "B"

      it "just some test" $ do
        let
          markov = fromList testTransitions
          index = imap markov
          matrix = hmatrix markov
          start = "A"
        gen <- getStdGen
        res <- walkFrom "B" 10 markov gen
        --putStrLn $ "index: " ++ ( show index)
        --putStrLn $ "matrix: " ++ ( show matrix)
        --putStrLn $ "result from " ++ start ++ ": " ++ ( show res)
        return ()

expectTrue :: HasCallStack => String -> Bool -> Expectation
expectTrue msg b = unless b (expectationFailure msg)

expectOnlyValidTransitions transitions actual = do
  let
    valid = map fst transitions
    actualTransitions = zip actual (drop 1 actual)
    actualUnique = Set.fromList actualTransitions
    invalid = Set.filter (\e -> not $ elem e valid) actualUnique
    len = length invalid
    invalidMsg inv = "invalid transitions encountered: " ++ ( intercalate ", " (Set.toList (Set.map (\(a,b) -> a ++ "->" ++ b) inv)))
  expectTrue (invalidMsg invalid) $ len == 0


{-# LANGUAGE BangPatterns, FlexibleContexts #-}
-- |
-- Module    : Statistics.Transform
-- Copyright : (c) 2011 Bryan O'Sullivan
-- License   : BSD3
--
-- Maintainer  : bos@serpentine.com
-- Stability   : experimental
-- Portability : portable
--
-- Fourier-related transformations of mathematical functions.
--
-- These functions are written for simplicity and correctness, not
-- speed.  If you need a fast FFT implementation for your application,
-- you should strongly consider using a library of FFTW bindings
-- instead.

module Transform ( kPointFFT ) where

import Control.Monad         (when)
import Control.Monad.ST      (ST)
import Data.Bits             (shiftL, shiftR, (.&.), (.|.))
import Data.Complex          (Complex(..), conjugate, realPart, magnitude)
import qualified Data.Vector.Generic         as G
import qualified Data.Vector.Generic.Mutable as M
import qualified Data.Vector.Unboxed         as U

type CD = Complex Float

-- | /O(log n)/ Compute the logarithm in base 2 of the given value.
log2 :: Int -> Int
log2 v0
    | v0 <= 0   = error "Statistics.Math.log2: invalid input"
    | otherwise = go 5 0 v0
  where
    go !i !r !v | i == -1        = r
                | v .&. b i /= 0 = let si = U.unsafeIndex sv i
                                   in go (i-1) (r .|. si) (v `shiftR` si)
                | otherwise      = go (i-1) r v
    b = U.unsafeIndex bv
    !bv = U.fromList [0x2, 0xc, 0xf0, 0xff00, 0xffff0000, 0xffffffff00000000]
    !sv = U.fromList [1,2,4,8,16,32]

hanning :: Float -> Float -> Float
hanning m n = 0.5 - 0.5 * cos(2 * pi * n / m)

kPointFFT :: U.Vector Float -> U.Vector Float
kPointFFT !vec = mag where
    winlen = realToFrac $ G.length vec - 1
    window = G.imap (\i x -> ((hanning winlen (realToFrac i)) * x) :+ 0.0) vec
    fftres = halfFFT window
    mag    = G.map (\x -> let val = 20 * (logBase 10 $ magnitude x) in if(val < 0) then 0 else val) fftres

halfFFT :: U.Vector CD -> U.Vector CD
halfFFT !vec = U.take (U.length vec `div` 2) (fft vec)

-- | Radix-2 decimation-in-time fast Fourier transform.
fft :: U.Vector CD -> U.Vector CD
fft !v = G.create $ do
          mv <- G.thaw v
          mfft mv
          return mv

mfft :: (M.MVector v CD) => v s CD -> ST s ()
mfft vec
    | 1 `shiftL` m /= len = error "Statistics.Transform.fft: bad vector size"
    | otherwise           = bitReverse 0 0
 where
  bitReverse i j | i == len-1 = stage 0 1
                 | otherwise  = do
    when (i < j) $ M.swap vec i j
    let inner k l | k <= l    = inner (k `shiftR` 1) (l-k)
                  | otherwise = bitReverse (i+1) (l+k)
    inner (len `shiftR` 1) j
  stage l !l1 | l == m    = return ()
              | otherwise = do
    let !l2 = l1 `shiftL` 1
        !e  = -6.283185307179586/fromIntegral l2
        flight j !a | j == l1   = stage (l+1) l2
                    | otherwise = do
          let butterfly i | i >= len  = flight (j+1) (a+e)
                          | otherwise = do
                let i1 = i + l1
                xi1 :+ yi1 <- M.read vec i1
                let !c = cos a
                    !s = sin a
                    d  = (c*xi1 - s*yi1) :+ (s*xi1 + c*yi1)
                ci <- M.read vec i
                M.write vec i1 (ci - d)
                M.write vec i (ci + d)
                butterfly (i+l2)
          butterfly j
    flight 0 0
  len = M.length vec
  m   = log2 len

fi :: Int -> CD
fi = fromIntegral

halve :: Int -> Int
halve = (`shiftR` 1)


{-# language OverloadedStrings #-}
module Report (generateReports) where

import Graphics.Vega.VegaLite hiding (Sum)
import Prelude hiding (filter, lookup, repeat)
import Data.Monoid (Sum(..))
import Numeric.LinearAlgebra ((#>), Vector, toList)
import Data.Maybe

import GA
import Fitness

-- | Plot the evolution of average fitness along the generations
plotEvo :: [Double] -> IO ()
plotEvo avgs =
  let gens = [1.0 .. fromIntegral (length avgs)]
      plotData = dataFromColumns []
               . dataColumn "generation" (Numbers gens)
               . dataColumn "avg" (Numbers avgs)
               $ []
      enc      = encoding
               . position X [ PName "generation", PmType Quantitative]
               . position Y [ PName "avg", PmType Quantitative]
               $ []
  in  toHtmlFile "evolution.html" $ toVegaLite [ plotData, mark Line [], enc, height 800, width 600 ]

plotPoly :: [[Double]] -> Vector Double -> Solution (Poly n) -> IO ()
plotPoly xss ys sol =
  let x0 = map head xss
      zss = decode xss (map _getPoly $ _chromo sol)
      ysHat = toList $ zss #> (fromJust . _coeffs) sol
      ys'   = toList ys
      clusters = replicate (length x0) "poly" ++ replicate (length x0) "real"
      plotData = dataFromColumns []
               . dataColumn "x0" (Numbers $ x0 ++ x0)
               . dataColumn "y" (Numbers $ ysHat ++ ys')
               . dataColumn "Cluster" (Strings clusters)
               $ []
      enc      = encoding
               . position X [ PName "x0", PmType Quantitative]
               . position Y [ PName "y", PmType Quantitative]
               . color [ MName "Cluster", MmType Nominal ]
               $ []
  in  toHtmlFile "poly.html" $ toVegaLite [ plotData, mark Line [], enc, height 800, width 600 ]

generateReports :: [[Double]] -> Vector Double -> [Double] -> Solution (Poly n) -> IO ()
generateReports xss ys avgs best = do plotEvo avgs
                                      plotPoly xss ys best


{-# LANGUAGE BangPatterns        #-}
{-# LANGUAGE CPP                 #-}
{-# LANGUAGE TemplateHaskell     #-}
{-# LANGUAGE DataKinds           #-}
{-# LANGUAGE KindSignatures      #-}
{-# LANGUAGE GADTs               #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# OPTIONS_GHC -fno-warn-missing-signatures #-}

module Test.Grenade.Layers.PadCrop where

import           Grenade

import           Hedgehog

import           Numeric.LinearAlgebra.Static ( norm_Inf )

import           Test.Hedgehog.Hmatrix

prop_pad_crop :: Property
prop_pad_crop =
  let net :: Network '[Pad 2 3 4 6, Crop 2 3 4 6] '[ 'D3 7 9 5, 'D3 16 15 5, 'D3 7 9 5 ]
      net = Pad :~> Crop :~> NNil
  in  property $
    forAll genOfShape >>= \(d :: S ('D3 7 9 5)) ->
      let (tapes, res)  = runForwards  net d
          (_    , grad) = runBackwards net tapes d
      in  do assert $ d ~~~ res
             assert $ grad ~~~ d

prop_pad_crop_2d :: Property
prop_pad_crop_2d =
  let net :: Network '[Pad 2 3 4 6, Crop 2 3 4 6] '[ 'D2 7 9, 'D2 16 15, 'D2 7 9 ]
      net = Pad :~> Crop :~> NNil
  in  property $
    forAll genOfShape >>= \(d :: S ('D2 7 9)) ->
      let (tapes, res)  = runForwards  net d
          (_    , grad) = runBackwards net tapes d
      in  do assert $ d ~~~ res
             assert $ grad ~~~ d

(~~~) :: S x -> S x -> Bool
(S1D x) ~~~ (S1D y) = norm_Inf (x - y) < 0.00001
(S2D x) ~~~ (S2D y) = norm_Inf (x - y) < 0.00001
(S3D x) ~~~ (S3D y) = norm_Inf (x - y) < 0.00001

tests :: IO Bool
tests = checkParallel $$(discover)


{-# LANGUAGE DataKinds #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE PolyKinds #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE UndecidableInstances #-}
{-# LANGUAGE NoMonomorphismRestriction #-}
{-# OPTIONS_GHC -fno-warn-type-defaults #-}
{-# OPTIONS_GHC -fno-warn-type-defaults -fno-warn-orphans -Wall #-}

module Main where

import Algebra.Bridge.Singular
import Algebra.Field.Prime
import Algebra.Prelude.Core
import Cases
import Control.Monad (void)
import qualified Data.Text as T
import qualified Data.Vector.Unboxed as V
import Statistics.Resampling
import Statistics.Resampling.Bootstrap
import Statistics.Types
import qualified System.Random.MWC as Rand
import Prelude (read)

benchIdeal :: IsSingularPolynomial poly => Text -> Ideal poly -> IO Double
benchIdeal fun i =
  fmap ((/ 1000) . read . T.unpack) $
    singular $
      toProg fun i

toProg :: (SingularOrder (Arity poly) (MOrder poly), SingularCoeff (Coefficient poly), IsOrderedPolynomial poly) => Text -> Ideal poly -> Text
toProg fun i =
  let expr = funE fun [idealE' i]
   in prettySingular $ do
        directC "system(\"--ticks-per-sec\", 1000000)"
        void $ ringC "R" expr
        declOnlyC IdealT "G"
        directC "timer=1"
        directC "int t=rtimer"
        letC "G" expr
        directC "print(rtimer-t)"
        directC "exit"

analyse :: (IsSingularPolynomial poly) => String -> Text -> Ideal poly -> IO ()
analyse lab fun ideal = do
  gen <- Rand.create
  i2Gr <- V.replicateM 50 $ benchIdeal fun ideal
  res <- resample gen [Mean, StdDev] 1000 i2Gr
  let [Estimate mn (ConfInt lbmn ubmn _), Estimate dv _] =
        bootstrapBCA cl95 i2Gr res
  putStrLn lab
  mapM_
    (putStrLn . ('\t' :))
    [ "Mean:\t" ++ show mn ++ "(ms)"
    , "MeanLB:\t" ++ show lbmn
    , "MeanUB:\t" ++ show ubmn
    , "StdDev:\t" ++ show dv
    ]

runTestCases ::
  (IsMonomialOrder n o, KnownNat n) =>
  String ->
  Ideal (OrderedPolynomial Rational o n) ->
  IO ()
runTestCases lab i = do
  analyse (lab ++ " (Q, Lex, Sing(gr))") "groebner" $ fmap (changeOrder Lex) i
  analyse (lab ++ " (Q, Lex, Sing(sba))") "sba" $ fmap (changeOrder Lex) i
  analyse (lab ++ " (Q, Grevlex, Sing(gr))") "groebner" $ fmap (changeOrder Grevlex) i
  analyse (lab ++ " (Q, Grevlex, Sing(sba))") "sba" $ fmap (changeOrder Grevlex) i
  analyse (lab ++ " (F_65521, Lex, Sing(gr))") "groebner" $ fmap (mapCoeff ratToF . changeOrder Lex) i
  analyse (lab ++ " (F_65521, Lex, Sing(sba))") "sba" $ fmap (mapCoeff ratToF . changeOrder Lex) i
  analyse (lab ++ " (F_65521, Grevlex, Sing(gr))") "groebner" $ fmap (mapCoeff ratToF . changeOrder Grevlex) i
  analyse (lab ++ " (F_65521, Grevlex, Sing(sba))") "sba" $ fmap (mapCoeff ratToF . changeOrder Grevlex) i

main :: IO ()
main = do
  runTestCases "Cyclic-4" $ cyclic (sNat :: SNat 4)
  runTestCases "Cyclic-5" $ cyclic (sNat :: SNat 5)
  runTestCases "Cyclic-6" $ cyclic (sNat :: SNat 6)
  runTestCases "Katsura-5" $ katsura (sNat :: SNat 5)
  runTestCases "Katsura-6" $ katsura (sNat :: SNat 6)
  runTestCases "Katsura-7" $ katsura (sNat :: SNat 7)

ratToF :: Rational -> F 65521
ratToF = modRat'


module Main where
-- package の import
-- CSVの読み込み
import qualified    Text.CSV                                as CSV
import qualified    Data.Text                               as T
import              Data.Attoparsec.Text                    as DAT
-- 統計のためのパッケージ
import qualified    Statistics.Correlation                  as SC
-- その他
import qualified    Control.Monad                           as CM
import qualified    Data.List                               as L
import qualified    Data.Vector.Unboxed                     as VU
import              Data.IORef

main = do
    -- csvファイルの読み込み
    xs <- read_csv "Data/multiTimeline.csv"
    -- データの表示
    print_data (L.head xs) (L.tail xs)
    -- Headerの抽出
    let headers = L.map T.unpack $ L.head xs
    -- 値のDoubleへの変換
    let body    = L.transpose
                $ L.map  (L.map parseDouble)
                $ L.tail xs
    -- 分析するデータの選択
    -- xとyのHeader名を書き換えよう
    let x = column body "AI" headers
    let y = column body "Python" headers
    -- 相関係数を求める
    let res = SC.pearson
            $ VU.zip (VU.fromList x) (VU.fromList y)
    -- 相関係数の表示
    putStrLn $ "相関係数:" ++ show res

-- ** 作業用関数
-- 別のモジュールにまとめておくのが良いが複雑になるのでここにまとめておく
-- | Doubleのパーサ
{-# INLINE parseDouble #-}
parseDouble :: T.Text -> Double
parseDouble tx = case DAT.parseOnly DAT.double tx of
        Right r   -> r
        Left  l   -> error $ "Error on parseDouble : " ++ show l

-- | データの読み込み
read_csv :: String -> IO [[T.Text]]
read_csv file
    = CSV.parseCSVFromFile file >>= \res
    -> case res of
        Right xs -> return $ L.map (L.map T.pack) xs
        Left  xs -> error $ "error on parseCSVFromFile: " ++ show xs

-- | data set っぽく表示する
print_data :: (Show a)  => [T.Text] {- header -}
                        -> [[a]]    {- body -}
                        -> IO ()
print_data hs bs = do
    let max_length  =  let hs_max = L.maximum (L.map T.length hs)
                    in let bs_max = L.maximum
                                  $ (flip L.map) bs
                                  $ \x  -> L.maximum
                                  $ L.map (L.length . show) x
                    in case compare hs_max bs_max  of
                        LT -> bs_max
                        _  -> hs_max

    let bs_digit_num =  L.length (show (L.length bs))
    putStr $ " " ++ L.replicate bs_digit_num ' '
    CM.forM_ hs $ \x -> let tx = T.unpack x
                     in case compare (L.length tx) max_length of
                        LT -> putStr $ tx ++ (L.replicate  (max_length -  (L.length tx)) ' ')  ++ " "
                        _  -> putStr $ tx ++ " "
    putStrLn ""

    currentNum <- newIORef 0 :: IO (IORef Int)

    let  total_col = L.length bs
    case total_col  > 10 of
        False -> CM.forM_ bs $  \xs -> readIORef currentNum >>= \num
                                    -> putStr ((show num)   ++ " "
                                                            ++ L.replicate
                                                                (bs_digit_num - (L.length (show num))) ' ')
                                    >> modifyIORef currentNum (\x -> x + 1)
                                    >> CM.forM_ xs  ( \x -> let tx = show x
                                                         in putStr  $ tx
                                                                    ++ (L.replicate  (max_length -  (L.length tx)) ' ')
                                                                    ++ " ")
                                    >> putStrLn ""
        True ->  CM.forM_ bs $  \xs -> readIORef currentNum >>= \num
                                    -> case (num <= 4 || (total_col - num) <= 4) of
                                        True  -> putStr ((show num)   ++ " " ++ L.replicate (bs_digit_num - (L.length (show num))) ' ')
                                              >> modifyIORef currentNum (\x -> x + 1)
                                              >> CM.forM_ xs  ( \x -> let tx = show x
                                                         in putStr  $ tx
                                                                    ++ (L.replicate  (max_length -  (L.length tx)) ' ')
                                                                    ++ " ")
                                              >> putStrLn ""
                                        False -> modifyIORef currentNum (\x -> x + 1)

-- | headerによる情報の抜き出し
column :: [[Double]] -> String -> [String] -> [Double]
column xs name hs = case (L.elemIndex name hs) of
                    Just i  -> xs !! i
                    Nothing -> error $ "error at column: " ++ name ++ "doesn't exest."