{-# LANGUAGE DataKinds, RankNTypes, TypeFamilies #-}
module TestCommon ((+\), (*\), (**\),
                   listsToParameters,
                   cmpTwo, cmpTwoSimple,
                   qcPropDom, quickCheckTest0, fquad, quad
                  ) where

import Prelude

import qualified Data.Array.DynamicS as OT
import qualified Data.Array.ShapedS as OS
import qualified Data.Vector.Generic as V
import           Numeric.LinearAlgebra (Vector)
import qualified Numeric.LinearAlgebra as HM
import           Test.Tasty
import           Test.Tasty.QuickCheck

import HordeAd hiding (sumElementsVectorOfDual)
import HordeAd.Core.DualClass (Dual)

(+\) :: DualMonad d r m
     => DualNumber d r -> DualNumber d r -> m (DualNumber d r)
(+\) u v = returnLet $ u + v

(*\) :: DualMonad d r m
     => DualNumber d r -> DualNumber d r -> m (DualNumber d r)
(*\) u v = returnLet $ u * v

(**\) :: DualMonad d r m
      => DualNumber d r -> DualNumber d r -> m (DualNumber d r)
(**\) u v = returnLet $ u ** v

-- Checks if 2 numbers are close enough.
close1 :: forall r. (Ord r, Fractional r)
          => r -> r -> Bool
close1 a b = abs (a - b) <= 1e-4

-- Checks if 2 number pairs are close enough.
close2 :: forall r. (Ord r, Fractional r)
          => (r,r) -> (r,r) -> Property
close2 (a1, b1) (a2, b2) = close1 a1 a2 .&&. close1 b1 b2

quad :: DualMonad d r m
     => DualNumber d r -> DualNumber d r -> m (DualNumber d r)
quad x y = do
  x2 <- returnLet $ square x
  y2 <- y *\ y
  tmp <- x2 +\ y2
  tmp +\ 5

fquad :: forall r d m. DualMonad d r m
      => DualNumberVariables d r -> m (DualNumber d r)
fquad variables = do
  let x = var0 variables 0
      y = var0 variables 1
  quad x y

listsToParameters :: forall r. (OT.Storable r)
                  => ([r], [r]) -> Domains r
listsToParameters (a0, a1) =
  (V.fromList a0, V.singleton $ V.fromList a1, V.empty, V.empty)

listsToParameters4 :: ([Double], [Double], [Double], [Double]) -> Domains Double
listsToParameters4 (a0, a1, a2, aX) =
  ( V.fromList a0
  , V.singleton $ V.fromList a1
  , if null a2 then V.empty else V.singleton $ HM.matrix 1 a2
  , if null aX then V.empty else V.singleton $ OT.fromList [length aX] aX )

quickCheckTest0 :: TestName
       -> (forall d r m. ( DualMonad d r m
                         , Floating (Out (DualNumber d (Vector r)))
                         , Floating (Out (DualNumber d (OS.Array '[2] r))) )
           => DualNumberVariables d r -> m (DualNumber d r))
       -> ((Double, Double, Double) -> ([Double], [Double], [Double], [Double]))
       -> TestTree
quickCheckTest0 txt f fArg =
  qcTestRanges txt f (listsToParameters4 . fArg) ((-2, -2, -2), (2, 2, 2)) ((-1e-7, -1e-7, -1e-7), (1e-7, 1e-7, 1e-7)) (-10, 10)

-- A quick check to compare the derivatives and values of 2 given functions.
cmpTwo
  :: ( m ~ DualMonadForward r, d ~ 'DModeDerivative, Dual d r ~ r
     , DualMonad d r m )
  => (DualNumberVariables d r -> m (DualNumber d r))
  -> (DualNumberVariables d r -> m (DualNumber d r))
  -> Domains r
  -> Domains r
  -> Domains r
  -> Domains r
  -> Property
cmpTwo f1 f2 params1 params2 ds1 ds2 =
  close2 (dFastForward f1 params1 ds1) (dFastForward f2 params2 ds2)

-- A quick check to compare the derivatives and values of 2 given functions.
cmpTwoSimple
  :: ( m ~ DualMonadForward r, d ~ 'DModeDerivative, Dual d r ~ r
     , DualMonad d r m )
  => (DualNumberVariables d r -> m (DualNumber d r))
  -> (DualNumberVariables d r -> m (DualNumber d r))
  -> Domains r
  -> Domains r
  -> Property
cmpTwoSimple f1 f2 parameters ds =
  cmpTwo f1 f2 parameters parameters ds ds

-- A quick consistency check of all the kinds of derivatives and gradients
-- and all kinds of computing the value of the objective function.
qcPropDom :: (forall d r m. ( DualMonad d r m
                            , r ~ Double
                            , Floating (Out (DualNumber d (Vector r)))
                            , Floating (Out (DualNumber d (OS.Array '[2] r))) )
              => DualNumberVariables d r -> m (DualNumber d r))
       -> Domains Double
       -> Domains Double
       -> Domains Double
       -> Double
       -> Property
qcPropDom f args ds perturbation dt =
      let ff@(derivative, ffValue) = dFastForward f args ds
          (derivativeAtPerturbation, valueAtPerturbation) = dFastForward f args perturbation
          (gradient, revValue) = dReverse dt f args
      in -- Two forward derivative implementations agree fully:
         dForward f args ds === ff
         -- Objective function value from gradients is the same.
         .&&. ffValue == revValue
         -- Gradients and derivatives agree.
         .&&. close1 (dt * derivative)
                     (dotParameters gradient ds)
         -- Objective function value is unaffected by perturbation.
         .&&. ffValue == valueAtPerturbation
         -- Derivative approximates the perturbation of value.
         .&&. close1 (primalValue
                                  f (addParameters
                                                   args perturbation))
                     (ffValue + derivativeAtPerturbation)

-- A quick consistency check of all the kinds of derivatives and gradients
-- and all kinds of computing the value of the objective function.
qcPropFArg :: (forall d r m. ( DualMonad d r m
                         , Floating (Out (DualNumber d (Vector r)))
                         , Floating (Out (DualNumber d (OS.Array '[2] r))) )
               => DualNumberVariables d r -> m (DualNumber d r))
       -> ((Double, Double, Double) -> Domains Double)
       -> (Double, Double, Double)
       -> (Double, Double, Double)
       -> (Double, Double, Double)
       -> Double
       -> Property
qcPropFArg f fArgDom xyz dsRaw perturbationRaw dt =
      let args = fArgDom xyz
          ds = fArgDom dsRaw
          perturbation = fArgDom perturbationRaw
      in qcPropDom f args ds perturbation dt

-- A quick consistency check of all the kinds of derivatives and gradients
-- and all kinds of computing the value of the objective function.
qcTestRanges :: TestName
       -> (forall d r m. ( DualMonad d r m
                         , Floating (Out (DualNumber d (Vector r)))
                         , Floating (Out (DualNumber d (OS.Array '[2] r))) )
           => DualNumberVariables d r -> m (DualNumber d r))
       -> ((Double, Double, Double) -> Domains Double)
       -> ((Double, Double, Double), (Double, Double, Double))
       -> ((Double, Double, Double), (Double, Double, Double))
       -> (Double, Double)
       -> TestTree
qcTestRanges txt f fArgDom dsRange perturbationRange dtRange =
  testProperty txt $
  forAll (choose dsRange) $ \xyz dsRaw ->
  forAll (choose perturbationRange) $ \perturbationRaw ->
  forAll (choose dtRange) $ \dt ->
  qcPropFArg f fArgDom xyz dsRaw perturbationRaw dt


{-# LANGUAGE ConstraintKinds #-}

module Element
( Element
, StructElem(..)
, getElementNodes
, getNumNodes
, getConnectivity
, getElementNumber
, computeJacobian
, computeJacobianDet
, dndx
) where

import qualified Numeric.LinearAlgebra         as L
import qualified Numeric.LinearAlgebra.HMatrix as HMat

import           Basis
import           Node
import           Quadrature
import           ShapeFcns

-- Utility functions
allCombinations :: [Int] -> [[Int]]
allCombinations x = mapM (const x) [1..(length x)]

-- Synonym for types required by HMatrix
type FracElNum a = (Fractional a, L.Element a, L.Numeric a)

-- Element types
data StructElem a = StructElem [Node a] Int deriving (Show,Eq)
data Simplex a    = Simplex    [Node a] Int deriving (Show,Eq)

-- This typeclass defines which functions operate for any element.  Different
-- element types are defined as instances of this typeclass.
class Element e where
  getElementNodes    :: (Fractional a) => e a -> [Node a]
  getNumNodes        :: (Fractional a) => e a -> Int
  getConnectivity    :: (Fractional a) => e a -> [Int]
  getElementNumber   :: (Fractional a) => e a -> Int
  faceNormals        :: (Fractional a) => e a -> Int -> [a]
  computeJacobian    :: (Basis b,ShapeFcn s,FracElNum a) => e a -> s b -> [a] -> L.Matrix a
  computeJacobianDet :: (Basis b,ShapeFcn s,FracElNum a,L.Field a) => e a -> s b -> [a] -> a
  dndx               :: (Basis b,ShapeFcn s,FracElNum a,L.Field a) => e a -> s b -> [a] -> Int -> L.Vector a

instance Element StructElem where

  getElementNodes  (StructElem nodes _)   = nodes
  getNumNodes      (StructElem nodes _)   = length nodes
  getConnectivity  (StructElem nodes _)   = map nodeNumber nodes
  getElementNumber (StructElem _ elemNum) = elemNum

  computeJacobian  (StructElem nodes elemNum) shpFcn coords = HMat.mul matA matB
    where
      dim  = getDimension shpFcn
      matA = L.tr $ L.fromLists $ map (dndXi shpFcn coords) [0..(getShapeFcnOrder shpFcn + 1)^dim - 1]
      matB = L.fromLists $ map nodeCoordinates nodes

  computeJacobianDet (StructElem nodes elemNum) shpFcn coords = L.det $ computeJacobian (StructElem nodes elemNum) shpFcn coords
  dndx elem shpFcn coords idx = HMat.app (HMat.inv (computeJacobian elem shpFcn coords)) (L.fromList $ dndXi shpFcn coords idx)

  faceNormals (StructElem nodes elemNum) 1 = [-1.0, 1.0]  -- Specialized for linear elements
  faceNormals _ _ = error "Face normals can only be computed for 1D elements for now."


{-# LANGUAGE FlexibleInstances #-}
{-# OPTIONS -Wall #-}
module NN.NeuralNetwork
  (
    NN (..)
  , NNVars
  , createNetworkW
  , Net (..)
  , module Util.Vars
  ) where

import AI.HNN.FF.Network
import Numeric.LinearAlgebra.HMatrix hiding (corr)
import System.Random (RandomGen, random)
import qualified Data.Vector as V
import Control.Monad.Reader
import RandomUtil.Random
import Util.Vars

import Matrix.Traversable
import Data.Traversable

type NNVars = Vars

class NN n where
  create :: NNVars -> IO n
  get :: n -> Vector Double -> Vector Double
  train :: n -> Sample Double -> n
  train net samp = trainV net [samp]
  trainV :: n -> Samples Double -> n
  randomize :: (RandomGen g) => g -> n -> (g,n)
  getWeights :: n -> V.Vector (Matrix Double)
  getVars :: n -> NNVars

data Net = Net
  {
    _network :: Network Double
  , _vars    :: NNVars
  }

instance NN Net where
  create vars = do
    n <- createNetworkW vars
    return $ Net n vars
  get net input = getNetwork (_vars net) (_network net) input
  trainV net io = let vars = _vars net in Net (trainNetwork vars (_network net) io) vars
  randomize g net = let vars = _vars net
                        (g',n) = randomizeNetwork vars g (_network net)
                    in (g', Net n vars)
  getWeights (Net (Network m) _) = m
  getVars (Net _ vars) = vars

createNetworkW :: NNVars -> IO (Network Double)
createNetworkW env = runReader createNetworkReader env

createNetworkReader :: Reader NNVars (IO (Network Double))
createNetworkReader = do
  numInput      <- asks (getVar numberInputsS)
  numHidden     <- asks (getVar numberHiddenS)
  numOutput     <- asks (getVar numberOutputsS)
  return $ createNetwork (round numInput) [(round numHidden)] (round numOutput)

getNetwork :: NNVars -> Network Double -> Vector Double -> Vector Double
getNetwork vars net input = runReader (getNetworkReader net input) vars

getNetworkReader :: Network Double -> Vector Double -> Reader NNVars (Vector Double)
getNetworkReader net input = do
  sigOrTanh <- asks (getVar sigmoidTanS)
  let (act, _) = getActivationFunctions sigOrTanh
  return $ output net act input

trainNetwork :: NNVars -> Network Double -> Samples Double -> Network Double
trainNetwork vars net samples = runReader (trainNetworkReader net samples) vars

trainNetworkReader :: Network Double -> Samples Double -> Reader NNVars (Network Double)
trainNetworkReader net samples = do
  timesToTrain  <- asks (getVar timesToTrainS)
  learningRate  <- asks (getVar learningRateS)
  sigOrTanh     <- asks (getVar sigmoidTanS)
  let (act, act') = getActivationFunctions sigOrTanh
  return $ trainNTimes (round timesToTrain) learningRate act act' net samples

type RandNet g = g -> Network Double -> (g, Network Double)

randomizeNetwork :: RandomGen g => NNVars -> RandNet g
randomizeNetwork vars g net = runReader (randomizeNetworkReader g net) vars

randomizeNetworkReader :: RandomGen g
  => g -> Network Double -> Reader NNVars (g, Network Double)
randomizeNetworkReader g net@(Network weights) = do
  mean <- asks (getVar meanS)
  std  <- asks (getVar stddevS)
  rate <- asks (getVar mutationRate)
  let mutF = probApply rate (addNormalNoise mean std)
      (g', weights') = mapAccumL (matrixTraverse mutF) g weights
  return $ (g',fromWeightMatrices weights')

type MutateElementF g a b = (g -> a -> (g, b))

-- With a probability of the provided double, apply the random function
-- Otherwise, apply nothing
probApply :: RandomGen g
  => Double -> MutateElementF g a a -> MutateElementF g a a
probApply p rF g a = let (r,g') = random g
                     in if r > p then randomID g' a else rF g' a

randomID :: MutateElementF g a a
randomID g a = (g,a)

getActivationFunctions :: (Floating a) => Double -> (ActivationFunction a, ActivationFunctionDerivative a)
getActivationFunctions inp = let asInt :: Integer
                                 asInt = round inp
                             in case asInt of
                                  1 -> (sigmoid, sigmoid')
                                  _ -> (tanh, tanh')


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)
-}


{-# LANGUAGE DataKinds, TypeApplications, RankNTypes, FlexibleContexts,
             PartialTypeSignatures, TypeFamilies, TypeOperators, ScopedTypeVariables #-}

import qualified Control.Foldl as L
import qualified Numeric.LinearAlgebra.Static as H
import Numeric.LinearAlgebra.Data
import GHC.TypeNats
import Data.List
import Numeric.Backprop
import System.Random
import System.Random.Shuffle
import Networks
import Optimierung
import MNIST
import Tuple
import Data.Proxy

main :: IO ()
main = testDigits

antwort :: KnownNat n => H.R n -> Int
antwort = maxIndex . H.extract

normalizeInput :: (KnownNat n, KnownNat m)
               => (H.R n, H.R m)
               -> (H.R n, H.R m)
normalizeInput (x, y) = (x / 255, y)

multilayerPerceptron :: forall i o. KN i o => Modell _ (H.R i) (H.R o)
multilayerPerceptron = sigmoidNoBias @30 <~ sigmoid

threeLayer :: forall i o n. (KN i o, KnownNat n) => Proxy n -> Modell _ (H.R i) (H.R o)
threeLayer _ = sigmoid @n <~ sigmoid

mP' :: (KN i o) => Modell _ (H.R i) (H.R o)
mP' = leakyReLU @30 <~ leakyReLU

emptyLn = putStr "\n"

randomOS :: (Fractional a, Random a) => a -> a -> IO a
randomOS x y = (* y) . subtract x <$> randomIO

trainableRBM s p0 = makeFoldAll (Minibatch (10, L2 (1.0, 0.000))) (rbmcd1 s) p0

testDigits :: IO ()
testDigits = do
    let importImage l i = fmap normalizeInput <$> unsafeGetBoth l i
    let inDownloads = (++) "/home/julian/Downloads/"
    let lpath  = inDownloads "train-labels-idx1-ubyte"
    let ipath  = inDownloads "train-images-idx3-ubyte"
    let tlpath = inDownloads "t10k-labels-idx1-ubyte"
    let tipath = inDownloads "t10k-images-idx3-ubyte"

    test  <- importImage tlpath tipath
    train <- importImage lpath  ipath
    trainData <- concat <$> (mapM shuffleM $ replicate 30 train)

    p0 <- randomIO

    let trained = L.fold (uberwachtesFold (Minibatch (10, SGD 1.5)) multilayerPerceptron
                          se' p0) trainData

    emptyLn
    print $ (/ 100) . sum . map (\(x, l) ->
        if ((== antwort l) . antwort . prediction trained $ x)
        then 1
        else 0) $ test

    print $ auswerten (antwort . evalBP2 multilayerPerceptron trained)
        (==) (map (\(x, a) -> (x, antwort a)) test)

    where pred p x i l = (antwort l == x) && ((== x) . antwort . prediction p $ i)
          prediction p i = evalBP2 multilayerPerceptron p i

testAND :: IO ()
testAND = do
    -- mische 100000 Samples
    daten <- shuffleM . take 100000 . cycle $ samps
    -- zufälliger Anfangswert, der resultierende Vektor von randomIO
    -- ist standard normal verteilt
    p0 <- (* 0.1) <$> randomIO
    -- Netz wird trainiert und optimiert mit SGD und der Lernrate 0.1
    let trainiert = L.fold (uberwachtesFold (SGD 0.1) linear2 se p0) daten
    -- Hilfsfunktion für Ausgabe
    let disp = printLogicOpWith (evalBP2 linear2 trainiert)

    emptyLn
    disp 0 0
    disp 1 0
    disp 0 1
    disp 1 1

    print $ auswerten (evalBP2 linear2 trainiert) (\x y -> round x == round y) daten

-- Die vier (Eingabe, Ziel) Paare, die gelernt werden
-- diese entsprechen dem logischen UND.
samps :: [(H.R 2, Double)]
samps = [(H.vec2 0 0, 0), (H.vec2 1 0, 0), (H.vec2 0 1, 0), (H.vec2 1 1, 1)]

printLogicOpWith :: (H.R 2 -> Double) -> Double -> Double -> IO ()
printLogicOpWith f a b = print $ disp (show . f $ H.vec2 a b)
    where disp ans = "[" ++ show a ++ ", " ++ show b  ++ "]: " ++ ans

auswerten :: Ord b
          => (a -> b)
          -> (b -> b -> Bool)
          -> [(a, b)]
          -> [Int]
auswerten f pred xs = map (foldr sumCorrect 0) . groupBy sndEq . sortBy comp $ xs
    where sndEq (_, b1) (_, b2) = b1 == b2
          comp  (_, b1) (_, b2) = compare b1 b2
          sumCorrect x n        = if (eval x) then n+1 else n
          eval (a, b)           = pred (f a) b


module Main where

import Math.FFT (dftRC)
import Data.Array.IArray (listArray, amap, elems)
import Data.Array.CArray (CArray)
import Data.Complex (Complex, magnitude)

import Sound.Pulse.Simple

import GHC.Conc (TVar,writeTVar,atomically,newTVar,forkIO,readTVar)
import Control.Monad (forever)

import Graphics.Gloss.Interface.IO.Animate (animateIO)
import Graphics.Gloss.Data.Color (black,white)
import Graphics.Gloss.Data.Controller (Controller)
import Graphics.Gloss.Data.Display (Display(..))
import Graphics.Gloss.Data.Picture

source = "alsa_output.usb-AudioQuest_AudioQuest_DragonFly_Black_v1.5_AQDFBL0100116667-00.analog-stereo.monitor"

sampleWindow = 1024

-- Copy samples from source into a TVar indefinitely
capture :: Maybe String -> TVar [Float] -> IO ()
capture source sink = do
    s <- simpleNew Nothing "spectraled" Record source "capture" sampleSpec Nothing bufferAttr
    forever $ copySamples s
    where
        bufferAttr = Just $ BufferAttr Nothing Nothing Nothing Nothing (Just $ sampleWindow * 4)
        sampleSpec = SampleSpec (F32 LittleEndian) 44100 1

        copySamples s = do
            samples <- simpleRead s sampleWindow :: IO [Float]
            atomically $ writeTVar sink samples

render :: TVar [Float] -> IO ()
render samples =
    animateIO display white renderFrame callback
    where
        display = InWindow "spectraled" (sampleWindow, 200) (0, 0)

        renderFrame :: Float -> IO Picture
        renderFrame _ = do
            ss <- atomically $ readTVar samples
            --return $ line $ samplesToPath (-256) $ fmap log $ fft ss
            return $ Pictures $ samplesToBars (-256) $ fft ss

        fft :: [Float] -> [Float]
        fft ss = elems $ amap magnitude $ dftRC $ listArray (0, (length ss - 1)) ss

        samplesToPath :: Float -> [Float] -> [Point]
        samplesToPath _ [] = []
        samplesToPath i (s:ss) = (i, s * 10) : samplesToPath (i + 1) ss

        samplesToBars :: Float -> [Float] -> [Picture]
        samplesToBars _ [] = []
        samplesToBars i (s:ss) = bin : samplesToBars (i + 1) ss
            where
                bin = translate (i * 6) (s * 5) $ rectangleSolid 5 (s * 10)

        callback :: Controller -> IO ()
        callback _ = return ()

main :: IO ()
main = do
    samples <- atomically $ newTVar $ take sampleWindow $ repeat 0
    captureTID <- forkIO $ capture (Just source) samples
    render samples


{-|
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)


{-# LANGUAGE AllowAmbiguousTypes, ConstraintKinds #-}

module PCA.GaussianProcess
       ( GaussianProcess (..)
       , GPTrainingData (..)
       , PosteriorSample (..)
       , gpToPosteriorSample
       , kernelGP
       ) where

import Universum hiding (transpose, Vector)

import Control.Lens (makeLenses)

import Data.Array.Repa
import Data.Array.Repa.Repr.Unboxed (Unbox)
import Numeric.LinearAlgebra.Repa hiding (Matrix, Vector)
import System.Random (Random, mkStdGen)

import PCA.Types (InputObservations(..), Matrix, Vector)
import PCA.Util

data GaussianProcess a = GaussianProcess
    { -- | we define the gaussian process by some kernel function
      _kernelGP :: Vector D a -> Vector D a -> Matrix D a
    }

data GPTrainingData a = GPTrainingData
    { _inputTrain  :: Vector D a  -- ^ an input training data
    , _outputTrain :: Vector D a  -- ^ an output training data
    }

makeLenses ''GaussianProcess
makeLenses ''GPTrainingData

newtype PosteriorSample a = PosteriorSample
    { unSample :: Matrix D a  -- ^ an posterior sample
    }

-- | The constraint kind required for getting a posterior sample by a given GP
type GPConstraint a =
  ( Field a
  , Random a
  , Unbox a
  , Floating a
  , Eq a
  )

-- | The Main GP function: get a posterior sample by some kernel function, input observations and the training data

gpToPosteriorSample
  :: GPConstraint a
  => InputObservations a        -- ^ input observations
  -> GaussianProcess a          -- ^ a kernel function
  -> GPTrainingData a           -- ^ a training data
  -> Int                        -- ^ the number of samples
  -> Maybe (PosteriorSample a)  -- ^ a posterior functional prior
gpToPosteriorSample (InputObservations observe@(ADelayed (Z :. len) _)) gP trainingData sampleNumber = do
  -- | The kernel applied to input test points (so-called K_ss)
  let covarianceMatrix = kernel observe observe

  -- | The kernel applied to input training points (so-called K)
  let trainingKernel = kernel inputTrain' inputTrain'

  -- | The Cholesky decomposition applied to kernel of training points (:), so-called L
  let cholK = cholSH $ trainingKernel +^
                (smap (* 0.00005) . identD . size . extent $ inputTrain')

  -- | a covariance between test points and input training points (so-called K_s)
  let testPointMean = kernel observe inputTrain'

  -- | (the roots of L * x = K_s)
  cholKSolve <- delay <$> linearSolveS cholK testPointMean

  -- | Here we solve  alinear system for output training points
  cholKSolveOut <- delay <$> linearSolveS cholK (transposeMatrix $ toMatrix outputTrain' len)

  -- | Here we compute a mean
  let mean = (transposeMatrix cholKSolveOut) `mulD` cholKSolve

  -- | A posterior
  let postF' = cholSH $
                     covarianceMatrix +^
                     ((smap (* 1.0e-6) (identD len)) -^
                     (transposeMatrix cholKSolve) `mulD` cholKSolve)

  -- | A posterior sample
  return $ (PosteriorSample $ mean +^ (functionalPrior postF' sampleNumber))
    where
       kernel = gP ^. kernelGP
       inputTrain' = trainingData ^. inputTrain
       outputTrain' = trainingData ^. outputTrain
       mulD m n = delay $ m `mulS` n
       functionalPrior matrix@(ADelayed (Z :. rows :. _) _) numberSample =
         delay $ matrix `mulS` randomCoeffs
         where
           randomCoeffs = randomMatrixD (mkStdGen (-4)) (rows, numberSample)


{-# LANGUAGE DataKinds           #-}
{-# LANGUAGE FlexibleContexts    #-}
{-# LANGUAGE ImplicitParams      #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeApplications    #-}

module Main ( main ) where

import           Data.Int
import           Data.List (sort)
import           Data.List.NonEmpty (NonEmpty(..))
import           Data.Massiv.Array (getComp, index')
import qualified Data.Set as S
import qualified Data.Vector as V
import qualified Data.Vector.Storable as VS
import qualified Data.Vector.Unboxed as U
import           Data.Word
import           Geography.MapAlgebra
import qualified Numeric.LinearAlgebra as LA
import           Prelude as P
import           Test.HUnit.Approx
import qualified Test.QuickCheck.Arbitrary as QC
import           Test.Tasty
import           Test.Tasty.HUnit
import           Test.Tasty.QuickCheck

---

main :: IO ()
main = do
  img <- fromRight <$> fromGray "data/gray512.tif"
  defaultMain $ suite img

suite :: Raster S p 512 512 Word8 -> TestTree
suite r = testGroup "Unit Tests"
  [ testGroup "Raster Creation"
    [ testCase "constant (256x256)"     $ length (lazy small) @?= 65536
    , testCase "constant (2^16 x 2^16)" $ length lazybig @?= 4294967296
    , testCase "Image Reading (RGBA)"   $ do
        i <- fileRGBA
        fmap (getComp . _array . _red) i @?= Right Par
    ]
  , testGroup "Typeclass Ops"
    [ testCase "(==)" $ assertBool "(==) doesn't work" (small == small)
    , testCase "(+)"  $ strict P (lazy one + lazy one) @?= two
    ]
  , testGroup "Folds"
    [ testCase "sum (small)" $ P.sum (lazy small) @?= 327680
    -- , testCase "sum (large)" $ P.sum lazybig @?= 21474836480
    ]
  , testGroup "Local Ops"
    [ testCase "(+)"       $ P.sum (lazy small + lazy small) @?= (327680 * 2)
    , testCase "lmin"      $ strict P (lmin one two) @?= one
    , testCase "lvariety"  $ (strict P . lvariety . fmap lazy $ one :| [two]) @?= two
    , testCase "lmajority" $ (strict P . lmajority . fmap lazy $ one :| [one, two]) @?= one
    , testCase "lminority" $ (strict P . lminority . fmap lazy $ one :| [one, two]) @?= two
    -- , testCase "(+) big"   $ strict P (lazy big + lazy big) @?= bog
    ]
  , testGroup "Focal Ops"
    [ testCase "fvariety" $ strict P (fvariety one) @?= one
    , testCase "fmax"     $ strict P (fmax one) @?= one
    , testCase "fmin"     $ strict P (fmin one) @?= one
    , testGroup "flinkage"
      [ testCase "single point" singlePoint
      , testCase "2x2 same" twoByTwoSame
      , testCase "2x2 diff" twoByTwoDiff
      , testCase "3x3" threeByThree
      ]
    , testCase "flength" flengthTest
    , testCase "fpartition" fpartitionTest
    , testCase "fshape" fshapeTest
    , testCase "ffrontage" ffrontageTest
    , testGroup "farea"
      [ testCase "3x3 Open" fareaOpen
      , testCase "3x3 Centre" fareaCentre
      , testCase "4x4 Complex" fareaComplex
      ]
    , testGroup "fvolume"
      [ testCase "3x3 Flat" fvolumeFlat
      , testCase "3x3 Hill" fvolumeHill
      ]
    , testProperty "Least Squares" leastSquares
    , testGroup "fgradient"
      [ testCase "3x3 Flat" fgradientFlat
      , testCase "3x3 (tau/8)" fgradient45
      ]
    , testGroup "faspect"
      [ testCase "3x3 Flat" faspectFlat
      , testCase "3x3 East" faspectEast
      , testCase "3x3 South" faspect45
      ]
    , testGroup "fdownstream"
      [ testCase "3x3 Spikey" fdownstream4
      , testCase "3x3 Flat" fdownstreamFlat
      , testCase "3x3 Peak" fdownstreamPeak
      , testCase "3x3 Pit"  fdownstreamPit
      ]
    , testGroup "fupstream"
      [ testCase "3x3 Peak" fupstreamPeak
      , testCase "3x3 Flat" fupstreamFlat
      ]
    ]
  , testGroup "Histograms"
    [ testCase "Total Sum"     $ VS.sum (_histogram $ histogram r) @?= 262144
    , testCase "10 Breaks"     $ length (breaks $ histogram r) @?= 10
    , testCase "Sorted Breaks" $ do
        let bs = breaks $ histogram r
        sort bs @?= bs
    ]
  ]

fromRight :: Either a b -> b
fromRight (Right b) = b
fromRight _         = error "Was Left"

one :: Raster P p 7 7 Word
one = constant P Seq 1

two :: Raster P p 7 7 Word
two = constant P Seq 2

small :: Raster P p 256 256 Int
small = constant P Seq 5

lazybig :: Raster D p 65536 65536 Int
lazybig = constant D Par 5

-- big :: Raster P p 65536 65536 Word8
-- big = constant P Par 5

-- bog :: Raster P p 65536 65536 Word8
-- bog = constant P Par 10

-- indices :: Raster S p 512 512 Word8
-- indices = fromFunction S Par (\(r :. c) -> fromIntegral $ r + c)

fileRGBA :: IO (Either String (RGBARaster p 512 512 Word8))
fileRGBA = fromRGBA "data/512x512.tif"

singlePoint :: Assertion
singlePoint = actual @?= expected
  where expected :: Raster B p 1 1 Line
        expected = constant B Seq (Line 0)
        actual :: Raster B p 1 1 Line
        actual = strict B . flinkage $ constant P Seq (1 :: Int)

twoByTwoSame :: Assertion
twoByTwoSame = actual @?= expected
  where expected :: Raster S p 2 2 Line
        expected = fromRight . fromVector Seq . VS.fromList
          $ P.map (Line . _drain . drainage . S.fromList) [ [ East, South ]
                                                          , [ West, South ]
                                                          , [ North,East ]
                                                          , [ West, North ] ]
        actual :: Raster S p 2 2 Line
        actual = fromRight . fmap (strict S . flinkage) . fromVector Seq $ U.fromList ([1,1,1,1] :: [Int])

twoByTwoDiff :: Assertion
twoByTwoDiff = actual @?= expected
  where expected :: Raster S p 2 2 Line
        expected = fromRight . fromVector Seq . VS.fromList
          $ P.map (Line . _drain . drainage . S.fromList) [ [ SouthEast ]
                                                          , [ SouthWest ]
                                                          , [ NorthEast ]
                                                          , [ NorthWest ] ]
        actual :: Raster S p 2 2 Line
        actual = fromRight . fmap (strict S . flinkage) . fromVector Seq $ U.fromList ([1,2,2,1] :: [Int])

threeByThree :: Assertion
threeByThree = actual @?= expected
  where expected :: Raster S p 3 3 Line
        expected = fromRight . fromVector Seq . VS.fromList
          $ P.map (Line . _drain . drainage . S.fromList) [ [ ]
                                                          , [ South ]
                                                          , [ ]
                                                          , [ East ]
                                                          , [ North, West, South, East ]
                                                          , [ West ]
                                                          , [ ]
                                                          , [ North ]
                                                          , [ ] ]
        actual :: Raster S p 3 3 Line
        actual = fromRight . fmap (strict S . flinkage) . fromVector Seq $ U.fromList ([1,2,1,2,2,2,1,2,1] :: [Int])

flengthTest :: Assertion
flengthTest = actual @?= expected
  where actual :: Raster U p 3 3 Double
        actual = strict U . flength . flinkage . fromRight . fromVector Seq $ VS.fromList ([1,2,1,2,2,2,1,2,1] :: [Int])
        expected :: Raster U p 3 3 Double
        expected = fromRight . fromVector Seq $ U.fromList [ 0, 0.5, 0, 0.5, 2, 0.5, 0, 0.5, 0 ]

fpartitionTest :: Assertion
fpartitionTest = actual @?= expected
  where expected :: Raster B p 2 2 Corners
        expected = fromRight . fromVector Seq $ V.fromList [ Corners Open Open Open Open
                                                           , Corners Open Open Open Open
                                                           , Corners OneSide Open OneSide Complete
                                                           , Corners Open Open Open Open ]
        actual :: Raster B p 2 2 Corners
        actual = strict B . fpartition . fromRight . fromVector Seq $ U.fromList ([1,1,2,1] :: [Int])

fshapeTest :: Assertion
fshapeTest = actual @?= expected
 where expected :: Raster B p 3 3 Corners
       expected = fromRight . fromVector Seq $ V.fromList [ Corners Open Open OutFlow Open
                                                          , Corners Open Open Open Open
                                                          , Corners Open OutFlow Open Open
                                                          , Corners Open Open Open Open
                                                          , Corners Complete Complete Complete Complete
                                                          , Corners Open Open Open Open
                                                          , Corners Open Open Open OutFlow
                                                          , Corners Open Open Open Open
                                                          , Corners OutFlow Open Open Open ]
       actual :: Raster B p 3 3 Corners
       actual = strict B . fshape . fromRight . fromVector Seq $ U.fromList ([1,1,1,1,0,1,1,1,1] :: [Int])

ffrontageTest :: Assertion
ffrontageTest = let ?epsilon = 0.001 in actual @?~ expected
  where expected :: Double
        expected = 1 + (1 / sqrt 2)
        actual :: Double
        actual = flip index' (1 :. 1) . _array . strict S $ ffrontage rast
        rast :: Raster DW p 4 4 Corners
        rast = fshape . fromRight . fromVector Seq $ U.fromList ( [1,1,1,0
                                                                  ,1,0,0,0
                                                                  ,1,0,0,1
                                                                  ,1,0,1,1] :: [Int] )

fareaOpen :: Assertion
fareaOpen = actual @?= expected
  where expected :: Raster U p 3 3 Double
        expected = fromRight . fromVector Seq $ U.fromList [1,1,1,1,1,1,1,1,1]
        actual :: Raster U p 3 3 Double
        actual = strict U . farea . fshape . fromRight . fromVector Seq $ U.fromList ([0,0,0,0,0,0,0,0,0] :: [Int])

fareaCentre :: Assertion
fareaCentre = actual @?= expected
  where expected :: Raster U p 3 3 Double
        expected = fromRight . fromVector Seq $ U.fromList [ 1 + 1/8, 1, 1 + 1/8
                                                           , 1, 1/2, 1
                                                           , 1 + 1/8, 1, 1 + 1/8 ]
        actual :: Raster U p 3 3 Double
        actual = strict U . farea . fshape . fromRight . fromVector Seq $ U.fromList ([0,0,0,0,1,0,0,0,0] :: [Int])

fareaComplex :: Assertion
fareaComplex = let ?epsilon = 0.001 in actual @?~ (7 / 8)
  where actual :: Double
        actual = flip index' (1 :. 1) . _array . strict P $ farea rast
        rast :: Raster DW p 4 4 Corners
        rast = fshape . fromRight . fromVector Seq $ U.fromList ( [1,1,1,0
                                                                  ,1,0,0,0
                                                                  ,1,0,0,1
                                                                  ,1,0,1,1] :: [Int] )

fvolumeFlat :: Assertion
fvolumeFlat = strict U (fvolume expected) @?= expected
  where expected :: Raster U p 3 3 Double
        expected = fromRight . fromVector Seq $ U.fromList [8,8,8,8,8,8,8,8,8]

fvolumeHill :: Assertion
fvolumeHill = index' (_array actual) (1 :. 1) @?= expected
  where expected :: Double
        expected = P.sum [20,20,16,20,16,16,16,16,12,16,12,12] / 12
        actual :: Raster U p 3 3 Double
        actual = strict U . fvolume @Double . fromRight . fromVector Seq $ U.fromList [24,24,24
                                                                                ,16,16,16
                                                                                ,8,8,8]

newtype Vec = Vec [Double] deriving (Show)

instance Arbitrary Vec where
  arbitrary = Vec <$> QC.vector 9

-- | A QuickCheck property to test whether my custom Least Squares is as
-- accurate as the one provided by HMatrix.
leastSquares :: Vec -> Bool
leastSquares (Vec vs) = f 0 && f 1 && f 2
  where m = head . LA.toColumns $ LA.linearSolveLS zing (LA.col vs)
        v = leftPseudo LA.#> LA.vector vs
        f i = (m LA.! i) =~ (v LA.! i)

-- | Approximate Equality.
(=~) :: Double -> Double -> Bool
a =~ b = abs (a - b) < 0.0001

zing :: LA.Matrix Double
zing = LA.matrix 3 [ -0.5, -0.5, 1
                   , -0.5, 0, 1
                   , -0.5, 0.5, 1
                   , 0, -0.5, 1
                   , 0, 0, 1
                   , 0, 0.5, 1
                   , 0.5, -0.5, 1
                   , 0.5, 0, 1
                   , 0.5, 0.5, 1 ]

fgradientFlat :: Assertion
fgradientFlat = actual @?= expected
  where expected :: Raster U p 3 3 Double
        expected = fromRight . fromVector Seq $ U.fromList [0,0,0,0,0,0,0,0,0]
        actual :: Raster U p 3 3 Double
        actual = strict U . fgradient . fromRight . fromVector Seq $ U.fromList ([1,1,1,1,1,1,1,1,1] :: [Double])

fgradient45 :: Assertion
fgradient45 = let ?epsilon = 0.0001 in index' (_array actual) (1 :. 1) @?~ (tau / 8)
  where actual :: Raster U p 3 3 Double
        actual = strict U . fgradient . fromRight . fromVector Seq $ U.fromList ([3,3,3,2,2,2,1,1,1] :: [Double])

faspectFlat :: Assertion
faspectFlat = index' (_array actual) (1 :. 1) @?= Nothing
  where actual :: Raster B p 3 3 (Maybe Double)
        actual = strict B . faspect . fromRight . fromVector Seq $ U.fromList ([1,1,1,1,1,1,1,1,1] :: [Double])

faspect45 :: Assertion
faspect45 = index' (_array actual) (1 :. 1) @?= Just (tau / 2)
  where actual :: Raster B p 3 3 (Maybe Double)
        actual = strict B . faspect . fromRight . fromVector Seq $ U.fromList ([3,3,3,2,2,2,1,1,1] :: [Double])

faspectEast :: Assertion
faspectEast = let ?epsilon = 0.0001 in index' (_array actual) (1 :. 1) @?~ (tau / 4)
  where actual :: Raster B p 3 3 Double
        actual = strict B . faspect' . fromRight . fromVector Seq $ U.fromList ([3,2,1,3,2,1,3,2,1] :: [Double])

fdownstream4 :: Assertion
fdownstream4 = index' (_array actual) (1 :. 1) @?= drainage (S.fromList [North,South,East,West])
  where actual :: Raster S p 3 3 Drain
        actual = strict S . fdownstream . fromRight . fromVector Seq $ U.fromList ([3,1,3,1,2,1,3,1,3] :: [Double])

fdownstreamFlat :: Assertion
fdownstreamFlat = index' (_array actual) (1 :. 1) @?= drainage (S.fromList [East ..])
  where actual :: Raster S p 3 3 Drain
        actual = strict S . fdownstream . fromRight . fromVector Seq $ U.fromList ([1,1,1,1,1,1,1,1,1] :: [Double])

fdownstreamPeak :: Assertion
fdownstreamPeak = index' (_array actual) (1 :. 1) @?= drainage (S.fromList [NorthEast, NorthWest, SouthWest, SouthEast])
  where actual :: Raster S p 3 3 Drain
        actual = strict S . fdownstream . fromRight . fromVector Seq $ U.fromList ([1,1,1,1,3,1,1,1,1] :: [Double])

fdownstreamPit :: Assertion
fdownstreamPit = index' (_array actual) (1 :. 1) @?= Drain 0
  where actual :: Raster S p 3 3 Drain
        actual = strict S . fdownstream . fromRight . fromVector Seq $ U.fromList ([2,2,2,2,1,2,2,2,2] :: [Double])

fupstreamFlat :: Assertion
fupstreamFlat = index' (_array actual) (1 :. 1) @?= drainage (S.fromList [East ..])
  where actual :: Raster S p 3 3 Drain
        actual = strict S . fupstream . strict S . fdownstream . fromRight . fromVector Seq $ U.fromList ([1,1,1,1,1,1,1,1,1] :: [Double])

fupstreamPeak :: Assertion
fupstreamPeak = index' (_array actual) (1 :. 1) @?= Drain 0
  where actual :: Raster S p 3 3 Drain
        actual = strict S . fupstream . strict S . fdownstream . fromRight . fromVector Seq $ U.fromList ([1,1,1,1,3,1,1,1,1] :: [Double])

{-
hists :: Raster S p 512 512 Word8 -> Assertion
hists r = do
  u  <- histU r
  um <- histMut r
  histU' r @?= u
  histMut' r @?= um
  histMut'' r @?= um
  histMut'' r @?= u
  um @?= u
-}


-- |
-- Module      :  Test.FFTW
-- Copyright   :  (c) 2016 Drexel University
-- License     :  BSD-style
-- Maintainer  :  mainland@drexel.edu

module Test.FFTW (
    fft
  ) where

import Data.Complex
import qualified Data.Vector.Storable as V
import Numeric.FFT.Vector.Unnormalized as FFTW

fft :: Int -> V.Vector (Complex Double) -> V.Vector (Complex Double)
fft n = FFTW.execute $ FFTW.plan FFTW.dft n


module Main (main) where

import Neural.Matrix
import Data.Ord (comparing)
import Data.List (scanl, foldl', maximumBy)
import Numeric.LinearAlgebra (toList, vector)
import Control.Monad (forM_)
import Codec.Compression.GZip (decompress)
import qualified Data.ByteString.Lazy as BS

getImage s n = fromIntegral . BS.index s . (n*28^2 + 16 +) <$> [0..28^2 - 1]
getX     s n = (vector (getImage s n)) / 256
getLabel s n = fromIntegral $ BS.index s (n + 8)
getY     s n = vector $ fromIntegral . fromEnum . (getLabel s n ==) <$> [0..9]

printImage imgs = putStrLn . unlines . take 28 . map (take 28) . iterate (drop 28)
                           . map pixel2Char . getImage imgs
  where pixel2Char n = let s = " ·:o◍O0@" in s !! (fromIntegral n * length s `div` 256)

display d prob = show d ++ ": " ++ show prob

bestOf :: [Float] -> Int
bestOf = fst . maximumBy (comparing snd) . zip [0..]

small = [[0..99], [100..299], [300..599], [600..999]]
medium = [[0..299], [300..999], [1000..2199], [2200..4599]]
large = [[0..999], [1000..2999], [3000..5999], [6000..9999]]
rapid = [[0..299], [300..599], [600..899], [900..1199], [1200..1499], [1500..1799], [1800..2099]]

main :: IO ()
main = do
  [trainI, trainL, testI, testL] <- mapM ((decompress  <$>) . BS.readFile . ("examples/mnistData/"++))
                                    ["train-images-idx3-ubyte.gz"
                                    ,"train-labels-idx1-ubyte.gz"
                                    ,"t10k-images-idx3-ubyte.gz"
                                    ,"t10k-labels-idx1-ubyte.gz"]
  net <- initNet 1 [784, 30, 10] [relu, relu]
  let n = 42
  printImage testI n

  let
    example = getX testI n
    randomUpdate net x = updateNet categoricalCrossEntropy 0.002 (getX trainI x) (getY trainL x) net
    trainingNet = scanl (foldl' randomUpdate) net large
    trainedNet = last trainingNet
  forM_ trainingNet $ putStrLn . unlines . zipWith display [0..9] . softmax
                               . toList . snd . last . forwardPass example


module MNIST.DataSet where

import MNIST.Prelude

import qualified Codec.Compression.GZip  as GZip
import qualified Data.ByteString.Lazy    as L
import qualified Data.Vector.Generic     as G
import qualified Numeric.LinearAlgebra   as HM
import qualified Numeric.LinearAlgebra.Static  as SM

rows :: Int
rows = 28

cols :: Int
cols = 28

nPixels :: Int
nPixels = cols * rows

nLabels :: Int
nLabels = 10

_loadMNIST
    :: FilePath
    -> FilePath
    -> IO [(Int, UVector Int)]
_loadMNIST dataPath labelPath =
  runMaybeT (do
    i <- MaybeT . fmap (decodeIDX       . GZip.decompress) . L.readFile $ dataPath
    l <- MaybeT . fmap (decodeIDXLabels . GZip.decompress) . L.readFile $ labelPath
    d <- MaybeT . pure $ labeledIntData l i
    return d)
  >>= \case
    Just mnist -> return mnist
    Nothing    -> throwString $
      "couldn't read gzipped MNIST at data file " <> dataPath
      <> " and label file " <> labelPath

loadMNISTTf :: FilePath -> FilePath -> IO [(Int, Vector Int)]
loadMNISTTf dp lp = _loadMNIST dp lp
  >>= pure . fmap (identity *** G.convert)

trainingDataTf :: IO [(Int, Vector Int)]
trainingDataTf =
  loadMNISTTf
    "data/train-images-idx3-ubyte.gz"
    "data/train-labels-idx1-ubyte.gz"

testDataTf :: IO [(Int, Vector Int)]
testDataTf =
  loadMNISTTf
    "data/t10k-images-idx3-ubyte.gz"
    "data/t10k-labels-idx1-ubyte.gz"

-- ========================================================================= --

loadMNISTBp
    :: FilePath
    -> FilePath
    -> IO [(R 784, R 9)]
loadMNISTBp dp lp = _loadMNIST dp lp
  >>= pure . fmap ((fromJust . mkImage *** fromJust . mkLabel) . swap)

  where
    mkImage :: UVector Int -> Maybe (R 784)
    mkImage u = SM.create . G.convert . G.map (\i -> fromIntegral i / 255) $ u

    mkLabel :: Int -> Maybe (R 9)
    mkLabel n = SM.create $ HM.build 9 (fromIntegral . fromEnum . (== n) . round)

trainingDataBp :: IO [(R 784, R 9)]
trainingDataBp =
  loadMNISTBp
    "data/train-images-idx3-ubyte.gz"
    "data/train-labels-idx1-ubyte.gz"

testDataBp :: IO [(R 784, R 9)]
testDataBp =
  loadMNISTBp
    "data/t10k-images-idx3-ubyte.gz"
    "data/t10k-labels-idx1-ubyte.gz"


-- 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 NamedFieldPuns, TemplateHaskell #-}

module School.Unit.Test.UnitBackward
( unitBackwardTest) where

import Conduit ((.|), sinkList, yield)
import Control.Monad.IO.Class (liftIO)
import Data.Either (isLeft)
import Numeric.LinearAlgebra (ident)
import School.TestUtils (assertRight, fromRight, randomAffineParams, randomMatrix, testState)
import School.Train.TrainState (TrainState(..), def)
import School.Types.FloatEq ((~=))
import School.Types.PingPong (pingPongSingleton, reversePingPong, toPingPong)
import School.Unit.Affine (affine)
import School.Unit.RecLin (recLin)
import School.Unit.Unit (Unit(..))
import School.Unit.UnitActivation (UnitActivation(..))
import School.Unit.UnitBackward
import School.Unit.UnitGradient (UnitGradient(..), isGradientFail)
import School.Unit.UnitParams (UnitParams(..))
import School.Utils.Tuple (snd3)
import Test.Tasty (TestTree)
import Test.Tasty.QuickCheck
import Test.Tasty.TH
import Test.QuickCheck.Monadic (assert, monadicIO)

prop_affine_input_fail :: Property
prop_affine_input_fail = monadicIO $ do
  let backward = unitBackward affine
  let acts = [ApplyFail "init"]
  let inGrad = BatchGradient $ ident 1
  let network =  yield (acts, inGrad, 0)
              .| backward
              .| sinkList
  result <- testState network def
  assert $ isLeft result

prop_affine_gradient_fail :: Property
prop_affine_gradient_fail = monadicIO $ do
  let backward = unitBackward affine
  let acts = [BatchActivation $ ident 1]
  let inGrad = GradientFail "init"
  let network =  yield (acts, inGrad, 0)
              .| backward
              .| sinkList
  result <- testState network def
  assert $ isLeft result

prop_affine_param_fail :: Positive Int -> Positive Int -> Property
prop_affine_param_fail (Positive fSize) (Positive oSize) = monadicIO $ do
  let backward = unitBackward affine
  actMat <- liftIO $ randomMatrix oSize fSize
  gradMat <- liftIO $ randomMatrix oSize fSize
  let acts = [BatchActivation actMat]
  let fStack = (acts, BatchGradient gradMat, 0)
  let network =  yield fStack
              .| backward
              .| sinkList
  result <- testState network def
  assertRight (isGradientFail . snd3 . head . fst)
              result

prop_reclin_gradient :: Positive Int -> Positive Int -> Property
prop_reclin_gradient (Positive bSize) (Positive fSize) = monadicIO $ do
  let backward = unitBackward recLin
  actMat <- liftIO $ randomMatrix bSize fSize
  let input = BatchActivation actMat
  gradMat <- liftIO $ randomMatrix bSize fSize
  let inGrad = BatchGradient gradMat
  let fStack = ([input], inGrad, 0)
  let network =  yield fStack
              .| backward
              .| sinkList
  result <- testState network def
  let (check, _) = deriv recLin EmptyParams inGrad input
  assertRight (((~=) check) . snd3 . head . fst)
              result

prop_affine_derivs :: Positive Int -> Positive Int -> Positive Int -> Property
prop_affine_derivs (Positive bSize) (Positive fSize) (Positive oSize) = monadicIO $ do
  let backward = unitBackward affine
  actMat <- liftIO $ randomMatrix bSize fSize
  let input = BatchActivation actMat
  gradMat <- liftIO $ randomMatrix bSize oSize
  let inGrad = BatchGradient gradMat
  let fStack = ([input], inGrad, 0)
  let network =  yield fStack
              .| backward
              .| sinkList
  params <- liftIO $ randomAffineParams fSize oSize
  let paramList = pingPongSingleton params
  let initState = def { paramList }
  result <- testState network initState
  let (_, check) = deriv affine params inGrad input
  assertRight (((~=) check) . head . paramDerivs . snd)
              result

prop_deriv_aff_rl_aff_rl :: Positive Int
                         -> Positive Int
                         -> Positive Int
                         -> Positive Int
                         -> Property
prop_deriv_aff_rl_aff_rl (Positive b)
                         (Positive f)
                         (Positive h)
                         (Positive o) = monadicIO $ do
  inMat <- liftIO $ randomMatrix b f
  act1Mat <- liftIO $ randomMatrix b h
  act2Mat <- liftIO $ randomMatrix b h
  act3Mat <- liftIO $ randomMatrix b o
  let acts = BatchActivation <$> [ act3Mat
                                 , act2Mat
                                 , act1Mat
                                 , inMat
                                 ]
  gradMat <- liftIO $ randomMatrix b o
  let inGrad = BatchGradient gradMat
  let network =  yield (acts, inGrad, 0)
              .| unitBackward recLin
              .| unitBackward affine
              .| unitBackward recLin
              .| unitBackward affine
              .| sinkList
  params1 <- liftIO $ randomAffineParams f h
  params2 <- liftIO $ randomAffineParams h o
  let allParams = toPingPong [ params1
                             , EmptyParams
                             , params2
                             , EmptyParams
                             ]
  let paramList = reversePingPong $ fromRight (pingPongSingleton EmptyParams) allParams
  let initState = def { paramList }
  result <- testState network initState
  let (grad1, _) = deriv recLin EmptyParams inGrad (head acts)
  let (grad2, dParams2) = deriv affine params2 grad1 (acts!!1)
  let (grad3, _) = deriv recLin EmptyParams grad2 (acts!!2)
  let (_, dParams1) = deriv affine params1 grad3 (last acts)
  let check = [ dParams1
              , EmptyParams
              , dParams2
              , EmptyParams
              ]
  assertRight (((~=) check) . paramDerivs . snd)
              result

unitBackwardTest :: TestTree
unitBackwardTest = $(testGroupGenerator)