import Graphics.Rendering.Plot
import Numeric.LinearAlgebra
import Numeric.GSL.Statistics
import qualified Graphics.Rendering.Cairo as C
import Graphics.Rendering.Pango.Enums
import Graphics.UI.Gtk hiding(Circle,Cross)
import Control.Monad.Trans

ln = 25
ts = linspace ln (0,1)
rs = randomVector 0 Gaussian ln

ss = sin (15*2*pi*ts)
ds = 0.25*rs + ss
es = constant (0.25*(stddev rs)) ln

fs :: Double -> Double
fs = sin . (15*2*pi*)

test_graph = do
       withTitle $ setText "Testing plot package:"
       withSubTitle $ do
                      setText "with 1 second of a 15Hz sine wave"
       withPointDefaults $ setPointSize 1.5 -- (A)
       setPlots 1 1
       withPlot (1,1) $ do
                        setLegend True SouthWest Inside -- (B)
                        addAxis XAxis (Side Lower) $ withAxisLabel $ setText "time (s)"
                        addAxis YAxis (Side Lower) $ do
                          withAxisLabel $ setText "amplitude"
                          withGridLine Major $ do
                            setDashStyle [Dash]
                            setLineColour lightgray
                        addAxis XAxis (Value 0) $ return ()
                        setRange YAxis Lower Linear (-1.25) 1.25
                        setDataset (ts,[point (ds,es,"dat") (Cross,red),line (fs,"sin") blue])
                        setRangeFromData XAxis Lower Linear

display :: ((Int,Int) -> C.Render ()) -> IO ()
display r = do
   initGUI       -- is start

   window <- windowNew
   set window [ windowTitle := "Cairo test window"
              , windowDefaultWidth := 600
              , windowDefaultHeight := 400
              , containerBorderWidth := 1
              ]

--   canvas <- pixbufNew ColorspaceRgb True 8 300 200
--   containerAdd window canvas
   frame <- frameNew
   containerAdd window frame
   canvas <- drawingAreaNew
   containerAdd frame canvas
   widgetModifyBg canvas StateNormal (Color 65535 65535 65535)

   widgetShowAll window

   on canvas exposeEvent $ tryEvent $ do s <- liftIO $ widgetGetSize canvas
                                         drw <- liftIO $ widgetGetDrawWindow canvas
                                         --dat <- liftIO $ takeMVar d
                                         --liftIO $ renderWithDrawable drw (circle 50 10)
                                         liftIO $ renderWithDrawable drw (r s)

   onDestroy window mainQuit
   mainGUI

main = display $ render test_graph

test_render :: (Int,Int) -> C.Render ()
test_render = render test_graph

--main = C.withSVGSurface "out.svg" 400 400 $ \surf -> C.renderWith surf $ test_render (400,400)


