{-# 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 DeriveDataTypeable #-}

module BBrotConf where

import Data.Complex
import System.Console.CmdArgs

loCorner, hiCorner :: Complex Double
loCorner = (-2.0) :+ (-1.5)
hiCorner =   1.0  :+   1.5

data ColorPalette = Gray | Reddish | Flames
                    deriving (Show, Data, Typeable)

data Curve = Line | Square | Root
             deriving (Show, Data, Typeable)

data BBrotConf = Compute { seed       :: Maybe Int
                         , samples    :: Int
                         , minIters   :: Int
                         , maxIters   :: Int
                         , ocachepath :: Maybe String
                         , gridStep   :: Double
                         }
               | Render { xpixels     :: Int
                        , ypixels     :: Int
                        , icachepath  :: String
                        , isComplex   :: Bool
                        , dontRender  :: Bool
                        , imagepath   :: Maybe String
                        , palette     :: ColorPalette
                        , curve       :: Curve
                        }
               | ShowCells { gridStep :: Double
                           , maxIters :: Int
                           , animpath :: String
                           }
                 deriving (Show, Data, Typeable)


{-# LANGUAGE RankNTypes, TupleSections, ConstraintKinds, LambdaCase, ViewPatterns,
             FlexibleContexts, DeriveGeneric, TypeOperators, PatternSynonyms #-}
{-# OPTIONS_GHC -fplugin GHC.TypeLits.KnownNat.Solver #-}
{-# OPTIONS_GHC -fplugin GHC.TypeLits.Normalise #-}

module Tuple
(pattern (:&&), (:&) (..), snd', fst', mkT) where

import Numeric.LinearAlgebra.Static.Backprop
import Numeric.LinearAlgebra.Static.Vector
import Numeric.OneLiner
import Numeric.Backprop
import Lens.Micro hiding ((&))
import Data.List
import Control.Monad.Trans.State
import GHC.Generics (Generic)
import GHC.TypeNats
import System.Random
import qualified Data.Vector.Storable.Sized as SVS
import Data.Random.Normal

data a :& b = !a :& !b
    deriving (Show, Generic)
infixr 2 :&

pattern (:&&) :: (Backprop a, Backprop b, Reifies z W)
              => BVar z a -> BVar z b -> BVar z (a :& b)
pattern x :&& y <- (\xy -> (xy ^^. t1, xy ^^. t2) -> (x, y))
    where (:&&) = isoVar2 (:&) (\case x :& y -> (x, y))
{-# COMPLETE (:&&) #-}

t1 :: Lens (a :& b) (a' :& b) a a'
t1 f (x :& y) = (:& y) <$> f x

t2 :: Lens (a :& b) (a :& b') b b'
t2 f (x :& y) = (x :&) <$> f y

fst' (a :& _) = a
snd' (_ :& b) = b

mkT = (:&)

instance (Num a, Num b) => Num (a :& b) where
    (+)         = gPlus
    (-)         = gMinus
    (*)         = gTimes
    negate      = gNegate
    abs         = gAbs
    signum      = gSignum
    fromInteger = gFromInteger

instance (Fractional a, Fractional b) => Fractional (a :& b) where
    (/) = gDivide
    recip = gRecip
    fromRational = gFromRational

instance (Floating a, Floating b) => Floating (a :& b) where
    (**)  = gPower
    sqrt  = gSqrt
    acos  = gAcos
    asin  = gAsin
    atan  = gAtan
    cos   = gCos
    sin   = gSin
    tan   = gTan
    acosh = gAcosh
    asinh = gAsinh
    atanh = gAtanh
    cosh  = gCosh
    sinh  = gSinh
    tanh  = gTanh
    exp   = gExp
    log   = gLog
    pi    = gPi

instance (Random a, Random b) => Random (a :& b) where
    random g0 = (x :& y, g2)
      where
        (x, g1) = random g0
        (y, g2) = random g1
    randomR (x0 :& y0, x1 :& y1) g0 = (x :& y, g2)
      where
        (x, g1) = randomR (x0, x1) g0
        (y, g2) = randomR (y0, y1) g1

instance (Backprop a, Backprop b) => Backprop (a :& b)

uncurryT
    :: (Backprop a, Backprop b, Reifies z W)
    => (BVar z a -> BVar z b -> BVar z c)
    -> BVar z (a :& b)
    -> BVar z c
uncurryT f x = f (x ^^. t1) (x ^^. t2)

{- TODO randomR ist undefiniert -}

instance (KnownNat n, KnownNat m) => Random (L n m) where
    random = runState $ vecL <$> SVS.replicateM (state normal)
    randomR (lo, hi) = undefined

instance (KnownNat n) => Random (R n) where
    random = runState $ vecR <$> SVS.replicateM (state normal)
    randomR (lo, hi) = undefined


module Plot (
    plot,
    plotDensity
) where

import Graphics.Rendering.Chart
import Graphics.Rendering.Chart.Backend.Cairo
import Control.Lens
import Data.Colour
import Data.Colour.Names
import Data.Default.Class

import Plot.Utils
import Statistics.Types
import Statistics.Distribution
import qualified Data.Vector.Unboxed as V

width, height :: Int
width = 800
height = 400
file_opts :: FileOptions
file_opts = fo_size .~ (width, height)
          $ fo_format .~ SVG
          $ def

plot :: Sample -> String -> FilePath -> IO (PickFn ())
plot vs title fn = renderableToFile file_opts fn (toRenderable layout)
  where
    n = V.length vs

    tv, ev, dv :: [(Double, Double)]
    tv = [(fromIntegral i, i /. n) | i <- [1 .. n]]
    ev = zip (map fromIntegral [1 .. n]) (V.toList vs)
    dv = zipWith (\ (x, t) (_, e) -> (x, t - e)) tv ev

    layout :: LayoutLR Double Double Double
    layout = layoutlr_title      .~ title
           $ layoutlr_background .~ solidFillStyle (opaque white)
           $ layoutlr_plots      .~ [ Left (toPlot theoretical),
                                      Left (toPlot empirical),
                                      Right (toPlot difference)]
           $ layoutlr_right_axis .~ yaxis
           $ layoutlr_foreground .~ opaque black
           $ def

    theoretical :: PlotLines Double Double
    theoretical = theoreticalWith tv

    empirical :: PlotLines Double Double
    empirical = empiricalWith ev

    difference :: PlotLines Double Double
    difference = differenceWith dv

yaxis :: LayoutAxis Double
yaxis  = laxis_title    .~ "theoretical - empirical"
       $ laxis_override .~ (axisGridHide . rStyleTicks)
       $ def

rStyleTicks     :: AxisData x -> AxisData x
rStyleTicks ad  = ad{ _axis_ticks = map invert_tick (_axis_ticks ad) }
    where
        invert_tick :: (x,Double) -> (x, Double)
        invert_tick (x, t) = (x, -t)

plotDensity :: (ContDistr d) => d -> Sample -> String -> FilePath -> IO (PickFn ())
plotDensity gd vs title fn = renderableToFile file_opts fn (toRenderable layout)
  where

    n  = V.length vs
    mx = V.maximum vs

    -- divide domain into bins
    bins :: Double
    bins = fromIntegral $ max (n `div` 100) 1000

    xs, ys :: [Double]
    xs = map (* (mx / bins)) [0 .. bins]
    -- count how many at each level
    ys =  let counts = count (V.toList vs) xs
              scale'  = fromIntegral n * (mx / bins)
          in  map ((/ scale') . fromIntegral) counts

    layout :: LayoutLR Double Double Double
    layout = stdLayout title
               [ Left (toPlot (  empiricalWith (zip xs ys)))
               , Left (toPlot (theoreticalWith [(x, density gd x) | x <- xs]) )]

(/.) :: (Real a, Real b, Fractional c) => a -> b -> c
(/.) x y = fromRational $ toRational x / toRational y

count :: [Double] -> [Double] -> [Int]
count samples breaks = map length (foldr go (:[]) breaks samples)
  where
    go x r s = let (ys, zs) = break (>= x) s in ys : r zs


import Test.Hspec
import Numeric.LinearAlgebra

import Statistics.GLM

main :: IO ()
main = hspec $ do
  describe "logReg" $ do
    it "can perform logistic regression" $ do
      fmap (truncate . (*10)) (toList (last $ take 5 $ glm logisticReg testA testX testB)) `shouldBe` fmap (truncate . (*10)) (toList $ vector [log 8, log 0.25])

testA :: Matrix Double
testA = (6><2) [1::Double, 1, 1, 1, 1, 1, 1, 2, 1, 2, 1, 2]

testB :: Vector Double
testB = vector [1, 1, 0, 1, 0, 0]

testX :: Vector Double
testX = vector [0, 0]


module School.Unit.RecLin
( recLin ) where

import Numeric.LinearAlgebra (R, cmap, cond)
import School.Unit.Unit (Unit(..))
import School.Unit.UnitActivation (UnitActivation(..))
import School.Unit.UnitGradient (UnitGradient(..))
import School.Unit.UnitParams (UnitParams(..))

recLin :: Unit R
recLin = Unit { apply = recLinApply
              , deriv = recLinDeriv
              }

recLinApply :: UnitParams R
            -> UnitActivation R
            -> UnitActivation R
recLinApply EmptyParams (BatchActivation inMatrix) =
  BatchActivation $ cmap (max 0) inMatrix
recLinApply _ f@(ApplyFail _) = f
recLinApply _ _ = ApplyFail "Failure when applying RecLin unit"

recLinDeriv :: UnitParams R
            -> UnitGradient R
            -> UnitActivation R
            -> ( UnitGradient R
               , UnitParams R)
recLinDeriv EmptyParams
            (BatchGradient inGrad)
            (BatchActivation input)
  = (BatchGradient grad, EmptyParams)
    where grad = cond input 0 0 0 inGrad
recLinDeriv EmptyParams f@(GradientFail _) _ = (f, EmptyParams)
recLinDeriv EmptyParams _ (ApplyFail msg) = (GradientFail msg, EmptyParams)
recLinDeriv _ _ _ = (GradientFail "Failure when differentiating RecLin unit" , EmptyParams)


{-# 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 FlexibleContexts, FlexibleInstances, NamedFieldPuns #-}

module School.FileIO.Source
( Source(..)
, pool
, source
) where

import Conduit ((.|), ConduitM, MonadResource, mapMC, nullC, sourceFileBS, takeCE)
import Control.Monad ((>=>), unless)
import Control.Monad.Except (MonadError)
import Data.ByteString (ByteString)
import Numeric.LinearAlgebra (Matrix, R)
import School.FileIO.BinConversion (binConversion)
import School.FileIO.ConduitHeader (conduitHeader)
import School.FileIO.FileHeader (FileHeader(..))
import School.FileIO.FilePath (FilePath)
import School.Types.DataType (DataType(..), getSize)
import School.Types.Decoding (binToDouble, binToListInt, binToMatrixDouble)
import School.Types.Error (Error)
import School.Types.LiftResult (LiftResult(..))

class Source a where
  decodeByteString :: FileHeader -> ByteString -> Either Error a

pool :: (Monad m)
     => Int
     -> (ByteString -> m a)
     -> ConduitM ByteString a m ()
pool chunkSize transformer = loop where
  loop = do
    takeCE chunkSize .| mapMC transformer
    isEmpty <- nullC
    unless isEmpty loop

source :: (LiftResult m, MonadError Error m, MonadResource m, Source a)
       => FileHeader
       -> FilePath
       -> ConduitM () a m ()
source header@FileHeader { dataType, cols, rows } path =
  let size = getSize dataType * cols * rows
      byteSource = sourceFileBS path
      cHeader = conduitHeader header
      trans = liftResult . decodeByteString header
  in byteSource .| cHeader .| pool size trans

instance Source (Matrix R) where
  decodeByteString FileHeader{ dataType, cols, rows } =
    binToMatrixDouble dataType rows cols

instance Source Double where
  decodeByteString FileHeader{ dataType } =
    binConversion dataType DBL64B >=> binToDouble

instance Source [Int] where
  decodeByteString FileHeader{ dataType } = binToListInt dataType


{-# LANGUAGE ExtendedDefaultRules #-}
{-# LANGUAGE ScopedTypeVariables #-}

module TensorSpec (spec) where

import Control.Arrow ((&&&))
import Control.Exception.Safe
import Numeric.Half
import Data.Complex
import Data.Int
import Data.Word
import Test.Hspec
import Test.QuickCheck
import Torch.DType
import Torch.Functional
import Torch.Tensor
import Torch.TensorFactories
import Torch.TensorOptions
import Test.QuickCheck.Arbitrary

instance Arbitrary Half where
  arbitrary = arbitrarySizedFractional
  shrink    = shrinkDecimal

spec :: Spec
spec = do
  describe "TensorLike" $ do
    it "TensorLike Bool" $
      property $
        \x -> asValue (asTensor x) `shouldBe` (x :: Bool)
    it "TensorLike Word8" $
      property $
        \x -> asValue (asTensor x) `shouldBe` (x :: Word8)
    it "TensorLike Int8" $
      property $
        \x -> asValue (asTensor x) `shouldBe` (x :: Int8)
    it "TensorLike Int16" $
      property $
        \x -> asValue (asTensor x) `shouldBe` (x :: Int16)
    it "TensorLike Int32" $
      property $
        \x -> asValue (asTensor x) `shouldBe` (x :: Int32)
    it "TensorLike Int" $
      property $
        \x -> asValue (asTensor x) `shouldBe` (x :: Int)
    it "TensorLike Int64" $
      property $
        \x -> asValue (asTensor x) `shouldBe` (x :: Int64)
    it "TensorLike Half" $
      property $
        \x -> asValue (asTensor x) `shouldBe` (x :: Half)
    it "TensorLike Float" $
      property $
        \x -> asValue (asTensor x) `shouldBe` (x :: Float)
    it "TensorLike Double" $
      property $
        \x -> asValue (asTensor x) `shouldBe` (x :: Double)
    it "TensorLike ComplexHalf" $
      property $
        \x -> asValue (asTensor x) `shouldBe` (x :: Complex Half)
    it "TensorLike Complex Float" $
      property $
        \x -> asValue (asTensor x) `shouldBe` (x :: Complex Float)
    it "TensorLike Complex Double" $
      property $
        \x -> asValue (asTensor x) `shouldBe` (x :: Complex Double)

    it "Compare internal expression of c++ with Storable expression of haskell" $ do
      show (asTensor [True, False, True, False])
        `shouldBe` "Tensor Bool [4] [ 1,  0,  1,  0]"
      show (asTensor ([1, 0, 1, 0] :: [Word8]))
        `shouldBe` "Tensor UInt8 [4] [ 1,  0,  1,  0]"
      show (asTensor ([1, 0, 1, 0] :: [Int8]))
        `shouldBe` "Tensor Int8 [4] [ 1,  0,  1,  0]"
      show (asTensor ([1, 0, 1, 0] :: [Int16]))
        `shouldBe` "Tensor Int16 [4] [ 1,  0,  1,  0]"
      show (asTensor ([1, 0, 1, 0] :: [Int32]))
        `shouldBe` "Tensor Int32 [4] [ 1,  0,  1,  0]"
      show (asTensor ([1, 0, 1, 0] :: [Int]))
        `shouldBe` "Tensor Int64 [4] [ 1,  0,  1,  0]"
      show (asTensor ([1, 0, 1, 0] :: [Int64]))
        `shouldBe` "Tensor Int64 [4] [ 1,  0,  1,  0]"
      show (asTensor ([1, 0, 1, 0] :: [Float]))
        `shouldBe` "Tensor Float [4] [ 1.0000   ,  0.0000,  1.0000   ,  0.0000]"
      show (asTensor ([1, 0, 1, 0] :: [Double]))
        `shouldBe` "Tensor Double [4] [ 1.0000   ,  0.0000,  1.0000   ,  0.0000]"
      show (asTensor ([[]] :: [[Int]]))
        `shouldBe` "Tensor Int64 [1,0] [[]]"
      show (asTensor ([[1]] :: [[Int]]))
        `shouldBe` "Tensor Int64 [1,1] [[ 1]]"

    it "TensorLike [Bool]" $
      property $
        \(NonEmpty (x :: [Bool])) -> do
          asValue (asTensor x) `shouldBe` x
          toDouble (select 0 0 (asTensor x)) `shouldBe` (if (head x) then 1 else 0)
          shape (asTensor x) `shouldBe` [length x]
          let xx = replicate 5 x
          asValue (asTensor xx) `shouldBe` xx
          let xxx = replicate 3 xx
          asValue (asTensor xxx) `shouldBe` xxx
    it "TensorLike [Word8]" $
      property $
        \(NonEmpty (x :: [Word8])) -> do
          asValue (asTensor x) `shouldBe` x
          toDouble (select 0 0 (asTensor x)) `shouldBe` fromIntegral (head x)
          shape (asTensor x) `shouldBe` [length x]
          let xx = replicate 5 x
          asValue (asTensor xx) `shouldBe` xx
          let xxx = replicate 3 xx
          asValue (asTensor xxx) `shouldBe` xxx
    it "TensorLike [Int8]" $
      property $
        \(NonEmpty (x :: [Int8])) -> do
          asValue (asTensor x) `shouldBe` x
          toDouble (select 0 0 (asTensor x)) `shouldBe` fromIntegral (head x)
          shape (asTensor x) `shouldBe` [length x]
          let xx = replicate 5 x
          asValue (asTensor xx) `shouldBe` xx
          let xxx = replicate 3 xx
          asValue (asTensor xxx) `shouldBe` xxx
    it "TensorLike [Int16]" $
      property $
        \(NonEmpty (x :: [Int16])) -> do
          asValue (asTensor x) `shouldBe` x
          toDouble (select 0 0 (asTensor x)) `shouldBe` fromIntegral (head x)
          shape (asTensor x) `shouldBe` [length x]
          let xx = replicate 5 x
          asValue (asTensor xx) `shouldBe` xx
          let xxx = replicate 3 xx
          asValue (asTensor xxx) `shouldBe` xxx
    it "TensorLike [Int32]" $
      property $
        \(NonEmpty (x :: [Int32])) -> do
          asValue (asTensor x) `shouldBe` x
          toDouble (select 0 0 (asTensor x)) `shouldBe` fromIntegral (head x)
          shape (asTensor x) `shouldBe` [length x]
          let xx = replicate 5 x
          asValue (asTensor xx) `shouldBe` xx
          let xxx = replicate 3 xx
          asValue (asTensor xxx) `shouldBe` xxx
    it "TensorLike [Int]" $
      property $
        \(NonEmpty (x :: [Int])) -> do
          asValue (asTensor x) `shouldBe` x
          toDouble (select 0 0 (asTensor x)) `shouldBe` fromIntegral (head x)
          shape (asTensor x) `shouldBe` [length x]
          let xx = replicate 5 x
          asValue (asTensor xx) `shouldBe` xx
          let xxx = replicate 3 xx
          asValue (asTensor xxx) `shouldBe` xxx
    it "TensorLike [Int64]" $
      property $
        \(NonEmpty (x :: [Int64])) -> do
          asValue (asTensor x) `shouldBe` x
          toDouble (select 0 0 (asTensor x)) `shouldBe` fromIntegral (head x)
          shape (asTensor x) `shouldBe` [length x]
          let xx = replicate 5 x
          asValue (asTensor xx) `shouldBe` xx
          let xxx = replicate 3 xx
          asValue (asTensor xxx) `shouldBe` xxx
    it "TensorLike [Half]" $
      property $
        \(NonEmpty (x :: [Half])) -> do
          asValue (asTensor x) `shouldBe` x
          toDouble (select 0 0 (asTensor x)) `shouldBe` realToFrac (head x)
          shape (asTensor x) `shouldBe` [length x]
          let xx = replicate 5 x
          asValue (asTensor xx) `shouldBe` xx
          let xxx = replicate 3 xx
          asValue (asTensor xxx) `shouldBe` xxx
    it "TensorLike [Float]" $
      property $
        \(NonEmpty (x :: [Float])) -> do
          asValue (asTensor x) `shouldBe` x
          toDouble (select 0 0 (asTensor x)) `shouldBe` realToFrac (head x)
          shape (asTensor x) `shouldBe` [length x]
          let xx = replicate 5 x
          asValue (asTensor xx) `shouldBe` xx
          let xxx = replicate 3 xx
          asValue (asTensor xxx) `shouldBe` xxx
    it "TensorLike [Double]" $
      property $
        \(NonEmpty (x :: [Double])) -> do
          asValue (asTensor x) `shouldBe` x
          toDouble (select 0 0 (asTensor x)) `shouldBe` realToFrac (head x)
          shape (asTensor x) `shouldBe` [length x]
          let xx = replicate 5 x
          asValue (asTensor xx) `shouldBe` xx
          let xxx = replicate 3 xx
          asValue (asTensor xxx) `shouldBe` xxx
    it "TensorLike [Complex Half]" $
      property $
        \(NonEmpty (x :: [Complex Half])) -> do
          asValue (asTensor x) `shouldBe` x
          asValue (select 0 0 (asTensor x)) `shouldBe` (head x)
          shape (asTensor x) `shouldBe` [length x]
          let xx = replicate 5 x
          asValue (asTensor xx) `shouldBe` xx
          let xxx = replicate 3 xx
          asValue (asTensor xxx) `shouldBe` xxx
    it "TensorLike [Complex Float]" $
      property $
        \(NonEmpty (x :: [Complex Float])) -> do
          asValue (asTensor x) `shouldBe` x
          asValue (select 0 0 (asTensor x)) `shouldBe` (head x)
          shape (asTensor x) `shouldBe` [length x]
          let xx = replicate 5 x
          asValue (asTensor xx) `shouldBe` xx
          let xxx = replicate 3 xx
          asValue (asTensor xxx) `shouldBe` xxx
    it "TensorLike [Complex Double]" $
      property $
        \(NonEmpty (x :: [Complex Double])) -> do
          asValue (asTensor x) `shouldBe` x
          asValue (select 0 0 (asTensor x)) `shouldBe` (head x)
          shape (asTensor x) `shouldBe` [length x]
          let xx = replicate 5 x
          asValue (asTensor xx) `shouldBe` xx
          let xxx = replicate 3 xx
          asValue (asTensor xxx) `shouldBe` xxx

    it "invalid cast of TensorLike a" $ do
      let x = asTensor (10 :: Int)
      (dtype x) `shouldBe` Int64
      (print (asValue x :: Double)) `shouldThrow` anyException
    it "invalid cast of TensorLike [a]" $ do
      let x = asTensor ([0 .. 10] :: [Int])
      (print (asValue x :: [Double])) `shouldThrow` anyException

    it "lists having different length" $ do
      (print (asTensor ([[1], [1, 2]] :: [[Double]]))) `shouldThrow` anyException
    it "cast of Tensor" $ do
      let x = asTensor ([0 .. 10] :: [Int])
      (dtype (toType Float x)) `shouldBe` Float

  describe "indexing" $ do
    it "pick up a value" $ do
      let x = asTensor ([[[0, 1, 2], [3, 4, 5]], [[6, 7, 8], [9, 10, 11]]] :: [[[Int]]])
          r = x ! (1, 0, 2)
      (dtype &&& shape &&& asValue) r `shouldBe` (Int64, ([], 8 :: Int))
    it "pick up a bottom tensor" $ do
      let x = asTensor ([[[0, 1, 2], [3, 4, 5]], [[6, 7, 8], [9, 10, 11]]] :: [[[Int]]])
          r = x ! (1, 0)
      (dtype &&& shape &&& asValue) r `shouldBe` (Int64, ([3], [6 :: Int, 7, 8]))
    it "make a slice of bottom values" $ do
      let x = asTensor ([[[0, 1, 2], [3, 4, 5]], [[6, 7, 8], [9, 10, 11]]] :: [[[Int]]])
          r = x ! ((), (), 1)
      (dtype &&& shape &&& asValue) r `shouldBe` (Int64, ([2, 2], [[1 :: Int, 4], [7, 10]]))
    it "ellipsis" $ do
      let x = asTensor ([[[0, 1, 2], [3, 4, 5]], [[6, 7, 8], [9, 10, 11]]] :: [[[Int]]])
          r = x ! (Ellipsis, 1)
      (dtype &&& shape &&& asValue) r `shouldBe` (Int64, ([2, 2], [[1 :: Int, 4], [7, 10]]))
    it "make a slice via muliple slices" $ do
      let x = asTensor ([[[0, 1, 2], [3, 4, 5]], [[6, 7, 8], [9, 10, 11]]] :: [[[Int]]])
          r = x ! ((), (Slice (1, None)))
      (dtype &&& shape &&& asValue) r `shouldBe` (Int64, ([2, 1, 3], [[[3 :: Int, 4, 5]], [[9, 10, 11]]]))
    it "make a slice via muliple slices" $ do
      let x = asTensor ([[[0, 1, 2], [3, 4, 5]], [[6, 7, 8], [9, 10, 11]]] :: [[[Int]]])
          r = x ! ((), (Slice (1, None)), (Slice (0, 1)))
      (dtype &&& shape &&& asValue) r `shouldBe` (Int64, ([2, 1, 1], [[[3 :: Int]], [[9]]]))
  describe "masked fill" $ do
    it "Fill a value" $ do
      let x = asTensor ([[[0, 1, 2], [3, 4, 5]], [[6, 7, 8], [9, 10, 11]]] :: [[[Int]]])
          r = maskedFill x (1, 0, 2) (9 :: Int)
      (dtype &&& shape &&& asValue) r `shouldBe` (Int64, ([2, 2, 3], [[[0, 1, 2], [3, 4, 5]], [[6, 7, 9], [9, 10, 11]]] :: [[[Int]]]))
    it "Fill a bottom tensor" $ do
      let x = asTensor ([[[0, 1, 2], [3, 4, 5]], [[6, 7, 8], [9, 10, 11]]] :: [[[Int]]])
          r = maskedFill x (1, 0) [8 :: Int, 8, 8]
      (dtype &&& shape &&& asValue) r `shouldBe` (Int64, ([2, 2, 3], [[[0, 1, 2], [3, 4, 5]], [[8, 8, 8], [9, 10, 11]]] :: [[[Int]]]))
    it "masked fill by boolean" $ do
      let m = asTensor ([[[True, True, False], [False, False, False]], [[True, False, False], [False, False, False]]] :: [[[Bool]]])
          x = asTensor ([[[0, 1, 2], [3, 4, 5]], [[6, 7, 8], [9, 10, 11]]] :: [[[Int]]])
          r = maskedFill x m [8 :: Int, 8, 8]
      (dtype &&& shape &&& asValue) r `shouldBe` (Int64, ([2, 2, 3], [[[8, 8, 2], [3, 4, 5]], [[8, 7, 8], [9, 10, 11]]] :: [[[Int]]]))


{-# LANGUAGE BangPatterns     #-}
{-# LANGUAGE TypeApplications #-}
module ArrayFire.IndexSpec where

import qualified ArrayFire         as A
import           Control.Exception
import           Data.Complex
import           Data.Int
import           Data.Proxy
import           Data.Word
import           Foreign.C.Types
import           Test.Hspec

spec :: Spec
spec =
  describe "Index spec" $ do
    it "Should index into an array" $ do
      let arr = A.vector @Int 10 [1..]
      A.index arr [A.Seq 0 4 1]
        `shouldBe`
           A.vector @Int 5 [1..]


module Dna.Search (
    dnaFromSepll, composingKMer,
    initDeBrujinGraphFromKMers,
    initDeBrujinGraphFromDna, fibrabbit) where

import Dna.Graph
import Dna.Nucleotides
import Data.List
import Control.Applicative
import Data.Maybe
import Control.Monad
import Numeric.LinearAlgebra.Devel

composingKMer ::  Int -> Dna -> DnaList
composingKMer n dna
    | lengthDna dna >= n = takeDna n dna : composingKMer n (tailDna dna)
    | otherwise = []

same one two = one `isSubsequenceOf` two && two `isSubsequenceOf` one

mergeHeadTail [] r = r
mergeHeadTail (lx:lxs) r@(rx:rxs)
    | lx == rx = r
    | otherwise = lx: mergeHeadTail lxs r

isConnected [] _ = False
isConnected _ [] = False
isConnected (_:lxs) r = lxs `isPrefixOf` r

dnaFromSepll :: DnaList -> Dna
dnaFromSepll [] = nullDna
dnaFromSepll [dna] = dna
dnaFromSepll (x:xs) = dnaSpellStep 1 x xs
    where
        dnaSpellStep dropStep growDna dnas = case find (isDnaPrefixOf $ dropDna dropStep growDna) dnas of
            Just newDna -> dnaSpellStep (dropStep+1) (takeDna dropStep growDna +:+ newDna) dnas
            Nothing -> growDna

initDeBrujinGraphFromDna :: Int -> Dna -> Graph Dna
initDeBrujinGraphFromDna n dnas = initOverlaGraph nodes $ \ m findIndex ->
        forM_ adjacentNodes $ \ (from, to) ->
            modifyMatrix m (findIndex from) (findIndex to) (1+)
    where
        nodes = composingKMer n dnas
        adjacentNodes = zip (init nodes) (tail nodes)

initDeBrujinGraphFromKMers :: DnaList -> Graph Dna
initDeBrujinGraphFromKMers kmers = initOverlaGraph nodes $ \m findIdxNode ->
    forM_ nodes $ \from ->
        forM_ (map tailDna (filter (from `isDnaPrefixOf`) kmers)) $ \to ->
            modifyMatrix m (findIdxNode from) (findIdxNode to) (1+)
        where nodes = nub $ map (liftDna init) kmers ++ map tailDna kmers

fibrabbit 1 _ = 1
fibrabbit 2 _ = 1
fibrabbit n k = fibrabbit (n-1) k + (k * fibrabbit (n-2) k)


{-# 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


module Proxy.PotentialFields.Internal.SimpleActions where
import Proxy.PotentialFields.Internal.Fields
import Proxy.Command.Commands
import Control.Monad.Reader
import Data.Complex
import Proxy.Types.Game
import Proxy.Server.Messages

type Map = [Location]
type FieldEnvironment = (Map,[UnitData],[UnitData])
type UnitAction = Reader FieldEnvironment Command

---------------------------------------------------------------------------------------------------------------------
--Simple UnitActions for a unit to perform. These are a single step that a unit might take at any given moment, they rely on the information from from potential field being generated.

toward :: UnitData -> Location -> UnitAction
toward ud l = mapReader (rightClickAt uid) field
  where uid = unitId ud
        field = zeta passAvU actAv (Just l) peaApp ud

attackTowards :: UnitData -> Location -> UnitAction
attackTowards ud l = mapReader (attackMove uid) field
 where uid = unitId ud
       field = zeta passAvU actAv (Just l) peaApp ud

-------------------------------------------------------------------------------------------------
--field accumulation zeta returns a computation which will return the resultant location for a unit to move to.

first (a,_,_) = a
second (_,a,_) = a
third (_,_,a) = a

zeta :: PotCalc UnitData -> PotCalc UnitData -> Maybe Location -> PotCalc Location -> UnitData -> Reader FieldEnvironment Location
zeta f h l lc u = do
  mapd <- asks (flip mapData u . first)
  friend <- asks (map (flip f u) . second)
  host <- asks (map (flip h u) . third)
  loc <- case l of
    Just a -> return [lc a u]
    otherwise -> return []
  return ( moveTo (unitLocation u) $ add mapd (vSum.concat $ [friend,host,loc]))

ro :: PotCalc UnitData -> PotCalc UnitData -> UnitData -> Reader FieldEnvironment Location
ro f h = zeta f h Nothing nofield

-----------------------------------------------------------------------------------------------------------
--Conversion from a float vector to a location for the unit to move to.
--NEED TO CONSIDER SPEED AND THE SIZE OF THE VECTOR RETURNED.

moveTo :: Location -> Vector -> Location
moveTo l v
  | magnitude cv == 0 = l
  | otherwise = l `add` xi 1 phi
  where phi = phase cv
        cv = uncurry (:+) v

xi :: (RealFloat a,Integral b) => b -> a -> (b,b)
xi b theta
  | negate ((1/4)*pi) < theta && (1/4)*pi >= theta = (b,chi $ theta)
  | (1/4)*pi < theta && (3/4)*pi >= theta = (chi.ref $ theta ,b)
  | negate ((3/4)*pi) < theta && negate ((1/4)*pi) >= theta = (chi $ theta + ((1/2)*pi),negate b)
  | otherwise = (negate b, chi.prep $ theta)
  where chi = round.(f.fromIntegral $ b)
        prep a = if a > 0
                 then pi - a
                 else negate (a + pi)
        ref a = (1/2)*pi - a

f :: (Floating a) => a -> a -> a
f a theta = (4 * a * theta)  / pi

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

allActions :: [UnitAction] -> FieldEnvironment -> [Command]
allActions ua  = runReader (sequence ua)


{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE Strict #-}
{-# LANGUAGE FlexibleContexts #-}

import CPVO.Numeric
import CPVO.IO
import CPVO.IO.Reader.Ecalj.Common
import CPVO.IO.Reader.Ecalj.DOS

import System.Environment (getArgs)
import qualified Data.Text as T
import qualified Data.Text.IO as T
import Data.List
-------------------------

import Numeric.LinearAlgebra
-------------------------

rydberg :: Double
rydberg=1/13.605

testArgs :: [String]
testArgs = concat [ [ "GWtable01.tex"
                   , "Total and partial DOS at *E*~F~ in states/eV/unit-cell, QSGW method."
                   , "|*N*$_\\uparrow$(*E*~F~)|*N*$_\\downarrow$(*E*~F~)|*N*(*E*~F~)"
                     ]
                   , words "@{}lSSS@{} -9:6 25 T o flipSpin nico2o4 extendedNiCo2O4.normal/nico2o4.6G10 Ni:Ni#3d:6:7:9:8:10 Co:Co#3d:6:7:8:9:10 O:O#2p:3:4:5"
                   ]

main :: IO ()
main = do
    allArgs <- getArgs
    --Right (invStat, ymax, xmin, xmax, ctrlAtoms, uniqAtoms, ctrlAtomicAOs,jdTable, jdHeads, foldernya, tailer, colAlign, texFile) <- readHeaderData allArgs
    Right (invStat,_,_,_,_,_, ctrlAtomicAOs,jdTable, jdHeads, foldernya, tailer,_, texFile) <- readHeaderData allArgs
    totalDOS <- readTotalDOSText tailer foldernya
    let resTotY0 = map (\i -> getY0 $ totalDOS ¿ [0,i]) $ flipBy invStat [1,2]
    putStrLn "===========NF total DOS========="
    putStrLn $ show resTotY0
    pdosAtAll <- readPDOS invStat tailer foldernya ctrlAtomicAOs
    let rs' = rendertable
              $ (:) jdHeads
              $ (:)  ("Total" : (map (showDouble 3) $ processCols $ map ((* rydberg) . last . toList) $ resTotY0))
              $ map (\a -> (((fst $ head a):(map (showDouble 3) $ processCols $ sumIt [0,0] $ map snd a))))
              $ groupBy (\a b -> fst a == fst b)
              $ (\[us,ds] -> zipWith (\(yu,(_,(j,_))) (yd,_) -> (j, [yu,yd])) us ds )
              $ groupBy (\(_,(s,_)) (_,(s',_)) ->  s == s')
              $ map (\(mp,b) -> ((* rydberg) $ last $ toList $ getY0 mp, b)) pdosAtAll
    let rs = unlines [ rs' , jdTable ]
    T.writeFile texFile $ T.pack rs
{-
    res' <- markdownToTex rs
    let res = T.replace "\\}" "}"
            $ T.replace "\\{" "{"
            $ T.pack $ unlines  [ "\\begin{longtable}[]{" ++ colAlign ++ "}"
                                , unlines $ tail $ lines $ T.unpack res'
                                ]
    T.writeFile texFile res
    -}
    putStrLn $ texFile ++ "=============done:t1dosEf.hs=========="
      where
        processCols [u,d] = [u,d,u+d]
        processCols _ = []
        sumIt x [] = x
        sumIt [u0,d0] ([up,dn]:xs) = sumIt [u0+up,d0+dn] xs
        sumIt _ _ = []


-------------------------------------------------------------------------------
-- |
-- Module    :  Data.CartPole
-- Copyright :  (c) Sentenai 2017
-- License   :  BSD3
-- Maintainer:  sam@sentenai.com
-- Stability :  experimental
-- Portability: non-portable
--
-- Shared datatypes between Gym environments and the haskell implementation of
-- CartPole.
-------------------------------------------------------------------------------
{-# LANGUAGE InstanceSigs #-}
{-# LANGUAGE DeriveGeneric #-}
{-# LANGUAGE DataKinds #-}
module Data.CartPole
  ( StateCP(..)
  , Action(..)
  , Event
  ) where

import qualified Reinforce.Spaces.State as Spaces

import Data.Aeson
import Data.Aeson.Types
import Data.Hashable
import GHC.Generics
import Control.Exception (AssertionFailed(..), throw)
import Reinforce.Spaces
import Reinforce.Spaces.Action (Size)
-- import Numeric.LinearAlgebra.Static
import qualified Data.Logger as Logger
import qualified Data.Vector as V

-- | Specific datatype for a CartPole event
type Event = Logger.Event Double StateCP Action

-- | Cartpole can only go left or right has an action space
-- of "discrete 2" containing {0..n-1}.
--
-- FIXME: Migrate this to either a more generic "directions" actions
-- (would need things like "up", "down" versions as well) or a "discrete
-- actions" version. I'm a fan of the former.
data Action
  = GoLeft
  | GoRight
  deriving (Show, Eq, Enum, Bounded, Ord, Generic)

instance Hashable Action

instance DiscreteActionSpace Action where
  type Size Action = 2

instance ToJSON Action where
  toJSON :: Action -> Value
  toJSON GoLeft  = toJSON (0 :: Int)
  toJSON GoRight = toJSON (1 :: Int)

-- | The state of a cart on a pole in a CartPole environment
data StateCP = StateCP
  { position  :: Float  -- ^ position of the cart on the track
  , angle     :: Float  -- ^ angle of the pole with the vertical
  , velocity  :: Float  -- ^ cart velocity
  , angleRate :: Float  -- ^ rate of change of the angle
  } deriving (Show, Eq, Generic, Ord)

instance Hashable StateCP

instance Semigroup StateCP where
  (StateCP a0 b0 c0 d0) <> (StateCP a1 b1 c1 d1)
    = StateCP (a0+a1) (b0+b1) (c0+c1) (d0+d1)

instance Monoid StateCP where
  mempty = StateCP 0 0 0 0
  mappend = (<>)

instance Spaces.StateSpace StateCP where
  toVector (StateCP p a v r) = Spaces.toVector [p, a, v, r]
  fromVector vec =
    case getVals of
      Nothing -> throw $ AssertionFailed "malformed vector found"
      Just s -> return s

    where
      getVals :: Maybe StateCP
      getVals = StateCP
        <$> findField 0
        <*> findField 1
        <*> findField 2
        <*> findField 3

      findField :: Int -> Maybe Float
      findField i = realToFrac <$> vec V.!? i

instance FromJSON StateCP where
  parseJSON :: Value -> Parser StateCP
  parseJSON arr@(Array _)= do
    (p, a, v, r) <- parseJSON arr :: Parser (Float, Float, Float, Float)
    return $ StateCP p a v r
  parseJSON invalid    = typeMismatch "StateCP" invalid

instance StateSpaceStatic StateCP where
  type Size StateCP = 4
  -- toR = vector . V.toList . Spaces.toVector
  -- fromR = Spaces.fromVector . V.fromList . LA.toList . unwrap


{-# 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)
-}