{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE PatternSynonyms #-}
{-# LANGUAGE ViewPatterns #-}
{-# LANGUAGE OverloadedStrings #-}

{-|
Module      : Balanced
Description : Amplitude-balanced path sums
Copyright   : (c) Matthew Amy, 2020
Maintainer  : matt.e.amy@gmail.com
Stability   : experimental
Portability : portable
-}

module Feynman.Algebra.Pathsum.Balanced where

import Data.List
import qualified Data.Set as Set
import Data.Ratio
import Data.Semigroup
import Control.Monad (mzero, msum)
import Data.Maybe (maybeToList)
import Data.Complex (Complex, mkPolar)
import Data.Bits (shiftL)
import Data.Map (Map, (!))
import qualified Data.Map as Map
import Data.String (IsString(..))
import Data.Tuple (swap)

import qualified Feynman.Util.Unicode as U
import Feynman.Algebra.Base
import Feynman.Algebra.Polynomial (degree)
import Feynman.Algebra.Polynomial.Multilinear

{-----------------------------------
 Variables
 -----------------------------------}

-- | Variables are either input variables or path variables. The distinction
--   is due to the binding structure of our pathsum representation, and moreover
--   improves readability
data Var = IVar !Int | PVar !Int | FVar !String deriving (Eq, Ord)

instance Show Var where
  show (IVar i) = U.sub "x" $ fromIntegral i
  show (PVar i) = U.sub "y" $ fromIntegral i
  show (FVar x) = x

instance IsString Var where
  fromString = FVar

-- | Convenience function for the string representation of the 'i'th input variable
ivar :: Int -> String
ivar = show . IVar

-- | Convenience function for the string representation of the 'i'th path variable
pvar :: Int -> String
pvar = show . PVar

-- | Construct an integer shift for input variables
shiftI :: Int -> (Var -> Var)
shiftI i = shiftAll i 0

-- | Construct an integer shift for path variables
shiftP :: Int -> (Var -> Var)
shiftP j = shiftAll 0 j

-- | General shift. Constructs a substitution from shift values for I and P
shiftAll :: Int -> Int -> (Var -> Var)
shiftAll i j = go
  where go (IVar i') = IVar (i + i')
        go (PVar j') = PVar (j + j')

-- | Check if a variable is an input
isI :: Var -> Bool
isI (IVar _) = True
isI _        = False

-- | Check if a variable is a path variable
isP :: Var -> Bool
isP (PVar _) = True
isP _        = False

-- | Check if a variable is a free variable
isF :: Var -> Bool
isF (FVar _) = True
isF _        = False

-- | Get the string of a free variable
unF :: Var -> String
unF (FVar s) = s
unF _        = error "Not a free variable"

{-----------------------------------
 Path sums
 -----------------------------------}

-- | Path sums of the form
--   \(\frac{1}{\sqrt{2}^k}\sum_{y\in\mathbb{Z}_2^m}e^{i\pi P(x, y)}|f(x, y)\rangle\)
data Pathsum g = Pathsum {
  sde       :: !Int,
  inDeg     :: !Int,
  outDeg    :: !Int,
  pathVars  :: !Int,
  phasePoly :: !(PseudoBoolean Var g),
  outVals   :: ![SBool Var]
  } deriving (Eq)

instance (Show g, Eq g, Periodic g, Real g) => Show (Pathsum g) where
  show sop = inputstr ++ scalarstr ++ sumstr ++ amplitudestr ++ statestr
    where inputstr = case inDeg sop of
            0 -> ""
            1 -> U.ket (ivar 0) ++ " " ++ U.mapsto ++ " "
            2 -> U.ket (ivar 0) ++ U.ket (ivar 1) ++ " " ++ U.mapsto ++ " "
            j -> U.ket (ivar 0) ++ U.dots ++ U.ket (ivar (j-1)) ++ " " ++ U.mapsto ++ " "
          scalarstr = case compare (sde sop) 0 of
            LT -> U.sup ("(" ++ U.rt2 ++ ")") (fromIntegral . abs $ sde sop)
            EQ -> if (inDeg sop == 0 && outDeg sop == 0 && phasePoly sop == 0) then "1" else ""
            GT -> U.sup ("1/(" ++ U.rt2 ++ ")") (fromIntegral $ sde sop)
          sumstr = case pathVars sop of
            0 -> ""
            1 -> U.sum ++ "[" ++ pvar 0 ++ "]"
            2 -> U.sum ++ "[" ++ pvar 0 ++ pvar 1 ++ "]"
            j -> U.sum ++ "[" ++ pvar 0 ++ U.dots ++ pvar (j-1) ++ "]"
          amplitudestr = case order (phasePoly sop) of
            0 -> U.e ++ "^" ++ U.i ++ U.pi ++ "{" ++ show (phasePoly sop) ++ "}"
            1 -> ""
            2 -> "(-1)^{" ++ show (makeIntegral 1 $ phasePoly sop) ++ "}"
            4 -> U.i ++ "^{" ++ show (makeIntegral 2 $ phasePoly sop) ++ "}"
            8 -> U.omega ++ "^{" ++ show (makeIntegral 4 $ phasePoly sop) ++ "}"
            j -> U.sub U.zeta j ++ "^{" ++ show (makeIntegral j $ phasePoly sop) ++ "}"
          statestr = concatMap (U.ket . show) $ outVals sop

-- | Convenience function for pretty printing
makeIntegral :: Real g => Integer -> PseudoBoolean v g -> PseudoBoolean v Integer
makeIntegral i = cast (\a -> numerator $ toRational a * toRational i)

-- | Retrieve the internal path variables
internalPaths :: Pathsum g -> [Var]
internalPaths sop = [PVar i | i <- [0..pathVars sop - 1]] \\ outVars
  where outVars = Set.toList . Set.unions . map vars $ outVals sop

-- | Retrieve the free variables
freeVars :: Pathsum g -> [String]
freeVars sop = map unF . Set.toList . Set.filter isF . foldr (Set.union) Set.empty $ xs
  where xs = (vars $ phasePoly sop):(map vars $ outVals sop)

-- | Checks if the path sum is (trivially) the identity
isTrivial :: (Eq g, Num g) => Pathsum g -> Bool
isTrivial sop = sop == identity (inDeg sop)

-- | (To be deprecated) Drops constant term from the phase polynomial
dropGlobalPhase :: (Eq g, Num g) => Pathsum g -> Pathsum g
dropGlobalPhase sop = sop { phasePoly = dropConstant $ phasePoly sop }

-- | (To be deprecated) Drops the phase polynomial
dropPhase :: (Eq g, Num g) => Pathsum g -> Pathsum g
dropPhase sop = sop { phasePoly = 0 }

-- | (To be deprecated) Drops normalization
dropAmplitude :: Pathsum g -> Pathsum g
dropAmplitude sop = sop { sde = 0 }

{----------------------------
 Constructors
 ----------------------------}

-- | Construct an 'n'-qubit identity operator
identity :: (Eq g, Num g) => Int -> Pathsum g
identity n = Pathsum 0 n n 0 0 [ofVar (IVar i) | i <- [0..n-1]]

-- | Construct a (symbolic) state
ket :: (Eq g, Num g) => [SBool String] -> Pathsum g
ket xs = Pathsum 0 0 (fromIntegral $ length xs) 0 0 $ map (rename FVar) xs

-- | Initialize a fresh ancilla
initialize :: (Eq g, Num g) => FF2 -> Pathsum g
initialize b = ket [constant b]

{-# INLINE initialize #-}

-- | Construct a uniform superposition of classical states of a given form
superposition :: (Ord v, Eq g, Num g) => [SBool v] -> Pathsum g
superposition xs = Pathsum k 0 n k 0 $ map (rename sub) xs
  where n   = length xs
        k   = Set.size fv
        fv  = Set.unions $ map vars xs
        sub = ((Map.fromList [(v, PVar i) | (v, i) <- zip (Set.toList fv) [0..]])!)

-- | Construct a classical transformation from the free variables of 'xs' to 'xs'
--   Effectively binds the free variables in 'state xs'
compute :: (Ord v, Eq g, Num g) => [SBool v] -> Pathsum g
compute xs = Pathsum 0 (Set.size fv) (length xs) 0 0 $ map (rename sub) xs
  where fv  = Set.unions $ map vars xs
        sub = ((Map.fromList [(v, IVar i) | (v, i) <- zip (Set.toList fv) [0..]])!)

-- | Breaks the connection between inputs and outputs by mapping
--   any input basis state to the maximally mixed state. Non-unitary
disconnect :: (Eq g, Num g) => Int -> Pathsum g
disconnect n = Pathsum n n n n 0 [ofVar $ PVar i | i <- [0..n-1]]

-- | Construct a permutation
permutation :: (Eq g, Num g) => [Int] -> Pathsum g
permutation xs
  | all (\i -> i >= 0 && i <= n) xs = Pathsum 0 n n 0 0 (map (ofVar . IVar) xs)
  | otherwise = error "permutation: Input not a permutation"
  where n = length xs

-- | Construct an n-ary unit
etaN :: (Eq g, Num g) => Int -> Pathsum g
etaN n = Pathsum 0 0 (2*n) n 0 $ xs ++ xs
  where xs = map (ofVar . PVar) [0..n-1]

{----------------------------
 Dual constructors
 ----------------------------}

-- | Construct a (symbolic) state destructor
bra :: (Eq g, Abelian g) => [SBool String] -> Pathsum g
bra xs = Pathsum (2*m) m 0 m (lift $ p) []
  where m         = fromIntegral $ length xs
        p         = foldr (+) 0 . map go $ zip [0..] xs
        go (i, v) = ofVar (PVar i) * (ofVar (IVar i) + (rename fromString v))

-- | Alternate state destructor with fewer paths but a more complicated polynomial
unstateAlt :: (Eq g, Abelian g) => [SBool String] -> Pathsum g
unstateAlt xs = Pathsum 2 (fromIntegral $ length xs) 0 1 (lift $ y*(1 + p)) []
  where y = ofVar (PVar 0)
        p = foldr (*) 1 . map valF $ zip xs [0..]
        valF (val, i) = 1 + (rename fromString val) + ofVar (IVar i)

-- | Dagger of initialize -- i.e. unnormalized post-selection
postselect :: (Eq g, Abelian g) => FF2 -> Pathsum g
postselect b = bra [constant b]

{-# INLINE postselect #-}

-- | Select on a classical state of a given form
unsuper :: (Ord v, Eq g, Abelian g) => [SBool v] -> Pathsum g
unsuper xs = Pathsum (2*m + n) m 0 (m+n) poly []
  where m    = length xs
        n    = Set.size fv
        fv   = Set.unions $ map vars xs
        sub  = ((Map.fromList [(v, PVar (m + i)) | (v, i) <- zip (Set.toList fv) [0..]])!)
        poly = foldr (+) zero $ map constructTerm [0..m-1]
        constructTerm i = lift $ ofVar (PVar i) * (ofVar (IVar i) + rename sub (xs!!i))

-- | Invert a classical transformation
uncompute :: (Ord v, Eq g, Abelian g) => [SBool v] -> Pathsum g
uncompute xs = Pathsum (2*m) m n (m+n) poly [ofVar (PVar $ m + i) | i <- [0..n-1]]
  where m    = length xs
        n    = Set.size fv
        fv   = Set.unions $ map vars xs
        sub  = ((Map.fromList [(v, PVar (m + i)) | (v, i) <- zip (Set.toList fv) [0..]])!)
        poly = foldr (+) zero $ map constructTerm [0..m-1]
        constructTerm i = lift $ ofVar (PVar i) * (ofVar (IVar i) + rename sub (xs!!i))

-- | Construct an inverse permutation
unpermutation :: (Eq g, Num g) => [Int] -> Pathsum g
unpermutation = permutation . snd . unzip . sort . map swap . zip [0..]

-- | Construct an n-ary co-unit
epsilonN :: (Eq g, Abelian g) => Int -> Pathsum g
epsilonN n = Pathsum (2*n) (2*n) 0 n (lift poly) []
  where poly = sum $ map f [0..n-1]
        f i  = ofVar (PVar i) * (ofVar (IVar i) + ofVar (IVar $ n+i))

{----------------------------
 Constants & gates
 ----------------------------}

-- | \(\sqrt{2}\)
root2 :: (Eq g, Abelian g, Dyadic g) => Pathsum g
root2 = Pathsum 0 0 0 1 ((-constant (half * half)) + scale half (lift $ ofVar (PVar 0))) []

-- | \(1/\sqrt{2}\)
roothalf :: (Eq g, Abelian g, Dyadic g) => Pathsum g
roothalf = Pathsum 1 0 0 0 0 []

-- | \(i\)
iunit :: (Eq g, Abelian g, Dyadic g) => Pathsum g
iunit = Pathsum 0 0 0 0 (constant half) []

-- | \(e^{i\pi/4}\)
omega :: (Eq g, Abelian g, Dyadic g) => Pathsum g
omega = Pathsum 0 0 0 0 (constant (half * half)) []

-- | A fresh, 0-valued ancilla
fresh :: (Eq g, Num g) => Pathsum g
fresh = Pathsum 0 0 1 0 0 [0]

-- | The dagger of fresh
unfresh :: (Eq g, Abelian g) => Pathsum g
unfresh = dagger fresh

-- | The unit, \(\eta\)
eta :: (Eq g, Num g) => Pathsum g
eta = Pathsum 0 0 2 1 0 [ofVar (PVar 0), ofVar (PVar 0)]

-- | The co-unit, \(\epsilon\)
epsilon :: (Eq g, Abelian g) => Pathsum g
epsilon = Pathsum 2 2 0 1 p []
  where p = lift $ ofVar (PVar 0) * (ofVar (IVar 0) + ofVar (IVar 1))

-- | X gate
xgate :: (Eq g, Num g) => Pathsum g
xgate = Pathsum 0 1 1 0 0 [1 + ofVar (IVar 0)]

-- | Z gate
zgate :: (Eq g, Abelian g) => Pathsum g
zgate = Pathsum 0 1 1 0 p [ofVar (IVar 0)]
  where p = lift $ ofVar (IVar 0)

-- | Y gate
ygate :: (Eq g, Abelian g, Dyadic g) => Pathsum g
ygate = Pathsum 0 1 1 0 p [1 + ofVar (IVar 0)]
  where p = constant half + (lift $ ofVar (IVar 0))

-- | S gate
sgate :: (Eq g, Abelian g, Dyadic g) => Pathsum g
sgate = Pathsum 0 1 1 0 p [ofVar (IVar 0)]
  where p = scale half (lift $ ofVar (IVar 0))

-- | S* gate
sdggate :: (Eq g, Abelian g, Dyadic g) => Pathsum g
sdggate = Pathsum 0 1 1 0 p [ofVar (IVar 0)]
  where p = scale (-half) (lift $ ofVar (IVar 0))

-- | T gate
tgate :: (Eq g, Abelian g, Dyadic g) => Pathsum g
tgate = Pathsum 0 1 1 0 p [ofVar (IVar 0)]
  where p = scale (half*half) (lift $ ofVar (IVar 0))

-- | T* gate
tdggate :: (Eq g, Abelian g, Dyadic g) => Pathsum g
tdggate = Pathsum 0 1 1 0 p [ofVar (IVar 0)]
  where p = scale (-half*half) (lift $ ofVar (IVar 0))

-- | R_k gate
rkgate :: (Eq g, Abelian g, Dyadic g) => Int -> Pathsum g
rkgate k = Pathsum 0 1 1 0 p [ofVar (IVar 0)]
  where p = scale (fromDyadic $ dyadic 1 k) (lift $ ofVar (IVar 0))

-- | R_z gate
rzgate :: (Eq g, Abelian g, Dyadic g) => DyadicRational -> Pathsum g
rzgate theta = Pathsum 0 1 1 0 p [ofVar (IVar 0)]
  where p = scale (fromDyadic theta) (lift $ ofVar (IVar 0))

-- | H gate
hgate :: (Eq g, Abelian g, Dyadic g) => Pathsum g
hgate = Pathsum 1 1 1 1 p [ofVar (PVar 0)]
  where p = lift $ (ofVar $ IVar 0) * (ofVar $ PVar 0)

-- | CH gate
chgate :: (Eq g, Abelian g, Dyadic g) => Pathsum g
chgate = Pathsum 1 2 2 1 p [x1, x2 + x1*x2 + x1*y]
  where p = lift $ x1 * x2 * y
        x1 = ofVar $ IVar 0
        x2 = ofVar $ IVar 1
        y = ofVar $ PVar 0

-- | CNOT gate
cxgate :: (Eq g, Num g) => Pathsum g
cxgate = Pathsum 0 2 2 0 0 [x0, x0+x1]
  where x0 = ofVar $ IVar 0
        x1 = ofVar $ IVar 1

-- | Toffoli gate
ccxgate :: (Eq g, Num g) => Pathsum g
ccxgate = Pathsum 0 3 3 0 0 [x0, x1, x2 + x0*x1]
  where x0 = ofVar $ IVar 0
        x1 = ofVar $ IVar 1
        x2 = ofVar $ IVar 2

-- | k-control Toffoli gate
mctgate :: (Eq g, Num g) => Int -> Pathsum g
mctgate k = Pathsum 0 (k+1) (k+1) 0 0 (controls ++ [t + foldr (*) 1 controls])
  where controls = [ofVar (IVar i) | i <- [0..k-1]]
        t        = ofVar $ IVar k

-- | SWAP gate
swapgate :: (Eq g, Num g) => Pathsum g
swapgate = Pathsum 0 2 2 0 0 [x1, x0]
  where x0 = ofVar $ IVar 0
        x1 = ofVar $ IVar 1

{----------------------------
 Channels
 ----------------------------}

-- | Choi matrix of computational basis measurement
measureChoi :: (Eq g, Abelian g) => Pathsum g
measureChoi = Pathsum 2 2 2 1 (lift $ y * (x0 + x1)) [x0, x1]
  where x0 = ofVar $ IVar 0
        x1 = ofVar $ IVar 1
        y  = ofVar $ PVar 0

-- | CPM operator of computational basis measurement
measure :: (Eq g, Abelian g) => Pathsum g
measure = unChoi measureChoi

{----------------------------
 Bind, unbind, and subst
 ----------------------------}

-- | Bind some collection of free variables in a path sum
bind :: (Foldable f, Eq g, Abelian g) => f String -> Pathsum g -> Pathsum g
bind = flip (foldr go)
  where go x sop =
          let v = IVar $ inDeg sop in
            sop { inDeg = (inDeg sop) + 1,
                  phasePoly = subst (FVar x) (ofVar v) (phasePoly sop),
                  outVals = map (subst (FVar x) (ofVar v)) (outVals sop) }

-- | Close a path sum by binding all free variables
close :: (Eq g, Abelian g) => Pathsum g -> Pathsum g
close sop = bind (freeVars sop) sop

-- | Unbind (instantiate) some collection of inputs
unbind :: (Foldable f, Eq g, Abelian g) => f Int -> Pathsum g -> Pathsum g
unbind xs (Pathsum a b c d e f) = Pathsum a (b - length xs) c d e' f' where
  e'  = substMany sub e
  f'  = map (substMany sub) f
  sub = \v -> Map.findWithDefault (ofVar v) v tmp
  tmp = snd $ foldr buildMap (0, Map.empty) [0..b-1]
  buildMap i (j, acc) = case i `elem` xs of
    True  -> (j, Map.insert (IVar i) (ofVar . FVar $ "#" ++ show (IVar i)) acc)
    False -> (j+1, Map.insert (IVar i) (ofVar . IVar $ j) acc)

-- | Open a path sum by instantiating all inputs
open :: (Eq g, Abelian g) => Pathsum g -> Pathsum g
open sop = unbind [0..(inDeg sop) - 1] sop

-- | Substitute a monomial with a symbolic Boolean expression throughout
--
--   This is generally not a very safe thing to do. Convenience for certain
--   local transformations
substitute :: (Eq g, Abelian g) => [Var] -> SBool Var -> Pathsum g -> Pathsum g
substitute xs p (Pathsum a b c d e f) = Pathsum a b c d e' f' where
  e' = substMonomial xs p e
  f' = map (substMonomial xs p) f

{----------------------------
 Operators
 ----------------------------}

-- | Return the dual of a path sum
dualize :: (Eq g, Abelian g) => Pathsum g -> Pathsum g
dualize sop@(Pathsum a b c d e f) = inSOP .> midSOP .> outSOP
  where inSOP  = tensor (identity c) (etaN b)
        midSOP = tensor (tensor (identity c) sop) (identity b)
        outSOP = tensor (epsilonN c) (identity b)

-- | Return the (column) vectorized path sum. By convention we place the inputs
--   first (i.e. f : A -> B becomes vectorize f : A* \otimes B)
vectorize :: (Eq g, Abelian g) => Pathsum g -> Pathsum g
vectorize sop@(Pathsum a b c d e f) = etaN b .> tensor (identity b) sop

-- | Return the (row) vectorized path sum. By convention we place the outputs
--   first (i.e. f : A -> B becomes vectorize f : A* \otimes B)
covectorize :: (Eq g, Abelian g) => Pathsum g -> Pathsum g
covectorize sop@(Pathsum a b c d e f) = tensor (identity c) sop .> epsilonN c

-- | Take the dagger of a path sum
dagger :: (Eq g, Abelian g) => Pathsum g -> Pathsum g
dagger (Pathsum a b c d e f) = dualize $ Pathsum a b c d (-e) f

-- | Take the conjugate (c.f., lower star) of a path sum
conjugate :: (Eq g, Abelian g) => Pathsum g -> Pathsum g
conjugate (Pathsum a b c d e f) = Pathsum a b c d (-e) f

-- | Trace a square morphism. Throws an error if the input and outputs are not
--   the same size
trace :: (Eq g, Abelian g) => Pathsum g -> Pathsum g
trace sop@(Pathsum a b c d e f)
  | b /= c = error "Can't trace a non-square operator"
  | otherwise = etaN b .> tensor (identity b) sop .> epsilonN b

-- | Trace out the first qubit. Throws an error if the input is a vector
ptrace :: (Eq g, Abelian g) => Pathsum g -> Pathsum g
ptrace sop@(Pathsum a b c d e f)
  | b < 1 || c < 1 = error "Can't partial trace a vector"
  | otherwise = tensor eta (identity $ b-1) .>
                tensor (identity 1) sop .>
                tensor epsilon (identity $ c-1)

-- | Turn a pure state into a density matrix. Throws an error if the input is not
--   a column vector
densify :: (Eq g, Abelian g) => Pathsum g -> Pathsum g
densify sop@(Pathsum a b c d e f)
  | b /= 0 = error "Can't densify an operator "
  | otherwise = dagger sop .> sop

-- | Turn a (single-qubit) Choi matrix into a CPM-style linear operator
unChoi :: (Eq g, Abelian g) => Pathsum g -> Pathsum g
unChoi sop@(Pathsum a b c d e f)
  | b /= 2 || c /= 2 = error "Only single-qubit channels currently supported"
  | otherwise        = tensor eta swapgate .>
                       tensor (identity 1) (tensor sop $ identity 1) .>
                       tensor swapgate epsilon

-- | Turn a unitary operator into a channel. That is, f becomes f_* \otimes f
channelize :: (Eq g, Abelian g) => Pathsum g -> Pathsum g
channelize sop = tensor (conjugate sop) sop

-- | Construct a controlled path sum
controlled :: (Eq g, Abelian g) => Pathsum g -> Pathsum g
controlled sop@(Pathsum a b c d e f) = Pathsum a (b+1) (c+1) d e' f' where
  shift   = shiftI 1
  x       = ofVar $ IVar 0
  e'      = (lift x)*(renameMonotonic shift e)
  f'      = [lift x] ++ (map g . zip [1..] . map (renameMonotonic shift) $ f)
  g (i,y) = (ofVar $ IVar i) + x*((ofVar $ IVar i) + y)

-- | Attempt to add two path sums. Only succeeds if the resulting sum is balanced
--   and the dimensions match.
plusMaybe :: (Eq g, Abelian g) => Pathsum g -> Pathsum g -> Maybe (Pathsum g)
plusMaybe sop sop'
  | inDeg sop  /= inDeg sop'                                       = Nothing
  | outDeg sop /= outDeg sop'                                      = Nothing
  | (sde sop) + 2*(pathVars sop') /= (sde sop') + 2*(pathVars sop) = Nothing
  | otherwise = Just $ Pathsum sde' inDeg' outDeg' pathVars' phasePoly' outVals'
  where sde'       = (sde sop) + 2*(pathVars sop')
        inDeg'     = inDeg sop
        outDeg'    = outDeg sop
        pathVars'  = (pathVars sop) + (pathVars sop') + 1
        y          = ofVar $ PVar (pathVars' - 1)
        phasePoly' = (lift y)*(phasePoly sop) +
                     (lift (1+y))*(renameMonotonic shift $ phasePoly sop')
        outVals'   = map (\(a,b) -> b + y*(a + b)) $
                       zip (outVals sop) (map (renameMonotonic shift) $ outVals sop')
        shift x    = case x of
          PVar i -> PVar $ i + (pathVars sop)
          _      -> x

-- | Construct the sum of two path sums. Raises an error if the sums are incompatible
plus :: (Eq g, Abelian g) => Pathsum g -> Pathsum g -> Pathsum g
plus sop sop' = case plusMaybe sop sop' of
  Nothing    -> error "Incompatible path sums"
  Just sop'' -> sop''

-- | Compose two path sums in parallel
tensor :: (Eq g, Num g) => Pathsum g -> Pathsum g -> Pathsum g
tensor sop sop' = Pathsum sde' inDeg' outDeg' pathVars' phasePoly' outVals'
  where sde'       = (sde sop) + (sde sop')
        inDeg'     = (inDeg sop) + (inDeg sop')
        outDeg'    = (outDeg sop) + (outDeg sop')
        pathVars'  = (pathVars sop) + (pathVars sop')
        phasePoly' = (phasePoly sop) + (renameMonotonic shift $ phasePoly sop')
        outVals'   = (outVals sop) ++ (map (renameMonotonic shift) $ outVals sop')
        shift x    = case x of
          IVar i -> IVar $ i + (inDeg sop)
          PVar i -> PVar $ i + (pathVars sop)
          _      -> x

