{-# LANGUAGE GeneralizedNewtypeDeriving #-}

module Valuation
    ( module Valuation
    , throwError
    ) where

import Control.Applicative
import Control.Monad.Except
import Control.Monad.Identity
import Control.Monad.Reader
import Statistics.Distribution
import Statistics.Distribution.Normal

type Year = Int

newtype Scenario a = Scenario
    { scenario :: a
    }

newtype Value s a = Value
    { runValue :: ReaderT (Scenario s) (ExceptT String Identity) a
    } deriving (Monad, Applicative, Alternative, Functor, MonadReader (Scenario s), MonadError String)

instance Num a => Num (Value s a) where
    negate = fmap negate
    (+) = liftM2 (+)
    (*) = liftM2 (*)
    fromInteger = pure . fromInteger
    abs = fmap abs
    signum = fmap signum

instance Fractional a => Fractional (Value s a) where
    (/) = liftM2 (/)
    fromRational = pure . fromRational
    recip x = 1 / x

type TimeValue s a = Year -> Value s a

runValuation :: s -> Value s a -> Either String a
runValuation s v = runIdentity (runExceptT (runReaderT (runValue v) (Scenario s)))

runV :: (Show a, RealFrac a) => TimeValue () a -> [Year] -> IO ()
runV v (y:ys) = do
    putStr $ show y
    putStr "  "
    case runValuation () $ v y of
        Left s -> print s
        Right result -> print result
    runV v ys
runV _ [] = return ()

discount :: TimeValue s Double -> Year -> TimeValue s Double -> TimeValue s Double
discount r t0 v t
    | t == t0 = v t
    | otherwise = discount r (t0 + 1) v t / (1 + r (t0 + 1))

npv :: TimeValue s Double -> TimeValue s Double -> TimeValue s Double
npv r v t = v t + (npv r v (t + 1) / (1 + r (t + 1)) <|> return 0)

linear :: Fractional a => TimeValue s a -> Year -> Year -> a -> TimeValue s a
linear v t0 t1 target t
    | t0 < t && t <= t1 = do
        base <- v t0
        return $ base + (target - base) * fromIntegral (t - t0) / fromIntegral (t1 - t0)
    | t1 < t = return target

pvAnnuity :: Double -> Double -> Int -> Double
pvAnnuity r pmt n = pmt * (1 - (1 + r) ^^ (-n)) / r

blackScholes :: Double -> Double -> Double -> Double -> Double -> Double
blackScholes sigma r t s k =
    let d1 = 1 / (sigma * sqrt t) * (log (s / k) + (r + sigma ** t / 2) * t)
        d2 = d1 - sigma * sqrt t
    in s * cumulative standard d1 - cumulative standard d2 * k * exp (-r * t)


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