{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE PatternGuards #-}
{-# LANGUAGE ScopedTypeVariables #-}

module Main
  ( main
  )
where

import Graphics.Gloss hiding (Vector)
import Graphics.Gloss.Interface.Pure.Game hiding (Vector)

import Numeric.LinearAlgebra.Data

import System.Random

import IFS.Base

-- | Main function, when we launch, just start and display the window
main :: IO ()
main = do
  play window white fps initialState render handleEvents stepWorld

-- | Window for the program
window :: Display
window = InWindow "Hacktals" (800, 600) (10, 10)

-- | FPS
fps :: Int
fps = 300

-- | Data structure to hold the state of the fractal
data FractalState = FractalState
  { pointLoc :: Vector R        -- ^ Location of the tracing point
  , points :: [(Float, Float)]  -- ^ Trace of all previous points
  , generator :: StdGen         -- ^ Random generator instance for the "chaos game"
  , autoPlay :: Bool            -- ^ If True, automatically generate the new datapoint
  } deriving Show

-- | Initial state for the program setup
initialState :: FractalState
initialState = FractalState
  { pointLoc = vector [0, 0] -- ^ Start the point at the origin
  , points = [(0, 0)]        -- ^ Starting point
  , generator = mkStdGen 3   -- ^ Start the generator with seed 3
  , autoPlay = False         -- ^ Have autoplay turned off at the start
  }

-- | Rendering function
render :: FractalState -- ^ Fractal info to render
       -> Picture      -- ^ Rendered fractal

render fs = pictures $ map printTrace pts
  where
    pts = points fs
    printTrace (x, y) = translate x y $ color (dark red) $ circleSolid 1

-- | Reactive event handler function
handleEvents :: Event -> FractalState -> FractalState
handleEvents (EventKey (Char 'n') Down _ _) fs = let pl = pointLoc fs
                                                     gen = generator fs
                                                     (pl', (x, y), gen') = chaosStep gen pl
                                                 in fs { pointLoc = pl'
                                                       , generator = gen'
                                                       , points = (x, y):points fs
                                                       }

handleEvents (EventKey (Char 'p') Down _ _) fs = fs { autoPlay = not $ autoPlay fs }
handleEvents _ fs = fs

-- | Stuff to perform at every frame of the animation
stepWorld :: Float -> FractalState -> FractalState
stepWorld _ fs = if autoPlay fs
                 then let pl = pointLoc fs
                          gen = generator fs
                          (pl', (x, y), gen') = chaosStep gen pl
                      in fs { pointLoc = pl'
                            , generator = gen'
                            , points = (x, y):points fs
                            }
                 else fs

chaosStep :: StdGen -> Vector R -> (Vector R, (Float, Float), StdGen)
chaosStep gen pl = let (pl', gen') = chaosIter differentFlake2IFS gen pl
                       (x:y:[]) = toList pl'
                       x' = realToFrac (300 * x)
                       y' = realToFrac (300 * y)
                   in (pl', (x', y'), gen')


import qualified Data.Vector as U
import Data.Complex

main = print . U.sum $ U.replicate (1000000000 :: Int) (1 :+ 1 ::Complex Double)


-- MandelBrot implementation in Haskell
import Data.Complex
import System.Environment
import Control.Parallel.Strategies
import Criterion.Measurement
--import Graphics.UI.SDL

--initSDL :: IO Surface
--initSDL = do
--    SDL.init [SDL.InitEverything]
--    SDL.setVideoMode (floor imageWidth) (floor imageHeight) 32 []
--    SDL.setCaption "Mandelbrot" "Mandelbrot"-   screen <- getVideoSurface
--    white <- SDL.mapRGB (SDL.surfaceGetPixelFormat screen) 255 255 255
--    fillRect screen Nothing white >> SDL.flip screen
--    return screen

mandelbrot :: Num a => a -> Int -> a
mandelbrot a n = iterate (\z -> z^2 + a) 0 !! n

mbrot_seq :: Int -> Int -> Int -> [[Bool]]
mbrot_seq iter w h = [[magnitude (mandelbrot (x :+ y) iter) < 2
                       | x <- [-2, (-2 + incW) ..   2]]
                       | y <- [ 2,  (2 - incH) ..  -2]]
 where incW = 4.0 / (fromIntegral w)
       incH = 4.0 / (fromIntegral h)

mbrot_par :: Int -> Int -> Int  -> [[Bool]]
mbrot_par iter w h = (mbrot_seq iter w h) `using` parList rdeepseq

draw_image :: Int -> Int -> [[Bool]] -> IO()
draw_image x y xs = print "to be implemented"

main = do
    initializeTime
    start_time <- getTime
    args <- getArgs
    if (length args < 4)
        then print "Expected \"par/seq\"iterations, image width, and image height"
        else do
            let nIter = read (args !! 1) :: Int
            let nWidth = read (args !! 2) :: Int
            let nHeight = read (args !! 3) :: Int
            case (args !! 0) of
                "seq" -> print $ mbrot_seq nIter nWidth nHeight
                "par" -> print $ mbrot_par nIter nWidth nHeight
                otherwise -> print "expected first argument to be `par` or `seq`"
            end_time <- getTime
            print $ "Time elapsed : " ++ show (floor $ 10000 * (end_time - start_time)) ++ "ms"


{-# LANGUAGE MonoLocalBinds #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE Strict           #-}
{-# LANGUAGE StrictData       #-}
module FourierPinwheel.Array where

import qualified Data.Array.Accelerate                  as A
import qualified Data.Array.Accelerate.LLVM.PTX as A
import           Data.Array.Repa                        as R
import           Data.Complex
import           Data.List                              as L
import           Data.Vector.Generic                    as VG
import           Data.Vector.Storable                   as VS
import           DFT.Plan
import           Filter.Utils
import           Graphics.Rendering.Chart.Backend.Cairo
import           Graphics.Rendering.Chart.Easy
import           Image.IO
import           Text.Printf
import           Utils.Array
import           Utils.BLAS
import           Utils.Parallel
import FourierMethod.FourierSeries2D

data FPArray vector = FPArray
  { getFPArrayNumXFreq     :: Int
  , getFPArrayNumYFreq     :: Int
  , getFPArrayNumRFreq     :: Int
  , getFPArrayNumThetaFreq :: Int
  , getFPArrayNumRhoFreq   :: Int
  , getFPArrayNumPhiFreq   :: Int
  , getFPArray             :: [vector]
  }

{-# INLINE toMatrixAcc #-}
toMatrixAcc ::
     (VG.Vector vector e, A.Elt e) => FPArray (vector e) -> A.Acc (A.Array A.DIM2 e)
toMatrixAcc (FPArray numXFreq numYFreq numRFreq numThetaFreq _ _ arr) =
  A.use .
  A.fromList (A.Z A.:. (numRFreq * numThetaFreq) A.:. (numXFreq * numYFreq)) .
  VG.toList . VG.concat $
  arr

{-# INLINE parMapFPArray #-}
parMapFPArray ::
     (NFData vector) => (vector -> vector) -> FPArray vector -> FPArray vector
parMapFPArray f (FPArray numXFreq numYFreq numRFreq numThetaFreq numRhoFreq numPhiFreq vecs) =
  FPArray numXFreq numYFreq numRFreq numThetaFreq numRhoFreq numPhiFreq .
  parMap rdeepseq f $
  vecs

{-# INLINE parZipWithFPArray #-}
parZipWithFPArray ::
     (NFData vector)
  => (vector -> vector -> vector)
  -> FPArray vector
  -> FPArray vector
  -> FPArray vector
parZipWithFPArray f (FPArray numXFreq numYFreq numRFreq numThetaFreq numRhoFreq numPhiFreq vecs1) =
  FPArray numXFreq numYFreq numRFreq numThetaFreq numRhoFreq numPhiFreq .
  parZipWith rdeepseq f vecs1 . getFPArray

plotFPArray ::
     DFTPlan
  -> FilePath
  -> FPArray (VS.Vector (Complex Double))
  -> IO (R.Array R.U R.DIM4 (Complex Double))
plotFPArray plan filePath arr = do
  let planIDBackward =
        DFTPlanID
          IDFT1DG
          [ getFPArrayNumThetaFreq arr
          , getFPArrayNumXFreq arr
          , getFPArrayNumYFreq arr
          ]
          [1, 2]
  vecsR2 <- dftExecuteBatchP plan planIDBackward . getFPArray $ arr
  arrR2 <-
    R.computeUnboxedP .
    makeFilter2DInverse .
    R.fromUnboxed
      (R.Z R.:. getFPArrayNumRFreq arr R.:. getFPArrayNumThetaFreq arr R.:.
       getFPArrayNumXFreq arr R.:.
       getFPArrayNumYFreq arr) .
    VS.convert . VS.concat $
    vecsR2
  arr1 <-
    R.sumP . R.sumS . R.map (\x -> magnitude x Prelude.^ 2) . rotate4D2 $ arrR2
  plotImageRepa filePath .
    ImageRepa 8 .
    R.fromUnboxed
      (R.Z R.:. (1 :: Int) R.:. getFPArrayNumXFreq arr R.:.
       getFPArrayNumYFreq arr) .
    -- VG.map (\x -> log (x + 1)) .
    R.toUnboxed $ arr1
  return arrR2

plotFPArrayAcc ::
     (A.Elt (Complex Double))
  => [A.PTX]
  -> FilePath
  -> Int
  -> Double
  -> Double
  -> Int
  -> FPArray (VS.Vector (Complex Double))
  -> IO (R.Array R.U R.DIM4 (Complex Double))
plotFPArrayAcc ptxs filePath numPoints delta periodR2 numBatch arr@(FPArray numXFreq numYFreq numRFreq numThetaFreq _ _ vecs) = do
  let rows = numRFreq * numThetaFreq
      cols = numXFreq * numYFreq
      arrMat =
        A.transpose .
        A.use .
        A.fromList (A.Z A.:. rows A.:. cols) .
        R.toList .
        makeFilter2DInverse .
        fromUnboxed (Z :. rows :. numXFreq :. numXFreq) . VG.convert . VG.concat $
        vecs
  arrR2 <-
    computeFourierSeriesR2StreamAcc
      ptxs
      numXFreq
      numPoints
      rows
      periodR2
      delta
      numBatch
      arrMat
  (sumP . R.map (\x -> magnitude x ** 2) . rotate3D $ arrR2) >>=
    plotImageRepa filePath .
    ImageRepa 8 .
    fromUnboxed (Z :. (1 :: Int) :. numPoints :. numPoints) . toUnboxed
  return .
    computeS .
    R.reshape (Z :. numRFreq :. numThetaFreq :. numPoints :. numPoints) $
    arrR2

plotFPArrayFreqency ::
     FilePath
  -> FPArray (VS.Vector (Complex Double))
  -> IO ()
plotFPArrayFreqency filePath arr = do
  arr1 <-
    R.sumP .
    R.sumS .
    R.map (\x -> magnitude x Prelude.^ 2) .
    rotate4D2 .
    makeFilter2DInverse .
    R.fromUnboxed
      (R.Z R.:. getFPArrayNumRFreq arr R.:. getFPArrayNumThetaFreq arr R.:.
       getFPArrayNumXFreq arr R.:.
       getFPArrayNumYFreq arr) .
    VS.convert . VS.concat . getFPArray $
    arr
  plotImageRepa filePath .
    ImageRepa 8 .
    R.fromUnboxed
      (R.Z R.:. (1 :: Int) R.:. getFPArrayNumXFreq arr R.:.
       getFPArrayNumYFreq arr) .
    -- VG.map (\x -> x^2) .
    R.toUnboxed $
    arr1

plotRThetaDist ::
     FilePath
  -> FilePath
  -> Int
  -> Int
  -> Int
  -> Double
  -> (Int, Int)
  -> R.Array R.U R.DIM4 (Complex Double)
  -> IO ()
plotRThetaDist filePathTheta filePathR numPoints numTheta numR periodEnv (x, y) arr = do
  let (Z :. numRFreq :. numThetaFreq :. _ :. _) = extent arr
      center = div numPoints 2
      centerR = div numR 2
      centerTheta = div numTheta 2
      centerRFreq = div numRFreq 2
      centerThetaFreq = div numThetaFreq 2
      vecRTheta =
        VG.convert . R.toUnboxed . R.computeUnboxedS . R.slice arr $
        (R.Z R.:. R.All R.:. R.All R.:. (x + center) R.:. (y + center))
      deltaTheta = 2 * pi / Prelude.fromIntegral numTheta
      deltaLogR = log periodEnv / Prelude.fromIntegral numR
      matrix =
        VG.convert .
        R.toUnboxed .
        R.computeUnboxedS .
        R.fromFunction (R.Z R.:. numR R.:. numTheta :. numRFreq :. numThetaFreq) $ \(Z :. r :. t :. rf :. tf) ->
          cis
            (fromIntegral ((r - centerR) * (rf - centerRFreq)) * 2 * pi /
             log periodEnv *
             deltaLogR +
             fromIntegral ((t - centerTheta) * (tf - centerThetaFreq)) *
             deltaTheta)
      m = numR * numTheta
      n = 1
      k = numRFreq * numThetaFreq
  vec <- VS.map magnitude <$> gemmBLAS m n k matrix vecRTheta
  let outputArr = fromUnboxed (Z :. numR :. numTheta) . VG.convert $ vec
      outputR = R.toList . sumS $ outputArr
      outputTheta =
        R.toList .
        sumS .
        R.backpermute (Z :. numTheta :. numR) (\(Z :. i :. j) -> (Z :. j :. i)) $
        outputArr
  toFile def filePathTheta $ do
    layout_title .= printf "Theta (%d , %d)" x y
    plot
      (line
         ""
         [ L.zip
             [ fromIntegral (i - centerTheta) * deltaTheta
             | i <- [0 .. numTheta - 1]
             ]
             outputTheta
         ])
  toFile def filePathR $ do
    layout_title .= printf "LogR (%d , %d)" x y
    plot
      (line
         ""
         [ L.zip
             [fromIntegral (i - centerR) * deltaLogR | i <- [0 .. numR - 1]]
             outputR
         ])


module Eventlog.Total(total) where

import Control.Monad.State.Strict (State(), execState, get, put, modify)
import Data.Map (Map, empty, alter)
import Prelude hiding (init, lookup, lines, words, drop, length, readFile)

import Eventlog.Types
import qualified Data.Vector as V
import Statistics.LinearRegression

data Parse =
  Parse
  { totals    :: !(Map Bucket (Double, Double, [(Double, Double)])) -- compute running totals and total of squares
  , count     :: !Int                         -- number of frames
  , times     :: [Double]
  }

parse0 :: Parse
parse0 = Parse{ totals = empty, count = 0, times = [] }

total :: [Frame] -> (Int, Map Bucket (Double, Double, Maybe (Double, Double, Double)))
total fs =
  let parse1 = flip execState parse0 . mapM_ parseFrame $ fs
  in  (
       count parse1
      , fmap (stddev $ fromIntegral (count parse1)) (totals parse1)
      )

stddev :: Double -> (Double, Double, [(Double, Double)]) -> (Double, Double, Maybe (Double, Double, Double))
stddev s0 (s1, s2, samples) = (s1, sqrt (s0 * s2 - s1 * s1) / s0, slope)
  where
    m = maximum values
    mt = maximum timesv
    (timesv, values) = unzip (reverse samples)
    yvect = V.fromList (map (/ m) values)
    xvect = V.fromList (map (/ mt) timesv)
    slope = -- TODO: Distinguish these cases
            case samples of
              [] -> Nothing
              -- Linear regression is meaningless with 1 sample
              [_] -> Nothing
              _
                -- All values are the same leads to NaN r2
                | V.all (1 ==) yvect -> Nothing
                | otherwise -> Just $! linearRegressionRSqr xvect yvect

parseFrame :: Frame -> State Parse ()
parseFrame (Frame time ls) = do
  mapM_ (inserter time) ls
  modify $ \p -> p{ count = count p + 1 }

inserter :: Double -> Sample -> State Parse Double
inserter t (Sample k v) = do
  p <- get
  put $! p { totals = alter (accum t v) k (totals p) }
  return $! v

accum :: Double -> Double -> Maybe (Double, Double, [(Double, Double)]) -> Maybe (Double, Double, [(Double, Double)])
accum t x Nothing  = Just $! ((((,,) $! x) $! (x * x)) $! [(t, x)])
accum t x (Just (y, yy, ys)) = Just $! ((((,,) $! (x + y)) $! (x * x + yy)) $! (t, x):ys)


{-# LANGUAGE DataKinds                                #-}
{-# LANGUAGE FlexibleContexts                         #-}
{-# LANGUAGE FlexibleInstances                        #-}
{-# LANGUAGE MultiParamTypeClasses                    #-}
{-# LANGUAGE PatternSynonyms                          #-}
{-# LANGUAGE RankNTypes                               #-}
{-# LANGUAGE TypeFamilies                             #-}
{-# LANGUAGE TypeOperators                            #-}
{-# LANGUAGE UndecidableInstances                     #-}
{-# OPTIONS_GHC -fplugin GHC.TypeLits.KnownNat.Solver #-}
{-# OPTIONS_GHC -fplugin GHC.TypeLits.Normalise       #-}

module Backprop.Learn.Model.Neural (
  -- * Fully connected
  -- ** Feed-forward
    FC, pattern FC, FCp, fcBias, fcWeights
  -- *** With activation function
  , FCA, pattern FCA, _fcaActivation
  -- ** Recurrent
  , FCR, pattern FCR, FCRp, fcrBias, fcrInputWeights, fcrStateWeights
  -- *** With activation function
  , FCRA, pattern FCRA, _fcraStore, _fcraActivation
  ) where

import           Backprop.Learn.Model.Combinator
import           Backprop.Learn.Model.Regression
import           Backprop.Learn.Model.State
import           Data.Tuple
import           GHC.TypeNats
import           Lens.Micro
import           Numeric.Backprop
import           Numeric.LinearAlgebra.Static.Backprop
import qualified Numeric.LinearAlgebra.Static          as H

-- | Fully connected feed-forward layer with bias.  Parameterized by its
-- initialization distribution.
--
-- Note that this has no activation function; to use as a model with
-- activation function, chain it with an activation function using 'RMap',
-- ':.~', etc.; see 'FCA' for a convenient type synonym and constructor.
--
-- Without any activation function, this is essentially a multivariate
-- linear regression.
--
-- With the logistic function as an activation function, this is
-- essentially multivariate logistic regression. (See 'Logistic')
type FC i o = LinReg i o

pattern FC :: FC i o
pattern FC = LinReg

-- | Convenient synonym for an 'FC' post-composed with a simple
-- parameterless activation function.
type FCA i o = RMap (R o) (R o) (FC i o)

-- | Construct an 'FCA' using a generating function and activation
-- function.
--
-- Some common ones include 'logistic' and @'vmap' 'reLU'@.
pattern FCA
    :: (forall s. Reifies s W => BVar s (R o) -> BVar s (R o))          -- ^ '_fcaActivation'
    -> FCA i o
pattern FCA { _fcaActivation } = RM _fcaActivation FC

type FCp = LRp

fcWeights :: Lens (FCp i o) (FCp i' o) (L o i) (L o i')
fcWeights = lrBeta

fcBias :: Lens' (FCp i o) (R o)
fcBias = lrAlpha

-- | Fully connected recurrent layer with bias.
--
-- Parameterized by its initialization distributions, and also the function
-- to compute the new state from previous input.
--
-- @
-- instance 'Learn' ('R' i) (R o) ('FCR' h i o) where
--     type 'LParamMaybe' (FCR h i o) = ''Just' ('FCRp' h i o)
--     type 'LStateMaybe' (FCR h i o) = 'Just (R h)
-- @
type FCR h i o = Recurrent (R (i + h)) (R i) (R h) (R o) (FC (i + h) o)

-- | Construct an 'FCR'
pattern FCR
    :: (KnownNat h, KnownNat i)
    => (forall s. Reifies s W => BVar s (R o) -> BVar s (R h))          -- ^ '_fcrSTore'
    -> FCR h i o
pattern FCR { _fcrStore } <-
    Rec { _recLoop  = _fcrStore
        , _recLearn = FC
        }
  where
    FCR s = Rec { _recSplit = H.split
                , _recJoin  = (H.#)
                , _recLoop  = s
                , _recLearn = FC
                }
{-# COMPLETE FCR #-}

-- | Convenient synonym for an 'FCR' post-composed with a simple
-- parameterless activation function.
type FCRA h i o = RMap (R o) (R o) (FCR h i o)

-- | Construct an 'FCRA' using a generating function and activation
-- function.
--
-- Some common ones include 'logistic' and @'vmap' 'reLU'@.
pattern FCRA
    :: (KnownNat h, KnownNat i)
    => (forall s. Reifies s W => BVar s (R o) -> BVar s (R h))          -- ^ '_fcraStore'
    -> (forall s. Reifies s W => BVar s (R o) -> BVar s (R o))          -- ^ '_fcraActivation'
    -> FCRA h i o
pattern FCRA { _fcraStore, _fcraActivation }
    = RM _fcraActivation (FCR _fcraStore)

type FCRp h i o = FCp (i + h) o

lensIso :: (s -> (a, x)) -> ((b, x) -> t) -> Lens s t a b
lensIso f g h x = g <$> _1 h (f x)

fcrInputWeights
    :: (KnownNat h, KnownNat i, KnownNat i', KnownNat o)
    => Lens (FCRp h i o) (FCRp h i' o) (L o i) (L o i')
fcrInputWeights = fcWeights . lensIso H.splitCols (uncurry (H.|||))

fcrStateWeights
    :: (KnownNat h, KnownNat h', KnownNat i, KnownNat o)
    => Lens (FCRp h i o) (FCRp h' i o) (L o h) (L o h')
fcrStateWeights = fcWeights . lensIso (swap . H.splitCols) (uncurry (H.|||) . swap)

fcrBias :: Lens' (FCRp h i o) (R o)
fcrBias = fcBias


{-# LANGUAGE FlexibleContexts #-}

module Main (
  hist,
  main
) where

-- | Global imports
import Data.List
import System.Random

import System.Environment (
    getArgs
  )

-- | Local imports
import Calculus
import Combinatorics
import Distribution
import Geometry
import Random
import Statistics
import Units
import Vec2
import Vec3

-- | Generate histogram
-- | xs - Input list
-- | n - Number of buckets
hist xs n = [floor (min (n - 1)) ((x - l) / i) | x <- xs]
  where
    s = sort xs
    l = head s
    r = last s
    range = r - l
    i = range / n

-- | Main entry point
main :: IO ()
main =  do
  [n] <- getArgs
  g <- getStdGen
  print n
  print . take (read n :: Int) $ rand (0 :: Float) (1 :: Float) g


module Statistics.Quantile.Bench.Gen(
    gauss
  , gamma
  , sampleGaussian
  , sampleGamma
) where

import Control.Applicative
import Control.Monad

import Data.Random (runRVar)
import Data.Random.Distribution
import Data.Random.Distribution.Normal
import qualified Data.Random.Distribution.Gamma as G

import Pipes

import Statistics.Quantile.Types

import System.Random.MWC

gauss :: Double -> Double -> Producer Double IO ()
gauss = runSampler sampleGaussian

gamma :: Double -> Double -> Producer Double IO ()
gamma = runSampler sampleGamma

runSampler :: (GenIO -> Double -> Double -> IO Double)
           -> Double
           -> Double
           -> Producer Double IO ()
runSampler sample a b = do
  gen <- liftIO createSystemRandom
  forever $ do
    x <- liftIO $ sample gen a b
    yield x

sampleGaussian :: GenIO -> Double -> Double -> IO Double
sampleGaussian gen mu sigma = runRVar (normal mu sigma) gen

sampleGamma :: GenIO -> Double -> Double -> IO Double
sampleGamma gen alpha beta = runRVar (G.gamma alpha beta) gen


-- file: Astro/PostNewtonian.hs
-- Post-Newtonian approximations for two-body problem

module Astro.PostNewtonian
	where

import Astro.Misc
import Numeric.LinearAlgebra

import Data.List (foldl')

-- Calculate Post-Newtonian acceleration functions for non-spin
-- dependent terms
mkNonSpinPNs :: Double->Double->Double->Double->
	[(Vector Double->Vector Double->Vector Double)]
mkNonSpinPNs gc c m1 m2 = 	[ a1
				, a2
				, a2_5
				]
			where
	-- Keplerian contribution
	--a0 :: Vector Double -> Vector Double -> Vector Double
	--a0 rv vv = scale (-gm / r^3) rv where r = norm rv

	a1 :: Vector Double -> Vector Double -> Vector Double
	a1 rv vv = {-# SCC "a1" #-}
		scale (-gmperc2 / r^2) $
			scale nbrac nv + scale vbrac vv
		where
			nbrac = -2*(2+eta)*gm/r + (1+3*eta)*v^2 - 3/2*eta*rd^2
			vbrac = -2*(2-eta)*rd
			r	= norm rv
			v	= norm vv
			nv	= unitV rv
			rd	= (rv <.> vv)/r

	a2 :: Vector Double -> Vector Double -> Vector Double
	a2 rv vv = {-# SCC "a2" #-}
		scale (- gmperc2 / (c^2 * r^2)) $
			scale (nbrac/r) rv + scale ((-rd/2) * vbrac) vv
		where
			nbrac = 3/4*(12+29*eta)*(gm/r)^2 + eta*(3-4*eta)*v^4
				+ 15/8*eta*(1-3*eta)*rd^4
				- 3/2*eta*(3-4*eta)*v^2*rd^2
				- 1/2*eta*(13-4*eta)*(gm/r)*v^2
				- (2+25*eta+2*eta^2)*(gm/r)*rd^2
			vbrac = eta*(15+4*eta)*v^2 - (eta+41*eta+8*eta^2)*(gm/r)
				- 3*eta*(3+2*eta)*rd^2
			r	= norm rv
			v	= norm vv
			rd	= (rv <.> vv)/r

	a2_5 :: Vector Double -> Vector Double -> Vector Double
	a2_5 rv vv = {-# SCC "a2_5" #-}
		scale (8/15*(gmperc2)^2*eta/(c*r^2)) $
			scale (nbrac*rd/r) rv + scale (-vbrac) vv
		where
			nbrac 	= 9*v^2 + 17*gm/r
			vbrac 	= 3*v^2 + 9*gm/r
			r	= norm rv
			v	= norm vv
			rd	= (rv <.> vv)/r
	gm 	= gc * m
	gmperc2 = gm / c ^2
	m 	= m1+m2
	eta 	= m1*m2/m^2

-- Calculate Post-Newtonian acceleration functions for spin
-- dependent terms
--mkSpinPNs :: Double->Double->Double->Double->Double->Double->
--	[(Vector Double->Vector Double->Vector Double->Vector Double)]
mkSpinPNs gc c m1 m2 chi q = 	[ spinOrbit
				, quadMonoInteraction
				] where
	--spinOrbit :: Vector Double -> Vector Double -> Vector Double
	--		-> Vector Double
	spinOrbit rv vv s = {-# SCC "SO" #-}
		scale (gm^2/(r^3*c^3)*(1+sqrt (1-4*eta))/4*chi) $
				scale nbrac nv + scale (nsbrac) (cross nv s')
					+ scale (-vsbrac) (cross vv s')
		where
			nbrac = 12 * s' <.> (cross nv vv)
			nsbrac = (9 + 3 * sqrt (1-4*eta))*rd
			vsbrac = 7 + sqrt(1-4*eta)
			r	= norm rv
			v	= norm vv
			nv	= unitV rv
			rd	= (rv <.> vv)/r
			s'	= spinCylToSpher s

	--quadMonoInteraction :: Vector Double -> Vector Double -> Vector Double
	--			-> Vector Double
	quadMonoInteraction rv vv s = {-# SCC "Q" #-}
		-- FIXME: Minus or no minus? Depends on whether
		-- coordinates fixed on m1 or m2, see OJ287 paper, minus
		-- in paper.
		scale (-q*chi^2*3*gc^3*m1^2*m/(2*c^4*r^4)) $
			scale nbrac nv + scale (-2 * nv <.> s') s'
		where
			nbrac = 5 * (nv <.> s')^2 - 1
			r	= norm rv
			v	= norm vv
			nv	= unitV rv
			s'	= spinCylToSpher s
	gm 	= gc * m
	m 	= m1+m2
	eta 	= m1*m2/m^2

-- Time derivative of spin unit vector, in spherical coordinates
spinDt :: Double->Double->Double->Double->Double->Double->
	Vector Double->Vector Double->Vector Double->Vector Double
spinDt gc c m1 m2 chi q rv vv s = cross omega s
	where
	omega = scale (gm*eta/(2*c^2*r^2)*(7+sqrt (1-4*eta))
				/(1+sqrt (1-4*eta))) (cross nv vv)
	r = norm rv
	nv = scale (1/r) rv
	gm 	= gc * m
	m 	= m1+m2
	eta 	= m1*m2/m^2

-- Time derivatives of spin unit vector components in cylindrical
-- coordinates, with x as the axisymmery axis
spinDtCyl :: Double->Double->Double->Double->Double->Double->
	Vector Double->Vector Double->[(Double,Double)->Double]
spinDtCyl gc c m1 m2 chi q rv vv = [dtheta, dxi]
	where
	dtheta (th,xi) =
		-xiPerRho*(dhs@>1*cos th + dhs@>2*sin th) + dhs@>0
		where
		xiPerRho = if abs xi >= 1 then 0 else xi/sqrt (1-xi^2)
	dxi (th,xi) =
		rho * (dhs@>1*sin th - dhs@>2 * cos th)
		where
		rho = if abs xi >= 1 then 0 else sqrt $ 1 - xi^2
	r = norm rv
	nv = scale (1/r) rv
	dhs = omega
	omega = scale (gm*eta/(2*c^2*r^2)*(7+sqrt (1-4*eta))
				/(1+sqrt (1-4*eta))) (cross nv vv)
	gm 	= gc * m
	m 	= m1+m2
	eta 	= m1*m2/m^2

-- Convert spin from cylindrical to cartesian representation
-- Axisymmetry axis is x-axis
spinCylToSpher :: (Double,Double)->Vector Double
spinCylToSpher (theta,xi) = --3|>[rho * cos theta, rho*sin theta, xi]
	3|>[xi, rho * cos theta, rho*sin theta]
	where rho = if abs xi >= 1 then 0 else sqrt $ 1 - xi^2

-- Post-Newtonian acceleration functions for highest supported degree
nonSpinPNAcc gc c m1 m2 rv vv =
	sum [f rv vv | f <- mkNonSpinPNs gc c m1 m2]
spinPNAcc gc c m1 m2 chi q rv vv s =
	sum [f rv vv s | f <- mkSpinPNs gc c m1 m2 chi q]

totalPNAcc gc c m1 m2 chi q rv vv s = ((f1 rv vv) + (f2 rv vv s))
	where (f1, f2) = (nonSpinPNAcc gc c m1 m2,
				spinPNAcc gc c m1 m2 chi q)

-- Explicitly calculated total acceleration
totalPNAccExp gc c m1 m2 chi q rv vv s =
	a1 + spinOrbit + quadMono + a2 + a2_5
	where
	a1 = {-# SCC "a1" #-}
		scale (-gmperc2 / r^2) $
			scale nbrac nv + scale vbrac vv
		where
			nbrac = -2*(2+eta)*gm/r + (1+3*eta)*v^2 - 3/2*eta*rd^2
			vbrac = -2*(2-eta)*rd

	a2 = {-# SCC "a2" #-}
		scale (- gmperc2 / (c^2 * r^2)) $
			scale nbrac nv + scale ((-rd/2) * vbrac) vv
		where
			nbrac = 3/4*(12+29*eta)*(gm/r)^2 + eta*(3-4*eta)*v^4
				+ 15/8*eta*(1-3*eta)*rd^4
				- 3/2*eta*(3-4*eta)*v^2*rd^2
				- 1/2*eta*(13-4*eta)*(gm/r)*v^2
				- (2+25*eta+2*eta^2)*(gm/r)*rd^2
			vbrac = eta*(15+4*eta)*v^2 - (eta+41*eta+8*eta^2)*(gm/r)
				- 3*eta*(3+2*eta)*rd^2

	a2_5 = {-# SCC "a2_5" #-}
		scale (8/15*(gmperc2)^2*eta/(c*r^2)) $
			scale (nbrac*rd) nv + scale (-vbrac) vv
		where
			nbrac 	= 9*v^2 + 17*gm/r
			vbrac 	= 3*v^2 + 9*gm/r

	spinOrbit = {-# SCC "SO" #-}
		scale (gm^2/(r^3*c^3)*(1+sqrt (1-4*eta))/4*chi) $
				scale nbrac nv + scale (nsbrac) (cross nv s')
					+ scale (-vsbrac) (cross vv s')
		where
			nbrac = 12 * s' <.> (cross nv vv)
			nsbrac = (9 + 3 * sqrt (1-4*eta))*rd
			vsbrac = 7 + sqrt(1-4*eta)

	quadMono = {-# SCC "Q" #-}
		scale (-q*chi^2*3*gc^3*m1^2*m/(2*c^4*r^4)) $
			scale nbrac nv + scale (-2 * nv <.> s') s'
		where
			nbrac = 5 * (nv <.> s')^2 - 1
	r	= norm rv
	v	= norm vv
	s'	= spinCylToSpher s
	nv	= scale (1/r) rv
	rd	= (rv <.> vv)/r
	gm 	= gc * m
	gmperc2 = gm / c ^2
	m 	= m1+m2
	eta 	= m1*m2/m^2


module MachineLearning.LogisticModelTest
(
  tests
)

where

import Test.Framework (testGroup)
import Test.Framework.Providers.HUnit
import Test.HUnit
import Test.HUnit.Approx
import Test.HUnit.Plus

import MachineLearning.DataSets (dataset2)

import qualified Numeric.LinearAlgebra as LA
import qualified MachineLearning as ML
import MachineLearning.Optimization (checkGradient)
import MachineLearning.Model
import MachineLearning.LogisticModel
import MachineLearning.Regularization (Regularization(..))

(x, y) = ML.splitToXY dataset2

x1 = ML.addBiasDimension $ ML.mapFeatures 6 x
onesTheta :: LA.Vector LA.R
onesTheta = LA.konst 1 (LA.cols x1)
zeroTheta :: LA.Vector LA.R
zeroTheta = LA.konst 0 (LA.cols x1)

gradientCheckingEps = 1e-3

checkGradientTest lambda theta = do
  let diffs = take 5 $ map (\e -> checkGradient Logistic lambda x1 y theta e) [1e-3, 1.1e-3 ..]
      diff = minimum diffs
  assertApproxEqual (show theta) gradientCheckingEps 0 diff

tests = [ testGroup "sigmoid" [
            testCase "zero" $ assertApproxEqual "" 1e-10 0.5 (sigmoid 0)
            , testCase "big positive value" $ assertApproxEqual "" 1e-10 1 (sigmoid 10e10)
            , testCase "big negative value" $ assertApproxEqual "" 1e-10 0 (sigmoid $ -10e10)
            , testCase "quite big positive value" $ assertApproxEqual "" 1e-2 1 (sigmoid 100)
            , testCase "quite big negative value" $ assertApproxEqual "" 1e-2 0 (sigmoid $ -100)
            ]
          , testGroup "sigmoidGradient" [
              testCase "zero" $ assertApproxEqual "" 1e-10 0.25 (sigmoidGradient 0)
              , testCase "big positive value" $ assertApproxEqual "" 1e-10 0 (sigmoidGradient 10e10)
              , testCase "big negative value" $ assertApproxEqual "" 1e-10 0 (sigmoidGradient $ -10e10)
              , testCase "quite big positive value" $ assertApproxEqual "" 1e-2 0 (sigmoidGradient 100)
              , testCase "quite big negative value" $ assertApproxEqual "" 1e-2 0 (sigmoidGradient $ -100)
              , testCase "small positive value" $ assertApproxEqual "" 1e-2 0.2 (sigmoidGradient 1)
              , testCase "small negative value" $ assertApproxEqual "" 1e-2 0.2 (sigmoidGradient $ -1)
              ]
          , testGroup "model" [
              testCase "cost, lambda = 0" $ assertApproxEqual "" 1e-3 2.020 (cost Logistic (L2 0) x1 y onesTheta)
              , testCase "cost, lambda = 1" $ assertApproxEqual "" 1e-3 2.135 (cost Logistic (L2 1) x1 y onesTheta)
              , testCase "cost, lambda = 1000" $ assertApproxEqual "" 1e-3 116.427 (cost Logistic (L2 1000) x1 y onesTheta)
              , testCase "gradient, lambda = 0" $ assertVector "" 1e-5 gradient_l0 (gradient Logistic (L2 0) x1 y onesTheta)
              , testCase "gradient, lambda = 1" $ assertVector "" 1e-5 gradient_l1 (gradient Logistic (L2 1) x1 y onesTheta)
              , testCase "gradient, lambda = 1000" $ assertVector "" 1e-5 gradient_l1000 (gradient Logistic (L2 1000) x1 y onesTheta)
              ]
          , testGroup "gradient checking" [
              testCase "non-zero theta, non-zero lambda" $ checkGradientTest (L2 2) onesTheta
              , testCase "zero theta, non-zero lambda" $ checkGradientTest (L2 2) zeroTheta
              , testCase "non-zero theta, zero lambda" $ checkGradientTest (L2 0) onesTheta
              , testCase "zero theta, zero lambda" $ checkGradientTest (L2 0) zeroTheta
              , testCase "non-zero theta, no reg" $ checkGradientTest RegNone onesTheta
              , testCase "zero theta, no reg" $ checkGradientTest RegNone zeroTheta

              ]
        ]

gradient_l0 = LA.vector [0.34604507367924525,7.660615656904722e-2,0.11004999290013262,0.14211701951318526,7.4399123914273965e-3,0.15963981400206023,5.864636026438637e-2,2.369595228045015e-2,1.7568631688383615e-2,9.87226947583087e-2,8.878427090266403e-2,2.509755728155993e-3,3.348199373717313e-2,1.0975419586624715e-3,0.11520318246520422,5.0480764463147594e-2,1.0229511332420786e-2,8.818652179760128e-3,1.5052078581608905e-2,6.655810974747264e-3,9.010665405440292e-2,6.480865498500844e-2,2.039892642614106e-3,1.4231095091583643e-2,5.737477462013599e-4,1.71609000731311e-2,-2.437841009425525e-4,9.753746346629336e-2]

gradient_l1 = LA.vector [0.34604507367924525,8.508073284023365e-2,0.11852456917131905,0.15059159578437167,1.5914488662613836e-2,0.16811439027324665,6.71209365355728e-2,3.217052855163659e-2,2.6043207959570054e-2,0.10719727102949513,9.725884717385046e-2,1.0984331999342433e-2,4.1956570008359576e-2,9.572118229848912e-3,0.12367775873639067,5.895534073433403e-2,1.8704087603607224e-2,1.729322845094657e-2,2.3526654852795346e-2,1.5130387245933704e-2,9.858123032558937e-2,7.328323125619488e-2,1.0514468913800546e-2,2.270567136277008e-2,9.048324017387801e-3,2.563547634431754e-2,8.230792170243889e-3,0.10601203973747979]

gradient_l1000 = LA.vector [0.34604507367924525,8.551182427755489,8.584626264086573,8.616693290699626,8.48201618357787,8.6342160851885,8.533222631450828,8.498272223466891,8.492144902874823,8.57329896594475,8.563360542089104,8.477086026914597,8.508058264923614,8.475673813145104,8.589779453651644,8.525057035649588,8.484805782518862,8.483394923366202,8.489628349768049,8.481232082161188,8.564682925240843,8.539384926171449,8.476616163829055,8.488807366278024,8.475150018932643,8.491737171259572,8.474332487085498,8.572113734652733]


{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE DeriveAnyClass   #-}
{-# LANGUAGE DeriveGeneric    #-}
{-# LANGUAGE FlexibleContexts #-}
module FokkerPlanck.FourierSeries
  ( fourierMellin
  , sampleScale
  -- , computeFourierCoefficients
  , computeHarmonicsArray
  , computeHarmonicsArraySparse
  -- , computeHarmonicsArrayGPU
  , getHarmonics
  , computeThetaRHarmonics
  , computeFourierSeriesThetaR
  -- , computeFourierSeriesR2
  , normalizeFreqArr
  -- , normalizeFreqArr'
  , plotThetaDimension
  , computeFourierSeriesOfLogPolarHarmonicsArray
  , computeRectangularInverseHarmonics
  ) where

import           Array.UnboxedArray                     as UA
import           Control.DeepSeq
import           Data.Array.IArray                      as IA
import           Data.Array.Repa                        as R
import           Data.Binary
import           Data.Complex
import           Data.DList                             as DL
import           Data.List                              as L
import           Data.Vector.Generic                    as VG
import           Data.Vector.Unboxed                    as VU
import           FokkerPlanck.BrownianMotion
import           FokkerPlanck.Histogram
import           GHC.Generics                           (Generic)
import           Text.Printf
import           Utils.Array
import           Utils.Parallel
import           Utils.Time
import           Graphics.Rendering.Chart.Backend.Cairo
import           Graphics.Rendering.Chart.Easy
import           Image.IO
import           System.FilePath
import           Utils.SimpsonRule
import Numeric.LinearAlgebra as NL
import Numeric.LinearAlgebra.Data as NL
import            Utils.Distribution
import Pinwheel.Base

-- {-# INLINE fourierMellin' #-}
-- fourierMellin' :: Int -> Int -> (Double, Double) -> Complex Double
-- fourierMellin' !angularFreq !radialFreq (!x, !y) =
--   if x == 0 && y == 0
--     then 0
--     else exp $ (-1 * log rho) :+ (tf * atan2 y x + rf * (log rho))

{-# INLINE sampleScale #-}
sampleScale ::
     Double
  -> Double
  -> Double
  -> Double
  -> Double
  -> [Particle]
  -> Complex Double
sampleScale !phiFreq !rhoFreq !thetaFreq !rFreq !halfLogPeriod =
  L.sum .
  L.map
    (\(Particle newPhi newRho theta newR _) ->
       if newRho == 0 -- || newR == 0
         then 0
         else (cos (phiFreq * newPhi + thetaFreq * (theta - newPhi)) :+ 0) *
              (cis $
               (-pi) * (rhoFreq * (newRho) + rFreq * (newR - newRho)) /
               halfLogPeriod))

-- {-# INLINE computeFourierCoefficients #-}
-- computeFourierCoefficients ::
--      [Double]
--   -> [Double]
--   -> [Double]
--   -> [Double]
--   -> Double
--   -> Double
--   -> [DList Particle]
--   -> Histogram (Complex Double)
-- computeFourierCoefficients !phiFreqs !rhoFreqs !thetaFreqs !rFreqs !halfLogPeriod !deltaLogRho !xs =
--   Histogram
--     [L.length phiFreqs, L.length rhoFreqs, L.length thetaFreqs, L.length rFreqs]
--     (L.length ys) .
--   toUnboxedVector .
--   UA.accum
--     (+)
--     0
--     ( (1, 1, 1, 1)
--     , ( L.length rFreqs
--       , L.length thetaFreqs
--       , L.length rhoFreqs
--       , L.length phiFreqs)) .
--   L.concat .
--   parMap
--     rdeepseq
--     (\(particle@(Particle phi rho theta r)) ->
--        let !n = floor $ (log r + halfLogPeriod) / deltaLogRho
--            !samples =
--              -- L.map moveParticle  $
--              particle :
--              [ -- (Particle
--              --      phi
--              --      rho
--              --      theta
--              --      (exp $ (fromIntegral i * deltaLogRho - halfLogPeriod)))
--              -- | i <- [0 .. n]
--              ]
--        in L.map
--             (\((rFreq, i), (thetaFreq, j), (rhoFreq, k), (phiFreq, l)) ->
--                ( (i, j, k, l)
--                , -- (deltaLogRho :+ 0) *
--                  1/ (16 * pi^2 * halfLogPeriod :+ 0) *
--                  sampleScale
--                    phiFreq
--                    rhoFreq
--                    thetaFreq
--                    rFreq
--                    halfLogPeriod
--                    samples))
--             freqs) $
--   ys
--   where
--     !ys = DL.toList . DL.concat $ xs
--     !freqs =
--       [ (rFreq', thetaFreq', rhoFreq', phiFreq')
--       | rFreq' <- L.zip rFreqs [1 ..]
--       , thetaFreq' <- L.zip thetaFreqs [1 ..]
--       , rhoFreq' <- L.zip rhoFreqs [1 ..]
--       , phiFreq' <- L.zip phiFreqs [1 ..]
--       ]

-- {-# INLINE normalizeFreqArr #-}
-- normalizeFreqArr ::
--      Double
--   -> [Double]
--   -> [Double]
--   -> R.Array U DIM4 (Complex Double)
--   -> R.Array U DIM4 (Complex Double)
-- normalizeFreqArr !std !phiFreqs !rhoFreqs arr =
--   computeUnboxedS .
--   R.traverse3
--     arr
--     (fromListUnboxed (Z :. L.length phiFreqs) phiFreqs)
--     (fromListUnboxed (Z :. L.length rhoFreqs) rhoFreqs)
--     (\sh _ _ -> sh) $ \fArr fPhi fRho idx@(Z :. r :. theta :. rho :. phi) ->
--     fArr idx *
--     ((exp $
--       (-1) *
--       ((fPhi (Z :. phi)) ^ 2 + (fPhi (Z :. theta)) ^ 2 + (fRho (Z :. rho)) ^ 2 +
--        (fRho (Z :. r)) ^ 2) /
--       2 /
--       (std ^ 2)) :+
--      0)

{-# INLINE normalizeFreqArr #-}
normalizeFreqArr ::
     Double
  -> Double
  -> [Double]
  -> [Double]
  -> R.Array U DIM4 (Complex Double)
  -> R.Array U DIM4 (Complex Double)
normalizeFreqArr !stdA !stdR !phiFreqs !rhoFreqs arr
  | stdA == 0 && stdR == 0 = arr
  | stdA == 0 =
    computeUnboxedS .
    R.traverse3
      arr
      (fromListUnboxed (Z :. L.length phiFreqs) phiFreqs)
      (fromListUnboxed (Z :. L.length rhoFreqs) rhoFreqs)
      (\sh _ _ -> sh) $ \fArr fPhi fRho idx@(Z :. r :. theta :. rho :. phi) ->
      fArr idx *
      (exp $
       ((-1) * ((fRho (Z :. rho)) ^ 2 + (fRho (Z :. r)) ^ 2) / (2 * stdR ^ 2)) :+
       0)
  | stdR == 0 =
    computeUnboxedS .
    R.traverse3
      arr
      (fromListUnboxed (Z :. L.length phiFreqs) phiFreqs)
      (fromListUnboxed (Z :. L.length rhoFreqs) rhoFreqs)
      (\sh _ _ -> sh) $ \fArr fPhi fRho idx@(Z :. _ :. theta :. rho :. phi) ->
      fArr idx *
      ((exp $
        (-1) * ((fPhi (Z :. phi)) ^ 2 + (fPhi (Z :. theta)) ^ 2) /
        (2 * stdA ^ 2)) :+
       0)
  | otherwise =
    computeUnboxedS .
    R.traverse3
      arr
      (fromListUnboxed (Z :. L.length phiFreqs) phiFreqs)
      (fromListUnboxed (Z :. L.length rhoFreqs) rhoFreqs)
      (\sh _ _ -> sh) $ \fArr fPhi fRho idx@(Z :. r :. theta :. rho :. phi) ->
      fArr idx *
      ((exp $
        (-1) * ((fPhi (Z :. phi)) ^ 2 + (fPhi (Z :. theta)) ^ 2) /
        (2 * stdA ^ 2) -
        ((fRho (Z :. rho)) ^ 2 + (fRho (Z :. r)) ^ 2) / (2 * stdR ^ 2)) :+
       0)

-- {-# INLINE computeHarmonicsArray #-}
computeHarmonicsArray ::
     (VG.Vector vector (Complex Double), NFData (vector (Complex Double)))
  => Int
  -> Double
  -> Int
  -> Double
  -> [Double]
  -> [Double]
  -> [Double]
  -> [Double]
  -> Double
  -> Double
  -> IA.Array (Int, Int) (vector (Complex Double))
computeHarmonicsArray !numRows !deltaRow !numCols !deltaCol !phiFreqs !rhoFreqs !thetaFreqs !rFreqs !halfLogPeriod !cutoff =
  let !centerRow = div numRows 2
      !centerCol = div numCols 2
      rangeFunc1 xs ys = [(L.head xs - L.last ys) .. (L.last xs - L.head ys)]
      rangeFunc2 xs ys =
        (round (L.head xs - L.last ys), round (L.last xs - L.head ys))
      (!tfLB, !tfUB) = rangeFunc2 phiFreqs thetaFreqs
      (!rfLB, !rfUB) = rangeFunc2 rhoFreqs rFreqs
      !xs =
        parMap
          rdeepseq
          (\(!rf, !tf) ->
             let !vec =
                   VG.convert .
                   toUnboxed . computeS . fromFunction (Z :. numCols :. numRows) $ \(Z :. c :. r) ->
                     let !x = fromIntegral (c - centerCol) * deltaCol
                         !y = fromIntegral (r - centerRow) * deltaRow
                         !rho = 0 + (sqrt $ x ^ 2 + y ^ 2)
                         !rho2 =
                           fromIntegral $
                           (c - centerCol) ^ 2 + (r - centerRow) ^ 2
                     in if rho2 > cutoff ^ 2 || (rho <= 0) -- || pi * rho < (abs tf) -- || log ((rho + 1) / rho) > (1 / (2 * rf))
                          then 0
                               -- (x :+ y) ** (tf :+ 0) *
                               -- ((x ^ 2 + y ^ 2) :+ 0) **
                               -- (((-tf - 0.5) :+ rf) / 2)
                          else ((rho :+ 0) ** ((-1) :+ rf)) *
                               (cis (tf * atan2 y x))
             in ((round rf, round tf), vec))
          [ (rf, tf)
          | rf <- rangeFunc1 rhoFreqs rFreqs
          , tf <- rangeFunc1 phiFreqs thetaFreqs
          ]
  in IA.array ((rfLB, tfLB), (rfUB, tfUB)) xs

-- {-# INLINE computeHarmonicsArraySparse #-}
computeHarmonicsArraySparse ::
     Int
  -> Double
  -> Int
  -> Double
  -> [Double]
  -> [Double]
  -> [Double]
  -> [Double]
  -> Double
  -> Double
  -> Double
  -> IA.Array (Int, Int) (R.Array U DIM2 (Complex Double))
computeHarmonicsArraySparse !numRows !deltaRow !numCols !deltaCol !phiFreqs !rhoFreqs !thetaFreqs !rFreqs !halfLogPeriod !cutoff !envelopeSigma =
  let !centerRow = div numRows 2
      !centerCol = div numCols 2
      rangeFunc1 xs ys =
        [(round $ L.head xs - L.last ys) .. (round $ L.last xs - L.head ys)]
      rangeFunc2 xs ys =
        ((round $ (L.head xs - L.last ys)), (round $ (L.last xs - L.head ys)))
      (!tfLB, !tfUB) = rangeFunc2 phiFreqs thetaFreqs
      (!rfLB, !rfUB) = rangeFunc2 rhoFreqs rFreqs
      !xs =
        parMap
          rseq
          (\(!rf, !tf) ->
             let !arr =
                   computeS . fromFunction (Z :. numCols :. numRows) $ \(Z :. c :. r) ->
                     let !x = fromIntegral (c - centerCol) * deltaCol
                         !y = fromIntegral (r - centerRow) * deltaRow
                         !rho = ((sqrt $ x ^ 2 + y ^ 2))
                         !rho2 =
                           fromIntegral $
                           (c - centerCol) ^ 2 + (r - centerRow) ^ 2
                     in if (rho <= 0) || rho2 > cutoff ^ 2
                          then 0
                          else (x :+ y) ** (fromIntegral tf :+ 0) *
                               (((x ^ 2 + y ^ 2) :+ 0) **
                                (((-(fromIntegral tf) - envelopeSigma) :+ (2 * pi / halfLogPeriod * fromIntegral rf)) /
                                 2))
             in deepSeqArray arr ((rf, tf), arr))
          [ (rf, tf)
          | rf <- rangeFunc1 rhoFreqs rFreqs
          , tf <- rangeFunc1 phiFreqs thetaFreqs
          ]
  in IA.array ((rfLB, tfLB), (rfUB, tfUB)) xs

-- {-# INLINE computeHarmonicsArrayGPU #-}
-- computeHarmonicsArrayGPU ::
--      (VG.Vector vector (Complex Double), NFData (vector (Complex Double)))
--   => Int
--   -> Double
--   -> Int
--   -> Double
--   -> [Double]
--   -> [Double]
--   -> [Double]
--   -> [Double]
--   -> Double
--   -> Double
--   -> IA.Array (Int, Int) (vector (Complex Double))
-- computeHarmonicsArrayGPU !numRows !deltaRow !numCols !deltaCol !phiFreqs !rhoFreqs !thetaFreqs !rFreqs !halfLogPeriod !cutoff =
--   let !centerRow = div numRows 2
--       !centerCol = div numCols 2
--       rangeFunc1 xs ys =
--         [round (L.head xs - L.last ys) .. round (L.last xs - L.head ys)]
--       rangeFunc2 xs ys =
--         (round (L.head xs - L.last ys), round (L.last xs - L.head ys))
--       (!tfLB, !tfUB) = rangeFunc2 phiFreqs thetaFreqs
--       (!rfLB, !rfUB) = rangeFunc2 rhoFreqs rFreqs
--       !xs =
--         parMap
--           rdeepseq
--           (\(!rf, !tf) ->
--              let !vec =
--                    VG.convert .
--                    toUnboxed . computeS . fromFunction (Z :. numCols :. numRows) $ \(Z :. c :. r) ->
--                      let !x = fromIntegral (c - centerCol) * deltaCol
--                          !y = fromIntegral (r - centerRow) * deltaRow
--                          !r2 = x ^ 2 + y ^ 2
--                      in if (x == 0 && y == 0) || r2 > cutoff ^ 2
--                           then 0
--                           else cis $
--                                fromIntegral tf * atan2 y x +
--                                (pi * (fromIntegral rf * 0.5 * log r2) /
--                                halfLogPeriod - pi)
--              in ((rf, tf), vec))
--           [ (rf, tf)
--           | rf <-  (rangeFunc1 rhoFreqs rFreqs)
--           , tf <- L.reverse (rangeFunc1 phiFreqs thetaFreqs)
--           ]
--   in IA.array ((rfLB, tfLB), (rfUB, tfUB)) xs

{-# INLINE getHarmonics #-}
getHarmonics ::
     (VG.Vector vector e)
  => IA.Array (Int, Int) (vector e)
  -> Double
  -> Double
  -> Double
  -> Double
  -> vector e
getHarmonics harmonicsArray phiFreq rhoFreq thetaFreq rFreq =
  harmonicsArray IA.! (round (rhoFreq - rFreq), round $ (phiFreq - thetaFreq))

-- {-# INLINE computeFourierSeriesR2 #-}
-- computeFourierSeriesR2 ::
--      Int
--   -> Double
--   -> Int
--   -> Double
--   -> [Double]
--   -> [Double]
--   -> [Double]
--   -> [Double]
--   -> Double
--   -> Double
--   -> IA.Array (Int, Int) (VU.Vector (Complex Double))
--   -> R.Array U DIM4 (Complex Double)
--   -> [VU.Vector (Complex Double)]
-- computeFourierSeriesR2 numRows deltaRow numCols deltaCol phiFreqs rhoFreqs thetaFreqs rFreqs halfLogPeriod cutoff harmonicsArray arr =
--   let (Z :. (!numRFreq) :. (!numThetaFreq) :. _ :. _) = extent arr
--       !initVec = VU.replicate (VU.length (harmonicsArray IA.! (0, 0))) 0
--   in parMap
--        rdeepseq
--        (\((!r, !rFreq), (!theta, !thetaFreq)) ->
--           L.foldl'
--             (\(!vec) ((!rho, !rhoFreq), (!phi, !phiFreq)) ->
--                VU.zipWith
--                  (+)
--                  vec
--                  (VU.map
--                     (* (arr R.! (Z :. r :. theta :. rho :. phi)))
--                     (getHarmonics harmonicsArray phiFreq rhoFreq thetaFreq rFreq)))
--             initVec $
--           (,) <$> (L.zip [0 ..] rhoFreqs) <*> (L.zip [0 ..] phiFreqs)) $
--      (,) <$> (L.zip [0 ..] rFreqs) <*> (L.zip [0 ..] thetaFreqs)

-- {-# INLINE computeThetaRHarmonics #-}
computeThetaRHarmonics ::
     Int -> Int -> [Double] -> [Double] -> Double -> [[(Complex Double)]]
computeThetaRHarmonics !numOrientation !numScale !thetaFreqs !rFreqs !halfLogPeriod =
  let !deltaTheta = 2 * pi / fromIntegral numOrientation
      !deltaScale = 2 * halfLogPeriod / fromIntegral numScale
      !numRFreq = L.length rFreqs
      !numThetaFreq = L.length thetaFreqs
  in if numScale == 1
       then parMap
              rdeepseq
              (\o ->
                 L.map
                   (\freq -> cis $ freq * (fromIntegral o * deltaTheta))
                   thetaFreqs) $
            [0 .. numOrientation - 1]
       else parMap
              rdeepseq
              (\(!o, !s) ->
                 R.toList .
                 R.traverse2
                   (fromListUnboxed (Z :. numRFreq) rFreqs)
                   (fromListUnboxed (Z :. numThetaFreq) thetaFreqs)
                   (\_ _ -> (Z :. numRFreq :. numThetaFreq)) $ \fRFreq fThetaFreq idx@(Z :. rFreq :. thetaFreq) ->
                   -- exp $
                   -- (-0.5 * (fromIntegral s * deltaScale - halfLogPeriod)) :+
                   -- ((fRFreq (Z :. rFreq)) *
                   --  (fromIntegral s * deltaScale - halfLogPeriod) +
                   --  fThetaFreq (Z :. thetaFreq) * fromIntegral o * deltaTheta)
                   cis
                     ((fRFreq (Z :. rFreq)) *
                      (pi * (fromIntegral s * deltaScale - halfLogPeriod) /
                       halfLogPeriod) +
                      fThetaFreq (Z :. thetaFreq) * fromIntegral o * deltaTheta)
               ) $
            (,) <$> [0 .. numOrientation - 1] <*> [0 .. numScale - 1]

{-# INLINE computeFourierSeriesThetaR #-}
computeFourierSeriesThetaR ::
     (VG.Vector vector (Complex Double), NFData (vector (Complex Double)))
  => [[(Complex Double)]]
  -> [vector (Complex Double)]
  -> [vector (Complex Double)]
computeFourierSeriesThetaR !harmonics !vecs =
  let !initVec = VG.replicate (VG.length . L.head $ vecs) 0
  in parMap
       rdeepseq
       (L.foldl'
          (\vec0 (!vec1, !v) -> VG.zipWith (+) vec0 . VG.map (* v) $ vec1)
          initVec .
        L.zip vecs) $
     harmonics

plotThetaDimension ::
     (R.Source s Double)
  => FilePath
  -> String
  -> (Int, Int)
  -> R.Array s DIM3 Double
  -> IO ()
plotThetaDimension folderPath prefix (x', y') inputArr = do
  let centerCol = div cols 2
      centerRow = div rows 2
      x = x' + centerCol
      y = y' + centerRow
      (Z :. numOrientation :. cols :. rows) = extent inputArr
      centerArr =
        fromFunction (Z :. (1 :: Int) :. cols :. rows) $ \(Z :. _ :. i :. j) ->
          if i == x && j == y
            then 1 :: Double
            else 0
  plotImageRepa (folderPath </> printf "%s(%d,%d)_center.png" prefix x' y') .
    ImageRepa 8 . computeS $
    centerArr
  let ys' = R.toList . R.slice inputArr $ (Z :. R.All :. x :. y)
      !maxY = L.maximum ys'
      -- ys = L.map (/maxY) ys'
      ys = ys'
      deltaTheta = (360 ::Double) / fromIntegral numOrientation
      xs = [fromIntegral x * deltaTheta | x <- [0 .. numOrientation - 1]]
  toFile def (folderPath </> printf "%s(%d,%d)_theta.png" prefix x' y') $ do
    layout_title .= printf "%s(%d,%d)" prefix x' y'
    layout_x_axis . laxis_generate .= scaledAxis def (0, 359)
    layout_y_axis . laxis_generate .= scaledAxis def (0, maxY)
    plot (line "" [L.zip xs ys])

-- The codes below compute Fourier series in R2

computeFourierSeriesOfLogPolarHarmonicsArray ::
     (VG.Vector vector (Complex Double), NFData (vector (Complex Double)))
  => Double
  -> Double
  -> Int
  -> Int
  -> Int
  -> Int
  -> Int
  -> Double
  -> IA.Array (Int, Int) (vector (Complex Double))
computeFourierSeriesOfLogPolarHarmonicsArray !radius !delta !r2Freq !phiFreq !rhoFreq !thetaFreq !rFreq !halfLogPeriod =
  let ((!angularFreqLB, !angularFreqUB), angularFreqs) =
        freqRange phiFreq thetaFreq
      ((!radialFreqLB, !radialFreqUB), radialFreqs) =
        freqRange phiFreq thetaFreq
      !numPoints' = round $ (2 * radius + 1) / delta
      !numPoints =
        if odd numPoints'
          then numPoints'
          else numPoints' - 1
      !center = div numPoints 2
      !period = 2 * radius + 1
      !periodConstant = delta * 2 * pi / period
      !numR2Freq = 2 * r2Freq + 1
      !std = fromIntegral $ div r2Freq 2
      !simpsonWeights =
        toUnboxed . computeS $
        (computeWeightArrFromListOfShape [numPoints, numPoints] :: R.Array D DIM2 (Complex Double))
      !weightedRectangularHarmonics =
        parMap
          rdeepseq
          (\(freq1, freq2) ->
             VU.zipWith (*) simpsonWeights .
             toUnboxed . computeS . fromFunction (Z :. numPoints :. numPoints) $ \(Z :. i :. j) ->
               cis
                 (-periodConstant *
                   fromIntegral (freq1 * (i - center) + freq2 * (j - center))))
          [ (freq1, freq2)
          | freq1 <- [-r2Freq .. r2Freq]
          , freq2 <- [-r2Freq .. r2Freq]
          ]
      !cartesianGrid =
        toUnboxed . computeS . fromFunction (Z :. numPoints :. numPoints) $ \(Z :. c :. r) ->
          let !x = fromIntegral (c - center) * delta
              !y = fromIntegral (r - center) * delta
          in (x, y)
      !simpsonNorm = (delta / 3) ^ 2 :+ 0
      !xs =
        parMap
          rdeepseq
          (\(!radialFreq, !angularFreq) ->
             let !logPolarHarmonics =
                   VU.map (fourierMellin 0.5 angularFreq radialFreq) cartesianGrid
                 coefficients =
                   fromListUnboxed (Z :. numR2Freq :. numR2Freq) .
                   L.map (VU.sum . VU.zipWith (*) logPolarHarmonics) $
                   weightedRectangularHarmonics
                 !coefficientsGaussian =
                   VG.convert .
                   toUnboxed . computeS . R.traverse coefficients id $ \f idx@(Z :. xFreq :. yFreq) ->
                     simpsonNorm * (f idx) *
                     (gaussian2D
                        (fromIntegral $ xFreq - r2Freq)
                        (fromIntegral $ yFreq - r2Freq)
                        std :+
                      0)
             in ((radialFreq, angularFreq), coefficientsGaussian))
          [ (radialFreq, angularFreq)
          | radialFreq <- radialFreqs
          , angularFreq <- angularFreqs
          ]
  in IA.array ((radialFreqLB, angularFreqLB), (radialFreqUB, angularFreqUB)) xs

computeRectangularInverseHarmonics ::
     (VG.Vector vector (Complex Double), NFData (vector (Complex Double)))
  => Int
  -> Int
  -> Double
  -> Double
  -> Int
  -> [vector (Complex Double)]
computeRectangularInverseHarmonics !numRows !numCols !delta !radius !maxFreq =
  let !period = 2 * radius + 1
      !periodConstant = delta * 2 * pi / period
      !centerRow = div numRows 2
      !centerCol = div numCols 2
  in parMap
       rdeepseq
       (\(col, row) ->
          VG.fromList
            [ cis
              (periodConstant *
               fromIntegral
                 (freq1 * (row - centerRow) + freq2 * (col - centerCol)))
            | freq1 <- [-maxFreq .. maxFreq]
            , freq2 <- [-maxFreq .. maxFreq]
            ])
       [(col, row) | col <- [0 .. numCols - 1], row <- [0 .. numRows - 1]]

-- Utilities

{-# INLINE freqRange #-}
freqRange :: Int -> Int -> ((Int, Int), [Int])
freqRange n m =
  let !x = n + m
  in ((-x, x), [-x .. x])


{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE FlexibleContexts #-}

module Data.AdaptiveGeneticAlgorithm
  ( adaptiveGeneticAlgorithm )
  where

import Control.Monad.State.Strict
import Data.AdaptiveCoordinateDescent.Internal
import Data.List ( sortBy )
import Data.Ord
import Data.Foldable
import Data.Traversable
import qualified Data.Vector as V
import Numeric.LinearAlgebra hiding ( toList )
import Pipes hiding ( for )
import qualified Pipes.Prelude as P
import System.Random.MWC

data GeneticAlgorithmState f = GeneticAlgorithmState
  { principalComponents :: (Vector Double, [Vector Double])
  , populationPool :: [f Double]
  , rng :: GenIO }

generateCauchy :: GenIO -> Double -> Double -> IO Double
generateCauchy rng mean scale = do
  r <- uniformR (0, 1) rng
  return $ (tan (pi * (r - 0.5)) * scale) + mean

principalComponentsFromPopulation :: Foldable f => [f Double] -> (Vector Double, [Vector Double])
principalComponentsFromPopulation population =
  let mat = normalize $ (length population><num_coordinates)
                        (mconcat $ fmap toList population)
      (evec, mat2) = rightSV mat

   in (evec, toColumns mat2)
 where
  num_coordinates = length $ toList $ head population

sortByFitness :: Monad m
              => (f Double -> m Double)
              -> [f Double]
              -> m [(f Double, Double)]
sortByFitness evaluate items = do
  evaluated_items <- for items $ \item -> (,) item <$> evaluate item
  return $ sortBy (comparing snd) evaluated_items

adaptiveGeneticAlgorithm :: (MonadIO m, Traversable f)
                         => (f Double -> m Double) -- ^ Cost function. Smaller cost = better model.
                         -> f Double
                         -> Int
                         -> Producer (f Double, Double) m void
adaptiveGeneticAlgorithm evaluate initial_params how_many_to_keep = do
  rrng <- liftIO createSystemRandom

  initial_items <- liftIO $ P.toListM (makePerturbed initial_params rrng (how_many_to_keep*2))

  best_initial_items' <- lift $ fmap (take how_many_to_keep) $ sortByFitness evaluate initial_items

  yield (head best_initial_items')

  let best_initial_items = fmap fst best_initial_items'
      pcomponents = principalComponentsFromPopulation best_initial_items

  evalStateT piper (GeneticAlgorithmState { principalComponents = pcomponents
                                          , rng = rrng
                                          , populationPool = best_initial_items })
 where
  piper = do
    current_pool <- V.fromList . populationPool <$> get
    (pcomponent_scale_vec, pcomponents) <- principalComponents <$> get
    for_ (zip [0..] pcomponents) $ \(index, pcomponent) -> do
      let pcomponent_scale = pcomponent_scale_vec `atIndex` index
      parent_index <- randomI (0, V.length current_pool-1)
      let parent = current_pool V.! parent_index
      scalef <- randomCauchy
      let final_scale = scalef * pcomponent_scale
      addCandidate (plus parent (fromVector (final_scale `scale` pcomponent) parent))

    current_candidates <- populationPool <$> get
    best_candidates <- lift $ lift $ sortByFitness evaluate current_candidates

    let truncated_candidates = take how_many_to_keep best_candidates
        pcomponents = principalComponentsFromPopulation $ fmap fst truncated_candidates

    modify $ \old -> old { principalComponents = pcomponents
                         , populationPool = fmap fst truncated_candidates }

    lift $ yield (head best_candidates)

    piper

  randomCauchy = do
    rrng <- rng <$> get
    liftIO $ generateCauchy rrng 0 1

  randomI range = do
    rrng <- rng <$> get
    liftIO $ uniformR range rrng

  addCandidate item = do
    st <- get
    put $ st { populationPool = item:populationPool st }


{-# LANGUAGE GADTs #-}
module Math.IRT.Model.FourPLM
  ( FourPLM (..)
  ) where

import Numeric.AD (Mode, Scalar, auto)
import Numeric.AD.Mode.Forward.Double
import qualified Numeric.AD.Mode.Tower as T

import Statistics.Distribution

import Math.IRT.Internal.Distribution
import Math.IRT.Internal.LogLikelihood
import Math.IRT.Model.Generic

data FourPLM = FourPLM { discrimination :: !Double
                       , difficulty     :: !Double
                       , pseudoGuessing :: !Double
                       , asymptote      :: !Double
                       } deriving (Show)

instance Distribution FourPLM where
    cumulative = cumulative4PL

instance ContDistr FourPLM where
    density  x = diff (cumulative4PL x)
    quantile _ = error "This shouldn't be needed"

instance DensityDeriv FourPLM where
    densityDeriv x = (!! 2) . T.diffs (cumulative4PL x)

instance GenericModel FourPLM where
    fromRasch         b = FourPLM 1.0 b 0.0 1.0
    fromOnePLM        b = FourPLM 1.7 b 0.0 1.0
    fromTwoPLM      a b = FourPLM   a b 0.0 1.0
    fromThreePLM  a b c = FourPLM   a b   c 1.0
    fromFourPLM         = FourPLM

instance LogLikelihood FourPLM where
    logLikelihood = logLikeFunc cumulative4PL

cumulative4PL :: (Mode a, Floating a, Scalar a ~ Double) =>
       FourPLM -- ^The IRT parameters
    -> a       -- ^Theta
    -> a
cumulative4PL params theta =
    let (a, b, c, d) = makeParamAutos params in
    c + (              (d - c)
         -----------------------------------
         / (1 + exp ((-a) * (theta - b))))
    where makeParamAutos (FourPLM sa sb sc sd) =
              let a = auto sa
                  b = auto sb
                  c = auto sc
                  d = auto sd
              in (a, b, c, d)


{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE OverloadedStrings #-}

module Utils where
import Control.Monad.Fix
import Control.Monad
import System.IO
import Signal
import Circuit
import AmpOp
    ( SignalState(SignalTwoNodes),
      AmpOp(openLoopGain),
      Impedance,
      Output,
      Input,
      dc100,
      sen,
      ground,
      lm741,
      r1,
      r2,
      eps )
import Data.List
import Data.Complex
import Simulation

type Iteration = Int

makePlotIterations = do
    outh <- openFile "outputIterations.txt" WriteMode
    mapM_ (\(i, o) -> hPutStrLn outh (show i ++ " " ++ show o)) (ampOpInvertingList' lm741 12 r1 r2)
    hClose outh

ampOpInvertingList' :: AmpOp -> Input -> Impedance -> Impedance -> [(Iteration, Output)]
ampOpInvertingList' model input r1 r2 = fix calculate (100, 0, [], 1)
    where calculate f (vXOld, vOutOld, l, it) = if abs(realPart vOutOld - realPart vOut) <= 0.000001 && abs(realPart vXOld - realPart vX) <= 0.000001 then l ++ [(it, vOut)] else f (vOut, vX, l ++ [(it, vOut)], it + 1)
            where vX = (- vOutOld) / openLoopGain model `at` 0
                  vOut = (vXOld * (r1 + r2) `at` 0 - input * r2 `at` 0) / r1 `at` 0

checkDataIterations :: IO ()
checkDataIterations = print (ampOpInvertingList' lm741 12 r1 r2)

ampOpInvertingList :: AmpOp -> Impedance -> Impedance -> SignalState -> Circuit Input [(Iteration, SignalState)]
ampOpInvertingList model r1 r2 (SignalTwoNodes (vXInit, vOutInit)) =
    ($ (vXInit, vOutInit, [], 1))<$> mfix
    (\f -> Circuit $ \vIn -> Signal $ \time (vXOld, vOutOld, l, it) ->
        let vX = (- vOutOld) / openLoopGain model
            vOut = ((vXOld * (r1 + r2)) - (vIn * r2)) / r1
        in if abs((realPart <$> vOutOld) - (realPart <$> vOut)) <= eps && abs((realPart <$> vXOld) - (realPart <$> vX)) <= eps
            then l ++ [(it, SignalTwoNodes (vX, vOut))]
            else f (vX, vOut, l ++ [(it, SignalTwoNodes (vX, vOut))], it + 1))

makeSamplesInput :: Signal Input -> [Metric] -> [(Metric, Input)]
makeSamplesInput input samples = do time <- samples
                                    let result = input `at` time
                                    return (time, result)

makeSamplesOutput :: SignalState -> (SignalState -> Circuit Input [(Iteration, SignalState)]) -> Signal Input -> [Metric] -> [[(Metric, Iteration, Output)]]
makeSamplesOutput _ _ _ [] = []
makeSamplesOutput initialState circuit input (t:ts) = map (\(i, SignalTwoNodes (_, o)) -> (t, i, o `at` t)) results : rest
    where results = circuit initialState `simulate` input `at` t
          rest = makeSamplesOutput (snd $ last results) circuit input ts

ampOpInvertingList2 :: AmpOp -> Impedance -> Impedance -> Circuit Input [(Iteration, Output)]
ampOpInvertingList2 model r1 r2 =
    ($ (dc100, ground, [], 1))<$> mfix
    (\f -> Circuit $ \vIn -> Signal $ \time (vXOld, vOutOld, l, it) ->
        let vX = (- vOutOld) / openLoopGain model
            vOut = ((vXOld * (r1 + r2)) - (vIn * r2)) / r1
        in if abs((realPart <$> vOutOld) - (realPart <$> vOut)) <= eps && abs((realPart <$> vXOld) - (realPart <$> vX)) <= eps
            then l ++ [(it, vOut `at` time)]
            else f (vX, vOut, l ++ [(it, vOut `at` time)], it + 1))

makeSamplesOutput2 :: Signal Input -> [Metric] -> [[(Metric, Iteration, Output)]]
makeSamplesOutput2 input samples = do time <- samples
                                      let results = (ampOpInvertingList2 lm741 r1 r2 `simulate` input) `at` time
                                      mapM (\(i, o) -> [(time, i, o)]) results

testCircuit = ampOpInvertingList lm741 r1 r2

makePlotMetricOutput = do
    outh <- openFile "outputMetric.txt" WriteMode
    mapM_ (\(t, i, o) -> hPutStrLn outh (show t ++ " " ++ show i ++ " " ++ show o)) (concat $ makeSamplesOutput iSignalState testCircuit sen (map (:+ 0) [0.1,0.2..10.0]))
    hClose outh

makePlotMetricInput = do
    outh <- openFile "inputMetric.txt" WriteMode
    mapM_ (\(t, o) -> hPutStrLn outh (show t ++ " " ++ show o)) (makeSamplesInput sen (map (:+ 0) [0.1,0.2..10.0]))
    hClose outh


{-# LANGUAGE TemplateHaskell #-}

module Shapes where

import Control.Lens hiding (element)
import Graphics.Gloss.Data.Picture
import Data.Complex

data ShapeType = Kite | Dart deriving (Show)

data Shape = Shape {_shape :: ShapeType
                   ,_translate :: Complex Double
                   ,_rotate :: Complex Double
} deriving (Show)

newtype Vertices = Vertices {_nodes :: [Complex Double]} deriving (Show)

makeLenses ''Shape
makeLenses ''Vertices

standardKite :: Shape
standardKite = Shape {_shape = Kite, _translate = 0, _rotate = 1}

standardDart :: Shape
standardDart = Shape {_shape = Dart, _translate = 0, _rotate = 1}

z :: Complex Double
z = cis (pi/5)

zPower :: Int -> Complex Double
zPower n = cis((pi/5) * fromIntegral (n `mod` 10))

longEdge :: Vertices
longEdge = Vertices {_nodes = [0,(0.5 :+ (-0.1))*z,z]}

shortEdge :: Vertices
shortEdge = Vertices {_nodes = [0,(0.5 :+ 0.1)*(1-z),1-z]}

rotranslate :: Complex Double -> Complex Double -> Vertices -> Vertices
rotranslate z w = over nodes (map (\u -> u*z+w))

reverseVertices :: Vertices -> Vertices
reverseVertices = over nodes reverse

concatVertices :: [Vertices] -> Vertices
concatVertices p = Vertices {_nodes = concatMap (tail._nodes) p}

shapeToVertices :: Shape -> Vertices
shapeToVertices (Shape Kite z x) = rotranslate x z kiteVertices
shapeToVertices (Shape Dart z x) = rotranslate x z dartVertices

complexToPoint :: Complex Double -> Point
complexToPoint z = (realToFrac x, realToFrac y)
  where
    x = realPart z
    y = imagPart z

verticesToPicture :: Vertices -> Picture
verticesToPicture = lineLoop . map complexToPoint . view nodes

shapeToPicture :: Shape -> Picture
shapeToPicture = verticesToPicture . shapeToVertices

refine :: Shape -> [Shape]
refine s@(Shape t w z) =
  let phi = (sqrt 5 - 1)/2 :+ 0
      tr t' u v = Shape t' (z*phi*v+w) (z*u*phi)
  in case t of
  Kite -> [tr Dart 1 0
          ,tr Kite (zPower 3) (1 - zPower 3)
          ,tr Kite (zPower 7) (1 + zPower 2)
          ]
  Dart -> [tr Kite (zPower 9) 0
          ,tr Dart (zPower 6) (1+phi)
          ]

tilings :: [[Shape]]
tilings = [standardKite] : map (concatMap refine) tilings

tilingPics :: [Picture]
tilingPics = map (pictures . map shapeToPicture) tilings

kiteVertices :: Vertices
kiteVertices = concatVertices
  [longEdge
  ,rotranslate 1 z shortEdge
  ,reverseVertices $ rotranslate (zPower 4) (zPower 9) shortEdge
  ,reverseVertices $ rotranslate (zPower 8) 0 longEdge
  ]

dartVertices :: Vertices
dartVertices = concatVertices
  [reverseVertices $ rotranslate (zPower 4) 1 longEdge
  ,reverseVertices $ rotranslate (zPower 3) x shortEdge
  ,rotranslate (zPower 9) x shortEdge
  ,rotranslate (zPower 2) (zPower 8) longEdge
  ]
  where
    x = (1 + (zPower 4) - (zPower 3))

printVertices :: Vertices -> IO ()
printVertices = mapM_ print . _nodes


{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE UndecidableInstances #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE FunctionalDependencies #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE TemplateHaskell #-}
{-# LANGUAGE ConstraintKinds #-}
{-# LANGUAGE InstanceSigs #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE PolyKinds #-}
{-# LANGUAGE StandaloneDeriving #-}

module HLearn.Models.Distributions.Multivariate.MultiNormal
    ( MultiNormal (..)
    )
    where

-- import qualified Control.ConstraintKinds as CK
import qualified Data.Foldable as F
import qualified Data.Vector as V
import qualified Data.Vector.Unboxed as VU
import Data.Vector.Unboxed.Deriving

import Control.DeepSeq
import GHC.ST
import GHC.TypeLits

import Foreign.Storable
import Numeric.LinearAlgebra hiding ((<>))
import qualified Numeric.LinearAlgebra as LA

import HLearn.Algebra
import HLearn.Models.Distributions.Common

-------------------------------------------------------------------------------
-- data types

data MultiNormalVec (n::Nat) prob = MultiNormalVec
    { q0 :: !prob
    , q1 :: !(VU.Vector prob)
    , q2 :: !(V.Vector (VU.Vector prob))
    }
    deriving (Read,Show,Eq,Ord)

instance NFData (MultiNormalVec n prob) where
    rnf mn = seq mn ()

newtype MultiNormal prob (xs::[*]) = MultiNormal (MultiNormalVec (Length xs) prob)
    deriving (Read,Show,Eq,Ord,NFData)

deriving instance (Monoid  (MultiNormalVec (Length xs) prob)) => Monoid (MultiNormal prob xs)
deriving instance (Abelian (MultiNormalVec (Length xs) prob)) => Abelian (MultiNormal prob xs)
deriving instance (Group   (MultiNormalVec (Length xs) prob)) => Group (MultiNormal prob xs)
deriving instance (Module  (MultiNormalVec (Length xs) prob)) => Module (MultiNormal prob xs)

-------------------------------------------------------------------------------
-- algebra

instance (Num prob, VU.Unbox prob, SingI n) => Abelian (MultiNormalVec n prob)
instance (Num prob, VU.Unbox prob, SingI n) => Monoid (MultiNormalVec n prob) where
    mempty = MultiNormalVec
        { q0 = 0
        , q1 = VU.replicate n 0
        , q2 = V.replicate n (VU.replicate n 0)
        }
        where
            n = fromIntegral $ fromSing (sing :: Sing n)
    mn1 `mappend` mn2 = MultiNormalVec
        { q0 = (q0 mn1) + (q0 mn2)
        , q1 = VU.zipWith (+) (q1 mn1) (q1 mn2)
        , q2 = V.zipWith (VU.zipWith (+)) (q2 mn1) (q2 mn2)
        }

instance (Num prob, VU.Unbox prob, SingI n) => Group (MultiNormalVec n prob) where
    inverse mn = MultiNormalVec
        { q0 = negate $ q0 mn
        , q1 = VU.map negate (q1 mn)
        , q2 = V.map (VU.map negate) (q2 mn)
        }

instance (Num prob) => HasRing (MultiNormalVec n prob) where
    type Ring (MultiNormalVec n prob) = prob

instance (Num prob, VU.Unbox prob, SingI n) => Module (MultiNormalVec n prob) where
    r .* mn = MultiNormalVec
        { q0 = r * q0 mn
        , q1 = VU.map (r*) (q1 mn)
        , q2 = V.map (VU.map (r*)) (q2 mn)
        }

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

instance (Num prob) => HasRing (MultiNormal prob xs) where
    type Ring (MultiNormal prob xs) = prob

-------------------------------------------------------------------------------
-- training

instance (SingI n, Num prob, VU.Unbox prob) => HomTrainer (MultiNormalVec n prob) where
    type Datapoint (MultiNormalVec n prob) = (VU.Vector prob)
    train1dp dp = MultiNormalVec
        { q0 = 1
        , q1 = dp
        , q2 = V.generate n (\i -> VU.generate n (\j -> (dp VU.! i)*(dp VU.! j)))
        }
        where
            n = fromIntegral $ fromSing (sing :: Sing n)

instance
    ( SingI (Length xs)
    , Num prob
    , VU.Unbox prob
    , HList2List (Datapoint (MultiNormal prob xs)) prob
    ) => HomTrainer (MultiNormal prob xs)
        where
    type Datapoint (MultiNormal prob xs) = HList xs
    train1dp dp = MultiNormal $ train1dp $ VU.fromList $ hlist2list dp

instance (Num prob) => NumDP (MultiNormal prob xs) where
    numdp (MultiNormal mn) = q0 mn

-------------------------------------------------------------------------------
-- distribution

class (Probabilistic dist) => Covariance dist where
    covar :: dist -> Matrix (Probability dist)

instance
    ( VU.Unbox prob
    , SingI k
    , Num prob
    ) => Probabilistic (MultiNormalVec k prob)
        where
    type Probability (MultiNormalVec k prob) = prob

instance
    ( VU.Unbox prob
    , SingI k
    , Fractional prob
    , Enum prob
    , Storable prob
    ) => Covariance (MultiNormalVec k prob)
        where
    covar mn = (k><k) $
            [ (1/(n-1))*( mij - ((q1 mn VU.! j)*((q1 mn) VU.! i))/n )
            | ((i,j),mij) <- assocs
            ]
        where
            assocs = zip [(i,j) | i<-[0..k-1],j<-[0..k-1]] (concat $ V.toList $ fmap VU.toList $ q2 mn)
            mean i = ((q1 mn) VU.! i)/n
            n = q0 mn
            k = fromIntegral $ fromSing (sing :: Sing k)

instance
    ( HList2List (HList dpL) prob
    , VU.Unbox prob
    , Num prob
    , SingI (FromNat1 (Length1 dpL))
    ) => Probabilistic (MultiNormal prob dpL)
        where
    type Probability (MultiNormal prob dpL) = prob

instance
    ( HList2List (HList dpL) prob
    , VU.Unbox prob
    , Floating prob
    , Field prob
    , Enum prob
    , SingI (FromNat1 (Length1 dpL))
--     , Covariance (MultiNormal dpL prob)
    , Storable prob
    ) => PDF (MultiNormal prob dpL)
        where
    pdf (MultiNormal dist) dpL = 1/(sqrt $ (2*pi)^(k)*(det sigma))*(exp $ (-1/2)*(top) )
        where
            top=minElement $ ((trans $ x `sub` mu) LA.<> (inv sigma) LA.<> (x `sub` mu))

            k = fromIntegral $ fromSing (sing :: Sing (Length dpL)) :: Int
            x = k><1 $ hlist2list dpL
            n = q0 dist
            sigma = covar dist
            mu = k><1 $ map (/n) $ VU.toList (q1 $ dist)
--             covarM = (k><k) $ elems $ covar dist

-------------------------------------------------------------------------------
-- tests

-- ds = map (listArray (0,2))
--     [[1,2,4]
--     ,[2,5,6]
--     ,[3,1,1]
--     ]

-- test = train ds :: MultiNormalArray Double 3

ds2 = map VU.fromList
    [[1,2,4]
    ,[2,5,6]
    ,[3,1,1]
    ]

test2 = train ds2 :: MultiNormalVec 3 Double

ds =
    [ 1:::2:::3:::HNil
    , 2:::5:::6:::HNil
    , 3:::1:::1:::HNil
    , 3:::2:::1:::HNil
    ]
test = train ds :: MultiNormal Double '[Double,Double,Double]


{-# LANGUAGE AllowAmbiguousTypes                  #-}
{-# LANGUAGE DataKinds                            #-}
{-# LANGUAGE PartialTypeSignatures                #-}
{-# LANGUAGE PolyKinds                            #-}
{-# LANGUAGE ScopedTypeVariables                  #-}
{-# LANGUAGE TupleSections                        #-}
{-# LANGUAGE TypeApplications                     #-}
{-# LANGUAGE TypeOperators                        #-}
{-# OPTIONS_GHC -fno-warn-partial-type-signatures #-}

import           Backprop.Learn
import           Control.DeepSeq
import           Control.Exception
import           Control.Monad
import           Control.Monad.IO.Class
import           Control.Monad.Trans.Class
import           Control.Monad.Trans.State
import           Data.Char
import           Data.Conduit
import           Data.Default
import           Data.List
import           Data.Proxy
import           Data.Time
import           Data.Type.Tuple
import           GHC.TypeNats
import           Numeric.LinearAlgebra.Static.Backprop
import           Numeric.Opto
import           System.Environment
import           Text.Printf
import qualified Conduit                               as C
import qualified Data.Conduit.Combinators              as C
import qualified Data.Set                              as S
import qualified Data.Text                             as T
import qualified Data.Vector.Sized                     as SV
import qualified Data.Vector.Storable                  as VS
import qualified Numeric.LinearAlgebra.Static          as H
import qualified System.Random.MWC                     as MWC

encoder :: forall n e. (KnownNat n, KnownNat e) => Model _ _ (R n) (R e)
encoder = fca logistic

decoder :: forall n e. (KnownNat n, KnownNat e) => Model _ _ (R e) (R n)
decoder = fca softMax

word2vec :: forall n e. (KnownNat n, KnownNat e) => Model _ _ (R n) (R n)
word2vec = decoder @n @e <~ encoder @n @e

oneHotWord
    :: KnownNat n
    => S.Set String
    -> String
    -> Maybe (R n)
oneHotWord ws = fmap (oneHotR . fromIntegral) . (`S.lookupIndex` ws)

makeCBOW :: KnownNat w => SV.Vector w (R n) -> (R n, R n)
makeCBOW v = (SV.index v mid, SV.sum (v SV.// [(mid, 0)]))
  where
    mid = maxBound `div` 2

main :: IO ()
main = MWC.withSystemRandom @IO $ \g -> do
    sourceFile:logFile:testFile:_  <- getArgs
    wordSet <- S.fromList . tokenize <$> readFile sourceFile
    SomeNat (Proxy :: Proxy n) <- pure $ someNatVal (fromIntegral (S.size wordSet))

    printf "%d unique words found.\n" (natVal (Proxy @n))

    let model = word2vec @n @100
        enc   = encoder  @n @100
    p0 <- initParamNormal model 0.2 g

    let report n b = do
          liftIO $ printf "(Batch %d)\n" (b :: Int)
          t0 <- liftIO getCurrentTime
          C.drop (n - 1)
          mp <- mapM (liftIO . evaluate . force) =<< await
          t1 <- liftIO getCurrentTime
          case mp of
            Nothing -> liftIO $ putStrLn "Done!"
            Just p@(_ :# pEnc) -> do
              chnk <- lift . state $ (,[])
              liftIO $ do
                printf "Trained on %d points in %s.\n"
                  (length chnk)
                  (show (t1 `diffUTCTime` t0))
                let trainScore = testModelAll maxIxTest model (TJust p) chnk
                printf "Training error:   %.3f%%\n" ((1 - trainScore) * 100)

                testWords <- tokenize <$> readFile testFile
                let tests = flip map testWords $ \w ->
                       let v = maybe 0 (runModelStateless enc (TJust pEnc))
                                $ oneHotWord wordSet w
                       in  intercalate "," $ map (printf "%0.4f") (VS.toList (H.extract v))

                writeFile logFile $ unlines tests

              report n (b + 1)

    C.runResourceT . flip evalStateT []
        . runConduit
        $ forever ( C.sourceFile sourceFile
                 .| C.decodeUtf8
                 .| C.concatMap (tokenize . T.unpack)
                 .| C.concatMap (oneHotWord wordSet)
                 .| C.slidingWindow 5
                 .| C.concatMap (SV.fromList @5)
                 .| C.map makeCBOW
                  )
       .| skipSampling 0.02 g
       .| C.iterM (modify . (:))
       .| optoConduit def p0
                (adam def (modelGradStoch crossEntropy noReg model g))
       .| report 500 0
       .| C.sinkNull

tokenize :: String -> [String]
tokenize = words . map (whitePunc . toLower)
  where
    whitePunc c
      | isPunctuation c = ' '
      | otherwise       = c


{-# LANGUAGE DataKinds
           , GADTs
           , TypeOperators #-}
module Aqua
    (
      module Aqua
    ) where

import Control.Arrow
import Data.Complex
import GHC.TypeLits

import QBasis
import Vect

type QReg n = Vect C (QBasis n)
type QCirc n = Vect C (QBasis n) -> Vect C (QBasis n)

runCircuit :: KnownNat n => QCirc n -> QReg n
runCircuit = ($ unit)

tensor :: QReg m -> QReg n -> QReg (m + n)
tensor vec1 vec2 = vec1 `te` vec2

(⊗) :: KnownNat m => QCirc m -> QCirc n -> QCirc (m + n)
c1 ⊗ c2 = c1 `tf` c2

unit :: KnownNat n => QReg n
unit = V [(unit', 1)]

idGate :: QCirc n
idGate = id

notGate :: QCirc 1
notGate = linear n
    where n :: QBasis 1 -> QReg 1
          n (Cons False Nil) = V [(Cons True Nil, 1)]
          n (Cons True Nil) = V [(Cons False Nil, 1)]

hGate :: QCirc 1
hGate = linear h
    where h :: QBasis 1 -> QReg 1
          h (Cons False Nil) = V [(Cons False Nil, 1 / sqrt 2)
                                 ,(Cons True Nil, 1 / sqrt 2)]
          h (Cons True Nil) = V [(Cons False Nil, 1 / sqrt 2)
                                ,(Cons True Nil, -1 / sqrt 2)]

rotGate :: Double -> QCirc 1
rotGate t = linear r
    where r :: QBasis 1 -> QReg 1
          r (Cons False Nil) = V [(Cons False Nil, 1)]
          r (Cons True Nil) = V [(Cons True Nil, cis t)]

cnotGate :: QCirc 2
cnotGate = linear n
    where n :: QBasis 2 -> QReg 2
          n (Cons False (Cons False Nil)) = V [(Cons False (Cons False Nil), 1)]
          n (Cons False (Cons True Nil)) = V [(Cons False (Cons True Nil), 1)]
          n (Cons True (Cons False Nil)) = V [(Cons True (Cons True Nil), 1)]
          n (Cons True (Cons True Nil)) = V [(Cons True (Cons False Nil), 1)]

swapGate :: QCirc 2
swapGate = linear s
    where s :: QBasis 2 -> QReg 2
          s (Cons False (Cons False Nil)) = V [(Cons False (Cons False Nil), 1)]
          s (Cons False (Cons True Nil)) = V [(Cons True (Cons False Nil), 1)]
          s (Cons True (Cons False Nil)) = V [(Cons False (Cons True Nil), 1)]
          s (Cons True (Cons True Nil)) = V [(Cons True (Cons True Nil), 1)]


{-# LANGUAGE FlexibleContexts #-}

module Blash where

import qualified Data.Vector.Generic as V
import Data.Vector.Generic ((!))
import Data.Complex
import BlashImpl (fZERO, stride, isample, map_reduce, Size, Inc, Index, copyM, axpyM)

copy :: (V.Vector v a)
         => Size
         -> v a
         -> Inc
         -> v a
         -> Inc
         -> v a
copy n dx incx dy incy = V.modify modifier dy
  where
    modifier ys = copyM n dx incx ys incy

axpy :: (Ord a, Num a, V.Vector v a)
         => Size
         -> a
         -> v a
         -> Inc
         -> v a
         -> Inc
         -> v a
axpy n da dx incx dy incy = V.modify modifier dy
  where
    modifier ys = axpyM n da dx incx ys incy

-- /* dot product dx dot dy. */
dot :: (Num a, V.Vector v a)
        => Size
        -> v a
        -> Inc
        -> v a
        -> Inc
        -> a
dot n _ _ _ _ | n <= 0 = fZERO
dot n dx incx dy incy = sum prods
  where
    prods = flip map [0..(n-1)] f
    f i = (dx ! ix i) * (dy ! iy i)
    ix = stride n incx
    iy = stride n incy

-- /* dot product dx dot dy where dx & dy are complex. */
dotc :: (RealFloat a, V.Vector v (Complex a))
        => Size
        -> v (Complex a)
        -> Inc
        -> v (Complex a)
        -> Inc
        -> Complex a
dotc n _ _ _ _ | n <= 0 = fZERO :+ fZERO
dotc n dx incx dy incy = sum prods
  where
    prods = flip map [0..(n-1)] f
    f i = (conjugate (dx ! ix i)) * (dy ! iy i)
    ix = stride n incx
    iy = stride n incy

-- /* compute the L2 norm of array DX of length N, stride INCX */
-- TODO should really be (Num a) - or split between Complex and Real
nrm2 :: (Ord a, Floating a, V.Vector v a)
        => Size
        -> v a
        -> Inc
        -> a
nrm2 n _  incx | n <= 0 || incx <= 0 = fZERO
nrm2 n dx _ | n == 1 = abs $ dx ! 0
nrm2 n dx incx =
  let dx' = isample n dx incx
      xmax = V.maximum $ V.map abs dx'
  in
    case xmax == fZERO of
    True -> fZERO
    False ->
      let scale = 1.0 / xmax
          scaled_dx' = V.map (* scale) dx'
          n' = V.length scaled_dx'
          scaled_drnm2 = sqrt $ dot n' scaled_dx' 1 scaled_dx' 1
      in
        xmax * scaled_drnm2

-- DASUM takes the sum of the absolute values.
asum :: (Num a, V.Vector v a)
        => Size
        -> v a
        -> Inc
        -> a
asum = map_reduce abs V.sum

-- DAMAX takes the max of the absolute values.
iamax :: (Ord a, Num a, V.Vector v a)
        => Size
        -> v a
        -> Inc
        -> Index
iamax = map_reduce abs V.maxIndex


{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE NoImplicitPrelude #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE TemplateHaskell #-}
{-# LANGUAGE DeriveGeneric #-}
{-# OPTIONS_GHC -funbox-strict-fields #-}

module Warden.Anomaly.Data (
    FeatureVector(..)
  , Features(..)
  , featuresFromList
  , Dimensionality(..)
  , dimensionality
  , fromDataRow
  , toMatrix
  ) where

import           Data.ByteString (ByteString)
import           Data.List.NonEmpty (NonEmpty(..))
import qualified Data.List.NonEmpty as NE
import qualified Data.Vector as V
import qualified Data.Vector.Unboxed as VU

import           Numeric.LinearAlgebra (R)
import qualified Numeric.LinearAlgebra as HM

import           P

import           Warden.Data.Numeric
import           Warden.Data.Row
import           Warden.Row (parseProbablyNumeric)

newtype FeatureVector =
  FeatureVector {
    unFeatureVector :: VU.Vector Double
  } deriving (Eq, Show)

numFeatures :: FeatureVector -> Dimensionality
numFeatures = Dimensionality . VU.length . unFeatureVector

-- | Contains one or more feature vectors. All feature vectors have the same
-- dimensionality.
newtype Features =
  Features {
    unFeatures :: V.Vector FeatureVector
  } deriving (Eq, Show)

newtype Dimensionality =
  Dimensionality {
    unDimensionality :: Int
  } deriving (Eq, Show)

dimensionality :: Features -> Dimensionality
dimensionality (Features fvs) =
  numFeatures $ V.head fvs

featuresFromList :: NonEmpty FeatureVector -> Maybe' Features
featuresFromList (v:|vs) =
  let
    dim = numFeatures v
    bads = filter (\x -> numFeatures x /= dim) vs
  in
  case bads of
    [] ->
      Just' . Features . V.fromList . NE.toList $ v :| vs
    _ ->
      Nothing'

getNumericField :: V.Vector ByteString -> Int -> Maybe' Double
getNumericField fs ix = do
  f <- strictMaybe $ fs V.!? ix
  case parseProbablyNumeric f of
    NoNumericField ->
      Nothing'
    MNumericField (NumericField x) ->
      pure x

fromDataRow :: [Int] -> Row -> Maybe' FeatureVector
fromDataRow nixes (SVFields fs) = do
  nums <- mapM (getNumericField fs) nixes
  pure . FeatureVector $ VU.fromList nums
fromDataRow _ (RowFailure _) =
  Nothing'

toMatrix :: [FeatureVector] -> HM.Matrix R
toMatrix =
  HM.fromRows . fmap (HM.vector . VU.toList . unFeatureVector)


{-# LANGUAGE ScopedTypeVariables, TypeFamilies #-}

module Fuml.Supervised where

import qualified Fuml.Base.LinearRegression as LR
import qualified Fuml.Base.Logistic as LG
import Fuml.Core
import Numeric.LinearAlgebra
import Control.Monad.Identity

ols :: Supervisor Identity Double (Vector Double) Double
ols = Supervisor $ \_ theData ->
              let beta = LR.ols theData
              in return $ linPredict beta

olsWeighted :: Vector Double -> Supervisor Identity Double (Vector Double) Double
olsWeighted wvs
   = Supervisor $ \_ theData ->
              let beta = LR.wols wvs theData
              in return $ linPredict beta

ridge :: Double -> Supervisor Identity Double (Vector Double) Double
ridge gammac = Supervisor $ \_ xys ->
  let npars = size $ fst $ head xys
      gamma = diag $ konst gammac npars
      beta = LR.ridge gamma xys
  in return $ linPredict beta

linPredict :: Vector Double -> Predict (Vector Double) Double
linPredict beta = Predict beta (`dot` beta)


module PlotLocalEigenVector where

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           FokkerPlanck.DomainChange
import           FokkerPlanck.MonteCarlo
import           FokkerPlanck.Pinwheel
import           Image.IO
import           STC
import           System.Directory
import           System.Environment
import           System.FilePath
import           Text.Printf
import           Types
import           Utils.Array
import           Utils.Parallel

main = do
  args@(numPointStr:numOrientationStr:numScaleStr:thetaSigmaStr:scaleSigmaStr:maxScaleStr:taoStr:numTrailStr:maxTrailStr:thetaFreqsStr:scaleFreqsStr:histFileName:numThreadStr:_) <-
    getArgs
  print args
  let numPoint = read numPointStr :: Int
      numOrientation = read numOrientationStr :: Int
      numScale = read numScaleStr :: Int
      thetaSigma = read thetaSigmaStr :: Double
      scaleSigma = read scaleSigmaStr :: Double
      maxScale = read maxScaleStr :: Double
      tao = read taoStr :: Double
      numTrail = read numTrailStr :: Int
      maxTrail = read maxTrailStr :: Int
      thetaFreq = read thetaFreqsStr :: Double
      thetaFreqs = [-thetaFreq .. thetaFreq]
      scaleFreq = read scaleFreqsStr :: Double
      scaleFreqs = [-scaleFreq .. scaleFreq]
      numThread = read numThreadStr :: Int
      folderPath = "output/test/PlotLocalEigenVector"
      histFilePath = folderPath </> histFileName
      imageFolder =
        printf
          "%d_%.2f_%.2f_%.0f_%.0f_%.0f_%.0f"
          numPoint
          thetaSigma
          scaleSigma
          maxScale
          tao
          thetaFreq
          scaleFreq
  removePathForcibly (folderPath </> imageFolder)
  createDirectoryIfMissing True (folderPath </> imageFolder)
  flag <- doesFileExist histFilePath
  radialArr <-
    if flag
      then R.map magnitude . getNormalizedHistogramArr <$>
           decodeFile histFilePath
      else do
        putStrLn "Couldn't find a Green's function data. Start simulation..."
        solveMonteCarloR2Z2T0S0Radial
          numThread
          numTrail
          maxTrail
          numPoint
          numPoint
          thetaSigma
          scaleSigma
          maxScale
          tao
          thetaFreqs
          thetaFreqs
          scaleFreqs
          scaleFreqs
          histFilePath
          (emptyHistogram
             [ (round . sqrt . fromIntegral $ 2 * (div numPoint 2) ^ 2)
             , L.length scaleFreqs
             , L.length thetaFreqs
             , L.length scaleFreqs
             , L.length thetaFreqs
             ]
             0)
  let localEigenVector =
        computeLocalEigenVector
          (ParallelParams numThread 0)
          (pinwheelHollow 8)
          (cutoff 64 radialArr)
          numPoint
          numPoint
          maxScale
          thetaFreqs
          scaleFreqs
      rotationDeg = 90
      rotatedLocalEigenVector =
        R.traverse2
          localEigenVector
          (fromListUnboxed (Z :. (L.length thetaFreqs)) thetaFreqs)
          const $ \f fFreq idx@(Z :. k :. l :. i :. j) ->
          f idx *
          ((exp (0 :+ (-1) * (fFreq (Z :. k)) * rotationDeg / 180 * pi)) --  +
           -- (exp (0 :+ (-1) * (fFreq (Z :. k)) * (rotationDeg + 90) / 180 * pi)) +
           -- (exp (0 :+ (-1) * (fFreq (Z :. k)) * (rotationDeg + 180) / 180 * pi)) +
           -- (exp (0 :+ (-1) * (fFreq (Z :. k)) * (rotationDeg + 270) / 180 * pi))
           )
  MP.mapM_
    (\(tf, sf) ->
       plotImageRepaComplex
         (folderPath </> imageFolder </> (printf "%d_%d.png" tf sf)) .
       ImageRepa 8 .
       computeS .
       R.extend (Z :. (1 :: Int) :. All :. All) . R.slice localEigenVector $
       (Z :. tf :. sf :. All :. All))
    [ (tf, sf)
    | tf <- [0 .. L.length thetaFreqs - 1]
    , sf <- [0 .. L.length scaleFreqs - 1]
    ]
  plotImageRepaComplex (folderPath </> imageFolder </> "sumFreq.png") .
    ImageRepa 8 .
    computeS .
    R.extend (Z :. (1 :: Int) :. All :. All) .
    R.sumS . R.sumS . rotate4D . rotate4D $
    localEigenVector
  plotImageRepaComplex (folderPath </> imageFolder </> "sumFreqRotation.png") .
    ImageRepa 8 .
    computeS .
    R.extend (Z :. (1 :: Int) :. All :. All) .
    R.sumS . R.sumS . rotate4D . rotate4D $
    rotatedLocalEigenVector
  let localEigenVectorR2S1RP =
        r2z2Tor2s1rp numOrientation thetaFreqs numScale scaleFreqs $
        localEigenVector
      rotatedLocalEigenVectorR2S1RP =
        r2z2Tor2s1rp numOrientation thetaFreqs numScale scaleFreqs $
        rotatedLocalEigenVector
  plotImageRepaComplex (folderPath </> imageFolder </> "sum.png") .
    ImageRepa 8 .
    computeS .
    R.extend (Z :. (1 :: Int) :. All :. All) .
    R.sumS . R.sumS . rotate4D . rotate4D $
    localEigenVectorR2S1RP
  plotImageRepaComplex (folderPath </> imageFolder </> "sumRotation.png") .
    ImageRepa 8 .
    computeS .
    R.extend (Z :. (1 :: Int) :. All :. All) .
    R.sumS . R.sumS . rotate4D . rotate4D $
    localEigenVectorR2S1RP
  MP.mapM_
    (\i ->
       plotImageRepaComplex
         (folderPath </> imageFolder </> printf "Theta%03d.png" i) .
       ImageRepa 8 .
       computeS .
       R.extend (Z :. (1 :: Int) :. All :. All) . R.slice localEigenVectorR2S1RP $
       (Z :. i :. (0 :: Int) :. All :. All))
    [0 .. numOrientation - 1]


{-# LANGUAGE FlexibleContexts, BangPatterns #-}

-- |
-- Module    : Statistics.Sample.Histogram
-- Copyright : (c) 2011 Bryan O'Sullivan
-- License   : BSD3
--
-- Maintainer  : bos@serpentine.com
-- Stability   : experimental
-- Portability : portable
--
-- Functions for computing histograms of sample data.

module Statistics.Sample.Histogram
    (
      histogram
    -- * Building blocks
    , histogram_
    , range
    ) where

import Numeric.MathFunctions.Constants (m_epsilon,m_tiny)
import Statistics.Function (minMax)
import qualified Data.Vector.Generic as G
import qualified Data.Vector.Generic.Mutable as GM

-- | /O(n)/ Compute a histogram over a data set.
--
-- The result consists of a pair of vectors:
--
-- * The lower bound of each interval.
--
-- * The number of samples within the interval.
--
-- Interval (bin) sizes are uniform, and the upper and lower bounds
-- are chosen automatically using the 'range' function.  To specify
-- these parameters directly, use the 'histogram_' function.
histogram :: (G.Vector v0 Double, G.Vector v1 Double, Num b, G.Vector v1 b) =>
             Int                -- ^ Number of bins (must be positive).
          -> v0 Double          -- ^ Sample data (cannot be empty).
          -> (v1 Double, v1 b)
histogram numBins xs = (G.generate numBins step, histogram_ numBins lo hi xs)
    where (lo,hi)    = range numBins xs
          step i     = lo + d * fromIntegral i
          d          = (hi - lo) / fromIntegral numBins
{-# INLINE histogram #-}

-- | /O(n)/ Compute a histogram over a data set.
--
-- Interval (bin) sizes are uniform, based on the supplied upper
-- and lower bounds.
histogram_ :: (Num b, RealFrac a, G.Vector v0 a, G.Vector v1 b) =>
              Int
           -- ^ Number of bins.  This value must be positive.  A zero
           -- or negative value will cause an error.
           -> a
           -- ^ Lower bound on interval range.  Sample data less than
           -- this will cause an error.
           -> a
           -- ^ Upper bound on interval range.  This value must not be
           -- less than the lower bound.  Sample data that falls above
           -- the upper bound will cause an error.
           -> v0 a
           -- ^ Sample data.
           -> v1 b
histogram_ numBins lo hi xs0 = G.create (GM.replicate numBins 0 >>= bin xs0)
  where
    bin xs bins = go 0
     where
       go i | i >= len = return bins
            | otherwise = do
         let x = xs `G.unsafeIndex` i
             b = truncate $ (x - lo) / d
         write' bins b . (+1) =<< GM.read bins b
         go (i+1)
       write' bins' b !e = GM.write bins' b e
       len = G.length xs
       d = ((hi - lo) / fromIntegral numBins) * (1 + realToFrac m_epsilon)
{-# INLINE histogram_ #-}

-- | /O(n)/ Compute decent defaults for the lower and upper bounds of
-- a histogram, based on the desired number of bins and the range of
-- the sample data.
--
-- The upper and lower bounds used are @(lo-d, hi+d)@, where
--
-- @d = (maximum sample - minimum sample) / ((bins - 1) * 2)@
--
-- If all elements in the sample are the same and equal to @x@ range
-- is set to @(x - |x|/10, x + |x|/10)@. And if @x@ is equal to 0 range
-- is set to @(-1,1)@. This is needed to avoid creating histogram with
-- zero bin size.
range :: (G.Vector v Double) =>
         Int                    -- ^ Number of bins (must be positive).
      -> v Double               -- ^ Sample data (cannot be empty).
      -> (Double, Double)
range numBins xs
    | numBins < 1 = error "Statistics.Histogram.range: invalid bin count"
    | G.null xs   = error "Statistics.Histogram.range: empty sample"
    | lo == hi    = case abs lo / 10 of
                      a | a < m_tiny -> (-1,1)
                        | otherwise  -> (lo - a, lo + a)
    | otherwise   = (lo-d, hi+d)
  where
    d | numBins == 1 = 0
      | otherwise    = (hi - lo) / ((fromIntegral numBins - 1) * 2)
    (lo,hi)          = minMax xs
{-# INLINE range #-}


-----------------------------------------------------------------------------
-- |
-- Copyright   :  (c) Edward Kmett 2010-2014
-- License     :  BSD3
-- Maintainer  :  ekmett@gmail.com
-- Stability   :  experimental
-- Portability :  GHC only
--
-- Root finding using Halley's rational method (the second in
-- the class of Householder methods). Assumes the function is three
-- times continuously differentiable and converges cubically when
-- progress can be made.
--
-----------------------------------------------------------------------------

module Numeric.AD.Rank1.Halley
  (
  -- * Halley's Method (Tower AD)
    findZero
  , inverse
  , fixedPoint
  , extremum
  ) where

import Prelude hiding (all)
import Numeric.AD.Internal.Forward (Forward)
import Numeric.AD.Internal.On
import Numeric.AD.Internal.Tower (Tower)
import Numeric.AD.Mode
import Numeric.AD.Rank1.Tower (diffs0)
import Numeric.AD.Rank1.Forward (diff)

-- $setup
-- >>> import Data.Complex

-- | The 'findZero' function finds a zero of a scalar function using
-- Halley's method; its output is a stream of increasingly accurate
-- results.  (Modulo the usual caveats.) If the stream becomes constant
-- ("it converges"), no further elements are returned.
--
-- Examples:
--
-- >>> take 10 $ findZero (\x->x^2-4) 1
-- [1.0,1.8571428571428572,1.9997967892704736,1.9999999999994755,2.0]
--
-- >>> last $ take 10 $ findZero ((+1).(^2)) (1 :+ 1)
-- 0.0 :+ 1.0
findZero :: (Fractional a, Eq a) => (Tower a -> Tower a) -> a -> [a]
findZero f = go where
  go x = x : if x == xn then [] else go xn where
    (y:y':y'':_) = diffs0 f x
    xn = x - 2*y*y'/(2*y'*y'-y*y'')
{-# INLINE findZero #-}

-- | The 'inverse' function inverts a scalar function using
-- Halley's method; its output is a stream of increasingly accurate
-- results.  (Modulo the usual caveats.) If the stream becomes constant
-- ("it converges"), no further elements are returned.
--
-- Note: the @take 10 $ inverse sqrt 1 (sqrt 10)@ example that works for Newton's method
-- fails with Halley's method because the preconditions do not hold!
inverse :: (Fractional a, Eq a) => (Tower a -> Tower a) -> a -> a -> [a]
inverse f x0 y = findZero (\x -> f x - auto y) x0
{-# INLINE inverse  #-}

-- | The 'fixedPoint' function find a fixedpoint of a scalar
-- function using Halley's method; its output is a stream of
-- increasingly accurate results.  (Modulo the usual caveats.)
--
-- If the stream becomes constant ("it converges"), no further
-- elements are returned.
--
-- >>> last $ take 10 $ fixedPoint cos 1
-- 0.7390851332151607
fixedPoint :: (Fractional a, Eq a) => (Tower a -> Tower a) -> a -> [a]
fixedPoint f = findZero (\x -> f x - x)
{-# INLINE fixedPoint #-}

-- | The 'extremum' function finds an extremum of a scalar
-- function using Halley's method; produces a stream of increasingly
-- accurate results.  (Modulo the usual caveats.) If the stream becomes
-- constant ("it converges"), no further elements are returned.
--
-- >>> take 10 $ extremum cos 1
-- [1.0,0.29616942658570555,4.59979519460002e-3,1.6220740159042513e-8,0.0]
extremum :: (Fractional a, Eq a) => (On (Forward (Tower a)) -> On (Forward (Tower a))) -> a -> [a]
extremum f = findZero (diff (off . f . On))
{-# INLINE extremum #-}


{-# LANGUAGE OverloadedLists #-}

module Main where

import Numeric.GSL.ODE
import Numeric.LinearAlgebra

-- Differential equation
f :: Double -> [Double] -> [Double]
f t [x, v] = [v, - x + mu * v * (1 - x ^ 2)]

-- Mu scalar, dampening strenth
mu :: Double
mu = 0.1

-- Boundary conditions
ts :: Vector Double
ts = linspace 1000 (0, 50)

-- Use default solver: Embedded Runge-Kutta-Fehlberg (4, 5) method.
vanderpol1 :: [Vector Double]
vanderpol1 = toColumns $ odeSolve f [1, 0] ts

-- Use Runge-Kutta (2,3) solver
vanderpol2 :: [Vector Double]
vanderpol2 = toColumns $ odeSolveV RK2 hi epsAbs epsRel (l2v f) [1, 0] ts
  where
    epsAbs = 1.49012e-08
    epsRel = epsAbs
    hi = (ts ! 1 - ts ! 0) / 100
    l2v f = \t -> fromList . f t . toList

main :: IO ()
main = do
  print vanderpol1
  print vanderpol2