module DecisionTree where

import Data.Function    ((&))
import Foreign.Storable (Storable)

import           Data.Map.Strict            (Map)
import qualified Data.Map.Strict            as M
import qualified Data.Vector.Storable       as V
import           Numeric.LinearAlgebra
import           Numeric.LinearAlgebra.Data as N

{-
building a decision tree:

- get info gains for each column
- group by top scoring column
- repeat, with smaller matrices

-}

{-
let v = fromList [0, 0, 1, 0, 0, 1, 0, 0, 2, 2, 0, 2] :: Vector I

entropy v
-}

entropy :: (Ord a, Storable a) => Vector a -> Double
entropy vec =
    let
        n =
            fromIntegral $ V.length vec

        proportions =
            count vec
                & M.elems
                & map fromIntegral
                & map (/ n)

        componentEntropies =
            [ -p * logBase 2 p | p <- proportions ]
    in
        sum componentEntropies

count :: (Ord a, Storable a) => Vector a -> Map a Int
count vec =
    V.foldl insertCount M.empty vec
    where
        insertCount :: Ord a => Map a Int -> a -> Map a Int
        insertCount countDict el = M.insertWith (+) el 1 countDict

{-
let v1 = fromList [0, 0, 1, 0, 0, 1, 0, 0, 2, 2, 0, 2] :: Vector I
let v2 = fromList [1, 2, 3, 4, 5, 5, 3, 5, 6, 8, 9, 5] :: Vector I

informationGain v1 v2
-}

informationGain :: (Ord a, Storable a) => Vector a -> Vector a -> Double
informationGain resultVector featureVector =
    let
        groupedResults =
            groupByVector featureVector resultVector

        n =
            fromIntegral $ V.length featureVector

        vecLen vec = fromIntegral (V.length vec)

        weightedEntropies =
            groupedResults
                & M.elems
                & map (\v -> (vecLen v / n) * entropy v)
    in
        sum weightedEntropies

informationGains :: Matrix I -> [Double]
informationGains mat =
    let
        featureVectors = N.toColumns mat
        resultVector = lastColumn mat
    in
        map (informationGain resultVector) featureVectors

groupByVector :: (Storable a, Storable b, Ord a) => Vector a -> Vector b -> Map a (Vector b)
groupByVector groupingVector targetVector =
    let
        groupedIndices =
            indicesMap groupingVector

        getAtIndices vec indices =
            V.ifilter (\i _ -> elem i indices) vec
    in
        groupedIndices
            & M.map (getAtIndices targetVector)

groupBy :: (Ord a, Element a) => Matrix a -> Int -> Map a (Matrix a)
groupBy mat columnIndex =
    let
        groupingVector =
            mat ¿ [columnIndex]
                & flatten

        groupedIndices =
            indicesMap groupingVector
    in
        groupedIndices
            & M.map (\v -> removeColumn (mat ? v) columnIndex)

indicesMap :: (Ord a, Storable a) => Vector a -> Map a [Int]
indicesMap vec =
    V.ifoldl insertIndex M.empty vec
    where
        insertIndex :: Ord a => Map a [Int] -> Int -> a -> Map a [Int]
        insertIndex indicesDict i el =
            M.insertWith (++) el [i] indicesDict

removeColumn :: (Element a) => Matrix a -> Int -> Matrix a
removeColumn m i =
    let
        m1 = m ?? (All, Take i)
        m2 = m ?? (All, Drop (i+1))
    in
    m1 ||| m2

lastColumn :: Matrix I -> Vector I
lastColumn m =
    let (_, width) = size m
    in flatten $ m ¿ [width-1]


{-|
Module      : Mandelbrot
Description : Algorithm of Mandelbrot functions

Module which get resolution of image and points and calculate
correct color to point based on Mandelbrot algorithm
-}
module Mandelbrot (drawMandelbrot, runMandelbrotTests, runQuickCheckMandelbrot) where
import Data.Complex
import Graphics.GD
import Plot
import Test.HUnit
import Test.QuickCheck

-- |Max number of iteration to approximate set
maxIteration = 150

-- |Main function, invoke draw Mandelbrot set
drawMandelbrot :: Size -- ^ Resolution of generating image
                -> Bound -- ^ Four coordinates of funtion bound
                -> String -- ^ Name of generating file
                -> IO() -- ^ Result: Save file
drawMandelbrot size bound fileName = drawPlot drawMandelbrot_ size bound fileName

-- |Generating colors for points
drawMandelbrot_ :: Coord -- ^ Point to draw
                -> Color -- ^ Color of the point calculated from mandelbrot formula
drawMandelbrot_ (x,y) = rgb (floor $ f * 255) (floor $ f * 120) 0
                        where i = mandelbrot (x :+ y)
                              f = fromIntegral i / fromIntegral maxIteration

-- |Recursive formula for points
mandelbrot :: Complex Double -- ^ Checking point
              -> Int -- ^ Number of iteration
mandelbrot z = length $ takeWhile (\c -> magnitude c <= 2) $ take maxIteration $ iterate (\w -> w^2 + z) 0

test1 = TestCase (assertEqual "color1" 16742400 (drawMandelbrot_ (0,0)))
test2 = TestCase (assertEqual "color2" 196864 (drawMandelbrot_ (1,1)))
test3 = TestCase (assertEqual "color3" 65536 (drawMandelbrot_ (10,10)))
test4 = TestCase (assertEqual "mandelbrot1" 2 (mandelbrot (1 :+1)))
test5 = TestCase (assertEqual "mandelbrot2" 1 (mandelbrot (0 :+ 20)))
test6 = TestCase (assertEqual "mandelbrot3" 1 (mandelbrot (123 :+ 0)))
test7 = TestCase (assertEqual "mandelbrot3" 150 (mandelbrot (0 :+ 0)))

tests = TestList [TestLabel "test1" test1, TestLabel "test2" test2, TestLabel "test3" test3, TestLabel "test4" test4, TestLabel "test5" test5, TestLabel "test6" test6, TestLabel "test7" test7]

-- |Run tests
runMandelbrotTests :: IO Counts
runMandelbrotTests = do runTestTT tests

max_Color :: Coord -> Bool
max_Color (x,y) =   drawMandelbrot_ (x,y)  <= 16777216

min_Color :: Coord -> Bool
min_Color (x,y) =   drawMandelbrot_ (x,y)  >= 0

-- |Run quick tests
runQuickCheckMandelbrot = do quickCheck max_Color >> quickCheck min_Color


import Data.Array.CArray
import Data.Complex
import Data.List (tails)
import Math.FFT (dft, idft)

convolution :: (Num a) => [a] -> [a] -> [a]
convolution x = map (sum . zipWith (*) (reverse x)) . spread
  where
    spread = init . tails . (replicate (length x - 1) 0 ++)

convolutionFFT :: [Complex Double] -> [Complex Double] -> [Complex Double]
convolutionFFT x y = elems $ idft $ liftArray2 (*) (fft x) (fft y)
  where
    fft a = dft $ listArray (1, length a) a

main :: IO ()
main = do
  let x = [1, 2, 1, 2, 1]
      y = [2, 1, 2, 1, 2]
  print $ convolution x y
  print $ convolutionFFT x y


--Mihai Ciobanu
--Student number : s1519734
import Graphics.UI.GLUT
import Graphics.Rendering.OpenGL
import Control.Monad
import Data.Complex -- we are using complex numbers
--some functions included in Data.Complex: x :+ y = x +yi, magnitude z = sqrt ((Re z)^2 + (Im z)^2)).

 ----------------------Settings-------------------------------------------------------------------------------------------------------------------------------------------------------------
type Saturation = GLfloat
shsv = 0.8 :: Saturation -- controls the saturation of the colors used + is included in the interval [0,1]

type ValueH = GLfloat
valueH = 1 :: ValueH -- value that influences the outcome of the coloring + is included in the interval [0,1]

type Chroma = GLfloat
chroma = shsv * valueH -- the chroma is calculated as the product of the saturation and the value

diff = valueH - chroma -- it's used to smooth out the colors

type Smoothcolor = GLfloat -- HSV value that will be converted to RGB

maxiter = 1000 :: Int -- maximum number of iterations - the bigger it is, the clearer the image obtained, but the program will render the image a lot slower.
constant = 0.32 :+ 0.42 :: Complex GLfloat --constant that will directly influence the formula used for generating the image - we will obtain different fractals for different constants
bound = 20 :: GLfloat -- the constant that decides when a value "escapes".
nx = 1.5 :: GLfloat -- controls the zoom of the image as explained before isJulia

wSize = Size 700 700 -- controls the size of the window
width = 1000 -- controls the width of the image
height = 1000 -- controls the height of the image
wxh = (fromIntegral width) * (fromIntegral height) -- constant we'll be needing in our calculations.

--We need to adjust the width and height of certain fractals in order to center the image.
valuewl = 0 :: Int
valuewr = 0 :: Int
valuehu = 0 :: Int
valuehd = 0 :: Int

red1 = 0.3 :: GLfloat
green1 = 1.3 :: GLfloat
blue1 = 1 :: GLfloat
red2 = 1 :: GLfloat
green2 = 1 :: GLfloat
blue2 = 1 :: GLfloat
red3 = 1 :: GLfloat
green3 = 1 :: GLfloat
blue3 = 1 :: GLfloat
red4 = 1 :: GLfloat
green4 = 1 :: GLfloat
blue4 = 1 :: GLfloat
red5 = 1 :: GLfloat
green5 = 1 :: GLfloat
blue5 = 1 :: GLfloat
red6 = 1 :: GLfloat
green6 = 1 :: GLfloat
blue6 = 1 :: GLfloat -- these variables allow us a greater freedom in coloring our fractal - must be bigger than 0, preferrably in the interval [0..1]
--Selects the function used to do the iterations.
function :: Complex GLfloat -> Complex GLfloat
function = f1

--Different functions with their settings for different fractals (specified in brackets)
f1 :: Complex GLfloat -> Complex GLfloat
f1 z = z ^ 2
-- Method of coloring with exp (nx = 2) : -0.835 + 0.2321i (constant1), (-0.8) :+ 0.156 (constant2) , 0.285 :+ 0.01 (constant3)
-- Method of coloring with log (nx = 1) : -0.79 :+ 0.15 (constant 1), .28 + .008 (constant2), (-0.163) :+ 1.04 (constant3), (-0.12) :+ (-0.77) (constant 4) (-0.75) :+ (0.1) (constant5 - r1 0.3 g1 1.3 b1 1), 0.21 :+ 0.58 (constant6),
-- 0.15 :+ 0.58 (constant 7) 0.28 :+ 0.49 (constant 8), 0.32 :+ 0.42 (constant 9), (-0.57) :+ 0.48 (constant 10), (-0.86) :+ 0.21 (constant 11)

--Failed attempts (with examples) :
f2 :: Complex GLfloat -> Complex GLfloat
f2 z = z ^ 5 -- constant = 0.544 :+ 0 bound = 8 nx = 2 shsv = 0.8 valueH = 1 maxiter = 1000 width = 1000 height = 1000 value__ = 0 red 1 green 0.8 blue 0.5

f3 :: Complex GLfloat -> Complex GLfloat
f3 z = sqrt (sinh (z ^ 2)) -- constant 0.544 :+ 0 bound = 5 nx = 2 shsv = 0.8 valueH = 1 maxiter = 1000 height = 1000 value__ = 0 red 1 green 1 blue 1
--constant 0.065,0.122i
f4 :: Complex GLfloat -> Complex GLfloat
f4 z = (z ^ 2 / 2) + 0.26 * (z ^ 4)

--OBS : keep constant 0 :+ 0 and bound 20 and maxiter 50 for the next ones.
g5 :: Complex GLfloat -> Complex GLfloat
g5 z = z ^ 5 - 1

g5' :: Complex GLfloat -> Complex GLfloat
g5' z = 5 * (z ^ 4)

