module Lib where

import qualified Numeric.LinearAlgebra as Matrix
import qualified Numeric.LinearAlgebra.Data as HM

import Control.Monad.Random
import Control.Monad.Random.Class

import Control.DeepSeq

import Control.Parallel.Strategies (parListChunk, rpar, rdeepseq, rseq, runEval, runEvalIO)

import Data.List (foldl', sort, sortBy, maximumBy, minimumBy)

import GHC.Conc (numCapabilities)

import qualified Data.ByteString.Lazy as BS (writeFile, readFile)
import Data.Serialize (encodeLazy, decodeLazy)

data Player = X | O deriving (Show, Ord, Eq, Enum, Read)

type Grid   = (Player, [Maybe Player])
data Result = Undecided | Tie | Win Player deriving Show

data Tree a = Node a [Tree a] deriving (Show)

showGrid :: Grid -> String
showGrid (_, xs) = " " ++ x00 ++ " | " ++ x01 ++ " | " ++ x02 ++ "\n" ++
                   "-----------" ++ "\n" ++
                   " " ++ x10 ++ " | " ++ x11 ++ " | " ++ x12 ++ "\n" ++
                   "-----------" ++ "\n" ++
                   " " ++ x20 ++ " | " ++ x21 ++ " | " ++ x22 ++ "\n"
    where
        [x00, x01, x02, x10, x11, x12, x20, x21, x22] =
            map (\x -> case x of
                          Nothing -> " "
                          Just X  -> "X"
                          Just O  -> "O") xs

moves :: (Grid, Result) -> [(Grid, Result)]
moves (g@(_, xs), Undecided) = gs
    where
        l  = length $ filter id $ map isNothing xs
        gs = map (\i -> move i g) [0..l-1]
moves _ = []

gameTree :: (Grid, Result) -> Tree (Grid, Result)
gameTree n = Node n fs
    where
        fs = map gameTree $ moves n

maxMin :: RandomGen g => Tree (Grid, Result) -> Rand g (Maybe ((Grid, Result), Int))
maxMin (Node _ []) = return Nothing
maxMin (Node n fs) = case mms of
                        [] -> return Nothing
                        _  -> do
                               i <- getRandomR (0, length maxScores - 1)
                               return (Just $ maxScores !! i)
    where
        mms = map minMax' fs
        maxScore  = maximum $ map snd mms
        maxScores = filter ((==maxScore) . snd) mms

maxMin' :: Tree (Grid, Result) -> ((Grid, Result), Int)
maxMin' (Node n@(_, Tie) []) = (n, 0)
maxMin' (Node n@(_, Win O) []) = (n, -1)
maxMin' (Node n@(_, Win X) []) = (n, 1)
maxMin' (Node (_, _) []) = error "Undecided Leaf"
maxMin' (Node n fs) = (n, snd $ maximumBy (\(_, r1) (_, r2) -> r1 `compare` r2) $ map minMax' fs)

minMax :: RandomGen g => Tree (Grid, Result) -> Rand g (Maybe ((Grid, Result), Int))
minMax (Node _ []) = return Nothing
minMax (Node n fs) = case mms of
                        [] -> return Nothing
                        _  -> do
                               i <- getRandomR (0, length minScores - 1)
                               return (Just $ minScores !! i)
    where
        mms = map maxMin' fs
        minScore  = minimum $ map snd mms
        minScores = filter ((==minScore) . snd) mms

minMax' :: Tree (Grid, Result) -> ((Grid, Result), Int)
minMax' (Node n@(_, Tie) []) = (n, 0)
minMax' (Node n@(_, Win O) []) = (n, -1)
minMax' (Node n@(_, Win X) []) = (n, 1)
minMax' (Node (_, _) []) = error "Undecided Leaf"
minMax' (Node n fs) = (n, snd $ minimumBy (\(_, r1) (_, r2) -> r1 `compare` r2) $ map maxMin' fs)

moveMM :: RandomGen g => (Grid, Result) -> Rand g (Maybe (Grid, Result))
moveMM gr@((X, _), _) = fmap (fmap fst) . maxMin $ gameTree gr
moveMM gr@((O, _), _) = fmap (fmap fst) . minMax $ gameTree gr

empty :: Grid
empty = (X, take 9 $ repeat Nothing)

chunksOf :: Int -> [a] -> [[a]]
chunksOf n [] = []
chunksOf n xs = take n xs : chunksOf n (drop n xs)

rows :: [a] -> [[a]]
rows xs      = rs ++
               map (\i -> map (!!i) rs) [0,1,2] ++
               pure (zipWith (!!) rs [0,1,2]) ++
               pure (zipWith (!!) rs [2,1,0])
    where
        rs = chunksOf 3 xs

allSame :: Eq a => [a] -> (Bool, Maybe a)
allSame []     = (True, Nothing)
allSame (x:xs) = case all (==x) xs of
                   True  -> (True, Just x)
                   False -> (False, Nothing)

isNothing :: Maybe a -> Bool
isNothing Nothing = True
isNothing _       = False

gridResult :: Grid -> Result
gridResult (_, xs)
        | not (null winners) = Win $ head winners
        | undecided          = Undecided
        | otherwise          = Tie
    where
        rs = rows xs
        rresults = map rowResult rs
        winners = map (\(Just x) -> x) $
                  filter (not . isNothing) $
                  map (\res -> case res of
                                 Win p -> Just p
                                 _     -> Nothing)
                  rresults
        undecided = any (\res -> case res of
                                 Undecided -> True
                                 _         -> False)
                  rresults

rowResult :: [Maybe Player] -> Result
rowResult r = case any isNothing r of
                True  -> Undecided
                False -> case allSame r of
                           (True, Just (Just p)) -> Win p
                           _                     -> Tie

emptyFields :: [Maybe a] -> [Int]
emptyFields xs = map snd $
                 filter (isNothing . fst) $
                 zip xs [0..]

move :: Int -> Grid -> (Grid, Result)
move m (p, xs) = ((nextPlayer, nextGrid), res)
    where
        nextPlayer = case p of
                       X -> O
                       O -> X
        empties = emptyFields xs
        numEmpties = length empties
        m' = m `mod` numEmpties
        pos = empties !! m'
        nextGrid = take pos xs ++ [Just p] ++ drop (pos+1) xs
        res = gridResult (nextPlayer, nextGrid)

moveAI :: DenseData -> Grid -> (Grid, Result)
moveAI net g = move m g
    where
        l = fromGrid g
        m = predict (net, [relu, softmax]) l

battle :: DenseData -> DenseData -> Maybe Player
battle pX pO = battle' pX pO empty

