module HashedExpression.Interp
  ( Evaluable (..),
    Approximable (..),
  )
where

import Data.Array
import Data.Complex
import qualified Data.IntMap.Strict as IM
import Data.List (intercalate)
import Data.Map (Map, fromList)
import qualified Data.Map as Map
import Debug.Trace (traceId, traceShowId)
import GHC.TypeLits (KnownNat)
import HashedExpression.Internal.Expression
  ( C,
    ET (..),
    Expression (..),
    ExpressionMap,
    Node (..),
    NodeID,
    NodeID,
    R,
    Scalar,
  )
import HashedExpression.Internal.Node
import HashedExpression.Internal.Utils
import HashedExpression.Prettify (prettify, showExp)
import HashedExpression.Value
import Text.Printf

-- | This operation emulates the mathematical operation
-- | Turn expression to the right type
expZeroR :: ExpressionMap -> NodeID -> Expression Scalar R
expZeroR = flip Expression

expZeroC :: ExpressionMap -> NodeID -> Expression Scalar C
expZeroC = flip Expression

-- | Choose branch base on 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)

-- | Approximable class
class
  Show a =>
  Approximable a where
  (~=) :: a -> a -> Bool
  prettifyShow :: a -> String

infix 4 ~=

-- |
relativeError :: Double -> Double -> Double
relativeError a b = abs (a - b) / max (abs a) (abs b)

instance Approximable Double where
  (~=) :: Double -> Double -> Bool
  a ~= b
    | abs (a - b) < 1.0e-5 = True
    | a == b = True
    | otherwise = relativeError a b < 0.01
  prettifyShow a
    | abs a < 1e-10 = "0"
    | otherwise = printf "%.2f" a

instance Approximable (Complex Double) where
  (~=) :: Complex Double -> Complex Double -> Bool
  a ~= b = (realPart a ~= realPart b) && (imagPart a ~= imagPart b)
  prettifyShow a =
    prettifyShow (realPart a) ++ " + " ++ prettifyShow (imagPart a) ++ "i"

instance Approximable (Array Int Double) where
  (~=) :: Array Int Double -> Array Int Double -> Bool
  a ~= b = (bounds a == bounds b) && and (zipWith (~=) (elems a) (elems b))
  prettifyShow a =
    "[" ++ (intercalate ", " . map prettifyShow . elems $ a) ++ "]"

instance Approximable (Array Int (Complex Double)) where
  (~=) :: Array Int (Complex Double) -> Array Int (Complex Double) -> Bool
  a ~= b = (bounds a == bounds b) && and (zipWith (~=) (elems a) (elems b))
  prettifyShow a =
    "[" ++ (intercalate ", " . map prettifyShow . elems $ a) ++ "]"

instance Approximable (Array (Int, Int) Double) where
  (~=) :: Array (Int, Int) Double -> Array (Int, Int) Double -> Bool
  a ~= b = (bounds a == bounds b) && and (zipWith (~=) (elems a) (elems b))
  prettifyShow a =
    "[" ++ (intercalate ", " . map prettifyShow . elems $ a) ++ "]"

instance Approximable (Array (Int, Int) (Complex Double)) where
  (~=) ::
    Array (Int, Int) (Complex Double) ->
    Array (Int, Int) (Complex Double) ->
    Bool
  a ~= b = (bounds a == bounds b) && and (zipWith (~=) (elems a) (elems b))
  prettifyShow a =
    "[" ++ (intercalate ", " . map prettifyShow . elems $ a) ++ "]"

instance Approximable (Array (Int, Int, Int) Double) where
  (~=) :: Array (Int, Int, Int) Double -> Array (Int, Int, Int) Double -> Bool
  a ~= b = (bounds a == bounds b) && and (zipWith (~=) (elems a) (elems b))
  prettifyShow a =
    "[" ++ (intercalate ", " . map prettifyShow . elems $ a) ++ "]"

instance Approximable (Array (Int, Int, Int) (Complex Double)) where
  (~=) ::
    Array (Int, Int, Int) (Complex Double) ->
    Array (Int, Int, Int) (Complex Double) ->
    Bool
  a ~= b = (bounds a == bounds b) && and (zipWith (~=) (elems a) (elems b))
  prettifyShow a =
    "[" ++ (intercalate ", " . map prettifyShow . elems $ a) ++ "]"

-- | These should be commented properly.
class Evaluable d rc output | d rc -> output where
  eval :: ValMaps -> Expression d rc -> output

-- |
instance Evaluable Scalar R Double where
  eval :: ValMaps -> Expression Scalar R -> Double
  eval valMap e@(Expression n mp)
    | [] <- retrieveShape n mp =
      case retrieveNode n mp of
        Var name ->
          case Map.lookup name valMap of
            Just (VScalar val) -> val
            _ -> error "no value associated with the variable"
        Const val -> val
        Sum R args -> sum . map (eval valMap . expZeroR mp) $ args
        Mul R args -> product . map (eval valMap . expZeroR mp) $ args
        Neg R arg -> - (eval valMap $ expZeroR mp arg)
        Scale R arg1 arg2 ->
          eval valMap (expZeroR mp arg1)
            * eval valMap (expZeroR mp arg2)
        Power x arg -> eval valMap (expZeroR mp arg) ^ x
        Div arg1 arg2 ->
          eval valMap (expZeroR mp arg1)
            / eval valMap (expZeroR mp arg2)
        Sqrt arg -> sqrt (eval valMap (expZeroR mp arg))
        Sin arg -> sin (eval valMap (expZeroR mp arg))
        Cos arg -> cos (eval valMap (expZeroR mp arg))
        Tan arg -> tan (eval valMap (expZeroR mp arg))
        Exp arg -> exp (eval valMap (expZeroR mp arg))
        Log arg -> log (eval valMap (expZeroR mp arg))
        Sinh arg -> sinh (eval valMap (expZeroR mp arg))
        Cosh arg -> cosh (eval valMap (expZeroR mp arg))
        Tanh arg -> tanh (eval valMap (expZeroR mp arg))
        Asin arg -> asin (eval valMap (expZeroR mp arg))
        Acos arg -> acos (eval valMap (expZeroR mp arg))
        Atan arg -> atan (eval valMap (expZeroR mp arg))
        Asinh arg -> asinh (eval valMap (expZeroR mp arg))
        Acosh arg -> acosh (eval valMap (expZeroR mp arg))
        Atanh arg -> atanh (eval valMap (expZeroR mp arg))
        RealPart arg -> realPart (eval valMap (expZeroC mp arg))
        ImagPart arg -> imagPart (eval valMap (expZeroC mp arg))
        InnerProd R arg1 arg2 ->
          case retrieveShape arg1 mp of
            [] ->
              eval valMap (expZeroR mp arg1)
                * eval valMap (expZeroR mp arg2)
            [size] ->
              let res1 = evaluate1DReal valMap $ (mp, arg1)
                  res2 = evaluate1DReal valMap $ (mp, arg2)
               in sum
                    [ x * y
                      | i <- [0 .. size - 1],
                        let x = res1 ! i,
                        let y = res2 ! i
                    ]
            [size1, size2] ->
              let res1 = evaluate2DReal valMap $ (mp, arg1)
                  res2 = evaluate2DReal valMap $ (mp, arg2)
               in sum
                    [ x * y
                      | i <- [0 .. size1 - 1],
                        j <- [0 .. size2 - 1],
                        let x = res1 ! (i, j),
                        let y = res2 ! (i, j)
                    ]
            [size1, size2, size3] ->
              let res1 = evaluate3DReal valMap $ (mp, arg1)
                  res2 = evaluate3DReal valMap $ (mp, arg2)
               in sum
                    [ x * y
                      | i <- [0 .. size1 - 1],
                        j <- [0 .. size2 - 1],
                        k <- [0 .. size3 - 1],
                        let x = res1 ! (i, j, k),
                        let y = res2 ! (i, j, k)
                    ]
            _ -> error "4D shape?"
        Piecewise marks conditionArg branchArgs ->
          let cdt = eval valMap $ expZeroR mp conditionArg
              branches = map (eval valMap . expZeroR mp) branchArgs
           in chooseBranch marks cdt branches
        _ ->
          error
            ("expression structure Scalar R is wrong " ++ prettify e)
    | otherwise = error "one r but shape is not [] ??"

instance Evaluable Scalar C (Complex Double) where
  eval :: ValMaps -> Expression Scalar C -> Complex Double
  eval valMap e@(Expression n mp)
    | [] <- retrieveShape n mp =
      case retrieveNode n mp of
        Sum C args -> sum . map (eval valMap . expZeroC mp) $ args
        Mul C args -> product . map (eval valMap . expZeroC mp) $ args
        Power x arg -> eval valMap (expZeroC mp arg) ^ x
        Neg C arg -> - (eval valMap $ expZeroC mp arg)
        Scale C arg1 arg2 ->
          case retrieveElementType arg1 mp of
            R ->
              fromR (eval valMap (expZeroR mp arg1))
                * eval valMap (expZeroC mp arg2)
            C ->
              eval valMap (expZeroC mp arg1)
                * eval valMap (expZeroC mp arg2)
        RealImag arg1 arg2 ->
          eval valMap (expZeroR mp arg1)
            :+ eval valMap (expZeroR mp arg2)
        InnerProd C arg1 arg2 ->
          case retrieveShape arg1 mp of
            [] ->
              eval valMap (expZeroC mp arg1)
                * conjugate (eval valMap (expZeroC mp arg2))
            [size] ->
              let res1 = evaluate1DComplex valMap $ (mp, arg1)
                  res2 = evaluate1DComplex valMap $ (mp, arg2)
               in sum
                    [ x * conjugate y
                      | i <- [0 .. size - 1],
                        let x = res1 ! i,
                        let y = res2 ! i
                    ]
            [size1, size2] ->
              let res1 = evaluate2DComplex valMap $ (mp, arg1)
                  res2 = evaluate2DComplex valMap $ (mp, arg2)
               in sum
                    [ x * conjugate y
                      | i <- [0 .. size1 - 1],
                        j <- [0 .. size2 - 1],
                        let x = res1 ! (i, j),
                        let y = res2 ! (i, j)
                    ]
            [size1, size2, size3] ->
              let res1 = evaluate3DComplex valMap $ (mp, arg1)
                  res2 = evaluate3DComplex valMap $ (mp, arg2)
               in sum
                    [ x * conjugate y
                      | i <- [0 .. size1 - 1],
                        j <- [0 .. size2 - 1],
                        k <- [0 .. size3 - 1],
                        let x = res1 ! (i, j, k),
                        let y = res2 ! (i, j, k)
                    ]
            _ -> error "4D shape?"
        Piecewise marks conditionArg branchArgs ->
          let cdt = eval valMap $ expZeroR mp conditionArg
              branches = map (eval valMap . expZeroC mp) branchArgs
           in chooseBranch marks cdt branches
        _ ->
          error
            ("expression structure Scalar C is wrong " ++ prettify e)
    | otherwise = error "One C but shape is not [] ??"

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

foldrElementwise :: Ix ix => (a -> a -> a) -> [Array ix a] -> Array ix a
foldrElementwise f [x] = x
foldrElementwise f (x : xs) = zipWithA f x (foldrElementwise f xs)