{-# LANGUAGE BangPatterns, FlexibleContexts #-}
-- |
-- Module    : Statistics.Transform
-- Copyright : (c) 2011 Bryan O'Sullivan
-- License   : BSD3
--
-- Maintainer  : bos@serpentine.com
-- Stability   : experimental
-- Portability : portable
--
-- Fourier-related transformations of mathematical functions.
--
-- These functions are written for simplicity and correctness, not
-- speed.  If you need a fast FFT implementation for your application,
-- you should strongly consider using a library of FFTW bindings
-- instead.

module Statistics.Transform
    (
    -- * Type synonyms
      CD
    -- * Discrete cosine transform
    , dct
    , dct_
    , idct
    , idct_
    -- * Fast Fourier transform
    , fft
    , ifft
    ) where

import Control.Monad (when)
import Control.Monad.ST (ST)
import Data.Bits (shiftL, shiftR)
import Data.Complex (Complex(..), conjugate, realPart)
import Numeric.SpecFunctions (log2)
import qualified Data.Vector.Generic         as G
import qualified Data.Vector.Generic.Mutable as M
import qualified Data.Vector.Unboxed         as U
import qualified Data.Vector                 as V

type CD = Complex Double

-- | Discrete cosine transform (DCT-II).
dct :: (G.Vector v CD, G.Vector v Double, G.Vector v Int) => v Double -> v Double
dct = dctWorker . G.map (:+0)
{-# INLINABLE  dct #-}
{-# SPECIAlIZE dct :: U.Vector Double -> U.Vector Double #-}
{-# SPECIAlIZE dct :: V.Vector Double -> V.Vector Double #-}

-- | Discrete cosine transform (DCT-II). Only real part of vector is
--   transformed, imaginary part is ignored.
dct_ :: (G.Vector v CD, G.Vector v Double, G.Vector v Int) => v CD -> v Double
dct_ = dctWorker . G.map (\(i :+ _) -> i :+ 0)
{-# INLINABLE  dct_ #-}
{-# SPECIAlIZE dct_ :: U.Vector CD -> U.Vector Double #-}
{-# SPECIAlIZE dct_ :: V.Vector CD -> V.Vector Double#-}

dctWorker :: (G.Vector v CD, G.Vector v Double, G.Vector v Int) => v CD -> v Double
{-# INLINE dctWorker #-}
dctWorker xs
  -- length 1 is special cased because shuffle algorithms fail for it.
  | G.length xs == 1 = G.map ((2*) . realPart) xs
  | vectorOK xs      = G.map realPart $ G.zipWith (*) weights (fft interleaved)
  | otherwise        = error "Statistics.Transform.dct: bad vector length"
  where
    interleaved = G.backpermute xs $ G.enumFromThenTo 0 2 (len-2) G.++
                                     G.enumFromThenTo (len-1) (len-3) 1
    weights = G.cons 2 . G.generate (len-1) $ \x ->
              2 * exp ((0:+(-1))*fi (x+1)*pi/(2*n))
      where n = fi len
    len = G.length xs

-- | Inverse discrete cosine transform (DCT-III). It's inverse of
-- 'dct' only up to scale parameter:
--
-- > (idct . dct) x = (* length x)
idct :: (G.Vector v CD, G.Vector v Double) => v Double -> v Double
idct = idctWorker . G.map (:+0)
{-# INLINABLE  idct #-}
{-# SPECIAlIZE idct :: U.Vector Double -> U.Vector Double #-}
{-# SPECIAlIZE idct :: V.Vector Double -> V.Vector Double #-}

-- | Inverse discrete cosine transform (DCT-III). Only real part of vector is
--   transformed, imaginary part is ignored.
idct_ :: (G.Vector v CD, G.Vector v Double) => v CD -> v Double
idct_ = idctWorker . G.map (\(i :+ _) -> i :+ 0)
{-# INLINABLE  idct_ #-}
{-# SPECIAlIZE idct_ :: U.Vector CD -> U.Vector Double #-}
{-# SPECIAlIZE idct_ :: V.Vector CD -> V.Vector Double #-}

idctWorker :: (G.Vector v CD, G.Vector v Double) => v CD -> v Double
{-# INLINE idctWorker #-}
idctWorker xs
  | vectorOK xs = G.generate len interleave
  | otherwise   = error "Statistics.Transform.dct: bad vector length"
  where
    interleave z | even z    = vals `G.unsafeIndex` halve z
                 | otherwise = vals `G.unsafeIndex` (len - halve z - 1)
    vals = G.map realPart . ifft $ G.zipWith (*) weights xs
    weights
      = G.cons n
      $ G.generate (len - 1) $ \x -> 2 * n * exp ((0:+1) * fi (x+1) * pi/(2*n))
      where n = fi len
    len = G.length xs

-- | Inverse fast Fourier transform.
ifft :: G.Vector v CD => v CD -> v CD
ifft xs
  | vectorOK xs = G.map ((/fi (G.length xs)) . conjugate) . fft . G.map conjugate $ xs
  | otherwise   = error "Statistics.Transform.ifft: bad vector length"
{-# INLINABLE  ifft #-}
{-# SPECIAlIZE ifft :: U.Vector CD -> U.Vector CD #-}
{-# SPECIAlIZE ifft :: V.Vector CD -> V.Vector CD #-}

-- | Radix-2 decimation-in-time fast Fourier transform.
fft :: G.Vector v CD => v CD -> v CD
fft v | vectorOK v  = G.create $ do mv <- G.thaw v
                                    mfft mv
                                    return mv
      | otherwise   = error "Statistics.Transform.fft: bad vector length"
{-# INLINABLE  fft #-}
{-# SPECIAlIZE fft :: U.Vector CD -> U.Vector CD #-}
{-# SPECIAlIZE fft :: V.Vector CD -> V.Vector CD #-}

-- Vector length must be power of two. It's not checked
mfft :: (M.MVector v CD) => v s CD -> ST s ()
{-# INLINE mfft #-}
mfft vec = bitReverse 0 0
 where
  bitReverse i j | i == len-1 = stage 0 1
                 | otherwise  = do
    when (i < j) $ M.swap vec i j
    let inner k l | k <= l    = inner (k `shiftR` 1) (l-k)
                  | otherwise = bitReverse (i+1) (l+k)
    inner (len `shiftR` 1) j
  stage l !l1 | l == m    = return ()
              | otherwise = do
    let !l2 = l1 `shiftL` 1
        !e  = -6.283185307179586/fromIntegral l2
        flight j !a | j == l1   = stage (l+1) l2
                    | otherwise = do
          let butterfly i | i >= len  = flight (j+1) (a+e)
                          | otherwise = do
                let i1 = i + l1
                xi1 :+ yi1 <- M.read vec i1
                let !c = cos a
                    !s = sin a
                    d  = (c*xi1 - s*yi1) :+ (s*xi1 + c*yi1)
                ci <- M.read vec i
                M.write vec i1 (ci - d)
                M.write vec i (ci + d)
                butterfly (i+l2)
          butterfly j
    flight 0 0
  len = M.length vec
  m   = log2 len

----------------------------------------------------------------
-- Helpers
----------------------------------------------------------------

fi :: Int -> CD
fi = fromIntegral

halve :: Int -> Int
halve = (`shiftR` 1)

vectorOK :: G.Vector v a => v a -> Bool
{-# INLINE vectorOK #-}
vectorOK v = (1 `shiftL` log2 n) == n where n = G.length v


{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeApplications #-}

-----------------------------------------------------------------------------
{-|
Module      : Math.Tensor.LinearAlgebra.Equations
Description : Linear tensor equations.
Copyright   : (c) Nils Alex, 2020
License     : MIT
Maintainer  : nils.alex@fau.de

Linear tensor equations.
-}
-----------------------------------------------------------------------------
module Math.Tensor.LinearAlgebra.Equations
  ( Equation
  , tensorToEquations
  , equationFromRational
  , equationsToSparseMat
  , equationsToMat
  , tensorsToSparseMat
  , tensorsToMat
  , systemRank
  , Solution
  , fromRref
  , fromRrefRev
  , fromRow
  , fromRowRev
  , applySolution
  , solveTensor
  , solveSystem
  , redefineIndets
  ) where

import Math.Tensor
  ( T
  , removeZerosT
  , toListT
  )
import Math.Tensor.LinearAlgebra.Scalar
  ( Poly (Const, Affine, NotSupported)
  , Lin (Lin)
  , polyMap
  , normalize
  )
import Math.Tensor.LinearAlgebra.Matrix
  ( rref
  , isrref
  , verify
  )

import qualified Numeric.LinearAlgebra.Data as HM
  ( Matrix
  , R
  , Z
  , fromLists
  , toLists
  )
import Numeric.LinearAlgebra (rank)

import Data.Maybe (mapMaybe)
import qualified Data.IntMap.Strict as IM
  ( IntMap
  , foldl'
  , map
  , assocs
  , null
  , findWithDefault
  , lookupMax
  , keys
  , fromList
  , (!)
  , difference
  , intersectionWith
  , mapKeys
  , empty
  )
import Data.List (nub, sort)
import Data.Ratio (numerator, denominator)

-- |A linear equation is a mapping from variable
-- indices to coefficients
type Equation a = IM.IntMap a

-- |Extract linear equations from tensor components.
-- The equations are normalized, sorted, and made unique.
tensorToEquations :: Integral a => T (Poly Rational) -> [Equation a]
tensorToEquations = nub . sort . filter (not . IM.null) . fmap (equationFromRational . normalize . snd) . toListT

-- |Extract linear equation with integral coefficients from polynomial
-- tensor component with rational coefficients.
-- Made made integral by multiplying with the @lcm@ of all denominators.
equationFromRational :: forall a.Integral a => Poly Rational -> Equation a
equationFromRational (Affine x (Lin lin))
    | x == 0 = lin'
    | otherwise = error "affine equation not supported for the moment!"
  where
    fac :: a
    fac = IM.foldl' (\acc v -> lcm (fromIntegral (denominator v)) acc) 1 lin
    lin' = IM.map (\v -> fromIntegral (numerator v) * (fac `div` fromIntegral (denominator v))) lin
equationFromRational (Const c)
    | c == 0    = IM.empty
equationFromRational _ = error "equation can only be extracted from linear scalar!"

-- |Convert list of equations to sparse matrix representation of the
-- linear system.
equationsToSparseMat :: [Equation a] -> [((Int,Int), a)]
equationsToSparseMat xs = concat $ zipWith (\i m -> fmap (\(j,v) -> ((i,j),v)) (IM.assocs m)) [1..] xs

-- |Convert list of equations to dense matrix representation of the
-- linear system.
equationsToMat :: Integral a => [Equation a] -> [[a]]
equationsToMat eqns = mapMaybe (\m -> if IM.null m
                                      then Nothing
                                      else Just $ fmap (\j -> IM.findWithDefault 0 j m) [1..maxVar]) eqns
  where
    maxVar = maximum $ mapMaybe (fmap fst . IM.lookupMax) eqns

-- |Extract sparse matrix representation for the linear system given
-- by a list of existentially quantified tensors with polynomial values.
tensorsToSparseMat :: Integral a => [T (Poly Rational)] -> [((Int,Int), a)]
tensorsToSparseMat = equationsToSparseMat . concatMap tensorToEquations

-- |Extract dense matrix representation for the linear system given
-- by a list of existentially quantified tensors with polynomial values.
tensorsToMat :: Integral a => [T (Poly Rational)] -> [[a]]
tensorsToMat = equationsToMat . concatMap tensorToEquations

-- |Rank of a linear system. Uses dense svd provided by hmatrix.
matRank :: forall a.Integral a => [[a]] -> Int
matRank []  = 0
matRank mat = rank (hmat :: HM.Matrix HM.R)
  where
    hmat = HM.fromLists $ fmap (fmap (fromIntegral @a @HM.R)) mat

-- |Rank of the linear system given by a list of existentially
-- quantified tensors with polynomial values.
systemRank :: [T (Poly Rational)] -> Int
systemRank = matRank . tensorsToMat @Int

-- |The solution to a linear system is represented as a list of
-- substitution rules, stored as @'IM.IntMap' ('Poly' 'Rational')@.
type Solution = IM.IntMap (Poly Rational)

-- |Read substitution rules from reduced row echelon form
-- of a linear system.
fromRref :: HM.Matrix HM.Z -> Solution
fromRref ref = IM.fromList assocs
  where
    rows   = HM.toLists ref
    assocs = mapMaybe fromRow rows

fromRrefRev :: HM.Matrix HM.Z -> Solution
fromRrefRev ref = IM.fromList assocs
  where
    rows   = fmap reverse $ HM.toLists ref
    assocs = mapMaybe fromRowRev rows

-- |Read single substitution rule from single
-- row of reduced row echelon form.
fromRow :: forall a.Integral a => [a] -> Maybe (Int, Poly Rational)
fromRow xs = case assocs of
               []             -> Nothing
               [(i,_)]        -> Just (i, Const 0)
               (i, v):assocs' -> let assocs'' = fmap (\(i',v') -> (i', - fromIntegral @a @Rational v' / fromIntegral @a @Rational v)) assocs'
                                 in Just (i, Affine 0 (Lin (IM.fromList assocs'')))
  where
    assocs = filter ((/=0). snd) $ zip [(1::Int)..] xs

fromRowRev :: forall a.Integral a => [a] -> Maybe (Int, Poly Rational)
fromRowRev xs = case assocs of
                  []             -> Nothing
                  [(i,_)]        -> Just (i, Const 0)
                  (i, v):assocs' -> let assocs'' = fmap (\(i',v') -> (i', - fromIntegral @a @Rational v' / fromIntegral @a @Rational v)) assocs'
                                    in Just (i, Affine 0 (Lin (IM.fromList assocs'')))
  where
    assocs = reverse $ filter ((/=0). snd) $ zip [(1::Int)..] xs

-- |Apply substitution rules to tensor component.
applySolution :: Solution -> Poly Rational -> Poly Rational
applySolution s (Affine x (Lin lin))
    | x == 0 = case p of
                 Affine xFin (Lin linFin) -> if IM.null linFin
                                             then Const xFin
                                             else p
                 _ -> p
    | otherwise = error "affine equations not yet supported"
  where
    s' = IM.intersectionWith (\row v -> if v == 0
                                        then error "value 0 encountered in linear scalar"
                                        else polyMap (v*) row) s lin
    lin' = IM.difference lin s
    p0 = if IM.null lin'
         then Const 0
         else Affine 0 (Lin lin')

    p = IM.foldl' (+) p0 s'

    {-
    lin' = IM.foldlWithKey'
             (\lin' i sub ->
                 case Affine 0 (Lin lin') + sub of
                   Affine 0 (Lin lin'') -> IM.delete i lin''
                   Const 0              -> IM.empty
                   _                    -> error "affine equations not yet supported")
             lin s'
    -}
applySolution _ _ = error "only linear equations supported"

-- |Apply substitution rules to all components of a tensor.
solveTensor :: Solution -> T (Poly Rational) -> T (Poly Rational)
solveTensor sol = removeZerosT . fmap (applySolution sol)

-- |Solve a linear system and apply solution to the tensorial
-- indeterminants.
solveSystem ::
     [T (Poly Rational)] -- ^ Tensorial linear system
  -> [T (Poly Rational)] -- ^ List of indeterminant tensors
  -> [T (Poly Rational)] -- ^ Solved indeterminant tensors
solveSystem system indets
    | wrongSolution = error "Wrong solution found. May be an Int64 overflow."
    | otherwise     = indets'
  where
    mat = HM.fromLists $ tensorsToMat @HM.Z system
    ref = rref mat
    wrongSolution = not (isrref ref && verify mat ref)
    sol = fromRref ref
    indets' = fmap (solveTensor sol) indets

-- |Relabelling of the indeterminants present in a list of tensors.
-- Redefines the labels of @n@ indeterminants as @[1..n]@, preserving
-- the previous order.
redefineIndets :: [T (Poly v)] -> [T (Poly v)]
redefineIndets indets = fmap (fmap (\case
                                       Const c -> Const c
                                       NotSupported -> NotSupported
                                       Affine a (Lin lin) ->
                                         Affine a (Lin (IM.mapKeys (varMap IM.!) lin)))) indets
  where
    comps = snd <$> concatMap toListT indets
    vars = nub $ concat $ mapMaybe (\case
                                       Affine _ (Lin lin) -> Just $ IM.keys lin
                                       _                  -> Nothing) comps
    varMap = IM.fromList $ zip vars [(1::Int)..]


module AoC2021.Day07 where

import           Data.Bifunctor               (bimap)
import           Data.List                    (sort)
import qualified Numeric.LinearAlgebra.Static as N
import qualified Parsers                      as P
import           Solver                       (Solver)

mean :: (Fractional n, Foldable f) => f n -> n
mean ns
  | null ns = 0
  | otherwise = uncurry (/) . foldl f (0,0) $ ns
  where f (t,c) n = (t+n,c+1)

median :: Ord a => [a] -> a
median as = let m = length as `div` 2
            in  sort as!!m

minFuelCost :: [Int] -> Int
minFuelCost ns = sum . map (abs . (median ns-)) $ ns

day07 :: ([Int] -> Int) -> Solver
day07 algo input = case P.runParser P.parseIntegerList () "" input of
  Left e          -> show e
  Right positions -> show . algo $ positions

solve01 :: Solver
solve01 = day07 minFuelCost


{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
-- This module uses the test-framework-quickcheck2 package.
module Main where

import Control.Monad
import qualified Data.Vector.Unboxed as V
import qualified Data.Vector.Storable as VS
import Data.Complex

import Test.Framework (defaultMain, testGroup)
import Test.Framework.Providers.QuickCheck2 (testProperty)
import Test.QuickCheck

import qualified Numeric.FFT.Vector.Invertible as I
import qualified Numeric.FFT.Vector.Invertible.Multi as IM
import qualified Numeric.FFT.Vector.Unitary as U
import qualified Numeric.FFT.Vector.Unitary.Multi as UM
import Numeric.FFT.Vector.Plan

main = defaultMain
            -- NB: There's no explicit tests for the Unnormalized package.
            -- However, its Planners are implicitly used by the other modules,
            -- so it's covered in the below tests.
            [ testGroup "invertibility"
              [ testProperty "I.dft" $ prop_invert I.dft I.idft
              , testProperty "I.dftR2C" $ prop_invert I.dftR2C I.dftC2R
              , testProperty "I.dct1" $ prop_invert I.dct1 I.idct1
              , testProperty "I.dct2" $ prop_invert I.dct2 I.idct2
              , testProperty "I.dct3" $ prop_invert I.dct3 I.idct3
              , testProperty "I.dct4" $ prop_invert I.dct4 I.idct4
              , testProperty "I.dst1" $ prop_invert I.dst1 I.idst1
              , testProperty "I.dst2" $ prop_invert I.dst2 I.idst2
              , testProperty "I.dst3" $ prop_invert I.dst3 I.idst3
              , testProperty "I.dst4" $ prop_invert I.dst4 I.idst4
              , testProperty "U.dft" $ prop_invert U.dft U.idft
              , testProperty "U.dftR2C" $ prop_invert U.dftR2C U.dftC2R
              , testProperty "U.dct2" $ prop_invert U.dct2 U.idct2
              ]
            , testGroup "orthogonality"
              [ testProperty "U.dft" $ prop_orthog U.dft
              , testProperty "U.idft" $ prop_orthog U.idft
              , testProperty "U.dftR2C" $ prop_orthog U.dftR2C
              , testProperty "U.dftC2R" $ prop_orthog U.dftR2C
              , testProperty "U.dct2" $ prop_orthog U.dct2
              , testProperty "U.idct2" $ prop_orthog U.idct2
              , testProperty "U.dct4" $ prop_orthog U.dct4
              ]
            , testGroup "invertibility ND"
              [ testProperty "IM.dft" $ prop_invertND IM.dft IM.idft
              , testProperty "IM.dftR2C" $ prop_invertND IM.dftR2C IM.dftC2R
              , testProperty "UM.dft" $ prop_invertND UM.dft UM.idft
              , testProperty "UM.dftR2C" $ prop_invertND UM.dftR2C UM.dftC2R
              ]
            , testGroup "orthogonality"
              [ testProperty "UM.dft" $ prop_orthogND UM.dft
              , testProperty "UM.idft" $ prop_orthogND UM.idft
              , testProperty "UM.dftR2C" $ prop_orthogND UM.dftR2C
              , testProperty "UM.dftC2R" $ prop_orthogND UM.dftR2C
              ]
            ]

-------------------
-- An instance of Arbitrary that probably belongs in another package.

instance (V.Unbox a, Arbitrary a) => Arbitrary (V.Vector a) where
    arbitrary = V.fromList `fmap` arbitrary

-------------------------
-- Support functions to compare Doubles for (near) equality.

class Num a => Mag a where
    mag :: a -> Double

instance Mag Double where
    mag = abs

instance Mag (Complex Double) where
    mag = magnitude

-- Robustly test whether two Doubles are nearly identical.
close :: Mag a => a -> a -> Bool
close x y = tol > mag (x-y) / max 1 (mag x + mag y)
  where
    tol = 1e-10

withinTol :: (Mag a, V.Unbox a) => V.Vector a -> V.Vector a -> Bool
withinTol a b
    | V.length a /= V.length b = False
    | otherwise = V.and $ V.zipWith close a b

---------------------
-- The actual properties

-- Test whether the inverse actually inverts the forward transform.
prop_invert f g a = let
                        p1 = plan f (V.length a)
                        p2 = plan g (V.length a)
                    in (V.length a > 1) ==> withinTol a $ execute p2 $ execute p1 a

-- Test whether the transform preserves the L2 (sum-of-squares) norm.
prop_orthog f a = let
                    p1 = plan f (V.length a)
                  in (V.length a > 1) ==> close (norm2 a) (norm2 $ execute p1 a)

data DimsAndValues a = DimsAndValues (VS.Vector Int) (V.Vector a)
  deriving (Show)

instance (Arbitrary a, V.Unbox a) => Arbitrary (DimsAndValues a) where
  arbitrary = do
    dims <- liftM (VS.fromList . map getPositive) arbitrary `suchThatMap` maybeReduceSize
    values <- V.replicateM (VS.product dims) arbitrary
    return (DimsAndValues dims values)
    where
      -- We use this to prevent test cases from growing too big
      maybeReduceSize ds =
        if VS.product ds < 1000 then Just ds else maybeReduceSize (VS.init ds)

prop_invertND f g (DimsAndValues ds a) = let
                        p1 = planND f ds
                        p2 = planND g ds
                    in (V.length a > 1) ==> withinTol a $ execute p2 $ execute p1 a

prop_orthogND f (DimsAndValues ds a) = let
                    p1 = planND f ds
                  in (V.length a > 1) ==> close (norm2 a) (norm2 $ execute p1 a)

norm2 a = sqrt $ V.sum $ V.map (\x -> x*x) $ V.map mag a


{-# LANGUAGE FlexibleContexts          #-}
{-# LANGUAGE FlexibleInstances         #-}
{-# LANGUAGE GADTs                     #-}
{-# LANGUAGE MultiParamTypeClasses     #-}
{-# LANGUAGE NoMonomorphismRestriction #-}
{-# LANGUAGE OverloadedStrings         #-}
{-# LANGUAGE PolyKinds                 #-}
{-# LANGUAGE RankNTypes                #-}
{-# LANGUAGE ScopedTypeVariables       #-}
{-# OPTIONS_GHC -fwarn-incomplete-patterns #-}

module Math.Regression.Regression where

import qualified Math.HMatrixUtils             as HU

import qualified Colonnade                     as C

import qualified Data.List                     as List
import qualified Data.Text                     as T
import qualified Data.Vector.Storable          as V
import qualified Lucid                         as LH
import qualified Lucid.Colonnade               as LC
import qualified Statistics.Distribution       as S
import qualified Statistics.Distribution.FDistribution
                                               as S
import qualified Statistics.Distribution.StudentT
                                               as S
import qualified Statistics.Types              as S
import qualified Text.Blaze.Colonnade          as BC
import qualified Text.Blaze.Html5              as BH
import qualified Text.Blaze.Html5.Attributes   as BHA
import qualified Text.Printf                   as TP

import           Numeric.LinearAlgebra          ( (#>)
                                                , (<.>)
                                                )
import qualified Numeric.LinearAlgebra         as LA
import           Numeric.LinearAlgebra.Data     ( Matrix
                                                , R
                                                , Vector
                                                )
import qualified Say
-- NB: for normal errors, we will request the ci and get it.  For interval errors, we may request it but we can't produce
-- what we request, we can only produce what we have.  So we make an interface for both.  We take CL as input to predict
-- *and* return it as output, in the prediction.

class Predictor e a b p where
  predict :: p -> b -> S.Estimate e a -- do we need a predict :: (b, b) -> (a, a) for uncertain inputs?

predictFromEstimateAtConfidence
  :: (RealFrac a, Predictor S.NormalErr a b p)
  => Double
  -> p
  -> S.CL Double
  -> b
  -> (a, a)
predictFromEstimateAtConfidence dof p cl b =
  let
    S.Estimate pt (S.NormalErr sigma) = predict p b
    prob =
      S.confidenceLevel $ S.mkCLFromSignificance (S.significanceLevel cl / 2)
    predCI =
      2 * S.quantile (S.studentTUnstandardized dof 0 (realToFrac sigma)) prob
  in
    (pt, realToFrac predCI)

data RegressionResult a = RegressionResult
                          {
                            parameterEstimates :: [S.Estimate S.NormalErr a]
                          , degreesOfFreedom   :: Double -- since N may be an effective N
                          , meanSquaredError   :: a
                          , rSquared           :: Double
                          , adjRSquared        :: Double
                          , fStatistic         :: Maybe Double -- we can't compute this for fits with one parameter
                          , covariances        :: Matrix Double
                          } deriving (Show)

instance (LA.Element a, LA.Numeric a, RealFloat a) => Predictor S.NormalErr a (Vector a) (RegressionResult a) where
  predict rr va =
    let mse = meanSquaredError rr
        sigmaPred = sqrt (mse + va <.> (LA.cmap realToFrac (covariances rr) #> va))
        prediction = va LA.<.> V.fromList (S.estPoint <$> parameterEstimates rr)
    in S.estimateNormErr prediction sigmaPred

data NamedEstimate a = NamedEstimate { regressorName :: T.Text, regressorEstimate :: a, regressorCI :: a, regressorPValue :: a }

namedEstimates :: [T.Text] -> RegressionResult R -> S.CL R -> [NamedEstimate R]
namedEstimates pNames res cl =
  let prob = S.confidenceLevel
        $ S.mkCLFromSignificance (S.significanceLevel cl / 2)
      dof = degreesOfFreedom res
      sigma e = S.normalError $ S.estError e
      dist e = S.studentTUnstandardized dof 0 (sigma e)
      ci e = 2 * S.quantile (dist e) prob
      pValue e = S.complCumulative (dist e) (abs $ S.estPoint e)
  in  List.zipWith (\n e -> NamedEstimate n (S.estPoint e) (ci e) (pValue e))
                   pNames
                   (parameterEstimates res)

namedEstimatesColonnade
  :: S.CL R -> C.Colonnade C.Headed (NamedEstimate R) T.Text
namedEstimatesColonnade cl =
  C.headed "parameter" regressorName
    <> C.headed "estimate" (T.pack . TP.printf "%4.3f" . regressorEstimate)
    <> C.headed
         (  T.pack (TP.printf "%.0f" (100 * S.confidenceLevel cl))
         <> "% confidence"
         )
         (T.pack . TP.printf "%4.3f" . regressorCI)
    <> C.headed "p-value" (T.pack . TP.printf "%4.3f" . regressorPValue)

namedSummaryStats :: RegressionResult R -> [(T.Text, T.Text)]
namedSummaryStats r =
  let
    p             = List.length (parameterEstimates r)
    effN          = degreesOfFreedom r + realToFrac p
    d1            = realToFrac $ p - 1
    d2            = effN - realToFrac p
    pValM = fmap (S.complCumulative (S.fDistributionReal d1 d2)) (fStatistic r)
    printNum      = T.pack . TP.printf "%4.3f"
    printMaybeNum = maybe "N/A" printNum
  in
    [ ("R-Squared"         , printNum (rSquared r))
    , ("Adj. R-squared"    , printNum (adjRSquared r))
    , ("F-stat"            , printMaybeNum (fStatistic r))
    , ("p-value"           , printMaybeNum pValM)
    , ("Mean Squared Error", printNum (meanSquaredError r))
    ]

namedSummaryStatsColonnade :: C.Colonnade C.Headed (T.Text, T.Text) T.Text
namedSummaryStatsColonnade =
  C.headed "Summary Stat." fst <> C.headed "Value" snd

prettyPrintRegressionResult
  :: T.Text -> [T.Text] -> RegressionResult R -> S.CL R -> T.Text
prettyPrintRegressionResult header xNames r cl =
  let nEsts = namedEstimates xNames r cl
      nSS   = namedSummaryStats r
  in  header
      <> "\n"
      <> T.pack (C.ascii (T.unpack <$> namedEstimatesColonnade cl) nEsts)
      <> T.pack (C.ascii (fmap T.unpack namedSummaryStatsColonnade) nSS)

prettyPrintRegressionResultLucid
  :: T.Text -> [T.Text] -> RegressionResult R -> S.CL R -> LH.Html ()
prettyPrintRegressionResultLucid header xNames r cl = do
  let nEsts = namedEstimates xNames r cl
      nSS   = namedSummaryStats r
      toCell t = LC.Cell [LH.style_ "border: 1px solid black"] (LH.toHtmlRaw t)
  LH.div_
      [ LH.style_
          "display: inline-block; padding: 7px; border-collapse: collapse"
      ]
    $ do
        LH.span_ (LH.toHtmlRaw header)
        LC.encodeCellTable
          [LH.style_ "border: 1px solid black; border-collapse: collapse"]
          (toCell <$> namedEstimatesColonnade cl)
          nEsts
        LC.encodeCellTable
          [LH.style_ "border: 1px solid black; border-collapse: collapse"]
          (fmap toCell namedSummaryStatsColonnade)
          nSS

prettyPrintRegressionResultBlaze
  :: T.Text -> [T.Text] -> RegressionResult R -> S.CL R -> BH.Html
prettyPrintRegressionResultBlaze header xNames r cl = do
  let nEsts = namedEstimates xNames r cl
      nSS   = namedSummaryStats r
      toCell t = BC.Cell (BHA.style "border: 1px solid black") (BH.toHtml t)
  BH.div
    BH.! BHA.style
           "display: inline-block; padding: 7px; border-collapse: collapse"
    $    do
           BH.span (BH.toHtml header)
           BC.encodeCellTable
             (BHA.style "border: 1px solid black; border-collapse: collapse")
             (toCell <$> namedEstimatesColonnade cl)
             nEsts
           BC.encodeCellTable
             (BHA.style "border: 1px solid black; border-collapse: collapse")
             (fmap toCell namedSummaryStatsColonnade)
             nSS

data FitStatistics a = FitStatistics { fsRSquared :: a, fsAdjRSquared :: a, fsFStatistic :: Maybe a}

goodnessOfFit
  :: MonadIO m
  => Int
  -> Vector R
  -> Maybe (Vector R)
  -> Vector R
  -> m (FitStatistics R)
goodnessOfFit pInt vB vWM vU = do
  let
    n  = LA.size vB
    p  = realToFrac pInt
    vW = maybe (V.fromList $ List.replicate n (1.0 / realToFrac n))
               (\v -> LA.scale (1 / realToFrac (LA.sumElements v)) v)
               vWM
    mW     = LA.diag vW
    vWB    = mW #> vB
--      vWU = mW #> vU
    meanWB = LA.sumElements vWB
    effN   = 1 / (vW <.> vW)
    ssTot  = let x = LA.cmap (\y -> y - meanWB) vB in x <.> (mW #> x)
    ssRes  = vU <.> (mW #> vU) -- weighted ssq of residuals
    rSq    = 1 - (ssRes / ssTot)
    arSq =
      1 - (1 - rSq) * realToFrac (effN - 1) / realToFrac (effN - p - 1)
    fStatM = if pInt > 1
      then Just (((ssTot - ssRes) / (p - 1.0)) / (ssRes / (effN - p)))
      else Nothing
  Say.say $ "n=" <> show n
  Say.say $ "p=" <> show p
--  Log.log Log.Diagnostic $ "vW=" <> show vW
  Say.say $ "effN=" <> show effN
  Say.say $ "ssTot=" <> show ssTot
  Say.say $ "ssRes=" <> show ssRes
  return $ FitStatistics rSq arSq fStatM

estimates :: Matrix R -> Vector R -> [S.Estimate S.NormalErr R]
estimates cov means =
  let sigmas = LA.cmap sqrt (LA.takeDiag cov)
  in  List.zipWith S.estimateNormErr
                   (LA.toList means)
                   (LA.toList sigmas)

eickerHeteroscedasticityEstimator
  :: MonadIO m
  => Matrix R
  -> Vector R
  -> Vector R
  -> m (Matrix R)
eickerHeteroscedasticityEstimator mA vB vB' = do
  HU.checkVectorMatrix "b" "A" vB mA
  HU.checkVectorMatrix "b'" "A" vB' mA
  let u     = vB - vB'
      diagU = LA.diag $ V.zipWith (*) u u
      mA'   = LA.tr mA
      mC    = LA.inv (mA' LA.<> mA)
  return $ mC LA.<> (mA' LA.<> diagU LA.<> mA) LA.<> mC


module Model where

import Control.Monad.State.Strict
--import Control.Monad.State
import System.Random (StdGen, mkStdGen, randomRs)
import Numeric.LinearAlgebra (Matrix(..), (><), toLists)

type Ising = State IsingState
type Spin = Double
type Model = Matrix Spin

data IsingState = IsingState
     { dim :: Int
     , j :: Double
     , h :: Double
     , step :: Int
     , nAccept :: Int
     , propFreq :: Int
     , model :: Model
     , rng :: StdGen
     }

instance Show IsingState where
    show state =
        "IsingState { d = " ++ show (dim state)
            ++ ", j = " ++ show (j state)
            ++ ", step = " ++ show (step state)
            ++ ", accept = " ++ show (nAccept state)
            ++ ", rng = " ++ show (rng state)
            ++ " }"
            ++ showModel (model state)

showModel :: Model -> String
showModel m =
    let rows = toLists m
        showRow r = "|" ++ foldl (\acc s -> acc ++ showSpin s) "" r ++ "|"
        showSpin s = if s == -1 then " " else "*"
        stringRows = map showRow rows
    in foldl (\acc r -> acc ++ "\n" ++ r) "" stringRows

newModel :: Int -> Int -> Double -> [Int] -> IsingState
newModel seed n j spins =
    let rng = mkStdGen seed
        model = (n><n) $ map fromIntegral spins
    in IsingState
        { dim = n
        , j = j
        , h = 0.0
        , step = 0
        , nAccept = 0
        , propFreq = 1000
        , model = model
        , rng = rng
        }

randModel seed n j =
    let r = mkStdGen seed
        spins = take (n*n) $ randomRs (0,1) r :: [Int]
    in newModel seed n j [if s == 0 then -1 else s | s <- spins]

upSpins   n = [ 1 | _ <- [1..n*n]]
downSpins n = [-1 | _ <- [1..n*n]]

setPropertyFreq :: Int -> IsingState -> IsingState
setPropertyFreq f model = model { propFreq = f }

setH :: Double -> IsingState -> IsingState
setH h model = model { h = h }


module Cas.QuickCheckTest where

import Test.QuickCheck
import Test.QuickCheck.Gen
import Test.QuickCheck.Random
import Control.Monad
import Data.Ratio
import PreludeCustom
import Cas.Interface
import Data.Complex.Generic

instance Arbitrary C where
    arbitrary = oneof [return Pi, return E]

instance Arbitrary T where
    arbitrary = sized genT

chooseQ :: Gen Rational
chooseQ = liftM2 (%) (choose (-1,1))  (choose (1,2))

maxTestCnt = 20000
maxDepth =13
scaleFactor = 2 -- maxTestCnt `div` 2^maxDepth + 1

genT, scaledGenT :: Int -> Gen T
genT n = scaledGenT $ n `div` scaleFactor

scaledGenT 0 = oneof [
     liftM f $ liftM2 (:+) chooseQ chooseQ
   , liftM x $ choose (0,1)
   , liftM TC arbitrary
   ]
scaledGenT n = do
    k <- choose (2,5)
    oneof [
          liftM add $ vectorOf k rec
        , liftM mul $ vectorOf k rec
        , liftM2 pow rec rec
        , liftM ln rec
        ]
   where
   rec =  scaledGenT =<< choose (0, n `div` 2)


{-# LANGUAGE ParallelListComp #-}
import Numeric.LinearAlgebra (linearSolve, toDense, (!), flatten)
import Data.Monoid ((<>), Sum(..))

rMesh n (ar, ac) (br, bc)
  | n < 2 = Nothing
  | any (\x -> x < 1 || x > n) [ar, ac, br, bc] = Nothing
  | otherwise = between a b <$> voltage
  where
    a = (ac - 1) + n*(ar - 1)
    b = (bc - 1) + n*(br - 1)

    between x y v = abs (v ! a - v ! b)

    voltage = flatten <$> linearSolve matrixG current

    matrixG = toDense $ concat [ element row col node
                               | row <- [1..n], col <- [1..n]
                               | node <- [0..] ]

    element row col node =
      let (Sum c, elements) =
            (Sum 1, [((node, node-n), -1)]) `when` (row > 1) <>
            (Sum 1, [((node, node+n), -1)]) `when` (row < n) <>
            (Sum 1, [((node, node-1), -1)]) `when` (col > 1) <>
            (Sum 1, [((node, node+1), -1)]) `when` (col < n)
      in [((node, node), c)] <> elements

    x `when` p = if p then x else mempty

    current  = toDense [ ((a, 0), -1) , ((b, 0),  1) , ((n^2-1, 0), 0) ]


{-# LANGUAGE NamedFieldPuns, TemplateHaskell #-}

module School.Unit.Test.Affine
( affineTest ) where

import Control.Monad.IO.Class (liftIO)
import Numeric.LinearAlgebra ((><), (|>), IndexOf, Matrix, R, Vector, accum,
                              fromRows, ident, size)
import School.TestUtils (CostFunc, diffInput, jTest, matIndexes,
                         randomAffineParams, randomMatrix)
import School.Types.FloatEq ((~=))
import School.Unit.Affine
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, compareDoubleVector,
                                   oneMatrix, zeroMatrix, zeroVector)
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 = 5e-2

prop_trivial_apply :: (Positive Int) -> (Positive Int) -> Property
prop_trivial_apply (Positive bSize) (Positive fSize) = monadicIO $ do
  let affineBias = zeroVector fSize
  m <- liftIO $ randomMatrix bSize fSize
  let input = BatchActivation m
  let affineWeights = ident fSize
  let params = AffineParams { affineBias
                            , affineWeights }
  let result = apply affine params input
  assert $ result ~= input

randomSetup :: Int
            -> Int
            -> Int
            -> IO (UnitParams R, UnitActivation R)
randomSetup bSize fSize oSize = do
  input <- randomMatrix bSize fSize
  affineParams <- randomAffineParams fSize oSize
  return ( affineParams
         , BatchActivation input
         )

prop_overwrite_apply :: (Positive Int) -> (Positive Int) -> Positive Int -> Property
prop_overwrite_apply (Positive bSize) (Positive fSize) (Positive oSize) = monadicIO $ do
  let affineWeights = zeroMatrix oSize fSize
  (params, input) <- liftIO $ randomSetup bSize fSize oSize
  let params' = params { affineWeights }
  let result = apply affine params' input
  let check = fromRows . (replicate bSize) $ affineBias params
  assert $ result ~= BatchActivation check

prop_deriv_dims :: (Positive Int) -> (Positive Int) -> Positive Int -> Property
prop_deriv_dims (Positive bSize) (Positive fSize) (Positive oSize) = monadicIO $ do
  (params, input) <- liftIO $ randomSetup bSize fSize oSize
  gradIn <- liftIO $ fmap BatchGradient $ randomMatrix bSize oSize
  let (BatchGradient gradOut, AffineParams{ affineBias, affineWeights }) = deriv affine params gradIn input
  assert $ (size gradOut == (bSize, fSize))
        && (size affineBias == oSize)
        && (size affineWeights == (oSize, fSize))

type AlterParams a = Double
                  -> IndexOf a
                  -> UnitParams R
                  -> UnitParams R

diffBias :: UnitParams R
          -> UnitActivation R
          -> CostFunc
          -> IndexOf Vector
          -> Double
diffBias params input cost idx =
  (jAdd - jSub) / (2*eps) where
    outAdd = apply affine (alterBias eps idx params) input
    outSub = apply affine (alterBias (-eps) idx params) input
    jAdd = cost outAdd
    jSub = cost outSub

diffWeights :: UnitParams R
            -> UnitActivation R
            -> CostFunc
            -> IndexOf Matrix
            -> Double
diffWeights params input cost idx =
  (jAdd - jSub) / (2*eps) where
    outAdd = apply affine (alterWeights eps idx params) input
    outSub = apply affine (alterWeights (-eps) idx params) input
    jAdd = cost outAdd
    jSub = cost outSub

alterBias :: AlterParams Vector
alterBias change idx AffineParams{ affineBias, affineWeights } =
  let newBias = accum affineBias (+) [(idx, change)]
  in AffineParams { affineBias = newBias, affineWeights }
alterBias _ _ _ = EmptyParams

alterWeights :: AlterParams Matrix
alterWeights change idx AffineParams{ affineBias, affineWeights } =
  let newWeights = accum affineWeights (+) [(idx, change)]
  in AffineParams { affineBias, affineWeights = newWeights }
alterWeights _ _ _ = EmptyParams

vecIndexes :: Int -> [IndexOf Vector]
vecIndexes n = [0..n-1]

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

prop_numerical_bias_deriv :: (Positive Int) -> (Positive Int) -> Positive Int -> Property
prop_numerical_bias_deriv (Positive bSize) (Positive fSize) (Positive oSize) = (mapSize (const 10)) . monadicIO $ do
  pre $ bSize < 50 && fSize < 50 && oSize < 50
  (params, input) <- liftIO $ randomSetup bSize fSize oSize
  let idxs = vecIndexes oSize
  let num = map (\idx -> diffBias params input jTest idx)
                idxs
  let numDeriv = oSize |> num
  let inGrad = BatchGradient $ oneMatrix bSize oSize
  let (_, AffineParams { affineBias }) = deriv affine params inGrad input
  assert $ compareDoubleVector prec affineBias numDeriv

prop_numerical_weights_deriv :: (Positive Int) -> (Positive Int) -> Positive Int -> Property
prop_numerical_weights_deriv (Positive bSize) (Positive fSize) (Positive oSize) = (mapSize (const 10)) . monadicIO $ do
  pre $ bSize < 50 && fSize < 50 && oSize < 50
  (params, input) <- liftIO $ randomSetup bSize fSize oSize
  let idxs = matIndexes oSize fSize
  let num = map (\idx -> diffWeights params input jTest idx)
                idxs
  let numDeriv = (oSize >< fSize) num
  let inGrad = BatchGradient $ oneMatrix bSize oSize
  let (_, AffineParams { affineWeights }) = deriv affine params inGrad input
  assert $ compareDoubleMatrix prec affineWeights numDeriv

affineTest :: TestTree
affineTest = $(testGroupGenerator)


{-# LANGUAGE FlexibleContexts    #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# OPTIONS_GHC -fno-warn-orphans #-}
-- |
-- Orphan instances for common data types
module Tests.Orphanage where

import Control.Applicative
import Statistics.Distribution.Beta            (BetaDistribution, betaDistr)
import Statistics.Distribution.Binomial        (BinomialDistribution, binomial)
import Statistics.Distribution.CauchyLorentz
import Statistics.Distribution.ChiSquared      (ChiSquared, chiSquared)
import Statistics.Distribution.Exponential     (ExponentialDistribution, exponential)
import Statistics.Distribution.FDistribution   (FDistribution, fDistribution)
import Statistics.Distribution.Gamma           (GammaDistribution, gammaDistr)
import Statistics.Distribution.Geometric
import Statistics.Distribution.Hypergeometric
import Statistics.Distribution.Laplace         (LaplaceDistribution, laplace)
import Statistics.Distribution.Lognormal       (LognormalDistribution, lognormalDistr)
import Statistics.Distribution.Normal          (NormalDistribution, normalDistr)
import Statistics.Distribution.Poisson         (PoissonDistribution, poisson)
import Statistics.Distribution.StudentT
import Statistics.Distribution.Transform       (LinearTransform, scaleAround)
import Statistics.Distribution.Uniform         (UniformDistribution, uniformDistr)
import Statistics.Distribution.Weibull         (WeibullDistribution, weibullDistr)
import Statistics.Distribution.DiscreteUniform (DiscreteUniform, discreteUniformAB)
import Statistics.Types

import Test.QuickCheck         as QC

----------------------------------------------------------------
-- Arbitrary instances for ditributions
----------------------------------------------------------------

instance QC.Arbitrary BinomialDistribution where
  arbitrary = binomial <$> QC.choose (1,100) <*> QC.choose (0,1)
instance QC.Arbitrary ExponentialDistribution where
  arbitrary = exponential <$> QC.choose (0,100)
instance QC.Arbitrary LaplaceDistribution where
  arbitrary = laplace <$> QC.choose (-10,10) <*> QC.choose (0, 2)
instance QC.Arbitrary GammaDistribution where
  arbitrary = gammaDistr <$> QC.choose (0.1,10) <*> QC.choose (0.1,10)
instance QC.Arbitrary BetaDistribution where
  arbitrary = betaDistr <$> QC.choose (1e-3,10) <*> QC.choose (1e-3,10)
instance QC.Arbitrary GeometricDistribution where
  arbitrary = geometric <$> QC.choose (0,1)
instance QC.Arbitrary GeometricDistribution0 where
  arbitrary = geometric0 <$> QC.choose (0,1)
instance QC.Arbitrary HypergeometricDistribution where
  arbitrary = do l <- QC.choose (1,20)
                 m <- QC.choose (0,l)
                 k <- QC.choose (1,l)
                 return $ hypergeometric m l k
instance QC.Arbitrary LognormalDistribution where
  -- can't choose sigma too big, otherwise goes outside of double-float limit
  arbitrary = lognormalDistr <$> QC.choose (-100,100) <*> QC.choose (1e-10, 20)
instance QC.Arbitrary NormalDistribution where
  arbitrary = normalDistr <$> QC.choose (-100,100) <*> QC.choose (1e-3, 1e3)
instance QC.Arbitrary PoissonDistribution where
  arbitrary = poisson <$> QC.choose (0,1)
instance QC.Arbitrary ChiSquared where
  arbitrary = chiSquared <$> QC.choose (1,100)
instance QC.Arbitrary UniformDistribution where
  arbitrary = do a <- QC.arbitrary
                 b <- QC.arbitrary `suchThat` (/= a)
                 return $ uniformDistr a b
instance QC.Arbitrary WeibullDistribution where
  arbitrary = weibullDistr <$> QC.choose (1e-3,1e3) <*> QC.choose (1e-3, 1e3)
instance QC.Arbitrary CauchyDistribution where
  arbitrary = cauchyDistribution
                <$> arbitrary
                <*> ((abs <$> arbitrary) `suchThat` (> 0))
instance QC.Arbitrary StudentT where
  arbitrary = studentT <$> ((abs <$> arbitrary) `suchThat` (>0))
instance QC.Arbitrary d => QC.Arbitrary (LinearTransform d) where
  arbitrary = do
    m <- QC.choose (-10,10)
    s <- QC.choose (1e-1,1e1)
    d <- arbitrary
    return $ scaleAround m s d
instance QC.Arbitrary FDistribution where
  arbitrary =  fDistribution
           <$> ((abs <$> arbitrary) `suchThat` (>0))
           <*> ((abs <$> arbitrary) `suchThat` (>0))

instance (Arbitrary a, Ord a, RealFrac a) => Arbitrary (PValue a) where
  arbitrary = do
    (_::Int,x) <- properFraction <$> arbitrary
    return $ mkPValue $ abs x

instance (Arbitrary a, Ord a, RealFrac a) => Arbitrary (CL a) where
  arbitrary = do
    (_::Int,x) <- properFraction <$> arbitrary
    return $ mkCLFromSignificance $ abs x

instance Arbitrary a => Arbitrary (NormalErr a) where
  arbitrary = NormalErr <$> arbitrary

instance Arbitrary a => Arbitrary (ConfInt a) where
  arbitrary = liftA3 ConfInt arbitrary arbitrary arbitrary

instance (Arbitrary (e a), Arbitrary a) => Arbitrary (Estimate e a) where
  arbitrary = liftA2 Estimate arbitrary arbitrary

instance (Arbitrary a) => Arbitrary (UpperLimit a) where
  arbitrary = liftA2 UpperLimit arbitrary arbitrary

instance (Arbitrary a) => Arbitrary (LowerLimit a) where
  arbitrary = liftA2 LowerLimit arbitrary arbitrary

instance QC.Arbitrary DiscreteUniform where
  arbitrary = discreteUniformAB <$> QC.choose (1,1000) <*> QC.choose(1,1000)


module Statistics.Covariance.Shrinkage.Strimmer (
    psmallSVD
  , nsmallSVD
  , posSVD
  , covShrink
  , corShrink
  , corCoef
  , varShrink
  , varCoef
  ) where

import Numeric.LinearAlgebra hiding ((<>))
import Numeric.Sum (kbn, sumVector)
import Statistics.Quantile (quantile, s)
import qualified Data.Vector.Storable as VS

covShrink :: Matrix Double -> Herm Double
covShrink x = trustSym $ (asColumn sc * unSym c) * asRow sc
  where
    sc = sqrt $ varShrink x
    c = corShrink x

psmallSVD :: Matrix Double -> (Matrix Double, Vector Double, Matrix Double)
psmallSVD x = (u', d', v')
  where
    b = unSym $ mTm x
    (_, d, v) = svd b
    tol = 1e-6
    positive = find (> tol) d
    u' = (x <> v') * asRow (recip d')
    d' = sqrt . vector $ map (atIndex d) positive
    v' = v ¿ positive

nsmallSVD :: Matrix Double -> (Matrix Double, Vector Double, Matrix Double)
nsmallSVD x = (u', d', v')
  where
    b = unSym $ mTm (tr x)
    (u, d, _) = svd b
    tol = 1e-6
    positive = find (> tol) d
    u' = u ¿ positive
    d' = sqrt . vector $ map (atIndex d) positive
    v' = (tr x <> u') * asRow (recip d')

posSVD :: Matrix Double -> (Matrix Double, Vector Double, Matrix Double)
posSVD x = (u', d', v')
  where
    (u, d, v) = svd x
    tol = 1e-6
    positive = find (> tol) d
    u' = u ¿ positive
    d' = vector $ map (atIndex d) positive
    v' = v ¿ positive

fastSVD :: Matrix Double -> (Matrix Double, Vector Double, Matrix Double)
fastSVD x
  | n > edgeRatio * p = psmallSVD x
  | edgeRatio * n < p = nsmallSVD x
  | otherwise = posSVD x
  where
    edgeRatio = 2
    n = rows x
    p = cols x

mpower :: Herm Double -> Double -> Herm Double
mpower x p = trustSym $ (evec * asRow (cmap exp $ scale p (cmap log eval))) <> tr evec
  where (eval, evec) = eigSH x

corShrink :: Matrix Double -> Herm Double
corShrink x
  | lambda == 1 || alpha == 0 = trustSym $ ident p
  | alpha == 1 = unitDiag $ scale (1 - lambda) r0
  | lambda == 0 = unitDiag . trustSym $ svdxsV <> unSym (mpower c2 alpha) <> tr svdxsV
  | otherwise = unitDiag . trustSym . scale (exp (alpha * log lambda)) $ ident p - (svdxsV <> f <> tr svdxsV)
  where
    unitDiag :: Herm Double -> Herm Double
    unitDiag powr = trustSym (diag (1 - takeDiag (unSym powr))) `add` powr
    n = rows x
    p = cols x
    alpha = 1 :: Double
    w2 = recip $ fromIntegral n -- w = 1 / n => sum(w^2) =  1 / n
    h1 = 1 / (1 - w2)
    xs = let (u, v) = meanCov x in (x - asRow u) / asRow (sqrt (takeDiag (unSym v)))
    lambda = corCoef x
    r0 = scale h1 $ mTm (xs / sqrt (fromIntegral n))
    (svdxsU, svdxsS, svdxsV) = fastSVD xs
    m = size svdxsS
    utwu = tr svdxsU <> svdxsU / fromIntegral n
    c0 = utwu * asRow svdxsS * asColumn svdxsS
    c1 = scale ((1 - lambda) * h1) c0
    c2 = sym c1
    f = ident m - unSym (mpower (trustSym $ unSym (scale (recip lambda) c2) + ident m) alpha)

corCoef :: Matrix Double -> Double
corCoef x
  | denominator == 0 = 1
  | otherwise = numerator / denominator * h1w2
  where
    n = rows x
    p = cols x
    xs = let (u, v) = meanCov x in (x - asRow u) / asRow (sqrt (takeDiag (unSym v)))
    w2 = recip $ fromIntegral n
    h1w2 = w2 / (1 - w2)
    sw = sqrt . recip $ fromIntegral n
    xsw = scale sw xs
    (xswsvdU, xswsvdS, xswsvdV) = fastSVD xsw
    colSums = vector . map (sumVector kbn) . toColumns
    -- colCumSums = fromColumns . map (VS.postscanl (+) 0) . toColumns
    rowCumSums = fromRows . map (VS.postscanl (+) 0) . toRows
    sE2R = sumVector kbn (flatten (xsw * ((xswsvdU * asRow (cb xswsvdS)) <> tr xswsvdV)))
      - sumVector kbn (sq (colSums (sq xsw)))
    xs2w = scale sw $ sq xs
    sER2 = 2 * (sumVector kbn . flatten
      $ fliprl (takeColumns (p - 1) xs2w) * (rowCumSums . fliprl $ dropColumns 1 xs2w))
    denominator = sE2R
    numerator = sER2 - sE2R

varShrink :: Matrix Double -> Vector Double
varShrink x = scalar (lambda * target) + scale (1 - lambda) v
  where
    lambda = varCoef x
    v = takeDiag . unSym . snd $ meanCov x
    target = quantile s 1 2 v

varCoef :: Matrix Double -> Double
varCoef x
  | denominator == 0 = 1
  | otherwise = max 0 . min 1 $ numerator / denominator * h1w2
  where
    n = rows x
    w2 = recip (fromIntegral n)
    h1 = 1 / (1 - w2)
    h1w2 = w2 / (1 - w2)
    xc = x - asRow (fst (meanCov x))
    v = scale h1 q1
    target = quantile s 1 2 v
    zz = xc * xc
    q1 = fst $ meanCov zz
    q2 = fst (meanCov (zz * zz) ) - q1 * q1
    numerator = sumVector kbn q2
    denominator = sumVector kbn . sq $ q1 - scalar (target / h1)

sq, cb :: Num a => a -> a
sq a = a * a
cb a = a * sq a


import Criterion.Main
import qualified Data.Vector.Unboxed as U
import Numeric.SpecFunctions
import Numeric.Polynomial
import Text.Printf

-- Uniformly sample logGamma performance between 10^-6 to 10^6
benchmarkLogGamma logG =
  [ bench (printf "%.3g" x) $ nf logG x
  | x <- [ m * 10**n | n <- [ -8 .. 8 ]
                     , m <- [ 10**(i / tics) | i <- [0 .. tics-1] ]
         ]
  ]
  where tics = 3
{-# INLINE benchmarkLogGamma #-}

-- Power of polynomial to be evaluated (In other words length of coefficients vector)
coef_size :: [Int]
coef_size = [ 1,2,3,4,5,6,7,8,9
            , 10,    30
            , 100,   300
            , 1000,  3000
            , 10000, 30000
            ]
{-# INLINE coef_size #-}

-- Precalculated coefficients
coef_list :: [U.Vector Double]
coef_list = [ U.replicate n 1.2 | n <- coef_size]
{-# NOINLINE coef_list #-}

main :: IO ()
main = defaultMain
  [ bgroup "logGamma" $
    benchmarkLogGamma logGamma
  , bgroup "logGammaL" $
    benchmarkLogGamma logGammaL
  , bgroup "incompleteGamma" $
      [ bench (show p) $ nf (incompleteGamma p) p
      | p <- [ 0.1
             , 1,   3
             , 10,  30
             , 100, 300
             , 999, 1000
             ]
      ]
  , bgroup "factorial"
    [ bench (show n) $ nf factorial n
    | n <- [ 0, 1, 3, 6, 9, 11, 15
           , 20, 30, 40, 50, 60, 70, 80, 90, 100
           ]
    ]
  , bgroup "incompleteBeta"
    [ bench (show (p,q,x)) $ nf (incompleteBeta p q) x
    | (p,q,x) <- [ (10,      10,      0.5)
                 , (101,     101,     0.5)
                 , (1010,    1010,    0.5)
                 , (10100,   10100,   0.5)
                 , (100100,  100100,  0.5)
                 , (1001000, 1001000, 0.5)
                 , (10010000,10010000,0.5)
                 ]
    ]
  , bgroup "log1p"
      [ bench (show x) $ nf log1p x
      | x <- [ -0.9
             , -0.5
             , -0.1
             ,  0.1
             ,  0.5
             ,  1
             ,  10
             ,  100
             ]
      ]
  , bgroup "poly"
      $  [ bench ("vector_"++show (U.length coefs)) $ nf (\x -> evaluatePolynomial x coefs) (1 :: Double)
         | coefs <- coef_list ]
      ++ [ bench ("unpacked_"++show n) $ nf (\x -> evaluatePolynomialL x (map fromIntegral [1..n])) (1 :: Double)
         | n <- coef_size ]
  ]


module SongId
( chunkIdentifier
, findMaxFreq
, mapChunk2Freqs
, ft
) where

import Debug.Trace

import Data.Array.Unboxed (elems, assocs, UArray)
import Data.Array.IArray (Array(..), listArray, array)
import Data.Array.CArray (CArray)

import Data.Complex (Complex(..), magnitude)
import Data.Int (Int64, Int32, Int16, Int8)
import Math.FFT (dft)

import Data.List (sort)

import Common (importantFrequencies, breakChunks, slice)

chunkIdentifier :: [(Int, Int)] -> ([Float], [Complex Float]) -> [Float]
chunkIdentifier freqRanges (freqs, amps) =
    let breakFreqs chunk = map (flip slice chunk) freqRanges
        unzippedChunksInFrequencyRanges = map unzip $ breakFreqs $ zip freqs amps
        maxFreqsInRanges = map findMaxFreq $ unzippedChunksInFrequencyRanges
    in map fst maxFreqsInRanges

findMaxFreq :: ([Float], [Complex Float]) -> (Float, Complex Float)
findMaxFreq (xs, ys) = foldl (\max@(_, maxamp) new@(_, amp) ->
        if (magnitude amp) > (magnitude maxamp)
        then new
        else max)
    (0,0) $ zip xs ys

mapChunk2Freqs :: (Fractional a, Fractional b) => [a] -> ([b], [a])
mapChunk2Freqs chunk =
    let l = length chunk
        halfOfChunk = slice (0, l `quot` 2) chunk
        amps = map (/ fromIntegral l) halfOfChunk

        xs1 = map fromIntegral $ [0,1..l]
        xs2 = replicate l (44100 / fromIntegral l)
        freqs = map (\(x1,x2) -> x1 * x2) $ zip xs1 xs2
    in (freqs, amps)

ft :: Int -> UArray Int Int32 -> [CArray Int (Complex Float)]
ft chunkSize samples = map dft chunks
    where chunks = map (listArray (0, chunkSize - 1)) $ breakChunks chunkSize elements -- enumerate the samples in each chunk
          elements = map convertToComplex $ elems samples
          convertToComplex v = (fromIntegral v) :+ 0


{-# LANGUAGE DataKinds              #-}
{-# LANGUAGE FlexibleContexts       #-}
{-# LANGUAGE ScopedTypeVariables    #-}
{-# LANGUAGE TypeFamilies           #-}

module Numeric.Optimization.TwoPhase.Tableau.Pivot
    ( index
    , allCells
    , pivot
    ) where

import           Data.Finite                        (Finite)
import           GHC.TypeLits                       (KnownNat)
import           Numeric.LinearAlgebra              ((!))
import qualified Numeric.LinearAlgebra as LA        (atIndex)
import           Numeric.LinearAlgebra.Static       (L)
import qualified Numeric.LinearAlgebra.Static as LS

type Index rows cols = (Finite rows, Finite cols)

index :: (KnownNat rows, KnownNat cols)
    => Index rows cols
    -> L rows cols
    -> Double
index (i, j) x =
    LA.atIndex (LS.unwrap x) ix
  where
    ix = (fromIntegral i, fromIntegral j)

allCells :: (KnownNat rows, KnownNat cols)
    => (Double -> Bool)
    -> [Index rows cols]
    -> L rows cols
    -> Bool
allCells f is x =
    all f $ fmap (`index` x) is

pivot :: (KnownNat rows, KnownNat cols)
    => (Finite rows, Finite cols) -> L rows cols -> L rows cols
pivot (x', y') table =
    LS.build (newValue $ LS.unwrap table)
  where
    (x, y)          = (fromIntegral x', fromIntegral y')

    newValue mat i' j'
        | i == x    = mat!i!j / mat!x!y
        | j == y    = 0
        | otherwise = mat!x!y * mat!i!j - mat!x!j * mat!i!y
      where
        (i, j) = (round i', round j')


{-# LANGUAGE FlexibleContexts #-}
module Network (initNet, feedForward, sgd, forwardPass, backprop, updateMiniBatch) where

import Control.Monad
import Data.List.Split
import Numeric.LinearAlgebra as L

import MnistLoader
import Aux

feedForward :: Network -> Matrix Float -> Matrix Float
feedForward n a = foldl f a $ zip (weights n) (biases n)
  where
    f a' (w, b) = sigmoid ((w L.<> a') + b)

sgd :: Network
  -> [(Matrix Float, Matrix Float)] -- traning data
  -> [(Matrix Float, Matrix Float)] -- Evaluation data
  -> Int -- Number of epochs
  -> Int -- Current Epoc
  -> Int -- Mini batch size
  -> Float -- Eta
  -> IO Network
sgd n exemplars testData maxEpoch epoch batchSize eta = do
  exemplars' <- shuffle exemplars

  let batches = chunksOf batchSize exemplars'
  let newNet = foldl (\accNet batch -> updateMiniBatch accNet batch eta) n batches
  let correct = evaluate newNet testData
  putStrLn $ "Epoch " ++ (show (epoch + 1)) ++ " Eval: " ++ (show correct) ++ " / " ++ (show $ length testData)

  if epoch == maxEpoch
    then return newNet
    else sgd newNet exemplars' testData maxEpoch (epoch + 1) batchSize eta

updateMiniBatch :: Network -> [(Matrix Float, Matrix Float)] -> Float -> Network
updateMiniBatch n l@((x, y) : rest) eta = let
    initAcc = backprop n x y
    nablaSums = foldl f initAcc rest
    nabla_b = map fst nablaSums
    nabla_w = map snd nablaSums
  in
    Network {
        sizes = sizes n
      , weights = zipWith wavg (weights n) nabla_w
      , biases = zipWith wavg (biases n) nabla_b
    }
  where
    f nablaSum (x, y) = let
        bp = backprop n x y
      in
        zipWith add nablaSum bp
    add (nb, nw) (nb', nw') = (nb+nb', nw+nw')
    wavg :: Matrix Float -> Matrix Float -> Matrix Float
    wavg m1 m2 = let
        nt = (1><1) [fromIntegral $ length l]
        eta' = (1><1) [eta]
      in m1 - (eta' / nt) * m2

-- Network -> x -> y -> (nabla_b, nabla_w)
backprop :: Network -> Matrix Float -> Matrix Float -> [(Matrix Float, Matrix Float)]
backprop n x y = let
    -- Reversed list of forward passes
    ((a:a':as), (z:zs)) = forwardPass n x
    revW = foldl (\acc w -> w:acc) [] (weights n)
    fp = zip as zs
    zippedFp = zip fp revW

    -- Handle last layer manually
    delta = (cost_derivative a y) * (sigmoid' z)
    nabla_b = delta
    nabla_w = delta L.<> (tr' a')
    res = foldl f [(nabla_b, nabla_w)] zippedFp
  in
    res
  where
    f (acc@((delta, _):_)) ((a, z), w) = let
        sp = sigmoid' z
        delta' = ((tr' w) L.<> delta) * sp
        nabla_b = delta'
        nabla_w = delta' L.<> (tr' a)
      in
        (nabla_b, nabla_w) : acc

-- Network -> activation -> (activations, z vectors)
forwardPass :: Network -> Matrix Float -> ([Matrix Float], [Matrix Float])
forwardPass n a = let
  pairs = zip (weights n) (biases n)
  fp = foldl f [(a, a)] pairs
  in
    (map fst fp, init $ map snd fp)
    where
      f (l@((a, _):_)) (w, b) = let
          z = (w L.<> a) + b
          activation = sigmoid z
        in (activation, z):l

evaluate :: Network -> [(Matrix Float, Matrix Float)] -> Int
evaluate net exemplars = sum $ map f exemplars
  where
    f (x, y) = let
        approx = argMax $ feedForward net x
        actual = argMax y
      in
        if approx == actual
          then 1
          else 0

cost_derivative output_activations y = output_activations - y

-- Miscellaneous
sigmoid z = 1.0 / (1.0 + (exp $ -z))

sigmoid' :: Matrix Float -> Matrix Float
sigmoid' z = (sigmoid z) * (1 - (sigmoid z))

argMax :: Matrix Float -> Int
argMax m = L.maxIndex $ L.flatten m

-- assuming row-major order
data Network = Network { sizes :: [Int]
                       , weights :: [Matrix Float]
                       , biases :: [Matrix Float]
                       } deriving (Show)

initNet :: [Int] -> IO (Network)
initNet sizes = do
  ws <- genWeights sizes
  bs <- genBiases sizes
  return Network { sizes = sizes
                  , weights = ws
                  , biases = bs
                  }
  where
    genWeights [s, s'] =
      do
        m <- L.randn s' s
        return [L.cmap realToFrac m]
    genWeights (s:s':ss) =
      do
        rest <- genWeights (s' : ss)
        m <- L.randn s' s
        return (L.cmap realToFrac m : rest)
    genWeights _ = error "Not enough layers"
    genBiases [s, s'] =
      do
        m <- L.randn s' 1
        return $ [L.cmap realToFrac m]
    genBiases (s:s':ss) =
      do
        rest <- genBiases (s' : ss)
        m <- L.randn s' 1
        return (L.cmap realToFrac m : rest)
    genBiases _ = error "Not enough layers"


module Statistics.Information.Mixed.Entropy where

import Data.Matrix
import qualified Statistics.Information.Continuous.Entropy as CE
import qualified Statistics.Information.Discrete.Entropy as DE
import Statistics.Information.Mixed.MutualInfo
import Statistics.Information.Utils.Random
import System.Random

centropydc :: (RandomGen g, Eq a) => g -> Int -> Int -> Matrix a ->
              Matrix Double -> Double
centropydc g k base xs ys = DE.entropy base xs - midc g k base xs ys

centropycd :: (RandomGen g, Eq a) => g -> Int -> Int -> Matrix Double ->
              Matrix a -> Double
centropycd g k base xs ys = CE.entropy g1 k base xs - micd g2 k base xs ys where
  (g1, g2) = split g

entropy :: (RandomGen g, Eq a) => g -> Int -> Int -> Matrix a ->
           Matrix Double -> Double
entropy g k base xs ys = hx + hy - ixy
  where
    (g1, g2) = split g
    hx = DE.entropy base xs
    hy = CE.entropy g1 k base ys
    ixy = midc g2 k base xs ys


{-# OPTIONS_GHC -Wall                     #-}
{-# OPTIONS_GHC -fno-warn-name-shadowing  #-}
{-# OPTIONS_GHC -fno-warn-type-defaults   #-}
{-# OPTIONS_GHC -fno-warn-unused-do-bind  #-}
{-# OPTIONS_GHC -fno-warn-missing-methods #-}
{-# OPTIONS_GHC -fno-warn-orphans         #-}

{-# LANGUAGE DataKinds                    #-}
{-# LANGUAGE ScopedTypeVariables          #-}
{-# LANGUAGE TypeOperators                #-}
{-# LANGUAGE TypeFamilies                 #-}

-- | An extended Kalman filter. Note that this could be generalized
-- further. If you need a Kalman filter and have a state update matrix
-- \(A\) then you can just use @const bigA@ as the function which
-- returns the Jacobian.

module Kalman (
  extKalman
  )where

import GHC.TypeLits
import Numeric.LinearAlgebra.Static hiding ( create )
import Data.Maybe ( fromJust )

-- | The model for the extended Kalman filter is given by
--
-- \[
-- \begin{aligned}
-- \boldsymbol{x}_i &= \boldsymbol{a}(\boldsymbol{x}_{i-1}) + \boldsymbol{\psi}_{i-1} \\
-- \boldsymbol{y}_i &= {H}\boldsymbol{x}_i + \boldsymbol{\upsilon}_i
-- \end{aligned}
-- \]
--
-- where
--
-- * \(\boldsymbol{a}\)\ is some non-linear vector-valued state update
-- function.
--
-- * \(\boldsymbol{\psi}_{i}\) are independent identically normally
-- distributed random variables with mean 0 representing the fact that
-- the state update is noisy: \(\boldsymbol{\psi}_{i} \sim {\cal{N}}(0,Q)\).
--
-- * \({H}\) is a matrix which represents how we observe the
-- hidden state process.
--
-- * \(\boldsymbol{\upsilon}_i\) are independent identically normally
-- distributed random variables with mean 0 represent the fact that
-- the observations are noisy: \(\boldsymbol{\upsilon}_{i} \sim {\cal{N}}(0,R)\).
--
-- We assume the whole process starts at 0 with our belief of the state (aka the
-- prior state) being given by
-- \(\boldsymbol{x}_0 \sim {\cal{N}}(\boldsymbol{\mu}_0, \Sigma_0)\)
extKalman ::  forall m n .
              (KnownNat m, KnownNat n, (1 <=? n) ~ 'True, (1 <=? m) ~ 'True) =>
              R n              -- ^ Prior mean \(\boldsymbol{\mu}_0\)
              -> Sq n          -- ^ Prior variance \(\Sigma_0\)
              -> L m n         -- ^ Observation map \(H\)
              -> Sq m          -- ^ Observation noise \(R\)
              -> (R n -> R n)  -- ^ State update function \(\boldsymbol{a}\)
              -> (R n -> Sq n) -- ^ A function which returns the
                               -- Jacobian of the state update
                               -- function at a given point
                               -- \(\frac{\partial \boldsymbol{a}}{\partial \boldsymbol{x}}\)
              -> Sq n          -- ^ State noise \(Q\)
              -> [R m]         -- ^ List of observations \(\boldsymbol{y}_i\)
              -> [(R n, Sq n)] -- ^ List of posterior means and variances \((\hat{\boldsymbol{x}}_i, \hat{\Sigma}_i)\)
extKalman muPrior sigmaPrior bigH bigSigmaY
  littleA bigABuilder bigSigmaX ys = result
  where
    result = scanl update (muPrior, sigmaPrior) ys

    update :: (R n, Sq n) -> R m -> (R n, Sq n)
    update (xHatFlat, bigSigmaHatFlat) y =
      (xHatFlatNew, bigSigmaHatFlatNew)
      where

        v :: R m
        v = y - (bigH #> xHatFlat)

        bigS :: Sq m
        bigS = bigH <> bigSigmaHatFlat <> (tr bigH) + bigSigmaY

        bigK :: L n m
        bigK = bigSigmaHatFlat <> (tr bigH) <> (inv bigS)

        xHat :: R n
        xHat = xHatFlat + bigK #> v

        bigSigmaHat :: Sq n
        bigSigmaHat = bigSigmaHatFlat - bigK <> bigS <> (tr bigK)

        bigA :: Sq n
        bigA = bigABuilder xHat

        xHatFlatNew :: R n
        xHatFlatNew = littleA xHat

        bigSigmaHatFlatNew :: Sq n
        bigSigmaHatFlatNew = bigA <> bigSigmaHat <> (tr bigA) + bigSigmaX

inv :: (KnownNat n, (1 <=? n) ~ 'True) => Sq n -> Sq n
inv m = fromJust $ linSolve m eye


{-# LANGUAGE DataKinds #-}

module Rotation
  ( Rotation
  , rotationMatrix
  , rotate
  , mkRotation
  , mkXYRotation
  , mkXYRotationCamera
  , rotateV
  ) where

import qualified Numeric.LinearAlgebra        as L
import           Numeric.LinearAlgebra.Static

import           Matrix

type Quaternion = R 4

type Rotation = Quaternion

rotate :: Quaternion -> Quaternion -> Quaternion
q1 `rotate` q2 =
  vec4
    (r1 * r2 - i1 * i2 - j1 * j2 - k1 * k2)
    (r1 * i2 + i1 * r2 + j1 * k2 - k1 * j2)
    (r1 * j2 + j1 * r2 + k1 * i2 - i1 * k2)
    (r1 * k2 + k1 * r2 + i1 * j2 - j1 * i2)
  where
    [r1, i1, j1, k1] = L.toList . unwrap $ q1
    [r2, i2, j2, k2] = L.toList . unwrap $ q2

mkRotation :: Double -> Vector3 -> Quaternion
mkRotation angle axis =
  normalize $
  vec4
    (cos (angle / 2))
    (x * sin (angle / 2))
    (y * sin (angle / 2))
    (z * sin (angle / 2))
  where
    [x, y, z] = L.toList . unwrap $ axis

rotationMatrix :: Quaternion -> Matrix4
rotationMatrix q = eye + 2 * m
  where
    [r, i, j, k] = L.toList . unwrap $ q
    m =
      row (vec4 (-j * j - k * k) (i * j - k * r) (i * k + j * r) 0) ===
      row (vec4 (i * j + k * r) (-i * i - k * k) (j * k - i * r) 0) ===
      row (vec4 (i * k - j * r) (j * k + i * r) (-i * i - j * j) 0) ===
      row (vec4 0 0 0 0)

mkXYRotationCamera x y = q2 `rotate` q1
  where
    q1 = mkRotation x yHat
    q2 = mkRotation y xHat

mkXYRotation x y = q2 `rotate` q1
  where
    q1 = mkRotation y xHat
    q2 = mkRotation x yHat

conj :: Quaternion -> Quaternion
conj q = q * vec4 1 (-1) (-1) (-1)

qmult = rotate

rotateV :: Vector3 -> Quaternion -> Vector3
rotateV v q = v'
  where
    (_, v') = split $ q `qmult` p `qmult` conj q
    p = vector [0] # v


{-# LANGUAGE RecordWildCards #-}

module Data.Matrix.CSC where

import Control.Arrow (second)
import Control.Exception
import Data.Maybe
import qualified Data.Vector.Storable as VS
import qualified Data.Vector as VB
import Foreign.C.Types
import Numeric.LinearAlgebra

data CSC a = CSC
    { cscVals :: VS.Vector a
    , cscCols :: VS.Vector CInt
    , cscRows :: VS.Vector CInt
    , cscNRows :: CInt
    , cscNCols :: CInt
    } deriving (Eq, Show)

fromMatrix :: (Element a, VS.Storable a, Eq a, Num a) => Matrix a -> CSC a
fromMatrix mat = CSC{..}
  where
    cscNRows = fromIntegral $ rows mat
    cscNCols = fromIntegral $ cols mat
    cscRows = VS.fromList $ concat irs
    cscVals = VS.fromList $ concat vals
    cscCols = VS.fromList $ scanl (\a b -> a + fromIntegral (length b)) 0 vals
    (irs, vals) = unzip $ map (go . VS.toList) (toColumns mat)
    go col = unzip $ catMaybes $ zipWith (\idx val -> if val == 0 then Nothing else Just (fromIntegral idx, val)) [0..] col

toMatrix :: (Element a, Num a, VS.Storable a) => CSC a -> Matrix a
toMatrix CSC {..} = fromColumns cols
  where
    cols = zipWith go rows vals
    go rs vs = VS.replicate (fromIntegral cscNRows) 0 VS.//
        (zipWith (\r v -> (fromIntegral r, v)) (VS.toList rs) (VS.toList vs))
    rows = split cscRows cscCols
    vals = split cscVals cscCols
    split v = snd . VS.foldr splitter (v, []) . VS.init
      where
        splitter idx (v, xs) = second (:xs) $ VS.splitAt (fromIntegral idx) v

(===) :: (VS.Storable a) => CSC a -> CSC a -> CSC a
a === b = assert (cscNCols a == cscNCols b) $ CSC
    { cscNCols = cscNCols a
    , cscNRows = cscNRows a + cscNRows b
    , cscRows = rows
    , cscVals = vals
    , cscCols = cols
    }
  where
    cols = VS.zipWith (+) (cscCols a) (cscCols b)
    rows = VS.concat $ concat $
        zipWith (\a b -> [a, b]) (split (cscRows a) (cscCols a))
        (split (VS.map (+ cscNRows a) $ cscRows b) (cscCols b))
    vals = VS.concat $ concat $
        zipWith (\a b -> [a, b]) (split (cscVals a) (cscCols a))
        (split (cscVals b) (cscCols b))
    split v = snd . VS.foldr splitter (v, [])
      where
        splitter idx (v, xs) = second (:xs) $ VS.splitAt (fromIntegral idx) v

(|||) :: (VS.Storable a) => CSC a -> CSC a -> CSC a
a ||| b = assert (cscNRows a == cscNRows b) $ CSC
    { cscNCols = cscNCols a + cscNCols b
    , cscNRows = cscNRows a
    , cscRows = cscRows a VS.++ cscRows b
    , cscVals = cscVals a VS.++ cscVals b
    , cscCols = cscCols a VS.++ (VS.map (+ VS.last (cscCols a)) $ VS.tail (cscCols b))
    }

diagBlock :: (VS.Storable a) => CSC a -> CSC a -> CSC a
diagBlock a b = CSC
    { cscNCols = cscNCols a + cscNCols b
    , cscNRows = cscNRows a + cscNRows b
    , cscRows = cscRows a VS.++ VS.map (+ cscNRows a) (cscRows b)
    , cscVals = cscVals a VS.++ cscVals b
    , cscCols = cscCols a VS.++ (VS.map (+ VS.last (cscCols a)) $ VS.tail (cscCols b))
    }

zero :: (Integral a, VS.Storable b) => a -> a -> CSC b
zero rows cols = CSC
    { cscNCols = fromIntegral cols
    , cscNRows = fromIntegral rows
    , cscRows = VS.empty
    , cscVals = VS.empty
    , cscCols = VS.replicate (fromIntegral cols + 1) 0
    }

fromBlocks :: (VS.Storable a) => [[CSC a]] -> CSC a
fromBlocks = foldr1 (===) . map (foldr1 (|||))

scale :: (VS.Storable a, Num a) => a -> CSC a -> CSC a
scale mult csc = csc { cscVals = VS.map (* mult) (cscVals csc) }


module Newton.Render where

-- this module comes from JuicyPixels package
import qualified Codec.Picture            as J
import           Codec.Picture.ColorQuant (defaultPaletteOptions)
import           Data.Complex
import           Newton.Utils             (root)

-- | Add brightness to an RGB8 Image
brightnessRGB8 :: Int -> J.Palette -> J.Palette
brightnessRGB8 add = J.pixelMap brightFunction
  where
    up v = fromIntegral (fromIntegral v + add)
    brightFunction (J.PixelRGB8 r g b) = J.PixelRGB8 (up r) (up g) (up b)

-- | Save an PNG image as a file
saveBasinsPic ::
     FilePath -- ^ Where to write the image file.
  -> J.Palette
  -> IO ()
saveBasinsPic = J.writePng

saveBasinsAnim ::
     FilePath -- ^ Filename of an image
  -> J.GifDelay -- ^ Delay between frames
  -> [J.Palette] -- ^ List of pictures to store within the GIF
  -> IO ()
saveBasinsAnim path delay images =
  case J.writeGifImages path J.LoopingForever finalImgs of
    Left msg  -> error msg
    Right ans -> ans
  where
    palettized = map (J.palettize defaultPaletteOptions) images
    finalImgs = map (\(img, plt) -> (plt, delay, img)) palettized

-- | Render attraction basins for zeros (roots) of a complex function
-- as a PNG image.
renderBasins ::
     (Fractional a, RealFloat a)
  => Int -- ^ Size of the square image in pixels
  -> (Complex a -> Complex a) -- ^ A function f to find zero for.
  -> (Complex a -> Complex a) -- ^ A derivative of function f.
  -> (Complex a -> J.PixelRGB8) -- ^ How to color roots.
  -> J.Palette
renderBasins commonSize f f' color = image
  where
    image = J.generateImage render width height
    width = commonSize -- width in pixels
    height = commonSize -- height in pixels
    (xfrom, xto) = (-2, 2) -- range for the real axis
    (yfrom, yto) = (-2, 2) -- range for the imaginary axis
    -- | Convert pixel coordinates to a complex number.
    pixelToComplex (i, j) = re :+ im
      where
        re = xfrom + (xto - xfrom) * fromIntegral i / fromIntegral width
        im = yfrom + (yto - yfrom) * fromIntegral j / fromIntegral height
    n = 100 -- maximum number of iterations
    p e = magnitude e < 1e-7 -- precision tolerance
    -- | Convert actual number of interations into
    -- into a grey pixel using logarithmic scale
    -- (this helps bring out more details in the image).
    logscale :: Integral a => Int -> a -> a
    logscale k c = floor (fromIntegral c * (1 - logBase n' k'))
      where
        n' = fromIntegral n + 1
        k' = fromIntegral k + 1
    -- | Render one pixel.
    render :: Int -> Int -> J.PixelRGB8
    render i j =
      case root n p f f' (pixelToComplex (i, j)) of
        Nothing -> J.PixelRGB8 0 0 0 -- black
        Just (k, xr) ->
          let J.PixelRGB8 r g b = color xr
           in J.PixelRGB8 (logscale k r) (logscale k g) (logscale k b)

-- | Simple (unstructured) complex number coloring.
simpleColoring :: Complex Double -> J.PixelRGB8
simpleColoring z = J.PixelRGB8 r g b
  where
    r = floor ((2 ^ 8 - 1) * r')
    g = floor ((2 ^ 8 - 1) * g')
    b = floor ((2 ^ 8 - 1) * b')
    (r', g', b') = hsl (phase z) 1 (1 - 0.5 ** magnitude z)

-- | Convert HSL color to RGB.
-- This function is copied with slight changes
-- from http://hackage.haskell.org/package/colour-2.3.4
hsl :: (RealFrac a, Floating a, Ord a) => a -> a -> a -> (a, a, a)
hsl h s l = (component tr, component tg, component tb)
  where
    hk = h / (2 * pi)
    (tr, tg, tb) = (mod1 (hk + 1 / 3), mod1 hk, mod1 (hk - 1 / 3))
    q
      | l < 0.5 = l * (1 + s)
      | otherwise = l + s - l * s
    p = 2 * l - q
    component t
      | t < 1 / 6 = p + ((q - p) * 6 * t)
      | t < 1 / 2 = q
      | t < 2 / 3 = p + ((q - p) * 6 * (2 / 3 - t))
      | otherwise = p
    mod1 x =
      if pf < 0
        then pf + 1
        else pf
      where
        (_, pf) = properFraction x


-- {-# DeriveDataTypeable #-}
module Main where

  import System.Environment
  import System.Random
  import Numeric.LinearAlgebra (Matrix, assoc, Z, toLists)
  import Control.Monad (join)
  import Data.Char (toLower)

  import System.Console.CmdArgs.Implicit

  import ODMatrix.Sim.BusState

  data SimParams = Params {
      size :: Int,
      capacity :: Int,
      seed :: Int,
      cases :: Int,
      lambdas :: Bool
  } deriving (Show, Data, Typeable)

  params = Params {
    size = 10 &= help "Size of the route",
    capacity = 50 &= help "Capacity of the unit",
    seed = 0 &= help "Seed for the simulation",
    cases = 10 &= help "Number of cases to simulate",
    lambdas = False &= help "Indicates if the lists of lambdas will be provided manually"
  }

  getLambdas :: RandomGen g => Bool -> g -> Int -> Int -> IO ([Double],[Double])
  getLambdas False g n k = do -- Random genetated
    let (bg,ag) = split g
        bs = take n $ randomRs (0,fromIntegral k) bg
        as = take n $ randomRs (0,fromIntegral k) ag
    return (bs, as)
  getLambdas True _ n _ = do  -- Manually provided
    input <- getContents
    let ls = lines input
        bs = take n ls
        as = take n $ drop n ls
    return (map read bs, map read as)

  main :: IO ()
  main = do
    -- ODMatrix Simulation
    putStrLn "{"

    args <- cmdArgs params

    let (Params n k seed cases manual_lambdas) = args
        (g1, g2) = split $ mkStdGen seed

    putStrLn "\"params\": {"
    putStrLn $ "\"size\":" ++ (show n) ++ ","
    putStrLn $ "\"capacity\":" ++ (show k) ++ ","
    putStrLn $ "\"seed\":" ++ (show seed) ++ ","
    putStrLn $ "\"cases\":" ++ (show cases) ++ ","
    putStrLn $ "\"manual_lambdas\":" ++ (map toLower $ show manual_lambdas)
    putStrLn "},"

    -- putStrLn $  if manual_lambdas
    --             then "# Reading lambdas..."
    --             else "# Generating randoms lambdas..."

    (bs,as) <- getLambdas manual_lambdas g1 n k

    putStrLn "\"lambdas\" : {"

    -- Boarding lambdas
    putStr "\"boards\": "
    putStr . show $ bs
    putStrLn ","

    -- Alighting lambdas
    putStr "\"alights\": "
    putStrLn . show $ as

    putStrLn "},"

    -- Generating origin destination matrices

    putStrLn "\"odms\": {"
    let ms = map (assoc (n,n) 0) (busSims k bs as g2) :: [Matrix Z]
        ms' = map (\(i,m) -> "\"" ++ (show i) ++ "\": " ++ (show m) ++ ",\n") .
              zip [0..] .
              take cases . map (join . toLists) $ ms

    putStrLn . init . init . join $ ms'
    putStrLn "}"

    putStrLn "}"


{-# OPTIONS_GHC -Wno-incomplete-patterns #-}

module CNVectorSpec where

import Data.Complex
import Data.Proxy

import Test.QuickCheck
import Test.QuickCheck.Poly

import CNVector
import Category
import Comonad
import Functor
import Unboxed
-- import Vec

type N = 10

prop_CNVector_Functor_id :: FnProp (CNVector N A)
prop_CNVector_Functor_id = checkFnEqual law_Functor_id

prop_CNVector_Functor_comp :: Fun B C -> Fun A B -> FnProp (CNVector N A)
prop_CNVector_Functor_comp f g = checkFnEqual (law_Functor_comp f g)

prop_CNVector_Semicomonad_comm ::
    Fun (CNVector N B) C -> Fun (CNVector N A) B -> FnProp (CNVector N A)
prop_CNVector_Semicomonad_comm f g = checkFnEqual (law_Semicomonad_comm f g)

prop_CNVector_Comonad_id :: FnProp (CNVector N A)
prop_CNVector_Comonad_id = checkFnEqual (law_Comonad_id (Proxy @(->)))

prop_CNVector_Comonad_apply :: Fun (CNVector N A) B -> FnProp (CNVector N A)
prop_CNVector_Comonad_apply f = checkFnEqual (law_Comonad_apply f)

prop_CNVector_Semicomonad1_comm ::
    Fun (CNVector N B) C -> Fun (CNVector N A) B -> FnProp (CNVector N A)
prop_CNVector_Semicomonad1_comm f g = checkFnEqual (law_Semicomonad1_comm f g)

prop_CNVector_Semicomonad1_comm' ::
    Fun (CNVector N B) C -> Fun (CNVector N A) B -> FnProp (CNVector N A)
prop_CNVector_Semicomonad1_comm' (Fn f) (Fn g) =
    checkFnEqual (law_Semicomonad1_comm' f g)

prop_CNVector_Comonad1_id :: FnProp (CNVector N A)
prop_CNVector_Comonad1_id = checkFnEqual (law_Comonad1_id (Proxy @(->)))

prop_CNVector_Comonad1_id' :: FnProp (CNVector N A)
prop_CNVector_Comonad1_id' = checkFnEqual law_Comonad1_id'

prop_CNVector_Comonad1_apply ::
    Fun (CNVector N A) B -> FnProp (CNVector N A)
prop_CNVector_Comonad1_apply f = checkFnEqual (law_Comonad1_apply f)

prop_CNVector_Comonad1_apply' ::
    Fun (CNVector N A) B -> FnProp (CNVector N A)
prop_CNVector_Comonad1_apply' (Fn f) = checkFnEqual (law_Comonad1_apply' f)

type UA = Int
type UB = Double
type UC = Complex Double

prop_CNUVector_Functor_id :: FnProp (CNUVector N UA)
prop_CNUVector_Functor_id = checkFnEqual law_Functor_id

prop_CNUVector_Functor_comp ::
    (UB -#> UC) -> (UA -#> UB) -> FnProp (CNUVector N UA)
prop_CNUVector_Functor_comp f g = checkFnEqual (law_Functor_comp f g)

-- prop_CNUVector_NatAdjunction_restrict ::
--     Proxy (CNUVector N) -> (UVec3 N) UA -> Property
-- prop_CNUVector_NatAdjunction_restrict = law_NatAdjunction_restrict

prop_CNUVector_Semicomonad1_comm' ::
    Fun (CNUVector N UB) UC -> Fun (CNUVector N UA) UB ->
    FnProp (CNUVector N UA)
prop_CNUVector_Semicomonad1_comm' (Fn f) (Fn g) =
    checkFnEqual (law_Semicomonad1_comm'
                     (f :: CNUVector N UB -> UC)
                     (g :: CNUVector N UA -> UB))

prop_CNUVector_Comonad1_id' :: FnProp (CNUVector N UA)
prop_CNUVector_Comonad1_id' = checkFnEqual law_Comonad1_id'

prop_CNUVector_Comonad1_apply' ::
    Fun (CNVector N UA) UB -> FnProp (CNVector N UA)
prop_CNUVector_Comonad1_apply' (Fn f) = checkFnEqual (law_Comonad1_apply' f)


{-# LANGUAGE BangPatterns #-}
{- |
   Description :  Functions to work with transition probability matrices on rooted trees
   Copyright   :  (c) Dominik Schrempf 2017
   License     :  GPLv3

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

Calculate transition probability matrices, map rate matrices on trees, populate
a tree with states according to a stationary distribution, etc.

The implementation of the Markov process is more than basic and can be improved in a lot of ways.

TODO: Try https://hackage.haskell.org/package/Dist-0.4.1.0/docs/Numeric-Probability-Distribution.html

* Changelog

-}

module Transition where
import           Control.Monad.Random.Strict hiding (fromList)
import           Data.Distribution           hiding (toList)
import           Numeric.LinearAlgebra       hiding (fromList)
import           RateMatrix
import           RTree

-- | This may be moved to a different module ProbMatrix or alike.
type ProbMatrix   = Matrix R

-- | The important matrix that gives the probabilities to move from one state to
-- another in a specific time (branch length).
probMatrix :: RateMatrix -> BranchLn -> ProbMatrix
probMatrix m t = exponential
  where !exponential = expm $ scale t m

-- | Convert a tree with branch lengths into a tree that has transition
-- probability matrices assigned to each of its branches.
branchLengthsToTransitionProbs :: RateMatrix -> RTree a Double -> RTree a ProbMatrix
branchLengthsToTransitionProbs m = fmap (probMatrix m)

-- | Move from a given state to a new one according to a transition probability
-- matrix (for performance reasons this probability matrix needs to be given as
-- a list of generators, see
-- https://hackage.haskell.org/package/distribution-1.1.0.0/docs/Data-Distribution-Sample.html).
-- This function is the bottleneck of the simulator and takes up most of the
-- computation time. However, I was not able to find a faster implementation
-- than the one from Data.Distribution.
jump :: (MonadRandom m) => State -> [Generator State] -> m State
jump (State s) p = target
  where !target = getSample $ p !! s

-- | Perform N jumps from a given state and according to a transition
-- probability matrix transformed to a list of generators. This implementation
-- uses 'foldM' and I am not sure how to access or store the actual chain. This
-- could be done by an equivalent of 'scanl' for general monads, which I was
-- unable to find. This function is neat, but will most likely not be needed.
-- However, it is instructive and is left in place.
jumpN :: (MonadRandom m) => State -> [Generator State] -> Int -> m State
jumpN s p n = foldM jump s (replicate n p)

-- | This is the heart of the simulation. Take a tree (most probably with node
-- labels of type a, but this could be anything) and a root state. If there is a
-- split (i.e., if we have a node and not a leaf), jump down the left branch and
-- populate and flatten the tree and jump down the right branch and populate and
-- flatten the tree and append the two results. This is what 'liftM2' is doing,
-- it lifts the append function of lists (++) to the 'Rand g' monad. If we
-- encounter a leaf, just return the node label (or whatever the type a is) and
-- the state that we ended up at.
populateAndFlattenTree :: (MonadRandom m) => RTree a [Generator State] -> State -> m [(a, State)]
populateAndFlattenTree (Leaf a) s = return [(a, s)]
populateAndFlattenTree (Node _ lp lc rp rc) s = liftM2 (++) (jumpDownBranch lp lc) (jumpDownBranch rp rc)
  where jumpDownBranch p t = jump s p >>= populateAndFlattenTree t

-- | Convert a stationary distribution to a generator which speeds up random
-- picks. Cf. 'jump'.
stationaryDistToGenerator :: StationaryDist -> Generator State
stationaryDistToGenerator f = fG
  -- This is a little complicated. I need to convert the vector to a list to be
  -- able to create a distribution.
  where !fL = toList f
        !fD = fromList $ zip (map State [0..]) fL
        !fG = fromDistribution fD

-- | See 'jump'.
treeProbMatrixToTreeGenerator :: RTree a ProbMatrix -> RTree a [Generator State]
treeProbMatrixToTreeGenerator t = tG
  where
    -- Create a tree with the probability matrices as list of row vectors.
    !tL = fmap toLists t
    -- A complicated double map. We need to create generators for each branch on
    -- the tree (fmap) and for each target state on each branch (map).
    !tG = (fmap . map) (fromDistribution . fromList . zip (map State [0..])) tL

-- | Simulate data (states at the leaves) for a tree with transition
-- probabilities on its branches and with the stationary distribution of states
-- at the root. The state at the root is randomly chosen from the stationary
-- distribution and the states at the nodes and leaves are randomly chosen
-- according to the transition probabilities.
simulateSite :: (MonadRandom m) =>
                Generator State
             -> RTree a [Generator State]
             -> m [(a, State)]
simulateSite f t = do
  !rootState <- getSample f
  populateAndFlattenTree t rootState

-- | Randomly draw an index according to a given generator. Use the stationary
-- distribution and rooted tree at the drawn index to simulate a site. This is
-- useful for simulation, e.g., Gamma rate heterogeneity models.
simulateSiteGen :: (MonadRandom m) =>
                   Generator Int
                -> [Generator State]
                -> [RTree a [Generator State]]
                -> m [(a, State)]
simulateSiteGen gen fs trs = do
  !i <- getSample gen
  let !f = fs  !! i
      !t = trs !! i
  simulateSite f t


{-# LANGUAGE GADTs #-}
module Math.IRT.Internal.LogLikelihood where

import Numeric.AD (Mode, Scalar)
import Statistics.Distribution (Distribution)

class (Distribution d) => LogLikelihood d where
    logLikelihood :: (Mode a, Floating a, Scalar a ~ Double) => Bool -> d -> a -> a

logLikeFunc :: (Distribution d, Mode a, Floating a, Scalar a ~ Double) => (d -> a -> a) -> Bool -> d -> a -> a
logLikeFunc f True  = (log .) . f
logLikeFunc f False = ((log . (1-)) .) . f


{-# Language DerivingVia                #-}
{-# Language FlexibleInstances          #-}
{-# Language GeneralizedNewtypeDeriving #-}
{-# Language MultiParamTypeClasses      #-}
{-# Language UndecidableInstances       #-}

module Geometry.Testing (
    DInt (..)
  , Parameter (..)
  , Position (..)
  , AnyPoint (..)
  , AnyCircle (..)
  , AnyAngle (..)
  , AnyLine (..)
  , AnyRay (..)
  , AnySegment (..)
  , Motion (..), appMotion
  , Nontrivial (..)
  , NonDegenerate (..)
  , (<==>)
) where

import Data.Complex
import Data.Foldable
import Data.Monoid
import Test.QuickCheck hiding (scale)
import Test.QuickCheck.Modifiers
import Data.Tuple.Extra (second)

import Geometry
import Geometry.Circle (mkCircle)
import Geometry.Base (roundUp)

infix 0 <==>
a <==> b = (a ==> b) .&&. (b ==> a)

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

{- |  A type for an interger double number.

>>>  sample (arbitrary :: Gen DInt)
DInt 0.0
DInt 2.0
DInt 0.0
DInt (-1.0)
DInt 7.0
...
-}
newtype DInt = DInt Double deriving Show

instance Arbitrary DInt where
  arbitrary = DInt . fromIntegral <$> (arbitrary :: Gen Int)
  shrink (DInt l) = DInt <$> shrink l

{- |  A type for a double number parameterizing the `Manifold`.

>>>  sample (arbitrary :: Gen Parameter)
Parameter 0.6977012699848539
Parameter 0.5321234247634057
Parameter 9.832672492232597e-2
Parameter 0.9510866292762408
...
-}
newtype Parameter = Parameter { getParameter :: Double} deriving Show

instance Arbitrary Parameter where
  arbitrary = Parameter <$> choose (0, 1)
  shrink (Parameter l) = Parameter <$> shrink l

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

{- |  A type for an affine position (`XY`, `Cmp`, `Point` etc).

>>>  sample (arbitrary :: Gen (Position XY))
Position (0.0,0.0)
Position (-1.0,2.0)
Position (-3.0,3.0)
Position (6.0,4.0)
Position (-1.0,-1.0
...
-}
newtype Position a = Position a
  deriving (Show, Affine, Trans, Metric, Eq)

instance (Eq a, Pnt a, Arbitrary a) =>
         Arbitrary (Position a) where
  arbitrary = Position . roundUp 1 <$> arbitrary
  shrink = shrinkPos 0.5

shrinkPos :: (Eq a, Affine a, Metric a) => Double -> a -> [a]
shrinkPos d x = res
  where res = filter (/= x) $
              map (roundUp d) $
              takeWhile (\p -> distance x p >= d/2) $
              map (\s -> asCmp $ (1 - s) * cmp x) $
              iterate (/2) 1

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

instance Arbitrary Direction where
  arbitrary = oneof [asDeg  <$> arbitrary, asCmp <$> arbitrary]
  shrink = shrinkPos 1

------------------------------------------------------------
instance Arbitrary Point where
  arbitrary = Point <$> arbitrary
  shrink = shrinkPos 1

{- |  A wrapper for a `Point`. -}
newtype AnyPoint = AnyPoint Point
  deriving (Show, Arbitrary)

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

instance Arbitrary Circle where
  arbitrary = do (Position c) <- arbitrary
                 (Position r) <- arbitrary :: Gen (Position Cmp)
                 return (mkCircle (norm r) c # rotateAt' c (angle r))

  shrink cir =
    do Position c <- shrink (Position (center cir))
       r <- shrink (radius cir)
       return $ mkCircle r c

{- |  A wrapper for nondegenative `Circle`. -}
newtype AnyCircle = AnyCircle Circle
  deriving Show
  deriving Arbitrary via Nontrivial Circle

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

instance Arbitrary Angle where
  arbitrary = do (Position p) <- arbitrary
                 Angle p <$> arbitrary <*> arbitrary

  shrink an =
    do Position p <- shrink (Position (refPoint an))
       s <- shrink (angleStart an)
       e <- shrink (angleEnd an)
       return $ Angle p s e

{- |  A wrapper for nontrivial `Angle`. -}
newtype AnyAngle = AnyAngle Angle
  deriving Show
  deriving Arbitrary via Nontrivial Angle

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

instance Arbitrary Segment where
  arbitrary =
    do Position p1 <- arbitrary
       Position p2 <- arbitrary
       return $ Segment (p1, p2)

  shrink l = let (p1, p2) = refPoints l
    in do Position p1' <- shrink (Position p1)
          Position p2' <- shrink (Position p2)
          return $ Segment (p1', p2')

instance Arbitrary Line where
  arbitrary = asLine <$> (arbitrary :: Gen Segment)
  shrink l = asLine <$> shrink (asSegment l)

instance Arbitrary Ray where
  arbitrary = asRay <$> (arbitrary :: Gen Segment)
  shrink l = asRay <$> shrink (asSegment l)

{- |  A wrapper for nontrivial `Line`. -}
newtype AnyLine = AnyLine Line
  deriving Show
  deriving Arbitrary via Nontrivial Line

{- |  A wrapper for nontrivial `Ray`. -}
newtype AnyRay = AnyRay Ray
  deriving Show
  deriving Arbitrary via Nontrivial Ray

{- |  A wrapper for nontrivial `Segment`. -}
newtype AnySegment = AnySegment Segment
  deriving Show
  deriving Arbitrary via Nontrivial Segment

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

{- |  A wrapper for a motion `translate`, `rotate` or `reflect`. -}
newtype Motion a = Motion (String, Endo a)

appMotion :: Motion a -> a -> a
appMotion (Motion (_, m)) = appEndo m

instance Show (Motion a) where
  show (Motion (s, _)) = s

instance Trans a => Arbitrary (Motion a) where
  arbitrary = Motion . fold <$> listOf (oneof motions)
    where
      motions = [ label "Tr " translate <$> (arbitrary :: Gen XY)
                , label "Rot " rotate <$> arbitrary
                , label "Ref " reflect <$> arbitrary ]
      label l t x = (l <> show x <> " ", Endo $ t x)

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

{- |  A wrapper for a nontrivial `Figure`. -}
newtype Nontrivial a = Nontrivial a deriving
  ( Eq
  , Show
  , Figure
  , Affine
  , Trans
  , Curve
  , ClosedCurve
  , Manifold
  , PiecewiseLinear
  , Polygonal
  , APoint
  , Linear
  , Circular
  , Angular
  )

instance (Arbitrary a, Figure a) => Arbitrary (Nontrivial a) where
  arbitrary = Nontrivial <$> arbitrary `suchThat` isNontrivial
  shrink (Nontrivial l) = Nontrivial <$> filter isNontrivial (shrink l)

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

{- |  A wrapper for a nondegenerate `Figure`. -}
newtype NonDegenerate a = NonDegenerate a deriving
  ( Eq
  , Show
  , Figure
  , Affine
  , Trans
  , Curve
  , ClosedCurve
  , Manifold
  , PiecewiseLinear
  , Polygonal
  , APoint
  , Linear
  , Circular
  , Angular
  )

instance (Arbitrary a, PiecewiseLinear a) => Arbitrary (NonDegenerate a) where
  arbitrary = NonDegenerate <$> arbitrary `suchThat` isNondegenerate
  shrink (NonDegenerate l) = NonDegenerate <$> filter isNondegenerate (shrink l)

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

instance Arbitrary Triangle where
  arbitrary = (`suchThat` isNontrivial) $ do
    Position p1 <- arbitrary
    Position p2 <- arbitrary
    Position p3 <- arbitrary
    pure $ Triangle [p1, p2, p3]

  shrink (Triangle [p1, p2, p3]) = filter isNontrivial $ do
    p1' <- shrink p1
    p2' <- shrink p2
    p3' <- shrink p3
    pure $ Triangle [p1', p2', p3']

instance Arbitrary RightTriangle where
  arbitrary = (`suchThat` isNondegenerate) $ do
    Positive a <- arbitrary
    Positive b <- arbitrary
    m <- arbitrary
    pure $ aRightTriangle # scaleX a # scaleY b # appMotion m
  shrink (RightTriangle t) = RightTriangle <$> shrink t


module Lib.Practice.P3_1 (
    f3_1
) where

import Lib.Practice.P3 (Input( Fire, NotFire ), input2Int, inputs)
import Numeric.LinearAlgebra

f3_1 = showOutput ("XOR", gXor)

gAnd = Single (SNeuron (vector [0.5, 0.5]) (-0.7))
gNand = Single (SNeuron (vector [-0.5, -0.5]) 0.7)
gOr = Single (SNeuron (vector [0.5, 0.5]) (-0.2))
gXor = Multi MNeuron { inputLayer = [gNand, gOr], intermediateLayer = gAnd }

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

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

data MNeuron = MNeuron {
    inputLayer :: [Neuron],
    intermediateLayer :: Neuron
} deriving Show

data Neuron = Single SNeuron | Multi MNeuron deriving Show
(-->) :: [Input] -> Neuron -> Input
inputs --> Single SNeuron { weight = w, bias = b } = result
    where inputV    = fromList $ map input2Int inputs
          sum'      = sumElements $ inputV * w
          result'   = b + sum'
          result    = case () of
            _ | result' >= 0 -> Fire
            _                -> NotFire
inputs --> Multi MNeuron { inputLayer = inpLs, intermediateLayer = intL } = result
    where inputs'   = map (inputs -->) inpLs
          result    = inputs' --> intL


{-
 - *Main Data.Complex> h = concat [[1..5], replicate 11 0]
 - *Main Data.Complex> x = concat [[1..10], replicate 6 0]
 - *Main Data.Complex> convFft h x
 - [1.0,4.000000000000007,10.0,20.0,35.0,49.99999999999999,65.0,80.0,95.0,110.0,114.0,106.0,85.0,50.00000000000001,7.105427357601002e-15,7.105427357601002e-15]
 - *Main Data.Complex>
 -}
import Data.Complex

convOne :: (Num a) => [a] -> [a] -> Int -> a
convOne h x n    = sum (zipWith (*) kernel input)
    where kernel = drop (-begin) (reverse h)
          input  = drop begin x
          begin  = n - length h + 1

conv :: (Num a) => [a] -> [a] -> [a]
conv h x = [convOne h x n | n <- [0..(length x + length h - 2)]]

isPowerOfTwo :: Int -> Bool
isPowerOfTwo 1 = True
isPowerOfTwo x
    | mod x 2 == 0 = isPowerOfTwo (div x 2)
    | otherwise = False

fftRaw :: (RealFloat a) => [a] -> Int -> Int -> [Complex a]
fftRaw _ 0 _ = []
fftRaw [] _  _ = []
fftRaw (x0:_) 1 _ = [x0 :+ 0]
fftRaw x n s = zipWith (+) x1 x2 ++ (zipWith (-) x1 x2)
    where x1 = fftRaw x (div n 2) (2 * s)
          x2 = zipWith (*) [exp (0 :+ (-2 * pi * fromIntegral k / (fromIntegral n))) | k <- [0..((div n 2) - 1)]] (fftRaw (drop s x) (div n 2) (2 * s))

fft :: (RealFloat a) => [a] -> [Complex a]
fft x
    | isPowerOfTwo n = fftRaw x n 1
    | otherwise = error "FFT works only for powers of two"
    where n = length x

ifftRaw :: (RealFloat a) => [Complex a] -> Int -> Int -> [Complex a]
ifftRaw _ 0 _ = []
ifftRaw [] _  _ = []
ifftRaw (x0:_) 1 _ = [x0]
ifftRaw x n s = zipWith (+) x1 x2 ++ (zipWith (-) x1 x2)
    where x1 = ifftRaw x (div n 2) (2 * s)
          x2 = zipWith (*) [exp (0 :+ (2 * pi * fromIntegral k / (fromIntegral n))) | k <- [0..((div n 2) - 1)]] (ifftRaw (drop s x) (div n 2) (2 * s))

ifft :: (RealFloat a) => [Complex a] -> [a]
ifft x
    | isPowerOfTwo n = [realPart v / (fromIntegral n) | v <- ifftRaw x n 1]
    | otherwise = error "IFFT works only for powers of two"
    where n = length x

convFft :: (RealFloat a) => [a] -> [a] -> [a]
convFft h x
    | length x == length h = ifft (zipWith (*) (fft h) (fft x))
    | otherwise = error "Kernel and input data must have the same length"