{-# LANGUAGE FlexibleContexts #-}

module AI.Util.Matrix where

import Control.Monad.Random hiding (fromList)
import Foreign.Storable (Storable)
import Numeric.LinearAlgebra

-- |Return the size of a matrix as a 2-tuple.
size :: Matrix a -> (Int,Int)
size x = (rows x, cols x)

-- |Concatenate matrices horizontally.
horzcat :: Element a => [Matrix a] -> Matrix a
horzcat = fromBlocks . return

-- |Concatenate matrices vertically.
vertcat :: Element a => [Matrix a] -> Matrix a
vertcat = fromBlocks . map return

-- |Add a column of ones to a matrix.
addOnes :: Matrix Double -> Matrix Double
addOnes x = fromBlocks [[1, x]]

-- |Create a row matrix.
row :: [Double] -> Matrix Double
row = asRow . fromList

-- |Create a column matrix.
column :: [Double] -> Matrix Double
column = asColumn . fromList

--------------------------
-- Functions on Vectors --
--------------------------

takeVector :: Storable a => Int -> Vector a -> Vector a
takeVector n v = subVector 0 n v

dropVector :: Storable a => Int -> Vector a -> Vector a
dropVector n v = subVector n (dim v - n) v

sumVector :: (Num a, Storable a) => Vector a -> a
sumVector xs = foldVector (+) 0 xs

prodVector :: (Num a, Storable a) => Vector a -> a
prodVector xs = foldVector (*) 1 xs

---------------------------
-- Functions on Matrices --
---------------------------

mapRows :: Element a => (Vector a -> b) -> Matrix a -> [b]
mapRows f m = map f (toRows m)

mapCols :: Element a => (Vector a -> b) -> Matrix a -> [b]
mapCols f m = map f (toColumns m)

eachRow :: (Element a, Element b) => (Vector a -> Vector b) -> Matrix a -> Matrix b
eachRow f = fromRows . mapRows f

eachCol :: (Element a, Element b) => (Vector a -> Vector b) -> Matrix a -> Matrix b
eachCol f = fromColumns . mapCols f

sumRows :: (Element a, Num (Vector a)) => Matrix a -> Vector a
sumRows m = sum $ toRows m

sumCols :: (Element a, Num (Vector a)) => Matrix a -> Vector a
sumCols m = sum $ toColumns m

sumMatrix :: Matrix Double -> Double
sumMatrix = sumVector . sum . toRows

------------------------
-- Subset Referencing --
------------------------

subRefVec :: Storable a => Vector a -> [Int] -> Vector a
subRefVec v is = fromList $ map (v@>) is

subRefRows :: Element a => Matrix a -> [Int] -> Matrix a
subRefRows m is = fromRows $ map (r!!) is where r = toRows m

subRefCols :: Element a => Matrix a -> [Int] -> Matrix a
subRefCols m is = fromColumns $ map (c!!) is where c = toColumns m


module Template where

import Data.Complex

-------------------------------------------------------------------------------
-- Utility functions

-- | Convert an 'Int' to a 'Double'
i2d :: Int -> Double
i2d = fromIntegral

-- | Convert a 'Double' to a 'Complex Double' with no imaginary part.
d2rc :: Double -> Complex Double
d2rc a = (a :+ 0.0)

-- | Convert a 'Double' to a 'Complex Double' with no real part.
d2ic :: Double -> Complex Double
d2ic a = (0.0 :+ a)

-- | A signal is a list of double precision samples
type Signal = [Double]

-- | Complete the definition of Fourier
data Fourier =  Fourier [Complex Double]

mkFourier :: [Complex Double] -> Fourier
mkFourier cl = (Fourier cl)

unFourier :: Fourier -> [Complex Double]
unFourier (Fourier cl) = cl

-------------------------------------------------------------------------------
-- Almost Equal

{-class AlmostEq a where
  (~=) :: a -> a -> Bool

infix 4 ~=

instance AlmostEq Double where
  a ~= b = abs(a-b) < 1e-14

instance AlmostEq (Complex Double) where
  a ~= b = magnitude (a-b) < 1e-14

instance AlmostEq a => AlmostEq [a] where
  xs ~= ys = (length xs == length ys) && equ xs ys where
    equ [] [] = True
    equ (x:xs) (y:ys) = x ~= y && equ xs ys

instance AlmostEq Fourier where
  xs ~= ys = equf xs ys where
    equf (Fourier a) (Fourier b) = a ~= b-}

-------------------------------------------------------------------------------
-- Naive DFT

-- | Naive discrete Fourier transform
dft :: Signal -> Fourier
dft [] = (Fourier [])
dft l = (Fourier (makehead : (maketail len))) where
  makehead = d2rc (sum l) --n=0
  len = (length l-1) --N-1
  maketail 0 = [] --n!=0
  maketail n = eoc : maketail (n-1) where
     eoc = d2rc (l !! 1) +  sum ([(d2rc (l !! k)) *  (exp (d2ic ((-2) * pi * (((i2d k) * (i2d n))/(i2d len))))) | k<-[1..len]])

makeTerm::Int->[Complex Double]->Int-> Double
makeTerm bn xs k= (1/(i2d bn)) * sum [realPart $ (xs !! n) * (exp $ d2ic $ pi * (i2d (2*n*k)) /(i2d bn)) | n <- [0..(bn-1)]]

-- | Naive inverse discrete Fourier transform
idft :: Fourier -> Signal
idft (Fourier list) = map (makeTerm (length list) list) [0..(length list -1)]

-------------------------------------------------------------------------------
-- Spectrum Analyser

rescale::[Double]->[Double]
rescale list = map (/maxi) list
  where maxi = maximum list

bar::Double->String
bar x = replicate (ceiling (x*40)) '#'

spectrumAnalyser :: Signal -> IO ()
spectrumAnalyser =   putStr . unlines . map bar . rescale . map magnitude . unFourier . dft

-------------------------------------------------------------------------------
-- Fast Fourier Transform

-- | Split a list in its even and oddly indexed elements.
scatter :: [a] -> ([a],[a])
scatter [] = ([],[])
scatter (x:y:xs) = ((x:l),(y:r)) where (l,r) = scatter xs

-- | Gather the Fourier coordinates into the correct order.
-- > gather [X0, X4, X1, X5, X2, X6, X3, X7] = [X0,X1,X2,X3,X4,X5,X6,X7]
gather :: [a] -> [a]
gather list = l ++ r where (l,r) = scatter list

-- | Combine even and odd Fourier coefficients with the twiddle factor.
twiddle :: Int -> Int -> Complex Double -> Complex Double -> [Complex Double]
twiddle m n c1 c2 = [c1 + number * c2 ,c1 - number * c2 ] where
  number = exp (d2ic ((-2)*pi*((i2d n)/(i2d m))))

-- | The butterfly step: construct Fourier coefficients from the
-- Fourier coefficients of the even an odd elements.
butterfly :: Int -> Fourier -> Fourier -> [Complex Double]
butterfly bn (Fourier even) (Fourier odd)= concat $ zipWith3 (twiddle bn)  [0..] even odd

-- | Cooley-Tukey Fast Fourier Transform algorithm.
fft :: Signal -> Fourier
fft xs
  | odd (length xs) = dft xs
  | otherwise = let (evens,odds) = scatter xs
                in Fourier $ gather $ butterfly (length xs) (fft evens) (fft odds)


module Validation
  ( trainValAndTest
  , valTrainAndTest
  , solution1_2
  , solution3_4
  ) where

import           Data.Maybe                    (fromMaybe)
import           LinearRegression
import           NonlinearTransform
import           Numeric.LinearAlgebra.Data
import           Numeric.LinearAlgebra.HMatrix as HMatrix
import           WeightDecay

-- | Make a nonlinear Transformation of the first n data points of the input data
makeTransformedFirstN k n inputdata = fmap (fmap makeMatrixX) inputdata
  where
    makeMatrixX lists =
      createTransformedXk k $
      map (fromMaybe (0, 0) . listToTuple') $ take n lists

-- | Make a nonlinear Transformation of the rest of the data points after the first n
makeTransformedLastN k n inputdata = fmap (fmap makeMatrixX) inputdata
  where
    makeMatrixX lists =
      createTransformedXk k $
      map (fromMaybe (0, 0) . listToTuple') $ drop n lists

trainXFirstn k n = makeTransformedFirstN k n traindata

trainXLastn k n = makeTransformedLastN k n traindata

testXTransformedk k = makeTransformedMatrix k testdata

trainYFirstn n = fmap (fmap makeVectorY) traindata
  where
    makeVectorY lists = vector $ take n $ map last lists

trainYLastn n = fmap (fmap makeVectorY) traindata
  where
    makeVectorY lists = vector $ drop n $ map last lists

-- | Trains and validate linear Regression from the data on the course website, using the first n points for training and the rest for validation, performing non-linear transformation with k-dimensional features
trainVal k n = do
  matrixXTrain <- trainXFirstn k n
  matrixXVal <- trainXLastn k n
  vectorYTrain <- trainYFirstn n
  vectorYVal <- trainYLastn n
  let weights = linearRegressionWeight <$> matrixXTrain <*> vectorYTrain
  let inSampleError =
        linRegClassificationError <$> matrixXTrain <*> vectorYTrain <*> weights
  let validationError =
        linRegClassificationError <$> matrixXVal <*> vectorYVal <*> weights
  return (weights, validationError, inSampleError)

-- | same as trainVal, but the first n points are used for validation and the rest for training
valTrain k n = do
  matrixXVal <- trainXFirstn k n
  matrixXTrain <- trainXLastn k n
  vectorYVal <- trainYFirstn n
  vectorYTrain <- trainYLastn n
  let weights = linearRegressionWeight <$> matrixXTrain <*> vectorYTrain
  let inSampleError =
        linRegClassificationError <$> matrixXTrain <*> vectorYTrain <*> weights
  let validationError =
        linRegClassificationError <$> matrixXVal <*> vectorYVal <*> weights
  return (weights, validationError, inSampleError)

-- | Tests linear Regression from the data on the course website
testLinRegTransformedk k weights = do
  matrixXTest <- testXTransformedk k
  vectorYTest <- testY
  let outOfSampleError =
        linRegClassificationError <$> matrixXTest <*> vectorYTest <*> weights
  return outOfSampleError

-- | train, validate and test with given data, with nonlinear transformation to k-dimensional feature space, with n triaining points and max-n validation points
trainValAndTest :: Int -> Int -> IO ()
trainValAndTest k n = do
  (weights, validationError, inSampleError) <- trainVal k n
  outOfSampleError <- testLinRegTransformedk k weights
  putStr "In Sample Error: "
  print inSampleError
  putStr "Validation Error: "
  print validationError
  putStr "Out of Sample Error: "
  print outOfSampleError

-- | same as trainValAndTest, but the first n points are used for validation and the rest for training
valTrainAndTest :: Int -> Int -> IO ()
valTrainAndTest k n = do
  (weights, validationError, inSampleError) <- valTrain k n
  outOfSampleError <- testLinRegTransformedk k weights
  putStr "In Sample Error: "
  print inSampleError
  putStr "Validation Error: "
  print validationError
  putStr "Out of Sample Error: "
  print outOfSampleError

-- | solution to problems 1 and 2
solution1_2 :: IO ()
solution1_2 =
  mapM_
    (\k -> do
       putStr " k ="
       print k
       trainValAndTest k 25)
    [3, 4, 5, 6, 7]

-- | solution to problems 3 and 4
solution3_4 :: IO ()
solution3_4 =
  mapM_
    (\k -> do
       putStr " k ="
       print k
       valTrainAndTest k 25)
    [3, 4, 5, 6, 7]


module Main where

import Numeric.LinearAlgebra
import Numeric.Morpheus

a = matrix 5 [
  71, 11, 3,  -9, -7,
  21, -7,  -2,  23, 11,
  -11, 32, 53, -49, 37,
  1,  -24, 78, 90, 17
  ]

main :: IO ()
main = do
  putStrLn "\nNumeric.Morpheus.MatrixReduce functions: "

  putStr "sum of elements of rows: "
  print $ rowSum a

  putStr "product of elements of rows: "
  print $ rowPredicate (*) a

  putStr "max values of columns: "
  print $ fst $ columnMaxIndex a

  putStr "indices of max values of columns: "
  print $ snd $ columnMaxIndex a

  putStrLn "\n\nNumeric.Morpheus.Activation functions: "

  putStrLn "\nSigmoid:"
  disp 3 (sigmoid a)

  putStrLn "ReLu:"
  disp 3 (relu a)

  putStrLn "ReLu gradient:"
  disp 3 (reluGradient a)

  putStrLn "\n\nNumeric.Morpheus.Statistics functions: "

  putStrLn "\nColumn Means:"
  let colmeans = columnMean a
  print colmeans

  putStrLn "Column Standard Deviations: "
  print (columnStddev_m colmeans a)


{-# LANGUAGE OverloadedLists #-}
module Ch4 where

import Data.List
import Control.Monad
import Numeric.LinearAlgebra
import Graphics.Gnuplot.Simple

import Perceptron
import NeuralNetwork
import qualified MNIST
import SimpleNet (SimpleNet (..))
import qualified SimpleNet

main :: IO ()
main = do
  displayFunction1
  displayTangents
  displayFunction2
  displayGradients
  displayDescents

-- ch4_3_2
function1 :: Double -> Double
function1 x = 0.01 * x * x + 0.1 * x

displayFunction1 :: IO ()
displayFunction1 = plotPath attrs $ zip xs ys
  where attrs = [(Title "Function 1"), (YRange (0, 6))]
        xs = linearScale 100 (0, 20) :: [Double]
        ys = function1 <$> xs

tangent :: (Double -> Double) -> Double -> Double -> Double
tangent f x0 x = a * (x - x0) + f x0
  where a = numericalDiff f x0

displayTangents :: IO ()
displayTangents = plotPathsStyle attrs paths
  where attrs = [(Title "Tangents of Function 1"), (YRange (-1, 6))]
        xs = linearScale 100 (0, 20) :: [Double]
        paths = do
          (fn, name) <- [ (function1, "0.01 * x^2 + 0.1 * x")
                        , (tangent function1 5, "tangent (x = 5)")
                        , (tangent function1 10, "tangent (x = 10)")
                        ]
          let style = PlotStyle Lines $ CustomStyle [(LineTitle name)]
          return $ (style, zip xs (fn <$> xs))

-- ch4_3_3
function2 :: Vector R -> Double
function2 v = dot v v

displayFunction2 :: IO ()
displayFunction2 = plotFunc3d attrs [] xs ys function2'
  where
    attrs = [Title "Function 2"]
    xs = linearScale 100 (-3, 3) :: [Double]
    ys = linearScale 100 (-3, 3) :: [Double]
    function2' x y = function2 $ vector [x, y]

displayGradients :: IO ()
displayGradients = plotListStyle attrs style vs
  where
    attrs = [Title "Gradients of Function 2"]
    style = PlotStyle Vectors $ CustomStyle [(LineTitle "gradients")]
    xs = [(-2.0), (-1.75) .. 2.0] :: [Double]
    ys = [(-2.0), (-1.75) .. 2.0] :: [Double]
    vs = do
      x <- xs
      y <- ys
      let [dx, dy] = toList $ numericalGradient function2 $ vector [x, y]
      return ((x, y), (- dx / 20, - dy / 20))

displayDescents :: IO ()
displayDescents = plotPathStyle attrs style vs
  where
    attrs = [Title "Gradient descents of Function 2", XRange (-4, 4), YRange (-4, 4)]
    style = PlotStyle Points $ CustomStyle [(LineTitle "descents")]
    v0 = vector [-3.0, 4.0]
    vs = do
      v <- take 100 $ gradientDescents function2 0.1 v0
      let [x, y] = toList v
      return (x, y)

-- ch4_4_2

functionLoss :: Vector R -> Vector R -> Matrix R -> Double
functionLoss input label w = SimpleNet.loss net' input label
  where net' = SimpleNet w

net :: SimpleNet
net = SimpleNet $ (2 >< 3) [ 0.47355232, 0.99773930, 0.84668094
                           , 0.85557411, 0.03563661, 0.69422093]

input :: Vector R
input = [0.6, 0.9]              -- input

answer :: Vector R
answer = [0, 0, 1]              -- correct answer

functionNet :: Matrix R -> Double
functionNet = functionLoss input answer

dw :: SimpleNet -> Matrix R
dw (SimpleNet w) = numericalGradient functionNet w

netDescents :: SimpleNet -> [SimpleNet]
netDescents net = iterate descend net
  where
    lr = 0.1
    descend net@(SimpleNet w) = SimpleNet $ w - scale lr (dw net)

displayLearnings :: IO ()
displayLearnings = plotPathsStyle attrs paths
  where
    attrs = [(Title "Learnings of simple net"), (YRange (0, 1))]
    xs = [0 .. 50] :: [Double]
    n = length xs
    nets = netDescents net
    outputs = map (toList . softmax . (flip SimpleNet.predict input)) nets
    paths = do
      (i, path) <- zip [0..] $ transpose $ take n outputs
      let style = PlotStyle Lines $ CustomStyle [(LineTitle (show i))]
      return (style, zip xs path)


module Numeric.Morpheus.MatrixReduceTest
(
  tests
)

where

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

import Numeric.LinearAlgebra
import Numeric.Morpheus.MatrixReduce

xRM = matrix 5 [
  71, 11, 3,  9, 7,
  21, 7,  2,  23, 11,
  11, 32, 53, 49, 37,
  1,  24, 78, 90, 17
  ]

xCM = tr xRM

columnSumRM = vector [104, 74, 136, 171, 72]
columnSumCM = vector [101, 64,182, 210]

rowSumRM = vector [101, 64,182, 210]
rowSumCM = vector [104, 74, 136, 171, 72]

columnMaxValuesRM = fromList [71, 32, 78, 90, 37]
columnMaxIndicesRM = fromList [0, 2, 3, 3, 2]
columnMaxValuesCM = fromList [71, 23, 53, 90]
columnMaxIndicesCM = fromList [0, 3, 2, 3]

columnMinValuesRM = fromList [1, 7, 2, 9, 7]
columnMinIndicesRM = fromList [3, 1, 1, 0, 0]
columnMinValuesCM = fromList [3, 2, 11, 1]
columnMinIndicesCM = fromList [2, 2, 0, 0]

rowMaxValuesRM = fromList [71, 23, 53, 90]
rowMaxIndicesRM = fromList [0, 3, 2, 3]
rowMaxValuesCM = fromList [71, 32, 78, 90, 37]
rowMaxIndicesCM = fromList [0, 2, 3, 3, 2]

rowMinValuesRM = fromList [3, 2, 11, 1]
rowMinIndicesRM = fromList [2, 2, 0, 0]
rowMinValuesCM = fromList [1, 7, 2, 9, 7]
rowMinIndicesCM = fromList [3, 1, 1, 0, 0]

assertValuesAndIndices (expectedR, expectedI) (actualR, actualI) = do
  assertVector "values" 1e-10 expectedR actualR
  assertVectorI "indices" expectedI actualI

tests = [testGroup "Column Sum" [
            testCase "row major" $ assertVector "" 1e-10 columnSumRM (columnSum xRM)
          , testCase "column major" $ assertVector "" 1e-10 columnSumCM (columnSum xCM)
          ]
        , testGroup "Row Sum" [
            testCase "row major" $ assertVector "" 1e-10 rowSumRM (rowSum xRM)
          , testCase "column major" $ assertVector "" 1e-10 rowSumCM (rowSum xCM)
          ]
        , testGroup "Column Max Index" [
            testCase "row major" $ assertValuesAndIndices (columnMaxValuesRM, columnMaxIndicesRM) (columnMaxIndex xRM)
          , testCase "column major" $ assertValuesAndIndices (columnMaxValuesCM, columnMaxIndicesCM) (columnMaxIndex xCM)
          ]
        , testGroup "Column Min Index" [
            testCase "row major" $ assertValuesAndIndices (columnMinValuesRM, columnMinIndicesRM) (columnMinIndex xRM)
          , testCase "column major" $ assertValuesAndIndices (columnMinValuesCM, columnMinIndicesCM) (columnMinIndex xCM)
          ]
        , testGroup "Row Max Index" [
            testCase "row major" $ assertValuesAndIndices (rowMaxValuesRM, rowMaxIndicesRM) (rowMaxIndex xRM)
          , testCase "column major" $ assertValuesAndIndices (rowMaxValuesCM, rowMaxIndicesCM) (rowMaxIndex xCM)
          ]
        , testGroup "Row Min Index" [
            testCase "row major" $ assertValuesAndIndices (rowMinValuesRM, rowMinIndicesRM) (rowMinIndex xRM)
          , testCase "column major" $ assertValuesAndIndices (rowMinValuesCM, rowMinIndicesCM) (rowMinIndex xCM)
          ]
        , testGroup "Column Predicate" [
            testCase "max - row major" $ assertVector "" 1e-10 columnMaxValuesRM (columnPredicate max xRM)
          , testCase "max - col major" $ assertVector "" 1e-10 columnMaxValuesCM (columnPredicate max xCM)
          ]
          , testGroup "Row Predicate" [
              testCase "min - row major" $ assertVector "" 1e-10 rowMinValuesRM (rowPredicate min xRM)
            , testCase "min - col major" $ assertVector "" 1e-10 rowMinValuesCM (rowPredicate min xCM)
            ]
        ]


--
-- Algebra: algebra library
--

module CNN.Algebra (
  vadd
, vsub
, vsum
, vscale
, vdot
, madd
, msub
, msum
, mscale
, mmul
, mdot
, mavg
, transpose
) where

import Data.List (foldl')
import Numeric.LinearAlgebra
--import Numeric.LinearAlgebra.Data

-- VECTOR FUNCTIONS

vadd :: Vector R -> Vector R -> Vector R
vadd = (+)

vsub :: Vector R -> Vector R -> Vector R
vsub = (-)

{- |
vsum

>>> let vs = [vector [1.0,2.0,3.0], vector [4.0,5.0,6.0]]
>>> vsum vs
[5.0,7.0,9.0]

-}

vsum :: [Vector R] -> Vector R
vsum [] = fromList [0.0]
vsum (v:vs) = v + (vsum vs)

vscale :: Double -> Vector R -> Vector R
vscale s a = scale s a

vdot :: Vector R -> Vector R -> Double
vdot a b = a `dot` b

-- MATRIX FUNCTIONS

{- |
madd

>>> let m1 = fromLists [[1.0,10.0,100.0],[2.0,20.0,200.0]]
>>> let m2 = fromLists [[3.0,30.0,300.0],[4.0,40.0,400.0]]
>>> madd m1 m2
(2><3)
 [ 4.0, 40.0, 400.0
 , 6.0, 60.0, 600.0 ]

-}

madd :: Matrix R -> Matrix R -> Matrix R
madd a b = a + b

{- |
msub

>>> let m1 = fromLists [[3.0,30.0,300.0],[5.0,50.0,500.0]]
>>> let m2 = fromLists [[1.0,10.0,100.0],[2.0,20.0,200.0]]
>>> msub m1 m2
(2><3)
 [ 2.0, 20.0, 200.0
 , 3.0, 30.0, 300.0 ]

-}

msub :: Matrix R -> Matrix R -> Matrix R
msub a b = a - b

{- |
msum

>>> let m1 = fromLists [[3.0,30.0,300.0],[5.0,50.0,500.0]]
>>> let m2 = fromLists [[3.0,30.0,300.0],[4.0,40.0,400.0]]
>>> let m3 = fromLists [[1.0,10.0,100.0],[2.0,20.0,200.0]]
>>> msum [m1,m2,m3]
(2><3)
 [  7.0,  70.0,  700.0
 , 11.0, 110.0, 1100.0 ]

-}

msum :: [Matrix R] -> Matrix R
--msum = foldr madd []
msum = foldl' madd (matrix 1 [0.0])

{- |
mscale

>>> let m0 = fromLists [[0.0,1.0,2.0],[3.0,4.0,5.0]]
>>> mscale 2.0 m0
(2><3)
 [ 0.0, 2.0,  4.0
 , 6.0, 8.0, 10.0 ]

-}

mscale :: Double -> Matrix R -> Matrix R
mscale s m = scale s m

{- |
mmul

>>> let v1 = vector [3.0,30.0,300.0]
>>> let m1 = fromLists [[1.0,10.0,100.0],[2.0,20.0,200.0]]
>>> mmul v1 m1
[30303.0,60606.0]

-}

mmul :: Vector R -> Matrix R -> Vector R
mmul v m = m #> v

{- |
mdot

>>> let m1 = fromLists [[1.0,2.0],[3.0,4.0]]
>>> let m2 = fromLists [[5.0,6.0],[7.0,8.0]]
>>> mdot m1 m2
70.0

-}

mdot :: Matrix R -> Matrix R -> Double
--mdot a b = sum $ concat $ toLists (a * b)
mdot a b = flatten a <.> flatten b

{- |
mavg

>>> let m1 = fromLists [[1.0,2.0],[3.0,4.0]]
>>> let m2 = fromLists [[4.0,3.0],[2.0,1.0]]
>>> mavg [m1,m2]
(2><2)
 [ 2.5, 2.5
 , 2.5, 2.5 ]
>>> let m1 = fromLists [[1.0,2.0,3.0],[4.0,5.0,6.0]]
>>> let m2 = fromLists [[6.0,5.0,4.0],[3.0,2.0,1.0]]
>>> let m3 = fromLists [[2.0,2.0,2.0],[2.0,2.0,2.0]]
>>> mavg [m1,m2,m3]
(2><3)
 [ 3.0, 3.0, 3.0
 , 3.0, 3.0, 3.0 ]

-}

mavg :: [Matrix R] -> Matrix R
mavg ms = mscale a ss
  where
    a = 1.0 / fromIntegral (length ms)
    ss = msum ms

{- |
transpose

>>> let m1 = fromLists [[1.0, 2.0], [3.0, 4.0]] :: Matrix R
>>> tr m1
(2><2)
 [ 1.0, 3.0
 , 2.0, 4.0 ]
>>> let m2 = fromLists [[1.0, 2.0, 3.0], [4.0, 5.0, 6.0], [7.0, 8.0, 9.0]] :: Matrix R
>>> tr m2
(3><3)
 [ 1.0, 4.0, 7.0
 , 2.0, 5.0, 8.0
 , 3.0, 6.0, 9.0 ]
>>> let m3 = fromLists [[1.0, 2.0, 3.0, 10.0], [4.0, 5.0, 6.0, 11.0], [7.0, 8.0, 9.0, 12.0]] :: Matrix R
>>> tr m3
(4><3)
 [  1.0,  4.0,  7.0
 ,  2.0,  5.0,  8.0
 ,  3.0,  6.0,  9.0
 , 10.0, 11.0, 12.0 ]

-}

transpose :: Eq a => [[a]] -> [[a]]
transpose fs
  | null l    = []
  | otherwise = l : transpose fs'
  where
    (ls, fs') = tr fs
    l = concat ls
    tr :: [[a]] -> ([[a]], [[a]])
    tr fs = unzip $ map (splitAt 1) fs


module Main where

import Svg
import Data.Complex
import Dft
import Data.Foldable
import System.Environment
import Data.Maybe
import Util
import qualified Data.Vector as V
import Numeric.FFT.Vector.Invertible

_N = 256 :: Int
_N_Minus1 = _N - 1 :: Int

square :: (Enum a, Floating a) => Int -> [a]
square n = [sum [sin ((fromIntegral (2 * (k * 2 + 1) * x)) * pi / (fromIntegral _N)) / (fromIntegral (k * 2 + 1)) | k <- [0..n-1]] | x <- [0.._N_Minus1]]

sawtooth = 0 : [1 - (2 / (fromIntegral _N)) * (fromIntegral x) | x <- [1.._N_Minus1]]

approxSawtooth :: (Enum a, Floating a) => Int -> [a]
approxSawtooth n = [foldr1 (+) [sin (fromIntegral (2 * k * x) * pi / (fromIntegral _N)) / (fromIntegral k) | k <- [1..n]] | x <- [0.._N_Minus1]]

signal = approxSawtooth 10

realToComplex :: (Floating a) => a -> Complex a
realToComplex a = a :+ 0

freqDomain :: [Complex Double]
freqDomain = let timeDomain = (V.fromList signal) in
  V.toList $ run dftR2C timeDomain

freqDomainAmp :: [Double]
freqDomainAmp = let (h : t) = freqDomain
                    (m, l) = splitAt (length t - 1) t in
  (magnitude h / 2) : (map magnitude m) ++ (map ((/2) . magnitude) l)

prefix :: Maybe String
prefix = Nothing

main :: IO ()
main = do
  args <- getArgs
  let f = (read . head) args
      graphStroke = Stroke (Just "black") (Just 0.03)
      graphScale = 8 / (fromIntegral _N)
      dasharray = fromMaybe "2" ((show . (*2)) <$> strokeWidthMaybe graphStroke) ++
                  fromMaybe " 1" (((' ' :) . show) <$> strokeWidthMaybe graphStroke)
      helperStroke = strokeWidthMap (/4) graphStroke
      signalGraphTransform = "translate(-4 3)"
  writeFile "background.svg" (
    prologue prefix "xlink" (Just (WidthHeight "600" "600")) (Just $ ViewBox "-5" "-5" "10" "10") ++
        "<" ++ maybePrefix prefix ++ "g transform=\"" ++ signalGraphTransform ++ "\">" ++
        foldr1 (++) (map (\i -> let x = fromIntegral (i * _N) / (fromIntegral f) in
                           "<" ++ maybePrefix prefix ++
                           "line" ++ strokeToString helperStroke ++
                           " stroke-dasharray=\"" ++
                           dasharray ++
                           "\" x1=\"0\" y1=\"1\" x2=\"0\" y2=\"-1\"" ++
                           " transform=\"translate(" ++ show (x * graphScale) ++ " 0)\" />") [0..f]) ++
        "<" ++ maybePrefix prefix ++ "line" ++ strokeToString helperStroke ++
        " stroke-dasharray=\"" ++ dasharray ++ "\" x1=\"" ++ show (-5 * graphScale) ++ "\" y1=\"0\"" ++
        " x2=\"" ++ show ((fromIntegral (_N + 5)) * graphScale) ++ "\" y2=\"0\" " ++
        " />" ++
        drawGraph prefix graphScale 1.0 graphStroke (zip [0..(fromIntegral _N_Minus1)] signal) ++
        "</" ++ maybePrefix prefix ++ "g>" ++
        "<" ++ maybePrefix prefix ++ "g transform=\"translate(-4 -4)\">" ++
        drawGraph prefix (graphScale * 2) freqDomainScale graphStroke (zip [0..(fromIntegral _N_Minus1)] freqDomainAmp) ++
        "<" ++ maybePrefix prefix ++ "circle cx=\"" ++ show ((fromIntegral f) * graphScale * 2) ++ "\" cy=\"" ++
        show ((freqDomainAmp !! f) * freqDomainScale)  ++ "\" r=\"" ++
        fromMaybe "1" (show <$> strokeWidthMaybe graphStroke) ++ "\"" ++
        (fromMaybe "" ((\s -> " fill=\"" ++ s ++ "\"") <$> (strokeColorMaybe graphStroke))) ++ "/>" ++
        "</" ++ maybePrefix prefix ++ "g>" ++
        epilogue prefix
    )
  for_ (zip3 (animateFreq (map realToComplex signal) f (_N `div` f)) [0..] signal)
    (\(content, num, s) ->
      writeFile ("graph" ++ show num ++ ".svg") (
        prologue prefix "xlink" (Just (WidthHeight "600" "600")) (Just $ ViewBox "-5" "-5" "10" "10") ++
        content ++
        "<" ++ maybePrefix prefix ++ "g transform=\"" ++ signalGraphTransform ++ "\">" ++
        "<" ++ maybePrefix prefix ++ "circle cx=\"" ++ show ((fromIntegral num) * graphScale) ++ "\" cy=\"" ++ show s ++ "\" r=\"" ++
        fromMaybe "1" (show <$> strokeWidthMaybe graphStroke) ++ "\"" ++
        (fromMaybe "" ((\s -> " fill=\"" ++ s ++ "\"") <$> (strokeColorMaybe graphStroke))) ++ "/>" ++
        "</" ++ maybePrefix prefix ++ "g>" ++
        epilogue prefix))

freqDomainScale :: RealFloat a => a
freqDomainScale = 1 / (2 * sqrt (fromIntegral (length signal)))

animateFreq :: (RealFloat a, Show a) => [Complex a] -> Int -> Int -> [String]
animateFreq signal f sumTraceN = let
  functionValues = ((map (\(s, wp, sum) -> (s, wp, sum * (freqDomainScale :+ 0)))) . tail . reverse) (correlateSignalAtFreq signal f)
  lastN = lastNWindow sumTraceN functionValues in
  map partToString (zip functionValues lastN) where
    partToString ((s, wavePart, sum), lastNTuples) =
      let opacities = [1.0 - (fromIntegral i) / (fromIntegral sumTraceN)  | i <- [0..sumTraceN - 1]]
          lastComplexes = zip lastNTuples opacities
          lastComplexesDraws =
            map
            (\((_, _, a :+ b), opacity) ->
              "<" ++ maybePrefix prefix ++ "circle fill=\"blue\" cx=\"" ++ show a ++
              "\" cy=\"" ++ show b ++ "\" r=\"0.01\" opacity=\"" ++ show opacity ++
              "\" />"
            )
            lastComplexes
      in
      drawComplex sum prefix (Stroke (Just "blue") (Just 0.03)) ++
      foldr1 (++) lastComplexesDraws ++
      drawComplex wavePart prefix (Stroke (Just "green") (Just 0.02)) ++
      drawComplex s prefix (Stroke (Just "red") (Just 0.01))


-- The magic of Lie Algebra

import Numeric.LinearAlgebra

disp = putStrLn . dispf 5

rot1 :: Double -> Matrix Double
rot1 a = (3><3)
 [ 1, 0, 0
 , 0, c, s
 , 0,-s, c ]
    where c = cos a
          s = sin a

g1,g2,g3 :: Matrix Double

g1 = (3><3) [0, 0,0
            ,0, 0,1
            ,0,-1,0]

rot2 :: Double -> Matrix Double
rot2 a = (3><3)
 [ c, 0, s
 , 0, 1, 0
 ,-s, 0, c ]
    where c = cos a
          s = sin a

g2 = (3><3) [ 0,0,1
            , 0,0,0
            ,-1,0,0]

rot3 :: Double -> Matrix Double
rot3 a = (3><3)
 [ c, s, 0
 ,-s, c, 0
 , 0, 0, 1 ]
    where c = cos a
          s = sin a

g3 = (3><3) [ 0,1,0
            ,-1,0,0
            , 0,0,0]

deg=pi/180

-- commutator
infix 8 &
a & b = a <> b - b <> a

infixl 6 |+|
a |+| b = a + b + a&b /2  + (a-b)&(a & b) /12

main = do
   let a =  45*deg
       b =  50*deg
       c = -30*deg
       exact = rot3 a <> rot1 b <> rot2 c
       lie = scalar a * g3 |+| scalar b * g1 |+| scalar c * g2
   putStrLn "position in the tangent space:"
   disp lie
   putStrLn "exponential map back to the group (2 terms):"
   disp (expm lie)
   putStrLn "exact position:"
   disp exact


{-# language TypeApplications #-}
{-# language FlexibleContexts #-}
{-# language DerivingVia #-}
{-# language StandaloneDeriving #-}
{-# language TypeFamilies #-}
{-# language ConstrainedClassMethods #-}

module Geometry.Figures
  ( origin
  , aPoint, aLabel
  , point, point'
  , pointOn, projectOn, intersectionPoints, closestTo
  , aligned
  , aLine, aRay, aSegment, oX, oY
  , line, line', ray, ray', segment, segment'
  , extendToLength, extendTo, normalSegment, heightFrom, clipBy
  , anAngle
  , angleBetween, angleWithin, innerAngle, bisectrisse
  , supplementary, vertical, reflex
  , aTriangle, triangle2a, triangle3s, aRightTriangle
  , aSquare, aRectangle, space
  , parametricPoly, polarPoly, regularPoly
  , aCircle, circle, circle'
  , linearScale, modularScale
  , translate, scaleAt, scaleXAt, scaleYAt, scaleFig
  , at', at, along', along, through', through
  , on, normalTo, flipAt
  , vertexAngle, altitude, median
 ) where

import Data.Complex
import Data.Maybe
import Data.List.Extra (minimumOn, sortOn)
import Control.Applicative
import Data.Fixed (mod')

import Geometry.Base
import Geometry.Point
import Geometry.Circle
import Geometry.Line
import Geometry.Angle
import Geometry.Polygon
import Geometry.Plot
import Geometry.Intersections
import Geometry.Decorations

------------------------------------------------------------
-- | The origin point. Equivalent to `aPoint`.
origin :: Point
origin = aPoint

-- | The x-axis. Equivalent to `aLine`.
oX = aLine

-- | The y-axis
oY = oX # rotate 90

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

-- | A point at the origin. Equivalent to `origin`.
aPoint :: Point
aPoint = asCmp 0

-- | A label: the invisible point which could be labeled using `#:` operator.
aLabel :: Label
aLabel = mkLabel origin

{- | The constructor for a point with given coordinates.

>>> point (1,2)
<Point (1.0 2.0)>
-}
point :: XY -> Point
point = point'

{- | The generalized version of `point`.

> point' 0 #: "O" <+>
> point' (45 :: Direction) #: "A" <+>
> point' (1 :: Cmp) #: "B"
<< figs/points.svg >>
-}
point' :: Affine a => a -> Point
point' p = mkPoint (cmp p)

{- | Returns the point on a given curve.

> let c = aCircle
> in c <+>
>    pointOn c 0 #: "A" <+>
>    pointOn c 0.25 #: "B" <+>
>    pointOn c 0.667 #: "C"
<< figs/pointOn.svg >>
-}
pointOn :: Curve c => c -> Double -> Decorated Point
pointOn c t = mkPoint (c @-> t) #: loffs (cmp (normal c t))

{- | Returns a normal projection of the given point on the curve.

> let c = Plot (\t -> (t, sin t)) (0,6) # asPolyline
>     pA = point (1,0) #: "A"
>     pB = point (2,0) #: "B"
>     pC = point (4,0) #: "C"
>     pD = point (6,1) #: "D"
> in c <+> (pA <+> pA # projectOn c #: "A'" <+>
>           pB <+> pB # projectOn c #: "B'" <+>
>           pB <+> pC # projectOn c #: "C'" <+>
>           pD <+> pD # projectOn c #: "D'")
<< figs/projectOn.svg >>
-}
projectOn :: (APoint p, Curve c, Affine p) => c -> p -> Maybe (Decorated Point)
projectOn c p = pointOn c <$> (p ->@? c)

{- | Returns a list of intersection points as `Point` objects.

> let p1 = regularPoly 7
>     c = aCircle
> in p1 <+> c <+> intersectionPoints c p1
<< figs/intersectionPoints.svg >>
-}
intersectionPoints :: ( Curve a, Curve b, Intersections a b )
                   => a -> b -> [Point]
intersectionPoints c1 c2 = point' <$> intersections c1 c2

-- | Returns a point from a list which is closest to a given one.
closestTo :: (Metric a, Affine a, Metric b, Affine b) => a -> [b] -> Maybe b
closestTo p [] = Nothing
closestTo p ps  = Just $ minimumOn (distance p) ps

{- | Returns @True@ if all points in a given list belong to one line
-}
aligned :: (AlmostEq a, Metric a, Affine a) => [a] -> Bool
aligned [] = False
aligned [_] = True
aligned [_, _] = True
aligned [x, y, z] | x ~= y = aligned [x, z]
                  | x ~= z || y ~= z = aligned [x, y]
                  | otherwise = line' x y `isContaining` z
aligned (x : y : ps) = (\p -> aligned [x, y, p]) `all` ps

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

-- | The generalized version of `circle`.
circle' :: Affine a => Double -> a -> Circle
circle' r p = mkCircle r (cmp p)

-- | The constructor for a circle with given radius and coordinates.
circle :: Double -> XY -> Circle
circle = circle'

-- | The template for a circle.
aCircle :: Circle
aCircle = circle' 1 origin

------------------------------------------------------------
-- |
line' :: (Affine a1, Affine a2) => a1 -> a2 -> Line
line' p1 p2 = mkLine (cmp p1, cmp p2)

-- |
line :: XY -> XY -> Line
line = line'

-- |
segment' :: (Affine a1, Affine a2) => a1 -> a2 -> Segment
segment' p1 p2 = mkSegment (cmp p1, cmp p2)

-- |
segment :: XY -> XY -> Segment
segment = segment'

-- |
ray' :: (Affine a1, Affine a2) => a1 -> a2 -> Ray
ray' p1 p2 = mkRay (cmp p1, cmp p2)

-- |
ray :: XY -> XY -> Ray
ray = ray'

-- | The template for a segment.
aSegment :: Segment
aSegment = asCmp 1

-- | The template for a line.
aLine :: Line
aLine = asCmp 1

-- | The template for a ray.
aRay :: Ray
aRay = asCmp 1

{- | Returns a line, ray or a segment with given unit, in case of a segment -- with given length.

> [ aSegment # rotate x # extendToLength r
> | x <- [0,5..360]
> , let r = 2 + sin (7 * rad x) ]
<< figs/extendToLength.svg >>
-}
extendToLength :: Double -> Segment -> Segment
extendToLength l s = s # through' (paramL l (asLine s))

{- | Returns a segment extended to a closest intersection point with a given curve.

> let t = aTriangle
>     s1 = aSegment # at (1,1)
>     s2 = aSegment # at (0.3,0.3)
> in t <+>
>    [ s1 # along a # extendTo t | a <- [0,10..360] ] <+>
>    [ s2 # along a # extendTo t | a <- [0,10..360] ]
<< figs/extendTo.svg >>
-}
extendTo :: (Curve c, Intersections Ray c)
         => c -> Segment -> Maybe Segment
extendTo c s = extend <$> closestTo (start s) (intersections (asRay s) c)
  where extend p = s # through' p

{- | Returns a segment normal to a given curve starting at given point.

> let t = aTriangle
>     pA = point (1,1) #: "A"
>     sa = t # heightFrom pA #: "a"
>     pB = point (1.2,0) #: "B"
>     sb = sa # heightFrom pB #: "b"
> in t <+> pA <+> sa <+> pB <+> sb
<< figs/heightFrom.svg >>
-}
heightFrom :: (Affine p, Curve c, Intersections Ray c)
           => p -> c -> Maybe Segment
heightFrom p c = (aSegment # at' p # normalTo c) >>= extendTo c

{- | Returns a list of segments as a result of clipping the line by a closed curve.

> let star = polarPoly (\x -> 2 + cos (5*x)) [0,0.1..1] # closePolyline
>     rs =  [ aSegment # scale 3 # rotate 35 # at (-1, x)
>           | x <- [-4,-3.5..4] ]
> in star <+> foldMap (clipBy star) rs
g
<< figs/clipBy.svg >>
-}
clipBy :: (Linear l, Intersections l c, ClosedCurve c)
       => c -> l -> [Segment]
clipBy c l = filter internal $ Segment <$> zip ints (tail ints)
  where
    ints = sortOn (->@ l) $ intersections l c <> ends
    internal s = c `isEnclosing` (s @-> 0.5)
    ends = cmp . (l @->) <$> case bounds l of
                               Unbound -> []
                               Semibound -> [0]
                               Bound -> [0,1]

-- | A generalized version of `through`.
through' :: (Metric p, Affine p, Linear l) => p -> (l -> l)
through' p l = l
               # along' (azimuth p0 p)
               # scaleAt' p0 (distance p0 p / unit l)
  where p0 = start l

{- | Turns and extends the line so that it passes through a given point.

> let pA = point (2,3) #: "A"
>     pB = point (3,2) #: "B"
> in aSegment # through' pA <+>
>    aRay # through (3,2) <+>
>    pA <+> pB <+> origin
<< figs/through.svg>>
-}
through :: Linear l => XY -> (l -> l)
through = through'

{- | If possible, turns the line so that it becomes normal to a given curve, pointing towards a curve.

> let c = Plot $ (\t -> t :+ sin t) . (*6)
> in c <+>
>     [ aSegment # at (x,0) # normalTo c
>     | x <- [0..7] ]
<< figs/normalTo.svg >>
-}
normalTo :: (Curve c, Linear l) => c -> l -> Maybe l
normalTo c l = turn <*> Just l
  where s = start l
        turn = if c `isContaining` s
               then along' . normal c <$> (s ->@? c)
               else along' . ray' s <$> (point' s # projectOn c)

{- | If possible, returns a line segment normal to the curve starting at a given parameter.

> let t = aTriangle
> in t
>    <+> t # normalSegment 0.2 #: "1"
>    <+> t # normalSegment 0.5 #: "2"
>    <+> t # normalSegment (2/3) #: "3"
>    <+> t # normalSegment 1 #: "4"
<< figs/normalSegment.svg >>
-}
normalSegment :: Curve c => Double -> c -> Maybe Segment
normalSegment x c =
  do p <- paramMaybe c x
     return $ aSegment # at' p # along' (normal c x)

{- | Reflects the curve  at a given parameter against the normal, if it exists, or does nothing otherwise.

> let p = plot (\t -> (t, sin t)) # range (0, 2*pi)
> in [ p # flipAt x #: thin | x <- [0,0.01..1] ]

<< figs/flipAt.svg >>
-}
flipAt :: Curve c => Double -> c -> Maybe c
flipAt x c = (`reflectAt` c) <$> normalSegment x c

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

{- $ang| The `Angle` mark looks like a small labeled arc and two segments.

>>> anAngle 30
<Angle 30 (0.0, 0.0)>

> writeSVG 400 "figs/angle1.svg" $
>   let t = triangle2a 30 60
>      a1 = anAngle 30 #: "#" <> loffs ((-1):+1)
>      a2 = anAngle 90 # on (side t 2) 0 #: "#"
>      a3 = vertexAngle t 1 #: "#"
>  in t <+> a1 <+> a2 <+> a3

<<figs/angle1.svg>>

Angle is a `Manifold Direction` instance:

>>> Angle 0 0 90 @-> 0.5
45°
>>>  Angle 0 0 90 @-> 0
0°
>>> Angle 0 0 90 @-> 1
90°

>>> 45 ->@ Angle 0 0 90
0.5
>>> 180 ->@ Angle 0 0 90
2.0
>>> 0 ->@ Angle 0 10 30
17.5$ang
-}

-- | The template for an angle with given value.
anAngle :: Direction -> Angle
anAngle = Angle 0 0

-- | Returns the angle equal to the angle between two lines, located on the first one.
angleBetween :: (Linear l1,  Linear l2) => l1 -> l2 -> Angle
angleBetween l1 l2 = anAngle (angle l2 - angle l1)
                     # at' (start l1)
                     # along' l1

-- | Returns the Angle mark for the angle, formed by three points.
angleWithin :: (Affine a1, Affine a2, Affine a3)
            => a1 -> a2 -> a3 -> Angle
angleWithin p1 p2 p3 = angleBetween (ray' p2 p1) (ray' p2 p3)

{- | Returns the Angle mark for the angle at given polyline vertex.

> let p = regularPoly 5 # rotate 20 # scaleX 2
> in p <+> [ vertexAngle i p #: "#" | i <- [1..5] ]

<< figs/vertexAngle.svg >>
-}
vertexAngle :: PiecewiseLinear p => Int -> p -> Angle
vertexAngle n p = angleWithin p3 p2 p1 # innerAngle
  where p1 = p # vertex (n-1)
        p2 = p # vertex n
        p3 = p # vertex (n+1)

{- | Returns a ray, representng the bisectrisse of a given angle.

> let t = triangle2a 50 70
>     r1 = t # vertexAngle 0 # bisectrisse #: thin <> white
>     r2 = t # vertexAngle 1 # bisectrisse #: thin <> white
>     r3 = t # vertexAngle 2 # bisectrisse #: thin <> white
>     p = head $ intersectionPoints r1 r2
>     c = circle' (p `distanceTo` side 0 t) p
> in t <+> r1 <+> r2 <+> r3 <+> p <+> c

<< figs/bisectrisse.svg >>
-}
bisectrisse :: Angular a => a -> Ray
bisectrisse an = aRay # at' (refPoint an) # along' (an @-> 0.5)

{- | The supplementary angle.

> let a = anAngle 60 #: "a"
>     b = a # supplementary #: "b"
> in aRay <+> aRay # rotate 60 <+> a <+> b
<<figs/angle2.svg>>
-}
supplementary :: Angular an => an -> an
supplementary = asAngle . (\(Angle p s e) -> Angle p e (s + 180)) . toAngle

{- | The vertical angle.

> let a = anAngle 60 #: "a"
>     b = a # vertical #: "b"
> in aRay <+> aRay # rotate 60 <+> a <+> b
<<figs/angle3.svg>>
-}
vertical :: Angular an => an -> an
vertical = rotate 180

{- | The reflex angle.

> let a = anAngle 60 #: "a"
>     b = a # reflex #: "b"
> in aRay <+> aRay # rotate 60 <+> a <+> b
<<figs/angle4.svg>>
-}
reflex :: Angular an => an -> an
reflex an = asAngle $ Angle (refPoint an) e s
  where s = angleStart an
        e = angleEnd an

-- | Returns the inner angle (less than `pi`) for a given one.
innerAngle :: Angular an => an -> an
innerAngle an | deg (angleValue an) > 180 = reflex an
              | otherwise = an

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

{- | Constructs a parametric graph as a `Polyline`.

<< figs/parametricPoly.svg >>
-}
parametricPoly :: (Double -> XY) -> [Double] -> Polyline
parametricPoly f range =
  mkPolyline [ x :+ y | t <- range , let (x,y) = f t ]

{- | Constructs a polar graph as a `Polyline`.

> polarPoly (\x -> 2 + cos (5*x)) [0,0.1..1] # closePolyline

<< figs/polarPoly.svg >>
-}
polarPoly :: (Double -> Double) -> [Double] -> Polyline
polarPoly rho range =
  mkPolyline [ mkPolar (rho phi) phi | x <- range
                                     , let phi = 2*pi*x ]

{- | Constructs a regular polygon with given number of sides, enscribed in a unit circle.

> row [ regularPoly n | n <- [3..7] ]

<< figs/row.svg >>
-}
regularPoly :: Int -> Polygon
regularPoly n' = rotate 90 $ closePolyline $
                 polarPoly (const 1) [0,1/n..1-1/n]
  where n = fromIntegral n'

-- | The template for a triangle.
aTriangle :: Triangle
aTriangle = asCmp 1

-- | The template for a square.
aSquare :: Rectangle
aSquare = asCmp 1

-- | An invisible square with given size, used for spacing aligned objects.
space :: Double -> Decorated Rectangle
space a = aSquare # scale a #: invisible

-- | The template for a Rectangle.
aRectangle :: Double -> Double -> Rectangle
aRectangle a b = aSquare # scaleX a . scaleY b

{- | Returns a triangle with base 1 and two given angles.

> -- a family of triangles, sharing one side and opposite vertex angle value (30°)
> [ triangle2a a (150-a) #: thin | a <- [10,20..170] ]

<< figs/triangle2a.svg >>
-}
triangle2a :: Direction -> Direction -> Triangle
triangle2a a1 a2 = case intersections r1 r2 of
                    [p] -> mkTriangle (0,0) (1,0) (xy p)
                    [] -> mkTriangle @XY (0,0) (1,0) (0,0)
  where r1 = aRay # along' a1
        r2 = aRay # at (1,0) # along' (180 - a2)

{- | Returns a triangle with given three sides.

> -- a family of triangles, sharing two sides
> [ triangle3s 10 a 9 #: thin | a <- [2..18] ]

<< figs/triangle3s.svg >>
-}
triangle3s :: Double -> Double -> Double -> Maybe Triangle
triangle3s a b c = case intersections c1 c2 of
                     [] -> Nothing
                     p:_ -> Just $ Triangle [0, a :+ 0, p]
  where c1 = aCircle # scale b
        c2 = aCircle # scale c # at (a, 0)

{- | A isosceles right triangle with unit side.

> [ aRightTriangle # scaleX a # scaleY (10-a) #: thin
> | a <- [0.5,1..10] ]

<< figs/aRightTriangle.svg >>
-}
aRightTriangle = RightTriangle $ triangle2a 90 45

------------------------------------------------------------
{- | The median of the polygon from given vertex to given side.

> let p = regularPoly 9 # rotate 20 # scaleX 2
> in p <+> [ p # median 0 i #: thin <> white | i <- [1..7] ]

<< figs/median.svg>>
-}
median :: PiecewiseLinear p
  => Int -- ^ vertex index
  -> Int -- ^ side index
  -> p -> Segment
median v s t = Segment (vertex v t, side s t @-> 0.5)

{- | The altitude of the polygon from given vertex to given side.

> let p = regularPoly 9 # rotate 20 # scaleX 2
> in p <+> [ p # altitude 0 i #: thin <> white <+>
>            p # side i # asLine #: thin <> dashed
>          | i <- [1..8] ]

<< figs/altitude.svg >>
-}
altitude :: PiecewiseLinear p
  => Int  -- ^ vertex index
  -> Int  -- ^ side index
  -> p -> Segment
altitude v s p = heightFrom (vertex v p) (asLine (side s p)) # fromJust

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

{- | Creates a scale as a list of labeled points on a given curve

> let p = polarPoly id [0,0.1..1]
>     sc = linearScale (round . (*10)) [0,0.1..1] p
> in p <+> sc

<< figs/linearScale.svg>>
-}
linearScale :: (Show s, Curve c)
            => (Double -> s) -- ^ labeling function
            -> [Double] -- ^ range of the curve parameter
            -> c -- ^ the curve
            -> [Decorated Point]
linearScale fn rng c = [ pointOn c x #:: label (show (fn x))
                       | x <- rng ]

{- | Creates a circular integer scale, representing modular arithmetics.

> let c  = aCircle # rotate 90
>     s1 = modularScale 12 c
>     t  = aTriangle # scale 2
>     s2 = modularScale 9 t
> in (c <+> s1) `beside` space 1 `beside` (s2 <+> t)
<< figs/modularScale.svg>>
-}
modularScale :: (Trans c, Curve c) => Int -> c -> [Decorated Point]
modularScale n = linearScale lf rng
  where n' = fromIntegral n
        rng = (/n') <$> [0..n'-1]
        lf = (`mod` n) . negate . round . (* n')

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

instance WithOptions Point where
  defaultOptions p = mkOptions
    [ LabelPosition $ cmp p
    , LabelOffset (0 :+ 1)
    , LabelCorner (0, 0)
    , LabelAngle 0
    , Stroke "#444"
    , Fill "red"
    , Thickness "1" ]

instance WithOptions Label where
  defaultOptions p = mkOptions
    [ LabelPosition $ cmp p
    , LabelOffset 0
    , LabelCorner (0, 0)
    , LabelAngle 0 ]

instance WithOptions Circle where
  defaultOptions c = mkOptions
    [ LabelPosition $ c @-> 0
    , LabelOffset $ cmp $ normal c 0
    , LabelCorner (-1,0)
    , LabelAngle 0
    , Stroke "orange"
    , Fill "none"
    , Thickness "2" ]

instance WithOptions Line where
  defaultOptions l = mkOptions
    [ LabelPosition $ l @-> 0.5
    , LabelOffset $ cmp $ normal l 0
    , Stroke "orange"
    , Fill "none"
    , Thickness "2" ]

deriving via Line instance WithOptions Ray
deriving via Line instance WithOptions Segment

instance WithOptions Polyline where
  defaultOptions _ = mkOptions
    [ Stroke "orange"
    , Fill "none"
    , Thickness "2" ]

deriving via Polyline instance WithOptions Polygon
deriving via Polyline instance WithOptions Triangle
deriving via Triangle instance WithOptions RightTriangle
deriving via Polyline instance WithOptions Rectangle

instance WithOptions Angle where
  defaultOptions an = mkOptions
    [ Stroke "white"
    , Fill "none"
    , Thickness "1.25"
    , MultiStroke 1
    , LabelPosition $ refPoint an + scale 30 (cmp (bisectrisse an)) ]

instance Pnt a => WithOptions (Plot a) where
  defaultOptions = defaultOptions . asPolyline

instance Pnt a => WithOptions (ClosedPlot a) where
  defaultOptions = defaultOptions . asPolyline

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


-- example 3.20

import Statistics.Distribution
import Statistics.Distribution.Poisson

-- Let a ten-square-yard sampling region be a unit area; the mean density is 5 per square yard:
pd :: PoissonDistribution
pd = poisson 5.0

p :: Int -> Double
p = probability pd

main = do
  putStrLn $ "the probability that none of the regions will contain: P(X<1) = P(X=0): "
          ++ show (p 0)


module Vector (
    tests_Vector
    ) where

import Data.AEq
import Debug.Trace
import Test.Framework
import Test.Framework.Providers.QuickCheck2
import Test.QuickCheck hiding ( vector )
import qualified Test.QuickCheck as QC

import Numeric.LinearAlgebra.Vector( Vector )
import qualified Numeric.LinearAlgebra.Vector as V

import Test.QuickCheck.LinearAlgebra( TestElem(..), Dim(..), Assocs(..),
    VectorPair(..) )
import qualified Test.QuickCheck.LinearAlgebra as Test

import Typed

tests_Vector = testGroup "Vector"
    [ testPropertyI "dim/fromList" prop_dim_fromList
    , testPropertyI "at/fromList" prop_at_fromList
    , testPropertyI "zero" prop_zero
    , testPropertyI "constant" prop_constant
    , testPropertyI "indices" prop_indices
    , testPropertyI "elems" prop_elems
    , testPropertyI "assocs" prop_assocs
    , testPropertyI "update" prop_update
    , testPropertyI "accum" prop_accum
    , testPropertyI "map" prop_map
    , testPropertyI "zipWith" prop_zipWith
    , testPropertyI "concat" prop_concat
    , testPropertyI "slice" prop_slice
    , testPropertyI "splitAt" prop_splitAt
    , testPropertyDZ "sumAbs" prop_sumAbs prop_sumAbs
    , testPropertyDZ "norm2" prop_norm2 prop_norm2
    , testPropertyDZ "whichMaxAbs1" prop_whichMaxAbs1 prop_whichMaxAbs1
    , testPropertyDZ "whichMaxAbs2" prop_whichMaxAbs1 prop_whichMaxAbs2
    , testPropertyDZ "dot" prop_dot prop_dot
    , testPropertyDZ "kronecker" prop_kronecker prop_kronecker
    , testPropertyDZ "add" prop_add prop_add
    , testPropertyDZ "addWithScale" prop_addWithScale prop_addWithScale
    , testPropertyDZ "sub" prop_sub prop_sub
    , testPropertyDZ "scale" prop_scale prop_scale
    , testPropertyDZ "mul" prop_mul prop_mul
    , testPropertyDZ "negate" prop_negate prop_negate
    , testPropertyDZ "conjugate" prop_conjugate prop_conjugate
    , testPropertyDZ "abs" prop_abs prop_abs
    , testPropertyDZ "signum" prop_signum prop_signum
    , testPropertyDZ "div" prop_div prop_div
    , testPropertyDZ "recip" prop_recip prop_recip
    , testPropertyDZ "sqrt" prop_sqrt prop_sqrt
    , testPropertyDZ "exp" prop_exp prop_exp
    , testPropertyDZ "log" prop_log prop_log
    , testPropertyDZ "pow" prop_pow prop_pow
    , testPropertyDZ "sin" prop_sin prop_sin
    , testPropertyDZ "cos" prop_cos prop_cos
    , testPropertyDZ "tan" prop_tan prop_tan
    , testPropertyDZ "asin" prop_asin prop_asin
    , testPropertyDZ "acos" prop_acos prop_acos
    , testPropertyDZ "atan" prop_atan prop_atan
    , testPropertyDZ "sinh" prop_sinh prop_sinh
    , testPropertyDZ "cosh" prop_cosh prop_cosh
    , testPropertyDZ "tanh" prop_tanh prop_tanh
    , testPropertyDZ "asinh" prop_asinh prop_asinh
    , testPropertyDZ "acosh" prop_acosh prop_acosh
    , testPropertyDZ "atanh" prop_atanh prop_atanh
    ]

------------------------- Vector Construction ------------------------------

prop_dim_fromList t (Dim n) =
    forAll (QC.vector n) $ \es -> let
        x = typed t $ V.fromList n es
        in V.dim x == n

prop_at_fromList t (Dim n) =
    forAll (QC.vector n) $ \es -> let
        x = typed t $ V.fromList n es
        in and [ V.at x i === e | (i,e) <- zip [ 0..n-1 ] es ]

prop_zero t (Dim n) = let
    x = typed t $ V.zero n
    in x === V.fromList n (replicate n 0)

prop_constant t (Dim n) e =
    V.constant n e === V.fromList n (replicate n e)
  where
    _ = typed t [e]

-------------------------- Accessing Vectors ------------------------------

prop_indices t x =
    V.indices x === [ 0..((V.dim x) - 1) ]
  where
    _ = immutableVector x
    _ = typed t x

prop_elems t x =
    V.elems x === [ V.at x i | i <- V.indices x ]
  where
    _ = typed t x

prop_assocs t x =
    V.assocs x === zip (V.indices x) (V.elems x)
  where
    _ = typed t x

------------------------- Incremental Updates ------------------------------

prop_update t (Assocs n ies) =
    forAll (typed t `fmap` Test.vector n) $ \x -> let
        x' = V.update x ies
        is = V.indices x
        is1 = (fst . unzip) ies
        is0 = [ i | i <- is, i `notElem` is1 ]
        in and $
            [ V.at x' i `elem` [ e | (i',e) <- ies, i' == i ]
            | i <- is1
            ] ++
            [ V.at x' i === V.at x i
            | i <- is0
            ]

prop_accum t (Blind f) (Assocs n ies) =
    forAll (typed t `fmap` Test.vector n) $ \x -> let
        x' = V.accum f x ies
        in x' === V.fromList n [ foldl f e [ e' | (i',e') <- ies, i' == i]
                               | (i,e) <- V.assocs x ]
  where
      _ = typed t $ (snd . unzip) ies

-------------------------- Derived Vectors ------------------------------

prop_map t (Blind f) x =
    V.map f x === V.fromList (V.dim x) (map f $ V.elems x)
  where
    _ = typed t x
    _ = typed t $ V.map f x

prop_zipWith t (Blind f) (VectorPair x y) =
    V.zipWith f x y === (V.fromList (V.dim x) $
                                zipWith f (V.elems x) (V.elems y))
  where
    _ = typed t x
    _ = typed t y
    _ = typed t $ V.zipWith f x y

prop_concat t xs =
    V.elems (V.concat xs) === concatMap V.elems xs
  where
    _ = typed t $ head xs

------------------------------ Vector Views-- --------------------------------

prop_slice t x =
    forAll (choose (0,n)) $ \n' ->
    forAll (choose (0,n-n')) $ \o ->
        V.slice o n' x === V.fromList n' (take n' $ drop o $ es)
  where
    n  = V.dim x
    es = V.elems x
    _  = typed t x

prop_splitAt t x =
    forAll (choose (0,n)) $ \k ->
        V.splitAt k x === (V.fromList k $ take k es,
                           V.fromList (n-k) $ drop k es)
  where
    n  = V.dim x
    es = V.elems x
    _  = typed t x

-------------------------- Num Vector Operations --------------------------

prop_add t (VectorPair x y) =
    x `V.add` y === V.zipWith (+) x y
  where
    _ = typed t x

prop_addWithScale t alpha (VectorPair x y) =
    V.addWithScale alpha x y ~== V.add (V.scale alpha x) y
  where
    _ = typed t x

prop_sub t (VectorPair x y) =
    x `V.sub` y === V.zipWith (-) x y
  where
    _ = typed t x

prop_scale t k x =
    V.scale k x ~== V.map (k*) x
  where
    _ = typed t x

prop_mul t (VectorPair x y) =
    x `V.mul` y === V.zipWith (*) x y
  where
    _ = typed t x

prop_negate t x =
    V.negate x === V.map negate x
  where
    _ = typed t x

prop_conjugate t x =
    V.conjugate x === V.map conjugate x
  where
    _ = typed t x

prop_abs t x =
    V.abs x === V.map abs x
  where
    _ = typed t x

prop_signum t x =
    V.signum x === V.map signum x
  where
    _ = typed t x

---------------------- Fractional Vector Operations ------------------------

prop_div t (VectorPair x y) =
    x `V.div` y ~== V.zipWith (/) x y
  where
    _ = typed t x

prop_recip t x =
    V.recip x ~== V.map (1/) x
  where
    _ = typed t x

---------------------- Floating Vector Operations ------------------------

prop_exp t x =
    V.exp x ~== V.map exp x
  where
    _ = typed t x

prop_sqrt t x =
    V.sqrt x ~== V.map sqrt x
  where
    _ = typed t x

prop_log t x =
    V.log x ~== V.map log x
  where
    _ = typed t x

prop_pow t (VectorPair x y) =
    x `V.pow` y ~== V.zipWith (**) x y
  where
    _ = typed t x

prop_sin t x =
    V.sin x ~== V.map sin x
  where
    _ = typed t x

prop_cos t x =
    V.cos x ~== V.map cos x
  where
    _ = typed t x

prop_tan t x =
    V.tan x ~== V.map tan x
  where
    _ = typed t x

prop_asin t x =
    -- trace (show (V.asin x) ++ "\n" ++ (show $ V.map asin x)) $
    V.asin x ~== V.map asin x
  where
    _ = typed t x

prop_acos t x =
    V.acos x ~== V.map acos x
  where
    _ = typed t x

prop_atan t x =
    V.atan x ~== V.map atan x
  where
    _ = typed t x

prop_sinh t x =
    V.sinh x ~== V.map sinh x
  where
    _ = typed t x

prop_cosh t x =
    V.cosh x ~== V.map cosh x
  where
    _ = typed t x

prop_tanh t x =
    V.tanh x ~== V.map tanh x
  where
    _ = typed t x

prop_asinh t x =
    -- trace (show (V.asinh x) ++ "\n" ++ (show $ V.map asinh x)) $
    V.asinh x ~== V.map asinh x
  where
    _ = typed t x

prop_acosh t x =
    V.acosh x ~== V.map acosh x
  where
    _ = typed t x

prop_atanh t x =
    V.atanh x ~== V.map atanh x
  where
    _ = typed t x

-------------------------- Vector Properties ---------------------------------

prop_sumAbs t x =
    V.sumAbs x ~== (sum $ map norm1 $ V.elems x)
  where
    _ = typed t x

prop_norm2 t x =
    V.norm2 x ~== (sqrt $ sum $ map (^^2) $ map norm $ V.elems x)
  where
    _ = typed t x

prop_whichMaxAbs1 t x =
    (V.dim x > 0) && all (not . isNaN) (map norm1 $ V.elems x) ==>
        V.at x i === e
  where
    (i,e) = V.whichMaxAbs x
    _     = typed t x

prop_whichMaxAbs2 t x =
    (V.dim x > 0) && all (not . isNaN) (map norm1 $ V.elems x) ==>
        all (<= norm1 e) $ map norm1 (V.elems x)
  where
    (_,e) = V.whichMaxAbs x
    _     = typed t x

prop_dot t (VectorPair x y) =
    V.dot x y ~== sum (V.elems (x * conj y))
  where
    conj = V.conjugate
    (*)  = V.mul
    _    = typed t x

prop_kronecker t x y =
    x `V.kronecker` y ~==
        V.fromList (V.dim x * V.dim y)
                   [ e*f | e <- V.elems x, f <- V.elems y ]
  where
    _ = typed t x


module Lib
    ( euclidianDistance
    , Point
    , Vector
    , makePoint
    , eye1
    , getEyeCenter
    , getLineFromPoints
    , projection
    ) where

import qualified Numeric.LinearAlgebra.HMatrix as HM

type Point = HM.Vector Double
type Vector = HM.Vector Double
data Line = Line { lineCoefficientMatrix :: HM.Matrix Double, lineTerms :: HM.Vector Double } deriving Show

makePoint :: Double -> Double -> Double -> Double -> Point
makePoint x y z t = HM.vector [x, y, z, t]

makeVector :: Double -> Double -> Double -> Double -> Vector
makeVector = makePoint

data EyeConfig = EyeConfig { retinaRadius :: Double, convergenceDistance :: Double }

data EyePosition = EyePosition { convergencePoint :: Point, direction :: Vector }

data Eye = Eye { config :: EyeConfig, position :: EyePosition }

eye1 :: Eye
eye1 = Eye {
    config = EyeConfig { retinaRadius = 10, convergenceDistance = 30 },
    position = EyePosition { convergencePoint = makePoint 1 1 1 1, direction = makeVector (-1) (-1) (-1) (-1) }
    }

euclidianDistance :: Point -> Point -> Double
euclidianDistance a b = sqrt . sum . map (**2) $ zipWith (-) (HM.toList a) (HM.toList b)

euclidianLength :: Vector -> Double
euclidianLength = euclidianDistance (makePoint 0 0 0 0)

getEyeCenter :: Eye -> Point
getEyeCenter Eye { config = eyeConfig, position = eyePosition } =
  let
    ratio = (/ (euclidianLength . direction $ eyePosition)) (convergenceDistance eyeConfig)
  in
    (+ convergencePoint eyePosition) . (* direction eyePosition) . HM.scalar $ ratio

getLineFromPoints :: Point -> Point -> Line
getLineFromPoints a b =
  let
    l = a - b
    a' = HM.toList a
    l' = HM.toList l
    coefficientXy = [l'!!1, -(head l'), 0, 0]
    coefficientYz = [0, l'!!2, -(l'!!1), 0]
    coefficientZt = [0, 0, l'!!3, -(l'!!2)]
    termXy = HM.vector coefficientXy HM.<.> a
    termYz = HM.vector coefficientYz HM.<.> a
    termZt = HM.vector coefficientZt HM.<.> a
  in
    Line {
        lineCoefficientMatrix = HM.matrix 4 (coefficientXy ++ coefficientYz ++ coefficientZt),
        lineTerms = HM.vector [termXy, termYz, termZt]
    }

projection :: Eye -> Point -> Maybe Point
projection eye point =
  let
    Eye { config = eyeConfig, position = eyePosition } = eye
    eyeCenter = getEyeCenter eye
    Line { lineCoefficientMatrix = coefficients, lineTerms = terms } =
        getLineFromPoints point (convergencePoint eyePosition)
    eyeDirection = direction eyePosition
    intersectionCoefficients = HM.fromRows . (++ [eyeDirection]) . HM.toRows $ coefficients
    intersectionTerms = HM.matrix 1 . (++ [eyeDirection HM.<.> eyeCenter]) . HM.toList $ terms
  in
    HM.flatten <$> HM.linearSolve intersectionCoefficients intersectionTerms


import Data.Complex (Complex, cis)

rootsOfUnity :: (Enum a, Floating a) => a -> [Complex a]
rootsOfUnity n =
  [ cis (2 * pi * k / n)
  | k <- [0 .. n - 1] ]

main :: IO ()
main = mapM_ print $ rootsOfUnity 3


module Fourier.FFTSpec (main,spec) where

import Test.Hspec
import Test.QuickCheck()
import Data.Complex(Complex((:+)))
import Fourier.FFT(fft)

main :: IO ()
main = hspec spec

spec :: Spec
spec = do

  describe "FFT" $ do

    it "transform steady function" $ do
      let x = fft [0.0, 0.0, 0.0]
        in x `shouldBe` ([0.0, 0.0, 0.0])

    it "transform impulse function" $ do
      let x = fft [0.0, 0.0, 1.0, 0.0, 0.0]
        in x `shouldBe` ([
          1.0 :+ (-0.0),
          (-0.8090169943749473) :+ (-0.5877852522924732),
          0.30901699437494723 :+ 0.9510565162951536,
          0.30901699437494773 :+ (-0.9510565162951534),
          (-0.8090169943749477) :+ 0.5877852522924728])

    it "transform sinusoid function" $ do
      let x = fft [sin x | x <- [0,0.3 .. pi]]
        in x `shouldBe` ([
          6.654043728627056 :+ (-0.0),
          (-2.495308482972513) :+ (-0.3737795890019122)
          ,(-0.46503997943364506) :+ (-0.1328809251599644),
          (-0.1886713255075324) :+ (-7.112462151227344e-2),
          (-0.10350867238552199) :+ (-3.701914259223908e-2),
          (-7.449340401431373e-2) :+ (-1.1604714311766413e-2),
          (-7.449340401431312e-2) :+ 1.1604714311764401e-2,
          (-0.10350867238552014) :+ 3.701914259223764e-2,
          (-0.18867132550753074) :+ 7.112462151227178e-2,
          (-0.46503997943364067) :+ 0.13288092515996228,
          (-2.4953084829725105) :+ 0.3737795890019047])


module Kantour.DropCalc where

import System.Environment
import Control.Monad
import Control.Monad.Primitive
import System.Random.MWC
import Text.Printf
import qualified Data.Vector as V
import Statistics.Sample
import Data.Vector.Algorithms.Intro
import Data.Function
import Numeric.SpecFunctions
import Kantour.Subcommand

data SubCmdDropCalc

instance Subcommand SubCmdDropCalc where
    name _ = "DropCalc"
    main _ = defaultMain

experiment :: PrimMonad m => Gen (PrimState m) -> Double -> m Int
experiment g r = experiment' 0
  where
    experiment' cnt = do
        let newCnt = cnt+1
        v <- uniform g
        if v <= r
          then pure newCnt
          else experiment' newCnt

computeCDF :: Int -> Int -> Double -> Double
computeCDF k r p = i (1-p) (fromIntegral r) (fromIntegral k)
  where
    i x a b = incompleteBeta a b x

analyzeResult :: Int -> Double -> [Int] -> IO ()
analyzeResult _ _ [] = putStrLn "Result is empty."
analyzeResult tot p xs = do
    let _xs = V.fromList xs :: V.Vector Int
    ys' <- V.unsafeThaw _xs
    sort ys'
    -- now ys is the sorted version of xs
    ys <- V.unsafeFreeze ys'
    let xMin = V.head ys
        xMax = V.last ys
        ysDb = fromIntegral <$> ys :: V.Vector Double
    printf "min: %d, max: %d\n" xMin xMax
    printf "mean: %6.4f, std dev.: %6.4f\n" (mean ysDb) (stdDev ysDb)
    (fix $ \indScan curInd ->
      if curInd >= V.length ys
        then pure ()
        else do
          let curE = ys V.! curInd
          -- find next index whose elem is not curE
          let nextInd = fix (\self curInd' ->
                             if curInd' >= V.length ys
                               then Nothing
                               else if ys V.! curInd' /= curE
                                      then Just curInd'
                                      else self (curInd'+1)) (curInd+1)
              fI = fromIntegral :: Int -> Double
              ppr n cr = when (n `mod` 10 == 0) $
                  printf "<= %d\tactual: %6.4f%%\tcomputed: %6.4f%%\n"
                      n
                      (cr * 100)
                      (computeCDF n 1 (1-p) * 100)
          case nextInd of
              Nothing ->
                  ppr curE (fI (V.length ys -1) / fI tot )
              Just ind -> do
                  ppr curE (fI (ind-1) / fI tot)
                  indScan ind
          ) 0

defaultMain :: IO ()
defaultMain = do
    as <- getArgs
    case as of
        [rawDropRate,rawCnt]
            | Just (dropRate,cnt) <-
                (,) <$> parseRate rawDropRate
                    <*> parseCount rawCnt
            -> do
                printf "Performing %d experiments with drop rate %6.4f%%:\n"
                  cnt (dropRate * 100)
                g <- createSystemRandom
                results <- replicateM cnt (experiment g dropRate)
                analyzeResult cnt dropRate results
        _ -> putStrLn "Usage: <rate> <# of experiments>"
             >> putStrLn "example of rate: '0.20' or '20%'"
  where
    -- either a number between 0 and 1
    -- or {x}% where x can be between 0 and 100
    parseRate :: String -> Maybe Double
    parseRate [] = Nothing
    parseRate xs
        | last xs == '%'
        , [(r,"")] <- reads (init xs) = Just (r/100) >>= verify
        | [(r,"")] <- reads xs = Just r >>= verify
        | otherwise = Nothing
      where
        verify x = x <$ guard (x >= 0 && x <= 1)

    -- must be positive number
    parseCount :: String -> Maybe Int
    parseCount xs
        | [(v,"")] <- reads xs
        , v > 0 = Just v
        | otherwise = Nothing


import Data.List (transpose)
import Data.Complex

type Matrix a = [[a]]

main :: IO ()
main =
    mapM_ (\a -> do
        putStrLn "\nMatrix:"
        mapM_ print a
        putStrLn "Conjugate Transpose:"
        mapM_ print (conjTranspose a)
        putStrLn $ "Hermitian? " ++ show (isHermitianMatrix a)
        putStrLn $ "Normal? " ++ show (isNormalMatrix a)
        putStrLn $ "Unitary? " ++ show (isUnitaryMatrix a))
        ([[[3,       2:+1],
           [2:+(-1), 1   ]],

          [[1, 1, 0],
           [0, 1, 1],
           [1, 0, 1]],

          [[sqrt 2/2:+0, sqrt 2/2:+0,     0   ],
           [0:+sqrt 2/2, 0:+ (-sqrt 2/2), 0   ],
           [0,           0,               0:+1]]] :: [Matrix (Complex Double)])

isHermitianMatrix, isNormalMatrix, isUnitaryMatrix :: RealFloat a => Matrix (Complex a) -> Bool
isHermitianMatrix a = a `approxEqualMatrix` conjTranspose a
isNormalMatrix a = (a `mmul` conjTranspose a) `approxEqualMatrix` (conjTranspose a `mmul` a)
isUnitaryMatrix a = (a `mmul` conjTranspose a) `approxEqualMatrix` ident (length a)

approxEqualMatrix :: (Fractional a, Ord a) => Matrix (Complex a) -> Matrix (Complex a) -> Bool
approxEqualMatrix a b = length a == length b && length (head a) == length (head b) &&
                        and (zipWith approxEqualComplex (concat a) (concat b))
    where approxEqualComplex (rx :+ ix) (ry :+ iy) = abs (rx - ry) < eps && abs (ix - iy) < eps
          eps = 1e-14

mmul :: Num a => Matrix a -> Matrix a -> Matrix a
mmul a b = [[sum (zipWith (*) row column) | column <- transpose b] | row <- a]

ident :: Num a => Int -> Matrix a
ident size = [[fromIntegral $ div a b * div b a | a <- [1..size]] | b <- [1..size]]

conjTranspose :: Num a => Matrix (Complex a) -> Matrix (Complex a)
conjTranspose = map (map conjugate) . transpose


module Number
    (
    parseNumber
    ) where

import Text.ParserCombinators.Parsec hiding (spaces)
import Control.Monad
import Numeric
import Data.Ratio
import Data.Complex

import LispType

-- Parse Numbers
-- |
-- >>> readExpr "123"
-- "Found value"
parseNumber :: Parser LispVal
parseNumber = parseDecimal
          <|> parseHex
          <|> parseOct
          <|> parseBin
          <|> parseRatio
          <|> parseComplex

parseDecimal = parseDecimal1 <|> parseDecimal2

parseDecimal1 :: Parser LispVal
parseDecimal1 = liftM (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 $ hex2dec x

-- |
-- >>> readExpr "#o12345"
-- "Found value"
parseOct :: Parser LispVal
parseOct = do
  _ <- try $ string "#o"
  x <- many1 octDigit
  return $ Number $ oct2dec x

-- |
-- >>> readExpr "#b1010"
-- "Found value"
parseBin :: Parser LispVal
parseBin = do
  _ <- try $ string "#b"
  x <- many1 $ oneOf "01"
  return $ Number $ bin2dec x

hex2dec :: String -> Integer
hex2dec x = fst $ head $ readHex x

oct2dec :: String -> Integer
oct2dec x = fst $ head $ readOct x

bin2dec :: String -> Integer
bin2dec = bin2dec' 0
  where
    bin2dec' digint "" = digint
    bin2dec' digint (x:xs) =
      let old = 2 * digint + (if x == '0' then 0 else 1) in
        bin2dec' old xs

-- |
-- >>> readExpr "10.123"
-- "Found value"
parseFloat :: Parser LispVal
parseFloat = do
  x <- many1 digit
  _ <- char '.'
  y <- many1 digit
  return $ Float $ fst . head $ readFloat $ x ++ "." ++ y

-- |
-- >> readExpr "/ 3 4"
-- "Found value"
parseRatio :: Parser LispVal
parseRatio = do
  x <- many1 digit
  _ <- char '/'
  y <- many1 digit
  return $ Ratio $ read x % read y

{- https://hackage.haskell.org/package/base-4.8.1.0/docs/Data-Complex.html -}

-- |
-- >> readExpr "44.2+33i"
-- "Found value"

{- on lisp
> (+ 12+1i 12+1i)
24+2i
>
-}

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

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


module Distr where

import MHMonad
import Statistics.Distribution
import Statistics.Distribution.Normal (normalDistr)
import Statistics.Distribution.Beta (betaDistr)
import Statistics.Distribution.Beta (betaDistr)
import Statistics.Distribution.Gamma (gammaDistr)
import qualified Statistics.Distribution.Poisson as Poisson
import Data.List

categorical :: [Double] -> Meas Integer
categorical rs = do r <- sample
                    return $ categ rs r

normal :: Double -> Double -> Meas Double
normal m s = do x <- sample
                return $ quantile (normalDistr m s) x

normalPdf :: Double -> Double -> Double -> Double
normalPdf m s = density $ normalDistr m s

exponential :: Double -> Meas Double
exponential rate =
  do x <- sample
     return $ - (log x / rate)

expPdf :: Double -> Double -> Double
expPdf rate x = exp (-rate*x) * rate

gamma :: Double -> Double -> Meas Double
gamma a b = do
  x <- sample
  return $ quantile (gammaDistr a b) x

beta :: Double -> Double -> Meas Double
beta a b = do
  x <- sample
  return $ quantile (betaDistr a b) x

poisson :: Double -> Meas Integer
poisson lambda = do
  x <- sample
  let cmf = scanl1 (+) $ map (probability $ Poisson.poisson lambda) $ [0,1..]
  let (Just n) = findIndex (\r -> r > x) cmf
  return $ fromIntegral n

poissonPdf :: Double -> Integer -> Double
poissonPdf rate n = probability (Poisson.poisson rate) (fromIntegral n)

dirichlet :: [Double] -> Meas[Double]
dirichlet as = do
  xs <- mapM exponential as
  let s = Prelude.sum xs
  let ys = map (/ s) xs
  return ys

bernoulli :: Double -> Meas Bool
bernoulli r = do
  x <- sample
  return $ x < r


-- tests/Data/Network/FullyConnected.hs
module Data.Network.FullyConnectedSpec (spec) where

import Test.Hspec

import Data.Network
import Data.Network.Activation
import Data.Network.FullyConnected as FC
import System.Random
import Numeric.LinearAlgebra

spec :: Spec
spec = do
  describe "A fully connected layer composed of a weight matrix, a bias vector and an activation function" $ do
    describe "instance Layer (FullyConnected (Matrix R) R) where" $ do
      describe "makeRandomLayer" $ do
        it "creates a randomly initialized layer" $ do
          g <- getStdGen
          let l = makeRandomLayer g id 10 20 (logistic :: ActivationFunction (Matrix Double)) :: FullyConnected (Matrix Double) Double
          (size (weights l)) `shouldBe` (10,20)
          (size (biases l)) `shouldBe` 20
          let (g',g'') = split g
          let rs = randoms g'
          let i = (10><20) rs
          (activate (activation l) i) `shouldBe` (activate logistic i)
      describe "compute" $ do
        it "Computes the linear combination of the layer's inputs" $ do
          g <- getStdGen
          let (g1,g2) = split g
          let l@(FC w b af) = makeRandomLayer g1 id 10 20 (logistic) :: FullyConnected (Matrix Double) Double
          let rs = randoms g2
          let i = (100><10) rs
          let o = compute l i
          size o `shouldBe` (100,20)
          o `shouldBe` ((fromRows $ replicate 100 b) + (i <> w))

    describe "instance Network i (Layers (FullyConnected i Double)) where" $ do
      describe "makeRandomNetwork" $ do
        it "Creates a randomly initialized network given the specification of the layers" $ do
          g <- getStdGen
          let (g1,g2) = split g
          let n = makeRandomNetwork g1 id 10 [(30, logistic)] (20, logistic) :: FeedForwardFC (Matrix R)
          pending


module Utils.Stats (
        pricesAtInterval
    ,   baseStats
    ,   pricePercChange
    ,   slope'
    ,   std'
    ,   mean
    ,   roundToNearest
) where

import qualified Data.List                   as L
import           Data.Maybe
import qualified Data.Vector.Unboxed         as U
import           Protolude
import           Statistics.LinearRegression
import           Types                       (BasicStats (..))

-- i wonder if this could be written with foldable fns
intervalPrices :: Num a  => Int -> [a] -> [[a]]
intervalPrices _ [] = []
intervalPrices i p = let (x, xs) = splitAt i p
                        in x : intervalPrices i xs

pricesAtInterval :: (Num a ) => Int -> [a] -> Maybe [a]
pricesAtInterval i p =
    let priceGroups = intervalPrices i $ reverse p
    in traverse head $ reverse priceGroups

mean :: [Double] -> Maybe Double
mean [] = Nothing
mean x  = Just $  sum x / fromIntegral (length x)

std' :: [Double] -> Maybe Double
std' [] = Nothing
std' x  =
    let maybeμ     = mean x
        length'    = fromIntegral (length x)
        maybeVar   = case maybeμ of
                        Nothing -> Nothing
                        Just μ  -> Just $ foldr(\c acc -> (μ - c)**2 + acc) 0 x
    in case maybeVar of
        Nothing  -> Nothing
        Just var -> Just $ sqrt (var / length')

slope' :: [Double] -> Maybe Double
slope' [] = Nothing
slope' y  =
    let length' = fromIntegral $ L.length y :: Double
        xs      = U.fromList [0..length']
        ys      = U.fromList y
        (_, beta) = linearRegression xs ys
    in Just beta

roundToNearest :: Int -> Double -> Double
roundToNearest n x = fromInteger (round (x * 10^n)) / 10^n

pricePercChange :: [Double] -> Maybe Double
pricePercChange [] = Nothing
pricePercChange x  =
    let maxPrice = maximum x -- maximum price for say past month
        currentAvgPrice' = mean $ reverse $ take 3 (reverse x)-- avg price for past 3 days
    in case currentAvgPrice' of
        Nothing              -> Nothing
        Just currentAvgPrice -> Just $ ((maxPrice - currentAvgPrice) / maxPrice) * 100

baseStats :: [Double] -> Maybe BasicStats
baseStats [] = Nothing
baseStats x  =
    let x'            = reverse x
        reverseTake y = reverse . take y -- i dont want to use drop
        direction     = roundToNearest 2 $ (* 1e5) $ fromMaybe 0 $ slope' (reverseTake 5 x')
        std''         = roundToNearest 2 $ (* 1e5) $ fromMaybe 0 $ std' (reverseTake 7 x')
        change'       = roundToNearest 2 $ fromMaybe 0  $ pricePercChange $ reverseTake 24 x'
    in Just $ BasicStats direction change' std''


{-# LANGUAGE TypeFamilies, PatternGuards #-}
-----------------------------------------------------------------------------
-- |
-- Module      :  Statistics.Order
-- Copyright   :  (C) 2012 Edward Kmett,
-- License     :  BSD-style (see the file LICENSE)
--
-- Maintainer  :  Edward Kmett <ekmett@gmail.com>
-- Stability   :  provisional
-- Portability :  Haskell 2011 + TypeFamilies
--
----------------------------------------------------------------------------

module Statistics.Order
  (
  -- * L-Estimator
    L(..)
  -- ** Applying an L-estimator
  , (@@), (@!)
  -- ** Analyzing an L-estimator
  , (@#)
  , breakdown
  -- ** Robust L-Estimators
  , trimean  -- Tukey's trimean
  , midhinge -- average of q1 and q3
  , iqr      -- interquartile range
  , iqm      -- interquartile mean
  , lscale   -- second L-moment
  -- ** L-Estimator Combinators
  , trimmed
  , winsorized, winsorised
  , jackknifed
  -- ** Trivial L-Estimators
  , mean
  , total
  , lmin, lmax
  , midrange
  -- ** Sample-size-dependent L-Estimators
  , nthSmallest
  , nthLargest
  -- ** Quantiles
  -- *** Common quantiles
  , quantile
  , median
  , tercile, t1, t2
  , quartile, q1, q2, q3
  , quintile, qu1, qu2, qu3, qu4
  , percentile
  , permille
  -- *** Harrell-Davis Quantile Estimator
  , hdquantile
  -- *** Compute a quantile using a specified quantile estimation strategy
  , quantileBy
  -- * Sample Quantile Estimators
  , Estimator
  , Estimate(..)
  , r1,r2,r3,r4,r5,r6,r7,r8,r9,r10
  ) where

import Data.Ratio
import Data.List (sort)
import Data.IntMap (IntMap)
import qualified Data.IntMap as IM
import Data.Vector (Vector, (!))
import qualified Data.Vector as V
import Data.VectorSpace
import Statistics.Distribution.Beta
import qualified Statistics.Distribution as D

-- | L-estimators are linear combinations of order statistics used by 'robust' statistics.
newtype L r = L { runL :: Int -> IntMap r }

-- TODO: Write code to test if the result of a given L-estimator will be always than or equal
-- to the result of another L-estimator at a given sample size.

-- | Calculate the result of applying an L-estimator after sorting list into order statistics
(@@) :: (Num r, Ord r) => L r -> [r] -> r
l @@ xs = l @! V.fromList (sort xs)

-- | Calculate the result of applying an L-estimator to a *pre-sorted* vector of order statistics
(@!) :: Num r => L r -> Vector r -> r
L f @! v = IM.foldrWithKey (\k x y -> (v ! k) * x + y) 0 $ f (V.length v)

-- | get a vector of the coefficients of an L estimator when applied to an input of a given length
(@#) :: Num r => L r -> Int -> [r]
L f @# n = map (\k -> IM.findWithDefault 0 k fn) [0..n-1] where fn = f n

-- | calculate the breakdown % of an L-estimator
breakdown :: (Num r, Eq r) => L r -> Int
breakdown (L f)
  | IM.null m = 50
  | otherwise = fst (IM.findMin m) `min` (100 - fst (IM.findMax m))
  where m = IM.filter (/= 0) $ f 101

instance Num r => AdditiveGroup (L r) where
  zeroV = L $ \_ -> IM.empty
  L fa ^+^ L fb = L $ \n -> IM.unionWith (+) (fa n) (fb n)
  negateV (L fa) = L $ fmap negate . fa

instance Num r => VectorSpace (L r) where
  type Scalar (L r) = r
  x *^ L y = L $ fmap (x *) . y

clamp :: Int -> Int -> Int
clamp n k
  | k <= 0 = 0
  | k >= n = n - 1
  | otherwise = k

-- | The average of all of the order statistics. Not robust.
--
-- > breakdown mean = 0%
mean :: Fractional r => L r
mean = L $ \n -> IM.fromList [ (i, 1 / fromIntegral n) | i <- [0..n-1]]

-- | The sum of all of the order statistics. Not robust.
--
-- > breakdown total = 0%
total :: Num r => L r
total = L $ \n -> IM.fromList [ (i, 1) | i <- [0..n-1]]

-- | Calculate a trimmed L-estimator. If the sample size isn't evenly divided, linear interpolation is used
-- as described in <http://en.wikipedia.org/wiki/Trimmed_mean#Interpolation>

-- Trimming can increase the robustness of a statistic by removing outliers.

trimmed :: Fractional r => Rational -> L r -> L r
trimmed p (L g) = L $ \n -> case properFraction (fromIntegral n * p) of
  (w, 0)               -> IM.fromDistinctAscList [ (k + w, v)                        | (k,v) <- IM.toAscList $ g (n - w*2)]
  (w, f) | w' <- w + 1 -> IM.fromListWith (+) $  [ (k + w, fromRational (1 - f) * v) | (k,v) <- IM.toList $ g (n - w*2)] ++
                                                 [ (k + w', fromRational f  * v)     | (k,v) <- IM.toList $ g (n - w'*2)]

-- | Calculates an interpolated winsorized L-estimator in a manner analogous to the trimmed estimator.
-- Unlike trimming, winsorizing replaces the extreme values.
winsorized, winsorised :: Fractional r => Rational -> L r -> L r
winsorised = winsorized
winsorized p (L g) = L $ \n -> case properFraction (fromIntegral n * p) of
  (w, 0)               -> IM.fromAscListWith (+) [ (w `max` min (n - 1 - w) k, v) | (k,v) <- IM.toAscList (g n) ]
  (w, f) | w' <- w + 1 -> IM.fromListWith (+) $ do
     (k,v) <- IM.toList (g n)
     [ (w  `max` min (n - 1 - w ) k, v * fromRational (1 - f)),
       (w' `max` min (n - 1 - w') k, v * fromRational f)]

-- | Jackknifes the statistic by removing each sample in turn and recalculating the L-estimator,
-- requires at least 2 samples!
jackknifed :: Fractional r => L r -> L r
jackknifed (L g) = L $ \n -> IM.fromAscListWith (+) $ do
  let n' = fromIntegral n
  (k,v) <- IM.toAscList (g (n - 1))
  let k' = fromIntegral k + 1
  [(k, (n' - k') * v / n'), (k + 1, k' * v / n')]

-- | The most robust L-estimator possible.
--
-- > breakdown median = 50
median :: Fractional r => L r
median = L go where
  go n
    | odd n        = IM.singleton (div (n - 1) 2) 1
    | k <- div n 2 = IM.fromList [(k-1, 0.5), (k, 0.5)]

-- | Sample quantile estimators
data Estimate r  = Estimate {-# UNPACK #-} !Rational (IntMap r)
type Estimator r = Rational -> Int -> Estimate r

-- | Compute a quantile using the given estimation strategy to interpolate when an exact quantile isn't available
quantileBy :: Num r => Estimator r -> Rational -> L r
quantileBy f p = L $ \n -> case f p n of
  Estimate h qp -> case properFraction h of
    (w, 0) -> IM.singleton (clamp n (w - 1)) 1
    _      -> qp

-- | Compute a quantile with traditional direct averaging
quantile :: Fractional r => Rational -> L r
quantile = quantileBy r2

tercile :: Fractional r => Rational -> L r
tercile n = quantile (n/3)

-- | terciles 1 and 2
--
-- > breakdown t1 = breakdown t2 = 33%
t1, t2 :: Fractional r => L r
t1 = quantile (1/3)
t2 = quantile (2/3)

quartile :: Fractional r => Rational -> L r
quartile n = quantile (n/4)

-- | quantiles, with breakdown points 25%, 50%, and 25% respectively
q1, q2, q3 :: Fractional r => L r
q1 = quantile 0.25
q2 = median
q3 = quantile 0.75

quintile :: Fractional r => Rational -> L r
quintile n = quantile (n/5)

-- | quintiles 1 through 4
qu1, qu2, qu3, qu4 :: Fractional r => L r
qu1 = quintile 1
qu2 = quintile 2
qu3 = quintile 3
qu4 = quintile 4

-- |
--
-- > breakdown (percentile n) = min n (100 - n)
percentile :: Fractional r => Rational -> L r
percentile n = quantile (n/100)

permille :: Fractional r => Rational -> L r
permille n = quantile (n/1000)

nthSmallest :: Num r => Int -> L r
nthSmallest k = L $ \n -> IM.singleton (clamp n k) 1

nthLargest :: Num r => Int -> L r
nthLargest k = L $ \n -> IM.singleton (clamp n (n - 1 - k)) 1

-- |
--
-- > midhinge = trimmed 0.25 midrange
-- > breakdown midhinge = 25%
midhinge :: Fractional r => L r
midhinge = (q1 ^+^ q3) ^/ 2

-- | Tukey's trimean
--
-- > breakdown trimean = 25
trimean :: Fractional r => L r
trimean = (q1 ^+^ 2 *^ q2 ^+^ q3) ^/ 4

-- | The maximum value in the sample
lmax :: Num r => L r
lmax = L $ \n -> IM.singleton (n-1) 1

-- | The minimum value in the sample
lmin :: Num r => L r
lmin = L $ \_ -> IM.singleton 0 1

-- |
-- > midrange = lmax - lmin
-- > breakdown midrange = 0%
midrange :: Fractional r => L r
midrange = L $ \n -> IM.fromList [(0,-1),(n-1,1)]

-- | interquartile range
--
-- > breakdown iqr = 25%
-- > iqr = trimmed 0.25 midrange
iqr :: Fractional r => L r
iqr = q3 ^-^ q1

-- | interquartile mean
--
-- > iqm = trimmed 0.25 mean
iqm :: Fractional r => L r
iqm = trimmed 0.25 mean

-- | Direct estimator for the second L-moment given a sample
lscale :: Fractional r => L r
lscale = L $ \n -> let
     r = fromIntegral n
     scale = 1 / (r * (r-1))
  in IM.fromList [ (i - 1, scale * (2 * fromIntegral i - 1 - r)) | i <- [1..n] ]

-- | The Harrell-Davis quantile estimate. Uses multiple order statistics to approximate the quantile
-- to reduce variance.
hdquantile :: Fractional r => Rational -> L r
hdquantile q = L $ \n ->
  let n' = fromIntegral n
      np1 = n' + 1
      q' = fromRational q
      d = betaDistr (q'*np1) (np1*(1-q')) in
  if q == 0 then IM.singleton 0 1
  else if q == 1 then IM.singleton (n - 1) 1
  else IM.fromAscList
    [ (i, realToFrac $ D.cumulative d ((fromIntegral i + 1) / n') -
                       D.cumulative d (fromIntegral i / n'))
    | i <- [0 .. n - 1]
    ]

-- More information on the individual estimators used below can be found in:
-- http://stat.ethz.ch/R-manual/R-devel/library/stats/html/quantile.html
-- and
-- http://en.wikipedia.org/wiki/Quantile#Estimating_the_quantiles_of_a_population

-- | Inverse of the empirical distribution function
r1 :: Num r => Estimator r
r1 p n = Estimate (np + 1%2) $ IM.singleton (clamp n (ceiling np - 1)) 1
  where np = fromIntegral n * p

-- | .. with averaging at discontinuities
r2 :: Fractional r => Estimator r
r2 p n =  Estimate (np + 1%2) $
  if p == 0      then IM.singleton 0       1
  else if p == 1 then IM.singleton (n - 1) 1
  else IM.fromList [(ceiling np - 1, 0.5), (floor np, 0.5)]
  where np = fromIntegral n * p

-- | The observation numbered closest to Np. NB: does not yield a proper median
r3 :: Num r => Estimator r
r3 p n = Estimate np $ IM.singleton (clamp n (round np - 1)) 1
  where np = fromIntegral n * p

-- continuous sample quantiles
continuousEstimator ::
  Fractional r =>
  (Rational -> (Rational, Rational)) -> -- take the number of samples, and return upper and lower bounds on 'p = k/n' for which this interpolation should be used
  (Rational -> Rational -> Rational) -> -- take p = k/q, and n the number of samples, and return the coefficient h which will be used for interpolation when h is not integral
  Estimator r
continuousEstimator bds f p n = Estimate h $
  if p < lo then IM.singleton 0 1
  else if p >= hi then IM.singleton (n - 1) 1
  else case properFraction h of
    (w,frac) | frac' <- fromRational frac -> IM.fromList [(w - 1, frac'), (w, 1 - frac')]
  where
    r = fromIntegral n
    h = f p r
    (lo, hi) = bds r

-- | Linear interpolation of the empirical distribution function. NB: does not yield a proper median.
r4 :: Fractional r => Estimator r
r4 = continuousEstimator (\n -> (1 / n, 1)) (*)

-- | .. with knots midway through the steps as used in hydrology. This is the simplest continuous estimator that yields a correct median
r5 :: Fractional r => Estimator r
r5 = continuousEstimator (\n -> let tn = 2 * n in (1 / tn, (tn - 1) / tn)) $ \p n -> p*n + 0.5

-- | Linear interpolation of the expectations of the order statistics for the uniform distribution on [0,1]
r6 :: Fractional r => Estimator r
r6 = continuousEstimator (\n -> (1 / (n + 1), n / (n + 1))) $ \p n -> p*(n+1)

-- | Linear interpolation of the modes for the order statistics for the uniform distribution on [0,1]
r7 :: Fractional r => Estimator r
r7 = continuousEstimator (\_ -> (0, 1)) $ \p n -> p*(n-1) + 1

-- | Linear interpolation of the approximate medans for order statistics.
r8 :: Fractional r => Estimator r
r8 = continuousEstimator (\n -> (2/3 / (n + 1/3), (n - 1/3)/(n + 1/3))) $ \p n -> p*(n + 1/3) + 1/3

-- | The resulting quantile estimates are approximately unbiased for the expected order statistics if x is normally distributed.
r9 :: Fractional r => Estimator r
r9 = continuousEstimator (\n -> (0.625 / (n + 0.25), (n - 0.375)/(n + 0.25))) $ \p n -> p*(n + 0.25) + 0.375

-- | When rounding h, this yields the order statistic with the least expected square deviation relative to p.
r10 :: Fractional r => Estimator r
r10 = continuousEstimator (\n -> (1.5 / (n + 2), (n + 0.5)/(n + 2))) $ \p n -> p*(n + 2) - 0.5


{- |
This module provides normalized versions of the transforms in @fftw@.

All of the transforms are normalized so that

 - Each transform is unitary, i.e., preserves the inner product and the sum-of-squares norm of its input.

 - Each backwards transform is the inverse of the corresponding forwards transform.

(Both conditions only hold approximately, due to floating point precision.)

For more information on the underlying transforms, see
<http://www.fftw.org/fftw3_doc/What-FFTW-Really-Computes.html>.
-}
module Numeric.FFT.Vector.Unitary(
                -- * Creating and executing 'Plan's
                run,
                plan,
                execute,
                -- * Complex-to-complex transforms
                dft,
                idft,
                -- * Real-to-complex transforms
                dftR2C,
                dftC2R,
                -- * Discrete cosine transforms
                -- $dct_size
                dct2,
                idct2,
                dct4,
                ) where

import Numeric.FFT.Vector.Base
import qualified Numeric.FFT.Vector.Unnormalized as U
import Data.Complex
import qualified Data.Vector.Storable.Mutable as MS
import Control.Monad.Primitive(RealWorld)

-- | A discrete Fourier transform. The output and input sizes are the same (@n@).
--
-- @y_k = (1\/sqrt n) sum_(j=0)^(n-1) x_j e^(-2pi i j k\/n)@
dft :: Transform (Complex Double) (Complex Double)
dft = U.dft {normalization = \n -> constMultOutput $ 1 / sqrt (toEnum n)}

-- | An inverse discrete Fourier transform.  The output and input sizes are the same (@n@).
--
-- @y_k = (1\/sqrt n) sum_(j=0)^(n-1) x_j e^(2pi i j k\/n)@
idft :: Transform (Complex Double) (Complex Double)
idft = U.idft {normalization = \n -> constMultOutput $ 1 / sqrt (toEnum n)}

-- | A forward discrete Fourier transform with real data.  If the input size is @n@,
-- the output size will be @n \`div\` 2 + 1@.
dftR2C :: Transform Double (Complex Double)
dftR2C = U.dftR2C {normalization = \n -> modifyOutput $
                    complexR2CScaling (sqrt 2) n
        }

-- | A normalized backward discrete Fourier transform which is the left inverse of
-- 'U.dftR2C'.  (Specifically, @run dftC2R . run dftR2C == id@.)
--
-- This 'Transform' behaves differently than the others:
--
--  - Calling @plan dftC2R n@ creates a 'Plan' whose /output/ size is @n@, and whose
--    /input/ size is @n \`div\` 2 + 1@.
--
--  - If @length v == n@, then @length (run dftC2R v) == 2*(n-1)@.
--
dftC2R :: Transform (Complex Double) Double
dftC2R = U.dftC2R {normalization = \n -> modifyInput $
                    complexR2CScaling (sqrt 0.5) n
        }

complexR2CScaling :: Double -> Int -> MS.MVector RealWorld (Complex Double) -> IO ()
complexR2CScaling !t !n !a = do
    let !s1 = sqrt (1/toEnum n)
    let !s2 = t * s1
    let len = MS.length a
    -- Justification for the use of unsafeModify:
    -- The output size is 2n+1; so if n>0 then the output size is >=1;
    -- and if n even then the output size is >=3.
    unsafeModify a 0 $ scaleByD s1
    if odd n
        then multC s2 (MS.unsafeSlice 1 (len-1) a)
        else do
            unsafeModify a (len-1) $ scaleByD s1
            multC s2 (MS.unsafeSlice 1 (len-2) a)

-- $dct_size
-- Some normalized real-even (DCT).  The input and output sizes
-- are the same (@n@).

-- | A type-4 discrete cosine transform.  It is its own inverse.
--
-- @y_k = (1\/sqrt n) sum_(j=0)^(n-1) x_j cos(pi(j+1\/2)(k+1\/2)\/n)@
dct4 :: Transform Double Double
dct4 = U.dct4 {normalization = \n -> constMultOutput $ 1 / sqrt (2 * toEnum n)}

-- | A type-2 discrete cosine transform.  Its inverse is 'dct3'.
--
-- @y_k = w(k) sum_(j=0)^(n-1) x_j cos(pi(j+1\/2)k\/n);@
-- where
-- @w(0)=1\/sqrt n@, and @w(k)=sqrt(2\/n)@ for @k>0@.
dct2 :: Transform Double Double
dct2 = U.dct2 {normalization = \n -> modifyOutput $ \a -> do
    let n' = toEnum n
    let !s1 = sqrt $ 1 / (4*n')
    let !s2 = sqrt $ 1 / (2*n')
    unsafeModify a 0 (*s1)
    multC s2 (MS.unsafeSlice 1 (MS.length a-1) a)
    }

-- | A type-3 discrete cosine transform which is the inverse of 'dct2'.
--
-- @y_k = (-1)^k w(n-1) x_(n-1) + 2 sum_(j=0)^(n-2) w(j) x_j sin(pi(j+1)(k+1\/2)/n);@
-- where
-- @w(0)=1\/sqrt(n)@, and @w(k)=1/sqrt(2n)@ for @k>0@.
idct2 :: Transform Double Double
idct2 = U.dct3 {normalization = \n -> modifyInput $ \a -> do
    let n' = toEnum n
    let !s1 = sqrt $ 1 / n'
    let !s2 = sqrt $ 1 / (2*n')
    unsafeModify a 0 (*s1)
    multC s2 (MS.unsafeSlice 1 (MS.length a-1) a)
    }


module Camera where

import Types
import Numeric.Vector
import Numeric.Scalar
import VecUtil (unpack3)

rayFrom :: Camera -> (Double, Double) -> Ray
rayFrom (OrthoCamera p forward up right size) (x, y) = let dir = forward
                                                           x' = right * fromScalar (scalar (x * size))
                                                           y' = up * fromScalar (scalar (y * size))
                                                           origin = p + x' + y'
                                                       in  Ray origin dir
rayFrom (Camera p forward up right fov) (x, y) = let mul = tan $ fov * 0.00872664625 -- Magic number is to convert degrees to radians and divide by two
                                                     (r, u, f) = unpack3 $ normalized $ vec3 (x * mul) (y * mul) 1
                                                     dir = forward * (fromScalar $ scalar f) + right * (fromScalar $ scalar r) + up * (fromScalar $ scalar u)
                                                 in  Ray p dir


{-# LANGUAGE PatternGuards #-}

module Math
  (getInfo, predict, subTrend, randoms)
  where

import Control.Monad
import System.Random hiding (randoms)
import Data.List (transpose,minimumBy)
import Data.Function (on)
import Data.Packed.Matrix
import Numeric.LinearAlgebra.Algorithms

-- import Data.Random.RVar
-- import Data.Random.Distribution.Normal
-- import Data.Random.Source.DevRandom

import Types

solve :: [[Double]] -> [Double] -> [Double]
solve a b = head $ transpose $ toLists $ linearSolve (fromLists a) (fromLists $ map (:[]) b)

type SystemBuilder = Int -> [AnyNumber] -> [AnyNumber] -> ([[AnyNumber]], [AnyNumber])

systemLinear :: SystemBuilder
systemLinear n xs ys = ([[sx2,sx],[sx,fromIntegral n]], [sxy, sy])
  where
    sx = sum xs
    sx2 = sum $ map (^2) xs
    sy = sum ys
    sxy = sum $ zipWith (*) xs ys

systemSquare :: SystemBuilder
systemSquare n xs ys = ([[sx4, sx3, sx2], [sx3, sx2, sx], [sx2, sx, fromIntegral n]],
                        [sxy2, sxy, sy])
  where
    sx = sum xs
    sy = sum ys
    sx2 = sum $ map (^2) xs
    sx3 = sum $ map (^3) xs
    sx4 = sum $ map (^4) xs
    sxy = sum $ zipWith (*) xs ys
    sxy2 = sum $ zipWith (*) (map (^2) xs) ys

systemExponent :: SystemBuilder
systemExponent n xs ys = systemLinear n xs (map log ys)

mls :: SystemBuilder -> Int -> [AnyNumber] -> [AnyNumber] -> (AnyNumber, AnyNumber,AnyNumber)
mls system n xs ys =
  let (ma,mb) = system n xs ys
      ans = solve (map (map toDouble) ma) (map toDouble mb)
      (a,b,c) = case ans of
                  (u:v:[]) -> (u,v,0)
                  (u:v:w:[]) -> (u,v,w)
                  _ -> error "Unexpected number of coefficients!"
  in  (Number a,Number b, Number c)

getInfo :: Int -> Formula -> [AnyNumber] -> [AnyNumber] -> Info
getInfo n f xs ys | Auto <- f =
  let infoL = getInfo' n Linear xs ys
      infoS = getInfo' n Square xs ys
      infoE = getInfo' n Exponent xs ys
  in  minimumBy (compare `on` stdDev) [infoL, infoS, infoE]
                  | otherwise = getInfo' n f xs ys

getInfo' :: Int -> Formula -> [AnyNumber] -> [AnyNumber] -> Info
getInfo' n f xs ys = Info xs ys ts a b c formula
  where
    (a,b,c) = mls system n xs ys
    ts      = map formula xs
    (system, formula) =
      case f of
        Linear -> (systemLinear, linear a b)
        Square -> (systemSquare, square a b c)
        Exponent -> (systemExponent, exponential a b)

avgStep :: Info -> AnyNumber
avgStep info =
  let s = zipWith (-) (tail $ xvals info) (xvals info)
  in  ((sum s)/(fromIntegral $ length s))

stdDev :: Info -> AnyNumber
stdDev info =
  let ts = trend info
      ys = yvals info
      n = length ts
  in  sqrt $ (sum $ map (^2) $ zipWith (-) ys ts)/(fromIntegral n-1)

-- randomsD :: Double -> Int -> IO [Double]
-- randomsD s n = do
--   replicateM n (runRVar (normal 0 s) DevRandom)
--
randoms :: AnyNumber -> Int -> IO [AnyNumber]
randoms s n = do
  lst <- randomsD (toDouble s) n
  return $ map Number lst

randomsD :: Double -> Int -> IO [Double]
randomsD s n = replicateM n $ do
  r1 <- randomIO
  r2 <- randomIO
  return $ s * cos(2*pi*r1)*sqrt(-2*log(r2))

predict :: Int -> Bool -> Info -> IO Info
predict n r info@(Info xs ys ts a b c func) = do
  let st = avgStep info
      xs' = take n $ tail $ iterate (+st) (last xs)
      ys' = map func xs'
  rs <- if r
         then randoms (stdDev info) n
         else return $ replicate n 0
  let ys'' = zipWith (+) ys' rs
  return $ Info (xs++xs') (ys++ys'') (ts++ys') a b c func

subTrend :: Info -> Info
subTrend (Info xs ys ts a b c f) = Info xs ys' ts a b c f
  where
    ys' = zipWith (-) ys ts

linear :: AnyNumber -> AnyNumber -> AnyNumber -> AnyNumber
linear a b x = a*x + b

square :: AnyNumber -> AnyNumber -> AnyNumber -> AnyNumber -> AnyNumber
square a b c x = a*x^2 + b*x + c

exponential :: AnyNumber -> AnyNumber -> AnyNumber -> AnyNumber
exponential a b x = exp (a*x + b)


{-#LANGUAGE BangPatterns #-}

module Tests.Correlation
    ( tests ) where

import Test.Framework
import Test.Framework.Providers.QuickCheck2
import Test.Framework.Providers.HUnit
import Test.HUnit (Assertion, (@=?))
import qualified Data.Vector as V
import Statistics.Correlation.Kendall

tests :: Test
tests = testGroup "Correlation"
    [ testProperty "Kendall test -- general" testKendall
    , testCase "Kendall test -- special cases" testKendallSpecial
    ]

testKendall :: [(Double, Double)] -> Bool
testKendall xy | isNaN r1 = isNaN r2
               | otherwise = r1 == r2
  where
    r1 = kendallBruteForce xy
    r2 = kendall $ V.fromList xy

testKendallSpecial :: Assertion
testKendallSpecial = ys @=? map (kendall.V.fromList) xs
  where
    (xs, ys) = unzip testData
    testData :: [([(Double, Double)], Double)]
    testData = [ ( [(1,1), (2,2), (3,1), (1,5), (2,2)], -0.375 )
               , ( [(1,3), (1,3), (1,3), (3,2), (3,5)], 0)
               ]

kendallBruteForce :: [(Double, Double)] -> Double
kendallBruteForce xy = (n_c - n_d) / sqrt ((n_0 - n_1) * (n_0 - n_2))
  where
    allPairs = f xy
    (n_c, n_d, n_1, n_2) = foldl g (0,0,0,0) allPairs
    n_0 = fromIntegral.length $ allPairs
    g (!nc, !nd, !n1, !n2) ((x1, y1), (x2, y2))
      | (x2 - x1) * (y2 - y1) > 0 = (nc+1, nd, n1, n2)
      | (x2 - x1) * (y2 - y1) < 0 = (nc, nd+1, n1, n2)
      | otherwise = if x1 == x2
                       then if y1 == y2
                               then (nc, nd, n1+1, n2+1)
                               else (nc, nd, n1+1, n2)
                       else (nc, nd, n1, n2+1)
    f (x:xs) = zip (repeat x) xs ++ f xs
    f _ = []