-- |
evaluate1DReal :: ValMaps -> (ExpressionMap, NodeID) -> Array Int Double
evaluate1DReal valMap (mp, n)
  | [size] <- retrieveShape n mp =
    case retrieveNode n mp of
      Var name ->
        case Map.lookup name valMap of
          Just (V1D val) -> val
          _ -> error "no value associated with the variable"
      Const val -> listArray (0, size - 1) $ replicate size val
      Sum R args ->
        foldrElementwise (+) . map (evaluate1DReal valMap . (mp,)) $
          args
      Mul R args ->
        foldrElementwise (*) . map (evaluate1DReal valMap . (mp,)) $
          args
      Power x arg -> fmap (^ x) (evaluate1DReal valMap $ (mp, arg))
      Neg R arg -> fmap negate . evaluate1DReal valMap $ (mp, arg)
      Scale R arg1 arg2 ->
        let scalar = eval valMap $ expZeroR mp arg1
         in fmap (scalar *) . evaluate1DReal valMap $ (mp, arg2)
      Div arg1 arg2 ->
        zipWithA
          (/)
          (evaluate1DReal valMap $ (mp, arg2))
          (evaluate1DReal valMap $ (mp, arg2))
      Sqrt arg -> fmap sqrt . evaluate1DReal valMap $ (mp, arg)
      Sin arg -> fmap sin . evaluate1DReal valMap $ (mp, arg)
      Cos arg -> fmap cos . evaluate1DReal valMap $ (mp, arg)
      Tan arg -> fmap tan . evaluate1DReal valMap $ (mp, arg)
      Exp arg -> fmap exp . evaluate1DReal valMap $ (mp, arg)
      Log arg -> fmap log . evaluate1DReal valMap $ (mp, arg)
      Sinh arg -> fmap sinh . evaluate1DReal valMap $ (mp, arg)
      Cosh arg -> fmap cosh . evaluate1DReal valMap $ (mp, arg)
      Tanh arg -> fmap tanh . evaluate1DReal valMap $ (mp, arg)
      Asin arg -> fmap asin . evaluate1DReal valMap $ (mp, arg)
      Acos arg -> fmap acos . evaluate1DReal valMap $ (mp, arg)
      Atan arg -> fmap atan . evaluate1DReal valMap $ (mp, arg)
      Asinh arg -> fmap asinh . evaluate1DReal valMap $ (mp, arg)
      Acosh arg -> fmap acosh . evaluate1DReal valMap $ (mp, arg)
      Atanh arg -> fmap atanh . evaluate1DReal valMap $ (mp, arg)
      RealPart arg -> fmap realPart . evaluate1DComplex valMap $ (mp, arg)
      ImagPart arg -> fmap imagPart . evaluate1DComplex valMap $ (mp, arg)
      -- Rotate rA arg ->
      Piecewise marks conditionArg branchArgs ->
        let cdt = evaluate1DReal valMap $ (mp, conditionArg)
            branches = map (evaluate1DReal valMap . (mp,)) branchArgs
         in listArray
              (0, size - 1)
              [ chosen ! i
                | i <- [0 .. size - 1],
                  let chosen = chooseBranch marks (cdt ! i) branches
              ]
      Rotate [amount] arg ->
        rotate1D size amount (evaluate1DReal valMap $ (mp, arg))
      TwiceReFT arg ->
        let innerRes = evaluate1DReal valMap $ (mp, arg)
            scaleFactor = fromIntegral size / 2
         in listArray
              (0, size - 1)
              [ scaleFactor
                  * (innerRes ! i + innerRes ! ((size - i) `mod` size))
                | i <- [0 .. size - 1]
              ]
      TwiceImFT arg ->
        let innerRes = evaluate1DReal valMap $ (mp, arg)
            scaleFactor = fromIntegral size / 2
         in listArray
              (0, size - 1)
              [ scaleFactor
                  * (innerRes ! i - innerRes ! ((size - i) `mod` size))
                | i <- [0 .. size - 1]
              ]
      ReFT arg ->
        case retrieveElementType arg mp of
          R ->
            let inner =
                  fmap (:+ 0) . evaluate1DReal valMap $ (mp, arg)
                ftResult = fourierTransform1D size inner
             in fmap realPart ftResult
          C ->
            let inner = evaluate1DComplex valMap $ (mp, arg)
                ftResult = fourierTransform1D size inner
             in fmap realPart ftResult
      ImFT arg ->
        case retrieveElementType arg mp of
          R ->
            let inner =
                  fmap (:+ 0) . evaluate1DReal valMap $ (mp, arg)
                ftResult = fourierTransform1D size inner
             in fmap imagPart ftResult
          C ->
            let inner = evaluate1DComplex valMap $ (mp, arg)
                ftResult = fourierTransform1D size inner
             in fmap imagPart ftResult
      _ -> error "expression structure One R is wrong"
  | otherwise = error "one r but shape is not [size] ??"

instance (KnownNat n) => Evaluable n R (Array Int Double) where
  eval :: ValMaps -> Expression n R -> Array Int Double
  eval valMap (Expression n mp) = evaluate1DReal valMap (mp, n)

-- |
evaluate1DComplex ::
  ValMaps -> (ExpressionMap, NodeID) -> Array Int (Complex Double)
evaluate1DComplex valMap (mp, n)
  | [size] <- retrieveShape n mp =
    case retrieveNode n mp of
      Sum C args ->
        foldrElementwise (+) . map (evaluate1DComplex valMap . (mp,)) $
          args
      Mul C args ->
        foldrElementwise (*) . map (evaluate1DComplex valMap . (mp,)) $
          args
      Power x arg -> fmap (^ x) (evaluate1DComplex valMap $ (mp, arg))
      Neg C arg -> fmap negate . evaluate1DComplex valMap $ (mp, arg)
      Scale C arg1 arg2 ->
        case retrieveElementType arg1 mp of
          R ->
            let scalar = fromR . eval valMap $ expZeroR mp arg1
             in fmap (scalar *) . evaluate1DComplex valMap $
                  (mp, arg2)
          C ->
            let scalar = eval valMap $ expZeroC mp arg1
             in fmap (scalar *) . evaluate1DComplex valMap $
                  (mp, arg2)
      RealImag arg1 arg2 ->
        zipWithA
          (:+)
          (evaluate1DReal valMap $ (mp, arg1))
          (evaluate1DReal valMap $ (mp, arg2))
      Piecewise marks conditionArg branchArgs ->
        let cdt = evaluate1DReal valMap $ (mp, conditionArg)
            branches =
              map (evaluate1DComplex valMap . (mp,)) branchArgs
         in listArray
              (0, size - 1)
              [ chosen ! i
                | i <- [0 .. size - 1],
                  let chosen = chooseBranch marks (cdt ! i) branches
              ]
      Rotate [amount] arg ->
        rotate1D size amount (evaluate1DComplex valMap $ (mp, arg))
      _ -> error "expression structure One C is wrong"
  | otherwise = error "one C but shape is not [size] ??"

--                         in chooseBranch marks cdt branches
instance (KnownNat n) => Evaluable n C (Array Int (Complex Double)) where
  eval :: ValMaps -> Expression n C -> Array Int (Complex Double)
  eval valMap (Expression n mp) = evaluate1DComplex valMap (mp, n)

-- |
evaluate2DReal :: ValMaps -> (ExpressionMap, NodeID) -> Array (Int, Int) Double
evaluate2DReal valMap (mp, n)
  | [size1, size2] <- retrieveShape n mp =
    case retrieveNode n mp of
      Var name ->
        case Map.lookup name valMap of
          Just (V2D val) -> val
          _ -> error $ "no value associated with the variable" ++ name
      Const val ->
        listArray ((0, 0), (size1 - 1, size2 - 1)) $
          replicate (size1 * size2) val
      Sum R args ->
        foldrElementwise (+) . map (evaluate2DReal valMap . (mp,)) $
          args
      Mul R args ->
        foldrElementwise (*) . map (evaluate2DReal valMap . (mp,)) $
          args
      Power x arg -> fmap (^ x) (evaluate2DReal valMap $ (mp, arg))
      Neg R arg -> fmap negate . evaluate2DReal valMap $ (mp, arg)
      Scale R arg1 arg2 ->
        let scalar = eval valMap $ expZeroR mp arg1
         in fmap (scalar *) . evaluate2DReal valMap $ (mp, arg2)
      Div arg1 arg2 ->
        zipWithA
          (/)
          (evaluate2DReal valMap $ (mp, arg2))
          (evaluate2DReal valMap $ (mp, arg2))
      Sqrt arg -> fmap sqrt . evaluate2DReal valMap $ (mp, arg)
      Sin arg -> fmap sin . evaluate2DReal valMap $ (mp, arg)
      Cos arg -> fmap cos . evaluate2DReal valMap $ (mp, arg)
      Tan arg -> fmap tan . evaluate2DReal valMap $ (mp, arg)
      Exp arg -> fmap exp . evaluate2DReal valMap $ (mp, arg)
      Log arg -> fmap log . evaluate2DReal valMap $ (mp, arg)
      Sinh arg -> fmap sinh . evaluate2DReal valMap $ (mp, arg)
      Cosh arg -> fmap cosh . evaluate2DReal valMap $ (mp, arg)
      Tanh arg -> fmap tanh . evaluate2DReal valMap $ (mp, arg)
      Asin arg -> fmap asin . evaluate2DReal valMap $ (mp, arg)
      Acos arg -> fmap acos . evaluate2DReal valMap $ (mp, arg)
      Atan arg -> fmap atan . evaluate2DReal valMap $ (mp, arg)
      Asinh arg -> fmap asinh . evaluate2DReal valMap $ (mp, arg)
      Acosh arg -> fmap acosh . evaluate2DReal valMap $ (mp, arg)
      Atanh arg -> fmap atanh . evaluate2DReal valMap $ (mp, arg)
      RealPart arg -> fmap realPart . evaluate2DComplex valMap $ (mp, arg)
      ImagPart arg -> fmap imagPart . evaluate2DComplex valMap $ (mp, arg)
      Piecewise marks conditionArg branchArgs ->
        let cdt = evaluate2DReal valMap $ (mp, conditionArg)
            branches = map (evaluate2DReal valMap . (mp,)) branchArgs
         in listArray
              ((0, 0), (size1 - 1, size2 - 1))
              [ chosen ! (i, j)
                | i <- [0 .. size1 - 1],
                  j <- [0 .. size2 - 1],
                  let chosen =
                        chooseBranch marks (cdt ! (i, j)) branches
              ]
      Rotate [amount1, amount2] arg ->
        rotate2D
          (size1, size2)
          (amount1, amount2)
          (evaluate2DReal valMap $ (mp, arg))
      TwiceReFT arg ->
        let innerRes = evaluate2DReal valMap $ (mp, arg)
            scaleFactor = fromIntegral size1 * fromIntegral size2 / 2
         in listArray
              ((0, 0), (size1 - 1, size2 - 1))
              [ scaleFactor
                  * ( innerRes ! (i, j)
                        + innerRes
                        ! ((size1 - i) `mod` size1, (size2 - j) `mod` size2)
                    )
                | i <- [0 .. size1 - 1],
                  j <- [0 .. size2 - 1]
              ]
      TwiceImFT arg ->
        let innerRes = evaluate2DReal valMap $ (mp, arg)
            scaleFactor = fromIntegral size1 * fromIntegral size2 / 2
         in listArray
              ((0, 0), (size1 - 1, size2 - 1))
              [ scaleFactor
                  * ( innerRes ! (i, j)
                        - innerRes
                        ! ((size1 - i) `mod` size1, (size2 - j) `mod` size2)
                    )
                | i <- [0 .. size1 - 1],
                  j <- [0 .. size2 - 1]
              ]
      ReFT arg ->
        case retrieveElementType arg mp of
          R ->
            let inner =
                  fmap (:+ 0) . evaluate2DReal valMap $ (mp, arg)
                ftResult = fourierTransform2D (size1, size2) inner
             in fmap realPart ftResult
          C ->
            let inner = evaluate2DComplex valMap $ (mp, arg)
                ftResult = fourierTransform2D (size1, size2) inner
             in fmap realPart ftResult
      ImFT arg ->
        case retrieveElementType arg mp of
          R ->
            let inner =
                  fmap (:+ 0) . evaluate2DReal valMap $ (mp, arg)
                ftResult = fourierTransform2D (size1, size2) inner
             in fmap imagPart ftResult
          C ->
            let inner = evaluate2DComplex valMap $ (mp, arg)
                ftResult = fourierTransform2D (size1, size2) inner
             in fmap imagPart ftResult
      _ -> error "expression structure Two R is wrong"
  | otherwise = error "Two r but shape is not [size1, size2] ??"

instance
  (KnownNat m, KnownNat n) =>
  Evaluable '(m, n) R (Array (Int, Int) Double)
  where
  eval :: ValMaps -> Expression '(m, n) R -> Array (Int, Int) Double
  eval valMap (Expression n mp) = evaluate2DReal valMap (mp, n)

-- |
evaluate2DComplex ::
  ValMaps -> (ExpressionMap, NodeID) -> Array (Int, Int) (Complex Double)
