{-# LANGUAGE DataKinds #-}
{-# LANGUAGE Arrows, FlexibleContexts #-}

module Examples.MultiTargetTracking where

import Control.Arrow (returnA)
import Control.Monad (forM)
import Control.Monad.Bayes.Class (MonadSample, MonadInfer)

import qualified Data.Map as M
import Data.Maybe (fromJust)

import Numeric.LinearAlgebra.Static hiding ((<>))

import Inference (zdsparticles)
import DelayedSampling (DelayedSample, DelayedInfer, Result (..))
import qualified SymbolicDistr as DS
import DSProg (DeepForce (..), Expr' (..), Expr, marginal, zdeepForce, deepForce', forgetE)
import Distributions
import Util.ZStream (ZStream)
import qualified Util.ZStream as ZS
import Metaprob ((~~), Gen, lift, (|->), obs)
import qualified Metaprob as MP
import Examples.Demo (Sampler, Delayed, Weighted)

data ObsType = Clutter | NewTrack | Track Int
  deriving (Eq, Show, Ord)

filterM :: Monad m => (k -> v -> m Bool) -> M.Map k v -> m (M.Map k v)
filterM f = M.traverseMaybeWithKey $ \k v -> do
  b <- f k v
  return (if b then Just v else Nothing)

type TrackMap = M.Map Int (Expr (R 4))

tdiff, birthRate, deathRate, clutterLambda, newTrackLambda, pd, survivalProb :: Double
tdiff = 1
birthRate = 0.1
deathRate = 0.02
clutterLambda = 3
newTrackLambda = birthRate * tdiff
pd = 0.8
survivalProb = exp (- tdiff * deathRate)

clutterDistr :: DS.Distr (Expr (R 2))
clutterDistr = DS.mvNormal (Const (0 :: R 2)) (10 * sym eye)

newTrackD :: DelayedSample m => m (Expr (R 4))
newTrackD = DS.sample (DS.mvNormal (Const mu) cov)
  where
  mu = 0 :: R 4
  cov = sym (((5 * eye :: Sq 2) ||| (0 :: Sq 2))
            ===
            ((0 :: Sq 2) ||| (0.1 * eye :: Sq 2)))

trackMotion :: DelayedSample m => Double -> Expr (R 4) -> m (Expr (R 4))
trackMotion tdiff track = DS.sample (DS.mvNormal (MVMul (Const motionMatrix) track) motionCov)
  where
  motionMatrix :: Sq 4
  motionMatrix = eye + konst tdiff *
    (((0 :: Sq 2) ||| (eye :: Sq 2))
    ===
    ((- 1 / 100 * eye :: Sq 2) ||| (-0.1 * eye :: Sq 2)))
  motionCov :: Sym 4
  motionCov = sym $ konst tdiff * (((0.01 * eye :: Sq 2) ||| (0 :: Sq 2))
                                  ===
                                  ((0 :: Sq 2) ||| (0.1 * eye :: Sq 2)))

trackMeasurement :: Expr (R 4) -> DS.Distr (Expr (R 2))
trackMeasurement posvel = DS.mvNormal (MVMul (Const posFromPosVel) posvel) (sym eye)
  where
  posFromPosVel :: L 2 4
  posFromPosVel = (eye :: Sq 2) ||| (0 :: Sq 2)

updateWithAssocs :: DelayedSample m => Int -> TrackMap -> [ObsType] -> m ((TrackMap, Int), [DS.Distr (Expr (R 2))])
updateWithAssocs nextTrackID updatedOldTracks assocs = do
  (obsDists, newTrackPVs) <- fmap mconcat . forM assocs $ \k -> case k of
    Clutter -> return ([clutterDistr], [])
    NewTrack -> do
      newTrackPV <- newTrackD
      return ([trackMeasurement newTrackPV], [newTrackPV])
    Track i -> return ([trackMeasurement (updatedOldTracks M.! i)], [])
  coasted <- filterM (\_ _ -> sample (bernoulli ((1 - pd) / (1 - survivalProb * pd)))) notObserved
  return ((mconcat (observed : coasted : zipWith M.singleton [nextTrackID..] newTrackPVs),
           nextTrackID + length newTrackPVs), obsDists)
  where
  (observed, notObserved) = M.partitionWithKey (\i _ -> Track i `elem` assocs) updatedOldTracks

sampleStep :: DelayedSample m => (TrackMap, Int) -> Gen m ((TrackMap, Int), [Expr (R 2)])
sampleStep (allOldTracks, nextTrackID) = do
  updatedOldTracks <- lift (mapM (trackMotion tdiff) allOldTracks)
  assocs <- "assocs" ~~ MP.prim (shuffleWithRepeats' countDistrs)
  (newTrackInfo, obsDists) <- lift (updateWithAssocs nextTrackID updatedOldTracks assocs)
  observations <- "observations" ~~ MP.isequence (map MP.dsPrim obsDists)
  return (newTrackInfo, observations)
  where
  countDistrs = M.singleton Clutter (poisson clutterLambda)
             <> M.singleton NewTrack (poisson newTrackLambda)
             <> M.mapKeys Track (fmap (\_ -> bernoulli01 (survivalProb * pd)) allOldTracks)

observeStep :: DelayedInfer m => (TrackMap, Int) -> [R 2] -> m ((TrackMap, Int), [Expr (R 2)])
observeStep (allOldTracks, nextTrackID) observations = do
  updatedOldTracks <- mapM (trackMotion tdiff) allOldTracks
  -- assocs <- proposeAssocs observations updatedOldTracks
  MP.observingWithProposal
    ("observations" |-> MP.trList observations) (sampleStep (allOldTracks, nextTrackID))
    "assocs" (proposeAssocs 0 observations updatedOldTracks)

proposeAssocs :: DelayedInfer m => Int -> [R 2] -> TrackMap -> Gen m [ObsType]
proposeAssocs j (obs : observations) remainingTracks = do
  newTrack <- lift $ newTrackD
  let distrs = M.singleton Clutter (clutterLambda, clutterDistr)
        <> M.singleton NewTrack (newTrackLambda, trackMeasurement newTrack)
        <> M.mapKeys Track (fmap (\t -> (survivalProb * pd, trackMeasurement t)) remainingTracks)
  likes <- lift $ mapM (\(intensity, d) -> (\ll -> intensity * exp ll) <$> DS.score d obs) distrs
  let assocDistr = let probs = fmap (/ sum likes) likes in categoricalM probs
  i <- ("assoc" ++ show j) ~~ MP.prim assocDistr
  let remainingTracks' = case i of
        Track k -> M.delete k remainingTracks
        _ -> remainingTracks
  (i :) <$> proposeAssocs (j + 1) observations remainingTracks'
proposeAssocs _ [] remainingTracks = return []

generateGroundTruth :: DelayedSample m => ZStream m () (TrackMap, [Expr (R 2)])
generateGroundTruth = ZS.fromStep stepf initState where
  initState = (mempty, 0)
  stepf state () = do
    (newState@(tracks', _), observations) <- MP.sim (sampleStep state)
    return (newState, (tracks', observations))

processObservations :: DelayedInfer m => ZStream m [R 2] (M.Map Int (Result (R 4)))
processObservations = ZS.fromStep stepf initState where
  initState = (mempty, 0)
  stepf state obs = do
    (newState@(tracks', _), _) <- observeStep state obs
    marginalTracks <- mapM (fmap fromJust . marginal) tracks'
    return (newState, marginalTracks)

runMTTPF :: Int -> ZStream Sampler () ([(Int, R 4)], [[(Int, Result (R 4))]], [R 2])
runMTTPF numParticles = proc () -> do
  (groundTruth, obs) <- zdeepForce generateGroundTruth -< ()
  particles <- zdsparticles numParticles processObservations -< obs
  returnA -< (M.assocs groundTruth, map M.assocs particles, obs)


{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE RecordWildCards  #-}

module Network.Layer
( LayerDefinition(..)
, Layer(..)
, Connectivity
, RandomTransform

, showableToLayer

, createLayer
, scaleLayer
, connectFully

, randomList
, boxMuller
, normals
, uniforms
, boundedUniforms
) where

import           Data.Binary           (Binary (..), decode, encode)
import           Network.Neuron
import           Numeric.LinearAlgebra
import           System.Random

-- | The LayerDefinition type is an intermediate type initialized by the
--   library user to define the different layers of the network.
data LayerDefinition = LayerDefinition { neuronDef   :: Neuron
                                       , neuronCount :: Int
                                       , connect     :: Connectivity
                                       }

-- | The Layer type, which stores the weight matrix, the bias matrix, and
--   a neuron type.
data Layer = Layer { weightMatrix :: Matrix Double
                   , biasVector   :: Vector Double
                   , neuron       :: Neuron
                   } deriving Show

instance Binary (Layer) where
  put Layer{..} = do put weightMatrix; put biasVector
  get = do weightMatrix <- get; biasVector <- get; return Layer{..}

-- | Connectivity is the type alias for a function that defines the connective
--   matrix for two layers (fully connected, convolutionally connected, etc.)
type Connectivity = Int -> Int -> Matrix Double

-- | A random transformation type alias. It is a transformation defined on an
--   infinite list of uniformly distributed random numbers, and returns a list
--   distributed on the transforming distribution.
type RandomTransform = [Double] -> [Double]

-- | The createLayer function takes in a random transformation on an infinite
--   stream of uniformly generated numbers, a source of entropy, and two
--   layer definitions, one for the previous layer and one for the next layer.
--   It returns a layer defined by the Layer type -- a weight matrix, a bias
--   vector, and a neuron type.
createLayer :: (RandomGen g)
  => RandomTransform -> g -> LayerDefinition -> LayerDefinition -> Layer
createLayer t g layerDef layerDef' =
  Layer (randomMatrix * (connectivity i j))
        (randomVector * bias)
        (neuronDef layerDef)
  where randomMatrix = (i >< j) (randomList t g')
        randomVector = i |> (randomList t g'')
        i = neuronCount layerDef'
        j = neuronCount layerDef
        connectivity = connect layerDef'
        bias = i |> (repeat 1) -- bias connectivity (full)
        (g', g'') = split g

scaleLayer :: Double -> Layer -> Layer
scaleLayer factor l =
  Layer (factor `scale` (weightMatrix l)) (factor `scale` (biasVector l)) (neuron l)

-- | The connectFully function takes the number of input neurons for a layer, i,
--   and the number of output neurons of a layer, j, and returns an i x j
--   connectivity matrix for a fully connected network.
connectFully :: Int -> Int -> Matrix Double
connectFully i j = (i >< j) (repeat 1)

-- | To go from a showable to a layer, we also need a neuron type,
--   which is an unfortunate restriction owed to Haskell's inability to
--   serialize functions.
showableToLayer :: (Layer, LayerDefinition) -> Layer
showableToLayer (s, d) = Layer (weightMatrix s) (biasVector s) (neuronDef d)

-- | Initialize an infinite random list given a random transform and a source
--   of entroy.
randomList :: RandomGen g => RandomTransform -> g -> [Double]
randomList transform = transform . randoms

-- | Define a transformation on the uniform distribution to generate
--   normally distributed numbers in Haskell (the Box-Muller transform)
boxMuller :: Double -> Double -> (Double, Double)
boxMuller x1 x2 = (z1, z2)
  where z1 = sqrt ((-2) * log x1) * cos (2 * pi * x2)
        z2 = sqrt ((-2) * log x1) * sin (2 * pi * x2)

-- | This is a function of type RandomTransform that transforms a list of
--   uniformly distributed numbers to a list of normally distributed numbers.
normals :: [Double] -> [Double]
normals (x1:x2:xs) = z1:z2:(normals xs)
  where (z1, z2) = boxMuller x1 x2
normals _ = []

-- | A non-transformation to return a list of uniformly distributed numbers
--   from a list of uniformly distributed numbers. It's really a matter of
--   naming consistency. It generates numbers on the range (0, 1]
uniforms :: [Double] -> [Double]
uniforms xs = xs

-- | An affine transformation to return a list of uniforms on the range
--   (a, b]
boundedUniforms :: (Double, Double) -> [Double] -> [Double]
boundedUniforms (lower, upper) xs = map affine xs
  where affine x = lower + x * (upper - lower)


{-# LANGUAGE CPP #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE TemplateHaskell #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE DeriveDataTypeable, DeriveGeneric #-}
-- |
-- Module    : Statistics.Types
-- Copyright : (c) 2009 Bryan O'Sullivan
-- License   : BSD3
--
-- Maintainer  : bos@serpentine.com
-- Stability   : experimental
-- Portability : portable
--
-- Data types common used in statistics
module Statistics.Types
    ( -- * Confidence level
      CL
      -- ** Accessors
    , confidenceLevel
    , significanceLevel
      -- ** Constructors
    , mkCL
    , mkCLE
    , mkCLFromSignificance
    , mkCLFromSignificanceE
      -- ** Constants and conversion to nσ
    , cl90
    , cl95
    , cl99
      -- *** Normal approximation
    , nSigma
    , nSigma1
    , getNSigma
    , getNSigma1
      -- * p-value
    , PValue
      -- ** Accessors
    , pValue
      -- ** Constructors
    , mkPValue
    , mkPValueE
      -- * Estimates and upper/lower limits
    , Estimate(..)
    , NormalErr(..)
    , ConfInt(..)
    , UpperLimit(..)
    , LowerLimit(..)
      -- ** Constructors
    , estimateNormErr
    , (±)
    , estimateFromInterval
    , estimateFromErr
      -- ** Accessors
    , confidenceInterval
    , asymErrors
    , Scale(..)
      -- * Other
    , Sample
    , WeightedSample
    , Weights
    ) where

import Control.Monad                ((<=<), liftM2, liftM3)
import Control.DeepSeq              (NFData(..))
import Data.Aeson                   (FromJSON(..), ToJSON)
import Data.Binary                  (Binary(..))
import Data.Data                    (Data,Typeable)
import Data.Maybe                   (fromMaybe)
import Data.Vector.Unboxed          (Unbox)
import Data.Vector.Unboxed.Deriving (derivingUnbox)
import GHC.Generics                 (Generic)

#if __GLASGOW_HASKELL__ == 704
import qualified Data.Vector.Generic
import qualified Data.Vector.Generic.Mutable
#endif

import Statistics.Internal
import Statistics.Types.Internal
import Statistics.Distribution
import Statistics.Distribution.Normal

----------------------------------------------------------------
-- Data type for confidence level
----------------------------------------------------------------

-- |
-- Confidence level. In context of confidence intervals it's
-- probability of said interval covering true value of measured
-- value. In context of statistical tests it's @1-α@ where α is
-- significance of test.
--
-- Since confidence level are usually close to 1 they are stored as
-- @1-CL@ internally. There are two smart constructors for @CL@:
-- 'mkCL' and 'mkCLFromSignificance' (and corresponding variant
-- returning @Maybe@). First creates @CL@ from confidence level and
-- second from @1 - CL@ or significance level.
--
-- >>> cl95
-- mkCLFromSignificance 0.05
--
-- Prior to 0.14 confidence levels were passed to function as plain
-- @Doubles@. Use 'mkCL' to convert them to @CL@.
newtype CL a = CL a
               deriving (Eq, Typeable, Data, Generic)

instance Show a => Show (CL a) where
  showsPrec n (CL p) = defaultShow1 "mkCLFromSignificance" p n
instance (Num a, Ord a, Read a) => Read (CL a) where
  readPrec = defaultReadPrecM1 "mkCLFromSignificance" mkCLFromSignificanceE

instance (Binary a, Num a, Ord a) => Binary (CL a) where
  put (CL p) = put p
  get        = maybe (fail errMkCL) return . mkCLFromSignificanceE =<< get

instance (ToJSON a)                 => ToJSON   (CL a)
instance (FromJSON a, Num a, Ord a) => FromJSON (CL a) where
  parseJSON = maybe (fail errMkCL) return . mkCLFromSignificanceE <=< parseJSON

instance NFData   a => NFData   (CL a) where
  rnf (CL a) = rnf a

-- |
-- >>> cl95 > cl90
-- True
instance Ord a => Ord (CL a) where
  CL a <  CL b = a >  b
  CL a <= CL b = a >= b
  CL a >  CL b = a <  b
  CL a >= CL b = a <= b
  max (CL a) (CL b) = CL (min a b)
  min (CL a) (CL b) = CL (max a b)

-- | Create confidence level from probability β or probability
--   confidence interval contain true value of estimate. Will throw
--   exception if parameter is out of [0,1] range
--
-- >>> mkCL 0.95    -- same as cl95
-- mkCLFromSignificance 0.05
mkCL :: (Ord a, Num a) => a -> CL a
mkCL
  = fromMaybe (error "Statistics.Types.mkCL: probability is out if [0,1] range")
  . mkCLE

-- | Same as 'mkCL' but returns @Nothing@ instead of error if
--   parameter is out of [0,1] range
--
-- >>> mkCLE 0.95    -- same as cl95
-- Just (mkCLFromSignificance 0.05)
mkCLE :: (Ord a, Num a) => a -> Maybe (CL a)
mkCLE p
  | p >= 0 && p <= 1 = Just $ CL (1 - p)
  | otherwise        = Nothing

-- | Create confidence level from probability α or probability that
--   confidence interval does not contain true value of estimate. Will
--   throw exception if parameter is out of [0,1] range
--
-- >>> mkCLFromSignificance 0.05    -- same as cl95
-- mkCLFromSignificance 0.05
mkCLFromSignificance :: (Ord a, Num a) => a -> CL a
mkCLFromSignificance = fromMaybe (error errMkCL) . mkCLFromSignificanceE

-- | Same as 'mkCLFromSignificance' but returns @Nothing@ instead of error if
--   parameter is out of [0,1] range
--
-- >>> mkCLFromSignificanceE 0.05    -- same as cl95
-- Just (mkCLFromSignificance 0.05)
mkCLFromSignificanceE :: (Ord a, Num a) => a -> Maybe (CL a)
mkCLFromSignificanceE p
  | p >= 0 && p <= 1 = Just $ CL p
  | otherwise        = Nothing

errMkCL :: String
errMkCL = "Statistics.Types.mkPValCL: probability is out if [0,1] range"

-- | Get confidence level. This function is subject to rounding
--   errors. If @1 - CL@ is needed use 'significanceLevel' instead
confidenceLevel :: (Num a) => CL a -> a
confidenceLevel (CL p) = 1 - p

-- | Get significance level.
significanceLevel :: CL a -> a
significanceLevel (CL p) = p

-- | 90% confidence level
cl90 :: Fractional a => CL a
cl90 = CL 0.10

-- | 95% confidence level
cl95 :: Fractional a => CL a
cl95 = CL 0.05

-- | 99% confidence level
cl99 :: Fractional a => CL a
cl99 = CL 0.01

----------------------------------------------------------------
-- Data type for p-value
----------------------------------------------------------------

-- | Newtype wrapper for p-value.
newtype PValue a = PValue a
               deriving (Eq,Ord, Typeable, Data, Generic)

instance Show a => Show (PValue a) where
  showsPrec n (PValue p) = defaultShow1 "mkPValue" p n
instance (Num a, Ord a, Read a) => Read (PValue a) where
  readPrec = defaultReadPrecM1 "mkPValue" mkPValueE

instance (Binary a, Num a, Ord a) => Binary (PValue a) where
  put (PValue p) = put p
  get            = maybe (fail errMkPValue) return . mkPValueE =<< get

instance (ToJSON a)                 => ToJSON   (PValue a)
instance (FromJSON a, Num a, Ord a) => FromJSON (PValue a) where
  parseJSON = maybe (fail errMkPValue) return . mkPValueE <=< parseJSON

instance NFData a => NFData (PValue a) where
  rnf (PValue a) = rnf a

-- | Construct PValue. Throws error if argument is out of [0,1] range.
--
mkPValue :: (Ord a, Num a) => a -> PValue a
mkPValue = fromMaybe (error errMkPValue) . mkPValueE

-- | Construct PValue. Returns @Nothing@ if argument is out of [0,1] range.
mkPValueE :: (Ord a, Num a) => a -> Maybe (PValue a)
mkPValueE p
  | p >= 0 && p <= 1 = Just $ PValue p
  | otherwise        = Nothing

-- | Get p-value
pValue :: PValue a -> a
pValue (PValue p) = p

-- | P-value expressed in sigma. This is convention widely used in
--   experimental physics. N sigma confidence level corresponds to
--   probability within N sigma of normal distribution.
--
--   Note that this correspondence is for normal distribution. Other
--   distribution will have different dependency. Also experimental
--   distribution usually only approximately normal (especially at
--   extreme tails).
nSigma :: Double -> PValue Double
nSigma n
  | n > 0     = PValue $ 2 * cumulative standard (-n)
  | otherwise = error "Statistics.Extra.Error.nSigma: non-positive number of sigma"

-- | P-value expressed in sigma for one-tail hypothesis. This correspond to
--   probability of obtaining value less than @N·σ@.
nSigma1 :: Double -> PValue Double
nSigma1 n
  | n > 0     = PValue $ cumulative standard (-n)
  | otherwise = error "Statistics.Extra.Error.nSigma1: non-positive number of sigma"

-- | Express confidence level in sigmas
getNSigma :: PValue Double -> Double
getNSigma (PValue p) = negate $ quantile standard (p / 2)

-- | Express confidence level in sigmas for one-tailed hypothesis.
getNSigma1 :: PValue Double -> Double
getNSigma1 (PValue p) = negate $ quantile standard p

errMkPValue :: String
errMkPValue = "Statistics.Types.mkPValue: probability is out if [0,1] range"

----------------------------------------------------------------
-- Point estimates
----------------------------------------------------------------

-- |
-- A point estimate and its confidence interval. It's parametrized by
-- both error type @e@ and value type @a@. This module provides two
-- types of error: 'NormalErr' for normally distributed errors and
-- 'ConfInt' for error with normal distribution. See their
-- documentation for more details.
--
-- For example @144 ± 5@ (assuming normality) could be expressed as
--
-- > Estimate { estPoint = 144
-- >          , estError = NormalErr 5
-- >          }
--
-- Or if we want to express @144 + 6 - 4@ at CL95 we could write:
--
-- > Estimate { estPoint = 144
-- >          , estError = ConfInt
-- >                       { confIntLDX = 4
-- >                       , confIntUDX = 6
-- >                       , confIntCL  = cl95
-- >                       }
--
-- Prior to statistics 0.14 @Estimate@ data type used following definition:
--
-- > data Estimate = Estimate {
-- >      estPoint           :: {-# UNPACK #-} !Double
-- >    , estLowerBound      :: {-# UNPACK #-} !Double
-- >    , estUpperBound      :: {-# UNPACK #-} !Double
-- >    , estConfidenceLevel :: {-# UNPACK #-} !Double
-- >    }
--
-- Now type @Estimate ConfInt Double@ should be used instead. Function
-- 'estimateFromInterval' allow to easily construct estimate from same inputs.
data Estimate e a = Estimate
    { estPoint           :: !a
      -- ^ Point estimate.
    , estError           :: !(e a)
      -- ^ Confidence interval for estimate.
    } deriving (Eq, Read, Show, Generic
#if __GLASGOW_HASKELL__ >= 708
               , Typeable, Data
#endif
               )

instance (Binary   (e a), Binary   a) => Binary   (Estimate e a) where
  get = liftM2 Estimate get get
  put (Estimate ep ee) = put ep >> put ee
instance (FromJSON (e a), FromJSON a) => FromJSON (Estimate e a)
instance (ToJSON   (e a), ToJSON   a) => ToJSON   (Estimate e a)
instance (NFData   (e a), NFData   a) => NFData   (Estimate e a) where
    rnf (Estimate x dx) = rnf x `seq` rnf dx

-- |
-- Normal errors. They are stored as 1σ errors which corresponds to
-- 68.8% CL. Since we can recalculate them to any confidence level if
-- needed we don't store it.
newtype NormalErr a = NormalErr
  { normalError :: a
  }
  deriving (Eq, Read, Show, Typeable, Data, Generic)

instance Binary   a => Binary   (NormalErr a) where
  get = fmap NormalErr get
  put = put . normalError
instance FromJSON a => FromJSON (NormalErr a)
instance ToJSON   a => ToJSON   (NormalErr a)
instance NFData   a => NFData   (NormalErr a) where
    rnf (NormalErr x) = rnf x

-- | Confidence interval. It assumes that confidence interval forms
--   single interval and isn't set of disjoint intervals.
data ConfInt a = ConfInt
  { confIntLDX :: !a
    -- ^ Lower error estimate, or distance between point estimate and
    --   lower bound of confidence interval.
  , confIntUDX :: !a
    -- ^ Upper error estimate, or distance between point estimate and
    --   upper bound of confidence interval.
  , confIntCL  :: !(CL Double)
    -- ^ Confidence level corresponding to given confidence interval.
  }
  deriving (Read,Show,Eq,Typeable,Data,Generic)

instance Binary   a => Binary   (ConfInt a) where
  get = liftM3 ConfInt get get get
  put (ConfInt l u cl) = put l >> put u >> put cl
instance FromJSON a => FromJSON (ConfInt a)
instance ToJSON   a => ToJSON   (ConfInt a)
instance NFData   a => NFData   (ConfInt a) where
    rnf (ConfInt x y _) = rnf x `seq` rnf y

----------------------------------------
-- Constructors

-- | Create estimate with normal errors
estimateNormErr :: a            -- ^ Point estimate
                -> a            -- ^ 1σ error
                -> Estimate NormalErr a
estimateNormErr x dx = Estimate x (NormalErr dx)

-- | Synonym for 'estimateNormErr'
(±) :: a      -- ^ Point estimate
    -> a      -- ^ 1σ error
    -> Estimate NormalErr a
(±) = estimateNormErr

-- | Create estimate with asymmetric error.
estimateFromErr
  :: a                     -- ^ Central estimate
  -> (a,a)                 -- ^ Lower and upper errors. Both should be
                           --   positive but it's not checked.
  -> CL Double             -- ^ Confidence level for interval
  -> Estimate ConfInt a
estimateFromErr x (ldx,udx) cl = Estimate x (ConfInt ldx udx cl)

-- | Create estimate with asymmetric error.
estimateFromInterval
  :: Num a
  => a                     -- ^ Point estimate. Should lie within
                           --   interval but it's not checked.
  -> (a,a)                 -- ^ Lower and upper bounds of interval
  -> CL Double             -- ^ Confidence level for interval
  -> Estimate ConfInt a
estimateFromInterval x (lx,ux) cl
  = Estimate x (ConfInt (x-lx) (ux-x) cl)

----------------------------------------
-- Accessors

-- | Get confidence interval
confidenceInterval :: Num a => Estimate ConfInt a -> (a,a)
confidenceInterval (Estimate x (ConfInt ldx udx _))
  = (x - ldx, x + udx)

-- | Get asymmetric errors
asymErrors :: Estimate ConfInt a -> (a,a)
asymErrors (Estimate _ (ConfInt ldx udx _)) = (ldx,udx)

-- | Data types which could be multiplied by constant.
class Scale e where
  scale :: (Ord a, Num a) => a -> e a -> e a

instance Scale NormalErr where
  scale a (NormalErr e) = NormalErr (abs a * e)

instance Scale ConfInt where
  scale a (ConfInt l u cl) | a >= 0    = ConfInt  (a*l)  (a*u) cl
                           | otherwise = ConfInt (-a*u) (-a*l) cl

instance Scale e => Scale (Estimate e) where
  scale a (Estimate x dx) = Estimate (a*x) (scale a dx)

----------------------------------------------------------------
-- Upper/lower limit
----------------------------------------------------------------

-- | Upper limit. They are usually given for small non-negative values
--   when it's not possible detect difference from zero.
data UpperLimit a = UpperLimit
    { upperLimit        :: !a
      -- ^ Upper limit
    , ulConfidenceLevel :: !(CL Double)
      -- ^ Confidence level for which limit was calculated
    } deriving (Eq, Read, Show, Typeable, Data, Generic)

instance Binary   a => Binary   (UpperLimit a) where
  get = liftM2 UpperLimit get get
  put (UpperLimit l cl) = put l >> put cl
instance FromJSON a => FromJSON (UpperLimit a)
instance ToJSON   a => ToJSON   (UpperLimit a)
instance NFData   a => NFData   (UpperLimit a) where
    rnf (UpperLimit x cl) = rnf x `seq` rnf cl

-- | Lower limit. They are usually given for large quantities when
--   it's not possible to measure them. For example: proton half-life
data LowerLimit a = LowerLimit {
    lowerLimit        :: !a
    -- ^ Lower limit
  , llConfidenceLevel :: !(CL Double)
    -- ^ Confidence level for which limit was calculated
  } deriving (Eq, Read, Show, Typeable, Data, Generic)

instance Binary   a => Binary   (LowerLimit a) where
  get = liftM2 LowerLimit get get
  put (LowerLimit l cl) = put l >> put cl
instance FromJSON a => FromJSON (LowerLimit a)
instance ToJSON   a => ToJSON   (LowerLimit a)
instance NFData   a => NFData   (LowerLimit a) where
    rnf (LowerLimit x cl) = rnf x `seq` rnf cl

----------------------------------------------------------------
-- Deriving unbox instances
----------------------------------------------------------------

derivingUnbox "CL"
  [t| forall a. Unbox a => CL a -> a |]
  [| \(CL a) -> a |]
  [| CL           |]

derivingUnbox "PValue"
  [t| forall a. Unbox a => PValue a -> a |]
  [| \(PValue a) -> a |]
  [| PValue           |]

derivingUnbox "Estimate"
  [t| forall a e. (Unbox a, Unbox (e a)) => Estimate e a -> (a, e a) |]
  [| \(Estimate x dx) -> (x,dx) |]
  [| \(x,dx) -> (Estimate x dx) |]

derivingUnbox "NormalErr"
  [t| forall a. Unbox a => NormalErr a -> a |]
  [| \(NormalErr a) -> a |]
  [| NormalErr           |]

derivingUnbox "ConfInt"
  [t| forall a. Unbox a => ConfInt a -> (a, a, CL Double) |]
  [| \(ConfInt a b c) -> (a,b,c) |]
  [| \(a,b,c) -> ConfInt a b c   |]

derivingUnbox "UpperLimit"
  [t| forall a. Unbox a => UpperLimit a -> (a, CL Double) |]
  [| \(UpperLimit a b) -> (a,b) |]
  [| \(a,b) -> UpperLimit a b   |]

derivingUnbox "LowerLimit"
  [t| forall a. Unbox a => LowerLimit a -> (a, CL Double) |]
  [| \(LowerLimit a b) -> (a,b) |]
  [| \(a,b) -> LowerLimit a b   |]


module GenerateSet where

import Data.Word
import Data.List.Split
import Data.Scientific as Scientific
import Data.Complex
import Control.Monad
import Codec.Picture
import Data.List
import Data.Time.Clock
import Data.Time.Calendar
import Graphics.X11.Xlib
import System.Exit (exitWith, ExitCode(..))
import Control.Concurrent
import Data.Bits
import GHC.Conc (numCapabilities)

import Mand

generateSet mv = do
  step <- promptForImput mv "enter step"
  sr <- promptForImput mv "enter starting real"
  nr <- promptForImput mv "enter no of real points"
  si <- promptForImput mv "starting im"
  ni <- promptForImput mv "no of im"
  name <- promptForImput mv "Enter name:"

  let stepf = read step :: Double
  let srf = read sr :: Double -- starting real
  let nrf = read nr :: Double --number of real points
  let sif = read si :: Double -- starting imaginary
  let nif = read ni :: Double -- number of imaginary points

  d <- getCurrentTime
  desc <- promptForImput mv "description:"
  appendFile "mand/log.txt" $ "\n"++(take 19 $ show d)++", "++name++", "++desc

  --writePng name $ generateImage (\x y -> PixelRGB8 (fromIntegral 0) (fromIntegral 0) (fromIntegral $ mand $ ((read sr :: S) + (read (show x) :: S) * (read step :: S)) :+ ((read si :: S) + (read (show y) :: S) * (read step :: S)))) (read nr :: Int) (read ni :: Int)
  writePng ("mand/"++name) $ generateImage (\x y -> generatingFunction stepf srf sif (sif+ nif * stepf) x y ) (read nr :: Int) (read ni :: Int)
--writePng "mandv1" $ generateImage (\x y -> PixelRGB8 (fromIntegral 100) (fromIntegral 100) (fromIntegral 0)) 300 200
  print "completed"
  comment <- promptForImput mv "add comments:"
  appendFile "mand/log.txt" $ ", "++comment

generatingFunction stepf srf sif realend x y = normal stepf srf sif realend x y
  where
    m = general mand_iteration Nothing ((srf + (fromIntegral x) * stepf) :+ (realend - ((fromIntegral y) * stepf))) 255
    normal stepf srf sif realend x y = PixelRGB8 (fromIntegral 0)  (fromIntegral 0) (fromIntegral m)
    fullcolour stepf srf sif realend x y = PixelRGB8 p1 p2 p3
    max = 10
    scale x = round $ 255 * ((fromIntegral x) / (fromIntegral (max * max * max -1)))
    --p1 = scale $ fromIntegral $ max * max * (div m max* max)
    --p2 = scale $ fromIntegral $ max * max * (mod (div m max) max * max)
    --p3 = scale $ max * max * (fromIntegral $ mod m max)
    p1 = fromIntegral $ div (scaleUp 5000 m) 65536
    p2 = fromIntegral $ mod (div (scaleUp 5000 m) 255) 256
    p3 = fromIntegral $ mod (scaleUp 5000 m) 255

scaleUp :: Integral a => a -> a -> a
scaleUp max i = round $ 16777215 * fromIntegral i / fromIntegral max
    --m :: Int
    --m = mand $ Right $ fromFloatDigits (srf + (fromIntegral x) * stepf) :+ fromFloatDigits (sif + (fromIntegral y) * stepf)
    --m = mand $ ( (read :: String -> Scientific) (show (srf + (fromIntegral x) * stepf)) :=  (read :: String -> Scientific) (show (sif+ (fromIntegral y) * stepf))
    --m = if (mand $ (srf + (fromIntegral x) * stepf) :+ (sif + (fromIntegral y) * stepf) )== 255  then 0 else 255

--test = foldl (\acc x -> 'b':acc) []  [1..100]

colCorect1 :: Int -> Int
colCorect1 x = floor $ 256*(1 - exp(fromIntegral x/40))

colCorect2 :: Int -> Int
colCorect2 x = floor . (/255) . fromIntegral  $  x * x

promptForImput :: MVar Bool -> String ->  IO String
promptForImput mv message = do
  takeMVar mv
  print message
  input <- getLine
  putMVar mv True
  if input == "*" then promptForImput mv message else return input


{-# LANGUAGE ScopedTypeVariables #-}

module AI.Learning.NeuralNetwork
    (
      -- * Representation
      NeuralNetwork
      -- * Prediction and Training
    , nnPredict
    , nnTrain
    , nnTrainIO
    ) where

import Control.Monad.Random hiding (fromList)
import Numeric.LinearAlgebra
import Numeric.LinearAlgebra.Util
import System.IO.Unsafe

import AI.Learning.Core
import AI.Util.Matrix

-----------------------
-- NN Representation --
----------------------

-- |Representation for a single-hidden-layer neural network. If the network
--  has K input nodes, H hidden nodes and L output layers then the dimensions
--  of the matrices theta0 and theta1 are
--
--  * size Layer1 = (K+1) x H
--  * size Layer2 = (H+1) x L
--
--  the (+1)s account for the addition of bias nodes in the input layer and
--  the hidden layer. Therefore the total number of parameters is
--  H * (K + L) + H + L
data NeuralNetwork = NN (Matrix Double) (Matrix Double)

-- |Three-tuple describing the shape of a network: (input, hidden, output).
type NNShape = (Int,Int,Int)

instance Show NeuralNetwork where
    show (NN theta0 theta1) = "Neural Net:\n\n" ++ dispf 3 theta0 ++
                                           "\n" ++ dispf 3 theta1

fromVector :: NNShape -> Vector Double -> NeuralNetwork
fromVector (k,h,l) vec = NN theta0 theta1
    where theta0 = reshape h $ takeVector ((k + 1) * h) vec
          theta1 = reshape l $ dropVector ((k + 1) * h) vec

toVector :: Matrix Double -> Matrix Double -> Vector Double
toVector theta0 theta1 = join [flatten theta0, flatten theta1]

----------------------
-- NN Train/Predict --
----------------------

-- |Make predictions using a neural network.
nnPredict :: NeuralNetwork -> Matrix Double -> Matrix Double
nnPredict nn x = h where (_,_,h) = nnForwardProp nn x

-- |Train a neural network from a training set. Note that you must supply
--  an initial vector of weights. Supplying initial weights all equal to
--  zero will generally give poor performance.
nnTrain :: Int              -- number of hidden neurons
        -> Matrix Double    -- y
        -> Matrix Double    -- x
        -> Double           -- lambda
        -> Vector Double    -- initial weights
        -> NeuralNetwork
nnTrain h y x lambda initialVec = fromVector shape vec
    where shape = (cols x, h, cols y)
          vec   = minimizeS cost grad initialVec
          cost  = fst . nnCostGradient shape y x lambda
          grad  = snd . nnCostGradient shape y x lambda

-- |Train a neural network from a training set.
--
--  Note that this is implemented as an IO action, because it randomly sets the
--  weights in the network before performing numerical optimization (this is to
--  break the symmetry that exists in the network when all weights are zero).
--  As a result, it is possible to get different results when training the
--  network multiple times with the same data.
nnTrainIO :: Int              -- number of hidden neurons
          -> Matrix Double    -- y
          -> Matrix Double    -- x
          -> Double           -- lambda
          -> IO NeuralNetwork
nnTrainIO h y x lambda = nnTrain h y x lambda `fmap` initialVec (cols x, h, cols y)

-- |Choose initial random weights for a neural network.
initialVec :: NNShape -> IO (Vector Double)
initialVec (k,h,l) = do
    let len = h * (k + l) + h + l
    xs <- getRandomRs (0.0, 0.01)
    return . fromList $ take len xs

------------------------------
-- Forward/Back Propagation --
------------------------------

-- |Perform forward propagation through a neural network, returning the matrices
--  created in the process.
nnForwardProp :: NeuralNetwork                                 -- neural net
              -> Matrix Double                                 -- design matrix (x)
              -> (Matrix Double, Matrix Double, Matrix Double) -- results of fwd prop
nnForwardProp (NN theta0 theta1) x = (a0,a1,a2)
    where a0 = addOnes $ x
          a1 = addOnes $ sigmoid (a0 <> theta0)
          a2 = sigmoid (a1 <> theta1)

-- |Perform backward propagiation through a neural network. You must supply the
--  target values and the results of forward propagation for each layer, and
--  the function returns the gradient matrices for the neural network.
nnBackProp :: NeuralNetwork                                 -- neural net
           -> Matrix Double                                 -- target (y)
           -> (Matrix Double, Matrix Double, Matrix Double) -- results of fwd prop
           -> (Matrix Double, Matrix Double)                -- gradient (delta0, delta1)
nnBackProp (NN _ theta1) y (a0,a1,a2) = (dropColumns 1 delta0, delta1)
    where
        d2     = a2 - y
        d1     = (d2 <> trans theta1) * a1 * (1 - a1)
        delta0 = trans a0 <> d1
        delta1 = trans a1 <> d2

-- |Perform back and forward propagation through a neural network, returning the
--  final predictions (variable /a2/) and the gradient matrices (variables
--  /delta0/ and /delta1/) produced.
nnFwdBackProp :: NeuralNetwork -> Matrix Double -> Matrix Double -> (Matrix Double, Matrix Double, Matrix Double)
nnFwdBackProp nn@(NN theta0 theta1) y x = (a2, delta0, delta1)
    where
        (a0,a1,a2)      = nnForwardProp nn x
        (delta0,delta1) = nnBackProp nn y (a0,a1,a2)

-- |Compute the penalty function and gradient vector for a neural network
--  given a training set.
nnCostGradient :: NNShape                   -- (K,H,L)
               -> Matrix Double             -- targets (y)
               -> Matrix Double             -- design matrix (x)
               -> Double                    -- regularization parameter (lambda)
               -> Vector Double             -- neural network
               -> (Double, Vector Double)   -- (cost, gradient)
nnCostGradient shape y x lambda vec = (cost, grad)
    where
        m = fromIntegral (rows x)
        nn@(NN theta0 theta1) = fromVector shape vec
        (h, delta0, delta1)   = nnFwdBackProp nn y x

        cost  = (cost1 + cost2) / m
        cost1 = negate $ sumMatrix $ y * log h + (1-y) * log (1-h)
        cost2 = lambda/2 * (normMatrix theta0 + normMatrix theta1)

        grad  = (1/m) `scale` (grad1 + grad2)
        grad1 = toVector delta0 delta1
        grad2 = lambda `scale` toVector (insertNils theta0) (insertNils theta1)

        normMatrix m = sumMatrix $ (dropRows 1 m) ^ 2
        insertNils m = vertcat [0, dropRows 1 m]

nnCost shape y x lambda = fst . nnCostGradient shape y x lambda
nnGrad shape y x lambda = snd . nnCostGradient shape y x lambda

-- |Use central differencing to compute an approximation to the gradient
--  vector for a neural network. This is mainly used for checking the
--  implementation of backprop.
nnGradApprox :: NNShape -> Matrix Double -> Matrix Double -> Double -> Vector Double -> Vector Double
nnGradApprox shape y x lambda vec = fromList $ g `map` [0..n-1]
    where
        h = 1e-6
        n = dim vec
        f v = nnCost shape y x lambda v
        g i = (f (vec + e i) - f (vec - e i)) / (2*h)
        e i = fromList $ replicate i 0 ++ [h] ++ replicate (n-i-1) 0

-------------
-- Testing --
-------------

testNN :: NeuralNetwork
testNN = NN t0 t1
    where t0 = fromLists [[11.9934, -5.1396], [-7.7162, 10.1512], [-7.668, 10.1835]]
          t1 = fromLists [[-16.8806], [10.0445], [8.7476]]

testFwdProp :: IO ()
testFwdProp = do
    putStrLn "***\nCompare forward propagation to the MATLAB implementation.\n"
    let theta0 = fromLists [[0], [-10]]
        theta1 = fromLists [[5],[-10]]
        nn = NN theta0 theta1
        x = fromLists [[-1],[0],[1]]
    -- sigmoid [10, 0, -10] = [1, 0.5, 0]
    -- sigmoid [-5, 0, 5]  = [0.0, 0.5, 1]
    let y = nnPredict nn x
    putStrLn "Predictions (should be roughly 0.0, 0.5, 1.0)"
    disp 2 y

testBackProp :: IO ()
testBackProp = do
    putStrLn "***\nCompare back propagation to the MATLAB implementation.\n"
    let nn = testNN
        x  = fromLists [[0, 0], [0, 1], [1, 0], [1, 1]]
        y  = fromLists [[0], [1], [1], [0]]
        (h, delta0, delta1) = nnFwdBackProp nn y x
    putStrLn "Predictions (should be 0.0011, 0.8487, 0.8476, 0.0004)"
    disp 4 h
    putStrLn "Delta 0 (should be -0.0401, -0.0171, -0.0205, -0.0088, -0.0195, -0.0084)"
    disp 4 delta0
    putStrLn "Delta 1 (should be -0.3022, -0.2985, -0.3013)"
    disp 4 delta1

testCostGradient :: IO ()
testCostGradient = do
    putStrLn "***\nTest cost/gradient against the MATLAB implementation.\n"
    let nn@(NN theta0 theta1) = testNN
        x = fromLists [[0, 0], [0, 1], [1, 0], [1, 1]]
        y = fromLists [[0], [1], [1], [0]]
        lambda = 1e-4
        (cost, grad) = nnCostGradient (2,2,1) y x lambda (toVector theta0 theta1)
    putStrLn "Cost (should be around 0.0890)"
    print cost
    putStrLn "Gradient (should be -0.0100, -0.0043, -0.0053, -0.0019, -0.0051, -0.0019, -0.0755, -0.0744, -0.0751)"
    disp 4 (asColumn grad)

test :: Int -> Double -> IO ()
test n lambda = do
    putStrLn "***\nLearning XOR function.\n"
    x <- rand n 2
    e <- fmap (0.01*) (rand n 1)
    let y = mapMatrix (\x -> if x > 0.5 then 1.0 else 0.0) (xor x)
    nn <- nnTrainIO 4 y x lambda
    let ypred = nnPredict nn x
    -- Show predictions
    putStrLn "Predictions:"
    disp 2 $ takeRows 10 $ horzcat [x, y, ypred]
    -- Show neural net
    print nn
    -- Final test; should approximately compute xor function
    let xx = fromLists [[0,0],[0,1],[1,0],[1,1]]
        yy = nnPredict nn xx
    putStrLn "Exclusive or:"
    disp 2 $ horzcat [xx,yy]

xor :: Matrix Double -> Matrix Double
xor x = let [u,v] = toColumns x in asColumn (u + v - 2 * u * v)


{-# LANGUAGE GeneralizedNewtypeDeriving #-}

module Valuation
    ( module Valuation
    , throwError
    ) where

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

type Year = Int

newtype Scenario a = Scenario
    { scenario :: a
    }

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

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

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

type TimeValue s a = Year -> Value s a

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

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

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

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

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

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

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


{-# LANGUAGE ApplicativeDo         #-}
{-# LANGUAGE EmptyCase             #-}
{-# LANGUAGE FlexibleContexts      #-}
{-# LANGUAGE GADTs                 #-}
{-# LANGUAGE InstanceSigs          #-}
{-# LANGUAGE KindSignatures        #-}
{-# LANGUAGE LambdaCase            #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE RankNTypes            #-}
{-# LANGUAGE ScopedTypeVariables   #-}
{-# LANGUAGE StandaloneDeriving    #-}
{-# LANGUAGE TypeFamilies          #-}
{-# LANGUAGE TypeInType            #-}
{-# LANGUAGE TypeOperators         #-}
{-# LANGUAGE UndecidableInstances  #-}

module TensorOps.Backend.BTensor
  ( BTensor
  , BTensorL, BTensorV
  , HMat
  , HMatD
  ) where

import           Control.Applicative
import           Control.DeepSeq
import           Data.Distributive
import           Data.Kind
import           Data.List.Util
import           Data.Monoid
import           Data.Nested hiding             (unScalar, unVector, gmul')
import           Data.Singletons
import           Data.Singletons.Prelude hiding (Reverse, Head, sReverse, (:-))
import           Data.Type.Combinator
import           Data.Type.Combinator.Util
import           Data.Type.Length               as TCL
import           Data.Type.Length.Util          as TCL
import           Data.Type.Nat
import           Data.Type.Product              as TCP
import           Data.Type.Product.Util         as TCP
import           Data.Type.Sing
import           Data.Type.Uniform
import           Statistics.Distribution
import           TensorOps.BLAS
import           TensorOps.BLAS.HMat
import           TensorOps.NatKind
import           TensorOps.Types
import           Type.Class.Higher
import           Type.Class.Higher.Util
import           Type.Class.Witness
import           Type.Family.List
import           Type.Family.List.Util
import           Type.Family.Nat
import qualified Data.Type.Vector               as TCV
import qualified Data.Type.Vector.Util          as TCV
import qualified Data.Vector.Sized              as VS

data BTensor :: (k -> Type -> Type) -> (BShape k -> Type) -> [k] -> Type where
    BTS :: { unScalar :: !(ElemB b)     } -> BTensor v b '[]
    BTV :: { unVector :: !(b ('BV n))   } -> BTensor v b '[n]
    BTM :: { unMatrix :: !(b ('BM n m)) } -> BTensor v b '[n,m]
    BTN :: { unNested :: !(v n (BTensor v b (o ': m ': ns))) }
        -> BTensor v b (n ': o ': m ': ns)

type BTensorL = BTensor (Flip2 TCV.VecT   I)
type BTensorV = BTensor (Flip2 VS.VectorT I)

instance (Nesting Proxy Show v, Show1 b, Show (ElemB b)) => Show (BTensor v b s) where
    showsPrec p = \case
      BTS x  -> showParen (p > 10) $ showString "BTS "
                                   . showsPrec 11 x
      BTV xs -> showParen (p > 10) $ showString "BTV "
                                   . showsPrec1 11 xs
      BTM xs -> showParen (p > 10) $ showString "BTM "
                                   . showsPrec1 11 xs
      BTN xs -> showParen (p > 10) $ showString "BTN "
                                   . showsPrec' 11 xs
      where
        showsPrec' :: forall n s'. Int -> v n (BTensor v b s') -> ShowS
        showsPrec' p' xs = showsPrec p' xs
            \\ (nesting Proxy :: Show (BTensor v b s')
                              :- Show (v n (BTensor v b s'))
               )

instance (Nesting Proxy Show v, Show1 b, Show (ElemB b)) => Show1 (BTensor v b)

instance (NFData (ElemB b), NFData1 b, Nesting Proxy NFData v) => Nesting1 Proxy NFData (BTensor v b) where
    nesting1 _ = Wit

instance (NFData (ElemB b), NFData1 b, Nesting Proxy NFData v) => NFData (BTensor v b js) where
    rnf = \case
      BTS x  -> rnf  x
      BTV xs -> rnf1 xs
      BTM xs -> rnf1 xs
      BTN (xs :: v n (BTensor v b (o ': m ': ns))) ->
        rnf xs \\ (nesting Proxy :: NFData (BTensor v b (o ': m ': ns))
                                 :- NFData (v n (BTensor v b (o ': m ': ns)))
                  )

instance (NFData (ElemB b), NFData1 b, Nesting Proxy NFData v) => NFData1 (BTensor v b)

instance ( BLAS b
         , Vec v
         , Nesting1 Proxy Functor v
         , Nesting1 Sing Applicative v
         , SingI ns
         , Num (ElemB b)
         )
        => Num (BTensor v b ns) where
    (+) = zipBase sing
                  (+)
                  (\_          xs ys -> axpy 1 xs (Just ys))
                  (\(SBM _ sM) xs ys -> gemm 1 xs (eye sM) (Just (1, ys)))
    {-# INLINE (+) #-}
    (-) = zipBase sing
                  (-)
                  (\_          xs ys -> axpy (-1) ys (Just xs))
                  (\(SBM _ sM) xs ys -> gemm 1 xs (eye sM) (Just (-1, ys)))
    {-# INLINE (-) #-}
    (*) = zipBTensorElems sing (*)
    {-# INLINE (*) #-}
    negate = mapBase sing
                     negate
                     (\_          xs -> axpy (-1) xs Nothing)
                     (\(SBM _ sM) xs -> gemm 1 xs (eye sM) Nothing)
    {-# INLINE negate #-}
    abs    = mapBTensorElems abs
    {-# INLINE abs #-}
    signum = mapBTensorElems signum
    {-# INLINE signum #-}
    fromInteger i = genBTensor sing $ \_ -> fromInteger i
    {-# INLINE fromInteger #-}

-- | TODO: add RULES pragmas so that this can be done without checking
-- lengths at runtime in the common case that the lengths are known at
-- compile-time.
--
-- Also, totally forgot about matrix-scalar multiplication here, but there
-- isn't really any way of making it work without a lot of empty cases.
-- should probably handle one level up.
dispatchBLAS
    :: forall b ms os ns v. (RealFloat (ElemB b), BLAS b)
    => MaxLength N1 ms
    -> MaxLength N1 os
    -> MaxLength N1 ns
    -> BTensor v b (ms         ++ os)
    -> BTensor v b (Reverse os ++ ns)
    -> BTensor v b (ms         ++ ns)
dispatchBLAS lM lO lN v r = case (lM, lO, lN) of
    (MLZ    , MLZ    , MLZ    ) -> case (v, r) of
      -- scalar-scalar
      (BTS x, BTS y) -> BTS $ x * y
    (MLZ    , MLZ    , MLS MLZ) -> case (v, r) of
      -- scalar-vector
      (BTS x, BTV y) -> BTV $ axpy x y Nothing
    (MLZ    , MLS MLZ, MLZ    ) -> case (v, r) of
      -- dot
      (BTV x, BTV y) -> BTS $ x `dot` y
    (MLZ    , MLS MLZ, MLS MLZ) -> case (v, r) of
      -- vector-matrix
      -- TODO: transpose?
      (BTV x, BTM y) -> BTV $ gemv 1 (transpB y) x Nothing
    (MLS MLZ, MLZ    , MLZ    ) -> case (v, r) of
      -- vector-scalar
      (BTV x, BTS y) -> BTV $ axpy y x Nothing
    (MLS MLZ, MLZ    , MLS MLZ) -> case (v, r) of
      -- vector-scalar
      (BTV x, BTV y) -> BTM $ ger x y
    (MLS MLZ, MLS MLZ, MLZ    ) -> case (v, r) of
      -- matrx-vector
      (BTM x, BTV y) -> BTV $ gemv 1 x y Nothing
    (MLS MLZ, MLS MLZ, MLS MLZ) -> case (v, r) of
      -- matrix-matrix
      (BTM x, BTM y) -> BTM $ gemm 1 x y Nothing
{-# INLINE dispatchBLAS #-}

mapRowsBTensor
    :: forall k (v :: k -> Type -> Type) ns ms os b. (Vec v, BLAS b)
    => Sing ns
    -> Length os
    -> (BTensor v b ms -> BTensor v b os)
    -> BTensor v b (ns ++ ms)
    -> BTensor v b (ns ++ os)
mapRowsBTensor sN lO f = getI . bRows sN lO (I . f)
{-# INLINE mapRowsBTensor #-}

bRows
    :: forall k (v :: k -> Type -> Type) ns ms os b f. (Applicative f, Vec v, BLAS b)
    => Sing ns
    -> Length os
    -> (BTensor v b ms -> f (BTensor v b os))
    -> BTensor v b (ns ++ ms)
    -> f (BTensor v b (ns ++ os))
bRows sN lO f = bIxRows sN lO (\_ -> f)
{-# INLINE bRows #-}

mapIxRows
    :: forall k (v :: k -> Type -> Type) ns ms os b. (Vec v, BLAS b)
    => Sing ns
    -> Length os
    -> (Prod (IndexN k) ns -> BTensor v b ms -> BTensor v b os)
    -> BTensor v b (ns ++ ms)
    -> BTensor v b (ns ++ os)
mapIxRows sN lO f = getI . bIxRows sN lO (\i -> I . f i)
{-# INLINE mapIxRows #-}

foldMapIxRows
    :: forall k (v :: k -> Type -> Type) ns ms m b. (Vec v, Monoid m, BLAS b)
    => Sing ns
    -> (Prod (IndexN k) ns -> BTensor v b ms -> m)
    -> BTensor v b (ns ++ ms)
    -> m
foldMapIxRows s f = getConst . bIxRows s LZ (\i -> Const . f i)
{-# INLINE foldMapIxRows #-}

bIxRows
    :: forall k (v :: k -> Type -> Type) ns ms os b f. (Applicative f, Vec v, BLAS b)
    => Sing ns
    -> Length os
    -> (Prod (IndexN k) ns -> BTensor v b ms -> f (BTensor v b os))
    -> BTensor v b (ns ++ ms)
    -> f (BTensor v b (ns ++ os))
bIxRows = \case
    SNil   -> \_  f -> f Ø
    s `SCons` ss -> \lO f -> \case
      BTV xs -> case ss of
        -- ns ~ '[n]
        -- ms ~ '[]
        SNil -> case lO of
          -- ns ++ os ~ '[n]
          LZ        -> BTV <$> iElemsB (\i -> fmap unScalar . f (pbvProd i) . BTS) xs
          -- ns ++ os ~ '[n,m]
          LS LZ     -> BTM <$> bgenRowsA (\i -> unVector <$> f (i :< Ø) (BTS $ indexB (PBV i) xs))
                         \\ s
          LS (LS _) -> BTN <$> vGenA s (\i -> f (i :< Ø) (BTS $ indexB (PBV i) xs))
      BTM xs -> case ss of
        -- ns ~ '[n]
        -- ms ~ '[m]
        SNil -> case lO of
          -- ns ++ os ~ '[n]
          LZ        -> BTV <$> bgenA (SBV s) (\(PBV i) -> unScalar <$> f (i :< Ø) (BTV (indexRowB i xs)))
          -- ns ++ os ~ '[n,o]
          LS LZ     -> BTM <$> iRowsB (\i -> fmap unVector . f (i :< Ø) . BTV) xs
          LS (LS _) -> BTN <$> vGenA s (\i -> f (i :< Ø) (BTV (indexRowB i xs)))
        -- ns ~ '[n,m]
        -- ms ~ '[]
        s' `SCons` ss' -> (\\ s') $ case ss' of
          SNil -> case lO of
            LZ   -> BTM <$> iElemsB (\i -> fmap unScalar . f (pbmProd i) . BTS) xs
            LS _ -> BTN <$>
                      vGenA s (\i ->
                          btn lO <$>
                            vGenA s' (\j ->
                                f (i :< j :< Ø) (BTS (indexB (PBM i j) xs))
                              )
                        )
      BTN xs -> (\\ s) $
          fmap (btn (singLength ss `TCL.append'` lO))
        . vITraverse (\i -> bIxRows ss lO (\is -> f (i :< is)))
        $ xs

indexRowBTensor
    :: forall k (b :: BShape k -> Type) v ns ms.
     ( BLAS b
     , Vec v
     )
    => Prod (IndexN k) ns
    -> BTensor v b (ns ++ ms)
    -> BTensor v b ms
indexRowBTensor = \case
    Ø       -> id
    i :< is -> \case
      BTV xs -> case is of
        Ø      -> BTS $ indexB    (PBV i)   xs
      BTM xs -> case is of
        Ø      -> BTV $ indexRowB i         xs
        j :< Ø -> BTS $ indexB    (PBM i j) xs
      BTN xs -> indexRowBTensor is (vIndex i xs)
{-# INLINE indexRowBTensor #-}

mapBTensorElems
    :: (Vec v, BLAS b)
    => (ElemB b -> ElemB b)
    -> BTensor v b ns
    -> BTensor v b ns
mapBTensorElems f = getI . bTensorElems (I . f)
{-# INLINE mapBTensorElems #-}

bTensorElems
    :: forall k (v :: k -> Type -> Type) ns b f. (Applicative f, Vec v, BLAS b)
    => (ElemB b -> f (ElemB b))
    -> BTensor v b ns
    -> f (BTensor v b ns)
bTensorElems f = \case
    BTS x  -> BTS <$> f x
    BTV xs -> BTV <$> elemsB f xs
    BTM xs -> BTM <$> elemsB f xs
    BTN xs -> BTN <$> vITraverse (\_ x -> bTensorElems f x) xs
{-# INLINE bTensorElems #-}

ifoldMapBTensor
    :: forall k (v :: k -> Type -> Type) ns m b. (Monoid m, Vec v, BLAS b)
    => (Prod (IndexN k) ns -> ElemB b -> m)
    -> BTensor v b ns
    -> m
ifoldMapBTensor f = getConst . bTensorIxElems (\i -> Const . f i)
{-# INLINE ifoldMapBTensor #-}

bTensorIxElems
    :: forall k (v :: k -> Type -> Type) ns b f. (Applicative f, Vec v, BLAS b)
    => (Prod (IndexN k) ns -> ElemB b -> f (ElemB b))
    -> BTensor v b ns
    -> f (BTensor v b ns)
bTensorIxElems f = \case
    BTS x  -> BTS <$> f Ø x
    BTV xs -> BTV <$> iElemsB (f . pbvProd) xs
    BTM xs -> BTM <$> iElemsB (f . pbmProd) xs
    BTN xs -> BTN <$> vITraverse (\i -> bTensorIxElems (\is -> f (i :< is))) xs
{-# INLINE bTensorIxElems #-}

zipBTensorElems
    :: forall v b ns. (BLAS b, Nesting1 Sing Applicative v)
    => Sing ns
    -> (ElemB b -> ElemB b -> ElemB b)
    -> BTensor v b ns
    -> BTensor v b ns
    -> BTensor v b ns
zipBTensorElems = \case
    SNil -> \f -> \case
      BTS x -> \case
        BTS y -> BTS (f x y)
    sN `SCons` SNil -> \f -> \case
      BTV xs -> \case
        BTV ys -> BTV (zipB (SBV sN) f xs ys)
    sN `SCons` (sM `SCons` SNil) -> \f -> \case
      BTM xs -> \case
        BTM ys -> BTM (zipB (SBM sN sM) f xs ys)
    (s :: Sing k) `SCons` ss@(_ `SCons` (_ `SCons` _)) -> \f -> \case
      BTN xs -> \case
        BTN ys -> BTN (zipBTensorElems ss f <$> xs <*> ys)
                    \\ (nesting1 s :: Wit (Applicative (v k)))
{-# INLINE zipBTensorElems #-}

liftBTensor
    :: forall v b ns n.
     ( BLAS b
     , Nesting1 Proxy Functor      v
     , Nesting1 Sing  Distributive v
     )
    => Sing ns
    -> (TCV.Vec n (ElemB b) -> ElemB b)
    -> TCV.Vec n (BTensor v b ns)
    -> BTensor v b ns
liftBTensor = \case
    SNil                         -> \f xs ->
        let xs' = unScalar <$> xs
        in  BTS $ f xs'
    sN `SCons` SNil              -> \f xs ->
        let xs' = unVector <$> xs
        in  BTV $ liftB (SBV sN) f xs'
    sN `SCons` (sM `SCons` SNil) -> \f xs ->
        let xs' = unMatrix <$> xs
        in  BTM $ liftB (SBM sN sM) f xs'
    (s :: Sing k) `SCons` ss@(_ `SCons` (_ `SCons` _)) -> \f xs ->
        let xs' = unNested <$> xs
        in  BTN $ TCV.liftVecD (liftBTensor ss f) xs'
              \\ (nesting1 s     :: Wit (Distributive (v k)))
{-# INLINE liftBTensor #-}

mapBTM
    :: forall k (v :: k -> Type -> Type) ns n m ms b. (Vec v, BLAS b)
    => Sing ns
    -> Length ms
    -> (b ('BM n m) -> BTensor v b ms)
    -> BTensor v b (ns ++ [n,m])
    -> BTensor v b (ns ++ ms)
mapBTM sN lM f = getI . traverseBTM sN lM (I . f)
{-# INLINE mapBTM #-}

foldMapBTM
    :: (Monoid a, Vec v, BLAS b)
    => Length ns
    -> (b ('BM n m) -> a)
    -> BTensor v b (ns ++ [n,m])
    -> a
foldMapBTM l f = ifoldMapBTM l (\_ -> f)
{-# INLINE foldMapBTM #-}

traverseBTM
    :: forall k (v :: k -> Type -> Type) ns n m ms b f. (Applicative f, Vec v, BLAS b)
    => Sing ns
    -> Length ms
    -> (b ('BM n m) -> f (BTensor v b ms))
    -> BTensor v b (ns ++ [n,m])
    -> f (BTensor v b (ns ++ ms))
traverseBTM = \case
    SNil         -> \_ f -> \case
      BTM x  -> f x
    s `SCons` ss -> \lM f -> \case
      BTV _  -> case ss of
      BTM _  -> case ss of
      BTN xs -> (\\ s) $
          fmap (btn (singLength ss `TCL.append'` lM))
        . vITraverse (\_ -> traverseBTM ss lM f)
        $ xs
{-# INLINE traverseBTM #-}

imapBTM
    :: forall k (v :: k -> Type -> Type) ns n m ms b. (Vec v, BLAS b)
    => Sing ns
    -> Length ms
    -> (Prod (IndexN k) ns -> b ('BM n m) -> BTensor v b ms)
    -> BTensor v b (ns ++ [n,m])
    -> BTensor v b (ns ++ ms)
imapBTM sN lM f = getI . itraverseBTM sN lM (\i -> I . f i)
{-# INLINE imapBTM #-}

ifoldMapBTM
    :: (Vec v, Monoid a, BLAS b)
    => Length ns
    -> (Prod (IndexN k) ns -> b ('BM n m) -> a)
    -> BTensor v b (ns ++ [n,m])
    -> a
ifoldMapBTM = \case
    LZ -> \f -> \case
      BTM xs -> f Ø xs
    LS l -> \f -> \case
      BTV _  -> case l of
      BTM _  -> case l of
      BTN xs -> vIFoldMap (\i -> ifoldMapBTM l (\is -> f (i :< is))) xs
{-# INLINE ifoldMapBTM #-}

itraverseBTM
    :: forall k (v :: k -> Type -> Type) ns n m ms b f. (Applicative f, Vec v, BLAS b)
    => Sing ns
    -> Length ms
    -> (Prod (IndexN k) ns -> b ('BM n m) -> f (BTensor v b ms))
    -> BTensor v b (ns ++ [n,m])
    -> f (BTensor v b (ns ++ ms))
itraverseBTM = \case
    SNil         -> \_ f -> \case
      BTM x  -> f Ø x
    s `SCons` ss -> \lM f -> \case
      BTV _  -> case ss of
      BTM _  -> case ss of
      BTN xs -> (\\ s) $
          fmap (btn (singLength ss `TCL.append'` lM))
        . vITraverse (\i -> itraverseBTM ss lM (\is ys -> f (i :< is) ys))
        $ xs
{-# INLINE itraverseBTM #-}

mapBase
    :: forall v b ns. (Nesting1 Proxy Functor v)
    => Sing ns
    -> (ElemB b -> ElemB b)
    -> (forall n. Sing n -> b ('BV n) -> b ('BV n))
    -> (forall n m. Sing ('BM n m) -> b ('BM n m) -> b ('BM n m))
    -> BTensor v b ns
    -> BTensor v b ns
mapBase = \case
    SNil -> \f _ _ -> \case
      BTS x  -> BTS (f x)
    sN `SCons` SNil -> \_ g _ -> \case
      BTV xs -> BTV (g sN          xs)
    sN `SCons` (sM `SCons` SNil) -> \_ _ h -> \case
      BTM xs -> BTM (h (SBM sN sM) xs)
    (_ :: Sing k) `SCons` ss@(_ `SCons` (_ `SCons` _)) -> \f g h -> \case
      BTN xs -> BTN (mapBase ss f g h <$> xs)
                  \\ (nesting1 Proxy :: Wit (Functor (v k)))
{-# INLINE mapBase #-}

zipBase
    :: forall v b ns. (Nesting1 Sing Applicative v)
    => Sing ns
    -> (ElemB b -> ElemB b -> ElemB b)
    -> (forall n. Sing n -> b ('BV n) -> b ('BV n) -> b ('BV n))
    -> (forall n m. Sing ('BM n m) -> b ('BM n m) -> b ('BM n m) -> b ('BM n m))
    -> BTensor v b ns
    -> BTensor v b ns
    -> BTensor v b ns
zipBase = \case
    SNil -> \f _ _ -> \case
      BTS x -> \case
        BTS y -> BTS (f x y)
    sN `SCons` SNil -> \_ g _ -> \case
      BTV xs -> \case
        BTV ys -> BTV (g sN          xs ys)
    sN `SCons` (sM `SCons` SNil) -> \_ _ h -> \case
      BTM xs -> \case
        BTM ys -> BTM (h (SBM sN sM) xs ys)
    (s :: Sing k) `SCons` ss@(_ `SCons` (_ `SCons` _)) -> \f g h -> \case
      BTN xs -> \case
        BTN ys -> BTN $ zipBase ss f g h <$> xs <*> ys
                    \\ (nesting1 s :: Wit (Applicative (v k)))
{-# INLINE zipBase #-}

genBTensorA
    :: forall k (b :: BShape k -> Type) v (ns :: [k]) f. (Applicative f, BLAS b, Vec v)
    => Sing ns
    -> (Prod (IndexN k) ns -> f (ElemB b))
    -> f (BTensor v b ns)
genBTensorA = \case
    SNil                                   -> \f ->
        BTS <$> f Ø
    sN `SCons` SNil                        -> \f ->
        BTV <$> bgenA (SBV sN)    (f . pbvProd)
    sN `SCons` (sM `SCons` SNil)           -> \f ->
        BTM <$> bgenA (SBM sN sM) (f . pbmProd)
    s `SCons` ss@(_ `SCons` (_ `SCons` _)) -> \f ->
        BTN <$> vGenA s (\i -> genBTensorA ss (\is -> f (i :< is)))
{-# INLINE genBTensorA #-}

genBTensor
    :: forall k (b :: BShape k -> Type) v (ns :: [k]). (BLAS b, Vec v)
    => Sing ns
    -> (Prod (IndexN k) ns -> ElemB b)
    -> BTensor v b ns
genBTensor s f = getI $ genBTensorA s (I . f)
{-# INLINE genBTensor #-}

indexBTensor
    :: forall k (b :: BShape k -> Type) v ns. (BLAS b, Vec v)
    => Prod (IndexN k) ns
    -> BTensor v b ns
    -> ElemB b
indexBTensor = \case
    Ø      -> \case
      BTS x  -> x
    i :< Ø -> \case
      BTV xs -> indexB (PBV i) xs
    i :< j :< Ø -> \case
      BTM xs -> indexB (PBM i j) xs
    i :< js@(_ :< _ :< _) -> \case
      BTN xs -> indexBTensor js (vIndex i xs)
{-# INLINE indexBTensor #-}

btn :: (BLAS b, Vec v, SingI n)
    => Length ns
    -> v n (BTensor v b ns)
    -> BTensor v b (n ': ns)
btn = \case
    LZ        -> \xs ->
      BTV $ bgen sing (unScalar . (`vIndex` xs) . unPBV)
    LS LZ     -> \xs ->
      BTM $ bgenRows  (unVector . (`vIndex` xs))
    LS (LS _) -> BTN
{-# INLINE btn #-}

gmul'
    :: forall v b ms os ns.
     ( SingI (ms ++ ns)
     , RealFloat (ElemB b)
     , Vec v
     , Nesting1 Proxy Functor     v
     , Nesting1 Sing  Applicative v
     , BLAS b
     )
    => Length ms
    -> Length os
    -> Length ns
    -> BTensor v b (ms         ++ os)
    -> BTensor v b (Reverse os ++ ns)
    -> BTensor v b (ms         ++ ns)
gmul' lM lO lN = gmulB sM lO lN \\ sN
  where
    sM :: Sing ms
    sN :: Sing ns
    (sM, sN) = splitSing lM sing
{-# INLINE[0] gmul' #-}

{-# RULES
"gmul'/SS"  gmul' = dispatchSS
"gmul'/SV"  gmul' = dispatchSV
"gmul'/dot" gmul' = dispatchDot
"gmul'/VM"  gmul' = dispatchVM
"gmul'/VS"  gmul' = dispatchVS
"gmul'/out" gmul' = dispatchOut
"gmul'/MV"  gmul' = dispatchMV
"gmul'/MM"  gmul' = dispatchMM
  #-}

-- | General strategy:
--
-- *   We can only outsource to BLAS (using 'dispatchBLAS') in the case
--     that @os@ and @ns@ have length 0 or 1.  Anything else, fall back to
--     the basic reverse-indexing method in "Data.Nested".
-- *   If @ms@ is length 2 or higher, "traverse down" to the length 0 or
--     1 tail...and then sum them up.
gmulB
    :: forall k (b :: BShape k -> Type) v ms os ns.
     ( RealFloat (ElemB b)
     , SingI ns
     , BLAS b
     , Vec v
     , Nesting1 Proxy Functor     v
     , Nesting1 Sing  Applicative v
     )
    => Sing ms
    -> Length os
    -> Length ns
    -> BTensor v b (ms         ++ os)
    -> BTensor v b (Reverse os ++ ns)
    -> BTensor v b (ms         ++ ns)
gmulB sM lO lN v r = case splitting (S_ Z_) (lengthProd lN) of
    Fewer mlN _ -> case splittingEnd (S_ (S_ Z_)) (lengthProd lO) of
      FewerEnd MLZ             _ -> gmulBLAS sM MLZ       mlN v r
      FewerEnd (MLS MLZ)       _ -> gmulBLAS sM (MLS MLZ) mlN v r
      FewerEnd (MLS (MLS MLZ)) _ -> case mlN of
        MLZ -> case r of
          BTM ys -> mapBTM sM LZ (\xs -> BTS $ traceB (gemm 1 xs ys Nothing)) v
        MLS MLZ -> naiveGMul sM lO lN v r
      SplitEnd _ _ _ -> naiveGMul sM lO lN v r
    Split _ _ _ -> naiveGMul sM lO lN v r
{-# INLINE[0] gmulB #-}

-- | Naive implementation of 'gmul' (based on the implementation for
-- 'NTensor') that does not utilize any BLAS capabilities.
naiveGMul
    :: forall k (b :: BShape k -> Type) v ms os ns.
     ( BLAS b
     , Vec v
     , Num (ElemB b)
     , Nesting1 Proxy Functor     v
     , Nesting1 Sing  Applicative v
     , SingI ns
     )
    => Sing ms
    -> Length os
    -> Length ns
    -> BTensor v b (ms         ++ os)
    -> BTensor v b (Reverse os ++ ns)
    -> BTensor v b (ms         ++ ns)
naiveGMul sM _ lN v r =
    mapRowsBTensor sM lN (getSum . ifoldMapBTensor (\i -> Sum . f i)) v
  where
    f  :: Prod (IndexN k) os
       -> ElemB b
       -> BTensor v b ns
    f is x = mapBase sing
                     (x *)
                     (\_ ys -> scaleB x ys)
                     (\_ ys -> scaleB x ys)
                     (indexRowBTensor (TCP.reverse' is) r)

-- | A 'gmul' that runs my dispatching BLAS commands when it can.
-- Contains the type-level constraint that @os@ and @ns@ have to have
-- either length 0 or 1.
--
-- TODO: no longer needs Sing ms
gmulBLAS
    :: forall b ms os ns v.
     ( RealFloat (ElemB b)
     , BLAS b
     , Vec v
     , SingI ns
     , Nesting1 Proxy Functor     v
     , Nesting1 Sing  Applicative v
     )
    => Sing ms
    -> MaxLength N1 os
    -> MaxLength N1 ns
    -> BTensor v b (ms         ++ os)
    -> BTensor v b (Reverse os ++ ns)
    -> BTensor v b (ms         ++ ns)
gmulBLAS sM mlO mlN v r = case mlO of
    MLZ -> case splittingEnd (S_ (S_ Z_)) spM of
      FewerEnd MLZ             _ -> dispatchBLAS MLZ       mlO  mlN v r
      FewerEnd (MLS MLZ)       _ -> dispatchBLAS (MLS MLZ) mlO mlN v r
      FewerEnd (MLS (MLS MLZ)) _ -> case v of
        BTM xs -> case mlN of
          MLZ     -> case r of
            BTS y  -> BTM $ scaleB y xs
          -- TODO: can this be made non-naive?
          -- ms ~ '[m1,m2]
          -- os ~ '[]
          -- ns ~ '[n]
          MLS MLZ -> naiveGMul sM LZ (fromMaxLength mlN) v r
      SplitEnd (ELS (ELS ELZ)) spM0 spM1 -> case mlN of
        MLZ -> case r of
          BTS y -> mapBTM (prodSing   spM0)
                          (prodLength spM1)
                          (\xs -> BTM $ scaleB y xs)
                          v
                     \\ appendNil lM
        -- TODO: can this be made non-naive?
        -- ms ~ (ms0 ++ '[m1,m2])
        -- os ~ '[]
        -- ns ~ '[n]
        MLS MLZ -> naiveGMul sM LZ (fromMaxLength mlN) v r
    MLS MLZ -> case splittingEnd (S_ Z_) spM of
      FewerEnd mlM       _         -> dispatchBLAS mlM mlO mlN v r
      SplitEnd (ELS ELZ) spM0 spM1 ->
        let sM0 = prodSing spM0
            lM0 = prodLength spM0
            lM1 = prodLength spM1
        in  (\\ appendAssoc (TCL.tail' lM0)
                            lM1
                            (LS LZ :: Length os)
            ) $ case mlN of
          MLZ -> case r of
            BTV ys -> mapBTM sM0 lM1 (\xs -> BTV $ gemv 1 xs ys Nothing) v
                        \\ appendNil lM
          MLS MLZ -> case r of
            BTM ys -> mapBTM sM0
                            (lM1 `TCL.append'` (LS LZ :: Length ns))
                            (\xs -> BTM $ gemm 1 xs ys Nothing)
                            v
                        \\ appendAssoc (TCL.tail' lM0)
                                       lM1
                                       (LS LZ :: Length ns)
  where
    spM = singProd sM
    lM  = singLength sM

diagBTensor
    :: forall k (b :: BShape k -> Type) v n ns.
     ( SingI (n ': ns)
     , BLAS b
     , Vec v
     , Num (ElemB b)
     , Eq (IndexN k n)
     )
    => Uniform n ns
    -> BTensor v b '[n]
    -> BTensor v b (n ': ns)
diagBTensor = \case
    UØ    -> id
    US UØ -> \case
      BTV xs -> BTM $ diagB xs
    u@(US (US _)) -> \(BTV xs) ->
      genBTensor sing (\i -> case TCV.uniformVec (prodToVec I (US u) i) of
                               Nothing -> 0
                               Just (I i') -> indexB (PBV i') xs
                      )
{-# INLINE diagBTensor #-}

transpBTensor
    :: (BLAS b, Vec v)
    => Sing ns
    -> BTensor v b ns
    -> BTensor v b (Reverse ns)
transpBTensor s = \case
    BTS x      -> BTS x
    BTV xs     -> BTV xs
    BTM xs     -> BTM $ transpB xs
    xs@(BTN _) -> (\\ reverseReverse (singLength s)) $
                    genBTensor (sReverse s) $ \i ->
                      indexBTensor (TCP.reverse' i) xs
{-# INLINE transpBTensor #-}

sumBTensor
    :: forall v b n ns.
     ( BLAS b
     , Vec v
     , Num (ElemB b)
     , Foldable (v n)
     , SingI ns
     , SingI n
     , Nesting1 Proxy Functor     v
     , Nesting1 Sing  Applicative v
     )
    => BTensor v b (n ': ns)
    -> BTensor v b ns
sumBTensor = \case
    BTV xs  -> BTS $ sumB xs
    BTM (xs :: b ('BM n m))
            -> BTV $ gemv 1 (transpB xs)
                          (bgen (SBV (sing :: Sing n)) (\_ -> 1))
                          Nothing
    BTN xs  -> sum xs

instance
      ( Vec (v :: k -> Type -> Type)
      , BLAS b
      , RealFloat (ElemB b)
      , Nesting1 Proxy Functor      v
      , Nesting1 Proxy Foldable     v
      , Nesting1 Sing  Applicative  v
      , Nesting1 Sing  Distributive v
      , Eq1 (IndexN k)
      )
  => Tensor (BTensor v b) where
    type ElemT (BTensor v b) = ElemB b

    liftT
        :: SingI ns
        => (TCV.Vec n (ElemB b) -> ElemB b)
        -> TCV.Vec n (BTensor v b ns)
        -> BTensor v b ns
    liftT = liftBTensor sing
    {-# INLINE liftT #-}

    sumT = sum'
    {-# INLINE sumT #-}

    scaleT α = mapBase sing (α*) (\_ -> scaleB α) (\_ -> scaleB α)
    {-# INLINE scaleT #-}

    gmul
        :: forall ms os ns. SingI (ms ++ ns)
        => Length ms
        -> Length os
        -> Length ns
        -> BTensor v b (ms         ++ os)
        -> BTensor v b (Reverse os ++ ns)
        -> BTensor v b (ms         ++ ns)
    gmul = gmul'
    {-# INLINE gmul #-}

    diag
        :: forall n ns. SingI (n ': ns)
        => Uniform n ns
        -> BTensor v b '[n]
        -> BTensor v b (n ': ns)
    diag = diagBTensor
             \\ (produceEq1 :: Eq1 (IndexN k) :- Eq (IndexN k n))
    {-# INLINE diag #-}

    getDiag
        :: SingI n
        => Uniform n ns
        -> BTensor v b (n ': n ': ns)
        -> BTensor v b '[n]
    getDiag = \case
      UØ   -> \case
        BTM xs -> BTV $ getDiagB xs
      u@(US _) -> \xs ->
        genBTensor sing $ \(i :< Ø) ->
          indexBTensor (TCP.replicate i (US (US u))) xs
    {-# INLINE getDiag #-}

    transp = transpBTensor sing
    {-# INLINE transp #-}

    generateA = genBTensorA sing
    {-# INLINE generateA #-}

    genRand d g = generateA (\_ -> realToFrac <$> genContVar d g)
    {-# INLINE genRand #-}

    ixRows
        :: forall f ms os ns. (Applicative f, SingI (ms ++ os))
        => Length ms
        -> Length os
        -> (Prod (IndexN k) ms -> BTensor v b ns -> f (BTensor v b os))
        -> BTensor v b (ms ++ ns)
        -> f (BTensor v b (ms ++ os))
    ixRows lM lO = bIxRows sM lO
      where
        sM :: Sing ms
        sM = takeSing lM lO (sing :: Sing (ms ++ os))
    {-# INLINE ixRows #-}

    (!) = flip indexBTensor
    {-# INLINE (!) #-}

    sumRows
        :: forall n ns. (SingI (n ': ns), SingI ns)
        => BTensor v b (n ': ns)
        -> BTensor v b ns
    sumRows = sumBTensor
                \\ (nesting1 Proxy :: Wit (Foldable (v n)))
                \\ sHead (sing :: Sing (n ': ns))
    {-# INLINE sumRows #-}

    mapRows :: forall ns ms. SingI (ns ++ ms)
            => Length ns
            -> (BTensor v b ms -> BTensor v b ms)
            -> BTensor v b (ns ++ ms)
            -> BTensor v b (ns ++ ms)
    mapRows l f = mapRowsBTensor sN (singLength sM) f
      where
        sN :: Sing ns
        sM :: Sing ms
        (sN, sM) = splitSing l (sing :: Sing (ns ++ ms))
    {-# INLINE mapRows #-}

-- * Boring dispatches

dispatchSS
    :: Num (ElemB b)
    => Length '[]
    -> Length '[]
    -> Length '[]
    -> BTensor v b '[]
    -> BTensor v b '[]
    -> BTensor v b '[]
dispatchSS _ _ _ (BTS x) (BTS y) = BTS (x * y)
{-# INLINE dispatchSS #-}

dispatchSV
    :: BLAS b
    => Length '[]
    -> Length '[]
    -> Length '[n]
    -> BTensor v b '[]
    -> BTensor v b '[n]
    -> BTensor v b '[n]
dispatchSV _ _ _ (BTS x) (BTV y) = BTV $ axpy x y Nothing
{-# INLINE dispatchSV #-}

dispatchDot
    :: BLAS b
    => Length '[]
    -> Length '[n]
    -> Length '[]
    -> BTensor v b '[n]
    -> BTensor v b '[n]
    -> BTensor v b '[]
dispatchDot _ _ _ (BTV x) (BTV y) = BTS $ x `dot` y
{-# INLINE dispatchDot #-}

dispatchVM
    :: (Num (ElemB b), BLAS b)
    => Length '[]
    -> Length '[n]
    -> Length '[m]
    -> BTensor v b '[n]
    -> BTensor v b '[n,m]
    -> BTensor v b '[m]
dispatchVM _ _ _ (BTV x) (BTM y) = BTV $ gemv 1 (transpB y) x Nothing
{-# INLINE dispatchVM #-}

dispatchVS
    :: BLAS b
    => Length '[n]
    -> Length '[]
    -> Length '[]
    -> BTensor v b '[n]
    -> BTensor v b '[]
    -> BTensor v b '[n]
dispatchVS _ _ _ (BTV x) (BTS y) = BTV $ axpy y x Nothing
{-# INLINE dispatchVS #-}

dispatchOut
    :: BLAS b
    => Length '[n]
    -> Length '[]
    -> Length '[m]
    -> BTensor v b '[n]
    -> BTensor v b '[m]
    -> BTensor v b '[n,m]
dispatchOut _ _ _ (BTV x) (BTV y) = BTM $ ger x y
{-# INLINE dispatchOut #-}

dispatchMV
    :: (Num (ElemB b), BLAS b)
    => Length '[n]
    -> Length '[m]
    -> Length '[]
    -> BTensor v b '[n,m]
    -> BTensor v b '[m]
    -> BTensor v b '[n]
dispatchMV _ _ _ (BTM x) (BTV y) = BTV $ gemv 1 x y Nothing
{-# INLINE dispatchMV #-}

dispatchMM
    :: (Num (ElemB b), BLAS b)
    => Length '[m]
    -> Length '[o]
    -> Length '[n]
    -> BTensor v b '[m,o]
    -> BTensor v b '[o,n]
    -> BTensor v b '[m,n]
dispatchMM _ _ _ (BTM x) (BTM y) = BTM $ gemm 1 x y Nothing
{-# INLINE dispatchMM #-}


{-# LANGUAGE ScopedTypeVariables #-}

module Clustering.Hierarchical where

import Numeric.LinearAlgebra
import Data.List.Lens
import Control.Lens
import Control.Applicative
import Debug.Trace ( trace )

import Clustering.Util
import Clustering.Data

findClosest2 :: (Ord t) => (a -> a -> t) -> [a] -> Maybe ((Int, Int), t)
findClosest2 _ []  = Nothing
findClosest2 f cls = getidx r1 r2
    where
        r1 = findClosest1s f cls
        r2 = argminval (snd <$>) r1
        getidx :: [Maybe (Int, t)] -> Maybe (Int, Maybe t) -> Maybe ((Int, Int), t)
        getidx l (Just (i, Just v1)) = case l !! i of
          Nothing -> Nothing
          Just (j, v2) -> Just ((i, j), v1)
        getidx l (Just (i, Nothing)) = Nothing
        getidx _ Nothing = Nothing

findClosest1s :: (Ord t) => (a -> a -> t) -> [a] -> [Maybe (Int, t)]
findClosest1s _ [] = []
findClosest1s f ls = reverse $ findClosest1s' [] ls []
    where
        findClosest1s' left [] cl = cl
        findClosest1s' left (x:right) cl =
            findClosest1s' (left ++ [x]) right $
                        selectMMin
                            (findClosest f x left)
                            ((_1 +~ 1 + length left) <$> findClosest f x right)
                            -- shift right index
                        : cl

findClosest :: (Ord t) =>  (a -> a -> t) -> a -> [a] -> Maybe (Int, t)
findClosest _ _ [] = Nothing
findClosest d x ls = argminval (d x) ls

selectMMin :: Ord b => Maybe (a, b) -> Maybe (a, b) -> Maybe (a, b)
selectMMin (Just (n, v1)) (Just (m, v2)) | v2 < v1 = Just (m, v2)
selectMMin (Just (n, v1)) (Just (m, v2)) | v2 > v1 = Just (n, v1)
selectMMin m1 m2 = m1 <|> m2

merge ::(Ord t) => t -> ([a] -> [a] -> t) -> [[a]] -> [[a]]
merge _ _ [] = []
merge m f ls = case findClosest2 f ls of
  Nothing          -> ls
  Just ((i, j), v) -> if v < m
                        then merge m f $ merge2 i j ls
                        else ls

merge2 :: Int -> Int -> [[a]] -> [[a]]
merge2 i j ls = deleteList j $ ix i %~ (\l -> l ++ ls !! j) $ ls

-- | Generalized aggregative scheme using Dendrogram
gasd :: forall a.
    (Dendrogram a -> Dendrogram a -> Float) -- ^ dendrogram distance
    -> (Dendrogram a -> Dendrogram a -> a)  -- ^ representative update
    -> [a] -> Dendrogram a
gasd dd ru xs = agglomerated' 0 $ map dLeaf xs
    where
        agglomerated' :: Int -> [Dendrogram a] -> Dendrogram a
        agglomerated' _ [] = DNil
        agglomerated' _ ds@(d:_) | length ds == 1 = d
        agglomerated' level ds = case findClosest2 dd ds of
            Nothing -> error "Nothing to merge!"
            Just ((i, j), v) -> let (newl, newr) = (ds !! i, ds !!j) in
                agglomerated' (level + 1) $ (DNode (ru newl newr) (level + 1) v newl newr) : dels [i, j] ds

-- | Generalized aggregative scheme
gas ::  Ord t => forall a. ([a] -> [a] -> t) -> [a] -> [[[a]]]
gas cd xs = agglomerate' [map (:[]) xs]
    where
        -- agglomerate' :: [[[a]]] -> [[[a]]]
        agglomerate' [] = []
        agglomerate' [[]] = []
        agglomerate' acs@(c:_) | length c == 1 = acs
        agglomerate' acs@(c:_) = case findClosest2 cd c of
            Nothing -> acs
            Just ((i, j), _) -> agglomerate' $ merge2 i j c:acs

buildSimilarityMtrx :: (a -> a -> Float) -> [a] -> Matrix Float
buildSimilarityMtrx d xs = l >< l $ concatMap (\x -> map (d x) xs) xs
    where
        l = length xs

-- | Matrix update aggregative scheme
muas :: Matrix Float -> (Matrix Float -> Int -> Int -> Int) -> [[Int]]
muas p d = undefined -- TODO optimized agglomerative scheme using matrix update

-- | Modified general divisive scheme
mgds :: (Ord t, Num t, Fractional t) =>
    (a -> a -> t) -- ^ dissimiliarity measure
    -> [a]
    -> [[[a]]]
mgds d [] = []
mgds d xs = divideAll [[xs]]
    where
        divideAll [] = []
        divideAll acs@(c:cs) = if length cs == length xs then acs
                                                         else divideAll $ (c >>= splitCluster) : acs
        splitCluster [] = []
        splitCluster xs = reassign $ outsider xs
                      -- map (((1/fromIntegral (length xs) * ) . sum) . (\x -> map (d x) xs)) xs
        outsider xs = case argmax id $ map (avg . flip map xs . d) xs of
                        Nothing -> error "What happend?!"
                        Just n -> case pop n xs of
                             (as, Nothing) -> ([], as)
                             (as, Just a) -> ([a], as)

        reassign ([], os) = [os]
        reassign (ns, []) = [ns]
        reassign (ns, as) = [ns, as]