{-# LANGUAGE DeriveDataTypeable, DeriveGeneric #-}
-- |
-- Module    : Statistics.Distribution.ChiSquared
-- Copyright : (c) 2010 Alexey Khudyakov
-- License   : BSD3
--
-- Maintainer  : bos@serpentine.com
-- Stability   : experimental
-- Portability : portable
--
-- The chi-squared distribution. This is a continuous probability
-- distribution of sum of squares of k independent standard normal
-- distributions. It's commonly used in statistical tests
module Statistics.Distribution.ChiSquared (
          ChiSquared
        -- Constructors
        , chiSquared
        , chiSquaredNDF
        ) where

import Data.Aeson (FromJSON, ToJSON)
import Data.Binary (Binary)
import Data.Data (Data, Typeable)
import GHC.Generics (Generic)
import Numeric.SpecFunctions (
  incompleteGamma,invIncompleteGamma,logGamma,digamma)

import qualified Statistics.Distribution         as D
import qualified System.Random.MWC.Distributions as MWC
import Data.Binary (put, get)

-- | Chi-squared distribution
newtype ChiSquared = ChiSquared Int
                     deriving (Eq, Read, Show, Typeable, Data, Generic)

instance FromJSON ChiSquared
instance ToJSON ChiSquared

instance Binary ChiSquared where
    get = fmap ChiSquared get
    put (ChiSquared x) = put x

-- | Get number of degrees of freedom
chiSquaredNDF :: ChiSquared -> Int
chiSquaredNDF (ChiSquared ndf) = ndf

-- | Construct chi-squared distribution. Number of degrees of freedom
--   must be positive.
chiSquared :: Int -> ChiSquared
chiSquared n
  | n <= 0    = error $
     "Statistics.Distribution.ChiSquared.chiSquared: N.D.F. must be positive. Got " ++ show n
  | otherwise = ChiSquared n

instance D.Distribution ChiSquared where
  cumulative = cumulative

instance D.ContDistr ChiSquared where
  density  = density
  quantile = quantile

instance D.Mean ChiSquared where
    mean (ChiSquared ndf) = fromIntegral ndf

instance D.Variance ChiSquared where
    variance (ChiSquared ndf) = fromIntegral (2*ndf)

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

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

instance D.Entropy ChiSquared where
  entropy (ChiSquared ndf) =
    let kHalf = 0.5 * fromIntegral ndf in
    kHalf
    + log 2
    + logGamma kHalf
    + (1-kHalf) * digamma kHalf

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

instance D.ContGen ChiSquared where
    genContVar (ChiSquared n) = MWC.chiSquare n

cumulative :: ChiSquared -> Double -> Double
cumulative chi x
  | x <= 0    = 0
  | otherwise = incompleteGamma (ndf/2) (x/2)
  where
    ndf = fromIntegral $ chiSquaredNDF chi

density :: ChiSquared -> Double -> Double
density chi x
  | x <= 0    = 0
  | otherwise = exp $ log x * (ndf2 - 1) - x2 - logGamma ndf2 - log 2 * ndf2
  where
    ndf  = fromIntegral $ chiSquaredNDF chi
    ndf2 = ndf/2
    x2   = x/2

quantile :: ChiSquared -> Double -> Double
quantile (ChiSquared ndf) p
  | p == 0         = 0
  | p == 1         = 1/0
  | p > 0 && p < 1 = 2 * invIncompleteGamma (fromIntegral ndf / 2) p
  | otherwise      =
    error $ "Statistics.Distribution.ChiSquared.quantile: p must be in [0,1] range. Got: "++show p


module Main where

import Graphics.Rendering.Chart.Backend.Diagrams (toFile)
import qualified Graphics.Rendering.Chart.Easy as Plt
import Numeric.GSL.ODE
import Numeric.LinearAlgebra

dynamics :: Double -> [Double] -> [Double]
dynamics _ [_, x1, w] =
  [ x1,
    - mass * g + k * w + cd * x1 ** 2,
    (torque - c * w) / j
  ]
  where
    -- timestep
    dt = 0.01
    -- Coefficient of drag
    cd = 0.1
    -- mass
    mass = 22
    -- gravity
    g = 9.8
    -- inertia
    j = 0.05
    -- Omega to thrust
    k = 5
    -- Coefficient of friction motor
    c = 0.1
    -- Calculate a physics based torque
    torque =
      clamp (-10, 10) $
        mass * g * j / (k * dt)
          - x1 / dt ** 2
          - w * j / dt
          + c * w
    -- Clamp between lower bound and upper bound
    clamp (lb, ub) x = max lb (min ub x)

dumbCopter :: IO ()
dumbCopter = toFile Plt.def "DumpCopter.svg" $ do
  let totalTime = 5
      dt = 0.01
      ts = linspace (round $ totalTime / dt) (0, totalTime)
      sol = odeSolve dynamics [0, -10, 0] ts
      [pos, vel, omega] = toColumns sol
      ts' = toList ts
  Plt.plot (Plt.line "pos" [zip ts' (toList pos)])
  Plt.plot (Plt.line "vel" [zip ts' (toList vel)])
  Plt.plot (Plt.line "omega" [zip ts' (toList omega)])

main :: IO ()
main = dumbCopter


{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE RankNTypes #-}
module Primitives.Math (primitives) where

import Data.Complex (Complex(..))
import Data.Functor ((<&>))
import Data.List (foldl', foldl1')
import Data.Ratio (approxRational)

import Val
import EvaluationMonad (throwError, panic)
import Types.Unwrappers (unwrapNum, unwrapRealNum, unwrapExactInteger)
import Primitives.Bool (predicate, predicateM)
import Primitives.Misc (valuesB)

addP, mulP, modP, quotP, remP :: Primitive
addP = numericBinop "+" (+) 0
mulP = numericBinop "*" (*) 1
modP = realBinop1 "mod" mod
quotP = realBinop1 "quotient" quot
remP = realBinop1 "remainder" rem

-- | Converts a (Haskell) builtin function on Integers to a 'Primitive'.
--   The primitive can only cause type errors;
--   Too many arguments will be folded over the operation;
--   none will use the default value.
numericBinop :: Symbol
             -> (Number -> Number -> Number)
             -> Number
             -> Primitive
numericBinop name op start = Prim name (AtLeast 0) $
  fmap (Number . foldl' op start) . mapM unwrapNum

-- | Same as numericBinop, except at least one arg must be given.
realBinop1 :: Symbol
              -> (RealNumber -> RealNumber -> RealNumber)
              -> Primitive
realBinop1 name op = Prim name (AtLeast 1) $
  fmap (Number . Real . foldl1' op) . mapM unwrapRealNum

guardDivZero :: Primitive -> Primitive
guardDivZero (Prim name arity f) = Prim name arity $ \args ->
    if any isExactZeroVal $ tail args
    then throwError $ Default "divide by zero"
    else f args
  where
    isExactZeroVal (Number n) = isExactZero n
    isExactZeroVal _ = False

-- | See the r7rs standard.
--
-- (- n) evaluates to -n, but (- x y z) evaluates to ((x - y) - z).
subP :: Primitive
subP = Prim "-" (AtLeast 1) $ \case
  [x]      -> Number . negate <$> unwrapNum x
  as@(_:_) -> Number . foldl1' (-) <$> mapM unwrapNum as
  _ -> panic "sub arity"

-- | r7rs: (/ n) is 1/n, (/ x y z) is ((x / y) / z)
divP :: Primitive
divP = guardDivZero $ Prim "/" (AtLeast 1) $ \case
  [] -> panic "div arity"
  [x] -> Number . (1 /) <$> unwrapNum x
  as  -> Number . foldl1' (/) <$> mapM unwrapNum as

negateP :: Primitive
negateP = Prim "negate" (Exactly 1) $
  \[x] -> Number . negate <$> unwrapNum x

numericPredicate :: Symbol -> (RealNumber -> Bool) -> Primitive
numericPredicate name p = predicateM name $ fmap p . unwrapRealNum

zeroCheck :: Primitive
zeroCheck = predicate "zero?" unwrapNum (== 0)

positiveCheck :: Primitive
positiveCheck = numericPredicate "positive?" (> 0)

negativeCheck :: Primitive
negativeCheck = numericPredicate "negative?" (< 0)

oddCheck :: Primitive
oddCheck = numericPredicate "odd?" $ (0 /=) . mod 2

evenCheck :: Primitive
evenCheck = numericPredicate "even?" $ (0 ==) . mod 2

-------------------------------------------------------------------------------
-- typically useful math operations
-------------------------------------------------------------------------------

inexactP :: Primitive
inexactP = Prim "inexact" (Exactly 1) $
  \[x] -> unwrapNum x <&> \case
    Real r -> Number $ Real $ realInexact r
    Complex (r :+ i) -> Number $ Complex $ realInexact r :+ realInexact i

rationalizeP :: Primitive
rationalizeP = Prim "rationalize" (Exactly 2) $
  \[x, eps] -> do
    rx <- unwrapRealNum x
    reps <- unwrapRealNum eps
    return $ Number $ Real $ Ratnum $ approxRational rx reps

squareP :: Primitive
squareP = Prim "square" (Exactly 1) $
  \[x] -> Number . (\n -> n*n) <$> unwrapNum x

-- | specialize ^ to Int powers
(^!) :: Num a => a -> Int -> a
x ^! n = x ^ n

exactIntegerSqrtP :: Primitive
exactIntegerSqrtP = Prim "exact-integer-sqrt" (Exactly 1) $
  \[k] -> do
    kN <- unwrapExactInteger k
    let root = squareRoot kN
        residual = kN - root ^! 2
    valuesB [makeBignum root, makeBignum residual]
  where
    -- This impl of the standard integer square root algorithm is from
    -- https://wiki.haskell.org/Generic_number_type#squareRoot
    squareRoot :: Integer -> Integer
    squareRoot 0 = 0
    squareRoot 1 = 1
    squareRoot n =
      let twopows = iterate (^! 2) 2
          (lowerRoot, lowerN) =
            last $ takeWhile ((n >=) . snd) $ zip (1:twopows) twopows
          newtonStep x = div (x + n `div` x) 2
          iters = iterate newtonStep (squareRoot (div n lowerN) * lowerRoot)
          isRoot r = r ^! 2 <= n && n < (r+1) ^! 2
       in head $ dropWhile (not . isRoot) iters

exptP :: Primitive
exptP = Prim "expt" (Exactly 2) $
  \[z1, z2] -> do
    z1' <- unwrapNum z1
    z2' <- unwrapNum z2
    return $ Number $ z1' ** z2'

-------------------------------------------------------------------------------
-- Complex functions
-------------------------------------------------------------------------------

makeRectangularP :: Primitive
makeRectangularP = Prim "make-rectangular" (Exactly 2) $
  \[x, y] -> do
    a <- unwrapRealNum x
    b <- unwrapRealNum y
    return $ Number $ Complex $ a :+ b

makePolarP :: Primitive
makePolarP = Prim "make-polar" (Exactly 2) $
  \[m, theta] -> do
    m' <- unwrapRealNum m
    theta' <- unwrapRealNum theta
    return $ Number $ Real m' * Complex (cos theta' :+ sin theta')

realPartP :: Primitive
realPartP = Prim "real-part" (Exactly 1) $
  \[z] -> unwrapNum z <&> \case
    Real r -> Number $ Real r
    Complex (r :+ _i) -> Number $ Real r

imagPartP :: Primitive
imagPartP = Prim "imag-part" (Exactly 1) $
  \[z] -> unwrapNum z <&> \case
    Real _r -> Number $ Real 0
    Complex (_r :+ i) -> Number $ Real i

magnitudeP :: Primitive
magnitudeP = Prim "magnitude" (Exactly 1) $
  \[z] -> Number . abs <$> unwrapNum z

angleP :: Primitive
angleP = Prim "angle" (Exactly 1) $
  \[z] -> unwrapNum z <&> \case
    Real r
      | r < 0 || isNegativeZero r -> Number pi
      | otherwise -> Number 0
    Complex (a :+ b) -> Number $ Real $ atan2 b a

liftFloatingMonop :: (Number -> Number) -> Builtin
liftFloatingMonop f = asBuiltin
  where asBuiltin [x] = do
          nx <- unwrapNum x
          return $ Number $ f nx
        asBuiltin _ = panic "liftFloatingMonop arity"

makeFloatingMonopPrim :: (Symbol, Number -> Number) -> Primitive
makeFloatingMonopPrim (name, f) = Prim name (Exactly 1) $ liftFloatingMonop f

floatingMonopPs :: [Primitive]
floatingMonopPs = map makeFloatingMonopPrim
  [ ("exp", exp)
  , ("sqrt", sqrt)
  , ("sin", sin)
  , ("cos", cos)
  , ("tan", tan)
  , ("asin", asin)
  , ("acos", acos)
  ]

logP :: Primitive
logP = Prim "log" (Between 1 2) $ \case
  [z] -> Number . log <$> unwrapNum z
  [z1, z2] -> do
    z1' <- unwrapNum z1
    z2' <- unwrapNum z2
    return $ Number $ logBase z2' z1'
  _ -> panic "logB arity"

atanP :: Primitive
atanP = Prim "atan" (Between 1 2) $ \case
  [z] -> Number . atan <$> unwrapNum z
  [y, x] -> do
    y' <- unwrapRealNum y
    x' <- unwrapRealNum x
    return $ Number $ Real $ atan2 y' x'
  _ -> panic "atanB arity"

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

primitives :: [Primitive]
primitives = [ addP
             , subP
             , mulP
             , divP
             , modP
             , quotP
             , remP
             , negateP
             , zeroCheck
             , positiveCheck
             , negativeCheck
             , oddCheck
             , evenCheck

             , inexactP
             , rationalizeP
             , atanP
             , exactIntegerSqrtP
             , squareP
             , exptP

             , makeRectangularP
             , makePolarP
             , realPartP, imagPartP
             , magnitudeP, angleP
             , logP
             ]
             ++ floatingMonopPs


{-# LANGUAGE DeriveDataTypeable, PatternGuards #-}
module Physics.Quaternion
  ( Quaternion(..)
  , Complicated(..)
  , Hamiltonian(..)
  , slerp
  , asinq
  , acosq
  , atanq
  , asinhq
  , acoshq
  , atanhq
  , absi
  , pow
  , rotate
  ) where

import Control.Applicative
import Control.Lens
import Control.Lens.Rep
import Data.Complex (Complex((:+)))
import Data.Data
import Data.Distributive
import Data.Foldable
import Data.Monoid
import Data.Traversable
import Physics.Epsilon
import Physics.Involutive
import Physics.Metric
import Physics.V3
import Physics.Vector
import Prelude hiding (any)

data Quaternion a = Quaternion a a a a deriving (Eq,Ord,Read,Show,Data,Typeable)

instance Functor Quaternion where
  fmap f (Quaternion e i j k) = Quaternion (f e) (f i) (f j) (f k)
  a <$ _ = Quaternion a a a a

instance Applicative Quaternion where
  pure a = Quaternion a a a a
  Quaternion f g h i <*> Quaternion a b c d = Quaternion (f a) (g b) (h c) (i d)

instance Monad Quaternion where
  return = pure
  (>>=) = bindRep

instance Representable Quaternion where
  rep f = Quaternion (f e) (f i) (f j) (f k)

instance Foldable Quaternion where
  foldMap f (Quaternion e i j k) = f e `mappend` f i `mappend` f j `mappend` f k
  foldr f z (Quaternion e i j k) = f e (f i (f j (f k z)))

instance Traversable Quaternion where
  traverse f (Quaternion e i j k) = Quaternion <$> f e <*> f i <*> f j <*> f k

instance RealFloat a => Num (Quaternion a) where
  {-# SPECIALIZE instance Num (Quaternion Float) #-}
  {-# SPECIALIZE instance Num (Quaternion Double) #-}
  (+) = liftA2 (+)
  (-) = liftA2 (-)
  negate = fmap negate
  Quaternion a1 b1 c1 d1 * Quaternion a2 b2 c2 d2 = Quaternion
    (a1*a2 - b1*b2 - c1*c2 - d1*d2)
    (a1*b2 + b1*a2 + c1*d2 - d1*c2)
    (a1*c2 - b1*d2 + c1*a2 + d1*b2)
    (a1*d2 + b1*c2 - c1*b2 + d1*a2)
  fromInteger x = Quaternion (fromInteger x) 0 0 0
  abs z = Quaternion (norm z) 0 0 0
  signum q@(Quaternion e i j k)
    | m == 0.0 = q
    | not (isInfinite m || isNaN m) = q ^/ sqrt m
    | any isNaN q = qNaN
    | not (ii || ij || ik) = Quaternion 1 0 0 0
    | not (ie || ij || ik) = Quaternion 0 1 0 0
    | not (ie || ii || ik) = Quaternion 0 0 1 0
    | not (ie || ii || ij) = Quaternion 0 0 0 1
    | otherwise = qNaN
    where
      m = quadrance q
      ie = isInfinite e
      ii = isInfinite i
      ij = isInfinite j
      ik = isInfinite k

  -- abs    = error "Quaternion.abs: use norm"
  -- signum = error "Quaternion.signum: use signorm"

qNaN :: RealFloat a => Quaternion a
qNaN = Quaternion fNaN fNaN fNaN fNaN where fNaN = 0/0

-- {-# RULES "abs/norm" abs x = Quaternion (norm x) 0 0 0 #-}
-- {-# RULES "signum/signorm" signum = signorm #-}

-- this will attempt to rewrite calls to abs to use norm intead when it is available.

instance RealFloat a => Fractional (Quaternion a) where
  {-# SPECIALIZE instance Fractional (Quaternion Float) #-}
  {-# SPECIALIZE instance Fractional (Quaternion Double) #-}
  Quaternion q0 q1 q2 q3 / Quaternion r0 r1 r2 r3 = Quaternion (r0*q0+r1*q1+r2*q2+r3*q3)
                                                               (r0*q1-r1*q0-r2*q3+r3*q2)
                                                               (r0*q2+r1*q3-r2*q0-r3*q1)
                                                               (r0*q3-r1*q2+r2*q1-r3*q0)
                                                 ^/ (r0*r0 + r1*r1 + r2*r2 + r3*r3)
  recip q = q ^/ quadrance q
  fromRational x = Quaternion (fromRational x) 0 0 0

instance Metric Quaternion where
  Quaternion e i j k `dot` Quaternion e' i' j' k' = e*e' + i*i' + j*j' + k*k'

class Complicated t where
  e :: Functor f => (a -> f a) -> t a -> f (t a)
  i :: Functor f => (a -> f a) -> t a -> f (t a)

instance Complicated Complex where
  e f (a :+ b) = (:+ b) <$> f a
  i f (a :+ b) = (a :+) <$> f b

instance Complicated Quaternion where
  e f (Quaternion a b c d) = (\a' -> Quaternion a' b c d) <$> f a
  i f (Quaternion a b c d) = (\b' -> Quaternion a b' c d) <$> f b

class Complicated t => Hamiltonian t where
  j :: Functor f => (a -> f a) -> t a -> f (t a)
  k :: Functor f => (a -> f a) -> t a -> f (t a)
  ijk :: Functor f => (V3 a -> f (V3 a)) -> t a -> f (t a)

instance Hamiltonian Quaternion where
  j f (Quaternion a b c d) = (\c' -> Quaternion a b c' d) <$> f c
  k f (Quaternion a b c d) = Quaternion a b c <$> f d

  ijk f (Quaternion a b c d) = (\(V3 b' c' d') -> Quaternion a b' c' d') <$> f (V3 b c d)

instance Distributive Quaternion where
  distribute f = Quaternion (fmap (^.e) f) (fmap (^.i) f) (fmap (^.j) f) (fmap (^.k) f)

instance (Involutive a, Num a) => Involutive (Quaternion a) where
  conjugate (Quaternion e i j k) = Quaternion (conjugate e) (-i) (-j) (-k)

reimagine :: RealFloat a => a -> a -> Quaternion a -> Quaternion a
reimagine r s (Quaternion _ i j k)
  | isNaN s || isInfinite s = Quaternion r
    (if i /= 0 then i * s else 0)
    (if j /= 0 then j * s else 0)
    (if k /= 0 then k * s else 0)
  | otherwise = Quaternion r (i * s)  (j * s) (k * s)

-- | quadrance of the imaginary component
qi :: Num a => Quaternion a -> a
qi (Quaternion _ i j k) = i*i + j*j + k*k

absi :: Floating a => Quaternion a -> a
absi = sqrt . qi

pow :: RealFloat a => Quaternion a -> a -> Quaternion a
pow q t = exp (t *^ log q)

-- ehh..
instance RealFloat a => Floating (Quaternion a) where
  {-# SPECIALIZE instance Floating (Quaternion Float) #-}
  {-# SPECIALIZE instance Floating (Quaternion Double) #-}
  pi = Quaternion pi 0 0 0
  exp q@(Quaternion e i j k)
    | qiq == 0 = Quaternion (exp e) i j k
    | ai <- sqrt qiq, ee <- exp e = reimagine (ee * cos ai) (ee * (sin ai / ai)) q
    where qiq = qi q
  log q@(Quaternion e i j k)
    | qiq == 0 = if e >= 0 then Quaternion (log e) i j k else Quaternion (log (negate e)) pi j k -- mmm, pi
    | ai <- sqrt qiq, m <- sqrt (e*e + qiq) = reimagine (log m) (atan2 m e / ai) q
    where qiq = qi q
  x ** y = exp (y * log x)
  sqrt q@(Quaternion e i j k)
    | m   == 0 = q
    | qiq == 0 = if e > 0 then Quaternion (sqrt e) 0 0 0 else Quaternion 0 (sqrt (negate e)) 0 0
    | im <- sqrt (0.5*(m-e)) / sqrt qiq = Quaternion (0.5*(m+e)) (i*im) (j*im) (k*im)
    where qiq = qi q
          m = sqrt (e*e + qiq)
  cos q@(Quaternion e i j k)
    | qiq == 0 = Quaternion (cos e) i j k
    | ai <- sqrt qiq = reimagine (cos e * cosh ai) (- sin e * (sinh ai / ai)) q
    where qiq = qi q
  sin q@(Quaternion e i j k)
    | qiq == 0 = Quaternion (sin e) i j k
    | ai <- sqrt qiq = reimagine (sin e * cosh ai) (cos e * (sinh ai / ai)) q
    where qiq = qi q
  tan q@(Quaternion e i j k)
    | qiq == 0 = Quaternion (tan e) i j k
    | ai <- sqrt qiq, ce <- cos e, sai <- sinh ai, d <- ce*ce + sai*sai =
      reimagine (ce * sin e / d) (cosh ai * (sai / ai) / d) q
    where qiq = qi q
  sinh q@(Quaternion e i j k)
    | qiq == 0 = Quaternion (sinh e) i j k
    | ai <- sqrt qiq = reimagine (sinh e * cos ai) (cosh e * (sin ai / ai)) q
    where qiq = qi q
  cosh q@(Quaternion e i j k)
    | qiq == 0 = Quaternion (cosh e) i j k
    | ai <- sqrt qiq = reimagine (cosh e * cos ai) ((sinh e * sin ai) / ai) q
    where qiq = qi q
  tanh q@(Quaternion e i j k)
    | qiq == 0 = Quaternion (tanh e) i j k
    | ai <- sqrt qiq, se <- sinh e, cai <- cos ai, d <- se*se + cai*cai =
      reimagine ((cosh e * se) / d) ((cai * (sin ai / ai)) / d) q
    where qiq = qi q

  asin q = cut asin q
  acos q = cut acos q
  atan q = cut atan q

  asinh q = cut asinh q
  acosh q = cut acosh q
  atanh q = cut atanh q

cut :: RealFloat a => (Complex a -> Complex a) -> Quaternion a -> Quaternion a
cut f q@(Quaternion e _ j k)
  | qiq == 0 = Quaternion a b j k
  | otherwise = reimagine a (b / ai) q
  where qiq = qi q
        ai = sqrt qiq
        a :+ b = f (e :+ ai)

cutWith :: RealFloat a => Complex a -> Quaternion a -> Quaternion a
cutWith (r :+ im) q@(Quaternion e i j k)
  | e /= 0 || qiq == 0 || isNaN qiq || isInfinite qiq = error "bad cut"
  | s <- im / sqrt qiq = Quaternion r (i*s) (j*s) (k*s)
  where qiq = qi q

asinq :: RealFloat a => Quaternion a -> Quaternion a -> Quaternion a
asinq q@(Quaternion e _ _ _) u
  | qiq /= 0.0 || e >= -1 && e <= 1 = asin q
  | otherwise = cutWith (asin (e :+ sqrt qiq)) u
  where qiq = qi q

acosq :: RealFloat a => Quaternion a -> Quaternion a -> Quaternion a
acosq q@(Quaternion e _ _ _) u
  | qiq /= 0.0 || e >= -1 && e <= 1 = acos q
  | otherwise = cutWith (acos (e :+ sqrt qiq)) u
  where qiq = qi q

atanq :: RealFloat a => Quaternion a -> Quaternion a -> Quaternion a
atanq q@(Quaternion e _ _ _) u
  | e /= 0.0 || qiq >= -1 && qiq <= 1 = atan q
  | otherwise = cutWith (atan (e :+ sqrt qiq)) u
  where qiq = qi q

asinhq :: RealFloat a => Quaternion a -> Quaternion a -> Quaternion a
asinhq q@(Quaternion e _ _ _) u
  | e /= 0.0 || qiq >= -1 && qiq <= 1 = asinh q
  | otherwise = cutWith (asinh (e :+ sqrt qiq)) u
  where qiq = qi q

acoshq :: RealFloat a => Quaternion a -> Quaternion a -> Quaternion a
acoshq q@(Quaternion e _ _ _) u
  | qiq /= 0.0 || e >= 1 = asinh q
  | otherwise = cutWith (acosh (e :+ sqrt qiq)) u
  where qiq = qi q

atanhq :: RealFloat a => Quaternion a -> Quaternion a -> Quaternion a
atanhq q@(Quaternion e _ _ _) u
  | qiq /= 0.0 || e > -1 && e < 1 = atanh q
  | otherwise = cutWith (atanh (e :+ sqrt qiq)) u
  where qiq = qi q

slerp :: RealFloat a => Quaternion a -> Quaternion a -> a -> Quaternion a
slerp q p t
  | 1.0 - cosphi < 1e-8 = q
  | phi <- acos cosphi, r <- recip (sin phi)
  = (sin ((1-t)*phi)*r *^ q ^+^ f (sin (t*phi)*r) *^ p) ^/ sin phi
  where
   dqp = dot q p
   (cosphi, f) = if dqp < 0 then (-dqp, negate) else (dqp, id)
{-# SPECIALIZE slerp :: Quaternion Float -> Quaternion Float -> Float -> Quaternion Float #-}
{-# SPECIALIZE slerp :: Quaternion Double -> Quaternion Double -> Double -> Quaternion Double #-}

--slerp :: RealFloat a => Quaternion a -> Quaternion a -> a -> Quaternion a
--slerp q0 q1 = let q10 = q1 / q0 in \t -> pow q10 t * q0

rotate :: Num a => Quaternion a -> V3 a -> V3 a
rotate (Quaternion a' b c d) (V3 x y z) = V3
  (2*((t8+t10)*x+(t6- t4)*y+(t3+t7)*z)+x)
  (2*((t4+ t6)*y+(t5+t10)*y+(t9-t2)*z)+y)
  (2*((t7- t3)*z+(t2+ t9)*z+(t5+t8)*z)+z)
  where
    a = -a'
    t2 = a*b
    t3 = a*c
    t4 = a*d
    t5 = -b*b
    t6 = b*c
    t7 = b*d
    t8 = -c*c
    t9 = c*d
    t10 = -d*d
{-# SPECIALIZE rotate :: Quaternion Float -> V3 Float -> V3 Float #-}
{-# SPECIALIZE rotate :: Quaternion Double -> V3 Double -> V3 Double #-}

instance (RealFloat a, Epsilon a) => Epsilon (Quaternion a) where
  nearZero = nearZero . quadrance


{-# LANGUAGE ApplicativeDo              #-}
{-# LANGUAGE DataKinds                  #-}
{-# LANGUAGE DeriveFoldable             #-}
{-# LANGUAGE DeriveFunctor              #-}
{-# LANGUAGE DeriveTraversable          #-}
{-# LANGUAGE FlexibleContexts           #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE KindSignatures             #-}
{-# LANGUAGE OverloadedStrings          #-}
{-# LANGUAGE QuantifiedConstraints      #-}
{-# LANGUAGE RankNTypes                 #-}
{-# LANGUAGE MultiParamTypeClasses #-}
module Main where
import           Control.Monad
import           Control.Monad.State
import           Data.Foldable
import           Data.List                   hiding (filter)
import           Data.RVar
import           Data.Random
import qualified Data.Text                   as T
import           Data.Time
import           Graphics.Vega.VegaLite      hiding (repeat, sample)
import qualified Numeric.LinearAlgebra       as V
import           Prelude                     hiding (filter)
import           Q.ContingentClaim.Options
import           Q.MonteCarlo
import           Q.Options.BlackScholes
import           Q.Stats.TimeSeries
import           Q.Stochastic.Discretize
import           Q.Stochastic.Process
import           Q.Time
import           Q.Time.DayCounter
import           Q.Types
import qualified Q.Util.File                 as QF
import           Statistics.Sample.Histogram

import qualified Data.Time                   as Time

dt :: Double
dt = 1/365
gbm = GeometricBrownian 0.0 0.15

plotPayouts :: IO ()
plotPayouts = do
  let strike      = 100.0
      lowStrike   = 95.0
      highStrike  = 105.0
      expiry      = fromGregorian 2015 3 31
      spots       = [90, 90.01..110]
      calls       = map (vanillaPayout Call strike) spots
      puts        = map (vanillaPayout Put strike) spots
      callSpreads = map (spreadPayout Call lowStrike highStrike) spots
      putSpreads  = map (spreadPayout Put lowStrike highStrike) spots
      straddles   = map (straddlePayout strike) spots

      payouts = dataFromColumns []
        . dataColumn "Spot" (Numbers spots)
        . dataColumn "Call" (Numbers calls)
        . dataColumn "Put"  (Numbers puts)
        . dataColumn "Call Spread"  (Numbers callSpreads)
        . dataColumn "Put Spread"  (Numbers putSpreads)
        . dataColumn "Straddle" (Numbers straddles)

      enc = encoding
        . position X [PName "Spot", PmType Quantitative]
      encVanillaCall = enc
        . position Y [PName "Call", PmType Quantitative]
      encVanillaPut  = enc
        . position Y [PName "Put", PmType Quantitative]
      encCallSpread = enc
        . position Y [PName "Call Spread", PmType Quantitative]
      encPutSpread = enc
        . position Y [PName "Put Spread", PmType Quantitative]
      encStraddle = enc
        . position Y [PName "Straddle", PmType Quantitative]

      vl = toVegaLite [payouts [], vlConcat [ asSpec [encVanillaCall [], mark Line []]
                                            , asSpec [encVanillaPut  [], mark Line []]
                                            , asSpec [encCallSpread  [], mark Line []]
                                            , asSpec [encPutSpread   [], mark Line []]
                                            , asSpec [encStraddle    [], mark Line []]
                                            ]
                      ]
  toHtmlFile "/tmp/payouts.html" vl

-- Option delta is the sensitivity of the option to the underlying:
-- \(dV/dS\).
--
-- for a fixed time and vol what is the delta per relative strike.
-- How does delta look as a function of time
-- How does delta look as a function of vol?

-- How does delta look like as a function of the spot?
-- Plot call delta against log relative strike.
-- Plot put  delta against log relative strike.
studyDelta :: IO ()
studyDelta = do
  let today           = fromGregorian 2015 3 31
      expiry          = YearFrac 1
      vol             = 0.6
      spot            =  100
      strikes         = map Strike [80, 81..120]
      vols            = [0.01,0.06..1]
      expiries        = map YearFrac [0.1,0.2..2]
      relativeStrikes = [log (k/(atmf bs expiry)) | k <- strikes]
      r               = 0.01
      bs              = BlackScholes spot r vol
      delta cp k      = vDelta (euOption bs expiry cp k)
      pv    cp k      = vPremium (euOption bs expiry cp k)
      row (vol,k,t)     = [show today,
                         show t,
                         show $log (k / atmf bs t),
                         show r,
                         show vol,
                         show $ pv Call k,
                         show $ pv Put k,
                         show $ delta Call k,
                         show $ delta Put k
                        ]
      header           = ["VD","Expiry","RelativeStrike", "IR", "Vol", "cPV", "pPV", "cDelta", "pDelta"]

  QF.write (map row [(vol, k, t) | vol <- vols, k <- strikes, t <- expiries]) header "/tmp/data/delta.csv"

  let mydata = dataFromUrl "delta.csv" [Parse [ ( "Expiry", FoDate "%Y-%m-%d" ), ( "VD", FoDate "%Y-%m-%d" ) ]]

  -- !This is \(dV/dS\) for different strikes, same expiry, interest rate,
  let l1 = toVegaLite [mydata , layers, width 800, height 600, res, trans] where
      calls = encoding
        . position X [PName "RelativeStrike"    , PmType Quantitative]
        . position Y [PName "cDelta" , PmType Quantitative, PTitle "Calls", PAxis [AxTitleColor "green"]]
      puts  = encoding
        . position X [PName "RelativeStrike"    , PmType Quantitative]
        . position Y [PName "pDelta"  , PmType Quantitative, PTitle "Puts", PAxis [AxTitleColor "red"]]
      layers = layer [ asSpec [calls [], mark Line [MColor "green"]]
                     , asSpec [puts  [], mark Line [MColor "red"]]
                     ]
      trans = transform . filter (FExpr "datum.Expiry == datetime(2015,3,15)") $ []
      res = resolve . resolution (RScale [(ChY, Independent)]) $ []

  toHtmlFile "/tmp/data/deltas.html" l1

-- | Plot call option premiums, deltas, payoffs across a differnt
-- range of strikes.
plotCallOptions :: IO ()
plotCallOptions = do
  let today   = fromGregorian 2015 3 31
      expiry  = YearFrac 1.0
      spot    = 100
      r       = 0
      vol     = 0.01
      bs      = BlackScholes spot r vol
      k       = atmf bs expiry
      cp      = Call
  let strikes = [95, 95.01..105]
      values  = [] -- map (vPremium . euOption cp bs expiry) strikes
      deltas  = [] --map (vDelta   . euOption cp bs expiry) strikes
      pvs = dataFromColumns []
        . dataColumn "Strikes" (Numbers strikes)
        . dataColumn "Premiums"     (Numbers values)
        . dataColumn "Deltas"  (Numbers deltas)
  let encPvs = encoding
        . position X [PName "Strikes", PmType Quantitative]
        . position Y [PName "Premiums"    , PmType Quantitative]
      encDeltas = encoding
        . position X [PName "Strikes", PmType Quantitative]
        . position Y [PName "Deltas", PmType Quantitative]
      layers = layer [ asSpec [ encPvs     [], mark Line [MColor "blue"]]
                     , asSpec [ encDeltas  [], mark Line [MColor "red"]]
                     ]
      res = resolve . resolution (RScale [(ChY, Independent)])
      vl     = toVegaLite [pvs [], layers, width 800, res [], height 400]
  toHtmlFile "/tmp/callOptions.html" vl

testPlot = do
  let dvals = dataSequenceAs 0 6.28 0.1 "x"
      trans = transform
              . calculateAs "sin(datum.x)" "sinX"
              . calculateAs "cos(datum.x)" "cosX"
      enc = encoding
            . position X [PName "x", PmType Quantitative]
      encCos = enc . position Y [PName "cosX", PmType Quantitative]
      encSin = enc . position Y [PName "sinX", PmType Quantitative]

      vl = toVegaLite [ dvals
                      , trans []
                      , vlConcat [ asSpec [encCos [], mark Line []]
                                 , asSpec [encSin [], mark Line []]
                                 ]
                      ]
  toHtmlFile "/tmp/test.html" vl

plotTrajectories :: IO ()
plotTrajectories = do
  let dt = 1/365
      today = fromGregorian 2015 3 31
      dates = [d | d <- fmap (flip addDays today) [0, 365]]
      dts = tail $ map (\d -> dcYearFraction ThirtyUSA today d) dates
      n = 64000
      ids = mconcat $ fmap (replicate $ length dates) [T.pack $ show id | id <- [1..n]]
  trajectories <- sample $ trajectories n (Euler dt) gbm  (100::Double) dts (repeat stdNormal)
  let trajectories' = dataFromColumns [ Parse [ ( "Date", FoDate "%Y%m%d" ) ] ]
                      . dataColumn "Date" (Strings  (mconcat (replicate n (map dayToString dates))))
                      . dataColumn "St" (Numbers $ mconcat trajectories)
                      . dataColumn "ID" (Strings ids)
      calls = mark Line
      enc = encoding
            . position X [ PName "Date", PmType Temporal]
            . position Y [ PName "St", PmType Quantitative, PScale [SDomain (DNumbers [50, 150])] ]
            . color [ MName "ID", MmType Nominal, MLegend []]
      vl = toVegaLite [trajectories' [], enc [], mark Line [], width 1000, height 600]
  toHtmlFile "/tmp/a.html" vl

main  = undefined


{- MDP.hs       Code for CS 533, Mini-Project 2.
 -
 - Copyright © 2012 Loren B. Davis.  ALL RIGHTS RESERVED.  I hereby release
 - this code under the MIT License: (http://www.opensource.org/licenses/MIT).
 - Since this is an assigned exercise, remember your academic integrity!
 -}

module MDP {- ( MarkovDP, State, Action, policyIterate, showUpdates, test4x3,
             showT, decodeT, whenToPass, whenToPark, parkingProblem,
             ParkingState, encodeState, decodeState ) -}
  where
import Data.Array.Unboxed (UArray, (!), accumArray, assocs, listArray)
import Data.List (maximumBy, (\\))
import Data.Ord ()
import Numeric.LinearAlgebra (Matrix, Vector, build, linearSolveSVD, toLists)
-- import Numeric.LinearAlgebra hiding (i)
-- import Numeric.LinearAlgebra.Algorithms hiding (i)
import Prelude hiding(pi)
{- This program uses the hmatrix package (cabal install hmatrix) rather than
 - attempt to reinvent the wheel.  You might also need to install libgsl,
 - liblapack, and possibly MKL or ATLAS.
 -}

type State = Int
type Action = Int
type MarkovDP = ( Int,
                  Int,
                  UArray State Double,
                  UArray (State,Action,State) Double
                )
{- n = The number of states.  S = {0,...,n-1}
 - m = The number of actions. A = {0,...,m-1}
 - r = An array of size n representing the reward function.  r[s] is the
 -     reward associated with state s.
 - t = A three-dimensional array representing the transition function, such
 -     that T(s,a,s') = t ! (s,a,s').  Note that the values are between 0 and
 -     1, each row sums to 1, and that the matrices will typically be sparse.
 -}
type PIUpdate = [(State, Action)]
{- The update of each step of the policy iteration algorithm is the list
 - of pairs (s,a) such that the new policy will take action a in state s.
 -}

policyIterate :: MarkovDP ->
                 Double ->
                 (UArray State Double, UArray State Action, [String])
{- Parameters:
 - process  = The Markov decision process to solve, in the form
 -        (|S|,|A|,R,T)
 - gamma    = The discount parameter, between 0 and 1.
 -
 - Returns (v, pi, updates) where
 - v        = A size-n array representing the value function.  The expect-
 -        ed value of being in state s is v[s].
 - pi       = A size-n array representing the policy.  The action that pi
 -        will take in state s is pi[s].
 - updates  = A human-readable list of update strings, one per step of the
 -        iteration.
 -
 - This implementation always starts with a policy of taking the first action
 - in any situation.  This is not really correct (the algorithm in the text
 - says to try a "random" policy), but saves me from dealing with stateful
 - computation in this program.  To specify a different initial policy, call
 - policyIterate1.  It does not attempt any asynchronous policy iteration.
 -}
policyIterate (n, m, r, t) gamma = (v', pi'', translatePIUpdates updates'' )
  where
    ( v', pi'', updates'' ) = seq ((n,m,r,t), gamma, naive)
                                  (policyIterate1 naive [])
    naive = (initialPolicy n)
    policyIterate1 :: UArray State Action ->
                      [PIUpdate] ->
                      (UArray State Double, UArray State Action, [PIUpdate])
{- Parameters:
 - pi       = The initial policy to attempt to improve.
 - updates  = The previous iterations' updates.
 -
 - Returns (v, pi) where
 - v    = A size-n array representing the value function.  The expected value
 -     of being in state s is v[s].
 - pi'  = A size-n array representing the policy.  The action that pi' will
 -       take in state s is pi[s].
 - updates' = The list of updates for all iterations of the algorithm.  When
 -            there are no further updates, the algorithm terminates.
 -}
    policyIterate1 pi updates | diff == [] = (v, pi', updates')
                              | otherwise  = policyIterate1 pi' $! updates'
      where
        states :: [State]
        states = [0..n-1]
        actions :: [Action]
        actions = [0..m-1]
        v :: UArray State Double
{- The updated value function.  This is based on eq. 17.10 in (Russell &
 - Norvig 2010), rather than the modified-policy-iteration algorithm in fig.
 - 17.7.  It uses LAPACK to solve the system of linear equations in O(n^3)
 - time.
 -}
        v = let
              updateHelper :: Double -> Double -> Double
              updateHelper w z = gamma * (t!(i,pi!i,j)) -
                                           (if i == j then 1.0 else 0.0)
                where
                  i :: Int
                  i = truncate w
                  j :: Int
                  j = truncate z
              a :: Matrix Double
              a = build (n,n) updateHelper
              b :: Matrix Double
              b = build (n,1) (\y _ -> let i = truncate y :: Int in - (r!i))
              x :: Matrix Double
              x = linearSolveSVD a b
            in
              listArray (0,n-1) (concat . toLists $ x)
        pi' :: UArray Int Action
{- The updated policy pi' is the best action to take in any state (determined
 - by one-step lookahead).  Q.v. eq. 17.4 in (Russell & Norvig 2010).
 -}
        pi' = seq v (listArray (0, n-1)
                        (map (\i -> maximumBy (argmax i) actions)
                             states))
        argmax :: State -> Action -> Action -> Ordering
{- A helper function that compares the expected value of actions a1 and a2
 - from state s0.
 -}
        argmax s0 a1 a2 = let
                            u1 = sum .
                                 map (\s' -> t!(s0,a1,s') * v!s') $
                                 states
                            u2 = sum .
                                 map (\s' -> t!(s0,a2,s') * v!s') $
                                 states
                          in compare u1 u2
        diff :: PIUpdate
        diff = seq pi' ((assocs pi') \\ (assocs pi))
        updates' :: [PIUpdate]
        updates' = seq diff (updates ++ [diff])

initialPolicy :: Int -> UArray Int Action
{- A policy for a MDP with n states that always takes action zero.  This is
 - not the same as the "random" policy I'm supposed to use for the algorithm,
 - but is simpler to write in this language.
 -}
initialPolicy n = listArray (0, n-1) (take n (repeat 0))

initialVF :: Int -> UArray Int Double
{- A value function initialized to n zeroes.  Unlike the above, this is corr-
 - ect.  (But unused in policy iteration.)
 -}
initialVF n = listArray (0, n-1) (take n (repeat 0.0))

translatePIUpdates :: [PIUpdate] -> [String]
translatePIUpdates u =
  map (\i -> "Step " ++ show (i+1) ++ ": " ++ translatePIUpdate (u !! i))
      [0 .. (length u) - 1]

translatePIUpdate :: PIUpdate -> String
translatePIUpdate [] = "Done."
translatePIUpdate [(s,a)] = "take action " ++
                            show a ++
                            " in state " ++
                            show s ++
                            "."
translatePIUpdate ((s,a):xs) = "take action " ++
                               show a ++
                               " in state " ++
                               show s ++
                               "; " ++
                               translatePIUpdate xs

-- Display the returned update messages:
showUpdates :: (UArray State Double, UArray State Action, [String]) ->
               IO ()
showUpdates (_,_,u) = putStrLn . unlines $ u

-- A useful means of displaying only the nonzero entries of the sparse array
-- t:
showT :: MarkovDP -> [((State,Action,State),Double)]
showT (_,_,_,t) = filter ((> 0.0) . snd) (assocs t)

-- The parking-lot world:

{- The decoded values of a state of the parking-lot world.  Note that Exit
 - in this implementation is a terminal state that the park action automat-
 - ically reaches, not a third action.  Crash is another world state, always
 - and only reached by attempting to park in a filled spot.
 -
 - A i True is the state in which the agent is in front of row A, space i,
 - which is filled.  A i False is similar, but the spot is empty; likewise
 - for row B.  I take advantage of the symmetry between parking in row A and
 - in row B to define P i as the state of being parked i spaces from the
 - store.  This reduces the size of the state space from 8n to 5n+1. The
 - range of i is 1 to n.
 -}
data ParkingState = A Int Bool | B Int Bool | P Int | Exit | Crash
  deriving (Eq, Show)

encodeState :: Int -> ParkingState -> State
{- Converts a ParkingState into an index suitable for the PolicyIterate algo-
 - rithm.  The encoding depends on the number of rows, n.
 -}
encodeState _ Exit = 0
encodeState _ Crash = 1
encodeState _ (P i) = 1 + i
encodeState n (A i True) = 1 + n + i
encodeState n (A i False) = 1 + 2*n + i
encodeState n (B i True) = 1 + 3*n + i
encodeState n (B i False) = 1 + 4*n + i

decodeState :: Int -> State -> ParkingState
{- The inverse of the mapping above.
 -}
decodeState _ 0 = Exit
decodeState _ 1 = Crash
decodeState n m | m <= 1 + n = P (m-1)
                | m <= 1 + 2*n = A (m-n-1) True
                | m <= 1 + 3*n = A (m-2*n-1) False
                | m <= 1 + 4*n = B (m-3*n-1) True
                | m <= 1 + 5*n = B (m-4*n-1) False
                | otherwise = error $ "Invalid state " ++ show m ++ "."

-- Decode an association from t:
decodeT' :: Int ->
           ((State,Action,State),Double) ->
           ((ParkingState,String,ParkingState),Double)
decodeT' c ((s,a,s'),p) = ((decodeState c s,
                            if (a == 1) then "Park" else "Move",
                            decodeState c s'),
                           p)

-- Show the condensed, human-readable transition function.
decodeT :: Int ->
           [((State,Action,State),Double)] ->
           [((ParkingState,String,ParkingState),Double)]
decodeT c = map (decodeT' c)

-- Display the states in which the policy tells the agent to keep looking.
whenToPass :: Int ->
              (UArray State Double, UArray State Action, [String]) ->
              [ParkingState]
whenToPass c (_,pi,_) = map (decodeState c) .
                        filter (> 1+c) .
                        map fst .
                        filter ((==) 0 . snd) .
                        assocs $
                        pi

-- And when to park.
whenToPark :: Int ->
              (UArray State Double, UArray State Action, [String]) ->
              [ParkingState]
whenToPark c (_,pi,_) = map (decodeState c) .
                        filter (> 1+c) .
                        map fst .
                        filter ((==) 1 . snd) .
                        assocs $
                        pi

parkingProblem :: Int ->
                  Double ->
                  Double ->
                  Double ->
                  Double ->
                  Double ->
                  Double ->
                  MarkovDP
{- Generates a MDP representing a parking problem with the given parameters:
 -
 - c = The number of parking spaces per row.
 - alpha = The probability that a handicap spot will be empty.  Should be be-
 -         low, but close to, 1.0.
 - beta = The penalty for parking in a handicap spot.  Should be a large neg-
 -        ative number.
 - delta = A parameter representing the rate at which a shopper loses utility
 -         for distance.  I represent the utility of parking as inversely
 -         proportional to the square of the distance.  Should be a positive
 -         number.
 - iota = The rate of utility loss for driving around.  Should be a
 -        small positive number.
 - kappa = The penalty for crashing into another car.  Should be a larger
 -         negative number than beta.
 - lambda = A parameter determining the probability that a non-handicap spot
 -          will be empty.  I represent this as a Poisson-like distribution, as
 -          we can model the spaces as filling in from the closest on out, with
 -          cars entering as independent events and leaving after a certain
 -          period of time.  Therefore, this parameter should be a positive
 -          number, equal to the expected number of cars parked in the lot.
 -}
parkingProblem c alpha beta delta iota kappa lambda = let
  n = c*5 + 2
  r = listArray (0,n-1) ([0, kappa, beta] ++
                         [delta/(fromIntegral i) | i <- [2..c]] ++
                         (repeat (-iota)))
  t = accumArray (+) 0 ((0,0,0),(n-1,1,n-1)) (
-- Exit is a terminal state.
    [((0,0,0),1.0), ((0,1,0),1.0)] ++
-- Crashing or parking always obviates any further reward.
    [((i,0,0),1.0) | i <- [1..c+1]] ++ [((i,1,0),1.0) | i <- [1..c+1]] ++
-- Attempting to park in any occupied space always crashes the car.
    [((encodeState c (A i True),1,1),1.0) | i <- [1..c]] ++
    [((encodeState c (B i True),1,1),1.0) | i <- [1..c]] ++
-- However, parking in a free space succeeds.
    [((encodeState c (A i False),1,(encodeState c (P i))),1.0) | i <- [1..c]] ++
    [((encodeState c (B i False),1,(encodeState c (P i))),1.0) | i <- [1..c]] ++
-- An agent at A[1] can only move to B[1], which could be free or not.  The
-- probability that it is free is alpha.
    [((encodeState c (A 1 False),0,encodeState c (B 1 True)),1.0-alpha),
     ((encodeState c (A 1 True),0,encodeState c (B 1 True)),1.0-alpha),
     ((encodeState c (A 1 False),0,encodeState c (B 1 False)),alpha),
     ((encodeState c (A 1 True),0,encodeState c (B 1 False)),alpha)] ++
-- An agent at A[2] can only move to A[1], which again is free with probabil-
-- ity alpha.
    [((encodeState c (A 2 False),0,encodeState c (A 1 True)),1.0-alpha),
     ((encodeState c (A 2 True),0,encodeState c (A 1 True)),1.0-alpha),
     ((encodeState c (A 2 False),0,encodeState c (A 1 False)),alpha),
     ((encodeState c (A 2 True),0,encodeState c (A 1 False)),alpha)] ++
-- An agent at any other A[i] can only move to A[i-1].
    [((encodeState c (A i False),0,encodeState c (A (i-1) False)),quasiPoisson lambda (i-1)) | i <- [3..c]] ++
    [((encodeState c (A i True),0,encodeState c (A (i-1) False)),quasiPoisson lambda (i-1)) | i <- [3..c]] ++
    [((encodeState c (A i False),0,encodeState c (A (i-1) True)),1.0 - quasiPoisson lambda (i-1)) | i <- [3..c]] ++
    [((encodeState c (A i True),0,encodeState c (A (i-1) True)),1.0 - quasiPoisson lambda (i-1)) | i <- [3..c]] ++
-- An agent at B[i], where i < c, can only move to B[i+1].
    [((encodeState c (B i False),0,encodeState c (B (i+1) False)),quasiPoisson lambda (i+1)) | i <- [1..(c-1)]] ++
    [((encodeState c (B i True),0,encodeState c (B (i+1) False)),quasiPoisson lambda (i+1)) | i <- [1..(c-1)]] ++
    [((encodeState c (B i False),0,encodeState c (B (i+1) True)),1.0-quasiPoisson lambda (i+1)) | i <- [1..(c-1)]] ++
    [((encodeState c (B i True),0,encodeState c (B (i+1) True)),1.0-quasiPoisson lambda (i+1)) | i <- [1..(c-1)]] ++
-- Finally, an agent at B[c] can only move to A[c].
    [((encodeState c (B c False),0,encodeState c (A c True)),1.0-quasiPoisson lambda c),
     ((encodeState c (B c True),0,encodeState c (A c True)),1.0-quasiPoisson lambda c),
     ((encodeState c (B c False),0,encodeState c (A c False)),quasiPoisson lambda c),
     ((encodeState c (B c True),0,encodeState c (A c False)),quasiPoisson lambda c)]
       )
  in (n, 2, r, t)

poisson :: Double -> Int -> Double
-- The pmf of a Poisson distribution.
poisson lambda i = (exp (-lambda)) *
                   (lambda^i) /
                   (fromIntegral (factorial i))

poissonCMF :: Double -> Int -> Double
-- The memoized cmf of a Poisson distribution.
poissonCMF lambda = ((scanl1 (+) (map (poisson lambda) [0..])) !!)

quasiPoisson :: Double -> Int -> Double
{- Take into account that there are two rows to fill.  The parameter gamma
 - should be the expected number of parked cars; the parameter i should be
 - at least 2.
 -}
quasiPoisson lambda i = ((poissonCMF lambda (i*2-4)) +
                         (poissonCMF lambda (i*2-3)))
                        / 2.0

factorial :: Int -> Integer
-- A memoized factorial function.
factorial = ((scanl (*) 1 [1..]) !!)

-- Test data:

-- The 4x3 world in fig. 17.3 of (Russell & Norvig 2010):

test4x3 :: MarkovDP
test4x3 = ( 12, -- The 11 states in the example, plus a terminal state 0.
            4,  -- Left, right, up, down
            listArray (0,11)
                      [0,
                       -0.04, -0.04, -0.04, 1.0,
                       -0.04,        -0.04, -1.0,
                       -0.04, -0.04, -0.04, -0.04],
            accumArray (+) 0.0 ((0,0,0),(11,3,11)) [
 -- The synthetic terminal state:
 ((0,0,0),1.0),((0,1,0),1.0),((0,2,0),1.0),((0,3,0),1.0),
 -- Position (1,1):
 ((1,0,1),1.0),((1,1,2),1.0),((1,2,1),1.0),((1,3,5),1.0),
 -- Position (1,2):
 ((2,0,1),1.0),((2,1,3),1.0),((2,2,2),1.0),((2,3,2),1.0),
 -- Position (1,3):
 ((3,0,2),1.0),((3,1,4),1.0),((3,2,3),1.0),((3,3,6),1.0),
 -- Position (1,4):
 ((4,0,0),1.0),((4,1,0),1.0),((4,2,0),1.0),((4,3,0),1.0),
 -- Position (2,1):
 ((5,0,5),1.0),((5,1,5),1.0),((5,2,1),1.0),((5,3,8),1.0),
 -- Position (2,2) does not exist.
 -- Position (2,3):
 ((6,0,6),1.0),((6,1,7),1.0),((6,2,3),1.0),((6,3,10),1.0),
 -- Position (2,4):
 ((7,0,0),1.0),((7,1,0),1.0),((7,2,0),1.0),((7,3,0),1.0),
 -- Position (3,1):
 ((8,0,8),1.0),((8,1,9),1.0),((8,2,5),1.0),((8,3,8),1.0),
 -- Position (3,2):
 ((9,0,8),1.0),((9,1,10),1.0),((9,2,9),1.0),((9,3,9),1.0),
 -- Position (3,3):
 ((10,0,9),1.0),((10,1,11),1.0),((10,2,6),1.0),((10,3,10),1.0),
 -- Position (3,4):
 ((11,0,10),1.0),((11,1,11),1.0),((11,2,7),1.0),((11,3,11),1.0)
            ]
          )

-- The three-state world in problem 17.10 in (Russell & Norvig 2010):
test3 :: MarkovDP
test3 = ( 3, -- There are three states in the example.  I rename 3 to 0.
          2, -- and two actions, a and b.
          listArray (0,2) [0,-1,-2],
          accumArray (+) 0.0 ((0,0,0),(2,1,2)) [
 -- State 0 is a terminal state.
 ((0,0,0),1.0), ((0,1,0),1.0),
 -- In state 1, action a goes to state 2 (p=0.8) or state 1 (p=0.2)
 -- Action b goes to state 0 (p=0.1) or state 1 (p=0.9)
 ((1,0,2),0.8), ((1,0,1),0.2), ((1,1,0),0.1), ((1,1,1),0.9),
 -- In state 2, action a goes to state 1 (p=0.8) or state 2 (p=0.2)
 -- Action b goes to state 0 (p=0.1) or state 1 (p=0.9)
 ((2,0,1),0.8), ((2,0,2),0.2), ((2,1,0),0.1), ((2,1,2),0.9)
          ]
        )

-- The 101 x 3 world from exercise 17.8 in (Russell & Norvig 2010):
{- This example is useful because I’ve already solved the critical values of
 - gamma for it, and because it contains a larger number of states.
 -}
test101 :: MarkovDP
test101 = ( 204, -- The start state, an immediate reward state on top followed
                 -- by 100 penalties, an immediate penalty state on the bottom
                 -- followed by 100 rewards, and a synthetic terminal state.
            2,   -- Up or down at the start.  Both move the agent right any-
                 -- where else.
            listArray (0,203)
  ((0:50.0:(take 100 (repeat (-1.0))))++(-50.0:(take 100 (repeat 1.0)))++[0]),
            accumArray (+) 0.0 ((0,0,0),(203,1,203))
-- In the start state, up goes to state 1, down to state 102.
  ([((0,0,1),1.0), ((0,1,102),1.0)] ++
-- In states 1-100, any action moves the agent to the next consecutive state.
  [((i,0,i+1), 1.0) | i <- [1..100]] ++ [((i,1,i+1), 1.0) | i <- [1..100]] ++
-- State 101 always sends the agent to terminal state 203.
  [((101,0,203),1.0),((101,1,203),1.0)] ++
-- States 102-202 also send the agent to the next consecutive state:
  [((i,0,i+1), 1.0) | i <- [102..202]] ++
  [((i,1,i+1), 1.0) | i <- [102..202]] ++
-- Finally, state 203 is terminal:
  [((203,0,203),1.0),((203,1,203),1.0)])
          )


import Data.Map.Strict ((!))
import Numeric.LinearAlgebra hiding ((!))
import LeastSquares

leastSquaresEstimate :: Vector R
leastSquaresEstimate = result ! "x"
  where result = minimize (SumSquares x) [c * x :==: d]
        c = Mat $ (2><3) [3.0, 7.0, -1.5, 2.8, 10.0, -5.3]
        d = Vec $ vector [-1.0, 3.5]
        x = Var "x" 3


module Fractals (pMandelbrot, pJulia) where

import           Data.Complex

data PointResult = Convergent | Divergent Int (Complex Double)

-- Calculates a point in the Mandelbrot set
pMandelbrot :: Bool -> Double -> Int -> Double -> Complex Double -> Maybe Double
pMandelbrot aa pixelSize maxIterations power point =
   fPoint aa pixelSize maxIterations power (\d z -> z ** (d :+ 0) + point) point

-- Calculates a point in the Julia set
pJulia :: Bool -> Double -> Int -> Double -> Complex Double -> Complex Double -> Maybe Double
pJulia aa pixelSize maxIterations power c point =
   fPoint aa pixelSize maxIterations power (\d z -> z ** (d :+ 0) + c) point

fPoint
   :: Bool
   -> Double
   -> Int
   -> Double
   -> (Double -> Complex Double -> Complex Double)
   -> Complex Double
   -> Maybe Double
fPoint bAAEnabled pixelSize maxIterations power f z0 = if not bAAEnabled
   then normalize $ iterPoint 0 z0
   else
      (\(sum, count', isConvergent) ->
            if isConvergent then Nothing else Just $ (sum / (numberOfSubpoints - count'))
         )
         $ foldl
              (\(sum, countConv, isConvergent) y -> if countConv > maxConvergent
                 then (sum, countConv, True)
                 else case normalize . iterPoint 0 $ z0 + y of
                    Nothing -> (sum, countConv + 1, False)
                    Just z  -> (sum + z, countConv, False)
              )
              (0, 0, False)
              [ (-size') :+ (-size'), (-size') :+ 0, (-size') :+ size'
              , 0 :+ (-size'),        0 :+ 0,        0 :+ size'
              , size' :+ (-size'),    size' :+ 0,    size' :+ size'
              ]
 where
  maxMagnitude = 50
  -- Main iteration function
  iterPoint n z | n > maxIterations          = Convergent
                | magnitude z > maxMagnitude = Divergent n z
                | otherwise                  = iterPoint (n + 1) (f power z)
  -- Prevents the aliasing present in generators which color based purely on escape time
  normalize Convergent      = Nothing
  normalize (Divergent 0 _) = Just 0
  normalize (Divergent iterations final) =
     Just . (fromIntegral iterations + 1 -) . logBase power . logBase maxMagnitude . magnitude $ final
  size'             = pixelSize / 3
  numberOfSubpoints = 9
  maxConvergent     = 4


-- |
-- Module      :  Statistics.Distribution.Dirichlet
-- Description :  Multivariate Dirichlet distribution
-- Copyright   :  (c) Dominik Schrempf, 2021
-- License     :  GPL-3.0-or-later
--
-- Maintainer  :  dominik.schrempf@gmail.com
-- Stability   :  unstable
-- Portability :  portable
--
-- Creation date: Tue Oct 20 10:10:39 2020.
module Statistics.Distribution.Dirichlet
  ( -- * Dirichlet distribution
    DirichletDistribution (ddGetParameters),
    dirichletDistribution,
    dirichletDensity,
    dirichletSample,

    -- * Symmetric Dirichlet distribution
    DirichletDistributionSymmetric (ddSymGetParameter),
    dirichletDistributionSymmetric,
    dirichletDensitySymmetric,
    dirichletSampleSymmetric,
  )
where

import Control.Monad.Primitive
import qualified Data.Vector.Unboxed as V
import Numeric.Log
import Numeric.SpecFunctions
import System.Random.MWC
import System.Random.MWC.Distributions

-- | The [Dirichlet distribution](https://en.wikipedia.org/wiki/Dirichlet_distribution).
data DirichletDistribution = DirichletDistribution
  { ddGetParameters :: V.Vector Double,
    _getDimension :: Int,
    _getNormConst :: Log Double
  }
  deriving (Eq, Show)

-- Check if vector is strictly positive.
isNegativeOrZero :: V.Vector Double -> Bool
isNegativeOrZero = V.any (<= 0)

-- Inverse multivariate beta function. Does not check if parameters are valid!
invBeta :: V.Vector Double -> Log Double
invBeta v = Exp $ logDenominator - logNominator
  where
    logNominator = V.sum $ V.map logGamma v
    logDenominator = logGamma (V.sum v)

-- | Create a Dirichlet distribution from the given parameter vector.
--
-- Return Left if:
--
-- - The parameter vector has less then two elements.
--
-- - One or more parameters are negative or zero.
dirichletDistribution :: V.Vector Double -> Either String DirichletDistribution
dirichletDistribution v
  | V.length v < 2 =
    Left "dirichletDistribution: Parameter vector is too short."
  | isNegativeOrZero v =
    Left "dirichletDistribution: One or more parameters are negative or zero."
  | otherwise = Right $ DirichletDistribution v (V.length v) (invBeta v)

-- Tolerance.
eps :: Double
eps = 1e-14

-- Check if vector is normalized with tolerance 'eps'.
isNormalized :: V.Vector Double -> Bool
isNormalized v
  | abs (V.sum v - 1.0) > eps = False
  | otherwise = True

-- | Density of the Dirichlet distribution evaluated at a given value vector.
--
-- Return 0 if:
--
-- - The value vector has a different length than the parameter vector.
--
-- - The value vector has elements being negative or zero.
--
-- - The value vector does not sum to 1.0 (with tolerance @eps = 1e-14@).
dirichletDensity :: DirichletDistribution -> V.Vector Double -> Log Double
dirichletDensity (DirichletDistribution as k c) xs
  | k /= V.length xs = 0
  | isNegativeOrZero xs = 0
  | not (isNormalized xs) = 0
  | otherwise = c * Exp logXsPow
  where
    logXsPow = V.sum $ V.zipWith (\a x -> log $ x ** (a - 1.0)) as xs

-- | Sample a value vector from the Dirichlet distribution.
dirichletSample :: PrimMonad m => DirichletDistribution -> Gen (PrimState m) -> m (V.Vector Double)
dirichletSample (DirichletDistribution as _ _) g = do
  ys <- V.mapM (\a -> gamma a 1.0 g) as
  let s = V.sum ys
  return $ V.map (/ s) ys

-- | See 'DirichletDistribution' but with parameter vector @replicate DIM VAL@.
data DirichletDistributionSymmetric = DirichletDistributionSymmetric
  { ddSymGetParameter :: Double,
    _symGetDimension :: Int,
    _symGetNormConst :: Log Double
  }
  deriving (Eq, Show)

-- Inverse multivariate beta function. Does not check if parameters are valid!
invBetaSym :: Int -> Double -> Log Double
invBetaSym k a = Exp $ logDenominator - logNominator
  where
    logNominator = fromIntegral k * logGamma a
    logDenominator = logGamma (fromIntegral k * a)

-- | Create a symmetric Dirichlet distribution of given dimension and parameter.
--
-- Return Left if:
--
-- - The given dimension is smaller than two.
--
-- - The parameter is negative or zero.
dirichletDistributionSymmetric :: Int -> Double -> Either String DirichletDistributionSymmetric
dirichletDistributionSymmetric k a
  | k < 2 =
    Left "dirichletDistributionSymmetric: The dimension is smaller than two."
  | a <= 0 =
    Left "dirichletDistributionSymmetric: The parameter is negative or zero."
  | otherwise = Right $ DirichletDistributionSymmetric a k (invBetaSym k a)

-- | Density of the symmetric Dirichlet distribution evaluated at a given value
-- vector.
--
-- Return 0 if:
--
-- - The value vector has a different dimension.
--
-- - The value vector has elements being negative or zero.
--
-- - The value vector does not sum to 1.0 (with tolerance @eps = 1e-14@).
dirichletDensitySymmetric :: DirichletDistributionSymmetric -> V.Vector Double -> Log Double
dirichletDensitySymmetric (DirichletDistributionSymmetric a k c) xs
  | k /= V.length xs = 0
  | isNegativeOrZero xs = 0
  | not (isNormalized xs) = 0
  | otherwise = c * Exp logXsPow
  where
    logXsPow = V.sum $ V.map (\x -> log $ x ** (a - 1.0)) xs

-- | Sample a value vector from the symmetric Dirichlet distribution.
dirichletSampleSymmetric ::
  PrimMonad m =>
  DirichletDistributionSymmetric ->
  Gen (PrimState m) ->
  m (V.Vector Double)
dirichletSampleSymmetric (DirichletDistributionSymmetric a k _) g = do
  ys <- V.replicateM k (gamma a 1.0 g)
  let s = V.sum ys
  return $ V.map (/ s) ys


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

import qualified Data.Vector.Unboxed           as VU
import qualified Data.Vector.Unboxed.Mutable   as VUM
import           Data.Complex
import           Data.Function
import           Control.Monad
import           Control.Monad.Primitive
import           Control.Monad.ST

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

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

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

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


{-# LANGUAGE FlexibleContexts, TypeFamilies #-}
module Main where

import Data.Complex (Complex(..), magnitude)
import qualified Diagrams.Prelude as D

import qualified Diagrams.Backend.Rasterific.CmdLine as CmdLine

main :: IO ()
main = CmdLine.mainWith (picture 32 2.0)

type C = Complex Double

-- calc
mandelbrotSet :: Traversable t =>
                 Int -> Double -> t C -> t Bool
mandelbrotSet n inf lst = fmap (testMandelbrot n inf) lst

testMandelbrot :: Int -> Double -> C -> Bool
testMandelbrot n inf c = isDiverged (mandelbrot c) n inf

isDiverged :: (C -> C) -> Int -> Double -> Bool
isDiverged fn n inf = filter (\z -> magnitude z > inf) seq == []
  where seq = take n $ iterate fn 0

mandelbrot :: C -> C -> C
mandelbrot c z = z*z + c

-- draw
toPixel :: Bool -> Int
toPixel True = 0
toPixel False = 1

toSquare n =
  D.opacity (fromIntegral n)
  $ D.fc D.black
  $ D.lw D.medium
  $ D.square 1

side n v0 v1 =
   let sv = (v1 - v0) / fromIntegral n
   in  [v0, (v0 + sv) .. v1]

grid :: Int -> Double -> Double -> [[C]]
grid n v0 v1 = map (\y -> map (:+ y) s) s
  where s = side n v0 v1

image n inf = map (map (toSquare . toPixel . (testMandelbrot n inf))) $ grid 64 (-2) 2

picture :: Int -> Double -> D.Diagram CmdLine.B
picture n inf = D.bgFrame 3 D.pink $ (D.vcat . map D.hcat $ image n inf)


module Scheme.Primitives.Numbers
  ( numericPrimitives,
    unwrapNumber,
  )
where

import Control.Exception (throw)
import Control.Monad.Except (MonadError (throwError))
import Data.Complex (imagPart, realPart)
import Data.Ratio (denominator, numerator)
import Data.Text (Text)
import Scheme.Operators (unOp)
import Scheme.Types (Number (..), SchemeError (..), SchemeResult, SchemeVal (..))

numericPrimitives :: [(Text, [SchemeVal] -> SchemeResult SchemeVal)]
numericPrimitives =
  [ ("number?", unOp isNumber),
    ("complex?", unOp isComplex),
    ("real?", unOp isReal),
    ("rational?", unOp isRational),
    ("integer?", unOp isInteger),
    ("exact?", unOp isExact),
    ("inexact?", unOp isInexact),
    ("exact-integer?", unOp isExactInteger),
    ("finite?", unOp isFinite),
    ("infinite?", unOp isInfinite'),
    ("nan?", unOp isNaN'),
    ("=", numericBoolOp (==)),
    ("<", numericBoolOp (<)),
    (">", numericBoolOp (>)),
    ("<=", numericBoolOp (<=)),
    (">=", numericBoolOp (>=)),
    ("+", add),
    ("-", sub),
    ("*", multiply),
    ("/", division),
    ("abs", unOp numAbs),
    ("floor/", floorInt),
    ("floor-quotient", floorQuot),
    ("floor-remainder", floorRem),
    ("truncate/", truncateInt),
    ("truncate-quotient", truncateQuot),
    ("truncate-remainder", truncateRem),
    ("quotient", truncateQuot),
    ("remainder", truncateRem),
    ("modulo", floorRem),
    ("gcd", numGcd),
    ("lcm", numLcm)
  ]

isNumber :: SchemeVal -> SchemeVal
isNumber (Number (Integer _)) = Boolean True
isNumber (Number (Real _)) = Boolean True
isNumber (Number (Rational _)) = Boolean True
isNumber (Number (Complex _)) = Boolean True
isNumber _ = Boolean False

isComplex :: SchemeVal -> SchemeVal
isComplex (Number (Integer _)) = Boolean True
isComplex (Number (Real _)) = Boolean True
isComplex (Number (Rational _)) = Boolean True
isComplex (Number (Complex _)) = Boolean True
isComplex _ = Boolean False

isReal :: SchemeVal -> SchemeVal
isReal (Number (Real _)) = Boolean True
isReal (Number (Rational _)) = Boolean True
isReal (Number (Integer _)) = Boolean True
isReal (Number (Complex x)) = if imagPart x == 0 then Boolean True else Boolean False
isReal _ = Boolean False

isRational :: SchemeVal -> SchemeVal
isRational (Number (Rational _)) = Boolean True
isRational (Number (Integer _)) = Boolean True
isRational (Number (Real x)) = Boolean (not $ isInfinite x)
isRational _ = Boolean False

isInteger :: SchemeVal -> SchemeVal
isInteger (Number (Integer _)) = Boolean True
isInteger (Number (Real x)) = Boolean (isDoubleInt x)
isInteger (Number (Rational x)) = Boolean (numerator x >= denominator x && numerator x `mod` denominator x == 0)
isInteger (Number (Complex x)) = Boolean (isDoubleInt (realPart x) && isDoubleInt (imagPart x))
isInteger _ = Boolean False

isExact :: SchemeVal -> SchemeVal
isExact (Number (Integer _)) = Boolean True
isExact (Number (Rational _)) = Boolean True
isExact _ = Boolean False

isInexact :: SchemeVal -> SchemeVal
isInexact (Number (Real _)) = Boolean True
isInexact _ = Boolean False

isExactInteger :: SchemeVal -> SchemeVal
isExactInteger x = case (isExact x, isInteger x) of
  (Boolean True, Boolean True) -> Boolean True
  _ -> Boolean False

isFinite :: SchemeVal -> SchemeVal
isFinite (Number (Real x)) = Boolean (not (isInfinite x || isNaN x))
isFinite (Number (Complex x)) = Boolean (not (isInfinite imag || isInfinite real || isNaN real || isNaN imag))
  where
    imag = imagPart x
    real = realPart x
isFinite _ = Boolean True

isInfinite' :: SchemeVal -> SchemeVal
isInfinite' (Number (Real x)) = Boolean (isInfinite x || isNaN x)
isInfinite' (Number (Complex x)) = Boolean (isInfinite imag || isInfinite real || isNaN real || isNaN imag)
  where
    imag = imagPart x
    real = realPart x
isInfinite' _ = Boolean False

isNaN' :: SchemeVal -> SchemeVal
isNaN' (Number (Real x)) = Boolean (isNaN x)
isNaN' (Number (Complex x)) = Boolean (isNaN (imagPart x) || isNaN (realPart x))
isNaN' _ = Boolean False

add :: [SchemeVal] -> SchemeResult SchemeVal
add [] = pure $ Number $ Integer 0
add xs = do
  nums <- mapM unwrapNumber xs
  return $ Number (sum nums)

multiply :: [SchemeVal] -> SchemeResult SchemeVal
multiply [] = pure $ Number $ Integer 1
multiply xs = do
  nums <- mapM unwrapNumber xs
  return $ Number (product nums)

sub :: [SchemeVal] -> SchemeResult SchemeVal
sub [] = throwError $ ArgumentLengthMismatch 1 []
sub [x] = do
  num <- unwrapNumber x
  return $ Number (negate num)
sub xs = do
  nums <- mapM unwrapNumber xs
  return $ Number (foldl1 (-) nums)

division :: [SchemeVal] -> SchemeResult SchemeVal
division [] = throwError $ ArgumentLengthMismatch 1 []
division [x] = do
  num <- unwrapNumber x
  return $ Number (recip num)
division xs = do
  nums <- mapM unwrapNumber xs
  return $ Number (foldl1 (/) nums)

numAbs :: SchemeVal -> SchemeVal
numAbs (Number n) = Number $ abs n
numAbs x = throw $ TypeMismatch "number" x

floorInt :: [SchemeVal] -> SchemeResult SchemeVal
floorInt [Number (Integer x), Number (Integer y)] = return $ List [Number $ Integer x', Number $ Integer y']
  where
    (x', y') = x `divMod` y
floorInt [x, Number (Integer _)] = throwError $ TypeMismatch "integer" x
floorInt [Number (Integer _), x] = throwError $ TypeMismatch "integer" x
floorInt _ = throwError $ ArgumentLengthMismatch 2 []

floorQuot :: [SchemeVal] -> SchemeResult SchemeVal
floorQuot [Number (Integer x), Number (Integer y)] = return $ Number $ Integer (x `div` y)
floorQuot [x, Number (Integer _)] = throwError $ TypeMismatch "integer" x
floorQuot [Number (Integer _), x] = throwError $ TypeMismatch "integer" x
floorQuot _ = throwError $ ArgumentLengthMismatch 2 []

floorRem :: [SchemeVal] -> SchemeResult SchemeVal
floorRem [Number (Integer x), Number (Integer y)] = return $ Number $ Integer (x `mod` y)
floorRem [x, Number (Integer _)] = throwError $ TypeMismatch "integer" x
floorRem [Number (Integer _), x] = throwError $ TypeMismatch "integer" x
floorRem _ = throwError $ ArgumentLengthMismatch 2 []

truncateInt :: [SchemeVal] -> SchemeResult SchemeVal
truncateInt [Number (Integer x), Number (Integer y)] = return $ List [Number $ Integer x', Number $ Integer y']
  where
    (x', y') = x `quotRem` y
truncateInt [x, Number (Integer _)] = throwError $ TypeMismatch "integer" x
truncateInt [Number (Integer _), x] = throwError $ TypeMismatch "integer" x
truncateInt _ = throwError $ ArgumentLengthMismatch 2 []

truncateQuot :: [SchemeVal] -> SchemeResult SchemeVal
truncateQuot [Number (Integer x), Number (Integer y)] = return $ Number $ Integer (x `quot` y)
truncateQuot [x, Number (Integer _)] = throwError $ TypeMismatch "integer" x
truncateQuot [Number (Integer _), x] = throwError $ TypeMismatch "integer" x
truncateQuot _ = throwError $ ArgumentLengthMismatch 2 []

truncateRem :: [SchemeVal] -> SchemeResult SchemeVal
truncateRem [Number (Integer x), Number (Integer y)] = return $ Number $ Integer (x `rem` y)
truncateRem [x, Number (Integer _)] = throwError $ TypeMismatch "integer" x
truncateRem [Number (Integer _), x] = throwError $ TypeMismatch "integer" x
truncateRem _ = throwError $ ArgumentLengthMismatch 2 []

numGcd :: [SchemeVal] -> SchemeResult SchemeVal
numGcd [] = return $ Number (Integer 0)
numGcd [x] = return x
numGcd xs = do
  ints <- mapM unwrapToInt xs
  return $ Number $ Integer (foldl1 gcd ints)

numLcm :: [SchemeVal] -> SchemeResult SchemeVal
numLcm [] = return $ Number (Integer 1)
numLcm [x] = return x
numLcm xs = do
  ints <- mapM unwrapToInt xs
  return $ Number $ Integer (foldl1 lcm ints)

numericBoolOp :: (Number -> Number -> Bool) -> [SchemeVal] -> Either SchemeError SchemeVal
numericBoolOp _ [] = throwError $ ArgumentLengthMismatch 1 []
numericBoolOp _ [x] = case isNaN' x of
  Boolean True -> return $ Boolean False
  _ -> return $ Boolean True
numericBoolOp op xs = do
  nums <- mapM unwrapNumber xs
  return $ Boolean $ and $ zipWith op nums $ drop 1 nums

isDoubleInt :: Double -> Bool
isDoubleInt d = (ceiling d :: Integer) == (floor d :: Integer)

unwrapToInt :: SchemeVal -> SchemeResult Integer
unwrapToInt (Number (Integer x)) = pure x
unwrapToInt (Number (Real x)) = pure $ floor x
unwrapToInt (Number (Rational x)) = pure $ floor x
unwrapToInt x = throwError $ TypeMismatch "integer" x

unwrapNumber :: SchemeVal -> SchemeResult Number
unwrapNumber (Number x) = pure x
unwrapNumber (List [x]) = unwrapNumber x
unwrapNumber x = throwError $ TypeMismatch "number" x


{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE DeriveDataTypeable, DeriveGeneric #-}
-----------------------------------------------------------------------------
-- |
-- Module      :  Statistics.Distribution.Beta
-- Copyright   :  (C) 2012 Edward Kmett,
-- License     :  BSD-style (see the file LICENSE)
--
-- Maintainer  :  Edward Kmett <ekmett@gmail.com>
-- Stability   :  provisional
-- Portability :  DeriveDataTypeable
--
----------------------------------------------------------------------------
module Statistics.Distribution.Beta
  ( BetaDistribution
    -- * Constructor
  , betaDistr
  , betaDistrE
  , improperBetaDistr
  , improperBetaDistrE
    -- * Accessors
  , bdAlpha
  , bdBeta
  ) where

import Control.Applicative
import Data.Aeson            (FromJSON(..), ToJSON, Value(..), (.:))
import Data.Binary           (Binary(..))
import Data.Data             (Data, Typeable)
import GHC.Generics          (Generic)
import Numeric.SpecFunctions (
  incompleteBeta, invIncompleteBeta, logBeta, digamma, log1p)
import Numeric.MathFunctions.Constants (m_NaN,m_neg_inf)
import qualified Statistics.Distribution as D
import Statistics.Internal

-- | The beta distribution
data BetaDistribution = BD
 { bdAlpha :: {-# UNPACK #-} !Double
   -- ^ Alpha shape parameter
 , bdBeta  :: {-# UNPACK #-} !Double
   -- ^ Beta shape parameter
 } deriving (Eq, Typeable, Data, Generic)

instance Show BetaDistribution where
  showsPrec n (BD a b) = defaultShow2 "improperBetaDistr" a b n
instance Read BetaDistribution where
  readPrec = defaultReadPrecM2 "improperBetaDistr" improperBetaDistrE

instance ToJSON BetaDistribution
instance FromJSON BetaDistribution where
  parseJSON (Object v) = do
    a <- v .: "bdAlpha"
    b <- v .: "bdBeta"
    maybe (fail $ errMsgI a b) return $ improperBetaDistrE a b
  parseJSON _ = empty

instance Binary BetaDistribution where
  put (BD a b) = put a >> put b
  get = do
    a <- get
    b <- get
    maybe (fail $ errMsgI a b) return $ improperBetaDistrE a b

-- | Create beta distribution. Both shape parameters must be positive.
betaDistr :: Double             -- ^ Shape parameter alpha
          -> Double             -- ^ Shape parameter beta
          -> BetaDistribution
betaDistr a b = maybe (error $ errMsg a b) id $ betaDistrE a b

-- | Create beta distribution. Both shape parameters must be positive.
betaDistrE :: Double             -- ^ Shape parameter alpha
          -> Double             -- ^ Shape parameter beta
          -> Maybe BetaDistribution
betaDistrE a b
  | a > 0 && b > 0 = Just (BD a b)
  | otherwise      = Nothing

errMsg :: Double -> Double -> String
errMsg a b = "Statistics.Distribution.Beta.betaDistr: "
          ++ "shape parameters must be positive. Got a = "
          ++ show a
          ++ " b = "
          ++ show b

-- | Create beta distribution. Both shape parameters must be
-- non-negative. So it allows to construct improper beta distribution
-- which could be used as improper prior.
improperBetaDistr :: Double             -- ^ Shape parameter alpha
                  -> Double             -- ^ Shape parameter beta
                  -> BetaDistribution
improperBetaDistr a b
  = maybe (error $ errMsgI a b) id $ improperBetaDistrE a b

-- | Create beta distribution. Both shape parameters must be
-- non-negative. So it allows to construct improper beta distribution
-- which could be used as improper prior.
improperBetaDistrE :: Double             -- ^ Shape parameter alpha
                   -> Double             -- ^ Shape parameter beta
                   -> Maybe BetaDistribution
improperBetaDistrE a b
  | a >= 0 && b >= 0 = Just (BD a b)
  | otherwise        = Nothing

errMsgI :: Double -> Double -> String
errMsgI a b
  =  "Statistics.Distribution.Beta.betaDistr: "
  ++ "shape parameters must be non-negative. Got a = " ++ show a
  ++ " b = " ++ show b

instance D.Distribution BetaDistribution where
  cumulative (BD a b) x
    | x <= 0    = 0
    | x >= 1    = 1
    | otherwise = incompleteBeta a b x
  complCumulative (BD a b) x
    | x <= 0    = 1
    | x >= 1    = 0
    -- For small x we use direct computation to avoid precision loss
    -- when computing (1-x)
    | x <  0.5  = 1 - incompleteBeta a b x
    -- Otherwise we use property of incomplete beta:
    --  > I(x,a,b) = 1 - I(1-x,b,a)
    | otherwise = incompleteBeta b a (1-x)

instance D.Mean BetaDistribution where
  mean (BD a b) = a / (a + b)

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

instance D.Variance BetaDistribution where
  variance (BD a b) = a*b / (apb*apb*(apb+1))
    where apb = a + b

instance D.MaybeVariance BetaDistribution where
  maybeVariance = Just . D.variance

instance D.Entropy BetaDistribution where
  entropy (BD a b) =
    logBeta a b
    - (a-1) * digamma a
    - (b-1) * digamma b
    + (a+b-2) * digamma (a+b)

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

instance D.ContDistr BetaDistribution where
  density (BD a b) x
    | a <= 0 || b <= 0 = m_NaN
    | x <= 0 = 0
    | x >= 1 = 0
    | otherwise = exp $ (a-1)*log x + (b-1) * log1p (-x) - logBeta a b
  logDensity (BD a b) x
    | a <= 0 || b <= 0 = m_NaN
    | x <= 0 = m_neg_inf
    | x >= 1 = m_neg_inf
    | otherwise = (a-1)*log x + (b-1)*log1p (-x) - logBeta a b

  quantile (BD a b) p
    | p == 0         = 0
    | p == 1         = 1
    | p > 0 && p < 1 = invIncompleteBeta a b p
    | otherwise      =
        error $ "Statistics.Distribution.Gamma.quantile: p must be in [0,1] range. Got: "++show p

instance D.ContGen BetaDistribution where
  genContVar = D.genContinuous


module Main where
import Plot
import Graphics.GD
import Data.Complex

imageDim :: Size
imageDim = (600,600)

-- How many iterations before  we declare the point doesn't diverge
maxIter = 255

-- Coordinates in the complex plane to graph
-- First coordinate must be top left
mbWindow :: Window
mbWindow = ( (-1.8, 0.02), (-1.7, -0.08) ) -- smaller ship on horizon

main :: IO ()
main = drawPlot drawShip imageDim mbWindow "burning-ship.png"

drawShip :: Coordinate -> Color
drawShip (x,y) = colorIterations $ ship (x :+ y) (0 :+ 0) 0

-- Count number of iterations for the given point to diverge
-- Non-diverging points return 0
ship :: Complex Double -- Coordinate to test
           -> Complex Double -- Iterating Z value
           -> Int -- Current iteration
           -> Int -- Iterations before diverging
ship c z iter
    | iter > maxIter = 0
    | otherwise = let z' = (abs(realPart z) :+ abs(imagPart z))^2 + c in
                  if magnitude z' > 2
                  then iter
                  else ship c z' (iter+1)

colorIterations :: Int -> Color
colorIterations x
    | x > maxIter = rgb 255 255 255
    | otherwise = rgb x x x


module Titanic where

import Numeric.LinearAlgebra                         (col, disp, Matrix, R, rand, scale, toLists, tr, Vector, (><))
import Numeric.LinearAlgebra.HMatrix                 (mul, matrix, sumElements)
import Text.Read                                     (readMaybe)

type Gender = Double
type Age = Double
type SibSp = Double
type Embarked = Double
type PClass = Double
type Survived = Double

data Row = Row Gender Age PClass Embarked

processRow :: [[Char]] -> [Double]
processRow row = [gender, age] ++ pclass ++ embarked where
    gender = getGender row
    age = getAge row
    pclass = getPClass row
    embarked = getEmbarked row

processFeatures :: [[[Char]]] -> Matrix Double
processFeatures inputs = matrix 8 $ concat [processRow x | x <- inputs]

processGender :: [[[Char]]] -> Matrix Gender
processGender inputs = do
    col [getGender x | x <- inputs]

processTargets :: [[[Char]]] -> Matrix Survived
processTargets inputs = do
    col [getTarget x | x <- inputs]

getAge :: [[Char]] -> Age
getAge row =  case (readMaybe (row !! 6)) of
    -- 80 is the maximum age in the dataset. We divide the age by it to
    -- keep it between 0.0 and 1.0
    Just x -> x / 80
    Nothing -> 0.0

getPClass :: [[Char]] -> [PClass]
getPClass row =  case (row !! 2) of
    "1" -> [1.0, 0.0, 0.0]
    "2" -> [0.0, 1.0, 0.0]
    "3" -> [0.0, 0.0, 1.0]
    "" -> [0.0, 0.0, 0.0]

getEmbarked :: [[Char]] -> [Embarked]
getEmbarked row = case (readMaybe (row !! 12)) of
    Just x -> case x of
        "C" -> [1.0, 0.0, 0.0]
        "Q" -> [0.0, 1.0, 0.0]
        "S" -> [0.0, 0.0, 1.0]
    Nothing -> [0.0, 0.0, 0.0]

getGender :: [[Char]] -> Gender
getGender row = case (row !! 5) of
    "male"   -> -1.0
    "female" -> 1.0

getTarget :: [[Char]] -> Survived
getTarget row = fromIntegral $ read $ row !! 1