{-|
Module      : Mandlebrot
Description :
Copyright   : (c) Jason Chen, 2021
License     : MIT
Maintainer  : chezbgone (at) gmail dot com
Stability   : experimental
-}
{-# LANGUAGE NamedFieldPuns #-}
module Mandlebrot where

import Data.Complex
import Data.List (iterate')
import Codec.Picture

import Interpolation

data Resolution = Resolution { width :: Int, height :: Int }

data Bounds = Bounds { top    :: Double
                     , bottom :: Double
                     , left   :: Double
                     , right  :: Double
                     }

data ComputationProps = ComputationProps
    { maxIterations :: Int
    , escapeRadius :: Double
    }

data RenderProps = RenderProps
    { resolution :: Resolution
    , bounds     :: Bounds
    }

data MandlebrotProps = MandlebrotProps
    { computationProps :: ComputationProps
    , renderProps      :: RenderProps
    }

magnitude2 :: RealFloat a => Complex a -> a
magnitude2 (x :+ y) = x*x + y*y

inCardioid :: Complex Double -> Bool
inCardioid (x :+ y) =
    let q = (x - 1/4)^2 + y^2
     in q * (q + x - 1/4) <= 1/4 * y^2

inBulb1 :: Complex Double -> Bool
inBulb1 (x :+ y) = (x + 1)^2 + y^2 <= 1/16

-- https://en.wikipedia.org/wiki/Plotting_algorithms_for_the_Mandelbrot_set#Cardioid_/_bulb_checking
mandlebrotEsc' :: ComputationProps -> Complex Double -> Int
mandlebrotEsc' ComputationProps{ maxIterations, escapeRadius } c =
    let f z = z^2 + c
        orbit = iterate' f 0
     in length $ takeWhile (\p -> magnitude2 p < escapeRadius^2)
              $ take maxIterations orbit

-- optimized
mandlebrotEsc :: ComputationProps -> Complex Double -> Int
mandlebrotEsc cp c = if inCardioid c || inBulb1 c
                         then maxIterations cp
                         else mandlebrotEsc' cp c

xs = (255*) <$> [0, 0.16, 0.42, 0.6425, 0.8575]
rs = [0, 32, 237, 255, 0]
gs = [7, 107, 255, 170, 2]
bs = [100, 203, 255, 0, 0]

rf = interpolate $ piecewiseMonotonicCurve $ zip xs rs
gf = interpolate $ piecewiseMonotonicCurve $ zip xs gs
bf = interpolate $ piecewiseMonotonicCurve $ zip xs bs

intPixel :: Int -> PixelRGB8
intPixel esc =
    let esc' = fromIntegral esc
     in PixelRGB8 (round $ rf esc') (round $ gf esc') (round $ bf esc')

mandlebrotInt :: MandlebrotProps -> DynamicImage
mandlebrotInt mp =
    let cp = computationProps mp
        RenderProps{ resolution=Resolution{ width, height }
                   , bounds=Bounds{ top, bottom, left, right } } = renderProps mp
        dx = (right - left) / fromIntegral width
        dy = (top - bottom) / fromIntegral height
        complex x y = (left + fromIntegral x * dx) :+ (top - fromIntegral y * dy)
     in ImageRGB8 $ generateImage (\x y -> intPixel $ mandlebrotEsc cp $ complex x y) width height


module Test.HUnit.Plus
(
    assertMaybeDouble
  , assertOnFunction
  , assertVector
  , assertMatrix
)
where

import Control.Monad
import Test.HUnit
import Test.HUnit.Approx
import Numeric.LinearAlgebra

assertMaybeDouble :: Maybe Double -> Maybe Double -> Double -> Assertion
assertMaybeDouble Nothing Nothing _ = assertString ""
assertMaybeDouble expected Nothing _ = assertString msg
  where msg = "expected: " ++ show expected ++ "\nbut got: Nothing"
assertMaybeDouble Nothing actual _ = assertString msg
  where msg = "expected: Nothing\nbit got: " ++ show actual
assertMaybeDouble (Just expected) (Just actual) eps = assertApproxEqual "Maybe Double" eps expected actual

assertOnFunction :: (Eq b, Show b) => (a -> b) -> a -> a -> Assertion
assertOnFunction func expected actual = func expected @=? func actual

assertVector :: String -> R -> Vector R -> Vector R -> Assertion
assertVector message eps expected actual =
  let diff = norm_2 (expected-actual)
      msg = message ++ "\nexpected: " ++ show expected ++ "\nbut got" ++ show actual
  in unless (diff < eps) (assertFailure msg)

assertMatrix :: String -> R -> Matrix R -> Matrix R -> Assertion
assertMatrix message eps expected actual =
  let diff = norm_2 (expected-actual)
      msg = message ++ "\nexpected: " ++ show expected ++ "\nbut got" ++ show actual
  in unless (diff < eps) (assertFailure msg)


{-# LANGUAGE RecordWildCards #-}
-- Chi square tests for random generators
module MWC.ChiSquare (
  tests
  ) where

import Control.Applicative
import Control.Monad

import Data.Typeable
import Data.Word
import Data.List (find)
import qualified Data.Vector.Unboxed              as U
import qualified Data.Vector.Unboxed.Mutable      as M
import qualified System.Random.MWC                as MWC
import qualified System.Random.MWC.Distributions  as MWC
import qualified System.Random.MWC.CondensedTable as MWC

import Statistics.Types
import Statistics.Test.ChiSquared
import Statistics.Distribution
import Statistics.Distribution.Poisson
import Statistics.Distribution.Binomial
import Statistics.Distribution.Geometric

import Test.HUnit hiding (Test)
import Test.Framework
import Test.Framework.Providers.HUnit

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

tests :: MWC.GenIO -> Test.Framework.Test
tests g = testGroup "Chi squared tests"
    -- Word8 tests
  [ uniformRTest (0,255 :: Word8) g
  , uniformRTest (0,254 :: Word8) g
  , uniformRTest (0,129 :: Word8) g
  , uniformRTest (0,126 :: Word8) g
  , uniformRTest (0,10  :: Word8) g
    -- * Tables
  , ctableTest   [1] g
  , ctableTest   [0.5,  0.5] g
  , ctableTest   [0.25, 0.25, 0.25, 0.25] g
  , ctableTest   [0.25, 0.5,  0.25] g
  , ctableTest   [1/3 , 1/3, 1/3] g
  , ctableTest   [0.1,  0.9] g
  , ctableTest   (replicate 10 0.1) g
    -- ** Poisson
  , poissonTest 0.2  g
  , poissonTest 1.32 g
  , poissonTest 6.8  g
  , poissonTest 100  g
    -- ** Binomial
  , binomialTest 4   0.5 g
  , binomialTest 10  0.1 g
  , binomialTest 10  0.6 g
  , binomialTest 10  0.8 g
  , binomialTest 100 0.3 g
    -- ** Geometric
  , geometricTest 0.1 g
  , geometricTest 0.5 g
  , geometricTest 0.9 g
  ]

----------------------------------------------------------------
-- | RNG and corresonding distribution
data Generator = Generator {
    generator    :: MWC.GenIO -> IO Int
  , probabilites :: U.Vector Double
  }

-- | Apply chi square test for a distribution
sampleTest :: String              -- ^ Name of test
           -> Generator           -- ^ Generator to test
           -> Int                 -- ^ N of events
           -> MWC.GenIO           -- ^ PRNG state
           -> Test.Framework.Test
sampleTest nm (Generator{..}) n g = testCase nm $ do
  let size = U.length $ probabilites
  h <- histogram (generator g) size n
  let w      = U.map (* fromIntegral n) probabilites
      Just t = chi2test 0 $ U.zip h w
  case isSignificant (mkPValue 0.01) t of
    Significant    -> assertFailure ("Significant: " ++ show t)
    NotSignificant -> return ()
{-# INLINE sampleTest #-}

-- | Fill histogram using supplied generator
histogram :: IO Int             -- ^ Random generator
          -> Int                -- ^ N of outcomes
          -> Int                -- ^ N of events
          -> IO (U.Vector Int)
histogram gen size n = do
  arr <- M.replicate size 0
  replicateM_ n $ do i <- gen
                     when (i < size) $ M.write arr i . (+1) =<< M.read arr i
  U.unsafeFreeze arr
{-# INLINE histogram #-}

-- | Test uniformR
uniformRTest :: (MWC.Variate a, Typeable a, Show a, Integral a) => (a,a) -> MWC.GenIO -> Test.Framework.Test
uniformRTest (a,b)
  = sampleTest ("uniformR: " ++ show (a,b) ++ " :: " ++ show (typeOf a)) gen (10^5)
  where
    n   = fromIntegral b - fromIntegral a + 1
    gen = Generator { generator    = \g -> fromIntegral . subtract a <$> MWC.uniformR (a,b) g
                    , probabilites = U.replicate n (1 / fromIntegral n)
                    }
{-# INLINE uniformRTest #-}

-- | Test for condensed tables
ctableTest :: [Double] -> MWC.GenIO -> Test.Framework.Test
ctableTest ps
  = sampleTest ("condensedTable: " ++ show ps) gen (10^4)
  where
    gen = Generator
          { generator    = MWC.genFromTable $ MWC.tableFromProbabilities $ U.fromList $ zip [0..] ps
          , probabilites = U.fromList ps
          }

-- | Test for condensed table for poissson distribution
poissonTest :: Double -> MWC.GenIO -> Test.Framework.Test
poissonTest lam
  = sampleTest ("poissonTest: " ++ show lam) gen (10^4)
  where
    pois      = poisson lam
    Just nMax = find (\n -> probability pois n < 2**(-33)) [floor lam ..]
    gen = Generator
          { generator    = MWC.genFromTable (MWC.tablePoisson lam)
          , probabilites = U.generate nMax (probability pois)
          }

-- | Test for condensed table for binomial distribution
binomialTest :: Int -> Double -> MWC.GenIO -> Test.Framework.Test
binomialTest n p
  = sampleTest ("binomialTest: " ++ show p ++ " " ++ show n) gen (10^4)
  where
    binom = binomial n p
    gen   = Generator
            { generator    = MWC.genFromTable (MWC.tableBinomial n p)
            , probabilites = U.generate (n+1) (probability binom)
            }

-- | Test for geometric distribution
geometricTest :: Double -> MWC.GenIO -> Test.Framework.Test
geometricTest gd
  = sampleTest ("geometricTest: " ++ show gd) gen (10^4)
  where
    n   = 1000
    gen = Generator
          { generator    = MWC.geometric1 gd
          , probabilites = U.generate (n+1) (probability $ geometric gd)
          }


--
-- LayerType : layer types
--

module CNN.LayerType (
  ActFunc
, Kernel
, Bias
, FilterC
, FilterF
, Delta
, Layer ( NopLayer, ActLayer, MaxPoolLayer, ConvLayer, FullConnLayer
        , FlattenLayer)
, LayerB
) where

import Numeric.LinearAlgebra.Data

import CNN.Image

-- for Activation Layer
type ActFunc = Matrix R -> Matrix R

-- filter for Convolution Layer
type Kernel = Matrix R
type Bias   = Double
type FilterC = (Kernel, Bias)

-- filter for Fully Connected Layer
type FilterF = Matrix R

-- for backward propergation
--type Delta = [Double]
type Delta = Image

-- LayerB is for converting from parsed data to Layer
type LayerB = (String, [Int])

data Layer =
    NopLayer
  | ActLayer ActFunc
  | MaxPoolLayer Int
  | ConvLayer Int [FilterC]
  | FullConnLayer FilterF
  | FlattenLayer Int Int

instance Eq Layer where
  NopLayer             == NopLayer             = True
  (ActLayer f1)        == (ActLayer f2)        = True
  (MaxPoolLayer s1)    == (MaxPoolLayer s2)    = s1 == s2
  (ConvLayer s1 fs1)   == (ConvLayer s2 fs2)   = s1 == s2 && fs1 == fs2
  (FullConnLayer fs1)  == (FullConnLayer fs2)  = fs1 == fs2
  (FlattenLayer x1 y1) == (FlattenLayer x2 y2) = x1 == x2 && y1 == y2
  _                    == _                    = False


module TextStatistics (
    TextStatistics,
    wordCount,
    lineCount,
    averageWordLength,
    mostCommonLetter
) where

import Prelude hiding (filter)
import qualified Statistics as S
import Control.Applicative ((<$>), (<*>))
import qualified Data.HashMap.Strict as HM
import Data.Char

-- Statistics on Char stream

type TextStatistics a = S.Statistics Char a

-- predicates

isWordChar :: Char -> Bool
isWordChar = not <$> isSpace

isNewLine :: Char -> Bool
isNewLine c = c == '\n'

-- word count

data WordCountState = WordCountState {
    charCount :: Int,
    isInWord :: Bool
}

wordCount :: TextStatistics Int
wordCount = charCount <$> S.fold f (WordCountState 0 False)
    where
        f w c = WordCountState (if not (isInWord w) && isWordChar c then (charCount w) + 1 else (charCount w)) (isWordChar c)

-- line count. Assuming empty file has 1 line of 0 length

lineCount :: TextStatistics Int
lineCount = S.count isNewLine 1

-- average word length

wordCharCount :: TextStatistics Int
wordCharCount = S.count isWordChar 0

averageWordLength :: TextStatistics (Maybe Float)
averageWordLength = maybeDivide <$> wordCharCount <*> wordCount
    where
        maybeDivide a b = if b /= 0 then Just (fromIntegral a / fromIntegral b) else Nothing

-- most common character (if any) with number of usages

mostCommonLetter :: TextStatistics (Maybe Char)
mostCommonLetter = fmap fst <$> ((HM.foldlWithKey' maxChar Nothing) <$> lettersMap)
    where
        lettersMap :: TextStatistics (HM.HashMap Char Int)
        lettersMap = S.filter isWordChar $ S.fold (\hm c -> if HM.member c hm then HM.insertWith (+) c 1 hm else HM.insert c 1 hm) HM.empty
        maxChar Nothing c n = Just (c, n)
        maxChar (Just (c0, n0)) c n = if n > n0 then Just (c, n) else Just (c0, n0)


{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE BangPatterns #-}

module Lib
    ( trilinear
    , resample
    , affine
    , array_correlation
    , fisher_transform
    , naive_12p_registration
    ) where

import Data.Array.Repa
import qualified Data.Array.Repa as Repa
import Data.Array.Repa.Algorithms.Matrix
import Statistics.Sample
import Data.Array.Repa.Repr.Unboxed

type DIM3D = (Z :. Double :. Double :. Double)

trilinear :: (Unbox a, Floating a, Fractional a) => Array U DIM3 a -> DIM3D -> a
trilinear !arr !(Z :. x :. y :. z) =
  let
    x0 = floor x
    x1 = ceiling x :: Int
    y0 = floor y
    y1 = ceiling y :: Int
    z0 = floor z
    z1 = ceiling z :: Int
    in let
      v000 = arr ! (Z :. x0 :. y0 :. z0)
      v100 = arr ! (Z :. x1 :. y0 :. z0)
      v010 = arr ! (Z :. x0 :. y1 :. z0)
      v110 = arr ! (Z :. x1 :. y1 :. z0)
      v001 = arr ! (Z :. x0 :. y0 :. z1)
      v101 = arr ! (Z :. x1 :. y0 :. y1)
      v011 = arr ! (Z :. x0 :. y1 :. y1)
      v111 = arr ! (Z :. x1 :. y1 :. y1)
      xd = if x0 /= x1 then
        ((fromIntegral x0) - (realToFrac x)) / (fromIntegral (x0 - x1)) else
        0
      yd = if y0 /= y1 then
        ((fromIntegral y0) - (realToFrac y)) / (fromIntegral (y0 - y1)) else
        0
      zd = if z0 /= z1 then
        ((fromIntegral z0) - realToFrac z) / (fromIntegral (z0 - z1)) else
        0
      in let
        c00 = v000 * (1 - xd) + v100 * xd
        c01 = v001 * (1 - xd) + v101 * xd
        c10 = v010 * (1 - xd) + v110 * xd
        c11 = v011 * (1 - xd) + v111 * xd
        in let
          c0 = c00 * (1 - yd) + c10 * yd
          c1 = c01 * (1 - yd) + c11 * yd
          in  c0 * (1 - zd) + c1 * zd

resample ::
  (Floating a, Floating b, Unbox a, Unbox b, Fractional a, Fractional b) =>
  Array U DIM3 a ->
  (DIM3 -> DIM3D) -> --fun from coordinate system to coordinate system
  (Array U DIM3 a -> DIM3D -> b) -> --interpolator
  Array D DIM3 b

resample !arr !xfm !interp =
  Repa.traverse arr id xfm_and_interp
  where
    xfm_and_interp = \_ d -> ((interp arr (xfm d)))

affine ::
  (Unbox a, Floating a, Fractional a, Real a) =>
  DIM3 -> Array U DIM2 a -> DIM3D

affine !(Z :. x :. y  :. z) !xfm =
  (Z :. x' :. y' :. z')
  where
    [x',y',z',_] = (toList (mmultS (computeS $ Repa.map realToFrac xfm) pnt_array))
      where
        pnt_array = (fromListUnboxed
                     (Z :. (3 :: Int) :. (1 :: Int))
                     (fmap fromIntegral [x,y,z])) -- :: Array U DIM2 Double

array_correlation ::
  (Unbox a, Real a, Fractional a) => Array U DIM3 a -> Array U DIM3 a -> a
array_correlation !arr1 !arr2 =
  realToFrac
   (correlation
    (toUnboxed
     (computeS
      (Repa.zipWith coerce_and_tuple
         arr1 arr2))))
  where
     coerce_and_tuple x y = (realToFrac x, realToFrac y)

fisher_transform = atanh

-- Test code
naive_12p_registration ::
  (Unbox a, Floating a, Fractional a, Real a) =>
  Array U DIM3 a -> [a] -> a

naive_12p_registration !arr !params =
  array_correlation arr (computeS $ resample arr (flip affine xfm) trilinear)
  where
    xfm = fromListUnboxed (Z :. 4 :. 4) params


{-# LANGUAGE ScopedTypeVariables, TypeFamilies #-}

module Fuml.Supervised.KNN where

import Numeric.LinearAlgebra
import Data.List (sortBy)
import Data.Ord (comparing)
import Fuml.Core
import Fuml.Base.KNN
import Control.Monad.Identity

data KNN a = KNN {
   knnPoints :: [(Vector Double, a)]
  }

generalKNN :: Int -> (Vector Double -> Vector Double -> Double) -> ([(Double, a)] -> b) -> Supervisor Identity a (KNN a) b
generalKNN k dist vote  = Supervisor $ \_ theData ->
  return $ Predict (KNN theData) $ \v ->
   let ptDist (v1, y) = (dist v v1,y)
       nns = take k $ sortBy (comparing fst) $ map ptDist $ theData
   in vote nns

binaryKNN :: Int -> Supervisor Identity Bool (KNN Bool) Double
binaryKNN k = generalKNN k euclideanDistance weightedBoolVote

catKNN :: Eq a => Int -> Supervisor Identity a (KNN a) a
catKNN k = generalKNN k euclideanDistance weightedMajorityVote

regressKNN :: Int -> Supervisor Identity Double (KNN Double) Double
regressKNN k = generalKNN k euclideanDistance weightedAvgVote


module AugLag where

import Numeric.LinearAlgebra
import Numeric.NLOPT

{-
See https://www.mat.univie.ac.at/~neum/glopt/coconut/Benchmark/Library1_new_v1.html

Alkyl problem in AMPL format:
var x2 := 1.745, >= 0, <= 2;
var x3 := 1.2, >= 0, <= 1.6;
var x4 := 1.1, >= 0, <= 1.2;
var x5 := 3.048, >= 0, <= 5;
var x6 := 1.974, >= 0, <= 2;
var x7 := 0.893, >= 0.85, <= 0.93;
var x8 := 0.928, >= 0.9, <= 0.95;
var x9 := 8, >= 3, <= 12;
var x10 := 3.6, >= 1.2, <= 4;
var x11 := 1.45, >= 1.45, <= 1.62;
var x12 := 1, >= 0.99, <= 1.01010101010101;
var x13 := 1, >= 0.99, <= 1.01010101010101;
var x14 := 1, >= 0.9, <= 1.11111111111111;
var x15 := 1, >= 0.99, <= 1.01010101010101;

minimize obj:  - 6.3*x5*x8 + 5.04*x2 + 0.35*x3 + x4 + 3.36*x6;

subject to
e2:  - 0.819672131147541*x2 + x5 - 0.819672131147541*x6 = 0;
e3: 0.98*x4 - x7*(0.01*x5*x10 + x4) = 0;
e4:  - x2*x9 + 10*x3 + x6 = 0;
e5: x5*x12 - x2*(1.12 + 0.13167*x9 - 0.0067*x9*x9) = 0;
e6: x8*x13 - 0.01*(1.098*x9 - 0.038*x9*x9) - 0.325*x7 - 0.57425 = 0;
e7: x10*x14 + 22.2*x11 - 35.82 = 0;
e8: x11*x15 - 3*x8 + 1.33 = 0;
-}
alkyl :: [Double] -> Either Result Solution
alkyl x0 =
  let lb = fromList [0, 0, 0, 0, 0, 0.85, 0.9, 3, 1.2, 1.45, 0.99, 0.99, 0.9, 0.99]
      ub = fromList [2, 1.6, 1.2, 5, 2, 0.93, 0.95, 12, 4, 1.62, 1.01010101010101, 1.01010101010101, 1.11111111111111, 1.01010101010101]
      objf x = - 6.3 * x ! 3 * x ! 6 + 5.04 * x ! 0 + 0.35 * x ! 1 + x ! 2 + 3.36 * x ! 4
      stop = ObjectiveRelativeTolerance 1e-12 :| []
      eqCons =
        [ EqualityConstraint
            { eqConstraintFunctions = Scalar (\x -> -0.819672131147541 * x ! 0 + x ! 3 - 0.819672131147541 * x ! 4),
              eqConstraintTolerance = 1e-10
            },
          EqualityConstraint
            { eqConstraintFunctions = Scalar (\x -> 0.98 * x ! 2 - x ! 5 * (0.01 * x ! 3 * x ! 8 + x ! 2)),
              eqConstraintTolerance = 1e-10
            },
          EqualityConstraint
            { eqConstraintFunctions = Scalar (\x -> - x ! 0 * x ! 7 + 10 * x ! 1 + x ! 4),
              eqConstraintTolerance = 1e-10
            },
          EqualityConstraint
            { eqConstraintFunctions = Scalar (\x -> x ! 3 * x ! 10 - x ! 0 * (1.12 + 0.13167 * x ! 7 - 0.0067 * x ! 7 * x ! 7)),
              eqConstraintTolerance = 1e-10
            },
          EqualityConstraint
            { eqConstraintFunctions = Scalar (\x -> x ! 6 * x ! 11 - 0.01 * (1.098 * x ! 7 - 0.038 * x ! 7 * x ! 7) - 0.325 * x ! 5 - 0.57425),
              eqConstraintTolerance = 1e-10
            },
          EqualityConstraint
            { eqConstraintFunctions = Scalar (\x -> x ! 8 * x ! 12 + 22.2 * x ! 9 - 35.82),
              eqConstraintTolerance = 1e-10
            },
          EqualityConstraint
            { eqConstraintFunctions = Scalar (\x -> x ! 9 * x ! 13 - 3 * x ! 6 + 1.33),
              eqConstraintTolerance = 1e-10
            }
        ]
      problem = AugLagProblem eqCons [] (AUGLAG_EQ_LOCAL subproblem)
      subproblem = LocalProblem (fromIntegral $ size lb) stop algorithm
      algorithm = SBPLX objf [LowerBounds lb, UpperBounds ub] Nothing
   in minimizeAugLag problem (fromList x0)


module Internal
  where
import Data.Complex

type Cplx = Complex Double

toCplx :: Double -> Cplx
toCplx x = x :+ 0.0

getPhiK :: Cplx -> (Cplx, Int)
getPhiK phi
  | realPart phi > pi/2 =
    until (\(x,_) -> realPart x <= pi/2) (\(x,k) -> (x-pi,k+1)) (phi,0)
  | realPart phi < -pi/2 =
    until (\(x,_) -> realPart x >= -pi/2) (\(x,k) -> (x+pi,k-1)) (phi,0)
  | otherwise = (phi,0)

atanC :: Cplx -> Cplx
atanC z = i * (log(1-i*z) - log(1+i*z)) / 2
  where
    i = 0.0 :+ 1.0


-----------------------------------------------------------------------------
-- |
-- Module      :  Data.Packed.Vector
-- Copyright   :  (c) Alberto Ruiz 2009
-- License     :  GPL
--
-- Maintainer  :  Alberto Ruiz <aruiz@um.es>
-- Stability   :  provisional
--
-- Random vectors and matrices.
--
-----------------------------------------------------------------------------

module Data.Packed.Random (
    Seed,
    RandDist(..),
    randomVector,
    gaussianSample,
    uniformSample
) where

import Numeric.GSL.Vector
import Data.Packed
import Numeric.ContainerBoot
import Numeric.LinearAlgebra.Algorithms

type Seed = Int

-- | Obtains a matrix whose rows are pseudorandom samples from a multivariate
-- Gaussian distribution.
gaussianSample :: Seed
               -> Int -- ^ number of rows
               -> Vector Double -- ^ mean vector
               -> Matrix Double -- ^ covariance matrix
               -> Matrix Double -- ^ result
gaussianSample seed n med cov = m where
    c = dim med
    meds = konst 1 n `outer` med
    rs = reshape c $ randomVector seed Gaussian (c * n)
    m = rs `mXm` cholSH cov `add` meds

-- | Obtains a matrix whose rows are pseudorandom samples from a multivariate
-- uniform distribution.
uniformSample :: Seed
               -> Int -- ^ number of rows
               -> [(Double,Double)] -- ^ ranges for each column
               -> Matrix Double -- ^ result
uniformSample seed n rgs = m where
    (as,bs) = unzip rgs
    a = fromList as
    cs = zipWith subtract as bs
    d = dim a
    dat = toRows $ reshape n $ randomVector seed Uniform (n*d)
    am = konst 1 n `outer` a
    m = fromColumns (zipWith scale cs dat) `add` am


import HLearn.Algebra
import HLearn.Models.Distributions

import Control.Monad
import Control.Monad.ST
import Data.List
import Data.STRef
import Data.Number.LogFloat
import qualified Control.ConstraintKinds as CK

import Statistics.Distribution
import Statistics.Distribution.Normal

import qualified Data.Vector.Unboxed as VU
import qualified Data.Vector.Generic as G
import qualified Data.Foldable as F

import qualified Data.Vector as V
import qualified Data.Foldable as F
size=10^7
main = do
    print $ (foldl' mappend mempty $ map train1dp [(0::Double)..size] :: Gaussian Double)
    putStrLn "done."
--     print $ (train (concat $ replicate (10^6) ["a","b"]) :: Categorical String Double)
--     print $ (train [1..fromIntegral size::Double] :: Gaussian Double)
--     print $ (train (VU.enumFromN (0::Double) size) :: Gaussian Double)
--     print $ normalFromSample $ VU.enumFromN 0 size


module Common (
    genWeights,
    initWeights,
    inputWithBias,
    weightWithoutBias,
    pickupSets
  ) where

import Numeric.LinearAlgebra
import System.Random.Shuffle

genWeights :: [Int] -> IO [Matrix R]
genWeights ns = mapM genWeight ts
  where ts = zip ns (tail ns)

genWeight :: (Int, Int) -> IO (Matrix R)
genWeight (i, o) = randn o (i + 1) -- +1: bias

initWeights :: [Int] -> [Matrix R]
initWeights ns = fmap initWeight ts
  where ts = zip ns (tail ns)

initWeight :: (Int, Int) -> Matrix R
initWeight (i, o) = matrix (i + 1) . replicate (o * (i + 1)) $ 0 -- +1: bias

inputWithBias :: Matrix R -> Matrix R
inputWithBias v = v === konst 1 (1, cols v) -- [1,1,..1]: bias

weightWithoutBias :: Matrix R -> Matrix R
weightWithoutBias w = w ?? (All, DropLast 1)

pickupSets :: Int -> (Matrix R, Matrix R) -> IO (Matrix R, Matrix R)
pickupSets n (x, y) = do
  s <- shuffleM [0..cols y - 1]
  let ps = take n s in
    return (x ¿ ps, y ¿ ps)


{-# LANGUAGE BangPatterns #-}
module Lib where

import Graphics.UI.GLUT
import Control.Monad
import Data.Int
import Data.Complex
import System.Exit ( exitWith, ExitCode(ExitSuccess) )
import qualified Data.Map as Map

start = 5
iterations = 7
maxIterations = 30

x // y = fromIntegral x / fromIntegral y

-- Divides [a] into [[a], [a], ...] with each sublist of length n,
-- except the last sublist which has length <= n.
chunkify n [] = []
chunkify n xs = let (xs', rest) = splitAt n xs
                in xs' : chunkify n rest

-- Converts a coordinate in screen space to a vertex.
pix2vert (Size w h) (x, y) =
  Vertex2 ((3 // w * fromIntegral x) - 2.0) ((2 // h * fromIntegral y) - 1.0)

-- List of all of the vertices that represent screen pixels.
vertices :: IO [Vertex2 GLfloat]
vertices = get windowSize >>= (\x -> return $ vertices' x)

vertices' (Size w h) = [pix2vert (Size w h) (x, y) | x <- [0..w-1], y <- [0..h-1]]

-- Gets the color for a number of iterations.
color3 r g b = Color3 r g b

getcolor :: Int -> Color3 Float
getcolor numberOfIterations
  | inMandelbrotSet = color3 0 0 0
  | otherwise = color3 (escapeNumber*0.5) (escapeNumber * 0.5) escapeNumber
  where escapeNumber = numberOfIterations // iterations
        inMandelbrotSet = numberOfIterations == iterations

-- Returns the number of iterations <= the maximum iterations of the
-- Mandelbrot set at the given vertex.
mandel numberOfIterations (Vertex2 r i) =
  length . takeWhile (\z -> magnitude z <= 2) .
  take numberOfIterations $ iterate (\z -> z^2 + (r :+ i)) 0

mandel' numberOfIterations (Vertex2 r i) =
  length . takeWhile (\z -> magnitude z <= 2) .
  take numberOfIterations $ iterate (\z -> z^2 + (r :+ i)) 0

mandelMap :: (RealFloat a1, Fractional a2) => Vertex2 a1 -> Map.Map Int (Color3 a2)
mandelMap (Vertex2 r i) = Map.fromList iterList
  where mandelNumbers = iterate (\z -> z^2 + (r :+ i)) 0

        escapeVelocity numIter = length . takeWhile (\z -> magnitude z <= 2) .
                                 take numIter $ mandelNumbers

        --iterList = makeMap (\iter -> getColor (escapeVelocity iter) iter) [1..maxIterations]
        iterList = makeMap (\iter -> getColor (escapeVelocity iter) iter) [start..iterations]

        getColor escapedAt totalNumIters
          | inMandelbrotSet = color3 0 0 0
          | otherwise = color3 (escapeNumber*0.5) (escapeNumber * 0.5) escapeNumber
          where escapeNumber = escapedAt // totalNumIters
                inMandelbrotSet = escapedAt == totalNumIters

colorPoint :: RealFloat a => Int -> Vertex2 a -> Color3 Float
colorPoint numIterations = getcolor . (mandel numIterations)

-- plots one point.
-- drawVert v = do
--   color . getcolor $ mandel v
--   vertex v

drawVert' (point, color_) = do
  color color_
  vertex point

-- draws all the vertices in slices (to update the display while drawing).
-- display' chunks = do
  -- mapM_ render chunks
  -- displayCallback $= display
  -- where render points = do
          -- renderPrimitive Points (mapM_ drawVert points)
          -- flush

getColor :: (Map.Map (Vertex2 GLfloat) (Map.Map Int (Color3 Double))) -> Int -> Vertex2 GLfloat -> Color3 Double
getColor pointMap numIter vertex =
  case Map.lookup vertex pointMap of
    Just colorMap -> case Map.lookup numIter colorMap of
                       Just r -> r
                       Nothing -> Color3 0 0 0
    Nothing -> Color3 0 0 0

drawAll :: Int ->  Map.Map (Vertex2 GLfloat) (Map.Map Int (Color3 Double)) -> [Vertex2 GLfloat] -> IO()
drawAll count pointMap points = do
  render
  displayCallback $= (display (count + 1) pointMap)
  where
        render = do
          renderPrimitive Points (mapM_ drawVert' toDraw)
          flush

        toDraw :: [(Vertex2 GLfloat, Color3 Double)]
        toDraw = makeMap (getColor pointMap count) points

-- draws the whole fractal
display :: Int -> Map.Map (Vertex2 GLfloat) (Map.Map Int (Color3 Double)) -> IO ()
display 200 _ = do exitWith ExitSuccess
display num pointMap = do
  clear [ ColorBuffer ]
  -- displayCallback $= (vertices >>= display' . chunkify 256)
  displayCallback $= (vertices >>= (drawAll num pointMap))
  get currentWindow >>= postRedisplay

keyboard :: KeyboardMouseCallback
keyboard (Char keyPressed) Down _ _ = case keyPressed of
   ' '   -> postRedisplay Nothing
   '\27' -> exitWith ExitSuccess
   _     -> return ()
keyboard _ _ _ _ = return ()

mandelMain = do
  getArgsAndInitialize
  initialDisplayMode $= [ SingleBuffered, RGBMode]
  initialWindowSize $= Size 1200 1024
  initialWindowPosition $= Position 100 100
  createWindow "Mandelbrot"
  --fullScreen
  -- renderString Roman "Test"
  clearColor $= Color4 0 0 0 0
  matrixMode $= Projection
  loadIdentity
  ortho (-2) 1 (-1) 1 (-1) 1
  keyboardMouseCallback $= Just keyboard
  points <- vertices
  putStrLn "PreLoading.."
  let !preload = Map.fromList $ makeMap mandelMap points
  putStrLn "Loaded"
  displayCallback $= (display start preload)
  mainLoop

makeMap :: (a -> b) -> [a] -> [(a, b)]
makeMap fn lst = zip lst (map fn lst)


module Main where

import Codec.Picture
import Codec.Picture.Gif
import Codec.Picture.Png
import Codec.Picture.Types
import Control.Parallel
import Control.Parallel.Strategies
import Data.Complex
import Data.List.Split
import Data.Vector.Storable
import System.Environment

main = do
  putStrLn "Resolution: "
  size <- read <$> getLine
  putStrLn "Iterations: "
  iters <- read <$> getLine
  writePng "out.png" $ genFractal size iters 2

genFractal :: Int -> Int -> Complex Double -> Image Pixel8
genFractal size iters comp = Image size size $ fromList pixels
  where
    pixels =
      [genPixel comp (i `mod` size) (i `div` size) | i <- [0 .. size * size `div` 2]] `using`
      parList rdeepseq :: [Pixel8]
    genPixel comp x y =
      fromIntegral $ div 0xff iters * (iters - boundedUntil comp) :: Pixel8
      where
        boundedUntil =
          isBounded iters $
          fromIntegral (x - div size 2) / (fromIntegral size / 3) :+
          fromIntegral (y - div size 2) / (fromIntegral size / 3)

isBounded iters c comp = isBounded' 1 c
  where
    isBounded' i z
      | i == iters = i
      | magnitude z > 2 = i
      | otherwise = isBounded' (i + 1) $ (z ** comp) + c


module CommonFunctions where

import Data.Complex

-- f(z) = z^2 + c
f1 :: Complex Double -> Complex Double -> Complex Double
f1 c z = z**2 + c

-- below is valid haskell
-- f1 :: Double -> Double -> Double
-- f1 a b = a**2 + b


{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE BangPatterns #-}
-- |
-- Module      : Statistics.Correlation.Pearson
--
module Statistics.Correlation
    ( -- * Pearson correlation
      pearson
    , pearsonMatByRow
      -- * Spearman correlation
    , spearman
    , spearmanMatByRow
    ) where

import qualified Data.Vector.Generic as G
import qualified Data.Vector.Unboxed as U
import Statistics.Matrix
import Statistics.Sample
import Statistics.Test.Internal (rankUnsorted)

----------------------------------------------------------------
-- Pearson
----------------------------------------------------------------

-- | Pearson correlation for sample of pairs. Exactly same as
-- 'Statistics.Sample.correlation'
pearson :: (G.Vector v (Double, Double), G.Vector v Double)
        => v (Double, Double) -> Double
pearson = correlation
{-# INLINE pearson #-}

-- | Compute pairwise pearson correlation between rows of a matrix
pearsonMatByRow :: Matrix -> Matrix
pearsonMatByRow m
  = generateSym (rows m)
      (\i j -> pearson $ row m i `U.zip` row m j)
{-# INLINE pearsonMatByRow #-}

----------------------------------------------------------------
-- Spearman
----------------------------------------------------------------

-- | compute spearman correlation between two samples
spearman :: ( Ord a
            , Ord b
            , G.Vector v a
            , G.Vector v b
            , G.Vector v (a, b)
            , G.Vector v Int
            , G.Vector v Double
            , G.Vector v (Double, Double)
            , G.Vector v (Int, a)
            , G.Vector v (Int, b)
            )
         => v (a, b)
         -> Double
spearman xy
  = pearson
  $ G.zip (rankUnsorted x) (rankUnsorted y)
  where
    (x, y) = G.unzip xy
{-# INLINE spearman #-}

-- | compute pairwise spearman correlation between rows of a matrix
spearmanMatByRow :: Matrix -> Matrix
spearmanMatByRow
  = pearsonMatByRow . fromRows . fmap rankUnsorted . toRows
{-# INLINE spearmanMatByRow #-}


module Util.Numeric where

import Control.Arrow (second)
import Numeric.Log
import Numeric.SpecFunctions (logGamma)

average :: Fractional a => [a] -> a
average = go 0 0 where
  go n x [] = x / fromIntegral n
  go n x (z : zs) = go (n + 1) (x + z) zs

shiftByMax :: (Precise d, RealFloat d, Ord d) => [(a, Log d)] -> [(a, Log d)]
shiftByMax xs = map (second (/ x')) xs
  where x' = maximum $ map snd xs

exponentiateWeights :: (Precise d, RealFloat d, Ord d) => [(a, Log d)] -> [(a, d)]
exponentiateWeights = map (second (exp . ln)) . shiftByMax

weightedAverage :: Fractional a => [(a, a)] -> a
weightedAverage = go 0 0 where
  go w x [] = x / w
  go w x ((y, ll) : ys) = go (w + ll) (ll * y + x) ys

weightedAverageGeneric :: Fractional a => (a -> b -> b) -> b -> (b -> b -> b) -> [(b, a)] -> b
weightedAverageGeneric scale zero plus = go 0 zero where
  go w x [] = scale (recip w) x
  go w x ((y, ll) : ys) = go (w + ll) (scale ll y `plus` x) ys

logFact :: Int -> Double
logFact n = logGamma (fromIntegral (n + 1))


{-# OPTIONS_GHC -XTemplateHaskell #-}

module Euler.Problem011Test (suite) where

import Numeric.LinearAlgebra.Data (Matrix, loadMatrix)
import Test.Tasty (TestTree)
import Test.Tasty.HUnit
import Test.Tasty.TH (testGroupGenerator)

import Euler.Problem011

suite :: TestTree
suite = $(testGroupGenerator)

matrix :: IO (Matrix Double)
matrix = loadMatrix "input/011.txt"

case_count_of_vectors_in_input :: Assertion
case_count_of_vectors_in_input =
  matrix >>= \m -> 10 @=? (length . vectors . head $ submatrices 4 m)

case_count_of_4x4_submatrices_in_input :: Assertion
case_count_of_4x4_submatrices_in_input =
  matrix >>= \m -> 17 * 17 @=? (length $ submatrices 4 m)


{-# LANGUAGE OverloadedLists     #-}
{-# LANGUAGE ScopedTypeVariables #-}

module Isumi.Math.FFT.BruunSpec
  ( spec
  ) where

import           Test.Hspec
import           Test.QuickCheck

import           Data.Complex
import qualified Data.Vector.Unboxed     as UV
import           Isumi.Hspec             (shouldBeAppx)
import           Isumi.Math.AppxEq       ((~=))
import           Isumi.Math.FFT.Bruun
import           Isumi.Math.FFT.SpecUtil

spec :: Spec
spec =
  describe "fftBruun/fftBruunC" $ do
    it "works for a simple sample" $
      fftBruun [1..4] `shouldBeAppx` Just
        [10 :+ 0, (-2) :+ 2, (-2) :+ 0, (-2) :+ (-2)]
    it "works for a simple complex sample" $
      fftBruunC [1 :+ 4, 2 :+ 3, 3 :+ 2, 4 :+ 1] `shouldBeAppx` Just
        [10 :+ 10, 0 :+ 4, (-2) :+ 2, (-4) :+ 0]
    it "works for input with trailing zeros" $
      fftBruun [1, 1, 0, 0] `shouldBeAppx` Just
        [2 :+ 0, 1 :+ (-1), 0, 1 :+ 1]
    it "gives same answer for all inputs as fftDirect" $ property $
      \(VectorOf2ExpSize xs) ->
        Just (fftDirect (realVToComplexV xs)) ~= fftBruun xs
    it "gives same answer for all complex inputs as fftDirect" $ property $
      \(VectorOf2ExpSize rs) (VectorOf2ExpSize cs) ->
        let xs = UV.zipWith (:+) rs cs
         in Just (fftDirect xs) ~= fftBruunC xs


{-# LANGUAGE UndecidableInstances #-} --For the BitPack instance
module Clash.DSP.Complex where

import Clash.Prelude

import GHC.Generics
import Test.QuickCheck

import qualified Data.Complex as C

{-| I defined my own complex type so that I can write a Num instance without the RealFloat constraint. TODO: think about whether this is really a good idea. -}
data Complex a = a :+ a deriving (Show, Lift, Generic, ShowX, NFDataX, Functor, Foldable, Traversable, Eq, Ord)

deriving instance (BitPack a, KnownNat (BitSize a)) => BitPack (Complex a)

instance Num a => Num (Complex a) where
    (a :+ b) + (c :+ d) = (a + c) :+ (b + d)
    (a :+ b) - (c :+ d) = (a - c) :+ (b - d)
    (a :+ b) * (c :+ d) = (a * c - b * d) :+ (a * d + b * c)
    fromInteger x       = (fromInteger x :+ 0)

instance Applicative Complex where
    pure a = a :+ a
    f :+ g <*> a :+ b = f a :+ g b

instance Arbitrary a => Arbitrary (Complex a)
    where
    arbitrary = liftA2 (:+) arbitrary arbitrary

instance Default a => Default (Complex a)
    where
    def = def :+ def

realPart :: Complex a -> a
realPart (x :+ _) = x

imagPart :: Complex a -> a
imagPart (_ :+ y) = y

fromComplex (a C.:+ b) = a :+ b
toComplex   (a :+ b)   = a C.:+ b

conjugate :: Num a => Complex a -> Complex a
conjugate (x :+ y) = x :+ (-y)

cMul
    :: ExtendingNum a b
    => ExtendingNum (MResult a b) (MResult a b)
    => Complex a
    -> Complex b
    -> Complex (AResult (MResult a b) (MResult a b))
cMul (xr :+ xi) (yr :+ yi)
    =  ((xr `mul` yr) `sub` (xi `mul` yi))
    :+ ((xr `mul` yi) `add` (xi `mul` yr))

cMul3
    :: ExtendingNum b b
    => ExtendingNum a a
    => ExtendingNum a (AResult b b)
    => ExtendingNum b (AResult a a)
    => ExtendingNum (MResult b (AResult a a)) (MResult a (AResult b b))
    => Complex a
    -> Complex b
    -> Complex (AResult (MResult b (AResult a a)) (MResult a (AResult b b)))
cMul3 (xr :+ xi) (yr :+ yi)
    =  (k1 `sub` k3)
    :+ (k1 `add` k2)
    where
    k1 = yr `mul` (xr `add` xi)
    k2 = xr `mul` (yi `sub` yr)
    k3 = xi `mul` (yr `add` yi)

cMulPipe
    :: HiddenClockResetEnable dom
    => ExtendingNum a b
    => ExtendingNum (MResult a b) (MResult a b)
    => NFDataX (MResult a b)
    => NFDataX (AResult (MResult a b) (MResult a b))
    => Signal dom Bool
    -> Signal dom (Complex a)
    -> Signal dom (Complex b)
    -> Signal dom (Complex (AResult (MResult a b) (MResult a b)))
cMulPipe en x y
    = liftA2 (:+) rr ri
    where
    --Products
    xryr = regEn (errorX "initial xryr") en $ liftA2 mul (realPart <$> x) (realPart <$> y)
    xiyi = regEn (errorX "initial xiyi") en $ liftA2 mul (imagPart <$> x) (imagPart <$> y)
    xryi = regEn (errorX "initial xryi") en $ liftA2 mul (realPart <$> x) (imagPart <$> y)
    xiyr = regEn (errorX "initial xiyr") en $ liftA2 mul (imagPart <$> x) (realPart <$> y)
    --Sums
    rr   = regEn (errorX "initial rr")   en $ liftA2 sub xryr xiyi
    ri   = regEn (errorX "initial ri")   en $ liftA2 add xryi xiyr

cMul3Pipe
    :: HiddenClockResetEnable dom
    => ExtendingNum b b
    => ExtendingNum a a
    => ExtendingNum a (AResult b b)
    => ExtendingNum b (AResult a a)
    => ExtendingNum (MResult b (AResult a a)) (MResult a (AResult b b))
    => NFDataX (AResult a a)
    => NFDataX (AResult b b)
    => NFDataX (MResult b (AResult a a))
    => NFDataX (MResult a (AResult b b))
    => NFDataX (AResult (MResult b (AResult a a)) (MResult a (AResult b b)))
    => NFDataX a
    => NFDataX b
    => Signal dom Bool
    -> Signal dom (Complex a)
    -> Signal dom (Complex b)
    -> Signal dom (Complex (AResult (MResult b (AResult a a)) (MResult a (AResult b b))))
cMul3Pipe en x y
    =  liftA2 (:+) rr ri
    where
    --Sums
    xSum    = regEn (errorX "initial xSum")    en $ liftA2 add (realPart <$> x) (imagPart <$> x)
    yDiff   = regEn (errorX "initial yDiff")   en $ liftA2 sub (imagPart <$> y) (realPart <$> y)
    ySum    = regEn (errorX "initial ySum")    en $ liftA2 add (realPart <$> y) (imagPart <$> y)
    --Delays
    xD      = regEn (errorX "initial xD") en x
    yD      = regEn (errorX "initial yD") en y
    --Products
    yrXSum  = regEn (errorX "initial yrXSum")  en $ liftA2 mul (realPart <$> yD) xSum
    xrYDiff = regEn (errorX "initial xrYDiff") en $ liftA2 mul (realPart <$> xD) yDiff
    xiYSum  = regEn (errorX "initial xiYSum")  en $ liftA2 mul (imagPart <$> xD) ySum
    --Sums
    rr      = regEn (errorX "initial rr")      en $ liftA2 sub yrXSum xiYSum
    ri      = regEn (errorX "initial ri")      en $ liftA2 add yrXSum xrYDiff


-- | Tests for Kernel density estimates.
module Tests.KDE (
  tests
  )where

import Data.Vector.Unboxed             ((!))
import Numeric.Sum                     (kbn, sumVector)
import Statistics.Sample.KernelDensity
import Test.Tasty                      (TestTree, testGroup)
import Test.Tasty.QuickCheck           (testProperty)
import Test.QuickCheck                 (Property, (==>), counterexample)
import Text.Printf                     (printf)
import qualified Data.Vector.Unboxed as U

tests :: TestTree
tests = testGroup "KDE"
  [ testProperty "integral(PDF) == 1" t_densityIsPDF
  ]

t_densityIsPDF :: [Double] -> Property
t_densityIsPDF vec
  = not (null vec) ==> test
  where
    (xs,ys)  = kde 4096 (U.fromList vec)
    step     = (xs ! 1) - (xs ! 0)
    integral = integratePDF step ys
    --
    test = counterexample (printf "Integral %f" integral)
         $ abs (1 - integral) <= 1e-3

integratePDF :: Double -> U.Vector Double -> Double
integratePDF step vec
  = step * sumVector kbn (U.zipWith (*) vec weights)
  where
    n       = U.length vec
    weights = U.generate n go
      where
        go i | i == 0    = 0.5
             | i == n-1  = 0.5
             | otherwise = 1


module NeuralNet.Net (
  NeuralNetDefinition,
  LayerDefinition (LayerDefinition),
  WeightsStream,
  NeuralNet,
  initNN,
  buildNNFromList,
  nnLayers,
  nnNumInputs,
  nnForward,
  nnForwardSet,
  isExampleSetCompatibleWithNN,
  isExampleSetCompatibleWithNNDef,
  updateNNLayers,
  createLogRegDefinition
) where

import NeuralNet.Matrix
import NeuralNet.Activation
import NeuralNet.Layer
import NeuralNet.Example
import Numeric.LinearAlgebra

type NumNeurons = Int
type NumInputs = Int
type NeuralNetDefinition = (NumInputs, [LayerDefinition])

type WeightsStream = [Double]

data LayerDefinition = LayerDefinition Activation NumNeurons
  deriving (Show, Eq)

data NeuralNet = NeuralNet [NeuronLayer]
  deriving (Show, Eq)

createLogRegDefinition :: NumInputs -> Activation -> NeuralNetDefinition
createLogRegDefinition n a = (n, [LayerDefinition a 1])

initNN :: WeightsStream -> NeuralNetDefinition -> NeuralNet
initNN _ (_, []) = error "No layer definitions provided"
initNN g (numInputs, layerDefs)
  | numInputs > 0 = NeuralNet (fst (initNeuronLayers g numInputs layerDefs))
  | otherwise     = error "Need positive num inputs"

initNeuronLayers :: WeightsStream -> Int -> [LayerDefinition] -> ([NeuronLayer], WeightsStream)
initNeuronLayers weightsStream _ [] = ([], weightsStream)
initNeuronLayers weightsStream numInputs (LayerDefinition a numNeurons : ds) = (NeuronLayer a w b : ls, ws2)
  where
    numSeeds = numInputs * numNeurons
    (weights, ws) = splitAt numSeeds weightsStream
    w = scale (sqrt (2.0 / fromIntegral numInputs)) ((numNeurons >< numInputs) weights)
    -- * sqrt (2 / numInputs) should be parameterisable
    b = col (replicate numNeurons 0)
    (ls, ws2) = initNeuronLayers ws numNeurons ds

buildNNFromList :: NeuralNetDefinition -> [Double] -> NeuralNet
buildNNFromList def@(numInputs, layerDefs) list
  | listSize /= defSize = error ("Wrong list size. Expected " ++ show defSize ++ " got " ++ show listSize)
  | otherwise           = NeuralNet (buildLayersFromList numInputs layerDefs list)
    where
      defSize = definitionToNNSize def
      listSize = length list

-- TODO: nnToList and include back and forward of buildNNFromList to ensure works

definitionToNNSize :: NeuralNetDefinition -> Int
definitionToNNSize (numInputs, layerDefs) = defToSizeStep numInputs layerDefs
  where
    defToSizeStep :: Int -> [LayerDefinition] -> Int
    defToSizeStep _ [] = 0
    defToSizeStep stepInputs (LayerDefinition _ numNeurons:ds) = (stepInputs * numNeurons) + numNeurons + defToSizeStep numNeurons ds

buildLayersFromList :: Int -> [LayerDefinition] -> [Double] -> [NeuronLayer]
buildLayersFromList _ [] _ = []
buildLayersFromList numInputs (LayerDefinition a numNeurons:ds) xs = layer : (buildLayersFromList numNeurons ds restXs)
  where
    numWeights = numInputs * numNeurons
    (layerXs, afterLayerXs) = splitAt numWeights xs
    (bs, restXs) = splitAt numNeurons afterLayerXs
    layer = NeuronLayer a ((numNeurons><numInputs) layerXs) (col bs)

nnLayers :: NeuralNet -> [NeuronLayer]
nnLayers (NeuralNet layers) = layers

nnForward :: NeuralNet -> [Double] -> [Double]
nnForward nn inputs = concat (toLists (forwardPropA (last setResult)))
  where
    set = createExampleSet [(inputs, error "Undefined")]
    setResult = nnForwardSet nn set

nnForwardSet :: NeuralNet -> ExampleSet -> [ForwardPropStep]
nnForwardSet nn examples = reverse foldResult
  where
    x = exampleSetX examples :: Matrix Double
    input = createInputForwardPropStep x
    foldResult = foldl (\i l -> layerForwardSet l (forwardPropA (head i)) : i) [input] (nnLayers nn)

nnNumInputs :: NeuralNet -> Int
nnNumInputs = layerNumInputs . head . nnLayers

isExampleSetCompatibleWithNNDef :: ExampleSet -> NeuralNetDefinition -> Bool
isExampleSetCompatibleWithNNDef examples def = isExampleSetCompatibleWithNN examples emptyNN
  where
    nums = replicate (definitionToNNSize def) 0
    emptyNN = buildNNFromList def nums

isExampleSetCompatibleWithNN :: ExampleSet -> NeuralNet -> Bool
isExampleSetCompatibleWithNN examples nn = nnNumInputs nn == exampleSetN examples

updateNNLayers :: NeuralNet -> [NeuronLayer] -> NeuralNet
updateNNLayers nn newLayers
  | layersShapes layers /= layersShapes newLayers = error "Must provide same shaped layers"
  | otherwise                                     = NeuralNet newLayers
    where
      layersShapes :: [NeuronLayer] -> [((Int, Int), (Int, Int))]
      layersShapes sLayers = map (\l -> (dims (layerW l), dims (layerB l))) sLayers
      layers = nnLayers nn


-- |
-- Defines Log-domain probability density functions

{-# LANGUAGE ViewPatterns, NoMonomorphismRestriction, FlexibleInstances #-}
module Math.Probably.PDF where

import qualified Math.Probably.Student as S
import Numeric.LinearAlgebra
import Control.Spoon
import Control.DeepSeq
import qualified Data.Vector.Storable as VS

-- | The type of probablility density functions
type PDF a = a->Double

instance Show (a->Double) where
    show f = error "Math.Prob.PDF: showing function"

instance Eq (a->Double) where
    f == g = error "Math.Prob.PDF: comparing functions!"

instance Num a => Num (a->Double) where
    f + g = \x-> f x + g x
    f * g = \x-> f x * g x
    f - g = \x-> f x - g x
    abs f = \x -> abs (f x)
    signum f = \x-> signum $ f x
    fromInteger = error $ "Math.Prob.PDF: frominteger on function"

-- | Uniform distributions
uniform :: (Real a, Fractional a) => a-> a-> PDF a
uniform from to = \x-> if x>=from && x <=to
                               then log $ realToFrac $ 1/(to-from )
                               else 0

--http://en.wikipedia.org/wiki/Normal_distribution
-- | Normal distribution
gauss :: (Real a, Floating a) => a-> a-> PDF a
gauss mean sd x = realToFrac $ negate $ (x-mean)**2/(2*sd*sd) + log(sd*sqrt(2*pi))

-- | Normal distribution by variance
normal :: (Real a, Floating a) => a-> a-> a -> a
normal = normalLogPdf

normalLogPdf :: (Real a, Floating a) => a-> a-> a -> a
normalLogPdf mean variance x
   = ((log 1)-(0.500*(log ((2.000*pi)*variance))))-(((x-mean)**2)/(2*variance))

-- | Normal distribution, specialised for Doubles
gaussD :: Double -> Double-> PDF Double
gaussD mean sd x = negate $ (x-mean)**2/(2*sd*sd) + log(sd*sqrt(2*pi))

gammafun = exp . S.gammaln

-- | <http://en.wikipedia.org/wiki/Gamma_distribution>
gammaD :: Double -> Double -> PDF Double
gammaD k theta x = log $ x**(k-1)*(exp(-x/theta)/(theta**k*(exp (S.gammaln k))))

invGammaD :: Double -> Double -> PDF Double
invGammaD a b x =log $ (b**a/gammafun a)*(1/x)**(a+1)*exp(-b/x)

logNormal m var x = negate $ log (x*sqrt(2*pi*var)) + (log x - m)**2 / (2*var)

logNormalD :: Double -> Double-> PDF Double
logNormalD = logNormal

-- | Beta distribution
beta a b x = log $ (recip $ S.beta a b) * x ** (a-1) * (1-x) ** (b-1)

-- | Binomial distribution. This is the Probability Mass Function, not PDF
binomial :: Int -> Double -> PDF Int
binomial n p k =
    let realk = realToFrac k
    in log (realToFrac (choose (toInteger n) (toInteger k))) + realk * log( p ) + (realToFrac $ n-k) * log( (1-p) )
 where --http://blog.plover.com/math/choose.html
   choose :: Integer -> Integer -> Integer
   choose n 0 = 1
   choose 0 k = 0
   choose (n) (k) |  k > n `div` 2 = choose n (n-k)
                  |  otherwise = (choose (n-1) (k-1)) * (n) `div` (k)

-- | the joint distribution of two independent distributions
zipPdfs :: Num a => PDF a -> PDF b -> PDF (a,b)
zipPdfs fx fy (x,y) = fx x + fy y

-- | the multiplication of two PDFS
mulPdf :: Num a => PDF a -> PDF a -> PDF a
mulPdf d1 d2 = \x -> (d1 x + d2 x)

--instance Num a => Num (PDF a) where

--instance NFData (Matrix Double)
--   where rnf mat = mapMatrix (\x-> x `seq` 1.0::Double) mat `seq` ()

-- | multivariate normal
multiNormal :: Vector Double -> Matrix Double -> PDF (Vector Double)
multiNormal mu sigma =
  let k = realToFrac $ dim mu
      invSigma =  inv sigma
      mat1 = head . head . toLists
  in \x-> log (recip ((2*pi)**(k/2) * sqrt(det sigma))) + (mat1 $ negate $ 0.5*(asRow $ x-mu) `multiply` invSigma `multiply` (asColumn $ x-mu) )

multiNormalByInv :: Double -> Matrix Double -> Vector Double -> PDF (Vector Double)
multiNormalByInv lndet invSigma mu =
  let k = realToFrac $ dim mu
      mat1 = head . head . toLists
  in \x-> log 1 - (k/2)*log (2*pi) - lndet/2 + (mat1 $ negate $ 0.5*(asRow $ x-mu) `multiply` invSigma `multiply` (asColumn $ x-mu) )

multiNormalByInvFixCov ::  Matrix Double -> Vector Double -> PDF (Vector Double)
multiNormalByInvFixCov invSigma mu =
  let k = realToFrac $ dim mu
      mat1 = head . head . toLists
  in \x-> (mat1 $ negate $ 0.5*(asRow $ x-mu) `multiply` invSigma `multiply` (asColumn $ x-mu) )

multiNormalIndep :: Vector Double -> Vector Double -> PDF (Vector Double)
multiNormalIndep vars mus xs= VS.sum $ VS.zipWith3 (\var mu x -> normalLogPdf mu var x) vars mus xs

{-mu1 = 2 |> [0, 0::Double]
sig1 = (2><2)[1::Double, 0,
              0, 1]

tst = multiNormal mu1 sig1 mu1
tsta = inv sig1
tstb = det sig1
tstc = let mu = mu1
           sigma = sig1
           x = mu1
           invSigma = inv sigma
       in (asRow $ x-mu) `multiply` invSigma `multiply` (asColumn $ x-mu)  -}

posdefify m =
   let (eigvals, eigvecM) = eigSH $ mkSym {- $ trace (show m) -}  m
       n = rows m
       eigValsVecs = map f $ zip (toList eigvals) (toColumns eigvecM)
       f (val,vec) = (abs val,vec)
       q = fromColumns $ map snd eigValsVecs
       bigLambda = diag $ fromList $ map fst eigValsVecs
   in mkSym $ q `multiply` bigLambda `multiply` inv q

mkSym m = buildMatrix (rows m) (cols m)$ \(i,j) ->if i>=j then m @@>(i,j)
                                                           else m @@>(j,i)


{-# LANGUAGE RankNTypes #-}

module Kalman
  ( movingAverage,
    System (..),
    Estimate (..),
    Measurement (..),
    genSigmas,
    ekf,
    ut,
    ukf,
    weightedCrossCov,
    runKF,
  )
where

import Control.Applicative
import Numeric.AD
import Numeric.AD.Internal.Reverse
import Numeric.LinearAlgebra
import Sim
import System.Random

type Sigmas = [[Double]]

data System = System
  { fA, fH :: forall a. Floating a => ([a] -> [a]),
    kQ, kR :: Matrix Double
  }

data Estimate = Estimate
  { sX :: Vector Double,
    sP :: Matrix Double
  }
  deriving (Show)

-- Simple 2 tap FIR to appx moving average filter
movingAverage :: [Vector Double] -> [Vector Double]
movingAverage = scanl1 (\v1 v2 -> (0.9 * v1) + (0.1 * v2))

-- kalman :: System -> Estimate -> Measurement -> Estimate
-- kalman (System a c q r) (Estimate x p) z = Estimate x' p' where
--   -- Convert lists to matrices since we have linear system
--   lA = fromLists a
--   lC = fromLists c
--   -- prediction step
--   px = lA #> x
--   pp = lA <> p <> tr lA + q
--   -- update step
--   g = pp <> tr lC <> inv (lC <> pp <> tr lC + r)
--   x' = px + g #> (z - lC #> px)
--   p' = (ident (size x) - g <> lC) <> pp

runKF :: (System -> Estimate -> Measurement -> Estimate) -> System -> Estimate -> [Measurement] -> [Estimate]
runKF kf sys = scanl (kf sys)

ekf :: System -> Estimate -> Measurement -> Estimate
ekf (System fA fH kQ kR) (Estimate x p) z = Estimate x' p'
  where
    -- prediction step
    a = fromLists (jacobian fA (toList x))
    px = fA $ toList x
    pp = a <> p <> tr a + kQ
    -- update step
    h = fromLists $ jacobian fH px
    k = pp <> tr h <> inv (h <> pp <> tr h + kR)
    x' = vector px + k #> (z - vector (take (size z) $ fH px))
    p' = (ident (size x) - k <> h) <> pp

genSigmas :: Estimate -> Double -> Sigmas
genSigmas (Estimate x p) c = sigmas
  where
    -- Cholesky Factorization to get sqrt of matrix
    cP = scale (sqrt c) $ chol $ sym p
    -- Add/sub the sqrt covariance matrix from the current estimate of x
    xs = fromLists $ replicate (size x) (toList x)
    pos = toLists $ xs + cP
    neg = toLists $ xs - cP
    -- Combine all sigma points including the original estimate of x at the head
    sigmas = [toList x] ++ pos ++ neg

-- Unscented Transformation
ut :: [Vector Double] -> [Double] -> [Double] -> Estimate
ut xs wm wc = Estimate x' p'
  where
    -- Weight each sigma by scaling factor
    xw = getZipList $ scale <$> ZipList wm <*> ZipList xs
    x' = sum xw
    -- Compute forecasted state statistics
    v1 = map (\v -> fromColumns [v - x']) xs
    v2 = map (\v -> fromRows [v - x']) xs
    cov = getZipList $ (<>) <$> ZipList v1 <*> ZipList v2
    covw = getZipList $ scale <$> ZipList wc <*> ZipList cov
    p' = sum covw

weightedCrossCov :: Vector Double -> [Vector Double] -> Vector Double -> [Vector Double] -> [Double] -> Matrix Double
weightedCrossCov v1 vs1 v2 vs2 wc = cP
  where
    vs1' = map (\v -> fromColumns [v - v1]) vs1
    vs2' = map (\v -> fromRows [v - v2]) vs2
    cov = getZipList $ (<>) <$> ZipList vs1' <*> ZipList vs2'
    covw = getZipList $ scale <$> ZipList wc <*> ZipList cov
    cP = sum covw

ukf :: System -> Estimate -> Measurement -> Estimate
ukf (System fA fH kQ kR) (Estimate x p) z = Estimate xh ph
  where
    l = fromIntegral $ size x
    m = fromIntegral $ size z
    alpha = 1e-3
    ki = 0
    beta = 2

    lambda = alpha ^ 2 * (l + ki) - l
    c = l + lambda

    wm = (lambda / c) : replicate (2 * fromIntegral (size x)) (0.5 / c)
    wc = ((lambda / c) + (1 - alpha ^ 2 + beta)) : replicate (2 * fromIntegral (size x)) (0.5 / c)

    -- Generate sigma points around current estimate
    sigmas = genSigmas (Estimate x p) c

    -- Run sigmas through non linear function to produce new points
    xs = map (fromList . fA) sigmas
    zs = map (fromList . fH) sigmas

    Estimate x' cX = ut xs wm wc
    Estimate z' cZ = ut zs wm wc

    -- Calculate the weighted cross covariance
    cXZ = weightedCrossCov x' xs z' zs wc

    kK = cXZ <> inv (cZ + kR)
    xh = x' + kK #> (z - z')
    ph = (cX + kQ) - kK <> (cZ + kR) <> tr kK


module OrientationAnalysisShape where

import           Control.Monad           as M
import           Data.Array.Repa         as R
import           Data.Complex
import           Data.List               as L
import           DFT.Plan
import           Filter.Pinwheel
import           FokkerPlanck.MonteCarlo
import           FokkerPlanck.Pinwheel
import           Image.IO
import           STC
import           STC.OrientationScaleAnalysis
import           System.Directory
import           System.Environment
import           System.FilePath

main = do
  args@(numPointStr:theta0FreqsStr:alphaStr:idxStr:numOrientationSampleStr:shapeStr:whiteBackgroundFlagStr:numThreadStr:_) <-
    getArgs
  print args
  let numPoint = read numPointStr :: Int
      theta0Freq = read theta0FreqsStr :: Double
      theta0Freqs = [-theta0Freq .. theta0Freq]
      alpha = read alphaStr :: Double
      idx = read idxStr :: (Int, Int)
      numOrientationSample = read numOrientationSampleStr :: Int
      whiteBackgroundFlag = read whiteBackgroundFlagStr :: Bool
      numThread = read numThreadStr :: Int
      pinwheelParams =
        PinwheelParams numPoint numPoint alpha (exp 1) theta0Freqs [0]
      shape2D = read shapeStr :: Shape2D
      shapeArr' =
        makeShape2DContrast numPoint numPoint (ContrastArea whiteBackgroundFlag shape2D)
      folderPath = "output/test/OrientationAnalysisShape"
  createDirectoryIfMissing True folderPath
  shapeArr <- computeUnboxedP shapeArr'

  (plan1, imgF) <-
    makeImagePlan emptyPlan .
    computeS .
    R.map fromIntegral .
    R.backpermute
      (Z :. (3 :: Int) :. numPoint :. numPoint)
      (\(Z :. _ :. i :. j) -> (Z :. i :. j)) $
    shapeArr
  (plan, filterF, filterPIF) <- pinwheelFilter plan1 pinwheelParams
  convolvedImg <-
    fmap (\x -> R.slice x (Z :. All :. (0 :: Int) :. All :. All)) $
    convolvePinwheel plan filterF (R.slice imgF (Z :. (0 :: Int) :. All :. All))
  idx' <- plotMagnitudeOrientation
            folderPath
            numOrientationSample
            theta0Freqs
            convolvedImg
            idx
  plotImageRepa (folderPath </> "Shape.png") .
    ImageRepa 8 .
    computeS .
    R.traverse
      shapeArr
      (\(Z :. cols :. rows) -> (Z :. (3 :: Int) :. cols :. rows)) $
    (\f (Z :. k :. i :. j) ->
       if (i, j) == idx'
         then if k == 0
                then 255
                else 0
         else fromIntegral $ f (Z :. i :. j))


module CubicBezier ( CubicBezier (..)
                   , pointAt
                   , bezierSplitAt
                   , isClockwise
                   , inflectionPoints
                   , realInflectionPoint
                   ) where

import Linear
import Control.Lens
import Data.Complex

data CubicBezier = CubicBezier { _p1 :: V2 Double
                               , _c1 :: V2 Double
                               , _c2 :: V2 Double
                               , _p2 :: V2 Double
                               } deriving Show

pointAt :: CubicBezier -> Double -> V2 Double
pointAt bezier t =  ((1 - t) ** 3) *^ _p1 bezier +
                    ((1 - t) ** 2) * 3 * t *^ _c1 bezier +
                    (t ** 2) * (1 - t) * 3 *^ _c2 bezier +
                    (t ** 3) *^ _p2 bezier

bezierSplitAt :: CubicBezier -> Double -> (CubicBezier, CubicBezier)
bezierSplitAt bezier t = (CubicBezier (_p1 bezier) p0 p01 dp, CubicBezier dp p12 p2 (_p2 bezier))
    where
        p0 = _p1 bezier + t *^ (_c1 bezier - _p1 bezier)
        p1 = _c1 bezier + t *^ (_c2 bezier - _c1 bezier)
        p2 = _c2 bezier + t *^ (_p2 bezier - _c2 bezier)

        p01 = p0 + t *^ (p1 - p0)
        p12 = p1 + t *^ (p2 - p1)

        dp = p01 + t *^ (p12 - p01)

isClockwise :: CubicBezier -> Bool
isClockwise bezier = s < 0
    where
        s = (_c1 bezier ^. _x - _p1 bezier  ^. _x) * (_c1 bezier ^. _y + _p1 bezier ^. _y)
          + (_c2 bezier ^. _x - _c1 bezier  ^. _x) * (_c2 bezier ^. _y + _c1 bezier ^. _y)
          + (_p2 bezier ^. _x - _c2 bezier  ^. _x) * (_p2 bezier ^. _y + _c2 bezier ^. _y)
          + (_p1 bezier ^. _x - _p2 bezier  ^. _x) * (_p1 bezier ^. _y + _p2 bezier ^. _y)

inflectionPoints :: CubicBezier -> (Complex Double, Complex Double)
inflectionPoints bezier = (t1, t2)
    where
        pa = _c1 bezier - _p1 bezier
        pb = _c2 bezier - _c1 bezier - pa
        pc = _p2 bezier - _c2 bezier - pa - 2 *^ pb

        a = (pb ^. _x * pc ^. _y - pb ^. _y * pc ^. _x) :+ 0
        b = (pa ^. _x * pc ^. _y - pa ^. _y * pc ^. _x) :+ 0
        c = (pa ^. _x * pb ^. _y - pa ^. _y * pb ^. _x) :+ 0

        t1 = (-b + sqrt (b * b  - 4 * a * c)) / (2 * a)
        t2 = (-b - sqrt (b * b  - 4 * a * c)) / (2 * a)

realInflectionPoint :: Complex Double -> Bool
realInflectionPoint c = imagPart c == 0 && realPart c > 0 && realPart c < 1


import Data.Complex

-- Cooley-Tukey
fft [] = []
fft [x] = [x]
fft xs = zipWith (+) ys ts ++ zipWith (-) ys ts
    where n = length xs
          ys = fft evens
          zs = fft odds
          (evens, odds) = split xs
          split [] = ([], [])
          split [x] = ([x], [])
          split (x:y:xs) = (x:xt, y:yt) where (xt, yt) = split xs
          ts = zipWith (\z k -> exp' k n * z) zs [0..]
          exp' k n = cis $ -2 * pi * (fromIntegral k) / (fromIntegral n)

main = mapM_ print $ fft [1,1,1,1,0,0,0,0]


module Main where

import Math.Probably.Sampler
import Math.Probably.RandIO
import Math.Probably.FoldingStats

import Numeric.LinearAlgebra

import Math.Probably.Unscented

import Graphics.Gnewplot.Exec
import Graphics.Gnewplot.Types
import Graphics.Gnewplot.Style
import Graphics.Gnewplot.Panels
import Graphics.Gnewplot.Instances
import Graphics.Gnewplot.Histogram

dt = 0.1

f v =  fromList [t1, x0 + dt * sin (t1) + noisex]
  where t0 = v @> 0
        x0 = v @> 1
        noiset = v @> 2
        noisex = v @> 3
        t1 = t0+dt + noiset

f' v =  fromList [t1, x0 + dt * sin (t1) ]
  where t0 = v @> 0
        x0 = v @> 1
        t1 = t0+dt

g v = fromList [t1+noiset, x1+noisex]
  where t1 = v @> 0
        x1 = v @> 1
        noiset = v @> 2
        noisex = v @> 3

g' v = fromList [t1, x1]
  where t1 = v @> 0
        x1 = v @> 1

procCov = (2><2) [0.001, 0,
                  0,     0.01]

obsCov = (2><2) [0.01, 0,
                  0,     0.01]

--obsCov = (1><1)  [0.01]

testUT = do
  let cov = (2><2) [0.1,0.08,0.08,0.1]
      mn = (2 |> [1,1])
      xform v = let [x,y] = toList v in fromList [x*x, y*y]
      mkpt = ((\[x,y] -> (x,y)) .toList)
  pts <- sample $ sequence $ replicate 1000 $ multiNormal mn cov
  let (mnU, covU, sigmapts) = unscentedTransform (mn,cov) xform
  let xformPts = map xform pts
  io $ gnuplotOnScreen $ map mkpt pts :+: map mkpt sigmapts --  :+: map mkpt xformPts
  io $ print mnU
  io $ print covU
  (mnE, covE) <- sample $ smellyTransform (multiNormal mn cov) 1000 xform
  io $ print mnE
  io $ print covE

main = runRIO $ do
     xyvs <- sample $ simDynamical 2000 f g procCov obsCov (fromList [0, 0.1])
     --let xyvs = flip map xyvs' $ \(xv,yv) -> (xv,negate yv)
--     pts <- sample $ sequence $ replicate 1000 $ lastSDESampler dt
     --w <- sample $ fmap (eulerMaruyama dt 0 a b) $ weiner dt 1
     let xhats = unscentedKalmanFilterAdditive f' g' procCov obsCov (fromList [0,0], ((2><2) [1,0,0,1])) $ map snd xyvs
     let get s f =  (s, Lines [] $ zip [(0::Double)..] $ map f xyvs)
     io $ gnuplotOnScreen $ get "y" ((@>1) . snd) :+: get "x" ((@>1) . fst)
                            :+: ("xhat", Lines [] $ zip [(0::Double)..] $ map ((@>1) . fst)  xhats)
     io $ print $ take 10 $ map snd xyvs

     testUT
--     io $ gnuplotOnScreen $ get "x" ((@>1) . fst)
--     io $ print $ runStat meanSDF $ pts


module NeuralA where

import Numeric.LinearAlgebra

x = vector [-5.0,-4.9..5.0]
y1 = cmap sigmoid x

sigmoid :: Double -> Double
sigmoid x = 1.0 / (1.0 + exp(-x))

step_func :: Double -> Double
step_func x
  | x > 0.0 = 1.0
  | otherwise = 0.0

relu :: Double -> Double
relu x
  | x > 0.0 = x
  | otherwise = 0.0


module Embadd where

import Control.Arrow (second)
import Data.Function (on)
import Data.List (foldl', sortBy)
import Data.Map.Strict (Map, empty, findWithDefault, insertWith, size)
import qualified Data.Map.Strict as M (fromList, toList, size, lookup)
import Data.Array (listArray, elems)

import Numeric.LinearAlgebra as LA
import System.Console.Readline

import Numeric.Statistics.PCA

import Graphics.Rendering.Chart hiding (Vector)
import Graphics.Rendering.Chart.Backend.Cairo
import Data.Default.Class
import Control.Lens hiding ((<.>))

type VDict = Map String (Vector R)

data Expr = Wrd String
          | Plus Expr Expr
          | UnMinus Expr
          | NullExpr

cosineD a b = (a LA.<.> b) / (norm_2 a * norm_2 b)

chart ws = toRenderable layout
  where
    layout = layout_title .~ ""
           $ layout_plots .~ [toPlot texts]
           $ def

    texts = plot_annotation_values .~ ws
          $ plot_annotation_angle .~ 0 -- degrees
          $ plot_annotation_style .~ fontStyle
          $ plot_annotation_hanchor .~ HTA_Centre -- the horizontal anchor around which the text is rotated
          $ plot_annotation_vanchor .~ VTA_Centre -- the vertical anchor
          $ def

    fontStyle = font_size .~ 16
              $ font_weight .~ FontWeightBold
              $ def

parseExpr :: Expr -> [String] -> Expr
parseExpr NullExpr ("-":w:ws) = parseExpr (UnMinus (Wrd w)) ws
parseExpr NullExpr (w:ws) = parseExpr (Wrd w) ws
parseExpr e ("+":w:ws) = parseExpr (Plus e (Wrd w)) ws
parseExpr e ("-":w:ws) = parseExpr (Plus e (UnMinus (Wrd w))) ws
parseExpr e [] = e
parseExpr _ _ = NullExpr

sumExpr vecDict (Plus e1 e2) = do
  a <- sumExpr vecDict e1
  b <- sumExpr vecDict e2
  return $ a + b
sumExpr vecDict (UnMinus e) = do
  x <- sumExpr vecDict e
  return $ - x
sumExpr vecDict (Wrd w) = M.lookup w vecDict

nearest vDict v1 = map fst . sortBy (flip compare `on` (cosineD v1 . snd)) $ M.toList vDict

eWrds (Wrd w) = [w]
eWrds (UnMinus e) = eWrds e
eWrds (Plus a b) = eWrds a ++ eWrds b

readEvalPrintLoop :: VDict -> IO ()
readEvalPrintLoop vecDict = do
  maybeLine <- readline "> "
  case words <$> maybeLine of
    Nothing     -> return () -- EOF / control-d
    Just ("PCA":wrds) -> do
      addHistory $ unwords ("PCA":wrds)
      case mapM (`M.lookup` vecDict) wrds :: Maybe [Vector R] of
        Nothing -> putStrLn "error: OOV"
        Just vcs -> do
          let vData = listArray (1, LA.size $ head vcs) . toRows $ fromColumns vcs
              reduced = toColumns . fromRows . elems . pcaTransform vData . snd $ pcaN vData 2
          print reduced
          renderableToFile def "pca.png" . chart . map (\(w,v) -> (v ! 0, v ! 1, w)) $ zip wrds reduced
          return ()
      readEvalPrintLoop vecDict
    Just wrds -> do
      addHistory $ unwords wrds
      case parseExpr NullExpr wrds of
        NullExpr -> putStrLn "bad parse"
        e -> case sumExpr vecDict e of
          Just eSum -> putStrLn . unlines . take 7 . filter (\x -> (x `notElem` eWrds e)) $ nearest vecDict eSum
          Nothing -> putStrLn "error: OOV"
      readEvalPrintLoop vecDict

computeNorm :: FilePath -> IO ()
computeNorm vectorFile = do
  vectorCont <- readFile vectorFile
  let vLine (w:ws) = (w, vector $ map read ws :: Vector Double)
      norms = map (second (show . norm_2) . vLine . words) . drop 1 . lines $ vectorCont
  mapM_ (putStrLn .  (\(a, b) -> a ++ "\t" ++ b)) norms


module Main where

import Numeric.LinearAlgebra
import Hoqus.Dirac

main = do
  let dim = 12
  let a = ket dim 3
  let b = ket dim 6
  let res0 = a <.> b
  let res1 = a <.> a
  print res0
  print res1


{-# LANGUAGE DataKinds                            #-}
{-# LANGUAGE FlexibleContexts                     #-}
{-# LANGUAGE FlexibleInstances                    #-}
{-# LANGUAGE MultiParamTypeClasses                #-}
{-# LANGUAGE PartialTypeSignatures                #-}
{-# LANGUAGE TupleSections                        #-}
{-# LANGUAGE TypeApplications                     #-}
{-# LANGUAGE TypeOperators                        #-}
{-# LANGUAGE TypeSynonymInstances                 #-}
{-# LANGUAGE UndecidableInstances                 #-}
{-# OPTIONS_GHC -fno-warn-orphans                 #-}
{-# OPTIONS_GHC -fno-warn-partial-type-signatures #-}

import           Backprop.Learn
import           Control.Monad.Trans.Maybe
import           Data.Bitraversable
import           Data.Default
import           Data.IDX
import           Data.Traversable
import           Data.Tuple
import           Numeric.LinearAlgebra.Static.Backprop
import           Numeric.Opto
import           Numeric.Opto.Run.Simple
import           System.Environment
import           System.FilePath
import           Text.Printf
import qualified Data.Vector.Generic                   as VG
import qualified Numeric.LinearAlgebra                 as HM
import qualified Numeric.LinearAlgebra.Static          as H
import qualified System.Random.MWC                     as MWC

mnistNet :: Model _ _ (R 784) (R 10)
mnistNet = fca     @300 softMax
        <~ dropout      0.25
        <~ fca     @784 logistic

main :: IO ()
main = MWC.withSystemRandom $ \g -> do
    datadir:_  <- getArgs
    Just train <- loadMNIST (datadir </> "train-images-idx3-ubyte")
                            (datadir </> "train-labels-idx1-ubyte")
    Just test  <- loadMNIST (datadir </> "t10k-images-idx3-ubyte")
                            (datadir </> "t10k-labels-idx1-ubyte")
    putStrLn "Loaded data."
    net0 <- initParamNormal mnistNet 0.2 g

    let so = def { soTestSet   = Just test
                 , soEvaluate  = runTest
                 , soSkipSamps = 2500
                 }
        runTest chnk net = printf "Error: %.2f%%" ((1 - score) * 100)
          where
            score = testModelAll maxIxTest mnistNet (TJust net) chnk

    simpleRunner so train SOSingle def net0
      (adam def $ modelGradStoch crossEntropy noReg mnistNet g)
      g

loadMNIST
    :: FilePath
    -> FilePath
    -> IO (Maybe [(R 784, R 10)])
loadMNIST fpI fpL = runMaybeT $ do
    i <- MaybeT          $ decodeIDXFile       fpI
    l <- MaybeT          $ decodeIDXLabelsFile fpL
    d <- MaybeT . return $ labeledIntData l i
    MaybeT . return $ for d (bitraverse mkImage mkLabel . swap)
  where
    mkImage = H.create . VG.convert . VG.map (\i -> fromIntegral i / 255)
    mkLabel n = H.create $ HM.build 10 (\i -> if round i == n then 1 else 0)


module ML.Unsupervised.GMM (
  fit,
  GMM(..),
  GMMComponent(..),
  FitOptions(..),
  defaultOptions,
  predict,
  empty
) where

import Numeric.LinearAlgebra
import qualified Data.Vector.Storable as V
import Statistics.Distribution
import qualified Data.List as L

import qualified ML.Unsupervised.Clustering.KMeans as KMeans

import Debug.Trace

data GMMComponent a = GMMComponent {
  gmmcPi :: Double,
  gmmcMean :: V.Vector a,
  gmmcSigma :: Matrix a
} deriving (Show)

data GMM a = GMM {
  mixtureComponents :: [GMMComponent a]
} deriving (Show)

empty :: GMM a
empty = GMM { mixtureComponents = [] }

normalPdf :: V.Vector Double -> Matrix Double -> V.Vector Double -> Double
normalPdf mu sigma x = (1 / sqrt (2 * pi * det sigma)) * exp (-0.5 * (x - mu) <.> ((inv sigma) #> (x - mu)))

vecsum (v:vs) = foldr (V.zipWith (+)) v vs
vecsum [] = error "Empty list in vecsum"

data FitOptions = FitOptions {
  foMaxIterations :: Int,
  foLikelihoodConvergenceRatio :: Double
}

defaultOptions = FitOptions {
  foMaxIterations = 100,
  foLikelihoodConvergenceRatio = 1.0e-12
}

likelihood :: GMM Double -> [V.Vector Double] -> Double
likelihood gmm points = sum $ fmap (logL gmm) points
  where
    logL gmm point = log $ sum $ fmap (logLFromComponent point) (mixtureComponents gmm)
    logLFromComponent point comp = (gmmcPi comp) * normalPdf (gmmcMean comp) (gmmcSigma comp) point

fit :: FitOptions -> Int -> [V.Vector Double] -> GMM Double
fit options nComponents points = fit' (GMM initialComponents) 0
  where
    fit' gmm itnum = let updatedGmm = update gmm in
                if (abs (likelihood updatedGmm points / likelihood gmm points - 1) < foLikelihoodConvergenceRatio options) || (itnum >= foMaxIterations options)
                  then gmm
                  else fit' updatedGmm (trace ("fit' " ++ show itnum) itnum + 1)

    calculateResponsibility gmm comp = zipWith (/) (fmap (respNumerator comp) points) (fmap (respDenominator gmm) points)

    respNumerator component point = (gmmcPi component) * normalPdf (gmmcMean component) (gmmcSigma component) point

    respDenominator gmm point = sum $ fmap (\x -> respNumerator x point) (mixtureComponents gmm)

    update gmm = GMM $ fmap (updateComponent gmm) (mixtureComponents gmm)

    updateComponent gmm comp = comp {gmmcPi = updatePi gmm comp, gmmcMean = updateMu gmm comp, gmmcSigma = updateSigma gmm comp }

    updatePi gmm comp = (sum $ calculateResponsibility gmm comp) / (fromIntegral $ length points)

    updateMu gmm comp = V.map (/ (sum $ calculateResponsibility gmm comp)) (vecsum $ zipWith (\x r -> V.map (* r) x) points (calculateResponsibility gmm comp))

    updateSigma gmm comp = (sum $ zipWith (zipR comp) (calculateResponsibility gmm comp) points) / scalar (sum $ calculateResponsibility gmm comp)

    zipR :: GMMComponent Double -> Double -> V.Vector Double -> Matrix Double
    zipR comp x p = scalar x * ((p - gmmcMean comp) `outer` (p - gmmcMean comp))

    dim = V.length $ head points

    initialComponents = let km = KMeans.fit [1..(nComponents)] points
                            predicted = fmap (\p -> (KMeans.predict km p, p)) points in
                          fmap (mkComponent predicted) (KMeans.centroids km)

    mkComponent predicted cent = let filteredPoints = filter (\x -> fst x == fst cent) predicted in
                         GMMComponent { gmmcPi = fromIntegral (length filteredPoints) / fromIntegral (length points),
                                        gmmcMean = snd cent,
                                        gmmcSigma = ident dim * scalar (sqrt ((foldr (\v x -> x + dist (snd cent) (snd v)) 0 filteredPoints) / fromIntegral (length points))) }

    dist :: V.Vector Double -> V.Vector Double -> Double
    dist v1 v2 = (v1 - v2) <.> (v1 - v2)

predict :: GMM Double -> V.Vector Double -> Int
predict gmm point = fst $ L.maximumBy (\x y -> snd x `compare` snd y) $ zip [0..] $ fmap (\comp -> (gmmcPi comp) * normalPdf (gmmcMean comp) (gmmcSigma comp) point) (mixtureComponents gmm)


{-# LANGUAGE DataKinds #-}
module Grenade.Utils.LinearAlgebra
    ( sumV
    , sumM
    , squareV
    , squareM
    ) where

import           GHC.TypeLits
import           Grenade.Types
import           Numeric.LinearAlgebra.Static

sumV :: (KnownNat n) => R n -> RealNum
sumV v = v <.> 1

sumM :: (KnownNat m, KnownNat n) => L m n -> RealNum
sumM m = (m #> 1) <.> 1

squareV :: (KnownNat n) => R n -> R n
squareV v = dvmap (^ (2 :: Int)) v

squareM :: (KnownNat m, KnownNat n) => L m n ->  L m n
squareM m = dmmap (^ (2 :: Int)) m

-- v :: R 10
-- v = vector [1..10]

-- m :: L 3 3
-- m = matrix [-3..5]


{-# LANGUAGE NoMonomorphismRestriction #-}

module Main where

import LinAlg (testCols, testRows)

import Numeric.LinearAlgebra.Data
import Numeric.LinearAlgebra.HMatrix

import Data.Vector.Storable

import Prelude hiding (sum)

testMult n m = sum $ mat #> vec
  where
    mat = matrix n $ fmap fromIntegral [1..(n*m)]
    vec = fromList $ fmap fromIntegral [1..n]

main = print $ testMult testCols testRows


{-# LANGUAGE TupleSections #-}
module SparseMatrix (
  pageRank
                    ) where

import Data.Graph (Graph, Edge, graphFromEdges, edges, outdegree)
import Data.Array ((!), range, bounds)

import qualified Numeric.LinearAlgebra as M
import Numeric.LinearAlgebra ((!#>), tr, mkSparse, cmap, size, konst)

import Numeric.Natural (Natural)

import qualified Data.Map.Strict as Map
import Data.Map.Strict (keys, mapWithKey)

type RealNum = Double
type Vector = M.Vector RealNum
type Matrix = M.GMatrix
type Map a b = Map.Map a b

pageRank :: (Ord a) => RealNum -> Natural -> Map a [a] -> Map a RealNum
pageRank damp its as = Map.fromList . zip (keys as) . M.toList . pageRank' damp its $ as

pageRank' :: (Ord a) => RealNum -> Natural -> Map a [a] -> Vector
pageRank' damp its m =
  (!! fromIntegral its)
  . flip iterate (uniform . length $ m) . (.) (dampen damp) . (!#>)
  . pageRankMatrix $ m

uniform :: Int -> Vector
uniform n = konst (recip . fromIntegral $ n) n

dampen :: RealNum -> Vector -> Vector
dampen factor v = cmap ((+) ((1-factor)/(fromIntegral . size $ v)) . (*) factor) v

pageRankMatrix :: (Ord a) => Map a [a] -> Matrix
pageRankMatrix = adjacency . normalize . unstick . connectionGraph

normalize :: Graph -> Map Edge RealNum
normalize g = mapWithKey (\ (f,_) x -> x / out f) . Map.fromAscList . fmap (,1) . edges $ g
  where
    out = fromIntegral . (outdegree g !)

adjacency :: Map Edge RealNum -> Matrix
adjacency = tr . mkSparse . Map.toList

connectionGraph :: (Ord a) => Map a [a] -> Graph
connectionGraph = (\ (x,_,_) -> x) . graphFromEdges . fmap (\ (x,xs) -> (x,x,xs)) . Map.toList

unstick :: Graph -> Graph
unstick g = fmap (\ l -> if null l then (range . bounds) g else l) g


{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE BangPatterns #-}

{-|
  This module implements PDF instances
  that are missing from random-fu.
  We also include functions to conveniently
  construct values of type SomeDist.
 -}

module Model.PDF where

import Data.Random
import Data.Random.Distribution
import Data.Random.Distribution.Bernoulli hiding ( bernoulli )
import Data.Random.Distribution.Beta hiding ( beta )
import Data.Random.Distribution.Binomial hiding ( binomial )
import Data.Random.Distribution.Categorical hiding ( categorical )
import Data.Random.Distribution.Exponential hiding ( exponential )
import Data.Random.Distribution.Gamma hiding ( gamma )
import Data.Random.Distribution.Normal hiding ( normal )
import Data.Random.Distribution.Poisson hiding ( poisson )
import Data.Random.Distribution.T hiding ( t )
import Data.Random.Distribution.Uniform hiding ( uniform )

import qualified Data.Vector as V

import Model.Types

import Numeric.SpecFunctions

bernoulli :: Double -> SomeDist Bool
bernoulli = SomeDist . Bernoulli

constant :: Eq a => a -> SomeDist a
constant a = categorical [(1, a)]

beta :: Double -> Double -> SomeDist Double
beta a b = SomeDist $ Beta a b

binomial :: Int -> Double -> SomeDist Int
binomial n p = SomeDist $ Binomial n p

normal :: Double -> Double -> SomeDist Double
normal m v = SomeDist $ Normal m v

categorical :: Eq a => [(Double, a)] -> SomeDist a
categorical = SomeDist . fromList

exponential :: Double -> SomeDist Double
exponential = SomeDist . Exp

gamma :: Double -> Double -> SomeDist Double
gamma a b = SomeDist $ Gamma a b

poisson :: Double -> SomeDist Int
poisson = SomeDist . Poisson

uniform :: Double -> Double -> SomeDist Double
uniform l u = SomeDist $ Uniform l u

t :: Integer -> SomeDist Double
t = SomeDist . T

instance PDF (Bernoulli Double) Bool where
  pdf (Bernoulli p) bool = if bool then p else 1 - p
  logPdf d = log . pdf d

instance Eq a => PDF (Categorical Double) a where
  pdf d val = scan l val / (sum $ fst <$> l)
    where scan [] x = 0
          scan ((p, x) : xs) v = if x == v then p else scan xs v
          l = toList d
  logPdf d = log . pdf d

instance PDF Exponential Double where
  pdf (Exp d) x = if x >= 0 then d * exp(-d * x) else 0
  logPdf d = log . pdf d

-- | TODO: Test which parameterization this is.
-- For now, hypothesize that it is the shape-rate parameterization
-- parameterization, where the expectation is (a/b)
-- and the pdf is $ \frac{\beta^\alpha}{\Gamma(\alpha)}  x^{\alpha-1} e^{-\beta x} $.
instance PDF Gamma Double where
  pdf d = exp . logPdf d
  logPdf (Gamma !a !b) !x = a * log b - logGamma a + (a - 1) * log x - b * x

instance PDF (Poisson Double) Int where
  pdf d = exp . logPdf d
  logPdf (Poisson !l) !x = -l + (fromIntegral x * log l) - log (fromIntegral (fact x))

-- | Factorial
fact :: Int -> Int
fact 0 = 1
fact !n = if n < 0 then error "Factorial called on negative" else fact' n 1
  where fact' 1 !m = m
        fact' !m' !m = fact' (m'-1) (m * m')

instance PDF Uniform Double where
  pdf (Uniform !l !u) !x = if x >= l && x <= u then (1 / (u - l)) else 0
  logPdf d = log . pdf d

instance PDF T Double where
  pdf d = exp . logPdf d
  logPdf (T !nu) !x = logGamma ((nud + 1) / 2) - (1/2) * log (nud * pi) - logGamma (nud / 2) + ((-nud-1) / 2) * logGamma (1 + ((x ^ 2) / nud))
    where nud = fromIntegral nu


{-|
Module      : Julia
Description : Computing Julia sets
Copyright   : (c) Jason Chen, 2021
License     : MIT
Maintainer  : chezbgone (at) gmail dot com
Stability   : experimental
-}
{-# LANGUAGE NamedFieldPuns #-}
module Julia where

import Data.Complex
import Data.List (iterate')
import Codec.Picture

import Interpolation

data Resolution = Resolution { width :: Int, height :: Int }

data Bounds = Bounds { top    :: Double
                     , bottom :: Double
                     , left   :: Double
                     , right  :: Double
                     }

data ComputationProps = ComputationProps
    { maxIterations :: Int
    , escapeRadius :: Double
    }

data RenderProps = RenderProps
    { resolution :: Resolution
    , bounds     :: Bounds
    }

data JuliaProps = JuliaProps
    { computationProps :: ComputationProps
    , renderProps      :: RenderProps
    }

escapeTime :: (Complex Double -> Complex Double) -- f
           -> ComputationProps
           -> Complex Double -- init
           -> Int
escapeTime f ComputationProps{ maxIterations, escapeRadius } init =
    length $ takeWhile (\z -> magnitude z < escapeRadius)
           $ take maxIterations
           $ iterate' f init

escapePixel :: Int -> PixelRGB8
escapePixel = undefined

gridToComplex :: RenderProps -> Int -> Int -> Complex Double
gridToComplex RenderProps{ resolution, bounds } x y
  = (left + xStep * fromIntegral x) :+ (top + yStep * fromIntegral y)
      where Resolution {width, height} = resolution
            Bounds {top, bottom, left, right} = bounds
            xStep = (right - left) / fromIntegral width
            yStep = (top - bottom) / fromIntegral height

generateJulia :: JuliaProps
              -> (Complex Double -> Complex Double) -- f
              -> FilePath
              -> IO ()
generateJulia JuliaProps{ computationProps, renderProps }
              f filepath = writePng filepath image
    where image = generateImage compute width height
          Resolution{ width, height } = resolution renderProps
          compute :: Int -> Int -> PixelRGB8
          compute x y = escapePixel $ escapeTime f computationProps
                                    $ gridToComplex renderProps x y


{-|
Module: MachineLearning.NeuralNetwork.WeightInitialization
Description: Weight Initialization
Copyright: (c) Alexander Ignatyev, 2017
License: BSD-3
Stability: experimental
Portability: POSIX

Various Weight Initialization algorithms.
-}

module MachineLearning.NeuralNetwork.WeightInitialization
(
  nguyen
  , he
)

where

import qualified Numeric.LinearAlgebra as LA
import qualified Control.Monad.Random as RndM
import MachineLearning.Types (Matrix)
import MachineLearning.Random (getRandomRMatrixM)

-- | Weight Initialization Algorithm discussed in Nguyen at al.: https://web.stanford.edu/class/ee373b/nninitialization.pdf
-- Nguyen, Derrick, Widrow, B. Improving the learning speed of 2-layer neural networks by choosing initial values of adaptive weights.
-- In Proc. IJCNN, 1990; 3: 21-26.
nguyen :: RndM.RandomGen g => (Int, Int) -> RndM.Rand g (Matrix, Matrix)
nguyen (r, c) = do
  let b = LA.konst 0 (1, r)
      eps = (sqrt 6) / (sqrt . fromIntegral $ r + c)
  w <- getRandomRMatrixM r c (-eps, eps)
  return (b, w)

-- | Weight Initialization Algorithm discussed in He at al.: https://arxiv.org/abs/1502.01852
-- Kaiming He, Xiangyu Zhang, Shaoqing Ren, Jian Sun.
-- Delving Deep into Rectifiers: Surpassing Human-Level Performance on ImageNet Classification.
he :: RndM.RandomGen g => (Int, Int) -> RndM.Rand g (Matrix, Matrix)
he (r, c) = do
  let b = LA.konst 0 (1, r)
      eps = sqrt (2/(fromIntegral $ r + c))
  w <- getRandomRMatrixM r c (-eps, eps)
  return (b, w)


-- | The method of moments can be used to estimate a number of commonly used distributions.  This module is still under construction as I work out the best way to handle morphisms from the Moments3 type to types of other distributions.  For more information, see the wikipedia entry: <https://en.wikipedia.org/wiki/Method_of_moments_(statistics)>

module HLearn.Models.Distributions.Univariate.Geometric
    ( Geometric
    )
    where

import Control.DeepSeq
import GHC.TypeLits
import qualified Data.Vector.Unboxed as U
import Data.Vector.Unboxed.Deriving

import qualified Statistics.Distribution as S
import Statistics.Distribution.Geometric

import HLearn.Algebra
import HLearn.Models.Distributions.Common
import HLearn.Models.Distributions.Univariate.Internal.Moments
import HLearn.Models.Distributions.Visualization.Gnuplot

-------------------------------------------------------------------------------
-- Geometric

newtype Geometric prob dp = Geometric {  moments :: (Moments3 dp) }
    deriving (Read,Show,Eq,Ord,Monoid,Group)

instance (Num dp) => HomTrainer (Geometric prob dp) where
    type Datapoint (Geometric prob dp) = dp
    train1dp dp = Geometric $ train1dp dp

instance (Num dp) => Probabilistic (Geometric prob dp) where
    type Probability (Geometric prob dp) = prob

instance (Integral dp, Floating prob) => PDF (Geometric prob dp) where
    pdf dist dp = p*(1-p)^^dp
        where
            p = geo_p dist

instance
    ( PDF (Geometric prob dp)
    , Show prob
    , Show dp
    , Ord dp
    , Ord prob
    , Fractional prob
    , RealFrac prob
    , Integral dp
    ) => PlottableDistribution (Geometric prob dp)
        where

    plotType _ = Points

    samplePoints dist = [min..max]
        where
            min = 0
            max = maximum [20,round $ 3*(fromIntegral $ mean dist)]

geo_p :: (Fractional prob, Integral dp) => Geometric prob dp -> prob
geo_p (Geometric dist) = 1/((fromIntegral $ m1 dist)/(fromIntegral $ m0 dist) +1)

instance (Integral dp, RealFrac prob) => Mean (Geometric prob dp) where
    mean dist = round $ 1/(geo_p dist)

instance (Integral dp, Fractional prob) => Variance (Geometric prob dp) where
    variance dist = (1-p)/(p*p)
        where
            p = geo_p dist


{-# LANGUAGE KindSignatures #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE BangPatterns #-}

module Main where

import NaiveBayes as NB
import StaticTypeNet
import Mnist (readMnistAndShow, readData, dat, label, readMnist, trainData, testData, LabeledData, dataLabel, render)
import Data.Maybe
import System.Environment (getArgs)
import Control.Monad.Random (evalRandIO)
import Text.Read (readMaybe)
import Numeric.LinearAlgebra (Vector, R, toRows, fromList, toList)

main :: IO ()
main = do
  -- readMnistAndShow
  -- now read MNIST and print out the prior distribution
  mnist <- readMnist
  let trainLabeledData :: LabeledData = trainData mnist  -- train set
      testLabeledData :: LabeledData = testData mnist  -- test set
      pri = NB.prior trainLabeledData
      lk = NB.likelihood trainLabeledData
      naiveBayesClassifier :: NB.NaiveBayesClassifier = NB.NBC pri lk  -- create a classifier
      predictor = predict naiveBayesClassifier

  -- show the prior
  putStrLn $ "Below is the prior distribution"
  putStrLn $ show pri  -- show prior values

  let datlab :: [(Vector R, R)] = dataLabel testLabeledData  -- data-label pairs
      predictedAndTruth ld = (predictor (fst ld), floor $ snd ld)
      -- the result of prediction
      predictedResult :: [(Int, Int)] = map predictedAndTruth datlab
      -- calculate the accuracy
      accuracy :: [(Int, Int)] -> Double
      accuracy res = NB.fractionDiv
          (foldl
              (\acc elem -> if fst elem == snd elem
                            then acc + 1
                            else acc)
              0 res)  -- accumulate correct predictions
          (length res)  -- ...divided by total size

  -- show predicted value and truth label
  mapM_ (putStrLn . show . predictedAndTruth) datlab

  putStrLn $ "The accuracy is: "
  putStrLn $ show $ accuracy predictedResult

  -- show the 100th sample
  putStrLn $ render $ fst $ datlab !! 100
  putStrLn $ show $ predictedResult !! 100

  args <- getArgs
  let n = readMaybe =<< (args !!? 0)
      rate = readMaybe =<< (args !!? 1)
  putStrLn "Training network .."
  putStrLn =<< evalRandIO (netTest (fromMaybe 0.25 rate) (fromMaybe 500000 n))