{-# language FlexibleInstances, FlexibleContexts, TypeFamilies #-}
{-# language ConstrainedClassMethods #-}
{-# language CPP #-}
-----------------------------------------------------------------------------
-- |
-- Module      :  Core.Numeric.BLAS.Class
-- Copyright   :  (c) Marco Zocca 2017-2018
-- License     :  see the file LICENSE
--
-- Maintainer  :  ocramz fripost org
-- Stability   :  experimental
-- Portability :  portable
--
-- Typeclasses for linear algebra and related concepts
--
-----------------------------------------------------------------------------
module Core.Numeric.BLAS.Class where

import Data.Complex
import Foreign.C.Types (CFloat, CDouble)

-- * Matrix and vector elements (optionally Complex)
--
-- Uses 'FlexibleContexts' and 'TypeFamilies'
class (Eq e , Fractional e, Floating e, Num (EltMag e), Ord (EltMag e)) => Elt e where
  type EltMag e :: *
  -- | Complex conjugate, or identity function if its input is real-valued
  conj :: e -> e
  conj = id
  -- | Magnitude
  mag :: e -> EltMag e

instance Elt Double where {type EltMag Double = Double ; mag = id}
instance Elt Float where {type EltMag Float = Float; mag = id}
instance (RealFloat e) => Elt (Complex e) where
  type EltMag (Complex e) = e
  conj = conjugate
  mag = magnitude

infixl 6 ^+^, ^-^

-- * Additive group
class AdditiveGroup v where
  -- | The zero element: identity for '(^+^)'
  zeroV :: v
  -- | Add vectors
  (^+^) :: v -> v -> v
  -- | Additive inverse
  negateV :: v -> v
  -- | Group subtraction
  (^-^) :: v -> v -> v
  (^-^) x y = x ^+^ negateV y

infixr 7 .*

-- * Vector space @v@.
class (AdditiveGroup v, Num (Scalar v)) => VectorSpace v where
  type Scalar v :: *
  -- | Scale a vector
  (.*) :: Scalar v -> v -> v

-- | Adds inner (dot) products.
class VectorSpace v => InnerSpace v where
  -- | Inner/dot product
  (<.>) :: v -> v -> Scalar v

-- | Inner product
dot :: InnerSpace v => v -> v -> Scalar v
dot = (<.>)

-- | Scale a vector by the reciprocal of a number (e.g. for normalization)
(./) :: (VectorSpace v, s ~ Scalar v, Fractional s) => v -> s -> v
v ./ s = recip s .* v

-- | Vector multiplied by scalar
(*.) :: (VectorSpace v, s ~ Scalar v) => v -> s -> v
(*.) = flip (.*)

-- | Convex combination of two vectors (NB: 0 <= `a` <= 1).
cvx :: VectorSpace v => Scalar v -> v -> v -> v
cvx a u v = a .* u ^+^ ((1-a) .* v)

-- ** Hilbert-space distance function

-- |`hilbertDistSq x y = || x - y ||^2` computes the squared L2 distance between two vectors
hilbertDistSq :: InnerSpace v => v -> v -> Scalar v
hilbertDistSq x y = t <.> t where
  t = x ^-^ y

-- * Normed vector spaces
--
-- Uses 'ConstrainedClassMethods'
class (InnerSpace v, Num (RealScalar v), Eq (RealScalar v), Epsilon (Magnitude v), Show (Magnitude v), Ord (Magnitude v)) => Normed v where
  type Magnitude v :: *
  type RealScalar v :: *
  -- | L1 norm
  norm1 :: v -> Magnitude v
  -- | Euclidean (L2) norm squared
  norm2Sq :: v -> Magnitude v
  -- | Lp norm (p > 0)
  normP :: RealScalar v -> v -> Magnitude v
  -- | Normalize w.r.t. Lp norm
  normalize :: RealScalar v -> v -> v
  -- | Normalize w.r.t. L2 norm
  normalize2 :: v -> v
  -- | Normalize w.r.t. norm2' instead of norm2
  normalize2' :: Floating (Scalar v) => v -> v
  normalize2' x = x ./ norm2' x
  -- | Euclidean (L2) norm
  norm2 :: Floating (Magnitude v) => v -> Magnitude v
  norm2 x = sqrt (norm2Sq x)
  -- | Euclidean (L2) norm; returns a Complex (norm :+ 0) for Complex-valued vectors
  norm2' :: Floating (Scalar v) => v -> Scalar v
  norm2' x = sqrt $ x <.> x
  -- | Lp norm (p > 0)
  norm :: Floating (Magnitude v) => RealScalar v -> v -> Magnitude v
  norm p v
    | p == 1 = norm1 v
    | p == 2 = norm2 v
    | otherwise = normP p v

-- | Infinity-norm (Real)
normInftyR :: (Foldable t, Ord a) => t a -> a
normInftyR x = maximum x

-- | Infinity-norm (Complex)
normInftyC :: (Foldable t, RealFloat a, Functor t) => t (Complex a) -> a
normInftyC x = maximum (magnitude <$> x)

-- | Lp inner product (p > 0)
dotLp :: (Set t, Foldable t, Floating a) => a -> t a -> t a ->  a
dotLp p v1 v2 = sum u**(1/p) where
  f a b = (a*b)**p
  u = liftI2 f v1 v2

-- | Reciprocal
reciprocal :: (Functor f, Fractional b) => f b -> f b
reciprocal = fmap recip

-- |Scale
scale :: (Num b, Functor f) => b -> f b -> f b
scale n = fmap (* n)

-- * Matrix ring

-- | A matrix ring is any collection of matrices over some ring R that form a ring under matrix addition and matrix multiplication

class (AdditiveGroup m, Epsilon (MatrixNorm m)) => MatrixRing m where
  type MatrixNorm m :: *
  -- | Matrix-matrix product
  (##) :: m -> m -> m
  -- | Matrix times matrix transpose (A B^T)
  (##^) :: m -> m -> m
  -- | Matrix transpose times matrix (A^T B)
  (#^#) :: m -> m -> m
  a #^# b = transpose a ## b
  -- | Matrix transpose (Hermitian conjugate in the Complex case)
  transpose :: m -> m
  -- | Frobenius norm
  normFrobenius :: m -> MatrixNorm m

-- * Linear vector space

class (VectorSpace v {-, MatrixRing (MatrixType v)-}) => LinearVectorSpace v where
  type MatrixType v :: *
  -- | Matrix-vector action
  (#>) :: MatrixType v -> v -> v
  -- | Dual matrix-vector action
  (<#) :: v -> MatrixType v -> v

-- * Typeclasses for internal use (not to be exposed to the user)

-- ** FiniteDim : finite-dimensional objects
class FiniteDim f where
  type FDSize f
  -- | Dimension (i.e. Int for SpVector, (Int, Int) for SpMatrix)
  dim :: f -> FDSize f

-- ** HasData : accessing inner data (do not export)
class HasData f where
  type HDData f
  -- | Number of nonzeros
  nnz :: f -> Int
  dat :: f -> HDData f

-- ** Sparse : sparse datastructures
class (FiniteDim f, HasData f) => Sparse f where
  -- | Sparsity (fraction of nonzero elements)
  spy :: Fractional b => f -> b

-- ** Set : types that behave as sets
--
-- Example : sparse vector addition operates on the union of the nonzero indices, whereas the sparse vector inner product uses
class Functor f => Set f where
  -- | Union binary lift : apply function on _union_ of two "sets"
  liftU2 :: (a -> a -> a) -> f a -> f a -> f a
  -- | Intersection binary lift : apply function on _intersection_ of two "sets"
  liftI2 :: (a -> a -> b) -> f a -> f a -> f b

-- * instances

-- | Instances for builtin types
--
-- Uses 'CPP'
#define ScalarType(t) \
instance AdditiveGroup (t) where {zeroV = 0; (^+^) = (+); negateV = negate};\
instance VectorSpace (t) where {type Scalar (t) = t; (.*) = (*) };

-- ScalarType(Int)
-- ScalarType(Integer)
ScalarType(Float)
ScalarType(Double)
ScalarType(Complex Float)
ScalarType(Complex Double)

#undef ScalarType

instance InnerSpace Float  where {(<.>) = (*)}
instance InnerSpace Double where {(<.>) = (*)}
instance InnerSpace (Complex Float)  where {x <.> y = x * conjugate y}
instance InnerSpace (Complex Double) where {x <.> y = x * conjugate y}

#define SimpleNormedInstance(t) \
instance Normed (t) where {type Magnitude (t) = t; type RealScalar (t) = t;\
 norm1 = abs; norm2Sq = (**2); normP _ = abs; normalize _ = signum;\
 normalize2 = signum; normalize2' = signum; norm2 = abs; norm2' = abs; norm _ = abs};

SimpleNormedInstance(Float)
SimpleNormedInstance(Double)

#undef SimpleNormedInstance

#define ComplexNormedInstance(t) \
instance Normed (Complex (t)) where {\
 type Magnitude  (Complex (t)) = t;\
 type RealScalar (Complex (t)) = t;\
 norm1   (r :+ i) = abs r + abs i;\
 norm2Sq (r :+ i) = r*r + i*i;\
 normP p (r :+ i) = (r**p + i**p)**(1/p);\
 normalize p c = toC (1 / normP p c) * c;\
 normalize2  c = (1 / norm2' c) * c;\
 norm2  = magnitude;\
 norm2' = toC . magnitude;};

ComplexNormedInstance(Float)
ComplexNormedInstance(Double)

#undef ComplexNormedInstance

-- * Utilities

-- | Lift a real number onto the complex plane
toC :: Num a => a -> Complex a
toC r = r :+ 0

-- | Provides a test to see if a quantity is near zero.
--
-- Appears originally in ekmett/linear
--
-- TODO[ocramz] Uses 'FlexibleInstances', so might be a good idea to move it in a separate module
--
-- >>> nearZero (1e-11 :: Double)
-- False
--
-- >>> nearZero (1e-17 :: Double)
-- True
--
-- >>> nearZero (1e-5 :: Float)
-- False
--
-- >>> nearZero (1e-7 :: Float)
-- True
class (Floating a, Num a) => Epsilon a where
  -- | Determine if a quantity is near zero.
  nearZero :: a -> Bool

-- | @'abs' a '<=' 1e-6@
instance Epsilon Float where
  nearZero a = abs a <= 1e-6

-- | @'abs' a '<=' 1e-12@
instance Epsilon Double where
  nearZero a = abs a <= 1e-12

-- | @'abs' a '<=' 1e-6@
instance Epsilon CFloat where
  nearZero a = abs a <= 1e-6

-- | @'abs' a '<=' 1e-12@
instance Epsilon CDouble where
  nearZero a = abs a <= 1e-12

-- | @'magnitude' a '<=' 1e-6@
instance Epsilon (Complex Float) where
  nearZero a = magnitude a <= 1e-6

-- | @'magnitude' a '<=' 1e-12@
instance Epsilon (Complex Double) where
  nearZero a = magnitude a <= 1e-12

-- | @'magnitude' a '<=' 1e-6@
instance Epsilon (Complex CFloat) where
  nearZero a = magnitude a <= 1e-6

-- | @'magnitude' a '<=' 1e-12@
instance Epsilon (Complex CDouble) where
  nearZero a = magnitude a <= 1e-12

-- * Rounding operations

-- | Is this quantity close to 1 ?
nearOne :: Epsilon a => a -> Bool
nearOne x = nearZero (1 - x)

-- | Is this quantity distinguishable from 0 ?
isNz :: Epsilon a => a -> Bool
isNz x = not (nearZero x)

withDefault :: (t -> Bool) -> t -> t -> t
withDefault q d x | q x = d
                  | otherwise = x

roundZero, roundOne, roundZeroOne :: Epsilon a => a -> a
roundZero = withDefault nearZero (fromIntegral (0 :: Int))
roundOne = withDefault nearOne (fromIntegral (1 :: Int))

with2Defaults :: (t -> Bool) -> (t -> Bool) -> t -> t -> t -> t
with2Defaults q1 q2 d1 d2 x | q1 x = d1
                            | q2 x = d2
                            | otherwise = x

-- | Round to respectively 0 or 1
roundZeroOne = with2Defaults nearZero nearOne (fromIntegral (0 :: Int)) (fromIntegral (1 :: Int))


module STCR2Z1T0Image 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           Data.Vector.Generic          as VG
import           Data.Vector.Storable         as VS
import           DFT.Plan
import           Filter.Gaussian
import           Filter.Pinwheel
import           FokkerPlanck.DomainChange
import           FokkerPlanck.MonteCarlo
import           FokkerPlanck.Pinwheel
import           Image.IO
import           Image.Transform
import           STC
import           STC.OrientationScaleAnalysis
import           System.Directory
import           System.Environment
import           System.FilePath
import           System.Random
import           Types
import           Utils.Array

main = do
  args@(numPointStr:numOrientationStr:sigmaStr:taoStr:lenStr:initialScaleStr:numTrailStr:maxTrailStr:theta0FreqsStr:thetaFreqsStr:histFilePath:alphaStr:pinwheelFlagStr:imagePath: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]
      alpha = read alphaStr :: Double
      pinwheelFlag = read pinwheelFlagStr :: Bool
      numIteration = read numIterationStr :: Int
      writeSourceFlag = read writeSourceFlagStr :: Bool
      numThread = read numThreadStr :: Int
      folderPath = "output/test/STCR2Z1T0Image"
  createDirectoryIfMissing True folderPath
  flag <- doesFileExist histFilePath
  (ImageRepa _ img'') <- readImageRepa imagePath False
  let img'
        -- resize25D (numPoint, numPoint) (0, 255) .
        -- R.backpermute (Z :. nf :. numPoint :. numPoint) id .
        -- upsample [4, 4, 1] . downsample [4, 4, 1] $
       = img''
      (Z :. nf :. _ :. _) = extent img''
      (Z :. _ :. cols :. rows) = extent img'
      pinwheelParams = PinwheelParams rows cols alpha (exp 1) theta0Freqs [0]
  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
          1
          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.5
      thetaFreqs
      theta0Freqs
  plan <- makeR2Z1T0Plan emptyPlan arrR2Z1T0
  -- (plan1, imgF) <- makeImagePlan plan0 . computeS $ img'
  -- (plan2, filterF, filterPIF) <- pinwheelFilter plan1 pinwheelParams
  -- (plan, gaussianFilterF) <-
  --   gaussian2DFilter plan2 (Gaussian2DParams 1 rows cols)
  -- -- img <- convolveGaussian2D plan gaussianFilterF . R.map (:+ 0) $ img'
  -- -- imgVec <- dropPixel (2 / 3) . toUnboxed . computeS . filterImage $ img'
  let img = R.map (:+ 0) img' -- = . fromUnboxed (extent img') $ imgVec
  plotImageRepa
    (folderPath </> "input.png") -- (ImageRepa 8 img)
    (ImageRepa 8 . computeS . R.map magnitude $ img)
  -- convolvedImg' <-
  --   convolvePinwheel
  --     plan
  --     filterPIF
  --     (R.slice imgF (Z :. (0 :: Int) :. All :. All))
  -- let convolvedImg =
  --       computeS . R.slice convolvedImg' $
  --       (Z :. All :. (0 :: Int) :. All :. All)
  -- let (initialDistSource, initialDistSink) =
  --       analyzeOrientation numOrientation theta0Freqs convolvedImg
  powerMethod1
    plan
    folderPath
    cols
    rows
    numOrientation
    thetaFreqs
    theta0Freqs
    arrR2Z1T0
    numIteration
    writeSourceFlag
    ""
    0.1
    (R.traverse img (const (Z :. (L.length theta0Freqs) :. cols :. rows)) $ \f (Z :. _ :. i :. j) ->
       if magnitude (f (Z :. (0 :: Int) :. i :. j)) > 0
         then 1
         else 0)
    (R.traverse
       img
       (const
          (Z :. (L.length thetaFreqs) :. (L.length theta0Freqs) :. cols :. rows)) $ \f (Z :. k :. _ :. i :. j) ->
       if k == div (L.length thetaFreqs) 2
         then f (Z :. (0 :: Int) :. i :. j) /
              (fromIntegral $ L.length theta0Freqs)
         else 0)


module Wigner where

import Operators
import BraKet
import LinAlg
import QUtil

import Data.Complex
import Data.List
import Data.List.Split

import qualified Data.Vector.Unboxed as U
import qualified Data.Vector.Generic.Mutable as M
import Numeric.FFT.Vector.Invertible

-- Expected value of displacement operator
-- for a given density matrix
-- Params:
-- - Number of terms to use for displacement operator expansion
-- - Displacement parameter
-- - Probabilities per state
-- - Basis in which to evaluate trace
-- - States for which to compute density matrix
expectDisp :: Int -> Complex Double
                -> [Complex Double] -> [Ket (Complex Double)] -> Complex Double
expectDisp n alpha = expectedValue (displacement n alpha)

-- 2D Wigner transform
-- fourier transform all rows, then inversely FT all columns
-- Then normalize
-- unfortunately this doesn't work that well, so manual transform is
-- implemented below with crude riemann integration
wt2d:: [[Complex Double]] -> [[Complex Double]]
wt2d = map ((/pi^2)<$>) . transpose . map (U.toList . run dft . U.fromList) .
    transpose . map (U.toList . run idft . U.fromList)

-- manual W transform
wignerTransform' :: Complex Double -> Double -> [Complex Double] -> [Complex Double] -> Complex Double
wignerTransform' alpha dbeta betas ds = sum $ map (*(dbeta :+ 0.0)) integrand
    where
        integrand = zipWith (*) ds $ (\beta -> exp $ alpha * conjugate beta - conjugate alpha * beta) <$> betas

wignerTransform :: [Complex Double] -> Double -> [Complex Double] -> [Complex Double] -> [Complex Double]
wignerTransform alphas dbeta betas ds = map (/pi^2) $ (\a -> wignerTransform' a dbeta betas ds) <$> alphas


module FinancialTimeseries.Optimize.Optimize where

import Data.Function (on)
import Data.Distributive (Distributive)

import qualified Data.List as List

import qualified Data.Vector as Vec

import Data.Maybe (catMaybes)

import qualified Statistics.Test.StudentT as StudentT
import qualified Statistics.Test.Types as TTypes
import qualified Statistics.Types as Types

import FinancialTimeseries.Algorithm.Evaluate (Profit, long)

import FinancialTimeseries.Statistics.Statistics (tradeStatistics, trMeanYield, moments, sampleSize, yield)

import FinancialTimeseries.Type.Long (Long(..))
import qualified FinancialTimeseries.Type.Strategy as Strategy
import FinancialTimeseries.Type.Strategy (Strategy(Strategy))
import FinancialTimeseries.Type.Timeseries (TimeseriesRaw, slice)
import FinancialTimeseries.Type.Types (TradeYield(..), Invested(..), partitionInvested)

data Metrics a = Metrics {
  name :: String
  , metrics :: a
  } deriving (Show, Eq, Ord)

data Config optParams price a = Config {
  strategy :: [Strategy optParams price a]
  , evaluationMetrics :: Vec.Vector a -> Metrics a
  }

timeseriesYield :: (Real a, Floating a) => Vec.Vector a -> Metrics a
timeseriesYield vs =
  let stats = tradeStatistics vs
  in Metrics {
    name = "timeseriesYield"
    , metrics = trMeanYield (moments stats) ** fromIntegral (sampleSize stats)
    }

tradeMeanYield :: (Real a, Floating a) => Vec.Vector a -> Metrics a
tradeMeanYield vs =
  let stats = tradeStatistics vs
  in Metrics {
    name = "tradeMeanYield"
    , metrics = trMeanYield (moments stats)
    }

evalStrategy ::
  (Distributive price, Profit price, Fractional a) =>
  TimeseriesRaw price a
  -> Strategy optParams price a
  -> Long (TradeYield (Invested (Vec.Vector a)))
evalStrategy ts (Strategy _ ps stgy) =
  let t = stgy ps ts
      lg = long (partitionInvested (slice t))
      k = fmap (Vec.fromList . map (yield . Vec.map snd))
  in fmap (fmap (snd . fmap k)) lg

equalDistribution ::
  (Real a) =>
  Int -> Types.PValue Double -> Long (TradeYield (Invested (Vec.Vector a))) -> Bool
equalDistribution n pval (Long (TradeYield (Invested vs))) =
  let x = ceiling ((fromIntegral (Vec.length vs) / fromIntegral n) :: Double)

      f xs | Vec.length xs == 0 = []
      f xs = chunk xs

      chunk = (\(as, bs) -> as : f bs) . Vec.splitAt x
      cs = chunk (Vec.map realToFrac vs)

      ds = catMaybes (zipWith (StudentT.studentTTest TTypes.SamplesDiffer) cs (tail cs))
      es = map (TTypes.isSignificant pval) ds

  in all (==TTypes.NotSignificant) es

optimize ::
  (Distributive price, Profit price, Fractional a, Real a, Floating a) =>
  Config optParams price a -> TimeseriesRaw price a -> [(optParams, Metrics a)]
optimize (Config strgies metr) ts =
  let ss = map (\stgy -> (Strategy.parameters stgy, evalStrategy ts stgy)) strgies

      n = 10
      pval = Types.mkPValue 0.05

      bs = filter (equalDistribution n pval . snd) ss

      f (Long (TradeYield (Invested vs))) = metr vs
      cs = map (fmap f) bs

      ds = List.sortBy (compare `on` snd) cs
  in ds


module Pineappl.Quantum where

import           Data.Bifunctor
import           Data.Complex
import qualified Data.Map                      as M
import           Pineappl                       ( Prob(P)
                                                , LogProb(LP)
                                                )
import           Pineappl.Deterministic         ( runBDDist
                                                , bddist
                                                , hist
                                                , WrappedBDDist
                                                )

fromAmplitude
  :: (RealFloat a, Ord b)
  => WrappedBDDist (LogProb (Complex a)) b
  -> WrappedBDDist (Prob a) b
fromAmplitude =
  bddist
    . fmap (second (\(LP psi) -> P (magnitude psi ^ 2)))
    . M.toList
    . hist
    . runBDDist


module Numeric.Sundials.ODEOpts where

import           Data.Int (Int32)
import qualified Data.Vector.Storable as VS

import           Numeric.LinearAlgebra.HMatrix (Vector, Matrix)

type Jacobian = Double -> Vector Double -> Matrix Double

-- | Adaptive step-size control
-- functions.
--
-- [GSL](https://www.gnu.org/software/gsl/doc/html/ode-initval.html#adaptive-step-size-control)
-- allows the user to control the step size adjustment using
-- \(D_i = \epsilon^{abs}s_i + \epsilon^{rel}(a_{y} |y_i| + a_{dy/dt} h |\dot{y}_i|)\) where
-- \(\epsilon^{abs}\) is the required absolute error, \(\epsilon^{rel}\)
-- is the required relative error, \(s_i\) is a vector of scaling
-- factors, \(a_{y}\) is a scaling factor for the solution \(y\) and
-- \(a_{dydt}\) is a scaling factor for the derivative of the solution \(dy/dt\).
--
-- [CVode](https://computation.llnl.gov/projects/sundials/cvode)
-- allows the user to control the step size adjustment using
-- \(\eta^{rel}|y_i| + \eta^{abs}_i\). For compatibility with
-- [hmatrix-gsl](https://hackage.haskell.org/package/hmatrix-gsl),
-- tolerances for \(y\) and \(\dot{y}\) can be specified but the latter have no
-- effect.
data StepControl = X     Double Double -- ^ absolute and relative tolerance for \(y\); in GSL terms, \(a_{y} = 1\) and \(a_{dy/dt} = 0\); in ARKode terms, the \(\eta^{abs}_i\) are identical
                 | X'    Double Double -- ^ absolute and relative tolerance for \(\dot{y}\); in GSL terms, \(a_{y} = 0\) and \(a_{dy/dt} = 1\); in ARKode terms, the latter is treated as the relative tolerance for \(y\) so this is the same as specifying 'X' which may be entirely incorrect for the given problem
                 | XX'   Double Double Double Double -- ^ include both via relative tolerance
                                                     -- scaling factors \(a_y\), \(a_{{dy}/{dt}}\); in ARKode terms, the latter is ignored and \(\eta^{rel} = a_{y}\epsilon^{rel}\)
                 | ScXX' Double Double Double Double (Vector Double) -- ^ scale absolute tolerance of \(y_i\); in ARKode terms, \(a_{{dy}/{dt}}\) is ignored, \(\eta^{abs}_i = s_i \epsilon^{abs}\) and \(\eta^{rel} = a_{y}\epsilon^{rel}\)
  deriving (Eq, Ord, Show, Read)

-- | Stepping functions
data ODEMethod = ADAMS
               | BDF
  deriving (Eq, Ord, Show, Read)

data ODEOpts = ODEOpts {
    maxNumSteps :: Int32
  , minStep     :: Double
  , maxFail     :: Int32
  , odeMethod   :: ODEMethod
  , stepControl :: StepControl
  , initStep    :: Maybe Double
    -- ^ initial step size - by default, CVode
    -- estimates the initial step size to be the
    -- solution \(h\) of the equation
    -- \(\|\frac{h^2\ddot{y}}{2}\| = 1\), where
    -- \(\ddot{y}\) is an estimated value of the second
    -- derivative of the solution at \(t_0\)
  } deriving (Read, Show, Eq, Ord)

data SundialsDiagnostics = SundialsDiagnostics {
    odeGetNumSteps               :: Int
  , odeGetNumStepAttempts        :: Int
  , odeGetNumRhsEvals_fe         :: Int
  , odeGetNumRhsEvals_fi         :: Int
  , odeGetNumLinSolvSetups       :: Int
  , odeGetNumErrTestFails        :: Int
  , odeGetNumNonlinSolvIters     :: Int
  , odeGetNumNonlinSolvConvFails :: Int
  , dlsGetNumJacEvals            :: Int
  , dlsGetNumRhsEvals            :: Int
  } deriving Show


{-# LANGUAGE KindSignatures #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE BangPatterns #-}

module UntypedNet where

import Data.Functor
import Data.Maybe
import Data.List (foldl')
import Numeric.LinearAlgebra (Matrix, Vector, randomVector, uniformSample, RandDist(Uniform), (#>), tr, outer, scale, norm_2, vector)
import Control.Monad.Random (MonadRandom, getRandom, replicateM)

-- Wx + b
data Weights = W { wBiases :: !(Vector Double)  -- n
                 , wNodes :: !(Matrix Double)  -- n x m
                 }  -- "m to n" layer

-- define a network data type
-- can build a network like : ih :&~ hh :&~ O ho
data Network :: * where
    O :: !Weights -> Network
    (:&~) :: !Weights -> !Network -> Network

infixr 5 :&~

{- equivalent to:
data Network = O Weights
             | Weights :&~ Network
-}

randomWeights :: MonadRandom m => Int -> Int -> m Weights
randomWeights i o = do
  seed1 :: Int <- getRandom  -- generate random seed
  seed2 :: Int <- getRandom
      -- randomVector :: Seed -> RandDist -> Int -> VectorDouble
      -- obtains a vector of random elements of provided size
  let wB = randomVector seed1 Uniform o * 2 - 1
      -- uniformSample :: Seed -> Int -> [(Double, Double)] -> Matrix Double
      wN = uniformSample seed2 o (replicate i (-1, 1))  -- o x i -sized matrix
  return $ W wB wN

-- build a random network composed of layers
-- input layer size -> [hidden layer sizes] -> output layer size
randomNet :: MonadRandom m => Int -> [Int] -> Int -> m Network
randomNet i [] o = O <$> randomWeights i o  -- final fully-connected layer
randomNet i (h:hs) o = (:&~) <$> randomWeights i h <*> randomNet h hs o

-- the logistic (sigmoid) function
logistic :: Floating a => a -> a
logistic x = 1 / (1 + exp (-x))

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

-- matrix-vector multiplication
-- (#>) :: Numeric t => Matrix t -> Vector t -> Vector t
runLayer :: Weights -> Vector Double -> Vector Double
runLayer (W wB wN) v = wB + wN #> v

-- feedforward the network with every activaciton function as sigmoid
runNet :: Network -> Vector Double -> Vector Double
runNet (O w) !v = logistic (runLayer w v)
runNet (w :&~ n') !v = let v' = logistic (runLayer w v)
                       in runNet n' v'

train :: Double  -- learning rate
      -> Vector Double  -- input vector
      -> Vector Double  -- target vector
      -> Network  -- network to train
      -> Network
train rate x0 target = fst . go x0
  where
    go :: Vector Double  -- input vector
       -> Network  -- network to train
       -> (Network, Vector Double)  -- (new network layer, back-propping derivative chain)
    -- handle the output layer
    go !x (O w@(W wB wN))
      = let y = runLayer w x
            o = logistic y
            -- the gradient (dE / dy)
            dEdy = logistic' y * (o - target)
            -- bias and node update
            wB' = wB - scale rate dEdy  -- multiply by scalar 'rate'
            wN' = wN - scale rate (dEdy `outer` x)
            w' = W wB' wN'
            -- bundle of derivatives for next step
            dWs = tr wN #> dEdy
         in (O w', dWs)
    -- handle the inner layers
    go !x (w@(W wB wN) :&~ n)
        = let y = runLayer w x
              o = logistic y
              -- get dWs', bundle of derivatives from rest of the net
              (n', dWs') = go o n
              -- the gradient (how much y affects the error)
              dEdy = logistic' y * dWs'
              -- new bias weights and node weights
              wB' = wB - scale rate dEdy
              wN' = wN - scale rate (dEdy `outer` x)
              w' = W wB' wN'
              -- bundle of derivatives for next step
              dWs = tr wN #> dEdy
          in (w' :&~ n', dWs)

netTest :: MonadRandom m => Double -> Int -> m String
netTest rate n = do
    inps <- replicateM n $ do
      s <- getRandom
      return $ randomVector s Uniform 2 * 2 - 1
    let outs = flip map inps $ \v ->
                 if v `inCircle` (fromRational 0.33, 0.33)
                      || v `inCircle` (fromRational (-0.33), 0.33)
                   then fromRational 1
                   else fromRational 0
    net0 <- randomNet 2 [16,8] 1
    let trained = foldl' trainEach net0 (zip inps outs)
          where
            trainEach :: Network -> (Vector Double, Vector Double) -> Network
            trainEach nt (i, o) = train rate i o nt

        outMat = [ [ render (norm_2 (runNet trained (vector [x / 25 - 1,y / 10 - 1])))
                   | x <- [0..50] ]
                 | y <- [0..20] ]
        render r | r <= 0.2  = ' '
                 | r <= 0.4  = '.'
                 | r <= 0.6  = '-'
                 | r <= 0.8  = '='
                 | otherwise = '#'

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

(!!?) :: [a] -> Int -> Maybe a
xs !!? i = listToMaybe (drop i xs)


{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE Strict       #-}
module STC.PowerMethod where

import           Control.Monad                  as M
import qualified Data.Array.Accelerate          as A
import           Data.Array.Accelerate.LLVM.PTX as A
import           Data.Array.IArray              as IA
import           Data.Array.Repa                as R
import           Data.Complex
import           Data.List                      as L
import           Data.Vector.Storable           as VS
import           Data.Vector.Unboxed            as VU
import           DFT.Plan
import           Filter.Utils
import           FokkerPlanck.FourierSeries
import           FourierMethod.BlockMatrixAcc
import           FourierMethod.FourierSeries2D
import           Graphics.Gnuplot.Simple
import           Image.IO
import           STC.CompletionField
import           STC.Convolution
import           STC.DFTArray
import           STC.Multiplication
import           STC.Plan
import           STC.Reversal
import           STC.Utils
import           System.FilePath                ((</>))
import           Text.Printf
import           Utils.Array
import           Utils.List
import           Utils.Parallel
import           Utils.Time
import FourierPinwheel as FP
import           Foreign.CUDA.Driver            as CUDA

powerMethod ::
     DFTPlan
  -> FilePath
  -> Bool
  -> R.Array U DIM4 (Complex Double)
  -> IA.Array (Int, Int) (VS.Vector (Complex Double))
  -> VS.Vector (Complex Double)
  -> Int
  -> DFTArray
  -> IO DFTArray
powerMethod _ _ _ _ _ _ 0 !arr = return arr
powerMethod !plan !folderPath !writeFlag !coefficients !harmonicsArray !bias !numStep !input@(DFTArray rows cols _ _ _) = do
  printCurrentTime (show numStep)
  convolvedArr <-
    STC.Convolution.convolve Source plan coefficients harmonicsArray input
  let !biasedConvolvedArr = parMapDFTArray (VS.zipWith (*) bias) convolvedArr
      !maxMag =
        L.sum .
        parMap rdeepseq (VS.sum . VS.map (\x -> (magnitude x)^2)) . getDFTArrayVector $
        biasedConvolvedArr
      !normalizedBiasedConvolvedArr =
        parMapDFTArray (VS.map (/ (maxMag :+ 0))) biasedConvolvedArr
  when writeFlag .
    plotDFTArrayPower
      (folderPath </> (printf "Source_%03d.png" numStep))
      cols
      rows $
    convolvedArr
  powerMethod
    plan
    folderPath
    writeFlag
    coefficients
    harmonicsArray
    bias
    (numStep - 1)
    normalizedBiasedConvolvedArr

{-# INLINE computeContour #-}
computeContour ::
     DFTPlan
  -> FilePath
  -> Bool
  -> R.Array U DIM4 (Complex Double)
  -> IA.Array (Int, Int) (VS.Vector (Complex Double))
  -> VS.Vector (Complex Double)
  -> Int
  -> String
  -> DFTArray
  -> IO DFTArray
computeContour !plan !folderPath !writeFlag !coefficients !harmonicsArray !bias !numIteration !suffix !input@(DFTArray rows cols _ _ _) = do
  eigenVector <-
    powerMethod
      plan
      folderPath
      writeFlag
      coefficients
      harmonicsArray
      bias
      numIteration
      input
  source <- STC.Convolution.convolve Source plan coefficients harmonicsArray eigenVector
  sink <- STC.Convolution.convolve Sink plan coefficients harmonicsArray eigenVector
  completion <- completionField plan source sink
  plotDFTArrayPower (folderPath </> "SinkPower.png") rows cols sink
  plotDFTArrayPower
    (folderPath </> (printf "CompletionPower_%s.png" suffix))
    rows
    cols
    completion
  return completion

powerMethod' ::
     DFTPlan
  -> FilePath
  -> Bool
  -> R.Array U DIM4 (Complex Double)
  -> IA.Array (Int, Int) (VS.Vector (Complex Double))
  -> VS.Vector (Complex Double)
  -> Int -> [[(Complex Double)]]
  -> DFTArray
  -> IO DFTArray
powerMethod' _ _ _ _ _ _ 0 _ !arr = return arr
powerMethod' !plan !folderPath !writeFlag !coefficients !harmonicsArray !bias !numStep thetaRHarmonics !input@(DFTArray rows cols thetaFreqs rFreqs _) = do
  printCurrentTime (show numStep)
  convolvedArr <- convolve' Source plan coefficients harmonicsArray input
  let -- sigma = 15
      -- gaussian =
      --   L.map (\x -> (exp (-(x ^ 2) / (2 * sigma ^ 2))) :+ 0) thetaFreqs
      convolvedArr' =
        fromUnboxed (Z :. 72 :. cols :. rows) .
        VS.convert .
        VS.concat .
        computeFourierSeriesThetaR thetaRHarmonics . getDFTArrayVector $
        convolvedArr
      -- convolvedArr1 =
      --   DFTArray rows cols thetaFreqs rFreqs .
      --   L.zipWith (\g vec -> VS.map (* g) vec) gaussian . getDFTArrayVector $
      --   convolvedArr
  plotThetaDimension folderPath (printf "FourierSeries_%d_" numStep) (29, 12) .
    R.map magnitude $
    convolvedArr'
  let !biasedConvolvedArr = parMapDFTArray (VS.zipWith (*) bias) convolvedArr
      !maxMag =
        L.sum .
        parMap rdeepseq (VS.sum . VS.map (\x -> (magnitude x) ^ 2)) .
        getDFTArrayVector $
        biasedConvolvedArr
      !normalizedBiasedConvolvedArr =
        parMapDFTArray (VS.map (/ (maxMag :+ 0))) biasedConvolvedArr
  when writeFlag .
    plotDFTArrayPower
      (folderPath </> (printf "Source_%03d.png" numStep))
      cols
      rows $
    convolvedArr
  powerMethod'
    plan
    folderPath
    writeFlag
    coefficients
    harmonicsArray
    bias
    (numStep - 1)
    thetaRHarmonics
    normalizedBiasedConvolvedArr

{-# INLINE computeContour' #-}
computeContour' ::
     DFTPlan
  -> FilePath
  -> Bool
  -> R.Array U DIM4 (Complex Double)
  -> IA.Array (Int, Int) (VS.Vector (Complex Double))
  -> VS.Vector (Complex Double)
  -> Int
  -> String -> [[Complex Double]]
  -> DFTArray
  -> IO DFTArray
computeContour' !plan !folderPath !writeFlag !coefficients !harmonicsArray !bias !numIteration !suffix thetaRHarmonics !input@(DFTArray rows cols _ _ _) = do
  eigenVector <-
    powerMethod'
      plan
      folderPath
      writeFlag
      coefficients
      harmonicsArray
      bias
      numIteration
      thetaRHarmonics
      input
  source <- convolve' Source plan coefficients harmonicsArray eigenVector
  sink <- convolve' Sink plan coefficients harmonicsArray eigenVector
  completion <- completionField' plan source sink
  plotDFTArrayPower (folderPath </> "Sink.png") rows cols sink
  plotDFTArrayPower
    (folderPath </> (printf "CompletionPower_%s.png" suffix))
    rows
    cols
    completion
  let point = (32, -4)
      source' =
        fromUnboxed (Z :. 72 :. cols :. rows) .
        VS.convert .
        VS.concat .
        computeFourierSeriesThetaR thetaRHarmonics .
        getDFTArrayVector $
        source
      sink' =
        fromUnboxed (Z :. 72 :. cols :. rows) .
        VS.convert .
        VS.concat .
        computeFourierSeriesThetaR thetaRHarmonics .
        getDFTArrayVector $
        sink
      completion' =
        fromUnboxed (Z :. 72 :. cols :. rows) .
        VS.convert .
        VS.concat .
        computeFourierSeriesThetaR thetaRHarmonics . getDFTArrayVector $
        completion
  plotThetaDimension folderPath ("FourierSeries_Completion_") point .
    R.map magnitude $
    completion'
  plotThetaDimension
    folderPath
    ("FourierSeries_Source_")
    point .
    R.map magnitude $
    source'
  plotThetaDimension
    folderPath
    ("FourierSeries_Sink_")
    point .
    R.map magnitude $
    sink'
  return completion

powerMethodG' ::
     DFTPlan
  -> FilePath
  -> Bool
  -> R.Array U DIM4 (Complex Double)
  -> IA.Array (Int, Int) (VS.Vector (Complex Double))
  -> VS.Vector (Complex Double)
  -> VS.Vector (Complex Double)
  -> Int
  -> DFTArray
  -> IO DFTArray
powerMethodG' _ _ _ _ _ _ _ 0 !arr = return arr
powerMethodG' !plan !folderPath !writeFlag !coefficients !harmonicsArray !bias !gaussian !numStep !input@(DFTArray rows cols _ _ _) = do
  printCurrentTime (show numStep)
  convolvedArr <- convolveG' Source plan coefficients harmonicsArray gaussian input
  let !biasedConvolvedArr = parMapDFTArray (VS.zipWith (*) bias) convolvedArr
      !maxMag =
        L.sum .
        parMap rdeepseq (VS.sum . VS.map (\x -> (magnitude x)^2)) .
        getDFTArrayVector $
        biasedConvolvedArr
      !normalizedBiasedConvolvedArr =
        parMapDFTArray (VS.map (/ (maxMag :+ 0))) biasedConvolvedArr
  when writeFlag .
    plotDFTArrayPower
      (folderPath </> (printf "Source_%03d.png" numStep))
      cols
      rows $
    convolvedArr
  powerMethodG'
    plan
    folderPath
    writeFlag
    coefficients
    harmonicsArray
    bias
    gaussian
    (numStep - 1)
    normalizedBiasedConvolvedArr

{-# INLINE computeContourG' #-}
computeContourG' ::
     DFTPlan
  -> FilePath
  -> Bool
  -> R.Array U DIM4 (Complex Double)
  -> IA.Array (Int, Int) (VS.Vector (Complex Double))
  -> VS.Vector (Complex Double)
  -> VS.Vector (Complex Double)
  -> Int
  -> String
  -> DFTArray
  -> IO DFTArray
computeContourG' !plan !folderPath !writeFlag !coefficients !harmonicsArray !bias !gaussian !numIteration !suffix !input@(DFTArray rows cols _ _ _) = do
  eigenVector <-
    powerMethodG'
      plan
      folderPath
      writeFlag
      coefficients
      harmonicsArray
      bias
      gaussian
      numIteration
      input
  source <- convolveG' Source plan coefficients harmonicsArray gaussian eigenVector
  sink <- convolveG' Sink plan coefficients harmonicsArray gaussian eigenVector
  completion <- completionField' plan source sink
  plotDFTArrayPower (folderPath </> "Sink.png") rows cols sink
  plotDFTArrayPower
    (folderPath </> (printf "CompletionPower_%s.png" suffix))
    rows
    cols
    completion
  return completion

powerMethodSparse ::
     Field
  -> R.Array U DIM4 (Complex Double)
  -> IA.Array (Int, Int) (R.Array U DIM2 (Complex Double))
  -> R.Array U DIM1 Double
  -> R.Array U DIM1 Double
  -> Double
  -> [(Double, Double)]
  -> Int
  -> [R.Array U DIM2 (Complex Double)]
  -> [R.Array U DIM2 (Complex Double)]
powerMethodSparse _ _ _ _ _ _ _ 0 !arr = arr
powerMethodSparse !field !coefficients !harmonicsArray !thetaFreqs !rFreqs !cutoff !xs !numStep !input =
  let !convolvedArr =
        convolveSparse
          field
          coefficients
          harmonicsArray
          thetaFreqs
          rFreqs
          cutoff
          xs
          input
      !maxMag =
        sqrt .
        L.sum .
        parMap rdeepseq (VU.sum . VU.map (\x -> (magnitude x) ^ 2) . toUnboxed) $
        convolvedArr
      !normalizedConvolvedArr =
        parMap rseq (computeS . R.map (/ (maxMag :+ 0))) convolvedArr
  in powerMethodSparse
       field
       coefficients
       harmonicsArray
       thetaFreqs
       rFreqs
       cutoff
       xs
       (numStep - 1)
       normalizedConvolvedArr

{-# INLINE computeContourSparse #-}
computeContourSparse ::
     DFTPlan
  -> FilePath
  -> R.Array U DIM4 (Complex Double)
  -> IA.Array (Int, Int) (R.Array U DIM2 (Complex Double))
  -> R.Array U DIM1 Double
  -> R.Array U DIM1 Double
  -> Double
  -> [(Double, Double)]
  -> Int
  -> String
  -> [R.Array U DIM2 (Complex Double)]
  -> IO [R.Array U DIM2 (Complex Double)]
  -- -> IO DFTArray
computeContourSparse !plan !folderPath !coefficients !harmonicsArray !thetaFreqs !rFreqs !cutoff !xs !numIteration !suffix !input = do
  printCurrentTime $ printf "Power Method %d iterations" numIteration
  let (Z :. cols :. rows) = extent (harmonicsArray IA.! (0, 0))
      ys =
        powerMethodSparse
          Source
          coefficients
          harmonicsArray
          thetaFreqs
          rFreqs
          cutoff
          xs
          numIteration
          input
      !eigenSourceSparse =
        sparseArrayToDFTArray
          rows
          cols
          (R.toList thetaFreqs)
          (R.toList rFreqs)
          xs
          ys
  harmonicsArrayDFT <-
    fmap (listArray (bounds harmonicsArray)) .
    dftExecuteBatchP plan (DFTPlanID DFT1DG [cols, rows] [0, 1]) .
    L.map (VS.convert . toUnboxed . computeUnboxedS . makeFilter2D) . IA.elems $
    harmonicsArray
  printCurrentTime "Source"
  source <-
    STC.Convolution.convolve Source plan coefficients harmonicsArrayDFT eigenSourceSparse
  plotDFTArrayPower (folderPath </> "SourcePower_Sparse.png") rows cols source
  printCurrentTime "Sink"
  let sink =
        parZipWithDFTArray
          (\vec (rFreq, thetaFreq) -> VS.map (* (cis (thetaFreq * pi))) vec)
          source
          ((,) <$> (R.toList rFreqs) <*> (R.toList thetaFreqs))
  -- sink <- convolve Sink plan coefficients harmonicsArrayDFT eigenSourceSparse
  plotDFTArrayPower (folderPath </> "SinkPower_Sparse.png") rows cols sink
  printCurrentTime "Completion"
  completion <- completionField plan source sink
  let mm =
        L.maximum . L.map (VS.maximum . VS.map magnitude) . getDFTArrayVector $
        completion
  print mm
  plotDFTArrayPower
    (folderPath </> ("CompletionPower_Sparse.png"))
    rows
    cols .
    parMapDFTArray (VS.map (/ (mm :+ 0))) $
    completion
  return ys

{-# INLINE computeContourSparseG #-}
computeContourSparseG ::
     DFTPlan
  -> FilePath
  -> R.Array U DIM4 (Complex Double)
  -> IA.Array (Int, Int) (R.Array U DIM2 (Complex Double))
  -> R.Array U DIM1 Double
  -> R.Array U DIM1 Double
  -> Double
  -> VS.Vector (Complex Double)
  -> [(Double, Double)]
  -> Int
  -> String
  -> [R.Array U DIM2 (Complex Double)]
  -> IO DFTArray
computeContourSparseG !plan !folderPath !coefficients !harmonicsArray !thetaFreqs !rFreqs !cutoff !gaussian !xs !numIteration !suffix !input = do
  printCurrentTime $ printf "Power Method %d iterations" numIteration
  let (Z :. cols :. rows) = extent (harmonicsArray IA.! (0, 0))
  blurredVecF <-
    fmap (L.map (VS.zipWith (*) gaussian)) .
    dftExecuteBatchP plan (DFTPlanID DFT1DG [cols, rows] [0, 1]) .
    L.map (VU.convert . toUnboxed) . IA.elems $
    harmonicsArray
  blurredHarmonicsArray <-
    (listArray (bounds harmonicsArray) .
     L.map
       (\vec ->
          let arr = fromUnboxed (Z :. cols :. rows) . VS.convert $ vec
          in computeS . R.traverse arr id $ \f idx@(Z :. i :. j) ->
               let r =
                     sqrt . fromIntegral $
                     (i - div cols 2) ^ 2 + (j - div rows 2) ^ 2
               in if r <= 0
                    then 0
                    else f idx)) <$>
    dftExecuteBatchP plan (DFTPlanID IDFT1DG [cols, rows] [0, 1]) blurredVecF
  let !eigenSourceSparse =
        sparseArrayToDFTArray
          rows
          cols
          (R.toList thetaFreqs)
          (R.toList rFreqs)
          xs .
        powerMethodSparse
          Source
          coefficients
          blurredHarmonicsArray
          thetaFreqs
          rFreqs
          cutoff
          xs
          numIteration $
        input
  harmonicsArrayDFT <-
    fmap (listArray (bounds blurredHarmonicsArray)) .
    dftExecuteBatchP plan (DFTPlanID DFT1DG [cols, rows] [0, 1]) .
    L.map (VS.convert . toUnboxed . computeUnboxedS . makeFilter2D) . IA.elems $
    blurredHarmonicsArray
  printCurrentTime "Source"
  source <-
    STC.Convolution.convolve Source plan coefficients harmonicsArrayDFT eigenSourceSparse
  plotDFTArrayPower (folderPath </> "SourcePower.png") rows cols source
  printCurrentTime "Sink"
  let sink =
        parZipWithDFTArray
          (\vec (rFreq, thetaFreq) -> VS.map (* (cis (thetaFreq * pi))) vec)
          source
          ((,) <$> (R.toList rFreqs) <*> (R.toList thetaFreqs))
  -- sink <- convolve Sink plan coefficients harmonicsArrayDFT eigenSourceSparse
  plotDFTArrayPower (folderPath </> "SinkPower.png") rows cols sink
  printCurrentTime "Completion"
  completion <- completionField plan source sink
  let mm =
        L.maximum . L.map (VS.maximum . VS.map magnitude) . getDFTArrayVector $
        completion
  print mm
  plotDFTArrayPower
    (folderPath </> (printf "CompletionPower_%s.png" suffix))
    rows
    cols .
    parMapDFTArray (VS.map (/ (mm :+ 0))) $
    completion
  return completion

powerMethodSparse' ::
     Field
  -> R.Array U DIM4 (Complex Double)
  -> IA.Array (Int, Int) (R.Array U DIM2 (Complex Double))
  -> R.Array U DIM1 Double
  -> R.Array U DIM1 Double
  -> Double
  -> [(Double, Double)]
  -> Int
  -> [R.Array U DIM1 (Complex Double)]
  -> IO [R.Array U DIM1 (Complex Double)]
powerMethodSparse' _ _ _ _ _ _ _ 0 !arr = return arr
powerMethodSparse' !field !coefficients !harmonicsArray !phiFreqs !rhoFreqs  !cutoff !xs !numStep !input =
  let !convolvedArr =
        convolveSparse'
          field
          coefficients
          harmonicsArray
          phiFreqs
          rhoFreqs
          cutoff
          xs
          input
      !maxMag =
        sqrt . L.sum .
        parMap rdeepseq (VU.sum . VU.map (\x -> (magnitude x) ^ 2) . toUnboxed) $
        convolvedArr
      !normalizedConvolvedArr =
        parMap rseq (computeS . R.map (\x -> x / (maxMag :+ 0))) convolvedArr
  in do
        powerMethodSparse'
          field
          coefficients
          harmonicsArray
          phiFreqs
          rhoFreqs
          cutoff
          xs
          (numStep - 1)
          normalizedConvolvedArr

{-# INLINE computeContourSparse' #-}
computeContourSparse' ::
     DFTPlan
  -> FilePath
  -> R.Array U DIM4 (Complex Double)
  -> IA.Array (Int, Int) (R.Array U DIM2 (Complex Double))
  -> [[Complex Double]]
  -> R.Array U DIM1 Double
  -> R.Array U DIM1 Double
  -> Double
  -> [(Double, Double)]
  -> Int
  -> String
  -> [R.Array U DIM1 (Complex Double)]
  -> IO [R.Array U DIM1 (Complex Double)]
  -- -> IO DFTArray
computeContourSparse' !plan !folderPath !coefficients !harmonicsArray thetaRHarmonics !phiFreqs !rhoFreqs !cutoff !xs !numIteration !suffix !input = do
  printCurrentTime $ printf "Power Method %d iterations" numIteration
  let (Z :. cols :. rows) = extent (harmonicsArray IA.! (0, 0))
  -- !eigenSourceSparse <-
  --   sparseArrayToDFTArray' rows cols (R.toList phiFreqs) (R.toList rhoFreqs) xs <$>
    -- powerMethodSparse'
    --   Source
    --   coefficients
    --   harmonicsArray
    --   phiFreqs
    --   rhoFreqs
    --   cutoff
    --   xs
    --   numIteration
    --   input
  ys <-
    powerMethodSparse'
      Source
      coefficients
      harmonicsArray
      phiFreqs
      rhoFreqs
      cutoff
      xs
      numIteration
      input
  let eigenSourceSparse =
        sparseArrayToDFTArray'
          rows
          cols
          (R.toList phiFreqs)
          (R.toList rhoFreqs)
          xs
          ys
  print (bounds harmonicsArray)
  print . L.length . IA.elems $ harmonicsArray
  harmonicsArrayDFT <-
    fmap (listArray (bounds harmonicsArray)) .
    dftExecuteBatchP plan (DFTPlanID DFT1DG [cols, rows] [0, 1]) .
    L.map (VS.convert . toUnboxed . computeUnboxedS . makeFilter2D) . IA.elems $
    harmonicsArray
  printCurrentTime "Source"
  source <-
    convolve' Source plan coefficients harmonicsArrayDFT eigenSourceSparse
  plotDFTArrayPower (folderPath </> "SourcePower_Sparse.png") rows cols source
  -- plotDFTArrayThetaR
  --   (folderPath </> "Source.png")
  --   rows
  --   cols
  --   thetaRHarmonics
  --   source
  printCurrentTime "Sink"
  -- sink <- convolve' Sink plan coefficients harmonicsArrayDFT eigenSourceSparse
  let sink =
        parZipWithDFTArray
          (\vec thetaFreq -> VS.map (* (cis (thetaFreq * pi))) vec)
          source
          (R.toList phiFreqs)
  plotDFTArrayPower (folderPath </> "SinkPower_Sparse.png") rows cols sink
  -- plotDFTArrayThetaR (folderPath </> "Sink.png") rows cols thetaRHarmonics sink
  printCurrentTime "Completion"
  completion <- completionField' plan source sink
  let mm =
        L.maximum . L.map (VS.maximum . VS.map magnitude) . getDFTArrayVector $
        completion
  print mm
  plotDFTArrayPower (folderPath </> ("CompletionPower_Sparse.png")) rows cols -- .
    -- parMapDFTArray (VS.map (/ (mm :+ 0))) $
    completion
  -- let point = (31, -4)
  --     source' =
  --       fromUnboxed (Z :. 72 :. cols :. rows) .
  --       VS.convert .
  --       VS.concat .
  --       computeFourierSeriesThetaR thetaRHarmonics .
  --       getDFTArrayVector $
  --       source
  --     sink' =
  --       fromUnboxed (Z :. 72 :. cols :. rows) .
  --       VS.convert .
  --       VS.concat .
  --       computeFourierSeriesThetaR thetaRHarmonics .
  --       getDFTArrayVector $
  --       sink
  --     completion' =
  --       fromUnboxed (Z :. 72 :. cols :. rows) .
  --       VS.convert .
  --       VS.concat .
  --       computeFourierSeriesThetaR thetaRHarmonics .
  --       getDFTArrayVector $
  --       completion
  -- plotThetaDimension
  --   folderPath
  --   ("FourierSeries_Completion_")
  --   point .
  --   R.map magnitude $
  --   completion'
  -- plotThetaDimension
  --   folderPath
  --   ("FourierSeries_Source_")
  --   point .
  --   R.map magnitude $
  --   source'
  -- plotThetaDimension
  --   folderPath
  --   ("FourierSeries_Sink_")
  --   point .
  --   R.map magnitude $
  --   sink'
  -- return completion
  return ys

{-# INLINE computeContourSparseG' #-}
computeContourSparseG' ::
     DFTPlan
  -> FilePath
  -> R.Array U DIM4 (Complex Double)
  -> IA.Array (Int, Int) (R.Array U DIM2 (Complex Double))
  -> [[Complex Double]]
  -> R.Array U DIM1 Double
  -> R.Array U DIM1 Double
  -> Double
  -> VS.Vector (Complex Double)
  -> Double
  -> [(Double, Double)]
  -> Int
  -> String
  -> [R.Array U DIM1 (Complex Double)]
  -> IO DFTArray
computeContourSparseG' !plan !folderPath !coefficients !harmonicsArray !thetaRHarmonics !phiFreqs !rhoFreqs !cutoff !gaussian !std !xs !numIteration !suffix !input = do
  printCurrentTime $ printf "Power Method %d iterations" numIteration
  let (Z :. cols :. rows) = extent (harmonicsArray IA.! (0, 0))
  blurredVecF <-
    fmap (L.map (VS.zipWith (*) gaussian)) .
    dftExecuteBatchP plan (DFTPlanID DFT1DG [cols, rows] [0, 1]) .
    L.map (VU.convert . toUnboxed) . IA.elems $
    harmonicsArray
  blurredHarmonicsArray <-
    (listArray (bounds harmonicsArray) .
     L.map
       (\vec ->
          let arr = fromUnboxed (Z :. cols :. rows) . VS.convert $ vec
          in computeS . R.traverse arr id $ \f idx@(Z :. i :. j) ->
               let r =
                     sqrt . fromIntegral $
                     (i - div cols 2) ^ 2 + (j - div rows 2) ^ 2
               in if r <= 0
                    then 0
                    else f idx)) <$>
    dftExecuteBatchP plan (DFTPlanID IDFT1DG [cols, rows] [0, 1]) blurredVecF
  !eigenSourceSparse <-
    sparseArrayToDFTArray'
      rows
      cols
      (R.toList phiFreqs)
      (R.toList rhoFreqs)
      xs <$>
    powerMethodSparse'
      Source
      coefficients
      blurredHarmonicsArray
      phiFreqs
      rhoFreqs
      cutoff
      xs
      numIteration
      input
  harmonicsArrayDFT <-
    fmap (listArray (bounds blurredHarmonicsArray)) .
    dftExecuteBatchP plan (DFTPlanID DFT1DG [cols, rows] [0, 1]) .
    L.map (VS.convert . toUnboxed . computeUnboxedS . makeFilter2D) . IA.elems $
    blurredHarmonicsArray
  printCurrentTime "Source"
  source <-
    convolve' Source plan coefficients harmonicsArrayDFT eigenSourceSparse
  plotDFTArrayPower (folderPath </> "SourcePower.png") rows cols source
  plotDFTArrayThetaR
    (folderPath </> "Source.png")
    rows
    cols
    thetaRHarmonics
    source
  printCurrentTime "Sink"
  sink <- convolve' Sink plan coefficients harmonicsArrayDFT eigenSourceSparse
  plotDFTArrayPower (folderPath </> "SinkPower.png") rows cols sink
  plotDFTArrayThetaR (folderPath </> "Sink.png") rows cols thetaRHarmonics sink
  printCurrentTime "Completion"
  completion <- completionField' plan source sink
  let mm =
        L.maximum . L.map (VS.maximum . VS.map magnitude) . getDFTArrayVector $
        completion
  print mm
  plotDFTArrayPower
    (folderPath </> (printf "CompletionPower_%s.png" suffix))
    rows
    cols .
    parMapDFTArray (VS.map (/ (mm :+ 0))) $
    completion
  return completion

-- # INLINE computeContourSparse''' #
computeContourSparse''' ::
     DFTPlan
  -> FilePath
  -> R.Array U DIM4 (Complex Double)
  -> IA.Array (Int, Int) (R.Array U DIM2 (Complex Double))
  -> [[Complex Double]]
  -> R.Array U DIM1 Double
  -> R.Array U DIM1 Double
  -> Double -> VS.Vector (Complex Double)
  -> [(Double, Double)]
  -> Int
  -> String
  -> [R.Array U DIM1 (Complex Double)]
  -> IO DFTArray
computeContourSparse''' !plan !folderPath !coefficients !harmonicsArray !thetaRHarmonics !phiFreqs !rhoFreqs !cutoff !gaussian !xs !numIteration !suffix !input = do
  printCurrentTime $ printf "Power Method %d iterations" numIteration
  let (Z :. cols :. rows) = extent (harmonicsArray IA.! (0, 0))
  blurredVecF <-
    fmap (L.map (VS.zipWith (*) gaussian)) .
    dftExecuteBatchP plan (DFTPlanID DFT1DG [cols, rows] [0, 1]) .
    L.map (VU.convert . toUnboxed) . IA.elems $
    harmonicsArray
  blurredHarmonicsArray <-
    (listArray (bounds harmonicsArray) .
     L.map
       (\vec ->
          let arr = fromUnboxed (Z :. cols :. rows) . VS.convert $ vec
          in computeS . R.traverse arr id $ \f idx@(Z :. i :. j) ->
               if (i == div cols 2) && (j == div rows 2)
                 then 0
                 else f idx)) <$>
    dftExecuteBatchP plan (DFTPlanID IDFT1DG [cols, rows] [0, 1]) blurredVecF
  !eigenSourceSparse <-
    sparseArrayToDFTArray'''
      rows
      cols
      (R.toList phiFreqs)
      (R.toList rhoFreqs)
      xs <$>
    powerMethodSparse'
      Source
      coefficients
      blurredHarmonicsArray
      phiFreqs
      rhoFreqs
      cutoff
      xs
      numIteration
      input
  harmonicsArrayDFT <-
    fmap (listArray (bounds blurredHarmonicsArray)) .
    dftExecuteBatchP plan (DFTPlanID DFT1DG [cols, rows] [0, 1]) .
    L.map (VS.convert . toUnboxed . computeUnboxedS . makeFilter2D) . IA.elems $
    blurredHarmonicsArray
  xs <-
    M.mapM
      (\(i, (p, q)) -> do
         source <- convolve' Source plan coefficients harmonicsArrayDFT q
         plotDFTArrayPower
           (folderPath </> (printf "SourcePower_%d.png" i))
           rows
           cols
           source
         sink <- convolve' Sink plan coefficients harmonicsArrayDFT p
         plotDFTArrayPower
           (folderPath </> (printf "SinkPower_%d.png" i))
           rows
           cols
           sink
         completion <- completionField' plan source sink
         plotDFTArrayPower
           (folderPath </> (printf "CompletionPower_%d.png" i))
           rows
           cols
           completion
         when
           (i == 0)
           (do let point = (31, -4)
                   source' =
                     fromUnboxed (Z :. 72 :. cols :. rows) .
                     VS.convert .
                     VS.concat .
                     computeFourierSeriesThetaR thetaRHarmonics .
                     getDFTArrayVector $
                     source
                   sink' =
                     fromUnboxed (Z :. 72 :. cols :. rows) .
                     VS.convert .
                     VS.concat .
                     computeFourierSeriesThetaR thetaRHarmonics .
                     getDFTArrayVector $
                     sink
                   completion' =
                     fromUnboxed (Z :. 72 :. cols :. rows) .
                     VS.convert .
                     VS.concat .
                     computeFourierSeriesThetaR thetaRHarmonics .
                     getDFTArrayVector $
                     completion
               plotThetaDimension
                 folderPath
                 ("FourierSeries_Completion_")
                 point .
                 R.map magnitude $
                 completion'
               plotThetaDimension
                 folderPath
                 ("FourierSeries_Source_")
                 point .
                 R.map magnitude $
                 source'
               plotThetaDimension
                 folderPath
                 ("FourierSeries_Sink_")
                 point .
                 R.map magnitude $
                 sink')
         return completion) .
    L.zip [0 :: Int ..] $
    eigenSourceSparse
  let (DFTArray rows cols thetaFreqs rFreqs _) = L.head xs
      completion =
        DFTArray rows cols thetaFreqs rFreqs .
        L.foldl1' (L.zipWith (VS.zipWith (+))) .
        L.map
          (\arr ->
             let ys = getDFTArrayVector arr
                 s = L.maximum . L.map (VS.maximum . VS.map magnitude) $ ys
             in L.map (VS.map (/ (s :+ 0))) ys) $
        xs
      mm =
        L.maximum . L.map (VS.maximum . VS.map magnitude) . getDFTArrayVector $
        completion
  print mm
  plotDFTArrayPower
    (folderPath </> (printf "CompletionPower_%s.png" suffix))
    rows
    cols .
    parMapDFTArray (VS.map (/ (mm :+ 0))) $
    completion
  return completion

powerMethodPinwheelBasis' ::
     DFTPlan
  -> FilePath
  -> [PTX]
  -> Bool
  -> R.Array U DIM4 (Complex Double)
  -> IA.Array (Int, Int) (VS.Vector (Complex Double))
  -> VS.Vector (Complex Double)
  -> Int
  -> Int
  -> Double
  -> Double
  -> Int
  -> Int
  -> DFTArray
  -> IO DFTArray
powerMethodPinwheelBasis' _ _ _ _ _ _ _ _ _ _ _ _ 0 arr = return arr
powerMethodPinwheelBasis' plan folderPath ptxs writeFlag coefficients harmonicsArray dftBias numPoints numR2Freqs delta periodR2 numBatch numStep input@(DFTArray rows cols _ _ _) = do
  printCurrentTime (show numStep)
  let convolvedArr = convoluvePinhweelBasis' coefficients harmonicsArray input
  biasedConvolvedArr <-
    multiplyPinwheelBasisBatch1 plan numBatch dftBias convolvedArr
  let s =
        L.maximum .
        parMap rdeepseq (VS.maximum . VS.map (\x -> (magnitude x) )) .
        getDFTArrayVector $
        biasedConvolvedArr
      normalizedBiasedConvolvedArr =
        parMapDFTArray (VS.map (/ (s :+ 0))) biasedConvolvedArr
      -- envelope =
      --   VU.convert .
      --   toUnboxed . computeS . fromFunction (Z :. numR2Freqs :. numR2Freqs) $ \(Z :. i :. j) ->
      --     let r =
      --           sqrt . fromIntegral $
      --           (i - div numR2Freqs 2) ^ 2 + (j - div numR2Freqs 2) ^ 2
      --     in if r == 0
      --          then 0
      --          else r ** (-0.75) :+ 0
      -- normalizedBiasedConvolvedArr =
      --   parMapDFTArray (VS.zipWith (*) envelope) normalizedBiasedConvolvedArr'
  when
    writeFlag
    (do let eigenMat =
              A.transpose . A.use . fromDFTArray . getDFTArrayVector $
              convolvedArr
        eigenR2 <-
          computeFourierSeriesR2StreamAcc
            ptxs
            numR2Freqs
            numPoints
            (L.length . getDFTArrayVector $ convolvedArr)
            periodR2
            delta
            numBatch
            eigenMat
        plotImageRepa (folderPath </> (printf "Source_%03d.png" numStep)) .
          ImageRepa 8 .
          fromUnboxed (Z :. (1 :: Int) :. numPoints :. numPoints) .
          -- VU.map sqrt .
          toUnboxed . sumS . R.map (\x -> (magnitude x) ** 2) . rotate3D $
          eigenR2
        -- let biasMat =
        --       A.transpose . A.use . fromDFTArray . getDFTArrayVector $
        --       normalizedBiasedConvolvedArr
        -- biasR2 <-
        --   computeFourierSeriesR2StreamAcc
        --     ptxs
        --     numR2Freqs
        --     numPoints
        --     (L.length . getDFTArrayVector $ convolvedArr)
        --     periodR2
        --     delta
        --     1
        --     biasMat
        -- plotImageRepa (folderPath </> (printf "Biased_%03d.png" numStep)) .
        --   ImageRepa 8 .
        --   fromUnboxed (Z :. (1 :: Int) :. numPoints :. numPoints) .
        --   VU.map sqrt .
        --   toUnboxed . sumS . R.map (\x -> (magnitude x) ** 2) . rotate3D $
        --   biasR2
    )
  powerMethodPinwheelBasis'
    plan
    folderPath
    ptxs
    writeFlag
    coefficients
    harmonicsArray
    dftBias
    numPoints
    numR2Freqs
    delta
    periodR2
    numBatch
    (numStep - 1)
    normalizedBiasedConvolvedArr

computeContourPinwheelBasis' ::
     DFTPlan
  -> FilePath
  -> [PTX]
  -> Bool
  -> R.Array U DIM4 (Complex Double)
  -> IA.Array (Int, Int) (VS.Vector (Complex Double))
  -> VS.Vector (Complex Double)
  -> Int
  -> Int
  -> Int
  -> Int
  -> Double
  -> Double
  -> Int
  -> Int
  -> DFTArray
  -> IO (R.Array U DIM4 (Complex Double))
computeContourPinwheelBasis' plan folderPath ptxs writeFlag coefficients harmonicsArray dftBias numStep numBatch numPoints numR2Freqs delta periodR2 thetaFreqs rFreqs input = do
  eigenSource' <-
    powerMethodPinwheelBasis'
      plan
      folderPath
      ptxs
      writeFlag
      coefficients
      harmonicsArray
      dftBias
      numPoints
      numR2Freqs
      delta
      periodR2
      numBatch
      numStep
      input
  let eigenSource =
        convoluvePinhweelBasis' coefficients harmonicsArray eigenSource'
      numLogPolarFreqs = thetaFreqs * rFreqs
      eigenSourceMat =
        A.transpose . A.use . fromDFTArray . getDFTArrayVector $ eigenSource
  eigenSourceR2' <-
    computeFourierSeriesR2StreamAcc
      ptxs
      numR2Freqs
      numPoints
      numLogPolarFreqs
      periodR2
      delta
      numBatch
      eigenSourceMat
  let eigenSourceR2 =
        R.reshape (Z :. rFreqs :. thetaFreqs :. numPoints :. numPoints) $
        eigenSourceR2'
      eigenSinkR2 =
        timeReversalRepa
          (L.map fromIntegral . getListFromNumber $ thetaFreqs)
          eigenSourceR2
  completionR2 <- completionFieldRepa plan eigenSourceR2 eigenSinkR2
  plotImageRepa (folderPath </> "Source.png") .
    ImageRepa 8 .
    fromUnboxed (Z :. (1 :: Int) :. numPoints :. numPoints) .
    -- VU.map sqrt .
    toUnboxed . sumS . sumS . R.map (\x -> (magnitude x) ** 2) . rotate4D2 $
    eigenSourceR2
  plotImageRepa (folderPath </> "Sink.png") .
    ImageRepa 8 .
    fromUnboxed (Z :. (1 :: Int) :. numPoints :. numPoints) .
    -- VU.map sqrt .
    toUnboxed . sumS . sumS . R.map (\x -> (magnitude x) ** 2) . rotate4D2 $
    eigenSinkR2
  plotImageRepa (folderPath </> "Completion.png") .
    ImageRepa 8 .
    fromUnboxed (Z :. (1 :: Int) :. numPoints :. numPoints) .
    -- VU.map sqrt .
    toUnboxed . sumS . sumS . R.map (\x -> (magnitude x) ** 2) . rotate4D2 $
    completionR2
  return completionR2

powerMethodPinwheelBasis ::
     DFTPlan
  -> FilePath
  -> [PTX]
  -> Bool
  -> R.Array U DIM4 (Complex Double)
  -> IA.Array (Int, Int) (VS.Vector (Complex Double))
  -> VS.Vector (Complex Double)
  -> Int
  -> Int
  -> Double
  -> Double
  -> Int
  -> Int
  -> DFTArray
  -> IO DFTArray
powerMethodPinwheelBasis _ _ _ _ _ _ _ _ _ _ _ _ 0 arr = return arr
powerMethodPinwheelBasis plan folderPath ptxs writeFlag coefficients harmonicsArray dftBias numPoints numR2Freqs delta periodR2 numBatch numStep input@(DFTArray rows cols _ _ _) = do
  printCurrentTime (show numStep)
  let convolvedArr = convoluvePinhweelBasis coefficients harmonicsArray input
  biasedConvolvedArr <-
    multiplyPinwheelBasisBatch1 plan numBatch dftBias convolvedArr
  let s =
        L.maximum .
        parMap rdeepseq (VS.maximum . VS.map (\x -> (magnitude x) )) .
        getDFTArrayVector $
        biasedConvolvedArr
      normalizedBiasedConvolvedArr =
        parMapDFTArray (VS.map (/ (s :+ 0))) biasedConvolvedArr
      -- envelope =
      --   VU.convert .
      --   toUnboxed . computeS . fromFunction (Z :. numR2Freqs :. numR2Freqs) $ \(Z :. i :. j) ->
      --     let r =
      --           sqrt . fromIntegral $
      --           (i - div numR2Freqs 2) ^ 2 + (j - div numR2Freqs 2) ^ 2
      --     in if r == 0
      --          then 0
      --          else r ** (-0.75) :+ 0
      -- normalizedBiasedConvolvedArr =
      --   parMapDFTArray (VS.zipWith (*) envelope) normalizedBiasedConvolvedArr'
  when
    writeFlag
    (do let eigenMat =
              A.transpose . A.use . fromDFTArray . getDFTArrayVector $
              convolvedArr
        eigenR2 <-
          computeFourierSeriesR2StreamAcc
            ptxs
            numR2Freqs
            numPoints
            (L.length . getDFTArrayVector $ convolvedArr)
            periodR2
            delta
            numBatch
            eigenMat
        plotImageRepa (folderPath </> (printf "Source_%03d.png" numStep)) .
          ImageRepa 8 .
          fromUnboxed (Z :. (1 :: Int) :. numPoints :. numPoints) .
          VU.map sqrt .
          toUnboxed . sumS . R.map (\x -> (magnitude x) ** 2) . rotate3D $
          eigenR2
        let biasMat =
              A.transpose . A.use . fromDFTArray . getDFTArrayVector $
              normalizedBiasedConvolvedArr
        biasR2 <-
          computeFourierSeriesR2StreamAcc
            ptxs
            numR2Freqs
            numPoints
            (L.length . getDFTArrayVector $ convolvedArr)
            periodR2
            delta
            numBatch
            biasMat
        plotImageRepa (folderPath </> (printf "Biased_%03d.png" numStep)) .
          ImageRepa 8 .
          fromUnboxed (Z :. (1 :: Int) :. numPoints :. numPoints) .
          VU.map sqrt .
          toUnboxed . sumS . R.map (\x -> (magnitude x) ** 2) . rotate3D $
          biasR2
    )
  powerMethodPinwheelBasis
    plan
    folderPath
    ptxs
    writeFlag
    coefficients
    harmonicsArray
    dftBias
    numPoints
    numR2Freqs
    delta
    periodR2
    numBatch
    (numStep - 1)
    normalizedBiasedConvolvedArr

computeContourPinwheelBasis ::
     DFTPlan
  -> FilePath
  -> [PTX]
  -> Bool
  -> R.Array U DIM4 (Complex Double)
  -> IA.Array (Int, Int) (VS.Vector (Complex Double))
  -> VS.Vector (Complex Double)
  -> Int
  -> Int
  -> Int
  -> Int
  -> Double
  -> Double
  -> Int
  -> Int
  -> DFTArray
  -> IO (R.Array U DIM4 (Complex Double))
computeContourPinwheelBasis plan folderPath ptxs writeFlag coefficients harmonicsArray dftBias numStep numBatch numPoints numR2Freqs delta periodR2 thetaFreqs rFreqs input = do
  eigenSource' <-
    powerMethodPinwheelBasis
      plan
      folderPath
      ptxs
      writeFlag
      coefficients
      harmonicsArray
      dftBias
      numPoints
      numR2Freqs
      delta
      periodR2
      numBatch
      numStep
      input
  let eigenSource =
        convoluvePinhweelBasis coefficients harmonicsArray eigenSource'
      numLogPolarFreqs = thetaFreqs * rFreqs
      eigenSourceMat =
        A.transpose . A.use . fromDFTArray . getDFTArrayVector $ eigenSource
  eigenSourceR2' <-
    computeFourierSeriesR2StreamAcc
      ptxs
      numR2Freqs
      numPoints
      numLogPolarFreqs
      periodR2
      delta
      numBatch
      eigenSourceMat
  let eigenSourceR2 =
        R.reshape (Z :. rFreqs :. thetaFreqs :. numPoints :. numPoints) $
        eigenSourceR2'
      eigenSinkR2 =
        timeReversalRepa
          (L.map fromIntegral . getListFromNumber $ thetaFreqs)
          eigenSourceR2
  completionR2 <- completionFieldRepa plan eigenSourceR2 eigenSinkR2
  plotImageRepa (folderPath </> "Source.png") .
    ImageRepa 8 .
    fromUnboxed (Z :. (1 :: Int) :. numPoints :. numPoints) .
    VU.map sqrt .
    toUnboxed . sumS . sumS . R.map (\x -> (magnitude x) ** 2) . rotate4D2 $
    eigenSourceR2
  plotImageRepa (folderPath </> "Sink.png") .
    ImageRepa 8 .
    fromUnboxed (Z :. (1 :: Int) :. numPoints :. numPoints) .
    VU.map sqrt .
    toUnboxed . sumS . sumS . R.map (\x -> (magnitude x) ** 2) . rotate4D2 $
    eigenSinkR2
  plotImageRepa (folderPath </> "Completion.png") .
    ImageRepa 8 .
    fromUnboxed (Z :. (1 :: Int) :. numPoints :. numPoints) .
    VU.map sqrt .
    toUnboxed . sumS . sumS . R.map (\x -> (magnitude x) ** 2) . rotate4D2 $
    completionR2
  return completionR2

powerMethodFourierPinwheel ::
     DFTPlan
  -> FilePath
  -> Bool
  -> FPData (VS.Vector (Complex Double))
  -> VS.Vector (Complex Double)
  -> Int
  -> Double
  -> Double
  -> Int
  -> Int
  -> FPArray (VS.Vector (Complex Double))
  -> IO (FPArray (VS.Vector (Complex Double)))
powerMethodFourierPinwheel _ _ _ _ _ _ _ _ _ 0 arr = return arr
powerMethodFourierPinwheel plan folderPath writeFlag harmonicsArray dftBias numPoints delta periodR2 numBatch numStep input = do
  printCurrentTime (show numStep)
  convolvedArr <- FP.convolve harmonicsArray input
  biasedConvolvedArr <- multiplyBias4D plan dftBias convolvedArr
  let s =
        L.sum .
        parMap rdeepseq (VS.sum . VS.map (\x -> (magnitude x) ^ 2)) . getFPArray $
        biasedConvolvedArr
      normalizedBiasedConvolvedArr =
        parMapFPArray (VS.map (/ (s :+ 0))) biasedConvolvedArr
  when
    (writeFlag && mod numStep 4 == 0)
    (do _ <-
          plotFPArray
            plan
            (folderPath </> (printf "Source_%03d.png" numStep))
            convolvedArr
        -- _ <- plotFPArray
        --        plan
        --        (folderPath </> (printf "Bias_%03d.png" numStep))
        --        normalizedBiasedConvolvedArr
        -- plotFPArrayFreqency
        --   (folderPath </> (printf "SourceFreq_%03d.png" numStep))
        --   convolvedArr
        -- plotFPArrayFreqency
        --   (folderPath </> (printf "BiasFreq1_%03d.png" numStep))
        --   normalizedBiasedConvolvedArr'
        -- plotFPArrayFreqency
        --   (folderPath </> (printf "BiasFreq2_%03d.png" numStep))
        --   normalizedBiasedConvolvedArr
        return ())
  powerMethodFourierPinwheel
    plan
    folderPath
    writeFlag
    harmonicsArray
    dftBias
    numPoints
    delta
    periodR2
    numBatch
    (numStep - 1)
    normalizedBiasedConvolvedArr

computeContourFourierPinwheel ::
     DFTPlan
  -> FilePath
  -> Bool
  -> FPData (VS.Vector (Complex Double))
  -> VS.Vector (Complex Double)
  -> Int
  -> Int
  -> Int
  -> Double
  -> Double -> Double
  -> FPArray (VS.Vector (Complex Double))
  -> String
  -> [Int]
  -> IO (R.Array U DIM4 (Complex Double))
computeContourFourierPinwheel plan folderPath writeFlag harmonicsArray dftBias numStep numBatch numPoints delta periodR2 periodEnv input suffix deviceIDs = do
  eigenSource' <-
    powerMethodFourierPinwheel
      plan
      folderPath
      writeFlag
      harmonicsArray
      dftBias
      numPoints
      delta
      periodR2
      numBatch
      numStep
      input
  eigenSource <- FP.convolve harmonicsArray eigenSource'
  initialise []
  devs <- M.mapM device deviceIDs
  ctxs <- M.mapM (\dev -> CUDA.create dev []) devs
  ptxs <- M.mapM createTargetFromContext ctxs
  -- eigenSourceR2 <-
  --   plotFPArray plan (folderPath </> printf "Source_%s.png" suffix) eigenSource
  eigenSourceR2 <-
    plotFPArrayAcc
      ptxs
      (folderPath </> printf "Source_%s.png" suffix)
      numPoints
      delta
      periodR2
      numBatch
      eigenSource
  let numRFreq = getFPArrayNumRFreq eigenSource
      numThetaFreq = getFPArrayNumThetaFreq eigenSource
  eigenSinkR2 <-
    computeUnboxedP $
    timeReversalRepa
      (L.map fromIntegral . getListFromNumber $ numThetaFreq)
      eigenSourceR2
  completionR2 <- completionFieldRepa plan eigenSourceR2 eigenSinkR2
  -- plotImageRepa (folderPath </> printf "Source_%s.png" suffix) .
  --   ImageRepa 8 .
  --   fromUnboxed (Z :. (1 :: Int) :. numPoints :. numPoints) .
  --   -- VU.map (\x -> x^2) .
  --   toUnboxed . sumS . sumS . R.map (\x -> (magnitude x) ** 2) . rotate4D2 $
  --   eigenSourceR2
  (sumP . R.map (\x -> magnitude x ** 2) . rotate4D2 $ eigenSinkR2) >>= sumP >>=
    plotImageRepa (folderPath </> printf "Sink_%s.png" suffix) .
    ImageRepa 8 . computeS . extend (Z :. (1 :: Int) :. All :. All)
  completionR2' <-
    (sumP . R.map (\x -> magnitude x ** 2) . rotate4D2 $ completionR2) >>= sumP
  plotImageRepa (folderPath </> printf "Completion_%s.png" suffix) .
    ImageRepa 8 . computeS . extend (Z :. (1 :: Int) :. All :. All) $
    completionR2'
  plotImageRepa (folderPath </> printf "Completion1_%s.png" suffix) .
    ImageRepa 8 . computeS . extend (Z :. (1 :: Int) :. All :. All) . R.map sqrt $
    completionR2'
  return completionR2

powerMethodFourierPinwheelDiscrete ::
     DFTPlan
  -> FilePath
  -> Bool
  -> FPData (VS.Vector (Complex Double))
  -> VS.Vector (Complex Double)
  -> Int
  -> Double
  -> Double
  -> Int
  -> (Int, Int)
  -> Int
  -> FPArray (VS.Vector (Complex Double))
  -> IO (FPArray (VS.Vector (Complex Double)))
powerMethodFourierPinwheelDiscrete  _ _ _ _ _ _ _ _ _ idx 0 arr = return arr
powerMethodFourierPinwheelDiscrete plan folderPath writeFlag harmonicsArray dftBias numPoints delta periodR2 numBatch idx numStep input' = do
  printCurrentTime (show numStep)
  -- input <- multiplyRFunction plan (periodR2^2 / 2) input'
  convolvedArr <- FP.convolve harmonicsArray input'
  biasedConvolvedArr <- multiplyBiasDiscrete plan dftBias convolvedArr
  let s =
        L.sum .
        parMap rdeepseq (VS.sum . VS.map (\x -> (magnitude x) ^ 2)) . getFPArray $
        biasedConvolvedArr
      normalizedBiasedConvolvedArr =
        parMapFPArray (VS.map (/ (s :+ 0))) biasedConvolvedArr
  when
    (writeFlag && mod numStep 2 == 0
    )
    (do arrR2 <-
          plotFPArray
            plan
            (folderPath </> (printf "Source_%03d.png" numStep))
            convolvedArr
        plotRThetaDist
          (folderPath </> (printf "ThetaDist_%03d.png" numStep))
          (folderPath </> (printf "RDist_%03d.png" numStep))
          (getFPArrayNumXFreq convolvedArr)
          72
          100
          periodR2
          idx
          arrR2
        _ <- plotFPArray
               plan
               (folderPath </> (printf "Bias_%03d.png" numStep))
               normalizedBiasedConvolvedArr
        return ()
        -- let biasMat = A.transpose . toMatrixAcc $ normalizedBiasedConvolvedArr
        -- biasR2 <-
        --   computeFourierSeriesR2StreamAcc
        --     ptxs
        --     (getFPArrayNumXFreq convolvedArr)
        --     numPoints
        --     (getFPArrayNumRFreq convolvedArr *
        --      getFPArrayNumThetaFreq convolvedArr)
        --     periodR2
        --     delta
        --     numBatch
        --     biasMat
        -- plotImageRepa (folderPath </> (printf "Biased_%03d.png" numStep)) .
        --   ImageRepa 8 .
        --   fromUnboxed (Z :. (1 :: Int) :. numPoints :. numPoints) .
        --   VU.map sqrt .
        --   toUnboxed . sumS . R.map (\x -> (magnitude x) ** 2) . rotate3D $
        --   biasR2
     )
  powerMethodFourierPinwheelDiscrete
    plan
    folderPath
    writeFlag
    harmonicsArray
    dftBias
    numPoints
    delta
    periodR2
    numBatch
    idx
    (numStep - 1)
    normalizedBiasedConvolvedArr

computeContourFourierPinwheelDiscrete ::
     DFTPlan
  -> FilePath
  -> Bool
  -> FPData (VS.Vector (Complex Double))
  -> VS.Vector (Complex Double)
  -> Int
  -> Int
  -> Int
  -> Double
  -> Double
  -> FPArray (VS.Vector (Complex Double))
  -> String
  -> (Int, Int)
  -> IO (R.Array U DIM4 (Complex Double))
computeContourFourierPinwheelDiscrete plan folderPath writeFlag harmonicsArray dftBias numStep numBatch numPoints delta periodR2 input suffix idx = do
  eigenSource' <-
    powerMethodFourierPinwheelDiscrete
      plan
      folderPath
      writeFlag
      harmonicsArray
      dftBias
      numPoints
      delta
      periodR2
      numBatch
      idx
      numStep
      input
  eigenSource <- FP.convolve harmonicsArray eigenSource'
  eigenSourceR2 <- plotFPArray plan (folderPath </> printf "Source_%s.png" suffix)  eigenSource
  let numThetaFreq = getFPArrayNumThetaFreq eigenSource
      eigenSinkR2 =
        timeReversalRepa
          (L.map fromIntegral . getListFromNumber $ numThetaFreq)
          eigenSourceR2
  completionR2 <- completionFieldRepa plan eigenSourceR2 eigenSinkR2
  plotImageRepa (folderPath </> printf "Sink_%s.png" suffix) .
    ImageRepa 8 .
    fromUnboxed (Z :. (1 :: Int) :. numPoints :. numPoints) .
    VU.map sqrt .
    toUnboxed . sumS . sumS . R.map (\x -> (magnitude x) ** 2) . rotate4D2 $
    eigenSinkR2
  plotImageRepa (folderPath </> printf "Completion_%s.png" suffix) .
    ImageRepa 8 .
    fromUnboxed (Z :. (1 :: Int) :. numPoints :. numPoints) .
    VU.map sqrt .
    toUnboxed . sumS . sumS . R.map (\x -> (magnitude x) ** 2) . rotate4D2 $
    completionR2
  return completionR2


module Stats (nintyFifth,  Estimate(..)) where

import qualified Data.Vector.Unboxed as U
import Statistics.Sample (mean)
import Statistics.Resampling (resample, fromResample)
import Statistics.Resampling.Bootstrap (bootstrapBCA, Estimate(..) )
import System.Random.MWC (create)

nintyFifth :: [Double] -> IO Estimate
nintyFifth sample = do
          g <- create
          resamples <- resample g [mean] 10000 sampleU -- (length sample^2) sampleU
--          print $ U.length $ fromResample $ head $ resamples
--          print resamples
          return $ head $ bootstrapBCA 0.95 sampleU [mean] resamples
  where
    sampleU = U.fromList sample


module Lib where

import qualified Data.Matrix.MatrixMarket as M

import qualified Numeric.LinearAlgebra.Sparse as LA

mmMatrixToSpMatrix (M.RMatrix (m,n) nz str ll) = LA.fromListSM (m,n) ll


module Main where

import App
import UKF

import Numeric.LinearAlgebra ((!))

app conn = do
  sensor <- recv conn
  app' conn (makeFilter sensor)

app' conn kf = do
  sensor <- recv conn

  let kf' = update sensor kf
  let px = (kf_x kf') ! 0
  let py = (kf_x kf') ! 1

  send conn Response { rX = px
                     , rY = py }
  app' conn kf'

main =
  runApp app


{-# LANGUAGE FlexibleInstances, ScopedTypeVariables,
    ViewPatterns #-}
module Tests.Distribution (tests) where

import Control.Applicative ((<$), (<$>), (<*>))
import qualified Control.Exception as E
import Data.List (find)
import Data.Typeable (Typeable)
import Numeric.MathFunctions.Constants (m_tiny,m_huge,m_epsilon)
import Numeric.MathFunctions.Comparison
import Statistics.Distribution
import Statistics.Distribution.Beta           (BetaDistribution)
import Statistics.Distribution.Binomial       (BinomialDistribution)
import Statistics.Distribution.CauchyLorentz
import Statistics.Distribution.ChiSquared     (ChiSquared)
import Statistics.Distribution.Exponential    (ExponentialDistribution)
import Statistics.Distribution.FDistribution  (FDistribution,fDistribution)
import Statistics.Distribution.Gamma          (GammaDistribution,gammaDistr)
import Statistics.Distribution.Geometric
import Statistics.Distribution.Hypergeometric
import Statistics.Distribution.Laplace        (LaplaceDistribution)
import Statistics.Distribution.Lognormal      (LognormalDistribution)
import Statistics.Distribution.Normal         (NormalDistribution)
import Statistics.Distribution.Poisson        (PoissonDistribution)
import Statistics.Distribution.StudentT
import Statistics.Distribution.Transform      (LinearTransform)
import Statistics.Distribution.Uniform        (UniformDistribution)
import Statistics.Distribution.Weibull        (WeibullDistribution)
import Statistics.Distribution.DiscreteUniform (DiscreteUniform)
import Test.Tasty                 (TestTree, testGroup)
import Test.Tasty.QuickCheck      (testProperty)
import Test.Tasty.ExpectedFailure (ignoreTest)
import Test.QuickCheck as QC
import Test.QuickCheck.Monadic as QC
import Text.Printf (printf)

import Tests.ApproxEq  (ApproxEq(..))
import Tests.Helpers   (T(..), Double01(..), testAssertion, typeName)
import Tests.Helpers   (monotonicallyIncreasesIEEE,isDenorm)
import Tests.Orphanage ()

-- | Tests for all distributions
tests :: TestTree
tests = testGroup "Tests for all distributions"
  [ contDistrTests (T :: T BetaDistribution        )
  , contDistrTests (T :: T CauchyDistribution      )
  , contDistrTests (T :: T ChiSquared              )
  , contDistrTests (T :: T ExponentialDistribution )
  , contDistrTests (T :: T GammaDistribution       )
  , contDistrTests (T :: T LaplaceDistribution     )
  , contDistrTests (T :: T LognormalDistribution   )
  , contDistrTests (T :: T NormalDistribution      )
  , contDistrTests (T :: T UniformDistribution     )
  , contDistrTests (T :: T WeibullDistribution     )
  , contDistrTests (T :: T StudentT                )
  , contDistrTests (T :: T (LinearTransform NormalDistribution))
  , contDistrTests (T :: T FDistribution           )

  , discreteDistrTests (T :: T BinomialDistribution       )
  , discreteDistrTests (T :: T GeometricDistribution      )
  , discreteDistrTests (T :: T GeometricDistribution0     )
  , discreteDistrTests (T :: T HypergeometricDistribution )
  , discreteDistrTests (T :: T PoissonDistribution        )
  , discreteDistrTests (T :: T DiscreteUniform            )

  , unitTests
  ]

----------------------------------------------------------------
-- Tests
----------------------------------------------------------------

-- Tests for continuous distribution
contDistrTests :: (Param d, ContDistr d, QC.Arbitrary d, Typeable d, Show d) => T d -> TestTree
contDistrTests t = testGroup ("Tests for: " ++ typeName t) $
  cdfTests t ++
  [ testProperty "PDF sanity"              $ pdfSanityCheck     t
  ] ++
  [ (if quantileIsInvCDF_enabled t then id else ignoreTest)
  $ testProperty "Quantile is CDF inverse" $ quantileIsInvCDF t
  , testProperty "quantile fails p<0||p>1" $ quantileShouldFail t
  , testProperty "log density check"       $ logDensityCheck    t
  , testProperty "complQuantile"           $ complQuantileCheck t
  ]

-- Tests for discrete distribution
discreteDistrTests :: (Param d, DiscreteDistr d, QC.Arbitrary d, Typeable d, Show d) => T d -> TestTree
discreteDistrTests t = testGroup ("Tests for: " ++ typeName t) $
  cdfTests t ++
  [ testProperty "Prob. sanity"         $ probSanityCheck       t
  , testProperty "CDF is sum of prob."  $ discreteCDFcorrect    t
  , testProperty "Discrete CDF is OK"   $ cdfDiscreteIsCorrect  t
  , testProperty "log probabilty check" $ logProbabilityCheck   t
  ]

-- Tests for distributions which have CDF
cdfTests :: (Param d, Distribution d, QC.Arbitrary d, Show d) => T d -> [TestTree]
cdfTests t =
  [ testProperty "C.D.F. sanity"        $ cdfSanityCheck         t
  , testProperty "CDF limit at +inf"    $ cdfLimitAtPosInfinity  t
  , testProperty "CDF limit at -inf"    $ cdfLimitAtNegInfinity  t
  , testProperty "CDF at +inf = 1"      $ cdfAtPosInfinity       t
  , testProperty "CDF at -inf = 1"      $ cdfAtNegInfinity       t
  , testProperty "CDF is nondecreasing" $ cdfIsNondecreasing     t
  , testProperty "1-CDF is correct"     $ cdfComplementIsCorrect t
  ]

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

-- CDF is in [0,1] range
cdfSanityCheck :: (Distribution d) => T d -> d -> Double -> Bool
cdfSanityCheck _ d x = c >= 0 && c <= 1
  where c = cumulative d x

-- CDF never decreases
cdfIsNondecreasing :: (Distribution d) => T d -> d -> Double -> Double -> Bool
cdfIsNondecreasing _ d = monotonicallyIncreasesIEEE $ cumulative d

-- cumulative d +∞ = 1
cdfAtPosInfinity :: (Distribution d) => T d -> d -> Bool
cdfAtPosInfinity _ d
  = cumulative d (1/0) == 1

-- cumulative d - ∞ = 0
cdfAtNegInfinity :: (Distribution d) => T d -> d -> Bool
cdfAtNegInfinity _ d
  = cumulative d (-1/0) == 0

-- CDF limit at +∞ is 1
cdfLimitAtPosInfinity :: (Param d, Distribution d) => T d -> d -> Bool
cdfLimitAtPosInfinity _ d
  = Just 1.0 == find (>=1) probs
  where
    probs = map (cumulative d)
          $ takeWhile (< (m_huge/2))
          $ iterate (*1.4) 1

-- CDF limit at -∞ is 0
cdfLimitAtNegInfinity :: (Param d, Distribution d) => T d -> d -> Bool
cdfLimitAtNegInfinity _ d
  = Just 0 == find (<=0) probs
  where
    probs = map (cumulative d)
          $ takeWhile (> (-m_huge/2))
          $ iterate (*1.4) (-1)

-- CDF's complement is implemented correctly
cdfComplementIsCorrect :: (Distribution d, Param d) => T d -> d -> Double -> Bool
cdfComplementIsCorrect _ d x
  = 1 - (cumulative d x + complCumulative d x) <= tol
  where
    tol = prec_complementCDF d

-- CDF for discrete distribution uses <= for comparison
cdfDiscreteIsCorrect :: (Param d, DiscreteDistr d) => T d -> d -> Property
cdfDiscreteIsCorrect _ d
  = counterexample (unlines badN)
  $ null badN
  where
    -- We are checking that:
    --
    -- > CDF(i) - CDF(i-e) = P(i)
    --
    -- Apporixmate equality is tricky here. Scale is set by maximum
    -- value of CDF and probability. Case when all proabilities are
    -- zero should be trated specially.
    badN = [ printf "N=%3i    p[i]=%g\tp[i+1]=%g\tdP=%g\trelerr=%g" i p p1 dp ((p1-p-dp) / max p1 dp)
           | i <- [0 .. 100]
           , let p      = cumulative d $ fromIntegral i - 1e-6
                 p1     = cumulative d $ fromIntegral i
                 dp     = probability d i
                 relerr = ((p1 - p) - dp) / max p1 dp
           , p  > m_tiny || p == 0
           , p1 > m_tiny
           , dp > m_tiny
           , relerr > tol
           ]
    tol = prec_discreteCDF d

logDensityCheck :: (ContDistr d) => T d -> d -> Double -> Property
logDensityCheck _ d x
  = not (isDenorm x)
  ==> ( counterexample (printf "density    = %g" p)
      $ counterexample (printf "logDensity = %g" logP)
      $ counterexample (printf "log p      = %g" (log p))
      $ counterexample (printf "eps        = %g" (abs (logP - log p) / max (abs (log p)) (abs logP)))
      $ or [ p == 0      && logP == (-1/0)
           , p <= m_tiny && logP < log m_tiny
             -- To avoid problems with roundtripping error in case
             -- when density is computed as exponent of logDensity we
             -- accept either inequality
           ,  (ulpDistance (log p) logP <= 32)
           || (ulpDistance p (exp logP) <= 32)
           ])
  where
    p    = density d x
    logP = logDensity d x

-- PDF is positive
pdfSanityCheck :: (ContDistr d) => T d -> d -> Double -> Bool
pdfSanityCheck _ d x = p >= 0
  where p = density d x

complQuantileCheck :: (ContDistr d) => T d -> d -> Double01 -> Property
complQuantileCheck _ d (Double01 p)
  = counterexample (printf "x0 = %g" x0)
  $ counterexample (printf "x1 = %g" x1)
  $ counterexample (printf "abs err = %g" $ abs (x1 - x0))
  $ counterexample (printf "rel err = %g" $ relativeError x1 x0)
  -- We avoid extreme tails of distributions
  --
  -- FIXME: all parameters are arbitrary at the moment
  $ and [ p > 0.01
        , p < 0.99
        , not $ isInfinite x0
        , not $ isInfinite x1
        ] ==> (if x0 < 1e6 then abs (x1 - x0) < 1e-6 else relativeError x1 x0 < 1e-12)
  where
    x0 = quantile      d (1 - p)
    x1 = complQuantile d p

-- Quantile is inverse of CDF
quantileIsInvCDF :: (Param d, ContDistr d) => T d -> d -> Double01 -> Property
quantileIsInvCDF _ d (Double01 p) =
  and [ p > m_tiny
      , p < 1
      , x > m_tiny
      , dens > 0
      ] ==>
    ( counterexample (printf "Quantile      = %g" x )
    $ counterexample (printf "Probability   = %g" p )
    $ counterexample (printf "Probability'  = %g" p')
    $ counterexample (printf "Rel. error    = %g" (relativeError p p'))
    $ counterexample (printf "Abs. error    = %e" (abs $ p - p'))
    $ counterexample (printf "Expected err. = %g" err)
    $ counterexample (printf "Distance      = %i" (ulpDistance p p'))
    $ counterexample (printf "Err/est       = %g" (fromIntegral (ulpDistance p p') / err))
    $ ulpDistance p p' <= round err
    )
  where
    -- Algorithm for error estimation is taken from here
    --
    -- http://sepulcarium.org/posts/2012-07-19-rounding_effect_on_inverse.html
    dens = density    d x
    err  = eps + eps' * abs (x / p) * dens
    --
    x    = quantile   d p
    p'   = cumulative d x
    (eps,eps') = prec_quantile_CDF d

-- Test that quantile fails if p<0 or p>1
quantileShouldFail :: (ContDistr d) => T d -> d -> Double -> Property
quantileShouldFail _ d p =
  p < 0 || p > 1 ==> QC.monadicIO $ do r <- QC.run $ E.catch
                                              (False <$ (return $! quantile d p))
                                              (\(_ :: E.SomeException) -> return True)
                                       QC.assert r

-- Probability is in [0,1] range
probSanityCheck :: (DiscreteDistr d) => T d -> d -> Int -> Bool
probSanityCheck _ d x = p >= 0 && p <= 1
  where p = probability d x

-- Check that discrete CDF is correct
discreteCDFcorrect :: (DiscreteDistr d) => T d -> d -> Int -> Int -> Property
discreteCDFcorrect _ d a b
  = counterexample (printf "CDF   = %g" p1)
  $ counterexample (printf "Sum   = %g" p2)
  $ counterexample (printf "Delta = %g" (abs (p1 - p2)))
  $ abs (p1 - p2) < 3e-10
  -- Avoid too large differeneces. Otherwise there is to much to sum
  --
  -- Absolute difference is used guard againist precision loss when
  -- close values of CDF are subtracted
  where
    n  = min a b
    m  = n + (abs (a - b) `mod` 100)
    p1 = cumulative d (fromIntegral m + 0.5) - cumulative d (fromIntegral n - 0.5)
    p2 = sum $ map (probability d) [n .. m]

logProbabilityCheck :: (DiscreteDistr d) => T d -> d -> Int -> Property
logProbabilityCheck _ d x
  = counterexample (printf "probability    = %g" p)
  $ counterexample (printf "logProbability = %g" logP)
  $ counterexample (printf "log p          = %g" (log p))
  $ counterexample (printf "eps            = %g" (abs (logP - log p) / max (abs (log p)) (abs logP)))
  $ or [ p == 0     && logP == (-1/0)
       , p < 1e-308 && logP < 609
         -- To avoid problems with roundtripping error in case
         -- when density is computed as exponent of logDensity we
         -- accept either inequality
       ,  (ulpDistance (log p) logP <= 32)
       || (ulpDistance p (exp logP) <= 32)
       ]
  where
    p    = probability d x
    logP = logProbability d x

-- | Parameters for distribution testing. Some distribution require
--   relaxing parameters a bit
class Param a where
  -- | Whether quantileIsInvCDF is enabled
  quantileIsInvCDF_enabled :: T a -> Bool
  quantileIsInvCDF_enabled _ = True
  -- | Precision for 'quantileIsInvCDF' test
  prec_quantile_CDF :: a -> (Double,Double)
  prec_quantile_CDF _ = (16,16)
  -- |
  prec_discreteCDF :: a -> Double
  prec_discreteCDF _ = 32 * m_epsilon
  -- | Precision of CDF's complement
  prec_complementCDF :: a -> Double
  prec_complementCDF _ = 1e-14

instance Param StudentT where
  -- FIXME: disabled unless incompleteBeta troubles are sorted out
  quantileIsInvCDF_enabled _ = False
instance Param BetaDistribution where
  -- FIXME: See https://github.com/bos/statistics/issues/161 for details
  quantileIsInvCDF_enabled _ = False
instance Param FDistribution where
  -- FIXME: disabled unless incompleteBeta troubles are sorted out
  quantileIsInvCDF_enabled _ = False

instance Param ChiSquared where
  prec_quantile_CDF _ = (32,32)

instance Param BinomialDistribution where
  prec_discreteCDF _ = 1e-13
instance Param CauchyDistribution
instance Param DiscreteUniform
instance Param ExponentialDistribution
instance Param GammaDistribution
instance Param GeometricDistribution
instance Param GeometricDistribution0
instance Param HypergeometricDistribution
instance Param LaplaceDistribution
instance Param LognormalDistribution where
  prec_quantile_CDF _ = (64,64)
instance Param NormalDistribution
instance Param PoissonDistribution
instance Param UniformDistribution
instance Param WeibullDistribution
instance Param a => Param (LinearTransform a)

----------------------------------------------------------------
-- Unit tests
----------------------------------------------------------------

unitTests :: TestTree
unitTests = testGroup "Unit tests"
  [ testAssertion "density (gammaDistr 150 1/150) 1 == 4.883311" $
      4.883311418525483 =~ density (gammaDistr 150 (1/150)) 1
    -- Student-T
  , testStudentPDF 0.3  1.34  0.0648215  -- PDF
  , testStudentPDF 1    0.42  0.27058
  , testStudentPDF 4.4  0.33  0.352994
  , testStudentCDF 0.3  3.34  0.757146   -- CDF
  , testStudentCDF 1    0.42  0.626569
  , testStudentCDF 4.4  0.33  0.621739
    -- Student-T General
  , testStudentUnstandardizedPDF 0.3    1.2  4      0.45 0.0533456  -- PDF
  , testStudentUnstandardizedPDF 4.3  (-2.4) 3.22 (-0.6) 0.0971141
  , testStudentUnstandardizedPDF 3.8    0.22 7.62   0.14 0.0490523
  , testStudentUnstandardizedCDF 0.3    1.2  4      0.45 0.458035   -- CDF
  , testStudentUnstandardizedCDF 4.3  (-2.4) 3.22 (-0.6) 0.698001
  , testStudentUnstandardizedCDF 3.8    0.22 7.62   0.14 0.496076
    -- F-distribution
  , testFdistrPDF  1  3   3     (1/(6 * pi)) -- PDF
  , testFdistrPDF  2  2   1.2   0.206612
  , testFdistrPDF  10 12  8     0.000385613179281892790166
  , testFdistrCDF  1  3   3     0.81830988618379067153 -- CDF
  , testFdistrCDF  2  2   1.2   0.545455
  , testFdistrCDF  10 12  8     0.99935509863451408041
  ]
  where
    -- Student-T
    testStudentPDF ndf x exact
      = testAssertion (printf "density (studentT %f) %f ~ %f" ndf x exact)
      $ eq 1e-5  exact  (density (studentT ndf) x)
    testStudentCDF ndf x exact
      = testAssertion (printf "cumulative (studentT %f) %f ~ %f" ndf x exact)
      $ eq 1e-5  exact  (cumulative (studentT ndf) x)
    -- Student-T General
    testStudentUnstandardizedPDF ndf mu sigma x exact
      = testAssertion (printf "density (studentTUnstandardized %f %f %f) %f ~ %f" ndf mu sigma x exact)
      $ eq 1e-5  exact  (density (studentTUnstandardized ndf mu sigma) x)
    testStudentUnstandardizedCDF ndf mu sigma x exact
      = testAssertion (printf "cumulative (studentTUnstandardized %f %f %f) %f ~ %f" ndf mu sigma x exact)
      $ eq 1e-5  exact  (cumulative (studentTUnstandardized ndf mu sigma) x)
    -- F-distribution
    testFdistrPDF n m x exact
      = testAssertion (printf "density (fDistribution %i %i) %f ~ %f [got %f]" n m x exact d)
      $ eq 1e-5  exact d
      where d = density (fDistribution n m) x
    testFdistrCDF n m x exact
      = testAssertion (printf "cumulative (fDistribution %i %i) %f ~ %f [got %f]" n m x exact d)
      $ eq 1e-5  exact d
      where d = cumulative (fDistribution n m) x


{-# LANGUAGE FlexibleContexts #-}

module Main where

import Wigner
import BraKet
import Operators
import QUtil
import LinAlg

import Data.Complex
import Data.List.Split

main :: IO ()
main = do
    let expd = computeExpD dim ps dZ (-bound, bound) states
    let n = round . sqrt . fromIntegral . length $ expd
    -- Wigner transform the expected D
    let cplane = getPlane dZ (-bound) bound
    writeComplex "wigner" $ wignerTransform cplane (dZ^2) cplane expd
    --writeComplex "wigner" . mconcat . wt2d . chunksOf n $ expd
    writeComplex "cplane" cplane

-- dimension of vector space to work in
dim = 40
-- number of steps in the complex plane to use
-- computation scales as O(nsteps^2)
nsteps = 50
-- bounds on the complex plane
bound = 5.0
-- step size in complex plane
dZ = 2.0 * bound / fromIntegral nsteps

-- list of states to use
-- really a list of functions, I use applicative to
-- convert to a list of sized state vectors
states :: [Int -> Ket (Complex Double)]
-- vacuum
states = [vacuum]
-- squeezed vacuum
--states = [\dim -> act (squeeze dim 0.75) $ vacuum dim]
-- one photon Fock
--states = [flip fockN 1]
-- Linear combination
--states = [\dim -> vacuum dim <+> fockN dim 1]
-- coherent state
--states = [flip coherent (2)]
-- coherent superposition (AKA schrodinger's cat)
--states = [\dim -> coherent dim 2 <+> coherent dim (-2.0)]
-- Mixed schrodinger cat
--states = [flip coherent ((-2.0) :+ 0.0), flip coherent (2.0 :+ 0.0)]
-- Absorbed
--states = [vacuum, flip fockN 1]
-- classical probabilities of mixed state
ps = [1.0]

-- Construct the complex plane from a range and step size
getPlane :: Double -> Double -> Double -> [Complex Double]
getPlane dz zMin zMax = (:+) <$> range <*> range
  where
    range = [zMin, zMin+dz..zMax]

-- Compute the expected value of the displacement operator
-- over the complex plane
computeExpD :: Int -> [Complex Double] -> Double -> (Double, Double)
                   -> [Int -> Ket (Complex Double)] -> [Complex Double]
computeExpD dim ps dz (zMin, zMax) states = (\z -> expectDisp dim z ps (states <*> [dim])) <$> plane
  where plane = getPlane dz zMin zMax

-- Write complex array to file
writeComplex :: FilePath -> [Complex Double] -> IO ()
writeComplex fpath = writeFile fpath . complexToPlane

-- Make list of complex numbers easy to parse by python
complexToPlane :: [Complex Double] -> String
complexToPlane [] = ""
complexToPlane ((x :+ y):zs) = show x ++ " " ++ show y ++ "\n" ++ complexToPlane zs