f5 :: Complex GLfloat -> Complex GLfloat
f5 z | magnitude z > 0.0001 = z - (g5 z) / (g5' z)
     | otherwise = (0 :+ 0)

g6 :: Complex GLfloat -> Complex GLfloat
g6 z = sin z

g6' :: Complex GLfloat -> Complex GLfloat
g6' z = cos z
--f6 creates two vertical "lines"
f6 :: Complex GLfloat -> Complex GLfloat
f6 z | magnitude (cos z) > 0.0001 = z - ((g6 z) / (g6' z))
     | otherwise = f1 z

f7 :: Complex GLfloat -> Complex GLfloat
f7 z | magnitude (z) > 0.00001 = (1 + 3 * z ^ 4) / (4 * z ^ 3)
     | otherwise = f1 z

f8 :: Complex GLfloat -> Complex GLfloat
f8 z | magnitude (cos z) > 0.005 = z - (sin z) / (cos z)
     | otherwise = f1 z

c1f9 = 0 :+ (-1) :: Complex GLfloat
c2f9 = 0.23 :+ 0.678 :: Complex GLfloat

f9 :: Complex GLfloat -> Complex GLfloat
f9 z = 1 + c1f9 * z ^ 5
--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------

----------------------OpenGL stuff----------------------------------------------------------------------------------------------------------------------------------------------------------
main :: IO ()
main = do
  (_progName, _args) <- getArgsAndInitialize
  _window <- createWindow "(Not Responding) (Not Responding) (Not Responding) Julia Set (Not Responding) (Not Responding)" --just a way to vent my frustration, ignore it
  keyboardMouseCallback $= Just keyboardMouse
  initialDisplayMode $= [DoubleBuffered]
  windowSize $= wSize
  reshapeCallback $= Just reshape
  displayCallback $= display
  actionOnWindowClose $= MainLoopReturns
  mainLoop

keyboardMouse :: KeyboardMouseCallback
keyboardMouse _key _state _modifiers _position = return ()

display = do
   clear [ColorBuffer]
   loadIdentity
   scale 0.001 0.001 (0.2::GLfloat) -- rescales the picture
   forM_ points $ \(x,y,z,t) -> -- plots the points and their color
    preservingMatrix $ do
      juliaa points
      swapBuffers
   flush

juliaa :: [(GLfloat,GLfloat,GLfloat,Color3 GLfloat)] -> IO ()
juliaa asdf = renderPrimitive Points $ do
            mapM_ drawing asdf
                where
                  drawing (x,y,z,c) = do
                  color c
                  vertex $ Vertex3 x y z

reshape :: ReshapeCallback
reshape size = do
  viewport $= (Position 0 0, size)
  postRedisplay Nothing

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

---------Functions used to generate the fractal-based image---------------------------------------------------------------------------------------------------------------------------------

-- This function decides if a point is going to be colored.
funcRec :: Complex GLfloat -> Int -> Smoothcolor -> (Bool,Smoothcolor)
funcRec z 0 sm = (False, 0)
funcRec z n sm = if (magnitude fz < bound) then funcRec (fz) (n-1) (sm) else (True, (fromIntegral (maxiter - n) - logBase 2 ((log (magnitude fz)) / (log bound))) / fromIntegral (maxiter)) where fz = function z + constant
--A smooth coloring method (will work for the majority of our functions)
-- funcRec z n sm = if ((magnitude z) < bound) then funcRec ((function z) + constant) (n-1) (sm + exp(-magnitude z)) else (True, (sm / fromIntegral (maxiter - n)))

--Calculates the mod 2 of a GLfloat - the function is needed for converting HSV to RGB in coloring.
modulo2 :: Smoothcolor -> Smoothcolor
modulo2 x | x < 2 = x
          | x > 2 = modulo2 (x-2)

--A small transformation to further smooth our coloring.
intermvalueX :: Smoothcolor -> GLfloat
intermvalueX sm = chroma * (1 - abs ((modulo2 sm) - 1))

--Converts HSV to RGB.
--We can have fun by making one of those lines black in order to see exactly which points are affected by the coloring.
coloring :: Smoothcolor -> Color3 GLfloat
coloring u | u == 0 = Color3 0 0 0
           | u >= 0 && u < 1 = Color3 ((valueH) * red1) (((intermvalueX u) + diff) * green1) ((diff) * blue1)
           | u >= 1 && u < 2 = Color3 (((intermvalueX u) + diff) * red2) ((valueH) * green2) ((diff) * blue2)
           | u >= 2 && u < 3 = Color3 ((diff) * red3) ((valueH) * green3) (((intermvalueX u) + diff) * blue3)
           | u >= 3 && u < 4 = Color3 ((diff) * red4) (((intermvalueX u) + diff) * green4) ((valueH) * blue4)
           | u >= 4 && u < 5 = Color3 (((intermvalueX u) + diff) * red5) ((diff) * green5) ((valueH) * blue5)
           | u >= 5 && u < 6 = Color3 ((valueH) * red6) ((diff) * green6) (((intermvalueX u) + diff) * blue6)

-- nx - bigger number => zoom out the picture, smaller => zoom in (I've done this before I discovered I could scale the picture in the OpenGL panel). A more correct name of the function would be isNotJulia, but we'll stick to this notation.
isJulia :: Complex GLfloat -> (Bool,Smoothcolor)
isJulia (x :+ y) = funcRec ((nx * x / (fromIntegral width)) :+ (nx * y / (fromIntegral height))) maxiter 0

-- Here we evaluate each point in the rectangle (-width, -height) (width, height) and color it if it's not in the Julia Set.
juliaP :: [(GLfloat,GLfloat,Smoothcolor)]
juliaP = [(x,y,snd j) | x <- [fromIntegral (-width + valuewl)..fromIntegral (width + valuewr)], y <- [fromIntegral (-height + valuehd)..fromIntegral (height + valuehu)], let j = isJulia (x :+ y), fst j]

--The points we have to plot and their color - "black" if the point is in the Julia Set, colored if it's not (if it escapes after a number of iterations)
points :: [(GLfloat,GLfloat,GLfloat,Color3 GLfloat)]
points =  [(x,y,0,(coloring (sm * 6))) | (x,y,sm) <- juliaP]
--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------

----------------------------References------------------------------------------------------------------------------------------------------------------------------------------------------
--https://en.wikipedia.org/wiki/HSL_and_HSV                                                   - for information on how HSV works and how to convert it to RGB;
--http://stackoverflow.com/questions/369438/smooth-spectrum-for-mandelbrot-set-rendering      - one smooth coloring algorithm for Julia Sets;
--http://hackage.haskell.org/package/GLUT-2.7.0.3/docs/Graphics-UI-GLUT-Callbacks-Window.html - information about OpenGL.
--https://wiki.haskell.org/OpenGLTutorial1#Lines                                              -
--https://wiki.haskell.org/OpenGLTutorial2                                                    - tutorial on OpenGL.
--http://www.juliasets.dk/Pictures_of_Julia_and_Mandelbrot_sets.htm                           - useful information about Julia Sets, Mandelbrot fractals and Newton fractals.
--http://www.karlsims.com/julia.html                                                          - interesting Julia Fractals


-- {-# LANGUAGE DeriveDataTypeable #-}

module Main where

-- import Lib

import Control.Monad.State
import Data.List (transpose)
import Data.Vector (snoc)
import Numeric.LinearAlgebra
  ( Linear (scale),
    Matrix,
    Vector,
    cmap,
    matrix,
    maxIndex,
    outer,
    vector,
    (#>),
  )
import Text.CSV (parseCSVFromFile)

type WeightValue = Matrix Double

type WeightState = Matrix Double

predict ::
  Matrix Double -> -- curried - Weights
  Vector Double -> -- curried - Input vector
  Vector Double -- Guess
predict w x = cmap sigmoid w #> x

predictEpoch ::
  Matrix Double -> -- w - Weights
  [Vector Double] -> -- x - An epoch of input vectorors
  [Vector Double] -- A guess for each input in the epoch
predictEpoch w x = map (cmap sigmoid) (fmap (w #>) x)

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

stocTrain ::
  Vector Double -> -- x - Single input vectoror
  Vector Double -> -- l - label
  State WeightState () -- Weights
stocTrain x l = do
  w <- get
  let error = predict w x - l
      sigDeriv = error * cmap (1 -) error
      delta = scale 0.01 $ sigDeriv `outer` x
      newWeights = w - delta
  put newWeights

train ::
  Int -> -- t - number of epochs to train on
  [Vector Double] -> -- x - An epoch of inputs
  [Vector Double] -> -- l - An epoch of labels
  State WeightState WeightValue -- Monad stuff
train t xs ls
  | t == 0 = get
  | otherwise = do
    trainEpoch xs ls
    train (t - 1) xs ls

trainEpoch ::
  [Vector Double] -> -- x - An epoch of inputs
  [Vector Double] -> -- l - An epoch of labels
  State WeightState WeightValue -- Weights
trainEpoch [] [] = get
trainEpoch (x : xs) (l : ls) = do
  stocTrain x l
  trainEpoch xs ls

test ::
  Matrix Double -> -- w - Weights
  [Vector Double] -> -- x - An epoch of input vectors
  [Vector Double] -> -- l - An epoch of labels
  Double -- A guess for each input in the epoch
test w x l = intDiv correct total
  where
    correct = sum (zipWith checkSame l (predictEpoch w x))
    total = length l

intDiv :: Int -> Int -> Double
intDiv n l = fromIntegral n / fromIntegral l

checkSame :: Vector Double -> Vector Double -> Int
checkSame x y
  | maxIndex x == maxIndex y = 1
  | otherwise = 0

run ::
  [[String]] -> -- tesCSV - read csv of test data
  [[String]] -> --  trainCSV - read csv of train data
  Double -- correct / total
run testCSV trainCSV = test w xTest lTest
  where
    (xTest, lTest) = csvToInputsLabels testCSV
    w = evalState (train 1 xTrain lTrain) initW
    (xTrain, lTrain) = csvToInputsLabels trainCSV
    initW = matrix 785 (replicate (785 * 10) 0)

csvToInputsLabels ::
  [[String]] -> -- read csv
  ([Vector Double], [Vector Double]) -- (x,l)
csvToInputsLabels csv = (finInputs, oneHotLabels)
  where
    oneHotLabels = fmap oneHotVector rawLabels
    finInputs = fmap vector scaledBiasedInputs
    scaledBiasedInputs = fmap (++ [1.0]) (scaleInput rawInputs)
    (rawLabels, rawInputs) = labelInputSplit (map (map read) (init csv))

scaleInput :: [[Double]] -> [[Double]]
scaleInput = map (map (/ 255))

labelInputSplit :: [[Double]] -> ([Double], [[Double]])
labelInputSplit xs = (head t, transpose (tail t))
  where
    t = transpose xs

oneHotVector :: Double -> Vector Double
oneHotVector x = vector $ replicate int 0 ++ [1] ++ replicate (9 - int) 0
  where
    int = round x

main :: IO ()
main = do
  -- testCSV <- parseCSVFromFile "./data/mnist_test.csv"
  -- trainCSV <- parseCSVFromFile "./data/mnist_train.csv"
  -- case testCSV of
  --   Left err -> print err
  --   Right testCSV ->
  --     case trainCSV of
  --       Left err -> print err
  --       Right trainCSV -> print (run testCSV testCSV)

  let inputs = map vector [[0, 0, 1], [0, 1, 1], [1, 0, 1], [1, 1, 1]]
  let labels = map vector [[0, 1], [0, 1], [0, 1], [1, 0]]
  let weights = matrix 3 (replicate 6 0)
  print "Initialized Weights"
  print weights
  print "target"
  print labels
  print "untrained prediction"
  print $ predictEpoch weights inputs
  print $ test weights inputs labels
  let trainedWeights = evalState (train 10 inputs labels) weights
  print "trained weights"
  print trainedWeights
  print "target"
  print labels
  print "trained prediciton"
  print $ predictEpoch trainedWeights inputs
  print "Test"
  print $ test trainedWeights inputs labels


module Main where

import Numeric.NLOPT
-- import Numeric.LinearAlgebra (dot)
import Numeric.LinearAlgebra (fromList, toList)

main :: IO ()
main = do
    let -- objf x = x `dot` x + 22
        -- objf x = sum (toList x)
        objf x = let [x1, x2, x3] = toList x
                 in 3 * x1 * x1 + 2 * x1 * x2 + x1 * x3
                    + 2.5 * x2 * x2 + 2 * x2 * x3 + 2 * x3 * x3
                    - 8 * x1 - 3 * x2 - 3 * x3

        stop = ObjectiveRelativeTolerance 1e-9 :| []
        -- algorithm = SBPLX objf [] Nothing
        algorithm = NELDERMEAD objf [] Nothing
        subproblem = LocalProblem 3 stop algorithm
        x0 = fromList [0, 0, 0]

        solU = minimizeLocal subproblem x0

    print solU

    -- define constraint function:
    let -- constraintf x = sum (toList x) - 1.0
        -- constraintf x = let [x1, x2] = toList x
        --                 in x1 * x1 + x2 * x2 - 2
        c1f x = let [x1, _, x3] = toList x
                in x1 + x3 - 3
        c2f x = let [_, x2, x3] = toList x
                in x2 + x3

    -- define constraint object to pass to the algorithm:
    let -- constraint = EqualityConstraint (Scalar constraintf) 1e-6
        c1 = EqualityConstraint (Scalar c1f) 1e-6
        c2 = EqualityConstraint (Scalar c2f) 1e-6
        -- problem = AugLagProblem [constraint] [] (AUGLAG_EQ_LOCAL subproblem)
        problem = AugLagProblem [c1, c2] [] (AUGLAG_EQ_LOCAL subproblem)

        -- x = (2, -1, 1)
        solC = minimizeAugLag problem x0

    print solC

    putStrLn "-- Done."


-- ch3 example 3.16

import Statistics.Distribution
import Statistics.Distribution.Hypergeometric

h :: HypergeometricDistribution
h = hypergeometric m l k
  where
    m =  5 --      r
    l = 20 --        N of population size
    k = 10 --          n of samplie size
--                     k elements chosen from
--                   l population, with
--                 m elements of one type and (l-m) of the other

main :: IO ()
main = do
  putStrLn "Among 20, 10 are randomely selected from; what is the probability that the 10 include 5 best."
  putStrLn $ "The syntax is m=r, l=N, k=r, y=k and:"
  putStrLn $ "p 5 = " ++ show (probability h 5)


module Fuml.Base.PCA where

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

sumWith :: Num a => (Vector Double -> a) -> [Vector Double] -> a
sumWith f (v:vs) = foldl' (\acc v -> acc + f v) (f v) vs

meanVs :: [Vector Double] -> Vector Double
meanVs vecs = scale (recip $ realToFrac $ length vecs) $ sumWith id vecs

covVs :: Vector Double -> [Vector Double] -> Herm Double
covVs xmn xs
   = let k = length xs
         f xi = outerSame $ xi - xmn
     in sym $ scale (recip $ realToFrac k - 1 ) $ sumWith f xs

outerSame v = outer v v

--following https://github.com/albertoruiz/hmatrix/blob/master/examples/pca2.hs
type Stat = (Vector Double, [Double], Matrix Double)

statVs :: [Vector Double] -> Stat
statVs vs = (m, toList s, tr' v) where
    m = meanVs vs
    (s,v) = eigSH (covVs m vs)

pcaN :: Int -> Stat -> (Vector Double -> Vector Double , Vector Double -> Vector Double )
pcaN n (m,s,v) = (encode,decode)
  where
    encode x = vp #> (x - m)
    decode x = x <# vp + m
    vp = takeRows n v


module Main where

import Plot
import Graphics.GD
import Data.Complex

imageDim :: Size
imageDim = (1200, 400)
-- 7200, 4800

--imageDim :: Size
--imageDim = (600, 600)

-- How many iterations before  we declare the point doesn't diverge
maxIter = 64

-- Coordinates in the complex plane to graph
-- First coordinate must be top left.
mbWindow :: Window
mbWindow = ( (-3,1), (3,-1) )
-- mbWindow = ( (-3,2), (3, -2))

--mbWindow :: Window
--mbWindow = ( (-1.5,1.5), (1.5,-1.5) ) -- Full view

-- C value
c :: Complex Double
c = (-0.423) :+ 0.745 -- Dust
-- c = (0.45) :+ (-0.1428)
-- c = (0.285 :+ 0.01)
-- c = (-0.45) :+ (-0.143) -- hazycloud
-- c = (-0.54) :+ (0.54) -- dragon
-- c = 0 :+ (0.8) -- veins
-- c = 0.355 :+ (0.355) -- swirlybois

main :: IO ()
main = drawPlot drawJulia imageDim mbWindow "juliaset.png"

drawJulia :: Coordinate -> Color
drawJulia (x,y) = colorIterations $ julia c (x :+ y) 0

-- Count number of iterations for the given point to diverge
-- Non-diverging points return 0
julia :: Complex Double -- Constant C
      -> Complex Double -- Iterating Z value
      -> Int -- Current iteration
      -> Int -- Iterations before diverging
julia c z iter
    | iter > maxIter = 0
    | otherwise = let z' = z^2 + c in
                  if magnitude z' > 2
                  then iter
                  else julia c z' (iter+1)

-- original colors
--colorIterations :: Int -> Color
--colorIterations x
--    | x > maxIter = rgb 255 255 255
--    | otherwise = let c = x*4
--                  in rgb x c x

-- the version we used for the purple julias
colorIterations :: Int -> Color
colorIterations x
    | x > maxIter = rgb 255 255 255
    | otherwise = let c = 50 + x
                      x'' = 50 + x*3 -- addition because to avoid fading to black
                      x' = 100 + x*2
                  in rgb x'' c x'


module FFT.Orig (
    fft,
    dft,
) where
import Data.Complex
import FFT.Samples
import System.Environment
import Control.Parallel
import Strategies

-- twiddle factors
tw :: Int -> Int -> Complex Float
tw n k = cis (-2 * pi * fromIntegral k / fromIntegral n)

-- Discrete Fourier Transform -- O(n^2)
dft :: [Complex Float] -> [Complex Float]
dft xs = [ sum [ xs!!j * tw n (j*k) | j <- [0..n']] | k <- [0..n']]
  where
    n = length xs
    n' = n-1

-- Fast Fourier Transform

-- In case you are wondering, this is the Decimation in Frequency (DIF)
-- radix 2 Cooley-Tukey FFT

fft :: [Complex Float] -> [Complex Float]
fft [a] = [a]
fft as = interleave ls rs
  where
    (cs,ds) = bflyS as
    ls = fft cs
    rs = fft ds

interleave [] bs = bs
interleave (a:as) bs = a : interleave bs as

bflyS :: [Complex Float] -> ([Complex Float], [Complex Float])
bflyS as = (los,rts)
  where
    (ls,rs) = halve as
    los = zipWith (+) ls rs
    ros = zipWith (-) ls rs
    rts = zipWith (*) ros [tw (length as) i | i <- [0..(length ros) - 1]]

-- split the input into two halves
halve as = splitAt n' as
  where
    n' = div (length as + 1) 2


module Lib where

import Data.List ()
import Numeric.LinearAlgebra
  ( Linear (scale),
    Matrix,
    Vector,
    cmap,
    maxIndex,
    outer,
    (#>),
  )
import Text.CSV ()

{-
Sooooooooooooo
What do I need to do
  Need to make the output a vector
  make the weights a matrix
  make the learning stage be an outer product thats then
  subracted from the weights

  I need to implement sigmoid?
-}

predict ::
  Matrix Double -> -- curried - Weights
  Vector Double -> -- curried - Input vector
  Vector Double -- Guess
predict w x = cmap sigmoid w #> x

predictEpoch ::
  Matrix Double -> -- w - Weights
  [Vector Double] -> -- x - An epoch of input vectorors
  [Vector Double] -- A guess for each input in the epoch
predictEpoch w x = map (cmap sigmoid) (fmap (w #>) x)

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

stocTrain ::
  Matrix Double -> -- w - Weights
  Vector Double -> -- x - Single input vectoror
  Vector Double -> -- l - label
  Matrix Double -- Trained Weights
stocTrain w x l = w - delta
  where
    error = predict w x - l
    sigDeriv = error * cmap (1 -) error
    delta = scale 0.001 $ sigDeriv `outer` x

trainEpoch ::
  Matrix Double -> -- w - Weights
  [Vector Double] -> -- x - An epoch of inputs
  [Vector Double] -> -- l - An epoch of labels
  Matrix Double -- Weights
trainEpoch w (x : xs) (l : ls) = trainEpoch (stocTrain w x l) xs ls
trainEpoch w [] [] = w

train ::
  Int -> -- t - number of epochs to train on
  Matrix Double -> -- w - Weights
  [Vector Double] -> -- x - An epoch of inputs
  [Vector Double] -> -- l - An epoch of labels
  Matrix Double -- Weights
train t w xs ls
  | t == 0 = w
  | otherwise = train (t - 1) (trainEpoch w xs ls) xs ls

successRate ::
  Matrix Double -> -- w - Weights
  [Vector Double] -> -- x - An epoch of input vectors
  [Vector Double] -> -- l - An epoch of labels
  Double -- A guess for each input in the epoch
successRate w x l = intDiv correct total
  where
    correct = sum (zipWith checkSame l (predictEpoch w x))
    total = length l

intDiv :: Int -> Int -> Double
intDiv n l = fromIntegral n / fromIntegral l

-- validate w x l = fmap maxIndex (predictEpoch w x)

checkSame :: Vector Double -> Vector Double -> Int
checkSame x y
  | maxIndex x == maxIndex y = 1
  | otherwise = 0

{-
so we need to:
stick the ending bias 1 onto each input array
turn the weights into a matrix
turn the input vectors in a list of vectors
turn the labels into a list of one hot encoded vectors
-}

-- csvToTrainingPred :: [[String]] -> [Double]
-- csvToTrainingPred s = predictEpoch trainedWeights scaledInputs
--   where
--     (labels, unscaledInputs) = labelInputSplit (map (map read) (init s))
--     scaledInputs = scaleInput unscaledInputs
--     trainedWeights = train 1 scaledInputs labels [0 ..]

-- scaleInput :: [[Double]] -> [[Double]]
-- scaleInput = map (map (/ 255))

-- labelInputSplit :: [[Double]] -> ([Double], [[Double]])
-- labelInputSplit xs = (head t, transpose (tail t))
--   where
--     t = transpose xs


module LogisticRegression where
import Util
import GradientDescent
import Data.List
import Numeric.LinearAlgebra
import Numeric.LinearAlgebra.Util
import Numeric.GSL.Minimization

{- Trains multiple logistic regression classifiers and returns all
  the classifiers in a list of matrices [theta], where the i-th element
  corresponds to the classifier for label i -}
classifyOneVsAll :: Matrix Double -> Matrix Double -> Double -> Double -> Int -> ([Matrix Double],[Double])
classifyOneVsAll x y lambda precision max_iters =  (map (\ target -> trainClassifier x target lambda precision max_iters) targetMatrixList,targets)
  where
    (targetMatrix,targets) = mapOutputsVectorToTargetMatrix y
    targetMatrixList = map asColumn $ toColumns targetMatrix

{- Trains a logistic regression classifier, returning the learned theta -}
trainClassifier :: Matrix Double -> Matrix Double -> Double -> Double -> Int -> Matrix Double
trainClassifier x y lambda precision max_iters = asColumn $ fst $ minimizeVDGeneric VectorBFGS2 x y (zeros (cols x) 1) max_iters lambda precision

{- The cost function J
   Compute cost and gradient for logistic regression with regularization.
   Computes the cost of using theta as the parameter for regularized logistic regression and the
   gradient of the cost w.r.t. to the parameters.
   Returns double of cost and gradient -}
computeCost :: Matrix Double -> Matrix Double -> Matrix Double -> Double -> (Double, Matrix Double)
computeCost x y theta lambda = (cost, gradient)
    where
      m = fromIntegral $ rows y
      prediction = sigmoid $ x <> theta
      tempTheta = accum theta (*) [((0,0),0)]
      cost = - (1/m) * sumElements ( y * (log prediction) + (1 - y) * (log (1 - prediction)) ) + (lambda/(2*m)) * sumElements (tempTheta ^ 2)
      gradient = ctrans ((scalar (1/m)) * ( sumColumns  ((prediction - y) * x) ) ) + (scalar (lambda/m)) * tempTheta

{- Variants that only compute cost OR gradient (for min functions) -}
costOnly :: Matrix Double -> Matrix Double -> Matrix Double -> Double -> Double
costOnly x y theta lambda = - (1/m) * sumElements ( y * (log prediction) + (1 - y) * (log (1 - prediction)) ) + (lambda/(2*m)) * sumElements (tempTheta ^ 2)
  where
      m = fromIntegral $ rows y
      prediction = sigmoid $ x <> theta
      tempTheta = accum theta (*) [((0,0),0)]

gradientOnly :: Matrix Double -> Matrix Double -> Matrix Double -> Double -> Matrix Double
gradientOnly x y theta lambda = ctrans ((scalar (1/m)) * ( sumColumns  ((prediction - y) * x) ) ) + (scalar (lambda/m)) * tempTheta
  where
      m = fromIntegral $ rows y
      prediction = sigmoid $ x <> theta
      tempTheta = accum theta (*) [((0,0),0)]

-- Minimize with gradient descent
minimizeGD x y initial_theta alpha num_iters lambda = gradientDescent gradFunc initial_theta alpha num_iters
    where
      gradFunc t = snd (computeCost x y t lambda)

{-
Minimization methods:
ConjugateFR
ConjugatePR
VectorBFGS
VectorBFGS2
SteepestDescent
-}

minimizeVDGeneric method x y initial_theta num_iters lambda precision = minimizeVD method precision num_iters first_step tol costFunction gradientFunction (flatten initial_theta)
    where
      first_step = precision * 100
      tol = 0.1 -- see http://www.gnu.org/software/gsl/manual/html_node/Multimin-Algorithms-with-Derivatives.html
      costFunction thetaVect = costOnly x y (asColumn thetaVect) lambda
      gradientFunction thetaVect = flatten $ gradientOnly x y (asColumn thetaVect) lambda

-- Predict probability that the label is 0 or 1 using learned logistic
predict xi theta = sigmoid (xi <> theta)

-- Predict whether the label is 0 or 1 using learned logistic
predictThreshold :: Matrix Double -> Matrix Double -> Double -> Matrix Double
predictThreshold xi theta threshold = cond (predict xi theta) (scalar threshold) 0 1 1

-- Gauge the accuracy of a prediction
accuracy p y = mean $ toList $ flatten (cond p y 0 1 0) * 100

{- Predict the label for a trained one-vs-all classifier. The labels
   are in the range 1..K, where K = size(all_theta, 1).
   p = PREDICTONEVSALL(all_theta, X) will return a vector of predictions
   for each example in the matrix X. Note that X contains the examples in
   rows. all_theta is a matrix where the i-th row is a trained logistic
   regression theta vector for the i-th class. You should set p to a vector
   of values from 1..K (e.g., p = [1; 3; 1; 2] predicts classes 1, 3, 1, 2
   for 4 examples) -}
predictOneVsAll :: [Matrix Double] -> Matrix Double -> [Double] -> Matrix Double
predictOneVsAll all_theta x targetMapping = selectTarget
  where
    at_matrix = fromBlocks [all_theta]
    maxIndex xs = snd $ maximum $ zip xs targetMapping
    selectTarget = asColumn $ fromList $ map maxIndex $ toLists $ predict x at_matrix


-- | Module for providing various functionals used in calculations involving
-- quantum mathematics. This includes the fidelity and trace norm.
module Hoqus.Fidelity where

import Numeric.LinearAlgebra.Data
import Numeric.LinearAlgebra

import Hoqus.MtxFun

-- | Function 'fidelity' calculates the fidelity between two quantum states (or
-- any two square matrices)
--fidelity :: Matrix C -> Matrix C -> C
--fidelity a b = (sum $ map sqrt $ toList $ eigenvalues $ a <> b ) ** 2

-- | Function 'superFidelity' calculates an upper bound for the fidelity using
-- only trace of the products of input matrices.
superFidelity :: Matrix C -> Matrix C -> C
superFidelity a b = (trace (a <> b)) + (sqrt (1 - trace (a <> a))) * (sqrt (1 - trace (b <> b)))

-- | Function 'subFidelity' calculates a lower bound for the fidelity using only
-- race of the products of input matrices.
subFidelity :: Matrix C -> Matrix C -> C
subFidelity a b = (trace p) +  (sqrt 2) * (sqrt ((trace p)*(trace p) - trace (p <> p)))
  where p = a <> b


import Data.Complex (Complex((:+)))
import qualified Data.Complex as C (realPart)
import qualified Data.Map as M (fromList, Map)

main = print $ mandelbrot (-2, 2) (-2, 2) 20 20 40

quadratic :: Num a => a -> a -> a
quadratic c z = z^2 + c

isBounded f lim z k
    | lim < k = True
    | C.realPart z <= 2 = isBounded f lim result (k+1)
    | otherwise = False
    where result = f z

isMandelbrot c lim = isBounded (quadratic c) lim 0 0

mandelbrot (x1, x2) (y1, y2) n1 n2 lim = M.fromList l
    where points = [ (x, y) | x <- linspace (x1, x2) n1,
                              y <- linspace (y1, y2) n2 ]
          pointMandel point = isMandelbrot (uncurry (:+) point) lim
          l = [ point | point <- points , pointMandel point ]

linspace (x1, x2) n = x1 : linspaceNext difference [x2] (n - 2)
    where difference = abs (x1 - x2) / (n - 1)

linspaceNext d li n
    | n == 0 = li
    | otherwise = linspaceNext d (ele : li) (n - 1)
    where ele = head li - d


-----------------------------------------------------------------------------
-- |
-- Module      :  Numeric.Particle
-- Copyright   :  (c) 2016 FP Complete Corporation
-- License     :  MIT (see LICENSE)
-- Maintainer  :  dominic@steinitz.org
--
-- =The Theory
--
-- The particle filter, `runPF`, in this library is applicable to the
-- state space model given by
--
-- \[
-- \begin{aligned}
-- \boldsymbol{x}_i &= \boldsymbol{a}_i(\boldsymbol{x}_{i-1}) + \boldsymbol{\psi}_{i-1} \\
-- \boldsymbol{y}_i &= \boldsymbol{h}_i(\boldsymbol{x}_i) + \boldsymbol{\upsilon}_i
-- \end{aligned}
-- \]
--
-- where
--
-- * \(\boldsymbol{a_i}\)\ is some non-linear vector-valued possibly
-- time-varying state update function.
--
-- * \(\boldsymbol{\psi}_{i}\) are independent normally
-- distributed random variables with mean 0 representing the fact that
-- the state update is noisy: \(\boldsymbol{\psi}_{i} \sim {\cal{N}}(0,Q_i)\).
--
-- * \(\boldsymbol{h}_i\)\ is some non-linear vector-valued possibly time-varying
-- function describing how we observe the hidden state process.
--
-- * \(\boldsymbol{\upsilon}_i\) are independent normally
-- distributed random variables with mean 0 represent the fact that
-- the observations are noisy: \(\boldsymbol{\upsilon}_{i} \sim {\cal{N}}(0,R_i)\).
--
-- Clearly this could be generalised further; anyone wishing for such
-- a generalisation is encouraged to contribute to the library.
--
-- The smoother, `oneSmoothingPath` implements [Forward filtering /
-- backward
-- smoothing](https://en.wikipedia.org/wiki/Particle_filter#Backward_particle_smoothers)
-- and returns just one path from the particle filter; in most cases,
-- this will need to be run many times to provide good estimates of
-- the past. Note that `oneSmoothingPath` uses /all/ observation up
-- until the current time. This could be generalised to select only a
-- window of observations up to the current time. Again, contributions
-- to implement this generalisation are welcomed.
--
-- = An Extended Example
--
-- The equation of motion for a pendulum of unit length subject to
-- [Gaussian white
-- noise](https://en.wikipedia.org/wiki/White_noise#Mathematical_definitions)
-- is
--
-- \[
-- \frac{\mathrm{d}^2\alpha}{\mathrm{d}t^2} = -g\sin\alpha + w(t)
-- \]
--
-- We can discretize this via the usual [Euler method](https://en.wikipedia.org/wiki/Euler_method)
--
-- \[
-- \begin{bmatrix}
-- x_{1,i} \\
-- x_{2,i}
-- \end{bmatrix}
-- =
-- \begin{bmatrix}
-- x_{1,i-1} + x_{2,i-1}\Delta t \\
-- x_{2,i-1} - g\sin x_{1,i-1}\Delta t
-- \end{bmatrix}
-- +
-- \mathbf{q}_{i-1}
-- \]
--
-- where \(q_i \sim {\mathcal{N}}(0,Q)\) and
--
-- \[
-- Q
-- =
-- \begin{bmatrix}
-- \frac{q^c \Delta t^3}{3} & \frac{q^c \Delta t^2}{2} \\
-- \frac{q^c \Delta t^2}{2} & {q^c \Delta t}
-- \end{bmatrix}
-- \]
--
-- Assume that we can only measure the horizontal position of the
-- pendulum and further that this measurement is subject to error so that
--
-- \[
-- y_i = \sin x_i + r_i
-- \]
--
-- where \(r_i \sim {\mathcal{N}}(0,R)\).
--
--
-- First let's set the time step and the acceleration caused by
-- earth's gravity and the covariance matrices for the state and
-- observation.
--
--
-- > {-# LANGUAGE DataKinds #-}
-- >
-- > import Numeric.Particle
-- >
-- > deltaT, g :: Double
-- > deltaT = 0.01
-- > g      = 9.81
-- >
-- > bigQ :: Sym 2
-- > bigQ = sym $ matrix bigQl
-- >
-- > qc :: Double
-- > qc = 0.01
-- >
-- > bigQl :: [Double]
-- > bigQl = [ qc * deltaT^3 / 3, qc * deltaT^2 / 2,
-- >           qc * deltaT^2 / 2,       qc * deltaT
-- >         ]
-- >
-- > bigR :: Sym 1
-- > bigR  = sym $ matrix [0.2]
-- >
-- > data SystemState a = SystemState { x1  :: a, x2  :: a }
-- >   deriving Show
-- >
-- > newtype SystemObs a = SystemObs { y1  :: a }
-- >   deriving Show
-- >
--
-- We use the data generated using code made available with Simo
-- Särkkä's
-- [book](http://www.cambridge.org/gb/academic/subjects/statistics-probability/applied-probability-and-stochastic-networks/bayesian-filtering-and-smoothing)
-- which starts with an angle of 1.5 radians and an angular velocity
-- of 0.0 radians per second. We set our prior to have a mean of 1.6, not too far from the actual value.
--
-- > {-# LANGUAGE DataKinds #-}
-- >
-- > m0 :: PendulumState
-- > m0 = vector [1.6, 0]
-- >
-- > bigP :: Sym 2
-- > bigP = sym $ diag 0.1
-- >
-- > initParticles :: R.MonadRandom m =>
-- >                  m (Particles (SystemState Double))
-- > initParticles = V.replicateM nParticles $ do
-- >   r <- R.sample $ R.rvar (Normal m0 bigP)
-- >   let x1 = fst $ headTail r
-- >       x2 = fst $ headTail $ snd $ headTail r
-- >   return $ SystemState { x1 = x1, x2 = x2}
-- >
-- > nObs :: Int
-- > nObs = 200
-- >
-- > nParticles :: Int
-- > nParticles = 20
-- >
-- > (.+), (.*), (.-) :: (Num a) => V.Vector a -> V.Vector a -> V.Vector a
-- > (.+) = V.zipWith (+)
-- > (.*) = V.zipWith (*)
-- > (.-) = V.zipWith (-)
-- >
-- > stateUpdateP :: Particles (SystemState Double) ->
-- >                 Particles (SystemState Double)
-- > stateUpdateP xPrevs = V.zipWith SystemState x1s x2s
-- >   where
-- >     ix = V.length xPrevs
-- >
-- >     x1Prevs = V.map x1 xPrevs
-- >     x2Prevs = V.map x2 xPrevs
-- >
-- >     deltaTs = V.replicate ix deltaT
-- >     gs = V.replicate ix g
-- >     x1s = x1Prevs .+ (x2Prevs .* deltaTs)
-- >     x2s = x2Prevs .- (gs .* (V.map sin x1Prevs) .* deltaTs)
--
--
-- <<diagrams/src_Numeric_Particle_diagV.svg#diagram=diagV&height=600&width=500>>
--
-- === Code for Plotting
--
-- The full code for plotting the results:
--
-- > {-# LANGUAGE ExplicitForAll #-}
-- > {-# LANGUAGE TypeOperators  #-}
-- > {-# LANGUAGE DataKinds      #-}
-- >
-- > import qualified Graphics.Rendering.Chart as C
-- > import Graphics.Rendering.Chart.Backend.Diagrams
-- > import Data.Colour
-- > import Data.Colour.Names
-- > import Data.Default.Class
-- > import Control.Lens
-- >
-- > import Data.Csv
-- > import System.IO hiding ( hGetContents )
-- > import Data.ByteString.Lazy ( hGetContents )
-- > import qualified Data.Vector as V
-- >
-- > import           Data.Random.Distribution.Static.MultivariateNormal ( Normal(..) )
-- > import qualified Data.Random as R
-- > import           Data.Random.Source.PureMT ( pureMT )
-- > import           Control.Monad.State ( evalState, replicateM )
-- >
-- > import Data.List ( transpose )
-- >
-- > import GHC.TypeLits ( KnownNat )
-- >
-- > import Numeric.LinearAlgebra.Static
-- >
-- > chartEstimated :: String ->
-- >               [(Double, Double)] ->
-- >               [(Double, Double)] ->
-- >               [(Double, Double)] ->
-- >               C.Renderable ()
-- > chartEstimated title acts obs ests = C.toRenderable layout
-- >   where
-- >
-- >     actuals = C.plot_lines_values .~ [acts]
-- >             $ C.plot_lines_style  . C.line_color .~ opaque red
-- >             $ C.plot_lines_title .~ "Actual Trajectory"
-- >             $ C.plot_lines_style  . C.line_width .~ 1.0
-- >             $ def
-- >
-- >     measurements = C.plot_points_values .~ obs
-- >                  $ C.plot_points_style  . C.point_color .~ opaque blue
-- >                  $ C.plot_points_title .~ "Measurements"
-- >                  $ def
-- >
-- >     estimas = C.plot_lines_values .~ [ests]
-- >             $ C.plot_lines_style  . C.line_color .~ opaque black
-- >             $ C.plot_lines_title .~ "Inferred Trajectory"
-- >             $ C.plot_lines_style  . C.line_width .~ 1.0
-- >             $ def
-- >
-- >     layout = C.layout_title .~ title
-- >            $ C.layout_plots .~ [C.toPlot actuals, C.toPlot measurements, C.toPlot estimas]
-- >            $ C.layout_y_axis . C.laxis_title .~ "Angle / Horizontal Displacement"
-- >            $ C.layout_y_axis . C.laxis_override .~ C.axisGridHide
-- >            $ C.layout_x_axis . C.laxis_title .~ "Time"
-- >            $ C.layout_x_axis . C.laxis_override .~ C.axisGridHide
-- >            $ def
-- >
-- > stateUpdateNoisy :: R.MonadRandom m =>
-- >                     Sym 2 ->
-- >                     Particles (SystemState Double) ->
-- >                     m (Particles (SystemState Double))
-- > stateUpdateNoisy bigQ xPrevs = do
-- >   let xs = stateUpdateP xPrevs
-- >
-- >       x1s = V.map x1 xs
-- >       x2s = V.map x2 xs
-- >
-- >   let ix = V.length xPrevs
-- >   etas <- replicateM ix $ R.sample $ R.rvar (Normal 0.0 bigQ)
-- >
-- >   let eta1s, eta2s :: V.Vector Double
-- >       eta1s = V.fromList $ map (fst . headTail) etas
-- >       eta2s = V.fromList $ map (fst . headTail . snd . headTail) etas
-- >
-- >   return (V.zipWith SystemState (x1s .+ eta1s) (x2s .+ eta2s))
-- >
-- > obsUpdate :: Particles (SystemState Double) ->
-- >              Particles (SystemObs Double)
-- > obsUpdate xs = V.map (SystemObs . sin . x1) xs
-- >
-- > weight :: forall a n . KnownNat n =>
-- >           (a -> R n) ->
-- >           Sym n ->
-- >           a -> a -> Double
-- > weight f bigR obs obsNew = R.pdf (Normal (f obsNew) bigR) (f obs)
-- >
-- > runFilter :: Particles (SystemObs Double) -> V.Vector (Particles (SystemState Double))
-- > runFilter pendulumSamples = evalState action (pureMT 19)
-- >   where
-- >     action = do
-- >       xs <- initParticles
-- >       scanMapM
-- >         (runPF (stateUpdateNoisy bigQ) obsUpdate (weight f bigR))
-- >         return
-- >         xs
-- >         pendulumSamples
-- >
-- > testSmoothing :: Particles (SystemObs Double) -> Int -> [Double]
-- > testSmoothing ss n = V.toList $ evalState action (pureMT 23)
-- >   where
-- >     action = do
-- >       xss <- V.replicateM n $ oneSmoothingPath (stateUpdateNoisy bigQ) (weight h bigQ) nParticles (runFilter ss)
-- >       let yss = V.fromList $ map V.fromList $
-- >                 transpose $
-- >                 V.toList $ V.map (V.toList) $
-- >                 xss
-- >       return $ V.map (/ (fromIntegral n)) $ V.map V.sum $ V.map (V.map x1) yss
-- >
-- > type PendulumState = R 2
-- >
-- > f :: SystemObs Double -> R 1
-- > f = vector . pure . y1
-- >
-- > h :: SystemState Double -> R 2
-- > h u = vector [x1 u , x2 u]
-- >
-- > diagV = do
-- >   h <- openFile "data/matlabRNGs.csv" ReadMode
-- >   cs <- hGetContents h
-- >   let df = (decode NoHeader cs) :: Either String (V.Vector (Double, Double))
-- >   case df of
-- >     Left _ -> error "Whatever"
-- >     Right generatedSamples -> do
-- >       let preObs = V.take nObs $ V.map fst generatedSamples
-- >       let obs = V.toList preObs
-- >       let acts = V.toList $ V.take nObs $ V.map snd generatedSamples
-- >       let nus = take nObs (testSmoothing (V.map SystemObs preObs) 50)
-- >       denv <- defaultEnv C.vectorAlignmentFns 600 500
-- >       let charte = chartEstimated "Particle Smoother"
-- >                                   (zip [0,1..] acts)
-- >                                   (zip [0,1..] obs)
-- >                                   (zip [0,1..] nus)
-- >       return $ fst $ runBackend denv (C.render charte (600, 500))
--
-----------------------------------------------------------------------------

{-# OPTIONS_GHC -Wall                     #-}

{-# LANGUAGE BangPatterns                 #-}

module Numeric.Particle (
    runPF
  , oneSmoothingPath
  , Particles
  , Path
  ) where

import           Data.Random hiding ( StdNormal, Normal )
import           Control.Monad
import qualified Data.Vector as V
import           Data.Vector ( Vector )
import           Data.Bits ( shiftR )

type Particles a = Vector a -- ^ As an aid for the reader
type Path a      = Vector a -- ^ As an aid for the reader

runPF
  :: MonadRandom m
  => (Particles a -> m (Particles a)) -- ^ System evolution at a point
                                      -- \(\boldsymbol{a}_i(\boldsymbol{x})\)
  -> (Particles a -> Particles b)     -- ^ Measurement operator at a point \(\boldsymbol{h}_i\)
  -> (b -> b -> Double)               -- ^ Observation probability density function \(p(\boldsymbol{y}_k|\boldsymbol{x}^{(i)}_k)\)
  -> Particles a                      -- ^ Current estimate
                                      -- \(\big\{\big(w^{(i)}_k, \boldsymbol{x}^{(i)}_k\big) : i \in \{1,\ldots,N\}\big\}\)
                                      -- where \(N\) is the number of particles
  -> b                                -- ^ New measurement \(\boldsymbol{y}_i\)
  -> m (Particles a)                  -- ^ New estimate
                                      -- \(\big\{\big(w^{(i)}_{k+1}, \boldsymbol{x}^{(i)}_{k+1}\big) : i \in \{1,\ldots,N\}\big\}\)
                                      -- where \(N\) is the number of particles

runPF stateUpdate obsUpdate weight statePrevs obs = do
  stateNews <- stateUpdate statePrevs
  let obsNews = obsUpdate stateNews
  let weights       = V.map (weight obs) obsNews
      cumSumWeights = V.tail $ V.scanl (+) 0 weights
      totWeight     = V.last cumSumWeights
      nParticles    = V.length statePrevs
  vs <- V.replicateM nParticles (sample $ uniform 0.0 totWeight)
  let js = indices cumSumWeights vs
      stateTildes = V.map (stateNews V.!) js
  return stateTildes

oneSmoothingPath
  :: MonadRandom m
  => (Particles a -> m (Particles a)) -- ^ System evolution at a point
                                      -- \(\boldsymbol{a}_i(\boldsymbol{x})\)
  -> (a -> a -> Double)               -- ^ State probability density function \(p(\boldsymbol{x}^{(i)}_k|\boldsymbol{x}^{(i)}_{k-1})\)
  -> Int                              -- ^ Number of particles \(N\)
  -> (Vector (Particles a))           -- ^ Filtering estimates
                                      -- \(\bigg\{\big\{\big(w^{(i)}_{k+1}, \boldsymbol{x}^{(i)}_{k+1}\big) : i \in \{1,\ldots,N\}\big\} : k \in \{1,\ldots,T\}\bigg\}\)
                                      -- where \(T\) is the number of timesteps
  -> m (Path a)
                                      -- ^ A path for the smoothed particle \(\{\boldsymbol{x}_k : k \in T\}\) where \(T\) is the number of timesteps
oneSmoothingPath stateUpdate weight nParticles filterEstss = do
  let ys = filterEstss
  ix <- sample $ uniform 0 (nParticles - 1)
  let xn = (V.head ys) V.! ix
  scanMapM (oneSmoothingStep stateUpdate weight) return xn (V.tail ys)

oneSmoothingStep :: MonadRandom m =>
         (Particles a -> m (Particles a)) ->
         (a -> a -> Double) ->
         a ->
         Particles a ->
         m a
oneSmoothingStep stateUpdate
                 stateDensity
                 smoothingSampleAtiPlus1
                 filterSamplesAti = do it
  where
    it = do
      mus <- stateUpdate filterSamplesAti
      let weights =  V.map (stateDensity smoothingSampleAtiPlus1) mus
          cumSumWeights = V.tail $ V.scanl (+) 0 weights
          totWeight     = V.last cumSumWeights
      v <- sample $ uniform 0.0 totWeight
      let ix = binarySearch cumSumWeights v
          xnNew = filterSamplesAti V.! ix
      return $ xnNew

indices :: Vector Double -> Vector Double -> Vector Int
indices bs xs = V.map (binarySearch bs) xs

binarySearch :: (Ord a) =>
                Vector a -> a -> Int
binarySearch vec x = loop 0 (V.length vec - 1)
  where
    loop !l !u
      | u <= l    = l
      | otherwise = let e = vec V.! k in if x <= e then loop l k else loop (k+1) u
      where k = l + (u - l) `shiftR` 1

scanMapM :: Monad m => (s -> a -> m s) -> (s -> m b) -> s -> Vector a -> m (Vector b)
scanMapM f g !s0 !xs
  | V.null xs = do
    r <- g s0
    return $ V.singleton r
  | otherwise = do
    s <- f s0 (V.head xs)
    r <- g s0
    liftM (r `V.cons`) (scanMapM f g s (V.tail xs))


module Main where

import Functions

import Math.FFT
import Control.Arrow
import Data.Array.CArray
import Data.Complex
import Foreign.Storable.Complex()

import Control.Lens
import Data.Colour
import Data.Colour.Names
import Data.Default.Class
import Graphics.Rendering.Chart
import Graphics.Rendering.Chart.Backend.Cairo

import Codec.Picture

import System.Random

import Control.Exception

----------

numOfData :: Int -- Frequency
numOfData = 256

rawData :: [Double]
rawData = take numOfData $ randomRs (-1, 1) $ mkStdGen 7

inputData :: [(Int, Double)]
inputData = zip [1..] rawData

----------

lowCutoffFrequency :: Double
lowCutoffFrequency = 40.0

highCutoffFrequency :: Double
highCutoffFrequency = 80.0

transBand :: Double
transBand = 10.0

filterDomain :: [Double]
filterDomain = [-(fromIntegral nof-1)/2 .. (fromIntegral nof-1)/2]
  where nof' = round $ 3.3 * fromIntegral numOfData / transBand
        nof = if odd nof' then nof' else nof'+1 :: Int

hamming :: Int -> Int -> Double -> Double
hamming n i v = v * (0.53836 - 0.46164 * cos((2.0 * pi * fromIntegral i) / (fromIntegral n - 1)));

lowPass :: [Double]
lowPass = zipWith (hamming (length filterDomain)) [1..] [ k * sinc (pi*k*n) | n <- filterDomain ]
  where k = 2 * lowCutoffFrequency / fromIntegral numOfData

lowPassedData :: [(Int, Double)]
lowPassedData = zip [1..] lpd
  where lpd = take numOfData $ drop (div (length filterDomain) 2) $ convolve lowPass rawData

highPass :: [Double]
highPass = zipWith (hamming (length filterDomain)) [1..] [ sinc (pi*n) - k * sinc (pi*k*n) | n <- filterDomain ]
  where k = 2 * highCutoffFrequency / fromIntegral numOfData

highPassedData :: [(Int, Double)]
highPassedData = zip [1..] hpd
  where hpd = take numOfData $ drop (div (length filterDomain) 2) $ convolve highPass rawData

bandPass :: [Double]
bandPass = zipWith (hamming (length filterDomain)) [1..] [ kh * sinc (pi*kh*n) - kl * sinc (pi*kl*n) | n <- filterDomain ]
  where kl = 2 * lowCutoffFrequency / fromIntegral numOfData
        kh = 2 * highCutoffFrequency / fromIntegral numOfData

bandPassedData :: [(Int, Double)]
bandPassedData = zip [1..] bpd
  where bpd = take numOfData $ drop (div (length filterDomain) 2) $ convolve bandPass rawData

bandStop :: [Double]
bandStop = zipWith (hamming (length filterDomain)) [1..] [ sinc (pi*n) - kh * sinc (pi*kh*n) + kl * sinc (pi*kl*n) | n <- filterDomain ]
  where kl = 2 * lowCutoffFrequency / fromIntegral numOfData
        kh = 2 * highCutoffFrequency / fromIntegral numOfData

bandStoppedData :: [(Int, Double)]
bandStoppedData = zip [1..] bsd
  where bsd = take numOfData $ drop (div (length filterDomain) 2) $ convolve bandStop rawData

----------

toDFT :: [(Int, Double)] -> [(Int, Complex Double)]
toDFT ds = take (div numOfData 2) $ assocs $ dft $ array (1, length ds) $ map (\(i, v) -> (i, v:+0)) ds

toAmplitude :: [(Int, Complex Double)] -> [(Int, Double)]
toAmplitude = map (second magnitude)

----------

timeChart :: Renderable ()
timeChart = toRenderable timeLayout
  where timeLayout = layout_title .~ "Time Scale Data"
                   $ layout_plots .~ [ toPlot inputLine ]
                   $ def

        inputLine = plot_lines_title .~ "Input Data"
                  $ plot_lines_style . line_color .~ opaque skyblue
                  $ plot_lines_values .~ [inputData]
                  $ def

timeChartLowPass :: Renderable ()
timeChartLowPass = toRenderable timeLayoutLowPass
  where timeLayoutLowPass = layout_title .~ "Time Scale Data(Low Pass)"
                          $ layout_plots .~ [ toPlot inputLine
                                            , toPlot lowPassLine ]
                          $ def

        inputLine = plot_lines_title .~ "Input Data"
                  $ plot_lines_style . line_color .~ opaque skyblue
                  $ plot_lines_values .~ [inputData]
                  $ def

        lowPassLine = plot_lines_title .~ "Low Passed Data"
                    $ plot_lines_style . line_color .~ opaque green
                    $ plot_lines_values .~ [lowPassedData]
                    $ def

timeChartHighPass :: Renderable ()
timeChartHighPass = toRenderable timeLayoutHighPass
  where timeLayoutHighPass = layout_title .~ "Time Scale Data(High Pass)"
                           $ layout_plots .~ [ toPlot inputLine
                                             , toPlot highPassLine ]
                           $ def

        inputLine = plot_lines_title .~ "Input Data"
                  $ plot_lines_style . line_color .~ opaque skyblue
                  $ plot_lines_values .~ [inputData]
                  $ def

        highPassLine = plot_lines_title .~ "High Passed Data"
                     $ plot_lines_style . line_color .~ opaque orange
                     $ plot_lines_values .~ [highPassedData]
                     $ def

timeChartBandPass :: Renderable ()
timeChartBandPass = toRenderable timeLayoutBandPass
  where timeLayoutBandPass = layout_title .~ "Time Scale Data(Band Pass)"
                           $ layout_plots .~ [ toPlot inputLine
                                             , toPlot bandPassLine ]
                           $ def

        inputLine = plot_lines_title .~ "Input Data"
                  $ plot_lines_style . line_color .~ opaque skyblue
                  $ plot_lines_values .~ [inputData]
                  $ def

        bandPassLine = plot_lines_title .~ "Band Passed Data"
                     $ plot_lines_style . line_color .~ opaque purple
                     $ plot_lines_values .~ [bandPassedData]
                     $ def

timeChartBandStop :: Renderable ()
timeChartBandStop = toRenderable timeLayoutBandStop
  where timeLayoutBandStop = layout_title .~ "Time Scale Data(Band Stop)"
                           $ layout_plots .~ [ toPlot inputLine
                                             , toPlot bandStopLine ]
                           $ def

        inputLine = plot_lines_title .~ "Input Data"
                  $ plot_lines_style . line_color .~ opaque skyblue
                  $ plot_lines_values .~ [inputData]
                  $ def

        bandStopLine = plot_lines_title .~ "Band Stopped Data"
                     $ plot_lines_style . line_color .~ opaque royalblue
                     $ plot_lines_values .~ [bandStoppedData]
                     $ def

frequencyChart :: Renderable ()
frequencyChart = toRenderable frequencyLayout
  where frequencyLayout = layout_title .~ "Frequency Scale Data"
                        $ layout_y_axis .~ yAxis'
                        $ layout_plots .~ [ toPlot inputLine' ]
                        $ def

        yAxis' = laxis_title .~ "Amplitude"
               $ def

        inputLine' = plot_lines_title .~ "Input Data"
                   $ plot_lines_style . line_color .~ opaque skyblue
                   $ plot_lines_values .~ [toAmplitude $ toDFT inputData]
                   $ def

frequencyChartLowPass :: Renderable ()
frequencyChartLowPass = toRenderable frequencyLayoutLowPass
  where frequencyLayoutLowPass = layout_title .~ "Frequency Scale Data(Low Pass)"
                               $ layout_y_axis .~ yAxis'
                               $ layout_plots .~ [ toPlot inputLine'
                                                 , toPlot lowPassLine' ]
                               $ def

        yAxis' = laxis_title .~ "Amplitude"
               $ def

        inputLine' = plot_lines_title .~ "Input Data"
                   $ plot_lines_style . line_color .~ opaque skyblue
                   $ plot_lines_values .~ [toAmplitude $ toDFT inputData]
                   $ def

        lowPassLine' = plot_lines_title .~ "Low Passed Data"
                     $ plot_lines_style . line_color .~ opaque green
                     $ plot_lines_values .~ [toAmplitude $ toDFT lowPassedData]
                     $ def

frequencyChartHighPass :: Renderable ()
frequencyChartHighPass = toRenderable frequencyLayoutHighPass
  where frequencyLayoutHighPass = layout_title .~ "Frequency Scale Data(High Pass)"
                                $ layout_y_axis .~ yAxis'
                                $ layout_plots .~ [ toPlot inputLine'
                                                  , toPlot highPassLine' ]
                                $ def

        yAxis' = laxis_title .~ "Amplitude"
               $ def

        inputLine' = plot_lines_title .~ "Input Data"
                   $ plot_lines_style . line_color .~ opaque skyblue
                   $ plot_lines_values .~ [toAmplitude $ toDFT inputData]
                   $ def

        highPassLine' = plot_lines_title .~ "High Passed Data"
                      $ plot_lines_style . line_color .~ opaque orange
                      $ plot_lines_values .~ [toAmplitude $ toDFT highPassedData]
                      $ def

frequencyChartBandPass :: Renderable ()
frequencyChartBandPass = toRenderable frequencyLayoutBandPass
  where frequencyLayoutBandPass = layout_title .~ "Frequency Scale Data(Band Pass)"
                                $ layout_y_axis .~ yAxis'
                                $ layout_plots .~ [ toPlot inputLine'
                                                  , toPlot bandPassLine' ]
                                $ def

        yAxis' = laxis_title .~ "Amplitude"
               $ def

        inputLine' = plot_lines_title .~ "Input Data"
                   $ plot_lines_style . line_color .~ opaque skyblue
                   $ plot_lines_values .~ [toAmplitude $ toDFT inputData]
                   $ def

        bandPassLine' = plot_lines_title .~ "Band Passed Data"
                      $ plot_lines_style . line_color .~ opaque purple
                      $ plot_lines_values .~ [toAmplitude $ toDFT bandPassedData]
                      $ def

frequencyChartBandStop :: Renderable ()
frequencyChartBandStop = toRenderable frequencyLayoutBandStop
  where frequencyLayoutBandStop = layout_title .~ "Frequency Scale Data(Band Stop)"
                                $ layout_y_axis .~ yAxis'
                                $ layout_plots .~ [ toPlot inputLine'
                                                  , toPlot bandStopLine' ]
                                $ def

        yAxis' = laxis_title .~ "Amplitude"
               $ def

        inputLine' = plot_lines_title .~ "Input Data"
                   $ plot_lines_style . line_color .~ opaque skyblue
                   $ plot_lines_values .~ [toAmplitude $ toDFT inputData]
                   $ def

        bandStopLine' = plot_lines_title .~ "Band Stopped Data"
                      $ plot_lines_style . line_color .~ opaque royalblue
                      $ plot_lines_values .~ [toAmplitude $ toDFT bandStoppedData]
                      $ def

----------

decode :: Either String DynamicImage -> IO (Image PixelRGB8)
decode (Right (ImageRGB8 image)) = return image
decode _ = throwIO $ userError "Image Parse Error"

----------

main :: IO ()
main = do
  _ <- renderableToFile def "./png/time_scale_1.png" timeChart
  _ <- renderableToFile def "./png/time_scale_2.png" timeChartLowPass
  _ <- renderableToFile def "./png/time_scale_3.png" timeChartBandPass
  _ <- renderableToFile def "./png/time_scale_4.png" timeChartHighPass
  _ <- renderableToFile def "./png/time_scale_5.png" timeChartBandStop
  _ <- renderableToFile def "./png/freq_scale_1.png" frequencyChart
  _ <- renderableToFile def "./png/freq_scale_2.png" frequencyChartLowPass
  _ <- renderableToFile def "./png/freq_scale_3.png" frequencyChartBandPass
  _ <- renderableToFile def "./png/freq_scale_4.png" frequencyChartHighPass
  _ <- renderableToFile def "./png/freq_scale_5.png" frequencyChartBandStop

  timeScaleImages <- sequence [ readImage ("./png/time_scale_" ++ show i ++ ".png") >>= decode | i <- [1..5] ]
  freqScaleImages <- sequence [ readImage ("./png/freq_scale_" ++ show i ++ ".png") >>= decode | i <- [1..5] ]

  either putStrLn id $ writeGifAnimation "time_scale.gif" 100 LoopingForever timeScaleImages
  either putStrLn id $ writeGifAnimation "freq_scale.gif" 100 LoopingForever freqScaleImages


-- file: Astro/Elements.hs
-- Orbital element calculations.

module Astro.Elements
	where

import Astro.Misc
import Astro.Kepler
import Numeric.LinearAlgebra

-- Standard orbital elements from position and velocity, with given
-- gravitational parameter
posvelToOrbelGp :: Vector Double -> Vector Double -> Double -> [Double]
	--(Monad m) => Vector Double -> Vector Double -> Double -> m [Double]
posvelToOrbelGp rv vv' gpar =
	{-
	if (vv'<.>vv')/2 - gpar / (norm rv) >= 0 then
		fail "Not an elliptic orbit"
	else
	-}
	let 	vv 	= scalar (1/sqrt gpar) * vv'
		r 	= norm rv
		alpha 	= 2/r - vv<.>vv
		eta 	= rv <.> vv
		zeta	= 1 - alpha * r
		hv	= cross rv vv
		h	= norm hv
		ev	= cross vv hv - scale (1/r) rv
		b	= sqrt $ (hv@>0)^2 + (hv@>1)^2
		m0	= atan2 (eta * sqrt alpha) zeta - eta * sqrt alpha
		inc	= atan2 b (hv@>2)
		ascnode = atan2 (hv@>0) (-hv@>1)
		argperi = atan2 ((ev@>2)*h) (ev@>1*hv@>0-ev@>0*hv@>1)
	in	[1/alpha, norm ev, inc, ascnode, argperi, m0]
		--return $ [1/alpha, norm ev, inc, ascnode, argperi, m0]
posvelToOrbel rv vv = posvelToOrbelGp rv vv 1

-- Position and velocity from orbital elements, with given gravitational
-- parameter
orbelToPosvelGp :: [Double] -> Double -> (Vector Double, Vector Double)
	--(Monad m) => [Double] -> Double -> m (Vector Double, Vector Double)
orbelToPosvelGp oe gpar =
	{-
	if (length oe < 6) then
		fail "Need 6 orbital elements"
	else
	-}
	let 	[a, ecc, inc, o, w, m0] = oe
		pv	= fromList [ cos w *cos o - sin w *sin o *cos inc
		 	    	   , cos w *sin o + sin w *cos o *cos inc
		 	    	   , sin w *sin inc ]
		qv	= fromList [ -sin w *cos o  - cos w *sin o *cos inc
		 	   	   , -sin w *sin o  + cos w *cos o *cos inc
		 	    	   ,  cos w *sin inc ]
		wv	= fromList [  sin o *sin inc
				   , -cos o *sin inc
				   ,  cos inc  ]
		ea	= solveEKepler ecc m0
		av	= scalar a * pv
		bv	= scalar (sqrt (1-ecc^2)) * cross wv av
		rv	= scalar (cos ea - ecc) * av
					+ scalar (sin ea) * bv
		vv	= scalar ((sqrt gpar)/(norm rv * sqrt a)) *
				( scalar (sin ea) * (-av)
					+ scalar (cos ea) * bv )
	in
		(rv, vv)
		--return $ (rv, vv)
orbelToPosvel oe = orbelToPosvelGp oe 1


module TB.Statistics.Examples (
) where

import Data.Vector
import Statistics.Sample

v1 = fromList [1,1,1,4,8,10,12,9,7,5,3,1,1,1,2,1]
weights = fromList [0.1,0.1,0.1,0.4,0.8,1,1.2,0.9,0.7,0.5,0.3,0.1,0.1,0.1,0.1,0.2,0.1]

stats =
  [
    ("stdDev", stdDev v1)
  , ("skewness", skewness v1)
  , ("range", range v1)
  , ("mean", mean v1)
  , ("welfordMean", welfordMean v1)
  , ("meanWeighted", meanWeighted $ Data.Vector.zip v1 weights)
  , ("harmonicMean", harmonicMean v1)
  , ("geometricMean", geometricMean v1)
  , ("kurtosis", kurtosis v1)
  , ("variance", variance v1)
  , ("varianceUnbiased", varianceUnbiased v1)
  , ("meanVariance fst", fst mv)
  , ("meanVariance snd", snd mv)
  ]
  where
    mv = meanVariance v1


import Data.Complex (Complex, realPart)

type CD = Complex Double

quadraticRoots :: (CD, CD, CD) -> (CD, CD)
quadraticRoots (a, b, c) =
  if realPart b > 0
    then ((2 * c) / (-b - d), (-b - d) / (2 * a))
    else ((-b + d) / (2 * a), (2 * c) / (-b + d))
  where
    d = sqrt $ b ^ 2 - 4 * a * c

main :: IO ()
main =
  mapM_
    (print . quadraticRoots)
    [(3, 4, 4 / 3), (3, 2, -1), (3, 2, 1), (1, -10e5, 1), (1, -10e9, 1)]


-- gaussian_integers.hs
-- to find prime numbers in Gaussian integers

import Data.Complex

-- merge two sorted sequence
-- f is a function that helps to compare
merge :: Ord b => (a -> b) -> [a] -> [a] -> [a]
merge _ xs [] = xs
merge _ [] ys = ys
merge f u@(x:xs) v@(y:ys)
  | f x > f y  = y : merge f u ys
  | f x <= f y = x : merge f xs v
merge _ _ _ = error "merge: this should not happen"

mergeSort :: Ord b => (a -> b) -> [a] -> [a]
mergeSort f us = let n = length us
               in if n < 2
                  then us
                  else let (xs, ys) = splitAt (n `div` 2) us
                       in merge f (mergeSort f xs) (mergeSort f ys)

-- ssq : "sum of squares"
ssq :: Num a => (a, a) -> a
ssq (a, b) = a*a + b*b

-- r'seq : "the integer points on the plane, sorted by distance from origin"
r'seq :: [(Integer, Integer)]
r'seq = loop 0 0 []
  where
    loop _ _ [] = loop 1 0 [(0,1), (1,1)]
    loop k l s@((x,y):_)
          | x < l     = let (us, vs) = splitAt (fromIntegral (l+1)) s
                        in us ++ loop k (l+1) vs
          | otherwise = let k' = k + 1
                            l' = floor $ sqrt $ fromIntegral k'
                            t  = zip (replicate (fromIntegral (k+2)) k') [0..k']
                            s' = merge ssq s $ t
                        in loop k' l' s'

-- multiplication of Complex (Integers), which is not provided
infix 7 `times`
times :: Num a => Complex a -> Complex a -> Complex a
(a:+b) `times` (c:+d) = (a*c-b*d) :+ (a*d+b*c)

-- return True if a+bi divides c+di
infix 4 `divides`
divides :: Integral a => Complex a -> Complex a -> Bool
(a:+b) `divides` (c:+d) = let (m:+n) = (a:+(-b)) `times` (c:+d)
                              p = ssq (a, b)
                          in (m `mod` p) == 0 && (n `mod` p) == 0

-- the Gaussian primes
g'primes :: [Complex Integer]
g'primes = sieve $ tail r'seq
  where sieve ((x,y):rest)
          = let ns = case () of _
                                  | y == 0 -> [x:+0,0:+x,(-x):+0,0:+(-x)]
                                  | x == y -> [x:+x,x:+(-x),(-x):+x,(-x):+(-x)]
                                  | otherwise ->  [x:+y,y:+x,(-y):+x,(-x):+y]
                                                 ++ [(-x):+(-y),(-y):+(-x)]
                                                 ++ [y:+(-x),x:+(-y)]
            in ns ++ sieve [(x,y)|(x,y) <- rest, not $ any (`divides` x:+y) ns]

main :: IO ()
--main = mapM_ (putStrLn . show) g'primes
main = print $ map (\(a:+b) -> (a,b)) g'primes


module Lib
    ( Mandelbrot(..)
    , makeMandelbrot
    , drawMandelbrot
    ) where

import Graphics.Image hiding (realPart, imagPart, scale)
import Data.Complex

maxDepth :: Int
maxDepth = 200

-- The function to test on
f :: Complex Double -> Complex Double -> Complex Double
f z c = z ** 2 + c

-- Return infinite list of all numbers in sequence
nextValues :: Complex Double -> Complex Double -> [Complex Double]
nextValues z c = z:nextValues (f z c) c

series :: Complex Double -> [Complex Double]
series pos = nextValues (0 :+ 0) pos

-- See if number is within radius of convergence (using the pythagorean theorem)
mightConverge :: Complex Double -> Bool
mightConverge z = (realPart z) ** 2 + (imagPart z) ** 2 <= 2 ** 2

-- Count how quickly the infinite list becomes guaranteed to never converge
speedOfBlowup :: Complex Double -> Int
speedOfBlowup pos = length $ takeWhile mightConverge $ take maxDepth $ series pos

-- A color value between 0 and 1, based on how quick the blowup is
valueOf :: Complex Double -> Double
valueOf pos = (fromIntegral $ speedOfBlowup pos) / (fromIntegral maxDepth)

-- Mandelbrot settings
data Mandelbrot = Mandelbrot
  { width :: Int
  , height :: Int
  , scale :: Double
  , x_offset :: Double
  , y_offset :: Double
  }

makeMandelbrot :: Int -> Int -> Mandelbrot
makeMandelbrot w h = Mandelbrot
  { width = w
  , height = h
  , scale = 1
  , x_offset = 0
  , y_offset = 0
  }

-- Pixel coordinates to complex numer
coordinate :: Mandelbrot -> (Int, Int) -> Complex Double
coordinate m (x, y) =
  let
    x_proportional = fromIntegral x / fromIntegral (width m)
    y_proportional = fromIntegral y / fromIntegral (height m)

    -- Instead of going from 0 to 1, go from -2 to 2
    x_normalised = (x_proportional - 0.5) * (2 / 0.5)
    y_normalised = (y_proportional - 0.5) * (2 / 0.5)
  in
    ((x_normalised - x_offset m) / scale m) :+ ((y_normalised - y_offset m) / scale m)

-- Make a mandelbrot image of the given size and zoom
drawMandelbrot :: Mandelbrot -> Image VS Y Double
drawMandelbrot m = makeImage (width m, height m) $ \(y, x) -> PixelY $ valueOf $ coordinate m (x, y)


import qualified Numeric.LinearAlgebra as LA
import qualified Numeric.LinearAlgebra.Data as D (tr,  (><))

type X_Train = LA.Matrix Double
type Y_Train = LA.Matrix Double
type W = Double
type B = Double
type Output = LA.Matrix Double
type LearningRate = Double
-- | y = w * x + b for a single training example

ones :: Int -> Int -> LA.Matrix Double
ones x y = (x D.>< y) [1,1 ..]

output' :: W -> X_Train -> B -> Output
output' w xTrain b = w `LA.scale` xTrain + b `LA.scale` ones (fst (LA.size xTrain)) 1

cost :: Output -> Y_Train -> Double
cost output yTrain = LA.sumElements $ (1/2) * (output - yTrain) ^ 2

deltaCW' :: W -> X_Train -> B -> Double
deltaCW' w xTrain b = LA.sumElements $ D.tr (output' w xTrain b) <> xTrain

deltaCB' :: W -> X_Train -> Y_Train -> B -> Double
deltaCB' w xTrain yTrain b = LA.sumElements $ (output' w xTrain b - yTrain)

updateW' :: W -> X_Train -> B -> LearningRate -> Double
updateW' w xTrain b alpha = w - alpha * deltaCW' w xTrain b

updateB' :: W -> X_Train -> Y_Train -> B -> LearningRate -> Double
updateB' w xTrain yTrain b alpha = b - alpha * deltaCB' w xTrain yTrain b

test = do
  xTrain <- LA.rand 3 1
  yTrain <- LA.rand 3 1
  let w = 0
  let b = 0
  let alpha = 0.00001
  let loop i w' b' | i > 600 = return ()
      loop i w' b' | i <= 600 = do
         let c = cost (output' w' xTrain b') yTrain
         putStrLn $ concat ["Cost: ", show c]
         let w'' = updateW' w' xTrain b' alpha
         let b'' = updateB' w' xTrain yTrain b' alpha
         loop (i + 1) w'' b''
  loop 0 w b

{- ==== BAD CODE ====
Only for my own learning

type W = LA.Matrix Double
type B = LA.Matrix Double

output :: W -> X_Train -> B -> Output
output w xTrain b = (D.tr (w) <> xTrain) + b

deltaCW :: W -> X_Train -> Y_Train -> B -> LA.Matrix Double
deltaCW w xTrain yTrain b = ((D.tr (w) <> xTrain) + b) <> xTrain

deltaCB :: W -> X_Train -> Y_Train -> B -> LA.Matrix Double
deltaCB w xTrain yTrain b = (D.tr (w) <> xTrain) + b

updateW :: W -> X_Train -> Y_Train -> B -> LearningRate -> W
updateW w xTrain yTrain b alpha = w - alpha `LA.scale` deltaCW w xTrain yTrain b

updateB :: W -> X_Train -> Y_Train -> B -> LearningRate -> W
updateB w xTrain yTrain b alpha = b - alpha `LA.scale` deltaCB w xTrain yTrain b

main = do
  xTrain <- LA.rand 3 1
  yTrain <- LA.rand 3 1
  w <- LA.rand 3 1
  b <- LA.rand 1 1
  let alpha = 0.0001
  let loop i w' b' | i > 1000 = return ()
      loop i w' b' | i <= 1000 = do
         let c = cost (output w' xTrain b') yTrain
         putStrLn $ concat ["Cost: ", show c]
  --       putStrLn "\n"
         let w'' = updateW w' xTrain yTrain b' alpha
         let b'' = updateB w' xTrain yTrain b' alpha
--         putStrLn $ concat ["Weights: ", show w'']
--         putStrLn "\n"
--         putStrLn $ concat ["Bias: ", show b'']
         loop (i + 1) w'' b''
  loop 0 w b
-}


module Day3 (solvePart1, solvePart2) where

import           Data.Complex
import           Data.Map     (Map, (!))
import qualified Data.Map     as Map

createMovementList :: [Complex Double] -> [Int] -> [Complex Double]
createMovementList (f:s:rest) (x:xs) = replicate x f ++ replicate x s ++ createMovementList rest xs

movementList :: [Complex Double]
movementList = createMovementList (iterate (*(0 :+ 1)) (1 :+ 0)) [1..]

positions :: [Complex Double]
positions = scanl (+) (0 :+ 0) movementList

manhattenDistance :: Complex Double -> Int
manhattenDistance c = round (sum (fmap abs c))

solvePart1 :: String -> Int
solvePart1 input = let
        n = read input
    in manhattenDistance (positions !! (n-1))

solvePart2 :: String ->  Int
solvePart2 input = let
        n = read input
    in populateMap n

populateMap :: Int -> Int
populateMap n = let
        foobar (p:ps) m = let
                v = eval p m
            in if v <= n
                then foobar ps (Map.insert p v m)
                else v
    in foobar positions Map.empty

eval :: Complex Double -> Map (Complex Double) Int -> Int
eval c m
    | Map.null m = 1
    | otherwise  = sum (map snd (neighbours c m))

neighbours :: Complex Double -> Map (Complex Double) Int -> [(Complex Double, Int)]
neighbours (a :+ b) m = let
        member k = Map.member k m
        entry k = (k, m ! k)
        neighbourhood = [(a+x) :+ (b+y) | x <- [-1, 0, 1], y <- [-1, 0, 1], x /= 0 || y /= 0]
    in map entry (filter member neighbourhood)

instance (Ord a) => Ord (Complex a) where
   (a :+ b) `compare` (c :+ d)
        | a > c = GT
        | a < c = LT
        | b > d = GT
        | b < d = LT
        | b == d = EQ


{-# LANGUAGE FlexibleContexts #-}

module School.Utils.LinearAlgebra
( compareDoubleMatrix
, compareDoubleVector
, constVector
, oneMatrix
, oneVector
, mapCols
, mapRows
, sumCols
, sumRows
, zeroMatrix
, zeroVector
) where

import Numeric.LinearAlgebra
import School.Types.FloatEq (compareDouble)

mapCols :: (Element a)
        => (Vector a -> Vector a)
        -> Matrix a
        -> Matrix a
mapCols f = fromColumns . map f . toColumns

mapRows :: (Element a)
        => (Vector a -> Vector a)
        -> Matrix a
        -> Matrix a
mapRows f = fromRows . map f . toRows

sumCols :: Matrix R -> Vector R
sumCols m = n |> summedCols where
  n = rows m
  summedCols = map sumElements $ toRows m

sumRows :: Matrix R -> Vector R
sumRows m = n |> summedRows where
  n = cols m
  summedRows = map sumElements $ toColumns m

oneVector :: Int -> Vector R
oneVector = konst 1

constVector :: Int -> Double -> Vector R
constVector d v = konst v d

zeroVector :: Int -> Vector R
zeroVector = konst 0

zeroMatrix :: (Container Vector a, Num a)
           => Int
           -> Int
           -> Matrix a
zeroMatrix r c = konst 0 (r, c)

oneMatrix :: Int -> Int -> Matrix R
oneMatrix r c = konst 1 (r, c)

compareDoubleMatrix :: Double
                    -> Matrix R
                    -> Matrix R
                    -> Bool
compareDoubleMatrix prec m1 m2 =
  and $ zipWith (compareDouble prec) l1 l2
  where l1 = concat . toLists $ m1
        l2 = concat . toLists $ m2

compareDoubleVector :: Double
                    -> Vector R
                    -> Vector R
                    -> Bool
compareDoubleVector prec v1 v2 =
  and $ zipWith (compareDouble prec) l1 l2
  where l1 = toList v1
        l2 = toList v2


module Formulas where

import           Numeric.LinearAlgebra
import           Control.Monad.IO.Class
import           Control.Monad (replicateM)
import           System.Random
import qualified Data.Vector as V
import           Data.List

type X          = Matrix R
type Beta       = Vector R
type Y          = Vector R
type Block      = V.Vector Int
type Powers     = V.Vector Block
type Input      = V.Vector Double
type Inplist    = V.Vector Input
type Polynomial = (Beta, X)

-- V.Vectors, Vectors and Matrices Transformators
vvToM :: V.Vector (V.Vector Double)
      -> Matrix Double
vvToM = fromLists . V.toList . V.map V.toList

mToVV :: Matrix Double
      -> V.Vector (V.Vector Double)
mToVV = V.fromList . map V.fromList . toLists
-----------------------------------------------------

-- (5) in the article --
f :: Powers
  -> Input
  -> V.Vector Double
f pss inp = V.map oneBlock pss
    where
        oneBlock ps = V.product $ V.zipWith (**) inp $ V.map fromIntegral ps

calcX :: Inplist
      -> Powers -> X
calcX inputs pss = vvToM (V.map (f pss) inputs)

-- (~7) in the article --
cost :: Beta
     -> X
     -> Y
     -> Double
cost b x y = abs $ sumElements $ e**2
    where e = y - x #> b

-- (8) in the article --
normEq :: X
       -> Y
       -> Beta
normEq x y  = (inv (tr x <> x) <> tr x) #> y :: Beta
---------------- Show polynomial ---------------------
showPolynomial :: Beta
               -> Powers
               -> String
showPolynomial coeffs powers =
    intercalate " +" $ zipWith showCoeff coeffs' powers'
        where
            coeffs' = filter isNotZero $ toList coeffs
            powers' = V.toList $ V.map V.toList powers
            showCoeff coeff powerList = " " ++ show_ (round3 coeff) ++ matchVarsPowers powerList
            isNotZero x = round3 x /= 0

matchVarsPowers :: [Int] -> String
-- Example: [1, 4, 2] => x1 * x2^4 * x3^2
matchVarsPowers powers =
    intercalate "" $ zipWith varToPower varNums powers
        where
            varNums = [1..length powers]
            varToPower :: Int -> Int -> String
            varToPower ind 0   = ""
            varToPower ind 1   = " * x" ++ show ind
            varToPower ind pow = " * x" ++ show ind ++ "^" ++ show_ pow

-- show_ :: Num n => n -> n
show_ num = if num > 0 then show num else "(" ++ show num ++ ")"

    -------------------------------------------------------
randomPowerGen :: Int      -- m
               -> Int      -- n
               -> IO (V.Vector (V.Vector Int)) -- m >< n Random Powers
randomPowerGen m n = do
    lists <- replicateM m $ replicateM n $ randomRIO (-4, 4)
    return $ V.fromList $ map V.fromList lists

round3 :: Double -> Double
round3 x = (fromInteger $ round $ x * (10^3)) / (10.0^^3)


module Numeric.Morpheus.ActivationTest
(
  tests
)

where

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

import Numeric.LinearAlgebra
import Numeric.Morpheus.Activation

x = row [10e10, 0, -10e10]
x_sigmoid = row [1, 0.5, 0]
x_sigmoid_gradient = row [0, 0.25, 0]
x_relu = row [10e10, 0, 0]
x_relu_gradient = row [1, 0, 0]
x_tanh = row [1, 0, -1]
x_tanh_gradient = row [0, 1, 0]

tests = [ testGroup "Sigmoid" [
            testCase "Sigmoid" $ assertMatrix "" 1e-10 x_sigmoid (sigmoid x)
            , testCase "Sigmoid Gradient" $ assertMatrix "" 1e-10 x_sigmoid_gradient (sigmoidGradient x)
            ]
        , testGroup "ReLu" [
            testCase "ReLu" $ assertMatrix "" 1e-10 x_relu (relu x)
            , testCase "ReLu Gradient" $ assertMatrix "" 1e-10 x_relu_gradient (reluGradient x)
            ]
        , testGroup "Tanh" [
            testCase "Tanh" $ assertMatrix "" 1e-10 x_tanh (tanh_ x)
            , testCase "Tanh Gradient" $ assertMatrix "" 1e-10 x_tanh_gradient (tanhGradient x)
            ]
        ]


module Y2017.M02.D24.Exercise where

-- below imports from cabal package 'statistics'

import Statistics.Distribution
import Statistics.Distribution.Binomial

{--
AHA!

Whenever you think you have to invent something, stop, pause, and think, ...

SOMEBODY ELSE HAS ALREADY DONE IT! USE THE DARK SIDE OF THE FORCE AND TOGETHER
WE CAN RULE THE GALAXY AS FATHER AND SON!

See? All you need is love, a psychopathic father that murders young children
at the Jedi Academy, and 'midichlorians.' Lots of 'midichlorians.'

Where was I?

Ah, yes: 'already invented.'

So, we know from Rosalind.Traits that any descendant from an Aa Bb/Aa Bb
pairing will have a 25% probability of itself being Aa Bb (you can verify
this if you'd like), so the question arises, what is the probability of at
least n offspring at generation x (or Generation X, if you prefer) also being
Aa Bb.

So.

Today's Haskell problem: with the help of the above imports, create a binomial
probability distribution at 25%, run a scenario where you are in generation
x, where the xth generation has x^2 offspring, and you are looking for at least
n of the offspring to have the desired characteristic (in the abstract case,
the what the characteristic is doesn't matter, we need what the probably outcome
of having it.
--}

type Trials = Int

aAbB :: Trials -> BinomialDistribution
aAbB trials = undefined

genx :: Int -> BinomialDistribution -> Double
genx atleastN bin = undefined

-- what is the result for gen = 2, n = at least 1?

-- Now solve for gen = 7, n = at least 35.

-- hint: what is 'trials' for gen = x?
-- hint-hint: remember solving for atLeast n outcomes from your stats classes?


module Mandelbrot where

import Data.List
import Data.Complex
import Data.Maybe

type Point = Complex Double

iterations = 1000

mandelbrot :: Point -> Point -> Point
mandelbrot c z = z ^^ 2 + c

series :: Point -> [Point]
series c = iterate (mandelbrot c) (0 :+ 0)

period :: Point -> Maybe Int
period c = helper 0 (take iterations (series c)) [] where
  helper count (z:zs) memory = if z `elem` memory
    then z `elemIndex` memory >>= (\i -> Just (i + 1))
    else helper (count + 1) zs (z:memory)
  helper count [] memory = Nothing

escaped :: Point -> Bool
escaped c = 4 < (realPart c * realPart c + imagPart c * imagPart c)

-- Does the series at 'c' eventually escape?
divergent :: Point -> Bool
divergent c = any escaped (take iterations (series c))

-- How many iterations does it take the series at 'c' to escape?
divergence :: Point -> Int
divergence c = helper 0 (0.0 :+ 0.0) where
  helper n acc = if n < iterations
    then
      let c' = mandelbrot c acc in
        if escaped c' then n else helper (n + 1) c'
    else
      n


module Ch02.Perceptron where

import Prelude hiding (and, or)
import Numeric.LinearAlgebra

and_ :: (Ord a, Fractional a, Num b) => a -> a -> b
and_ x1 x2 = if (x1*w1 + x2*w2) <= theta then 0 else 1
  where
    w1    = 0.5
    w2    = 0.5
    theta = 0.7
-- λ> and_ 0 0
-- 0
-- λ> and_ 1 0
-- 0
-- λ> and_ 0 1
-- 0
-- λ> and_ 1 1
-- 1

perceptron :: R -> R -> R -> R -> R -> R
perceptron w1 w2 b = \x1 x2 ->
  if (sum $ toList $ (vector [x1,x2]) * vector [w1,w2]) + b <= 0 then 0 else 1

and       = perceptron   0.5    0.5  (-0.7)
nand      = perceptron (-0.5) (-0.5)   0.7
or        = perceptron   0.5    0.5  (-0.2)
xor x1 x2 = and (nand x1 x2) (or x1 x2)


-- example 3.21

import Statistics.Distribution
import Statistics.Distribution.Binomial
import Statistics.Distribution.Poisson

bi = binomial n p
  where
    n = 20   -- trials
    p =  0.1 -- probability

pd = poisson l
  where
    l = n*p
    n = 20
    p =  0.1

main = do
  putStrLn $ "with binomial, P(X<=3):              " ++ show (cumulative bi 3.0)
  putStrLn $ "with poisson approximation, P(X<=3): " ++ show (cumulative pd 3.0)


-- | Simple linear regression example for the README.

import Lib(someFunc)

import Control.Monad (replicateM, replicateM_, zipWithM)
import System.Random (randomIO)
import Test.HUnit (assertBool)

import qualified TensorFlow.Core as TF
import qualified TensorFlow.GenOps.Core as TF
import qualified TensorFlow.Gradient as TF
import qualified TensorFlow.Ops as TF

import qualified Data.Random.Normal as Normal
import qualified Statistics.Sample as Stat
import qualified Data.Vector as Vector
import Text.Printf

-- | R² of the model Ŷ = aX+b
-- Note that 'TF.gradient' doesn't support Double
linearModelRSquared :: [Float] -> [Float] -> Float -> Float
                    -> Float
linearModelRSquared xs ys a b = 1 - ssRes / ssTot
        where
                es    = zipWith (-) ys yHats
                yHats = [ a*x+b | x <- xs ]
                -- realToFrac is only ok in this situation
                yMean = realToFrac $ Stat.mean $ Vector.fromList $ (map realToFrac ys)
                ssTot = sum [ (y - yMean) ^ 2 | y <- ys ]
                ssReg = sum [ (y - yMean) ^ 2 | y <- yHats ]
                ssRes = sum [ e ^ 2 | e <- es ]

main :: IO ()
main = do
        -- Generate a list of Gaussian errors
        noise <- replicateM 100 (Normal.normalIO' (0, 0.1))

        let xData = tail [0,0.01..1]
            yData = [ 3*x + 8 + e | (x,e) <- zip xData noise]

       -- Fit the linear regression model.
        (w, b) <- fit xData yData

       -- Calculate and print the R-squared for the model
        let rSquared = linearModelRSquared xData yData w b
        printf "R^2 = %f\n" rSquared

fit :: [Float] -> [Float] -> IO (Float, Float)
fit xData yData = TF.runSession $ do
    -- Create tensorflow constants for x and y.
    let x = TF.vector xData
        y = TF.vector yData
    -- Create scalar variables for slope and intercept.
    w <- TF.initializedVariable 0
    b <- TF.initializedVariable 0
    -- Define the loss function.
    let yHat = (x `TF.mul` w) `TF.add` b
        loss = TF.square (yHat `TF.sub` y)
    -- Optimize with gradient descent.
    trainStep <- gradientDescent 0.001 loss [w, b]
    replicateM_ 1000 (TF.run trainStep)
    -- Return the learned parameters.
    (TF.Scalar w', TF.Scalar b') <- TF.run (w, b)
    return (w', b')

gradientDescent :: Float
                -> TF.Tensor TF.Build Float
                -> [TF.Tensor TF.Ref Float]
                -> TF.Session TF.ControlNode
gradientDescent alpha loss params = do
    let applyGrad param grad =
            TF.assign param (param `TF.sub` (TF.scalar alpha `TF.mul` grad))
    TF.group =<< zipWithM applyGrad params =<< TF.gradients loss params


{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE RecordWildCards #-}
module LinearRegression where

import qualified Data.Vector as V
import Numeric.LinearAlgebra
import Prelude hiding ((<>))
import Util

data HyperParameters = HyperParams { theta :: Matrix Double
                                   , alpha :: Double
                                   , maxIterations :: Int
                                   , tolerance :: Double }

computeCost :: Matrix Double -- Features
            -> Matrix Double -- Labels
            -> Matrix Double -- Theta
            -> Double
computeCost x y theta = (sumElements squareDifference) / (2.0 * m)
  where m = fromIntegral ((fst . size) y)
        h = x <> theta
        squareDifference = (h - y) ^ 2

-- A simple gradient descent algorithm that specifies an arbitrary number of iterations.
gradientDescent :: Matrix Double -- Features
                -> Matrix Double -- Labels
                -> Matrix Double -- Parameters
                -> Double        -- Learning Rate
                -> Int           -- Number of iterations
                -> ([Double], Matrix Double)
gradientDescent x y theta alpha numIters = gradientDescent' x y theta alpha numIters []

gradientDescent' :: Matrix Double -- Features
                 -> Matrix Double -- Labels
                 -> Matrix Double -- Parameters
                 -> Double        -- Learning Rate
                 -> Int           -- Number of Iterations
                 -> [Double]      -- Cost history
                 -> ([Double], Matrix Double)
gradientDescent' x y theta alpha iterations hist
  | iterations > 0   = gradientDescent' x y theta' alpha (iterations - 1) hist'
  | otherwise        = (hist' , theta)
    where cost  = computeCost x y theta
          hist' = cost : hist
          theta' = theta - step
          step = (scalar alpha) * ((tr' x) <> ((x <> theta) - y)) / (scalar m)
          m = fromIntegral ((fst . size) y)

featureNormalize :: Matrix Double -> (Matrix Double, Matrix Double, Matrix Double)
featureNormalize x = ((x - mu) / sigma, mu, sigma)
  where mu = mean x
        sigma = std x

normalEquation :: Matrix Double -> Matrix Double -> Matrix Double
normalEquation x y = (pinvTol 1e-2 (tr' x <> x)) <> (tr' x) <> y


{-# LANGUAGE Rank2Types #-}
-----------------------------------------------------------------------------
-- |
-- Module     : Numeric.LinearAlgebra.Vector
-- Copyright  : Copyright (c) , Patrick Perry <patperry@gmail.com>
-- License    : BSD3
-- Maintainer : Patrick Perry <patperry@gmail.com>
-- Stability  : experimental
--
-- Immutable dense vectors.

module Numeric.LinearAlgebra.Vector (
    -- * Immutable vectors
    Vector,
    dim,

    -- * Vector construction
    fromList,
    zero,
    constant,

    -- * Accessing vectors
    at,
    unsafeAt,
    indices,
    elems,
    assocs,

    -- * Incremental vector updates
    update,
    unsafeUpdate,
    accum,
    unsafeAccum,

    -- * Derived vectors
    map,
    zipWith,
    unsafeZipWith,
    concat,

    -- * Vector views
    slice,
    unsafeSlice,
    splitAt,
    drop,
    take,

    -- * Vector properties
    sumAbs,
    norm2,
    whichMaxAbs,
    dot,
    unsafeDot,
    kronecker,

    -- * Vector math operations
    -- ** Num
    add,
    addWithScale,
    sub,
    scale,
    mul,
    negate,
    conjugate,
    abs,
    signum,

    -- ** Fractional
    div,
    recip,

    -- ** Floating
    sqrt,
    exp,
    log,
    pow,
    sin,
    cos,
    tan,
    asin,
    acos,
    atan,
    sinh,
    cosh,
    tanh,
    asinh,
    acosh,
    atanh,

    -- * Conversions between foreign pointers
    unsafeFromForeignPtr,
    unsafeToForeignPtr,

    -- * Mutable interface
    module Numeric.LinearAlgebra.Vector.ST,

    -- * Basic multivariate statistics
    module Numeric.LinearAlgebra.Vector.Statistics,

    ) where

import Prelude( Int, Double, ($), (*), return )
import Control.Monad.ST( runST, ST )
import Foreign( Storable )
import Foreign.BLAS( BLAS1, BLAS2 )
import Foreign.VMath( VNum, VFractional, VFloating )

import Numeric.LinearAlgebra.Vector.Base
import Numeric.LinearAlgebra.Vector.ST
import Numeric.LinearAlgebra.Vector.Statistics

infixr 8 `pow`
infixl 7 `div`
infixl 7 `mul`, `scale`, `kronecker`
infixl 6 `add`, `sub`

-- | Compute the sum of absolute values of entries in the vector.
sumAbs :: (BLAS1 e) => Vector e -> Double
sumAbs v = runST $ getSumAbs v
{-# INLINE sumAbs #-}

-- | Compute the 2-norm (Euclidean norm) of a vector.
norm2 :: (BLAS1 e) => Vector e -> Double
norm2 v = runST $ getNorm2 v
{-# INLINE norm2 #-}

-- | Get the index and norm of the element with absulte value.  Not valid
-- if any of the vector entries are @NaN@.  Raises an exception if the
-- vector has length @0@.
whichMaxAbs :: (BLAS1 e) => Vector e -> (Int, e)
whichMaxAbs v = runST $ getWhichMaxAbs v
{-# INLINE whichMaxAbs #-}

-- | Compute the dot product of two vectors.
dot :: (BLAS1 e) => Vector e -> Vector e -> e
dot v v' = runST $ getDot v v'
{-# INLINE dot #-}

unsafeDot :: (BLAS1 e) => Vector e -> Vector e -> e
unsafeDot v v' = runST $ unsafeGetDot v v'
{-# INLINE unsafeDot #-}

-- | Compute the kronecker product of two vectors.
kronecker :: (BLAS2 e) => Vector e -> Vector e -> Vector e
kronecker x y = create $ do
    z <- new_ (dim x * dim y)
    kroneckerTo z x y
    return z

-- | @add x y@ returns @x + y@.
add :: (VNum e) => Vector e -> Vector e -> Vector e
add = result2 addTo

-- | @sub x y@ returns @x - y@.
sub :: (VNum e) => Vector e -> Vector e -> Vector e
sub = result2 subTo

-- | @scale k x@ returns @k * x@.
scale :: (BLAS1 e) => e -> Vector e -> Vector e
scale k x = create $ do
    x' <- newCopy x
    scaleM_ k x'
    return x'

-- | @addWithScale alpha x y@ return @alpha * x + y@.
addWithScale :: (BLAS1 e) => e -> Vector e -> Vector e -> Vector e
addWithScale alpha x y = create $ do
    y' <- newCopy y
    addWithScaleM_ alpha x y'
    return y'

-- | @mul x y@ returns @x * y@.
mul :: (VNum e) => Vector e -> Vector e -> Vector e
mul = result2 mulTo

-- | @negate x@ returns @-x@.
negate :: (VNum e) => Vector e -> Vector e
negate = result negateTo

-- | @conjugate x@ returns @conjugate(x)@.
conjugate :: (VNum e) => Vector e -> Vector e
conjugate = result conjugateTo

-- | @abs x@ returns @abs(x)@.
abs :: (VNum e) => Vector e -> Vector e
abs = result absTo

-- | @signum x@ returns @signum(x)@.
signum :: (VNum e) => Vector e -> Vector e
signum = result signumTo

-- | @div x y@ returns @x / y@.
div :: (VFractional e) => Vector e -> Vector e -> Vector e
div = result2 divTo

-- | @recip x y@ returns @1 / x@.
recip :: (VFractional e) => Vector e -> Vector e
recip = result recipTo

-- | @sqrt x@ returns @sqrt(x)@.
sqrt :: (VFloating e) => Vector e -> Vector e
sqrt = result sqrtTo

-- | @exp x@ returns @exp(x)@.
exp :: (VFloating e) => Vector e -> Vector e
exp = result expTo

-- | @log x@ returns @log(x)@.
log :: (VFloating e) => Vector e -> Vector e
log = result logTo

-- | @pow x y@ returns @x ** y@.
pow :: (VFloating e) => Vector e -> Vector e -> Vector e
pow = result2 powTo

-- | @sin x@ returns @sin(x)@.
sin :: (VFloating e) => Vector e -> Vector e
sin = result sinTo

-- | @cos x@ returns @cos(x)@.
cos :: (VFloating e) => Vector e -> Vector e
cos = result cosTo

-- | @tan x@ returns @tan(x)@.
tan :: (VFloating e) => Vector e -> Vector e
tan = result tanTo

-- | @asin x@ returns @asin(x)@.
asin :: (VFloating e) => Vector e -> Vector e
asin = result asinTo

-- | @acos x@ returns @acos(x)@.
acos :: (VFloating e) => Vector e -> Vector e
acos = result acosTo

-- | @atan x@ returns @atan(x)@.
atan :: (VFloating e) => Vector e -> Vector e
atan = result atanTo

-- | @sinh x@ returns @sinh(x)@.
sinh :: (VFloating e) => Vector e -> Vector e
sinh = result sinhTo

-- | @cosh x@ returns @cosh(x)@.
cosh :: (VFloating e) => Vector e -> Vector e
cosh = result coshTo

-- | @tanh x@ returns @tanh(x)@.
tanh :: (VFloating e) => Vector e -> Vector e
tanh = result tanhTo

-- | @asinh x@ returns @asinh(x)@.
asinh :: (VFloating e) => Vector e -> Vector e
asinh = result asinhTo

-- | @acosh x@ returns @acosh(x)@.
acosh :: (VFloating e) => Vector e -> Vector e
acosh = result acoshTo

-- | @atanh x@ returns @atanh(x)@.
atanh :: (VFloating e) => Vector e -> Vector e
atanh = result atanhTo

result :: (Storable e, Storable f)
       => (forall s . STVector s f -> Vector e -> ST s a)
       -> Vector e
       -> Vector f
result f v = create $ newResult f v
{-# INLINE result #-}

result2 :: (Storable e, Storable f, Storable g)
        => (forall s . STVector s g -> Vector e -> Vector f -> ST s a)
        -> Vector e
        -> Vector f
        -> Vector g
result2 f v1 v2 = create $ newResult2 f v1 v2
{-# INLINE result2 #-}

newResult :: (RVector v, Storable e, Storable f)
          => (STVector s f -> v e -> ST s a)
          -> v e
          -> ST s (STVector s f)
newResult f v = do
    n <- getDim v
    z <- new_ n
    _ <- f z v
    return z
{-# INLINE newResult #-}

newResult2 :: (RVector v1, RVector v2, Storable e, Storable f, Storable g)
           => (STVector s g -> v1 e -> v2 f -> ST s a)
           -> v1 e
           -> v2 f
           -> ST s (STVector s g)
newResult2 f v1 v2 = do
    n <- getDim v1
    z <- new_ n
    _ <- f z v1 v2
    return z
{-# INLINE newResult2 #-}