{-#Language TypeFamilies#-}
module Morse where

import Stratification as S hiding (shiftCeil)
import Data.Matrix (Matrix(..))
import Control.Monad
import Control.Arrow
import SymbolicImage (
    fromRad, toRad, fromCeil, toCeil, myNub, mySequence, (#=))
import Jac (jac, e)
import Data.Array as A
import Graph
import Data.Function
import Data.List as L
import Data.Map as M
import Data.IntMap as IM
import Data.Complex
import Data
import Point
import Data.Graph

morse :: F Point -> Imagination -> Int -> [(Double, Double, Int)]
morse fp im i = mrsElem <$> cyclics
    ((shiftCeil d1 d2 fp #=) <$> ls)
  where
    mrsElem :: [(Ceil3, [(Ceil3, Double)])] -> (Double, Double, Int)
    mrsElem a = (
                         minOptZ gr  (minBazeCircuit gr),
                negate $ minOptZ gr' (maxBazeCircuit gr'),
                length a
--                       maximum $ plat a,
--                       minimum $ plat a
            )
        where
            gr  = formGraph                    a
            gr' = formGraph $ modifyIns negate a

    (Stratification d1 d2 ls) = str fp im i

plat :: [(Ceil3, [(Ceil3, Double)])] -> [Double]
plat = concatMap (\x -> snd <$> snd x)

-- XXX Дзен
modifyIns :: (a -> b) -> [(c, [(d, a)])] -> [(c, [(d, b)])]
modifyIns f l = second (second f <$>) <$> l

-- выделяем области связности
--cyclics :: (Ord t0, Ord k0) => [(k0, t0, [(t0, b0)])] -> [[(k0, [(t0, b0)])]]
cyclics g = M.toList <$> do
    let g' = (\(a, b, l) -> (a, b, L.filter (\x -> fst x /= a) l)) <$> g
    CyclicSCC a <- stronglyConnComp
        [(a, b, fst <$> c) | (a, b, c) <- g']
    return $ elect a $ M.fromList [(a, b) | (a, _, b) <- g']

shiftCeil :: Diameter -> Diameter -> F Point -> Ceil3 -> [(Ceil3, Double)]
shiftCeil d1 d2 fp (Ceil3 c l) = myNub $ do
    let Point x y = fromCeil d1 c
    ret <$> mySequence x d1
        <*> mySequence y d1
        <*> mySequence (rad d2 l) (pi/2/d2)
  where
    ret x' y' t' = (
        func d1 d2 fp x' y' t',
        log $ x $ abs $ toPoint $ jacobian fp x' y' * e t')

elect :: Ord a => [a] -> F (Map a b)
elect a m = M.intersection m $ M.fromList $ (\a -> (a,a)) <$> a


module Neural.Util where

import Control.Monad
import Control.Monad.ST
import Control.Parallel.Strategies
import Data.Array.ST
import Data.STRef
import Numeric.LinearAlgebra
import System.Random

 -- shamelessly nicked from https://wiki.haskell.org/Random_shuffle

shuffle' :: [a] -> StdGen -> ([a],StdGen)
shuffle' xs gen = runST (do
        g <- newSTRef gen
        let randomRST lohi = do
              (a,s') <- liftM (randomR lohi) (readSTRef g)
              writeSTRef g s'
              return a
        ar <- newArray n xs
        xs' <- forM [1..n] $ \i -> do
                j <- randomRST (i,n)
                vi <- readArray ar i
                vj <- readArray ar j
                writeArray ar j vi
                return vj
        gen' <- readSTRef g
        return (xs',gen'))
  where
    n = length xs
    newArray :: Int -> [a] -> ST s (STArray s Int a)
    newArray n xs =  newListArray (1,n) xs

shuffleIO :: [a] -> IO [a]
shuffleIO xs = getStdRandom (shuffle' xs)

parZipWith :: Strategy c -> (a -> b -> c) -> [a] -> [b] -> [c]
parZipWith strat f x y = zipWith f x y `using` parList strat

parZipWith3 :: Strategy d -> (a -> b -> c -> d) -> [a] -> [b] -> [c] -> [d]
parZipWith3 strat f x y z = zipWith3 f x y z `using` parList strat

fullyFlatten :: [Matrix Double] -> Vector Double
fullyFlatten = vjoin . (parMap rdeepseq flatten)


--------------------------------------------------------------------------------
-- |
-- Module    :  Datasets.Iris
-- Copyright :  (c) Sam Stites 2017
-- License   :  BSD3
-- Maintainer:  sam@stites.io
-- Stability :  experimental
-- Portability: non-portable
--
-- === Iris Data Set
--
-- Data Set Characteristics  :  Multivariate
-- Number of Instances       : 150
-- Area                      : Life
-- Attribute Characteristics : Real
-- Number of Attributes      : 4
-- Date Donated              : 1988-07-01
-- Associated Tasks          : Classification
-- Missing Values?           : No
-- Creator                   : R. A. Fisher
-- Source                    : https://archive.ics.uci.edu/ml/datasets/Iris
--
-- === Data Set Information
--
-- This is perhaps the best known database to be found in the pattern
-- recognition literature. Fisher's paper is a classic in the field and is
-- referenced frequently to this day. (See Duda & Hart, for example.) The data
-- set contains 3 classes of 50 instances each, where each class refers to a
-- type of iris plant. One class is linearly separable from the other 2; the
-- latter are NOT linearly separable from each other.
--
-- Predicted attribute: class of iris plant.
--
-- This is an exceedingly simple domain.
--
-- This data differs from the data presented in Fishers article (identified by
-- Steve Chadwick, spchadwick '@' espeedaz.net ). The 35th sample should be:
-- 4.9,3.1,1.5,0.2,"Iris-setosa" where the error is in the fourth feature. The
-- 38th sample: 4.9,3.6,1.4,0.1,"Iris-setosa" where the errors are in the second
-- and third features.
--
--
-- === Attribute Information
--
-- 1. Sepal length in cm
-- 2. Sepal width in cm
-- 3. Petal length in cm
-- 4. Petal width in cm
-- 5. Class: Iris Setosa | Iris Versicolour | Iris Virginica
--
-- === Summary Statistics:
--
--               Min  Max   Mean    SD   Class Correlation
-- sepal length: 4.3  7.9   5.84  0.83    0.7826
-- sepal width : 2.0  4.4   3.05  0.43   -0.4194
-- petal length: 1.0  6.9   3.76  1.76    0.9490  (high!)
-- petal width : 0.1  2.5   1.20  0.76    0.9565  (high!)
--
-- === Class Distribution
--
-- 33.3% for each of 3 classes.
--------------------------------------------------------------------------------
{-# LANGUAGE DeriveGeneric #-}
module Datasets.Iris where

import Prelude hiding (readFile)
import Data.Csv
import Control.Arrow
import GHC.Generics
import Data.ByteString.Lazy (readFile)
import Control.Monad.IO.Class
import Control.Exception.Safe

import Data.Vector (Vector, toList, fromList)
import qualified Data.Vector as V
import Numeric.LinearAlgebra (Matrix, fromLists)
import Numeric.LinearAlgebra.Data ((??), Extractor(..))

data Datum = Datum
  { sepalLength :: Double
  , sepalWidth  :: Double
  , petalLength :: Double
  , petalWidth  :: Double
  , irisClass   :: IrisClass
  } deriving (Eq, Show, Generic)

instance FromRecord Datum

toDoubles :: Datum -> [Double]
toDoubles (Datum a b c d e) = [a, b, c, d, fromIntegral $ fromEnum e]
-- ========================================================================= --

data IrisClass
  = Setosa
  | Versicolour
  | Virginica
  deriving (Eq, Bounded, Enum, Generic)

instance FromRecord IrisClass
instance FromField  IrisClass where
  parseField s
    | s == "Iris-setosa"     = pure Setosa
    | s == "Iris-versicolor" = pure Versicolour
    | s == "Iris-virginica"  = pure Virginica

instance Show IrisClass where
  show Setosa      = "Iris Setosa"
  show Versicolour = "Iris Versicolour"
  show Virginica   = "Iris Virginica"

-- ========================================================================= --

loadIrisRaw :: (MonadThrow m, MonadIO m) => m (Vector Datum)
loadIrisRaw = do
  iris <- liftIO $ readFile "data/Datasets/Iris.csv"
  Right xs <- pure $ decode NoHeader iris
  return xs

loadIris :: (MonadThrow m, MonadIO m) => m (Matrix Double, Vector IrisClass)
loadIris = featureLabelSplit <$> loadIrisRaw
  where
    featureLabelSplit :: Vector Datum -> (Matrix Double, Vector IrisClass)
    featureLabelSplit v = foo $ fmap ((init . toDoubles) &&& irisClass) v

    foo :: Vector ([Double], IrisClass) -> (Matrix Double, Vector IrisClass)
    foo vs = (fromLists $ toList (fmap fst vs), fmap snd vs)

type Dataset = (Matrix Double, Vector IrisClass)

splitCV :: Float -> Dataset -> (Dataset, Dataset)
splitCV percent (features, labels) =
  ( (features ?? (Take splitpoint, All), V.take splitpoint labels)
  , (features ?? (Drop splitpoint, All), V.drop splitpoint labels)
  )

  where
    splitpoint :: Int
    splitpoint = truncate (fromIntegral (length labels) * percent) :: Int


{-# LANGUAGE FlexibleContexts #-}
module R2Z2T0S0ToR2S1RPT0S0 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           System.Directory
import           System.Environment
import           System.FilePath
import           Text.Printf
import           Utils.Array
import           Utils.Time

main = do
  args@(numPointStr:numOrientationStr:numScaleStr:thetaSigmaStr:scaleSigmaStr:maxScaleStr:taoStr:lenStr:initStr:numTrailStr:maxTrailStr:theta0FreqsStr:thetaFreqsStr:scale0FreqsStr:scaleFreqsStr:histPath:alphaStr:pinwheelFlagStr: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
      len = read lenStr :: Int
      init = read initStr :: (Double, Double, Double, Double, Double, Double)
      numTrail = read numTrailStr :: Int
      maxTrail = read maxTrailStr :: Int
      theta0Freq = read theta0FreqsStr :: Double
      thetaFreq = read thetaFreqsStr :: Double
      scale0Freq = read scale0FreqsStr :: Double
      scaleFreq = read scaleFreqsStr :: Double
      theta0Freqs = [-theta0Freq .. theta0Freq]
      thetaFreqs = [-thetaFreq .. thetaFreq]
      scale0Freqs = [-scale0Freq .. scale0Freq]
      scaleFreqs = [-scaleFreq .. scaleFreq]
      alpha = read alphaStr :: Double
      pinwheelFlag = read pinwheelFlagStr :: Bool
      numThread = read numThreadStr :: Int
      folderPath = "output/test/R2Z2T0S0ToR2S1RPT0S0"
  flag <- doesFileExist histPath
  -- arrR2Z2T0S0 <-
  --   -- if flag
  --   --   then getNormalizedHistogramArr <$> decodeFile histPath
  --   --   else
  --     solveMonteCarloR2Z2T0S0
  --            numThread
  --            numTrail
  --            maxTrail
  --            numPoint
  --            numPoint
  --            thetaSigma
  --            scaleSigma
  --            maxScale
  --            tao
  --            len
  --            theta0Freqs
  --            thetaFreqs
  --            scale0Freqs
  --            scaleFreqs
  --            histPath
  --            (emptyHistogram
  --               [ numPoint
  --               , numPoint
  --               , L.length scale0Freqs
  --               , L.length theta0Freqs
  --               , L.length scaleFreqs
  --               , L.length thetaFreqs
  --               ]
  --               0)
  radialArr <-
    if flag
      then R.map magnitude . getNormalizedHistogramArr <$> decodeFile histPath
      else do
        putStrLn "Couldn't find a Green's function data. Start simulation..."
        printCurrentTime
        solveMonteCarloR2Z2T0S0Radial
          numThread
          numTrail
          maxTrail
          numPoint
          numPoint
          thetaSigma
          scaleSigma
          maxScale
          tao
          theta0Freqs
          thetaFreqs
          scale0Freqs
          scaleFreqs
          ""
          (emptyHistogram
             [ (round . sqrt . fromIntegral $ 2 * (div numPoint 2) ^ 2)
             , L.length scale0Freqs
             , L.length theta0Freqs
             , L.length scaleFreqs
             , L.length thetaFreqs
             ]
             0)
  printCurrentTime
  arrR2Z2T0S0 <-
    computeUnboxedP $
    computeR2Z2T0S0ArrayRadial
      radialArr
      numPoint
      numPoint
      1
      maxScale
      thetaFreqs
      scaleFreqs
      theta0Freqs
      scale0Freqs
  arrR2Z2 <- R.sumP . R.sumS . rotateR2Z2T0S0Array $ arrR2Z2T0S0
      -- arr =
      --   r2z2t0s0Tor2s1rpt0s0
      --     numOrientation
      --     thetaFreqs
      --     numScale
      --     scaleFreqs
      --     arrR2Z2T0S0
      -- arr4d =
      --   R.slice arr $
      --   (Z :. All :. All :. (L.length theta0Freqs - 1) :.
      --    (L.length scale0Freqs - 1) :.
      --    All :.
      --    All)
  let arr4d = r2z2Tor2s1rp numOrientation thetaFreqs numScale scaleFreqs arrR2Z2
  createDirectoryIfMissing True folderPath
  printCurrentTime
  MP.mapM_
    (\(i, j) ->
       plotImageRepaComplex
         (folderPath </> show (i + 1) L.++ "_" L.++ show (j + 1) L.++ ".png") .
       ImageRepa 8 .
       computeS . R.extend (Z :. (1 :: Int) :. All :. All) . R.slice arr4d $
       (Z :. i :. j :. All :. All))
    [(i, j) | i <- [0 .. numOrientation - 1], j <- [0 .. numScale - 1]]
  arr2D <- R.sumP . R.sumS . rotate4D . rotate4D $ arr4d
  plotImageRepaComplex (folderPath </> "sum.png") .
    ImageRepa 8 . computeS . R.extend (Z :. (1 :: Int) :. All :. All) $
    arr2D


module Statistics.Quantile.Bench.Accuracy where

import qualified Data.Vector as V

import Statistics.Quantile.Types
import Statistics.Quantile.Util

import Statistics.Sample

import System.IO

err :: Double
    -> Double
    -> Double
err true estimate =
  let e = abs (true - estimate)
  in e * e

selectorAccuracy :: Stream IO
                 -> Quantile
                 -> Double
                 -> Selector IO
                 -> IO Double
selectorAccuracy src q true (Selector select) =
  err true <$> select q src

benchAccuracy :: Int
              -> FilePath
              -> Quantile
              -> Double
              -> Selector IO
              -> IO Deviation
benchAccuracy n fp q true s = do
  as <- V.replicateM n accuracy
  pure $ Deviation (mean as) (stdDev as)
  where accuracy = do
          h <- openFile fp ReadMode
          r <- selectorAccuracy (streamHandle h) q true s
          hClose h
          pure r


{-# LANGUAGE OverloadedStrings #-}

module Statistics.Classification.ConfusionMatrix.Binary (
    ConfusionMatrix
  , trueConfMatrix
  , resultCount, classCount, totalCount
  , actualPositive, actualNegative, predPositive, predNegative
  , truePosRate, falseNegRate, falsePosRate, trueNegRate
  , prevalence, accuracy, errorRate, recall, precision, specificity
  , oddsRatio
  , phiCorrelation
  , fScore, f1Score
  , actualEntropy, predEntropy, condActualEntropy, condPredEntropy, totalCondEntropy
  , mutualInformation
  , asBinaryConfusion
  , asMultiConfusion
  ) where

import Data.Monoid (Sum(..))

import Data.Csv ((.:), (.=), DefaultOrdered(..), FromNamedRecord(..), ToNamedRecord(..), header, namedRecord)
import qualified Data.Map.Strict as Map
import qualified Data.Set as Set

import Statistics.Classification.ConfusionMatrix.Multi (MultiConfusionMatrix(..), PosNeg
  , truePositiveCount, falsePositiveCount, falseNegativeCount, trueNegativeCount)
import Statistics.Classification.Types (ClassificationResult(..))

data ConfusionMatrix = ConfusionMatrix {
    truePositive :: {-# UNPACK #-}!(Sum Int)
  , falsePositive :: {-# UNPACK #-}!(Sum Int)
  , falseNegative :: {-# UNPACK #-}!(Sum Int)
  , trueNegative :: {-# UNPACK #-}!(Sum Int)
  } deriving (Eq, Show, Read)

instance Semigroup ConfusionMatrix where
  cm0 <> cm1 = ConfusionMatrix {
      truePositive  = truePositive  cm0 <> truePositive  cm1
    , falsePositive = falsePositive cm0 <> falsePositive cm1
    , falseNegative = falseNegative cm0 <> falseNegative cm1
    , trueNegative  = trueNegative  cm0 <> trueNegative  cm1
    }

instance Monoid ConfusionMatrix where
  mempty = ConfusionMatrix mempty mempty mempty mempty
  mappend = (<>)

instance DefaultOrdered ConfusionMatrix where
  headerOrder _ = header [
      "true_positive"
    , "false_positive"
    , "true_negative"
    , "false_negative"
    ]

instance ToNamedRecord ConfusionMatrix where
  toNamedRecord conf = namedRecord [
      "true_positive"  .= getSum (truePositive conf)
    , "false_positive" .= getSum (falsePositive conf)
    , "true_negative"  .= getSum (trueNegative conf)
    , "false_negative" .= getSum (falseNegative conf)
    ]

instance FromNamedRecord ConfusionMatrix where
  parseNamedRecord r = ConfusionMatrix <$>
        fmap Sum (r .: "true_positive")
    <*> fmap Sum (r .: "false_positive")
    <*> fmap Sum (r .: "false_negative")
    <*> fmap Sum (r .: "true_negative")

-- create confusion matrix given assumed true results
-- i.e., classifiedPredicted is the number predicted correctly
--       classifiedActual    is the number of actual samples in class
-- Give numbers for both true and false labels
trueConfMatrix :: ClassificationResult Int -> ClassificationResult Int -> ConfusionMatrix
trueConfMatrix pos neg = ConfusionMatrix {
    truePositive  = Sum $ classifiedPredicted pos
  , falseNegative = Sum $ classifiedActual pos - classifiedPredicted pos
  , falsePositive = Sum $ classifiedActual neg - classifiedPredicted neg
  , trueNegative  = Sum $ classifiedPredicted neg
  }

resultCount :: ConfusionMatrix -> ClassificationResult Bool -> Int
resultCount cm (ClassificationResult actual prd) = classCount cm prd actual

classCount :: ConfusionMatrix -> Bool -> Bool -> Int
classCount cm prd actual = getSum $ if prd
  then (if actual then  truePositive cm else falsePositive cm)
  else (if actual then falseNegative cm else  trueNegative cm)

totalCount :: ConfusionMatrix -> Int
totalCount cm = getSum $ truePositive cm <> falsePositive cm <> falseNegative cm <> trueNegative cm

actualPositive, actualNegative, predPositive, predNegative :: ConfusionMatrix -> Int
actualPositive cm = getSum $  truePositive cm <> falseNegative cm
actualNegative cm = getSum $ falsePositive cm <>  trueNegative cm
predPositive   cm = getSum $  truePositive cm <> falsePositive cm
predNegative   cm = getSum $ falseNegative cm <>  trueNegative cm

truePosRate, falseNegRate, falsePosRate, trueNegRate :: ConfusionMatrix -> Double
truePosRate = recall
falseNegRate cm = fromIntegral (getSum (falseNegative cm)) / fromIntegral (actualPositive cm)
falsePosRate cm = fromIntegral (getSum (falsePositive cm)) / fromIntegral (actualNegative cm)
trueNegRate = specificity

prevalence, accuracy, errorRate, recall, precision, specificity :: ConfusionMatrix -> Double
prevalence cm = fromIntegral (actualPositive cm) / fromIntegral (totalCount cm)
accuracy cm = fromIntegral (getSum (truePositive cm <> trueNegative cm)) / fromIntegral (totalCount cm)
errorRate cm = fromIntegral (getSum (falsePositive cm <> falseNegative cm)) / fromIntegral (totalCount cm)
recall cm = fromIntegral (getSum (truePositive cm)) / fromIntegral (actualPositive cm)
precision cm = fromIntegral (getSum (truePositive cm)) / fromIntegral (predPositive cm)
specificity cm = fromIntegral (getSum (trueNegative cm)) / fromIntegral (actualNegative cm)

oddsRatio :: ConfusionMatrix -> Double
oddsRatio (ConfusionMatrix (Sum tp) (Sum fp) (Sum fn) (Sum tn)) =
  fromIntegral (tp * tn) / fromIntegral (fn * fp)

phiCorrelation :: ConfusionMatrix -> Double
phiCorrelation (ConfusionMatrix (Sum tp) (Sum fp) (Sum fn) (Sum tn)) =
  fromIntegral (tp * tn - fp * fn) / sqrt (fromIntegral $ (tp + fn) * (tn + fp) * (tp + fp) * (tn + fn))

fScore :: Double -> ConfusionMatrix -> Double
fScore alpha cm = recip $ alpha * recip (precision cm) + (1 - alpha) * recip (recall cm)

f1Score :: ConfusionMatrix -> Double
f1Score = fScore 0.5

-- Entropy calculations
binomEntropy :: Double -> Double
binomEntropy p = p * log p + (1 - p) * log (1 - p)

discreteEntropy :: Int -> Int -> Double
discreteEntropy n tot = binomEntropy (fromIntegral n / fromIntegral tot)

actualEntropy, predEntropy :: ConfusionMatrix -> Double
actualEntropy cm = discreteEntropy (actualPositive cm) (totalCount cm)
predEntropy cm = discreteEntropy (predPositive cm) (totalCount cm)

condActualEntropy :: ConfusionMatrix -> Bool -> Double
condActualEntropy cm plbl = if plbl
  then discreteEntropy (getSum (truePositive cm)) (predPositive cm)
  else discreteEntropy (getSum (trueNegative cm)) (predNegative cm)

condPredEntropy :: ConfusionMatrix -> Bool -> Double
condPredEntropy cm albl = if albl
  then discreteEntropy (getSum (truePositive cm)) (actualPositive cm)
  else discreteEntropy (getSum (trueNegative cm)) (actualNegative cm)

totalCondEntropy :: ConfusionMatrix -> Double
totalCondEntropy cm = ppos * condActualEntropy cm True + (1 - ppos) * condActualEntropy cm False
  where ppos = fromIntegral (predPositive cm) / fromIntegral (totalCount cm)

mutualInformation :: ConfusionMatrix -> Double
mutualInformation cm = actualEntropy cm - totalCondEntropy cm

asBinaryConfusion :: Ord a => PosNeg a -> MultiConfusionMatrix a -> ConfusionMatrix
asBinaryConfusion posneg mcm = ConfusionMatrix {
    truePositive = Sum $ truePositiveCount mcm posneg
  , falsePositive = Sum $ falsePositiveCount mcm posneg
  , falseNegative = Sum $ falseNegativeCount mcm posneg
  , trueNegative = Sum $ trueNegativeCount mcm posneg
  }

asMultiConfusion :: ConfusionMatrix -> MultiConfusionMatrix Bool
asMultiConfusion cm = MultiConfusionMatrix {
    confCounts = Map.fromList [
        (ClassificationResult True  True,  getSum (truePositive cm))
      , (ClassificationResult False True,  getSum (falsePositive cm))
      , (ClassificationResult True  False, getSum (falseNegative cm))
      , (ClassificationResult False False, getSum (trueNegative cm))
      ]
  , confClasses = Set.fromList [True, False]
  }


{-# LANGUAGE TemplateHaskell #-}
{-# LANGUAGE Strict          #-}

module Optimization.Types where

import Control.Lens
import Numeric.LinearAlgebra            ( Numeric )
import Numeric.LinearAlgebra.Data
import Data.Vector.Storable       as VS ( map
                                        , singleton
                                        , zipWith )

{- |Basic "zippy" Num instance for storable vectors. Note that, while this is defined
    for some types of Vector by hmatrix, it unfortunately does not account for any
    Num (Vector a). -}
instance Numeric a => Num (Vector a) where
  a + b    = VS.zipWith (+)   a b
  a - b    = VS.zipWith (-)   a b
  a * b    = VS.zipWith (*)   a b
  abs a    = VS.map  (abs)    a
  signum a = VS.map  (signum) a
  fromInteger i = VS.singleton (fromInteger i)

{- |Termination criterion for a numerical optimization method.
    We terminate either on maxit iterations or once our guesses are within
    threshold of one another. -}
data TerminationCriterion a = TerminateAt {
  _maxit     :: Maybe Int, -- ^Maximum number of iterations (if any).
  _threshold :: Maybe a    -- ^Threshold used to declare convergence (if any).
  } deriving (Show, Eq)

-- |Step parameters for Nelder-Mead simplex search. a must be Storable and Num.
data NMStepParameters a = NMStepParameters {
  _epsilon :: a, -- ^Offset around initial guess to construct starting simplex.
  _alpha   :: a, -- ^Reflection coefficient for Nelder-Mead algorithm.
  _gamma   :: a, -- ^Expansion coefficient for Nelder-Mead algorithm.
  _rho     :: a, -- ^Contraction coefficient for Nelder-Mead algorithm.
  _sigma   :: a  -- ^Shrink coefficient for Nelder-Mead algorithm.
  } deriving (Show, Eq)

-- |Standard Nelder-Mead parameters.
defaultStep :: NMStepParameters Double
defaultStep = NMStepParameters 1.0 1.0 2.0 0.5 0.5

-- |The Nelder-Mead simplex.
type Simplex a = Matrix a

-- |Create a TerminationCriterion with just a convergence threshold.
terminateAt ::
  a ->                   -- ^Convengence threshold
  TerminationCriterion a -- ^A TerminationCriterion for optimization.
terminateAt t = TerminateAt Nothing (Just t)

-- |Create a TerminationCriterion with just a maximum iteration threshold.
terminateAfter ::
  Int               ->   -- ^Maximum number of iterations.
  TerminationCriterion a -- ^A TerminationCriterion for optimization.
terminateAfter m = TerminateAt (Just m) Nothing

-- |Create a TerminationCriterion with a convergence threshold and maximum iteration
-- |limit.
terminateAtOrAfter ::
  a      ->              -- ^Convergence threshold
  Int    ->              -- ^Maximum number of iterations
  TerminationCriterion a -- ^A TerminationCriterion for optimization.
terminateAtOrAfter t m = TerminateAt (Just m) (Just t)

makeLenses ''TerminationCriterion
makeLenses ''NMStepParameters


{-# LANGUAGE CPP                  #-}
{-# LANGUAGE FlexibleContexts     #-}
{-# LANGUAGE FlexibleInstances    #-}
{-# LANGUAGE ScopedTypeVariables  #-}
{-# LANGUAGE UndecidableInstances #-}

{-# OPTIONS_GHC -fno-warn-orphans #-}
{-# OPTIONS_GHC -fno-warn-missing-signatures #-}

-----------------------------------------------------------------------------
{- |
Module      :  Numeric.LinearAlgebra.Tests.Instances
Copyright   :  (c) Alberto Ruiz 2008
License     :  BSD3
Maintainer  :  Alberto Ruiz
Stability   :  provisional

Arbitrary instances for vectors, matrices.

-}

module Test.Numeric.LinearAlgebra.Tests.Instances(
    Sq(..),     rSq,cSq,
    Sq2WC(..),  rSq2WC,cSq2WC,
    Rot(..),    rRot,cRot,
                rHer,cHer,
    WC(..),     rWC,cWC,
    SqWC(..),   rSqWC, cSqWC, rSymWC, cSymWC,
    PosDef(..), rPosDef, cPosDef,
    Consistent(..), rConsist, cConsist,
    RM,CM, rM,cM,
    FM,ZM, fM,zM
) where

import           System.Random

import           Control.Monad                 (replicateM)
import           Numeric.LinearAlgebra.HMatrix hiding (vector)
import           Test.QuickCheck               (Arbitrary, arbitrary, choose, shrink,
                                                sized, vector)

import           Data.Proxy                    (Proxy (..))
import           GHC.TypeLits
import qualified Numeric.LinearAlgebra.Static  as Static
#if MIN_VERSION_base(4,11,0)
import           Prelude                       hiding ((<>))
#endif

shrinkListElementwise :: (Arbitrary a) => [a] -> [[a]]
shrinkListElementwise []     = []
shrinkListElementwise (x:xs) = [ y:xs | y  <- shrink x                 ]
                            ++ [ x:ys | ys <- shrinkListElementwise xs ]

shrinkPair :: (Arbitrary a, Arbitrary b) => (a,b) -> [(a,b)]
shrinkPair (a,b) = [ (a,x) | x <- shrink b ] ++ [ (x,b) | x <- shrink a ]

chooseDim = sized $ \m -> choose (1,max 1 m)

instance (Field a, Arbitrary a) => Arbitrary (Vector a) where
    arbitrary = do m <- chooseDim
                   l <- vector m
                   return $ fromList l
    -- shrink any one of the components
    shrink = map fromList . shrinkListElementwise . toList

instance KnownNat n => Arbitrary (Static.R n) where
    arbitrary = do
      l <- vector n
      return (Static.fromList l)

      where
        n :: Int
        n = fromIntegral (natVal (Proxy :: Proxy n))

    shrink _v = []

instance (Element a, Arbitrary a) => Arbitrary (Matrix a) where
    arbitrary = do
        m <- chooseDim
        n <- chooseDim
        l <- vector (m*n)
        return $ (m><n) l

    -- shrink any one of the components
    shrink a = map (rows a >< cols a)
               . shrinkListElementwise
               . concat . toLists
                     $ a

instance (KnownNat n, KnownNat m) => Arbitrary (Static.L m n) where
    arbitrary = do
      l <- vector (m * n)
      return (Static.fromList l)

      where
        m :: Int
        m = fromIntegral (natVal (Proxy :: Proxy m))

        n :: Int
        n = fromIntegral (natVal (Proxy :: Proxy n))

    shrink _mat = []

-- a square matrix
newtype (Sq a) = Sq (Matrix a) deriving Show
instance (Element a, Arbitrary a) => Arbitrary (Sq a) where
    arbitrary = do
        n <- chooseDim
        l <- vector (n*n)
        return $ Sq $ (n><n) l

    shrink (Sq a) = [ Sq b | b <- shrink a ]

-- a pair of square matrices
newtype (Sq2WC a) = Sq2WC (Matrix a, Matrix a) deriving Show
instance (ArbitraryField a, Numeric a) => Arbitrary (Sq2WC a) where
    arbitrary = do
        n <- chooseDim
        l <- vector (n*n)
        r <- vector (n*n)
        l' <- makeWC $ (n><n) l
        r' <- makeWC $ (n><n) r
        return $ Sq2WC (l', r')
        where
            makeWC m = do
              let (u,_,v) = svd m
                  n = rows m
              sv' <- replicateM n (choose (1,100))
              let s = diag (fromList sv')
              return $ u <> real s <> tr v

-- a unitary matrix
newtype (Rot a) = Rot (Matrix a) deriving Show
instance (Field a, Arbitrary a) => Arbitrary (Rot a) where
    arbitrary = do
        Sq m <- arbitrary
        let (q,_) = qr m
        return (Rot q)

-- a complex hermitian or real symmetric matrix
instance (Field a, Arbitrary a, Num (Vector a)) => Arbitrary (Herm a) where
    arbitrary = do
        Sq m <- arbitrary
        let m' = m/2
        return $ sym m'

class (Field a, Arbitrary a, Element (RealOf a), Random (RealOf a)) => ArbitraryField a
instance ArbitraryField Float
instance ArbitraryField (Complex Float)

-- a well-conditioned general matrix (the singular values are between 1 and 100)
newtype (WC a) = WC (Matrix a) deriving Show
instance (Numeric a, ArbitraryField a) => Arbitrary (WC a) where
    arbitrary = do
        m <- arbitrary
        let (u,_,v) = svd m
            r = rows m
            c = cols m
            n = min r c
        sv' <- replicateM n (choose (1,100))
        let s = diagRect 0 (fromList sv') r c
        return $ WC (u <> real s <> tr v)

-- a well-conditioned square matrix (the singular values are between 1 and 100)
newtype (SqWC a) = SqWC (Matrix a) deriving Show
instance (ArbitraryField a, Numeric a) => Arbitrary (SqWC a) where
    arbitrary = do
        Sq m <- arbitrary
        let (u,_,v) = svd m
            n = rows m
        sv' <- replicateM n (choose (1,100))
        let s = diag (fromList sv')
        return $ SqWC (u <> real s <> tr v)

-- a positive definite square matrix (the eigenvalues are between 0 and 100)
newtype (PosDef a) = PosDef (Matrix a) deriving Show
instance (Numeric a, ArbitraryField a, Num (Vector a))
    => Arbitrary (PosDef a) where
    arbitrary = do
        m <- arbitrary
        let (_,v) = eigSH m
            n = rows (unSym m)
        l <- replicateM n (choose (0,100))
        let s = diag (fromList l)
            p = v <> real s <> tr v
        return $ PosDef (0.5 * p + 0.5 * tr p)

-- a pair of matrices that can be multiplied
newtype (Consistent a) = Consistent (Matrix a, Matrix a) deriving Show
instance (Field a, Arbitrary a) => Arbitrary (Consistent a) where
    arbitrary = do
        n <- chooseDim
        k <- chooseDim
        m <- chooseDim
        la <- vector (n*k)
        lb <- vector (k*m)
        return $ Consistent ((n><k) la, (k><m) lb)

    shrink (Consistent (x,y)) = [ Consistent (u,v) | (u,v) <- shrinkPair (x,y) ]

type RM = Matrix Float
type CM = Matrix (Complex Float)
type FM = Matrix Float
type ZM = Matrix (Complex Float)

rM m = m :: RM
cM m = m :: CM
fM m = m :: FM
zM m = m :: ZM

rHer m = unSym m :: RM
cHer m = unSym m :: CM

rRot (Rot m) = m :: RM
cRot (Rot m) = m :: CM

rSq  (Sq m)  = m :: RM
cSq  (Sq m)  = m :: CM

rSq2WC (Sq2WC (a, b)) = (a, b) :: (RM, RM)
cSq2WC (Sq2WC (a, b)) = (a, b) :: (CM, CM)

rWC (WC m) = m :: RM
cWC (WC m) = m :: CM

rSqWC (SqWC m) = m :: RM
cSqWC (SqWC m) = m :: CM

rSymWC (SqWC m) = sym m :: Herm R
cSymWC (SqWC m) = sym m :: Herm C

rPosDef (PosDef m) = m :: RM
cPosDef (PosDef m) = m :: CM

rConsist (Consistent (a,b)) = (a,b::RM)
cConsist (Consistent (a,b)) = (a,b::CM)


module School.Types.Encoding
( Put
, putDouble
, putInt
, putWord
, doubleToBin
, intToBin
, wordToBin
, matrixDoubleToBin
, matrixIntToBin
, runPut
) where

import Data.ByteString (ByteString)
import Data.Serialize.IEEE754 (putFloat64be)
import Data.Serialize.Put (Put, putInt32be, putWord8, runPut)
import Data.Word (Word8)
import Numeric.LinearAlgebra (Element, I, Matrix, R, toLists)
import School.Types.DataType (DataType(..))

putDouble :: R -> Put
putDouble = putFloat64be

doubleToBin :: R -> ByteString
doubleToBin = runPut . putDouble

putInt :: I -> Put
putInt = putInt32be . fromIntegral

intToBin :: I -> ByteString
intToBin = runPut . putInt

putWord :: Word8 -> Put
putWord = putWord8

wordToBin :: Word8 -> ByteString
wordToBin = runPut . putWord

writer :: (Element b)
       => (a -> Put)
       -> ([b] -> [a])
       -> (Matrix b -> ByteString)
writer put trans matrix = runPut $
  mapM_ put ( trans
            . concat
            . toLists
            $ matrix
            )

matrixDoubleToBin :: DataType
                  -> Matrix R
                  -> ByteString
matrixDoubleToBin DBL64B = writer putDouble id
matrixDoubleToBin INT32B = undefined
matrixDoubleToBin INT08B = undefined

matrixIntToBin :: DataType
               -> Matrix I
               -> ByteString
matrixIntToBin DBL64B =
  writer putDouble (map fromIntegral)
matrixIntToBin INT32B =
  writer putInt id
matrixIntToBin INT08B = undefined


{-# LANGUAGE TypeOperators, CPP #-}
{-# LANGUAGE FlexibleInstances  #-}
{-# LANGUAGE FlexibleContexts   #-}
{-# LANGUAGE DefaultSignatures   #-}
{-# LANGUAGE ScopedTypeVariables #-}
----------------------------------------------------------------------
-- |
-- Module      :   Data.AdditiveGroup
-- Copyright   :  (c) Conal Elliott and Andy J Gill 2008
-- License     :  BSD3
--
-- Maintainer  :  conal@conal.net, andygill@ku.edu
-- Stability   :  experimental
--
-- Groups: zero, addition, and negation (additive inverse)
----------------------------------------------------------------------

module Data.AdditiveGroup
  (
    AdditiveGroup(..), sumV
  , Sum(..), inSum, inSum2
  ) where

import Prelude hiding (foldr)

import Control.Applicative
#if !(MIN_VERSION_base(4,8,0))
import Data.Monoid (Monoid(..))
import Data.Foldable (Foldable)
#endif
import Data.Foldable (foldr)
import Data.Complex hiding (magnitude)
import Data.Ratio
#if !(MIN_VERSION_base(4,11,0))
import Data.Semigroup (Semigroup(..))
#endif
import Foreign.C.Types (CSChar, CInt, CShort, CLong, CLLong, CIntMax, CFloat, CDouble)

import Data.MemoTrie

import Data.VectorSpace.Generic
import qualified GHC.Generics as Gnrx
import GHC.Generics (Generic, (:*:)(..))

infixl 6 ^+^, ^-^

-- | Additive group @v@.
class AdditiveGroup v where
  -- | The zero element: identity for '(^+^)'
  zeroV :: v
  default zeroV :: (Generic v, AdditiveGroup (VRep v)) => v
  zeroV = Gnrx.to (zeroV :: VRep v)
  -- | Add vectors
  (^+^) :: v -> v -> v
  default (^+^) :: (Generic v, AdditiveGroup (VRep v)) => v -> v -> v
  v ^+^ v' = Gnrx.to (Gnrx.from v ^+^ Gnrx.from v' :: VRep v)
  -- | Additive inverse
  negateV :: v -> v
  default negateV :: (Generic v, AdditiveGroup (VRep v)) => v -> v
  negateV v = Gnrx.to (negateV $ Gnrx.from v :: VRep v)
  -- | Group subtraction
  (^-^) :: v -> v -> v
  v ^-^ v' = v ^+^ negateV v'

-- | Sum over several vectors
sumV :: (Foldable f, AdditiveGroup v) => f v -> v
sumV = foldr (^+^) zeroV

instance AdditiveGroup () where
  zeroV     = ()
  () ^+^ () = ()
  negateV   = id

-- For 'Num' types:
--
-- instance AdditiveGroup n where {zeroV=0; (^+^) = (+); negateV = negate}

#define ScalarTypeCon(con,t) \
  instance con => AdditiveGroup (t) where {zeroV=0; (^+^) = (+); negateV = negate}

#define ScalarType(t) ScalarTypeCon((),t)

ScalarType(Int)
ScalarType(Integer)
ScalarType(Float)
ScalarType(Double)
ScalarType(CSChar)
ScalarType(CInt)
ScalarType(CShort)
ScalarType(CLong)
ScalarType(CLLong)
ScalarType(CIntMax)
ScalarType(CFloat)
ScalarType(CDouble)
ScalarTypeCon(Integral a,Ratio a)

instance (RealFloat v, AdditiveGroup v) => AdditiveGroup (Complex v) where
  zeroV   = zeroV :+ zeroV
  (^+^)   = (+)
  negateV = negate

-- Hm.  The 'RealFloat' constraint is unfortunate here.  It's due to a
-- questionable decision to place 'RealFloat' into the definition of the
-- 'Complex' /type/, rather than in functions and instances as needed.

instance (AdditiveGroup u,AdditiveGroup v) => AdditiveGroup (u,v) where
  zeroV             = (zeroV,zeroV)
  (u,v) ^+^ (u',v') = (u^+^u',v^+^v')
  negateV (u,v)     = (negateV u,negateV v)

instance (AdditiveGroup u,AdditiveGroup v,AdditiveGroup w)
    => AdditiveGroup (u,v,w) where
  zeroV                  = (zeroV,zeroV,zeroV)
  (u,v,w) ^+^ (u',v',w') = (u^+^u',v^+^v',w^+^w')
  negateV (u,v,w)        = (negateV u,negateV v,negateV w)

instance (AdditiveGroup u,AdditiveGroup v,AdditiveGroup w,AdditiveGroup x)
    => AdditiveGroup (u,v,w,x) where
  zeroV                       = (zeroV,zeroV,zeroV,zeroV)
  (u,v,w,x) ^+^ (u',v',w',x') = (u^+^u',v^+^v',w^+^w',x^+^x')
  negateV (u,v,w,x)           = (negateV u,negateV v,negateV w,negateV x)

-- Standard instance for an applicative functor applied to a vector space.
instance AdditiveGroup v => AdditiveGroup (a -> v) where
  zeroV   = pure   zeroV
  (^+^)   = liftA2 (^+^)
  negateV = fmap   negateV

-- Maybe is handled like the Maybe-of-Sum monoid
instance AdditiveGroup a => AdditiveGroup (Maybe a) where
  zeroV = Nothing
  Nothing ^+^ b'      = b'
  a' ^+^ Nothing      = a'
  Just a' ^+^ Just b' = Just (a' ^+^ b')
  negateV = fmap negateV

{-

Alexey Khudyakov wrote:

  I looked through vector-space package and found lawless instance. Namely Maybe's AdditiveGroup instance

  It's group so following relation is expected to hold. Otherwise it's not a group.
  > x ^+^ negateV x == zeroV

  Here is counterexample:

  > let x = Just 2 in x ^+^ negateV x == zeroV
  False

  I think it's not possible to sensibly define group instance for
  Maybe a at all.

I see that the problem here is in distinguishing 'Just zeroV' from
Nothing. I could fix the Just + Just line to use Nothing instead of Just
zeroV when a' ^+^ b' == zeroV, although doing so would require Eq a and
hence lose some generality. Even so, the abstraction leak would probably
show up elsewhere.

Hm.

-}

-- Memo tries
instance (HasTrie u, AdditiveGroup v) => AdditiveGroup (u :->: v) where
  zeroV   = pure   zeroV
  (^+^)   = liftA2 (^+^)
  negateV = fmap   negateV

-- | Monoid under group addition.  Alternative to the @Sum@ in
-- "Data.Monoid", which uses 'Num' instead of 'AdditiveGroup'.
newtype Sum a = Sum { getSum :: a }
  deriving (Eq, Ord, Read, Show, Bounded)

instance Functor Sum where
  fmap f (Sum a) = Sum (f a)

-- instance Applicative Sum where
--   pure a = Sum a
--   Sum f <*> Sum x = Sum (f x)

instance Applicative Sum where
  pure  = Sum
  (<*>) = inSum2 ($)

instance AdditiveGroup a => Semigroup (Sum a) where
  (<>) = liftA2 (^+^)

instance AdditiveGroup a => Monoid (Sum a) where
  mempty  = Sum zeroV
#if !(MIN_VERSION_base(4,11,0))
  mappend = (<>)
#endif

-- | Application a unary function inside a 'Sum'
inSum :: (a -> b) -> (Sum a -> Sum b)
inSum = getSum ~> Sum

-- | Application a binary function inside a 'Sum'
inSum2 :: (a -> b -> c) -> (Sum a -> Sum b -> Sum c)
inSum2 = getSum ~> inSum

instance AdditiveGroup a => AdditiveGroup (Sum a) where
  zeroV   = mempty
  (^+^)   = mappend
  negateV = inSum negateV

---- to go elsewhere

(~>) :: (a' -> a) -> (b -> b') -> ((a -> b) -> (a' -> b'))
(i ~> o) f = o . f . i

-- result :: (b -> b') -> ((a -> b) -> (a -> b'))
-- result = (.)

-- argument :: (a' -> a) -> ((a -> b) -> (a' -> b))
-- argument = flip (.)

-- g ~> f = result g . argument f

instance AdditiveGroup a => AdditiveGroup (Gnrx.Rec0 a s) where
  zeroV = Gnrx.K1 zeroV
  negateV (Gnrx.K1 v) = Gnrx.K1 $ negateV v
  Gnrx.K1 v ^+^ Gnrx.K1 w = Gnrx.K1 $ v ^+^ w
  Gnrx.K1 v ^-^ Gnrx.K1 w = Gnrx.K1 $ v ^-^ w
instance AdditiveGroup (f p) => AdditiveGroup (Gnrx.M1 i c f p) where
  zeroV = Gnrx.M1 zeroV
  negateV (Gnrx.M1 v) = Gnrx.M1 $ negateV v
  Gnrx.M1 v ^+^ Gnrx.M1 w = Gnrx.M1 $ v ^+^ w
  Gnrx.M1 v ^-^ Gnrx.M1 w = Gnrx.M1 $ v ^-^ w
instance (AdditiveGroup (f p), AdditiveGroup (g p)) => AdditiveGroup ((f :*: g) p) where
  zeroV = zeroV :*: zeroV
  negateV (x:*:y) = negateV x :*: negateV y
  (x:*:y) ^+^ (ξ:*:υ) = (x^+^ξ) :*: (y^+^υ)
  (x:*:y) ^-^ (ξ:*:υ) = (x^-^ξ) :*: (y^-^υ)


{-# LANGUAGE ForeignFunctionInterface #-}
{-# LANGUAGE FlexibleContexts         #-}
{-# LANGUAGE FlexibleInstances        #-}
{-# LANGUAGE BangPatterns             #-}
-----------------------------------------------------------------------------
-- |
-- Module      :  Data.Packed.Internal.Matrix
-- Copyright   :  (c) Alberto Ruiz 2007
-- License     :  GPL-style
--
-- Maintainer  :  Alberto Ruiz <aruiz@um.es>
-- Stability   :  provisional
-- Portability :  portable (uses FFI)
--
-- Internal matrix representation
--
-----------------------------------------------------------------------------
-- #hide

module Data.Packed.Internal.Matrix(
    Matrix(..), rows, cols, cdat, fdat,
    MatrixOrder(..), orderOf,
    createMatrix, mat,
    cmat, fmat,
    toLists, flatten, reshape,
    Element(..),
    trans,
    fromRows, toRows, fromColumns, toColumns,
    matrixFromVector,
    subMatrix,
    liftMatrix, liftMatrix2,
    (@@>), atM',
    saveMatrix,
    singleton,
    size, shSize, conformVs, conformMs, conformVTo, conformMTo
) where

import Data.Packed.Internal.Common
import Data.Packed.Internal.Signatures
import Data.Packed.Internal.Vector

import Foreign.Marshal.Alloc(alloca, free)
import Foreign.Marshal.Array(newArray)
import Foreign.Ptr(Ptr, castPtr)
import Foreign.Storable(Storable, peekElemOff, pokeElemOff, poke, sizeOf)
import Data.Complex(Complex)
import Foreign.C.Types
import Foreign.C.String(newCString)
import System.IO.Unsafe(unsafePerformIO)
import Control.DeepSeq

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

{- Design considerations for the Matrix Type
   -----------------------------------------

- we must easily handle both row major and column major order,
  for bindings to LAPACK and GSL/C

- we'd like to simplify redundant matrix transposes:
   - Some of them arise from the order requirements of some functions
   - some functions (matrix product) admit transposed arguments

- maybe we don't really need this kind of simplification:
   - more complex code
   - some computational overhead
   - only appreciable gain in code with a lot of redundant transpositions
     and cheap matrix computations

- we could carry both the matrix and its (lazily computed) transpose.
  This may save some transpositions, but it is necessary to keep track of the
  data which is actually computed to be used by functions like the matrix product
  which admit both orders.

- but if we need the transposed data and it is not in the structure, we must make
  sure that we touch the same foreignptr that is used in the computation.

- a reasonable solution is using two constructors for a matrix. Transposition just
  "flips" the constructor. Actual data transposition is not done if followed by a
  matrix product or another transpose.

-}

data MatrixOrder = RowMajor | ColumnMajor deriving (Show,Eq)

transOrder RowMajor = ColumnMajor
transOrder ColumnMajor = RowMajor
{- | Matrix representation suitable for GSL and LAPACK computations.

The elements are stored in a continuous memory array.

-}

data Matrix t = Matrix { irows :: {-# UNPACK #-} !Int
                       , icols :: {-# UNPACK #-} !Int
                       , xdat :: {-# UNPACK #-} !(Vector t)
                       , order :: !MatrixOrder }
-- RowMajor: preferred by C, fdat may require a transposition
-- ColumnMajor: preferred by LAPACK, cdat may require a transposition

cdat = xdat
fdat = xdat

rows :: Matrix t -> Int
rows = irows

cols :: Matrix t -> Int
cols = icols

orderOf :: Matrix t -> MatrixOrder
orderOf = order

-- | Matrix transpose.
trans :: Matrix t -> Matrix t
trans Matrix {irows = r, icols = c, xdat = d, order = o } = Matrix { irows = c, icols = r, xdat = d, order = transOrder o}

cmat :: (Element t) => Matrix t -> Matrix t
cmat m@Matrix{order = RowMajor} = m
cmat Matrix {irows = r, icols = c, xdat = d, order = ColumnMajor } = Matrix { irows = r, icols = c, xdat = transdata r d c, order = RowMajor}

fmat :: (Element t) => Matrix t -> Matrix t
fmat m@Matrix{order = ColumnMajor} = m
fmat Matrix {irows = r, icols = c, xdat = d, order = RowMajor } = Matrix { irows = r, icols = c, xdat = transdata c d r, order = ColumnMajor}

-- C-Haskell matrix adapter
-- mat :: Adapt (CInt -> CInt -> Ptr t -> r) (Matrix t) r

mat :: (Storable t) => Matrix t -> (((CInt -> CInt -> Ptr t -> t1) -> t1) -> IO b) -> IO b
mat a f =
    unsafeWith (xdat a) $ \p -> do
        let m g = do
            g (fi (rows a)) (fi (cols a)) p
        f m
-- | Creates a vector by concatenation of rows. If the matrix is ColumnMajor, this operation requires a transpose.
--
-- @\> flatten ('ident' 3)
-- 9 |> [1.0,0.0,0.0,0.0,1.0,0.0,0.0,0.0,1.0]@
flatten :: Element t => Matrix t -> Vector t
flatten = xdat . cmat

type Mt t s = Int -> Int -> Ptr t -> s
-- not yet admitted by my haddock version
-- infixr 6 ::>
-- type t ::> s = Mt t s

-- | the inverse of 'Data.Packed.Matrix.fromLists'
toLists :: (Element t) => Matrix t -> [[t]]
toLists m = splitEvery (cols m) . toList . flatten $ m

-- | Create a matrix from a list of vectors.
-- All vectors must have the same dimension,
-- or dimension 1, which is are automatically expanded.
fromRows :: Element t => [Vector t] -> Matrix t
fromRows vs = case compatdim (map dim vs) of
    Nothing -> error "fromRows applied to [] or to vectors with different sizes"
    Just c  -> reshape c . join . map (adapt c) $ vs
  where
    adapt c v | dim v == c = v
              | otherwise = constantD (v@>0) c

-- | extracts the rows of a matrix as a list of vectors
toRows :: Element t => Matrix t -> [Vector t]
toRows m = toRows' 0 where
    v = flatten m
    r = rows m
    c = cols m
    toRows' k | k == r*c  = []
              | otherwise = subVector k c v : toRows' (k+c)

-- | Creates a matrix from a list of vectors, as columns
fromColumns :: Element t => [Vector t] -> Matrix t
fromColumns m = trans . fromRows $ m

-- | Creates a list of vectors from the columns of a matrix
toColumns :: Element t => Matrix t -> [Vector t]
toColumns m = toRows . trans $ m

-- | Reads a matrix position.
(@@>) :: Storable t => Matrix t -> (Int,Int) -> t
infixl 9 @@>
m@Matrix {irows = r, icols = c} @@> (i,j)
    | safe      = if i<0 || i>=r || j<0 || j>=c
                    then error "matrix indexing out of range"
                    else atM' m i j
    | otherwise = atM' m i j
{-# INLINE (@@>) #-}

--  Unsafe matrix access without range checking
atM' Matrix {icols = c, xdat = v, order = RowMajor} i j = v `at'` (i*c+j)
atM' Matrix {irows = r, xdat = v, order = ColumnMajor} i j = v `at'` (j*r+i)
{-# INLINE atM' #-}

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

matrixFromVector o c v = Matrix { irows = r, icols = c, xdat = v, order = o }
    where (d,m) = dim v `quotRem` c
          r | m==0 = d
            | otherwise = error "matrixFromVector"

-- allocates memory for a new matrix
createMatrix :: (Storable a) => MatrixOrder -> Int -> Int -> IO (Matrix a)
createMatrix ord r c = do
    p <- createVector (r*c)
    return (matrixFromVector ord c p)

{- | Creates a matrix from a vector by grouping the elements in rows with the desired number of columns. (GNU-Octave groups by columns. To do it you can define @reshapeF r = trans . reshape r@
where r is the desired number of rows.)

@\> reshape 4 ('fromList' [1..12])
(3><4)
 [ 1.0,  2.0,  3.0,  4.0
 , 5.0,  6.0,  7.0,  8.0
 , 9.0, 10.0, 11.0, 12.0 ]@

-}
reshape :: Storable t => Int -> Vector t -> Matrix t
reshape c v = matrixFromVector RowMajor c v

singleton x = reshape 1 (fromList [x])

-- | application of a vector function on the flattened matrix elements
liftMatrix :: (Storable a, Storable b) => (Vector a -> Vector b) -> Matrix a -> Matrix b
liftMatrix f Matrix { icols = c, xdat = d, order = o } = matrixFromVector o c (f d)

-- | application of a vector function on the flattened matrices elements
liftMatrix2 :: (Element t, Element a, Element b) => (Vector a -> Vector b -> Vector t) -> Matrix a -> Matrix b -> Matrix t
liftMatrix2 f m1 m2
    | not (compat m1 m2) = error "nonconformant matrices in liftMatrix2"
    | otherwise = case orderOf m1 of
        RowMajor    -> matrixFromVector RowMajor    (cols m1) (f (xdat m1) (flatten m2))
        ColumnMajor -> matrixFromVector ColumnMajor (cols m1) (f (xdat m1) ((xdat.fmat) m2))

compat :: Matrix a -> Matrix b -> Bool
compat m1 m2 = rows m1 == rows m2 && cols m1 == cols m2

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

{- | Supported matrix elements.

    This class provides optimized internal
    operations for selected element types.
    It provides unoptimised defaults for any 'Storable' type,
    so you can create instances simply as:
    @instance Element Foo@.
-}
class (Storable a) => Element a where
    subMatrixD :: (Int,Int) -- ^ (r0,c0) starting position
               -> (Int,Int) -- ^ (rt,ct) dimensions of submatrix
               -> Matrix a -> Matrix a
    subMatrixD = subMatrix'
    transdata :: Int -> Vector a -> Int -> Vector a
    transdata = transdataP -- transdata'
    constantD  :: a -> Int -> Vector a
    constantD = constantP -- constant'

instance Element Float where
    transdata  = transdataAux ctransF
    constantD  = constantAux cconstantF

instance Element Double where
    transdata  = transdataAux ctransR
    constantD  = constantAux cconstantR

instance Element (Complex Float) where
    transdata  = transdataAux ctransQ
    constantD  = constantAux cconstantQ

instance Element (Complex Double) where
    transdata  = transdataAux ctransC
    constantD  = constantAux cconstantC

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

transdata' :: Storable a => Int -> Vector a -> Int -> Vector a
transdata' c1 v c2 =
    if noneed
        then v
        else unsafePerformIO $ do
                w <- createVector (r2*c2)
                unsafeWith v $ \p ->
                    unsafeWith w $ \q -> do
                        let go (-1) _ = return ()
                            go !i (-1) = go (i-1) (c1-1)
                            go !i !j = do x <- peekElemOff p (i*c1+j)
                                          pokeElemOff      q (j*c2+i) x
                                          go i (j-1)
                        go (r1-1) (c1-1)
                return w
  where r1 = dim v `div` c1
        r2 = dim v `div` c2
        noneed = r1 == 1 || c1 == 1

-- {-# SPECIALIZE transdata' :: Int -> Vector Double -> Int ->  Vector Double #-}
-- {-# SPECIALIZE transdata' :: Int -> Vector (Complex Double) -> Int -> Vector (Complex Double) #-}

-- I don't know how to specialize...
-- The above pragmas only seem to work on top level defs
-- Fortunately everything seems to work using the above class

-- C versions, still a little faster:

transdataAux fun c1 d c2 =
    if noneed
        then d
        else unsafePerformIO $ do
            v <- createVector (dim d)
            unsafeWith d $ \pd ->
                unsafeWith v $ \pv ->
                    fun (fi r1) (fi c1) pd (fi r2) (fi c2) pv // check "transdataAux"
            return v
  where r1 = dim d `div` c1
        r2 = dim d `div` c2
        noneed = r1 == 1 || c1 == 1

transdataP :: Storable a => Int -> Vector a -> Int -> Vector a
transdataP c1 d c2 =
    if noneed
       then d
       else unsafePerformIO $ do
          v <- createVector (dim d)
          unsafeWith d $ \pd ->
              unsafeWith v $ \pv ->
                  ctransP (fi r1) (fi c1) (castPtr pd) (fi sz) (fi r2) (fi c2) (castPtr pv) (fi sz) // check "transdataP"
          return v
   where r1 = dim d `div` c1
         r2 = dim d `div` c2
         sz = sizeOf (d @> 0)
         noneed = r1 == 1 || c1 == 1

foreign import ccall unsafe "transF" ctransF :: TFMFM
foreign import ccall unsafe "transR" ctransR :: TMM
foreign import ccall unsafe "transQ" ctransQ :: TQMQM
foreign import ccall unsafe "transC" ctransC :: TCMCM
foreign import ccall unsafe "transP" ctransP :: CInt -> CInt -> Ptr () -> CInt -> CInt -> CInt -> Ptr () -> CInt -> IO CInt

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

constant' v n = unsafePerformIO $ do
    w <- createVector n
    unsafeWith w $ \p -> do
        let go (-1) = return ()
            go !k = pokeElemOff p k v >> go (k-1)
        go (n-1)
    return w

-- C versions

constantAux fun x n = unsafePerformIO $ do
    v <- createVector n
    px <- newArray [x]
    app1 (fun px) vec v "constantAux"
    free px
    return v

constantF :: Float -> Int -> Vector Float
constantF = constantAux cconstantF
foreign import ccall unsafe "constantF" cconstantF :: Ptr Float -> TF

constantR :: Double -> Int -> Vector Double
constantR = constantAux cconstantR
foreign import ccall unsafe "constantR" cconstantR :: Ptr Double -> TV

constantQ :: Complex Float -> Int -> Vector (Complex Float)
constantQ = constantAux cconstantQ
foreign import ccall unsafe "constantQ" cconstantQ :: Ptr (Complex Float) -> TQV

constantC :: Complex Double -> Int -> Vector (Complex Double)
constantC = constantAux cconstantC
foreign import ccall unsafe "constantC" cconstantC :: Ptr (Complex Double) -> TCV

constantP :: Storable a => a -> Int -> Vector a
constantP a n = unsafePerformIO $ do
    let sz = sizeOf a
    v <- createVector n
    unsafeWith v $ \p -> do
       alloca $ \k -> do
                      poke k a
                      cconstantP (castPtr k) (fi n) (castPtr p) (fi sz) // check "constantP"
    return v
foreign import ccall unsafe "constantP" cconstantP :: Ptr () -> CInt -> Ptr () -> CInt -> IO CInt

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

-- | Extracts a submatrix from a matrix.
subMatrix :: Element a
          => (Int,Int) -- ^ (r0,c0) starting position
          -> (Int,Int) -- ^ (rt,ct) dimensions of submatrix
          -> Matrix a -- ^ input matrix
          -> Matrix a -- ^ result
subMatrix (r0,c0) (rt,ct) m
    | 0 <= r0 && 0 < rt && r0+rt <= (rows m) &&
      0 <= c0 && 0 < ct && c0+ct <= (cols m) = subMatrixD (r0,c0) (rt,ct) m
    | otherwise = error $ "wrong subMatrix "++
                          show ((r0,c0),(rt,ct))++" of "++show(rows m)++"x"++ show (cols m)

subMatrix'' (r0,c0) (rt,ct) c v = unsafePerformIO $ do
    w <- createVector (rt*ct)
    unsafeWith v $ \p ->
        unsafeWith w $ \q -> do
            let go (-1) _ = return ()
                go !i (-1) = go (i-1) (ct-1)
                go !i !j = do x <- peekElemOff p ((i+r0)*c+j+c0)
                              pokeElemOff      q (i*ct+j) x
                              go i (j-1)
            go (rt-1) (ct-1)
    return w

subMatrix' (r0,c0) (rt,ct) (Matrix { icols = c, xdat = v, order = RowMajor}) = Matrix rt ct (subMatrix'' (r0,c0) (rt,ct) c v) RowMajor
subMatrix' (r0,c0) (rt,ct) m = trans $ subMatrix' (c0,r0) (ct,rt) (trans m)

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

-- | Saves a matrix as 2D ASCII table.
saveMatrix :: FilePath
           -> String     -- ^ format (%f, %g, %e)
           -> Matrix Double
           -> IO ()
saveMatrix filename fmt m = do
    charname <- newCString filename
    charfmt <- newCString fmt
    let o = if orderOf m == RowMajor then 1 else 0
    app1 (matrix_fprintf charname charfmt o) mat m "matrix_fprintf"
    free charname
    free charfmt

foreign import ccall unsafe "matrix_fprintf" matrix_fprintf :: Ptr CChar -> Ptr CChar -> CInt -> TM

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

conformMs ms = map (conformMTo (r,c)) ms
  where
    r = maximum (map rows ms)
    c = maximum (map cols ms)

conformVs vs = map (conformVTo n) vs
  where
    n = maximum (map dim vs)

conformMTo (r,c) m
    | size m == (r,c) = m
    | size m == (1,1) = reshape c (constantD (m@@>(0,0)) (r*c))
    | size m == (r,1) = repCols c m
    | size m == (1,c) = repRows r m
    | otherwise = error $ "matrix " ++ shSize m ++ " cannot be expanded to (" ++ show r ++ "><"++ show c ++")"

conformVTo n v
    | dim v == n = v
    | dim v == 1 = constantD (v@>0) n
    | otherwise = error $ "vector of dim=" ++ show (dim v) ++ " cannot be expanded to dim=" ++ show n

repRows n x = fromRows (replicate n (flatten x))
repCols n x = fromColumns (replicate n (flatten x))

size m = (rows m, cols m)

shSize m = "(" ++ show (rows m) ++"><"++ show (cols m)++")"

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

instance (Storable t, NFData t) => NFData (Matrix t)
  where
    rnf m | d > 0     = rnf (v @> 0)
          | otherwise = ()
      where
        d = dim v
        v = xdat m


{-# LANGUAGE FlexibleInstances     #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE PackageImports        #-}

module FPNLA.Matrix.Instances.HMatrix () where

import           Control.DeepSeq               (NFData (rnf))
import           Control.Parallel.Strategies   (rseq, withStrategy)
import "hmatrix" Numeric.LinearAlgebra hiding  (Matrix, Vector)
--import           Numeric.LinearAlgebra.HMatrix
--import "hmatrix" Numeric.LinearAlgebra         (buildVector)
import           Numeric.LinearAlgebra.HMatrix hiding (Matrix, Vector)
import qualified Numeric.LinearAlgebra.HMatrix  as HM (Matrix, Vector)
import           FPNLA.Matrix                  (Matrix (generate_m, fromList_m, transpose_m, dim_m, elem_m, subMatrix_m, fromBlocks_m, toBlocks_m, map_m), MatrixVector (fromCols_vm, toCols_vm, row_vm, col_vm), Vector (generate_v, fromList_v, concat_v, elem_v, length_v, foldr_v, map_v, zipWith_v))
{-
instance (NFData e) => NFData (Vector e) where
    -- Asumo que Vector es estricto
    -- http://haskell.1045720.n5.nabble.com/NFData-instance-for-Numeric-LinearAlgebra-Matrix-td4265725.html
    rnf v = (withStrategy rseq v) `seq` ()

instance (NFData e) => NFData (Matrix e) where
    -- Matrix es estricto
    -- http://haskell.1045720.n5.nabble.com/NFData-instance-for-Numeric-LinearAlgebra-Matrix-td4265725.html
    rnf m = withStrategy rseq m `seq` ()
-}

buildVector :: Int -> (Int -> e) -> HM.Vector e
buildVector size f = undefined

buildMatrix :: Int -> Int -> (Int -> Int -> e) -> HM.Matrix e
buildMatrix rows cols f = undefined

instance (Element e) => Vector HM.Vector e where
    generate_v = buildVector
    fromList_v = fromList
    concat_v = undefined -- joinsas
    elem_v pos v = undefined -- v @> pos
    length_v = undefined -- dim
    foldr_v = undefined -- foldVector
    map_v = undefined -- mapVector
    zipWith_v = undefined -- zipVectorWith

instance (Element e) => Matrix HM.Matrix e where
    generate_m = buildMatrix
    fromList_m = (><)
    transpose_m = undefined -- trans
    dim_m m = (rows m, cols m)
    elem_m i j m = undefined -- m @@> (i, j)
    map_m = undefined -- mapMatrix
    --zipWith_m
    subMatrix_m posI posJ cantRows cantCols = subMatrix (posI, posJ) (cantRows, cantCols)
    fromBlocks_m = fromBlocks
    toBlocks_m = toBlocksEvery

instance (Element e) => MatrixVector HM.Matrix HM.Vector e where
    row_vm pos m = toRows m !! pos
    col_vm pos m = toColumns m !! pos
    fromCols_vm = fromColumns
    toCols_vm = toColumns


{-# LANGUAGE AllowAmbiguousTypes                      #-}
{-# LANGUAGE ConstraintKinds                          #-}
{-# LANGUAGE DataKinds                                #-}
{-# LANGUAGE DefaultSignatures                        #-}
{-# LANGUAGE FlexibleContexts                         #-}
{-# LANGUAGE FlexibleInstances                        #-}
{-# LANGUAGE GADTs                                    #-}
{-# LANGUAGE RankNTypes                               #-}
{-# LANGUAGE ScopedTypeVariables                      #-}
{-# LANGUAGE TypeApplications                         #-}
{-# LANGUAGE TypeFamilies                             #-}
{-# LANGUAGE TypeOperators                            #-}
{-# LANGUAGE UndecidableInstances                     #-}
{-# OPTIONS_GHC -fplugin GHC.TypeLits.KnownNat.Solver #-}
{-# OPTIONS_GHC -fplugin GHC.TypeLits.Normalise       #-}

module Backprop.Learn.Initialize (
    Initialize(..)
  , gInitialize
  , initializeNormal
  , initializeSingle
  -- * Reshape
  , reshapeR
  , reshapeLRows
  , reshapeLCols
  ) where

import           Control.Monad.Primitive
import           Data.Complex
import           Data.Proxy
import           Data.Type.Equality
import           Data.Type.Tuple
import           Data.Vinyl
import           GHC.TypeLits.Compare
import           GHC.TypeNats
import           Generics.OneLiner
import           Numeric.LinearAlgebra.Static.Vector
import           Statistics.Distribution
import           Statistics.Distribution.Normal
import qualified Data.Vector.Generic                 as VG
import qualified Data.Vector.Generic.Sized           as SVG
import qualified Numeric.LinearAlgebra.Static        as H
import qualified System.Random.MWC                   as MWC

-- | Class for types that are basically a bunch of 'Double's, which can be
-- initialized with a given identical and independent distribution.
class Initialize p where
    initialize
        :: (ContGen d, PrimMonad m)
        => d
        -> MWC.Gen (PrimState m)
        -> m p

    default initialize
        :: (ADTRecord p, Constraints p Initialize, ContGen d, PrimMonad m)
        => d
        -> MWC.Gen (PrimState m)
        -> m p
    initialize = gInitialize

-- | 'initialize' for any instance of 'Generic'.
gInitialize
    :: (ADTRecord p, Constraints p Initialize, ContGen d, PrimMonad m)
    => d
    -> MWC.Gen (PrimState m)
    -> m p
gInitialize d g = createA' @Initialize (initialize d g)

-- | Helper over 'inititialize' for a gaussian distribution centered around
-- zero.
initializeNormal
    :: (Initialize p, PrimMonad m)
    => Double                               -- ^ standard deviation
    -> MWC.Gen (PrimState m)
    -> m p
initializeNormal = initialize . normalDistr 0

-- | 'initialize' definition if @p@ is a single number.
initializeSingle
    :: (ContGen d, PrimMonad m, Fractional p)
    => d
    -> MWC.Gen (PrimState m)
    -> m p
initializeSingle d = fmap realToFrac . genContVar d

instance Initialize Double where
    initialize = initializeSingle
instance Initialize Float where
    initialize = initializeSingle

-- | Initializes real and imaginary components identically
instance Initialize a => Initialize (Complex a) where

instance Initialize T0
instance Initialize a => Initialize (TF a)
instance (Initialize a, Initialize b) => Initialize (a :# b)

instance RPureConstrained Initialize as => Initialize (T as) where
    initialize d g = rtraverse (fmap TF)
                   $ rpureConstrained @Initialize (initialize d g)

-- instance (Initialize a, ListC (Initialize <$> as), Known Length as) => Initialize (NETup (a ':| as)) where
--     initialize d g = NET <$> initialize d g
--                          <*> initialize d g

instance Initialize ()
instance (Initialize a, Initialize b) => Initialize (a, b)
instance (Initialize a, Initialize b, Initialize c) => Initialize (a, b, c)
instance (Initialize a, Initialize b, Initialize c, Initialize d) => Initialize (a, b, c, d)
instance (Initialize a, Initialize b, Initialize c, Initialize d, Initialize e) => Initialize (a, b, c, d, e)

instance (VG.Vector v a, KnownNat n, Initialize a) => Initialize (SVG.Vector v n a) where
    initialize d = SVG.replicateM . initialize d

instance KnownNat n => Initialize (H.R n) where
    initialize d = fmap vecR . initialize d
instance KnownNat n => Initialize (H.C n) where
    initialize d = fmap vecC . initialize d

instance (KnownNat n, KnownNat m) => Initialize (H.L n m) where
    initialize d = fmap vecL . initialize d
instance (KnownNat n, KnownNat m) => Initialize (H.M n m) where
    initialize d = fmap vecM . initialize d

-- | Reshape a vector to have a different amount of items  If the matrix is
-- grown, new weights are initialized according to the given distribution.
reshapeR
    :: forall i j d m. (ContGen d, PrimMonad m, KnownNat i, KnownNat j)
    => d
    -> MWC.Gen (PrimState m)
    -> H.R i
    -> m (H.R j)
reshapeR d g x = case Proxy @j %<=? Proxy @i of
    LE  Refl      -> pure . vecR . SVG.take @_ @j @(i - j) . rVec $ x
    NLE Refl Refl -> (x H.#) <$> initialize @(H.R (j - i)) d g

-- | Reshape a matrix to have a different amount of rows  If the matrix
-- is grown, new weights are initialized according to the given
-- distribution.
reshapeLRows
    :: forall i j n d m. (ContGen d, PrimMonad m, KnownNat n, KnownNat i, KnownNat j)
    => d
    -> MWC.Gen (PrimState m)
    -> H.L i n
    -> m (H.L j n)
reshapeLRows d g x = case Proxy @j %<=? Proxy @i of
    LE Refl       -> pure . rowsL . SVG.take @_ @j @(i - j) . lRows $ x
    NLE Refl Refl -> (x H.===) <$> initialize @(H.L (j - i) n) d g

-- | Reshape a matrix to have a different amount of columns.  If the matrix
-- is grown, new weights are initialized according to the given
-- distribution.
reshapeLCols
    :: forall i j n d m. (ContGen d, PrimMonad m, KnownNat n, KnownNat i, KnownNat j)
    => d
    -> MWC.Gen (PrimState m)
    -> H.L n i
    -> m (H.L n j)
reshapeLCols d g x = case Proxy @j %<=? Proxy @i of
    LE Refl       -> pure . colsL . SVG.take @_ @j @(i - j) . lCols $ x
    NLE Refl Refl -> (x H.|||) <$> initialize @(H.L n (j - i)) d g


module Mnist where

  import Numeric.LinearAlgebra
  import System.Environment

  import Neural
  import Test
  import Mnist.Load

  main :: IO ()
  main = do
    mnist <- loadExamples
    learningRate <- (getArgs >>= (\args -> (return.read) $ head args))
    epoch <- (getArgs >>= (\args -> (return.read) $ args!!1))
    let dataSet = splitDataSet mnist
    runSuite epoch learningRate networkSetting dataSet

  -- one hidden layer with 300 nodes
  networkSetting = [10, 300]


module TwoMoon where

  import Numeric.LinearAlgebra
  import System.Environment

  import Neural
  import Test
  import TwoMoon.Load

  main :: IO ()
  main = do
    twomoon <- loadExamples
    learningRate <- (getArgs >>= (\args -> (return.read) $ head args))
    epoch <- (getArgs >>= (\args -> (return.read) $ args!!1))
    let dataSet = splitDataSet twomoon
    runSuite epoch learningRate networkSetting dataSet

  -- two hidden layer with 20, 10 nodes each
  networkSetting = [2, 4]


{-# LANGUAGE DataKinds             #-}
{-# LANGUAGE DeriveAnyClass        #-}
{-# LANGUAGE DeriveGeneric         #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE OverloadedLabels      #-}
{-# LANGUAGE OverloadedStrings     #-}
{-# LANGUAGE RankNTypes            #-}
{-# LANGUAGE ScopedTypeVariables   #-}
{-# LANGUAGE TypeFamilies          #-}
{-# LANGUAGE TypeOperators         #-}
{-|
Module      : Grenade.Layers.LeakyRelu
Description : Rectifying linear unit layer
Copyright   : (c) Manuel Schneckenreither, 2020
License     : BSD2
Stability   : experimental
-}
module Grenade.Layers.LeakyRelu (
    LeakyRelu (..)
  ) where

import           Control.DeepSeq                (NFData (..))
import           Data.Proxy                     (Proxy(Proxy))
import           Data.Maybe                     (fromJust)
import           Data.Serialize
import           GHC.Generics                   (Generic)
import           GHC.TypeLits
import qualified Numeric.LinearAlgebra.Static   as H

import           Grenade.Core
import           Grenade.Onnx
import           Grenade.Types
import           Grenade.Layers.Internal.Activations

-- | A rectifying linear unit.
--   A layer which can act between any shape of the same dimension, acting as a
--   diode on every neuron individually.
newtype LeakyRelu = LeakyRelu RealNum
  deriving (Generic, NFData, Show)

instance UpdateLayer LeakyRelu where
  type Gradient LeakyRelu = ()
  runUpdate _ x _ = x
  reduceGradient _ = ()

instance RandomLayer LeakyRelu where
  createRandomWith _ _ = return $ LeakyRelu 0.01

instance Serialize LeakyRelu where
  put (LeakyRelu alpha) = put alpha
  get = LeakyRelu <$> get

instance (KnownNat i) => Layer LeakyRelu ('D1 i) ('D1 i) where
  type Tape LeakyRelu ('D1 i) ('D1 i) = S ('D1 i)

  runForwards (LeakyRelu alpha) (S1D y) = (S1D y, (S1D . fromJust . H.create) relu)
    where
      w    = fromIntegral $ natVal (Proxy :: Proxy i)
      relu = leakyRelu1d w alpha (H.extract y)

  runBackwards (LeakyRelu alpha) (S1D y) (S1D dEdy) = ((), S1D (relu' y * dEdy))
    where
      relu' = H.dvmap (\a -> if a < 0 then alpha else 1)

instance (KnownNat i, KnownNat j) => Layer LeakyRelu ('D2 i j) ('D2 i j) where
  type Tape LeakyRelu ('D2 i j) ('D2 i j) = S ('D2 i j)

  runForwards (LeakyRelu alpha) (S2D y) = (S2D y, (S2D . fromJust . H.create) relu)
    where
      h    = fromIntegral $ natVal (Proxy :: Proxy i)
      w    = fromIntegral $ natVal (Proxy :: Proxy j)
      y'   = H.extract y
      relu = leakyRelu 1 h w alpha y'

  runBackwards (LeakyRelu alpha) (S2D y) (S2D dEdy) = ((), S2D (relu' y * dEdy))
    where
      relu' = H.dmmap (\a -> if a < 0 then alpha else 1)

instance (KnownNat i, KnownNat j, KnownNat k) => Layer LeakyRelu ('D3 i j k) ('D3 i j k) where

  type Tape LeakyRelu ('D3 i j k) ('D3 i j k) = S ('D3 i j k)

  runForwards (LeakyRelu alpha) (S3D y) = (S3D y, (S3D . fromJust . H.create) relu)
    where
      c    = fromIntegral $ natVal (Proxy :: Proxy k)
      h    = fromIntegral $ natVal (Proxy :: Proxy i)
      w    = fromIntegral $ natVal (Proxy :: Proxy j)
      y'   = H.extract y
      relu = leakyRelu c h w alpha y'

  runBackwards (LeakyRelu alpha) (S3D y) (S3D dEdy) = ((), S3D (relu' y * dEdy))
    where
      relu' = H.dmmap (\a -> if a < 0 then alpha else 1)

instance OnnxOperator LeakyRelu where
  onnxOpTypeNames _ = ["LeakyRelu"]

instance OnnxLoadable LeakyRelu where
  loadOnnxNode _ node = do
    alpha <- readFloatAttributeToRealNum "alpha" node

    return $ LeakyRelu alpha


module Scene.Material.Skylike where

import Numeric.LinearAlgebra

import Scene.Base
import Math.Color
import Math.Ray
import Scene.Monad

import Control.Applicative
import Control.Monad.Random

skylike :: Color -> Color -> Material
skylike near far (Ray p normal) (Ray q incidence) = do
    let t = abs $ (normal <.> incidence)/(norm_2 normal * norm_2 incidence)
    return $! weigh [(t,near),(1-t,far)]


{-# LANGUAGE ConstraintKinds        #-}
{-# LANGUAGE FlexibleContexts       #-}
{-# LANGUAGE FlexibleInstances      #-}
{-# LANGUAGE FunctionalDependencies #-}
{-# LANGUAGE GADTs                  #-}
{-# LANGUAGE InstanceSigs           #-}
{-# LANGUAGE MultiParamTypeClasses  #-}
{-# LANGUAGE QuantifiedConstraints  #-}
{-# LANGUAGE RankNTypes             #-}
{-# LANGUAGE TemplateHaskell        #-}
{-# LANGUAGE TypeOperators          #-}
{-# LANGUAGE UndecidableInstances   #-}

module Q.Stats.Arima where
import           Control.Monad.State
import           Data.Foldable
import           Data.Functor.Identity
import           Data.Random
import           Data.Random.Source
import           Data.RVar
import           Data.Time
import           Numeric.LinearAlgebra
import           Q.Stats.TimeSeries
import           System.Random.Mersenne.Pure64
import           Data.Random.Distribution
import           Data.Random.Distribution.Poisson
import           Data.Random.Distribution.T
import           Data.RVar
import           Statistics.Sample

data Ewma d = Ewma Double d

--ll :: (Ewma d) -> [DataPoint Double] -> (Double -> Double)
ll (Ewma lambda d) datapoints = mapM ll_ datapoints where
  ll_ :: DataPoint LocalTime Double -> State Double Double
  ll_ x@(DataPoint _ v) = do
    vart <- get
    let vart2 = lambda * vart + (1 - lambda) * v * v
    put vart2
    return $ logPdf d (sqrt (v  * v / vart))

--forecast :: (Distribution d Double) => (Ewma d) -> Int ->
forecast :: forall d. (Distribution d Double) => Ewma (d Double) -> StateT Double RVar Double
forecast (Ewma lambda d) = do
  y <- lift $ rvar d
  vart <- get
  let vart2 = lambda * vart + (1 - lambda) * y * y
  put vart2
  return (y * sqrt vart)

--forecastN :: Distribution d Double => Ewma (d Double) -> Int -> Double -> RVar ([Double], Double)
forecastN ewma var0 n =  sample $ runStateT (replicateM n (forecast ewma)) var0


{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE OverloadedStrings #-}
-- | Various tests
module Main where

import AI.Network.RNN.LSTM
import AI.Network.RNN.RNN
import AI.Network.RNN.Genetic
import AI.Network.RNN.Data
import AI.Network.RNN.Types
import AI.Network.RNN.Util
import AI.Network.RNN.Expr

import Test.Tasty
import Test.Tasty.HUnit
import Test.Tasty.QuickCheck

import qualified Numeric.LinearAlgebra.HMatrix as M

import Control.Monad.Random
import qualified Data.Text as T
import qualified Data.Set as DS

import Data.Char
import Data.List
import Numeric

import Debug.Trace

import Numeric.AD
import qualified Data.IntMap as I
import Text.PrettyPrint.HughesPJClass

main :: IO()
main = defaultMain tests

tests :: TestTree
tests = testGroup "Tests"
    [  testGroup "Utils" [
          testProperty "Matrix/vector product" prop_product
        , testCase "takes" $ do
            let l=[1,2,3,4,5,6]::[Int]
            [[1],[2,3],[4,5,6]] @=? takes ([1,2,3]::[Int]) l
            [[],[],[]] @=? takes [1,2,3] ([]::[Int])
        , testProperty "takes keep all data in order" prop_takes_concat
        , testProperty "takes has proper number of lists" prop_takes_length
        , testProperty "binary digits" prop_binary_digits
        , testCase "softmax" $ do
            let l=[1,2,3,4,5,6]::[Double]
                sf=softmax l
            sum sf @?= 1
        , testProperty "euclidian" prop_euclidian
        , testProperty "equilateral size" prop_equilateral_size
        , testProperty "equilateral distance" prop_equilateral_distance
        , testProperty "equilateral decoding" prop_equilateral_decoding
       ]
     , testGroup "RNN" [
        testCase "Check Steps without Back" $ checkSteps False
        , testCase "Check Steps with Back" $ checkSteps True
        , testCase "Check array conversion without Back" $ checkArray False
        , testCase "Check array conversion with Back" $ checkArray True
        , testCase "Check vector conversion without Back" $ checkArray False
        , testCase "Check vector conversion with Back" $ checkArray True
        ]
     , testGroup "Data" [
        testCase "Text" $ do
            testTextData "hello"
            testTextData "hello world!"
        , testCase "Text Sparse" $ do
            testTextDataS "hello"
            testTextDataS "hello world!"
     , testGroup "Genetic" [
            testCase "MixVector" $ do
                let v1 = M.fromList [1,1,1,1]
                let v2 = M.fromList [2,2,2,2]
                (v3,v4) <- evalRandIO $ mixVector v1 v2 0.5
                (v3 /= v4) @? "v3 == v4"
                M.size v1 @=? M.size v3
                M.size v1 @=? M.size v4
            , testCase "crossNetworkFull" $ testCrossover crossNetworkFull
            , testCase "crossNetworkHalf" $ testCrossover crossNetworkHalf
            , testCase "pointMutation" $ testMutation pointMutation
            , testCase "swapMutation" $ testMutation swapMutation
            , testCase "insertMutation" $ testMutation insertMutation
        ]
      , testGroup "LSTM" [
            testCase "Check steps" checkLSTMSteps
          , testProperty "Property steps" prop_lstm_eval_steps
          , testCase "Check vector conversion" checkLSTMVector
          , testProperty "list and vector match" prop_lstm_step
          , testCase "Check list of LSTMs" checkLSTMList
          , testProperty "Property list" prop_lstm_list
        ]
       , testGroup "LSTMIO" [
            testCase "LSTMIO basics" checkLSTMIO
        ]
       , testGroup "Expr" [
            testCase "Expr basics" testExpr
        ]
     ]
    ]

prop_product :: MatrixVector -> Bool
prop_product (MatrixVector sz ms vs) =
    let l=listMProd ms vs
        m=M.toList (M.matrix sz ms M.#> M.vector vs)
    in l == m -- (trace (show l) l)  == (trace (show m) m)

data MatrixVector = MatrixVector Int [Double] [Double]
    deriving (Show,Read,Eq,Ord)

instance Arbitrary MatrixVector where
    arbitrary = do
        Positive rows<-arbitrary
        Positive cols<-arbitrary
        ms <- vector (rows*cols)
        vs <- vector cols
        return $ MatrixVector cols ms vs

prop_takes_concat ::  String -> [Positive Int] -> Bool
prop_takes_concat xs ps = let
    idxs = map getPositive ps
    tot = sum idxs
    in concat (takes idxs xs) == take tot xs

prop_takes_length ::  String -> [Positive Int] -> Bool
prop_takes_length xs ps = let
    idxs = map getPositive ps
    in length (takes idxs xs) == length ps

prop_binary_digits :: Positive Int -> Bool
prop_binary_digits (Positive a) = binaryDigits a == length (showIntAtBase 2 intToDigit a "")

data EuclidianData = EuclidianData [Double] [Double]
    deriving Show

instance Arbitrary EuclidianData where
    arbitrary = do
        Positive sz<-arbitrary
        is1 <- vector sz
        is2 <- vector sz
        return $ EuclidianData is1 is2

prop_euclidian :: EuclidianData -> Bool
prop_euclidian (EuclidianData ds1 ds2) =
    euclidian (M.fromList ds1) (M.fromList ds2)
        == sum (map (\x->x*x) $ zipWith (-) ds2 ds1)

prop_equilateral_size :: Positive Int -> Bool
prop_equilateral_size (Positive n) =
    let m = equilateralEncoding n
    in (n,n-1) == (M.rows m,M.cols m)

prop_equilateral_distance :: Positive Int -> Bool
prop_equilateral_distance (Positive n) =
    let m = equilateralEncoding (n+1)
        ts = M.toRows m
        pairs = concatMap (\(x:ys)->map (\y->(x,y)) ys) $ init $ tails ts
        dis = map (roundTo 5 . uncurry euclidian) pairs
    in 1 == length (ordNub dis)

prop_equilateral_decoding :: Positive Int -> Bool
prop_equilateral_decoding (Positive n) =
    let rn = n + 1
        m  = equilateralEncoding rn
        vs= M.toRows m
        rs = map (equilateralDecoding m) vs
    in rs == [0..n]

testTextData :: T.Text -> IO()
testTextData t = do
    let (TrainData isSt is sz m) = textToTrainData t
    let charSet = T.foldl (flip DS.insert) DS.empty t
    DS.size charSet @=? sz+1
    t @=? dataToText m is
    T.init t @=? dataToText m (tail isSt)

testTextDataS :: T.Text -> IO()
testTextDataS t = do
    let (TrainData isSt is sz m) = textToTrainDataS t
    -- let charSet = T.foldl (flip DS.insert) DS.empty t
    -- sparseSize (DS.size charSet) @=? sz
    t @=? dataToTextS m is
    T.init t @=? dataToTextS m (tail isSt)

checkSteps :: Bool -> IO ()
checkSteps back = do
    (n::RNNetwork) <- evalRandIO $ randomNetwork (RNNDimensions 1 2 3 back) totalDataLength
    let (n2,out)=evalStep n $ M.fromList [1::Double]
        (n3,out1)=evalStep n2 $ M.fromList [3]
        (n4,out2)=evalSteps n [M.fromList [1],M.fromList [3]]
    n3 @=? n4
    out @=? head out2
    out1 @=? last out2

checkLSTMSteps :: IO ()
checkLSTMSteps = do
    (n::LSTMNetwork) <- evalRandIO $ randomNetwork 2 lstmFullSize
    let (n2,out)=evalStep n $ M.fromList [1::Double,2]
        (n3,out1)=evalStep n2 $ M.fromList [3,4]
        (n4,out2)=evalSteps n [M.fromList [1,2],M.fromList [3,4]]
    n3 @=? n4
    out @=? head out2
    out1 @=? last out2

prop_lstm_eval_steps :: LSTMData2 -> Bool
prop_lstm_eval_steps (LSTMData2 n1 is1 _ is2) =
    let
        (n2,out)=evalStep n1 $ M.fromList is1
        (n3,out1)=evalStep n2 $ M.fromList is2
        (n4,out2)=evalSteps n1 [M.fromList is1,M.fromList is2]
    in n3==n4 && out==head out2 && out1 == last out2

checkLSTMList :: IO ()
checkLSTMList = do
    (n1::LSTMNetwork) <- evalRandIO $ randomNetwork 2 lstmFullSize
    (n2::LSTMNetwork) <- evalRandIO $ randomNetwork 2 lstmFullSize
    let (n1_2,out)=evalStep n1 $ M.fromList [1::Double,2]
        (n2_2,out1)=evalStep n2 out
        (n3,out2)=evalStep [n1,n2] $ M.fromList [1,2]
    n3 @=? [n1_2,n2_2]
    out1 @=? out2

prop_lstm_list :: LSTMData2 -> Bool
prop_lstm_list (LSTMData2 n1 is1 n2 _) =
    let (n1_2,out)=evalStep n1 $ M.fromList is1
        (n2_2,out1)=evalStep n2 out
        (n3,out2)=evalStep [n1,n2] $ M.fromList is1
    in n3 == [n1_2,n2_2] && out1 == out2

checkLSTMVector :: IO ()
checkLSTMVector = do
    (n::LSTMNetwork) <- evalRandIO $ randomNetwork 2 lstmFullSize
    let arr = toVector n
        n2  = fromVector 2 arr
    n @=? n2

prop_lstm_step :: LSTMData -> Bool
prop_lstm_step (LSTMData n is)= snd (evalStep n (M.fromList is)) == M.fromList (snd $ lstmList 10 (M.toList $ toVector n) is)

checkLSTMIO :: IO()
checkLSTMIO = do
    print $ lstmFullSize 4
    print $ lstmioFullSize (2,4,3,1)
    (n1::LSTMIO) <- evalRandIO $ randomNetwork (2,4,3,1) lstmioFullSize
    let sz=rnnsize n1
        v1=toVector n1
        n2=fromVector sz v1
    n1 @=? n2
    let is=M.fromList [0.3,0.4]
        r1@(_,o1) = evalStep n1 is
    M.size o1 @?= 1
    let r2 = evalStep n2 is
    r1 @=? r2

data LSTMData = LSTMData LSTMNetwork [Double]
    deriving Show

instance Arbitrary LSTMData where
    arbitrary = do
        let sz=10
        ls <- vector (lstmFullSize sz)
        is <- vector sz
        let n = fromVector sz $ M.fromList ls
        return $ LSTMData n is

data LSTMData2 = LSTMData2 LSTMNetwork [Double] LSTMNetwork [Double]
    deriving Show

instance Arbitrary LSTMData2 where
    arbitrary = do
        let sz=10
        ls1 <- vector (lstmFullSize sz)
        is1 <- vector sz
        let n1 = fromVector sz $ M.fromList ls1
        ls2 <- vector (lstmFullSize sz)
        is2 <- vector sz
        let n2 = fromVector sz $ M.fromList ls2
        return $ LSTMData2 n1 is1 n2 is2

checkArray :: Bool -> IO ()
checkArray back = do
    let dim = RNNDimensions 1 2 3 back
    (n::RNNetwork) <- evalRandIO $ randomNetwork dim totalDataLength
    let arr = networkToArray n
        n2  = createNetworkFromArray dim arr
    case n2 of
        Right rn2 -> n @=? rn2
        Left err  -> assertFailure err

checkVector :: Bool -> IO ()
checkVector back = do
    let dim = RNNDimensions 1 2 3 back
    (n::RNNetwork) <- evalRandIO $ randomNetwork dim totalDataLength
    let arr = toVector n
        n2  = createNetworkFromVector dim arr
    case n2 of
        Right rn2 -> n @=? rn2
        Left err  -> assertFailure err

testCrossover :: (RNNetwork
                        -> RNNetwork -> Rand StdGen [RNNetwork])
                       -> IO ()
testCrossover f = do
    let dim = RNNDimensions 1 2 3 True
    (n1::RNNetwork) <- evalRandIO $ randomNetwork dim totalDataLength
    (n2::RNNetwork) <- evalRandIO $ randomNetwork dim totalDataLength
    rnns <- evalRandIO $ f n1 n2
    2 @=? length rnns
    notElem n1 rnns @? "n1 in result"
    notElem n2 rnns @? "n2 in result"

testMutation :: (LSTMNetwork -> Rand StdGen LSTMNetwork) -> IO()
testMutation f = do
    (n1::LSTMNetwork) <- evalRandIO $ randomNetwork 5 lstmFullSize
    n2 <- evalRandIO $ f n1
    toVector n1 /= toVector n2 @? "n1==n2"

testExpr :: IO()
testExpr = do
    let
        f :: (Num a,Floating a)=> [a] -> a
        f= \[x,y,_] -> x * sin (x + log y)
        i = I.fromList [(1,2::Double),(2,3),(3,4)]
        gs0 = grad f [Var 1,Var 2,Var 3]
        gs02 = grad f [2::Double,3,4]
        cexpr = map (\g-> close g i) gs0
        eexpr = map eval cexpr
    gs02 @=? eexpr
    let
        f1 :: (Num a,Floating a)=> a -> [a] -> a
        f1 = \x [y,_] -> x * sin (x + log y)
        i1 = I.fromList [(2,3),(3,4)]
        x1 :: Double
        x1 = 2
        gs1 = grad (f1 (autoEval (Lit $ Lit x1) $ I.empty)) [Var 2,Var 3]
        gs12 = grad (f1 2) [3,4]
        cexpr1 = map (\g-> fullSimplify $ close g i1) gs1
        eexpr1 = map eval cexpr1
    gs12 @=? eexpr1
    let
        f2 :: (Num a,Floating a)=> [a] -> [a] -> a
        f2= \[x] [y,_] -> x * sin (x + log y)
        i2 = I.fromList [(2,3),(3,4)]
        --ex :: [Expr (Expr Double)]
        --ex = map (\x->autoEval (Lit x) I.empty) [2::Double]
        x2 :: [Double]
        x2 = [2]
        gs2 = grad (f2 (map (\x->autoEval (Lit $ Lit x) I.empty) x2)) [Var 2,Var 3]
        gs22 = grad (f2 [2]) [3,4]
        cexpr2 = map (\g-> fullSimplify $ close g i2) gs2
        eexpr2 = map eval cexpr2
    gs22 @=? eexpr2
    let
        fs = lstmFullSize 2
        ixs = [1..fs]
        ixsF = map fromIntegral ixs
        gs3::[Expr Double] = grad costT $ map Var ixs
        i3 = I.fromList $ zip ixs ixsF
        gs32 = grad (costT) ixsF
        cexpr3 = map (\g-> fullSimplify $ close g i3) gs3
        eexpr3 = map eval cexpr3
    (map (rnd 6) gs32) @=? (map (rnd 6) eexpr3)
    where rnd n f = (fromInteger $ round $ f * (10^n)) / (10.0^^n)
--  print gs3
--  mapM_ (print . prettyShow . fullSimplify) (take 1 gs3)
--  print $ length gs3
--  print fs

costT :: (Num b,Floating b,Fractional b) => [b] -> b
costT lstm = let
    res = snd $ mapAccumL (lstmList 2) lstm [[1,2],[2,3]]
    in sum $ zipWith err [[3,4],[4,5]] res
    where
      err :: (Num b,Floating b) => [b] -> [b] -> b
      err a b  = sum (zipWith (\c d -> (c- d)**2 ) a b)


{-# LANGUAGE FlexibleContexts #-}
module STC.DFTFilter where

import           Control.Monad                as M
import           Control.Monad.IO.Class
import           Control.Monad.Parallel       as MP
import           Control.Monad.Trans.Resource
import           Data.Array.Repa              as R
import           Data.Binary
import           Data.ByteString.Lazy         as BL
import           Data.Complex
import           Data.Conduit
import           Data.Conduit.List            as C
import           Data.Int
import           Data.List                    as L
import           Data.Vector.Storable         as VS
import           Data.Vector.Unboxed          as VU
import           DFT.Plan
import           Filter.Utils
import           FokkerPlanck.Interpolation
import           System.IO
import           Utils.Array
import           Utils.Parallel               hiding ((.|))

{-# INLINE sourceFile #-}
sourceFile :: FilePath -> ConduitT () (VS.Vector (Complex Double)) (ResourceT IO) ()
sourceFile filePath = do
  h <- liftIO $ openBinaryFile filePath ReadMode
  sourceFunc h
  liftIO $ hClose h
  where
    sourceFunc handle = do
      flag <- liftIO $ hIsEOF handle
      if flag
        then sourceNull
        else do
          lenBS <- liftIO $ BL.hGet handle 8
          let len = fromIntegral (decode lenBS :: Int64) :: Int
          dataBS <- liftIO $ BL.hGet handle len
          yield . VS.fromList $ (decode dataBS :: [Complex Double])
          sourceFunc handle

{-# INLINE sinkFile #-}
sinkFile :: FilePath -> ConduitT ByteString Void (ResourceT IO) ()
sinkFile filePath = do
  h <- liftIO $ openBinaryFile filePath WriteMode
  sinkFunc h
  where
    sinkFunc handle = do
      x <- await
      case x of
        Nothing -> liftIO $ hClose handle
        Just y -> do
          let len = BL.length y
          liftIO $ BL.hPut handle (encode len)
          liftIO $ BL.hPut handle y
          sinkFunc handle

{-# INLINE interpolation #-}
interpolation ::
     (R.Source r Double)
  => DFTPlan
  -> (Double -> Double -> Double -> Double -> Int -> Int -> Complex Double)
  -> R.Array r DIM5 Double
  -> Int
  -> Int
  -> Double
  -> Double
  -> (VU.Vector Double)
  -> (VU.Vector Double)
  -> (VU.Vector Double)
  -> (VU.Vector Double)
  -> (Int, Int)
  -> IO ByteString
interpolation plan pinwheelFunc radialArr xLen yLen scaleFactor rMax thetaFreqs scaleFreqs theta0Freqs scale0Freqs (t, s) = do
  let tf = thetaFreqs VU.! t
      sf = scaleFreqs VU.! s
      pinwheelArr =
        traverse2
          (fromUnboxed (Z :. VU.length theta0Freqs) theta0Freqs)
          (fromUnboxed (Z :. VU.length scale0Freqs) scale0Freqs)
          (\(Z :. numTheta0Freq) (Z :. numScale0Freq) ->
             (Z :. numTheta0Freq :. numScale0Freq :. xLen :. yLen)) $ \ft0 fs0 (Z :. t0 :. s0 :. i :. j) ->
          pinwheelFunc
            (tf - ft0 (Z :. t0))
            (sf + fs0 (Z :. s0))
            rMax
            0
            (i - center xLen)
            (j - center yLen)
      filter =
        radialCubicInterpolation
          (R.slice radialArr (Z :. t :. s :. All :. All :. All))
          scaleFactor
          pinwheelArr
  filterF <-
    dftExecute
      plan
      (DFTPlanID
         DFT1DG
         [VU.length theta0Freqs, VU.length scale0Freqs, xLen, yLen]
         [2, 3]) .
    VU.convert . toUnboxed . computeS $
    filter
  return . encode . VS.toList $ filterF

{-# INLINE dftFilter2File #-}
dftFilter2File ::
     (R.Source r Double)
  => ParallelParams
  -> DFTPlan
  -> (Double -> Double -> Double -> Double -> Int -> Int -> Complex Double)
  -> R.Array r DIM5 Double
  -> Int
  -> Int
  -> Double
  -> Double
  -> (VU.Vector Double)
  -> (VU.Vector Double)
  -> (VU.Vector Double)
  -> (VU.Vector Double)
  -> FilePath
  -> IO ()
dftFilter2File parallelParams plan pinwheelFunc radialArr xLen yLen scaleFactor rMax thetaFreqs scaleFreqs theta0Freqs scale0Freqs filePath =
  runConduitRes $
  sourceList
    [ (t, s)
    | t <- [0 .. VU.length thetaFreqs]
    , s <- [0 .. VU.length scaleFreqs]
    ] .|
  parConduitIO
    parallelParams
    (interpolation
       plan
       pinwheelFunc
       radialArr
       xLen
       yLen
       scaleFactor
       rMax
       thetaFreqs
       scaleFreqs
       theta0Freqs
       scale0Freqs) .|
  sinkFile filePath

{-# INLINE convolutionHelper #-}
convolutionHelper ::
     DFTPlan
  -> Int
  -> Int
  -> Int
  -> Int
  -> VS.Vector (Complex Double)
  -> VS.Vector (Complex Double)
  -> IO (VU.Vector (Complex Double))
convolutionHelper plan numThetaFreq numScaleFreq xLen yLen vecF1 vecF2 =
  toUnboxed .
  sumS .
  sumS .
  rotate4D .
  rotate4D .
  fromUnboxed (Z :. numThetaFreq :. numScaleFreq :. xLen :. yLen) . VS.convert <$>
  dftExecute
    plan
    (DFTPlanID IDFT1DG [numThetaFreq, numScaleFreq, xLen, yLen] [2, 3])
    (VS.zipWith (*) vecF1 vecF2)

{-# INLINE convolveSlowly #-}
convolveSlowly ::
     ParallelParams
  -> FilePath
  -> DFTPlan
  -> R.Array U DIM4 (Complex Double)
  -> IO (R.Array U DIM4 (Complex Double))
convolveSlowly parallelParams filePath plan arr = do
  let (Z :. numThetaFreq :. numScaleFreq :. xLen :. yLen) = extent arr
  xs <-
    runConduitRes $
    sourceFile filePath .|
    parConduitIO
      parallelParams
      (convolutionHelper
         plan
         numThetaFreq
         numScaleFreq
         xLen
         yLen
         (VU.convert . toUnboxed $ arr)) .|
    C.consume
  return . fromUnboxed (extent arr) . VU.concat $ xs

{-# INLINE convolutionHelperSink #-}
convolutionHelperSink ::
     DFTPlan
  -> Int
  -> Int
  -> Int
  -> Int
  -> VU.Vector Double
  -> VS.Vector (Complex Double)
  -> ((Int, Int), VS.Vector (Complex Double))
  -> IO (VU.Vector (Complex Double))
convolutionHelperSink plan numThetaFreq numScaleFreq xLen yLen thetaFreqs vecF1 ((t, _), vecF2) = do
  let tf = thetaFreqs VU.! t
  arr <-
    fromUnboxed (Z :. numThetaFreq :. numScaleFreq :. xLen :. yLen) . VS.convert <$>
    dftExecute
      plan
      (DFTPlanID IDFT1DG [numThetaFreq, numScaleFreq, xLen, yLen] [2, 3])
      (VS.zipWith (*) vecF1 vecF2)
  return . toUnboxed . sumS . sumS . rotate4D . rotate4D . R.traverse arr id $ \f idx@(Z :. t0 :. _ :. i :. j) ->
    f idx * exp (0 :+ ((thetaFreqs VU.! t0) + tf) * pi)

{-# INLINE convolveSlowlySink #-}
convolveSlowlySink ::
     ParallelParams
  -> FilePath
  -> DFTPlan
  -> VU.Vector Double
  -> VU.Vector Double
  -> R.Array U DIM4 (Complex Double)
  -> IO (R.Array U DIM4 (Complex Double))
convolveSlowlySink parallelParams filePath plan thetaFreqs scaleFreqs arr = do
  let (Z :. numThetaFreq :. numScaleFreq :. xLen :. yLen) = extent arr
  xs <-
    runConduitRes $
    sourceFile filePath .|
    mergeSource
      (sourceList
         [ (t, s)
         | t <- [0 .. VU.length thetaFreqs]
         , s <- [0 .. VU.length scaleFreqs]
         ]) .|
    parConduitIO
      parallelParams
      (convolutionHelperSink
         plan
         numThetaFreq
         numScaleFreq
         xLen
         yLen
         thetaFreqs
         (VU.convert . toUnboxed $ arr)) .|
    C.consume
  return . fromUnboxed (extent arr) . VU.concat $ xs


module MnistLoader (loadData, loadDataWrapper, loadTestWrapper, investigate) where

import Data.IDX
import qualified Data.Vector.Unboxed as V
import Numeric.LinearAlgebra as L

vect n = (n><1)

loadData :: IO ([(Int, V.Vector Double)])
loadData = do
  mData <- decodeIDXFile "data/train-images-idx3-ubyte"
  let d = getSomething mData
  mLabels <- decodeIDXLabelsFile "data/train-labels-idx1-ubyte"
  let l = getSomething mLabels
  return $ getSomething $ labeledDoubleData l d

loadTest :: IO ([(Int, V.Vector Double)])
loadTest = do
  mData <- decodeIDXFile "data/t10k-images.idx3-ubyte"
  let d = getSomething mData
  mLabels <- decodeIDXLabelsFile "data/t10k-labels.idx1-ubyte"
  let l = getSomething mLabels
  return $ getSomething $ labeledDoubleData l d

-- Converts the label component to a vector representation where
-- n'th component is 1
loadDataWrapper :: IO ([(Matrix Float, Matrix Float)])
loadDataWrapper = do
  theData <- loadData
  return $ map (\(l, d) -> (mVect d, makeVector l)) theData
  where
    makeVector l = vect 10 $ (map (\i -> if i == l then 1 else 0) [0..10])
    mVect l = vect (V.length l) (map realToFrac $ V.toList l)

loadTestWrapper :: IO ([(Matrix Float, Matrix Float)])
loadTestWrapper = do
  theData <- loadTest
  return $ map (\(l, d) -> (mVect d, makeVector l)) theData
  where
    makeVector l = vect 10 $ (map (\i -> if i == l then 1 else 0) [0..10])
    mVect l = vect (V.length l) $ map realToFrac $ V.toList l

investigate :: IO ()
investigate = do
    putStrLn "Inspecting File"
    training_data_mb <- decodeIDXFile "data/train-images-idx3-ubyte"
    let td = getSomething training_data_mb

    putStrLn "Image Data"
    putStr "idxType: "
    putStrLn $ show $ idxType td

    putStr "idxDimensions: "
    putStrLn $ show $ idxDimensions td

    putStr "isIDXReal: "
    putStrLn $ show $ isIDXReal td

    putStr "isIDXIntegral: "
    putStrLn $ show $ isIDXIntegral td

    let training_vector = idxDoubleContent td
    putStrLn $ show $ V.length training_vector

    loaded_data <- loadData
    let (l, d) = loaded_data !! 1
    putStrLn $ show $ V.length $ d
    putStrLn $ show $ l

getSomething :: Maybe a -> a
getSomething (Just x) = x
getSomething Nothing = error "Bah"


{-# LANGUAGE LambdaCase #-}
module AI.NeuroCar where

import           AI.GeneticAlgorithm   as GA
import qualified AI.NeuralNetwork      as NN
import qualified Car                   as C
import           Control.Lens
import           Control.Monad.Loops
import           Control.Monad.State
import           Data.Foldable
import           Data.Functor.Identity
import           Data.Word             (Word32)
import           Game
import qualified Geometry              as G
import qualified Numeric.LinearAlgebra as LA
import           System.Random
import qualified Track                 as T
import           World

rayAngle :: Double
rayAngle = pi

rayCount :: Int
rayCount = 7

evolveCar :: Int -> Int -> Int -> Double -> Double -> Double -> Word32 ->
    C.CarParams -> T.Track Double -> [Population Double]
evolveCar seed generations popSize mutChance mutStrength time deltaTicks carParams track = do
    let gen          = mkStdGen seed
    let drawFunc     = return . const ()
    let inputFunc nn = Identity . getNetworkInput nn
    let timeFunc     = return deltaTicks
    let world        = initWorld carParams track time
    let evolveLoop nn = gameLoop drawFunc (inputFunc nn) timeFunc
    let runGame nn = runIdentity $ iterateUntilM (\w -> w^.gameState /= GameRunning) (evolveLoop nn) world
    let fitfunc nn = fromIntegral $ view score (runGame nn)
    let mutfunc = GA.mutate mutChance mutStrength
    let indGen = NN.newNetwork [3+rayCount, 15, 15, 2]
    let evoFunc = evolves generations popSize indGen fitfunc mutfunc
    evalState evoFunc gen

getNetworkInput :: NN.Network -> World -> [Input]
getNetworkInput nn w =
    let push a        = (state $ \xs -> ((),a:xs)) :: State [Input] ()
        rays          = C.shootRays rayAngle rayCount $ w^.car
        origin        = w^.car.C.position
        intersections = T.rayIntersection (w^.track) <$> rays
        distances     = (\case Nothing -> 0
                               Just v  -> G.segLength (origin, v)) <$> intersections
        velocity      = C.localVelocity $ w^.car
        rotation      = w^.car.C.rotation
        x             = rotation : toList velocity ++ distances
        y             = LA.toList $ NN.feedforward (LA.fromList x) nn
     in execState (do { when (head y < 1/3) (push GoBackward)
                      ; when (head y > 2/3) (push GoForward)
                      ; when (y !! 1 < 1/3) (push GoLeft)
                      ; when (y !! 1 > 2/3) (push GoRight) }) []


{- | Example of sampling

Two samplers are availables : the 'discreteAncestralSampler' and the 'gibbsSampler'.
Only the 'gibbsSampler' can be used with evidence.

In this example, we have a very simple network.

@
    simple :: (['TDV' Bool],'SBN' 'CPT')
    simple = 'runBN' $ do
        a <- 'variable' \"a\" ('t' :: Bool)
        b <- 'variable' \"b\" ('t' :: Bool)
--
        'proba' a '~~' [0.4,0.6]
        'cpt' b [a] '~~' [0.8,0.2,0.2,0.8]
--
        return [a,b]
@

This network is representing a sensor b. We observe the value of b and we want to infer the value of a.

We use the 'gibbsSampler' for this with an initial period of 200 samples which are dropped. The 'gibbsSampler' is
generate a stream of samples. From this stream, we need to compute a probability distribution. For this, we use
the 'samplingHistograms' histogram function which is generating a list : the probability values of each vertex.

@
    let (vars\@[a,b],exampleG) = simple
    n <- 'runSampling' 5000 200 ('gibbsSampler' exampleG [b '=:' True])
    let h = 'samplingHistograms' n
    print $ h
@

Then, we compare this result with the exact one we get with a junction tree.

@
    let jt = 'createJunctionTree' 'nodeComparisonForTriangulation' exampleG
        jt' = 'changeEvidence' [b '=:' True] jt
    mapM_ (\x -> print . 'posterior' jt' $ [x]) vars
@

We can also use the 'discreteAncestralSampler' to compute the posterior but it is not supporting the use of evidence in this
version. The syntax is similar.

@
    n <- 'runSampling' 500 ('discreteAncestralSampler' exampleG)
@

-}
module Bayes.Examples.Sampling(
    -- * Test function
	  testSampling
) where

import Bayes.Sampling
import System.Random.MWC.CondensedTable
import qualified Data.Vector as V
import Bayes.Factor
import Bayes
import Bayes.FactorElimination
import Data.Maybe(fromJust)
import Bayes.BayesianNetwork
import Bayes.Factor.CPT
import Statistics.Sample.Histogram

simple :: ([TDV Bool],SBN CPT)
simple = runBN $ do
    a <- variable "a" (t :: Bool)
    b <- variable "b" (t :: Bool)

    proba a ~~ [0.4,0.6]
    cpt b [a] ~~ [0.8,0.2,0.2,0.8]

    return [a,b]

testSampling = do
    let (vars@[a,b],exampleG) = simple
        jt = createJunctionTree nodeComparisonForTriangulation exampleG
    --n <- runSampling 500 (discreteAncestralSampler exampleG)
    putStrLn "The bayesian network"
    print exampleG
    putStrLn "\nThe values of the Bayesian network"
    printGraphValues exampleG

    n <- runSampling 5000 200 (gibbsSampler exampleG [b =: True])
    let h = samplingHistograms 2 n
    putStrLn "\nThe result of the sampling"
    print $ h

    n <- runSampling 5000 200 (gibbsMCMCSampler exampleG [b =: True])
    let h = samplingHistograms 2 n
    putStrLn "\nThe result of the sampling with MCMC"
    print $ h

    let jt' = changeEvidence [b =: True] jt
    putStrLn "\nThe result of the junction tree inference"
    mapM_ (\x -> print . posterior jt' $ [x]) vars