-- | Attempt to compose two path sums in sequence. Only succeeds if the dimensions
--   are compatible (i.e. if the out degree of the former is the in degree of the
--   latter)
timesMaybe :: (Eq g, Abelian g) => Pathsum g -> Pathsum g -> Maybe (Pathsum g)
timesMaybe sop sop'
  | outDeg sop /= inDeg sop' = Nothing
  | otherwise = Just $ Pathsum sde' inDeg' outDeg' pathVars' phasePoly' outVals'
  where sde'       = (sde sop) + (sde sop')
        inDeg'     = inDeg sop
        outDeg'    = outDeg sop'
        pathVars'  = (pathVars sop) + (pathVars sop')
        phasePoly' = (phasePoly sop) +
                     (substMany sub . renameMonotonic shift $ phasePoly sop')
        outVals'   = (map (substMany sub . renameMonotonic shift) $ outVals sop')
        shift x    = case x of
          PVar i -> PVar $ i + (pathVars sop)
          _      -> x
        sub x      = case x of
          IVar i -> (outVals sop)!!i
          _      -> ofVar x

-- | Compose two path sums in sequence. Throws an error if the dimensions are
--   not compatible
times :: (Eq g, Abelian g) => Pathsum g -> Pathsum g -> Pathsum g
times sop sop' = case timesMaybe sop sop' of
  Nothing    -> error "Incompatible path sum dimensions"
  Just sop'' -> sop''

