module Proginduction where

import MHMonad
import Distr
import Data.Monoid
import Statistics.Distribution
import Statistics.Distribution.Normal (normalDistr)
import Graphics.Rendering.Chart
import Graphics.Rendering.Chart.Backend.Diagrams
import Graphics.Rendering.Chart.State
import Data.Colour
import Data.Colour.Names
import Data.Default.Class
import Control.Lens

-- A grammar for a simple expression language
data BinOp = Add | Multiply
data UnOp = Neg
data Expr = BinOp BinOp Expr Expr
          | UnOp UnOp Expr
          | Number Double
          | Var
          | Cond BoolExpr Expr Expr
data BoolExpr = Leq Expr Expr

-- Simple evaluation function for expressions
eval :: Expr -> Double -> Double
eval (BinOp Add e1 e2) x = (eval e1 x) + (eval e2 x)
eval (BinOp Multiply e1 e2) x = (eval e1 x) * (eval e2 x)
eval (UnOp Neg e ) x = - (eval e x)
eval (Number r) x = r
eval (Cond (Leq e1 e2) e3 e4) x = if (eval e1 x) <= (eval e2 x) then (eval e3 x) else (eval e4 x)
eval Var x = x

-- a probabilistic context free grammar
pcfgbinop :: Meas BinOp
pcfgbinop = do n <- categorical [0.5,0.5]
               case n of 0 -> return Add
                         1 -> return Multiply

pcfgexpr :: Meas Expr
pcfgexpr = do n <- categorical [0.2,0.1,0.3,0.39,0.01]
              case n of 0 -> do {o <- pcfgbinop ; e1 <- pcfgexpr ; e2 <- pcfgexpr ; return $ BinOp o e1 e2 }
                        1 -> do {e <- pcfgexpr ; return $ UnOp Neg e }
                        2 -> do {r <- normal 0 3 ; return $ Number r }
                        3 -> return $ Var
                        4 -> do {e1 <- return Var ; r <- normal 0 4 ; e3 <- pcfgexpr ; e4 <- pcfgexpr ; return $ Cond (Leq e1 (Number r)) e3 e4 }

-- length of an expression
len :: Expr -> Integer
len (BinOp _ e1 e2) = len e1 + len e2 + 1
len (UnOp _ e1) = len e1 + 1
len (Number r) = 1
len Var = 1
len (Cond (Leq e1 e2) e3 e4) = len e1 + len e2 + len e3 + len e4 + 1

-- a random expression is drawn from the context free grammar,
-- but then we weight to get a more reasonable length
randexpr :: Meas (Double -> Double)
randexpr = do { e <- pcfgexpr ;
                score ((normalPdf 25 5) (fromIntegral $ len e)) ;
                return $ \x -> eval e x }

-- functions for plotting a histogram
totalweight :: [(Integer,Double)] -> Double
totalweight rxs = sum $ map (\(_,r) -> r) rxs

toHistogram :: [(Integer,Double)] -> Integer -> Integer -> [Double]
toHistogram rxs from to = if from > to then [] else
  ((sum $ map (\(x,r) -> if x == from then r else 0) rxs) / (totalweight rxs)) : (toHistogram rxs (from + 1) to)

chart :: [Double] -> Renderable ()
chart rs = toRenderable layout
 where
  layout =
        layout_title .~ "Sample Bars" ++ btitle
      $ layout_title_style . font_size .~ 10
      $ layout_x_axis . laxis_generate .~ autoIndexAxis alabels
      $ layout_y_axis . laxis_override .~ axisGridHide
      $ layout_left_axis_visibility . axis_show_ticks .~ False
      $ layout_plots .~ [ plotBars bars2 ]
      $ def :: Layout PlotIndex Double

  bars2 = plot_bars_titles .~ [""]
      $ plot_bars_values .~ addIndexes (map (\x -> [x]) rs) -- [45,30],[30,20],[70,25]]
--      $ plot_bars_style .~ BarsClustered
      $ plot_bars_spacing .~ BarsFixGap 30 5
      $ plot_bars_item_styles .~ map mkstyle (cycle defaultColorSeq)
      $ def

  alabels = map show [0 .. ((length rs)-1)]

  btitle = ""
  bstyle = Just (solidLine 1.0 $ opaque black)
  mkstyle c = (solidFillStyle c, bstyle)

-- plot a histogram of lengths when the weighted pcfg grammar is used
test1 =
  let test =
        do
          e <- pcfgexpr
          -- score ((density $ normalDistr 25 5) (fromIntegral $ len e))
          return $ len e
  in
  do
    rxs' <- weightedsamples test
    let rxs = map (\(x,r) -> (x,r)) rxs'
    _ <- renderableToFile def "histogram.svg" (chart (toHistogram (take 100000 rxs) 0 40))
    return ()

sample_fun f =
  [ (x, f x) | x <- [(-0.25),(-0.25+0.1)..6.2]]

plot_funs :: String -> [(Double,Double)] -> [Double -> Double] -> IO ()
plot_funs filename dataset funs =
  let graphs  = map sample_fun funs                 in
  let my_lines  = plot_lines_style . line_color .~ blue `withOpacity` 0.1
                $ plot_lines_values .~ graphs $ def in
  let my_dots = plot_points_style .~ filledCircles 4 (opaque black)
              $ plot_points_values .~ dataset
              $ def in
  let my_layout = layout_plots .~ [toPlot my_lines , toPlot my_dots]
                $ layout_x_axis .
                  laxis_generate .~ scaledAxis def (0,6)
                $ layout_y_axis . laxis_generate .~ scaledAxis def (-2,10)
                $ def in
  let graphic =  toRenderable my_layout in
  do
     putStr ("Generating " ++ filename ++ "...")
     renderableToFile def filename graphic;
     putStrLn (" Done!")
     return ()

dataset :: [(Double, Double)]
dataset = [(0,0.6), (1, 0.7), (2,1.2), (3,3.2), (4,6.8), (5, 8.2), (6,8.4)]

fit prior dataset = do f <- prior
                       mapM (\(x,y) -> score $ (density $ normalDistr (f x) 0.3) y) dataset
                       return f

every n xs = case drop (n-1) xs of
              (y:ys) -> y : every n ys
              [] -> []

-- Plot some fitted expressions
test3 =
  do
    fs' <- mh $ fit randexpr dataset
    let fs = map (\(f,_)->f) $ take 10 $ every 1000 $ drop 1000000 $ fs'
    plot_funs "prog-ind.svg" dataset fs
    return ()

-- Histogram of the posterior length
test4 =
  do
    let posteriorLength =
          do e <- pcfgexpr
             mapM (\(x,y) -> score $ (density $ normalDistr (eval e x) 0.3) y) dataset
             return (len e)
    ls <- mh $ posteriorLength
    let rxs = map (\(x,r) -> (x,getProduct r)) ls
    _ <- renderableToFile def "histogram.svg" (chart (toHistogram (take 100000 $ drop 100000 rxs) 0 40))
    return ()


