{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE BangPatterns #-}
{- Test distribution by comparing output with theoretical results
 -
 - Kolmogorov–Smirnov
 -
 -}

import           System.Environment (getArgs)
import           Data.List (intersperse)
import           Text.Printf

import           Control.Monad ( (>=>) )

import qualified Data.Vector.Unboxed as V
import           Data.Vector.Algorithms.Intro (sort)
import qualified Random.CRI.MWC as R
import           Statistics.Distribution
import           Statistics.Distribution.Gamma
import           Numeric.MathFunctions.Constants
import           Statistics.Test.KolmogorovSmirnov

import           Statistics.Distribution.Random.Gamma
import           Statistics.Sampler.Slice

import           Plot

--import Test.Framework

main :: IO()
main = mainSlice
--    args <- filter (/= "--tests") `fmap` getArgs
--    flip defaultMainWithArgs args
--        [ Statistics.Distribution.Random.Gamma.testSuite
--        ]

-- slice test
mainSlice :: IO ()
mainSlice = do

    args <- getArgs
    let size :: Int
        title :: String
        outfile :: String
        (size, sampler, gd, initial, title, outfile) =
          case args of
            [size', "gamma", shape', rate' ] ->
                 (read size', slice slst (log . (density  gammad)), gammad, shape/rate, title', outfile')
                    where
                       shape :: Double = read shape'
                       rate  :: Double = read rate'
                       gammad  = gammaDistr shape (1/rate)
                       slst   = adaptOff (newSlicerState 0.0001 m_pos_inf 1)
                       title' = printf "shape = %f, rate = %f" shape rate
                       outfile' =  printf "slice-%f-%f.csv" shape rate

            _ -> error "Usage: ./test size gamma shape rate"

    rng <- R.create

    --  get samples
    let sample res = do
        !(_, ns) <- sampler (head res) rng
        return (ns:res)
    samples <- foldr1 (>=>) (replicate (size-1) sample) [initial]
    let v = V.fromList samples

    -- print out samples
    let pv = "slice\n" ++ (concat $ intersperse "\n" $ map show $ V.toList v)
    writeFile outfile pv
    --  benchmark times

    --  calc K-S test statistic
    let d = kolmogorovSmirnovD gd v
    let p = kolmogorovSmirnovProbability (V.length v) d

    putStrLn $ printf "D = %f, p = %f" d p

    --  generate plot(s)
    let vs = V.modify sort v
    let cm = V.map (cumulative gd) vs

    _ <- plot cm title "test-slice.svg"

    _ <- plotDensity gd vs title "test-slice-density.svg"

    putStrLn $ "Done "++ (show $ V.length vs)

-- original test
mainGamma :: IO ()
mainGamma = do

    args <- getArgs

    let size  :: Int
        shape :: Double
        rate  :: Double
        (size, shape, rate) =
          case args of
            [size', shape', rate'] -> (read size', read shape', read rate')
            _                      -> error "Usage: ./test size shape rate"

    rng <- R.create
    let gs = gamma shape (1/rate)

    --  get samples
    v <- V.replicateM size (gs rng)

    -- print out samples
    let pv = "gammas\n" ++ (concat $ intersperse "\n" $ map show $ V.toList v)
    writeFile (printf "gammas-%f-%f.csv" shape rate) pv
    --  benchmark times

    --  calc K-S test statistic
    let gd = gammaDistr shape (1/rate)
    let d = kolmogorovSmirnovD gd v
    let p = kolmogorovSmirnovProbability (V.length v) d

    putStrLn $ printf "D = %f, p = %f" d p

    --  generate plot(s)
    let vs = V.modify sort v
    let cm = V.map (cumulative gd) vs
    let title = printf "D = %f, p-value = %f" d p

    _ <- plot cm title "test.svg"

    let titledensity = printf "shape = %f, rate = %f" shape rate
    _ <- plotDensity gd vs titledensity "test-density.svg"

    putStrLn $ "Done "++ (show $ V.length vs)


{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE DeriveDataTypeable, DeriveGeneric #-}
-- |
-- Module    : Statistics.Distribution.Geometric
-- Copyright : (c) 2009 Bryan O'Sullivan
-- License   : BSD3
--
-- Maintainer  : bos@serpentine.com
-- Stability   : experimental
-- Portability : portable
--
-- The Geometric distribution. There are two variants of
-- distribution. First is the probability distribution of the number
-- of Bernoulli trials needed to get one success, supported on the set
-- [1,2..] ('GeometricDistribution'). Sometimes it's referred to as
-- the /shifted/ geometric distribution to distinguish from another
-- one.
--
-- Second variant is probability distribution of the number of
-- failures before first success, defined over the set [0,1..]
-- ('GeometricDistribution0').
module Statistics.Distribution.Geometric
    (
      GeometricDistribution
    , GeometricDistribution0
    -- * Constructors
    , geometric
    , geometricE
    , geometric0
    , geometric0E
    -- ** Accessors
    , gdSuccess
    , gdSuccess0
    ) where

import Control.Applicative
import Control.Monad       (liftM)
import Data.Aeson          (FromJSON(..), ToJSON, Value(..), (.:))
import Data.Binary         (Binary(..))
import Data.Data           (Data, Typeable)
import GHC.Generics        (Generic)
import Numeric.MathFunctions.Constants (m_pos_inf, m_neg_inf)
import qualified System.Random.MWC.Distributions as MWC

import qualified Statistics.Distribution as D
import Statistics.Internal

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

-- | Distribution over [1..]
newtype GeometricDistribution = GD {
      gdSuccess :: Double
    } deriving (Eq, Typeable, Data, Generic)

instance Show GeometricDistribution where
  showsPrec i (GD x) = defaultShow1 "geometric" x i
instance Read GeometricDistribution where
  readPrec = defaultReadPrecM1 "geometric" geometricE

instance ToJSON GeometricDistribution
instance FromJSON GeometricDistribution where
  parseJSON (Object v) = do
    x <- v .: "gdSuccess"
    maybe (fail $ errMsg x) return  $ geometricE x
  parseJSON _ = empty

instance Binary GeometricDistribution where
  put (GD x) = put x
  get = do
    x <- get
    maybe (fail $ errMsg x) return  $ geometricE x

instance D.Distribution GeometricDistribution where
    cumulative = cumulative

instance D.DiscreteDistr GeometricDistribution where
    probability (GD s) n
      | n < 1     = 0
      | otherwise = s * (1-s) ** (fromIntegral n - 1)
    logProbability (GD s) n
       | n < 1     = m_neg_inf
       | otherwise = log s + log (1-s) * (fromIntegral n - 1)

instance D.Mean GeometricDistribution where
    mean (GD s) = 1 / s

instance D.Variance GeometricDistribution where
    variance (GD s) = (1 - s) / (s * s)

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

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

instance D.Entropy GeometricDistribution where
  entropy (GD s)
    | s == 0 = m_pos_inf
    | s == 1 = 0
    | otherwise = negate $ (s * log s + (1-s) * log (1-s)) / s

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

instance D.DiscreteGen GeometricDistribution where
  genDiscreteVar (GD s) g = MWC.geometric1 s g

instance D.ContGen GeometricDistribution where
  genContVar d g = fromIntegral `liftM` D.genDiscreteVar d g

cumulative :: GeometricDistribution -> Double -> Double
cumulative (GD s) x
  | x < 1        = 0
  | isInfinite x = 1
  | isNaN      x = error "Statistics.Distribution.Geometric.cumulative: NaN input"
  | otherwise    = 1 - (1-s) ^ (floor x :: Int)

-- | Create geometric distribution.
geometric :: Double                -- ^ Success rate
          -> GeometricDistribution
geometric x = maybe (error $ errMsg x) id $ geometricE x

-- | Create geometric distribution.
geometricE :: Double                -- ^ Success rate
           -> Maybe GeometricDistribution
geometricE x
  | x >= 0 && x <= 1 = Just (GD x)
  | otherwise        = Nothing

errMsg :: Double -> String
errMsg x = "Statistics.Distribution.Geometric.geometric: probability must be in [0,1] range. Got " ++ show x

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

-- | Distribution over [0..]
newtype GeometricDistribution0 = GD0 {
      gdSuccess0 :: Double
    } deriving (Eq, Typeable, Data, Generic)

instance Show GeometricDistribution0 where
  showsPrec i (GD0 x) = defaultShow1 "geometric0" x i
instance Read GeometricDistribution0 where
  readPrec = defaultReadPrecM1 "geometric0" geometric0E

instance ToJSON GeometricDistribution0
instance FromJSON GeometricDistribution0 where
  parseJSON (Object v) = do
    x <- v .: "gdSuccess0"
    maybe (fail $ errMsg x) return  $ geometric0E x
  parseJSON _ = empty

instance Binary GeometricDistribution0 where
  put (GD0 x) = put x
  get = do
    x <- get
    maybe (fail $ errMsg x) return  $ geometric0E x

instance D.Distribution GeometricDistribution0 where
    cumulative (GD0 s) x = cumulative (GD s) (x + 1)

instance D.DiscreteDistr GeometricDistribution0 where
    probability    (GD0 s) n = D.probability    (GD s) (n + 1)
    logProbability (GD0 s) n = D.logProbability (GD s) (n + 1)

instance D.Mean GeometricDistribution0 where
    mean (GD0 s) = 1 / s - 1

instance D.Variance GeometricDistribution0 where
    variance (GD0 s) = D.variance (GD s)

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

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

instance D.Entropy GeometricDistribution0 where
  entropy (GD0 s) = D.entropy (GD s)

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

instance D.DiscreteGen GeometricDistribution0 where
  genDiscreteVar (GD0 s) g = MWC.geometric0 s g

instance D.ContGen GeometricDistribution0 where
  genContVar d g = fromIntegral `liftM` D.genDiscreteVar d g

-- | Create geometric distribution.
geometric0 :: Double                -- ^ Success rate
           -> GeometricDistribution0
geometric0 x = maybe (error $ errMsg0 x) id $ geometric0E x

-- | Create geometric distribution.
geometric0E :: Double                -- ^ Success rate
            -> Maybe GeometricDistribution0
geometric0E x
  | x >= 0 && x <= 1 = Just (GD0 x)
  | otherwise        = Nothing

errMsg0 :: Double -> String
errMsg0 x = "Statistics.Distribution.Geometric.geometric0: probability must be in [0,1] range. Got " ++ show x


{-# LANGUAGE NamedFieldPuns, TemplateHaskell #-}

module School.Unit.Test.RecLin
( recLinTest ) where

import Control.Monad.IO.Class (liftIO)
import Numeric.LinearAlgebra ((><), fromLists, size, toLists)
import School.TestUtils (diffInput, matIndexes, randomMatrix)
import School.Types.FloatEq ((~=))
import School.Unit.RecLin
import School.Unit.Unit (Unit(..))
import School.Unit.UnitActivation (UnitActivation(..))
import School.Unit.UnitGradient (UnitGradient(..))
import School.Unit.UnitParams (UnitParams(..))
import School.Utils.LinearAlgebra (compareDoubleMatrix, oneMatrix)
import Test.Tasty (TestTree)
import Test.Tasty.QuickCheck hiding ((><), scale)
import Test.Tasty.TH
import Test.QuickCheck.Monadic (assert, monadicIO, pre)

eps :: Double
eps = 1e-5

prec :: Double
prec = 1e-4

prop_apply :: (Positive Int) -> (Positive Int) -> Property
prop_apply (Positive bSize) (Positive fSize) = monadicIO $ do
  m <- liftIO $ randomMatrix bSize fSize
  let input = BatchActivation m
  let params = EmptyParams
  let (BatchActivation result) = apply recLin params input
  let check = fromLists $ map (map (max 0)) $ toLists m
  assert $ result ~= check

prop_gradient_dims :: (Positive Int) -> Positive Int -> Property
prop_gradient_dims (Positive bSize) (Positive fSize) = monadicIO $ do
  input <- liftIO $ fmap BatchActivation $ randomMatrix bSize fSize
  gradIn <- liftIO $ fmap BatchGradient $ randomMatrix bSize fSize
  let params = EmptyParams
  let (BatchGradient gradOut, _) = deriv recLin params gradIn input
  assert $ (size gradOut == (bSize, fSize))

prop_numerical_gradient :: (Positive Int) -> Positive Int -> Property
prop_numerical_gradient (Positive bSize) (Positive fSize) = monadicIO $ do
  pre $ bSize < 25 && fSize < 25
  input <- liftIO $ fmap BatchActivation $ randomMatrix bSize fSize
  let params = EmptyParams
  let idxs = matIndexes bSize fSize
  let num = map (\idx -> diffInput recLin params input eps idx)
                idxs
  let numGrad = (bSize >< fSize) num
  let inGrad = BatchGradient $ oneMatrix bSize fSize
  let (BatchGradient result, _) = deriv recLin params inGrad input
  assert $ compareDoubleMatrix prec result numGrad

recLinTest :: TestTree
recLinTest = $(testGroupGenerator)


module Lec7 where

-- Функторы
--
-- Функция map была бы логична не только для списков. Как бы она работала для Maybe?
-- f x = x + 1
-- x1 = Just 42
-- x2 = Nothing
--
-- -хочу применить f к x1, x2, точнее, к значению внутри
-- maybeMap fun (Just n) = Just (fun n)
-- maybeMap _ Nothing = Nothing
-- print $ maybeMap f x1  -> Just 43
-- print $ maybeMap f x2  -> Nothing
--
-- Допустим, я хочу уметь так же применять функции ко вторым элементам пары:
-- f x = x + 1
-- x1 = ("hello", 42)
-- x2 = ([1, 2, 3], 42.0)
-- хочу применить f ко второму элементу
-- pairMap fun (x, y) = (x, fun y)
-- print $ pairMap f x1  -> ("hello", 43)
-- print $ pairMap f x2  -> ([1, 2, 3], 43.0)
--
-- Есть и другие аналогичные ситуации. Когда значение (или значения) хранятся в какой-то структуре с дополнительными
-- данными.
-- Список - несколько значений и информация об их порядке
-- Maybe - одно значение или ноль
-- Пара - значение вместе с дополнительной информацией
--
-- Функции тоже можно понимать как такие структуры:
-- fun :: a -> b
-- получается, функция хранит значение типа B, доступ к которым осуществляется по аргументу типа a:
-- Например,
-- f 1 = "один"
-- f 2 = "два"
-- функция "хранит" внутри себя значения "один", "два", получить их можно, подставив соответствующий аргумент в функцию.
-- f x = x + 1
-- x1 = \x -> 2 * x
-- x2 = \x -> 3 * x
-- funMap f x1 - должна получиться функция, которая сопадает с x1, но после применения x1 дополнительно вызывает f
-- funMap fun x = \t -> fun (x t)
--
-- xx1 = funMap f x1
-- xx2 = funMap f x2
-- print $ xx1 42  -> 85
-- print $ xx2 42  -> 127
--
-- На самом деле все эти map функции уже есть. Они называются fmap или <$>, и они делают ровно то же, что делаем мы:
-- print $ fmap (+1) (Just 42)  -> Just 43
-- print $ (+1) <$> Just 42  -> Just 43
-- print $ (+1) <$> [10, 20, 30]  -> [11, 21, 31]
-- print $ (+1) <$> ("abc", 123)  -? ("abc", 124)
--
-- Это работает, потому что в Haskell есть класс типов Functor - функтор. Он определен так:
-- class Functor f where
--   fmap :: (a -> b) -> f a -> f b
--
-- У него есть реализации, например, примерно так выглядит реализация для Maybe.
-- instance Functor Maybe where
--   fmap f (Just x) = Just (f x)
--   fmap _ Nothing = Nothing
--
-- Если вы определяете функтор, вы должны гарантировать выполнение следующего правила. Функция id - это функция, которая
-- ничего не делает: id x = x, её тип a -> a
-- **Правило**:
-- 1) fmap id x == x (или эквивалентно) fmap id == id
-- 2) fmap (g . f) x == fmap g . fmap f
--
-- ## Аппликативные функторы
-- Расширение обычных функторов дополнительными операциями. Представьте опять же Maybe.
-- Было в функторах:
-- 1) Обычная функция
-- 2) Данные внутри структуры
-- 3) Применяем функцию к каждому данному внутри структуры
-- Теперь в аппликативных функторах можно дополнительно:
-- 1) Обычные функция внутри структуры
-- 2) Данные внутри структуры
-- 3) Применяем функции из структуры к данным. Как именно применять - очень сильно зависит от ситуации.
--
-- Пример с Maybe:
-- f = Just (+1) -- увеличение на 1, но внутри Maybe
-- x = Just 42 -- значение внутри Maybe
-- применение:
-- maybeApply (Just fun) (Just x) = Just (fun x)
-- maybeApply _ _ = Nothing
--
-- print $ maybeApply f x  -> Just 43
-- print $ maybeApply Nothing x  -> Nothing
-- print $ maybeApply (Just (*2)) Nothing  -> Nothing
--
-- А для списков?
-- f = [(+1), (*2)]
-- x = [10, 20, 30]
-- что может означать listApply, когда есть список функций и список значений?
-- Применить каждую функцию к каждому элементу, получить список новых значений.
--
-- listApply fs xs = concatMap (\f -> map f xs) fs
-- print $ listApply f x  -> [11, 21, 31, 20, 40, 60]
-- Немного философии, почему для списков это логично делать именно так. Можно считать, что Maybe - это вычисления,
-- в которых иногда отсутствуют результаты. А в списках можно считать, что мы делаем вычисления, в которых может быть
-- несколько результатов.
-- В Haskell есть класс типов Аппликативные функторы:
-- class Functor f => Applicative f where
--   pure :: a -> f a -- вставить значение в функтор
--   (<*>) :: f (a -> b) -> f a -> f b -- та самая функция применения, рассмотренная выше
--   -- и ещё несколько, см. дальше
--
-- pure - это функция, которая помещает значение внутрь функции, "самым естественным способом". Что это значит будет
-- понятно из требований к аппликативному функтору, см. дальше.
-- Для Maybe: pure x = Just x, для списка pure x = [x]
--
-- Допустим, хотим сложить значения в двух Maybe. Или в двух списках сложить попарно значения. Как это будет написано?
-- x1 = Just 110
-- x2 = Just 654
-- надо сложить. Сначала давайте получим Just (+110), потом применим её к Just 654:
-- (+) <*> x1 <*> x2
--
-- Аппликативные функторы должны уметь делать не только <*>, но и функцию pure, которая вводит обычное значение в функтор.
-- 'pure 42' в типе Maybe это 'Just 42' . 'pure 42' в типе списка: [42]
--
-- Правила для функторов:
--   pure id <*> v = v
--   pure f <*> pure x = pure (f x)
--   и др.
--
--
-- #Монады
-- Вычисления с Maybe - это вычисления, в которых может не получиться результата. Если мы делаем последовательность
-- вычислений, в которой в какой-то момент результата не получилось, то и дальше считаем, что результата нет.
-- mSqrt :: Double -> Maybe Double
-- mSqrt x | x < 0 = Nothing
--         | otherwise = Just $ sqrt x
-- [хотим вычислить (sqrt(x) + 1) / 2]
-- eval x = let sx = mSqrt x
--              y = (+1) <$> sx
--            in (/2) <$> y
-- eval 9 -> Just 2.0
-- eval (-9) -> Nothing
--
-- sqrt(x) + sqrt(y) + 1
-- eval2 x y = let sx = mSqrt x
--                 sy = mSqrt y
--                 sq_plus_sq = Just (+) <*> sx <*> sy -- liftA2 (+) sx sy
--                in
--                 (+1) <$> sq_plus_sq
-- eval2 9 16 -> Just 8.0
-- eval2 (-9) 16 -> Nothing
--
-- sqrt(sqrt(x) - 3)
-- eval3 x = let sx = mSqrt x
--               sx_3 = (\t -> t - 3) <$> sx
--               -- ?? не получается сделать ни с <*>, ни с <$>
--               -- есть Just y и есть mSqrt :: Double -> Maybe Double
--               z = case sx_3 of
--                      Just n -> mSqrt n
--                      Nothing -> Nothing
--              in
--               z
-- eval3 100 -> Just ...
-- eval3 (-100) -> Nothing
-- eval3 4 -> Nothing
--
-- Хочется вспомогательную функцию типа <*> и <$>, чтобы она имела такой заголовок
-- bind :: Maybe a -> (a -> Maybe b) -> Maybe b
--
-- Если подумать, то
-- 1) fmap - частный случай этой функции. Сравним
--    bind :: Maybe a -> (a -> Maybe b) -> Maybe b
--    fmap :: Maybe a -> (a -> b) -> Maybe b
--    fmap f t = bind t (\x -> pure $ f x) или fmap f t = bind t (pure . f)
-- 2) <*> - тоже частный случай. Попробуйте её сам выразить через bind
--
-- Реализуем bind для Maybe:
-- maybeBind :: Maybe a -> (a -> Maybe b) -> Maybe b
-- maybeBind Nothing _ = Nothing
-- maybeBind (Just x) f = f x
--
-- снова sqrt(sqrt(x) - 3)
-- eval3 x = let sx = mSqrt(x)
--               sx_3 = maybeBind sx (\t -> Just (t - 3))
--              in
--               maybeBind sx_3 mSqrt
-- eval3 100 -> Just ...
-- eval3 (-100) -> Nothing
-- eval3 4 -> Nothing
--
-- короткий синтаксис
-- eval3 x = do
--             sx <- mSqrt x
--             sx_3 <- sx - 3
--             mSqrt sx_3
-- ответы те же
--
-- Давайте остановимся на одной вспомогательной функции
-- maybeBind :: Maybe a -> (a -> Maybe b) -> Maybe b
-- Функции даётся значение, очередной шаг вычисления, она возвращает новое значение.
-- maybeBind :: Maybe a -> (a -> Maybe b) -> Maybe b
-- maybeBind Nothing _ = Nothing
-- maybeBind (Just x) f = f x
--
-- взятие корня
-- mSqrt :: Double -> Maybe Double
-- mSqrt x | x < 0 = Nothing
--         | otherwise = Just $ sqrt x
-- sqrt(sqrt(x) - 3)
--
-- eval4 :: Double -> Maybe Double
-- maybeBind (Just x) mSqrt - это sqrt(x) внутри maybe
-- eval4 x = maybeBind (mSqrt x) (\a -> mSqrt (a - 3))
--
-- Эта программа похожа на последловательность шагов:
-- 1. mSqrt
-- 2. \a -> mSqrt (a - 3)
--
-- Посчитаем теперь sqrt x + sqrt y + 1
-- Шаги:
-- 1. sqrt x
-- 2. sqrt y
-- 3. sqrt x + sqrt y + 1
--
-- eval5 :: Double -> Double -> Maybe Double
-- eval5 x y = maybeBind (mSqrt x) (\a -> maybeBind (mSqrt y) (\b -> Just (a + b + 1)))
--
-- maybeBind - это функция '>>=', которая существует для Maybe и других типов.
-- eval5 :: Double -> Double -> Maybe Double
-- eval5 x y = (mSqrt x) >>= (\a -> (mSqrt y) >>= (\b -> Just (a + b + 1)))
--
-- для подообной записи есть do-нотация
-- eval5 :: Double -> Double -> Maybe Double
-- eval5 x y = do
--              a <- mSqrt x
--              b <- mSqrt y
--              return (a + b + 1)
-- Получился "императивный" код. return для Maybe это Just.
-- Перепишем теперь eval4
-- eval4 :: Double -> Maybe Double
-- eval4 x = (mSqrt x) >>= (\a -> mSqrt (a - 3))
-- или в do-нотации
-- eval4 x = do
--             a <- mSqrt x
--             mSqrt (a - 3) -- сразу возвращает Maybe
--
-- Комплексное число 1 :+ 2 это 1 + 2i
-- import Data.Complex
-- lSqrt :: Complex Double -> [Complex Double]
-- lSqrt x = [sqrt x, -sqrt x]
--
-- eval5 :: Complex Double -> Complex Double -> [Complex Double]
-- eval5 x y = do
--              a <- lSqrt x
--              b <- lSqrt y
--              return (a + b + 1)
--
-- Монада в хаскеле:
-- class Applicative a => Monad (m a) where
-- (>>=) :: m a -> (a -> m b) -> m b
-- return :: a -> m a -- превращает обычное значение в значение в монаде
-- (>>) :: m a -> m b -> m b
-- (>>) x y = x >>= (\_ -> y)
--
-- Пример монады. Вычисления с журналированием действий.
data Log a =
  Log a [String]
  deriving (Show) -- храним значение и список сообщений