-- | Left-to-right composition
(.>) :: (Eq g, Abelian g) => Pathsum g -> Pathsum g -> Pathsum g
(.>) = times

infixr 5 .>

-- | Scale the normalization factor
renormalize :: Int -> Pathsum g -> Pathsum g
renormalize k (Pathsum a b c d e f) = Pathsum (a + k) b c d e f

-- | Embed a path sum into a larger space with a specified input and
--   output embedding.
embed :: (Eq g, Abelian g) => Pathsum g -> Int -> (Int -> Int) -> (Int -> Int) -> Pathsum g
embed sop n embedIn embedOut
  | n < 0     = error "Can't embed in smaller space"
  | otherwise = inPerm .> tensor (identity n) sop .> outPerm where
      mIn = inDeg sop
      ins = map embedIn [0..mIn-1]
      inPerm  = permutation $ ([0..mIn+n-1] \\ ins) ++ ins
      mOut = outDeg sop
      outs = map embedOut [0..mOut-1]
      outPerm = unpermutation $ ([0..mOut+n-1] \\ outs) ++ outs

-- | Drop a qubit
discard :: Eq g => Int -> Pathsum g -> Pathsum g
discard i sop@(Pathsum a b c d e f) = Pathsum a b' c' d e f' where
  b' = if i < b then b-1 else b
  c' = if i < c then c-1 else c
  f' = snd . unzip . filter (\(j,_) -> i /= j) $ zip [0..] f

{--------------------------
 Type class instances
 --------------------------}

instance (Eq g, Num g) => Semigroup (Pathsum g) where
  (<>) = tensor

instance (Eq g, Num g) => Monoid (Pathsum g) where
  mempty  = Pathsum 0 0 0 0 0 []
  mappend = tensor

instance (Eq g, Abelian g) => Num (Pathsum g) where
  (+)                          = plus
  (*)                          = (flip times)
  negate (Pathsum a b c d e f) = Pathsum a b c d (lift 1 + e) f
  abs (Pathsum a b c d e f)    = Pathsum a b c d (dropConstant e) f
  signum sop                   = sop
  fromInteger                  = identity . fromInteger

instance Functor Pathsum where
  fmap g (Pathsum a b c d e f) = Pathsum a b c d (cast g e) f

{--------------------------
 Reduction rules
 --------------------------}

{-
class RewriteRule rule g where
  matchAll :: Pathsum g -> [rule]
  matchOne :: Pathsum g -> rule
  apply :: rule -> Pathsum g -> Pathsum g

data E = E !Var {-# UNBOX #-}

instance (Eq g, Periodic g) => RewriteRule E g where
  match = matchElim
  apply = applyElim
-}

-- | Maps the order 1 and order 2 elements of a group to FF2
injectFF2 :: Periodic g => g -> Maybe FF2
injectFF2 a = case order a of
  1 -> Just 0
  2 -> Just 1
  _ -> Nothing

-- | Gives a Boolean polynomial equivalent to the current polynomial, if possible
toBooleanPoly :: (Eq g, Periodic g) => PseudoBoolean v g -> Maybe (SBool v)
toBooleanPoly = castMaybe injectFF2

-- | Elim rule. \(\dots(\sum_y)\dots = \dots 2 \dots\)
matchElim :: (Eq g, Periodic g) => Pathsum g -> [Var]
matchElim sop = msum . (map go) $ internalPaths sop
  where go v = if Set.member v (vars $ phasePoly sop) then [] else [v]

-- | Generic HH rule. \(\dots(\sum_y (-1)^{y\cdot f})\dots = \dots|_{f = 0}\)
matchHH :: (Eq g, Periodic g) => Pathsum g -> [(Var, SBool Var)]
matchHH sop = msum . (map (maybeToList . go)) $ internalPaths sop
  where go v = toBooleanPoly (quotVar v $ phasePoly sop) >>= \p -> return (v, p)

-- | Solvable instances of the HH rule.
--   \(\dots(\sum_y (-1)^{y(z \oplus f)})\dots = \dots[z \gets f]\)
matchHHSolve :: (Eq g, Periodic g) => Pathsum g -> [(Var, Var, SBool Var)]
matchHHSolve sop = do
  (v, p)   <- matchHH sop
  (v', p') <- solveForX p
  case v' of
    PVar _ -> return (v, v', p')
    _      -> mzero

-- | Instances of the HH rule with a linear substitution
matchHHLinear :: (Eq g, Periodic g) => Pathsum g -> [(Var, Var, SBool Var)]
matchHHLinear sop = do
  (v, p)   <- filter (\(_, p) -> degree p <= 1) $ matchHH sop
  (v', p') <- solveForX p
  return (v, v', p')

-- | Instances of the (\omega\) rule
matchOmega :: (Eq g, Periodic g, Dyadic g) => Pathsum g -> [(Var, SBool Var)]
matchOmega sop = do
  v <- internalPaths sop
  p <- maybeToList . toBooleanPoly . addFactor v $ phasePoly sop
  return (v, p)
  where addFactor v p = constant (fromDyadic $ dyadic 3 1) + quotVar v p

{--------------------------
 Pattern synonyms
 --------------------------}

-- | Pattern synonym for Elim
pattern Triv :: (Eq g, Num g) => Pathsum g
pattern Triv <- (isTrivial -> True)

-- | Pattern synonym for Elim
pattern Elim :: (Eq g, Periodic g) => Var -> Pathsum g
pattern Elim v <- (matchElim -> (v:_))

-- | Pattern synonym for HH
pattern HH :: (Eq g, Periodic g) => Var -> SBool Var -> Pathsum g
pattern HH v p <- (matchHH -> (v, p):_)

-- | Pattern synonym for solvable HH instances
pattern HHSolved :: (Eq g, Periodic g) => Var -> Var -> SBool Var -> Pathsum g
pattern HHSolved v v' p <- (matchHHSolve -> (v, v', p):_)

-- | Pattern synonym for linear HH instances
pattern HHLinear :: (Eq g, Periodic g) => Var -> Var -> SBool Var -> Pathsum g
pattern HHLinear v v' p <- (matchHHLinear -> (v, v', p):_)

-- | Pattern synonym for HH instances where the polynomial is strictly a
--   function of input variables
pattern HHKill :: (Eq g, Periodic g) => Var -> SBool Var -> Pathsum g
pattern HHKill v p <- (filter (all (not . isP) . vars . snd) . matchHH -> (v, p):_)

-- | Pattern synonym for Omega instances
pattern Omega :: (Eq g, Periodic g, Dyadic g) => Var -> SBool Var -> Pathsum g
pattern Omega v p <- (matchOmega -> (v, p):_)

{--------------------------
 Applying reductions
 --------------------------}

-- | Apply an elim rule. Does not check if the instance is valid
applyElim :: Var -> Pathsum g -> Pathsum g
applyElim (PVar i) (Pathsum a b c d e f) = Pathsum (a-2) b c (d-1) e' f'
  where e' = renameMonotonic varShift e
        f' = map (renameMonotonic varShift) f
        varShift (PVar j)
          | j > i     = PVar $ j - 1
          | otherwise = PVar $ j
        varShift v = v

