-- | This module implements the Fast Fourier Transformation purely in Haskell.
module Sound.Hommage.FFT
 (
 -- * FFT
   fft
 , fft'
 , fftt

 -- * FFT
 , fftc
 , fftc'

 , ffttv

 , fftco
 , fftco'

 , rect_transform
 , rect_filter
 , rect_filter'

 , i
 , w
 , ws
 , map2
 , zip2
 , evens
 , odds
 , drop_odds
 , appendpair
-- , fitlastlength
 , reorder
 , reorder_init
 , reorder_init'
-- , reorder_overlap
 , butterfly
 , butterfly'
 , fft_level
 , fft_merge
 , fft_init_overlap
 , fft_level'
 , fft_merge'
 , fft_init
 , fft_overlap_loop
 , expandComplex
 , overlap_curve
 , mix_overlap

 )
 where
import Data.Complex
import Sound.Hommage.Misc
-------------------------------------------------------------------------------
-- | The complex value 'i' = @(0 :+ 1)@
i :: Complex Double
i = 0.0 :+ 1.0

-- | The n-th root of 1
w :: Int -> Complex Double
w n = exp (-2 * i * pi / fromIntegral n)

-- | the 2n-th root with exponents 0, 1, .. n. False=inverse (exponents are negated)
ws :: Bool -> Int -> [Complex Double]
ws b n = take n $ map (w1^^) ix
 where
  w1 = exp (-i * pi / fromIntegral n)
  ix | b         = [0, 1 ..]
     | otherwise = [0, (-1) ..]
-------------------------------------------------------------------------------
map2 :: (a -> a -> b) -> [a] -> [b]
map2 f (x:y:r) = f x y : map2 f r
map2 _ _       = []

map2opt :: (a -> a -> b) -> a -> [a] -> [b]
map2opt f a = loop
 where
  loop (x:y:r) = f x y : loop r
  loop [x]     = [f x a]
  loop _       = []

zip2 :: [a] -> [a] -> [a]
zip2 (x:xs) (y:ys) = x : y : zip2 xs ys
zip2 _ _ = []

evens (x:xs) = x : odds xs
evens []     = []

odds (_:xs) = evens xs
odds []     = []

-- | returns [] if argument has zero or one element.
drop_odds (x:y:r) = x : evens r
drop_odds _       = []

appendpair :: ([a], [a]) -> [a]
appendpair (x,y) = x++y

--fitlastlength :: Int -> a -> [[a]] -> [[a]]
--fitlastlength n a = loop
-- where
--  loop (x:xs) = case loop xs of
--                [] -> [take n (x ++ repeat a)]
--                xs -> x : xs
--  loop []     = []
-------------------------------------------------------------------------------
-- | list is grouped into sublists with length N (must be power of 2) and bitwise reverse order
reorder :: Int -> [[a]] -> [[a]]
reorder 0 = id
reorder n = map2 zip2 . reorder (n-1)

reorder_r2 :: Int -> a -> [[a]] -> [[a]]
reorder_r2 n a = loop n
 where
  loop 0 = id
  loop n = map2opt zip2 (repeat a) . loop (n-1)

reorder_init (x1:x2:x3:x4:xs) = [x1 :+ x2, x3 :+ x4] : reorder_init xs
reorder_init [x1,x2,x3]       = [[x1 :+ x2, x3 :+ 0]]
reorder_init [x1,x2]          = [[x1 :+ x2, 0]]
reorder_init [x1]             = [[x1 :+ 0, 0]]
reorder_init _                = []

reorder_init' :: [a] -> [[a]]
reorder_init' (x:y:r) = [x,y] : reorder_init' r
reorder_init' _       = []

reorder_init_r2' :: a -> [a] -> [[a]]
reorder_init_r2' a = loop
 where
  loop (x:y:r) = [x,y] : loop r
  loop [x]     = [[x,a]]
  loop _       = []

--reorder_overlap :: [[a]] -> [a]
--reorder_overlap (xs:ys:r) = (zip2 xs ys) ++ reorder_overlap (ys:r)
--reorder_overlap _       = []
-------------------------------------------------------------------------------
butterfly :: Complex Double ->  Complex Double -> Complex Double -> (Complex Double, Complex Double)
butterfly (wr :+ wi) (xr :+ xi) (yr :+ yi) = (a, b)
 where
  wyr = wr*yr-wi*yi
  wyi = wr*yi+yr*wi
  a = ( (xr + wyr) * 0.5 :+ (xi + wyi) * 0.5 )
  b = ( (xr - wyr) * 0.5 :+ (xi - wyi) * 0.5 )

butterfly' :: Complex Double ->  Complex Double -> Complex Double -> (Complex Double, Complex Double)
butterfly' (wr :+ wi) (xr :+ xi) (yr :+ yi) = (a, b)
 where
  wyr = wr*yr-wi*yi
  wyi = wr*yi+yr*wi
  a = ( (xr + wyr) :+ (xi + wyi) )
  b = ( (xr - wyr) :+ (xi - wyi) )
-------------------------------------------------------------------------------
-- | FFT transformation. Input is grouped into overlapping parts of 2^(N+2) reals and mapped to sublists
--   with 2^(N+1) complex numbers.
fft :: Int -> [Double] -> [[Complex Double]]
fft n = fft_level (ws True (2^n)) . fft_init_overlap . reorder_r2 (n-1) 0 . reorder_init

fft_level :: [Complex Double] -> [[Complex Double]] -> [[Complex Double]]
fft_level [_] = id
fft_level ws  = map2 (\x -> appendpair . fft_merge ws x) . fft_level (drop_odds ws)

fft_merge :: [Complex Double] -> [Complex Double] -> [Complex Double] -> ([Complex Double], [Complex Double])
fft_merge (w:ws) (x:xs) (y:ys) = let (xs', ys') = fft_merge ws xs ys
                                     (x', y') = butterfly w x y
                                 in (x' : xs', y' : ys')
fft_merge _ _ _ = ([], [])

fft_init_overlap (x:y:r) = (zipWith (\x y -> [(x + y)*0.5, (x - y)*0.5]) x y) ++ fft_init_overlap (y:r)
fft_init_overlap _       = []

fft_init (x:y:r) = (zipWith (\x y -> [(x + y)*0.5, (x - y)*0.5]) x y) ++ fft_init r
fft_init _       = []
-------------------------------------------------------------------------------
-- | Inverse FFT transformation. Complex input is grouped into parts with length 2^(N+1) and mapped to
--   sublists with 2^(N+2) reals, which are overlapped and mixed.
fft' :: Int -> [Complex Double] -> [Double]
fft' n = mix_overlap (2^(n+2)) . map expandComplex . fft_level' (ws False (2^n)) . fft_init' . reorder (n-1) . reorder_init'

fft_level' :: [Complex Double] -> [[Complex Double]] -> [[Complex Double]]
fft_level' [_] = id
fft_level' ws  = map2 (\x -> appendpair . fft_merge' ws x) . fft_level' (drop_odds ws)

fft_merge' :: [Complex Double] -> [Complex Double] -> [Complex Double] -> ([Complex Double], [Complex Double])
fft_merge' (w:ws) (x:xs) (y:ys) = let (xs', ys') = fft_merge' ws xs ys
                                      (x', y') = butterfly' w x y
                                  in (x' : xs', y' : ys')
fft_merge' _ _ _ = ([], [])

fft_init' (x:y:r) = (zipWith (\x y -> [x + y, x - y]) x y) ++ fft_init' r
fft_init' _       = []

fft_overlap_loop (x:r@(y:_)) = (concat $ zipWith (\x y -> [x + y, x - y]) x y) : fft_overlap_loop r
fft_overlap_loop _           = []

expandComplex ((x:+y):r) = x : y : expandComplex r
expandComplex _ = []
-------------------------------------------------------------------------------
-- | Creates a fade with length N
overlap_curve :: Int -> [Double]
overlap_curve n = take n $ map (\k -> 0.5 - 0.5 * cos (fromIntegral k * 2 * pi / fromIntegral n )) [0..]

-- | Overlaps a sequence of parts of length N (overlaps by N\/2).
mix_overlap :: Int -> [[Double]] -> [Double]
mix_overlap n xs = merge (+) l r
  where
  c = cycle $ overlap_curve n
  l = zipWith (*) c $ concat $ evens xs
  r = replicate (div n 2) 0.0 ++ (zipWith (*) c $ concat $ odds xs)
-------------------------------------------------------------------------------
-- | FFT transformation with overlapping and windowing.
--   Argument function maps the coefficients. Windowsize: 2^(N+2), i. e. 2^(N+1) coefficients
fftt :: Int -> ([[Complex Double]] -> [Complex Double]) -> [Double] -> [Double]
fftt n f =
 mix_overlap (2^(n+2)) .
 map expandComplex .
 fft_level' (ws False (2^n)) .
 fft_init' .
 reorder (n-1) .
 reorder_init' .
 f .
 fft_level (ws True (2^n)) .
 fft_init_overlap .
 reorder_r2 (n-1) 0 .
 reorder_init

ffttv :: Int -> [Double] -> [Double] -> [Double]
ffttv n vs = fftt n (zipWith (\v (x:+y) -> (v*x) :+ (v*y)) vs . concat . map (\((a:+_):r) -> (a:+0):r))
-------------------------------------------------------------------------------
-- | FFT transformation for complex input (segments of length 2^n).
--   No overlapping or windowing.
fftc :: Int -> [Complex Double] -> [[Complex Double]]
fftc n = fft_level (ws True (2^n)) . fft_init . reorder_r2 (n-1) 0 . reorder_init_r2' 0

-- | Inverse FFT transformation for complex input (segments of length 2^n).
--   No overlapping or windowing.
fftc' :: Int -> [Complex Double] -> [[Complex Double]]
fftc' n = fft_level' (ws False (2^n)) . fft_init' . reorder (n-1) . reorder_init'
-------------------------------------------------------------------------------
-- | FFT for complex input with overlapping. Segment-size: 2^(n+1)
fftco :: Int -> [Complex Double] -> [[Complex Double]]
fftco n = fft_level (ws True (2^n)) . fft_init_overlap . reorder_r2 (n-1) 0 . reorder_init_r2' 0

fftco' :: Int -> [Complex Double] -> [Complex Double]
fftco' n = mix_overlapc (2^(n+1)) . fft_level' (ws False (2^n)) . fft_init' . reorder (n-1) . reorder_init'

-- | Overlaps a sequence of parts of length N (overlaps by N\/2).
mix_overlapc :: Int -> [[Complex Double]] -> [Complex Double]
mix_overlapc n xs = merge (+) l r
  where
  c = cycle $ map (:+0) $ overlap_curve n
  l = zipWith (*) c $ concat $ evens xs
  r = replicate (div n 2) 0.0 ++ (zipWith (*) c $ concat $ odds xs)
-------------------------------------------------------------------------------

-- | A self-inverse transformation similar to FFT but with a simple butterfly
--   operation that uses always W=1. Modifying the data between application
--   and inverse is similar to filtering but the result will be built up from
--   rectangle waves instead of sinus waves.
rect_transform :: Floating a => Int -> [a] -> [[a]]
rect_transform n = run n . map (:[])
 where
  sq = 1 / sqrt 2
  add x y = (x+y) * sq
  sub x y = (x-y) * sq
  run 0 = id
  run n = loop . run (n-1)
  loop (x:y:r) = iter True x y : loop r
  loop [x]     = [iter True x (repeat 0)]
  loop _       = []
  iter True  (x:xs) (y:ys) = add x y : sub x y : iter False xs ys
  iter False (x:xs) (y:ys) = sub x y : add x y : iter True  xs ys
  iter _     _      _      = []

rect_filter :: Floating a => Int -> ([[a]] -> [[a]]) -> [a] -> [a]
rect_filter n f =  concat . rect_transform n . concat . f . rect_transform n

rect_filter' :: Floating a => Int -> [[a]] -> [a] -> [a]
rect_filter' n vs =  concat . rect_transform n . concat . zipWith f vs . rect_transform n
 where
  f v c = zipWith (*) (v ++ repeat 0) c


import Numeric.LinearAlgebra

and' x1 x2 =
    let w1 = 0.5; w2 = 0.5
        theta = 0.7
        tmp = x1*w1 + x2*w2
    in if tmp > theta then 1 else 0

and'' :: Num t => Vector R -> t
and'' x =
    let w = vector [0.5, 0.5]
        b = (-0.7)
        tmp = (x <.> w) + b
    in if tmp > 0 then 1 else 0

nand :: Num t => Vector R -> t
nand x =
    let w = vector [(-0.5), (-0.5)]
        b = 0.7
        tmp = (x <.> w) + b
    in if tmp > 0 then 1 else 0

or' :: Num t => Vector R -> t
or' x =
    let w = vector [0.5, 0.5]
        b = (-0.2)
        tmp = (x <.> w) + b
    in if tmp > 0 then 1 else 0

xor :: Num t => Vector R -> t
xor x =
    let s1 = nand x
        s2 = or' x
    in and'' $ vector [s1, s2]


-- file:Astro/AccretionDisk.hs
-- Numerical tools for accretion disk analysis

module Astro.AccretionDisk
	where

import Control.Monad
import System.IO.Unsafe
import System.Random
import Numeric.LinearAlgebra

import Astro.Misc

seed = 1337

rndDoubles :: Int -> (Double,Double) -> IO [Double]
rndDoubles n (min,max) = replicateM n $ randomRIO (min,max)

rndDouble :: (Double, Double) -> IO Double
rndDouble (min, max) = randomRIO (min,max) :: IO Double