evaluate2DComplex valMap (mp, n)
  | [size1, size2] <- retrieveShape n mp =
    case retrieveNode n mp of
      Sum C args ->
        foldrElementwise (+) . map (evaluate2DComplex valMap . (mp,)) $
          args
      Mul C args ->
        foldrElementwise (*) . map (evaluate2DComplex valMap . (mp,)) $
          args
      Power x arg -> fmap (^ x) (evaluate2DComplex valMap $ (mp, arg))
      Neg C arg -> fmap negate . evaluate2DComplex valMap $ (mp, arg)
      Scale C arg1 arg2 ->
        case retrieveElementType arg1 mp of
          R ->
            let scalar = fromR . eval valMap $ expZeroR mp arg1
             in fmap (scalar *) . evaluate2DComplex valMap $
                  (mp, arg2)
          C ->
            let scalar = eval valMap $ expZeroC mp arg1
             in fmap (scalar *) . evaluate2DComplex valMap $
                  (mp, arg2)
      RealImag arg1 arg2 ->
        zipWithA
          (:+)
          (evaluate2DReal valMap $ (mp, arg1))
          (evaluate2DReal valMap $ (mp, arg2))
      Piecewise marks conditionArg branchArgs ->
        let cdt = evaluate2DReal valMap $ (mp, conditionArg)
            branches =
              map (evaluate2DComplex valMap . (mp,)) branchArgs
         in listArray
              ((0, 0), (size1 - 1, size2 - 1))
              [ chosen ! (i, j)
                | i <- [0 .. size1 - 1],
                  j <- [0 .. size2 - 1],
                  let chosen =
                        chooseBranch marks (cdt ! (i, j)) branches
              ]
      Rotate [amount1, amount2] arg ->
        rotate2D
          (size1, size2)
          (amount1, amount2)
          (evaluate2DComplex valMap $ (mp, arg))
      _ -> error "expression structure Two C is wrong"
  | otherwise = error "Two C but shape is not [size1, size2] ??"

instance
  (KnownNat m, KnownNat n) =>
  Evaluable '(m, n) C (Array (Int, Int) (Complex Double))
  where
  eval ::
    ValMaps -> Expression '(m, n) C -> Array (Int, Int) (Complex Double)
  eval valMap (Expression n mp) = evaluate2DComplex valMap (mp, n)

evaluate3DReal ::
  ValMaps -> (ExpressionMap, NodeID) -> Array (Int, Int, Int) Double
evaluate3DReal valMap (mp, n)
  | [size1, size2, size3] <- retrieveShape n mp =
    case retrieveNode n mp of
      Var name ->
        case Map.lookup name valMap of
          Just (V3D val) -> val
          _ -> error "no value associated with the variable"
      Const val ->
        listArray ((0, 0, 0), (size1 - 1, size2 - 1, size3 - 1)) $
          replicate (size1 * size2 * size3) val
      Sum R args ->
        foldrElementwise (+) . map (evaluate3DReal valMap . (mp,)) $
          args
      Mul R args ->
        foldrElementwise (*) . map (evaluate3DReal valMap . (mp,)) $
          args
      Power x arg -> fmap (^ x) (evaluate3DReal valMap $ (mp, arg))
      Neg R arg -> fmap negate . evaluate3DReal valMap $ (mp, arg)
      Scale R arg1 arg2 ->
        let scalar = eval valMap $ expZeroR mp arg1
         in fmap (scalar *) . evaluate3DReal valMap $ (mp, arg2)
      Div arg1 arg2 ->
        zipWithA
          (/)
          (evaluate3DReal valMap $ (mp, arg2))
          (evaluate3DReal valMap $ (mp, arg2))
      Sqrt arg -> fmap sqrt . evaluate3DReal valMap $ (mp, arg)
      Sin arg -> fmap sin . evaluate3DReal valMap $ (mp, arg)
      Cos arg -> fmap cos . evaluate3DReal valMap $ (mp, arg)
      Tan arg -> fmap tan . evaluate3DReal valMap $ (mp, arg)
      Exp arg -> fmap exp . evaluate3DReal valMap $ (mp, arg)
      Log arg -> fmap log . evaluate3DReal valMap $ (mp, arg)
      Sinh arg -> fmap sinh . evaluate3DReal valMap $ (mp, arg)
      Cosh arg -> fmap cosh . evaluate3DReal valMap $ (mp, arg)
      Tanh arg -> fmap tanh . evaluate3DReal valMap $ (mp, arg)
      Asin arg -> fmap asin . evaluate3DReal valMap $ (mp, arg)
      Acos arg -> fmap acos . evaluate3DReal valMap $ (mp, arg)
      Atan arg -> fmap atan . evaluate3DReal valMap $ (mp, arg)
      Asinh arg -> fmap asinh . evaluate3DReal valMap $ (mp, arg)
      Acosh arg -> fmap acosh . evaluate3DReal valMap $ (mp, arg)
      Atanh arg -> fmap atanh . evaluate3DReal valMap $ (mp, arg)
      RealPart arg -> fmap realPart . evaluate3DComplex valMap $ (mp, arg)
      ImagPart arg -> fmap imagPart . evaluate3DComplex valMap $ (mp, arg)
      Piecewise marks conditionArg branchArgs ->
        let cdt = evaluate3DReal valMap $ (mp, conditionArg)
            branches = map (evaluate3DReal valMap . (mp,)) branchArgs
         in listArray
              ((0, 0, 0), (size1 - 1, size2 - 1, size3 - 1))
              [ chosen ! (i, j, k)
                | i <- [0 .. size1 - 1],
                  j <- [0 .. size2 - 1],
                  k <- [0 .. size3 - 1],
                  let chosen =
                        chooseBranch marks (cdt ! (i, j, k)) branches
              ]
      Rotate [amount1, amount2, amount3] arg ->
        rotate3D
          (size1, size2, size3)
          (amount1, amount2, amount3)
          (evaluate3DReal valMap $ (mp, arg))
      TwiceReFT arg ->
        let innerRes = evaluate3DReal valMap $ (mp, arg)
            scaleFactor =
              fromIntegral size1 * fromIntegral size2
                * fromIntegral size3
                / 2
         in listArray
              ((0, 0, 0), (size1 - 1, size2 - 1, size3 - 1))
              [ scaleFactor
                  * ( innerRes ! (i, j, k)
                        + innerRes
                        ! ( (size1 - i) `mod` size1,
                            (size2 - j) `mod` size2,
                            (size3 - k) `mod` size3
                          )
                    )
                | i <- [0 .. size1 - 1],
                  j <- [0 .. size2 - 1],
                  k <- [0 .. size3 - 1]
              ]
      TwiceImFT arg ->
        let innerRes = evaluate3DReal valMap $ (mp, arg)
            scaleFactor =
              fromIntegral size1 * fromIntegral size2
                * fromIntegral size3
                / 2
         in listArray
              ((0, 0, 0), (size1 - 1, size2 - 1, size3 - 1))
              [ scaleFactor
                  * ( innerRes ! (i, j, k)
                        - innerRes
                        ! ( (size1 - i) `mod` size1,
                            (size2 - j) `mod` size2,
                            (size3 - k) `mod` size3
                          )
                    )
                | i <- [0 .. size1 - 1],
                  j <- [0 .. size2 - 1],
                  k <- [0 .. size3 - 1]
              ]
      ReFT arg ->
        case retrieveElementType arg mp of
          R ->
            let inner =
                  fmap (:+ 0) . evaluate3DReal valMap $ (mp, arg)
                ftResult =
                  fourierTransform3D (size1, size2, size3) inner
             in fmap realPart ftResult
          C ->
            let inner = evaluate3DComplex valMap $ (mp, arg)
                ftResult =
                  fourierTransform3D (size1, size2, size3) inner
             in fmap realPart ftResult
      ImFT arg ->
        case retrieveElementType arg mp of
          R ->
            let inner =
                  fmap (:+ 0) . evaluate3DReal valMap $ (mp, arg)
                ftResult =
                  fourierTransform3D (size1, size2, size3) inner
             in fmap imagPart ftResult
          C ->
            let inner = evaluate3DComplex valMap $ (mp, arg)
                ftResult =
                  fourierTransform3D (size1, size2, size3) inner
             in fmap imagPart ftResult
      _ -> error "expression structure Three R is wrong"
  | otherwise = error "Three r but shape is not [size1, size2, size3] ??"

instance
  (KnownNat m, KnownNat n, KnownNat p) =>
  Evaluable '(m, n, p) R (Array (Int, Int, Int) Double)
  where
  eval :: ValMaps -> Expression '(m, n, p) R -> Array (Int, Int, Int) Double
  eval valMap (Expression n mp) = evaluate3DReal valMap (mp, n)

evaluate3DComplex ::
  ValMaps -> (ExpressionMap, NodeID) -> Array (Int, Int, Int) (Complex Double)
evaluate3DComplex valMap (mp, n)
  | [size1, size2, size3] <- retrieveShape n mp =
    case retrieveNode n mp of
      Sum C args ->
        foldrElementwise (+) . map (evaluate3DComplex valMap . (mp,)) $
          args
      Mul C args ->
        foldrElementwise (*) . map (evaluate3DComplex valMap . (mp,)) $
          args
      Power x arg -> fmap (^ x) (evaluate3DComplex valMap $ (mp, arg))
      Neg C arg -> fmap negate . evaluate3DComplex valMap $ (mp, arg)
      Scale C arg1 arg2 ->
        case retrieveElementType arg1 mp of
          R ->
            let scalar = fromR . eval valMap $ expZeroR mp arg1
             in fmap (scalar *) . evaluate3DComplex valMap $
                  (mp, arg2)
          C ->
            let scalar = eval valMap $ expZeroC mp arg1
             in fmap (scalar *) . evaluate3DComplex valMap $
                  (mp, arg2)
      RealImag arg1 arg2 ->
        zipWithA
          (:+)
          (evaluate3DReal valMap $ (mp, arg1))
          (evaluate3DReal valMap $ (mp, arg2))
      Piecewise marks conditionArg branchArgs ->
        let cdt = evaluate3DReal valMap $ (mp, conditionArg)
            branches =
              map (evaluate3DComplex valMap . (mp,)) branchArgs
         in listArray
              ((0, 0, 0), (size1 - 1, size2 - 1, size3 - 1))
              [ chosen ! (i, j, k)
                | i <- [0 .. size1 - 1],
                  j <- [0 .. size2 - 1],
                  k <- [0 .. size3 - 1],
                  let chosen =
                        chooseBranch marks (cdt ! (i, j, k)) branches
              ]
      Rotate [amount1, amount2, amount3] arg ->
        rotate3D
          (size1, size2, size3)
          (amount1, amount2, amount3)
          (evaluate3DComplex valMap $ (mp, arg))
      _ -> error "expression structure Three C is wrong"
  | otherwise = error "Three C but shape is not [size1, size2, size3] ??"

instance
  (KnownNat m, KnownNat n, KnownNat p) =>
  Evaluable '(m, n, p) C (Array (Int, Int, Int) (Complex Double))
  where
  eval ::
    ValMaps ->
    Expression '(m, n, p) C ->
    Array (Int, Int, Int) (Complex Double)
  eval valMap (Expression n mp) = evaluate3DComplex valMap (mp, n)

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

-- | One dimension rotation
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
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]
    ]

-- |
fourierTransform1D ::
  Int -> Array Int (Complex Double) -> Array Int (Complex Double)
fourierTransform1D size arr =
  listArray (0, size - 1) [computeX i | i <- [0 .. size - 1]]
  where
    computeX i = sum $ zipWithA (*) arr (fourierBasis i)
    fourierBasis i =
      let frequency n = 2 * pi * fromIntegral (i * n) / fromIntegral size
       in listArray
            (0, size - 1)
            [ cos (frequency n) :+ (- sin (frequency n))
              | n <- [0 .. size - 1]
            ]

-- |
fourierTransform2D ::
  (Int, Int) ->
  Array (Int, Int) (Complex Double) ->
  Array (Int, Int) (Complex Double)