-- | Apply a (solvable) HH rule. Does not check if the instance is valid
applyHHSolved :: (Eq g, Abelian g) => Var -> Var -> SBool Var -> Pathsum g -> Pathsum g
applyHHSolved (PVar i) v p (Pathsum a b c d e f) = Pathsum a b c (d-1) e' f'
  where e' = renameMonotonic varShift . subst v p . remVar (PVar i) $ e
        f' = map (renameMonotonic varShift . subst v p) f
        varShift (PVar j)
          | j > i     = PVar $ j - 1
          | otherwise = PVar $ j
        varShift v = v

-- | Apply an (\omega\) rule. Does not check if the instance is valid
applyOmega :: (Eq g, Abelian g, Dyadic g) => Var -> SBool Var -> Pathsum g -> Pathsum g
applyOmega (PVar i) p (Pathsum a b c d e f) = Pathsum (a-1) b c (d-1) e' f'
  where e' = renameMonotonic varShift $ p' + remVar (PVar i) e
        f' = map (renameMonotonic varShift) f
        p' = constant (fromDyadic $ dyadic 1 2) + distribute (fromDyadic $ dyadic 3 1) (lift p)
        varShift (PVar j)
          | j > i     = PVar $ j - 1
          | otherwise = PVar $ j
        varShift v = v

-- | Finds and applies the first elimination instance
rewriteElim :: (Eq g, Periodic g) => Pathsum g -> Pathsum g
rewriteElim sop = case sop of
  Elim v -> applyElim v sop
  _      -> sop

-- | Finds and applies the first hh instance
rewriteHH :: (Eq g, Periodic g) => Pathsum g -> Pathsum g
rewriteHH sop = case sop of
  HHSolved v v' p -> applyHHSolved v v' p sop
  _               -> sop

-- | Finds and applies the first omega instance
rewriteOmega :: (Eq g, Periodic g, Dyadic g) => Pathsum g -> Pathsum g
rewriteOmega sop = case sop of
  Omega v p -> applyOmega v p sop
  _         -> sop

{--------------------------
 Reduction procedures
 --------------------------}

-- | Performs basic simplifications
simplify :: (Eq g, Periodic g, Dyadic g) => Pathsum g -> Pathsum g
simplify sop = case sop of
  Elim y         -> grind $ applyElim y sop
  HHLinear y z p -> grind $ applyHHSolved y z p sop
  Omega y p      -> grind $ applyOmega y p sop
  _              -> sop

-- | A complete normalization procedure for Clifford circuits. Originally described in
--   the paper M. Amy,
--   / Towards Large-Scaled Functional Verification of Universal Quantum Circuits /, QPL 2018.
grind :: (Eq g, Periodic g, Dyadic g) => Pathsum g -> Pathsum g
grind sop = case sop of
  Elim y         -> grind $ applyElim y sop
  HHSolved y z p -> grind $ applyHHSolved y z p sop
  Omega y p      -> grind $ applyOmega y p sop
  _              -> sop

-- | A single step of 'grind'
grindStep :: (Eq g, Periodic g, Dyadic g) => Pathsum g -> Pathsum g
grindStep sop = case sop of
  Elim y         -> applyElim y sop
  HHSolved y z p -> applyHHSolved y z p sop
  Omega y p      -> applyOmega y p sop
  _              -> sop

{--------------------------
 Simulation
 --------------------------}

-- | Gets the cofactors of some path variable
expand :: (Eq g, Abelian g) => Pathsum g -> Var -> (Pathsum g, Pathsum g)
expand (Pathsum a b c d e f) v = (p0, p1) where
  p0  = Pathsum a b c (d-1) (subst v 0 e) (map (subst v 0) f)
  p1  = Pathsum a b c (d-1) (subst v 1 e) (map (subst v 1) f)

-- | Simulates a pathsum on a given input
simulate :: (Eq g, Periodic g, Dyadic g, Real g, RealFloat f) => Pathsum g -> [FF2] -> Map [FF2] (Complex f)
simulate sop xs = go $ sop * ket (map constant xs)
  where go      = go' . grind
        go' ps  = case ps of
          (Pathsum k 0 _ 0 p xs) ->
            let phase     = fromRational . toRational $ getConstant p
                base      = case k `mod` 2 of
                  0 -> fromInteger $ 1 `shiftL` (abs k)
                  1 -> sqrt(2.0) * (fromInteger $ 1 `shiftL` (abs (k-1)))
                magnitude = base**(fromIntegral $ signum k)
            in
              Map.singleton (map getConstant xs) (mkPolar magnitude (pi * phase))
          (Pathsum k 0 n i p xs) ->
            let v     = PVar $ i-1
                left  = go (Pathsum k 0 n (i-1) (subst v zero p) (map (subst v zero) xs))
                right = go (Pathsum k 0 n (i-1) (subst v one p) (map (subst v one) xs))
            in
              Map.unionWith (+) left right
          _                      -> error "Incompatible dimensions"

-- | Evaluates a pathsum on a given input and output
amplitude :: (Eq g, Periodic g, Dyadic g, Real g, RealFloat f) => [FF2] -> Pathsum g -> [FF2] -> Complex f
amplitude o sop i = (simulate (bra (map constant o) * sop) i)![]

-- | Checks identity by checking inputs iteratively
isIdentity :: (Eq g, Periodic g, Dyadic g) => Pathsum g -> Bool
isIdentity sop
  | isTrivial sop = True
  | otherwise     = case inDeg sop of
      0 -> False
      i -> let p0 = (grind $ identity (i-1) <> ket [0] .> sop .> identity (i-1) <> bra [0])
               p1 = (grind $ identity (i-1) <> ket [1] .> sop .> identity (i-1) <> bra [1])
           in
             isIdentity p0 && isIdentity p1

{--------------------------
 Examples
 --------------------------}

-- | A symbolic state |x>
sstate :: Pathsum DMod2
sstate = open $ identity 1

-- | A bell state
bellstate :: Pathsum DMod2
bellstate = fresh <> fresh .> hgate <> (identity 1) .> cxgate

-- | Teleportation circuit
teleport :: Pathsum DMod2
teleport = (identity 1) <> bellstate .>
           cxgate <> (identity 1) .>
           hgate <> cxgate .>
           swapgate <> hgate .>
           (identity 1) <> cxgate .>
           swapgate <> hgate

-- | Teleportation channel
teleportChannel :: Pathsum DMod2
teleportChannel = channelize ((identity 1) <> bellstate) .>
                  channelize (cxgate <> (identity 1)) .>
                  channelize (hgate <> cxgate) .>
                  embed measure 4 (* 3) (* 3) .>
                  embed measure 4 (\i -> i*3 + 1) (\j -> j*3 + 1) .>
                  channelize (swapgate <> hgate) .>
                  channelize ((identity 1) <> cxgate) .>
                  channelize (swapgate <> hgate) .>
                  embed epsilon 4 (* 3) (* 3) .> -- trace out first qubit
                  embed epsilon 2 (* 2) (* 2)    -- trace out second

-- | Verify teleportation
verifyTele :: () -> IO ()
verifyTele _ = case (densify sstate == grind (ptrace . ptrace . densify $ sstate .> teleport)) of
  True -> putStrLn "Identity"
  False -> putStrLn "Not identity"

-- | Verify teleportation channel
verifyTeleC :: () -> IO ()
verifyTeleC _ = case (rho == grind (rho .> teleportChannel)) of
  True -> putStrLn "Identity"
  False -> putStrLn "Not identity"
  where rho = grind $ vectorize $ densify sstate

-- | The |A> = T|+> state
aState :: Pathsum DMod2
aState = fresh .> hgate .> tgate

-- | T gate teleportation channel
teleportTChannel :: Pathsum DMod2
teleportTChannel = channelize ((identity 1) <> aState) .>
                   channelize (cxgate) .>
                   embed measure 2 (\i -> i*2 + 1) (\j -> j*2 + 1) .>
                   channelize (controlled sgate) .>
                   embed epsilon 2 (\i -> i*2 + 1) (\j -> j*2 + 1) -- trace out the resource state

-- | Verify teleportation channel
verifyTeleT :: () -> IO ()
verifyTeleT _ = case (channelize tgate == grind (teleportTChannel)) of
  True -> putStrLn "Identity"
  False -> putStrLn "No identity"


