{-# LANGUAGE FlexibleContexts, Rank2Types, NoMonomorphismRestriction,
             StandaloneDeriving #-}

module Numeric.MaxEnt.Linear where

import Control.Applicative

import Data.List (transpose)
import qualified Data.Vector.Storable as S

import Numeric.MaxEnt.ConjugateGradient (minimize, dot)
import Numeric.Optimization.Algorithms.HagerZhang05 (Result, Statistics)
import Numeric.AD

multMV :: (Num a) => [[a]] -> [a] -> [a]
multMV mat vec = map (\row -> dot row vec) mat

probs :: (Floating a) => [[a]] -> [a] -> [a]
probs matrix ls = result where
    norm = partitionFunc matrix ls
    result = map (\x -> exp x / norm ) $ (transpose matrix) `multMV` ls

partitionFunc :: (Floating a) => [[a]] -> [a] -> a
partitionFunc matrix ws = sum . map exp . multMV (transpose matrix) $ ws

-- This is almost the sam as the objectiveFunc

objectiveFunc :: (Floating a) => [[a]] -> [a] -> [a] -> a
objectiveFunc as moments ls = log $ partitionFunc as ls - dot ls moments

data LinearConstraints = LC
  { unLC :: forall a. (Floating a) => ([[a]], [a]) }

-- These instances default the underlying numeric type of `LC` to `Double`,
-- which may be problematic for some usages.
deriving instance Eq LinearConstraints
deriving instance Show LinearConstraints

-- | This is for the linear case Ax = b
--   @x@ in this situation is the vector of probablities.
--
--  Consider the 1 dimensional circular convolution using hmatrix.
--
--  >>> import Numeric.LinearAlgebra
--  >>> fromLists [[0.68, 0.22, 0.1], [0.1, 0.68, 0.22], [0.22, 0.1, 0.68]] <> fromLists [[0.2], [0.5], [0.3]]
--  (3><1) [0.276, 0.426, 0.298]
--
--   Now if we were given just the convolution and the output, we can use 'linear' to infer the input.
--
--   >>> linear 3.0e-17 $ LC ([[0.68, 0.22, 0.1], [0.1, 0.68, 0.22], [0.22, 0.1, 0.68]], [0.276, 0.426, 0.298])
--   Right (fromList [0.2000000000000001,0.49999999999999983,0.30000000000000004])
--
--   I fell compelled to point out that we could also just invert the original
--   convolution matrix. Supposedly using maxent can reduce errors from noise if
--   the convolution matrix is not properly estimated.
linear :: Double
       -- ^ Tolerance for the numerical solver
       -> LinearConstraints
       -- ^ The matrix A and column vector b
       -> Either (Result, Statistics) (S.Vector Double)
       -- ^ Either a description of what went wrong or the probability
       --   distribution
linear tolerance constraints  =
    let (matrix, output) = unLC constraints
        obj = objectiveFunc matrix output
        n = length output
    in (S.fromList . probs matrix . S.toList) <$> minimize tolerance n obj

--------------------------------------------------------------------------------
-- I updated everything below to work with the new types, but it's not clear to
-- me what it's for.  -- EP
--------------------------------------------------------------------------------

linear' :: (Floating a, Ord a)
        => LinearConstraints
        -- ^ The matrix A and column vector b
        -> [[a]]
        -- ^ Either a description of what went wrong or the probability
        --   distribution
linear' constraints =
    let (matrix, output) = unLC constraints
        obj = objectiveFunc matrix output
        guess = 1 : replicate (length output - 1) 0
    in map (probs matrix) . gradientDescent obj $ guess

linear'' :: (Floating a, Ord a)
         => LinearConstraints
         -- ^ The matrix A and column vector b
         -> [[a]]
         -- ^ Either a description of what went wrong or the probability
         --   distribution
linear'' constraints =
    let (matrix, output) = unLC constraints
        obj = objectiveFunc matrix output
        guess = 1 : replicate (length output - 1) 0
    in map (probs matrix) . conjugateGradientDescent obj $ guess

--test1 = LC ( [ [0.892532,0.003851,0.063870,0.001593,0.038155]
--             , [0.237713,0.111149,0.326964,0.271535,0.052639]
--             , [0.133708,0.788233,0.051543,0.003976,0.022539]
--             , [0.238064,0.263171,0.112279,0.270452,0.116034]
--             , [0.844155,0.011312,0.001470,0.001826,0.141237]
--             ]
--           ,
--             [0.246323,0.235600,0.071699,0.211339,0.238439]
--           )


{-# LANGUAGE FlexibleInstances #-}
-----------------------------------------------------------------------------
-- |
-- Module     : Foreign.LAPACK
-- Copyright  : Copyright (c) 2010, Patrick Perry <patperry@gmail.com>
-- License    : BSD3
-- Maintainer : Patrick Perry <patperry@gmail.com>
-- Stability  : experimental
--
-- Types with a Linear Algebra PACKage.
--
module Foreign.LAPACK (
    -- * LAPACK typeclass
    LAPACK(..),
    -- * Enums
    module Foreign.LAPACK.Types,
    ) where

import Control.Monad( forM_, when )
import Control.Exception( assert )
import Data.Complex( Complex )
import Foreign
import Foreign.BLAS
import Foreign.LAPACK.Types
import Foreign.LAPACK.Double
import Foreign.LAPACK.Zomplex

-- | Types with LAPACK operations.
class (BLAS3 e) => LAPACK e where
    geqrf :: Int -> Int -> Ptr e -> Int -> Ptr e -> IO ()
    gelqf :: Int -> Int -> Ptr e -> Int -> Ptr e -> IO ()

    heevr :: EigJob -> EigRange -> Uplo -> Int -> Ptr e -> Int -> Double
          -> Ptr Double -> Ptr e -> Int -> Ptr Int -> IO Int

    larfg :: Int -> Ptr e -> Ptr e -> Int -> IO e
    potrf :: Uplo -> Int -> Ptr e -> Int -> IO Int
    potrs :: Uplo -> Int -> Int -> Ptr e -> Int -> Ptr e -> Int -> IO ()
    pptrf :: Uplo -> Int -> Ptr e -> IO Int
    pptrs :: Uplo -> Int -> Int -> Ptr e -> Ptr e -> Int -> IO ()

    unmqr :: Side -> Trans -> Int -> Int -> Int -> Ptr e -> Int -> Ptr e -> Ptr e -> Int -> IO ()
    unmlq :: Side -> Trans -> Int -> Int -> Int -> Ptr e -> Int -> Ptr e -> Ptr e -> Int -> IO ()

callWithWork :: (Storable e) => (Ptr e -> Ptr LAInt -> Ptr LAInt -> IO ()) -> IO LAInt
callWithWork call =
    alloca $ \pquery ->
    with (-1) $ \pnone ->
    alloca $ \pinfo -> do
        call pquery pnone pinfo
        lwork <- (ceiling . max 1) `fmap` (peek (castPtr pquery) :: IO Double)
        allocaArray lwork $ \pwork ->
            withEnum lwork $ \plwork -> do
                call pwork plwork pinfo
                peek pinfo

callWithWorkIWork :: (Storable e)
    => (Ptr e -> Ptr LAInt -> Ptr LAInt -> Ptr LAInt -> Ptr LAInt -> IO ())
    -> IO LAInt
callWithWorkIWork call =
    alloca $ \pquery ->
    alloca $ \piquery ->
    with (-1) $ \pnone ->
    alloca $ \pinfo -> do
        call pquery pnone piquery pnone pinfo
        lwork <- (ceiling . max 1) `fmap` (peek (castPtr pquery) :: IO Double)
        liwork <- max 1 `fmap` peek piquery
        allocaArray lwork $ \pwork ->
            withEnum lwork $ \plwork ->
                allocaArray (fromEnum liwork) $ \piwork ->
                    with liwork $ \pliwork -> do
                        call pwork plwork piwork pliwork pinfo
                        peek pinfo

callWithWorkRWorkIWork :: (Storable e)
    => (Ptr e -> Ptr LAInt -> Ptr Double -> Ptr LAInt -> Ptr LAInt -> Ptr LAInt -> Ptr LAInt -> IO ())
    -> IO LAInt
callWithWorkRWorkIWork call =
    alloca $ \pquery ->
    alloca $ \prquery ->
    alloca $ \piquery ->
    with (-1) $ \pnone ->
    alloca $ \pinfo -> do
        call pquery pnone prquery pnone piquery pnone pinfo
        lwork <- (ceiling . max 1) `fmap` (peek (castPtr pquery) :: IO Double)
        lrwork <- (ceiling . max 1) `fmap` peek prquery
        liwork <- max 1 `fmap` peek piquery
        allocaArray lwork $ \pwork ->
            withEnum lwork $ \plwork ->
                allocaArray lrwork $ \prwork ->
                    withEnum lrwork $ \plrwork ->
                        allocaArray (fromEnum liwork) $ \piwork ->
                            with liwork $ \pliwork -> do
                                call pwork plwork prwork plrwork piwork pliwork pinfo
                                peek pinfo

checkInfo :: (Num info) => info -> IO ()
checkInfo info = assert (info == 0) $ return ()

withEnum :: (Enum a, Storable a) => Int -> (Ptr a -> IO b) -> IO b
withEnum = with . toEnum
{-# INLINE withEnum #-}

instance LAPACK Double where
    geqrf m n pa lda ptau =
        withEnum m $ \pm ->
        withEnum n $ \pn ->
        withEnum lda $ \plda ->
            checkInfo =<< callWithWork (dgeqrf pm pn pa plda ptau)

    gelqf m n pa lda ptau =
        withEnum m $ \pm ->
        withEnum n $ \pn ->
        withEnum lda $ \plda ->
            checkInfo =<< callWithWork (dgelqf pm pn pa plda ptau)

    heevr jobz range uplo n pa lda abstol pw pz ldz psuppz =
        withEigJob jobz $ \pjobz ->
        withEigRange range $ \prange pvl pvu pil piu ->
        withUplo uplo $ \puplo ->
        withEnum n $ \pn ->
        withEnum lda $ \plda ->
        with abstol $ \pabstol ->
        alloca $ \pm ->
        withEnum ldz $ \pldz ->
        allocaArray nsuppz $ \psuppz' -> do
            checkInfo =<<
                (callWithWorkIWork $
                    dsyevr pjobz prange puplo pn pa plda pvl pvu pil piu
                           pabstol pm pw pz pldz psuppz')
            m <- fromEnum `fmap` peek pm
            when (psuppz /= nullPtr) $
                forM_ [ 0..2*m - 1 ] $ \i -> do
                    s <- peekElemOff psuppz' i
                    pokeElemOff psuppz i (fromEnum s)
            return m
      where
        nsuppz = if psuppz == nullPtr then 0 else 2 * max 1 n

    unmqr side trans m n k pa lda ptau pc ldc =
        withSide side $ \pside ->
        withTrans trans $ \ptrans ->
        withEnum m $ \pm ->
        withEnum n $ \pn ->
        withEnum k $ \pk ->
        withEnum lda $ \plda ->
        withEnum ldc $ \pldc ->
            checkInfo =<< callWithWork (dormqr pside ptrans pm pn pk pa plda ptau pc pldc)

    unmlq side trans m n k pa lda ptau pc ldc =
        withSide side $ \pside ->
        withTrans trans $ \ptrans ->
        withEnum m $ \pm ->
        withEnum n $ \pn ->
        withEnum k $ \pk ->
        withEnum lda $ \plda ->
        withEnum ldc $ \pldc ->
            checkInfo =<< callWithWork (dormlq pside ptrans pm pn pk pa plda ptau pc pldc)

    larfg n palpha px incx =
        withEnum n $ \pn ->
        withEnum incx $ \pincx ->
        alloca $ \ptau -> do
            dlarfg pn palpha px pincx ptau
            peek ptau

    potrf uplo n pa lda =
        withUplo uplo $ \puplo ->
        withEnum n $ \pn ->
        withEnum lda $ \plda ->
        alloca $ \pinfo -> do
            dpotrf puplo pn pa plda pinfo
            info <- fromEnum `fmap` peek pinfo
            assert (info >= 0) $ return info

    potrs uplo n nrhs pa lda pb ldb =
        withUplo uplo $ \puplo ->
        withEnum n $ \pn ->
        withEnum nrhs $ \pnrhs ->
        withEnum lda $ \plda ->
        withEnum ldb $ \pldb ->
        alloca $ \pinfo -> do
            dpotrs puplo pn pnrhs pa plda pb pldb pinfo
            checkInfo =<< peek pinfo

    pptrf uplo n pa =
        withUplo uplo $ \puplo ->
        withEnum n $ \pn ->
        alloca $ \pinfo -> do
            dpptrf puplo pn pa pinfo
            info <- fromEnum `fmap` peek pinfo
            assert (info >= 0) $ return info

    pptrs uplo n nrhs pa pb ldb =
        withUplo uplo $ \puplo ->
        withEnum n $ \pn ->
        withEnum nrhs $ \pnrhs ->
        withEnum ldb $ \pldb ->
        alloca $ \pinfo -> do
            dpptrs puplo pn pnrhs pa pb pldb pinfo
            checkInfo =<< peek pinfo

instance LAPACK (Complex Double) where
    geqrf m n pa lda ptau =
        withEnum m $ \pm ->
        withEnum n $ \pn ->
        withEnum lda $ \plda ->
            checkInfo =<< callWithWork (zgeqrf pm pn pa plda ptau)

    gelqf m n pa lda ptau =
        withEnum m $ \pm ->
        withEnum n $ \pn ->
        withEnum lda $ \plda ->
            checkInfo =<< callWithWork (zgelqf pm pn pa plda ptau)

    heevr jobz range uplo n pa lda abstol pw pz ldz psuppz =
        withEigJob jobz $ \pjobz ->
        withEigRange range $ \prange pvl pvu pil piu ->
        withUplo uplo $ \puplo ->
        withEnum n $ \pn ->
        withEnum lda $ \plda ->
        with abstol $ \pabstol ->
        alloca $ \pm ->
        withEnum ldz $ \pldz ->
        allocaArray nsuppz $ \psuppz' -> do
            checkInfo =<<
                (callWithWorkRWorkIWork $
                    zheevr pjobz prange puplo pn pa plda pvl pvu pil piu
                           pabstol pm pw pz pldz psuppz')
            m <- fromEnum `fmap` peek pm
            when (psuppz /= nullPtr) $
                forM_ [ 0..2*m - 1 ] $ \i -> do
                    s <- peekElemOff psuppz' i
                    pokeElemOff psuppz i (fromEnum s)
            return m
      where
        nsuppz = if psuppz == nullPtr then 0 else 2 * max 1 n

    unmqr side trans m n k pa lda ptau pc ldc =
        withSide side $ \pside ->
        withTrans trans $ \ptrans ->
        withEnum m $ \pm ->
        withEnum n $ \pn ->
        withEnum k $ \pk ->
        withEnum lda $ \plda ->
        withEnum ldc $ \pldc ->
            checkInfo =<< callWithWork (zunmqr pside ptrans pm pn pk pa plda ptau pc pldc)

    unmlq side trans m n k pa lda ptau pc ldc =
        withSide side $ \pside ->
        withTrans trans $ \ptrans ->
        withEnum m $ \pm ->
        withEnum n $ \pn ->
        withEnum k $ \pk ->
        withEnum lda $ \plda ->
        withEnum ldc $ \pldc ->
            checkInfo =<< callWithWork (zunmlq pside ptrans pm pn pk pa plda ptau pc pldc)

    larfg n palpha px incx =
        withEnum n $ \pn ->
        withEnum incx $ \pincx ->
        alloca $ \ptau -> do
            zlarfg pn palpha px pincx ptau
            peek ptau

    potrf uplo n pa lda =
        withUplo uplo $ \puplo ->
        withEnum n $ \pn ->
        withEnum lda $ \plda ->
        alloca $ \pinfo -> do
            zpotrf puplo pn pa plda pinfo
            info <- fromEnum `fmap` peek pinfo
            assert (info >= 0) $ return info

    potrs uplo n nrhs pa lda pb ldb =
        withUplo uplo $ \puplo ->
        withEnum n $ \pn ->
        withEnum nrhs $ \pnrhs ->
        withEnum lda $ \plda ->
        withEnum ldb $ \pldb ->
        alloca $ \pinfo -> do
            zpotrs puplo pn pnrhs pa plda pb pldb pinfo
            checkInfo =<< peek pinfo

    pptrf uplo n pa =
        withUplo uplo $ \puplo ->
        withEnum n $ \pn ->
        alloca $ \pinfo -> do
            zpptrf puplo pn pa pinfo
            info <- fromEnum `fmap` peek pinfo
            assert (info >= 0) $ return info

    pptrs uplo n nrhs pa pb ldb =
        withUplo uplo $ \puplo ->
        withEnum n $ \pn ->
        withEnum nrhs $ \pnrhs ->
        withEnum ldb $ \pldb ->
        alloca $ \pinfo -> do
            zpptrs puplo pn pnrhs pa pb pldb pinfo
            checkInfo =<< peek pinfo


{-# LANGUAGE FlexibleContexts #-}
module Filter.Pinwheel
  ( module Filter.Utils
  , PinwheelParams(..)
  , pinwheel
  , PinwheelType(..)
  , pinwheelFunc
  , pinwheelFilter
  , convolvePinwheel
  , frequencyDomainMultiply
  ) where

import           Data.Array.Repa     as R
import           Data.Complex
import           Data.List           as L (length, reverse)
import           Data.Vector.Generic as VG (convert)
import           DFT.Plan
import           Filter.Utils
import           Utils.Coordinates

data PinwheelParams = PinwheelParams
  { pinwheelParamsRows         :: !Int
  , pinwheelParamsCols         :: !Int
  , pinwheelParamsAlpha        :: !Double
  , pinwheelParamsRadialPeriod :: !Double
  , pinwheelParamsAngularFreqs :: ![Double]
  , pinwheelParamsRadialFreqs  :: ![Double]
  }

data PinwheelType
  = Pinwheel
  | PinwheelHollow0 !Double
  | PinwheelHollow1 !Double
  deriving (Read, Show)

{-# INLINE rFunc #-}
rFunc :: Int -> Int -> Double
rFunc x y = sqrt . fromIntegral $ x ^ (2 :: Int) + y ^ (2 :: Int)

{-# INLINE thetaFunc #-}
thetaFunc :: Int -> Int -> Double
thetaFunc x y = angleFunctionRad (fromIntegral x) (fromIntegral y)

{-# INLINE pinwheel #-}
pinwheel :: Double -> Double -> Double -> Double -> Int -> Int -> Complex Double
pinwheel af rf maxR alpha x y
  | r == 0 && af == 0 && rf == 0 = 1
  | r == 0 = 0
  | otherwise =
    (r :+ 0) ** (alpha :+ rf * 2 * pi / log maxR) *
    exp (0 :+ af * (thetaFunc x y))
  where
    r = rFunc x y

{-# INCLUDE pinwheelFunc #-}
pinwheelFunc ::
     PinwheelType
  -> Double
  -> Double
  -> Double
  -> Double
  -> Int
  -> Int
  -> Complex Double
pinwheelFunc Pinwheel af rf maxR alpha x y
  | r == 0 && af == 0 && rf == 0 = 1
  | r == 0 = 0
  | otherwise =
    (r :+ 0) ** (alpha :+ rf * 2 * pi / logMaxR) *
    exp (0 :+ af * (thetaFunc x y))
  where
    r = rFunc x y
    logMaxR =
      if maxR == 1
        then 1
        else log maxR
pinwheelFunc (PinwheelHollow0 radius) af rf maxR alpha x y
  | r < radius = 0.0
  | otherwise =
    (r :+ 0) ** (alpha :+ rf * 2 * pi / logMaxR) *
    exp (0 :+ af * (thetaFunc x y))
  where
    r = rFunc x y
    logMaxR =
      if maxR == 1
        then 1
        else log maxR
pinwheelFunc (PinwheelHollow1 radius) af rf maxR alpha x y
  | r == 0 && af == 0 && rf == 0 = 1
  | r < radius = 0.0
  | otherwise =
    (r :+ 0) ** (alpha :+ rf * 2 * pi / logMaxR) *
    exp (0 :+ af * (thetaFunc x y))
  where
    r = rFunc x y
    logMaxR =
      if maxR == 1
        then 1
        else log maxR

{-# INLINE pinwheelPI #-}
pinwheelPI :: Double -> Double -> Double -> Double -> Int -> Int -> Complex Double
pinwheelPI af rf maxR alpha x y
  | r == 0 = 0
  | otherwise =
    (((r) :+ 0) ** (alpha :+ (rf * 2 * pi / (log maxR)))) *
    exp (0 :+ ((af) * theta))
  where
    r = (sqrt . fromIntegral $ x ^ (2 :: Int) + y ^ (2 :: Int))
    theta = pi  + angleFunctionRad (fromIntegral x) (fromIntegral y)

pinwheelFilter ::
     DFTPlan
  -> PinwheelParams
  -> IO (DFTPlan, Array U DIM4 (Complex Double), Array U DIM4 (Complex Double))
pinwheelFilter plan (PinwheelParams rows cols alpha radialPeriod angularFreqs radialFreqs) = do
  let numAngularFreq = L.length angularFreqs
      numRadialFreq = L.length radialFreqs
      filterVec =
        VG.convert . toUnboxed . computeS . makeFilter2D $
        R.traverse2
          (fromListUnboxed (Z :. L.length angularFreqs) angularFreqs)
          (fromListUnboxed (Z :. L.length radialFreqs) radialFreqs)
          (\(Z :. numAngularFreq') (Z :. numRadialFreq') ->
             (Z :. numAngularFreq' :. numRadialFreq' :. cols :. rows))
          (\f1 f2 (Z :. af :. rf :. i :. j) ->
             pinwheel
               (f1 (Z :. af))
               (f2 (Z :. rf))
               radialPeriod
               alpha
               (i - center cols)
               (j - center rows))
      filterPIVec =
        VG.convert . toUnboxed . computeS . makeFilter2D $
        R.traverse2
          (fromListUnboxed (Z :. L.length angularFreqs) angularFreqs)
          (fromListUnboxed (Z :. L.length radialFreqs) radialFreqs)
          (\(Z :. numAngularFreq') (Z :. numRadialFreq') ->
             (Z :. numAngularFreq' :. numRadialFreq' :. cols :. rows))
          (\f1 f2 (Z :. af :. rf :. i :. j) ->
             pinwheelPI
               (f1 (Z :. af))
               (f2 (Z :. rf))
               radialPeriod
               alpha
               (i - div cols 2)
               (j - div rows 2))
      forwardPlanID =
        DFTPlanID DFT1DG [numAngularFreq, numRadialFreq, cols, rows] [2, 3]
      backwardPlanID =
        DFTPlanID IDFT1DG [numAngularFreq, numRadialFreq, cols, rows] [2, 3]
  lock <- getFFTWLock
  (plan1, filterVecF) <-
    dft1dGPlan
      lock
      plan
      [numAngularFreq, numRadialFreq, cols, rows]
      [2, 3]
      filterVec
  (newPlan, _) <-
    idft1dGPlan
      lock
      plan1
      [numAngularFreq, numRadialFreq, cols, rows]
      [2, 3]
      filterVecF
  filterPIVecF <-
    dftExecute
      newPlan
      (DFTPlanID DFT1DG [numAngularFreq, numRadialFreq, cols, rows] [2, 3])
      filterPIVec
  return
    ( newPlan
    , fromUnboxed (Z :. numAngularFreq :. numRadialFreq :. cols :. rows) .
      VG.convert $
      filterVecF
    , fromUnboxed (Z :. numAngularFreq :. numRadialFreq :. cols :. rows) .
      VG.convert $
      filterPIVecF)

{-# INLINE frequencyDomainMultiply #-}
frequencyDomainMultiply ::
     (Source r (Complex Double))
  => Array U DIM4 (Complex Double)
  -> Array r DIM2 (Complex Double)
  -> Array D DIM4 (Complex Double)
frequencyDomainMultiply filterF input =
  R.traverse2
    filterF
    input
    const
    (\f1 f2 idx@(Z :. af :. rf :. i :. j) -> f1 idx * f2 (Z :. i :. j))

{-# INLINE convolvePinwheel #-}
convolvePinwheel ::
     (Source r (Complex Double))
  => DFTPlan
  -> Array U DIM4 (Complex Double)
  -> Array r DIM2 (Complex Double)
  -> IO (Array U DIM4 (Complex Double))
convolvePinwheel plan filterF input =
  fmap (fromUnboxed (extent filterF) . VG.convert) .
  dftExecute
    plan
    (DFTPlanID IDFT1DG (L.reverse . listOfShape . extent $ filterF) [2, 3]) .
  VG.convert . toUnboxed . computeUnboxedS $
  frequencyDomainMultiply filterF input


{-# LANGUAGE AllowAmbiguousTypes, ConstraintKinds #-}

module GPLVM.GaussianProcess
       ( GaussianProcess (..)
       , GPTrainingData (..)
       , PosteriorSample (..)
       , gpToPosteriorSample
       , kernelGP
       ) where

import Universum hiding (transpose, Vector)

import Control.Lens (makeLenses)

import Data.Array.Repa
import Data.Array.Repa.Repr.Unboxed (Unbox)
import Numeric.LinearAlgebra.Repa hiding (Matrix, Vector)
import System.Random (Random, RandomGen, mkStdGen)

import GPLVM.Types (InputObservations(..), Matrix, Vector, unInputObs)
import GPLVM.Util

data GaussianProcess a = GaussianProcess
    { -- | we define gaussian process by some kernel function
      _kernelGP :: Vector D a -> Vector D a -> Matrix D a
    }

data GPTrainingData a = GPTrainingData
    { _inputTrain  :: Vector D a  -- ^ input training data
    , _outputTrain :: Vector D a  -- ^ output training data
    }

makeLenses ''GaussianProcess
makeLenses ''GPTrainingData

newtype PosteriorSample a = PosteriorSample
    { unSample :: Matrix D a  -- ^ posterior sample
    }

-- | Constraint kind required to getting a posterior sample by a given GP
type GPConstraint a =
  ( Field a
  , Random a
  , Unbox a
  , Floating a
  , Eq a
  )

-- | Main GP function: get a posterior sample by some kernel function, input observations and training data

gpToPosteriorSample
  :: forall a g.
  GPConstraint a
  => InputObservations a        -- ^ input observations
  -> GaussianProcess a          -- ^ kernel function
  -> GPTrainingData a           -- ^ training data
  -> Int                        -- ^ number of samples
  -> Maybe (PosteriorSample a)  -- ^ posterior functional prior
gpToPosteriorSample (InputObservations observe@(ADelayed (Z :. len) _)) gP trainingData sampleNumber = do
  -- | kernel applied to input test points (so-called K_ss)
  let covarianceMatrix = kernel observe observe

  -- | kernel applied to input training points (so-called K)
  let trainingKernel = kernel inputTrain' inputTrain'

  -- | Cholesky decomposition applied to kernel of training points (:), so-called L
  let cholK = cholSH $ trainingKernel +^
                (smap (* 0.00005) . identD . size . extent $ inputTrain')

  -- | covariance between test points and input training points (so-called K_s)
  let testPointMean = kernel observe inputTrain'

  -- | (roots of L * x = K_s)
  cholKSolve <- delay <$> linearSolveS cholK testPointMean

  -- | solve linear system for output training points
  cholKSolveOut <- delay <$> linearSolveS cholK (transposeMatrix $ toMatrix outputTrain' len)

  -- | compute mean
  let mean = (transposeMatrix cholKSolveOut) `mulD` cholKSolve

  -- | posterior
  let postF' = cholSH $
                     covarianceMatrix +^
                     ((smap (* 1.0e-6) (identD len)) -^
                     (transposeMatrix cholKSolve) `mulD` cholKSolve)

  -- | posterior sample
  return $ (PosteriorSample $ mean +^ (functionalPrior postF' (mkStdGen 1) sampleNumber))
    where
       kernel = gP ^. kernelGP
       inputTrain' = trainingData ^. inputTrain
       outputTrain' = trainingData ^. outputTrain
       mulD m n = delay $ m `mulS` n
       functionalPrior matrix@(ADelayed (Z :. rows :. cols) _) gen sampleNumber =
         delay $ matrix `mulS` randomCoeffs
         where
           randomCoeffs = randomMatrixD (mkStdGen (-4)) (rows, sampleNumber)


module Fuml.Optimisation.SGD where

import Numeric.LinearAlgebra
import Data.Random

data SGDOpts = SGDOpts
  { learning_rate :: Double
  , batch_size :: Int
  , epochs :: Int
  }

sgdBatch :: SGDOpts -> (Vector Double -> Vector Double) -> Vector Double -> Vector Double
sgdBatch opts grad p = p - scale (learning_rate opts) (grad p)

sgdEpoch :: SGDOpts -> ([a] -> Vector Double -> Vector Double) -> [a] -> Vector Double -> RVar (Vector Double)
sgdEpoch opts grad pts' p' = do
  pts <- shuffle pts'
  return $ go p' $ inGroupsOf (batch_size opts) pts
    where go p [] = p
          go p (batch:bs) = go (sgdBatch opts (grad batch) p) bs

inGroupsOf :: Int -> [a] -> [[a]]
inGroupsOf _ [] = []
inGroupsOf n xs = let (these,those) = splitAt n xs
                  in these : inGroupsOf n those


module Main where

import Net
import Data.Traversable
import Numeric.LinearAlgebra
import Numeric.LinearAlgebra.Devel
import qualified Data.Vector.Storable as V
import System.Random
import Data.List
import Control.Monad

testNet :: DoubleNet -> [Vector Double] -> [Vector Double] -> Double
testNet n ins ans =
  let outs :: [Vector Double]
      outs = map (runNet n) ins
  in sum (zipWith (\x y -> norm_1 $ add x (scale (-1) y)) outs ans)

rankNets :: [DoubleNet] -> IO [(Double,DoubleNet)]
rankNets nets =
  do values <- mapM (\_ -> do x <- randomRIO (0,10)
                              return $ vector [x]) [1..100]
     let answers = map sin values
         scoredNets = map (\x -> (testNet x values answers,x)) nets
     return $ sortOn fst scoredNets

plotSome :: DoubleNet -> IO ()
plotSome n =
  do let pts = map (\x -> (x/1000)*10) $ [1..1000]
     mapM_ (\x -> putStrLn $ (show $ runNet n (vector [x])) ++ " " ++ (show $ (sin x)/2)) pts

sinFunc = DiffFunction sin (\x y -> y*cos x)

errFunc :: Double -> DifferentiableFunction (Vector Double) Double
errFunc c = DiffFunction (\v -> ((atIndex v 0) - c)^2) (\v w -> 2*((atIndex v 0) - c)*(atIndex w 0))

iterateNet :: Int -> DoubleNet -> IO DoubleNet
iterateNet 0 net = return net
iterateNet steps net =
  do let n = 10
     xs <- replicateM n $ randomRIO (0,20)

     --let x = 0.5
     let loop loopSteps currNet =
           if loopSteps ==0 then
             return currNet
           else
             do grads <- forM xs (\x -> do let (val,err,grad) = runAndDiffWithError currNet (vector [x]) (errFunc $ (sin x)/2)
                                         --print (x,val,(sin x)/2,err)
                                           return grad)
                let grad = foldl1 (zipWith (\a b -> (add b a))) grads
                let newNet = updateNet net grad (-0.0005)
                loop (loopSteps - 1) newNet
     loop 5 net >>= iterateNet (steps-1)

main :: IO ()
main = do randNet <- randomSimpleNet 1 1 100
          finalNet <- iterateNet 10000 randNet
          plotSome finalNet


{-# LANGUAGE OverloadedStrings #-}

module Main where

import Numeric.LinearAlgebra

main :: IO ()
main = putStrLn $ show $ vector [1,2,3] * vector [3,0,-2]


{-# LANGUAGE OverloadedStrings #-}

module Main where

import Prelude hiding (readFile, drop)
import Control.Monad
import Control.Monad.IO.Class
import Control.Exception
import Foreign.C.Types
import System.Random
import Numeric.LinearAlgebra
import Numeric.LinearAlgebra.Data
import SDL
import Debug.Trace
import qualified Linear
import Data.List.Split (chunksOf)

import Neural
import MNIST

-- Generates a layer with a given number of neurons with inputCount weights.
-- Weights are randomised between 0 and 1.
genLayer :: Int -> Int -> IO Layer
genLayer neurons inputCount = do
  weights <- replicateM (neurons * inputCount) (randomRIO (-0.2, 0.2) :: IO R)
  biases <- replicateM (neurons) (randomRIO (-0.2, 0.2) :: IO R)
  return $ Layer (matrix inputCount weights) (vector biases)

-- Generates a network with the specified number of neurons in each layer,
-- and the number of inputs to the network.
-- Initialised with random weights and biases as in genLayer.
genNetwork :: [Int] -> Int -> IO Network
genNetwork layers inputCount = sequence . map (uncurry genLayer) $ layerDefs
  where
    -- Layer definitions in the form (neuronCount, inputCount)
    layerDefs = zip layers (inputCount:layers)

-- A test network
testNetwork :: IO Network
testNetwork = do
  l1 <- genLayer 5 3
  l2 <- genLayer 5 5
  l3 <- genLayer 1 5
  return [l1, l2, l3]

-- Another test network
testNetwork2 :: Network
testNetwork2 = [hiddenLayer, outputLayer]
  where
    hiddenLayer = Layer (matrix 2 [0.15, 0.2, 0.25, 0.3]) (vector [0.35, 0.35])
    outputLayer = Layer (matrix 2 [0.4, 0.45, 0.5, 0.55]) (vector [0.6, 0.6])

-- Test inputs for testNetwork2
testInputs = vector [0.05, 0.1]

-- Training data for testInputs
trainingData = vector [0.01, 0.99]

-- Train neural network
trainNet = do
  (imageCount, width, height, imageData) <- readMNISTImages "train-images-idx3-ubyte"
  (labelCount, labels) <- readMNISTLabels "train-labels-idx1-ubyte"

  let image = head imageData
  let input = vector (map fromIntegral image)
  let correctNumber = fromIntegral . head $ labels
  let training = vector $ map (\i -> if i == correctNumber then 1.0 else 0.0) [1..10]

  hiddenLayer <- genLayer 20 (width * height)
  outputLayer <- genLayer 10 20

  let initialNetwork = [hiddenLayer, outputLayer]

  let loop net i err =
        case i > 1000 || err < 0.1 of
          True  -> net
          False -> trace (show $ sumElements (last error)) $ loop nextNet (i+1) (sumElements (last error))
            where
              output = evalNetwork net input
              error = evalNetworkError net output input training
              nextNet = gradientDescent net input (map fst output) error 0.1

  let trainedNet = loop initialNetwork 0 1000

  return (input, initialNetwork, trainedNet, training)

-- Train neural network
trainNetM = do
  (imageCount, width, height, imageData) <- readMNISTImages "train-images-idx3-ubyte"
  (labelCount, labels) <- readMNISTLabels "train-labels-idx1-ubyte"

  let image = head imageData
  let input = fromRows [vector (map fromIntegral image)]
  let correctNumber = fromIntegral . head $ labels
  let training = fromRows [vector $ map (\i -> if i == correctNumber then 1.0 else 0.0) [1..10]]

  hiddenLayer <- genLayer 20 (width * height)
  outputLayer <- genLayer 10 20

  let initialNetwork = [hiddenLayer, outputLayer]

  let loop net i err =
        case i > 100000 || err < 0.1 of
          True  -> net
          False -> trace (show $ sumElements (last error)) $ loop nextNet (i+1) (sumElements (last error))
            where
              output = evalNetworkM net input
              error = evalNetworkErrorM net output input training
              nextNet = gradientDescentM net input (map fst output) error 0.1

  let trainedNet = loop initialNetwork 0 1000

  return (input, initialNetwork, trainedNet, training)

-- Make a window and show MNIST image
window image width height = do
  initializeAll

  window <- createWindow "SDL Application" defaultWindow
  renderer <- createRenderer window (-1) defaultRenderer

  -- Convert to sdl texture
  texture <- loadTextureMNIST renderer width height image

  SDL.rendererDrawColor renderer $= Linear.V4 255 255 255 255
  SDL.clear renderer

  copy renderer texture Nothing Nothing

  present renderer

  let loop = do
        pollEvents

        keyState <- getKeyboardState

        case keyState ScancodeEscape of
          True -> (return ()) :: IO ()
          False -> loop

  loop

  destroyTexture texture
  destroyRenderer renderer
  destroyWindow window

-- Train a network using a given set of inputs and training values for the given
-- number of epochs using the given learning rate.
trainNetwork :: Network -> [Vector R] -> [Vector R] -> Int -> Int -> R -> Network
trainNetwork initialNet input trainingValues epochs batchSize learnRate =
  train initialNet 0
    where
      train net epoch = case epoch >= epochs of
        True  -> net
        False -> train nextNet (epoch + 1)
          where
            --inputBatch = fromRows input
            --trainingBatch = fromRows trainingValues
            inputBatches = map fromRows . chunksOf batchSize $ input
            trainingBatches = map fromRows . chunksOf batchSize $ trainingValues
            nextNet = foldl runBatch net (zip inputBatches trainingBatches)
            --nextNet = net
            runBatch network (inputBatch, trainingBatch) =
              gradientDescentM net inputBatch (map fst output) error learnRate
                where
                  output = evalNetworkM net inputBatch
                  error = evalNetworkErrorM net output inputBatch trainingBatch

toOutput = \n -> vector . map (\x -> if n == x then 1.0 else 0.0) $ [0..9]

fromOutput :: Vector R -> Int
fromOutput v = guess
  where
    list = zip [0..] . toList $ v
    possibles = filter ((>=0.5) . snd) list
    guess = case length possibles of
              0 -> -1
              1 -> fst . head $ possibles
              _ -> -2

runTest :: Network -> ([CUChar], Int) -> Bool
runTest network (image, label) = guess == label
  where output = fst . last $ evalNetwork network $ vector . map fromIntegral $ image
        guess = fromOutput output

--runTest :: Network -> ([CUChar], Int) -> IO ()
--runTest network (image, label) = do
--  let output = fst . last $ evalNetwork network $ vector . map fromIntegral $ image
--  let guess = fromOutput output
--  let expected = label
--  putStr "Guess: "
--  putStr . show $ guess
--  putStr ",\tExpected: "
--  putStr . show $ label
--  putStr $ if guess == expected then ".\tCorrect! " else ".\tINCORRECT!"
--  putStrLn ""

main = do

  (trainingImageCount, trainingWidth,
   trainingHeight, trainingImageData)  <- readMNISTImages "train-images-idx3-ubyte"
  (trainingLabelCount, trainingLabels) <- readMNISTLabels "train-labels-idx1-ubyte"

  (testImageCount, testWidth,
   testHeight, testImageData)  <- readMNISTImages "t10k-images-idx3-ubyte"
  (testLabelCount, testLabels) <- readMNISTLabels "t10k-labels-idx1-ubyte"

  assert (trainingWidth == testWidth && trainingHeight == testHeight) (return ())

  let input = take 60000 . map (vector . map fromIntegral) $ trainingImageData
  let trainingData = take 60000 . map toOutput $ trainingLabels

  initialNetwork <- genNetwork [20,10] (trainingWidth * trainingHeight)

  let trainedNetwork = trainNetwork initialNetwork input trainingData 2 10 3.0

  let testImage = head testImageData
  let testInput = vector . map fromIntegral $ testImage
  let testOutput = head testLabels --toOutput $ head testLabels

  --window testImage testWidth testHeight

  --let output = fst . last $ evalNetwork trainedNetwork testInput
  --let filteredOutput = map (>=0.5) . toList $ output

  --print $ fromOutput output
  --print $ testOutput

  let tests = map (runTest trainedNetwork) $ zip testImageData testLabels

  let success = length . filter (==True) $ tests
  let total = length tests

  putStr . show $ success
  putStr "/"
  putStrLn . show $ total

  --flip mapM_ (take 100 $ zip testImageData testLabels) runTest
  --  where
  --    runTest (imageData, label) = output
  --      where
  --        output2 = evalNetwork trainedNetwork
  --        output = putStrLn "test"
  --      --let output = evalNetwork trainedNetwork (vector . map fromIntegral $ imageData)
  --      --in do
  --      --  putStr "Guess: "
  --      --  putStr . show . fromOutput $ output


module GPLVM.PPCA
       ( PPCA (..)
       , makePPCA
       ) where

import Prelude (isNaN, log)

import Universum hiding (All, Any, Vector, map, toList, transpose, (++))

import qualified Universum as U

import GPLVM.Types
import GPLVM.Util

import Data.Array.Repa
import Data.Array.Repa.Repr.ForeignPtr
import Data.List ((\\))
import Numeric.LinearAlgebra.Repa hiding (Matrix, Vector, diag)
import System.Random

data PPCA = PPCA
  {  _noMissedData        :: Bool
   , _learningData        :: Matrix D Double
   , desiredDimension     :: Int
   , stopParameter        :: Either Int Double
   , _variance            :: Double
   , _W                   :: Matrix D Double
   , _finalExpLikelihood  :: Double
   , _restoredMatrix      :: Maybe (Matrix D Double)
   }

makePPCA
  :: (HasCallStack) => RandomGen gen
  => Matrix D Double   -- ^ learning vectors
  -> Int               -- ^ desired dimention
  -> Either Int Double -- ^ number of iterations or
                       --   the value of stop difference of the likelihood expectation between interations of EM.
  -> gen
  -> PPCA
makePPCA leaningVectors desiredDimension stopParameter generator =
  let _learningData@(ADelayed (Z :. n :. _) _) = leaningVectors
      initMatrix = randomMatrixD generator (n, desiredDimension)
      initDispersion = fromIntegral $ fst . next $ generator
      _noMissedData = not $ isNaN $ runIdentity $ sumAllP _learningData
      (_W, _variance, _finalExpLikelihood, _restoredMatrix) =
        if _noMissedData
        then emStepsFast _learningData initMatrix initDispersion stopParameter
        else emStepsMissed _learningData initMatrix initDispersion stopParameter
  in PPCA{..}

-- TODO: add comments to the most of following local binds
-- Most of them does not have informative names because they are just intermediate calculations.

emStepsFast
  :: (HasCallStack) => Matrix D Double
  -> Matrix D Double
  -> Double
  -> Either Int Double
  -> (Matrix D Double, Double, Double, Maybe (Matrix D Double))
emStepsFast learnMatrix@(ADelayed (Z :. _ :. n) _) initMatrix@(ADelayed (Z :. d :. _) _) initVariance stopParam =
  let learnMatrixCentered = transpose $ substractMean (transpose learnMatrix)

      stepOne :: (HasCallStack) => Array D DIM2 Double -> Double -> Int -> (Matrix D Double, Double, Double, Maybe (Matrix D Double))
      stepOne (!oldW) oldVariance iteration =
        let m = delay $ mapDiagonal (oldVariance +) $ (transpose oldW) `mulS` oldW
            invM = invS m
            expEznZ = delay $ invM `mul` (transpose oldW `mulS` learnMatrixCentered)
            expXnEZnZ = learnMatrixCentered `mulS` (transpose expEznZ)
            expZtrZ = computeS $ (map (*((fromIntegral n)*oldVariance)) $ delay invM) +^ (expEznZ `mulS` (transpose $ expEznZ)) :: Matrix F Double
        in stepTwo (expEznZ, expZtrZ, expXnEZnZ) oldW oldVariance iteration

      stepTwo :: (HasCallStack) => (Matrix D Double, Matrix F Double, Matrix F Double) -> Matrix D Double -> Double -> Int -> (Matrix D Double, Double, Double, Maybe (Matrix D Double))
      stepTwo (expEznZ, expZtrZ, expXnEZnZ) oldW oldVariance iteration =
        let newW = delay $ expXnEZnZ `mul` (inv expZtrZ)
            u = chol $ trustSym expZtrZ
            wr = newW `mulS` (transpose u)
            totalSum = sumAllS $ map (^ (2 :: Int)) learnMatrixCentered
            secondDenominatorInVariance = (*2) $ sumAllS $ expEznZ *^ (delay $ (transpose newW) `mulS` learnMatrixCentered)
            thirdDenominatorInVariance  = sumAllS $ map (^(2 :: Int)) wr
            newVariance = (totalSum - secondDenominatorInVariance + thirdDenominatorInVariance)/(fromIntegral $ n*d)
            maxDiffNewOldW = foldAllS max 0.0 $ map abs $ oldW -^ newW
            diffVariance = abs $ newVariance - oldVariance
            newC = mapDiagonal ((+) newVariance) $ newW `mulS` (transpose newW)
            newInvM = delay $ invS $ mapDiagonal (newVariance +) $ (transpose newW) `mulS` newW
            newInvC = (mapDiagonal (+(1/newVariance))
                        (map (\x -> -x/newVariance) $ (delay $ newW `mulS` newInvM) `mulS` (transpose newW)))
            expLogLikelihood = (-1.0)*((fromIntegral n)/2.0)*((fromIntegral d)*(log $ 2*pi)
              + (log $ detS newC)
              + (trace2S $ computeS $ map (/(fromIntegral $ n-1)) $ newInvC `mulS` (delay $ learnMatrixCentered `mulS` (transpose learnMatrixCentered))))
        in case stopParam of
              Left maxIteration -> if maxIteration > iteration then stepOne newW newVariance (iteration + 1) else (newW,newVariance, expLogLikelihood, Nothing)
              Right stopDiffBetweenIterations -> if (max diffVariance maxDiffNewOldW) > stopDiffBetweenIterations then stepOne newW newVariance (iteration + 1) else (newW,newVariance, expLogLikelihood, Nothing)

  in stepOne initMatrix initVariance 0

emStepsMissed
  :: (HasCallStack) => Matrix D Double
  -> Matrix D Double
  -> Double
  -> Either Int Double
  -> (Matrix D Double, Double, Double, Maybe (Matrix D Double))
emStepsMissed learnMatrix@(ADelayed (Z :. _ :. n) _) initMatrix@(ADelayed (Z :. d :. _) _) initVariance stopParam =
  let initMu = delay $ fromListUnboxed (Z:.d:.1) $ replicate d 0.0 -- extend (Any :. (1 :: Int)) $ meanColumnWithNan learnMatrix

      stepOne :: (HasCallStack) => Array D DIM2 Double -> Array D DIM2 Double -> Double -> Int -> (Matrix D Double, Double, Double, Maybe (Matrix D Double))
      stepOne oldMu (!oldW) oldVariance iteration =
        let yi i = extend (Any :. (1 :: Int)) $ slice learnMatrix (Any :. (i :: Int))
            yList i = toList (yi i)
            unknownIndices i =
              let zipped = zip [0..] (yList i)
              in U.map fst $ filter (isNaN . snd) zipped
            yPi i  = deleteRows (unknownIndices i) (yi i)
            muP i =  deleteRows (unknownIndices i) oldMu
            oldWP i = deleteRows (unknownIndices i) oldW
            mP i = mapDiagonal (+ oldVariance) $ (transpose (oldWP i)) `mulS` (oldWP i)
            invMP i = delay $ invS $ mP i
            expXi i = delay $ (delay $ (delay $ invMP i) `mulS` (transpose (oldWP i))) `mulS` (yPi i -^ muP i)
        in stepTwo ([yPi, expXi, invMP], unknownIndices, oldW) oldVariance iteration

      stepTwo :: (HasCallStack) => ([Int -> Matrix D Double], Int -> [Int], Matrix D Double) -> Double -> Int -> (Matrix D Double, Double, Double, Maybe (Matrix D Double))
      stepTwo ([yP, expXi, invMP], unknownIndices, oldW) oldVariance iteration =
        let expX = foldl1 (\acc i -> acc ++ i) $ U.map expXi [0..(n-1)]
            newMu = extend (Any :. (1 :: Int)) $ meanColumnWithNan $ (learnMatrix -^ (oldW `mulS` expX))

            newW =
              let yj j = extend (Any :. (1 :: Int) :. All) $ slice learnMatrix (Any :. (j :: Int) :. All)
                  yJList j = toList (yj j)
                  unknownColumns j =
                    let zipped = zip [0..] (yJList j)
                    in U.map fst $ filter (isNaN . snd) zipped
                  knownColumns j = [0..(n-1)] \\ (unknownColumns j)
                  expXPj j = deleteColumns (unknownColumns j) expX
                  sumInvMP j = sumListMatrices $ U.map invMP (knownColumns j)
                  gj j = ((expXPj j) `mulS` (transpose $ expXPj j)) +^ (map (*oldVariance) (sumInvMP j))
                  yPj j = deleteColumns (unknownColumns j) (yj j)
                  expXtrXj j = (expXPj j) `mulS` (transpose (map (\x -> x - (newMu ! (Z:.j:.0))) (yPj j)))
                  newWj j = (gj j) `solveS` (delay $ expXtrXj j)
              in transpose $ foldl1 (\acc j -> acc ++ j) $ U.map (\x -> delay $ newWj x) [0..(d-1)]

            newVariance =
              let newWPi i = deleteRows (unknownIndices i) newW
                  newMuP i = deleteRows (unknownIndices i) newMu
                  varianceStep i = sumAllS $
                    (map (^(2 :: Int)) (yP i -^ ((newWPi i) `mulS` (expXi i)) -^ (newMuP i)))
                    +^ (map (*oldVariance) $ diag $ delay $ (delay $ (delay $ newWPi i) `mulS` (invMP i)) `mulS` (transpose $ newWPi i))
                  knownVars = foldAllS (\acc x -> if isNaN x then acc else acc + 1.0) 0.0 learnMatrix
              in (sum $ U.map varianceStep [0..(n-1)])/knownVars

            expLogLikelihood =
              let newMuP i = deleteRows (unknownIndices i) newMu
                  newWPi i = deleteRows (unknownIndices i) newW
                  yCenteredP i = (yP i) -^ (newMuP i)
                  yCenteredProd i = delay $ (yCenteredP i) `mulS` (transpose $ yCenteredP i)
                  invMY i = mapDiagonal (+newVariance) $ (newWPi i) `mulS` (transpose $ newWPi i)
                  knownIndices i = [0..(d-1)] \\ (unknownIndices i)
                  expLogLikelihoodStep i = (fromIntegral $ length $ knownIndices i)*(log $ 2*pi) + (log $ detS $ invMY i) +
                    (trace2S $ computeS $ delay $ (invMY i) `solveS` (yCenteredProd i))

              in (-1.0)*(sum $ U.map expLogLikelihoodStep [0..(n-1)])/2.0

            restoredData =
              let muX = extend (Any :. (1 :: Int)) $ meanColumn (transpose expX)
                  finalMu = slice (newMu +^ (newW `mulS` muX)) (Any :. (0 :: Int))
                  wtrw = (transpose newW) `mulS` newW
                  factor1 = newW `mulS` (delay $ inv wtrw)
                  factor2 = computeS $ mapDiagonal (+newVariance) wtrw
                  restoredCentered =  factor1 `mul` factor2 `mul` (computeS $ expX)
              in (delay restoredCentered) +^ (extend (Any :. (n :: Int)) finalMu)

            diffVariance = abs $ newVariance - oldVariance
            maxDiffNewOldW = foldAllS max 0.0 $ map abs $ oldW -^ newW

        in case stopParam of
              Left maxIteration ->
                if maxIteration > iteration
                then stepOne newMu newW newVariance (iteration + 1)
                else (newW,newVariance, expLogLikelihood, Just restoredData)
              Right stopDiffBetweenIterations ->
                if (max diffVariance maxDiffNewOldW) > stopDiffBetweenIterations
                then stepOne newMu newW newVariance (iteration + 1)
                else (newW,newVariance, expLogLikelihood, Just restoredData)
  in stepOne initMu initMatrix initVariance 0


{-# LANGUAGE LambdaCase #-}
-- |
-- Module: Analysis.PerfModel
--
-- An experimental approach to obtaining strengths, which takes Elo (or
-- Elo-like) ratings as derived from a performance model. With a performance
-- model in hand, we can use https://stats.stackexchange.com/a/44142 to
-- obtain the victory probability of a 1500-rated racer given a field of
-- racers, and use that as a measure of strength.
--
-- According to the model, player performances follow a gamma distribution,
-- with the outcome 0 corresponding to an ideal performance. The shape
-- parameter of the distribution is fixed, while the rate parameter varies
-- across players. A higher rate parameter amounts to results that are both
-- more consistent and closer to the ideal. (Pinning both mean and standard
-- deviation of player performance to a single parameter is the main
-- simplifying assumption of the model.)
--
-- Note that the performance model winning probabilities for pairs of players
-- agree exactly with the conventional Elo formula if the shape parameter of
-- the gamma distribution is taken to be 1.
module Analysis.PerfModel
    ( perfModelStrength
    , perfModelTopStrength
    , example
    ) where

import Orbital (orbitalDistr, initialRating)
import qualified Util.Combinations as Util

import qualified Numeric.Integration.TanhSinh as Integration
import Statistics.Distribution

import Data.Profunctor
import qualified Control.Foldl as L
import qualified Data.IntMap.Strict as IntMap
import Data.IntMap.Strict (IntMap, (!))
import qualified Data.IntSet as IntSet
import Data.IntSet (IntSet)
import qualified Control.Foldl as L

-- | PDF for a race victory against a field of racers.
raceWinPDF
    :: Int       -- ^ Gamma shape.
    -> [Double]  -- ^ Ratings of the field of opponents.
    -> Double    -- ^ Rating of the racer.
    -> Double    -- ^ Performance.
    -> Double
raceWinPDF gsh rs r t = density (orbitalDistr gsh r) t
    * L.fold (lmap winVersus L.product) rs
    where
    winVersus r' = 1 - cumulative (orbitalDistr gsh r') t

-- | Integrating the race win PDF gives the likelihood of victory, whose
-- reciprocal we use as a strength metric.
perfModelStrength
    :: Int       -- ^ Gamma shape.
    -> [Double]  -- ^ Ratings of the field of opponents.
    -> Double
perfModelStrength gsh rs = 1 / integ
    where
    integ = Integration.result . Integration.absolute 1e-6 $
        Integration.nonNegative Integration.trap (raceWinPDF gsh rs initialRating)

-- TODO: positionPDF and positionPDFPre are still the same as in
-- Analysis.PerfModel.Reference.Precompute . Conside redefining them so that
-- topPDF can share code with them.

-- | Like 'raceWinPDF', but for an arbitrary position on the scoreboard.
positionPDF
    :: Int       -- ^ Gamma shape.
    -> Int       -- ^ Target position.
    -> [Double]  -- ^ Ratings of the field of opponents.
    -> Double    -- ^ Rating of the racer.
    -> Double    -- ^ Performance.
    -> Double
positionPDF gsh pos rs r t = density (orbitalDistr gsh r) t
    * L.fold L.sum (L.fold L.product . toFactors <$> combos)
    where
    n = length rs
    combos = Util.combinations n (pos-1) [1..n]
    loseVersus r' = cumulative (orbitalDistr gsh r') t
    -- winVersus r' = 1 - loseVersus r'
    lvs = IntMap.fromList (zip [1..n] (loseVersus <$> rs))
    wvs = (1 -) <$> lvs
    toFactors (above, below) = map (lvs !) above ++ map (wvs !) below

-- | A variant of 'positionPDF' that takes precomputed probabilities.
positionPDFPre :: Int                            -- ^ Target position.
               -> (IntMap Double, IntMap Double) -- ^ Win and loss probabilities
                                                 -- (given a probe laptime) for
                                                 -- for each racer on the
                                                 -- scoreboard.
               -> Double                         -- ^ Probability density of the
                                                 -- probe laptime.
               -> Double
positionPDFPre pos (lvs, wvs) p =
    p * L.fold L.sum (L.fold L.product . toFactors <$> combos)
    where
    n = IntMap.size lvs
    combos = Util.combinations n (pos-1) [1..n]
    toFactors (above, below) = map (lvs !) above ++ map (wvs !) below

-- | Like 'positionPDF', but for a top-N finish.
topPDF
    :: Int       -- ^ Gamma shape.
    -> Int       -- ^ Target position.
    -> [Double]  -- ^ Ratings of the field of opponents.
    -> Double    -- ^ Rating of the racer.
    -> Double    -- ^ Performance.
    -> Double
topPDF gsh pos rs r t = (probeDensity *) . L.fold L.sum $
    L.fold L.sum . partials <$> validPositions
    where
    validPositions = zipWith const [1..pos] (r : rs)
    loseVersus r' = cumulative (orbitalDistr gsh r') t
    lvs = IntMap.fromList (zip [1..] (loseVersus <$> rs))
    wvs = (1 -) <$> lvs
    probeDensity = density (orbitalDistr gsh r) t
    partials = Util.processCombsInt alg (length rs) . subtract 1
    alg = \case
        Util.LeafF Nothing -> 0
        Util.LeafF (Just rest) ->
            L.fold L.product ((wvs !) <$> IntSet.toList rest)
        Util.FlowerF a rest -> (lvs ! a)
            * L.fold L.product ((wvs !) <$> IntSet.toList rest)
        Util.BranchF a bs -> (lvs ! a) * L.fold L.sum bs

-- | Like 'perfModelStrength', but for a top-N finish. Note that the
-- computational cost grows quickly with the length of the list of ratings.
perfModelTopStrength
    :: Int       -- ^ Gamma shape.
    -> Int       -- ^ Target position.
    -> [Double]  -- ^ Ratings of the field of racers.
    -> Double
perfModelTopStrength gsh pos rs = 1 / integ
    where
    integ = Integration.result . Integration.absolute 1e-6 $
        Integration.nonNegative Integration.trap (topPDF gsh pos rs initialRating)

example :: [Double]
example = [2200,2100,1900,1870,1850,1600]
--example = [2200,2100,1900,1870,1850,1600,1590,1570,1510,1420,1370,1350,1225]

-- >$> perfModelStrength [2200,2100,1900,1870,1850]
--
-- >$> :set +s
--

-- >$> perfModelTopStrength 5 example


module Statistics.ANOVA (anova, fisherLSD) where
import Data.List
import Statistics.Distribution hiding (mean, variance)
import Statistics.Distribution.FDistribution
import Statistics.Distribution.StudentT
import Statistics.Sample
import qualified Data.Vector as V

{-# INLINE treatments #-}
treatments :: (Ord a,Eq a) => V.Vector (a, Double) -> V.Vector a
treatments input = V.fromList $! nub $! V.toList $! V.map fst input --not good

treatmentMeans :: (Ord a, Eq a) => V.Vector a -> V.Vector (a, Double) -> V.Vector (a,Double)
treatmentMeans treatments input = V.map (\x -> x input) meanFuncs
  where meanFuncs = V.map treatmentMean $! treatments

treatmentMean :: (Ord a, Eq a) => a -> V.Vector (a,Double) -> (a,Double)
treatmentMean treatment input = (treatment, avg)
  where avg = mean $! V.map snd $! filterTreatment treatment $! input

treatmentSize treatment input = V.length $! filterTreatment treatment $! input

{-# INLINE filterTreatment #-}
filterTreatment :: (Ord a, Eq a) => a -> V.Vector (a, Double) -> V.Vector (a, Double)
filterTreatment treatment = V.filter (\(x,y) -> x == treatment)

sampleMean input = mean $! V.map snd input

ssTreatment :: (Ord a, Eq a) => V.Vector a -> V.Vector (a, Double) -> Double
ssTreatment treatments input = V.sum $! V.zipWith (\size mean -> size*(mean-sMean)*(mean-sMean)) tSizes tMeans
  where sMean = sampleMean input
        tMeans = V.map snd $! treatmentMeans treatments input
        tSizes = V.map (\x -> fromIntegral $! x input) $! V.map treatmentSize $! treatments --input

ssError :: (Ord a, Eq a) => V.Vector (a, Double) -> Double
ssError input = V.sum $! V.map (\x -> x input) $! V.map treatmentError $! treatments input

treatmentError treatment input = V.sum $! V.map (\x -> (x-tMean)*(x-tMean)) filteredInput
  where filteredInput = V.map snd $! filterTreatment treatment input
        tMean = snd $! treatmentMean treatment input

fValue :: (Ord a, Eq a) => V.Vector a -> V.Vector (a, Double) -> Double
fValue treatments input = ((ssTreatment treatments input)/(ssError input)) * ((sampleSize - numberOfTreatments)/(numberOfTreatments -1))
  where numberOfTreatments = fromIntegral $! V.length $! treatments
        sampleSize = fromIntegral $! V.length input

fDist treatments input = fDistribution (numberOfTreatments-1) (sampleSize-numberOfTreatments)
  where numberOfTreatments = fromIntegral $! V.length treatments
        sampleSize = fromIntegral $! V.length input

pValue input = 1 - cumulative (fDist t input) (fValue t input)
  where t = treatments input

anova pvalue input
  | pValue input < pvalue = True
  | otherwise = False

tTest :: (Ord a, Eq a) => Double -> a -> a -> V.Vector (a, Double) -> Bool
tTest pvalue treatment1 treatment2 inputs = abs (tTestT treatment1 treatment2 inputs) > critical
  where critical = quantile (studentT (degreesOfFreedom treatment1 treatment2 inputs)) (1 - pvalue/2)

tTestT :: (Ord a, Eq a) => a -> a -> V.Vector (a, Double) -> Double
tTestT treatment1 treatment2 inputs = (mean1-mean2)/sqrt(var1/size1+var2/size2)
  where mean1 = snd $! treatmentMean treatment1 inputs
        mean2 = snd $! treatmentMean treatment2 inputs
        size1 = fromIntegral $! V.length $! filterTreatment treatment1 inputs
        size2 = fromIntegral $! V.length $! filterTreatment treatment2 inputs
        var1 = variance $! V.map snd $! filterTreatment treatment1 inputs
        var2 = variance $! V.map snd $! filterTreatment treatment2 inputs

degreesOfFreedom
  :: (Ord a, Eq a) => a -> a -> V.Vector (a, Double) -> Double
degreesOfFreedom treatment1 treatment2 inputs = numerator / denominator
  where size1 = fromIntegral $! V.length $! filterTreatment treatment1 inputs
        size2 = fromIntegral $! V.length $! filterTreatment treatment2 inputs
        var1 = variance $! V.map snd $! filterTreatment treatment1 inputs
        var2 = variance $! V.map snd $! filterTreatment treatment2 inputs
        numerator = (var1/size1+var2/size2)^2
        denominator = (((var1/size1)^2)/(size1-1)) + (((var2/size2)^2)/(size2-1))

noDifferenceInMean pvalue treatment input = (:) treatment $!
                                            map fst $!
                                            filter (\(x,y) -> not y) $!
                                            map (\x -> (x, tTest pvalue treatment x input)) othertreatments
  where othertreatments = delete treatment $! V.toList $! treatments input

equivalenceList pvalue input = nub $! V.toList $! V.map (\x -> sort $! noDifferenceInMean pvalue x input) $! treatments input

fisherLSD :: (Eq a, Ord a) => Double -> V.Vector (a, Double) -> [[a]]
fisherLSD pvalue input = case (anova pvalue input) of
  False -> [V.toList $! treatments input]
  True -> equivalenceList pvalue input


{-# LANGUAGE DeriveDataTypeable, DeriveGeneric #-}
-- |
-- Module    : Statistics.Distribution.FDistribution
-- Copyright : (c) 2011 Aleksey Khudyakov
-- License   : BSD3
--
-- Maintainer  : bos@serpentine.com
-- Stability   : experimental
-- Portability : portable
--
-- Fisher F distribution
module Statistics.Distribution.FDistribution (
    FDistribution
  , fDistribution
  , fDistributionNDF1
  , fDistributionNDF2
  ) where

import Data.Aeson (FromJSON, ToJSON)
import Data.Binary (Binary)
import Data.Data (Data, Typeable)
import Numeric.MathFunctions.Constants (m_neg_inf)
import GHC.Generics (Generic)
import qualified Statistics.Distribution as D
import Statistics.Function (square)
import Numeric.SpecFunctions (
  logBeta, incompleteBeta, invIncompleteBeta, digamma)
import Data.Binary (put, get)
import Control.Applicative ((<$>), (<*>))

-- | F distribution
data FDistribution = F { fDistributionNDF1 :: {-# UNPACK #-} !Double
                       , fDistributionNDF2 :: {-# UNPACK #-} !Double
                       , _pdfFactor        :: {-# UNPACK #-} !Double
                       }
                   deriving (Eq, Show, Read, Typeable, Data, Generic)

instance FromJSON FDistribution
instance ToJSON FDistribution

instance Binary FDistribution where
    get = F <$> get <*> get <*> get
    put (F x y z) = put x >> put y >> put z

fDistribution :: Int -> Int -> FDistribution
fDistribution n m
  | n > 0 && m > 0 =
    let n' = fromIntegral n
        m' = fromIntegral m
        f' = 0.5 * (log m' * m' + log n' * n') - logBeta (0.5*n') (0.5*m')
    in F n' m' f'
  | otherwise =
    error "Statistics.Distribution.FDistribution.fDistribution: non-positive number of degrees of freedom"

instance D.Distribution FDistribution where
  cumulative = cumulative

instance D.ContDistr FDistribution where
  density d x
    | x <= 0    = 0
    | otherwise = exp $ logDensity d x
  logDensity d x
    | x <= 0    = m_neg_inf
    | otherwise = logDensity d x
  quantile = quantile

cumulative :: FDistribution -> Double -> Double
cumulative (F n m _) x
  | x <= 0       = 0
  | isInfinite x = 1            -- Only matches +∞
  | otherwise    = let y = n*x in incompleteBeta (0.5 * n) (0.5 * m) (y / (m + y))

logDensity :: FDistribution -> Double -> Double
logDensity (F n m fac) x
  = fac + log x * (0.5 * n - 1) - log(m + n*x) * 0.5 * (n + m)

quantile :: FDistribution -> Double -> Double
quantile (F n m _) p
  | p >= 0 && p <= 1 =
    let x = invIncompleteBeta (0.5 * n) (0.5 * m) p
    in m * x / (n * (1 - x))
  | otherwise =
    error $ "Statistics.Distribution.Uniform.quantile: p must be in [0,1] range. Got: "++show p

instance D.MaybeMean FDistribution where
  maybeMean (F _ m _) | m > 2     = Just $ m / (m - 2)
                      | otherwise = Nothing

instance D.MaybeVariance FDistribution where
  maybeStdDev (F n m _)
    | m > 4     = Just $ 2 * square m * (m + n - 2) / (n * square (m - 2) * (m - 4))
    | otherwise = Nothing

instance D.Entropy FDistribution where
  entropy (F n m _) =
    let nHalf = 0.5 * n
        mHalf = 0.5 * m in
    log (n/m)
    + logBeta nHalf mHalf
    + (1 - nHalf) * digamma nHalf
    - (1 + mHalf) * digamma mHalf
    + (nHalf + mHalf) * digamma (nHalf + mHalf)

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

instance D.ContGen FDistribution where
  genContVar = D.genContinous


{-# OPTIONS_GHC -Wall #-}

-- | A magma heirarchy for addition. The basic magma structure is repeated and prefixed with 'Additive-'.
module Plankton.Additive
  ( AdditiveMagma(..)
  , AdditiveUnital(..)
  , AdditiveAssociative
  , AdditiveCommutative
  , AdditiveInvertible(..)
  , AdditiveIdempotent
  , sum
  , Additive(..)
  , AdditiveRightCancellative(..)
  , AdditiveLeftCancellative(..)
  , AdditiveGroup(..)
  ) where

import Data.Complex (Complex(..))
import qualified Protolude as P
import Protolude (Bool(..), Double, Float, Int, Integer)

-- | 'plus' is used as the operator for the additive magma to distinguish from '+' which, by convention, implies commutativity
--
-- > ∀ a,b ∈ A: a `plus` b ∈ A
--
-- law is true by construction in Haskell
class AdditiveMagma a where
  plus :: a -> a -> a

instance AdditiveMagma Double where
  plus = (P.+)

instance AdditiveMagma Float where
  plus = (P.+)

instance AdditiveMagma Int where
  plus = (P.+)

instance AdditiveMagma Integer where
  plus = (P.+)

instance AdditiveMagma Bool where
  plus = (P.||)

instance (AdditiveMagma a) => AdditiveMagma (Complex a) where
  (rx :+ ix) `plus` (ry :+ iy) = (rx `plus` ry) :+ (ix `plus` iy)

-- | Unital magma for addition.
--
-- > zero `plus` a == a
-- > a `plus` zero == a
class AdditiveMagma a =>
      AdditiveUnital a where
  zero :: a

instance AdditiveUnital Double where
  zero = 0

instance AdditiveUnital Float where
  zero = 0

instance AdditiveUnital Int where
  zero = 0

instance AdditiveUnital Integer where
  zero = 0

instance AdditiveUnital Bool where
  zero = False

instance (AdditiveUnital a) => AdditiveUnital (Complex a) where
  zero = zero :+ zero

-- | Associative magma for addition.
--
-- > (a `plus` b) `plus` c == a `plus` (b `plus` c)
class AdditiveMagma a =>
      AdditiveAssociative a

instance AdditiveAssociative Double

instance AdditiveAssociative Float

instance AdditiveAssociative Int

instance AdditiveAssociative Integer

instance AdditiveAssociative Bool

instance (AdditiveAssociative a) => AdditiveAssociative (Complex a)

-- | Commutative magma for addition.
--
-- > a `plus` b == b `plus` a
class AdditiveMagma a =>
      AdditiveCommutative a

instance AdditiveCommutative Double

instance AdditiveCommutative Float

instance AdditiveCommutative Int

instance AdditiveCommutative Integer

instance AdditiveCommutative Bool

instance (AdditiveCommutative a) => AdditiveCommutative (Complex a)

-- | Invertible magma for addition.
--
-- > ∀ a ∈ A: negate a ∈ A
--
-- law is true by construction in Haskell
class AdditiveMagma a =>
      AdditiveInvertible a where
  negate :: a -> a

instance AdditiveInvertible Double where
  negate = P.negate

instance AdditiveInvertible Float where
  negate = P.negate

instance AdditiveInvertible Int where
  negate = P.negate

instance AdditiveInvertible Integer where
  negate = P.negate

instance AdditiveInvertible Bool where
  negate = P.not

instance (AdditiveInvertible a) => AdditiveInvertible (Complex a) where
  negate (rx :+ ix) = negate rx :+ negate ix

-- | Idempotent magma for addition.
--
-- > a `plus` a == a
class AdditiveMagma a =>
      AdditiveIdempotent a

instance AdditiveIdempotent Bool

-- | sum definition avoiding a clash with the Sum monoid in base
--
sum :: (Additive a, P.Foldable f) => f a -> a
sum = P.foldr (+) zero

-- | Additive is commutative, unital and associative under addition
--
-- > zero + a == a
-- > a + zero == a
-- > (a + b) + c == a + (b + c)
-- > a + b == b + a
class (AdditiveCommutative a, AdditiveUnital a, AdditiveAssociative a) =>
      Additive a where
  infixl 6 +
  (+) :: a -> a -> a
  a + b = plus a b

instance Additive Double

instance Additive Float

instance Additive Int

instance Additive Integer

instance Additive Bool

instance (Additive a) => Additive (Complex a)

-- | Non-commutative left minus
--
-- > negate a `plus` a = zero
class (AdditiveUnital a, AdditiveAssociative a, AdditiveInvertible a) =>
      AdditiveLeftCancellative a where
  infixl 6 ~-
  (~-) :: a -> a -> a
  (~-) a b = negate b `plus` a

-- | Non-commutative right minus
--
-- > a `plus` negate a = zero
class (AdditiveUnital a, AdditiveAssociative a, AdditiveInvertible a) =>
      AdditiveRightCancellative a where
  infixl 6 -~
  (-~) :: a -> a -> a
  (-~) a b = a `plus` negate b

-- | Minus ('-') is reserved for where both the left and right cancellative laws hold.  This then implies that the AdditiveGroup is also Abelian.
--
-- Syntactic unary negation - substituting "negate a" for "-a" in code - is hard-coded in the language to assume a Num instance.  So, for example, using ''-a = zero - a' for the second rule below doesn't work.
--
-- > a - a = zero
-- > negate a = zero - a
-- > negate a + a = zero
-- > a + negate a = zero
class (Additive a, AdditiveInvertible a) =>
      AdditiveGroup a where
  infixl 6 -
  (-) :: a -> a -> a
  (-) a b = a `plus` negate b

instance AdditiveGroup Double

instance AdditiveGroup Float

instance AdditiveGroup Int

instance AdditiveGroup Integer

instance (AdditiveGroup a) => AdditiveGroup (Complex a)


{-# LANGUAGE BangPatterns        #-}
{-# LANGUAGE DataKinds           #-}
{-# LANGUAGE DeriveGeneric       #-}
{-# LANGUAGE GADTs               #-}
{-# LANGUAGE KindSignatures      #-}
{-# LANGUAGE LambdaCase          #-}
{-# LANGUAGE PolyKinds           #-}
{-# LANGUAGE RankNTypes          #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE StandaloneDeriving  #-}
{-# LANGUAGE TypeApplications    #-}
{-# LANGUAGE TypeOperators       #-}

module Lib where

import           Control.Monad.Random
import           Data.Binary                  as B
import           Data.Singletons
import           Data.Singletons.Prelude
import           Data.Singletons.TypeLits
import           GHC.Generics                 (Generic)
import           GHC.Natural
import           Numeric.LinearAlgebra.Static
import           Prelude

{-# ANN module ("HLint: ignore Eta reduce" :: Prelude.String) #-}
{-# ANN module ("HLint: ignore Redundant lambda" :: Prelude.String) #-}

------------------------------------------------------------------------------
--import           Test.Hspec
------------------------------------------------------------------------------
{-
https://blog.jle.im/entry/practical-dependent-types-in-haskell-2.html

Practical Dependent Types in Haskell 2: Existential Neural Networks and Types at Runtime
by Justin Le ♦ Thursday June 30, 2016

using dependent types to write type-safe neural networks!

Part 1 dealt with types that are fixed at compile-time.

Part 2 types that depend runtime info
-}
randomWeights
  :: (MonadRandom m, KnownNat i, KnownNat o)
  => m (Weights i o)
randomWeights = do
  s1 :: Int <- getRandom
  s2 :: Int <- getRandom
  let wB = randomVector  s1 Uniform * 2 - 1
      wN = uniformSample s2 (-1) 1
  return $ W wB wN

runLayer
  :: (KnownNat i, KnownNat o)
  => Weights i o
  -> R i
  -> R o
runLayer (W wB wN) v = wB + wN #> v

-- * Network
--

data Network :: Nat -> [Nat] -> Nat -> * where
  O     :: !(Weights i o)
        -> Network i '[] o
  (:&~) :: KnownNat h
        => !(Weights i h)
        -> !(Network h hs o)
        -> Network i (h ': hs) o
infixr 5 :&~

deriving instance (KnownNat i, KnownNat o) => Show (Network i hs o)

logistic :: Floating a => a -> a
logistic x = 1 / (1 + exp (-x))

runNet
  :: (KnownNat i, KnownNat o)
  => Network i hs o
  -> R i
  -> R o
runNet (O w)      !v = logistic (runLayer w v)
runNet (w :&~ n') !v = let v' = logistic (runLayer w v)
                       in  runNet n' v'
{-
------------------------------------------------------------------------------
Types at Runtime

type for neural networks:

:k Network
Network :: Nat -> [Nat] -> Nat -> *
           ^      ^        ^
           |      |        size of the output vector
           |      list of hidden layer sizes
           size of input vector

above requires knowing entire structure of the network at compile-time

------------------------------------------------------------------------------
An Existential Crisis

two main ways to solve this issue : actually just two styles of doing the same thing.

------------------------------------------------------------------------------
Types hiding behind constructors

a user only needs to know input and output vectors, so Neural Network is R i -> R o

downside of having the structure in the type
- can’t store them in the same list or data structure
  Network 10 '[5,3] 1  and Network 10 '[5,2] 1

If written without the internal structure in the type then could:

data OpaqueNet i o

OpaqueNet a useful type for users
- only exposes types relevant to usage/API
- can store in list or MVar : [OpaqueNet 10 3] and MVar (OpaqueNet 10 3)
- serialize/deserialize

OpaqueNet as an "existential" wrapper over Network ('hs' is existential)
-}
data OpaqueNet :: Nat -> Nat -> * where
  ONet :: Network i hs o -> OpaqueNet i o
{-
Given        Network 6 '[10,6,3] 2
can create OpaqueNet 6           2
When using ONet constructor, the structure of the hidden layers disappears from the type.

Use the Network inside by pattern matching on ONet:
-}
runOpaqueNet
  :: (KnownNat i, KnownNat o)
  => OpaqueNet i o
  -> R i
  -> R o
runOpaqueNet (ONet n) x = runNet n x

numHiddens :: OpaqueNet i o -> Int
numHiddens (ONet n) = go n
 where
  go :: Network i hs o -> Int
  go  = \case
    O     _  -> 0
    _ :&~ n' -> 1 + go n'
{-
ScopedTypeVariables
- enables bringing hs back into scope:

case oN of
  ONet (n :: Network i hs o) -> ...

Above pattern is sometimes called dependent pair:
- pattern matching on ONet yields the hidden existentially quantified type : hs
  and also a value whose type is based on it (Network i hs o)
- like hs “paired” with Network i hs o

Pattern match on the results to give both the type 'hs' and the data structure.

(could have implemented it as ONet :: Sing hs -> Network i hs o -> OpaqueNet i o)

key to making this work
- after pattern matching on ONet, must handle 'hs' in a completely polymorphic way
- cannot assume anything about 'hs'

e.g., not ok:

bad :: OpaqueNet i o -> Network i hs o
bad (ONet n) = n

because the type signature means the caller can decide what 'hs' can be.
but the function has a specific 'hs' : the network that ONet hides.
The caller must accommodate whatever is inside ONet.

------------------------------------------------------------------------------
The Universal and the Existential

central to using existential types
- who has the power to decide what the types will be instantiated as

Most polymorphic functions in Haskell are "universally qualified", e.g.,

    map :: (a -> b) -> [a] -> [b]

the caller of map decides what a and b are

map is defined such that it will work for any a and b that the caller chooses

But for

foo :: [Int] -> OpaqueNet i o

caller of 'foo' chooses i and o
'foo' implementation chooses 'hs'
To use what 'foo' returns, must being able to handle anything 'foo'returns.

summary:

universal   : caller chooses
              function’s implementation must accommodate any choice

existential : function’s implementation chooses
              caller must accommodate any choice

earlier showed to Network i hs o inside the OpaqueNet i o
must deal with it in a parametrically polymorphic way (i.e., handle any 'hs')

------------------------------------------------------------------------------
A familiar friend

OpaqueNet i o  is a "dependent pair"
- pairs 'hs' with 'Network i hs o'

another common term : DEPENDENT SUM

"sum types" are Either-like types that can be one thing or another.

Dependent pairs/existential types actually are similar to Either/sum types

working with existential types is no different than working with Either

Given:

foo :: String -> Either Int Bool

then must handle the result for both cases.
- the function gets to pick which

(can view OpaqueNet i o as being an infinite Either over '[], '[1], '[1,2], ...)

handle Either via pattern matching
OpaqueNet i o is the same
- do not know the type of Network i hs o until pattern matching (a "dependent pattern match")

------------------------------------------------------------------------------
REIFICATION

re-write randomNet with explicit singleton input style
-}
randomNet'
  :: forall m i hs o. (MonadRandom m, KnownNat i, KnownNat o)
  => Sing hs -> m (Network i hs o)
randomNet' = \case
  SNil            ->     O <$> randomWeights
  SNat `SCons` ss -> (:&~) <$> randomWeights <*> randomNet' ss

randomNet
  :: forall m i hs o. (MonadRandom m, KnownNat i, SingI hs, KnownNat o)
  => m (Network i hs o)
randomNet = randomNet' sing
{-
use sing :: SingI hs => Sing hs to call the Sing hs ->-style function
from the SingI hs => one

still need to get the list of integers to the type level
- to create a Network i hs o
- to put into ONet

Use 'SomeSing' (a lot like OpaqueNet)
- wrapping the Sing a inside an existential data constructor
- toSing
  - takes term-level value (here [Integer])
  - returns 'SomeSing' wrapping the type-level value (here [Nat])
- pattern match on SomeSing constructor brings 'a' into scope

SomeSing implementation:

data SomeSing :: * -> * where
    SomeSing :: Sing (a :: k) -> SomeSing k

so,

foo :: SomeSing Bool
foo  = SomeSing STrue

bar :: SomeSing Nat
bar  = SomeSing (SNat :: Sing 10)
-}

mSing :: IO ()
mSing  = do
  putStrLn "How many cats do you own?"
  c <- readLn :: IO Natural
  case toSing c of SomeSing sc -> print (fromSing sc + 2)

xIn :: Int -> SomeSing Nat
xIn  = toSing . intToNatural

xOut :: SomeSing Nat -> Int
xOut  = \case SomeSing n -> naturalToInt (fromSing n)

xio :: Int -> Int
xio  = xOut . xIn
{-
inside the case statement branch the type n :: Nat is in scope.
pattern matching on the SNat constructor gives a KnownNat n instance

toSing can convert value x of type a to a singleton representing type x with kind a

Can use this to write randomONet:
-}
randomONet
  :: (MonadRandom m, KnownNat i, KnownNat o)
  => [Natural]
  -> m (OpaqueNet i o)
randomONet hs = case toSing hs of
  SomeSing ss -> ONet <$> randomNet' ss
{-
converting term-level value to type level type known as reification

original goal is now within reach:
-}
main :: IO ()
main  = do
  putStrLn "What hidden layer structure do you want?"
  hs <- readLn
  n  <- randomONet hs
  case n of
    ONet (net :: Network 10 hs 3) ->
      print net
      -- more stuff with dynamically generated net
{-
------------------------------------------------------------------------------
The Boundary

existentially quantified types (e.g., SomeSing),
enables ability to work with types that depend on runtime results

toSing and SomeSing is a BOUNDARY between untyped world and typed world
reification : cleanly separates the two

Reification boundary similar to boundary between pure code and impure (IO, etc.) code.

working with typed/untyped worlds as the same thing
- write as as possible in typed world
- write only what absolutely must be in untyped world
  (e.g., values from runtime environment)

------------------------------------------------------------------------------
Continuation-Based Existentials

another way to work with existentials that can be more natural to use

when pattern matching on existential, must work with ANY values in constructor

oNetToFoo :: OpaqueNet i o -> Foo
oNetToFoo (ONet n) = f n

f take 'Network i hs o' and work for all 'hs'

could skip the constructor and represent an existential type
as something taking the continuation f and giving it what it needs
-}

type OpaqueNet' i o r = (forall hs. Network i hs o -> r) -> r

{-
"Tell me how you would make an r
- given a 'Network i hs o' that works for any 'hs'
- then I will make it for you”

Uses Rank-N types

Using above is similar to using constructor-style:
-}
runOpaqueNet'
  :: (KnownNat i, KnownNat o)
  => OpaqueNet' i o (R o)
  -> R i
  -> R o
runOpaqueNet' oN x = oN (`runNet` x)
--            :: ((forall hs. Network i hs o -> R o) -> R o)
--            -> R i
--            -> R o

numHiddens' :: OpaqueNet' i o Int -> Int
numHiddens' oN = oN go
 where
  go :: Network i hs o -> Int
  go  = \case
    O     _  -> 0
    _ :&~ n' -> 1 + go n'
--          :: ((forall hs. Network i hs o -> Int) -> Int)
--          -> Int
{-
Above "continuation transformation" formally known as skolemization

wrap
-}

oNet' :: Network i hs o -> OpaqueNet' i o r
oNet' n = \f -> f n

-- randomONet that returns a continuation-style existential:

withRandomONet'0
  :: (MonadRandom m, KnownNat i, KnownNat o)
  => [Natural]
  -> (forall hs. Network i hs o -> m r)
  -> m r
-- aka
-- => [Integer]
-- -> OpaqueNet' i o (m r)
withRandomONet'0 hs f = case toSing hs of
  SomeSing ss -> do
    net <- randomNet' ss
    f net
{-
-- version of `toSing` that returns a skolemized `SomeSing`
withSomeSing :: [Integer]
             -> (forall (hs :: [Nat]). Sing hs -> r)
             -> r
-}

withRandomONet'
  :: (MonadRandom m, KnownNat i, KnownNat o)
  => [Natural]
  -> (forall hs. Network i hs o -> m r)
  -> m r
--aka,
-- => [Integer]
-- -> OpaqueNet' i o (m r)
withRandomONet' hs f = withSomeSing hs $ \ss -> do
  net <- randomNet' ss
  f net

main' :: IO ()
main' = do
  putStrLn "What hidden layer structure do you want?"
  hs <- readLn
  withRandomONet' hs $ \(net :: Network 10 hs 3) ->
    print net
    -- blah blah stuff with our dynamically generated net
{-
Just like the case statement pattern match represented the lexical boundary
between the untyped and typed world when using constructor-style existentials,
the ... $ \net -> ... is also a BOUNDARY for continuation-style existentials.

------------------------------------------------------------------------------
A Tale of Two Styles

two equivalent ways representing/working with existential types
can always convert between them

which one to use or offer can make a difference in code clarity

PROS/CONS

continuation-style
- does not require a data type/constructor
  - can use newtypes, but they force users to learn those types/constructors
    for each existential type returned
  - if there are many existentials, then that means a lot of un/wrapping
- easier to use when functions return existentials
  - especially if intended to immediately use them
    - saves an extraneous pattern match
- when using several existentials at once, continuation-style is better
  because each nested existential doesn’t force another level of indentation:

foo = withSomeSing x $ \sx ->
      withSomeSing y $ \sy ->
      withSomeSing z $ \sz ->
        -- ...
vs.

foo = case toSing x of
        SomeSing sx ->
          case toSing y of
            SomeSing sy ->
              case toSing z of
                SomeSing sz ->
                  -- ...
In monadic context, use do notation and ScopedTypeVariables for a non-nested style

main = do
  ONet n1 <- randomONet [7,5,3] :: IO (OpaqueNet 10 1)
  ONet n2 <- randomONet [5,5,5] :: IO (OpaqueNet 10 1)
  ONet n3 <- randomONet [5,4,3] :: IO (OpaqueNet 10 1)
  hs <- readLn
  ONet (n4 :: Network 10 hs 1) <- randomONet hs
    -- ...

nicer than

main = withRandomONet' [7,5,3] $ \n1 ->
       withRandomONet' [5,5,5] $ \n2 ->
       withRandomONet' [5,4,3] $ \n3 -> do
         hs <- readLn
         withRandomONet' hs $ \(n4 :: Network 10 hs 1) -> do
           -- ...
trick less useful for functions like toSing where things are not returned in a monad
You could wrap it in Identity, but that’s kind of silly:

foo = runIdentity $ do
        SomeSing sx <- Identity $ toSing x
        SomeSing sy <- Identity $ toSing y
        SomeSing sz <- Identity $ toSing z
        return $ -- ...

Constructor-style
- necessary for writing typeclass instances
  - i.e., cannowrite a Show instance for (forall hs. Network i hs o -> r) -> r
    but can write one for OpaqueNet i o
  - i.e., Binary de/serialization instances

Haskell does not support using Rank-N types as arguments to type constructors,
so [OpaqueNet i o] is possible,
but not [OpaqueNet' i o r] or [(forall hs. Network i hs o -> r) -> r]

MVar (OpaqueNet i o) is OK
but not MVar ((forall hs. Network i hs o -> r) -> r)

(latter are known as impredicative types)

If the type constructor is a Monad, you can use ContT-style skolemization,
like (forall hs. Network i hs o -> [r]) -> [r] and (forall hs. Network i hs o -> IO r) -> IO r.
But doesn’t work for MVar and other useful type constructors.

When writing functions that take existentials as inputs,
the constructor-style is arguably more natural. But barely.

For example, we wrote a function to find the number of hidden layers in a network earlier:

numHiddens :: OpaqueNet i o -> Int

But the continuation-style version has a slightly messier type:

numHiddens' :: ((forall hs. Network i hs o -> Int) -> Int)
            -> Int

Even with with the type synonym, it’s still a little awkward:

numHiddens' :: OpaqueNet' i o Int -> Int

This is why you’ll encounter many more functions returning continuation-style existentials
in libraries than taking them, for the most part.

------------------------------------------------------------------------------
Serializing Networks

apply existential types and reification to serialization

------------------------------------------------------------------------------
Recap on the Binary Library

Serializing networks whose sizes are statically in
their types — is straightforward --- an advantages of having sizes in types

binary library : typeclass-based de/serializing

monads (Get, Put) for describing serialization schemes and also a
typeclass used to provide serialization instructions for different
types.

usually do not write instances
use GHC generics:
-}
data Weights i o = W
  { wBiases :: !(R o)
  , wNodes  :: !(L o i)
  } deriving (Show, Generic)

instance (KnownNat i, KnownNat o) => Binary (Weights i o)
{-
gives get and put:

get :: (KnownNat i, KnownNat o) => Get (Weights i o)
put :: (KnownNat i, KnownNat o) => Weights i o -> Put
However, for GADTs like Network, we have to things manually.

------------------------------------------------------------------------------
Serializing Network
-}
putNet :: (KnownNat i, KnownNat o) => Network i hs o -> Put
putNet  = \case
  O w     -> put w
  w :&~ n -> put w *> putNet n
{-
Usually a flag is needed  to tell what constructor is being serialized,
so the deserializer can know what constructor to deserialize.
But need needed for Network because it is known know how many :&~ layers to expect
via the type: deserializing Network 5 '[10,6,3] 2 means three (:&~)’s and one O.
-}
getNet :: forall i hs o. (KnownNat i, KnownNat o) => Sing hs -> Get (Network i hs o)
getNet  = \case
  SNil            ->     O <$> get
  SNat `SCons` ss -> (:&~) <$> get <*> getNet ss
{-
getNet similar to randomNet': pattern match on 'hs' to get constructors
and ollow what the singleton’s structure says.

a Network Binary instance : cannot have get take a Sing hs input
- change the arity/type of the function
- switch to SingI-style and have their Binary instances require a SingI hs constraint.
-}
instance (KnownNat i, SingI hs, KnownNat o) => Binary (Network i hs o) where
  put = putNet
  get = getNet sing
{-
------------------------------------------------------------------------------
Serializating OpaqueNet

to de/serialize, must know complete structure

the complete structure is not in the type
- need to encode it as a flag so deserializer will which constructors
-}
hiddenStruct :: Network i hs o -> [Natural]
hiddenStruct = \case
    O _                           -> []
    _ :&~ (n' :: Network h hs' o) -> natVal (Proxy @h) : hiddenStruct n'

{-
natVal :: KnownNat n => Proxy n -> Natural
returns the value-level Natural corresponding to type-level n :: Nat

natVal and hiddenStruct : turn type-level infor ('n', 'hs') into term-level values

inverse of reification functions (like toSing)

REFECTION: type level to the value level
-}
putONet :: (KnownNat i, KnownNat o) => OpaqueNet i o -> Put
putONet (ONet net) = do
  put (hiddenStruct net) -- put the structure (as a binary flag)
  putNet net             -- then put the network itself

-- deserialize: load the list of Integers and reify it
getONet :: (KnownNat i, KnownNat o) => Get (OpaqueNet i o)
getONet = do
  hs <- get              -- load flag
  withSomeSing hs        -- reify it
    (fmap ONet . getNet) -- create it

instance (KnownNat i, KnownNat o) => Binary (OpaqueNet i o) where
  put = putONet
  get = getONet
{-
used constructor-style existential (instead of continuation-style)
because typeclass instances for the latter are not possible

------------------------------------------------------------------------------
An Existence For All

learned
- reification : untyped world to typed world to create tyypes depend on runtime values
- reflection : de/serialization shows type safety in dynamic context

------------------------------------------------------------------------------
Exercises

1. Implement putONet' and getONet' using continuation-style existentials
https://github.com/mstksg/inCode/tree/master/code-samples/dependent-haskell/NetworkTyped2.hs#L188-L193
https://github.com/mstksg/inCode/tree/master/code-samples/dependent-haskell/NetworkTyped2.hs#L195-L203

2. Work with an existential wrapper over the entire network structure (inputs and outputs, too):

data SomeNet where
  SNet :: (KnownNat i, KnownNat o)
       => Network i hs o
       -> SomeNet

(Need KnownNat because of type erasure, to recover original input/output dimensions
back on pattern matching)

And write:

- A function to convert SomeNets to OpaqueNets.
  Return OpaqueNet with existentially quantified i and o in continuation-style.
  (Also try a data type to return it in constructor-style)
https://github.com/mstksg/inCode/tree/master/code-samples/dependent-haskell/NetworkTyped2.hs#L231-L234

- randomSNet, returning m SomeNet.
https://github.com/mstksg/inCode/tree/master/code-samples/dependent-haskell/NetworkTyped2.hs#L236-L245

While you’re at it, write it to return a random continuation-style SomeNet, too! (See the type of withRandomONet' for reference on how to write the type)

The binary instance for SomeNet.

Hint: Remember natVal :: KnownNat n => Proxy n -> Integer!

Hint: Remember that toSomeSing also works for Integers, to get Sings for Nats, too!

A bit of a stretch, because the set of all [Nat]s is non-enumerable and uncountable, but hopefully you get the picture!↩

Recall that I recommend (personally, and subjectively) a style where your external API functions and typeclass instances are implemented in SingI a => style, and your internal ones in Sing a -> style. This lets all of your internal functions fit together more nicely (Sing a -> style tends to be easier to write in, especially if you stay in it the entire time, because Sings are normal first-class values, unlike those global and magical typeclasses) while at the same time removing the burden of calling with explicit singletons from people using the functionality externally.↩

In older versions of singletons, before GHC 8 and TypeInType, we had to implement it using “kind proxies”. Don’t worry if you’re following along in 7.10; the basic usage of SomeSing is essentially identical either way.↩

Skolemization is probably one of the coolest words you’ll encounter when learning/using Haskell, and sometimes just knowing that you’re “skolemizing” something makes you feel cooler. Thank you Thoralf Skolem. If you ever see a “rigid, skolem” error in GHC, you can thank him for that too! He is also the inspiration behind my decision to name my first-born son Thoralf. (My second son’s name will be Curry)↩

It even lets you write Storable instances!↩
-}


{-# LANGUAGE FlexibleInstances #-}
{-|
Module      : Math
Description : Auxiliary math functions
Copyright   : (c) Julian Kopka Larsen, 2015
Stability   : experimental

Definition of math functions like 'scatterMatrix' and the multivariate log Gamma function 'mlgamma'
-}
{-# LANGUAGE TypeSynonymInstances #-}
module Math (
    X,
    dimension,
    Sstat,
    add_means,
    sub_means,
    add_scatter,
    sub_scatter,
    fromX,
    AbelianGroup((<+>), (<->)),
    mlgamma,
    scatterMatrix,
    scatterMatrixAlt,
    mean,
    meanv,
    lnDeterminant
) where

import Numeric.LinearAlgebra (Matrix, Vector, trans, (<>), fromRows, dim, meanCov, invlndet, outer, konst, scale)
import Numeric.LinearAlgebra.Data (size)
import Debug.Trace
import Data.List ((\\))
import Partition (group)

qtr x = trace ("value: " ++ show x) x

-- | Class for the
class AbelianGroup a where
    (<+>) :: a -> a -> a
    (<->) :: a -> a -> a

-- | Type synonym for data objects that can be anything as long as it is compatible with the likelihood
--   could be extended to a type class with sufficient statistics member functions.
type X = [Vector Double]

-- | Computes the scatter matrix defined as:
scatterMatrixAlt :: X -> Matrix Double
scatterMatrixAlt x = sum $ map (\v -> outer (v-xm) (v-xm)) x
                where xm = meanv x

scatterMatrix :: X -> Matrix Double
scatterMatrix x = sum (map (\ v -> outer v v) x) - scale n (outer xm xm)
                    where xm = meanv x
                          n = fromIntegral $ length x

-- | Mean of a dataset
mean :: [Vector Double] -> Matrix Double
mean x = fromRows [meanv x]

meanv :: [Vector Double] -> Vector Double
meanv x = scale (1/(fromIntegral $ length x))(sum x)

dimension :: [Vector Double] -> Int
dimension = dim . head

type Sstat = (Int, Matrix Double, Vector Double)

fromX :: X -> Sstat
fromX x = (length x, scatterMatrix x, meanv x)

instance AbelianGroup Sstat where
    s <+> (0,_,_) = s
    (n1, s1, m1) <+> (n2, s2, m2) = (n1+n2, add_scatter n1 s1 m1 n2 s2 m2 , add_means n1 m1 n2 m2)
    s <-> (0,_,_) = s
    (n1, s1, m1) <-> (n2, s2, m2) = (n1-n2, sub_scatter n1 s1 m1 n2 s2 m2, sub_means n1 m1 n2 m2)

instance (AbelianGroup a) => AbelianGroup [a] where
    x <+> y = map (\(a,b) -> a <+> b) $ zip x y
    x <-> y = map (\(a,b) -> a <-> b) $ zip x y

add_means :: Int -> Vector Double -> Int -> Vector Double -> Vector Double
add_means _ m1 0 _ = m1
add_means 0 _ _ m2 = m2
add_means i1 m1 i2 m2 = scale w1 m1 + scale w2 m2
        where n1 = fromIntegral i1
              n2 = fromIntegral i2
              w1 = (n1/(n1+n2))
              w2 = (n2/(n1+n2))

sub_means :: Int -> Vector Double -> Int -> Vector Double -> Vector Double
sub_means _ m1 0 _ = m1
sub_means i1 _ i2 _ | i1==i2 = meanv []
sub_means i1 m1 i2 m2 = scale w1 m1 - scale w2 m2
        where n1 = fromIntegral i1
              n2 = fromIntegral i2
              w1 = (n1/(n1-n2))
              w2 = (n2/(n1-n2))

add_scatter :: Int -> Matrix Double -> Vector Double -> Int -> Matrix Double -> Vector Double -> Matrix Double
add_scatter _ sa _ 0 _ _ = sa
add_scatter 0 _ _ _ sb _ = sb
add_scatter ia sa ma ib sb mb = sa + sb + (scale na xaa) + (scale nb xbb) - (scale n xx)
        where na = fromIntegral ia
              nb = fromIntegral ib
              xaa = outer ma ma
              xbb = outer mb mb
              mx = add_means ia ma ib mb
              xx = outer mx mx
              n = na + nb

sub_scatter :: Int -> Matrix Double -> Vector Double -> Int -> Matrix Double -> Vector Double -> Matrix Double
sub_scatter 0 _ _ n _ _ = error "Nothing to subtract from."
sub_scatter _ sa _ 0 _ _ = sa
sub_scatter ia sa ma ib sb mb | ia==ib = scatterMatrix [ma]
sub_scatter ia sa ma ib sb mb | ia-ib==1 = konst 0 (size sa)
sub_scatter ia sa ma ib sb mb = sa - sb + (scale na xaa) - (scale nb xbb) - (scale n xx)
        where na = fromIntegral ia
              nb = fromIntegral ib
              xaa = outer ma ma
              xbb = outer mb mb
              mx = sub_means ia ma ib mb
              xx = outer mx mx
              n = na - nb

lnDeterminant :: Matrix Double -> Double
lnDeterminant x = absLnDet * signDet
        where (_,(absLnDet,signDet)) = invlndet x

cofG :: [Double]
cofG = [76.18009172947146,-86.50532032941677,24.01409824083091,-1.231739572450155,0.001208650973866179,-0.000005395239384953]

serG :: Double
serG = 1.000000000190015

lgamma :: Double -> Double
lgamma xx = let tmp' = (xx+5.5) - (xx+0.5)*log(xx+5.5)
                ser' = serG + (sum $ zipWith (/) cofG [xx+1..])
            in -tmp' + log(2.5066282746310005 * ser' / xx)

mlgamma :: Double -> Double -> Double
mlgamma p a = (log pi)*(p*(p-1)/4) + sum (map (\j -> lgamma (a + ((1-j)/2))) [1..p])


module Raycast (intersect, findHits, raycast) where

import           Control.Monad (guard)
import           Data.Maybe (isJust, fromJust, listToMaybe)
import Types
import Ray (createHit)
import Numeric.Vector
import Numeric.Scalar

-- https://courses.cs.washington.edu/courses/csep557/10au/lectures/triangle_intersection.pdf
intersect :: Ray -> Object -> Tri -> Maybe RayHit
intersect ray@(Ray pos dir) obj tri = do
  let a = v_getPos $ t_getP0 tri
      b = v_getPos $ t_getP1 tri
      c = v_getPos $ t_getP2 tri
      n = normalized $ (b - a) × (c - a)
      d = n · a
      t = (d - n · pos) / (n · dir)
      q = pos + dir * (fromScalar t)

      areaQAB = ((b - a) × (q - a)) · n
      areaQBC = ((c - b) × (q - b)) · n
      areaQCA = ((a - c) × (q - c)) · n
      areaABC = ((b - a) × (c - a)) · n

  guard $ areaQAB >= 0
  guard $ areaQBC >= 0
  guard $ areaQCA >= 0
  guard $ t > 0
  -- Only counts as a hit if the ray is traveling towards the front of the triangle
  -- (when direction and normal are closer to antiparallel)
  guard $ (dir · n) < 0

  let alpha = unScalar $ areaQBC / areaABC
      beta = unScalar $ areaQCA / areaABC
      gamma = unScalar $ areaQAB / areaABC

  let hitPos = pos + dir * (fromScalar t)

  return $ createHit obj tri ray (unScalar t) (alpha, beta, gamma) hitPos

findHits :: Scene -> Ray -> [RayHit]
findHits s ray = let objs = s_getObjects s
                     intersections = concatMap (\o -> map (intersect ray o) $ o_getTris o) objs
                 -- It's safe to use `fromJust` here because the filter guarantees all values are `Just`
                 in  map fromJust $ filter isJust intersections

minimumBy :: (Ord b) => (a -> b) -> [a] -> a
minimumBy f xs = foldr1 (\m x -> if f x < f m then x else m) xs

raycast :: Scene -> Ray -> Maybe RayHit
raycast s ray@(Ray origin _) = let hs = findHits s ray
                               in  if   length hs == 0 then Nothing
                                   else Just $ minimumBy rh_getDistance hs


module Main where

import           Control.Monad.Par
import           Data.Monoid                              ( (<>) )
import           Types
import qualified Text.Printf                   as Printf
import           Graphics.Gloss
import           Graphics.Gloss.Data.Color
import           Graphics.Gloss.Data.ViewPort
import qualified Numeric.LinearAlgebra         as LinAlg

windowSize :: (Int, Int)
windowSize = (1200, 800)

fps :: Int
fps = 10

data Model = Model
  { earth :: Point
  , earthSpeed :: LinAlg.Vector Float
  , jupiter :: Point
  , jupiterSpeed :: LinAlg.Vector Float
  } deriving (Show, Eq)

earthPosString :: Model -> String
earthPosString m =
  Printf.printf "earth (%.1f, %.1f)" (fst . earth $ m) (snd . earth $ m)

jupiterPosString :: Model -> String
jupiterPosString m =
  Printf.printf "jupiter (%.1f, %.1f)" (fst . jupiter $ m) (snd . jupiter $ m)

initialModel = Model
  { earth        = (-200, 0)
  , earthSpeed   = LinAlg.fromList [0, 20]
  , jupiter      = (0, 0)
  , jupiterSpeed = LinAlg.fromList [0, 0]
  }

main :: IO ()
main = simulate (InWindow "Newtons Planets" windowSize (200, 500))
                black
                fps
                initialModel
                drawModel
                stepFunction

radiusScaleFactor :: Integer -> Float
radiusScaleFactor diameter = fromIntegral diameter / distanceScale

earthRadius :: Float
earthRadius = radiusScaleFactor . diameter $ Earth

jupiterRadius :: Float
jupiterRadius = radiusScaleFactor . diameter $ Jupiter

drawModel :: Model -> Picture
drawModel model =
  let earth'               = color blue (circle earthRadius)
      jupiter'             = color cyan (circle jupiterRadius)
      (earthX  , earthY  ) = earth model
      (jupiterX, jupiterY) = jupiter model
      earthPos             = translate earthX earthY earth'
      jupiterPos           = translate jupiterX jupiterY jupiter'
  in  earthPos <> jupiterPos <> modelInfo (earthPosString model) (-580) (-350) <>
        modelInfo (jupiterPosString model) (-580) (-300)

modelInfo :: String -> Integer -> Integer -> Picture
modelInfo text a b =
  let text'  = Text text
      scaled = scale 0.1 0.1 text'
      moved  = translate (fromIntegral a) (fromIntegral b) scaled
  in  color red moved

stepFunction :: ViewPort -> Float -> Model -> Model
stepFunction vp secs model =
  let earthPoint    = earth model
      jupiterPoint  = jupiter model
      earthSpeed'   = earthSpeed model
      jupiterSpeed' = jupiterSpeed model
      (earthPulled, jupiterPulled)  = runPar $ do
        p1 <- spawnP $ pullOn Earth earthPoint Jupiter jupiterPoint
        p2 <- spawnP $ pullOn Jupiter jupiterPoint Earth earthPoint
        p1' <- get p1
        p2' <- get p2
        return (p1', p2')

  in  Model
        { earth        = newPoint earthPoint earthSpeed'
        , earthSpeed   = alterSpeed earthSpeed' earthPulled
        , jupiter      = newPoint jupiterPoint jupiterSpeed'
        , jupiterSpeed = alterSpeed jupiterSpeed' jupiterPulled
        }