-- Будем пользоваться функциями, которые при вычислениях могут выдать несколько сообщений:
add1 :: Int -> Log Int
add1 x = Log (x + 1) ["added 1"]

--
mul2 :: Int -> Log Int
mul2 x = Log (2 * x) ["multiplied by 2", "that was hard"]

--
-- Скажем, что Log это монада и объясним, как комбинировать вычисления.
instance Functor Log where
  fmap f (Log v msg) = Log (f v) msg

--
instance Applicative Log where
  pure x = Log x []
  Log f newMsg <*> Log v msg = Log (f v) (msg ++ newMsg)

--
instance Monad Log where
-- (>>=) :: Log a -> (a -> Log b) -> Log b
  (Log x messages) >>= f =
    let (Log y newMessages) = f x
      in Log y (messages ++ newMessages)

  -- return :: a -> Log a
  return x = Log x []

--
-- 2 + x + 1 через функции add1 и mul2
eval6 :: Int -> Log Int
eval6 x = do
  a <- mul2 x
  add1 a -- или b <- add1 a; return b
-- eval6 10 -- должен вернуть 21 и журнал ["умножила на 2", "добавила 1"]

--sumIO :: IO ()
--sumIO = do
--  a <- getLine
--  b <- getLine
--  print (read a + read b)
sumIO :: IO Int
sumIO = do
    a <- getLine
    b <- getLine
    return (read a + read b)


import Framework
import Data.Complex
import Numeric.LinearAlgebra

-- wave: array in either time or frequency domain
-- type Wave = [Double] -- might change to int or unsigned int, still dont know where i will get data from and how it will be formatted
type Wave = Vector R -- Double Vector

-- sample: the target wave and the components used to try and reconstruct it
data Sample = Sample {target::Wave, components::Matrix R}

type Result = Vector R

-- compsToWave :: Result -> [Wave] -> Wave
-- compsToWave ms cs =
--   do
--     let compsList = zipWith multComp ms cs
--     addComps (replicate (length compsList) 0) compsList

-- multComp :: Double -> Wave -> Wave
-- multComp k w = map ((*) k) w

-- addComps :: Wave -> [Wave] -> Wave
-- addComps base [] = base
-- addComps base (w:ws) = addComps (zipWith (+) base w) ws

compsToWave :: Result -> Matrix R -> Wave
compsToWave weights components = components #> weights

test :: (Num a) => a -> a -> a
test q w = q + w

test2 :: (Num a) => a -> a -> a
test2 q w = q * w

main =
  do
    putStrLn (show (evalAlgs [test, test2] [1,2] [[1,4],[2,5],[3,6]]))

-- result: target::wave, predicted::wave or maybe just multiplicities of the components. but what to do about phase?
-- algorithm: input::instance -> predicted::wave/result
-- errcalc: input:result || predicted::wave, target::wave -> some insightful output, IO String maybe?