{-# Language BangPatterns #-}
{-# Language ScopedTypeVariables #-}

module Neuralnetwork where

import           Data.Bifunctor                 ( bimap )
import           Data.Functor.Identity
import           Data.IDX
import           Data.Maybe                     ( fromMaybe )
-- import           Debug.Trace
import           Control.Monad.Random
import qualified Data.Vector.Unboxed           as V
import           Numeric.LinearAlgebra          ( Matrix
                                                , Vector
                                                , (<>)
                                                , (#>)
                                                , (><)
                                                , (|>)
                                                )
import qualified Numeric.LinearAlgebra         as LA

-- The main crux of my flaw seems to be that my weights increase without bound.
--
-- I've run into this issue _every_ single time I have written up this neural
-- network.  I've rewritten the weight updates countless times now and it's
-- starting to get to the point where my being past the deadline is a bit
-- ridiculous and if I keep it up I won't even get to the second programming
-- assignment at all.
--
-- See the delta functions and the weight update function for the most likely
-- causes of me screwing up.

-- Weights :
--   - biases = n vector of length = num of inputs to that layer
--   - nodes  = n x m where m = num of outputs to that layer
--   => m -> n layer
data Weights = W
  { wBiases :: !(Vector Double) -- length = num of inputs
  , wNodes  :: !(Matrix Double)
  } deriving Show

-- Type aliases offer no type safety but help make type signatures more
-- readable.
type Network    = [Weights]
type Input      = Vector Double
type Target     = Vector Double
type Output     = Vector Double
type Activation = Vector Double
type Prev a     = a -- Cute trick to help make type signatures more readable
type Diff a     = a

-- Here we see that we have, behold, dreaded globals.
-- However, this is fine because
--  a) I don't care
--  b) The reader pattern is overly complicated for this particular usecase
η :: Double -- Learning Rate
η = 0.1

α :: Double -- Momentum Rate
α = 0.9

-- Layer + Network generation
-- Generates a layer of random numbers between -0.05 and 0.05
randLayer :: MonadRandom m => Int -> Int -> m Weights
randLayer i o = do
  r1 :: Int <- getRandom
  r2 :: Int <- getRandom
  let bias    = (LA.randomVector r1 LA.Uniform o - 0.5) / 10
      weights = LA.uniformSample r2 o (replicate i (-0.05, 0.05))
  pure $ W bias weights

-- Generates a layer consisting of all zeroed weights
zeroLayer :: Int -> Int -> Identity Weights
zeroLayer i o = pure $ W (o |> repeat 0) (o >< i $ repeat 0)

-- A network is a list of layers; this simply steps through the list
randNet :: MonadRandom m => Int -> [Int] -> Int -> m Network
randNet i []       o = (:) <$> randLayer i o <*> pure []
randNet i (h : hs) o = (:) <$> randLayer i h <*> randNet h hs o

zeroNet :: Int -> [Int] -> Int -> Identity Network
zeroNet i []       o = (:) <$> zeroLayer i o <*> pure []
zeroNet i (h : hs) o = (:) <$> zeroLayer i h <*> zeroNet h hs o

-- Forward propagate a neural network
-- a #> b is matrix vector product where vector b is Nx1 dimensional
-- i : a 784 length vector
-- wN : (20,784) shaped matrix
runLayer :: Weights -> Input -> Activation
runLayer (W !wB !wN) !i = (wN #> i) + wB
-- runLayer (W !wB !wN) !i = trace "runLayer" $ (wN #> i) + wB

-- Gets the answer from a neural network by propagating the input forward
runNet :: Network -> Input -> Output
runNet !n !i = last $ runNet' n i
-- runNet !n !i = trace "runNet" $ last $ runNet' n i

-- Collects all intermediate activations into a list of activations
runNet' :: Network -> Input -> [Activation]
runNet' []        !i = [i]
runNet' (w : net) !i = i : runNet' net α where α = activation w i
-- runNet' (w : net) !i = trace "runNet'" $ i : runNet' net α where α = activation w i

-- Runs the neural network over a list of inputs and collects every "answer"
collectOutputs :: Network -> [Input] -> [Output]
collectOutputs net = fmap (runNet net)

logistic :: Floating a => a -> a
logistic x = 1 / (1 + exp (-x))

logistic' :: Floating a => a -> a
logistic' x = o * (1 - o) where o = logistic x

activation :: Weights -> Input -> Activation
activation !w !v = logistic $ runLayer w v

-- (Not used when actually running the neural network. The slides didn't
-- mention using a loss function at all to update the weights...)
-- Target -> Output -> Error
loss :: Target -> Output -> Double
loss = (* 0.5) ... LA.sumElements ... (^ 2) ... (-)

-- Not used anywhere in the codebase
lossO :: Floating a => a -> a -> a -> a
lossO t y α = logistic' y * (α - t)

-- If there's an error in this code, it is likely to be in the delta functions
-- or the weight updates
δo :: Target -> Output -> Vector Double
δo !t !α = α * (1 - α) * (t - α)
-- δo !t !α = trace "δo" $ α * (1 - α) * (t - α)

δh :: Weights -> Activation -> Vector Double -> Vector Double
δh (W !wB !wN) !h !δ' = h * (1 - h) * (LA.tr' wN #> δ')
-- δh (W !wB !wN) !h !δ' = trace "δh" $ h * (1 - h) * (LA.tr' wN #> δ')
-- Can't be this because it's not an inconsistent matrix error...

-- The error terms; algorithm for this and the deltas given by the slides
errors :: Network -> Target -> [Input] -> [Vector Double]
errors [!ih, !ho] !t [!αH, !αO] = [δhidden, δoutput]
-- errors [!ih, !ho] !t [!αH, !αO] = trace "errors" [δhidden, δoutput]
 where
  δoutput = δo t αO
  δhidden = δh ho αH δoutput

-- Weight deltas with momentum (only the deltas, not the actual change)
-- scalar is a function that allows me to multiply a matrix (or vector) by a single number
-- δ and i need to be the same size. δ should be 20
-- if wNΔ is (n,m) size, LA.asRow δ*i needs to be (x,m), LA.asColumn δ*i needs to be (n,x)
weightΔ :: Vector Double -> Input -> Prev (Diff Weights) -> (Weights, Diff Weights)
weightΔ !δ !i w@(W !wBΔ !wNΔ) = (w, W wΔb wΔn)
-- weightΔ !δ !i w@(W !wBΔ !wNΔ) = trace "weightΔ" (w, W wΔb wΔn)
 where
  wΔb = LA.scalar η * δ + LA.scalar α * wBΔ
  wΔn = LA.asColumn (LA.scalar η * δ * i) + LA.scalar α * wNΔ
  -- tracePrint = "δ: "      ++ show (LA.size δ)
  --           ++ " i: "     ++ show (LA.size i)
  --           ++ " wNΔ: "   ++ show (LA.size wNΔ)
  --           ++ " wBΔ: "   ++ show (LA.size wBΔ)
            -- ++ " asCol: " ++ show (LA.size $ LA.scalar η * δ * i)
            -- δ: 20 i: 784 wNΔ: (20,784) wBΔ: 20
            -- passing in wrong δ?

wgtUpd :: Weights -> Diff Weights -> Weights
wgtUpd (W wB wN) (W wB' wN') = W (wB + wB') (wN + wN')
-- wgtUpd (W wB wN) (W wB' wN') = trace "wgtUpd" $ W (wB + wB') (wN + wN')

backprop -- Given the list of prev diff weights for every layer as well
  :: (Prev Network, Network)
  -> Target
  -> Input
  -> [Prev (Diff Weights)]
  -> ([Diff Weights], Network)
backprop (net', net) !target !input w' = (wΔs, zipWith wgtUpd net wΔs)
-- backprop (net', net) !target !input w' = trace "Backprop" (wΔs, zipWith wgtUpd net wΔs)
 where
  αs  = tail $ runNet' net input -- the activations, excluding the input vector
  δs  = errors net target αs     -- length 2 lst
  wΔs = snd <$> zipWith3 weightΔ δs αs w'

-- Iterate through a list of inputs (and list of targets) and successively
-- train on those inputs.
epoch
  :: (Prev Network, Network)   -- Old network (zeros in first run), current network
  -> [Target]                  -- List of input vectors
  -> [Input]                   -- List of target vectors
  -> [Diff Weights]
  -> ([Diff Weights], Network) -- Resulting network
epoch (_    , net ) []       []       wΔs   = (wΔs, net)
epoch (!net', !net) (t : ts) (i : is) !wΔs' = epoch (net, next) ts is wΔs
-- epoch (!net', !net) (t : ts) (i : is) !wΔs' = trace "Epoch happening" $ epoch (net, next) ts is wΔs
  where (wΔs, next) = backprop (net', net) t i wΔs'

train
  :: Int -- Epochs to train for
  -> (Prev Network, Network) -- (Rest of params are identical to epoch)
  -> ([Target], [Target])
  -> ([Input], [Input])
  -> [Diff Weights]
  -> IO ()
train 0 nets (tstT, t) (tstI, i) wΔs' = do
  let final    = snd $ epoch nets t i wΔs'
      outputs  = collectOutputs final i
      oTest    = collectOutputs final tstI
      accuracy = totalError (t, outputs)
      tstAcc   = totalError (tstT, oTest)
  print $ "Epoch 50 accuracy: " ++ show accuracy
  print $ "Epoch 50 test accuracy: " ++ show tstAcc
  -- print out stuff
train n (net', net) t'@(tstT, t) i'@(tstI, i) wΔs' = do
  let (wΔs, next) = epoch (net', net) t i wΔs'
      outputs     = collectOutputs next i
      oTest       = collectOutputs next tstI
      accuracy    = totalError (t, outputs)
      tstAcc      = totalError (tstT, oTest)
  print $ "Epoch " ++ show (50-n) ++ " accuracy: "      ++ show accuracy
  print $ "Epoch " ++ show (50-n) ++ " test accuracy: " ++ show tstAcc
  train (n - 1) (net, next) t' i' wΔs

-- Simple debugging function that tells me the shape of a particular network.
-- Useful for ensuring my matrices are the right shape
sizeOf :: Network -> [(Int, (Int, Int))]
sizeOf [W hB hN, W oB oN] = [(LA.size hB, LA.size hN), (LA.size oB, LA.size oN)]

totalError :: ([Target], [Output]) -> Double
totalError (!ts, !os) = sum (zipWith loss ts os) / fromIntegral (length ts)

main :: IO ()
main = do
  -- Yes, this is slightly ugly...
  !file <- fromMaybe (error "file decoding failed")
    <$> decodeIDXFile "/home/jaredweakly/Documents/Classes/CS445/train-images-idx3-ubyte"
  !labels <- fromMaybe (error "label decoding failed") <$> decodeIDXLabelsFile
    "/home/jaredweakly/Documents/Classes/CS445/train-labels-idx1-ubyte"
  !testFile <- fromMaybe (error "test file decoding failed")
    <$> decodeIDXFile "/home/jaredweakly/Documents/Classes/CS445/t10k-images-idx3-ubyte"
  !testLabels <- fromMaybe (error "test label decoding failed") <$> decodeIDXLabelsFile
    "/home/jaredweakly/Documents/Classes/CS445/t10k-labels-idx1-ubyte"

  -- mnist : 60_000 [(target, Vector 784 Double)]
  let !mnist = fromMaybe (error "labeling failed") $ labeledDoubleData labels file
      !mnist' =
        fromMaybe (error "labeling failed") $ labeledDoubleData testLabels testFile
      !initN20   = runIdentity $ zeroNet 784 [20] 10
      !initN50   = runIdentity $ zeroNet 784 [50] 10
      !initN100  = runIdentity $ zeroNet 784 [100] 10
      !testset = mnist' ++ mnist

  print "Loaded up data"
  -- inputs : (60000><785) Matrix Double
  -- targets : Vector 10 Double. 0.1 everywhere but target which is 0.9
  let mkSparseVectors [] = [] :: [Vector Double]
      mkSparseVectors (x : xs) =
        LA.assoc 10 0.1 [(x, 0.9)] : mkSparseVectors xs :: [Vector Double]

      normalize :: V.Vector Double -> Vector Double
      normalize            = (/ 255) . V.convert

      !(!targets, !inputs) = bimap mkSparseVectors id . unzip $ normalize <$$> mnist
      !(!tstT   , !tstI  ) = bimap mkSparseVectors id . unzip $ normalize <$$> testset

  !weights20  <- randNet 784 [20] 10
  !weights50  <- randNet 784 [50] 10
  !weights100 <- randNet 784 [100] 10

  print "Beginning training with n=20"
  train 50 (initN20, weights20) (tstT, targets) (tstI, inputs) initN20

  print "Beginning training with n=50"
  train 50 (initN50, weights50) (tstT, targets) (tstI, inputs) initN50

  print "Beginning training with n=100"
  train 100 (initN100, weights100) (tstT, targets) (tstI, inputs) initN100

-- Helper functions I was too lazy to import
-- These are not essential to understanding any of the actual code.

-- Lets me write pointfree code using functions that takes two arguments
-- instead of just one
infixr 9 ...
(...) :: (b -> c) -> (a1 -> a2 -> b) -> a1 -> a2 -> c
(...) = (.) . (.)

-- Lets me map over something inside another structure.
-- eg: (+1) <$$> [[a]] would map over all of the inner lists
infixl 4 <$$>
(<$$>) :: (Functor f2, Functor f1) => (a -> b) -> f1 (f2 a) -> f1 (f2 b)
(<$$>) = fmap fmap fmap

both :: (a -> b) -> (a, a) -> (b, b)
both f (a, a') = (f a, f a')


{-# LANGUAGE DeriveDataTypeable #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE MultiParamTypeClasses #-}

module Scheme.Types where

import Control.Exception (Exception, throw)
import Control.Monad.Except (ExceptT)
import Data.Array (Array, elems)
import qualified Data.ByteString as BS
import Data.Complex (Complex ((:+)), imagPart, realPart)
import Data.IORef (IORef)
import Data.Ratio (denominator, numerator)
import Data.Text (Text)
import qualified Data.Text as T
import Data.Typeable (Typeable)
import System.IO (Handle)
import Text.Pretty.Simple (pPrint, pPrintLightBg)
import Text.Show.Functions ()

type Env = IORef [(Text, IORef SchemeVal)]

type SchemeResult = Either SchemeError

type IOSchemeResult = ExceptT SchemeError IO

extractValue :: Either a p -> p
extractValue (Right val) = val
extractValue _ = error "Called with Left"

data Number
  = Integer Integer
  | Real Double
  | Rational Rational
  | Complex (Complex Double)
  deriving (Show, Typeable, Eq)

showNum :: Number -> Text
showNum (Integer i) = T.pack $ show i
showNum (Real d) = T.pack $ show d
showNum (Rational r) = T.pack $ show (numerator r) <> "/" <> show (denominator r)
showNum (Complex p) = T.pack $ show (realPart p) <> "+" <> show (imagPart p) <> "i"

instance Ord Number where
  compare (Integer x) (Integer y) = compare x y
  compare (Integer x) (Real y) = compare (fromInteger x) y
  compare (Integer x) (Rational y) = compare (fromInteger x) y
  compare (Real x) (Integer y) = compare x (fromInteger y)
  compare (Rational x) (Integer y) = compare x (fromInteger y)
  compare (Real x) (Real y) = compare x y
  compare (Real x) (Rational y) = compare (toRational x) y
  compare (Rational x) (Real y) = compare x (toRational y)
  compare (Rational x) (Rational y) = compare x y
  compare _ _ = throw $ InvalidOperation "Cannot compare complex numbers"

instance Num Number where
  fromInteger x = Integer x

  (Integer x) + (Integer y) = Integer (x + y)
  (Integer x) + (Real y) = Real (fromInteger x + y)
  (Integer x) + (Rational y) = Rational (fromInteger x + y)
  (Integer x) + (Complex y) = Complex (fromInteger x + y)
  (Real x) + (Integer y) = Real (x + fromInteger y)
  (Rational x) + (Integer y) = Rational (x + fromInteger y)
  (Complex x) + (Integer y) = Complex (x + fromInteger y)
  (Real x) + (Real y) = Real (x + y)
  (Real x) + (Rational y) = Rational (toRational x + y)
  (Real x) + (Complex y) = Complex ((x :+ 0.0) + y)
  (Rational x) + (Real y) = Rational (x + toRational y)
  (Complex x) + (Real y) = Complex (x + (y :+ 0.0))
  (Rational x) + (Rational y) = Rational (x + y)
  (Complex x) + (Complex y) = Complex (x + y)
  _ + _ = throw $ InvalidOperation "Cannot add rationals and complex numbers"

  (Integer x) * (Integer y) = Integer (x * y)
  (Integer x) * (Real y) = Real (fromInteger x * y)
  (Integer x) * (Rational y) = Rational (fromInteger x * y)
  (Integer x) * (Complex y) = Complex (fromInteger x * y)
  (Real x) * (Integer y) = Real (x * fromInteger y)
  (Rational x) * (Integer y) = Rational (x * fromInteger y)
  (Complex x) * (Integer y) = Complex (x * fromInteger y)
  (Real x) * (Real y) = Real (x * y)
  (Real x) * (Rational y) = Rational (toRational x * y)
  (Real x) * (Complex y) = Complex ((x :+ 0.0) * y)
  (Rational x) * (Real y) = Rational (x * toRational y)
  (Complex x) * (Real y) = Complex (x * (y :+ 0.0))
  (Rational x) * (Rational y) = Rational (x * y)
  (Complex x) * (Complex y) = Complex (x * y)
  _ * _ = throw $ InvalidOperation "Cannot multiply rationals and complex numbers"

  (Integer x) - (Integer y) = Integer (x - y)
  (Integer x) - (Real y) = Real (fromInteger x - y)
  (Integer x) - (Rational y) = Rational (fromInteger x - y)
  (Integer x) - (Complex y) = Complex (fromInteger x - y)
  (Real x) - (Integer y) = Real (x - fromInteger y)
  (Rational x) - (Integer y) = Rational (x - fromInteger y)
  (Complex x) - (Integer y) = Complex (x - fromInteger y)
  (Real x) - (Real y) = Real (x - y)
  (Real x) - (Rational y) = Rational (toRational x - y)
  (Real x) - (Complex y) = Complex ((x :+ 0.0) - y)
  (Rational x) - (Real y) = Rational (x - toRational y)
  (Complex x) - (Real y) = Complex (x - (y :+ 0.0))
  (Rational x) - (Rational y) = Rational (x - y)
  (Complex x) - (Complex y) = Complex (x - y)
  _ - _ = throw $ InvalidOperation "Cannot subtract rationals and complex numbers"

  negate (Integer x) = Integer (negate x)
  negate (Real x) = Real (negate x)
  negate (Rational x) = Rational (negate x)
  negate (Complex x) = Complex (negate x)

  abs (Integer x) = Integer (abs x)
  abs (Real x) = Real (abs x)
  abs (Rational x) = Rational (abs x)
  abs (Complex x) = Complex (abs x)

  signum (Integer x) = Integer (signum x)
  signum (Real x) = Real (signum x)
  signum (Rational x) = Rational (signum x)
  signum (Complex x) = Complex (signum x)

instance Fractional Number where
  fromRational x = Rational x

  (Integer x) / (Integer y) = Rational (fromInteger x / fromInteger y)
  (Integer x) / (Real y) = Real (fromInteger x / y)
  (Integer x) / (Rational y) = Rational (fromInteger x / y)
  (Integer x) / (Complex y) = Complex (fromInteger x / y)
  (Real x) / (Integer y) = Real (x / fromInteger y)
  (Rational x) / (Integer y) = Rational (x / fromInteger y)
  (Complex x) / (Integer y) = Complex (x / fromInteger y)
  (Real x) / (Real y) = Real (x / y)
  (Real x) / (Rational y) = Rational (toRational x / y)
  (Real x) / (Complex y) = Complex ((x :+ 0.0) / y)
  (Rational x) / (Real y) = Rational (x / toRational y)
  (Complex x) / (Real y) = Complex (x / (y :+ 0.0))
  (Rational x) / (Rational y) = Rational (x / y)
  (Complex x) / (Complex y) = Complex (x / y)
  _ / _ = throw $ InvalidOperation "Cannot divide rationals and complex numbers"

  recip (Integer x) = Real (recip $ fromInteger x)
  recip (Real x) = Real (recip x)
  recip (Rational x) = Rational (recip x)
  recip (Complex x) = Complex (recip x)

data Fn = Fn {macro :: Bool, params :: [Text], vararg :: Maybe Text, body :: [SchemeVal], closure :: Env}

instance Show Fn where
  show _ = "<fn>"

data SchemeVal
  = List [SchemeVal]
  | PairList [SchemeVal] SchemeVal
  | Vector (Array Int SchemeVal)
  | Bytevector BS.ByteString
  | String Text
  | Character Char
  | Symbol Text
  | Boolean Bool
  | Number Number
  | Nil
  | Primitive ([SchemeVal] -> SchemeResult SchemeVal)
  | Fun Fn
  | IOFun ([SchemeVal] -> IOSchemeResult SchemeVal)
  | Port Handle
  deriving (Show, Typeable)

instance Eq SchemeVal where
  (==) (List x) (List y) = length x == length y && all (uncurry (==)) (zip x y)
  (==) (PairList xs x) (PairList ys y) = List (xs ++ [x]) == List (ys ++ [y])
  (==) (Vector x) (Vector y) = x == y
  (==) (Bytevector x) (Bytevector y) = x == y
  (==) (String x) (String y) = x == y
  (==) (Character x) (Character y) = x == y
  (==) (Symbol x) (Symbol y) = x == y
  (==) (Boolean x) (Boolean y) = x == y
  (==) (Number x) (Number y) = x == y
  (==) _ _ = False

unwordVals :: [SchemeVal] -> Text
unwordVals xs = T.unwords $ showVal <$> xs

showVal :: SchemeVal -> Text
showVal (List (Symbol "quote" : xs)) = "'" <> unwordVals xs
showVal (List (Symbol "quasiquote" : xs)) = "`" <> unwordVals xs
showVal (List (Symbol "unquote" : xs)) = "," <> unwordVals xs
showVal (List (Symbol "unquote-splicing" : xs)) = ",@" <> unwordVals xs
showVal (List contents) = "(" <> unwordVals contents <> ")"
showVal (PairList contents cdr) = "(" <> unwordVals contents <> " . " <> showVal cdr <> ")"
showVal (Vector vec) = T.pack $ "#(" <> unwords (map show $ elems vec) <> ")"
showVal (Bytevector vec) = T.pack $ "#u8(" <> unwords (map show $ BS.unpack vec) <> ")"
showVal (String s) = "\"" <> s <> "\""
showVal (Character a) = T.pack $ "#\\" <> [a]
showVal (Symbol s) = s
showVal (Boolean True) = "#t"
showVal (Boolean False) = "#f"
showVal (Number n) = showNum n
showVal Nil = "nil"
showVal Primitive {} = "<prim>"
showVal Fun {} = "<fun>"
showVal IOFun {} = "<io>"
showVal (Port _) = "<port>"

dumpAST :: SchemeVal -> IO ()
dumpAST = dumpAST' True

dumpAST' :: Bool -> SchemeVal -> IO ()
dumpAST' True = pPrint
dumpAST' False = pPrintLightBg

data SchemeError
  = Generic Text
  | ArgumentLengthMismatch Int [SchemeVal]
  | TypeMismatch Text SchemeVal
  | UnboundSymbol Text
  | ParserError Text
  | NotFunction SchemeVal
  | InvalidOperation Text
  | ReservedName Text
  | EmptyList
  deriving (Show)

instance Exception SchemeError

showError :: SchemeError -> Text
showError (TypeMismatch err vap) = "Invalid type: expected " <> err <> ", but found " <> T.pack (show vap)
showError (Generic err) = "Unexpectec error: " <> err
showError (UnboundSymbol sym) = "Unbound symbol: " <> sym
showError (ArgumentLengthMismatch ex act) = "Expected " <> T.pack (show ex) <> " but found " <> T.pack (show $ length act)
showError (ParserError err) = "Parsing error, could not parse input: " <> err
showError (NotFunction err) = "Attempt at calling " <> showVal err <> " as a function"
showError (InvalidOperation err) = "Invalid: " <> err
showError (ReservedName err) = "Cannot define : " <> err <> " because it is a reserved name"
showError EmptyList = "Empty list is not allowed"


{-# LANGUAGE RankNTypes, BangPatterns, ScopedTypeVariables, FlexibleInstances #-}

module Math.Probably.MALA where

import qualified Math.Probably.PDF as PDF
import Math.Probably.Sampler
import qualified Math.Probably.PDF as PDF
import Control.Applicative

import Numeric.LinearAlgebra
import Text.Printf
import System.IO
import Data.Maybe

import Statistics.Test.KolmogorovSmirnov
import Statistics.Test.MannWhitneyU
import qualified Data.Vector.Unboxed as U
import qualified Control.Monad.State.Strict as S
import qualified Data.Vector.Storable as VS

import Debug.Trace
import Data.IORef
import Control.Spoon

data MalaPar = MalaPar { mpXi :: !(Vector Double),
                         mpPi :: !Double,
                         mpGradLast :: !(Vector Double),
                         mpSigma :: !Double,
                         mpCount :: !Double,
                         mpAccept :: !Double,
                         mpFreezeSigma :: !Bool,
                         mpThin :: !Int,
                         mpMaxGradLen :: !Double } --,
--                         mpLastRatio :: !Double }
  deriving Show

trunc t v = scale (t/(max t (norm2 v))) v

mala1 :: Covariance a => a -> (Vector Double -> (Double,Vector Double))
         -> Bool
         -> MalaPar
         -> Sampler MalaPar
mala1 cov postgrad useCache
      (MalaPar xi piCached gradientiCached sigma tr tracc freeze thinn maxGrad) = do
  let (pi,gradienti) = if useCache
                          then (piCached, gradientiCached)
                          else let (p, gr_untr) = postgrad xi
                               in (p, trunc maxGrad gr_untr)
  let xstarMean = xi + scale (sigma/2) (cov `covMul` gradienti)
  xstar <- mvnSampler xstarMean sigma cov
  u <- unitSample
  let (!pstar, gradientStarUntrunc) = postgrad xstar
      gradientStar = trunc maxGrad gradientStarUntrunc
  let !revJumpMean = xstar + scale (sigma/2) (cov `covMul` gradientStar)
      ptop = mvnPDF revJumpMean sigma cov xi
      pbot = mvnPDF xstarMean sigma cov xstar
      ratio = exp $   pstar -pi + ptop - pbot
      tr' = max 1 tr
      sigmaNext = case () of
         _ | freeze -> sigma
      accept = tracc / tr
      freezeNext = freeze {-| freeze = True
                 | not freeze = tr > 100 && accept > 0.5 && accept < 0.6  -}
  if trace (show $ (tr, pstar ,pi , ratio, sigma)) $ u < ratio
     then return $ MalaPar xstar pstar (gradientStar)
                           (if freezeNext then sigma else (min 1.4 $ 1+kmala/tr')*sigma)
                           (tr+1) (tracc+1) freezeNext  thinn maxGrad
--                           sigma (tr+1) (tracc+1)
     else return $ MalaPar xi pi ( gradienti)
                           (if freezeNext then sigma else (max 0.7143 $ 1-kmala/tr')**1.3*sigma)
                           (tr+1) tracc freezeNext thinn maxGrad
--                           sigma (tr+1) tracc

runMalaMP :: Covariance a => a -> (Vector Double -> (Double,Vector Double))
         ->  Int -> MalaPar -> [(Double,Vector Double)] -> RIO (MalaPar, [(Double,Vector Double)])
runMalaMP cov  pdf nsam init xs0 = go nsam init xs0 where
  go 0 mpar xs = do io $ putStrLn $ "MALA accept = "++show (mpAccept mpar/mpCount mpar)
                    io $ putStrLn $ "MALA sigma = "++show (mpSigma mpar)
                    return (mpar, xs)
  go n y xs = do y1 <- sample $ mala1 cov pdf True y
--                 io $ do putStrLn $ show (mpCount y1, mpPi y1, mpSigma y1)
--                         hFlush stdout
                 let newChainRes = if mpThin y1 == 0 || round (mpCount y1) `mod` mpThin y1 ==0
                                      then let !xi = mpXi y1
                                               !pi = mpPi y1
                                               !more = (pi,xi)
                                           in more:xs
                                      else xs
                 go (n-1) y1 newChainRes

runMalaUntilBetter :: Covariance a => a ->(Vector Double -> (Double,Vector Double))
         ->  Int -> MalaPar -> RIO (MalaPar, [(Double,Vector Double)])
runMalaUntilBetter cov  pdf nsam init = go nsam init [] where
  pTarget = mpPi init
  go 0 mpar xs = do io $ putStrLn $ "MALA accept = "++show (mpAccept mpar/mpCount mpar)
                    io $ putStrLn $ "MALA sigma = "++show (mpSigma mpar)
                    return (mpar, xs)
  go n y xs = do y1 <- sample $ mala1 cov  pdf True y
                 io $ do putStrLn $ show (mpCount y1, mpPi y1, mpSigma y1)
                         hFlush stdout
                 let newChainRes = if mpThin y1 == 0 || round (mpCount y1) `mod` mpThin y1 ==0
                                      then (mpPi y1, mpXi y1):xs
                                      else xs
                 if mpPi y1 > pTarget
                    then return (y1, newChainRes)
                    else go (n-1) y1 $ newChainRes

{-runMalaMPaccept' :: Matrix Double -> (Vector Double -> (Double,Vector Double))
         ->  Int -> MalaPar -> RIO (MalaPar, [(Double,Vector Double)])
runMalaMPaccept'  cov pdf nsam init  = do
  stseed <- S.get
  res <- io $ do
   seedref <- newIORef stseed
   resRef <- newIORef Nil
   let go !mpar
        | (round $ mpAccept mpar) >= nsam
          = return mpar
        | otherwise
          = do seed <- readIORef seedref
               putStr ((show (round $ mpAccept mpar)) ++".") >> hFlush stdout
               let (!mpar1, !seed1) = unSam (mala1 cov pdf mpar) seed
               writeIORef seedref seed1
               modifyIORef resRef (Cons (Pair (mpPi mpar1) ( mpXi mpar1)))
               go mpar1
   mpar <- go init
   seed <- readIORef seedref
   reslist <- readIORef resRef
   return (seed, (mpar,reslist))
  S.put $ fst res
  let mp = fst $ snd res
  return $ (mp, map unPair $ unList $ snd $ snd res)

--  pi = pdf $ toList $ mpXi init -}

runMalaMPaccept :: Covariance a => a -> (Vector Double -> (Double,Vector Double))
         -> Int -> MalaPar ->  RIO (MalaPar, [(Double,Vector Double)])
runMalaMPaccept cov pdf nsam init = go init [] where
  go !mpar !xs
    | (round $ mpAccept mpar) >= nsam = do
         io $ putStrLn $ "MALA accept = "++show (mpAccept mpar/mpCount mpar)
         io $ putStrLn $ "MALA sigma = "++show (mpSigma mpar)
         return (mpar, xs)
    | mpAccept mpar < 0.5 && mpCount mpar > 20 =
         return (mpar, [])
    | otherwise = do
         !mpar1 <- sample $ mala1 cov  pdf True mpar
         io $ putStr ((show (round $ mpAccept mpar)) ++".") >> hFlush stdout
         go mpar1 $ (mpPi mpar1, mpXi mpar1):xs

kmala = 5

{-runMalaRioESS ::  Matrix Double -> (Vector Double -> (Double,Vector Double))
                  ->  Int -> Vector Double -> RIO [Vector Double]
runMalaRioESS cov pdf want_ess xi = do
    let (p0,grad0) = pdf xi
    let sigma0 = 1.5 / (realToFrac $ dim xi) -- determined empirically
    let mp0 = MalaPar xi p0 grad0 sigma0 0 0 False
        nsam0 = want_ess*20
    io $ putStrLn $ "initial sigma = "++show (mpSigma mp0)
    (mp1, xs1) <- runMalaMP cov pdf nsam0 mp0 []
    let have_ess = min (mpAccept mp1) $ calcESSprim $ map snd xs1
    if have_ess > realToFrac want_ess
       then return $ map snd xs1
       else do let need_ess =  max 1 $ realToFrac want_ess - have_ess
                   samples_per_es = realToFrac nsam0/have_ess
                   to_do = round $ samples_per_es * need_ess
               (mp2, xs2) <- runMalaMP cov pdf to_do mp1 xs1
               return  $ map snd  xs2 -}

runMalaRioCodaESS ::  Covariance a => a-> (Vector Double -> (Double,Vector Double))
                  ->  Int -> Double -> Vector Double -> RIO [Vector Double]
runMalaRioCodaESS cov pdf want_ess truncN xi = do
    let (p0,grad0) = pdf xi
    let sigma0 = 1.5 / (realToFrac $ dim xi) -- determined empirically
    let mp0 = MalaPar xi p0 (trunc truncN grad0) sigma0 0 0 False 0 truncN
        nsam0 = want_ess*1
    let converged mp  xs = do
         let have_ess = min (mpAccept mp) $ calcESSprim $ map snd xs
         io $ putStrLn $ "ESS=" ++show have_ess
         if have_ess > realToFrac want_ess
            then return $ map snd xs
            else do let need_ess =  max 1 $ realToFrac want_ess - have_ess
                        samples_per_es = realToFrac nsam0/have_ess
                        to_do = round $ samples_per_es * need_ess
                    io $ putStrLn $ "running converged for "++show to_do
                    (mp2, xs2) <- runMalaMP cov pdf to_do (mp {mpFreezeSigma = True}) xs
                    io $ putStrLn $ "All done"
                    return $ map snd xs2
    let go mp  n xs = do
            (mp2, xs2) <- runMalaMP cov  pdf n mp []
            let testres = mannWhitneyUtest TwoTailed 0.05 (U.fromList $ map fst xs2)
                                            (U.fromList $ map fst xs)

            if testres/= Just NotSignificant
               then do io$ putStrLn $ "not converged: "++show testres++" at "++show (mpPi mp2)
                       go mp2  (round $ realToFrac n*2) xs2
               else do io$ putStrLn "converged!"
                       converged mp2  (xs2++xs)

    io $ putStrLn $ "initial sigma = "++show (mpSigma mp0)
    (mp1, xs1) <- runMalaMPaccept cov  pdf nsam0 mp0
    case xs1 of
            [] -> return []
            _ -> go mp1  (round $ mpCount mp1) xs1

{-    case (spoon (invlndet cov), mbCholSH cov) of
       (Just (covInv, (lndt,_)), Just covChol) -> go_rest (covInv, covChol)
       _ -> let cov' = PDF.posdefify cov in
            case (spoon (invlndet cov'), mbCholSH cov') of
              (Just (covInv, (lndt,_)), Just covChol) -> go_rest (covInv, covChol)
              _ ->  do io $ putStrLn "non-invertible covariance matrix"
                       return [] -}

runMalaRioSimple ::  Covariance a => a -> (Vector Double -> (Double,Vector Double))
                  ->  Int -> Double -> Int -> Vector Double -> RIO [Vector Double]
runMalaRioSimple cov pdf samples maxGrad thinN xi = do
    let (p0,grad0) = pdf xi
    let sigma0 = 2.7e-4 --1.5 / (realToFrac $ dim xi) -- determined empirically
    let mp0 = MalaPar xi p0 (trunc maxGrad grad0) sigma0 0 0 False thinN maxGrad

    (mp2, xs2) <- runMalaMP cov  pdf samples mp0 []
    io $ putStrLn $ "All done"
    return $ map snd xs2

calcCovariance :: Vector Double ->
                  Vector Double ->
                  (Vector Double -> (Double,Vector Double)) ->
                  (Vector Double -> Double) ->
                  Either (Vector Double) (Matrix Double, Matrix Double, Matrix Double)
calcCovariance vinit vnear postgrad posterior = finalcov where
   ndim = dim vinit
   finalcov
     | ndim > 0 -- > 20000 --FIXME
        = Left $ calcFDindepVars vinit vnear posterior --Left $ iCov vinit --
     | otherwise
        = hessToCov (calcFDhess vinit vnear postgrad) Nothing

iCov v = VS.replicate (VS.length v) 1 -- in (m,m,m)

calcFDindepVars v v' post = trace ("FDVars = "++show (VS.take 10 vars))  $ vars where
   hv =  mapVector (*1e-4) v --mapVector (max 1e-9 .  abs) $ v - v'
   n = dim v
   postv = post v
   postPlus i = post $ v VS.// [(i,v @>i + hv @> i)]
   postMinus i = post $ v VS.// [(i,v @>i - hv@> i)]
   vars = buildVector n fvar
   fvar i = negate $ recip $ (postPlus i - 2*postv + postMinus i)/((hv @> i)*(hv @> i))

calcFDhess v v' postgrad = hess2 where
   grad =  snd . postgrad
   gradi i = (@>i) . grad
   hv = mapVector (max 1e-9 . abs) $ v - v'
   n = dim v
   gradv = grad v
   grads =  fromRows $ flip map [0..(n-1)] $ \i ->
               grad (v VS.// [(i,v @>i + hv @> i)])
   fhess (i,j) | i<j = 0
               | otherwise = (grads @@>(j,i) - (gradv @> i))
                                       /(2*(hv @> j)) +
                             (grads @@>(i,j)- (gradv @> j))
                                       /(2*(hv @> i))

--   hess3 = scale (recip $ realToFrac n) $ sum $ map outerSelf $ grads

   hess1 = buildMatrix n n fhess
   hess2 = buildMatrix n n $ \(i,j) -> if i>=j then hess1 @@> (i,j)
                                               else hess1 @@> (j,i)

outerSelf v= v `outer` v

hessToCov hess mOriginal =
   let mTryPosdefify = case mOriginal of
            Nothing -> hessToCov (PDF.posdefify hess) (Just hess)
            Just horiginal -> let ds =  mapVector (negate . recip)
                                         $ takeDiag horiginal
                              in trace ("USING DIAGS"++show (VS.take 10 ds)) $ Left $ ds
   in
   case spoon $ inv $ negate $ hess of
     Just cov -> case mbCholSH cov of
                   Just cholm ->  trace ("invert success:"++show (VS.take 10 $ takeDiag cov) ++ "det="++show (det cov)) $ Right (cov, negate hess, cholm)
                   Nothing -> trace ("chol fail") mTryPosdefify
     Nothing -> trace ("inv fail") mTryPosdefify
                  {- Just cov -> case mbCholSH cov of
                                Just cholm ->  Right (cov, negate hessToCov,cholm)
                                Nothing -> case mbCholSH $ PDF.posdefify cov of
                                            Just cholm -> Right (cov, negate hess,cholm)
                                            Nothing -> vars
                  Nothing -> vars -}

acceptSM ampar  | mpCount ampar == 0 = "0/0"
               | otherwise = printf "%.3g" (rate::Double) ++ " ("++show yes++"/"++show total++")" where
   rate = realToFrac (yes) / realToFrac (total)
   yes = mpAccept ampar
   total = mpCount ampar


-- program: S.A.R.A.H. jam simulator
-- author: Maurício C. Antunes
-- e-mail: mauricio.antunes@gmail.com
-- license: public domain

module Main where

import Control.Applicative
import Prelude

import Data.Maybe
import Graphics.UI.Gtk
import Graphics.Rendering.Cairo
import Control.Monad
import Data.IORef
import Data.List
import Data.Time
import Data.Complex

-- Constants

accelerator = 0.7*carSize :: Double
brake = 10*accelerator:: Double
carSize = 2*pi/59 :: Double
responseTime = 0.24 :: Double
drawSide = 5/2 :: Double

-- A few conveniences

eventWindowSize = do
    dr <- eventWindow
    w <- liftIO $ drawWindowGetWidth dr
    h <- liftIO $ drawWindowGetHeight dr
    return $ if w*h > 1
        then (fromIntegral w, fromIntegral h)
        else (1,1)

eventPolarCoordinates = do
    (w,h) <- eventWindowSize
    (x,y) <- eventCoordinates
    let (origX, origY) = (w/2, h/2)
    let (scaleX, scaleY) = (drawSide/w, drawSide/h)
    let (x',y') = (scaleX*(x-origX), scaleY*(y-origY))
    let (radius,theta) = polar $ x' :+ y'
    return $ (radius,theta)

getAndSet :: a -> IO (IO a, a -> IO ())
getAndSet a = do
    ior <- newIORef a
    let get = readIORef ior
    let set = writeIORef ior
    return (get,set)

diffTime :: UTCTime -> UTCTime -> Double
diffTime = (realToFrac .) . diffUTCTime

moveToLineTo :: Double -> Double
 -> Double -> Double -> Render ()
moveToLineTo a b c d = moveTo a b >> lineTo c d

-- Car list handling

-- Each car is represented by a pair of Doubles. The first
-- Double is its position in a circular road, represented by
-- an angle. The second is its angular velocity. The general
-- idea behind the simulation is that in a list of cars each
-- one will try to keep a safe speed to avoid a crash in the
-- event of a sudden brake of the next car.

newCarList nCars = take nCars $ zip [0,2*pi/nCars'..] (repeat 0)
    where nCars' = fromIntegral nCars

-- This resizes car lists by copying or keeping those
-- at lower speeds.

newCarListFromList nCars [] = newCarListFromList nCars [(0,0)]
newCarListFromList nCars list = sortBy ((. fst).(compare . fst)) $
    take nCars $ cycle $ sortBy ((. snd).(compare . snd)) list

-- Safe speed for car, given data from itself and the next
-- and, possibly, a forced (by the user) jam. Speed changes
-- are limited by accelerator and brake maxima.

newSpeed dt jam (p1,s1) (p2,s2) = min cv $ max bv $ ds - br
    where
        pd = (p2-p1-carSize) - responseTime*(s2-s1)
        pj = maybe pd ((subtract $ carSize/2)
         . (until (>0) (+2*pi)) . (subtract p1)) jam
        dd = brake*(max 0 $ min pd pj)
        br = brake*responseTime
        ds = sqrt $ br^2 + 2*dd
        cv = s1 + accelerator*dt
        bv = s1 - brake*dt

-- Update positions and speeds based on a timestep and maybe
-- taking a forced congestion into account

updateCarList _ _ [] = []
updateCarList timestep jam list = zip newPositions' newSpeeds
    where
        fakeCar = (p+2*pi,s) where (p,s) = head list
        newSpeeds = zipWith ns list (tail list ++ [fakeCar])
            where ns = newSpeed timestep jam
        newPositions = zipWith3 mean fsts snds newSpeeds
            where
                mean a b c = a + timestep*(b+c)/2
                fsts = map fst list
                snds = map snd list
        newPositions' = map (subtract base) newPositions
        base = (*(2*pi)) $ fromIntegral $ floor $ (/ (2*pi)) $
            head newPositions

about = do
    ad <- aboutDialogNew
    set ad [ aboutDialogName := "S.A.R.A.H."
           , aboutDialogVersion := "1.0"
           , aboutDialogAuthors := ["Maurício C. Antunes "
                                ++ "<mauricio.antunes@gmail.com>"]
           , aboutDialogComments := "Software Automation of "
                                ++ "Road Automobile Headache"]
    dialogRun ad
    widgetDestroy ad

main :: IO ()
main = do

    initGUI

    mainWindow <- windowNew
    drawingArea <- drawingAreaNew

    (getTimeStamp,setTimeStamp) <- getCurrentTime >>= getAndSet
    (getCars,setCars) <- getAndSet $ newCarList 20
    (getJam,setJam) <- getAndSet Nothing
    (getTimeoutId,setTimeoutId) <- getAndSet Nothing

    -- If 'resume' is called, 'step' will be called at small
    -- timesteps to update car data. If 'pause' is called, 'step'
    -- calls are stoped.  'resume' is called at program startup,
    -- and then the pause button alternates 'resume' and 'pause'.

    let step = do
         time <- getCurrentTime
         dt <- getTimeStamp >>= return . (diffTime time)
         setTimeStamp time
         liftM2 (updateCarList dt) getJam getCars >>= setCars
    let pause = do
         maybe (return ()) timeoutRemove =<< getTimeoutId
         setTimeoutId Nothing
    let resume = do
         setTimeoutId . Just =<< flip timeoutAdd 33
          (step >> widgetQueueDraw drawingArea >> return True)
         getCurrentTime >>= setTimeStamp

    -- The elements of the graphic interface are the set of
    -- buttons, the scale to set the number of cars and the
    -- car track. They are named as 'buttons', 'howMany' and
    -- 'track'. Each of them contains other widgets inside, but
    -- there's no reason to expose their names to the main IO.

    buttons <- do

        qr <- buttonNewFromStock stockClear
        on qr buttonActivated $ do
            (liftM length) getCars >>= setCars . newCarList
            getCurrentTime >>= setTimeStamp
            widgetQueueDraw drawingArea

        qp <- toggleButtonNewWithLabel stockMediaPause
        buttonSetUseStock qp True
        on qp toggled $ do
            p <- toggleButtonGetActive qp
            case p of
                True -> pause
                False -> resume

        qa <- buttonNewFromStock stockAbout
        on qa buttonActivated $ about

        qq <- buttonNewFromStock stockQuit
        on qq buttonActivated (do
                       widgetDestroy mainWindow
                       mainQuit)

        bb <- hButtonBoxNew
        containerAdd bb qr
        containerAdd bb qp
        containerAdd bb qa
        containerAdd bb qq
        return bb

    howMany <- do

        sc <- vScaleNewWithRange 1 40 1
        after sc valueChanged $ do
            v <- liftM floor $ rangeGetValue sc
            c <- getCars
            setCars $ newCarListFromList v c
            widgetQueueDraw drawingArea

        scaleSetValuePos sc PosTop
        scaleSetDigits sc 0
--        rangeSetUpdatePolicy sc UpdateDiscontinuous
        rangeSetValue sc =<< liftM (fromIntegral . length) getCars

        al <- alignmentNew 0.5 0.5 0 1
        alignmentSetPadding al 15 15 15 15
        containerAdd al sc
        return al

    track <- do

        let dr = drawingArea
        widgetAddEvents dr [PointerMotionMask]

        on dr motionNotifyEvent $ do
            (r,t) <- eventPolarCoordinates
            liftIO $ if (0.8<r && r<1.2)
                then setJam (Just t)
                else setJam Nothing
            liftIO $ widgetQueueDraw dr
            return True

        on dr leaveNotifyEvent $ liftIO $
            setJam Nothing >> return True

        on dr draw $ do
            w <- liftIO $ (fromIntegral <$> widgetGetAllocatedWidth dr)
            h <- liftIO $ (fromIntegral <$> widgetGetAllocatedHeight dr)
            jam <- liftIO getJam
            cars <- liftIO getCars
            translate (w/2) (h/2)
            scale (w/drawSide) (h/drawSide)
            road2render jam cars
            -- return True

        af <- aspectFrameNew 0.5 0.5 (Just 1)
        frameSetShadowType af ShadowNone
        containerAdd af dr
        return af

    -- 'layout' is a widget that contains all interface elements
    -- properly arranged.

    layout <- do
        vb <- vBoxNew False 0
        hb <- hBoxNew False 0
        boxPackStart vb track PackGrow 0
        boxPackStart vb buttons PackNatural 0
        boxPackStart hb howMany PackNatural 0
        boxPackStart hb vb PackGrow 0
        return hb

    set mainWindow [ windowTitle := "S.A.R.A.H."
                   , windowDefaultWidth := 400
                   , windowDefaultHeight := 400 ]
    on mainWindow objectDestroy mainQuit
    containerAdd mainWindow layout
    widgetShowAll mainWindow

    resume

    mainGUI

-- As the name says, this takes road info, in the form of a
-- possible jam and a list of cars, and make it into a Cairo
-- render.  Road will have radius 1.

road2render :: Maybe Double -> [(Double,Double)] -> Render ()
road2render jam cars = do
    newPath
    setSourceRGB 0 0 0
    drawRoad
    when (isJust jam) drawJam
    setSourceRGBA 0 0 0 0.55
    let cars' = map fst cars
    let rotations = zipWith subtract (0:cars') cars'
    sequence_ $ map ((>> drawCar) . rotate) rotations
 where
    drawRoad = setLineWidth 0.01 >> setDash [2*pi/34,2*pi/34]
     (pi/34) >> arc 0.0 0.0 1.0 0.0 (2*pi) >> stroke
    drawJam = setLineWidth 0.005 >> setDash [0.03,0.02] 0.04 >>
     save >> rotate (fromJust jam) >> moveToLineTo 0.8 0 1.2
     0 >> stroke >> setDash [] 0 >> moveToLineTo 0.8 (-0.015)
     0.8 0.015 >> moveToLineTo 1.2 (-0.015) 1.2 0.015 >> stroke
     >> restore
    drawCar = arc 1 0 (carSize/2) 0 (2*pi) >> fill