module LeastSquares
    ( VarName
    , LinExpr(..)
    , LinContr(..)
    , QuadExpr(..)
    , minimize
    ) where

import Data.List (elemIndex, foldl1')
import qualified Data.Map.Strict as M
import Data.Maybe (fromJust)
import qualified Data.Set as S
import Numeric.LinearAlgebra

type VarName = String

type CanonicalMap = M.Map VarName (Matrix R)

infix 4 :==:

data LinExpr = Var VarName Int
             | Coeff R
             | Vec (Vector R)
             | Mat (Matrix R)
             | Neg LinExpr
             | Sum LinExpr LinExpr
             | Prod LinExpr LinExpr
             deriving (Show, Eq)

instance Num LinExpr where
    (+) = Sum
    a - b = Sum a (Neg b)
    (*) = Prod
    negate = Neg
    signum = undefined
    abs = undefined
    fromInteger a = Coeff (fromInteger a)

data LinContr = LinExpr :==: LinExpr

data QuadExpr = SumSquares LinExpr
              | CoeffQuad R
              | ProdQuad QuadExpr QuadExpr
              | SumQuad QuadExpr QuadExpr

instance Num QuadExpr where
    (+) = SumQuad
    (*) = ProdQuad
    (-) = undefined
    negate = undefined
    signum = undefined
    abs = undefined
    fromInteger a = CoeffQuad (fromInteger a)

class CanonExpr e  where
    varSet :: e -> S.Set (VarName, Int)
    canonicalize :: e -> CanonicalMap

-- Canonicalize a generic expression
-- The vector is stored as a column matrix under the variable "".
instance CanonExpr LinExpr where
    varSet (Var v s) = S.singleton (v, s)
    varSet (Coeff _) = S.empty
    varSet (Vec _) = S.empty
    varSet (Mat _) = S.empty
    varSet (Neg e) = varSet e
    varSet (Prod _ e) = varSet e
    varSet (Sum e1 e2) = varSet e1 `S.union` varSet e2
    canonicalize (Var v s) = M.singleton v (ident s)
    canonicalize (Vec v) = M.singleton "" (asColumn v)
    canonicalize (Neg e) = M.map negate $ canonicalize e
    canonicalize (Prod (Coeff c) e) = M.map (scale c) $ canonicalize e
    canonicalize (Prod (Mat m) e) = M.map (m <>) $ canonicalize e
    canonicalize (Sum e1 e2) =
        M.unionWith (+) (canonicalize e1) (canonicalize e2)
    canonicalize _ = error "Expression is not well-formed"

-- Canonicalize an expression in the constraint
instance CanonExpr LinContr where
    varSet (e1 :==: e2) = varSet e1 `S.union` varSet e2
    canonicalize (e1 :==: e2) = canonicalize (e1 - e2)

-- Canonicalize an expression in the objective, which consists of SumSquares.
instance CanonExpr QuadExpr where
    varSet (SumSquares e) = varSet e
    varSet (CoeffQuad _) = S.empty
    varSet (ProdQuad _ e) = varSet e
    varSet (SumQuad e1 e2) = varSet e1 `S.union` varSet e2
    canonicalize (SumSquares e) = canonicalize e
    canonicalize (ProdQuad (CoeffQuad c) e) =
        M.map (scale (sqrt c)) $ canonicalize e
    canonicalize (SumQuad e1 e2) = canonicalize e1 `vstack` canonicalize e2
    canonicalize _ = error "Expression is not well-formed"

-- Combine two canonical maps by stacking one on top of another.
vstack :: CanonicalMap -> CanonicalMap -> CanonicalMap
vstack m1 m2 = M.fromSet combine (M.keysSet m1 `S.union` M.keysSet m2)
  where
    combine k = lookupDefault k m1 === lookupDefault k m2
    lookupDefault k m = M.findWithDefault (konst 0 (firstDim m, 1)) k m
    firstDim = rows . snd . fromJust . M.lookupGE "" :: CanonicalMap -> Int

minimize :: QuadExpr -> [LinContr] -> M.Map VarName (Vector R)
minimize obj [] = unpack vars $ a <\> (-b)
  where
    (a,b) = pack vars . canonicalize $ obj
    vars = varSet obj
minimize obj constraints = unpack vars $ mat <\> (-vec)
  where
    mat = fromBlocks [[2 * tr a <> a, tr c], [c, 0]]
    vec = vjoin [2 * tr a #> b, d]
    (a,b) = pack vars . canonicalize $ obj
    (c,d) = pack vars $ foldl1' vstack cds
    cds = map canonicalize constraints
    vars = varSet obj `S.union` S.unions (map varSet constraints)

pack :: S.Set (VarName, Int) -> CanonicalMap -> (Matrix R, Vector R)
pack vars m = (a, b)
  where
    a = fromBlocks [map getCoefficient (S.toAscList vars)]
    b = flatten $ M.findWithDefault (konst 0 (rows a, 1)) "" m
    getCoefficient (varName,s) = M.findWithDefault (konst 0 (s, s)) varName m

-- The vector given to unpack might be longer than the total size of the
-- variables. It might contain the dual variable values, for example.
unpack :: S.Set (VarName, Int) -> Vector R -> M.Map VarName (Vector R)
unpack vars result = M.fromList $ zip varNames coefficients
  where
    coefficients = map getCoefficient varList
    getCoefficient (varName,s) = subVector (startIndex varName) s result
    startIndex varName = cumsum !! fromJust (elemIndex varName varNames)
    cumsum = scanl (+) 0 varSizes
    (varNames,varSizes) = unzip varList
    varList = S.toAscList vars


{-# LANGUAGE DataKinds                                #-}

module Main where
import Control.Monad.State                           (execState)
import Data.List.Split                               (splitOn)
import           Control.DeepSeq
import           Control.Exception.Base
import           Control.Monad
import           Control.Monad.IO.Class
import           Control.Monad.Trans.Maybe
import           Control.Monad.Trans.State
import           Data.Bitraversable
import           Data.IDX
import           Data.Maybe
import           Data.Traversable
import           Data.Tuple
import qualified Data.Vector                         as V
import qualified Data.Vector.Generic                 as VG
import qualified Data.Vector.Unboxed                 as VU
import           GHC.TypeLits
import qualified Numeric.LinearAlgebra               as HM  -- non-backprop hmatrix
import qualified Numeric.LinearAlgebra.Data          as HMD
import qualified Numeric.LinearAlgebra.Static        as HMS -- hmatrix with type-checked operations :D
import System.IO

import MNIST
import Neural
import Regression                                    (activate, forward, loss, runTrain)
import Titanic                                       (getAge, processFeatures, processTargets)

readCSV :: [Char] -> [[[Char]]]
readCSV file = [splitOn "," x | x <- (lines file)]

accuracy :: [Double] -> [Double] -> Int -> Double
accuracy targets predictions batchSize = (fromIntegral correct) / (fromIntegral batchSize) where
    correct = length $ filter (\idx -> targets !! idx == (activate predictions) !! idx) [0..batchSize-1]

main = do
    -- Logistic Regression
    --
    -- contents <- readFile "data/titanic/train.csv"
    -- let weights = matrix 1 [0.0 | _ <- [0..7]]
    -- let inputs = processFeatures $ tail (readCSV contents)
    -- let targets = processTargets $ tail (readCSV contents)

    -- -- subtract the header row
    -- let batchSize = (length $ toLists inputs) - 1
    -- let finalWeights = (execState (runTrain inputs targets batchSize) weights)
    -- let finalPredictions = forward inputs finalWeights

    -- print $ accuracy (concat $ toLists targets) (concat $ toLists finalPredictions) batchSize
    -- print $ loss targets finalPredictions batchSize

    -- MNIST Multilayer Perceptron
    --
    -- Just train <- loadMNIST "data/mnist/train-images-idx3-ubyte" "data/mnist/train-labels-idx1-ubyte"
    -- Just test  <- loadMNIST "data/mnist/t10k-images-idx3-ubyte"  "data/mnist/t10k-labels-idx1-ubyte"
    weightsInit <- HM.rand 728 300
    inputsInit <- HM.rand 1 728
    biasInit <- HM.rand 1 300
    yInit <- HM.rand 1 10
    let weights = HMS.matrix (concat $ HM.toLists weightsInit) :: HMS.L 728 300
    let inputs = HMS.matrix (concat $ HM.toLists inputsInit) :: HMS.L 1 728
    let bias = HMS.matrix (concat $ HM.toLists biasInit) :: HMS.L 1 300
    let y = HMS.vector (concat $ HM.toLists yInit) :: HMS.R 10
    print $ relu $ linear weights bias inputs


module Scene.Object.SimpleSphere where

import Scene.Base
import Math.Color
import Math.Ray
import Numeric.LinearAlgebra

-- Sphere
sphere :: Double -> Vec3 -> Material -> Scene
sphere r p mat = [sphere' r p mat]

sphere' :: Double -> Vec3 -> Material -> Object
sphere' r p = Object (sphere_intercept r p)

sphere_intercept :: Double -> Vec3 -> Ray -> Maybe (Double, Ray)
sphere_intercept r p ray@(Ray o v) =
    let
        o' = o-p
        a = v <.> v
        b = v <.> o'
        c = o' <.> o' - r**2
        discriminant = b**2 - a*c
    in if discriminant > 0
        then
            let
                x1 = (-b + sqrt(discriminant))/a
                x2 = (-b - sqrt(discriminant))/a
                x = case (x1 > 0, x2 > 0) of
                    (True,True) -> Just (min x1 x2)
                    (False,True) -> Just x2
                    (True, False) -> Just x1
                    (False, False) -> Nothing

            in (\lambda -> (lambda, Ray (parametric ray lambda) (parametric ray lambda - p))) <$> x
        else Nothing


-- |
-- Module    : Statistics.Matrix.Algorithms
-- Copyright : 2014 Bryan O'Sullivan
-- License   : BSD3
--
-- Useful matrix functions.

module Statistics.Matrix.Algorithms
    (
      qr
    ) where

import Control.Applicative
import Control.Monad.ST (ST, runST)
import Statistics.Matrix (Matrix, column, dimension, for, norm)
import qualified Statistics.Matrix.Mutable as M
import Statistics.Sample.Internal (sum)
import qualified Data.Vector.Unboxed as U
import Prelude hiding (sum, replicate)

-- | /O(r*c)/ Compute the QR decomposition of a matrix.
-- The result returned is the matrices (/q/,/r/).
qr :: Matrix -> (Matrix, Matrix)
qr mat = runST $ do
  let (m,n) = dimension mat
  r <- M.replicate n n 0
  a <- M.thaw mat
  for 0 n $ \j -> do
    cn <- M.immutably a $ \aa -> norm (column aa j)
    M.unsafeWrite r j j cn
    for 0 m $ \i -> M.unsafeModify a i j (/ cn)
    for (j+1) n $ \jj -> do
      p <- innerProduct a j jj
      M.unsafeWrite r j jj p
      for 0 m $ \i -> do
        aij <- M.unsafeRead a i j
        M.unsafeModify a i jj $ subtract (p * aij)
  (,) <$> M.unsafeFreeze a <*> M.unsafeFreeze r

innerProduct :: M.MMatrix s -> Int -> Int -> ST s Double
innerProduct mmat j k = M.immutably mmat $ \mat ->
  sum $ U.zipWith (*) (column mat j) (column mat k)


module Dft where

import Data.Complex

nthPartOfCycle :: Floating a => Int -> Int -> Int -> Complex a
nthPartOfCycle n f _N =
  cis (2 * pi * ((fromIntegral (n * f)) / (fromIntegral _N)))

correlateSignalAtFreq :: RealFloat a => [Complex a] -> Int -> [(Complex a, Complex a, Complex a)]
correlateSignalAtFreq signal f =
  let _N = length signal
      wavePart = [nthPartOfCycle i f _N | i <- [0.._N - 1]]
  in
   scanr (\(s, w) (_, _, sum) -> (s, w, sum + s * w)) (0 :+ 0, 0 :+ 0, 0 :+ 0) (reverse $ zip signal wavePart)


import           Advection1D
import           Basis
import           Element
import           Mesh
import           Node
import           Numeric.LinearAlgebra.HMatrix as HMat
import           Projection
import           ShapeFcns

linBasis   = Lagrange 1
linShpFcn  = TensorProduct linBasis 1
linShpFcn2 = TensorProduct linBasis 2
nodes      = [Node 0 [0.0], Node 1 [0.25]]
nodes2d    = [Node 0 [-1.0, -1.0], Node 1 [1.0, -1.0], Node 2 [1.0, 1.0], Node 3 [-1.0, 1.0]]
lineElem   = StructElem nodes 0
quadElem   = StructElem nodes2d 0
someCoords = [0.0, 0.0]
mesh2D     = generateMesh [0.0, 0.0] [1.0, 1.0] [3, 3] StructElem
mesh1D     = generateMesh [0.0] [1.0] [6] StructElem
meshDiscon = genDMesh (0.0::Double) (3.0::Double) 4 StructElem
lineElemMats = elemMatrices lineElem linShpFcn 1
globalMats = assembleGlobalMatrices meshDiscon linShpFcn 1 1.0

someFunction :: [Double] -> Double
someFunction [x] = sin x

nodalInitialConds = map (\e -> projectL2 e linShpFcn 2 someFunction) (getMeshElements meshDiscon) -- Should be [Maybe (Matrix Double)]
nodeCs = map (concatMap nodeCoordinates . getElementNodes) (getMeshElements meshDiscon)
dispElemIC :: [Double] -> Maybe (Matrix Double) -> String
dispElemIC nodeCoords (Just nodeICs) = concat [show (nodeCoords !! i) ++ " " ++ show (atIndex nodeICs (i,0)) ++ "\n" | i <- [0..1]]
dispElemIC _ Nothing = error "Linear solve failed."
icStrings = concat (zipWith dispElemIC nodeCs nodalInitialConds)

main = do
  print $ computeJacobian lineElem linShpFcn [0.0 :: Double]
  print $ show quadElem
  print $ show $ boundaryElements mesh2D !! 2
  print $ map (fst . elemMatrices lineElem linShpFcn) [1,2,3]
  writeMesh meshDiscon "mesh.dat"
  print $ HMat.toDense $ head globalMats
  print $ HMat.toDense $ globalMats !! 1
  print $ HMat.toDense $ last globalMats
  writeFile "sine_ic.dat" icStrings


-- Take a Ten Minute Walk
-- http://www.codewars.com/kata/54da539698b8a2ad76000228/

module Codewars.Kata.TenMinuteWalk where

import Data.Complex (magnitude, Complex( (:+) ))

isValidWalk :: String -> Bool
isValidWalk walk = case splitAt 10 walk of (ds, []) -> (length walk == 10) &&
                                                       ((==0) . magnitude . sum . map (\c -> case c of
                                                           'n' -> 0 :+ 1
                                                           's' -> 0 :+ (-1)
                                                           'w' -> (-1) :+ 0
                                                           'e' -> 1 :+ 0) $ walk)
                                           (_, rest) -> False


module Colors
   (
   -- * Normalization functions
     normLinear
   , normSigmoid
   , normPeriodic
   , normSine
   -- * Colorization functions
   , colorGrey
   , colorHue
   , colorGrad
   )
where

import           Data.Complex
import qualified Graphics.Image                as I

mod' :: RealFloat a => a -> a -> a
mod' x y = (-) x . (*) y . fromIntegral . truncate $ x / y

normLinear :: Int -> Double -> Double
normLinear max = (/ fromIntegral max)

normSigmoid :: Double -> Double -> Double -> Double
normSigmoid center power = (1 /) . (1 +) . (power **) . (/ (center / 4)) . (center -)

normPeriodic :: Double -> Double -> Double
normPeriodic period = (/ period) . flip mod' period

normSine :: Double -> Double -> Double
normSine period = (/ 2) . (1 -) . cos . (* 6.2831853) . (/ period) . (flip mod' period)

colorGrey :: I.Pixel I.RGB Double -> Maybe Double -> I.Pixel I.RGB Double
colorGrey setColor Nothing    = setColor
colorGrey _        (Just val) = I.PixelRGB val val val

colorGrad :: I.Pixel I.RGB Double -> Bool -> [I.Pixel I.RGB Double] -> Maybe Double -> I.Pixel I.RGB Double
colorGrad convergentC _ _      Nothing  = convergentC
colorGrad _           isCircular colorsIn (Just x) = interColors (colors !! section) (colors !! (section + 1))
 where
  colors  = if isCircular then colorsIn ++ [head colorsIn] else colorsIn
  count   = fromIntegral $ (subtract 1) $ length colors
  section = truncate $ x * count
  p       = (* count) $ (x `mod'`) $ 1 / count
  inter i f = (i +) $ (* p) $ (f - i)
  interColors (I.PixelRGB iR iG iB) (I.PixelRGB fR fG fB) =
     I.PixelRGB (inter iR fR) (inter iG fG) (inter iB fB)

colorHue :: I.Pixel I.RGB Double -> Maybe Double -> I.Pixel I.RGB Double
colorHue setColor Nothing = setColor
colorHue _        (Just x) | x' <= 1   = I.PixelRGB x'       0        1
                           | x' <= 2   = I.PixelRGB 1        0        (2 - x')
                           | x' <= 3   = I.PixelRGB 1        (x' - 2) 0
                           | x' <= 4   = I.PixelRGB (4 - x') 1        0
                           | x' <= 5   = I.PixelRGB 0        1        (x' - 4)
                           | otherwise = I.PixelRGB 0        (6 - x') 1
   where x' = 6 * x


module LogHarmonicsNormalization where

import Data.Complex
import Data.List as L

main = do
  let delta = 0.00001 :: Double
      maxScale = 16 :: Double
      freq = 5
      xs =
        L.map
          (\x -> cis $ (-pi) * freq * (log x) /  (log maxScale))
          [1 / maxScale,1 / maxScale + delta .. maxScale]
      y =
        (exp (log maxScale :+ (-pi) * freq) -
         exp (((-1) * log maxScale) :+ pi * freq)) / (1 :+ (-pi) * freq / (log maxScale))
  print . polar $ ((delta :+ 0) * (L.sum xs))
  print . polar $ y


-- | This module implements various functions that return a probabilistic result,
-- defined as unitary operators, and quantum computations.
module QIO.QIORandom where

import Data.Monoid as Monoid
import QIO.QioSyn
import QIO.Qdata
import QIO.Qio
import Data.Complex

-- | The exponentiated Pauli-X rotation
rX :: RR -> Rotation
rX r (x,y) = if x==y then (cos (r/2):+0) else (0:+ (-(sin (r/2))))

-- | The exponentiated Pauli-Y rotation
rY :: RR -> Rotation
rY r (x,y) = if x==y then (cos (r/2):+0) else (s * sin (r/2):+0) where s = if x then 1 else -1

-- | Applies a Hadamard rotation to each qubit in the given list of qubits
hadamards :: [Qbit] -> U
hadamards [] = mempty
hadamards (q:qs) = uhad q `mappend` hadamards qs

-- | returns the highest integer power of 2 that is less than or equal to \x\
pow2 :: Int -> Int
pow2 x = pow2' 0
  where pow2' y | 2^(y+1) > x = 2^y
            | otherwise = pow2' (y+1)

-- | A rotation that, given a qubit in state 0, leaves it in a super-position of
-- 0 and 1, such that the probability of measuring as state 0 is \ps\.
weightedU :: RR -> Qbit -> U
weightedU ps q | sqrt ps <= 1 = rot q (rX (2*(acos (sqrt ps))))
           | otherwise = error ("weightedU: Invalid Probability: " ++ show ps)
-- | A QIO computation that uses the "weightedU" unitary, to return a Bool that
-- has a probablity of \pf\ of being False.
weightedBool :: RR -> QIO Bool
weightedBool pf = do q <- mkQbit False
                     applyU (weightedU pf q)
                     measQ q

-- | removes any leading Falses from a list of booleans
rlf :: [Bool] -> [Bool]
rlf (False:bs) = rlf bs
rlf bs = bs

-- | removes any leading Falses from the (big-endian) bit-wise representation
-- of the given Int.
rlf_l :: Int -> [Bool]
rlf_l x = rlf (reverse (int2bits x))

-- | returns the number of bits left after calling the "flf_l" function
rlf_n :: Int -> Int
rlf_n x = length (rlf_l x)

-- | Given an Int \max\ that is the largest number required to be represented in
-- a quantum register, this function trims the front off the given register, to
-- leave the number of qubits required to represent \max\.
trim :: Int -> [Qbit] -> [Qbit]
trim max qbs = drop ((length qbs)-(rlf_n max)) qbs

-- | Given an Int \max\, and a quantum register in the state \max\, this function
-- defines a unitary operation that will leave the quantum register in state that
-- has equal probability of being measured in any of the states 0 to \max\.
randomU :: Int -> [Qbit] -> U
randomU max qbs = randomU' max (trim max qbs)
 where
  randomU' _ [] = mempty
  randomU' 0 _ = mempty
  randomU' max (q:qbs) = weightedU (fromIntegral ((max+1)-p)/fromIntegral (max+1)) q
                 `mappend`
                 condQ q (\x -> if x then (randomU (max-p) qbs)
                                else (hadamards qbs))
                  where p = pow2 max

-- | A quantum computation that will return a quantum integer in a state that
-- has equal probabilities of being measured in any of the state 0 to \max\.
randomQInt :: Int -> QIO QInt
randomQInt max = do
  qbs <- mkQ (reverse (int2bits max))
  applyU (randomU max qbs)
  return (QInt (reverse qbs))

-- | A quantum computation that will return a quantum integer in a state that
-- has equal probabilities of being measured in any of the state \min\ to \max\.
randomQIO :: (Int,Int) -> QIO Int
randomQIO (min,max) = do q <- randomInt (max-min)
                         return (q + min)

-- | A quantum computation that measures the outcome of "randomQInt"
randomInt :: Int -> QIO Int
randomInt max = do
  q <- randomQInt max
  measQ q

-- | A quantum computation that returns an integer that is equally likely to be
-- any number in the range 0 to \x\-1
random :: Int -> QIO Int
random x = randomInt (x-1)

-- | This function uses a Quantum computation to simulate the roll of a dice
dice :: IO Int
dice = do
  x <- run (randomInt 5)
  return (x+1)

-- | This function simulates the given number of repitions of dice rolls
dice_rolls :: Int -> IO [Int]
dice_rolls 0 = return []
dice_rolls y = do
  x <- dice
  xs <- dice_rolls (y-1)
  return (x:xs)

-- | Returns the number of occurences of 1 through 6 in the given list of Ints
occs :: [Int] -> (Int,Int,Int,Int,Int,Int)
occs rs = (rs' 1,rs' 2,rs' 3,rs' 4,rs' 5,rs' 6)
  where
-- "unterminated quasiquotation" -> space needed
   rs' x = length [ y | y <-rs, y == x]

-- | Returns the number of occurences of 1 through 6 in the given number of
-- rolls of the dice.
probs' :: Int -> IO (Int,Int,Int,Int,Int,Int)
probs' x = do
  xs <- dice_rolls x
  return (occs xs)

-- | Returns the percentage of occurences of 1 through 6, after the given number
-- of rolls of the dice.
probs :: Int -> IO (RR,RR,RR,RR,RR,RR)
probs x = do
  (a,b,c,d,e,f) <- probs' x
  return (fromIntegral a/x',fromIntegral b/x',fromIntegral c/x',fromIntegral d/x',fromIntegral e/x',fromIntegral f/x')
    where x' = fromIntegral x


{-# OPTIONS_GHC -Wall #-}

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

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

import           Control.Lens

import           Test.Hspec

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

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

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

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

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

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

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

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

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

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

main :: IO ()
main = do

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

  hspec $ describe "Compare results" $ do
    it "for SDIRK_5_3_4' and TRBDF2_3_3_2'" $ maxDiffA < 1.0e-6
    it "for SDIRK_5_3_4' and BDF" $ maxDiffB < 1.0e-6
    it "for TRBDF2_3_3_2' and BDF" $ maxDiffC < 1.0e-6
    it "for CV and ARK for the Predator Prey model" $ maxDiffPpA < 1.0e-3


{-# LANGUAGE Strict #-}
module STC.Bias where

import           Array.UnboxedArray       as AU
import           Data.Array.Repa          as R
import           Data.Complex
import           Data.List                as L
import           Data.Vector.Storable     as VS
import           Data.Vector.Unboxed      as VU
import Data.Vector.Generic as VG
import           DFT.Plan
import           Filter.Utils
import           Pinwheel.FourierSeries2D
import           STC.Point
import           STC.Utils
import           Utils.List
import FourierPinwheel.AsteriskGaussian
import Filter.Utils
import           Data.Array.Accelerate.LLVM.PTX as A
import qualified Data.Array.Accelerate          as A
import FourierMethod.FourierSeries2D
import Image.IO
import Utils.Array
import System.FilePath
import Utils.Parallel
import FourierPinwheel.GaussianEnvelopePinwheel

{-# INLINE computeBias #-}
computeBias ::
     (Num a, Unbox a, VG.Vector vector a) => Int -> Int -> [Point] -> vector a
computeBias rows cols =
  let (minR, maxR) = computeRange rows
      (minC, maxC) = computeRange cols
   in VG.convert .
      toUnboxedVector .
      AU.accum (+) 0 ((minC, minR), (maxC, maxR)) .
      L.map (\(Point x y theta scale) -> ((round x, round y), 1))

computeBiasPinwheelBasis ::
     DFTPlan -> Int -> Double -> Double -> [Point] -> IO (VS.Vector (Complex Double))
computeBiasPinwheelBasis plan numR2Freq sigma period points =
  let r2Freqs = L.map fromIntegral . getListFromNumber $ numR2Freq
      envelope =
        computeUnboxedS $
        analyticalFourierCoefficients2 numR2Freq 1 0 0 sigma period (period * sqrt 2)
  in dftExecute plan (DFTPlanID DFT1DG [numR2Freq, numR2Freq] [0, 1]) .
     VS.convert .
     toUnboxed .
     computeS .
     makeFilter2D .
     R.zipWith (*) envelope .
     fromListUnboxed (Z :. numR2Freq :. numR2Freq) $
     [ L.foldl'
       (\b (Point x y theta _) ->
          b + (cis $ (freqX * x + freqY * y) * (-2) * pi / period))
       0
       points
     | freqY <- r2Freqs
     , freqX <- r2Freqs
     ]

computeBiasPinwheelBasis1 ::
     DFTPlan -> Int -> Double -> Double -> [Point] -> VU.Vector (Complex Double)
computeBiasPinwheelBasis1 plan numR2Freq sigma period points =
  let r2Freqs = L.map fromIntegral . getListFromNumber $ numR2Freq
      envelope =
        toUnboxed . computeS $
        analyticalFourierCoefficients2 numR2Freq 1 0 0 sigma period (period * sqrt 2)
  in VU.zipWith (*) envelope . VU.fromList $
     [ L.foldl'
       (\b (Point x y theta _) ->
          b + (cis $ (freqX * x + freqY * y) * (-2) * pi / period))
       0
       points
     | freqY <- r2Freqs
     , freqX <- r2Freqs
     ]

computeBiasGaussian ::
     DFTPlan
  -> Int
  -> Double
  -> Double
  -> [Point]
  -> IO (VS.Vector (Complex Double))
computeBiasGaussian plan numR2Freq sigma period points =
  let r2Freqs = L.map fromIntegral . getListFromNumber $ numR2Freq
      freqCenter = div numR2Freq 2
      envelope =
        fromFunction (Z :. numR2Freq :. numR2Freq) $ \(Z :. i :. j) ->
          (exp $
           (fromIntegral $ (i - freqCenter) ^ 2 + (j - freqCenter) ^ 2) *
           (sigma ^ 2) /
           (-2)) *
          (sigma ^ 2) /
          (2 * pi) :+
          0
  in dftExecute plan (DFTPlanID DFT1DG [numR2Freq, numR2Freq] [0, 1]) .
     VS.convert .
     toUnboxed .
     computeS .
     makeFilter2D .
     R.zipWith (*) envelope . fromListUnboxed (Z :. numR2Freq :. numR2Freq) $
     [ L.foldl'
       (\b (Point x y theta _) ->
          b + (cis $ (freqX * x + freqY * y) * (-2) * pi / period))
       0
       points
     | freqY <- r2Freqs
     , freqX <- r2Freqs
     ]

computeBiasGaussian1 ::
     DFTPlan -> Int -> Double -> Double -> [Point] -> VU.Vector (Complex Double)
computeBiasGaussian1 plan numR2Freq sigma period points =
  let r2Freqs = L.map fromIntegral . getListFromNumber $ numR2Freq
      freqCenter = div numR2Freq 2
      envelope =
        toUnboxed . computeS . fromFunction (Z :. numR2Freq :. numR2Freq) $ \(Z :. i :. j) ->
          (exp $
           (fromIntegral $ (i - freqCenter) ^ 2 + (j - freqCenter) ^ 2) *
           (sigma ^ 2) /
           (-2)) *
          (sigma ^ 2) /
          (2 * pi) :+
          0
  in VU.zipWith (*) envelope . VU.fromList $
     [ L.foldl'
       (\b (Point x y theta _) ->
          b + (cis $ (freqX * x + freqY * y) * (-2) * pi / period))
       0
       points
     | freqY <- r2Freqs
     , freqX <- r2Freqs
     ]

computeBiasFourierPinwheel ::
     DFTPlan
  -> Int
  -> Int
  -> Double
  -> Double
  -> Double
  -> Double
  -> Double
  -> [Point]
  -> IO (VS.Vector (Complex Double), VS.Vector (Complex Double))
computeBiasFourierPinwheel plan numR2Freqs thetaFreq alpha periodR2 periodEnv radius stdTheta points = do
  asteriskGaussianVec <-
    asteriskGaussianLowPass
      plan
      numR2Freqs
      thetaFreq
      alpha
      periodR2
      periodEnv
      stdTheta
      radius
  let r2Freqs = L.map fromIntegral . getListFromNumber $ numR2Freqs
      numThetaFreq = 2 * thetaFreq + 1
      shiftVec =
        VS.concat . L.replicate numThetaFreq . VS.fromList $
        [ L.foldl'
          (\b (Point x y theta scale) ->
             b + cis (-(freqX * x + freqY * y) * 2 * pi / periodR2))
          0
          points
        | freqY <- r2Freqs
        , freqX <- r2Freqs
        ]
      biasVec = VS.zipWith (*) asteriskGaussianVec shiftVec
      -- centerFreq = div numR2Freqs 2
      -- gaussian2D =
      --   VU.convert .
      --   toUnboxed . computeS . fromFunction (Z :. numR2Freqs :. numR2Freqs) $ \(Z :. i' :. j') ->
      --     let i = i' - centerFreq
      --         j = j' - centerFreq
      --      in exp
      --           (pi * fromIntegral (i ^ 2 + j ^ 2) /
      --            ((-1) * periodR2 ^ 2 * stdR2 ^ 2)) /
      --         (2 * pi * stdR2 ^ 2) :+
      --         0
      -- zeroVec = VS.replicate (numR2Freqs ^ 2) 0
      -- shiftVec =
      --   VS.concat .
      --   L.map
      --     (\af ->
      --        if af == 0
      --          then VS.zipWith (*) gaussian2D . VS.fromList $
      --               [ L.foldl'
      --                 (\b (Point x y theta scale) ->
      --                    b + cis (-(freqX * x + freqY * y) * 2 * pi / periodR2))
      --                 0
      --                 points
      --               | freqY <- r2Freqs
      --               , freqX <- r2Freqs
      --               ]
      --          else zeroVec) $
      --   [-thetaFreq .. thetaFreq]
  biasVecF <-
    dftExecute
      plan
      (DFTPlanID DFT1DG [numThetaFreq, numR2Freqs, numR2Freqs] [0, 1, 2]) .
    VU.convert .
    toUnboxed .
    computeS .
    makeFilter3D .
    fromUnboxed (Z :. numThetaFreq :. numR2Freqs :. numR2Freqs) . VS.convert $
    -- shiftVec
    biasVec
  return (biasVec, biasVecF)
  -- return (shiftVec, biasVecF)

computeBiasFourierPinwheelFull ::
     DFTPlan
  -> Int
  -> Int
  -> Int
  -> Double
  -> Double
  -> Double
  -> Double
  -> Double
  -> Double
  -> Double
  -> (VU.Vector (Complex Double))
  -> [Point]
  -> IO (VS.Vector (Complex Double), VS.Vector (Complex Double))
computeBiasFourierPinwheelFull plan numR2Freqs thetaFreq rFreq alpha periodR2 periodEnv radius stdTheta stdR stdR2 asteriskGaussianVec points = do
  let -- asteriskGaussianVec =
      --   gaussianPinwheel
      --     numR2Freqs
      --     periodR2
      --     stdR2
      --     10
      --     thetaFreq
      --     rFreq
      --     periodEnv
      --     stdTheta
      --     stdR
      r2Freqs = L.map fromIntegral . getListFromNumber $ numR2Freqs
      numThetaFreq = 2 * thetaFreq + 1
      numRFreq = 2 * rFreq + 1
      shiftArr =
        fromListUnboxed (Z :. numR2Freqs :. numR2Freqs) .
        parMap
          rdeepseq
          (\(freqY, freqX) ->
             L.foldl'
               (\b (Point x y theta scale) ->
                  b + cis (-(freqX * x + freqY * y) * 2 * pi / periodR2))
               0
               points) $
        [(freqY, freqX) | freqY <- r2Freqs, freqX <- r2Freqs]
      asteriskGaussianArr =
        fromUnboxed
          (Z :. numRFreq :. numThetaFreq :. numR2Freqs :. numR2Freqs)
          asteriskGaussianVec
  biasArr <-
    computeUnboxedP . R.traverse2 asteriskGaussianArr shiftArr const $ \fA fS idx@(Z :. _ :. _ :. i :. j) ->
      fA idx * fS (Z :. i :. j)
  biasVecF <-
    dftExecute
      plan
      (DFTPlanID
         DFT1DG
         [numRFreq, numThetaFreq, numR2Freqs, numR2Freqs]
         [0, 1, 2, 3]) .
    VU.convert . toUnboxed . computeS . makeFilter4D $
    biasArr
  return
    ( VU.convert . toUnboxed . computeUnboxedS . makeFilter2D $ biasArr
    , biasVecF)

computeBiasFourierPinwheelKoffkaCorss ::
     DFTPlan
  -> Int
  -> Int
  -> Int
  -> Double
  -> Double
  -> Double
  -> Double
  -> Double
  -> Double
  -> Double
  -> (VU.Vector (Complex Double))
  -> [Point]
  -> IO (VS.Vector (Complex Double), VS.Vector (Complex Double))
computeBiasFourierPinwheelKoffkaCorss plan numR2Freqs thetaFreq rFreq alpha periodR2 periodEnv radius stdTheta stdR stdR2 asteriskGaussianVec points
      -- asteriskGaussianVec =
      --   gaussianPinwheel1
      --     numR2Freqs
      --     periodR2
      --     stdR2
      --     10
      --     thetaFreq
      --     rFreq
      --     periodEnv
      --     stdTheta
      --     stdR
 = do
  let r2Freqs = L.map fromIntegral . getListFromNumber $ numR2Freqs
      thetaFreqs = [fromIntegral (-thetaFreq) .. fromIntegral thetaFreq]
      numThetaFreq = 2 * thetaFreq + 1
      numRFreq = 2 * rFreq + 1
      halfPI = 0.5 * pi
      centerFreqR2 = div numR2Freqs 2
      shiftArr =
        fromFunction (Z :. numThetaFreq :. numR2Freqs :. numR2Freqs) $ \(Z :. tf :. yf :. xf) ->
          let freqTheta = fromIntegral (tf - thetaFreq)
              freqY = fromIntegral (yf - centerFreqR2)
              freqX = fromIntegral (xf - centerFreqR2)
           in L.foldl'
                (\b (Point x y theta scale) ->
                   let phi =
                         if abs x < abs y
                           then atan2 0 x
                           else if abs x > abs y
                                  then atan2 y 0
                                  else error
                                         "computeBiasFourierPinwheelKoffkaCorss: x == y"
                    in b +
                       cis (-(freqX * x + freqY * y) * 2 * pi / periodR2) *
                       (cis (-freqTheta * (phi + halfPI)) +
                        cis (-freqTheta * (phi - halfPI))))
                0
                points
      -- shiftArr =
      --   fromListUnboxed (Z :. numThetaFreq :. numR2Freqs :. numR2Freqs) .
      --   parMap
      --     rdeepseq
      --     (\(thetaFreq, freqY, freqX) ->
      --        L.foldl'
      --          (\b (Point x y theta scale) ->
      --             let phi =
      --                   if abs x < abs y
      --                     then atan2 0 x
      --                     else if abs x > abs y
      --                            then atan2 y 0
      --                            else error
      --                                   "computeBiasFourierPinwheelKoffkaCorss: x == y"
      --              in b +
      --                 cis (-(freqX * x + freqY * y) * 2 * pi / periodR2) *
      --                 (cis (-thetaFreq * (phi + halfPI)) +
      --                  cis (-thetaFreq * (phi - halfPI))))
      --          0
      --          points) $
      --   [ (thetaFreq, freqY, freqX)
      --   | thetaFreq <- thetaFreqs
      --   , freqY <- r2Freqs
      --   , freqX <- r2Freqs
      --   ]
      asteriskGaussianArr =
        fromUnboxed
          (Z :. numRFreq :. numThetaFreq :. numR2Freqs :. numR2Freqs)
          asteriskGaussianVec
  biasArr <-
    computeUnboxedP . R.traverse2 asteriskGaussianArr shiftArr const $ \fA fS idx@(Z :. _ :. k :. i :. j) ->
      fA idx * fS (Z :. k :. i :. j)
  biasVecF <-
    (computeUnboxedP . makeFilter4D $ biasArr) >>=
    dftExecute
      plan
      (DFTPlanID
         DFT1DG
         [numRFreq, numThetaFreq, numR2Freqs, numR2Freqs]
         [0, 1, 2, 3]) .
    VU.convert . toUnboxed
  biasVec <- VU.convert . toUnboxed <$> computeUnboxedP (makeFilter2D biasArr)
  return (biasVec, biasVecF)

computeBiasFourierPinwheelDiscrete ::
     DFTPlan
  -> Int
  -> Int
  -> Int
  -> Double
  -> Double
  -> Double
  -> Double
  -> Double
  -> Double
  -> Double
  -> [Point]
  -> IO (VS.Vector (Complex Double), VS.Vector (Complex Double))
computeBiasFourierPinwheelDiscrete plan numR2Freqs thetaFreq rFreq alpha periodR2 periodEnv radius stdTheta stdR stdR2 points = do
  let biasVec =
        VU.convert .
        toUnboxed .
        computeS . makeFilter2D . fromUnboxed (Z :. numR2Freqs :. numR2Freqs) $
        computeBias numR2Freqs numR2Freqs points
      bias = VS.concat . L.replicate (2 * thetaFreq + 1) $ biasVec
      zeroVec = VS.replicate (VS.length biasVec) 0
  biasVecF <-
    dftExecute plan (DFTPlanID DFT1DG [numR2Freqs, numR2Freqs] [0, 1]) biasVec
  let biasF =
        VS.concat .
        L.map
          (\angularFreq ->
             if angularFreq == 0
               then biasVecF
               else zeroVec) $
        [-thetaFreq .. thetaFreq]
  return (bias, biasF)

computeBiasFourierPinwheelDiscrete' ::
     DFTPlan
  -> Int
  -> Int
  -> Int
  -> Double
  -> Double
  -> Double
  -> Double
  -> Double
  -> Double
  -> Double
  -> [Point]
  -> IO (VS.Vector (Complex Double), VS.Vector (Complex Double))
computeBiasFourierPinwheelDiscrete' plan numR2Freqs thetaFreq rFreq alpha periodR2 periodEnv radius stdTheta stdR stdR2 points = do
  let r2Freqs = L.map fromIntegral . getListFromNumber $ numR2Freqs
      shiftArr =
        fromListUnboxed (Z :. numR2Freqs :. numR2Freqs) .
        parMap
          rdeepseq
          (\(freqY, freqX) ->
             L.foldl'
               (\b (Point x y theta scale) ->
                  b + cis (-(freqX * (fromIntegral (round x :: Int)) + freqY * (fromIntegral (round y :: Int))) * 2 * pi / periodR2))
               0
               points) $
        [(freqY, freqX) | freqY <- r2Freqs, freqX <- r2Freqs]
      shiftVec =
        VU.convert . toUnboxed . computeUnboxedS . makeFilter2D $ shiftArr
      centerFreq = div numR2Freqs 2
      envelope =
        fromFunction (Z :. numR2Freqs :. numR2Freqs) $ \(Z :. i :. j) ->
          if i == centerFreq && j == centerFreq
            then 0
            else let r =
                       sqrt . fromIntegral $
                       (i - centerFreq) ^ 2 + (j - centerFreq) ^ 2
                  in (1 / (r * periodR2)) :+ 0
      envelopeShiftVec =
        VU.convert . toUnboxed . computeS . makeFilter2D $
        shiftArr -- *^ (centerHollowArray numR2Freqs envelope)
      zeroVec = VS.replicate (VS.length shiftVec) 0
  biasVec <-
    dftExecute
      plan
      (DFTPlanID IDFT1DG [numR2Freqs, numR2Freqs] [0, 1])
      envelopeShiftVec
  let bias = VS.concat . L.replicate (2 * thetaFreq + 1) $ biasVec
      biasF =
        VS.concat .
        L.map
          (\angularFreq ->
             if angularFreq == 0
               then shiftVec
               else zeroVec) $
        [-thetaFreq .. thetaFreq]
  return (bias, biasF)

-- computeBiasFourierPinwheel ::
--      DFTPlan
--   -> Int
--   -> Int
--   -> Double
--   -> Double
--   -> Double
--   -> Double
--   -> Double
--   -> [Point]
--   -> FilePath
--   -> [PTX]
--   -> Int
--   -> Int
--   -> IO (VS.Vector (Complex Double))
-- computeBiasFourierPinwheel plan numR2Freqs thetaFreq alpha periodR2 periodEnv stdR2 stdTheta points folderPath ptxs numPoints numBatch = do
--   let r2Freqs = L.map fromIntegral . getListFromNumber $ numR2Freqs
--       numThetaFreq = 2 * thetaFreq + 1
--       gaussianVec =
--         asteriskGaussian1 numR2Freqs thetaFreq alpha periodR2 periodEnv stdR2
--       shiftCoef =
--         [ L.foldl'
--           (\b (Point x y theta scale) ->
--              b + cis (-(freqX * x + freqY * y) * 2 * pi / periodR2))
--           0
--           points
--         | freqY <- r2Freqs
--         , freqX <- r2Freqs
--         ]
--       shiftVec = VS.concat . L.replicate numThetaFreq . VS.fromList $ shiftCoef
--       biasVec = VS.zipWith (*) (VS.concat gaussianVec) shiftVec
--       biasMat =
--         A.transpose .
--         A.use .
--         A.fromList (A.Z A.:. numThetaFreq A.:. (numR2Freqs ^ 2)) . VS.toList $
--         biasVec
--   biasR2 <-
--     computeFourierSeriesR2StreamAcc
--       ptxs
--       numR2Freqs
--       numPoints
--       numThetaFreq
--       periodR2
--       1
--       numBatch
--       biasMat
--   plotImageRepa (folderPath </> "AsteriskBias.png") .
--     ImageRepa 8 .
--     fromUnboxed (Z :. (1 :: Int) :. numPoints :. numPoints) .
--     VU.map sqrt . toUnboxed . sumS . R.map (\x -> (magnitude x) ** 2) . rotate3D $
--     biasR2
--   let centerFreq = div numR2Freqs 2
--       shiftArr = fromListUnboxed (Z :. numR2Freqs :. numR2Freqs) shiftCoef
--       gaussian2D =
--         computeUnboxedS . fromFunction (Z :. numR2Freqs :. numR2Freqs) $ \(Z :. i' :. j') ->
--           let i = i' - centerFreq
--               j = j' - centerFreq
--            in (exp
--                  ((pi * (fromIntegral $ i ^ 2 + j ^ 2)) /
--                   ((-1) * periodR2 ^ 2 * stdR2 ^ 2))) /
--               (2 * pi * stdR2 ^ 2) :+
--               0
--       gaussian2DMat =
--         A.use .
--         A.fromList (A.Z A.:. (numR2Freqs ^ 2) A.:. (1 :: Int)) .
--         R.toList . R.zipWith (*) shiftArr $
--         gaussian2D
--   gaussian2DR2 <-
--     computeFourierSeriesR2StreamAcc
--       ptxs
--       numR2Freqs
--       numPoints
--       1
--       periodR2
--       1
--       numBatch
--       gaussian2DMat
--   plotImageRepa (folderPath </> "Gaussian2D.png") .
--     ImageRepa 8 .
--     fromUnboxed (Z :. (1 :: Int) :. numPoints :. numPoints) .
--     VU.map sqrt . toUnboxed . sumS . R.map (\x -> (magnitude x) ** 2) . rotate3D $
--     gaussian2DR2
--   let planID = DFTPlanID DFT1DG [numR2Freqs, numR2Freqs] [0, 1]
--       inversePlanID = DFTPlanID IDFT1DG [numR2Freqs, numR2Freqs] [0, 1]
--   gaussian2DF <-
--     dftExecute plan planID . VU.convert . toUnboxed . computeS . makeFilter2D $
--     gaussian2D
--   asteriskGaussianF <- dftExecuteBatchP plan planID gaussianVec
--   outputVec <-
--     fmap VS.concat .
--     dftExecuteBatchP plan inversePlanID .
--     parMap rdeepseq (VS.zipWith (*) gaussian2DF) $
--     asteriskGaussianF
--   let outputMat =
--         A.transpose .
--         A.use .
--         A.fromList (A.Z A.:. numThetaFreq A.:. (numR2Freqs ^ 2)) . VS.toList $
--         outputVec
--   outputR2 <-
--     computeFourierSeriesR2StreamAcc
--       ptxs
--       numR2Freqs
--       numPoints
--       numThetaFreq
--       periodR2
--       1
--       numBatch
--       outputMat
--   plotImageRepa (folderPath </> "AsteriskGaussian.png") .
--     ImageRepa 8 .
--     fromUnboxed (Z :. (1 :: Int) :. numPoints :. numPoints) .
--     VU.map sqrt . toUnboxed . sumS . R.map (\x -> (magnitude x) ** 2) . rotate3D $
--     outputR2
--   dftExecute
--     plan
--     (DFTPlanID DFT1DG [numThetaFreq, numR2Freqs, numR2Freqs] [0, 1, 2]) .
--     VU.convert .
--     toUnboxed .
--     computeS .
--     makeFilter3D .
--     fromUnboxed (Z :. numThetaFreq :. numR2Freqs :. numR2Freqs) . VS.convert $
--     -- shiftVec
    -- VS.zipWith (*) outputVec shiftVec

-- zeroVec = VS.replicate (numR2Freqs ^ 2) 0
-- shiftVec =
--   VS.concat .
--   L.map
--     (\angularFreq ->
--        if angularFreq == 0
--          then VS.fromList $
--               [ L.foldl'
--                 (\b (Point x y _ _) ->
--                    b + cis (-(freqX * x + freqY * y) * 2 * pi / periodR2))
--                 0
--                 points
--               | freqY <- r2Freqs
--               , freqX <- r2Freqs
--               ]
--          else zeroVec) $
--   [-thetaFreq .. thetaFreq]


module Main where
import           Control.Monad.State
import           Data.List           (sortBy, (\\))
import           Data.Matrix         ((<->), (<|>))
import qualified Data.Matrix         as M
import           Data.Monoid         ((<>))
import           Data.Ord            (comparing)
import           Data.Random         (sampleState, shuffleNofM)
import           Data.Text           (Text, unpack)
import qualified Data.Text.IO        as T (readFile)
import qualified Data.Vector         as V
import           Statistics.Sample   (mean, stdDev)
import           System.Environment  (getArgs)
import           System.IO           (hFlush, stdout)
import           System.Random       (mkStdGen)

{- massage the big string into a good dataset
-}

parse :: String -> M.Matrix Double
parse =
    M.fromLists . map (map read . words) . lines

{- Split the initial matrix into a vector of labels and a matrix of features
-}

splitData :: M.Matrix Double -> (V.Vector Double, M.Matrix Double)
splitData m =
    ( M.getCol 1 m
    , M.submatrix 1 (M.nrows m) 2 (M.ncols m) m
    )

prompt :: IO (String, String) -- The filepath and the algorithm number
prompt = do
    args <- getArgs
    putStrLn "Welcome to Nick Lawler's Feature Selection Algorithm"
    case args of
        [filename, algorithm] -> do
            putStrLn "Filename and Algorithm provided at command line\n"
            return (filename,algorithm)
        _ -> do
            putStr "Type in the name of the file to test : "
            hFlush stdout
            filename <- getLine
            putStrLn ""
            putStrLn "Type the number of the algorithm you want to run."
            putStrLn ""
            putStrLn "        1) Forward Selection"
            putStrLn "        2) Backward Elimination"
            putStrLn "        3) Nick’s Special Algorithm."
            putStr "                                    "
            hFlush stdout
            algorithm <- getLine
            return (filename,algorithm)

{- Z normalize the matrix one column at a time
-}

normalize :: M.Matrix Double -> M.Matrix Double
normalize =
    let normalizeCol col = V.map (\x -> (x - mean col) / stdDev col ) col
    in V.foldl1' (<|>) . V.map (M.colVector . normalizeCol) . toCols

main :: IO ()
main = do
    (filename, algorithm) <- prompt
    raw <- T.readFile filename
    let (labels, rawDataset) = splitData (parse (unpack raw))
    putStrLn $
          "This dataset has "
          ++ show (M.ncols rawDataset)
          ++ " features (not including the class attribute) with "
          ++ show (M.nrows rawDataset)
          ++ " instances.\n"
    putStr "Please wait while I normalize the data... "
    hFlush stdout
    let dataset = M.forceMatrix (normalize rawDataset)
    putStrLn "Done!"
    let algo = case algorithm of
                    "1" -> forwardSelection
                    "2" -> backwardElimination
                    "3" -> specialAlgorithm
                    _ -> error "Choose a proper algorithm!"
    putStrLn $ "Running nearest neighbor with all "
            ++ show (M.ncols dataset)
            ++ " features, using “leaving-one-out” evaluation, I get an accuracy of "
            ++ show (leaveOneOutCV [1.. M.ncols dataset] labels dataset)
            ++ "%\n"
    putStrLn "Beginning Search\n"
    (accuracy,featureSet) <- algo labels dataset
    putStrLn $ "Finished Search!! The best feature subset is "
            ++ show featureSet
            ++ ", which has an accuracy of "
            ++ show accuracy
            ++ "%"

{- The full cross validation pipeline

   After computing the nearest neighbor for each validation run, pair them with
   the true labels, count the number of times you computed the correct one and
   return the proportion of correct classifications
-}

leaveOneOutCV :: [Int]           -- features to consider
              -> V.Vector Double -- labels
              -> M.Matrix Double -- dataset
              -> Double          -- the percentage of computed labels that matched the true lable
leaveOneOutCV features labels dataset =
    let correctlyLabled = V.length
                        $ V.filter (== True)
                        $ V.zipWith (==) labels
                        $ allNearestNeighbors labels (getCols (V.fromList features) dataset)
    in 100 * (fromIntegral correctlyLabled / fromIntegral (M.nrows dataset))

{- Picks out the columns of 'matrix' specified by their 1 based indices in the array cols
-}

getCols :: V.Vector Int -> M.Matrix a -> M.Matrix a
getCols cols matrix =
    if   V.null cols then M.fromLists [[]]
    else V.foldl1' (<|>) $ V.map (M.colVector . (`M.getCol` matrix)) cols

{- Apply nearest neighbor to all (true label, example, training set) permutations
-}

allNearestNeighbors :: V.Vector Double -> M.Matrix Double -> V.Vector Double
allNearestNeighbors labels dataset =
    V.zipWith3 nearestNeighbor (leaveOneOutVec labels) (toRows dataset) (leaveOneOut dataset)

{- Returns a vector containing all possible removals of a single element from
   the argument vector
-}

leaveOneOutVec :: V.Vector a -> V.Vector (V.Vector a)
leaveOneOutVec vec =
    V.map (removeElement vec) (V.enumFromN 0 (V.length vec))

{- Returns a vector equivalent to the argument matrix except that 0-based
   element i has been removed
-}

removeElement :: V.Vector a -> Int -> V.Vector a
removeElement vec i =
    let (top, bottom) = V.splitAt i vec
    in  top <> V.slice 1 (V.length bottom - 1) bottom

{- Returns a vector containing all possible removals of a single row from the
   argument matrix
-}

leaveOneOut :: M.Matrix a -> V.Vector (M.Matrix a)
leaveOneOut matrix =
    V.map (removeRow matrix) (V.enumFromN 1 (M.nrows matrix))

{- Returns a copy of the argument matrix missing the 1-indexed row provided as
   the second argumetn
-}

removeRow :: M.Matrix a -> Int -> M.Matrix a
removeRow matrix rowNum
    | rowNum == 1              = rowsBelow
    | rowNum == M.nrows matrix = rowsAbove
    | otherwise                = rowsAbove <-> rowsBelow
    where rowsAbove = M.submatrix 1 (rowNum - 1) 1 (M.ncols matrix) matrix
          rowsBelow = M.submatrix (rowNum+1) (M.nrows matrix) 1 (M.ncols matrix) matrix

{- The core nearest neighbor algorithm.

   After turning the dataset into a vector of rows, it calculates the euclidian
   distance between each row and the example, pairs that distance with the
   provided label, then extracts the label of the minimum-scoring example

-}
nearestNeighbor :: V.Vector Double -- the labels
                -> V.Vector Double -- the example to be classified
                -> M.Matrix Double -- the dataset
                -> Double          -- the computed label
nearestNeighbor labels example =
      fst
    . V.minimumBy (comparing snd)
    . V.zip labels
    . V.map (euclidianDistance example)
    . toRows

{- Calculate the euclidian distance between two numeric vectors
-}

euclidianDistance :: Num a => V.Vector a -> V.Vector a -> a
euclidianDistance v1 v2 =
      V.sum (V.map (^(2 :: Int)) (V.zipWith (-) v1 v2))

{- Represent a matrix as a vector of row vectors
-}

toRows :: M.Matrix a -> V.Vector (V.Vector a)
toRows matrix =
    V.map (`M.getRow` matrix) (V.enumFromN 1 (M.nrows matrix))

{- Represent a matrix as a vector of Column vectors
-}

toCols :: M.Matrix a -> V.Vector (V.Vector a)
toCols = toRows . M.transpose

{- It turns out that the vast majority of the code for the forward and backward
   searches is essentially the same. We can reuse the main "featureSearch" section
   by abstracting out a function to choose which features are examined in the
   inner loop, another function to determine whether we add or remove a feature
   on each iteration, and the inital set of features.
-}

featureSearch :: ([Int] -> [Int] -> [Int]) -- how to choose which features to consider
              -> ([Int] -> [Int] -> [Int]) -- how to add the next feature
              -> [Int]                     -- starting features
              -> V.Vector Double
              -> M.Matrix Double
              -> IO (Double, [Int])
featureSearch choose add initial labels dataset = do
    (_,_,acc,features) <- execStateT computeFeatures (50,initial,50,initial)
    return (acc,features)
    where
        -- Haskell forces us to be clear about our intention to use 4 mutable variables
        -- even then, its only a functional simulation of mutable state
        computeFeatures :: StateT (Double, [Int], Double, [Int]) IO ()
        computeFeatures =
            forM_ [1.. M.ncols dataset] $ \_ -> do
                (lastBestAcc, lastBestFeatures, bestAcc, bestFeatures) <- get
                let featuresToConsider = [1.. M.ncols dataset] `choose` lastBestFeatures
                accSets <- forM featuresToConsider $ \k -> do
                            let candidateSet = lastBestFeatures `add` [k]
                                (accuracy,set) = case candidateSet of
                                                [] -> (leaveOneOutCV [1..M.ncols dataset] labels dataset,[1..M.ncols dataset])
                                                _  -> (leaveOneOutCV candidateSet labels dataset,candidateSet)
                            putStrLnM ("       Using feature(s) " ++ show set ++ " accuracy is " ++ show accuracy ++ "%")
                            return (accuracy, set)
                let (newAcc, newFeatures) = maximum accSets -- maximum looks at left entry of tuple
                putStrLnM  ""
                when (newAcc < lastBestAcc) $ -- possible corner case when equal
                    putStrLnM "(Warning, Accuracy has decreased! Continuing search in case of local maxima)"
                putStrLnM ("Feature set " ++ show newFeatures ++ " was best, accuracy " ++ show newAcc ++ "%\n")
                put (newAcc, newFeatures, max newAcc bestAcc, if newAcc > bestAcc then newFeatures else bestFeatures)

-- to choose the features considered, we remove the ones we've already looked at
forwardSelection :: V.Vector Double -> M.Matrix Double -> IO (Double, [Int])
forwardSelection = featureSearch (\\) (++) []

-- we're only interested in the features that are left after the last iteration
-- so we use `seq`, which always returns its second argument.
backwardElimination :: V.Vector Double -> M.Matrix Double -> IO (Double,[Int])
backwardElimination labels dataset = featureSearch seq (\\) [1..M.ncols dataset] labels dataset

{- in each inner loop, try out the top 3 scoring feature subsets on the full dataset.
   pick the one with the highest score and maintain that score going forward
-}
specialAlgorithm :: V.Vector Double -> M.Matrix Double -> IO (Double,[Int])
specialAlgorithm labels dataset = do
    putStrLn "Taking a random sample of the original dataset"
    let sampleSize = M.nrows dataset `div` 10
    (_,_,acc,features) <- execStateT  (uncurry computeFeatures (randomSample sampleSize labels dataset)) (50,[],50,[])
    return (acc,features)
    where
        computeFeatures :: V.Vector Double -> M.Matrix Double -> StateT (Double, [Int], Double, [Int]) IO ()
        computeFeatures labels' dataset' =
            forM_ [1.. M.ncols dataset'] $ \_ -> do
                (lastBestAcc, lastBestFeatures, bestAcc, bestFeatures) <- get
                let featuresToConsider = [1.. M.ncols dataset'] \\ lastBestFeatures
                accSets <- forM featuresToConsider $ \k -> do
                            let candidateSet = lastBestFeatures ++ [k]
                                (accuracy,set) = case candidateSet of
                                                [] -> (leaveOneOutCV [1..M.ncols dataset'] labels' dataset',[1..M.ncols dataset'])
                                                _  -> (leaveOneOutCV candidateSet labels' dataset',candidateSet)
                            putStrLnM ("       Using feature(s) " ++ show set ++ " accuracy is " ++ show accuracy ++ "%")
                            return (accuracy, set)
                let accSets' = topNOnFullDataset 3 labels dataset accSets
                    (newAcc, newFeatures) = maximum accSets' -- maximum looks at left entry of tuple
                putStrLnM  ""
                when (newAcc < lastBestAcc) $ -- possible corner case when equal
                    putStrLnM "(Warning, Accuracy has decreased! Continuing search in case of local maxima)"
                putStrLnM ("Feature set " ++ show newFeatures ++ " was best, full dataset accuracy " ++ show newAcc ++ "%\n")
                put (newAcc, newFeatures, max newAcc bestAcc, if newAcc > bestAcc then newFeatures else bestFeatures)

{- Compute the full-dataset validation score of the features subsets that scored
   the best on the reduced-size dataset
-}

topNOnFullDataset :: Int -> V.Vector Double -> M.Matrix Double -> [(Double,[Int])] -> [(Double,[Int])]
topNOnFullDataset n labels dataset =
    let scoreWithFeatures labels' dataset' (_,features) = (leaveOneOutCV features labels' dataset', features)
    in map (scoreWithFeatures labels dataset) . take n . sortBy (flip compare)

{- Return a random n-length subset of the supplied label vector and dataset matrix
-}

randomSample :: Int -> V.Vector Double -> M.Matrix Double -> (V.Vector Double, M.Matrix Double)
randomSample n labels dataset =
    let (labels',dataset') = unzip . fst $ sampleState (shuffleNofM n (M.nrows dataset) (zip (V.toList labels) (M.toLists dataset))) (mkStdGen 366)
    in (V.fromList labels', M.fromLists dataset')

-- quick dataset access in repl
small34 :: IO Text
small34 = T.readFile "data/cs_170_small34.txt"

large34 :: IO Text
large34 = T.readFile "data/cs_170_large34.txt"

small80 :: IO Text
small80 = T.readFile "data/cs_170_small80.txt"

large80 :: IO Text
large80 = T.readFile "data/cs_170_large34.txt"

-- easy printing in IO Monads
putStrLnM :: MonadIO m => String -> m ()
putStrLnM = liftIO . putStrLn


{-# LANGUAGE FlexibleContexts, ConstraintKinds, DataKinds #-}
import GHC.TypeLits
import Numeric.LinearAlgebra.Static
import qualified Numeric.LinearAlgebra as LA
import System.Environment
import Text.Printf (printf)

-- TODO: Learn about statically typed matrices.
-- TODO: Figure out how to avoid converting b/t lists and matrices.
--       (Will basically require slicing and mapping over rows/colmns).
--
--       Solution: extractRows and extractCols after the predictions line.
--       Could get rid of the toColumns entry as well and separately extract.
--       (Not really worth the time.)

-- Create design matrix for polynomial powers.
-- (or, just use Vector Double. But we want to keep generality.)
-- expand :: (Element t, Field t, Num (Vector t)) => Vector t -> Int -> Matrix t
expand :: (KnownNat n, KnownNat d) => R n -> L n d
expand x deg = fromColumns $ map (x^) [0..deg]

regress :: (Element t, Field t, Num (Vector t)) => Vector t -> Vector t -> Int -> Vector t
regress x y deg = (expand x deg) <\> y
--                                  ^ efficient pseudoinverse

evalPoly :: (Element t, Field t, Product t, Num (Vector t)) => Vector t -> Vector t -> Vector t
evalPoly coeff x = expand x (dim coeff - 1) <> coeff

-- sumVector :: (Element t, Field t) => Vector t -> t
squareDiff :: Vector Double -> Vector Double -> Vector Double
squareDiff x y = mapVector (flip (^) 2) (x - y)
sumVector :: Vector Double -> Double
sumVector x    = x <.> ones where ones = konst 1 (dim x)
sumSquareDiff x y = sumVector $ squareDiff x y

-- NB: This requires concrete types due to subroutines called:
-- readFile   :: FilePath -> IO String
-- readMatrix :: String -> Matrix Double
readData :: String -> IO [Vector Double]
readData filename = fmap (toColumns . readMatrix) $ readFile filename

maxDeg = 5
main = do
  args <- getArgs
  if null args
     then putStrLn "Usage: polyreg datafile"
     else do [x,y] <- readData "polyreg_data.txt"
             let models      = map (regress x y) [1..maxDeg]
                 predictions = map (flip evalPoly x) models
                 residuals   = map (squareDiff y) predictions
                 predsview   = fromColumns $ [x, y] ++ predictions
                 residsview  = fromColumns $ [x, y] ++ residuals
                 totalerrs   = map sumVector residuals
                 putFormat   = putStrLn . format "  " (printf "%.2f")

             putStrLn "Modeling predictions"
             putFormat predsview

             putStrLn "Modeling residuals"
             putFormat residsview

             putStrLn "Total errors"
             putFormat $ fromRows [fromList totalerrs]


---------------------------------------------------------------------------------
-- |
-- Utility functions for matrix operations
---------------------------------------------------------------------------------
module Utils where

import           Numeric.LinearAlgebra (Container, Element, Matrix, R, cmap,
                                        sumElements)
import           Prelude               (Num, Ord, (**), (.), (*), (-), abs, flip,
                                        max, signum)

frobeniusNorm :: Matrix R -> R
frobeniusNorm = flip (**) 0.5 . sumElements . cmap ((flip (**) 2) . abs)

absL1Norm :: Matrix R -> R
absL1Norm = sumElements . cmap abs

softThreshold
  :: (Ord e, Num e, Element e, Container c e)
  => c e
  -> e
  -> c e
softThreshold xs mu' = cmap f xs
  where f x = (signum x) * max ((abs x) - mu') 0


{-# OPTIONS_GHC -Wall #-}

-- | Field classes
module Plankton.Field
  ( Semifield
  , Field
  , ExpField(..)
  , QuotientField(..)
  , BoundedField(..)
  , infinity
  , neginfinity
  , TrigField(..)
  ) where

import Data.Complex (Complex(..))
import Plankton.Additive
import Plankton.Multiplicative
import Plankton.Ring
import Protolude (Bool, Double, Float, Integer, (||))
import qualified Protolude as P

-- | A Semifield is a Field without Commutative Multiplication.
class (MultiplicativeInvertible a, Ring a) =>
      Semifield a

instance Semifield Double

instance Semifield Float

instance (Semifield a) => Semifield (Complex a)

-- | A Field is a Ring plus additive invertible and multiplicative invertible operations.
--
-- A summary of the rules inherited from super-classes of Field
--
-- > zero + a == a
-- > a + zero == a
-- > (a + b) + c == a + (b + c)
-- > a + b == b + a
-- > a - a = zero
-- > negate a = zero - a
-- > negate a + a = zero
-- > a + negate a = zero
-- > one * a == a
-- > a * one == a
-- > (a * b) * c == a * (b * c)
-- > a * (b + c) == a * b + a * c
-- > (a + b) * c == a * c + b * c
-- > a * zero == zero
-- > zero * a == zero
-- > a * b == b * a
-- > a / a = one
-- > recip a = one / a
-- > recip a * a = one
-- > a * recip a = one
class (AdditiveGroup a, MultiplicativeGroup a, Ring a) =>
      Field a

instance Field Double

instance Field Float

instance (Field a) => Field (Complex a)

-- | A hyperbolic field class
--
-- > sqrt . (**2) == identity
-- > log . exp == identity
-- > for +ive b, a != 0,1: a ** logBase a b ≈ b
class (Field a) =>
      ExpField a where
  exp :: a -> a
  log :: a -> a
  logBase :: a -> a -> a
  logBase a b = log b / log a
  (**) :: a -> a -> a
  (**) a b = exp (log a * b)
  sqrt :: a -> a
  sqrt a = a ** (one / (one + one))

instance ExpField Double where
  exp = P.exp
  log = P.log
  (**) = (P.**)

instance ExpField Float where
  exp = P.exp
  log = P.log
  (**) = (P.**)

-- | todo: bottom is here somewhere???
instance (TrigField a, ExpField a) => ExpField (Complex a) where
  exp (rx :+ ix) = exp rx * cos ix :+ exp rx * sin ix
  log (rx :+ ix) = log (sqrt (rx * rx + ix * ix)) :+ atan2 ix rx

-- | quotient fields explode constraints if they allow for polymorphic integral types
--
-- > a - one < floor a <= a <= ceiling a < a + one
-- > round a == floor (a + one/(one+one))
class (Field a) =>
      QuotientField a where
  round :: a -> Integer
  ceiling :: a -> Integer
  floor :: a -> Integer
  (^^) :: a -> Integer -> a

instance QuotientField Float where
  round = P.round
  ceiling = P.ceiling
  floor = P.floor
  (^^) = (P.^^)

instance QuotientField Double where
  round = P.round
  ceiling = P.ceiling
  floor = P.floor
  (^^) = (P.^^)

-- | A bounded field includes the concepts of infinity and NaN, thus moving away from error throwing.
--
-- > one / zero + infinity == infinity
-- > infinity + a == infinity
-- > isNaN (infinity - infinity)
-- > isNaN (infinity / infinity)
-- > isNaN (nan + a)
-- > zero / zero != nan
--
-- Note the tricky law that, although nan is assigned to zero/zero, they are never-the-less not equal. A committee decided this.
class (Field a) =>
      BoundedField a where
  maxBound :: a
  maxBound = one / zero
  minBound :: a
  minBound = negate (one / zero)
  nan :: a
  nan = zero / zero
  isNaN :: a -> Bool

-- | prints as `Infinity`
infinity :: BoundedField a => a
infinity = maxBound

-- | prints as `-Infinity`
neginfinity :: BoundedField a => a
neginfinity = minBound

instance BoundedField Float where
  isNaN = P.isNaN

instance BoundedField Double where
  isNaN = P.isNaN

-- | todo: work out boundings for complex
-- as it stands now, complex is different eg
--
-- > one / (zero :: Complex Float) == nan
instance (BoundedField a) => BoundedField (Complex a) where
  isNaN (rx :+ ix) = isNaN rx || isNaN ix

-- | Trigonometric Field
class (P.Ord a, Field a) =>
      TrigField a where
  pi :: a
  sin :: a -> a
  cos :: a -> a
  tan :: a -> a
  tan x = sin x / cos x
  asin :: a -> a
  acos :: a -> a
  atan :: a -> a
  sinh :: a -> a
  cosh :: a -> a
  tanh :: a -> a
  tanh x = sinh x / cosh x
  asinh :: a -> a
  acosh :: a -> a
  atanh :: a -> a
  atan2 :: a -> a -> a
  atan2 y x
    | x P.> zero = atan (y / x)
    | x P.== zero P.&& y P.> zero = pi / (one + one)
    | x P.< one P.&& y P.> one = pi + atan (y / x)
    | (x P.<= zero P.&& y P.< zero) || (x P.< zero) =
      negate (atan2 (negate y) x)
    | y P.== zero = pi -- must be after the previous test on zero y
    | x P.== zero P.&& y P.== zero = y -- must be after the other double zero tests
    | P.otherwise = x + y -- x or y is a NaN, return a NaN (via +)

instance TrigField Double where
  pi = P.pi
  sin = P.sin
  cos = P.cos
  asin = P.asin
  acos = P.acos
  atan = P.atan
  sinh = P.sinh
  cosh = P.cosh
  asinh = P.sinh
  acosh = P.acosh
  atanh = P.atanh

instance TrigField Float where
  pi = P.pi
  sin = P.sin
  cos = P.cos
  asin = P.asin
  acos = P.acos
  atan = P.atan
  sinh = P.sinh
  cosh = P.cosh
  asinh = P.sinh
  acosh = P.acosh
  atanh = P.atanh


-----------------------------------------------------------------------------
-- |
-- Module     : Numeric.LinearAlgebra.Matrix.Statistics
-- Copyright  : Copyright (c) 2010, Patrick Perry <patperry@gmail.com>
-- License    : BSD3
-- Maintainer : Patrick Perry <patperry@gmail.com>
-- Stability  : experimental
--
-- Basic multivariate statistics.
--

module Numeric.LinearAlgebra.Matrix.Statistics (
    defaultCovUplo,

    -- * Immutable interface
    cov,
    covWithMean,
    weightedCov,
    weightedCovWithMean,

    -- * Mutable interface
    covTo,
    covWithMeanTo,
    weightedCovTo,
    weightedCovWithMeanTo,

    ) where

import Control.Monad( when )
import Control.Monad.ST( ST )
import Data.List( foldl' )
import Text.Printf( printf )

import Numeric.LinearAlgebra.Types

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

import Numeric.LinearAlgebra.Matrix.Base( Matrix )
import Numeric.LinearAlgebra.Matrix.STBase( STMatrix )
import qualified Numeric.LinearAlgebra.Matrix.STBase as M
import qualified Numeric.LinearAlgebra.Matrix.Herm as M

-- | Returns the default storage scheme for covariance matrices.
defaultCovUplo :: Uplo
defaultCovUplo = Lower

-- | Returns the sample covariance matrix as a hermitian matrix with storage
-- scheme equal to 'defaultCovUplo'.  The first argument gives the dimension
-- of the vectors.
cov :: (BLAS3 e)
          => Int -> CovMethod -> [Vector e] -> Herm Matrix e
cov p t xs = M.hermCreate $ do
    c <- Herm uplo `fmap` M.new_ (p,p)
    covTo c t xs
    return c
  where
    uplo = defaultCovUplo

-- | Given the pre-computed mean, returns the sample covariance matrix
-- with storage scheme equal to 'defaultCovUplo'.
covWithMean :: (BLAS3 e)
                  => Vector e -> CovMethod -> [Vector e] -> Herm Matrix e
covWithMean mu t xs = M.hermCreate $ do
    p <- V.getDim mu
    c <- Herm uplo `fmap` M.new_ (p,p)
    covWithMeanTo c mu t xs
    return c
  where
    uplo = defaultCovUplo

-- | Returns the weighed sample covariance matrix with storage scheme equal
-- to 'defaultCovUplo'. The first argument gives the dimension of the vectors.
weightedCov :: (BLAS3 e)
                  => Int -> CovMethod -> [(Double, Vector e)] -> Herm Matrix e
weightedCov p t wxs = M.hermCreate $ do
    c <- Herm uplo `fmap` M.new_ (p,p)
    weightedCovTo c t wxs
    return c
  where
    uplo = defaultCovUplo

-- | Given the pre-computed mean, returns the weighed sample covariance matrix
-- with storage scheme equal to 'defaultCovUplo'.
weightedCovWithMean :: (BLAS3 e)
                    => Vector e -> CovMethod -> [(Double, Vector e)]
                    -> Herm Matrix e
weightedCovWithMean mu t wxs = M.hermCreate $ do
    p <- V.getDim mu
    c <- Herm uplo `fmap` M.new_ (p,p)
    weightedCovWithMeanTo c mu t wxs
    return c
  where
    uplo = defaultCovUplo

-- | Computes and copies the sample covariance matrix to the given
-- destination.
covTo :: (RVector v, BLAS3 e)
      => Herm (STMatrix s) e -> CovMethod -> [v e] -> ST s ()
covTo c@(Herm _ a) t xs = do
    (p,_) <- M.getDim a
    mu <- V.new p 1
    V.meanTo mu xs
    covWithMeanTo c mu t xs

-- | Given the pre-computed mean, computes and copies the sample covariance
-- matrix to the given destination.
covWithMeanTo :: (RVector v1, RVector v2, BLAS3 e)
              => Herm (STMatrix s) e -> v1 e -> CovMethod -> [v2 e] -> ST s ()
covWithMeanTo c@(Herm _ a) mu t xs = do
    (ma,na) <- M.getDim a
    p <- V.getDim mu

    when ((ma,na) /= (p,p)) $ error $
        printf ("covWithMeanTo"
                ++ " (Herm _ <matrix with dim (%d,%d)>)"
                ++ " <vector with dim %d>"
                ++ " _ _"
                ++ ": dimension mismatch")
               ma na p

    xt <- M.new_ (p,n)
    M.withColsM xt $ \xs' ->
        sequence_ [ V.subTo x' mu x
                  | (x,x') <- zip xs xs'
                  ]
    M.hermRankKUpdateM_ (1/df) NoTrans xt 0 c
  where

    n = length xs
    df = fromIntegral $ case t of { MLCov -> n ; UnbiasedCov -> n - 1 }

-- | Computes and copies the weighed sample covariance matrix to the
-- given destination.
weightedCovTo :: (RVector v, BLAS3 e)
              => Herm (STMatrix s) e -> CovMethod -> [(Double, v e)] -> ST s ()
weightedCovTo c@(Herm _ a) t wxs = do
    (p,_) <- M.getDim a
    mu <- V.new p 1
    V.weightedMeanTo mu wxs
    weightedCovWithMeanTo c mu t wxs

-- | Given the pre-computed mean, computes and copies the weighed sample
-- covariance matrix to the given destination.
weightedCovWithMeanTo :: (RVector v1, RVector v2, BLAS3 e)
                      => Herm (STMatrix s) e
                      -> v1 e -> CovMethod -> [(Double, v2 e)]
                      -> ST s ()
weightedCovWithMeanTo c@(Herm _ a) mu t wxs = do
    (ma,na) <- M.getDim a
    p <- V.getDim mu

    when ((ma,na) /= (p,p)) $ error $
        printf ("weightedCovWithMeanTo"
                ++ " (Herm _ <matrix with dim (%d,%d)>):"
                ++ " <vector with dim %d>"
                ++ " _ _"
                ++ " dimension mismatch")
               ma na p

    xt <- M.new_ (p,n)
    M.withColsM xt $ \xs' ->
        sequence_ [  V.subTo x' mu x
                  >> V.scaleM_ (realToFrac $ sqrt (w / invscale)) x'
                  |  (w,x,x') <- zip3 ws xs xs'
                  ]
    M.hermRankKUpdateM_ 1 NoTrans xt 0 c
  where
    (ws0,xs) = unzip wxs
    w_sum = foldl' (+) 0 ws0
    ws = if w_sum == 0 then ws0 else map (/w_sum) ws0
    w2s_sum = foldl' (+) 0 $ map (^^(2::Int)) ws
    invscale = case t of
                   MLCov -> 1
                   UnbiasedCov -> (1 - w2s_sum)
    n = length ws0


{-# LANGUAGE MonoLocalBinds #-}
{-# LANGUAGE FlexibleContexts #-}
module DigitRecognition where

import Numeric.LinearAlgebra as LA
import Codec.Compression.GZip (decompress)
import qualified Data.ByteString.Lazy as BS
import Data.Functor
import qualified Foreign.Storable as ForeignStorable
import GHC.Float (float2Double)
--import System.Random

render :: Integral a => a -> Char
render n = let s = " .:oO@" in s !! (fromIntegral n * length s `div` 256)

--main :: IO ()
--main = do
--  s <- decompress <$> BS.readFile "train-images-idx3-ubyte.gz"
--  l <- decompress <$> BS.readFile "train-labels-idx1-ubyte.gz"
--  n <- (`mod` 60000) <$> randomIO
--  putStr . unlines $
--    [render . BS.index s . (n * 28 ^ 2 + 16 + r * 28 +) <$> [0 .. 27] | r <- [0..27]]
-- print $ BS.index l (n + 8)

clip :: Ord c => c -> c -> c -> c
clip a b = max a . min b

computeProbabilities :: Matrix R -> Matrix R -> R -> Matrix R
computeProbabilities x theta temp = cmap (\x -> if x < 1.0e-16 then 0 else x) p / divisor
    where
        c = (rows theta >< rows x) $ concat $ repeat $ map (maximum . toList) $ toColumns $ (theta LA.<> tr x) / scalar temp
        p = cmap exp $ (theta LA.<> tr x) / scalar temp - c
        divisor = (rows theta >< rows x) $ concat $ repeat $ map (sum . toList) $ toColumns p

contributors :: (ForeignStorable.Storable R, Fractional R) => Int -> Int -> Vector Int -> Matrix R
contributors n k y = (n >< k) $ contributor k y
    where
        contributor k = concatMap (contributor' k) . toList
        contributor' k v = l ++ [1.0] ++ r
            where
                (l, r) = splitAt v $ replicate (k - 1) 0.0

computeCostFunction :: Matrix R -> Vector Int -> Matrix R -> R -> R -> R
computeCostFunction x y theta lambda temp = loss (rows x) (rows theta) y + regularisation
    where
        logProbabilities = cmap log $ cmap (clip 1e-15 (1 - 1e-15)) (computeProbabilities x theta temp)
        regularisation = lambda / 2 * sumElements (cmap (** 2) theta)
        loss n k y' = -sumElements (tr (contributors n k y') * logProbabilities) / fromIntegral n

runGradientDescentIteration :: Matrix R -> Vector Int -> Matrix R -> R -> R -> R -> Matrix R
runGradientDescentIteration x y theta alpha lambda temp = theta - scale alpha grad
    where
        n = rows x
        k = rows theta
        m = tr $ contributors n k y
        probs = cmap (clip 1e-15 (1 - 1e-15)) (computeProbabilities x theta temp)
        grad = scale lambda theta - (m - probs) LA.<> x / (scalar temp * fromIntegral n)

augmentFeatures :: (Element R, Fractional R) => Matrix R -> Matrix R
augmentFeatures x = (rows x >< 1) (repeat 1.0) ||| x

softmaxRegression :: Matrix R -> Vector Int -> R -> R -> R -> Int -> Int -> (Matrix R, [R])
softmaxRegression x y temp alpha lambda k t = iterate step ((k >< cols x') $ repeat 0.0, []) !! t
    where
        x' = augmentFeatures x
        step (theta, cs) =
            (runGradientDescentIteration x' y theta alpha lambda temp, cs ++ [computeCostFunction x' y theta lambda temp])

getClassifification :: Matrix R -> Matrix R -> R -> [Int]
getClassifification x theta temp = map maxIndex $ toColumns probs
    where
        x' = augmentFeatures x
        probs = computeProbabilities x' theta temp

computeTestError :: Matrix R -> Vector Int -> Matrix R -> R -> R
computeTestError x y theta temp = 1 - sum (compare y assignedLabels) / fromIntegral (length assignedLabels)
    where
        assignedLabels = getClassifification x theta temp
        compare y z = zipWith (\x y -> if x == y then 1 else 0) (toList y) z

buildData :: (ForeignStorable.Storable Float, Num Float) => Int -> BS.ByteString -> Matrix R
buildData length ds = cmap float2Double $ (length >< (28 ^ 2)) $ map ((/ 256) . fromIntegral) $ BS.unpack $ BS.drop 16 ds

buildLabels :: (ForeignStorable.Storable Int, Num Int) => BS.ByteString -> Vector Int
buildLabels bs = fromList $ map fromIntegral $ BS.unpack $ BS.drop 8 bs

main :: IO ()
main = do

    tr <- decompress <$> BS.readFile "train-images-idx3-ubyte.gz"
    trl <- decompress <$> BS.readFile "train-labels-idx1-ubyte.gz"
    te <- decompress <$> BS.readFile "t10k-images-idx3-ubyte.gz"
    tel <- decompress <$> BS.readFile "t10k-labels-idx1-ubyte.gz"

    let trainData = buildData 60000 tr
    let trainLabels = buildLabels trl
    let testData = buildData 10000 te
    let testLabels = buildLabels tel
    let temp = 1
    let alpha = 0.3
    let lambda = 0.0001
    let k = 10
    let iter = 150
    let (theta, cs) = softmaxRegression trainData trainLabels temp alpha lambda k iter
--    print $ trainData ? [0]
--    print $ toList trainLabels !! 0
--    print $ (computeProbabilities' (augmentFeatures trainData) theta temp) ¿ [0]
--    let theta' = runGradientDescentIteration (augmentFeatures trainData) trainLabels theta alpha lambda temp
--    print $ computeCostFunction (augmentFeatures trainData) trainLabels theta lambda temp
--    print $ theta ¿ [0]
--    print $ getClassifification testData theta temp
    print $ computeTestError testData testLabels theta temp
    -- print $ take 5 $ toList trainLabels


{-# LANGUAGE NoMonomorphismRestriction, Rank2Types #-}
module STMatrix (
    tests_STMatrix,
    mutatesToMatrix,
    readOnlyMatrix,

    ) where

import Control.Monad
import Control.Monad.ST
import Data.AEq
import Data.Complex( magnitude )
import Debug.Trace
import Test.Framework
import Test.Framework.Providers.QuickCheck2
import Test.QuickCheck hiding ( vector )
import qualified Test.QuickCheck as QC

import Numeric.LinearAlgebra
import qualified Numeric.LinearAlgebra.Vector as V
import qualified Numeric.LinearAlgebra.Matrix as M

import Test.QuickCheck.LinearAlgebra( TestElem(..), Dim2(..), Index2(..),
    Assocs2(..), MatrixPair(..), MatrixTriple(..) )
import qualified Test.QuickCheck.LinearAlgebra as Test

import STVector( readOnlyVector )
import Typed

tests_STMatrix = testGroup "STMatrix"
    [ testPropertyI "new_" prop_new_
    , testPropertyI "new" prop_new
    , testPropertyI "newCopy" prop_newCopy
    , testPropertyI "copyTo" prop_copyTo
    , testPropertyI "read" prop_read
    , testPropertyI "write" prop_write
    , testPropertyI "modify" prop_modify
    , testPropertyI "getElems" prop_getElems
    , testPropertyI "getElems'" prop_getElems'
    , testPropertyI "getAssocs" prop_getAssocs
    , testPropertyI "getAssocs'" prop_getAssocs'
    , testPropertyI "setElems" prop_setElems
    , testPropertyI "setAssocs" prop_setAssocs
    , testPropertyI "mapTo" prop_mapTo
    , testPropertyI "zipWithTo" prop_zipWithTo
    , testPropertyDZ "shiftDiagM_" prop_shiftDiagM_ prop_shiftDiagM_
    , testPropertyDZ "shiftDiagWithScaleToM_"
        prop_shiftDiagWithScaleM_ prop_shiftDiagWithScaleM_
    , testPropertyDZ "addTo" prop_addTo prop_addTo
    , testPropertyDZ "subTo" prop_subTo prop_subTo
    , testPropertyDZ "scaleM_" prop_scaleM_ prop_scaleM_
    , testPropertyDZ "addWithScaleM_" prop_addWithScaleM_ prop_addWithScaleM_
    , testPropertyDZ "scaleRowsM" prop_scaleRowsM_ prop_scaleRowsM_
    , testPropertyDZ "scaleColsM_" prop_scaleColsM_ prop_scaleColsM_
    , testPropertyDZ "negateTo" prop_negateTo prop_negateTo
    , testPropertyDZ "conjugateTo" prop_conjugateTo prop_conjugateTo
    , testPropertyDZ "swapRows" prop_swapRows prop_swapRows
    , testPropertyDZ "swapCols" prop_swapCols prop_swapCols
    ]

prop_new_ t (Dim2 n) =
    (M.dim $ typed t $ M.create $ M.new_ n) === n

prop_new t (Dim2 n) e =
    (M.create $ M.new n e) === (typed t $ M.constant n e)

prop_newCopy t x =
    (M.create $ M.newCopy x) === x
  where
    _ = typed t x

prop_copyTo t (MatrixPair x y) = runST $
    x `readOnlyMatrix` \mx ->
    y `mutatesToMatrix` x $ \my ->
        M.copyTo my mx
  where
    _ = typed t x

prop_read t (Index2 n i) =
    forAll (typed t `fmap` Test.matrix n) $ \x -> runST $
        x `readOnlyMatrix` \mx -> do
            e <- M.read mx i
            return $ e === M.at x i

prop_write t (Index2 n i) e =
    forAll (Test.matrix n) $ \x -> runST $
        x `mutatesToMatrix` (x `M.update` [(i,e)]) $ \mx -> do
            M.write mx i e
  where
    _ = e == t

prop_modify t (Index2 n i) (Blind f) =
    forAll (Test.matrix n) $ \x -> runST $
        x `mutatesToMatrix`
            (typed t $ x `M.update` [(i, f $ M.at x i)]) $ \mx ->
                M.modify mx i f

prop_getElems t x =
    forAll arbitrary $ \(Blind f) -> runST $
        x `mutatesToMatrix` (M.map f x) $ \mx -> do
            es <- M.getElems mx
            M.mapTo mx f mx
            return $ es === M.elems (M.map f x)
  where
    _ = typed t x

prop_getElems' t x =
    forAll arbitrary $ \(Blind f) -> runST $
        x `mutatesToMatrix` (M.map f x) $ \mx -> do
            es <- M.getElems' mx
            M.mapTo mx f mx
            return $ es === M.elems x
  where
    _ = typed t x

prop_getAssocs t x =
    forAll arbitrary $ \(Blind f) -> runST $
        x `mutatesToMatrix` (M.map f x) $ \mx -> do
            ies <- M.getAssocs mx
            M.mapTo mx f mx
            return $ ies === M.assocs (M.map f x)
  where
    _ = typed t x

prop_getAssocs' t x =
    forAll arbitrary $ \(Blind f) -> runST $
        x `mutatesToMatrix` (M.map f x) $ \mx -> do
            ies <- M.getAssocs' mx
            M.mapTo mx f mx
            return $ ies === M.assocs x
  where
    _ = typed t x

prop_setElems t x =
    forAll (QC.vector $ m*n) $ \es -> runST $
        x `mutatesToMatrix` (M.fromList (m,n) es) $ \mx ->
            M.setElems mx es
  where
    (m,n) = M.dim x
    _ = typed t x

prop_setAssocs t (Assocs2 mn ies) =
    forAll (Test.matrix mn) $ \x -> runST $
        x `mutatesToMatrix` (typed t $ M.update x ies) $ \mx ->
            M.setAssocs mx ies

prop_mapTo t (Blind f) = binaryProp t
    (\x -> M.map f x)
    (\dst mx -> M.mapTo dst f mx)

prop_zipWithTo t (Blind f) = ternaryProp t
    (\x y -> M.zipWith f x y)
    (\dst mx my -> M.zipWithTo dst f mx my)

prop_shiftDiagM_ t a =
    forAll (Test.vector (min m n)) $ \s -> runST $
        s `readOnlyVector` \ms ->
        a `mutatesToMatrix` (M.shiftDiag s a) $ \ma ->
            M.shiftDiagM_ ms ma
  where
    (m,n) = M.dim a
    _ = typed t a

prop_shiftDiagWithScaleM_ t e a =
    forAll (Test.vector (min m n)) $ \s -> runST $
        s `readOnlyVector` \ms ->
        a `mutatesToMatrix` (M.shiftDiagWithScale e s a) $ \ma ->
            M.shiftDiagWithScaleM_ e ms ma
  where
    (m,n) = M.dim a
    _ = typed t a

prop_addTo t = ternaryProp t M.add M.addTo

prop_subTo t = ternaryProp t M.sub M.subTo

prop_scaleM_ t e a = runST $
    a `mutatesToMatrix` (M.scale e a) $ \ma ->
        M.scaleM_ e ma
  where
    _ = typed t a

prop_addWithScaleM_ t alpha (MatrixPair x y) = runST $
    x `readOnlyMatrix` \mx ->
    y `mutatesToMatrix` (M.addWithScale alpha x y) $ \my ->
        M.addWithScaleM_ alpha mx my
  where
    _ = typed t x

prop_scaleRowsM_ t a =
    forAll (Test.vector m) $ \s -> runST $
        s `readOnlyVector` \ms ->
        a `mutatesToMatrix` (M.scaleRows s a) $ \ma ->
            M.scaleRowsM_ ms ma
  where
    (m,n) = M.dim a
    _ = typed t a

prop_scaleColsM_ t a =
    forAll (Test.vector n) $ \s -> runST $
        s `readOnlyVector` \ms ->
        a `mutatesToMatrix` (M.scaleCols s a) $ \ma ->
            M.scaleColsM_ ms ma
  where
    (m,n) = M.dim a
    _ = typed t a

prop_negateTo t = binaryProp t M.negate M.negateTo

prop_conjugateTo t = binaryProp t M.conjugate M.conjugateTo

prop_swapRows t =
    forAll arbitrary $ \(Index2 (m,n) (i1,_)) ->
    forAll (choose (0,m-1)) $ \i2 ->
    forAll (typed t `fmap` Test.matrix (m,n)) $ \a ->
        let a' = M.update a $ [ ((i2,j),e) | (j,e) <- V.assocs (M.row a i1)
                              ] ++
                              [ ((i1,j),e) | (j,e) <- V.assocs (M.row a i2)
                              ]
        in runST $
            a `mutatesToMatrix` a' $ \ma ->
            M.swapRows ma i1 i2

prop_swapCols t =
    forAll arbitrary $ \(Index2 (m,n) (_,j1)) ->
    forAll (choose (0,n-1)) $ \j2 ->
    forAll (typed t `fmap` Test.matrix (m,n)) $ \a ->
        let a' = M.update a $ [ ((i,j2),e) | (i,e) <- V.assocs (M.col a j1)
                              ] ++
                              [ ((i,j1),e) | (i,e) <- V.assocs (M.col a j2)
                              ]
        in runST $
            a `mutatesToMatrix` a' $ \ma ->
            M.swapCols ma j1 j2

binaryProp :: (AEq e, Arbitrary e, Show e, Storable e, Testable a)
            => e
            -> (Matrix e -> Matrix e)
            -> (forall s . STMatrix s e -> STMatrix s e -> ST s a)
            -> Property
binaryProp t imm_f f = let
    prop1 t x = runST $
        x `mutatesToMatrix` (imm_f x) $ \mx -> f mx mx
      where _ = typed t x
    prop2 t (MatrixPair x y) = runST $ let _ = typed t x in
        x `readOnlyMatrix` \mx ->
        y `mutatesToMatrix` (imm_f x) $ \my -> f my mx
      where _ = typed t x
    in (    label "1" prop1
        .&. label "2" prop2
       )

ternaryProp :: (AEq e, Arbitrary e, Show e, Storable e, Testable a)
            => e
            -> (Matrix e -> Matrix e -> Matrix e)
            -> (forall s . STMatrix s e -> STMatrix s e -> STMatrix s e -> ST s a)
            -> Property
ternaryProp t imm_f f = let
    prop1 t x = runST $
        x `mutatesToMatrix` (imm_f x x) $ \mx -> f mx mx mx
      where _ = typed t x
    prop2a t (MatrixPair x y) = runST $
        x `readOnlyMatrix` \mx ->
        y `mutatesToMatrix` (imm_f x y) $ \my -> f my mx my
      where _ = typed t x
    prop2b t (MatrixPair x y) = runST $
        x `mutatesToMatrix` (imm_f x y) $ \mx ->
        y `readOnlyMatrix` \my -> f mx mx my
      where _ = typed t x
    prop3 t (MatrixTriple x y z) = runST $
        x `readOnlyMatrix` \mx ->
        y `readOnlyMatrix` \my ->
        z `mutatesToMatrix` (imm_f x y) $ \mz -> f mz mx my
      where _ = typed t x
    in (    label "1" prop1
        .&. label "2a" prop2a
        .&. label "2b" prop2b
        .&. label "3" prop3
       )

readOnlyMatrix :: (Storable e, AEq e, Testable prop, Show e)
               => Matrix e
               -> (STMatrix s e ->  ST s prop)
               -> ST s Property
readOnlyMatrix x = mutatesToMatrix x x

mutatesToMatrix :: (Storable e, AEq e, Testable prop, Show e)
                => Matrix e
                -> Matrix e
                -> (STMatrix s e -> ST s prop)
                -> ST s Property
mutatesToMatrix x x_new f = do
    mx <- M.newCopy x
    prop <- f mx
    x' <- M.freeze mx
    return $ prop .&. (x' === x_new)


{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE OverloadedStrings #-}

module MLP.DataSet (
   toPat,
   normPat
) where

import Numeric.LinearAlgebra.Data (cols, Matrix, R, takeColumns, dropColumns, toRows)
import Numeric.LinearAlgebra.Static (create, zipWithVector)
import MLP.Types (Pattern(..), Topology(..), zipWithPat)
import MLP.Network (scalarR)
import qualified Data.Vector.Storable as V
import qualified Data.Text as T
import Control.Monad (when)
import GHC.TypeLits (KnownNat)
import Data.Singletons ( SingI(sing), SingKind(fromSing) )
import Data.Maybe (mapMaybe)

toPat :: forall i o. (KnownNat i, KnownNat o) => Matrix R -> Either String [Pattern i o]
toPat m
   | cols m == (dimI + dimO) = Right $ zipWith Pattern is os
   | otherwise = Left "Mismatch on number of columns"

   where
      [dimI, dimO] = map fromIntegral . fromSing $ sing @'[i,o]
      is = mapMaybe create . toRows . takeColumns dimI $ m
      os = mapMaybe create . toRows . dropColumns dimI $ m

normPat :: forall i o. (KnownNat i, KnownNat o) => [Pattern i o] -> [Pattern i o]
normPat [] = []
normPat xxs@(x:xs) = map (\p -> (p - minp) / scalep) xxs
   where
      zeroPat = Pattern (scalarR 0) (scalarR 0)
      minp = foldl (zipWithPat min) x xs
      maxp = foldl (zipWithPat max) x xs
      scalep = zipWithPat max zeroPat (maxp - minp)


-- |
-- Module:       Math.NumberTheory.RootsOfUnityTests
-- Copyright:    (c) 2018 Bhavik Mehta
-- License:      MIT
-- Maintainer:   Andrew Lelechenko <andrew.lelechenko@gmail.com>
--
-- Tests for Math.NumberTheory.RootsOfUnity
--

module Math.NumberTheory.RootsOfUnityTests where

import Test.Tasty

import Data.Complex
import Data.Ratio
import Data.Semigroup

import Math.NumberTheory.DirichletCharacters (toRootOfUnity, toComplex)
import Math.NumberTheory.TestUtils (testSmallAndQuick, Positive(..))

rootOfUnityTest :: Integer -> Positive Integer -> Bool
rootOfUnityTest n (Positive d) = toComplex ((d `div` gcd n d) `stimes` toRootOfUnity (n % d)) == (1 :: Complex Double)

testSuite :: TestTree
testSuite = testSmallAndQuick "RootOfUnity contains roots of unity" rootOfUnityTest


{-# LANGUAGE OverloadedStrings     #-}
{-# LANGUAGE QuasiQuotes #-}

module Taiji.Utils.Plot.ECharts.Boxplot (boxplot) where

import Language.Javascript.JMacro
import Statistics.Quantile
import qualified Data.Text as T
import qualified Data.Vector.Unboxed as U

import Taiji.Utils.Plot.ECharts.Types

boxplot :: [(T.Text, [Double])] -> EChart
boxplot input = mkEChart [jmacroE| {
    grid: { containLabel: true },
    xAxis: {
        type: "category",
        data: `names`,
        axisLabel: { rotate: 45 }
    },
    yAxis: {
        axisLine: {onZero: false},
        type: "value"
    },
    series: [{
        name: "boxplot",
        type: "boxplot",
        data: `dat'`
    }]
    } |]
  where
    (names, dat) = unzip input
    dat' = flip map dat $ \xs -> fst $ range $ U.fromList xs

range :: U.Vector Double -> ([Double], U.Vector Double)
range xs = ([U.minimum xs', q1, q2, q3, U.maximum xs'], outliers)
  where
    (xs', outliers) = U.partition (\x -> x >= lower && x <= upper) xs
    lower = q1 - 1.5 * iqr
    upper = q3 + 1.5 * iqr
    iqr = q3 - q1
    [q1, q2, q3] = quantiles def [1, 2, 3] 4 xs


module Main (main) where

import           Control.Monad
import           System.IO
import           Text.Printf (printf)

import           Data.Number.LogFloat (LogFloat)
import qualified Data.Number.LogFloat as LF
import           Statistics.Distribution
import           Statistics.Distribution.Binomial

-- | Compute the probability of seeing more than @k@ blocks in @n@ slots
moreThanK ::
     Double  -- ^ Active slot coefficient
  -> Int     -- ^ Security parameter (@k@)
  -> Int     -- ^ Number of slots
  -> LogFloat
moreThanK f k n =
    LF.sum [LF.logToLogFloat $ logProbability d i | i <- [k + 1 .. n]]
  where
    d :: BinomialDistribution
    d = binomial n f

defaultS ::
     Double  -- ^ Active slot coefficient
  -> Int     -- ^ Security parameter (@k@)
  -> Int
defaultS f k = floor (fromIntegral k / f) `div` 4

main :: IO ()
main = do
    forM_ [s .. 4 * s] $ \n -> do
      putStrLn $ show n ++ "\t" ++ showLogFloat (moreThanK f k n)
      hFlush stdout
  where
    f = 0.05
    k = 2160
    s = defaultS f k

{-------------------------------------------------------------------------------
  LogFloat util
-------------------------------------------------------------------------------}

showLogFloat :: LogFloat -> String
showLogFloat lf = printf "%6.4f * 10 ^ %d" m e
  where
    (m, e) = logFloatToScientific lf

logFloatToScientific :: LogFloat -> (Double, Int)
logFloatToScientific lf = (m, e)
  where
    l :: Double
    l = LF.logFromLogFloat lf

    e :: Int
    e = floor $ l / log10

    m :: Double
    m = exp $ l - log10 * fromIntegral e

    log10 :: Double
    log10 = log 10


module Arrow where

import           Graphics.Gloss
import           Data.Complex

type Arrow = Complex Float

vectorize :: Arrow -> Vector
vectorize a = (realPart a, imagPart a)

-- Rotate an arrow according to its frequency in cycles per second
updateArrow :: Arrow -> Integer -> Float -> Arrow
updateArrow arrow frequency t =
  cis (2 * pi * fromInteger frequency * t) * arrow

-- Create a path through a list of arrows arranged tip to tail
arrange :: [Arrow] -> Path
arrange = map vectorize . scanl (+) (0 :+ 0)

endpoint :: [Arrow] -> Point
endpoint = last . arrange


module Fuml.Optimisation where

import Fuml.Core
import Fuml.Optimisation.BFGS
import Fuml.Optimisation.SGD
import Numeric.LinearAlgebra
import qualified Data.Vector.Storable as VS
import Fuml.Base.Logistic
import Data.Maybe (fromMaybe)
import Control.Monad.Identity
import Data.Random

type BfgsModel = (Vector Double, Matrix Double)

data CostGradientModel a b = CostGradientModel
  { costFunction :: ([(Vector Double, a)] -> Vector Double -> Double)
  , costGradient :: ([(Vector Double, a)] -> Vector Double -> Vector Double)
  , generatePrediction :: Vector Double -> Vector Double -> b
  }

bfgsSupervisor :: BFGSOpts
               -> CostGradientModel a b
               -> Supervisor Identity a BfgsModel b
bfgsSupervisor bopts (CostGradientModel ll gradll predf) = Supervisor $ \mp theData ->
  let start = fromMaybe (VS.map (const 0) $ fst $ head theData) $ fmap fst mp
      hessInit = fromMaybe (ident $ size start) $ fmap snd mp
  in case bfgsWith bopts (ll theData) (gradll theData) start hessInit of
       Left s -> error s
       Right (p, h) -> return $ Predict (p, h) (predf p)

defBfgs = BFGSOpts 1e-7 1e-7 200

data SGDModel = SGDModel
  { sgdFit :: Vector Double
  , lossCurve :: [Double]
  }

sgdSupervisor :: SGDOpts
              -> CostGradientModel a b
              -> Supervisor RVar a SGDModel b
sgdSupervisor opts (CostGradientModel ll gradll predf) = Supervisor $ \mp theData ->
  let start = fromMaybe (VS.map (const 0) $ fst $ head theData) $ fmap sgdFit mp
      mkPredict sgdModel = Predict sgdModel (predf (sgdFit sgdModel))
      go 0 p losses = return $ SGDModel p $ reverse losses
      go n p losses = do
        nextp <- sgdEpoch opts gradll theData p
        let loss = ll theData p
        go (n-1) nextp (loss:losses)
  in fmap mkPredict $ go (epochs opts) start []

logisticCost :: Double -> CostGradientModel Bool Double
logisticCost delta = CostGradientModel (logLikelihood delta) (gradLogLikelihood delta) (\beta v -> logit $ v `dot` beta)

logisticBfgs :: Double -> Supervisor Identity Bool BfgsModel Double
logisticBfgs delta = bfgsSupervisor defBfgs (logisticCost delta)


module Binary where

import Nnaskell
import Numeric.LinearAlgebra.Data

binaryNN :: IO NN
binaryNN = loadNN "./data/binary.txt"

binaryTD :: [(Vector R, Vector R)]
binaryTD = zip (map binaryInput ns) (map binaryOutput ns)
    where ns = [0 .. 15]

binaryInput :: Int -> Vector R
binaryInput x = vector $ map fromIntegral [ x `mod` 2
                                          , (x `div` 2) `mod` 2
                                          , (x `div` 4) `mod` 2
                                          , (x `div` 8) `mod` 2]

binaryOutput :: Int -> Vector R
binaryOutput x = accum (vector $ replicate 16 0.0) (\a _ -> a) [(x, 1.0)]