battle' :: DenseData -> DenseData -> Grid -> Maybe Player
battle' pX pO g@(p, _) = case moveAI ai g of
                            (_, Tie)        -> Nothing
                            (_, Win w)      -> Just w
                            (g', Undecided) -> battle' pX pO g'
    where
        ai = case p of
                X -> pX
                O -> pO

battleMM' :: RandomGen g => (Grid, Result) -> Rand g [(Grid, Result)]
battleMM' g = do
                g' <- moveMM g
                case g' of
                  Nothing  -> return [g]
                  Just g'' -> fmap (g:) $ battleMM' g''

battleMM :: RandomGen g => Rand g [(Grid, Result)]
battleMM = battleMM' (empty, Undecided)

loses :: DenseData -> DenseData -> Bool
loses p1 p2 = case w1 of
                Just O -> True
                _      -> case w2 of
                            Just X -> True
                            _      -> False
    where
        w1 = battle p1 p2
        w2 = battle p2 p1

beats :: DenseData -> DenseData -> Bool
beats p1 p2 = case w1 of
                Just X -> case w2 of
                            Nothing -> True
                            Just O  -> True
                            _       -> False
                _      -> False
    where
        w1 = battle p1 p2
        w2 = battle p2 p1

type Layer  = HM.Vector HM.R
type Bias   = HM.Vector HM.R
type Weight = HM.Matrix HM.R
type Activation = Layer -> Layer

type DenseData = [(Bias, Weight)]
type DenseNet  = (DenseData, [Activation])

type CrossingData = [(Bias, Bias, Weight, Weight)]

toFile :: String -> DenseData -> IO ()
toFile fileName net = BS.writeFile fileName $
                      encodeLazy $
                      fmap (\(b, w) -> (HM.toList b, HM.toLists w)) net

fromFile :: String -> IO DenseData
fromFile fileName = fmap (fmap (\(b, w) -> (HM.fromList b, HM.fromLists w)) . (\(Right x) -> x) . decodeLazy) $
                    BS.readFile fileName

randBias :: RandomGen g => Int -> Rand g Bias
randBias n = (pure . HM.fromList . take n) =<< getRandomRs (-1, 1)

randWeight :: RandomGen g => Int -> Int -> Rand g Weight
randWeight m n = (pure . (m HM.>< n)) =<< getRandomRs (-1, 1)

repeatM :: Monad m => Int -> m a -> m [a]
repeatM 0 m = pure []
repeatM n m =
    do
        x  <- m
        xs <- repeatM (n-1) m
        pure (x:xs)

toBinary :: (Num a, Ord a) => a -> a
toBinary x
    | x < 0 = 0
    | otherwise = 1

invert :: Num a => a -> a
invert x = 1-x

randCrossing :: RandomGen g => Rand g CrossingData
randCrossing =
  do
    ~[(b1, w01), (b2, w12)] <- randPlayer
    let b1' = HM.cmap toBinary b1
    let b2' = HM.cmap toBinary b2
    let w01' = HM.cmap toBinary w01
    let w12' = HM.cmap toBinary w12
    return [(b1', HM.cmap invert b1', w01', HM.cmap invert w01'),
            (b2', HM.cmap invert b2', w12', HM.cmap invert w12')]

randPlayer :: RandomGen g => Rand g DenseData
randPlayer =
  do
    b1 <- randBias 18
    b2 <- randBias 9
    w01 <- randWeight 18 18
    w12 <- randWeight 9 18
    return [(b1, w01), (b2, w12)]

fromGrid :: Grid -> Layer
fromGrid (_, xs) = HM.fromList $ lX ++ lO
    where
        lX = map (\x -> case x of
                          Just X -> 1
                          _      -> 0)
             xs
        lO = map (\x -> case x of
                          Just O -> 1
                          _      -> 0)
             xs

dense :: Bias -> Weight -> Activation -> Layer -> Layer
dense b w a l = a $ b + (w Matrix.#> l)

relu :: Activation
relu = HM.cmap relu'
    where
        relu' x = if x < 0 then 0 else x

softmax :: Activation
softmax l = HM.cmap (/s) exps
    where
        m    = HM.maxElement l
        exps = HM.cmap (\x -> exp (x - m)) l
        s    = sum $ HM.toList exps

runDenseNet :: DenseNet -> Layer -> Layer
runDenseNet (dat, acts) l =
    foldl' (\l' ((b, w), a) -> dense b w a l') l $ zip dat acts

predict :: DenseNet -> Layer -> Int
predict net = HM.maxIndex . runDenseNet net

select :: [DenseData] -> (Int, [DenseData])
select ps = (maximum scores,
             take n' $
             map fst $
             sortBy (\(_, score1) (_, score2) -> score2 `compare` score1) $
             zip ps scores)
    where
        scores = runEval $
                 parListChunk (n `div` numCapabilities) rdeepseq $
                 fmap (\x -> length $ filter not $ fmap (loses x) ps) ps
        n = length ps
        n' = n `div` 100

breed2' :: DenseData -> DenseData -> CrossingData -> DenseData
breed2' m f cd = map (\((b, w), (b', w'), (cb, cb', cw, cw')) -> (b*cb + b' * cb', w*cw + w'*cw')) $ zip3 m f cd

breed2 :: RandomGen g => DenseData -> DenseData -> Rand g [DenseData]
breed2 m f =
    do
        cds <- repeatM 10 randCrossing
        return $ fmap (breed2' m f) cds

breed :: RandomGen g => [DenseData] -> Rand g [DenseData]
breed ps = fmap concat $ sequence $ breed2 <$> ps <*> ps

initPopulation :: RandomGen g => Int -> Rand g [DenseData]
initPopulation n = repeatM n randPlayer

someFunc :: IO ()
someFunc = do
            ps <- evalRandIO $ initPopulation 1000
            go ps
    where
        go ps = let (score, fittest) = select ps
                in do
                     putStrLn $ show score ++ "\t(" ++ show (length ps) ++ ")"
                     if score == 1000
                       then return () >> (toFile "weights.dat" $ head fittest)
                       else go =<< evalRandIO (breed fittest)


module Numeric.Network where
import Data.Kind
import Dependent.Size
import Numeric.Layer
import Numeric.Vector.Sized
import Numeric.Loss

infixr 0 :~>

type family NetworkOutput (layers :: [Type]) where
  NetworkOutput '[x] = Outputs x
  NetworkOutput (x ': y ': rest) = NetworkOutput (y ': rest)

type family NetworkInput (layers :: [Type]) where
  NetworkInput (x ': _) = Inputs x

data Network (layers :: [Type]) where
  Last :: Layer x => x -> Network '[x]
  (:~>) :: (Layer x, Outputs x ~ Inputs y) => x -> Network (y:xs) -> Network (x : y : xs)

instance (Show x) => Show (Network '[x]) where
  showsPrec n (Last g) = showsPrec n g
instance (Show x, Show (Network (y:xs))) => Show (Network (x:y:xs)) where
  showsPrec n (layer :~> net) = showsPrec n layer . showString " :~> " . showsPrec n net

infixr 0 :|>
data Tapes (layers :: [Type]) where
  LastTape :: Layer x => Tape x -> Tapes '[x]
  (:|>) :: (Layer x, Outputs x ~ Inputs y) => Tape x -> Tapes (y:xs) -> Tapes (x:y:xs)

instance (Show (Tape x)) => Show (Tapes '[x]) where
  showsPrec n (LastTape g) = showsPrec n g
instance (Show (Tape x), Show (Tapes (y:xs))) => Show (Tapes (x:y:xs)) where
  showsPrec n (tape :|> tapes) = showsPrec n tape . showString " :|> " . showsPrec n tapes

infixr 0 :<|
data Gradients (layers :: [Type]) where
  LastGradient :: Layer x => Gradient x -> Gradients '[x]
  (:<|) :: (Layer x, Outputs x ~ Inputs y) => Gradient x -> Gradients (y:xs) -> Gradients (x : y : xs)

instance (Show (Gradient x)) => Show (Gradients '[x]) where
  showsPrec n (LastGradient g) = showsPrec n g
instance (Show (Gradient x), Show (Gradients (y:xs))) => Show (Gradients (x:y:xs)) where
  showsPrec n (grad :<| grads) = showsPrec n grad . showString " :<| " . showsPrec n grads

predict
  :: Network layers
  -> SizedArray (NetworkInput layers)
  -> SizedArray (NetworkOutput layers)
predict (Last layer ) x = fst (forward layer x)
predict (layer:~>net) x = predict net y where y = fst (forward layer x)

predict'
  :: Network layers
  -> SizedArray (NetworkInput layers)
  -> (SizedArray (NetworkOutput layers), Tapes layers)
predict' (Last layer) x = let (y, tape) = forward layer x in (y, LastTape tape)
predict' (layer:~>net) x =
  let (y, tape) = forward layer x
  in  let (z, tapes) = predict' net y in (z, tape :|> tapes)

backprop
  :: Network layers
  -> Tapes layers
  -> SizedArray (NetworkOutput layers)
  -> (Gradients layers, SizedArray (NetworkInput layers))
backprop (Last layer) (LastTape tape) dy =
  let (grad, dx) = backward layer tape dy in (LastGradient grad, dx)
backprop (layer:~>net) (tape:|>tapes) dz =
  let (grads, dy) = backprop net tapes dz
  in  let (grad, dx) = backward layer tape dy in (grad :<| grads, dx)

backpropagation
  :: (Loss (NetworkOutput layers) loss)
  => Network layers
  -> loss
  -> SizedArray (NetworkInput layers)
  -> SizedArray (NetworkOutput layers)
  -> (Gradients layers, SizedArray (NetworkInput layers))
backpropagation net loss x target = backprop net tapes dy
  where
    (y, tapes) = predict' net x
    dy         = lossDerivative loss y target

learn :: Network layers -> Gradients layers -> Network layers
learn (Last layer ) (LastGradient grad) = Last (applyGradient layer grad)
learn (layer:~>net) (grad:<|grads     ) = applyGradient layer grad :~> learn net grads

train
  :: (Loss (NetworkOutput layers) loss)
  => Network layers
  -> loss
  -> SizedArray (NetworkInput layers)
  -> SizedArray (NetworkOutput layers)
  -> Network layers
train net loss x t = f $ backpropagation net loss x t
  where f (gradients, _) = learn net gradients

-- A network is a layer
instance Layer (Network layers) where
  type Inputs (Network layers) = NetworkInput layers
  type Outputs (Network layers) = NetworkOutput layers
  type Gradient (Network layers) = Gradients layers
  type Tape (Network layers) = Tapes layers
  forward = predict'
  backward = backprop
  applyGradient = learn


-- {-# LANGUAGE AllowAmbiguousTypes                      #-}
{-# LANGUAGE BangPatterns                             #-}
{-# LANGUAGE DataKinds                                #-}
{-# LANGUAGE DeriveGeneric                            #-}
{-# LANGUAGE FlexibleContexts                         #-}
{-# LANGUAGE FlexibleInstances                        #-}
{-# LANGUAGE GADTs                                    #-}
{-# LANGUAGE LambdaCase                               #-}
{-# LANGUAGE PartialTypeSignatures                    #-}
{-# LANGUAGE PatternSynonyms                          #-}
{-# LANGUAGE RankNTypes                               #-}
{-# LANGUAGE ScopedTypeVariables                      #-}
{-# LANGUAGE TupleSections                            #-}
{-# LANGUAGE TypeApplications                         #-}
{-# LANGUAGE TypeOperators                            #-}
{-# LANGUAGE ViewPatterns                             #-}
{-# OPTIONS_GHC -Wno-incomplete-patterns              #-}
{-# OPTIONS_GHC -Wno-orphans                          #-}
{-# OPTIONS_GHC -Wno-unused-top-binds                 #-}
{-# OPTIONS_GHC -fno-warn-orphans                     #-}
{-# OPTIONS_GHC -fno-warn-partial-type-signatures     #-}
{-# OPTIONS_GHC -fwarn-redundant-constraints          #-}
{-# OPTIONS_GHC -fplugin GHC.TypeLits.Extra.Solver    #-}
-- {-# OPTIONS_GHC -fplugin GHC.TypeLits.KnownNat.Solver #-}
-- {-# OPTIONS_GHC -fplugin GHC.TypeLits.Normalise       #-}

module Main where
import           Control.DeepSeq
import           Control.Exception
import           Control.Monad
import           Control.Monad.IO.Class
import           Control.Monad.Trans.Maybe
import           Control.Monad.Trans.State
import           Data.Bitraversable
import           Data.Foldable
import           Data.IDX
import           Data.List.Split
import           Data.Maybe
import           Data.Proxy
import           Data.Time.Clock
import           Data.Traversable
import           Data.Tuple
import           GHC.Generics                        (Generic)
import           GHC.TypeLits
import           GHC.TypeLits.Extra
import           Lens.Micro                          hiding ((&))
import           Mnist.Internal.Convolution
import           Numeric.Backprop                    hiding ((:&))
import           Numeric.LinearAlgebra.Static.Backprop                          -- hmatrix-backprop
import           Numeric.LinearAlgebra.Static.Vector (vecL, vecR, lVec, rVec)
import           Numeric.OneLiner
import           Text.Printf
import qualified Data.Vector                         as V
import qualified Data.Vector.Generic                 as VG
import qualified Data.Vector.Storable.Sized          as SVS
import qualified Data.Vector.Unboxed                 as VU
import qualified Numeric.LinearAlgebra               as HM  -- non-backprop hmatrix
import qualified Numeric.LinearAlgebra.Data          as HMD
import qualified Numeric.LinearAlgebra.Static        as HMS -- hmatrix with type-checked operations :D
import qualified System.Random.MWC                   as MWC
import qualified System.Random.MWC.Distributions     as MWC
import Debug.Trace

type Model p a b = forall z. Reifies z W
  => BVar z p
  -> BVar z a
  -> BVar z b

-- Custom data type for a tuple with strict values
data a :& b = !a :& !b
  deriving (Show, Generic)
infixr 2 :&

-- Do we actually need this? BVar already defines a similar instance
instance (NFData a, NFData b) => NFData (a :& b)
  where rnf (a :& b) = force a `seq` force b `seq` ()

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

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

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

instance (MWC.Variate a, MWC.Variate b, Num a, Num b) => MWC.Variate (a :& b) where
    uniform g = (:&) <$> MWC.uniform g <*> MWC.uniform g
    uniformR (l, h) g = (\x -> x * (h - l) + l) <$> MWC.uniform g

instance KnownNat n => MWC.Variate (R n) where
    uniform g = HMS.randomVector <$> MWC.uniform g <*> pure HMS.Uniform
    uniformR (l, h) g = (\x -> x * (h - l) + l) <$> MWC.uniform g

instance (KnownNat m, KnownNat n) => MWC.Variate (L m n) where
    uniform g = HMS.uniformSample <$> MWC.uniform g <*> pure 0 <*> pure 1
    uniformR (l, h) g = (\x -> x * (h - l) + l) <$> MWC.uniform g

-- So if we have a BVar z (a :& b) (a BVar containing a tuple), then matching on (x :&& y) will give us x :: BVar z a and y :: BVar z b.
pattern (:&&) :: (Backprop a, Backprop b, Reifies z W)
              => BVar z a -> BVar z b -> BVar z (a :& b)
pattern x :&& y <- (\xy -> (xy ^^. t1, xy ^^. t2)->(x, y))
  where
    (:&&) = isoVar2 (:&) (\case x :& y -> (x, y))
{-# COMPLETE (:&&) #-}

-- Just some lenses to help us get BVars out of the tuple
t1 :: Lens (a :& b) (a' :& b) a a'
t1 f (x :& y) = (:& y) <$> f x
{-# INLINE t1 #-}

t2 :: Lens (a :& b) (a :& b') b b'
t2 f (x :& y) = (x :&) <$> f y
{-# INLINE t2 #-}

feedForward
    :: (KnownNat i, KnownNat o)
    => Model (L o i :& R o) (R i) (R o)
feedForward (w :&& b) x = w #> x + b
{-# INLINE feedForward #-}

convolution
  :: ( KnownNat kernelSize
     , KnownNat filters
     , KnownNat stride
     , KnownNat inputRows
     , KnownNat inputCols
     , KnownNat (kernelSize * kernelSize * channels)
     , KnownNat (inputRows * channels)
     , KnownNat (((Div (inputRows - filters) stride) + 1) * filters)
     , KnownNat (((Div (inputCols - filters) stride) + 1))
     )
     -- | These proxies introduce type variables into the scope
     --   so GHC can do type level arithmetic to infer the Model types
     -- | Example: convolution (Proxy @3) (Proxy @12) (Proxy @1) (Proxy @28) (Proxy @28) (Proxy @1)
     --   will return a function of type Model (L 9 12) (L 28 28) (L 204 17)
     => Proxy kernelSize
     -> Proxy filters
     -> Proxy stride
     -> Proxy inputRows
     -> Proxy inputCols
     -> Proxy channels
     -> Model
        (L (kernelSize * kernelSize * channels) filters) -- TODO: Add bias
        (L (inputRows * channels) inputCols)
        (L (((Div (inputRows - filters) stride) + 1) * filters) ((Div (inputCols - filters) stride) + 1))
convolution k' fs' st' ix' iy' cs' =
  liftOp2 . op2 $ \kernel input ->

    -- natVal "reflects" the type level Nat into a term level Int.
    -- Example: for ix' (input rows) this works because we told the type checker that the function argument ix' is a Proxy of a type
    -- that has a KnownNat constraint. So the type level Nat is brought into scope when we pass Proxy @28 (which is sugar for Proxy :: Proxy 28)
    -- Hyperparameters as types. Neat.
    let ix = fromIntegral $ natVal ix'
        iy = fromIntegral $ natVal iy'
        kx = fromIntegral $ natVal k'
        ky = fromIntegral $ natVal k'
        sx = fromIntegral $ natVal st'
        sy = fromIntegral $ natVal st'
        ox = ((ix - (fromIntegral $ natVal fs')) `div` sx) + 1
        oy = ((iy - (fromIntegral $ natVal fs')) `div` sx) + 1
        ex = HMS.extract input
        ek = HMS.extract kernel

        -- Transform the 3D input matrix into a (W*H, K*K*D) matrix
        c  = vid2col kx ky sx sy ix iy ex
    in

    (
      -- FORWARD PASS
      -- This is the actual "convolution"
      let mt = c HM.<> ek
          -- Stretch the image back to the output dimensions
          r  = col2vid 1 1 1 1 ox oy mt
      in
          fromJust . HMS.create $ r

    , -- BACKWARD PASS (see backprop ops for more info)
      \dzdy ->
        let eo = HMS.extract dzdy

            -- what is this actually doing?
            -- should be taking output from forward pass and reshaping it
            -- so we can do convolutions via matrix mult
            vs = vid2col 1 1 1 1 ox oy eo

            -- TODO: Gradient for weights -- I cannot get this to work -- WHY?
            -- It currently only works when filters === kernelSize -- wat
            dW = HM.tr c HM.<> vs

            -- Gradient for input -- This seems to work fine
            -- convolve (via matrix mult) output with transposed weights matrix http://soumith.ch/ex/pages/2014/08/07/why-rotate-weights-convolution-gradient/
            dX' = vs HM.<> HM.tr ek
            dX = col2vid kx ky sx sy ix iy dX' -- stretch back into image dimensions
        in
            -- trace ("convolution backwards") (fromJust . HMS.create $ dW, fromJust . HMS.create $ dX)
            (fromJust . HMS.create $ dW, fromJust . HMS.create $ dX)
    )
{-# INLINE convolution #-}

flattenLayer :: (KnownNat o, KnownNat i, KnownNat (o * i), Reifies s W) => BVar s (L o i) -> BVar s (R (o * i))
flattenLayer = liftOp1 . op1 $ \input ->
  (
    let ex = HMS.extract input
        flattened = HMD.flatten ex
        result = fromJust . HMS.create $ flattened
      in
        -- trace "flatten forwards" result
        result
  ,
    \dzdy ->
      let ex = HMS.extract dzdy
          ei = HMS.extract input
        in
          -- trace ("flatten backwards") fromJust . HMS.create $ HMD.reshape (HMD.cols ei) ex
          fromJust . HMS.create $ HMD.reshape (HMD.cols ei) ex
  )
{-# INLINE flattenLayer #-}

-- TODO Add type annotation?
convLayer k fs st ix io cs p = vmap reLU . flattenLayer . convolution k fs st ix io cs p
{-# INLINE convLayer #-}

logistic :: Floating a => a -> a
logistic x = 1 / (1 + exp (-x))
{-# INLINE logistic #-}

feedForwardLog
    :: (KnownNat i, KnownNat o)
    => Model (L o i :& R o) (R i) (R o)
feedForwardLog wb = logistic . feedForward wb
{-# INLINE feedForwardLog #-}

reLU :: (Num a, Ord a) => a -> a
reLU x | x < 0     = 0
       | otherwise = x
{-# INLINE reLU #-}

feedForwardReLU
    :: (KnownNat i, KnownNat o)
    => Model (L o i :& R o) (R i) (R o)
feedForwardReLU wb = vmap reLU . feedForward wb
{-# INLINE feedForwardReLU #-}

softMax :: (KnownNat n, Reifies s W) => BVar s (R n) -> BVar s (R n)
softMax x = konst (1 / sumElements expx) * expx
   where
      expx = exp x
{-# INLINE softMax #-}

feedForwardSoftMax
    :: (KnownNat i, KnownNat o)
    => Model (L o i :& R o) (R i) (R o)
feedForwardSoftMax wb = softMax . feedForward wb
{-# INLINE feedForwardSoftMax #-}

crossEntropy
   :: (KnownNat n, Reifies s W)
   => R n
   -> BVar s (R n)
   -> BVar s Double
crossEntropy !targ !res = -(log res <.> constVar targ)
{-# INLINE crossEntropy #-}

netErr
   :: forall m n o p s. (KnownNat o, Reifies s W)
   => Model p (L m n) (R o)
   -> L m n
   -> R o
   -> BVar s p
   -> BVar s Double
netErr f !x !targ !p = crossEntropy targ $ f p $ auto x
{-# INLINE netErr #-}

trainModel
   :: forall m n o p. (KnownNat o, Backprop p, Fractional p)
   => Double                -- ^ learning rate
   -> Model p (L m n) (R o)
   -> p                     -- ^ initial params
   -> [(L m n, R o)]        -- ^ input and target pairs
   -> p                     -- ^ trained params
trainModel r f = foldl' (\p (x,y) -> p - realToFrac r * gradBP (netErr f x y) p)
{-# INLINE trainModel #-}

testNet
   :: forall m n o p. (KnownNat o)
   => Model p (L m n) (R o)
   -> p
   -> [(L m n, R o)]
   -> Double
testNet f !p !xs = sum (map (uncurry test) xs) / fromIntegral (length xs)
   where
      test :: L m n -> R o -> Double -- test if the max index is correct
      -- second argument here is using ViewPatterns extension of GHC
      test x (HMS.extract->t)
         | HM.maxIndex t == HM.maxIndex (HMS.extract r) = 1
         | otherwise                                    = 0
         where
            r :: R o
            r = evalBP2 f p x

-- Given two Models, we can define the composition of them both
-- as long as their input/output shapes match
(<~)
    :: (Backprop p, Backprop q)
    => Model  p       b c
    -> Model       q  a b
    -> Model (p :& q) a c
(f <~ g) (p :&& q) = f p . g q
infixr 8 <~
{-# INLINE (<~) #-}

-- The type wildcard here means GHC should infer the type of the model params
-- We could define it but then we'd have to update it manually every time we add a new "layer" here
-- Each layer's parameters (weights + bias) are tupled together so we can pattern match on them (see <~ operator above)
model :: Model _ (L 28 28) (R 10)
model =
   feedForwardSoftMax @100 @10
   <~ feedForwardLog @500 @100
   <~ feedForwardLog @2028 @500
   <~ convLayer (Proxy @3) (Proxy @3) (Proxy @1) (Proxy @28) (Proxy @28) (Proxy @1)
{-# INLINE model #-}

main :: IO ()
main = MWC.withSystemRandom $ \g -> do
   Just train <- loadMNIST "data/train-images-idx3-ubyte" "data/train-labels-idx1-ubyte"
   Just test  <- loadMNIST "data/t10k-images-idx3-ubyte"  "data/t10k-labels-idx1-ubyte"
   putStrLn "Loaded data."

   p0 <- MWC.uniformR (-0.5, 0.5) g
   -- print (show $ take 10 train)

   flip evalStateT p0 . forM_ [1..] $ \e -> do
      train' <- liftIO . fmap V.toList $ MWC.uniformShuffle (V.fromList train) g
      liftIO $ printf "[Epoch %d]\n" (e :: Int)

      forM_ ([1..] `zip` chunksOf batch train') $ \(b, chnk) -> StateT $ \ps0 -> do
         printf "(Batch %d)\n" (b :: Int)

         t0 <- getCurrentTime
         newP <- evaluate . force $ trainModel (rate ^ fromIntegral e) model ps0 chnk
         t1 <- getCurrentTime
         printf "Trained on %d points in %s.\n" batch (show (t1 `diffUTCTime` t0))

         let trainScore = testNet model newP chnk
             testScore  = testNet model newP test
         printf "Training error:   %.2f%%\n" ((1 - trainScore) * 100)
         printf "Validation error: %.2f%%\n" ((1 - testScore ) * 100)

         -- TODO: Serialize trained params so we can save/load them

         -- Because we are in the StateT monad, the next iteration of the loop
         -- will be passed newP (the updated parameters for the model)
         return ((), newP)
   where
      rate = 0.01
      batch = 30

loadMNIST
   :: FilePath
   -> FilePath
   -> IO (Maybe [(L 28 28, R 10)])
loadMNIST fpI fpL = runMaybeT $ do
   i <- MaybeT          $ decodeIDXFile       fpI
   l <- MaybeT          $ decodeIDXLabelsFile fpL
   d <- MaybeT . return $ labeledIntData l i
   r <- MaybeT . return $ for d (bitraverse mkImage mkLabel . swap)
   liftIO . evaluate $ force r
      where
         mkImage :: VU.Vector Int -> Maybe (L 28 28)
         mkImage = HMS.create . HMD.reshape 28 . VG.convert . VG.map (\i -> fromIntegral i / 255)
         mkLabel :: Int -> Maybe (R 10)
         mkLabel n = HMS.create $ HM.build 10 (\i -> if round i == n then 1 else 0)


module Statistics.Information.Mixed.TotalCorrelation where

import Data.Matrix
import qualified Statistics.Information.Continuous.Entropy as CE
import Statistics.Information.Continuous.MutualInfo
import qualified Statistics.Information.Continuous.TotalCorrelation as CTC
import qualified Statistics.Information.Discrete.Entropy as DE
import qualified Statistics.Information.Discrete.TotalCorrelation as DTC
import Statistics.Information.Mixed.Entropy
import Statistics.Information.Mixed.MutualInfo
import Statistics.Information.Utils.Matrix
import Statistics.Information.Utils.Random
import System.Random

ctcdc :: (RandomGen g, Eq a) => g -> Int -> Int -> Matrix a -> Matrix Double ->
         Double
ctcdc g k base xs ys = sum hxs - hx where
  hxs = [centropydc gcol k base col ys | (col, gcol) <- cols]
  hx = centropydc g2 k base xs ys
  (g1, g2) = split g
  cols = (map colVector (columns xs)) `zip` (splitN (ncols xs) g1)

ctccd :: (RandomGen g, Eq a) => g -> Int -> Int -> Matrix Double -> Matrix a ->
         Double
ctccd g k base xs ys = sum hxs - hx where
  hxs = [centropycd gcol k base col ys | (col, gcol) <- cols]
  hx = centropycd g2 k base xs ys
  (g1, g2) = split g
  cols = (map colVector (columns xs)) `zip` (splitN (ncols xs) g1)

corexcd_tcs :: (RandomGen g, Eq a) => g -> Int -> Int -> Matrix Double ->
               Matrix a -> Double
corexcd_tcs g k base xs ys = CTC.tc g1 k base xs - ctccd g2 k base xs ys where
  (g1, g2) = split g

corexcd_mis :: (RandomGen g, Eq a) => g -> Int -> Int -> Matrix Double ->
               Matrix a -> Double
corexcd_mis g k base xs ys = sum mixs - mi_all where
  mixs = [micd gcol k base col ys | (col, gcol) <- cols]
  mi_all = micd g2 k base xs ys
  (g1, g2) = split g
  cols = (map colVector (columns xs)) `zip` (splitN (ncols xs) g1)

corexdc_tcs :: (RandomGen g, Eq a) => g -> Int -> Int -> Matrix a ->
               Matrix Double -> Double
corexdc_tcs g k base xs ys = DTC.tc base xs - ctcdc g k base xs ys

corexdc_mis :: (RandomGen g, Eq a) => g -> Int -> Int -> Matrix a ->
               Matrix Double -> Double
corexdc_mis g k base xs ys = sum mixs - mi_all where
  mixs = [midc gcol k base col ys | (col, gcol) <- cols]
  mi_all = midc g2 k base xs ys
  (g1, g2) = split g
  cols = (map colVector (columns xs)) `zip` (splitN (ncols xs) g1)


{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE StandaloneDeriving #-}
{-# LANGUAGE OverlappingInstances #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE ConstraintKinds #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE TypeOperators #-}

module Numeric.LinearAlgebra.Dimensional.DK.Internal
{-
(
   -- * Data.Packed.Vector
   (@>),
   -- * Data.Packed.Matrix
   -- ** dimension
   cols, rows,
   colsNT, rowsNT,
   hasRows, hasCols,
   -- (><),
   trans,
   -- reshape, flatten, fromLists, toLists, buildMatrix,
   (@@>),

   -- asRow, asColumn, fromRows, toRows, fromColumns, toColumns
   -- fromBlocks
   -- diagBlock,
   -- toBlocks, toBlocksEvery, repmat, flipud, fliprl
   -- subMatrix, takeRows, dropRows, takeColumns, dropColumns,
   -- extractRows, diagRect, takeDiag, mapMatrix,
   -- mapMatrixWithIndexM, mapMatrixWithIndexM_, liftMatrix,
   -- liftMatrix2, liftMatrix2Auto, fromArray2D,

   ident, -- where to put this?
   -- * Numeric.Container
   -- constant, linspace,
   diag,
   ctrans,
   -- ** Container class
   scalar,
   conj,
   scale, scaleRecip,
   recipMat,
   add,
   sub,
   mul,
   divide,
   equal,
   arctan2,
   hconcat,
   vconcat,
   concat,
   konst,
   zeroes,
   -- build, atIndex, minIndex, maxIndex, minElement, maxElement,
   -- sumElements, prodElements, step, cond, find, assoc, accum,
   -- Convert
   -- ** Product class
   multiply,
   -- dot, absSum, norm1, norm2, normInf,
   -- norm1, normInf,
   -- optimiseMult, mXm, mXv, vXm, (<.>),
   -- (<>), (<\>), outer, kronecker,
   (<>),
   -- ** Random numbers
   -- ** Element conversion
   -- ** Input/Output
   -- ** Experimental

   -- * Numeric.LinearAlgebra.Algorithms
   -- | incomplete wrapper for "Numeric.LinearAlgebra.Algorithms"

   -- ** Linear Systems
   -- linearSolve, luSolve, cholSolve, linearSolveLS, linearSolveSVD,
   inv,
   PInv(pinv),
   pinvTol,
   det,
   -- invlndet,
   rank,
   -- rcond,
   -- ** Matrix factorizations

   -- *** Singular value decomposition
   -- *** Eigensystems
   -- $eigs
   --wrapEig, wrapEigOnly,
   --EigCxt,
   -- **** eigenvalues and eigenvectors
   --eig,
   --eigC,
   --eigH,
   --eigH',
   --eigR,
   --eigS,
   --eigS',
   --eigSH,
   --eigSH',

   -- **** eigenvalues
   --eigOnlyC,
   --eigOnlyH,
   --eigOnlyR,
   --eigOnlyS,
   --eigenvalues,
   --eigenvaluesSH,
   --eigenvaluesSH',

   -- *** QR
   -- *** Cholesky
   -- *** Hessenberg
   -- *** Schur
   -- *** LU

   -- ** Matrix functions
   -- sqrtm, matFunc
   expm,

   -- ** Nullspace
   -- ** Norms
   -- ** Misc
   -- ** Util

   -- * actually internal
   toDM,
   DimMat(..),
  )
  -}
   where
import Foreign.Storable (Storable)
import GHC.Exts (Constraint)
import Numeric.Units.Dimensional.DK.Prelude hiding (concat)
import qualified Prelude as P
import qualified Numeric.NumType.DK.Integers as N

import Numeric.LinearAlgebra.Dimensional.DK.Shapes as S
import Data.Proxy

import qualified Numeric.Matrix as M
import Numeric.Matrix ((<|>), (<->))

data DimMat (shape :: Shape) a where
  DimMat :: M.Matrix a -> DimMat ('MatrixShape g rs cs) a
  DimVec :: M.Matrix a -> DimMat ('VectorShape d ds) a -- as a column vector

type ValidElement = M.MatrixElement

deriving instance (Show a, ValidElement a) => Show (DimMat s a)

(@>) :: (NN.KnownNat n, KnownDimension (VectorElement ('VectorShape d ds) n), Fractional a, ValidElement a)
     => DimMat ('VectorShape d ds) a
     -> Proxy n
     -> Quantity (VectorElement ('VectorShape d ds) n) a
(DimVec v) @> n = (v `M.at` (asIntFrom1 n,1)) *~ siUnit

(@@>) :: (NN.KnownNat nr, NN.KnownNat nc, KnownDimension (MatrixElement ('MatrixShape g rs cs) nr nc), Fractional a, ValidElement a)
    => DimMat ('MatrixShape g rs cs) a
    -> (Proxy nr, Proxy nc)
    -> Quantity (MatrixElement ('MatrixShape g rs cs) nr nc) a
DimMat m @@> (nr,nc) = (m `M.at` (asIntFrom1 nr, asIntFrom1 nc)) *~ siUnit

asInt, asIntFrom1 :: (NN.KnownNat n) => Proxy n -> Int
asInt = fromInteger . NN.natVal
asIntFrom1 = (P.+ 1) . asInt

{-
norm1 :: (sh ~ [r11 ': rs,ci], rs ~ MapConst r11 rs, ci ~ MapConst DOne ci, a ~ H.RealOf a)
         => DimMat sh a
         -> Quantity r11 a
norm1 (DimMat a) = Dimensional (H.pnorm H.PNorm1 a)

normInf :: (sh ~ [r11 ': rs,ci], rs ~ MapConst r11 rs, ci ~ MapConst DOne ci, a ~ H.RealOf a)
           => DimMat sh a
           -> Quantity r11 a
normInf (DimMat a) = Dimensional (H.pnorm H.Infinity a)
-}

{- | does H.'H.mXm' and H.'H.mXv'.

vXm and vXv (called dot) might be supported in the future too
-}
multiply :: (HasProduct s1 s2, ValidElement a)
    => DimMat s1 a -> DimMat s2 a
    -> DimMat (ShapeProduct s1 s2) a
multiply (DimMat a) (DimMat b) = DimMat (M.times a b)
multiply (DimMat a) (DimVec b) = DimVec (M.times a b)

infixl 7 <>
(<>) :: (HasProduct s1 s2, ValidElement a)
    => DimMat s1 a -> DimMat s2 a
    -> DimMat (ShapeProduct s1 s2) a
(<>) = multiply

trans :: (ValidElement a) => DimMat s a -> DimMat (ShapeTranspose s) a
trans (DimMat m) = DimMat $ M.transpose m

{-
pinvTol :: (PInv sh sh',
            a ~ Double,
           sh' ~ [ri2 ': _1 , DOne ': ci2]) => Double -> DimMat sh a -> DimMat sh' a
pinvTol tol (DimMat a) = DimMat (H.pinvTol tol a)

-}

det :: (Square s, KnownDimension (ShapeDeterminant s), Fractional a, ValidElement a) => DimMat s a -> Quantity (ShapeDeterminant s) a
det (DimMat m) = M.det m *~ siUnit

expm :: (s ~ ShapeProduct s s, HasProduct s s) => DimMat s a -> DimMat s a
expm = undefined

scale :: (Fractional a, ValidElement a, KnownDimension d)
         => Quantity d a -> DimMat s a -> DimMat (ShapeScale d s) a
scale x (DimMat m) = DimMat $ M.scale m (x /~ siUnit)
scale x (DimVec v) = DimVec $ M.scale v (x /~ siUnit)

add :: (ValidElement a) => DimMat s a -> DimMat s a -> DimMat s a
add (DimMat x) (DimMat y) = DimMat (M.plus x y)

sub :: (ValidElement a) => DimMat s a -> DimMat s a -> DimMat s a
sub (DimMat x) (DimMat y) = DimMat (M.minus x y)

equal :: (Eq a, ValidElement a) => DimMat s a -> DimMat s a -> Bool
equal (DimMat m1) (DimMat m2) = m1 == m2
equal (DimVec v1) (DimVec v2) = v1 == v2

hconcat :: (HorizontallyConcatenable s1 s2, ValidElement a) => DimMat s1 a -> DimMat s2 a -> DimMat (HorizontalConcatenation s1 s2) a
hconcat (DimMat m1) (DimMat m2) = DimMat (m1 <|> m2)
hconcat (DimMat m1) (DimVec v2) = DimMat (m1 <|> v2)
hconcat (DimVec v1) (DimMat m2) = DimMat (v1 <|> m2)
hconcat (DimVec v1) (DimVec v2) = DimMat (v1 <|> v2)

vconcat :: (VerticallyConcatenable s1 s2, ValidElement a) => DimMat s1 a -> DimMat s2 a -> DimMat (VerticalConcatenation s1 s2) a
vconcat (DimMat m1) (DimMat m2) = DimMat (m1 <-> m2)
vconcat (DimMat m1) (DimVec v2) = undefined --DimMat (m1 <-> M.transpose v2)
vconcat (DimVec v1) (DimMat m2) = undefined --DimMat (M.transpose v1 <-> m2)
vconcat (DimVec v1) (DimVec v2) = DimMat (M.transpose v1 <-> M.transpose v2)

fromRowVector :: (ValidElement a) => DimMat (VectorShape d ds) a -> DimMat (MatrixShape d '[] (MapDiv d ds)) a
fromRowVector (DimVec v) = DimMat (M.transpose v)

vconcat'  :: (ValidElement a,
              MapMulEq ds d cs
              )
             => DimMat (VectorShape d ds) a
             -> DimMat (MatrixShape g rs cs) a
             -> DimMat (MatrixShape d ((g/d) ': (MapMul (g/d) rs)) cs) a
vconcat' (DimVec v1) (DimMat m2) = DimMat (M.transpose v1 <-> m2)

-- I can't figure out why this is needed. But I also can't make a version of vconcat' that seems to work when used with fromRowVector. I don't get it.
vconcat'' :: (ValidElement a,
              MapMulEq ds2 (d2/d1) ds1
             )
             => DimMat (VectorShape d1 ds1) a
             -> DimMat (VectorShape d2 ds2) a
             -> DimMat (MatrixShape d1 '[d2/d1] (MapDiv d1 ds1)) a
vconcat'' (DimVec v1) (DimVec v2) = DimMat (M.transpose v1 <-> M.transpose v2)

concat :: (ValidElement a) => DimMat s1 a -> DimMat s2 a -> DimMat (VectorConcatenation s1 s2) a
concat (DimVec v1) (DimVec v2) = DimVec (v1 <-> v2)

--vecSingleton :: (KnownDimension d, Fractional a, ValidElement a) => Quantity d a -> DimMat (VectorShape d '[]) a
--vecSingleton x = DimVec . M.fromList $ [[x /~ siUnit]]

--vecCons :: (KnownDimension d', Fractional a, ValidElement a) => Quantity d' a -> DimMat (VectorShape d ds) a -> DimMat (VectorShape d' (d ': ds)) a
--vecCons x = concat (vecSingleton x)

vecSingleton :: (Fractional a, ValidElement a) => Quantity d a -> DimMat (VectorShape d '[]) a
vecSingleton x = DimVec . M.fromList $ [[12345]]

vecCons :: (Fractional a, ValidElement a) => Quantity d' a -> DimMat (VectorShape d ds) a -> DimMat (VectorShape d' (d ': ds)) a
vecCons x = concat (vecSingleton x)

rank :: DimMat s a -> Integer
rank = undefined

rows :: forall s a.(NN.KnownNat (ShapeRows s)) => DimMat s a -> Integer
rows _ = NN.natVal (Proxy :: Proxy (ShapeRows s))

cols :: forall s a.(NN.KnownNat (ShapeCols s)) => DimMat s a -> Integer
cols _ = NN.natVal (Proxy :: Proxy (ShapeCols s))

-- TODO: add a constraint that the row exists for better error message?
row :: forall n s a d ds.(MatrixRow s n ~ 'VectorShape d ds) => Proxy n -> DimMat s a -> DimMat ('VectorShape d ds) a
row _ _ = DimVec undefined

-- TODO: add a constraint that the column exists for better error message?
col :: forall n s a d ds.(MatrixColumn s n ~ 'VectorShape d ds) => Proxy n -> DimMat s a -> DimMat ('VectorShape d ds) a
col _ _ = DimVec undefined

{-
scalar :: (H.Field a,
          sh ~ ['[u], '[DOne]]) => Quantity u a -> DimMat sh a
scalar (Dimensional a) = DimMat (H.scalar a)
-}

{- | Numeric.Container.'H.konst', but the size is determined by the type.

>>> let n = hSucc (hSucc hZero) -- 2
>>> konst ((1::Double) *~ second) `hasRows` n `hasCols` n
2><2 1   1
s    1.0 1.0
s    1.0 1.0

-}
{-
konst :: forall u us ones a _1.
    (H.Field a,
     HNat2Integral (HLength ones),
     HNat2Integral (HLength us),
     ones ~ (DOne ': _1),
     AllEq DOne _1,
     AllEq u us)
    => Quantity u a -> DimMat [us, ones] a
konst (Dimensional a) = DimMat (H.konst a
    (hNat2Integral (proxy :: Proxy (HLength us)),
     hNat2Integral (proxy :: Proxy (HLength ones))))

-}

-- | identity matrix. The size is determined by the type.
ident :: forall g rs cs s a.(s ~ 'MatrixShape g rs cs, HasIdentity s, NN.KnownNat (ShapeRows s), ValidElement a) => DimMat ('MatrixShape g rs cs) a
ident = DimMat $ M.unit $ asInt (Proxy :: Proxy (ShapeRows s))

-- | zero matrix. The size and dimension is determined by the type.
zeroes :: forall g rs cs s a.(s ~ 'MatrixShape g rs cs, NN.KnownNat (ShapeRows s), NN.KnownNat (ShapeCols s), ValidElement a) => DimMat ('MatrixShape g rs cs) a
zeroes = DimMat $ M.matrix (r,c) (\(i,j) -> 0)
           where
             r = asInt (Proxy :: Proxy (ShapeRows s))
             c = asInt (Proxy :: Proxy (ShapeCols s))

trace :: (HasTrace s, Fractional a, ValidElement a, KnownDimension (ShapeTrace s)) => DimMat s a -> Quantity (ShapeTrace s) a
trace (DimMat m) = (P.sum $ M.trace m) *~ siUnit

conj :: DimMat s a -> DimMat s a
conj = undefined

-- | conjugate transpose
ctrans :: DimMat s a -> DimMat (ShapeTranspose s) a
ctrans = undefined

{-

diag :: (MapConst DOne v ~ c,
        c ~ (DOne ': _1)
        ) => DimMat '[v] t -> DimMat '[v,c] t
diag (DimVec a) = DimMat (H.diag a)
-}

{- $eigs

The Hmatrix eig factors A into P and D where A = P D inv(P) and D is diagonal.

The units for eigenvalues can be figured out:

>               _____
>      -1       |  c
> P D P  = A =  |r
>               |

>       _______
>       |   d
> P   = |c
>       |

>       _______
>       |   -1
>       |  c
>  -1   |
> P   = | -1
>       |d

So we can see that the dimension labeled `d-1` in P inverse is actually the
same `c` in `A`. The actual units of `d` don't seem to matter because the
`inv(d)` un-does any units that the `d` adds. So `d` can be all DOne. But
another choice, such as 1/c would be more appropriate, since then you can
expm your eigenvectors (not that that seems to be something people do)?

To get the row-units of A to match up, sometimes `D` will have units.
The equation ends up as D/c = r

Please ignore the type signatures on 'eig' 'eigC' etc. instead look at the type of
'wrapEig' 'wrapEigOnly' together with the hmatrix documentation (linked).

Perhaps the convenience definitions `eig m = wrapEig H.eig m` should be in
another module.
-}

{-

-- | 'wrapEig' H.'H.eig'
eig m = wrapEig H.eig m
-- | 'wrapEig' H.'H.eigC'
eigC m = wrapEig H.eigC m
-- | 'wrapEig' H.'H.eigH'
eigH m = wrapEig H.eigH m
-- | 'wrapEig' H.'H.eigH''
eigH' m = wrapEig H.eigH' m
-- | 'wrapEig' H.'H.eigR'
eigR m = wrapEig H.eigR m
-- | 'wrapEig' H.'H.eigS'
eigS m = wrapEig H.eigS m
-- | 'wrapEig' H.'H.eigS''
eigS' m = wrapEig H.eigS' m
-- | 'wrapEig' H.'H.eigSH'
eigSH m = wrapEig H.eigSH m
-- | 'wrapEig' H.'H.eigSH''
eigSH' m = wrapEig H.eigSH' m

-- | 'wrapEigOnly' H.'H.eigOnlyC'
eigOnlyC m = wrapEigOnly H.eigOnlyC m
-- | 'wrapEigOnly' H.'H.eigOnlyH'
eigOnlyH m = wrapEigOnly H.eigOnlyH m
-- | 'wrapEigOnly' H.'H.eigOnlyR'
eigOnlyR m = wrapEigOnly H.eigOnlyR m
-- | 'wrapEigOnly' H.'H.eigOnlyS'
eigOnlyS m = wrapEigOnly H.eigOnlyS m
-- | 'wrapEigOnly' H.'H.eigenvalues'
eigenvalues m = wrapEigOnly H.eigenvalues m
-- | 'wrapEigOnly' H.'H.eigenvaluesSH'
eigenvaluesSH m = wrapEigOnly H.eigenvaluesSH m
-- | 'wrapEigOnly' H.'H.eigenvaluesSH''
eigenvaluesSH' m = wrapEigOnly H.eigenvaluesSH' m

wrapEig :: (ones ~ (DOne ': _1), EigCxt [r,c] eigVal [cinv,ones],
    H.Field y, H.Field z)
    => (H.Matrix x -> (H.Vector y, H.Matrix z)) ->
    DimMat [r,c] x ->
    (DimMat '[eigVal] y, DimMat [cinv,ones] z)
wrapEig hmatrixFun (DimMat a) = case hmatrixFun a of
    (e,v) -> (DimVec e, DimMat v)

wrapEigOnly :: (EigCxt [r,c] eigVal '(), H.Field y)
    => (H.Matrix x -> H.Vector y) ->
    DimMat [r,c] x -> DimMat '[eigVal] y
wrapEigOnly hmatrixFun (DimMat a) = case hmatrixFun a of
    (e) -> DimVec e

-}


{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE FlexibleContexts #-}

module FEECa.Internal.Form (

  -- * Generic form types
    Dim, Form (Form, arity, dimVec, terms), split

  -- * Predefined primitive constructors
  , zeroForm, nullForm, oneForm

  -- * Form operations
  , apply, refine, refineBasis, inner, contract, trace

  ) where

import            Data.List ( intersect, elemIndex )
import            Data.Maybe ( fromJust )
import qualified  Numeric.LinearAlgebra.HMatrix as M
-- import qualified Numeric.LinearAlgebra.Data as M

import            FEECa.Utility.Combinatorics
import            FEECa.Utility.Discrete
import qualified  FEECa.Utility.Print           as P  ( Pretty(..), printForm,
                                                        text, (<>), (<+>), int )
import            FEECa.Utility.Utility               ( pairM, sumV, expSign, sign )

import            FEECa.Internal.Spaces     hiding    ( inner )
import qualified  FEECa.Internal.Spaces         as S  ( inner )

-- * General form: does not depend on the underlying vector space it works on
--   in any way.

type Dim = Int
type Idx = Int
type Prd = [Idx] -- a product of projections / indexed 1-forms
-- type LinComb f = [(f, Prd)]

-- | Bilinear, alternating forms over vectorspaces
data Form f =  -- we lose dependency on the type of vector!
    Form  { arity   :: Dim            -- ^ For complete evaluation
          , dimVec  :: Dim            -- ^ Of the underlying vector space
          , terms   :: [(f, Prd)] }   -- ^ List of terms of (coeff,wedge)'s
  deriving Eq

split :: (Module w, Scalar w ~ v) => Form w -> ([w], [Form v])
split (Form k n cs) = unzip $ map split' cs
  where split' (a,b) = (a, Form k n [(mulId, b)])

-- terms [(17, [1,2]), (38, [1,3])] = 17*dx1/\dx2 + 38*dx1/\dx3

-- | Invariant for terms of a defined form: all terms have the same arity
--   ie: each is the result of the exterior product of the same number of 1-forms
termsInv :: [(f, Prd)] -> Bool
termsInv []          = True
termsInv ((_,xs):ys) = all (\(_,xs') -> length xs == length xs') ys

-- NB: because of ((,) t) 's functorial nature, maybe it could make sense to
--   rearrange our terms so as to have them be (inds,coeff) like they used to be?
-- XXX: change so as to inspect result (in case f zeroes out a coeff?)
instance Functor Form where
  fmap f (Form k n cs) = Form k n (map (pairM f id) cs)

instance Show f => Show (Form f) where
  show (Form k n cs) = show k ++ "-form in " ++ show n ++ " dimensions: " ++
                        show cs

instance P.Pretty f => P.Pretty (Form f) where
  pPrint (Form k n cs) =
    P.int k P.<> P.text "-form in" P.<+> P.int n P.<+> P.text "dimensions:"
    P.<+> P.printForm "dx" "0" P.pPrint cs

-- NB: will be (i -> i -> Ordering) once we normalise all products to be in
--      their normal form - an increasing list
combineWithBy :: (a -> a -> a) -> (i -> i -> Bool)
              -> [(a, i)] -> [(a, i)] -> [(a,i)]
combineWithBy f p = foldl ins
  where ins [] x = [x]
        ins (y:ys) x | p (snd x) (snd y) = (fst x `f` fst y, snd y) : ys
                     | otherwise         = y : ins ys x

-- XXX: have combinator
--        aggregate f p cs1 cs2
--      where f :: a -> a -> a (or b -> c)
--            p :: MultiIndex -> MultiIndex -> Bool
--            cs1, cs2 :: [(a, MultiIndex)]
--      such that combines elements from cs2 with those in cs1 (sequentially,
--      no repeat, or filtering those appropriate) with p defining the "quotient"
--      class, f defining the transformation on coeff's
--      OPT: also a MultiIndex transformation? or a separte combinator for it?

-- | Sum of forms
-- Shall we do: permutation simplification/identification
(+++) :: Ring f => Form f -> Form f -> Form f
omega +++ eta
    | degNEq omega eta = errForm "(+++)" BiDegEq
    | spaNEq omega eta = errForm "(+++)" BiSpaEq
    | otherwise = Form (arity eta) (dimVec eta)
                       (step (terms omega) (terms eta))
  where step [] ys = ys
        step xs [] = xs
        step (x:xs) (y:ys)
          | snd x == snd y = let z = add (fst x) (fst y) in
              if fst x /= (addInv . fst) y then (z, snd x) : step xs ys
                            else step xs ys
          | snd x < snd y  = x : step xs (y:ys)
          | otherwise      = y : step (x:xs) ys

-- | Scaling of forms
(***) :: Ring f => f -> Form f -> Form f
(***) a | a == addId = \(Form k n _) -> zeroForm k n
        | otherwise  = fmap (mul a)

-- | (Exterior) Product of forms
(//\\) :: Ring f => Form f -> Form f -> Form f
omega //\\ eta
    | spaNEq omega eta = errForm "(//\\\\)" BiSpaEq
    | otherwise = Form (arity omega + arity eta) (dimVec eta)
                       (concatMap (\d -> map (`combine` d) (dxs d)) (terms eta))
  where dxs     (_,ys) = filter (null . intersect ys . snd) (terms omega)
        combine (a,xs) = pairM (mul a) (xs++)

-- | Forms over a 'Ring' form a 'Module'.
instance Ring f => Module (Form f) where
  type Scalar (Form f) = f
  addV = (+++)
  sclV = (***)

-- | Forms over a 'Field' form a 'VectorSpace'.
instance Field f => VectorSpace (Form f)

-- | For 'Form's defined over a 'Ring' we associate an 'Algebra': the exterior
-- algebra.
-- This instance is valid this way since we do not have a restrictive typing
-- for forms and hence addition is blind to arity *type-wise* - however,
-- runtime errors will take place if invalid addition is attempted.
instance Ring f => Algebra (Form f) where
  addA = addV
  (/\) = (//\\)
  sclA = sclV

-- | Basic abstract 1-form
oneForm :: Ring f => Dim -> Dim -> Form f
oneForm i n | i < 0 || i > n  = errForm "oneForm" MoProjBd
            | otherwise       = Form 1 n [ (mulId,[i]) ]

-- TODO: shall we have something special for these? no need to state dimension
-- n since they will be constantly zero anyway

-- | The (normalised) == 0 form
zeroForm :: Dim -> Dim -> Form f
zeroForm k n = Form k n []

-- | The k-arity == 0 form
nullForm :: Dim -> f -> Form f
nullForm n f = Form 0 n [(f, [])]

-- Necesitamos una función de pinchado
--  y así pinchar las consecutivas componentes
-- If the function was actually a field, this part would be simplified
contract :: (Ring f, Module v, Dimensioned v)
         => (Idx -> v -> f) -> Form f -> v -> Form f
contract proj omega v
    | vecNEq omega v = errForm "contract" MoVecEq
    | otherwise      = foldl (+++) (zeroForm k n) $
        map (\c -> Form k n $ map (pinchado c) [1..arity omega])
            (terms omega)
        {- concatMap (\c -> map (pinchado c) [1..arity omega]) -}
  where k = max 0 (arity omega - 1)
        n = dimVec omega
        pinchado (f,[]) _ = (f, []) -- error ??
        pinchado (f,ds) i = let (ds1,j:ds2) = splitAt (i-1) ds in
                              (expSign i $ mul f (proj j v), ds1 ++ ds2)
  {- TODO:  optimise
            error handling: proj indexing beyond dimension of v -}

-- list_to_index :: Integral a => a -> [a] -> a
-- list_to_index n (l:ls)

-- list_to_index' :: Integral a => a -> a -> [a] -> a
-- list_to_index' acc n (l:ls) = list_to_index' ((acc * n) + l) n ls
-- list_to_index' acc   _ []     = acc

-- make_lookup  :: (Ring r, EuclideanSpace v, Scalar v ~ r)
--              => Int -> Int -> [v] -> [[v]] -> Array [r]
-- make_lookup n k ds vvs
apply :: (EuclideanSpace v, Ring w, Module w, Scalar v ~ Scalar w)
      => [v] -> [v] -> Form w -> w
apply ds vs (Form k _ cs) = foldl addV addId (map (apply' k ds vs) cs)

apply' :: (EuclideanSpace v, Module w, Scalar v ~ Scalar w)
       => Int -> [v] -> [v] -> (w,[Int]) -> w
apply' _ _  _  (p, []) = p
apply' k ds vs (p, cs) = sclV c p
  where projections    = [toDouble $ dot (ds !! i) v | v <- vs, i <- cs]
        c              = fromDouble $ M.det $ M.matrix k projections

refineBasis :: (EuclideanSpace v, Scalar v ~ r)
            => [v] -> [[v]] -> [[r]]
refineBasis ds vvs = map (map (fromDouble . M.det)) submatrices
  where projections = [[[toDouble $ dot d v | v <- vs] | d <- ds] | vs <- vvs]
        matrices    = map (kSublists k) projections
        submatrices = map (map (M.matrix k . concat)) matrices
        k           = length (head vvs)

-- | Run function for 'Form's: given (an appropriate number of) vector arguments
--   and a 1-form basis (given as a basis-element indexing function 'proj'), it
--   evaluates the form on those arguments
refine :: (Ring w, Module w, Module v, Scalar v ~ Scalar w)
       => (Idx -> v -> Scalar w)  -- ^ The definition for the projection function
                                  --   for the specific vector space
       -> Form w
       -> [v] -> w
refine proj (Form _ _ cs) vs = {-#SCC "Form.refine" #-} sumV (map (($ vs) . formify proj) cs')
  where cs' | null cs   = [(addId,[])]
            | otherwise = cs
-- TODO: capture inconsistency between k and length vs here??
-- ALSO: 0-forms... not evaluating correctly now! Cfr: formify does not accept
--    empty cs
-- XXX: for now proj should take care of the error... change later when settled

-- | Helper function in evaluation: given a 1-form basis, converts a single
--   'Form' term into an actual function on vectors
formify :: (Ring w, Module w, Module v, Scalar w ~ Scalar v)
        => (i -> v -> Scalar v) -> (w,[i]) -> [v] -> w
formify _    (s, [])   _  = s
formify proj (s, i:is) vs
    | null is   = {-#SCC "Form.formify"  #-} sclV (proj i (head vs)) s
    | otherwise = {-#SCC "Form.formifyR" #-}
        foldl addV addId
              (map (\(w,e) -> sclV
                                (mul (sign (w,e)) ((proj i . head) (pick' w vs)))
                                (formify proj (s,is) (pick' e vs)))
                   (permutationPairs (length is + 1) 1 (length is)))
  where pick' ns = pick (differences ns)

-- We need a basis here
inner :: (InnerProductSpace w, EuclideanSpace v, Scalar w ~ Scalar v)
      => (Idx -> v -> Scalar w)  -- ^ Projection function in the specific vector space
      -> Form w -> Form w -> Scalar w
inner proj omega eta
    | degNEq omega eta = errForm "inner" BiDegEq -- TODO (??)
    | otherwise = foldl
          (flip $ \vs -> add (S.inner (app omega vs) (app eta vs)))
          addId
          (map pick' (permutations n (arity omega)))
  where pick' is  = pick (differences is) (map (unitVector n) [0..n-1])
        app       = refine proj
        n         = dimVec omega

trace :: [Int] -> Form w -> Form w
trace sigma (Form k n ts)
    | k' < k    = zeroForm k' n
    | otherwise = Form k n (map (pairM id restrict') ts')
  where ts'       = filter (is_in_range' sigma . snd) ts
        restrict' = map (fromJust . (`elemIndex` sigma))
        k'        = length sigma
        -- is_in_range' :: [Int] -> [Int] -> Bool
        is_in_range' sigma = all (`elem` sigma)

-- | Checks arity equality
degNEq :: Form f -> Form f -> Bool
degNEq omega eta = arity omega /= arity eta

-- | Checks combined arity bound
degNBd :: Form f -> Form f -> Bool
degNBd  omega eta = (arity omega + arity eta) <= dimVec omega

-- | Checks compatible underlying vector space dimensions between forms
spaNEq :: Form f -> Form f -> Bool
spaNEq omega eta = dimVec omega /= dimVec eta

-- | Checks compatible underlying vector space dimensions between a form and a
-- 'Dimensioned' type value
vecNEq :: Dimensioned v => Form f -> v -> Bool
vecNEq omega v = dimVec omega /= dim v

errForm :: String -> FormMust -> t
errForm callee obligation = error $ "Form." ++ callee ++
                                    ": forms must " ++ show obligation

-- | Kinds of enforcements to the definitions and operations between/for 'Form'
data FormMust = BiDegEq | BiDegBd | BiSpaEq | MoProjBd | MoVecEq

instance Show FormMust where
  show BiDegEq  = "be of the same degree"
  show BiDegBd  = "have joint degree bounded by the working vector space dimension"
  show BiSpaEq  = "act on the same vector space"
  show MoProjBd = "project components of the underlying vector space"
  show MoVecEq  = "act on vectors of the working vectors space"


{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE ViewPatterns #-}
{-# LANGUAGE MultiWayIf #-}

module Data.AdaptiveCoordinateDescent
  ( adaptiveCoordinateDescent )
  where

import Control.Monad.State.Strict
import Data.AdaptiveCoordinateDescent.Internal
import Data.Foldable
import Data.List ( sortBy )
import Data.Ord
import Data.Traversable
import Numeric.LinearAlgebra hiding ( toList )
import Pipes hiding ( for )
import qualified Pipes.Prelude as P
import System.Random.MWC

-- | Performs adaptive coordinate descent.
adaptiveCoordinateDescent :: (MonadIO m, Traversable f, Eq (f Double))
                          => (f Double -> m Double)  -- ^ Cost function. Smaller cost = better model.
                          -> f Double                -- ^ Initial parameters from which to start optimization.
                          -> Double                  -- ^ Tolerance on how small changes are allowed.
                                                     --   If no changes larger
                                                     --   than this value are
                                                     --   found, then this
                                                     --   descent function
                                                     --   returns. You probably
                                                     --   want something small
                                                     --   like 0.0000001
                          -> Int                     -- ^ How many best N candidates to keep for PCA phase.
                                                     --   Good values are
                                                     --   probably around the
                                                     --   size of number of
                                                     --   your parameters.
                                                     --   Maybe more if you
                                                     --   only have like 2
                                                     --   parameters. Probably.
                                                     --   We don't know.
                          -> Producer (f Double, Double) m (f Double, Double)
adaptiveCoordinateDescent evaluate initial_params tolerance how_many_to_keep = do
  score <- lift $ evaluate initial_params
  yield (initial_params, score)

  loop_it initial_params initial_params score (toColumns $ ident $ length $ toList initial_params) [] (replicate (length $ toList initial_params) 1.0) []
 where
  k_succ = 2.0 :: Double
  k_unsucc = 0.5 :: Double

  loop_it original_params params score principal_components last_n_params step_sizes new_step_sizes | length last_n_params < 5 = do
    rng <- liftIO createSystemRandom
    items <- liftIO $ P.toListM (makePerturbed params rng 5)
    scored_items <- lift $ for (initial_params:items) $ \item -> (,) item <$> evaluate item
    loop_it original_params params score principal_components scored_items step_sizes new_step_sizes

  loop_it original_params params score (principal_component:principal_components) last_n_params (step_size':step_sizes) new_step_sizes = do
    let step_size = max tolerance step_size'
        scaled_pcomponent = cmap (*step_size) principal_component
        top_candidate    = params `plus` fromVector scaled_pcomponent initial_params
        bottom_candidate = params `minus` fromVector scaled_pcomponent initial_params

    top_score    <- lift $ evaluate top_candidate
    bottom_score <- lift $ evaluate bottom_candidate

    yield (top_candidate, top_score)
    yield (bottom_candidate, bottom_score)

    if | top_score < score
         -> loop_it original_params top_candidate top_score principal_components ((top_candidate, top_score):last_n_params) step_sizes (step_size*k_succ:new_step_sizes)
       | bottom_score < score
         -> loop_it original_params bottom_candidate bottom_score principal_components ((bottom_candidate, bottom_score):last_n_params) step_sizes (step_size*k_succ:new_step_sizes)
       | otherwise
         -> loop_it original_params params score principal_components last_n_params step_sizes (step_size*k_unsucc:new_step_sizes)

  loop_it original_params params score [] _last_n_params _step_sizes new_step_sizes
    | original_params == params &&
      all (\x -> x <= tolerance) new_step_sizes = yield (params, score) >> return (params, score)
  loop_it _original_params params score [] last_n_params _step_sizes new_step_sizes = do
    let num_coordinates = length (toList params)
        sorted_items = sortBy (comparing snd) last_n_params
        picked_items = take how_many_to_keep sorted_items

        items_mat = normalize $
                    ( (length picked_items><num_coordinates)
                      (mconcat $ fmap (toList . fst) picked_items) )

    let (_evec, mat2) = rightSV items_mat
    loop_it params params score (toColumns mat2) picked_items (reverse new_step_sizes) []

  loop_it _ _ _ _ _ _ _ = error "adaptiveCoordinateDescent: impossible case."


{-# LANGUAGE Rank2Types #-}
-----------------------------------------------------------------------------
-- |
-- Module     : Numeric.LinearAlgebra.Matrix.Tri
-- Copyright  : Copyright (c) 2010, Patrick Perry <patperry@gmail.com>
-- License    : BSD3
-- Maintainer : Patrick Perry <patperry@gmail.com>
-- Stability  : experimental
--
-- Triangular views of matrices.
--

module Numeric.LinearAlgebra.Matrix.Tri (
    -- * Immutable interface

    -- ** Vector multiplication
    triMulVector,

    -- ** Matrix multiplication
    triMulMatrix,
    triMulMatrixWithScale,

    -- ** Vector solving
    triSolvVector,

    -- ** Matrix solving
    triSolvMatrix,
    triSolvMatrixWithScale,

    -- * Mutable interface
    triCreate,

    -- ** Vector multiplication
    triMulVectorM_,

    -- ** Matrix multiplication
    triMulMatrixM_,
    triMulMatrixWithScaleM_,

    -- ** Vector solving
    triSolvVectorM_,

    -- ** Matrix solving
    triSolvMatrixM_,
    triSolvMatrixWithScaleM_,

    ) where

import Control.Monad( when )
import Control.Monad.ST( ST, runST, unsafeIOToST )
import Text.Printf( printf )

import Numeric.LinearAlgebra.Vector( Vector, STVector )
import qualified Numeric.LinearAlgebra.Vector as V
import Numeric.LinearAlgebra.Matrix.Base( Matrix )
import Numeric.LinearAlgebra.Matrix.STBase( STMatrix, RMatrix )
import qualified Numeric.LinearAlgebra.Matrix.STBase as M
import Numeric.LinearAlgebra.Types
import qualified Foreign.BLAS as BLAS

-- | A safe way to create and work with a mutable Tri Matrix before returning
-- an immutable one for later perusal.
triCreate :: (Storable e)
           => (forall s. ST s (Tri (STMatrix s) e))
           -> Tri Matrix e
triCreate mt = runST $ do
    (Tri u d ma) <- mt
    a <- M.unsafeFreeze ma
    return $ Tri u d a

-- | @triMulVector trans a x@ returns @op(a) * x@, where @op(a)@ is
-- determined by @trans@.
triMulVector :: (BLAS2 e)
             => Trans
             -> Tri Matrix e
             -> Vector e
             -> Vector e
triMulVector trans a x =
    V.create $ do
        x' <- V.newCopy x
        triMulVectorM_ trans a x'
        return x'

-- | @triMulMatrix side a b@
-- returns @alpha * op(a) * b@ when @side@ is @LeftSide@ and
-- @alpha * b * op(a)@ when @side@ is @RightSide@.  Operation
-- @op(a)@ is determined by @trans@.
triMulMatrix :: (BLAS3 e)
              => Side
              -> Trans -> Tri Matrix e
              -> Matrix e
              -> Matrix e
triMulMatrix side trans a b =
    M.create $ do
        b' <- M.newCopy b
        triMulMatrixM_ side trans a b'
        return b'

-- | @triMulMatrixWithScale alpha side trans a b@
-- returns @alpha * op(a) * b@ when @side@ is @LeftSide@ and
-- @alpha * b * op(a)@ when @side@ is @RightSide@.  Operation
-- @op(a)@ is determined by @trans@.
triMulMatrixWithScale :: (BLAS3 e)
                       => e
                       -> Side
                       -> Trans -> Tri Matrix e
                       -> Matrix e
                       -> Matrix e
triMulMatrixWithScale alpha side trans a b =
    M.create $ do
        b' <- M.newCopy b
        triMulMatrixWithScaleM_ alpha side trans a b'
        return b'

-- | @triMulVectorM_ a x@ sets @x := op(a) * x@, where @op(a)@ is determined
-- by @trans@.
triMulVectorM_ :: (RMatrix m, BLAS2 e)
               => Trans -> Tri m e
               -> STVector s e
               -> ST s ()
triMulVectorM_ trans (Tri uplo diag a) x = do
    (ma,na) <- M.getDim a
    nx <- V.getDim x
    let n = nx

    when (ma /= na) $ error $
        printf ("triMulVectorM_"
                ++ " _"
                ++ " (Tri _ _ <matrix with dim (%d,%d)>)"
                ++ " _"
                ++ ": matrix is not square")
               ma na

    when ((not . and) [ (ma,na) == (n,n)
                      , nx == n
                      ]) $ error $
        printf ("triMulVectorM_"
                ++ " _"
                ++ " (Tri _ _ <matrix with dim (%d,%d)>)"
                ++ " <vector with dim %d>"
                ++ ": dimension mismatch")
               ma na
               nx

    unsafeIOToST $
        M.unsafeWith a $ \pa lda ->
        V.unsafeWith x $ \px ->
            BLAS.trmv uplo trans diag n pa lda px 1

-- | @triMulMatrixM_ side trans a b@
-- sets @b := op(a) * b@ when @side@ is @LeftSide@ and
-- @b := b * op(a)@ when @side@ is @RightSide@.  Operation
-- @op(a)@ is determined by @trans@.
triMulMatrixM_ :: (RMatrix m, BLAS3 e)
               => Side
               -> Trans -> Tri m e
               -> STMatrix s e
               -> ST s ()
triMulMatrixM_ = triMulMatrixWithScaleM_ 1

-- | @triMulMatrixWithScaleM_ alpha side trans a b@
-- sets @b := alpha * op(a) * b@ when @side@ is @LeftSide@ and
-- @b := alpha * b * op(a)@ when @side@ is @RightSide@.  Operation
-- @op(a)@ is determined by @trans@.
triMulMatrixWithScaleM_ :: (RMatrix m, BLAS3 e)
                         => e
                         -> Side
                         -> Trans -> Tri m e
                         -> STMatrix s e
                         -> ST s ()
triMulMatrixWithScaleM_ alpha side trans (Tri uplo diag a) b = do
    (ma,na) <- M.getDim a
    (mb,nb) <- M.getDim b
    let (m,n) = (mb,nb)

    when (ma /= na) $ error $
        printf ("triMulMatrixWithScaleM_"
                ++ " _"
                ++ " _"
                ++ " _"
                ++ " (Tri _ _ <matrix with dim (%d,%d)>)"
                ++ " _"
                ++ ": matrix is not square")
               ma na

    when ((not . and) [ case side of LeftSide  -> (ma,na) == (m,m)
                                     RightSide -> (ma,na) == (n,n)
                      , (mb, nb ) == (m,n)
                      ]) $ error $
        printf ("triMulMatrixWithScaleM_"
                ++ " _"
                ++ " %s"
                ++ " _"
                ++ " (Tri _ _ <matrix with dim (%d,%d)>)"
                ++ " <matrix with dim (%d,%d)>"
                ++ ": dimension mismatch")
               (show side)
               ma na
               mb nb

    unsafeIOToST $
        M.unsafeWith a $ \pa lda ->
        M.unsafeWith b $ \pb ldb ->
            BLAS.trmm side uplo trans diag m n alpha pa lda pb ldb

-- | @triSolvVector trans a x@ returns @op(a) \\ x@, where @op(a)@ is
-- determined by @trans@.
triSolvVector :: (BLAS2 e)
             => Trans
             -> Tri Matrix e
             -> Vector e
             -> Vector e
triSolvVector trans a x =
    V.create $ do
        x' <- V.newCopy x
        triSolvVectorM_ trans a x'
        return x'

-- | @triSolvMatrix side a b@
-- returns @alpha * op(a) \\ b@ when @side@ is @LeftSide@ and
-- @alpha * b * op(a)@ when @side@ is @RightSide@.  Operation
-- @op(a)@ is determined by @trans@.
triSolvMatrix :: (BLAS3 e)
              => Side
              -> Trans -> Tri Matrix e
              -> Matrix e
              -> Matrix e
triSolvMatrix side trans a b =
    M.create $ do
        b' <- M.newCopy b
        triSolvMatrixM_ side trans a b'
        return b'

-- | @triSolvMatrixWithScale alpha side trans a b@
-- returns @alpha * op(a) \\ b@ when @side@ is @LeftSide@ and
-- @alpha * b * op(a)@ when @side@ is @RightSide@.  Operation
-- @op(a)@ is determined by @trans@.
triSolvMatrixWithScale :: (BLAS3 e)
                       => e
                       -> Side
                       -> Trans -> Tri Matrix e
                       -> Matrix e
                       -> Matrix e
triSolvMatrixWithScale alpha side trans a b =
    M.create $ do
        b' <- M.newCopy b
        triSolvMatrixWithScaleM_ alpha side trans a b'
        return b'

-- | @triSolvVectorM_ a x@ sets @x := op(a) \\ x@, where @op(a)@ is determined
-- by @trans@.
triSolvVectorM_ :: (RMatrix m, BLAS2 e)
               => Trans -> Tri m e
               -> STVector s e
               -> ST s ()
triSolvVectorM_ trans (Tri uplo diag a) x = do
    (ma,na) <- M.getDim a
    nx <- V.getDim x
    let n = nx

    when (ma /= na) $ error $
        printf ("triSolvVectorM_"
                ++ " _"
                ++ " (Tri _ _ <matrix with dim (%d,%d)>)"
                ++ " _"
                ++ ": matrix is not square")
               ma na

    when ((not . and) [ (ma,na) == (n,n)
                      , nx == n
                      ]) $ error $
        printf ("triSolvVectorM_"
                ++ " _"
                ++ " (Tri _ _ <matrix with dim (%d,%d)>)"
                ++ " <vector with dim %d>"
                ++ ": dimension mismatch")
               ma na
               nx

    unsafeIOToST $
        M.unsafeWith a $ \pa lda ->
        V.unsafeWith x $ \px ->
            BLAS.trsv uplo trans diag n pa lda px 1

-- | @triSolvMatrixM_ side trans a b@
-- sets @b := op(a) \\ b@ when @side@ is @LeftSide@ and
-- @b := b * op(a)@ when @side@ is @RightSide@.  Operation
-- @op(a)@ is determined by @trans@.
triSolvMatrixM_ :: (RMatrix m, BLAS3 e)
               => Side
               -> Trans -> Tri m e
               -> STMatrix s e
               -> ST s ()
triSolvMatrixM_ = triSolvMatrixWithScaleM_ 1

-- | @triSolvMatrixWithScaleM_ alpha side trans a b@
-- sets @b := alpha * op(a) \\ b@ when @side@ is @LeftSide@ and
-- @b := alpha * b * op(a)@ when @side@ is @RightSide@.  Operation
-- @op(a)@ is determined by @trans@.
triSolvMatrixWithScaleM_ :: (RMatrix m, BLAS3 e)
                         => e
                         -> Side
                         -> Trans -> Tri m e
                         -> STMatrix s e
                         -> ST s ()
triSolvMatrixWithScaleM_ alpha side trans (Tri uplo diag a) b = do
    (ma,na) <- M.getDim a
    (mb,nb) <- M.getDim b
    let (m,n) = (mb,nb)

    when (ma /= na) $ error $
        printf ("triSolvMatrixWithScaleM_"
                ++ " _"
                ++ " _"
                ++ " _"
                ++ " (Tri _ _ <matrix with dim (%d,%d)>)"
                ++ " _"
                ++ ": matrix is not square")
               ma na

    when ((not . and) [ case side of LeftSide  -> (ma,na) == (m,m)
                                     RightSide -> (ma,na) == (n,n)
                      , (mb, nb ) == (m,n)
                      ]) $ error $
        printf ("triSolvMatrixWithScaleM_"
                ++ " _"
                ++ " %s"
                ++ " _"
                ++ " (Tri _ _ <matrix with dim (%d,%d)>)"
                ++ " <matrix with dim (%d,%d)>"
                ++ ": dimension mismatch")
               (show side)
               ma na
               mb nb

    unsafeIOToST $
        M.unsafeWith a $ \pa lda ->
        M.unsafeWith b $ \pb ldb ->
            BLAS.trsm side uplo trans diag m n alpha pa lda pb ldb


{-|
Module      : HABQTlib.MeasurementProcessing

Functions that deal with optimising and simulating measurements that take place
during tomography.
-}
module HABQTlib.MeasurementProcessing
  ( PurePOVM
  , SingleQbParam
  , measurementProbs
  , svToAngles
  , blochAnglesToSV
  , mkAntipodalPOVM
  , productPOVM
  , simulateMeasuremet
  , optimiseSingleQbPOVM
  ) where

import Control.Applicative (liftA2)
import Control.Newtype.Generics (unpack)
import Data.Complex (mkPolar, polar)
import Data.List (unfoldr)
import Data.Maybe (fromJust)
import qualified Data.Vector as V
import HABQTlib.Data
import HABQTlib.Data.Particle
import qualified Numeric.GSL as GSL
import Numeric.LinearAlgebra (Complex((:+)))
import qualified Numeric.LinearAlgebra as LA
import qualified System.Random.MWC as MWC
import System.Random.MWC.Distributions (categorical)

-- | A POVM consisting of projections onto pure states.
type PurePOVM = V.Vector WeighedPureStateVector

-- | Spherical coordinates of a pure single qubit state on Bloch sphere.
type SingleQbParam = (Double, Double)

-- | Probabilities to measure elements of POVM when performing the measurement
-- over a mixed state.
measurementProbs :: PurePOVM -> DensityMatrix -> V.Vector Double
measurementProbs povm dm = V.map (weighedProb dm) povm

weighedProb :: DensityMatrix -> WeighedPureStateVector -> Double
weighedProb dm (WeighedPureStateVector (w, sv)) = w * pureStateLikelihood sv dm

entropy :: V.Vector Double -> Double
entropy = negate . V.sum . V.map (\p -> p * log p)

povmPointEntropy :: PurePOVM -> DensityMatrix -> Double
povmPointEntropy povm dm = entropy (measurementProbs povm dm)

povmMeanEntropy :: PurePOVM -> ParticleHierarchy -> Double
povmMeanEntropy povm ph =
  let meanEnt = foldOverPts (povmPointEntropy povm) (*) (+) 0
      totalWeight = V.sum . V.map ptsWeight $ ph
      rawEntropy = V.sum . V.map (liftA2 (*) ptsWeight meanEnt) $ ph
   in rawEntropy / totalWeight

-- | Recover a state vector from spherical coordinates.
blochAnglesToSV :: SingleQbParam -> PureStateVector
blochAnglesToSV (th, phi) =
  let z = LA.fromList [1, 0]
      o = LA.fromList [0, 1]
   in PureStateVector . LA.fromColumns . pure $
      LA.scalar (cos (th / 2) :+ 0) * z +
      LA.scalar (mkPolar (sin (th / 2)) phi) * o

-- | Return spherical coordinates of a single qubit pure state (on Bloch
-- sphere).
svToAngles :: PureStateVector -> Maybe SingleQbParam
svToAngles (PureStateVector sv) =
  let s = LA.size sv
      (m0, ph0) = polar $ LA.atIndex sv (0, 0)
      (_, ph1) = polar $ LA.atIndex sv (1, 0)
      ph = ph1 - ph0
      th = 2 * acos m0
   in if s == (2, 1)
        then Just (th, ph)
        else Nothing

-- | Given spherical coordinates, construct a POVM from the given vector and
-- one orthogonal to it.
mkAntipodalPOVM :: SingleQbParam -> PurePOVM
mkAntipodalPOVM c@(th, phi) =
  let sv = blochAnglesToSV c
      sv' = blochAnglesToSV (pi - th, phi + pi)
   in V.fromList $ WeighedPureStateVector <$> [(1, sv), (1, sv')]

reshapeList :: Int -> [a] -> [[a]]
reshapeList n =
  unfoldr
    (\b ->
       if length b < n
         then Nothing
         else Just (splitAt n b))

listToPairs :: [a] -> [(a, a)]
listToPairs = fmap (\(a:[b]) -> (a, b)) . reshapeList 2

pairToList :: (a, a) -> [a]
pairToList (a, b) = [a, b]

-- | Given a list of POVM measurements on sub-systems, construct a POVM over
-- the composite system that includes all of them.
productPOVM :: [PurePOVM] -> PurePOVM
productPOVM sqbPovms =
  let pr ::
           WeighedPureStateVector
        -> WeighedPureStateVector
        -> WeighedPureStateVector
      pr wsv0 wsv1 =
        WeighedPureStateVector (w0 * w1, PureStateVector $ LA.kronecker sv0 sv1)
        where
          up = fmap unpack . unpack
          (w0, sv0) = up wsv0
          (w1, sv1) = up wsv1
   in V.foldl1' (liftA2 pr) $ V.fromList sqbPovms

-- | Approximate most informative separable POVM over a composite system of
-- quantum bits, given a list of single-qubit starting points and a particle
-- distribution.
optimiseSingleQbPOVM ::
     OptIter -- ^ Number of optimisation steps to perform
  -> [PureStateVector] -- ^ single qubit initial states
  -> ParticleHierarchy
  -> PurePOVM
optimiseSingleQbPOVM iter sv0s ph =
  let nq = length sv0s
      method = GSL.NMSimplex2
      precision = 1e-6
      iterations = iter
      initialBox = replicate (2 * nq) (pi / 2)
      estimateDM = getMixedEstimate ph
      obj params =
        negate $ povmPointEntropy povm estimateDM - povmMeanEntropy povm ph
        where
          povm = productPOVM . fmap mkAntipodalPOVM . listToPairs $ params
      start = concatMap (pairToList . fromJust . svToAngles) sv0s
      result = GSL.minimize method precision iterations initialBox obj start
   in productPOVM . fmap mkAntipodalPOVM . listToPairs . fst $ result

-- | Simulate a POVM over a mixed state and return the state vector on which
-- the projection was obtained.
simulateMeasuremet ::
     DensityMatrix -> PurePOVM -> MWC.GenIO -> IO PureStateVector
simulateMeasuremet dm povm gen = do
  let probs = measurementProbs povm dm
      svs = V.map (\(WeighedPureStateVector (_, sv)) -> sv) povm
  idx <- categorical probs gen
  return $ svs V.! idx


{-# LANGUAGE FlexibleInstances     #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE DataKinds #-}
{-# OPTIONS_GHC -Wno-simplifiable-class-constraints #-}
module Data.Matrix.Static.LinearAlgebra.Types
    ( Numeric(..)
    , Matrix
    , MMatrix
    , SparseMatrix
    ) where

import Data.Vector.Storable (Vector, Storable)
import Data.Vector.Storable.Mutable (MVector)
import Data.Complex (Complex)
import Foreign.C.Types

import qualified Data.Matrix.Static.Dense as D
import qualified Data.Matrix.Static.Dense.Mutable as DM
import qualified Data.Matrix.Static.Sparse as S

class (S.Zero a, Storable a, Num a) => Numeric a where
    foreignType :: a -> CInt

instance Numeric Float where foreignType _ = 0
instance Numeric CFloat where foreignType _ = 0
instance Numeric Double where foreignType _ =1
instance Numeric (Complex Float) where foreignType _ = 2
instance Numeric (Complex Double) where foreignType _ = 3

type Matrix r c a = D.Matrix r c Vector a
type MMatrix r c s a = DM.MMatrix r c MVector s a

type SparseMatrix r c a = S.SparseMatrix r c Vector a


#!/usr/bin/env stack
-- stack --resolver lts-15.04 runghc --package reanimate
module Main where

import           Codec.Picture.Types
import           Control.Lens
import           Control.Monad
import           Data.List
import           Data.Maybe
import qualified Data.Text                     as T
import           Data.Tuple
import qualified Data.Vector                   as V
import           Linear.V2
import           Linear.Vector
import qualified Numeric.LinearAlgebra         as Matrix
import           Numeric.LinearAlgebra.HMatrix (Matrix, linearSolve, toLists,
                                                (><))
import           Reanimate
import           Reanimate.Math.Common
import           Reanimate.Math.Polygon

bgColor :: PixelRGBA8
bgColor = PixelRGBA8 252 252 252 0xFF

type Points = V.Vector (V2 Double)
type Edges = [(Int, Int, Int)]
data Mesh = Mesh { meshPoints :: Points, meshEdges :: Edges }
data MeshPair = MeshPair Points Points Edges
-- The points in a RelMesh are:
--   relMeshStatic ++ x where Ax = B
data RelMesh = RelMesh
  { relMeshStatic :: Points
  , relMeshEdges  :: Edges
  , relMeshA      :: Matrix Double
  , relMeshB      :: Matrix Double
  }
data RelMeshPair = RelMeshPair Points Edges (Matrix Double) (Matrix Double) (Matrix Double) (Matrix Double)
-- Linear interpolation on RelMesh gives smooth morph.
-- solveMesh :: RelMesh -> Mesh
-- mkRelative :: Mesh -> RelMesh
-- mkRelativePair :: MeshPair -> RelMeshPair
-- triangulate :: Polygon -> Polygon -> MeshPair ?
-- embed :: MeshPair -> MeshPair
-- compatible :: Mesh -> Mesh -> Maybe MeshPair
-- linearInterpolate :: MeshPair -> Double -> Mesh
-- convexInterpolate :: RelMeshPair -> Double -> RelMesh

meshToPolygon :: Mesh -> Polygon
meshToPolygon mesh = pScale 2 $ mkPolygon $ V.fromList
    [ realToFrac <$> (meshPoints mesh V.! (n-1))
    | n <- polygonNodes ]
  where
    polygonNodes :: [Int]
    polygonNodes = [5,8,9,7,4,6]

main :: IO ()
main = reanimate morphAnimation

genTrails :: (Double -> Mesh) -> SVG
genTrails mkMesh =
    withFillOpacity 0 $
    withStrokeWidth (defaultStrokeWidth*0.5) $
    withStrokeColor "black" $
    withStrokeDashArray [0.1,0.1] $
    mkGroup $ map mkTrail $ transpose
      [ V.toList $ V.map (fmap realToFrac) $ polygonPoints $ meshToPolygon mesh
      | n <- [0..steps]
      , let mesh = mkMesh (fromIntegral n / fromIntegral steps)
      ]
  where
    steps = 100 :: Int
    mkTrail lst =mkLinePath [ (x,y) | V2 x y <- lst ]

morphAnimation :: Animation
morphAnimation = playThenReverseA $ pauseAround 1 1 $ addStatic (mkBackgroundPixel bgColor) $
  signalA (curveS 2) $ mkAnimation 5 $ \t -> lowerTransformations $ scale 3 $ pathify $
  mkGroup
  [ translate (-1.5) 0 $ withFillColor "lightgreen" $ mkGroup
    [ withStrokeWidth defaultStrokeWidth $
      withStrokeColor "black" $
      polygonShape $ meshToPolygon $ linearInterpolate meshPair t
    , linearTrails
    , polygonNumDots (meshToPolygon $ linearInterpolate meshPair t) 0
    ]
  , translate 1 0 $ withFillColor "cyan" $ mkGroup
    [ withStrokeWidth defaultStrokeWidth $
      withStrokeColor "black" $
      polygonShape $ meshToPolygon $ solveMesh $ convexInterpolate relPair t
    , curveTrails
    , polygonNumDots (meshToPolygon $ solveMesh $ convexInterpolate relPair t) 0
    ]
  ]
  where
    linearTrails = genTrails (linearInterpolate meshPair)
    curveTrails = genTrails (solveMesh . convexInterpolate relPair)
    relPair = mkRelativePair meshPair
    meshPair = fromJust $ compatible example1 example2

mkLineP :: P -> P -> SVG
mkLineP (V2 x1 y1) (V2 x2 y2) = mkLine (x1,y1) (x2,y2)

-- FIXME: Check that the triangles are all anticlockwise.
-- FIXME: Check that the edges connect all the points.
-- FIXME: Check that the edgse leave no gaps.
compatible :: Mesh -> Mesh -> Maybe MeshPair
compatible a b =
  if meshEdges a == meshEdges b && V.length (meshPoints a) == V.length (meshPoints b)
    then Just $ MeshPair (meshPoints a) (meshPoints b) (meshEdges a)
    else Nothing

linearInterpolate :: MeshPair -> Double -> Mesh
linearInterpolate (MeshPair aP bP edges) t = Mesh
    { meshPoints = V.zipWith (lerp (1-t)) aP bP
    , meshEdges = edges }

example1 :: Mesh
example1 = Mesh points edges
  where
    points = V.fromList
      [ V2 1 0
      , V2 (-1/2) (sqrt 3 / 2)
      , V2 (-1/2) (-sqrt 3 / 2)
      , 3 * points V.! 0 ^/ 4
      , 3 * points V.! 1 ^/ 4
      , 3 * points V.! 2 ^/ 4
      , points V.! 0 ^/ 2
      , points V.! 1 ^/ 2
      , points V.! 2 ^/ 2
      ]
    edges =
      [ (1,5,4), (1,2,5), (2,6,5), (2,3,6), (3,4,6), (3,1,4)
      , (4,8,7), (4,5,8), (5,9,8), (5,6,9), (6,7,9), (6,4,7), (7,8,9)]

example2 :: Mesh
example2 = Mesh points edges
  where
    points = V.fromList
      [ V2 1 0
      , V2 (-1/2) (sqrt 3 / 2)
      , V2 (-1/2) (-sqrt 3 / 2)
      , 3 * points V.! 2 ^/ 4
      , 3 * points V.! 0 ^/ 4
      , 3 * points V.! 1 ^/ 4
      , points V.! 1  ^/ 2
      , points V.! 2 ^/ 2
      , points V.! 0 ^/ 2
      ]
    edges = meshEdges example1

-- T = (U, G)
-- G = [Polygon]
-- U = nub $ concat G

findStarNeighbours :: Eq a => [(a,a,a)] -> a -> [(a, a)]
findStarNeighbours allTrig self =
  [ (b,c)
  | (a,b,c) <- allTrig
  , self == a
  ] ++
  [ (c,a)
  | (a,b,c) <- allTrig
  , self == b
  ] ++
  [ (a,b)
  | (a,b,c) <- allTrig
  , self == c
  ]

isInterior :: Eq a => [(a, a)] -> Bool
isInterior = isJust . getExteriorPoly

getExteriorPoly :: Eq a => [(a, a)] -> Maybe [a]
getExteriorPoly [] = Nothing
getExteriorPoly ((a,b):rest) = worker [a] a b rest
  where
    worker acc start this [] = do
      guard (start == this)
      return (reverse acc)
    worker acc start this xs =
      case lookup this xs of
        Just next -> worker (this:acc) start next (delete (this,next) xs)
        Nothing   ->
          case lookup this (map swap xs) of
            Just next -> worker (this:acc) start next (delete (next, this) xs)
            Nothing   -> Nothing

convexInterpolate :: RelMeshPair -> Double -> RelMesh
convexInterpolate (RelMeshPair static edges leftM leftB rightM rightB) t =
  RelMesh
  { relMeshStatic = static
  , relMeshEdges  = edges
  , relMeshA      = Matrix.scale (1-t) leftM +
                    Matrix.scale t rightM
  , relMeshB      = Matrix.scale (1-t) leftB +
                    Matrix.scale t rightB
  }

solveMesh :: RelMesh -> Mesh
solveMesh (RelMesh static edges m b) =
  case linearSolve m b of
    Nothing -> error "Failed to solve mesh"
    Just ret ->
      Mesh (static <> V.fromList (worker (toLists ret))) edges
  where
    worker []             = []
    worker ([x]:[y]:rest) = V2 x y : worker rest
    worker _              = error "invalid result"

mkRelative :: Mesh -> RelMesh
mkRelative (Mesh points edges) = RelMesh (V.fromList exteriorPoints) edges mM bM
  where
    mM = (s><s) (concat m)
    bM = (s><1) b
    (s,exterior, (m, b)) = toParameters points edges
    exteriorPoints =
      [ points V.! (i-1)
      | i <- exterior
      ]

mkRelativePair :: MeshPair -> RelMeshPair
mkRelativePair (MeshPair p1 p2 edges) =
  let RelMesh static _ leftM leftB = mkRelative (Mesh p1 edges)
      RelMesh _ _ rightM rightB = mkRelative (Mesh p2 edges)
  in RelMeshPair static edges leftM leftB rightM rightB

toParameters :: (Ord b, Fractional b) => V.Vector (V2 b) -> [(Int, Int, Int)] -> (Int, [Int], ([[b]], [b]))
toParameters points groups = (length interior*2,exterior,unzip $ concat
  [ let lst = [(if i == j then -1 else t)
              | j <- interior
              , let t = fromMaybe 0 $ lookup (i,j) lam_ij_cache
              ]
        pos = negate $ sum
           [ pj ^* t
           | j <- exterior
           , let t = fromMaybe 0 $ lookup (i,j) lam_ij_cache
                 pj = points V.! (j-1)
           ]
    in [ (dupX lst, pos ^. _x)
       , (dupY lst, pos ^. _y)]
  | i <- interior ])
  where
    lam_ij_cache = lam_ij points groups
    dupX []     = []
    dupX (x:xs) = x:0:dupX xs
    dupY []     = []
    dupY (x:xs) = 0:x:dupY xs
    (interior, exterior) =
      partition (isInterior . findStarNeighbours groups) [1 .. length points]

lam_ij :: (Ord b, Fractional b) => V.Vector (V2 b) -> [(Int, Int, Int)] -> [((Int, Int), b)]
lam_ij points groups =
  [ ((i, j), t)
  | i <- [1..V.length points]
  , (j, t) <- lam_j points groups i
  ]

lam_j :: (Ord b, Fractional b) => V.Vector (V2 b) -> [(Int, Int, Int)] -> Int -> [(Int, b)]
lam_j points groups p =
  [ (nP, sum [ t | (j,_k,t) <- mu, j == nP ] / fromIntegral (length nPoints))
  | let n = findStarNeighbours groups p
        nPoints = fromMaybe [] $ getExteriorPoly n
        mu = calcMu points groups p
  , nP <- nPoints ]

calcMu :: (Ord c, Fractional c) => V.Vector (V2 c) -> [(Int, Int, Int)] -> Int -> [(Int, Int, c)]
calcMu points groups p = concat
    [ [ (nP, nP, t1)
      , (a, nP, t2)
      , (b, nP, t3) ]
      -- (nP, a, b)
    | {-p <- [1..length points]-}
      let selfVert = points V.! (p-1)
          n = findStarNeighbours groups p
          nPoints :: [Int]
          nPoints = fromMaybe [] $ getExteriorPoly n
    , nP <- nPoints
    , let vert = points V.! (nP-1)
    , let line = (vert, selfVert)
    , let (a,b,aP,bP) = head $
            [ (_a,_b,_aP,_bP)
            | (_a,_b) <- n
            , let _aP = points V.! (_a-1)
                  _bP = points V.! (_b-1)
                  segment = (points V.! (_a-1), points V.! (_b-1))
            , case rayIntersect line segment of
                Nothing -> False
                Just u  -> isBetween u segment
            , _a /= nP
            , _b /= nP ]
    -- , b == (nPoints ++ nPoints) !! i
    , let (t1,t2,t3) = barycentricCoords vert aP bP selfVert
    ]

genColor :: Int -> Int -> PixelRGBA8
genColor n m =
    promotePixel $ parula (fromIntegral (n+offset) / fromIntegral (m+offset))
  where
    offset = 5

polygonNumDots :: Polygon -> Double -> SVG
polygonNumDots p t = mkGroup $ reverse
    [ mkGroup
      [ colored n $ withStrokeWidth (defaultStrokeWidth*0.5) $ withStrokeColor "black" $
        translate x y $ mkCircle circR
      , withFillColor "black" $
        translate x y $ ppNum n ]
    | n <- [0..pSize p-1]
    , let a = realToFrac <$> pAccess p n
          b = realToFrac <$> pAccess p (pNext p n)
          V2 x y = lerp t b a ]
  where
    circR = 0.1
    colored n =
      let c = genColor n (pSize p-1)
      in withFillColorPixel c
    ppNum n = scaleToHeight (circR*1.5) $ center $ latex $ T.pack $ "\\texttt{" ++ show n ++ "}"

polygonShape :: Polygon -> SVG
polygonShape p = mkLinePathClosed
  [ (x,y) | V2 x y <- map (fmap realToFrac) $ V.toList (polygonPoints p) ]


{-# LANGUAGE DataKinds             #-}
{-# LANGUAGE DeriveAnyClass        #-}
{-# LANGUAGE DeriveGeneric         #-}
{-# LANGUAGE FlexibleInstances     #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE RankNTypes            #-}
{-# LANGUAGE ScopedTypeVariables   #-}
{-# LANGUAGE TypeFamilies          #-}
{-# LANGUAGE TypeOperators         #-}
{-# LANGUAGE OverloadedStrings     #-}
{-|
Module      : Grenade.Layers.GlobalAvgPool
Description : Global Average Pooling
Maintainer  : Theo Charalambous
License     : BSD2
Stability   : experimental

The layer follows the implementation as introduced in the paper:
Min Lin, Qiang Chen, Shuicheng Yan. Network In Network
<https://arxiv.org/abs/1312.4400>
-}

module Grenade.Layers.GlobalAvgPool
  (
  -- * Layer definition
    GlobalAvgPool (..)
  )
where

import           Control.DeepSeq                (NFData (..))
import           Data.Serialize
import           Data.Singletons
import           GHC.Generics                   (Generic)
import           GHC.TypeLits

import qualified Numeric.LinearAlgebra.Static   as H

import           Grenade.Core
import           Grenade.Utils.LinearAlgebra

import           Grenade.Onnx

data GlobalAvgPool = GlobalAvgPool
  deriving (Generic, NFData, Show)

instance UpdateLayer GlobalAvgPool where
  type Gradient GlobalAvgPool = ()
  runUpdate _ _ _ = GlobalAvgPool
  reduceGradient _ = ()

instance RandomLayer GlobalAvgPool where
  createRandomWith _ _ = return GlobalAvgPool

instance Serialize GlobalAvgPool where
  put _ = return ()
  get   = return GlobalAvgPool

-- | Performing a global average pool on a 1d vector will take the average of the elements and
--   return a vector of size 1
instance (KnownNat i) => Layer GlobalAvgPool ('D1 i) ('D1 1) where
  type Tape GlobalAvgPool ('D1 i) ('D1 1) = S ('D1 i)

  runForwards _ x =
    let n   = fromIntegral $ natVal (Proxy :: Proxy i)
        avg = nsum x / n
        vec = listToVector [avg]
    in  (x, S1D vec)

  runBackwards = undefined

-- | Performing a global average pool on a 2d matrix will take the average of the elements
--   and return a matrix of size 1x1
instance (KnownNat i, KnownNat j) => Layer GlobalAvgPool ('D2 i j) ('D2 1 1) where
  type Tape GlobalAvgPool ('D2 i j) ('D2 1 1) = S ('D2 i j)

  runForwards _ x =
    let n   = fromIntegral $ natVal (Proxy :: Proxy i)
        m   = fromIntegral $ natVal (Proxy :: Proxy j)
        avg = nsum x / (n * m)
        mat = H.fromList [avg]
    in  (x, S2D mat)

  runBackwards = undefined

-- | Performing a global average pool on a 3d matrix with k channels will take the average of the elements
--   channel-wise and return a 3d matrix with k channels each of size 1x1
instance (KnownNat i, KnownNat j, KnownNat k) => Layer GlobalAvgPool ('D3 i j k) ('D3 1 1 k) where

  type Tape GlobalAvgPool ('D3 i j k) ('D3 1 1 k) = S ('D3 i j k)

  runForwards _ x =
    let n   = fromIntegral $ natVal (Proxy :: Proxy i)
        m   = fromIntegral $ natVal (Proxy :: Proxy j)
        cs  = splitChannels x
        ys  = map (\c -> nsum c / (n * m)) cs
        mat = H.fromList ys
    in  (x, S3D mat)

  runBackwards = undefined

instance OnnxOperator GlobalAvgPool where
  onnxOpTypeNames _ = ["GlobalAveragePool"]

instance OnnxLoadableActivation GlobalAvgPool where
  activationLayer = GlobalAvgPool