{-# LANGUAGE CPP #-}
#if __GLASGOW_HASKELL__ >= 708

{-# LANGUAGE DataKinds #-}
{-# LANGUAGE KindSignatures #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE FunctionalDependencies #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE Rank2Types #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE ViewPatterns #-}
{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE DeriveGeneric #-}

{-# OPTIONS_GHC -fno-warn-missing-signatures #-}
{-# OPTIONS_GHC -fno-warn-simplifiable-class-constraints #-}

{- |
Module      :  Internal.Static
Copyright   :  (c) Alberto Ruiz 2006-14
License     :  BSD3
Stability   :  provisional

-}

module Internal.Static where

import GHC.TypeLits
import qualified Numeric.LinearAlgebra as LA
import Numeric.LinearAlgebra hiding (konst,size,R,C)
import Internal.Vector as D hiding (R,C)
import Internal.ST
import Control.DeepSeq
import Data.Proxy(Proxy)
import Foreign.Storable(Storable)
import Text.Printf

import Data.Binary
import GHC.Generics (Generic)
import Data.Proxy (Proxy(..))

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

type ℝ = Float
type ℂ = Complex Float

newtype Dim (n :: Nat) t = Dim t
  deriving (Show, Generic)

instance (KnownNat n, Binary a) => Binary (Dim n a) where
  get = do
    k <- get
    let n = natVal (Proxy :: Proxy n)
    if n == k
      then Dim <$> get
      else fail ("Expected dimension " ++ (show n) ++ ", but found dimension " ++ (show k))

  put (Dim x) = do
    put (natVal (Proxy :: Proxy n))
    put x

lift1F
  :: (c t -> c t)
  -> Dim n (c t) -> Dim n (c t)
lift1F f (Dim v) = Dim (f v)

lift2F
  :: (c t -> c t -> c t)
  -> Dim n (c t) -> Dim n (c t) -> Dim n (c t)
lift2F f (Dim u) (Dim v) = Dim (f u v)

instance NFData t => NFData (Dim n t) where
    rnf (Dim (force -> !_)) = ()

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

newtype R n = R (Dim n (Vector ℝ))
  deriving (Num,Fractional,Floating,Generic,Binary)

newtype C n = C (Dim n (Vector ℂ))
  deriving (Num,Fractional,Floating,Generic)

newtype L m n = L (Dim m (Dim n (Matrix ℝ)))
  deriving (Generic, Binary)

newtype M m n = M (Dim m (Dim n (Matrix ℂ)))
  deriving (Generic)

mkR :: Vector ℝ -> R n
mkR = R . Dim

mkC :: Vector ℂ -> C n
mkC = C . Dim

mkL :: Matrix ℝ -> L m n
mkL x = L (Dim (Dim x))

mkM :: Matrix ℂ -> M m n
mkM x = M (Dim (Dim x))

instance NFData (R n) where
    rnf (R (force -> !_)) = ()

instance NFData (C n) where
    rnf (C (force -> !_)) = ()

instance NFData (L n m) where
    rnf (L (force -> !_)) = ()

instance NFData (M n m) where
    rnf (M (force -> !_)) = ()

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

type V n t = Dim n (Vector t)

ud :: Dim n (Vector t) -> Vector t
ud (Dim v) = v

mkV :: forall (n :: Nat) t . t -> Dim n t
mkV = Dim

vconcat :: forall n m t . (KnownNat n, KnownNat m, Numeric t)
    => V n t -> V m t -> V (n+m) t
(ud -> u) `vconcat` (ud -> v) = mkV (vjoin [u', v'])
  where
    du = fromIntegral . natVal $ (undefined :: Proxy n)
    dv = fromIntegral . natVal $ (undefined :: Proxy m)
    u' | du /= 1 && LA.size u == 1 = LA.konst (u D.@> 0) du
       | otherwise = u
    v' | dv /= 1 && LA.size v == 1 = LA.konst (v D.@> 0) dv
       | otherwise = v

gvec2 :: Storable t => t -> t -> V 2 t
gvec2 a b = mkV $ runSTVector $ do
    v <- newUndefinedVector 2
    writeVector v 0 a
    writeVector v 1 b
    return v

gvec3 :: Storable t => t -> t -> t -> V 3 t
gvec3 a b c = mkV $ runSTVector $ do
    v <- newUndefinedVector 3
    writeVector v 0 a
    writeVector v 1 b
    writeVector v 2 c
    return v

gvec4 :: Storable t => t -> t -> t -> t -> V 4 t
gvec4 a b c d = mkV $ runSTVector $ do
    v <- newUndefinedVector 4
    writeVector v 0 a
    writeVector v 1 b
    writeVector v 2 c
    writeVector v 3 d
    return v

gvect :: forall n t . (Show t, KnownNat n, Numeric t) => String -> [t] -> V n t
gvect st xs'
    | ok = mkV v
    | not (null rest) && null (tail rest) = abort (show xs')
    | not (null rest) = abort (init (show (xs++take 1 rest))++", ... ]")
    | otherwise = abort (show xs)
  where
    (xs,rest) = splitAt d xs'
    ok = LA.size v == d && null rest
    v = LA.fromList xs
    d = fromIntegral . natVal $ (undefined :: Proxy n)
    abort info = error $ st++" "++show d++" can't be created from elements "++info

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

type GM m n t = Dim m (Dim n (Matrix t))

gmat :: forall m n t . (Show t, KnownNat m, KnownNat n, Numeric t) => String -> [t] -> GM m n t
gmat st xs'
    | ok = Dim (Dim x)
    | not (null rest) && null (tail rest) = abort (show xs')
    | not (null rest) = abort (init (show (xs++take 1 rest))++", ... ]")
    | otherwise = abort (show xs)
  where
    (xs,rest) = splitAt (m'*n') xs'
    v = LA.fromList xs
    x = reshape n' v
    ok = null rest && ((n' == 0 && dim v == 0) || n'> 0 && (rem (LA.size v) n' == 0) && LA.size x == (m',n'))
    m' = fromIntegral . natVal $ (undefined :: Proxy m) :: Int
    n' = fromIntegral . natVal $ (undefined :: Proxy n) :: Int
    abort info = error $ st ++" "++show m' ++ " " ++ show n'++" can't be created from elements " ++ info

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

class Num t => Sized t s d | s -> t, s -> d
  where
    konst     ::  t   -> s
    unwrap    ::  s   -> d t
    fromList  :: [t]  -> s
    extract   ::  s   -> d t
    create    ::  d t -> Maybe s
    size      ::  s   -> IndexOf d

singleV v = LA.size v == 1
singleM m = rows m == 1 && cols m == 1

instance KnownNat n => Sized ℂ (C n) Vector
  where
    size _ = fromIntegral . natVal $ (undefined :: Proxy n)
    konst x = mkC (LA.scalar x)
    unwrap (C (Dim v)) = v
    fromList xs = C (gvect "C" xs)
    extract s@(unwrap -> v)
      | singleV v = LA.konst (v!0) (size s)
      | otherwise = v
    create v
        | LA.size v == size r = Just r
        | otherwise = Nothing
      where
        r = mkC v :: C n

instance KnownNat n => Sized ℝ (R n) Vector
  where
    size _ = fromIntegral . natVal $ (undefined :: Proxy n)
    konst x = mkR (LA.scalar x)
    unwrap (R (Dim v)) = v
    fromList xs = R (gvect "R" xs)
    extract s@(unwrap -> v)
      | singleV v = LA.konst (v!0) (size s)
      | otherwise = v
    create v
        | LA.size v == size r = Just r
        | otherwise = Nothing
      where
        r = mkR v :: R n

instance (KnownNat m, KnownNat n) => Sized ℝ (L m n) Matrix
  where
    size _ = ((fromIntegral . natVal) (undefined :: Proxy m)
             ,(fromIntegral . natVal) (undefined :: Proxy n))
    konst x = mkL (LA.scalar x)
    fromList xs = L (gmat "L" xs)
    unwrap (L (Dim (Dim m))) = m
    extract (isDiag -> Just (z,y,(m',n'))) = diagRect z y m' n'
    extract s@(unwrap -> a)
        | singleM a = LA.konst (a `atIndex` (0,0)) (size s)
        | otherwise = a
    create x
        | LA.size x == size r = Just r
        | otherwise = Nothing
      where
        r = mkL x :: L m n

instance (KnownNat m, KnownNat n) => Sized ℂ (M m n) Matrix
  where
    size _ = ((fromIntegral . natVal) (undefined :: Proxy m)
             ,(fromIntegral . natVal) (undefined :: Proxy n))
    konst x = mkM (LA.scalar x)
    fromList xs = M (gmat "M" xs)
    unwrap (M (Dim (Dim m))) = m
    extract (isDiagC -> Just (z,y,(m',n'))) = diagRect z y m' n'
    extract s@(unwrap -> a)
        | singleM a = LA.konst (a `atIndex` (0,0)) (size s)
        | otherwise = a
    create x
        | LA.size x == size r = Just r
        | otherwise = Nothing
      where
        r = mkM x :: M m n

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

instance (KnownNat n, KnownNat m) => Transposable (L m n) (L n m)
  where
    tr a@(isDiag -> Just _) = mkL (extract a)
    tr (extract -> a) = mkL (tr a)
    tr' = tr

instance (KnownNat n, KnownNat m) => Transposable (M m n) (M n m)
  where
    tr a@(isDiagC -> Just _) = mkM (extract a)
    tr (extract -> a) = mkM (tr a)
    tr' a@(isDiagC -> Just _) = mkM (extract a)
    tr' (extract -> a) = mkM (tr' a)

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

isDiag :: forall m n . (KnownNat m, KnownNat n) => L m n -> Maybe (ℝ, Vector ℝ, (Int,Int))
isDiag (L x) = isDiagg x

isDiagC :: forall m n . (KnownNat m, KnownNat n) => M m n -> Maybe (ℂ, Vector ℂ, (Int,Int))
isDiagC (M x) = isDiagg x

isDiagg :: forall m n t . (Numeric t, KnownNat m, KnownNat n) => GM m n t -> Maybe (t, Vector t, (Int,Int))
isDiagg (Dim (Dim x))
    | singleM x = Nothing
    | rows x == 1 && m' > 1 || cols x == 1 && n' > 1 = Just (z,yz,(m',n'))
    | otherwise = Nothing
  where
    m' = fromIntegral . natVal $ (undefined :: Proxy m) :: Int
    n' = fromIntegral . natVal $ (undefined :: Proxy n) :: Int
    v = flatten x
    z = v `atIndex` 0
    y = subVector 1 (LA.size v-1) v
    ny = LA.size y
    zeros = LA.konst 0 (max 0 (min m' n' - ny))
    yz = vjoin [y,zeros]

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

instance KnownNat n => Show (R n)
  where
    show s@(R (Dim v))
      | singleV v = "(" ++ show (v!0) ++ " :: R " ++ show d ++ ")"
      | otherwise = "(vector " ++ show v ++ " :: R " ++ show d ++")"
      where
        d = size s

instance KnownNat n => Show (C n)
  where
    show s@(C (Dim v))
      | singleV v = "(" ++ show (v!0) ++ " :: C " ++ show d ++ ")"
      | otherwise = "(vector " ++ show v ++ " :: C " ++ show d ++")"
      where
        d = size s

instance (KnownNat m, KnownNat n) => Show (L m n)
  where
    show (isDiag -> Just (z,y,(m',n'))) = printf "(diag %s %s :: L %d %d)" (show z) (show y) m' n'
    show s@(L (Dim (Dim x)))
       | singleM x = printf "(%s :: L %d %d)" (show (x `atIndex` (0,0))) m' n'
       | otherwise = "(matrix" ++ dropWhile (/='\n') (show x) ++ " :: L " ++ show m' ++ " " ++ show n' ++ ")"
      where
        (m',n') = size s

instance (KnownNat m, KnownNat n) => Show (M m n)
  where
    show (isDiagC -> Just (z,y,(m',n'))) = printf "(diag %s %s :: M %d %d)" (show z) (show y) m' n'
    show s@(M (Dim (Dim x)))
       | singleM x = printf "(%s :: M %d %d)" (show (x `atIndex` (0,0))) m' n'
       | otherwise = "(matrix" ++ dropWhile (/='\n') (show x) ++ " :: M " ++ show m' ++ " " ++ show n' ++ ")"
      where
        (m',n') = size s

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

instance (Num (Vector t), Numeric t )=> Num (Dim n (Vector t))
  where
    (+) = lift2F (+)
    (*) = lift2F (*)
    (-) = lift2F (-)
    abs = lift1F abs
    signum = lift1F signum
    negate = lift1F negate
    fromInteger x = Dim (fromInteger x)

instance (Num (Vector t), Num (Matrix t), Fractional t, Numeric t) => Fractional (Dim n (Vector t))
  where
    fromRational x = Dim (fromRational x)
    (/) = lift2F (/)

instance (Fractional t, Floating (Vector t), Numeric t) => Floating (Dim n (Vector t)) where
    sin   = lift1F sin
    cos   = lift1F cos
    tan   = lift1F tan
    asin  = lift1F asin
    acos  = lift1F acos
    atan  = lift1F atan
    sinh  = lift1F sinh
    cosh  = lift1F cosh
    tanh  = lift1F tanh
    asinh = lift1F asinh
    acosh = lift1F acosh
    atanh = lift1F atanh
    exp   = lift1F exp
    log   = lift1F log
    sqrt  = lift1F sqrt
    (**)  = lift2F (**)
    pi    = Dim pi

instance (Num (Matrix t), Numeric t) => Num (Dim m (Dim n (Matrix t)))
  where
    (+) = (lift2F . lift2F) (+)
    (*) = (lift2F . lift2F) (*)
    (-) = (lift2F . lift2F) (-)
    abs = (lift1F . lift1F) abs
    signum = (lift1F . lift1F) signum
    negate = (lift1F . lift1F) negate
    fromInteger x = Dim (Dim (fromInteger x))

instance (Num (Vector t), Num (Matrix t), Fractional t, Numeric t) => Fractional (Dim m (Dim n (Matrix t)))
  where
    fromRational x = Dim (Dim (fromRational x))
    (/) = (lift2F.lift2F) (/)

instance (Num (Vector t), Floating (Matrix t), Fractional t, Numeric t) => Floating (Dim m (Dim n (Matrix t))) where
    sin   = (lift1F . lift1F) sin
    cos   = (lift1F . lift1F) cos
    tan   = (lift1F . lift1F) tan
    asin  = (lift1F . lift1F) asin
    acos  = (lift1F . lift1F) acos
    atan  = (lift1F . lift1F) atan
    sinh  = (lift1F . lift1F) sinh
    cosh  = (lift1F . lift1F) cosh
    tanh  = (lift1F . lift1F) tanh
    asinh = (lift1F . lift1F) asinh
    acosh = (lift1F . lift1F) acosh
    atanh = (lift1F . lift1F) atanh
    exp   = (lift1F . lift1F) exp
    log   = (lift1F . lift1F) log
    sqrt  = (lift1F . lift1F) sqrt
    (**)  = (lift2F . lift2F) (**)
    pi    = Dim (Dim pi)

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

adaptDiag f a@(isDiag -> Just _) b | isFull b = f (mkL (extract a)) b
adaptDiag f a b@(isDiag -> Just _) | isFull a = f a (mkL (extract b))
adaptDiag f a b = f a b

isFull m = isDiag m == Nothing && not (singleM (unwrap m))

lift1L f (L v) = L (f v)
lift2L f (L a) (L b) = L (f a b)
lift2LD f = adaptDiag (lift2L f)

instance (KnownNat n, KnownNat m) =>  Num (L n m)
  where
    (+) = lift2LD (+)
    (*) = lift2LD (*)
    (-) = lift2LD (-)
    abs = lift1L abs
    signum = lift1L signum
    negate = lift1L negate
    fromInteger = L . Dim . Dim . fromInteger

instance (KnownNat n, KnownNat m) => Fractional (L n m)
  where
    fromRational = L . Dim . Dim . fromRational
    (/) = lift2LD (/)

instance (KnownNat n, KnownNat m) => Floating (L n m) where
    sin   = lift1L sin
    cos   = lift1L cos
    tan   = lift1L tan
    asin  = lift1L asin
    acos  = lift1L acos
    atan  = lift1L atan
    sinh  = lift1L sinh
    cosh  = lift1L cosh
    tanh  = lift1L tanh
    asinh = lift1L asinh
    acosh = lift1L acosh
    atanh = lift1L atanh
    exp   = lift1L exp
    log   = lift1L log
    sqrt  = lift1L sqrt
    (**)  = lift2LD (**)
    pi    = konst pi

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

adaptDiagC f a@(isDiagC -> Just _) b | isFullC b = f (mkM (extract a)) b
adaptDiagC f a b@(isDiagC -> Just _) | isFullC a = f a (mkM (extract b))
adaptDiagC f a b = f a b

isFullC m = isDiagC m == Nothing && not (singleM (unwrap m))

lift1M f (M v) = M (f v)
lift2M f (M a) (M b) = M (f a b)
lift2MD f = adaptDiagC (lift2M f)

instance (KnownNat n, KnownNat m) =>  Num (M n m)
  where
    (+) = lift2MD (+)
    (*) = lift2MD (*)
    (-) = lift2MD (-)
    abs = lift1M abs
    signum = lift1M signum
    negate = lift1M negate
    fromInteger = M . Dim . Dim . fromInteger

instance (KnownNat n, KnownNat m) => Fractional (M n m)
  where
    fromRational = M . Dim . Dim . fromRational
    (/) = lift2MD (/)

instance (KnownNat n, KnownNat m) => Floating (M n m) where
    sin   = lift1M sin
    cos   = lift1M cos
    tan   = lift1M tan
    asin  = lift1M asin
    acos  = lift1M acos
    atan  = lift1M atan
    sinh  = lift1M sinh
    cosh  = lift1M cosh
    tanh  = lift1M tanh
    asinh = lift1M asinh
    acosh = lift1M acosh
    atanh = lift1M atanh
    exp   = lift1M exp
    log   = lift1M log
    sqrt  = lift1M sqrt
    (**)  = lift2MD (**)
    pi    = M pi

instance Additive (R n) where
    add = (+)

instance Additive (C n) where
    add = (+)

instance (KnownNat m, KnownNat n) => Additive (L m n) where
    add = (+)

instance (KnownNat m, KnownNat n) => Additive (M m n) where
    add = (+)

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

class Disp t
  where
    disp :: Int -> t -> IO ()

instance (KnownNat m, KnownNat n) => Disp (L m n)
  where
    disp n x = do
        let a = extract x
        let su = LA.dispf n a
        printf "L %d %d" (rows a) (cols a) >> putStr (dropWhile (/='\n') $ su)

instance (KnownNat m, KnownNat n) => Disp (M m n)
  where
    disp n x = do
        let a = extract x
        let su = LA.dispcf n a
        printf "M %d %d" (rows a) (cols a) >> putStr (dropWhile (/='\n') $ su)

instance KnownNat n => Disp (R n)
  where
    disp n v = do
        let su = LA.dispf n (asRow $ extract v)
        putStr "R " >> putStr (tail . dropWhile (/='x') $ su)

instance KnownNat n => Disp (C n)
  where
    disp n v = do
        let su = LA.dispcf n (asRow $ extract v)
        putStr "C " >> putStr (tail . dropWhile (/='x') $ su)

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

overMatL' :: (KnownNat m, KnownNat n)
          => (LA.Matrix ℝ -> LA.Matrix ℝ) -> L m n -> L m n
overMatL' f = mkL . f . unwrap
{-# INLINE overMatL' #-}

overMatM' :: (KnownNat m, KnownNat n)
          => (LA.Matrix ℂ -> LA.Matrix ℂ) -> M m n -> M m n
overMatM' f = mkM . f . unwrap
{-# INLINE overMatM' #-}

#else

module Numeric.LinearAlgebra.Static.Internal where

#endif


{-# LANGUAGE DataKinds             #-}
{-# LANGUAGE FlexibleContexts      #-}
{-# LANGUAGE GADTs                 #-}
{-# LANGUAGE InstanceSigs          #-}
{-# LANGUAGE KindSignatures        #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE PolyKinds             #-}
{-# LANGUAGE RecordWildCards       #-}
{-# LANGUAGE ScopedTypeVariables   #-}
{-# LANGUAGE TypeFamilies          #-}
{-# LANGUAGE TypeOperators         #-}
{-# LANGUAGE UndecidableInstances  #-}

module Grenade.Recurrent.Layers.BasicRecurrent (
    BasicRecurrent (..)
  ) where

import           Data.Proxy
import           Data.Singletons.TypeLits hiding (natVal)

import           Numeric.LinearAlgebra.Static

import           GHC.TypeLits

import           Grenade.Core
import           Grenade.Recurrent.Core

data BasicRecurrent :: Nat -- Input layer size
                    -> Nat -- Output layer size
                    -> * where
  BasicRecurrent :: ( KnownNat input
                    , KnownNat output
                    , KnownNat matrixCols
                    , matrixCols ~ (input + output))
                 => !(R output)   -- Bias neuron weights
                 -> !(R output)   -- Bias neuron momentum
                 -> !(L output matrixCols) -- Activation
                 -> !(L output matrixCols) -- Momentum
                 -> BasicRecurrent input output

data BasicRecurrent' :: Nat -- Input layer size
                     -> Nat -- Output layer size
                     -> * where
  BasicRecurrent' :: ( KnownNat input
                     , KnownNat output
                     , KnownNat matrixCols
                     , matrixCols ~ (input + output))
                  => !(R output)   -- Bias neuron gradients
                  -> !(L output matrixCols)
                  -> BasicRecurrent' input output

instance Show (BasicRecurrent i o) where
  show BasicRecurrent {} = "BasicRecurrent"

instance (KnownNat i, KnownNat o, KnownNat (i + o)) => UpdateLayer (BasicRecurrent i o) where
  type Gradient (BasicRecurrent i o) = (BasicRecurrent' i o)

  runUpdate LearningParameters {..} (BasicRecurrent oldBias oldBiasMomentum oldActivations oldMomentum) (BasicRecurrent' biasGradient activationGradient) =
    let newBiasMomentum = konst learningMomentum * oldBiasMomentum - konst learningRate * biasGradient
        newBias         = oldBias + newBiasMomentum
        newMomentum     = konst learningMomentum * oldMomentum - konst learningRate * activationGradient
        regulariser     = konst (learningRegulariser * learningRate) * oldActivations
        newActivations  = oldActivations + newMomentum - regulariser
    in BasicRecurrent newBias newBiasMomentum newActivations newMomentum

instance (KnownNat i, KnownNat o, KnownNat x, KnownNat (x*o), x ~ (i+o)) => RandomLayer (BasicRecurrent i o) where
  createRandomWith m gen = do
    wB <- getRandomVector i o m gen
    wN <- getRandomMatrix i o m gen
    let bm = konst 0
        mm = konst 0
    return $ BasicRecurrent wB bm wN mm
      where i = natVal (Proxy :: Proxy i)
            o = natVal (Proxy :: Proxy o)

instance (KnownNat i, KnownNat o, KnownNat (i + o), i <= (i + o), o ~ ((i + o) - i)) => RecurrentUpdateLayer (BasicRecurrent i o) where
  type RecurrentShape (BasicRecurrent i o) = S ('D1 o)

instance (KnownNat i, KnownNat o, KnownNat (i + o), i <= (i + o), o ~ ((i + o) - i)) => RecurrentLayer (BasicRecurrent i o) ('D1 i) ('D1 o) where

  type RecTape (BasicRecurrent i o) ('D1 i) ('D1 o) = (S ('D1 o), S ('D1 i))
  -- Do a matrix vector multiplication and return the result.
  runRecurrentForwards (BasicRecurrent wB _ wN _) (S1D lastOutput) (S1D thisInput) =
    let thisOutput = S1D $ wB + wN #> (thisInput # lastOutput)
    in ((S1D lastOutput, S1D thisInput), thisOutput, thisOutput)

  -- Run a backpropogation step for a full connected layer.
  runRecurrentBackwards (BasicRecurrent _ _ wN _) (S1D lastOutput, S1D thisInput) (S1D dRec) (S1D dEdy) =
    let biasGradient        = (dRec + dEdy)
        layerGrad           = (dRec + dEdy) `outer` (thisInput # lastOutput)
        -- calcluate derivatives for next step
        (backGrad, recGrad) = split $ tr wN #> (dRec + dEdy)
    in  (BasicRecurrent' biasGradient layerGrad, S1D recGrad, S1D backGrad)

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

instance (KnownNat i,KnownNat o,KnownNat (i+o)) => GNum (BasicRecurrent i o) where
  n |* (BasicRecurrent wB mB a nM) = (BasicRecurrent (fromRational n * wB) mB a nM)
  (BasicRecurrent wB mB a nM) |+ (BasicRecurrent wB2 mB2 a2 nM2)  = (BasicRecurrent (0.5 * (wB + wB2)) (0.5* (mB+mB2)) (0.5 * (a+a2)) (0.5 * (nM + nM2)))
  gFromRational r = (BasicRecurrent (fromRational r) 0 (fromRational r) 0)


{-# language TypeSynonymInstances #-}
{-# language LiberalTypeSynonyms #-}

module FrequencyResponse
  (Analyzer (..),
  AnalyzerAcyclic (..),
  Spectrum (..),
  SpecLin (..),
  SpecWide ,
  getResponse,
  getResponseLin,
  getResponseAcyclic)

where

import CLaSH.Prelude (Default(..))
import Prelude
import CLaSH.Signal.MultiSignal
import Data.Complex as C
import Data.Maybe
import Control.Applicative
import Data.Map.Strict as Map

-- | Type representing input and output signal in circuitry
-- p of kind (* -> *) represents underlaying Signal of type Prependable (ZipList for example)
-- s is type of Spectrum
data Analyzer p s = Analyzer {unAnalyzer :: ([Maybe (s -> s)] , p s)}

instance (Functor p, Prependable p) => Prependable (Analyzer p) where
     prepend a da = Analyzer (Nothing : f , (prepend a $ fmap fx d))
      where
        (f,d) = unAnalyzer da
        fx = fromMaybe id $ foldr1 (<|>) f

instance (Applicative p,Num n) => Num (Analyzer p n) where
    (+) = liftA2FA (+)
    (-) = liftA2FA (-)
    (*) = liftA2FA (*)
    fromInteger = makeAnalyzer . fromInteger
    negate = liftA1FA negate
    abs = liftA1FA abs
    signum = liftA1FA signum

instance (Applicative p,Fractional n) => Fractional (Analyzer p n) where
    (/) = liftA2FA (/)
    fromRational = makeAnalyzer . fromRational

liftA2FA op da db = Analyzer (zipWith (<|>) fa fb, liftA2 op a b) where
       (fa,a) = unAnalyzer da
       (fb,b) = unAnalyzer db

liftA1FA op f = Analyzer (repeat Nothing , fmap op s)   where
   (_,s) = unAnalyzer f

makeAnalyzer s = Analyzer ( repeat Nothing , pure s )
{- inline liftA2FA, liftA1FA-}

-- | spectrum is result of frequency mixing two spectrums
-- n represents type represented as harmonics (Int or better, Rational)
-- and a is type of complex number (frequently Double)
data Spectrum n a = Spectrum {getSpectrum :: Map.Map n (Complex a)}

instance (Ord n,Num n,Default a,RealFloat a) => Default (Spectrum n a) where
    def = dcSpectrum def

instance (Ord n,Num n,RealFloat a) => Num (Spectrum n a) where
    (+) = dotMergeSpec (+)
    (-) = dotMergeSpec (-)
    (*) = mixSpectrum
    negate = fmapSpec negate
    abs = error "no abs function for spectrum available"
    signum = error "no sig num function for spectrum available"
    fromInteger = dcSpectrum . fromInteger

instance (Show n, Show a, Ord n,Num n,Fractional a,RealFloat a) => Fractional (Spectrum n a) where
    (/) = spectrumDiv
    fromRational = dcSpectrum . fromRational

dotMergeSpec op (Spectrum sa) (Spectrum sb) = Spectrum $ (Map.unionWith op) sa sb
fmapSpec f = Spectrum . fmap f . getSpectrum
constSpectrum freq a = Spectrum $ Map.singleton freq (a :+ 0 )
dcSpectrum a = constSpectrum 0 a

mixSpectrum (Spectrum sa) (Spectrum sb) = Spectrum
        $ Map.fromListWith (+)
        $ concat [mixFeq a b | a <- Map.toList sa , b <- Map.toList sb]

mixFeq (fa,pa) (fb,pb)
   | fa > fb = mixFeqSorted (fb,pb) (fa,pa)
   | True =  mixFeqSorted (fa,pa) (fb,pb)

mixFeqSorted (0,pa) (fb,pb) = [(fb,pa*pb)]
mixFeqSorted (fa,pa) (fb,pb) = [(fb - fa,r1),(fb+fa,r2)] where
   (ma,pha) = polar pa
   (mb,phb) = polar pb
   k = ma * mb / 2
   r1
      | fa == fb  = mkPolar (k * sin (phb - pha + pi/2)) 0
      | otherwise = mkPolar k (phb - pha + pi/2)
   r2 = mkPolar ((-1)*k) (phb + pha + pi/2)

delaySpec ph (Spectrum s) = Spectrum $ Map.mapWithKey f s where
    f 0 a = a
    f k a = a * mkPolar 1 (fromRational k * ph)

spectrumDiv (Spectrum a) sb = Spectrum (fmap (/ dc) a) where
    dc = getDcOrError sb

getDcOrError (Spectrum a) =
     case toList a of
            ((0,dc):[]) -> dc
            s  -> error ("spectrum is not dc, but" ++ show s )

makeDcSig :: (Applicative p, Num n, Num a) => a -> Analyzer p (Spectrum n a)
makeDcSig = makeAnalyzer . dcSpectrum

-- | spectrum for linear system, faster version of Spectrum
-- underlaying type is tuple insted of Map
newtype SpecLin = SpecLin {getSpecLin :: (Double,Complex Double)}

instance  Default SpecLin where
    def = SpecLin (def,def :+ def)

instance  Num SpecLin where
    (+) = dotMergeSpecLin (+)
    (-) = dotMergeSpecLin (-)
    (*) = mixSpecLin
    negate = fmapSpecLin negate
    abs = error "no abs function for spectrum available"
    signum = error "no signum function for spectrum available"
    fromInteger = dcSpecLin . fromInteger

instance Fractional SpecLin where
    (/) = undefined -- dcSpecLin
    fromRational = dcSpecLin . fromRational

dotMergeSpecLin op (SpecLin (dcA,acA)) (SpecLin (dcB,acB)) = SpecLin (op dcA dcB , acr) where
    acr = op (realPart acA) (realPart acB) :+ op (imagPart acA) (imagPart acB)

mixSpecLin (SpecLin (dcA,acA)) (SpecLin (dcB,acB))
   | (acA /= 0 ) && (acB /= 0) = error "can not mix two linear spectrum"
   | otherwise = SpecLin (dcr , acr) where
       dcr = dcA * dcB
       acr = fmap (* dcA) acB + fmap (* dcB) acA

dcSpecLin dc = SpecLin (dc,0)

fmapSpecLin f (SpecLin (dc,ac)) =  SpecLin (f dc,fmap f ac)

delayLinSpec ph (SpecLin (dc,ac)) = SpecLin (dc,ac * mkPolar 1 ph)
-- end of SpecLin

-- for acyclic circuits we can implement faster version of Prependable
-- It has nice property that first value of output is correct.
data AnalyzerAcyclic s = AnalyzerAcyclic {unAnalyzerAcyclic :: (Maybe (s -> s), s)}

-- Prependable instance has funny laws (no laws?)
-- it does not follow list isomorphism, because it is single value
instance Prependable AnalyzerAcyclic where
     prepend _ da = AnalyzerAcyclic (f , fx d)
      where
        (f,d) = unAnalyzerAcyclic da
        fx = fromMaybe id f

instance (Num n) => Num (AnalyzerAcyclic n) where
    (+) = opAcyclic (+)
    (-) = opAcyclic (-)
    (*) = opAcyclic (*)
    fromInteger = pureAcyclic . fromInteger
    negate = liftAcyclic negate
    abs = liftAcyclic abs
    signum = liftAcyclic signum

instance (Fractional n) => Fractional (AnalyzerAcyclic n) where
    (/) = opAcyclic (/)
    fromRational = pureAcyclic . fromRational

pureAcyclic n = AnalyzerAcyclic (Nothing,n)
opAcyclic op (AnalyzerAcyclic (fa,a)) (AnalyzerAcyclic (fb,b)) = (AnalyzerAcyclic (fa <|> fb ,op a b))
liftAcyclic f (AnalyzerAcyclic (fa,a)) = (AnalyzerAcyclic (fa,f a))

-- end of AnalyzerAcyclic

type SpecWide = Spectrum Rational Double

-- | run analysis for both acyclic\/cyclic and both linear\/nonlinear circuits
-- It is required to skip first N resulting values
-- where N represent delay of circuitry (approx total number of registers)
getResponse :: (Analyzer ZipList SpecWide -> Analyzer ZipList SpecWide) -- ^ circuitry
            -> Double  -- ^ phase in radians that represents Nyquist frequency. (phase pi ≡ 1.0 of sampling frequency)
            -> [Spectrum Rational Double] -- ^ response as stream
getResponse f ph = getZipList $ snd $ unAnalyzer $ f $ Analyzer ([ Just (delaySpec ph)] , pure (constSpectrum 1 1))

-- | run analysis for both acyclic\/cyclic and only linear circuits
-- faster version of getResponse for acyclic linear ciruits
getResponseLin :: (Analyzer ZipList SpecLin -> Analyzer ZipList SpecLin) -- ^ circuitry
          -> Double -- ^ phase in radians that represents Nyquist frequency. (phase pi ≡ 1.0 of sampling frequency)
          -> [SpecLin] -- ^ response as stream
getResponseLin f ph = getZipList $ snd $ unAnalyzer $ f $ Analyzer ([Just (delayLinSpec ph)] , pure (SpecLin (0,1 :+ 0)))

-- | run analysis both linear\/nonlinear and only acyclic circuit
-- calculate response of acyclic circuit
getResponseAcyclic :: (AnalyzerAcyclic SpecWide -> AnalyzerAcyclic SpecWide) -- ^ circuitry
                      -> Double -- ^ phase in radians that represents Nyquist frequency. (phase pi ≡ 1.0 of sampling frequency)
                      -> Spectrum Rational Double -- ^ response as stream
getResponseAcyclic f ph =  snd $ unAnalyzerAcyclic $ f $ AnalyzerAcyclic (Just (delaySpec ph) ,constSpectrum 1 1)


{-# LANGUAGE DataKinds             #-}
{-# LANGUAGE GADTs                 #-}
{-# LANGUAGE TypeOperators         #-}
{-# LANGUAGE TypeFamilies          #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE FlexibleContexts      #-}
{-# LANGUAGE RankNTypes            #-}
{-# LANGUAGE FlexibleInstances     #-}
{-# LANGUAGE ScopedTypeVariables   #-}

module GrenadeExtras.Zip (
  Zip (..),
) where

import           Data.Serialize (Serialize, put, get)

--import           Data.Singletons
import           GHC.TypeLits

import           Grenade.Core (UpdateLayer, Layer, Shape(D1), S(S1D), Gradient,
                               runForwards, runBackwards, runUpdate, createRandom, Tape)
import           GrenadeExtras.GradNorm (GradNorm, normSquared)

import           Numeric.LinearAlgebra.Static ((#), split, R) -- row, (===), splitRows, unrow

data Zip :: Shape -> Shape -> * -> Shape -> Shape -> * -> * where
  Zip :: x -> y -> Zip ix ox x iy oy y

instance (Show x, Show y) => Show (Zip ix ox x iy oy y) where
  show (Zip x y) = "Zip\n" ++ show x ++ "\n" ++ show y

instance (UpdateLayer x, UpdateLayer y) => UpdateLayer (Zip ix ox x iy oy y) where
  type Gradient (Zip ix ox x iy oy y) = (Gradient x, Gradient y)
  runUpdate lr (Zip x y) (x', y') = Zip (runUpdate lr x x') (runUpdate lr y y')
  createRandom = Zip <$> createRandom <*> createRandom

instance ( Layer x ('D1 p) ('D1 m)
         , Layer y ('D1 q) ('D1 n)
         , KnownNat r
         , KnownNat p
         , KnownNat q
         , r ~ (p + q)
         , q ~ (r - p)
         , (p <=? r) ~ 'True
         , KnownNat o
         , KnownNat m
         , KnownNat n
         , o ~ (m + n)
         , n ~ (o - m)
         , (m <=? o) ~ 'True
         ) => Layer (Zip ('D1 p) ('D1 m) x ('D1 q) ('D1 n) y) ('D1 r) ('D1 o) where
   type Tape (Zip ('D1 p) ('D1 m) x ('D1 q) ('D1 n) y) ('D1 r) ('D1 o) = (Tape x ('D1 p) ('D1 m), Tape y ('D1 q) ('D1 n))

   runForwards (Zip x y) (S1D input) =
     let (inputX :: R p, inputY :: R q) = split input
         (xT, xOut :: S ('D1 m)) = runForwards x (S1D inputX)
         (yT, yOut :: S ('D1 n)) = runForwards y (S1D inputY)
     in case (xOut, yOut) of
         (S1D xOut', S1D yOut') ->
             ((xT, yT), S1D (xOut' # yOut'))

   runBackwards (Zip x y) (xTape, yTape) (S1D o) =
     let (ox :: R m , oy :: R n) = split o
         (x', xB :: S ('D1 p)) = runBackwards x xTape (S1D ox)
         (y', yB :: S ('D1 q)) = runBackwards y yTape (S1D oy)
      in case (xB, yB) of
           (S1D xB', S1D yB') ->
             ((x', y'), S1D (xB' # yB'))

instance (Serialize a, Serialize b) => Serialize (Zip sia sa a sib sb b) where
  put (Zip a b) = put a *> put b
  get = Zip <$> get <*> get

instance (GradNorm a, GradNorm b) => GradNorm (Zip sia sa a sib sb b) where
  normSquared (Zip a b) = normSquared a + normSquared b

instance (Num a, Num b) => Num (Zip sia sa a sib sb b) where
  (Zip w1 m1) + (Zip w2 m2) = Zip (w1+w2) (m1+m2)
  (Zip w1 m1) - (Zip w2 m2) = Zip (w1-w2) (m1-m2)
  (Zip w1 m1) * (Zip w2 m2) = Zip (w1*w2) (m1*m2)
  abs    (Zip w m) = Zip (abs w) (abs m)
  signum (Zip w m) = Zip (signum w) (signum m)
  negate (Zip w m) = Zip (negate w) (negate m)
  fromInteger i = Zip (fromInteger i) (fromInteger i)


-- |
-- Module      : AOC.Challenge.Day11
-- Copyright   : (c) Justin Le 2018
-- License     : BSD3
--
-- Maintainer  : justin@jle.im
-- Stability   : experimental
-- Portability : non-portable
--
-- Day 11.  See "AOC.Solver" for the types used in this module!

module AOC.Challenge.Day11 (
    day11a
  , day11b
  ) where

import           AOC.Common                       (Point, meanVar)
import           AOC.Solver                       ((:~>)(..))
import           Control.DeepSeq                  (force)
import           Data.Foldable                    (maximumBy)
import           Data.Ix                          (range)
import           Data.Map                         (Map)
import           Data.Maybe                       (catMaybes)
import           Data.Ord                         (comparing)
import           Data.Profunctor                  (dimap)
import           Linear                           (V2(..))
import           Text.Read                        (readMaybe)
import qualified Control.Foldl                    as F
import qualified Data.Map                         as M
import qualified Data.Set                         as S
import qualified Statistics.Distribution          as D
import qualified Statistics.Distribution.Binomial as D
import qualified Statistics.Distribution.Normal   as D

powerLevel :: Int -> Point -> Int
powerLevel sid (V2 x y) = hun ((rid * y + sid) * rid) - 5
  where
    hun = (`mod` 10) . (`div` 100)
    rid = x + 10

findMaxThree :: Map Point Int -> Point
findMaxThree mp = fst
                . maximumBy (comparing snd)
                . map (\x -> (x, go x))
                $ range (V2 1 1, V2 298 298)
  where
    go p = sum [ mp M.! (p + shift)
               | shift <- range (V2 0 0, V2 2 2)
               ]

mkMap :: Int -> Map Point Int
mkMap i = M.fromSet (powerLevel i) . S.fromList $ range (V2 1 1, V2 300 300)

day11a :: Int :~> Point
day11a = MkSol
    { sParse = readMaybe
    , sShow  = \(V2 x y) -> show x ++ "," ++ show y
    , sSolve = Just . findMaxThree . mkMap
    }

day11b :: Int :~> (Point, Int)
day11b = MkSol
    { sParse = readMaybe
    , sShow  = \(V2 x y, s) -> show x ++ "," ++ show y ++ "," ++ show s
    , sSolve = Just . findMaxAny . mkMap
    }

findMaxAny :: Map Point Int -> (Point, Int)
findMaxAny mp = fst $ go 1
  where
    go n
       | goOn > 0.001 = maximumBy (comparing snd) [((pMax, n), oMax), go (n + 1)]
       | otherwise    = ((pMax, n), oMax)
                      -- & traceShow (n, oMax, goOn)
      where
        (pMax, oMax) = maximumBy (comparing snd)
            [ (p, fromIntegral (fromSAT sat p n))
            | p <- range (V2 1 1, V2 (300 - n + 1) (300 - n + 1))
            ]
        goOn = sum
            [ probGreaterThan oMax n'
            | n' <- [n + 1 .. 300]
            ]
    !sat = summedAreaTable mp
    σ :: Double
    σ = sqrt $ ((4 + 5)**2)/12   -- stdev of uniform distribution between -5 and 4
    probGreaterThan o n
        | prob2 == 0 = 0
        | otherwise  = probIn2
      where
        n' = fromIntegral n
        distr2  = D.normalDistr (-(n' ** 2) * 0.5) (n' * σ)
        prob2   = D.complCumulative distr2 o
        numIn2  = ( (300 `div` n) ^ (2 :: Int)    -- we compensate for dependence
                  + (300 - n + 1)^(2 :: Int)
                  ) `div` 2
        probIn2 = 1 - D.probability (D.binomial numIn2 prob2) 0
                    -- this is technically not a bernoulli process,
                    -- because our items are dependent. but we fudge this
                    -- by tweaking the number of trials

fromSAT :: Map Point Int -> Point -> Int -> Int
fromSAT sat (subtract (V2 1 1)->p) n = sum . catMaybes $
    [            M.lookup p            sat
    ,            M.lookup (p + V2 n n) sat
    , negate <$> M.lookup (p + V2 0 n) sat
    , negate <$> M.lookup (p + V2 n 0) sat
    ]

summedAreaTable :: Map Point Int -> Map Point Int
summedAreaTable mp = force sat
  where
    sat = M.mapWithKey go mp
    go p0 v = (+ v) . sum . catMaybes $
      [ negate <$> M.lookup (p0 - V2 1 1) sat
      ,            M.lookup (p0 - V2 1 0) sat
      ,            M.lookup (p0 - V2 0 1) sat
      ]

-- | Debug: find the variance of the map at every square size
_chunkyVars :: Map Point Int -> Map Int Double
_chunkyVars mp = flip M.fromSet (S.fromList [1..300]) $ \n ->
    F.fold (dimap fromIntegral snd meanVar)
      [ fromSAT sat p n
      | p <- range (V2 1 1, V2 (300 - n + 1) (300 - n + 1))
      ]
  where
    !sat = summedAreaTable mp


{-# OPTIONS_GHC -fglasgow-exts #-}
{-# LANGUAGE UndecidableInstances #-}
{- OPTIONS_GHC -fwarn-incomplete-patterns -}

module GeneralizedSignals where

--import
import qualified Data.StorableVector as SV
import qualified Data.StorableVector.Base as SVB
import qualified Data.StorableVector.ST.Strict as SVST
import Foreign.Storable
--import EvalM
import Foreign.Storable.Tuple
import Data.List
import Control.Monad.ST
import Control.Monad
import TNUtils
import VectorsL
import Nats
--import PlotGnuplot
import Data.Ix
import Data.Complex
import Numeric.FFT
import Math.FFT
import qualified Data.Array.CArray as CA
import System.IO.Unsafe

newtype Time = Time Double
    deriving (Eq, Show, Num, Fractional, Floating, Discretizable, IsDouble, Storable)
newtype Freq = Freq Double
    deriving (Eq, Show, Num, Fractional, Floating, Discretizable, IsDouble, Storable)
newtype Length = Length Double
    deriving (Eq, Show, Num, Fractional, Floating, Discretizable, IsDouble, Storable)
newtype SpaceFreq = SpaceFreq Double
    deriving (Eq, Show, Num, Fractional, Floating, Discretizable, IsDouble, Storable)

class (Eq a,Show a,Num a,Fractional a,Floating a)=> IsDouble a where
    toDouble :: a -> Double
    fromDouble :: Double -> a

instance IsDouble Double where
    toDouble = id
    fromDouble = id

class Discretizable a where
    discreteIndex ::  (a,a) -> a -> a -> Int
    discreteRange :: (a,a) -> a -> [a]
    discreteRngLength :: (a,a) -> a -> Int
    foldRangeM :: Monad m => (a,a) -> a -> b -> (a -> b -> m b) -> m b

instance Discretizable Double where
    discreteIndex (lo,hi) dx x = round $ (x-lo)/dx
    discreteRange (lo,hi) dx = [lo, lo+dx..hi]
    discreteRngLength (lo,hi) dx = round $ (hi-lo)/dx
    foldRangeM (lo,hi) dx init f =
        let nmax = discreteRngLength (lo,hi) dx
            go n acc | n > nmax = return acc
                     | otherwise = do
              next <- f ((realToFrac n)*dx+lo) acc
              go (n+1) next
        in go 0 init

instance Discretizable Int where
    discreteIndex (lo,hi) 1 x = x-lo
    discreteIndex (lo,hi) dx x = (x-lo) `div` dx
    discreteRange (lo,hi) dx = [lo,lo+dx..hi]
    discreteRngLength (lo,hi) dx = (hi-lo) `div` dx
    foldRangeM (lo,hi) dx init f =
        let go n acc | n > hi = return acc
                     | otherwise = do
              next <- f n acc
              go (n+dx) next
        in go lo init

{-instance Discretizable a => Discretizable (Vec Z a) where
    discreteIndex (lo,hi) dx x = 0
    discreteRange (lo,hi) dx = [vnil]
    discreteRngLength (lo,hi) dx = 1
    foldRangeM (lo,hi) dx init f = return init -}

instance Discretizable a => Discretizable (Vec (S Z) a) where
    discreteIndex (lo,hi) dx x = discreteIndex (vcar lo,vcar hi) (vcar dx) (vcar x)
    discreteRange (lo,hi) dx = map (`vcons` vnil) $ discreteRange (vcar lo,vcar hi) (vcar dx)
    discreteRngLength (lo,hi) dx = discreteRngLength (vcar lo,vcar hi) (vcar dx)
    foldRangeM (lo,hi) dx init f = foldRangeM (vcar lo,vcar hi) (vcar dx) init f'
        where f' x acc = f (x `vcons` vnil) acc

instance (Discretizable a, Discretizable (Vec (S n) a)) => Discretizable (Vec (S (S n)) a) where
    discreteIndex (vx, vy) (vd) (vw)
        = discreteIndex (vcar vx,vcar vy) (vcar vd) (vcar vw) +
         (discreteIndex (vcar vx,vcar vy) (vcar vd) (vcar vy) + 1) *
         discreteIndex (vcdr vx,vcdr vy) (vcdr vd) (vcdr vw)
    discreteRange (vx, vy) vd
        = [ vcons xys vxsys | vxsys <- discreteRange (vcdr vx, vcdr vy) (vcdr vd),
                              xys <- discreteRange (vcar vx,vcar vy) (vcar vd)]
    discreteRngLength (vx, vy) (vd) = discreteRngLength (vcar vx,vcar vy) (vcar vd) *
                                      discreteRngLength (vcdr vx,vcdr vy) (vcdr vd)
--    foldRangeM :: Monad m => (a,a) -> a -> b -> (a -> b -> m b) -> m b
    foldRangeM (vlo, vhi) vd  init f =
      let f' x acc =
             foldRangeM (vcdr vlo, vcdr vhi) (vcdr vd) acc $ \ix acc' -> f (x `vcons` ix) acc'
      in foldRangeM (vcar vlo, vcar vhi) (vcar vd) init f'

f v (a, i) = return $ (v:a, i+1)

tf = print =<< foldRangeM (0::Vec Three Int, 2) (1) ([],0) f

tf2 = discreteIndex (0::Vec Three Int, 2) (1) (2 `vcons` 2 `vcons` 2 `vcons` vnil)

{-$ \x-> do
         let thisf ix acc = f (x `vcons` ix) acc
         foldRangeM (vcdr vlo, vcdr vhi) (vcdr vd) init thisf -}

type family Volume a :: *
type instance Volume Time = (Time,Time)
type instance Volume Freq = (Freq,Freq)

--type family FInverses a b :: * -> *

--type instance FInverses Time Freq

--type family FInv a :: *

--type instance FInv Time = Freq
--type instance (FInv t ~ s) => FInv s = t

{-class Inverses a b | a -> b, b-> a where
    invertFwd :: a -> b
    invertBack :: b -> a

instance Inverses Time Freq where
    invertFwd (Time t) = Freq $ recip t
    invertBack (Freq f) = Time $ recip f

instance Inverses Length SpaceFreq where
    invertFwd (Length t) = SpaceFreq $ recip t
    invertBack (SpaceFreq f) = Length $ recip f-}

class HasInverse a where
      type Inv a :: *
      invertValue :: a -> Inv a

instance HasInverse Int where
      type Inv Int = Int
      invertValue = negate -- haha. FIXME

instance HasInverse Time where
      type Inv Time = Freq
      invertValue (Time t) = Freq $ recip t

instance HasInverse Freq where
      type Inv Freq = Time
      invertValue (Freq t) = Time $ recip t

instance HasInverse Length where
      type Inv Length = SpaceFreq
      invertValue (Length t) = SpaceFreq $ recip t

instance HasInverse SpaceFreq where
      type Inv SpaceFreq = Length
      invertValue (SpaceFreq t) = Length $ recip t

instance HasInverse a => HasInverse (Vec n a) where
      type Inv (Vec n a) = Vec n (Inv a)
      invertValue = fmap invertValue

data Signal a b where
    Signal  :: (Discretizable a, {-Ix (Discretized a),-} Storable b, Storable a) =>
               a -> (a,a) -> (Inv a, Inv a) -> SV.Vector b -> Signal a b
    SigFmap :: (b->b') -> Signal a b -> Signal a b'
    SigFun  :: (a->b) -> Signal a b

sigLims :: Signal a b -> (a,a)
sigLims (Signal d1 lims1 invlims1 arr1) = lims1
sigLims (SigFmap _ s) = sigLims s

instance Functor (Signal a) where
    fmap f (SigFmap g s) = SigFmap (f . g) s
    fmap f s = SigFmap f s

sOfVToVofS :: Signal a (Vec n b) -> Vec n (Signal a b)
sOfVToVofS s = undefined

vOfSToSofV :: Vec n (Signal a b) -> Signal a (Vec n b)
vOfSToSofV voss = let sigs = vecToList voss
                  in undefined

type Events a b = [(a,b)]
type Region a b = [(Volume a,b)]

at :: Signal a b -> a -> b
at (Signal delta lims _ arr) x =  arr `SV.index` (discreteIndex lims delta x)
at (SigFmap f s) x = f $ s `at` x
at (SigFun f) x = f x

fill :: (Discretizable a, Storable b, Storable a) =>
        a -> (a,a) -> (Inv a, Inv a) -> (a->b) -> Signal a b
fill delta lims invlims f =
     Signal delta lims invlims $ SV.pack $ map f $ discreteRange lims delta

trans :: Storable c => (a->c) -> Signal a b -> Signal a c
trans f (Signal delta lims ilims arr) = Signal delta lims ilims newarr where
    newarr = SV.pack $ map f $ discreteRange lims delta

trans' :: Storable c => (a->b->c) -> Signal a b -> Signal a c
trans' f (Signal delta lims ilims arr) = Signal delta lims ilims newarr where
    newarr = SV.pack $ map (uncurry f) $ zip (discreteRange lims delta) (SV.unpack arr)
--    newarr = SV.zipWith f (SV.pack $ discreteRange lims delta) (arr)

--DOESNT WORK -- problem may be in foldRangeM implementation
fillIO' :: (Discretizable a, Storable b, Show a, Storable a) =>
          a -> (a,a) -> (Inv a, Inv a) -> (a->IO b) -> IO (Signal a b)
fillIO' delta lims invlims mf = do
  let n = discreteRngLength lims delta
  print n
  stv <- stToIO $ SVST.new_ (n+1800) --FUDGE
  let f ixv nix = do v <- mf ixv
                     --print (ixv, nix)
                     stToIO $ SVST.write stv nix v
                     return $ nix+1
  _ <- foldRangeM lims delta 0 f
  sv <- stToIO $ SVST.unsafeFreeze stv
  return $ Signal delta lims invlims sv

fillIO :: (Discretizable a, Storable b, Show a, Storable a) =>
          a -> (a,a) -> (Inv a, Inv a) -> (a->IO b) -> IO (Signal a b)
fillIO delta lims invlims mf = do
  let rng = discreteRange lims delta
  let n = discreteRngLength lims delta
  xs <- mapM mf rng
  return $ Signal delta lims invlims $ SV.pack xs

zipSignalsWith :: Storable c => (a -> b -> c) -> Signal i a -> Signal i b -> Signal i c
zipSignalsWith f (Signal d1 lims1 invlims1 arr1) (Signal d2 lims2 invlims2 arr2)
               = Signal d1 lims1 invlims2 $ SV.zipWith f arr1 arr2

--toCArray :: Signal a b -> CArray

class X a b where
    fourier :: Signal a b -> Signal (Inv a) b

instance X Time (Complex Double) where
    fourier (Signal d lims invlims arr) = Signal (invertValue d) invlims lims $ SV.pack $ fft $ SV.unpack arr

instance X Freq (Complex Double) where
    fourier (Signal d lims invlims arr) = Signal (invertValue d) invlims lims $ SV.pack $ ifft $ SV.unpack arr

instance X Length (Complex Double) where
    fourier (Signal d lims invlims arr) = Signal (invertValue d) invlims lims $ SV.pack $ fft $ SV.unpack arr

instance X SpaceFreq (Complex Double) where
    fourier (Signal d lims invlims arr) = Signal (invertValue d) invlims lims $ SV.pack $ ifft $ SV.unpack arr

--instance X (Vec Z Length) (Complex Double) where
--    fourier = undefined

instance (Nat n, X (Vec n Length) (Complex Double), Storable (Vec n SpaceFreq),
          Discretizable (Vec n SpaceFreq), Ix (Vec n Int), CA.Shapable (Vec n Int))
    => X (Vec n Length) (Complex Double) where
    fourier s@(Signal d lims invlims arr) =
        let ixdims :: Nat n => Signal (Vec n a) b -> n
            ixdims = undefined
            dims = toInt $ ixdims s
            (oldp, oldn, oldsomething) = SVB.toForeignPtr arr
            cixs = discreteLims s
            carray = unsafePerformIO (CA.unsafeForeignPtrToCArray oldp cixs)
            newcarray = dftN [0..dims] carray
            (newn, newp) = CA.toForeignPtr newcarray
            newarr = SVB.fromForeignPtr newp newn
        in Signal (invertValue d) invlims lims newarr

--instance IsDouble d => Shapable (Vec Z d) where

discreteLims :: (Nat n, IsDouble d) => Signal (Vec n d) a -> (Vec n Int, Vec n Int)
discreteLims s@(Signal d (lo,hi) invlims arr)
    = ( fmap (round . toDouble) $ lo/d,  fmap (round . toDouble) $ hi/d)

instance (Storable (b,c), X a b, X a c) => X a (b,c) where
    fourier s = zipSignalsWith (,) (fourier $ fmap fst s) (fourier $ fmap snd s)

instance (X a b) => X a (Vec n b) where
    fourier sOfVecs =  vOfSToSofV $ fmap fourier $ sOfVToVofS sOfVecs

class XRC a where
    fourierRC :: Signal a Double -> Signal (Inv a) (Complex Double)
    fourierCR :: Signal (Inv a) (Complex Double) -> Signal a Double


{-# LANGUAGE DataKinds             #-}
{-# LANGUAGE DeriveAnyClass        #-}
{-# LANGUAGE DeriveGeneric         #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE RankNTypes            #-}
{-# LANGUAGE ScopedTypeVariables   #-}
{-# LANGUAGE TypeFamilies          #-}
{-# LANGUAGE TypeOperators         #-}
{-|
Module      : Grenade.Layers.Logit
Description : Exponential linear unit layer
Copyright   : (c) Huw Campbell, 2016-2017
License     : BSD2
Stability   : experimental
-}
module Grenade.Layers.Elu
  ( Elu(..)
  , SpecElu (..)
  , specElu1D
  , specElu2D
  , specElu3D
  , elu
  ) where

import           Control.DeepSeq                (NFData)
import           Data.Constraint                (Dict (..))
import           Data.Reflection                (reifyNat)
import           Data.Serialize
import           Data.Singletons
import           GHC.Generics                   (Generic)
import           GHC.TypeLits
import qualified Numeric.LinearAlgebra.Static   as LAS
import           Unsafe.Coerce                  (unsafeCoerce)

import           Grenade.Core
import           Grenade.Dynamic
import           Grenade.Dynamic.Internal.Build

-- | An exponential linear unit.
--   A layer which can act between any shape of the same dimension, acting as a
--   diode on every neuron individually.
data Elu = Elu
  deriving (Generic, NFData, Show)

instance UpdateLayer Elu where
  type Gradient Elu = ()
  runUpdate _ _ _ = Elu

instance RandomLayer Elu where
  createRandomWith _ _ = return Elu

instance Serialize Elu where
  put _ = return ()
  get = return Elu

instance ( KnownNat i) => Layer Elu ('D1 i) ('D1 i) where
  type Tape Elu ('D1 i) ('D1 i) = LAS.R i

  runForwards _ (S1D y) = (y, S1D (elu y))
    where
      elu = LAS.dvmap (\a -> if a <= 0 then exp a - 1 else a)
  runBackwards _ y (S1D dEdy) = ((), S1D (elu' y * dEdy))
    where
      elu' = LAS.dvmap (\a -> if a <= 0 then exp a else 1)

instance (KnownNat i, KnownNat j) => Layer Elu ('D2 i j) ('D2 i j) where
  type Tape Elu ('D2 i j) ('D2 i j) = S ('D2 i j)

  runForwards _ (S2D y) = (S2D y, S2D (elu y))
    where
      elu = LAS.dmmap (\a -> if a <= 0 then exp a - 1 else a)
  runBackwards _ (S2D y) (S2D dEdy) = ((), S2D (elu' y * dEdy))
    where
      elu' = LAS.dmmap (\a -> if a <= 0 then exp a else 1)

instance (KnownNat i, KnownNat j, KnownNat k) => Layer Elu ('D3 i j k) ('D3 i j k) where

  type Tape Elu ('D3 i j k) ('D3 i j k) = S ('D3 i j k)

  runForwards _ (S3D y) = (S3D y, S3D (elu y))
    where
      elu = LAS.dmmap (\a -> if a <= 0 then exp a - 1 else a)
  runBackwards _ (S3D y) (S3D dEdy) = ((), S3D (elu' y * dEdy))
    where
      elu' = LAS.dmmap (\a -> if a <= 0 then exp a else 1)

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

instance FromDynamicLayer Elu where
  fromDynamicLayer inp _ _ = SpecNetLayer $ SpecElu (tripleFromSomeShape inp)

instance ToDynamicLayer SpecElu where
  toDynamicLayer _ _ (SpecElu (rows, cols, depth)) =
     reifyNat rows $ \(_ :: (KnownNat rows) => Proxy rows) ->
     reifyNat cols $ \(_ :: (KnownNat cols) => Proxy cols) ->
     reifyNat depth $ \(_ :: (KnownNat depth) => Proxy depth) ->
     case (rows, cols, depth) of
         (_, 1, 1)    -> return $ SpecLayer Elu (sing :: Sing ('D1 rows)) (sing :: Sing ('D1 rows))
         (_, _, 1) -> return $ SpecLayer Elu (sing :: Sing ('D2 rows cols)) (sing :: Sing ('D2 rows cols))
         _    -> case (unsafeCoerce (Dict :: Dict()) :: Dict (KnownNat (rows GHC.TypeLits.* depth))) of
           Dict -> return $ SpecLayer Elu (sing :: Sing ('D3 rows cols depth)) (sing :: Sing ('D3 rows cols depth))

-- | Create a specification for a elu layer.
specElu1D :: Integer -> SpecNet
specElu1D i = specElu3D (i, 1, 1)

-- | Create a specification for a elu layer.
specElu2D :: (Integer, Integer) -> SpecNet
specElu2D (i,j) = specElu3D (i,j,1)

-- | Create a specification for a elu layer.
specElu3D :: (Integer, Integer, Integer) -> SpecNet
specElu3D = SpecNetLayer . SpecElu

-- | Add a Elu layer to your build.
elu :: BuildM ()
elu = buildGetLastLayerOut >>= buildAddSpec . SpecNetLayer . SpecElu

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

instance GNum Elu where
  _ |* Elu = Elu
  _ |+ Elu = Elu


{-# LANGUAGE BangPatterns  #-}
{-# LANGUAGE UnboxedTuples #-}
{-----------------------------------------------------------------------------------------

          Computing the Minimum Message Length function
                   Rawle Prince, February 2012

             calculating the mixture models from the MML2ds functions
------------------------------------------------------------------------------------------}
module RprMix2 (alph, sigD, sigE, tau, r_n, aBar,sim_eq,
               mmlMixture, initModel,  initModelR , --- initModelGD,
               oMod, oModMl, negGradients, negGradientsM, negGradientsOM,
               updateModel, tDupdate,  validOM, valid,nextModelR',dataError,
               vmean,vvar,vstDiv,vvarM,vunStdNorm,printElm,vstNorm,invS,invNeg,
               -------- datatpes and modules ----
               Predicted, Data,  RNs, OModel (Omod),Mmodel (MkModel),LPars,
               module System.Random
              ) where
--
import Control.Concurrent (forkIO)
import ListStats (normal, mtone3,cdf)
import LINreg2 (foldl'Rnf)
import LMisc (myTake,printElm,printElmChr,similarBy,findCuts, ascending,manyInvNS,sumBy)
--import Prelude
import System.IO.Unsafe
import Numeric.LinearAlgebra
import Numeric.GSL.Statistics (mean,variance,variance_m,stddev)
import Control.Parallel (pseq)
import Numeric
import System.Random
import Foreign.Storable
import Control.DeepSeq
import Data.Time
import Data.List
import Data.Maybe (isJust, fromJust, listToMaybe)
import Data.Ord (comparing)
import qualified Data.ByteString.Lazy.Char8 as L
import System.Directory (doesFileExist)
import IO
-----------------------------------------------------------------------------------------------
instance (Storable a, NFData a) => NFData (Vector a) where
    rnf va = rnf (toList va)

instance NFData (Mmodel) where
    rnf (MkModel a b c vs d) = rnf (a, b, c ,vs, d)

instance NFData (OModel) where
    rnf (Omod m a) = rnf (m , a)

-----------------------------------------------------------------------------------------------
--- some statistics using vectors instead of lists
----------------------------------------------------------------------------------------
--- mean, variance and correlations
----------------------------------------------------------------------------------------
vmean  ::  Vector Double -> Double
vmean  = mean
---
vvar  :: Vector Double -> Double
vvar    = variance

-- variance of a list of values with the mean given
vvarM  :: Double -> Vector Double -> Double
vvarM    = variance_m

--- the standard deviation without the Mean given
vstDiv   = stddev

------ standardise by subtracting the mean and dividing by the standard deviation
vstNorm       xs            = mapVector (\ a -> (a - m) / std) xs
                 where m    = vmean xs
                       std  = sqrt $ vvarM m xs

-- unstandardising
vunStdNorm m std vs  = mapVector (\a -> std * a + m)  vs

-- building the mixture model
-----------------------------------------------------------------------------------------------

--- the model: alphs, sigE, sigD,  linear parms
data Mmodel = MkModel !Double !Double !Double (Vector Double) !Double deriving (Eq)
-- the data and the error class cannot have the same mean, so we need to learn that as well

-- if two models are close to each other
sim_eq  (MkModel a b c _ z) (MkModel p q r _ z1) =
        near a p && near b q &&  near c r  && near z z1 -- && etc
                  where near x y  =  abs (x - y) <= 0.008

--- the data values
type Weight      = Vector Double
type Predicted   = Vector Double
type Data        = Vector Double
type RNs         = Vector Double -- [Double]
type LPars       = Vector Double
type Rn          =  Double
type Cost        =  Double
type Lambda      =  Double
type Alpha       =  Double

-- projecting the compponents of the model
alph, sigD, sigE, tau :: Mmodel -> Double
alph   (MkModel a _ _ _ _)   =  a
sigD   (MkModel _ s _ _ _)   =  s
sigE   (MkModel _ _ e _ _)   =  e
tau    (MkModel _ _ _ _ z)   =  z

aBar :: Mmodel -> Vector Double
aBar   (MkModel _ _ _ a _)   =  a

negGradients :: Vector Double -> Bool
negGradients xs = True -- all negGD (zip xs [1..])

negGradientsM m = negGradients (aBar m)
negGradientsOM om = negGradientsM (oMod om)

-- update a model by a fraction of another - gradient decent update
updateModel :: Mmodel -> Double -> Mmodel -> Mmodel
updateModel  mO   lambda   mN  = MkModel nalph nsigD' nsigE' naBar ntau
                 where !nalph   = alph mO + lambda * (alph mN)
                       !nsigD   = sigD mO + lambda * (sigD mN)
                       !nsigE   = sigE mO + lambda * (sigE mN)
                       !ntau    = tau mO  + lambda * (tau mN)
                       !naBar   = zipVectorWith (+)  (aBar mO) (mapVector (lambda *)  (aBar mN))
                       --dehC     =  alpha * (c1 - c0)
                       nsigD'   = min nsigD nsigE
                       nsigE'   = max nsigD nsigE

--   update--- need to fix this to include lambda
tDupdate:: Mmodel -> Alpha -> Lambda -> Cost -> Cost ->  Mmodel -> Mmodel
tDupdate  mO alpha lam c0 c1    mN  = MkModel nalph nsigD nsigE naBar ntau
                 where nalph   =  alph mO + dehC * (alph mN - alph mO)
                       nsigD   =  sigD mO + dehC * (sigD mN  -  sigD mO )
                       nsigE   =  sigE mO + dehC * (sigE mN  -  sigE mO)
                       -- !nsigE   =  lam * sigE mO + dehC * (sigE mN - sigE mO)
                       ntau    =   tau mO  + dehC * (tau mN  -  tau mO)
                       naBar   = zipVectorWith (+) abar (mapVector (dehC *) (zipVectorWith (-) (aBar mN) abar))
                       dehC     =  alpha * (c1 - c0)
                       nsigD'   = min nsigD nsigE
                       nsigE'   = max nsigD nsigE
                       abar     = aBar mO
----

valid:: Mmodel -> Bool
valid (MkModel a b c _ e) = valAlph && valSigs && all ((/= "NaN") . show) [a,b,c]
     where
           !valAlph = a >= 0 && a <= 1
           !valSigs =  b > 0 && b < c

-- calculating probalilities for the data r_n.
r_n :: Mmodel -> Double -> Double -> Double
r_n    m  mu  x                     =  data_given_model  / p_data
     where       data_given_model   =  alpha * nxmu
                 p_data             =  alpha * nxmu + (1 - alpha) * nxmuE
                 !alpha             =  alph m
                 !nxmu              =  normal x mu (sigD m)
                 !nxmuE             =  normal x mu (sigE m)

-------- weighted sums for the emAlgorithm
rn_mean'' ::  Predicted  -> RNs -> ((Double, Double), (Double, Double))
rn_mean''  ps rns =  ( (dataMean , errMean) , (nd, ne) )
             where
                   ((nd, ne),(sumW, sumWE))  =  foldl'Rnf  fF ((0,0),(0,0)) pdrs -- $ zipVector (,)  ps  rns
                   mkD (a,r) (x, y)   =  (x + r  , y + (1 - r) )
                   mkN (a,r) (x, y)  =  (a + x, (1 - a) + y) -- (1 - r)
                   fF (b,b1) a       =  (mkD a b, mkN a b1)
                   pdrs              =  zip (toList ps) (toList rns)
                   --
                   mkMn f n p  b     =   p * (f b) /n
                   nn                =  nd + ne
                   dataMean          = nd / nn
                   errMean           = ne / nn

-- variances: ran_vars takes the result of rn_mean', which returns the
-- totals for the casses and the weighted sums
--rn_vars :: Mmodel   -> ((Double, Double), (Double, Double)) -> Data -> RNs  -> (Double, Double)
rn_vars  nD nE rns ps ds   =  (fst squares / nD , snd squares / nE )
            where squares    =  foldl'Rnf  fF (0, 0) (zip rns (zip (toList ps) (toList ds)))
                  fF (a,a1) (r ,(p,d)) =  (a + r * (p - d)^2 , a1 + (1-r) *(p - d)^2)

---------------------------------------------------------------------------------------------------------------
--- calculate the mixture model: takes a model and the number of cases, n
--------------------------------------------------------------------------------------------------------------
ml_alpha  ::  Mmodel -> Double -> Double
ml_alpha     ms      n             =  0.5 *  log n  - (n_alpha + 0.5) * log alpHa - ml_alpha_etc
            where    !ml_alpha_etc  =  (n - n_alpha + 0.5) * log (1 - alpHa) - 0.742
                     !n_alpha       =  alpHa * (n + 1)  - 0.5
                     !alpHa         =  alph ms

mmlSgAlpha :: Mmodel ->  RNs -> Double
mmlSgAlpha     m      rns  =  foldVector sF 0 rns
    where
        sF !a !r   =  let vr = val r in
                        a + vr * log  (vr / alpHa)  +  (1 - vr) * log  ((1 - vr) / alpH1)
        !alpHa     =  alph m
        !alpH1     =  1 - alph m
        val  a  | a >= 0 && a <= 1 = a
                | a < 0  = 0.0000000001
                | otherwise = 0.9999999999

----
mmlMixture :: Mmodel ->  RNs -> Double
mmlMixture      ms    rns    =     mla + mlSa
              where  n      = fromIntegral $ dim rns
                     !mla    = ml_alpha ms n
                     !mlSa   = mmlSgAlpha ms rns

------------------------------------------------------------------------------------------------
-----              learning the RPR
------------------------------------------------------------------------------------------------

-- a type for models and their message lengths
data OModel = Omod !Mmodel !Double  --deriving (Show)

-- averaging the oMods
oMod_mean    []   =  Nothing
alpha_mean (x:xs) =  Just (mDiv n ms)
  where alpha_mean_aux y ys = foldl'Rnf  aVg (1,y) ys
        aVg (!n, Omod (MkModel a b c ds  t) ml) (Omod (MkModel a1 b1 c1 ds1  t1) ml1) =
                   (n+1 , Omod (MkModel (a+a1) (b+b1) (c+c1) (zipVectorWith (+) ds  ds1) (t+t1)) (ml+ml1) )
        (n, ms) = alpha_mean_aux x xs
        mDiv n (Omod (MkModel a b c ds  t) ml) =
               Omod (MkModel (a/n) (b/n) (c/n) (mapVector (/n) ds) (t/n)) (ml/n)

instance Ord OModel where
   Omod _ ds <= Omod _ ds2  =   ds <= ds2
--
instance Eq OModel where
    Omod _ ds == Omod _ ds2  =  ds == ds2

--mkOmod  ms  ds  = Omod ms ds

oMod (Omod m _ ) = m
oModMl (Omod _ ml) = ml

--- valid OMod
validOM :: OModel -> Bool
validOM  (Omod m mgl)  =  noNan mgl && valid m
         where noNan a =   a > 0 - 10000000 &&  a < 10000000

--------------------------------------------------------------------------------------------
-- learn the parameters around the predicted value using the em algorithm
--nextModelR :: Predicted -> Data -> Vector Double -> Mmodel -> Mmodel
nextModelR ::  Data -> Predicted -> RNs -> Mmodel -> Mmodel
nextModelR ds ps rns m   =   MkModel aa1 bb cc1 (aBar m) (tau m) -- newMod, mixture, newRns)
                  where
                  ----------- quasi EM algoritm -------------------
                     !aa     =  (nD  + 0.5) / (nD + nE + 1)
                     aa1     =  if aa <  0.55 then (7 / 9.7) else aa -- then 1 - aa else aa
                     !bb     =  sqrt vd -- sd  of the data class
                     !cc     =  sqrt ve -- sd  of the error class
                     !cc1    =  max cc (bb * 1.5)
                     -- the new parametes are determined by values of the updated predictions
                     (!vd,!ve) =   rn_vars nD nE (toList rns) ps ds
                     !m1      =   rn_mean'' ps rns --
                     !nD      =   fst (snd m1)
                     !nE      =   snd (snd m1)

--- EM of sorts. Ideally, should iterate until the parameters are equal
-- nextModelR' :: Bool -> Matrix Double -> Vector Double ->  Mmodel
--nextModelR' :: Predicted -> Vector Double -> Vector Double -> Mmodel -> (Mmodel, Double)
--nextModelR' ps ys m  = nextModelR'' (invSlopes ps) ys m -- plst `pseq` (minimumBy (comparing snd)  plst) --
-- -- ys
{-# NOINLINE nextModelR' #-}
nextModelR' :: (RandomGen g) => g -> [Int] -> Maybe Bool -> Matrix Double -> Vector Double ->  Mmodel -> IO (Mmodel, Double)--
nextModelR' g cts restrictSlopes matrix ys inMod =  do
           let ks = (plst1 gg)
           ks `pseq` (return . ssd . (!! 0) $ ks)
    where
        (gg,_)      = split g
        ssd (_,b,c) = (b,c)
        trd (_,b,c) = c
        --
        ps          =  invS restrictSlopes . (matrix <> ) . aBar --
        ps1         =   (matrix <> ) . aBar
        --------------------------------------------------------------------------------------------
        ----- updating the errors
        rns m       =   zipVectorWith (r_n m) (ps m) ys
        faA         =  fromIntegral . (\ a -> if a == 2 then 1 else a `div` 2)  . dim -- . aBar
        --------------------------------------------------------------------------------------------
        applyMM g   = iterate  (unsafePerformIO . gMod gg) (1 , inMod , mmlMixture inMod (rns inMod) )
        plst1  g'   =  [ maybe ks id (listToMaybe xs)  |
                          let ys = myTake 35  . takeWhile (\(_,md,_) -> sigD md >= 0.01) $ applyMM g'
                        , let (front, ks) = (init ys , last ys)
                        , let xs  = [fromJust y | let y = similarBy simEqML  front, isJust y]
                        ]
        --- compare the message lengths  ---
        simEqML :: (a, Mmodel, Double) -> (a, Mmodel, Double) -> Bool
        simEqML (_, a ,mla) (_,b ,mlb)  =  abs (mla - mlb)  < 0.00001 -- sim_eq a b --
        ---

        gMod :: (RandomGen g) => g -> (Int, Mmodel, Double) ->  IO (Int , Mmodel, Double)
        gMod g (n , md, _)   =  do
            -- print ("iteration: " ++ show n ) -- ", probabilities: " ++ (show . (!! 8) . toList $ rns md) ++
            let dms mx = (rows mx, cols mx)
            -- print ("matrix dims: " ++ show (dms matrix) ++ " ; wWW dim: " ++ show (dms wWW))
            --       " alpha: " ++ show (alph md) ++ " std: " ++ show (sigD md)  )
            -- print ("num linear parms: " ++  show cts )
            let points = "\ndata points: "++ (printElm $ toList ys )
            let predictions = "\n predictions : "++ (printElm . toList $ ps md )
            let otherPredictions = "\n predictions : "++ (show . findCuts . toList $ ps1 md )
            let errors = "\nerrors : "++ (printElm . toList $ rns md )
            let dataStd = "\nData std  : "++ (show $ sigD md )
            let errorStd = "\nError std  : "++ (show $ sigE md )
            let abund = "\nabundance  : "++ (show $ alph md )
            let iter n' = "\n---- \niteration: " ++ show n' ++ "\n-----"
            let output n' = L.pack ((iter n') ++ points ++ predictions ++ errors ++ dataStd ++ errorStd ++ abund)
            let file = "probs1.txt"
            let mixR = mmlMixture nmd rnN
            return (n+1 , nmd,  mmlMixture nmd rnN )
           -- | otherwise =  (nmdd, mmlMixture nmdd rnN')

            where
                ------ restart with the original model if there is an invallid model ----
                md1  =   md
                ---------------------------------------------------------------------------
                !md' =  nextModelR ys (ps md1) (rns md1) md1
                nmd  = nextModelR ys (ps mdd) (rns mdd) mdd
                rnN  = rns nmd
                ------------ updating the linear parameters ---
                siG  =  sigD md' -- model
                siGE =  sigE md' -- model
                abar =  aBar md'
                rN   =  rns md'
                lnR  =  dim rN
                taU  =  tau md'
                ------
                wWW  =   diagRect 0  rN lnR lnR
                xTW  =  (trans matrix) <> wWW
                xTWx =  xTW <> matrix
                ---------------re calcuate the model with the new parameters ----------------------
                pli b (i,j) v = if i == j then v  else v -- + b
                zZ   =  mapMatrixWithIndex  (pli (1/taU^2)) xTWx
                --
                arN_aux =   head . toColumns . linearSolveLS zZ $ asColumn (xTW <> ys)
                arN  =   arN_aux
                ---
                nn   = faA arN
                tuN  =  sqrt ((arN <.> arN) /(nn + 1))
                mdd  =  MkModel (alph md') (sigD md') (sigE md') arN tuN

-------- inverting negative slopes --------------------------
invNeg :: Vector Double -> Vector Double
invNeg acs  = zipVectorWith flipSlp acs (fromList [1 .. dim acs] :: Vector Int)
    where
        flipSlp y a    = if even a && y < 0 then  invS1 y else y
        invS1 n
           | n >= -1   =  n *  (-0.4)
           | otherwise =  (- 0.6) / n

-- inverting the predictions
invS nSlp  =
    case nSlp  of
        Just tf ->  fromList . head  . manyInvNS tf . toList
        anyElse -> id

----------   INITIALIZE A MODEL
-----------------------------------------------------------------------------------------------
initModelR  gen acs       = initModelGD  gen (dim acs)

----
initModel  gen s acss    = MkModel alpha sgD sgE acs taU
            where
                   [alpha, kk]  =   take 2 (randomRs (0.6, 0.9) gen)
                   sgD    =   max 0.2 s -- (0.7 * s)   -- minimum values -- (values !! 0) (values !! 2)
                   sgE    =  65 * kk * sgD -- maximum  values
                   acs    =  invNeg acss --  acs -- invNeg acs
                   taU    =  vstDiv  acs

---- for reinitializing the mizture model
initModelMix  gen aA acs   = MkModel alpha sgD sgE asc1 taU
            where  --(g1, _)  =  split gen
                   sgD    =   head (randomRs (0.15, 4.5) gen)
                   alpha  =   head (randomRs (0.68, 0.95) gen)
                   --sgD     =  minimum values -- (values !! 0) (values !! 2)
                   sgE    =  4 * sgD -- maximum  values
                   asc1   =   invNeg acs
                   taU    =   vstDiv  asc1 -- sqrt ((asc1 <.> asc1) /(aA + 2)) -- vstDiv  asc1
                   ----

--- initializing the model for gradient decent
initModelGD  gen abarLen    = MkModel alpha (min sgD sgE) (max sgD sgE) abar taU
            where  (g1, _)  =  split gen
                   values  =  myTake 3 (randomRs (1, 50.0) gen)
                   alpha   =  (myTake 1 (randomRs (65.0, 95.0) gen) !! 0)/100
                   sgD     =  min (values !! 1) (values !! 2)
                   sgE     =  max (values !! 1) (values !! 2)
                   taU     =  vstDiv  abar -- values !! 3
                   abar    =  fromList (myTake abarLen  [ x | (y,a) <- zip (randomRs (-10, 20) gen) [1 .. ]
                                                              , let x =  if even a && a >= 0 then y else (-1) * y  ] )

------------------------------------------------------------------------------------------------
--- calculating probability values
filterV :: (Double -> Bool) -> Vector Double -> Vector Double
filterV  f  va  = initV $ foldVector (\a va -> if f a then vcons a va else va) dummy va
    where vcons x xs = join [fromList [x],xs]
          dummy      = fromList [-10] :: Vector Double
          initV vs   = if n > 1 then subVector 0 (n-1) vs else dummy
               where n = dim vs

-- error and data
-- Vector of Points -> Vector of means -> alpha -> (a -> Bool)
dataError :: Mmodel -> Vector Double -> Vector Double -> (Vector Double,Vector Double)
dataError  mod actual  predicted =  (mkProbs (>= am)  , mkProbs (< am) )
    where
        --mn                 =
        am                 =  let a = alph mod in a  -- - (0.05 * a)
        --am'                = let a = alph mod in a + (0.05 * a)
        mkProbs    f       =  zipVectorWith (\a b -> if f (calProbs mod a b) then a else -10) actual predicted
        calProbs mod x mu  = let dataProb = am * normal x mu (sigD mod)
                                 errorProb = (1- am) * normal x mu (sigE mod)
                             in
                                 dataProb / (dataProb + errorProb)

instance Show Mmodel where
    show (MkModel a b c ds ss)  = "\nalp:"++show a ++  "\nstd:" ++ rest1
             where lparms e e1  = "\nabr:" ++  printElm e  ++ "\ntau:" ++ show e1
                   rest1        = show b ++ "\nste:" ++ show c ++ lparms (toList ds) ss -- ++ "\n"
instance Show OModel where
   show (Omod m ml) = show m -- ++ "mln:" ++ show ml
-------------------------------------------------------------------------------------------------


{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE FlexibleContexts #-}
module STC.CompletionFieldR2S1 where

import           Array.UnboxedArray          as AU
import           Control.Monad               as M
import           Control.Monad.Parallel      as MP
import           Control.Parallel.Strategies
import           Data.Array.Repa             as R
import           Data.Complex
import           Data.List                   as L
import           Data.Vector.Storable        as VS
import           Data.Vector.Unboxed         as VU
import           DFT.Plan
import           Filter.Utils
import           Image.IO
import           Image.Transform
import           System.FilePath
import           System.Random
import           Text.Printf
import           Types
import           Utils.Array
import           Utils.Time
import           FokkerPlanck.FourierSeries
import FokkerPlanck.Analytic
import qualified Numeric.LinearAlgebra as NL
import qualified Numeric.LinearAlgebra.HMatrix as NL
import            STC.Utils
import FokkerPlanck.BrownianMotion (thetaPlus)
import Utils.Diagram

-- Implementation of [Williams and Jacbos, 1997]

{-# INLINE checkN #-}
checkN :: Int -> Int -> Int
checkN maxN n
  | n < 0 = checkN maxN (n + maxN)
  | n >= maxN = checkN maxN (n - maxN)
  | otherwise = n

{-# INLINE rotateST #-}
rotateST ::
     (R.Source r Double)
  => R.Array r DIM3 Double
  -> Int
  -> R.Array U DIM3 Double
rotateST arr 0 = computeUnboxedS . delay $ arr
rotateST arr n' =
  let (Z :. nf :. nx :. ny) = extent arr
      n = checkN nf n'
      deg = (-2) * pi / (fromIntegral nf) * fromIntegral n
  in rotate25D deg (fromIntegral $ center nx, fromIntegral $ center ny) $
     R.backpermute
       (extent arr)
       (\(Z :. k :. i :. j) -> (Z :. (mod (k + nf - n) nf) :. i :. j))
       arr

{-# INLINE timeReversal #-}
timeReversal :: (R.Source r e) => Array r DIM3 e -> Array D DIM3 e
timeReversal arr =
  let (Z :. nf :. _ :. _) = extent arr
      n = div nf 2
  in R.backpermute
       (extent arr)
       (\(Z :. k :. i :. j) -> (Z :. (mod (k + n) nf) :. i :. j))
       arr

{-# INLINE rotateSTTR #-}
rotateSTTR ::
     (R.Source r Double)
  => R.Array r DIM3 Double
  -> R.Array U DIM3 Double
rotateSTTR arr =
  let (Z :. _ :. nx :. ny) = extent arr
  in rotate25D pi (fromIntegral $ center nx, fromIntegral $ center ny) arr

{-# INLINE makeR2S1Plan #-}
makeR2S1Plan :: (R.Source r e) => DFTPlan -> R.Array r DIM3 e -> IO DFTPlan
makeR2S1Plan oldPlan arr = do
  let (Z :. orientations :. rows :. cols) = extent arr
  vec3D <-
    VS.map (:+ 0) . VS.fromList <$>
    M.replicateM (orientations * rows * cols) randomIO
  lock <- getFFTWLock
  fst <$>
    (dft1dGPlan lock oldPlan [orientations, rows, cols] [1, 2] vec3D >>= \(plan, vec) ->
       idft1dGPlan lock plan [orientations, rows, cols] [1, 2] vec)

{-# INLINE convolve #-}
convolve ::
     (R.Source r (Complex Double))
  => DFTPlan
  -> Array U DIM3 Double
  -> Array r DIM2 (Complex Double)
  -> IO (VU.Vector Double)
convolve dftPlan arr3D arr2DF = do
  let (Z :. orientations :. rows :. cols) = extent arr3D
      planID3D = DFTPlanID DFT1DG [orientations, rows, cols] [1, 2]
      inversePlanID = DFTPlanID IDFT1DG [orientations, rows, cols] [1, 2]
  vec3DF <-
    dftExecute dftPlan planID3D . VU.convert . VU.map (:+ 0) . toUnboxed $ arr3D
  let arr3DF =
        fromUnboxed (Z :. orientations :. rows :. cols) . VS.convert $ vec3DF
  vec <-
    dftExecute
      dftPlan
      inversePlanID
      (VU.convert . toUnboxed . computeS . R.traverse2 arr3DF arr2DF const $ \f3 f2 idx@(Z :. k :. i :. j) ->
         f3 idx * f2 (Z :. i :. j))
  return . VU.map realPart . VS.convert $ vec

{-# INLINE shareWeightST #-}
shareWeightST ::
     (R.Source r Double)
  => DFTPlan
  -> Array U DIM3 Double
  -> Array r DIM3 Double
  -> IO (Array D DIM3 Double)
shareWeightST dftPlan arr arrG = do
  let (Z :. nf :. _ :. _) = extent arrG
      (Z :. orientations :. rows :. cols) = extent arr
  arrF <-
    fmap (fromUnboxed (extent arr) . VS.convert) .
    dftExecute dftPlan (DFTPlanID DFT1DG [orientations, rows, cols] [1, 2]) .
    VU.convert . VU.map (:+ 0) . toUnboxed $
    arr
  xs <-
    MP.mapM
      (\i ->
         convolve
           dftPlan
           (computeS . makeFilter2D . rotateST arrG $ i)
           (R.slice arrF $ (Z :. i :. All :. All)))
      [0 .. nf - 1]
  return . delay . fromUnboxed (extent arrG) . L.foldl1' (VU.zipWith (+)) $ xs

{-# INLINE computeInitialDistribution #-}
computeInitialDistribution ::
     Int -> Int -> Int -> Double -> [R2S1RPPoint] -> R.Array U DIM3 Double
computeInitialDistribution xLen yLen numOrientation delta xs =
  let deltaTheat = 2 * pi / fromIntegral numOrientation
      xShift = div xLen 2
      (xMin, xMax) =
        if odd xLen
          then (-xShift, xShift)
          else (-xShift, xShift - 1)
      yShift = div yLen 2
      (yMin, yMax) =
        if odd yLen
          then (-yShift, yShift)
          else (-yShift, yShift - 1)
      vec =
        toUnboxedVector .
        AU.accum (+) 0 ((0, xMin, yMin), (numOrientation - 1, xMax, yMax)) .
        L.map
          (\(R2S1RPPoint (x, y, theta, _)) ->
             ( ((floor $ theta * pi / 180 / deltaTheat), round (fromIntegral x / delta), round (fromIntegral y / delta))
             , (1 / (fromIntegral . L.length $ xs)))) $
        xs
   in fromUnboxed (Z :. numOrientation :. xLen :. yLen) vec

-- Implementation of [Williams and Zweck, 2003]

{-# INLINE computeBias #-}
computeBias ::
     Int -> Int -> Int -> [R2S1RPPoint] -> R.Array D DIM3 Double
computeBias xLen yLen numOrientation xs =
  let xShift = div xLen 2
      (xMin, xMax) =
        if odd xLen
          then (-xShift, xShift)
          else (-xShift, xShift - 1)
      yShift = div yLen 2
      (yMin, yMax) =
        if odd yLen
          then (-yShift, yShift)
          else (-yShift, yShift - 1)
      vec =
        toUnboxedVector .
        AU.accum (+) 0 ((xMin, yMin), (xMax, yMax)) .
        L.map (\(R2S1RPPoint (x, y, _, _)) -> ((x, y), 1)) $
        xs
   in R.traverse
        (fromUnboxed (Z :. xLen :. yLen) vec)
        (const (Z :. numOrientation :. xLen :. yLen)) $ \f (Z :. _ :. i :. j) ->
        f (Z :. i :. j)

{-# INLINE computeInitialEigenVec #-}
computeInitialEigenVec ::
     Int -> Int -> Int -> [R2S1RPPoint] -> R.Array D DIM3 Double
computeInitialEigenVec xLen yLen numOrientation xs =
  let xShift = div xLen 2
      (xMin, xMax) =
        if odd xLen
          then (-xShift, xShift)
          else (-xShift, xShift - 1)
      yShift = div yLen 2
      (yMin, yMax) =
        if odd yLen
          then (-yShift, yShift)
          else (-yShift, yShift - 1)
      vec =
        toUnboxedVector .
        AU.accum (+) 0 ((xMin, yMin), (xMax, yMax)) .
        L.map
          (\(R2S1RPPoint (x, y, _, _)) ->
             ((x, y), (1 / (fromIntegral . L.length $ xs)))) $
        xs
   in R.traverse
        (fromUnboxed (Z :. xLen :. yLen) vec)
        (const (Z :. numOrientation :. xLen :. yLen)) $ \f (Z :. _ :. i :. j) ->
        f (Z :. i :. j) / fromIntegral numOrientation

{-# INLINE computeInitialEigenVecGaussian #-}
computeInitialEigenVecGaussian ::
     Int -> Int -> Int -> [(Int, Int, Double)] -> R.Array D DIM3 Double
computeInitialEigenVecGaussian xLen yLen numOrientation xs =
  let xShift = div xLen 2
      (xMin, xMax) =
        if odd xLen
          then (-xShift, xShift)
          else (-xShift, xShift - 1)
      yShift = div yLen 2
      (yMin, yMax) =
        if odd yLen
          then (-yShift, yShift)
          else (-yShift, yShift - 1)
      vec =
        toUnboxedVector .
        AU.accum (+) 0 ((xMin, yMin), (xMax, yMax)) .
        L.map (\(x, y, v) -> ((x, y), v)) $
        xs
  in R.traverse
       (fromUnboxed (Z :. xLen :. yLen) vec)
       (const (Z :. numOrientation :. xLen :. yLen)) $ \f (Z :. _ :. i :. j) ->
       f (Z :. i :. j)

{-# INLINE eigenVectorR2S1 #-}
eigenVectorR2S1 ::
     (R.Source s Double, R.Source s1 Double, R.Source s2 Double)
  => DFTPlan
  -> FilePath
  -> R.Array s DIM3 Double
  -> R.Array s DIM3 Double
  -> Int
  -> Bool
  -> R.Array s1 DIM3 Double
  -> R.Array s2 DIM3 Double
  -> IO (R.Array D DIM3 Double)
eigenVectorR2S1 plan folderPath filterSource filterSink n writeFlag bias inputSource = do
  let sourceDist = R.zipWith (*) bias inputSource
      -- s = R.sumAllS sourceDist
      s = L.maximum . R.toList $ sourceDist
      sourceVec = computeS $ R.map (/ s) sourceDist
      (Z :. numOrientation :. _ :. _) = extent inputSource
  printCurrentTime (show n)
  sourceEigenVec <- shareWeightST plan sourceVec filterSource
  -- plotThetaDimension folderPath (printf "R2S1_%d_" n) (31,-4) .
  --   R.backpermute
  --     (extent sourceEigenVec)
  --     (\(Z :. k :. i :. j) ->
  --        (Z :.
  --         (mod (k + numOrientation - (div numOrientation 2)) numOrientation :: Int) :.
  --         i :.
  --         j)) $
  --   sourceEigenVec
  -- sinkEigenVec <-
  --   timeReversal <$>
  --   shareWeightST
  --     plan
  --     (computeS . R.zipWith (*) bias $ sourceEigenVec')
  --     filterSink
  -- let sourceEigenVec =
  --       R.zipWith (\x y -> x + 0.2 * y) sourceEigenVec' sinkEigenVec
  when
    writeFlag
    (plotImageRepa
       (folderPath </> printf "Source_%d.png" n)
       (ImageRepa 8 .
        computeS . R.extend (Z :. (1 :: Int) :. All :. All) . R.map (\(x) -> (abs x) ** (1/3)) . R.sumS . rotate3D $
        sourceEigenVec))
  return sourceEigenVec

powerMethod ::
     (R.Source s1 Double, R.Source s2 Double)
  => DFTPlan
  -> FilePath
  -> (R.Array U DIM3 Double)
  -> Int
  -> Bool
  -> String
  -> Double
  -> (R.Array s1 DIM3 Double)
  -> (R.Array s2 DIM3 Double)
  -> IO (R.Array U DIM3 Double)
powerMethod oldPlan folderPath filterSource numIteration writeFlag idStr threshold bias' initialEigenVec = do
  let (Z :. numOrientation :. xLen :. yLen) = extent filterSource
      filterSink = rotateSTTR filterSource
      bias = delay bias'
  plotImageRepa
    (folderPath </> "Bias.png")
    (ImageRepa 8 .
     computeS . R.extend (Z :. (1 :: Int) :. All :. All) . R.sumS . rotate3D $
     bias)
  plan <- makeR2S1Plan oldPlan filterSource
  source <-
    M.foldM
      (\input n ->
         eigenVectorR2S1
           plan
           folderPath
           filterSource
           filterSink
           n
           writeFlag
           bias
           input)
      (delay initialEigenVec)
      [1 .. numIteration]
  -- print . R.toList . R.slice source $ (Z :. All :. (div xLen 2) :. (div yLen 2))
  -- plotThetaDimension folderPath "R2S1_Source_" (31,-4) .
  --   R.backpermute
  --     (extent source)
  --     (\(Z :. k :. i :. j) ->
  --        (Z :.
  --         (mod (k + numOrientation - (div numOrientation 2)) numOrientation :: Int) :.
  --         i :.
  --         j)) $
  --   source
  sink <- shareWeightST plan (rotateSTTR $ bias *^ source) filterSink
  plotImageRepa
    (folderPath </> "Sink.png")
    (ImageRepa 8 .
     computeS . R.extend (Z :. (1 :: Int) :. All :. All) . R.sumS . rotate3D $
     sink)
  -- plotThetaDimension folderPath "R2S1_Sink_" (31,-4) .
  --   R.backpermute
  --     (extent sink)
  --     (\(Z :. k :. i :. j) ->
  --        (Z :.
  --         (mod (k + numOrientation - (div numOrientation 2)) numOrientation :: Int) :.
  --         i :.
  --         j)) $
  --   sink
  let completion = computeS $ R.zipWith (*) source sink
      completionR2 = R.sumS . rotate3D $ completion
      completionR2Vec = toUnboxed completionR2
      completionMax = VU.maximum completionR2Vec
      completionMin = VU.minimum completionR2Vec
  plotImageRepa
    (folderPath </> printf "Completion%s.png" idStr)
    (ImageRepa 8 .
     computeS .
     R.extend (Z :. (1 :: Int) :. All :. All) .
     R.map (\x -> sqrt $ (x - completionMin) / (completionMax - completionMin)) $
     completionR2)
  -- plotThetaDimension folderPath "R2S1_Completion_" (31, -4) .
  --   R.backpermute
  --     (extent completion)
  --     (\(Z :. k :. i :. j) ->
  --        (Z :.
  --         (mod (k + numOrientation - (div numOrientation 2)) numOrientation :: Int) :.
  --         i :.
  --         j)) $
  --   completion
  return completion

eigenVectorR2S1Analytic ::
     Int -> Double -> Double -> Double -> Double -> [(Double, Double)] -> IO [[Double]]
eigenVectorR2S1Analytic !oris !sigma !tau !gamma !delta !xs = do
  let !numPoints = L.length xs
      !locArr = fromListUnboxed (Z :. L.length xs) xs
      !deltaTheta = 2 * pi / fromIntegral oris
      -- !weights =
      --   [1 / 16, 1 / 8, 1 / 16, 1 / 8, 1 / 4, 1 / 8, 1 / 16, 1 / 8, 1 / 16]
      -- !origins =
      --   L.zipWith
      --     (\w (i, j) -> (i * delta, j * delta, w))
      --     weights
      --     [(i, j) | i <- [-1 .. 1], j <- [-1 .. 1]]
      transitionMatrixArr' =
        R.traverse locArr (\_ -> (Z :. numPoints :. oris :. numPoints :. oris)) $ \f (Z :. i1 :. j1 :. i2 :. j2) ->
          let rowIdx = i1 * oris + j1
              colIdx = i2 * oris + j2
          in if rowIdx /= colIdx && i1 /= i2
             -- then L.foldl'
             --        (\s (i, j, w) ->
             --           s +
             --           w *
             --           computePji
             --             sigma
             --             tau
             --             (R2S1RP
             --                (fst . f $ (Z :. i2))
             --                (snd . f $ (Z :. i2))
             --                (deltaTheta * fromIntegral j2)
             --                gamma)
             --             (R2S1RP
             --                (i + (fst . f $ (Z :. i1)))
             --                (j + (snd . f $ (Z :. i1)))
             --                (deltaTheta * fromIntegral j1)
             --                gamma))
             --        0
             --        origins
             then computePji
                    sigma
                    tau
                    (R2S1RP
                       ((fst . f $ (Z :. i2)))
                       ((snd . f $ (Z :. i2)))
                       (deltaTheta * fromIntegral j2)
                       gamma)
                    (R2S1RP
                       (fst . f $ (Z :. i1))
                       (snd . f $ (Z :. i1))
                       (deltaTheta * fromIntegral j1)
                       gamma)
             else 0
  transitionMatrixArr <- computeUnboxedP transitionMatrixArr'
      -- transitionMatrixArr = computeUnboxedS transitionMatrixArr'
  let !size = numPoints * oris
      transitionMatrix =
        deepSeqArray transitionMatrixArr . (size NL.>< size) . R.toList $
        transitionMatrixArr
      (eigVal, eigVec) = NL.eig transitionMatrix
      pairs =
        L.reverse .
        L.sortOn (realPart . fst) .
        L.filter
          (\(c, _) ->
             let (m, p) = polar c
             in abs (p) < 1e-10) $
        L.zip (NL.toList eigVal) (L.map NL.toList . NL.toColumns $ eigVec)
      eigenVector =
        fromListUnboxed (Z :. numPoints :. oris) .
        L.map magnitude . snd . L.head $
        pairs
  print . L.map (polar . fst) . L.take 10 $ pairs
  print . L.map (polar . fst) . L.take 10 . L.reverse $ pairs
  print . L.maximum . L.map realPart . snd . L.head $ pairs
  print . L.minimum . L.map realPart . snd . L.head $ pairs
  return . L.map (\i -> R.toList . R.slice eigenVector $ (Z :. i :. All)) $
    [0 .. numPoints - 1]

computeContourR2S1Analytic ::
     (R.Source s Double)
  => FilePath
  -> (R.Array U DIM3 Double)
  -> Int
  -> Int
  -> Int
  -> Double
  -> Double
  -> Double
  -> Double
  -> Double
  -> R.Array s DIM3 Double
  -> [(Double, Double)]
  -> IO (R.Array U DIM3 Double)
computeContourR2S1Analytic !folderPath !filterSource !oris !rows !cols !sigma !tau !gamma !std !delta !bias !xs' = do
  let xs =
        L.map
          (\(x', y') -> (fromIntegral . round $ x', fromIntegral . round $ y'))
          xs'
  eigenVector <-
    eigenVectorR2S1Analytic
      oris
      sigma
      tau
      gamma
      delta
      (L.map (\(a, b) -> (a * delta, b * delta)) xs)
  let ys =
        L.concat $
        L.zipWith
          (\(x, y) thetas ->
             L.zipWith
               (\k theta -> ((k, round x, round y), theta))
               [0 ..]
               thetas)
          xs
          eigenVector
      (!minR, !maxR) = computeRange rows
      (!minC, !maxC) = computeRange cols
      eigenSourceArr =
        fromUnboxed (Z :. oris :. cols :. rows) .
        toUnboxedVector .
        AU.accum (+) 0 ((0, minC, minR), (oris - 1, maxC, maxR)) $
        ys
      eigenSinkArr = rotateSTTR eigenSourceArr
      completionArr = computeUnboxedS $ eigenSourceArr *^ eigenSinkArr
      filterSink = rotateSTTR filterSource
  plotR2S1Array
    (folderPath </> "EigenSource.eps")
    (fromIntegral rows)
    (fromIntegral cols)
    (fromIntegral rows / 16)
    xs
    eigenSourceArr
  plotR2S1Array
    (folderPath </> "EigenSink.eps")
    (fromIntegral rows)
    (fromIntegral cols)
    (fromIntegral rows / 16)
    xs
    eigenSinkArr
  plotR2S1Array
    (folderPath </> "EigenCompletion.eps")
    (fromIntegral rows)
    (fromIntegral cols)
    (fromIntegral rows / 16)
    xs
    completionArr
  plan <- makeR2S1Plan emptyPlan filterSource
  source <- shareWeightST plan eigenSourceArr filterSource
  plotImageRepa
    (folderPath </> "Source.png")
    (ImageRepa 8 .
     computeS . R.extend (Z :. (1 :: Int) :. All :. All) . R.sumS . rotate3D $
     source)
  plotR2S1Array
    (folderPath </> "Source.eps")
    (fromIntegral rows)
    (fromIntegral cols)
    (fromIntegral rows / 16)
    xs
    source
  sink <- shareWeightST plan eigenSinkArr filterSink
  plotImageRepa
    (folderPath </> "Sink.png")
    (ImageRepa 8 .
     computeS . R.extend (Z :. (1 :: Int) :. All :. All) . R.sumS . rotate3D $
     sink)
  plotR2S1Array
    (folderPath </> "Sink.eps")
    (fromIntegral rows)
    (fromIntegral cols)
    (fromIntegral rows / 16)
    xs
    sink
  let completion = computeS $ R.zipWith (*) source sink
      completionR2 = R.sumS . rotate3D $ completion
      completionR2Vec = toUnboxed completionR2
      completionMax = VU.maximum completionR2Vec
      completionMin = VU.minimum completionR2Vec
  print (completionMin, completionMax)
  plotImageRepa
    (folderPath </> printf "Completion.png")
    (ImageRepa 8 .
     computeS .
     R.extend (Z :. (1 :: Int) :. All :. All) .
     -- R.map (\x -> (abs x) ** (1 / 2)) .
     R.map (\x -> sqrt $ (x - completionMin) / (completionMax - completionMin)) $
     completionR2)
  plotR2S1Array
    (folderPath </> "Completion.eps")
    (fromIntegral rows)
    (fromIntegral cols)
    (fromIntegral rows / 16)
    xs
    completion
  return completion

computeContourR2S1Tangent ::
     FilePath
  -> (R.Array U DIM3 Double)
  -> Int
  -> Int
  -> Int
  -> Double
  -> [(Double, Double)]
  -> IO (R.Array U DIM3 Double)
computeContourR2S1Tangent !folderPath !filterSource !oris !rows !cols !delta !xs = do
  let (cMin, cMax) = computeRange cols
      (rMin, rMax) = computeRange rows
      !deltaTheta = 2 * pi / fromIntegral oris
      initArr =
        fromUnboxed (Z :. oris :. cols :. rows) .
        toUnboxedVector .
        AU.accum (+) 0 ((0, cMin, rMin), (oris - 1, cMax, rMax)) .
        L.concatMap
          (\(x, y) ->
             let phi = atan2 y x
                 idx1 = floor $ ((phi `thetaPlus` (pi / 2)) + pi) / deltaTheta
                 idx2 = floor $ ((phi `thetaPlus` ((-pi) / 2)) + pi) / deltaTheta
             in [((idx1, round x, round y), 1), ((idx2, round x, round y), 1)]) $
        xs
      filterSink = rotateSTTR filterSource
  plan <- makeR2S1Plan emptyPlan filterSource
  source <- shareWeightST plan initArr filterSource
  plotImageRepa
    (folderPath </> "Source.png")
    (ImageRepa 8 .
     computeS .
     R.extend (Z :. (1 :: Int) :. All :. All) .
     R.map (\x -> (abs x) ** (1 / 3)) . R.sumS . rotate3D $
     source)
  sink <- shareWeightST plan initArr filterSink
  plotImageRepa
    (folderPath </> "Sink.png")
    (ImageRepa 8 .
     computeS .
     R.extend (Z :. (1 :: Int) :. All :. All) .
     R.map (\x -> (abs x) ** (1 / 3)) . R.sumS . rotate3D $
     sink)
  let completion = computeS $ R.zipWith (*) source sink
      completionR2 = R.sumS . rotate3D $ completion
      completionR2Vec = toUnboxed completionR2
      completionMax = VU.maximum completionR2Vec
      completionMin = VU.minimum completionR2Vec
  print (completionMin, completionMax)
  plotImageRepa
    (folderPath </> printf "Completion.png")
    (ImageRepa 8 .
     computeS .
     R.extend (Z :. (1 :: Int) :. All :. All) .
     R.map (\x -> (abs x) ** (1 / 2)) .
     R.map (\x -> sqrt $ (x - completionMin) / (completionMax - completionMin)) $
     completionR2)
  return completion