{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE RecordWildCards #-}
module Q.Options.ImpliedVol.InterpolatingSmile where

import Q.Types
import Q.SortedVector (SortedVector)
import Numeric.LinearAlgebra (Vector)
import Q.Options.ImpliedVol.TimeSlice
import Q.Options.ImpliedVol.StrikeInterpolation
import Q.Interpolation
data InterpolatingSmile = StrikeSmile
  {
    smileForward :: Forward
  , smileTenor   :: YearFrac
  , smileStrikes :: SortedVector Strike
  , smileVols    :: Vector Vol
  , smileInterpolation :: StrikeInterpolation
  , smileExtrapolation :: StrikeExtrapolation
  , smileMinStrike :: Strike
  , smileMaxStrike :: Strike
  }

instance TimeSlice InterpolatingSmile Strike where
  totalVar smile@StrikeSmile{..} k = TotalVar $ scale smileTenor (sigma * sigma) where
    (Vol sigma) = impliedVol smile k

impliedVol StrikeSmile{..} = interpolateV smileInterpolation smileStrikes smileVols


{-# LANGUAGE BangPatterns #-}

{- |
Module      : Main
Description : Haskell translation of Neural Networks and Deep Learning code sample
Copyright   : (c) FUJII Satohi, 2016
License     : BSD3
Maintainer  : fujiisat@gmaill.com
Stability   : experimental
Portability : portable

My personal Haskell translation of Neural Networks and Deep Learning
code sample
(https://github.com/mnielsen/neural-networks-and-deep-learning,
Copyright (c) 2012-2015 Michael Nielsen)

-}

module Main where

import           Control.Monad         (replicateM, zipWithM)
import           Data.List             (foldl')
import qualified Data.Vector           as V
import           Numeric.LinearAlgebra (( #> ))
import qualified Numeric.LinearAlgebra as H
import           System.Environment    (getArgs)
import           System.Random         (newStdGen, randomIO)
import           System.Random.Shuffle (shuffle')
import           Text.Printf           (printf)

import           MnistData

type Vec = H.Vector H.R
type Mat = H.Matrix H.R
type Data = (Vec, Vec)
data Network = Network
               { numLayers :: !Int
               , sizes     :: V.Vector Int
               , biases    :: V.Vector Vec
               , weights   :: V.Vector Mat
               } deriving (Eq, Show)

main :: IO ()
main = do
  -- ([30, 10, 3.0], [784, 30, 10])
  ([epo, mbsize, e], networkArgs) <-  splitAt 3  <$> getArgs
  putStrLn $ "Epoch: " ++ epo
  putStrLn $ "Minibatch: " ++ mbsize
  putStrLn $ "Eta: " ++ e
  putStrLn $ "Network: " ++ show networkArgs
  (trainingData, validateDate, testData) <- loadDataWrapper
  network <- networkInit (map read networkArgs)
  !result <- sgd network (V.fromList trainingData) (read epo) (read mbsize) (read e) (Just $ V.fromList testData)
  -- print result
  return ()

networkInit :: [Int] -> IO Network
networkInit sizes = do
  b <- map H.flatten <$> mapM (`H.randn` 1) (tail sizes)
  w <- zipWithM H.randn sizes (tail sizes)
  return Network { numLayers = length sizes
                 , sizes     = V.fromList sizes
                 , biases    = V.fromList b
                 , weights   = V.fromList w
                 }

feedforward :: Network -> Vec -> Vec
feedforward (Network n s biases weights) a =
  V.foldl' (\ !a (w, b) -> sigmoidVec ((H.tr w #> a) + b) ) a $ V.zip weights biases

sgd :: Network
    -> V.Vector Data               -- Training data
    -> Int                         -- Epoch
    -> Int                         -- Minibatch size
    -> Double                      -- Eta
    -> Maybe (V.Vector (Vec, Int)) -- Test data
    -> IO Network
sgd network trainingData 0 miniBatchSize eta testData = return network
sgd network trainingData epohcs miniBatchSize eta testData = do
  miniBatches <- divide miniBatchSize <$> shuffleList trainingData
  let network' = V.foldl' (\ !n b -> updateMiniBatch n b eta) network miniBatches
  case testData of
    Just t  -> putStrLn $ printf "Epoch %d: %d / %d" epohcs (evaluate network' t) (length t)
    Nothing -> putStrLn $ printf "Epoch %d complete" epohcs
  sgd network' trainingData (epohcs - 1) miniBatchSize eta testData

updateMiniBatch :: Network -> V.Vector Data -> Double -> Network
updateMiniBatch network miniBatch eta = network { weights=weights', biases=biases'}
  where
    f :: (V.Vector Vec, V.Vector Mat) -> Data -> (V.Vector Vec, V.Vector Mat)
    f nabla mb = nablaNew
      where
        !deltaNabla = backprop network mb
        !nablaNew = ( V.zipWith (+) (fst nabla) (fst deltaNabla)
                    , V.zipWith (+) (snd nabla) (snd deltaNabla))

    (!nablaB, !nablaW) = V.foldl' f zeros miniBatch

    !zeros = ( V.map (\a -> a - a) $ biases network
             , V.map (\a -> a - a) $ weights network )

    !etaN = eta / fromIntegral (length miniBatch)
    !weights' = V.map (\(w, nw) -> w - H.cmap (etaN *) nw) $ V.zip (weights network) nablaW
    !biases'  = V.map (\(b, nb) -> b - H.cmap (etaN *) nb) $ V.zip (biases network) nablaB

backprop :: Network -> Data -> (V.Vector Vec, V.Vector Mat)
backprop network (x, y) = (V.fromList nablaB, V.fromList nablaW)
  where
    (ac0:ac1:activations, z:zs) = V.foldl' forward ([x], []) $ V.zip (biases network) (weights network)
--    !delta = costDerivative ac0 y * sigmoidPrimeVec z -- quadratic
    !delta = costDerivative ac0 y  -- cross entropy
    !nablaIni = ([delta], [H.tr (delta `H.outer` ac1)], delta)
    (nablaB, nablaW, _) = foldl' backward nablaIni
      $ zip3 zs (V.toList $ V.reverse $ weights network) activations

backward :: ([Vec], [Mat], Vec) -> (Vec, Mat, Vec) -> ([Vec], [Mat], Vec)
backward  (nabB, nabW, del) (z, w, ac) = (delta : nabB, H.tr (delta `H.outer` ac) : nabW, delta)
  where
    !sp = sigmoidPrimeVec z
    !delta = (w #> del) * sp

forward :: ([Vec], [Vec]) -> (Vec, Mat) -> ([Vec], [Vec])
forward (a@(ac:acs), zzs) (b, w) = (activation:a, z:zzs)
  where
    !z = (H.tr w #> ac) + b
    !activation = sigmoidVec z

costDerivative :: Vec -> Vec -> Vec
costDerivative outputActivations y = outputActivations - y

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

sigmoidVec :: Vec -> Vec
sigmoidVec = H.cmap sigmoid

sigmoidPrime :: Double -> Double
sigmoidPrime z = sigmoid z * (1 - sigmoid z)

sigmoidPrimeVec :: Vec -> Vec
sigmoidPrimeVec = H.cmap sigmoidPrime

evaluate :: Network -> V.Vector (Vec, Int) -> Int
evaluate network testData = length $ V.filter (uncurry (==)) testResults
  where
    testResults = V.map (\(x, y) -> (H.maxIndex (feedforward network x), y)) testData

divide :: Int -> V.Vector a -> V.Vector (V.Vector a)
divide n xs | V.null xs = V.empty
            | otherwise =  V.cons hd (divide n tl)
  where
    (hd, tl) = V.splitAt n xs

shuffleList :: V.Vector a -> IO (V.Vector a)
shuffleList xs = do
  let len = length xs
  gen <- newStdGen
  return $ V.fromList $ shuffle' (V.toList xs) len gen


-- | Vector is an alias to a DataFrame with order 1.
module Numeric.Vector
    (  -- * Type aliases
      Vector
    , Vec2f, Vec3f, Vec4f, Vec2d, Vec3d, Vec4d
    , Vec2i, Vec3i, Vec4i, Vec2w, Vec3w, Vec4w
      -- * Vector constructors
    , Vector2 (..), Vector3 (..), Vector4 (..)
    , DataFrame(Vec2, Vec3, Vec4)
      -- * Common operations
    , (.*.), dot, (·)
    , normL1, normL2, normLPInf, normLNInf, normLP
    , normalized
    , det2, cross, (×)
    ) where

import Numeric.DataFrame.Internal.Backend ()
import Numeric.Vector.Internal


{-# LANGUAGE ApplicativeDo                        #-}
{-# LANGUAGE FlexibleContexts                     #-}
{-# LANGUAGE PartialTypeSignatures                #-}
{-# LANGUAGE RankNTypes                           #-}
{-# LANGUAGE ScopedTypeVariables                  #-}
{-# LANGUAGE TypeApplications                     #-}
{-# LANGUAGE TypeFamilies                         #-}
{-# OPTIONS_GHC -fno-warn-partial-type-signatures #-}

module Backprop.Learn.Test (
  -- * Tests
    Test
  , maxIxTest, rmseTest
  , squaredErrorTest, absErrorTest, totalSquaredErrorTest, squaredErrorTestV
  , crossEntropyTest
  -- ** Manipulate tests
  , lossTest, lmapTest
  -- * Run tests
  , testLearn, testLearnStoch, testLearnAll, testLearnStochAll
  -- ** Correlation tests
  , testLearnCov, testLearnCorr
  , testLearnStochCov, testLearnStochCorr
  ) where

import           Backprop.Learn.Loss
import           Backprop.Learn.Model
import           Control.Monad.Primitive
import           Data.Bifunctor
import           Data.Bitraversable
import           Data.Function
import           Data.Profunctor
import           Data.Proxy
import           GHC.TypeNats
import           Numeric.Backprop
import qualified Control.Foldl                as L
import qualified Numeric.LinearAlgebra        as HU
import qualified Numeric.LinearAlgebra.Static as H
import qualified System.Random.MWC            as MWC

-- TODO: support non-double results?

type Test o = o -> o -> Double

-- | Create a 'Test' from a 'Loss'
lossTest :: Loss a -> Test a
lossTest l x = evalBP (l x)

maxIxTest :: KnownNat n => Test (H.R n)
maxIxTest x y
    | match x y = 1
    | otherwise = 0
  where
    match = (==) `on` (HU.maxIndex . H.extract)

rmseTest :: forall n. KnownNat n => Test (H.R n)
rmseTest x y = H.norm_2 (x - y) / sqrt (fromIntegral (natVal (Proxy @n)))

squaredErrorTest :: Real a => Test a
squaredErrorTest x y = e * e
  where
    e = realToFrac (x - y)

absErrorTest :: Real a => Test a
absErrorTest x y = realToFrac . abs $ x - y

totalSquaredErrorTest :: (Applicative t, Foldable t, Real a) => Test (t a)
totalSquaredErrorTest x y = realToFrac (sum e)
  where
    e = do
      x' <- x
      y' <- y
      pure ((x' - y') ^ (2 :: Int))

squaredErrorTestV :: KnownNat n => Test (H.R n)
squaredErrorTestV x y = e `H.dot` e
  where
    e = x - y

crossEntropyTest :: KnownNat n => Test (H.R n)
crossEntropyTest targ res = -(log res H.<.> targ)

lmapTest
    :: (a -> b)
    -> Test b
    -> Test a
lmapTest f t x y = t (f x) (f y)

testLearn
    :: (Learn a b l, NoState l)
    => Test b
    -> l
    -> LParam_ I l
    -> a
    -> b
    -> Double
testLearn t l mp x y = t y $ runLearnStateless_ l mp x

testLearnStoch
    :: (Learn a b l, NoState l, PrimMonad m)
    => Test b
    -> l
    -> MWC.Gen (PrimState m)
    -> LParam_ I l
    -> a
    -> b
    -> m Double
testLearnStoch t l g mp x y = t y <$> runLearnStochStateless_ l g mp x

cov :: Fractional a => L.Fold (a, a) a
cov = do
    x  <- lmap fst           L.sum
    y  <- lmap snd           L.sum
    xy <- lmap (uncurry (*)) L.sum
    n  <- fromIntegral <$> L.length
    pure (xy / n - (x * y) / n / n)

corr :: Floating a => L.Fold (a, a) a
corr = do
    x  <- lmap fst           L.sum
    x2 <- lmap ((**2) . fst) L.sum
    y  <- lmap snd           L.sum
    y2 <- lmap ((**2) . snd) L.sum
    xy <- lmap (uncurry (*)) L.sum
    n  <- fromIntegral <$> L.length
    pure $ (xy / n - (x * y) / n / n)
         / sqrt ( x2 / n - (x / n)**2 )
         / sqrt ( y2 / n - (y / n)**2 )

testLearnCov
    :: (Learn a b l, NoState l, Foldable t, Fractional b)
    => l
    -> LParam_ I l
    -> t (a, b)
    -> b
testLearnCov l p = L.fold $ (lmap . first) (runLearnStateless_ l p) cov

testLearnCorr
    :: (Learn a b l, NoState l, Foldable t, Floating b)
    => l
    -> LParam_ I l
    -> t (a, b)
    -> b
testLearnCorr l p = L.fold $ (lmap . first) (runLearnStateless_ l p) corr

testLearnAll
    :: (Learn a b l, NoState l, Foldable t)
    => Test b
    -> l
    -> LParam_ I l
    -> t (a, b)
    -> Double
testLearnAll t l p = L.fold $ lmap (uncurry (testLearn t l p)) L.mean

-- newtype M m a = M { getM :: m a }
-- instance (Semigroup a, Applicative m) => Semigroup (M m a) where
--     M x <> M y = M $ liftA2 (<>) x y
-- instance (Monoid a, Applicative m) => Monoid (M m a) where
--     mappend = (<>)
--     mempty = M (pure mempty)

testLearnStochAll
    :: (Learn a b l, NoState l, PrimMonad m, Foldable t)
    => Test b
    -> l
    -> MWC.Gen (PrimState m)
    -> LParam_ I l
    -> t (a, b)
    -> m Double
testLearnStochAll t l g p = L.foldM $ L.premapM (uncurry (testLearnStoch t l g p))
                                      (L.generalize L.mean)

testLearnStochCov
    :: (Learn a b l, NoState l, PrimMonad m, Foldable t, Fractional b)
    => l
    -> MWC.Gen (PrimState m)
    -> LParam_ I l
    -> t (a, b)
    -> m b
testLearnStochCov l g p = L.foldM $ (L.premapM . flip bitraverse pure)
                                      (runLearnStochStateless_ l g p)
                                      (L.generalize cov)

testLearnStochCorr
    :: (Learn a b l, NoState l, PrimMonad m, Foldable t, Floating b)
    => l
    -> MWC.Gen (PrimState m)
    -> LParam_ I l
    -> t (a, b)
    -> m b
testLearnStochCorr l g p = L.foldM $ (L.premapM . flip bitraverse pure)
                                       (runLearnStochStateless_ l g p)
                                       (L.generalize corr)


module Main where

import Quadratic
import Data.Complex

main :: IO ()
main = do
  putStrLn $ show $ roots (Quadratic 0 1 2)
  putStrLn $ show $ roots (Quadratic 1 3 4)
  putStrLn $ show $ roots (Quadratic 1 3 4)
  putStrLn $ show $ roots (Quadratic 1 4 4)
  putStrLn $ show $ roots (Quadratic 1 0 4)


-- By Christian Oliveros and Minmin Chen
module PingPong.Player.Stig.Kinematics where

import Control.Lens (view)
import Data.Bool (bool)
import Data.Geometry hiding (Vector, dot, head)
import Data.Maybe
import GHC.Float
import Numeric.LinearAlgebra hiding (Element)
import PingPong.Model
import PingPong.Player.Stig.General
import PingPong.Player.Stig.GeometryHelpers
import Prelude hiding ((<>))

-- Inverse and Forward Kinematics

-- | Calculate the forward kinematic matrix for a joint / link
calFwdMatrix :: Element -> Element -> [[Float]]
calFwdMatrix (Joint _ a) (Link _ l) =
  [ [cos a, - sin a, l * cos a],
    [sin a, cos a, l * sin a],
    [0, 0, 1]
  ]
calFwdMatrix _ _ = [[1, 0, 0], [0, 1, 0], [0, 0, 1]]

data ArmKinematicPart
  = ArmBase
      { footParam :: R
      }
  | ArmJoint
      { translationParam :: R
      }
  | ArmBat
      { translationParam :: R
      }
  deriving (Show)

type ArmKinematic = [ArmKinematicPart]

-- | Goes from rough representation to a more refined one
fromFootArmToArmKinematic :: Float -> Arm -> ArmKinematic
fromFootArmToArmKinematic foot arm = ArmBase (float2Double foot) : processArm 0 arm
  where
    processArm :: R -> Arm -> ArmKinematic
    processArm t [Link _ t2] = [ArmBat $ t + float2Double t2]
    processArm t (Link _ t2 : as) = processArm (t + float2Double t2) as
    processArm t (Joint _ _ : as) = ArmJoint t : processArm 0 as

-- | Goes from rough representation to a more refined one.
-- As well, get the motion vector
getArmKinematicAndMotion :: Float -> Arm -> (ArmKinematic, Motion)
getArmKinematicAndMotion foot arm =
  (fromFootArmToArmKinematic foot simpleArm, getCurrentJoints simpleArm)
  where
    simpleArm = simplifyArm arm

-- | Converts a motion list to a vector of joint coordinates
motionToJointVector :: Motion -> Vector R
motionToJointVector m = fromList $ map float2Double m

-- | Converts vector of joint coordinates to a motion list
jointVectorToMotion :: Vector R -> Motion
jointVectorToMotion v = map double2Float $ toList v

-- | Translation in the X axis matrix
translateXTrans :: R -> Matrix R
translateXTrans t =
  fromLists [[1, 0, t], [0, 1, 0], [0, 0, 1]] :: Matrix R

-- | Translation in the X axis inverse matrix
translateXInvTrans :: R -> Matrix R
translateXInvTrans t =
  fromLists [[1, 0, - t], [0, 1, 0], [0, 0, 1]] :: Matrix R

-- | Rotation matrix
rotateTrans :: R -> Matrix R
rotateTrans a =
  fromLists [[c, - s, 0], [s, c, 0], [0, 0, 1]] :: Matrix R
  where
    c = cos a
    s = sin a

-- | Rotation inverse matrix
rotateInvTrans :: R -> Matrix R
rotateInvTrans a =
  fromLists [[c, s, 0], [- s, c, 0], [0, 0, 1]] :: Matrix R
  where
    c = cos a
    s = sin a

-- | Create an homogeneous 2D point
homogeneousPoint :: R -> R -> Vector R
homogeneousPoint x y = fromList [x, y, 1]

-- | Zero Point in Homogenous Coordinates
homogeneousZero :: Vector R
homogeneousZero = homogeneousPoint 0 0

-- | 2d Point to Homogeneous Coordinates
pointToHomogenousPoint :: Point 2 Float -> Vector R
pointToHomogenousPoint p = homogeneousPoint (float2Double $ view xCoord p) (float2Double $ view yCoord p)

-- | Create an homogeneous 2D vector
homogeneousVector :: R -> R -> Vector R
homogeneousVector x y = fromList [x, y, 0]

-- | X Vector
homogeneousVectorX :: Vector R
homogeneousVectorX = homogeneousVector 1 0

-- | Y Vector
homogeneousVectorY :: Vector R
homogeneousVectorY = homogeneousVector 0 1

-- | Homogeneous 2D Identity Matrix
homogeneousIdent :: Matrix R
homogeneousIdent = ident 3

data ArmKinematicPartMatrix
  = ArmBaseMatrix
      { footMatrix :: Matrix R
      }
  | ArmJointMatrix
      { translationMatrix :: Matrix R,
        jointMatrix :: Matrix R
      }
  | ArmBatMatrix
      { batMatrix :: Matrix R
      }
  deriving (Show)

type ArmKinematicMatrix = [ArmKinematicPartMatrix]

-- | Get Transform from Part Matrix
getTrans :: ArmKinematicPartMatrix -> Matrix R
getTrans (ArmBaseMatrix t) = t
getTrans (ArmJointMatrix t r) = t <> r
getTrans (ArmBatMatrix t) = t

-- | Get Inverse Transform from Part Matrix
getTransInv :: ArmKinematicPartMatrix -> Matrix R
getTransInv (ArmBaseMatrix t) = t
getTransInv (ArmJointMatrix t r) = r <> t
getTransInv (ArmBatMatrix t) = t

-- | Calculate the transforms for forward kinematics
applyForwardKinematicTrans ::
  ArmKinematic -> Vector R -> ArmKinematicMatrix
applyForwardKinematicTrans arm v = toTrans arm motion
  where
    motion = toList v :: [R]

    toTrans :: ArmKinematic -> [R] -> ArmKinematicMatrix
    toTrans (ArmBase t : as) js = (ArmBaseMatrix $ translateXTrans t <> rotateTrans (pi / 2)) : toTrans as js
    toTrans (ArmJoint t : as) (j : js) = ArmJointMatrix (translateXTrans t) (rotateTrans j) : toTrans as js
    toTrans [ArmBat t] [] = [ArmBatMatrix $ translateXTrans t]

-- | Calculate the inverse transforms for forward kinematics
applyForwardKinematicTransInv ::
  ArmKinematic -> Vector R -> ArmKinematicMatrix
applyForwardKinematicTransInv arm v = reverse $ toTrans arm motion
  where
    motion = toList v :: [R]

    toTrans :: ArmKinematic -> [R] -> ArmKinematicMatrix
    toTrans (ArmBase t : as) js = (ArmBaseMatrix $ rotateInvTrans (pi / 2) <> translateXInvTrans t) : toTrans as js
    toTrans (ArmJoint t : as) (j : js) = ArmJointMatrix (translateXInvTrans t) (rotateInvTrans j) : toTrans as js
    toTrans [ArmBat t] [] = [ArmBatMatrix $ translateXInvTrans t]

-- | Apply Forward Kinematic Transformations
applyForwardKinematicMatrixTrans :: ArmKinematicMatrix -> Matrix R
applyForwardKinematicMatrixTrans = foldr ((<>) . getTrans) homogeneousIdent

-- | Apply Forward Kinematic Inverse Transformations (the list must go from end effecto to base)
applyForwardKinematicMatrixTransInv :: ArmKinematicMatrix -> Matrix R
applyForwardKinematicMatrixTransInv = foldr ((<>) . getTransInv) homogeneousIdent

-- | Compress the transforms of a Forward Kinematic Arm to only the joints and the end effector
compressForwardKinematicsJointsTrans :: ArmKinematicMatrix -> ArmKinematicMatrix
compressForwardKinematicsJointsTrans a = compress a homogeneousIdent
  where
    compress :: ArmKinematicMatrix -> Matrix R -> ArmKinematicMatrix
    compress [ArmBatMatrix t] m = [ArmBatMatrix (m <> t)]
    compress (ArmJointMatrix t r : as) m = ArmJointMatrix (m <> t) r : compress as homogeneousIdent
    compress (a : as) m = compress as (m <> getTrans a)

-- | Compress the transforms of a Forward Kinematic Arm to only the joints and the end base
compressForwardKinematicsJointsInvTrans :: ArmKinematicMatrix -> ArmKinematicMatrix
compressForwardKinematicsJointsInvTrans a = compress a homogeneousIdent
  where
    compress :: ArmKinematicMatrix -> Matrix R -> ArmKinematicMatrix
    compress [ArmBaseMatrix t] m = [ArmBaseMatrix (t <> m)]
    compress (ArmJointMatrix t r : as) m = ArmJointMatrix (t <> m) r : compress as homogeneousIdent
    compress (a : as) m = compress as (getTransInv a <> m)

-- | Calculate the 2D Homogeneous Jacobian of an arm [J^T|0]^T .
-- fwdMatTrans : Kinematic Forward Transforms (From base to end effector).
-- fwdMatTransInv : Kinematic Forward Inverse Transforms (From end effector to base).
-- xLocal : Position of Tool on End Effector Local Coordinates
calculateJacobian ::
  ArmKinematicMatrix -> ArmKinematicMatrix -> Vector R -> Matrix R
calculateJacobian fwdMatTrans fwdMatTransInv xLocal = jH
  where
    -- Compressed form guarantees that the links were applied to the joints
    fwdTransCompressed = compressForwardKinematicsJointsTrans fwdMatTrans
    fwdTransInvCompressed = compressForwardKinematicsJointsInvTrans fwdMatTransInv

    -- Derivative of a revolute Joint
    revoluteDeriv = fromLists [[0, -1, 0], [1, 0, 0], [0, 0, 0]] :: Matrix R

    rollingTransInv :: ArmKinematicMatrix -> Matrix R -> [Matrix R]
    rollingTransInv [ArmBaseMatrix t] m = [t <> m]
    rollingTransInv (a : as) m = roll : rollingTransInv as roll
      where
        roll = getTransInv a <> m
    calculate ::
      ArmKinematicMatrix -> Matrix R -> [Matrix R] -> [Vector R]
    calculate [ArmBatMatrix _] _ [_] = [] -- We reach the bat, we no longer need to process
    calculate (j@(ArmJointMatrix _ _) : as) m (endToJ : mIs) =
      (baseToJ <> revoluteDeriv <> endToJ) #> xLocal : calculate as baseToJ mIs
      where
        baseToJ = m <> getTrans j

    -- Jacobian in this form: [J^T|0]^T there is a row of 0s at the end due to 2d homogeneous coordinates
    jH =
      fromColumns $
        calculate fwdTransCompressed homogeneousIdent (rollingTransInv fwdTransInvCompressed homogeneousIdent)

-- | Get the Jacobian from the Homogeneous Jacobian (Note this doesn't work for the inverse of the Jacobian)
getJacobianFromHomogeneousJacobian :: Matrix R -> Matrix R
getJacobianFromHomogeneousJacobian jH = jH ?? (DropLast 1, All)

-- | Newton Raphson Threshold
newtonRaphsonThreshold :: (Num r, Ord r, Fractional r) => r -> r -> r
newtonRaphsonThreshold = threshold 0.0000001

-- | Calculates a Newton-Raphson IK step (if singular it fails). Returns the new Joints and its error agains target
-- a : Arm Description.
-- q : Current Joints.
-- xLocal : Tool in End-Effector Local Coordinates
-- xTargetGlobal : Target in Global Coordinates
newtonRaphsonStep ::
  ArmKinematic ->
  Vector R ->
  Vector R ->
  Vector R ->
  Maybe (Vector R, R)
newtonRaphsonStep a q xLocal xTargetGlobal
  | singular = Nothing
  | otherwise = Just (qn, norm_2 eN)
  where
    -- Forward Transforms
    fwdT = applyForwardKinematicTrans a q
    fwdTI = applyForwardKinematicTransInv a q
    jH = calculateJacobian fwdT fwdTI xLocal

    -- Bat Global Position
    batGlobal = applyForwardKinematicMatrixTrans fwdT #> xLocal

    -- Error
    e = xTargetGlobal - batGlobal

    -- Drop Homogeneous part
    er = subVector 0 2 e
    j = getJacobianFromHomogeneousJacobian jH

    -- Calculate delta q
    dq = - pinv j #> er

    dqNorm = norm_2 dq

    qnUnormalized = q + dq
    -- Re normalizes angles for more accuracy
    qn = fromList $ map normalizeAngle $ toList qnUnormalized

    -- New Bat Position
    batN = applyForwardKinematicMatrixTrans (applyForwardKinematicTrans a qn) #> xLocal

    -- New Error
    eN = xTargetGlobal - batN

    -- If the move was singular
    singular = rank jH < 2 && newtonRaphsonThreshold 0 dqNorm == 0

-- | Maximum Newton Raphson Step
newtonRaphsonMaxStep :: Int
newtonRaphsonMaxStep = 10000

-- | Really Bad Step
newtonRaphsonBadStep :: forall a. Floating a => a
newtonRaphsonBadStep = 2

-- | When to do a random Restart Newton Raphson Step
newtonRaphsonMaxRandomRestartStep :: Int
newtonRaphsonMaxRandomRestartStep = round (fromIntegral newtonRaphsonMaxStep / 10)

-- | NewtonRaphson Loop Iteration
newtonRaphsonIKIter ::
  Int ->
  Int ->
  ArmKinematic ->
  (Vector R, Vector R) ->
  Vector R ->
  (Vector R, R) ->
  (Vector R, R)
newtonRaphsonIKIter i j a (xLocal, xTargetGlobal) q (qBest, eBest)
  | i >= newtonRaphsonMaxStep = (qBest, eBest)
  | newtonRaphsonThreshold 0 eBest == 0 = (qBest, eBest)
  | singular = newtonRaphsonIKIter (i + 1) 0 a (xLocal, xTargetGlobal) qR (qBest, eBest)
  | needReset = newtonRaphsonIKIter (i + 1) 0 a (xLocal, xTargetGlobal) qR (qBest, eBest)
  | eN < eBest = newtonRaphsonIKIter (i + 1) (j + 1) a (xLocal, xTargetGlobal) qN (qN, eN)
  | otherwise = newtonRaphsonIKIter (i + 1) (j + 1) a (xLocal, xTargetGlobal) qN (qBest, eBest)
  where
    -- Random Vector
    qSize = size q
    pseudoInt = round $ fromIntegral i ** (dot q qBest * bool (eBest + 1) (1 / eBest) (eBest < 1 && 0 < eBest))
    qR = pi / 2 * (2 * randomVector pseudoInt Uniform qSize - 1)

    -- Perform the step
    step = newtonRaphsonStep a q xLocal xTargetGlobal

    singular = case step of
      Nothing -> True
      Just _ -> False

    (qN, eN) = fromJust step

    needReset =
      ( eN >= ((eBest / (1 + fromIntegral j)) * newtonRaphsonBadStep)
          || (newtonRaphsonThreshold 0 (norm_2 (qN - q)) == 0)
      )
        && j >= newtonRaphsonMaxRandomRestartStep

-- | Newton Raphson IK Algorithm
newtonRaphsonIK ::
  ArmKinematic ->
  (Vector R, Vector R) ->
  Vector R ->
  (Vector R, R)
newtonRaphsonIK a (xLocal, xTargetGlobal) q = newtonRaphsonIKIter 0 0 a (xLocal, xTargetGlobal) q (q, eNorm)
  where
    -- Forward Transforms
    fwdT = applyForwardKinematicTrans a q

    -- Bat Global Position
    batGlobal = applyForwardKinematicMatrixTrans fwdT #> xLocal

    -- Error
    e = xTargetGlobal - batGlobal
    eNorm = norm_2 e

-- Calculates newton raphson step to approximate acos
newtonRaphsonAcosStep :: Int -> Int -> Double -> (Double, Double) -> Double -> (Double, Double)
newtonRaphsonAcosStep i j q (qB, eB) t
  | i == newtonRaphsonMaxStep = (qB, eB)
  | newtonRaphsonThreshold 0 eB == 0 = (qB, eB)
  | isSingular || needReset = newtonRaphsonAcosStep (i + 1) 0 qR (qB, eB) t
  | eN < eB = newtonRaphsonAcosStep (i + 1) (j + 1) qN (qN, eN) t
  | otherwise = newtonRaphsonAcosStep (i + 1) (j + 1) qN (qB, eB) t
  where
    -- Calculate derivative and function
    f = t - cos q
    fDeriv = sin q

    -- Check if we are in a singularity
    isSingular = newtonRaphsonThreshold 0 fDeriv == 0

    -- Step
    qN = normalizeAngle $ q - f / fDeriv

    -- New Error
    eN = t - cos qN

    -- Random Restarts
    needReset = eN >= eB * newtonRaphsonBadStep && j >= newtonRaphsonMaxRandomRestartStep

    pseudoInt = round $ fromIntegral i ** (q * qB * bool (eB + 1) (1 / eB) (eB < 1 && 0 < eB))
    randNum = atIndex (randomVector pseudoInt Uniform 1) 0
    qR = pi / 2 * (2 * randNum - 1)

-- | Approximate Acos with newton raphson
newtonRaphsonAcos :: Double -> Double -> (Double, Double)
newtonRaphsonAcos q t = newtonRaphsonAcosStep 0 0 qN (qN, eN) t
  where
    qN = normalizeAngle q
    eN = t - cos qN


{-|
Module: Science.QuantumChemistry.NumericalTools.EigenValues
Description: Compute all the eigenvalues of the real symmetrical matrix
Copyright: @2016 Felipe Zapata
-}

module Science.QuantumChemistry.NumericalTools.EigenValues (
                   eigenSolve
                   ) where

-- =============================> Standard and third party libraries <===============================
import Control.Arrow ((***))
import Data.Array.Repa as R
import Data.Vector.Storable (convert)
import qualified Data.Vector.Unboxed as U
import qualified Numeric.LinearAlgebra.Data as ND
import Numeric.LinearAlgebra.HMatrix (eigSH')

-- =================> Internal Modules <======================
import Science.QuantumChemistry.GlobalTypes (VecUnbox)
import Science.QuantumChemistry.NumericalTools.JacobiMethod (jacobiP)
import Science.QuantumChemistry.NumericalTools.VectorTools (sortEigenData)

-- | Compute all the eigenvalues and eigenvectors of symmetrical matrix.
eigenSolve :: Array U DIM2 Double -> (VecUnbox, Array U DIM2 Double)
eigenSolve arr = (vs, fromUnboxed (extent arr) mtx )
  where (vs,mtx) = sortEigenData $ hv2VecUnbox *** hm2VecUnbox $ eigSH' $ repa2HM arr

vecUnbox2HV :: VecUnbox -> ND.Vector Double
vecUnbox2HV = convert

hv2VecUnbox :: ND.Vector Double -> VecUnbox
hv2VecUnbox = convert

repa2HM :: Array U DIM2 Double -> ND.Matrix Double
repa2HM arr = ND.reshape dim . convert . toUnboxed $ arr

 where (Z :. dim :. _) = extent arr

hm2Repa :: ND.Matrix Double -> Array U DIM2 Double
hm2Repa mtx = fromUnboxed (ix2 dim dim) . convert . ND.flatten $ mtx
 where dim = ND.rows mtx

hm2VecUnbox :: ND.Matrix Double -> VecUnbox
hm2VecUnbox = convert . ND.flatten


module OrientationAnalysis where

import           Control.Monad                as M
import           Data.Array.Repa              as R
import           Data.Binary                  (decodeFile)
import           Data.Complex
import           Data.List                    as L
import           DFT.Plan
import           Filter.Pinwheel
import           FokkerPlanck.MonteCarlo
import           FokkerPlanck.Pinwheel
import           Image.IO
import           STC.CompletionField          (makeImagePlan, makeR2Z1T0Plan)
import           STC.CompletionField
import           STC.OrientationScaleAnalysis
import           System.Directory
import           System.Environment
import           System.FilePath

main = do
  args@(numPointStr:numOrientationStr:sigmaStr:taoStr:lenStr:initStr:numTrailStr:maxTrailStr:theta0FreqsStr:thetaFreqsStr:histFilePath:alphaStr:pinwheelFlagStr:imagePath:idxStr:numOrientationSampleStr:numThreadStr:_) <-
    getArgs
  print args
  let numPoint = read numPointStr :: Int
      numOrientation = read numOrientationStr :: Int
      sigma = read sigmaStr :: Double
      tao = read taoStr :: Double
      len = read lenStr :: Int
      init = read initStr :: (Double, Double, Double, Double, Double, Double)
      numTrail = read numTrailStr :: Int
      maxTrail = read maxTrailStr :: Int
      theta0Freq = read theta0FreqsStr :: Double
      theta0Freqs = [-theta0Freq .. theta0Freq]
      thetaFreq = read thetaFreqsStr :: Double
      thetaFreqs = [-thetaFreq .. thetaFreq]
      alpha = read alphaStr :: Double
      pinwheelFlag = read pinwheelFlagStr :: Bool
      idx = read idxStr :: (Int, Int)
      numOrientationSample = read numOrientationSampleStr :: Int
      numThread = read numThreadStr :: Int
      pinwheelParams =
        PinwheelParams numPoint numPoint alpha (exp 1) theta0Freqs [0]
      folderPath = "output/test/OrientationAnalysis"
  createDirectoryIfMissing True folderPath
  flag <- doesFileExist histFilePath
  arrR2Z1T0 <-
    if pinwheelFlag
      then error
             "Using pinwheels to construct the Green's function has not been implemented yet."
           --computeR2Z1T0Array numPoint numPoint alpha thetaFreqs theta0Freqs
      else if flag
             then do
               getNormalizedHistogramArr <$> decodeFile histFilePath
             else solveMonteCarloR2Z1T0
                    numThread
                    numTrail
                    maxTrail
                    numPoint
                    numPoint
                    sigma
                    tao
                    len
                    theta0Freqs
                    thetaFreqs
                    ""
                    (emptyHistogram
                       [ numPoint
                       , numPoint
                       , L.length theta0Freqs
                       , L.length thetaFreqs
                       ]
                       0)
  (ImageRepa _ img) <- readImageRepa imagePath False
  plotImageRepa (folderPath </> "input.png") (ImageRepa 8 img)
  plan0 <- makeR2Z1T0Plan emptyPlan arrR2Z1T0
  (plan1, imgF) <- makeImagePlan plan0 img
  (plan, filterF, filterPIF) <- pinwheelFilter plan1 pinwheelParams
  arrR2Z1T0F <- dftR2Z1T0 plan . makeFilterR2Z1T0 $ arrR2Z1T0
  convolvedImg <-
    fmap (\x -> R.slice x (Z :. All :. (0 :: Int) :. All :. All)) $
    convolvePinwheel plan filterF (R.slice imgF (Z :. (0 :: Int) :. All :. All))
  plotMagnitudeOrientationSource
    plan
    folderPath
    numOrientationSample
    numOrientation
    thetaFreqs
    arrR2Z1T0F
    convolvedImg
    idx
  plotMagnitudeOrientation
    folderPath
    numOrientationSample
    numOrientation
    thetaFreqs
    convolvedImg
    idx


{-# LANGUAGE ExistentialQuantification,
             DataKinds #-}
module MVDistributions
where

import Prelude hiding ((<>))

import Control.Arrow (first)
import Control.Monad (replicateM, zipWithM)
import Control.Monad.Bayes.Class (MonadSample, normal, shuffle)
import Data.Aeson (ToJSON (toJSON))
import Data.List (permutations, sortBy, maximumBy)
import Data.Ord (comparing, Down (Down))
import qualified Data.Vector as V
import Data.Proxy (Proxy (Proxy))
import GHC.TypeLits (KnownNat, natVal)

import Numeric.LinearAlgebra.Static
import qualified Numeric.LinearAlgebra.Data as LAD
import Numeric.SpecFunctions (logGamma)

import Distributions (Distr, Distr' (..))

instance KnownNat n => ToJSON (R n) where
  toJSON x = toJSON (extract x :: LAD.Vector Double)

instance (KnownNat m, KnownNat n) => ToJSON (L m n) where
  toJSON x = toJSON (LAD.toLists (extract x :: LAD.Matrix Double))

instance KnownNat n => ToJSON (Sym n) where
  toJSON x = toJSON (unSym x)

eyeSym :: KnownNat n => Sym n
eyeSym = sym eye

sampleNormalDist :: forall n m. MonadSample m => KnownNat n => R n -> Sym n -> m (R n)
sampleNormalDist mu sig = do
  xs <- replicateM (fromInteger (natVal (Proxy :: Proxy n))) (normal 0 1)
  let x = vector xs
  return (mu + chol sig #> x)

log2pi :: Double
log2pi = log (2 * pi)

mvNormal_ll0 :: forall n. KnownNat n => Sym n -> R n -> Double
mvNormal_ll0 s x = -1/2 * (x <.> (sinv #> x) + lndets + fromInteger n * log2pi)
  where
  (sinv, (lndets, _)) = invlndet (unSym s)
  sinv' :: Sym n
  sinv' = sym sinv
  n = natVal (Proxy :: Proxy n)

mvNormal_ll :: KnownNat n => R n -> Sym n -> R n -> Double
mvNormal_ll mu s x = mvNormal_ll0 s (x - mu)

mvNormal :: KnownNat n => R n -> Sym n -> Distr (R n)
mvNormal mu s = Distr (sampleNormalDist mu s) (mvNormal_ll mu s)

uniformNSphere :: forall n. KnownNat n => Distr (R n)
uniformNSphere = Distr (do { x <- sampleNormalDist 0 (sym eye); pure (konst (1 / (norm_2 x)) * x) })
  (\_ -> - logNSphereVolume (natVal (Proxy :: Proxy n)))

-- https://en.wikipedia.org/wiki/N-sphere#Closed_forms
logNSphereVolume :: Integer -> Double
logNSphereVolume n = log (fromIntegral n) + fromIntegral n / 2 * log pi
  - logGamma (fromIntegral n / 2 + 1)

conjugate :: (KnownNat m, KnownNat n) => L m n -> Sym n -> Sym m
conjugate f p = sym (f <> unSym p <> tr f)

conjugate' :: (KnownNat m, KnownNat n) => L n m -> Sym n -> Sym m
conjugate' f p = sym (tr f <> unSym p <> f)

kalmanPredict :: KnownNat n => Sq n -> Sym n -> (R n, Sym n) -> (R n, Sym n)
kalmanPredict f q (x, p) = (f #> x, conjugate f p + q)

symEachSide :: KnownNat n => Sym n -> Sym n -> Sym n
symEachSide p x = sym (p' <> unSym x <> p')
  where p' = unSym p

-- The below is not ideal numerically
kalmanUpdate :: forall m n. (KnownNat m, KnownNat n) => R m -> L m n -> Sym m
  -> (R n, Sym n) -> ((R n, Sym n), Double)
kalmanUpdate z h r (x, p) =
  ((x + k #> ytilde, p - symEachSide p (conjugate' h sinv'))
   , loglike)
  where
  ytilde :: R m
  ytilde = z - h #> x
  trh :: L n m
  trh = tr h
  s :: Sym m
  s = r + conjugate h p
  (sinv, (lndets, _)) = invlndet (unSym s)
  sinv' :: Sym m
  sinv' = sym sinv
  k = unSym p <> trh <> sinv
  n = natVal (Proxy :: Proxy n)
  loglike = -1/2 * (ytilde <.> (sinv #> ytilde) + lndets + fromInteger n * log2pi)

shuffleList :: MonadSample m => [a] ->  m [a]
shuffleList xs = V.toList <$> shuffle (V.fromList xs)

printVector :: KnownNat n => R n -> IO ()
printVector x = print (extract x)

scaleSym :: KnownNat n => Double -> Sym n -> Sym n
scaleSym c m = sym $ konst c * unSym m


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

direction :: Char -> Complex Float
direction 'L' = 0 :+ 1
direction 'R' = 0 :+ (-1)

blocks :: String -> Float
blocks str =
  let [(count, _)] = reads str
  in count

parse :: String -> [(Complex Float, Float)]
parse = map (\(x:xs) -> (direction x, blocks xs)) . words

advance :: (Complex Float, Complex Float) -> (Complex Float, Float) -> (Complex Float, Complex Float)
advance (dir, pos) (turn, steps) =
  let dir' = dir * turn
      pos' = pos + dir' * (steps :+ 0)
  in (dir', pos')

part1 :: String -> Int
part1 input =
  let parts = parse input
      (_, (x :+ y)) = foldl advance (0 :+ 1, 0 :+ 0) parts
  in grid x + grid y
  where grid = round . abs

firstTwice :: (Complex Float, Complex Float) -> [(Complex Float, Float)] -> [Complex Float] -> Complex Float
firstTwice (dir, pos) (x:xs) prev =
  let (turn, steps) = x
      (dir':_, poss@(pos':_)) = unzip $ map (advance (dir, pos)) $ map (\s -> (turn, s)) $ reverse [1..steps]
  in fromMaybe (firstTwice (dir', pos') xs (poss ++ prev)) (listToMaybe $ intersect poss prev)

part2 :: String -> Int
part2 input =
  let parts = parse input
      (x :+ y) = firstTwice (0 :+ 1, 0 :+ 0) parts []
  in grid x + grid y
  where grid = round . abs


module STCR2S1Edge where

import           Control.Monad           as M
import           Control.Monad.Parallel  as MP
import           Data.Array.Repa         as R
import           Data.Binary             (decodeFile)
import           Data.Complex
import           Data.List               as L
import           Data.Vector.Unboxed     as VU
import           DFT.Plan
import           Filter.Utils
import           FokkerPlanck.MonteCarlo
import           Image.Edge
import           Image.IO
import           STC.CompletionFieldR2S1
import           STC.PointSet
import           System.Directory
import           System.Environment
import           System.FilePath
import           Text.Printf
import           Types
import           Utils.Array

main = do
  args@(numPointStr:numOrientationStr:sigmaStr:taoStr:numTrailStr:maxTrailStr:initOriStr:initSpeedStr:writeFlagStr:numIterationStr:thresholdStr:histFilePath:edgeFilePath:numThreadStr:_) <-
    getArgs
  print args
  let numPoint = read numPointStr :: Int
      numOrientation = read numOrientationStr :: Int
      sigma = read sigmaStr :: Double
      tao = read taoStr :: Double
      numTrail = read numTrailStr :: Int
      maxTrail = read maxTrailStr :: Int
      initOri = read initOriStr :: Double
      initSpeed = read initSpeedStr :: Double
      writeFlag = read writeFlagStr :: Bool
      numIteration = read numIterationStr :: Int
      threshold = read thresholdStr :: Double
      numThread = read numThreadStr :: Int
      folderPath = "output/test/STCR2S1Edge"
  createDirectoryIfMissing True folderPath
  flag <- doesFileExist histFilePath
  arrG <-
    if flag
      then getNormalizedHistogramArr .
           mapHistogram
             (\x ->
                let y = x :: Int
                 in fromIntegral y) <$>
           decodeFile histFilePath
      else do
        putStrLn "Couldn't find a Green's function data. Start simulation..."
        computeS . R.map magnitude <$>
          solveMonteCarloR2S1
            numThread
            numTrail
            numPoint
            numPoint
            numOrientation
            sigma
            tao
            histFilePath
            (0, 0, 0, 0, initOri / 180 * pi, initSpeed)
  xs <- parseEdgeFile edgeFilePath
  randomPonintSet <- generateRandomPointSet 40 numPoint numPoint
  let ys =
        L.map
          (\(R2S1RPPoint (a, b, c, d)) -> (R2S1RPPoint (a - 150, b - 150, c, d)))
          xs
      zs =
        L.map
          (\(R2S1RPPoint (a, b, c, d)) ->
             (R2S1RPPoint (a - center numPoint, b - center numPoint, c, d)))
          randomPonintSet
      pointset = ys L.++ zs
      bias = computeBias numPoint numPoint numOrientation pointset
      initialEigenVec =
        computeInitialEigenVec numPoint numPoint numOrientation pointset
  powerMethod
    emptyPlan
    folderPath
    arrG
    numIteration
    writeFlag
    ""
    threshold
    bias
    initialEigenVec


{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE NamedFieldPuns #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeApplications #-}

module Statistics.Array
  ( median
  , minimum
  , maximum
  , quartiles
  , listDerivative
  , mad
  , bowleySkew
  ) where

import Prelude hiding (minimum,maximum)

import Control.Monad (forM_)
import Data.Int (Int64)
import Data.Primitive (Prim,PrimArray)
import Data.Primitive.Contiguous (index, write)
import Statistics.Array.Types (Quartiles(..),AscList(..))

import qualified Data.Primitive.Contiguous as Arr
import qualified Data.Primitive.Sort as Arr
import qualified Statistics.Array.Types as Asc

-- | @O(1)@ Median element.
median :: Prim a
  => AscList a -> a
{-# inlineable median #-}
{-# specialize median :: AscList Int -> Int #-}
{-# specialize median :: AscList Int64 -> Int64 #-}
{-# specialize median :: AscList Double -> Double #-}
median (AscList arr) = Arr.index arr (Arr.size arr `div` 2)

-- | @O(1)@ Minimum element.
minimum :: Prim a
  => AscList a -> a
{-# inlineable minimum #-}
{-# specialize minimum :: AscList Int -> Int #-}
{-# specialize minimum :: AscList Int64 -> Int64 #-}
{-# specialize minimum :: AscList Double -> Double #-}
minimum (AscList arr) = Arr.index arr 0

-- | @O(1)@ Maximum element.
maximum :: Prim a
  => AscList a -> a
{-# inlineable maximum #-}
{-# specialize maximum :: AscList Int -> Int #-}
{-# specialize maximum :: AscList Int64 -> Int64 #-}
{-# specialize maximum :: AscList Double -> Double #-}
maximum (AscList arr) = Arr.index arr (Arr.size arr - 1)

quartiles :: Prim a
  => AscList a -> Quartiles a
quartiles (AscList arr) = Quartiles
  { q1 = Arr.index arr (len `div` 4)
  , q2 = Arr.index arr (len `div` 2)
  , q3 = Arr.index arr (3 * len `div` 4)
  }
  where
  len = Arr.size arr

-- hmmm, I guess I'm going with Int64 here to match with Chronos
-- times will be stored in number of nanoseconds after unix epoch
-- which allows me times up to about the year 2262

-- so:
--   check skew of time diffs
--   check dispersion of time diffs (median absolute deviation)
--   check skew of data sizes
--   check dispersion of time diffs (median absolute deviation)
--   check size of data (at the mode)

-- | Median of absolute deviations.
-- This is a measure of dispersion that is more robust than standard deviation.
mad :: forall a.
     (Prim a, Ord a, Num a)
  => AscList a -> a
{-# inlineable mad #-}
{-# specialize mad :: AscList Int -> Int #-}
{-# specialize mad :: AscList Int64 -> Int64 #-}
{-# specialize mad :: AscList Double -> Double #-}
mad asc@(AscList arr) =
  let m = median asc
      devs = (\x -> abs (x - m)) `Arr.map` arr :: PrimArray a
      devsAsc = Arr.sort devs
   in median $ Asc.unsafeFromAscendingArray devsAsc

-- | Take the difference between adjacent elements. Resulting array has
-- length one less than the argument array.
listDerivative :: (Prim a, Num a)
  => AscList a -> PrimArray a
{-# inlineable listDerivative #-}
{-# specialize listDerivative :: AscList Int -> PrimArray Int #-}
{-# specialize listDerivative :: AscList Int64 -> PrimArray Int64 #-}
{-# specialize listDerivative :: AscList Double -> PrimArray Double #-}
listDerivative (AscList arr) = Arr.create $ do
  let len' = Arr.size arr - 1
  diffs <- Arr.new len'
  let go !dstIx !prevElement = if dstIx < len'
        then do
          let !currentElement = index arr (dstIx + 1)
          write diffs dstIx (currentElement - prevElement)
          go (dstIx + 1) currentElement
        else pure diffs
  go 0 (index arr 0)

bowleySkew :: (Integral a) => Quartiles a -> Double
{-# inline bowleySkew #-}
bowleySkew Quartiles{q1,q2,q3}
  -- the limit as q1 approaches q3 is of course dependant on the direction of approach in the 3d space {q1,q2,q3}
  -- realistically, I'm going to call this zero, since real data has no approach, and the dirac distribution has zero skew
  | q1 == q3 = 0.0
  -- if `q1 == q2`, then skew simplifies to `(q3 - q1) / (q3 - q1)`, or `1`
  | q2 == q1 = 1.0
  -- similarly, if `q2 == q3`, then skew simplifies to `-1`
  | q2 == q3 = -1.0
  -- I'm not sure if either of those situations are reliable indicators of skew, though
  | otherwise = (realToFrac $ q3 + q1 - 2*q2) / (realToFrac $ q3 - q1)


module Physics (module Data.Complex, module Physics) where

import Data.Complex

type Vector = Complex Float
j :: Vector
j = 0:+1

normalize :: RealFloat a => Complex a -> Complex a
normalize (0:+0) = 0:+0
normalize a = a / (sqrt $ a * conjugate a)

toxy :: Complex a -> (a, a)
toxy (x :+ y) = (x, y)

toVector :: Float -> Vector
toVector x = x:+0


module NeuralNetwork where

import           Numeric.LinearAlgebra as LA
import           Text.Printf ( printf )
import           Numeric.Morpheus.Activation ( relu, reluGradient, sigmoid,
                    sigmoidGradient, tanh_, tanhGradient)

data Activation = Relu | Sigmoid | Tanh | Id

-- Neural network layer: weights, biases, and activation
data Layer a = Layer (Matrix a) (Matrix a) Activation

type NeuralNetwork a = [Layer a]

-- | Weight and bias gradients
data Gradients a = Gradients (Matrix a) (Matrix a)

getActivation :: Activation -> (Matrix Double -> Matrix Double)
getActivation Id = id
getActivation Sigmoid = sigmoid
getActivation Relu = relu
getActivation Tanh = tanh_

getActivation' :: Activation -> Matrix R -> Matrix R -> Matrix R
getActivation' Id = flip const
getActivation' Sigmoid = \x dY -> sigmoidGradient x * dY
getActivation' Relu = \x dY -> reluGradient x * dY
getActivation' Tanh = \x dY -> tanhGradient x * dY

bias' :: Matrix R -> Matrix R
bias' dY = cmap (/ m) r
  where
    r = matrix (cols dY) $ map sumElements (toColumns dY)
    m = fromIntegral $ rows dY

linearW' :: (Numeric t, Fractional t) => Matrix t -> Matrix t -> Matrix t
linearW' x dy = cmap (/ m) (tr' x LA.<> dy)
  where
    m = fromIntegral $ rows x

linearX' :: Numeric t => Matrix t -> Matrix t -> Matrix t
linearX' w dy = dy LA.<> tr' w

-- | Gradient descent optimization
optimize :: R -> Int -> [Layer R] -> (Matrix R, Matrix R) -> [Layer R]
optimize lr iterN net0 dataSet = last $ take iterN (iterate step net0)
  where
    step net = zipWith f net dW
        where
            (_, dW) = pass net dataSet

    f :: Layer R -> Gradients R -> Layer R
    f (Layer w b act) (Gradients dW dB) =
        Layer (w - lr `scale` dW) (b - lr `scale` dB) act

pass :: NeuralNetwork R -> (Matrix R, Matrix R) -> (Matrix R, [Gradients R])
pass net (x, tgt) = (pred, grads)
  where
    (_, pred, grads) = _pass x net

    _pass inp [] = (loss', pred, [])
      where
        pred = sigmoid inp
        loss' = pred - tgt

    _pass inp (Layer w b sact:layers) = (dX, pred, Gradients dW dB:t)
      where
        lin = (inp LA.<> w) + b
        y = getActivation sact lin

        (dZ, pred, t) = _pass y layers

        dY = getActivation' sact lin dZ
        dW = linearW' inp dY
        dB = bias' dY
        dX = linearX' w dY

data AdamParameters = AdamParameters { _beta1 :: R, _beta2 :: R, _epsilon :: R, _lr :: R }

adamParams :: AdamParameters
adamParams = AdamParameters { _beta1 = 0.9
                            , _beta2 = 0.999
                            , _epsilon = 1e-8
                            , _lr = 0.001
                            }

-- | Adam optimization
optimizeAdam :: AdamParameters -> Int -> [Layer R] -> (Matrix R, Matrix R) -> [Layer R]
optimizeAdam p iterN w0 dataSet = w
  where
    s0 = map zf w0
    v0 = map zf w0
    zf (Layer a b _) = (zerosLike a, zerosLike b)
    zerosLike m = matrix c (replicate (r*c) 0.0)
      where
        r = rows m
        c = cols m
    (w, _, _) = adam p iterN (w0, s0, v0) dataSet

adam :: AdamParameters -> Int -> ([Layer R], [(Matrix R, Matrix R)], [(Matrix R, Matrix R)]) -> (Matrix R, Matrix R) -> ([Layer R], [(Matrix R, Matrix R)], [(Matrix R, Matrix R)])
adam p@AdamParameters { _lr = lr, _beta1 = beta1, _beta2 = beta2, _epsilon = epsilon } iterN (w0, s0, v0) dataSet = last $ take iterN (iterate step (w0, s0, v0))
      where
        step (w, s, v) = (wN, sN, vN)
          where
            (_, dW) = pass w dataSet

            sN = zipWith f2 s dW
            vN = zipWith f3 v dW
            wN = zipWith3 f w vN sN

        f (Layer w_ b_ sf) (vW, vB) (sW, sB) =
           Layer (w_ - lr `scale` vW / (sqrt sW `addC` epsilon))
                 (b_ - lr `scale` vB / (sqrt sB `addC` epsilon))
                 sf

        addC m c = cmap (+ c) m

        f2 (sW, sB) (Gradients dW dB) =
          ( beta2 `scale` sW + (1 - beta2) `scale` (dW^2)
          , beta2 `scale` sB + (1 - beta2) `scale` (dB^2))

        f3 (vW, vB) (Gradients dW dB) =
          ( beta1 `scale` vW + (1 - beta1) `scale` dW
          , beta1 `scale` vB + (1 - beta1) `scale` dB)

genWeights :: (Int, Int) -> IO (Matrix Double, Matrix Double)
genWeights (nin, nout) = do
  w <- _genWeights (nin, nout)
  b <- _genWeights (1, nout)
  return (w, b)
    where
      _genWeights (nin, nout) = do
          let k = sqrt (1.0 / fromIntegral nin)
          w <- randn nin nout
          return (k `scale` w)

makeCircles :: Int -> R -> R -> IO (Matrix R, Matrix R)
makeCircles m factor noise = do
  let rand' n = scale (2 * pi) <$> rand n 1
      m1 = m `div` 2
      m2 = m - (m `div` 2)

  r1 <- rand' m1
  r2 <- rand' m2
  ns <- scale noise <$> randn m 2

  let outerX = cos r1
      outerY = sin r1
      innerX = scale factor $ cos r2
      innerY = scale factor $ sin r2
      -- Merge them all
      x = fromBlocks [[outerX, outerY], [innerX, innerY]]

      -- Labels
      y1 = m1 >< 1 $ repeat 0
      y2 = m2 >< 1 $ repeat 1
      y = y1 === y2

  return (x + ns, y)

makeSpirals :: Int -> R -> IO (Matrix R, Matrix R)
makeSpirals m noise = do
  r0 <- scale (780 * 2 * pi / 360) . sqrt <$> rand m 1
  d1x0 <- scale noise <$> rand m 1
  d1y0 <- scale noise <$> rand m 1

  let d1x = d1x0 - cos r0 * r0
  let d1y = d1y0 + sin r0 * r0

  let x = fromBlocks [[d1x, d1y], [- d1x, - d1y]] / 10.0
  let y1 = m >< 1 $ repeat 0
  let y2 = m >< 1 $ repeat 1
  let y = y1 === y2
  return (x, y)

accuracy :: [Layer R] -> (Matrix R, Matrix R) -> R
accuracy net (dta, tgt) = 100 * (1 - e / m)
  where
    pred = net `inferBinary` dta
    e = sumElements $ abs (tgt - pred)
    m = fromIntegral $ rows tgt

inferBinary :: NeuralNetwork R -> Matrix R -> Matrix R
inferBinary net dta =
  let pred = forward net dta
  in cmap (\a -> if a < 0.5 then 0 else 1) pred

forward :: NeuralNetwork R -> Matrix R -> Matrix R
forward net dta = fst $ pass net (dta, undefined)

genNetwork :: [Int] -> [Activation] -> IO [Layer R]
genNetwork nodes activations = do
    weights <- mapM genWeights nodes'
    return (zipWith (\(w, b) a -> Layer w b a) weights activations)
  where
    nodes' = zip nodes (tail nodes)

experiment1 :: IO ()
experiment1 = do
  trainSet <- makeCircles 200 0.6 0.1
  testSet <- makeCircles 100 0.6 0.1

  let (dta, tgt) = trainSet

  (w1_rand, b1_rand) <- genWeights (2, 128)
  (w2_rand, b2_rand) <- genWeights (128, 1)

  let net = [ Layer w1_rand b1_rand Relu
            , Layer w2_rand b2_rand Id ]

  -- -- Alternatively:
  -- net0 <- genNetwork [2, 128, 1] [Relu, Id]

  let epochs = 1000
      lr = 0.001  -- Learning rate
      net' = optimize lr epochs net trainSet
      netA = optimizeAdam adamParams epochs net trainSet

  putStrLn $ printf "Circles problem, 1 hidden layer of 128 neurons, %d epochs" epochs
  putStrLn "---"

  putStrLn $ printf "Training accuracy (gradient descent) %.1f" (net' `accuracy` trainSet)
  putStrLn $ printf "Validation accuracy (gradient descent) %.1f\n" (net' `accuracy` testSet)

  putStrLn $ printf "Training accuracy (Adam) %.1f" (netA `accuracy` trainSet)
  putStrLn $ printf "Validation accuracy (Adam) %.1f\n" (netA `accuracy` testSet)

  putStrLn ""

experiment2 :: IO ()
experiment2 = do
  trainSet <- makeSpirals 200 0.5
  testSet <- makeSpirals 100 0.5

  let (dta, tgt) = trainSet
      epochs = 700

  putStrLn $ printf "Spirals problem, Adam, %d epochs" epochs
  putStrLn "---"
  putStrLn "1 hidden layer, 128 neurons (513 parameters)"
  net0 <- genNetwork [2, 128, 1] [Relu, Id]
  let net0' = optimizeAdam adamParams epochs net0 trainSet

  putStrLn $ printf "Training accuracy %.1f" (net0' `accuracy` trainSet)
  putStrLn $ printf "Validation accuracy %.1f\n" (net0' `accuracy` testSet)

  putStrLn "1 hidden layer, 512 neurons (2049 parameters)"
  net1 <- genNetwork [2, 512, 1] [Relu, Id]
  let net1' = optimizeAdam adamParams epochs net1 trainSet

  putStrLn $ printf "Training accuracy %.1f" (net1' `accuracy` trainSet)
  putStrLn $ printf "Validation accuracy %.1f\n" (net1' `accuracy` testSet)

  putStrLn "3 hidden layers, 40, 25, and 10 neurons (1416 parameters)"
  net2 <- genNetwork [2, 40, 25, 10, 1] [Relu, Relu, Relu, Id]
  let net2' = optimizeAdam adamParams epochs net2 trainSet

  putStrLn $ printf "Training accuracy %.1f" (net2' `accuracy` trainSet)
  putStrLn $ printf "Validation accuracy %.1f\n" (net2' `accuracy` testSet)

main :: IO ()
main = experiment1 >> experiment2


{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE OverloadedLabels #-}
{-# LANGUAGE TemplateHaskell #-}
{-# LANGUAGE UndecidableInstances #-}

module Day12 where

import Control.Monad.State (MonadState, State, execState)
import Data.ByteString qualified as B
import Data.Complex (Complex((:+)), cis)
import Data.Functor (($>))
import Data.Text qualified as T
import Data.Text.Encoding qualified as TE
import Data.Void (Void)
import Flow ((.>))
import Optics
    ( Optic, A_Setter, Is, view, noPrefixFieldLabels, makeFieldLabelsWith )
import Optics.State (use)
import Optics.State.Operators ((%=))
import Text.Megaparsec qualified as P
import Text.Megaparsec.Char qualified as PC
import Text.Megaparsec.Char.Lexer qualified as PCL
-- import Text.Pretty.Simple (pPrint)

type Parser = P.Parsec Void T.Text

data NavAction
    = MoveNorth
    | MoveSouth
    | MoveEast
    | MoveWest
    | TurnLeft
    | TurnRight
    | MoveForward
    deriving (Bounded, Enum, Eq, Ord, Read, Show)

data NavInstr = MkNavInstr
    { navAction :: !NavAction
    , navValue :: !Int
    } deriving (Eq, Ord, Read, Show)

data ShipState = MkShipState
    { shipPhase :: !Double
    , shipPosition :: !(Complex Double)
    , wayPosition :: !(Complex Double)
    } deriving (Eq, Read, Show)

makeFieldLabelsWith noPrefixFieldLabels ''ShipState

main :: IO ()
main = do
    file <- readFileUtf8 "day-12/input.txt"
    case P.parse (parseNavInstrs <* P.eof) "day 12 input" file of
        Left err -> putStrLn $ P.errorBundlePretty err
        Right navInstrs -> do
            print $ part1 navInstrs
            print $ part2 navInstrs

part2 :: [NavInstr] -> Double
part2 = pilotShip2 .> view #shipPosition .> manhattanNorm

part1 :: [NavInstr] -> Double
part1 = pilotShip1 .> view #shipPosition .> manhattanNorm

pilotShip2 :: [NavInstr] -> ShipState
pilotShip2 instrs = execState (applyInstrs2 instrs) $
    MkShipState
        { shipPhase = 0
        , shipPosition = pure 0
        , wayPosition = 10 :+ 1
        }

pilotShip1 :: [NavInstr] -> ShipState
pilotShip1 instrs = execState (applyInstrs1 instrs) $
    MkShipState
        { shipPhase = 0
        , shipPosition = pure 0
        , wayPosition = pure 0
        }

applyInstrs1 :: [NavInstr] -> State ShipState ()
applyInstrs1 = foldr (\instr st -> applyInstr1 instr *> st) (pure ())

applyInstr1 :: NavInstr -> State ShipState ()
applyInstr1 (MkNavInstr action intVal) = case action of
    MoveNorth -> #shipPosition += (0 :+ val)
    MoveSouth -> #shipPosition -= (0 :+ val)
    MoveEast -> #shipPosition += (val :+ 0)
    MoveWest -> #shipPosition -= (val :+ 0)
    TurnLeft -> #shipPhase %= (+ val) .> (`modulo` 360)
    TurnRight -> #shipPhase %= subtract val .> (`modulo` 360)
    MoveForward -> do
        phase <- use #shipPhase
        #shipPosition += fmap (* val) (cisDegrees phase)
  where
    val = fromIntegral intVal

applyInstrs2 :: [NavInstr] -> State ShipState ()
applyInstrs2 = foldr (\instr st -> applyInstr2 instr *> st) (pure ())

applyInstr2 :: NavInstr -> State ShipState ()
applyInstr2 (MkNavInstr action intVal) = case action of
    MoveNorth -> #wayPosition += (0 :+ val)
    MoveSouth -> #wayPosition -= (0 :+ val)
    MoveEast -> #wayPosition += (val :+ 0)
    MoveWest -> #wayPosition -= (val :+ 0)
    TurnLeft -> #wayPosition *= cisDegrees val
    TurnRight -> #wayPosition *= cisDegrees (-val)
    MoveForward -> do
        wayPos <- use #wayPosition
        #shipPosition += fmap (* val) wayPos
  where
    val = fromIntegral intVal

cisDegrees :: Double -> Complex Double
cisDegrees degRaw
    | deg == 0 = 1 :+ 0
    | deg == 90 = 0 :+ 1
    | deg == 180 = (-1) :+ 0
    | deg == 270 = 0 :+ (-1)
    | otherwise = cis rad
  where
    deg = degRaw `modulo` 360
    rad = deg * pi / 180

modulo :: Double -> Double -> Double
modulo x m = x - floorDbl (x / m) * m
  where
    floorDbl = floor .> fromIntegral @Int

manhattanDist :: RealFloat a => Complex a -> Complex a -> a
manhattanDist u v = sum $ fmap abs $ v - u

manhattanNorm :: RealFloat a => Complex a -> a
manhattanNorm = manhattanDist (pure 0)

parseNavInstrs :: Parser [NavInstr]
parseNavInstrs = parseNavInstr `P.sepEndBy` PC.space

parseNavInstr :: Parser NavInstr
parseNavInstr = do
    action <- parseNavAction
    val <- PC.space *> PCL.decimal
    pure $ MkNavInstr action val

parseNavAction :: Parser NavAction
parseNavAction = P.choice
    [ PC.char 'N' $> MoveNorth
    , PC.char 'S' $> MoveSouth
    , PC.char 'E' $> MoveEast
    , PC.char 'W' $> MoveWest
    , PC.char 'L' $> TurnLeft
    , PC.char 'R' $> TurnRight
    , PC.char 'F' $> MoveForward
    ]

readFileUtf8 :: FilePath -> IO T.Text
readFileUtf8 path = TE.decodeUtf8 <$> B.readFile path

(+=)
    :: (Is k A_Setter, MonadState s m, Num a)
    => Optic k is s s a a
    -> a
    -> m ()
optic += x = optic %= (+ x)

(-=)
    :: (Is k A_Setter, MonadState s m, Num a)
    => Optic k is s s a a
    -> a
    -> m ()
optic -= x = optic %= subtract x

(*=)
    :: (Is k A_Setter, MonadState s m, Num a)
    => Optic k is s s a a
    -> a
    -> m ()
optic *= x = optic %= (* x)


{-# LANGUAGE ViewPatterns #-}
{-# OPTIONS_GHC -Wall #-}

import qualified Numeric.Sundials.ARKode.ODE as ARK
import qualified Numeric.Sundials.CVode.ODE  as CV
import           Numeric.LinearAlgebra
import           Numeric.Sundials.ODEOpts (ODEOpts(..))

import           Plots as P
import qualified Diagrams.Prelude as D
import           Diagrams.Backend.Rasterific

import           Control.Lens

import           Data.Functor.Compose

import           Test.Hspec

lorenz :: Double -> [Double] -> [Double]
lorenz _t u = [ sigma * (y - x)
              , x * (rho - z) - y
              , x * y - beta * z
              ]
  where
    rho = 28.0
    sigma = 10.0
    beta = 8.0 / 3.0
    x = u !! 0
    y = u !! 1
    z = u !! 2

_lorenzJac :: Double -> Vector Double -> Matrix Double
_lorenzJac _t u = (3><3) [ (-sigma), rho - z, y
                        , sigma   , -1.0   , x
                        , 0.0     , (-x)   , (-beta)
                        ]
  where
    rho = 28.0
    sigma = 10.0
    beta = 8.0 / 3.0
    x = u ! 0
    y = u ! 1
    z = u ! 2

brusselator :: Double -> [Double] -> [Double]
brusselator _t x = [ a - (w + 1) * u + v * u * u
                   , w * u - v * u * u
                   , (b - w) / eps - w * u
                   ]
  where
    a = 1.0
    b = 3.5
    eps = 5.0e-6
    u = x !! 0
    v = x !! 1
    w = x !! 2

brussJac :: Double -> Vector Double -> Matrix Double
brussJac _t x = tr $
  (3><3) [ (-(w + 1.0)) + 2.0 * u * v, w - 2.0 * u * v, (-w)
         , u * u                     , (-(u * u))     , 0.0
         , (-u)                      , u              , (-1.0) / eps - u
         ]
  where
    y = toList x
    u = y !! 0
    v = y !! 1
    w = y !! 2
    eps = 5.0e-6

brusselatorWithJacobian :: Vector Double -> Bool -> CV.SolverResult
brusselatorWithJacobian ts usejac = CV.odeSolveRootVWith' opts
                      (\t v -> vector $ brusselator t (toList v))
                      (if usejac then Just brussJac else Nothing)
                      (vector [1.2, 3.1, 3.0])
                      [] 0
                      ts
  where
    opts = ODEOpts { maxNumSteps = 10000
                   , minStep     = 1.0e-12
                   , maxFail     = 10
                   , odeMethod = CV.BDF
                   , stepControl = CV.XX' 1.0e-6 1.0e-10 1 1
                   , initStep = Nothing
                   }

brussRoot :: CV.SolverResult
brussRoot = CV.odeSolveRootVWith' opts
                      (\t v -> vector $ brusselator t (toList v))
                      Nothing
                      (vector [1.2, 3.1, 3.0])
                      events 100
                      (vector [0.0, 0.1 .. 10.0])
  where
    events =
      [ CV.EventSpec { CV.eventCondition = brussRootFn
                     , CV.eventUpdate =
                         \_ev x -> let y = toList x in vector [(y!!0) + 0.5 , (y!!1), (y!!2)]
                     , CV.eventDirection = CV.AnyDirection
                     }
      ]
    opts = ODEOpts { maxNumSteps = 10000
                   , minStep     = 1.0e-12
                   , maxFail     = 10
                   , odeMethod = CV.BDF
                   , stepControl = CV.XX' 1.0e-6 1.0e-10 1 1
                   , initStep = Nothing
                   }

brussRootFn :: Double -> Vector Double -> Double
brussRootFn _ v = case xs of
                    [y1, _y2, y3] -> y1 - y3
                    _            -> error "brusselator root function RHS not defined"
  where
    xs = toList v

exponential :: CV.SolverResult
exponential = CV.odeSolveRootVWith' opts
                      (\t y -> vector [y ! 0])
                      Nothing
                      (vector [1])
                      events 100
                      (vector [ fromIntegral k / 100 | k <- [0..(22::Int)]])
  where
    events =
      [ CV.EventSpec { CV.eventCondition = \t y -> y ! 0 - 1.1
                     , CV.eventUpdate = \ev y -> vector [ 2 ]
                     , CV.eventDirection = CV.Upwards
                     }
      ]
    opts = ODEOpts { maxNumSteps = 10000
                   , minStep     = 1.0e-12
                   , maxFail     = 10
                   , odeMethod = CV.BDF
                   , stepControl = CV.XX' 1.0e-6 1.0e-10 1 1
                   , initStep = Nothing
                   }

-- A sine wave that only changes direction once it reaches ±0.9.
-- Illustrates event-specific reset function
boundedSine :: CV.SolverResult
boundedSine = CV.odeSolveRootVWith'
  opts
  (\_t y -> vector [y ! 1, - y ! 0]) -- ODE RHS
  Nothing
  (vector [0, 1]) -- initial conditions
  events
  100 -- maximum number of events
  (vector [ 2 * pi * k / 360 | k <- [0..360]]) -- solution times
  where
    opts = ODEOpts { maxNumSteps = 10000
                   , minStep     = 1.0e-12
                   , maxFail     = 10
                   , odeMethod = CV.ADAMS
                   , stepControl = CV.XX' 1.0e-6 1.0e-10 1 1
                   , initStep = Nothing
                   }
    events =
      [ CV.EventSpec { CV.eventCondition = \_t y -> y ! 0 - 0.9
                     , CV.eventUpdate = \_ y -> vector [ y ! 0, - abs (y ! 1) ]
                     , CV.eventDirection = CV.Upwards
                     }
      , CV.EventSpec { CV.eventCondition = \_t y -> y ! 0 + 0.9
                     , CV.eventUpdate = \_ y -> vector [ y ! 0, abs (y ! 1) ]
                     , CV.eventDirection = CV.Downwards
                     }
      ]

stiffish :: Double -> [Double] -> [Double]
stiffish t v = [ lamda * u + 1.0 / (1.0 + t * t) - lamda * atan t ]
  where
    lamda = -100.0
    u = v !! 0

stiffishV :: Double -> Vector Double -> Vector Double
stiffishV t v = fromList [ lamda * u + 1.0 / (1.0 + t * t) - lamda * atan t ]
  where
    lamda = -100.0
    u = v ! 0

_stiffJac :: Double -> Vector Double -> Matrix Double
_stiffJac _t _v = (1><1) [ lamda ]
  where
    lamda = -100.0

predatorPrey :: Double -> [Double] -> [Double]
predatorPrey _t v = [ x * a - b * x * y
                    , d * x * y - c * y - e * y * z
                    , (-f) * z + g * y * z
                    ]
  where
    x = v!!0
    y = v!!1
    z = v!!2
    a = 1.0
    b = 1.0
    c = 1.0
    d = 1.0
    e = 1.0
    f = 1.0
    g = 1.0

roberts :: Double -> Vector Double -> Vector Double
roberts t v = vector $ robertsAux t (toList v)
  where
    robertsAux _ [y1, y2, y3] =
      [ -0.04 * y1 + 1.0e4 * y2 * y3
      , 0.04 * y1 - 1.0e4 * y2 * y3 - 3.0e7 * (y2)^(2 :: Int)
      , 3.0e7 * (y2)^(2 :: Int)
      ]
    robertsAux _ _ = error "roberts RHS not defined"

robertsJac :: Double -> Vector Double -> Matrix Double
robertsJac _t (toList -> [y1, y2, y3]) = (3 >< 3)
  [ -0.04, 1.0e4 * y3, 1.0e4 * y2
  , 0.04, -1.0e4*y3 - 3.0e7*2*y2, -1.0e4*y2
  , 0, 3.0e7*2*y2, 0
  ]

ts :: [Double]
ts = take 12 $ map (* 10.0) (0.04 : ts)

solve :: CV.SolverResult
solve = CV.odeSolveRootVWith' opts
                      roberts Nothing (vector [1.0, 0.0, 0.0])
                      events 100
                      (vector (0.0 : ts))
  where
    opts = ODEOpts { maxNumSteps = 10000
                   , minStep     = 1.0e-12
                   , maxFail     = 10
                   , odeMethod = CV.BDF
                   , stepControl = CV.ScXX' 1.0 1.0e-4 1.0 1.0 (vector [1.0e-8, 1.0e-14, 1.0e-6])
                   , initStep = Nothing
                   }
    events =
      [ CV.EventSpec { CV.eventCondition = \_t y -> y ! 0 - 0.0001
                     , CV.eventUpdate = const id
                     , CV.eventDirection = CV.AnyDirection
                     }
      , CV.EventSpec { CV.eventCondition = \_t y -> y ! 2 - 0.01
                     , CV.eventUpdate = const id
                     , CV.eventDirection = CV.AnyDirection
                     }
      ]

solve2 :: CV.SolverResult
solve2 = CV.odeSolveRootVWith' opts
                      roberts Nothing (vector [1.0, 0.0, 0.0])
                      events 100
                      (vector (0.0 : ts))
  where
    opts = ODEOpts { maxNumSteps = 10000
                   , minStep     = 1.0e-12
                   , maxFail     = 10
                   , odeMethod = CV.BDF
                   , stepControl = CV.ScXX' 1.0 1.0e-4 1.0 1.0 (vector [1.0e-8, 1.0e-14, 1.0e-6])
                   , initStep = Nothing
                   }
    events =
      [ CV.EventSpec { CV.eventCondition = \_t y -> y ! 0 - 0.0001
                     , CV.eventUpdate = upd
                     , CV.eventDirection = CV.AnyDirection
                     }
      , CV.EventSpec { CV.eventCondition = \_t y -> y ! 2 - 0.01
                     , CV.eventUpdate = upd
                     , CV.eventDirection = CV.AnyDirection
                     }
      ]
    upd _ _ = vector [1.0, 0.0, 0.0]

solve1 :: CV.SolverResult
solve1 = CV.odeSolveRootVWith' opts
                      roberts Nothing (vector [1.0, 0.0, 0.0])
                      events 100
                      (vector (0.0 : ts))
  where
    opts = ODEOpts { maxNumSteps = 10000
                   , minStep     = 1.0e-12
                   , maxFail     = 10
                   , odeMethod = CV.BDF
                   , stepControl = CV.ScXX' 1.0 1.0e-4 1.0 1.0 (vector [1.0e-8, 1.0e-14, 1.0e-6])
                   , initStep = Nothing
                   }
    events =
      [ CV.EventSpec { CV.eventCondition = \t _y -> t - 1.0
                     , CV.eventUpdate = \t y -> vector [2.0, y!1, y!2]
                     , CV.eventDirection = CV.AnyDirection
                     }
      ]

robertsonWithJacobian :: Vector Double -> Bool -> CV.SolverResult
robertsonWithJacobian ts usejac = CV.odeSolveRootVWith' opts
                      roberts (if usejac then Just robertsJac else Nothing) (vector [1.0, 0.0, 0.0])
                      [] 0
                      ts
  where
    opts = ODEOpts { maxNumSteps = 10000
                   , minStep     = 1.0e-12
                   , maxFail     = 10
                   , odeMethod = CV.BDF
                   , stepControl = CV.ScXX' 1.0 1.0e-4 1.0 1.0 (vector [1.0e-8, 1.0e-14, 1.0e-6])
                   , initStep = Nothing
                   }

lSaxis :: [[Double]] -> P.Axis B D.V2 Double
lSaxis xs = P.r2Axis &~ do
  let zs = xs!!0
      us = xs!!1
      vs = xs!!2
      ws = xs!!3
  P.linePlot' $ zip zs us
  P.linePlot' $ zip zs vs
  P.linePlot' $ zip zs ws

lSaxis2 :: [[Double]] -> P.Axis B D.V2 Double
lSaxis2 xs = P.r2Axis &~ do
  let zs = xs!!0
      us = xs!!1
      vs = xs!!2
  P.linePlot' $ zip zs us
  P.linePlot' $ zip zs vs

kSaxis :: [(Double, Double)] -> P.Axis B D.V2 Double
kSaxis xs = P.r2Axis &~ do
  P.linePlot' xs

main :: IO ()
main = do
  let res1 = ARK.odeSolve brusselator [1.2, 3.1, 3.0] (fromList [0.0, 0.1 .. 10.0])
  renderRasterific "diagrams/brusselator.png"
                   (D.dims2D 500.0 500.0)
                   (renderAxis $ lSaxis $ [0.0, 0.1 .. 10.0]:(toLists $ tr res1))

  let res1a = ARK.odeSolve brusselator [1.2, 3.1, 3.0] (fromList [0.0, 0.1 .. 10.0])
  renderRasterific "diagrams/brusselatorA.png"
                   (D.dims2D 500.0 500.0)
                   (renderAxis $ lSaxis $ [0.0, 0.1 .. 10.0]:(toLists $ tr res1a))

  let res2 = ARK.odeSolve stiffish [0.0] (fromList [0.0, 0.1 .. 10.0])
  renderRasterific "diagrams/stiffish.png"
                   (D.dims2D 500.0 500.0)
                   (renderAxis $ kSaxis $ zip [0.0, 0.1 .. 10.0] (concat $ toLists res2))

  let res2a = ARK.odeSolveV (ARK.SDIRK_5_3_4') Nothing 1e-6 1e-10 stiffishV (fromList [0.0]) (fromList [0.0, 0.1 .. 10.0])

  let res2b = ARK.odeSolveV (ARK.TRBDF2_3_3_2') Nothing 1e-6 1e-10 stiffishV (fromList [0.0]) (fromList [0.0, 0.1 .. 10.0])

  let maxDiffA = maximum $ map abs $
                 zipWith (-) ((toLists $ tr res2a)!!0) ((toLists $ tr res2b)!!0)

  let res2c = CV.odeSolveV (CV.BDF) Nothing 1e-6 1e-10 stiffishV (fromList [0.0]) (fromList [0.0, 0.1 .. 10.0])

  let maxDiffB = maximum $ map abs $
                 zipWith (-) ((toLists $ tr res2a)!!0) ((toLists $ tr res2c)!!0)

  let maxDiffC = maximum $ map abs $
                 zipWith (-) ((toLists $ tr res2b)!!0) ((toLists $ tr res2c)!!0)

  let res3 = ARK.odeSolve lorenz [-5.0, -5.0, 1.0] (fromList [0.0, 0.01 .. 20.0])

  renderRasterific "diagrams/lorenz.png"
                   (D.dims2D 500.0 500.0)
                   (renderAxis $ kSaxis $ zip ((toLists $ tr res3)!!0) ((toLists $ tr res3)!!1))

  renderRasterific "diagrams/lorenz1.png"
                   (D.dims2D 500.0 500.0)
                   (renderAxis $ kSaxis $ zip ((toLists $ tr res3)!!0) ((toLists $ tr res3)!!2))

  renderRasterific "diagrams/lorenz2.png"
                   (D.dims2D 500.0 500.0)
                   (renderAxis $ kSaxis $ zip ((toLists $ tr res3)!!1) ((toLists $ tr res3)!!2))

  let res4 = CV.odeSolve predatorPrey [0.5, 1.0, 2.0] (fromList [0.0, 0.01 .. 10.0])

  renderRasterific "diagrams/predatorPrey.png"
                   (D.dims2D 500.0 500.0)
                   (renderAxis $ kSaxis $ zip ((toLists $ tr res4)!!0) ((toLists $ tr res4)!!1))

  renderRasterific "diagrams/predatorPrey1.png"
                   (D.dims2D 500.0 500.0)
                   (renderAxis $ kSaxis $ zip ((toLists $ tr res4)!!0) ((toLists $ tr res4)!!2))

  renderRasterific "diagrams/predatorPrey2.png"
                   (D.dims2D 500.0 500.0)
                   (renderAxis $ kSaxis $ zip ((toLists $ tr res4)!!1) ((toLists $ tr res4)!!2))

  let res4a = ARK.odeSolve predatorPrey [0.5, 1.0, 2.0] (fromList [0.0, 0.01 .. 10.0])

  let maxDiffPpA = maximum $ map abs $
                   zipWith (-) ((toLists $ tr res4)!!0) ((toLists $ tr res4a)!!0)

  let cond5 =
        case solve of
          CV.SolverSuccess events _ _ -> do
            length events `shouldBe` 2
            (abs (CV.eventTime (events!!0) - 0.2640208751331032) / 0.2640208751331032 < 1.0e-8) `shouldBe` True
            (abs (CV.eventTime (events!!1) - 2.0786731062254436e7) / 2.0786731062254436e7 < 1.0e-8) `shouldBe` True
          CV.SolverError _ _ ->
            error "Root finding error!"

  let cond6 =
        case solve1 of
          CV.SolverSuccess events _ _ -> do
            length events `shouldBe` 1
            (abs (CV.eventTime (events!!0) - 1.0) / 1.0 < 1.0e-10) `shouldBe` True
          CV.SolverError _ _ ->
            error "Root finding error!"

  let cond7 =
        case solve2 of
          CV.SolverSuccess {} ->
            error "Solver returned Success"
          CV.SolverError _ _ ->
            True

  case brussRoot of
    CV.SolverSuccess events m _diagn -> do
      renderRasterific
        "diagrams/brussRoot.png"
        (D.dims2D 500.0 500.0)
        (renderAxis $ lSaxis $ toLists $ tr m)
    CV.SolverError m n ->
      expectationFailure $ show n

  let boundedSineSpec = do
        case boundedSine of
          CV.SolverSuccess events m _ -> do
            renderRasterific
              "diagrams/boundedSine.png"
              (D.dims2D 500.0 500.0)
              (renderAxis $ lSaxis2 $ toLists $ tr m)
            length events `shouldBe` 3
            map CV.rootDirection events `shouldBe` [CV.Upwards, CV.Downwards, CV.Upwards]
            map CV.eventIndex events `shouldBe` [0, 1, 0]
            all ((< 1e-8) . abs) (zipWith (-)
              (map CV.eventTime events)
              [1.1197660081724263,3.3592952656818404,5.5988203973243])
                `shouldBe` True
          CV.SolverError m n ->
            expectationFailure "Solver error"
  let exponentialSpec = do
        case exponential of
          CV.SolverSuccess events _m _diagn -> do
            length events `shouldBe` 1
            (abs (CV.eventTime (events!!0) - log 1.1) < 1e-4) `shouldBe` True
            CV.rootDirection (events!!0) `shouldBe` CV.Upwards
            CV.eventIndex (events!!0) `shouldBe` 0
          CV.SolverError m n ->
            expectationFailure $ show n

      robertsonJac = do
        let ts = vector [0, 1 .. 10]
            CV.SolverSuccess _ m1 _ = robertsonWithJacobian ts True
            CV.SolverSuccess _ m2 _ = robertsonWithJacobian ts False
        (norm_2 (m1-m2) < 1e-4) `shouldBe` True

      brusselatorJac = do
        let ts = [0.0, 0.1 .. 10.0]
            CV.SolverSuccess _ m1 _ = brusselatorWithJacobian (vector ts) True
            CV.SolverSuccess _ m2 _ = brusselatorWithJacobian (vector ts) False
        (norm_2 (m1-m2) < 1e-3) `shouldBe` True

  hspec $ describe "Compare results" $ do
    it "Robertson should fail" $ cond7
    it "Robertson time only" $ cond6
    it "Robertson from SUNDIALS manual" $ cond5
    it "Robertson with explicit Jacobian up to t=10" robertsonJac
    it "Brusselator with explicit Jacobian" brusselatorJac
    it "for SDIRK_5_3_4' and TRBDF2_3_3_2'" $ maxDiffA < 1.0e-6
    it "for SDIRK_5_3_4' and BDF" $ maxDiffB < 1.0e-6
    it "for TRBDF2_3_3_2' and BDF" $ maxDiffC < 1.0e-6
    it "for CV and ARK for the Predator Prey model" $ maxDiffPpA < 1.0e-3
    it "Bounded sine events" $ boundedSineSpec
    it "Exponential events" $ exponentialSpec