-- Building blocks for making fully connected neural networks (FCNs).
--
-- Original author: David Banas <capn.freako@gmail.com>
-- Original date:   January 18, 2018
--
-- Copyright (c) 2018 David Banas; all rights reserved World wide.

{-# OPTIONS_GHC -Wall #-}
{-# OPTIONS_GHC -Wno-unused-top-binds #-}

{-# LANGUAGE AllowAmbiguousTypes #-}
{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE DeriveGeneric #-}
{-# LANGUAGE ExplicitForAll #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE KindSignatures #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE TypeOperators #-}

{-|
Module      : Haskell_ML.FCN
Description : Allows: creation, training, running, saving, and loading,
              of multi-layer, fully connected neural networks.
Copyright   : (c) David Banas, 2018
License     : BSD-3
Maintainer  : capn.freako@gmail.com
Stability   : experimental
Portability : ?
-}
module Haskell_ML.FCN
  ( FCNet(), TrainEvo(..)
  , randNet, runNet, runNet', netTest, hiddenStruct
  , getWeights, getBiases
  , trainNTimes, classificationAccuracy
  ) where

import Control.Monad.Random
import Data.Binary
import Data.List
import Data.Singletons.Prelude
import Data.Singletons.TypeLits
import Data.Vector.Storable (toList, maxIndex)
import GHC.Generics (Generic)
import Numeric.LinearAlgebra.Static hiding (mean)  -- ^ @R@ & @L@ come from here.
-- import Internal.Static (lift1F)  -- Doesn't work; module is hidden.

import Haskell_ML.Util  hiding (classificationAccuracy, maxIndex, getWeights, getBiases)

-- | A fully connected, multi-layer network with fixed input/output
-- widths, but variable (and existentially hidden!) internal structure.
data FCNet :: Nat -> Nat -> * where
  FCNet :: Network i hs o -> FCNet i o

-- | Returns a value of type `FCNet`, filled with random weights
-- ready for training, tucked inside the appropriate Monad, which must
-- be an instance of `MonadRandom` . (IO is such an instance.)
--
-- The input/output widths are determined by the compiler automatically,
-- via type inferencing.
--
-- The internal structure of the network is determined by the list of
-- integers passed in. Each integer in the list indicates the output
-- width of one hidden layer, with the first entry in the list
-- corresponding to the hidden layer nearest to the input layer.
randNet :: (KnownNat i, KnownNat o, MonadRandom m)
        => [Integer]
        -> m (FCNet i o)
randNet hs = withSomeSing hs (fmap FCNet . randNetwork')

-- | Data type for holding training evolution data.
data TrainEvo = TrainEvo
  { accs  :: [Double]                   -- ^ training accuracies
  , diffs :: [([[Double]],[[Double]])]  -- ^ differences of weights/biases, by layer
  }

-- | Train a network on several epochs of the training data, keeping
-- track of accuracy and weight/bias changes per layer, after each.
trainNTimes :: (KnownNat i, KnownNat o)
            => Int           -- ^ Number of epochs
            -> Double        -- ^ learning rate
            -> FCNet i o     -- ^ the network to be trained
            -> [(R i, R o)]  -- ^ the training pairs
            -> (FCNet i o, TrainEvo)
trainNTimes = trainNTimes' [] []

trainNTimes' :: (KnownNat i, KnownNat o)
             => [Double]                    -- accuracies
             -> [([[Double]], [[Double]])]  -- weight/bias differences
             -> Int -> Double -> FCNet i o -> [(R i, R o)] -> (FCNet i o, TrainEvo)
trainNTimes' accs diffs 0 _    net _   = (net, TrainEvo accs diffs)
trainNTimes' accs diffs n rate net prs = trainNTimes' (accs ++ [acc]) (diffs ++ [diff]) (n-1) rate net' prs
  where net'  = trainNet rate net prs
        acc   = classificationAccuracy res ref
        res   = runNet net' $ map fst prs
        ref   = map snd prs
        diff  = ( zipWith (zipWith (-)) (getWeights net') (getWeights net)
                , zipWith (zipWith (-)) (getBiases  net') (getBiases  net) )

-- | Calculate the classification accuracy, given:
--
--   - a list of results vectors, and
--   - a list of reference vectors.
classificationAccuracy :: (KnownNat n) => [R n] -> [R n] -> Double
classificationAccuracy us vs = mean $ cmpr us vs
  where cmpr :: (KnownNat n) => [R n] -> [R n] -> [Double]
        cmpr xs ys = for (zipWith maxComp xs ys) $ \case
                       True  -> 1.0
                       False -> 0.0

        maxComp :: (KnownNat n) => R n -> R n -> Bool
        maxComp u v = maxIndex (extract u) == maxIndex (extract v)

-- | Run a network on a list of inputs.
runNet :: (KnownNat i, KnownNat o)
       => FCNet i o  -- ^ the network to run
       -> [R i]      -- ^ the list of inputs
       -> [R o]      -- ^ the list of outputs
runNet (FCNet n) = map (runNetwork n)

-- | Run a network on a list of inputs,
-- enforcing 4-bit precision for activation outputs.
runNet' :: (KnownNat i, KnownNat o)
       => FCNet i o  -- ^ the network to run
       -> [R i]      -- ^ the list of inputs
       -> [R o]      -- ^ the list of outputs
runNet' (FCNet n) = map (runNetwork' n)

-- | `Binary` instance definition for `FCNet`.
--
-- With this definition, the user of our library is able to use standard
-- `put` and `get` calls, to serialize his created/trained network for
-- future use. And we don't need to provide auxilliary `saveNet` and
-- `loadNet` functions in the API.
instance (KnownNat i, KnownNat o) => Binary (FCNet i o) where
    put = putFCNet
    get = getFCNet

-- | Basic sanity test of our code, taken from Justin's repository.
--
-- Printed output should contain two offset solid circles.
netTest :: MonadRandom m => Double -> Int -> m String
netTest rate n = do
    inps <- replicateM n $ do
      s <- getRandom
      return $ randomVector s Uniform * 2 - 1
    let outs = flip map inps $ \v ->
                 if v `inCircle` (fromRational 0.33, 0.33)
                      || v `inCircle` (fromRational (-0.33), 0.33)
                   then fromRational 1
                   else fromRational 0
    net0 :: Network 2 '[16, 8] 1 <- randNetwork
    let trained = sgd rate (zip inps outs) net0

        outMat = [ [ render (norm_2 (runNetwork trained (vector [x / 25 - 1,y / 10 - 1])))
                   | x <- [0..50] ]
                 | y <- [0..20] ]

        render r | r <= 0.2  = ' '
                 | r <= 0.4  = '.'
                 | r <= 0.6  = '-'
                 | r <= 0.8  = '='
                 | otherwise = '#'

    return $ unlines outMat
  where
    inCircle :: KnownNat n => R n -> (R n, Double) -> Bool
    v `inCircle` (o, r) = norm_2 (v - o) <= r

-- | Returns a list of integers corresponding to the widths of the hidden
-- layers of a `FCNet`.
hiddenStruct :: FCNet i o -> [Integer]
hiddenStruct (FCNet net) = hiddenStruct' net

hiddenStruct' :: Network i hs o -> [Integer]
hiddenStruct' = \case
    W _    -> []
    _ :&~ (n' :: Network h hs' o)
           -> natVal (Proxy @h)
            : hiddenStruct' n'

-- | Returns a list of lists of Doubles, each containing the weights of
-- one layer of the network.
getWeights :: (KnownNat i, KnownNat o) => FCNet i o -> [[Double]]
getWeights (FCNet net) = getWeights' net

getWeights' :: (KnownNat i, KnownNat o) => Network i hs o -> [[Double]]
getWeights' (W Layer{..})       = [concatMap (toList . extract) (toRows nodes)]
getWeights' (Layer{..} :&~ net) = concatMap (toList . extract) (toRows nodes) : getWeights' net

-- | Returns a list of lists of Doubles, each containing the biases of
-- one layer of the network.
getBiases :: (KnownNat i, KnownNat o) => FCNet i o -> [[Double]]
getBiases (FCNet net) = getBiases' net

getBiases' :: (KnownNat i, KnownNat o) => Network i hs o -> [[Double]]
getBiases' (W Layer{..})       = [toList $ extract biases]
getBiases' (Layer{..} :&~ net) = toList (extract biases) : getBiases' net

-----------------------------------------------------------------------
-- All following functions are for internal library use only!
-- They are not exported through the API.
-----------------------------------------------------------------------

-- A single network layer mapping an input of width `i` to an output of
-- width `o`, via simple matrix/vector mult.
data Layer i o = Layer { biases :: !(R o)
                       , nodes  :: !(L o i)
                       }
  deriving (Show, Generic)

instance (KnownNat i, KnownNat o) => Binary (Layer i o)

-- Generates a value of type `Layer i o`, filled with uniformly
-- distributed random values, tucked inside the appropriate Monad, which
-- must be an instance of `MonadRandom`.
--
-- Note: normally distributed values were tried (See commented code, below.)
--       and found not to perform as well as uniformly distributed values.
--
--       This could be due to my use of `logistic`, as opposed to `relu`,
--       as my activation function; I'm not sure.
--       (See code and comments near end of this file.)
randLayer :: forall m i o. (MonadRandom m, KnownNat i, KnownNat o)
          => m (Layer i o)
randLayer = do
  s1 :: Int <- getRandom
  s2 :: Int <- getRandom
  let b = randomVector  s1 Uniform * 2 - 1  -- Need to change to Gaussian equivalent.
      n = uniformSample s2 (-1) 1           -- Need to change to Gaussian equivalent.
  -- PLEASE, KEEP THIS COMMENTED OUT CODE!
  -- let m = eye
  --     b = randomVector s2 Gaussian
  --     n = gaussianSample s1 (takeDiag m) (sym m)
  return $ Layer b n

-- This is the network structure that `FCNet i o` wraps, hiding its
-- internal structure existentially, outside of the library.
data Network :: Nat -> [Nat] -> Nat -> * where
  W     :: !(Layer i o)
        -> Network i '[] o

  (:&~) :: KnownNat h
        => !(Layer i h)
        -> !(Network h hs o)
        -> Network i (h ': hs) o

infixr 5 :&~

-- Generates a value of type `Network i hs o`
-- filled with random weights, ready to begin training.
--
-- Note: `hs` is determined explicitly, via the first argument, while
--       `i` and `o` are determined implicitly, via type inference.
randNetwork :: forall m i hs o. (MonadRandom m, KnownNat i, SingI hs, KnownNat o)
            => m (Network i hs o)
randNetwork = randNetwork' sing

randNetwork' :: forall m i hs o. (MonadRandom m, KnownNat i, KnownNat o)
             => Sing hs -> m (Network i hs o)
randNetwork' = \case
  SNil            -> W     <$> randLayer
  SNat `SCons` ss -> (:&~) <$> randLayer <*> randNetwork' ss

-- Binary instance definition for `Network i hs o`.
putNet :: (KnownNat i, KnownNat o)
       => Network i hs o
       -> Put
putNet = \case
    W w    -> put w
    w :&~ n -> put w *> putNet n

getNet :: forall i hs o. (KnownNat i, KnownNat o)
       => Sing hs
       -> Get (Network i hs o)
getNet = \case
    SNil            -> W    <$> get
    SNat `SCons` ss -> (:&~) <$> get <*> getNet ss

instance (KnownNat i, SingI hs, KnownNat o) => Binary (Network i hs o) where
    put = putNet
    get = getNet sing

putFCNet :: (KnownNat i, KnownNat o)
         => FCNet i o
         -> Put
putFCNet (FCNet net) = do
  put (hiddenStruct' net)
  putNet net

getFCNet :: (KnownNat i, KnownNat o)
         => Get (FCNet i o)
getFCNet = do
  hs <- get
  withSomeSing hs (fmap FCNet . getNet)

runLayer :: (KnownNat i, KnownNat o)
         => Layer i o
         -> R i
         -> R o
runLayer (Layer b n) v = b + n #> v

runNetwork :: (KnownNat i, KnownNat o)
           => Network i hs o
           -> R i
           -> R o
runNetwork = \case
  W w        -> \(!v) -> logistic (runLayer w v)
  (w :&~ n') -> \(!v) -> let v' = logistic (runLayer w v)
                         in runNetwork n' v'

-- 4-bit quantization experiment
runNetwork' :: (KnownNat i, KnownNat o)
            => Network i hs o
            -> R i
            -> R o
runNetwork' = \case
  -- W w -> \(!v) ->
  W Layer{..} -> \(!v) ->
    let vq = quantizeV v
        wq = Layer (quantizeV biases) (quantizeM nodes)
     in quantizeV $ logistic (runLayer wq vq)
  -- (w :&~ n') -> \(!v) ->
  (Layer{..} :&~ n') -> \(!v) ->
    let v' = quantizeV $ logistic (runLayer wq vq)
        vq = quantizeV v
        wq = Layer (quantizeV biases) (quantizeM nodes)
     in runNetwork' n' v'

quantizeV :: KnownNat n => R n -> R n
quantizeV v = dvmap (fromIntegral . floor) (abs v' * 15 + 0.5) / 15
  where v' = dvmap (max 0 . min 1) v

quantizeM :: (KnownNat i, KnownNat o) => L o i -> L o i
quantizeM m = dmmap (fromIntegral . floor) (abs m' * 15 + 0.5) / 15
  where m' = dmmap (max 0 . min 1) m

-- Trains a value of type `FCNet i o`, using the supplied list of
-- training pairs (i.e. - matched input/output vectors).
trainNet :: (KnownNat i, KnownNat o)
         => Double        -- learning rate
         -> FCNet i o     -- the network to be trained
         -> [(R i, R o)]  -- the training pairs
         -> FCNet i o     -- the trained network
trainNet rate (FCNet net) trn_prs = FCNet $ sgd rate trn_prs net

-- Train a network of type `Network i hs o` using a list of training
-- pairs and the Stochastic Gradient Descent (SGD) approach.
sgd :: forall i hs o. (KnownNat i, KnownNat o)
    => Double           -- learning rate
    -> [(R i, R o)]     -- training pairs
    -> Network i hs o   -- network to train
    -> Network i hs o   -- trained network
sgd rate trn_prs net = foldl' (sgdStep rate) net trn_prs

-- Train a network of type `Network i hs o` using a single training pair.
--
-- This code was taken directly from Justin Le's public GitHub archive:
-- https://github.com/mstksg/inCode/blob/43adae31b5689a95be83a72866600033fcf52b50/code-samples/dependent-haskell/NetworkTyped.hs#L77
-- and modified only slightly.
sgdStep :: forall i hs o. (KnownNat i, KnownNat o)
         => Double           -- learning rate
         -> Network i hs o   -- network to train
         -> (R i, R o)       -- training pair
         -> Network i hs o   -- trained network
sgdStep rate net trn_pr = fst $ go x0 net
  where
    x0     = fst trn_pr
    target = snd trn_pr
    go  :: forall j js. KnownNat j
        => R j              -- input vector
        -> Network j js o   -- network to train
        -> (Network j js o, R j)
    go !x (W w@(Layer wB wN))
        = let y    = runLayer w x
              o    = logistic y
              -- the gradient (how much y affects the error)
              --   (logistic' is the derivative of logistic)
              dEdy = logistic' y * (o - target)
              -- new bias weights and node weights
              wB'  = wB - konst rate * dEdy
              wN'  = wN - konst rate * (dEdy `outer` x)
              w'   = Layer wB' wN'
              -- bundle of derivatives for next step
              dWs  = tr wN #> dEdy
          in  (W w', dWs)
    -- Handle the inner layers.
    go !x (w@(Layer wB wN) :&~ n)
        = let y          = runLayer w x
              o          = logistic y
              -- get dWs', bundle of derivatives from rest of the net
              (n', dWs') = go o n
              -- the gradient (how much y affects the error)
              dEdy       = logistic' y * dWs'
              -- new bias weights and node weights
              wB'  = wB - konst rate * dEdy
              wN'  = wN - konst rate * (dEdy `outer` x)
              w'   = Layer wB' wN'
              -- bundle of derivatives for next step
              dWs  = tr wN #> dEdy
          in  (w' :&~ n', dWs)

-- Doesn't work, because the "constructors of R are not in scope."
-- What am I to do, here?!
-- Orphan `Ord` instance, for R n.
-- deriving instance (KnownNat n) => Ord (R n)

-- | Normalize a vector to a probability vector, via softmax.
-- softMax :: (KnownNat n)
--         => R n  -- ^ vector to be normalized
--         -> R n
-- softMax v = exp v / norm_0 v

-- Rectified Linear Unit
-- relu :: (KnownNat n)
--      => R n
--      -> R n
-- relu = max 0

-- relu' :: (KnownNat n)
--       => R n
--       -> R n
-- relu' v = if v > 0 then 1
--                    else 0

-- Logistic non-linear activation function.
logistic :: Floating a => a -> a
logistic x = 1 / (1 + exp (-x))

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


{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE Strict           #-}
module FourierPinwheel.GaussianEnvelopePinwheel where

import           Data.Array.Repa             as R
import           Data.Complex
import           Data.Vector.Generic         as VG
import           Data.Vector.Unboxed         as VU
import           FourierPinwheel.Hypergeo1F1
import           Math.Gamma
import           Pinwheel.FourierSeries2D
import           Utils.Distribution
import           Utils.Parallel

-- The Fourier coefficients of e^{-a^2 r^2} r^\mu
{-# INLINE gaussianPinwheelFourierCoefficients #-}
gaussianPinwheelFourierCoefficients ::
     (RealFloat a, Gamma (Complex a), Enum a)
  => Int
  -> a
  -> a
  -> a
  -> Int
  -> Int
  -> a
  -> a
  -> a
  -> Complex a
gaussianPinwheelFourierCoefficients numR2Freqs periodR2 a sigma angularFreq radialFreq periodEnv phi rho =
  let periodConst = (-2) * pi / log periodEnv
      piRhoPConst = pi * rho / periodR2
      real = pi / periodR2 ^ 2 * piRhoPConst ^ abs angularFreq
      mu =
        (2 + sigma + fromIntegral (abs angularFreq)) :+
        (periodConst * fromIntegral radialFreq)
      alpha = mu / 2
      beta = fromIntegral (1 + abs angularFreq) :+ 0
      z = ((-1) * (piRhoPConst / a) ^ 2) :+ 0
      img =
        (0 :+ (-1)) ^ abs angularFreq * cis (fromIntegral (-angularFreq) * phi) *
        gamma alpha /
        gamma beta *
        (a :+ 0) ** (-mu) *
        hypergeom alpha beta z
   in (real :+ 0) * img

gaussianPinwheel ::
     ( RealFloat a
     , Gamma (Complex a)
     , Enum a
     , Unbox a
     , VG.Vector vector (Complex a)
     , NFData (vector (Complex a))
     )
  => Int
  -> a
  -> a
  -> a
  -> Int
  -> Int
  -> a
  -> a
  -> a
  -> vector (Complex a)
gaussianPinwheel numR2Freqs periodR2 stdR2 sigma thetaFreq rFreq periodEnv stdTheta stdR =
  let zeroVec = VG.replicate (numR2Freqs ^ 2) 0
      a = 1 / (stdR2 * sqrt 2)
   in VG.concat .
      parMap
        rdeepseq
        (\(radialFreq, angularFreq) ->
           if angularFreq == 0
             then let pinwheel =
                        centerHollowArray numR2Freqs $
                        createFrequencyArray
                          numR2Freqs
                          (gaussianPinwheelFourierCoefficients
                             numR2Freqs
                             periodR2
                             a
                             sigma
                             angularFreq
                             radialFreq
                             periodEnv)
                      arr =
                        R.map
                          (* ((gaussian1DFourierCoefficients
                                 (fromIntegral radialFreq)
                                 (log periodEnv)
                                 stdR) :+
                              0)) $
                        centerHollowArray numR2Freqs pinwheel
                   in VG.convert . toUnboxed . computeS $ arr
             else zeroVec) $
      [ (radialFreq, angularFreq)
      | radialFreq <- [-rFreq .. rFreq]
      , angularFreq <- [-thetaFreq .. thetaFreq]
      ]

-- This one has a orientation preference at 0 degree. It is used for Koffka cross problem.
gaussianPinwheel1 ::
     ( RealFloat a
     , Gamma (Complex a)
     , Enum a
     , Unbox a
     , VG.Vector vector (Complex a)
     , NFData (vector (Complex a))
     )
  => Int
  -> a
  -> a
  -> a
  -> Int
  -> Int
  -> a
  -> a
  -> a
  -> vector (Complex a)
gaussianPinwheel1 numR2Freqs periodR2 stdR2 sigma thetaFreq rFreq periodEnv stdTheta stdR =
  let a = 1 / (stdR2 * sqrt 2)
   in VG.concat .
      parMap
        rdeepseq
        (\(radialFreq, angularFreq) ->
           let pinwheel =
                 centerHollowArray numR2Freqs $
                 createFrequencyArray
                   numR2Freqs
                   (gaussianPinwheelFourierCoefficients
                      numR2Freqs
                      periodR2
                      a
                      sigma
                      0
                      radialFreq
                      periodEnv)
               arr =
                 R.map
                   (* ((gaussian1DFreq (fromIntegral angularFreq) stdTheta *
                        gaussian1DFourierCoefficients
                          (fromIntegral radialFreq)
                          (log periodEnv)
                          stdR) :+
                       0)) $
                 centerHollowArray numR2Freqs pinwheel
            in VG.convert . toUnboxed . computeS $ arr) $
      [ (radialFreq, angularFreq)
      | radialFreq <- [-rFreq .. rFreq]
      , angularFreq <- [-thetaFreq .. thetaFreq]
      ]


module Quantum where

--import Data.Kind
--import Control.Applicative
--import Data.Typeable
import Prelude hiding ((^))
--import Data.Singletons
--import Numeric.LinearAlgebra hiding (toInt) -- hmatrix library
import GHC.TypeLits

import Prelim hiding (One)
import Types
import Interface
import DeepEmbedding

import LinTrans2 hiding (cnot)

-- Signature
data QuantumSig exp = MkQubit
type Qubit = MkLType MkQubit

class HasMILL exp => HasQuantum exp where
  new :: Bool -> exp '[] Qubit
  meas :: exp γ Qubit -> exp γ (Lower Bool)
  unitary :: KnownType σ => Unitary σ -> exp γ σ -> exp γ σ

controlBy :: (HasQuantum exp,KnownType σ)
          => Unitary σ -> exp γ (Qubit ⊗ σ) -> exp γ (Qubit ⊗ σ)
controlBy u e = unitary (control u) e

-- -- first element is the control
cnot :: HasQuantum exp
     => exp γ (Qubit ⊗ Qubit) -> exp γ (Qubit ⊗ Qubit)
cnot = controlBy PauliX

-------------------------------------------------------
-- Deep embedding -------------------------------------
-------------------------------------------------------

type instance Effect Deep = DensityMonad
data instance LVal Deep Qubit = QId Int

data QuantumExp :: Sig where
  New     :: Bool -> QuantumExp '[] Qubit
  Meas    :: Deep γ Qubit -> QuantumExp γ (Lower Bool)
  Unitary :: KnownType σ => Unitary σ -> Deep γ σ -> QuantumExp γ σ

instance HasQuantum Deep where
  new = Dom . New
  meas = Dom . Meas
  unitary u = Dom . Unitary u

instance Domain QuantumExp where
  evalDomain (New b) _ = QId <$> newM b
  evalDomain (Meas e) ρ = do QId i <- eval e ρ
                             VPut <$> measM i
  evalDomain (Unitary u e) ρ = do v <- eval e ρ
                                  applyU u v
                                  return v

applyU :: forall σ. KnownType σ => Unitary σ -> LVal Deep σ -> DensityMonad ()
applyU u v = applyMatrix (interpU u) (valToQubits v)

class KnownType σ where
  numQubits :: Int
  valToQubits :: LVal Deep σ -> [Int]
instance KnownType Qubit where
  numQubits = 1
  valToQubits (QId i) = [i]
instance KnownType One where
  numQubits = 1
  valToQubits VUnit = []
instance (KnownType σ1,KnownType σ2) => KnownType (σ1 ⊗ σ2) where
  numQubits = numQubits @σ1 + numQubits @σ2
  valToQubits (VPair v1 v2) = valToQubits v1 ++ valToQubits v2
instance KnownType (Lower α) where
  numQubits = 0
  valToQubits (VPut _) = []

instance Show (Unitary σ) where
  show Identity = "I"
  show Hadamard = "H"
  show PauliX   = "X"
  show PauliY   = "Y"
  show PauliZ   = "Z"
  show (R m)    = "R " ++ show m
  show (Alt u0 u1) = "(" ++ show u0 ++ " ⊕ " ++ show u1 ++ ")"
  show (Transpose u) = show u ++ "†"
-- instance Show (QuantumLExp lang g τ) where
--   show (New b)  = "New(" ++ show b ++ ")"
--   show (Meas q) = "Meas(" ++ show q ++ ")"
--   show (Unitary u e) = "Unitary (" ++ show u ++ ") " ++ show e
-- --  show (ControlBy _ e e') = show e ++ "`ControlBy`" ++ show e'

-- Quantum Data

-- Add more?
data Unitary (σ :: LType) where
  Identity  :: KnownType σ => Unitary σ
  Hadamard  :: Unitary Qubit
  PauliX    :: Unitary Qubit -- (NOT)
  PauliY    :: Unitary Qubit
  PauliZ    :: Unitary Qubit
  R         :: Int -> Unitary Qubit
  Alt       :: KnownType σ => Unitary σ -> Unitary σ -> Unitary (Qubit ⊗ σ)
  Transpose :: Unitary σ -> Unitary σ

control :: KnownType σ  => Unitary σ -> Unitary (Qubit ⊗ σ)
control = Alt Identity

{-
type KnownQubits σ = SingI (NumQubits σ)

unitarySing :: forall σ. Unitary σ -> Sing (NumQubits σ)
unitarySing Identity = sing
unitarySing Hadamard = one
unitarySing PauliX   = one
unitarySing PauliY   = one
unitarySing PauliZ   = one
unitarySing (R _)    = one
unitarySing (Alt u0 _)    = SS $ unitarySing u0
unitarySing (Transpose u) = unitarySing u
-}

-- instance HasQuantumEffect ('Sig DensityMonad sigs) where
-- --  type QUnitary ('Sig DensityMonad sigs) = Density
--type QUnitary σ = Squared (NumQubits σ)

interpU :: forall σ. KnownType σ => Unitary σ -> Matrix
interpU Identity = ident $ numQubits @σ
interpU Hadamard = hadamard
interpU PauliX   = pauliX
interpU PauliY   = pauliY
interpU PauliZ   = pauliZ
interpU (R m)    = undefined -- rotation about the z axis by 2πi/2^m?
interpU (Alt (u0 :: Unitary σ') u1) =
    (newD False `kron` interpU u0) + (newD True `kron` interpU u1)
interpU (Transpose u) = transpose $ interpU u

type family   NumQubits (τ :: LType) :: Nat
type instance NumQubits One            = 0
type instance NumQubits Qubit          = 1
type instance NumQubits (τ1 ⊗ τ2) = NumQubits τ1 + NumQubits τ2

{-
instance HasQuantumDom lang => Domain QuantumDom (lang :: Lang sig) where

  evalDomain _ (New b)   = do
    i <- newQubit @sig b
    return $ vqubit i
  evalDomain ρ (Meas e)  = do
    VQubit i <- toDomain @QuantumDom <$> eval' ρ e
    b <- measQubit @sig i
    return $ vput b
  evalDomain ρ (Unitary u e) = do
    v  <- eval' ρ e
    qs <- valToQubits @sig v
    applyU @sig u qs
    return v

type Qubits (τ :: LType sig) = [QId]

valToQubits :: forall sig (lang :: Lang sig) τ.
              HasQuantumDom lang => LVal lang τ -> SigEffect sig (Qubits τ)
valToQubits v = case toDomain' @QuantumDom v of
    Just (VQubit i) -> return [i]
    Nothing -> case toDomain' @OneDom v of
      Just VUnit -> return []
      Nothing -> case toDomain' @TensorDom v of
        Just (VPair v1 v2) -> liftA2 (++) (valToQubits v1) (valToQubits v2)
        Nothing -> return []
--        case fromLVal' proxyLower v of
--          Just (VPut _) -> return []
--          Nothing       -> error "Cannot extract qubits from the given value"

type family   NumQubits (τ :: LType sig) :: Nat
type instance NumQubits ('LType _ 'OneSig)            = 'Z
type instance NumQubits ('LType _ 'QubitSig)          = 'S 'Z
type instance NumQubits ('LType _ ('TensorSig τ1 τ2)) = NumQubits τ1 `Plus` NumQubits τ2
-}

--type instance NumQubits ('LType _ ('PlusSig τ1 τ2))   = NumQubits τ1 `Plus` NumQubits τ2
type instance NumQubits (Lower _)      = 0

-- EXAMPLES

-- Flip -------------------------------------------

qflip :: Lin Deep Bool
qflip = suspend $ meas (unitary Hadamard (new False))

-- ----------------------------------------------------
-- -- Teleportation -----------------------------------
-- ----------------------------------------------------

plus_minus :: HasQuantum exp => Bool -> Lift exp Qubit
plus_minus b = suspend $ unitary Hadamard $ new b

share :: HasQuantum exp => Lift exp (Qubit ⊸ Qubit ⊗ Qubit)
share = suspend . λ $ \q -> cnot (q ⊗ new False)

bell00 :: HasQuantum exp => Lift exp (Qubit ⊗ Qubit)
bell00 = suspend $ force (plus_minus False) `letin` \a ->
                   force share ^ a

alice :: HasQuantum exp => Lift exp (Qubit ⊸ Qubit ⊸ Lower (Bool,Bool))
alice = suspend . λ $ \q -> λ $ \a ->
            cnot (q ⊗ a) `letPair` \(q,a) ->
            meas (unitary Hadamard q) >! \x ->
            meas a >! \y ->
            put (x,y)

bob :: HasQuantum exp => Bool -> Bool -> Lift exp (Qubit ⊸ Qubit)
bob x y = suspend . λ $ \b ->
    if y then unitary PauliX b else b `letin` \b ->
    if x then unitary PauliZ b else b

teleport :: HasQuantum exp => Lift exp (Qubit ⊸ Qubit)
teleport = suspend . λ $ \q ->
    force bell00 `letPair` \(a,b) ->
    force alice ^ q ^ a >! \(x,y) ->
    force (bob x y) ^ b

teleport0 :: Lin Deep Bool
teleport0 = suspend . meas $ force teleport ^ new False

main' :: Lin Deep Bool
main' = suspend $ force bell00 `letPair` \(q1,q2) ->
                 meas q1 >! \x ->
                 meas q2 >! \y ->
                 put (x == y)

main :: IO ()
main = print $ run main'

-- Dependent fourier transform

type family (σ :: LType) ⊗⊗ (n :: Nat) :: LType where
    σ ⊗⊗ n = CompareOrd (CmpNat n 1)
                        One                -- n = 0
                        σ                  -- n = 1
                        (σ ⊗ (σ ⊗⊗ (n-1))) -- n >= 2
--  σ ⊗⊗ Z = One
--  σ ⊗⊗ S Z = σ
--  σ ⊗⊗ (S (S n)) = σ ⊗ (σ ⊗⊗ S n)

{- TODO: fix
rotations :: forall (m :: Nat) (n :: Nat) exp. (HasQuantum exp)
          => Sing m -> Sing n -> Lift exp (Qubit ⊗⊗ S n ⊸ Qubit ⊗⊗ S n)
rotations _ SZ      = idL
rotations _ (SS SZ) = idL
rotations m (SS n'@(SS _)) = suspend . λpair $ \(c,qs) -> qs `letPair` \(q,qs') ->
    force (rotations m n') ^ (c ⊗ qs') `letPair` \(c,qs') ->
    controlBy (R $ 2 + toInt m - toInt n') (c ⊗ q) `letPair` \(c,q) ->
    c ⊗ (q ⊗ qs')

fourier :: forall n exp. (HasQuantum exp)
        => Sing n -> Lift exp (Qubit ⊗⊗ n ⊸ Qubit ⊗⊗ n)
fourier SZ = idL
fourier (SS SZ) = suspend . λ $ unitary Hadamard
fourier (SS n'@(SS _)) = Suspend . λpair $ \(q,w) ->
        force (fourier n') ^ w `letin` \w ->
        force (rotations (SS n') n') ^ (q ⊗ w)
-}


{-|
Module      : Example.Mnist
Description : MNIST traning example
Copyright   : (c) Anatoly Yakovenko, 2015-2016
License     : MIT
Maintainer  : aeyakovenko@gmail.com
Stability   : experimental
Portability : POSIX

This module implements an example traning the MNIST data set.  Using
the dataset parsing code from:
<https://github.com/mhwombat/backprop-example/blob/master/Mnist.hs>
-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE FlexibleContexts #-}
module Examples.Mnist (generateTrainBatches
                      ,generateTrainLabels
                      ,generateTestBatches
                      ,mnist
                      )where

import Control.Applicative((<|>))
import Control.Monad.Trans(liftIO)
import Control.Monad(when,forever,forM_, foldM_)
import qualified Data.ByteString.Lazy as BL
import Data.Binary.Get hiding(label)
import qualified Data.Binary as B
import Data.Word
import qualified Data.List.Split as S
import qualified Data.Array.Repa as R
import Codec.Compression.GZip as GZ
import Data.List.Split(chunksOf)
import Statistics.LinearRegression as S

import qualified Data.DNN.Trainer as T
import qualified Data.RBM as RB
import qualified Data.Matrix as M
import qualified Data.ImageUtils as I
import Data.Matrix(Matrix(..)
                  ,U
                  ,B
                  ,I
                  ,H
                  )
{-|
Train a network to recognize digits from the MNIST dataset.  Start by training
each layer as an RBM with the Contrastive Divergence algorithm.  Then finish
up the traning with Back-propagation.  Animate the weights as the traning
progresses.
-}
mnist :: IO ()
mnist = do
   let r1 = RB.new 0 785 530 --dimentions include the bias node
       r2 = RB.new 0 530 530 --sqrt(529) == 23
       r3 = RB.new 0 530 11

   -- train the first layer
   tr1 <- B.decodeFile "dist/rbm1" --decode if this layer is already been trained
      <|> do tr1 <- snd <$> (T.run [r1] $ trainCD "dist/rbm1.gif" 0.01)
             B.encodeFile "dist/rbm1" tr1 --save the layer to a file
             return tr1

   -- train the second layer
   tr2 <- B.decodeFile "dist/rbm2"
      <|> do tr2 <- snd <$> (T.run (tr1++[r2]) $ trainCD "dist/rbm2.gif" 0.001)
             B.encodeFile "dist/rbm2" tr2
             return tr2

   -- train the third layer
   tr3 <- B.decodeFile "dist/rbm3"
      <|> do tr3 <- snd <$> (T.run (tr2++[r3]) $ trainCD "dist/rbm3.gif" 0.001)
             B.encodeFile "dist/rbm3" tr3
             return tr3

   -- backprop
   let train pbp xx = do
            let name = "dist/bp" ++ (show xx)
                gif = name ++ ".gif"
            bp <- B.decodeFile name
               <|> do bp <- snd <$> (T.run pbp $ trainBP gif 0.01 0.001)
                      B.encodeFile name bp
                      return bp
            mapM_ (testBatch bp) [0..9] --print how well our network recognizes the images
            return bp
   foldM_ train tr3 [1::Int ..]  --train forever

-- max number of minibatches
maxCount :: Int
maxCount = 25000
-- how often to check the progress
testCount :: Int
testCount = 1000
-- size of the minibatch
rowCount :: Int
rowCount = 5

-- |train the last layer in the DNN via the CD algorithm
trainCD :: String -> Double ->  T.Trainer IO ()
trainCD file mine = forever $ do
  T.setLearnRate 0.001
  let batchids = [0..468::Int]
  forM_ batchids $ \ ix -> do
     big <- liftIO $ readBatch ix
     small <- mapM M.d2u $ M.splitRows rowCount big
     forM_ small $ \ batch -> do
        T.contraDiv batch
        cnt <- T.getCount
        when (0 == cnt `mod` testCount) $ do
           nns <- T.getDNN
           ww <- M.cast1 <$> M.transpose (last nns)
           liftIO $ I.appendGIF file ww
        when (0 == cnt `mod` testCount) $ do
           err <- T.reconErr big
           liftIO $ print (cnt, err)
           when (cnt >= maxCount || err < mine) $ T.finish_

-- |train the entire DNN via backprop
trainBP :: String -> Double -> Double -> T.Trainer IO ()
trainBP file lr mine = forever $ do
  T.setLearnRate lr
  let batchids = [0..468::Int]
  forM_ batchids $ \ ix -> do
     bbatch <- liftIO $ readBatch ix
     blabel <- liftIO $ readLabel ix
     sbatch <- mapM M.d2u $ M.splitRows rowCount bbatch
     slabel <- mapM M.d2u $ M.splitRows rowCount blabel
     forM_ (zip sbatch slabel) $ \ (batch,label) -> do
        T.backProp batch label
        cnt <- T.getCount
        when (0 == cnt `mod` testCount) $ do
           gen <- T.backward (Matrix $ toLabelM [0..9])
           liftIO $ I.appendGIF file gen
        when (0 == cnt `mod` testCount) $ do
           err <- T.forwardErr bbatch blabel
           liftIO $ print (cnt, err)
           when (cnt >= maxCount || err < mine) $ T.finish_

-- |compute the correlation between the labels and the DNN output
testBatch :: [Matrix U I H] -> Int -> IO ()
testBatch nns ix = do
   let name = "dist/test" ++ (show ix)
   bxi <- Matrix <$> readArray name
   let bxh = M.fromList (M.row bxi, 11) $ concat $ replicate (M.row bxi) $ labelVector ix
   (bxh',_) <- T.run nns $ T.feedForward bxi
   let cor = S.correl (M.toUnboxed bxh') (M.toUnboxed bxh)
   print (ix, cor)

-- |Generate a batch of images with bias values for training.
generateTrainBatches :: IO ()
generateTrainBatches = do
   images <- readImages "dist/train-images-idx3-ubyte.gz"
   let batches = map toMatrix $ chunksOf 128 images
   (flip mapM_) (zip [0::Integer ..] batches) $ \ (ix, bb) -> do
      let name = "dist/train" ++ (show ix)
      writeArray name bb

-- |Generate a batch of of matching labels with bias values.
generateTrainLabels :: IO ()
generateTrainLabels = do
   labels <- readLabels "dist/train-labels-idx1-ubyte.gz"
   let batches = map toLabelM $ chunksOf 128 labels
   (flip mapM_) (zip [0::Integer ..] batches) $ \ (ix, bb) -> do
      let name = "dist/label" ++ (show ix)
      writeArray name bb

-- |Generate a the test images with bias values sorted by label.
generateTestBatches :: IO ()
generateTestBatches = do
   images <- readImages "dist/t10k-images-idx3-ubyte.gz"
   labels <- readLabels "dist/t10k-labels-idx1-ubyte.gz"
   (flip mapM_) ([0..9]) $ \ ix -> do
      let name = "dist/test" ++ (show ix)
      let batch = filter (((==) ix) . fst) $ zip labels images
      let bb = toMatrix $ snd $ unzip batch
      writeArray name bb

-- |generate a 2d array from the list of images
-- |add the bias node as the first node to each image
toMatrix :: [Image] -> R.Array R.U R.DIM2 Double
toMatrix images = m
  where
        m = R.fromListUnboxed (R.Z R.:. len R.:. maxsz) (concatMap pixels images)
        maxsz = 1 + (maximum $ map (\ ii -> (iRows ii) * (iColumns ii)) images)
        len = length images
        pixels im = take maxsz $ 1:((normalisedData im) ++ [0..])

-- |file IO utils
readImages :: FilePath -> IO [Image]
readImages filename = do
  content <- GZ.decompress <$> BL.readFile filename
  let (_, _, r, c, unpackedData) = runGet deserialiseHeader content
  return (map (Image (fromIntegral r) (fromIntegral c)) unpackedData)

writeArray :: String -> R.Array R.U R.DIM2 Double -> IO ()
writeArray fileName array = do
   let (R.Z R.:. r R.:. c) = R.extent array
   B.encodeFile fileName (r,c,R.toList array)

readArray ::String -> IO (R.Array R.U R.DIM2 Double)
readArray fileName = do
   (r,c,ls) <- B.decodeFile fileName
   return $ R.fromListUnboxed  (R.Z R.:. r R.:. c) ls

readBatch :: Int -> IO (Matrix U B I)
readBatch ix = Matrix <$> readArray name
   where name = "dist/train" ++ (show ix)

readLabel :: Int -> IO (Matrix U B H)
readLabel ix = Matrix <$> readArray name
   where name = "dist/label" ++ (show ix)

-- |MNIST image parsing code
-- |from https://github.com/mhwombat/backprop-example/blob/master/Mnist.hs
data Image = Image {
      iRows :: Int
    , iColumns :: Int
    , iPixels :: [Word8]
    } deriving (Eq, Show)

normalisedData :: Image -> [Double]
normalisedData image = map normalisePixel (iPixels image)

normalisePixel :: Word8 -> Double
normalisePixel p = (fromIntegral p) / 255.0

toLabelM :: [Int] -> R.Array R.U R.DIM2 Double
toLabelM labels = m
  where
        m = R.fromListUnboxed (R.Z R.:. len R.:. 11) (concatMap labelVector labels)
        len = length labels

labelVector :: Int -> [Double]
labelVector ll = take 11 $ 1.0:(start ++ end)
   where start = take ll $ repeat 0.0
         end = 1.0 : repeat 0.0

{-|
MNIST label file format

[offset] [type]          [value]          [description]
0000     32 bit integer  0x00000801(2049) magic number (MSB first)
0004     32 bit integer  10000            number of items
0008     unsigned byte   ??               label
0009     unsigned byte   ??               label
........
xxxx     unsigned byte   ??               label

The labels values are 0 to 9.
-}
deserialiseLabels :: Get (Word32, Word32, [Word8])
deserialiseLabels = do
  magicNumber <- getWord32be
  count <- getWord32be
  labelData <- getRemainingLazyByteString
  let labels = BL.unpack labelData
  return (magicNumber, count, labels)

readLabels :: FilePath -> IO [Int]
readLabels filename = do
  content <- GZ.decompress <$> BL.readFile filename
  let (_, _, labels) = runGet deserialiseLabels content
  return (map fromIntegral labels)

{-|
MNIST Image file format

[offset] [type]          [value]          [description]
0000     32 bit integer  0x00000803(2051) magic number
0004     32 bit integer  ??               number of images
0008     32 bit integer  28               number of rows
0012     32 bit integer  28               number of columns
0016     unsigned byte   ??               pixel
0017     unsigned byte   ??               pixel
........
xxxx     unsigned byte   ??               pixel

Pixels are organized row-wise. Pixel values are 0 to 255. 0 means background (white), 255
means foreground (black).
-}
deserialiseHeader :: Get (Word32, Word32, Word32, Word32, [[Word8]])
deserialiseHeader = do
  magicNumber <- getWord32be
  imageCount <- getWord32be
  r <- getWord32be
  c <- getWord32be
  packedData <- getRemainingLazyByteString
  let len = fromIntegral (r * c)
  let unpackedData = S.chunksOf len (BL.unpack packedData)
  return (magicNumber, imageCount, r, c, unpackedData)


-- This file is part of Quipper. Copyright (C) 2011-2014. Please see the
-- file COPYRIGHT for a list of authors, copyright holders, licensing,
-- and other details. All rights reserved.
--
-- ======================================================================

-- | This module provides the Jordan-Wigner transformation and
-- symbolic derivation of circuit templates for second quantized
-- interaction terms. It is essentially a fully automated version of
-- the calculations from
--
-- * James D. Whitfield, Jacob Biamonte, and Alán
-- Aspuru-Guzik. \"Simulation of electronic structure Hamiltonians
-- using quantum computers.\"
-- /Molecular Physics/ 109(5):735–750, 2011.
-- See also <http://arxiv.org/abs/1001.3855v3>.

module Algorithms.GSE.JordanWigner where

import Quipper

import QuipperLib.Decompose

import Data.Complex
import qualified Data.Map as Map
import Data.Map (Map)

import Libraries.Auxiliary (sequence_right_)

import Text.Printf

-- ----------------------------------------------------------------------
-- * Overview

-- $ For a given tuple of orbital indices, (/p/,/q/) in case of
-- one-electron interactions, or (/p/,/q/,/r/,/s/) in case of two-electron
-- interactions, we first calculate the Jordan-Wigner transformation
-- of the second quantized hermitian interaction terms
--
-- /a/[sub /p/][sup †]/a/[sub /p/],
--
-- /a/[sub /p/][sup †]/a/[sub /q/] + /a/[sub /q/][sup †]/a/[sub /p/],
--
-- /a/[sub /p/][sup †]/a/[sub /q/][sup †]/a/[sub /q/]/a/[sub /p/],
--
-- /a/[sub /p/][sup †]/a/[sub /q/][sup †]/a/[sub /r/]/a/[sub /s/] +
-- /a/[sub /s/][sup †]/a/[sub /r/][sup †]/a/[sub /q/]/a/[sub /p/].
--
-- Next, we decompose each operator into a linear combination /H/ =
-- λ[sub 1]/M/[sub 1] + ... + λ[sub /n/]/M/[sub /n/] of mutually
-- commuting hermitian tensors. At this point, each summand /M/[sub /j/]
-- in the linear combination will be a tensor product of the following
-- operators (not necessarily in this order):
--
-- * an even number (possibly zero) of Pauli /X/ operators;
--
-- * an even number (possibly zero) of Pauli /Y/ operators;
--
-- * zero or more Pauli /Z/ operators, and
--
-- * zero or more /D/ operators, where /D/ = σ[sup −]σ[sup +] = (/I/−/Z/)\/2.
--
-- Note that there may be zero terms in the summation; this happens,
-- for example, for
-- /a/[sub /p/][sup †]/a/[sub /p/][sup †]/a/[sub /r/]/a/[sub /s/],
-- because two electrons cannot occupy the same spin orbital due to
-- their fermionic nature. In this case, /H/ = 0 and [exp -/i/θ/H/] = /I/.
--
-- Next, we calculate [exp -/i/θ/H/]. Because the summands /M/[sub /j/]
-- commute, we can exponentiate each summand separately, using the formula
-- [exp -/i/θ/H/] = [exp -/i/θλ[sub 1]/M/[sub 1]]⋯[exp -/i/θλ[sub /n/]/M/[sub /n/]].
--
-- We then generate the circuit for [exp -/i/θλ[sub /j/]/M/[sub /j/]]
-- by applying a sequence of basis changes until the problem is
-- reduced to a controlled rotation. The basis changes are, in this
-- order:
--
-- 1. Change each Pauli /X/ operator in /M/[sub /j/] to a Pauli /Z/
-- operator, and apply a Hadamard basis change to the corresponding
-- qubit. This uses the relation /HXH/ = /Z/.
--
-- 2. Change each Pauli /Y/ operator in /M/[sub /j/] to a Pauli /Z/
-- operator, and apply a [bold Y] basis change to the corresponding
-- qubit. Note: the [bold Y] basis change operator is defined in
-- [Whitfield et al.] as /R/[sub /x/](-π\/2) = (/I/+/iX/)\/√2, or
-- equivalently [bold Y] = /SHS/, and satisfies [bold Y][super †]/Y/[bold Y] =
-- /Z/. It should not be confused with the Pauli /Y/ operator.
--
-- 3. If the operator /M/[sub /j/] contains one or more Pauli /Z/ operators
-- (including those obtained in steps 1 and 2), then do a basis change
-- by a cascade of controlled-not gates to reduce this to a single /Z/
-- operator. This uses the relation /CNot/ (/Z/⊗/Z/) /CNot/ = /I/⊗/Z/.
--
-- After these basis changes, the operator /M/[sub /j/] consists of
-- exactly zero or one Pauli /Z/ operator, together with zero or more
-- /D/ operators.  To see how to translate this into a controlled
-- rotation, note that for any operator /A/, we have
--
-- \[image expDA.png]
--
-- Therefore, each /D/ operator in /M/[sub /j/] turns into a control
-- after exponentiation. The final rotation is then computed by
-- distinguishing two cases:
--
-- * If /M/[sub /j/] contains a Pauli /Z/ operator, then use the
-- relation [exp -/i/θ/Z/] = /R/[sub /z/](2θ). In this case, the circuit
-- for [exp -/i/θ/M/[sub /j/]] is a controlled /R/[sub /z/](2θ) gate in
-- the position of the /Z/ operator, with zero or more controls in the
-- positions of any /D/ operators.
--
-- * If /M/[sub /j/] does not contain a Pauli /Z/ operator, then the
-- operation to be performed is a phase change [exp -/i/θ], controlled
-- by the qubits in the positions of the /D/ operators. Note that
-- there must be at least one /D/ operator in this case. Also note
-- that a controlled [exp -/i/θ] gate is identical to a /T/(θ) gate.

-- ----------------------------------------------------------------------
-- * Correctness of the templates

-- $ As outlined above, the functions in this module generate each
-- circuit from first principles, based on the Jordan-Wigner
-- representation of operators and on algebraic transformations. They
-- do not rely on pre-fabricated circuit templates.
--
-- Based on the automated calculations provided by this module, we
-- have found small typos in the 5 templates provided by [Whitfield
-- et al.] (Table 3, or Table A1 in the arXiv version).
--
-- * The template for the number-excitation operator is missing a
-- control on its rotation gate.
--
-- * In the template for the Coulomb operator, the angles are wrong.
-- Moreover, this program finds a simpler template.
--
-- * In the template for the double excitation operator, the angles
-- are wrong; they should be ±θ\/4 instead of θ.
--
-- The corrected templates generated by our code are as follows:
--
-- * Number operator /h/[sub /pp/] /a/[sub /p/][sup †]/a/[sub /p/].
--
-- > [image b0-template.png]
--
-- * Excitation operator /h/[sub /pq/] /a/[sub /p/][sup †]/a/[sub /q/].
--
-- > [image b1-template.png]
--
-- * Coulomb and exchange operators /h/[sub /pqqp/]
-- /a/[sub /p/][sup †]/a/[sub /q/][sup †]/a/[sub /q/]/a/[sub /p/].
--
-- > [image b2-template.png]
--
-- * Number-excitation operator /h/[sub /pqqr/]
-- (/a/[sub /p/][sup †]/a/[sub /q/][sup †]/a/[sub /q/]/a/[sub /r/] +
-- /a/[sub /r/][sup †]/a/[sub /q/][sup †]/a/[sub /q/]/a/[sub /p/]).
-- The sign of ±θ depends on the relative ordering of the indices /p,q,r/.
--
-- > [image b3-template.png]
--
-- * Double excitation operator
-- /h/[sub /pqrs/]
-- (/a/[sub /p/][sup †]/a/[sub /q/][sup †]/a/[sub /r/]/a/[sub /s/] +
-- /a/[sub /s/][sup †]/a/[sub /r/][sup †]/a/[sub /q/]/a/[sub /p/]).
-- The sign of ±θ\/4 in each of the eight terms depends on the relative ordering of the indices /p,q,r,s/.
--
-- > [image b4-template.png]

-- ----------------------------------------------------------------------
-- * Alternate Coulomb templates

-- $ As noted above, our algorithm found the following template for the
-- Coulomb operator
-- /a/[sub /p/][sup †]/a/[sub /q/][sup †]/a/[sub /q/]/a/[sub /p/]:
--
-- > [image b2-template.png]
--
-- This is simpler than the template given in [Whitfield et al.], even
-- after one accounts for the cost of decomposing the additional
-- controlled /T/(θ) gate into elementary gates. However, an
-- equivalent circuit can also be given that is more similar to the
-- one in [Whitfield et al.] (but with corrected rotation angles):
--
-- > [image b2-orthodox.png]
--
-- We call this the \"orthodox\" template, because it is closer to the
-- one specified by Whitfield et al. The program will use the orthodox
-- template if the command line option @--orthodox@ is given, and it
-- will use the simplified template otherwise.

-- ----------------------------------------------------------------------
-- * General-purpose auxiliary functions

-- | Construct a list consisting of /n/ repetitions of some element.
power :: Int -> a -> [a]
power n x = take n $ repeat x

-- | Extract a list of /n/-1 consecutive pairs from an /n/-element list:
--
-- > consecutive_pairs [] = []
-- > consecutive_pairs [1] = []
-- > consecutive_pairs [1,2] = [(1,2)]
-- > consecutive_pairs [1,2,3] = [(1,2),(2,3)]
-- > consecutive_pairs [1,2,3,4] = [(1,2),(2,3),(3,4)]
consecutive_pairs :: [a] -> [(a,a)]
consecutive_pairs [] = []
consecutive_pairs [h] = []
consecutive_pairs (h1:h2:t) = (h1,h2) : consecutive_pairs (h2:t)

-- ----------------------------------------------------------------------
-- * Scalars

-- | The type of complex numbers. Here, we use a floating point
-- representation, although a symbolic representation would also be
-- possible. Since for the purpose of this algorithm, all denominators
-- are powers of 2, the floating point representation is in fact exact.
type Scalar = Complex Double

-- | The complex number /i/.
i :: Scalar
i = 0 :+ 1

-- ----------------------------------------------------------------------
-- * Basic Gates

-- | Apply a /R/[sub z](θ)=[exp -/i/θ/Z/\/2] gate. The parameter θ is a
-- Bloch sphere angle.
--
-- \[image Rz.png]
rotZ_at :: Double -> Qubit -> Circ ()
rotZ_at theta q = named_rotation_at "Rz(%)" theta q

-- | Apply a /G/(θ) gate. This is a global phase change of [exp -/i/θ],
-- so this gate only \"does\" something when it is controlled.
-- Although it is logically a 0-ary gate, we give it a qubit argument
-- to specify where the gate can be drawn in circuit diagrams.
--
-- \[image G.png]
gse_G_at :: Double -> Qubit -> Circ ()
gse_G_at theta q = named_rotation_at "G(%)" theta q

-- | Apply a /T/(θ) gate. This is a /Z/-rotation, but differs
-- from /R/[sub z](-θ) by a global phase.
--
-- \[image T.png]
gse_T_at :: Double -> Qubit -> Circ ()
gse_T_at theta q = named_rotation_at "T(%)" theta q

-- | Apply a [bold Y] basis change gate. This is defined as [bold Y] = /SHS/,
-- or equivalently,
--
-- \[image Y.png]
--
-- This should not be confused with the Pauli /Y/ gate.
gse_Y_at :: Qubit -> Circ ()
gse_Y_at q = named_gate_at "YY" q

-- ----------------------------------------------------------------------
-- * Basic operators

-- | This type provides a symbolic representation of certain
-- operators, generated by the Pauli operators, /P/ = σ[sup +], and
-- /M/ = σ[sup −]. For lack of a better term, we call these the
-- \"basic\" operators. Note that apart from /P/ and /M/, all of these
-- are hermitian.
data Op =
  I -- ^ Identity operator.
  | X  -- ^ Pauli /X/ operator.
  | Y  -- ^ Pauli /Y/ operator.
  | Z  -- ^ Pauli /Z/ operator.
  | P  -- ^ σ[sup +] operator = (0,1;0,0).
  | M  -- ^ σ[sup −] operator = (0,0;1,0).
  | A  -- ^ σ[sup +]σ[sup −] operator = (1,0;0,0).
  | D  -- ^ σ[sup −]σ[sup +] operator = (0,0;0,1).
  deriving (Show, Eq, Ord)

-- | A type to represent scalar multiples. An element of ('Scaled'
-- /a/) is a pair (λ, /x/) of a complex scalar λ and an element /x/ ∈
-- /a/.
data Scaled a = Scaled Scalar a

-- | Multiplication of basic operators. Note that the product of two
-- basic operators is not usually itself a basic operator, but a
-- scalar multiple thereof. This multiplication encodes the algebraic
-- laws of basic operators in symbolic form.

-- Implementation note: the multiplication laws are currently
-- implemented as a long case distinction. Perhaps it could be done
-- more cleverly.
mult :: Op -> Op -> Scaled Op

-- The Pauli group
mult I x = Scaled 1 x
mult x I = Scaled 1 x

mult X X = Scaled 1 I
mult X Y = Scaled i Z
mult X Z = Scaled (-i) Y

mult Y X = Scaled (-i) Z
mult Y Y = Scaled 1 I
mult Y Z = Scaled i X

mult Z X = Scaled i Y
mult Z Y = Scaled (-i) X
mult Z Z = Scaled 1 I

-- P, M, D, and A
mult X P = Scaled 1 D
mult X M = Scaled 1 A
mult X A = Scaled 1 M
mult X D = Scaled 1 P

mult Y P = Scaled i D
mult Y M = Scaled (-i) A
mult Y A = Scaled i M
mult Y D = Scaled (-i) P

mult Z P = Scaled 1 P
mult Z M = Scaled (-1) M
mult Z A = Scaled 1 A
mult Z D = Scaled (-1) D

mult P X = Scaled 1 A
mult M X = Scaled 1 D
mult A X = Scaled 1 P
mult D X = Scaled 1 A

mult P Y = Scaled i A
mult M Y = Scaled (-i) D
mult A Y = Scaled (-i) P
mult D Y = Scaled i M

mult P Z = Scaled (-1) P
mult M Z = Scaled 1 M
mult A Z = Scaled 1 A
mult D Z = Scaled (-1) D

mult P P = Scaled 0 I
mult P M = Scaled 1 A
mult P A = Scaled 0 I
mult P D = Scaled 1 P

mult M P = Scaled 1 D
mult M M = Scaled 0 I
mult M A = Scaled 1 M
mult M D = Scaled 0 I

mult A P = Scaled 1 P
mult A M = Scaled 0 I
mult A A = Scaled 1 A
mult A D = Scaled 0 I

mult D P = Scaled 0 I
mult D M = Scaled 1 M
mult D A = Scaled 0 I
mult D D = Scaled 1 D

-- ----------------------------------------------------------------------
-- * Tensors of basic operators

-- | We use a list of basic operators to represent a tensor
-- product. The convention is that infinitely many identity operators
-- are implicitly appended at the end of the list.
type Tensor = [Op]

-- | Normalize a tensor, by stripping away trailing identities.
normalize_tensor :: Tensor -> Tensor
normalize_tensor [] = []
normalize_tensor (h:t) =
  if h == I && null n
  then []
  else (h:n)
    where n = normalize_tensor t

-- | The identity tensor.
tensor_id :: Tensor
tensor_id = []

-- | Multiply two tensors. This returns a scaled tensor.
mult_tensor :: Tensor -> Tensor -> Scaled Tensor
mult_tensor [] bs = Scaled 1 bs
mult_tensor as [] = Scaled 1 as
mult_tensor (a:as) (b:bs) = Scaled (x*y) (c:cs) where
  Scaled x c = mult a b
  Scaled y cs = mult_tensor as bs

-- | Multiply two scaled tensors.
mult_scaled_tensor :: Scaled Tensor -> Scaled Tensor -> Scaled Tensor
mult_scaled_tensor (Scaled x a) (Scaled y b) = Scaled (x*y*z) c where
  Scaled z c = a `mult_tensor` b

-- ----------------------------------------------------------------------
-- * Linear combinations of tensors

-- | A type to represent complex linear combinations of tensors.
type TensorLC = Map Tensor Scalar

-- | The origin.
lc_zero :: TensorLC
lc_zero = Map.empty

-- | Add a tensor to a linear combination.
lc_insert :: TensorLC -> Scaled Tensor -> TensorLC
lc_insert lc (Scaled lambda t) =
  if newvalue == 0
  then
    Map.delete m lc
  else
    Map.insert m newvalue lc
      where
        m = normalize_tensor t
        newvalue = case Map.lookup m lc of
          Nothing -> lambda
          Just x -> lambda + x

-- | Turn a list of scaled tensors into a 'TensorLC'.
lc_from_list :: [Scaled Tensor] -> TensorLC
lc_from_list = foldl lc_insert lc_zero

-- | Turn a 'TensorLC' into a list of scaled tensors.
lc_to_list :: TensorLC -> [Scaled Tensor]
lc_to_list lc = [Scaled x y | (y,x) <- Map.toList lc]

-- ----------------------------------------------------------------------
-- * Jordan-Wigner representation

-- $ The next two functions provide the Jordan-Wigner representation of
-- (Fock-space) annihilation and creation operators.

-- | Construct the Jordan-Wigner annihilation operator /a/[sub /p/] =
-- /IIIIPZZZZZ.../ for spin-orbital index /p/. The first parameter is
-- /p/, and the second one is /M/ (the number of spin-orbitals).
-- Precondition: 0 ≤ /p/ < /M/.
jw :: Int -> Int -> Scaled Tensor
jw p m = Scaled 1 (power p I ++ [P] ++ power (m-p-1) Z)

-- | Construct the Jordan-Wigner creation operator /a/[sub /p/][sup †]
-- = /IIIIMZZZZ.../ for spin-orbital index /p/.  The first parameter
-- is /p/, and the second one is /M/ (the number of spin-orbitals).
-- Precondition: 0 ≤ /p/ < /M/.
jw_dagger :: Int -> Int -> Scaled Tensor
jw_dagger p m = Scaled 1 (power p I ++ [M] ++ power (m-p-1) Z)

-- ----------------------------------------------------------------------
-- * Second quantized interaction terms

-- ** Simple interaction terms

-- | Construct the one-electron second quantized non-hermitianized
-- interaction term /a/[sub /p/][sup †]/a/[sub /q/].  The parameters
-- are /p,q/.
one_electron_operator_simple :: Int -> Int -> Scaled Tensor
one_electron_operator_simple p q = ap * aq
  where
    ap = jw_dagger p m
    aq = jw q m
    m = maximum [p,q] + 1
    (*) = mult_scaled_tensor

-- | Construct the two-electron second quantized non-hermitianized
-- interaction term
-- /a/[sub /p/][sup †]/a/[sub /q/][sup †]/a/[sub /r/]/a/[sub /s/].
-- The parameters are /p,q,r,s/.
two_electron_operator_simple :: Int -> Int -> Int -> Int -> Scaled Tensor
two_electron_operator_simple p q r s = ap * aq * ar * as
  where
    ap = jw_dagger p m
    aq = jw_dagger q m
    ar = jw r m
    as = jw s m
    m = maximum [p,q,r,s] + 1
    (*) = mult_scaled_tensor

-- ** Hermitian interaction terms

-- | Construct
-- /a/[sub /p/][sup †]/a/[sub /q/]
-- if /p/ = /q/, and
-- /a/[sub /p/][sup †]/a/[sub /q/] + /a/[sub /q/][sup †]/a/[sub /p/]
-- otherwise.
one_electron_operator :: Int -> Int -> TensorLC
one_electron_operator p q =
  if p == q
  then lc_from_list [a_pq]
  else lc_from_list [a_pq, a_qp]
    where
      a_pq = one_electron_operator_simple p q
      a_qp = one_electron_operator_simple q p

-- | Construct
-- /a/[sub /p/][sup †]/a/[sub /q/][sup †]/a/[sub /r/]/a/[sub /s/]
-- if (/p/,/q/) = (/s/,/r/), and
-- /a/[sub /p/][sup †]/a/[sub /q/][sup †]/a/[sub /r/]/a/[sub /s/] +
-- /a/[sub /s/][sup †]/a/[sub /r/][sup †]/a/[sub /q/]/a/[sub /p/]
-- otherwise.
two_electron_operator :: Int -> Int -> Int -> Int -> TensorLC
two_electron_operator p q r s =
  if (p,q) == (s,r)
  then lc_from_list [a_pqrs]
  else lc_from_list [a_pqrs, a_srqp]
    where
    a_pqrs = two_electron_operator_simple p q r s
    a_srqp = two_electron_operator_simple s r q p

-- ----------------------------------------------------------------------
-- * /XYZD/ decomposition

-- | Decompose a basic operator into linear combinations of hermitian basic operators.
-- This uses the relations /P/ = 1\/2 /X/ + /i/\/2 /Y/ and
-- /M/ = 1\/2 /X/ - /i/\/2 /Y/.
decompose_basis :: Op -> [Scaled Op]
decompose_basis P = [Scaled (1/2) X, Scaled (i/2) Y]
decompose_basis M = [Scaled (1/2) X, Scaled (-i/2) Y]
decompose_basis x = [Scaled 1 x]  -- default case: the remaining operators are already hermitian

-- | Decompose a tensor into a linear combination of hermitian
-- tensors. Due to sign alternation, the individual tensors all come
-- out to commute with each other.
decompose_tensor :: Tensor -> TensorLC
decompose_tensor [] = lc_from_list [Scaled 1 tensor_id]
decompose_tensor (h:t) =
  lc_from_list [Scaled (x*y) (g:gs) |
                Scaled x g <- decompose_basis h,
                Scaled y gs <- lc_to_list (decompose_tensor t)]

-- | Decompose a linear combination of tensors into a linear
-- combination of hermitian tensors.
decompose_tensor_lc :: TensorLC -> TensorLC
decompose_tensor_lc lc =
  lc_from_list [ Scaled (x*y) g |
                 Scaled x gs <- lc_to_list lc,
                 Scaled y g <- lc_to_list (decompose_tensor gs)]

-- ----------------------------------------------------------------------
-- * Exponentiation and circuit generation

-- | Given a simple hermitian tensor /H/ and an angle θ, generate a
-- circuit for [exp -/i/θ/H/].  The given list of input qubits is in
-- the same order as the operators in /H/. Precondition: /H/ is made
-- up of zero or more identity operators and one or more of the
-- operators /X/, /Y/, /Z/, and /D/. The last parameter is a list of
-- additional controls.

exponentiate_simple :: Scaled Tensor -> Double -> [Qubit] -> [Qubit] -> Circ ()
exponentiate_simple (Scaled s ms) theta qs ctl = do
  -- First analyze the tensor:
  let
    -- Find all X positions
    xs = [ i | (m, i) <- zip ms [0,1..], m == X ]
    -- Find all Y positions
    ys = [ i | (m, i) <- zip ms [0,1..], m == Y ]
    -- Find all X, Y, Z positions
    zs = [ i | (m, i) <- zip ms [0,1..], m `elem` [X, Y, Z] ]
    -- Find all D positions
    ds = [ i | (m, i) <- zip ms [0,1..], m == D ]
  basischange xs ys zs qs
  rotation alpha ds zs `controlled` ctl
  reverse_generic_imp (basischange xs ys zs) qs
  where
    alpha = theta * realPart s
    basischange :: [Int] -> [Int] -> [Int] -> [Qubit] -> Circ ()
    basischange xs ys zs qs = do
      -- for every X or Y in the operator, apply the appropriate basis
      -- change to change it to Z.
      sequence_ [ hadamard_at (qs !! i) | i <- xs ]
      sequence_ [ gse_Y_at (qs !! i) | i <- ys ]
      -- apply a cascade of c-not operators to all X, Y, or Z in the
      -- operator
      sequence_right_ [ qnot_at (qs !! i0) `controlled` (qs !! i1) | (i0,i1) <- consecutive_pairs zs]
    rotation :: Timestep -> [Int] -> [Int] -> Circ ()
    rotation alpha ds zs = do
      case zs of
        [] -> -- if there are no Z operators, produce a controlled
              -- e^{-iα} gate.
          case ds of
            [] -> error "exponentiate_simple: precondition violated"
            d:ds' -> gse_T_at alpha (qs !! d) `controlled` [qs !! d' | d' <- ds']
        z:zs' -> -- if there are Z operators, produce a controlled
                 -- e^{-iαZ} gate.
          rotZ_at (2*alpha) (qs !! z) `controlled` [qs !! d | d <- ds]

-- | Given a tensor /H/ (already decomposed into commuting simple
-- tensors) and an angle θ, generate a circuit for [exp -/i/θ/H/].
-- The given list of input qubits is in the same order as the
-- operators in /H/.
exponentiate :: TensorLC -> Double -> [Qubit] -> [Qubit] -> Circ ()
exponentiate lc theta qs ctl =
  sequence_ [ exponentiate_simple a theta qs ctl | a <- lc_to_list lc ]

-- ----------------------------------------------------------------------
-- * Generate top-level templates

-- | @'one_electron_circuit' theta p q@: Generate the circuit for the
-- hermitianized one-electron interaction with spin-orbital indices
-- /p/, /q/. More precisely, generate [exp -/i/θ/H/], where
-- /H/ = /a/[sub /p/][sup †]/a/[sub /q/] if /p/ = /q/ and
-- /H/ = /a/[sub /p/][sup †]/a/[sub /q/]
-- + /a/[sub /q/][sup †]/a/[sub /p/] otherwise.
--
-- This function recognizes an important special case: if θ=0.0, don't
-- generate any gates at all. The case θ=0.0 frequently arises because
-- of the conversion from spatial orbitals to spin orbitals.
one_electron_circuit :: Double -> Int -> Int -> [Qubit] -> [Qubit] -> Circ ()
one_electron_circuit 0.0 p q qs ctl = return ()
one_electron_circuit theta p q qs ctl = do
  comment_with_label (printf "ENTER: one_electron_circuit (theta=%.3e, p=%d, q=%d)" theta p q) (qs,ctl) ("qs","ctl")
  exponentiate op theta qs ctl
  comment_with_label "EXIT: one_electron_circuit" (qs,ctl) ("qs","ctl")
  where
    op = decompose_tensor_lc (one_electron_operator p q)

-- | @'two_electron_circuit' theta p q r s@:
-- Generate the circuit for the hermitianized two-electron interaction
-- with spin-orbital indices /p/, /q/, /r/, /s/. More precisely, generate
-- [exp -/i/θ/H/], where
-- /H/ = /a/[sub /p/][sup †]/a/[sub /q/][sup †]/a/[sub /r/]/a/[sub /s/] if
-- (/p/,/q/) = (/s/,/r/) and /H/ =
-- /a/[sub /p/][sup †]/a/[sub /q/][sup †]/a/[sub /r/]/a/[sub /s/] +
-- /a/[sub /s/][sup †]/a/[sub /r/][sup †]/a/[sub /q/]/a/[sub /p/]
-- otherwise.
--
-- This function recognizes an important special case: if θ=0.0, don't
-- generate any gates at all. The case θ=0.0 frequently arises because
-- of the conversion from spatial orbitals to spin orbitals.
two_electron_circuit :: Double -> Int -> Int -> Int -> Int -> [Qubit] -> [Qubit] -> Circ ()
two_electron_circuit 0.0 p q r s qs ctl = return ()
two_electron_circuit theta p q r s qs ctl = do
  comment_with_label (printf "ENTER: two_electron_circuit (theta=%.3e, p=%d, q=%d, r=%d, s=%d)" theta p q r s) (qs,ctl) ("qs","ctl")
  exponentiate op theta qs ctl
  comment_with_label "EXIT: two_electron_circuit" (qs,ctl) ("qs","ctl")
  where
    op = decompose_tensor_lc (two_electron_operator p q r s)

-- | Like 'two_electron_circuit', but use the \"orthodox\" circuit
-- template for the Coulomb operator /a/[sub /p/][sup †]/a/[sub
-- /q/][sup †]/a/[sub /q/]/a/[sub /p/]. This generates a circuit using
-- three rotations, similar to [Whitfield et al.], but with corrected
-- angles,
--
-- > [image b2-orthodox.png]
--
-- instead of the simpler circuit that 'two_electron_circuit' would
-- normally generate:
--
-- > [image b2-template.png]

two_electron_circuit_orthodox :: Double -> Int -> Int -> Int -> Int -> [Qubit] -> [Qubit] -> Circ ()
two_electron_circuit_orthodox 0.0 p q r s qs ctl = return ()
two_electron_circuit_orthodox theta p q r s qs ctl | q==r && p==s && p/=q = do
  comment_with_label (printf "ENTER: two_electron_circuit_orthodox(theta=%.3e, p=%d, q=%d, r=%d, s=%d)" theta p q r s) (qs,ctl) ("qs","ctl")
  let pp = qs !! p
      qq = qs !! q
  gse_G_at (theta/4) pp `controlled` ctl
  rotZ_at (-theta/2) pp `controlled` ctl
  rotZ_at (-theta/2) qq `controlled` ctl
  qnot_at qq `controlled` pp  -- control ctl is not needed
  rotZ_at (theta/2) qq `controlled` ctl
  qnot_at qq `controlled` pp  -- control ctl is not needed
  comment_with_label "EXIT: two_electron_circuit_orthodox" (qs,ctl) ("qs","ctl")

-- all other cases aren't Coulomb operators, so fall back to
-- two_electron_circuit.
two_electron_circuit_orthodox theta p q r s qs ctl = do
  two_electron_circuit theta p q r s qs ctl

-- ----------------------------------------------------------------------
-- * Testing

-- $ We provide two functions, accessible via command line options,
-- that allow the user to display individual templates.

-- | Display the circuit for the hermitianized one-electron interaction,
-- with θ=1.
show_one_electron :: Format -> GateBase -> Int -> Int -> IO ()
show_one_electron format gatebase p q =
  print_generic format (decompose_generic gatebase circuit) (replicate (n-m) qubit) where
    circuit qs = one_electron_circuit 1.0 (p-m) (q-m) qs []
    n = maximum [p,q] + 1
    m = minimum [p,q]

-- | Display the circuit for the hermitianized two-electron interaction,
-- with θ=1.
show_two_electron :: Format -> GateBase -> Int -> Int -> Int -> Int -> IO ()
show_two_electron format gatebase p q r s =
  print_generic format (decompose_generic gatebase circuit) (replicate (n-m) qubit) where
    circuit qs = two_electron_circuit 1.0 (p-m) (q-m) (r-m) (s-m) qs []
    n = maximum [p,q,r,s] + 1
    m = minimum [p,q,r,s]

-- | Like 'show_two_electron', but use the \"orthodox\" template for
-- the Coulomb operator.
show_two_electron_orthodox :: Format -> GateBase -> Int -> Int -> Int -> Int -> IO ()
show_two_electron_orthodox format gatebase p q r s =
  print_generic format (decompose_generic gatebase circuit) (replicate (n-m) qubit) where
    circuit qs = two_electron_circuit_orthodox 1.0 (p-m) (q-m) (r-m) (s-m) qs []
    n = maximum [p,q,r,s] + 1
    m = minimum [p,q,r,s]


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

module Examples.MultiTargetTracking where

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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


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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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


{-# LANGUAGE ScopedTypeVariables #-}

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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