{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE EmptyCase #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE DefaultSignatures #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE DeriveGeneric #-}

-- {-# OPTIONS_GHC -O -ddump-rule-firings #-}

-- | Provides the class `ParamSet` which is used to represent the set of
-- parameters of a particular model.  The goal of SGD is then to find the
-- parameter values which minimize a given objective function.

module Numeric.SGD.ParamSet
  (
  -- * Class
    ParamSet(..)
  -- * Generics
  , GPMap
  , GAdd
  , GSub
  , GDiv
  , GMul
  , GNorm2
  ) where

import           GHC.Generics
import           GHC.TypeNats (KnownNat)

import           Prelude hiding (div)

import qualified Data.Map.Strict as M

import qualified Numeric.LinearAlgebra.Static as LA

-- | Class of types that can be treated as parameter sets.  It provides basic
-- element-wise operations (addition, multiplication, mapping) which are
-- required to perform stochastic gradient descent.  Many of the operations
-- (`add`, `mul`, `sub`, `div`, etc.) have the same interpretation and follow
-- the same laws (e.g. associativity) as the corresponding operations in `Num`
-- and `Fractional`.
--
-- `zero` takes a parameter set as argument and "zero out"'s all its elements
-- (as in the backprop library).  This allows instances for `Maybe`, `M.Map`,
-- etc., where the structure of the parameter set is dynamic.  This leads to
-- the following property:
--
--     @add (zero x) x = x@
--
-- However, `zero` does not have to obey @(add (zero x) y = y)@.
--
-- A `ParamSet` can be also seen as a (structured) vector, hence `pmap` and
-- `norm_2`.  The latter is not strictly necessary to perform SGD, but it is
-- useful to control the training process.
--
-- `pmap` should obey the following law:
--
--     @pmap id x = x@
--
-- If you leave the body of an instance declaration blank, GHC Generics will be
-- used to derive instances if the type has a single constructor and each field
-- is an instance of `ParamSet`.
class ParamSet a where
  -- | Element-wise mapping
  pmap :: (Double -> Double) -> a -> a

  -- | Zero-out all elements
  zero :: a -> a
  zero = pmap (const 0.0)

--   -- | Element-wise negation
--   neg :: a -> a
--   neg = pmap (\x -> -x)

  -- | Element-wise addition
  add :: a -> a -> a
  -- | Elementi-wise substruction
  sub :: a -> a -> a

  -- | Element-wise multiplication
  mul :: a -> a -> a
  -- | Element-wise division
  div :: a -> a -> a

  -- | L2 norm
  norm_2 :: a -> Double

--   default zero :: (Generic a, GZero (Rep a)) => a -> a
--   zero = genericZero
--   {-# INLINE zero #-}

  default pmap
    :: (Generic a, GPMap (Rep a))
    => (Double -> Double) -> a -> a
  pmap = genericPMap
  {-# INLINE pmap #-}

  default add :: (Generic a, GAdd (Rep a)) => a -> a -> a
  add = genericAdd
  {-# INLINE add #-}

  default sub :: (Generic a, GSub (Rep a)) => a -> a -> a
  sub = genericSub
  {-# INLINE sub #-}

  default mul :: (Generic a, GMul (Rep a)) => a -> a -> a
  mul = genericMul
  {-# INLINE mul #-}

  default div :: (Generic a, GDiv (Rep a)) => a -> a -> a
  div = genericDiv
  {-# INLINE div #-}

  default norm_2 :: (Generic a, GNorm2 (Rep a)) => a -> Double
  norm_2 = genericNorm2
  {-# INLINE norm_2 #-}

{-# RULES
"ParamSet pmap/pmap" forall f g p. pmap f (pmap g p) = pmap (f . g) p
  #-}

-- {-# RULES
-- "ParamSet pmap/add/pmap" forall f g p h q.
--   pmap f (add (pmap g p) (pmap h q))
--   = add (pmap (f . g) p) (pmap (f . h) q)
--   #-}

-- -- | 'add' using GHC Generics; works if all fields are instances of
-- -- 'ParamSet', but only for values with single constructors.
-- genericZero :: (Generic a, GZero (Rep a)) => a -> a
-- genericZero x = to $ gzero (from x)
-- {-# INLINE genericZero #-}

-- | 'add' using GHC Generics; works if all fields are instances of
-- 'ParamSet', but only for values with single constructors.
genericAdd :: (Generic a, GAdd (Rep a)) => a -> a -> a
genericAdd x y = to $ gadd (from x) (from y)
{-# INLINE genericAdd #-}

-- | 'sub' using GHC Generics; works if all fields are instances of
-- 'ParamSet', but only for values with single constructors.
genericSub :: (Generic a, GSub (Rep a)) => a -> a -> a
genericSub x y = to $ gsub (from x) (from y)
{-# INLINE genericSub #-}

-- | 'div' using GHC Generics; works if all fields are instances of
-- 'ParamSet', but only for values with single constructors.
genericDiv :: (Generic a, GDiv (Rep a)) => a -> a -> a
genericDiv x y = to $ gdiv (from x) (from y)
{-# INLINE genericDiv #-}

-- | 'mul' using GHC Generics; works if all fields are instances of
-- 'ParamSet', but only for values with single constructors.
genericMul :: (Generic a, GMul (Rep a)) => a -> a -> a
genericMul x y = to $ gmul (from x) (from y)
{-# INLINE genericMul #-}

-- | 'norm_2' using GHC Generics; works if all fields are instances of
-- 'ParamSet', but only for values with single constructors.
genericNorm2 :: (Generic a, GNorm2 (Rep a)) => a -> Double
genericNorm2 x = gnorm_2 (from x)
{-# INLINE genericNorm2 #-}

-- | 'pmap' using GHC Generics; works if all fields are instances of
-- 'ParamSet', but only for values with single constructors.
genericPMap :: (Generic a, GPMap (Rep a)) => (Double -> Double) -> a -> a
genericPMap f x = to $ gpmap f (from x)
{-# INLINE genericPMap #-}

--------------------------------------------------
-- Generics
--
-- Partially borrowed from the backprop library
--------------------------------------------------

-- -- | Helper class for automatically deriving 'add' using GHC Generics.
-- class GZero f where
--     gzero :: f t -> f t
--
-- instance ParamSet p => GZero (K1 i p) where
--     gzero (K1 x) = K1 (zero x)
--     {-# INLINE gzero #-}
--
-- instance (GZero f, GZero g) => GZero (f :*: g) where
--     gzero (x1 :*: y1) = x2 :*: y2
--       where
--         !x2 = gzero x1
--         !y2 = gzero y1
--     {-# INLINE gzero #-}
--
-- instance GZero V1 where
--     gzero = \case {}
--     {-# INLINE gzero #-}
--
-- instance GZero U1 where
--     gzero _ = U1
--     {-# INLINE gzero #-}
--
-- instance GZero f => GZero (M1 i c f) where
--     gzero (M1 x) = M1 (gzero x)
--     {-# INLINE gzero #-}
--
-- -- instance GZero f => GZero (f :.: g) where
-- --     gzero = Comp1 gzero
-- --     {-# INLINE gzero #-}

-- | Helper class for automatically deriving 'add' using GHC Generics.
class GAdd f where
    gadd :: f t -> f t -> f t

instance ParamSet a => GAdd (K1 i a) where
    gadd (K1 x) (K1 y) = K1 (add x y)
    {-# INLINE gadd #-}

instance (GAdd f, GAdd g) => GAdd (f :*: g) where
    gadd (x1 :*: y1) (x2 :*: y2) = x3 :*: y3
      where
        !x3 = gadd x1 x2
        !y3 = gadd y1 y2
    {-# INLINE gadd #-}

instance GAdd V1 where
    gadd = \case {}
    {-# INLINE gadd #-}

instance GAdd U1 where
    gadd _ _ = U1
    {-# INLINE gadd #-}

instance GAdd f => GAdd (M1 i c f) where
    gadd (M1 x) (M1 y) = M1 (gadd x y)
    {-# INLINE gadd #-}

-- instance GAdd f => GAdd (f :.: g) where
--     gadd (Comp1 x) (Comp1 y) = Comp1 (gadd x y)
--     {-# INLINE gadd #-}

-- | Helper class for automatically deriving 'sub' using GHC Generics.
class GSub f where
    gsub :: f t -> f t -> f t

instance ParamSet a => GSub (K1 i a) where
    gsub (K1 x) (K1 y) = K1 (sub x y)
    {-# INLINE gsub #-}

instance (GSub f, GSub g) => GSub (f :*: g) where
    gsub (x1 :*: y1) (x2 :*: y2) = x3 :*: y3
      where
        !x3 = gsub x1 x2
        !y3 = gsub y1 y2
    {-# INLINE gsub #-}

instance GSub V1 where
    gsub = \case {}
    {-# INLINE gsub #-}

instance GSub U1 where
    gsub _ _ = U1
    {-# INLINE gsub #-}

instance GSub f => GSub (M1 i c f) where
    gsub (M1 x) (M1 y) = M1 (gsub x y)
    {-# INLINE gsub #-}

-- instance GSub f => GSub (f :.: g) where
--     gsub (Comp1 x) (Comp1 y) = Comp1 (gsub x y)
--     {-# INLINE gsub #-}

-- | Helper class for automatically deriving 'mul' using GHC Generics.
class GMul f where
    gmul :: f t -> f t -> f t

instance ParamSet a => GMul (K1 i a) where
    gmul (K1 x) (K1 y) = K1 (mul x y)
    {-# INLINE gmul #-}

instance (GMul f, GMul g) => GMul (f :*: g) where
    gmul (x1 :*: y1) (x2 :*: y2) = x3 :*: y3
      where
        !x3 = gmul x1 x2
        !y3 = gmul y1 y2
    {-# INLINE gmul #-}

instance GMul V1 where
    gmul = \case {}
    {-# INLINE gmul #-}

instance GMul U1 where
    gmul _ _ = U1
    {-# INLINE gmul #-}

instance GMul f => GMul (M1 i c f) where
    gmul (M1 x) (M1 y) = M1 (gmul x y)
    {-# INLINE gmul #-}

-- instance GMul f => GMul (f :.: g) where
--     gmul (Comp1 x) (Comp1 y) = Comp1 (gmul x y)
--     {-# INLINE gmul #-}

-- | Helper class for automatically deriving 'div' using GHC Generics.
class GDiv f where
    gdiv :: f t -> f t -> f t

instance ParamSet a => GDiv (K1 i a) where
    gdiv (K1 x) (K1 y) = K1 (div x y)
    {-# INLINE gdiv #-}

instance (GDiv f, GDiv g) => GDiv (f :*: g) where
    gdiv (x1 :*: y1) (x2 :*: y2) = x3 :*: y3
      where
        !x3 = gdiv x1 x2
        !y3 = gdiv y1 y2
    {-# INLINE gdiv #-}

instance GDiv V1 where
    gdiv = \case {}
    {-# INLINE gdiv #-}

instance GDiv U1 where
    gdiv _ _ = U1
    {-# INLINE gdiv #-}

instance GDiv f => GDiv (M1 i c f) where
    gdiv (M1 x) (M1 y) = M1 (gdiv x y)
    {-# INLINE gdiv #-}

-- instance GDiv f => GDiv (f :.: g) where
--     gdiv (Comp1 x) (Comp1 y) = Comp1 (gdiv x y)
--     {-# INLINE gdiv #-}

-- | Helper class for automatically deriving 'norm_2' using GHC Generics.
class GNorm2 f where
    gnorm_2 :: f t -> Double

instance ParamSet a => GNorm2 (K1 i a) where
    gnorm_2 (K1 x) = norm_2 x
    {-# INLINE gnorm_2 #-}

instance (GNorm2 f, GNorm2 g) => GNorm2 (f :*: g) where
    gnorm_2 (x1 :*: y1) =
      sqrt ((x2 ^ (2 :: Int)) + (y2 ^ (2 :: Int)))
      where
        !x2 = gnorm_2 x1
        !y2 = gnorm_2 y1
    {-# INLINE gnorm_2 #-}

instance GNorm2 V1 where
    gnorm_2 = \case {}
    {-# INLINE gnorm_2 #-}

instance GNorm2 U1 where
    gnorm_2 _ = 0
    {-# INLINE gnorm_2 #-}

instance GNorm2 f => GNorm2 (M1 i c f) where
    gnorm_2 (M1 x) = gnorm_2 x
    {-# INLINE gnorm_2 #-}

-- -- TODO: Make sure this makes sense
-- instance GNorm2 f => GNorm2 (f :.: g) where
--     gnorm_2 (Comp1 x) = gnorm_2 x
--     {-# INLINE gnorm_2 #-}

-- | Helper class for automatically deriving 'pmap' using GHC Generics.
class GPMap f where
    gpmap :: (Double -> Double) -> f t -> f t

instance ParamSet a => GPMap (K1 i a) where
    gpmap f (K1 x) = K1 (pmap f x)
    {-# INLINE gpmap #-}

instance (GPMap f, GPMap g) => GPMap (f :*: g) where
    gpmap f (x1 :*: y1) = x2 :*: y2
      where
        !x2 = gpmap f x1
        !y2 = gpmap f y1
    {-# INLINE gpmap #-}

instance GPMap V1 where
    gpmap _ = \case {}
    {-# INLINE gpmap #-}

instance GPMap U1 where
    gpmap _ _ = U1
    {-# INLINE gpmap #-}

instance GPMap f => GPMap (M1 i c f) where
    gpmap f (M1 x) = M1 (gpmap f x)
    {-# INLINE gpmap #-}

-- instance GPMap f => GPMap (f :.: g) where
--     gpmap f (Comp1 x) = Comp1 (gpmap f x)
--     {-# INLINE gpmap #-}

--------------------------------------------------
-- Basic instances
--------------------------------------------------

instance ParamSet Double where
  zero = const 0
  pmap = id
  add = (+)
  sub = (-)
  mul = (*)
  div = (/)
  norm_2 = abs

instance (ParamSet a, ParamSet b) => ParamSet (a, b) where
  pmap f (x, y) = (pmap f x, pmap f y)
  add (x1, y1) (x2, y2) = (x1 `add` x2, y1 `add` y2)
  sub (x1, y1) (x2, y2) = (x1 `sub` x2, y1 `sub` y2)
  mul (x1, y1) (x2, y2) = (x1 `mul` x2, y1 `mul` y2)
  div (x1, y1) (x2, y2) = (x1 `div` x2, y1 `div` y2)
  norm_2 (x, y)
    = sqrt . sum . map ((^(2::Int)))
    $ [norm_2 x, norm_2 y]

instance (KnownNat n) => ParamSet (LA.R n) where
  zero = const 0
  pmap = LA.dvmap
  add = (+)
  sub = (-)
  mul = (*)
  div = (/)
  norm_2 = LA.norm_2

instance (KnownNat n, KnownNat m) => ParamSet (LA.L n m) where
  zero = const 0
  pmap = LA.dmmap
  add = (+)
  sub = (-)
  mul = (*)
  div = (/)
  norm_2 = LA.norm_2

-- | `Nothing` represents a deactivated parameter set component. If `Nothing`
-- is given as an argument to one of the `ParamSet` operations, the result is
-- `Nothing` as well.
--
-- This differs from the corresponding instance in the backprop library, where
-- `Nothing` is equivalent to `Just 0`.  However, the implementation below
-- seems to correspond adequately enough to the notion that a particular
-- component is either active or not in both the parameter set and the
-- gradient, hence it doesn't make sense to combine `Just` with `Nothing`.
instance (ParamSet a) => ParamSet (Maybe a) where
  zero = fmap zero
  pmap = fmap . pmap

  add (Just x) (Just y) = Just (add x y)
  add _ _ = Nothing

  sub (Just x) (Just y) = Just (sub x y)
  sub _ _ = Nothing

  mul (Just x) (Just y) = Just (mul x y)
  mul _ _ = Nothing

  div (Just x) (Just y) = Just (div x y)
  div _ _ = Nothing

  norm_2 = maybe 0 norm_2

-- | A map with different parameter sets (of the same type) assigned to the
-- individual keys.
--
-- When combining two maps with different sets of keys, only their intersection
-- is preserved.
instance (Ord k, ParamSet a) => ParamSet (M.Map k a) where
  zero = fmap zero
  pmap f = fmap (pmap f)
  add = M.intersectionWith add
  sub = M.intersectionWith sub
  mul= M.intersectionWith mul
  div= M.intersectionWith div
  norm_2 = sqrt . sum . map ((^(2::Int)) . norm_2)  . M.elems


module STCLocalEigenVectorEndPoint where

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

main = do
  args@(numPointStr:numOrientationStr:numScaleStr:thetaSigmaStr:scaleSigmaStr:maxScaleStr:taoStr:numTrailStr:maxTrailStr:thetaFreqsStr:scaleFreqsStr:histFileName:numIterationStr:writeSourceFlagStr:cutoffRadiusStr:patchNormFlagStr:patchNormSizeStr:shape2DStr:useFFTWWisdomFlagStr:fftwWisdomFileName:numThreadStr:_) <-
    getArgs
  print args
  let numPoint = read numPointStr :: Int
      numOrientation = read numOrientationStr :: Int
      numScale = read numScaleStr :: Int
      thetaSigma = read thetaSigmaStr :: Double
      scaleSigma = read scaleSigmaStr :: Double
      maxScale = read maxScaleStr :: Double
      tao = read taoStr :: Double
      numTrail = read numTrailStr :: Int
      maxTrail = read maxTrailStr :: Int
      thetaFreq = read thetaFreqsStr :: Double
      thetaFreqs = [-thetaFreq .. thetaFreq]
      scaleFreq = read scaleFreqsStr :: Double
      scaleFreqs = [-scaleFreq .. scaleFreq]
      numIteration = read numIterationStr :: Int
      writeSourceFlag = read writeSourceFlagStr :: Bool
      cutoffRadius = read cutoffRadiusStr :: Int
      patchNormFlag = read patchNormFlagStr :: Bool
      patchNormSize = read patchNormSizeStr :: Int
      shape2D@(Points _ minDist _) = read shape2DStr :: Points Shape2D
      useFFTWWisdomFlag = read useFFTWWisdomFlagStr :: Bool
      numThread = read numThreadStr :: Int
      folderPath = "output/test/STCLocalEigenVectorEndPoint"
      histFilePath = folderPath </> histFileName
      fftwWisdomFilePath = folderPath </> fftwWisdomFileName
  createDirectoryIfMissing True folderPath
  flag <- doesFileExist histFilePath
  radialArr <-
    if flag
      then R.map magnitude . getNormalizedHistogramArr <$>
           decodeFile histFilePath
      else do
        putStrLn "Couldn't find a Green's function data. Start simulation..."
        solveMonteCarloR2Z2T0S0Radial
          numThread
          numTrail
          maxTrail
          numPoint
          numPoint
          thetaSigma
          scaleSigma
          maxScale
          tao
          thetaFreqs
          thetaFreqs
          scaleFreqs
          scaleFreqs
          histFilePath
          (emptyHistogram
             [ (round . sqrt . fromIntegral $ 2 * (div numPoint 2) ^ 2)
             , L.length scaleFreqs
             , L.length thetaFreqs
             , L.length scaleFreqs
             , L.length thetaFreqs
             ]
             0)
  let localEigenVector =
        computeLocalEigenVector
          (ParallelParams numThread 0)
          (pinwheelHollow 2)
          (cutoff cutoffRadius radialArr)
          numPoint
          numPoint
          maxScale
          thetaFreqs
          scaleFreqs
      pinwheelArr =
        -- computeS .
        R.traverse2
          (fromListUnboxed (Z :. (L.length thetaFreqs)) thetaFreqs)
          (fromListUnboxed (Z :. (L.length scaleFreqs)) scaleFreqs)
          (\(Z :. numThetaFreq) (Z :. numScaleFreq) ->
             (Z :. numThetaFreq :. numScaleFreq :. numPoint :. numPoint)) $ \ft fs (Z :. k :. l :. i :. j) ->
          let r =
                sqrt . fromIntegral $
                (i - center numPoint) ^ 2 + (j - center numPoint) ^ 2
           in if r < fromIntegral cutoffRadius
                then 0
                else pinwheelHollow
                       1
                       (ft (Z :. k))
                       (fs (Z :. l))
                       maxScale
                       0
                       (i - center numPoint)
                       (j - center numPoint)
      -- localEigenVector = (computeS $
      --                     R.traverse2
      --                       localEigenVector'
      --                       (fromListUnboxed (Z :. (L.length thetaFreqs)) thetaFreqs)
      --                       const $ \f fFreq idx@(Z :. k :. _ :. _ :. _) ->
      --                       f idx * exp (0 :+ fFreq (Z :. k) * pi))
      pointSet = makeShape2D shape2D
      shapeArr =
        getShape2DRepaArray
          numPoint
          numPoint
          (L.map
             (\(x, y) ->
                (x + fromIntegral numPoint / 2, y + fromIntegral numPoint / 2))
             pointSet)
      xs =
        L.map (\(x, y) -> R2S1RPPoint (x, y, 0, 1)) . getShape2DIndexList $
        pointSet
      bias = computeBiasR2T0S0 numPoint numPoint thetaFreqs scaleFreqs xs
  arrR2Z2T0S0 <-
    computeUnboxedP $
    computeR2Z2T0S0ArrayRadial
      (pinwheelHollow 4)
      (cutoff cutoffRadius radialArr)
      numPoint
      numPoint
      1
      maxScale
      thetaFreqs
      scaleFreqs
      thetaFreqs
      scaleFreqs
  plan' <-
    make4DPlan emptyPlan useFFTWWisdomFlag fftwWisdomFilePath localEigenVector
  plan <-
    makeR2Z2T0S0Plan plan' useFFTWWisdomFlag fftwWisdomFilePath arrR2Z2T0S0
  plotImageRepa (folderPath </> "Shape.png") . ImageRepa 8 $ shapeArr
  filterF <- dft4D plan . computeS . makeFilter2D $ localEigenVector
  initialVec <-
    R.zipWith (*) bias <$>
    convolve4D
      plan
      filterF
      (computeInitialDistributionR2T0S0'
         numPoint
         numPoint
         thetaFreqs
         scaleFreqs
         maxScale
         xs)
  -- eigenVec <-
  --   powerMethodR2Z2
  --     plan
  --     folderPath
  --     numPoint
  --     numPoint
  --     numOrientation
  --     thetaFreqs
  --     numScale
  --     scaleFreqs
  --     maxScale
  --     pinwheelArr
  --     -- localEigenVector
  --     PowerMethodGlobal
  --     numIteration
  --     writeSourceFlag
  --     (printf
  --        "_%d_%d_%d_%d_%d_%d_%.2f_%.2f"
  --        numPoint
  --        (round thetaFreq :: Int)
  --        (round scaleFreq :: Int)
  --        (round maxScale :: Int)
  --        (round tao :: Int)
  --        cutoffRadius
  --        thetaSigma
  --        scaleSigma)
  --     bias
  --     initialVec
  --     -- bias
  --     -- (R.traverse bias id $ \f idx@(Z :. k :. l :. _ :. _) ->
  --     --    if k == div (L.length thetaFreqs) 2 && l == div (L.length scaleFreqs) 2
  --     --      then f idx
  --     --      else 0)
  pathNormMethod <-
    if patchNormFlag
      then do
        let points =
              createIndex2D .
              L.map (\(i, j) -> (i + div numPoint 2, j + div numPoint 2)) .
              getShape2DIndexList $
              pointSet
            ys =
              pointCluster
                (connectionMatrixP
                   (ParallelParams numThread 1)
                   (minDist + 1)
                   points) $
              points
        M.zipWithM_
          (\i ->
             plotImageRepa (folderPath </> (printf "Cluster%03d.png" i)) .
             ImageRepa 8)
          [1 :: Int ..] .
          cluster2Array numPoint numPoint $
          ys
        return . PowerMethodConnection $ ys
      else return PowerMethodGlobal
  sourceR2Z2 <-
    powerMethodR2Z2T0S0Reversal
      plan
      folderPath
      numPoint
      numPoint
      numOrientation
      thetaFreqs
      thetaFreqs
      numScale
      scaleFreqs
      scaleFreqs
      maxScale
      arrR2Z2T0S0
      pathNormMethod
      0
      writeSourceFlag
      (printf
         "_%d_%d_%d_%d_%d_%d_%.2f_%.2f_Final"
         numPoint
         (round thetaFreq :: Int)
         (round scaleFreq :: Int)
         (round maxScale :: Int)
         (round tao :: Int)
         cutoffRadius
         thetaSigma
         scaleSigma)
      0.0
      0.0
      bias
      -- (computeInitialEigenVectorR2T0S0
      --    numPoint
      --    numPoint
      --    thetaFreqs
      --    scaleFreqs
      --    thetaFreqs
      --    scaleFreqs
      --    xs)
      (R.traverse
         -- eigenVec
         initialVec
         (\(Z :. numThetaFreq :. numScaleFreq :. x :. y) ->
            (Z :. numThetaFreq :. numScaleFreq :. numThetaFreq :. numScaleFreq :.
             x :.
             y)) $ \f idx@(Z :. tf :. sf :. t0f :. s0f :. i :. j) ->
         f (Z :. tf :. sf :. i :. j))
  powerMethodR2Z2
    plan
    folderPath
    numPoint
    numPoint
    numOrientation
    thetaFreqs
    numScale
    scaleFreqs
    maxScale
    localEigenVector
    PowerMethodGlobal
    numIteration
    writeSourceFlag
    (printf
       "_%d_%d_%d_%d_%d_%d_%.2f_%.2f_TEST"
       numPoint
       (round thetaFreq :: Int)
       (round scaleFreq :: Int)
       (round maxScale :: Int)
       (round tao :: Int)
       cutoffRadius
       thetaSigma
       scaleSigma)
    bias
    (R.zipWith (*) bias sourceR2Z2)


{-# LANGUAGE DataKinds #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE TypeOperators #-}
{-# OPTIONS_GHC -fno-warn-orphans #-}

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

module Test.Codegen (
    codegenTests,
    ditCodegenTests,
    difCodegenTests,
    splitRadixCodegenTests,
    improvedSplitRadixCodegenTests,
    searchCodegenTests,
    modCodegenTests
  ) where

import Control.Monad (mzero,
                      replicateM)
import Data.Complex
import Data.Modular
import Data.Proxy (Proxy(..))
import Data.Typeable (Typeable)
import qualified Data.Vector.Storable as VS
import Data.Word (Word32)
import GHC.TypeLits (KnownNat,
                     natVal)
import Test.Hspec
import Test.Hspec.Core.Spec (Result)
import Test.QuickCheck ((===),
                        Arbitrary(..),
                        Gen,
                        Property,
                        counterexample,
                        forAll)
import Test.QuickCheck.Modifiers (Positive(..))

import Spiral (Config(..))
import Spiral.Driver
import Spiral.Exp
import Spiral.FFT.CooleyTukey
import Spiral.Monad
import Spiral.RootOfUnity
import Spiral.SPL
import Spiral.Search
import Spiral.Search.FFTBreakdowns
import Spiral.Search.OpCount

import qualified Test.FFTW as FFTW
import Test.Gen
import Test.Instances ()

codegenTests :: Config
             -> [Int]
             -> Spec
codegenTests conf sizes = do
    ditCodegenTests conf sizes
    difCodegenTests conf sizes
    splitRadixCodegenTests conf sizes
    improvedSplitRadixCodegenTests conf sizes
    modCodegenTests conf

ditCodegenTests :: Config -> [Int] -> Spec
ditCodegenTests conf sizes =
    describe "Generated DIT" $
    mkCodegenTests conf "DIT" (return . Re . dit) sizes

difCodegenTests :: Config -> [Int] -> Spec
difCodegenTests conf sizes =
    describe "Generated DIF" $
    mkCodegenTests conf "DIF" (return . Re . dif) sizes

splitRadixCodegenTests :: Config -> [Int] -> Spec
splitRadixCodegenTests conf sizes =
    describe "Generated splitRadix DFT" $
    mkCodegenTests conf "SplitRadix" (runSearch () f . Re . DFT) sizes
  where
    f :: (Typeable a, Typed a, MonadSpiral m)
      => SPL (Exp a)
      -> S s m (SPL (Exp a))
    f (F n w) = splitRadixBreakdown n w
    f _       = mzero

improvedSplitRadixCodegenTests :: Config -> [Int] -> Spec
improvedSplitRadixCodegenTests conf sizes =
    describe "Generated improved split radix DFT" $
    mkCodegenTests conf "ImprovedSplitRadix" (runSearch () f . Re . DFT) sizes
  where
    f :: (Typeable a, Typed a, MonadSpiral m)
      => SPL (Exp a)
      -> S s m (SPL (Exp a))
    f (F n w) = improvedSplitRadixBreakdown n w
    f _       = mzero

searchCodegenTests :: Config -> [Int] -> Spec
searchCodegenTests conf sizes =
    describe "Generated opcount-optimized DFT" $
    mkCodegenTests conf "OpcountSearch" (searchOpCount . Re . DFT) sizes

mkCodegenTests :: Config
               -> String
               -> (Int -> Spiral (SPL (Exp Double)))
               -> [Int]
               -> Spec
mkCodegenTests conf desc f = mapM_ dftTest
  where
    dftTest :: Int -> Spec
    dftTest n = it (desc ++ "(" ++ show n ++ ")") $ \(k :: Property -> IO Result) -> do
        Re e <- runSpiralWith mempty $ f n
        withComplexTransform conf ("dft" ++ show n) e $ \dft ->
          k $ forAll (uniformVectorsOfSize n) $ \v -> epsDiff (dft v) (FFTW.fft n v)

-- | Generate vectors of a given size with uniformly distributed elements.
uniformVectorsOfSize :: (Arbitrary (Uniform01 a), VS.Storable a)
                     => Int
                     -> Gen (VS.Vector a)
uniformVectorsOfSize n = do
    xs <- replicateM n arbitrary
    return $ VS.fromList (map unUniform01 xs)

-- | A uniform value i the range [0,1]
newtype Uniform01 a = Uniform01 { unUniform01 :: a }
  deriving (Eq, Ord, Show)

instance Arbitrary (Uniform01 a) => Arbitrary (Uniform01 (Complex a)) where
    arbitrary = do
        Uniform01 r <- arbitrary
        Uniform01 i <- arbitrary
        return $ Uniform01 $ r :+ i

instance Arbitrary (Uniform01 Double) where
    -- This code take from the dsp Haskell library, authored by Matt Donadio.
    -- 53 bits in [0,1], i.e., 64-bit IEEE 754 in [0,1]
    -- 67108864 = 2^26
    -- 9007199254740991 = 2^53 - 1
    arbitrary = do
        u1 :: Word32 <- arbitrary
        u2 :: Word32 <- arbitrary
        -- 27 bits
        let a = fromIntegral u1 / 32.0
        -- 26 bits
        let b = fromIntegral u2 / 64.0
        return $ Uniform01 $ (a * 67108864.0 + b) / 9007199254740991.0

-- | Return 'True' if the maximum pairwise difference between two vectors is
-- less than epsilon.
epsDiff :: forall a . (Show a, RealFloat a, VS.Storable a)
        => VS.Vector (Complex a)
        -> VS.Vector (Complex a)
        -> Property
epsDiff v1 v2 =
    counterexample ("Max delta: " ++ show maxDelta) $
    maxDelta < eps
  where
    maxDelta :: a
    maxDelta = VS.maximum $ VS.zipWith (\x y -> magnitude (x - y)) v1 v2

    eps :: a
    eps = 1e-12

modCodegenTests :: Config
                -> Spec
modCodegenTests conf =
    describe "Generated modular DFT" $ do
    describe "ℤ/17" $
        mkModCodegenTests (Proxy :: Proxy 17) conf "DIT" (return . dit) [2^n | n <- [1..4::Int]]
    describe "ℤ/2013265921" $
        mkModCodegenTests (Proxy :: Proxy 2013265921) conf "DIT" (return . dit) [2^n | n <- [1..12::Int]]

mkModCodegenTests :: forall p . KnownNat p
                  => Proxy p
                  -> Config
                  -> String
                  -> (Int -> Spiral (SPL (Exp (ℤ/p))))
                  -> [Int]
                  -> Spec
mkModCodegenTests proxy conf desc f = mapM_ dftTest
  where
    p :: Integer
    p = natVal proxy

    dftTest :: Int -> Spec
    dftTest n = it (desc ++ "(" ++ show n ++ ")") $ \(k :: Property -> IO Result) -> do
        e <- runSpiralWith mempty $ f n
        withModularTransform conf ("moddft_" ++ show p ++ "_" ++ show n) e $ \moddft ->
          k $ forAll basis $ \i -> moddft (e_i i) === res i
      where
        basis :: Gen Int
        basis = do
            Positive x <- arbitrary
            return $ x `mod` n

        e_i :: Int -> VS.Vector (ℤ/p)
        e_i i = VS.generate n $ \j -> if i == j then 1 else 0

        res :: Int -> VS.Vector (ℤ/p)
        res i = VS.generate n $ \j -> a^j
          where
            a :: ℤ/p
            a = w^i

        w :: ℤ/p
        w = omega n


{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE RebindableSyntax #-}
{-# LANGUAGE TypeSynonymInstances #-}
module DSLsofMath.W06 where
import DSLsofMath.FunExp hiding (eval, derive)
import DSLsofMath.W05
import DSLsofMath.Algebra
import Prelude hiding (Num(..),(/),(^),Fractional(..),Floating(..))
import Prelude (abs)
import Data.Complex ()

evalAll :: Transcendental a => FunExp -> [a -> a]
evalAll e = (evalFunExp e) : evalAll (derive e)

help a b as bs = as * (b : bs) + (a : as) * bs

mulStream :: Ring a => Stream a -> Stream a -> Stream a
mulStream []        _         = []
mulStream _         []        = []
mulStream (a : as)  (b : bs)  = (a*b) :  (as * (b : bs) + (a : as) * bs)

type Stream a = [a]
instance Additive a => Additive (Stream a) where
  zero  = repeat zero
  (+)   = addStream
instance AddGroup a => AddGroup (Stream a) where
  negate = negStream
instance Ring a => Multiplicative (Stream a) where
  one   = one : zero
  (*)   = mulStream

addStream :: Additive a => Stream a -> Stream a -> Stream a
addStream = zipWithLonger (+)

negStream :: AddGroup a => Stream a -> Stream a
negStream = map negate

type Taylor a = Stream a

toMaclaurin :: Ring a => PowerSeries a -> Taylor a
toMaclaurin (Poly as) = zipWith (*) as factorials

fromMaclaurin :: Field a => Taylor a -> PowerSeries a
fromMaclaurin as = Poly (zipWith (/) as factorials)

factorials :: Ring a => [a]
factorials = factorialsFrom 0 1

factorialsFrom :: Ring a => a -> a -> [a]
factorialsFrom n factn = factn : factorialsFrom (n+1) (factn * (n + 1))

ex3, ex4 :: (Eq a, Field a) => Taylor a
ex3 = toMaclaurin (x^3 + two * x)
ex4 = toMaclaurin sinx

ida a = toMaclaurin (evalP (X :+: Const a))

d f a = take 10 (toMaclaurin (evalP (f (X :+: Const a))))

dP f a = toMaclaurin (f (idx + Poly [a]))

integT :: a -> Taylor a -> Taylor a
integT = (:)

integ  ::  Field a => a -> PowerSeries a -> PowerSeries a
integ  a0 (Poly as)  =  Poly (integL a0 as)

integL :: Field a => a -> [a] -> [a]
integL c cs = c : zipWith (/) cs oneUp

type PS a = PowerSeries a

solve :: Field a => a -> (PS a -> PS a) -> PS a
solve f0 g = f              -- solves |f' = g f|, |f 0 = f0|
  where f = integ f0 (g f)

idx  ::  Field a => PS a
idx  =   solve 0 (\_f -> 1)         -- \(f'(x) = 1\), \(f(0) = 0\)

expx  ::  Field a => PS a
expx  =   solve 1 (\f -> f)         -- \(f'(x) = f(x)\), \(f(0) = 1\)
expf  ::  Field a => a -> a
expf  =   evalPS 100 expx

testExp :: Double
testExp = maximum (map diff [0,0.001..1::Double])
  where diff = abs . (expf - exp)  -- using the function instance for |exp|

testExpUnits :: Double
testExpUnits =  testExp / epsilon

epsilon :: Double  -- one bit of |Double| precision
epsilon = last (takeWhile (\x -> 1 + x /= 1) (iterate (/2) 1))

sinx,  cosx  :: Field a =>  PS a
sinx  =  integ 0 cosx
cosx  =  integ 1 (-sinx)

sinf,  cosf  :: Field a =>  a -> a
sinf  =  evalPS 100 sinx
cosf  =  evalPS 100 cosx

sx,cx::[Double]
sx = 0  :  1      :  neg 0           :  frac (neg 1) 6  :  error "TODO"
cx = 1  :  neg 0  :  frac (neg 1) 2  :  0               :  error "TODO"

instance (Eq a, Transcendental a) => Transcendental (PowerSeries a) where
   pi   =  Poly [pi]
   exp  =  expPS
   sin  =  sinPS
   cos  =  cosPS

expPS, sinPS, cosPS :: (Eq a, Transcendental a) => PS a -> PS a
expPS  as  = integ  (exp  (val as))  (expPS as   * deriv as)
sinPS  as  = integ  (sin  (val as))  (cosPS as   * deriv as)
cosPS  as  = integ  (cos  (val as))  (-sinPS as  * deriv as)

val ::  Additive a => PS a  ->  a
val (Poly (a:_))   =   a
val _              =   zero

evalP :: (Eq r, Transcendental r) => FunExp -> PS r
evalP (Const x)    =  Poly [fromRational (toRational x)]
evalP (e1 :+: e2)  =  evalP e1 + evalP e2
evalP (e1 :*: e2)  =  evalP e1 * evalP e2
evalP X            =  idx
evalP (Negate e)   =  negate  (evalP e)
evalP (Recip e)    =  recip   (evalP e)
evalP (Exp e)      =  exp     (evalP e)
evalP (Sin e)      =  sin     (evalP e)
evalP (Cos e)      =  cos     (evalP e)

evalFunExp  ::  Transcendental a => FunExp -> a -> a
evalFunExp  (Const alpha)  =   const (fromRational (toRational alpha))
evalFunExp  X              =   id
evalFunExp  (e1 :+: e2)    =   evalFunExp e1  +  evalFunExp e2
evalFunExp  (e1 :*: e2)    =   evalFunExp e1  *  evalFunExp e2
evalFunExp  (Exp e)        =   exp     (evalFunExp e)
evalFunExp  (Sin e)        =   sin     (evalFunExp e)
evalFunExp  (Cos e)        =   cos     (evalFunExp e)
evalFunExp  (Recip e)      =   recip   (evalFunExp e)
evalFunExp  (Negate e)     =   negate  (evalFunExp e)

derive     (Const _)      =  Const 0
derive     X              =  Const 1
derive     (e1 :+: e2)    =  derive e1  :+:  derive e2
derive     (e1 :*: e2)    =  (derive e1  :*:  e2)  :+:  (e1  :*:  derive e2)
derive     (Recip e)      =  let re = Recip e in Negate (re:*:re) :*: derive e
derive     (Negate e)     =  Negate (derive e)
derive     (Exp e)        =  Exp e :*: derive e
derive     (Sin e)        =  Cos e :*: derive e
derive     (Cos e)        =  Const (-1) :*: Sin e :*: derive e

instance Additive FunExp where
  (+)   =  (:+:)
  zero  = Const 0

instance AddGroup FunExp where
  negate x  = Const (-1) * x

instance Multiplicative FunExp where
  (*)  =  (:*:)
  one  = Const 1

instance MulGroup FunExp where
  recip = Recip

instance Transcendental FunExp where
  pi   =  Const pi
  exp  =  Exp
  sin  =  Sin
  cos  =  Cos

instance Additive a => Additive (a, a) where
  (f, f') + (g, g')  = (f + g,  f' + g')
  zero               = (zero,   zero)

instance AddGroup a => AddGroup (a, a) where
  negate (f, f')     = (negate f, negate f')

instance Ring a => Multiplicative (a,a) where
  (f, f') * (g, g')  = (f * g,  f' * g + f * g')
  one                = (one,zero)

instance Field a => MulGroup (a, a) where
  (f, f') / (g, g')  = (f / g,  (f' * g - g' * f) / (g * g))

instance Transcendental a => Transcendental (a, a) where
  pi = (pi, zero)
  exp  (f, f')       = (exp f,  (exp f) * f')
  sin  (f, f')       = (sin f,  cos f * f')
  cos  (f, f')       = (cos f,  -(sin f) * f')


module Numeric.Tensile.Operations.Linear.Internal where

import Data.Vector.Sized (Vector)
import Numeric.Tensile.Dimensions
import Numeric.Tensile.Tensor.Internal
import Unsafe.Coerce (unsafeCoerce)

import qualified Data.Finite as F
import qualified Data.Vector.Sized as N
import qualified Data.Vector.Storable as S
import qualified Data.Vector.Storable.Mutable as M

-- import qualified Numeric.LinearAlgebra.HMatrix as Ma

-- f :: (Dims d' -> Perm n -> Perm n) -> Perm n -> Tensor d e -> Tensor d e'
-- f dim2Idx perm t = Tensor $ reifySomeDims (permuteDims perm (dims @d)) $ \p ->
--   modifyIdx (reflect p) (modify (permuteIdxs (dim2Idx (reflect p) perm) _)) (reflect p)) t -- basically make user derive the Idxs d' -> Idxs d'

-- dim2Idx :: Rank d ~ n => Dims d -> Perm n -> Perm n
-- takes a perm on dimensions and derives a perm in indices, eg
-- dim2Idx d p = lowerPerm' ...
-- --
-- otherwise consider using the raw index fold and lowerPerm???
-- could also create :  Perm d d'

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

transpose
  :: Elt e
  => Permutable d d'
  => Dims d -> Perm (Rank d) -> Tensor d e -> Tensor d' e
transpose d p (Tensor v) = Tensor v'
  where v' = modifyIdxs d v $ \i m ->
               remapIdxs p d i $ \d' i' ->
                 M.modify m (const $ v S.! fromIdxs d' (unsafePermute p i)) (fromIdxs d' i')

{-
gen#
  :: forall s t d. PrimBytes t
  => Int# -- ^ number of elements, not checked!
             --   Avoid using this argument if possible.
  -> (s -> (# s, t #))
  -> s -> (# s, ArrayBase t d #)
gen# n f z0 = go (byteSize @e undefined *# n)
  where
    go bsize = case runRW#
     ( \s0 -> case newByteArray# bsize s0 of
         (# s1, mba #) -> case loop0 mba 0# z0 s1 of
           (# s2, z1 #) -> case unsafeFreezeByteArray# mba s2 of
             (# s3, ba #) -> (# s3, (# z1, ba #) #)
     ) of (# _, (# z1, ba #) #) -> (# z1, ArrayBase (# | (# 0# , n , ba #) #) #)
    {-# NOINLINE go #-}
    loop0 mba i z s
      | isTrue# (i ==# n) = (# s, z #)
      | otherwise = case f z of
          (# z', x #) -> loop0 mba (i +# 1#) z' (writeArray mba i x s)

3-D tensor `a` w shape=[2, 2, 3]
[[[ 1,  2,  3],
  [ 4,  5,  6]],
 [[ 7,  8,  9],
  [10, 11, 12]]]

3-D tensor `b` w shape=[2, 3, 2]
[[[13, 14],
  [15, 16],
  [17, 18]],
 [[19, 20],
  [21, 22],
  [23, 24]]]
                )

mul `a` `b` has shape=[2,2,2]
[[[ 94, 100],
  [229, 244]],
 [[508, 532],
  [697, 730]]]
-}

{-

-- <#
-- same as tf.matmul
productN
  :: forall a b c x. All KnownDim '[a, b, c]
  => KnownDims x
  => T (x +: a +: b)
  -> T (x +: b +: c)
  -> T (x +: a +: c)
productN = undefined
  where
    d = dims @x
    --s = unsafeCoerce a

product
  :: forall m x y. KnownDim m
  => KnownDims x
  => KnownDims y
  => T (x +: m) -> T (m :+ y) -> T (x ++ y)
product t u
    | I# m <- fromIntegral $ dimVal @m
    , I# n <- fromIntegral $ totalDim' @x
    , I# k <- fromIntegral $ totalDim' @y
    , nk <- n *# k
    = let loop1 i j l r | isTrue# (l ==# m) = r
                        | otherwise = loop1 i j (l +# 1#)
                          (r + ix# (i +# n *# l) t * ix# (l +# m *# j) u)

          loop2 (T# i j) | isTrue# (j ==# k) = (# T# i j, 0 #)
                         | isTrue# (i ==# n) = loop2 (T# 0# (j +# 1#))
                         | otherwise = (# T# (i +# 1#) j, loop1 i j 0# 0 #)
      in case gen# nk loop2 (T# 0# 0#) of
          (# _, r #) -> r

data T# = T# Int# Int#

-- mode-i tensor-matrix product
-- see http://citeseerx.ist.psu.edu/viewdoc/download;jsessionid=4A0663C7848627DADDBA6A243BC43E78?doi=10.1.1.130.782&rep=rep1&type=pdf
product
  :: forall m x y. KnownDim m
  => KnownDims x
  => KnownDims y
  => T (x ++ [m] ++ y)
  -> T '[m, n]
  -> T (x ++ n ++ y)

-- #>
productR
  :: All KnownDim '[a, b, c]
  => KnownDims x
  => T (a :+ b :+ x)
  -> T (b :+ c :+ x)
  -> T (a :+ c :+ x)
productR = undefined

-- <#
-- same as tf.matmul
productN
  :: All KnownDim '[a, b, c]
  => KnownDims x
  => T (x +: a +: b)
  -> T (x +: b +: c)
  -> T (x +: a +: c)
productN = undefined

transpose
  :: forall m n x. All KnownDim '[m, n]
  => T '[n, m] --(n :+ m :+ x)
  -> T '[m, n] --(m :+ n :+ x)
transpose t = case elemSize0 t of
  0# -> broadcast (ix# 0# t)
  nm | I# m <- fromIntegral $ dimVal @m
     , I# n <- fromIntegral $ dimVal @n
     -> let f ( I# i,  I# j )
              | isTrue# (i ==# m) = f ( 0 , I# (j +# 1#) ) -- skip to next col
              | otherwise         = (# ( I# (i +# 1#), I# j ), ix# (i *# n +# j) t #) --col-major indexing
        in case gen# nm f (0,0) of
          (# _, r #) -> r

-}


{-# LANGUAGE FlexibleContexts, UndecidableInstances #-}

module School.Unit.UnitParams
( UnitParams(..) ) where

import Numeric.LinearAlgebra (Container, Matrix, Vector)

data UnitParams a =
   AffineParams { affineWeights :: Matrix a
                , affineBias :: Vector a
                }
 | EmptyParams deriving (Show)

instance (Container Vector a, Eq a, Num a) => Eq (UnitParams a) where
  AffineParams { affineBias = b1, affineWeights = w1 } == AffineParams { affineBias = b2, affineWeights = w2 } =
    b1 == b2 && w1 == w2
  EmptyParams == EmptyParams = True
  _ == _ = False


module Numeric.LinearAlgebra.Matrix.Mat44 where

import Numeric.LinearAlgebra.Matrix.Class
import Numeric.LinearAlgebra.Matrix.Mat33 ( det33 )

import Data.List ( foldl1' )
import Data.NumInstances() -- Add function instances to Num

data Mat44 a = Mat44
  { m00 :: !a, m01 :: !a, m02 :: !a, m03 :: !a
  , m10 :: !a, m11 :: !a, m12 :: !a, m13 :: !a
  , m20 :: !a, m21 :: !a, m22 :: !a, m23 :: !a
  , m30 :: !a, m31 :: !a, m32 :: !a, m33 :: !a
  }
  deriving (Read, Show, Eq, Ord)

instance Matrix Mat44 where
  {-# INLINE mDim #-}
  mDim _ = 4
  {-# INLINE mElement #-}
  mElement m 0 0 = m00 m
  mElement m 0 1 = m01 m
  mElement m 0 2 = m02 m
  mElement m 0 3 = m03 m
  mElement m 1 0 = m10 m
  mElement m 1 1 = m11 m
  mElement m 1 2 = m12 m
  mElement m 1 3 = m13 m
  mElement m 2 0 = m20 m
  mElement m 2 1 = m21 m
  mElement m 2 2 = m22 m
  mElement m 2 3 = m23 m
  mElement m 3 0 = m30 m
  mElement m 3 1 = m31 m
  mElement m 3 2 = m32 m
  mElement m 3 3 = m33 m
  mElement _ i j = error ("Index " ++ show i ++ ", " ++ show j ++ ": out of range, must be 0,0 to 3,3")
  {-# INLINE mZip #-}
  mZip f m n = Mat44 (f (m00 m) (m00 n)) (f (m01 m) (m01 n)) (f (m02 m) (m02 n)) (f (m03 m) (m03 n))
                     (f (m10 m) (m10 n)) (f (m11 m) (m11 n)) (f (m12 m) (m12 n)) (f (m13 m) (m13 n))
                     (f (m20 m) (m20 n)) (f (m21 m) (m21 n)) (f (m22 m) (m22 n)) (f (m23 m) (m23 n))
                     (f (m30 m) (m30 n)) (f (m31 m) (m31 n)) (f (m32 m) (m32 n)) (f (m33 m) (m33 n))
  {-# INLINE mFold #-}
  mFold f m = foldl1' f [ mElement m i j | i <- [ 0 .. 3 ], j <- [ 0 .. 3 ] ]
  {-# INLINE mIndexOf #-}
  mIndexOf p m = fst (foldl1' p' [ ((i,j), mElement m i j) | j <- [ 3, 2 .. 0 ], i <- [ 3, 2 .. 0 ] ])
    where
    p' x@(_, a) y@(_, a') | a `p` a'  = x
                          | otherwise = y
  {-# INLINE det #-}
  det = det44

instance Functor Mat44 where
  {-# INLINE fmap #-}
  fmap f m = Mat44 (f (m00 m)) (f (m01 m)) (f (m02 m)) (f (m03 m))
                   (f (m10 m)) (f (m11 m)) (f (m12 m)) (f (m13 m))
                   (f (m20 m)) (f (m21 m)) (f (m22 m)) (f (m23 m))
                   (f (m30 m)) (f (m31 m)) (f (m32 m)) (f (m33 m))

{-# SPECIALIZE det44 :: Mat44 Double -> Double #-}
{-# SPECIALIZE det44 :: Mat44 Float  -> Float #-}
det44 :: Num a => Mat44 a -> a
det44 = m00 * det33 m11 m12 m13
                    m21 m22 m23
                    m31 m32 m33
      - m01 * det33 m10 m12 m13
                    m20 m22 m23
                    m30 m32 m33
      + m02 * det33 m10 m11 m13
                    m20 m21 m23
                    m30 m31 m33
      - m03 * det33 m10 m11 m12
                    m20 m21 m22
                    m30 m31 m32

{-# SPECIALIZE inv44 :: Mat44 Double -> Mat44 Double #-}
{-# SPECIALIZE inv44 :: Mat44 Float  -> Mat44 Float  #-}
inv44 :: Fractional a => Mat44 a -> Mat44 a
inv44 m = mApply m' m
  where
  d  = const $ det44 m
  m' = Mat44 (det33 m11 m12 m13
                    m21 m22 m23
                    m31 m32 m33 / d)
           (-(det33 m01 m02 m03
                    m21 m22 m23
                    m31 m32 m33 / d))
             (det33 m01 m02 m03
                    m11 m12 m13
                    m31 m32 m33 / d)
           (-(det33 m01 m02 m03
                    m11 m12 m13
                    m21 m22 m23 / d))

           (-(det33 m10 m12 m13
                    m20 m22 m23
                    m30 m32 m33 / d))
             (det33 m00 m02 m03
                    m20 m22 m23
                    m30 m32 m33 / d)
           (-(det33 m00 m02 m03
                    m10 m12 m13
                    m30 m32 m33 / d))
             (det33 m00 m02 m03
                    m10 m12 m13
                    m20 m22 m23 / d)

             (det33 m10 m11 m13
                    m20 m21 m23
                    m30 m31 m33 / d)
           (-(det33 m00 m01 m03
                    m20 m21 m23
                    m30 m31 m33 / d))
             (det33 m00 m01 m03
                    m10 m11 m13
                    m30 m31 m33 / d)
           (-(det33 m00 m01 m03
                    m10 m11 m13
                    m20 m21 m23 / d))

           (-(det33 m10 m11 m12
                    m20 m21 m22
                    m30 m31 m32 / d))
             (det33 m00 m01 m02
                    m20 m21 m22
                    m30 m31 m32 / d)
           (-(det33 m00 m01 m02
                    m10 m11 m12
                    m30 m31 m32 / d))
             (det33 m00 m01 m02
                    m10 m11 m12
                    m20 m21 m22 / d)


module Main (main) where

import Data.Complex (Complex(..))
import Data.Functor.Classes (readPrec1, showsPrec1)
import Text.ParserCombinators.ReadPrec (readPrec_to_S)
import Text.Read (Read(..))

comp :: Complex Int
comp = 1 :+ 1

compareInstances :: Int -> IO ()
compareInstances p = do
  let precBanner = " (at precedence " ++ show p ++ ")"
  putStrLn $ "Read vs. Read1" ++ precBanner
  print (readPrec_to_S readPrec  p "1 :+ 1" :: [(Complex Int, String)])
  print (readPrec_to_S readPrec1 p "1 :+ 1" :: [(Complex Int, String)])
  putStrLn ""
  putStrLn $ "Show vs. Show1" ++ precBanner
  putStrLn $ showsPrec  p comp ""
  putStrLn $ showsPrec1 p comp ""
  putStrLn ""

main :: IO ()
main = do
  compareInstances 6
  compareInstances 7


{-# language TypeFamilies #-}
module FFT where

import Data.Complex
import Data.Vector.Generic
import Numeric.FFT.Vector.Invertible

fft :: (Vector v a, a ~ Complex Double) => v a -> v a
fft = run dft

-- vector-fftw hace la normalizacion por si mismo
ifft :: (Vector v a, a ~ Complex Double) => v a -> v a
ifft = run idft


module Numbers where

import Control.Monad.Except (throwError)
import Data.Complex
import Data.Ratio
import DataTypes
       (Arity(..), LispError(..), LispVal(..), PrimitiveFunc, ThrowsError)

numPrimitives :: [(String, PrimitiveFunc)]
numPrimitives =
  [ ("+", numAdd)
  , ("-", numSub)
  , ("*", numMul)
  , ("/", numDiv)
  , ("modulo", numMod)
  , ("Integer?", isNumber)
  , ("complex?", isComplex)
  , ("real?", isReal)
  , ("rational?", isRational)
  , ("integer?", isInteger)
  , ("=", numBoolBinopEq)
  , ("/=", numBoolBinopNeq)
  , (">", numBoolBinopGt)
  , ("<", numBoolBinopLt)
  , (">=", numBoolBinopGte)
  , ("<=", numBoolBinopLte)
  , ("quotient", numQuotient)
  , ("remainder", numRem)
  , ("sin", numSine)
  , ("cos", numCos)
  ]

foldlM :: Monad m => (a -> b -> m a) -> a -> [b] -> m a
foldlM f v (x:xs) = (f v x) >>= \a -> foldlM f a xs
foldlM _ v [] = return v

foldl1M :: Monad m => (a -> a -> m a) -> [a] -> m a
foldl1M f (x:xs) = foldlM f x xs
foldl1M _ _ = error "Unexpected error in foldl1M"

numAdd :: PrimitiveFunc
numAdd [] = return $ Integer 0
numAdd a = foldlM (\b c -> doAdd =<< (numCast [b, c])) (Integer 0) a
  where
    doAdd :: LispVal -> ThrowsError LispVal
    doAdd (List [Integer c, Integer d]) = return $ Integer (c + d)
    doAdd (List [Rational c, Rational d]) = return $ Rational (c + d)
    doAdd (List [Float c, Float d]) = return $ Float (c + d)
    doAdd (List [Complex c, Complex d]) = return $ Complex (c + d)
    doAdd _ = throwError $ Default "Unexpected error in +"

numSub :: PrimitiveFunc
numSub [] = throwError $ NumArgs (Min 1) 0 []
numSub a = foldl1M (\b c -> doSub =<< (numCast [b, c])) a -- TODO zero args check
  where
    doSub :: LispVal -> ThrowsError LispVal
    doSub (List [Integer c, Integer d]) = return $ Integer (c - d)
    doSub (List [Rational c, Rational d]) = return $ Rational (c - d)
    doSub (List [Float c, Float d]) = return $ Float (c - d)
    doSub (List [Complex c, Complex d]) = return $ Complex (c - d)
    doSub _ = throwError $ Default "Unexpected error in -"

numMul :: PrimitiveFunc
numMul [] = return $ Integer 1
numMul a = foldl1M (\b c -> doMul =<< (numCast [b, c])) a
  where
    doMul :: LispVal -> ThrowsError LispVal
    doMul (List [Integer c, Integer d]) = return $ Integer (c * d)
    doMul (List [Rational c, Rational d]) = return $ Rational (c * d)
    doMul (List [Float c, Float d]) = return $ Float (c * d)
    doMul (List [Complex c, Complex d]) = return $ Complex (c * d)
    doMul _ = throwError $ Default "Unexpected error in -"

numDiv :: PrimitiveFunc
numDiv [] = throwError $ NumArgs (Min 1) 0 []
numDiv a = foldl1M (\b c -> doDiv =<< (numCast [b, c])) a -- TODO: Zero division error
  where
    doDiv :: LispVal -> ThrowsError LispVal
    doDiv (List [Integer c, Integer d]) =
      if d == 0
        then throwError $ Default "Zero division error"
        else return $ Rational (c % d)
    doDiv (List [Rational c, Rational d]) =
      if d == 0
        then throwError $ Default "Zero division error"
        else return $ Rational (c / d)
    doDiv (List [Float c, Float d]) =
      if d == 0.0
        then throwError $ Default "Zero division error"
        else return $ Float (c / d)
    doDiv (List [Complex c, Complex d]) =
      if d == 0
        then throwError $ Default "Zero division error"
        else return $ Complex (c / d)
    doDiv _ = throwError $ Default "Unexpected error in /"

numMod :: PrimitiveFunc
numMod [] = throwError $ NumArgs (MinMax 2 2) 0 []
numMod (a:[]) = throwError $ NumArgs (MinMax 2 2) 1 [a]
numMod [a, b] = do
  c <- numCast $ [a, b]
  doMod c
  where
    doMod :: LispVal -> ThrowsError LispVal
    doMod (List [Integer c, Integer d]) = return $ Integer (c `mod` d)
    doMod (List [c@(Rational _), d@(Rational _)]) = do
      Integer c' <- numToInt c
      Integer d' <- numToInt d
      return $ Rational ((c' `mod` d') % 1)
    doMod (List [c@(Float _), d@(Float _)]) = do
      Integer c' <- numToInt c
      Integer d' <- numToInt d
      return $ Float (fromInteger (c' `mod` d'))
    doMod (List [c@(Complex _), d@(Complex _)]) = do
      Integer c' <- numToInt c
      Integer d' <- numToInt d
      return $ Complex (fromInteger (c' `mod` d') :+ 0)
    doMod _ = throwError $ Default "Unexpected error in modulo"
numMod a = throwError $ NumArgs (MinMax 2 2) (length a) a

numRem :: PrimitiveFunc
numRem [] = throwError $ NumArgs (MinMax 2 2) 0 []
numRem (a:[]) = throwError $ NumArgs (MinMax 2 2) 1 [a]
numRem [a, b] = do
  c <- numCast $ [a, b]
  doRem c
  where
    doRem :: LispVal -> ThrowsError LispVal
    doRem (List [Integer c, Integer d]) = return $ Integer (c `rem` d)
    doRem (List [c@(Rational _), d@(Rational _)]) = do
      Integer c' <- numToInt c
      Integer d' <- numToInt d
      return $ Rational ((c' `rem` d') % 1)
    doRem (List [c@(Float _), d@(Float _)]) = do
      Integer c' <- numToInt c
      Integer d' <- numToInt d
      return $ Float (fromInteger (c' `rem` d'))
    doRem (List [c@(Complex _), d@(Complex _)]) = do
      Integer c' <- numToInt c
      Integer d' <- numToInt d
      return $ Complex (fromInteger (c' `rem` d') :+ 0)
    doRem _ = throwError $ Default "Unexpected error in remainder"
numRem a = throwError $ NumArgs (MinMax 2 2) (length a) a

isInteger :: PrimitiveFunc
isInteger [Integer _] = return $ Bool True
isInteger [_] = return $ Bool False
isInteger a = throwError $ NumArgs (MinMax 1 1) (length a) a

isRational :: PrimitiveFunc
isRational [Rational _] = return $ Bool True
isRational a = isInteger a

isReal :: PrimitiveFunc
isReal [Float _] = return $ Bool True
isReal a = isRational a

isComplex :: PrimitiveFunc
isComplex [Complex _] = return $ Bool True
isComplex a = isReal a

isNumber :: PrimitiveFunc
isNumber = isComplex

numToInt :: LispVal -> ThrowsError LispVal
numToInt a@(Integer _) = return a
numToInt (Rational a) =
  if denominator a == 1
    then return $ Integer (numerator a)
    else throwError $ Default "Could not convert rational to integer"
numToInt (Float a) =
  if a == fromInteger (round a)
    then return $ Integer (round a)
    else throwError $ Default "Could not convert float to integer"
numToInt (Complex a) =
  let rp = realPart a
  in if imagPart a == 0 && rp == fromInteger (round rp)
       then return $ Integer (round rp)
       else throwError $ Default "Could not convert complex to integer"
numToInt _ = throwError $ Default "Could not convert non-number to integer"

numCast :: [LispVal] -> ThrowsError LispVal
numCast [a@(Integer _), b@(Integer _)] = return $ List [a, b]
numCast [a@(Rational _), b@(Rational _)] = return $ List [a, b]
numCast [a@(Float _), b@(Float _)] = return $ List [a, b]
numCast [a@(Complex _), b@(Complex _)] = return $ List [a, b]
-- Integer
numCast [(Integer a), b@(Rational _)] = return $ List [Rational (a % 1), b]
numCast [(Integer a), b@(Float _)] = return $ List [Float (fromInteger a), b]
numCast [(Integer a), b@(Complex _)] =
  return $ List [Complex (fromInteger a :+ 0), b]
-- Rational
numCast [a@(Rational _), (Integer b)] = return $ List [a, Rational (b % 1)]
numCast [(Rational a), b@(Float _)] = return $ List [Float (fromRational a), b]
numCast [(Rational a), b@(Complex _)] =
  return $ List [Complex (fromRational a :+ 0), b]
-- Float
numCast [a@(Float _), (Rational b)] = return $ List [a, Float (fromRational b)]
numCast [a@(Float _), (Integer b)] = return $ List [a, Float (fromInteger b)]
numCast [(Float a), b@(Complex _)] = return $ List [Complex (a :+ 0), b]
-- Complex
numCast [a@(Complex _), (Rational b)] =
  return $ List [a, Complex (fromRational b :+ 0)]
numCast [a@(Complex _), (Float b)] = return $ List [a, Complex (b :+ 0)]
numCast [a@(Complex _), (Integer b)] =
  return $ List [a, Complex (fromInteger b :+ 0)]
numCast [a, b] =
  case a of
    Integer _ -> doThrowError b
    Float _ -> doThrowError b
    Rational _ -> doThrowError b
    Complex _ -> doThrowError b
    _ -> doThrowError a
  where
    doThrowError :: LispVal -> ThrowsError LispVal
    doThrowError num = throwError $ TypeMismatch "Integer" num
numCast _ = throwError $ Default "Unexpected error in numCast"

numBoolBinop ::
     String
  -> (LispVal -> LispVal -> ThrowsError LispVal)
  -> LispVal
  -> [LispVal]
  -> ThrowsError LispVal
numBoolBinop name' op b (c:d) = do
  List [b', c'] <- numCast [b, c]
  result <- op b' c'
  case result of
    Bool True -> numBoolBinop name' op c' d
    Bool False -> return $ Bool False
    _ -> throwError $ Default $ "Unexpected error in " ++ name'
numBoolBinop _ _ _ _ = return $ Bool True

numBoolBinopEq :: [LispVal] -> ThrowsError LispVal
numBoolBinopEq [] = throwError $ NumArgs (Min 1) 0 []
numBoolBinopEq (x:xs) = numBoolBinop "=" fn x xs
  where
    fn :: LispVal -> LispVal -> ThrowsError LispVal
    fn (Integer c) (Integer d) = return $ Bool (c == d)
    fn (Rational c) (Rational d) = return $ Bool (c == d)
    fn (Float c) (Float d) = return $ Bool (c == d)
    fn (Complex c) (Complex d) = return $ Bool (c == d)
    fn _ _ = throwError $ Default "Unexpected error in ="

numBoolBinopNeq :: [LispVal] -> ThrowsError LispVal
numBoolBinopNeq [] = throwError $ NumArgs (Min 1) 0 []
numBoolBinopNeq (x:xs) = numBoolBinop "/=" fn x xs
  where
    fn :: LispVal -> LispVal -> ThrowsError LispVal
    fn (Integer c) (Integer d) = return $ Bool (c /= d)
    fn (Rational c) (Rational d) = return $ Bool (c /= d)
    fn (Float c) (Float d) = return $ Bool (c /= d)
    fn (Complex c) (Complex d) = return $ Bool (c /= d)
    fn _ _ = throwError $ Default "Unexpected error in /="

numBoolBinopLt :: [LispVal] -> ThrowsError LispVal
numBoolBinopLt [] = throwError $ NumArgs (Min 1) 0 []
numBoolBinopLt (x:xs) = numBoolBinop "<" fn x xs
  where
    fn :: LispVal -> LispVal -> ThrowsError LispVal
    fn (Integer c) (Integer d) = return $ Bool (c < d)
    fn (Rational c) (Rational d) = return $ Bool (c < d)
    fn (Float c) (Float d) = return $ Bool (c < d)
    fn (Complex _) (Complex _) =
      throwError $ Default "< not defined for complex numbers"
    fn _ _ = throwError $ Default "Unexpected error in <"

numBoolBinopLte :: [LispVal] -> ThrowsError LispVal
numBoolBinopLte [] = throwError $ NumArgs (Min 1) 0 []
numBoolBinopLte (x:xs) = numBoolBinop "<=" fn x xs
  where
    fn :: LispVal -> LispVal -> ThrowsError LispVal
    fn (Integer c) (Integer d) = return $ Bool (c <= d)
    fn (Rational c) (Rational d) = return $ Bool (c <= d)
    fn (Float c) (Float d) = return $ Bool (c <= d)
    fn (Complex _) (Complex _) =
      throwError $ Default "<= not defined for complex numbers"
    fn _ _ = throwError $ Default "Unexpected error in <="

numBoolBinopGt :: [LispVal] -> ThrowsError LispVal
numBoolBinopGt [] = throwError $ NumArgs (Min 1) 0 []
numBoolBinopGt (x:xs) = numBoolBinop ">" fn x xs
  where
    fn :: LispVal -> LispVal -> ThrowsError LispVal
    fn (Integer c) (Integer d) = return $ Bool (c > d)
    fn (Rational c) (Rational d) = return $ Bool (c > d)
    fn (Float c) (Float d) = return $ Bool (c > d)
    fn (Complex _) (Complex _) =
      throwError $ Default "> not defined for complex numbers"
    fn _ _ = throwError $ Default "Unexpected error in >"

numBoolBinopGte :: [LispVal] -> ThrowsError LispVal
numBoolBinopGte [] = throwError $ NumArgs (Min 1) 0 []
numBoolBinopGte (x:xs) = numBoolBinop ">=" fn x xs
  where
    fn :: LispVal -> LispVal -> ThrowsError LispVal
    fn (Integer c) (Integer d) = return $ Bool (c >= d)
    fn (Rational c) (Rational d) = return $ Bool (c >= d)
    fn (Float c) (Float d) = return $ Bool (c >= d)
    fn (Complex _) (Complex _) =
      throwError $ Default ">= not defined for complex numbers"
    fn _ _ = throwError $ Default "Unexpected error in >="

numQuotient :: PrimitiveFunc
numQuotient args =
  if length args /= 2
    then throwError $ NumArgs (MinMax 2 2) (length args) args
    else do
      nums <- numCast args
      case nums of
        List [Integer a, Integer b] -> return $ Integer (a `quot` b)
        _ -> throwError $ Default "Unexpected error in <=" -- TODO better errors

unaryTrig ::
     (Double -> Double) -> (Double -> Double -> Complex Double) -> PrimitiveFunc
unaryTrig op complexOp args =
  if length args /= 1
    then throwError $ NumArgs (MinMax 1 1) (length args) args
    else case args of
           [Integer a] -> return $ Float (op $ fromInteger a)
           [Rational a] -> return $ Float (op $ fromRational a)
           [Float a] -> return $ Float (op a)
           [Complex a] -> return $ Complex $ complexOp (realPart a) (imagPart a)
           _ -> throwError $ Default "Numerical input expected"

numSine :: PrimitiveFunc
numSine = unaryTrig sin (\r i -> ((sin r) * (cosh i)) :+ ((cos r) * (sinh i)))

numCos :: PrimitiveFunc
numCos =
  unaryTrig cos (\r i -> ((cos r) * (cosh i)) :+ (-1 * ((sin r) * (sinh i))))


module HollowPinwheel where

import           Control.Monad                  as M
import qualified Data.Array.Accelerate          as A
import           Data.Array.Accelerate.LLVM.PTX
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           FokkerPlanck.FourierSeries
import           FokkerPlanck.GreensFunction
import           FokkerPlanck.Histogram
import           Foreign.CUDA.Driver            as CUDA
import           FourierMethod.BlockMatrixAcc
import           FourierMethod.FourierSeries2D
import           Image.IO
import           Pinwheel.FourierSeries2D
import           STC
import           System.Directory
import           System.Environment
import           System.FilePath
import           Text.Printf
import           Utils.Array
import           Utils.List
import           Utils.Time
import Graphics.Rendering.Chart.Easy
import Utils.Parallel
import Graphics.Rendering.Chart.Backend.Cairo
import Utils.SimpsonRule

main = do
  args@(deviceIDsStr:numPointsStr:numR2FreqStr:periodR2Str:phiFreqsStr:rhoFreqsStr:thetaFreqsStr:scaleFreqsStr:stdStr:radiusStr:sStr:idxStr:numBatchStr:aStr:bStr:deltaStr:radialFreqStr:numThreadStr:_) <-
    getArgs
  let deviceIDs = read deviceIDsStr :: [Int]
      numPoints = read numPointsStr :: Int
      numR2Freq = read numR2FreqStr :: Int
      periodR2 = read periodR2Str :: Double
      phiFreq = read phiFreqsStr :: Int
      phiFreqs = L.map fromIntegral [-phiFreq .. phiFreq]
      rhoFreq = read rhoFreqsStr :: Int
      rhoFreqs = L.map fromIntegral [-rhoFreq .. rhoFreq]
      thetaFreq = read thetaFreqsStr :: Int
      thetaFreqs = L.map fromIntegral [-thetaFreq .. thetaFreq]
      scaleFreq = read scaleFreqsStr :: Int
      scaleFreqs = L.map fromIntegral [-scaleFreq .. scaleFreq]
      std = read stdStr :: Double
      radius = read radiusStr :: Double
      s = read sStr :: Double
      idx = read idxStr :: (Int, Int)
      numBatch = read numBatchStr :: Int
      numThread = read numThreadStr :: Int
      folderPath = "output/test/HollowPinwheel"
  createDirectoryIfMissing True folderPath
  initialise []
  devs <- M.mapM device deviceIDs
  ctxs <- M.mapM (\dev -> CUDA.create dev []) devs
  ptxs <- M.mapM createTargetFromContext ctxs
  let harmonicsArray' =
        pinwheelFourierCoefficientsAnatical
          numR2Freq
          phiFreq
          rhoFreq
          thetaFreq
          scaleFreq
          (-s)
          periodR2
          (periodR2 * sqrt 2)
      hpf = idealLowPassFilter radius periodR2 numR2Freq
      hpfMat =
        A.use .
        A.fromList (A.Z A.:. (numR2Freq ^ 2) A.:. (1 :: Int)) . VS.toList $
        hpf
  plotImageRepaComplex (folderPath </> "Filter_Coef.png") .
    ImageRepa 8 .
    fromUnboxed (Z :. (1 :: Int) :. numR2Freq :. numR2Freq) . VS.convert $
    hpf
  plan <-
    makePlan
      folderPath
      emptyPlan
      numPoints
      numPoints
      numR2Freq
      (2 * thetaFreq + 1)
      (2 * scaleFreq + 1)
  harmonicsArray <- convolveFrequency plan numR2Freq hpf harmonicsArray'
  -- let harmonicsArray = centerHollow numR2Freq harmonicsArray'
  let pinwheelMat =
        A.use .
        A.fromList (A.Z A.:. (numR2Freq ^ 2) A.:. (1 :: Int)) . VS.toList $
        (harmonicsArray' IA.! idx)
      pinwheelHPMat =
        A.use .
        A.fromList (A.Z A.:. (numR2Freq ^ 2) A.:. (1 :: Int)) . VS.toList $
        (harmonicsArray IA.! idx)
  plotImageRepaComplex (folderPath </> "FilteredPinhweel_Coef.png") .
    ImageRepa 8 .
    fromUnboxed (Z :. (1 :: Int) :. numR2Freq :. numR2Freq) . VS.convert $
    (harmonicsArray IA.! idx)
  pinwheelR2 <-
    computeFourierSeriesR2StreamAcc
      ptxs
      numR2Freq
      numPoints
      1
      periodR2
      1
      numBatch
      pinwheelMat
  pinwheelHPR2 <-
    computeFourierSeriesR2StreamAcc
      ptxs
      numR2Freq
      numPoints
      1
      periodR2
      1
      numBatch
      pinwheelHPMat
  hpfR2' <-
    computeFourierSeriesR2StreamAcc
      ptxs
      numR2Freq
      numPoints
      1
      periodR2
      1
      numBatch
      hpfMat
  plotImageRepaComplex (folderPath </> "Pinwheel.png") .
    ImageRepa 8 .
    computeS . extend (Z :. (1 :: Int) :. All :. All) . sumS . rotate3D $
    pinwheelR2
  printf
    "pinwheel center: %f\n"
    (magnitude $
     (pinwheelR2 R.!
      (Z :. (0 :: Int) :. (div numPoints 2) :. (div numPoints 2 - 1))))
  plotImageRepaComplex (folderPath </> "Pinwheel_HighPass.png") .
    ImageRepa 8 .
    computeS . extend (Z :. (1 :: Int) :. All :. All) . sumS . rotate3D $
    pinwheelHPR2
  printf
    "lowpassed pinwheel center: %f\n"
    (magnitude $
     (pinwheelHPR2 R.!
      (Z :. (0 :: Int) :. (div numPoints 2) :. (div numPoints 2 - 1))))
  plotImageRepaComplex (folderPath </> "Pinwheel_Coef.png") .
    ImageRepa 8 .
    fromUnboxed (Z :. (1 :: Int) :. numR2Freq :. numR2Freq) . VS.convert $
    (harmonicsArray' IA.! idx)
  printf
    "filter center: %f\n"
    (magnitude $
     (hpfR2' R.! (Z :. (0 :: Int) :. (div numPoints 2) :. (div numPoints 2 - 1))))
  plotImageRepaComplex (folderPath </> "HighPass.png") .
    ImageRepa 8 .
    computeS . extend (Z :. (1 :: Int) :. All :. All) . sumS . rotate3D $
    hpfR2'
  plotImageRepaComplex (folderPath </> "PinwheelHollow.png") .
    ImageRepa 8 .
    computeS .
    extend (Z :. (1 :: Int) :. All :. All) .
    sumS . rotate3D . R.zipWith (-) pinwheelR2 $
    pinwheelHPR2


{-# OPTIONS_GHC -Wall #-}
{-# LANGUAGE CPP #-}

{-|

  Module      : AutoBench.Internal.Regression
  Description : Generating and fitting models for regression analysis.
  Copyright   : (c) 2018 Martin Handley
  License     : BSD-style
  Maintainer  : martin.handley@nottingham.ac.uk
  Stability   : Experimental
  Portability : GHC

  This module is responsible for generating and fitting regression models.
  See 'LinearType', 'LinearCandidate' and 'LinearFit' for more details.

  The system uses ridge regression because it has a closed-form solution so
  is fast. Previously used LASSO implemented using the HVX library but
  it was too computationally expensive, (see bottom of this file).

-}

{-
   ----------------------------------------------------------------------------
   <TO-DO>:
   ----------------------------------------------------------------------------
   - Too much white space;
   -
-}

module AutoBench.Internal.Regression
  (

  -- * Candidate generation
    generateLinearCandidate    -- Generate a 'LinearCandidate' from a 'LinearType' by filling
                               -- in the details of the model necessary to fit it to a given dataset.
  -- * Regression analysis
  , fitOLS                     -- Given a data set (a list of xy-'Coord's) and a 'LinearCandidate',
                               -- fit the model to the data using the OLS method.
  , fitRidgeRegress            -- Given a data set (a list of xy-'Coord's) and a 'LinearCandidate',
                               -- fit the model to the data using ridge regression.

  ) where

import           Data.List            (minimumBy, nub)
import           Data.Ord             (comparing)
import qualified Data.Vector.Storable as V

import Numeric.LinearAlgebra
  ( Element
  , Matrix
  , Vector
  , (<>)
  , (<\>)
  , (#>)
  , asColumn
  , cmap
  , diag
  , dot
  , fromColumns
  , fromRows
  , konst
  , norm_2
  , norm_Inf
  , qr
  , rows
  , scale
  , subVector
  , sumElements
  , takeRows
  , toColumns
  , tr'
  , vjoin
  )

import qualified AutoBench.Internal.Expr as E
import AutoBench.Internal.Types
  ( Coord
  , Exp
  , LinearCandidate(..)
  , LinearType(..)
  , numPredictors
  )

#ifdef MIN_VERSION_base
#if MIN_VERSION_base(4,11,0)
import Prelude hiding ((<>))
#endif
#endif

-- * 'LinearCandidate' generation

-- | Generate a 'LinearCandidate' from a 'LinearType' by filling in the
-- details of the model necessary to fit it to a given dataset.
-- Makes use of the generation helper functions 'generateLinearCandidateExpr'
-- and 'generateLinearCandidateYHat'.
generateLinearCandidate :: LinearType -> LinearCandidate
generateLinearCandidate t@(Poly _) =
  LinearCandidate
   {
     _lct   = t
   , _fxs   = id
   , _fex   = generateLinearCandidateExpr t
   , _fyhat = generateLinearCandidateYHat t
   }
generateLinearCandidate t@(Log b _) =
  LinearCandidate
   {
     _lct   = t
   , _fxs   = cmap $ logBase (fromIntegral b)
   , _fex   = generateLinearCandidateExpr t
   , _fyhat = generateLinearCandidateYHat t
   }
generateLinearCandidate t@(PolyLog b _) =
  LinearCandidate
   {
     _lct   = t
   , _fxs   = cmap (\x -> x * logBase (fromIntegral b) x)
   , _fex   = generateLinearCandidateExpr t
   , _fyhat = generateLinearCandidateYHat t
   }
generateLinearCandidate t@(Exp n) =
  LinearCandidate
   {
     _lct   = t
   , _fxs   = cmap (flip (**) $ fromIntegral n)
   , _fex   = generateLinearCandidateExpr t
   , _fyhat = generateLinearCandidateYHat t
   }

-- * Linear fitting

-- | Given a data set (a list of xy-'Coord's) and a 'LinearCandidate',
-- fit the model to the data using the ordinary least squares (OLS) method. The
-- result is the model's coefficients.
fitOLS :: [Coord] -> LinearCandidate -> Vector Double
fitOLS coords c = head $ toColumns $ (_preds <\> _resps)
  where (_preds, _resps) = genFitMatrices coords c

-- | Given a data set (a list of xy-'Coord's) and a 'LinearCandidate',
-- fit the model to the data using ridge regression. The result is the model's
-- coefficients.
fitRidgeRegress :: [Coord] -> LinearCandidate -> Vector Double
fitRidgeRegress coords c =
  -- Select the lambda that minimises mean squared error.
  fst $ minimumBy (comparing snd) $ zip fits errs
  where
    n        = length coords
    (xs, ys) = unzip  coords
    vXs      = V.fromList xs
    vYs      = V.fromList ys
    -- Fitting error is mean squared error.
    calcMse m vec = ((norm_2 $ V.zipWith (-) vYs vec) ** 2.0) / fromIntegral m
    -- Generate the matrices to solve.
    (_preds, _resps) = genFitMatrices coords c
    ---------------------------------------------------------------------------
    ---------------------------------------------------------------------------
    -- Using ridge regression with maximum tuning parameter:
    -- https://stats.stackexchange.com/questions/289075/what-is-the-
    -- smallest-lambda-that-gives-a-0-component-in-lasso?noredirect=1&lq=1
    ---------------------------------------------------------------------------
    maxLam   = (1.0 / fromIntegral n) * norm_Inf (tr' _preds <> _resps)
    -- Minimise 'mse' to select best tuning parameter.
    lamStep  = maxLam / 10
    lamRange = [lamStep, 2 * lamStep .. maxLam]
    fits     = nub $ fmap (head
                            . toColumns
                            . ridgeRegress _preds _resps
                          ) lamRange
    -- ** Do not use vFxs here, use vXs **
    errs     = fmap (calcMse n . flip (_fyhat c) vXs) fits
    ---------------------------------------------------------------------------
    ---------------------------------------------------------------------------
    -- Closed form solution of ridge regression.
    -- https://github.com/chris-taylor/aima-haskell/blob/master/src/AI/
    -- Learning/LinearRegression.hs
    ---------------------------------------------------------------------------
    ridgeRegress
      :: Matrix Double  -- X
      -> Matrix Double  -- Y
      -> Double         -- λ
      -> Matrix Double  -- (X'X + λI)^−1 X'Y.
    ridgeRegress x y lambda =
      let m      = rows x
          (_, r) = qr x
          rr     = takeRows m r
          ww     = diag $ vjoin [0, konst 1 (m - 1)]
      in (tr' rr <> rr + (lambda `scale` ww)) <\> (tr' x <> y)

-- * Helpers

-- | Generate an expression ('Exp') for a 'LinearCandidate' that will allow
-- its model's equation to be pretty printed at a later time. Important note:
-- the vectors are assumed to be of the correct form, /no/  checking is
-- performed currently.
generateLinearCandidateExpr
   :: LinearType      -- Model.
   -> Vector Double   -- Coefficients of the model.
   -> Exp             -- Model's equation as an expression for pretty printing.
--
-- Given coefficients:
--
--   a_0, a_1, ..., a_n
--
-- Generate:
--
--   a_0 + a_1 * x^1 + a_2 * x^2 + ... + a_n * x^n
--
generateLinearCandidateExpr (Poly n) coeffs =
  foldr1 E.Add . take (n + 1) $ zipWith ($) terms (V.toList coeffs)
  where terms = E.Num                                         -- a_0
                 : ( \a_k -> E.Mul (E.Num a_k) (E.Var "x")    -- a_1 * x
                   )
                 : fmap ( \p a_k -> E.Mul (E.Num a_k)         -- a_k * x^k, k > 1
                                          (E.Pow (E.Var "x")
                                                 (E.Num p)
                                          )
                        ) [2.0..]
--
-- Given coefficients:
--
-- a_0, a_1, ..., a_n
--
-- Generate:
--
-- a_0 + a_1 * log_b^1(x) + a_2 * log_b^2(x) + ... + a_n * log_b^n(x)
--
generateLinearCandidateExpr (Log b n) coeffs =
  foldr1 E.Add $ take (n + 1) $ zipWith ($) terms (V.toList coeffs)
  where terms = E.Num                                                -- a_0
                 : ( \a_k -> E.Mul (E.Num a_k)                       -- a_1 * log_b(x)
                                   (E.Log (E.Num $ fromIntegral b)
                                          (E.Var "x")
                                   )
                   )
                 : fmap ( \p a_k -> E.Mul (E.Num a_k)                -- a_k * log_b^k(x), k > 1
                                          (E.Pow (E.Log (E.Num $ fromIntegral b)
                                                        (E.Var "x")
                                                 )
                                                 (E.Num p)
                                          )
                        ) [2.0..]
--
-- Given coefficients:
--
-- a_0, a_1, ..., a_n
--
-- Generate:
--
-- a_0 + a_1 * x^1 * log_b^1(x) + a_2 * x^2 * log_b^2(x) + ... + a_n * x^n * log_b^n(x)
--
generateLinearCandidateExpr (PolyLog b n) coeffs =
  foldr1 E.Add $ take (n + 1) $ zipWith ($) terms (V.toList coeffs)
  where terms = E.Num                                                -- a_0
                 : ( \a_k -> E.Mul (E.Num a_k)                       -- a_1 * x * log_b(x)
                                   (E.Mul (E.Var "x")
                                          (E.Log (E.Num $ fromIntegral b)
                                                 (E.Var "x")
                                          )
                                   )
                   )
                 : fmap ( \p a_k -> E.Mul (E.Num a_k)                -- a_k * x^k * log_b^k(x), k > 1
                                          (E.Mul (E.Pow (E.Var "x")
                                                        (E.Num p)
                                                 )
                                                 (E.Pow (E.Log (E.Num $ fromIntegral b)
                                                               (E.Var "x")
                                                        )
                                                        (E.Num p)
                                                 )
                                          )
                        ) [2.0..]
--
-- Given coefficients:
--
-- a_0, a_1
--
-- Generate:
--
-- a_0 + a_1 * n^x
--
generateLinearCandidateExpr (Exp n) coeffs =
  foldr1 E.Add $ zipWith ($) terms (V.toList coeffs)
  where terms = [ E.Num  -- a_0
                , \a_1 -> E.Mul (E.Num a_1)   -- a_1 *
                                (E.Pow (E.Num $ fromIntegral n) -- n^x
                                        (E.Var "x")
                                )
                ]

-- | For a given set of x-coordinates, generate the y-coordinates that are
-- predicted by the model of a 'LinearCandidate'. Important note: the vectors
-- are assumed to be of the correct form, /no/ checking is performed currently.
generateLinearCandidateYHat
  :: LinearType
  -> Vector Double   -- Coefficients of the model.
  -> Vector Double   -- x-coordinates.
  -> Vector Double   -- Predicted y-coordinates.
--
-- Given coefficients and x-coords:
--
--   a_0, a_1, ..., a_n     x_0, x_1, ..., x_m
--
-- For each x_i, generate:
--
--   a_0 + a_1 * x_i^1 + a_2 * x_i^2 + ... + a_n * x_i^n
--
generateLinearCandidateYHat (Poly n) coeffs xs =
  (fromColumns $ V.fromList (replicate m 1.0) : xs') #> coeffs
  where
    m   = V.length xs
    xs' = zipWith vecPow [1.0..] (replicate n xs)
--
-- Given coefficients and x-coords:
--
-- a_0, a_1, ..., a_n     x_0, x_1, ..., x_m
--
-- For each x_i, generate:
--
-- a_0 + a_1 * log_b^1(x_i) + a_2 * log_b^2(x_i) + ... + a_n * log_b^n(x_i)
--
generateLinearCandidateYHat (Log b n) coeffs xs =
  (fromColumns $ V.fromList (replicate m 1.0) : xs') #> coeffs
  where
    m   = V.length xs
    xs' = zipWith (vecLogPow $ fromIntegral b) [1.0..] (replicate n xs)
--
-- Given coefficients and x-coords:
--
-- a_0, a_1, .., a_n     x_0, x_1, ..., x_m
--
-- For each x_i, generate:
--
-- a_0 + a_1 * x_i^1 * log_b^1(x_i) + a_2 * x_i^2 * log_b^2(x_i) + ... + a_n * x_i^n * log_b^n(x_i)
--
generateLinearCandidateYHat (PolyLog b n) coeffs xs =
  (fromColumns $ V.fromList (replicate m 1.0) : xs') #> coeffs
  where
    m   = V.length xs
    xs' = zipWith (vecPolyLogPow $ fromIntegral b) [1.0..] (replicate n xs)
--
-- Given coefficients and x-coords:
--
-- a_0, a_1      x_0, x_1, ..., x_m
--
-- For each x_i, generate:
--
-- a_0 + a_1 * n^x_i
--
generateLinearCandidateYHat (Exp n) coeffs xs =
  (fromColumns $ V.fromList (replicate m 1.0) : xs') #> coeffs
  where
    m   = V.length xs
    xs' = zipWith (vecBasePow $ fromIntegral n) [1.0..] (replicate 1 xs)

-- | Constructs predictor (@preds@) and responder (@resps@) matrices so that
-- regression equations can be solved in their matrix form.
genFitMatrices :: [Coord] -> LinearCandidate -> (Matrix Double, Matrix Double)
genFitMatrices coords c = gen (numPredictors (_lct c) - 1)
  where

    gen :: Int -> (Matrix Double, Matrix Double)
    gen k = let
               _xPows    = xPows k
               _xSumPows = xSumPows _xPows
               _ySums    = ySums _xPows k
               _preds    = preds _xSumPows k
               _resps    = resps _ySums
            in (_preds, asColumn _resps)

    ---------------------------------------------------------------------------
    --
    --
    (xs, ys) = unzip coords
    -- Transform the x-coords using the candidate's 'fxs' function. This allows
    -- us to fit different models to the data in a generic way. For example, a
    -- 'Log b n' candidate will first take 'LogBase b' of the x-coords. We can
    -- then use (a variant of) linear least squares to fit a straight line
    -- through the resulting data points. Although this is a straight line on
    -- the /resulting data/, it corresponds to a logarithmic function on the
    -- /initial/ data. This technique is standard.
    vFxs = (_fxs c) (V.fromList xs)
    vYs  = V.fromList ys
    -- Number of data points.
    n    = length coords
    -- As the x-'Coords' have been transformed by the candidate's 'fxs'
    -- function, this method works generically for any model and 'k'
    -- coefficients.
    ---------------------------------------------------------------------------
    --
    -- Helper functions to construct the matrices to solve using linear
    -- regression. Note: these functions are to be used in a particular way so
    -- their names reflect their use.
    --
    -- x_i^j for j = 1, 2, .., 2 * k
    xPows :: Int -> [Vector Double]
    xPows k = zipWith vecPow [1.0..] (replicate (2 * k) vFxs)
    -- SUM x_i^j for j = 1, 2, .., 2 * k
    xSumPows :: [Vector Double] -> [Double]
    xSumPows  = fmap sumElements
    -- SUM y_i, SUM x_i * y_i, SUM x_i^2 * y_i, .., SUM x_i^j * y_i
    ySums :: [Vector Double] -> Int -> [Double]
    ySums vs k = zipWith dot (V.fromList (replicate n 1.0) : vs) (replicate (k + 1) vYs)
    --
    -- Convert xSumPows into a matrix by 'chunking' off rows from the initial vector.
    --
    -- Rows are length k + 1, the initial rows are:
    --
    --   n      , SUM x_i  , SUM x_i^2, .., SUM x_i^k
    --
    --   SUM x_i, SUM x_i^2, SUM x_i^3, .., SUM x_i^(k + 1)
    --
    -- The last row is:
    --
    --   SUM x_i^j, SUM x_i^(j + 1), .., SUM x_i^(2 * k)
    --
    -- These values are the /predictors/.
    --
    preds :: [Double] -> Int -> Matrix Double
    preds ds k = vecToMatrix (V.fromList $ fromIntegral n : ds) (k + 1)
    -- 'ySums' are the /responders/.
    resps :: [Double] -> Vector Double
    resps  = V.fromList

-- | Raise each double in a vector of doubles to a given power.
vecPow :: Double -> Vector Double -> Vector Double
vecPow pow = cmap (** pow)

-- | For each double 'x' in a vector of doubles, apply the following
-- calculation:
--
-- > \x -> base ** (x * pow)
--
-- for a given 'base' and 'pow'.
vecBasePow :: Double -> Double -> Vector Double -> Vector Double
vecBasePow base pow = cmap (\x -> base ** (x * pow))

-- | For each double 'x' in a vector of doubles, apply the following
-- calculation:
--
-- > \x -> (x ** pow) * (logBase b x ** pow)
--
-- for a given 'b' and 'pow'.
vecPolyLogPow :: Double -> Double -> Vector Double -> Vector Double
vecPolyLogPow b pow = cmap (\x -> (x ** pow) * ((logBase b x) ** pow))

-- | For each double 'x' in a vector of doubles, apply the following
-- calculation:
--
-- > \x -> logBase b x ** pow
--
-- for a given 'b' and 'pow'.
vecLogPow :: Double -> Double -> Vector Double -> Vector Double
vecLogPow b pow = cmap (\x -> (logBase b x) ** pow)

-- | Create a matrix by 'chunking' off rows of a given length from a vector.
-- Works similarly to 'tails'.
vecToMatrix :: (V.Storable a, Element a) => Vector a -> Int -> Matrix a
vecToMatrix v j = fromRows [ subVector i j v | i <- [0..j - 1] ]

{-

LASSO has been retired for the time being due to computational expense.

{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE DataKinds        #-}

import           HVX.Internal.Constraints ((<=~))
import           HVX.Primitives           ((+~), (*~), neg, norm, quadform)
import           HVX.Internal.Solvers     (ellipsoidMinimize)
import qualified HVX.Internal.Primitives  as HVX

-- ** LASSO (Least Absolute Shrinkage and Selection Operator)

fitLASSO :: [Coord] -> LinearCandidate -> Vector Double
fitLASSO coords c =
  -- Select the lambda that minimises mean squared error.
  fst $ minimumBy (comparing snd) $ zip fits errs
  where
    n        = length coords
    (xs, ys) = unzip  coords
    vXs      = V.fromList xs
    vYs      = V.fromList ys
    -- Fitting error is mean squared error.
    calcMse m vec = ((norm_2 $ V.zipWith (-) vYs vec) ** 2.0) / fromIntegral m
    -- Generate the matrices to solve.
    (_preds, _resps) = genFitMatrices coords c
    ---------------------------------------------------------------------------
    ---------------------------------------------------------------------------
    initCoeffs = _preds <\> _resps
    maxT = V.sum $ cmap abs $ flatten initCoeffs
    -- Minimise 'mse' to select best tuning parameter.
    tStep  = maxT / 10
    tRange = [0 + tStep, 0 + (2 * tStep) .. maxT]
    fits   = filter (all (\m -> not (isNaN m || isInfinite m)) . V.toList)
              . nub
              $ fmap (head
                       . toColumns
                       . lassoRegress _preds _resps
                     ) tRange
    -- ** Do not use vFxs here, use vXs **
    errs   = fmap (calcMse n . flip (fyhat c) vXs) fits
    ---------------------------------------------------------------------------
    ---------------------------------------------------------------------------
    -- LASSO regression (convex equation using the subgradient method in HVX).
    lassoRegress
      :: Matrix Double      -- ^ n >< n matrix X.
      -> Matrix Double      -- ^ n >< 1 matrix Y.
      -> Double             -- ^ Constraint on sum of coefficients.
      -> Matrix Double      -- ^ Optimal variable values.
    lassoRegress x y t = snd . head $ res
      where
        (res, _, _) = ellipsoidMinimize expr [constr] [("a", m)] 1e-20 (2 * maxT)
        m      = rows x
        constr = norm 1 (HVX.EVar "a") <=~ (HVX.EConst $ konst t (1, 1))
        expr   = quadform (HVX.EConst $ ident m) $ ((HVX.EConst y) +~ neg (HVX.EConst x *~ HVX.EVar "a"))

-}


{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE TemplateHaskell   #-}
module DBPnet.Correlation
    ( cisCorMat
    , transCorMat
    ) where

import           Bio.Data.Bed
import           Conduit
import Lens.Micro
import           Control.Monad
import qualified Data.ByteString.Char8       as B
import           Data.List.Split
import qualified Data.Matrix.Unboxed         as MU
import qualified Data.Matrix.Unboxed.Mutable as MUM
import qualified Data.Vector                 as V
import qualified Data.Vector.Unboxed         as U
import           Statistics.Correlation      (spearman)

import           DBPnet.Utils                (readLoops)

cisCorMat :: [(String, FilePath)]
          -> IO ([B.ByteString], MU.Matrix Double)
cisCorMat dat = do
    counts <- fmap V.fromList $ forM dat $ \(i, fl) -> do
        beds <- readBed fl
        return (i, U.fromList $ map (^.bdgValue) beds)

    let l = V.length counts
        header = map B.pack $ V.toList $ fst . V.unzip $ counts
        mat :: MU.Matrix Double
        mat = MU.create $ do
            m <- MUM.new (l,l)
            forM_ [0..l-1] $ \i -> do
                let (_, tfA) = counts V.! i
                forM_ [0..l-1] $ \j -> case () of
                    _ | j < i -> MUM.read m (j,i) >>= MUM.write m (i,j)
                      | otherwise -> do
                          let (_, tfB) = counts V.! j
                              vec = U.filter (\(a,b) -> a /= 0 || b /= 0) $ U.zip tfA tfB
                              cor | U.length vec < 500 = 0
                                  | otherwise = spearman vec
                          MUM.write m (i,j) cor
            return m
    return (header, mat)

transCorMat :: [(String, FilePath)] -> FilePath
            -> IO ([B.ByteString], MU.Matrix Double)
transCorMat dat hic = do
    loops <- concatMap (\(a,b) -> [a,b]) <$> readLoops hic

    counts <- fmap V.fromList $ forM dat $ \(i,fl) -> do
        beds <- readBed fl :: IO [BEDGraph]
        let rs = runIdentity $ runConduit $ yieldMany loops .|
                intersectBedWith fn beds .| sinkList
            fn _ [] = 0
            fn _ x = maximum $ map (^.bdgValue) x
        return (i, U.fromList $ map (\[a,b] -> (a,b)) $ chunksOf 2 rs)

    let l = V.length counts
        header = map B.pack $ V.toList $ fst . V.unzip $ counts
        mat :: MU.Matrix Double
        mat = MU.create $ do
            m <- MUM.new (l,l)
            forM_ [0..l-1] $ \i -> do
                let (_, tfA) = counts V.! i
                forM_ [0..l-1] $ \j -> case () of
                    _ | j < i -> MUM.read m (j,i) >>= MUM.write m (i,j)
                      | otherwise -> do
                          let (_, tfB) = counts V.! j
                              vec = U.filter (\(a,b) -> a /= 0 || b /= 0) $ U.zipWith f tfA tfB
                              cor | U.length vec < 500 = 0
                                  | otherwise = spearman vec
                          MUM.write m (i,j) cor
            return m
    return (header, mat)
  where
    f (x1,x2) (y1,y2) | x1 > x2 = (x1, y2)
                      | otherwise = (x2, y1)


module Forward (
    forwardClassification,
    forwardRegression,
    forward
  ) where

import Numeric.LinearAlgebra
import Common
import ActivFunc

forwardClassification :: (Matrix R -> Matrix R) -> [Matrix R] -> Matrix R -> Matrix R
forwardClassification f ws v = softmaxC $ forwardRegression f ws v

forwardRegression :: (Matrix R -> Matrix R) -> [Matrix R] -> Matrix R -> Matrix R
forwardRegression f ws v = forward id (last ws) nv
  where nv = forwards f (init ws) v

forwards :: (Matrix R -> Matrix R) -> [Matrix R] -> Matrix R -> Matrix R
forwards f ws v = foldl (flip $ forward f) v ws

forward :: (Matrix R -> Matrix R) -> Matrix R -> Matrix R -> Matrix R
forward f w v = f $ w <> inputWithBias v


module Q.Risk.VaR where

import Statistics.Autocorrelation
import Data.Csv
import Data.Text (Text)
import qualified Data.ByteString.Lazy as B
import qualified Data.Vector.Storable as VS
import qualified Data.Vector as V
import qualified Numeric.GSL.Statistics as GSL

y :: IO [Double]
y = do
  bytes <- B.readFile "/home/ghais/Downloads/TableF20-1.csv"
  return $ case (decode HasHeader bytes) of
    Left err -> error err
    Right v -> V.toList (V.map fromOnly v)

ar = do
  xs <- demean . VS.fromList <$> y
  let xs2 = VS.map (^2) xs
  return $ GSL.lag1auto  xs2
  where demean xs = let mean = GSL.mean xs
                    in VS.map (\x -> x - mean) xs


module Main where

import RayLib
import Numeric.LinearAlgebra
import qualified Graphics.Image as G

main :: IO ()
main = do let screen = Screen 256 256
          let ka = MatParam (fromList [0,1,1])
          let ks = MatParam (fromList [1,1,0])
          let kd = MatParam (fromList [0.4,0.4,0])
          let ke = MatParam (fromList [0,0,0])
          let kr = MatParam (fromList [0,0,0])
          let kt = MatParam (fromList [0,0,0])
          let shininess = 64
          let material = Material ka ks kd ke kr kt shininess
          let sphere1 = Sphere (fromList [0,0,0]) 1.5 material
          let scene = Scene [] [sphere1] (fromList [1,0,0.1])
          putStrLn . show $ scene
          let camera = Camera (fromList [0,0,-4]) (fromList [0,0,1]) (fromList [0,1,0]) 1 90
          putStrLn . show $ camera
          putStrLn . show $ screen
          let image = traceImage scene screen camera
          let toWrite = render image
          G.writeImage "test.png" toWrite


{-# LANGUAGE OverloadedStrings, ExistentialQuantification, FlexibleContexts #-}

module BNN where

import System.Random.MWC (createSystemRandom)

--import Data.Aeson (ToJSON(toJSON), Value)
--import Data.Aeson (encode)
import qualified Data.ByteString.Lazy.Char8 as BL

import Data.List (sort, partition)
import qualified Data.Histogram as DH
import qualified Data.Histogram.Fill as DF
import qualified Data.Histogram.Bin.Bin2D as Bin2D
import Data.Vector.Unboxed.Base (Unbox)
import qualified Data.Vector.Unboxed as Vec

import Control.Monad (liftM2, replicateM, forM, forM_)
import Control.Monad.IO.Class (liftIO)
import Control.Monad.Bayes.Class
import Control.Monad.Bayes.Sampler
import Control.Monad.Bayes.Traced
import Control.Monad.Bayes.Weighted
import Control.Monad.Bayes.Inference.SMC as SMC
import Control.Monad.Bayes.Inference.RMSMC as RMSMC
import Control.Monad.Bayes.Sequential
import Control.Monad.Bayes.Population
import Control.Monad.Bayes.Traced.Static (Traced)
import Numeric.LinearAlgebra (Matrix, Vector, toList, vector, matrix, dot, (#>), (!), cmap, scalar, size)

import Numeric.Log

-- Plotting

histo2D :: (Foldable f, Unbox val, Num val) =>
    (Double, Int, Double)
    -> (Double, Int, Double)
    -> f (Double, Double)
    -> DH.Histogram (Bin2D.Bin2D DH.BinD DH.BinD) val
histo2D (xmin, xn, xmax) (ymin, yn, ymax) = DF.fillBuilder buildr
  where
    binX = DH.binD xmin xn xmax
    binY = DH.binD ymin yn ymax
    bins = Bin2D.Bin2D binX binY
    buildr = DF.mkSimple bins

histo :: (Foldable v, Unbox a, Num a) =>
         (Double, Int, Double)
      -> v Double
      -> DH.Histogram DF.BinD a
histo (xmin, n, xmax) = DF.fillBuilder buildr
  where
    bins = DH.binD xmin n xmax
    buildr = DF.mkSimple bins

-- A SIMPLE NEURAL NETWORK
-- Model Setup

data NN
  = NN
      { biass :: Vector Double,
        weights :: Vector Double,
        sigma :: Double
      }
  deriving (Eq, Show)

data Data
  = Data
      { xValue :: Double,
        yValue :: Double
      }
  deriving (Eq, Show)

forwardNN :: NN -> Double -> Double
forwardNN (NN bs ws _) x =
    ws `dot` cmap activation (scalar x - bs)
  where activation x = if x < 0 then 0 else 1

errorModel :: Double -> Double -> Double -> Log Double
errorModel mean std = normalPdf mean std

likelihood :: NN -> Data -> Log Double
likelihood nn (Data xObs yObs) =
    errorModel yMean ySigma yObs
  where
    ySigma = sigma nn
    yMean = forwardNN nn xObs

-- Prior, Posterior and Predictive Distribution

postNN :: MonadInfer m => m NN -> [Data] ->  m NN
postNN pr obs = do
  nn <- pr
  forM_ obs (score . likelihood nn)
  return nn

uniformVec :: MonadSample m => (Double, Double) -> Int -> m (Vector Double)
uniformVec (wmin, wmax) nelements =
  vector <$> replicateM nelements (uniform wmin wmax)

priorNN :: MonadSample m => Int -> m NN
priorNN nnodes = do
  bias <- uniformVec (0, 10) nnodes
  weight <- uniformVec (-10, 10) nnodes
  sigma <- uniform 0.5 1.5
  return $ NN bias weight sigma

predDist :: MonadInfer m => m NN -> m (NN, Data)
predDist pr = do
  nn <- pr
  x <- uniform 0 10
  y <- uniform (-5) 10
  score $ likelihood nn (Data x y)
  return (nn, Data x y)


module AI.Learning.Bootstrap where

import Control.Monad.Random
import Foreign.Storable (Storable)
import Numeric.LinearAlgebra
import qualified Data.List as L

import AI.Util.Matrix
import AI.Util.Util

---------------
-- Bootstrap --
---------------

-- |Generate a bootstrap sample of size @sz@.
genBootstrapSample :: RandomGen g => Int -> Rand g [Int]
genBootstrapSample sz = go sz []
    where go 0 accum = return accum
          go n accum = do
            i <- getRandomR (0,sz-1)
            go (n - 1) (i:accum)

sampleVector :: (Storable a, RandomGen g) => Vector a -> Rand g (Vector a)
sampleVector v = do
    idx <- genBootstrapSample (dim v)
    return (v `subRefVec` idx)

sampleMatrixRows :: (Element a, RandomGen g) => Matrix a -> Rand g (Matrix a)
sampleMatrixRows m = do
    idx <- genBootstrapSample (rows m)
    return $ m `subRefRows` idx

sampleMatrixCols :: (Element a, RandomGen g) => Matrix a -> Rand g (Matrix a)
sampleMatrixCols m = do
    idx <- genBootstrapSample (cols m)
    return $ m `subRefCols` idx

-- |Generate a bootstrap sample of a statistic from a data set of type /a/.
bootStrapResample :: RandomGen g =>
                     (a -> Rand g a)    -- Sampling function
                  -> (a -> b)           -- Statistic to be resampled
                  -> Int                -- Number of resamples
                  -> a                  -- Data
                  -> Rand g [b]
bootStrapResample sample func nSamples x = go [] nSamples
    where
        go samples 0 = return samples
        go samples n = do
            x' <- sample x
            go (func x' : samples) (n-1)


{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE DeriveDataTypeable, DeriveGeneric #-}
-- |
-- Module    : Statistics.Distribution.Normal
-- Copyright : (c) 2009 Bryan O'Sullivan
-- License   : BSD3
--
-- Maintainer  : bos@serpentine.com
-- Stability   : experimental
-- Portability : portable
--
-- The normal distribution.  This is a continuous probability
-- distribution that describes data that cluster around a mean.

module Statistics.Distribution.Normal
    (
      NormalDistribution
    -- * Constructors
    , normalDistr
    , normalDistrE
    , normalDistrErr
    , standard
    ) where

import Control.Applicative
import Data.Aeson            (FromJSON(..), ToJSON, Value(..), (.:))
import Data.Binary           (Binary(..))
import Data.Data             (Data, Typeable)
import GHC.Generics          (Generic)
import Numeric.MathFunctions.Constants (m_sqrt_2, m_sqrt_2_pi)
import Numeric.SpecFunctions (erfc, invErfc)
import qualified System.Random.MWC.Distributions as MWC
import qualified Data.Vector.Generic as G

import qualified Statistics.Distribution as D
import qualified Statistics.Sample as S
import Statistics.Internal

-- | The normal distribution.
data NormalDistribution = ND {
      mean       :: {-# UNPACK #-} !Double
    , stdDev     :: {-# UNPACK #-} !Double
    , ndPdfDenom :: {-# UNPACK #-} !Double
    , ndCdfDenom :: {-# UNPACK #-} !Double
    } deriving (Eq, Typeable, Data, Generic)

instance Show NormalDistribution where
  showsPrec i (ND m s _ _) = defaultShow2 "normalDistr" m s i
instance Read NormalDistribution where
  readPrec = defaultReadPrecM2 "normalDistr" normalDistrE

instance ToJSON NormalDistribution
instance FromJSON NormalDistribution where
  parseJSON (Object v) = do
    m  <- v .: "mean"
    sd <- v .: "stdDev"
    either fail return $ normalDistrErr m sd
  parseJSON _ = empty

instance Binary NormalDistribution where
    put (ND m sd _ _) = put m >> put sd
    get = do
      m  <- get
      sd <- get
      either fail return $ normalDistrErr m sd

instance D.Distribution NormalDistribution where
    cumulative      = cumulative
    complCumulative = complCumulative

instance D.ContDistr NormalDistribution where
    logDensity    = logDensity
    quantile      = quantile
    complQuantile = complQuantile

instance D.MaybeMean NormalDistribution where
    maybeMean = Just . D.mean

instance D.Mean NormalDistribution where
    mean = mean

instance D.MaybeVariance NormalDistribution where
    maybeStdDev   = Just . D.stdDev
    maybeVariance = Just . D.variance

instance D.Variance NormalDistribution where
    stdDev = stdDev

instance D.Entropy NormalDistribution where
  entropy d = 0.5 * log (2 * pi * exp 1 * D.variance d)

instance D.MaybeEntropy NormalDistribution where
  maybeEntropy = Just . D.entropy

instance D.ContGen NormalDistribution where
    genContVar d = MWC.normal (mean d) (stdDev d)

-- | Standard normal distribution with mean equal to 0 and variance equal to 1
standard :: NormalDistribution
standard = ND { mean       = 0.0
              , stdDev     = 1.0
              , ndPdfDenom = log m_sqrt_2_pi
              , ndCdfDenom = m_sqrt_2
              }

-- | Create normal distribution from parameters.
--
-- IMPORTANT: prior to 0.10 release second parameter was variance not
-- standard deviation.
normalDistr :: Double            -- ^ Mean of distribution
            -> Double            -- ^ Standard deviation of distribution
            -> NormalDistribution
normalDistr m sd = either error id $ normalDistrErr m sd

-- | Create normal distribution from parameters.
--
-- IMPORTANT: prior to 0.10 release second parameter was variance not
-- standard deviation.
normalDistrE :: Double            -- ^ Mean of distribution
             -> Double            -- ^ Standard deviation of distribution
             -> Maybe NormalDistribution
normalDistrE m sd = either (const Nothing) Just $ normalDistrErr m sd

-- | Create normal distribution from parameters.
--
normalDistrErr :: Double            -- ^ Mean of distribution
               -> Double            -- ^ Standard deviation of distribution
               -> Either String NormalDistribution
normalDistrErr m sd
  | sd > 0    = Right $ ND { mean       = m
                           , stdDev     = sd
                           , ndPdfDenom = log $ m_sqrt_2_pi * sd
                           , ndCdfDenom = m_sqrt_2 * sd
                           }
  | otherwise = Left $ errMsg m sd

errMsg :: Double -> Double -> String
errMsg _ sd = "Statistics.Distribution.Normal.normalDistr: standard deviation must be positive. Got " ++ show sd

-- | Variance is estimated using maximum likelihood method
--   (biased estimation).
--
--   Returns @Nothing@ if sample contains less than one element or
--   variance is zero (all elements are equal)
instance D.FromSample NormalDistribution Double where
  fromSample xs
    | G.length xs <= 1 = Nothing
    | v == 0           = Nothing
    | otherwise        = Just $! normalDistr m (sqrt v)
    where
      (m,v) = S.meanVariance xs

logDensity :: NormalDistribution -> Double -> Double
logDensity d x = (-xm * xm / (2 * sd * sd)) - ndPdfDenom d
    where xm = x - mean d
          sd = stdDev d

cumulative :: NormalDistribution -> Double -> Double
cumulative d x = erfc ((mean d - x) / ndCdfDenom d) / 2

complCumulative :: NormalDistribution -> Double -> Double
complCumulative d x = erfc ((x - mean d) / ndCdfDenom d) / 2

quantile :: NormalDistribution -> Double -> Double
quantile d p
  | p == 0         = -inf
  | p == 1         = inf
  | p == 0.5       = mean d
  | p > 0 && p < 1 = x * ndCdfDenom d + mean d
  | otherwise      =
    error $ "Statistics.Distribution.Normal.quantile: p must be in [0,1] range. Got: "++show p
  where x          = - invErfc (2 * p)
        inf        = 1/0

complQuantile :: NormalDistribution -> Double -> Double
complQuantile d p
  | p == 0         = inf
  | p == 1         = -inf
  | p == 0.5       = mean d
  | p > 0 && p < 1 = x * ndCdfDenom d + mean d
  | otherwise      =
    error $ "Statistics.Distribution.Normal.complQuantile: p must be in [0,1] range. Got: "++show p
  where x          = invErfc (2 * p)
        inf        = 1/0