{-# LANGUAGE CPP                   #-}
{-# LANGUAGE UndecidableInstances  #-}
{-# LANGUAGE PolyKinds             #-}
{-# LANGUAGE DataKinds             #-}
{-# LANGUAGE GADTs                 #-}
{-# LANGUAGE RankNTypes            #-}
{-# LANGUAGE TypeOperators         #-}
{-# LANGUAGE TypeFamilies          #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE FlexibleContexts      #-}
{-# LANGUAGE ViewPatterns          #-}
{-# LANGUAGE ScopedTypeVariables   #-}
{-# OPTIONS_GHC -Wno-incomplete-uni-patterns #-}

module Grenade.Recurrent.Layers.LSTM (
    LSTM (..)
  , LSTMWeights (..)
  , randomLSTM
  ) where

import           Control.Monad.Primitive           (PrimBase, PrimState)
import           System.Random.MWC hiding (create)
import GHC.TypeLits

-- import           Data.List ( foldl1' )
import           Data.Proxy
import           Data.Serialize

#if MIN_VERSION_base(4,9,0)
import           Data.Kind (Type)
#endif

import qualified Numeric.LinearAlgebra as LA
import           Numeric.LinearAlgebra.Static

import           Grenade.Core
import           Grenade.Utils.LinearAlgebra
import           Grenade.Recurrent.Core
import           Grenade.Layers.Internal.Update

-- | Long Short Term Memory Recurrent unit
--
--   This is a Peephole formulation, so the recurrent shape is
--   just the cell state, the previous output is not held or used
--   at all.
data LSTM :: Nat -> Nat -> Type where
  LSTM :: ( KnownNat input
          , KnownNat output
          ) => !(LSTMWeights input output) -- Weights
            -> !(LSTMWeights input output) -- Momentums
            -> LSTM input output

data LSTMWeights :: Nat -> Nat -> Type where
  LSTMWeights :: ( KnownNat input
                 , KnownNat output
                 ) => {
                   lstmWf :: !(L output input)  -- Weight Forget     (W_f)
                 , lstmUf :: !(L output output) -- Cell State Forget (U_f)
                 , lstmBf :: !(R output)        -- Bias Forget       (b_f)
                 , lstmWi :: !(L output input)  -- Weight Input      (W_i)
                 , lstmUi :: !(L output output) -- Cell State Input  (U_i)
                 , lstmBi :: !(R output)        -- Bias Input        (b_i)
                 , lstmWo :: !(L output input)  -- Weight Output     (W_o)
                 , lstmUo :: !(L output output) -- Cell State Output (U_o)
                 , lstmBo :: !(R output)        -- Bias Output       (b_o)
                 , lstmWc :: !(L output input)  -- Weight Cell       (W_c)
                 , lstmBc :: !(R output)        -- Bias Cell         (b_c)
                 } -> LSTMWeights input output

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

instance FoldableGradient (LSTMWeights input output) where
  mapGradient f (LSTMWeights wf uf bf wi ui bi wo uo bo wc bc) =
    LSTMWeights (dmmap f wf) (dmmap f uf) (dvmap f bf) (dmmap f wi) (dmmap f ui) (dvmap f bi) (dmmap f wo) (dmmap f uo) (dvmap f bo) (dmmap f wc) (dvmap f bc)
  squaredSums (LSTMWeights wf uf bf wi ui bi wo uo bo wc bc) =
    [ sumM . squareM $ wf
    , sumM . squareM $ uf
    , sumV . squareV $ bf
    , sumM . squareM $ wi
    , sumM . squareM $ ui
    , sumV . squareV $ bi
    , sumM . squareM $ wo
    , sumM . squareM $ uo
    , sumV . squareV $ bo
    , sumM . squareM $ wc
    , sumV . squareV $ bc
    ]

instance (KnownNat i, KnownNat o) => UpdateLayer (LSTM i o) where
  -- The gradients are the same shape as the weights and momentum
  -- This seems to be a general pattern, maybe it should be enforced.
  type Gradient (LSTM i o) = (LSTMWeights i o)

  -- Run the update function for each group matrix/vector of weights, momentums and gradients.
  -- Hmm, maybe the function should be used instead of passing in the learning parameters.
  runUpdate opt@OptSGD{} (LSTM w m) g =
    let MatrixResultSGD wf wf' = u lstmWf w m g
        MatrixResultSGD uf uf' = u lstmUf w m g
        VectorResultSGD bf bf' = v lstmBf w m g
        MatrixResultSGD wi wi' = u lstmWi w m g
        MatrixResultSGD ui ui' = u lstmUi w m g
        VectorResultSGD bi bi' = v lstmBi w m g
        MatrixResultSGD wo wo' = u lstmWo w m g
        MatrixResultSGD uo uo' = u lstmUo w m g
        VectorResultSGD bo bo' = v lstmBo w m g
        MatrixResultSGD wc wc' = u lstmWc w m g
        VectorResultSGD bc bc' = v lstmBc w m g
    in LSTM (LSTMWeights wf uf bf wi ui bi wo uo bo wc bc) (LSTMWeights wf' uf' bf' wi' ui' bi' wo' uo' bo' wc' bc')
      where
    -- Utility function for updating with the momentum, gradients, and weights.
    u :: forall x ix out. (KnownNat ix, KnownNat out) => (x -> L out ix) -> x -> x -> x -> MatrixResult out ix
    u e (e -> weights) (e -> momentum) (e -> gradient) =
      descendMatrix opt (MatrixValuesSGD weights gradient momentum)

    v :: forall x ix. (KnownNat ix) => (x -> R ix) -> x -> x -> x -> VectorResult ix
    v e (e -> weights) (e -> momentum) (e -> gradient) =
      descendVector opt (VectorValuesSGD weights gradient momentum)
  runUpdate _ d g = runUpdate defOptimizer d g

  -- There's a lot of updates here, so to try and minimise the number of data copies
  -- we'll create a mutable bucket for each.
  -- runUpdates rate lstm gs =
  --   let combinedGradient = foldl1' uu gs
  --   in  runUpdate rate lstm combinedGradient
  --     where
  --   uu :: (KnownNat i, KnownNat o) => LSTMWeights i o -> LSTMWeights i o -> LSTMWeights i o
  --   uu a b =
  --     let wf = u lstmWf a b
  --         uf = u lstmUf a b
  --         bf = v lstmBf a b
  --         wi = u lstmWi a b
  --         ui = u lstmUi a b
  --         bi = v lstmBi a b
  --         wo = u lstmWo a b
  --         uo = u lstmUo a b
  --         bo = v lstmBo a b
  --         wc = u lstmWc a b
  --         bc = v lstmBc a b
  --     in LSTMWeights wf uf bf wi ui bi wo uo bo wc bc
  --   u :: forall x ix out. (KnownNat ix, KnownNat out) => (x -> (L out ix)) -> x -> x -> L out ix
  --   u e (e -> a) (e -> b) = tr $ tr a + tr b

  --   v :: forall x ix. (x -> (R ix)) -> x -> x -> R ix
  --   v e (e -> a) (e -> b) = a + b

instance (KnownNat i, KnownNat o, KnownNat (i*o), KnownNat (o*o)) => RandomLayer (LSTM i o) where
  createRandomWith = randomLSTM

instance (KnownNat i, KnownNat o) => RecurrentUpdateLayer (LSTM i o) where
  -- The recurrent shape is the same size as the output.
  -- It's actually the cell state however, as this is a peephole variety LSTM.
  type RecurrentShape (LSTM i o) = S ('D1 o)

instance (KnownNat i, KnownNat o) => RecurrentLayer (LSTM i o) ('D1 i) ('D1 o) where

  -- The tape stores essentially every variable we calculate,
  -- so we don't have to run any forwards component again.
  type RecTape (LSTM i o) ('D1 i) ('D1 o) = (R o, R i, R o, R o, R o, R o, R o, R o, R o, R o, R o)
  -- Forward propagation for the LSTM layer.
  -- The size of the cell state is also the size of the output.
  runRecurrentForwards (LSTM lw _) (S1D cell) (S1D input) =
    let -- Forget state vector
        f_s = lstmBf lw + lstmWf lw #> input + lstmUf lw #> cell
        f_t = sigmoid f_s
        -- Input state vector
        i_s = lstmBi lw + lstmWi lw #> input + lstmUi lw #> cell
        i_t = sigmoid i_s
        -- Output state vector
        o_s = lstmBo lw + lstmWo lw #> input + lstmUo lw #> cell
        o_t = sigmoid o_s
        -- Cell input state vector
        c_s = lstmBc lw + lstmWc lw #> input
        c_x = tanh c_s
        -- Cell state
        c_t = f_t * cell + i_t * c_x
        -- Output (it's sometimes recommended to use tanh c_t)
        h_t = o_t * c_t
    in ((cell, input, f_s, f_t, i_s, i_t, o_s, o_t, c_s, c_x, c_t), S1D c_t, S1D h_t)

  -- Run a backpropogation step for an LSTM layer.
  -- We're doing all the derivatives by hand here, so one should
  -- be extra careful when changing this.
  --
  -- There's a test version using the AD library without hmatrix in the test
  -- suite. These should match always.
  runRecurrentBackwards (LSTM lw _) (cell, input, f_s, f_t, i_s, i_t, o_s, o_t, c_s, c_x, c_t) (S1D cellGrad) (S1D h_t') =
    let -- Reverse Mode AD Derivitives
        c_t' = h_t' * o_t + cellGrad

        f_t' = c_t' * cell
        f_s' = sigmoid' f_s * f_t'

        o_t' = h_t' * c_t
        o_s' = sigmoid' o_s * o_t'

        i_t' = c_t' * c_x
        i_s' = sigmoid' i_s * i_t'

        c_x' = c_t' * i_t
        c_s' = tanh' c_s * c_x'

        -- The derivatives to pass sideways (recurrent) and downwards
        cell'  = tr (lstmUf lw) #> f_s' + tr (lstmUo lw) #> o_s' + tr (lstmUi lw) #> i_s' + c_t' * f_t
        input' = tr (lstmWf lw) #> f_s' + tr (lstmWo lw) #> o_s' + tr (lstmWi lw) #> i_s' + tr (lstmWc lw) #> c_s'

        -- Calculate the gradient Matricies for the input
        lstmWf' = f_s' `outer` input
        lstmWi' = i_s' `outer` input
        lstmWo' = o_s' `outer` input
        lstmWc' = c_s' `outer` input

        -- Calculate the gradient Matricies for the cell
        lstmUf' = f_s' `outer` cell
        lstmUi' = i_s' `outer` cell
        lstmUo' = o_s' `outer` cell

        -- The biases just get the values, but we'll write it so it's obvious
        lstmBf' = f_s'
        lstmBi' = i_s'
        lstmBo' = o_s'
        lstmBc' = c_s'

        gradients = LSTMWeights lstmWf' lstmUf' lstmBf' lstmWi' lstmUi' lstmBi' lstmWo' lstmUo' lstmBo' lstmWc' lstmBc'
    in  (gradients, S1D cell', S1D input')

-- | Generate an LSTM layer with random Weights
--   one can also just call createRandom from UpdateLayer
--
--   Has forget gate biases set to 1 to encourage early learning.
--
--   https://github.com/karpathy/char-rnn/commit/0dfeaa454e687dd0278f036552ea1e48a0a408c9
--
randomLSTM :: forall m i o. (PrimBase m, KnownNat i, KnownNat o, KnownNat (i*o),KnownNat (o*o))
           => NetworkInitSettings -> Gen (PrimState m) -> m (LSTM i o)
randomLSTM (NetworkInitSettings m HMatrix _) gen = do
  let w = getRandomMatrix i o m gen
  let u = getRandomMatrix i o m gen
  let v = getRandomVector i o m gen

  let w0 = konst 0
      u0 = konst 0
      v0 = konst 0

  LSTM <$> (LSTMWeights <$> w <*> u <*> pure (konst 1) <*> w <*> u <*> v <*> w <*> u <*> v <*> w <*> v)
         <*> pure (LSTMWeights w0 u0 v0 w0 u0 v0 w0 u0 v0 w0 v0)

  where i = natVal (Proxy :: Proxy i)
        o = natVal (Proxy :: Proxy o)
randomLSTM (NetworkInitSettings _ cpu _) _ = error $ "CPU backend " ++ show cpu ++ " not supported by LSTM layer"

-- | Maths
--
-- TODO: Move to not here
--       Optimise backwards derivative
sigmoid :: Floating a => a -> a
sigmoid x = 1 / (1 + exp (-x))

sigmoid' :: Floating a => a -> a
sigmoid' x = logix * (1 - logix)
  where
    logix = sigmoid x

tanh' :: (Floating a) => a -> a
tanh' t = 1 - s ^ (2 :: Int)  where s = tanh t

instance (KnownNat i, KnownNat o) => Serialize (LSTM i o) where
  put (LSTM lw _) = do
      u (lstmWf lw)
      u (lstmUf lw)
      v (lstmBf lw)
      u (lstmWi lw)
      u (lstmUi lw)
      v (lstmBi lw)
      u (lstmWo lw)
      u (lstmUo lw)
      v (lstmBo lw)
      u (lstmWc lw)
      v (lstmBc lw)
    where
      u :: forall a b. (KnownNat a, KnownNat b) => Putter  (L b a)
      u = putListOf put . LA.toList . LA.flatten . extract
      v :: forall a. (KnownNat a) => Putter (R a)
      v = putListOf put . LA.toList . extract

  get = do
      w <- LSTMWeights <$> u <*> u <*> v <*> u <*> u <*> v <*> u <*> u <*> v <*> u <*> v
      return $ LSTM w (LSTMWeights w0 u0 v0 w0 u0 v0 w0 u0 v0 w0 v0)
    where
      u :: forall a b. (KnownNat a, KnownNat b) => Get  (L b a)
      u = let f = fromIntegral $ natVal (Proxy :: Proxy a)
          in  maybe (fail "Vector of incorrect size") return . create . LA.reshape f . LA.fromList =<< getListOf get
      v :: forall a. (KnownNat a) => Get (R a)
      v = maybe (fail "Vector of incorrect size") return . create . LA.fromList =<< getListOf get

      w0 = konst 0
      u0 = konst 0
      v0 = konst 0

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

instance (KnownNat i, KnownNat o) => GNum (LSTM i o) where
  n |* (LSTM w m) = LSTM (n |* w) (n |* m)
  (LSTM w1 m1) |+ (LSTM w2 m2) = LSTM (w1 |+ w2) (m1 |+ m2)

instance (KnownNat i, KnownNat o) => GNum (LSTMWeights i o) where
  n |* (LSTMWeights wf uf bf wi ui bi wo uo bo wc bc) =
    LSTMWeights
      (fromRational n * wf)
      (fromRational n * uf)
      (fromRational n * bf)
      (fromRational n * wi)
      (fromRational n * ui)
      (fromRational n * bi)
      (fromRational n * wo)
      (fromRational n * uo)
      (fromRational n * bo)
      (fromRational n * wc)
      (fromRational n * bc)
  (LSTMWeights wf1 uf1 bf1 wi1 ui1 bi1 wo1 uo1 bo1 wc1 bc1) |+ (LSTMWeights wf2 uf2 bf2 wi2 ui2 bi2 wo2 uo2 bo2 wc2 bc2) =
    LSTMWeights (wf1 + wf2) (uf1 + uf2) (bf1 + bf2) (wi1 + wi2) (ui1 + ui2) (bi1 + bi2) (wo1 + wo2) (uo1 + uo2) (bo1 + bo2) (wc1 + wc2) (bc1 + bc2)


{-# LANGUAGE CPP                   #-}
{-# LANGUAGE ConstraintKinds       #-}
{-# LANGUAGE FlexibleContexts      #-}
{-# LANGUAGE FlexibleInstances     #-}
{-# LANGUAGE GADTs                 #-}
{-# LANGUAGE MagicHash             #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE PatternSynonyms       #-}
{-# LANGUAGE RebindableSyntax      #-}
{-# LANGUAGE ScopedTypeVariables   #-}
{-# LANGUAGE TypeApplications      #-}
{-# LANGUAGE TypeFamilies          #-}
{-# LANGUAGE TypeSynonymInstances  #-}
{-# LANGUAGE UndecidableInstances  #-}
{-# LANGUAGE ViewPatterns          #-}
{-# OPTIONS_GHC -fno-warn-orphans #-}
-- |
-- Module      : Data.Array.Accelerate.Data.Complex
-- Copyright   : [2015..2019] The Accelerate Team
-- License     : BSD3
--
-- Maintainer  : Trevor L. McDonell <trevor.mcdonell@gmail.com>
-- Stability   : experimental
-- Portability : non-portable (GHC extensions)
--
-- Complex numbers, stored in the usual C-style array-of-struct representation,
-- for easy interoperability.
--
module Data.Array.Accelerate.Data.Complex (

  -- * Rectangular from
  Complex(..), pattern (::+),
  real,
  imag,

  -- * Polar form
  mkPolar,
  cis,
  polar,
  magnitude, magnitude',
  phase,

  -- * Conjugate
  conjugate,

) where

import Data.Array.Accelerate.Array.Sugar
import Data.Array.Accelerate.Classes
import Data.Array.Accelerate.Data.Functor
import Data.Array.Accelerate.Pattern
import Data.Array.Accelerate.Prelude
import Data.Array.Accelerate.Smart
import Data.Array.Accelerate.Type

import Data.Complex                                                 ( Complex(..) )
import Prelude                                                      (($))
import qualified Data.Complex                                       as C
import qualified Prelude                                            as P

infix 6 ::+
pattern (::+) :: Elt a => Exp a -> Exp a -> Exp (Complex a)
pattern r ::+ i <- (deconstructComplex -> (r, i))
  where (::+) = constructComplex
{-# COMPLETE (::+) #-}

-- Use an array-of-structs representation for complex numbers if possible.
-- This matches the standard C-style layout, but we can use this representation only at
-- specific types (not for any type 'a') as we can only have vectors of primitive type.
-- For other types, we use a structure-of-arrays representation. This is handled by the
-- ComplexRepr. We use the GADT ComplexR and function complexR to reconstruct
-- information on how the elements are represented.
--
instance Elt a => Elt (Complex a) where
  type EltRepr (Complex a) = ComplexRepr (EltRepr a)
  {-# INLINE eltType     #-}
  {-# INLINE [1] toElt   #-}
  {-# INLINE [1] fromElt #-}
  eltType = case complexR tp of
      ComplexRvec s -> TupRsingle $ VectorScalarType $ VectorType 2 s
      ComplexRtup   -> TupRunit `TupRpair` tp `TupRpair` tp
    where
      tp = eltType @a
  toElt = case complexR $ eltType @a of
    ComplexRvec _ -> \(Vec2 r i)   -> toElt r :+ toElt i
    ComplexRtup   -> \(((), r), i) -> toElt r :+ toElt i
  fromElt (r :+ i) = case complexR $ eltType @a of
    ComplexRvec _ -> Vec2 (fromElt r) (fromElt i)
    ComplexRtup   -> (((), fromElt r), fromElt i)

type family ComplexRepr a where
  ComplexRepr Half   = Vec2 Half
  ComplexRepr Float  = Vec2 Float
  ComplexRepr Double = Vec2 Double
  ComplexRepr Int    = Vec2 Int
  ComplexRepr Int8   = Vec2 Int8
  ComplexRepr Int16  = Vec2 Int16
  ComplexRepr Int32  = Vec2 Int32
  ComplexRepr Int64  = Vec2 Int64
  ComplexRepr Word   = Vec2 Word
  ComplexRepr Word8  = Vec2 Word8
  ComplexRepr Word16 = Vec2 Word16
  ComplexRepr Word32 = Vec2 Word32
  ComplexRepr Word64 = Vec2 Word64
  ComplexRepr a      = Tup2 a a

data ComplexR a c where
  ComplexRvec :: VecElt a => SingleType a -> ComplexR a (Vec2 a)
  ComplexRtup :: ComplexR a (Tup2 a a)

complexR :: TupleType a -> ComplexR a (ComplexRepr a)
complexR (TupRsingle (SingleScalarType (NumSingleType (FloatingNumType TypeHalf  )))) = ComplexRvec singleType
complexR (TupRsingle (SingleScalarType (NumSingleType (FloatingNumType TypeFloat )))) = ComplexRvec singleType
complexR (TupRsingle (SingleScalarType (NumSingleType (FloatingNumType TypeDouble)))) = ComplexRvec singleType
complexR (TupRsingle (SingleScalarType (NumSingleType (IntegralNumType TypeInt   )))) = ComplexRvec singleType
complexR (TupRsingle (SingleScalarType (NumSingleType (IntegralNumType TypeInt8  )))) = ComplexRvec singleType
complexR (TupRsingle (SingleScalarType (NumSingleType (IntegralNumType TypeInt16 )))) = ComplexRvec singleType
complexR (TupRsingle (SingleScalarType (NumSingleType (IntegralNumType TypeInt32 )))) = ComplexRvec singleType
complexR (TupRsingle (SingleScalarType (NumSingleType (IntegralNumType TypeInt64 )))) = ComplexRvec singleType
complexR (TupRsingle (SingleScalarType (NumSingleType (IntegralNumType TypeWord  )))) = ComplexRvec singleType
complexR (TupRsingle (SingleScalarType (NumSingleType (IntegralNumType TypeWord8 )))) = ComplexRvec singleType
complexR (TupRsingle (SingleScalarType (NumSingleType (IntegralNumType TypeWord16)))) = ComplexRvec singleType
complexR (TupRsingle (SingleScalarType (NumSingleType (IntegralNumType TypeWord32)))) = ComplexRvec singleType
complexR (TupRsingle (SingleScalarType (NumSingleType (IntegralNumType TypeWord64)))) = ComplexRvec singleType
complexR (TupRsingle (SingleScalarType (NonNumSingleType TypeChar)))                  = ComplexRtup
complexR (TupRsingle (SingleScalarType (NonNumSingleType TypeBool)))                  = ComplexRtup
complexR (TupRsingle (VectorScalarType (_)))                                          = ComplexRtup
complexR TupRunit                                                                     = ComplexRtup
complexR TupRpair{}                                                                   = ComplexRtup

constructComplex :: forall a. Elt a => Exp a -> Exp a -> Exp (Complex a)
constructComplex r i = case complexR $ eltType @a of
  ComplexRvec _ ->
    let
      r', i' :: Exp (EltRepr a)
      r' = coerce @a @(EltRepr a) r
      i' = coerce i
      v :: Exp (Vec2 (EltRepr a))
      v = V2 r' i'
    in
      coerce @(Vec2 (EltRepr a)) @(Complex a) $ v
  ComplexRtup   -> coerce $ T2  r i

deconstructComplex :: forall a. Elt a => Exp (Complex a) -> (Exp a, Exp a)
deconstructComplex c = case complexR $ eltType @a of
  ComplexRvec _ -> let V2 r i = coerce @(Complex a) @(Vec2 (EltRepr a)) c in (coerce r, coerce i)
  ComplexRtup   -> let T2 r i = coerce c in (r, i)

coerce :: EltRepr a ~ EltRepr b => Exp a -> Exp b
coerce (Exp e) = Exp e

instance (Lift Exp a, Elt (Plain a)) => Lift Exp (Complex a) where
  type Plain (Complex a) = Complex (Plain a)
  lift (r :+ i) = lift r ::+ lift i

instance Elt a => Unlift Exp (Complex (Exp a)) where
  unlift (r ::+ i) = r :+ i

instance Eq a => Eq (Complex a) where
  r1 ::+ c1 == r2 ::+ c2 = r1 == r2 && c1 == c2
  r1 ::+ c1 /= r2 ::+ c2 = r1 /= r2 || c1 /= c2

instance RealFloat a => P.Num (Exp (Complex a)) where
  (+)    = lift2 ((+) :: Complex (Exp a) -> Complex (Exp a) -> Complex (Exp a))
  (-)    = lift2 ((-) :: Complex (Exp a) -> Complex (Exp a) -> Complex (Exp a))
  (*)    = lift2 ((*) :: Complex (Exp a) -> Complex (Exp a) -> Complex (Exp a))
  negate = lift1 (negate :: Complex (Exp a) -> Complex (Exp a))
  signum z@(x ::+ y) =
    if z == 0
       then z
       else let r = magnitude z
             in x/r ::+ y/r
  abs z         = magnitude z ::+ 0
  fromInteger n = fromInteger n ::+ 0

instance RealFloat a => P.Fractional (Exp (Complex a)) where
  fromRational x  = fromRational x ::+ 0
  z / z'          = (x*x''+y*y'') / d ::+ (y*x''-x*y'') / d
    where
      x  :+ y   = unlift z
      x' :+ y'  = unlift z'
      --
      x'' = scaleFloat k x'
      y'' = scaleFloat k y'
      k   = - max (exponent x') (exponent y')
      d   = x'*x'' + y'*y''

instance RealFloat a => P.Floating (Exp (Complex a)) where
  pi                = pi ::+ 0
  exp (x ::+ y)     = let expx = exp x
                       in expx * cos y ::+ expx * sin y
  log z             = log (magnitude z) ::+ phase z
  sqrt z@(x ::+ y)  =
    if z == 0
      then 0
      else u ::+ (y < 0 ? (-v, v))
    where
      T2 u v = x < 0 ? (T2 v' u', T2 u' v')
      v'     = abs y / (u'*2)
      u'     = sqrt ((magnitude z + abs x) / 2)

  x ** y =
    if y == 0 then 1 else
    if x == 0 then if exp_r > 0 then 0 else
                   if exp_r < 0 then inf ::+ 0
                                else nan ::+ nan
              else if isInfinite r || isInfinite i
                     then if exp_r > 0 then inf ::+ 0 else
                          if exp_r < 0 then 0
                                       else nan ::+ nan
                     else exp (log x * y)
    where
      r     ::+ i  = x
      exp_r ::+ _  = y
      --
      inf = 1 / 0
      nan = 0 / 0

  sin (x ::+ y)  = sin x * cosh y ::+ cos x * sinh y
  cos (x ::+ y)  = cos x * cosh y ::+ (- sin x * sinh y)
  tan (x ::+ y)  = (sinx*coshy ::+ cosx*sinhy) / (cosx*coshy ::+ (-sinx*sinhy))
    where
      sinx  = sin x
      cosx  = cos x
      sinhy = sinh y
      coshy = cosh y

  sinh (x ::+ y) = cos y * sinh x ::+ sin  y * cosh x
  cosh (x ::+ y) = cos y * cosh x ::+ sin y * sinh x
  tanh (x ::+ y) = (cosy*sinhx ::+ siny*coshx) / (cosy*coshx ::+ siny*sinhx)
    where
      siny  = sin y
      cosy  = cos y
      sinhx = sinh x
      coshx = cosh x

  asin z@(x ::+ y) = y' ::+ (-x')
    where
      x' ::+ y' = log (((-y) ::+ x) + sqrt (1 - z*z))

  acos z                    = y'' ::+ (-x'')
    where
      x'' ::+ y''  = log (z + ((-y') ::+ x'))
      x'  ::+ y'   = sqrt (1 - z*z)

  atan z@(x ::+ y) = y' ::+ (-x')
    where
      x' ::+ y' = log (((1-y) ::+ x) / sqrt (1+z*z))

  asinh z =  log (z + sqrt (1+z*z))
  acosh z =  log (z + (z+1) * sqrt ((z-1)/(z+1)))
  atanh z =  0.5 * log ((1.0+z) / (1.0-z))

instance (FromIntegral a b, Num b, Elt (Complex b)) => FromIntegral a (Complex b) where
  fromIntegral x = fromIntegral x ::+ 0

-- | @since 1.2.0.0
--
instance Functor Complex where
  fmap f (r ::+ i) = f r ::+ f i

-- | The non-negative magnitude of a complex number
--
magnitude :: RealFloat a => Exp (Complex a) -> Exp a
magnitude (r ::+ i) = scaleFloat k (sqrt (sqr (scaleFloat mk r) + sqr (scaleFloat mk i)))
  where
    k     = max (exponent r) (exponent i)
    mk    = -k
    sqr z = z * z

-- | As 'magnitude', but ignore floating point rounding and use the traditional
-- (simpler to evaluate) definition.
--
-- @since 1.3.0.0
--
magnitude' :: RealFloat a => Exp (Complex a) -> Exp a
magnitude' (r ::+ i) = sqrt (r*r + i*i)

-- | The phase of a complex number, in the range @(-'pi', 'pi']@. If the
-- magnitude is zero, then so is the phase.
--
phase :: RealFloat a => Exp (Complex a) -> Exp a
phase z@(r ::+ i) =
  if z == 0
    then 0
    else atan2 i r

-- | The function 'polar' takes a complex number and returns a (magnitude,
-- phase) pair in canonical form: the magnitude is non-negative, and the phase
-- in the range @(-'pi', 'pi']@; if the magnitude is zero, then so is the phase.
--
polar :: RealFloat a => Exp (Complex a) -> Exp (a,a)
polar z =  T2 (magnitude z) (phase z)

-- | Form a complex number from polar components of magnitude and phase.
--
#if __GLASGOW_HASKELL__ <= 708
mkPolar :: forall a. RealFloat a => Exp a -> Exp a -> Exp (Complex a)
#else
mkPolar :: forall a. Floating a => Exp a -> Exp a -> Exp (Complex a)
#endif
mkPolar = lift2 (C.mkPolar :: Exp a -> Exp a -> Complex (Exp a))

-- | @'cis' t@ is a complex value with magnitude @1@ and phase @t@ (modulo
-- @2*'pi'@).
--
#if __GLASGOW_HASKELL__ <= 708
cis :: forall a. RealFloat a => Exp a -> Exp (Complex a)
#else
cis :: forall a. Floating a => Exp a -> Exp (Complex a)
#endif
cis = lift1 (C.cis :: Exp a -> Complex (Exp a))

-- | Return the real part of a complex number
--
real :: Elt a => Exp (Complex a) -> Exp a
real (r ::+ _) = r

-- | Return the imaginary part of a complex number
--
imag :: Elt a => Exp (Complex a) -> Exp a
imag (_ ::+ i) = i

-- | Return the complex conjugate of a complex number, defined as
--
-- > conjugate(Z) = X - iY
--
conjugate :: Num a => Exp (Complex a) -> Exp (Complex a)
conjugate z = real z ::+ (- imag z)


{-# LANGUAGE CPP #-}
{-# LANGUAGE KindSignatures #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE GADTs #-}

module HVX.Internal.Primitives
  ( Var
  , Fun(..)
  , Expr(..)
  , getFun
  , getJacobian
  , getProperties
  ) where

import Numeric.LinearAlgebra

import HVX.Internal.Matrix
import HVX.Internal.DCP

type Var = String

data Expr vex mon where
  EConst :: Mat -> Expr 'Affine 'Const
  EVar   :: Var -> Expr 'Affine 'Nondec
  EFun   ::
#ifndef DISABLE_EXPR_CXT
      ValidVex vex =>
#endif
      Fun v1 m1 -> Expr v2 m2 -> Expr vex mon
  EAdd   ::
#ifndef DISABLE_EXPR_CXT
      ValidVex vex =>
#endif
    Expr v1 m1 -> Expr v2 m2 -> Expr vex mon

instance Show (Expr vex mon) where
  show (EConst mat) = show mat
  show (EVar   var) = var
  show (EFun   f a) = show f ++ "(" ++ show a ++ ")"
  show (EAdd   a b) = show a ++ " + " ++ show b

getProperties :: (GetVex vex, GetMon mon) => Expr vex mon -> String
getProperties expr = getVex expr ++ " " ++ getMon expr

data Fun (vex :: Vex) (mon :: Mon) where
  Mul                :: Mat -> Fun 'Affine 'Nonmon
  Abs                :: Fun 'Convex 'Nonmon
  Neg                :: Fun 'Affine 'Noninc
  Log                :: Fun 'Concave 'Nondec
  Exp                :: Fun 'Convex 'Nondec
  LogSumExp          :: Fun 'Convex 'Nondec
  Max                :: Fun 'Convex 'Nondec
  Min                :: Fun 'Concave 'Nondec
  Norm               :: Double -> Fun 'Convex 'Nonmon
  Berhu              :: Double -> Fun 'Convex 'Nonmon
  Huber              :: Double -> Fun 'Convex 'Nonmon
  Quadform           :: Mat -> Fun 'Convex 'Nonmon
  PowBaseP0          :: Fun 'Affine 'Const
  PowBaseP01         :: Double -> Fun 'Concave 'Nondec
  PowBaseP1          :: Fun 'Affine 'Nondec
  PowBaseP1InfEven   :: Integer -> Fun 'Convex 'Nonmon
  PowBaseP1InfNotInt :: Double -> Fun 'Convex 'Nondec

instance Show (Fun vex mon) where
  show (Mul _) = "mul"
  show Abs = "abs"
  show Neg = "-"
  show Log = "log"
  show Exp = "exp"
  show LogSumExp = "log_sum_exp"
  show Max = "max"
  show Min = "min"
  show (Norm p) = "norm" ++ show p
  show (Berhu m) = "berhu" ++ show m
  show (Huber m) = "huber" ++ show m
  show (Quadform m) = "quadform" ++ show m
  show PowBaseP0 = "pow0"
  show (PowBaseP01 p) = "pow" ++ show p
  show PowBaseP1 = "pow1"
  show (PowBaseP1InfEven p) = "pow" ++ show p
  show (PowBaseP1InfNotInt p) = "pow" ++ show p

getFun :: Fun vex mon -> Mat -> Mat
getFun (Mul a) x = a <> x
getFun Abs x = abs x
getFun Neg x = negate x
getFun Exp x = exp x
getFun Log x
  | anyMat (<= 0) x = error "Cannot take log of number <= 0."
  | otherwise = log x
getFun LogSumExp x = reduceMat (log . sum) $ exp x
getFun Max x = reduceMat maximum x
getFun Min x = reduceMat minimum x
getFun (Norm p) x = scalarMat $ lpnorm p x
getFun (Berhu m) x =
  cmap (\y -> if abs y <= m then abs y else (abs y ** 2 + m ** 2) / (2 * m)) x
getFun (Huber m) x =
  cmap (\y -> if abs y <= m then abs y ** 2 else 2 * m * abs y - m ** 2) x
getFun (Quadform m) x = tr x <> m <> x
getFun PowBaseP0 x = konst 1.0 ((rows x), (cols x))
getFun (PowBaseP01 p) x
  | allMat (0 <=) x = matrixPow p x
  | otherwise = error "Cannot raise negative number to power p with 0 < p < 1."
getFun PowBaseP1 x = x
getFun (PowBaseP1InfEven p) x = matrixPow (fromIntegral p) x
getFun (PowBaseP1InfNotInt p) x
  | allMat (0 <=) x = matrixPow p x
  | otherwise = error "Cannot raise negative number to a nonintegral power of p for p > 1."

getJacobian :: Fun vex mon -> Mat -> Mat
getJacobian (Mul a) _ = a
getJacobian Abs x = diagMat $ signum x
getJacobian Neg x = (-1) * ident (rows x)
getJacobian Log x = diagMat $ 1 / x
getJacobian Exp x = diagMat $ exp x
getJacobian LogSumExp x = tr $ scale (1 / sumElements (exp x)) $ exp x
getJacobian Max x = tr $ ei (rows x) i
  where (i, _) = maxIndex x
getJacobian Min x = tr $ ei (rows x) i
  where (i, _) = minIndex x
getJacobian (Norm p) x = tr $ if x == zeroVec n
  then
    zeroVec n
  else
    scale (1 / denominator) numerator
  where denominator = lpnorm p x ** (p - 1)
        numerator = x * abs x ** (pMat - 2)
        pMat = (1><1) [p]
        n = rows x
getJacobian (Berhu m) x = diagMat $
  cmap (\y -> if abs y <= m then signum y else y / m) x
getJacobian (Huber m) x = diagMat $
  cmap (\y -> if abs y <= m then 2 * y else 2 * m * signum y) x
getJacobian (Quadform m) x = scale 2 $ tr x <> m
getJacobian PowBaseP0 x = konst 0.0 (n, n)
  where n = rows x
getJacobian (PowBaseP01 p) x = scale p $ diagMat $ matrixPow (p - 1) x
getJacobian PowBaseP1 x = ident n
  where n = rows x
getJacobian (PowBaseP1InfEven intP) x = scale p $ diagMat $ matrixPow (p - 1) x
  where p = fromIntegral intP
getJacobian (PowBaseP1InfNotInt p) x = scale p $ diagMat $ matrixPow (p - 1) x


{-# LANGUAGE Strict #-}
module FourierPinwheel.Multiplication where

import           Control.Concurrent.Async
import           Data.Array.Repa          as R
import           Data.Complex
import           Data.List                as L
import           Data.Vector.Storable     as VS
import           DFT.Plan
import           FourierPinwheel.Array
import           Utils.Parallel
import Filter.Utils
import Utils.Distribution

{-# INLINE multiplyBias #-}
multiplyBias ::
     DFTPlan
  -> VS.Vector (Complex Double)
  -> FPArray (VS.Vector (Complex Double))
  -> IO (FPArray (VS.Vector (Complex Double)))
multiplyBias plan biasF (FPArray numXFreq numYFreq numRFreq numThetaFreq numRhoFreq numPhiFreq vecs2) = do
  let planIDForward =
        DFTPlanID DFT1DG [numThetaFreq, numXFreq, numYFreq] [0, 1, 2]
      planIDBackward =
        DFTPlanID IDFT1DG [numThetaFreq, numXFreq, numYFreq] [0, 1, 2]
  vecs2F <- dftExecuteBatchP plan planIDForward vecs2
  let vecs3F = parMap rdeepseq (VS.zipWith (*) biasF) vecs2F
  FPArray numXFreq numYFreq numRFreq numThetaFreq numRhoFreq numPhiFreq <$>
    dftExecuteBatchP plan planIDBackward vecs3F

{-# INLINE multiplyBias4D #-}
multiplyBias4D ::
     DFTPlan
  -> VS.Vector (Complex Double)
  -> FPArray (VS.Vector (Complex Double))
  -> IO (FPArray (VS.Vector (Complex Double)))
multiplyBias4D plan biasF (FPArray numXFreq numYFreq numRFreq numThetaFreq numRhoFreq numPhiFreq vecs2) = do
  let planIDForward =
        DFTPlanID
          DFT1DG
          [numRFreq, numThetaFreq, numXFreq, numYFreq]
          [0, 1, 2, 3]
      planIDBackward =
        DFTPlanID
          IDFT1DG
          [numRFreq, numThetaFreq, numXFreq, numYFreq]
          [0, 1, 2, 3]
  vecs2F <- dftExecute plan planIDForward . VS.concat $ vecs2
  let vecs3F = VS.zipWith (*) biasF vecs2F
  arr <-
    fromUnboxed (Z :. numRFreq :. numThetaFreq :. numXFreq :. numYFreq) .
    VS.convert <$>
    dftExecute plan planIDBackward vecs3F
  return .
    FPArray numXFreq numYFreq numRFreq numThetaFreq numRhoFreq numPhiFreq .
    parMap
      rdeepseq
      (\rFreq ->
         VS.convert . toUnboxed . computeS . R.slice arr $ (Z :. rFreq :. All :. All :. All)) $
    [0 .. numRFreq - 1]

{-# INLINE multiplyBiasDiscrete #-}
multiplyBiasDiscrete ::
     DFTPlan
  -> VS.Vector (Complex Double)
  -> FPArray (VS.Vector (Complex Double))
  -> IO (FPArray (VS.Vector (Complex Double)))
multiplyBiasDiscrete plan bias (FPArray numXFreq numYFreq numRFreq numThetaFreq numRhoFreq numPhiFreq vecs2F) = do
  let planIDForward = DFTPlanID DFT1DG [numThetaFreq, numXFreq, numYFreq] [1, 2]
      planIDBackward =
        DFTPlanID IDFT1DG [numThetaFreq, numXFreq, numYFreq] [1, 2]
  vecs2 <-
    parMap rdeepseq (VS.zipWith (*) bias) <$>
    dftExecuteBatchP plan planIDBackward vecs2F
  FPArray numXFreq numYFreq numRFreq numThetaFreq numRhoFreq numPhiFreq <$>
    dftExecuteBatchP plan planIDForward vecs2

{-# INLINE multiplyRFunction #-}
multiplyRFunction ::
     DFTPlan
  -> Double
  -> FPArray (VS.Vector (Complex Double))
  -> IO (FPArray (VS.Vector (Complex Double)))
multiplyRFunction plan periodEnv (FPArray numXFreq numYFreq numRFreq numThetaFreq numRhoFreq numPhiFreq vecs) = do
  let freqConst = 2 * pi / log periodEnv
      filterVec =
        computeUnboxedS . makeFilter1D . fromFunction (Z :. numRFreq) $ \(Z :. i') ->
          let i = fromIntegral (i' - div numRFreq 2)
           in 1 / (1 :+ freqConst * fromIntegral i) -- gaussian1D i 10 :+ 0 --1 / (1 :+ (freqConst * fromIntegral i)^2)
      vec = VS.convert . VS.concat $ vecs
  filterF <-
    dftExecute plan (DFTPlanID DFT1DG [numRFreq] [0]) . VS.convert . toUnboxed $
    filterVec
  vecF <-
    dftExecute
      plan
      (DFTPlanID DFT1DG [numRFreq, numThetaFreq, numXFreq, numYFreq] [0])
      vec
  let arrF =
        fromUnboxed (Z :. numRFreq :. numThetaFreq :. numXFreq :. numYFreq) .
        VS.convert $
        vecF
      convolvedArrFR =
        R.traverse2
          arrF
          (fromUnboxed (Z :. numRFreq) . VS.convert $ filterF)
          const $ \fArr fVec idx@(Z :. r :. _ :. _ :. _) ->
          fArr idx * fVec (Z :. r)
  convolvedArrF <- computeUnboxedP convolvedArrFR
  convolvedVec <-
    dftExecute
      plan
      (DFTPlanID IDFT1DG [numRFreq, numThetaFreq, numXFreq, numYFreq] [0]) .
    VS.convert . toUnboxed $
    convolvedArrF
  let convolvedArr =
        fromUnboxed (Z :. numRFreq :. numThetaFreq :. numXFreq :. numYFreq) .
        VS.convert $
        convolvedVec
      convolvedVecs =
        parMap
          rdeepseq
          (\i ->
             VS.convert . toUnboxed . computeS . R.slice convolvedArr $
             (Z :. i :. All :. All :. All))
          [0 .. numRFreq - 1]
  return
    (FPArray
       numXFreq
       numYFreq
       numRFreq
       numThetaFreq
       numRhoFreq
       numPhiFreq
       convolvedVecs)

-- {-# INLINE multiplyRFunction #-}
-- multiplyRFunction :: FPArray (VS.Vector (Complex Double)) -> FPArray (VS.Vector (Complex Double))
-- multiplyRFunction (FPArray numXFreq numYFreq numRFreq numThetaFreq numRhoFreq numPhiFreq vecs') =
--   let vecs = L.reverse vecs'
--       zeroVec = VS.replicate (numThetaFreq * numXFreq * numYFreq) 0
--       convolvedSGNPositive = L.init . L.scanl' (VS.zipWith (+)) zeroVec $ vecs
--       convolvedSGNNegative = L.tail . L.scanr (VS.zipWith (+)) zeroVec $ vecs
--       convolvedSGN =
--         parZipWith
--           rdeepseq
--           (VS.zipWith (\x y -> (x - y) * (0 :+ (-pi))))
--           convolvedSGNPositive
--           convolvedSGNNegative
--    in FPArray numXFreq numYFreq numRFreq numThetaFreq numRhoFreq numPhiFreq .
--       L.reverse $
--       convolvedSGN


{-# LANGUAGE KindSignatures, DataKinds #-}
{-# LANGUAGE FlexibleContexts, FlexibleInstances, MultiParamTypeClasses #-}
{-# LANGUAGE UndecidableInstances #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE TypeFamilies #-}

module DSProg (
  module DSProg,
  Result
) where

import Prelude hiding ((<>))

import Control.Monad (join, (<=<), forM_, void, replicateM)
import Control.Monad.State (StateT, MonadState (get, put), evalStateT)
import qualified Control.Monad.State as State
import Control.Monad.Bayes.Class (MonadSample, MonadInfer)
import Control.Monad.Bayes.Sampler
import Control.Monad.Bayes.Weighted
import Control.Monad.Trans (liftIO)
import Control.Monad.Fix

import Data.Aeson (ToJSON (toJSON))
import Data.Maybe (isNothing, fromJust)
import qualified Data.Map
import Data.Proxy (Proxy (Proxy))

import GHC.TypeLits

import Numeric.LinearAlgebra.Static hiding (Gaussian, M)

import Unsafe.Coerce (unsafeCoerce)

import qualified Util.AssocList as M
import Util.AssocList (Map)
import qualified Util.AssocListSet as S
import Util.AssocListSet (Set)
import qualified Util.MStream as MS
import Util.MStream (MStream)
import Distributions
import MVDistributions
import DelayedSampling
import Inference (particles, Weight, zunheap)
import Util.Numeric (average)
import Util.Ref

import qualified Util.ZStream as ZS

import Numeric.Log (Log (Exp))

instance ToJSON a => ToJSON (Result a) where
  toJSON (RConst x) = toJSON x
  toJSON (RMarginal d) = toJSON d

instance MonadState s m => MonadState s (Weighted m) where
  get = State.lift get
  put = State.lift . put
  state = State.lift . State.state

meanMDistr :: Double ~ MType a => MDistr a -> Double
meanMDistr (MGaussian mu var) = mu
meanMDistr (MBeta a b) = a / (a + b)

meanResult :: Result Double -> Double
meanResult (RConst x) = x
meanResult (RMarginal m) = meanMDistr m

addResult :: Num a => Result a -> Result a -> Maybe (Result a)
addResult (RConst x) (RConst y) = Just $ RConst (x + y)
addResult (RConst x) (RMarginal (MGaussian mu var)) = Just $ RMarginal (MGaussian (mu + x) var)
addResult (RMarginal (MGaussian mu var)) (RConst x) = addResult (RConst x) (RMarginal (MGaussian mu var))
addResult (RMarginal (MGaussian mu var)) (RMarginal (MGaussian mu' var')) = Nothing
  --  Just $ RMarginal (MGaussian (mu + mu') (var + var')) -- only valid when variables are independent
addResult _ _ = Nothing

negateResult :: Num a => Result a -> Maybe (Result a)
negateResult (RConst x) = Just $ RConst (- x)
negateResult (RMarginal (MGaussian mu var)) = Just $ RMarginal (MGaussian (- mu) var)
negateResult _ = Nothing

multResult :: Num a => Result a -> Result a -> Maybe (Result a)
multResult (RConst x) (RConst y) = Just $ RConst (x * y)
multResult (RConst x) (RMarginal (MGaussian mu var)) = Just $ RMarginal (MGaussian (x * mu) (x^2 * var))
multResult (RMarginal (MGaussian mu var)) (RConst x) = multResult (RConst x) (RMarginal (MGaussian mu var))
multResult _ _ = Nothing

data Expr' var a where
  Const :: a -> Expr' var a
  Var :: var a -> Expr' var a
  Plus :: Num a => Expr' var a -> Expr' var a -> Expr' var a
  Times :: Num a => Expr' var a -> Expr' var a -> Expr' var a
  MVMul :: (KnownNat m, KnownNat n) => Expr' var (L m n) -> Expr' var (R n) -> Expr' var (R m)

instance Show (RefNodeToT a) where
  show (RefNodeToT r) = show r

deriving instance (Show a) => Show (Expr' RefNodeToT a)

class DeepForce e where
  type Forced e :: *
  deepForce :: MonadState Heap m => MonadSample m => e -> m (Forced e)
  deepConst :: Forced e -> e

deepForce' :: MonadState Heap m => MonadSample m => DeepForce e => e -> m e
deepForce' = fmap deepConst . deepForce

instance DeepForce Double where
  type Forced Double = Double
  deepForce = pure
  deepConst = id

instance DeepForce Int where
  type Forced Int = Int
  deepForce = pure
  deepConst = id

instance DeepForce (R n) where
  type Forced (R n) = R n
  deepForce = pure
  deepConst = id

instance DeepForce (L m n) where
  type Forced (L m n) = L m n
  deepForce = pure
  deepConst = id

instance DeepForce (Expr a) where
  type Forced (Expr a) = a
  deepForce = force
  deepConst = Const

instance DeepForce e => DeepForce [e] where
  type Forced [e] = [Forced e]
  deepForce = mapM deepForce
  deepConst = map deepConst

instance DeepForce e => DeepForce (Data.Map.Map k e) where
  type Forced (Data.Map.Map k e) = Data.Map.Map k (Forced e)
  deepForce = mapM deepForce
  deepConst = fmap deepConst

instance (DeepForce a, DeepForce b) => DeepForce (a, b) where
  type Forced (a, b) = (Forced a, Forced b)
  deepForce (x, y) = (,) <$> deepForce x <*> deepForce y
  deepConst (x, y) = (deepConst x, deepConst y)

instance (DeepForce a, DeepForce b, DeepForce c) => DeepForce (a, b, c) where
  type Forced (a, b, c) = (Forced a, Forced b, Forced c)
  deepForce (x, y, z) = (,,) <$> deepForce x <*> deepForce y <*> deepForce z
  deepConst (x, y, z) = (deepConst x, deepConst y, deepConst z)

instance (DeepForce a, DeepForce b, DeepForce c, DeepForce d) => DeepForce (a, b, c, d) where
  type Forced (a, b, c, d) = (Forced a, Forced b, Forced c, Forced d)
  deepForce (x, y, z, w) = (,,,) <$> deepForce x <*> deepForce y <*> deepForce z <*> deepForce w
  deepConst (x, y, z, w) = (deepConst x, deepConst y, deepConst z, deepConst w)

type M = Weighted (StateT Heap SamplerIO)

evalM :: M a -> SamplerIO (Double, a)
evalM x = (\(x, Exp w) -> (w, x)) <$> evalStateT (runWeighted x) emptyHeap

forgettableVar :: MonadState Heap m => Ref (Node a b) -> m (Expr (MType b))
forgettableVar nref = do
  --  addFinalizer v (forget nref)
  pure v
  where
  v = Var (RefNodeToT nref)

type Expr = Expr' RefNodeToT

data RefNodeToT b where
  RefNodeToT :: Ref (Node a b) -> RefNodeToT (MType b)

instance Eq (RefNodeToT a) where
  RefNodeToT (Ref x) == RefNodeToT (Ref y) = x == unsafeCoerce y

instance Num a => Num (Expr' var a) where
  fromInteger = Const . fromInteger
  x + y = Plus x y
  x * y = Times x y

fmapExpr' :: forall f var var' a. Applicative f => (forall x. var x -> f (var' x)) -> Expr' var a -> f (Expr' var' a)
fmapExpr' varMap = f where
  f :: forall a. Expr' var a -> f (Expr' var' a)
  f (Const x) = pure $ Const x
  f (Var i) = Var <$> varMap i
  f (Plus x y) = Plus <$> f x <*> f y
  f (Times x y) = Times <$> f x <*> f y
  f (MVMul a x) = MVMul <$> f a <*> f x

read :: MonadState Heap m => Expr' RefNodeToT a -> m (Expr' NodeTo a)
read = fmapExpr' (\(RefNodeToT nref) -> NodeTo <$> readRef nref)

data AffineMVMult (m :: Nat) where
  AffineMVMult :: KnownNat n => RefNodeToT (R n) -> L m n -> AffineMVMult m

getAffineMV :: KnownNat m => Expr (R m) -> Maybe ([AffineMVMult m], R m)
getAffineMV (Const x) = Just ([], x)
getAffineMV (Var i) = Just ([AffineMVMult i eye], 0)
getAffineMV (Plus x y) = do
  (m1, b1) <- getAffineMV x
  (m2, b2) <- getAffineMV y
  Just (m1 ++ m2, b1 + b2)
getAffineMV (Times x y) = Nothing
getAffineMV (MVMul (Const f) x) = do
  (m, b) <- getAffineMV x
  Just ([ AffineMVMult z (f <> g) | AffineMVMult z g <- m ], f #> b)
getAffineMV (MVMul _ _) = Nothing

getAffineMV1 :: KnownNat m => Expr (R m) -> Maybe (R m, Maybe (AffineMVMult m))
getAffineMV1 e = getAffineMV e >>= toAffine
  where
  toAffine :: KnownNat m => ([AffineMVMult m], R m) -> Maybe (R m, Maybe (AffineMVMult m))
  toAffine (m, b) = do
    mx <- case m of
      [] -> Just Nothing
      [x] -> Just (Just x)
      _ -> Nothing
    pure (b, mx)

getAffine :: Num a => Expr a -> Maybe (Map (Maybe (RefNodeToT a)) a)
getAffine (Const x) = Just (M.singleton Nothing x)
getAffine (Var i) = Just (M.singleton (Just i) 1)
getAffine (Plus x y) = M.unionWith (+) <$> getAffine x <*> getAffine y
getAffine (Times x y) = do
  mx <- getAffine x
  my <- getAffine y
  terms <- sequence [ mult x1 x2 | x1 <- M.toList mx, x2 <- M.toList my ]
  Just (M.fromList terms)
  where
  mult (Nothing, x) (Nothing, y) = Just (Nothing, x * y)
  mult (Nothing, x) (Just i, y) = Just (Just i, x * y)
  mult (Just i, x) (Nothing, y) = Just (Just i, x * y)
  mult (Just i, x) (Just j, y) = Nothing
getAffine (MVMul f x) = Nothing

getAffine1 :: Num a => Expr a -> Maybe (a, Maybe (RefNodeToT a, a))
getAffine1 e = getAffine e >>= toAffine

toAffine :: Num v => Map (Maybe k) v -> Maybe (v, Maybe (k, v))
toAffine mp = let (mconst, mvar) = M.partition Data.Maybe.isNothing mp in do
  let b = case mconst of
            M.Map [] -> 0
            M.Map [(Nothing, v)] -> v
            _ -> error "impossible"
  mx <- case mvar of
    M.Map [] -> Just Nothing
    M.Map [(Just k, v)] -> Just (Just (k, v))
    _ -> Nothing
  pure (b, mx)

realizedToConst' :: forall var f a. Applicative f => (forall x. var x -> f (Maybe x)) -> Expr' var a -> f (Expr' var a)
realizedToConst' varMap = f where
  f :: forall a. Expr' var a -> f (Expr' var a)
  f (Const x) = pure $ Const x
  f (Var i) = flip fmap (varMap i) $ \mk -> case mk of
      Nothing -> Var i
      Just k -> Const k
  f (Plus x y) = Plus <$> f x <*> f y
  f (Times x y) = Times <$> f x <*> f y
  f (MVMul a x) = MVMul <$> f a <*> f x

realizedToConst :: MonadState Heap m => Expr a -> m (Expr a)
realizedToConst = realizedToConst' f where
  f :: MonadState Heap m => RefNodeToT b -> m (Maybe b)
  f (RefNodeToT nref) = do
    n <- readRef nref
    pure $ case n of
      RealizedNode x -> Just x
      _ -> Nothing

allVars :: Expr a -> Set (SomeRefNode)
allVars (Const x) = S.empty
allVars (Var (RefNodeToT i)) = S.singleton (SomeRefNode i)
allVars (Plus x y) = allVars x `S.union` allVars y
allVars (Times x y) = allVars x `S.union` allVars y
allVars (MVMul f x) = allVars f `S.union` allVars x

initializedMarginal :: MonadState Heap m => Ref (Node a b) -> m (MDistr b)
initializedMarginal nref = do
  n <- readRef nref
  case distr n of
    UDistr d -> pure d
    CDistr par cd -> do
      Just pard <- marginalRef par
      pure (makeMarginal pard cd)

marginalRef :: MonadState Heap m => Ref (Node a b) -> m (Maybe (MDistr b))
marginalRef nref = do
  n <- readRef nref
  case n of
    RealizedNode x -> error "shouldn't occur"
    _ -> case state n of
      Marginalized d -> do
        isStale <- stale nref
        pure $ if isStale then Nothing else Just d
      Initialized -> Just <$> initializedMarginal nref

-- This is currently unsound
-- Consider
-- x <- normal(0, 1)
-- pure (x + (- x))
marginal :: MonadState Heap m => Expr a -> m (Maybe (Result a))
marginal (Const x) = pure (Just (RConst x))
marginal (Var (RefNodeToT nref)) = do
  n <- readRef nref
  case n of
    RealizedNode x -> pure $ Just (RConst x)
    _ -> case state n of
      Marginalized d -> do
        isStale <- stale nref
        pure $ if isStale then Nothing else Just (RMarginal d)
      Initialized -> Just . RMarginal <$> initializedMarginal nref
marginal (Plus x y) = pure Nothing -- could do independence ananlysis, etc.
marginal (Times x y) = pure Nothing -- could do independence ananlysis, etc.
marginal (MVMul f x) = pure Nothing

marginal' :: MonadState Heap m => Expr a -> m (Result a)
marginal' = fmap fromJust . marginal

data NodeTo a where
  NodeTo :: Node z a -> NodeTo (MType a)

force :: MonadSample m => MonadState Heap m => Expr a -> m a
force (Const x) = pure x
force (Var (RefNodeToT nref)) = getValue nref
force (Plus x y) = (+) <$> force x <*> force y
force (Times x y) = (*) <$> force x <*> force y
force (MVMul a x) = (#>) <$> force a <*> force x

typeOfRefNodeToT :: MonadState Heap m => Ref (Node a b) -> m (SMarginalT b)
typeOfRefNodeToT nref = typeOfDSDistr . distr <$> readRef nref

forgetE :: MonadState Heap m => Expr a -> m ()
forgetE e = forM_ (allVars e) forgetSomeRefNode

forgetSomeRefNode :: MonadState Heap m => SomeRefNode -> m ()
forgetSomeRefNode (SomeRefNode nref) = forget nref

justRun :: Show p => MonadState Heap m => State.MonadIO m => MStream m (Expr p) (Expr a) -> m (Maybe (Result a))
justRun s = do
  res <- MS.runStream (\x -> marginal x >>= liftIO . print) s
  res' <- marginal res
  pure res'

runParticle :: Int -> MStream M (Maybe (Result Double)) (Maybe (Result Double)) -> StateT Heap SamplerIO Double
runParticle numParticles =
  fmap avg . MS.runStream (liftIO . print . avg)
  . Inference.particles numParticles
  where
  avg = average . map (meanResult . fromJust)
  fromJust (Just x) = x
  fromJust Nothing = RConst $ Prelude.read "NaN"

zdeepForce :: DeepForce b => MonadSample m => ZS.ZStream (StateT Heap m) a b -> ZS.ZStream m a (Forced b)
zdeepForce model = zunheap $ ZS.zconstM deepForce ZS.<<< model

zdeepForce' :: DeepForce b => MonadSample m => ZS.ZStream (StateT Heap m) a b -> ZS.ZStream m a b
zdeepForce' model = fmap deepConst (zdeepForce model)


{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE FlexibleContexts #-}
module Math.HiddenMarkovModel.Distribution (
   State(..),
   Emission, Probability, Trained,
   Info(..), Generate(..), EmissionProb(..), Estimate(..),

   Discrete(..), DiscreteTrained(..),
   Gaussian(..), GaussianTrained(..), gaussian,

   CSV(..), HMMCSV.CSVParser, CSVSymbol(..),
   ) where

import qualified Math.HiddenMarkovModel.CSV as HMMCSV
import Math.HiddenMarkovModel.Utility (randomItemProp, normalizeProb)

import qualified Numeric.LinearAlgebra.HMatrix as HMatrix
import qualified Numeric.LinearAlgebra.Algorithms as Algo
import qualified Numeric.Container as NC
import qualified Data.Packed.Matrix as Matrix
import qualified Data.Packed.Vector as Vector
import Numeric.Container ((<>))
import Data.Packed.Matrix (Matrix)
import Data.Packed.Vector (Vector)
import Foreign.Storable (Storable)

import qualified System.Random as Rnd

import qualified Text.CSV.Lazy.String as CSV
import Text.Read.HT (maybeRead)
import Text.Printf (printf)

import qualified Control.Monad.Exception.Synchronous as ME
import qualified Control.Monad.Trans.Class as MT
import qualified Control.Monad.Trans.State as MS
import Control.DeepSeq (NFData, rnf)
import Control.Monad (liftM2)

import qualified Data.NonEmpty as NonEmpty
import qualified Data.Foldable as Fold
import qualified Data.Map as Map
import qualified Data.Set as Set
import qualified Data.Array as Array
import qualified Data.List as List
import Data.Foldable (foldMap)
import Data.Tuple.HT (mapFst)
import Data.Array (Array, Ix, listArray, (!))
import Data.Map (Map)
import Data.Maybe (listToMaybe)

newtype State = State Int
   deriving (Eq, Ord, Show, Read, Ix)

instance Enum State where
   toEnum = State
   fromEnum (State n) = n

instance NFData State where
   rnf (State n) = rnf n

type family Probability distr
type family Emission distr
type family Trained distr

class
   (NC.Container Vector (Probability distr), NC.Product (Probability distr)) =>
      Info distr where
   numberOfStates :: distr -> Int

class
   (NC.Container Vector (Probability distr), NC.Product (Probability distr)) =>
      Generate distr where
   generate ::
      (Rnd.RandomGen g, Probability distr ~ prob, Emission distr ~ emission) =>
      distr -> State -> MS.State g emission

class
   (NC.Container Vector (Probability distr), NC.Product (Probability distr)) =>
      EmissionProb distr where
   {-
   This function could be implemented generically in terms of emissionStateProb
   but that would require an Info constraint.
   -}
   emissionProb :: distr -> Emission distr -> Vector (Probability distr)
   emissionStateProb :: distr -> Emission distr -> State -> Probability distr
   emissionStateProb distr e (State s) = NC.atIndex (emissionProb distr e) s

class
   (EmissionProb (Distribution tdistr),
    Trained (Distribution tdistr) ~ tdistr) =>
      Estimate tdistr where
   type Distribution tdistr
   accumulateEmissions ::
      (Distribution tdistr ~ distr, Probability distr ~ prob) =>
      [[(Emission distr, prob)]] -> tdistr
   -- could as well be in Semigroup class
   combine :: tdistr -> tdistr -> tdistr
   normalize :: (Distribution tdistr ~ distr) => tdistr -> distr

newtype Discrete prob symbol = Discrete (Map symbol (Vector prob))
   deriving (Show)

newtype DiscreteTrained prob symbol = DiscreteTrained (Map symbol (Vector prob))
   deriving (Show)

type instance Probability (Discrete prob symbol) = prob
type instance Emission (Discrete prob symbol) = symbol

type instance Trained (Discrete prob symbol) = DiscreteTrained prob symbol

instance (NFData prob, NFData symbol) => NFData (Discrete prob symbol) where
   rnf (Discrete m) = rnf m

instance
   (NFData prob, NFData symbol) =>
      NFData (DiscreteTrained prob symbol) where
   rnf (DiscreteTrained m) = rnf m

instance
   (NC.Container Vector prob, NC.Product prob, Ord symbol) =>
      Info (Discrete prob symbol) where
   numberOfStates (Discrete m) = Vector.dim $ snd $ Map.findMin m

instance
   (NC.Container Vector prob, NC.Product prob, Ord symbol,
    Ord prob, Rnd.Random prob) =>
      Generate (Discrete prob symbol) where
   generate (Discrete m) (State state) =
      randomItemProp $ Map.toAscList $ fmap (flip NC.atIndex state) m

instance
   (NC.Container Vector prob, NC.Product prob, Ord symbol) =>
      EmissionProb (Discrete prob symbol) where
   emissionProb (Discrete m) =
      mapLookup "emitDiscrete: unknown emission symbol" m

instance
   (NC.Container Vector prob, NC.Product prob, Ord symbol) =>
      Estimate (DiscreteTrained prob symbol) where
   type Distribution (DiscreteTrained prob symbol) = Discrete prob symbol
   accumulateEmissions grouped =
      let set = Set.toAscList $ foldMap (Set.fromList . map fst) grouped
          emi = Map.fromAscList $ zip set [0..]
      in  DiscreteTrained $ Map.fromAscList $ zip set $
          transposeVectorList $
          map
             (NC.accum (NC.konst 0 (length set)) (+) .
              map (mapFst
                 (mapLookup "estimateDiscrete: unknown emission symbol" emi)))
             grouped
   combine (DiscreteTrained distr0) (DiscreteTrained distr1) =
      DiscreteTrained $ Map.unionWith NC.add distr0 distr1
   normalize (DiscreteTrained distr) =
      Discrete $ Map.fromAscList $ zip (Map.keys distr) $
      transposeVectorList $ map normalizeProb $
      transposeVectorList $ Map.elems distr

transposeVectorList :: (Matrix.Element a) => [Vector a] -> [Vector a]
transposeVectorList = Matrix.toRows . Matrix.fromColumns

mapLookup :: (Ord k) => String -> Map.Map k a -> k -> a
mapLookup msg dict x =
   Map.findWithDefault (error msg) x dict

newtype Gaussian a = Gaussian (Array State (Vector a, Matrix a, a))
   deriving (Show)

newtype GaussianTrained a = GaussianTrained (Map State (Vector a, Matrix a, a))
   deriving (Show)

type instance Probability (Gaussian a) = a
type instance Emission (Gaussian a) = Vector a

type instance Trained (Gaussian a) = GaussianTrained a

instance (NFData a, Storable a) => NFData (Gaussian a) where
   rnf (Gaussian params) = rnf params

instance (NFData a, Storable a) => NFData (GaussianTrained a) where
   rnf (GaussianTrained params) = rnf params

instance (Algo.Field a) => Info (Gaussian a) where
   numberOfStates (Gaussian params) = Array.rangeSize $ Array.bounds params

instance (Algo.Field a) => Generate (Gaussian a) where
   generate (Gaussian allParams) state = do
      let (center, covarianceChol, _c) = allParams ! state
      seed <- MS.state Rnd.random
      return $
         NC.add center $
         NC.cmap realToFrac
               (NC.randomVector seed NC.Gaussian (Vector.dim center))
            <> covarianceChol

instance (HMatrix.Numeric a, Algo.Field a) => EmissionProb (Gaussian a) where
   emissionProb (Gaussian allParams) x =
      Vector.fromList $ map (emissionProbGen x) $ Array.elems allParams
   emissionStateProb (Gaussian allParams) x s =
      emissionProbGen x $ allParams ! s

emissionProbGen ::
   (HMatrix.Numeric a, Algo.Field a) =>
   Vector a -> (Vector a, Matrix a, a) -> a
emissionProbGen =
   let cholSolve m x = Matrix.flatten $ Algo.cholSolve m $ Matrix.asColumn x
   in  \x (center, covarianceChol, c) ->
         let x0 = NC.sub x center
         in  c * exp ((-1/2) * NC.dot x0 (cholSolve covarianceChol x0))

instance (HMatrix.Numeric a, Algo.Field a) => Estimate (GaussianTrained a) where
   type Distribution (GaussianTrained a) = Gaussian a
   accumulateEmissions =
      let params xs =
             let center =
                    NonEmpty.foldl1Map NC.add (uncurry $ flip NC.scale) xs
                 covariance =
                    NonEmpty.foldl1Map NC.add (\(x,c) -> NC.scale c $ NC.outer x x) xs
             in  (center, covariance, Fold.sum $ fmap snd xs)
      in  GaussianTrained . fmap params . Map.mapMaybe NonEmpty.fetch .
          Map.fromList . zip [State 0 ..]
   combine (GaussianTrained distr0) (GaussianTrained distr1) =
      let comb (center0, covariance0, weight0)
               (center1, covariance1, weight1) =
             (NC.add center0 center1,
              NC.add covariance0 covariance1,
              weight0 + weight1)
      in  GaussianTrained $ Map.unionWith comb distr0 distr1
   {-
     Sum_i (xi-mi) * (xi-mi)^T
   = Sum_i xi*xi^T + Sum_i mi*mi^T - Sum_i xi*mi^T - Sum_i mi*xi^T
   = Sum_i xi*xi^T - Sum_i mi*mi^T
   = Sum_i xi*xi^T - n * mi*mi^T
   -}
   normalize (GaussianTrained distr) =
      let params (centerSum, covarianceSum, weight) =
             let c = recip weight
                 center = NC.scale c centerSum
             in  (center,
                  NC.sub (NC.scale c covarianceSum) (NC.outer center center))
      in  Gaussian $
          Array.array (fst $ Map.findMin distr, fst $ Map.findMax distr) $
          Map.toList $ fmap (gaussianParameters . params) distr

gaussian ::
   (Algo.Field prob) =>
   [(Vector prob, Matrix prob)] -> Gaussian prob
gaussian =
   consGaussian . map gaussianParameters

gaussianParameters ::
   (Algo.Field prob) =>
   (Vector prob, Matrix prob) -> (Vector prob, Matrix prob, prob)
gaussianParameters (center, covariance) =
   gaussianFromCholesky center $ Algo.chol covariance

consGaussian :: [(Vector a, Matrix a, a)] -> Gaussian a
consGaussian xs =
   Gaussian $ listArray (State 0, State $ length xs - 1) xs

gaussianFromCholesky ::
   (Algo.Field prob) =>
   Vector prob -> Matrix prob -> (Vector prob, Matrix prob, prob)
gaussianFromCholesky center covarianceChol =
   let covarianceSqrtDet = NC.prodElements $ Matrix.takeDiag covarianceChol
   in  (center, covarianceChol,
        recip (sqrt (2*pi) ^ Vector.dim center * covarianceSqrtDet))

class CSV distr where
   toCells :: distr -> [[String]]
   parseCells :: Int -> HMMCSV.CSVParser distr

class (Ord symbol) => CSVSymbol symbol where
   cellFromSymbol :: symbol -> String
   symbolFromCell :: String -> Maybe symbol

instance CSVSymbol Char where
   cellFromSymbol = (:[])
   symbolFromCell = listToMaybe

instance CSVSymbol Int where
   cellFromSymbol = show
   symbolFromCell = maybeRead

instance
   (Algo.Field prob, Show prob, Read prob, CSVSymbol symbol) =>
      CSV (Discrete prob symbol) where
   toCells (Discrete m) =
      map
         (\(symbol, probs) ->
            cellFromSymbol symbol : HMMCSV.cellsFromVector probs) $
      Map.toAscList m
   parseCells n =
      fmap (Discrete . Map.fromList) $
      HMMCSV.manyRowsUntilEnd $ parseSymbolProb n

parseSymbolProb ::
   (Algo.Field prob, Read prob, CSVSymbol symbol) =>
   Int -> CSV.CSVRow -> HMMCSV.CSVParser (symbol, Vector prob)
parseSymbolProb n row =
   case row of
      [] -> MT.lift $ ME.throw "missing symbol"
      c:cs ->
         liftM2 (,)
            (let str = CSV.csvFieldContent c
             in  MT.lift $ ME.fromMaybe (printf "unknown symbol %s" str) $
                 symbolFromCell str)
            (do v <- HMMCSV.parseVectorFields cs
                HMMCSV.assert (n == Vector.dim v)
                   (printf "number of states (%d) and size of probability vector (%d) mismatch"
                      n (Vector.dim v))
                return v)

instance (Algo.Field a, Eq a, Show a, Read a) => CSV (Gaussian a) where
   toCells (Gaussian params) =
      List.intercalate [[]] $
      map
         (\(center, covarianceChol, _) ->
            HMMCSV.cellsFromVector center :
            HMMCSV.cellsFromMatrix covarianceChol) $
      Array.elems params
   parseCells n = do
      gs <- HMMCSV.manySepUntilEnd parseSingleGaussian
      HMMCSV.assert (length gs == n) $
         printf "number of states (%d) and number of Gaussians (%d) mismatch"
            n (length gs)
      return $ consGaussian gs

parseSingleGaussian ::
   (Algo.Field prob, Eq prob, Read prob) =>
   HMMCSV.CSVParser (Vector prob, Matrix prob, prob)
parseSingleGaussian = do
   center <- HMMCSV.parseNonEmptyVectorCells
   covarianceChol <- HMMCSV.parseSquareMatrixCells $ Vector.dim center
   HMMCSV.assert (isUpperTriang covarianceChol) $
      "matrices must be upper triangular"
   return $ gaussianFromCholesky center covarianceChol

{-
Maybe this test is too strict.
It would also be ok, and certainly more intuitive
to use an orthogonal but not normalized matrix.
We could get such a matrix from the eigensystem.
-}
isUpperTriang :: (Algo.Field a, Eq a) => Matrix a -> Bool
isUpperTriang m =
   let cleared = Matrix.mapMatrixWithIndex (\(i,j) x -> if i>j then x else 0) m
   in  NC.minElement cleared == 0 &&
       NC.maxElement cleared == 0


{-# LANGUAGE DeriveDataTypeable #-}
{-# LANGUAGE PatternGuards #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE CPP #-}
{-# LANGUAGE TypeFamilies #-}
#if defined(__GLASGOW_HASKELL__) && __GLASGOW_HASKELL__ >= 702
{-# LANGUAGE Trustworthy #-}
{-# LANGUAGE DeriveGeneric #-}
#endif

#ifndef MIN_VERSION_vector
#define MIN_VERSION_vector(x,y,z) 1
#endif
-----------------------------------------------------------------------------
-- |
-- Copyright   :  (C) 2012-2015 Edward Kmett
-- License     :  BSD-style (see the file LICENSE)
--
-- Maintainer  :  Edward Kmett <ekmett@gmail.com>
-- Stability   :  experimental
-- Portability :  non-portable
--
-- Quaternions
----------------------------------------------------------------------------
module Linear.Quaternion
  ( Quaternion(..)
  , Complicated(..)
  , Hamiltonian(..)
  , ee, ei, ej, ek
  , slerp
  , asinq
  , acosq
  , atanq
  , asinhq
  , acoshq
  , atanhq
  , absi
  , pow
  , rotate
  , axisAngle
  ) where

import Control.Applicative
import Control.DeepSeq (NFData(rnf))
import Control.Monad (liftM)
import Control.Monad.Fix
import Control.Monad.Zip
import Control.Lens hiding ((<.>))
import Data.Binary as Binary
import Data.Bytes.Serial
import Data.Complex (Complex((:+)))
import Data.Data
import Data.Distributive
import Data.Foldable
import Data.Functor.Bind
import Data.Functor.Classes
import Data.Functor.Rep
import Data.Hashable
import Data.Serialize as Cereal
import GHC.Arr (Ix(..))
import qualified Data.Foldable as F
import Data.Monoid
import Foreign.Ptr (castPtr, plusPtr)
import Foreign.Storable (Storable(..))
#if defined(__GLASGOW_HASKELL__) && __GLASGOW_HASKELL__ >= 702
import GHC.Generics (Generic)
#endif
#if defined(__GLASGOW_HASKELL__) && __GLASGOW_HASKELL__ >= 706
import GHC.Generics (Generic1)
#endif
import qualified Data.Vector.Generic.Mutable as M
import qualified Data.Vector.Generic as G
import qualified Data.Vector.Unboxed.Base as U
import Linear.Epsilon
import Linear.Conjugate
import Linear.Metric
import Linear.V3
import Linear.Vector
import Prelude hiding (any)

{-# ANN module "HLint: ignore Reduce duplication" #-}

-- | Quaternions
data Quaternion a = Quaternion !a {-# UNPACK #-}!(V3 a)
                    deriving (Eq,Ord,Read,Show,Data,Typeable
#if defined(__GLASGOW_HASKELL__) && __GLASGOW_HASKELL__ >= 702
                             ,Generic
#endif
#if defined(__GLASGOW_HASKELL__) && __GLASGOW_HASKELL__ >= 706
                             ,Generic1
#endif
                             )

instance Functor Quaternion where
  fmap f (Quaternion e v) = Quaternion (f e) (fmap f v)
  {-# INLINE fmap #-}
  a <$ _ = Quaternion a (V3 a a a)
  {-# INLINE (<$) #-}

instance Apply Quaternion where
  Quaternion f fv <.> Quaternion a v = Quaternion (f a) (fv <.> v)
  {-# INLINE (<.>) #-}

instance Applicative Quaternion where
  pure a = Quaternion a (pure a)
  {-# INLINE pure #-}
  Quaternion f fv <*> Quaternion a v = Quaternion (f a) (fv <*> v)
  {-# INLINE (<*>) #-}

instance Additive Quaternion where
  zero = pure 0
  {-# INLINE zero #-}
  liftU2 = liftA2
  {-# INLINE liftU2 #-}
  liftI2 = liftA2
  {-# INLINE liftI2 #-}

instance Bind Quaternion where
  Quaternion a (V3 b c d) >>- f = Quaternion a' (V3 b' c' d') where
    Quaternion a' _          = f a
    Quaternion _ (V3 b' _ _) = f b
    Quaternion _ (V3 _ c' _) = f c
    Quaternion _ (V3 _ _ d') = f d
  {-# INLINE (>>-) #-}

instance Monad Quaternion where
  return = pure
  {-# INLINE return #-}
  -- the diagonal of a sedenion is super useful!
  Quaternion a (V3 b c d) >>= f = Quaternion a' (V3 b' c' d') where
    Quaternion a' _          = f a
    Quaternion _ (V3 b' _ _) = f b
    Quaternion _ (V3 _ c' _) = f c
    Quaternion _ (V3 _ _ d') = f d
  {-# INLINE (>>=) #-}

instance Ix a => Ix (Quaternion a) where
    {-# SPECIALISE instance Ix (Quaternion Int) #-}

    range (Quaternion l1 l2, Quaternion u1 u2) =
      [ Quaternion i1 i2 | i1 <- range (l1,u1), i2 <- range (l2,u2) ]
    {-# INLINE range #-}

    unsafeIndex (Quaternion l1 l2, Quaternion u1 u2) (Quaternion i1 i2) =
      unsafeIndex (l1,u1) i1 * unsafeRangeSize (l2,u2) + unsafeIndex (l2,u2) i2
    {-# INLINE unsafeIndex #-}

    inRange (Quaternion l1 l2, Quaternion u1 u2) (Quaternion i1 i2) =
      inRange (l1,u1) i1 && inRange (l2,u2) i2
    {-# INLINE inRange #-}

instance Representable Quaternion where
  type Rep Quaternion = E Quaternion
  tabulate f = Quaternion (f ee) (V3 (f ei) (f ej) (f ek))
  {-# INLINE tabulate #-}
  index xs (E l) = view l xs
  {-# INLINE index #-}

instance FunctorWithIndex (E Quaternion) Quaternion where
  imap f (Quaternion a (V3 b c d)) = Quaternion (f ee a) $ V3 (f ei b) (f ej c) (f ek d)
  {-# INLINE imap #-}

instance FoldableWithIndex (E Quaternion) Quaternion where
  ifoldMap f (Quaternion a (V3 b c d)) = f ee a `mappend` f ei b `mappend` f ej c `mappend` f ek d
  {-# INLINE ifoldMap #-}

instance TraversableWithIndex (E Quaternion) Quaternion where
  itraverse f (Quaternion a (V3 b c d)) = Quaternion <$> f ee a <*> (V3 <$> f ei b <*> f ej c <*> f ek d)
  {-# INLINE itraverse #-}

type instance Index (Quaternion a) = E Quaternion
type instance IxValue (Quaternion a) = a

instance Ixed (Quaternion a) where
  ix = el
  {-# INLINE ix #-}

instance Each (Quaternion a) (Quaternion b) a b where
  each = traverse
  {-# INLINE each #-}

instance Foldable Quaternion where
  foldMap f (Quaternion e v) = f e `mappend` foldMap f v
  {-# INLINE foldMap #-}
  foldr f z (Quaternion e v) = f e (F.foldr f z v)
  {-# INLINE foldr #-}

instance Traversable Quaternion where
  traverse f (Quaternion e v) = Quaternion <$> f e <*> traverse f v
  {-# INLINE traverse #-}

instance Storable a => Storable (Quaternion a) where
  sizeOf _ = 4 * sizeOf (undefined::a)
  {-# INLINE sizeOf #-}
  alignment _ = alignment (undefined::a)
  {-# INLINE alignment #-}
  poke ptr (Quaternion e v) = poke (castPtr ptr) e >>
                              poke (castPtr (ptr `plusPtr` sz)) v
    where sz = sizeOf (undefined::a)
  {-# INLINE poke #-}
  peek ptr = Quaternion <$> peek (castPtr ptr)
                        <*> peek (castPtr (ptr `plusPtr` sz))
    where sz = sizeOf (undefined::a)
  {-# INLINE peek #-}

instance RealFloat a => Num (Quaternion a) where
  {-# SPECIALIZE instance Num (Quaternion Float) #-}
  {-# SPECIALIZE instance Num (Quaternion Double) #-}
  (+) = liftA2 (+)
  {-# INLINE (+) #-}
  (-) = liftA2 (-)
  {-# INLINE (-) #-}
  negate = fmap negate
  {-# INLINE negate #-}
  Quaternion s1 v1 * Quaternion s2 v2 = Quaternion (s1*s2 - (v1 `dot` v2)) $
                                        (v1 `cross` v2) + s1*^v2 + s2*^v1
  {-# INLINE (*) #-}
  fromInteger x = Quaternion (fromInteger x) 0
  {-# INLINE fromInteger #-}
  abs z = Quaternion (norm z) 0
  {-# INLINE abs #-}
  signum q@(Quaternion e (V3 i j k))
    | m == 0.0 = q
    | not (isInfinite m || isNaN m) = q ^/ sqrt m
    | any isNaN q = qNaN
    | not (ii || ij || ik) = Quaternion 1 (V3 0 0 0)
    | not (ie || ij || ik) = Quaternion 0 (V3 1 0 0)
    | not (ie || ii || ik) = Quaternion 0 (V3 0 1 0)
    | not (ie || ii || ij) = Quaternion 0 (V3 0 0 1)
    | otherwise = qNaN
    where
      m = quadrance q
      ie = isInfinite e
      ii = isInfinite i
      ij = isInfinite j
      ik = isInfinite k
  {-# INLINE signum #-}

instance Hashable a => Hashable (Quaternion a) where
  hashWithSalt s (Quaternion a b) = s `hashWithSalt` a `hashWithSalt` b
  {-# INLINE hashWithSalt #-}

qNaN :: RealFloat a => Quaternion a
qNaN = Quaternion fNaN (V3 fNaN fNaN fNaN) where fNaN = 0/0
{-# INLINE qNaN #-}

-- {-# RULES "abs/norm" abs x = Quaternion (norm x) 0 #-}
-- {-# RULES "signum/signorm" signum = signorm #-}

-- this will attempt to rewrite calls to abs to use norm intead when it is available.

instance RealFloat a => Fractional (Quaternion a) where
  {-# SPECIALIZE instance Fractional (Quaternion Float) #-}
  {-# SPECIALIZE instance Fractional (Quaternion Double) #-}
  Quaternion q0 (V3 q1 q2 q3) / Quaternion r0 (V3 r1 r2 r3) =
    Quaternion (r0*q0+r1*q1+r2*q2+r3*q3)
               (V3 (r0*q1-r1*q0-r2*q3+r3*q2)
                   (r0*q2+r1*q3-r2*q0-r3*q1)
                   (r0*q3-r1*q2+r2*q1-r3*q0))
               ^/ (r0*r0 + r1*r1 + r2*r2 + r3*r3)
  {-# INLINE (/) #-}
  recip q = q ^/ quadrance q
  {-# INLINE recip #-}
  fromRational x = Quaternion (fromRational x) 0
  {-# INLINE fromRational #-}

instance Metric Quaternion where
  Quaternion e v `dot` Quaternion e' v' = e*e' + (v `dot` v')
  {-# INLINE dot #-}

-- | A vector space that includes the basis elements '_e' and '_i'
class Complicated t where
  _e, _i :: Lens' (t a) a

ee, ei :: Complicated t => E t
ee = E _e
ei = E _i

instance Complicated Complex where
  _e f (a :+ b) = (:+ b) <$> f a
  {-# INLINE _e #-}
  _i f (a :+ b) = (a :+) <$> f b
  {-# INLINE _i #-}

instance Complicated Quaternion where
  _e f (Quaternion a v) = (`Quaternion` v) <$> f a
  {-# INLINE _e #-}
  _i f (Quaternion a v) = Quaternion a <$> _x f v
  {-# INLINE _i #-}

-- | A vector space that includes the basis elements '_e', '_i', '_j' and '_k'
class Complicated t => Hamiltonian t where
  _j, _k :: Lens' (t a) a
  _ijk :: Lens' (t a) (V3 a)

ej, ek :: Hamiltonian t => E t
ej = E _j
ek = E _k

instance Hamiltonian Quaternion where
  _j f (Quaternion a v) = Quaternion a <$> _y f v
  {-# INLINE _j #-}
  _k f (Quaternion a v) = Quaternion a <$> _z f v
  {-# INLINE _k #-}
  _ijk f (Quaternion a v) = Quaternion a <$> f v
  {-# INLINE _ijk #-}

instance Distributive Quaternion where
  distribute f = Quaternion (fmap (\(Quaternion x _) -> x) f) $ V3
    (fmap (\(Quaternion _ (V3 y _ _)) -> y) f)
    (fmap (\(Quaternion _ (V3 _ z _)) -> z) f)
    (fmap (\(Quaternion _ (V3 _ _ w)) -> w) f)
  {-# INLINE distribute #-}

instance (Conjugate a, RealFloat a) => Conjugate (Quaternion a) where
  conjugate (Quaternion e v) = Quaternion (conjugate e) (negate v)
  {-# INLINE conjugate #-}

reimagine :: RealFloat a => a -> a -> Quaternion a -> Quaternion a
reimagine r s (Quaternion _ v)
  | isNaN s || isInfinite s = let aux 0 = 0
                                  aux x = s * x
                              in Quaternion r (aux <$> v)
  | otherwise = Quaternion r (v^*s)
{-# INLINE reimagine #-}

-- | quadrance of the imaginary component
qi :: Num a => Quaternion a -> a
qi (Quaternion _ v) = quadrance v
{-# INLINE qi #-}

-- | norm of the imaginary component
absi :: Floating a => Quaternion a -> a
absi = sqrt . qi
{-# INLINE absi #-}

-- | raise a 'Quaternion' to a scalar power
pow :: RealFloat a => Quaternion a -> a -> Quaternion a
pow q t = exp (t *^ log q)
{-# INLINE pow #-}

sqrte2pqiq :: (Floating a, Ord a) => a -> a -> a
sqrte2pqiq e qiq -- = sqrt (e*e) + qiq
  | e < - 1.5097698010472593e153 = -(qiq/e) - e
  | e < 5.582399551122541e57      = sqrt ((e*e) + qiq) -- direct definition
  | otherwise                     = (qiq/e) + e
-- {-# SPECIALIZE sqrte2pqiq :: Double -> Double -> Double #-}
-- {-# SPECIALIZE sqrte2pqiq :: Float -> Float -> Float #-}
#ifdef HERBIE
{-# ANN sqrte2pqiq "NoHerbie" #-}
#endif

tanrhs :: (Floating a, Ord a) => a -> a -> a -> a
tanrhs sai ai d -- = cosh ai * (sai / ai) / d -- improved from 6.04 bits of error to 0.19 bits
  | sai < -4.618902267687042e-52 = (sai / d / ai) * cosh ai
  | sai < 1.038530535935153e-39 = (cosh ai * sai) / ai / d
  | otherwise = (sai / d / ai) * cosh ai
-- {-# SPECIALIZE tanrhs :: Double -> Double -> Double -> Double #-}
-- {-# SPECIALIZE tanrhs :: Float -> Float -> Float -> Float #-}
#ifdef HERBIE
{-# ANN tanrhs "NoHerbie" #-}
#endif

-- ehh..
instance RealFloat a => Floating (Quaternion a) where
  {-# SPECIALIZE instance Floating (Quaternion Float) #-}
  {-# SPECIALIZE instance Floating (Quaternion Double) #-}
  pi = Quaternion pi 0
  {-# INLINE pi #-}
  exp q@(Quaternion e v)
    | qiq == 0 = Quaternion (exp e) v
    | ai <- sqrt qiq, exe <- exp e = reimagine (exe * cos ai) (exe * (sin ai / ai)) q
    where qiq = qi q
  {-# INLINE exp #-}
  log q@(Quaternion e v@(V3 _i j k))
    | qiq == 0 = if e >= 0
                 then Quaternion (log e) v
                 else Quaternion (log (negate e)) (V3 pi j k) -- mmm, pi
    | ai <- sqrt qiq = reimagine (log m) (atan2 m e / ai) q
    where qiq = qi q
          m = sqrte2pqiq e qiq
  {-# INLINE log #-}

  x ** y = exp (y * log x)
  {-# INLINE (**) #-}

  sqrt q@(Quaternion e v)
    | m   == 0 = q
    | qiq == 0 = if e > 0
                 then Quaternion (sqrt e) 0
                 else Quaternion 0 (V3 (sqrt (negate e)) 0 0)
    | im <- sqrt (0.5*(m-e)) / sqrt qiq = Quaternion (0.5*(m+e)) (v^*im)
    where qiq = qi q
          m = sqrte2pqiq e qiq
  {-# INLINE sqrt #-}

  cos q@(Quaternion e v)
    | qiq == 0 = Quaternion (cos e) v
    | ai <- sqrt qiq = reimagine (cos e * cosh ai) (- sin e / ai / sinh ai) q -- 0.15 bits error
                    -- reimagine (cos e * cosh ai) (- sin e * sinh ai / ai) q -- 13.5 bits worse
    where qiq = qi q
  {-# INLINE cos #-}

  sin q@(Quaternion e v)
    | qiq == 0 = Quaternion (sin e) v
    | ai <- sqrt qiq = reimagine (sin e * cosh ai) (cos e * sinh ai / ai) q
    where qiq = qi q
  {-# INLINE sin #-}

  tan q@(Quaternion e v)
    | qiq == 0 = Quaternion (tan e) v
    | ai <- sqrt qiq, ce <- cos e, sai <- sinh ai, d <- ce*ce + sai*sai =
      reimagine (ce * sin e / d) (tanrhs sai ai d) q
    where qiq = qi q
  {-# INLINE tan #-}

  sinh q@(Quaternion e v)
    | qiq == 0 = Quaternion (sinh e) v
    | ai <- sqrt qiq = reimagine (sinh e * cos ai) (cosh e * sin ai / ai) q
    where qiq = qi q
  {-# INLINE sinh #-}

  cosh q@(Quaternion e v)
    | qiq == 0 = Quaternion (cosh e) v
    | ai <- sqrt qiq = reimagine (cosh e * cos ai) (sin ai * (sinh e / ai)) q
    where qiq = qi q
  {-# INLINE cosh #-}

  tanh q@(Quaternion e v)
    | qiq == 0 = Quaternion (tanh e) v
    | ai <- sqrt qiq, se <- sinh e, cai <- cos ai, d <- se*se + cai*cai =
      reimagine (cosh e * se / d) (tanhrhs cai ai d) q
    where qiq = qi q
  {-# INLINE tanh #-}

  asin = cut asin
  {-# INLINE asin #-}
  acos = cut acos
  {-# INLINE acos #-}
  atan = cut atan
  {-# INLINE atan #-}

  asinh = cut asinh
  {-# INLINE asinh #-}
  acosh = cut acosh
  {-# INLINE acosh #-}
  atanh = cut atanh
  {-# INLINE atanh #-}

tanhrhs :: (Floating a, Ord a) => a -> a -> a -> a
tanhrhs cai ai d -- = cai * (sin ai / ai) / d
  | d >= -4.2173720203427147e-29 && d < 4.446702369113811e64 = cai / (d * (ai / sin ai))
  | otherwise                                                = cai * (1 / ai / sin ai) / d
-- {-# SPECIALIZE tanhrhs :: Double -> Double -> Double -> Double #-}
-- {-# SPECIALIZE tanhrhs :: Float -> Float -> Float -> Float #-}
#ifdef HERBIE
{-# ANN tanhrhs "NoHerbie" #-}
#endif

-- | Helper for calculating with specific branch cuts
cut :: RealFloat a => (Complex a -> Complex a) -> Quaternion a -> Quaternion a
cut f q@(Quaternion e (V3 _ y z))
  | qiq == 0 = Quaternion a (V3 b y z)
  | otherwise = reimagine a (b / ai) q
  where qiq = qi q
        ai = sqrt qiq
        a :+ b = f (e :+ ai)
{-# INLINE cut #-}

-- | Helper for calculating with specific branch cuts
cutWith :: RealFloat a => Complex a -> Quaternion a -> Quaternion a
cutWith (r :+ im) q@(Quaternion e v)
  | e /= 0 || qiq == 0 || isNaN qiq || isInfinite qiq = error "bad cut"
  | s <- im / sqrt qiq = Quaternion r (v^*s)
  where qiq = qi q
{-# INLINE cutWith #-}

-- | 'asin' with a specified branch cut.
asinq :: RealFloat a => Quaternion a -> Quaternion a -> Quaternion a
asinq q@(Quaternion e _) u
  | qiq /= 0.0 || e >= -1 && e <= 1 = asin q
  | otherwise = cutWith (asin (e :+ sqrt qiq)) u
  where qiq = qi q
{-# INLINE asinq #-}

-- | 'acos' with a specified branch cut.
acosq :: RealFloat a => Quaternion a -> Quaternion a -> Quaternion a
acosq q@(Quaternion e _) u
  | qiq /= 0.0 || e >= -1 && e <= 1 = acos q
  | otherwise = cutWith (acos (e :+ sqrt qiq)) u
  where qiq = qi q
{-# INLINE acosq #-}

-- | 'atan' with a specified branch cut.
atanq :: RealFloat a => Quaternion a -> Quaternion a -> Quaternion a
atanq q@(Quaternion e _) u
  | e /= 0.0 || qiq >= -1 && qiq <= 1 = atan q
  | otherwise = cutWith (atan (e :+ sqrt qiq)) u
  where qiq = qi q
{-# INLINE atanq #-}

-- | 'asinh' with a specified branch cut.
asinhq :: RealFloat a => Quaternion a -> Quaternion a -> Quaternion a
asinhq q@(Quaternion e _) u
  | e /= 0.0 || qiq >= -1 && qiq <= 1 = asinh q
  | otherwise = cutWith (asinh (e :+ sqrt qiq)) u
  where qiq = qi q
{-# INLINE asinhq #-}

-- | 'acosh' with a specified branch cut.
acoshq :: RealFloat a => Quaternion a -> Quaternion a -> Quaternion a
acoshq q@(Quaternion e _) u
  | qiq /= 0.0 || e >= 1 = asinh q
  | otherwise = cutWith (acosh (e :+ sqrt qiq)) u
  where qiq = qi q
{-# INLINE acoshq #-}

-- | 'atanh' with a specified branch cut.
atanhq :: RealFloat a => Quaternion a -> Quaternion a -> Quaternion a
atanhq q@(Quaternion e _) u
  | qiq /= 0.0 || e > -1 && e < 1 = atanh q
  | otherwise = cutWith (atanh (e :+ sqrt qiq)) u
  where qiq = qi q
{-# INLINE atanhq #-}

-- | Spherical linear interpolation between two quaternions.

slerp :: RealFloat a => Quaternion a -> Quaternion a -> a -> Quaternion a
slerp q p t
  | 1.0 - cosphi < 1e-8 = q
  | otherwise           = ((sin ((1-t)*phi) *^ q) + sin (t*phi) *^ f p) ^/ sin phi
  where
    dqp = dot q p
    (cosphi, f) = if dqp < 0 then (-dqp, negate) else (dqp, id)
    phi = acos cosphi
{-# SPECIALIZE slerp :: Quaternion Float -> Quaternion Float -> Float -> Quaternion Float #-}
{-# SPECIALIZE slerp :: Quaternion Double -> Quaternion Double -> Double -> Quaternion Double #-}

-- | Apply a rotation to a vector.
rotate :: (Conjugate a, RealFloat a) => Quaternion a -> V3 a -> V3 a
rotate q v = ijk where
  Quaternion _ ijk = q * Quaternion 0 v * conjugate q
{-# SPECIALIZE rotate :: Quaternion Float -> V3 Float -> V3 Float #-}
{-# SPECIALIZE rotate :: Quaternion Double -> V3 Double -> V3 Double #-}

instance (RealFloat a, Epsilon a) => Epsilon (Quaternion a) where
  nearZero = nearZero . quadrance
  {-# INLINE nearZero #-}

-- | @'axisAngle' axis theta@ builds a 'Quaternion' representing a
-- rotation of @theta@ radians about @axis@.
axisAngle :: (Epsilon a, Floating a) => V3 a -> a -> Quaternion a
axisAngle axis theta = Quaternion (cos half) (sin half *^ normalize axis)
  where half = theta / 2
{-# INLINE axisAngle #-}

data instance U.Vector    (Quaternion a) =  V_Quaternion !Int (U.Vector    a)
data instance U.MVector s (Quaternion a) = MV_Quaternion !Int (U.MVector s a)
instance U.Unbox a => U.Unbox (Quaternion a)

instance U.Unbox a => M.MVector U.MVector (Quaternion a) where
  basicLength (MV_Quaternion n _) = n
  basicUnsafeSlice m n (MV_Quaternion _ v) = MV_Quaternion n (M.basicUnsafeSlice (4*m) (4*n) v)
  basicOverlaps (MV_Quaternion _ v) (MV_Quaternion _ u) = M.basicOverlaps v u
  basicUnsafeNew n = liftM (MV_Quaternion n) (M.basicUnsafeNew (4*n))
  basicUnsafeRead (MV_Quaternion _ v) i =
    do let o = 4*i
       x <- M.basicUnsafeRead v o
       y <- M.basicUnsafeRead v (o+1)
       z <- M.basicUnsafeRead v (o+2)
       w <- M.basicUnsafeRead v (o+3)
       return (Quaternion x (V3 y z w))
  basicUnsafeWrite (MV_Quaternion _ v) i (Quaternion x (V3 y z w)) =
    do let o = 4*i
       M.basicUnsafeWrite v o     x
       M.basicUnsafeWrite v (o+1) y
       M.basicUnsafeWrite v (o+2) z
       M.basicUnsafeWrite v (o+3) w
#if MIN_VERSION_vector(0,11,0)
  basicInitialize (MV_Quaternion _ v) = M.basicInitialize v
#endif

instance U.Unbox a => G.Vector U.Vector (Quaternion a) where
  basicUnsafeFreeze (MV_Quaternion n v) = liftM ( V_Quaternion n) (G.basicUnsafeFreeze v)
  basicUnsafeThaw   ( V_Quaternion n v) = liftM (MV_Quaternion n) (G.basicUnsafeThaw   v)
  basicLength       ( V_Quaternion n _) = n
  basicUnsafeSlice m n (V_Quaternion _ v) = V_Quaternion n (G.basicUnsafeSlice (4*m) (4*n) v)
  basicUnsafeIndexM (V_Quaternion _ v) i =
    do let o = 4*i
       x <- G.basicUnsafeIndexM v o
       y <- G.basicUnsafeIndexM v (o+1)
       z <- G.basicUnsafeIndexM v (o+2)
       w <- G.basicUnsafeIndexM v (o+3)
       return (Quaternion x (V3 y z w))

instance MonadZip Quaternion where
  mzipWith = liftA2

instance MonadFix Quaternion where
  mfix f = Quaternion (let Quaternion a _ = f a in a)
                      (V3 (let Quaternion _ (V3 a _ _) = f a in a)
                          (let Quaternion _ (V3 _ a _) = f a in a)
                          (let Quaternion _ (V3 _ _ a) = f a in a))

instance NFData a => NFData (Quaternion a) where
  rnf (Quaternion a b) = rnf a `seq` rnf b

instance Serial1 Quaternion where
  serializeWith f (Quaternion a b) = f a >> serializeWith f b
  deserializeWith f = Quaternion <$> f <*> deserializeWith f

instance Serial a => Serial (Quaternion a) where
  serialize = serializeWith serialize
  deserialize = deserializeWith deserialize

instance Binary a => Binary (Quaternion a) where
  put = serializeWith Binary.put
  get = deserializeWith Binary.get

instance Serialize a => Serialize (Quaternion a) where
  put = serializeWith Cereal.put
  get = deserializeWith Cereal.get

instance Eq1 Quaternion where eq1 = (==)
instance Ord1 Quaternion where compare1 = compare
instance Show1 Quaternion where showsPrec1 = showsPrec
instance Read1 Quaternion where readsPrec1 = readsPrec


{-# LANGUAGE AllowAmbiguousTypes  #-}
{-# LANGUAGE UndecidableInstances #-}
{-# OPTIONS_GHC -fno-warn-deprecations #-}

module Numeric.Tensile.Tensor.Internal where

import Data.Bits
import Data.ByteString (ByteString())
import Data.Complex
import Data.Int
import Data.Kind
import Data.Proxy
import Data.Word
import Data.Vector.Storable (Storable(..))
import Unsafe.Coerce (unsafeCoerce)

import Numeric.Tensile.Dimensions

import Data.Vector.Sized (Vector)
import TensorFlow.Types
import System.IO.Unsafe (unsafePerformIO)

import qualified Data.Finite as F
import qualified Data.Vector as V
import qualified Data.Vector.Sized as N
import qualified Data.Vector.Storable as S
import qualified Data.Vector.Storable.Mutable as M
import qualified TensorFlow.Tensor as T
import qualified TensorFlow.Build as T
import qualified TensorFlow.GenOps.Core as O
import qualified TensorFlow.Ops as O
import qualified TensorFlow.Session as Ss

import Control.Monad.IO.Class (liftIO)
import Control.Monad.Trans.Reader

{-
fromTensor :: Elt e => Tensor d e -> [e]
fromTensor t = unsafePerformIO $ Ss.runSession $ do
  a <- T.render $ unTensor t
  b <- Ss.run a
  return $ V.toList b

:: OneOf '[Int32, Int64] tidx
=> Tensor v'1 Bool
input

-> Tensor v'2 tidx
reduction_indices

-> Tensor Build Bool

transpose
  :: Elt e
  => Permutable d d'
  => Dims d -> Perm (Rank d) -> Tensor d e -> Tensor d' e
transpose d (Perm p) (Tensor t) = Tensor $ (flip O.transpose) (O.vector w) t
  where v = [1.. fromIntegral $ rank d] :: [IVal]
        w = P.permuteList p v

-}
-- | Generally specialized as in 'TVal' or 'IVal'.
type Elt e = OneOf '[Int32, Int64, Word32, Word64, Float, Double, Complex Float, Complex Double] e

-- TODO: move to application / test stanza
--type Elt = TensorType -- TODO deal with Variant and exclude ResourceHandle
type TVal = Float
type IVal = Int32
type BVal = Bool

type Dynamic d = Reader (Dims d)

newtype Tensor (d :: [Nat]) (e :: Type) = Tensor { unTensor :: T.Tensor T.Build e }

-- | A real or complex-valued tensor of shape 'd'.
type T d = Tensor d TVal

-- | An integer or non-negative integer-valued tensor of shape 'd'.
type I d = Tensor d IVal

-- | A boolean-valued tensor of shape 'd'.
type B d = Tensor d BVal

--TODO update Show instance
instance (Elt e, Show e) => Show (Tensor d e) where show _ = "No Show instance."

--TODO implement
instance (KnownDims d, Elt e, Eq e) => Eq (Tensor d e) where
    (==) ta tb = all $ ta `eq` tb
      where all :: B d -> Bool
            all (Tensor b) = Prelude.head $ toList $ Tensor $ O.all b (O.vector ([] :: [IVal]))

instance (KnownDims d, Num e, Elt e) => Num (Tensor d e) where
    (+) (Tensor v1) (Tensor v2) = Tensor $ (+) v1 v2
    {-# INLINE (+) #-}
    (-) (Tensor v1) (Tensor v2) = Tensor $ (-) v1 v2
    {-# INLINE (-) #-}
    (*) (Tensor v1) (Tensor v2) = Tensor $ (*) v1 v2
    {-# INLINE (*) #-}
    negate (Tensor v1) = Tensor $ negate v1
    {-# INLINE negate #-}
    abs (Tensor v1) = Tensor $ abs v1
    {-# INLINE abs #-}
    signum (Tensor v1) = Tensor $ signum v1
    {-# INLINE signum #-}
    fromInteger = constant (dims @d) . fromInteger
    {-# INLINE fromInteger #-}

{-
instance (KnownDim (Size d), Fractional e, Elt e) => Fractional (Tensor d e)  where
    (/) = _lift2 (/)
    {-# INLINE (/) #-}
    recip = _lift recip
    {-# INLINE recip #-}
    fromRational = _replicate (fromIntegral . dimVal $ (dim :: Dim (Size d))) . fromRational
    {-# INLINE fromRational #-}

instance (KnownDim (Size d), Floating e, Elt e) => Floating (Tensor d e) where
    pi = Tensor $ S.singleton pi  --TODO make this dim safe
    {-# INLINE pi #-}
    exp = _lift exp
    {-# INLINE exp #-}
    log = _lift log
    {-# INLINE log #-}
    sqrt = _lift sqrt
    {-# INLINE sqrt #-}
    sin = _lift sin
    {-# INLINE sin #-}
    cos = _lift cos
    {-# INLINE cos #-}
    tan = _lift tan
    {-# INLINE tan #-}
    asin = _lift asin
    {-# INLINE asin #-}
    acos = _lift acos
    {-# INLINE acos #-}
    atan = _lift atan
    {-# INLINE atan #-}
    sinh = _lift sinh
    {-# INLINE sinh #-}
    cosh = _lift cosh
    {-# INLINE cosh #-}
    tanh = _lift tanh
    {-# INLINE tanh #-}
    (**) = _lift2 (**)
    {-# INLINE (**) #-}
    logBase = _lift2 logBase
    {-# INLINE logBase #-}
    asinh = _lift asinh
    {-# INLINE asinh #-}
    acosh = _lift acosh
    {-# INLINE acosh #-}
    atanh = _lift atanh
    {-# INLINE atanh #-}

instance (KnownDim (Size d), Elt e, Eq e, Bits e, Num e) => Bits (Tensor d e) where
    (.&.) = _lift2 (.&.)
    {-# INLINE (.&.) #-}
    (.|.) = _lift2 (.|.)
    {-# INLINE (.|.) #-}
    xor = _lift2 xor
    {-# INLINE xor #-}
    complement = _lift complement
    shift t i = _lift (flip shift i) t
    rotate t i = _lift (flip rotate i) t
    bit = _replicate (fromIntegral . dimVal $ (dim :: Dim (Size d))) . bit
    testBit = testBitDefault
    bitSizeMaybe _ = bitSizeMaybe @e undefined
    bitSize _ = bitSize @e undefined
    isSigned _ = isSigned @e undefined
    popCount = popCountDefault
-}

{-

instance (KnownDim (Size d), Elt e, Real e) => Real (Tensor d e) where
  toRational = undefined --TODO find a reasonable sum-based implementation or scrap the typeclass

instance (KnownDim (Size d), Elt e, Enum e) => Enum (Tensor d e) where
  toEnum = Tensor . S.replicate (fromIntegral . dimVal $ (dim :: Dim (Size d))) . toEnum
  {-# INLINE toEnum #-}
  fromEnum = undefined --TODO find a reasonable sum-based implementation or scrap the typeclass

instance (KnownDim (Size d), Elt e, Integral e) => Integral (Tensor d e) where
  quot (Tensor a) (Tensor b) = _lift2 quot a b
  rem (Tensor a) (Tensor b) = _lift2 rem a b
  div (Tensor a) (Tensor b) = _lift2 div a b
  mod (Tensor a) (Tensor b) = _lift2 mod a b
  quotRem ta tb = (quot ta tb, rem ta tb)
  divMod ta tb = (div ta tb, mod ta tb)
  toInteger _ = undefined --TODO find a reasonable sum-based implementation or scrap the typeclass

-}

eq :: Elt e => Eq e => Tensor d e -> Tensor d e -> Tensor d BVal
eq (Tensor a) (Tensor b) = Tensor $ O.equal a b

neq :: Elt e => Eq e => Tensor d e -> Tensor d e -> Tensor d BVal
neq (Tensor a) (Tensor b) = Tensor $ O.notEqual a b

lt :: Elt e => Ord e => Tensor d e -> Tensor d e -> Tensor d BVal
--lt :: Ord TVal => T d -> T d -> B d
lt (Tensor a) (Tensor b) = Tensor $ O.less a b

lte :: Elt e => Ord e => Tensor d e -> Tensor d e -> Tensor d BVal
lte (Tensor a) (Tensor b) = Tensor $ O.lessEqual a b

gt :: Elt e => Ord e => Tensor d e -> Tensor d e -> Tensor d BVal
gt (Tensor a) (Tensor b) = Tensor $ O.greater a b

gte :: Elt e => Ord e => Tensor d e -> Tensor d e -> Tensor d BVal
gte (Tensor a) (Tensor b) = Tensor $ O.greaterEqual a b

maximum :: Elt e => Ord e => Tensor d e -> Tensor d e -> Tensor d e
maximum (Tensor a) (Tensor b) = Tensor $ O.maximum a b

minimum :: Elt e => Ord e => Tensor d e -> Tensor d e -> Tensor d e
minimum (Tensor a) (Tensor b) = Tensor $ O.minimum a b

toList :: TensorDataType V.Vector a => Tensor d a -> [a]
toList t = unsafePerformIO $ Ss.runSession $ do
  a <- T.render $ unTensor t
  b <- Ss.run a
  return $ V.toList b

fromList :: forall d e. Elt e => Dims d -> [e] -> Maybe (Tensor d e)
fromList d l = if size d == (fromIntegral $ length l) then Just $ f l else Nothing
  where f = Tensor . O.constant (toShape d)

fromList'' :: forall d e. Elt e => Dynamic d ([e] -> Maybe (Tensor d e))
fromList'' = reader fromList

fromList' :: forall d e. Elt e => KnownDims d => [e] -> Maybe (Tensor d e)
fromList' = withDyn @d fromList''

-- reflectDims @d $ \d -> runReader fromList'' d $ l

liftDyn :: (Dims d -> r) -> Dynamic d r
liftDyn f = reader f

withDyn :: forall d r. KnownDims d => Dynamic d r -> r
withDyn dyn = reflectDims @d $ runReader dyn

{-
 - http://hackage.haskell.org/package/contravariant-1.5/docs/Data-Functor-Contravariant.html
 -
=> f (Dims d)

TODO:
 - do lowerPerm and withDyn have anything in common? (e.g. operator constraints 'lowering' to Predicate (Dims d)). both are functors?
 - type-level Predicate for '[Nat] ? Equivalence '[Nat] ~ ('[Nat],'[Nat]) ~> Bool   Could also include >1 req?
 -
 -
withDyn :: forall d r. KnownDims d => Predicate (Dims d) -> Dynamic d r -> Maybe r

withDyn2
  :: forall x y r. (KnownDims x, KnownDims y)
  => Predicate (Dims x,)
  => Predicate (Dims y)
  -> Dynamic d r -> Maybe r

divided :: Divisible f => f a -> f b -> f (a, b) Source#
divided = divide id

Predicate (Tensor d e)
Equivalence (Tensor d e)
Predicate (Dims d)
Equivalence (Dims d)

divided :: Divisible f => f a -> f b -> f (a, b)

note in general we have

Dims x -> Bool
Dims x -> Dims y -> Bool

type Dynamic d = Reader (Dims d)

fromList :: forall d e. Elt e => KnownDims d => [e] -> Maybe (Tensor d e)
fromList v = if size (dims @d) == (fromIntegral $ length v) then Just $ f v else Nothing
  where f = Tensor . O.constant (toShape (dims @d))

i = 0 :+ 1 :+ 2 :+ S :: Idxs '[1,2,3]
j = 0 :+ 1 :+ 1 :+ S :: Idxs '[1,2,3]

-- fromVector :: Elt e => Dims d -> Vector e -> Maybe (Tensor d e)
-- fromVector d v = undefined
--
-- toVector :: Elt e => Tensor d e -> Vector e
-- toVector = unTensor
--
fromSizedVector :: Elt e => Vector (Size d) e -> Tensor d e
fromSizedVector = Tensor . S.convert . N.fromSized

toSizedVector :: Elt e => Tensor d e -> Vector (Size d) e
toSizedVector = coerce . S.convert . unTensor
  where coerce :: S.Vector e -> N.Vector n e
        coerce = unsafeCoerce

fill :: Elt e => Dims d -> (Idxs d -> e) -> Tensor d e
fill d f = Tensor $ S.create $ do
  mv <- M.new $ fromIntegral $ product $ listDims d
  let act ix = M.write mv (minorToMajor d ix) $ f ix
  forMIdxs_ d act
  return mv

-- constant :: TensorType a => Shape -> [a] -> Tensor Build a

constant :: Elt e => Dims d -> e -> Tensor d e
constant d t = fill d $ const t

-}

constant :: Elt e => Dims d -> e -> Tensor d e
constant d = Tensor . O.constant (toShape d) . pure

toShape :: Dims d -> Shape
toShape = Shape . fmap fromIntegral . listDims

toShape' :: Idxs d -> Shape
toShape' = Shape . fmap fromIntegral . listIdxs

toShapeTensor :: Idxs d -> I '[Rank d]
toShapeTensor i = Tensor . O.constant (toShape' i) $ l
  where l = fmap fromIntegral . listIdxs $ i

-- TODO use more efficient density list approach
fill :: Elt e => Dims d -> (Idxs d -> e) -> Tensor d e
fill d k = Tensor . O.constant (toShape d) $ Prelude.reverse l
  where l = foldIdxs d (\i x -> k i : x) []

{-

fromIntegral' :: Int -> Int32
fromIntegral' = fromIntegral

toList $ fill (dims @'[2,4]) (fromIntegral' . fromEnum)

t1 = fill (dims @'[2,4]) $ const 1 . fromEnum
t2 = fill (dims @'[4,2]) $ const 1 . fromEnum

t2 = fill (dims @'[2,4]) ((+1) . fromIntegral' . fromEnum)

f :: Elt e => Tensor '[2,4] e -> Tensor '[4,2] e
f = transpose (dims @'[2,4]) (reversal (dims @'[2,4]))

t1' = f t1

foo :: Dims d -> (Idxs d -> TVal) -> IO ()
foo d f = printTensor $ fill d f

> d = dims @'[2,4]
> foo d (minorToMajor d)
[7,6,5,4,3,2,1,0]
> foo d (transposeIdxs $ majorToMinor d)
[7,5,3,1,6,4,2,0]

pred_sum_idxs :: forall ds. Dims ds -> Bool
pred_sum_idxs ds = foldIdxs ds (\_ a -> a + 1) 0 == (product . listDims $ ds)

-}

{-

Splits a tensor into sub tensors.

If num_or_size_splits is an integer type, then value is split along dimension axis into num_split smaller tensors. Requires that num_split evenly divides value.shape[axis].

If num_or_size_splits is not an integer type, it is presumed to be a Tensor size_splits, then splits value into len(size_splits) pieces. The shape of the i-th piece has the same size as the value except along dimension axis where the size is size_splits[i].

slice:
t = tf.constant([[[1, 1, 1], [2, 2, 2]],
                 [[3, 3, 3], [4, 4, 4]],
                 [[5, 5, 5], [6, 6, 6]]]) 3 x 2 x 3
tf.slice(t, [1, 0, 0], [1, 1, 3])  # [[[3, 3, 3]]]
tf.slice(t, [1, 0, 0], [1, 2, 3])  # [[[3, 3, 3],
                                   #   [4, 4, 4]]]
tf.slice(t, [1, 0, 0], [2, 1, 3])  # [[[3, 3, 3]],
                                   #  [[5, 5, 5]]]

-- want d' < d,
slice :: Idxs d -> Dims d' -> Tensor d e -> Tensor d' e

slice :: (TensorType t, OneOf '[Int32, Int64] index)	=> Tensor v'1 t	-> Tensor v'2 index	-> Tensor v'3 index	-> Tensor Build t

Idxs d -> I '[Rank d]

i = toEnum 6 :: Idxs '[3,2,3]
j = toEnum 9 :: Idxs '[3,2,3]

diffIdxs (dims @'[3,2,3]) i j

stepIdx (dims @'[3,2,3]) 3 i
i = 1 :+ 0 :+ 0 :+ S :: Idxs '[3,2,3]

-}

{-

pack
  :: Elt e
  => Vector n (Tensor (x ++ y) e) -> Tensor (x ++ n :+ y) e
pack = undefined

pack0
  :: forall d e n. Elt e
  => KnownDims d
  => KnownDim n
  => Vector n (Tensor d e) -> Tensor (n :+ d) e
pack0 v = Tensor res
  where d = dims @d
        n = fromIntegral $ dimVal $ dim @n
        size = product $ listDims d
        res = S.create $ do
          mv <- M.new $ fromIntegral $ size * n
          flip N.imapM_ v $ \i t ->
            let i' = idxToWord . idxFromFinite $ i
                off = fromIntegral $ i' * size
                v' = unTensor t
                act ix = M.write mv (off + fromEnum ix) $ v' S.! (fromEnum ix) -- could use a tensor op instead here
            in forMIdxs_ d act
          return mv

-- TODO use http://hackage.haskell.org/package/reflection-2.1.4/docs/Data-Reflection.html#t:reifyNat
unpack0
  :: forall d e n. Elt e
  => KnownDims d
  => KnownNat n
  => Tensor (n :+ d) e -> Vector n (Tensor d e)
unpack0 t = N.generate f
  where d = dims @d
        size = fromIntegral $ product $ listDims d
        f i = fill d $ \ix ->
          let i' = fromIntegral $ F.getFinite i
              off = i' * size
              v = unTensor t
          in v S.! (off + fromEnum ix)
-}

{-

see example http://jeankossaifi.com/blog/unfolding.html

t :: Vector 4 (Tensor '[2,2] Word)
t = N.generate $ \f ->
  let d = dims @'[2,2]
      i' = idxToWord . idxFromFinite $ f
  in fill d (const i')

t' :: Tensor '[4,2,2] Word
t' = pack0 t

t'' :: Vector 4 (Tensor '[2,2] Word)
t'' = unpack0 t'

generate :: forall n a. KnownNat n => (Finite n -> a) -> Vector n a

t :: Data.Vector.Sized.Vector 4 (Tensor '[2,2] Word)
t = generate $ \f ->
  let d = dims @'[2,2]
      i' = idxToWord . idxFromFinite $ f
  in fill d (const i')

t' :: Tensor '[4,2,2] Word
t' = pack0 t
N.imapM_ :: Monad m => (Finite n -> a -> m b) -> Vector n a -> m ()

overDim_ :: Monad m
         => Dims ds -- ^ Shape of a space
         -> (Idxs ds -> Int -> m ()) -- ^ Function to call on each dimension
         -> Int -- ^ Initial offset
         -> Int -- ^ Offset step
         -> m ()

stack ::  Dim n -> Vector n (Tensor (x ++ y) e) -> Tensor (x +: n :+ y) e
unstack :: (KnownDims x, Elt e) => Dim n -> Tensor (x +: n :+ y) e -> Vector n (Tensor (x ++ y) e)
-}