-- Random position and corresponding velocity in a flat xy-plane Keplerian disk
-- with CCW rotation when viewed towards negative z-axis
rndDiskPosVel :: (Double, Double) -> Double -> IO (Vector Double, Vector Double)
rndDiskPosVel (minR,maxR) gpar =
	do
	r <- rndDouble (minR,maxR)
	th <- rndDouble (0, 2*pi)
	let rv = fromList [ r * cos th, r * sin th, 0]
	    vv = fromList [-rv@>1, rv@>0, 0.0]
	    vv' = scale (sqrt gpar/norm rv) vv
	return (rv, vv')

-- Create a flat disk (planar), rotating in a CCW sense when viewed
-- towards negative z-axis
flatDisk :: (Double, Double) -> Double -> Int
		-> [(Vector Double, Vector Double)]
flatDisk (minR, maxR) gpar nparts = unsafePerformIO $
	do
		setStdGen $ mkStdGen seed
		replicateM nparts $ rndDiskPosVel (minR,maxR) gpar


{-# LANGUAGE ViewPatterns #-}
import Numeric.GSL.ODE
import Numeric.LinearAlgebra
import Graphics.Plot
import Debug.Trace(trace)
debug x = trace (show x) x

vanderpol mu = do
    let xdot mu t [x,v] = [v, -x + mu * v * (1-x^2)]
        ts = linspace 1000 (0,50)
        sol = toColumns $ odeSolve (xdot mu) [1,0] ts
    mplot (ts : sol)
    mplot sol

harmonic w d = do
    let xdot w d t [x,v] = [v, a*x + b*v] where a = -w^2; b = -2*d*w
        ts = linspace 100 (0,20)
        sol = odeSolve (xdot w d) [1,0] ts
    mplot (ts : toColumns sol)

kepler v a = mplot (take 2 $ toColumns sol) where
    xdot t [x,y,vx,vy] = [vx,vy,x*k,y*k]
        where g=1
              k=(-g)*(x*x+y*y)**(-1.5)
    ts = linspace 100 (0,30)
    sol = odeSolve xdot [4, 0, v * cos (a*degree), v * sin (a*degree)] ts
    degree = pi/180

main = do
    vanderpol 2
    harmonic 1 0
    harmonic 1 0.1
    kepler 0.3 60
    kepler 0.4 70
    vanderpol' 2

-- example of odeSolveV with jacobian
vanderpol' mu = do
    let xdot mu t (toList->[x,v]) = fromList [v, -x + mu * v * (1-x^2)]
        jac t (toList->[x,v]) = (2><2) [ 0          ,          1
                                       , -1-2*x*v*mu, mu*(1-x**2) ]
        ts = linspace 1000 (0,50)
        hi = (ts@>1 - ts@>0)/100
        sol = toColumns $ odeSolveV (MSBDF jac) hi 1E-8 1E-8 (xdot mu) (fromList [1,0]) ts
    mplot sol


module Nnaskell ( NN(..)
                , randomNN
                , activateNN
                , cost
                , optimizeCost
                , loadNN
                , saveNN
                ) where

import Data.Function
import Data.List
import Control.Monad
import System.Random
import Numeric.LinearAlgebra.Data
import Numeric.LinearAlgebra.HMatrix

data NN = NN { nnWs :: [Matrix R]
             , nnBs :: [Vector R]
             } deriving (Show, Read)

randomVec :: Int -> IO (Vector R)
randomVec n = vector <$> (replicateM n $ randomRIO (-1.0, 1.0))

randomMatrix :: Int -> Int -> IO (Matrix R)
randomMatrix n m = matrix m <$> (replicateM ((*) n m) $ randomRIO (-1.0, 1.0))

randomLayer :: Int -> Int -> IO (Matrix R, Vector R)
randomLayer n m = liftM2 (\ws bs -> (ws, bs)) (randomMatrix m n) (randomVec m)

randomNN :: [Int] -> IO NN
randomNN ns = do (ws, bs) <- unzip <$> (sequence $ map (uncurry randomLayer) $ pairScan ns)
                 return $ NN { nnWs = ws
                             , nnBs = bs
                             }

sigmoid :: R -> R
sigmoid t = 1 / (1 + exp (-1 * t))

sigmoidVec :: Vector R -> Vector R
sigmoidVec = fromList . map sigmoid . toList

activateLayer :: Vector R -> (Matrix R, Vector R) -> Vector R
activateLayer as (ws, bs) = sigmoidVec (ws #> as + bs)

activateNN :: NN -> Vector R -> Vector R
activateNN nn as = foldl activateLayer as $ zip ws bs
    where ws = nnWs nn
          bs = nnBs nn

pairScan :: [a] -> [(a, a)]
pairScan (x:y:rest) = (x, y) : pairScan (y:rest)
pairScan _ = []

cost :: [(Vector R, Vector R)] -> NN -> R
cost td nn = (sum $ map inputCost td) / n
    where inputCost (input, output) = sumElements $ cmap (** 2) (activateNN nn input - output)
          n = fromIntegral $ length td

stepBias :: NN -> Int -> R -> NN
stepBias nn updateIdx s = nn { nnBs = map (\(v, idx) -> if idx <= updateIdx && updateIdx < idx + size v
                                                        then accum v (+) [(updateIdx - idx, s)]
                                                        else v)
                                      $ zip bs
                                      $ scanl (\idx v -> idx + size v) 0 bs
                             }
    where bs = nnBs nn

stepWeight :: NN -> Int -> R -> NN
stepWeight nn updateIdx s = nn { nnWs = map (\(mx, idx) -> let (n, m) = size mx
                                                               effIdx = updateIdx - idx
                                                           in if idx <= updateIdx && updateIdx < idx + (n * m)
                                                              then accum mx (+) [((effIdx `div` m, effIdx `mod` m), s)]
                                                              else mx)
                                        $ zip ws
                                        $ scanl (\idx m -> idx + uncurry (*) (size m)) 0 ws
                               }
    where ws = nnWs nn

class Domain d where
    countArgs :: d -> Int
    stepArg :: d -> Int -> R -> d

instance Domain NN where
    countArgs nn = bsCount + wsCount
        where bsCount = sum $ map (length . toList) $ nnBs nn
              wsCount = sum $ map length $ concatMap toLists $ nnWs nn
    stepArg nn idx s = if idx < bsCount
                       then stepBias nn idx s
                       else stepWeight nn (idx - bsCount) s
        where bsCount = sum $ map (length . toList) $ nnBs nn

optimizeCost :: Domain d => (d -> R) -> d -> [d]
optimizeCost cost d = scanl optimizeStep d $ cycle [0 .. n - 1]
    where n = countArgs d
          optimizeStep d idx = minimumBy (compare `on` cost) $ map (stepArg d idx) [-step, 0, step]
          step = 0.1

loadNN :: FilePath -> IO NN
loadNN filePath = read <$> readFile filePath

saveNN :: FilePath -> NN -> IO ()
saveNN filePath nn = writeFile filePath $ show nn


{- |
Description :  Some handy (mathematical, string manipulation, and so on) tools
Copyright   :  (c) Dominik Schrempf 2017
License     :  GPLv3

Maintainer  :  dominik.schrempf@gmail.com
Stability   :  unstable
Portability :  non-portable (not tested)

Please see function definitions and documentations.

* Changelog

-}

module Tools
  ( allValues
  , harmonic
  , matrixSeparateSymSkew
  , matrixSetDiagToZero
  , nearlyEq
  , left
  , right
  ) where

import Numeric.LinearAlgebra

-- | Get all values of a bounded enumerated type.
allValues :: (Bounded a, Enum a) => [a]
allValues = [minBound..]

-- | Calculate the nth harmonic number.
harmonic :: Int -> Double
harmonic 1 = 1.0
harmonic n = 1.0 / fromIntegral n + harmonic (n-1)

-- | Separate a square matrix into a symmetric and a skew-symmetric matrix.
matrixSeparateSymSkew :: Matrix R -> (Matrix R, Matrix R)
matrixSeparateSymSkew m = (mSym, mSkew)
  where trM = tr m
        mSym  = scale 0.5 $ m + trM
        mSkew = scale 0.5 $ m - trM

-- | Set the diagonal entries of a matrix to zero.
matrixSetDiagToZero :: Matrix R -> Matrix R
matrixSetDiagToZero m = m - diag (takeDiag m)

-- | Test for equality with tolerance (needed because of machine precision).
nearlyEq :: Double -> Double -> Double -> Bool
nearlyEq tol a b = tol > abs (a-b)

-- Functions that fill a string 's' to a given width 'n' by adding a pad
-- character 'c' (c) to align right.
fillDiff :: Int -> String -> Int
fillDiff width entry =
  if l >= width then 0 else width - l
  where l = length entry

fillLeft :: Char -> Int -> String -> String
fillLeft c n s = s ++ replicate (fillDiff n s) c

fillRight :: Char -> Int -> String -> String
fillRight c n s = replicate (fillDiff n s) c ++ s

-- | Fill a string to a given width by adding spaces. Align left.
left :: Int -> String -> String
left = fillLeft ' '

-- | Fill a string to a given width by adding spaces. Align right.
right :: Int -> String -> String
right = fillRight ' '


module Lib.Utility.Differentiation (
    numericalGradient
 ) where

import Numeric.LinearAlgebra
import Numeric.LinearAlgebra.Devel

h = 1e-4

type F = X -> Double
type X = Matrix Double
numericalGradient :: F -> X -> X
numericalGradient f xs = runSTMatrix $ do
    xs' <- thawMatrix xs
    let (rows, columns) = size xs
    let grad = newMatrix 0 rows columns
    let allIndex = [(r,c) | r <- [0..rows-1], c <- [0..columns-1]]
    foldr (\(r, c) grad -> do
        x <- readMatrix xs' r c
        grad' <- grad
        writeMatrix xs' r c $ x+h
        xs'' <- freezeMatrix xs'
        let fxHPlus = f xs''
        writeMatrix xs' r c $ x-h
        xs'' <- freezeMatrix xs'
        let fxHMinus = f xs''
        let delta = (fxHPlus - fxHMinus) / (2 * h)
        writeMatrix xs' r c x
        writeMatrix grad' r c delta
        return grad'
        ) grad allIndex


module Main where

import Numeric.LinearAlgebra
import Common
import Forward
import BackProp
import AutoEncoder
import ActivFunc
import Other

main :: IO ()
main = do
  regression
  classification

regression :: IO ()
regression = do
  ws <- genWeights [2, 4, 8, 1]
  let x = matrix 4 [0, 0, 1, 1,
                    0, 1, 0, 1]
  let y = matrix 4 [0, 1, 1, 0]
  let i = matrix 4 [0, 0, 1, 1,
                    0, 1, 0, 1] -- example
  -- let nws = last . take 500 $ iterate (backPropRegression sigmoids (x, y)) ws
  nws <- last . take 500 $ iterateM (sgdMethod 2 (x, y) $ backPropRegression 0.1 sigmoids) ws
  let pws = preTrains 0.1 500 sigmoids x ws
  npws <- last . take 500 $ iterateM (sgdMethod 2 (x, y) $ backPropRegression 0.1 sigmoids) pws
  putStrLn "training inputs"
  print x
  putStrLn "training outputs"
  print y
  putStrLn "inputs"
  print i
  putStrLn "not trained outputs"
  print $ forwardRegression sigmoidC ws i
  putStrLn "trainined outputs"
  print $ forwardRegression sigmoidC nws i
  putStrLn "pretrainined outputs"
  print $ forwardRegression sigmoidC npws i

classification :: IO ()
classification = do
  ws <- genWeights [2, 4, 8, 3]
  let x = matrix 4 [0, 0, 1, 1,
                    0, 1, 0, 1]
  let y = matrix 4 [1, 0, 0, 0,
                    0, 1, 1, 0,
                    0, 0, 0, 1]
  let i = matrix 4 [0, 0, 1, 1,
                    0, 1, 0, 1] -- example
  -- let nws = last . take 500 $ iterate (backPropClassification sigmoids (x, y)) ws
  nws <- last . take 500 $ iterateM (sgdMethod 2 (x, y) $ backPropClassification 0.1 sigmoids) ws
  let pws = preTrains 0.1 500 sigmoids x ws
  npws <- last . take 500 $ iterateM (sgdMethod 2 (x, y) $ backPropClassification 0.1 sigmoids) pws
  putStrLn "training inputs"
  print x
  putStrLn "training outputs"
  print y
  putStrLn "inputs"
  print i
  putStrLn "not trained outputs"
  print $ forwardClassification sigmoidC ws i
  putStrLn "trainined outputs"
  print $ forwardClassification sigmoidC nws i
  putStrLn "pretrainined outputs"
  print $ forwardClassification sigmoidC npws i


module Perceptron where

import Numeric.LinearAlgebra

v00 = vector [0.0,0.0]
v01 = vector [0.0,1.0]
v10 = vector [1.0,0.0]
v11 = vector [1.0,1.0]

p_and :: Vector Double -> Double
p_and x
  | tmp <= 0.0 = 0.0
  | otherwise = 1.0
  where
    tmp = w <.> x + b
    b = -0.7
    w = vector [0.5,0.5]

p_nand :: Vector Double -> Double
p_nand x
  | tmp <= 0.0 = 0.0
  | otherwise = 1.0
  where
    tmp = w <.> x + b
    b = 0.7
    w = vector [-0.5,-0.5]

p_or :: Vector Double -> Double
p_or x
  | tmp <= 0.0 = 0.0
  | otherwise = 1.0
  where
    tmp = w <.> x + b
    b = -0.2
    w = vector [0.5,0.5]

p_xor :: Vector Double -> Double
p_xor x = p_and $ vector [s1,s2]
  where
    s1 = p_nand x
    s2 = p_or x


{-# LANGUAGE CPP #-}
-- |
-- Module    : Numeric.SpecFunctions
-- Copyright : (c) 2009, 2011, 2012 Bryan O'Sullivan
-- License   : BSD3
--
-- Maintainer  : bos@serpentine.com
-- Stability   : experimental
-- Portability : portable
--
-- Special functions and factorials.
module Numeric.SpecFunctions (
    -- * Error function
    erf
  , erfc
  , invErf
  , invErfc
    -- * Gamma function
  , logGamma
  , logGammaL
  , incompleteGamma
  , invIncompleteGamma
  , digamma
    -- * Beta function
  , logBeta
  , incompleteBeta
  , incompleteBeta_
  , invIncompleteBeta
    -- * Sinc
  , sinc
    -- * Logarithm
    -- $log1p
  , log1p
  , log1pmx
  , log2
    -- * Exponent
  , expm1
    -- * Factorial
  , factorial
  , logFactorial
  , stirlingError
    -- * Combinatorics
  , choose
  , logChoose
    -- * References
    -- $references
  ) where

import Numeric.SpecFunctions.Internal

-- $log1p
--
-- Base starting from @4.9.0@ (GHC 8.0) provides 'log1p' and 'expm1'
-- as method of class 'Floating'. In this case we simply reexport
-- these function. Otherwise we provide our own with more restrictive
-- signature @Double → Double@.

-- $references
--
-- * Bernardo, J. (1976) Algorithm AS 103: Psi (digamma)
--   function. /Journal of the Royal Statistical Society. Series C
--   (Applied Statistics)/ 25(3):315-317.
--   <http://www.jstor.org/stable/2347257>
--
-- * Cran, G.W., Martin, K.J., Thomas, G.E. (1977) Remark AS R19
--   and Algorithm AS 109: A Remark on Algorithms: AS 63: The
--   Incomplete Beta Integral AS 64: Inverse of the Incomplete Beta
--   Function Ratio. /Journal of the Royal Statistical Society. Series
--   C (Applied Statistics)/ Vol. 26, No. 1 (1977), pp. 111-114
--   <http://www.jstor.org/pss/2346887>
--
-- * Lanczos, C. (1964) A precision approximation of the gamma
--   function.  /SIAM Journal on Numerical Analysis B/
--   1:86–96. <http://www.jstor.org/stable/2949767>
--
-- * Loader, C. (2000) Fast and Accurate Computation of Binomial
--   Probabilities. <http://projects.scipy.org/scipy/raw-attachment/ticket/620/loader2000Fast.pdf>
--
-- * Macleod, A.J. (1989) Algorithm AS 245: A robust and reliable
--   algorithm for the logarithm of the gamma function.
--   /Journal of the Royal Statistical Society, Series C (Applied Statistics)/
--   38(2):397–402. <http://www.jstor.org/stable/2348078>
--
-- * Majumder, K.L., Bhattacharjee, G.P. (1973) Algorithm AS 63: The
--   Incomplete Beta Integral. /Journal of the Royal Statistical
--   Society. Series C (Applied Statistics)/ Vol. 22, No. 3 (1973),
--   pp. 409-411. <http://www.jstor.org/pss/2346797>
--
-- * Majumder, K.L., Bhattacharjee, G.P. (1973) Algorithm AS 64:
--   Inverse of the Incomplete Beta Function Ratio. /Journal of the
--   Royal Statistical Society. Series C (Applied Statistics)/
--   Vol. 22, No. 3 (1973), pp. 411-414
--   <http://www.jstor.org/pss/2346798>
--
-- * Temme, N.M. (1992) Asymptotic inversion of the incomplete beta
--   function. /Journal of Computational and Applied Mathematics
--   41(1992) 145-157.
--
-- * Temme, N.M. (1994) A set of algorithms for the incomplete gamma
--   functions. /Probability in the Engineering and Informational
--   Sciences/, 8, 1994, 291-307. Printed in the U.S.A.


module Main where

import           BuckinghamPi
import           FileOperations
import           Numeric.LinearAlgebra.HMatrix

rho = DimensionalVariable "rho" [(Mass   , 1.0)  , (Length , -3.0)]
vel = DimensionalVariable "v"   [(Length , 1.0)  , (Time   , -1.0)]
l   = DimensionalVariable "l"   [(Length , 1.0)]
mu  = DimensionalVariable "mu"  [(Mass   , 1.0)  , (Length , -1.0)  , (Time , -1.0)]
a   = DimensionalVariable "a"   [(Length , 1.0)  , (Time   , -1.0)]

dimVars = [rho, vel, l, mu, a]
funUnits = [Mass, Length, Time]

--writeDimensionlessGroups :: [DimensionlessVariable] -> FilePath -> IO ()
main :: IO ()
main = fmap (\d -> generatePiGroups d [0..(numFundamentalUnits dimVars - 1)]) (readVariablesFromFile "./app/dimensional_variables") >>= (\v -> writeDimensionlessGroups v "dimensionless_groups")


{-|
Module      : Marvin.API.Algorithms.Internal.GradientDescent
Description : Optimization algorithm.
-}
module Marvin.API.Algorithms.Internal.GradientDescent where

import Marvin.API.Table.Internal as Table
import Marvin.API.Fallible

import Numeric.LinearAlgebra.HMatrix
import Numeric.LinearAlgebra.Data as LA

import qualified Data.Vector as Vec
import qualified Data.Vector.Unboxed as UVec

import Control.Arrow
import Control.Monad.Except
import Data.List (intercalate)

-- | Resulting model
data LinearModel = LM {
  featureColumnNames :: [Fallible ColumnName]
  , targetVariableName :: Fallible ColumnName
  , coefs :: LA.Vector Double
  , intercept_ :: Double
} deriving Eq

instance Show LinearModel where
  show model = intercalate " +\n" $ map ((++) "\t") $ show (intercept_ model) : columnStrings
    where
      columnStrings = map (\(idx, (coef, fallibleName)) ->
        let name = case fallibleName of
              Right n -> n
              Left _ -> "column #" ++ show idx
        in show coef ++ " * " ++ name) $
        zip [0..] $ zip (LA.toList (coefs model)) (featureColumnNames model)

fit' :: GradientDescent -> (NumericTable, NumericColumn) -> LinearModel
fit' gradDesc (nt, nc) = LM {
    featureColumnNames = columnNames nt
    , targetVariableName = columnName nc
    , intercept_ = intercept'
    , coefs = theta'
  }
  where
    featureMtx = tableToMatrix nt
    targetVar = columnToVector nc
    modelParams = gradientDescent gradDesc featureMtx targetVar
    theta = modelParams
    (intercept', theta') = if (addIntercept gradDesc)
      then (theta ! 0, LA.subVector 1 (LA.size theta - 1) theta)
      else (0, theta)

predict' :: LinearModel -> NumericTable -> Fallible NumericColumn
predict' model test = do
      dimCheck
      return $ unsafePredict model test
    where
      modelDim = LA.size $ coefs model
      testDim = numberOfColumns test
      dimCheck = if modelDim == testDim
          then return ()
          else throwError $
            RowLengthMismatch $
              "When attempting to make predictions based on a linear model. " ++
              "coefficients: " ++ (show modelDim) ++
              ", columns: " ++ (show testDim) ++ "."

-- | Predicts without checks.
unsafePredict :: LinearModel -> NumericTable -> NumericColumn
unsafePredict model nt = pred
  where
    pred = unsafeFromVector $ fromLinAlgVector $ batchPredict theta c x
    x = tableToMatrix nt
    theta = coefs model
    c = intercept_ model

-- | Predicts with hmatrix types.
batchPredict :: LA.Vector R -> R -> FeatureMatrix -> TargetVariable
batchPredict theta intercept x = cmap (+ intercept) (x #> theta)

-- * Type aliases
type TargetVariable   = Vector R
type FeatureVector    = Vector R

type FeatureMatrix    = Matrix R
type ModelParameters  = Vector R

type CostFunction = FeatureMatrix -> TargetVariable -> ModelParameters -> (R, Vector R)
type Regularization = R -> ModelParameters -> (R, Vector R)

-- * Parameters
data GradientDescent = GradientDescent {
  learningRate :: Double
  , addIntercept :: Bool
  , lambda :: Double
  , cost :: CostFunction
  , initModelParams :: FeatureMatrix -> TargetVariable -> ModelParameters
  , numIter :: Int
}

defaultGradientDescent = GradientDescent {
  learningRate = 0.1
  , initModelParams = \x y -> cols x |> [1,1..]
  , numIter = 10
  , lambda = 0.1
  , addIntercept = True
  , cost = linearRegressionCost
}
-- * Cost functions

logisticRegressionCost :: FeatureMatrix -> TargetVariable -> ModelParameters
                                                     -> (R, Vector R)
logisticRegressionCost x@features y@targetVariable theta@modelParameters =
  (cost, gradient)
  where
    cost            = (1 / m) * sumElements (-y * log h - (1-y) * log(1-h))
    gradient        = scale (1 / m) $ tr x #> error
    m               = fromIntegral (rows x) :: R
    error           = h - y
    h@hypothesis    = cmap sigmoid $ x #> theta

linearRegressionCost :: FeatureMatrix -> TargetVariable -> ModelParameters -> (R, Vector R)
linearRegressionCost x@features y@targetVariable theta@modelParameters =
  (cost, gradient)
  where
    cost            = (1 / (2 * m)) * sumElements (error ^ 2)
    gradient        = scale (1 / m) $ tr x #> error
    m               = fromIntegral (rows x) :: R
    error           = h - y
    h@hypothesis    = x #> theta

addRegularization :: R -> CostFunction -> Regularization -> CostFunction
addRegularization lambda costFunction regularization x y theta =
  (cost + regCost, gradient + regGradient)
  where
    (cost, gradient) = costFunction x y theta
    (regCost, regGradient) = regularization lambda theta

l2 :: R -> ModelParameters -> (R, Vector R)
l2 lambda@regularizationCoefficient theta@modelParameters =
  (cost, gradient)
  where
    cost = lambda * (norm_2 modelParameters)^2
    gradient = scale (2 * lambda) modelParameters

-- * Algorithm

gradientDescent :: GradientDescent -> FeatureMatrix -> TargetVariable ->
  ModelParameters
gradientDescent params featureMtx targetVariable =
  runIteration (numIter params) modelUpdater initModel
    where
      features = if addIntercept params
        then addInterceptFeature featureMtx
        else featureMtx
      lambda' = lambda params
      cost' = cost params
      costFunction = if lambda' /= 0
        then addRegularization lambda' cost' l2
        else cost'
      alpha                       = learningRate params
      initModel                   = initModelParams params features targetVariable
      gradient model              = snd $ costFunction features targetVariable model
      modelUpdater theta          = theta - scale alpha (gradient theta)

-- * Helpers

addInterceptFeature :: FeatureMatrix -> FeatureMatrix
addInterceptFeature features = LA.fromColumns $ constantFeature:toColumns x
  where
    x                   = features
    (m,_)               = size x
    constantFeature     = m |> repeat 1

sigmoid :: R -> R
sigmoid z = 1 / (1 + exp(-z))

runIteration :: Int -> (a -> a) -> a -> a
runIteration n f init = iterate f init !! n

-- * Conversions

columnToVector :: NumericColumn -> LA.Vector R
columnToVector = LA.fromList . UVec.toList . values

tableToMatrix :: NumericTable -> LA.Matrix R
tableToMatrix =
  columnsVec >>>
  Vec.toList >>>
  map columnToVector >>>
  LA.fromColumns

fromLinAlgVector :: LA.Vector R -> UVec.Vector R
fromLinAlgVector = LA.toList >>> UVec.fromList

toLinAlgVector :: Vec.Vector R -> LA.Vector R
toLinAlgVector = Vec.toList >>> LA.fromList


-- Gaussian primes
-- http://www.codewars.com/kata/54d6abf84a35017d30000b26/

module Data.Complex.Gaussian.Prime where

import Control.Arrow ((&&&))
import Data.Complex.Gaussian (Gaussian (..), norm)

isGaussianPrime :: Gaussian -> Bool
isGaussianPrime g@(Gaussian a b) | a == 0 = f b
                                 | b == 0 = f a
                                 | otherwise = isPrime . norm $ g
                                 where f = uncurry (&&) . ((==3) . (`mod` 4) &&& isPrime) . abs
                                       isPrime n | n `elem` [2, 3, 5] = True
                                                 | n `mod` 2 == 0 || n `mod` 3 == 0 = False
                                                 | otherwise = notElem 0 . map (n `mod`) $ [5, 7 .. floor . sqrt . fromIntegral $ n]


-----------------------------------------------------------------------------
{- |
This module defines the monad of sampling functions. See Park, Pfenning and Thrun:
A probabilistic language based upon sampling functions, Principles of programming languages 2005

Sampling functions allow the composition of both discrete and continuous
probability distributions.

The implementation and interface are similar to those in the random-fu, monte-carlo
and monad-mersenne-random packages.

Example -- a biased coin:

@
data Throw = Head | Tail

throw bias = do
  b <- bernoulli bias
  return $ if b then Head else Tail

tenThrowsCrooked = replicateM 10 $ throw 0.3

countHeads = do
   throws <- tenThrowsCrooked
   return $ length [ () | Head <- throws]

main = do
 print =<< sampleIO tenThrowsCrooked
 print =<< eval ((\<4) \`fmap\` countHeads)
@

-}

{-# LANGUAGE BangPatterns, ScopedTypeVariables #-}

module Math.Probably.Sampler where

import Control.Monad
import Control.Applicative
import qualified Math.Probably.PDF as PDF
import Numeric.LinearAlgebra hiding (find)
import System.Random.Mersenne.Pure64
import System.Environment
import Data.List
import Data.Maybe
import Data.Ord
import Control.Spoon

--import Debug.Trace

type Seed = PureMT

data Sampler a = Sam {unSam :: Seed -> (a, Seed) }
               | Samples [a]

instance Functor Sampler where
    fmap f (Sam sf) = Sam $ \rs -> let (x,rs') = sf rs in
                                   (f x, rs')
    fmap f (Samples xs) = Samples $ map f xs

instance Applicative Sampler where
    pure x = Sam (\rs-> (x, rs))
    (Sam sff) <*> (Sam sfx) = Sam $ \rs-> let (f ,rs') = sff rs
                                              (x, rs'') = sfx rs' in
                                          (f x, rs'')

instance Monad Sampler where
    return = pure
    (Sam sf) >>= f = Sam $ \rs-> let (x, rs'::Seed) = sf rs
                                     nextProb = f x
                                 in case nextProb of
                                      Sam g -> g rs'
                                      Samples xs -> primOneOf xs rs'
    (Samples xs) >>= f = Sam $ \rs-> let (x, rs'::Seed) = primOneOf xs rs
                                         nextProb = f x
                                     in case nextProb of
                                          Sam g -> g rs'
                                          Samples ys -> primOneOf ys rs'

-- | given a seed, return an infinite list of draws from sampling function
runSampler :: Seed -> Sampler a -> [a]
runSampler pmt s@(Sam sf)
   = let (x, pmt') = sf pmt
     in x : runSampler pmt' s
runSampler _ (Samples xs) = xs

-- | Get a seed
getSeedIO :: IO Seed
getSeedIO = do
   args <- getArgs
   case mapMaybe (stripPrefix "--seed=") args of
      [] ->  newPureMT
      sdStr:_ -> return $ pureMT $ read sdStr

-- | Return an infinite list of draws from sampling function in the IO monad
runSamplerIO :: Sampler a -> IO [a]
runSamplerIO s =
   fmap (`runSampler` s) $ getSeedIO

-- | Return a singe draw from sampling function
sampleIO :: Sampler a -> IO a
sampleIO s = head `fmap` runSamplerIO s

-- | Return a list of n draws from sampling function
sampleNIO :: Int -> Sampler a -> IO [a]
sampleNIO n s = take n `fmap` runSamplerIO s

-- | Estimate the probability that a hypothesis is true (in the IO monad)
eval :: Sampler Bool -> Sampler Double
eval s = do
  bs <- replicateM 1000 s
  return $ realToFrac (length (filter id bs)) / 1000

{-mu :: Vector Double
sigma :: Matrix Double

mystery = -1
mu = fromList [0,0,0]
sigma =  (3><3) [ 1,    1,       0,
                  1,    1,     mystery,
                  0,  mystery,   1]

samIt = sampleNIO 2 $ multiNormal mu sigma
 -}
-- | The joint distribution of two independent distributions
joint :: Sampler a -> Sampler b -> Sampler (a,b)
joint sf1 sf2 = liftM2 (,) sf1 sf2

-- | The joint distribution of two distributions where one depends on the other
jointConditional :: Sampler a -> (a-> Sampler b) -> Sampler (a,b)
jointConditional sf1 condsf
    = do x <- sf1
         y <- condsf x
         return (x,y)

--replicateM :: Monad m => Int -> m a -> m [a]
--replicateM n ma = forM [1..n] $ const ma

-- * Uniform distributions

-- | The unit interval U(0,1)
unitSample :: Sampler Double
unitSample = Sam randomDouble

-- | for x and y, the uniform distribution between x and y
uniform :: (Fractional a) => a -> a -> Sampler a
uniform a b = (\x->(realToFrac x)*(b-a)+a) `fmap` unitSample

-- * Normally distributed sampling function

--http://en.wikipedia.org/wiki/Box-Muller_transform
-- | The univariate gaussian (normal) distribution defined by mean and standard deviation
gauss :: (Floating b) => b -> b -> Sampler b
gauss m sd =
    do (u1,u2) <- (mapPair realToFrac) `fmap` joint unitSample unitSample
       return $ sqrt(-2*log(u1))*cos(2*pi*u2)*sd+m
  where mapPair f (x,y) = (f x, f y)

-- | Gaussians specialised for doubles
gaussD :: Double -> Double -> Sampler Double
gaussD m sd =
    do (u1,u2) <- joint unitSample unitSample
       return $ sqrt(-2*log(u1))*cos(2*pi*u2)*sd+m

gaussMany :: Floating b => [(b,b)] -> Sampler [b]
gaussMany means_sds = do gus <- gaussManyUnit (length means_sds)
                         return $ map f $ zip gus means_sds
    where f (gu, (mean, sd)) = gu*sd+mean

gaussManyD :: [(Double,Double)] -> Sampler [Double]
gaussManyD means_sds = do gus <- gaussManyUnitD (length means_sds)
                          return $ zipWith f gus means_sds
    where f gu (mean, sd) = gu*sd+mean

gaussManyUnit :: Floating b => Int -> Sampler [b]
gaussManyUnit 0 = return []
gaussManyUnit n | odd n = liftM2 (:) (gauss 0 1) (gaussManyUnit (n-1))
                | otherwise = do us <- forM [1..n] $ const $ unitSample
                                 return $ gaussTwoAtATime $ map realToFrac us
  where
    gaussTwoAtATime :: Floating a =>  [a] -> [a]
    gaussTwoAtATime (u1:u2:rest) = sqrt(-2*log(u1))*cos(2*pi*u2) : sqrt(-2*log(u1))*sin(2*pi*u2) : gaussTwoAtATime rest
    gaussTwoAtATime _ = []

gaussManyUnitD :: Int -> Sampler [Double]
gaussManyUnitD 0 = return []
gaussManyUnitD n | odd n = liftM2 (:) (gauss 0 1) (gaussManyUnit (n-1))
                | otherwise = do us <- forM [1..n] $ const $ unitSample
                                 return $ gaussTwoAtATimeD us
 where
   gaussTwoAtATimeD :: [Double] -> [Double]
   gaussTwoAtATimeD (u1:u2:rest) = sqrt(-2*log(u1))*cos(2*pi*u2) : sqrt(-2*log(u1))*sin(2*pi*u2) : gaussTwoAtATimeD rest
   gaussTwoAtATimeD _ = []

-- | Multivariate normal distribution
multiNormal :: Vector Double -> Matrix Double -> Sampler (Vector Double)
multiNormal mu sigma =
  let c =   cholSH sigma
      a = trans c
      k = dim mu
  in do z <- fromList `fmap` gaussManyUnitD k
--        return $ mu + (head $ toColumns $ a*asRow z)
        let c = asColumn z
        let r = asRow z
        return $ (mu + (head $ toColumns $ a `multiply` asColumn z))

multiNormalByChol :: Vector Double -> Matrix Double -> Sampler (Vector Double)
multiNormalByChol mu cholSigma =
  let a = trans $ cholSigma
      k = dim mu
  in do z <- fromList `fmap` gaussManyUnitD k
--        return $ mu + (head $ toColumns $ a*asRow z)
        let c = asColumn z
        let r = asRow z
        return $ (mu + (head $ toColumns $ a `multiply` asColumn z))

multiNormalIndep  :: Vector Double -> Vector Double -> Sampler (Vector Double)
multiNormalIndep vars mus = do
   let k = dim mus
   gs <- gaussManyUnitD k
   return $ fromList $ zipWith3 (\var mu g -> g*sqrt(var) + mu) (toList vars) (toList mus) gs

--http://en.wikipedia.org/wiki/Log-normal_distribution#Generating_log-normally-distributed_random_variates

-- | log-normal distribution <http://en.wikipedia.org/wiki/Log-normal_distribution>
logNormal :: (Floating b) => b -> b -> Sampler b
logNormal m sd =
    do n <- gauss 0 1
       return $ exp $ m + sd * n

-- * Other distribution

-- | Bernoulli distribution. Returns a Boolean that is 'True' with probability 'p'
bernoulli :: Double -> Sampler Bool
bernoulli p = (<p) `fmap` unitSample

discrete :: [(Double,a)] -> Sampler a
discrete weightedSamples =
   let sumWeights = sum $ map fst weightedSamples
       cummWeightedSamples = scanl (\(csum,_) (w,x) -> (csum+w,x)) (0,undefined) $ sortBy (comparing fst) weightedSamples
   in do u <- unitSample
         return . snd . fromJust $ find ((>=u*sumWeights) . fst) cummWeightedSamples

primOneOf :: [a] -> Seed -> (a, Seed)
primOneOf xs seed
   = let (u, nextSeed) = randomDouble seed
         idx = floor $ (realToFrac u)*(realToFrac $ length xs )
     in (xs !! idx, nextSeed)

oneOf :: [a] -> Sampler a
oneOf xs = do idx <- floor `fmap` uniform (0::Double) (realToFrac $ length xs )
              return $ xs !! idx

nOf :: Int -> [a] -> Sampler [a]
nOf n xs = sequence $ replicate n $ oneOf xs

{-main = do
  rnds <- take 1000 `fmap` runSamplerSysRan (oneOf [1,2,3])
  let diff = sort $ nub rnds
  print $ map (\x->(x, length $ filter (==x) rnds )) $ diff -}

-- | Bayesian inference from likelihood and prior using rejection sampling.
bayesRejection :: (PDF.PDF a) -> Double -> Sampler a -> Sampler a
bayesRejection p c q = bayes
    where bayes = do x <- q
                     u <- unitSample
                     if u < p x / c
                        then return x
                        else bayes

{-
expectation :: Fractional a =>  Int -> Sampler a -> IO a
expectation n sf =
    (mean . take n) `fmap` runSamplerIO sf

expectSD :: Floating a =>  Int -> Sampler a -> IO (a,a)
expectSD n sf =
    (meanSD . take n) `fmap` runSamplerIO sf
-}

--mapPair :: (a->b) -> (a,a) -> (b,b)

{-
--http://cgi.cse.unsw.edu.au/~dons/blog/2008/05/16#fast
mean :: Fractional a =>  [a] -> a
mean = go 0 0
        where
            -- go ::  -> Int -> [Double] -> Double
            go s n []     = s / fromIntegral n
            go !s !n (x:xs) = go (s+x) (n+1) xs

meanSD :: Floating a => [a] -> (a,a)
meanSD = go 0 0 0
    where go sq s n [] = let len = fromIntegral n in
                         (s/len, (recip len)*sqrt (len*sq-s*s))
          go !sq !s !n (x:xs) = go (sq+x*x) (s+x) (n+1) xs

test_mean = mean [0..1e8]
test_meanVar = meanSD [0..1e8]

main = do u <- expectSD 1000000 $ gauss 0 1
          print u

-}

--poisson ::  ::  Double -> [Double] -> IO Double
-- | Exponential distribution
expDist rate =  (\u-> negate $ (log(1-u))/rate) `fmap` unitSample

-- | Homogeneous poisson process defined by rate and duration
poissonMany :: Double -> Double -> Sampler [Double]
poissonMany rate tmax = aux 0
    where aux last = do
            next <- (+last) `fmap` expDist rate
            if next > tmax
               then return []
               else liftM2 (:) (return next) $ aux next

-- | binomial distribution
binomial :: Int -> Double -> Sampler Int
binomial n p = do
  bools <- forM [1..n] $ const $ fmap (<p) unitSample
  return $ length $ [t | t@True <- bools]

-- from random-fu
-- | Gamma distribution
gamma :: Double -> Double -> Sampler Double
gamma a b
     | a < 1
    = do
        u <- unitSample
        x <- gamma (1 + a) b
        return (x * u ** recip a)
    | otherwise
    = go
        where
            d = a - (1 / 3)
            c = recip (3 * sqrt d) -- (1 / 3) / sqrt d

            go = do
                x <- gaussD 0 1

                let cx = c * x
                    v = (1 + cx) ^ 3

                    x_2 = x * x
                    x_4 = x_2 * x_2

                if cx <= (-1)
                    then go
                    else do
                        u <- unitSample

                        if         u < 1 - 0.0331 * x_4
                            || log u < 0.5 * x_2  + d * (1 - v + log v)
                            then return (b * d * v)
                            else go

-- | inverse gamma distribution
invGamma :: Double -> Double -> Sampler Double
invGamma a b = recip `fmap` gamma a b

--http://en.wikipedia.org/wiki/Multivariate_normal_distribution#Drawing_values_from_the_distribution
--multiNormal :: Vector Double -> Matrix Double -> Sampler (Vector Double)

--http://www.xycoon.com/beta_randomnumbers.htm
-- | beta distribution
beta :: Int -> Int -> Sampler Double
beta a b =
    let gam n = do us <- forM [1..n] $ const unitSample
                   return $ log $ product us
    in do gama1 <- gamma (realToFrac a) 1
--          gama2 <- gamma (realToFrac a) 1

          gamb <- gamma (realToFrac b) 1
          return $ gama1/(gama1+gamb)

tbeta = sampleNIO 100 $ beta 1 1


module NumberParser
  ( parseNumber
  , parseFloat
  , parseComplex
  ) where

import Data.Complex
import LispVal
import Numeric
import Text.ParserCombinators.Parsec

parseNumber :: Parser LispVal
parseNumber = parseDecimal1 <|> parseDecimal2 <|> parseHex <|> parseOct <|> parseBin

parseDecimal1 :: Parser LispVal
parseDecimal1 = Number . read <$> many1 digit

parseDecimal2 :: Parser LispVal
parseDecimal2 = do
  try $ string "#d"
  x <- many1 digit
  (return . Number . read) x

parseHex :: Parser LispVal
parseHex = do
  try $ string "#x"
  x <- many1 hexDigit
  return $ Number (hex2dig x)

parseOct :: Parser LispVal
parseOct = do
  try $ string "#o"
  x <- many1 octDigit
  return $ Number (oct2dig x)

parseBin :: Parser LispVal
parseBin = do
  try $ string "#b"
  x <- many1 (oneOf "10")
  return $ Number (bin2dig x)

oct2dig x = fst $ head (readOct x)

hex2dig x = fst $ head (readHex x)

bin2dig = bin2dig' 0

bin2dig' digint "" = digint
bin2dig' digint (x:xs) =
  let old =
        2 * digint +
        (if x == '0'
           then 0
           else 1)
   in bin2dig' old xs

toDouble :: LispVal -> Double
toDouble (Float f) = realToFrac f
toDouble (Number n) = fromIntegral n

parseFloat :: Parser LispVal
parseFloat = do
  x <- many1 digit
  char '.'
  y <- many1 digit
  return $ Float (fst . head $readFloat (x ++ "." ++ y))

parseComplex :: Parser LispVal
parseComplex = do
  x <- try parseFloat <|> parseDecimal1 <|> parseDecimal2
  char '+'
  y <- try parseFloat <|> parseDecimal1 <|> parseDecimal2
  char 'i'
  return $ Complex (toDouble x :+ toDouble y)


module Main where

import Numeric.LinearAlgebra
import LeastSquares

main :: IO ()
main = do
  let a = Mat $ (3><2) [1.0, 1.0, 1.0, 4.0, 2.0, -1.0]
  let b = Vec $ vector [13, 27, 1]
  let x = Var "x" 2
  print $ minimize (SumSquares (a * x - b)) []
  -- let a1 = Mat $ (3><1) [1.0, 1.0, 1.0]
  -- let a2 = Mat $ (3><1) [2.0, 4.0, -1.0]
  -- let x1 = Var "x1" 1
  -- let x2 = Var "x2" 1
  print $ minimize (SumSquares (a * x - b) + 2 * SumSquares x) []
  let c = Mat $ (1><2) [1.0, -1]
  let d = Vec $ vector [2]
  print $ minimize (SumSquares (a * x - b)) [c * x :==: d]


{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE ScopedTypeVariables #-}

module Fields (C, Field (..), Fp, Prime (..), Q, R, Val (..), modInv) where

import Data.Complex (Complex (..))
import Data.Ratio (denominator, numerator)

------------------------------------------------------------------------------
-- Fields

-- Value information
data Val = None | NumDen Integer Integer | ReIm Double Double

-- Field type class
class (Eq f, Fractional f, Show f) => Field f where
  {-# MINIMAL char, val #-}
  char :: f -> Integer
  val :: f -> Val

------------------------------------------------------------------------------
-- Characteristic zero

-- Rational numbers Q
newtype Q = Q Rational deriving (Eq, Fractional, Num, Ord)
instance Show Q where
  show q = show n ++ if d == 1 then "" else "/" ++ show d
    where
      NumDen n d = val q
instance Field Q where
  char = const 0
  val (Q q) = NumDen (numerator q) (denominator q)

-- Real numbers R
newtype R = R Double deriving (Eq, Fractional, Num, Ord)
instance Show R where
  show (R r) = if r == fromIntegral r' then show r' else show r
    where
      r' = floor r
instance Field R where
  char = const 0
  val = const None

-- Complex numbers C
newtype C = C (Complex Double) deriving (Eq, Fractional, Num)
instance Show C where
  show (C (r :+ i))
    | signum i == -1 = "(" ++ show (R r) ++ "-" ++ show (R (abs i)) ++ "i)"
    | signum i == 1 = "(" ++ show (R r) ++ "+" ++ show (R i) ++ "i)"
    | otherwise = show (R r)
instance Field C where
  char = const 0
  val (C (r :+ i)) = ReIm r i

------------------------------------------------------------------------------
-- Characteristic prime

-- Prime numbers
class (Enum p, Show p) => Prime p where
  prime :: p -> Integer
  prime _ = read . tail . show $ (toEnum 0 :: p)

-- Prime subfields Fp
newtype Fp p = Fp Integer

-- Fp field instance
instance Prime p => Field (Fp p) where
  char = const $ prime (undefined :: p)
  val = const None

-- Fp standard instances
instance Prime p => Bounded (Fp p) where
  minBound = 0
  maxBound = -1
instance Prime p => Enum (Fp p) where
  toEnum = fromIntegral
  fromEnum (Fp n) = fromInteger n
instance Eq (Fp p) where
  Fp n == Fp n' = fromInteger n == fromInteger n'
instance Prime p => Fractional (Fp p) where
  fromRational n = fromInteger (numerator n) / fromInteger (denominator n)
  recip (Fp n) = case modInv n $ prime (undefined :: p) of
    Right m -> fromIntegral m
    Left m -> error $ show m
instance Prime p => Integral (Fp p) where
  quotRem (Fp n) (Fp n') = (fromInteger q, fromInteger r)
    where
      (q, r) = quotRem n n'
  toInteger (Fp n) = fromInteger n
instance Prime p => Num (Fp p) where
  Fp n + Fp n' = fromInteger $ n + n'
  Fp n * Fp n' = fromInteger $ n * n'
  abs n = n
  signum n = if n == 0 then 0 else 1
  fromInteger n = Fp . mod n $ prime (undefined :: p)
  negate (Fp n) = fromInteger $ (-n)
instance Ord (Fp p) where
  Fp n <= Fp n' = n <= n'
instance Prime p => Real (Fp p) where
  toRational (Fp n) = fromInteger n
instance Show (Fp p) where
  show (Fp n) = show n

-- Extended Euclidean algorithm
extGCD :: Integral a => a -> a -> ((a, a), a)
extGCD 0 y = ((0, 1), y)
extGCD x y = ((t - s * q, s), g)
  where
    (q, r) = quotRem y x
    ((s, t), g) = extGCD r x

-- Modular inverse
modInv :: Integral a => a -> a -> Either a a
modInv x p = if g == 1 then return (mod y p) else Left g
  where
    ((y, _), g) = extGCD x p


module Main where

import Graphics.Rendering.Chart.Easy
import Graphics.Rendering.Chart.Backend.Cairo

import Statistics.Distribution.Normal
import Statistics.Distribution

-- pOak = class1
-- pBeech = class2

stdDevOak :: NormalDistribution
stdDevOak = normalDistr (- 2.25) 1

stdDevBeech :: NormalDistribution
stdDevBeech = normalDistr 1.596 1

type Priori = Double

pOak :: Priori
pOak = 0.25

pBeech :: Priori
pBeech = 0.75

oak :: [Double] -> [(Double,Double)]
oak xs = [ (x,(density stdDevOak x) * pOak) | x <- xs ]

beech :: [Double] -> [(Double, Double)]
beech xs = [ (x,(density stdDevBeech x) * pBeech) | x <- xs ]

identit :: [Double] -> [(Double, Double)]
identit xs = [(x, x) | x <- xs]

increments :: Double -> Double -> Double -> [Double]
increments from to inc
    | from > to = []
    | otherwise = (from):(increments (from + inc) to inc)

--main = print $ (increments (-5) 5 0.1)

main = toFile def "output.png" $ do
    layout_title .= "KMeans"
    setShapes [PointShapeCircle, PointShapePlus, PointShapeStar]
    setColors [opaque blue, opaque red]
    plot (points "class2" $ beech (increments (-5) 5 0.1))
    plot (line "class1" [oak (increments (-5) 5 0.1)])


{-# OPTIONS_GHC -Wall #-}

module Numeric.FFT.HMatrix.Vector (
  -- * Complex-to-complex transforms
  dft
  , idft
    -- * Real-to-complex transforms
  , dftR2C
  , dftC2R
    -- * Real-to-real transforms
    -- ** Discrete cosine transforms
  , dct1
  , dct2
  , dct3
  , dct4
    -- ** Discrete sine transforms
  , dst1
  , dst2
  , dst3
  , dst4
  ) where

import Foreign.Storable ( Storable )
import Numeric.LinearAlgebra.Data ( Vector, Complex )
import Numeric.LinearAlgebra ( Element )
import qualified Numeric.LinearAlgebra.Data as M
import qualified Numeric.FFT.Vector.Unnormalized as FFT

dft :: Vector (Complex Double) -> Vector (Complex Double)
dft = FFT.run FFT.dft

idft :: Vector (Complex Double) -> Vector (Complex Double)
idft = FFT.run FFT.idft

dftR2C :: Vector Double -> Vector (Complex Double)
dftR2C = FFT.run FFT.dftR2C

dftC2R :: Vector (Complex Double) -> Vector Double
dftC2R = FFT.run FFT.dftC2R

dct1 :: Vector Double -> Vector Double
dct1 = FFT.run FFT.dct1

dct2 :: Vector Double -> Vector Double
dct2 = FFT.run FFT.dct2

dct3 :: Vector Double -> Vector Double
dct3 = FFT.run FFT.dct3

dct4 :: Vector Double -> Vector Double
dct4 = FFT.run FFT.dct4

dst1 :: Vector Double -> Vector Double
dst1 = FFT.run FFT.dst1

dst2 :: Vector Double -> Vector Double
dst2 = FFT.run FFT.dst2

dst3 :: Vector Double -> Vector Double
dst3 = FFT.run FFT.dst3

dst4 :: Vector Double -> Vector Double
dst4 = FFT.run FFT.dst4


-- | authors: Connor Ford, Jake Hauser

module Persistence where

import Data.List (intercalate)
import Numeric.LinearAlgebra
import Data.List.HT (mapAdjacent)
import ExplicitSimplexStream

-- Betti Vector
-- This data type is equivalent to a list of Ints. It represents the betti vector.
-- For a betti vector of length n, all betti numbers b_i such that i >= n are assumed to be 0.
data BettiVector = BettiVector [Int] deriving (Eq)

instance Show BettiVector where
    show (BettiVector []) = "()"
    show (BettiVector vs) = "(" ++ intercalate ", " (map show vs) ++ ")"

-- getBoundaryMapHelper (see: getBoundaryMap)
-- First argument is SimplexListByDegree representing C_k
-- Second argument is SimplexListByDegree representing C_k+1
-- Third argument should initially be the length of the x in (x:xs) for C_k+1. It it a counter for removing elements from x.
-- Fourth argument should initially be zero. It is a counter for the index of x in (x:xs).
-- Fifth argument should be the zero matrix for _correct_ boundary map calculations.
getBoundaryMapHelper :: (Ord a) => SimplexListByDegree a -> SimplexListByDegree a -> Int -> Int -> Matrix Double -> Matrix Double
getBoundaryMapHelper _     (SimplexListByDegree _ [])     _ _ mat = mat
getBoundaryMapHelper list1 (SimplexListByDegree n (_:xs)) 0 k mat = getBoundaryMapHelper list1 (SimplexListByDegree n (xs)) n (k+1) mat
getBoundaryMapHelper list1 (SimplexListByDegree n (x:xs)) m k mat =
    let
        x_dim = indexOfSimplex (Simplex (removeSimplexElement x (m-1))) list1
        y_dim = k
        val = (-1)^(m-1)
        updatedMatrix = accum mat (\a _ -> a) [((x_dim,y_dim), val)] -- use accumulator to change values of matrix
    in
        getBoundaryMapHelper list1 (SimplexListByDegree n (x:xs)) (m-1) k updatedMatrix

-- getBoundaryMap
-- First argument C_k
-- Second argument C_k+1
-- return matrix dimensions: len(C_k) = num rows, len(C_k+1) = num cols
-- algorithm:
-- Given element of C_k+1, y, remove i'th element from y (starting with index zero) to get element in C_k called x.
-- Then the value in the matrix at row idxOf(x) and column idxOf(y) is (-1)^i.
getBoundaryMap :: (Ord a) => SimplexListByDegree a -> SimplexListByDegree a -> Matrix Double
getBoundaryMap list1@(SimplexListByDegree _ simps1) list2@(SimplexListByDegree n simps2) =
    let
        zero_matrix = matrix (length simps2) (replicate ((length simps1)*(length simps2)) 0)
    in
        getBoundaryMapHelper list1 list2 n 0 zero_matrix

-- getHomologyDimension
-- first arg is D_i+1
-- second arg is D_i
-- third argument k (or i in the below illustration)
-- want to find: dim ker D_i - dim im D_i+1
-- algorithm to find ker and im:
-- (i) rref matrix
-- (ii) # of non-zero rows correspond to rnk of original matrix.
-- (iii) # of columns with leading 1’s correspond to the columns of original matrix that span image.
-- (iv) dim ker = (ii) - (iii)
getHomologyDimension :: Matrix Double -> Matrix Double -> Int -> Int
getHomologyDimension m1 m2 k =
    let
        m1_rank = rank m1
        m2_columns = cols m2
        m2_rank = rank m2
        m2_nullity = m2_columns - m2_rank
    in
        if k > 0 then
            m2_nullity - m1_rank
        else
            m2_columns - m1_rank

-- persistenceHelper (see: persistence)
-- First argument is a list of matrices representing a sequence of boundary maps.
-- Second argument is k representing the k'th homology.
persistenceHelper :: [Matrix Double] -> Int -> [Int]
persistenceHelper []       _ = []
persistenceHelper [x]      k = [getHomologyDimension ((ident 1) - (ident 1)) x k] -- base case for last boundary map
persistenceHelper (x:y:xs) k = (getHomologyDimension y x k):(persistenceHelper (y:xs) (k+1))

-- persistence
-- The persistence algorithm computes the BettiVector for s simplex stream. The algorithm is as follows:
-- Given a simplex stream, and a field coefficient p corresponding to Z/pZ we will return betti numbers (b_0, b_1, ..., b_n) where n+1 equals the degree of the highest order simplex in the stream.
-- (1) Get simplex lists ordered by the degree of simplices
-- (2) Get boundary maps D_0, ..., D_{n-1} such that:
-- 0 <-- C_0 <-- C_1 <-- ... <-- C_{n-1} <-- C_n <-- 0
--  D_{-1}   D_0     D_1                D_{n-1}
-- (3) Compute dimensions of the Homologies:
-- dim H_0 = dim Ker 0 - dim Im D_0
-- dim H_1 = dim Ker D_0 - dim Im D_1
-- ...
-- dim H_n = dim Ker D_{n-1} - dim Im 0
-- (4) return betti profile (dim H_0, dim H_1, ..., dim H_n)
-- TODO: implement with rref mod p where Field = Z/pZ.
persistence :: (Ord a) => Stream a -> Int -> BettiVector
persistence stream field =
    let
        (OrderedSimplexList l) = streamToOrderedSimplexList stream
    in
        BettiVector $ persistenceHelper (mapAdjacent getBoundaryMap l) 0


{-|
Module      : ControlSystems.DynamicSystems.Conversion
Description : Converts continuous time systems to discrete time
Copyright   : (c) Ryan Orendorff, 2020
License     : BSD3
Stability   : experimental
-}
module ControlSystems.DynamicSystems.Conversions (
  c2d
)

where

import Numeric.LinearAlgebra

-- | Convert a continuous time linear differential system to a discrete form,
-- using state space models.
c2d :: Double                         -- ^ The sampling time
    -> Matrix Double                  -- ^ The state transition matrix, aka A
    -> Matrix Double                  -- ^ The state input matrix, aka B
    -> (Matrix Double, Matrix Double) -- ^ The discretized A & B system
c2d t a b = (a_d, b_d)
  where
    -- First we have to extract the sizes of all the matrices at run time.
    (m_a, n_a) = (rows a, cols a)
    (_  , n_b) = (rows b, cols b)

    -- Then manually make the correct block. This must be a square matrix
    -- because the matrix exponential function (`expm`) expects to only have
    -- square inputs (given that it is a series expansion of matrix to higher
    -- powers). We need to form this matrix
    -- ⌈ A B ⌉
    -- ⌊ 0 0 ⌋
    block = a ||| b
              ===
              konst 0 (m_a, n_a + n_b) -- I left in a bug here! ;-D

    -- Calculate the discrete time block matrix, which is the following
    -- ⌈ A_d B_d ⌉
    -- ⌊  0   I  ⌋
    exp_block = expm (scale t block)

    -- Finally we can extract the correct submatrices assuming we pull out the
    -- right pieces.
    a_d = subMatrix (0, 0) (m_a, n_a) exp_block
    b_d = subMatrix (0, n_a) (m_a, n_b) exp_block


-- | RVM regression model
module Numeric.BMML.RVMR
    ( fit
    , predict
    ) where

import Numeric.LinearAlgebra.Data (Matrix, Vector)
import qualified Numeric.LinearAlgebra.HMatrix as H

data RVMRModel = RVMRModel
    { mean :: Vector Double
    , covariance :: Matrix Double
    , optbeta :: Double
    }

initAlpha :: Int -> Vector Double
initAlpha m = H.konst 1.0e-3 m

initBeta :: Double
initBeta = 1.0e-3

fit :: Matrix Double -> Vector Double -> RVMRModel
fit x t =
    let (_,d) = H.size x
        alpha = initAlpha d
        beta = initBeta
        sigma = calcSigma alpha beta x
        m = calcMean beta sigma x t
    in go 10000 alpha beta sigma m
  where
    go 0 _ beta sigma m =
        RVMRModel
        { mean = m
        , covariance = sigma
        , optbeta = beta
        }
    go k alpha beta sigma m = if sNorm alpha na < 0.01
                                 then go 0 na nb ns nm
                                else go (k - 1) na nb ns nm
      where
        (n,d) = H.size x
        gamma :: Vector Double
        gamma = (H.konst 1 d) - (alpha * (H.takeDiag sigma))
        na :: Vector Double
        na =
            H.fromList
                (zipWith
                     (\gi mi ->
                           gi / (mi * mi))
                     (H.toList gamma)
                     (H.toList m))
        nb :: Double
        nb = normSq (t - x H.#> m) / (fromIntegral n - H.sumElements gamma)
        ns = calcSigma na nb x
        nm = calcMean nb ns x t

sNorm :: Vector Double -> Vector Double -> Double
sNorm x y =
    sqrt
        (foldr
             (\(xi,yi) res ->
                   res + (yi - xi) * (yi - xi))
             0
             (filter
                  (\(xi,yi) ->
                        xi < 10000 && yi < 10000)
                  (zip (H.toList x) (H.toList y))))
{-# INLINE sNorm #-}

calcSigma :: Vector Double -> Double -> Matrix Double -> Matrix Double
calcSigma a b x = H.inv (H.diag a + b `H.scale` (H.tr x H.<> x))
{-# INLINE calcSigma #-}

calcMean :: Double
         -> Matrix Double
         -> Matrix Double
         -> Vector Double
         -> Vector Double
calcMean b s x t = b `H.scale` ((s H.<> H.tr x) H.#> t)
{-# INLINE calcMean #-}

normSq :: Vector Double -> Double
normSq v = v `H.dot` v
{-# INLINE normSq #-}

predict :: RVMRModel -> Vector Double -> (Double, Double)
predict (RVMRModel m sigma ob) v =
    (m `H.dot` v, 1 / ob + v `H.dot` (sigma H.#> v))


{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE DeriveDataTypeable, DeriveGeneric #-}
-- |
-- Module    : Statistics.Distribution.Exponential
-- Copyright : (c) 2009 Bryan O'Sullivan
-- License   : BSD3
--
-- Maintainer  : bos@serpentine.com
-- Stability   : experimental
-- Portability : portable
--
-- The exponential distribution.  This is the continuous probability
-- distribution of the times between events in a poisson process, in
-- which events occur continuously and independently at a constant
-- average rate.

module Statistics.Distribution.Exponential
    (
      ExponentialDistribution
    -- * Constructors
    , exponential
    , exponentialE
    -- * Accessors
    , edLambda
    ) where

import Control.Applicative
import Data.Aeson                      (FromJSON(..),ToJSON,Value(..),(.:))
import Data.Binary                     (Binary, put, get)
import Data.Data                       (Data, Typeable)
import GHC.Generics                    (Generic)
import Numeric.SpecFunctions           (log1p)
import Numeric.MathFunctions.Constants (m_neg_inf)
import qualified System.Random.MWC.Distributions as MWC
import qualified Data.Vector.Generic as G

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

newtype ExponentialDistribution = ED {
      edLambda :: Double
    } deriving (Eq, Typeable, Data, Generic)

instance Show ExponentialDistribution where
  showsPrec n (ED l) = defaultShow1 "exponential" l n
instance Read ExponentialDistribution where
  readPrec = defaultReadPrecM1 "exponential" exponentialE

instance ToJSON ExponentialDistribution
instance FromJSON ExponentialDistribution where
  parseJSON (Object v) = do
    l <- v .: "edLambda"
    maybe (fail $ errMsg l) return $ exponentialE l
  parseJSON _ = empty

instance Binary ExponentialDistribution where
  put = put . edLambda
  get = do
    l <- get
    maybe (fail $ errMsg l) return $ exponentialE l

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

instance D.ContDistr ExponentialDistribution where
    density (ED l) x
      | x < 0     = 0
      | otherwise = l * exp (-l * x)
    logDensity (ED l) x
      | x < 0     = m_neg_inf
      | otherwise = log l + (-l * x)
    quantile      = quantile
    complQuantile = complQuantile

instance D.Mean ExponentialDistribution where
    mean (ED l) = 1 / l

instance D.Variance ExponentialDistribution where
    variance (ED l) = 1 / (l * l)

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

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

instance D.Entropy ExponentialDistribution where
  entropy (ED l) = 1 - log l

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

instance D.ContGen ExponentialDistribution where
  genContVar = MWC.exponential . edLambda

cumulative :: ExponentialDistribution -> Double -> Double
cumulative (ED l) x | x <= 0    = 0
                    | otherwise = 1 - exp (-l * x)

complCumulative :: ExponentialDistribution -> Double -> Double
complCumulative (ED l) x | x <= 0    = 1
                         | otherwise = exp (-l * x)

quantile :: ExponentialDistribution -> Double -> Double
quantile (ED l) p
  | p >= 0 && p <= 1 = - log1p(-p) / l
  | otherwise        =
    error $ "Statistics.Distribution.Exponential.quantile: p must be in [0,1] range. Got: "++show p

complQuantile :: ExponentialDistribution -> Double -> Double
complQuantile (ED l) p
  | p == 0          = 0
  | p >= 0 && p < 1 = -log p / l
  | otherwise       =
    error $ "Statistics.Distribution.Exponential.quantile: p must be in [0,1] range. Got: "++show p

-- | Create an exponential distribution.
exponential :: Double            -- ^ Rate parameter.
            -> ExponentialDistribution
exponential l = maybe (error $ errMsg l) id $ exponentialE l

-- | Create an exponential distribution.
exponentialE :: Double            -- ^ Rate parameter.
             -> Maybe ExponentialDistribution
exponentialE l
  | l > 0     = Just (ED l)
  | otherwise = Nothing

errMsg :: Double -> String
errMsg l = "Statistics.Distribution.Exponential.exponential: scale parameter must be positive. Got " ++ show l

-- | Create exponential distribution from sample. Returns @Nothing@ if
--   sample is empty or contains negative elements. No other tests are
--   made to check whether it truly is exponential.
instance D.FromSample ExponentialDistribution Double where
  fromSample xs
    | G.null xs       = Nothing
    | G.all (>= 0) xs = Nothing
    | otherwise       = Just $! ED (S.mean xs)


{-|
Module: MachineLearning.Optimization.GradientDescent
Description: Gradient Descent
Copyright: (c) Alexander Ignatyev, 2016
License: BSD-3
Stability: experimental
Portability: POSIX

-}

module MachineLearning.Optimization.GradientDescent
(
  gradientDescent
)

where

import MachineLearning.Types (R, Vector, Matrix)
import MachineLearning.Regularization (Regularization)
import qualified Data.Vector.Storable as V
import qualified Numeric.LinearAlgebra as LA

import qualified MachineLearning.Model as Model

-- | Gradient Descent method implementation. See "MachineLearning.Regression" for usage details.
gradientDescent :: Model.Model a => R-> a -> R -> Int -> Regularization -> Matrix -> Vector -> Vector -> (Vector, Matrix)
gradientDescent alpha model eps maxIters lambda x y theta = helper theta maxIters []
  where gradient = Model.gradient model lambda
        cost = Model.cost model lambda
        helper theta nIters optPath =
          let theta' = theta - (LA.scale alpha (gradient x y theta))
              j = cost x y theta'
              gradientTest = LA.norm_2 (theta' - theta) < eps
              optPathRow = V.concat [LA.vector [(fromIntegral $ maxIters - nIters), j], theta']
              optPath' = optPathRow : optPath
          in if gradientTest || nIters <= 1
             then (theta, LA.fromRows $ reverse optPath')
             else helper theta' (nIters - 1) optPath'


module Lib.Practice.P3 where

import Control.Monad
import Numeric.LinearAlgebra

f3 = do
    let gates = [("AND", gAnd), ("NAND", gNand), ("OR", gOr)]
    mapM_ showOutput gates

gAnd = Neuron { weight = vector [0.5, 0.5], bias = -0.7 }
gNand = Neuron { weight = vector [-0.5, -0.5], bias = 0.7 }
gOr = Neuron { weight = vector [0.5, 0.5], bias = -0.2 }

inputs = replicateM 2 [NotFire, Fire]
showOutput (name, gate) = do
    putStrLn $ name ++ "ゲート"
    mapM_ (showOutput' gate) inputs
    putStrLn ""
    where
        showOutput' gate input = do
            let result = input --> gate
            putStrLn $ show input ++ " -> " ++ show result

data Neuron = Neuron {
    weight :: Vector R,
    bias :: R
} deriving Show

data Input = Fire | NotFire deriving Show
input2Int Fire = 1
input2Int NotFire = 0

(-->) = output
output :: [Input] -> Neuron -> Input
output inputs Neuron { weight = w, bias = b } = result
    where inputV    = fromList $ map input2Int inputs
          sum'      = sumElements $ inputV * w
          result'   = b + sum'
          result    = case () of
            _ | result' >= 0    -> Fire
            otherwise           -> NotFire


module Quantum where

import Data.AdditiveGroup
import Data.Complex hiding ((:+))
import qualified Data.Complex as C (Complex ((:+)))
import Data.Foldable (foldl')
import Data.VectorSpace hiding (magnitude)
import Lib
import Matrix
import Vector

-- base type for qubits
type CDouble = Complex Double

type CVVec n = VVec n CDouble

type Qubit = CVVec Two

type NQubits n = CVVec (Exp Two n)

-- for ease of use
sqtwo :: Floating a => a
sqtwo = 1 / sqrt 2

i :: Num a => Complex a
i = 0 C.:+ 1

-- define some basic qubits/vectors, and qubit collections
basisVec :: forall n a. (Num a, KnownNat n) => Fin n -> VVec n a
basisVec f = generateMat (\x _ -> fromIntegral $ fromEnum $ x == f)

zero' :: forall n a. (Num a, KnownNat n) => VVec n a
zero' = basisVec FZero

zero :: Qubit
zero = zero'

one' :: forall n a m. (Num a, KnownNat n, n ~ 'Succ m) => VVec n a
one' = basisVec (FSucc FZero)

one :: Qubit
one = one'

plus :: Qubit
plus = sqtwo *^ zero ^+^ sqtwo *^ one

minus :: Qubit
minus = sqtwo *^ zero ^-^ sqtwo *^ one

qubits :: [Qubit]
qubits = [zero, one, plus, minus]

makeTwoQubits :: [Qubit] -> [NQubits Two]
makeTwoQubits qs = (.*.) <$> qs <*> qs

twoQubits :: [NQubits Two]
twoQubits = makeTwoQubits [zero, one]

consQubit :: CDouble -> CDouble -> Qubit
consQubit a b = numMatFromList [[a], [b]]

consVVec :: forall n a. (KnownNat n, Num a) => [a] -> VVec n a
consVVec ns = numMatFromList @n @One (fmap (: []) ns)

-- define some transformation matrices
pauliX :: Num a => Matrix Two Two a
pauliX = numMatFromList [[0, 1], [1, 0]]

pauliY :: RealFloat a => Matrix Two Two (Complex a)
pauliY = numMatFromList [[0, - i], [i, 0]]

pauliZ :: Num a => Matrix Two Two a
pauliZ = numMatFromList [[1, 0], [0, -1]]

hadamard :: Floating a => Matrix Two Two a
hadamard = (* sqtwo) <$> numMatFromList [[1, 1], [1, -1]]

-- a rotation matrix. produces a matrix which rotates by n radians
rotation :: Floating a => a -> Matrix Two Two a
rotation n = numMatFromList [[c, - s], [s, c]]
  where
    c = cos n
    s = sin n

cnot :: Num a => Matrix Four Four a
cnot = numMatFromList [[1, 0, 0, 0], [1, 0, 0, 0], [0, 0, 0, 1], [0, 0, 1, 0]]

cnot' :: Num a => Matrix Four Four a
cnot' = numMatFromList [[1, 0, 0, 0], [0, 0, 0, 1], [0, 0, 1, 0], [0, 1, 0, 0]]

iden2 :: Num a => Matrix Two Two a
iden2 = identity

swap :: Num a => Matrix Four Four a
swap = numMatFromList [[1, 0, 0, 0], [0, 0, 1, 0], [0, 1, 0, 0], [0, 0, 0, 1]]

-- define the tensor product
tensorProd ::
  (Num a) => Matrix n m a -> Matrix i j a -> Matrix (Mul n i) (Mul m j) a
tensorProd n m = expandNested $ fmap (\a -> fmap (a *) m) n

-- ease of use function for the tensor product
(.*.) :: (Num a) => Matrix n m a -> Matrix i j a -> Matrix (Mul n i) (Mul m j) a
(.*.) = tensorProd

infixr 8 .*.

-- apply the tensor product a number of times.
-- to use this function, use `tensorPower @(type number) matrix`
-- for example, to tensor product matrix `a` gate three times, do `tensorPower @Three a`
-- If used in relation to other matrices, tensorPower can infer
--  the `i` value, which is pretty neat
-- Only works on non-vector matrices - see tensorPowerVVec for column vec variant
tensorPower ::
  forall i a n m.
  (Num a, KnownNat i, i ~ GetExp ('Succ n) i, i ~ GetExp ('Succ m) i) =>
  Matrix ('Succ n) ('Succ m) a ->
  Matrix (Exp ('Succ n) i) (Exp ('Succ m) i) a
tensorPower m = tensorPower' @i m natSing

-- apply the tensor product a number of times on a column vector
-- very similar to tensorPower, but only works on column vectors
tensorPowerVVec ::
  forall i a n.
  (Num a, KnownNat i, i ~ GetExp ('Succ n) i) =>
  Matrix ('Succ n) 'One a ->
  Matrix (Exp ('Succ n) i) (Exp 'One i) a
tensorPowerVVec m = tensorPower' @i m natSing

-- helper function for tensorPower
tensorPower' ::
  forall i n m a.
  Num a =>
  Matrix n m a ->
  NatS i ->
  Matrix (Exp n i) (Exp m i) a
tensorPower' m OneS = m
tensorPower' m (SuccS s) = m .*. tensorPower' m s

-- transpose and get the conjugate of the given qubit
-- effectively, find the bra
conjTrans :: CVVec n -> HVec n CDouble
conjTrans v = conjugate <$> Matrix.transpose v

-- using conjTrans, find the inner product of two qubits
innerProduct :: CVVec n -> CVVec n -> CDouble
innerProduct v v' = getVal $ conjTrans v *.* v'

-- generate all basis vectors for a given dimension
compBasis' :: forall n a. (Num a, KnownNat n) => Vector n (VVec n a)
compBasis' = generateVec basisVec

-- define the computational basis
compBasis :: Vector Two Qubit
compBasis = compBasis'

-- apply a transformation matrix to the computational basis
transformBasis ::
  forall n a. (Num a, KnownNat n) => Matrix n n a -> Vector n (VVec n a)
transformBasis m = fmap (m *.*) compBasis'

-- measure in a given basis, returning probabilities for each
-- basis vector
measureIn :: CVVec m -> Vector n (CVVec m) -> Vector n Double
measureIn v = fmap ((** 2) . magnitude . Quantum.innerProduct v)

-- measure in the computational basis, returning probabilities
-- for each basis vector
measure :: Qubit -> Vector Two Double
measure v = measureIn v compBasis

-- given a list of vector-probability pairs, return Just the density matrix
--  if the probabilities sum up to 1; else return Nothing
getDensityMatrix ::
  (KnownNat m) => [(CVVec m, CDouble)] -> Maybe (Matrix m m CDouble)
getDensityMatrix qs
  | sum (map snd qs) == 1 =
    Just $
      foldr (^+^) zeroed [p *^ f tq | (tq, p) <- qs]
  | otherwise = Nothing
  where
    zeroed = generateMat (\_ _ -> 0)
    f m = m *.* conjTrans m

-- given a density matrix (as seen in previous function),
--  and a complex number vector, return the probability of
--  seeing that vector
calcProbability :: Matrix m m CDouble -> CVVec m -> CDouble
calcProbability densityMatrix base =
  getVal $ conjTrans base *.* densityMatrix *.* base

chi_f :: CVVec m -> CVVec m -> CDouble
chi_f s x = (-1) ** Quantum.innerProduct s x

-- vecToFunc :: (KnownNat m) => CVVec m -> Matrix m m CDouble
-- vecToFunc v =
--   generateMat (\f f' -> if f == f' then getAtMatrix f FZero v else 0)

-- fhat :: (KnownNat n) => (CVVec n -> CDouble) -> CVVec n -> CDouble
-- fhat f s =
--   (foldl' (\b a -> b + (apply a)) 0 compBasis') / (2 ** (fst $ Matrix.size s))
--  where
--   chi_s = chi s
--   apply x = (f x) * chi_s x

createInput :: Num a => Vector n (VVec m a) -> VVec (Exp m n) a
createInput (VecSing v) = v
createInput (v :+ vs) = v .*. createInput vs

fourierBasis :: (KnownNat n) => Vector n (CVVec n)
fourierBasis = fmap chi compBasis'

chi :: (KnownNat n) => CVVec n -> CVVec n
chi s = normalise $ toVVec $ fmap (chi_f s) compBasis'

fhat_f :: (KnownNat n) => CVVec n -> CVVec n -> CDouble
fhat_f = Quantum.innerProduct

fhat :: (KnownNat n) => CVVec n -> CVVec n
fhat f = toVVec $ fmap (fhat_f f) fourierBasis

normalise :: (KnownNat n) => CVVec n -> CVVec n
normalise v = v ^/ sqrt (sum $ fmap (abs) v)

groverDiffusion ::
  forall i a.
  (KnownNat (Exp Two i), KnownNat i, i ~ GetExp Two i, Floating a) =>
  Matrix (Exp Two i) (Exp Two i) a
groverDiffusion =
  tensorPower hadamard *.* generateMat f *.* tensorPower hadamard
  where
    f a b
      | a == b && a == FZero = 1
      | a == b = -1
      | otherwise = 0

-- fhat :: (KnownNat n) => CVVec n -> CVVec n -> CDouble
-- fhat f s


{-# LANGUAGE Strict #-}

module HEP.Util.Polynomial (quarticEqSol, quadEqSol) where

import Numeric.GSL.Polynomials (polySolve)

import Data.Complex            (imagPart, realPart)
import Data.List               (nub, sort)

-- | returns the real roots of quartic polynomial equation in ascending order.
--
-- It does not return zero roots. Thus, if the roots are empty, it implies
-- that the polynomial has only zero roots.
quarticEqSol :: (Double -> Double)
             -> [Double]               -- ^ four distinct nonzero inputs
             -> Double                 -- ^ the cut value for complex roots
             -> Maybe (Int, [Double])  -- ^ (the number of real roots, roots)
quarticEqSol f xs eps
    | length (filter (/= 0) xs) /= 4 || xs /= nub xs = Nothing
    | otherwise = do
          let roots = filter (\r -> abs (imagPart r) < eps) $
                      polySolve (coeffQuartic f xs)
              realroots = (sort . filter ((> eps) . abs)) $ realPart <$> roots
          return (length realroots, realroots)

-- | returns the coefficients of quartic polynomial
--   using four distinct nonzero inputs.
--
--   a * x^4 + b * x^3 + c * x^2 + d * x + e
--
--   the output is the list of the coefficients [e, d, c, b, a].
coeffQuartic :: (Double -> Double) -> [Double] -> [Double]
coeffQuartic f xs =
    let [x1, x2, x3, x4] = xs
        [f1, f2, f3, f4] = f <$> xs

        d1 = x1 * (x2 - x1) * (x3 - x1) * (x4 - x1)
        d2 = x2 * (x2 - x1) * (x3 - x2) * (x4 - x2)
        d3 = x3 * (x3 - x1) * (x3 - x2) * (x4 - x3)
        d4 = x4 * (x4 - x1) * (x4 - x2) * (x4 - x3)

        p12 = x1 * x2
        p13 = x1 * x3
        p14 = x1 * x4
        p23 = x2 * x3
        p24 = x2 * x4
        p34 = x3 * x4
        p1234 = p12 * p34

        e = f 0

        a = - f1 / d1 + f2 / d2 - f3 / d3 + f4 / d4 + e / p1234

        b =   f1 * (x2 + x3 + x4) / d1
            - f2 * (x4 + x3 + x1) / d2
            + f3 * (x4 + x2 + x1) / d3
            - f4 * (x3 + x2 + x1) / d4
            - e * (x1 + x2 + x3 + x4) / p1234

        c = - f1 * (p34 + p24 + p23) / d1
            + f2 * (p34 + p14 + p13) / d2
            - f3 * (p24 + p14 + p12) / d3
            + f4 * (p23 + p13 + p12) / d4
            + e * (p34 + p24 + p14 + p23 + p13 + p12) / p1234

        d =   f1 * x2 * p34 / d1
            - f2 * x3 * p14 / d2
            + f3 * x4 * p12 / d3
            - f4 * x1 * p23 / d4
            - e * (x2 * p34 + x3 * p14 + x4 * p12 + x1 * p23) / p1234
    in
        [e, d, c, b, a]

-- | the real roots of quadratic equation: a x^2 + b x + c = 0.
--
-- If the root x0 is complex and |Im(x0)| < cut, it returns Re(x0).
quadEqSol :: Double  -- ^ the coefficient a
          -> Double  -- ^ the coefficient b
          -> Double  -- ^ the coefficient c
          -> Double  -- ^ the cut value for complex roots
          -> Maybe (Double, Double)
quadEqSol a b c epsScale = do
    let d = b * b - 4 * a * c
        eps = 1.0e-12
    if d >= 0
        -- then return ((-b + sqrt d) / (2*a), (-b - sqrt d) / (2*a))
        then do let q = -0.5 * (b + signum b * sqrt d)
                return (q / (a + eps), c / (q  + eps))
        else do let r = sqrt $ c / (a + eps)
                    th = acos $ - b / (2 * sqrt (a * c) + eps)
                    x = r * cos th
                if abs (r * sin th) < epsScale
                    then return (x, x)
                    else Nothing


-- might not work if module is not installed, possible fix is to place the file in
-- hnn-master root folder

import AI.HNN.FF.ComplexNetwork
import Numeric.LinearAlgebra
import Data.Complex
import System.Random.MWC
import Foreign.Storable           (Storable)

samplesXOR :: Samples (Complex Double)
samplesXOR = [ (fromList [(-1):+ (-1)], fromList [1:+0])
          , (fromList [(-1):+1], fromList [0:+0])
          , (fromList [1:+(-1)], fromList [1:+1])
          , (fromList [1:+1], fromList [0:+1])
          ]

real' :: (Floating a) => Complex a -> Complex a
real' (x :+ y) = x :+ 0

mapComplex ::  (Complex a -> a) -> (Complex a -> a) -> Complex a -> Complex a
mapComplex f g z = (f z) :+ (g z)

quarters :: (Floating a, Ord a, RealFloat a) => Complex a -> Complex a
quarters (x :+ y) = (if x > 0 then 1 else 0) :+ (if y > 0 then 1 else 0)

makeSamples :: (Floating a, Variate a, Storable a) => Int -> (Complex a -> Complex a) -> IO (Samples (Complex a))
makeSamples 0 _ = return []
makeSamples n f = do
  z <- randComplex
  lst <- makeSamples (n-1) f
  return ((fromList [z], fromList [f(z)]) : lst)

-- uncomment relevant sections for different tests
main :: IO ()
main = do
  -- Complex XOR:
  --mapM_ (putStrLn . show ) samplesXOR
  --n <- createComplexNetwork 1 [2] 1 :: IO (ComplexNetwork (Complex Double))
  --putStrLn $ show n
  --putStrLn "------------------"
  --let n' = trainNTimes 1000 0.8 complexSigmoid complexSigmoid' n samplesXOR
  --mapM_ (putStrLn . show . output n' complexSigmoid . fst) samplesXOR
  --putStrLn "------------"
  --putStrLn $ show n'

  -- Quarters on plane:
  --samples <- makeSamples 50 quarters
  --test <- makeSamples 10 quarters
  --mapM_ (putStrLn . show ) test
  --n <- createComplexNetwork 1 [2] 1 :: IO (ComplexNetwork (Complex Double))
  --putStrLn $ show n
  --putStrLn "------------------"
  --let n' = trainNTimes 1000 0.8 complexSigmoid complexSigmoid' n samples
  --mapM_ (putStrLn . show . output n' complexSigmoid . fst) test
  --putStrLn "------------"
  --putStrLn $ show n'

  -- Parabolas on plane:
  let h1 (x :+ y) = if x > 0.5*x^2 -0.5 then 1 else 0
  let h2 (x :+ y) = if y > -0.5*x^2 +0.5 then 1 else 0

  samples <- makeSamples 1000 (mapComplex h1 h2)
  test <- makeSamples 10 (mapComplex h1 h2 )

  mapM_ (putStrLn . show ) test
  n <- createComplexNetwork 1 [20] 1 :: IO (ComplexNetwork (Complex Double))
  let n' = trainNTimes 1 0.5 complexSigmoid complexSigmoid' n samples

  mapM_ (putStrLn . show . output n' complexSigmoid . fst) test


{-# LANGUAGE FlexibleInstances #-}
{-|
Module      : Algorithm.ShapeConstraint
Description : evaluation of shape-constraints for symbolic regression
Copyright   : (c) Fabricio Olivetti de Franca, 2022
License     : GPL-3
Maintainer  : fabricio.olivetti@gmail.com
Stability   : experimental
Portability : POSIX

This package provides support functions to evaluate different
shape-constraints for symbolic regression using modal arithmetic.

See https://direct.mit.edu/evco/article/30/1/75/99840 for more details.

@
@article{kronberger2022shape,
  title={Shape-Constrained Symbolic Regression—Improving Extrapolation with Prior Knowledge},
  author={Kronberger, Gabriel and de Fran{\c{c}}a, Fabricio Olivetti and Burlacu, Bogdan and Haider, Christian and Kommenda, Michael},
  journal={Evolutionary computation},
  volume={30},
  number={1},
  pages={75--98},
  year={2022},
  publisher={MIT Press One Rogers Street, Cambridge, MA 02142-1209, USA journals-info~…}
}
@

-}
module Algorithm.ShapeConstraint
    ( getViolationFun
    , Evaluator(..)
    , Shape(..)
    )
      where

import Numeric.ModalInterval
import Numeric.ModalInterval.Algorithms
import Data.SRTree

import Data.Map.Strict (Map(..))
import qualified Data.Map.Strict as M
import qualified Data.Vector as V
import Data.Bool (bool)
import Data.Maybe (fromMaybe)
import qualified Numeric.LinearAlgebra.Data as LA

-- | Data type describing the type of constraint:
--
-- * Range of the function
-- * Range of a partial derivative
-- * Monotonicity (non-increasing and non-decreasing)
-- * Partial monotonicity
-- * Inflection, Convex or Concave
--
-- The monotonicity constraints require the index of the
-- variable, the ranges require a tuple of the minimum and
-- maximum values, the partial monotonicity also requires
-- the range of the domain that has this property,
-- inflection, convexity and concavity requires the first
-- and second index of the partial derivatives.
data Shape   = Range (Double, Double)                     -- ^ f(x) \in [a, b]
             | DiffRng Int (Double, Double)               -- ^ d f(x) / dx \in [a, b]
             | NonIncreasing Int                          -- ^ d f(x) / dx \in [-inf, 0]
             | NonDecreasing Int                          -- ^ d f(x) / dx \in [0, inf]
             | PartialNonIncreasing Int (Double, Double)  -- ^ d f(x) / dx \in [-inf, 0], a <= x <= b
             | PartialNonDecreasing Int (Double, Double)  -- ^ d f(x) / dx \in [0, inf], a <= x <= b
             | Inflection Int Int                         -- ^ d^2 f(x)/dx^2 == 0
             | Convex Int Int                             -- ^ d^2 f(x)/dx^2 \in (0, Infinity)
             | Concave Int Int                            -- ^ d^2 f(x)/dx^2 \in (-Infinitiy,0)
                    deriving (Show, Read)

-- inner (a, b) no sharper than that
-- outer (a, b) no wider than that

-- | The `Evaluator` indicates the algorithm used to calculate the constraint:
--
-- * InnerInterval can result in false positives (i.e., it says it is feasible but it is not)
-- * OuterInterval can result in false negatives (i.e., it says it is infeasible but it is not)
-- * Sampling can result in false positives and usually requires a lot of samples
-- * Hybrid - not yet implemented
-- * Bisection can improve the results of OuterInterval
data Evaluator = InnerInterval | OuterInterval | Sampling Int | Hybrid | Bisection Int
                   deriving (Show, Read)

-- | Function that calculates how much of the constraints a symbolic tree vioaltes
type ConstraintFun = SRTree Int Double -> Double

-- | A list of domains of the variables
type Domains       = [(Double, Double)]

-- | A list of domains using Modal Arithmetic
type KDomains      = Map (Kaucher Double)

-- | Calculates the partial derivatives of a tree when needed.
ofShape :: (Floating a, Eq a, OptIntPow a) => SRTree Int a -> Shape -> SRTree Int a
t `ofShape` Range _ = t
t `ofShape` DiffRng ix _ = simplify $ deriveBy ix t
t `ofShape` NonIncreasing ix = simplify $ deriveBy ix t
t `ofShape` NonDecreasing ix = simplify $ deriveBy ix t
t `ofShape` PartialNonIncreasing ix _ = simplify $ deriveBy ix t
t `ofShape` PartialNonDecreasing ix _ = simplify $ deriveBy ix t
t `ofShape` Inflection ix iy = simplify $ deriveBy iy $ simplify $ deriveBy ix t
t `ofShape` Convex ix iy = simplify $ deriveBy iy $ simplify $ deriveBy ix t
t `ofShape` Concave ix iy = simplify $ deriveBy iy $ simplify $ deriveBy ix t
{-# INLINE ofShape #-}

toTuple :: Kaucher Double -> (Double, Double)
toTuple k = (fromMaybe (-1/0) $ inf k, fromMaybe (1/0) $ sup k)
{-# INLINE toTuple #-}

eps :: Double
eps = 1e-6
{-# INLINE eps #-}

calcPenalty :: Bool -> Double -> Double
calcPenalty b pnlty = bool 0 pnlty b
{-# INLINE calcPenalty #-}

inRng (a, b) (x, y) = calcPenalty (x < a - eps) (a - x) + calcPenalty (y > b + eps) (y - b)
{-# INLINE inRng #-}
isPositive (x, y)   = calcPenalty (x <= -eps) (abs x) + calcPenalty (y <= -eps) (abs y)
{-# INLINE isPositive #-}
isNegative (x, y)   = calcPenalty (x >= eps) (abs x) + calcPenalty (y >= eps) (abs y)
{-# INLINE isNegative #-}
isZero (x, y)       = calcPenalty (x <= -eps) (abs x) + calcPenalty (y >= eps) y
{-# INLINE isZero #-}

-- | Evaluates the constraint
evalConstraints :: Shape -> (Double, Double) -> Double
evalConstraints (Range ab) xy                 = inRng ab xy
evalConstraints (DiffRng _ ab) xy             = inRng ab xy
evalConstraints (NonIncreasing _) xy          = isNegative xy
evalConstraints (NonDecreasing _) xy          = isPositive xy
evalConstraints (PartialNonIncreasing _ _) xy = isNegative xy
evalConstraints (PartialNonDecreasing _ _) xy = isPositive xy
evalConstraints (Inflection _ _) xy           = isZero xy
evalConstraints (Convex _ _) xy               = isPositive xy
evalConstraints (Concave _ _) xy              = isNegative xy
{-# INLINE evalConstraints #-}

-- | Convert the domains for the partial monotonicity constraint
convertDomains :: Shape -> Map Int (Kaucher Double) -> Map Int (Kaucher Double)
convertDomains (PartialNonIncreasing ix (a, b)) ds = M.insert ix (a <.< b) ds
convertDomains (PartialNonDecreasing ix (a, b)) ds = M.insert ix (a <.< b) ds
convertDomains _ ds = ds
{-# INLINE convertDomains #-}

-- | Sum a list of constraint violations
sumCnstr :: (SRTree Int (Kaucher Double) -> Map Int (Kaucher Double) -> Kaucher Double) -> Map Int (Kaucher Double) -> [SRTree Int (Kaucher Double)] -> [Shape] -> Double -> Double
sumCnstr f ds [] []         acc = acc
sumCnstr f ds (x:xs) (z:zs) acc = sumCnstr f ds xs zs (acc + cnstr)
   where
     cnstr = evalConstraints z $ toTuple $ f x $ convertDomains z ds
{-# INLINE sumCnstr #-}

toMap :: Domains -> Map Int (Kaucher Double)
toMap    = M.fromList . zip [0..] . map (uncurry (<.<))
{-# INLINE toMap #-}

-- | Returns a constraint function given an `Evaluator` a list of `Shape` and the
-- `Domains` of the input variables.
getViolationFun :: Evaluator -> [Shape] -> Domains -> ConstraintFun
getViolationFun InnerInterval shapes domains t = sumCnstr innerApprox ds ts shapes 0.0
  where
    t'       = fmap singleton t
    ts       = map (t' `ofShape`) shapes
    ds       = toMap domains

getViolationFun OuterInterval shapes domains t = sumCnstr outerApprox ds ts shapes 0.0
  where
    t'       = fmap singleton t
    ts       = map (t' `ofShape`) shapes
    ds       = toMap domains

getViolationFun (Sampling nSamples) shapes domains t = go shapes ts samples 0.0
  where
    ds        = toMap domains
    samples   = map (M.fromList . zip [0..] . makeSamples nSamples . (`convertDomains` ds)) shapes
    getSize s = LA.size $ s M.! 0
    ts        = map (t `ofShape`) shapes

    evalSamples :: Maybe (LA.Vector Double) -> (Double, Double)
    evalSamples mxs = (fromMaybe (-1/0) mmin, fromMaybe (1/0) mmax)
      where
        mmax = LA.maxElement <$> mxs
        mmin = LA.minElement <$> mxs

    go [] _ _               acc = acc
    go (x:xs) (y:ys) (z:zs) acc = go xs ys zs (acc + cnstr)
      where
        cnstr = evalConstraints x rng
        y'    = fmap (LA.fromList . replicate (getSize z)) y
        rng   = evalSamples $ evalTreeWithMap y' z

getViolationFun e shapes domains t = error $ show e <> " evaluator not implemented"

instance OptIntPow (LA.Vector Double) where
    (^.) = (^^)
    {-# INLINE (^.) #-}

-- | Creates the samples of the `Sampling` approach
makeSamples :: Int -> Map Int (Kaucher Double) -> [LA.Vector Double]
makeSamples n domains = LA.toColumns $ LA.fromLists $ map (map midpoint) $ head $ dropWhile (\x -> length x < n) samples
  where
    step x  = let (a,b) = bisect x in [a,b]
    f       = map (concatMap step)
    samples = map sequence $ iterate f $ map (pure.snd) $ M.toAscList domains
{-# INLINE makeSamples #-}


{-# LANGUAGE DataKinds #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE PolyKinds #-}
{-# LANGUAGE TupleSections #-}
module Main where

import Map
import Sector
-- import RepaMap
import Map.PixelMap
-- import SectorTree
-- import SectorMap
import Map.Dimension

import Linear.V as L
import Data.Complex
import qualified Data.Vector as V
import Codec.Picture
-- import Data.Kind

-- mandelbrot :: Int -> Complex Double -> Bool
{-# INLINE mandelbrot #-}
mandelbrot :: (Ord t, Num t, RealFloat a) => t -> Complex a -> Bool
mandelbrot n x = mandelbrot' x x n where
  mandelbrot' z _ _ | (sqr . magnitude $ z) > 4 = False
  mandelbrot' _ _ i | i <= 0 = True
  mandelbrot' z c i = mandelbrot' (z*z + c) c (i - 1)
  sqr a = a * a

{-# INLINE toComplexCoords #-}
toComplexCoords :: V 2 Double -> Complex Double
toComplexCoords (V v) = v V.! 1 :+ v V.! 0

{-# INLINE mandelMap #-}
mandelMap :: Int -> PixelMap Double
mandelMap n = fmap toGrayScale
            . changeCoordinates toComplexCoords
            $ getPoint >>= return . mandelbrot n

{-# INLINE mandelMap' #-}
mandelMap' :: Int -> Sector 2 Double -> Int -> PixelMap Double
mandelMap' n = optimizeMap (mandelMap n)

viewSector :: Zone -> Double -> Sector 2 Double
viewSector (Zone x y r) ratio = V . V.fromList $ [(x - r, x + r), (y - r * ratio, y + r * ratio)]

data Zone = Zone { zoneX :: Double
                 , zoneY :: Double
                 , zoneR :: Double }

interestingZones :: [Zone]
interestingZones =
  [ Zone 0 0 2
  , Zone 0.1 (-0.7) (5e-4)
  , Zone (0.109)  (-0.7443) (0.005)
  -- , Zone (0.1102) (-0.7463) (0.005)
  , Zone (0.1127) (-0.7453) (6.5e-4)
  , Zone (1.0405) (-0.16)   (0.026)
  ]

main :: IO ()
main = do
  let resX          = 1920 * resMultiplier -- x and y dimensions
      resY          = 1080 * resMultiplier
      resMultiplier = 1 -- increment this to get 4k, 8k, etc.
      ratio         = (fromIntegral resX) / (fromIntegral resY)
      zone          = Zone 0 0 2
      depth         = 400
      sec           = viewSector zone ratio
      img           = bakePixelMap
                        sec
                        (resolution resX resY)
                        (mandelMap depth)
  putStrLn "Execution started."
  savePngImage ("/tmp/mandelbrot.png") img
  putStrLn "Execution finished"