{-# LANGUAGE BangPatterns        #-}
{-# LANGUAGE CPP                 #-}
{-# LANGUAGE TemplateHaskell     #-}
{-# LANGUAGE DataKinds           #-}
{-# LANGUAGE KindSignatures      #-}
{-# LANGUAGE TypeApplications    #-}
{-# LANGUAGE GADTs               #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# OPTIONS_GHC -fno-warn-missing-signatures #-}

module Test.Grenade.Layers.PadCrop where

import           Grenade

import           Hedgehog

import           Numeric.LinearAlgebra.Static ( norm_Inf )

import           Test.Hedgehog.Hmatrix

import           Data.Serialize
import           Data.Either
import           System.Random.MWC             (create)

type PadCropNet3D = Network '[Pad 2 3 4 6, Crop 2 3 4 6] '[ 'D3 7 9 5, 'D3 16 15 5, 'D3 7 9 5 ]
type PadCropNet2D = Network '[Pad 2 3 4 6, Crop 2 3 4 6] '[ 'D2 7 9, 'D2 16 15, 'D2 7 9 ]

prop_pad_crop :: Property
prop_pad_crop =
  let net :: PadCropNet3D
      net = Pad :~> Crop :~> NNil
  in  property $
    forAll genOfShape >>= \(d :: S ('D3 7 9 5)) ->
      let (tapes, res)  = runForwards  net d
          (_    , grad) = runBackwards net tapes d
      in  do assert $ d ~~~ res
             assert $ grad ~~~ d

prop_pad_crop_2d :: Property
prop_pad_crop_2d =
  let net :: Network '[Pad 2 3 4 6, Crop 2 3 4 6] '[ 'D2 7 9, 'D2 16 15, 'D2 7 9 ]
      net = Pad :~> Crop :~> NNil
  in  property $
    forAll genOfShape >>= \(d :: S ('D2 7 9)) ->
      let (tapes, res)  = runForwards  net d
          (_    , grad) = runBackwards net tapes d
      in  do assert $ d ~~~ res
             assert $ grad ~~~ d

prop_pad_crop_is_serializable :: Property
prop_pad_crop_is_serializable = withTests 1 $ property $ do
  gen <- evalIO create
  pad  :: Pad 2 3 4 6  <- evalIO $ createRandomWith UniformInit gen
  crop :: Crop 2 3 4 6 <- evalIO $ createRandomWith UniformInit gen
  let net :: PadCropNet3D
      net = pad :~> crop :~> NNil
      bs = encode net
      dec = decode bs :: Either String PadCropNet3D
  assert $ isRight dec

prop_can_show_pad_crop :: Property
prop_can_show_pad_crop =
  let net :: PadCropNet3D
      net = Pad :~> Crop :~> NNil
  in (withTests 1 . property) $ show net `seq` success

prop_can_update_pad_crop_and_use_in_batches :: Property
prop_can_update_pad_crop_and_use_in_batches =
  let net :: PadCropNet3D
      net = Pad :~> Crop :~> NNil
  in  property $
    forAll genOfShape >>= \(d :: S ('D3 7 9 5)) ->
      let (tapes, _) = runForwards  net d
          (v    , _)   = runBackwards net tapes d
      in do
        runUpdate defSGD net v `seq` success
        runUpdate defAdam net v `seq` success
        reduceGradient @PadCropNet3D [v] `seq` success

(~~~) :: S x -> S x -> Bool
(S1D x) ~~~ (S1D y) = norm_Inf (x - y) < 0.00001
(S2D x) ~~~ (S2D y) = norm_Inf (x - y) < 0.00001
(S3D x) ~~~ (S3D y) = norm_Inf (x - y) < 0.00001

tests :: IO Bool
tests = checkParallel $$(discover)


{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE PolyKinds #-}
{-# LANGUAGE NamedFieldPuns #-}
{-# LANGUAGE NamedFieldPuns #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE TypeSynonymInstances #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE DeriveGeneric #-}
{-# LANGUAGE DeriveFunctor #-}
{-# LANGUAGE ConstraintKinds #-}

module Matrix where

import Prelude
import qualified Prelude
import qualified Data.Matrix
import qualified Data.Eigen.Matrix
import Foreign.C.Types
import qualified Numeric.LinearAlgebra.HMatrix
import qualified Numeric.LinearAlgebra.Data
import Control.Parallel.Strategies
import qualified Data.Array as Array
import Data.Array(Array, (!), range, Ix)
import GHC.Generics
import Control.Arrow
import Control.Monad.Writer
import Control.Monad.Identity

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

data Matrix a b v = Matrix (Array (a, b) v) deriving (Show, Eq, Ord, Functor, Generic)
instance (NFData a, NFData b, NFData v) => NFData (Matrix a b v)
type Vector a v = Matrix a () v

type Ix' a = (Ix a, Bounded a, Enum a)

instance (Ix' a, Ix' b) => Enum (a, b) where
  fromEnum (x, (y :: b)) = fromEnum x * Array.rangeSize (fullRange :: (b, b)) + fromEnum y
  toEnum i =
    let (d, m) = i `divMod` (Array.rangeSize (fullRange :: (b, b))) in
      (toEnum d, (toEnum m :: b))

instance (Linear v, Ix' a, Ix' b) => Linear (Matrix a b v) where
  zero = Matrix (f_array (\(_a, _b) -> zero))
  add = matrixZipWith add
  minus = fmap minus

fullRange :: Bounded a => (a, a)
fullRange = (minBound, maxBound)

vector_of_function f = Matrix (f_array (\(x, ()) -> f x))

{-# INLINE matrix_mult #-}
matrix_mult :: (Ix' a, Ix' b, Ix' c, Num v) => Matrix a b v -> Matrix b c v -> Matrix a c v
matrix_mult (Matrix x) (Matrix y) = Matrix (Array.array fullRange [ ((a, c), sum [ x ! (a, b) * y ! (b, c) | b <- range fullRange]) | a <- range fullRange, c <- range fullRange])

{-# INLINE matrixZipWith #-}
matrixZipWith f (Matrix a) (Matrix b) =
  Matrix (f_array (\k -> f (a ! k) (b ! k)))

{-# INLINE f_array #-}
f_array :: Ix' k => (k -> v) -> Array k v
f_array f = Array.array fullRange [(k, f k) | k <- range fullRange]

type Solution b v = Map b (Map b v)

type Solution_with_trace a b v =
  -- decomposes as much products as it can, providing for each product:
  -- 1. how much of non-decomposable products are necessary
  -- 2. how much of input recipes ([a]) you need to craft
  Map b (Map b v, Map a v)

class Linear a where
  zero :: a
  add :: a -> a -> a
  minus :: a -> a

instance (Linear a, Linear b) => Linear (a, b) where
  zero = (zero, zero)
  add (a1, b1) (a2, b2) = (add a1 a2, add b1 b2)
  minus (a, b) = (minus a, minus b)

instance Linear Double where
  zero = 0
  add = (+)
  minus x = -x

instance Linear Rational where
  zero = 0
  add = (+)
  minus x = -x

instance VectorSpace Rational where
  type Scalar Rational = Rational
  scale = (*)

instance (Eq v, Ord k, VectorSpace v) => VectorSpace (Map k v) where
  type Scalar (Map k v) = Scalar v
  scale s = fmap (scale s)

instance (VectorSpace a, VectorSpace b, Scalar a ~ Scalar b) => VectorSpace (a, b) where
  type Scalar (a, b) = Scalar a
  scale s = (scale s *** scale s)

instance (Ord k, Linear v, Eq v) => Linear (Map k v) where
  zero = Map.empty
  add a b = Map.filter (/= zero) (Map.unionWith add a b)
  minus = fmap minus

find_kernel ::
  forall v1 b s.
  ( Linear v1
  , Eq b
  , Eq v1
  , Ord b)
  =>
  (v1 -> v1 -> s)
  -> (s -> v1 -> v1)
  -> [Map b v1] -> Solution b s
find_kernel divide mult_v1 rows = go rows where

  lookupLhs :: (Map b v1) -> b -> v1
  lookupLhs lhs b = case Map.lookup b lhs of
    Nothing -> zero
    Just x -> x

  addRow a1 a2 = (Map.unionWith add a1 a2)

  scaleRow (s :: s) a1 = fmap (mult_v1 s) a1

  minusRow m = fmap minus m

  remove_b :: (b, (Map b v1)) -> (Map b v1) -> (Map b v1)
  remove_b (b, row0) row1 =
    (\m -> if m Map.! b == zero then Map.delete b m else error "should be zero") $ addRow row1 (minusRow $ scaleRow (lookupLhs row1 b `divide` lookupLhs row0 b) row0)

  go :: [(Map b v1)] -> Solution b s
  go [] = Map.empty
  go (row0 : rest) = case [(b, v) | (b, v) <- Map.toList row0, v /= zero] of
    [] -> -- all coefficients are 0, equation is useless
      go rest
    (chosen_b, chosen_v) : _ ->
      case go (map (remove_b (chosen_b, row0)) rest) of
        solutionDecompose ->
          case map (\(b, v) ->
                           if b == chosen_b
                           then Map.empty
                           else
                             case (Map.lookup b solutionDecompose) of
                               Nothing ->
                                 Map.singleton b v
                               Just vs ->
                                 fmap (`mult_v1` v) vs
                        ) (Map.toList row0) of
            lhs ->
              let sum_of_lhs = Map.unionsWith add lhs in
              (Map.insert chosen_b (fmap ((`divide` chosen_v) . minus) sum_of_lhs) solutionDecompose)

class Linear a => VectorSpace a where
  type Scalar a :: *
  scale :: Scalar a -> a -> a

find_kernel_with_trace ::
  forall v a b s.
  ( VectorSpace v
  , Scalar v ~ v
  , Ord a
  , Fractional v
  , Eq b
  , Eq v
  , Num v
  , Ord b)
  =>
  (v -> v -> v)
  -> (v -> v -> v)
  -> Map a (Map b v) -> Solution_with_trace a b v
find_kernel_with_trace divide mult_v1 rows = go (map (\(a, row) -> (row, Map.singleton a 1)) $ Map.toList rows) where

  lookupLhs :: (Map b v, Map a v) -> b -> v
  lookupLhs (lhs, _) b = case Map.lookup b lhs of
    Nothing -> zero
    Just x -> x

  addRow = add

  scaleRow = scale

  minusRow = minus

  remove_b :: (b, (Map b v, Map a v)) -> (Map b v, Map a v) -> (Map b v, Map a v)
  remove_b (b, row0) row1 =
    (\(m, r) -> if Map.member b m then error "should be absent" else (m, r)) $ addRow row1 (minusRow $ scaleRow (lookupLhs row1 b `divide` lookupLhs row0 b) row0)

  go :: [(Map b v, Map a v)] -> Solution_with_trace a b v
  go [] = Map.empty
  go (row0 : rest) = case [(b, v) | (b, v) <- Map.toList (fst row0), v /= zero] of
    [] -> -- all coefficients are 0, equation is useless
      go rest
    (chosen_b, chosen_v) : _ -> runIdentity $ do
      row0 <- return $ scaleRow (recip chosen_v) row0
      chosen_v <- return $ ()
      rest <- return $ map (remove_b (chosen_b, row0)) rest
      (lhs, rhs) <- return $ row0
      row0 <- return $ ()
      return $
        case go rest of
          solutionDecompose ->
            case map (\(b, v) ->
                             if b == chosen_b
                             then zero
                             else
                               case Map.lookup b solutionDecompose of
                                 Nothing -> do
                                   (Map.singleton b v, zero) -- don't need any recipes to make the raw material
                                 Just results ->
                                   scale v results
                          ) (Map.toList lhs) of
              lhs ->
                let sum_of_lhs = foldr add zero lhs in
                (Map.insert chosen_b (add (minus sum_of_lhs) (zero, rhs)) solutionDecompose)


module STCR2Z1T0PointSet where

import           Data.Array.Repa         as R
import           Data.Binary             (decodeFile)
import           Data.Complex
import           Data.List               as L
import           DFT.Plan
import           FokkerPlanck.MonteCarlo
import           FokkerPlanck.Pinwheel
import           Image.IO
import           STC
import           System.Directory
import           System.Environment
import           System.FilePath
import           Types
import           Utils.Array

main = do
  args@(numPointStr:numOrientationStr:sigmaStr:taoStr:lenStr:initialScaleStr:numTrailStr:maxTrailStr:theta0FreqsStr:thetaFreqsStr:histFilePath:pinwheelFlagStr:numIterationStr:writeSourceFlagStr:numThreadStr:_) <-
    getArgs
  print args
  let numPoint = read numPointStr :: Int
      numOrientation = read numOrientationStr :: Int
      sigma = read sigmaStr :: Double
      tao = read taoStr :: Double
      len = read lenStr :: Int
      initialScale = read initialScaleStr :: Double
      numTrail = read numTrailStr :: Int
      maxTrail = read maxTrailStr :: Int
      theta0Freq = read theta0FreqsStr :: Double
      theta0Freqs = [-theta0Freq .. theta0Freq]
      thetaFreq = read thetaFreqsStr :: Double
      thetaFreqs = [-thetaFreq .. thetaFreq]
      pinwheelFlag = read pinwheelFlagStr :: Bool
      numIteration = read numIterationStr :: Int
      writeSourceFlag = read writeSourceFlagStr :: Bool
      numThread = read numThreadStr :: Int
      folderPath = "output/test/STCR2Z1T0PointSet"
  createDirectoryIfMissing True folderPath
  flag <- doesFileExist histFilePath
  radialArr <-
    if flag
      then R.map magnitude . getNormalizedHistogramArr <$>
           decodeFile histFilePath
      else do
        putStrLn "Couldn't find a Green's function data. Start simulation..."
        solveMonteCarloR2Z1T0Radial
          numThread
          numTrail
          maxTrail
          numPoint
          numPoint
          sigma
          tao
          initialScale
          theta0Freqs
          thetaFreqs
          histFilePath
          (emptyHistogram
             [ (round . sqrt . fromIntegral $ 2 * (div numPoint 2) ^ 2)
             , L.length theta0Freqs
             , L.length thetaFreqs
             ]
             0)
  arrR2Z1T0 <-
    computeUnboxedP $
    computeR2Z1T0ArrayRadial
      radialArr
      numPoint
      numPoint
      1
      thetaFreqs
      theta0Freqs
  plan <- makeR2Z1T0Plan emptyPlan arrR2Z1T0
  let n = 30
      m = round $ (fromIntegral n) * (sqrt 2) / 2
      a = 10
      b = -10
      c = 10
      a' = round $ (fromIntegral a) * (sqrt 2) / 2
      b' = round $ (fromIntegral b) * (sqrt 2) / 2
      c' = round $ (fromIntegral c) * (sqrt 2) / 2
      r = 30
      numTheta = 7
      deltaTheta = (1 * pi) / numTheta
      xs =
        ((L.map
            (\(i, j) -> R2S1RPPoint (round i, round j, 0, 1))
            [ (r * cos (k * deltaTheta) + 0, r * sin (k * deltaTheta) + 0)
            | k <- [0 .. numTheta]
            ]) -- L.++
                -- (L.map
                --    (\(i, j) -> R2S1RPPoint (round i, round j, 0, 1))
                --    [ ((r + 30) * cos (k * deltaTheta), (r + 30) * sin (k * deltaTheta))
                --    | k <- [0 .. numTheta - 1]
                --    ])
               -- L.++
                -- (L.map
                --    (\(i, j) -> R2S1RPPoint (round i, round j, 0, 1))
                --    [ ( (r + 15 * 2) * cos (k * deltaTheta)
                --      , (r + 15 * 2) * sin (k * deltaTheta))
                --    | k <- [0 .. numTheta - 1]
                --    ])
               -- L.++ [R2S1RPPoint (3, 5, 0, 1)]
         )
               -- ([R2S1RPPoint (i, i, 0, 1) | i <- [a',a' + b' .. c']] L.++
               --  [R2S1RPPoint (i, -i, 0, 1) | i <- [-a',-(a' + b') .. -c']] L.++
               --  [R2S1RPPoint (i, i, 0, 1) | i <- [-a',-(a' + b') .. -c']] L.++
               --  [R2S1RPPoint (i, -i, 0, 1) | i <- [a',a' + b' .. c']] L.++
               --  [R2S1RPPoint (i, 0, 0, 1) | i <- [a,a + b .. c]] L.++
               --  [R2S1RPPoint (0, i, 0, 1) | i <- [-a,-(a + b) .. -c]] L.++
               --  [R2S1RPPoint (i, 0, 0, 1) | i <- [-a,-(a + b) .. -c]] L.++
               --  [R2S1RPPoint (0, i, 0, 1) | i <- [a,a + b .. c]] L.++
               --  [R2S1RPPoint (3, 5, 0, 1)])
                -- [ R2S1RPPoint (n, 0, 0, 1)
                -- , R2S1RPPoint (0, n, 0, 1)
                -- , R2S1RPPoint (-n, 0, 0, 1)
                -- , R2S1RPPoint (0, -n, 0, 1)
                -- , R2S1RPPoint (m, m, 0, 1)
                -- , R2S1RPPoint (-m, m, 0, 1)
                -- , R2S1RPPoint (m, -m, 0, 1)
                -- , R2S1RPPoint (-m, -m, 0, 1)
                -- ]
  let bias = computeBiasR2T0 numPoint numPoint theta0Freqs xs
      eigenVec =
        computeInitialEigenVectorR2T0
          numPoint
          numPoint
          theta0Freqs
          thetaFreqs
          xs
  powerMethod1
    plan
    folderPath
    numPoint
    numPoint
    numOrientation
    thetaFreqs
    theta0Freqs
    arrR2Z1T0
    numIteration
    writeSourceFlag
    ""
    0.5
    bias
    eigenVec


module STCLocalEigenVector where

import           Control.Monad             as M
import           Control.Monad.Parallel             as MP
import           Data.Array.Repa           as R
import           Data.Binary               (decodeFile)
import           Data.Complex
import           Data.List                 as L
import           DFT.Plan
import           FokkerPlanck.DomainChange
import           FokkerPlanck.MonteCarlo
import           FokkerPlanck.Pinwheel
import           Image.IO
import           STC
import           System.Directory
import           System.Environment
import           System.FilePath
import           Text.Printf
import           Types
import           Utils.Array
import           Utils.Parallel

{-# INLINE timeReversal #-}
timeReversal :: (R.Source r e) => Array r DIM4 e -> Array D DIM4 e
timeReversal arr =
  let (Z :. nf :. _ :. _ :. _) = extent arr
      n = div nf 2
   in R.backpermute
        (extent arr)
        (\(Z :. k :. l :. i :. j) ->
           let x = nf - n
            in if k >= x
                 then (Z :. k - x :. l  :. i :. j)
                 else (Z :. k + n :. l :. i :. j))
        arr

main = do
  args@(numPointStr:numOrientationStr:numScaleStr:thetaSigmaStr:scaleSigmaStr:maxScaleStr:taoStr:numTrailStr:maxTrailStr:thetaFreqsStr:scaleFreqsStr:histFileName:cutoffRadiusStr:initDistStr:useFFTWWisdomFlagStr:fftwWisdomFileName:numThreadStr:_) <-
    getArgs
  print args
  let numPoint = read numPointStr :: Int
      numOrientation = read numOrientationStr :: Int
      numScale = read numScaleStr :: Int
      thetaSigma = read thetaSigmaStr :: Double
      scaleSigma = read scaleSigmaStr :: Double
      maxScale = read maxScaleStr :: Double
      tao = read taoStr :: Double
      numTrail = read numTrailStr :: Int
      maxTrail = read maxTrailStr :: Int
      thetaFreq = read thetaFreqsStr :: Double
      thetaFreqs = [-thetaFreq .. thetaFreq]
      scaleFreq = read scaleFreqsStr :: Double
      scaleFreqs = [-scaleFreq .. scaleFreq]
      cutoffRadius = read cutoffRadiusStr :: Int
      initDist = read initDistStr :: [R2S1RPPoint]
      useFFTWWisdomFlag = read useFFTWWisdomFlagStr :: Bool
      numThread = read numThreadStr :: Int
      folderPath = "output/test/STCLocalEigenVector"
      histFilePath = folderPath </> histFileName
      fftwWisdomFilePath = folderPath </> fftwWisdomFileName
      sourceDist = L.take 1 initDist
      sinkDist = L.drop 1 initDist
  createDirectoryIfMissing True folderPath
  flag <- doesFileExist histFilePath
  radialArr <-
    if flag
      then R.map magnitude . getNormalizedHistogramArr <$>
           decodeFile histFilePath
      else do
        putStrLn "Couldn't find a Green's function data. Start simulation..."
        solveMonteCarloR2Z2T0S0Radial
          numThread
          numTrail
          maxTrail
          numPoint
          numPoint
          thetaSigma
          scaleSigma
          maxScale
          tao
          thetaFreqs
          thetaFreqs
          scaleFreqs
          scaleFreqs
          histFilePath
          (emptyHistogram
             [ (round . sqrt . fromIntegral $ 2 * (div numPoint 2) ^ 2)
             , L.length scaleFreqs
             , L.length thetaFreqs
             , L.length scaleFreqs
             , L.length thetaFreqs
             ]
             0)
  let localEigenVector =
        computeLocalEigenVector
          (ParallelParams numThread 0)
          (pinwheelHollow 2)
          (cutoff cutoffRadius radialArr)
          numPoint
          numPoint
          maxScale
          thetaFreqs
          scaleFreqs
      localEigenVectorSink =
        computeLocalEigenVectorSink
          (ParallelParams numThread 0)
          (pinwheelHollow 2)
          (cutoff cutoffRadius radialArr)
          numPoint
          numPoint
          maxScale
          thetaFreqs
          scaleFreqs
  plan <-
    make4DPlan emptyPlan useFFTWWisdomFlag fftwWisdomFilePath localEigenVector
  sourceDistArr <-
    computeInitialDistributionR2T0S0
      plan
      numPoint
      numPoint
      thetaFreqs
      scaleFreqs
      maxScale
      sourceDist
  sinkDistArr <-
    computeInitialDistributionR2T0S0
      plan
      numPoint
      numPoint
      thetaFreqs
      scaleFreqs
      maxScale
      sinkDist
  filterF <- dft4D plan . computeS . makeFilter4D $ localEigenVector
  filterSinkF <- dft4D plan . computeS . makeFilter4D $ localEigenVector
  -- Source
  sourceArr <- convolve4D plan filterF sourceDistArr
  sourceR2 <-
    R.sumP .
    R.sumS .
    rotate4D .
    rotate4D . r2z2Tor2s1rp numOrientation thetaFreqs numScale scaleFreqs $
    sourceArr
  plotImageRepaComplex (folderPath </> "Source.png") .
    ImageRepa 8 . computeS . extend (Z :. (1 :: Int) :. All :. All) $
    sourceR2
  -- Sink
  sinkArr <- convolve4D plan filterSinkF sinkDistArr
  -- let sinkArr =
  --       R.traverse2
  --         sinkArr'
  --         (fromListUnboxed (Z :. (L.length thetaFreqs)) thetaFreqs)
  --         const $ \f fFreq idx@(Z :. k :. l :. i :. j) ->
  --         f idx * ((exp (0 :+ (-1) * (fFreq (Z :. k)) * pi)))
  sinkR2 <-
    R.sumP .
    R.sumS .
    rotate4D .
    rotate4D . r2z2Tor2s1rp numOrientation thetaFreqs numScale scaleFreqs $
    sinkArr
  plotImageRepaComplex (folderPath </> "Sink.png") .
    ImageRepa 8 . computeS . extend (Z :. (1 :: Int) :. All :. All) $
    sinkR2
  -- Completion
  -- plotImageRepa (folderPath </> "Completion.png") .
  --   ImageRepa 8 .
  --   computeS .
  --   extend (Z :. (1 :: Int) :. All :. All) .
  --   R.sumS . R.sumS . rotate4D . rotate4D . R.map magnitude $
  --   (timeReversal $  r2z2Tor2s1rp numOrientation thetaFreqs numScale scaleFreqs sinkArr) *^
  --   (r2z2Tor2s1rp numOrientation thetaFreqs numScale scaleFreqs sourceArr)
  -- completionFieldR2 plan folderPath "" sourceR2 sinkR2
  completionFieldR2Z2
    plan
    folderPath
    ""
    numOrientation
    thetaFreqs
    numScale
    scaleFreqs
    sourceArr
    sinkArr
    -- (computeS $
    --  R.traverse2
    --    sinkArr
    --    (fromListUnboxed (Z :. (L.length thetaFreqs)) thetaFreqs)
    --    const $ \f fFreq idx@(Z :. k :. l :. i :. j) ->
    --    f idx * ((exp (0 :+ (-1) * (fFreq (Z :. k)) * pi))))
    -- (computeS . R.map (\x -> let (m,p) = polar x
    --                          in mkPolar m (p + pi)) $ sinkArr)
    -- (computeS $ flip4D sinkArr)
    -- (computeS $ timeReverse4D thetaFreqs sinkArr)
  -- MP.mapM_
  --   (\(tf, sf) ->
  --      plotImageRepaComplex (folderPath </> printf "Source_%d_%d.png" tf sf) .
  --      ImageRepa 8 .
  --      computeS .
  --      extend (Z :. (1 :: Int) :. All :. All) .
  --      R.slice
  --        (r2z2Tor2s1rp numOrientation thetaFreqs numScale scaleFreqs sourceArr) $
  --      (Z :. tf :. sf :. All :. All))
  --   [ (tf, sf)
  --   | tf <- [0 .. numOrientation - 1]
  --   , sf <- [0 .. numScale - 1]
  --   ]
    -- [ (tf, sf)
    -- | tf <- [0 .. L.length thetaFreqs - 1]
    -- , sf <- [0 .. L.length scaleFreqs - 1]
    -- ]


{-# LANGUAGE AllowAmbiguousTypes                      #-}
{-# LANGUAGE ApplicativeDo                            #-}
{-# LANGUAGE DeriveDataTypeable                       #-}
{-# LANGUAGE DeriveGeneric                            #-}
{-# LANGUAGE FlexibleContexts                         #-}
{-# LANGUAGE FlexibleInstances                        #-}
{-# LANGUAGE GADTs                                    #-}
{-# LANGUAGE KindSignatures                           #-}
{-# LANGUAGE MultiParamTypeClasses                    #-}
{-# LANGUAGE PatternSynonyms                          #-}
{-# LANGUAGE RankNTypes                               #-}
{-# LANGUAGE RecordWildCards                          #-}
{-# LANGUAGE ScopedTypeVariables                      #-}
{-# LANGUAGE TypeApplications                         #-}
{-# LANGUAGE TypeFamilies                             #-}
{-# LANGUAGE TypeInType                               #-}
{-# LANGUAGE TypeOperators                            #-}
{-# LANGUAGE UndecidableInstances                     #-}
{-# OPTIONS_GHC -fplugin GHC.TypeLits.KnownNat.Solver #-}
{-# OPTIONS_GHC -fplugin GHC.TypeLits.Normalise       #-}

module Backprop.Learn.Model.Regression (
    LinReg(..)
  , LogReg, pattern LogReg
  , LRp(..), lrBeta, lrAlpha, runLRp
  , expandInput, expandOutput, reshapeInput, reshapeOutput
  , ARIMA(..), ARIMAp(..), ARIMAs(..)
  , ARIMAUnroll, arimaUnroll
  , ARIMAUnrollFinal, arimaUnrollFinal
  , AR, MA, ARMA
  ) where

import           Backprop.Learn.Initialize
import           Backprop.Learn.Model.Class
import           Backprop.Learn.Model.Combinator
import           Backprop.Learn.Model.Function
import           Backprop.Learn.Model.Parameter
import           Backprop.Learn.Model.State
import           Control.DeepSeq
import           Control.Monad.Primitive
import           Data.Finite
import           Data.Kind
import           Data.List
import           Data.Maybe
import           Data.Proxy
import           Data.Type.Equality
import           Data.Typeable
import           GHC.Generics                          (Generic)
import           GHC.TypeLits.Compare
import           GHC.TypeLits.Extra
import           GHC.TypeNats
import           Lens.Micro
import           Numeric.Backprop
import           Numeric.LinearAlgebra.Static.Backprop
import           Numeric.LinearAlgebra.Static.Vector
import           Numeric.OneLiner
import           Numeric.Opto.Ref
import           Numeric.Opto.Update hiding            ((<.>))
import           Statistics.Distribution
import           Unsafe.Coerce
import qualified Data.Binary                           as Bi
import qualified Data.Type.Tuple                       as T
import qualified Data.Vector.Generic.Sized             as SVG
import qualified Data.Vector.Sized                     as SV
import qualified Data.Vector.Storable.Sized            as SVS
import qualified Numeric.LinearAlgebra                 as HU
import qualified Numeric.LinearAlgebra.Static          as H
import qualified System.Random.MWC                     as MWC

-- | Multivariate linear regression, from an i-vector to an o-vector.
data LinReg (i :: Nat) (o :: Nat) = LinReg
  deriving Typeable

-- | Mutivariate Logistic regression, from an i-vector to an o-vector.
--
-- Essentially a linear regression postcomposed with the logistic
-- function.
type LogReg i o = RMap (R o) (R o) (LinReg i o)

-- | Constructor for a 'LogReg'
pattern LogReg :: LogReg i o
pattern LogReg <- RM _ LinReg
  where
    LogReg = RM logistic LinReg
{-# COMPLETE LogReg #-}

-- | Linear Regression parameter
data LRp i o = LRp
    { _lrAlpha :: !(R o)
    , _lrBeta  :: !(L o i)
    }
  deriving (Generic, Typeable, Show)

instance NFData (LRp i o)
instance (KnownNat i, KnownNat o) => Initialize (LRp i o)
instance (KnownNat i, KnownNat o) => Additive (LRp i o)
instance (KnownNat i, KnownNat o) => Scaling Double (LRp i o)
instance (KnownNat i, KnownNat o) => Metric Double (LRp i o)
instance (KnownNat i, KnownNat o, Ref m (LRp i o) v) => AdditiveInPlace m v (LRp i o)
instance (KnownNat i, KnownNat o, Ref m (LRp i o) v) => ScalingInPlace m v Double (LRp i o)
instance (KnownNat i, KnownNat o) => Bi.Binary (LRp i o)
instance (KnownNat i, KnownNat o) => Backprop (LRp i o)

lrBeta :: Lens (LRp i o) (LRp i' o) (L o i) (L o i')
lrBeta f lrp = (\w -> lrp { _lrBeta = w }) <$> f (_lrBeta lrp)

lrAlpha :: Lens' (LRp i o) (R o)
lrAlpha f lrp = (\b -> lrp { _lrAlpha = b }) <$> f (_lrAlpha lrp)

runLRp
    :: (KnownNat i, KnownNat o, Reifies s W)
    => BVar s (LRp i o)
    -> BVar s (R i)
    -> BVar s (R o)
runLRp lrp x = (lrp ^^. lrBeta) #> x + (lrp ^^. lrAlpha)

instance (KnownNat i, KnownNat o) => Num (LRp i o) where
    (+)         = gPlus
    (-)         = gMinus
    (*)         = gTimes
    negate      = gNegate
    abs         = gAbs
    signum      = gSignum
    fromInteger = gFromInteger

instance (KnownNat i, KnownNat o) => Fractional (LRp i o) where
    (/)          = gDivide
    recip        = gRecip
    fromRational = gFromRational

instance (KnownNat i, KnownNat o) => Floating (LRp i o) where
    pi    = gPi
    sqrt  = gSqrt
    exp   = gExp
    log   = gLog
    sin   = gSin
    cos   = gCos
    tan   = gTan
    asin  = gAsin
    acos  = gAcos
    atan  = gAtan
    sinh  = gSinh
    cosh  = gCosh
    asinh = gAsinh
    acosh = gAcosh
    atanh = gAtanh

instance (KnownNat i, KnownNat o) => Learn (R i) (R o) (LinReg i o) where
    type LParamMaybe (LinReg i o) = 'Just (LRp i o)

    runLearn _ (J_ p) = stateless (runLRp p)

-- | Adjust an 'LRp' to take extra inputs, initialized randomly.
--
-- Initial contributions to each output is randomized.
expandInput
    :: (PrimMonad m, ContGen d, KnownNat i, KnownNat j, KnownNat o)
    => LRp i o
    -> d
    -> MWC.Gen (PrimState m)
    -> m (LRp (i + j) o)
expandInput LRp{..} d g = LRp _lrAlpha . (_lrBeta H.|||) <$> initialize d g

-- | Adjust an 'LRp' to return extra ouputs, initialized randomly
expandOutput
    :: (PrimMonad m, ContGen d, KnownNat i, KnownNat o, KnownNat p)
    => LRp i o
    -> d
    -> MWC.Gen (PrimState m)
    -> m (LRp i (o + p))
expandOutput LRp{..} d g = do
    newAlpha <- initialize d g
    newBeta  <- initialize d g
    pure (LRp (_lrAlpha H.# newAlpha) (_lrBeta H.=== newBeta))

-- | Premute (or remove) inputs
--
-- Removed inputs will simply have their contributions removed from each
-- output.
reshapeInput
    :: KnownNat i
    => SV.Vector i' (Finite i)
    -> LRp i o
    -> LRp i' o
reshapeInput is p = p { _lrBeta = colsL . fmap (β `SV.index`) $ is }
  where
    β = lCols (_lrBeta p)

-- | Premute (or remove) outputs
reshapeOutput
    :: KnownNat o
    => SV.Vector o' (Finite o)
    -> LRp i o
    -> LRp i o'
reshapeOutput is LRp{..} =
    LRp { _lrAlpha = vecR . SVG.convert . fmap (α `SVS.index`) $ is
        , _lrBeta  = rowsL . fmap (β `SV.index`) $ is
        }
  where
    α = rVec _lrAlpha
    β = lRows _lrBeta

-- | Auto-regressive integrated moving average model.
--
-- ARIMA(p,d,q) is an ARIMA model with p autoregressive history terms on
-- the d-th order differenced history and q error (innovation) history
-- terms.
--
-- It is a @'Learn' Double Double@ instance, and is meant to predict the
-- "next step" of a time sequence.
--
-- In this state, it is a runnable stateful model.  To train, use with
-- 'ARIMAUnroll' to unroll and fix the initial state.
data ARIMA :: Nat -> Nat -> Nat -> Type where
    ARIMA :: { _arimaGenYHist :: forall m. PrimMonad m => MWC.Gen (PrimState m) -> m Double
             , _arimaGenEHist :: forall m. PrimMonad m => MWC.Gen (PrimState m) -> m Double
             }
          -> ARIMA p d q

-- | The "unrolled" and "destated" 'ARIMA' model, which takes a vector of
-- sequential inputs and outputs a vector of the model's expected next
-- steps.
--
-- Useful for actually training 'ARIMA' using gradient descent.
--
-- This /fixes/ the initial error history to be zero (or a fixed stochastic
-- sample), and treats the initial output history to be a /learned
-- parameter/.
--
-- @
-- instance 'Learn' ('SV.Vector' n 'Double') (SV.Vector n Double) ('ARIMAUnroll' p q) where
--     -- | Initial state is a parameter, but initial error history is fixed
--     type 'LParamMaybe' (ARIMAUnroll p q) = 'Just (T2 (ARIMAp p q) (ARIMAs p q))
--     type 'LStateMaybe' (ARIMAUnroll p q) = 'Nothing
-- @
type ARIMAUnroll p d q = DeParamAt (T.T2 (ARIMAp p q) (ARIMAs p d q))
                                   (T.T2 (ARIMAp p q) (T.T2 Double (R (p + d))))
                                   (R q)
                                   (UnrollTrainState (Max (p + d) q) (ARIMA p d q))

-- | 'ARIMAUnroll', but only looking at the final output after running the
-- model on all inputs.
--
-- @
-- instance 'Learn' ('SV.Vector' n 'Double') Double ('ARIMAUnrollFinal' p q) where
--     -- | Initial state is a parameter, but initial error history is fixed
--     type 'LParamMaybe' (ARIMAUnrollFinal p q) = 'Just (T2 (ARIMAp p q) (ARIMAs p q))
--     type 'LStateMaybe' (ARIMAUnrollFinal p q) = 'Nothing
-- @
type ARIMAUnrollFinal p d q = DeParamAt (T.T2 (ARIMAp p q) (ARIMAs p d q))
                                        (T.T2 (ARIMAp p q) (T.T2 Double (R (p + d))))
                                        (R q)
                                        (UnrollFinalTrainState (Max (p + d) q) (ARIMA p d q))

splitHist
    :: T.T2 (ARIMAp p q) (ARIMAs p d q)
    -> (T.T2 (ARIMAp p q) (T.T2 Double (R (p + d))), R q)
splitHist (T.T2 p ARIMAs{..}) = (T.T2 p (T.T2 _arimaYPred _arimaYHist), _arimaEHist)

joinHist
    :: T.T2 (ARIMAp p q) (T.T2 Double (R (p + d)))
    -> R q
    -> T.T2 (ARIMAp p q) (ARIMAs p d q)
joinHist (T.T2 p (T.T2 _arimaYPred _arimaYHist)) _arimaEHist = T.T2 p ARIMAs{..}

-- | Constructor for 'ARIMAUnroll'
arimaUnroll
    :: KnownNat q
    => ARIMA p d q
    -> ARIMAUnroll p d q
arimaUnroll a@ARIMA{..} = DPA
    { _dpaSplit      = splitHist
    , _dpaJoin       = joinHist
    , _dpaParam      = 0
    , _dpaParamStoch = fmap vecR . SVS.replicateM . _arimaGenEHist
    , _dpaLearn      = UnrollTrainState a
    }

-- | Constructor for 'ARIMAUnrollFinal'
arimaUnrollFinal
    :: KnownNat q
    => ARIMA p d q
    -> ARIMAUnrollFinal p d q
arimaUnrollFinal a@ARIMA{..} = DPA
    { _dpaSplit      = splitHist
    , _dpaJoin       = joinHist
    , _dpaParam      = 0
    , _dpaParamStoch = fmap vecR . SVS.replicateM . _arimaGenEHist
    , _dpaLearn      = UnrollFinalTrainState a
    }

-- | 'ARIMA' parmaeters
data ARIMAp :: Nat -> Nat -> Type where
    ARIMAp :: { _arimaPhi      :: !(R p)
              , _arimaTheta    :: !(R q)
              , _arimaConstant :: !Double
              }
           -> ARIMAp p q
  deriving (Generic, Show)

-- | 'ARIMA' state
data ARIMAs :: Nat -> Nat -> Nat -> Type where
    ARIMAs :: { _arimaYPred :: !Double
              , _arimaYHist :: !(R (p + d))
              , _arimaEHist :: !(R q)
              }
          -> ARIMAs p d q
  deriving (Generic, Show)

instance (KnownNat p, KnownNat d, KnownNat q) => Learn Double Double (ARIMA p d q) where
    type LParamMaybe (ARIMA p d q) = 'Just (ARIMAp p q)
    type LStateMaybe (ARIMA p d q) = 'Just (ARIMAs p d q)

    runLearn ARIMA{..} (J_ p) x (J_ s) = (y, J_ s')
      where
        d :: L p (p + d)
        d  = difference
        e  = x - (s ^^. arimaYPred)
        y  = (p ^^. arimaConstant)
           + (p ^^. arimaPhi  ) <.> (constVar d #> (s ^^. arimaYHist))
           + (p ^^. arimaTheta) <.> (s ^^. arimaEHist)
        yHist' = case Proxy @1 %<=? Proxy @(p + d) of
          LE Refl -> single y # constVar dropLast #> (s ^^. arimaYHist)
          NLE _ _ -> 0
        eHist' = case Proxy @1 %<=? Proxy @q of
          LE Refl -> single e # constVar dropLast #> (s ^^. arimaEHist)
          NLE _ _ -> 0
        s' = isoVar3 ARIMAs (\(ARIMAs pr yh eh) -> (pr,yh,eh))
                y
                yHist'
                eHist'

monosquare :: forall n. (n <=? (n ^ 2)) :~: 'True
monosquare = unsafeCoerce Refl

dropLast :: forall n. (KnownNat n, 1 <= n) => L (n - 1) n
dropLast = case monosquare @n of
    Refl -> vecL . SVS.generate $ \ij ->
      let i :: Finite n
          j :: Finite (n - 1)
          (i, j) = separateProduct ij
      in  if fromIntegral @_ @Int i == fromIntegral j
            then 1
            else 0

single :: Reifies s W => BVar s Double -> BVar s (R 1)
single = konst

difference'
    :: Int                  -- ^ initial
    -> Int                  -- ^ target
    -> HU.Matrix Double     -- ^ target x initial
difference' n m = foldl' go (HU.ident m) [m + 1 .. n]
  where
    go x k = x HU.<> d k
    d k = HU.build (k-1, k) $ \i j ->
        case round @_ @Int (j - i) of
          0 -> 1
          1 -> -1
          _ -> 0

difference :: forall n m. (KnownNat n, KnownNat m) => L n (n + m)
difference = fromJust . H.create $ difference' (n + m) n
  where
    n = fromIntegral $ natVal (Proxy @n)
    m = fromIntegral $ natVal (Proxy @m)

instance NFData (ARIMAp p q)
instance NFData (ARIMAs p d q)

instance Num (ARIMAp p q) where
    (+)         = gPlus
    (-)         = gMinus
    (*)         = gTimes
    negate      = gNegate
    abs         = gAbs
    signum      = gSignum
    fromInteger = gFromInteger

instance Num (ARIMAs p d q) where
    (+)         = gPlus
    (-)         = gMinus
    (*)         = gTimes
    negate      = gNegate
    abs         = gAbs
    signum      = gSignum
    fromInteger = gFromInteger

instance Fractional (ARIMAp p q) where
    (/)          = gDivide
    recip        = gRecip
    fromRational = gFromRational

instance Fractional (ARIMAs p d q) where
    (/)          = gDivide
    recip        = gRecip
    fromRational = gFromRational

instance Floating (ARIMAp p q) where
    pi    = gPi
    sqrt  = gSqrt
    exp   = gExp
    log   = gLog
    sin   = gSin
    cos   = gCos
    asin  = gAsin
    acos  = gAcos
    atan  = gAtan
    sinh  = gSinh
    cosh  = gCosh
    asinh = gAsinh
    acosh = gAcosh
    atanh = gAtanh

instance Floating (ARIMAs p d q) where
    pi    = gPi
    sqrt  = gSqrt
    exp   = gExp
    log   = gLog
    sin   = gSin
    cos   = gCos
    asin  = gAsin
    acos  = gAcos
    atan  = gAtan
    sinh  = gSinh
    cosh  = gCosh
    asinh = gAsinh
    acosh = gAcosh
    atanh = gAtanh

instance Additive (ARIMAp p q) where
    (.+.)   = gAdd
    addZero = gAddZero
instance Additive (ARIMAs p d q) where
    (.+.)   = gAdd
    addZero = gAddZero

instance (KnownNat p, KnownNat q) => Scaling Double (ARIMAp p q)
instance (KnownNat p, KnownNat d, KnownNat q) => Scaling Double (ARIMAs p d q)

instance (KnownNat p, KnownNat q) => Metric  Double (ARIMAp p q)
instance (KnownNat p, KnownNat d, KnownNat q) => Metric  Double (ARIMAs p d q)

instance (KnownNat p, KnownNat q, Ref m (ARIMAp p q) v) => AdditiveInPlace m v (ARIMAp p q)
instance (KnownNat p, KnownNat d, KnownNat q, Ref m (ARIMAs p d q) v) => AdditiveInPlace m v (ARIMAs p d q)

instance (KnownNat p, KnownNat q, Ref m (ARIMAp p q) v) => ScalingInPlace m v Double (ARIMAp p q)
instance (KnownNat p, KnownNat d, KnownNat q, Ref m (ARIMAs p d q) v) => ScalingInPlace m v Double (ARIMAs p d q)

instance (KnownNat p, KnownNat q) => Initialize (ARIMAp p q)
instance (KnownNat p, KnownNat d, KnownNat q) => Initialize (ARIMAs p d q)

instance (KnownNat p, KnownNat q) => Bi.Binary (ARIMAp p q)
instance (KnownNat p, KnownNat d, KnownNat q) => Bi.Binary (ARIMAs p d q)

instance (KnownNat p, KnownNat q) => Backprop (ARIMAp p q)
instance (KnownNat p, KnownNat d, KnownNat q) => Backprop (ARIMAs p d q)

arimaPhi :: Lens (ARIMAp p q) (ARIMAp p' q) (R p) (R p')
arimaPhi f a = (\x' -> a { _arimaPhi = x' } ) <$> f (_arimaPhi a)

arimaTheta :: Lens (ARIMAp p q) (ARIMAp p q') (R q) (R q')
arimaTheta f a = (\x' -> a { _arimaTheta = x' } ) <$> f (_arimaTheta a)

arimaConstant :: Lens' (ARIMAp p q) Double
arimaConstant f a = (\x' -> a { _arimaConstant = x' } ) <$> f (_arimaConstant a)

arimaYPred :: Lens' (ARIMAs p d q) Double
arimaYPred f a = (\x' -> a { _arimaYPred = x' } ) <$> f (_arimaYPred a)

arimaYHist :: Lens' (ARIMAs p d q) (R (p + d))
arimaYHist f a = (\x' -> a { _arimaYHist = x' } ) <$> f (_arimaYHist a)

arimaEHist :: Lens (ARIMAs p d q) (ARIMAs p d q') (R q) (R q')
arimaEHist f a = (\x' -> a { _arimaEHist = x' } ) <$> f (_arimaEHist a)

-- | Autoregressive model
type AR p = ARIMA p 0 0

-- | Moving average model
type MA = ARIMA 0 0

-- | Autoregressive Moving average model
type ARMA p = ARIMA p 0


{-# LANGUAGE BangPatterns,
             ScopedTypeVariables,
             RecordWildCards,
             FlexibleContexts,
             TypeFamilies,
             DeriveGeneric #-}
module Main where

import Data.IDX ( decodeIDXFile, idxIntContent, IDXData )
{------------------------}
import Numeric.LinearAlgebra as NL ( Vector, fromList )
{------------------------}
import Data.Vector.Unboxed as DTU ( toList )

import Prelude hiding (readFile)
import Neuro
    ( Samples, Sample, createNetwork, tanh', saveNetwork, trainNTimes )

main :: IO ()
main = do
  n <- createNetwork 784 [64] 10
  samples <- importTrain
  {-----------------------------}
  let n' = trainNTimes 50 5.0 tanh tanh' n samples
  saveNetwork "smartNet5.nn" n'

importTrain :: IO (Samples Double)
importTrain = do
    Just idxTrain  <- decodeIDXFile "train-images.idx3-ubyte" -- image
    Just idxResult <- decodeIDXFile "train-labels.idx1-ubyte" -- result
    return $ samples idxTrain idxResult

samples :: IDXData -> IDXData -> [Sample Double]
samples idxTrain idxResult = Prelude.zip (image idxTrain) (result idxResult) :: [Sample Double]

image  = matrix2x2 . _0_1to0_255 . DTU.toList . idxIntContent
result = unitar . DTU.toList . idxIntContent

unitar :: [Int] -> [NL.Vector Double]
unitar []     = []
unitar (0:xs) = NL.fromList [1,-1,-1,-1,-1,-1,-1,-1,-1,-1] : unitar xs
unitar (1:xs) = NL.fromList [-1,1,-1,-1,-1,-1,-1,-1,-1,-1] : unitar xs
unitar (2:xs) = NL.fromList [-1,-1,1,-1,-1,-1,-1,-1,-1,-1] : unitar xs
unitar (3:xs) = NL.fromList [-1,-1,-1,1,-1,-1,-1,-1,-1,-1] : unitar xs
unitar (4:xs) = NL.fromList [-1,-1,-1,-1,1,-1,-1,-1,-1,-1] : unitar xs
unitar (5:xs) = NL.fromList [-1,-1,-1,-1,-1,1,-1,-1,-1,-1] : unitar xs
unitar (6:xs) = NL.fromList [-1,-1,-1,-1,-1,-1,1,-1,-1,-1] : unitar xs
unitar (7:xs) = NL.fromList [-1,-1,-1,-1,-1,-1,-1,1,-1,-1] : unitar xs
unitar (8:xs) = NL.fromList [-1,-1,-1,-1,-1,-1,-1,-1,1,-1] : unitar xs
unitar (9:xs) = NL.fromList [-1,-1,-1,-1,-1,-1,-1,-1,-1,1] : unitar xs

_0_1to0_255 :: [Int] -> [Double]
_0_1to0_255 [] = []
_0_1to0_255 (x:xs) = (fromIntegral x / 255) : _0_1to0_255 xs

matrix2x2 :: [Double] -> [NL.Vector Double]
matrix2x2 [] = []
matrix2x2 xs = NL.fromList (Prelude.take 784 xs) : matrix2x2 (Prelude.drop 784 xs)


import System.Exit

import Quipper
import Quipper.Internal

import Quantum.Synthesis.Matrix
import Quantum.Synthesis.Ring
import Quipper.Libraries.Synthesis

import Quipper.Libraries.Decompose.GateBase
import Quipper.Libraries.Decompose

import Data.Complex
import Data.Ratio
import Data.Tuple

import Quipper.Utils.RandomSource
import System.Random

-- declare sample_oracle's data type
data Oracle = Oracle {
   qubit_num :: Int,
   function :: ([Qubit], Qubit) -> Circ ([Qubit], Qubit)
}

-- declare circuit function
circuit_function :: Oracle -> Circ ([Bit], Bit)
circuit_function oracle = do
     -- initialize string of qubits
     top_qubits <- qinit (replicate (qubit_num oracle) False)
     bottom_qubit <- qinit False
     label (top_qubits, bottom_qubit) ("|0>","|0>")

     --prepare the initial states
     --mapUnary hadamard top_qubits
     --mapUnary hadamard bottom_qubit

     comment "before oracle"
     -- call oracle
     function oracle (top_qubits, bottom_qubit)
     comment "after oracle"

     -- measure qubits
     (top_qubits, bottom_qubit) <- measure (top_qubits, bottom_qubit)
     -- discard unnecessary output and return result
     return (top_qubits, bottom_qubit)

-- * From a matrix
type Sixteen = Ten_and Six

mymatrix :: Matrix Sixteen Sixteen (Integer)
mymatrix = matrix [[1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], -- ([],0)
                   [0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], -- ([],1)
                   [0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], -- ([0],0)
                   [0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], -- ([0],1)
                   [0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], -- ([0,0],0)
                   [0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], -- ([0,0],1)
                   [0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0], -- ([1,0],0)
                   [0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0], -- ([1,0],1)
                   [0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0], -- ([1],0)
                   [0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0], -- ([1],1)
                   [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0], -- ([0,1],0)
                   [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0], -- ([0,1],1)
                   [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0], -- ([1,1],0)
                   [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0], -- ([1,1],1)
                   [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0], -- ([0,0,0],0)
                   [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1]] -- ([0,0,0],1)

synthesized = exact_synthesis mymatrix

circuit ::([Qubit], Qubit)-> Circ ([Qubit], Qubit)
circuit ([a, b, c], d) = do
  synthesized [a,b,c,d]
  return ([a,b,c],d)

-- Standard decompose in X, Y , Z, H , S, S*, T, T* and CNOT
-- needs : precision and randomSource
rand = RandomSource(fst(split(mkStdGen 10)))
prec = 20 * bits
circuit_decompose = decompose_generic (Standard prec rand) circuit

--
original_pdf :: IO ()
original_pdf = do
  print_generic Preview (circuit_function my_oracle)
    where
    -- declare empty_oracle's data type
    my_oracle :: Oracle
    my_oracle = Oracle {
    -- set the length of qubit string
        qubit_num = 3,
        function = circuit
    }

--
original_ASCII :: IO ()
original_ASCII = do
  print_generic ASCII (circuit_function my_oracle)
    where
     -- declare empty_oracle's data type
     my_oracle :: Oracle
     my_oracle = Oracle {
      -- set the length of qubit string
       qubit_num = 3,
       function = circuit
     }

decompose_pdf :: IO ()
decompose_pdf = do
  print_generic Preview (circuit_function my_oracle)
    where
   -- declare empty_oracle's data type
   my_oracle :: Oracle
   my_oracle = Oracle {
     -- set the length of qubit string
      qubit_num = 3,
      function = circuit_decompose
   }
 -- print mymatrix

decompose_ASCII :: IO ()
decompose_ASCII = do
  print_generic ASCII (circuit_function my_oracle)
    where
   -- declare empty_oracle's data type
   my_oracle :: Oracle
   my_oracle = Oracle {
     -- set the length of qubit string
      qubit_num = 3,
      function = circuit_decompose
   }

main_menu = do
  putStrLn "\n \n \n \n \n \n \n \n choose an option:"
  putStrLn " 1 - PDF original circuit\n"
  putStrLn " 2 - text description of original circuit\n"
  putStrLn " 3 - PDF decomposed circuit\n"
  putStrLn " 4 - text description of decomposed circuit\n"
  putStrLn " 5 - exit\n"
  line <- getLine
  case line of
    "1" -> do original_pdf
              main
    "2" -> do original_ASCII
              main
    "3" -> do decompose_pdf
              main
    "4" -> do decompose_ASCII
              main
    "5" -> do exitSuccess
    _ -> do main

main_qfold_cnot_7x2 = do
  decompose_ASCII

main = do
  --main_menu

  -- or
  -- if you want the circuit description a .txt file
  -- comment the main_menu above
  -- replacing it by:

  main_qfold_cnot_7x2

  -- and run:
  -- $ ./qfold_cnot_7x2_quipper > circuit_cnot_7x2_quipper.txt


-----------------------------------------------------------------------------
-- |
-- Module      :  RL.Util
-- Description :  Misc. Utilities
-- Copyright   :  (c) David Banas, 2018
-- License     :  BSD-style (see the file LICENSE)
--
-- Maintainer  :  David.Banas@target.com
-- Stability   :  experimental
-- Portability :  ?
--
-- Common utilities used throughout the RL package.
-----------------------------------------------------------------------------

{-# OPTIONS_GHC -Wall #-}

{-# LANGUAGE AllowAmbiguousTypes #-}
{-# LANGUAGE TypeApplications #-}

module RL.Util where

import qualified Prelude as P
import Prelude (Show(..))
import Protolude  hiding (show, for)

import GHC.TypeLits

import qualified Data.Vector.Sized   as VS
import Data.Vector.Sized             (Vector)

import Control.Monad.Writer
import Data.Finite
import Data.Finite.Internal
import Data.List                     ((!!))
import Statistics.Distribution       (density)
import Statistics.Distribution.Gamma (gammaDistr)
import Text.Printf

import ConCat.TArr

-- | Convenient abbreviation of: `natVal (Proxy @n)`.
nat :: forall n. KnownNat n => Integer
nat = natVal (Proxy @n)

-- | Convenient abbreviation of: `fromIntegral (nat @n)`.
int :: forall n. KnownNat n => Int
int = fromIntegral (nat @n)

-- | Apply the matrix representation of a two argument function.
--
-- The first argument is assumed to index the rows of the matrix.
appFm
  :: ( HasFin' a,   HasFin' b )
  => Vector (Card a) (Vector (Card b) c)  -- ^ matrix representation of @f(r,c)@
  -> a                                    -- ^ row
  -> b                                    -- ^ column
  -> c
appFm f x y = f `VS.index` toFin x `VS.index` toFin y

-- | Apply the vector representation of a function.
appFv
  :: ( HasFin' x )
  => Vector (Card x) a  -- ^ vector representation of @f(x)@
  -> x                  -- ^ function argument
  -> a
appFv f x = f `VS.index` toFin x

-- | Convert an action-value matrix to a value vector.
--
-- (i.e. - \(Q(s,a) -> V(s)\))
qToV
  :: ( Ord a
     , KnownNat m, KnownNat n, KnownNat k
     , n ~ (k + 1)
     )
  => Vector m (Vector n a)  -- ^ matrix representation of @Q(s,a)@
  -> Vector m a             -- ^ vector representation of @V(s)@
qToV = VS.map VS.maximum

-- | Convert an action-value matrix to a policy vector.
--
-- (i.e. - \(Q(s,a) -> A(s)\))
qToP
  :: ( Ord a, HasFin' act
     , KnownNat m, KnownNat n, KnownNat k
     , n ~ (k + 1), Card act ~ n
     )
  => Vector m (Vector n a)  -- ^ matrix representation of Q(s,a)
  -> Vector m act           -- ^ vector representation of P(s)
qToP = VS.map (unFin . VS.maxIndex)

{----------------------------------------------------------------------
  Misc.
----------------------------------------------------------------------}

-- | To control the formatting of printed floats in output matrices.
newtype Pfloat = Pfloat { unPfloat :: Float}
  deriving (Eq)

instance Show Pfloat where
  show x = printf "%4.1f" (unPfloat x)

-- | To control the formatting of printed doubles in output matrices.
newtype Pdouble = Pdouble { unPdouble :: Double }
  deriving (Eq, Ord)

instance Show Pdouble where
  show x = printf "%4.1f" (unPdouble x)

poisson :: Finite 5 -> Finite 12 -> Float
poisson (Finite lambda) (Finite n') =
  lambda' ^ n * exp (-lambda') / fromIntegral (fact n)
 where lambda' = fromIntegral lambda
       n       = fromIntegral n'

fact :: Int -> Int
fact 0 = 1
fact n = product [1..n]

poissonVals :: VS.Vector 5 (VS.Vector 12 Float)
poissonVals = VS.generate (VS.generate . poisson)

poisson' :: Finite 5 -> Finite 21 -> Float
poisson' n x@(Finite x') =
  if x > 11  -- The Python code enforces this limit. And we're trying
    then 0   -- for an "apples-to-apples" performance comparison.
    else poissonVals `VS.index` n `VS.index` finite x'

-- | Gamma pdf
--
-- Assuming `scale = 1`, `shape` should be: 1 + mean.
gammaPdf :: Double -> Double -> Double -> Double
gammaPdf shape scale = density (gammaDistr shape scale)

-- | Gamma pmf
--
-- Scale assumed to be `1`, so as to match the calling signature of
-- `poisson`.
gamma :: Finite 5 -> Finite 12 -> Double
gamma (Finite expect') (Finite n') =
  0.1 * sum [gammaPdf' (n + x) | x <- [0.1 * m | m <- [-4..5]]]
    where gammaPdf' = gammaPdf (1 + expect) 1
          expect    = fromIntegral expect'
          n         = fromIntegral n'

gammaVals :: VS.Vector 5 (VS.Vector 12 Double)
gammaVals = VS.generate (VS.generate . gamma)

gamma' :: Finite 5 -> Finite 21 -> Double
gamma' n (Finite x') =
  if x' > 11
    then 0
    else gammaVals `VS.index` n `VS.index` finite x'

-- | Monadically search list for first element less than
-- given threshold under the given function, and return the last element
-- if the threshold was never met.
-- Return 'Nothing' if the input list was empty.
withinOnM :: Monad m
          => Double
          -> (a -> m Double)
          -> [a]
          -> m (Maybe a)
withinOnM _   _ [] = return Nothing
withinOnM eps f xs = do
  n <- withinIxM eps $ map f xs
  case n of
    Nothing -> return Nothing
    Just n' -> return $ Just (xs !! n')

-- | Monadically find index of first list element less than or equal to
-- given threshold, or the index of the last element if the threshold
-- was never met.
--
-- A return value of 'Nothing' indicates an empty list was given.
withinIxM :: Monad m
          => Double
          -> [m Double]
          -> m (Maybe Int)
withinIxM _   [] = return Nothing
withinIxM eps xs = withinIxM' 0 xs
 where withinIxM' n []     = return $ Just (n - 1)
       withinIxM' n (y:ys) = do
         y' <- y
         if y' < eps then return (Just n)
                     else withinIxM' (n+1) ys

-- | Return the maximum value of a set, as well as a count of the number
-- of non-zero elements in the set.
--
-- (See documentation for `chooseAndCount` function.)
maxAndNonZero :: (Foldable t, Num a, Ord a) => t a -> Writer [Int] a
maxAndNonZero = chooseAndCount max (/= 0)

-- | Choose a value from the set using the given comparison function,
-- and provide a count of the number of elements in the set meeting the
-- given criteria.
chooseAndCount :: (Foldable t, Num a)
               => (a -> a -> a)  -- ^ choice function
               -> (a -> Bool)    -- ^ counting predicate
               -> t a            -- ^ foldable set of elements to count/compare
               -> Writer [Int] a
chooseAndCount f p xs = do
  let (val, cnt::Int) =
        foldl' ( \ (v, c) x ->
                   ( f v x
                   , if p x
                       then c + 1
                       else c
                   )
               ) (0,0) xs
  tell [cnt]
  return val

for :: [a] -> (a -> b) -> [b]
for = flip map

vsFor :: Vector n a -> (a -> b) -> Vector n b
vsFor = flip VS.map

-- | Mean value of a collection
mean :: (Foldable f, Fractional a) => f a -> a
mean = uncurry (/) . second fromIntegral . foldl' (\ (!s, !n) x -> (s+x, n+1)) (0,0::Integer)

-- | Find the mean square of a list of lists.
arrMeanSqr :: (Functor f, Foldable f, Functor g, Foldable g, Fractional a) => f (g a) -> a
arrMeanSqr = mean . fmap mean . fmap (fmap sqr)

-- | Take the square of a numerical type.
sqr :: Num a => a -> a
sqr x = x * x

-- | Convert any showable type to a string, avoiding the introduction
-- of extra quotation marks when that type is a string to begin with.
toString :: (Show a, Typeable a) => a -> P.String
toString x = fromMaybe (show x) (cast x)

-- | Convert a Bool to a Double.
boolToDouble :: Bool -> Double
boolToDouble True = 1
boolToDouble _    = 0

-- | Take every nth element from a list.
takeEvery :: Int -> [a] -> [a]
takeEvery _ [] = []
takeEvery n xs = P.head xs : (takeEvery n $ drop n xs)


-- Copyright © 2016 Bart Massey
-- [This work is made available under the "MIT License".
-- Please see the file LICENSE in this distribution for
-- license details.]

-- Analysis tool for a Hardware Random Number Generator

import Control.Lens
import Data.Char (toLower)
import Data.Colour
import Data.Colour.Names
import Data.Colour.SRGB
import Data.Default.Class
import Data.List.Split (splitOn)
import Data.Maybe
import Graphics.Rendering.Chart
import Graphics.Rendering.Chart.Backend.Cairo

import Data.Complex
import qualified Data.Vector.Unboxed as V
import Numeric.FFT.Vector.Unnormalized

import Control.Monad
import Data.Bits
import qualified Data.ByteString as B
import Data.List
import System.Console.ParseArgs
import System.Directory
import System.Exit (exitSuccess)
import System.FilePath
import System.Random
import Text.Printf

defaultAnalysisDir :: FilePath
defaultAnalysisDir = "analysis"

wWindowSize :: Int
wWindowSize = 4096

analysisFile :: FilePath -> FilePath -> String -> String-> FilePath
analysisFile dir what suff ext =
    joinPath [dir, addExtension (what ++ "-" ++ suff) ext]

plotSuffix :: Maybe FileFormat -> String
plotSuffix (Just PDF) = "pdf"
plotSuffix (Just SVG) = "svg"
plotSuffix _ = error "no suffix for (non)plot format"

plotRender :: FileFormat -> String -> Renderable a -> IO ()
plotRender ff name r = do
    _ <- renderableToFile (fo_format .~ ff $ def) name r
    return ()

readSamples :: Int -> B.ByteString -> [Int]
readSamples sampleCount stuff =
    take sampleCount $ makeInts $ B.unpack stuff
    where
      makeInts [] = []
      makeInts [_] = error "odd byte array"
      makeInts (b2 : b1 : bs) =
          let i1 = fromIntegral b1
              i2 = fromIntegral b2
          in
          ((i1 `shiftL` 8) .|. i2) : makeInts bs

takeSpan :: Int -> Int -> [a] -> [a]
takeSpan start len xs =
    take len $ drop start xs

rawHist :: [Int] -> [(Int, Int)]
rawHist samples =
    map histify $ group $ sort $ samples ++ [smallest..largest]
    where
      smallest = minimum samples
      largest = maximum samples
      histify [] = error "bad hist group"
      histify (x : xs) = (x, length xs)

showHist :: [(Int, Int)] -> String
showHist bins =
    unlines $ map showBin bins
    where
      showBin (x, y) = unwords [show x, show y]

type ColourPair = (Colour Double, Colour Double)

plotTimeSeries :: ColourPair -> [Int]
               -> Renderable (LayoutPick Int Int Int)
plotTimeSeries (color, color2) samples = do
  let nSamples = length samples
  let samplePoints = zip [(1::Int)..] samples
  let zoomedPoints =
          zip [0, nSamples `div` 100 ..] $
              takeSpan (nSamples `div` 3) 100 samples
  let allPlot =
          toPlot $
          plot_points_style .~ filledCircles 0.5 (opaque color2) $
          plot_points_values .~ samplePoints $
          plot_points_title .~ "all" $
          def
  let zoomedPlot =
          toPlot $
          plot_lines_style . line_color .~ opaque color $
          plot_lines_values .~ [zoomedPoints] $
          plot_lines_title .~ "zoomed" $
          def
  let tsLayout =
          layout_x_axis .~ (laxis_title .~ "sample" $ def) $
          layout_y_axis .~ (laxis_title .~ "value" $ def) $
          layout_plots .~ [allPlot, zoomedPlot] $
          def
  layoutToRenderable tsLayout

plotSampleHist :: ColourPair -> Int -> [Int]
               -> Renderable (LayoutPick Int Int Int)
plotSampleHist (color, _) nBins samples = do
  let histPlot =
          plotBars $
          plot_bars_values .~  sampleBars ++ [(nBins, [0])] $
          plot_bars_item_styles .~ [barStyle] $
          plot_bars_spacing .~ BarsFixGap 0 0 $
          plot_bars_alignment .~ BarsLeft $
          def
  let histLayout =
          layout_x_axis .~ (laxis_title .~ "value" $ def) $
          layout_y_axis .~ (laxis_title .~ "frequency" $ def) $
          layout_plots .~ [histPlot] $
          def
  layoutToRenderable histLayout
  where
    sampleBars =
        map (\(x, y) -> (x, [y])) $ rawHist samples
    barStyle = (FillStyleSolid (opaque gray),
                Just (line_width .~ 0.05 $ line_color .~ opaque color $ def))

plotSampleDFT :: ColourPair -> [Double]
              -> Renderable (LayoutPick Double Double Double)
plotSampleDFT (color, _) dftBins = do
  let dftPlot =
          plotBars $
          plot_bars_spacing .~ BarsFixWidth 0.3 $
          plot_bars_values .~ dftBars $
          plot_bars_item_styles .~ [barStyle] $
          def
  let dftLayout =
          layout_x_axis .~ (laxis_title .~ "frequency" $ def) $
          layout_y_axis .~ (laxis_title .~ "amplitude" $ def) $
          layout_plots .~ [dftPlot] $
          def
  layoutToRenderable dftLayout
  where
    dftBars = map (\(x, y) -> (x, [y])) $ zip [0..] dftBins
    barStyle = (FillStyleSolid (opaque color), Nothing)

data EntropyMode = EntropyModeRaw | EntropyModeNormalized

entropy :: Real a => EntropyMode -> [a] -> Double
entropy mode samples =
    negate $ sum $ map binEntropy samples
    where
      weight =
          case mode of
            EntropyModeRaw -> 1.0
            EntropyModeNormalized -> realToFrac $ sum samples
      binEntropy 0 = 0
      binEntropy count =
          p * logBase 2 p
          where
            p = realToFrac count / weight

hannWindow :: Int -> [Double]
hannWindow nSamples =
    map hannFunction [0 .. nSamples - 1]
    where
      hannFunction n =
          0.5 * (1.0 - cos(2 * pi * fromIntegral n / nn))
          where
            nn = fromIntegral (nSamples - 1)

sampleDFT :: [Double] -> [Double]
sampleDFT samples =
    map magnitude $ V.toList $ run dftR2C $ V.fromList samples

data Bias = BiasDebiased | BiasNominal Int
data DFTMode = DFTModeRaw |
               DFTModeProper Int (Int -> [Double])

splitSamples :: Int -> [a] -> [[a]]
splitSamples n xs
    | length first < n = []
    | otherwise =
        first : splitSamples n rest
    where
      (first, rest) = splitAt n xs

processDFT :: Bias -> DFTMode -> [Int] -> [Double]
processDFT bias dftMode samples =
    case dftMode of
      DFTModeRaw ->
          sampleDFT dftSamples
          where
            dftStart = (nSamples - dftLength) `div` 3
            dftLength = min 10000 nSamples
            dftSamples = takeSpan dftStart dftLength normedSamples
      DFTModeProper windowSize window ->
        avgBins $ map (sampleDFT . applyWindow) $
          splitSamples windowSize normedSamples
        where
          applyWindow xs =
              zipWith (*) xs $ window windowSize
          avgBins bins =
              map average $ transpose bins
    where
      nSamples = length samples
      normedSamples =
          map norm samples
          where
            norm sample =
                (fromIntegral sample - dftDC) / fromIntegral nSamples
                where
                  dftDC =
                      case bias of
                        BiasNominal nBits ->
                            fromIntegral (2 ^ (nBits - 1) - 1 :: Integer)
                        BiasDebiased ->
                            average samples

average :: Real a => [a] -> Double
average samples =
    realToFrac (sum samples) / fromIntegral (length samples)

stdDeviation :: Real a => [a] -> Double
stdDeviation samples =
    sqrt (sum devs / fromIntegral (length devs))
    where
      devs = map (square . (`subtract` mean) . realToFrac) samples
      square x = x * x
      mean = average samples

spectralFlatness :: [Double] -> Double
spectralFlatness rDFT =
    10.0 * (gMeanDB - aMeanDB)
    where
      xDFT = tail rDFT
      nxDFT = fromIntegral (length xDFT)
      gMeanDB = sum (map (logBase 10) xDFT) / nxDFT
      aMeanDB = logBase 10 (sum xDFT) - logBase 10 nxDFT

secondMoment64 :: [Int] -> (Double, Double)
secondMoment64 samples =
    (minimum devs, maximum devs)
    where
      devs = map stdDeviation $ splitSamples 64 samples

showStats :: Int -> [Int] -> [Double] -> [Double] -> String
showStats nBits samples rDFT wDFT = unlines [
  printf "min: %d" (minimum samples),
  printf "max: %d" (maximum samples),
  printf "mean: %0.3g" (average samples),
  printf "byte-entropy: %0.3g"
      (entropyAdj * entropy EntropyModeNormalized hist),
  printf "second-moment-64s-min: %0.3g" smMin,
  printf "second-moment-64s-max: %0.3g" smMax,
  printf "spectral-entropy: %0.3g"
      (spectralEntropyAdj * entropy EntropyModeNormalized rDFT2),
  printf "spectral-flatness-db: %0.3g" (spectralFlatness rDFT),
  printf "avg-spectral-flatness-db: %0.3g" (spectralFlatness wDFT) ]
  where
    (smMin, smMax) = secondMoment64 samples
    hist = map snd $ rawHist samples
    rDFT2 = map (**2.0) $ tail rDFT
    entropyAdj =
        max 1.0 $ 8.0 / fromIntegral nBits
    spectralEntropyAdj =
        fromIntegral (max 8 nBits)  / logBase 2.0 (fromIntegral (length rDFT2))

analyze :: Maybe FileFormat -> ColourPair
        -> String -> String -> Int -> [Int]
        -> IO ()
analyze plotMode colors dir what nBits samples = do
  let rDFT = processDFT BiasDebiased DFTModeRaw samples
  let wDFT = processDFT BiasDebiased (DFTModeProper wWindowSize hannWindow) samples
  writeFile (af "stats" "txt") $
    showStats nBits samples rDFT wDFT
  case plotMode of
    Just ff -> do
      writeFile (af "hist" "txt") $ showHist $ rawHist samples
      let pr = plotRender ff
      pr (af' "ts") $ plotTimeSeries colors samples
      pr (af' "hist") $ plotSampleHist colors (2 ^ nBits) samples
      pr (af' "dft") $ plotSampleDFT colors rDFT
      pr (af' "wdft") $ plotSampleDFT colors wDFT
    Nothing -> return ()
  where
    af = analysisFile dir what
    af' name = af name (plotSuffix plotMode)

data ArgIndex = ArgIndexBitsFile
              | ArgIndexPlotFormat
              | ArgIndexAnalysisDir
              | ArgIndexSampleCount
              | ArgIndexColor
              | ArgIndexColor2
              | ArgIndexTests
                deriving (Eq, Ord, Show)

maybeReadS :: ReadS a -> String -> Maybe a
maybeReadS f s =
    case f s of
      [(v, "")] -> Just v
      _ -> Nothing

readColor :: String -> Colour Double
readColor name =
    head $ mapMaybe id [
               maybeReadS sRGB24reads name,
               readColourName name,
               error "illegal color" ]

argd :: [Arg ArgIndex]
argd = [
  Arg {
    argIndex = ArgIndexPlotFormat,
    argName = Just "plot-format",
    argAbbr = Just 'p',
    argData = argDataDefaulted "type" ArgtypeString "pdf",
    argDesc = "Plot format (\"pdf\", \"svg\", or \"none\" for just stats)" },
  Arg {
    argIndex = ArgIndexAnalysisDir,
    argName = Just "analysis-dir",
    argAbbr = Just 'a',
    argData = argDataDefaulted "dir" ArgtypeString defaultAnalysisDir,
    argDesc = "Directory for analysis results" },
  Arg {
    argIndex = ArgIndexSampleCount,
    argName = Just "sample-count",
    argAbbr = Just 's',
    argData = argDataOptional "count" ArgtypeInt,
    argDesc = "Number of samples to analyze" },
  Arg {
    argIndex = ArgIndexColor,
    argName = Just "color",
    argAbbr = Nothing,
    argData = argDataDefaulted "name/rgb" ArgtypeString "black",
    argDesc = "Plot color" },
  Arg {
    argIndex = ArgIndexColor2,
    argName = Just "secondary-color",
    argAbbr = Nothing,
    argData = argDataDefaulted "name/rgb" ArgtypeString "gray",
    argDesc = "Plot secondary color" },
  Arg {
    argIndex = ArgIndexTests,
    argName = Just "tests",
    argAbbr = Just 't',
    argData = argDataOptional "test-list" ArgtypeString,
    argDesc = "Comma-separated list of tests (\"help\" for help)" },
  Arg {
    argIndex = ArgIndexBitsFile,
    argName = Nothing,
    argAbbr = Nothing,
    argData = argDataOptional "bits-file" ArgtypeString,
    argDesc = "Random bits." } ]

testSet :: [String]
testSet = ["raw","prng","twobit","mid","mid7","low"]

testHelp :: IO ()
testHelp = do
  printf "Available Tests:\n"
  mapM_ (\t -> printf "  %s\n" t) testSet
  exitSuccess

main :: IO ()
main = do
  args <- parseArgsIO ArgsComplete argd

  let tests =
          case getArg args ArgIndexTests of
            Nothing -> testSet
            Just(ts) -> splitOn "," ts
  when (tests == ["help"]) testHelp

  let plotFormat =
          case map toLower $ getRequiredArg args ArgIndexPlotFormat of
            "none" -> Nothing
            "pdf" -> Just PDF
            "svg" -> Just SVG
            _ -> error "unrecognized plot format"

  bitsFile <- getArgStdio args ArgIndexBitsFile ReadMode
  rawSamples <- B.hGetContents bitsFile
  let nSamples = B.length rawSamples `div` 2
  let sampleCount = case getArg args ArgIndexSampleCount of
                      Nothing -> nSamples
                      Just n -> nSamples `min` n
  let samples = readSamples sampleCount rawSamples

  let analysisDir = getRequiredArg args ArgIndexAnalysisDir
  createDirectoryIfMissing True analysisDir

  let plotColor = readColor $ getRequiredArg args ArgIndexColor
  let plotColor2 = readColor $ getRequiredArg args ArgIndexColor2

  let aa name bits source =
          when (name `elem` tests) $
               analyze plotFormat (plotColor, plotColor2) analysisDir
                       name bits source

  aa "raw" 12 samples

  prng12Samples <- replicateM (length samples) (randomRIO (0, 4095) :: IO Int)
  aa "prng12" 12 prng12Samples

  prngSamples <- replicateM (length samples) (randomRIO (0, 255) :: IO Int)
  aa "prng" 8 prngSamples

  let twoBitSamples = map (.&. 0x03) samples
  aa "twobit" 2 twoBitSamples

  let midSamples = map ((.&. 0xff) . (`shiftR` 1)) samples
  aa "mid" 8 midSamples

  let mid7Samples = map ((.&. 0x7f) . (`shiftR` 2)) samples
  aa "mid7" 7 mid7Samples

  let lowSamples = map (.&. 0xff) samples
  aa "low" 8 lowSamples

  let hmSamples = map hmBits samples
                  where
                    hmBits sample =
                        let mid = (sample `shiftR` 1) .&. 0xff in
                        mid `xor` (sample .&. 0x01) `xor`
                          ((sample `shiftR` 9) .&. 0x01) `xor`
                          ((sample `shiftR` 10) .&. 0x01) `xor`
                          ((sample `shiftR` 11) .&. 0x01)
  aa "hm" 8 hmSamples


{-# LANGUAGE CPP                 #-}
{-# LANGUAGE FlexibleContexts    #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeApplications    #-}
{-# LANGUAGE TypeOperators       #-}

#if MIN_VERSION_base(4,12,0)
{-# LANGUAGE NoStarIsType         #-}
#endif

-- |
-- Module      : Numeric.LinearAlgebra.Static.Vector
-- Copyright   : (c) Justin Le 2018
-- License     : BSD3
--
-- Maintainer  : justin@jle.im
-- Stability   : experimental
-- Portability : non-portable
--
-- Conversions between statically sized types in
-- "Numeric.LinearAlgebra.Static" from /hmatrix/ and /vector-sized/.
--
-- This module is intentionally minimal, exporting only functions that
-- cannot be written without "unsafe" operations.  With these, however, you
-- can easily write other useful combinators by using type-safe operations
-- like 'fmap', 'VS.map', 'liftA2', 'Data.Vector.Generic.Sized.convert',
-- etc.
--

module Numeric.LinearAlgebra.Static.Vector (
  -- * Vector
  -- ** Real
    rVec
  , grVec
  , vecR
  , gvecR
  -- ** Complex
  , cVec
  , gcVec
  , vecC
  , gvecC
  -- * Matrix
  -- ** Real
  , lRows
  , rowsL
  , lCols
  , colsL
  , lVec
  , glVec
  , vecL
  , gvecL
  -- ** Complex
  , mRows
  , rowsM
  , mCols
  , colsM
  , mVec
  , gmVec
  , vecM
  , gvecM
  ) where

import           Data.Foldable
import           Data.Proxy
import           GHC.TypeLits
import           Unsafe.Coerce
import qualified Data.Vector                  as UV
import qualified Data.Vector.Generic          as UVG
import qualified Data.Vector.Generic.Sized    as VG
import qualified Data.Vector.Sized            as V
import qualified Data.Vector.Storable.Sized   as VS
import qualified Numeric.LinearAlgebra        as HU
import qualified Numeric.LinearAlgebra.Static as H

-- | Convert an /hmatrix/ vector (parameterized by its lenth) to
-- a /vector-sized/ storable vector of 'Double's.
--
-- This is normally /O(1)/, but will be /O(n)/ if the 'H.R' was contructed
-- with 'H.konst' or any other replicated-value constructor (like literals
-- and 'fromInteger'/'fromRational').
rVec :: KnownNat n => H.R n -> VS.Vector n H.ℝ
rVec = unsafeCoerce . H.extract

-- | 'rVec', but generalized to work for all types of sized vectors.
--
-- Usually /O(n)/, but if using this with storable vectors, should have the
-- same characteristics as 'rVec' due to rewrite rules.
--
-- @since 0.1.3.0
grVec :: (KnownNat n, UVG.Vector v H.ℝ) => H.R n -> VG.Vector v n H.ℝ
grVec = VG.convert . rVec
{-# NOINLINE[1] grVec #-}
{-# RULES "grVec" grVec = rVec #-}

-- | Convert a /vector-sized/ storable vector to an /hmatrix/ vector
-- (parameterized by its lenth).
--
-- /O(1)/
vecR :: VS.Vector n H.ℝ -> H.R n
vecR = unsafeCoerce

-- | 'vecR', but generalized to work for all types of sized vectors.
--
-- Usually /O(n)/, but if using this with storable vectors, should be /O(1)/
-- due to rewrite rules (but don't rely on this).
--
-- @since 0.1.3.0
gvecR :: UVG.Vector v H.ℝ => VG.Vector v n H.ℝ -> H.R n
gvecR = vecR . VG.convert
{-# NOINLINE[1] gvecR #-}
{-# RULES "gvecR" gvecR = vecR #-}

-- | Convert an /hmatrix/ complex vector (parameterized by its lenth) to
-- a /vector-sized/ storable vector of 'Complex Double's, preserving the
-- length in the type.
--
-- This is normally /O(1)/, but will be /O(n)/ if the 'H.C' was contructed
-- with 'H.konst' or any other replicated-value constructor (like literals
-- and 'fromInteger'/'fromRational').
cVec :: KnownNat n => H.C n -> VS.Vector n H.ℂ
cVec = unsafeCoerce . H.extract

-- | 'cVec', but generalized to work for all types of sized vectors.
--
-- Usually /O(n)/, but if using this with storable vectors, should have the
-- same characteristics as 'cVec' due to rewrite rules.
--
-- @since 0.1.3.0
gcVec :: (KnownNat n, UVG.Vector v H.ℂ) => H.C n -> VG.Vector v n H.ℂ
gcVec = VG.convert . cVec
{-# NOINLINE[1] gcVec #-}
{-# RULES "gcVec" gcVec = cVec #-}

-- | Convert a /vector-sized/ storable vector to an /hmatrix/ complex
-- vector (parameterized by its lenth), preserving the length in the type.
--
-- /O(1)/
vecC :: VS.Vector n H.ℂ -> H.C n
vecC = unsafeCoerce

-- | 'vecC', but generalized to work for all types of sized vectors.
--
-- Usually /O(n)/, but if using this with storable vectors, should be /O(1)/
-- due to rewrite rules (but don't rely on this).
--
-- @since 0.1.3.0
gvecC :: UVG.Vector v H.ℂ => VG.Vector v n H.ℂ -> H.C n
gvecC = vecC . VG.convert
{-# NOINLINE[1] gvecC #-}
{-# RULES "gvecC" gvecC = vecC #-}

-- | Split an /hmatrix/ matrix (parameterized by its dimensions) to
-- a /vector-sized/ boxed vector of its rows (as /hmatrix/ vectors).
--
-- This is normally /O(m*n)/, but can sometimes be /O(m)/ depending on the
-- representation of the 'H.L' being used.
lRows
    :: (KnownNat m, KnownNat n)
    => H.L m n
    -> V.Vector m (H.R n)
lRows = unsafeCoerce
      . UV.fromList
      . HU.toRows
      . H.extract

-- | Join together a /vector-sized/ boxed vector of /hmatrix/ vectors to an
-- /hmatrix/ matrix as its rows.
--
-- /O(m*n)/
rowsL
    :: forall m n. KnownNat n
    => V.Vector m (H.R n)
    -> H.L m n
rowsL = (unsafeCoerce :: HU.Matrix Double -> H.L m n)
      . HU.fromRows
      . map H.extract
      . toList

-- | Split an /hmatrix/ matrix (parameterized by its dimensions) to
-- a /vector-sized/ boxed vector of its columns (as /hmatrix/ vectors).
--
-- This is normally /O(m*n)/, but can sometimes be /O(n)/ depending on the
-- representation of the 'H.L' being used.
lCols
    :: forall m n. (KnownNat m, KnownNat n)
    => H.L m n
    -> V.Vector n (H.R m)
lCols = unsafeCoerce
      . UV.fromList
      . HU.toColumns
      . H.extract

-- | Join together a /vector-sized/ boxed vector of /hmatrix/ vectors to an
-- /hmatrix/ matrix as its columns.
--
-- /O(m*n)/
colsL
    :: forall m n. KnownNat m
    => V.Vector n (H.R m)
    -> H.L m n
colsL = (unsafeCoerce :: HU.Matrix Double -> H.L m n)
      . HU.fromColumns
      . map H.extract
      . toList

-- | Split an /hmatrix/ complex matrix (parameterized by its dimensions) to
-- a /vector-sized/ boxed vector of its rows (as /hmatrix/ complex
-- vectors).
--
-- This is normally /O(m*n)/, but can sometimes be /O(m)/ depending on the
-- representation of the 'H.C' being used.
mRows
    :: forall m n. (KnownNat m, KnownNat n)
    => H.M m n
    -> V.Vector m (H.C n)
mRows = unsafeCoerce
      . UV.fromList
      . HU.toRows
      . H.extract

-- | Join together a /vector-sized/ boxed vector of /hmatrix/ complex
-- vectors to an /hmatrix/ complex matrix as its rows.
--
-- /O(m*n)/
rowsM
    :: forall m n. KnownNat n
    => V.Vector m (H.C n)
    -> H.M m n
rowsM = (unsafeCoerce :: HU.Matrix H.ℂ -> H.M m n)
      . HU.fromRows
      . map H.extract
      . toList

-- | Split an /hmatrix/ complex matrix (parameterized by its dimensions) to
-- a /vector-sized/ boxed vector of its columns (as /hmatrix/ complex
-- vectors).
--
-- This is normally /O(m*n)/, but can sometimes be /O(n)/ depending on the
-- representation of the 'H.C' being used.
mCols
    :: forall m n. (KnownNat m, KnownNat n)
    => H.M m n
    -> V.Vector n (H.C m)
mCols = unsafeCoerce
      . UV.fromList
      . HU.toColumns
      . H.extract

-- | Join together a /vector-sized/ boxed vector of /hmatrix/ complex
-- vectors to an /hmatrix/ complex matrix as its columns.
--
-- /O(m*n)/
colsM
    :: forall m n. KnownNat m
    => V.Vector n (H.C m)
    -> H.M m n
colsM = (unsafeCoerce :: HU.Matrix H.ℂ -> H.M m n)
      . HU.fromColumns
      . map H.extract
      . toList

-- | Shape a /vector-sized/ storable vector of elements into an /hmatrix/
-- matrix.
--
-- /O(1)/
--
-- @since 0.1.1.0
vecL
    :: forall m n. KnownNat n
    => VS.Vector (m * n) H.ℝ
    -> H.L m n
vecL = (unsafeCoerce :: HU.Matrix H.ℝ -> H.L m n)
     . HU.reshape (fromIntegral (natVal (Proxy @n)))
     . unsafeCoerce

-- | 'vecL', but generalized to work for all types of sized vectors.
--
-- Usually /O(n)/, but if using this with storable vectors, should be /O(1)/
-- due to rewrite rules (but don't rely on this).
--
-- @since 0.1.3.0
gvecL
    :: (KnownNat n, UVG.Vector v H.ℝ)
    => VG.Vector v (m * n) H.ℝ
    -> H.L m n
gvecL = vecL . VG.convert
{-# NOINLINE[1] gvecL #-}
{-# RULES "gvecL" gvecL = vecL #-}

-- | Flatten an /hmatrix/ matrix into a /vector-sized/ storable vector of
-- its items.
--
-- This is normally /O(m*n)/, but can sometimes be /O(1)/ depending on the
-- representation of the 'H.L' being used.
--
-- @since 0.1.1.0
lVec
    :: forall m n. (KnownNat m, KnownNat n)
    => H.L m n
    -> VS.Vector (m * n) H.ℝ
lVec = unsafeCoerce
     . HU.flatten
     . H.extract

-- | 'lVec', but generalized to work for all types of sized vectors.
--
-- Usually /O(m*n)/, but if using this with storable vectors, should have the
-- same characteristics as 'lVec' due to rewrite rules.
--
-- @since 0.1.3.0
glVec
    :: (KnownNat m, KnownNat n, UVG.Vector v H.ℝ)
    => H.L m n
    -> VG.Vector v (m * n) H.ℝ
glVec = VG.convert . lVec
{-# NOINLINE[1] glVec #-}
{-# RULES "glVec" glVec = lVec #-}

-- | Shape a /vector-sized/ storable vector of elements into an /hmatrix/
-- complex matrix.
--
-- /O(1)/
--
-- @since 0.1.1.0
vecM
    :: forall m n. KnownNat n
    => VS.Vector (m * n) H.ℂ
    -> H.M m n
vecM = (unsafeCoerce :: HU.Matrix H.ℂ -> H.M m n)
     . HU.reshape (fromIntegral (natVal (Proxy @n)))
     . unsafeCoerce

-- | 'vecM', but generalized to work for all types of sized vectors.
--
-- Usually /O(m*n)/, but if using this with storable vectors, should be /O(1)/
-- due to rewrite rules (but don't rely on this).
--
-- @since 0.1.3.0
gvecM
    :: (KnownNat n, UVG.Vector v H.ℂ)
    => VG.Vector v (m * n) H.ℂ
    -> H.M m n
gvecM = vecM . VG.convert
{-# NOINLINE[1] gvecM #-}
{-# RULES "gvecM" gvecM = vecM #-}

-- | Flatten an /hmatrix/ complex matrix into a /vector-sized/ storable
-- vector of its items.
--
-- This is normally /O(m*n)/, but can sometimes be /O(1)/ depending on the
-- representation of the 'H.M' being used.
--
-- @since 0.1.1.0
mVec
    :: forall m n. (KnownNat m, KnownNat n)
    => H.M m n
    -> VS.Vector (m * n) H.ℂ
mVec = unsafeCoerce
     . HU.flatten
     . H.extract

-- | 'mVec', but generalized to work for all types of sized vectors.
--
-- Usually /O(m*n)/, but if using this with storable vectors, should have the
-- same characteristics as 'mVec' due to rewrite rules.
--
-- @since 0.1.3.0
gmVec
    :: (KnownNat m, KnownNat n, UVG.Vector v H.ℂ)
    => H.M m n
    -> VG.Vector v (m * n) H.ℂ
gmVec = VG.convert . mVec
{-# NOINLINE[1] gmVec #-}
{-# RULES "gmVec" gmVec = mVec #-}


module Noether.Algebra.Actions.Compatible where

import           Data.Complex

import           Noether.Lemmata.Prelude
import           Noether.Lemmata.TypeFu

import           Noether.Algebra.Single
import           Noether.Algebra.Tags

import           Noether.Algebra.Actions.Acts

{-| A strategy-parameterized typeclass for a compatible action, where compatibility
    is defined in the group action sense.

    A compatible action satisfies
    a `act` (a' `act` b) = (a `op` a') `act` b
-}

class CompatibleK (lr :: Side) (op :: k1) (act :: k2) a b (s :: CompatibleE)

data CompatibleE = Compatible_Acts_Semigroup
  { compatible_actor           :: Type
  , compatible_action          :: ActsE
  , compatible_actor_semigroup :: SemigroupE
  }

type family CompatibleS (lr :: Side) (op :: k1) (act :: k2) (a :: Type) (b :: Type) = (r :: CompatibleE)

instance (ActsK lr act a b za, SemigroupK op a zs) =>
         CompatibleK lr op act a b (Compatible_Acts_Semigroup a za zs)


{-# LANGUAGE FlexibleContexts #-}
-- |
-- Module      : PSO.NeuralSpec
-- Description : Tests for PSO.Neural
-- Copyright   : (c) Tom Westerhout, 2017
-- License     : BSD3
-- Maintainer  : t.westerhout@student.ru.nl
-- Stability   : experimental

module PSO.NeuralSpec where

import           Control.Monad
import           Control.Monad.ST
import           Control.Monad.Reader
import           Data.Complex
import qualified Data.Vector.Storable as V
import           Foreign.Storable
import qualified Numeric.LinearAlgebra as LA
import qualified Numeric.LinearAlgebra.Devel as LA.Devel

import           Test.Hspec

import           PSO.Random
import           PSO.Neural

main :: IO ()
main = hspec spec

roundTo :: (RealFrac a) => Int -> a -> a
roundTo n x = (fromInteger . round $ x * (10^n)) / (10.0^^n)

roundToC :: (RealFrac a) => Int -> Complex a -> Complex a
roundToC n (x :+ y) = (roundTo n x) :+ (roundTo n y)

class AlmostEq a where
  eq' :: a -> a -> Bool

instance AlmostEq Float where
  eq' x y = let epsilon = 1.0E-5
             in abs (x - y) <= epsilon * max (abs x) (abs y)

instance AlmostEq a => AlmostEq (Complex a) where
  eq' (xr :+ xi) (yr :+ yi) = xr `eq'` yr && xi `eq'` yi

instance (Storable a, AlmostEq a) => AlmostEq (V.Vector a) where
  eq' x y = V.and $ V.zipWith eq' x y

toW n m w = LA.Devel.matrixFromVector LA.Devel.ColumnMajor m n w

mkTheta a b w' σ = b + (toW (V.length a) (V.length b) w') LA.#> σ

logWF' :: (Storable a, RealFloat a, Num (V.Vector (Complex a)), LA.Numeric (Complex a))
       => V.Vector (Complex a)
       -> V.Vector (Complex a)
       -> V.Vector (Complex a)
       -> V.Vector (Complex a)
       -> Complex a
logWF' a b w' σ =
  let n = V.length a
      m = V.length b
      θ = mkTheta a b w' σ
   in V.sum (V.zipWith (*) a σ) + V.sum (V.map (log . cosh) θ)

logQuotient1' ::
     (Storable a, RealFloat a, Num (V.Vector (Complex a)), LA.Numeric (Complex a))
  => V.Vector (Complex a)
  -> V.Vector (Complex a)
  -> V.Vector (Complex a)
  -> V.Vector (Complex a)
  -> Int -> Complex a
logQuotient1' a b w' σ flip =
  let σ' = σ V.// [(flip, (-1) * (σ V.! flip))]
   in logWF' a b w' σ' - logWF' a b w' σ

logQuotient2' ::
     (Storable a, RealFloat a, Num (V.Vector (Complex a)), LA.Numeric (Complex a))
  => V.Vector (Complex a)
  -> V.Vector (Complex a)
  -> V.Vector (Complex a)
  -> V.Vector (Complex a)
  -> Int -> Int -> Complex a
logQuotient2' a b w' σ flip1 flip2 =
  let σ' = σ V.// [ (flip1, (-1) * (σ V.! flip1))
                  , (flip2, (-1) * (σ V.! flip2)) ]
   in logWF' a b w' σ' - logWF' a b w' σ

locEnergyHH1DOpen' ::
     (Storable α, RealFloat α, Eq α, Num (V.Vector (Complex α)), LA.Numeric (Complex α))
  => V.Vector (Complex α)
  -> V.Vector (Complex α)
  -> V.Vector (Complex α)
  -> V.Vector (Complex α)
  -> Complex α
locEnergyHH1DOpen' a b w' σ =
  let zipper i x y
        | x == y    = 1
        | otherwise = -1 + 2 * exp (logQuotient2' a b w' σ i (i + 1))
   in V.sum $ V.izipWith zipper σ (V.tail σ)

a1 = V.fromList [   3.7e-01  :+   1.5e-01
                ,   6.5e-01  :+   3.0e-01
                , (-2.8e-01) :+ (-3.0e-01)
                ] :: V.Vector (Complex Float)

ra1 :: UniformDist m (Complex Float) => m (V.Vector (Complex Float))
ra1 = uniformVector 10 ((-1.0) :+ (-1.0), 1.0 :+ 1.0)

b1 = V.fromList [   6.0e-01  :+   8.4e-02
                , (-2.6e-01) :+   9.2e-01
                ,   5.5e-01  :+ (-8.2e-02)
                ,   4.3e-01  :+   3.7e-01
                ,   9.3e-01  :+   5.1e-01
                ,   6.0e-02  :+   8.9e-01
                ] :: V.Vector (Complex Float)

rb1 :: UniformDist m (Complex Float) => m (V.Vector (Complex Float))
rb1 = uniformVector 20 ((-1.0) :+ (-1.0), 1.0 :+ 1.0)

w1 = V.fromList $
      [   0.22 :+ (-0.93),   (-0.16) :+ (-0.47), (-0.57) :+ (-0.83),   (-0.16) :+ 0.86,    (-0.44) :+ 0.16,       0.43 :+ 0.96
      , (-3.6e-2) :+ 0.36, (-1.5e-2) :+ (-0.14),    (-0.15) :+ 0.57,   0.79 :+ (-0.26),     0.95 :+ (-0.6),       0.56 :+ 0.98
      ,      0.82 :+ 0.34,      0.85 :+ (-0.35),    0.64 :+ (-0.24), (-0.73) :+ (-0.8), (-0.34) :+ (-0.24), (-0.94) :+ (-0.46)
      ] :: V.Vector (Complex Float)

rw1 :: UniformDist m (Complex Float) => m (V.Vector (Complex Float))
rw1 = uniformVector (10 * 20) ((-1.0) :+ (-1.0), 1.0 :+ 1.0)

σ1  = V.fromList [ -1, -1, -1 ] :: V.Vector (Complex Float)
σ2  = V.fromList [ 1,   1,  1 ] :: V.Vector (Complex Float)
σ3  = V.fromList [ 1,  -1,  1 ] :: V.Vector (Complex Float)

rσ1 :: Randomisable m Bool => m (V.Vector (Complex Float))
rσ1 = randomSpin 10

rbm1  = runST $ mkRbm a1 b1 w1

mcmc1 = runST $ newMcmc rbm1 σ1 >>= unsafeFreezeMcmc
mcmc2 = runST $ newMcmc rbm1 σ2 >>= unsafeFreezeMcmc
mcmc3 = runST $ newMcmc rbm1 σ3 >>= unsafeFreezeMcmc

spec :: Spec
spec = do
  describe "mkRbm" $ do
    it "Constructs a new RBM given a, b, and w" $
      debugPrintRbm rbm1
  describe "newMcmc" $ do
    it "Constructs a new MCMC given RBM and σ" $
      do debugPrintMcmc mcmc1
         print $ (b1 + (toW 3 6 w1) LA.#> σ1)
  describe "logWF" $ do
    it "1) Calculates ln(ψ(S)) " $
      do putStrLn $ "logWF:  " ++ show (logWF mcmc1)
         putStrLn $ "logWF': " ++ show (logWF' a1 b1 w1 σ1)
    it "2) Calculates ln(ψ(S)) " $
      do putStrLn $ "logWF:  " ++ show (logWF mcmc2)
         putStrLn $ "logWF': " ++ show (logWF' a1 b1 w1 σ2)
    it "3) Calculates ln(ψ(S)) " $
      do putStrLn $ "logWF:  " ++ show (logWF mcmc3)
         putStrLn $ "logWF': " ++ show (logWF' a1 b1 w1 σ3)
  describe "logQuotient1" $ do
    it "1) Calculates ln(ψ(S')/ψ(S))" $
      do putStrLn $ "logQuotient1:  " ++ show (logQuotient1 mcmc1 0)
         putStrLn $ "logQuotient1': " ++ show (logQuotient1' a1 b1 w1 σ1 0)
    it "2) Calculates ln(ψ(S')/ψ(S))" $
      do putStrLn $ "logQuotient1:  " ++ show (logQuotient1 mcmc1 1)
         putStrLn $ "logQuotient1': " ++ show (logQuotient1' a1 b1 w1 σ1 1)
    it "3) Calculates ln(ψ(S')/ψ(S))" $
      do putStrLn $ "logQuotient1:  " ++ show (logQuotient1 mcmc1 2)
         putStrLn $ "logQuotient1': " ++ show (logQuotient1' a1 b1 w1 σ1 2)
    it "4) Calculates ln(ψ(S')/ψ(S))" $
      do putStrLn $ "logQuotient1:  " ++ show (logQuotient1 mcmc2 0)
         putStrLn $ "logQuotient1': " ++ show (logQuotient1' a1 b1 w1 σ2 0)
    it "5) Calculates ln(ψ(S')/ψ(S))" $
      do putStrLn $ "logQuotient1:  " ++ show (logQuotient1 mcmc2 2)
         putStrLn $ "logQuotient1': " ++ show (logQuotient1' a1 b1 w1 σ2 2)
  describe "logQuotient2" $ do
    it "1) Calculates ln(ψ(S')/ψ(S))" $
      do putStrLn $ "logQuotient2:  " ++ show (logQuotient2 mcmc1 0 1)
         putStrLn $ "logQuotient2': " ++ show (logQuotient2' a1 b1 w1 σ1 0 1)
    it "2) Calculates ln(ψ(S')/ψ(S))" $
      do putStrLn $ "logQuotient2:  " ++ show (logQuotient2 mcmc1 0 2)
         putStrLn $ "logQuotient2': " ++ show (logQuotient2' a1 b1 w1 σ1 0 2)
  describe "locEnergyHH1DOpen'" $ do
    it "1) Calculates Eloc(ψ)" $
      do putStrLn $ "locEnergyHH1DOpen:  " ++ show (locEnergyHH1DOpen mcmc3)
         putStrLn $ "locEnergyHH1DOpen': " ++ show (locEnergyHH1DOpen' a1 b1 w1 σ3)
  describe "Random!" $ do
    it "1) Calculates Eloc(ψ)" $
      do let job = do
               a <- ra1
               b <- rb1
               w <- rw1
               σ <- rσ1
               rbm <- mkRbm a b w
               mcmc <- newMcmc rbm σ >>= unsafeFreezeMcmc
               lift $ putStrLn
                    $ "locEnergyHH1DOpen:  " ++ show (locEnergyHH1DOpen mcmc)
               lift $ putStrLn
                    $ "locEnergyHH1DOpen': " ++ show (locEnergyHH1DOpen' a b w σ)
         runReaderT job =<< mkMWCGen (Just 135)
         runReaderT job =<< mkMWCGen (Just 136)
         runReaderT job =<< mkMWCGen (Just 137)


module Budget.Budgeter(
  budgeter
  )where

import Numeric.LinearAlgebra
import Budget.Types
import Budget.Parse
import Budget.Printer
import Text.Parsec
import Finance.FutureValue

getNeededExpense :: BudgetItem -> Expense
getNeededExpense bs = case needed bs of
  True -> Expense (name bs) (minPrice bs)
  False -> Expense (name bs) 0.0

reduceBs :: [BudgetItem] -> [Expense] -> [BudgetItem]
reduceBs bs xs = map removeMin $ zip bs xs

removeMin :: (BudgetItem, Expense) -> BudgetItem
removeMin ((BudgetItem s i l h n), Expense _ m) = BudgetItem s i (l-m) (h-m) n

aDiag :: BudgetItem -> Double
aDiag (BudgetItem _ i low high _)
  | low == high = 100000
  | otherwise = (-2.0) * i * k**2
  where k = 1.0 / (high-low)

aDiags :: [BudgetItem] -> [Double]
aDiags = map aDiag

type CurrentIndex = Integer
type Rank = Integer
aList :: CurrentIndex -> Rank -> [Double] -> [Double]
aList c r (x:xs)
  |c `mod` (r+1) == 0 = x : aList (c+1) r xs
aList c r xs
  | c == r^2-1 = [0]
  | ((c+1) `mod` r == 0) && (c /= r^2-1) = -1 : aList (c+1) r xs
  | (c `mod` (r+1) /= 0 && c >= r^2-r) = 1 : aList (c+1) r xs
  | otherwise = 0 : aList (c+1) r xs

availableCosts :: Double -> Double -> Double -> Double -> Double -> Double
availableCosts cpm y v ipm apr
  | cpm <= 0 = 0
  | retirement > sustainableRetirement = cpm
  | otherwise = availableCosts (cpm-10) y v ipm apr
  where retirement = futureValue numberMonths v (ipm-cpm) ratePerMonth
        sustainableRetirement = costsPerYear / apr
        numberMonths = floor (12 * y)
        ratePerMonth = apr/12
        costsPerYear = cpm * 12

aMatrix :: [BudgetItem] -> Matrix Double
aMatrix bs = let as = aDiags bs
                 r = (length as) + 1
             in (r><r) $ aList 0 (toInteger r) as

bElem :: BudgetItem -> Double
bElem (BudgetItem _ i low high _)
  | low == high = 100000
  | otherwise = (-2.0) * i * k**2 * high
  where k = 1.0 / (high-low)

bElems :: [BudgetItem] -> [Double]
bElems = map bElem

type TotalCost = Double
bVector :: [BudgetItem] -> TotalCost -> Vector Double
bVector bs tc = let vecElems = bElems bs ++ [tc]
                in vector vecElems

removeNeeded :: BudgetInfo -> (BudgetInfo, [Expense])
removeNeeded (BudgetInfo i s y a bs) = (BudgetInfo (i - totalNeeded) s y a newBs, needs)
  where needs = map getNeededExpense bs
        newBs = reduceBs bs needs
        totalNeeded = sum [x | (Expense _ x) <- needs]

type ItemNames = [String]
calculateBudget :: BudgetInfo -> [Expense]
calculateBudget (BudgetInfo i s y a bs)
  | i > totalMaxCosts = expenses
  | otherwise = calculateBudget' [] (BudgetInfo i s y a bs)
  where maxCosts = map maxPrice bs
        names = map name bs
        expenses = [Expense n m | (n, m) <- zip names maxCosts]
        totalMaxCosts = sum maxCosts

inefficientNames :: [BudgetItem] -> Costs -> ItemNames
inefficientNames bs costs  = [(name b)| (b,c) <- (zip bs costs), c < 0]

optimizedCosts :: Double -> [BudgetItem] -> [Double]
optimizedCosts i bs = init $ toList $ (aMatrix bs) <\> (bVector bs (realToFrac i))

calculateBudget' :: ItemNames -> BudgetInfo -> [Expense]
calculateBudget' unusedItems (BudgetInfo i s y a bs)
  | validCosts = expenses
  | otherwise = calculateBudget' newUnusedNames cleanedBI
  where costs = optimizedCosts i bs
        validCosts = length [c | c <- costs, c < 0] == 0
        allNames = (map name bs) ++ unusedItems
        newUnusedNames = unusedItems ++ inefficientNames bs costs
        cleanedBI = BudgetInfo i s y a [b | (b,c) <- (zip bs costs), c > 0]
        appendedCosts = costs ++ (replicate ((length allNames) - (length costs)) 0)
        expenses = [Expense n m | (n, m) <- zip allNames appendedCosts]

storeSavings (BudgetInfo ipm v ytr apr bs) = BudgetInfo cpm v ytr apr bs
  where cpm = availableCosts ipm ytr v ipm apr

unsafe :: Either ParseError BudgetInfo -> BudgetInfo
unsafe (Right x) = x

yearlySavings (BudgetInfo ipm _ _ _ _) cs = ipm - (sum cs)

combineExpenses :: [Expense] -> [Expense] -> [Expense]
combineExpenses e1s e2s = [Expense n1 (cost1 + cost2) | (Expense n1 cost1) <- e1s, (Expense n2 cost2) <- e2s,
                           n1 == n2]

getBudget :: BudgetInfo -> String
getBudget bi = let saved_bi = storeSavings bi
                   (available_bi, mandatory_costs) = removeNeeded saved_bi
                   distro = calculateBudget available_bi
                   finalExpenses = combineExpenses mandatory_costs distro
               in fullPrint bi finalExpenses

budgeter :: String -> String
budgeter s = let bi = parse getBudgetInfo "Fetching budget info" s
             in case fmap getBudget bi of
                Left parsecError -> "Encountered Input Parsing errror:\n" ++ show parsecError
                Right finalString -> finalString


{-# LANGUAGE BangPatterns        #-}
{-# LANGUAGE CPP                 #-}
{-# LANGUAGE DataKinds           #-}
{-# LANGUAGE FlexibleContexts    #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TupleSections       #-}
{-# LANGUAGE TypeFamilies        #-}
{-# LANGUAGE TypeOperators       #-}

import           Control.Applicative
import           Control.DeepSeq
import           Control.Monad
import           Control.Monad.Random
import           Control.Monad.Trans.Except

import qualified Data.Attoparsec.Text         as A
import           Data.List                    (foldl')
import qualified Data.Text                    as T
import qualified Data.Text.IO                 as T
import qualified Data.Vector.Storable         as V

import           Numeric.LinearAlgebra        (maxIndex)
import qualified Numeric.LinearAlgebra.Static as SA

import           Options.Applicative

import           Grenade
import           Grenade.Utils.OneHot

--
-- Note: Input files can be downloaded at https://www.kaggle.com/scolianni/mnistasjpg
--

-- It's logistic regression!
--
-- This network is used to show how we can embed a Network as a layer in the larger MNIST
-- type.
type FL i o =
  Network
    '[ FullyConnected i o, Logit ]
    '[ 'D1 i, 'D1 o, 'D1 o ]

-- The definition of our convolutional neural network.
-- In the type signature, we have a type level list of shapes which are passed between the layers.
-- One can see that the images we are inputing are two dimensional with 28 * 28 pixels.

-- It's important to keep the type signatures, as there's many layers which can "squeeze" into the gaps
-- between the shapes, so inference can't do it all for us.

-- With the mnist data from Kaggle normalised to doubles between 0 and 1, learning rate of 0.01 and 15 iterations,
-- this network should get down to about a 1.3% error rate.
--
-- /NOTE:/ This model is actually too complex for MNIST, and one should use the type given in the readme instead.
--         This one is just here to demonstrate Inception layers in use.
--
type MNIST =
  Network
    '[ Reshape,
       Concat ('D3 28 28 1) Trivial ('D3 28 28 14) (InceptionMini 28 28 1 5 9),
       Pooling 2 2 2 2, Relu,
       Concat ('D3 14 14 3) (Convolution 15 3 1 1 1 1) ('D3 14 14 15) (InceptionMini 14 14 15 5 10), Crop 1 1 1 1, Pooling 3 3 3 3, Relu,
       Reshape, FL 288 80, FL 80 10 ]
    '[ 'D2 28 28, 'D3 28 28 1,
       'D3 28 28 15, 'D3 14 14 15, 'D3 14 14 15, 'D3 14 14 18,
       'D3 12 12 18, 'D3 4 4 18, 'D3 4 4 18,
       'D1 288, 'D1 80, 'D1 10 ]

randomMnist :: IO MNIST
randomMnist = randomNetwork

convTest :: Int -> FilePath -> FilePath -> Optimizer opt -> ExceptT String IO ()
convTest iterations trainFile validateFile opt = do
  net0         <- lift randomMnist
  trainData    <- readMNIST trainFile
  validateData <- readMNIST validateFile
  lift $ foldM_ (runIteration trainData validateData) net0 [1..iterations]

    where
  trainEach !opt' !network (!i, !o) = force (train opt' network i o)

  runIteration !trainRows !validateRows !net !i = do
    putStrLn $ "Number of training rows: " ++ show (length trainRows)
    let !trained' = foldl' (trainEach (sgdUpdateLearningParamters opt)) net trainRows
    let !res      = fmap (\(rowP,rowL) -> (rowL,) $ runNet trained' rowP) validateRows
    let !res'     = fmap (\(S1D label, S1D prediction) -> (maxIndex (SA.extract label), maxIndex (SA.extract prediction))) res
    print trained'
    putStrLn $ "Iteration " ++ show i ++ ": " ++ show (length (filter ((==) <$> fst <*> snd) res')) ++ " of " ++ show (length res')
    return trained'
  sgdUpdateLearningParamters :: Optimizer opt -> Optimizer opt
  sgdUpdateLearningParamters (OptSGD rate mom reg) = OptSGD rate mom (reg * 10)
  sgdUpdateLearningParamters o                     = o

data MnistOpts = MnistOpts FilePath FilePath Int Bool (Optimizer 'SGD) (Optimizer 'Adam)

mnist' :: Parser MnistOpts
mnist' = MnistOpts <$> argument str (metavar "TRAIN")
                   <*> argument str (metavar "VALIDATE")
                   <*> option auto (long "iterations" <> short 'i' <> value 15)
                 <*> flag False True (long "use-adam" <> short 'a')
                 <*> (OptSGD
                       <$> option auto (long "train_rate" <> short 'r' <> value 0.01)
                       <*> option auto (long "momentum" <> value 0.9)
                       <*> option auto (long "l2" <> value 0.0005)
                       )
                 <*> (OptAdam
                       <$> option auto (long "alpha" <> short 'r' <> value 0.001)
                       <*> option auto (long "beta1" <> value 0.9)
                       <*> option auto (long "beta2" <> value 0.999)
                       <*> option auto (long "epsilon" <> value 1e-4)
                       <*> option auto (long "lambda" <> value 1e-3)
                      )

main :: IO ()
main = do
    MnistOpts mnist vali iter useAdam sgd adam <- execParser (info (mnist' <**> helper) idm)
    putStrLn "Training convolutional neural network..."
    res <- if useAdam
      then runExceptT $ convTest iter mnist vali adam
      else runExceptT $ convTest iter mnist vali sgd

    case res of
      Right () -> pure ()
      Left err -> putStrLn err

readMNIST :: FilePath -> ExceptT String IO [(S ('D2 28 28), S ('D1 10))]
readMNIST mnist = ExceptT $ do
  mnistdata <- T.readFile mnist
  return $ traverse (A.parseOnly parseMNIST) (tail $ T.lines mnistdata)

parseMNIST :: A.Parser (S ('D2 28 28), S ('D1 10))
parseMNIST = do
  Just lab <- oneHot <$> A.decimal
  pixels   <- many (A.char ',' >> A.double)
  image    <- maybe (fail "Parsed row was of an incorrect size") pure (fromStorable . V.fromList $ map realToFrac pixels)
  return (image, lab)


module Main
    ( main
    ) where

-------------------------------------------------------------------------------
import           Criterion
import           Criterion.Main
import           Data.Foldable
-------------------------------------------------------------------------------
import           Statistics.RollingAverage
-------------------------------------------------------------------------------

main :: IO ()
main = defaultMain [
  bgroup "ravgAdd" $ flip map [1, 100, 1000] $ \n ->
    bgroup (show n) [
        bench "Int" $ whnf (avgList n) ([1..fromIntegral n] :: [Int])
      , bench "Integer" $ whnf (avgList n) ([1..fromIntegral n] :: [Integer])
      ]
  , bgroup "ravg" $ flip map [1, 100, 1000] $ \n ->
      bgroup (show n) [
          bench "Int" $ whnf ravg' (avgList n ([1..fromIntegral n] :: [Int]))
        , bench "Integer" $ whnf ravg' (avgList n ([1..fromIntegral n] :: [Integer]))
      ]
  ]

-------------------------------------------------------------------------------
avgList :: Num a => Int -> [a] -> RollingAvg a
avgList lim = foldl' ravgAdd (mkRavg lim)

-------------------------------------------------------------------------------
ravg' :: (Integral a) => RollingAvg a -> Double
ravg' = ravg


{-# LANGUAGE OverloadedStrings #-}

import Control.Applicative ((<$>))
import Statistics.Sample.KernelDensity (kde)
import Text.Hastache (MuType(..), defaultConfig, hastacheFile)
import Text.Hastache.Context (mkStrContext)
import qualified Data.Attoparsec as B
import qualified Data.Attoparsec.Char8 as A
import qualified Data.ByteString as B
import qualified Data.ByteString.Lazy as L
import qualified Data.Vector.Unboxed as U

csv = do
  B.takeTill A.isEndOfLine
  (A.double `A.sepBy` A.char ',') `A.sepBy` A.endOfLine

main = do
  waits <- (either error (U.fromList . map last . filter (not.null)) .
            A.parseOnly csv) <$> B.readFile "data/faithful.csv"
  let xs = map (\(a,b) -> [a,b]) . U.toList . uncurry U.zip . kde 64 $ waits
      context "data" = MuVariable . show $ xs
  s <- hastacheFile defaultConfig "kde.tpl" (mkStrContext context)
  L.writeFile "kde.html" s