module Statistics.Information.Utils.Random where

import Data.Matrix
import Statistics.Information.Utils.List
import System.Random

splitN :: RandomGen g => Int -> g -> [g]
splitN 0 _ = []
splitN 1 g = [g]
splitN n g = let (g', g'') = split g in
  g' : splitN (n-1) g''

rands :: (RandomGen g, Random a) => g -> Int -> (a, a) -> [a]
rands g n (a, b) = take n $ randomRs (a, b) g

noisy :: RandomGen g => g -> Matrix Double -> Matrix Double
noisy g xs =
  let noise = [repeatN d (intens * r) | r <- rands g n (0.0, fromIntegral n)] in
    elementwise (+) xs (fromLists noise)
  where
    intens = 1e-10
    n = nrows xs
    d = ncols xs


{-# LANGUAGE CPP #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE DefaultSignatures #-}
{-# LANGUAGE PolyKinds #-}
{-# LANGUAGE Trustworthy #-}
---------------------------------------------------------------------------
-- |
-- Copyright   :  (C) 2012-2015 Edward Kmett
-- License     :  BSD-style (see the file LICENSE)
--
-- Maintainer  :  Edward Kmett <ekmett@gmail.com>
-- Stability   :  experimental
-- Portability :  non-portable
--
-- Simple matrix operation for low-dimensional primitives.
---------------------------------------------------------------------------
module Linear.Trace
  ( Trace(..)
  , frobenius
  ) where

import Control.Monad as Monad
import Linear.V0
import Linear.V1
import Linear.V2
import Linear.V3
import Linear.V4
import Linear.Plucker
import Linear.Quaternion
import Linear.V
import Linear.Vector
import Data.Complex
import Data.Distributive
import Data.Foldable as Foldable
import Data.Functor.Bind as Bind
import Data.Functor.Compose
import Data.Functor.Product
import Data.Hashable
import Data.HashMap.Lazy
import Data.IntMap (IntMap)
import Data.Map (Map)

-- $setup
-- >>> import Data.Complex
-- >>> import Debug.SimpleReflect.Vars
-- >>> import Linear.V2

class Functor m => Trace m where
  -- | Compute the trace of a matrix
  --
  -- >>> trace (V2 (V2 a b) (V2 c d))
  -- a + d
  trace :: Num a => m (m a) -> a
#ifndef HLINT
  default trace :: (Foldable m, Num a) => m (m a) -> a
  trace = Foldable.sum . diagonal
  {-# INLINE trace #-}
#endif

  -- | Compute the diagonal of a matrix
  --
  -- >>> diagonal (V2 (V2 a b) (V2 c d))
  -- V2 a d
  diagonal :: m (m a) -> m a
#ifndef HLINT
  default diagonal :: Monad m => m (m a) -> m a
  diagonal = Monad.join
  {-# INLINE diagonal #-}
#endif

instance Trace IntMap where
  diagonal = Bind.join
  {-# INLINE diagonal #-}

instance Ord k => Trace (Map k) where
  diagonal = Bind.join
  {-# INLINE diagonal #-}

instance (Eq k, Hashable k) => Trace (HashMap k) where
  diagonal = Bind.join
  {-# INLINE diagonal #-}

instance Dim n => Trace (V n)
instance Trace V0
instance Trace V1
instance Trace V2
instance Trace V3
instance Trace V4
instance Trace Plucker
instance Trace Quaternion

instance Trace Complex where
  trace ((a :+ _) :+ (_ :+ b)) = a + b
  {-# INLINE trace #-}
  diagonal ((a :+ _) :+ (_ :+ b)) = a :+ b
  {-# INLINE diagonal #-}

instance (Trace f, Trace g) => Trace (Product f g) where
  trace (Pair xx yy) = trace (pfst <$> xx) + trace (psnd <$> yy) where
    pfst (Pair x _) = x
    psnd (Pair _ y) = y
  {-# INLINE trace #-}
  diagonal (Pair xx yy) = diagonal (pfst <$> xx) `Pair` diagonal (psnd <$> yy) where
    pfst (Pair x _) = x
    psnd (Pair _ y) = y
  {-# INLINE diagonal #-}

instance (Distributive g, Trace g, Trace f) => Trace (Compose g f) where
  trace = trace . fmap (fmap trace . distribute) . getCompose . fmap getCompose
  {-# INLINE trace #-}
  diagonal = Compose . fmap diagonal . diagonal . fmap distribute . getCompose . fmap getCompose
  {-# INLINE diagonal #-}

-- | Compute the <http://mathworld.wolfram.com/FrobeniusNorm.html Frobenius norm> of a matrix.
frobenius :: (Num a, Foldable f, Additive f, Additive g, Distributive g, Trace g) => f (g a) -> a
frobenius m = trace $ fmap (\ f' -> Foldable.foldl' (^+^) zero $ liftI2 (*^) f' m) (distribute m)


--
-- FullConnLayer : fully connected layer
--

module CNN.FullConnLayer (
  initFilterF
, zeroFilterF
, connect
, deconnect
, reverseFullConnFilter
, updateFullConnFilter
) where

import Control.Monad hiding (msum)
--import Data.List
import Data.Maybe
import Debug.Trace
import Numeric.LinearAlgebra
import System.Random.Mersenne as MT

import CNN.Algebra
import CNN.Image
import CNN.LayerType

{- |
initFilterF

  IN : #kernel
       #channel

  OUT: filter of fully connected layer

  OUT: filter of fully connected layer

>>> f <- initFilterF 3 20
>>> rows f
3
>>> cols f
21

-}

initFilterF :: Int -> Int -> IO FilterF
initFilterF k c = do
  rs <- forM [1..k] $ \i -> do
    w <- forM [1..c] $ \j -> do
      r <- MT.randomIO :: IO Double
      return ((r * 2.0 - 1.0) * a)
    return (0.0:w)
  return $ fromLists rs
  where
    --a = 1.0 / fromIntegral c
    a = 4.0 * sqrt (6.0 / fromIntegral (c + k))
    --a = sqrt (6.0 / fromIntegral (c + k))

zeroFilterF :: Int -> Int -> IO FilterF
zeroFilterF k c = return $ (k><(c+1)) $ repeat 0.0

--

{- |
connect

  IN : filter of fully connected layer
       image

  OUT: updated image

>>> let fs = fromLists [[0.5,1.0,2.0,3.0],[0.1,4.0,5.0,6.0]]
>>> let im = [fromLists [[9.0,8.0,7.0]]]
>>> connect fs im
[(1><2)
 [ 46.5, 118.1 ]]
>>> connect fs []
*** Exception: invalid Image
CallStack (from HasCallStack):
  error, called at CNN/FullConnLayer.hs:82:16 in main:CNN.FullConnLayer

-}

connect :: FilterF -> Image -> Image
connect _ [] = error "invalid Image"
connect fs [im] = [reshape (rows fs) $ fs #> v]
  where
    v = vjoin [konst 1.0 1, flatten im]

-- back prop

{- |
deconnect

  IN : filter of fully connected layer
       image
       difference from previous layer

  OUT: difference and updated layer

>>> let im = [fromLists [[1.0, 2.0, 3.0]]]
>>> let delta = [fromLists [[1.0, 2.0]]]
>>> let fs = fromLists [[1.0, 2.0],[3.0,4.0],[5.0,6.0]]
>>> let (d,l) = deconnect fs im delta
>>> d
[(1><2)
 [ 11.0, 17.0 ]]
>>> l
Just FullConnLayer:(2><4)
 [ 1.0, 1.0, 2.0, 3.0
 , 2.0, 2.0, 4.0, 6.0 ]

-}

deconnect :: FilterF -> Image -> Delta -> (Delta, Maybe Layer)
deconnect fs im delta = ([reshape (rows fs') $ fs' #> dl]
                        , Just (FullConnLayer $ outer dl im'))
  where
    dl  = flatten $ head delta    -- to Vector
    fs' = dropRows 1 fs
    im' = vjoin [konst 1.0 1, flatten $ head im]

-- reverse

reverseFullConnFilter :: FilterF -> Layer
reverseFullConnFilter fs = FullConnLayer $ tr' fs

-- update filter

updateFullConnFilter :: FilterF -> Double -> [Maybe Layer] -> Layer
updateFullConnFilter fs lr dl
  | dl' == [] = FullConnLayer fs
  | otherwise = FullConnLayer fs'
  where
    dl' = catMaybes dl
    delta = mscale (lr / fromIntegral (length dl')) (msum $ strip dl')
    fs' = msub fs delta

strip :: [Layer] -> [FilterF]
strip [] = []
strip (FullConnLayer fs:ds) = fs:strip ds
strip (_:ds) = strip ds


module STCR2S1 where

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

main = do
  args@(numPointStr:numOrientationStr:sigmaStr:taoStr:thresholdStr:numTrailStr:maxTrailStr:initDistStr:initOriStr:initSpeedStr:rStr:locationStr:deltaStr:cornerWeightStr:numThreadStr:_) <-
    getArgs
  print args
  let numPoint = read numPointStr :: Int
      numOrientation = read numOrientationStr :: Int
      sigma = read sigmaStr :: Double
      tao = read taoStr :: Double
      threshold = read thresholdStr :: Double
      numTrail = read numTrailStr :: Int
      maxTrail = read maxTrailStr :: Int
      initDist = read initDistStr :: [R2S1RPPoint]
      initOri = read initOriStr :: Double
      initSpeed = read initSpeedStr :: Double
      r = read rStr :: Double
      location = read locationStr :: (Int, Int)
      delta = read deltaStr :: Double
      cornerWeight = read cornerWeightStr :: Double
      numThread = read numThreadStr :: Int
      sourceDist =
        computeInitialDistribution numPoint numPoint numOrientation delta . L.take 1 $
        initDist
      sinkDist =
        computeInitialDistribution numPoint numPoint numOrientation delta . L.drop 1 $
        initDist
      dist =
        computeInitialDistribution numPoint numPoint numOrientation delta initDist
      folderPath = "output/test/STCR2S1"
  createDirectoryIfMissing True folderPath
  -- Compute the Green's function
  -- arrG <-
  --   solveMonteCarloR2S1
  --     numThread
  --     numTrail
  --     maxTrail
  --     numPoint
  --     numPoint
  --     numOrientation
  --     sigma
  --     tao
  --     r
  --     initSpeed
  --     ""
  let deltaTheta = 2 * pi / fromIntegral numOrientation
  -- arrG <-
  --   sampleR2S1
  --     numPoint
  --     numPoint
  --     delta
  --     deltaTheta
  --     initSpeed
  --     sigma
  --     tao
  --     [0 .. numOrientation - 1]
  arrG <-
    sampleR2S1Corner
      numPoint
      numPoint
      delta
      deltaTheta
      initSpeed
      sigma
      tao
      threshold
      [0 .. numOrientation - 1]
      cornerWeight
  printCurrentTime "Done."
  plan <- makeR2S1Plan emptyPlan arrG
  -- Source Field
  source <- shareWeightST plan sourceDist arrG
  plotImageRepa
    (folderPath </> "Source.png")
    (ImageRepa 8 .
     computeS .
     R.extend (Z :. (1 :: Int) :. All :. All) .
     reduceContrast 100 . R.sumS . rotate3D $
     source)
  plotThetaDimension folderPath "R2S1_" location .
    R.backpermute
      (extent source)
      (\(Z :. k :. i :. j) ->
         (Z :.
          (mod (k + numOrientation - (div numOrientation 2)) numOrientation :: Int) :.
          i :.
          j)) $
    source
  -- MP.mapM_
  --   (\i ->
  --      plotImageRepa
  --        (folderPath </> printf "Source_%d.png" i)
  --        (ImageRepa 8 .
  --         computeS . R.extend (Z :. (1 :: Int) :. All :. All) . R.slice source $
  --         (Z :. i :. All :. All)))
  --   [0 .. numOrientation - 1]
  -- Sink Field
  sink <- shareWeightST plan sinkDist arrG -- . rotateST arrG $ div numOrientation 2
  plotImageRepa
    (folderPath </> "Sink.png")
    (ImageRepa 8 .
     computeS . R.extend (Z :. (1 :: Int) :. All :. All) . reduceContrast 100 . R.sumS . rotate3D $
     sink)
  -- MP.mapM_
  --   (\i ->
  --      plotImageRepa
  --        (folderPath </> printf "Sink_%d.png" i)
  --        (ImageRepa 8 .
  --         computeS . R.extend (Z :. (1 :: Int) :. All :. All) . R.slice sink $
  --         (Z :. i :. All :. All)))
  --   [0 .. numOrientation - 1]
  -- Completion Field
  let completion = R.zipWith (*) source . timeReversal $ sink
  plotImageRepa
    (folderPath </> "Completion.png")
    (ImageRepa 8 .
     computeS . R.extend (Z :. (1 :: Int) :. All :. All) . R.sumS . rotate3D $
     completion)


-- |
-- This module defines the @Density@ type, which is the result of
-- doing a kernel density estimate on a sample of data. The idea
-- is that the distribution of metrics can be recorded using the @kde@
-- function, and then sampled in the model using @sampleDensity@.
{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE DeriveGeneric #-}
{-# LANGUAGE DeriveAnyClass #-}
module Hypersphere.Density where

import Control.DeepSeq
import Control.Monad.Bayes.Class
import Data.Aeson (ToJSON, FromJSON)
import Data.Foldable
import Data.List.NonEmpty (NonEmpty)
import qualified Data.List.NonEmpty as NonEmpty
import qualified Data.Vector.Unboxed as UV
import Data.Monoid (Any)
import qualified Diagrams.Core as D
import GHC.Generics
import Graphics.Rendering.Chart.Backend.Diagrams
import Graphics.Rendering.Chart.Easy
import           Linear.V2 (V2)
import qualified Linear.V2 as V2
import qualified Statistics.Sample.KernelDensity as S

-- | Represents a Kernel Density Estimate.
--
-- TODO: Should we store the cummulutive density to prevent recalculating
-- it all the time when sampling?
data Density = Density
    { dMesh :: !(UV.Vector Double)
    , dDensity :: !(UV.Vector Double)
    } deriving (Eq, Ord, Show, Generic, NFData)

instance ToJSON Density
instance FromJSON Density

-- | Find the area under the curve of the probability density. Should
-- be pretty close to 1.
--
-- Useful for debug.
integrate :: Density -> Double
integrate d@Density{..} =
    let
        spacing = UV.foldl1 subtract $ UV.take 2 dMesh
        zoids = getTrapezoids d
    in UV.sum $ UV.map (uncurry $ trapezoidalArea spacing) zoids

-- | The area of a trapezoid.
-- `trapezoid width height1 height2`
trapezoidalArea :: Double -> Double -> Double -> Double
trapezoidalArea w h1 h2 = (h1 + h2) * 0.5 * w

-- | Get the trapezoids that make up the area under the density curve.
getTrapezoids :: Density -> UV.Vector (Double,Double)
getTrapezoids Density{..} = UV.zip dDensity $ UV.tail dDensity

-- | Cumulative area under the density curve.
scanIntegral :: Density -> UV.Vector Double
scanIntegral d@Density{..} =
    let
        spacing = UV.foldl1 subtract $ UV.take 2 dMesh
        zoids = getTrapezoids d
    in
        UV.scanl (\area trap -> area + uncurry (trapezoidalArea spacing) trap) 0 zoids

-- | Create a density from a series of observations. There must be at least
-- one observation.
kde :: NonEmpty Double -> Density
kde = uncurry Density . S.kde 512 . UV.fromList . toList

-- | Plot the @Density@ to a file with the given name and title.
plotDensity :: FilePath -> String -> Density -> IO ()
plotDensity file title Density{..} = toFile def file $ do
    let
        maxDensity = 1.1 * UV.maximum dDensity
    layout_title .= title
    layout_y_axis . laxis_generate .= scaledAxis def (0,maxDensity)
    plot $ line "" [UV.toList $ UV.zip dMesh dDensity]

-- | Sample a value from the @Density@.
sampleDensity :: MonadSample m => Density -> m Double
sampleDensity d@Density{..} = do
    s <- random
    let
        mi = UV.findIndex (>=s) $ scanIntegral d
        o = maybe UV.last (\i -> (UV.! i)) mi dMesh
    return o


{-# LANGUAGE DataKinds #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE TypeOperators #-}

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

module Test.Convolution (
    convolutionTests,
    linearConvolutionTests,
    cyclicConvolutionTests
  ) where

import Data.Complex
import Test.HUnit ((@?=))
import Test.Hspec

import Spiral.Convolution
import Spiral.Exp
import Spiral.SPL

convolutionTests :: Spec
convolutionTests = describe "Convolution" $ do
    linearConvolutionTests
    cyclicConvolutionTests

linearConvolutionTests :: Spec
linearConvolutionTests = describe "Linear Convolution" $ do
  sequence_ [linearConvolutionTest (Standard i) i | i <- [1..16::Int]]
  sequence_ [linearConvolutionTest (ToomCook i) i | i <- [1..16::Int]]
  linearConvolutionTest (Tensor [2,2] [Standard 2, ToomCook 2]) 4
  linearConvolutionTest (Tensor [3,5] [Standard 3, ToomCook 5]) 15
  linearConvolutionTest (Tensor [2,2,2,2] [ToomCook 2, ToomCook 2, ToomCook 2, ToomCook 2]) 16
  sequence_ [linearConvolutionTest (Lift 5 i (Standard i)) 5 | i <- [6..10::Int]]
  sequence_ [linearConvolutionTest (FromCyclic i (ConvolutionTheorem (2*i - 2))) i | i <- [2..8::Int]]

linearConvolutionTest :: LinearConvolution (Exp (Complex Double)) -> Int -> Spec
linearConvolutionTest lin n =
      it (show lin) $ toMatrix (convolve lin vector) @?= expected_result
  where
      l :: Int
      l = 2 * n - 1

      vector :: [Exp (Complex Double)]
      vector = [fromIntegral i | i <- [1..n]]

      -- expected_result :: Matrix M (Exp (Complex Double))
      expected_result = toMatrix $ transpose $ fromLists
          [replicate i 0 ++ vector ++ replicate (l - n - i) 0 | i <- [0..n-1]]

cyclicConvolutionTests :: Spec
cyclicConvolutionTests = describe "Cyclic Convolution" $ do
  sequence_ [cyclicConvolutionTest (ConvolutionTheorem i) | i <- [3..8::Int]]
  cyclicConvolutionTest (Winograd 2 [Standard 1, Standard 1])
  cyclicConvolutionTest (Winograd 3 [Standard 1, Standard 2])
  cyclicConvolutionTest (Winograd 3 [Standard 1, Lift 2 4 (Standard 4)])
  cyclicConvolutionTest (Winograd 3 [Standard 1, Lift 2 3 (Standard 3)])
  cyclicConvolutionTest (AgarwalCooley 2 3 (ConvolutionTheorem 2) (Winograd 3 [Standard 1, Standard 2]))
  cyclicConvolutionTest (AgarwalCooley 4 3 (ConvolutionTheorem 4) (Winograd 3 [Standard 1, Standard 2]))
  cyclicConvolutionTest (SplitNesting [2, 3] [(Winograd 2 [Standard 1, Standard 1]), (Winograd 3 [Standard 1, Standard 2])])
  cyclicConvolutionTest (SplitNesting [3, 2] [(Winograd 3 [Standard 1, Standard 2]), (Winograd 2 [Standard 1, Standard 1])])

cyclicConvolutionTest :: CyclicConvolution (Exp (Complex Double)) -> Spec
cyclicConvolutionTest cyc = it (show cyc) $
    toMatrix (convolve cyc vector) @?= toMatrix (circ vector)
  where
    n :: Int
    n = getSize cyc

    vector :: [Exp (Complex Double)]
    vector = [fromIntegral i | i <- [1..n]]


{-# LANGUAGE ViewPatterns #-}
module AsteriskGaussian where

import           Control.Parallel.Strategies
import           Data.Array.Repa                        as R
import           Data.Complex                           as C
import           Data.List                              as L
import           Data.Vector.Storable                   as VS
import           Data.Vector.Unboxed                    as VU
import           Graphics.Rendering.Chart.Backend.Cairo
import           Graphics.Rendering.Chart.Easy
import           Math.Gamma
import           Pinwheel.FourierSeries2D
-- import           Pinwheel.Gaussian
import qualified FourierPinwheel.Filtering              as FP
import           System.Directory
import           System.Environment
import           System.FilePath
import           Text.Printf
import           Utils.BLAS
import           Utils.List
import Utils.Distribution

main = do
  args@(numR2FreqStr:numThetaFreqStr:numRFreqStr:numPointStr:numThetaStr:numRStr:xStr:yStr:periodStr:std1Str:std2Str:stdR2Str:_) <-
    getArgs
  let numR2Freq = read numR2FreqStr :: Int
      numThetaFreq = read numThetaFreqStr :: Int
      numRFreq = read numRFreqStr :: Int
      numPoint = read numPointStr :: Int
      numTheta = read numThetaStr :: Int
      numR = read numRStr :: Int
      x = read xStr :: Double
      y = read yStr :: Double
      period = read periodStr :: Double
      std1 = read std1Str :: Double
      std2 = read std2Str :: Double
      stdR2 = read stdR2Str :: Double
      folderPath = "output/test/AsteriskGaussian"
  let centerR2Freq = div numR2Freq 2
      centerThetaFreq = div numThetaFreq 2
      periodEnv = period ^ 2 / 4
      a = std1
      b = std2
      asteriskGaussianFreqTheta =
        VS.concat .
        parMap
          rdeepseq
          (\tFreq' ->
             let tFreq = tFreq' - centerThetaFreq
                 arr =
                   R.map (* (gaussian1DFreq (fromIntegral tFreq) b :+ 0)) $
                   analyticalFourierCoefficients1
                     numR2Freq
                     1
                     tFreq
                     0
                     a
                     period
                     periodEnv
              in VS.convert . toUnboxed . computeS $
                 centerHollowArray numR2Freq arr) $
        [0 .. numThetaFreq - 1]
      deltaTheta = 2 * pi / Prelude.fromIntegral numTheta
      harmonicsThetaD =
        fromFunction (Z :. numTheta :. numThetaFreq :. numR2Freq :. numR2Freq) $ \(Z :. theta' :. tFreq' :. xFreq' :. yFreq') ->
          let tFreq = fromIntegral $ tFreq' - centerThetaFreq
              xFreq = fromIntegral $ xFreq' - centerR2Freq
              yFreq = fromIntegral $ yFreq' - centerR2Freq
              theta = fromIntegral theta' * deltaTheta
           in (1 / (2 * pi) :+ 0) *
              cis (2 * pi / period * (x * xFreq + y * yFreq) + tFreq * theta)
  harmonicsTheta <- computeUnboxedP harmonicsThetaD
  xs <-
    VS.toList . VS.map magnitude <$>
    gemmBLAS
      numTheta
      1
      (numThetaFreq * numR2Freq ^ 2)
      (VS.convert . toUnboxed $ harmonicsTheta)
      asteriskGaussianFreqTheta
  toFile def (folderPath </> "theta.png") $ do
    layout_title .=
      printf "(%d , %d)" (Prelude.round x :: Int) (Prelude.round y :: Int)
    plot
      (line
         ""
         [ L.zip
             [fromIntegral i * deltaTheta / pi | i <- [0 .. numTheta - 1]]
             xs
         ])
  let centerRFreq = div numRFreq 2
      centerR2Freq = div numR2Freq 2
      gaussian2D =
        computeUnboxedS . fromFunction (Z :. numR2Freq :. numR2Freq) $ \(Z :. i' :. j') ->
          let i = i' - centerR2Freq
              j = j' - centerR2Freq
           in exp
                (pi * fromIntegral (i ^ 2 + j ^ 2) /
                 ((-1) * period ^ 2 * stdR2 ^ 2)) /
              (2 * pi * stdR2 ^ 2) :+
              0
      asteriskGaussianFreqR =
        VS.concat .
        parMap
          rdeepseq
          (\rFreq' ->
             let rFreq = rFreq' - centerRFreq
                 arr =
                   R.map
                     (* ((gaussian1DFourierCoefficients
                            (fromIntegral rFreq)
                            (log periodEnv)
                            b :+
                          0) -- *
                         -- cis
                         --   ((-2) * pi / log periodEnv * fromIntegral rFreq *
                         --    log 0.25)
                        )) . centerHollowArray numR2Freq $
                   analyticalFourierCoefficients1
                     numR2Freq
                     1
                     0
                     rFreq
                     a
                     period
                     periodEnv
              in VS.convert . toUnboxed . computeS . centerHollowArray numR2Freq $
                 arr -- *^ gaussian2D
           ) $
        [0 .. numRFreq - 1]
      deltaR = log periodEnv / Prelude.fromIntegral numR
      harmonicsRD =
        fromFunction (Z :. numR :. numRFreq :. numR2Freq :. numR2Freq) $ \(Z :. r' :. rFreq' :. xFreq' :. yFreq') ->
          let rFreq = fromIntegral $ rFreq' - centerRFreq
              xFreq = fromIntegral $ xFreq' - centerR2Freq
              yFreq = fromIntegral $ yFreq' - centerR2Freq
              r = fromIntegral r' * deltaR - log periodEnv / 2
           in (1 / log periodEnv :+ 0) *
              cis
                (2 * pi / period * (x * xFreq + y * yFreq) +
                 2 * pi / log periodEnv * rFreq * r)
  harmonicsR <- computeUnboxedP harmonicsRD
  ys <-
    VS.toList . VS.map magnitude <$>
    gemmBLAS
      numR
      1
      (numRFreq * numR2Freq ^ 2)
      (VS.convert . toUnboxed $ harmonicsR)
      asteriskGaussianFreqR
  toFile def (folderPath </> "R.png") $ do
    layout_title .= printf "(%.3f , %.3f)" x y
    plot
      (line
         ""
         [ L.zip
             [ fromIntegral i * deltaR - log periodEnv / 2
             | i <- [0 .. numR - 1]
             ]
             ys
         ])


--stolen from kmeans!

module Math.Probably.KMeans (kmeansOn, kmeans', kmeansLargestCentroids)
    where
import Data.List (transpose, sort, groupBy, minimumBy, sortBy)
import Data.Function (on)
import Data.Ord (comparing)
--import Numeric.LinearAlgebra
--import qualified Data.Vector.Storable as VS

type Vector a = [a]

dist :: (Floating b) => Vector b -> Vector b -> b
dist a b = sqrt . sum $ zipWith (\x y-> (x-y) ^ 2) a b

centroid :: (Fractional b) => (a -> Vector b) -> [a] -> Vector b
centroid f points = map (flip (/) l . sum) $ transpose $ map f points
    where l = fromIntegral $ length points

closest :: (Ord b, Floating b) => [Vector b] -> Vector b -> Vector b
closest points point = minimumBy (comparing $ dist point) points

recluster'
  :: (Ord b, Ord a, Floating b) =>
     (a -> Vector b) -> [Vector b] -> [a] -> [[a]]
recluster' f centroids points = map (map snd) $ groupBy ((==) `on` fst) reclustered
    where reclustered = sort [(closest centroids $ f a, a) | a <- points]

recluster :: (Floating a, Ord a, Ord b) => (b->Vector a) -> [[b]] -> [[b]]
recluster f clusters = recluster' f centroids $ concat clusters
    where centroids = map (centroid f) clusters

part :: (Eq a) => Int -> [a] -> [[a]]
part x ys
     | zs' == [] = [zs]
     | otherwise = zs : part x zs'
    where (zs, zs') = splitAt x ys
-- | Recluster points
kmeans' :: (Floating a, Ord a, Eq b, Ord b) => (b->Vector a) -> [[b]] -> [[b]]
kmeans' f clusters
    | clusters == clusters' = clusters
    | otherwise             = kmeans' f clusters'
    where clusters' = recluster f clusters
-- | Cluster points into k clusters.
-- |
-- | The initial clusters are chosen arbitrarily
kmeansOn :: (Floating a, Ord a, Eq b, Ord b) => (b->Vector a) -> Int -> [b] -> [[b]]
kmeansOn f k points = kmeans' f $ part l points
    where l = (length points + k - 1) `div` k

kmeansLargestCentroids :: Int -> [Vector Double] -> [Vector Double]
kmeansLargestCentroids k data0 =
  let clusters = reverse $ sortBy (comparing length) $ kmeansOn id k data0
  in map (centroid id) clusters


module HVX.Internal.Matrix
  ( Mat
  , allMat
  , anyMat
  , diagMat
  , ei
  , fpequalsMat
  , lpnorm
  , matrixPow
  , reduceMat
  , scalarMat
  , zeroMat
  , zeroVec
  ) where

import Numeric.LinearAlgebra hiding (i)
import Numeric.LinearAlgebra.Util

import HVX.Internal.Util

type Mat = Matrix Double

allMat :: (Double -> Bool) -> Mat -> Bool
allMat f x = all f (toList . flatten $ x)

anyMat :: (Double -> Bool) -> Mat -> Bool
anyMat f x = any f (toList . flatten $ x)

diagMat :: Mat -> Mat
diagMat = diag . flatten

ei :: Int -> Int -> Mat
ei n i = buildMatrix n 1 (\(j, _) -> if i == j then 1 else 0)

fpequalsMat :: Mat -> Mat -> Bool
fpequalsMat a b
  | ra == rb && ca == cb = all (uncurry fpequals) $ zip alist blist
  | otherwise = error "Two matrices with different dimensions cannot possibley be equal!"
  where
    ra = rows a
    rb = rows b
    ca = cols a
    cb = cols b
    alist = toList . flatten $ a
    blist = toList . flatten $ b

lpnorm :: Double -> Mat -> Double
lpnorm p x = sumElements y ** (1/p)
  where pMat = (1><1) [p]
        y = abs x ** pMat

matrixPow :: Double -> Mat -> Mat
matrixPow p = mapMatrix (** p)

reduceMat :: ([Double] -> Double) -> Mat -> Mat
reduceMat f = (1><1) . (:[]) . f . toList . flatten

scalarMat :: Double -> Mat
scalarMat x = (1><1) [x]

zeroMat :: Int -> Mat
zeroMat n = zeros n n

zeroVec :: Int -> Mat
zeroVec n = zeros n 1


{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE Strict #-}
module FokkerPlanck.GreensFunction where

import           Control.DeepSeq
import           Control.Monad                  as M
import           Data.Array.Accelerate.LLVM.PTX
import           Data.Array.IArray              as IA
import           Data.Array.Repa                as R
import           Data.Binary
import           Data.Complex
import           Data.List                      as L
import           Data.Vector.Storable           as VS
import           Data.Vector.Unboxed            as VU
import           FokkerPlanck.Analytic
import           FokkerPlanck.FourierSeriesGPU
import           FokkerPlanck.Histogram
import           Foreign.CUDA.Driver            as CUDA
import           FourierMethod.FourierSeries2D
import           Image.IO
import           Pinwheel.FourierSeries2D
import           Pinwheel.Transform
import           Sparse.Vector
import           STC.Utils
import           System.FilePath
import           Utils.Array
import           Utils.List
import           Utils.Parallel
import           Utils.SimpsonRule
import           Utils.Time
import Text.Printf
import           Utils.Distribution

sampleCartesian ::
     FilePath
  -> FilePath
  -> [PTX]
  -> Int
  -> Double
  -> Double
  -> Int
  -> Double
  -> Double
  -> Double
  -> Double
  -> Double
  -> Int
  -> Int
  -> Int
  -> Int -> Double
  -> IO (Histogram (Complex Double))
sampleCartesian filePath folderPath ptxs numPoints period delta oris gamma thetaSigma tau threshold r2Sigma maxPhiFreqs maxRhoFreqs maxThetaFreqs maxRFreqs stdR2 = do
  let center = div numPoints 2
      deltaTheta = 2 * pi / fromIntegral oris
      origin = R2S1RP 0 0 0 gamma
      idx =
        L.filter
          (\(R2S1RP c r _ g) -> (sqrt (c ^ 2 + r ^ 2) > 0) && g > 0)
          [ R2S1RP
            ((fromIntegral $ c - center) * delta)
            ((fromIntegral $ r - center) * delta)
            0
            gamma
          | c <- [0 .. numPoints - 1]
          , r <- [0 .. numPoints - 1]
          ]
      logGamma = 0 -- log gamma
      -- simpsonWeights = weightsSimpsonRule oris
      xs =
        L.concat -- .
        -- L.zipWith
        --   (\w -> L.map (\(a, b, c, d, v) -> (a, b, c, d, v * w)))
        --   simpsonWeights
         $
        parMap
          rdeepseq
          (\o ->
             let ori = fromIntegral o * deltaTheta
              in L.filter (\(_, _, _, _, v) -> v > threshold) .
                 L.map
                   (\point@(R2S1RP x y _ g) ->
                      let v =
                            computePji thetaSigma tau origin (R2S1RP x y ori g)
                          phi = atan2 y x
                          rho = sqrt $ (x ^ 2 + y ^ 2)
                       in ( phi
                          , log rho
                          , ori
                          , logGamma
                          , v * (1 - gaussian2DPolar rho stdR2)
                          )) $
                 idx)
          [0 .. oris - 1]
  -- arr <-
  --   computeUnboxedP . fromFunction (Z :. oris :. numPoints :. numPoints) $ \(Z :. o :. i :. j) ->
  --     let x = (fromIntegral $ i - center) * delta
  --         y = (fromIntegral $ j - center) * delta
  --         phi = atan2 y x
  --         rho = sqrt $ (x ^ 2 + y ^ 2)
  --         theta = fromIntegral o * deltaTheta
  --         v = computePji thetaSigma tau origin (R2S1RP x y theta gamma)
  --     in if rho <= 1.5
  --          then (1, 1, 1, 1, 0)
  --          else (phi, log rho, theta, logGamma, v)
  -- plotImageRepa (folderPath </> "GreensFunction.png") .
  --   ImageRepa 8 .
  --   computeS .
  --   -- reduceContrast 20 .
  --   extend (Z :. (1 :: Int) :. All :. All) .
  --   sumS . rotate3D . R.map (\(_, _, _, _, v) -> v) $
  --   arr
      -- simpsonWeights =
      --   computeWeightArrFromListOfShape [numPoints, numPoints, oris]
      -- xs =
      --   L.filter (\(_, _, _, _, v) -> v > threshold) .
      --   R.toList -- .
      --   -- R.zipWith (\w (a, b, c, d, v) -> (a, b, c, d, w * v)) simpsonWeights
      --   $
      --   arr
  let
  printCurrentTime $
    printf
      "Sparsity: %f%%\n"
      (100 * (fromIntegral . L.length $ xs) /
       (fromIntegral $ oris * numPoints ^ 2) :: Double)
  let phiFreqs = L.map fromIntegral [-maxPhiFreqs .. maxPhiFreqs]
      rhoFreqs = L.map fromIntegral [-maxRhoFreqs .. maxRhoFreqs]
      thetaFreqs = L.map fromIntegral [-maxThetaFreqs .. maxThetaFreqs]
      rFreqs = L.map fromIntegral [-maxRFreqs .. maxRFreqs]
      hist =
        L.foldl1' (addHistogram) .
        parZipWith
          rdeepseq
          (\ptx ys ->
             computeFourierCoefficientsGPU'
               r2Sigma
               period
               phiFreqs
               rhoFreqs
               thetaFreqs
               rFreqs
               ptx
               ys)
          ptxs .
        divideListN (L.length ptxs) $
        xs
  encodeFile filePath hist
  return hist

sampleCartesianCorner ::
     FilePath
  -> FilePath
  -> [PTX]
  -> Int
  -> Double
  -> Double
  -> Int
  -> Double
  -> Double
  -> Double
  -> Double
  -> Double
  -> Int
  -> Int
  -> Int
  -> Int
  -> Double
  -> Double
  -> IO (Histogram (Complex Double))
sampleCartesianCorner filePath folderPath ptxs numPoints period delta oris gamma thetaSigma tau threshold r2Sigma maxPhiFreqs maxRhoFreqs maxThetaFreqs maxRFreqs stdR2 weight = do
  let center = div numPoints 2
      deltaTheta = 2 * pi / fromIntegral oris
  --     origin = R2S1RP 0 0 0 gamma
      logGamma = 0
      idxFunc c r o =
        ( fromIntegral o * deltaTheta
        , fromIntegral (c - center) * delta
        , fromIntegral (r - center) * delta)
  --     arr1 =
  --       fromFunction (Z :. oris :. numPoints :. numPoints) $ \(Z :. o :. c :. r) ->
  --         if c == center && r == center
  --           then 0
  --           else let (ori, x, y) = idxFunc c r o
  --                 in computePji thetaSigma tau origin (R2S1RP x y ori gamma)
  --     orientations = [fromIntegral i * deltaTheta | i <- [0 .. oris - 1]]
  --     arr2 =
  --       fromFunction (Z :. oris :. numPoints :. numPoints) $ \(Z :. o :. c :. r) ->
  --         if c == center && r == center
  --           then 0
  --           else let (ori, x, y) = idxFunc c r o
  --                 in computePjiCorner'
  --                      thetaSigma
  --                      tau
  --                      threshold
  --                      orientations
  --                      (R2S1RP x y ori gamma)
  -- arr <- computeUnboxedP $ arr1 +^ (R.map (* weight) arr2)
  arr <- sampleR2S1Corner numPoints numPoints delta deltaTheta gamma thetaSigma tau threshold [0 .. oris - 1] weight
  let idx =
        [ idxFunc c r o
        | o <- [0 .. oris - 1]
        , c <- [0 .. numPoints - 1]
        , r <- [0 .. numPoints - 1]
        ]
      xs =
        L.map
          (\((ori, x, y), v) ->
             let phi = atan2 y x
                 rho = sqrt (x ^ 2 + y ^ 2)
              in ( phi
                 , log rho
                 , ori
                 , logGamma
                 , v * (1 - gaussian2DPolar rho stdR2))) .
        L.filter (\((_, x, y), v) -> v > threshold && (x /= 0 || y /= 0)) .
        L.zip idx . R.toList $
        arr
  printCurrentTime $
    printf
      "Sparsity: %f%%\n"
      (100 * (fromIntegral . L.length $ xs) /
       (fromIntegral $ oris * numPoints ^ 2) :: Double)
  let phiFreqs = L.map fromIntegral [-maxPhiFreqs .. maxPhiFreqs]
      rhoFreqs = L.map fromIntegral [-maxRhoFreqs .. maxRhoFreqs]
      thetaFreqs = L.map fromIntegral [-maxThetaFreqs .. maxThetaFreqs]
      rFreqs = L.map fromIntegral [-maxRFreqs .. maxRFreqs]
      hist =
        L.foldl1' (addHistogram) .
        parZipWith
          rdeepseq
          (\ptx ys ->
             computeFourierCoefficientsGPU'
               r2Sigma
               period
               phiFreqs
               rhoFreqs
               thetaFreqs
               rFreqs
               ptx
               ys)
          ptxs .
        divideListN (L.length ptxs) $
        xs
  encodeFile filePath hist
  return hist

-- sampleLogpolar ::
--      FilePath
--   -> [PTX]
--   -> Int
--   -> Int
--   -> Int
--   -> Double
--   -> Double
--   -> Double
--   -> Double
--   -> Double
--   -> Double
--   -> Int
--   -> Int
--   -> Int
--   -> Int
--   -> IO (Histogram (Complex Double))
-- sampleLogpolar filePath ptxs phis rhos oris maxScale gamma sigma tau threshold r2Sigma maxPhiFreqs maxRhoFreqs maxThetaFreqs maxRFreqs = do
--   let deltaPhi = 2 * pi / fromIntegral phis
--       deltaTheta = 2 * pi / fromIntegral oris
--       logMaxScale = log maxScale
--       deltaLogRho = logMaxScale / fromIntegral rhos
--       logGamma = 0 -- log gamma
--       origin = R2S1RP 0 0 0 gamma
--       -- xs =
--       --   parMap
--       --     rdeepseq
--       --     (\(phi, logRho, theta) ->
--       --        let x = (exp logRho) * cos phi
--       --            y = (exp logRho) * sin phi
--       --            v = computePji sigma tau origin (R2S1RP x y theta gamma)
--       --        in (phi, logRho, theta, 0, v))
--       --     [ ( deltaTheta * fromIntegral iPhi
--       --       , deltaLogRho * fromIntegral iLogRho - logMaxScale
--       --       , deltaTheta * fromIntegral iTheta)
--       --     | iPhi <- [0 .. oris - 1]
--       --     , iLogRho <- [0 .. rhos - 1]
--       --     , iTheta <- [0 .. oris - 1]
--       --     ]
--   arr <-
--     computeUnboxedP . R.fromFunction (Z :. oris :. rhos :. phis) $ \(Z :. iTheta :. iLogRho :. iPhi) ->
--       let theta = deltaTheta * fromIntegral iTheta
--           phi = deltaPhi * fromIntegral iPhi
--           logRho = deltaLogRho * fromIntegral iLogRho -- logMaxScale
--           x = (exp logRho) * cos phi
--           y = (exp logRho) * sin phi
--           v = computePji sigma tau origin (R2S1RP x y theta gamma)
--       in (phi, logRho, theta, logGamma, v)
--   let simpsonWeights = computeWeightArrFromListOfShape [phis, rhos, oris]
--       xs =
--         L.filter (\(_, _, _, _, v) -> v > threshold) .
--         R.toList .
--         R.zipWith (\w (a, b, c, d, v) -> (a, b, c, d, w * v)) simpsonWeights $
--         arr
--   printCurrentTime $
--     printf
--       "Sparsity: %f%%\n"
--       (100 * (fromIntegral . L.length $ xs) / (fromIntegral $ oris * phis * rhos) :: Double)
--   let phiFreqs = L.map fromIntegral [-maxPhiFreqs .. maxPhiFreqs]
--       rhoFreqs = L.map fromIntegral [-maxRhoFreqs .. maxRhoFreqs]
--       thetaFreqs = L.map fromIntegral [-maxThetaFreqs .. maxThetaFreqs]
--       rFreqs = L.map fromIntegral [-maxRFreqs .. maxRFreqs]
--       hist =
--         L.foldl1' (addHistogram) .
--         parZipWith
--           rdeepseq
--           (\ptx ys ->
--              computeFourierCoefficientsGPU'
--                r2Sigma
--                phiFreqs
--                rhoFreqs
--                thetaFreqs
--                rFreqs
--                ptx
--                ys)
--           ptxs .
--         divideListN (L.length ptxs) $
--         xs
--   encodeFile filePath hist
--   return hist

{-# INLINE sampleR2S1 #-}
sampleR2S1 ::
     Int
  -> Int
  -> Double
  -> Double
  -> Double
  -> Double
  -> Double
  -> [Int]
  -> IO (R.Array U DIM3 Double)
sampleR2S1 rows cols delta deltaTheta gamma sigma tau oris = do
  let centerRow = div rows 2
      centerCol = div cols 2
      origin = R2S1RP 0 0 0 gamma
  computeUnboxedP .
    R.traverse
      (fromListUnboxed (Z :. (L.length oris)) oris)
      (\(Z :. o) -> (Z :. o :. cols :. rows)) $ \f (Z :. o :. c :. r) ->
    if (c == centerCol) && (r == centerRow)
      then 0
      else computePji
             sigma
             tau
             origin
             (R2S1RP
                (delta * fromIntegral (c - centerCol))
                (delta * fromIntegral (r - centerRow))
                (fromIntegral (f (Z :. o)) * deltaTheta)
                gamma)

{-# INLINE sampleR2S1' #-}
sampleR2S1' ::
     Int
  -> Int
  -> Double
  -> Double
  -> Double
  -> Double
  -> Double
  -> R.Array U DIM2 (Int, Double)
sampleR2S1' rows cols delta gamma sigma tau ori =
  let centerRow = div rows 2
      centerCol = div cols 2
      origin = R2S1RP 0 0 0 gamma
  in computeUnboxedS . R.fromFunction (Z :. cols :. rows) $ \(Z :. c :. r) ->
       let x = c - centerCol
           y = r - centerRow
       in if x == 0 && y == 0
            then (0, 0)
            else ( getVectorIndex2D rows c r
                 , computePji
                     sigma
                     tau
                     origin
                     (R2S1RP
                        (delta * fromIntegral (c - centerCol))
                        (delta * fromIntegral (r - centerRow))
                        ori
                        gamma))

{-# INLINE sampleR2S1Corner #-}
sampleR2S1Corner ::
     Int
  -> Int
  -> Double
  -> Double
  -> Double
  -> Double
  -> Double
  -> Double
  -> [Int]
  -> Double
  -> IO (R.Array U DIM3 Double)
sampleR2S1Corner rows cols delta deltaTheta gamma sigma tau threshold oris weight = do
  let centerRow = div rows 2
      centerCol = div cols 2
      origin = R2S1RP 0 0 0 gamma
      orientations = L.map (\o -> deltaTheta * fromIntegral o) oris
  arr1 <-
    computeUnboxedP .
    R.traverse
      (fromListUnboxed (Z :. (L.length oris)) oris)
      (\(Z :. o) -> (Z :. o :. cols :. rows)) $ \f (Z :. o :. c :. r) ->
      if (c == centerCol) && (r == centerRow)
        then 0
        else computePji
               sigma
               tau
               origin
               (R2S1RP
                  (delta * fromIntegral (c - centerCol))
                  (delta * fromIntegral (r - centerRow))
                  (fromIntegral (f (Z :. o)) * deltaTheta)
                  gamma)
  arr2 <-
    computeUnboxedP .
    R.traverse
      (fromListUnboxed (Z :. (L.length oris)) oris)
      (\(Z :. o) -> (Z :. o :. cols :. rows)) $ \f (Z :. o :. c :. r) ->
      if (c == centerCol) && (r == centerRow)
        then 0
        else computePjiCorner'
               delta
               sigma
               tau
               threshold
               orientations
               (R2S1RP
                  (delta * fromIntegral (c - centerCol))
                  (delta * fromIntegral (r - centerRow))
                  (fromIntegral (f (Z :. o)) * deltaTheta)
                  gamma)
  computeUnboxedP $ (arr1 +^ (R.map (* weight) arr2))

cornerDistribution ::
     Double
  -> Double
  -> Double
  -> Double
  -> Double
  -> [Double]
  -> Double
  -> (Double, Double)
  -> [Double]
cornerDistribution delta gamma sigma tau threshold oris weight (x, y) =
  let origin = R2S1RP 0 0 0 gamma
   in L.map
        (\ori ->
           computePji sigma tau origin (R2S1RP x y ori gamma) +
           weight *
           computePjiCorner'
             delta
             sigma
             tau
             threshold
             oris
             (R2S1RP x y ori gamma))
        oris

{-# INLINE cornerDistribution' #-}
cornerDistribution' ::
     Double
  -> Double
  -> Double
  -> Double
  -> Double
  -> [Double]
  -> Double
  -> (Double, Double)
  -> Double
  -> IO (Double, Double, Double)
cornerDistribution' delta gamma sigma tau threshold oris weight (x, y) ori = do
  let origin = R2S1RP 0 0 0 gamma
      a = computePji sigma tau origin (R2S1RP x y ori gamma)
      b = computePjiCorner' delta sigma tau threshold oris (R2S1RP x y ori gamma)
  printf "sigma = %f\ntau = %f\nthreshold = %f\noris = %s\n(R2S1RP %f %f %f %f)\n" sigma tau threshold (show oris) x y ori gamma
  return (a, b, a + weight * b)

computeFourierCoefficients ::
     FilePath
  -> [Int]
  -> [PTX]
  -> Int
  -> Int
  -> Double
  -> Double
  -> Double
  -> Double
  -> Double
  -> Int
  -> Double
  -> Double
  -> Int
  -> Int
  -> Int
  -> Int
  -> Int
  -> Int
  -> IO (Histogram (Complex Double))
computeFourierCoefficients histFilePath deviceIDs ptxs numPoints oris delta deltaFreq gamma sigma tau numR2Freqs period std numBatchR2Freqs numBatchOri maxPhiFreq maxRhoFreq maxThetaFreq maxRFreq = do
  when
    (even oris)
    (error "computeFourierCoefficients: the number of orientations is not odd.")
  let deltaTheta = 2 * pi / fromIntegral oris
      simpsonWeights =
        computeUnboxedS $
        (computeWeightArrFromListOfShape [numPoints, numPoints] :: R.Array D DIM2 (Complex Double))
      oriIdxs = divideListN numBatchOri [0 .. oris - 1]
  xs <-
    M.mapM
      (sampleR2S1 numPoints numPoints delta deltaTheta gamma sigma tau)
      oriIdxs
  let weightedVecs =
        parMap
          rdeepseq
          (\arr ->
             let (Z :. n :. m :. c) = extent arr
             in CuMat (n * m) c .
                CuVecHost .
                VU.convert .
                toUnboxed . computeS . R.traverse2 arr simpsonWeights const $ \f fW idx@(Z :. i :. j :. _) ->
                  (f idx :+ 0) * fW (Z :. i :. j))
          xs
  coefficients <-
    computeFourierCoefficientsR2Stream
      deviceIDs
      ptxs
      numR2Freqs
      numPoints
      period
      delta
      deltaFreq
      numBatchR2Freqs
      weightedVecs
  print "coefficients done"
  let centerR2Freq = div numR2Freqs 2
      centerPoints = div numPoints 2
      logGamma = log gamma
      gaussianWeightedCoef =
        computeUnboxedS $ applyGaussian deltaFreq std coefficients
      -- gaussianWeightedCoef = coefficients
      simpsonWeightsOri =
        computeUnboxedS $
        (computeWeightArrFromListOfShape [oris] :: R.Array D DIM1 (Complex Double))
      weightedCoef =
        R.traverse2 gaussianWeightedCoef simpsonWeightsOri const $ \fG fS idx@(Z :. o :. _ :. _) ->
          fG idx * fS (Z :. o)
      folderPath = "output/test/STCPinwheel"
  -- print . extent $ gaussianWeightedCoef
  -- inverseR2Harmonics <-
  --   createInverseHarmonicMatriesGPU ptxs 1 numPoints numR2Freqs period delta
  -- sourceR2 <-
  --   computeFourierSeriesR2
  --     deviceIDs
  --     numR2Freqs
  --     numPoints
  --     period
  --     inverseR2Harmonics
  --     [ CuMat oris (numR2Freqs ^ 2) . CuVecHost . VU.convert . toUnboxed $
  --       gaussianWeightedCoef
  --     ]
  plotImageRepa (folderPath </> "Green.png") .
    ImageRepa 8 .
    fromUnboxed (Z :. (1 :: Int) :. numPoints :. numPoints) .
    VU.map sqrt . toUnboxed . sumS . R.map (\x -> (x) ** 2) . L.head $
    xs
  plotImageRepa (folderPath </> "Coefficients.png") .
    ImageRepa 8 .
    fromUnboxed (Z :. (1 :: Int) :. numR2Freqs :. numR2Freqs) .
    VU.map sqrt . toUnboxed . sumS . R.map (\x -> (magnitude x) ** 2) . rotate3D $
    gaussianWeightedCoef
  -- plotImageRepa (folderPath </> "GreenRecon.png") .
  --   ImageRepa 8 .
  --   fromUnboxed (Z :. (1 :: Int) :. numPoints :. numPoints) .
  --   VU.map sqrt . toUnboxed . sumS . R.map (\x -> (magnitude x) ** 2) . rotate3D $
  --   sourceR2
  ys <-
    (computeUnboxedP . R.traverse weightedCoef id $ \f idx@(Z :. o :. i :. j) ->
       if i == centerR2Freq && j == centerR2Freq
         then (0, 0, 0, 0, 0)
         else let x = deltaFreq * (fromIntegral $ i - centerR2Freq)
                  y = deltaFreq * (fromIntegral $ j - centerR2Freq)
                  phi = atan2 y x
                  logRho = log . sqrt $ (x ^ 2 + y ^ 2)
              in (phi, logRho, fromIntegral o * deltaTheta, logGamma, f idx))
  -- ys <-
  --   (computeUnboxedP .
  --    R.traverse
  --      (R.traverse2 (L.head xs) simpsonWeightsOri const $ \fG fS idx@(Z :. _ :. _ :. o) ->
  --         (fG idx :+ 0) * fS (Z :. o))
  --      id $ \f idx@(Z :. i :. j :. o) ->
  --      if i == centerPoints && j == centerPoints
  --        then (0, 0, 0, 0, 0)
  --        else let x = delta * (fromIntegral $ i - centerPoints)
  --                 y = delta * (fromIntegral $ j - centerPoints)
  --                 phi = atan2 x y
  --                 logRho = log . sqrt $ (x ^ 2 + y ^ 2)
  --             in ( phi
  --                , logRho
  --                , fromIntegral o * deltaTheta
  --                , logGamma
  --                , f idx))
  let hist =
        L.foldl1' (addHistogram) .
        parZipWith
          rdeepseq
          (\ptx ->
             computePinwheelCoefficients
               maxPhiFreq
               maxRhoFreq
               maxThetaFreq
               maxRFreq
               ptx)
          ptxs .
        divideListN (L.length ptxs) . R.toList $
        ys
  encodeFile histFilePath hist
  return hist

computePinwheelTransformCoefficients ::
     Int
  -> FilePath
  -> [Int]
  -> [PTX]
  -> Int
  -> Int
  -> Double
  -> Double
  -> Double
  -> Double
  -> Double
  -> Int
  -> Double
  -> Double
  -> Int
  -> Int
  -> Int
  -> Int
  -> Int
  -> Int
  -> Int
  -> Double
  -> IO (Histogram (Complex Double))
computePinwheelTransformCoefficients numThread histFilePath deviceIDs ptxs numPoints oris delta threshold gamma sigma tau numR2Freqs periodR2 std numBatchR2 numBatchPinwheelFreqs maxPhiFreq maxRhoFreq maxThetaFreq maxRFreq batchSize s = do
  printCurrentTime "compute pinwheelArr"
  pinwheelArr <-
    pinwheelFourierSeries
      deviceIDs
      ptxs
      numPoints -- numR2Freqs
      numPoints
      delta
      periodR2
      maxPhiFreq
      maxRhoFreq
      maxThetaFreq
      maxRFreq
      s
      numBatchR2
      numBatchPinwheelFreqs :: IO (IA.Array (Int, Int) (VU.Vector (Complex Double)))
  printCurrentTime "pinwheelArr done"
  let folderPath = "output/test/STCPinwheel"
  let deltaTheta = 2 * pi / (fromIntegral oris)
      greensFunc =
        parMap
          rdeepseq
          (\o ->
             let theta = fromIntegral o * deltaTheta
                 vec =
                   toUnboxed $
                   sampleR2S1' numPoints numPoints delta gamma sigma tau theta
             in PinwheelTransformData (log gamma) theta .
                uncurry SparseVector .
                VU.unzip .
                VU.map (\(i, v) -> (i, v :+ 0)) .
                VU.filter (\(_, v) -> v > threshold) $
                vec)
          [0 .. oris - 1]
      len =
        L.foldl'
          (\s (PinwheelTransformData _ _ (SparseVector vec _)) ->
             s + VU.length vec)
          0
          greensFunc
  printCurrentTime $
    printf
      "Sparsity: %f%%\n"
      (100 * (fromIntegral len) / (fromIntegral $ oris * numPoints ^ 2) :: Double)
  let hist =
        pinwheelTransform
          numThread
          maxPhiFreq
          maxRhoFreq
          maxThetaFreq
          maxRFreq
          s
          pinwheelArr
          greensFunc
  encodeFile histFilePath hist
  return hist


{-|
Module      : Qubism.StateVec
Description : Types and functions for quantum state vectors
Copyright   : (c) Keith Pearson, 2018
License     : MIT
Maintainer  : keith@qubitrot.org
-}

{-# LANGUAGE DataKinds           #-}
{-# LANGUAGE KindSignatures      #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeOperators       #-}

module Qubism.StateVec
  ( StateVec (UnsafeMkStateVec)
  , mkStateVec
  , mkStateVec'
  , mkQubit
  , normalize
  , tensor
  , collapse
  , measureQubit
  , measure
  , dimension
  ) where

-- For dependent typing
import GHC.TypeLits
import Data.Singletons
import Data.Singletons.TypeLits
import Data.Finite

import           Control.Monad.Random
import           Control.Monad.Trans.State.Strict
import           Data.Complex
import           Numeric.LinearAlgebra ((|>))
import qualified Numeric.LinearAlgebra as LA
import           Text.Printf

import Qubism.Algebra
import Qubism.CReg

newtype StateVec (n :: Nat) =
  UnsafeMkStateVec (LA.Vector C)

-- | To provide "close enough" equality testing.
instance Eq (StateVec n) where
  UnsafeMkStateVec zs == UnsafeMkStateVec ws =
     LA.norm_2 (zs - ws) < 0.000001

instance KnownNat n => VectorSpace (StateVec n) where
  zero = UnsafeMkStateVec $ internalLen (sing :: Sing n) |> repeat 0
  z .: (UnsafeMkStateVec a) = UnsafeMkStateVec $ LA.scalar z * a
  (UnsafeMkStateVec a) +: (UnsafeMkStateVec b) = UnsafeMkStateVec $ a + b
  neg (UnsafeMkStateVec a) = UnsafeMkStateVec $ -a

instance KnownNat n => HilbertSpace (StateVec n) where
  (UnsafeMkStateVec zs) <.> (UnsafeMkStateVec ws) = zs LA.<.> ws

instance KnownNat n => Show (StateVec n) where
  show (UnsafeMkStateVec zs) = foldl show' "" $ zip [0..] (LA.toList zs)
    where show' str (i,z) = str ++ pf z ++ "  " ++ basis i ++ "\n"
          pf z    =  printf "% 6.4f" (realPart z) ++ "  + "
                  ++ printf "% 6.4f" (imagPart z) ++ "i"
          basis i = let bit j = if i `quot` 2^(n-j-1) `mod` 2 == 0
                                then '0' else '1'
                    in  "|" ++ fmap bit (take n [0..]) ++ ">"
          n = fromIntegral $ fromSing (sing :: Sing n)

internalLen :: (KnownNat n, Num a) => Sing n -> a
internalLen = (2 ^) . fromIntegral . fromSing

-- | Number of qubits in the StateVec
dimension :: forall (n :: Nat) a. (KnownNat n, Num a) => StateVec n -> a
dimension _ = fromIntegral $ fromSing (sing :: SNat n)

-- | StateVec's are intialized to |0>
mkStateVec :: forall n . KnownNat n => StateVec n
mkStateVec = UnsafeMkStateVec $ internalLen (sing :: Sing n) |> (1 : repeat 0)

-- | Make a statevector with an explicity specified size.
-- StateVec's are intialized to |0>
mkStateVec' :: Sing n -> StateVec n
mkStateVec' sn = UnsafeMkStateVec $ l |> (1 : repeat 0)
  where l = (2^) . fromIntegral $ fromSing sn

-- | A qubit is just a StateVec 1, initalized to |0>
mkQubit :: StateVec 1
mkQubit = UnsafeMkStateVec $ LA.fromList [1, 0]

normalize :: StateVec n -> StateVec n
normalize (UnsafeMkStateVec zs) = UnsafeMkStateVec $ LA.normalize zs

adjoint :: KnownNat n => StateVec n -> (StateVec n -> C)
adjoint qr = (qr <.>)

-- | The tensor product on elements of our Hilbert space
tensor :: StateVec n -> StateVec m -> StateVec (n + m)
tensor (UnsafeMkStateVec zs) (UnsafeMkStateVec ws) =
  UnsafeMkStateVec $ LA.flatten (zs `LA.outer` ws)

-- | Collapse a StateVec to a state compatable with qubit <n> being measured <bit>.
-- Only occurs physically with measurement, so internal use only.
collapse
  :: forall n . KnownNat n
  => Finite n -> Bit -> StateVec n -> StateVec n
collapse i b (UnsafeMkStateVec zs) = normalize . UnsafeMkStateVec $ zs * mask
  where
    mask   = l |> altseq
    altseq = replicate m ifZero ++ replicate m ifOne ++ altseq
    ifZero = if b == Zero then 1 else 0
    ifOne  = if b == One  then 1 else 0
    l      = internalLen (sing :: Sing n)
    m      = l `quot` (2 ^ (getFinite i + 1))

-- | Preforms a measurement of an induvidual qubit (indexed Finite n)
-- in a quantum register.
measureQubit
  :: (MonadRandom m, KnownNat n)
  => Finite n -> StateT (StateVec n) m Bit
measureQubit i = do
  qr <- get
  r  <- getRandomR (0, 1)
  let qrZero = collapse i Zero qr
      qrOne  = collapse i One qr
      pOne   = realPart $ qrOne <.> qr -- guaranteed to be real, so this
  if r < pOne                          -- is just a type-cast.
    then put qrOne  >> pure One
    else put qrZero >> pure Zero

-- | Preform a measurement on all qubits in a quantum register, returning a
-- computational basis state.
measure
  :: forall m n . (MonadRandom m, KnownNat n)
  => StateT (StateVec n) m CReg
measure = mkCReg <$> traverse measureQubit (take n [0 ..])
  where n = fromIntegral $ fromSing (sing :: Sing n)


{-# LANGUAGE OverloadedStrings #-}

module Main where

import Control.Monad (forever)
import Control.Monad.IO.Class (MonadIO (liftIO))
import Control.Monad.State.Strict
  ( MonadState (get, put),
    execStateT,
    forever,
    lift,
    runStateT,
  )
import Data.Complex (Complex (..), magnitude)
import Data.List (foldl')
import qualified Data.Text as T
import System.Random (Random (randomR), newStdGen)
import UI.NCurses
  ( Color (..),
    CursorMode (..),
    defaultWindow,
    drawText,
    moveCursor,
    render,
    runCurses,
    setColor,
    setCursorMode,
    setEcho,
    updateWindow,
  )
import UI.NCurses.HalfHeight
  ( Buffer,
    colorId,
    drawBuffer,
    initHexColors,
    mkBuffer,
    setXY,
  )

-- Set up a 15-color greyscale pallette
greyScale :: [String]
greyScale =
  [ "000000",
    "111111",
    "222222",
    "333333",
    "444444",
    "666666",
    "777777",
    "888888",
    "999999",
    "aaaaaa",
    "bbbbbb",
    "cccccc",
    "dddddd",
    "eeeeee",
    "ffffff"
  ]

-- Calculate the Mandelbrot steps-to-divergence and plot as Color 1 to Color 15
mandelbrot :: (Fractional a, RealFloat a) => a -> a -> Color
mandelbrot x y = Color (16 - ((+) 1 $ steps (0 :+ 0) (x' :+ y') 0))
  where
    x' = (x - 100) * 0.025
    y' = (y - 50) * 0.025
    steps z c i
      | i > 13 || magnitude z' > 2 = i
      | otherwise = steps z' c (i + 1)
      where
        z' = z * z + c

-- Draw the Mandelbrot set to the buffer point by point
drawMandelbrot :: Buffer -> Buffer
drawMandelbrot buffer =
  foldl'
    (\b (x, y) -> setXY x y (mandelbrot (fromIntegral x) (fromIntegral y)) b)
    buffer
    [(x, y) | x <- [0 .. 149], y <- [0 .. 99]]

mandelbrotMain :: IO ()
mandelbrotMain = do
  -- Register the 15 greyscale hex colors as Color 1 through Color 15.
  -- Every combination of these colors (fg and bg) is registered with Curses
  -- with a unique ID, which we return in the colorMap.
  colorMap <- runCurses $ initHexColors nord

  -- Creates a new buffer and draws the Mandelbrot set in each cell using the
  -- greyscale colors 1 through 15
  let buffer = drawMandelbrot $ mkBuffer 150 100 (Color 1)

  runCurses $ do
    -- Hide the cursor
    setEcho False
    setCursorMode CursorInvisible

    -- Enter an infinite drawing loop displaying the buffer.
    forever $ do
      let drawOp = do
            -- First draw the buffer to the screen.
            drawBuffer colorMap 0 0 buffer
            -- Now draw some text in the centre.
            -- We set the color to curses ID corresponding to fg 15 (white) and bg 1 (black)
            let (Just c) = colorId colorMap (Color 15) (Color 1)
            setColor c
            moveCursor 25 50
            drawText "| Text demonstrating that each cell is half the row height |"
      w <- defaultWindow
      updateWindow w drawOp
      render

nord :: [String]
nord =
  [ "3b4252",
    "bf616a",
    "a3be8c",
    --"ebcb8b",
    "81a1c1",
    "b48ead",
    "88c0d0",
    "e5e9f0",
    "4c566a",
    "bf616a",
    "a3be8c",
    "ebcb8b",
    "81a1c1",
    "b48ead",
    "8fbcbb",
    "eceff4"
  ]

performanceMain :: IO ()
performanceMain = do
  colorMap <- runCurses $ initHexColors greyScale
  -- Precompute the buffers
  let buffer = mkBuffer 100 100 (Color 1)
      col x y t =
        let x' = fromIntegral (x - 50)
            y' = fromIntegral (y - 50)
            t' = fromIntegral t
         in abs $ round $ (14 * ((sin $ 0.1 * x') * (cos $ 0.1 * y')) / (sin $ 0.1 * t')) + 1
      nextBuf b t = foldl' (\b (x, y) -> setXY x y (Color $ col x y t) b) b [(x, y) | x <- [0 .. 99], y <- [0 .. 99]]
      bufs = cycle $ scanl nextBuf buffer [0 .. 63]
  runCurses $ do
    _ <- (flip runStateT) (bufs, 0) $ do
      lift $ setEcho False
      lift $ setCursorMode CursorInvisible
      forever $ do
        (bufs, t) <- get
        let b = head bufs
            drawOp = do
              drawBuffer colorMap 0 0 b
              let (Just c) = colorId colorMap (Color 15) (Color 1)
              setColor c
              moveCursor 0 0
              drawText $ T.pack $ show t
        w <- lift defaultWindow
        lift $ updateWindow w drawOp
        lift render
        put (tail bufs, t + 1)
    return ()

main :: IO ()
main = do
  --mandelbrotMain
  performanceMain


module Problem215 where
{-
  Task description:
  Consider the problem of building a wall out of 2×1 and 3×1 bricks (horizontal×vertical dimensions) such that,
  for extra strength, the gaps between horizontally-adjacent bricks never line up in consecutive layers, i.e. never form a "running crack".

  For example, the following 9×3 wall is not acceptable due to the running crack shown in red:

  There are eight ways of forming a crack-free 9×3 wall, written W(9,3) = 8.

  Calculate W(32,10).
-}

import Control.Monad
import Data.Function (on)
import Data.Graph (Graph, Vertex)
import Data.Maybe (fromJust)
import Numeric.LinearAlgebra (Matrix, (<>))
import qualified GHC.Arr as Arr
import qualified Data.Graph as Graph
import qualified Data.HashSet as HashSet
import qualified Data.List as List
import qualified Data.Map as Map
import qualified Data.Tree as Tree
import qualified Numeric.LinearAlgebra as LA

{-
A row may only ever depend on the preceding row.
We may have cycles of rows that are allowed to follow each other.
An interesting question will be wether collapsing such cycles will be easier than simply 'exploring them'.
In both ways memoization may come in handy here.
-}

type Brick = Int
type Line = [Brick]
type Width = Int
type Height = Int
type Wall = [Line]

bricks :: [Brick]
bricks = [2, 3]

buildLines :: Width -> [Line]
buildLines 0 = [[]]
buildLines w = do
  brick <- bricks
  guard $ brick <= w
  lineRest <- buildLines $ w - brick
  return $ brick:lineRest

matchingLines :: Line -> Line -> Bool
matchingLines l1 l2 = not $ (matchingCracks `on` crackPositions) l1 l2
  where
    crackPositions :: Line -> [Int]
    crackPositions = fmap sum . init . tail . List.inits

    matchingCracks :: [Int] -> [Int] -> Bool
    matchingCracks [] _ = False
    matchingCracks _ [] = False
    matchingCracks xss@(x:xs) yss@(y:ys)
      | x == y = True
      | x < y = matchingCracks xs yss
      | x > y = matchingCracks xss ys

{-
With the matchingLines predicate I can generate a graph out of the possible lines.
We're searching for all routes of a given length that traverse this graph
and that start in any of it's nodes.
-}

mkLineEdges :: [Line] -> [(Line, Line)]
mkLineEdges [] = []
mkLineEdges (l:ls) = findMatchingLines l ls `mappend` mkLineEdges ls
  where
    findMatchingLines :: Line -> [Line] -> [(Line, Line)]
    findMatchingLines l lines = concat $ do
      l' <- lines
      guard $ matchingLines l l'
      return [(l,l'), (l',l)]

mkLineToVertex :: [Line] -> (Line -> Vertex)
mkLineToVertex lines = let lvMap = Map.fromList $ zip lines [0..]
                       in fromJust . (`Map.lookup` lvMap)

graphForWidth :: Width -> Graph
graphForWidth w = let lines = buildLines w
                      lEdges = mkLineEdges lines
                      bounds = (0, length lines - 1)
                      toVertex = mkLineToVertex lines
                      edges = [(toVertex a, toVertex b)|(a,b) <- lEdges]
                  in Graph.buildG bounds edges

graphToDenseMatrix :: Graph -> Matrix Double
graphToDenseMatrix = LA.toDense . graphToAssocMatrix
  where
    graphToAssocMatrix :: Graph -> LA.AssocMatrix
    graphToAssocMatrix g = do
      (v, targets) <- Arr.assocs g
      [((v, t), 1)| t <- targets]

dissect :: Graph -> [Graph]
dissect g = do
  tree <- Graph.components g
  let vertices = Tree.flatten tree
      (renumber, bounds) = mkRenumber vertices
  return . Graph.buildG bounds $ do
    v <- vertices
    w <-  Arr.unsafeAt g v
    return (renumber v, renumber w)
  where
    mkRenumber ls = let mapping = Map.fromList $ zip ls [0..]
                        λ = fromJust . (`Map.lookup` mapping)
                    in (λ, (0, Map.size mapping))

mPow :: Matrix Double -> Int -> Matrix Double
mPow m 0 = LA.ident $ LA.rows m
mPow m 1 = m
mPow m n
  | even n = mPow (m <> m) (n `div` 2)
  | otherwise = m <> mPow m (n - 1)

countWalls :: Width -> Height -> Int
countWalls w h
  | h < 1 = 0
  | h == 1 = length $ buildLines w
-- Using dissect reduced computation time from ~140s to ~9s.
  | otherwise = let g = graphForWidth w
                    ms = graphToDenseMatrix <$> dissect g
                    ms' = fmap (mPow `flip` (h - 1)) ms
                    calcSum = sum . concat . LA.toLists
                in round . sum $ fmap calcSum ms'

example = countWalls 9 3

-- 806844323190414
solution = countWalls 32 10

main = print solution


-- comment
module Main where

import qualified Graphics.Rendering.OpenGL as GL
import Numeric.LinearAlgebra.HMatrix as H hiding (reshape)
import Prelude hiding (init)
import GLCode
import Points

import Sketch

rrand :: Int -> IO (Matrix Double)
rrand n = randn n n

crand :: Int -> IO (Matrix (Complex Double))
crand n = curry toComplex <$> randn n n <*> randn n n

unit :: Floating a => Int -> [a]
unit n = [fromIntegral j/fromIntegral n | j <- [0..n-1]]

initialPoints :: Int -> Vector (Complex Double)
initialPoints num = fromList (n ++ s) where
    n = [(-12.5+25*j) :+ (-20) | j <- unit num]
    s = [(-12.5+25*j) :+ 20 | j <- unit num]

finalPoints :: Int -> Vector (Complex Double)
finalPoints num = fromList (e ++ w) where
    e = [(-20) :+ (-12.5+25*j) | j <- unit num]
    w = [20 :+ (-12.5+25*j) | j <- unit num]

main :: IO ()
main = do
    let num = 20

    u' <- crand (num * 2)
    let u = scale 1.5 u'

    let d = initialPoints num
    let d' = finalPoints num

    let p = diag d
    let p' = u `mul` diag d' `mul` inv u

    mainGifLoop "xxx" 12.0 512 512 1 0 100 $ \time -> do
        let t = time/200

        let eigs = H.toList $ eigenvalues $ scale (1-realToFrac t) p + scale (realToFrac t) p'

        GL.clearColor GL.$= GL.Color4 0.0 0.0 0.0 1
        io $ GL.clear [GL.ColorBuffer]

        let points = [GL.Vertex2 (0.04*realToFrac x) (0.04*realToFrac y) |
                        x :+ y <- eigs] :: [GL.Vertex2 Float]
        setUniform "shader" "pointSize" (8.0 :: Float)
        drawPoint "shader" "vPosition" points

        io GL.flush


{-# LANGUAGE BangPatterns    #-}
{-# LANGUAGE NamedFieldPuns  #-}
{-# LANGUAGE RecordWildCards #-}
module Network
    ( NeuralNet (..),
      train,
      initWB,
      loadData,
      loadTestData,
    ) where

import           Codec.Compression.GZip (decompress)
import           Control.Monad
import qualified Data.Array.IO          as A
import qualified Data.ByteString        as BS
import qualified Data.ByteString.Lazy   as BL
import           Data.Int               (Int64)
import           Data.List
import           GHC.Base               (build)
import           Numeric.LinearAlgebra  hiding (build, normalize)
import           Prelude                hiding ((<>))
import           System.Directory       (getCurrentDirectory)
import qualified System.Random          as R
import           Text.Printf
import           Util

type Image        = Vector Double
type Label        = Vector Double
type Bias         = Vector Double
type Weight       = Matrix Double
type Activation   = Matrix Double
type Nabla        = Matrix Double
type ImagesMat    = Matrix Double
type LabelsMat    = Matrix Double

data NeuralNet = NN
    { weights   :: ![Weight]   -- ^ List of all weights
    , biases    :: ![Bias]    -- ^ List of all biases
    , eta       :: !Double    -- ^ Learning rate
    , epochs    :: !Int       -- ^ No. of epochs
    , layers    :: !Int       -- ^ No. of layers
    , layerSize :: !Int       -- ^ Size of hidden layers
    , batchSize :: !Int       -- ^ Size of mini-batch
    , trainData :: ![(Image, Label)]
    , testData  :: ![(ImagesMat, Int)]
    } deriving (Show, Eq)

test :: NeuralNet -> IO ()
test NN{..} = do
    let (imgs, labels) = unzip testData
    let guesses = [guess img (zip weights $ toMatrix <$> biases) | img <- imgs]
    let correct = sum $ (\(a,b) -> if a == b then 1 else 0) <$> zip guesses labels
    printf "Epoch %d: %.2f%%\n" (30-epochs :: Int)  (correct/100 :: Float)
  where
    guess :: ImagesMat -> [(Weight, Matrix Double)] -> Int
    guess i wb = maxIndex . flatten . head $ feedforward i wb

train :: NeuralNet -> IO NeuralNet
train net@NN{..} = do
    shuffledData <- shuffle trainData
    let miniBatches = getMiniBatch <$> chunkList batchSize shuffledData
    let newNet = foldl trainBatch net miniBatches
    test newNet
    case epochs of
        0 -> return newNet
        _ -> train newNet { epochs=epochs-1 }

trainBatch :: NeuralNet -> (ImagesMat, LabelsMat) -> NeuralNet
trainBatch net@NN{weights=ws, biases=bs, ..} (mX, mY) = newNet
    where newNet = let activations = feedforward mX (zip ws $ bTm bs batchSize)
                      in updateMiniBatch net $
                          backprop mX mY net activations

-- | Make biases suitable for full-matrix backprop.
bTm :: [Bias] -> Int -> [Matrix Double]
bTm biases batchSize = fmap (fromColumns . replicate batchSize) biases

-- | Turn mini-batch matrices into one matrix.
getMiniBatch :: [(Image, Label)] -> (ImagesMat, LabelsMat)
getMiniBatch mBatchData = (fromColumns imgs, fromColumns labels)
  where
    (imgs, labels) = unzip mBatchData

feedforward :: ImagesMat -> [(Weight, Matrix Double)] -> [Activation]
feedforward = scanr (\(w, b) a -> sigmoid $ w <> a + b)

--  | Stochastic gradient descent
sgd = id

--  | Updates mini-batch
--  NOTE: use 'seq' or 'deepseq' for full evaluation,
--  instead of just reducing to WHNF
updateMiniBatch :: NeuralNet -> ([Vector Double], [Matrix Double])  -> NeuralNet
updateMiniBatch net@NN{..} (nablaB, nablaW) = net {weights=wnew, biases=bnew}
    where
        wnew :: [Weight]
        !wnew = [ w - scale (eta/fromIntegral batchSize) nw |
            (w, nw)  <- zip weights nablaW ]
        bnew :: [Bias]
        !bnew = [ b - scale (eta/fromIntegral batchSize) nb |
            (b, nb)  <- zip biases nablaB ]

--  | Returns nablaW and nablaB
backprop :: ImagesMat -> LabelsMat -> NeuralNet -> [Activation] -> ([Vector Double], [Matrix Double])
backprop x y NN{..} (aL:as) = (nablaB, nablaW)
    where
        nablaW :: [Matrix Double]
        nablaW = zipWithSafe (<>) delta (tr <$> as) -- THIS IS WHERE THE MEMORY LEAK IS IF I DO JUST zipWith
        nablaB :: [Vector Double]
        nablaB = sumRows <$> delta
        delta :: [Matrix Double]
        delta = init $ deltas [crossEntropy y aL] y as weights

deltas :: [Matrix Double] -> LabelsMat -> [Activation] -> [Weight] -> [Matrix Double]
deltas deltal _ [] [] = deltal
deltas deltal y (a:as) (wlp1:ws) = deltas (deltal++[(tr wlp1 <> last deltal) * sigmoid' a]) y as ws

quadraticLoss :: LabelsMat -> Activation -> Matrix Double
quadraticLoss y aL = (aL - y) * sigmoid' aL

crossEntropy :: LabelsMat -> Activation -> Matrix Double
crossEntropy y aL = (aL - y)

{-- Arithmetic Functions --}
cost :: Matrix Double -> Matrix Double -> Matrix Double
cost a y = ((a-y)**2) / 2

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

sigmoid' :: Floating a => a -> a
sigmoid' x = x * (1 - x)

normalize :: (Integral a, Floating b) => a -> b
normalize x = (fromIntegral x) / 255

{-- Data Processing Functions --}
initWB :: NeuralNet -> IO NeuralNet
initWB net@NN{batchSize = bs, layerSize = ls, layers = n} = do
    headW <- randn ls 784
    midW  <- replicateM (n-2) (randn ls ls)
    lastW <- randn 10 ls

    initB <- replicateM (n-1) (randn ls 1)
    lastB <- randn 10 1

    let weights = reverse $ headW : midW ++ [lastW]
    let biases = reverse . fmap flatten $ initB ++ [lastB]

    return net { weights = weights
               , biases  = biases }

loadData :: IO [(Vector Double, Vector Double)]
loadData = do
    trainImgs <- getData "train-images-idx3-ubyte.gz"
    trainLabels <- getData "train-labels-idx1-ubyte.gz"
    let !labels = BL.toStrict trainLabels
    let !imgs = BL.toStrict trainImgs
    let !l = [vectorizeLabel $ getLabel n labels | n <- [0..49999]] :: [Label]
    let !i = [getImage n imgs | n <- [0..49999]] :: [Image]
    return $ zip i l

loadTestData :: IO [(Matrix Double, Int)]
loadTestData = do
    testImgs <- getData "t10k-images-idx3-ubyte.gz"
    testLabels <- getData "t10k-labels-idx1-ubyte.gz"
    let !labels = BL.toStrict testLabels
    let !imgs = BL.toStrict testImgs
    let !l = [getLabel n labels | n <- [0..9999]] :: [Int]
    let !i = asColumn <$> [getImage n imgs | n <- [0..9999]] :: [Matrix Double]
    return $ zip i l

getData :: FilePath -> IO BL.ByteString
getData path = do
    currentDir <- getCurrentDirectory
    fileData <- decompress <$> BL.readFile (currentDir ++ "/data/mnist_dataset/" ++ path)
    return fileData

getImage :: Int -> BS.ByteString -> Image
getImage n imgs = fromList [normalize $ BS.index imgs (16 + n*784 + s) | s <- [0..783]]

-- | `getImage` but lazy
getImage' :: Int64 -> BL.ByteString -> Image
getImage' n imgs = fromList [normalize $ BL.index imgs (16 + n*784 + s) | s <- [0..783]]

getLabel :: Num a => Int -> BS.ByteString -> a
getLabel n labels = fromIntegral $ BS.index labels (n+8)

getGuess :: [Activation] -> Int
getGuess = maxIndex . flatten . head

vectorizeLabel :: Int -> Vector Double
vectorizeLabel l = fromList $ x ++ 1 : y
    where (x,y) = splitAt l $ replicate 9 0

{-- HELPER FUNCTIONS --}
toMatrix :: Vector Double -> Matrix Double
toMatrix = reshape 1

-- | Means of rows in matrix as a column vector
meanRows :: Matrix Double -> Vector Double
meanRows = fst . meanCov . tr

sumRows :: (Num a, Element a) => Matrix a -> Vector a
sumRows m = fromList $ sum <$> toLists m

-- | zipWith, but throws error if lengths don't match
zipWithSafe :: (a -> b -> c) -> [a] -> [b] -> [c]
zipWithSafe f as bs
    | length as == length bs = zipWith f as bs
    | otherwise = error $ "zipWith: lengths don't match (" ++ (show $ length as) ++ " " ++ (show $ length bs) ++ ")"

-- | Randomly shuffle a list
--   /O(N)/
shuffle :: [a] -> IO [a]
shuffle xs = do
        ar <- newArray n xs
        forM [1..n] $ \i -> do
            j <- R.randomRIO (i,n)
            vi <- A.readArray ar i
            vj <- A.readArray ar j
            A.writeArray ar j vi
            return vj
  where
    n = length xs
    newArray :: Int -> [a] -> IO (A.IOArray Int a)
    newArray n xs =  A.newListArray (1,n) xs


{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE FlexibleContexts #-}

module Neural.Matrix(
    initNet, updateNet, feedLayer, forwardPass, propagate,
    backPropagate, relu, softmax, categoricalCrossEntropy
) where

import Control.Monad (zipWithM)
import Data.List (foldl', scanl', transpose)
import Control.Arrow ((&&&))
import Numeric.AD (grad)
import Numeric.AD.Mode (auto)
import Numeric.AD.Internal.Reverse (Reverse, Tape)
import Data.Reflection (Reifies)
import Numeric.LinearAlgebra

type Biases = Vector R
type Weights = Matrix R
type Activation = [R] -> [R]
type Layer = (Activation, Biases, Weights)

parameters :: Layer -> (Biases, Weights)
parameters (_, bs, ws) = (bs, ws)

weights :: Layer -> Weights
weights (_, _, ws) = ws

initNet :: R -> [Int] -> [Activation] -> IO [Layer]
initNet b szs@(_:lns) acts = layer <$> zipWithM randn szs lns
  where layer = zip3 acts $ vector . flip replicate b <$> lns

relu :: Activation
relu = map $ max 0

relu' :: Activation
relu' = map f
  where f x | x < 0     = 0
            | otherwise = 1

categoricalCrossEntropy :: RealFloat a => [a] -> [a] -> a
categoricalCrossEntropy ys = negate . sum . zipWith (*) ys . map log . softmax

softmax :: RealFloat a => [a] -> [a]
softmax ts = let exps = exp <$> ts in map (/ sum exps) exps

feedLayer :: (Vector R, Vector R) -> Layer -> (Vector R, Vector R)
feedLayer (_, an) (aF, bs, ws) = (id &&& fromList . aF . toList) $ bs + an <# ws

forwardPass :: Vector R -> [Layer] -> [(Vector R, Vector R)]
forwardPass = scanl' feedLayer . (id &&& id)

propagate :: (Vector R, Matrix R) -> Vector R -> Vector R
propagate (zn, wn) dldn = wn #> dldn * dadz
  where dadz = fromList . relu' . toList $ zn

backPropagate :: Vector R -> [Vector R] -> [Matrix R] -> [Vector R]
backPropagate dldL tls lws = tail $ scanr propagate dldL $ zip tls lws

updateLayer :: R -> (Vector R, Vector R, Layer) -> Layer
updateLayer s (a, e, (aF, bs, ws)) = (aF, newBs, newWs)
  where newWs = ws - scalar s * outer a e
        newBs = bs - scalar s * e

updateNet :: (forall a. Reifies a Tape => [Reverse a R]
                                       -> [Reverse a R]
                                       -> Reverse a R) ->
             R -> Vector R -> Vector R -> [Layer] -> [Layer]
updateNet c s xs ys ls = map (updateLayer s) $ zip3 (init ans) dls ls
  where fps  = forwardPass xs ls
        ans  = snd <$> fps
        yhat = auto <$> (toList . last) ans
        dldL = fromList $ grad (c (auto <$> (toList ys))) yhat
        dls  = backPropagate dldL (fst <$> init fps) (weights <$> ls)


module Lib
( Ising
, randModel
, runBatch
, runMC
, totalEnergy
, totalMagnetization
, flipRandSpin
, flipSpin
) where

import Model

import Data.Maybe
import Control.Monad.State
import System.Random (Random, random, randomR)
import Numeric.LinearAlgebra (accum, atIndex, toLists)

runBatch :: Show a => Int -> [Ising a] -> Ising ([Maybe [a]], IsingState)
runBatch steps prop = do
    props <- replicateM steps $ runMC prop
    state <- get
    return (props, state)

runMC :: [Ising a] -> Ising (Maybe [a])
runMC properties = do
    incrementStep
    state_i <- get
    dE <- flipRandSpin >>= getSpinEnergy >>= return . negate
    state_f <- get

    if dE <= 0
        then put $ state_f { nAccept = nAccept state_f + 1 }
    else do
        r <- randR (0,1)

        if r <= exp (-beta * dE)
            then put $ state_f { nAccept = nAccept state_f + 1 }
        else do
            put state_i            -- reset the state to before the spin flip
            r <- rand :: Ising Int -- increment RNG
            return ()

    state <- get
    if step state `mod` propFreq state == 0
        then liftM Just $ sequence properties
        else return Nothing

totalEnergy :: Ising Double
totalEnergy = do
    state <- get
    let n = dim state
    let indices = [(row, col) | row <- [0..n-1], col <- [0..n-1]]
    spinE <- liftM sum $ mapM getSpinEnergy indices
    let fieldE = sum $ map (* h state) (spins state)
    normalize $ spinE - fieldE

totalMagnetization :: Ising Double
totalMagnetization = do
        state <- get
        normalize $ sum (spins state)

getSpinEnergy :: (Int, Int) -> Ising Double
getSpinEnergy (row, col) = do
    state <- get
    neighbors <- getNeighbors (row, col)
    let spin = model state `atIndex` (row,col)
    return $ 2 * j state * spin * (sum neighbors)

getNeighbors :: (Int, Int) -> Ising [Spin]
getNeighbors (row, col) = do
    state <- get
    let n = dim state
    let m = model state
    let ns = map (atIndex m) [ ((row - 1) `mod` n, col)
                             , ((row + 1) `mod` n, col)
                             , (row, (col - 1) `mod` n)
                             , (row, (col + 1) `mod` n)
                             ]
    return ns
flipRandSpin :: Ising (Int, Int)
flipRandSpin = do
    state <- get
    let n = dim state
    row <- randR (0,n-1)
    col <- randR (0,n-1)
    flipSpin (row, col)
    return (row, col)

flipSpin :: (Int, Int) -> Ising ()
flipSpin (row, col) = do
    state <- get
    let spin = (model state) `atIndex` (row, col)
    let model' = accum (model state) (\a b -> a) [((row,col), -spin)]
    put $ state { model = model' }

-----
--
-- Helper functions
--
-----

beta = 1.0

normalize :: Fractional a => a -> Ising a
normalize x = do
    state <- get
    return $ x / (fromIntegral $ dim state)^2

spins :: IsingState -> [Spin]
spins state = concat . toLists $ model state

randR :: (Random a) => (a,a) -> Ising a
randR (a,b) = do
    state <- get
    let (r, rng') = randomR (a,b) $ rng state
    put $ state { rng = rng' }
    return r

rand :: (Random a) => Ising a
rand = do
    state <- get
    let (r, rng') = random $ rng state
    put $ state { rng = rng' }
    return r

incrementStep :: Ising Int
incrementStep = do
    state <- get
    let step' = step state + 1
    put $ state { step = step' }
    return step'


module Main where

import Control.Monad.Par
import Data.Complex
import qualified Data.Vector as V

import Graphics.Gloss.Data.Display
import Graphics.Gloss.Data.Color
import Graphics.Gloss.Data.Picture
import Graphics.Gloss.Interface.Pure.Display

windowSpec :: Display
windowSpec = InWindow "hello" (800,800) (0,0)

blackColor :: Color
blackColor = makeColor 0.0 0.0 0.0 1.0

whiteColor :: Color
whiteColor = makeColor 1.0 1.0 1.0 1.0

pseudoPixel :: Picture
pseudoPixel = rectangleSolid 1.0 1.0

blockf :: ((Float, Float), Int) -> Picture
blockf ((x,y),t) = let b = ((fromIntegral t) / 25) :: Float
                       c = makeColor b b b 1.0
                   in color c (translate x y pseudoPixel)

fractalRadius :: Float
fractalRadius = 2.0

fractalCalc :: (Complex Float, Complex Float) -> ((Float,Float),Int)
fractalCalc n = let a1 = V.iterateN 50 (\(start,c) -> (start*start + c,c)) (fmap (*0.005) n) :: V.Vector (Complex Float, Complex Float)
                    b =  V.length . V.filter (> fractalRadius) . V.map (realPart . abs . fst)  $ a1 :: Int
                    c = snd n
                    d = (realPart c, imagPart c) :: (Float, Float)
                in (d, b)

complexBlock :: V.Vector (Complex Float, Complex Float)
complexBlock = let x = [-400.0..400.0]
                   xv = V.fromList x
                   c = 0 :+ 0
               in V.concatMap (\x -> V.map (\z -> (c, x :+ z)) xv) xv

fractalP :: V.Vector (Par ((Float, Float), Int))
fractalP = V.map (\x -> (spawnP . fractalCalc $ x) >>= get) complexBlock

fractal :: Picture
fractal =  mconcat . V.toList . V.map blockf . runPar . sequence $ fractalP

main :: IO ()
main = display windowSpec whiteColor fractal


-- |
-- Module    : Statistics.Test.MannWhitneyU
-- Copyright : (c) 2010 Neil Brown
-- License   : BSD3
--
-- Maintainer  : bos@serpentine.com
-- Stability   : experimental
-- Portability : portable
--
-- Mann-Whitney U test (also know as Mann-Whitney-Wilcoxon and
-- Wilcoxon rank sum test) is a non-parametric test for assessing
-- whether two samples of independent observations have different
-- mean.
module Statistics.Test.MannWhitneyU (
    -- * Mann-Whitney U test
    mannWhitneyUtest
  , mannWhitneyU
  , mannWhitneyUCriticalValue
  , mannWhitneyUSignificant
    -- ** Wilcoxon rank sum test
  , wilcoxonRankSums
  , module Statistics.Test.Types
    -- * References
    -- $references
  ) where

import Control.Applicative ((<$>))
import Data.List (findIndex)
import Data.Ord (comparing)
import Numeric.SpecFunctions (choose)
import Prelude hiding (sum)
import Statistics.Distribution (quantile)
import Statistics.Distribution.Normal (standard)
import Statistics.Function (sortBy)
import Statistics.Sample.Internal (sum)
import Statistics.Test.Internal (rank, splitByTags)
import Statistics.Test.Types (TestResult(..), PositionTest(..), significant)
import Statistics.Types (PValue,pValue)
import qualified Data.Vector.Unboxed as U

-- | The Wilcoxon Rank Sums Test.
--
-- This test calculates the sum of ranks for the given two samples.
-- The samples are ordered, and assigned ranks (ties are given their
-- average rank), then these ranks are summed for each sample.
--
-- The return value is (W₁, W₂) where W₁ is the sum of ranks of the first sample
-- and W₂ is the sum of ranks of the second sample.  This test is trivially transformed
-- into the Mann-Whitney U test.  You will probably want to use 'mannWhitneyU'
-- and the related functions for testing significance, but this function is exposed
-- for completeness.
wilcoxonRankSums :: (Ord a, U.Unbox a) => U.Vector a -> U.Vector a -> (Double, Double)
wilcoxonRankSums xs1 xs2 = (sum ranks1, sum ranks2)
  where
    -- Ranks for each sample
    (ranks1,ranks2) = splitByTags $ U.zip tags (rank (==) joinSample)
    -- Sorted and tagged sample
    (tags,joinSample) = U.unzip
                      $ sortBy (comparing snd)
                      $ tagSample True xs1 U.++ tagSample False xs2
    -- Add tag to a sample
    tagSample t = U.map (\x -> (t,x))

-- | The Mann-Whitney U Test.
--
-- This is sometimes known as the Mann-Whitney-Wilcoxon U test, and
-- confusingly many sources state that the Mann-Whitney U test is the same as
-- the Wilcoxon's rank sum test (which is provided as 'wilcoxonRankSums').
-- The Mann-Whitney U is a simple transform of Wilcoxon's rank sum test.
--
-- Again confusingly, different sources state reversed definitions for U₁
-- and U₂, so it is worth being explicit about what this function returns.
-- Given two samples, the first, xs₁, of size n₁ and the second, xs₂,
-- of size n₂, this function returns (U₁, U₂)
-- where U₁ = W₁ - (n₁(n₁+1))\/2
-- and U₂ = W₂ - (n₂(n₂+1))\/2,
-- where (W₁, W₂) is the return value of @wilcoxonRankSums xs1 xs2@.
--
-- Some sources instead state that U₁ and U₂ should be the other way round, often
-- expressing this using U₁' = n₁n₂ - U₁ (since U₁ + U₂ = n₁n₂).
--
-- All of which you probably don't care about if you just feed this into 'mannWhitneyUSignificant'.
mannWhitneyU :: (Ord a, U.Unbox a) => U.Vector a -> U.Vector a -> (Double, Double)
mannWhitneyU xs1 xs2
  = (fst summedRanks - (n1*(n1 + 1))/2
    ,snd summedRanks - (n2*(n2 + 1))/2)
  where
    n1 = fromIntegral $ U.length xs1
    n2 = fromIntegral $ U.length xs2

    summedRanks = wilcoxonRankSums xs1 xs2

-- | Calculates the critical value of Mann-Whitney U for the given sample
-- sizes and significance level.
--
-- This function returns the exact calculated value of U for all sample sizes;
-- it does not use the normal approximation at all.  Above sample size 20 it is
-- generally recommended to use the normal approximation instead, but this function
-- will calculate the higher critical values if you need them.
--
-- The algorithm to generate these values is a faster, memoised version of the
-- simple unoptimised generating function given in section 2 of \"The Mann Whitney
-- Wilcoxon Distribution Using Linked Lists\"
mannWhitneyUCriticalValue
  :: (Int, Int)     -- ^ The sample size
  -> PValue Double  -- ^ The p-value (e.g. 0.05) for which you want the critical value.
  -> Maybe Int      -- ^ The critical value (of U).
mannWhitneyUCriticalValue (m, n) p
  | m < 1 || n < 1 = Nothing    -- Sample must be nonempty
  | p' <= 1        = Nothing    -- p-value is too small. Null hypothesis couldn't be disproved
  | otherwise      = findIndex (>= p')
                   $ take (m*n)
                   $ tail
                   $ alookup !! (m+n-2) !! (min m n - 1)
  where
    mnCn = (m+n) `choose` n
    p'   = mnCn * pValue p

{-
-- Original function, without memoisation, from Cheung and Klotz:
-- Double is needed to avoid integer overflows.
a :: Int -> Int -> Int -> Double
a u bigN m
  | u < 0            = 0
  | u >= m * n       = bigN `choose` m
  | m == 1 || n == 1 = fromIntegral (u + 1)
  | otherwise        = a  u      (bigN - 1)  m
                     + a (u - n) (bigN - 1) (m-1)
  where
    n = bigN - m
-}

-- Memoised version of the original a function, above.
--
-- Doubles are stored to avoid integer overflow. 32-bit Ints begin to
-- overflow for bigN as small as 33 (64-bit one at 66) while Double to
-- go to infinity till bigN=1029
--
--
-- outer list is indexed by big N - 2
-- inner list by (m-1) (we know m < bigN)
-- innermost list by u
--
-- So: (alookup !! (bigN - 2) !! (m - 1) ! u) == a u bigN m
alookup :: [[[Double]]]
alookup = gen 2 [1 : repeat 2]
  where
    gen bigN predBigNList
       = let bigNlist = [ [ amemoed u m
                          | u <- [0 .. m*(bigN-m)]
                          ] ++ repeat (bigN `choose` m)
                        | m <- [1 .. (bigN-1)]] -- has bigN-1 elements
         in bigNlist : gen (bigN+1) bigNlist
      where
        amemoed :: Int -> Int -> Double
        amemoed u m
          | m == 1 || n == 1 = fromIntegral (u + 1)
          | otherwise        = mList !! u
                             + if u < n then 0 else predmList !! (u-n)
          where
            n = bigN - m
            (predmList : mList : _) = drop (m-2) predBigNList
            -- Lists for m-1 and m respectively. i-th list correspond to m=i+1
            --
            -- We know that predBigNList has bigN - 2 elements
            -- (and we know that n > 1 therefore bigN > m + 1)
            -- So bigN - 2 >= m, i.e. predBigNList must have at least m elements
            -- elements, so dropping (m-2) must leave at least 2

-- | Calculates whether the Mann Whitney U test is significant.
--
-- If both sample sizes are less than or equal to 20, the exact U critical value
-- (as calculated by 'mannWhitneyUCriticalValue') is used.  If either sample is
-- larger than 20, the normal approximation is used instead.
--
-- If you use a one-tailed test, the test indicates whether the first sample is
-- significantly larger than the second.  If you want the opposite, simply reverse
-- the order in both the sample size and the (U₁, U₂) pairs.
mannWhitneyUSignificant
  :: PositionTest     -- ^ Perform one-tailed test (see description above).
  -> (Int, Int)       -- ^ The samples' size from which the (U₁,U₂) values were derived.
  -> PValue Double    -- ^ The p-value at which to test (e.g. 0.05)
  -> (Double, Double) -- ^ The (U₁, U₂) values from 'mannWhitneyU'.
  -> Maybe TestResult -- ^ Return 'Nothing' if the sample was too
                      --   small to make a decision.
mannWhitneyUSignificant test (in1, in2) pVal (u1, u2)
  -- Use normal approximation
  | in1 > 20 || in2 > 20 =
    let mean  = n1 * n2 / 2     -- (u1+u2) / 2
        sigma = sqrt $ n1*n2*(n1 + n2 + 1) / 12
        z     = (mean - u1) / sigma
    in Just $ case test of
                AGreater      -> significant $ z     < quantile standard p
                BGreater      -> significant $ (-z)  < quantile standard p
                SamplesDiffer -> significant $ abs z > abs (quantile standard (p/2))
  -- Use exact critical value
  | otherwise = do crit <- fromIntegral <$> mannWhitneyUCriticalValue (in1, in2) pVal
                   return $ case test of
                              AGreater      -> significant $ u2        <= crit
                              BGreater      -> significant $ u1        <= crit
                              SamplesDiffer -> significant $ min u1 u2 <= crit
  where
    n1 = fromIntegral in1
    n2 = fromIntegral in2
    p  = pValue pVal

-- | Perform Mann-Whitney U Test for two samples and required
-- significance. For additional information check documentation of
-- 'mannWhitneyU' and 'mannWhitneyUSignificant'. This is just a helper
-- function.
--
-- One-tailed test checks whether first sample is significantly larger
-- than second. Two-tailed whether they are significantly different.
mannWhitneyUtest
  :: (Ord a, U.Unbox a)
  => PositionTest     -- ^ Perform one-tailed test (see description above).
  -> PValue Double    -- ^ The p-value at which to test (e.g. 0.05)
  -> U.Vector a       -- ^ First sample
  -> U.Vector a       -- ^ Second sample
  -> Maybe TestResult -- ^ Return 'Nothing' if the sample was too small to
                      --   make a decision.
mannWhitneyUtest ontTail p smp1 smp2 =
  mannWhitneyUSignificant ontTail (n1,n2) p $ mannWhitneyU smp1 smp2
    where
      n1 = U.length smp1
      n2 = U.length smp2

-- $references
--
-- * Cheung, Y.K.; Klotz, J.H. (1997) The Mann Whitney Wilcoxon
--   distribution using linked lists. /Statistica Sinica/
--   7:805–813.
-- <http://www3.stat.sinica.edu.tw/statistica/oldpdf/A7n316.pdf>.


{-# LANGUAGE DataKinds #-}
{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE TypeFamilies #-}

module Song.Grenade (
  SongSD
, song2SD
, song2Rn
, song2TupleRn
) where

import Song (Song(Song, _genre, _loudness, _tempo, _timeSignature, _key, _mode, _duration, _avgTimbre, _varTimbre))
import Grenade (S(S1D), Shape(D1))
import Numeric.LinearAlgebra.Static (R, konst, vector, (#), (&))
import Control.Arrow ((***), first)

type SongSD = (S ('D1 30), S ('D1 1))

song2SD :: (String -> Double) -> Song -> SongSD
song2SD tag = (S1D *** S1D) . song2Rn tag

song2Rn :: (String -> Double) -> Song -> (R 30, R 1)
song2Rn tag song = first (uncurry (#)) $ song2TupleRn tag song

song2TupleRn :: (String -> Double) -> Song -> ((R 3, R 27), R 1)
song2TupleRn tag Song{..} = ((discreteFeats, floatFeats), konst $ tag _genre)
  where
    discreteFeats = konst (fromIntegral _timeSignature)
                  & fromIntegral _key
                  & fromIntegral _mode

    floatFeats = (konst _loudness & _tempo & _duration)
               # (vector _avgTimbre::R 12)
               # (vector _varTimbre::R 12)


module STCR2Z1T0Edge where

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

main = do
  args@(numPointStr:numOrientationStr:sigmaStr:taoStr:lenStr:initialScaleStr:numTrailStr:maxTrailStr:theta0FreqsStr:thetaFreqsStr:histFilePath:pinwheelFlagStr:numIterationStr:writeSourceFlagStr:edgeFilePath:numThreadStr:_) <-
    getArgs
  print args
  let numPoint = read numPointStr :: Int
      numOrientation = read numOrientationStr :: Int
      sigma = read sigmaStr :: Double
      tao = read taoStr :: Double
      len = read lenStr :: Int
      initialScale = read initialScaleStr :: Double
      numTrail = read numTrailStr :: Int
      maxTrail = read maxTrailStr :: Int
      theta0Freq = read theta0FreqsStr :: Double
      theta0Freqs = [-theta0Freq .. theta0Freq]
      thetaFreq = read thetaFreqsStr :: Double
      thetaFreqs = [-thetaFreq .. thetaFreq]
      pinwheelFlag = read pinwheelFlagStr :: Bool
      numIteration = read numIterationStr :: Int
      writeSourceFlag = read writeSourceFlagStr :: Bool
      numThread = read numThreadStr :: Int
      folderPath = "output/test/STCR2Z1T0Edge"
  createDirectoryIfMissing True folderPath
  flag <- doesFileExist histFilePath
  radialArr <-
    if flag
      then R.map magnitude . getNormalizedHistogramArr <$>
           decodeFile histFilePath
      else do
        putStrLn "Couldn't find a Green's function data. Start simulation..."
        solveMonteCarloR2Z1T0Radial
          numThread
          numTrail
          maxTrail
          numPoint
          numPoint
          sigma
          tao
          1
          theta0Freqs
          thetaFreqs
          histFilePath
          (emptyHistogram
             [ (round . sqrt . fromIntegral $ 2 * (div numPoint 2) ^ 2)
             , L.length theta0Freqs
             , L.length thetaFreqs
             ]
             0)
  arrR2Z1T0 <-
    computeUnboxedP $
    computeR2Z1T0ArrayRadial
      radialArr
      numPoint
      numPoint
      1.5
      thetaFreqs
      theta0Freqs
  plan <- makeR2Z1T0Plan emptyPlan arrR2Z1T0
  xs <- parseEdgeFile edgeFilePath
  randomPonintSet <- generateRandomPointSet 40 numPoint numPoint
  let ys =
        L.map
          (\(R2S1RPPoint (a, b, c, d)) -> (R2S1RPPoint (a - 150, b - 150, c, d)))
          xs
      zs =
        L.map
          (\(R2S1RPPoint (a, b, c, d)) ->
             (R2S1RPPoint (a - center numPoint, b - center numPoint, c, d)))
          randomPonintSet
      points = ys L.++ zs
  let bias = computeBiasR2T0 numPoint numPoint theta0Freqs points
      eigenVec =
        computeInitialEigenVectorR2T0
          numPoint
          numPoint
          theta0Freqs
          thetaFreqs
          points
  powerMethod1
    plan
    folderPath
    numPoint
    numPoint
    numOrientation
    thetaFreqs
    theta0Freqs
    arrR2Z1T0
    numIteration
    writeSourceFlag
    ""
    0.5
    bias
    eigenVec