{-# LANGUAGE TypeApplications #-}
{-# OPTIONS_GHC -fno-warn-type-defaults #-}

module SSSFN
  ( -- * Domain
    np,

    -- * Grid
    Grid (..),
    gunit,
    (!),
    gcoerce,

    -- * Operators
    Op (..),
    ounit,
    ocoerce,

    -- * Grids are functors etc.
    gmap,
    gpure,
    gzipWith,
    gfoldMap,
    gcount,
    gmaxabs,
    gsum,

    -- * Operators are functors etc.
    omap,
    opure,
    ozipWith,
    ofoldMap,
    ocount,
    omaxabs,
    osum,

    -- * Operators
    (#>),
    transpose,
    inv,

    -- * Coordinates
    isOrigin,
    coordU,
    coordV,
    coordT,
    coordR,
    coordR1, -- remove?

    -- * Sampling functions
    gsample,

    -- * Evaluating functions
    geval,
    gevalCoeffs,

    -- * Derivatives
    vandermonde,
    expfilter,
    delta,
    symmetrizeU,
    symmetrizeV,
    symmetrizeR,
    derivU,
    derivV,
    derivUU,
    derivUV,
    derivVV,
    derivT,
    derivR,
    derivTT,
    derivRR,
    r1derivR,
    r1derivU,
    r1derivV,
    r2project,
    zproject,
    bndU,
    bndV,
    metric,

    -- * Systems of equations
    solve,
    solveLS,

    -- * MaxAbs
    maxabs,
    getMaxAbs,
  )
where

import Control.Exception (assert)
import Data.Maybe
import Data.Poly hiding (scale)
import Data.Semigroup hiding ((<>))
import qualified Data.Vector.Unboxed as U
import Data.VectorSpace
import Foreign
import qualified GHC.Exts as GHC
import Numeric.LinearAlgebra hiding ((!), (#>), inv)
import qualified Numeric.LinearAlgebra as L
import Numeric.LinearAlgebra.Devel
import qualified Test.QuickCheck as QC
import Prelude hiding ((<>))
import qualified Prelude as P

default (Int)

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

-- | Grid size
np :: Int
np = 33

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

-- | Grid datatype
newtype Grid a = Grid {getGrid :: Vector a}
  deriving (Eq, Ord, Read, Show)

instance (QC.Arbitrary a, Storable a) => QC.Arbitrary (Grid a) where

  arbitrary = Grid . fromList <$> QC.vector (np ^ 2)

  shrink = map (Grid . fromList) . traverse QC.shrink . toList . getGrid

instance (Storable a) => GHC.IsList (Grid a) where

  type Item (Grid a) = a

  fromList = Grid . fromList

  toList = toList . getGrid

(!) :: (Container Vector a) => Grid a -> (Int, Int) -> a
(!) (Grid x) (i, j) = x `atIndex` (i * np + j)

gunit :: (Num a, Storable a) => (Int, Int) -> Grid a
gunit (i, j) =
  let ij = i * np + j
   in Grid $ fromList [if k == ij then 1 else 0 | k <- [0 .. np ^ 2 -1]]

gcoerce :: forall a b. (Storable a, Storable b) => Grid a -> Grid b
gcoerce (Grid xss) =
  assert (na * sa `mod` sb == 0) $
    (Grid $ fromByteString $ toByteString xss)
  where
    na = cols (asRow xss)
    sa = sizeOf @a undefined
    sb = sizeOf @b undefined

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

-- | Operator datatype
newtype Op a = Op {getOp :: Matrix a}
  deriving (Read, Show)

instance (QC.Arbitrary a, Element a) => QC.Arbitrary (Op a) where

  arbitrary = Op . reshape (np ^ 2) . fromList <$> QC.vector (np ^ 4)

  shrink =
    map (Op . reshape (np ^ 2) . fromList)
      . traverse QC.shrink
      . (toList . flatten . getOp)

instance Element a => GHC.IsList (Op a) where

  type Item (Op a) = a

  fromList = Op . reshape (np ^ 2) . fromList

  toList = toList . flatten . getOp

ounit :: (Element a, Num a) => Op a
ounit = Op $ ident (np ^ 2)

ocoerce :: forall a b. (Container Vector a, Num a, Storable b) => Op a -> Op b
ocoerce (Op x) =
  assert (ma * sa `mod` sb == 0 && na * sa `mod` sb == 0) $
    (Op $ reshape nb $ getGrid $ gcoerce $ Grid $ flatten x)
  where
    (ma, na) = size x
    sa = sizeOf @a undefined
    sb = sizeOf @b undefined
    nb = na * sa `div` sb

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

-- | Grids are functors etc.
-- gmap :: (Container Vector a, Element b) => (a -> b) -> Grid a -> Grid b
-- gmap f (Grid x) = Grid (cmap f x)
gmap :: (Storable a, Storable b) => (a -> b) -> Grid a -> Grid b
gmap f (Grid x) = Grid (mapVectorWithIndex (\i a -> f a) x)

gpure :: Container Vector a => a -> Grid a
gpure a = Grid $ konst a (np ^ 2)

gzipWith ::
  (Storable a, Storable b, Storable c) =>
  (a -> b -> c) ->
  Grid a ->
  Grid b ->
  Grid c
gzipWith f (Grid x) (Grid y) = Grid (zipVectorWith f x y)

gfoldMap :: (Storable a, Monoid b) => (a -> b) -> Grid a -> b
gfoldMap f (Grid x) = foldVector (\a b -> f a P.<> b) mempty x

gcount :: Container Vector a => Grid a -> Int
gcount (Grid x) = size x

gmaxabs :: (Container Vector a, Num a, Ord a) => Grid a -> a
gmaxabs = getMaxAbs . gfoldMap maxabs

gsum :: (Container Vector a) => Grid a -> a
gsum (Grid x) = sumElements x

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

-- | Operators are functors etc.
-- omap :: (Container Vector a, Num a, Element b) => (a -> b) -> Op a -> Op b
-- omap f (Op x) = Op (cmap f x)
omap :: (L.Element a, Storable b) => (a -> b) -> Op a -> Op b
omap f (Op x) = Op $ reshape n $ getGrid $ gmap f $ Grid $ flatten x
  where
    n = cols x

opure :: (Container Vector a, Num a) => a -> Op a
opure a = Op (konst a (np ^ 2, np ^ 2))

ozipWith ::
  (Element a, Element b, Storable c) =>
  (a -> b -> c) ->
  Op a ->
  Op b ->
  Op c
ozipWith f (Op x) (Op y) =
  Op $ reshape (np ^ 2) $ zipVectorWith f (flatten x) (flatten y)

ofoldMap :: (Element a, Monoid b) => (a -> b) -> Op a -> b
ofoldMap f (Op x) = foldVector (\a b -> f a P.<> b) mempty (flatten x)

ocount :: (Container Vector a, Num a) => Op a -> Int
ocount (Op x) = let (ni, nj) = size x in ni * nj

omaxabs :: (Container Vector a, Num a, Ord a) => Op a -> a
omaxabs = getMaxAbs . ofoldMap maxabs

osum :: (Container Vector a, Num a) => Op a -> a
osum (Op x) = sumElements x

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

-- | Grids are an additive group
instance
  (Container Vector a, Num a) =>
  AdditiveGroup (Grid a)
  where

  zeroV = gpure 0

  (^+^) = gzipWith (+)

  negateV = gmap negate

-- | Grids are a vector space
instance
  (Container Vector a, Num a) =>
  VectorSpace (Grid a)
  where

  type Scalar (Grid a) = a

  (*^) a = gmap (a *)

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

-- | Operators form an additive group
instance (Container Vector a, Num a) => AdditiveGroup (Op a) where

  zeroV = opure 0

  (^+^) = ozipWith (+)

  negateV = omap negate

-- | Operators form a vector space
instance (Container Vector a, Num a) => VectorSpace (Op a) where

  type Scalar (Op a) = a

  (*^) a = omap (a *)

-- | Operators form a ring
instance (Container Vector a, Num a, Numeric a) => Num (Op a) where

  (+) = (^+^)

  negate = negateV

  (-) = (^-^)

  Op x * Op y = Op (x <> y)

  abs = omap abs

  signum = omap signum

  fromInteger i = fromInteger i *^ ounit

-- | Matrix-vector multiplication
infixr 8 #>

(#>) :: Numeric a => Op a -> Grid a -> Grid a
Op x #> Grid y = Grid (x L.#> y)

-- | Matrix transpose
transpose :: Transposable (Matrix a) (Matrix a) => Op a -> Op a
transpose (Op x) = Op (tr x)

-- | Matrix inverse
inv :: Field a => Op a -> Op a
inv (Op x) = Op (L.inv x)

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

-- | Coordinate of a grid point
coord :: Floating a => Int -> a
coord i =
  let imin = 0
      imax = np - 1
      umin = -1
      umax = 1
      u =
        fromIntegral (i - imax) / fromIntegral (imin - imax) * umin
          + fromIntegral (i - imin) / fromIntegral (imax - imin) * umax
   in sin (pi / 2 * u)

-- | Coordinates of all grid points
coords :: (Floating a, Storable a) => Vector a
coords = fromList [coord i | i <- [0 .. np -1]]

poly2vector :: (Num a, Storable a, U.Unbox a) => UPoly a -> Vector a
poly2vector p =
  let v = unPoly p
      n = U.length v
   in fromList [if i < n then v U.! i else 0 | i <- [0 .. np -1]]

-- | Evalute a polynomial at a grid point
evalcoeffs :: (Eq a, Num a, Storable a, U.Unbox a) => a -> Vector a
evalcoeffs u = fromList [eval (mon i) u | i <- [0 .. np -1]]
  where
    mon :: (Eq a, Num a, U.Unbox a) => Int -> UPoly a
    mon i = monomial (fromIntegral i) 1

-- | Evaluate a polynomial at all grid points
gvandermonde :: (Element a, Eq a, Floating a, U.Unbox a) => Matrix a
gvandermonde = fromRows [evalcoeffs u | u <- toList coords]

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

-- Note: hmatrix is row-major by default
mkGrid :: (Num a, Storable a) => (Int -> a) -> (Int -> a) -> Grid a
mkGrid fx fy =
  Grid $ fromList [fx i * fy j | i <- [0 .. np -1], j <- [0 .. np -1]]

isOrigin :: Grid Bool
isOrigin =
  Grid $
    fromList
      [ i == j
        | i <- [0 .. np -1],
          j <- [0 .. np -1]
      ]

-- | Coordinates of all grid points
coordU :: (Container Vector a, Floating a) => Grid a
coordU = mkGrid (coords `atIndex`) (const 1)

-- | Coordinates of all grid points
coordV :: (Container Vector a, Floating a) => Grid a
coordV = mkGrid (const 1) (coords `atIndex`)

-- | Combined coordinates
coordT :: (Container Vector a, Floating a) => Grid a
coordT = coordV ^+^ coordU

coordR :: (Container Vector a, Floating a) => Grid a
coordR = coordV ^-^ coordU

coordR1 :: (Container Vector a, RealFloat a) => Grid a
coordR1 = gzipWith (\isOrig r -> if isOrig then 0 else recip r) isOrigin coordR

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

-- | Sample a function at grid points
gsample ::
  (Container Vector a, Floating a, Storable b) =>
  ((a, a) -> b) ->
  Grid b
gsample f = gzipWith (curry f) (coordU) (coordV)

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

-- | Evaluate a grid at a point
-- f / Identity
geval ::
  ( Container Vector a,
    Eq a,
    Field a,
    U.Unbox a,
    Container Vector b,
    VectorSpace b,
    Scalar b ~ a
  ) =>
  Grid b ->
  (a, a) ->
  b
geval g (u, v) =
  let ucoeffs = tr (L.inv gvandermonde) L.#> evalcoeffs u
      vcoeffs = tr (L.inv gvandermonde) L.#> evalcoeffs v
   in sumV
        [ ((ucoeffs `atIndex` i) * (vcoeffs `atIndex` j)) *^ (g ! (i, j))
          | i <- [0 .. np - 1],
            j <- [0 .. np - 1]
        ]

gevalCoeffs ::
  ( Container Vector a,
    Eq a,
    Field a,
    U.Unbox a
  ) =>
  (a, a) ->
  Grid a
gevalCoeffs (u, v) =
  let ucoeffs = tr (L.inv gvandermonde) L.#> evalcoeffs u
      vcoeffs = tr (L.inv gvandermonde) L.#> evalcoeffs v
   in Grid $
        fromList
          [ (ucoeffs `atIndex` i) * (vcoeffs `atIndex` j)
            | i <- [0 .. np - 1],
              j <- [0 .. np - 1]
          ]

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

-- | Derivative of a basis function
deriv1 :: (Eq a, Num a, Storable a, U.Unbox a) => Int -> Vector a
deriv1 i = poly2vector $ deriv $ monomial (fromIntegral i) 1

-- | Derivatives of all basis functions
derivs1 :: (Element a, Eq a, Num a, U.Unbox a) => Matrix a
derivs1 = fromColumns [deriv1 i | i <- [0 .. np -1]]

derivs :: (Eq a, Field a, U.Unbox a) => Matrix a
derivs = gvandermonde <> derivs1 <> L.inv gvandermonde

-- derivs2 :: (Eq a, Field a, U.Unbox a) => Matrix a
-- derivs2 = derivs <> derivs

expfilter1 :: (Container Vector a, RealFloat a) => Matrix a
expfilter1 =
  diag $
    fromList
      [ let n0 = np `div` 2
         in if i <= n0
              then 1
              else peps ** (fromIntegral (i - n0) / fromIntegral (np - n0))
        | i <- [0 .. np -1]
      ]

expfilters :: (Container Vector a, Field a, RealFloat a, U.Unbox a) => Matrix a
expfilters = gvandermonde <> expfilter1 <> L.inv gvandermonde

bnd :: (Element a, Fractional a) => Matrix a
bnd =
  diag $
    fromList
      [ if i == 0
          then -1 / fromIntegral (np - 1)
          else if i == np - 1 then 1 / fromIntegral (np - 1) else 0
        | i <- [0 .. np - 1]
      ]

weights :: (Element a, Fractional a) => Matrix a
weights =
  diag $
    fromList
      [ if i == 0 || i == np - 1
          then 1 / fromIntegral (np - 1)
          else 2 / fromIntegral (np - 1)
        | i <- [0 .. np - 1]
      ]

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

mkOp :: (Element a, Num a) => ((Int, Int) -> a) -> ((Int, Int) -> a) -> Op a
mkOp fx fy =
  Op $
    fromLists
      [ [ fx (ir, ic) * fy (jr, jc)
          | ic <- [0 .. np - 1],
            jc <- [0 .. np - 1]
        ]
        | ir <- [0 .. np - 1],
          jr <- [0 .. np - 1]
      ]

vandermonde :: (Container Vector a, Eq a, Floating a, U.Unbox a) => Op a
vandermonde = mkOp (gvandermonde `atIndex`) (gvandermonde `atIndex`)

expfilter :: (Container Vector a, Field a, RealFloat a, U.Unbox a) => Op a
expfilter = mkOp (expfilters `atIndex`) (expfilters `atIndex`)

delta :: Field a => Op a
delta = mkOp delta1 delta1
  where
    delta1 (i, j) = if i == j then 1 else 0

diagonal :: (Element a, Num a) => Grid a -> Op a
diagonal (Grid xs) = Op $ diag xs

originR :: Field a => Op a
originR =
  Op $
    fromLists
      [ [ if ic == ir && jc == jr && ic == jc then 1 else 0
          | ic <- [0 .. np - 1],
            jc <- [0 .. np - 1]
        ]
        | ir <- [0 .. np - 1],
          jr <- [0 .. np - 1]
      ]

valueAtOriginR :: (Eq a, Field a, U.Unbox a) => Op a
valueAtOriginR =
  Op $
    fromRows
      [ let t = coordT ! (ir, jr)
            r = 0
            v = (t + r) / 2
            u = (t - r) / 2
         in getGrid (gevalCoeffs (u, v))
        | ir <- [0 .. np - 1],
          jr <- [0 .. np - 1]
      ]

reverseU :: (Element a, Num a) => Op a
reverseU = mkOp reverse1 delta1
  where
    reverse1 (i, j) = if i == np -1 - j then 1 else 0
    delta1 (i, j) = if i == j then 1 else 0

reverseV :: (Element a, Num a) => Op a
reverseV = mkOp delta1 reverse1
  where
    reverse1 (i, j) = if i == np -1 - j then 1 else 0
    delta1 (i, j) = if i == j then 1 else 0

-- change (t,r) -> (t,-r); (u,v) -> (v,u)
reverseR :: (Element a, Num a) => Op a
reverseR =
  Op $
    fromLists
      [ [ if ic == jr && jc == ir then 1 else 0
          | ic <- [0 .. np - 1],
            jc <- [0 .. np - 1]
        ]
        | ir <- [0 .. np - 1],
          jr <- [0 .. np - 1]
      ]

symmetrizeU :: Field a => Op a
symmetrizeU = (1 / 2) *^ (delta ^+^ reverseU)

symmetrizeV :: Field a => Op a
symmetrizeV = (1 / 2) *^ (delta ^+^ reverseV)

symmetrizeR :: Field a => Op a
symmetrizeR = (1 / 2) *^ (delta ^+^ reverseR)

-- | Derivatives of a grid
derivU :: (Eq a, Field a, U.Unbox a) => Op a
derivU = mkOp (derivs `atIndex`) delta1
  where
    delta1 (i, j) = if i == j then 1 else 0

derivV :: (Eq a, Field a, U.Unbox a) => Op a
derivV = mkOp delta1 (derivs `atIndex`)
  where
    delta1 (i, j) = if i == j then 1 else 0

derivUU :: (Eq a, Field a, U.Unbox a) => Op a
derivUU = derivU * derivU

derivUV :: (Eq a, Field a, U.Unbox a) => Op a
derivUV = derivU * derivV

derivVV :: (Eq a, Field a, U.Unbox a) => Op a
derivVV = derivV * derivV

-- | Combined derivatives
-- > t = v + u
-- > r = v - u
-- > t^2 = v^2 + 2 v u + u^2
-- > r^2 = v^2 - 2 v u + u^2
derivT :: (Eq a, Field a, U.Unbox a) => Op a
derivT = (derivV ^+^ derivU) ^/ 2

derivR :: (Eq a, Field a, U.Unbox a) => Op a
derivR = (derivV ^-^ derivU) ^/ 2

derivTT :: (Eq a, Field a, U.Unbox a) => Op a
derivTT = derivT * derivT

derivRR :: (Eq a, Field a, U.Unbox a) => Op a
derivRR = derivR * derivR

-- | 1/r dr
r1derivR :: (Field a, RealFloat a, U.Unbox a) => Op a
r1derivR =
  ((delta - originR) * diagonal coordR1 * derivR + originR * derivRR)
    * symmetrizeR

-- t = v + u   v = (t + r) / 2
-- r = v - u   u = (t - r) / 2

-- | 1/r du
r1derivU :: (Field a, RealFloat a, U.Unbox a) => Op a
r1derivU =
  ((delta - originR) * diagonal coordR1 * derivU - originR * derivUU)
    * symmetrizeR

-- | 1/r dv
r1derivV :: (Field a, RealFloat a, U.Unbox a) => Op a
r1derivV =
  ((delta - originR) * diagonal coordR1 * derivV + originR * derivVV)
    * symmetrizeR

-- | 1/r^2
r2project :: (Field a, RealFloat a, U.Unbox a) => Op a
r2project =
  ( (delta - originR) * diagonal (gmap (^ 2) coordR1)
      + originR * ((1 / 2) *^ derivRR)
  )
    * (delta - valueAtOriginR)
    * symmetrizeR

-- | zero at origin
zproject :: (Eq a, Field a, U.Unbox a) => Op a
zproject = (delta - valueAtOriginR) * symmetrizeR

-- | Boundary locations
bndU :: (Container Vector a, Fractional a) => Op a
bndU = mkOp (bnd `atIndex`) delta1
  where
    delta1 (i, j) = if i == j then 1 else 0

bndV :: (Container Vector a, Fractional a) => Op a
bndV = mkOp delta1 (bnd `atIndex`)
  where
    delta1 (i, j) = if i == j then 1 else 0

metric :: (Container Vector a, Fractional a) => Op a
metric = mkOp (weights `atIndex`) (weights `atIndex`)

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

solve :: Field a => Op a -> Grid a -> Grid a
solve (Op a) (Grid x) =
  Grid $ head $ toColumns $ fromJust $ linearSolve a (asColumn x)

solveLS :: Field a => Op a -> Grid a -> Grid a
solveLS (Op a) (Grid x) =
  Grid $ head $ toColumns $ linearSolveLS a (asColumn x)

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

newtype MaxAbs a = MaxAbsInternal a
  deriving (Eq, Ord, Read, Show)

maxabs :: Num a => a -> MaxAbs a
maxabs x = MaxAbsInternal (abs x)

getMaxAbs :: MaxAbs a -> a
getMaxAbs (MaxAbsInternal x) = x

instance Ord a => Semigroup (MaxAbs a) where
  MaxAbsInternal x <> MaxAbsInternal y = MaxAbsInternal $ max x y

instance (Num a, Ord a) => Monoid (MaxAbs a) where
  mempty = MaxAbsInternal 0


module Scheme.Eval.Primitives (ioPrimitives, basicPrimitives) where

import Control.Monad
import Data.Maybe
import Data.Complex as C
import Data.Ratio
import Control.Monad.Error
import Scheme.Types
import System.IO

import Scheme.Eval.DynamicTypes

prim = undefined

convert typ wrap = expect typ >=> return . wrap

foldlOnType typ f = foldM (\acc e -> if check typ e
    then return (acc `f` e)
    else errorNeedType typ (show e))
foldlOnType1 typ f initPrep l = do expect typ (head l)
                                   foldlOnType typ f (initPrep $ head l) (init l)

basicPrimitives :: [(String, [LispVal] -> ThrowsErrorIO LispVal)]
basicPrimitives =
    let toBool = return . Bool
        dol = dottedOrPlainListType
    in
    -- number type checking
    [ ("number?",
        oneArgOnly >=> toBool . isNum)
    , ("real?",
        oneArgOnly >=> toBool . checkNum isDouble)
    , ("rational?",
        oneArgOnly >=> toBool . checkNum isRational)
    , ("integer?",
        oneArgOnly >=> toBool . checkNum isInt)
    , ("complex?",
        oneArgOnly >=> toBool . checkNum isComplex)

    -- only nice rationals and integers are exact
    , ("exact?",
        oneArgOnly >=> \n -> toBool $ any ($n) [checkNum isRational, checkNum isInt])

    -- arithmetic, preserve exactness in this order: integer, rational, float, complex
    , ("+",
        minArgs 2 >=> expect (listOf numType) . List >=> return . Num . sum)
    , ("-",
        minArgs 2 >=> expect (listOf numType) . List >=> return . Num . foldl1 (-))
    , ("*",
        minArgs 2 >=> expect (listOf numType) . List >=> return . Num . product)
    , ("/",
        onlyArgs 2 >=> \[_a, _b] -> do a <- expect numType _a
                                       b <- expect numType _b
                                       return . Num $ a / b)
    , ("negate",
        oneArgOnly >=> expect numType >=> return . Num . negate)
    , ("signum",
        oneArgOnly >=> expect numType >=> return . Num . signum)
    , ("abs",
        oneArgOnly >=> expect numType >=> return . Num . abs)
    , ("recip",
        oneArgOnly >=> expect numType >=> return . Num . recip)

    -- integer arithmetic functions
    , ("mod",
        minArgs 2 >=> foldlOnType1 numIntType (\(Num (Int a)) (Num (Int b))
                                              -> Num . Int $ a `mod` b) id)
    , ("quotient",
        minArgs 2 >=> foldlOnType1 numIntType (\(Num (Int a)) (Num (Int b))
                                              -> Num . Int $ a `quot` b) id)
    , ("remainder",
        minArgs 2 >=> foldlOnType1 numIntType (\(Num (Int a)) (Num (Int b))
                                              -> Num . Int $ a `rem` b) id)

    -- boolean
    , ("boolean?", oneArgOnly >=> toBool . isBool)
    , ("&&",
        minArgs 1 >=> liftM Bool . foldlOnType boolType (\a b -> a && fromJust (getBool b)) True)
    , ("||",
        minArgs 1 >=> liftM Bool . foldlOnType boolType (\a b -> a || fromJust (getBool b)) False)

    -- list processing !
    , ("list?",
        oneArgOnly >=> \l -> toBool $ any ($l) [isList, isDottedList, isVector])
    , ("dottedlist?",
        oneArgOnly >=> toBool . isDottedList)

    , ("car",
        oneArgOnly >=> expect dol >=> return .
                        (return . either id fst >=> safeHead)
                   >=> maybe (throwError $ Default "head-less list given to car") return)

    , ("cdr",
        oneArgOnly >=> expect dol >=> return .
                        --extract the dotted/normal list
                        either (liftM List . safeTail)
                            -- handle two cases for dotted lists
                            (\(xs, dot) -> Just $ if null xs then dot else DottedList (tail xs) dot)

                    -- handle bad tail
                   >=> maybe (throwError $ Default "tail-less list given to cdr") return)

    , ("cons", onlyArgs 2 >=> (\(a:b:[]) ->
                maybe (return $ List [a,b]) return
                    (get dol b >>= either (Just . List . (a:)) (\(xs, dot) -> Just $ DottedList (a:xs) dot))
                ))

    , ("vector?",
        oneArgOnly >=> toBool . isVector)
    , ("vector-ref", prim)

    , ("vector-set", prim)
    , ("list->vector",
        oneArgOnly >=> prim)
    , ("vector->list",
        oneArgOnly >=> prim)

    -- symbol processing
    , ("symbol?",
        oneArgOnly >=> toBool . isAtom)

    , ("string->symbol",
        oneArgOnly >=> convert stringType Atom)
    , ("symbol->string",
        oneArgOnly >=> convert atomType String)

    -- string processing
    , ("string?",
        oneArgOnly >=> toBool . isString)
    , ("string-ref",
        onlyArgs 2 >=> \([lispstr, lispix]) -> do
            str <- expect stringType lispstr
            ix  <- expect numIntType lispix
            char           <- maybe (throwError $ BadExpr "index out of bounds" lispix) return (safeIndex (fromIntegral ix) str)
            return . Character $ char)

    -- string construction
    , ("string->list",
        oneArgOnly >=> convert stringType (List . map Character))
    , ("list->string",
        convert (listOf charType) String . List)

    -- characters
    , ("character?",
        oneArgOnly >=> toBool . isCharacter)
    , ("character=", minArgs 1 >=> liftM (Bool . fst) . foldlOnType charType (\(res, prev) c -> (isNothing prev || getCharacter c == prev, getCharacter c)) (True, Nothing))

    -- equivalence
    , ("eq?",    onlyArgs 2 >=> return . eqv)
    , ("eqv?",   onlyArgs 2 >=> return . eqv)
    , ("=",     onlyArgs 2 >=> return . eqv)
    , ("equal?", onlyArgs 2 >=> return . equal)
    , ("==",    onlyArgs 2 >=> return . equal)
    ]

ioPrimitives =
    [ ("open-input-file", oneArgOnly >=> expect stringType >=> makePort ReadMode)
    , ("open-output-file", oneArgOnly >=> expect stringType >=> makePort WriteMode)
    , ("close-port", closePort)
    , ("close-input-port", closePort)
    , ("close-output-port", closePort)
    , ("read-contents", oneArgOnly >=> expect portType >=> liftM String . liftIO . hGetContents)
    , ("write", maxArgs 2 >=> minArgs 1 >=> \l -> case l of
        [obj] -> safeLiftIO (print obj) >> lispNull
        [obj, p] -> expect portType p >>= \port -> safeLiftIO (hPrint port (show obj)) >> lispNull)
    ]

closePort = oneArgOnly >=> liftIO . hClose . fromJust . getPort >=> const lispNull

makePort mode = liftM Port . safeLiftIO . flip openFile mode

eqv [a,b]
    | a == List [] && a == b = Bool True
    | all (\x -> check (dottedListType `orType` listType `orType` vectorType) x) [a,b] = Bool False
    | otherwise = equal [a,b]

equal [a,b] = Bool $ a == b


{-# LANGUAGE FlexibleContexts #-}
-----------------------------------------------------------------------------
-- |
-- Module      :  Numeric.Statistics.ICA
-- Copyright   :  (c) A. V. H. McPhail 2010, 2014
-- License     :  BSD3
--
-- Maintainer  :  haskell.vivian.mcphail <at> gmail <dot> com
-- Stability   :  provisional
-- Portability :  portable
--
-- Independent Components Analysis
--
--  implements the FastICA algorithm found in:
--
-- * Aapo Hyvärinen and Erkki Oja,
--   Independent Component Analysis: Algorithms and Applications,
--   /Neural Networks/, 13(4-5):411-430, 2000
--
--   <http://www.google.com/url?sa=t&source=web&cd=2&ved=0CBgQFjAB&url=http%3A%2F%2Fciteseerx.ist.psu.edu%2Fviewdoc%2Fdownload%3Fdoi%3D10.1.1.79.7003%26rep%3Drep1%26type%3Dpdf&ei=RQozTJb6L4_fcbCV6cMD&usg=AFQjCNGClLIB9MAvbrEj45SyUx9cYubLyA&sig2=hg5Wnfy3dLPkoIc1hqSfjg>
--
-----------------------------------------------------------------------------

module Numeric.Statistics.ICA (
                               sigmoid, sigmoid',
                               demean, whiten,
                               ica, icaDefaults,
                               --
                               NormType(..)
                          ) where

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

import qualified Data.Array.IArray as I

import Prelude hiding ((<>))

import Numeric.LinearAlgebra

import qualified Data.Vector.Generic as GV

import Numeric.GSL.Statistics

import Numeric.Statistics

import System.Random

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

data NormType = NormZero | NormOne | NormTwo | NormInf

pnorm :: Normed (Vector a) => NormType -> Vector a -> R
pnorm NormZero = norm_0
pnorm NormOne  = norm_1
pnorm NormTwo  = norm_2
pnorm NormInf  = norm_Inf

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

-- | sigmoid transfer function
sigmoid :: Double -> Double
sigmoid u = u * exp((-u**2)/2)

-- | derivative of sigmoid transfer function
sigmoid' :: Double -> Double
sigmoid' u = -u**2 * exp((-u**2)/2)

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

-- preprocessing:
--   demean
--   whiten
--        eigenvalue decomposition of covariance matrix E{xx^T} = EDE^T
--                   E orthogonal matrix of eigenvectors
--                   D diagonal matrix of eigenvalues, D = diag(d_1,...,d_n)
--                   x_white = ED^{-1/2}E^Tx
--                   D^{-1/2} = diag{d_1^{-1/2},...}
--

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

-- | remove the mean from data
demean :: I.Array Int (Vector Double)                  -- ^ the data
       -> (I.Array Int (Vector Double),Vector Double)  -- ^ (demeaned data,mean)
demean d = let u = I.elems $ fmap mean d
               d' = I.listArray (I.bounds d) (zipWith (-) (I.elems d) (map scalar u))
               u' = fromList u
               in (d',u')

-- | whiten data
whiten :: I.Array Int (Vector Double)                 -- ^ the data
       -> Double                                      -- ^ eigenvalue threshold
       -> (I.Array Int (Vector Double),Matrix Double) -- ^ (whitened data,transform)
whiten d q = let cv = covarianceMatrix d
                 (val',vec') = eigSH $ trustSym cv  -- the covariance matrix is real symmetric
                 val = toList val'
                 vec = toColumns vec'
                 v' = zip val vec
                 v = filter ((> q) . fst) v'        -- keep only eigens > than parameter
                 (dd',e') = unzip v
                 dd = diag $ (** (-0.5)) $ fromList dd'  -- square root of eigenvalues diagonalised
                 e = fromColumns e'
                 x = fromRows $ I.elems d
                 t = e <> dd <> tr' e          -- the actual mathematics
                 x' = t <> x                     -- the actual mathematics
                 d' = I.listArray (I.bounds d) (toRows x')
             in (d',t)

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

-- assuming that a weight vector is a row

-- algorithm:
-- 1  initial random weight vectors w_i
-- 2  w_i^+ = E{xg(w^Tx)} - E{g'(w^Tx)}w     (newton phase)
-- 3  W = W = (WW^T)^{-1/2)W                 (decorrelation) W = ( ..., w_i, ...)^T
--                                           WW^T = FDF^T (eigenvalue decomposition)
-- 4  w_i = w^+/norm(w^+)                    (normalisation) (almost any norm but not Frobenius)
-- 5  if not converged (dot w w^+ ~ 1 implies convergence) go to step 2
--
-- in matrix form, 2 becomes:
--    W^+ = W + (diag a_i)[(diag b_i) + E{g(y)y^T}]W
--
--      where
--            y = Wx
--            b_i = -E{y_ig(y_i)}
--            a_i = -1/(b_i-E{g'(y_i)})
--
--    g(u) = tanh(au) 0<=a<=2, often a = 1
--    g(u) = u exp(-u^2/2)

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

unconcat 0 _ _  = []
unconcat r c xs = [take c xs] ++ unconcat (r-1) c (drop c xs)

random_vector :: Int -> (Int,Int) -> Matrix Double
random_vector s (r,c) = fromLists $ unconcat r c $ randomRs (-1,1) (mkStdGen s)

-- g g' w x -> w'
update :: (Double -> Double) -> (Double -> Double) -> Matrix Double -> Matrix Double -> Matrix Double
update g g' w x = let y = w <> x
                      ys = toRows y
                      bis = map (\y' -> - mean (y' * (GV.map g y'))) ys
                      ais = zipWith (\b y' -> -1 / (b - mean (GV.map g y'))) bis ys
                      r = rows y
                      ix = ((1,1),(r,r))
                      cov = fromArray2D $ I.listArray ix $ map (\(m,n) -> covariance (GV.map g' (ys!!(m-1))) (ys!!(n-1))) $ I.range ix
                  in w + (diag $ fromList ais) <> ((diag $ fromList bis) + cov) <> w

decorrelate :: Matrix Double -> Matrix Double
decorrelate m = let (d',v') = eig m
                    d = fst $ fromComplex d'
                    v = fst $ fromComplex v'
                in v <> (diag (d ** (-0.5))) <> tr' v <> m
{-decorrelate n t w = let w' = w / (scalar $ sqrt $ pnorm n (w <> tr' w))
                    in decorrelate' t w w'
    where decorrelate' t' m m'
              | converged t' m m' = m'
              | otherwise         = decorrelate' t' m' ((scale 1.5 m') - (scale 0.5 (m' <> tr' m' <> m')))
-}

normalise :: NormType -> Matrix Double -> Matrix Double
normalise t m = fromRows $ map (\v -> v / (scalar $ pnorm t v)) (toRows m)

converged :: Double -> Matrix Double -> Matrix Double -> Bool
converged t m m' = let d' = map ((-) 1) $ zipWith dot (toRows m) (toRows m')
                   in maximum d' <= t

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

ica' :: (Double -> Double)          -- ^ transfer function (tanh,u exp(u^2/2), etc...)
     -> (Double -> Double)          -- ^ derivative of transfer function
     -> NormType                    -- ^ type of normalisation: Infinity, PNorm1, PNorm2
     -> Double                      -- ^ convergence tolerance for feature vectors
     -> Matrix Double               -- ^ weight matrix
     -> [Matrix Double]             -- ^ input data in chunks
     -> Matrix Double               -- ^ ica transform (weight matrix)
ica' _ _  _ _ _ []     = error "no sample data"
ica' g g' n t w (x:xs) = let w' = normalise n $ decorrelate $ update g g' w x
                             in if converged t w w'
                                then w'
                                else ica' g g' n t w' (xs ++ [x])

-- | perform an ICA transform
ica :: Int                         -- ^ random seed
    -> (Double -> Double)          -- ^ transfer function (tanh,u exp(u^2/2), etc...)
    -> (Double -> Double)          -- ^ derivative of transfer function
    -> NormType                    -- ^ type of normalisation: Infinity, PNorm1, PNorm2
    -> Double                      -- ^ convergence tolerance for feature vectors
--    -> Int                         -- ^ output dimensions
    -> Int                         -- ^ sampling size (must be smaller than length of data)
    -> I.Array Int (Vector Double) -- ^ data
    -> (I.Array Int (Vector Double),Matrix Double) -- ^ transformed data, ica transform
ica r g g' n t s a = let j = I.rangeSize $ I.bounds a
                         w = random_vector r (j,j)
                         x' = fromRows $ I.elems a
                         -- next line is BAD if distribution not stationary
                         x = concat $ toBlocksEvery j s x'
                         w' = ica' g g' n t w x
                         y = w' <> x'
                     in (I.listArray (1,1) $ toRows y,w')

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

-- | ICA with default values: no dimension reduction, euclidean norms, 16 sample groups, sigmoid
icaDefaults :: Int                         -- ^ random seed
            -> I.Array Int (Vector Double) -- ^ data
            -> (I.Array Int (Vector Double),Matrix Double) -- ^ transformed data, ica transform
icaDefaults r a = let c = I.rangeSize $ I.bounds a
                      s = (GV.length $ (a I.! 1)) `div` 16
                  in ica r sigmoid sigmoid' NormInf 0.0000001 s a

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


-- |
-- Module    : Numeric.ApproxEq
-- Copyright : (c) 2011 Aleksey Khudyakov, Bryan O'Sullivan
-- License   : BSD3
--
-- Maintainer  : Aleksey Khudyakov <alexey.skladnoy@gmail.com>
-- Stability   : experimental
-- Portability : portable
--
-- Different implementations of approximate equality for floating
-- point values. There are multiple ways to implement approximate
-- equality. They have different semantics and it's up to programmer
-- to choose right one.
module Numeric.ApproxEq (
    eqRelative
  , eqRelCompl
  , eqAbsolute
  , within
  ) where

import Control.Monad.ST         (runST)

import Data.Complex             (Complex(..))
import Data.Primitive.ByteArray (newByteArray, readByteArray, writeByteArray)
import Data.Int                 (Int64)

-- | Relative difference between two numbers are less than predefined
--   value.  For example 1 is approximately equal to 1.0001 with 1e-4
--   precision. Same is true for 10000 and 10001.
--
--   This method of camparison doesn't work for numbers which are
--   approximately 0. 'eqAbsolute' should be used instead.
eqRelative :: (Fractional a, Ord a)
           => a            -- ^ Relative precision
           -> a -> a -> Bool
eqRelative eps a b = abs (a - b) <= abs (eps * a)
{-# INLINE eqRelative #-}

-- | Relative equality for comlex numbers.
eqRelCompl :: (RealFloat a, Ord a)
           => a                 -- ^ Relative precision
           -> Complex a -> Complex a -> Bool
eqRelCompl eps a b = d <= eps * r
  where
    (d :+ _) = abs (a - b)
    (r :+ _) = abs a
{-# INLINE eqRelCompl #-}

-- | Difference between values is less than specified precision.
eqAbsolute :: (Num a, Ord a)
           => a                 -- ^ Absolute precision
           -> a -> a -> Bool
eqAbsolute eps a b = abs (a - b) <= eps
{-# INLINE eqAbsolute #-}

-- | Compare two 'Double' values for approximate equality, using
-- Dawson's method.
--
-- The required accuracy is specified in ULPs (units of least
-- precision).  If the two numbers differ by the given number of ULPs
-- or less, this function returns @True@.
--
-- Algorithm is based on Bruce Dawson's \"Comparing floating point
-- numbers\":
-- <http://www.cygnus-software.com/papers/comparingfloats/comparingfloats.htm>
within :: Int                   -- ^ Number of ULPs of accuracy desired.
       -> Double -> Double -> Bool
within ulps a b = runST $ do
  buf <- newByteArray 8
  ai0 <- writeByteArray buf 0 a >> readByteArray buf 0
  bi0 <- writeByteArray buf 0 b >> readByteArray buf 0
  let big  = 0x8000000000000000 :: Int64
      ai | ai0 < 0   = big - ai0
         | otherwise = ai0
      bi | bi0 < 0   = big - bi0
         | otherwise = bi0
  return $! abs (ai - bi) <= fromIntegral ulps


{-# LANGUAGE FlexibleContexts #-}
-- |
-- Module    : Statistics.Test.KolmogorovSmirnov
-- Copyright : (c) 2011 Aleksey Khudyakov
-- License   : BSD3
--
-- Maintainer  : bos@serpentine.com
-- Stability   : experimental
-- Portability : portable
--
-- Kolmogov-Smirnov tests are non-parametric tests for assessing
-- whether given sample could be described by distribution or whether
-- two samples have the same distribution. It's only applicable to
-- continuous distributions.
module Statistics.Test.KolmogorovSmirnov (
    -- * Kolmogorov-Smirnov test
    kolmogorovSmirnovTest
  , kolmogorovSmirnovTestCdf
  , kolmogorovSmirnovTest2
    -- * Evaluate statistics
  , kolmogorovSmirnovCdfD
  , kolmogorovSmirnovD
  , kolmogorovSmirnov2D
    -- * Probablities
  , kolmogorovSmirnovProbability
    -- * References
    -- $references
  , module Statistics.Test.Types
  ) where

import Control.Monad (when)
import Prelude hiding (exponent, sum)
import Statistics.Distribution (Distribution(..))
import Statistics.Function (gsort, unsafeModify)
import Statistics.Matrix (center, for, fromVector)
import qualified Statistics.Matrix as Mat
import Statistics.Test.Types
import Statistics.Types (mkPValue)
import qualified Data.Vector          as V
import qualified Data.Vector.Storable as S
import qualified Data.Vector.Unboxed  as U
import qualified Data.Vector.Generic  as G
import           Data.Vector.Generic    ((!))
import qualified Data.Vector.Unboxed.Mutable as M

----------------------------------------------------------------
-- Test
----------------------------------------------------------------

-- | Check that sample could be described by distribution. Returns
--   @Nothing@ is sample is empty
--
--   This test uses Marsaglia-Tsang-Wang exact algorithm for
--   calculation of p-value.
kolmogorovSmirnovTest :: (Distribution d, G.Vector v Double)
                      => d        -- ^ Distribution
                      -> v Double -- ^ Data sample
                      -> Maybe (Test ())
{-# INLINE kolmogorovSmirnovTest #-}
kolmogorovSmirnovTest d
  = kolmogorovSmirnovTestCdf (cumulative d)

-- | Variant of 'kolmogorovSmirnovTest' which uses CDF in form of
--   function.
kolmogorovSmirnovTestCdf :: (G.Vector v Double)
                         => (Double -> Double) -- ^ CDF of distribution
                         -> v Double           -- ^ Data sample
                         -> Maybe (Test ())
{-# INLINE kolmogorovSmirnovTestCdf #-}
kolmogorovSmirnovTestCdf cdf sample
  | G.null sample = Nothing
  | otherwise     = Just Test
      { testSignificance = mkPValue $ 1 - prob
      , testStatistics   = d
      , testDistribution = ()
      }
  where
    d    = kolmogorovSmirnovCdfD cdf sample
    prob = kolmogorovSmirnovProbability (G.length sample) d

-- | Two sample Kolmogorov-Smirnov test. It tests whether two data
--   samples could be described by the same distribution without
--   making any assumptions about it. If either of samples is empty
--   returns Nothing.
--
--   This test uses approximate formula for computing p-value.
kolmogorovSmirnovTest2 :: (G.Vector v Double)
                       => v Double -- ^ Sample 1
                       -> v Double -- ^ Sample 2
                       -> Maybe (Test ())
kolmogorovSmirnovTest2 xs1 xs2
  | G.null xs1 || G.null xs2 = Nothing
  | otherwise                = Just Test
      { testSignificance = mkPValue $ 1 - prob d
      , testStatistics   = d
      , testDistribution = ()
      }
  where
    d    = kolmogorovSmirnov2D xs1 xs2
         * (en + 0.12 + 0.11/en)
    -- Effective number of data points
    n1   = fromIntegral (G.length xs1)
    n2   = fromIntegral (G.length xs2)
    en   = sqrt $ n1 * n2 / (n1 + n2)
    --
    prob z
      | z <  0    = error "kolmogorovSmirnov2D: internal error"
      | z == 0    = 0
      | z <  1.18 = let y = exp( -1.23370055013616983 / (z*z) )
                    in  2.25675833419102515 * sqrt( -log y ) * (y + y**9 + y**25 + y**49)
      | otherwise = let x = exp(-2 * z * z)
                    in  1 - 2*(x - x**4 + x**9)
{-# INLINABLE  kolmogorovSmirnovTest2 #-}
{-# SPECIALIZE kolmogorovSmirnovTest2 :: U.Vector Double -> U.Vector Double -> Maybe (Test ()) #-}
{-# SPECIALIZE kolmogorovSmirnovTest2 :: V.Vector Double -> V.Vector Double -> Maybe (Test ()) #-}
{-# SPECIALIZE kolmogorovSmirnovTest2 :: S.Vector Double -> S.Vector Double -> Maybe (Test ()) #-}
-- FIXME: Find source for approximation for D

----------------------------------------------------------------
-- Kolmogorov's statistic
----------------------------------------------------------------

-- | Calculate Kolmogorov's statistic /D/ for given cumulative
--   distribution function (CDF) and data sample. If sample is empty
--   returns 0.
kolmogorovSmirnovCdfD :: G.Vector v Double
                      => (Double -> Double) -- ^ CDF function
                      -> v Double           -- ^ Sample
                      -> Double
kolmogorovSmirnovCdfD cdf sample
  | G.null sample = 0
  | otherwise     = G.maximum
                  $ G.zipWith3 (\p a b -> abs (p-a) `max` abs (p-b))
                    ps steps (G.tail steps)
  where
    xs = gsort sample
    n  = G.length xs
    --
    ps    = G.map cdf xs
    steps = G.map (/ fromIntegral n)
          $ G.generate (n+1) fromIntegral
{-# INLINABLE  kolmogorovSmirnovCdfD #-}
{-# SPECIALIZE kolmogorovSmirnovCdfD :: (Double -> Double) -> U.Vector Double -> Double #-}
{-# SPECIALIZE kolmogorovSmirnovCdfD :: (Double -> Double) -> V.Vector Double -> Double #-}
{-# SPECIALIZE kolmogorovSmirnovCdfD :: (Double -> Double) -> S.Vector Double -> Double #-}

-- | Calculate Kolmogorov's statistic /D/ for given cumulative
--   distribution function (CDF) and data sample. If sample is empty
--   returns 0.
kolmogorovSmirnovD :: (Distribution d, G.Vector v Double)
                   => d         -- ^ Distribution
                   -> v Double  -- ^ Sample
                   -> Double
kolmogorovSmirnovD d = kolmogorovSmirnovCdfD (cumulative d)
{-# INLINE kolmogorovSmirnovD #-}

-- | Calculate Kolmogorov's statistic /D/ for two data samples. If
--   either of samples is empty returns 0.
kolmogorovSmirnov2D :: (G.Vector v Double)
                    => v Double   -- ^ First sample
                    -> v Double   -- ^ Second sample
                    -> Double
kolmogorovSmirnov2D sample1 sample2
  | G.null sample1 || G.null sample2 = 0
  | otherwise                        = worker 0 0 0
  where
    xs1 = gsort sample1
    xs2 = gsort sample2
    n1  = G.length xs1
    n2  = G.length xs2
    en1 = fromIntegral n1
    en2 = fromIntegral n2
    -- Find new index
    skip x i xs = go (i+1)
      where go n | n >= G.length xs = n
                 | xs ! n == x      = go (n+1)
                 | otherwise        = n
    -- Main loop
    worker d i1 i2
      | i1 >= n1 || i2 >= n2 = d
      | otherwise            = worker d' i1' i2'
      where
        d1  = xs1 ! i1
        d2  = xs2 ! i2
        i1' | d1 <= d2  = skip d1 i1 xs1
            | otherwise = i1
        i2' | d2 <= d1  = skip d2 i2 xs2
            | otherwise = i2
        d'  = max d (abs $ fromIntegral i1' / en1 - fromIntegral i2' / en2)
{-# INLINABLE  kolmogorovSmirnov2D #-}
{-# SPECIALIZE kolmogorovSmirnov2D :: U.Vector Double -> U.Vector Double -> Double #-}
{-# SPECIALIZE kolmogorovSmirnov2D :: V.Vector Double -> V.Vector Double -> Double #-}
{-# SPECIALIZE kolmogorovSmirnov2D :: S.Vector Double -> S.Vector Double -> Double #-}

-- | Calculate cumulative probability function for Kolmogorov's
--   distribution with /n/ parameters or probability of getting value
--   smaller than /d/ with n-elements sample.
--
--   It uses algorithm by Marsgalia et. al. and provide at least
--   7-digit accuracy.
kolmogorovSmirnovProbability :: Int    -- ^ Size of the sample
                             -> Double -- ^ D value
                             -> Double
kolmogorovSmirnovProbability n d
  -- Avoid potentially lengthy calculations for large N and D > 0.999
  | s > 7.24 || (s > 3.76 && n > 99) = 1 - 2 * exp( -(2.000071 + 0.331 / sqrt n' + 1.409 / n') * s)
  -- Exact computation
  | otherwise = fini $ KSMatrix 0 matrix `power` n
  where
    s  = n' * d * d
    n' = fromIntegral n

    size = 2*k - 1
    k    = floor (n' * d) + 1
    h    = fromIntegral k - n' * d
    -- Calculate initial matrix
    matrix =
      let m = U.create $ do
            mat <- M.new (size*size)
            -- Fill matrix with 0 and 1s
            for 0 size $ \row ->
              for 0 size $ \col -> do
                let val | row + 1 >= col = 1
                        | otherwise      = 0 :: Double
                M.write mat (row * size + col) val
            -- Correct left column/bottom row
            for 0 size $ \i -> do
              let delta = h ^^ (i + 1)
              unsafeModify mat (i    * size)         (subtract delta)
              unsafeModify mat (size * size - 1 - i) (subtract delta)
            -- Correct corner element if needed
            when (2*h > 1) $ do
              unsafeModify mat ((size - 1) * size) (+ ((2*h - 1) ^ size))
            -- Divide diagonals by factorial
            let divide g num
                  | num == size = return ()
                  | otherwise   = do for num size $ \i ->
                                       unsafeModify mat (i * (size + 1) - num) (/ g)
                                     divide (g * fromIntegral (num+2)) (num+1)
            divide 2 1
            return mat
      in fromVector size size m
    -- Last calculation
    fini (KSMatrix e m) = loop 1 (center m) e
      where
        loop i ss eQ
          | i  > n       = ss * 10 ^^ eQ
          | ss' < 1e-140 = loop (i+1) (ss' * 1e140) (eQ - 140)
          | otherwise    = loop (i+1)  ss'           eQ
          where ss' = ss * fromIntegral i / fromIntegral n

data KSMatrix = KSMatrix Int Mat.Matrix

multiply :: KSMatrix -> KSMatrix -> KSMatrix
multiply (KSMatrix e1 m1) (KSMatrix e2 m2) = KSMatrix (e1+e2) (Mat.multiply m1 m2)

power :: KSMatrix -> Int -> KSMatrix
power mat 1 = mat
power mat n = avoidOverflow res
  where
    mat2 = power mat (n `quot` 2)
    pow  = multiply mat2 mat2
    res | odd n     = multiply pow mat
        | otherwise = pow

avoidOverflow :: KSMatrix -> KSMatrix
avoidOverflow ksm@(KSMatrix e m)
  | center m > 1e140 = KSMatrix (e + 140) (Mat.map (* 1e-140) m)
  | otherwise        = ksm

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

-- $references
--
-- * G. Marsaglia, W. W. Tsang, J. Wang (2003) Evaluating Kolmogorov's
--   distribution, Journal of Statistical Software, American
--   Statistical Association, vol. 8(i18).


{-|
Module      : Distributions
Description : Definition of the Distributions used in Models and associated functions.
Copyright   : (c) Julian Kopka Larsen, 2015
Stability   : experimental
-}
module Distributions (
    lnormalInvWishart,
    lnormalInvWishartSS,
    lMixNormWish,
    delta,
    Expr,
    dlNormW,
    lNormW
) where

import Math (X,
             mlgamma,
             scatterMatrix,
             meanv,
             lnDeterminant)

import Numeric.LinearAlgebra (Matrix,
                              Vector,
                              (<>),
                              ident,
                              dim,
                              scale,
                              trans,
                              takeDiag,
                              outer)

import Numeric.LinearAlgebra.Data (konst, size)

import Partition (Partition, Component, groups, group)

import LikSpecLang

-- | Delta log Normal-Wishard likelihood for a mixture of components.
dlNormW :: X -> (Partition -> Partition -> Double)
dlNormW x = delta (lMixNormWish x)

-- | Log Normal-Wishard likelihood for a mixture of components.
lNormW :: X -> (Partition -> Double)
lNormW x = compile (lMixNormWish x)

-- | Log Normal-Wishard likelihood for a mixture of components. Same as 'lNormW' but defined without the "LikSpecLang" (LSL)
lMixNormWish ::X -> Expr Partition Double
lMixNormWish x = Mixture (Fun lnormalInvWishart) % selection x -- The '%' operator "pipes" expressions together

selection :: X -> Expr Partition [X]
selection x = Fun (groups x)

-- | Log Normal-Wishard likelihood for a single component
lnormalInvWishartSS :: Int -> Matrix Double -> Vector Double-> Double
lnormalInvWishartSS n s mu = - fromIntegral n * d * 2 * log pi
                      + mlgamma d (vn/2)
                      - mlgamma d (v0/2)
                      + (v0/2) * ldalpha0
                      - (vn/2) * ldalphaN
                      + (d/2) * ((log k0)-(log kn))
                    where   d = fromIntegral $ fst $ size s :: Double
                            v0 = 100 -- hyperparameter must be (v0 > d)
                            vn = v0 + fromIntegral n
                            alpha0 = ident $ fst $ size s -- hyperparameter for the shape of the covariance matrix
                            alphaN = alpha0 + s + scale ((k0*fromIntegral n)/(k0+kn)) ((mu-mu0) `outer` (mu-mu0))
                            ldalpha0 = lnDeterminant alpha0
                            ldalphaN = lnDeterminant alphaN
                            k0 = 0.2 -- hyperparameter
                            kn = k0 + fromIntegral n
                            mu0 = (konst 0.0 $ fst $ size s ):: Vector Double -- hyperparameter

-- | Log Normal-Wishard likelihood for a single component
lnormalInvWishart :: X -> Double
lnormalInvWishart x = - n * d * 2 * log pi
                      + mlgamma d (vn/2)
                      - mlgamma d (v0/2)
                      + (v0/2) * ldalpha0
                      - (vn/2) * ldalphaN
                      + (d/2) * ((log k0)-(log kn))
                    where   n = fromIntegral $ length x
                            d = fromIntegral $ dim $ head x
                            v0 = 100 -- hyperparameter must be (v0 > d)
                            vn = v0 + n
                            alpha0 = ident $ dim $ head x -- hyperparameter for the shape of the covariance matrix
                            alphaN = alpha0 + s + scale ((k0*n)/(k0+kn)) ((xm-mu0) `outer` (xm-mu0))
                            ldalpha0 = lnDeterminant alpha0
                            ldalphaN = lnDeterminant alphaN
                            k0 = 0.2 -- hyperparameter
                            kn = k0 + n
                            s = scatterMatrix x
                            mu0 = (konst 0.0 $ dim $ head x ):: Vector Double -- hyperparameter
                            xm = meanv x


module Filter (FilterOptions(..), parseOpts, lowPass, convolve) where

import           Control.Monad.State.Strict
import           Data.Coerce
import           Data.Complex
import qualified Data.Vector.Storable as VS
import qualified Data.Vector.Storable.Mutable as VSM
import           Foreign.C.Types
import           Foreign.ForeignPtr
import           Foreign.Ptr
import           Foreign.Storable
import           IQ
import           Options.Applicative
import           Pipes

data FilterOptions = FilterOptions
  { cutoff :: Int
  , transitionWidth :: Int }

parseOpts :: Parser FilterOptions
parseOpts = FilterOptions
         <$> option auto
             ( long "cutoff-frequency"
            <> short 'c'
            <> value 15000
            <> showDefault
            <> metavar "CUTOFF_FREQ"
            <> help "Cutoff frequency, value in Hz")
         <*> option auto
             ( long "transition-width"
            <> short 't'
            <> value 10000
            <> showDefault
            <> metavar "TRANS_WIDTH"
            <> help "Transition width, value in Hz")

calculateNumTaps :: Int -> Int -> Int
calculateNumTaps sr tw = round $ 70.0 * sr' / (22.0 * tw')
  where sr' = fromIntegral sr
        tw' = fromIntegral tw

hann :: Floating a => Int -> [a]
hann n = map (\x -> sin (pi * (fromIntegral x / n')) ** 2) vec
  where vec = [0..n-1]
        n' =  fromIntegral n

-- sin (pi * x) / x. The cutoff frequency controls the sampling frequency of ths
-- sinc function: Sf = 2 * Fc, where Fc is expressed as a fraction of the
-- original signal sampling rate.
sinc :: Floating a =>  Int -> Int -> Int -> [a]
sinc len sr fc = map ((fac *) . sinc') [-(len `div` 2)..(len `div` 2)]
  where sinc' x | x == 0 = 1
                | otherwise =  let n = fromIntegral x in sin (fac * n * pi) / (fac * n * pi)
        fac = 2 * (fromIntegral fc / fromIntegral sr)

lowPass :: Int -> FilterOptions -> VS.Vector Float
lowPass sr opts = VS.fromList x
  where nTaps = calculateNumTaps sr (transitionWidth opts)
        window = hann nTaps
        x = zipWith (*) window $ sinc nTaps sr (cutoff opts)

convolve :: VS.Vector Float -> Pipe IQ IQ IO ()
convolve kernel = do
  let klen = VS.length kernel
      bufSz = 1024
      mkCInt = CInt . fromIntegral

  inBuf <- lift $ VSM.replicate bufSz 0.0
  outBuf <- lift $ VSM.replicate bufSz 0.0
  stateBuf <- lift $ VSM.replicate klen (0.0 :+ 0.0 :: IQ)
  stateHeadPos <- lift mallocForeignPtr
  lift . withForeignPtr stateHeadPos $ flip poke 0

  forever $ do
    forM_ [0..bufSz-1] $ (\i -> await >>= lift . VSM.write inBuf i)

    lift $ VS.unsafeWith (coerce kernel) $ \kBuf ->
      VSM.unsafeWith (coerce stateBuf) $ \stateBuf ->
        VSM.unsafeWith (coerce inBuf) $ \inBuf ->
          VSM.unsafeWith (coerce outBuf) $ \outBuf ->
            withForeignPtr stateHeadPos $ \headPos ->
              conv_c (mkCInt klen) (mkCInt bufSz) headPos kBuf stateBuf inBuf outBuf

    forM_ [0..bufSz-1] $ lift . VSM.read outBuf >=> yield
{-# INLINE convolve #-}

foreign import ccall "convolve" conv_c :: CInt -> CInt -> Ptr CInt -> Ptr CFloat -> Ptr CFloat -> Ptr CFloat -> Ptr CFloat -> IO ()


{-# LANGUAGE ScopedTypeVariables #-}

module CSVMatrixParser
  (
    getData
  ) where

import Numeric.LinearAlgebra
import qualified Data.ByteString.Lazy as BL
import Data.Csv
import qualified Data.Vector as V

getCSV :: String -> IO (V.Vector [Double])
getCSV path = do
  csvData <- BL.readFile path
  case decode NoHeader csvData of
    Left err -> return V.empty
    Right v -> return v

splitIO :: [Double] -> ([Double], [Double])
splitIO list = splitAt 5 list

toVectorRow :: ([Double], [Double]) -> (Vector R, Vector R)
toVectorRow (xs,ys) = (vector xs, vector ys)

concatMats :: V.Vector (Vector R, Vector R) -> (Matrix R, Matrix R)
concatMats vec = let list = V.toList vec
                 in (fromRows (map fst list), fromRows (map snd list))

getData :: String -> IO (Matrix R, Matrix R)
getData path = getCSV path >>= return . concatMats . fmap (toVectorRow . splitIO)

concatTupleMats :: (Matrix R, Matrix R) -> (Matrix R, Matrix R) -> (Matrix R, Matrix R)
concatTupleMats (a,x) (b,y) = (a === b, x === y)


{-|
Module: MachineLearning.Clustering
Description: Clustering
Copyright: (c) Alexander Ignatyev, 2017
License: BSD-3
Stability: experimental
Portability: POSIX

Cluster Analysis a.k.a. Clustering - grouping data into coherent subsets.
-}

module MachineLearning.Clustering
(
  Cluster(..)
  , kmeans

  -- * Exported for testing purposes only.
  , kmeansIterM
)

where

import MachineLearning.Types (R, Vector, Matrix)
import Data.List (sortOn, groupBy, minimumBy)
import Control.Applicative ((<$>))
import Control.Monad (forM)
import qualified Control.Monad.Random as RndM
import qualified Data.Vector as V
import qualified Numeric.LinearAlgebra as LA
import MachineLearning.Random (sampleM)

-- | Cluster type (list of samples associtaed with the cluster).
type Cluster = V.Vector Vector

-- | Gets list if the nearest centroid to the sample.
nearestCentroidIndex :: V.Vector Vector  -- ^ list of cluster cetroids;
                     -> Vector            -- ^ sample;
                     -> Int               -- ^ index of the nearest centroid.
nearestCentroidIndex centroids v =
  let distances = V.map (\centroid -> LA.norm_2 (v-centroid)) centroids
  in V.minIndex distances

-- | Calculates cost associated with one claster.
calcClusterCost :: Cluster  -- ^ cluster;
                -> Vector   -- ^ cluster centroid;
                -> R        -- ^ cost value.
calcClusterCost cluster centroid = sum $ fmap (\sample -> LA.norm_2 $ sample-centroid) cluster

-- | Calculates cost value for all clusters.
calcCost :: V.Vector Cluster  -- ^ cluster list;
         -> V.Vector Vector   -- ^ list of cluster centroids;
         -> R                  -- ^ cost value.
calcCost clusters centroids = sum $ V.zipWith calcClusterCost clusters centroids

-- | Calculates centroid (middle point) of the given cluster.
getNewCentroid :: Cluster      -- ^ cluster;
               -> Vector       -- ^ centroid.
getNewCentroid cluster =
  let n = length cluster
      centroid = (sum cluster) / (fromIntegral n)
  in centroid

-- | Calculates new cluster centroids for each cluster.
moveCentroids :: V.Vector Cluster    -- ^ list of clusters;
              -> V.Vector Vector     -- ^ list of cluster centroids.
moveCentroids clusters = fmap getNewCentroid clusters

-- | Build cluster list from list of clusters indices.
buildClusterList :: V.Vector Vector   -- ^ list of samples;
                 -> V.Vector Int      -- ^ list of cluster indices (associated cluster index for each sample);
                 -> V.Vector Cluster  -- ^ list of clusters.
buildClusterList samples clusterIndicesList = V.fromList $ fmap getClusterSamples clusters''
  where clusters' = groupBy (\l r -> snd l == snd r) $ sortOn snd $ zip [0..] $ V.toList clusterIndicesList
        clusters'' = map (map fst) clusters'
        getClusterSamples clusterIndices = V.fromList $ fmap (samples V.!) clusterIndices

-- -- | Run K-Means algorithm once.
kmeansIter :: V.Vector Vector           -- ^ list of samples;
              -> Int                    -- ^ number of clusters (`K`);
              -> V.Vector Vector        -- ^ list of initial centroids;
              -> (V.Vector Cluster, [R])  -- ^ (list of clusters, cost values).
kmeansIter samples k initialCentroids =
  let iter centroids js =
        let clusterIndicesList = fmap (nearestCentroidIndex centroids) samples
            clusters = buildClusterList samples clusterIndicesList
            centroids' = moveCentroids clusters
            j = calcCost clusters centroids'
            diff = sum . fmap LA.norm_2 $ V.zipWith (-) centroids centroids'
        in if diff < 0.001 then (clusters, j:js)
           else iter centroids' (j:js)
  in iter initialCentroids []

-- | Run K-Means algorithm once inside Random Monad.
kmeansIterM :: RndM.RandomGen g =>
               V.Vector Vector  -- ^ list of samples;
               -> Int           -- ^ number of clusters (`K`);
               -> Int           -- ^ iteration number;
               -> RndM.Rand g (V.Vector Cluster, [R])  -- ^ (list of clusters, cost values) inside Random Monad.
kmeansIterM samples k _ = do
  centroids <- sampleM k samples
  return (kmeansIter samples k centroids)

-- | Clusters data using K-Means Algorithm inside Random Monad.
-- Runs K-Means algorithm `N` times, returns the clustering with smaller cost.
kmeans :: RndM.RandomGen g =>
           Int                     -- ^ number of K-Means Algorithm runs (`N`);
           -> Matrix                  -- ^ data to cluster;
           -> Int                     -- ^ desired number of clusters (`K`);
           -> RndM.Rand g (V.Vector Cluster)  -- ^ list of clusters inside Random Monad.
kmeans nIters x k = fst <$>
    (minimumBy (\(_, js1) (_, js2) -> compare (head js1) (head js2))) <$>
    forM [1..nIters] (kmeansIterM samples k)
  where samples = V.fromList $ LA.toRows x


{-# LANGUAGE FlexibleContexts, NamedFieldPuns #-}

module School.Unit.WeightDecay
( weightDecay ) where

import Conduit (ConduitM, mapMC)
import Numeric.LinearAlgebra (Container, Matrix, Vector,
                              cmap, scale, sumElements)
import School.Types.Slinky (Slinky(..), slinkyAppend)
import School.Unit.CostFunction (CostFunction(..), defSetupCost)
import School.Unit.CostParams (CostParams(..))
import School.Unit.UnitActivation (UnitActivation(..))
import School.Unit.UnitForward (ForwardStack)
import School.Unit.UnitGradient (UnitGradient(..))

square :: (Num a) => a -> a
square a = a * a

weight :: (Container Vector a, Num a)
       => a -> Matrix a -> a
weight coeff input =
  (*coeff) . sumElements $ cmap square input

weightDeriv :: (Container Vector a, Num a)
            => a -> Matrix a -> UnitGradient a
weightDeriv coeff input =
  BatchGradient . scale coeff $
    cmap (*2) input

errorMsg :: String
errorMsg = "Weight decay expects batch activation and no cost params"

compute :: (Container Vector a, Num a)
        => a
        -> UnitActivation a
        -> Slinky CostParams
        -> Either String a
compute coeff
        (BatchActivation input)
        (SNode NoCostParams _) =
  Right $ weight coeff input
compute coeff
        (BatchActivation input)
        SNil =
  Right $ weight coeff input
compute _ _ _ = Left errorMsg

deriv :: (Container Vector a, Num a)
      => a
      -> UnitActivation a
      -> Slinky CostParams
      -> Either String (UnitGradient a)
deriv coeff
      (BatchActivation input)
      (SNode NoCostParams _) =
  Right $ weightDeriv coeff input
deriv coeff
      (BatchActivation input)
      SNil =
  Right $ weightDeriv coeff input
deriv _ _ _ = Left errorMsg

prepare :: (Monad m)
        => ConduitM (ForwardStack a)
                    (ForwardStack a)
                    m
                    ()
prepare = mapMC $ \(activations, cParams) ->
  return (activations, slinkyAppend NoCostParams cParams)

weightDecay :: (Container Vector a, Num a, Monad m)
            => a
            -> CostFunction a m
weightDecay coeff =
  CostFunction { computeCost = compute coeff
               , derivCost = deriv coeff
               , prepareCost = prepare
               , setupCost = defSetupCost
               }


{-# LANGUAGE BangPatterns #-}

module Statistics.Matrix.Fast (
    multiply,
    norm,
    multiplyV,
    transpose
    ) where

import Prelude hiding (exponent, map)
import Control.Monad.ST
import qualified Data.Vector.Unboxed as U
import qualified Data.Vector.Unboxed.Mutable as UM

import Statistics.Matrix (row)
import Statistics.Matrix.Function
import Statistics.Matrix.Types
import Statistics.Matrix.Mutable  (unsafeNew,unsafeWrite,unsafeFreeze)

-- | Matrix-matrix multiplication in a more imperative fashion. Matrices must be of compatible
-- sizes (/note: not checked/). Faster but less accurate than Statistics.Matrix.multiply
multiply :: Matrix -> Matrix -> Matrix
multiply m1@(Matrix r1 _ _) m2@(Matrix _ c2 _) = runST $ do
  m3 <- unsafeNew r1 c2
  for 0 c2 $ \j -> do
    for 0 r1 $ \i -> do
      let
        z = accum i m1 j m2
      unsafeWrite m3 i j z
  unsafeFreeze m3

accum :: Int -> Matrix -> Int -> Matrix -> Double
accum ithrow (Matrix _ c1 v1) jthcol (Matrix _ c2 v2) = sub 0 0
  where sub !acc !ij | ij == c1 = acc
                     | otherwise = sub ( valRow*valCol + acc ) (ij+1)
                                   where
                                    valRow = U.unsafeIndex v1 (ithrow*c1 + ij)
                                    valCol = U.unsafeIndex v2 (ij*c2+jthcol)

-- | Matrix-vector multiplication, with better performances but not as accurate as
-- Statistics.Matrix.multiplyV
multiplyV :: Matrix -> Vector -> Vector
multiplyV m v
  | cols m == c = U.generate (rows m) (U.sum . U.zipWith (*) v . row m)
  | otherwise   = error $ "matrix/vector unconformable " ++ show (cols m,c)
  where c = U.length v

-- | Norm of a vector. Faster but less accurate than Statistics.Matrix.norm
norm :: Vector -> Double
norm = sqrt . U.sum . U.map square

transpose :: Matrix -> Matrix
transpose (Matrix r0 c0 v0)
  = Matrix c0 r0 $ runST $ do
    vec <- UM.unsafeNew (r0*c0)
    for 0 r0 $ \i -> do
      UM.unsafeWrite vec (i + i * c0) $ v0 `U.unsafeIndex` (i + i * c0)
      for (i+1) c0 $ \j -> do
        let tmp = v0 `U.unsafeIndex` (j + i * c0)
            tmp2 = v0 `U.unsafeIndex` (i + j * c0)
        UM.unsafeWrite vec (j + i * c0) tmp2
        UM.unsafeWrite vec (i + j * c0) tmp
    U.unsafeFreeze vec


module Deanie.Measure (
    measure

  -- * queries

  , integrate
  , expectation
  , variance
  , mgf
  , cgf
  , cdf
  ) where

import Control.Monad
import Data.List (foldl')
import Deanie.Language
import Control.Foldl (Fold)
import Numeric.Integration.TanhSinh
import Numeric.SpecFunctions

newtype Measure a = Measure ((a -> Double) -> Double)

integrate :: (a -> Double) -> Measure a -> Double
integrate f (Measure nu) = nu f

expectation :: Measure Double -> Double
expectation = integrate id

variance :: Measure Double -> Double
variance nu = integrate (^ 2) nu - expectation nu ^ 2

mgf :: Measure Double -> Double -> Double
mgf mu t = integrate (\x -> exp (t * x)) mu

cgf :: Measure Double -> Double -> Double
cgf mu = log . mgf mu

cdf :: Measure Double -> Double -> Double
cdf nu x = integrate (negativeInfinity `to` x) nu where
  negativeInfinity :: Double
  negativeInfinity = negate (1 / 0)

  to :: (Num a, Ord a) => a -> a -> a -> a
  to a b x
    | x >= a && x <= b = 1
    | otherwise        = 0

instance Functor Measure where
  fmap f nu = Measure $ \g ->
    integrate (g . f) nu

instance Applicative Measure where
  pure x = Measure (\f -> f x)
  Measure h <*> Measure g = Measure $ \f ->
    h (\k -> g (f . k))

instance Monad Measure where
  return x  = Measure (\f -> f x)
  rho >>= g = Measure $ \f ->
    integrate (\nu -> integrate f (g nu)) rho

fromMassFunction :: Foldable f => (a -> Double) -> f a -> Measure a
fromMassFunction f support = Measure $ \g ->
  foldl' (\acc x -> acc + f x * g x) 0 support

fromDensityFunction :: (Double -> Double) -> Measure Double
fromDensityFunction d = Measure $ \f ->
    quadratureTanhSinh (\x -> f x * d x)
  where
    quadratureTanhSinh = result . last . everywhere trap

mbernoulli :: Double -> Measure Bool
mbernoulli p = fromMassFunction (pmf p) [False, True] where
  pmf p x
    | p < 0 || p > 1 = 0
    | otherwise      = if x then p else 1 - p

mbeta :: Double -> Double -> Measure Double
mbeta a b = fromDensityFunction (density a b) where
  density a b p
    | p < 0 || p > 1 = 0
    | otherwise      = 1 / exp (logBeta a b) * p ** (a - 1) * (1 - p) ** (b - 1)

mgamma :: Double -> Double -> Measure Double
mgamma a b = fromDensityFunction (density a b) where
  density a b x
    | a < 0 || b < 0 = 0
    | otherwise  =
       b ** a / exp (logGamma a) * x ** (a - 1) * exp (negate (b * x))

mgaussian :: Double -> Double -> Measure Double
mgaussian m s = fromDensityFunction (density m s) where
  density m s x
    | s <= 0    = 0
    | otherwise =
        1 / (s * sqrt (2 * pi)) *
          exp (negate ((x - m) ^^ 2) / (2 * (s ^^ 2)))

measure :: Program a -> Measure a
measure = iterM $ \case
    ProgramF (InL term) -> evalAlg term
    ProgramF (InR term) -> join (runAp measure term)
  where
    evalAlg = \case
      BernoulliF p k  -> mbernoulli p >>= k
      BetaF a b k     -> mbeta a b >>= k
      GammaF a b k    -> mgamma a b >>= k
      GaussianF m s k -> mgaussian m s >>= k


{-# LANGUAGE BangPatterns #-}
module LINreg2 (mmlLinearRegLP, diagl, rowM, colM,(<+>), (<->), foldl'Rnf
                , mkZ,realignedNpieces,lineFun, mmlCorr-- ,vmean,vvar,vstDiv,vvarM,vunStdNorm,vstNorm, corrRoundCentre
                , standardizedFit,relFit, corr,nCorrV,logg
                ) where

import Numeric.LinearAlgebra
import Numeric.GSL.Statistics as S
import Control.DeepSeq (NFData, rnf)
import Numeric.GSL.Statistics (mean,variance,variance_m,stddev)
import Control.Parallel (pseq)
import Data.Ord (comparing)
import Data.List (transpose, nub,minimumBy,foldl')
import LMisc
import Data.Time
import IO

ident n = (diag . fromList) (take n $ repeat 1) :: Matrix Double

vector xs = fromList xs :: Vector Double
diagl = diag . vector
rowM = asRow . vector
colM = asColumn . vector
-------------------------------------------------------
foldl'Rnf :: (NFData a) => (a -> b -> a) -> a -> [b] -> a
foldl'Rnf f z xs = lgo z xs
    where
        lgo z   []      = z
        lgo z (x:xs)    = lgo z' xs
            where
                z' = rnf (f z x) `pseq` (f z x)

------ standardise by subtracting the mean and dividing by the standard deviation
vstNorm       xs            = mapVector (\ a -> (a - m) / std) xs
                 where m    = mean xs
                       std  = sqrt $ variance_m m xs

-- unstandardising
vunStdNorm m std vs  = mapVector (\a -> std * a + m)  vs

--mean,variance,variance_m,stddev)
vunStdNorm' vs = vunStdNorm (mean vs) (stddev vs)  vs
-----------------------------------------------------------------------------------------------
logg a      --     = log a
      | a <= 0     = a
      | otherwise  = log a
-------------------------------------------------------
infix 1  <+>
(<+>) :: Matrix Double -> Matrix Double -> Matrix Double
m1  <+>    m2   = fromRows $ zipWith (zipVectorWith (+)) (toRows m1) (toRows m2)

(<->) :: Matrix Double -> Matrix Double -> Matrix Double
m1  <->    m2  =  fromRows $ zipWith (zipVectorWith (-)) (toRows m1) (toRows m2)
--}

-- adds a to the diadonal elements of ma
mkZ :: Double -> Matrix Double -> Matrix Double
mkZ  a   ma  = mapMatrixWithIndex (pli a) (trans ma <>  ma)
    where  pli b (!i,!j) v = if i == j then v + b else v

mmlLinearRegLP :: Vector Double ->  Double -> Double -> Matrix Double -> (Double, (Vector Double, (Double, Vector Double))) -- Vector Double
mmlLinearRegLP     xs     m   r  ys  =  (ans, (xTy, (sig, abar)))
                where
                  z       =   mkZ m ys -- make the Z matrix
                  n       =   fromIntegral $ rows  ys
                  !mM     =   n + k  + 1
                  !rest   =   log r + 0.5 * log (2 * mM * abs (det z)) + (n / 2) -- cd (k+1)
                  --abar   =   (inv z) <> (trans ys) <>  xs --
                  abar_aux    =    toColumns $ linearSolveLS z (asColumn ((trans ys) <>  xs))
                  abar    =  head abar_aux
                  !k      =   fromIntegral $ cols z
                  sigSq   =    (xs <.> xs)  - ( xs <.> (ys <> abar))
                  !sig    =   sqrt (sigSq / n)
                  xTy     =     (ys <>  abar) -- predictions fromList . manyInvNS1 Nothing $ toList
                  ans     =   (mM/2) * log (2 * pi) + n * log sig - 0.5 * (k + 1) * log k + rest -- k is 2 here
----
mmlMLwithAbar ::Vector Double ->  Vector Double ->  Double -> Double -> Matrix Double -> Double
mmlMLwithAbar abar     xs     m   r  ys  =  ans
                where
                  z       =   mkZ m ys -- make the Z matrix
                  n       =   fromIntegral $ rows  ys
                  !mM     =   n + k  + 1
                  !k      =   fromIntegral $ cols z
                  !rest   =   log r + 0.5 * logg (2 * mM * abs (det z)) + (n / 2) -- cd (k+1)
                  sigSq   =   (xs <.> xs)  - ( xs <.> (ys <> abar))
                  !sig    =   sqrt (sigSq / n)
                  xTy     =     (ys <>  abar) -- predictions fromList . manyInvNS1 Nothing $ toList
                  ans     =   (mM/2) * log (2 * pi) + n * log sig - 0.5 * (k + 1) * log m + rest

--- realigning a matrix
--- map mkMatrix : [[(Matrix Double,[Int])]]
realignedNpieces :: Int ->  --- number of interventions
    Int ->                  --- minimum length of each sub group
   -- Int ->                  -- maximum length of each subgroup
    [[Double]] ->             --- the list of values
    [(Matrix Double,[Int])] --- the matrix for analysis, along with the locations of interventions
realignedNpieces n m  =  nub . concatMap combMatrix . transpose . map  mkMatrix -- map (\ks -> (newMX, ks)) ctts
    where
        ---
        combMatrix :: [(Matrix Double,[Int])] -> [(Matrix Double, [Int]) ]
        combMatrix mcI  = map (\ks -> (newMX, ks)) cts
            where
                (mxx, cts) = unzip mcI
                newMX      = appendByColumns mxx
        mkMatrix :: [Double] ->  [(Matrix Double,[Int])]
        mkMatrix xs      =   [(reAlign ps ,ks) |
                                let nPnms  =  nPiecesP n m Nothing xs
                                , (ps,ks) <-   nPnms
                                ]
            where
                nPns             = []
                len              = length xs
                cons a vs        = join [fromList [a] , vs]
                app2rows f       = fromRows . map f  . toRows
                tS'              = app2rows (cons 1) . trans . app2rows vstNorm . fromLists --vstNorm
                reAlign          = tS' . adjustMatrix . fillZeroes_aux len 0 1
appendByColumns :: Element t => [Matrix t] -> Matrix t
appendByColumns  = fromRows . concatMap toRows

-- draw a best fit line through a list of points
lineFun :: Vector Double -> (Double,Double)
lineFun vs = (ks @> 1, ks @> 0)
    where
        mx          =  (fst . head . realignedNpieces 0 1 . (: [])) [1 .. fromIntegral $ dim  vs]
        (_,(_,ks))  =  snd $ mmlLinearRegLP (vstNorm vs) 1 1  mx

--
-- finding mormalized correlation using  using MML
mmlCorr :: [Double] -> [Double] -> Double
mmlCorr xs ys =  nCorrV (fromList xs) (fromList ys)

mmlCorr1  xs ys =  (mmlCorr_aux (Just n) xs ys)/ (mmlCorr_aux Nothing [1 .. fromIntegral n] [1 .. fromIntegral n])
    where
        n           = (minimum . map length) [xs, ys]
        mmlCorr_aux mN xs' ys' = (ks @> 1)
            where
                stNorm      =  toList . vstNorm . fromList
                yss         =  (stNorm . maybe ys' (flip take ys')) mN
                vss         =  (vstNorm . fromList . maybe ys' (flip take xs')) mN
                (_,(_,ks))  =  snd $ mmlLinearRegLP vss 1 1  mx
                mx          =  (fst . head . realignedNpieces 0 1 . (: [])) yss

-- relative fit of twoe straight line functions
relFit :: [Double] -> [Double] -> Double
relFit xs ys   = relFit_1 (vstNorm $ fromList xs) (vstNorm $ fromList ys)
relFit_1  xs ys =   abs ((mly  - mlx)  /  mlx)
    where
        mly    = mmlMLwithAbar ab  ys  1 1  (mx ys)
        (mlx,(_,(_,ab))) =   mmlFun  xs
        mx   xs'     =  (fst . head . realignedNpieces 0 1 . (: [])) [1 .. fromIntegral $ dim xs']
        mmlFun vs    =  mmlLinearRegLP vs 1 1  (mx vs)

---pearsons correlation
corr xs  ys  =  S.correlation (fromList xs) (fromList ys)

-- normalised correlation
nCorrV xs ys  = (S.correlation (subVector 0 n xs) (subVector 0 n ys))
    where
         n    =  min (dim xs) (dim ys)

---
standardizedFit  :: [Vector Double] -> Double
standardizedFit [] = 0
standardizedFit ks@(_:_) = S.stddev . fromList $ map mSgLn ks
    where
        mSgLn  ys    = mmlMLwithAbar ab  ys  1 1  (mx ys)
        (_,(_,(_,ab)))   = minimumBy (comparing fst) $ map mmlFun ks -- (fst . snd . snd)
        mx   xs'     =  (fst . head . realignedNpieces 0 1 . (: [])) [1 .. fromIntegral $ dim xs']
        mmlFun vs    =  mmlLinearRegLP vs 1 1  (mx vs)


{-# LANGUAGE CPP, TypeOperators, TypeFamilies #-}
{-# LANGUAGE DeriveFunctor, DeriveDataTypeable #-}
{-# OPTIONS_GHC -Wall #-}

-- {-# OPTIONS_GHC -fno-warn-unused-imports #-} -- TEMP
-- {-# OPTIONS_GHC -fno-warn-unused-binds   #-} -- TEMP

#define CustomComplex

#if !defined CustomComplex
{-# OPTIONS_GHC -fno-warn-orphans #-} -- for HasRep
#endif

----------------------------------------------------------------------
-- |
-- Module      :  ConCat.Complex
-- Copyright   :  (c) 2015 Conal Elliott and David Banas
-- License     :  BSD3
--
-- Maintainer  :  conal@conal.net
-- Stability   :  experimental
--
-- Simplified Complex type
----------------------------------------------------------------------

module ConCat.Complex (Complex(..),cis) where

#if defined CustomComplex
import Control.Applicative (liftA2)
import Data.Typeable
import Data.Data
import Test.QuickCheck (Arbitrary(..),CoArbitrary(..))
#else
import Data.Complex
#endif

#if defined CustomComplex
infixl 1 :+
data Complex a = a :+ a deriving (Functor,Eq,Show,Typeable,Data,Ord)

-- | The nonnegative magnitude of a complex number.
-- {-# SPECIALISE magnitude :: Complex Double -> Double #-}
magnitude :: (RealFloat a) => Complex a -> a

magnitude (x:+y) = sqrt (x*x + y*y)

-- From Data.Complex

-- magnitude (x:+y) =  scaleFloat k
--                      (sqrt (sqr (scaleFloat mk x) + sqr (scaleFloat mk y)))
--                     where k  = max (exponent x) (exponent y)
--                           mk = - k
--                           sqr z = z * z

-- | The phase of a complex number, in the range @(-'pi', 'pi']@.
-- If the magnitude is zero, then so is the phase.
-- {-# SPECIALISE phase :: Complex Double -> Double #-}
phase :: (RealFloat a) => Complex a -> a
phase (0 :+ 0)   = 0            -- SLPJ July 97 from John Peterson
phase (x:+y)     = atan2 y x

-- | Complex Identity Scaler
-- Returns a point on the unit circle, such that the angle between the
-- vector to that point and the positive x-axis is equal to the function
-- argument, given in radians.
cis :: (RealFloat a) => a -> Complex a
cis w = cos w :+ sin w

instance  (RealFloat a) => Num (Complex a)  where
    -- {-# SPECIALISE instance Num (Complex Float) #-}
    -- {-# SPECIALISE instance Num (Complex Double) #-}
    (x:+y) + (x':+y')   =  (x+x') :+ (y+y')
    (x:+y) - (x':+y')   =  (x-x') :+ (y-y')
    (x:+y) * (x':+y')   =  (x*x'-y*y') :+ (x*y'+y*x')
    negate (x:+y)       =  negate x :+ negate y
    abs z               =  magnitude z :+ 0
    signum (0:+0)       =  0
    signum z@(x:+y)     =  x/r :+ y/r  where r = magnitude z
    fromInteger n       =  fromInteger n :+ 0
    {-# INLINE (+)         #-}
    {-# INLINE (-)         #-}
    {-# INLINE (*)         #-}
    {-# INLINE negate      #-}
    {-# INLINE abs         #-}
    {-# INLINE signum      #-}
    {-# INLINE fromInteger #-}

instance  (RealFloat a) => Fractional (Complex a)  where
    -- {-# SPECIALISE instance Fractional (Complex Float) #-}
    -- {-# SPECIALISE instance Fractional (Complex Double) #-}
    recip z@(x:+y) = x/mag :+ (-y)/mag where mag = magnitude z
--     (x:+y) / (x':+y')   =  (x*x''+y*y'') / d :+ (y*x''-x*y'') / d
--                            where x'' = scaleFloat k x'
--                                  y'' = scaleFloat k y'
--                                  k   = - max (exponent x') (exponent y')
--                                  d   = x'*x'' + y'*y''

    fromRational a      =  fromRational a :+ 0

instance  (RealFloat a) => Floating (Complex a) where
    -- {-# SPECIALISE instance Floating (Complex Float) #-}
    -- {-# SPECIALISE instance Floating (Complex Double) #-}
    pi             =  pi :+ 0
    exp (x:+y)     =  expx * cos y :+ expx * sin y
                      where expx = exp x
    log z          =  log (magnitude z) :+ phase z

    sqrt (0:+0)    =  0
    sqrt z@(x:+y)  =  u :+ (if y < 0 then -v else v)
                      where (u,v) = if x < 0 then (v',u') else (u',v')
                            v'    = abs y / (u'*2)
                            u'    = sqrt ((magnitude z + abs x) / 2)

    sin (x:+y)     =  sin x * cosh y :+ cos x * sinh y
    cos (x:+y)     =  cos x * cosh y :+ (- sin x * sinh y)
    tan (x:+y)     =  (sinx*coshy:+cosx*sinhy)/(cosx*coshy:+(-sinx*sinhy))
                      where sinx  = sin x
                            cosx  = cos x
                            sinhy = sinh y
                            coshy = cosh y

    sinh (x:+y)    =  cos y * sinh x :+ sin  y * cosh x
    cosh (x:+y)    =  cos y * cosh x :+ sin y * sinh x
    tanh (x:+y)    =  (cosy*sinhx:+siny*coshx)/(cosy*coshx:+siny*sinhx)
                      where siny  = sin y
                            cosy  = cos y
                            sinhx = sinh x
                            coshx = cosh x

    asin z@(x:+y)  =  y':+(-x')
                      where  (x':+y') = log (((-y):+x) + sqrt (1 - z*z))
    acos z         =  y'':+(-x'')
                      where (x'':+y'') = log (z + ((-y'):+x'))
                            (x':+y')   = sqrt (1 - z*z)
    atan z@(x:+y)  =  y':+(-x')
                      where (x':+y') = log (((1-y):+x) / sqrt (1+z*z))

    asinh z        =  log (z + sqrt (1+z*z))
    acosh z        =  log (z + (z+1) * sqrt ((z-1)/(z+1)))
    atanh z        =  0.5 * log ((1.0+z) / (1.0-z))

instance Arbitrary a => Arbitrary (Complex a) where
  arbitrary = liftA2 (:+) arbitrary arbitrary
  shrink (x :+ y) = map (uncurry (:+)) (shrink (x,y))

instance CoArbitrary a => CoArbitrary (Complex a) where
  coarbitrary (x :+ y) = coarbitrary x . coarbitrary y

#endif


module Mandelbrot where

import Graphics.UI.GLUT
import Control.Monad
import Data.Int
import Data.Complex
import System.Exit ( exitWith, ExitCode(ExitSuccess) )
import Data.IORef
import qualified Data.ByteString as B
import           Graphics.Rendering.OpenGL.Capture (capturePPM)
import           Graphics.Image
import           System.Directory (removeFile)
import Data.List.Split (splitOn)

convert :: FilePath -> FilePath -> Bool -> IO ()
convert input output remove = do
  ppm <- readImageRGBA VU input
  writeImage output ppm
  when remove $ removeFile input

iterations = 40

x // y = fromIntegral x / fromIntegral y

-- Divides [a] into [[a], [a], ...] with each sublist of length n,
-- except the last sublist which has length <= n.
chunkify n [] = []
chunkify n xs = let (xs', rest) = splitAt n xs
                in xs' : chunkify n rest

-- Converts a coordinate in screen space to a vertex.
pix2vert (Size w h) (x, y) = Vertex2 ((3 // w * fromIntegral x) - 2.0)
                             ((2 // h * fromIntegral y) - 1.0)

-- List of all of the vertices that represent screen pixels.
vertices :: IO [Vertex2 GLfloat]
vertices = get windowSize >>= \(Size w h) ->
           return $ [pix2vert (Size w h) (x, y) | x <- [0..w-1], y <- [0..h-1]]

-- Gets the color for a number of iterations.
color3 r g b = Color3 r g b
getcolor :: Int -> Int -> Color3 Float
getcolor iterations iter
  | iter == iterations = color3 0.1 0.1 0.1
  | otherwise          = color3 (amt*0.5) (amt * 0.5) (amt*0.5)
  where amt = iter // iterations

-- Returns the number of iterations <= the maximum iterations of the
-- Mandelbrot set at the given vertex.
mandel numIterations horizontalShift verticalShift zoom (Vertex2 r i) =
  length . takeWhile (\z -> Data.Complex.magnitude z <= 2) .
  take numIterations $
  iterate (\z -> z^2 + (((r / zoom) + horizontalShift) :+ ((i / zoom) + verticalShift))) 0

-- plots one point.
drawVert iter horShift verShift zoom v = do
  horShift' <- get horShift
  verShift' <- get verShift
  zoom' <- get zoom
  color . (getcolor iter) $ mandel iter horShift' verShift' zoom' v
  vertex v

-- draws the whole fractal
display :: IORef Int -> IORef GLfloat -> IORef GLfloat -> IORef GLfloat ->  IO ()
display numIter horShift verShift zoom = do
  -- clear [ ColorBuffer ]
  displayCallback $= (vertices >>= (display' numIter horShift verShift zoom) . (\x -> [x]))
  get currentWindow >>= postRedisplay

-- draws all the vertices in slices (to update the display while drawing).
display' :: IORef Int -> IORef GLfloat -> IORef GLfloat -> IORef GLfloat -> [[Vertex2 GLfloat]] -> IO ()
display' numIter horShift verShift zoom chunks = do
  iter <- get numIter
  mapM_ (render iter) chunks
  displayCallback $= display numIter horShift verShift zoom
  -- exitWith ExitSuccess
  where render iter' points = do
          renderPrimitive Points (mapM_ (drawVert iter' horShift verShift zoom) points)
          flush

-- keyboard :: Key -> KeyState -> p1 -> p2 -> IO ()
-- keyboard :: IORef GLfloat -> IORef GLfloat -> Key -> KeyState -> p3 -> p4 -> IO ()
keyboard numIter horShift verShift zoom (Char keyPressed) Down _ _ = case keyPressed of
   ' '   -> postRedisplay Nothing
   '\27' -> exitWith ExitSuccess
   '\100' -> do
     val <- shiftAmount zoom
     horShift $~! (+ val)
     displayCallback $= display numIter horShift verShift zoom
   '\97' -> do
     val <- shiftAmount zoom
     horShift $~! (+ (-val))
     displayCallback $= display numIter horShift verShift zoom
   '\115' -> do
     val <- shiftAmount zoom
     verShift $~! (+ (-val))
     displayCallback $= display numIter horShift verShift zoom
   '\119' -> do
     val <- shiftAmount zoom
     verShift $~! (+ val)
     displayCallback $= display numIter horShift verShift zoom
   '\61' -> do
     val <- get zoom
     zoom $~! (+ val)
     displayCallback $= display numIter horShift verShift zoom
   '\45' -> do
     val <- get zoom
     zoom $~! (+ (-(val / 2)))
     displayCallback $= display numIter horShift verShift zoom
   '\49' -> do
     resetIter numIter
     displayCallback $= display numIter horShift verShift zoom
   '\50' -> do
     resetIter numIter
     numIter $~! (+ (20))
     displayCallback $= display numIter horShift verShift zoom
   '\51' -> do
     resetIter numIter
     numIter $~! (+ (100))
     displayCallback $= display numIter horShift verShift zoom
   '\52' -> do
     resetIter numIter
     numIter $~! (+ (400))
     displayCallback $= display numIter horShift verShift zoom
   '\53' -> do
     resetIter numIter
     numIter $~! (+ (1000))
     displayCallback $= display numIter horShift verShift zoom

   'c' -> do
     horShift' <- get horShift
     verShift' <- get verShift
     zoom' <- get zoom
     let name = show horShift' ++ "," ++
                show verShift' ++ "," ++
                show zoom'

     putStrLn ""
     putStrLn ""
     putStrLn "Writing out file:"
     putStrLn $ "images/" ++ name ++ ".png"
     putStrLn ""
     let ppm = "images/" ++ name ++ ".ppm"
         png = "images/" ++ name ++ ".png"
     (>>=) capturePPM (B.writeFile ppm)
     Mandelbrot.convert ppm png True

   _     -> return ()
keyboard _ _ _ _ _ _ _ _ = return ()

shiftAmount zoom = do
  val <- get zoom
  pure (0.1 / val)

resetIter numIter = do
  val <- get numIter
  numIter $~! (+ (iters + (-val)))

parseFile' (file) = return
  where [x, y, z] = splitOn "," $ (splitOn ".png" file) !! 0
        return = ((read x :: GLfloat), (read y :: GLfloat), (read z :: GLfloat))

parseFile :: [String] -> (GLfloat, GLfloat, GLfloat)
parseFile [justOne] = parseFile' $ (splitOn "images/" justOne) !! 1
parseFile _ = (0.0, 0.0, 1.0)

iters = 20
mandelMain = do
  (_, file) <- getArgsAndInitialize
  mapM putStrLn file
  let (initHor, initVer, initZoom) = parseFile file

  initialDisplayMode $= [ SingleBuffered, RGBMode ]
  initialWindowSize $= Size 1200 1024
  initialWindowPosition $= Position 100 100
  createWindow "Mandelbrot"
  clearColor $= Color4 0 0 0 0
  matrixMode $= Projection
  loadIdentity
  ortho (-2) 1 (-1) 1 (-1) 1
  -- fullScreen
  numIter <- newIORef iters
  horShift <- newIORef initHor
  verShift <- newIORef initVer
  zoom <- newIORef initZoom
  keyboardMouseCallback $= Just (keyboard numIter horShift verShift zoom)
  displayCallback $= (display numIter horShift verShift zoom)
  idleCallback $= Just (idle numIter)
  mainLoop

idle :: IORef Int -> IdleCallback
idle numIter = do
  val <- get numIter
  let newVal = if val < 50 then 5 else 100
  numIter $~! (+ 0)
  postRedisplay Nothing


module Main where

import Test.Hspec
import Test.QuickCheck
import System.Console.Ansigraph
import System.Console.Ansigraph.Examples (wave)
import Data.Complex
import Data.Char (isSpace)
-- for GHC <= 7.8
import Control.Applicative

---- For horizontal graphing tests ----

bars = " ▁▂▃▄▅▆▇█"

eighths :: [Double]
eighths = (/8) <$> [0..8]

{- These two vectors have a maximum element of 8, to maintain a particular normalization,
   while displacing other entries by just under one half. Both vectors should display
   as █▇▆▅▄▃▂▁. Same idea for m1, m_hi, m_low below. -}

v_hi :: [Double]
v_hi = 8 : map (+0.48) (reverse [0..7])

v_low :: [Double]
v_low = 8 : map (\x -> x - 0.48) (reverse [0..7])

{- These functions are for testing the property that:
   simpleRender v == complement <$> simpleRenderR v-}

-- the 'complement' of a particular bar character is the one representing the negation of the
-- corresponding numeric value (when displayed with reversed foreground/background colors).
complement :: Char -> Char
complement c = match c bars (reverse bars)

match :: Eq a => a -> [a] -> [a] -> a
match x (y:ys) (z:zs) = if x == y then z else match x ys zs

rwave = realPart <$> wave

barInvert :: (String,String) -> (String,String)
barInvert (x,y)  = (complement <$> y, complement <$> x)

barInvertC :: (String,String,String,String) -> (String,String,String,String)
barInvertC (x1,y1,x2,y2) = (complement <$> y1,
                            complement <$> x1,
                            complement <$> y2,
                            complement <$> x2)

---- For matrix graphing tests ----

fourths :: [Double]
fourths = (/4) <$> [1..4]

mmap :: (a -> b) -> [[a]] -> [[b]]
mmap = map . map

m1 :: [[Double]]
m1 = [[1,2,3],[3,4,4]]

m_low :: [[Double]]
m_low = [[0.55,1.55,2.55], [2.55,3.55,4]]

m_hi :: [[Double]]
m_hi = [[1.45,2.45,3.45], [3.45,4,4]]

---- The spec ----

main :: IO ()
main = hspec $ do

  describe "renderPV" $ do

    it "maps exact multiples of 1/8 to the correct characters" $
      renderPV eighths `shouldBe` bars

    it "maps scaled multiples of 1/8 to the correct characters" $
      renderPV ((* 1337) <$> eighths) `shouldBe` bars

    it "maps values near tops of rounding regions to the correct characters" $
      renderPV v_hi `shouldBe` reverse bars

    it "maps values near bottoms of rounding regions to the correct characters" $
      renderPV v_low `shouldBe` reverse bars

    it "maps zero vectors to whitespace" $
      renderPV [0,0,0] `shouldBe` "   "

    it "maps non-positive vectors to whitespace" $
      renderPV [-1,-23,-0.02] `shouldBe` "   "

    it "maps null vectors to null strings" $
      renderPV [] `shouldBe` ""

  describe "renderRV" $ do

    it "inverts 'rwave' consistently" $
      let (p,n) = renderRV rwave
      in  renderRV (negate <$> rwave) `shouldBe` (complement <$> n, complement <$> p)

    it "inverts arbitrary vectors consistently" $
      property $ \xs -> renderRV xs == barInvert (renderRV $ negate <$> xs)

    it "maps arbitrary-dimensional zero vectors to whitespace, for positive components" $
      property $ \n -> let (p,m) = renderRV $ replicate n 0
                       in  all isSpace p

    it "maps arbitrary-dimensional zero vectors to solid bocks █, for negative components" $
      property $ \n -> let (p,m) = renderRV $ replicate n 0
                       in  all (== '█') m

  describe "renderCV" $ do

    it "inverts 'wave' consistently" $ do
      let (rp,rn,ip,im) = renderCV wave
      renderCV (negate <$> wave) `shouldBe` (complement <$> rn,
                                             complement <$> rp,
                                             complement <$> im,
                                             complement <$> ip)

    it "inverts arbitrary complex vectors consistently" $
      property $ \zs -> renderCV zs == barInvertC (renderCV $ negate <$> zs)

    it "maps arbitrary dimensional zero vectors to whitespace, for positive components" $
      property $ \n -> let (r1,_,i1,_) = renderCV $ replicate n (0.0 :+ 0.0)
                       in  all isSpace $ r1 ++ i1

    it "maps arbitrary dimensional zero vectors to solid bocks █, for negative components" $
      property $ \n -> let (_,r2,_,i2) = renderCV $ replicate n (0.0 :+ 0.0)
                       in  all (== '█') $ r2 ++ i2

    it "transforms appropriately under exchange of real and complex components" $
      property $ \v -> let v'  = (\(x :+ y) -> (y :+ x)) <$> v
                           (x1,y1,x2,y2)  = renderCV v
                           (x3,y3,x4,y4)  = renderCV v'
                       in  and [x1 == x4,
                                y1 == y4,
                                x2 == x3,
                                y2 == y3]

  describe "matShow" $ do

    it "maps exact multiples of 1/4 to the correct characters" $
      matShow [[1/4,2/4],[3/4,4/4]] `shouldBe` ["░▒", "▓█"]

    it "maps scaled multiples of 1/4 to the correct characters" $ do
      matShow (mmap (* 3097) [[1,2],[3,4]]) `shouldBe` ["░▒", "▓█"]
      matShow m1 `shouldBe` ["░▒▓", "▓██"]

    it "maps values near top of rounding regions to the correct characters" $
      matShow m_hi `shouldBe` ["░▒▓", "▓██"]

    it "maps values near bottom of rounding regions to the correct characters" $
      matShow m_low `shouldBe` ["░▒▓", "▓██"]


{-# LANGUAGE FlexibleInstances #-}

module TheField.Plot where

import Data.Complex
import Data.List as L
import Graphics.Rendering.Chart.Easy

class Plotable a where
  plotable :: RealFloat b => a -> (b, b)

instance RealFloat a => Plotable (Complex a) where
  plotable c = (realToFrac . realPart $ c, realToFrac . imagPart $ c)

instance RealFloat a => Plotable (a, a) where
  plotable (a, b) = (realToFrac a, realToFrac b)

plot' :: (Plotable a, RealFloat b, Show b) => [a] -> b -> b -> EC (Layout b b) ()
plot' s w h = do
  setColors [opaque red]
  plot (points "s" $ L.map plotable s)
  vectorPlot w h

plotComplexAddition :: Complex Double -> Complex Double -> EC (Layout Double Double) ()
plotComplexAddition a b = do
  vectorPlot (maxReal * 1.1) (maxImag * 1.1)
  plot . liftEC $ do
    plot_lines_style . line_color .= opaque red
    plot_lines_style . line_width .= 3
    plot_lines_values .= L.map (L.map plotable) [[0, a], [a, a + b]]
  plot . liftEC $ do
    plot_lines_style . line_color .= opaque green
    plot_lines_style . line_width .= 2
    plot_lines_style . line_dashes .= [5, 5]
    plot_lines_values .= [ L.map plotable [0, a + b] ]
  where
    maxReal = maximum [ realPart a, realPart b, realPart (a + b) ]
    maxImag = maximum [ imagPart a, imagPart b, imagPart (a + b) ]

vectorPlot :: (RealFloat a, Show a) => a -> a -> EC (Layout a a) ()
vectorPlot w h = do
  layout_x_axis . laxis_generate .= scaledAxis def (-w,w)
  layout_y_axis . laxis_generate .= scaledAxis def (-h,h)
  layout_bottom_axis_visibility . axis_show_line .= False
  layout_bottom_axis_visibility . axis_show_ticks .= False
  layout_left_axis_visibility . axis_show_line .= False
  layout_left_axis_visibility . axis_show_ticks .= False
  layout_legend .= Nothing
  setColors [opaque black, opaque black]
  plot (line "x" [x])
  plot (line "y" [y])
    where
      x = [(-w*2, 0), (w*2, 0)]
      y = [(0, -h*2), (0, h*2)]


{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE NoMonomorphismRestriction #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE ConstraintKinds #-}
{-# LANGUAGE FlexibleContexts #-}

{-# OPTIONS -Wall #-}
module Statistics.Iteratee.Sample (
  minMaxNBy
, range
, mean
, harmonicMean
, variance
, stdDev
) where

import Statistics.Iteratee.Compat
import Control.Arrow
import Control.Monad
import Data.Iteratee as I

import Data.Heap as Heap

--  | /O(n)/ minMaxNBy. Calculate the 'n' highest and lowest elements of the
--  stream, according to the given priority function.
--  Returns /([minimum],[maximum]/ with the /(minimum,maximum)/
--  elements listed first.
minMaxNBy
  :: forall m prio s el. (Monad m, Ord prio, ListLikey s el)
  => Int
  -> (el -> prio)
  -> Iteratee s m ([(prio,el)],[(prio,el)])
minMaxNBy ns prio = finalize `liftM` I.foldl' step (Heap.empty,Heap.empty)
  where
    finalize :: (MaxPrioHeap prio el, MinPrioHeap prio el)
             -> ([(prio,el)],[(prio,el)])
    finalize = Heap.toDescList *** Heap.toDescList
    addHeap val = Heap.insert (prio val, val)
    step :: (MaxPrioHeap prio el, MinPrioHeap prio el) -> el
            -> (MaxPrioHeap prio el, MinPrioHeap prio el)
    step (!mins,!maxes) val = let sz = Heap.size mins
                                  adj hp = if sz >= ns
                                             then Heap.drop 1 hp
                                             else hp
                              in (adj $ addHeap val mins
                                 ,adj $ addHeap val maxes)
{-# INLINE minMaxNBy #-}

-- | /O(n)/ Range. The difference between the largest and smallest elements of
-- a stream.
range :: (Monad m, ListLikey s el, Num el, Ord el)
      => Iteratee s m el
range = finalize `liftM` minMaxNBy 1 id
  where
    finalize ([mins],[maxes]) = snd maxes - snd mins
    finalize _ = 0
{-# INLINE range #-}

-- | /O(n)/ Arithmetic mean.  Uses Welford's algorithm.
mean :: forall s m el. (Fractional el, Monad m, ListLikey s el)
     => Iteratee s m el
mean = fst `liftM` I.foldl' step (0,0)
  where
    step :: (el,Integer) -> el -> (el,Integer)
    step (!m,!n) x = let m' = m + (x-m) / fromIntegral n'
                         n' = n + 1
                     in  (m',n')
{-# INLINE mean #-}

-- | /O(n)/ Harmonic mean.
harmonicMean :: (Fractional el, Monad m, ListLikey s el) => Iteratee s m el
harmonicMean = finalize `liftM` I.foldl' step (0,0 :: Integer)
  where
    finalize (m,n) = fromIntegral n / m
    step (!m,!n) val = (m+(1/val),n+1)
{-# INLINE harmonicMean #-}

-- | /O(n)/ variance, using Knuth's algorithm.
var :: (Fractional el, Integral t, Monad m, ListLikey s el)
    => Iteratee s m (t, el, el)
var = I.foldl' step (0,0,0)
  where
    step (!n,!m,!s) x = let n' = n+1
                            m' = m+d/fromIntegral n'
                            s' = s+d* (x-m')
                            d  = x-m
                        in (n',m',s')
{-# INLINE var #-}

-- | /O(n)/ Maximum likelihood estimate of a sample's variance, using Knuth's
--   algorithm.
variance :: (Fractional b, Monad m, ListLikey s b) => Iteratee s m b
variance = finalize `liftM` var
  where
    finalize (n,_,s)
      | n > 1 = s / fromInteger n
      | otherwise = 0
{-# INLINE variance #-}

-- | /O(n) Standard deviation, using Knuth's algorithm.
stdDev :: (Floating b, Monad m, Functor m, ListLikey s b) => Iteratee s m b
stdDev = sqrt `liftM` variance
{-# INLINE stdDev #-}


-- example 3.11

import Statistics.Distribution
import Statistics.Distribution.Geometric

g :: GeometricDistribution
g = geometric 0.02

main = do
  putStrLn $ "the rate of engines which survive two hours: "
          ++ show (1 - cumulative g 2.0)
  putStrLn $ "the rate of engines which survive two hours: "
          ++ show (complCumulative g 2.0)
  putStrLn $ "the rate of engines which survive two hours: "
          ++ show (1 - sum (map (probability g) [0 .. 2]))