{-# LANGUAGE BinaryLiterals #-}
{-# LANGUAGE OverloadedLists #-}

module Main (main) where

-- import Data.Yaml.Aeson
import qualified Data.Aeson as Aeson
import Data.Bits
import Data.Complex
import Data.Ratio ((%))
import Data.Type.Equality
import qualified Data.Vector as V
import qualified Data.Vector.Storable as S
import qualified Data.Vector.Unboxed as U
import Foreign.Storable
import GHC.Exts (IsList (..))
-- import LatticeSymmetries
import LatticeSymmetries.Algebra
import LatticeSymmetries.Basis
import LatticeSymmetries.Benes
-- import qualified LatticeSymmetries.CSR as CSR
import LatticeSymmetries.ComplexRational
import LatticeSymmetries.Dense
import LatticeSymmetries.Generator
import LatticeSymmetries.Group
-- import LatticeSymmetries.IO
import LatticeSymmetries.NonbranchingTerm
import LatticeSymmetries.Operator
import LatticeSymmetries.Parser
-- import LatticeSymmetries.Sparse
-- import LatticeSymmetries.Types
import Test.Hspec
import Text.Parsec (parse)
import Text.PrettyPrint.ANSI.Leijen (hPutDoc, hardline, pretty, putDoc)
import Prelude hiding (Product, Sum, toList)

-- anySpinEDException :: Selector SpinEDException
-- anySpinEDException = const True

-- anyLatticeSymmetriesException :: Selector LatticeSymmetriesException
-- anyLatticeSymmetriesException = const True

-- parseLines :: FromJSON a => [Text] -> Either ParseException a
-- parseLines xs = decodeEither' . encodeUtf8 . unlines $ xs

-- checkRight s y = case s of
--   Left _ -> s `shouldSatisfy` isRight
--   Right x -> x `shouldBe` y

-- checkJust s y = case s of
--   Just x -> x `shouldBe` y
--   Nothing -> s `shouldSatisfy` isJust

type SpinPolynomial =
  Sum (Scaled ComplexRational (Product (Generator Int SpinGeneratorType)))

type FermionicPolynomial =
  Sum (Scaled ComplexRational (Product (Generator Int FermionGeneratorType)))

main :: IO ()
main = hspec $ do
  describe "ToJSON/FromJSON" $ do
    it "parses Permutation" $ do
      Aeson.decode "[3, 0, 2, 1]" `shouldBe` Just (mkPermutation [3, 0, 2, 1])
      Aeson.decode "[3, 0, 2, 1, 4, 5, 6]" `shouldBe` Just (mkPermutation [3, 0, 2, 1, 4, 5, 6])
      Aeson.decode "[]" `shouldBe` Just (mkPermutation [])
      Aeson.decode "{ \"oops\": [2, 1, 0] }" `shouldBe` (Nothing :: Maybe Permutation)
    it "encodes Permutation" $ do
      Aeson.encode (mkPermutation [0, 1, 2, 3]) `shouldBe` "[0,1,2,3]"
      Aeson.encode (mkPermutation []) `shouldBe` "[]"
      Aeson.encode (mkPermutation [3, 5, 0, 1, 2, 4]) `shouldBe` "[3,5,0,1,2,4]"
    it "parses Symmetry" $ do
      Aeson.decode "{ \"permutation\": [1, 2, 3, 0], \"sector\": 2 }" `shouldBe` Just (mkSymmetry [1, 2, 3, 0] 2)
      Aeson.decode "{ \"permutation\": [0, 1], \"sector\": 0 }" `shouldBe` Just (mkSymmetry [0, 1] 0)
      Aeson.decode "{ \"permutation\": [], \"sector\": 0 }" `shouldBe` Just (mkSymmetry [] 0)
    -- it "encodes Symmetry" $ do
    --   Aeson.encode (mkSymmetry [0, 1, 2, 3, 4] 0) `shouldBe` "{\"sector\":0,\"permutation\":[0,1,2,3,4]}"
    --   Aeson.encode (mkSymmetry [] 0) `shouldBe` "{\"sector\":0,\"permutation\":[]}"
    --   Aeson.encode (mkSymmetry [1, 2, 0] 2) `shouldBe` "{\"sector\":2,\"permutation\":[1,2,0]}"
    it "parses Symmetries" $ do
      Aeson.decode "[{\"permutation\": [1, 2, 0], \"sector\": 1}]" `shouldBe` Just (mkSymmetries [mkSymmetry [1, 2, 0] 1])
      Aeson.decode
        "[{\"permutation\": [1, 2, 3, 0], \"sector\": 0}, \
        \ {\"permutation\": [3, 2, 1, 0], \"sector\": 0}]"
        `shouldBe` Just (mkSymmetries [mkSymmetry [1, 2, 3, 0] 0, mkSymmetry [3, 2, 1, 0] 0])
    -- it "encodes Symmetries" $ do
    --   Aeson.encode (mkSymmetries [mkSymmetry [1, 2, 0] 1])
    --     `shouldBe` "[{\"sector\":0,\"permutation\":[0,1,2]},{\"sector\":2,\"permutation\":[1,2,0]},{\"sector\":1,\"permutation\":[2,0,1]}]"
    it "parses/encodes ParticleTy" $ do
      forM_ ([SpinTy, SpinlessFermionTy, SpinfulFermionTy] :: [ParticleTy]) $ \tp ->
        Aeson.decode (Aeson.encode tp) `shouldBe` Just tp
      Aeson.decode "\"spin\"" `shouldBe` Just SpinTy
      Aeson.decode "\"spinless-fermion\"" `shouldBe` Just SpinlessFermionTy
      Aeson.decode "\"spinful-fermion\"" `shouldBe` Just SpinfulFermionTy
    it "parses/encodes Basis 'SpinTy" $ do
      let s₁ = mkSpinBasis 10 (Just 5) Nothing emptySymmetries
          s₂ = mkSpinBasis 4 (Just 2) (Just 1) (mkSymmetries [mkSymmetry [1, 2, 3, 0] 0])
      print $ Aeson.encode s₂
      forM_ ([s₁, s₂] :: [Basis 'SpinTy]) $ \s ->
        Aeson.decode (Aeson.encode s) `shouldBe` Just s
    it "parses/encodes Basis 'SpinfulFermionTy" $ do
      let s₁ = mkSpinfulFermionicBasis 5 (SpinfulTotalParticles 4)
          s₂ = mkSpinfulFermionicBasis 3 (SpinfulPerSector 2 1)
          s₃ = mkSpinfulFermionicBasis 30 SpinfulNoOccupation
      forM_ ([s₁, s₂, s₃] :: [Basis 'SpinfulFermionTy]) $ \s -> do
        print $ Aeson.encode s
        Aeson.decode (Aeson.encode s) `shouldBe` Just s
    it "parses/encodes Basis 'SpinlessFermionTy" $ do
      let s₁ = mkSpinlessFermionicBasis 5 Nothing
          s₂ = mkSpinlessFermionicBasis 10 (Just 5)
      forM_ ([s₁, s₂] :: [Basis 'SpinlessFermionTy]) $ \s -> do
        print $ Aeson.encode s
        Aeson.decode (Aeson.encode s) `shouldBe` Just s

  --   (mkdecode "[{\"permutation\": [1, 2, 0], \"sector\": 1}]" `shouldBe` Just
  --   `shouldBe` Just (mkSymmetries [mkSymmetry [1, 2, 3, 0] 0, mkSymmetry [3, 2, 1, 0] 0])
  -- describe "SymmetrySpec" $ do
  --   it "parses JSON specifications" $ do
  --     let s₁ =
  --           parseLines @SymmetrySpec
  --             [ "permutation: [3, 1, 2, 4]",
  --               "sector: 8"
  --             ]
  --         s₂ =
  --           parseLines @SymmetrySpec
  --             [ "permutation: [1, 1, 2, 4]",
  --               "sector: -5"
  --             ]
  --         s₃ =
  --           parseLines @SymmetrySpec
  --             [ "permutation: []",
  --               "sector: 5"
  --             ]
  --         s₄ =
  --           parseLines @SymmetrySpec
  --             [ "permutation: [1]",
  --               "sector: null"
  --             ]
  --     case s₁ of
  --       Left _ -> s₁ `shouldSatisfy` isRight
  --       Right x -> x `shouldBe` SymmetrySpec (fromList [3, 1, 2, 4]) 8
  --     case s₂ of
  --       Left _ -> s₂ `shouldSatisfy` isRight
  --       Right x -> x `shouldBe` SymmetrySpec (fromList [1, 1, 2, 4]) (-5)
  --     s₃ `shouldSatisfy` isLeft
  --     s₄ `shouldSatisfy` isLeft
  -- describe "denseMatrixFromList" $ do
  --   it "handles empty lists" $ do
  --     denseMatrixFromList []
  --       `shouldBe` Right (DenseMatrix S.empty :: DenseMatrix S.Vector Double)
  --     denseMatrixFromList [[], []]
  --       `shouldBe` Right (DenseMatrix S.empty :: DenseMatrix S.Vector Double)
  --   it "handles square matrices" $ do
  --     denseMatrixFromList [[1, 2], [3, 4]]
  --       `shouldBe` Right (DenseMatrix (fromList [1, 2, 3, 4]) :: DenseMatrix S.Vector Double)
  --   it "handles rectangular matrices" $ do
  --     denseMatrixFromList [[1, 2, 3], [4, 5, 6]]
  --       `shouldBe` Right (DenseMatrix (fromList [1, 2, 3, 4, 5, 6]) :: DenseMatrix S.Vector Double)
  --     denseMatrixFromList [[1, 2], [3, 4], [5, 6]]
  --       `shouldBe` Right (DenseMatrix (fromList [1, 2, 3, 4, 5, 6]) :: DenseMatrix S.Vector 3 2 Double)
  -- describe "DenseMatrixSpec" $ do
  --   it "parses JSON specifications" $ do
  --     parseLines @DenseMatrixSpec
  --       [ "[[1, 2,  1],",
  --         " [4, 2,  1],",
  --         " [1, -2, 3]]"
  --       ]
  --       `checkRight` DenseMatrixSpec [[1, 2, 1], [4, 2, 1], [1, -2, 3]]
  --     parseLines @DenseMatrixSpec
  --       [ "[[1, 2],",
  --         " [4, 2],",
  --         " [1, [-2, 3]]]"
  --       ]
  --       `checkRight` DenseMatrixSpec [[1, 2], [4, 2], [1, (-2) :+ 3]]
  --     parseLines @DenseMatrixSpec ["[]"] `checkRight` DenseMatrixSpec []
  --     parseLines @DenseMatrixSpec ["[[], []]"] `checkRight` DenseMatrixSpec [[], []]
  -- describe "denseToCSRSquare" $ do
  --   it "works" $ do
  --     denseToSparse (fromList [[1, 0, 2], [0, 0, -3]]) `shouldBe` Nothing
  --     print $ denseToSparse (fromList [[1, 0, 2], [0, 4, 0], [0, 0, -3]])
  --     True `shouldBe` True
  --   it ".." $ do
  --     let m =
  --           fromList
  --             [ [1, 0, 0.2, 0],
  --               [0, 1, (-2) :+ 8, 0],
  --               [0, -2.1, 1, 0],
  --               [0 :+ 0.1, 0, 0, 3]
  --             ]
  --         expr = denseToSparse m
  --      in case expr of
  --           Just m' -> sparseToDense m' `shouldBe` m
  --           Nothing -> expr `shouldSatisfy` isJust
  --     let m =
  --           fromList
  --             [[3]]
  --         expr = denseToSparse m
  --      in case expr of
  --           Just m' -> sparseToDense m' `shouldBe` m
  --           Nothing -> expr `shouldSatisfy` isJust
  --     let m =
  --           fromList
  --             []
  --         expr = denseToSparse m
  --      in case expr of
  --           Just m' -> sparseToDense m' `shouldBe` m
  --           Nothing -> expr `shouldSatisfy` isJust
  -- describe "Storable instances" $ do
  --   it "has correct sizeOf" $ do
  --     sizeOf (undefined :: Csparse_matrix) `shouldBe` trueCsparse_matrixSizeOf
  --     sizeOf (undefined :: Cbit_index) `shouldBe` trueCbit_indexSizeOf
  --     sizeOf (undefined :: Cterm) `shouldBe` trueCtermSizeOf
  --     sizeOf (undefined :: Coutput_buffer) `shouldBe` trueCoutput_bufferSizeOf
  --     sizeOf (undefined :: Csparse_operator) `shouldBe` trueCsparse_operatorSizeOf
  --   it "has correct alignment" $ do
  --     alignment (undefined :: Csparse_matrix) `shouldBe` trueCsparse_matrixAlignment
  --     alignment (undefined :: Cbit_index) `shouldBe` trueCbit_indexAlignment
  --     alignment (undefined :: Cterm) `shouldBe` trueCtermAlignment
  --     alignment (undefined :: Coutput_buffer) `shouldBe` trueCoutput_bufferAlignment
  --     alignment (undefined :: Csparse_operator) `shouldBe` trueCsparse_operatorAlignment
  -- describe "applyOperatorTerm'" $ do
  --   it ".." $ do
  --     let spec =
  --           InteractionSpec
  --             [ [1, 0, 0, 0],
  --               [0, -1, 2, 0],
  --               [0, 2, -1, 0],
  --               [0, 0, 0, 1]
  --             ]
  --             [[0, 1]]
  --         (Right term) = toOperatorTerm spec
  --     print =<< term `applyOperatorTerm'` [0x1]
  --     True `shouldBe` True
  -- describe "binarySearch" $ do
  --   it ".." $ do
  --     binarySearch (S.fromList [1, 3, 4, 5, 7, 9]) (4 :: Int) `shouldBe` (Just 2)
  --     binarySearch (S.fromList [1, 3, 4, 5, 7, 9]) (2 :: Int) `shouldBe` Nothing
  --     binarySearch (S.fromList []) (2 :: Int) `shouldBe` Nothing
  --     binarySearch (S.fromList [1]) (1 :: Int) `shouldBe` Just 0
  -- describe "cooToCSR" $ do
  --   it ".." $ do
  --     let (m :: CSR S.Vector Int 3 4 Int) = cooToCsr @U.Vector @S.Vector . fromList $ [(2, 2, -3), (1, 1, 1), (2, 3, 8), (1, 1 :: Int, 3 :: Int)]
  --     print m
  --     m `csrIndex` (1, 1) `shouldBe` 4
  -- describe "denseToCSR" $ do
  --   it ".." $ do
  --     let (m :: CSR S.Vector Int 3 4 Int) = denseToCsr (fromList [[1, 0, 2, 0], [0, 4, 0, 0], [0, 0, -3, 8]])
  --     print m
  --     m `csrIndex` (0, 0) `shouldBe` 1
  --     m `csrIndex` (0, 1) `shouldBe` 0
  --     m `csrIndex` (0, 2) `shouldBe` 2
  --     m `csrIndex` (2, 3) `shouldBe` 8
  -- describe "Num CSR" $ do
  --   it "(+)" $ do
  --     let (a :: CSR S.Vector Int 3 4 Int) =
  --           denseToCsr (fromList [[1, 0, 2, 0], [0, 4, 0, 0], [0, 0, -3, 8]])
  --         (b :: CSR S.Vector Int 3 4 Int) =
  --           denseToCsr (fromList [[3, -8, 0, 0], [1, 5, 0, 0], [0, 0, -3, 0]])
  --         (c :: CSR S.Vector Int 3 4 Int) =
  --           denseToCsr (fromList [[4, -8, 2, 0], [1, 9, 0, 0], [0, 0, -6, 8]])
  --     print a
  --     print b
  --     print c
  --     a + b `shouldBe` c
  --   it "(+)" $ do
  --     let a₁' =
  --           fromList @(SomeCSR S.Vector Int Double) $
  --             [ (0, 0, 0.05052698254161314),
  --               (0, 1, 0.34388186562432865),
  --               (0, 2, 0.3736125370797053),
  --               (0, 4, 0.37422323470438656),
  --               (1, 1, 0.08956717735368347),
  --               (1, 3, 0.20028803087704727),
  --               (2, 0, 0.010520730797154432),
  --               (2, 1, 0.7913535778080222),
  --               (2, 3, 0.6557777513743904),
  --               (2, 4, 0.9609587020099947)
  --             ]
  --         b₁' =
  --           fromList @(SomeCSR S.Vector Int Double) $
  --             [ (2, 3, 0.3942046813633675),
  --               (2, 4, 0.0008746925379253812)
  --             ]
  --         (c₁ :: [(Int, Int, Double)]) =
  --           [ (0, 0, 0.05052698254161314),
  --             (0, 1, 0.34388186562432865),
  --             (0, 2, 0.3736125370797053),
  --             (0, 4, 0.37422323470438656),
  --             (1, 1, 0.08956717735368347),
  --             (1, 3, 0.20028803087704727),
  --             (2, 0, 0.010520730797154432),
  --             (2, 1, 0.7913535778080222),
  --             (2, 3, 1.049982432737758),
  --             (2, 4, 0.96183339454792)
  --           ]
  --         a₂ =
  --           fromList @(CSR S.Vector Int 5 2 Double) $
  --             [(1, 1, 0.9913900269818051), (2, 0, 0.701927418412059)]
  --         b₂ =
  --           fromList @(CSR S.Vector Int 5 2 Double) $
  --             [(2, 0, 0.5563116167541139), (2, 1, 0.5507459794849707), (3, 1, 0.761606773358582)]
  --         (c₂ :: [(Int, Int, Double)]) =
  --           [ (1, 1, 0.9913900269818051),
  --             (2, 0, 1.2582390351661727),
  --             (2, 1, 0.5507459794849707),
  --             (3, 1, 0.761606773358582)
  --           ]
  --     withSomeCsr a₁' $ \a₁ -> withSomeCsr b₁' $ \b₁ ->
  --       case sameShape a₁ b₁ of
  --         Just Refl -> toList (a₁ + b₁) `shouldBe` c₁
  --         x -> x `shouldSatisfy` isJust
  --     toList (a₂ + b₂) `shouldBe` c₂
  -- describe "pPrimitiveOperator" $ do
  --   it ".." $ do
  --     parse pPrimitiveOperator "" ("c†↓₁₀" :: Text)
  --       `shouldBe` Right (SpinfulFermionicOperator FermionicCreationOperator 'c' SpinDown 10)
  --     parse pPrimitiveOperator "" ("f₀" :: Text)
  --       `shouldBe` Right (SpinlessFermionicOperator FermionicAnnihilationOperator 'f' 0)
  --     parse pPrimitiveOperator "" ("n↑₃₈" :: Text)
  --       `shouldBe` Right (SpinfulFermionicOperator FermionicNumberCountingOperator 'n' SpinUp 38)
  --     parse pPrimitiveOperator "" ("Sᶻ₃₈" :: Text)
  --       `shouldBe` Right (SpinOperator SpinZOperator 'S' 38)
  --     parse pOperatorString "" ("n₃₈ f↓₁₅" :: Text)
  --       `shouldBe` Right
  --         ( SpinlessFermionicOperator FermionicNumberCountingOperator 'n' 38
  --             :| [SpinfulFermionicOperator FermionicAnnihilationOperator 'f' SpinDown 15]
  --         )
  --     parse pOperatorString "" ("n↑₃₈f↓₁₅" :: Text)
  --       `shouldBe` Right
  --         ( SpinfulFermionicOperator FermionicNumberCountingOperator 'n' SpinUp 38
  --             :| [SpinfulFermionicOperator FermionicAnnihilationOperator 'f' SpinDown 15]
  --         )
  -- describe "sortByWithParity" $ do
  --   it ".." $ do
  --     sortByWithParity compare [3 :: Int, 1, 4, 8, 0] `shouldBe` (Odd, [0, 1, 3, 4, 8])
  --     sortByWithParity compare [3 :: Int, 1] `shouldBe` (Odd, [1, 3])
  --     sortByWithParity compare [1 :: Int, 3] `shouldBe` (Even, [1, 3])
  --     sortByWithParity compare ([] :: [Int]) `shouldBe` (Even, [])
  -- describe "csrMatMul" $ do
  --   it ".." $ do
  --     let (a :: StorableDenseMatrix 3 4 Float) =
  --           fromList $
  --             [ [0.0, 0.0, 0.0, 0.87383316],
  --               [0.0, 0.0, 0.0, 0.13542224],
  --               [0.24292645, 0.0, 0.0, 0.0]
  --             ]
  --         (b :: StorableDenseMatrix 4 5 Float) =
  --           fromList $
  --             [ [0.0, 0.0, 0.0, 0.67190817, 0.0],
  --               [0.80930562, 0.0, 0.0, 0.94956449, 0.51363394],
  --               [0.63021572, 0.94316889, 0.0, 0.89156391, 0.0],
  --               [0.0, 0.0, 0.64573977, 0.05604065, 0.28140917]
  --             ]
  --         (c :: StorableDenseMatrix 3 5 Float) =
  --           fromList $
  --             [ [0.0, 0.0, 0.56426882, 0.048970178, 0.24590467],
  --               [0.0, 0.0, 0.08744753, 0.00758915, 0.03810906],
  --               [0.0, 0.0, 0.0, 0.16322428, 0.0]
  --             ]
  --     csrMatMul (denseToCsr @S.Vector @Int a) (denseToCsr b) `shouldBe` (denseToCsr c)
  --   it "multiplies matrices" $ do
  --     let (a :: CSR S.Vector Int 5 2 Double) = fromList $ [(0, 0, 0.8213116953652885), (0, 1, 0.2721720739350296), (1, 1, 0.0004725096934642403), (2, 0, 0.03350833939255393), (4, 0, 0.2598645975559749)]
  --         (b :: CSR S.Vector Int 2 6 Double) = fromList $ [(0, 2, 0.7175422685574337), (0, 3, 0.2603509882468078), (0, 5, 0.20233694825026882), (1, 0, 0.9695204042199307), (1, 1, 0.25913885822625227), (1, 3, 0.3241512404858632)]
  --         (c :: CSR S.Vector Int 5 6 Double) = fromList $ [(0, 1, 0.07053036048059469), (0, 0, 0.2638763791388668), (0, 5, 0.16618170200246693), (0, 3, 0.3020542269386639), (0, 2, 0.589325857085161), (1, 3, 0.00015316460327802846), (1, 1, 0.00012244562246515968), (1, 0, 0.0004581077890052858), (2, 5, 0.006779975133623629), (2, 3, 0.008723929275360856), (2, 2, 0.02404364986332557), (4, 5, 0.05258020962776023), (4, 3, 0.06765600478405706), (4, 2, 0.1864638328480788)]
  --     csrMatMul a b `shouldBe` c

  -- describe "combineNeighbors" $ do
  --   it ".." $ do
  --     combineNeighbors (==) (+) (S.fromList [1 :: Int, 2, 3]) `shouldBe` (S.fromList [1, 2, 3])
  --     combineNeighbors (==) (+) (S.fromList [1 :: Int, 1, 3]) `shouldBe` (S.fromList [2, 3])
  --     combineNeighbors (==) (const) (S.fromList [1 :: Int, 1, 2, 1, 1, 1, 3, 3])
  --       `shouldBe` (S.fromList [1, 2, 1, 3])
  --     combineNeighbors (==) (+) (S.fromList ([] :: [Int])) `shouldBe` (S.fromList [])
  describe "pSpinOperator" $ do
    it "parses simple operators" $ do
      parse pSpinOperator "" ("σᶻ₁₀" :: Text)
        `shouldBe` Right (fromList [Scaled 1 (Generator 10 SpinZ)])
      parse pSpinOperator "" ("S⁺₀" :: Text)
        `shouldBe` Right (fromList [Scaled (fromRational (1 % 2)) (Generator 0 SpinPlus)])
    it "parses composite operators" $ do
      parse pSpinOperator "" ("σˣ₅" :: Text)
        `shouldBe` Right
          ( fromList
              [ Scaled 1 (Generator 5 SpinPlus),
                Scaled 1 (Generator 5 SpinMinus)
              ]
          )
      parse pSpinOperator "" ("σʸ₈₉₄₃" :: Text)
        `shouldBe` Right
          ( fromList
              [ Scaled (ComplexRational 0 (-1)) (Generator 8943 SpinPlus),
                Scaled (ComplexRational 0 1) (Generator 8943 SpinMinus)
              ]
          )
  describe "pOperatorString" $ do
    it ".." $ do
      parse (pOperatorString pSpinOperator) "" ("σᶻ₁₀σ⁺₁₁" :: Text)
        `shouldBe` Right [Scaled 1 [Generator 10 SpinZ, Generator 11 SpinPlus]]
      parse (pOperatorString pSpinOperator) "" ("σᶻ₁₀ σ⁺₁₁" :: Text)
        `shouldBe` Right [Scaled 1 [Generator 10 SpinZ, Generator 11 SpinPlus]]
      parse (pOperatorString pSpinOperator) "" ("σᶻ₁₀ Sˣ₅ σ⁺₁₁" :: Text)
        `shouldBe` Right
          ( [ Scaled
                (fromRational (1 % 2))
                [Generator 10 SpinZ, Generator 5 SpinPlus, Generator 11 SpinPlus],
              Scaled
                (fromRational (1 % 2))
                [Generator 10 SpinZ, Generator 5 SpinMinus, Generator 11 SpinPlus]
            ]
          )
  describe "sumToCanonical" $ do
    it ".." $ do
      -- let (Right [x₀]) =
      --       toList <$> parse (pOperatorString pSpinOperator) "" ("σᶻ₁₀ σ⁺₅ σ⁺₁₁" :: Text)
      -- productToCanonical x₀

      let (Right x₀) = parse (pOperatorString pSpinOperator) "" ("σᶻ₁₀ Sˣ₅ σ⁺₁₁" :: Text)
          y₀ :: SpinPolynomial
          y₀ =
            [ Scaled
                (fromRational (1 % 2))
                [Generator 5 SpinPlus, Generator 10 SpinZ, Generator 11 SpinPlus],
              Scaled
                (fromRational (1 % 2))
                [Generator 5 SpinMinus, Generator 10 SpinZ, Generator 11 SpinPlus]
            ]
          x₁ :: SpinPolynomial
          x₁ =
            [ Scaled
                5
                [Generator 5 SpinMinus, Generator 10 SpinZ, Generator 11 SpinPlus],
              Scaled
                3
                [Generator 11 SpinMinus, Generator 5 SpinPlus, Generator 10 SpinZ],
              Scaled
                (-1)
                [Generator 11 SpinMinus, Generator 5 SpinZ]
            ]
          y₁ :: SpinPolynomial
          y₁ =
            [ Scaled
                (-1)
                [Generator 5 SpinZ, Generator 11 SpinMinus],
              Scaled
                3
                [Generator 5 SpinPlus, Generator 10 SpinZ, Generator 11 SpinMinus],
              Scaled
                5
                [Generator 5 SpinMinus, Generator 10 SpinZ, Generator 11 SpinPlus]
            ]
      simplify x₀ `shouldBe` y₀
      simplify x₁ `shouldBe` y₁

  describe "simplify (spin)" $ do
    let parseString ::
          Text ->
          Polynomial ComplexRational (Generator Int SpinGeneratorType)
        -- Sum (Scaled ComplexRational (Product (Generator Int SpinGeneratorType)))
        parseString s = case parse (pOperatorString pSpinOperator) "" s of
          Left e -> error (show e)
          Right x -> x
    it "computes spin commutators" $ do
      simplify (parseString "σ⁺₁₀ σᶻ₁₀" - parseString "σᶻ₁₀ σ⁺₁₀")
        `shouldBe` [Scaled (-2) [Generator 10 SpinPlus]]
      simplify (parseString "σ⁺₁₀ σ⁻₁₀" - parseString "σ⁻₁₀ σ⁺₁₀")
        `shouldBe` [Scaled 1 [Generator 10 SpinZ]]
      simplify (parseString "σ⁻₁₀ σᶻ₁₀" - parseString "σᶻ₁₀ σ⁻₁₀")
        `shouldBe` [Scaled 2 [Generator 10 SpinMinus]]
    it "simplifies products of primitive operators" $ do
      simplify (parseString "σᶻ₁₀ σᶻ₁₀") `shouldBe` [Scaled 1 [Generator 10 SpinIdentity]]
      simplify (parseString "σ⁺₁₀ σ⁺₁₀") `shouldBe` []
      simplify (parseString "σ⁻₁₀ σ⁻₁₀") `shouldBe` []
      simplify (parseString "σ⁺₁₀ σᶻ₁₀") `shouldBe` [Scaled (-1) [Generator 10 SpinPlus]]
      simplify (parseString "σ⁻₁₀ σᶻ₁₀") `shouldBe` [Scaled 1 [Generator 10 SpinMinus]]
      simplify (parseString "σᶻ₁₀ σ⁺₁₀") `shouldBe` [Scaled 1 [Generator 10 SpinPlus]]
      simplify (parseString "σᶻ₁₀ σ⁻₁₀") `shouldBe` [Scaled (-1) [Generator 10 SpinMinus]]
      simplify (parseString "σ⁺₁₀ σ⁻₁₀")
        `shouldBe` [ Scaled (fromRational (1 % 2)) [Generator 10 SpinIdentity],
                     Scaled (fromRational (1 % 2)) [Generator 10 SpinZ]
                   ]
      simplify (parseString "σ⁻₁₀ σ⁺₁₀")
        `shouldBe` [ Scaled (fromRational (1 % 2)) [Generator 10 SpinIdentity],
                     Scaled (fromRational (-1 % 2)) [Generator 10 SpinZ]
                   ]
    it "simplifies products with identities" $ do
      forM_ ([SpinIdentity, SpinPlus, SpinMinus, SpinZ] :: [SpinGeneratorType]) $ \g -> do
        simplify [Scaled 1 [Generator 5 SpinIdentity, Generator 5 g]]
          `shouldBe` [Scaled (1 :: Rational) [Generator (5 :: Int) g]]
        simplify [Scaled 1 [Generator 5 g, Generator 5 SpinIdentity]]
          `shouldBe` [Scaled (1 :: Rational) [Generator (5 :: Int) g]]

  -- describe "CSR.csrMatrixFromDense" $ do
  --   it "converts dense matrices to sparse form" $ do
  --     CSR.csrMatrixFromDense ([[1, 2], [3, 4]] :: Dense.DenseMatrix S.Vector (Complex Double))
  --       `shouldBe` (CSR.CsrMatrix [0, 1, 2] [1, 0] [2, 3] [1, 4])
  -- describe "Num CSR" $ do
  --   it "(+)" $
  --     do
  --       let a = CSR.csrMatrixFromDense ([[1, 2], [0, 2]] :: Dense.DenseMatrix S.Vector (Complex Double))
  --           b = CSR.csrMatrixFromDense ([[0, 0], [1, 3]] :: Dense.DenseMatrix S.Vector (Complex Double))
  --           c = CSR.csrMatrixFromDense ([[1, 2], [1, 5]] :: Dense.DenseMatrix S.Vector (Complex Double))
  --       (a + b) `shouldBe` c

  -- describe "builds operators" $ do
  --   it ".." $ do
  --     print $ mkSpinOperator "σ⁺₁ σ⁻₂" [[0, 1], [1, 2], [2, 0]]
  --     True `shouldBe` True

  describe "NonbranchingTerm" $ do
    it ".." $ do
      let a = nonbranchingRepresentation (Generator (3 :: Int) FermionCreate)
          b = nonbranchingRepresentation (Generator (4 :: Int) FermionAnnihilate)
          c = nonbranchingRepresentation (Generator (3 :: Int) FermionAnnihilate)
          d = nonbranchingRepresentation (Generator (4 :: Int) FermionIdentity)
      print a
      print b
      case a <> a of
        (NonbranchingTerm v _ _ _ _ _) -> v `shouldBe` 0
      case b <> b of
        (NonbranchingTerm v _ _ _ _ _) -> v `shouldBe` 0
      (a <> c) `shouldBe` (nonbranchingRepresentation (Generator (3 :: Int) FermionCount))
      (d <> d) `shouldBe` d
      (a <> d) `shouldBe` a
      (d <> a) `shouldBe` a
      (b <> d) `shouldBe` b
      (d <> b) `shouldBe` b
      print (a <> b)
      print (b <> a)

  -- describe "binomialCoefficient" $ do
  --   it "computes binomial(n, k)" $ do
  --     binomialCoefficient 10 5 `shouldBe` 252
  --     binomialCoefficient 10 0 `shouldBe` 1
  --     binomialCoefficient 10 24 `shouldBe` 0
  --     binomialCoefficient 4 1 `shouldBe` 4
  describe "stateIndex" $ do
    it "computes indices of representatives" $ do
      -- isStateIndexIdentity (basisHeader (mkSpinBasis 10 Nothing)) `shouldBe` True
      -- isStateIndexIdentity (basisHeader (mkSpinBasis 10 (Just 8))) `shouldBe` False
      -- stateIndex (mkSpinBasis 10 Nothing) "|0000000101⟩" `shouldBe` Just 5
      -- stateIndex (mkSpinBasis 10 (Just 2)) "|0000000101⟩" `shouldBe` Just 1
      stateIndex (mkSpinfulFermionicBasis 3 (SpinfulPerSector 2 1)) "|001101⟩" `shouldBe` Just 1
      stateIndex (mkSpinfulFermionicBasis 3 (SpinfulPerSector 2 1)) "|100101⟩" `shouldBe` Just 7
  describe "BenesNetwork" $ do
    it "should permute bits" $ do
      let p₁ = toBenesNetwork (mkPermutation [0, 1, 2])
      permuteBits p₁ 0b100 `shouldBe` 0b100
      let p₂ = toBenesNetwork (mkPermutation [1, 2, 0])
      permuteBits p₂ 0b100 `shouldBe` 0b010
      permuteBits p₂ 0b101 `shouldBe` 0b110
      let v₃ = mkPermutation [1, 2, 3, 0, 5, 6, 7, 4, 9, 10, 11, 8]
          p₃ = toBenesNetwork v₃
      permuteBits p₃ 0b100111010010 `shouldBe` 0b110011100001
      permuteBits p₃ 0b100111010010 `shouldBe` permuteBits' v₃ 0b100111010010
  describe "mkSymmetries" $ do
    it "should build symmetry groups" $ do
      print $
        mkSymmetries
          [ mkSymmetry (mkPermutation [1, 2, 3, 0]) 0,
            mkSymmetry (mkPermutation [3, 2, 1, 0]) 1
          ]
      True `shouldBe` True

-- describe "applyOperator" $ do
--   it ".." $ do
--     let basis1 = LatticeSymmetries.Basis.SpinBasis 10 Nothing
--         operator1 = operatorFromString basis1 "σ⁺₀" [[5]]
--     -- hPutDoc stdout (pretty (opTerms operator1) <> hardline)
--     applyOperator operator1 "|0000000000⟩" `shouldBe` [(1, "|0000100000⟩")]
--     applyOperator operator1 "|0000100000⟩" `shouldBe` []
--     applyOperator operator1 "|0111011111⟩" `shouldBe` [(1, "|0111111111⟩")]
--     let basis2 = LatticeSymmetries.Basis.SpinBasis 2 Nothing
--         operator2 =
--           operatorFromString basis2 "σˣ₀ σˣ₁" [[0, 1]]
--             + operatorFromString basis2 "σʸ₀ σʸ₁" [[0, 1]]
--             + operatorFromString basis2 "σᶻ₀ σᶻ₁" [[0, 1]]
--     -- hPutDoc stdout (pretty (opTerms operator2) <> hardline)
--     -- print $ getNonbranchingTerms operator2
--     applyOperator operator2 "|00⟩" `shouldBe` [(1, "|00⟩")]
--     applyOperator operator2 "|01⟩" `shouldBe` [(2, "|10⟩"), (-1, "|01⟩")]
--     applyOperator operator2 "|10⟩" `shouldBe` [(2, "|01⟩"), (-1, "|10⟩")]
--     applyOperator operator2 "|11⟩" `shouldBe` [(1, "|11⟩")]

--     let basis3 = LatticeSymmetries.Basis.SpinfulFermionicBasis 4 SpinfulNoOccupation
--         operator3 =
--           operatorFromString basis3 "c†↑₀ c↑₁" [[0, 1], [3, 2]]
--     hPutDoc stdout (pretty (opTerms operator3) <> hardline)
--     applyOperator operator3 "|00000010⟩" `shouldBe` [(1, "|00000001⟩")]
--     applyOperator operator3 "|00000100⟩" `shouldBe` [(-1, "|00001000⟩")]


{-# LANGUAGE FlexibleContexts #-}
module Sound.Hommage.Sample
 (
 -- * Sample Datastructure
   Sample
 , arraySample
 , sizeSample
 , newSample
 , newSampleFromList

 -- * File Operations
 , openWavSample
 , saveWavSampleMono
 , saveWavSampleStereo

 -- * Pitching, Playing and Scratching
 , pitchWavSignal
 , scratchWavSignal
 , playSample
 , pitchSample
 , scratchSample
 , playCoeffs

 -- * Transforming Samples
 , mapSample
 , foldlSample
 , updateSample
 , normaliseSample
 , fadeinoutSample
 , stretchSample

 -- * Filter and Fourier Coeffs
 , Coeffs
 , arrayCoeffs
 , sizeCoeffs
 , wavesizeCoeffs
 , syntheseSample
 , analyseSample
 , filterCoeffs
 , mkPlayCoeffs
 , readCoeffs
 , writeCoeffs

 -- * Playing a Sample
 , mkTriggerSample
 , mkPlaySample
 , mkScratchSample
 , mkPitchSample

 -- * Array functions
 , interpolArray
 , mkLoopArray
 , mkPlayArray

 )
 where

import Foreign.Storable
import Data.Array.Storable
import Data.Array.IO
import Data.Int
import Data.IORef
import Data.Complex

import Sound.Hommage.Misc
import Sound.Hommage.WavFile
import Sound.Hommage.DFTFilter
import Sound.Hommage.Signal

---------------------------------------------------------------------------------------------------
-- | Open a Wav-File, but play it with variable (positive) speed\/ frequency.
pitchWavSignal :: FilePath -> [Double] -> Signal
pitchWavSignal fp = either Mono Stereo . inListE''
 (openWavSample fp >>= either (fmap Left . mkPitchSample) (fmap Right . mkPitchSample) )

-- | Open a Wav-File, but play it with variable (positive and\/or negative) speed\/ frequency.
--   Backward playing is possible.
scratchWavSignal :: FilePath -> [Double] -> Signal
scratchWavSignal fp = either Mono Stereo . inListE'
 (openWavSample fp >>= either (fmap Left . mkScratchSample) (fmap Right . mkScratchSample) )
---------------------------------------------------------------------------------------------------
-- | Creates a finite list with the content of the Sample.
playSample :: Sample a -> [a]
playSample = toList' . mkPlaySample

-- | Creates a list with the content of the Sample, but played with variable speed\/frequency.
pitchSample :: Fractional a => Sample a -> [Double] -> [a]
pitchSample s = inList'' $ mkPitchSample s

-- | Creates a list with the content of the Sample, but played with variable speed\/frequency,
--   backward playing possible.
scratchSample :: Fractional a => Sample a -> [Double] -> [a]
scratchSample s = inList' $ mkScratchSample s

-- | Creates a list with the content of 'Coeffs' (a representation of the fourier coefficients of
--   a 'Sample').
playCoeffs :: Coeffs -> [Complex Double]
playCoeffs = toList' . mkPlayCoeffs
---------------------------------------------------------------------------------------------------

---------------------------------------------------------------------------------------------------
--interpolArray :: MArray array Double IO => array Int Double -> Int -> Double -> IO Double
interpolArray :: (Fractional a, MArray array a IO) => array Int a -> Int -> Double -> IO a
interpolArray arr size =
 let f  = fromRational . toRational
 in \pos -> do
  let p0 = floor pos
      p1 = p0 + 1
      d :: Double
      d  = pos - fromIntegral p0
      d' = 1.0 - d
  v1 <- readArray arr $ mod p0 size
  v2 <- readArray arr $ mod p1 size
  let r = v1 * f d' + v2 * f d
  return r

mkLoopArray :: (Fractional a, MArray arr a IO) => arr Int a -> Int -> IO (Double -> IO a)
mkLoopArray arr size = do
 r <- newIORef 0.0
 let f  = fromRational . toRational
 return $ \dp ->
  do pos <- readIORef r
     let p  = floor pos
         p0 = mod p size
         p1 = mod (p + 1) size
         d :: Double
         d  = pos - fromIntegral p
         pos' = d + fromIntegral p0
         d' = 1.0 - d
     v1 <- readArray arr p0
     v2 <- readArray arr p1
     writeIORef r (pos' + dp)
     return (v1 * f d' + v2 * f d)

mkPlayArray :: (Fractional a, MArray arr a IO) => arr Int a -> Int -> IO (Double -> IO (Maybe a))
mkPlayArray arr size = do
 r <- newIORef 0.0
 let f  = fromRational . toRational
 return $ \dp ->
  do pos <- readIORef r
     let p0 = floor pos
         p1 = p0 + 1
         d :: Double
         d  = pos - fromIntegral p0
         d' = 1.0 - d
     if p1 >= size then return Nothing else do
      v1 <- readArray arr p0
      v2 <- readArray arr p1
      writeIORef r (pos + abs dp)
      return $ Just (v1 * f d' + v2 * f d)

---------------------------------------------------------------------------------------------------
data Sample a = SAMPLE
 { arraySample :: !(IOArray Int a)
 , sizeSample  :: !Int
 }
---------------------------------------------------------------------------------------------------
newSample :: Int -> a -> IO (Sample a)
newSample n a = newArray (0, n-1) a >>= \arr -> return $ SAMPLE arr n

newSampleFromList :: [a] -> IO (Sample a)
newSampleFromList xs = do
 let n = length xs
 arr <- newListArray (0, n - 1) xs
 return $ SAMPLE arr n

mapSample :: (a -> b) -> Sample a -> IO (Sample b)
mapSample f (SAMPLE arr n) = do
 brr <- newArray_ (0, n - 1)
 for 0 (<n) (+1) (\i -> readArray arr i >>= writeArray brr i . f)
 return $ SAMPLE brr n
---------------------------------------------------------------------------------------------------
openWavSample :: FilePath -> IO (Either (Sample Mono) (Sample Stereo))
openWavSample filepath = do
 (close, read, len, stereo) <- openInputWavFile filepath
 let bsize = 4000
 buf <- newArray_ (0, bsize-1)
 if stereo
  then do let lenh = div len 2
              hsize = div bsize 2
          arr <- newArray_ (0, lenh - 1)
          let step o i = do let i1=i*2                  -- i: maximal 1999
                            x1 <- readArray buf i1      --
                            x2 <- readArray buf (i1+1)
                            writeArray arr (o + i) (wavInt16ToDouble x1 :><: wavInt16ToDouble x2)
              loop o k | k < bsize = do read buf k
                                        for 0 (<(div k 2)) (+1) (step o)
                                        close
                                        return $ Right $ SAMPLE arr lenh
                       | otherwise = do read buf bsize
                                        for 0 (<hsize) (+1) (step o)
                                        loop (o+hsize) (k-bsize)
          loop 0 len
  else do arr <- newArray_ (0, len - 1)
          let loop o k | k < bsize = do read buf k
                                        for 0 (<k) (+1) (\i -> readArray buf i >>= writeArray arr (o + i) . wavInt16ToDouble )
                                        close
                                        return $ Left $ SAMPLE arr len
                       | otherwise = do read buf bsize
                                        for 0 (<bsize) (+1) (\i -> readArray buf i >>= writeArray arr (o + i) . wavInt16ToDouble )
                                        loop (o+bsize) (k-bsize)
          loop 0 len
---------------------------------------------------------------------------------------------------
saveWavSampleMono :: FilePath -> Sample Double -> IO ()
saveWavSampleMono filepath (SAMPLE arr len) = do
 (close, write) <- openOutputWavFileMono filepath
 let bsize = 4000
 buf <- newArray_ (0, bsize-1)
 let step o i = do x <- readArray arr (o+i)
                   writeArray buf i (wavDoubleToInt16 x)
     loop o k | k < bsize = do for 0 (<k) (+1) (step o)
                               write buf k
                               close
              | otherwise = do for 0 (<bsize) (+1) (step o)
                               write buf bsize
                               loop (o + bsize) (k - bsize)
 loop 0 len
---------------------------------------------------------------------------------------------------
saveWavSampleStereo :: FilePath -> Sample Stereo -> IO ()
saveWavSampleStereo filepath (SAMPLE arr len) = do
 (close, write) <- openOutputWavFileStereo filepath
 let bsize = 4000
     hsize = div bsize 2
 buf <- newArray_ (0, bsize-1)
 let step o i = do let i1 = i * 2
                       i2 = i1 + 1
                   x <- readArray arr (o+i)
                   writeArray buf i1 (wavDoubleToInt16 $ leftStereo x)
                   writeArray buf i2 (wavDoubleToInt16 $ rightStereo x)
     loop o k | k < hsize = do for 0 (<k) (+1) (step o)
                               write buf (2*k)
                               close
              | otherwise = do for 0 (<hsize) (+1) (step o)
                               write buf bsize
                               loop (o + hsize) (k - hsize)
 loop 0 len

--12000    348

---------------------------------------------------------------------------------------------------
mkTriggerSample :: Sample a -> a -> IO (Bool -> IO a)
mkTriggerSample s a = do
 rp <- newIORef 0
 let n = sizeSample s
     arr = arraySample s
     fun b = if b
              then readIORef rp >>= \p ->
                   if p < n
                    then writeIORef rp (p+1) >>
                         readArray arr p
                    else return a
              else writeIORef rp 0 >>
                   return a
 return fun

mkPlaySample :: Sample a -> IO (IO (Maybe a))
mkPlaySample (SAMPLE arr n) = do
 rpos <- newIORef 0
 let read = readIORef rpos >>= \pos ->
            if pos >= n then return Nothing
                        else writeIORef rpos (pos+1) >> readArray arr pos >>= return . Just
 return read
---------------------------------------------------------------------------------------------------
foldlSample :: Sample a -> (b -> a -> b) -> b -> IO b
foldlSample (SAMPLE arr n) f x =
 for' 0 (<n) (+1) x (\i x -> readArray arr i >>= \y -> return (f x y))

updateSample :: Sample a -> (a -> a) -> IO ()
updateSample (SAMPLE arr n) f =
 for 0 (<n) (+1) (\i -> readArray arr i >>= writeArray arr i . f)
---------------------------------------------------------------------------------------------------
normaliseSample :: (Ord a, Fractional a) => Sample a -> IO ()
normaliseSample s = do
 m <- foldlSample s (\b a -> max b (abs a)) 0.001
 updateSample s ((*)(1.0 / m))
---------------------------------------------------------------------------------------------------
fadeinoutSample :: Fractional a => (Int, Int) -> Sample a -> IO ()
fadeinoutSample (li,lo) sa = do
 let n = sizeSample sa
     arr = arraySample sa
     li' = min li n
     lo' = min lo n
     ol  = n - lo'
     si  = 1.0 / fromIntegral li'
     so  = 1.0 / fromIntegral lo'
 for' 0 (<li') (+1) 0.0 (\i s -> readArray arr i >>= \v -> writeArray arr i (v * s) >> return (s+si))
 for' ol (<n) (+1) 1.0 (\i s -> readArray arr i >>= \v -> writeArray arr i (v * s) >> return (s-so))
 return ()

---------------------------------------------------------------------------------------------------
stretchSample :: Fractional a => Int -> Sample a -> IO (Sample a)
stretchSample n sam = do
 let s   = sizeSample sam
     arr = arraySample sam
     d   = fromIntegral s / fromIntegral n
 ram <- newSample n 0.0
 for 0 (<n) (+1) (\i -> interpolArray arr s (d * fromIntegral i) >>= writeArray (arraySample ram) i)
 return ram

mkScratchSample :: Fractional a => Sample a -> IO (Double -> IO a)
mkScratchSample (SAMPLE arr n) = mkLoopArray arr n

mkPitchSample :: Fractional a => Sample a -> IO (Double -> IO (Maybe a))
mkPitchSample (SAMPLE arr n) = mkPlayArray arr n
---------------------------------------------------------------------------------------------------
data Coeffs = COEFFS
 { arrayCoeffs    :: IOArray Int CoeffArr
 , sizeCoeffs     :: Int
 , wavesizeCoeffs :: Int
 }

mkPlayCoeffs :: Coeffs -> IO (IO (Maybe (Complex Double)))
mkPlayCoeffs (COEFFS arr siz wsiz) = do
 r <- newIORef (0,0)
 let read = readIORef r >>= \(n,k) ->
            if n >= siz then return Nothing else
            if k >= 512 then writeIORef r (n+1,0) >> read else
            readArray arr n >>= \cs ->
            readCoeffArr cs k >>= \c ->
            writeIORef r (n,k+1) >>
            return (Just c)
 return read

analyseSample :: Sample Mono -> IO Coeffs
analyseSample s = do
 let arr = arraySample s
     siz   = sizeSample s
     n     = div siz 512
     r     = mod siz 512
     n'    = if r == 0 then n-1 else n
     n''   = n' - 1
     r'    = if r == 0 then 1024 else r+512
     p'    = n'' * 512
 dummy <- newArray_ (0,1023)
 brr <- newArray (0, n'') undefined
 for 0 (<n'') (+1) (\i -> do let p = i * 512
                             for 0 (<1024) (+1) (\i -> readArray arr (i+p) >>= writeArray dummy i)
                             crr <- newArray_ (0, 1023)
                             writeArray brr i crr
                             analyseDFT dummy crr)
 crr <- newArray_ (0, 1023)
 writeArray brr n'' crr
 for 0 (<r') (+1) (\i -> readArray arr (i+p') >>= writeArray dummy i)
 for r' (<1024) (+1) (\i -> writeArray dummy i 0.0)
 analyseDFT dummy crr
 return $ COEFFS brr n' siz

syntheseSample :: Coeffs -> IO (Sample Mono)
syntheseSample (COEFFS arr n siz) = do
 so <- newSample siz 0.0
 let n' = n-1
 dummy <- newArray (0, 1023) 0.0
-- sa :: StoraArray Int Double) <- newArray (0, 1023) 0.0
 for 0 (<n') (+1) (\i -> do crr <- readArray arr i
                            syntheseDFT crr dummy
                            kurveDFT dummy
                            let pl = i * 512
                            for 0 (<512) (+1) (\j -> let jpl = j + pl
                                                         jpr = jpl + 512
                                                     in
                                                     readArray (arraySample so) jpl >>= \x ->
                                                     readArray dummy j >>= \y ->
                                                     writeArray (arraySample so) jpl (x+y) >>
                                                     readArray dummy (j+512) >>= \z ->
                                                     writeArray (arraySample so) jpr z
                                              )
                  )
 crr <- readArray arr n'
 syntheseDFT crr dummy
 kurveDFT dummy
 let pl = n' * 512
     pm = siz - (pl + 512)
 for 0 (<512) (+1) (\j -> let jpl = j + pl
                          in
                          readArray (arraySample so) jpl >>= \x ->
                          readArray dummy j >>= \y ->
                          writeArray (arraySample so) jpl (x+y) )
 for 0 (<pm) (+1) (\j -> readArray dummy (j+512) >>= \z ->
                          writeArray (arraySample so) (j+pl+512) z )
 return so

filterCoeffs :: IO (CoeffMap ()) -> Coeffs -> IO Coeffs
filterCoeffs mkf (COEFFS arrci n siz) = do
 arrco <- newArray (0, n-1) undefined
 f   <- mkf
 for 0 (<n) (+1) (\i -> do ci <- readArray arrci i
                           co <- newArray_ (0, 1023)
                           f ci co
                           writeArray arrco i co)
 return $ COEFFS arrco n siz

-- | usage: @ readCoeffs obj nrOfCoeffarr nrOfCoefficient @
readCoeffs :: Coeffs -> Int -> Int -> IO (Complex Double)
readCoeffs cf a c = readArray (arrayCoeffs cf) a >>= \ca ->
                    readCoeffArr ca c

writeCoeffs :: Coeffs -> Int -> Int -> (Complex Double) -> IO ()
writeCoeffs cf a c v = readArray (arrayCoeffs cf) a >>= \ca ->
                       writeCoeffArr ca c v

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


module ReadSpec (spec) where

import           Data.Complex
import           Data.Ratio
import           Read
import           Test.Hspec
import           Types

spec :: Spec
spec =
  let r = extractValue . readExpr in
  describe "readExpr" $ do
    it "parseString" $ do
      r "\"foo\"" `shouldBe` String "foo"
      r "\"foo\\\"bar\\\"baz\"" `shouldBe` String "foo\"bar\"baz"

    it "parseBool" $ do
      r "#t" `shouldBe` Bool True
      r "#f" `shouldBe` Bool False

    it "parseAtom" $
      r "foo" `shouldBe` Atom "foo"

    it "parseHex" $
      r "#xFF" `shouldBe` Number 255

    it "parseDigits" $ do
      r "#o101" `shouldBe` Number 65
      r "#d123" `shouldBe` Number 123
      r "#b101" `shouldBe` Number 5

    it "parseDecimal" $
      r "123" `shouldBe` Number 123

    it "parseCharacter" $ do
      r "#\\newline" `shouldBe` Character '\n'
      r "#\\space" `shouldBe` Character ' '
      r "#\\a" `shouldBe` Character 'a'

    it "parseFloat" $ do
      r "1.0" `shouldBe` Float 1.0
      r "0.1" `shouldBe` Float 0.1
      -- readExpr ".0" `shouldStartWith` String "No match:"

    it "parseRatio" $ do
      r "1/2" `shouldBe` Ratio (1 % 2)
      r "2/4" `shouldBe` Ratio (1 % 2)

    it "parseComplex" $ do
      r "1+2i" `shouldBe` Complex (1.0 :+ 2.0)
      r "1.23+2.46i" `shouldBe` Complex (1.23 :+ 2.46)

    it "parseList" $ do
      r "(a test)" `shouldBe` List [Atom "a", Atom "test"]
      r "(a (nested) test)" `shouldBe` List[Atom "a", List[Atom "nested"], Atom "test"]
      -- readExpr "(a test" `shouldStartWith` "No match:"

    it "parseDottedList" $ do
      r "(a b . c)" `shouldBe` DottedList [Atom "a", Atom "b"] (Atom "c")
      r "(a (b . c) d)" `shouldBe` List [Atom "a", DottedList [Atom "b"] (Atom "c"), Atom "d"]

    it "parseQuoted" $ do
      r "'(a b)" `shouldBe` List [Atom "quote", List [Atom "a", Atom "b"]]
      r "(a 'b c)" `shouldBe` List [Atom "a", List [Atom "quote", Atom "b"], Atom "c"]


-- | A basic interface between Vectors and the fftw library.
module Numeric.FFT.Vector.Base(
            -- * Transforms
            Transform(..),
            planOfType,
            PlanType(..),
            plan,
            run,
            -- * multi-demensional Transforms
            TransformND(..),
            planOfTypeND,
            planND,
            runND,
            -- * Plans
            Plan(),
            planInputSize,
            planOutputSize,
            execute,
            executeM,
            withPlanner,
            -- * Unsafe C stuff
            CFlags,
            CPlan,
            -- * Normalization helpers
            Scalable(..),
            modifyInput,
            modifyOutput,
            constMultOutput,
            multC,
            unsafeModify,
            ) where

import qualified Data.Vector.Storable as VS
import qualified Data.Vector.Storable.Mutable as MS
import Data.Vector.Generic as V hiding (forM_)
import Data.Vector.Generic.Mutable as M hiding (unsafeModify)
import Data.List as L
import Control.Concurrent.MVar
import Control.Monad.Primitive (RealWorld,PrimMonad(..), PrimBase,
            unsafePrimToPrim, unsafePrimToIO)
import Control.Monad(forM_)
import Foreign (Storable(..), Ptr, FunPtr,
                ForeignPtr, withForeignPtr, newForeignPtr)
import Foreign.C (CInt(..), CUInt, CSize(..))
import Data.Bits ( (.|.) )
import Data.Complex(Complex(..))
import Foreign.Storable.Complex()
import System.IO.Unsafe (unsafePerformIO)

#include <fftw3.h>

---------------------
-- Creating FFTW plans

-- First, the Transform flags:
data PlanType = Estimate | Measure | Patient | Exhaustive
data Preservation = PreserveInput | DestroyInput

type CFlags = CUInt

-- | Marshal the Transform flags for use by fftw.
planInitFlags :: PlanType -> Preservation -> CFlags
planInitFlags pt pr = planTypeInt .|. preservationInt
  where
    planTypeInt = case pt of
                    Estimate -> #const FFTW_ESTIMATE
                    Measure -> #const FFTW_MEASURE
                    Patient -> #const FFTW_PATIENT
                    Exhaustive -> #const FFTW_EXHAUSTIVE
    preservationInt = case pr of
                    PreserveInput -> #const FFTW_PRESERVE_INPUT
                    DestroyInput -> #const FFTW_DESTROY_INPUT

newtype CPlan = CPlan {unCPlan :: ForeignPtr CPlan}

withPlan :: CPlan -> (Ptr CPlan -> IO a) -> IO a
withPlan = withForeignPtr . unCPlan

foreign import ccall unsafe fftw_execute :: Ptr CPlan -> IO ()
foreign import ccall "&" fftw_destroy_plan :: FunPtr (Ptr CPlan -> IO ())

newPlan :: Ptr CPlan -> IO CPlan
newPlan = fmap CPlan . newForeignPtr fftw_destroy_plan

----------------------------------------
-- vector-fftw plans

-- | A 'Plan' can be used to run an @fftw@ algorithm for a specific input/output size.
data Plan a b = Plan {
                    planInput :: {-# UNPACK #-} !(VS.MVector RealWorld a),
                    planOutput :: {-# UNPACK #-} !(VS.MVector RealWorld b),
                    planExecute :: IO ()
                }

-- | The (only) valid input size for this plan.
planInputSize :: Storable a => Plan a b -> Int
planInputSize = MS.length . planInput

-- | The (only) valid output size for this plan.
planOutputSize :: Storable b => Plan a b -> Int
planOutputSize = MS.length . planOutput

-- | Run a plan on the given 'Vector'.
--
-- If @'planInputSize' p /= length v@, then calling
-- @execute p v@ will throw an exception.
execute :: (Vector v a, Vector v b, Storable a, Storable b)
            => Plan a b -> v a -> v b
execute Plan{..} = \v -> -- fudge the arity to make sure it's always inlined
    if n /= V.length v
        then error $ "execute: size mismatch; expected " L.++ show n
                    L.++ ", got " L.++ show (V.length v)
        else unsafePerformIO $ do
                        forM_ [0..n-1] $ \k -> M.unsafeWrite planInput k
                                                $ V.unsafeIndex v k
                        planExecute
                        v' <- unsafeNew m
                        forM_ [0..m-1] $ \k -> M.unsafeRead planOutput k
                                                >>= M.unsafeWrite v' k
                        V.unsafeFreeze v'
  where
    n = MS.length planInput
    m = MS.length planOutput
{-# INLINE execute #-}

-- TODO: decide whether this is actually unsafe.
-- | Run a plan on the given mutable vectors.  The same vector may be used for both
-- input and output.
--
-- If @'planInputSize' p \/= length vIn@ or @'planOutputSize' p \/= length vOut@,
-- then calling @unsafeExecuteM p vIn vOut@ will throw an exception.
executeM :: forall m v a b .
        (PrimBase m, MVector v a, MVector v b, Storable a, Storable b)
            => Plan a b -- ^ The plan to run.
            -> v (PrimState m) a  -- ^ The input vector.
                    -> v (PrimState m) b -- ^ The output vector.
                    -> m ()
executeM Plan{..} = \vIn vOut ->
    if n /= M.length vIn || m /= M.length vOut
        then error $ "executeM: size mismatch; expected " L.++ show (n,m)
                    L.++ ", got " L.++ show (M.length vIn, M.length vOut)
        else unsafePrimToPrim $ act vIn vOut
  where
    n = MS.length planInput
    m = MS.length planOutput

    act :: v (PrimState m) a -> v (PrimState m) b -> IO ()
    act vIn vOut = do
            forM_ [0..n-1] $ \k -> unsafePrimToIO (M.unsafeRead vIn k :: m a)
                                    >>= M.unsafeWrite planInput k
            unsafePrimToPrim planExecute
            forM_ [0..n-1] $ \k -> M.unsafeRead planOutput k
                                    >>= unsafePrimToIO . (M.unsafeWrite vOut k
                                                            :: b -> m ())
{-# INLINE executeM #-}

------------------
-- Malloc/free of fftw array

foreign import ccall unsafe fftw_malloc :: CSize -> IO (Ptr a)
foreign import ccall "&" fftw_free :: FunPtr (Ptr a -> IO ())

newFFTVector :: forall a . Storable a => Int -> IO (MS.MVector RealWorld a)
newFFTVector n = do
    p <- fftw_malloc $ toEnum $ n * sizeOf (undefined :: a)
    fp <- newForeignPtr fftw_free p
    return $ MS.MVector n fp
{-# INLINE newFFTVector #-}

-----------------------
-- Transforms: methods of plan creation.

-- | A transform which may be applied to vectors of different sizes.
data Transform a b = Transform {
                        inputSize :: Int -> Int,
                        outputSize :: Int -> Int,
                        creationSizeFromInput :: Int -> Int,
                        makePlan :: CInt -> Ptr a -> Ptr b -> CFlags -> IO (Ptr CPlan),
                        normalization :: Int -> Plan a b -> Plan a b
                    }

-- | Create a 'Plan' of a specific size for this transform.
planOfType :: (Storable a, Storable b) => PlanType
                                -> Transform a b -> Int -> Plan a b
planOfType ptype Transform{..} n
  | m_in <= 0 || m_out <= 0 = error "Can't (yet) plan for empty arrays!"
  | otherwise  = unsafePerformIO $ do
    planInput <- newFFTVector m_in
    planOutput <- newFFTVector m_out
    MS.unsafeWith planInput $ \inP -> MS.unsafeWith planOutput $ \outP -> do
        pPlan <- makePlan (toEnum n) inP outP $ planInitFlags ptype DestroyInput
        cPlan <- newPlan pPlan
        -- Use unsafeWith here to ensure that the Storable MVectors' ForeignPtrs
        -- aren't released too soon:
        let planExecute = MS.unsafeWith planInput $ \_ ->
                            MS.unsafeWith planOutput $ \_ ->
                                withPlan cPlan fftw_execute
        return $ normalization n $ Plan {..}
  where
    m_in = inputSize n
    m_out = outputSize n
{-# INLINE planOfType #-}

-- | Create a 'Plan' of a specific size.  This function is equivalent to
-- @'planOfType' 'Estimate'@.
plan :: (Storable a, Storable b) => Transform a b -> Int -> Plan a b
plan = planOfType Estimate
{-# INLINE plan #-}

-- | Create and run a 'Plan' for the given transform.
run :: (Vector v a, Vector v b, Storable a, Storable b)
            => Transform a b -> v a -> v b
run p = \v -> execute
            (planOfType Estimate p $ creationSizeFromInput p $ V.length v)
            v
{-# INLINE run #-}

-----------------------
-- TransformND: methods of plan creation for multi-dimensional plans.

-- | A transform which may be applied to vectors of different sizes.
--
-- @since 0.2
data TransformND a b = TransformND {
                        inputSizeND :: Int -> Int,
                        outputSizeND :: Int -> Int,
                        creationSizeFromInputND :: Int -> Int,
                        makePlanND :: CInt -> Ptr CInt -> Ptr a -> Ptr b -> CFlags -> IO (Ptr CPlan),
                        normalizationND :: VS.Vector Int -> Plan a b -> Plan a b
                    }

-- | Create a 'Plan' of a specific size for this transform.
-- 'dims' must have rank greater or equal to 1
--
-- @since 0.2
planOfTypeND :: (Storable a, Storable b) => PlanType
                                -> TransformND a b -> VS.Vector Int -> Plan a b
planOfTypeND ptype TransformND{..} dims
  | m_in <= 0 || m_out <= 0 = error "Can't (yet) plan for empty arrays!"
  | otherwise  = unsafePerformIO $ do
    mdims <- unsafeThaw $ V.map toEnum dims
    planInput <- newFFTVector m_in
    planOutput <- newFFTVector m_out
    MS.unsafeWith mdims $ \dimsP -> MS.unsafeWith planInput $ \inP -> MS.unsafeWith planOutput $ \outP -> do
        pPlan <- makePlanND (toEnum $ V.length dims) dimsP inP outP $ planInitFlags ptype DestroyInput
        cPlan <- newPlan pPlan
        -- Use unsafeWith here to ensure that the Storable MVectors' ForeignPtrs
        -- aren't released too soon:
        let planExecute = MS.unsafeWith planInput $ \_ ->
                                MS.unsafeWith planOutput $ \_ ->
                                withPlan cPlan fftw_execute
        return $ normalizationND dims $ Plan {..}
  where
    m = V.product $ V.init dims
    m_in = m * inputSizeND (V.last dims)
    m_out = m * outputSizeND (V.last dims)
{-# INLINE planOfTypeND #-}

-- | Create a 'Plan' of a specific size.  This function is equivalent to
-- @'planOfType' 'Estimate'@.
--
-- @since 0.2
planND :: (Storable a, Storable b) => TransformND a b -> VS.Vector Int -> Plan a b
planND = planOfTypeND Estimate
{-# INLINE planND #-}

-- | Create and run a 'Plan' for the given transform.
--
-- @since 0.2
runND :: (Vector v a, Vector v b, Storable a, Storable b)
            => TransformND a b -> VS.Vector Int ->  v a -> v b
runND p = \dims v ->
  let creationSize = V.init dims `V.snoc` creationSizeFromInputND p (V.last dims) in
    execute
    (planOfTypeND Estimate p creationSize)
    v
{-# INLINE runND #-}

---------------------------
-- For scaling input/output:

class Scalable a where
    scaleByD :: Double -> a -> a

instance Scalable Double where
    scaleByD = (*)
    {-# INLINE scaleByD #-}

instance Scalable (Complex Double) where
    scaleByD s (x:+y) = s*x :+ s*y
    {-# INLINE scaleByD #-}

{-# INLINE modifyInput #-}
modifyInput :: (MS.MVector RealWorld a -> IO ()) -> Plan a b -> Plan a b
modifyInput f p@Plan{..} = p {planExecute = f planInput >> planExecute}

{-# INLINE modifyOutput #-}
modifyOutput :: (MS.MVector RealWorld b -> IO ()) -> Plan a b -> Plan a b
modifyOutput f p@Plan{..} = p {planExecute = planExecute >> f planOutput}

{-# INLINE constMultOutput #-}
constMultOutput :: (Storable b, Scalable b) => Double -> Plan a b -> Plan a b
constMultOutput !s = modifyOutput (multC s)

{-# INLINE multC #-}
multC :: (Storable a, Scalable a) => Double -> MS.MVector RealWorld a -> IO ()
multC !s v = forM_ [0..n-1] $ \k -> unsafeModify v k (scaleByD s)
  where !n = MS.length v

-- | Helper function; seems like it should be in the vector package...
{-# INLINE unsafeModify #-}
unsafeModify :: (Storable a)
                => MS.MVector RealWorld a -> Int -> (a -> a) -> IO ()
unsafeModify v k f = MS.unsafeRead v k >>= MS.unsafeWrite v k . f

plannerLock :: MVar ()
plannerLock = unsafePerformIO $ newMVar ()
{-# NOINLINE plannerLock #-}

-- | Calls to the FFTW planner are non-reentrant. Here we take a mutex to
-- ensure thread safety.
withPlanner :: IO a -> IO a
withPlanner action = do
    takeMVar plannerLock
    res <- action
    putMVar plannerLock ()
    return res


{-# LANGUAGE TemplateHaskell, RankNTypes #-}
{-# LANGUAGE ScopedTypeVariables         #-}
{-# LANGUAGE MultiParamTypeClasses       #-}
{-# LANGUAGE TypeOperators               #-}
{-# LANGUAGE BangPatterns                #-}

module Chapter2Bandit
    ( Bandit(..)
    , Policy(..)
    , mkBandit
    , loopSteps
    ) where

import           Control.Monad
import           Control.Monad.IO.Class
import           Control.Monad.Trans.State

import           Control.Lens (makeLenses, over, element, (+~), (&), (.~), (%~))
import           Data.List(take, repeat)

import           Data.Random
import           Data.Random.Distribution
import           Data.Random.Distribution.Uniform
import           Data.Random.Shuffle.Weighted
import           Data.Random.Source.IO

import           Language.Haskell.TH
import           Numeric.LinearAlgebra (Vector, Matrix)
import qualified Numeric.LinearAlgebra as LA

-- project
import           Utils

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

-- | Bandit Problem
--   Three methods (Policies): epsilon-greedy, UCB, gradient bandit.
--   Policy includes two parts: action selection & state-value estimation update.

-- | For action selections (exploit & explore) of above methods:
--   1. epsilon-greedy: randomly do exploration with probability epsilon;
--   2. UCB using: Qt(a) + c sqrt(log t / Nt(a)) to balance exploit and explore
--   3. Gradient using Qt(a) as Preference value, select action by softmax;

-- | For Qt(a) (Value Estimate) Update:
--   1. 'epsilon-greedy' & 'UCB' methods using form of:
--         newVal = oldVal + stepSize * (curReward - oldVal)
--      here, 'stepSize' parameter in configuration file:
--        if < 0, use 'sample average', then it is stationary;
--        if > 0, use 'exponential, recency-weighted average',
--                then it targets non-stationary problem, weight recent reward more.
--   2. Gradient bandit using gradient approximation, estimate value as preference:
--           newValue = oldValue + stepSize * (curReward - baselineVal) * ( 1 - softmax)
--           for the selected action's state value update
--       and newValue = oldValue - stepSize * (curReward - baselineVal) * softmax
--           for all other non-selected actions'
--       where 'baseline' could be 0, or the average value of all action rewards until the time
--       and 'stepSize' is the same as above.

------------------------------------------------------------------------------------------
data Policy = EGreedy (RVar Bool) -- Bernoulli distribution with p = epsilon
            | UCB Double -- ucb explore parameter 'c'
            | Gradient Bool -- gradient baseline param, if < 0, use average reward value

data Bandit = Bandit {
     _kArms :: Int
    ,_bestValueIdx :: Int -- it is pre-known, for OptimalAction statistics
    ,_qValues :: ![Double] -- estimate value of each action
    ,_nActions :: ![Int] -- count of each atcion has taken
    ,_stepSize :: Double
    ,_totalReward :: !Double
    ,_curStep :: !Int
    ,_bestTakes :: ![Double]
    -- input randoms
    ,_srcRVars :: ![RVar Double]
    -- exploit & explore methods
    ,_policy :: Policy
    }

makeLenses ''Bandit

mkBandit :: Int -> Int -> Double -> Double -> [Double] -> Policy -> Bandit
mkBandit karm totalStep initValue stepSize trueValues policy =
  let (maxValueIdx, _) = maxWithIndex (zip [0..] trueValues)
  in  (Bandit karm maxValueIdx (take karm $ repeat initValue)
              (take karm $ repeat 0) stepSize 0.0 0 []
              (map (flip normal 1.0) trueValues) policy
      )

--------------------------------------------------------------------------------
---- learning

loopSteps :: Int -> StateT Bandit IO [Double]
loopSteps times = replicateM times step

step :: StateT Bandit IO Double
step = selectOneAction >>= takeOneAction

selectOneAction :: StateT Bandit IO Int
selectOneAction = do
  bandit <- get
  actionN <- case _policy bandit of
    EGreedy epsilonRVar -> do
      bExplore <- liftIO $ sample epsilonRVar
      if bExplore then liftIO $ sample (randomElement [0..((_kArms bandit) - 1)])
         else pure . fst . maxWithIndex $ (zip [0..] (_qValues bandit))
    UCB c -> do
      let !ucbEstValues = zipWith (calcUCB c (_curStep bandit)) (_nActions bandit) (_qValues bandit)
      pure $ fst $ maxWithIndex (zip [0..] ucbEstValues)
    Gradient _ -> do
      let gradientValues = map exp (_qValues bandit)
          gradientProbs = map ( / sum gradientValues) gradientValues
          -- Sampling according to distribution (element probability)
          -- The following line will produce good result, even 0.4 without baseline. why?
          -- Is it because we do weighted random extract twice ?
          weightedRVar = weightedShuffle $ zip gradientProbs [0..]
      liftIO $ head <$> sample weightedRVar
      -- pure $ fst $ maxWithIndex (zip [0..] gradientProbs)
  pure actionN
  where
  calcUCB :: Double -> Int -> Int -> Double -> Double
  calcUCB c total n val = val + c * sqrt ((log $ fromIntegral total) / fromIntegral n)

takeOneAction :: Int -> StateT Bandit IO Double
takeOneAction actionN = do
  bandit <- get
  reward <- liftIO $ sample (_srcRVars bandit !! actionN)
  let !bestTake = (_bestValueIdx bandit == actionN) ? (1.0, 0.0)
      !bandit' = bandit & (nActions . element actionN +~ 1)
                        & (totalReward +~ reward)
                        & (curStep +~ 1)
                        & (bestTakes %~ (++ [bestTake]))
      !bandit'' = stateUpdate bandit' actionN reward
  put bandit''
  pure reward
  where
  stateUpdate :: Bandit -> Int -> Double -> Bandit
  stateUpdate bandit actionN reward =
    let !actionTakes = (_nActions bandit) !! actionN
        !ss = (_stepSize bandit < 0) ? (1.0 / fromIntegral actionTakes, _stepSize bandit)
        !oldEstValue = (_qValues bandit) !! actionN
    in  case _policy bandit of
          Gradient bBaseline ->
            let baseline = bBaseline ? (_totalReward bandit / (fromIntegral $ _curStep bandit), 0.0)
                !gradientValues = map exp (_qValues bandit)
                !gradientProbs = map ( / sum gradientValues) gradientValues
                !newEstValues = zipWith3 (updateGradientPreference reward baseline ss)
                                         (_qValues bandit) gradientProbs [0..]
            in  bandit & (qValues .~ newEstValues)
          -- epsilone greedy & ucb use the same updating formular
          _ -> let !newEstValue = oldEstValue + (reward - oldEstValue) * ss
               in  bandit & (qValues . element actionN .~ newEstValue)

  updateGradientPreference :: Double -> Double -> Double -> Double -> Double -> Int -> Double
  updateGradientPreference reward baseline ss oldVal prob actionIdx =
    (actionN == actionIdx) ? (oldVal + ss * (reward - baseline) * (1 - prob),
                              oldVal - ss * (reward - baseline) * prob)
--------------------------------------------------------------------------------


{-# LANGUAGE NamedFieldPuns #-}
{-# LANGUAGE OverloadedStrings #-}

module Evaluate where

import AlphaBeta
import Control.Monad
import Control.Parallel.Strategies (parMap, rpar)
import Data.ByteString.Lazy.Char8 (unpack)
import Data.Csv
import Data.Maybe (fromMaybe)
import Data.Text (Text)
import Data.Vector (fromList)
import Debug.Trace
import Flow ((<|), (|>))
import GHC.Float (float2Double, int2Double)
import GHC.Generics (Generic)
import Game
import Helpers
import Minimax
import Statistics.Sample (mean, stdDev)
import System.IO (writeFile)
import System.Random (StdGen, mkStdGen, split)
import Text.Printf (printf)

evaluate =
  let results =
        map
          ( \(s1, ai1) ->
              map
                ( \(s2, ai2) ->
                    let (utilities, lengths) =
                          playOften
                            ( \a ->
                                case a of
                                  Gold -> ai1
                                  Silver -> ai2
                            )
                        vUtilities = utilities |> map float2Double |> fromList
                        vLengths = lengths |> map int2Double |> fromList
                     in ( \_ ->
                            ( s1,
                              s2,
                              vUtilities |> mean,
                              vUtilities |> stdDev,
                              vLengths |> mean,
                              vLengths |> stdDev
                            )
                        )
                )
                ais
          )
          ais
          |> concat
          |> parMap rpar (\a -> a ())
   in let path = "evaluation/results.csv"
       in do
            ( writeFile
                path
                ( "Gold,Silver,Utility,Utility (StdDev),Ply-Depth,Ply-Depth (StdDev)\n"
                    ++ (unpack (encode (results)))
                )
              )
            putStrLn ("Written evaluation to ./" ++ path)
            putStrLn "Run `cd evaluation && pipenv run python __main__.py` to create visualizations."

ais :: [(Text, (StdGen -> Ai))]
ais =
  [ ("a Random", random),
    ("b Minimax 1", minimax 1),
    --("c AlphaBeta 1", alphaBeta 1),
    ("c Minimax 2", minimax 2),
    --("e AlphaBeta 2", alphaBeta 2),
    ("d Minimax 3", minimax 3)
    --("g AlphaBeta 3", alphaBeta 3)
  ]

-- | Run multiple AIs against each other.
playOften :: (Player -> StdGen -> Ai) -> ([Float], [Int])
playOften ais =
  [10]
    |> parMap rpar (\i -> play (mkStdGen i) ais [] (initial breakthru))
    |> foldl
      ( \(us, ls) (Utility u, history) ->
          (u : us, length history : ls)
      )
      ([], [])

-- | Play a game to the end. Takes a random number generator, a specification of a (possibly random-number-generator dependent) AI for each player, the history of states, and the current states. Returns the utility of player Gold and the history of states.
play :: StdGen -> (Player -> StdGen -> Ai) -> [State] -> State -> (Utility, [State])
play g ai history state@State {player} =
  let Game {result, utility} = breakthru
   in case utility state of
        Just u -> (u, reverse history)
        Nothing ->
          let (g1, g2) = split g
           in play
                g2
                ai
                ((traceShowId state) : history)
                ( (ai player g1 state)
                    |> fmap
                      ( \a ->
                          let r = result state a
                           in (if r == Nothing then traceShow ("result error", state, a) else id) r
                      )
                    |> (\a -> (if a == Nothing then traceShow ("ai action error", state, a) else id) a)
                    |> join
                    |> fromMaybe initial_
                )


{-# options_ghc -Wno-type-defaults #-}

{- |
Module      : Task.Block
Description : Blocks encapsulate the task of rendering a picture of a
              dynamical system to one or more tiles.
-}
module Task.Block
  ( Block(..)
  , fillBlock
  , progressively
  ) where

import Task.Concurrent

import Control.Concurrent.MVar
import Control.Monad
import Data.Word
import Foreign.ForeignPtr
import Control.Concurrent
import Data.Array.IO hiding (index)
import System.Random
import Data.Time (diffUTCTime, getCurrentTime)
import Data.Complex
import Foreign (Ptr, peekByteOff, pokeByteOff, allocaArray)

-- import Utilities (groupsOf)

-- | A Block carries the information required to go from a
--   runnable dynamical system and choice of color scheme
--   to a buffer filled with the resulting color data.
data Block =
  Block { coordToModel  :: (Double, Double) -> (Double, Double)
        , compute       :: Word32 -> Word32 -> Complex Double -> Complex Double -> Ptr Word8 -> IO ()
          -- ^ The arguments of the compute function are:
          --     * Block width, in subsamples
          --     * Subsamples per point
          --     * Array of two doubles: initial (x,y) values
          --     * Array of two doubles: (dx,dy)
          --     * Output color samples
        , logSampleRate :: Int
          -- ^ The rate of over- or under-sampling to use.
          --      * logSampleRate == 0: draw one point per pixel.
          --      * logSampleRate == N < 0: draw 2^-N by 2^-N pixel
          --          blocks per model point.
          --      * logSampleRate == N > 0: subsample each pixel
          --          on a 2^N by 2^N subgrid and average the
          --          results, for a smoother picture.
        , blockBuffer   :: Synchronizable (ForeignPtr Word8)
          -- ^ The pixel buffer to write into.
        , xStride       :: Int
          -- ^ The width of the pixel buffer.
        , x0            :: Int
          -- ^ The upper-left x coordinate of this block in the pixel buffer.
        , y0            :: Int
          -- ^ The upper-left y coordinate of this block in the pixel buffer.
        , deltaX        :: Double
          -- ^ The model size of a pixel in the X direction
        , deltaY        :: Double
          -- ^ The model size of a pixel in the Y direction
        , xSize         :: Int
          -- ^ The width of the block.
        , ySize         :: Int
          -- ^ The height of the block
        , shouldRedraw  :: MVar ()
          -- ^ A variable used to signal that this block is complete
          --   and the pixel buffer should be redrawn.
        }

-- | Render the block into its output buffer.
fillBlock :: Block -> IO ()
fillBlock Block{..} = do
    let skip = if logSampleRate < 0 then 2^(negate logSampleRate) else 1
        k    = if logSampleRate > 0 then 2^logSampleRate          else 1

    let points    = [ (x,y) | y <- [0, skip .. ySize - 1]
                            , x <- [0, skip .. xSize - 1] ]
        uv_points = [ (fromIntegral (x0 + x), fromIntegral (y0 + y))
                    | (x,y) <- points ]
        indexOf (u,v) = u + v * xStride

    -- Run the computation on each subsampled point.
    allocaArray (length points * 3) $ \tmp -> do
      compute (fromIntegral (xSize `div` skip)) (floor k)
               (deltaX :+ deltaY)
               (let (u,v) = coordToModel (fromIntegral x0, fromIntegral y0)
                in u :+ v)
               tmp

      -- Resample the results
      --let rgbs = resampleBy averageColor nSubsamples results

      -- Fill target buffer with result colors.
      with blockBuffer $ \buffer -> withForeignPtr buffer $ \ptr -> do
        forM_ (zip uv_points [0..]) $ \((u,v), ix) -> do
          let index = floor $ u + v * fromIntegral xStride
          forM_ [indexOf (du,dv) | dv <- [0 .. skip - 1]
                                 , du <- [0 .. skip - 1] ] $ \offset -> do
            peekByteOff @Word8 tmp (3*ix + 0) >>= pokeByteOff ptr (3*index + offset + 0)
            peekByteOff @Word8 tmp (3*ix + 1) >>= pokeByteOff ptr (3*index + offset + 1)
            peekByteOff @Word8 tmp (3*ix + 2) >>= pokeByteOff ptr (3*index + offset + 2)

      -- Completed the block, signal for a redraw
      void (tryPutMVar shouldRedraw ())

{-
-- | Render the block into its output buffer.
fillBlock :: Block -> IO ()
fillBlock Block{..} = do
    let skip = if logSampleRate < 0 then 2^(negate logSampleRate) else 1
        k    = if logSampleRate > 0 then 2^logSampleRate          else 1
        subsamples  = [s / k | s <- [0..k-1]]
        nSubsamples = floor (k^2)

    let points    = [ (x,y) | y <- [0, skip .. ySize - 1]
                            , x <- [0, skip .. xSize - 1] ]
        uv_points = [ (fromIntegral (x0 + x), fromIntegral (y0 + y))
                    | (x,y) <- points ]
        indexOf (u,v) = u + v * xStride

    let subsample (u,v) = [ (u + du, v + dv) | du <- subsamples, dv <- subsamples ]

    let samples = map coordToModel $ concatMap subsample uv_points

    -- Run the computation on each subsampled point.
    results <- compute samples

    -- Resample the results
    let rgbs = resampleBy averageColor nSubsamples results

    -- Fill target buffer with result colors.
    with blockBuffer $ \buffer -> withForeignPtr buffer $ \ptr -> do
      forM_ (zip uv_points rgbs) $ \((u,v), rgb) -> do
        let index = floor $ u + v * fromIntegral xStride
        forM_ [indexOf (du,dv) | dv <- [0 .. skip - 1]
                               , du <- [0 .. skip - 1] ] $ \offset -> do
          pokeColor ptr (index + offset) rgb

    -- Completed the block, signal for a redraw
    void (tryPutMVar shouldRedraw ())

resampleBy :: ([a] -> b) -> Int -> [a] -> [b]
resampleBy f n
  | n > 0     = map f . groupsOf n
  | otherwise = map f . groupsOf 1
-}

-- | Chop up a block into sub-blocks, delegate rendering
--   tasks for sub-blocks, and blit the results back
--   into an image.
progressively :: (Block -> IO ()) -> (Block -> IO ())

progressively render block = do

    let subBlockSize = 16
        width   = xSize block
        height  = ySize block
        xBlocks =  case width `divMod` subBlockSize of
            (z, 0) -> z
            (z, _) -> z + 1
        yBlocks = case height `divMod` subBlockSize of
            (z, 0) -> z
            (z, _) -> z + 1

    blockIDs <- shuffle [(x,y) | x <- [0..xBlocks - 1], y <- [0..yBlocks - 1]]

    poolSize <- subtract 1 <$> getNumCapabilities
    caps <- getNumCapabilities
    putStrLn $ show caps ++ " capabilities, pool size " ++ show poolSize ++ " (w=" ++ show (xSize block) ++ ", h=" ++ show (ySize block) ++ ")"

    let rates = filter (<= logSampleRate block) [0] -- [-4, -2, 0] --, logSampleRate block]

    let todo = [(rate, x, y) | rate <- rates, (x,y) <- blockIDs]
    putStrLn $ "***** start @ rates=" ++ show rates
    start <- getCurrentTime

    forPool_ poolSize todo $ \(rate,x,y) -> do
      render $ block { xSize = subBlockSize, ySize = subBlockSize,
                       x0 = subBlockSize * x, y0 = subBlockSize * y,
                       logSampleRate = rate }

    end <- getCurrentTime
    putStrLn $ "***** " ++ show width ++ " x " ++ show height ++ " @ rates " ++ show rates
               ++ " rendered in " ++ show (diffUTCTime end start)

shuffle :: [a] -> IO [a]
shuffle xs = do
        ar <- makeArray xs
        forM_ [2..n] $ \i -> do
            j <- randomRIO (1,i)
            vi <- readArray ar i
            vj <- readArray ar j
            writeArray ar j vi
            writeArray ar i vj
        forM [1..n] (readArray ar)
  where
    n = length xs
    makeArray :: [a] -> IO (IOArray Int a)
    makeArray = newListArray (1,n)


-- 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.
--
-- ======================================================================

{-# LANGUAGE TypeSynonymInstances #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE UndecidableInstances #-}

-- ----------------------------------------------------------------------
-- | This program reads an execution trace produced by QCL, and turns
-- it into a Quipper circuit.

module Programs.QCLParser.Main where

import Quipper hiding (cnot, Format)

import Data.IntMap (IntMap)
import qualified Data.IntMap as IntMap
import qualified Data.Traversable as Trav
import Control.Monad.State
import Prelude hiding (not)
import qualified Prelude
import System.IO
import Data.Complex
import Data.Maybe

import Text.ParserCombinators.ReadP hiding (get)
import Data.Char (isAlpha, isAlphaNum, isDigit)

-- ----------------------------------------------------------------------
-- * A monad for a QCL state

-- | In QCL, qubits are identified by integers. We have to map these
-- to Quipper's native qubits. A 'QCLState' holds such a map.
-- Implicitly, it also holds the set of qubits currently defined.
type QCLState = IntMap Qubit

-- | The 'QCLCirc' monad is like the 'Circ' monad, except that it also
-- keeps track of an additional 'QCLState'. The 'lift' function must
-- be used to lift any command for the 'Circ' monad to the 'QCLCirc'
-- monad.
type QCLCirc a = StateT QCLState Circ a

-- ----------------------------------------------------------------------
-- * Auxiliary state functions

-- | Look up the qubit corresponding to a QCL register, or allocate a
-- new qubit if it doesn't already exist.
provide :: Int -> QCLCirc Qubit
provide r = do
  s <- get
  case IntMap.lookup r s of
    Just q -> return q
    Nothing -> do
      q <- lift $ qinit False
      let s' = IntMap.insert r q s
      put s'
      return q

-- ----------------------------------------------------------------------
-- * Implementation of the QCL primitives

-- | Reset all qubits to state 0.
qcl_reset :: QCLCirc ()
qcl_reset = do
  s <- get
  lift $ Trav.mapM qdiscard s
  let s' = IntMap.empty
  put s'
  return ()

-- | Apply a controlled-not operation to the first argument.
qcl_cnot :: Int -> [Int] -> QCLCirc ()
qcl_cnot r ctrls = do
  q <- provide r
  cs <- Trav.mapM provide ctrls
  lift $ qnot_at q `controlled` cs
  return ()

-- | Apply an uncontrolled not operation.
qcl_not :: Int -> QCLCirc ()
qcl_not r = qcl_cnot r []

-- | @'qcl_fanout' ins outs ctrls@: Copy the qubit register /ins/ to
-- the qubit register /outs/ by means of c-not operations, provided
-- that /outs/ was previously 0. The whole operation is controlled by
-- /ctrls/.
qcl_fanout :: [Int] -> [Int] -> [Int] -> QCLCirc ()
qcl_fanout ins outs ctrls = do
  qins <- Trav.mapM provide ins
  qouts <- Trav.mapM provide outs
  qctrls <- Trav.mapM provide ctrls
  lift $ with_controls qctrls $ do
    let zips = zip qins qouts
    Trav.mapM (\(x,y) -> qnot_at y `controlled` x) zips
  return ()

-- | Calculate the square distance between two vectors, which must be
-- of the same length.
sqdist :: [Complex Double] -> [Complex Double] -> Double
sqdist v1 v2 = d2 where
  w = zipWith (-) v1 v2
  w2 = map (\x -> x * x) $ map magnitude w
  d2 = sum w2

-- | If the matrix looks like a /W/-gate, return 'True'.
matrix_w :: [Complex Double] -> Bool
matrix_w amps | length amps == 16 =
  let v = sqrt 0.5
      m = [ 1, 0, 0, 0,
            0, v, v, 0,
            0, v, -v, 0,
            0, 0, 0, 1]
  in
   if sqdist amps m <= 0.001 then
     True
   else
     False
matrix_w _ = False

-- | If the matrix looks like a controlled /e/^/tiZ/-gate, return the
-- angle /t/.
matrix_exp :: [Complex Double] -> Maybe Double
matrix_exp amps | length amps == 16 =
  let t = phase ((amps !! 10) + conjugate (amps !! 15))
      u = cis t
      m = [ 1, 0, 0, 0,
            0, 1, 0, 0,
            0, 0, u, 0,
            0, 0, 0, conjugate u]
  in
   if sqdist amps m <= 0.001 then
     Just t
   else
     Nothing
matrix_exp _ = Nothing

-- | @'qcl_matrix' n amps regs@: Apply the /n/-by-/n/ unitary gate
-- whose matrix is given in /amps/, to the qubit list /regs/.  This
-- function must guess, based on the complex entries of the matrix,
-- what the name of the gate should be. This guessing is crude at the
-- moment, and must be extended to include additional gates as
-- required by each algorithm. If the first argument is 'True', invert
-- the matrix.
qcl_matrix :: Bool -> Int -> [Complex Double] -> [Int] -> QCLCirc ()
qcl_matrix inv n amps [x,y] | matrix_w amps = do
  [q,r] <- Trav.mapM provide [x,y]
  lift $ gate_W_at q r  -- this gate is self-inverse
qcl_matrix inv n amps [x,y] | isJust t = do
  [q,r] <- Trav.mapM provide [x,y]
  lift $ expZt_at (sign * fromJust t) r `controlled` q
    where
      t = matrix_exp amps
      sign = if inv then -1 else 1
qcl_matrix inv n amps regs = do
  qregs <- Trav.mapM provide regs
  lift $ named_gate_at (if inv then "Gate*" else "Gate") qregs

-- | The inverse of 'qcl_cnot'.
qcl_cnot_inv = qcl_cnot

-- | The inverse of 'qcl_not'.
qcl_not_inv = qcl_not

-- | The inverse of 'qcl_fanout'.
qcl_fanout_inv = qcl_fanout

-- | A sample circuit to illustrate how to use the primitives.
testcircuit1 :: QCLCirc ()
testcircuit1 = do
  qcl_reset
  qcl_not 5
  qcl_not 0
  qcl_cnot 31 [4]
  qcl_cnot 31 [4,30]
  qcl_cnot 29 [31]
  qcl_cnot_inv 31 [4,30]
  qcl_cnot_inv 31 [4]

-- ----------------------------------------------------------------------
-- * Unpacking QCLCirc

-- | Run function for the 'QCLCirc' monad: execute the actions and
-- produce a circuit.
run :: QCLCirc a -> Circ a
run f = do
  (x,_) <- runStateT f IntMap.empty
  return x

-- ----------------------------------------------------------------------
-- * An abstract syntax for QCL output

-- | A data type to hold a QCL gate.
data QCLGate =
  Comment
  | Reset
  | Not Int
  | XNot Int
  | CNot Int [Int]
  | XCNot Int [Int]
  | Fanout [Int] [Int] [Int]
  | XFanout [Int] [Int] [Int]
  | Matrix Int [Complex Double] [Int]
  | XMatrix Int [Complex Double] [Int]
  deriving (Show)

-- | Take a gate from the abstract syntax and execute it in the
-- 'QCLCirc' monad.
do_qcl_gate :: QCLGate -> QCLCirc ()
do_qcl_gate Comment = return ()
do_qcl_gate Reset = qcl_reset
do_qcl_gate (Not x) = qcl_not x
do_qcl_gate (XNot x ) = qcl_not_inv x
do_qcl_gate (CNot x ctrl) = qcl_cnot x ctrl
do_qcl_gate (XCNot x ctrl) = qcl_cnot_inv x ctrl
do_qcl_gate (Fanout a b c) = qcl_fanout a b c
do_qcl_gate (XFanout a b c) = qcl_fanout_inv a b c
do_qcl_gate (Matrix n amps regs) = qcl_matrix False n amps regs
do_qcl_gate (XMatrix n amps regs) = qcl_matrix True n amps regs

-- ----------------------------------------------------------------------
-- * Parsing

-- $ The output of QCL consists of lines of the following forms. Lines
-- not starting with \"\@\" are comments or other non-circuit output.
--
-- > @ RESET
-- > @ NOT(qureg q=<5>)
-- > @ CNOT(qureg q=<31>,quconst c=<4,30>)
-- > @ !CNOT(qureg q=<31>,quconst c=<4,30>)
-- > @ Fanout(quconst a=<47,48,49,50,51,52>,quvoid b=<40,41,42,43,44,45>;cond=<>)
-- > @ !Fanout(quconst a=<47,48,49,50,51,52>,quvoid b=<40,41,42,43,44,45>;cond=<>)
-- > @ Matrix4x4(complex u00=1,complex u01=0,complex u02=0,complex u03=0,complex u10=0,complex u11=1,complex u12=0,complex u13=0,complex u20=0,complex u21=0,complex u22=(0.995004,-0.0998334),complex u23=0,complex u30=0,complex u31=0,complex u32=0,complex u33=(0.995004,0.0998334),qureg q=<12,13>)
-- > @ !Matrix4x4(complex u00=1,complex u01=0,complex u02=0,complex u03=0,complex u10=0,complex u11=0.707107,complex u12=0.707107,complex u13=0,complex u20=0,complex u21=0.707107,complex u22=-0.707107,complex u23=0,complex u30=0,complex u31=0,complex u32=0,complex u33=1,qureg q=<0,6>)
--
-- We use Koen Claessen's parser combinators (see
-- "Text.ParserCombinators.ReadP") to implement the parser.

-- | Parse a QCL identifier, which we take to be a non-empty string of
-- alphanumeric characters, starting with a letter
identifier :: ReadP String
identifier = do
  satisfy isAlpha
  munch isAlphaNum

-- | Parse a signless integer. We avoid the usual trick ('readS_to_P'
-- 'reads'), because this introduces backtracking errors.
int :: ReadP Int
int = do
  s <- munch1 isDigit
  return $ (read s :: Int)

-- | Parse a floating point number. We avoid the usual trick
-- ('readS_to_P' 'reads'), because this introduces backtracking
-- errors.
double :: ReadP Double
double = do
  (s, _) <- gather parse_double
  return $ (read s :: Double)
  where
    parse_double = do
      option '+' (char '+' +++ char '-')
      munch isDigit
      optional (char '.' >> munch1 isDigit)

-- | Parse a comma separated list.
commalist :: ReadP a -> ReadP [a]
commalist elt = sepBy elt (skipSpaces >> char ',' >> skipSpaces)

-- | Parse a QCL quantum register of the form
--
-- > q=<31,32>
-- > c=<4,31>
-- > b=<40,41,42,43,44,45>.
qureg :: ReadP (String, [Int])
qureg = do
  id <- identifier
  skipSpaces
  char '='
  skipSpaces
  rs <- between (char '<' >> skipSpaces) (skipSpaces >> char '>') $ do
    commalist int
  return (id, rs)

-- | Consume an optional \"!\". Return 'True' if consumed, and 'False'
-- otherwise.
inversechar :: ReadP Bool
inversechar = do
  c <- option '+' (char '!')
  return (c == '!')

-- | Parse a complex number declaration of the format
--
-- > complex u00=1
--
-- or
--
-- > complex u22=(0.995004,-0.0998334).
complex :: ReadP (Complex Double)
complex = do
  string "complex"
  skipSpaces
  identifier
  skipSpaces
  char '='
  skipSpaces
  choice [
    do
      x <- double
      return (x :+ 0)
    ,
    do
      char '('
      skipSpaces
      x <- double
      skipSpaces
      char ','
      skipSpaces
      y <- double
      skipSpaces
      char ')'
      return (x :+ y)
    ]

-- | Parse a single line of QCL output into a 'QCLGate'.
qcl_line :: ReadP QCLGate
qcl_line = choice [
  do -- @ RESET
    char '@'
    skipSpaces
    string "RESET"
    return Reset
  ,
  do -- @ NOT(qureg q=<5>)
    char '@'
    skipSpaces
    inv <- inversechar
    skipSpaces
    string "NOT"
    skipSpaces
    char '('
    skipSpaces
    string "qureg"
    skipSpaces
    (_,[r]) <- qureg
    skipSpaces
    char ')'
    return (if inv then XNot r else Not r)
  ,
  do -- @ CNOT(qureg q=<31>,quconst c=<4,30>)
    char '@'
    skipSpaces
    inv <- inversechar
    skipSpaces
    string "CNOT"
    skipSpaces
    char '('
    skipSpaces
    string "qureg"
    skipSpaces
    (_,[r]) <- qureg
    skipSpaces
    char ','
    skipSpaces
    string "quconst"
    skipSpaces
    (_,ctrls) <- qureg
    skipSpaces
    char ')'
    return (if inv then XCNot r ctrls else CNot r ctrls)
  ,
  do -- @ Fanout(quconst a=<47,48>,quvoid b=<40,41>;cond=<>)
    char '@'
    skipSpaces
    inv <- inversechar
    skipSpaces
    string "Fanout"
    skipSpaces
    char '('
    skipSpaces
    string "quconst"
    skipSpaces
    (_,ins) <- qureg
    skipSpaces
    char ','
    skipSpaces
    string "quvoid"
    skipSpaces
    (_,outs) <- qureg
    skipSpaces
    char ';'
    skipSpaces
    (_,ctrls) <- qureg
    skipSpaces
    char ')'
    return (if inv then XFanout ins outs ctrls else Fanout ins outs ctrls)
  ,
  do -- @ Matrix4x4(complex u00=1,complex u01=0,complex u02=0,complex u03=0,complex u10=0,complex u11=1,complex u12=0,complex u13=0,complex u20=0,complex u21=0,complex u22=(0.995004,-0.0998334),complex u23=0,complex u30=0,complex u31=0,complex u32=0,complex u33=(0.995004,0.0998334),qureg q=<12,13>)
    char '@'
    skipSpaces
    inv <- inversechar
    skipSpaces
    string "Matrix"
    skipSpaces
    dim1 <- int
    skipSpaces
    char 'x'
    skipSpaces
    dim2 <- int
    skipSpaces
    char '('
    skipSpaces
    amps <- commalist complex
    skipSpaces
    char ','
    skipSpaces
    string "qureg"
    skipSpaces
    (_,r) <- qureg
    skipSpaces
    char ')'
    when (dim1 /= dim2) $ do
      error "Non-square matrix"
    return (if inv then XMatrix dim1 amps r else Matrix dim1 amps r)
  ,
  do -- any line not starting with '@' is a comment
    satisfy ((/=) '@')
    munch (\x -> True)
    return Comment
  ,
  do -- empty lines are comments
    eof
    return Comment
  ]

-- | String version of 'qcl_line': parse a string and turn it into a
-- 'QCLGate'.
parse_qcl_line :: String -> QCLGate
parse_qcl_line s =
  case readP_to_S readline s of
    (h, _):_ -> h
    _ -> error ("Unrecognized line: " ++ s)
  where
    readline = do
      skipSpaces
      l <- qcl_line
      skipSpaces
      eof
      return l

-- | Monad version of 'parse_qcl_line': parse a string and execute the
-- resulting gate directly in the 'QCLCirc' monad.
run_qcl_line :: String -> QCLCirc ()
run_qcl_line = do_qcl_gate . parse_qcl_line

-- | Parse a stream consisting of many lines of QCL output and execute
-- the parsed gates in the 'QCLCirc' monad.
run_qcl_lines :: String -> QCLCirc ()
run_qcl_lines s =
  mapM_ run_qcl_line (lines s)

-- | A sample circuit to illustrate the parser.
testcircuit2 :: QCLCirc ()
testcircuit2 = do
  run_qcl_line "@ RESET"
  run_qcl_line "@ NOT(qureg q=<5>)"
  run_qcl_line "@ NOT(qureg q=<0>)"
  run_qcl_line "@ CNOT(qureg q=<31>,quconst c=<4>)"

-- ----------------------------------------------------------------------
-- * Main function

-- | Main function: read a circuit in QCL format from 'stdin', and
-- preview the translated Quipper circuit.
main :: IO ()
main = do
  lines <- hGetContents stdin
  print_simple Preview (run (run_qcl_lines lines))


{-# LANGUAGE DeriveAnyClass    #-}
{-# LANGUAGE DeriveGeneric     #-}
{-# LANGUAGE GADTs             #-}
{-# LANGUAGE NoImplicitPrelude #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE TypeFamilies      #-}
-- Ops.hs ---
--
-- Filename: Ops.hs
-- Description:
-- Author: Manuel Schneckenreither
-- Maintainer:
-- Created: Wed Nov  1 17:40:44 2017 (+0100)
-- Version:
-- Package-Requires: ()
-- Last-Updated:
--           By:
--     Update #: 155
-- URL:
-- Doc URL:
-- Keywords:
-- Compatibility:
--
--

-- Commentary:
--
--
--
--

-- Change Log:
--
--
--
--
--
--
--

-- Code:

module SimSim.Order.Ops
  ( generateOrdersContDueDateDistr
  , generateOrdersDiscDueDateDistr
  , generateOrdersFixedDueDateSlack
  , generateOrdersUniform
  ) where

import           ClassyPrelude
import           Control.DeepSeq
import           Data.Serialize
import           GHC.Generics
import           Statistics.Distribution
import           Statistics.Distribution.DiscreteUniform
import           Statistics.Distribution.Uniform
import           System.Random.MWC

import           SimSim.Order.Type
import           SimSim.ProductType
import           SimSim.Simulation.Type
import           SimSim.Time

type MinimumTime = Time

generateOrdersContDueDateDistr ::
     (ContGen interarrivalTime, DiscreteGen productType, ContGen dueDateSlack) => SimSim -> interarrivalTime -> productType -> (MinimumTime, dueDateSlack) -> IO [Order]
generateOrdersContDueDateDistr sim dInterarrival dProductType (minTime, dueDateSlack) = do
  let repeatUntilTime acc = do
        interArr <- genContVar dInterarrival gen
        if acc + interArr < periodLen
          then ((acc + interArr) :) <$> repeatUntilTime (acc + interArr)
          else return []
  arrivals <- map ((+currentTime) . timeFromDouble) <$> repeatUntilTime 0
  productTypes <- map (Product . (+ 1) . (`mod` length pts)) <$> mapM (\_ -> genDiscreteVar dProductType gen) arrivals
  dueDates <- mapM (\a -> (\x -> timeFromDouble x + a + minTime) <$> genContVar dueDateSlack gen) arrivals
  return $ zipWith3 newOrder productTypes arrivals dueDates
  where
    currentTime = simCurrentTime sim
    gen = simRandGenDemand $ simInternal sim
    periodLen = timeToDouble $ simPeriodLength sim
    pts = productTypes sim

generateOrdersDiscDueDateDistr ::
     (ContGen interarrivalTime, DiscreteGen productType, DiscreteGen dueDateSlack) => SimSim -> interarrivalTime -> productType -> (MinimumTime, dueDateSlack) -> IO [Order]
generateOrdersDiscDueDateDistr sim dInterarrival dProductType (minTime, dueDateSlack) = do
  let repeatUntilTime acc = do
        interArr <- genContVar dInterarrival gen
        if acc + interArr < periodLen
          then ((acc + interArr) :) <$> repeatUntilTime (acc + interArr)
          else return []
  arrivals <- map ((+currentTime) . timeFromDouble) <$> repeatUntilTime 0
  productTypes <- map (Product . (+ 1) . (`mod` length pts)) <$> mapM (\_ -> genDiscreteVar dProductType gen) arrivals
  dueDates <- mapM (\_ -> (\x -> x + minTime + currentTime) . Time . fromIntegral <$> genDiscreteVar dueDateSlack gen) arrivals
  return $ zipWith3 newOrder productTypes arrivals dueDates
  where
    currentTime = simCurrentTime sim
    gen = simRandGenDemand $ simInternal sim
    periodLen = timeToDouble $ simPeriodLength sim
    pts = productTypes sim

generateOrdersFixedDueDateSlack :: (ContGen interarrivalTime, ContGen productType) => SimSim -> interarrivalTime -> productType -> Time -> IO [Order]
generateOrdersFixedDueDateSlack sim dInterarrival dProductType dueDateSlack = do
  let repeatUntilTime acc = do
        interArr <- genContVar dInterarrival gen
        if acc + interArr < periodLen
          then ((acc + interArr) :) <$> repeatUntilTime (acc + interArr)
          else return []
  arrivals <- map ((+currentTimeMinusPeriod) . timeFromDouble) <$> repeatUntilTime 0
  -- print currentTimeMinusPeriod
  -- print $ map timeToDouble arrivals
  -- productTypes <- map (Product . (+ 1) . (`mod` length pts)) <$> mapM (\_ -> genDiscreteVar dProductType gen) arrivals
  productTypes <- map (Product . (+ 1) . (`mod` length pts) . round) <$> mapM (\_ -> genContVar dProductType gen) arrivals
  let dueDates = replicate (length arrivals) (currentTime + dueDateSlack)
  return $ zipWith3 newOrder productTypes arrivals dueDates
  where
    currentTimeMinusPeriod = currentTime - timeFromDouble periodLen
    currentTime = simCurrentTime sim
    gen = simRandGenDemand $ simInternal sim
    periodLen = timeToDouble $ simPeriodLength sim
    pts = productTypes sim

generateOrdersUniform :: SimSim -> Int -> Int -> Time -> IO [Order]
generateOrdersUniform sim minOrders maxOrders dueDateSlack = do
  nrOrders <- round <$> genContVar (uniformDistr (fromIntegral minOrders-0.4999) (fromIntegral maxOrders+0.4999)) gen
  let arrivals = replicate nrOrders currentTime
  -- productTypes <- map Product <$> mapM (\_ -> genDiscreteVar (discreteUniform (length pts)) gen) arrivals
  productTypes <- map (Product . round) <$> mapM (\_ -> genContVar (uniformDistr 0.5001 (fromIntegral (length pts)+0.4999)) gen) arrivals
  let dueDates = replicate (length arrivals) (currentTime + dueDateSlack)
  return $ zipWith3 newOrder productTypes arrivals dueDates
  where
    currentTime = simCurrentTime sim
    gen = simRandGenDemand $ simInternal sim
    pts = productTypes sim

--
-- Ops.hs ends here


module HashedExpression.Utils where

import Data.Array
import Data.Complex (Complex (..))
import qualified Data.Complex as Complex
import Data.Function ((&))
import qualified Data.IntMap.Strict as IM
import Data.List (foldl')
import Data.List.Split (splitOn)
import Data.Map (Map, fromList)
import qualified Data.Map as Map
import Data.Maybe
import qualified Data.Set as Set
import qualified Data.Text as T
import Data.Time (diffUTCTime, getCurrentTime)
import GHC.IO.Unsafe (unsafePerformIO)
import GHC.Stack (HasCallStack)
import HashedExpression.Internal.Base
import HashedExpression.Internal.Node
import Prelude hiding ((^))
import qualified Prelude

-- | Forward pipe operator
(|>) :: a -> (a -> b) -> b
(|>) = (&)

infixl 1 |>

-- |
measureTime :: IO a -> IO ()
measureTime action = do
  beforeTime <- getCurrentTime
  action
  afterTime <- getCurrentTime
  putStrLn $ "Took " ++ show (diffUTCTime afterTime beforeTime) ++ " seconds"

-- | Check if all elements of the list is equal
allEqual :: (Eq a) => [a] -> Bool
allEqual [] = True
allEqual [x] = True
allEqual (x : y : xs) = x == y && allEqual (y : xs)

fromR :: Double -> Complex Double
fromR x = x :+ 0

-- |
pullConstant :: ExpressionMap -> NodeID -> Maybe (Shape, Double)
pullConstant mp n
  | (shape, R, Const c) <- retrieveNode n mp = Just (shape, c)
  | otherwise = Nothing

-- |
pullConstants :: ExpressionMap -> [NodeID] -> Maybe (Shape, [Double])
pullConstants mp ns
  | xs@(x : _) <- mapMaybe (pullConstant mp) ns = Just (fst x, map snd xs)
  | otherwise = Nothing

-- |
isZero :: ExpressionMap -> NodeID -> Bool
isZero mp nId
  | Const 0 <- retrieveOp nId mp = True
  | RealImag arg1 arg2 <- retrieveOp nId mp,
    Const 0 <- retrieveOp arg1 mp,
    Const 0 <- retrieveOp arg2 mp =
    True
  | otherwise = False

-- |
isOne :: ExpressionMap -> NodeID -> Bool
isOne mp nId
  | Const 1 <- retrieveOp nId mp = True
  | RealImag arg1 arg2 <- retrieveOp nId mp,
    Const 1 <- retrieveOp arg1 mp,
    Const 0 <- retrieveOp arg2 mp =
    True
  | otherwise = False

-- |
isConstant :: ExpressionMap -> NodeID -> Bool
isConstant mp nId
  | Const _ <- retrieveOp nId mp = True
  | otherwise = False

-- |
pullSumOperands :: ExpressionMap -> NodeID -> [NodeID]
pullSumOperands mp nId
  | Sum operands <- retrieveOp nId mp = operands
  | otherwise = [nId]

-- |
pullProdOperands :: ExpressionMap -> NodeID -> [NodeID]
pullProdOperands mp nId
  | Mul operands <- retrieveOp nId mp = operands
  | otherwise = [nId]

toRange :: DimSelector -> Int -> (Int, Int, Int)
toRange (At i) size = (i, i, 1)
toRange (Range start end step) size = (start, if end >= start then end else end + size, step)

mkIndices :: DimSelector -> Int -> [Int]
mkIndices (At i) size = [i]
mkIndices (Range start end step) size =
  map (`mod` size) [start, start + step .. if end >= start then end else end + size]

breakSub :: Int -> [a] -> [([a], [a], [a])]
breakSub len [] = []
breakSub len xs | len > length xs = []
breakSub len xs@(y : ys) =
  ([], take len xs, drop len xs) :
  map (\(prefix, sub, suffix) -> (y : prefix, sub, suffix)) (breakSub len ys)

coercible :: Shape -> Shape -> Bool
coercible sml target =
  any
    ( \(prefix, sub, suffix) ->
        sub == sml && all (== 1) prefix && all (== 1) suffix
    )
    $ breakSub (length sml) target

-------------------------------------------------------------------------------
zipMp :: forall a b c. Ord a => Map a b -> Map a c -> Map a (b, c)
zipMp mp1 mp2 = foldl' f Map.empty $ Map.keys mp1
  where
    f acc k = case (Map.lookup k mp1, Map.lookup k mp2) of
      (Just v1, Just v2) -> Map.insert k (v1, v2) acc
      _ -> acc

-------------------------------------------------------------------------------
zipMp3 :: forall a b c d. Ord a => Map a b -> Map a c -> Map a d -> Map a (b, c, d)
zipMp3 mp1 mp2 mp3 = foldl' f Map.empty $ Map.keys mp1
  where
    f acc k = case (Map.lookup k mp1, Map.lookup k mp2, Map.lookup k mp3) of
      (Just v1, Just v2, Just v3) -> Map.insert k (v1, v2, v3) acc
      _ -> acc

instance PowerOp Double Int where
  (^) x y = x Prelude.^ y

instance ComplexRealOp Double (Complex Double) where
  (+:) = (:+)
  xRe = Complex.realPart
  xIm = Complex.imagPart
  conjugate = Complex.conjugate

-------------------------------------------------------------------------------
class ToText a where
  tt :: a -> T.Text
  ttq :: a -> T.Text
  ttq s = "\"" <> tt s <> "\""

instance ToText Double where
  tt s = T.pack $ show s

instance ToText Int where
  tt s = T.pack $ show s

instance ToText String where
  tt = T.pack

instance ToText T.Text where
  tt = id