fourierTransform2D (size1, size2) arr =
  listArray
    ((0, 0), (size1 - 1, size2 - 1))
    [computeX i j | i <- [0 .. size1 - 1], j <- [0 .. size2 - 1]]
  where
    computeX i j = sum $ zipWithA (*) arr (fourierBasis i j)
    fourierBasis i j =
      let frequency m n =
            2 * pi * fromIntegral (i * m) / fromIntegral size1
              + 2 * pi * fromIntegral (j * n) / fromIntegral size2
       in listArray
            ((0, 0), (size1 - 1, size2 - 1))
            [ cos (frequency m n) :+ (- sin (frequency m n))
              | m <- [0 .. size1 - 1],
                n <- [0 .. size2 - 1]
            ]

-- |
fourierTransform3D ::
  (Int, Int, Int) ->
  Array (Int, Int, Int) (Complex Double) ->
  Array (Int, Int, Int) (Complex Double)
fourierTransform3D (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
    computeX i j k = sum $ zipWithA (*) arr (fourierBasis i j k)
    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
       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 TypeFamilies #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE RecordWildCards, PatternGuards, BangPatterns #-}
-----------------------------------------------------------------------------
--
-- Module      :  AI.Network.RNN.RNN
-- Copyright   :  (c) JP Moresmau
-- License     :  BSD3
--
-- Maintainer  :  JP Moresmau <jp@moresmau.fr>
-- Stability   :  experimental
-- Portability :
--
-- | Recurrent neural network
--
-----------------------------------------------------------------------------

module AI.Network.RNN.RNN  where

import Numeric.LinearAlgebra.HMatrix hiding ((|>))
import Control.DeepSeq
import Control.Monad.Random hiding (fromList)

import AI.Network.RNN.Types
import AI.Network.RNN.Util

data RNNDimensions = RNNDimensions
  { rnndInput    :: Int
  , rnndInternal :: Int
  , rnndOutput   :: Int
  , rnndWithBack :: Bool
  }  deriving (Show,Read,Eq,Ord)

instance NFData RNNDimensions where
    rnf RNNDimensions{..} = rnf (rnndInput,rnndInternal,rnndOutput,rnndWithBack)

data RNNetwork = RNNetwork
  { rnnDimensions :: RNNDimensions
  , rnnMIn        :: Matrix Double
  , rnnM          :: Matrix Double
  , rnnMOut       :: Matrix Double
  , rnnMBack      :: Maybe (Matrix Double)
  , rnnState      :: Vector Double
  , rnnOutput     :: Vector Double
  } deriving (Show,Read,Eq)

networkDimensions :: RNNetwork -> RNNDimensions
networkDimensions = rnnDimensions

instance NFData RNNetwork where
    rnf RNNetwork{..} = rnf (rnnDimensions,rnnMIn,rnnM,rnnMOut,rnnMBack,rnnState,rnnOutput)

createNetwork
    :: RNNDimensions
    -> Matrix Double
    -> Matrix Double
    -> Matrix Double
    -> Maybe (Matrix Double)
    -> Vector Double
    -> Vector Double
    -> Either [String] RNNetwork
createNetwork dim mIn m mOut mmback st out =
    case checkNetworkDimensions dim mIn m mOut mmback st of
        []  -> Right $ RNNetwork dim mIn m mOut mmback st out
        err -> Left err

createNetworkFromArray
    :: RNNDimensions
    -> [Double]
    -> Either String RNNetwork
createNetworkFromArray dim@RNNDimensions{..} v =
    let
        enough = length v >= totalDataLength dim
    in if enough
        then Right $ fromVector dim $ fromList v
        else Left $ "Not enough data in array (needs at least"++show (totalDataLength dim)++")"

networkToArray :: RNNetwork -> [Double]
networkToArray RNNetwork{..} =
       concat (toLists rnnMIn)
    ++ concat (toLists rnnM)
    ++ concat (toLists rnnMOut)
    ++ (case rnnMBack of
            Just b -> concat $ toLists b
            _      -> [])
    ++ toList rnnState
    ++ toList rnnOutput

createNetworkFromVector
    :: RNNDimensions
    -> Vector Double
    -> Either String RNNetwork
createNetworkFromVector dim@RNNDimensions{..} v =
    let
        enough = size v >= totalDataLength dim
    in if enough
        then {-# SCC "createNetworkFromVector" #-} Right $ fromVector dim v
        else Left $ "Not enough data in vector (needs at least"++ show (totalDataLength dim) ++")"

inputMatrixLength :: RNNDimensions -> Int
inputMatrixLength RNNDimensions{..} = rnndInput * rnndInternal

internalMatrixLength :: RNNDimensions -> Int
internalMatrixLength RNNDimensions{..} = rnndInternal * rnndInternal

outputMatrixLength :: RNNDimensions -> Int
outputMatrixLength RNNDimensions{..} = rnndOutput * (rnndInput + rnndInternal)

backMatrixLength :: RNNDimensions -> Int
backMatrixLength RNNDimensions{..} = if rnndWithBack then rnndInternal * rnndOutput else 0

totalDataLength :: FullSize RNNetwork
totalDataLength dim@RNNDimensions{..} =
    inputMatrixLength dim + internalMatrixLength dim + outputMatrixLength dim
    + backMatrixLength dim
    + rnndInternal + rnndOutput

randNetwork :: (Monad m,RandomGen g) =>  RNNDimensions -> RandT g m RNNetwork
randNetwork dim@RNNDimensions{..} = do
    s <- getRandom
    -- vs <- sequence (replicate (totalDataLength dim) getRandom)
    return $ fromVector dim $ randomVector s Gaussian (totalDataLength dim)

collectErrors :: [(Bool,a)] ->  [a]
collectErrors = map snd . filter fst

checkDimensions
    :: RNNDimensions
    -> [String]
checkDimensions RNNDimensions{..} = collectErrors
    [(rnndInput    < 0, "input must be >=0")
    ,(rnndInternal < 1, "internal must be >0")
    ,(rnndOutput   < 1, "output must be >0")]

checkNetworkDimensions
    :: RNNDimensions
    -> Matrix Double
    -> Matrix Double
    -> Matrix Double
    -> Maybe (Matrix Double)
    -> Vector Double
    -> [String]
checkNetworkDimensions dim@RNNDimensions{..} mIn m mOut mmback st = checkDimensions dim ++ collectErrors (
    [(cols mIn /= rnndInput, "input matrix column count doesn't match input neurons count")
    ,(rows mIn /= rnndInternal, "input matrix row count doesn't match internal neurons count")
    ,(cols m /= rnndInternal, "internal matrix column count doesn't match internal neurons count")
    ,(rows m /= rnndInternal, "internal matrix row count doesn't match internal neurons count")
    ,(cols mOut /= (rnndInput + rnndInternal), "internal matrix column count doesn't match input + internal neurons count")
    ,(rows mOut /= rnndOutput, "output matrix row count doesn't match output neurons count")
    ,(size st /= rnndInternal,"internal state length doesn't match internal neurons count")]
    ++(case mmback of
        Nothing -> []
        Just mback ->
            [(rows mback /= rnndInternal,"back matrix row count doesn't match internal neurons count")
            ,(cols m /= rnndOutput,"back matrix column count doesn't match output neurons count")
            ]))

instance RNNEval RNNetwork where
    type Size RNNetwork = RNNDimensions
    evalStep rnn@RNNetwork{..} iv =
        let
            sum1 = (rnnMIn #> iv) + (rnnM #> rnnState)
            sum2 = case rnnMBack of
                Just mback -> sum1 + (mback #> rnnOutput)
                _          -> sum1
            s2 = cmap tanh sum2
            out = cmap sigmoid (rnnMOut #> vjoin [iv,s2])
        in (rnn{rnnState=s2,rnnOutput=out},out)
    toVector RNNetwork{..} = vjoin
        [ flatten rnnMIn
        , flatten rnnM
        , flatten rnnMOut
        , maybe (fromList []) flatten rnnMBack
        , rnnState
        , rnnOutput
        ]
    fromVector dim@RNNDimensions{..} v =
        let
            lenIn   = inputMatrixLength dim
            len     = internalMatrixLength dim
            lenOut  = outputMatrixLength dim
            lenBack = backMatrixLength dim
            [v1,v2,v3,v4,v5,v6] = takesV [lenIn,len,lenOut,lenBack,rnndInternal,rnndOutput] v
            mIn     = reshape rnndInput v1
            m       = reshape rnndInternal v2
            mOut    = reshape (rnndInput + rnndInternal) v3
            mback = if rnndWithBack
                            then Just $ reshape rnndOutput v4
                            else Nothing
        in RNNetwork dim mIn m mOut mback v5 v6
    rnnsize = rnnDimensions
    fullSize = totalDataLength . rnnDimensions


{-# LANGUAGE CPP                   #-}
{-# LANGUAGE DataKinds             #-}
{-# LANGUAGE FlexibleContexts      #-}
{-# LANGUAGE FlexibleInstances     #-}
{-# LANGUAGE GADTs                 #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE RankNTypes            #-}
{-# LANGUAGE ScopedTypeVariables   #-}
{-# LANGUAGE TypeFamilies          #-}
{-# LANGUAGE TypeOperators         #-}
{-# LANGUAGE UndecidableInstances  #-}
{-|
Module      : Grenade.Layers.Deconvolution
Description : Deconvolution layer
Copyright   : (c) Huw Campbell, 2016-2017
License     : BSD2
Stability   : experimental

A deconvolution layer is in many ways a convolution layer in reverse.
It learns a kernel to apply to each pixel location, spreading it out
into a larger layer.

This layer is important for image generation tasks, such as GANs on
images.
-}
module Grenade.Layers.Deconvolution (
    Deconvolution (..)
  , Deconvolution' (..)
  ) where

import           Control.DeepSeq                     (NFData (..))
import           Data.Kind                           (Type)
import           Data.List                           (foldl1')
import           Data.Maybe
import           Data.Proxy
import           Data.Serialize
import           Data.Singletons.TypeLits            hiding (natVal)
import           GHC.TypeLits

import           Numeric.LinearAlgebra               hiding (konst,
                                                      uniformSample)
import qualified Numeric.LinearAlgebra               as LA
import           Numeric.LinearAlgebra.Static        hiding (build, toRows,
                                                      (|||))

import           Grenade.Core
import           Grenade.Layers.Internal.Convolution
import           Grenade.Layers.Internal.Update
import           Grenade.Utils.LinearAlgebra
import           Grenade.Utils.ListStore

-- | A Deconvolution layer for a neural network.
--   This uses the im2col Convolution trick popularised by Caffe.
--
--   The Deconvolution layer is a way of spreading out a single response
--   into a larger image, and is useful in generating images.
--
data Deconvolution :: Nat -- Number of channels, for the first layer this could be RGB for instance.
                   -> Nat -- Number of filters, this is the number of channels output by the layer.
                   -> Nat -- The number of rows in the kernel filter
                   -> Nat -- The number of column in the kernel filter
                   -> Nat -- The row stride of the Deconvolution filter
                   -> Nat -- The columns stride of the Deconvolution filter
                   -> Type where
  Deconvolution :: ( KnownNat channels
                   , KnownNat filters
                   , KnownNat kernelRows
                   , KnownNat kernelColumns
                   , KnownNat strideRows
                   , KnownNat strideColumns
                   , KnownNat kernelFlattened
                   , kernelFlattened ~ (kernelRows * kernelColumns * filters))
                 => !(L kernelFlattened channels) -- The kernel filter weights
                 -> !(ListStore (L kernelFlattened channels)) -- The last kernel update (or momentum)
                 -> Deconvolution channels filters kernelRows kernelColumns strideRows strideColumns

instance NFData (Deconvolution channels filters kernelRows kernelColumns strideRows strideColumns) where
  rnf (Deconvolution a b) = rnf a `seq` rnf b

instance ( KnownNat channels
         , KnownNat filters
         , KnownNat kernelRows
         , KnownNat kernelColumns
         , KnownNat strideRows
         , KnownNat strideColumns
         , KnownNat (kernelRows * kernelColumns * filters)
         ) => Serialize (Deconvolution channels filters kernelRows kernelColumns strideRows strideColumns) where
  put (Deconvolution w store) = do
    putListOf put . toList . flatten . extract $ w
    put (fmap (toList . flatten . extract) store)
  get = do
      let f  = fromIntegral $ natVal (Proxy :: Proxy channels)
      wN    <- maybe (fail "Vector of incorrect size") return . create . reshape f . LA.fromList =<< getListOf get
      store <- fmap (fromMaybe (error "Vector of incorrect size") . create . reshape f . LA.fromList)  <$> get
      return $ Deconvolution wN store

data Deconvolution' :: Nat -- Number of channels, for the first layer this could be RGB for instance.
                    -> Nat -- Number of filters, this is the number of channels output by the layer.
                    -> Nat -- The number of rows in the kernel filter
                    -> Nat -- The number of column in the kernel filter
                    -> Nat -- The row stride of the Deconvolution filter
                    -> Nat -- The columns stride of the Deconvolution filter
                    -> Type where
  Deconvolution' :: ( KnownNat channels
                    , KnownNat filters
                    , KnownNat kernelRows
                    , KnownNat kernelColumns
                    , KnownNat strideRows
                    , KnownNat strideColumns
                    , KnownNat kernelFlattened
                    , kernelFlattened ~ (kernelRows * kernelColumns * filters))
                 => !(L kernelFlattened channels) -- The kernel filter gradient
                 -> Deconvolution' channels filters kernelRows kernelColumns strideRows strideColumns

instance NFData (Deconvolution' channels filters kernelRows kernelColumns strideRows strideColumns) where
  rnf (Deconvolution' a) = rnf a `seq` ()

instance ( KnownNat channels
         , KnownNat filters
         , KnownNat kernelRows
         , KnownNat kernelColumns
         , KnownNat strideRows
         , KnownNat strideColumns
         , KnownNat (kernelRows * kernelColumns * filters)
         ) =>
         Serialize (Deconvolution' channels filters kernelRows kernelColumns strideRows strideColumns) where
  put (Deconvolution' w) = putListOf put . toList . flatten . extract $ w
  get = do
    let f = fromIntegral $ natVal (Proxy :: Proxy channels)
    wN <- maybe (fail "Vector of incorrect size") return . create . reshape f . LA.fromList =<< getListOf get
    return $ Deconvolution' wN

instance (KnownNat channels
         , KnownNat filters
         , KnownNat kernelRows
         , KnownNat kernelColumns
         , KnownNat strideRows
         , KnownNat strideColumns) => FoldableGradient (Deconvolution' channels filters kernelRows kernelColumns strideRows strideColumns) where
  mapGradient f (Deconvolution' kernelGradient) = Deconvolution' (dmmap f kernelGradient)
  squaredSums (Deconvolution' kernelGradient) = [sumM . squareM $ kernelGradient]

instance Show (Deconvolution c f k k' s s') where
  show (Deconvolution a _) = renderConv a
    where
      renderConv mm =
        let m = extract mm
            ky = fromIntegral $ natVal (Proxy :: Proxy k)
            rs = LA.toColumns m
            ms = map (take ky) $ toLists . reshape ky <$> rs
            render n'
              | n' <= 0.2 = ' '
              | n' <= 0.4 = '.'
              | n' <= 0.6 = '-'
              | n' <= 0.8 = '='
              | otherwise = '#'
            px = (fmap . fmap . fmap) render ms
         in unlines $ foldl1 (zipWith (\a' b' -> a' ++ "   |   " ++ b')) px

instance ( KnownNat channels
         , KnownNat filters
         , KnownNat kernelRows
         , KnownNat kernelColumns
         , KnownNat strideRows
         , KnownNat strideColumns
         , KnownNat ((kernelRows * kernelColumns) * filters)
         , KnownNat ((kernelRows * kernelColumns) * channels)
         , KnownNat (channels * ((kernelRows * kernelColumns) * filters))
         ) =>
         RandomLayer (Deconvolution channels filters kernelRows kernelColumns strideRows strideColumns) where
  createRandomWith m gen = do
    wN <- getRandomMatrix i i m gen
    return $ Deconvolution wN mkListStore
    where
      i = natVal (Proxy :: Proxy ((kernelRows * kernelColumns) * channels))

instance ( KnownNat channels
         , KnownNat filters
         , KnownNat kernelRows
         , KnownNat kernelColumns
         , KnownNat strideRows
         , KnownNat strideColumns
         , KnownNat (kernelRows * kernelColumns * filters)
         ) => UpdateLayer (Deconvolution channels filters kernelRows kernelColumns strideRows strideColumns) where
  type Gradient (Deconvolution channels filters kernelRows kernelColumns strideRows strideColumns) = (Deconvolution' channels filters kernelRows kernelColumns strideRows strideColumns)

  type MomentumStore (Deconvolution channels filters kernelRows kernelColumns strideRows strideColumns)  = ListStore (L (kernelRows * kernelColumns * filters) channels)

  runUpdate opt@OptSGD{} x@(Deconvolution oldKernel store) (Deconvolution' kernelGradient) =
    let oldMomentum = getData opt x store
        result = descendMatrix opt (MatrixValuesSGD oldKernel kernelGradient oldMomentum)
        newStore = setData opt x store (matrixMomentum result)
    in Deconvolution (matrixActivations result) newStore
  runUpdate opt@OptAdam{} x@(Deconvolution oldKernel store) (Deconvolution' kernelGradient) =
    let (m, v) = toTuple $ getData opt x store
        result = descendMatrix opt (MatrixValuesAdam (getStep store) oldKernel kernelGradient m v)
        newStore = setData opt x store [matrixM result, matrixV result]
    in Deconvolution (matrixActivations result) newStore
    where toTuple [m ,v] = (m, v)
          toTuple xs = error $ "unexpected input of length " ++ show (length xs) ++ "in toTuple in Convolution.hs"

  reduceGradient grads = Deconvolution' $ dmmap (/ (fromIntegral $ length grads)) (foldl1' add (map (\(Deconvolution' x) -> x) grads))

instance ( KnownNat channels
         , KnownNat filters
         , KnownNat kernelRows
         , KnownNat kernelColumns
         , KnownNat strideRows
         , KnownNat strideColumns
         , KnownNat (kernelRows * kernelColumns * filters)
         ) =>
         LayerOptimizerData (Deconvolution channels filters kernelRows kernelColumns strideRows strideColumns) (Optimizer 'SGD) where
  type MomentumExpOptResult (Deconvolution channels filters kernelRows kernelColumns strideRows strideColumns) (Optimizer 'SGD) = L (kernelRows * kernelColumns * filters) channels
  type MomentumDataType (Deconvolution channels filters kernelRows kernelColumns strideRows strideColumns) (Optimizer 'SGD) = L (kernelRows * kernelColumns * filters) channels
  getData opt x store = head $ getListStore opt x store
  setData opt x store = setListStore opt x store . return
  newData _ _ = konst 0

instance ( KnownNat channels
         , KnownNat filters
         , KnownNat kernelRows
         , KnownNat kernelColumns
         , KnownNat strideRows
         , KnownNat strideColumns
         , KnownNat (kernelRows * kernelColumns * filters)
         ) => LayerOptimizerData (Deconvolution channels filters kernelRows kernelColumns strideRows strideColumns) (Optimizer 'Adam) where
  type MomentumExpOptResult (Deconvolution channels filters kernelRows kernelColumns strideRows strideColumns) (Optimizer 'Adam)  = [L (kernelRows * kernelColumns * filters) channels]
  type MomentumDataType (Deconvolution channels filters kernelRows kernelColumns strideRows strideColumns) (Optimizer 'Adam) = L (kernelRows * kernelColumns * filters) channels
  getData = getListStore
  setData = setListStore
  newData _ _ = konst 0

-- | A two dimentional image may have a Deconvolution filter applied to it
instance ( KnownNat kernelRows
         , KnownNat kernelCols
         , KnownNat filters
         , KnownNat strideRows
         , KnownNat strideCols
         , KnownNat inputRows
         , KnownNat inputCols
         , KnownNat outputRows
         , KnownNat outputCols
         , ((inputRows - 1) * strideRows) ~ (outputRows - kernelRows)
         , ((inputCols - 1) * strideCols) ~ (outputCols - kernelCols)
         , KnownNat (kernelRows * kernelCols * filters)
         , KnownNat (outputRows * filters)
         ) => Layer (Deconvolution 1 filters kernelRows kernelCols strideRows strideCols) ('D2 inputRows inputCols) ('D3 outputRows outputCols filters) where
  type Tape (Deconvolution 1 filters kernelRows kernelCols strideRows strideCols) ('D2 inputRows inputCols) ('D3 outputRows outputCols filters) = S ('D3 inputRows inputCols 1)
  runForwards c (S2D input) =
    runForwards c (S3D input :: S ('D3 inputRows inputCols 1))

  runBackwards c tape grads =
    case runBackwards c tape grads of
      (c', S3D back :: S ('D3 inputRows inputCols 1)) ->  (c', S2D back)

-- | A two dimentional image may have a Deconvolution filter applied to it
instance ( KnownNat kernelRows
         , KnownNat kernelCols
         , KnownNat strideRows
         , KnownNat strideCols
         , KnownNat inputRows
         , KnownNat inputCols
         , KnownNat outputCols
         , ((inputRows - 1) * strideRows) ~ (outputRows - kernelRows)
         , ((inputCols - 1) * strideCols) ~ (outputCols - kernelCols)
         , KnownNat (kernelRows * kernelCols * 1)
         , KnownNat (outputRows * 1)
         ) => Layer (Deconvolution 1 1 kernelRows kernelCols strideRows strideCols) ('D2 inputRows inputCols) ('D2 outputRows outputCols) where
  type Tape (Deconvolution 1 1 kernelRows kernelCols strideRows strideCols) ('D2 inputRows inputCols) ('D2 outputRows outputCols) = S ('D3 inputRows inputCols 1)
  runForwards c (S2D input) =
    case runForwards c (S3D input :: S ('D3 inputRows inputCols 1)) of
      (tps, S3D fore :: S ('D3 outputRows outputCols 1)) ->  (tps, S2D fore)

  runBackwards c tape (S2D grads) =
    case runBackwards c tape (S3D grads :: S ('D3 outputRows outputCols 1)) of
      (c', S3D back :: S ('D3 inputRows inputCols 1)) ->  (c', S2D back)

-- | A two dimentional image may have a Deconvolution filter applied to it
instance ( KnownNat kernelRows
         , KnownNat kernelCols
         , KnownNat strideRows
         , KnownNat strideCols
         , KnownNat inputRows
         , KnownNat inputCols
         , KnownNat outputCols
         , ((inputRows - 1) * strideRows) ~ (outputRows - kernelRows)
         , ((inputCols - 1) * strideCols) ~ (outputCols - kernelCols)
         , KnownNat (kernelRows * kernelCols * 1)
         , KnownNat (outputRows * 1)
         , KnownNat channels
         ) => Layer (Deconvolution channels 1 kernelRows kernelCols strideRows strideCols) ('D3 inputRows inputCols channels) ('D2 outputRows outputCols) where
  type Tape (Deconvolution channels 1 kernelRows kernelCols strideRows strideCols) ('D3 inputRows inputCols channels) ('D2 outputRows outputCols) = S ('D3 inputRows inputCols channels)
  runForwards c input =
    case runForwards c input of
      (tps, S3D fore :: S ('D3 outputRows outputCols 1)) ->  (tps, S2D fore)

  runBackwards c tape (S2D grads) =
    runBackwards c tape (S3D grads :: S ('D3 outputRows outputCols 1))

-- | A three dimensional image (or 2d with many channels) can have
--   an appropriately sized Deconvolution filter run across it.
instance ( KnownNat kernelRows
         , KnownNat kernelCols
         , KnownNat filters
         , KnownNat strideRows
         , KnownNat strideCols
         , KnownNat inputRows
         , KnownNat inputCols
         , KnownNat outputRows
         , KnownNat outputCols
         , KnownNat channels
         , ((inputRows - 1) * strideRows) ~ (outputRows - kernelRows)
         , ((inputCols - 1) * strideCols) ~ (outputCols - kernelCols)
         , KnownNat (kernelRows * kernelCols * filters)
         , KnownNat (outputRows * filters)
         ) => Layer (Deconvolution channels filters kernelRows kernelCols strideRows strideCols) ('D3 inputRows inputCols channels) ('D3 outputRows outputCols filters) where

  type Tape (Deconvolution channels filters kernelRows kernelCols strideRows strideCols) ('D3 inputRows inputCols channels) ('D3 outputRows outputCols filters) = S ('D3 inputRows inputCols channels)

  runForwards (Deconvolution kernel _) (S3D input) =
    let ex = extract input
        ek = extract kernel
        ix = fromIntegral $ natVal (Proxy :: Proxy inputRows)
        iy = fromIntegral $ natVal (Proxy :: Proxy inputCols)
        kx = fromIntegral $ natVal (Proxy :: Proxy kernelRows)
        ky = fromIntegral $ natVal (Proxy :: Proxy kernelCols)
        sx = fromIntegral $ natVal (Proxy :: Proxy strideRows)
        sy = fromIntegral $ natVal (Proxy :: Proxy strideCols)
        ox = fromIntegral $ natVal (Proxy :: Proxy outputRows)
        oy = fromIntegral $ natVal (Proxy :: Proxy outputCols)

        c  = vid2col 1 1 1 1 ix iy ex

        mt = c LA.<> tr ek

        r  = col2vid kx ky sx sy ox oy mt
        rs = fromJust . create $ r
    in  (S3D input, S3D rs)
  runBackwards (Deconvolution kernel _) (S3D input) (S3D dEdy) =
    let ex = extract input
        ix = fromIntegral $ natVal (Proxy :: Proxy inputRows)
        iy = fromIntegral $ natVal (Proxy :: Proxy inputCols)
        kx = fromIntegral $ natVal (Proxy :: Proxy kernelRows)
        ky = fromIntegral $ natVal (Proxy :: Proxy kernelCols)
        sx = fromIntegral $ natVal (Proxy :: Proxy strideRows)
        sy = fromIntegral $ natVal (Proxy :: Proxy strideCols)
        ox = fromIntegral $ natVal (Proxy :: Proxy outputRows)
        oy = fromIntegral $ natVal (Proxy :: Proxy outputCols)

        c  = vid2col 1 1 1 1 ix iy ex

        eo = extract dEdy
        ek = extract kernel

        vs = vid2col kx ky sx sy ox oy eo

        kN = fromJust . create . tr $ tr c LA.<> vs

        dW = vs LA.<> ek

        xW = col2vid 1 1 1 1 ix iy dW
    in  (Deconvolution' kN, S3D . fromJust . create $ xW)


{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE ExistentialQuantification #-}
{-# LANGUAGE FlexibleContexts #-}
-----------------------------------------------------------------------------
--
-- Module      :  AI.Network.RNN.Genetic
-- Copyright   :  (c) JP Moresmau
-- License     :  BSD3
--
-- Maintainer  :  JP Moresmau <jp@moresmau.fr>
-- Stability   :  experimental
-- Portability :
--
-- | Code for genetic evolution of a network
--
-----------------------------------------------------------------------------

module AI.Network.RNN.Genetic where

import Control.Monad
import Control.Monad.Random hiding (fromList)
import qualified Control.Monad.Random as MR
import Control.DeepSeq
import Data.IORef
import AI.GeneticAlgorithm.Simple
import AI.Network.RNN.Types
import AI.Network.RNN.Util
import Numeric.LinearAlgebra.HMatrix

import Debug.Trace

-- | Keep all relevant data together
data RNNData a b = (RNNEval a) => RNNData a !(TrainData b) !Double
   -- deriving (Show,Read,Eq)

-- | evaluate network
instance NFData (RNNData a b) where
    rnf (RNNData rnn _ _) = rnf rnn

-- | Error threshold (not currently used)
errorThreshold :: Double
errorThreshold = 0.1

-- | Chromosome instance
instance Chromosome (RNNData a b) where
    -- | cross network by taking half the values of one and half of the other
    crossover (RNNData rnn1 td l) (RNNData rnn2 _ _) = do
        rnns <- crossNetworkFull rnn1 rnn2
        --idx <- getRandomR (0, size rnn1 - 1)
        --let (v3,v4) = crossVector rnn1 rnn2 idx
        return $ map (\r->RNNData r td l) rnns -- [v3,v4]

    mutation (RNNData rnn1 td l) = do
        rnn2 <- mutateNetwork rnn1
        --rnn2 <- stdNormalMutate rnn1
        return $ RNNData rnn2 td l
    -- | fitness is "higher is better", inverse of cost
    fitness (RNNData rnn1 td _) = 1 / cost rnn1 td -- (fromVector sz rnn1)

-- | Cross 2 networks by taking the first half of one, the second half of the other (not currently used)
crossNetworkFull :: (RandomGen g,RNNEval a) =>  a -> a -> Rand g [a]
crossNetworkFull rnn1 rnn2 = do
    let v1 = toVector rnn1
        v2 = toVector rnn2
    idx <- getRandomR (0, size v1 - 1)
    let (v3,v4) = crossVector v1 v2 idx
        rnn3 = fromVector (rnnsize rnn1) v3
        rnn4 = fromVector (rnnsize rnn1) v4
    return $[rnn3,rnn4]

-- | cross 2 networks by taking half the values of one and half of the other
crossNetworkHalf :: (RandomGen g,RNNEval a) =>  a -> a -> Rand g [a]
crossNetworkHalf rnn1 rnn2 = do
    let v1 = toVector rnn1
        v2 = toVector rnn2
    (v3,v4) <- mixVector v1 v2 0.5
    let
        rnn3 = fromVector (rnnsize rnn1) v3
        rnn4 = fromVector (rnnsize rnn1) v4
    return $ force [rnn3,rnn4]

crossMutate :: (RandomGen g,RNNEval a) =>  a -> a -> Rand g [a]
crossMutate rnn1 rnn2 = do
    let v1 = toVector rnn1
        v2 = toVector rnn2
    v3 <- stdNormalMutate v1
    v4 <- stdNormalMutate v2
    let
        rnn3 = fromVector (rnnsize rnn1) v3
        rnn4 = fromVector (rnnsize rnn1) v4
    return [rnn3,rnn4]

-- | Cross networks by taking the average of their values (not currently used)
avgNetwork :: (RandomGen g,RNNEval a) =>  a -> a -> Rand g [a]
avgNetwork rnn1 rnn2 = do
    let v1 = toVector rnn1
        v2 = toVector rnn2
        v3 = fromList $ zipWith (\x y -> (x+y)/2) (toList v1) (toList v2)
        rnn3 = fromVector (rnnsize rnn1) v3
      --  rnn4 = fromVector (rnnsize rnn1) v4 of
      --      Right n -> n
      --      Left e  -> error e
    return [rnn3]

crossMatrixEq :: Matrix Double -> Matrix Double -> (Matrix Double,Matrix Double)
crossMatrixEq m1 m2 =
    let row1      = upperHalf $ rows m1
        col1      = upperHalf $ cols m1
        [[a1,a2],[a3,a4]] = toBlocksEvery row1 col1 m1
        [[b1,b2],[b3,b4]] = toBlocksEvery row1 col1 m2
    in (fromBlocks [[a1,b2],[a3,b4]],fromBlocks [[b1,a2],[b3,a4]])
    where
        upperHalf a = (if odd a then a+1 else a) `div` 2

crossVectorEq :: Vector Double -> Vector Double -> (Vector Double,Vector Double)
crossVectorEq v1 v2 = crossVector v1 v2 (size v1 `div` 2)

-- | Cross vector at the given point
crossVector :: Vector Double -> Vector Double -> Int -> (Vector Double,Vector Double)
crossVector v1 v2 idx =
--    let m1 = asRow v1
--        m2 = asRow v2
--        [[a1,a2]] = toBlocks [1] [idx,size v1-idx]  m1
--        [[b1,b2]] = toBlocks [1] [idx,size v1-idx] m2
--    in (head $ toRows $ fromBlocks [[a1,b2]],head $ toRows $ fromBlocks [[b1,a2]])
--    let sz2 = size v2 - idx - 1
--        [a1,a2] = takesV [idx,sz2] v1
--        [b1,b2] = takesV [idx,sz2] v2
--    in (vjoin [a1,b2],vjoin [a2,b1])
 let sz = size v1-idx
 in force $ (vjoin [subVector 0 idx v1,subVector idx sz v2],
                    vjoin [subVector 0 idx v2,subVector idx sz v1])

-- | Mix vector values given a probability to take a value from the first one and not the second
mixVector :: (Monad m,RandomGen g) =>  Vector Double -> Vector Double -> Double -> RandT g m (Vector Double,Vector Double)
mixVector v1 v2 prob = do
    -- rs <- replicateM (size v1) (getRandomR (0, 1))
    let (l1,l2) = swapL (toList v1) (toList v2) True
        -- unzip $ map swapR $ zip3 (toList v1) (toList v2) [0..]
    return (fromList l1,fromList l2)
    where
        swapR (a,b,idx) = if (idx `mod` 2) == 0 then (a,b) else (b,a)
        --if idx <= prob then (a,b) else (b,a):
        swapL :: [a] -> [a] -> Bool -> ([a],[a])
        swapL [] _ _ = ([],[])
        swapL (x:xs) (y:ys) True = let
            (rs1,rs2)=swapL xs ys False in (x:rs1,y:rs2)
        swapL (x:xs) (y:ys) False = let
            (rs1,rs2)=swapL xs ys True in (y:rs1,x:rs2)

-- | Mix vector values given a probability to take a value from the first one and not the second
mixList :: (Monad m,RandomGen g) =>  [Double] -> [Double] -> Double -> RandT g m ([Double],[Double])
mixList v1 v2 prob = do
    --rs <- getRandomRs (0,1)
    return $ unzip $ map swapR $ zip3 v1 v2 [0..]
    where
        swapR (a,b,idx) = if (idx `mod` 2) == 0 then (a,b) else (b,a)

stdNormalMutate :: (Monad m,RandomGen g) =>  Vector Double -> RandT g m (Vector Double)
stdNormalMutate v1 = do
    --rs <- replicateM (size v1) stdNormal
    -- let mx = maxElement v1
    -- let mn = minElement v1
    --let sz= trace ("v1:"++ show (mn,mx)) $ size v1
    seed <- getRandom
    let rs = randomVector seed Gaussian $ size v1
    let v2 = (rs) + v1
--    let nv2 = cmap (\x->
--                            if x < 0
--                                then 0
--                                else if x > 1
--                                    then 1
--                                    else x)
--         v2
    -- let mx2 = maxElement v2
    -- let mn2 = minElement v2
    return v2
    -- return $ trace ("v2:"++show (mn2,mx2)) v2

-- | network mutation: select randomly from several mutation algorithms
mutateNetwork :: (RandomGen g,RNNEval a) =>  a -> Rand g a
mutateNetwork rnn = do
    f <- MR.fromList [(stdNormalMutation,1),(pointMutation,1),(swapMutation,1),(insertMutation,1),(flipMutation,0.5)]
    f rnn

-- | Mutate values around a standard deviation
stdNormalMutation :: (RandomGen g,RNNEval a) =>  a -> Rand g a
stdNormalMutation rnn = do
    let v1 = toVector rnn
    v2 <- stdNormalMutate v1
    return $ force $ fromVector (rnnsize rnn) v2

-- | Mutate a single value of the vector, taking a random value
pointMutation :: (RandomGen g,RNNEval a) =>  a -> Rand g a
pointMutation rnn = do
    let v1 = toVector rnn
    idx <- getRandomR (0, size v1 - 1)
    dbl <- getRandomR (0,1)
    let v2 = accum v1 const [(idx,dbl)]
    return $ fromVector (rnnsize rnn) v2

-- | Swap two values in the vector
swapMutation :: (RandomGen g,RNNEval a) =>  a -> Rand g a
swapMutation rnn = do
    let v1 = toVector rnn
    idx1 <- getRandomR (0, size v1 - 1)
    idx2 <- getRandomR (0, size v1 - 1)
    let v2 = accum v1 const [(idx1,atIndex v1 idx2),(idx2,atIndex v1 idx1)]
    return $ fromVector (rnnsize rnn) v2

-- | Insert a random value at a random point in the vector, discarding the last value
insertMutation :: (RandomGen g,RNNEval a) =>  a -> Rand g a
insertMutation rnn = do
    let v1 = toVector rnn
    idx <- getRandomR (0, size v1 - 1)
    dbl <- getRandomR (0,1)
    let v2 = vjoin [subVector 0 idx v1,fromList [dbl],subVector idx (size v1 - idx -1) v1]
    return $ fromVector (rnnsize rnn) v2

flipMutation :: (RandomGen g,RNNEval a) =>  a -> Rand g a
flipMutation rnn = do
    let v1 = toVector rnn
    let v2 = cmap (\a->(-a)) v1
    return $ fromVector (rnnsize rnn) v2

-- | Stop function
-- we take a list of past fitnesses and stop when the fitness 5 generations ago was not worse that the current one
stopf :: IORef [Double] -> Int -> (RNNData a b,Double) -> Int -> IO Bool
stopf fitnessList maxGen (nd,fit) gen = do
    mfit <- atomicModifyIORef' fitnessList (\l->
        let l2 = take 5 (fit:l)
        in if length l2 == 5
            then (l2, Just $ last l2)
            else (l2, Nothing)
        )

    print $ "Fitness (" ++ show gen ++"): " ++ show fit
    let converged = case mfit of
                        Nothing -> False
                        Just f  -> f >= fit
    when converged $ print "Converged!"
    return ( gen >= maxGen)

-- | Build a random network data
buildNetworkData :: (RNNEval re,Monad m,RandomGen g) => TrainData b -> Size re -> FullSize re -> RandT g m (RNNData re b)
buildNetworkData td@(TrainData is _ _ _) dim fullSz = do
  --n <- randomNetwork dim fullSz
  let n = startNetwork dim fullSz
  return $ RNNData n td (fromIntegral $ length is)


module BondCalculus.ODEExtraction
  (IVP(..), PrintStyle(..), ODE(..), AsSage(..),
  matlabExpr, sympyExpr, sageExpr, matlabODE, vectorFieldToODEs, sageODE, sageExprAbst,
  sageExprAbst', extractIVP, sympyODE, solveODEPython, printODEPython,
  sympySimplify, runPython, generateSage) where

import BondCalculus.Base
import BondCalculus.Symbolic
import BondCalculus.Processes
import BondCalculus.Vector
import BondCalculus.AST (pretty)
import qualified BondCalculus.AST as AST
import qualified Data.HashMap.Strict as H
import BondCalculus.Simulation (Trace)
import Data.String.Utils
import Text.Printf

-- import qualified Control.Exception as X
import qualified Data.Map as M
import qualified Data.List as L
import Data.Maybe

--import qualified Numeric.LinearAlgebra as LA

import System.IO.Unsafe
-- import qualified System.Process as OS
import System.Process (readProcess)
-- import Data.Maybe

newtype IVP a = IVP ([String], [Expr (Atom a)], [a]) deriving (Show, Eq)
newtype ODE a = ODE ([String], [Expr (Atom a)])

-------------------------------
-- MATLAB script output:
-------------------------------

extractODE :: (ExprConstant a) => AST.Env -> ConcreteAffinityNetwork a -> P' a -> ODE a
extractODE env network p = vectorFieldToODEs p'
  where -- p' must be expressed in the same basis of p
        -- otherwise the variable order becomes
        -- inconsistient, and we get extra variables
        -- p' :: Vect AST.Species (SymbolicExpr a)
        p' = fromList [(p'coeff i, i) | (_,i) <- toList p]
        -- p'coeff :: AST.Species -> SymbolicExpr a
        p'coeff i = fromMaybe 0.0 (H.lookup i v)
        Vect v = dPdt' tr network p
        tr = tracesGivenNetwork network env

extractIVP :: ExprConstant a => AST.Env -> ConcreteAffinityNetwork a -> P' a -> [a] -> IVP a
extractIVP env network p = fromODEToIVP (extractODE env network p)

vectorFieldToODEs :: ExprConstant a => P' a -> ODE a
vectorFieldToODEs v = ODE (vars, rhss)
  where
    vars = map (pretty.snd) kis
    rhss = map (simplify.fst) kis
    kis = toList v

fromODEToIVP :: ODE a -> [a] -> IVP a
fromODEToIVP (ODE (x,y)) z = IVP (x, y, take (length x) z)

-- modelToIVP :: BondCalculusModel SymbolicExpr Double -> IVP

matlabODE :: IVP Double -> (Int,(Double,Double)) -> Either String String
matlabODE ivp (n,(t0,tn)) =
  if any isNothing eqns
  then Left "An ODE equation has an unbound variable"
  else Right $
  "init = [" ++ L.intercalate ";" (map show inits) ++ "];\n" ++
    "t = linspace(" ++ show t0 ++ "," ++ show tn ++ "," ++ show n ++ ");\n" ++
    "function xdot = f(x,t) \n\n" ++
    L.intercalate ";\n" (catMaybes eqns) ++
    ";\nendfunction;\n\n" ++
    -- Need new support for calculating jacobian
    -- "function jac = jj(x,t)\n\n" ++
    -- matlabJac env p' ++ ";\nendfunction;\n\n" ++
    "x = lsode(f, init, t);\n" ++
    "save (\"-ascii\", \"-\", \"x\");"
      where
        xdots = ["xdot(" ++ show i ++ ")" | i <- [(1::Integer)..]]
        xvars = ["x(" ++ show i ++ ")" | i <- [(1::Integer)..]]
        varmap = M.fromList $ zip vars xvars
        IVP (vars,rhss,inits) = ivp
        eqns = [fmap (\y -> xdot ++ " = " ++ y) (matlabExpr varmap rhs)
               | (xdot,rhs) <- zip xdots rhss]

class Show a => AsSage a where
    asSage :: a -> String

instance AsSage Double where
    asSage = show

instance AsSage Interval where
    asSage x = printf "RIF('[%.20f .. %.20f]')" l u
        where (l, u) = endpoints x
              -- Do some fiddling with the endpoints to ensure sound enclosure
            --   l' = l - 6e-21
            --   u' = u + 6e-21

sageExpr :: AsSage a => M.Map String String -> Expr (Atom a) -> Maybe String
sageExpr mp (Atom (Var x)) = M.lookup x mp
sageExpr mp (Atom (Const x)) = return $ asSage x
sageExpr mp (x `Sum` y) = do
    x' <- sageExpr mp x
    y' <- sageExpr mp y
    return $ x' ++ " + " ++ y'
sageExpr mp (x `Prod` y) = do
    x' <- sageExpr mp x
    y' <- sageExpr mp y
    return $ "(" ++ x' ++ ") * (" ++ y' ++ ")"
sageExpr mp (x `Frac` y) = do
    x' <- sageExpr mp x
    y' <- sageExpr mp y
    return $ "(" ++ x' ++ ") / (" ++ y' ++ ")"
sageExpr mp (x `Pow` y) = do
    x' <- sageExpr mp x
    y' <- sageExpr mp y
    return $ "(" ++ x' ++ ") ** (" ++ y' ++ ")"
sageExpr mp (Abs x) = do
    x' <- sageExpr mp x
    return $ "abs(" ++ x' ++ ")"
sageExpr mp (Sin x) = do
    x' <- sageExpr mp x
    return $ "sin(" ++ x' ++ ")"
sageExpr mp (Cos x) = do
    x' <- sageExpr mp x
    return $ "sin(" ++ x' ++ ")"
sageExpr mp (Tan x) = do
    x' <- sageExpr mp x
    return $ "tan(" ++ x' ++ ")"
sageExpr mp (Exp x) = do
    x' <- sageExpr mp x
    return $ "exp(" ++ x' ++ ")"
sageExpr mp (Log x) = do
    x' <- sageExpr mp x
    return $ "log(" ++ x' ++ ")"

-- Convert a sage expression into a string with all non-simple double constants (e.g. intervals)
-- abstracted by symbolic variables
sageExprAbst' :: (AsSage a, Boundable a) => M.Map String String -> Expr (Atom a) -> Maybe (M.Map String String, String)
sageExprAbst' = sageExprAbst M.empty

sageExprAbst :: (AsSage a, Boundable a) => M.Map String String -> M.Map String String -> Expr (Atom a) -> Maybe (M.Map String String, String)
sageExprAbst cmp mp (Atom (Var x)) = do
    x' <- M.lookup x mp
    return (cmp, x')
sageExprAbst cmp mp (Atom (Const x)) = return $ case singleValue x of
    Just 1.0    -> (cmp, "1")
    Just (-1.0) -> (cmp, "-1")
    Just 0.0    -> (cmp, "0")
    Just v      -> (cmp, asSage v)
    Nothing     -> (cmp', s)
        where n = M.size cmp
              s = "a" ++ show n
              cmp' = M.insert s (asSage x) cmp
sageExprAbst cmp mp (x `Sum` y) = f2 cmp mp x y $ \x' y' -> x' ++ " + " ++ y'
sageExprAbst cmp mp (x `Prod` y) = f2 cmp mp x y $ \x' y' -> "(" ++ x' ++ ") * (" ++ y' ++ ")"
sageExprAbst cmp mp (x `Frac` y) = f2 cmp mp x y $ \x' y' -> "(" ++ x' ++ ") / (" ++ y' ++ ")"
sageExprAbst cmp mp (x `Pow` y) = f2 cmp mp x y $ \x' y' -> "(" ++ x' ++ ") ** (" ++ y' ++ ")"
sageExprAbst cmp mp (Abs x) = f1 cmp mp x $ \x' -> "abs(" ++ x' ++ ")"
sageExprAbst cmp mp (Sin x) = f1 cmp mp x $ \x' -> "sin(" ++ x' ++ ")"
sageExprAbst cmp mp (Cos x) = f1 cmp mp x $ \x' -> "cos(" ++ x' ++ ")"
sageExprAbst cmp mp (Tan x) = f1 cmp mp x $ \x' -> "tan(" ++ x' ++ ")"
sageExprAbst cmp mp (Exp x) = f1 cmp mp x $ \x' -> "exp(" ++ x' ++ ")"
sageExprAbst cmp mp (Log x) = f1 cmp mp x $ \x' -> "log(" ++ x' ++ ")"
sageExprAbst cmp mp x = error $ "Support for converting expression " ++ show x ++ " to sage/sympy not implemented!"

f2 cmp mp x y g = do
    (cmp', x') <- sageExprAbst cmp mp x
    (cmp'', y') <- sageExprAbst cmp' mp y
    return (cmp'', g x' y')

f1 cmp mp x g = do
    (cmp', x') <- sageExprAbst cmp mp x
    return (cmp', g x')

-- This assumes we can reduce the ODEs to sage
sageODE :: (AsSage a, Boundable a) => IVP a -> String -> Either String String
sageODE (IVP (vars, rhss, inits)) network = case odeExprs of
        Nothing -> Left "An ODE equation has an unbound variable"
        Just (cmp, odes) -> Right $
            "from sage.all import *\n" ++
            "import sympy as sym\n\n" ++
            "R, x = PolynomialRing(RIF, " ++ show nvars ++ ", ', '.join(map('x{}'.format, range(0, " ++ show nvars ++ ")))).objgens()\n" ++
            "xsr = [SR.var(str(x1)) for x1 in x]\n" ++
            "for v in xsr:\n"++
            "   assume(v, 'real')\n"++
            "xsym = sym.var(','.join(map('x{}'.format, range(0," ++ show nvars ++ "))))\n" ++
            (if (M.size cmp) > 0
             then "asym = sym.var(','.join(map('a{}'.format, range(0," ++ show (M.size cmp) ++ "))))\n"
             else "asym = []\n") ++
            "y0 = [" ++ L.intercalate ", " (map asSage inits) ++ "]\n" ++
            "ysymraw = [" ++ L.intercalate ", " odes ++ "]\n" ++
            "ysym = [sym.simplify(y1) for y1 in ysymraw]\n" ++
            "ysr = [y1._sage_().substitute(" ++ subsExpr ++ ") for y1 in ysym]\n" ++
            "varmap = {" ++ L.intercalate ", " varmappings ++ "}\n" ++
            "varmapsr = {" ++ L.intercalate ", " varmappingssr ++ "}\n" ++
            "affinity_network = \"" ++ network ++ "\"\n" ++
            "try:\n" ++
            "    y = vector([R(y1) for y1 in ysr])\n" ++
            "except TypeError:\n" ++
            "    y = None\n"
            where subsExpr = L.intercalate ","
                             [x ++ "=" ++ y | (x, y) <- M.toList cmp]
                  varmappings = [ "'" ++ p ++ "': x[" ++ show i ++ "]"
                                | p <- vars | i <- [0..] ]
                  varmappingssr = [ "'" ++ p ++ "': xsr[" ++ show i ++ "]"
                                  | p <- vars | i <- [0..] ]
    where
    --   odeExprs' :: AsSage a => M.Map String String -> [Expr (Atom a)] -> Maybe (M.Map String String, [String])
      odeExprs' cmp (x:xs) = do
        (cmp', expr) <- sageExprAbst cmp varmap x
        (cmp'', exprs) <- odeExprs' cmp' xs
        return (cmp'', expr:exprs)
      odeExprs' cmp [] = Just (cmp, [])
      odeExprs :: Maybe (M.Map String String, [String])
      odeExprs = odeExprs' M.empty rhss
      nvars = length vars
      xvars = ["x" ++ show n | (n, _) <- zip [0..] vars]
      varmap = M.fromList $ zip vars xvars
    --   odes = map odeExpr rhss
    --   odeExpr rhs = do
    --     expr <- sageExpr varmap rhs
    --     return $ "sym.simplify(" ++ expr ++ ")"

sympyODE :: IVP Double -> (Int,(Double,Double)) -> Either String String
sympyODE ivp (n,(t0,tn)) =
  if any isNothing eqns
  then Left "An ODE equation has an unbound variable"
  else Right $
  "import sympy as sym\n" ++
  "import numpy as np\n" ++
  "from sympy.abc import t\n" ++
  "from scipy.integrate import odeint\n" ++
  "import sys\n" ++
  "sys.setrecursionlimit(100000)\n\n" ++
  "xs = [" ++ L.intercalate ","
    ["sym.Function('" ++ show x ++ "')" | x <- vars ] ++ "]\n" ++
  "xts = [x(t) for x in xs]\n" ++
  "odes = [" ++ L.intercalate ", " (catMaybes eqns) ++ "]\n" ++
  "y0 = [" ++ L.intercalate ", " (map show inits) ++ "]\n" ++
  "ts = np.linspace(" ++ show t0 ++ "," ++ show tn ++ "," ++ show n ++ ")\n" ++
  "rhss = [eqn.rhs for eqn in odes]\n" ++
  "Jac = sym.Matrix(rhss).jacobian(xts)\n" ++
  "f = sym.lambdify((xts, t), rhss, modules='numpy')\n" ++
  "J = sym.lambdify((xts, t), Jac, modules='numpy')\n" ++
  -- Try to solve ODEs using symbolic Jacobian
  "try: ys = odeint(f, y0, ts, (), J)\n" ++
  -- Fallback to using method without Jacobian
  "except NameError: ys = odeint(f, y0, ts, ())\n" ++
  -- "print(ys)\n" ++
  "print('\\n'.join([' '.join([('%.18e' % y).replace('nan', '0') for y in ysa]) for ysa in ys]))"
    where
      xvars = ["xts[" ++ show i ++ "]" | i <- [(0::Integer)..]]
      xdots = [xvar ++ ".diff()" | xvar <- xvars]
      varmap = M.fromList $ zip vars xvars
      IVP (vars,rhss,inits) = ivp
      mkeqn xdot y = "sym.Eq(" ++ xdot ++ ", sym.simplify(" ++ y ++ "))"
      eqns = [fmap (mkeqn xdot) (sympyExpr varmap rhs)
             | (xdot,rhs) <- zip xdots rhss]

data PrintStyle = Plain | Pretty | LaTeX | MathML deriving (Eq)

sympyODEPrint :: ODE Double -> PrintStyle -> Either String String
sympyODEPrint ode style =
  if any isNothing eqns
  then Left "An ODE equation has an unbound variable"
  else Right $
  "import sympy as sym\n" ++
  "import numpy as np\n" ++
  "from sympy.abc import t\n" ++
  "import sys\n" ++
  "sys.setrecursionlimit(100000)\n\n" ++
  -- "sym.init_printing(" ++ printingoptions ++ ")\n\n" ++
  "names = [" ++ L.intercalate "," ["'[" ++ (if style == LaTeX then replace "->" "\\\\rightarrow " (replace " " "\\\\ " x) else x) ++ "]'" | x <- vars] ++ "]\n" ++
  "xs = list(map(sym.Function, names))\n" ++
  "xts = [x(t) for x in xs]\n" ++
  "odes = [" ++ L.intercalate ", " (catMaybes eqns) ++ "]\n" ++
  "rhss = [eqn.rhs for eqn in odes]\n" ++
  case style of
    Pretty -> "for ode in odes: print(sym.pretty(ode))\n"
    Plain -> "for ode in odes: print(ode)\n"
    -- print system of odes as latex (formatted for use in an align environment)
    LaTeX -> "print(sym.latex(odes)[7:-7].replace('=', '& =').replace(', \\\\  ', ' \\\\\\\\\\n').replace('{\\\\left(t \\\\right)}', '').replace('\\\\frac{d}{d t}', '\\\\frac{\\\\mathrm d}{\\\\mathrm d t}').replace('operatorname', 'mathrm'))\n"
    MathML -> "from sympy.printing import print_mathml\nprint_mathml(odes)\n"
    where
      xvars = ["xts[" ++ show i ++ "]" | i <- [(0::Integer)..]]
      xdots = [xvar ++ ".diff()" | xvar <- xvars]
      varmap = M.fromList $ zip vars xvars
      ODE (vars,rhss) = ode
      mkeqn xdot y = "sym.Eq(" ++ xdot ++ ", sym.simplify(" ++ y ++ "))"
      eqns = [fmap (mkeqn xdot) (sympyExpr varmap rhs)
             | (xdot,rhs) <- zip xdots rhss]

matlabExpr :: M.Map String String -> SymbolicExpr Double -> Maybe String
matlabExpr mp (Atom (Var x)) = M.lookup x mp
matlabExpr _ (Atom (Const x)) = return $ show x
matlabExpr mp (x `Sum` y) = do
  x' <- matlabExpr mp x
  y' <- matlabExpr mp y
  return $ x' ++ " .+ " ++ y'
matlabExpr mp (x `Prod` y) = do
  x' <- matlabExpr mp x
  y' <- matlabExpr mp y
  return $ "(" ++ x' ++ ") .* (" ++ y' ++ ")"
matlabExpr mp (x `Pow` y) = do
  x' <- matlabExpr mp x
  y' <- matlabExpr mp y
  return $ "(" ++ x' ++ ") .** (" ++ y' ++ ")"
matlabExpr mp (Abs x) = do
  x' <- matlabExpr mp x
  return $ "abs(" ++ x' ++ ")"

sympyExpr :: M.Map String String -> SymbolicExpr Double -> Maybe String
sympyExpr mp (Atom (Var x)) = M.lookup x mp
sympyExpr _ (Atom (Const x)) = return $ show x
sympyExpr mp (x `Sum` y) = do
  x' <- sympyExpr mp x
  y' <- sympyExpr mp y
  return $ x' ++ " + " ++ y'
sympyExpr mp (x `Prod` y) = do
  x' <- sympyExpr mp x
  y' <- sympyExpr mp y
  return $ "(" ++ x' ++ ") * (" ++ y' ++ ")"
sympyExpr mp (x `Frac` y) = do
  x' <- sympyExpr mp x
  y' <- sympyExpr mp y
  return $ "(" ++ x' ++ ") / (" ++ y' ++ ")"
sympyExpr mp (x `Pow` y) = do
  x' <- sympyExpr mp x
  y' <- sympyExpr mp y
  return $ "(" ++ x' ++ ") ** (" ++ y' ++ ")"
sympyExpr mp (Abs x) = do
  x' <- sympyExpr mp x
  -- return $ x'
  return $ "abs(" ++ x' ++ ")"
sympyExpr mp (Sin x) = do
  x' <- sympyExpr mp x
  return $ "sin(" ++ x' ++ ")"
sympyExpr mp (Cos x) = do
  x' <- sympyExpr mp x
  return $ "sin(" ++ x' ++ ")"
sympyExpr mp (Tan x) = do
  x' <- sympyExpr mp x
  return $ "tan(" ++ x' ++ ")"
sympyExpr mp (Exp x) = do
  x' <- sympyExpr mp x
  return $ "exp(" ++ x' ++ ")"
sympyExpr mp (Log x) = do
  x' <- sympyExpr mp x
  return $ "log(" ++ x' ++ ")"

---------------------------------------
-- Using Octave to execute the scripts
---------------------------------------

runPython :: String -> IO String
runPython = readProcess "python3" ["-q"]

callSolveODEPython :: AST.Env -> ConcreteAffinityNetwork Double -> P' Double -> [Double] -> (Int, (Double, Double)) -> IO String
callSolveODEPython env network p inits tr = case sympyODE (extractIVP env network p inits) tr of
  Right script -> do
    putStrLn $ "Python script:\n\n" ++ script
    writeFile "script.py" script
    runPython script
  Left _ -> undefined

generateSage :: (ExprConstant a, AsSage a)
             => String
             -> AST.Env
             -> ConcreteAffinityNetwork a
             -> P' a
             -> [a]
             -> IO String
generateSage filename env network p inits = case sageODE (extractIVP env network p inits) (pretty network) of
  Right script -> do
    putStrLn $ "Sage script:\n\n" ++ script
    writeFile filename script
    return script
  Left _ -> undefined

callPrintODEPython :: AST.Env -> ConcreteAffinityNetwork Double -> P' Double -> PrintStyle -> IO String
callPrintODEPython env network p style = case sympyODEPrint (extractODE env network p) style of
  Right script -> do
    -- putStrLn $ "Python script:\n\n" ++ script
    writeFile "print_script.py" script
    runPython script
  Left _ -> undefined

solveODEPython :: AST.Env -> ConcreteAffinityNetwork Double -> P' Double -> [Double] -> (Int, (Double, Double)) -> Trace
solveODEPython env network p inits tr@(n,(t0,tn))
  = let raw = unsafePerformIO (callSolveODEPython env network p inits tr)
        ts = [t0 + fromIntegral i*(tn-t0)/fromIntegral n | i <- [0..n]]
        yss = (map (map read.words) $ lines raw) :: [[Double]]
        ys = [fromList (xs `zip` pbasis) | xs <- yss]
        pbasis = map snd $ toList p
    in ts `zip` ys

sympySimplify :: String -> String
sympySimplify s = unsafePerformIO (runPython script)
  where script = "from sympy import simplify, sympify\nimport re\n" ++
                 "print(simplify(sympify(re.sub(r'([0-9]+)\\.0(?![0-9]*[1-9])', r'\\1', " ++ show s ++ "))))"

printODEPython :: AST.Env -> ConcreteAffinityNetwork Double -> P' Double -> PrintStyle -> String
printODEPython env network p style
  = let raw = unsafePerformIO (callPrintODEPython env network p style)
    in raw
-- callOctave env p mts p' ts = let
--     script = matlabODE env (wholeProc env p mts) p' ts
--   in do
--     putStrLn $ "Octave Script: \n" ++ script
--     OS.readProcess
--       "octave" ["-q", "--eval", script] []

-- | Solver which calculates the symbolic Jacobian, writes MATLAB code, and executes it with GNU Octave. (General purpose, deals with stiff systems, uses LSODE.)
-- solveODEoctave :: Solver
-- solveODEoctave env p mts p' ts@(n,(t0,tn))
--     = let raw = unsafePerformIO (callOctave env p mts p' ts)
--       in (n>< Map.size p') $ map s2d $ words raw

-- Return the MATLAB script for ODEs
-- matlabScript :: Env
--              -> Process
--              -> MTS
--              -> P'
--              -> (Int, (Double, Double))
--              -> String
-- matlabScript env p mts = matlabODE env (wholeProc env p mts)