{-# LANGUAGE CPP #-}
{- MatInv.hs
 - 4x4 Matrix invert function.
 - Direct implementation obtained from Maple's CodeGeneration
 -
 - Timothy A. Chagnon
 - CS 636 - Spring 2009
 -}
module MatInv where

#ifdef USING_HMATRIX
import MathHmatrix
import Numeric.LinearAlgebra
inverse4f :: Mat4f -> Mat4f
inverse4f = inv

#else
import MathDirect

-- Analytical solution via Maple (see minv.mw)
inverse4f :: Mat4f -> Mat4f
inverse4f   (Mat4f  (Vec4f m11 m12 m13 m14)
                    (Vec4f m21 m22 m23 m24)
                    (Vec4f m31 m32 m33 m34)
                    (Vec4f m41 m42 m43 m44)) =
    let t1 = m22 * m43 in
    let t3 = m22 * m44 in
    let t5 = m42 * m24 in
    let t7 = m44 * m32 in
    let t9 = m43 * m32 in
    let t11 = m42 * m23 in
    let t14 = m11 * m22 in
    let t15 = m43 * m34 in
    let t17 = m44 * m33 in
    let t19 = m11 * m42 in
    let t20 = m24 * m33 in
    let t22 = m11 * m44 in
    let t23 = m32 * m23 in
    let t25 = m11 * m43 in
    let t26 = m32 * m24 in
    let t28 = m23 * m34 in
    let t30 = m31 * m14 in
    let t32 = m22 * m41 in
    let t33 = m13 * m34 in
    let t35 = m14 * m33 in
    let t37 = m31 * m13 in
    let t39 = m43 * m31 in
    let t40 = m12 * m24 in
    let t42 = m21 * m12 in
    let t44 = -t14 * t15 + t14 * t17 - t19 * t20 - t22 * t23 + t25 * t26 + t19 * t28 + t1 * t30 + t32 * t33 - t32 * t35 - t3 * t37 - t39 * t40 - t17 * t42 in
    let t46 = m42 * m21 in
    let t50 = m41 * m12 in
    let t52 = m21 * m13 in
    let t54 = m41 * m14 in
    let t57 = m44 * m31 in
    let t58 = m12 * m23 in
    let t60 = m41 * m13 in
    let t62 = m21 * m14 in
    let t65 = t5 * t37 - t46 * t33 + t46 * t35 + t15 * t42 + t50 * t20 + t7 * t52 + t54 * t23 - t50 * t28 + t57 * t58 - t60 * t26 - t9 * t62 - t11 * t30 in
    let t67 = 1.0 / (t44 + t65) in
    let t121 = m44 * m21 in
    let t125 = m43 * m21 in
    let t146 = m41 * m32 in
    let t152 = m42 * m31 in
    let t172 = m32 * m11 in
    let t174 = m32 * m21 in
    let t176 = m31 * m12 in
    let r00 = (-t1 * m34 + t3 * m33 - t5 * m33 - t7 * m23 + t9 * m24 + t11 * m34) * t67 in
    let r01 = -(t17 * m12 + m42 * m13 * m34 - m42 * m14 * m33 - t15 * m12 - t7 * m13 + t9 * m14) * t67 in
    let r02 = (t1 * m14 - t3 * m13 - m43 * m12 * m24 + t5 * m13 + m44 * m12 * m23 - t11 * m14) * t67 in
    let r03 = -(-m22 * m13 * m34 + m22 * m14 * m33 - t40 * m33 - m14 * m32 * m23 + t58 * m34 + m13 * m32 * m24) * t67 in
    let r10 = -(t17 * m21 + m41 * m23 * m34 - t15 * m21 - t57 * m23 + t39 * m24 - m41 * m24 * m33) * t67 in
    let r11 = (-t25 * m34 + t22 * m33 + t39 * m14 + t60 * m34 - t54 * m33 - t57 * m13) * t67 in
    let r12 = -(t22 * m23 - t25 * m24 - t121 * m13 - t54 * m23 + t60 * m24 + t125 * m14) * t67 in
    let r13 = (-m11 * m24 * m33 + m11 * m23 * m34 + m24 * m31 * m13 - t52 * m34 + t62 * m33 - m23 * m31 * m14) * t67 in
    let r20 = (t32 * m34 - t3 * m31 + t5 * m31 - t46 * m34 + t7 * m21 - t146 * m24) * t67 in
    let r21 = -(t22 * m32 - t19 * m34 + t152 * m14 + t50 * m34 - t54 * m32 - t57 * m12) * t67 in
    let r22 = (t22 * m22 - t121 * m12 - t54 * m22 - t19 * m24 + t46 * m14 + t50 * m24) * t67 in
    let r23 = -(m34 * m11 * m22 - m34 * m21 * m12 - t30 * m22 - t172 * m24 + t174 * m14 + t176 * m24) * t67 in
    let r30 = -(-t1 * m31 + t32 * m33 - t46 * m33 + t11 * m31 - t146 * m23 + t9 * m21) * t67 in
    let r31 = (-t19 * m33 + t25 * m32 - t39 * m12 + t152 * m13 + t50 * m33 - t60 * m32) * t67 in
    let r32 = -(t25 * m22 - t125 * m12 - t60 * m22 - t19 * m23 + t46 * m13 + t50 * m23) * t67 in
    let r33 = t67 * (m33 * m11 * m22 - m33 * m21 * m12 - t37 * m22 - t172 * m23 + t174 * m13 + t176 * m23) in
    Mat4f   (Vec4f r00 r01 r02 r03)
            (Vec4f r10 r11 r12 r13)
            (Vec4f r20 r21 r22 r23)
            (Vec4f r30 r31 r32 r33)

#endif


{-# LANGUAGE DataKinds #-}
{-# LANGUAGE FlexibleContexts #-}
module Main where

import Lib
import Numeric.Vector.Sized
import Numeric.Layer.Affine
import Numeric.Network
import Numeric.Layer.Affine
import Numeric.Layer.Sigmoid
import Numeric.Loss.MeanSquaredError
import qualified Data.Array.Accelerate.Interpreter as AI
import qualified Data.Array.Accelerate.LLVM.Native as Native

v :: SizedVector 3
v = use (fromList [1, 2, 3]) .+. use (fromList [2, 3, 4])

xs :: SizedMatrix 2 3
xs = use (fromList [1 .. 6])
-- 1 2 3
-- 4 5 6

ys :: SizedMatrix 3 4
ys = use (fromList [1 .. 12])
-- 1 2 3 4
-- 5 6 7 8
-- 9 10 11 12
--
u = use (fromList [1, 2]) :: SizedVector 2
zs = use (fromList [1, 2, 3, 4]) :: SizedMatrix 2 2

aff = affine xs v u

net :: Network '[Affine 3 2, Sigmoid 2, Affine 2 2, Sigmoid 2]
net = Affine ws b :~> Sigmoid () :~> Affine ws' b' :~> Last (Sigmoid ())
  where
    ws  = fromList [1, 2, 3, 4, 5, 6]
    b   = fromList [1, 2]
    ws' = fromList [-2, 2, 0, 1]
    b'  = fromList [0, 1]

main :: IO ()
main = do
  print net'
  where
    net' = train
      net
      (MeanSquaredError () :: MeanSquaredError 2)
      (use $ fromList [-1, -2, -3])
      (use $ fromList [1, 2])

  -- print (runUsing Native.run v)
  -- print (runUsing Native.run (xs >< ys))
  -- print (runUsing Native.run (xs #> v))
  -- print (runUsing Native.run (u <# zs))
  -- print (runUsing Native.run aff)


-- | Hamiltonian Monte Carlo.  See, ex: Neal (2012)
--   http://arxiv.org/pdf/1206.1901.pdf.

module Strategy.Hamiltonian (hamiltonian) where

import Control.Monad.State.Strict
import Data.Maybe
import qualified Data.Vector.Storable as V
import Math.Probably.Sampler
import Math.Probably.Types
import Math.Probably.Utils
import Numeric.LinearAlgebra

-- | A Hamiltonian transition with step size and leapfrog increment tunables.
hamiltonian :: Maybe Double -> Maybe Int -> Transition (Double, Int)
hamiltonian e l = do
  Chain (ds, q0) target _ (te, tl) <- get
  let stepSize = fromMaybe te e
      nDisc    = fromMaybe tl l
      lTarget  = curry (logObjective target) ds
      glTarget = handleGradient $ gradient target

  r0 <- V.replicateM (dim q0) (lift unormal)
  zc <- lift unit
  let (q, r) = leapfrogIntegrator glTarget (q0, r0) stepSize nDisc
      next   = (ds, nextState lTarget (q0, q) (r0, r) zc)

  put $ Chain next target (logObjective target next) (stepSize, nDisc)
  return next

-- | The leapfrog or Stormer-Verlet integrator.
leapfrogIntegrator :: Gradient -> Particle -> Double -> Int -> Particle
leapfrogIntegrator glTarget (q0, r0) e = go q0 r0 where
  go q r 0 = (q, r)
  go q r n = let (q1, r1) = leapfrog glTarget (q, r) e
             in  go q1 r1 (pred n)

-- | The acceptance probability of a move.
acceptProb :: (Vector Double -> Double) -> Particle -> Particle -> Double
acceptProb lTarget (q0, q1) (r0, r1) = exp . min 0 $
  auxilliaryTarget lTarget (q1, r1) - auxilliaryTarget lTarget (q0, r0)

-- | The next state, given a proposal and random value from (0, 1).
nextState
  :: (Vector Double -> Double)
  -> Particle
  -> Particle
  -> Double
  -> ContinuousParams
nextState lTarget position momentum z
    | z < pAccept = snd position
    | otherwise   = fst position
  where
    pAccept = acceptProb lTarget position momentum


{-# OPTIONS_GHC -fno-warn-type-defaults #-}
{-# LANGUAGE DataKinds, FlexibleContexts, FlexibleInstances, GADTs #-}
{-# LANGUAGE MultiParamTypeClasses, NoMonomorphismRestriction      #-}
{-# LANGUAGE OverloadedStrings, PolyKinds, TypeFamilies            #-}
{-# LANGUAGE UndecidableInstances                                  #-}
{-# OPTIONS_GHC -fno-warn-type-defaults -fno-warn-orphans -Wall #-}
module Main where
import           Algebra.Bridge.Singular
import           Algebra.Prelude.Core
import           Control.Monad                   (void)
import           Control.Parallel.Strategies
import qualified Data.Text                       as T
import qualified Data.Vector.Unboxed             as V
import           Numeric.Field.Fraction          (Fraction)
import           Prelude                         (read)
import           Statistics.Resampling
import           Statistics.Resampling.Bootstrap
import           Statistics.Types
import qualified System.Random.MWC               as Rand

i2 :: [OrderedPolynomial (Fraction Integer) Grevlex 5]
i2 =  [35 * y^4 - 30*x*y^2 - 210*y^2*z + 3*x^2 + 30*x*z - 105*z^2 +140*y*t - 21*u
      ,5*x*y^3 - 140*y^3*z - 3*x^2*y + 45*x*y*z - 420*y*z^2 + 210*y^2*t -25*x*t + 70*z*t + 126*y*u
      ]
     where [t,u,x,y,z] = vars

i3 :: [OrderedPolynomial (Fraction Integer) Grevlex 4]
i3 = [ x^31 - x^6 - x- y, x^8 - z, x^10 -t]
  where
    [t,x,y,z] = vars

i4 :: [OrderedPolynomial (Fraction Integer) Grevlex 4]
i4 = [ w+x+y+z, w*x+x*y+y*z+z*w, w*x*y + x*y*z + y*z*w + z*w*x, w*x*y*z - 1]
  where
    [x,y,z,w] = vars

benchIdeal :: IsSingularPolynomial poly => Text -> Ideal poly -> IO Double
benchIdeal fun i = fmap ((/1000) . read . T.unpack) $ singular $
  toProg fun i

toProg :: (SingularOrder (Arity poly) (MOrder poly), SingularCoeff (Coefficient poly), IsOrderedPolynomial poly) => Text -> Ideal poly -> Text
toProg fun i =
  let expr = funE fun [idealE' i]
  in prettySingular $ do
    directC "system(\"--ticks-per-sec\", 1000000)"
    void $ ringC "R" expr
    declOnlyC IdealT "G"
    directC "timer=1"
    directC "int t=rtimer"
    letC "G" expr
    directC "print(rtimer-t)"
    directC "exit"

analyse :: (IsSingularPolynomial poly) => String -> Text -> Ideal poly -> IO ()
analyse lab fun ideal = do
  gen <- Rand.create
  i2Gr <- V.replicateM 200 $ benchIdeal fun ideal
  res  <- resample gen [Mean, StdDev] 1000 i2Gr
  let [Estimate mn (ConfInt lbmn ubmn _) ,Estimate dv _]
        = bootstrapBCA cl95 i2Gr res
  putStrLn lab
  putStrLn $ unlines $ map ('\t':)
    ["Mean:\t" ++ show mn ++ "(ms)"
    ,"MeanLB:\t" ++ show lbmn
    ,"MeanUB:\t" ++ show ubmn
    ,"StdDev:\t" ++ show dv
    ]

runTestCases :: (IsMonomialOrder n o, SingularCoeff k, KnownNat n)
             => String -> Ideal (OrderedPolynomial k o n) -> IO ()
runTestCases lab i = do
  analyse (lab ++ " (Lex, Sing(gr))") "groebner" $ fmap (changeOrder Lex) i
  analyse (lab ++ " (Lex, Sing(sba))") "sba" $ fmap (changeOrder Lex) i
  analyse (lab ++ " (Grevlex, Sing(gr))") "groebner" $ fmap (changeOrder Grevlex) i
  analyse (lab ++ " (Grevlex, Sing(sba))") "sba" $ fmap (changeOrder Grevlex) i

main :: IO ()
main = do
  ideal2 <- return $! (toIdeal i2 `using` rdeepseq)
  ideal3 <- return $! (toIdeal i3 `using` rdeepseq)
  ideal4 <- return $! (toIdeal i4 `using` rdeepseq)

  runTestCases "I1" ideal2
  runTestCases "Cyclic4" ideal4
  analyse "I3 (Grevlex, Sing(gr))" "groebner" ideal3
  analyse "I3 (Grevlex, Sing(sba))" "sba" ideal3
  -- analyse "groebner" ideal4
  -- analyse "sba" ideal4
  -- analyse "groebner" ideal3
  -- analyse "sba" ideal3
  -- defaultMainWith defaultConfig $
  --      mkTestCases (Left 1) ideal2
  --   ++ mkTestCases (Right "Cyclic-4") ideal4
  --   ++ [mkTC "grevlex03" ideal3]


#!/usr/bin/env stack
-- stack runghc --resolver lts-16 --package ad --package hmatrix-vector-sized-0.1.3.0 --package vector-sized --package hmatrix --package comonad --package free

-- | Source file accompanying
-- https://blog.jle.im/entry/hamiltonian-dynamics-in-haskell.html

{-# LANGUAGE DataKinds           #-}
{-# LANGUAGE RankNTypes          #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE ViewPatterns        #-}

import           Data.Foldable
import           Data.Maybe
import           GHC.TypeLits
import           Numeric.AD
import           Numeric.LinearAlgebra.Static
import           Numeric.LinearAlgebra.Static.Vector
import qualified Control.Comonad                     as C
import qualified Control.Comonad.Cofree              as C
import qualified Data.Vector.Sized                   as V

-- | A @'System' m n@ represents a system parameterized by @n@ generalized
-- coordinates, moving in an "underlying" @m@-dimensional cartesian
-- coordinate system.
data System m n = System
    { sysInertia       :: R m                         -- ^ 'm' vector
    , sysCoords        :: R n -> R m                  -- ^ f
    , sysJacobian      :: R n -> L m n                -- ^ J_f
    , sysHessian       :: R n -> V.Vector n (L m n)   -- ^ H_f
    , sysPotential     :: R n -> Double               -- ^ U
    , sysPotentialGrad :: R n -> R n                  -- ^ grad U
    }

-- | A system's position in configuration space
data Config n = Config
    { confPositions  :: R n
    , confVelocities :: R n
    }
  deriving Show

-- | A system's position in phase space
data Phase n = Phase
    { phasePositions :: R n
    , phaseMomenta   :: R n
    }
  deriving Show

-- | Given a position in @n@-dimensional generalized coordinates, return
-- the position in the underlying @m@-dimensional cartesian coordinate
-- system.
underlyingPosition
    :: System m n
    -> R n
    -> R m
underlyingPosition = sysCoords

-- | Give the system's generalized momenum, given its position in
-- configuration space.
momenta
    :: (KnownNat n, KnownNat m)
    => System m n
    -> Config n
    -> R n
momenta s (Config q v) = tr j #> mHat #> j #> v
  where
    j    = sysJacobian s q
    mHat = diag (sysInertia s)

-- | Convert a position in configuration space to a position in phase space
toPhase
    :: (KnownNat n, KnownNat m)
    => System m n
    -> Config n
    -> Phase n
toPhase s c = Phase (confPositions c) (momenta s c)

-- | Shift around the Hessian into the shape expected
tr2 :: (KnownNat m, KnownNat n)
    => V.Vector m (L n n)
    -> V.Vector n (L m n)
tr2 = fmap rowsL . traverse lRows
{-# INLINE tr2 #-}

-- | Vector of vectors to matrix
vec2l
    :: KnownNat n
    => V.Vector m (V.Vector n Double)
    -> L m n
vec2l = rowsL . fmap gvecR
{-# INLINE vec2l #-}

-- | Make a system given its inertias, coordinate functions, and potential
-- energy function
mkSystem
    :: (KnownNat m, KnownNat n)
    => R m
    -> (forall a. RealFloat a => V.Vector n a -> V.Vector m a)
    -> (forall a. RealFloat a => V.Vector n a -> a)
    -> System m n
mkSystem m f u = System
                    -- < convert from      | actual thing | convert to >
    { sysInertia       =                     m
    , sysCoords        =        gvecR      . f            . grVec
    , sysJacobian      = tr   . vec2l      . jacobianT f  . grVec
    , sysHessian       = tr2  . fmap vec2l . hessianF f   . grVec
    , sysPotential     =                     u            . grVec
    , sysPotentialGrad =        gvecR      . grad u       . grVec
    }

-- | Equations of motion for a system at a given position in phase space
hamilEqns
    :: (KnownNat n, KnownNat m)
    => System m n
    -> Phase n
    -> (R n, R n)       -- dq/dt and dp/dt
hamilEqns s (Phase q p) = (dqdt, dpdt)
  where
    j       = sysJacobian s q
    trj     = tr j
    mHat    = diag (sysInertia s)
    kHat    = trj `mul` mHat `mul` j
    kHatInv = inv kHat
    dqdt    = kHatInv #> p
    dpdt    = gvecR bigUglyThing - sysPotentialGrad s q
      where
        bigUglyThing =
          fmap (\j2 -> -p <.> kHatInv #> trj #> mHat #> j2 #> kHatInv #> p)
               (sysHessian s q)

-- | Step a system's position through phase space using Euler integration
stepEuler
    :: (KnownNat n, KnownNat m)
    => System m n       -- ^ the system
    -> Double           -- ^ dt
    -> Phase n          -- ^ q(t) and p(t)
    -> Phase n          -- ^ q(t + dt) and p(t + dt)
stepEuler s dt ph@(Phase q p) = Phase (q + konst dt * dq) (p + konst dt * dp)
  where
    (dq, dp) = hamilEqns s ph

-- | Iterate Euler's method to repeatedly step a system through phase space
runSystem
    :: (KnownNat n, KnownNat m)
    => System m n       -- ^ the system
    -> Double           -- ^ dt
    -> Phase n          -- ^ initial phase
    -> [Phase n]        -- ^ progression of the system using Euler integration
runSystem s dt = go
  where
    go p0 = p0 : go (stepEuler s dt p0)

-- | A simple particle in 2D cartesian space under gravity
simpleSystem :: System 2 2
simpleSystem = mkSystem (vec2 5 5) id pot
  where
    -- potential energy of a gravity field
    -- U(x,y) = 9.8 * y
    pot :: RealFloat a => V.Vector 2 a -> a
    pot xy = 9.8 * (xy `V.index` 1)

-- | An initial position in configuration space, representing a particle
-- at <0,0> with initial velocity <1,3>
simpleConfig0 :: Config 2
simpleConfig0 = Config
    { confPositions  = vec2 0 0
    , confVelocities = vec2 1 3
    }

simpleMain :: IO ()
simpleMain =
    mapM_ (disp 2 . phasePositions)  -- position with 2 digits of precision
  . take 25                          -- 25 steps
  $ runSystem simpleSystem 0.1 (toPhase simpleSystem simpleConfig0)

-- | A pendulum system, parameterized by its angle clockwise from
-- equilibrium
pendulum :: System 2 1
pendulum = mkSystem (vec2 5 5) coords pot      -- 5kg particle
  where
    -- <x,y> = <-0.5 sin(theta), -0.5 cos(theta)>
    -- pendulum of length 0.25
    coords :: RealFloat a => V.Vector 1 a -> V.Vector 2 a
    coords (V.head->theta) = V.fromTuple (- 0.25 * sin theta, - 0.25 * cos theta)
    -- potential energy of gravity field
    -- U(x,y) = 9.8 * y
    pot :: RealFloat a => V.Vector 1 a -> a
    pot q = 9.8 * (coords q `V.index` 1)

-- | An initial pendulum position in configuration space, representing
-- a pendulum at equilibrium with initial angular velocity 0.1 rad/sec
-- clockwise
pendulumConfig0 :: Config 1
pendulumConfig0 = Config
    { confPositions  = 0
    , confVelocities = 0.1
    }

pendulumMain :: IO ()
pendulumMain =
    mapM_ (disp 3 . phasePositions)  -- position with 2 digits of precision
  . take 25                          -- 25 steps
  $ runSystem pendulum 0.1 (toPhase pendulum pendulumConfig0)

main :: IO ()
main = pendulumMain


{-# LANGUAGE BangPatterns, DeriveDataTypeable, DeriveGeneric #-}
-- |
-- Module    : Statistics.Distribution.Normal
-- Copyright : (c) 2009 Bryan O'Sullivan
-- License   : BSD3
--
-- Maintainer  : bos@serpentine.com
-- Stability   : experimental
-- Portability : portable
--
-- The normal distribution.  This is a continuous probability
-- distribution that describes data that cluster around a mean.

module Statistics.Distribution.Normal
    (
      NormalDistribution
    -- * Constructors
    , normalDistr
    , normalFromSample
    , standard
    ) where

import Data.Aeson (FromJSON, ToJSON)
import Control.Applicative ((<$>), (<*>))
import Data.Binary (Binary)
import Data.Binary (put, get)
import Data.Data (Data, Typeable)
import GHC.Generics (Generic)
import Numeric.MathFunctions.Constants (m_sqrt_2, m_sqrt_2_pi)
import Numeric.SpecFunctions (erfc, invErfc)
import qualified Statistics.Distribution as D
import qualified Statistics.Sample as S
import qualified System.Random.MWC.Distributions as MWC

-- | The normal distribution.
data NormalDistribution = ND {
      mean       :: {-# UNPACK #-} !Double
    , stdDev     :: {-# UNPACK #-} !Double
    , ndPdfDenom :: {-# UNPACK #-} !Double
    , ndCdfDenom :: {-# UNPACK #-} !Double
    } deriving (Eq, Read, Show, Typeable, Data, Generic)

instance FromJSON NormalDistribution
instance ToJSON NormalDistribution

instance Binary NormalDistribution where
    put (ND w x y z) = put w >> put x >> put y >> put z
    get = ND <$> get <*> get <*> get <*> get

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

instance D.ContDistr NormalDistribution where
    logDensity = logDensity
    quantile   = quantile

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

instance D.Mean NormalDistribution where
    mean = mean

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

instance D.Variance NormalDistribution where
    stdDev = stdDev

instance D.Entropy NormalDistribution where
  entropy d = 0.5 * log (2 * pi * exp 1 * D.variance d)

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

instance D.ContGen NormalDistribution where
    genContVar d = MWC.normal (mean d) (stdDev d)

-- | Standard normal distribution with mean equal to 0 and variance equal to 1
standard :: NormalDistribution
standard = ND { mean       = 0.0
              , stdDev     = 1.0
              , ndPdfDenom = log m_sqrt_2_pi
              , ndCdfDenom = m_sqrt_2
              }

-- | Create normal distribution from parameters.
--
-- IMPORTANT: prior to 0.10 release second parameter was variance not
-- standard deviation.
normalDistr :: Double            -- ^ Mean of distribution
            -> Double            -- ^ Standard deviation of distribution
            -> NormalDistribution
normalDistr m sd
  | sd > 0    = ND { mean       = m
                   , stdDev     = sd
                   , ndPdfDenom = log $ m_sqrt_2_pi * sd
                   , ndCdfDenom = m_sqrt_2 * sd
                   }
  | otherwise =
    error $ "Statistics.Distribution.Normal.normalDistr: standard deviation must be positive. Got " ++ show sd

-- | Create distribution using parameters estimated from
--   sample. Variance is estimated using maximum likelihood method
--   (biased estimation).
normalFromSample :: S.Sample -> NormalDistribution
normalFromSample xs
  = normalDistr m (sqrt v)
  where
    (m,v) = S.meanVariance xs

logDensity :: NormalDistribution -> Double -> Double
logDensity d x = (-xm * xm / (2 * sd * sd)) - ndPdfDenom d
    where xm = x - mean d
          sd = stdDev d

cumulative :: NormalDistribution -> Double -> Double
cumulative d x = erfc ((mean d - x) / ndCdfDenom d) / 2

complCumulative :: NormalDistribution -> Double -> Double
complCumulative d x = erfc ((x - mean d) / ndCdfDenom d) / 2

quantile :: NormalDistribution -> Double -> Double
quantile d p
  | p == 0         = -inf
  | p == 1         = inf
  | p == 0.5       = mean d
  | p > 0 && p < 1 = x * ndCdfDenom d + mean d
  | otherwise      =
    error $ "Statistics.Distribution.Normal.quantile: p must be in [0,1] range. Got: "++show p
  where x          = - invErfc (2 * p)
        inf        = 1/0


module Chapter08 where

import Numeric.LinearAlgebra.Data
--import Numeric.LinearAlgebra.Algorithms
import Numeric.LinearAlgebra.HMatrix
import Data.List as L

import Chapter04
import Chapter06
import Chapter07

stopWords :: [String]
stopWords = ["a", "about", "above", "above", "across", "after",
    "afterwards", "again", "against", "all", "almost", "alone", "along",
    "already", "also", "although", "always", "am", "among", "amongst",
    "amoungst", "amount", "an", "and", "another", "any", "anyhow",
    "anyone", "anything", "anyway", "anywhere", "are", "around", "as",
    "at", "back", "be", "became", "because", "become", "becomes",
    "becoming", "been", "before", "beforehand", "behind", "being",
    "below", "beside", "besides", "between", "beyond", "bill", "both",
    "bottom", "but", "by", "call", "can", "cannot", "cant", "co", "con",
    "could", "couldnt", "cry", "de", "describe", "detail", "do", "done",
    "dont", "down", "due", "during", "each", "eg", "eight", "either",
    "eleven", "else", "elsewhere", "empty", "enough", "etc", "even",
    "ever", "every", "everyone", "everything", "everywhere", "except",
    "few", "fifteen", "fifty", "fill", "find", "fire", "first", "five",
    "for", "former", "formerly", "forty", "found", "four", "from",
    "front", "full", "further", "get", "give", "go", "got", "had", "has",
    "hasnt", "have", "he", "hence", "her", "here", "hereafter", "hereby",
    "herein", "hereupon", "hers", "herself", "him", "himself", "his",
    "how", "however", "hundred", "i", "ie", "if", "im", "in", "inc",
    "indeed", "interest", "into", "is", "it", "its", "itself", "just",
    "keep", "last", "latter", "latterly", "least", "less", "ltd", "made",
    "many", "may", "me", "meanwhile", "might", "mill", "mine", "more",
    "moreover", "most", "mostly", "move", "much", "must", "my", "myself",
    "name", "namely", "neither", "need", "never", "nevertheless",
    "next", "nine", "no", "nobody", "none", "noone", "nor", "not",
    "nothing", "now", "nowhere", "of", "off", "often", "on", "once",
    "one", "only", "onto", "or", "other", "others", "otherwise", "our",
    "ours", "ourselves", "out", "over", "own", "part", "per", "perhaps",
    "please", "put", "rather", "re", "same", "see", "seem", "seemed",
    "seeming", "seems", "serious", "several", "she", "should", "show",
    "side", "since", "sincere", "six", "sixty", "so", "some", "somehow",
    "someone", "something", "sometime", "sometimes", "somewhere", "still",
    "such", "system", "take", "ten", "than", "that", "the", "their",
    "them", "themselves", "then", "thence", "there", "thereafter",
    "thereby", "therefore", "therein", "thereupon", "these", "they",
    "thick", "thin", "third", "this", "those", "though", "three",
    "through", "throughout", "thru", "thus", "to", "together", "too",
    "top", "toward", "towards", "twelve", "twenty", "two", "un", "under",
    "until", "up", "upon", "us", "very", "via", "want", "was", "we",
    "well", "were", "what", "whatever", "when", "whence", "whenever",
    "where", "whereafter", "whereas", "whereby", "wherein", "whereupon",
    "wherever", "whether", "which", "while", "whither", "who", "whoever",
    "whole", "whom", "whose", "why", "will", "with", "within", "without",
    "would", "yet", "you", "your", "youre", "yours", "yourself",
    "yourselves", "rt", "mt", "u"]

principalComponentAnalysis :: [[Double]] -> Integer -> Matrix Double
principalComponentAnalysis records top =
    tr $ mul (tr topOrderedEvectors) (tr featureVectors)
  where featureVectors = fromLists records
        (_, covMatrix) = meanCov featureVectors
        (_, evectors) = eigSH covMatrix
        topOrderedEvectors = takeColumns (fromInteger top) evectors

euclidianDistanceSqrdFromRow :: Matrix Double -> Integer -> Vector Double
euclidianDistanceSqrdFromRow records row =
    sumOfSquaresVector
  where d = cols records
        copyMatrix = repmat (subMatrix (fromIntegral row, 0) (1,d) records) (rows records) 1
        diffMatrix = records - copyMatrix
        diffSqrdMatrix = diffMatrix * diffMatrix
        sumOfSquaresVector = sum $ toColumns diffSqrdMatrix

order :: Eq a => [a] -> [a] -> [Integer]
order unsorted nubSorted =
    concatMap
        (\x -> L.map toInteger $ L.elemIndices x unsorted)
        nubSorted

findClosestFeaturesToRow :: Matrix Double -> Integer -> Integer -> [Integer]
findClosestFeaturesToRow records row knn =
    take (fromIntegral knn) orderOfFeatures
  where sumOfSquares = toList $ euclidianDistanceSqrdFromRow records row
        orderOfFeatures = order sumOfSquares . nub $ sort sumOfSquares


{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE NoImplicitPrelude #-}
module Main where

import qualified Prelude as P
import Data.Array.Accelerate.Data.Complex as A
import Data.Array.Accelerate as A
-- import qualified Data.Array.Accelerate.Fourier.Adhoc as A
import qualified Data.Array.Accelerate.Interpreter as I
import Data.Array.Accelerate.Interpreter as I
import Data.Array.Accelerate.Data.Fold
import Data.Array.Accelerate.Data.Monoid
import qualified Data.List as L
import qualified Data.Array.Accelerate.LLVM.PTX as PTX

import qualified Data.Csv as CSV
import qualified Data.ByteString.Lazy as BS
import qualified Statistics.Sample as S
import qualified Data.Vector as V

type Arr a = Acc (Vector a)

mean :: (Floating a, FromIntegral Int a) => Arr a -> Acc (Scalar a)
mean xs = unit $ the (fold (+) 0 xs) / (fromIntegral $ length xs)

smooth :: (Floating a, FromIntegral Int a) => Arr a -> Exp Int -> Arr a
smooth xs res = error "LOL"

sumF :: (Elt a, Lift Exp a, P.Num (Exp a)) => Fold (Exp a) (Exp a)
sumF = Fold (lift . Sum) (getSum . unlift)

lengthF :: (Elt a, Lift Exp a, P.Num (Exp a)) => Fold (Exp a) (Exp a)
lengthF = Fold (\_ -> 1) (getSum . unlift)

meanF :: (Elt a, Lift Exp a, P.Fractional (Exp a)) => Fold (Exp a) (Exp a)
meanF = sumF / lengthF

valF :: (Elt a, Monoid (Exp a)) => Fold (Exp a) (Exp a)
valF = Fold P.id P.id

-- | E(X^2)
sumOfSquaresF :: (Elt a, Lift Exp a, P.Num (Exp a)) => Fold (Exp a) (Exp a)
sumOfSquaresF = Fold (lift . Sum . (^(2 :: Exp Int))) (getSum . unlift)

-- | sqrt (E(X^2) - E(X) ^ 2)
standardDeviationF :: (Elt a, Lift Exp a, P.Floating (Exp a)) => Fold (Exp a) (Exp a)
standardDeviationF = sqrt $ (sumOfSquaresF / lengthF) - (meanF * meanF)

-- | E(X - mu) ^ 4
-- differenceQuadsF :: (Elt a, Lift Exp a, P.Fractional (Exp a), Monoid (Exp a)) => Fold (Exp a) (Exp a)
-- differenceQuadsF = (\m -> Fold (\x -> lift $ Sum ((x - m) ^ e)) (getSum . unlift)) P.<$> meanF
--   where
--     e = 4 :: Exp Int

-- | E((X - u)^2)^2
-- otherThing :: (Elt a, Lift Exp a, P.Fractional (Exp a), Monoid (Exp a)) => Fold (Exp a) (Exp a)
-- otherThing = Fold (lift . Sum . (^(2 :: Exp Int))) (getSum . unlift)

-- kurtosisF :: (Elt a, Lift Exp a, P.Fractional (Exp a), Monoid (Exp a)) => Fold (Exp a) (Exp a)
-- kurtosisF = (valF - meanF)

-- | Kurt(X) = E(X - u)^4/(E((X - u)^2)^2)
-- (mean (map (\x -> x - u) xs)) ^ four / ((mean (map (\x -> (x - u) ^ two) xs))^two)
kurtosis :: (Floating a, Elt a, FromIntegral Int a) => Arr a -> Exp a
kurtosis xs = (u4 / (variance P.^ 2))
  where
    u4 = the $ compute $ mean $ map (\x -> (x - u) ^ (4 :: Exp Int)) xs
    variance = the $ compute $ mean $ map (\x -> (x - u) ^ (2 :: Exp Int)) xs
    u = the $ compute $ mean xs

diffs :: (Elt a, P.Num (Exp a)) => Arr a -> Arr a
diffs xs = zipWith (-) (tail xs) xs

roughness :: (Elt a, P.Floating (Exp a), FromIntegral Int a, Lift Exp a) => Arr a -> Acc (Scalar a)
roughness xs = compute $ runFold standardDeviationF $ diffs xs

-- smoothSimple :: (Elt a, Floating a, FromIntegral Int a, Ord a) => Arr a -> Exp Int -> Arr a
-- smoothSimple xs resolution =
--   ifThenElse (resolution < length xs)
--     smaLargeResolution
--     (sma bestWindowSize 1 xs)
--   where
--     -- THIS IS WRONG. IT NEEDS TO MINIMIZE THE ROUGHNESS
--     bestWindowSize =  the $ minimum $ map fst $ afst $ filter acceptableWindow $ map (go xs) ws
--     acceptableWindow b =
--       (r < originalRoughness) && (k >= originalKurt)
--       where
--         r = fst $ snd b
--         k = snd $ snd b
--     smaLargeResolution = sma (length xs `div` resolution) (length xs `div` resolution) xs
--     originalKurt = kurtosis xs
--     originalRoughness = the $ roughness xs
--     end :: Exp Int
--     end = length xs `div` 10
--     ws :: Arr Int
--     ws = enumFromN (index1 $ end - 2) 2

-- smoothSimple' :: (Elt a, Floating a, FromIntegral Int a, Ord a) => Arr a -> Exp Int -> Arr a
-- smoothSimple' xs resolution = sma bestWindowSize 1 xs
--   where
--     bestWindowSize =
--       fst $ while (\a -> fst a >= 2) godude $ lift ((length xs `div` resolution), (originalRoughness, 1 :: Exp Int))
--     originalKurt = kurtosis xs
--     originalRoughness = the $ roughness xs
--     godude a =
--       let
--         w = fst a
--         minObj = fst $ snd a
--         smoothed = sma w 1 xs
--         r = the $ roughness smoothed
--         kurt = kurtosis smoothed
--       in ifThenElse (r < minObj && kurt >= originalKurt) (lift (w-1, (r, w))) (lift (w-1, snd a))

-- go :: (Elt a, P.Floating (Exp a), FromIntegral Int a) => Arr a -> Exp Int -> Exp (Int, (a, a))
-- go xs w = lift (w,(the $ roughness smoothed, kurtosis smoothed))
--   where
--     smoothed = sma w 1 xs

sma :: (Floating a, FromIntegral Int a) => Exp Int -> Exp Int -> Arr a -> Arr a
sma range slide xs = mavg range xs
  where
    skip ys =
      map snd. afst
        $ filter (\y -> fst y `mod` slide == 0)
          $ imap (\i y -> lift (unindex1 i, y)) ys

computeStuff :: (Floating a, FromIntegral Int a, Lift Exp a) => Exp Int -> Exp Int -> Arr a -> Exp (a, a)
computeStuff resolution window xs = lift (the $ roughness smoothed, kurtosis smoothed)
  where
    smoothed = sma window resolution xs

kurtAndRoughness :: (Floating a, FromIntegral Int a, Lift Exp a) => Arr a -> Exp (a, a)
kurtAndRoughness xs = lift (the $ roughness xs, kurtosis xs)

mavg :: (Floating a, FromIntegral Int a) => Exp Int -> Arr a -> Arr a
mavg k xs =
  map (/ fromIntegral k)
    $ scanl (+) (the $ sum h)
      $ zipWith (-) t xs
  where
    h = take k xs
    t = drop k xs

-- transform :: Arr (Complex Double) -> Arr (Complex Double)
-- transform = A.transform A.forward

-- inverseTransform :: Arr (Complex Double) -> Arr (Complex Double)
-- inverseTransform = A.transform A.inverse

binarySearch :: Int -> Int -> Arr (Complex Double) -> a -> a -> a -> ()
binarySearch h t xs minObj originalKurt windowSize = ()

main :: P.IO ()
main = do
  P.Right xs <- CSV.decode CSV.HasHeader P.<$> BS.readFile "data.csv" :: P.IO (P.Either P.String (V.Vector (P.String, Double)))

  let
    -- use accelerate-io
    v = V.toList $ P.snd P.<$> xs :: [Double]
    vec = use $ fromList (Z :. (V.length xs)) v
    ws = P.reverse [2..V.length xs `div` 10]
    [(minObj, originalKurt)] = toList $ PTX.runN $ unit $ kurtAndRoughness vec
    runOnWindow w = P.head $ toList $ PTX.runN $ unit $ computeStuff 1 (constant w) vec
    roughsAndKurts = (\w -> (w, runOnWindow w)) P.<$> ws
    bestWindow@(_, w) = L.foldl' goman (minObj, 1) roughsAndKurts
    goman old@(obj, _) (w,(r, k)) = if (r P.< obj P.&& k P.>= originalKurt) then (r,w) else old
    smoothed = PTX.run $ sma (constant w) 1 vec
  P.print $ L.take 100 roughsAndKurts
  P.print $ (minObj, originalKurt)
  P.print $ V.length xs
  P.print smoothed
  -- P.print $ S.kurtosis $ (V.fromList v :: V.Vector Double)
    -- bestWindowSize =  the $ minimum $ map fst $ afst $ filter acceptableWindow $ map (go xs) ws
    -- acceptableWindow b =
    --   (r < originalRoughness) && (k >= originalKurt)
    --   where
    --     r = fst $ snd b
    --     k = snd $ snd b
    -- smaLargeResolution = sma (length xs `div` resolution) (length xs `div` resolution) xs
    -- originalKurt = kurtosis xs
    -- originalRoughness = the $ roughness xs
    -- end :: Exp Int
    -- end = length xs `div` 10
    -- ws :: Arr Int
    -- ws = enumFromN (index1 $ end - 2) 2

mavg' :: P.Fractional b => Int -> [b] -> [b]
mavg' k lst = P.map (/ P.fromIntegral k) $ L.scanl' (+) (P.sum h) $ P.zipWith (-) t lst
  where (h, t) = L.splitAt k lst


{-# LANGUAGE Strict, StrictData #-}
{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE DeriveGeneric, DeriveAnyClass #-}
{-# LANGUAGE TypeSynonymInstances #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE StandaloneDeriving #-}
{-# LANGUAGE UndecidableInstances #-}
{-# LANGUAGE TypeFamilies, TypeFamilyDependencies #-}

------------------------------------------------
-- |
-- Module    :  Numeric.Algebra.Homology
-- Copyright :  (c) Jun Yoshida 2019
-- License   :  BSD3
--
-- Chain complexes over integers
--
------------------------------------------------

module Numeric.Algebra.Homology
  (ChainEliminable, Homology(..), homology) where

import GHC.Generics (Generic, Generic1)
import Control.DeepSeq (NFData, force, deepseq)

import Control.Parallel.Strategies

import Control.Monad
import Control.Monad.ST

import Control.Applicative

import Data.STRef

import Data.Bifunctor
import Data.Traversable
import Data.Tuple (swap)
import Data.Maybe
import Data.List

import qualified Numeric.LinearAlgebra as LA
import qualified Numeric.LinearAlgebra.Devel as LA

import qualified Data.Vector as V
import qualified Data.Vector.Mutable as MV

import Numeric.Algebra.FreeModule
import Numeric.Algebra.Presentation
import Numeric.Matrix.Integral
import Numeric.Matrix.Integral.NormalForms
import Numeric.F2 (F2)
import qualified Numeric.Matrix.F2 as F2
import qualified Numeric.Matrix.F2.Mutable as F2

{-- for Debug
import Debug.Trace

traceCond :: Bool -> String -> a -> a
traceCond False _  = id
traceCond True msg = trace msg
--}

--------------
-- * Classes
--------------
class (Normalized a) => ChainEliminable a where
  elimChainHead :: Mat a -> Mat a -> (NormalForm a, Mat a)
  -- ^ If
  -- @ let (a',v) = elimChainHead u a @
  -- Then
  -- @ a' == toNF (a <> u^{-1}) @
  -- and the vectors of v (i.e. toVecs v) induces a basis for the cokernel of a⊗ℚ.

-------------------------------------
-- * Computation of homology groups
-------------------------------------
data Homology a = Homology {
  freeCycs :: [Vec a],
  torsions :: [(a,Vec a)],
  bndries :: [Vec a]
  } deriving (Generic)

deriving instance (Coefficient a, Show a, Show (Vec a)) => Show (Homology a)
deriving instance (Coefficient a) => NFData (Homology a)

-- | Compute the kernel and the image of a matrix from its normal form.
homologyOfNF :: (Normalized a) => NormalForm a -> (Homology a, Homology a)
homologyOfNF nf
  = let invs = invariants nf ++ repeat 0
        cycs = map snd $ filter ((==0) . fst) $ zip invs (toVecs (basisDom nf))
        img = filter ((/=0) . fst) $ zip invs (toVecs (basisCod nf))
        (tors,bnds') = partition ((/=1) . fst) img
        bnds = map (uncurry scale) bnds'
    in (Homology [] tors bnds, Homology cycs [] [])

-- | Decompose complexes into two-term complexes with differential in a Smith normal form.
chainToNormals :: (ChainEliminable a, Traversable t) => t (Mat a) -> (Mat a,t (NormalForm a))
chainToNormals = first (maybe (ident 0) id) . mapAccumL chainToNormals' Nothing
  where
    chainToNormals' mayU mat
      = let (u,mat') = case mayU of
                         Just u -> (u, mat <> u)
                         Nothing -> (ident (rankDom mat), mat)
            !(!matNF, !coker) = force (elimChainHead u mat')
        in (Just coker, matNF)

-- | Compute homology group with integral coefficient.
homology :: (ChainEliminable a, Traversable t, Alternative t, NFData (t (NormalForm a)), NFData (t (Homology a, Homology a))) => t (Mat a) -> t (Homology a)
homology diffs =
  let !(!coker, !nfs) = force (chainToNormals diffs)
      !hs = force (fmap homologyOfNF nfs `using` parTraversable (rparWith rdeepseq))
  in (uncurry ((<|>) . pure)) (mapAccumR pile (Homology (toVecs coker) [] []) hs)
  where
    pile (Homology cycs _ _) (thisH, nextCycH) =
      (nextCycH, thisH {freeCycs = cycs})

-----------------
-- * Instances
-----------------
-- | Enables chain-level Gaussian eliminations on integer matrices.
instance ChainEliminable LA.Z where
  elimChainHead basis matA = runST $ do
    stMatA <- LA.thawMatrix matA
    stBasis <- LA.thawMatrix basis
    stU <- LA.thawMatrix (LA.ident (LA.rows matA))
    -- Compute the Smith normal form
    smithRepST stU stMatA stBasis
    -- Write the new basis of the source
    resBasis <- LA.unsafeFreezeMatrix stBasis
    -- Take the invariant factor.
    invFs <- LA.takeDiag <$!> LA.unsafeFreezeMatrix stMatA
    let rkQ = length (takeWhile (/=0) (LA.toList invFs)) -- rank over the field of fractions; e.g. Q for Z
    -- Free closure of the image of matA & its orthogonal.
    fimage <- if LA.rows matA > 0 && rkQ > 0
              then LA.extractMatrix stU LA.AllRows (LA.ColRange 0 (rkQ-1))
              else return $ (LA.rows matA LA.>< 0) []
    fcoker <- if LA.rows matA > 0 && rkQ < LA.rows matA
              then LA.extractMatrix stU LA.AllRows (LA.FromCol rkQ)
              else return $ (LA.rows matA LA.>< 0) []
    return (SmithNF invFs resBasis fimage, fcoker)

instance ChainEliminable F2 where
  elimChainHead basis matA = runST $ do
    stMatA <- F2.thawMatrixF2 matA
    stBasis <- F2.thawMatrixF2 basis
    stU <- F2.newIdentF2 (F2.rows matA)
    -- Compute the Smith normal form
    rk <- F2.mdiagRepF2 stU stMatA stBasis
    -- Write the new basis of the source
    resBasis <- F2.unsafeFreezeMatrixF2 stBasis
    (fimage,fcoker) <- F2.splitColumnsAt rk <$> F2.unsafeFreezeMatrixF2 stU
    return (DiagNF rk resBasis fimage, fcoker)


{-# LANGUAGE BangPatterns #-}
module Main where

import qualified Data.Vector                   as V
import qualified Data.Vector.Mutable           as VM
import qualified Data.Vector.Unboxed           as VU
import qualified Data.Vector.Unboxed.Mutable   as VUM
import qualified Data.Vector.Generic           as VG
import qualified Data.Vector.Generic.Mutable   as VGM
import qualified Data.ByteString.Char8         as BS
import qualified Data.ByteString.Lazy.Char8    as BSL
import qualified Data.ByteString.Builder       as BSB
import           Data.Char                      ( isSpace )
import           Data.List                      ( unfoldr )
import           Data.Function
import           Data.Complex
import           Control.Monad
import           Control.Monad.Primitive
import           Control.Monad.ST

main :: IO ()
main = do
  n   <- readLn :: IO Int
  abs <-
    VU.replicateM n
    $   (\v -> (v VU.! 0, v VU.! 1))
    .   VU.unfoldr (BS.readInt . BS.dropWhile isSpace)
    <$> BS.getLine
  let v = uncurry multiplyCT $ VU.unzip abs
  print 0
  VU.mapM_ print v

-- O(n^2) multiply polynomials
naiveMultiply :: VU.Vector Int -> VU.Vector Int -> VU.Vector Int
naiveMultiply f g = VU.create $ do
  let len = VU.length f + VU.length g - 1
  h <- VUM.new len
  flip VU.imapM_ f $ \i x ->
    flip VU.imapM_ g $ \j y -> VUM.modify h (\z -> z + x * y) (i + j)
  return h

-- O(n log n) multiply polynomials by FFT （Immutable Vector)
-- 859 ms in atc001-c
-- c.f. https://qiita.com/ageprocpp/items/0d63d4ed80de4a35fe79
multiply :: VU.Vector Int -> VU.Vector Int -> VU.Vector Int
{-# INLINE multiply #-}
multiply !f !g =
  let !len = VU.length f + VU.length g - 1
      !sz  = fix (\rec !k -> if k >= len then k else rec (k * 2)) 1
      !f'  = VU.map fromIntegral f
        VU.++ VU.replicate (sz - VU.length f) (0 :: Complex Double)
      !g' = VU.map fromIntegral g
        VU.++ VU.replicate (sz - VU.length g) (0 :: Complex Double)
  in  VU.map (round . magnitude . (/ fromIntegral sz))
        . VU.take len
        . dft (-1)
        $ VU.zipWith (*) (dft 1 f') (dft 1 g')

-- DFT algorithm
-- if sign == 1 then DFT
-- if sign == -1 then inverse DFT
dft :: Double -> VU.Vector (Complex Double) -> VU.Vector (Complex Double)
{-# INLINE dft #-}
dft !sign !f
  | VU.length f == 1
  = f
  | otherwise
  = let !n    = VU.length f
        !nd2  = n `div` 2
        !f0   = dft sign $ VU.map (f VU.!) $ VU.enumFromStepN 0 2 nd2
        !f1   = dft sign $ VU.map (f VU.!) $ VU.enumFromStepN 1 2 nd2
        -- f0 = dft sign $ VU.ifilter (\i _ -> i `mod` 2 == 0) f
        -- f1 = dft sign $ VU.ifilter (\i _ -> i `mod` 2 == 1) f
        !zeta = cis (sign * 2 * acos (-1) / fromIntegral n) :: Complex Double
    in  VU.imap (\i z -> (f0 VU.! (i `mod` nd2)) + z * (f1 VU.! (i `mod` nd2)))
          $ VU.iterateN n (* zeta) 1

-- O(n log n) multiply polynomials by FFT
-- c.f. https://qiita.com/ageprocpp/items/0d63d4ed80de4a35fe79
-- multiply :: VU.Vector Int -> VU.Vector Int -> VU.Vector Int
-- multiply f g = runST $ do
--   let len = VU.length f + VU.length g - 1
--   let sz = fix (\rec k -> if k >= len then k else rec (k*2)) 1
--   tf <- VU.thaw (VU.map fromIntegral f :: VU.Vector (Complex Double))
--   nf <- VUM.grow tf (sz - VU.length f)
--   tg <- VU.thaw (VU.map fromIntegral g :: VU.Vector (Complex Double))
--   ng <- VUM.grow tg (sz - VU.length g)
--   dft nf 1 sz
--   dft ng 1 sz
--   forM_ [0..sz-1] $ \i -> do
--     y <- VUM.read ng i
--     VUM.modify nf (* y) i
--   dft nf (-1) sz
--   ret <- VU.freeze nf
--   return $ VU.map ((`div` sz) . round . realPart) $ VU.take len ret

-- DFT algorithm
-- if sign == 1 then DFT
-- if sign == -1 then inverse DFT
-- dft :: PrimMonad m => VUM.MVector (PrimState m) (Complex Double) -> Double -> Int -> m ()
-- {-# INLINE dft #-}
-- dft f sign 1 = return ()
-- dft f sign n = do
--   let nd2 = n `div` 2
--   f0 <- VUM.new nd2
--   f1 <- VUM.new nd2
--   forM_ [0..nd2-1] $ \i -> do
--     x0 <- VUM.read f (2*i)
--     VUM.write f0 i x0
--     x1 <- VUM.read f (2*i+1)
--     VUM.write f1 i x1
--   dft f0 sign nd2
--   dft f1 sign nd2
--   let zeta = cis (sign * 2 * acos (-1) / fromIntegral n)
--   let zetas = VU.iterateN n (*zeta) 1
--   flip VU.imapM_ zetas $ \i z -> do
--     x0 <- VUM.read f0 (i `mod` nd2)
--     x1 <- VUM.read f1 (i `mod` nd2)
--     VUM.write f i (x0 + z * x1)

-- O(n log n) multiply polynomials by Cooley-Tukey FFT （Mutable Vector)
-- 448 ms in atc001-c
-- c.f. http://wwwa.pikara.ne.jp/okojisan/stockham/cooley-tukey.html

-- This function computes polynomial coefficient of (f * g)(x)
-- f & g are polynomial coefficient array of f(x) & g(x)
multiplyCT :: VU.Vector Int -> VU.Vector Int -> VU.Vector Int
multiplyCT f g = runST $ do
  let len = VU.length f + VU.length g - 1
  let sz  = fix (\rec k -> if k >= len then k else rec (k * 2)) 1
  tf <- VU.thaw (VU.map fromIntegral f :: VU.Vector (Complex Double))
  nf <- VUM.grow tf (sz - VU.length f)
  tg <- VU.thaw (VU.map fromIntegral g :: VU.Vector (Complex Double))
  ng <- VUM.grow tg (sz - VU.length g)
  fft sz nf
  fft sz ng
  forM_ [0 .. sz - 1] $ \i -> do
    y <- VUM.read ng i
    VUM.modify nf (* y) i
  ifft sz nf
  ret <- VU.freeze nf
  return $ VU.map (round . realPart) $ VU.take len ret
{-# INLINE multiplyCT #-}

fft
  :: PrimMonad m
  => Int -- length
  -> VUM.MVector (PrimState m) (Complex Double) -- target data of FFT
  -> m ()
fft n = cooleyTukey n 1 0 0
{-# INLINE fft #-}

ifft
  :: PrimMonad m
  => Int -- length
  -> VUM.MVector (PrimState m) (Complex Double) -- target data of FFT
  -> m ()
ifft n x = do
  forM_ [0 .. n - 1] $ VUM.modify x conjugate
  cooleyTukey n 1 0 0 x
  forM_ [0 .. n - 1] $ VUM.modify x ((/ fromIntegral n) . conjugate)
{-# INLINE ifft #-}

cooleyTukey
  :: PrimMonad m
  => Int -- length
  -> Int -- stride
  -> Int -- start position of block
  -> Int -- temp var for bit reverse position in start position of block
  -> VUM.MVector (PrimState m) (Complex Double) -- target data of FFT
  -> m ()
cooleyTukey n s q d x
  | n > 1 = do
    let m  = n `div` 2
    let θ0 = 2 * pi / fromIntegral n
    forM_ [0 .. m - 1] $ \p -> do
      let wp = cos (fromIntegral p * θ0) :+ (-sin (fromIntegral p * θ0))
      a <- VUM.read x (q + p + 0)
      b <- VUM.read x (q + p + m)
      VUM.write x (q + p + 0) (a + b)
      VUM.write x (q + p + m) ((a - b) * wp)
    cooleyTukey m (2 * s) (q + 0) (d + 0) x
    cooleyTukey m (2 * s) (q + m) (d + s) x
  | q > d = VUM.swap x q d
  | otherwise = return ()
{-# INLINE cooleyTukey #-}


{-# LANGUAGE DataKinds             #-}
{-# LANGUAGE DeriveDataTypeable    #-}
{-# LANGUAGE FlexibleContexts      #-}
{-# LANGUAGE FlexibleInstances     #-}
{-# LANGUAGE GADTs                 #-}
{-# LANGUAGE LambdaCase            #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE PatternSynonyms       #-}
{-# LANGUAGE RankNTypes            #-}
{-# LANGUAGE ScopedTypeVariables   #-}
{-# LANGUAGE TypeApplications      #-}
{-# LANGUAGE TypeFamilies          #-}
{-# LANGUAGE TypeInType            #-}
{-# LANGUAGE TypeOperators         #-}
{-# LANGUAGE UndecidableInstances  #-}
{-# LANGUAGE ViewPatterns          #-}

module Backprop.Learn.Model.Function (
  -- * Types
    ParamFunc(..)
  , ParamFuncP, pattern PFP, runParamFuncP
  , FixedFunc, pattern FF, runFixedFunc
  , paramMap, learnParam
  , idLearn
  -- ** Combinators
  , dimapPF, lmapPF, rmapPF, compPF, parPF
  -- *** Chain
  , (.-), nilPF, onlyPF
  -- * Activation functions
  -- | See <https://en.wikipedia.org/wiki/Activation_function>
  --
  -- ** Maps
  -- *** Unparameterized
  , step
  , logistic
  , softsign
  , reLU
  , softPlus
  , bentIdentity
  , siLU
  , softExponential
  , sinc
  , gaussian
  , tanh
  , atan
  , sin
  , vmap
  , vmap'
  -- *** Parameterized
  , liftUniform
  , isru
  , preLU
  , sreLU
  , sreLUPFP
  , eLU
  , isrLU
  , apl
  , aplPFP
  -- ** Mixing
  , softMax
  , maxout
  , kSparse
  ) where

import           Backprop.Learn.Model.Class
import           Control.Category
import           Data.Foldable
import           Data.Proxy
import           Data.Type.Length
import           Data.Type.Mayb hiding                        (type (<$>))
import           Data.Type.Tuple hiding                       (T2(..), T3(..))
import           Data.Typeable
import           GHC.TypeNats
import           Numeric.Backprop
import           Numeric.LinearAlgebra.Static.Backprop hiding (tr)
import           Prelude hiding                               ((.), id)
import           Type.Class.Known
import           Type.Family.List
import qualified Data.Type.Tuple                              as T
import qualified Data.Vector.Sized                            as SV
import qualified Data.Vector.Storable.Sized                   as SVS
import qualified Numeric.LinearAlgebra                        as HU
import qualified Numeric.LinearAlgebra.Static                 as H
import qualified Numeric.LinearAlgebra.Static.Vector          as H

-- | An unparameterized function.  Has a 'Category' instance.
--
-- A @'FixedFunc' a b@ essentially the same as a:
--
-- @
-- forall s. 'Reifies' s 'W' => 'BVar' s a -> 'BVar' s b
-- @
--
-- And the 'FF' pattern (and 'runFixedFunc' extractor) witness this.
--
-- It is usually better to just use the functions directly, with
-- combinators like 'Dimap', 'LMap', 'RMap', 'dimapPF', 'lmapPF', 'rmapPF',
-- etc.  This is just provided to let you work nicely with 'ParamFunc'
-- combinators.
type FixedFunc = ParamFunc 'Nothing

-- | 'FF' and 'runFixedFunc' witness the fact that a @'FixedFunc' a b@ is
-- just a function from @'BVar' s a@ to @'BVar' s b@.
pattern FF :: (forall s. Reifies s W => BVar s a -> BVar s b) -> FixedFunc a b
pattern FF { runFixedFunc } <- (getFF->runFixedFunc) where
    FF f = PF (const f)
{-# COMPLETE FF #-}

getFF :: forall a b. FixedFunc a b -> (forall s. Reifies s W => BVar s a -> BVar s b)
getFF ff = runParamFunc ff N_

-- | Identity model, useful for using with other combinators.
idLearn :: FixedFunc a a
idLearn = FF id

-- | A @'ParamFunc' p a b@ is a parameterized function from @a@ to @b@,
-- potentially with trainable parameter @p@.
--
-- A utility wrapper for a deterministic and stateless model.
newtype ParamFunc p a b =
    PF { runParamFunc :: forall s. Reifies s W => Mayb (BVar s) p -> BVar s a -> BVar s b
       }
  deriving (Typeable)

instance ( Num a, Num b, MaybeC Num p, KnownMayb p)
      => Learn a b (ParamFunc p a b) where
    type LParamMaybe (ParamFunc p a b) = p

    runLearn p = stateless . runParamFunc p

instance Category (ParamFunc p) where
    id = PF (const id)
    f . g = PF $ \p -> runParamFunc f p
                     . runParamFunc g p

-- | Convenient type synonym for a 'ParamFunc' with parameters.
--
-- Mostly made to be easy to construct/deconstruct with 'PFP' and
-- 'runParamFuncP'.
type ParamFuncP p = ParamFunc ('Just p)

pattern PFP :: (forall s. Reifies s W => BVar s p -> BVar s a -> BVar s b)
            -> ParamFuncP p a b
pattern PFP { runParamFuncP } <- (getPFP->(getWF->runParamFuncP))
  where
    PFP f = PF $ \(J_ p) -> f p
{-# COMPLETE PFP #-}

newtype WrapFun p a b = WF { getWF :: forall s. Reifies s W => BVar s p -> BVar s a -> BVar s b }

getPFP :: ParamFuncP p a b -> WrapFun p a b
getPFP pf = WF (\case p -> runParamFunc pf (J_ p))

-- | Create a 'ParamFunc' from any instance of 'Learn' that does not have
-- state.
learnParam
    :: forall l a b. (Learn a b l, NoState l)
    => l
    -> ParamFunc (LParamMaybe l) a b
learnParam l = PF (runLearnStateless l)

-- | 'ParamFuncP' taking a singleton list; meant to be used with '.-'
onlyPF
    :: forall p a b. (KnownMayb p, MaybeC Backprop p)
    => ParamFunc p a b
    -> ParamFuncP (T (MaybeToList p)) a b
onlyPF f = PFP $ \ps -> case knownMayb @p of
                          N_   -> runParamFunc f N_
                          J_ _ -> runParamFunc f (J_ (isoVar tOnly onlyT ps))

-- | Compose two 'ParamFuncP's on lists.
(.-)
    :: forall ps qs a b c. (ListC (Backprop <$> ps), ListC (Backprop <$> qs), Known Length ps)
    => ParamFuncP (T ps) b c
    -> ParamFuncP (T qs) a b
    -> ParamFuncP (T (ps ++ qs)) a c
f .- g = PFP $ \ps -> runParamFuncP f (ps ^^. tTake @ps @qs known)
                    . runParamFuncP g (ps ^^. tDrop @ps @qs known)
infixr 9 .-

-- | The identity of '.-'
nilPF :: ParamFuncP (T '[]) a a
nilPF = id

-- | Pre- and post-compose a 'ParamFunc'
dimapPF
    :: (forall s. Reifies s W => BVar s a -> BVar s b)
    -> (forall s. Reifies s W => BVar s c -> BVar s d)
    -> ParamFunc p b c
    -> ParamFunc p a d
dimapPF f g h = PF $ \p -> g . runParamFunc h p . f

-- | Precompose a 'ParamFunc'
lmapPF
    :: (forall s. Reifies s W => BVar s a -> BVar s b)
    -> ParamFunc p b c
    -> ParamFunc p a c
lmapPF f = dimapPF f id

-- | Postcompose a 'ParamFunc'
rmapPF
    :: (forall s. Reifies s W => BVar s b -> BVar s c)
    -> ParamFunc p a b
    -> ParamFunc p a c
rmapPF = dimapPF id

-- | Compose two 'ParamFunc's sequentially
compPF
    :: forall p q a b c. (MaybeC Backprop p, MaybeC Backprop q, KnownMayb p, KnownMayb q)
    => ParamFunc p a b
    -> ParamFunc q b c
    -> ParamFunc (TupMaybe p q) a c
compPF f g = PF $ \pq ->
    let (p, q) = splitTupMaybe @_ @p @q (\(T2B p' q') -> (p', q')) pq
    in  runParamFunc g q . runParamFunc f p

-- | Compose two 'ParamFunc's in parallel
parPF
    :: forall p q a b c d.
     ( MaybeC Backprop p
     , MaybeC Backprop q
     , KnownMayb p
     , KnownMayb q
     , Backprop a
     , Backprop b
     , Backprop c
     , Backprop d
     )
    => ParamFunc p a c
    -> ParamFunc q b d
    -> ParamFunc (TupMaybe p q) (T.T2 a b) (T.T2 c d)
parPF f g = PF $ \pq (T2B x y) ->
    let (p, q) = splitTupMaybe @_ @p @q (\(T2B p' q') -> (p', q')) pq
    in  T2B (runParamFunc f p x)
            (runParamFunc g q y)

-- TODO: replace all of these with manual ops?

-- | Softmax normalizer
softMax
    :: (KnownNat i, Reifies s W)
    => BVar s (R i)
    -> BVar s (R i)
softMax x = expx / konst (norm_1V expx)
  where
    expx = exp x

-- | Logistic function
--
-- \[
-- \sigma(x) = \frac{1}{1 + e^{-x}}
-- \]
logistic :: Floating a => a -> a
logistic x = 1 / (1 + exp (-x))

-- | Binary step / heaviside step
--
-- To use with vectors ('R'), use 'vmap''.
--
-- \[
-- \begin{cases}
-- 0 & \text{for } x < 0 \\
-- 1 & \text{for } x \ge 0
-- \end{cases}
-- \]
step :: (Ord a, Num a) => a -> a
step x | x < 0     = 0
       | otherwise = 1

-- | Softsign activation function
--
-- \[
-- \frac{x}{1 + \lvert x \rvert}
-- \]
softsign :: Fractional a => a -> a
softsign x = x / (1 + abs x)

-- | Inverse square root unit
--
-- \[
-- \frac{x}{\sqrt{1 + \alpha x^2}}
-- \]
--
-- See 'liftUniform' to make this compatible with 'PFP'.
--
-- You can also just use this after partially applying it, to fix the
-- parameter (and not have it trained).
isru
    :: Floating a
    => a        -- ^ α (scaling parameter)
    -> a        -- ^ x
    -> a
isru α x = x / sqrt (1 + α * x * x)

-- | Rectified linear unit.
--
-- To use with vectors ('R'), use 'vmap''.
--
--
-- \[
-- \max(0,x)
-- \]
--
-- @
-- 'reLU' = 'preLU' 0
-- @
--
reLU :: (Num a, Ord a) => a -> a
reLU x | x < 0     = 0
       | otherwise = x

-- | Parametric rectified linear unit
--
-- To use with vectors ('R'), use 'vmap''.
--
-- If scaling parameter is a fixed (and not learned) parameter, this is
-- typically called a leaky recitified linear unit (typically with
-- α = 0.01).
--
-- To use as a learned parameter:
--
-- @
-- 'vmap' . 'preLU' :: 'BVar' s Double -> 'BVar' s ('R' n) -> BVar s (R n)
-- @
--
-- This can be give directly to 'PFP'.
--
-- To fix the paramater ("leaky"), just partially apply a parameter:
--
-- @
-- 'preLU' 0.01           :: 'BVar' s ('R' n) -> BVar s (R n)
-- preLU ('realToFrac' α) :: BVar s (R n) -> BVar s (R n)
-- @
--
-- See also 'rreLU'.
--
-- \[
-- \begin{cases}
-- \alpha x & \text{for } x < 0 \\
-- x        & \text{for } x \ge 0
-- \end{cases}
-- \]
preLU
    :: (Num a, Ord a)
    => a        -- ^ α (scaling parameter)
    -> a        -- ^ x
    -> a
preLU α x | x < 0     = α * x
          | otherwise = x

-- | Exponential linear unit
--
-- To use with vectors ('R'), use 'vmap''.
--
-- To use as a learned parameter:
--
-- @
-- 'vmap' . 'eLU' :: 'BVar' s Double -> 'BVar' s ('R' n) -> BVar s (R n)
-- @
--
-- This can be give directly to 'PFP'.
--
-- To fix the paramater, just partially apply a parameter:
--
-- @
-- 'vmap'' ('eLU' 0.01) :: 'BVar' s ('R' n) -> BVar s (R n)
-- @
--
-- \[
-- \begin{cases}
-- \alpha (e^x - 1) & \text{for } x < 0 \\
-- x                & \text{for } x \ge 0
-- \end{cases}
-- \]
eLU :: (Floating a, Ord a)
    => a        -- ^ α (scaling parameter)
    -> a        -- ^ x
    -> a
eLU α x | x < 0     = α * (exp x - 1)
        | otherwise = x

-- | S-shaped rectified linear activiation unit
--
-- See 'sreLUPFP' for an uncurried and uniformly lifted version usable with
-- 'PFP'.
--
-- \[
-- \begin{cases}
-- t_l + a_l (x - t_l) & \text{for } x \le t_l  \\
-- x                   & \text{for } t_l < x < t_r  \\
-- t_r + a_r (x - t_r) & \text{for } x \ge t_r
-- \end{cases}
-- \]
sreLU
    :: (Num a, Ord a)
    => a        -- ^ @t_l@
    -> a        -- ^ @a_l@
    -> a        -- ^ @t_r@
    -> a        -- ^ @a_l@
    -> a        -- ^ @x@
    -> a
sreLU tl al tr ar x
    | x < tl    = tl + al * (x - tl)
    | x > tr    = tr + ar * (x - tr)
    | otherwise = x

-- | An uncurried and uniformly lifted version of 'sreLU' directly usable
-- with 'PFP'.
sreLUPFP
    :: (KnownNat n, Reifies s W)
    => BVar s (T.T2 (T.T2 Double Double) (T.T2 Double Double))
    -> BVar s (R n)
    -> BVar s (R n)
sreLUPFP (T2B (T2B tl al) (T2B tr ar)) = vmap (sreLU tl al tr ar)

-- | Inverse square root linear unit
--
-- To use with vectors ('R'), use 'vmap''.
--
-- \[
-- \begin{cases}
-- \frac{x}{1 + \alpha x^2} & \text{for } x < 0 \\
-- x                        & \text{for } x \ge 0
-- \end{cases}
-- \]
isrLU
    :: (Floating a, Ord a)
    => a        -- ^ α
    -> a        -- ^ x
    -> a
isrLU α x
    | x < 0     = x / sqrt (1 + α * x * x)
    | otherwise = x

-- | Adaptive piecewise linear activation unit
--
-- See 'aplPFP' for an uncurried version usable with 'PFP'.
--
-- \[
-- \max(0, x_i) + \sum_j^M a_i^j \max(0, -x_i + b_i^j)
-- \]
apl :: (KnownNat n, KnownNat m, Reifies s W)
    => BVar s (L n m)     -- ^ a
    -> BVar s (L n m)     -- ^ b
    -> BVar s (R m)       -- ^ x
    -> BVar s (R m)
apl as bs x = vmap' (max 0) x
            + sum (toRows (as * (bs - fromRows (SV.replicate x))))

-- | 'apl' uncurried, to be directly usable with 'PFP'.
aplPFP
    :: (KnownNat n, KnownNat m, Reifies s W)
    => BVar s (T.T2 (L n m) (L n m))
    -> BVar s (R m)
    -> BVar s (R m)
aplPFP (T2B a b) = apl a b

-- | SoftPlus
--
-- \[
-- \ln(1 + e^x)
-- \]
softPlus :: Floating a => a -> a
softPlus x = log (1 + exp x)

-- | Bent identity
--
-- \[
-- \frac{\sqrt{x^2 + 1} - 1}{2} + x
-- \]
bentIdentity :: Floating a => a -> a
bentIdentity x = (sqrt (x * x + 1) - 1) / 2 + x

-- | Sigmoid-weighted linear unit.  Multiply 'logistic' by its input.
--
-- \[
-- x \sigma(x)
-- \]
siLU :: Floating a => a -> a
siLU x = x * logistic x

-- | SoftExponential
--
-- To use with vectors ('R'), use 'vmap''.
--
-- \[
-- \begin{cases}
-- - \frac{\ln(1 - \alpha (x + \alpha))}{\alpha} & \text{for } \alpha < 0  \\
-- x                                             & \text{for } \alpha = 0  \\
-- \frac{e^{\alpha x} - 1}{\alpha} + \alpha      & \text{for } \alpha > 0
-- \end{cases}
-- \]
softExponential
    :: (Floating a, Ord a)
    => a            -- ^ α
    -> a            -- ^ x
    -> a
softExponential α x = case compare α x of
    LT -> - log (1 - α * (x + α)) / α
    EQ -> x
    GT -> (exp (α * x) - 1) / α + α

-- | Sinc
--
-- \[
-- \begin{cases}
-- 1                 & \text{for } x = 0  \\
-- \frac{\sin(x)}{x} & \text{for } x \ne 0
-- \end{cases}
-- \]
sinc :: (Floating a, Eq a) => a -> a
sinc 0 = 1
sinc x = sin x / x

-- | Gaussian
--
-- \[
-- e^{-x^2}
-- \]
gaussian :: Floating a => a -> a
gaussian x = exp (- (x * x))

-- | Maximum of vector.
--
-- Compare to 'norm_InfV', which gives the maximum absolute value.
maxout :: (KnownNat n, Reifies s W) => BVar s (R n) -> BVar s Double
maxout = liftOp1 . op1 $ \x ->
    let n = HU.maxElement . H.extract $ x
    in  ( n
        , \d -> H.vecR . SVS.map (\e -> if e == n then d else 0) . H.rVec $ x
        )

-- | Usable with functions like '*', 'isru', etc. to turn them into a form
-- usable with 'PFP':
--
-- @
-- 'liftUniform' ('*')  :: 'BVar' s 'Double' -> BVar s ('R' n) -> BVar s (R n)
-- liftUniform 'isru' :: BVar s Double -> BVar s (R n) -> BVar s (R n)
-- @
--
-- Basically turns a parmaeterized function on individual elements of
-- into one that shares the same parameter across all elements of the
-- vector.
liftUniform
    :: (Reifies s W, KnownNat n)
    => (BVar s (R n) -> r)
    -> BVar s Double
    -> r
liftUniform f = f . konst

-- | Utility function to make a 'ParamFunc' that maps a parameterized
-- function over every item in the 'R'.  The parameter is shared across the
-- entire map, and trained.
paramMap
    :: KnownNat i
    => (forall s. Reifies s W => Mayb (BVar s) p -> BVar s Double -> BVar s Double)
    -> ParamFunc p (R i) (R i)
paramMap f = PF (vmap . f)

-- -- TODO: vmap can do better.

-- | Keep only the top @k@ values, and zero out all of the rest.
--
-- Useful for postcomposing in between layers (with a logistic function
-- before) to encourage the number of "activated" neurons is kept to be
-- around @k@.  Used in k-Sprase autoencoders (see 'KAutoencoder').
--
-- <http://www.ericlwilkinson.com/blog/2014/11/19/deep-learning-sparse-autoencoders>
kSparse
    :: forall n s. (Reifies s W, KnownNat n)
    => Int                      -- ^ number of items to keep
    -> BVar s (R n)
    -> BVar s (R n)
kSparse k = liftOp1 . op1 $ \xs ->
    let xsSort = HU.sortVector (H.extract xs)
        thresh = xsSort `HU.atIndex` (n - k)
        mask   = H.dvmap (\x -> if x >= thresh then 1 else 0) xs
    in  ( H.dvmap (\x -> if x >= thresh then x else 0) xs
        , (* mask)
        )
  where
    n = fromIntegral $ natVal (Proxy @n)


{-# LANGUAGE FlexibleContexts #-}
module FourierPinwheel.Filtering where

import           Data.Array.Repa            as R
import           Data.Complex
import           Data.List                  as L
import           Data.Vector.Generic        as VG
import           Data.Vector.Storable       as VS
import           Data.Vector.Unboxed        as VU
import           DFT.Plan
import           Filter.Utils
import           Numeric.GSL.Special.Bessel
import           Utils.List

{-# INLINE idealLowPassFilter #-}
idealLowPassFilter ::
     (VG.Vector vector (Complex Double))
  => Double
  -> Double
  -> Int
  -> Int
  -> vector (Complex Double)
idealLowPassFilter radius period numRhoFreqs numRFreqs =
  let centerRhoFreq = div numRhoFreqs 2
      centerRFreq = div numRFreqs 2
      sinc x =
        if x == 0
          then 1
          else sin (pi * x) / (pi * x)
   in VG.convert .
      toUnboxed . computeS . fromFunction (Z :. numRFreqs :. numRhoFreqs) $ \(Z :. i' :. j') ->
        let i = i' - centerRFreq
            j = j' - centerRhoFreq
         in -- if i == (j)
            --   then
          (2 * radius) / period *
                   sinc (fromIntegral j  / period * (2 * radius)) :+
                   0
              -- else 0

{-# INLINE applyFilterRho #-}
applyFilterRho ::
     DFTPlan
  -> VU.Vector (Complex Double)
  -> R.Array U DIM4 (Complex Double)
  -> IO (R.Array U DIM4 (Complex Double))
applyFilterRho plan filterF arr' = do
  let (Z :. numRFreqs :. numThetaFreqs :. numRhoFreqs :. numPhiFreqs) =
        extent arr'
      arr =
        computeUnboxedS $
        R.backpermute
          (Z :. numRFreqs :. numRhoFreqs :. numThetaFreqs :. numPhiFreqs)
          (\(Z :. r :. rho :. theta :. phi) -> (Z :. r :. theta :. rho :. phi))
          arr'
  arrF <-
    fmap (fromUnboxed (extent arr) . VG.convert) .
    dftExecute
      plan
      (DFTPlanID
         DFT1DG
         [numRFreqs, numRhoFreqs, numThetaFreqs, numPhiFreqs]
         [0, 1]) .
    VG.convert . toUnboxed $
    arr
  let convolvedArrF =
        computeUnboxedS .
        R.traverse2
          arrF
          (fromUnboxed (Z :. numRFreqs :. numRhoFreqs) filterF)
          const $ \fArr fFilter idx@(Z :. r :. rho :. theta :. phi) ->
          fArr idx * fFilter (Z :. r :. rho)
  fmap
    (computeS .
     R.backpermute
       (Z :. numRFreqs :. numThetaFreqs :. numRhoFreqs :. numPhiFreqs)
       (\(Z :. r :. theta :. rho :. phi) -> (Z :. r :. rho :. theta :. phi)) .
     fromUnboxed (extent arr) . VG.convert) .
    dftExecute
      plan
      (DFTPlanID
         IDFT1DG
         [numRFreqs, numRhoFreqs, numThetaFreqs, numPhiFreqs]
         [0, 1]) .
    VG.convert . toUnboxed $
    convolvedArrF

hollowCoefficients ::
     DFTPlan
  -> Double
  -> Double
  -> R.Array U DIM4 (Complex Double)
  -> IO (R.Array U DIM4 (Complex Double))
hollowCoefficients plan radius period coefficients = do
  let (Z :. numRFreqs :. _ :. numRhoFreqs :. _) = extent coefficients
      lpf = idealLowPassFilter (log radius) (log period) numRhoFreqs numRFreqs
      filter =
        VG.convert .
        toUnboxed .
        computeS . makeFilter2D . fromUnboxed (Z :. numRFreqs :. numRhoFreqs) $
        lpf
  filterF <-
    VG.convert <$>
    dftExecute plan (DFTPlanID DFT1DG [numRFreqs, numRhoFreqs] [0, 1]) filter
  -- computeS . R.zipWith (-) coefficients <$>
  applyFilterRho plan filterF coefficients


{-# LANGUAGE DataKinds #-}
{-# LANGUAGE GADTs #-}

module Neural(
             Example,
             initNetwork,
             train,
             predict
             ) where

  import Numeric.LinearAlgebra
  import System.Random
  import Debug.Trace
  import Data.List

  import Tool

  type Example = (Vector Double, Vector Double)

  -- Fully connected neural network described in ADT
  data Network where
    -- Start layer : data layer. consists of just one data vector. (homogeneous)
    StartLayer :: Vector Double -> Network
    -- Hidden layer : consists of three elements
    -- 1. network subset that comes before this layer
    -- 2. weight matrix of size ((curr hidden layer dimension), (prev layer output dimension+1)).
    -- 3. current forward values (homogeneous)
    HiddenLayer :: Network -> Matrix Double -> Vector Double -> Network
    -- End layer : softmax layer. consists of three elements
    -- 1. network subset that comes before this layer
    -- 2. current softmax value
    -- 3. current loss of training (not used)
    EndLayer :: Network -> Vector Double -> Double -> Network
    deriving (Show)

  {- Activations and loss -}

  softmax :: Vector Double -> Vector Double
  softmax v = normalizer `seq` ((\val -> val / normalizer) `cmap` (exp v))
    where normalizer = sumElements (exp v)

  crossEntropy :: Vector Double -> Vector Double -> Double
  crossEntropy t y = -(sumElements $ (t * (log y)))

  activation x = tanh x

  {- Derivation computations -}

  {-
   - jaNodeHidden is a matrix of derivatives
   -
   - dy1/dx1 dy1/dx2 ... dy1/dxr
   -          ...
   - dyk/dx1 dyk/dx2 ... dyk/dxr
   -
   - where yi = activation (sum (weight * xj)), dimension of x is r, dimension of y is k.
   -}
  jaNodeHidden :: Matrix Double -> Vector Double -> Matrix Double
  jaNodeHidden weights y = weights * (repeatCols (cols weights) jaActivation)
    where
      jaActivation = 1 - (y^2)

  {-
   - jaWeightHidden is a matrix of derivatives
   -
   - dy1/dw11 dy1/dw12 ... dy1/dw1r
   -          ...
   - dyk/dwk1 dyk/dwk2 ... dyk/dwkr
   -
   - where yi = activation (sum (wij * xj)), dimension of x is r, dimension of y is k.
   -}
  jaWeightHidden :: Vector Double -> Vector Double -> Matrix Double
  jaWeightHidden x y = jaActivation `outer` jaLinear
    where
      jaActivation = 1 - (y^2)
      jaLinear = x

  -- Derivatives of softmax-crossentropy loss function
  jaLoss t y = y - t

  {- Network handling -}

  -- Initialize network with random weights and zero values
  -- The first argument is dimension of hidden layers.
  --   0th element represents dimension of the last hidden layer.
  --   0th element should be equal to the number of classes
  -- The second argument is dimension of input layer.
  -- length of hiddenDims should be larger than or equal to 1
  initNetwork :: [Int] -> Int -> IO Network
  initNetwork hiddenDims inputDim = do
    inner <- initInnerNetwork hiddenDims inputDim
    return $ EndLayer
      inner
      (zeros $ head hiddenDims)
      0.0
      where
        initInnerNetwork :: [Int] -> Int -> IO Network
        initInnerNetwork hiddenDims inputDim =
          if 1 == length hiddenDims then do
            randWeight <- randmat (-0.5, 0.5) (head hiddenDims) (inputDim+1) -- col+1 for bias
            return $ HiddenLayer
              (StartLayer (homo $ ones inputDim))
              randWeight
              (homo $ zeros (head hiddenDims)) -- append 1 for bias
          else do
            randWeight <- randmat (-0.5, 0.5) (head hiddenDims) (head (tail hiddenDims) + 1) -- col+1 for bias
            prevLayer <- (initInnerNetwork (tail hiddenDims) inputDim)
            return $ HiddenLayer
              prevLayer
              randWeight
              (homo $ zeros (last hiddenDims)) -- append 1 for bias

  -- Get a vector value of the last layer of the network
  lastLayerOf :: Network -> Vector Double
  lastLayerOf network = case network of
    StartLayer d -> d
    HiddenLayer _ _ value -> value
    EndLayer _ value _ -> value

  -- Train FC neural network 'epoch' times using given examples
  train :: Int -> Double -> [Example] -> Network -> Network
  train epoch learningRate examples network = examples `seq` ((iterate (trainEpoch learningRate examples) network)!!epoch)

  -- Train FC neural network with examples
  trainEpoch :: Double -> [Example] -> Network -> Network
  trainEpoch learningRate examples network =
    foldl' (\acc (input, gt) ->
      ((backward learningRate gt) . forward . (feedData input)) acc
      ) network examples

  -- Predict softmax values of network
  predict :: Network -> Vector Double -> Vector Double
  predict network input = (lastLayerOf . forward . (feedData input)) network

  -- Prepare Network with input
  feedData :: Vector Double -> Network -> Network
  feedData input network = input `seq` case network of
    StartLayer d -> StartLayer $ homo input
    HiddenLayer prevLayer weights value -> HiddenLayer (feedData input prevLayer) weights value
    EndLayer prevLayer value loss -> EndLayer (feedData input prevLayer) value loss

  -- Forward data to update values of Network. Leaves weights of layers unchanged.
  -- Use this for prediction
  forward :: Network -> Network
  forward network = case network of
    StartLayer d -> StartLayer d
    HiddenLayer prevLayer weights _ ->
      value `seq` HiddenLayer prevResult weights value
        where prevResult = forward prevLayer
              value = (homo (activation (weights #> lastLayerOf prevResult)))
    EndLayer prevLayer _ loss ->
      lastVector `seq` EndLayer prevResult (softmax $ invhomo lastVector) loss -- remove the last element since it is meaningless
        where prevResult = forward prevLayer
              lastVector = lastLayerOf prevResult

  -- Backpropagate errors to update weights of Network. Leaves values of layers unchanged
  -- Uses stochastic gradient descent
  backward :: Double -> Vector Double -> Network -> Network
  backward learningRate gt (EndLayer prevLayer value _) =
    EndLayer (innerBackward prevLayer lastDiff) value loss
      where
        {-
         - lastDiff is a vector of partial derivatives
         -
         - [dL/dY1 ... dL/dYk]
         -
         - where last softmax layer is k-dimension.
         -}
        lastDiff = jaLoss gt value
        loss = crossEntropy gt value

        innerBackward :: Network -> Vector Double -> Network
        innerBackward network nextDiffs = case network of
          StartLayer d -> StartLayer d
          HiddenLayer prevLayer weights values -> newWeights `seq` HiddenLayer (innerBackward prevLayer currDiff) newWeights values
            where
              {-
               - currDiff is a vector of partial derivatives
               -
               - [dL/dX1 ... dL/dXr]
               -
               - where X is the output of previous layer
               -}
              currDiff = (invhomo . flatten) $ ((jaNodeHidden weights (invhomo values)) * (repeatCols (cols weights) nextDiffs)) ? [0]
              newWeights = weightDiff `seq` (weights - ((\w -> w*learningRate) `cmap` weightDiff))
                where
                  {-
                   - weightDiff is a matrix filled with partial derivatives
                   -
                   -  dL/dW11  dL/dW12 ...  dL/dW1r
                   -               ...
                   -  dL/dWk1  dL/dWk2 ...  dL/dWkr
                   -
                   - where this hidden layer have 'k' nodes and previous layer has 'r' nodes
                   -}
                  weightDiff = pjacobian `seq` (pjacobian * (repeatCols (cols pjacobian) nextDiffs))
                    where
                      {-
                       - pjacobian is a matrix filled with partial derivatives (p is for pseudo)
                       -
                       -  dY1/dW11  dY1/dW12 ...  dY1/dW1r
                       -               ...
                       -  dYk/dWk1  dYk/dWk2 ...  dYk/dWkr
                       -
                       - where this hidden layer have 'k' nodes and previous layer has 'r' nodes
                       -}
                      pjacobian = let x = lastLayerOf prevLayer in x `seq` jaWeightHidden x (invhomo values)


{-# LANGUAGE TemplateHaskell #-}
module NewTTRS.Law where

import Control.Lens
import Data.Array.Unboxed
import Data.List
import Data.Monoid
import Statistics.Distribution (complCumulative, cumulative)
import Statistics.Distribution.Normal (normalDistr)

import NewTTRS.Outcome

data Law = Law { lawRaw :: !(UArray Int Double)
               , lawMean, lawStddev :: !Double }

data LawUpdate = LawUpdate
  { playerLaw , opponentLaw :: Law
  , updateOutcome :: Outcome
  }

-- | Law assigned to unrated players
defaultLaw :: Law
defaultLaw = normalLaw 1800 450

-- | The list of discrete scores characterized by a law
omega :: [Int]
omega = [0,10..3600]

-- | The parameter used to characterize upset probability.
alpha :: Double
alpha = 0.0148540595817432

-- | Generate a normalized law from a list of probabilities
-- which correspond to the elements of 'omega'.
lawFromList :: [Double] -> Law
lawFromList xs = Law (listArray (0,360) normalized) mean (sqrt variance)
  where
  normalized = fmap (/ sum xs) xs
  mean     = sum [ fromIntegral p       * x | (p,x) <- zip omega normalized ]
  variance = sum [ fromIntegral (p * p) * x | (p,x) <- zip omega normalized ]
           - mean * mean

-- | Find the probability that the score is in the range of
-- [i-5,i+5] given a law.
lawAt :: Law -> Int -> Double
lawAt law i = lawRaw law ! (i `div` 10)
{-# INLINE lawAt #-}

-- | Probability of an upset given the difference in two ratings.
upsetProbability :: Int -> Double
upsetProbability d = recip (1 + exp (alpha * fromIntegral d))

-- | Perform a bayesian inference given a player's law, the opponent's law
-- and the outcome of the matches between the two.
lawUpdate ::
  Law {- ^ Player's law to be updated -} ->
  Law {- ^ Opponent's law -} ->
  Outcome {- ^ Player's outcome against opponent -} ->
  Law
lawUpdate lp lq outcome
    = lawFromList
      [ sum [ upsetProbability (q - p) ^ w
            * upsetProbability (p - q) ^ l
            * lawAt lq q
            | q <- omega ]
      * lawAt lp p
      | p <- omega
      ]
  where
  w = view outcomeWins outcome
  l = view outcomeLosses outcome

-- | Perform the inverse computation of 'lawUpdate'
lawUnupdate ::
  Law {- ^ Player's law to be updated -} ->
  Law {- ^ Opponent's law -} ->
  Outcome {- ^ Player's outcome against opponent -} ->
  Law
lawUnupdate lp lq outcome
    = lawFromList
      [ lawAt lp p
      / sum [ upsetProbability (q - p) ^ w
            * upsetProbability (p - q) ^ l
            * lawAt lq q
            | q <- omega ]
      | p <- omega
      ]
  where
  w = view outcomeWins outcome
  l = view outcomeLosses outcome

-- | Chance that a player with the first law will defeat a player with the second law
chanceToWin :: Law -> Law -> Double
chanceToWin lp lq =
  sum [ upsetProbability (q - p)
      * lawAt lp p
      * lawAt lq q
      | q <- omega
      , p <- omega
      ]

-- | Generate a discretized normal law given a mean and standard deviation.
normalLaw ::
   Double {- ^ mean -} ->
   Double {- ^ standard deviation -} ->
   Law
normalLaw mean stddev = lawFromList $ mkNormal 0 3600 mean stddev

-- | Generate a discrete approximation of a normal distribution.
mkNormal ::
  Int    {- ^ lower bound -} ->
  Int    {- ^ upper bound -} ->
  Double {- ^ mean -} ->
  Double {- ^ standard deviation -} ->
  [Double]
mkNormal lo hi mean stddev =
    [ c (lo + 5)  ]
 ++ [ c (p+10) - c p | p <- [lo+5, lo+15 .. hi - 15]]
 ++ [ c' (hi - 5) ]
  where
  distr = normalDistr mean stddev
  c     = cumulative      distr . fromIntegral
  c'    = complCumulative distr . fromIntegral

-- | Degrade a law given a certain number of days. When days is
-- less than 1, no degrading is done.
timeEffect :: Int -> Law -> Law
timeEffect days law
  | days <= 0 = law
  | otherwise = lawFromList
              $ sum [ lawAt law x * timeAt y | x <- omega, y <- [-3600,-3590.. -x]]
              : [ sum [ lawAt law x * timeAt (r - x) | x <- omega ]
                | r <- omega \\ [0,3600] ]
             ++ [ sum [lawAt law x * timeAt y | x <- omega, y <- [3600-x,3610-x..3600]] ]
  where
  timeArray :: UArray Int Double
  timeArray = listArray (-360,360) $ mkNormal (-3600) 3600 0 (70 * sqrt (fromIntegral days / 365))
  timeAt y = timeArray ! (y `div` 10)

-- | Compute a metric of the law used to order players
lawScore :: Law -> Double
lawScore law = lawMean law - 2 * lawStddev law

-- | Return the raw law discrete approximations
lawElems :: Law -> [Double]
lawElems = elems . lawRaw


{-# LANGUAGE NoImplicitPrelude #-}
module Algebra.Additive (
    -- * Class
    C,
    zero,
    (+), (-),
    negate, subtract,

    -- * Complex functions
    sum, sum1,
    sumNestedAssociative,
    sumNestedCommutative,

    -- * Instance definition helpers
    elementAdd, elementSub, elementNeg,
    (<*>.+), (<*>.-), (<*>.-$),

    -- * Instances for atomic types
    propAssociative,
    propCommutative,
    propIdentity,
    propInverse,
  ) where

import qualified Algebra.Laws as Laws

import Data.Int  (Int,  Int8,  Int16,  Int32,  Int64,  )
import Data.Word (Word, Word8, Word16, Word32, Word64, )

import qualified NumericPrelude.Elementwise as Elem
import Control.Applicative (Applicative(pure, (<*>)), )
import Data.Tuple.HT (fst3, snd3, thd3, )
import qualified Data.List.Match as Match

import qualified Data.Complex as Complex98
import qualified Data.Ratio as Ratio98
import qualified Prelude as P
import Prelude (Integer, Float, Double, fromInteger, )
import NumericPrelude.Base

infixl 6  +, -

{- |
Additive a encapsulates the notion of a commutative group, specified
by the following laws:

@
          a + b === b + a
    (a + b) + c === a + (b + c)
       zero + a === a
   a + negate a === 0
@

Typical examples include integers, dollars, and vectors.

Minimal definition: '+', 'zero', and ('negate' or '(-)')
-}

class C a where
    {-# MINIMAL zero, (+), ((-) | negate) #-}
    -- | zero element of the vector space
    zero     :: a
    -- | add and subtract elements
    (+), (-) :: a -> a -> a
    -- | inverse with respect to '+'
    negate   :: a -> a

    {-# INLINE negate #-}
    negate a = zero - a
    {-# INLINE (-) #-}
    a - b    = a + negate b

{- |
'subtract' is @(-)@ with swapped operand order.
This is the operand order which will be needed in most cases
of partial application.
-}
subtract :: C a => a -> a -> a
subtract = flip (-)

{- |
Sum up all elements of a list.
An empty list yields zero.

This function is inappropriate for number types like Peano.
Maybe we should make 'sum' a method of Additive.
This would also make 'lengthLeft' and 'lengthRight' superfluous.
-}
sum :: (C a) => [a] -> a
sum = foldl (+) zero

{- |
Sum up all elements of a non-empty list.
This avoids including a zero which is useful for types
where no universal zero is available.
-}
sum1 :: (C a) => [a] -> a
sum1 = foldl1 (+)

{- |
Sum the operands in an order,
such that the dependencies are minimized.
Does this have a measurably effect on speed?

Requires associativity.
-}
sumNestedAssociative :: (C a) => [a] -> a
sumNestedAssociative [] = zero
sumNestedAssociative [x] = x
sumNestedAssociative xs = sumNestedAssociative (sum2 xs)

{-
Make sure that the last entries in the list
are equally often part of an addition.
Maybe this can reduce rounding errors.
The list that sum2 computes is a breadth-first-flattened binary tree.

Requires associativity and commutativity.
-}
sumNestedCommutative :: (C a) => [a] -> a
sumNestedCommutative [] = zero
sumNestedCommutative xs@(_:rs) =
   let ys = xs ++ Match.take rs (sum2 ys)
   in  last ys

_sumNestedCommutative :: (C a) => [a] -> a
_sumNestedCommutative [] = zero
_sumNestedCommutative xs@(_:rs) =
   let ys = xs ++ take (length rs) (sum2 ys)
   in  last ys

{-
[a,b,c, a+b,c+(a+b)]
[a,b,c,d, a+b,c+d,(a+b)+(c+d)]
[a,b,c,d,e, a+b,c+d,e+(a+b),(c+d)+e+(a+b)]
[a,b,c,d,e,f, a+b,c+d,e+f,(a+b)+(c+d),(e+f)+((a+b)+(c+d))]
-}

sum2 :: (C a) => [a] -> [a]
sum2 (x:y:rest) = (x+y) : sum2 rest
sum2 xs = xs

{- |
Instead of baking the add operation into the element function,
we could use higher rank types
and pass a generic @uncurry (+)@ to the run function.
We do not do so in order to stay Haskell 98
at least for parts of NumericPrelude.
-}
{-# INLINE elementAdd #-}
elementAdd ::
   (C x) =>
   (v -> x) -> Elem.T (v,v) x
elementAdd f =
   Elem.element (\(x,y) -> f x + f y)

{-# INLINE elementSub #-}
elementSub ::
   (C x) =>
   (v -> x) -> Elem.T (v,v) x
elementSub f =
   Elem.element (\(x,y) -> f x - f y)

{-# INLINE elementNeg #-}
elementNeg ::
   (C x) =>
   (v -> x) -> Elem.T v x
elementNeg f =
   Elem.element (negate . f)

-- like <*>
infixl 4 <*>.+, <*>.-, <*>.-$

{- |
> addPair :: (Additive.C a, Additive.C b) => (a,b) -> (a,b) -> (a,b)
> addPair = Elem.run2 $ Elem.with (,) <*>.+  fst <*>.+  snd
-}
{-# INLINE (<*>.+) #-}
(<*>.+) ::
   (C x) =>
   Elem.T (v,v) (x -> a) -> (v -> x) -> Elem.T (v,v) a
(<*>.+) f acc =
   f <*> elementAdd acc

{-# INLINE (<*>.-) #-}
(<*>.-) ::
   (C x) =>
   Elem.T (v,v) (x -> a) -> (v -> x) -> Elem.T (v,v) a
(<*>.-) f acc =
   f <*> elementSub acc

{-# INLINE (<*>.-$) #-}
(<*>.-$) ::
   (C x) =>
   Elem.T v (x -> a) -> (v -> x) -> Elem.T v a
(<*>.-$) f acc =
   f <*> elementNeg acc

-- * Instances for atomic types

instance C Integer where
   {-# INLINE zero #-}
   {-# INLINE negate #-}
   {-# INLINE (+) #-}
   {-# INLINE (-) #-}
   zero   = P.fromInteger 0
   negate = P.negate
   (+)    = (P.+)
   (-)    = (P.-)

instance C Float   where
   {-# INLINE zero #-}
   {-# INLINE negate #-}
   {-# INLINE (+) #-}
   {-# INLINE (-) #-}
   zero   = P.fromInteger 0
   negate = P.negate
   (+)    = (P.+)
   (-)    = (P.-)

instance C Double  where
   {-# INLINE zero #-}
   {-# INLINE negate #-}
   {-# INLINE (+) #-}
   {-# INLINE (-) #-}
   zero   = P.fromInteger 0
   negate = P.negate
   (+)    = (P.+)
   (-)    = (P.-)

instance C Int     where
   {-# INLINE zero #-}
   {-# INLINE negate #-}
   {-# INLINE (+) #-}
   {-# INLINE (-) #-}
   zero   = P.fromInteger 0
   negate = P.negate
   (+)    = (P.+)
   (-)    = (P.-)

instance C Int8    where
   {-# INLINE zero #-}
   {-# INLINE negate #-}
   {-# INLINE (+) #-}
   {-# INLINE (-) #-}
   zero   = P.fromInteger 0
   negate = P.negate
   (+)    = (P.+)
   (-)    = (P.-)

instance C Int16   where
   {-# INLINE zero #-}
   {-# INLINE negate #-}
   {-# INLINE (+) #-}
   {-# INLINE (-) #-}
   zero   = P.fromInteger 0
   negate = P.negate
   (+)    = (P.+)
   (-)    = (P.-)

instance C Int32   where
   {-# INLINE zero #-}
   {-# INLINE negate #-}
   {-# INLINE (+) #-}
   {-# INLINE (-) #-}
   zero   = P.fromInteger 0
   negate = P.negate
   (+)    = (P.+)
   (-)    = (P.-)

instance C Int64   where
   {-# INLINE zero #-}
   {-# INLINE negate #-}
   {-# INLINE (+) #-}
   {-# INLINE (-) #-}
   zero   = P.fromInteger 0
   negate = P.negate
   (+)    = (P.+)
   (-)    = (P.-)

instance C Word    where
   {-# INLINE zero #-}
   {-# INLINE negate #-}
   {-# INLINE (+) #-}
   {-# INLINE (-) #-}
   zero   = P.fromInteger 0
   negate = P.negate
   (+)    = (P.+)
   (-)    = (P.-)

instance C Word8   where
   {-# INLINE zero #-}
   {-# INLINE negate #-}
   {-# INLINE (+) #-}
   {-# INLINE (-) #-}
   zero   = P.fromInteger 0
   negate = P.negate
   (+)    = (P.+)
   (-)    = (P.-)

instance C Word16  where
   {-# INLINE zero #-}
   {-# INLINE negate #-}
   {-# INLINE (+) #-}
   {-# INLINE (-) #-}
   zero   = P.fromInteger 0
   negate = P.negate
   (+)    = (P.+)
   (-)    = (P.-)

instance C Word32  where
   {-# INLINE zero #-}
   {-# INLINE negate #-}
   {-# INLINE (+) #-}
   {-# INLINE (-) #-}
   zero   = P.fromInteger 0
   negate = P.negate
   (+)    = (P.+)
   (-)    = (P.-)

instance C Word64  where
   {-# INLINE zero #-}
   {-# INLINE negate #-}
   {-# INLINE (+) #-}
   {-# INLINE (-) #-}
   zero   = P.fromInteger 0
   negate = P.negate
   (+)    = (P.+)
   (-)    = (P.-)

-- * Instances for composed types

instance (C v0, C v1) => C (v0, v1) where
   {-# INLINE zero #-}
   {-# INLINE negate #-}
   {-# INLINE (+) #-}
   {-# INLINE (-) #-}
   zero   = (,) zero zero
   (+)    = Elem.run2 $ pure (,) <*>.+  fst <*>.+  snd
   (-)    = Elem.run2 $ pure (,) <*>.-  fst <*>.-  snd
   negate = Elem.run  $ pure (,) <*>.-$ fst <*>.-$ snd

instance (C v0, C v1, C v2) => C (v0, v1, v2) where
   {-# INLINE zero #-}
   {-# INLINE negate #-}
   {-# INLINE (+) #-}
   {-# INLINE (-) #-}
   zero   = (,,) zero zero zero
   (+)    = Elem.run2 $ pure (,,) <*>.+  fst3 <*>.+  snd3 <*>.+  thd3
   (-)    = Elem.run2 $ pure (,,) <*>.-  fst3 <*>.-  snd3 <*>.-  thd3
   negate = Elem.run  $ pure (,,) <*>.-$ fst3 <*>.-$ snd3 <*>.-$ thd3

instance (C v) => C [v] where
   zero   = []
   negate = map negate
   (+) (x:xs) (y:ys) = (+) x y : (+) xs ys
   (+) xs     []     = xs
   (+) []     ys     = ys
   (-) (x:xs) (y:ys) = (-) x y : (-) xs ys
   (-) xs     []     = xs
   (-) []     ys     = negate ys

instance (C v) => C (b -> v) where
   {-# INLINE zero #-}
   {-# INLINE negate #-}
   {-# INLINE (+) #-}
   {-# INLINE (-) #-}
   zero       _ = zero
   (+)    f g x = (+) (f x) (g x)
   (-)    f g x = (-) (f x) (g x)
   negate f   x = negate (f x)

-- * Properties

propAssociative :: (Eq a, C a) => a -> a -> a -> Bool
propCommutative :: (Eq a, C a) => a -> a -> Bool
propIdentity    :: (Eq a, C a) => a -> Bool
propInverse     :: (Eq a, C a) => a -> Bool

propCommutative  =  Laws.commutative (+)
propAssociative  =  Laws.associative (+)
propIdentity     =  Laws.identity (+) zero
propInverse      =  Laws.inverse (+) negate zero

-- legacy

instance (P.Integral a) => C (Ratio98.Ratio a) where
   {-# INLINE zero #-}
   {-# INLINE negate #-}
   {-# INLINE (+) #-}
   {-# INLINE (-) #-}
   zero                =  P.fromInteger 0
   (+)                 =  (P.+)
   (-)                 =  (P.-)
   negate              =  P.negate

instance (P.RealFloat a) => C (Complex98.Complex a) where
   {-# INLINE zero #-}
   {-# INLINE negate #-}
   {-# INLINE (+) #-}
   {-# INLINE (-) #-}
   zero                =  P.fromInteger 0
   (+)                 =  (P.+)
   (-)                 =  (P.-)
   negate              =  P.negate