{-# LANGUAGE RecordWildCards #-}
module Q.Options.Black76
  (
    module Q.Options
  , Black76(..)
  , atmf
  , euOption
  , eucall
  , euput
  )
  where

import           Q.Options
import           Q.Types
import           Statistics.Distribution        (cumulative, density)
import           Statistics.Distribution.Normal (standard)

data Black76 = Black76 {
    b76F   :: Forward
  , b76DF  :: DF
  , b76T   :: YearFrac
  , b76Vol :: Vol
}

-- | At the money forward strike.
atmf :: Black76 -> Strike
atmf Black76{..} = Strike f
  where (Forward f) = b76F

-- | European option valuation with black 76
euOption :: Black76 -> OptionType -> Strike -> Valuation
euOption b76@Black76{..} cp k = Valuation premium delta vega gamma where
  (Forward f) = b76F
  n           = cumulative standard
  (Vol sigmaSqt) = scale b76T b76Vol
  d1          = (dPlus  b76F b76Vol k b76T)
  d2          = (dMinus b76F b76Vol k b76T)
  nd1         = n d1
  nd2         = n d2
  callDelta   = b76DF `discount` nd1
  putDelta    = b76DF `discount` (- (n (-d1)))
  vega        = Vega  $ b76DF `discount` (density standard d1 ) * f * sigmaSqt
  gamma       = Gamma $ b76DF `discount` (density standard d1) / (f * sigmaSqt)
  premium  = Premium $ case cp of
    Call -> b76DF `discount` (f * nd1 - nd2 * k')
    Put  -> b76DF `discount` (n (-d2) * k' - n (-d1) * f)
    where (Strike k') = k
  delta | cp == Call = Delta $ callDelta
        | cp == Put  = Delta $ putDelta

-- | see 'euOption'
euput b76 =  euOption b76 Put

-- | see 'euOption'
eucall b76 = euOption b76 Call

dPlus (Forward f) (Vol sigma) (Strike k) (YearFrac t) =
  recip (sigma * sqrt t) * (log (f/k) + (0.5 * sigma * sigma) * t)
dMinus (Forward f) (Vol sigma) (Strike k) (YearFrac t) =
  recip (sigma * sqrt t) * (log (f/k) - (0.5 * sigma * sigma) * t)


{-# LANGUAGE FlexibleContexts, TypeOperators, RankNTypes, GeneralizedNewtypeDeriving, DeriveFunctor, PolyKinds, GADTs #-}

module Qubit where

import Linear.V2
import Linear.V1
import Linear.Vector
import Vec
import Data.Complex
import Data.Functor.Compose
import Data.Functor.Product
import Data.Functor.Identity
import Cat
import Data.Coerce

type Kron = Compose
type Spin = V2
type QuBit = V2
type C = Complex Double

type DSum = Product

newtype FKron f g a = FKron [(f a, g a)]

type Fock a =  Q [a] -- a "fock space" over a. No assumption of anti symmettric or symmettric.
newtype Anti a = Anti (Fock a)
-- for antisymmettric space, keep list sorted (canonical form).
-- liftAnti1 :: (a -> Q a) -> ([a] -> Q [a])
-- liftAnti2 :: (a,a) -> Q (a,a) -> ([a] -> Q [a])
--

-- this is fock system as a sparse indictaor matrix. We can also do the form of occupation number  Q (V Bool)

vac :: Fock a
vac = pure []

type SHO = Q Int

shovac :: SHO
shovac = pure 0
adag :: Int -> SHO
adag n = W [(n + 1, sqrt (fromInteger (toInteger n) + 1))]
agad 0 = mempty
agad n = W [(n - 1, sqrt n)]

raise = LinOp adag
lower = LinOp agad

-- boson system can be seen as collection of oscillators. Q (V Int).  Positive integers really.

type Boson v = Q (v Int) -- accepts a functor describing all the oscillators. A Vec for example.
-- bvac = pure zero -- if v is a Vector type, it probably has a zero.
--

-- traversable ordering. The vector is put into canocnical form with adag appearing first in the traversal coming first.
type Fermion v = Q (v Bool) -- fermion needs to prefix count to figure out signs. This means that v has to be traversable? Foldable?
fvac :: Applicative v => Fermion v
fvac = pure (pure False)
{-
-- what about systems with different types of particles.
particle-hole  Q ([a], [a])

AntiOp a = AntiOp [a] -> Q [a]
-- indexed linear operators. a -> AntiOp a
psi :: Ord a => a -> (a -> Anti a)
psidag :: Ord a => a -> (a -> Anti a)

phi :: Ord a => a -> (a -> Sym a)
phidag :: Ord a => a -> (a -> Sym a)

-- hmm. Something kind of funny about these.
lift1 :: (i -> Double) -> (i -> Fock i)
lift2 :: (i -> i -> Double) -> (i -> Fock i)
symlift -- mention that there is implcit symmettrization
2 natural lifts. cnnected by logarithms. UxUxUxU vs HIII + IHII + IIHI + IIIH

sum1 :: (i -> i -> Q i) -> (i -> Q i) -- takes an indexed operator and turns it into an operator
sum1 :: (i -> AntiOp i) -> (AntiOp i)
sum2
sum :: Addable b => (i -> b) -> b

turn all of this into a typeclass functions. Then do this form, but also full initial form for manipulation.

dyson

-}
-- instance (Additive f, Additive g) => Additive (Compose f g) where
-- instance Num f (g a) => Num (Compose f g a)

-- (Kron (Kron V2 V2) V2) Complex Double  -- a 3 qubit vector space
swap :: (Traversable f, Applicative g) => Kron f g a -> Kron g f a
swap = Compose . sequenceA . getCompose

-- kron :: (Num a, Functor f, Functor g) => f a -> g a -> Kron f g a
-- kron f g = Compose $ fmap (\amp1 -> fmap (\amp2 -> amp1 * amp2) g) f
kron' :: (Num a, Functor f, Functor g) => f a -> g a -> Kron f g a
kron' x y = Compose (outer x y)

densify ::  (Additive (Kron f g), Num a, Functor f, Functor g) => FKron f g a -> Kron f g a
densify (FKron xs) = sumV $ map (uncurry kron') xs

kron''' :: (Applicative f', Applicative g', Traversable f, Traversable g') =>
    (f a -> f' a) -> (g a -> g' a) -> (Kron f g a -> Kron f' g' a)
kron''' lf lg (Compose fga) = Compose $ sequenceA $ (fmap lf) $ sequenceA  $ (fmap lg fga)

type LinOp' f g a = f a -> g a
-- SVD let's you re free
-- but also [indicatro vectro, g a]
-- and perhaps randomize sampling.
    {-
sigx :: Spin C -> Spin C
sigx (V2 up down) = V2 down up

sigz :: Spin C -> Spin C
sigz (V2 (a :+ b) (c :+ d)) = V2 ()

sigz :: Spin C -> Spin C
sigz (V2 up down) = V2 up (-down)
-}
{-
class Add a where
    (+++) :: a -> a -> a
    negate' :: a -> a

instance Num a => Add a where
    +++ = +
instance    => Additive f
-}

flip' :: Spin a -> Spin a
flip' (V2 up down) = V2 down up

sigx = flip'

sigz :: Num a => Spin a -> Spin a
sigz (V2 up down) = V2 up (negate down)

-- phase :: Complex a => Spin a -> Spin a
type f ~> g = forall a. (Functor f, Functor g) => f a -> g a
type f *** g = Compose f g
-- (Functor f, Functor g, Functor h, Functor k)
-- we only need g to be a functor
-- see below. If we remove the type signature, this implementation generalizes to another type
par :: Functor g => (forall a. f a -> g a) -> (forall b. h b -> k b) -> (forall c. (Compose f h c) -> (Compose g k c))
par f g = Compose . (fmap g) . f . getCompose -- insufficient though. We can't write linear operators this way

-- I should used V1, not Identity
-- These are all also coerce.
rightUnitor :: Functor f => Compose f Identity a -> f a
rightUnitor (Compose f) = fmap runIdentity f

rightUnitor' :: f ~> Compose f Identity
rightUnitor' = Compose . fmap Identity

leftUnitor' :: f ~> Compose Identity f
leftUnitor' = Compose . Identity

leftUnitor :: Identity *** f ~> f
leftUnitor = runIdentity . getCompose

-- should be coerce. Confusing role problems
assoc :: Functor f => (Kron (Kron f g) h) ~> (Kron f (Kron g h))
assoc = Compose . (fmap Compose) . getCompose . getCompose

assoc' :: Functor f =>  (Kron f (Kron g h)) ~> (Kron (Kron f g) h)
assoc' (Compose x)  = Compose $ Compose $ (fmap getCompose x)

-- if the functor in question happens to be representable
    -- then we can use our previous defitions of monoidal catgoeries.
    -- Reader Functors
-- Num c => prod :: Reader a c -> Reader b c -> Reader (a,b) c
-- Num
-- newtype ReadCat c a b = ReadCat (Reader a c -> Reader b c)
-- newtype LinOp c a b = LinOp ((a -> c) -> (b -> c))
-- par (LinOp f) (LinOp g) = LinOp \h -> h (\(a,b) ->   )
-- FMapN
-- fmapN @n
{-Compose f g

assoc :: Compose (Compose f g) h a -> Compose f (Compose g h) a

assoc = Compose . (fmap Compose) . getCompose . getCompose

Compose f Id a -> f a

rightunit = (fmap runIdentity) . getCompose

The morphisms of the functor category are natural transformations

par :: (f ~> g) -> (h~> k) -> ((Compose f h) ~> (Compose g k))

-}

newtype Kron' k l f a = Kron' ((k (l f)) a) deriving Functor

-- deriving via compose
{-
newtype V2' f a = V2' (Compose V2 f a)  deriving Functor
newtype V1' f a = V1' (Compose V1 f a) deriving Functor
-}
{-
newtype V2' f a = V2' (V2 (f a))  deriving Functor
newtype V1' f a = V1' (f a) deriving Functor
-}

type V2' f a = Kron V2 f a
type V1' f a = Kron V1 f a

lowerv1 (Compose (V1 f)) = f -- It's a coerce. V1 is assuredly a newtype
{-
exvsa :: (Additive f, Num a) => V2' f a
exvsa = zero
-}

-- constant vecotrs a morphisms from unit
-- Using f ~ one = (V1 1)
exconst :: (a ~ Double, Functor f) => V1' f a -> V2' f a
exconst x = let x' = lowerv1 x in Compose $ V2 (1.4 *^ x') (7.3 *^ x')

one = V1 1
attemp = exconst (Compose (V1 one))

-- Applicative composes, so Additive probably should too.
-- Wait is this in the library?
-- Why is applicative g necessary?
{-
Ok. There doesn't appear to be a compose
But V2 etc have a Num instance
and liftU2 has a default implementation that uses Applicative
I probably don't need Applcative g if I define a liftU2 using g's liftU2

-}
-- I should check to make sure the lift makes sense
-- Why isn't this in the library? There must be a good reason.
-- Is there an additive instance for Product?

instance (Additive g, Additive f) => Additive (Compose f g) where
    (Compose v) ^+^ (Compose w) = Compose (liftU2 (^+^) v w) -- I  we might be able to used
    zero = zero -- what?
    liftU2 f (Compose x) (Compose y) = Compose $ liftU2 (liftU2 f) x y
    liftI2 f (Compose x) (Compose y) = Compose $ liftI2 (liftI2 f) x y

type LinOp1 f g a = forall k. Additive k => Kron f k a -> Kron g k a

-- This feels less right to me. But why?
type LinOp1' f g a = forall k. Functor k => Kron k f a -> Kron k g a
{-
It feels less right perhaps because it will change how we do things so little.
We just have to fmap.
No, but saying we always have direct access to the scalar isn't right.
-}

-- reassoc :: Kron f (Kron g k) a -> Kron (Kron f g) k a
-- reassoc = coerce

--reassoc' ::  Kron (Kron f g) k a -> Kron f (Kron g k) a
--reassoc' = -- coerce

-- yikes. coerce will help a bit. Quite a bit.
-- Also Fixing the unnecessary applicative instances.
-- Additive k' is expected. Applicative is coming from Additve (Kron )  constraint.

{-kron'' :: (Applicative g, Applicative k', Additive k', Functor f', Functor f, Functor g') =>
  (forall k. Additive k => Kron f k a -> Kron f' k a) -> (forall k. Additive k => Kron g k a -> Kron g' k a) -> (Kron (Kron f g) k') a -> (Kron (Kron f' g') k') a
-}

type Qubit = V2

kron'' :: (Additive f, Additive g, Additive k, Additive f', Additive g') =>
  LinOp1 f f' a -> LinOp1 g g' a -> Kron (Kron f g) k a -> Kron (Kron f' g') k a
kron'' lf lg fgk = let v = (Qubit.assoc fgk) in Qubit.assoc' (Compose $ fmap lg $ getCompose (lf v))

sigx' :: LinOp1 Qubit Qubit C
sigx' (Compose (V2 up down)) = Compose $ V2 down up

sigz' :: LinOp1 Qubit Qubit C
sigz' (Compose (V2 up down)) = Compose  $ V2 up ((-1) *^ down)

sigy' :: LinOp1 Qubit Qubit C
sigy' (Compose (V2 up down)) = Compose $ V2 ((-i) *^ down) (i *^ up) where i = 0 :+ 1

swap' :: (Traversable f, Applicative g) => LinOp1 (Kron f g) (Kron g f) a
swap' (Compose (Compose fgk)) = Compose $ Compose $ sequenceA fgk

cnot :: LinOp1 (Kron Qubit Qubit) (Kron Qubit Qubit) a
cnot (Compose (Compose (V2 (V2 up1 down1) v))) = Compose $ Compose $ V2 (V2 down1 up1) v

phase :: Double -> LinOp1 Qubit Qubit C
phase phi (Compose (V2 up down)) = Compose $ V2 up ((cis phi) *^ down)

mtest :: LinOp1 (Kron Qubit Qubit) (Kron Qubit Qubit) C
mtest = (left sigx') . (left sigy') . (right sigz') . swap'

left :: (Additive f, Additive k, Additive g) => LinOp1 f g a -> LinOp1 (Kron f k) (Kron g k) a
left l = kron'' l id -- Qubit.assoc' . l . Qubit.assoc

right :: (Additive k, Additive f, Additive g) => LinOp1 f g a -> LinOp1 (Kron k f) (Kron k g) a
right l = kron'' id l -- (Compose (Compose fkk)) = Compose $ Compose $ (fmap (getCompose . l . Compose) fkk)

leftUnit :: Kron V1 (Kron f k) a -> Kron f k a
leftUnit (Compose (V1 x)) = x

{-
(f a -> g a) is a natural transfromation
(Fractional a => f a -> g a) is a numerical transfromation

Horizontal composition in Functor land. Sure.
Even in applicative
The issue is that Good linear ops are going to hold numbers. So ti feels wrong to quanitfy over them
One option is to add a type class constraint
(Num a) or (Fractional a)
This will work because linear lifts these

-- this is the same as par written above.

-- Functor f => (t -> f a) -> (a -> b) -> t -> f b
The oppotite ordering
--  Functor f' => (f' b -> t') -> (a -> b) -> f' a -> t'
The more general type as inferred is
f b ~ t'

(f' a -> f a) -> (a -> b) -> (f' a -> f b)
vs
(f' b -> f b) -> (a -> b) -> (f' a -> f b)
vs.
forall

kron l l' x = l' (fmap l x)

(forall f. Additive f => Kron V2 f a -> Kron V3 f a) -> () -> (forall p. Kron f (Kron k p) -> Kron f' (Kron k' p) )
We want to allow the field a to unify across them all.
The scalar is generic in a certain sense, but not in another.

This is the lifted Kron'
We don't want Additive1, something went screwy there.
Additive (Kron' )

The problem with Num based Kron: If we define LinOp
(LinOp f g = forall Num a => f a -> g a
Why Num, Why not Fractional. This is a problem.
And really we want to be threading the scalar a throughout the entire system.

kron

And what if we unified the two orderings

Hmm. These are pretty similar to van Laarhoven lenses. This observation is porbably poison
Swtich the arguments
(a -> b) -> Lens
Functor f => (t -> f a) -> (t -> f b)
type Lens s t a b = forall f. Functor f => (a -> f b) -> s -> f t
Lens t a t b

kron :: (Fractional a => f a -> g a) -> (Fractional k) -> LinOp
kron l l' (Compose fk) =  Compose (fmap l' (l fk))

Constrained natural transformation
kron  forall const. (const a => f a -> g a) -> (forall b. const (k b) => k b -> l b)

(f Double -> g Double) ->
(forall f. V2' f Double -> V2' f Double)

kron :: LinOp f g -> LinOp k l -> LinOp (Kron f k) (Kron g l)

-}

{-

-- can't build. Well we can build vzero
exv2' :: (Additive f, Num a) => V2' f a
exv2' = V2' (Compose ( ?  ))

-}
-- constants are given as morphisms
-- exv2' :: (Additive f, Num a) => V1' f a -> V2' f a
-- exv2' (V1' f) = (V2' (V2   )

{-
There is no additive instance for Compose?
Huh.

class (forall f. Additive f => Additive k f) => Additive1 k where

I'm doing something insane.
Yea, The 1 pattern. I've been here before and decided it was insane.
Something in purescript. Surely it wasn't Additive1 was it?

-}
{-
class Additive1 k where
    vadd :: (Additive f, Num a) => k f a -> k f a -> k f a
    vzero :: (Additive f, Num a) => k f a
instance Additive1 V2' where -- (Functor (V2' f)) =>

    vadd (V2' (V2 v w)) (V2' (V2 x y)) = V2' (V2 (v ^+^ x)  (w ^+^ y)   )
    vzero = V2' (V2 zero zero)

    --vadd (V2' v) (V2' w) = V2' (v ^+^ w)
    --vzero = V2' zero

instance Additive1 V1' where -- (Functor (V2' f)) =>
    vadd (V1' v) (V1' w) = V1' (v ^+^ w)
    vzero = V1' zero

instance (Additive1 k, Additive1 l) => Additive1 (Kron' k l) where
   vadd (Kron' x) (Kron' y) = Kron' (vadd x y)
   vzero = Kron' vzero
   -}
{-
Deriv functor for V2'
-}


{-# LANGUAGE GADTs #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE DataKinds, PolyKinds #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE KindSignatures #-}
{-# LANGUAGE FlexibleContexts, FlexibleInstances #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE StandaloneDeriving #-}
{-# LANGUAGE ConstraintKinds #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE InstanceSigs #-}
{-# LANGUAGE UndecidableInstances #-}

module TensorHMatrix where

import Misc.Nats
import Numeric.LinearAlgebra as H
import Data.Vector.Storable as V
import Data.Default
import Data.Singletons.Prelude
import Generics.SOP.Constraint as C
import Misc.Constraints
import GenericTensor

import Prelude hiding (zipWith)

-- the kitchen sink of constraints
type Usable a = (Element a, Num a, Numeric a, Num (Vector a), Container Vector a, Floating a, Floating (Vector a))

data HTensor a (dims :: [k]) where
  Scalar :: !a -> HTensor a '[]
  Vector :: Vector a -> HTensor a '[n]
  Matrix :: Matrix a -> HTensor a '[n, m]

deriving instance (Show a, Element a) => Show (HTensor a dims)

instance (IntegralL dims, Default a, Usable a) => Default (HTensor a dims) where
  def = fill' def

--type instance IndexOf (HTensor l) = Rec LT l
--type instance IndexOf (HTensor l) = Rec (Const Int) l
--type instance IndexOf (HTensor l) = Vec (Length l) Int

{- Fails backwards fundep :(
instance Indexable (HTensor (n ': l) a) (HTensor l a) where
  Matrix m ! i = Vector (m ! i)
  Vector v ! i = Scalar (v ! i)
-}

instance Usable a => Tensor (HTensor a :: [k] -> *) where
  type N (HTensor a) = a

  scalar (Scalar a) = a

  transpose (Matrix m) = Matrix (tr' m)

  scale (Scalar a) (Vector v) = Vector (H.scale a v)

  dot (Vector a) (Vector b) = Scalar (a <.> b)

  outer (Vector a) (Vector b) = Matrix (H.outer a b)

  mv (Matrix m) (Vector v) = Vector $ m #> v

  mm (Matrix m1) (Matrix m2) = Matrix $ m1 <> m2

  select i (Vector v) = Scalar (v V.! i)

  oneHot :: forall (n :: k). (SingI n, IntegralN n) => Int -> HTensor a '[n]
  oneHot i = Vector $ (konst 0 (natVal' (sing::Sing n))) // [(i, 1)]

  --fill :: forall (dims :: [k]). (IntegralL dims) => a -> HTensor a dims
  fill sDims a = case sDims of
    SNil -> Scalar a
    SCons n SNil -> Vector $ konst a (natVal' n)
    SCons m (SCons n SNil) -> Matrix $ konst a (natVal' m, natVal' n)
    _ -> error "HTensors only go up to dimension 2."

-- pretty much all of these instances should be automatically derivable
-- the only issue is that fromInteger needs to use natVal'

instance (Usable a, IntegralL dims) => Num (HTensor a dims) where
  Scalar a + Scalar b = Scalar (a + b)
  Vector a + Vector b = Vector (a + b)
  Matrix a + Matrix b = Matrix (a + b)

  Scalar a - Scalar b = Scalar (a - b)
  Vector a - Vector b = Vector (a - b)
  Matrix a - Matrix b = Matrix (a - b)

  Scalar a * Scalar b = Scalar (a * b)
  Vector a * Vector b = Vector (a * b)
  Matrix a * Matrix b = Matrix (a * b)

  negate (Scalar a) = Scalar (negate a)
  negate (Vector a) = Vector (negate a)
  negate (Matrix a) = Matrix (negate a)

  abs (Scalar a) = Scalar (abs a)
  abs (Vector a) = Vector (abs a)
  abs (Matrix a) = Matrix (abs a)

  signum (Vector a) = Vector (signum a)
  signum (Scalar a) = Scalar (signum a)
  signum (Matrix a) = Matrix (signum a)

  fromInteger = fill' . fromInteger

instance (IntegralL dims, Usable a) => Fractional (HTensor a dims) where
  Scalar a / Scalar b = Scalar (a / b)
  Vector a / Vector b = Vector (a / b)
  Matrix a / Matrix b = Matrix (a / b)

  recip (Scalar a) = Scalar (recip a)
  recip (Vector a) = Vector (recip a)
  recip (Matrix a) = Matrix (recip a)

  fromRational = fill' . fromRational

instance (IntegralL dims, Usable a) => Floating (HTensor a dims) where
  pi = fill' pi

  exp (Scalar a) = Scalar (exp a)
  exp (Vector a) = Vector (exp a)
  exp (Matrix a) = Matrix (exp a)

  log (Scalar a) = Scalar (log a)
  log (Vector a) = Vector (log a)
  log (Matrix a) = Matrix (log a)

  sqrt (Scalar a) = Scalar (sqrt a)
  sqrt (Vector a) = Vector (sqrt a)
  sqrt (Matrix a) = Matrix (sqrt a)

  Scalar a ** Scalar b = Scalar (a ** b)
  Vector a ** Vector b = Vector (a ** b)
  Matrix a ** Matrix b = Matrix (a ** b)

  logBase (Scalar a) (Scalar b) = Scalar (logBase a b)
  logBase (Vector a) (Vector b) = Vector (logBase a b)
  logBase (Matrix a) (Matrix b) = Matrix (logBase a b)

  sin (Scalar a) = Scalar (sin a)
  sin (Vector a) = Vector (sin a)
  sin (Matrix a) = Matrix (sin a)

  cos (Scalar a) = Scalar (cos a)
  cos (Vector a) = Vector (cos a)
  cos (Matrix a) = Matrix (cos a)

  tan (Scalar a) = Scalar (tan a)
  tan (Vector a) = Vector (tan a)
  tan (Matrix a) = Matrix (tan a)

  asin (Scalar a) = Scalar (asin a)
  asin (Vector a) = Vector (asin a)
  asin (Matrix a) = Matrix (asin a)

  acos (Scalar a) = Scalar (acos a)
  acos (Vector a) = Vector (acos a)
  acos (Matrix a) = Matrix (acos a)

  atan (Scalar a) = Scalar (atan a)
  atan (Vector a) = Vector (atan a)
  atan (Matrix a) = Matrix (atan a)

  sinh (Scalar a) = Scalar (sinh a)
  sinh (Vector a) = Vector (sinh a)
  sinh (Matrix a) = Matrix (sinh a)

  cosh (Scalar a) = Scalar (cosh a)
  cosh (Vector a) = Vector (cosh a)
  cosh (Matrix a) = Matrix (cosh a)

  tanh (Scalar a) = Scalar (tanh a)
  tanh (Vector a) = Vector (tanh a)
  tanh (Matrix a) = Matrix (tanh a)

  asinh (Scalar a) = Scalar (asinh a)
  asinh (Vector a) = Vector (asinh a)
  asinh (Matrix a) = Matrix (asinh a)

  acosh (Scalar a) = Scalar (acosh a)
  acosh (Vector a) = Vector (acosh a)
  acosh (Matrix a) = Matrix (acosh a)

  atanh (Scalar a) = Scalar (atanh a)
  atanh (Vector a) = Vector (atanh a)
  atanh (Matrix a) = Matrix (atanh a)

instance Usable a => ImpliesC (IntegralL dims) (CompC Floating (HTensor a) dims) where
  impliesC = Sub Dict

instance Usable a => ForallC (ImpliesC1 IntegralL (CompC Floating (HTensor a))) where
  forallC = Forall Dict

tmap :: (Container Vector a, Num a, Element b) => (a -> b) -> HTensor a dims -> HTensor b dims
tmap f (Scalar a) = Scalar $ f a
tmap f (Vector v) = Vector $ cmap f v
tmap f (Matrix m) = Matrix $ cmap f m

tZipWith :: (Storable a, Storable b, Storable c) => (a -> b -> c) -> HTensor a '[n] -> HTensor b '[n] -> HTensor c '[n]
tZipWith f (Vector a) (Vector b) = Vector (zipWith f a b)

{-
transpose :: (Numeric a) => HTensor '[n, m] a -> HTensor '[m, n] a
transpose (Matrix m) = Matrix (tr' m)

asCol :: Storable a => HTensor '[n] a -> HTensor '[n, FromInteger 1] a
asCol (Vector v) = Matrix $ asColumn v

asRow' :: Storable a => HTensor '[n] a -> HTensor '[FromInteger 1, n] a
asRow' (Vector v) = Matrix $ asRow v

dot :: Numeric a => HTensor '[n] a -> HTensor '[n] a -> a
dot (Vector a) (Vector b) = a <.> b

mv :: Numeric a => HTensor '[n, m] a -> HTensor '[m] a -> HTensor '[n] a
mv (Matrix m) (Vector v) = Vector $ m #> v

mm :: Numeric a => HTensor '[n, m] a -> HTensor '[m, k] a -> HTensor '[n, k] a
mm (Matrix m1) (Matrix m2) = Matrix $ m1 <> m2

gradDot :: Numeric a => HTensor '[n] a -> HTensor '[n] a -> a -> (HTensor '[n] a, HTensor '[n] a)
gradDot (Vector a) (Vector b) g = (Vector $ scale g b, Vector $ scale g a)

gradMV :: Numeric a => HTensor '[n, m] a -> HTensor '[m] a -> HTensor '[n] a -> (HTensor '[n, m] a, HTensor '[m] a)
gradMV m v g = (mm (asCol g) (asRow' v), mv (transpose m) g)

gradMM :: Numeric a => HTensor '[n, m] a -> HTensor '[m, k] a -> HTensor '[n, k] a -> (HTensor '[n, m] a, HTensor '[m, k] a)
gradMM a b g = (mm g (transpose b), mm (transpose a) g)

gradSelect :: forall a n. (SingI n, IntegralN n, Usable a) => Int -> HTensor '[n] a -> a -> HTensor '[n] a
gradSelect i _ a = Vector $ (konst 0 (natVal' (sing::Sing n))) // [(i, a)]
-}


{-# OPTIONS_GHC -Wall #-}
{-# LANGUAGE BangPatterns #-}

module Numeric.MCMC.RiemannianLangevin (
            MarkovChain(..)
          , Options, createOptions
          , runChain, localMean, perturb
          ) where

import Numeric.MCMC.Langevin hiding (Options, runChain)
import Numeric.LinearAlgebra
import Control.Arrow
import Control.Monad
import Control.Monad.Primitive
import Control.Monad.Reader
import System.Random.MWC
import System.Random.MWC.Distributions

-- | Options for the chain.  The target (expected to be a log density), the
--   right matrix product of its curvature and gradient, the inverse Fisher metric tensor,
--   its (Cholesky) square root, and the step size.
data Options = Options {
    _target              :: [Double] -> Double        -- Target (log density)
  , _curvatureXgradient  :: [Double] -> [Double]      -- Curvature right-multiplied by gradient
  , _invFisherMetric     :: [Double] -> Matrix Double -- Inverse Fisher metric tensor
  , _sqrtInvFisherMetric :: [Double] -> Matrix Double -- Square root of the tensor
  , _eps                 :: {-# UNPACK #-} !Double    -- Step size
  }

-- | A result with this type has a view of the chain options.
type ViewsOptions = ReaderT Options

-- | Construct Options (data constructor not exported).
-- FIXME sqrtm is not converging; alternatives?
createOptions :: ([Double] -> Double)     -- Target (log density)
              -> ([Double] -> [Double])   -- Gradient
              -> ([Double] -> [[Double]]) -- Hessian
              -> Double                   -- Step size
              -> Options                  -- Options
createOptions t g h =
    Options t curvatureXgradient invFisherMetric sqrtInvFisherMetric
  where curvatureXgradient xs =
            let mat = invFisherMetric xs <> fromColumns [fromList (g xs)]
            in  concat . toLists . trans $ mat
        invFisherMetric       = pinv . fromRows . map (fromList . map (* (-1))) . h
        sqrtInvFisherMetric x = let z = schur (invFisherMetric x)
                                in  fst z <> sqrtm (snd z) <> trans (fst z)
{-# INLINE createOptions #-}

-- | Non-isotropic Gaussian density.
nonIsoGauss :: [Double] -> [Double] -> Matrix Double -> Double
nonIsoGauss xs mu sig = exp val
  where val = -0.5*p*log (2*pi) - 0.5*ldet - 0.5*
            (trans (xsM `sub` muM) <> invSig <> (xsM `sub` muM)) @@> (0, 0)
        (xsM, muM) = (\f (a, b) -> (f a, f b)) (\l -> fromColumns [fromList l]) (xs, mu)
        p                   = fromIntegral $ cols sig
        (invSig, (ldet, _)) = invlndet sig
{-# INLINE nonIsoGauss #-}

-- | Mean function for the discretized Riemannian Langevin diffusion.
localMean :: Monad m
          => [Double]                -- Current state
          -> ViewsOptions m [Double] -- Localized mean of proposal distribution
localMean t = do
    Options _ c _ _ e <- ask
    return $! zipWith (+) t (map (* (0.5 * e^(2 :: Int))) (c t))
{-# INLINE localMean #-}

-- | Perturb the state, creating a new proposal.
perturb :: PrimMonad m
        => [Double]                   -- Current state
        -> Gen (PrimState m)          -- MWC PRNG
        -> ViewsOptions m [Double]    -- Resulting perturbation.
perturb t g = do
    Options _ _ _ s _ <- ask
    zs <- replicateM (length t) (lift $ standard g)
    t0 <- localMean t
    let adjustedBrownianMotion = s t <> fromColumns [fromList zs]
        abmList = concat . toLists . trans $ adjustedBrownianMotion
        perturbedState = zipWith (+) t0 t1
        t1 = map (* eps) abmList
    return $! perturbedState
{-# INLINE perturb #-}

-- | Perform a Metropolis accept/reject step.
metropolisStep :: PrimMonad m
               => MarkovChain                -- Current state
               -> Gen (PrimState m)          -- MWC PRNG
               -> ViewsOptions m MarkovChain -- New state
metropolisStep state g = do
    Options target _ iF _ e <- ask
    let (t0, nacc) = (theta &&& accepts) state
    zc           <- lift $ uniformR (0, 1) g
    proposal     <- perturb t0 g
    t0Mean       <- localMean t0
    proposalMean <- localMean proposal
    let mc = if   zc < acceptProb
             then (proposal, 1)
             else (t0,       0)

        acceptProb = if isNaN val then 1 else val where val = arRatio

        arRatio = exp . min 0 $
            target proposal + log (nonIsoGauss proposal t0Mean ((e^(2::Int)) `scale` iF t0))
          - target t0       - log (nonIsoGauss t0 proposalMean ((e^(2::Int)) `scale` iF proposal))

    return $! MarkovChain (fst mc) (nacc + snd mc)
{-# INLINE metropolisStep #-}

-- | Diffuse through states.
runChain :: Options         -- Options of the Markov chain.
         -> Int             -- Number of epochs to iterate the chain.
         -> Int             -- Print every nth iteration
         -> MarkovChain     -- Initial state of the Markov chain.
         -> Gen RealWorld   -- MWC PRNG
         -> IO MarkovChain  -- End state of the Markov chain, wrapped in IO.
runChain = go
  where go o n t !c g | n == 0 = return c
                      | n `rem` t /= 0 = do
                            r <- runReaderT (metropolisStep c g) o
                            go o (n - 1) t r g
                      | otherwise = do
                            r <- runReaderT (metropolisStep c g) o
                            print r
                            go o (n - 1) t r g
{-# INLINE runChain #-}


{-# LANGUAGE AllowAmbiguousTypes, RankNTypes, UnicodeSyntax, NoMonomorphismRestriction #-}
{-# LANGUAGE BangPatterns, FlexibleInstances #-}

module Optimizer where

import Utils
import Style
import GenOptProblem
import Numeric.AD
import Numeric.AD.Internal.On
import Numeric.AD.Internal.Reverse
import Numeric.AD.Internal.Sparse
import qualified Numeric.LinearAlgebra as L
import Debug.Trace
import System.Random
import System.Console.ANSI
import Data.List (foldl')

default (Int, Float)

------ Opt types, util functions, and params

type ObjFn1 a = forall a . (Autofloat a) => [a] -> a

-- used for duf
type ObjFn2 a = forall a . (Autofloat a) => [a] -> [a] -> a

type GradFn a = forall a . (Autofloat a) => [a] -> [a]

----- Various consts

nanSub :: (Autofloat a) => a
nanSub = 0

infinity :: Floating a => a
infinity = 1/0 -- x/0 == Infinity for any x > 0 (x = 0 -> Nan, x < 0 -> -Infinity)
-- all numbers are smaller than infinity except infinity, to which it's equal

----- Hyperparameters

weightGrowthFactor :: (Autofloat a) => a -- for EP weight
weightGrowthFactor = 10

epsUnconstr :: Floating a => a
epsUnconstr = 10 ** (-2)

epStop :: Floating a => a -- for EP diff
epStop = 10 ** (-5)
-- epStop = 60 ** (-3)
-- epStop = 10 ** (-1)
-- epStop = 0.05

-- Parameters for Armijo-Wolfe line search
-- NOTE: must maintain 0 < c1 < c2 < 1
c1 :: Floating a => a
c1 = 0.001 -- for Armijo, corresponds to alpha in backtracking line search (see below for explanation)
-- smaller c1 = shallower slope = less of a decrease in fn value needed = easier to satisfy
-- turn Armijo off: c1 = 0
-- Nocedal p38: "In practice, c1 is chosen to be quite small, say 10−4."

c2 :: Floating a => a
c2 = 0.9 -- for Wolfe, is the factor decrease needed in derivative value
-- new directional derivative value / old DD value <= c2
-- smaller c2 = smaller new derivative value = harder to satisfy
-- turn Wolfe off: c1 = 1 (basically backatracking line search only)
-- Nocedal p39: "Typical values of c2 are 0.9 when the search direction pk is chosen by a Newton or quasi-Newton method"

-- true = force linesearch halt if interval gets too small; false = no forced halt
intervalMin = True

useLineSearch :: Bool
useLineSearch = True

useAutodiff :: Bool
useAutodiff = True

constT :: Floating a => a
constT = 0.001

debugOpt = True
debugLineSearch = False
debugBfgs = True

trb :: String -> a -> a
trb s x = if debugBfgs then trace "---" $ trace s x else x -- prints in left to right order

tro :: String -> a -> a
tro s x = if debugOpt then trace "---" $ trace s x else x -- prints in left to right order

trl :: Show a => String -> a -> a
trl s x = if debugLineSearch then trace "---" $ trace s $ traceShowId x else x -- prints in left to right order

----- Convergence criteria

-- convergence criterion for EP
-- if you want to use it for UO, needs a different epsilon
epStopCond :: (Autofloat a) => [a] -> [a] -> a -> a -> Bool
epStopCond x x' fx fx' =
           tro ("EP: \n||x' - x||: " ++ (show $ norm (x -. x'))
           ++ "\n|f(x') - f(x)|: " ++ (show $ abs (fx - fx'))) $
           (norm (x -. x') <= epStop) || (abs (fx - fx') <= epStop)

unconstrainedStopCond :: (Autofloat a) => [a] -> Bool
unconstrainedStopCond gradEval = norm gradEval < epsUnconstr

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

-- Policies

stepPolicy :: State -> (Params, PolicyParams)
stepPolicy s =
    -- Check overall convergence first
    let epStatus = optStatus $ (paramsr s) in
    let pparams = policyParams s in
    case epStatus of

    -- Generate new objective function and replace the optimization and policy params accordingly
    EPConverged ->
                -- TODO: clean up the step incrementing
        let pparams' = pparams { policySteps = 1 + policySteps pparams } in
        let (policyRes, psNew) = (policyFn s) (objFns s) (constrFns s) pparams' in -- See what the policy function wants
        case policyRes of
            Nothing     -> (paramsr s, pparams' { policyState = psNew }) -- steps incremented, policy done

            Just newFns -> -- Policy keeps going
                let objFnNew = genObjfn (transr s) (filter isObjFn newFns) (filter isConstr newFns) (varyingPaths s)
                    -- TODO: check that these inputs are right
                    -- Change obj function and restart optimization
                    pparamsNew = pparams' { policyState = psNew,
                                            currFns = newFns }
                    paramsNew = Params { weight = initWeight,
                                         optStatus = NewIter,
                                         overallObjFn = objFnNew,
                                         bfgsInfo = defaultBfgsParams }
                in tro ("Step policy, EP converged, new params:\n" ++ show (paramsNew, pparamsNew, newFns)) $ (paramsNew, pparamsNew)

    -- If not converged, optimize as usual, don't change policy mid-optimization
    _ -> tro ("Step policy, EP not converged, new params:\n" ++ show (paramsr s, pparams)) $ (paramsr s, pparams)

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

-- Main optimization functions

step :: State -> State
step s = let (state', params') = stepShapes (oConfig s) (paramsr s) (varyingState s) (rng s)
             s'                = s { varyingState = state',
                                     paramsr = params' }
             -- NOTE: we intentionally discard the random generator here because
             -- we want to have consistent computation output in a single
             -- optimization session
             -- For the same reason, all subsequent step* functions such as
             -- stepShapes do not return the new random generator
             (!shapes', _, _)     = evalTranslation s'

             -- Check the state and see if the overall objective function should be changed
             -- The policy may change EPConverged to a new iteration before the frontend sees it
             (paramsNew, pparamsNew) = stepPolicy s'

             -- For debugging
             oldParams = paramsr s

         in tro ("Params: \n" ++ show oldParams ++ "\n:") $
            s' { shapesr = shapes',
                 paramsr = paramsNew,
                 policyParams = pparamsNew }
            -- note: trans is not updated in state

-- Note use of realToFrac to generalize type variables (on the weight and on the varying state)

-- implements exterior point algo as described on page 6 here:
-- https://www.me.utexas.edu/~jensen/ORMM/supplements/units/nlp_methods/const_opt.pdf
stepShapes :: OptConfig -> Params -> [Float] -> StdGen -> ([Float], Params)
stepShapes config params vstate g = -- varying state
         -- if null vstate then error "empty state in stepshapes" else
         let (epWeight, epStatus) = (weight params, optStatus params) in
         case epStatus of

         -- start the outer EP optimization and the inner unconstrained optimization, recording initial EPstate
         NewIter -> let status' = UnconstrainedRunning (map realToFrac vstate) in
                    (vstate', params { weight = initWeight, optStatus = status', bfgsInfo = defaultBfgsParams } )
         -- check *weak* convergence of inner unconstrained opt.
         -- if UO converged, set opt state to converged and update UO state (NOT EP state)
         -- if not, keep running UO (inner state implicitly stored)
         -- note convergence checks are only on the varying part of the state
         UnconstrainedRunning lastEPstate ->  -- doesn't use last EP state
           -- let unconstrConverged = unconstrainedStopCond gradEval in
           let unconstrConverged = epStopCond vstate vstate' (objFnApplied vstate) (objFnApplied vstate') in
               -- Two stopping conditions
               -- unconstrainedStopCond gradEval in
           if unconstrConverged then
              let status' = UnconstrainedConverged lastEPstate in -- update UO state only!
              (vstate', params { optStatus = status', bfgsInfo = defaultBfgsParams }) -- note vstate' (UO converged), not vstate
           else (vstate', params { bfgsInfo = bfgs' }) -- update UO state but not EP state; UO still running

         -- check EP convergence. if converged then stop, else increase weight, update states, and run UO again
         -- TODO some trickiness about whether unconstrained-converged has updated the correct state
         -- and whether to check WRT the updated state or not
         UnconstrainedConverged lastEPstate ->
           let epConverged = epStopCond lastEPstate (map r2f vstate) -- lastEPstate is last state for converged UO
                                   (objFnApplied lastEPstate) (objFnApplied (map r2f vstate)) in
           if epConverged then
              let status' = EPConverged in -- no more EP state
              (vstate, params { optStatus = status', bfgsInfo = defaultBfgsParams }) -- do not update UO state
           -- update EP state: to be the converged state from the most recent UO
           else let status' = UnconstrainedRunning (map realToFrac vstate) in -- increase weight
                let epWeight' = weightGrowthFactor * epWeight in
                -- trace ("Unconstrained converged. New weight: " ++ show epWeight') $
                      (vstate, params { weight = epWeight', optStatus = status', bfgsInfo = defaultBfgsParams })

         -- done; don't update obj state or params; user can now manipulate
         EPConverged -> (vstate, params { bfgsInfo = defaultBfgsParams } )

         -- TODO: implement EPConvergedOverride (for when the magnitude of the gradient is still large)

         -- TODO factor out--only unconstrainedRunning needs to run stepObjective, but EPconverged needs objfn
        where (vstate', gradEval, bfgs') = stepWithObjective config g params vstate
              objFnApplied = (overallObjFn params) g (r2f $ weight params)

-- Given the time, state, and evaluated gradient (or other search direction) at the point,
-- return the new state

stepT :: Float -> Float -> Float -> Float
stepT dt x dfdx = x - dt * dfdx

-- Calculates the new state by calculating the directional derivatives (via autodiff)
-- and timestep (via line search), then using them to step the current state.
-- Also partially applies the objective function.
stepWithObjective :: OptConfig -> StdGen -> Params -> [Float] -> ([Float], [Float], BfgsParams)
stepWithObjective config g params state =
          -- get timestep via line search, and evaluated gradient at the state
          let (t', gradEval, gradToUse, bfgs') = timeAndGrad config params objFnApplied state
              -- step each parameter of the state with the time and gradient
              state' = map (\(v, dfdv) -> stepT t' v dfdv) (zip state $ gradToUse)
              (fx, fx') = (objFnApplied state, objFnApplied state')
          in -- if fx' > fx then error ("Error: new energy is greater than old energy: " ++ show (fx', fx)) else
             tro ("\n----------------------------------------\n"
                   ++ "\nopt params: \n" ++ (show params)
                   ++ "\n||x' - x||: " ++ (show $ norm (state -. state'))
                   ++ "\n|f(x') - f(x)|: " ++
                  (show $ abs (fx' - fx))
                   ++ "\nf(x'): \n" ++ (show fx')
                   ++ "\ngradEval: \n" ++ (show gradEval)
                   ++ "\n||gradEval||: \n" ++ (show $ norm gradEval)
                   ++ "\ngradToUse: \n" ++ (show gradToUse)
                   ++ "\n||gradToUse||: \n" ++ (show $ norm gradToUse)
                   -- ++ "\nhessian: \n" ++ (show $ h)
                   ++ "\nbfgs': \n" ++ (show bfgs') -- TODO: use trb
                   ++ "\n timestep: \n" ++ (show t')
                   ++ "\n original state: \n" ++ (show state)
                   ++ "\n new state: \n" ++ (show state')
                  )
             (state', gradEval, bfgs')

          where objFnApplied :: ObjFn1 b
                objFnApplied = (overallObjFn params) g cWeight

                cWeight = r2f $ weight params
                -- realToFrac generalizes the type variable `a` to the type variable `b`, which timeAndGrad expects

-- a version of grad with a clearer type signature
appGrad :: (Autofloat a) => (forall b . (Autofloat b) => [b] -> b) -> [a] -> [a]
appGrad f l = grad f l

instance Show (Numeric.AD.Internal.On.On
                             (Numeric.AD.Internal.Reverse.Reverse
                                s (Numeric.AD.Internal.Sparse.Sparse a))) where
         show a = "error: not sure how to derive show for hessian element"

appHess :: (Autofloat a) => (forall b . (Autofloat b) => [b] -> b) -> [a] -> [[a]]
appHess f l = hessian f l

-- Precondition the gradient
gradP :: OptConfig -> BfgsParams -> [Double] -> ObjFn1 a -> [Float] -> ([Double], BfgsParams)
gradP config bfgsParams gradEval f state =
      let x_k = L.vector $ map r2f state
          grad_fx_k = L.vector gradEval
      in case optMethod config of
      GradientDescent -> (gradEval, bfgsParams)

      Newton -> -- Precondition gradient with the pseudoinverse of the hessian
          let h = appHess f state
              h_list = (map r2f $ concat h) :: [Double]
              hinv = L.pinv $ L.matrix (length gradEval) $ h_list
              gradPreconditioned = hinv L.#> (L.vector gradEval) in
          (L.toList gradPreconditioned, bfgsParams)

      BFGS -> -- Approximate inverse of hessian with the change in gradient (see Nocedal S9.1 p224)
           let grad_val = lastGrad bfgsParams
               h_val = invH bfgsParams
               state_val = lastState bfgsParams
           in case (h_val, grad_val, state_val) of
              (Nothing, Nothing, Nothing) -> -- First step. Initialize the approximation to the identity (not clear how else to approximate it)
                    -- k=0 steps from x_0 to x_1: so on the first step, we take a normal gradient descent step
                    let h_0 = L.ident $ length gradEval
                    in (gradEval, bfgsParams { lastState = Just x_k, lastGrad = Just grad_fx_k, invH = Just h_0 })

              (Just h_km1, Just grad_fx_km1, Just x_km1) -> -- For x_{k+1}, to compute H_k, we need the (k-1) info
                    -- Our convention is that we are always working "at" k to compute k+1
                    -- x_0 doesn't require any H; x_1 (the first step) with k = 0 requires H_0
                    -- x_2 (the NEXT step) with k=1 requires H_1. For example>
                    -- x_2 = x_1 - alpha_1 H_1 grad f(x_1)   [GD step]
                    -- H_1 = V_0 H_0 V_0 + rho_0 s_0 s_0^T   [This is confusing because the book adds an extra +1 to the H index]
                    -- V_0 = I - rho_0 y_0 s_0^T
                    -- rho_0 = 1 / y_0^T s_0
                    -- s_0 = x_1 - x_0
                    -- y_0 = grad f(x_1) - grad f(x_0)

                    let s_km1 = x_k - x_km1
                        y_km1 = grad_fx_k - grad_fx_km1
                        rho_km1 = 1 / (y_km1 `L.dot` s_km1) -- Scalar
                        v_km1 = L.ident (length gradEval) - (rho_km1 `L.scale` y_km1 `L.outer` s_km1) -- Scaling can happen before outer
                        h_k = (L.tr (v_km1) L.<> h_km1 L.<> v_km1) + (rho_km1 `L.scale` s_km1 `L.outer` s_km1)
                        gradPreconditioned = h_k L.#> grad_fx_k
                    in (L.toList gradPreconditioned,
                        bfgsParams { lastState = Just x_k, lastGrad = Just grad_fx_k, invH = Just h_k })

              _ -> error "invalid BFGS state"

      LBFGS -> -- Approximate the inverse of the Hessian times the gradient
               -- Only using the last `m` gradient/state difference vectors, not building the full h_k matrix (Nocedal p226)
            let grad_prev = lastGrad bfgsParams
                x_prev = lastState bfgsParams
                ss_val = s_list bfgsParams
                ys_val = y_list bfgsParams
                km1 = numUnconstrSteps bfgsParams -- Our current step is k; the last step is km1 (k_minus_1)
                m = memSize bfgsParams
            in case (grad_prev, x_prev, ss_val, ys_val, km1) of

               -- Perform normal gradient descent on first step
               (Nothing, Nothing, [], [], 0) ->
                   -- Store x_k, grad f(x_k) so we can compute s_k, y_k on next step
                   let bfgsParams' = bfgsParams { lastState = Just x_k, lastGrad = Just grad_fx_k,
                                                  s_list = [], y_list = [], numUnconstrSteps = km1 + 1 } in
                   (gradEval, bfgsParams')

               (Just grad_fx_km1, Just x_km1, ss, ys, km1) ->
                   -- Compute s_{k-1} = x_k - x_{k-1} and y_{k-1} = (analogous with grads)
                   -- Unlike Nocedal, compute the difference vectors first instead of last (same result, just a loop rewrite)
                   -- Use the updated {s_i} and {y_i}. (If k < m, this reduces to normal BFGS, i.e. we use all the vectors so far)
                   let (s_km1, y_km1) = (x_k - x_km1, grad_fx_k - grad_fx_km1) -- Newest vectors added to front
                       (ss', ys') = (take m $ s_km1 : ss, take m $ y_km1 : ys) -- The limited-memory part: drop stale vectors
                       gradPreconditioned = lbfgs grad_fx_k ss' ys'
                       descentDirCheck = -gradPreconditioned `L.dot` grad_fx_k

                   -- Reset L-BFGS if the result is not a descent direction, and use steepest descent direction
                   -- https://github.com/JuliaNLSolvers/Optim.jl/issues/143 https://github.com/JuliaNLSolvers/Optim.jl/pull/144

                   in if descentDirCheck < epsd
                      then (L.toList gradPreconditioned, bfgsParams { lastState = Just x_k, lastGrad = Just grad_fx_k,
                                                                      s_list = ss', y_list = ys', numUnconstrSteps = km1 + 1 })
                      else tro "L-BFGS did not find a descent direction. Resetting correction vectors." $
                           (gradEval, bfgsParams { lastState = Just x_k, lastGrad = Just grad_fx_k,
                                                   s_list = [], y_list = [], numUnconstrSteps = km1 + 1 })

                   -- TODO: check the curvature condition y_k^T s_k > 0 (8.7) (Nocedal 201)
                   -- https://github.com/JuliaNLSolvers/Optim.jl/issues/26

               _ -> error "invalid L-BFGS state"

type LVector = L.Vector L.R
type LMatrix = L.Matrix L.R

-- See Nocedal p225
-- expects ys and ss to be in order from most recent to oldest (k-1 ... k-m)
-- expects length ys == length ss, length ys > 0, length ys > 0
lbfgs :: LVector -> [LVector] -> [LVector] -> LVector
lbfgs grad_fx_k ss ys =
    let rhos = map calculate_rho $ zip ss ys -- The length of any list should be the number of stored vectors
        q_k = grad_fx_k
        (q_k_minus_m, alphas) = foldl' pull_q_back (q_k, []) (zip3 rhos ss ys) -- backward: for i = k-1 ... k-m
        -- Note the order of alphas will be from k-m through k-1 for the push_r_forward loop
        h_0_k = estimate_hess (head ys) (head ss)
        r_k_minus_m = h_0_k L.#> q_k_minus_m
        r_k = foldl' push_r_forward r_k_minus_m (zip (zip (reverse rhos) alphas) (zip (reverse ss) (reverse ys)))
        -- forward: for i = k-m .. k-1 (TODO: optimize out the reverses)
    in r_k -- is H_k * grad f(x_k)

           where calculate_rho :: (LVector, LVector) -> L.R
                 calculate_rho (s, y) = 1 / ((y `L.dot` s) + epsd)

                 pull_q_back :: (LVector, [L.R]) -> (L.R, LVector, LVector) -> (LVector, [L.R]) -- from i+1 to i
                 pull_q_back (q_i_plus_1, alphas) (rho_i, s_i, y_i) =
                        let alpha_i = rho_i * (s_i `L.dot` q_i_plus_1) -- scalar
                            q_i = q_i_plus_1 - (alpha_i `L.scale` y_i) -- scalar * vector
                        in (q_i, alpha_i : alphas) -- Note the order of alphas

                 -- Scale I by an estimate of the size of the Hessian along the most recent search direction (Nocedal p226)
                 estimate_hess :: LVector -> LVector -> LMatrix
                 estimate_hess y_km1 s_km1 =
                        let gamma_k = (s_km1 `L.dot` y_km1) / ((y_km1 `L.dot` y_km1) + epsd)
                        in gamma_k `L.scale` (L.ident (L.size y_km1))

                 push_r_forward :: LVector -> ((L.R, L.R), (LVector, LVector)) -> LVector -- from i to i+1
                 push_r_forward r_i ((rho_i, alpha_i), (s_i, y_i)) =
                        let beta_i = rho_i * (y_i `L.dot` r_i) -- scalar
                            r_i_plus_1 = r_i + (alpha_i - beta_i) `L.scale` s_i -- scalar * vector
                        in r_i_plus_1

estimateGradient :: ObjFn1 a -> [Float] -> [Float]
estimateGradient f state =
      let len = length state
          h = 0.001 -- Choice of h really matters!!! This is not an accurate estimate.
          fx = f state
      -- time is O(|state|^2)
          dfx i = ((f $ replace i ((state !! i) + h) state) - fx) / h
          dfxs = map dfx [0..(len-1)]
      in dfxs
      where replace pos newVal list = take pos list ++ newVal : drop (pos+1) list

-- Given the objective function, gradient function, timestep, and current state,
-- return the timestep (found via line search) and evaluated gradient at the current state.
-- the autodiff library requires that objective functions be polymorphic with Floating a
timeAndGrad :: OptConfig -> Params -> ObjFn1 a -> [Float] -> (Float, [Float], [Float], BfgsParams)
timeAndGrad config params f state =
            let gradEval = if useAutodiff then gradF state else estimateGradient f state
                (gradToUse_d, bfgs') = gradP config (bfgsInfo params) (map r2f gradEval :: [Double]) f state
                gradToUse = map r2f gradToUse_d
                -- Use line search to find a good timestep. If we use Newton's method, the descent direction uses the preconditioned gradient.
                descentDir = negL $ gradToUse
                timestep = let resT = if useLineSearch && useAutodiff
                                      then awLineSearch config f (duf descentDir) descentDir state
                                      else constT in -- hardcoded timestep
                           if isNaN resT then error "returned timestep is NaN" else resT

            in tr "timeAndGrad: " (timestep, gradEval, gradToUse, bfgs')

            where gradF :: GradFn a
                  gradF = appGrad f

                  -- directional derivative of f at x in the direction of u (descent direction, which may not have unit norm)
                  duf :: (Autofloat a) => [a] -> [a] -> a
                  duf u x = u `dotL` gradF x

linesearch_max :: Int
linesearch_max = 100 -- TODO what's a reasonable limit (if any)?

-- Implements Armijo-Wolfe line search as specified in Keenan's notes, converges on nonconvex fns as well
-- based off Lewis & Overton, "Nonsmooth optimization via quasi-Newton methods
-- duf = D_u(f), the directional derivative of f at descent direction u
-- D_u(x) = <gradF(x), u>. If u = -gradF(x) (as it is here), then D_u(x) = -||gradF(x)||^2

awLineSearch :: OptConfig -> ObjFn1 a -> ObjFn1 a -> [Float] -> [Float] -> Float
awLineSearch config f duf descentDir x0 =

     let (a0, b0, t0) = (0, infinity, 1) -- Unit step length should be tried first for quasi-Newton methods. (Nocedal 201)

         -- drop while a&w are not satisfied OR the interval is large enough
         (numUpdates, (af, bf, tf)) = head $ dropWhile intervalOK_or_notArmijoAndWolfe
                                      $ zip [0..] $ iterate update (a0, b0, t0)

     in tro ("Line search # updates: " ++ show numUpdates) $ tf

          where update :: (Float, Float, Float) -> (Float, Float, Float)
                update (a, b, t) =
                       let (a', b', sat) | not $ armijo t    = trl "not armijo" (a, t, False)
                                         | not $ wolfe t     = trl "not wolfe" (t, b, False)
                                         | otherwise         = (a, b, True) in
                       if sat then (a, b, t) -- if armijo and wolfe, then we use (a, b, t) as-is
                       else if b' < infinity then trl "b' < infinity" (a', b', (a' + b') / 2)
                       else trl "b' = infinity" (a', b', 2 * a')

                intervalOK_or_notArmijoAndWolfe :: (Int, (Float, Float, Float)) -> Bool
                intervalOK_or_notArmijoAndWolfe (numUpdates, (a, b, t)) =
                      not $
                      if armijo t && wolfe t then
                           trl ("stop: both sat. |descentDir at x0| = " ++ show (norm descentDir)) True
                      else if abs (b - a) < minInterval then
                           trl ("stop: interval too small. |descentDir at x0| = " ++ show (norm descentDir)) True
                      else if numUpdates > linesearch_max then
                           trl ("stop: number of line search updates exceeded max") True
                      else False

                armijo :: Float -> Bool
                armijo t = (f (x0 +. t *. descentDir)) <= (fAtx0 + c1 * t * dufAtx0)

                weakWolfe :: Float -> Bool -- Better for nonsmooth functions
                weakWolfe t = (duf (x0 +. t *. descentDir)) >= (c2 * dufAtx0)

                strongWolfe :: Float -> Bool
                strongWolfe t = (abs (duf (x0 +. t *. descentDir))) <= (c2 * abs dufAtx0)

                -- Descent direction (at x0) must have a negative dot product with the gradient of f at x0.
                dufAtx0 :: Float -- TODO: this is redundant, cache it
                dufAtx0 = let res = duf x0 in
                        trl ("<grad f(x0), descent direction at x0>: " ++ show res) $
                        if res > 0 then tro "WARNING: descent direction doesn't satisfy condition" res else res

                fAtx0 :: Float
                fAtx0 = f x0

                minInterval :: Float -- stop if the interval gets too small; might not terminate
                minInterval = if intervalMin then 10 ** (-10) else 0
                -- TODO: the line search is very sensitive to this parameter. Blows up with 10**(-5). Why?

                wolfe :: Float -> Bool
                wolfe = weakWolfe


{-# LANGUAGE RankNTypes, FlexibleContexts, TypeFamilies #-}

module Numeric.Trainee.Gradee (
	Gradee(..), gradee, ad,
	Unary, Binary,
	unary, binary,
	dup, vdup, vdupWith,
	conjoin, plus, vconjoin, vsum, vfold,
	swap,

	matMat, matVec, odot,
	corrVec, corrMat,
	flattenMat, reshapeVec, concatVecs,
	transposeMat,
	vecRow, vecCol,
	biasVec, biasMat,
	maxPoolVec, maxPoolMat,
	padVec, padMat
	) where

import Prelude hiding (id, (.))
import Prelude.Unicode

import Control.Category
import Control.Lens (lens)
import Data.Reflection (Reifies, reify)
import qualified Data.Vector as V
import Numeric.AD (grad, auto)
import Numeric.AD.Internal.Reverse (Reverse, primal, Tape)
import Numeric.LinearAlgebra

import Numeric.Trainee.Types

-- | Make Gradee like lens
gradee ∷ (a → b) → (a → b → a) → Gradee a b
gradee g s = Gradee $ lens g s

-- | Make like lens from binary op
gradee2 ∷ (a → b → c) → (a → b → c → (a, b)) → Gradee (a, b) c
gradee2 g s = gradee (uncurry g) (uncurry s)

-- | Make Gradee from any function
ad ∷ (Traversable f, Num a) ⇒ (forall s . Reifies s Tape ⇒ f (Reverse s a) → Reverse s a) → Gradee (f a) a
ad f = gradee f' (\x dx → fmap (* dx) (grad f x)) where
	f' = reify undefined (\p → primal ∘ spec p f ∘ fmap auto)
	spec ∷ Reifies t Tape ⇒ proxy t → (forall s . Reifies s Tape ⇒ g (Reverse s a) → Reverse s a) → g (Reverse t a) → Reverse t a
	spec _ h = h

type Unary a = forall s . Reifies s Tape ⇒ Reverse s a → Reverse s a
type Binary a = forall s . Reifies s Tape ⇒ Reverse s a → Reverse s a → Reverse s a

-- | Make Gradee from unary function
unary ∷ Num a ⇒ Unary a → Gradee a a
unary f = ad (\[x] → f x) . gradee return (const head)

-- | Make @Gradee@ from binary function
binary ∷ Num a ⇒ Binary a → Gradee (a, a) a
binary f = ad (\[x, y] → f x y) . gradee g s where
	g (x, y) = [x, y]
	s _ [x, y] = (x, y)
	s _ _ = error "binary"

dup ∷ Num a ⇒ Gradee a (a, a)
dup = gradee (\x → (x, x)) (\_ (dx', dx'') → dx' + dx'')

vdup ∷ Num a ⇒ Int → Gradee a (V.Vector a)
vdup n = gradee (V.replicate n) (const V.sum)

vdupWith ∷ (a → a → a) → Int → Gradee a (V.Vector a)
vdupWith fn n
	| n ≤ 0 = error "vdupWith: negative argument"
	| otherwise = gradee (V.replicate n) (const $ V.foldr1 fn)

conjoin ∷ Num a ⇒ Gradee (a, a) a
conjoin = binary (+)

plus ∷ Num a ⇒ Gradee (a, a) a
plus = conjoin

vconjoin ∷ Num a ⇒ Gradee (V.Vector a) a
vconjoin = ad V.sum

vsum ∷ Num a ⇒ Gradee (V.Vector a) a
vsum = vconjoin

vfold ∷ (a → a → a) → Gradee (V.Vector a) a
vfold fn = gradee (V.foldr1 fn) (V.replicate ∘ V.length)

swap ∷ Gradee (a, b) (b, a)
swap = gradee (\(x, y) → (y, x)) (\_ (dy, dx) → (dx, dy))

matMat ∷ Numeric a ⇒ Gradee (Matrix a, Matrix a) (Matrix a)
matMat = gradee2 (<>) backprop where
	backprop a b dc = (dc <> tr b, tr a <> dc)

matVec ∷ Numeric a ⇒ Gradee (Matrix a, Vector a) (Vector a)
matVec = gradee2 (#>) backprop where
	backprop a b dc = (outer dc b, tr a #> dc)

odot ∷ Num (Vector a) ⇒ Gradee (Vector a, Vector a) (Vector a)
odot = gradee2 (+) backprop where
	backprop _ _ dc = (dc, dc)

corrVec ∷ Numeric a ⇒ Gradee (Vector a, Vector a) (Vector a)
corrVec = gradee2 corr backprop where
	backprop a b dc = (corr dc b, conv a dc)

corrMat ∷ (Numeric a, Num (Vector a)) ⇒ Gradee (Matrix a, Matrix a) (Matrix a)
corrMat = gradee2 corr2 backprop where
	backprop a b dc = (corr2 dc b, conv2 a dc)

flattenMat ∷ Numeric a ⇒ Gradee (Matrix a) (Vector a)
flattenMat = gradee flatten backprop where
	backprop a = reshape (cols a)

reshapeVec ∷ Numeric a ⇒ Int → Gradee (Vector a) (Matrix a)
reshapeVec cols' = gradee (reshape cols') (const flatten)

concatVecs ∷ Numeric a ⇒ Gradee (V.Vector (Vector a)) (Vector a)
concatVecs = gradee
	(vjoin ∘ V.toList)
	(\vs → V.fromList ∘ takesV (V.toList (V.map size vs)))

transposeMat ∷ Numeric a ⇒ Gradee (Matrix a) (Matrix a)
transposeMat = gradee tr (const tr)

vecRow ∷ Numeric a ⇒ Gradee (Vector a) (Matrix a)
vecRow = gradee asRow (const flatten)

vecCol ∷ Numeric a ⇒ Gradee (Vector a) (Matrix a)
vecCol = gradee asColumn (const flatten)

biasVec ∷ Numeric a ⇒ Gradee (a, Vector a) (Vector a)
biasVec = gradee2 (\b v → cmap (+ b) v) backprop where
	backprop _ _ dv = (sumElements dv, dv)

biasMat ∷ Numeric a ⇒ Gradee (a, Matrix a) (Matrix a)
biasMat = gradee2 (\b m → cmap (+ b) m) backprop where
	backprop _ _ dm = (sumElements dm, dm)

maxPoolVec ∷ (RealFrac a, Numeric a) ⇒ Int → Gradee (Vector a) (Vector a)
maxPoolVec n = gradee pool' unpool' where
	pool' v = build (size v `div` n) $ \i → maxElement (subVector (round i * n) n v)
	unpool' v dv = accum zero (+) (zip maxIndices (toList dv)) where
		maxIndices = map (\i → maxIndex (subVector (i * n) n v) + i * n) [0 .. size dv - 1]
		zero = build (size v) (const 0)

maxPoolMat ∷ (RealFrac a, Numeric a) ⇒ Int → Int → Gradee (Matrix a) (Matrix a)
maxPoolMat w h = gradee pool' unpool' where
	pool' m = build (rows m `div` h, cols m `div` w) $ \i j → maxElement (subMatrix (round i * h, round j * w) (h, w) m)
	unpool' m dm = accum zero (+) (zip maxIndices (toList $ flatten dm)) where
		maxIndices = map (\(i, j) → maxIndex (subMatrix (i * h, j * w) (h, w) m) `biplus` (i * h, j * w)) [(i, j) | i ← [0 .. rows dm - 1], j ← [0 .. cols dm - 1]]
		zero = build (size m) (const $ const 0)
		biplus (lx, ly) (rx, ry) = (lx + rx, ly + ry)

padVec ∷ Numeric a ⇒ Int → Gradee (Vector a) (Vector a)
padVec n = gradee pad' unpad' where
	pad' v = vjoin [zeroes', v, zeroes'] where
		zeroes' = build n (const 0)
	unpad' = subVector n ∘ size

padMat ∷ Numeric a ⇒ Int → Int → Gradee (Matrix a) (Matrix a)
padMat w h = gradee pad' unpad' where
	pad' m = diagBlock [zeroes', m, zeroes'] where
		zeroes' = build (h, w) (const $ const 0)
	unpad' = subMatrix (h, w) ∘ size


-- Copyright (c) 2015 James M. Lawrence
--
-- Permission is hereby granted, free of charge, to any person obtaining
-- a copy of this software and associated documentation files (the
-- "Software"), to deal in the Software without restriction, including
-- without limitation the rights to use, copy, modify, merge, publish,
-- distribute, sublicense, and/or sell copies of the Software, and to
-- permit persons to whom the Software is furnished to do so, subject to
-- the following conditions:
--
-- The above copyright notice and this permission notice shall be included
-- in all copies or substantial portions of the Software.
--
-- THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
-- EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
-- MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.
-- IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY
-- CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT,
-- TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE
-- SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
--
-----------------------------------------------------------------------------
-- |
-- Module      :  Math.CayleyDickson
-- Copyright   :  (c) James M. Lawrence
-- License     :  MIT
--
-- Maintainer  :  James M. Lawrence <llmjjmll@gmail.com>
-- Stability   :  provisional
-- Portability :  portable
--
-- Cayley-Dickson constructions (complex numbers, quaternions,
-- octonions, sedenions, etc.) over general scalars without limit to
-- the number of dimensions.
--
-- An element of this structure is composed of an m-dimensional
-- /scalar part/ and an m*(2^n - 1)-dimensional /pure part/ (unrelated
-- to Haskell's uses of "pure"). An element whose scalar part is zero
-- is called a /pure/. Construction with real scalars yields the
-- Cayley-Dickson algebras, in which case the scalar part is also
-- called the /real part/. Other structures may be obtained by
-- considering general scalars, for instance the quaternions over
-- complex scalars.
-----------------------------------------------------------------------------

module Math.CayleyDickson (
    -- * Types
    Nion, Complex, Quaternion, Octonion, Sedenion,

    -- * Construction
    nion, fromScalar, complex, quaternion, octonion, sedenion,

    -- * Operations
    dot, cross, sqnorm, norm, polar,
    -- ** Operations with scalars
    --
    -- | The mnemonic is that the period (".") is on the side of the
    -- scalar.
    (**.),
    (.+), (+.), (.-), (-.), (.*), (*.), (/.),

    -- * Accessors
    coord, coords, setCoord, scalarPart, purePart,

    -- * Constants
    basisElement,

    -- * Classes
    Conjugable(..),

    -- ** Tags
    Tag(tagVal),
    Tag0,  Tag1,  Tag2,  Tag3,  Tag4,  Tag5,  Tag6,  Tag7,  Tag8,  Tag9,
    Tag10, Tag11, Tag12, Tag13, Tag14, Tag15, Tag16, Tag17, Tag18, Tag19,
    Tag20, Tag21, Tag22, Tag23, Tag24, Tag25, Tag26, Tag27, Tag28, Tag29,
    Tag30,

  ) where

----------------------------------------------------------
-- import

import Data.List (genericSplitAt, genericTake, genericReplicate, genericLength)
import Data.Bits (Bits, testBit)
import Data.Proxy (Proxy(Proxy))
import qualified Data.Int as Z
import qualified Data.Ratio as Q
import qualified Data.Complex as C
import qualified Data.Fixed as F
import qualified Data.Word as W

----------------------------------------------------------
-- infix

infix 6 :@

infix 6 .+
infix 6 +.
infix 6 .-
infix 6 -.

infix 7 .*
infix 7 *.
infix 7 /.

infixr 8 **.

----------------------------------------------------------
-- Nion

-- | General Cayley-Dickson construction producing \"N-ions\". The
-- first parameter is a 'Tag' instance that determines the dimension,
-- which is 2 raised to 'tagVal'. The second parameter is the scalar
-- type.
data Nion n a = Scalar a | Nion n a :@ Nion n a

----------------------------------------------------------
-- basic operations

-- | Equivalent to @'coord' x 0@.
scalarPart :: Nion n a -> a
scalarPart (Scalar x) = x
scalarPart (x :@ _) = scalarPart x

-- | Equivalent to @'setCoord' x 0 0@.
purePart :: Num a => Nion n a -> Nion n a
purePart (Scalar _) = Scalar 0
purePart (x :@ y) = purePart x :@ y

-- | Dot product. For real scalars, @(1 \`dot\`)@ is equivalent to
-- @'scalarPart'@.
dot :: Conjugable a => Nion n a -> Nion n a -> a
-- (y * conj x + x * conj y) / 2
Scalar x `dot` Scalar y = scalarPart' $ y * conj x
x@(Scalar _) `dot` (y1 :@ _) = x `dot` y1
(x1 :@ _) `dot` y@(Scalar _) = x1 `dot` y
(x1 :@ x2) `dot` (y1 :@ y2) = (x1 `dot` y1) + (x2 `dot` y2)

-- | Cross product. The cross product of two pures yields a pure that
-- is orthogonal to both operands. For real scalars, @(1 \`cross\`)@
-- is equivalent to @'purePart'@.
cross :: Conjugable a => Nion n a -> Nion n a -> Nion n a
-- (y * conj x - x * conj y) / 2
x `cross` y = y * conj x -. x `dot` y

-- | Squared norm: the dot product of an element with itself.
sqnorm :: Conjugable a => Nion n a -> a
sqnorm x = x `dot` x

-- | Square root of @sqnorm@.
norm :: (Conjugable a, Floating a) => Nion n a -> a
norm = sqrt . sqnorm

-- | Promote a scalar, returning an element whose scalar part is the
-- argument and whose pure part is zero. The element behaves as if it
-- were padded with zeros, but no actual padding is done.
fromScalar :: a -> Nion n a
fromScalar = Scalar

----------------------------------------------------------
-- operations with scalars

leftScalarOp :: (Nion n a -> Nion n a -> Nion n a) -> a -> Nion n a -> Nion n a
leftScalarOp f x y = f (Scalar x) y

rightScalarOp :: (Nion n a -> Nion n a -> Nion n a) -> Nion n a -> a -> Nion n a
rightScalarOp f x y = f x (Scalar y)

-- | Equivalent to @'fromScalar' x + y@.
(.+) :: Conjugable a => a -> Nion n a -> Nion n a
(.+) = leftScalarOp (+)

-- | Equivalent to @'fromScalar' x - y@.
(.-) :: Conjugable a => a -> Nion n a -> Nion n a
(.-) = leftScalarOp (-)

-- | Equivalent to @'fromScalar' x * y@.
(.*) :: Conjugable a => a -> Nion n a -> Nion n a
(.*) = leftScalarOp (*)

-- | Equivalent to @x + 'fromScalar' y@.
(+.) :: Conjugable a => Nion n a -> a -> Nion n a
(+.) = rightScalarOp (+)

-- | Equivalent to @x - 'fromScalar' y@.
(-.) :: Conjugable a => Nion n a -> a -> Nion n a
(-.) = rightScalarOp (-)

-- | Equivalent to @x * 'fromScalar' y@.
(*.) :: Conjugable a => Nion n a -> a -> Nion n a
(*.) = rightScalarOp (*)

-- | Equivalent to @x / 'fromScalar' y@.
(/.) :: (Conjugable a, Fractional a) => Nion n a -> a -> Nion n a
(/.) = rightScalarOp (/)

-- | Raise to a scalar power.
(**.) :: (Tag n, Conjugable a, RealFloat a) => Nion n a -> a -> Nion n a
Scalar x **. y = Scalar $ x ** y
x **. y = exp (y .* log x)

----------------------------------------------------------
-- polar form and complex function application

realPolar :: RealFloat a => Nion n a -> a -> (a, a, Nion n a)
realPolar sqrtMinus1 r
  | r >= 0 = (r, 0, sqrtMinus1)
  | otherwise = (-r, pi, sqrtMinus1)

polarUsing :: (Conjugable a, RealFloat a) =>
              Nion n a -> Nion n a -> (a, a, Nion n a)
polarUsing sqrtMinus1 (Scalar r) = realPolar sqrtMinus1 r
polarUsing sqrtMinus1 x
  | sqnormp == 0 = realPolar sqrtMinus1 r
  | otherwise = (normx, acos (r / normx), u)
  where
    p = purePart x
    sqnormp = sqnorm p
    u = p /. sqrt sqnormp
    r = scalarPart x
    normx = sqrt $ sqnormp + r * conj r

polar' :: (Tag n, Conjugable a, RealFloat a) =>
          Proxy n -> Nion n a -> (a, a, Nion n a)
polar' n x
  | tagVal n == 0 = error "polar: no polar form in 1 dimension"
  | otherwise = polarUsing basisElement1 x

-- | Return @(s, t, u)@ such that (approximately)
--
--     @x == s .* 'exp' (t .* u)@
--
-- where @s@ and @t@ are scalars, @s >= 0@, and @u@ is a unit pure.
--
-- If @x@ has no pure part then @u@ is arbitrarily chosen to be the
-- first pure basis element.
polar :: (Tag n, Conjugable a, RealFloat a) => Nion n a -> (a, a, Nion n a)
polar = polar' Proxy

applyUsing' :: (Tag n, Conjugable a, RealFloat a) =>
               Proxy n ->
               Nion n a -> (a -> a) -> (C.Complex a -> C.Complex a) ->
               Nion n a -> Nion n a
applyUsing' n sqrtMinus1 fr f z
  | tagVal n == 0 = Scalar $ fr $ scalarPart z
  | otherwise = x .+ u *. y
  where (s, t, u) = polarUsing sqrtMinus1 z
        -- handle special cases for a little more accuracy
        x C.:+ y | t == 0 = f $ c s 0
                 | t == pi = f $ c (-s) 0
                 | otherwise = f $ c s 0 * exp (c t 0 * c 0 1)
                 where c = (C.:+)

applyUsing :: (Tag n, Conjugable a, RealFloat a) =>
              Nion n a -> (a -> a) -> (C.Complex a -> C.Complex a) ->
              Nion n a -> Nion n a
applyUsing = applyUsing' Proxy

----------------------------------------------------------
-- constants

fill' :: Tag n => Proxy n -> a -> Nion n a
fill' n s = f $ tagVal n where
  f 0 = Scalar s
  f k = f k' :@ f k' where k' = k - 1

fill :: Tag n => a -> Nion n a
fill = fill' Proxy

paddedZero :: (Tag n, Num a) => Nion n a
paddedZero = fill 0

validIndex :: (Tag n, Num b, Ord b) => Proxy n -> b -> Bool
validIndex n i = i >= 0 && i < 2 ^ tagVal n

basisElement' :: (Tag n, Conjugable a, Bits i, Integral i) =>
                 Proxy n -> i -> Nion n a
basisElement' _ 0 = Scalar 1
basisElement' n index
  | validIndex n index = setCoord paddedZero index 1
  | otherwise = error "basisElement: out of range"

-- | The nth basis element.
basisElement :: (Tag n, Conjugable a, Bits i, Integral i) => i -> Nion n a
basisElement = basisElement' Proxy

basisElement1 :: (Tag n, Conjugable a) => Nion n a
basisElement1 = basisElement (1 :: Integer)

----------------------------------------------------------
-- util

-- Proper Functor and Foldable instances would have to translate
-- top-level scalars to their equivalent representation with padded
-- zeros. The machinations needed for this indirection are rather
-- cumbersome relative to the benefits of having the instances, whose
-- use would seem uncommon.

smap :: (a -> a) -> Nion n a -> Nion n a
smap f (Scalar s) = Scalar $ f s
smap f (x :@ y) = smap f x :@ smap f y

sfoldr :: (a -> b -> b) -> b -> Nion n a -> b
sfoldr f acc (Scalar x) = f x acc
sfoldr f acc (x :@ y) = sfoldr f (sfoldr f acc y) x

----------------------------------------------------------
-- accessors

coords' :: (Tag n, Num a) => Proxy n -> Nion n a -> [a]
coords' n (Scalar x) = x : genericReplicate k 0 where
                         k = 2 ^ tagVal n - 1 :: Integer
coords' _ x = sfoldr (:) [] x

-- | List of coordinates for this element.
coords :: (Tag n, Num a) => Nion n a -> [a]
coords = coords' Proxy

coord' :: (Tag n, Num a, Integral b, Bits b) => Proxy n -> Nion n a -> b -> a
coord' _ (Scalar x) 0 = x
coord' n elt index
  | validIndex n index = case elt of
                           Scalar _ -> 0
                           _ -> f elt $ fromInteger $ tagVal n - 1
  | otherwise = error "coord: out of range"
  where
    f (Scalar x) _ = x
    f (x :@ y) k = case testBit index k of
                     False -> f x k'
                     True  -> f y k'
                   where k' = k - 1

-- | Get the nth coordinate.
coord :: (Tag n, Num a, Integral b, Bits b) => Nion n a -> b -> a
coord = coord' Proxy

setCoord' :: (Tag n, Conjugable a, Integral b, Bits b) =>
             Proxy n -> Nion n a -> b -> a -> Nion n a
setCoord' _ (Scalar _) 0 value = Scalar value
setCoord' n elt index value
  | validIndex n index = case elt of
                           Scalar x -> setCoord (x .+ paddedZero) index value
                           _ -> f elt $ fromInteger $ tagVal n - 1
  | otherwise = error "setCoord: out of range"
  where
    f (Scalar _) _ = Scalar value
    f (x :@ y) k = case testBit index k of
                     False -> f x k' :@ y
                     True  -> x :@ f y k'
                   where k' = k - 1

-- | Set the nth coordinate, returning a new element.
setCoord :: (Tag n, Conjugable a, Integral b, Bits b) =>
            Nion n a -> b -> a -> Nion n a
setCoord = setCoord' Proxy

----------------------------------------------------------
-- construction

fromList :: Integer -> [a] -> Nion n a
fromList _ (x:[]) = Scalar x
fromList k xs = fromList k' l :@ fromList k' r where
                  k' = k `div` 2
                  (l, r) = genericSplitAt k' xs

nion' :: (Tag n, Num a) => Proxy n -> [a] -> Nion n a
nion' n elems = fromList d $ taken ++ padding where
                  d = 2 ^ tagVal n
                  taken = genericTake d elems
                  padding = genericReplicate (d - genericLength taken) 0

-- | Construct an element from a list of coordinates. If the list is
-- too small then the remaining coordinates are padded with zeros. If
-- the list is too large then the extra values are ignored.
nion :: (Tag n, Num a) => [a] -> Nion n a
nion = nion' Proxy

----------------------------------------------------------
-- instances

instance (Tag n, Show a, Num a) => Show (Nion n a) where
  show x = "nion " ++ show (coords x)

instance (Conjugable a, Eq a) => Eq (Nion n a) where
  Scalar x == Scalar y = x == y
  x@(Scalar _) == y1 :@ y2 = x == y1 && y2 == 0
  x1 :@ x2 == y@(Scalar _) = x1 == y && x2 == 0
  x1 :@ x2 == y1 :@ y2 = x1 == y1 && x2 == y2

instance Conjugable a => Num (Nion n a) where
  Scalar x + Scalar y = Scalar $ x + y
  x@(Scalar _) + (y1 :@ y2) = (x + y1) :@ y2
  (x1 :@ x2) + y@(Scalar _) = (x1 + y) :@ x2
  (x1 :@ y1) + (x2 :@ y2) = (x1 + x2) :@ (y1 + y2)

  Scalar x - Scalar y = Scalar $ x - y
  x@(Scalar _) - (y1 :@ y2) = (x - y1) :@ negate y2
  (x1 :@ x2) - y@(Scalar _) = (x1 - y) :@ x2
  (x1 :@ y1) - (x2 :@ y2) = (x1 - x2) :@ (y1 - y2)

  Scalar x * Scalar y = Scalar $ x * y
  x@(Scalar _) * (y1 :@ y2) = (x * y1) :@ (x * y2)
  (x1 :@ x2) * y@(Scalar _) = (x1 * y) :@ (x2 * y)
  (x1 :@ x2) * (y1 :@ y2) = (x1 * y1 - conj y2 * x2) :@ (y2 * x1 + x2 * conj y1)

  negate = smap negate
  fromInteger = fromScalar . fromInteger
  abs = doNotUse
  signum = doNotUse

instance (Conjugable a, Fractional a) => Fractional (Nion n a) where
  Scalar x / Scalar y = Scalar $ x / y
  x@(_ :@ _) / Scalar y = smap (/ y) x
  x / y = (x * conj y) /. sqnorm y

  recip x = conj x /. sqnorm x
  fromRational = fromScalar . fromRational

-- | The first pure basis element is arbitrarily chosen as sqrt (-1).
instance (Tag n, Conjugable a, RealFloat a) => Floating (Nion n a) where
  pi    = Scalar pi
  exp   = applyUsing basisElement1 exp exp
  log   = applyUsing basisElement1 log log
  sqrt  = applyUsing basisElement1 sqrt sqrt
  sin   = applyUsing basisElement1 sin sin
  cos   = applyUsing basisElement1 cos cos
  tan   = applyUsing basisElement1 tan tan
  asin  = applyUsing basisElement1 asin asin
  acos  = applyUsing basisElement1 acos acos
  atan  = applyUsing basisElement1 atan atan
  sinh  = applyUsing basisElement1 sinh sinh
  cosh  = applyUsing basisElement1 cosh cosh
  tanh  = applyUsing basisElement1 tanh tanh
  asinh = applyUsing basisElement1 asinh asinh
  acosh = applyUsing basisElement1 acosh acosh
  atanh = applyUsing basisElement1 atanh atanh

----------------------------------------------------------
-- convenience types

-- | Complex numbers, the 2^1-dimensional construction.
type Complex a = Nion Tag1 a

-- | Quaternions, the 2^2-dimensional construction.
type Quaternion a = Nion Tag2 a

-- | Octonions, the 2^3-dimensional construction.
type Octonion a = Nion Tag3 a

-- | Sedenions, the 2^4-dimensional construction.
type Sedenion a = Nion Tag4 a

-- | Construct a complex number.
complex :: a -> a -> Complex a
complex x y = (:@) (Scalar x) (Scalar y)

-- | Construct a quaternion.
quaternion :: a -> a -> a -> a -> Quaternion a
quaternion w x y z = (:@) ((:@) (Scalar w) (Scalar x))
                          ((:@) (Scalar y) (Scalar z))

-- | Construct an octonion.
octonion :: a -> a -> a -> a ->
            a -> a -> a -> a -> Octonion a
octonion s t u v
         w x y z = (:@) ((:@) ((:@) (Scalar s) (Scalar t))
                              ((:@) (Scalar u) (Scalar v)))
                        ((:@) ((:@) (Scalar w) (Scalar x))
                              ((:@) (Scalar y) (Scalar z)))

-- | Construct a sedenion.
sedenion :: a -> a -> a -> a ->
            a -> a -> a -> a ->
            a -> a -> a -> a ->
            a -> a -> a -> a -> Sedenion a
sedenion k l m n
         o p q r
         s t u v
         w x y z = (:@) ((:@) ((:@) ((:@) (Scalar k) (Scalar l))
                                    ((:@) (Scalar m) (Scalar n)))
                              ((:@) ((:@) (Scalar o) (Scalar p))
                                    ((:@) (Scalar q) (Scalar r))))
                        ((:@) ((:@) ((:@) (Scalar s) (Scalar t))
                                    ((:@) (Scalar u) (Scalar v)))
                              ((:@) ((:@) (Scalar w) (Scalar x))
                                    ((:@) (Scalar y) (Scalar z))))

----------------------------------------------------------
-- Conjugable

-- | The /conjugate/ of an element is obtained by negating the pure
-- part and conjugating the scalar part. The conjugate of a real
-- number is the identity ('id'), which is the default implementation.
class Num a => Conjugable a where
  -- | Conjugate.
  conj :: a -> a
  conj = id

  -- | Equivalent to @(x + conj x) / 2@.
  scalarPart' :: a -> a
  scalarPart' = id

instance Conjugable a => Conjugable (Nion n a) where
  conj (Scalar x) = Scalar $ conj x
  conj (x :@ y) = conj x :@ negate y

  scalarPart' (Scalar x) = Scalar $ scalarPart' x
  scalarPart' (x :@ _) = scalarPart' x

instance (Conjugable a, RealFloat a) => Conjugable (C.Complex a) where
  conj = C.conjugate
  scalarPart' (x C.:+ _) = scalarPart' x C.:+ 0

instance Conjugable Int
instance Conjugable Integer
instance Conjugable Float
instance Conjugable Double
instance Conjugable Z.Int8
instance Conjugable Z.Int16
instance Conjugable Z.Int32
instance Conjugable Z.Int64
instance Conjugable W.Word8
instance Conjugable W.Word16
instance Conjugable W.Word32
instance Conjugable W.Word64
instance Integral a => Conjugable (Q.Ratio a)
instance F.HasResolution a => Conjugable (F.Fixed a)

-----------------------------------------------------------------------------
-- doNotUse

rant :: String
rant = unlines $
  ["",
   "The Num class is a bit messed up, having tied (+), (-), and (*) to abs",
   "and signum. Number systems that have no appropriate definition for abs",
   "or signum must either invent their own operators for addition,",
   "subtraction, and multiplication, else break the contract with Num by",
   "raising an error such as this one when someone uses abs or signum.",
   "",
   "For some time I resisted hijacking Num, but eventually the replacement",
   "operators became too cumbersome and, coupled with the lack of numeric",
   "promotion, significantly detracted from the usability of the package.",
   "So here we are. Good luck, and stay away from abs and signum, which",
   "officially have cooties."]

doNotUse :: a -> a
doNotUse _ = error rant

----------------------------------------------------------
-- Tag

-- | Tags serve to determine a type's dimension, which is 2 raised to
-- `tagVal`. Tag instances are included for convenience only, as you
-- may create your own tag.
class Tag n where
  tagVal :: Proxy n -> Integer

data     Tag0
data     Tag1
data     Tag2
data     Tag3
data     Tag4
data     Tag5
data     Tag6
data     Tag7
data     Tag8
data     Tag9
data    Tag10
data    Tag11
data    Tag12
data    Tag13
data    Tag14
data    Tag15
data    Tag16
data    Tag17
data    Tag18
data    Tag19
data    Tag20
data    Tag21
data    Tag22
data    Tag23
data    Tag24
data    Tag25
data    Tag26
data    Tag27
data    Tag28
data    Tag29
data    Tag30

instance Tag    Tag0 where tagVal _ =    0
instance Tag    Tag1 where tagVal _ =    1
instance Tag    Tag2 where tagVal _ =    2
instance Tag    Tag3 where tagVal _ =    3
instance Tag    Tag4 where tagVal _ =    4
instance Tag    Tag5 where tagVal _ =    5
instance Tag    Tag6 where tagVal _ =    6
instance Tag    Tag7 where tagVal _ =    7
instance Tag    Tag8 where tagVal _ =    8
instance Tag    Tag9 where tagVal _ =    9
instance Tag   Tag10 where tagVal _ =   10
instance Tag   Tag11 where tagVal _ =   11
instance Tag   Tag12 where tagVal _ =   12
instance Tag   Tag13 where tagVal _ =   13
instance Tag   Tag14 where tagVal _ =   14
instance Tag   Tag15 where tagVal _ =   15
instance Tag   Tag16 where tagVal _ =   16
instance Tag   Tag17 where tagVal _ =   17
instance Tag   Tag18 where tagVal _ =   18
instance Tag   Tag19 where tagVal _ =   19
instance Tag   Tag20 where tagVal _ =   20
instance Tag   Tag21 where tagVal _ =   21
instance Tag   Tag22 where tagVal _ =   22
instance Tag   Tag23 where tagVal _ =   23
instance Tag   Tag24 where tagVal _ =   24
instance Tag   Tag25 where tagVal _ =   25
instance Tag   Tag26 where tagVal _ =   26
instance Tag   Tag27 where tagVal _ =   27
instance Tag   Tag28 where tagVal _ =   28
instance Tag   Tag29 where tagVal _ =   29
instance Tag   Tag30 where tagVal _ =   30


-- |
-- Copyright  : (c) 2012 Daniël de Kok
-- License    : BSD3
--
-- Maintainer : Daniël de Kok <me@danieldk.eu>
-- Stability  : experimental
--
-- This module provides functionality to perform approximate randomization
-- tests (Noreen, 1989).

{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE DoAndIfThenElse #-}
{-# LANGUAGE FlexibleContexts #-}

module Statistics.Test.ApproxRand (
  -- * Description
  -- $description

  -- * Examples
  -- $examples

  -- * Data types
  TestOptions(..),
  TestResult(..),
  Significance(..),
  RandWithError,

  -- * Approximate randomization tests
  approxRandTest,
  approxRandStats,

  approxRandPairTest,
  approxRandPairStats,

  -- * Test statistics
  TestStatistic,
  differenceMean,
  meanDifference,
  varianceRatio
) where

import           Prelude hiding ((++))
import           Control.Monad (liftM, replicateM, when)
import           Control.Monad.Error (ErrorT)
import           Control.Monad.Error.Class (throwError)
import           Control.Monad.Mersenne.Random (R(..), Rand(..), getBool)
import           Control.Monad.ST (runST)
import           Control.Monad.Trans.Class (lift)
import           Data.Vector.Generic ((++))
import qualified Data.Vector.Generic as VG
import qualified Data.Vector.Generic.Mutable as GM
import           Data.Word (Word)
import           Statistics.Sample (variance)
import           Statistics.Test.Types (TestType(..))
import           Statistics.Types
import           System.Random.Mersenne.Pure64 (PureMT, randomInt, randomWord)

-- $description
--
-- Approximate randomization tests rely on a simple premise: given a test
-- statistic, if the null-hypothesis (the samples do not differ) is true,
-- we can randomly swap values between samples without an (extreme) impact
-- on the test statistic. Otherwise, the null-hypothesis must be rejected.
--
-- The test works by generating a given number of sample shuffles and computing
-- the test statistic for each shuffle. If /r/ is the number of shuffled
-- samples where the test statistic is at least as high as the test statistic
-- applied on the original samples; and /N/ the number of shuffles, then
-- the null-hypothesis is rejected iff /(r + 1):(N + 1) < p-value/ (for
-- one-sided tests).
--
-- Two kinds of test are supported:
--
-- * /Paired sample/ ('approxRandPairTest'): values from samples are shuffled
--   pair-wise. This requires the samples to have an equal length.
--
-- * /Unpaired sample/ ('approxRandTest'): values from samples are shuffled
--   among both samples. Consequently the i-th element of one sample does not
--   bear a relationship with the i-th element of the other sample. The
--   shuffled samples retain the sizes of the original samples.
--
-- Both tests can be performed as a one-tailed or two-tailed test.

-- $examples
-- Both unpaired and paired sample tests use the 'Rand' monad to obtain
-- random numbers. We can obtain a pseudo-random number generator that
-- is seeded using the system clock using the
-- 'System.Random.Mersenne.Pure64.newPureMT' function (please refer to
-- the documentation of 'System.Random.Mersenne.Pure64' for more
-- information):
--
-- > prng <- newPureMT
--
-- Suppose that we have the samples 's1' and 's2'. We could now perform
-- a Two-Tailed randomization test with 10,000 shuffles and the mean
-- difference as the test statistic, by running 'approxRandTest' in the 'Rand'
-- monad (at the /p = 0.01/ level):
--
-- > evalRandom (approxRandTest TwoTailed meanDifference 10000 0.01 s1 s2) prng
--
-- It is also possible to obtain the test statistics of the shuffled samples
-- directly (e.g. to inspect the distribution of test statistics) using the
-- 'approxRandStats'/'approxRandPiarStats' functions:
--
-- > evalRandom (approxRandStats meanDifference 10000 0.01 s1 s2) prng

-- | Computations with random numbers that can fail.
type RandWithError a = ErrorT String Rand a

-- | Options for randomization tests
--
data TestOptions = TestOptions {
    toTestType      :: TestType,      -- ^ Type of test ('OneTailed' or 'TwoTailed')
    toTestStatistic :: TestStatistic, -- ^ Test statistic
    toIterations    :: Int,           -- ^ Number of shuffled samples to create
    toPValue        :: Double         -- ^ he p-value at which to test (e.g. 0.05)
}

-- |
-- The result of hypothesis testing.
data TestResult = TestResult {
    trSignificance     :: Significance, -- ^ Significance
    trStat             :: Double,       -- ^ Test statistic for the samples
    trRandomizedStats  :: Sample        -- ^ Test statistics for the
                                        --   randomized samples
  } deriving (Eq, Ord, Show)

-- |
-- Significance.
data Significance =
    Significant    Double -- ^ The null hypothesis should be rejected
  | NotSignificant Double -- ^ Data is compatible with the null hypothesis
  deriving (Eq, Ord, Show)

-- |
-- Apply a pair-wise approximate randomization test.
--
-- In pair-wise approximate randomization tests the data points at a given
-- index are swapped between samples with a probability of 0.5. Since
-- swapping is pairwise, the samples should have the same length.
approxRandPairTest ::
     TestOptions              -- ^ Options for the test
  -> Sample                   -- ^ First sample
  -> Sample                   -- ^ Second sample
  -> RandWithError TestResult -- ^ The test result
approxRandPairTest (TestOptions testType stat n pTest) s1 s2 = do
  stats <- approxRandPairStats stat n s1 s2
  let tOrig = stat s1 s2
  let sig = significance testType pTest n $ countExtremes tOrig stats
  return $ TestResult sig tOrig stats

-- |
-- Apply an approximate randomization test.
--
-- In approximate randomization tests, the values of two samples are
-- shuffled among those samples. A test statistic is calculated for
-- the original samples and the shuffled samples, to detect whether the
-- difference of the samples is extreme or not.
approxRandTest ::
     TestOptions     -- ^ Options for the test
  -> Sample          -- ^ First sample
  -> Sample          -- ^ Second sample
  -> Rand TestResult -- ^ The test result
approxRandTest (TestOptions testType stat n pTest) s1 s2 = do
  stats <- approxRandStats stat n s1 s2
  let tOrig = stat s1 s2
  let sig = significance testType pTest n $ countExtremes tOrig stats
  return $ TestResult sig tOrig stats

-- | Determine the significance.
significance ::
     TestType     -- ^ Type of test ('OneTailed' or 'TwoTailed')
  -> Double       -- ^ The p-value at which to test (e.g. 0.05)
  -> Int          -- ^ Number of sample shuffles
  -> (Int, Int)   -- ^ Extreme statistic counts
  -> Significance -- ^ The test result
significance TwoTailed pTest n =
  significant (pTest / 2) . pValue n . uncurry min
significance OneTailed pTest n =
  significant pTest . pValue n . snd

-- | Wrap a p-value in a 'TestResult'.
significant ::
     Double       -- ^ The p-value at which to test
  -> Double       -- ^ The p-value
  -> Significance -- ^ The test result
significant pTail p =
  if p < pTail then Significant p else NotSignificant p

-- | Calculate a p-value
pValue ::
     Int    -- ^ Number of extreme outcomes
  -> Int    -- ^ Number of shuffles
  -> Double -- ^ The p-value
pValue n r = (fromIntegral r + 1) / (fromIntegral n + 1)

-- |
-- Count extreme test statistic values. If the test statistic value of the
-- original sample is in the right tail, we want to count values equal to
-- or larger than that value. If the value is in the left tail, we want to
-- count value smaller than or equal to that value. Since we do not know
-- the tail (yet), we count both.
--
-- Note: we can determine the tail by (1) averaging the test statistics of
-- the randomized samples, or (2) taking the smaller of the two counts.
countExtremes ::
     Double     -- ^ Test statistic value of the original samples
  -> Sample     -- ^ Test statistic values of the randomized samples.
  -> (Int, Int) -- ^ Count of left- and right-tail extremes.
countExtremes tOrig =
  VG.foldl' count (0, 0)
  where
    count (left, right) tPerm =
      let !newLeft = if tPerm <= tOrig then succ left else left in
      let !newRight = if tPerm >= tOrig then succ right else right in
        (newLeft, newRight)

-- |
-- Generate a given number of pairwise shuffled samples, and calculate
-- the test statistic for each shuffle.
--
-- Since the data points at a given index are swapped (with a probability of
-- 0.5), the samples should have the same length.
approxRandPairStats ::
     TestStatistic          -- ^ Test statistic
  -> Int                    -- ^ Number of shuffled samples to create
  -> Sample                 -- ^ First sample
  -> Sample                 -- ^ Second sample
  -> RandWithError Sample   -- ^ The statistics of the shuffles
approxRandPairStats stat n s1 s2 = do
  when (VG.length s1 /= VG.length s2) $
    throwError "Cannot calculate pairwise statistic: samples have different sizes"
  lift $ liftM VG.fromList $ replicateM n $
    uncurry stat `liftM` shuffleVectorsPairwise s1 s2

-- |
-- Generate a given number of shuffled samples, and calculate the test
-- statistic for each shuffle.
--
-- This function does not require the samples to have an equal length.
approxRandStats ::
     TestStatistic -- ^ Test statistic
  -> Int           -- ^ Number of shuffled samples to create
  -> Sample        -- ^ First sample
  -> Sample        -- ^ Second sample
  -> Rand Sample   -- ^ The statistics of the shuffles
approxRandStats stat n s1 s2 =
  liftM VG.fromList $ replicateM n $ uncurry stat `liftM` shuffleVectors s1 s2

-- | Pair-wise shuffle of two vectors.
shuffleVectorsPairwise :: (VG.Vector v a, VG.Vector v Bool) =>
  v a -> v a -> Rand (v a, v a)
shuffleVectorsPairwise vec1 vec2 = do
  randomVec <- randomVector (VG.length vec1)
  let pv1 = VG.zipWith3 permute vec1 vec2 randomVec
  let pv2 = VG.zipWith3 permute vec2 vec1 randomVec
  return (pv1, pv2)
  where
    permute val1 val2 coin =
      if coin then val1 else val2

randomVector :: VG.Vector v Bool => Int -> Rand (v Bool)
randomVector len =
  VG.replicateM len getBool

-- Shuffle values amongst two vectors, keeping the original vector lengths.
shuffleVectors :: VG.Vector v a => v a -> v a -> Rand (v a, v a)
shuffleVectors v1 v2 = do
  shuffledVectors <- shuffleVector $ v1 ++ v2
  return (VG.slice 0 (VG.length v1) shuffledVectors,
    VG.slice (VG.length v1) (VG.length v2) shuffledVectors)

-- Fisher-Yates shuffle in the Rand monad
shuffleVector :: VG.Vector v a => v a -> Rand (v a)
shuffleVector v =
  Rand $ \s -> case shuffleVector' s v of (sv, s') -> R sv s'

-- Fisher-Yates shuffle
shuffleVector' :: VG.Vector v a => PureMT -> v a -> (v a, PureMT)
shuffleVector' gen v = runST $ do
  let maxIdx = VG.length v - 1
  vm   <- VG.thaw v
  gen' <- swaps vm 0 maxIdx gen
  vmf  <- VG.unsafeFreeze vm
  return (vmf, gen')
  where
    swaps vm idx maxIdx gen'
      | idx < maxIdx = do
          let (newIdx, gen'') = randomIntR gen' (idx, maxIdx)
          GM.unsafeSwap vm idx newIdx
          swaps vm (idx + 1) maxIdx gen''
      | otherwise = return gen'

-- |
-- A test stastic calculates the difference between two samples. See
-- 'meanDifference' and 'varianceRatio' for examples.
type TestStatistic = Sample -> Sample -> Double

-- |
-- Calculates the difference mean of two samples (/mean(s1 - s2)/). When the
-- two samples do not have an equal length, the trailing elements of the
-- longer vector are ignored.
differenceMean :: TestStatistic
differenceMean v1 v2 =
  VG.sum (subVector v1 v2) / fromIntegral (VG.length v1)

-- | Calculates the mean difference of two samples (/mean(s1) - mean(s2)/).
meanDifference :: TestStatistic
meanDifference s1 s2 =
  mean s1 - mean s2

-- | Calculate the mean of a sample.
mean :: Sample -> Double
mean = do
  t <- VG.sum
  l <- VG.length
  return $ t / fromIntegral l

-- | Calculate the ratio of sample variances (/var(s1) : var(s2)/).
varianceRatio :: TestStatistic
varianceRatio v1 v2 =
  variance v1 / variance v2

-- | Subtract two vectors.
subVector :: (VG.Vector v n, Num n) => v n -> v n -> v n
subVector = VG.zipWith (-)

subIIW :: Int -> Int -> Word
subIIW a b = fromIntegral a - fromIntegral b
{-# INLINE subIIW #-}

addIWI :: Int -> Word -> Int
addIWI a b = a + fromIntegral b
{-# INLINE addIWI #-}

-- | Generate Int numbers within a range
randomIntR :: PureMT -> (Int, Int) -> (Int, PureMT)
randomIntR gen (a, b)
  | n == 0    = randomInt gen
  | otherwise = loop gen
  where
    (a', b') = if a < b then (a, b) else (b, a)
    -- Number of different Ints that should be generated
    n = 1 + subIIW b' a'
    -- The total range of Word can hold x complete n ranges
    x = maxBound `div` n
    -- Pick from a range the is dividable by n without remainders
    s = x * n
    loop gen'
      | r >= s    = loop gen'' -- r is outside the range, discard it...
      | otherwise = (addIWI a' (r `div` x), gen'')
      where
        (!r, !gen'') = randomWord gen'
{-# INLINE randomIntR #-}