{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE DeriveDataTypeable, DeriveGeneric #-}
-- |
-- Module    : Statistics.Distribution.Hypergeometric
-- Copyright : (c) 2009 Bryan O'Sullivan
-- License   : BSD3
--
-- Maintainer  : bos@serpentine.com
-- Stability   : experimental
-- Portability : portable
--
-- The Hypergeometric distribution.  This is the discrete probability
-- distribution that measures the probability of /k/ successes in /l/
-- trials, without replacement, from a finite population.
--
-- The parameters of the distribution describe /k/ elements chosen
-- from a population of /l/, with /m/ elements of one type, and
-- /l/-/m/ of the other (all are positive integers).

module Statistics.Distribution.Hypergeometric
    (
      HypergeometricDistribution
    -- * Constructors
    , hypergeometric
    , hypergeometricE
    -- ** Accessors
    , hdM
    , hdL
    , hdK
    ) where

import Control.Applicative
import Data.Aeson           (FromJSON(..), ToJSON, Value(..), (.:))
import Data.Binary          (Binary(..))
import Data.Data            (Data, Typeable)
import GHC.Generics         (Generic)
import Numeric.MathFunctions.Constants (m_epsilon,m_neg_inf)
import Numeric.SpecFunctions (choose,logChoose)

import qualified Statistics.Distribution as D
import Statistics.Internal

data HypergeometricDistribution = HD {
      hdM :: {-# UNPACK #-} !Int
    , hdL :: {-# UNPACK #-} !Int
    , hdK :: {-# UNPACK #-} !Int
    } deriving (Eq, Typeable, Data, Generic)

instance Show HypergeometricDistribution where
  showsPrec i (HD m l k) = defaultShow3 "hypergeometric" m l k i
instance Read HypergeometricDistribution where
  readPrec = defaultReadPrecM3 "hypergeometric" hypergeometricE

instance ToJSON HypergeometricDistribution
instance FromJSON HypergeometricDistribution where
  parseJSON (Object v) = do
    m <- v .: "hdM"
    l <- v .: "hdL"
    k <- v .: "hdK"
    maybe (fail $ errMsg m l k) return $ hypergeometricE m l k
  parseJSON _ = empty

instance Binary HypergeometricDistribution where
  put (HD m l k) = put m >> put l >> put k
  get = do
    m <- get
    l <- get
    k <- get
    maybe (fail $ errMsg m l k) return $ hypergeometricE m l k

instance D.Distribution HypergeometricDistribution where
    cumulative = cumulative

instance D.DiscreteDistr HypergeometricDistribution where
    probability    = probability
    logProbability = logProbability

instance D.Mean HypergeometricDistribution where
    mean = mean

instance D.Variance HypergeometricDistribution where
    variance = variance

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

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

instance D.Entropy HypergeometricDistribution where
  entropy = directEntropy

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

variance :: HypergeometricDistribution -> Double
variance (HD m l k) = (k' * ml) * (1 - ml) * (l' - k') / (l' - 1)
  where m' = fromIntegral m
        l' = fromIntegral l
        k' = fromIntegral k
        ml = m' / l'

mean :: HypergeometricDistribution -> Double
mean (HD m l k) = fromIntegral k * fromIntegral m / fromIntegral l

directEntropy :: HypergeometricDistribution -> Double
directEntropy d@(HD m _ _)
  = negate . sum
  $ takeWhile (< negate m_epsilon)
  $ dropWhile (not . (< negate m_epsilon))
    [ let x = probability d n in x * log x | n <- [0..m]]

hypergeometric :: Int               -- ^ /m/
               -> Int               -- ^ /l/
               -> Int               -- ^ /k/
               -> HypergeometricDistribution
hypergeometric m l k
  = maybe (error $ errMsg m l k) id $ hypergeometricE m l k

hypergeometricE :: Int               -- ^ /m/
                -> Int               -- ^ /l/
                -> Int               -- ^ /k/
                -> Maybe HypergeometricDistribution
hypergeometricE m l k
  | not (l > 0)            = Nothing
  | not (m >= 0 && m <= l) = Nothing
  | not (k > 0  && k <= l) = Nothing
  | otherwise              = Just (HD m l k)

errMsg :: Int -> Int -> Int -> String
errMsg m l k
  =  "Statistics.Distribution.Hypergeometric.hypergeometric: "
  ++ "m=" ++ show m
  ++ "l=" ++ show l
  ++ "k=" ++ show k
  ++ " should hold: l>0 & m in [0,l] & k in (0,l]"

-- Naive implementation
probability :: HypergeometricDistribution -> Int -> Double
probability (HD mi li ki) n
  | n < max 0 (mi+ki-li) || n > min mi ki = 0
    -- No overflow
  | li < 1000 = choose mi n * choose (li - mi) (ki - n)
              / choose li ki
  | otherwise = exp $ logChoose mi n
                    + logChoose (li - mi) (ki - n)
                    - logChoose li ki

logProbability :: HypergeometricDistribution -> Int -> Double
logProbability (HD mi li ki) n
  | n < max 0 (mi+ki-li) || n > min mi ki = m_neg_inf
  | otherwise = logChoose mi n
              + logChoose (li - mi) (ki - n)
              - logChoose li ki

cumulative :: HypergeometricDistribution -> Double -> Double
cumulative d@(HD mi li ki) x
  | isNaN x      = error "Statistics.Distribution.Hypergeometric.cumulative: NaN argument"
  | isInfinite x = if x > 0 then 1 else 0
  | n <  minN    = 0
  | n >= maxN    = 1
  | otherwise    = D.sumProbabilities d minN n
  where
    n    = floor x
    minN = max 0 (mi+ki-li)
    maxN = min mi ki


{-|
Module      : kMeans
Description : k-Means
Copyright   : (c) Fabrício Olivetti, 2017
License     : GPL-3
Maintainer  : fabricio.olivetti@gmail.com

k-Means Clustering algorithm
-}

module KMeansPar where

import Control.Parallel.Strategies
import Data.List
import Data.Ord
import Data.List.Split (chunksOf)
--import qualified Numeric.LinearAlgebra as H

import Vector
import Dados

-- |'parseFile' parses a space separated file
-- to a list of lists of Double
parseFile :: String -> [[Double]]
parseFile file = map parseLine (lines file)
  where
    parseLine l = map toDouble (words l)
    toDouble  w = read w :: Double

firstClusters idxs points = [ points !! i | i <- idxs ]

kmeans :: Int -> ChunksOf [[Double]] -> [[Double]]
       -> [[Double]]
kmeans it points clusters
  | it == 0               = clusters
  | clusters' == clusters = clusters
  | otherwise             = kmeans (it-1) points clusters'
  where
    clusters' = emStep points clusters

emStep :: ChunksOf [[Double]] -> [[Double]] -> [[Double]]
emStep points clusters = map (\(idx, (xs,n)) -> xs ./ n)
                       $ mapReduceByKey (\pi -> (assign pi clusters, (pi,1.0))) soma points
  where
    soma (xs1, y1) (xs2, y2) = (xs1 .+. xs2, y1+y2)

assign :: [Double] -> [[Double]] -> Integer
assign point clusters = closestTo point
  where
    closestTo p = argmin $ map (euclid p) clusters
    euclid pi ci = sum $ map (^2) $ zipWith (-) pi ci
    argmin xs = fst $ head $ sortByValue $ zip [0..] xs


{-# LANGUAGE BangPatterns, CPP #-}
import System.Environment
import Control.Monad
import Control.Seq
import Data.Complex
import System.IO
import Debug.Trace
import Control.DeepSeq
import Control.Monad.Par
import Control.Exception

import PortablePixmap
import Control.Monad.Par.AList as A

mandel :: Int -> Complex Double -> Int
mandel max_depth c = loop 0 0
  where
   fn = magnitude
   loop i !z
    | i == max_depth = i
    | fn(z) >= 2.0   = i
    | otherwise      = loop (i+1) (z*z + c)

threshold = 1

runMandel :: Double -> Double -> Double -> Double -> Int -> Int -> Int -> Par (AList [Int])
runMandel minX minY maxX maxY winX winY max_depth = do

  parBuildThreshM threshold (InclusiveRange 0 (winY-1)) $ \y ->
       do
          let l = [ mandelStep y x | x <- [0.. winX-1] ]
          deepseq l (return l)

  where
    mandelStep i j = mandel max_depth z
        where z = ((fromIntegral j * r_scale) / fromIntegral winY + minY) :+
                  ((fromIntegral i * c_scale) / fromIntegral winX + minX)

    r_scale  =  maxY - minY  :: Double
    c_scale =   maxX - minX  :: Double

makeImage :: Integer -> Integer -> Int -> AList [Int] -> PixMap
makeImage x y depth ls =
  createPixmap x y depth
   (map prettyRGB (concat (toList ls)))
 where
   prettyRGB :: Int -> (Int,Int,Int)
   prettyRGB s = let t = (depth - s) in (s,t,t)

simple x y depth =
  runMandel 0 0 x' y' x y depth
 where
  x' = fromIntegral x
  y' = fromIntegral y

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

-- A meaningless checksum.
mandelCheck :: AList [Int] -> Int -> Int -> Int
mandelCheck als max_col max_depth = loop 0 als 0
 where
 loop i als !sum | A.null als = sum
 loop i als !sum = loop (i+1) (A.tail als)
		        (loop2 i 0 (A.head als) sum)
 loop2 i j []    !sum = sum
 loop2 i j (h:t) !sum | h == max_depth = loop2 i (j+1) t (sum + i*max_col + j)
		      | otherwise      = loop2 i (j+1) t  sum

main = do args <- getArgs

          let (x,y,depth) =
		case args of
		 [] ->
    --                 runPar $ simple 3 3 3
		       (3,3,3)

		 [x,y,depth] ->
    --		   simple (read x) (read y) (read depth)
		       (read x, read y, read depth)

		 -- [minX,minY,maxX,maxY,winX,winY,depth] ->
		 --    runPar $
		 -- 	    runMandel (read minX) (read minY)
		 -- 		      (read maxX) (read maxY)
		 -- 		      (read winX) (read winY) (read depth)

          let ls = runPar$ simple x y depth
          when (False) $ do
	     hnd <- openFile "mandel_image.ppm" WriteMode
	     hSetBinaryMode hnd True
	     hPrint hnd (makeImage (fromIntegral x) (fromIntegral y) depth ls)
	     hClose hnd

          putStrLn$ "Spot check: " ++ show (mandelCheck ls y depth)


module BackProp (
    backPropRegression,
    backPropClassification,
    backProp,
    sgdMethod
  ) where

import Numeric.LinearAlgebra
import Common
import ActivFunc

sgdMethod :: Int -> (Matrix R, Matrix R) -> ((Matrix R, Matrix R) -> [Matrix R] -> [Matrix R]) -> [Matrix R] -> IO [Matrix R]
sgdMethod n xy f ws = do
  nxy <- pickupSets n xy
  return $ f nxy ws

backPropRegression :: R -> (Matrix R -> Matrix R, Matrix R -> Matrix R) -> (Matrix R, Matrix R) -> [Matrix R] -> [Matrix R]
backPropRegression = backProp id

backPropClassification :: R -> (Matrix R -> Matrix R, Matrix R -> Matrix R) -> (Matrix R, Matrix R) -> [Matrix R] -> [Matrix R]
backPropClassification = backProp softmaxC

backProp :: (Matrix R -> Matrix R) -> R -> (Matrix R -> Matrix R, Matrix R -> Matrix R) -> (Matrix R, Matrix R) -> [Matrix R] -> [Matrix R]
backProp lf rate (f, df) (x, y) ws@(m:ms) = zipWith (-) ws (fmap (scalar rate *) dws)
  where
    dws = fmap (/ len) . zipWith (<>) ds $ fmap (tr . inputWithBias) vs
    len = fromIntegral $ cols x
    vs = x : fmap f rs -- length: L - 1
    ds = calcDeltas df (zip rs ms) dInit -- length: L - 1
    dInit = lf r - y
    (rs, r) = (init us, last us)
    us = forwards' f ms uInit -- length: L - 1
    uInit = m <> inputWithBias x

forwards' :: (Matrix R -> Matrix R) -> [Matrix R] -> Matrix R -> [Matrix R]
forwards' f ws u = scanl (flip $ forward' f) u ws

forward' :: (Matrix R -> Matrix R) -> Matrix R -> Matrix R -> Matrix R
forward' f w u = w <> inputWithBias (f u)

calcDeltas :: (Matrix R -> Matrix R) -> [(Matrix R, Matrix R)] -> Matrix R -> [Matrix R]
calcDeltas df uws d = scanr (calcDelta df) d uws

calcDelta :: (Matrix R -> Matrix R) -> (Matrix R, Matrix R) -> Matrix R -> Matrix R
calcDelta df (u, w) d = df u * tr (weightWithoutBias w) <> d


------------------------------------------------
-- |
-- Module    :  Numeric.Matrix.Integral
-- Copyright :  (c) Jun Yoshida 2019
-- License   :  BSD3
--
-- Operations on integer matrices
--
------------------------------------------------

module Numeric.Matrix.Integral (
  -- Hermite normal form
  hermiteNF,
  -- Smith normal form
  smithNF,
  smithRep,
  ) where

import Control.Monad
import Control.Monad.ST (ST, runST)
import Control.Monad.Loops (whileM_)
import Data.STRef

import qualified Numeric.LinearAlgebra as LA

--import Numeric.Matrix.Integral.HNFLLL
--import Numeric.Matrix.Integral.SmithNF
import Numeric.Matrix.Integral.NormalForms

-- Compute the submatrix containing all the non-zero diagonal entries
extractMaxNZDiag :: LA.Matrix LA.Z -> (Int,LA.Matrix LA.Z)
extractMaxNZDiag mx = runST $ do
  dRef <- newSTRef 0
  let p d = d < uncurry min (LA.size mx) && (mx LA.! d LA.! d) /= 0
  whileM_ (p <$> readSTRef dRef) $ modifySTRef' dRef (+1)
  d <- readSTRef dRef
  return (d,LA.subMatrix (0,0) (d,d) mx)

-- D = P <> A <> Q
-- | Return (d,ker,im) where
-- d: the diagonal entries of Smith normal form
-- ker: basis for the kernel
-- im: basis for the image
kerImOf :: LA.Matrix LA.Z -> (LA.Vector LA.Z,[LA.Vector LA.Z],[LA.Vector LA.Z])
kerImOf mt =
  let (ul,h,ur) = smithNF mt
      (ull,_,ulr) = smithNF ul
      (rk,dmt) = extractMaxNZDiag h
      ker' = ur LA.¿ [rk..(LA.cols ur-1)]
      (_,kert) = hermiteNF (LA.tr' ker')
      im' = (ulr LA.<> ull LA.<> h) LA.¿ [0..rk-1]
      (_,imt) = hermiteNF (LA.tr' im')
  in (LA.takeDiag dmt, LA.toRows kert, LA.toRows imt)


module EndPointMarkovChain where

import           Control.Monad           as M
import           Data.Array.Repa         as R
import           Data.List               as L
import           Graphics.Gnuplot.Simple
import           Numeric.LinearAlgebra   as NL
import           System.Directory
import           System.Environment
import           System.FilePath         ((</>))
import           Text.Printf

{-# INLINE splitList #-}
splitList :: Int ->  [a] -> [[a]]
splitList _ [] = []
splitList n xs = (L.take n xs) : splitList n (L.drop n xs)

{-# INLINE makeMatrix #-}
makeMatrix :: Int -> Double -> R.Array D DIM2 Double
makeMatrix n p =
  fromFunction (Z :. (2 * n) :. (2 * n)) $ \(Z :. i :. j) ->
    let a = div i 2
        b = div j 2
        c = mod i 2
        d = mod j 2
     in if abs (a - b) == 1
          then if a == 0
                 then if d == 1
                        then if c == 0
                               then p
                               else 1
                        else 0
                 else if a == n - 1
                        then if d == 0
                               then if c == 0
                                      then 1
                                      else p
                               else 0
                        else if b < a
                               then if d == 0
                                      then if c == 0
                                             then 1
                                             else p
                                      else if c == 0 -- 0
                                              then p
                                              else 0
                               else if d == 0
                                      then  if c == 0 -- 0
                                               then 0
                                               else p
                                      else if c == 0
                                             then p
                                             else 1
          else 0

{-# INLINE powerMethodIO #-}
powerMethodIO :: Matrix Double -> Vector Double -> Int -> IO (Vector Double)
powerMethodIO mat vec n = do
  let (m,_) = NL.size mat
      x =  ((inv (mat - (0.1) * ident m)) #> vec)
      maxV = sumElements x -- maxElement x
      eigenVec = x / scalar maxV
  printf "%d: %f\n" n (rayleighQuotient mat eigenVec)
  printf "%03d: %s" n (printList . sumVec  . NL.toList $ eigenVec)
  return eigenVec

{-# INLINE powerMethod #-}
powerMethod :: Matrix Double -> Vector Double -> (Double, Vector Double)
powerMethod mat vec =
  let x = mat #> vec
      maxV = maxElement x
      eigenVec = x / scalar maxV
      rq = rayleighQuotient mat eigenVec
   in (rq, eigenVec)

{-# INLINE sumVec #-}
sumVec :: [Double] -> [Double]
sumVec []       = []
sumVec (x:y:xs) = (x + y) : sumVec xs

{-# INLINE printList #-}
printList :: [Double] -> String
printList []     = "\n"
printList (x:xs) = printf "%.3f " x L.++ printList xs

{-# INLINE rayleighQuotient #-}
rayleighQuotient :: Matrix Double -> Vector Double -> Double
rayleighQuotient mat vec = vec <.> (mat #> vec) / vec <.> vec

{-# INLINE inverse #-}
inverse :: Vector Double -> Vector Double
inverse vec = cmap (\x -> sumElements vec - x ) vec

{-# INLINE timeReverse #-}
timeReverse :: [a] -> [a]
timeReverse [] = []
timeReverse (x:y:xs) = y : (x : timeReverse xs)

main = do
  args <- getArgs
  let (nStr:aStr:mStr:_) = args
      n = read nStr :: Int
      a = read aStr :: Double
      m = read mStr :: Int
      mat = (tr . ((2 * n) >< (2 * n)) . R.toList . makeMatrix n $ a) + (0.0) * ident (2 * n)
      -- mat = NL.fromLists . timeReverse . L.map timeReverse . NL.toLists $ mat'
      vec = vector . L.take (2 * n) $ [1,1 ..]
      (eigVal', eigVec') = eig mat
      (maxEigVal, maxEigVec) =
        (\(x, y) -> (realPart x, sumVec $ L.map realPart y)) .
        L.maximumBy (\x y -> compare (realPart $ fst x) (realPart $ fst y)) .
        L.zip (NL.toList eigVal') . L.map NL.toList . NL.toColumns $
        eigVec'
      (maxEigVal', maxEigVec') =
        (\(x, y) -> (realPart x, L.map realPart y)) .
        L.maximumBy (\x y -> compare (realPart $ fst x) (realPart $ fst y)) .
        L.zip (NL.toList eigVal') . L.map NL.toList . NL.toColumns $
        eigVec'
      maxV = L.maximumBy (\x y -> compare (abs x) (abs y)) maxEigVec
      (rqs, eigVecs) =
        L.unzip $ L.scanl' (\(_, v) _ -> powerMethod mat v) (0, vec) [1 .. m]
      -- eigVec = L.last eigVecs
      folderPath = "output/test/EndPointMarkovChain"
  createDirectoryIfMissing True folderPath
  -- eigVec <- M.foldM (powerMethodIO mat) vec [1 .. m]
  -- -- print . rayleighQuotient mat $ eigVec
  -- -- putStr . printList . sumVec . NL.toList $ eigVec
  -- -- print maxEigVal
  -- -- putStr . printList . L.map (/ maxV) $ maxEigVec
  -- printf
  --   "%.3f: %s"
  --   (rayleighQuotient mat eigVec)
  --   (printList . sumVec . NL.toList $ eigVec)
  -- printf
  --   "%.3f: %s"
  --   maxEigVal
  --   -- (L.sum . L.map (/ maxV) $ maxEigVec)
  --   (printList . L.map (/ maxV) $ maxEigVec)
  -- plotPathStyle
  --   [ PNG (folderPath </> "RayleighQuotient.png")
  --   , Title ("Rayleigh Quotient")
  --   , XLabel "Iteration"
  --   ]
  --   (defaultStyle {plotType = Lines, lineSpec = CustomStyle [LineTitle ""]}) .
  --   L.zip [1 ..] . L.tail $
  --   rqs
  -- let m = 10
  --     croppedVec =
  --       L.take (2 * m) . L.drop (div (2 * n - 2 * m) 2) $
  --       L.map (/ maxV) $ maxEigVec'
  -- print . L.sum $ croppedVec
  M.mapM_ print .
    L.map
      (\(x, ys) ->
         let v = L.maximumBy (\a b -> compare (abs a) (abs b)) ys
          in (x, L.map (/ v) ys)) .
    L.map (\(x, y) -> (realPart x, sumVec $ L.map realPart y)) .
    L.sortBy (\x y -> compare (realPart $ fst x) (realPart $ fst y)) .
    L.zip (NL.toList eigVal') . L.map NL.toList . NL.toColumns $
    eigVec'
  -- M.mapM_  print . NL.toLists $ (mat + 1 * ident 20)
  -- M.mapM_  print . NL.toLists . inv $ (mat + 1 * ident 20)
  -- M.mapM_ print . NL.toLists $ mat - (0.017) * ident (2 * n)
  -- putStrLn ""
  M.mapM_ (putStr . printList) . NL.toLists $ (inv $ (mat - (0.6) * ident (2*n))) / 59.592--2261.598 --494.987 -- 4962.166
  --   (mat + 0.001 * ident 20) NL.<> (tr $ (mat + 0.001 * ident 20))


{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE DeriveDataTypeable, DeriveGeneric #-}
-- |
-- Module    : Statistics.Distribution.Lognormal
-- Copyright : (c) 2020 Ximin Luo
-- License   : BSD3
--
-- Maintainer  : infinity0@pwned.gg
-- Stability   : experimental
-- Portability : portable
--
-- The Weibull distribution.  This is a continuous probability
-- distribution that describes the occurrence of a single event whose
-- probability changes over time, controlled by the shape parameter.

module Statistics.Distribution.Weibull
    (
      WeibullDistribution
      -- * Constructors
    , weibullDistr
    , weibullDistrErr
    , weibullStandard
    , weibullDistrApproxMeanStddevErr
    ) where

import Control.Applicative
import Data.Aeson            (FromJSON(..), ToJSON, Value(..), (.:))
import Data.Binary           (Binary(..))
import Data.Data             (Data, Typeable)
import GHC.Generics          (Generic)
import Numeric.MathFunctions.Constants (m_eulerMascheroni)
import Numeric.SpecFunctions (expm1, log1p, logGamma)
import qualified Data.Vector.Generic as G

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

-- | The Weibull distribution.
data WeibullDistribution = WD {
      wdShape  :: {-# UNPACK #-} !Double
    , wdLambda :: {-# UNPACK #-} !Double
    } deriving (Eq, Typeable, Data, Generic)

instance Show WeibullDistribution where
  showsPrec i (WD k l) = defaultShow2 "weibullDistr" k l i
instance Read WeibullDistribution where
  readPrec = defaultReadPrecM2 "weibullDistr" $
    (either (const Nothing) Just .) . weibullDistrErr

instance ToJSON WeibullDistribution
instance FromJSON WeibullDistribution where
  parseJSON (Object v) = do
    k <- v .: "wdShape"
    l <- v .: "wdLambda"
    either fail return $ weibullDistrErr k l
  parseJSON _ = empty

instance Binary WeibullDistribution where
  put (WD k l) = put k >> put l
  get = do
    k <- get
    l <- get
    either fail return $ weibullDistrErr k l

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

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

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

instance D.Mean WeibullDistribution where
  mean (WD k l) = l * exp (logGamma (1 + 1 / k))

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

instance D.Variance WeibullDistribution where
  variance (WD k l) = l * l * (exp (logGamma (1 + 2 * invk)) - q * q)
   where
    invk = 1 / k
    q    = exp (logGamma (1 + invk))

instance D.Entropy WeibullDistribution where
  entropy (WD k l) = m_eulerMascheroni * (1 - 1 / k) + log (l / k) + 1

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

instance D.ContGen WeibullDistribution where
  genContVar d = D.genContinuous d

-- | Standard Weibull distribution with scale factor (lambda) 1.
weibullStandard :: Double -> WeibullDistribution
weibullStandard k = weibullDistr k 1.0

-- | Create Weibull distribution from parameters.
--
-- If the shape (first) parameter is @1.0@, the distribution is equivalent to a
-- 'Statistics.Distribution.Exponential.ExponentialDistribution' with parameter
-- @1 / lambda@ the scale (second) parameter.
weibullDistr
  :: Double            -- ^ Shape
  -> Double            -- ^ Lambda (scale)
  -> WeibullDistribution
weibullDistr k l = either error id $ weibullDistrErr k l

-- | Create Weibull distribution from parameters.
--
-- If the shape (first) parameter is @1.0@, the distribution is equivalent to a
-- 'Statistics.Distribution.Exponential.ExponentialDistribution' with parameter
-- @1 / lambda@ the scale (second) parameter.
weibullDistrErr
  :: Double            -- ^ Shape
  -> Double            -- ^ Lambda (scale)
  -> Either String WeibullDistribution
weibullDistrErr k l | k <= 0     = Left $ errMsg k l
                    | l <= 0     = Left $ errMsg k l
                    | otherwise = Right $ WD k l

errMsg :: Double -> Double -> String
errMsg k l =
  "Statistics.Distribution.Weibull.weibullDistr: both shape and lambda must be positive. Got shape "
    ++ show k
    ++ " and lambda "
    ++ show l

-- | Create Weibull distribution from mean and standard deviation.
--
-- The algorithm is from "Methods for Estimating Wind Speed Frequency
-- Distributions", C. G. Justus, W. R. Hargreaves, A. Mikhail, D. Graber, 1977.
-- Given the identity:
--
-- \[
-- (\frac{\sigma}{\mu})^2 = \frac{\Gamma(1+2/k)}{\Gamma(1+1/k)^2} - 1
-- \]
--
-- \(k\) can be approximated by
--
-- \[
-- k \approx (\frac{\sigma}{\mu})^{-1.086}
-- \]
--
-- \(\lambda\) is then calculated straightforwardly via the identity
--
-- \[
-- \lambda = \frac{\mu}{\Gamma(1+1/k)}
-- \]
--
-- Numerically speaking, the approximation for \(k\) is accurate only within a
-- certain range. We arbitrarily pick the range \(0.033 \le \frac{\sigma}{\mu} \le 1.45\)
-- where it is good to ~6%, and will refuse to create a distribution outside of
-- this range. The paper does not cover these details but it is straightforward
-- to check them numerically.
weibullDistrApproxMeanStddevErr
  :: Double            -- ^ Mean
  -> Double            -- ^ Stddev
  -> Either String WeibullDistribution
weibullDistrApproxMeanStddevErr m s = if r > 1.45 || r < 0.033
    then Left msg
    else weibullDistrErr k l
  where r = s / m
        k = (s / m) ** (-1.086)
        l = m / exp (logGamma (1 + 1/k))
        msg = "Statistics.Distribution.Weibull.weibullDistr: stddev-mean ratio "
          ++ "outside approximation accuracy range [0.033, 1.45]. Got "
          ++ "stddev " ++ show s ++ " and mean " ++ show m

-- | Uses an approximation based on the mean and standard deviation in
--   'weibullDistrEstMeanStddevErr', with standard deviation estimated
--   using maximum likelihood method (unbiased estimation).
--
--   Returns @Nothing@ if sample contains less than one element or
--   variance is zero (all elements are equal), or if the estimated mean
--   and standard-deviation lies outside the range for which the
--   approximation is accurate.
instance D.FromSample WeibullDistribution Double where
  fromSample xs
    | G.length xs <= 1 = Nothing
    | v == 0           = Nothing
    | otherwise        = either (const Nothing) Just $
      weibullDistrApproxMeanStddevErr m (sqrt v)
    where
      (m,v) = S.meanVarianceUnb xs

logDensity :: WeibullDistribution -> Double -> Double
logDensity (WD k l) x
  | x < 0     = 0
  | otherwise = log k + (k - 1) * log x - k * log l - (x / l) ** k

cumulative :: WeibullDistribution -> Double -> Double
cumulative (WD k l) x | x < 0     = 0
                      | otherwise = -expm1 (-(x / l) ** k)

complCumulative :: WeibullDistribution -> Double -> Double
complCumulative (WD k l) x | x < 0     = 1
                           | otherwise = exp (-(x / l) ** k)

quantile :: WeibullDistribution -> Double -> Double
quantile (WD k l) p
  | p == 0         = 0
  | p == 1         = inf
  | p > 0 && p < 1 = l * (-log1p (-p)) ** (1 / k)
  | otherwise      =
    error $ "Statistics.Distribution.Weibull.quantile: p must be in [0,1] range. Got: " ++ show p
  where inf = 1 / 0

complQuantile :: WeibullDistribution -> Double -> Double
complQuantile (WD k l) q
  | q == 0         = inf
  | q == 1         = 0
  | q > 0 && q < 1 = l * (-log q) ** (1 / k)
  | otherwise      =
    error $ "Statistics.Distribution.Weibull.complQuantile: q must be in [0,1] range. Got: " ++ show q
  where inf = 1 / 0


module STCR2Z1T0_1 where

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

main = do
  args@(numPointStr:numOrientationStr:sigmaStr:taoStr:lenStr:initStr:numTrailStr:maxTrailStr:theta0FreqsStr:thetaFreqsStr:initDistStr:histFilePath:alphaStr:pinwheelFlagStr:numThreadStr:_) <-
    getArgs
  print args
  let numPoint = read numPointStr :: Int
      numOrientation = read numOrientationStr :: Int
      sigma = read sigmaStr :: Double
      tao = read taoStr :: Double
      len = read lenStr :: Int
      init = read initStr :: (Double, Double, Double, Double, Double, Double)
      numTrail = read numTrailStr :: Int
      maxTrail = read maxTrailStr :: Int
      theta0Freq = read theta0FreqsStr :: Double
      theta0Freqs = [-theta0Freq .. theta0Freq]
      thetaFreq = read thetaFreqsStr :: Double
      thetaFreqs = [-thetaFreq .. thetaFreq]
      initDist = read initDistStr :: [R2S1RPPoint]
      alpha = read alphaStr :: Double
      pinwheelFlag = read pinwheelFlagStr :: Bool
      numThread = read numThreadStr :: Int
      sourceDist = L.take 1 initDist
      sinkDist = L.drop 1 initDist
      folderPath = "output/test/STCR2Z1T0_1"
  flag <- doesFileExist histFilePath
  arrR2Z1T0' <-
    if pinwheelFlag
      then computeR2Z1T0Array numPoint numPoint alpha thetaFreqs theta0Freqs
      else if flag
             then getNormalizedHistogramArr <$> decodeFile histFilePath
             else solveMonteCarloR2Z1T0
                    numThread
                    numTrail
                    maxTrail
                    numPoint
                    numPoint
                    sigma
                    tao
                    len
                    theta0Freqs
                    thetaFreqs
                    histFilePath
                    (emptyHistogram
                       [ numPoint
                       , numPoint
                       , L.length theta0Freqs
                       , L.length thetaFreqs
                       ]
                       0)
  let arrR2Z1T0 =
        computeUnboxedS .
        pad [numPoint, numPoint, L.length theta0Freqs, L.length thetaFreqs] 0 .
        downsample [2, 2, 1, 1] $
        arrR2Z1T0'
  createDirectoryIfMissing True folderPath
  plan <- makeR2Z1T0Plan emptyPlan arrR2Z1T0
  sourceDistArr <-
    computeInitialDistributionR2Z1T0
      plan
      numPoint
      numPoint
      thetaFreqs
      theta0Freqs
      sourceDist
  sinkDistArr <-
    computeInitialDistributionR2Z1T0
      plan
      numPoint
      numPoint
      thetaFreqs
      theta0Freqs
      sinkDist
  arrR2Z1T0F <- dftR2Z1T0 plan . makeFilterR2Z1T0 $ arrR2Z1T0
  arrR2Z1T0TRF <-
    dftR2Z1T0 plan . makeFilterR2Z1T0 . timeReverseR2Z1T0 thetaFreqs theta0Freqs $
    arrR2Z1T0
  -- Source field
  sourceArr <- convolveR2Z1T0 plan arrR2Z1T0F sourceDistArr
  sourceR2 <- R.sumP . R.sumS . rotate4D . rotate4D $ sourceArr
  let sourceField =
        computeS . R.extend (Z :. (1 :: Int) :. All :. All) $ sourceR2
  plotImageRepaComplex (folderPath </> "Source.png") . ImageRepa 8 $ sourceField
  -- Sink field
  sinkArr <- convolveR2Z1T0 plan arrR2Z1T0TRF sinkDistArr
  sinkR2 <- R.sumP . R.sumS . rotate4D . rotate4D $ sinkArr
  let sinkField = computeS . R.extend (Z :. (1 :: Int) :. All :. All) $ sinkR2
  plotImageRepaComplex (folderPath </> "Sink.png") . ImageRepa 8 $ sinkField
  -- Completion Filed
  completionFiled <- convolveR2Z1' plan thetaFreqs theta0Freqs sourceArr sinkArr
  completionFiledR2 <- R.sumP . R.sumS . rotate4D . rotate4D $ completionFiled
  plotImageRepa (folderPath </> "Completion.png") .
    ImageRepa 8 .
    computeS . R.extend (Z :. (1 :: Int) :. All :. All) . R.map magnitude $
    completionFiledR2


module IQ where

import Data.Complex

type IQ = Complex Float


{-# LANGUAGE DataKinds, ExistentialQuantification, FlexibleContexts         #-}
{-# LANGUAGE FlexibleInstances, MultiParamTypeClasses, TypeSynonymInstances #-}
{-# LANGUAGE UndecidableInstances                                           #-}
{-# OPTIONS_GHC -fno-warn-orphans #-}
module Numeric.Algebra.Ring.Noetherian (Noetherian) where
import qualified Data.Complex              as C
import           Data.Function
import           Data.Ord
import           Data.Ratio
import           Numeric.Algebra
import qualified Numeric.Algebra.Complex   as NA
import           Numeric.Algebra.Instances
import           Prelude                   hiding (negate, subtract, (*), (+),
                                            (-))
import qualified Prelude                   as P
-- | Noetherian ring (i.e. every ideal is finitely generated).
class (Commutative r, Ring r) => Noetherian r where

instance Noetherian Int where

instance Noetherian Integer where

instance (Commutative (NA.Complex r), Ring (NA.Complex r)) => Noetherian (NA.Complex r) where
instance (Commutative (C.Complex r), Ring (C.Complex r)) => Noetherian (C.Complex r) where
instance Integral n => Noetherian (Ratio n)


{-# LANGUAGE BangPatterns        #-}
{-# LANGUAGE DataKinds           #-}
{-# LANGUAGE GADTs               #-}
{-# LANGUAGE KindSignatures      #-}
{-# LANGUAGE LambdaCase          #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE StandaloneDeriving  #-}
{-# LANGUAGE TypeOperators       #-}

module BPTT where

import Control.Monad
import Control.Monad.Random
import Data.Char
import Data.List
import Data.Maybe
import Data.Singletons
import Data.Singletons.Prelude
import Data.Singletons.TypeLits
import GHC.Show
import Numeric (showIntAtBase)
import qualified Numeric.LinearAlgebra as LA
import Numeric.LinearAlgebra.Static

-- dependently typed weights
data Weights i o = W {
    wBiases :: !(R o)
  , wNodes  :: !(L o i)
  }

-- an simple 3-layer RNN, including one recurrent layer,
-- is not like the Network type discussed in Justin Le's
-- blog series (because not all layers are just stacked
-- one on top of the other.)
-- Instead:
data RNN i h o = MkRNN {
    getU    :: !(Weights i h)
  , getW    :: !(Weights h h)
  , getV    :: !(Weights h o)
  }

-- If one were to leverage a.g. ad the Hackage library for
-- autodifferentiable functions, then functions being used for
-- activation could be provided in a symbolic form amenable to
-- being automatically differentiated (for use in the backprop
-- process.) Instead, each activation function used & its derivative
-- are provided separately.

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

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

softmax :: KnownNat v => R v -> R v
softmax v =
  let vexp = exp v
  in  vexp / (konst $ norm_1 vexp)

-- The formula for the (x, y)-entry of softmax',
-- which is a matrix, is
--
-- softmax'[x, y] = -softmax[x]    * softmax[y],    x /= y
--                  (1-softmax[x]) * softmax[y],    x == y
--
-- which is the same as adding the vector
-- 'softmax v' diagonally down the negative of
-- said vector's outer product with itself.
-- I would prefer to build this using 'build', but
-- I don't know how right now.
softmax' :: KnownNat v => R v -> L v v
softmax' v =
  let s = softmax v
  in  (diag s) + konst (-1) * (s `outer` s)

-- runLayer is for performing the linear mapping
-- due to a particular layer's set of parameters,
-- on an appropriately sized vector
runLayer :: (KnownNat i, KnownNat o)
         => Weights i o
         -> R i -> R o
runLayer (W wB wN) v_in =
  wB + wN #> v_in

-- Verbose 'run RNN', for exposing all intermediate computations.
-- (R i, R h,  R h, R o, R o)
--  x_t  h_t-1 h_t  z_t  y_t

-- runRNN can be expressed in terms of this, if I locate / find
-- generalized tuple accessors
runRNN_verbose  :: (KnownNat i, KnownNat h, KnownNat o)
                => RNN i h o
                -> R i -> R h
                -> (R i, R h, R h, R o, R o)
runRNN_verbose (MkRNN rnn_U rnn_W rnn_V)
  x_t h_t_1 =
    let h_t = logistic $
                -- recurrent parameters, weighting the prior hidden state
                (runLayer rnn_W h_t_1) +
                -- input parameters, weighting the current input
                (runLayer rnn_U x_t)
        -- output parameters, weighting the current hidden state
        z_t = runLayer rnn_V h_t
        -- the final activation provides the prediction
        y_t = logistic z_t
    in  (x_t, h_t_1, h_t, z_t, y_t)

-- Forward phase of training an RNN.
-- No adjustments to the parameters (weights) are made.
-- Instead, foldl is used to accumulate a list of tuples,
-- each tuple bearing all input, intermediate and output vectors
-- for one of the steps (t := 1, 2, ..., n) in the forward direction.
forward_phase :: (KnownNat i, KnownNat h, KnownNat o)
              => RNN i h o    -- RNN
              -> [(R i, R o)] -- list of training inputs paired with target outputs
              -> R h          -- initial hidden state

                              -- All tuples of resulting vectors, from every step carried out:
                              -- For t ranging from 1 to n, n the number of training samples:
              -> [( R i,      -- x_t
                    R h,      -- h_t-1
                    R h,      -- h_t
                    R o,      -- z_t
                    R o,      -- y_t
                    R o)]     -- tgt_t
forward_phase
  rnn
  pairs
  h_0 =
    snd $ foldl unroll (h_0, []) pairs
    where
      -- unroll :: (R h, [(R i, R h, R h, R o, R o, R o)]) -> (R i, R o)
      --        -> (R h, [(R i, R h, R h, R o, R o, R o)])

      unroll (h_t_1, results) pair@(x_t, tgt_t) =
        --                               ^ tgt_t plays no role in the forward direction,
        --                                 but we pattern-match on it so that we can re-package it
        --                                 up into this step's tuple for the result list
        let result@(_, _, h_t, z_t, y_t) = runRNN_verbose rnn x_t h_t_1
        in  (h_t, results ++ [(x_t, h_t_1, h_t, z_t, y_t, tgt_t)])

-- Backward phase of training an RNN.
-- Backpropagation is required.
-- There's a lot that goes into backprop!
-- (Basically, everything that is used or
-- produced at any point in the forward step.)
backward_phase  :: (KnownNat i, KnownNat h, KnownNat o)
                => RNN i h o   -- an RNN to train

                                -- All tuples of resulting vectors, from every step carried out:
                                -- For t ranging from 1 to n, n the number of training samples:
                -> [( R i,      -- x_t
                      R h,      -- h_t-1
                      R h,      -- h_t
                      R o,      -- z_t
                      R o,      -- y_t
                      R o)]     -- tgt_t

                -- At every step #t of BPTT, there are *two* sources of "delta-h_t" the bump
                -- to h_t required for the sake of trying to bump down the error at that same step:
                -- a) a bump due to what W is at the current time, "delta-h given W", and
                -- b) a bump *propagated backwards* from the next step #(t+1).
                -- But in the base case, this dh_t_dh_t' is chosen as 0 (because it's undefined.)

                -- Note the *symmetry* with forward_phase:
                --    just as we 'primed' forward_phase  with an initial hidden state,    h_0        : 0,
                --    we also    'prime'  backward_phase with an initial back-prop'ed dh, dh_n_dh_n' : 0.
                -> R h        -- dh_n_dh_n'

                -> Double     -- a training rate
                -> RNN i h o  -- the trained RNN
backward_phase rnn intermediates dh_n_dh_n' rate =
  snd $ foldr bptt (dh_n_dh_n', rnn) intermediates
    where
      -- -- A BPTT step (at time t) consumes the following data:
      -- --       x_t, h_t, h_t_1, z_t,  y_t, tgt_t
      -- bptt :: (R i, R h, R h,   R o,  R o, R o)
      -- -- It also "accumulates" by swapping out,
      -- -- based upon time-t data  <dh_t   | dh_t+1>, rnn
      --          -> (R h, RNN i h o)
      -- -- the time-(t-1) data     <dh_t_1 | dh_t>,   rnn';
      -- -- in this way, the rnn' "accumulates" training bumps to the parameters
      --          -> (R h, RNN i h o)

      bptt    (x_t, h_t, h_t_1, z_t, y_t, tgt_t)
              (dh_t_dh_t', MkRNN
                -- We'll pattern match into each Weights term u, w, v for convenience
                -- later in deriving the desired infinitesimal bumps.
                u@(W uB uN)
                w@(W wB wN)
                v@(W vB vN)) =
                -- We are hard-coding the use of squared-error for the error function:
                -- e(y_t) := (y_t - tgt_t) ^ 2
                -- whose derivative with respect to the pre-activation z_t is:

                {-
                  de/dz = de/dy           * dy/dz
                        = (e' at y_t)     * (y' at z_t)
                        = 2 (y_t - tgt_t) * (logistic' z_t)
                  These 2 factors appear (in reverse order) as 'dedz' below;
                  the factor of 2 is discarded & we defer all control of the
                  desired training rate to the scalar input param 'rate'.
                -}
            let dedz          = logistic' z_t * (y_t - tgt_t)
                              -- = softmax' z_t #> (y_t - tgt_t)

                {-
                  Informally:
                  dedz is a vector that reports,
                    for a unit "nudge" to each component z_t(j) ->
                      size / direction of resulting "nudge" for error component e_t(j).
                  The choice of dz_t below corresponds to the following training strategy:
                    a) bump each z-component in the direction that will *decrease* its
                       matching e-component (thus, factor of '-1')
                    b) but re-scaled based on the learning rate desired (thus, 'scale rate')
                -}
                dz_t          = -1 * konst rate * dedz

                {-
                  Informally:
                  Having set out to make these infinitesimal changes to the z_t components,
                  we need to figure out
                    what changes to V,    *given* what h_t is,     and
                    what changes to h_t,  *given* what V is
                  will help us accomplish that.

                  Observe that
                    z  = vN h              + vB
                  therefore (product rule on 1st summand)
                    dz = (dvN h) + (vN dh) + dvB

                  we do *not* try to formally solve for dvN, dh, dvB in here, but rather use these
                  as instructions for how to recombine the vectors dz_t / h_t and matrix V
                  so as to arrive at dvN_h ("bump to v's matrix part, given h") and
                  dh_vN ("bump to h, given v's matrix part").

                  As for v's bias-vector part, vB, we simply nudge it directly by the dz (z-nudge)
                  we worked out above. Why? Look at the infinitesimals equation above - a nudge
                  dvB performed on the bias vector will show up directly in dz.
                -}
                dvB           = dz_t
                {-
                  Instead of trying to solve for dvN, dimensional considerations lead us to this formula
                  for dvN in terms of dz_t, h_t:
                    vN   goes    from h-vectors  to o-vectors; i.e.    o * h -matrix
                    h_t            an h-vector
                    dz_t                         an o-vector
                    so the way to combine these to get a "bumps matrix" compatible with vN is
                    dz_t `outer` h_t   <- "column" o-vector times "row" h-vector
                -}
                dvN_h         = dz_t `outer` h_t
                {-
                  Similar to the above, but for building a "bumps h-vector" out of o * h matrix
                  and o-vector.
                -}
                dh_vN         = tr vN #> dz_t
                {-
                  See the note in the type signature of 'backward_phase' about the 2 sources of
                  training bumps to h_t at every step t.
                  Read dh_t_dh_t' as (bump to h at current time t,
                                      propped backwards from bump to h at future time t+1)
                -}
                dh_t          = dh_vN + dh_t_dh_t'

                -- Having now characterized in what way we are going to bump the
                -- entries in the current hidden state h_t, so as to achieve
                -- a desired reduction in error (i.e. dh_t), we now use it to
                -- determine how we should modify the intermediate preactivation,
                -- calling it q, q_t := (W * h_t_1 + U * x_t) to achieve the same goal:

                -- de / dq = (de / dh) * (dh / dq)  -- chain rule
                --         = dh_t      * (dh / dq)  -- back propagation
                --         = dh_t      * (g' at q)  -- h is the activation (g) of q

                dedq          = dh_t * logistic' q
                                  where q = runLayer w h_t_1 + runLayer u x_t

                -- and work out a desired perturbation to q, based upon scaling dedq
                dq_t          = (-1) * konst rate * dedq

                {-
                  q  = W * h_t_1 + U * x_t
                  dq = (dW * h_t_1) + (W * dh_t_1)
                        + (dU * x_t) + [the term in dx_t -> 0, since x constant]
                -}

                {-
                  Recurrent contribution to time-t hidden state:
                  q_t  = wN h_t_1                   + wB        + ...
                  dq_t = (dwN h_t_1) + (wN dh_t_1)  + dwB       + ...
                -}
                dwB           = dq_t
                -- dq_t and h_t_1 both happen to be h-vectors having the same
                -- dimensionality, so here the driving intuition is that such
                -- "matrix-bumps" take the form [destination] `outer` [origin]
                dwN_h         = dq_t `outer` h_t_1

                -- This is what will be given back from the current fold instance
                -- to the next one
                dh_t_1_dh_t   = tr wN #> dq_t
                {-
                  Input contribution to time-t hidden state:
                  q_t  = ... + uN x_t                 + uB
                  dq_t = ... + (duN x_t) + (uN dx_t)  + duB
                       = ... + (duN x_t) + ZERO       + duB
                                        -- ^ x_t, the training inputs, are constant!
                -}
                duB           = dq_t
                duN_x         = dq_t `outer` x_t

                -- Don't forget to apply all the RNN-training nudges!
                vB'           = vB + dvB
                vN'           = vN + dvN_h
                wB'           = wB + dwB
                wN'           = wN + dwN_h
                uB'           = uB + duB
                uN'           = uN + duN_x

            in  (dh_t_1_dh_t, MkRNN (W uB' uN')
                                    (W wB' wN')
                                    (W vB' vN'))

{-
 Making some simple training data.

 Here's a simple sentential grammar for producing "sentences"
 of ints, that later can go in an overall "document":

  <sentence>        ::= <sentence-start> <sentence-middle> <sentence-finish>
  <sentence-start>  ::= a multiple of 20, 20 * s, for some s : 1 <= s <= 10
  <sentence-middle> ::= a sequence of l integers, for some l : 0 <= l <= 10
                        that are derived from the sentence-start as
                          s - m, (s - m) + n, (s - m + n) - m, (s - m + n - m) + n, ...
                        for some choice of parameters      m, n : 1 <= m, n <= 5
  <sentence-finish> ::= a multiple of 17, 17 * f, for some f : 1 <= f <= 10
-}

-- Generate a "biarithmetic" sequence for sentence-middle, given
-- a base number s, a down size m, an up size n
biarithmetic :: Int -> Int -> Int -> [Int]
biarithmetic    s      m      n   =  concat $
 [ [ s - (j * m) + (k * n) | k <- [j-1 .. j] ]
                           | j <- [1..] ]

-- For use as a multiple of 20
genStartParam :: MonadRandom m => m Int
genStartParam = getRandomR (3, 10)

-- For use as a mid-sequence length
genMiddleLenParam :: MonadRandom m => m Int
genMiddleLenParam = getRandomR (0, 10)

-- For use as the down size and up size of a biarithmetic sequence
genMiddleDownUpParams :: MonadRandom m => m (Int, Int)
genMiddleDownUpParams = do
  m <- getRandomR (1, 5)
  n <- getRandomR (1, 5)
  return (m, n)

-- For use as a multiple of 17
genFinishParam :: MonadRandom m => m Int
genFinishParam = getRandomR (1, 10)

-- For use as a number of "sentences".
genSentenceCount :: MonadRandom m => m Int
genSentenceCount = getRandomR (500, 1000)

type Sentence = [Int]

genSentence :: MonadRandom m => m Sentence
genSentence = do
  -- 1. Get sentence-start
  s       <- genStartParam
  let sent_0 = s * 20
  -- 2. Get sentence-middle length
  l       <- genMiddleLenParam
  -- 3. Get sentence up & down size params
  (m, n)  <- genMiddleDownUpParams
  -- 4. Get sentence-middle
  let sent_mid = take l $ biarithmetic sent_0 m n
  -- 5. Get sentence-finish
  f       <- genFinishParam
  let sent_f = f * 17
  -- 6. Return random sentence
  return (sent_0 : sent_mid ++ [sent_f])

type Document = [Sentence]

genDocument :: MonadRandom m => m Document
genDocument = do
  -- 1. Get number of sentences
  sent_count <- genSentenceCount
  -- 2. For as many sentences as 'sent_count', generate one
  --    random sentence each & then combine the results inside
  --    MonadRandom context
  let rand_sent_stream = repeat genSentence
  sentences <- sequenceA $ take sent_count rand_sent_stream
  return sentences

{-
 Now that we have a way of generating "documents" made up out of
  numerical "sentences", we need to convert this into a form that
  would be characteristically used to train a neural net.
  Convert each integer to a (possibly sparse) vector. How?
  Every integer present in a (d :: Document) should always lie
  within the range (0 <= j <= 256), so we could use a binary representation
  taking up 8 bits (tossing in an extra 2 to get 10, for good measure)
-}

-- The full source module citation is probably forgivable if I only do it this once.
-- Best practice is to 'import qualified' with an alias intToVector :: Int -> R 10
intToVector m =
  let binary_string = reverse $ showIntAtBase 2 intToDigit m "" -- :: [Char]
      binary_padded = binary_string ++
                        (take (10 - length binary_string) $ repeat '0')
      doubles_list  = (fromIntegral . digitToInt) <$> binary_padded
      doubles_vect  = Numeric.LinearAlgebra.Static.fromList doubles_list
  in  doubles_vect

-- The other way round, for casting RNN predictions back into Int's.
-- This is unsafe in case the entries of any y_t end up rounding to
-- numbers other than 0, 1.
-- (Does the hmatrix API support a more direct way of casting from
-- R 10 to the list of rounded numbers?)
vectorToInt :: R 10 -> Int
vectorToInt v =
  let bins = (fmap toInteger) <$> LA.toList
                                $ LA.toZ
                                $ LA.roundVector
                                $ extract v
  in  fromIntegral $
        foldr (\p@(bit,  power)  total ->
                   bit * power + total)
        0
        (zip bins [ 2^j | j <- [0..] ])

zero_vec :: KnownNat n => R n
zero_vec = 0

-- via Justin Le's blog series on neural nets
randomWeights :: (MonadRandom m, KnownNat i, KnownNat o)
              => m (Weights i o)
randomWeights = do
    s1 :: Int <- getRandom
    s2 :: Int <- getRandom
    let wB = randomVector  s1 Uniform * 2 - 1
        wN = uniformSample s2 (-1) 1
    return $ W wB wN

randomRNN :: (MonadRandom m, KnownNat i, KnownNat h, KnownNat o)
          => m (RNN i h o)
randomRNN =
  (pure MkRNN)  <*> randomWeights
                <*> randomWeights
                <*> randomWeights

trainRNN :: (KnownNat i, KnownNat h, KnownNat o)
         => RNN i h o     -- an RNN
         -> [(R i, R o)]  -- the training data: pairs of inputs & training outputs
         -> Double
         -> RNN i h o     -- an RNN with trained parameters
trainRNN rnn train_data rate =
  let intermediates = forward_phase rnn train_data zero_vec
      trained       = backward_phase rnn intermediates zero_vec rate
  in  trained

-- Before re-formatting into a vector form suitable for use
-- in training, we have Document ~ [Sentence] ~ [[Int]];
-- now, turn each Int into a binary-repr vector, R 10
genSampleData :: MonadRandom m => m [[ R 10 ]]
genSampleData = do
  doc <- genDocument
  return $ (intToVector <$>) <$> doc

-- Every sentence (sublist) in the document (list of lists) will
-- furnish one sequence of training data for the RNN to "unfold"
-- over and perform BPTT on. So, BPTT is performed once per sentence
-- in the document overall.
pairEachWithNext :: [x] -> [(x, x)]
pairEachWithNext xs =
  zip (init xs) (tail xs)

-- In the special case of an 'RNN i h o' where o = i,
-- we can pipe back a prediction y_t due to x_t and an
-- initial state h, to furnish the next x_(t+1) for further
-- predictions. We'll use this to observe the way that a
-- trained RNN behaves.
rnnStep :: (KnownNat i, KnownNat h)
        => RNN i h i
        -> (R h, R i)
        -> (R h, R i)
rnnStep rnn (h_t_1, x_t) =
  let full_output = runRNN_verbose rnn x_t h_t_1
      third five@(_, _, x3, _, _)  = x3
      fifth five@(_, _, _,  _, x5) = x5
  in  (third full_output, fifth full_output)

getDocumentIO :: IO Document
getDocumentIO = evalRandIO genDocument

takePredictions  :: (KnownNat i, KnownNat h)
                 => RNN i h i
                 -> (R h, R i)
                 -> Int
                 -> [R i]
takePredictions rnn (h_0, x_1) n =
  let hxs_stream = take (n+1) $ iterate (rnnStep rnn) (h_0, x_1)
  in  snd <$> hxs_stream

takePredsFromN :: KnownNat h
               => RNN 10 h 10
               -> Int      -- word of a sentence
               -> Int      -- number of predictions
               -> [R 10]
takePredsFromN rnn x_1 n =
  takePredictions rnn (zero_vec, intToVector x_1) n

trainRNNOnSentence  :: KnownNat h
                    => RNN 10 h 10
                    -> Sentence
                    -> Double
                    -> RNN 10 h 10
trainRNNOnSentence rnn sent rate =
  let training_data = pairEachWithNext (intToVector <$> sent)
  in  trainRNN rnn training_data rate

-- trainRNNOnDocument has to be structured like a fold
-- so that the trained RNN due to each sentence can be
-- passed to the following sentence (at the same time,
-- I think that hidden state "restarts" with the beginning
-- of a new sentence.)
trainRNNOnDocument  :: KnownNat h
                    => RNN 10 h 10
                    -> Document
                    -> Double
                    -> RNN 10 h 10
trainRNNOnDocument rnn sents rate =
  foldl (\rnn sent -> trainRNNOnSentence rnn sent rate)
        rnn
        sents

observeUntrainedPredictions :: Int    -- initial word
                            -> Int    -- number of predictions
                            -> IO ()  -- action printing info to console
observeUntrainedPredictions n_0 num = do
  -- get an RNN of random parameters
  putStrLn "Getting a random RNN ..."
  (rnn :: RNN 10 300 10) <- evalRandIO randomRNN
  -- get a series of predictions from the untrained RNN
  putStrLn $ "Producing: " ++ show num ++
             " predictions, starting from: " ++ show n_0
  let vs = takePredsFromN rnn n_0 num
  let ns = vectorToInt <$> vs
  putStrLn "The RNN has predicted: "
  print ns

observeTrainedPredictions :: Int    -- initial word
                          -> Int    -- number of predictions
                          -> Double -- learning rate
                          -> IO ()  -- action printing info to console
observeTrainedPredictions n_0 num rate = do
  -- get an RNN of random parameters
  putStrLn "Getting a random RNN ..."
  (rnn :: RNN 10 300 10) <- evalRandIO randomRNN
  -- get a random document
  putStrLn "Getting a random document to train on ..."
  document <- evalRandIO genDocument
  -- train RNN on the document
  putStrLn "Training the RNN on the document ..."
  let rnn' = trainRNNOnDocument rnn document rate
  -- get a series of predictions from the trained RNN
  putStrLn $ "Producing: " ++ show num ++
             " predictions, starting from: " ++ show n_0 ++
             " and using a trained RNN"
  let vs = takePredsFromN rnn' n_0 num
  let ns = vectorToInt <$> vs
  putStrLn "The trained RNN has predicted: "
  print ns

observeUntrainedVsTrained :: Int    -- initial word
                          -> Int    -- number of predictions
                          -> Double -- learning rate
                          -> IO ()  -- action printing info to console
observeUntrainedVsTrained n_0 num rate = do
  -- get an RNN of random parameters
  putStrLn "Getting a random RNN ..."
  (rnn :: RNN 10 300 10) <- evalRandIO randomRNN
  -- get a series of predictions from the untrained RNN
  putStrLn $ "Producing: " ++ show num ++
             " predictions, starting from: " ++ show n_0
  let vs = takePredsFromN rnn n_0 num
  let ns = vectorToInt <$> vs
  putStrLn "The RNN has predicted: "
  print ns
  -- get a random document
  putStrLn "Getting a random document to train on ..."
  document <- evalRandIO genDocument
  -- train RNN on the document
  putStrLn "Training the RNN on the document ..."
  let rnn' = trainRNNOnDocument rnn document rate
  -- get a series of predictions from the trained RNN
  putStrLn $ "Producing: " ++ show num ++
             " predictions, starting from: " ++ show n_0 ++
             " and using a trained RNN"
  let vs' = takePredsFromN rnn' n_0 num
  let ns' = vectorToInt <$> vs'
  putStrLn "The trained RNN has predicted: "
  print ns'


module Marvin.API.Algorithms.Internal.GradientDescentSpec where

import Test.Hspec

import Marvin.API.Algorithms.Internal.GradientDescent hiding (learningRate)
import qualified Marvin.API.Algorithms.Internal.GradientDescent as GD

import Marvin.API.Fallible
import Marvin.API.Table.Internal
import Marvin.Test.TestUtils

import qualified Numeric.LinearAlgebra as LA

-- todo lambda, addIntercept
spec :: Spec
spec =
  describe "gradient descent" $ do
    it "computes cost" $
      (linearRegressionCost fm rv tt) `shouldSatisfy` (isAround 50.376562500000006) . fst
    it "does linear regression with lin.alg. lib" $
      (let theta = gradientDescent gradDesc fm rv in
        ((fst . linearRegressionCost fm rv) theta, theta))
      `shouldSatisfy` \(cost, theta) ->
          cost `isAround` 0.293648087854516 &&
          theta `isAround` (LA.fromList expectedTheta)
  where
    tt = initModelParams gradDesc fm rv
    fm = addInterceptFeature $ LA.fromRows $ fmap LA.vector xData
    rv = LA.vector yData
    xData = map init trainMatrix
    yData = map last trainMatrix
    gradDesc = defaultGradientDescent {
      numIter = 30000
      , GD.learningRate = 0.005
      , cost = linearRegressionCost
      , lambda = 0
      , addIntercept = False
    }
    expectedTheta =
      [13.781001191549892,
      3.900308511353248e-2,
      -8.360245732382168e-2,
      2.8699265280611834e-2,
      2.7797762929208453e-2,
      0.13913675368346703,
      3.900308511353248e-2,
      0.11340373429004338,
      0.1166082670688832,
      4.446542254295697e-2,
      0.14211000816593214,
      4.446542254295697e-2]


{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE Strict           #-}
module FourierPinwheel.AsteriskGaussian where

import           Data.Array.Repa             as R
import           Data.Complex
import           Data.List                   as L
import           Data.Vector.Generic         as VG
import           Data.Vector.Storable        as VS
import           Data.Vector.Unboxed         as VU
import           DFT.Plan
import           Filter.Utils
import           FourierPinwheel.Array
import           FourierPinwheel.Hypergeo1F1
import           Math.Gamma
import           Pinwheel.FourierSeries2D
import           Utils.Distribution
import           Utils.Parallel

-- The envelope in R2 is r ^ alpha, where -2 < alpha < -0.5
{-# INLINE asteriskGaussian #-}
asteriskGaussian ::
     ( VG.Vector vector (Complex e)
     , RealFloat e
     , NFData (vector (Complex e))
     , Gamma (Complex e)
     , Unbox e
     )
  => Int
  -> Int
  -> e
  -> e
  -> e
  -> e
  -> [vector (Complex e)]
asteriskGaussian numR2Freqs thetaFreq alpha periodR2 periodEnv std =
  parMap
    rdeepseq
    (\angularFreq ->
       let arr =
             R.map
               (* (sqrt (pi / std) *
                   exp (fromIntegral angularFreq ^ 2 / (-4) / std) :+
                   0)) .
             centerHollowArray numR2Freqs $
             analyticalFourierCoefficients1
               numR2Freqs
               1
               angularFreq
               0
               alpha
               periodR2
               periodEnv
        in VG.convert . toUnboxed . computeS $ arr)
    [-thetaFreq .. thetaFreq]

{-# INLINE asteriskGaussianEnvelope #-}
asteriskGaussianEnvelope ::
     (R.Source s (Complex Double))
  => DFTPlan
  -> Int
  -> Int
  -> Double
  -> Double
  -> Double
  -> Double
  -> R.Array s DIM2 (Complex Double)
  -> IO (VS.Vector (Complex Double))
asteriskGaussianEnvelope plan numR2Freqs thetaFreq alpha periodR2 periodEnv stdTheta filter = do
  let planID = DFTPlanID DFT1DG [numR2Freqs, numR2Freqs] [0, 1]
      inversePlanID = DFTPlanID IDFT1DG [numR2Freqs, numR2Freqs] [0, 1]
      vecs =
        asteriskGaussian numR2Freqs thetaFreq alpha periodR2 periodEnv stdTheta
  filterF <-
    dftExecute plan planID . VU.convert . toUnboxed . computeS . makeFilter2D $
    filter
  asteriskGaussianF <- dftExecuteBatchP plan planID vecs
  fmap VS.concat .
    dftExecuteBatchP plan inversePlanID .
    parMap rdeepseq (VS.zipWith (*) filterF) $
    asteriskGaussianF

-- The envelope is r^alpha X Gaussian
asteriskGaussian2 ::
     DFTPlan
  -> Int
  -> Int
  -> Double
  -> Double
  -> Double
  -> Double
  -> Double
  -> IO (VS.Vector (Complex Double))
asteriskGaussian2 plan numR2Freqs thetaFreq alpha periodR2 periodEnv stdR2 stdTheta = do
  let centerFreq = div numR2Freqs 2
      gaussian2D =
        fromFunction (Z :. numR2Freqs :. numR2Freqs) $ \(Z :. i' :. j') ->
          let i = i' - centerFreq
              j = j' - centerFreq
           in exp
                (pi * fromIntegral (i ^ 2 + j ^ 2) /
                 ((-1) * periodR2 ^ 2 * stdR2 ^ 2)) /
              (2 * pi * stdR2 ^ 2) :+
              0
  asteriskGaussianEnvelope
    plan
    numR2Freqs
    thetaFreq
    alpha
    periodR2
    periodEnv
    stdTheta
    gaussian2D

-- The envelope is r^alpha X LowPass
asteriskGaussianLowPass ::
     DFTPlan
  -> Int
  -> Int
  -> Double
  -> Double
  -> Double
  -> Double
  -> Double
  -> IO (VS.Vector (Complex Double))
asteriskGaussianLowPass plan numR2Freqs thetaFreq alpha periodR2 periodEnv stdTheta radius = do
  let lpf =
        fromUnboxed (Z :. numR2Freqs :. numR2Freqs) $
        idealLowPassFilter radius periodR2 numR2Freqs
  asteriskGaussianEnvelope
    plan
    numR2Freqs
    thetaFreq
    alpha
    periodR2
    periodEnv
    stdTheta
    lpf

{-# INLINE asteriskGaussianFull #-}
asteriskGaussianFull ::
     ( VG.Vector vector (Complex e)
     , RealFloat e
     , NFData (vector (Complex e))
     , Gamma (Complex e)
     , Unbox e
     )
  => Int
  -> Int
  -> Int
  -> e
  -> e
  -> e
  -> e
  -> e
  -> [vector (Complex e)]
asteriskGaussianFull numR2Freqs thetaFreq rFreq alpha periodR2 periodEnv stdTheta stdR =
  let zeroVec = VG.replicate (numR2Freqs ^ 2) 0
   in parMap
        rdeepseq
        (\(radialFreq, angularFreq) ->
           let pinwheel =
                 analyticalFourierCoefficients1
                   numR2Freqs
                   1
                   angularFreq
                   radialFreq
                   alpha
                   periodR2
                   periodEnv
               -- arr =
               --   R.map
               --     (* ((-- gaussian1DFreq (fromIntegral angularFreq) stdTheta *
               --          gaussian1DFourierCoefficients
               --            (fromIntegral radialFreq)
               --            (log periodEnv)
               --            stdR
               --          :+
               --          0) -- *
               --         -- cis
               --         --   (2 * pi / log periodEnv * fromIntegral radialFreq *
               --         --    log 0.5)
               --        )) $
               --   -- if radialFreq == 0 && angularFreq == 0
               --   --   then pinwheel
               --   --   else
                   -- centerHollowArray numR2Freqs pinwheel
               arr =
                 R.map
                   (* ((1 / (1 + (2 * pi / log periodEnv * fromIntegral radialFreq )^2)) :+ 0)
                      ) $
                   centerHollowArray numR2Freqs pinwheel
            in if angularFreq == 0
                 then VG.convert . toUnboxed . computeS $ arr
                 else zeroVec)
        [ (radialFreq, angularFreq)
        | radialFreq <-  [-rFreq .. rFreq]
        , angularFreq <- [-thetaFreq .. thetaFreq]
        ]

{-# INLINE gaussianFull #-}
gaussianFull ::
     ( VG.Vector vector (Complex e)
     , RealFloat e
     , NFData (vector (Complex e))
     , Gamma (Complex e)
     , Unbox e
     )
  => Int
  -> Int
  -> Int
  -> e
  -> e
  -> [vector (Complex e)]
gaussianFull numR2Freqs thetaFreq rFreq periodR2 stdR2 =
  let zeroVec = VG.replicate (numR2Freqs ^ 2) 0
      periodEnv = periodR2 ^ 2 / 4
   in parMap
        rdeepseq
        (\(radialFreq, angularFreq) ->
           if angularFreq == 0 -- && radialFreq == 0
             then VG.convert . toUnboxed . computeS $
                  fromFunction (Z :. numR2Freqs :. numR2Freqs) $ \(Z :. i' :. j') ->
                    let i = fromIntegral $ i' - div numR2Freqs 2
                        j = fromIntegral $ j' - div numR2Freqs 2
                     in (gaussian2DFourierCoefficients i j periodR2 stdR2 :+ 0) /
                        (1 :+ (2 * pi / log periodEnv * fromIntegral radialFreq)^2)
             else zeroVec)
        [ (radialFreq, angularFreq)
        | radialFreq <- [-rFreq .. rFreq]
        , angularFreq <- [-thetaFreq .. thetaFreq]
        ]

{-# INLINE asteriskGaussianFullEnvelope #-}
asteriskGaussianFullEnvelope ::
     (R.Source s (Complex Double))
  => DFTPlan
  -> Int
  -> Int
  -> Int
  -> Double
  -> Double
  -> Double
  -> Double
  -> Double
  -> R.Array s DIM2 (Complex Double)
  -> IO (VS.Vector (Complex Double))
asteriskGaussianFullEnvelope plan numR2Freqs thetaFreq rFreq alpha periodR2 periodEnv stdTheta stdR filter = do
  let planID = DFTPlanID DFT1DG [numR2Freqs, numR2Freqs] [0, 1]
      inversePlanID = DFTPlanID IDFT1DG [numR2Freqs, numR2Freqs] [0, 1]
      vecs =
        asteriskGaussianFull
          numR2Freqs
          thetaFreq
          rFreq
          alpha
          periodR2
          periodEnv
          stdTheta
          stdR
      -- vecs =
      --  gaussianFull
      --      numR2Freqs
      --      thetaFreq
      --      rFreq
      --      periodR2
      --      stdR
  filterF <-
    dftExecute plan planID . VU.convert . toUnboxed . computeS . makeFilter2D $
    filter
  asteriskGaussianF <- dftExecuteBatchP plan planID vecs
  fmap VS.concat .
    dftExecuteBatchP plan inversePlanID .
    parMap rdeepseq (VS.zipWith (*) filterF) $
    asteriskGaussianF
  -- return . VS.concat $ vecs

-- The envelope is r^alpha X LowPass
asteriskGaussianFullLowPass ::
     DFTPlan
  -> Int
  -> Int
  -> Int
  -> Double
  -> Double
  -> Double
  -> Double
  -> Double
  -> Double
  -> IO (VS.Vector (Complex Double))
asteriskGaussianFullLowPass plan numR2Freqs thetaFreq rFreq alpha periodR2 periodEnv stdTheta stdR radius = do
  let lpf =
        fromUnboxed (Z :. numR2Freqs :. numR2Freqs) $
        idealLowPassFilter radius periodR2 numR2Freqs
  asteriskGaussianFullEnvelope
    plan
    numR2Freqs
    thetaFreq
    rFreq
    alpha
    periodR2
    periodEnv
    stdTheta
    stdR
    lpf

-- The envelope is r^alpha X LowPass
asteriskGaussian2Full ::
     DFTPlan
  -> Int
  -> Int
  -> Int
  -> Double
  -> Double
  -> Double
  -> Double
  -> Double
  -> Double
  -> IO (VS.Vector (Complex Double))
asteriskGaussian2Full plan numR2Freqs thetaFreq rFreq alpha periodR2 periodEnv stdTheta stdR stdR2 = do
  let centerFreq = div numR2Freqs 2
      gaussian2D =
        fromFunction (Z :. numR2Freqs :. numR2Freqs) $ \(Z :. i' :. j') ->
          let i = fromIntegral $ i' - centerFreq
              j = fromIntegral $ j' - centerFreq
           in gaussian2DFourierCoefficients i j periodR2 stdR2 :+ 0
  asteriskGaussianFullEnvelope
    plan
    numR2Freqs
    thetaFreq
    rFreq
    alpha
    periodR2
    periodEnv
    stdTheta
    stdR
    gaussian2D

asteriskGaussianRFull ::
     DFTPlan
  -> Int
  -> Int
  -> Int
  -> Double
  -> Double
  -> Double
  -> Double
  -> Double
  -> Double
  -> IO (VS.Vector (Complex Double))
asteriskGaussianRFull plan numR2Freqs thetaFreq rFreq alpha periodR2 periodEnv stdTheta stdR stdR2 = do
  let planID = DFTPlanID DFT1DG [numR2Freqs, numR2Freqs] [0, 1]
      inversePlanID = DFTPlanID IDFT1DG [numR2Freqs, numR2Freqs] [0, 1]
      centerFreq = div numR2Freqs 2
      gaussian2D =
        fromFunction (Z :. numR2Freqs :. numR2Freqs) $ \(Z :. i' :. j') ->
          let i = i' - centerFreq
              j = j' - centerFreq
           in exp
                (pi * fromIntegral (i ^ 2 + j ^ 2) /
                 ((-1) * periodR2 ^ 2 * stdR2 ^ 2)) /
              (2 * pi * stdR2 ^ 2) :+
              0
      r =
        centerHollowArray' numR2Freqs $ analyticalFourierCoefficients3 numR2Freqs 1 0 0 alpha periodR2 periodEnv
      vecs =
        asteriskGaussianFull
          numR2Freqs
          thetaFreq
          rFreq
          alpha
          periodR2
          periodEnv
          stdTheta
          stdR
  gaussian2DF <-
    dftExecute plan planID . VU.convert . toUnboxed . computeS . makeFilter2D $
    gaussian2D
  rF <-
    dftExecute plan planID . VU.convert . toUnboxed . computeS . makeFilter2D $
    r
  let filterF = VS.zipWith (\x y -> x + y ^ 2) gaussian2DF rF
  asteriskGaussianF <- dftExecuteBatchP plan planID vecs
  fmap VS.concat .
    dftExecuteBatchP plan inversePlanID .
    parMap rdeepseq (VS.zipWith (*) filterF) $
    asteriskGaussianF


{-# LANGUAGE RankNTypes #-}

module Observable.MCMC.MetropolisHastings (metropolisHastings) where

import Control.Monad.Primitive
import Control.Monad.State.Strict
import Data.HashMap.Strict (HashMap)
import qualified Data.HashMap.Strict as HashMap
import Data.Vector.Unboxed (Vector)
import qualified Data.Vector.Unboxed as V
import Observable.Core
import Observable.Types
import Statistics.Distribution
import Statistics.Distribution.Normal

sphericalGaussian :: Vector Double -> Vector Double -> Double -> Double
sphericalGaussian xs mu sd = product $ zipWith density normalDists xsAsList
  where
    xsAsList    = V.toList xs
    muAsList    = V.toList mu
    normalDists = map (`normalDistr` sd) muAsList

perturb
  :: PrimMonad m
  => Vector Double
  -> Double
  -> Observable m (Vector Double)
perturb q sd = V.mapM (`normal` sd) q

acceptRatio
  :: Target Double -> Vector Double -> Vector Double -> Double -> Double
acceptRatio target current proposed sd = exp . min 0 $
    logObjective target proposed + log (sphericalGaussian current proposed sd)
  - logObjective target current  - log (sphericalGaussian proposed current sd)

nextState
  :: Target Double
  -> Vector Double
  -> Vector Double
  -> Double
  -> Double
  -> Vector Double
nextState target current proposed sd z
    | z < acceptProb = proposed
    | otherwise      = current
  where
    ratio = acceptRatio target current proposed sd
    acceptProb | isNaN ratio = 0
               | otherwise   = ratio

getStandardDeviation :: Maybe Double -> OptionalStore -> Double
getStandardDeviation (Just sd) _   = sd
getStandardDeviation Nothing store = sd where
  (ODouble sd) = HashMap.lookupDefault (ODouble 1.0) MH store

updateStandardDeviation :: Double -> OptionalStore -> OptionalStore
updateStandardDeviation sd = HashMap.insert MH (ODouble sd)

metropolisHastings :: PrimMonad m => Maybe Double -> Transition m Double
metropolisHastings e = do
  Chain current target _ store <- get
  let sd = getStandardDeviation e store
  proposed <- lift $ perturb current sd
  zc       <- lift unit
  let next     = nextState target current proposed sd zc
      newStore = updateStandardDeviation sd store
  put $ Chain next target (logObjective target next) newStore
  return next


module Hw4
    (
    Vec(..),
    range,
    absolute,
    rd,
    dft,
    fft) where

import Data.Complex
import GHC.Base

range :: Double -> Double -> Double -> [Double]
absolute :: Vec (Complex Double) -> Vec Double
rd :: Int -> Vec Double -> Vec Double
dft :: [Double] -> Vec (Complex Double)
fft :: [Double] -> Vec (Complex Double)

newtype Vec a = Vec {runVec :: [a]}

instance Show a => Show (Vec a) where
    show (Vec a) = "[" ++ show' a ++ "]"

show' :: Show a => [a] -> [Char]
show' [] = ""
show' [x] = show x
show' (x:xs) = show x ++ " " ++ show' xs

instance Functor Vec where
    fmap g x = Vec $ g <$> runVec x

instance Applicative Vec where
    pure x = Vec $ repeat x
    liftA2 f a b =  Vec $ zipWith f (runVec a) (runVec b)

instance Semigroup (Vec a) where
    (<>) (Vec a) (Vec b) = Vec $ a ++ b

instance Monoid (Vec a) where
    mempty = Vec []

instance Num a => Num (Vec a) where
    fromInteger x = Vec $ pure $ fromInteger x
    (+) = liftA2 (+)
    (-) = liftA2 (-)
    (*) = liftA2 (*)
    negate (Vec a) = Vec $ fmap negate a
    abs (Vec a) = Vec $ fmap abs a
    signum (Vec a) = Vec $ fmap signum a

instance Fractional a => Fractional (Vec a) where
    (/) = liftA2 (/)
    fromRational x = Vec $ pure $ fromRational x

instance Floating a => Floating (Vec a) where
    exp (Vec a) = Vec (fmap exp a)
    log (Vec a) = Vec (fmap log a)
    sin (Vec a) = Vec (fmap sin a)
    cos (Vec a) = Vec (fmap cos a)
    asin (Vec a) = Vec (fmap asin a)
    acos (Vec a) = Vec (fmap acos a)
    atan (Vec a) = Vec (fmap atan a)
    sinh (Vec a) = Vec (fmap sinh a)
    cosh (Vec a) = Vec (fmap cosh a)
    asinh (Vec a) = Vec (fmap asinh a)
    acosh (Vec a) = Vec (fmap acosh a)
    atanh (Vec a) = Vec (fmap atanh a)
    pi = Vec $ pure pi

range from to count = fmap (\x -> from + x * step) [0..count-1]
    where step = (to - from)/count

absolute = fmap (\x -> sqrt(realPart x * realPart x + imagPart x * imagPart x))

rd n = fmap (\x -> fromIntegral (round $ c * x) / c)
  where c = 10^n

dft x =
    let n = fromIntegral $ length x
        index = range 0 n n
        xn = x `zip` index

        f k = sum y
                where
                    y = map factor xn
                    factor (xi, j) = let y = 2 * pi * j * k / n
                                        in xi * cos y :+ (-xi * sin y)
    in
        Vec $ fmap f index

fft x
    | n <= 16 = dft x
    | otherwise =
        let (even, odd) = split x
            (e, o) = (fft even, fft odd)
            t = -2 * pi / fromIntegral n

            g (e', o', k) =
                        let
                            y = t * fromIntegral k
                            (fr, fi) = (cos y, sin y)
                            (pr, pi) = (fr * realPart o' - fi * imagPart o', fr * imagPart o' + fi * realPart o')
                        in
                            ((realPart e' + pr) :+ (imagPart e' + pi), (realPart e' - pr) :+ (imagPart e' - pi))

            (lower, upper) = unzip $ fmap g $ zip3 (runVec e) (runVec o) [0..]
        in
            Vec (lower ++ upper)

    where
            n = length x
            split [] = ([], [])
            split [a] = ([a], [])
            split (a:b:c) = (a:x, b:y)
                where (x,y) = split c


module Main where

{-# LANGUAGE ScopedTypeVariables #-}

import Data.Maybe

import Graphics.UI.WX hiding (Event)
import Reactive.Banana
import Reactive.Banana.WX

import AI.HNN.FF.Network
import Numeric.LinearAlgebra

import Data.List.Split
import Data.List (replicate)

main :: IO ()
main = start gui

gui :: IO ()
gui = do
    f       <- frame [ text := "Neural Network" ]

    lEpocs          <- staticText f [ text := "Epocs"]
    lMiddleLayers   <- staticText f [ text := "Middle & output layers"]
    lLearningRate   <- staticText f [ text := "Learning rate"]

    iEpocs          <- textCtrlRich f [ text := "1000" ]
    iMiddleLayers   <- textCtrlRich f [ text := "2,2,1"]
    iLearningRate   <- textCtrlRich f [ text := "0.8"]

    bStart  <- button f [ text := "Start" ]
    bQuit   <- button f [ text := "Quit" ]

    set bStart [ on command := parseInputAndTrain f iEpocs iMiddleLayers iLearningRate]
    set bQuit  [ on command := close f]

    set f [ layout := column 25 [ widget lEpocs, widget iEpocs,
                                  widget lMiddleLayers, widget iMiddleLayers,
                                  widget lLearningRate, widget iLearningRate,
                                  widget bStart, widget bQuit ] ]

parseInputAndTrain :: (Paint w, Textual w1, Textual w2, Textual w3) => w -> w1 -> w2 -> w3 -> IO ()
parseInputAndTrain f iEpocs iMiddleLayers iLearningRate = do
    epocsString         <- get iEpocs text
    middleLayersString  <- get iMiddleLayers text
    learningRateString  <- get iLearningRate text
    let epocs = read epocsString :: Int
    let middleLayersSplited = splitOn "," middleLayersString
    let middleLayer = map (\s -> read s :: Int) middleLayersSplited
    let learningRate = read learningRateString :: Double
    trainNeuralNetwork epocs 2 middleLayer 1 learningRate

trainNeuralNetwork :: Int -> Int -> [Int] -> Int -> Double -> IO()
trainNeuralNetwork epocs firstLayer middleLayers outputLayer learningRate =  do

    xys <- readFile "xy.txt"
    zs <- readFile "z.txt"

    let list1 = parse xys
    let list2 = parse zs
    let list11 = map fromList list1
    let list22 = map fromList list2
    let samples2 = zip list11 list22

    putStrLn "1------------------"
    n <- createNetwork firstLayer middleLayers outputLayer :: IO (Network Double)
    mapM_ (putStrLn . show) samples2
    mapM_ (putStrLn . show . output n tanh . fst) samples2
    putStrLn "2------------------"
    let n' = trainNTimes epocs learningRate tanh tanh' n samples2
    mapM_ (putStrLn . show . output n' tanh . fst) samples2

parse :: String -> [[Double]]
parse a = map (map read . words) (lines a)

oneList :: [a] -> [b] -> [(a, b)]
oneList []     _      = []
oneList (x:xs) (y:ys) = (x, y) : oneList xs ys

samples :: Samples Double
samples = [ (fromList [0, 0], fromList [0])
          , (fromList [0, 1], fromList [1])
          , (fromList [1, 0], fromList [1])
          , (fromList [1, 1], fromList [0])
          ]


{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE CPP          #-}
{-# LANGUAGE ViewPatterns #-}
-- | Tests for Statistics.Math
module Tests.SpecFunctions (
  tests
  ) where

import Control.Monad
import Data.List
import Data.Maybe
import qualified Data.Vector as V
import           Data.Vector   ((!))
import qualified Data.Vector.Unboxed as U

import Test.QuickCheck  hiding (choose,within)
import Test.Tasty
import Test.Tasty.QuickCheck   (testProperty)
import Test.Tasty.HUnit

import Tests.Helpers
import Tests.SpecFunctions.Tables
import Numeric.SpecFunctions
import Numeric.SpecFunctions.Internal   (factorialTable)
import Numeric.MathFunctions.Comparison (within,relativeError,ulpDistance)
import Numeric.MathFunctions.Constants  (m_epsilon,m_tiny)

erfTol,erfcTol,erfcLargeTol :: Int
#if USE_SYSTEM_ERF && !defined(__GHCJS__)
erfTol       = 1
erfcTol      = 2
erfcLargeTol = 2
#else
erfTol       = 4
erfcTol      = 4
erfcLargeTol = 64
#endif

isGHCJS :: Bool
#if defined(__GHCJS__)
isGHCJS = True
#else
isGHCJS = False
#endif

tests :: TestTree
tests = testGroup "Special functions"
  [ testGroup "erf"
    [ -- implementation from numerical recipes loses presision for
      -- large arguments
      testCase "erfc table" $
        forTable "tests/tables/erfc.dat" $ \[x, exact] ->
          checkTabularPure erfcTol (show x) exact (erfc x)
    , testCase "erfc table [large]" $
        forTable "tests/tables/erfc-large.dat" $ \[x, exact] ->
          checkTabularPure erfcLargeTol (show x) exact (erfc x)
      --
    , testCase "erf table" $
        forTable "tests/tables/erf.dat" $ \[x, exact] -> do
          checkTabularPure erfTol (show x) exact (erf x)
    , testProperty "id = erfc . invErfc" invErfcIsInverse
    , testProperty "id = invErfc . erfc" invErfcIsInverse2
    , testProperty "invErf  = erf^-1"    invErfIsInverse
    ]
  --
  , testGroup "log1p & Co"
    [ testCase "expm1 table" $
        forTable "tests/tables/expm1.dat" $ \[x, exact] ->
          checkTabularPure 2 (show x) exact (expm1 x)
    , testCase "log1p table" $
        forTable "tests/tables/log1p.dat" $ \[x, exact] ->
          checkTabularPure 1 (show x) exact (log1p x)
    ]
  ----------------
  , testGroup "gamma function"
    [ testCase "logGamma table [fractional points" $
        forTable "tests/tables/loggamma.dat" $ \[x, exact] -> do
          checkTabularPure 2 (show x) exact (logGamma x)
    , testProperty "Gamma(x+1) = x*Gamma(x)" $ gammaReccurence
    , testCase     "logGamma is expected to be precise at 1e-15 level" $
        forM_ [3..10000::Int] $ \n -> do
          let exact = logFactorial (n-1)
              val   = logGamma (fromIntegral n)
          checkTabular 8 (show n) exact val
    ]
  ----------------
  , testGroup "incomplete gamma"
    [ testCase "incompleteGamma table" $
        forTable "tests/tables/igamma.dat" $ \[a,x,exact] -> do
          let err | a < 10    = 16
                  | a <= 101  = if isGHCJS then 64 else 32
                  | a == 201  = 200
                  | otherwise = 32
          checkTabularPure err (show (a,x)) exact (incompleteGamma a x)
    , testProperty "incomplete gamma - increases" $
        \(abs -> s) (abs -> x) (abs -> y) -> s > 0 ==> monotonicallyIncreases (incompleteGamma s) x y
    , testProperty "0 <= gamma <= 1"               incompleteGammaInRange
    , testProperty "gamma(1,x) = 1 - exp(-x)"      incompleteGammaAt1Check
    , testProperty "invIncompleteGamma = gamma^-1" invIGammaIsInverse
    ]
  ----------------
  , testGroup "beta function"
    [ testCase "logBeta table" $
        forTable "tests/tables/logbeta.dat" $ \[p,q,exact] ->
          let errEst
                -- For Stirling approx. errors are very good
                | b > 10          = 2
                -- Partial Stirling approx
                | a > 10 = case () of
                    _| b >= 1    -> 4
                     | otherwise -> 2 * est
                -- sum of logGamma
                | otherwise = case () of
                    _| a <= 1 && b <= 1 -> 8
                     | a >= 1 && b >= 1 -> 8
                     | otherwise        -> 2 * est
                where
                  a = max p q
                  b = min p q
                  --
                  est = ceiling
                      $ abs (logGamma a) + abs (logGamma b) + abs (logGamma (a + b))
                      / abs (logBeta a b)
          in checkTabularPure errEst (show (p,q)) exact (logBeta p q)
    , testCase "logBeta factorial" betaFactorial
    , testProperty "beta(1,p) = 1/p"   beta1p
    -- , testProperty "beta recurrence"   betaRecurrence
    ]
  ----------------
  , testGroup "incomplete beta"
    [ testCase "incompleteBeta table" $
        forM_ tableIncompleteBeta $ \(p,q,x,exact) ->
          checkTabular 64 (show (x,p,q)) (incompleteBeta p q x) exact
    , testCase "incompleteBeta table with p > 3000 and q > 3000" $
        forM_ tableIncompleteBetaP3000 $ \(x,p,q,exact) ->
          checkTabular 7000 (show (x,p,q)) (incompleteBeta p q x) exact
    --
    , testProperty "0 <= I[B] <= 1" incompleteBetaInRange
    , testProperty "ibeta symmetry" incompleteBetaSymmetry
    -- XXX FIXME DISABLED due to failures
    -- , testProperty "invIncompleteBeta  = B^-1" $ invIBetaIsInverse
    ]
  ----------------
  , testGroup "digamma"
    [ testAssertion "digamma is expected to be precise at 1e-14 [integers]"
        $ digammaTestIntegers 1e-14
      -- Relative precision is lost when digamma(x) ≈ 0
    , testCase "digamma is expected to be precise at 1e-12" $
      forTable "tests/tables/digamma.dat" $ \[x, exact] ->
        checkTabularPure 2048
          (show x) (digamma x) exact
    ]
  ----------------
  , testGroup "factorial"
    [ testCase "Factorial table" $
      forM_ [0 .. 170] $ \n -> do
        checkTabular 1
          (show n)
          (fromIntegral (factorial' n))
          (factorial (fromIntegral n :: Int))
      --
    , testCase "Log factorial from integer" $
      forM_ [2 .. 170] $ \n -> do
        checkTabular 1
          (show n)
          (log $ fromIntegral $ factorial' n)
          (logFactorial (fromIntegral n :: Int))
    , testAssertion "Factorial table is OK"
    $ U.length factorialTable == 171
    , testCase "Log factorial table" $
      forTable "tests/tables/factorial.dat" $ \[i,exact] ->
        checkTabularPure 3
          (show i) (logFactorial (round i :: Int)) exact
    ]
  ----------------
  , testGroup "combinatorics"
    [ testCase "choose table" $
      forM_ [0 .. 1000] $ \n ->
        forM_ [0 .. n]  $ \k -> do
          checkTabular 2048
            (show (n,k))
            (fromIntegral $ choose' n k)
            (choose (fromInteger n) (fromInteger k))
    --
    , testCase "logChoose == log . choose" $
      forM_ [0 .. 1000] $ \n ->
        forM_ [0 .. n]  $ \k -> do
          checkTabular 2
            (show (n,k))
            (log $ choose n k)
            (logChoose n k)
    ]
    ----------------------------------------------------------------
    -- Self tests
  , testGroup "self-test"
    [ testProperty "Self-test: 0 <= range01 <= 1" $ \x -> let f = range01 x in f <= 1 && f >= 0
    ]
  ]

----------------------------------------------------------------
-- efr tests
----------------------------------------------------------------

roundtrip_erfc_invErfc,
  roundtrip_invErfc_erfc,
  roundtrip_erf_invErf
  :: (Double,Double)
#if USE_SYSTEM_ERF && !defined(__GHCJS__)
roundtrip_erfc_invErfc = (2,2)
roundtrip_invErfc_erfc = (2,2)
roundtrip_erf_invErf   = (1,1)
#else
roundtrip_erfc_invErfc = (2,8)
roundtrip_invErfc_erfc = (8,4)
roundtrip_erf_invErf   = (128,128)
#endif

-- id ≈ erfc . invErfc
invErfcIsInverse :: Double -> Property
invErfcIsInverse ((*2) . range01 -> x)
  = (not $ isInfinite x) ==>
  ( counterexample ("x        = " ++ show x )
  $ counterexample ("y        = " ++ show y )
  $ counterexample ("x'       = " ++ show x')
  $ counterexample ("calc.err = " ++ show (delta, delta-e'))
  $ counterexample ("ulps     = " ++ show (ulpDistance x x'))
  $ ulpDistance x x' <= round delta
  )
  where
    (e,e') = roundtrip_erfc_invErfc
    delta  = e' + e * abs ( y / x  *  2 / sqrt pi * exp( -y*y ))
    y  = invErfc x
    x' = erfc y

-- id ≈ invErfc . erfc
invErfcIsInverse2 :: Double -> Property
invErfcIsInverse2 x
  = (not $ isInfinite x') ==>
    (y > m_tiny)          ==>
    (x /= 0) ==>
    counterexample ("x        = " ++ show x )
  $ counterexample ("y        = " ++ show y )
  $ counterexample ("x'       = " ++ show x')
  $ counterexample ("calc.err = " ++ show delta)
  $ counterexample ("ulps     = " ++ show (ulpDistance x x'))
  $ ulpDistance x x' <= delta
  where
    (e,e') = roundtrip_invErfc_erfc
    delta  = round
           $ e' + e * abs (y / x  /  (2 / sqrt pi * exp( -x*x )))
    y  = erfc x
    x' = invErfc y

-- id ≈ erf . invErf
invErfIsInverse :: Double -> Property
invErfIsInverse a
  = (x /= 0) ==>
    counterexample ("x        = " ++ show x )
  $ counterexample ("y        = " ++ show y )
  $ counterexample ("x'       = " ++ show x')
  $ counterexample ("calc.err = " ++ show delta)
  $ counterexample ("ulps     = " ++ show (ulpDistance x x'))
  $ ulpDistance x x' <= delta
  where
    (e,e') = roundtrip_erf_invErf
    delta  = round
           $ e + e' * abs (y / x  *  2 / sqrt pi * exp ( -y * y ))
    x  | a < 0     = - range01 a
       | otherwise =   range01 a
    y  = invErf x
    x' = erf y

----------------------------------------------------------------
-- QC tests
----------------------------------------------------------------

-- B(p,q) = (x - 1)!(y-1)! / (x + y - 1)!
betaFactorial :: IO ()
betaFactorial = do
  forM_ prod $ \(p,q,facP,facQ,facProd) -> do
    let exact = fromIntegral (facQ * facP)
              / fromIntegral facProd
    checkTabular 16 (show (p,q))
      (logBeta (fromIntegral p) (fromIntegral q))
      (log exact)
  where
    prod    = [ (p,q,facP,facQ, factorial' (p + q - 1))
              | (p,facP) <- facList
              , (q,facQ) <- facList
              , p + q < 170
              , not (p == 1 && q== 1)
              ]
    facList = [(p,factorial' (p-1)) | p <- [1 .. 170]]

-- B(1,p) = 1/p
beta1p :: Double -> Property
beta1p (abs -> p)
  = p > 2 ==>
    counterexample ("p    = " ++ show p)
  $ counterexample ("logB = " ++ show lb)
  $ counterexample ("err  = " ++ show d)
  $ d <= 24
  where
    lb = logBeta 1 p
    d  = ulpDistance lb (- log p)

{-
-- B(p+1,q) = B(p,q) · p/(p+q)
betaRecurrence :: Double -> Double -> Property
betaRecurrence (abs -> p) (abs -> q)
  = p > 0  &&  q > 0  ==>
    counterexample ("p          = " ++ show p)
  $ counterexample ("q          = " ++ show q)
  $ counterexample ("log B(p,q) = " ++ show (logBeta p q))
  $ counterexample ("log B(p+1,q) = " ++ show (logBeta (p+1) q))
  $ counterexample ("err        = " ++ show d)
  $ d <= 128
  where
    logB  = logBeta p q + log (p / (p + q))
    logB' = logBeta (p + 1) q
    d     = ulpDistance logB logB'
-}

-- Γ(x+1) = x·Γ(x)
gammaReccurence :: Double -> Property
gammaReccurence x
  = x > 0  ==>  err < errEst
    where
      g1     = logGamma x
      g2     = logGamma (x+1)
      err    = abs (g2 - g1 - log x)
      -- logGamma apparently is not as precise for small x. See #59 for details
      errEst = max 1e-14
             $ 2 * m_epsilon * sum (map abs [ g1 , g2 , log x ])

-- γ(s,x) is in [0,1] range
incompleteGammaInRange :: Double -> Double -> Property
incompleteGammaInRange (abs -> s) (abs -> x) =
  x >= 0 && s > 0  ==> let i = incompleteGamma s x in i >= 0 && i <= 1

-- γ(1,x) = 1 - exp(-x)
-- Since Γ(1) = 1 normalization doesn't make any difference
incompleteGammaAt1Check :: Double -> Bool
incompleteGammaAt1Check (abs -> x) =
  ulpDistance (incompleteGamma 1 x) (-expm1(-x)) < 16

-- invIncompleteGamma is inverse of incompleteGamma
invIGammaIsInverse :: Double -> Double -> Property
invIGammaIsInverse (abs -> a) (range01 -> p) =
  a > m_tiny && p > m_tiny && p < 1 && x > m_tiny  ==>
    ( counterexample ("a    = " ++ show a )
    $ counterexample ("p    = " ++ show p )
    $ counterexample ("x    = " ++ show x )
    $ counterexample ("p'   = " ++ show p')
    $ counterexample ("err  = " ++ show (ulpDistance p p'))
    $ counterexample ("est  = " ++ show est)
    $ ulpDistance p p' <= est
    )
  where
    x  = invIncompleteGamma a p
    f' = exp ( log x * (a-1) - x - logGamma a)
    p' = incompleteGamma    a x
    -- FIXME: Test should be rechecked when #42 is fixed
    (e,e') = (32,32)
    est    = round
           $ e' + e * abs (x / p * f')

-- I(x;p,q) is in [0,1] range
incompleteBetaInRange :: Double -> Double -> Double -> Property
incompleteBetaInRange (abs -> p) (abs -> q) (range01 -> x) =
  p > 0 && q > 0  ==> let i = incompleteBeta p q x in i >= 0 && i <= 1

-- I(0.5; p,p) = 0.5
incompleteBetaSymmetry :: Double -> Property
incompleteBetaSymmetry (abs -> p) =
    p > 0 ==>
    counterexample ("p   = " ++ show p)
  $ counterexample ("ib  = " ++ show ib)
  $ counterexample ("err = " ++ show d)
  $ counterexample ("est = " ++ show est)
  $ d <= est
  where
    est | p < 1     = 80
        | p < 10    = 200
        | otherwise = round $ 6 * p
    d  = ulpDistance ib 0.5
    ib = incompleteBeta p p 0.5

-- invIncompleteBeta is inverse of incompleteBeta
invIBetaIsInverse :: Double -> Double -> Double -> Property
invIBetaIsInverse (abs -> p) (abs -> q) (range01 -> x) =
  p > 0 && q > 0  ==> ( counterexample ("p   = " ++ show p )
                      $ counterexample ("q   = " ++ show q )
                      $ counterexample ("x   = " ++ show x )
                      $ counterexample ("x'  = " ++ show x')
                      $ counterexample ("a   = " ++ show a)
                      $ counterexample ("err = " ++ (show $ abs $ (x - x') / x))
                      $ abs (x - x') <= 1e-12
                      )
  where
    x' = incompleteBeta    p q a
    a  = invIncompleteBeta p q x

-- Table for digamma function:
--
-- Uses equality ψ(n) = H_{n-1} - γ where
--   H_{n} = Σ 1/k, k = [1 .. n]     - harmonic number
--   γ     = 0.57721566490153286060  - Euler-Mascheroni number
digammaTestIntegers :: Double -> Bool
digammaTestIntegers eps
  = all (uncurry $ eq eps) $ take 3000 digammaInt
  where
    ok approx exact = approx
    -- Harmonic numbers starting from 0
    harmN = scanl (\a n -> a + 1/n) 0 [1::Rational .. ]
    gam   = 0.57721566490153286060
    -- Digamma values
    digammaInt = zipWith (\i h -> (digamma i, realToFrac h - gam)) [1..] harmN

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

-- Lookup table for fact factorial calculation. It has fixed size
-- which is bad but it's OK for this particular case
factorial_table :: V.Vector Integer
factorial_table = V.generate 2000 (\n -> product [1..fromIntegral n])

-- Exact implementation of factorial
factorial' :: Integer -> Integer
factorial' n = factorial_table ! fromIntegral n

-- Exact albeit slow implementation of choose
choose' :: Integer -> Integer -> Integer
choose' n k = factorial' n `div` (factorial' k * factorial' (n-k))

-- Truncate double to [0,1]
range01 :: Double -> Double
range01 = abs . (snd :: (Integer, Double) -> Double) . properFraction

readTable :: FilePath -> IO [[Double]]
readTable
  = fmap (fmap (fmap read . words) . lines)
  . readFile

forTable :: FilePath -> ([Double] -> Maybe String) -> IO ()
forTable path fun = do
  rows <- readTable path
  case mapMaybe fun rows of
    []   -> return ()
    errs -> assertFailure $ intercalate "---\n" errs

checkTabular :: Int -> String -> Double -> Double -> IO ()
checkTabular prec x exact val =
  case checkTabularPure prec x exact val of
    Nothing -> return ()
    Just s  -> assertFailure s

checkTabularPure :: Int -> String -> Double -> Double -> Maybe String
checkTabularPure prec x exact val
  | within prec exact val = Nothing
  | otherwise             = Just $ unlines
      [ " x         = " ++ x
      , " expected  = " ++ show exact
      , " got       = " ++ show val
      , " ulps diff = " ++ show (ulpDistance exact val)
      , " err.est.  = " ++ show prec
      ]


{-# OPTIONS_GHC -Wno-incomplete-patterns #-}

module NVectorSpec where

import Data.Complex

import Test.QuickCheck
import Test.QuickCheck.Poly

import Category
import Comonad
import Functor
import NVector
import Unboxed

type N = 10

prop_NVector_Functor_id :: FnProp (NVector N A)
prop_NVector_Functor_id = checkFnEqual law_Functor_id

prop_NVector_Functor_comp :: Fun B C -> Fun A B -> FnProp (NVector N A)
prop_NVector_Functor_comp f g = checkFnEqual (law_Functor_comp f g)

prop_NVector_Semicomonad_comm ::
    Fun (NVector N B) C -> Fun (NVector N A) B -> FnProp (NVector N A)
prop_NVector_Semicomonad_comm f g = checkFnEqual (law_Semicomonad_comm f g)

prop_NVector_Semicomonad1_comm ::
    Fun (NVector N B) C -> Fun (NVector N A) B -> FnProp (NVector N A)
prop_NVector_Semicomonad1_comm f g = checkFnEqual (law_Semicomonad1_comm f g)

prop_NVector_Semicomonad1_comm' ::
    Fun (NVector N B) C -> Fun (NVector N A) B -> FnProp (NVector N A)
prop_NVector_Semicomonad1_comm' (Fn f) (Fn g) =
    checkFnEqual (law_Semicomonad1_comm' f g)

type UA = Int
type UB = Double
type UC = Complex Double

prop_NUVector_Functor_id :: FnProp (NUVector N UA)
prop_NUVector_Functor_id = checkFnEqual law_Functor_id

prop_NUVector_Functor_comp ::
    (UB -#> UC) -> (UA -#> UB) -> FnProp (NUVector N UA)
prop_NUVector_Functor_comp f g = checkFnEqual (law_Functor_comp f g)

prop_NUVector_Semicomonad1_comm' ::
    Fun (NUVector N UB) UC -> Fun (NUVector N UA) UB ->
    FnProp (NUVector N UA)
prop_NUVector_Semicomonad1_comm' (Fn f) (Fn g) =
    checkFnEqual (law_Semicomonad1_comm' f g)


module MarkovChain
    (
    markov
    ) where

import Numeric.LinearAlgebra
import Text.ParserCombinators.Parsec

matParse = sepBy line (char ';')
line = sepBy cell (char ',' )
cell = many ( noneOf ",\n;" )

stringToTranstionMatrix :: String -> Either ParseError [[String]]
stringToTranstionMatrix str = parse matParse "unknown" str

markov matStr = stringToTranstionMatrix matStr


module SceneRandom where

import Numeric.Vector
import Numeric.Scalar (fromScalar, scalar)
import Control.Monad.State
import System.Random
import Types
import VecUtil

randDouble :: State SceneContext Double
randDouble = do
  ss <- get
  let (v, gen) = random (ss_getGen ss) :: (Double, StdGen)
  put $ ss { ss_getGen = gen }
  return v

randRange :: Double -> Double -> State SceneContext Double
randRange l h = randDouble >>= (\v -> return $ v * (h - l) + l)

randNeg1To1 :: State SceneContext Double
randNeg1To1 = randRange (-1) 1

randVecSphere :: State SceneContext Vec3d
randVecSphere = do
  x <- randNeg1To1
  y <- randNeg1To1
  z <- randNeg1To1
  return $ normalized $ vec3 x y z

randVecHemisphere :: Vec3d -> State SceneContext Vec3d
randVecHemisphere normal = do
  let (tangent, bitangent) = tangentBitangent normal

  theta <- randRange 0 (2 * pi)
  phi <- randRange 0 (pi / 2)

  let t = fromScalar $ scalar $ sin theta * cos phi
  let b = fromScalar $ scalar $ cos theta * cos phi
  let n = fromScalar $ scalar $ sin phi
  return $ tangent * t + bitangent * b + normal * n

randVecDisc :: Vec3d -> Vec3d -> Double -> State SceneContext Vec3d
randVecDisc normal center radius = do
  let (tangent, bitangent) = tangentBitangent normal
  theta <- randRange 0 (2 * pi)
  let t = fromScalar $ scalar $ sin theta * radius
  let b = fromScalar $ scalar $ cos theta * radius
  return $ tangent * t + bitangent * b + center

-- https://math.stackexchange.com/a/182936
randVecCone :: Double -> Vec3d -> State SceneContext Vec3d
randVecCone angle vector = do
  let (tangent, bitangent) = tangentBitangent vector
  z <- randRange (cos angle) 1
  phi <- randRange 0 (2 * pi)
  let theta = acos z
  let t = fromScalar $ scalar $ cos phi
  let b = fromScalar $ scalar $ sin phi
  let v = fromScalar $ scalar $ cos theta
  let sinTheta = fromScalar $ scalar $ sin theta
  return $ sinTheta * (t * tangent + b * tangent) + v * vector

randColor :: State SceneContext Vec3d
randColor = do
  r <- randDouble
  g <- randDouble
  b <- randDouble
  return $ vec3 r g b