{-# LANGUAGE CPP                   #-}
{-# LANGUAGE DataKinds             #-}
{-# LANGUAGE DeriveGeneric         #-}
{-# LANGUAGE ScopedTypeVariables   #-}
{-# LANGUAGE TypeOperators         #-}
{-# LANGUAGE TypeFamilies          #-}
{-# LANGUAGE FlexibleContexts      #-}
{-# LANGUAGE OverloadedStrings     #-}

module Main where

import           Control.Applicative
import           Control.Monad
import           Control.Monad.Random

import           Data.List ( foldl' )
import           Data.Maybe ( mapMaybe )
#if ! MIN_VERSION_base(4,13,0)
import           Data.Semigroup ( (<>) )
#endif
import           Data.Csv ( FromField, FromRecord, decode, parseField, HasHeader(..) )
import           Data.ByteString.Lazy.Char8 (pack)
import qualified Data.Vector as V
import qualified Data.Vector.Storable as SV

import           Numeric.LinearAlgebra ( maxIndex )
import qualified Numeric.LinearAlgebra.Static as SA

import           Options.Applicative
import           System.FilePath ( (</>) )
import           System.Random.Shuffle (shuffleM)
import           GHC.Generics (Generic)
import           GHC.Float ( float2Double)

import           Grenade
import           Grenade.Utils.OneHot

-- Adapted from https://mmhaskell.com/machine-learning/deep-learning

-- It's logistic regression!
--
-- This network is used to show how we can embed a Network as a layer in the larger IrisNetwork
-- type.

type IrisNetwork
  = Network
      '[FullyConnected 4 10, Relu, FullyConnected 10 3]
      '[ 'D1 4, 'D1 10, 'D1 10, 'D1 3]

type IrisRow = (S ( 'D1 4), S ( 'D1 3))

randomIris :: MonadRandom m => m IrisNetwork
randomIris = randomNetwork

runIris :: Int -> FilePath -> Int -> LearningParameters -> IO ()
runIris iterations dataDir nSamples rate = do
  records <- readIrisFromFile (dataDir </> "iris.data")
  let numRecords = V.length records
  shuffledRecords <- chooseRandomRecords records numRecords
  let (trainRecords, validateRecords) = V.splitAt nSamples shuffledRecords

  let trainData = mapMaybe parseRecord (V.toList trainRecords)
  let validateData = mapMaybe parseRecord (V.toList validateRecords)

  if length trainData
       /= length trainRecords
       || length validateData
       /= length validateRecords
    then putStrLn
      "Parsing train data or validation data could not be fully parsed"
    else do
      initialNetwork <- randomIris
      foldM_ (run trainData validateData) initialNetwork [1 .. iterations]
 where

  run :: [IrisRow] -> [IrisRow] -> IrisNetwork -> Int -> IO IrisNetwork
  run trainData validateData network iterationNum = do
    sampledData <- V.toList
      <$> chooseRandomRecords (V.fromList trainData) (nSamples * 3 `div` 4)
    -- Slower drop the learning rate
    let rate' = rate { learningRate = learningRate rate * 0.99 ^ iterationNum }
    let newNetwork     = foldl' (trainRow rate') network sampledData
    let labelVectors   = fmap (testRow newNetwork) validateData
    let labelValues    = fmap getLabels labelVectors
    let total          = length labelValues
    let correctEntries = length $ filter ((==) <$> fst <*> snd) labelValues
    putStrLn $ "Iteration: " ++ show iterationNum
    putStrLn $ show correctEntries ++ " correct out of: " ++ show total
    return newNetwork

  trainRow :: LearningParameters -> IrisNetwork -> IrisRow -> IrisNetwork
  trainRow lp network (input, output) = train lp network input output

  -- Takes a test row, returns predicted output and actual output from the network.
  testRow :: IrisNetwork -> IrisRow -> (S ( 'D1 3), S ( 'D1 3))
  testRow net (rowInput, predictedOutput) =
    (predictedOutput, runNet net rowInput)

  -- Goes from probability output vector to label
  getLabels :: (S ( 'D1 3), S ( 'D1 3)) -> (Int, Int)
  getLabels (S1D predictedLabel, S1D actualOutput) =
    (maxIndex (SA.extract predictedLabel), maxIndex (SA.extract actualOutput))

data IrisOpts = IrisOpts FilePath Int Int LearningParameters

iris' :: Parser IrisOpts
iris' =
  IrisOpts
    <$> argument str (metavar "DATADIR")
        -- How many samples from the dataset should be used for training?
        -- (The rest are used for validation)
    <*> option auto (long "training_samples" <> short 't' <> value 100)
    <*> option auto (long "iterations" <> short 'i' <> value 15)
    <*> (   LearningParameters
        <$> option auto (long "train_rate" <> short 'r' <> value 0.01)
        <*> option auto (long "momentum" <> value 0.9)
        <*> option auto (long "l2" <> value 0.0005)
        )

main :: IO ()
main = do
  IrisOpts dataDir nSamples iter rate <- execParser
    (info (iris' <**> helper) idm)
  putStr "Training convolutional neural network with "
  putStr $ show nSamples
  putStrLn " samples..."

  runIris iter dataDir nSamples rate

data IrisClass = Setosa | Versicolor | Virginica
  deriving (Show, Read, Eq, Ord, Generic, Enum, Bounded)

data IrisRecord = IrisRecord {
    sepalLength :: Float,
    sepalWidth  :: Float,
    petalLength :: Float,
    petalWidth  :: Float,
    specie      :: IrisClass
} deriving (Generic, Show, Read)

instance FromRecord IrisRecord

instance FromField IrisClass where
  parseField "Iris-setosa"     = return Setosa
  parseField "Iris-versicolor" = return Versicolor
  parseField "Iris-virginica"  = return Virginica
  parseField _                 = fail "unknown iris class"

parseRecord :: IrisRecord -> Maybe IrisRow
parseRecord record = case (input, output) of
  (Just i, Just o) -> Just (i, o)
  _                -> Nothing
 where
  input =
    fromStorable
      $   SV.fromList
      $   float2Double
      <$> [ sepalLength record / 8.0
          , sepalWidth record / 8.0
          , petalLength record / 8.0
          , petalWidth record / 8.0
          ]
  output = oneHot (fromEnum $ specie record)

readIrisFromFile :: FilePath -> IO (V.Vector IrisRecord)
readIrisFromFile fp = do
  contents <- readFile fp
  let contentsAsBs = pack contents
  let results =
        decode HasHeader contentsAsBs :: Either String (V.Vector IrisRecord)
  case results of
    Left  err     -> error err
    Right records -> return records

-- A function that takes this vector of records, and selects sampleSize of them at random.
chooseRandomRecords :: V.Vector a -> Int -> IO (V.Vector a)
chooseRandomRecords records sampleSize = do
  let numRecords = V.length records
  chosenIndices <- take sampleSize <$> shuffleM [0 .. (numRecords - 1)]
  return . V.fromList $ (records V.!) <$> chosenIndices


{-# LANGUAGE DataKinds           #-}
{-# LANGUAGE DeriveGeneric       #-}
{-# LANGUAGE FlexibleContexts    #-}
{-# LANGUAGE GADTs               #-}
{-# LANGUAGE OverloadedStrings   #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TemplateHaskell     #-}
{-# LANGUAGE TypeApplications    #-}
{-# LANGUAGE TypeOperators       #-}
{-# LANGUAGE QuasiQuotes         #-}
{-# LANGUAGE RankNTypes          #-}
{-# LANGUAGE PolyKinds           #-}
{-# LANGUAGE TypeFamilies        #-}
{-# LANGUAGE ConstraintKinds     #-}
{-# LANGUAGE MultiParamTypeClasses     #-}
{-# LANGUAGE FlexibleInstances         #-}
{-# LANGUAGE UndecidableInstances      #-}
{-# LANGUAGE NoMonomorphismRestriction #-}
{-# LANGUAGE AllowAmbiguousTypes       #-}
{-# LANGUAGE UndecidableSuperClasses #-}
module Frames.Regression where

import qualified Frames.Misc                  as FU
import qualified Frames.VegaLite.Utils         as FV
import qualified Math.Regression.Regression    as MR
import qualified Math.Regression.LeastSquares  as MR
import qualified Polysemy                      as P
import qualified Knit.Effect.Logger            as Log

import qualified Lucid                         as LH
import qualified Text.Blaze.Html               as BH
import qualified Control.Foldl                 as FL
import qualified Data.List                     as List
import qualified Data.Text                     as T
import qualified Data.Vinyl                    as V
import qualified Data.Vinyl.Functor            as V
import qualified Data.Vinyl.TypeLevel          as V
import           Data.Profunctor               as PF
import qualified Frames                        as F
import qualified Frames.Melt                   as F
import qualified Data.Vector.Storable          as VS

import qualified MachineLearning               as ML
import qualified MachineLearning.Regression    as ML
import qualified Numeric.LinearAlgebra         as LA
import           Numeric.LinearAlgebra.Data     ( R )
import qualified Statistics.Types              as S

import           GHC.TypeLits                   ( Symbol )

type Unweighted = "unweighted" F.:-> Void

type family NonVoidField (f :: (Symbol, Type)) :: Bool where
  NonVoidField '(a,Void) = 'False
  NonVoidField _         = 'True

class BoolVal (w :: Bool) where
  asBool :: Bool

instance BoolVal 'True where
  asBool = True

instance BoolVal 'False where
  asBool = False

data FrameRegressionResult (y :: (Symbol, Type)) (wc :: Bool) (xs :: [(Symbol, Type)]) (w :: (Symbol,Type)) (rs :: [(Symbol, Type)]) where
  FrameUnweightedRegressionResult :: (w ~ Unweighted)
    => MR.RegressionResult R -> FrameRegressionResult y wc xs Unweighted rs
  FrameWeightedRegressionResult :: (V.KnownField w, F.ElemOf rs w)
    => (V.Snd w -> R) -> MR.RegressionResult R -> FrameRegressionResult y wc xs w rs

realRecToList
  :: ( RealFrac b
     , V.RMap as
     , V.RecordToList as
     , V.ReifyConstraint Real F.ElField as
     , V.NatToInt (V.RLength as)
     )
  => F.Record as
  -> [b]
realRecToList =
  V.recordToList
    . V.rmap (\(V.Compose (V.Dict x)) -> V.Const $ realToFrac x)
    . V.reifyConstraint @Real

instance ( MR.Predictor S.NormalErr a (LA.Vector a) (MR.RegressionResult R)
         , BoolVal wc
         , LA.Element a
         , LA.Numeric a
         , RealFloat a
         , xs F.⊆ rs
         , V.RMap xs
         , V.RecordToList xs
         , V.ReifyConstraint Real F.ElField xs
         , V.NatToInt (V.RLength xs)
         ) => MR.Predictor S.NormalErr a (F.Record rs) (FrameRegressionResult y wc xs w as) where
  predict frr r =
    let rr = regressionResult frr
        addBias l = if asBool @wc then 1 : l else l -- if we regressed with a constant term, we need to put it into the xs for prediction
        va :: LA.Vector a = VS.fromList $ addBias $ realRecToList (F.rcast @xs r)
    in MR.predict rr va

regressionResult :: FrameRegressionResult y wc xs w rs -> MR.RegressionResult R
regressionResult (FrameUnweightedRegressionResult x) = x
regressionResult (FrameWeightedRegressionResult _ x) = x

withConstant
  :: forall y wc xs w rs
   . (BoolVal wc)
  => FrameRegressionResult y wc xs w rs
  -> Bool
withConstant _ = asBool @wc

weightedRegression
  :: forall y wc xs w rs
   . (BoolVal (NonVoidField w))
  => FrameRegressionResult y wc xs w rs
  -> Bool
weightedRegression _ = asBool @(NonVoidField w)

-- make X, y from the data
prepRegression
  :: forall y as f rs
   . ( Foldable f
     , as F.⊆ rs
     , F.ElemOf rs y
     , FU.RealField y
     , V.AllConstrained (FU.RealFieldOf rs) as
     , V.RMap as
     , V.RecordToList as
     , V.ReifyConstraint Real F.ElField as
     , V.NatToInt (V.RLength as)
     )
  => f (F.Record rs)
  -> (LA.Matrix R, LA.Vector R)
prepRegression dat =
  let
    nCols  = V.natToInt @(V.RLength as)
    yListF = PF.dimap (realToFrac . F.rgetField @y) List.reverse FL.list
    toListOfDoubles :: F.Record as -> [Double] = realRecToList
    mListF = PF.dimap (toListOfDoubles . F.rcast)
                      (List.concat . List.reverse)
                      FL.list
    (yList, mList) = FL.fold ((,) <$> yListF <*> mListF) dat
  in
    (LA.matrix nCols mList, LA.vector yList)

-- make X, y, w from the data
prepWeightedRegression
  :: forall y as w f rs
   . ( Foldable f
     , as F.⊆ rs
     , F.ElemOf rs y
     , FU.RealField y
     , F.ElemOf rs w
     , FU.RealField w
     , V.AllConstrained (FU.RealFieldOf rs) as
     , V.RMap as
     , V.RecordToList as
     , V.ReifyConstraint Real F.ElField as
     , V.NatToInt (V.RLength as)
     )
  => f (F.Record rs)
  -> (LA.Matrix R, LA.Vector R, LA.Vector R)
prepWeightedRegression dat =
  let
    nCols  = V.natToInt @(V.RLength as)
    yListF = PF.dimap (realToFrac . F.rgetField @y) List.reverse FL.list
    wListF = PF.dimap (realToFrac . F.rgetField @w) List.reverse FL.list
    toListOfDoubles :: F.Record as -> [Double] = realRecToList
    mListF = PF.dimap (toListOfDoubles . F.rcast)
                      (List.concat . List.reverse)
                      FL.list
    (yList, mList, wList) = FL.fold ((,,) <$> yListF <*> mListF <*> wListF) dat
  in
    (LA.matrix nCols mList, LA.vector yList, LA.vector wList)

prettyPrintRegressionResult
  :: forall y wc as w rs
   . ( F.ColumnHeaders '[y]
     , F.ColumnHeaders '[w]
     , F.ColumnHeaders as
     , BoolVal wc
     )
  => (T.Text -> T.Text -> T.Text)
  -> FrameRegressionResult y wc as w rs
  -> S.CL Double
  -> T.Text
prettyPrintRegressionResult headerF res cl =
  let yName   = FV.colName @y
      wName   = FV.colName @w
      xNames' = T.pack <$> F.columnHeaders (Proxy :: Proxy (F.Record as))
      xNames  = if withConstant res then "intercept" : xNames' else xNames'
  in  MR.prettyPrintRegressionResult (headerF yName wName)
                                     xNames
                                     (regressionResult res)
                                     cl

prettyPrintRegressionResultLucid
  :: forall y wc as w rs
   . ( F.ColumnHeaders '[y]
     , F.ColumnHeaders '[w]
     , F.ColumnHeaders as
     , BoolVal wc
     )
  => (T.Text -> T.Text -> T.Text)
  -> FrameRegressionResult y wc as w rs
  -> S.CL Double
  -> LH.Html ()
prettyPrintRegressionResultLucid headerF res cl =
  let yName   = FV.colName @y
      wName   = FV.colName @w
      xNames' = T.pack <$> F.columnHeaders (Proxy :: Proxy (F.Record as))
      xNames  = if withConstant res then "intercept" : xNames' else xNames'
  in  MR.prettyPrintRegressionResultLucid (headerF yName wName)
                                          xNames
                                          (regressionResult res)
                                          cl

prettyPrintRegressionResultBlaze
  :: forall y wc as w rs
   . ( F.ColumnHeaders '[y]
     , F.ColumnHeaders '[w]
     , F.ColumnHeaders as
     , BoolVal wc
     )
  => (T.Text -> T.Text -> T.Text)
  -> FrameRegressionResult y wc as w rs
  -> S.CL Double
  -> BH.Html
prettyPrintRegressionResultBlaze headerF res cl =
  let yName   = FV.colName @y
      wName   = FV.colName @w
      xNames' = T.pack <$> F.columnHeaders (Proxy :: Proxy (F.Record as))
      xNames  = if withConstant res then "intercept" : xNames' else xNames'
  in  MR.prettyPrintRegressionResultBlaze (headerF yName wName)
                                          xNames
                                          (regressionResult res)
                                          cl

keyRecordText
  :: (V.ReifyConstraint Show F.ElField ks, V.RecordToList ks, V.RMap ks)
  => F.Record ks
  -> T.Text
keyRecordText keyRec =
  let keyValuesAsText rks =
        V.recordToList
          . V.rmap (\(V.Compose (V.Dict x)) -> V.Const $ T.pack $ show x)
          $ V.reifyConstraint @Show rks
  in  T.intercalate "; " $ keyValuesAsText keyRec

prettyPrintRegressionResults
  :: forall y wc as w rs k f a
   . ( F.ColumnHeaders '[y]
     , F.ColumnHeaders '[w]
     , F.ColumnHeaders as
     , BoolVal wc
     , BoolVal (NonVoidField w)
     , Foldable f
     , Monoid a
     )
  => (k -> T.Text)
  -> f (k, FrameRegressionResult y wc as w rs)
  -> S.CL R
  -> (  (T.Text -> T.Text -> T.Text)
     -> FrameRegressionResult y wc as w rs
     -> S.CL Double
     -> a
     )
  -> a
  -> a
prettyPrintRegressionResults keyText keyed cl printOne sepEach =
  let headerF res key yName wName = if weightedRegression res
                                    then
                                      "Explaining "
                                        <> yName
                                        <> " ("
                                        <> keyText key
                                        <> "; weights from "
                                        <> wName
                                        <> ")"
                                    else "Explaining " <> yName <> " (" <> keyText key <> ")"
  in  FL.fold
        (FL.Fold (\t (rk, res) -> t <> printOne (headerF res rk) res cl)
                 sepEach
                 id
        )
        keyed

-- explain y in terms of as
leastSquaresByMinimization
  :: forall y as rs f
   . ( Foldable f
     , as F.⊆ rs
     , F.ElemOf rs y
     , FU.RealField y
     , V.AllConstrained (FU.RealFieldOf rs) as
     , V.RMap as
     , V.RecordToList as
     , V.ReifyConstraint Real F.ElField as
     , V.NatToInt (V.RLength as)
     )
  => Bool
  -> [R]
  -> f (F.Record rs)
  -> [R]
leastSquaresByMinimization wc guess dat =
  let (mX, y)       = prepRegression @y @as dat
      mX1           = if wc then ML.addBiasDimension mX else mX
      (solution, _) = ML.minimize (ML.ConjugateGradientFR 0.1 0.1)
                                  ML.LeastSquares
                                  0.001
                                  20
                                  ML.RegNone
                                  mX1
                                  y
                                  (LA.fromList guess)
  in  LA.toList solution

ordinaryLeastSquares
  :: forall effs y wc as rs f
   . ( Log.LogWithPrefixesLE effs
     , Foldable f
     , as F.⊆ rs
     , F.ElemOf rs y
     , FU.RealField y
     , V.AllConstrained (FU.RealFieldOf rs) as
     , V.RMap as
     , V.RecordToList as
     , V.ReifyConstraint Real F.ElField as
     , BoolVal wc
     , V.NatToInt (V.RLength as)
     )
  => f (F.Record rs)
  -> P.Sem effs (FrameRegressionResult y wc as Unweighted rs)
ordinaryLeastSquares dat = do
  let (mA, vB)  = prepRegression @y @as dat
      withConst = asBool @wc
  FrameUnweightedRegressionResult <$> MR.ordinaryLS withConst mA vB

weightedLeastSquares
  :: forall y wc as w rs f effs
   . ( Log.LogWithPrefixesLE effs
     , Foldable f
     , as F.⊆ rs
     , F.ElemOf rs y
     , FU.RealField y
     , F.ElemOf rs w
     , FU.RealField w
     , BoolVal wc
     , V.AllConstrained (FU.RealFieldOf rs) as
     , V.RMap as
     , V.RecordToList as
     , V.ReifyConstraint Real F.ElField as
     , V.NatToInt (V.RLength as)
     )
  => f (F.Record rs)
  -> P.Sem effs (FrameRegressionResult y wc as w rs)
weightedLeastSquares dat = do
  let (mA, vB, vW) = prepWeightedRegression @y @as @w dat
      withConst    = asBool @wc
  FrameWeightedRegressionResult realToFrac <$> MR.weightedLS withConst mA vB vW

-- special case when weights come from observations being population averages of different populations
popWeightedLeastSquares
  :: forall effs y wc as w rs f
   . ( Log.LogWithPrefixesLE effs
     , Foldable f
     , as F.⊆ rs
     , F.ElemOf rs y
     , FU.RealField y
     , F.ElemOf rs w
     , FU.RealField w
     , BoolVal wc
     , V.AllConstrained (FU.RealFieldOf rs) as
     , V.RMap as
     , V.RecordToList as
     , V.ReifyConstraint Real F.ElField as
     , V.NatToInt (V.RLength as)
     )
  => f (F.Record rs)
  -> P.Sem effs (FrameRegressionResult y wc as w rs)
popWeightedLeastSquares dat = do
  let (mA, vB, vW) = prepWeightedRegression @y @as @w dat
      withConst    = asBool @wc
      vWpop        = LA.cmap sqrt vW -- this is the correct weight for population average, the sqrt of the number averaged in that sample
  FrameWeightedRegressionResult (sqrt . realToFrac)
    <$> MR.weightedLS withConst mA vB vWpop

-- special case when we know residuals are heteroscedastic with variances proportional to given numbers
varWeightedLeastSquares
  :: forall effs y wc as w rs f
   . ( Log.LogWithPrefixesLE effs
     , Foldable f
     , as F.⊆ rs
     , F.ElemOf rs y
     , FU.RealField y
     , F.ElemOf rs w
     , FU.RealField w
     , BoolVal wc
     , V.AllConstrained (FU.RealFieldOf rs) as
     , V.RMap as
     , V.RecordToList as
     , V.ReifyConstraint Real F.ElField as
     , V.NatToInt (V.RLength as)
     )
  => f (F.Record rs)
  -> P.Sem effs (FrameRegressionResult y wc as w rs)
varWeightedLeastSquares dat = do
  let (mA, vB, vW) = prepWeightedRegression @y @as @w dat
      withConst    = asBool @wc
      vWvar        = LA.cmap (\x -> 1 / sqrt x) vW -- this is the correct weight for given variance
  FrameWeightedRegressionResult (\x -> 1 / sqrt (realToFrac x))
    <$> MR.weightedLS withConst mA vB vWvar

totalLeastSquares
  :: forall effs y wc as rs f
   . ( Log.LogWithPrefixesLE effs
     , Foldable f
     , as F.⊆ rs
     , F.ElemOf rs y
     , FU.RealField y
     , BoolVal wc
     , V.AllConstrained (FU.RealFieldOf rs) as
     , V.RMap as
     , V.RecordToList as
     , V.ReifyConstraint Real F.ElField as
     , V.NatToInt (V.RLength as)
     )
  => f (F.Record rs)
  -> P.Sem effs (FrameRegressionResult y wc as Unweighted rs)
totalLeastSquares dat = do
  let (mA, vB)  = prepRegression @y @as dat
      withConst = asBool @wc
  FrameUnweightedRegressionResult <$> MR.totalLS withConst mA vB

weightedTLS
  :: forall effs y wc as w rs f
   . ( Log.LogWithPrefixesLE effs
     , Foldable f
     , as F.⊆ rs
     , F.ElemOf rs y
     , FU.RealField y
     , BoolVal wc
     , F.ElemOf rs w
     , FU.RealField w
     , V.AllConstrained (FU.RealFieldOf rs) as
     , V.RMap as
     , V.RecordToList as
     , V.ReifyConstraint Real F.ElField as
     , V.NatToInt (V.RLength as)
     )
  => f (F.Record rs)
  -> P.Sem effs (FrameRegressionResult y wc as w rs)
weightedTLS dat = do
  let (mA, vB, vW) = prepWeightedRegression @y @as @w dat
      withConst    = asBool @wc
  FrameWeightedRegressionResult realToFrac <$> MR.weightedTLS withConst mA vB vW

popWeightedTLS
  :: forall effs y wc as w rs f
   . ( Log.LogWithPrefixesLE effs
     , Foldable f
     , as F.⊆ rs
     , F.ElemOf rs y
     , FU.RealField y
     , BoolVal wc
     , F.ElemOf rs w
     , FU.RealField w
     , V.AllConstrained (FU.RealFieldOf rs) as
     , V.RMap as
     , V.RecordToList as
     , V.ReifyConstraint Real F.ElField as
     , V.NatToInt (V.RLength as)
     )
  => f (F.Record rs)
  -> P.Sem effs (FrameRegressionResult y wc as w rs)
popWeightedTLS dat = do
  let (mA, vB, vW) = prepWeightedRegression @y @as @w dat
      withConst    = asBool @wc
      vWpop        = LA.cmap sqrt vW -- this is the correct weight for population average, the sqrt of the number averaged in that sample
  FrameWeightedRegressionResult (sqrt . realToFrac)
    <$> MR.weightedTLS withConst mA vB vWpop

varWeightedTLS
  :: forall effs y wc as w rs f
   . ( Log.LogWithPrefixesLE effs
     , Foldable f
     , as F.⊆ rs
     , F.ElemOf rs y
     , FU.RealField y
     , BoolVal wc
     , F.ElemOf rs w
     , FU.RealField w
     , V.AllConstrained (FU.RealFieldOf rs) as
     , V.RMap as
     , V.RecordToList as
     , V.ReifyConstraint Real F.ElField as
     , V.NatToInt (V.RLength as)
     )
  => f (F.Record rs)
  -> P.Sem effs (FrameRegressionResult y wc as w rs)
varWeightedTLS dat = do
  let (mA, vB, vW) = prepWeightedRegression @y @as @w dat
      withConst    = asBool @wc
      vWvar        = LA.cmap (\x -> 1 / sqrt x) vW -- this is the correct weight for given variance
  FrameWeightedRegressionResult (\x -> 1 / sqrt (realToFrac x))
    <$> MR.weightedTLS withConst mA vB vWvar

{-
-- this is sort of ugly but I think it will work.  And it has a pretty narrow purpose
-- But what will we do when we want logistic regression or whatever?
class (V.AllConstrained RealFrac (V.Unlabeled rs), Real a) => RealFracRecordFromList a rs where
  recordFromList :: [a] -> Maybe (F.Record rs)

instance Real a => RealFracRecordFromList a '[] where
  recordFromList _ = Just $ V.RNil

instance ( V.AllConstrained RealFrac (V.Unlabeled (r : rs))
         , V.AllConstrained RealFrac (V.Unlabeled rs)
         , RealFracRecordFromList a rs
         , V.KnownField r
         , RealFrac (V.Snd r)
         , Real a) => RealFracRecordFromList a (r : rs) where
  recordFromList [] = Nothing
  recordFromList (a : as) = case recordFromList as of
    Nothing -> Nothing
    Just xs -> let x = (realToFrac a) :: V.Snd r in Just $ x &: xs
-}


{-# LANGUAGE BangPatterns         #-}
{-# LANGUAGE DataKinds            #-}
{-# LANGUAGE FlexibleContexts     #-}
{-# LANGUAGE GADTs                #-}
{-# LANGUAGE KindSignatures       #-}
{-# LANGUAGE LambdaCase           #-}
{-# LANGUAGE PolyKinds            #-}
{-# LANGUAGE RankNTypes           #-}
{-# LANGUAGE ScopedTypeVariables  #-}
{-# LANGUAGE TypeApplications     #-}
{-# LANGUAGE TypeInType           #-}
{-# LANGUAGE TypeOperators        #-}
{-# LANGUAGE UndecidableInstances #-}

module TensorOps.Learn.NeuralNet.FeedForward
  ( Network(..)
  , buildNet
  , runNetwork
  , trainNetwork
  , induceNetwork
  , nmap
  , (~*)
  , (*~)
  , liftNet
  , netParams
  , networkGradient
  , genNet
  , ffLayer
  ) where

import           Control.Category
import           Control.DeepSeq
import           Control.Monad.Primitive
import           Data.Kind
import           Data.Singletons
import           Data.Singletons.Prelude        (Sing(..))
import           Data.Type.Conjunction
import           Data.Type.Length
import           Data.Type.Product              as TCP
import           Data.Type.Product.Util         as TCP
import           Data.Type.Sing
import           Prelude hiding                 ((.), id)
import           Statistics.Distribution.Normal
import           System.Random.MWC
import           TensorOps.Learn.NeuralNet
import           TensorOps.NatKind
import           TensorOps.TOp                  as TO
import           TensorOps.Tensor               as TT
import           TensorOps.Types
import           Type.Class.Higher
import           Type.Class.Higher.Util
import           Type.Class.Known
import           Type.Class.Witness
import           Type.Family.List
import           Type.Family.List.Util

data Network :: ([k] -> Type) -> k -> k -> Type where
    N :: { _nsPs    :: !(Sing ps)
         , _nOp     :: !(TOp ('[i] ': ps) '[ '[o] ])
         , _nParams :: !(Prod t ps)
         } -> Network t i o

instance NFData1 t => NFData (Network t i o) where
    rnf = \case
      N _ o p -> o `seq` p `deepseq1` ()
    {-# INLINE rnf #-}

buildNet
    :: SingI ps
    => TOp ('[i] ': ps) '[ '[o] ]
    -> Prod t ps
    -> Network t i o
buildNet = N sing

netParams
    :: Network t i o
    -> (forall ps. SingI ps => Prod t ps -> r)
    -> r
netParams n f = case n of
    N o _ p -> f p \\ o

(~*~)
    :: Network t a b
    -> Network t b c
    -> Network t a c
N sPs1 o1 p1 ~*~ N sPs2 o2 p2 =
    N (sPs1 %:++ sPs2) (o1 *>> o2) (p1 `TCP.append'` p2)
        \\ singLength sPs1
infixr 4 ~*~
{-# INLINE (~*~) #-}

instance Category (Network t) where
    id = N SNil idOp Ø
    (.) = flip (~*~)

(~*) :: TOp '[ '[a] ] '[ '[b] ]
     -> Network t b c
     -> Network t a c
f ~* N sO o p = N sO (f *>> o) p
infixr 4 ~*
{-# INLINE (~*) #-}

(*~) :: Network t a b
     -> TOp '[ '[b] ] '[ '[c] ]
     -> Network t a c
N sO o p *~ f = N sO (o >>> f) p
infixl 5 *~
{-# INLINE (*~) #-}

liftNet
     :: TOp '[ '[i] ] '[ '[o] ]
     -> Network t i o
liftNet o = buildNet o Ø

nmap
     :: SingI o
     => (forall a. RealFloat a => a -> a)
     -> Network t i o
     -> Network t i o
nmap f n = n *~ TO.map f
{-# INLINE nmap #-}

runNetwork
    :: (RealFloat (ElemT t), Tensor t)
    => Network t i o
    -> t '[i]
    -> t '[o]
runNetwork (N _ o p) = head' . runTOp o . (:< p)
{-# INLINE runNetwork #-}

trainNetwork
    :: forall i o t. (Tensor t, RealFloat (ElemT t))
    => TOp '[ '[o], '[o] ] '[ '[] ]
    -> ElemT t
    -> t '[i]
    -> t '[o]
    -> Network t i o
    -> Network t i o
trainNetwork loss r x y = \case
    N s o p ->
      let p' = map1 (\(!(s1 :&: o1 :&: g1)) -> TT.zip stepFunc o1 g1 \\ s1)
             $ zipProd3 (singProd s) p (tail' $ netGrad loss x y s o p)
      in  N s o p'
  where
    stepFunc :: ElemT t -> ElemT t -> ElemT t
    stepFunc !o' !g' = o' - r * g'
    {-# INLINE stepFunc #-}
{-# INLINE trainNetwork #-}

induceNetwork
    :: forall i o t. (Tensor t, RealFloat (ElemT t), SingI i)
    => TOp '[ '[o], '[o] ] '[ '[] ]
    -> ElemT t
    -> t '[o]
    -> Network t i o
    -> t '[i]
    -> t '[i]
induceNetwork loss r y = \case
    N s o p -> \x -> TT.zip stepFunc x (head' $ netGrad loss x y s o p)
  where
    stepFunc :: ElemT t -> ElemT t -> ElemT t
    stepFunc o' g' = o' - r * g'
    {-# INLINE stepFunc #-}
{-# INLINE induceNetwork #-}

networkGradient
    :: forall i o t r. (Tensor t, RealFloat (ElemT t))
    => TOp '[ '[o], '[o] ] '[ '[] ]
    -> t '[i]
    -> t '[o]
    -> Network t i o
    -> (forall ps. SingI ps => Prod t ps -> r)
    -> r
networkGradient loss x y = \case
    N s o p -> \f -> f (tail' $ netGrad loss x y s o p) \\ s
{-# INLINE networkGradient #-}

netGrad
    :: forall i o ps t. (Tensor t, RealFloat (ElemT t))
    => TOp '[ '[o], '[o] ] '[ '[] ]
    -> t '[i]
    -> t '[o]
    -> Sing ps
    -> TOp ('[i] ': ps) '[ '[o] ]
    -> Prod t ps
    -> Prod t ('[i] ': ps)
netGrad loss x y s o p = (\\ appendSnoc lO (Proxy @'[o])) $
                         (\\ lO                         ) $
                         takeProd @'[ '[o] ] (LS lO)
                       $ gradTOp o' inp
  where
    lO  :: Length ps
    lO = singLength s
    o'  :: ((ps ++ '[ '[o] ]) ~ (ps >: '[o]), Known Length ps)
        => TOp ('[i] ': ps >: '[o]) '[ '[]]
    o' = o *>> loss
    inp  :: Prod t ('[i] ': ps >: '[o])
    inp = x :< p >: y
{-# INLINE netGrad #-}

ffLayer
    :: forall i o m t. (SingI i, SingI o, PrimMonad m, Tensor t)
    => Gen (PrimState m)
    -> m (Network t i o)
ffLayer g = (\w b -> buildNet ffLayer' (w :< b :< Ø))
          <$> genRand (normalDistr 0 0.5) g
          <*> genRand (normalDistr 0 0.5) g
  where
    ffLayer'
        :: TOp '[ '[i], '[o,i], '[o]] '[ '[o] ]
    ffLayer' = firstOp (TO.swap >>> TO.matVec)
           >>> TO.add
    {-# INLINE ffLayer' #-}
{-# INLINE ffLayer #-}

genNet
    :: forall k o i m (t :: [k] -> Type). (SingI o, SingI i, PrimMonad m, Tensor t)
    => [(Integer, Activation k)]
    -> Activation k
    -> Gen (PrimState m)
    -> m (Network t i o)
genNet xs0 f g = go sing xs0
  where
    go  :: forall (j :: k). ()
        => Sing j
        -> [(Integer, Activation k)]
        -> m (Network t j o)
    go sj = (\\ sj) $ \case
      []        -> (*~ getAct f) <$> ffLayer g
      (x,f'):xs -> withNatKind x $ \sl -> (\\ sl) $ do
        n <- go sl xs
        l <- ffLayer g
        return $ l *~ getAct f' ~*~ n
    {-# INLINE go #-}
{-# INLINE genNet #-}


{-# LANGUAGE BangPatterns        #-}
{-# LANGUAGE FlexibleContexts    #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeOperators       #-}
{-# LANGUAGE ViewPatterns        #-}
module FokkerPlanck.GPUKernel where

import           Data.Array.Accelerate              as A
import           Data.Array.Accelerate.Data.Complex as A
import qualified Data.Complex                       as C
import Debug.Trace

{-# INLINE moveParticle #-}
moveParticle ::
     (A.RealFloat a, A.Num a, Elt a-- , FromIntegral Int a
     )
  => Exp (a, a, a, a)
  -> Exp (a, a, a, a)
moveParticle (unlift -> (phi, rho, theta, r)) =
  let !x = theta - phi
      !cosX = A.cos x
      !newPhi = phi + (A.atan2 (r * A.sin x) (rho + r * cosX))
      !newRho = A.sqrt $ (rho * rho + r * r + 2 * r * rho * cosX)
      -- !xx = A.round $ newRho * A.cos newPhi :: Exp Int
      -- !y = A.round $ newRho * A.sin newPhi :: Exp Int
      -- !newRho' = A.sqrt . A.fromIntegral $ xx * xx + y * y
      -- !newPhi' = A.atan2 (A.fromIntegral y) (A.fromIntegral xx)
  in lift (newPhi, newRho, theta, r)

{-# INLINE normalizationTerm #-}
normalizationTerm ::
     (A.Floating a, A.Num a, A.RealFloat a, A.Elt (Complex a))
  => Exp a
  -> Exp a
  -> Exp (Complex a)
normalizationTerm halfLogPeriod freq =
  (lift $ 1 A.:+ ((-A.pi) * freq / halfLogPeriod)) /
  (A.exp (lift $ halfLogPeriod A.:+ (-A.pi * freq)) -
   A.exp (lift $ (-halfLogPeriod) A.:+ (A.pi * freq)))

-- {-# INLINE coefficient #-}
-- coefficient ::
--      forall a. (A.Floating a, A.Num a, A.RealFloat a, A.Elt (Complex a), Prelude.Fractional a)
--   => Exp a
--   -> Exp a
--   -> Exp a
--   -> Exp a
--   -> Exp a
--   -> Exp (a, a, a, a)
--   -> Exp (A.Complex a)
-- coefficient halfLogPeriod rFreq thetaFreq rhoFreq phiFreq particle -- (unlift -> (phi, rho, theta, r))
--  =
--   let (phi, rho, theta, r) = unlift particle :: (Exp a, Exp a, Exp a, Exp a)
--      -- (normalizationTerm halfLogPeriod rhoFreq) *
--      -- (normalizationTerm halfLogPeriod rFreq) *
--      -- (lift (A.cos (phiFreq * phi + thetaFreq * (theta - phi)) A.:+ 0)) *
--                            -- (A.cis $
--                            --  (-2 * A.pi) * (rhoFreq * rho + rFreq * (r - rho)) /
--                            --  (A.exp halfLogPeriod))

--   in (lift $
--       ((A.exp $ (-0.5) * (rho + r)) *
--        (A.cos (phiFreq * phi  + thetaFreq * (theta - phi)))
--       ) A.:+
--       0) *
--      (A.cis $ (-1) * ((rhoFreq * rho + rFreq * (r - rho)
--                       ) -- * (2 * A.pi) / (A.exp halfLogPeriod)
--                       -- + phiFreq * phi  + thetaFreq * (theta - phi)
--                      ))
--      -- (lift $ ((A.cos (phiFreq * phi + thetaFreq * (theta - phi) / 2))) A.:+ 0) *
--      -- (A.cis $ (-2 * A.pi) * (rhoFreq * rho + rFreq * (r - rho)) / (A.exp halfLogPeriod))
--      -- ((lift $ rho A.:+ 0) A.** (lift $ (-0.5) A.:+ (rFreq - rhoFreq))) *
--      -- ((lift $ r A.:+ 0) A.** (lift $ (-0.5) A.:+ (-rFreq)))
--      -- (A.cis $ (-A.pi) * (rhoFreq * rho + rFreq * (r - rho)
--      --                        ) / (halfLogPeriod))
--      -- (A.cis $
--      --  (-A.pi) * (rhoFreq * (A.log rho) + rFreq * (A.log (r / rho))) /
--      --  halfLogPeriod)
--   -- in ((lift $ rho A.:+ 0) A.**
--   --     (lift $ 0 A.:+ (-A.pi) * (rhoFreq - rFreq) / halfLogPeriod)) *
--   --    ((lift $ r A.:+ 0) A.** (lift $ 0 A.:+ (-A.pi) * rFreq / halfLogPeriod)) *
--   --    (A.cis $ (-phiFreq) * phi - thetaFreq * (theta - phi))

coefficient ::
     forall a. (A.Floating a, A.Num a, A.RealFloat a, A.Elt (Complex a), Prelude.Fractional a)
  => Exp a
  -> Exp a
  -> Exp a
  -> Exp a
  -> Exp a
  -> Exp (a, a, a, a)
  -> Exp (A.Complex a)
coefficient halfLogPeriod rFreq thetaFreq rhoFreq phiFreq particle =
  let (phi, rho, theta, r) = unlift particle :: (Exp a, Exp a, Exp a, Exp a)
  in (lift $
      ((A.exp $ (-0.5) * (rho + r)) *
       (A.cos (phiFreq * phi + thetaFreq * (theta - phi)))) A.:+
      0) *
     (A.cis $ (-1) * (rhoFreq * rho + rFreq * (r + rho)))

-- {-# INLINE gpuKernel #-}
gpuKernel ::
     forall a. (A.Eq a, A.Floating a, A.Num a, A.RealFloat a, A.Elt (Complex a), A.FromIntegral Int a, Prelude.Fractional a)
  => Exp a
  -> Exp a
  -> Exp (A.Complex a)
  -> Acc (A.Vector (a, a, a, a))
  -> Acc (A.Vector (a, a, a, a))
  -> Acc (A.Vector (A.Complex a))
gpuKernel !maxScaleExp !halfLogPeriodExp !deltaLogRhoComplexExp freqArr particles =
  let -- delta = constant 0.01
      movedParticleArr =
        compute .
        A.map
          (\particle ->
             let (phi, rho, theta, r) =
                   unlift particle :: (Exp a, Exp a, Exp a, Exp a)
                 logRho = A.log rho
                   -- (A.fromIntegral $
                   --  (A.round (((A.log rho) + halfLogPeriodExp) / delta) :: Exp Int)) *
                   -- delta -
                   -- halfLogPeriodExp
                 -- newRho =
                 --   (A.fromIntegral $ (A.round (rho / delta) :: Exp Int)) * delta :: Exp a
             in lift (phi, logRho, theta, (A.log r) :: Exp a)) -- .
        -- afst .
        -- A.filter
        --   (\particle ->
        --      let (_, rho, _, _) =
        --            unlift particle :: (Exp a, Exp a, Exp a, Exp a)
        --      in rho A.> (A.constant 0) ) -- .
        -- A.map moveParticle
        $
        particles
  in A.map
       (\(unlift -> (rFreq, thetaFreq, rhoFreq, phiFreq))
          -- (lift $ delta A.:+ 0) /
          -- (lift $ (8 * A.pi * A.pi * halfLogPeriodExp) A.:+ 0) *
         ->
          (sfoldl
             (\s particle ->
                s +
                (coefficient
                   halfLogPeriodExp
                   rFreq
                   thetaFreq
                   rhoFreq
                   phiFreq
                   particle))
             0
             (constant Z)
             movedParticleArr) * deltaLogRhoComplexExp
          -- (lift $
          --  (16 * A.pi * A.pi * halfLogPeriodExp * halfLogPeriodExp) A.:+ 0)
             -- movedParticleArr
        ) $
     freqArr

{-# INLINE convolveKernel #-}
convolveKernel ::
     (A.Num a, A.RealFloat a, A.Elt (Complex a))
  => Acc (Array DIM4 (Complex a))
  -> Acc (Array DIM4 (Complex a))
  -> Acc (Scalar Int)
  -> Acc (Scalar Int)
  -> Acc (Array DIM4 (Complex a))
  -> Acc (Vector (Complex a))
convolveKernel coefficients harmonics thetaIdx rIdx input =
  let (Z :. numRhoFreq :. numPhiFreq :. cols :. rows) =
        unlift . shape $ input :: (Z :. Exp Int :. Exp Int :. Exp Int :. Exp Int)
      coefficientsArr =
        A.replicate (lift (Z :. All :. All :. cols :. rows)) .
        slice coefficients $
        (lift (Z :. (the rIdx) :. (the thetaIdx) :. All :. All))
      harmonicsArr =
        backpermute
          (shape input)
          (\(unlift -> Z :. rho :. phi :. col :. row :: (Z :. Exp Int :. Exp Int :. Exp Int :. Exp Int)) ->
             lift (Z :. (rho + the rIdx) :. (phi + the thetaIdx) :. col :. row))
          harmonics
  in flatten .
     A.sum .
     A.sum .
     backpermute
       (lift (Z :. cols :. rows :. numRhoFreq :. numPhiFreq))
       (\(unlift -> Z :. col :. row :. rho :. phi :: (Z :. Exp Int :. Exp Int :. Exp Int :. Exp Int)) ->
          lift (Z :. rho :. phi :. col :. row)) .
     A.zipWith (*) harmonicsArr . A.zipWith (*) coefficientsArr $
     input

{-# INLINE coefficient' #-}
coefficient' ::
     forall a.
     ( A.Floating a
     , A.Num a
     , A.RealFloat a
     , A.Elt (Complex a)
     , Prelude.Fractional a
     )
  => Exp a
  -> Exp a
  -> Exp a
  -> Exp a
  -> Exp a
  -> Exp (a, a, a, a, a)
  -> Exp (A.Complex a)
coefficient' sigma rFreq thetaFreq rhoFreq phiFreq particle =
  let (phi, rho, theta, r, v) =
        unlift particle :: (Exp a, Exp a, Exp a, Exp a, Exp a)
  in (lift $ (v * (A.exp $ (sigma - 1) * (rho + r))) A.:+ 0) *
     (A.cis $
      (-1) *
      (rhoFreq * rho + rFreq * (r - rho) + phiFreq * phi +
       thetaFreq * (theta - phi)))

gpuKernel' ::
     forall a.
     ( A.Eq a
     , A.Floating a
     , A.Num a
     , A.RealFloat a
     , A.Elt (Complex a)
     , A.FromIntegral Int a
     , Prelude.Fractional a
     )
  => Exp a
  -> Acc (A.Vector (a, a, a, a))
  -> Acc (A.Vector (a, a, a, a, a))
  -> Acc (A.Vector (A.Complex a))
gpuKernel' sigma freqArr xs =
  A.map
    (\(unlift -> (rFreq, thetaFreq, rhoFreq, phiFreq)) ->
       A.sfoldl
         (\s particle ->
            s + (coefficient' sigma rFreq thetaFreq rhoFreq phiFreq particle))
         0
         (constant Z)
         xs)
    freqArr

{-# INLINE coefficient'' #-}
coefficient'' ::
     forall a.
     ( A.Floating a
     , A.Num a
     , A.RealFloat a
     , A.Elt (Complex a)
     , Prelude.Fractional a
     )
  => Exp a
  -> Exp a
  -> Exp a
  -> Exp a
  -> Exp a
  -> Exp a
  -> Exp (a, a, a, a, a)
  -> Exp (A.Complex a)
coefficient'' sigma period rFreq thetaFreq rhoFreq phiFreq particle =
  let (phi, rho, theta, r, v) =
        unlift particle :: (Exp a, Exp a, Exp a, Exp a, Exp a)
   in lift
        ((v * A.exp ((sigma - 1) * rho) *
          A.cos (phiFreq * phi + thetaFreq * (theta - phi))
         ) :+
         0) *
      A.cis ((-1) * (2 * A.pi / period * (rhoFreq * rho + rFreq * (r - rho))))

gpuKernel'' ::
     forall a.
     ( A.Eq a
     , A.Floating a
     , A.Num a
     , A.RealFloat a
     , A.Elt (Complex a)
     , A.FromIntegral Int a
     , Prelude.Fractional a
     )
  => Exp a
  -> Exp a
  -> Acc (A.Vector (a, a, a, a))
  -> Acc (A.Vector (a, a, a, a, a))
  -> Acc (A.Vector (A.Complex a))
gpuKernel'' sigma period freqArr xs =
  A.map
    (\(unlift -> (rFreq, thetaFreq, rhoFreq, phiFreq)) ->
       A.sfoldl
         (\s particle ->
            s +
            coefficient'' sigma period rFreq thetaFreq rhoFreq phiFreq particle)
         0
         (constant Z)
         xs)
    freqArr

{-# INLINE pinwheelAcc #-}
pinwheelAcc ::
     forall a. (A.Floating a, A.Num a, A.RealFloat a, A.Elt (Complex a), Prelude.Fractional a)
  => Exp a
  -> Exp a
  -> Exp a
  -> Exp a
  -> Exp (a, a, a, a, A.Complex a)
  -> Exp (A.Complex a)
pinwheelAcc rFreq thetaFreq rhoFreq phiFreq particle =
  let (phi, rho, theta, r, v) =
        unlift particle :: (Exp a, Exp a, Exp a, Exp a, Exp (A.Complex a))
  in v * (lift $ ((A.exp $ (-0.5) * (rho + r))) A.:+ 0) *
     (A.cis $
      (-1) *
      (rhoFreq * rho + rFreq * (r - rho) + phiFreq * phi +
       thetaFreq * (theta - phi)))

pinwheelCoefficientsAcc ::
     forall a.
     ( A.Eq a
     , A.Floating a
     , A.Num a
     , A.RealFloat a
     , A.Elt (Complex a)
     , A.FromIntegral Int a
     , Prelude.Fractional a
     )
  => Acc (A.Vector (a, a, a, a))
  -> Acc (A.Vector (a, a, a, a, A.Complex a))
  -> Acc (A.Vector (A.Complex a))
pinwheelCoefficientsAcc freqArr xs =
  A.map
    (\(unlift -> (rFreq, thetaFreq, rhoFreq, phiFreq)) ->
       A.sfoldl
         (\s particle ->
            s + (pinwheelAcc rFreq thetaFreq rhoFreq phiFreq particle))
         0
         (constant Z)
         xs)
    freqArr


{-# LANGUAGE TypeFamilies, GeneralizedNewtypeDeriving, DeriveGeneric #-}

module BayesStack.Models.Topic.SharedTaste
  ( -- * Primitives
    STData(..)
  , ItemSource(..)
  , Node(..), Item(..), Topic(..)
  , NodeItem, setupNodeItems
  , Friendship(..), otherFriend, isFriend, getFriends
    -- * Initialization
  , ModelInit
  , randomInitialize, smartInitialize
    -- * Model
  , STModel(..), ItemUnit
  , model, modelLikelihood
  , STModelState (..), getModelState
  , sortTopics
  , ItemVars(..)
    -- * Quantities of interest
  , friendInfluence
  , theta
  ) where

import Prelude hiding (mapM, sum, product)

import Data.EnumMap (EnumMap)
import qualified Data.EnumMap as EM

import Data.Sequence (Seq)
import qualified Data.Sequence as SQ

import Data.Set (Set)
import qualified Data.Set as S

import qualified Data.EnumSet as ES
import qualified Data.Vector as V

import Data.Traversable
import Data.Foldable
import Data.Monoid
import Data.Function (on)
import Data.List (sortBy)

import Control.Monad (liftM)
import Control.Monad.Trans.State
import Control.Monad.Trans.Class

import Data.Random
import Data.Random.Distribution.Bernoulli
import Data.Random.Sequence
import Data.Number.LogFloat hiding (realToFrac)

import BayesStack.Core
import BayesStack.Categorical
import BayesStack.DirMulti
import BayesStack.TupleEnum
import BayesStack.Models.Topic.Types

import Control.Monad (when)
import Control.Monad.IO.Class

import Statistics.Sample
import Data.Serialize (Serialize)
import GHC.Generics

data ItemSource = Shared | Own deriving (Show, Eq, Generic, Enum, Ord)
instance Serialize ItemSource

data STData = STData { stAlphaGammaShared, stAlphaGammaOwn :: Double
                     , stAlphaOmega :: Double
                     , stAlphaPsi :: Double
                     , stAlphaLambda :: Double
                     , stAlphaPhi :: Double
                     , stNodes :: Set Node
                     , stFriendships :: Set Friendship
                     , stItems :: Set Item
                     , stTopics :: Set Topic
                     , stNodeItems :: EnumMap NodeItem (Node, Item)
                     }
               deriving (Show, Eq, Generic)
instance Serialize STData

data ItemVars = ItemVars { ivS :: !ItemSource
                         , ivF :: !Node
                         , ivT :: !Topic
                         }
                deriving (Show, Eq, Generic)
instance Serialize ItemVars

data STModel = STModel { mData :: STData
                       , mGammas :: SharedEnumMap Node (Multinom ItemSource)
                       , mOmegas :: SharedEnumMap Node (Multinom Topic)
                       , mPsis :: SharedEnumMap Node (Multinom Node)
                       , mLambdas :: SharedEnumMap Friendship (Multinom Topic)
                       , mPhis :: SharedEnumMap Topic (Multinom Item)
                       , mVars :: SharedEnumMap NodeItem ItemVars
                       , mSortedTopics :: SharedEnumMap Item [Topic]
                       }

data STModelState = STModelState { msData :: STData
                                 , msGammas :: EnumMap Node (Multinom ItemSource)
                                 , msOmegas :: EnumMap Node (Multinom Topic)
                                 , msPsis :: EnumMap Node (Multinom Node)
                                 , msLambdas :: EnumMap Friendship (Multinom Topic)
                                 , msPhis :: EnumMap Topic (Multinom Item)
                                 , msVars :: EnumMap NodeItem ItemVars
                                 , msLogLikelihood :: Double
                                 } deriving (Show, Generic)
instance Serialize STModelState

type ModelInit = EnumMap NodeItem ItemVars

randomInitialize' :: STData -> ModelInit -> RVar ModelInit
randomInitialize' d init =
  let unset = EM.keysSet (stNodeItems d) `ES.difference` EM.keysSet init
      topics = S.toList $ stTopics d
      randomInit :: NodeItem -> RVar ModelInit
      randomInit ni = do t <- randomElement topics
                         let (n,_) = EM.findWithDefault (error "Can't find nodeItem") ni (stNodeItems d)
                             friends = getFriends (S.toList $ stFriendships d) n
                         f <- randomElement friends
                         s <- bernoulli $ stAlphaGammaShared d
                         let s' = case s of
                                        True  -> Shared
                                        False -> Own
                         return $ EM.singleton ni $ ItemVars s' f t
  in liftM ((init `mappend`) . mconcat) $ forM (ES.toList unset) randomInit

randomInitialize :: STData -> RVar ModelInit
randomInitialize = (flip randomInitialize') EM.empty

smartInitialize :: STData -> RVar ModelInit
smartInitialize d =
  let STData {stTopics=topics, stNodes=nodes, stItems=items, stNodeItems=nodeItems} = d
      STData {stFriendships=friendships, stNodeItems=nis} = d
      nisInv :: EnumMap (Node,Item) [NodeItem]
      nisInv = EM.fromListWith (++) $ map (\(ni,(n,i))->((n,i),[ni])) $ EM.toList nis

      sharedTopics :: Friendship -> StateT (EnumMap Item Topic, EnumMap Friendship (Set Topic)) RVar ModelInit
      sharedTopics fs@(Friendship (a,b)) =
        let findItems n = S.fromList $ map snd $ filter (\(n',i)->n==n') $ toList nodeItems
            sharedItems = toList $ S.intersection (findItems a) (findItems b)
        in liftM mconcat $ forM sharedItems $ \x -> do
             (itemTopics,_) <- get
             t <- if x `EM.member` itemTopics
                    then return $ EM.findWithDefault (error "Item has no topic") x itemTopics
                    else do (_, friendshipTopics) <- get
                            --let possTopics = EM.findWithDefault topics fs friendshipTopics
                            let possTopics = topics -- FIXME
                            t <- lift $ randomElement $ S.toList possTopics
                            modify $ \(a,b)->(EM.insert x t a, b)
                            return t
             modify $ \(a,b)->(a, EM.insertWith S.union fs (S.singleton t) b)
             return $ EM.fromList $ do ax <- EM.findWithDefault (error "ouch") (a,x) nisInv
                                       return (ax, ItemVars Shared b t)
                                 ++ do bx <- EM.findWithDefault (error "ouch") (b,x) nisInv
                                       return (bx, ItemVars Shared a t)
  in do a <- evalStateT (mapM sharedTopics $ S.toList friendships) (EM.empty, EM.empty)
        randomInitialize' d $ mconcat a

data ItemUnit = ItemUnit { iuModel :: STModel
                         , iuNodeItem :: NodeItem
                         , iuFriends :: Set Node
                         , iuN :: Node
                         , iuVars :: Shared ItemVars
                         , iuX :: Item
                         , iuGamma :: Shared (Multinom ItemSource)
                         , iuOmega :: Shared (Multinom Topic)
                         , iuLambdas :: SharedEnumMap Friendship (Multinom Topic)
                         , iuPhis :: SharedEnumMap Topic (Multinom Item)
                         , iuState :: Shared GibbsUpdateState
                         }

model :: STData -> ModelInit -> ModelMonad (Seq ItemUnit, STModel)
model d init =
  do let STData {stTopics=topics, stNodes=nodes, stItems=items, stNodeItems=nis} = d
         STData {stFriendships=friendships} = d
         friends :: EnumMap Node (Set Node)
         friends = foldMap (\n->EM.singleton n $ S.fromList $ getFriends (S.toList friendships) n) nodes
     gammas <- newSharedEnumMap (S.toList nodes) $ \n ->
       --return $ dirMulti [ (Shared, stAlphaGammaShared d)
       --                  , (Own, stAlphaGammaOwn d) ]
       return $ multinom [ (Shared, stAlphaGammaShared d)
                         , (Own, stAlphaGammaOwn d) ]
     omegas <- newSharedEnumMap (S.toList nodes) $ \n ->
       return $ symDirMulti (stAlphaOmega d) (S.toList topics)
     psis <- newSharedEnumMap (S.toList nodes) $ \n ->
       return $ symDirMulti (stAlphaPsi d) (S.toList nodes)
     lambdas <- newSharedEnumMap (S.toList friendships) $ \n ->
       return $ symDirMulti (stAlphaLambda d) (S.toList topics)
     phis <- newSharedEnumMap (S.toList topics) $ \t ->
       return $ symDirMulti (stAlphaPhi d) (S.toList items)

     ivs <- newSharedEnumMap (EM.keys nis) $ \ni ->
       return $ EM.findWithDefault (error "Incomplete initialization") ni init

     sortedTopics <- newSharedEnumMap (S.toList items) $ \x -> return $ S.toList topics

     let model = STModel { mData = d
                         , mGammas = gammas
                         , mOmegas = omegas
                         , mPsis = psis
                         , mLambdas = lambdas
                         , mPhis = phis
                         , mVars = ivs
                         , mSortedTopics = sortedTopics
                         }

     itemUnits <- forM (EM.toList ivs) $ \(ni,iv) ->
       do state <- newGibbsUpdateState
          let (n,x) = nis EM.! ni
          let unit = ItemUnit { iuModel = model
                              , iuNodeItem = ni
                              , iuFriends = friends EM.! n
                              , iuN = n
                              , iuVars = iv
                              , iuX = x
                              , iuGamma = gammas EM.! n
                              , iuOmega = omegas EM.! n
                              , iuLambdas = EM.filterWithKey (\k _->isFriend n k) lambdas
                              , iuPhis = phis
                              , iuState = state
                              }
          getShared iv >>= guSet unit
          return unit
     return (SQ.fromList itemUnits, model)

sortTopics :: STModel -> ModelMonad ()
sortTopics model =
  forM_ (EM.toList $ mSortedTopics model) $ \(x,topics)->do
    d <- getShared topics
    weights <- forM d $ \t->do phi <- getShared $ mPhis model EM.! t
                               return $ sampleProb phi x -- FIXME
    setShared topics $ map snd $ sortBy (flip (compare `on` fst)) $ zip weights d

modelLikelihood :: STModelState -> Probability
modelLikelihood model =
  product $ map likelihood (EM.elems $ msGammas model)
         ++ map likelihood (EM.elems $ msPhis model)
         ++ map likelihood (EM.elems $ msLambdas model)
         ++ map likelihood (EM.elems $ msOmegas model)
         ++ map likelihood (EM.elems $ msPsis model)

instance GibbsUpdateUnit ItemUnit where
  type GUValue ItemUnit = ItemVars
  guProb unit (ItemVars s f t) =
    do gamma <- getShared $ iuGamma unit
       omega <- getShared $ iuOmega unit
       psi <- getShared $ mPsis (iuModel unit) EM.! iuN unit
       phi <- getShared $ iuPhis unit EM.! t
       lambda <- getShared $ iuLambdas unit EM.! Friendship (iuN unit, f)
       case s of
            Shared -> return $ sampleProb gamma s
                             * sampleProb psi f
                             * sampleProb lambda t
                             * sampleProb phi (iuX unit)
            Own -> return $ sampleProb gamma s
                          * sampleProb omega t
                          * sampleProb phi (iuX unit)

  guDomain unit = return $ (do t <- S.toList $ stTopics $ mData $ iuModel unit
                               f <- S.toList $ iuFriends unit
                               return $ ItemVars Shared f t)
                        ++ (do t <- S.toList $ stTopics $ mData $ iuModel unit
                               let f = head $ S.toList $ iuFriends unit
                               return $ ItemVars Own f t)

  guUnset unit =
    do ItemVars s f t <- getShared $ iuVars unit
       let x = iuX unit
           u = iuN unit
           m = iuModel unit
           gamma = iuGamma unit
           omega = iuOmega unit
           lambda = iuLambdas unit EM.! Friendship (iuN unit, f)
           phi = iuPhis unit EM.! t
       gamma `updateShared` decMultinom s
       case s of
            Shared -> do (mPsis m EM.! u) `updateShared` decMultinom f
                         (mPsis m EM.! f) `updateShared` decMultinom u
                         lambda `updateShared` decMultinom t
            Own -> omega `updateShared` decMultinom t
       phi `updateShared` decMultinom x
       return $ ItemVars s f t

  guSet unit iv@(ItemVars s f t) =
    do iuVars unit `setShared` iv
       let x = iuX unit
           u = iuN unit
           m = iuModel unit
           gamma = iuGamma unit
           omega = iuOmega unit
           lambda = iuLambdas unit EM.! Friendship (iuN unit, f)
           phi = iuPhis unit EM.! t
       gamma `updateShared` incMultinom s
       case s of
            Shared -> do (mPsis m EM.! u) `updateShared` incMultinom f
                         (mPsis m EM.! f) `updateShared` incMultinom u
                         lambda `updateShared` incMultinom t
            Own -> omega `updateShared` incMultinom t
       phi `updateShared` incMultinom x

  guState = iuState

getModelState :: STModel -> ModelMonad STModelState
getModelState model =
  do gammas <- getSharedEnumMap $ mGammas model
     omegas <- getSharedEnumMap $ mOmegas model
     psis <- getSharedEnumMap $ mPsis model
     lambdas <- getSharedEnumMap $ mLambdas model
     phis <- getSharedEnumMap $ mPhis model
     vars <- getSharedEnumMap $ mVars model
     let state = STModelState { msData = mData model
                              , msGammas = gammas
                              , msOmegas = omegas
                              , msPsis = psis
                              , msLambdas = lambdas
                              , msPhis = phis
                              , msVars = vars
                              , msLogLikelihood = logFromLogFloat $ modelLikelihood state
                              }
     return state

-- | The analogue of theta in LDA
theta :: STModelState -> Node -> Topic -> Double
theta state u t =
  sampleProb gamma Own * sampleProb omega t
  + sampleProb gamma Shared
  * sum (map (\f->let lambda = msLambdas state EM.! Friendship (u,f)
                  in sampleProb psi f * sampleProb lambda t
             )
         $ getFriends (S.toList $ stFriendships $ msData state) u
        )
  where psi = msPsis state EM.! u
        gamma = msGammas state EM.! u
        omega = msOmegas state EM.! u

friendInfluence :: STModelState -> Node -> Node -> Double
friendInfluence state u f =
  let lambda = msLambdas state EM.! Friendship(u,f)
      vars = map snd
             $ filter (\(ni,iv)->let (u',x) = stNodeItems (msData state) EM.! ni in u==u')
             $ EM.assocs $ msVars state
      tProbF = map (realToFrac . sampleProb lambda . ivT) vars
  in case tProbF of
       [] -> error "friendInfluence: vars is null"
       otherwise -> geometricMean $ V.fromList tProbF


{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE FlexibleContexts #-}
module Math.HiddenMarkovModel (
   T(..), Distr.State, state,
   Discrete, DiscreteTrained,
   Gaussian, GaussianTrained,
   uniform,
   generate,
   probabilitySequence,
   Normalized.logLikelihood,
   Normalized.reveal,
   Normalized.posterior,

   Trained(..),
   trainSupervised,
   Normalized.trainUnsupervised,
   mergeTrained, finishTraining, trainMany,
   deviation,

   toCSV,
   fromCSV,
   ) where

import qualified Math.HiddenMarkovModel.Distribution as Distr
import qualified Math.HiddenMarkovModel.Normalized as Normalized
import qualified Math.HiddenMarkovModel.CSV as HMMCSV
import Math.HiddenMarkovModel.Private
          (T(..), Trained(..), mergeTrained, toCells, parseCSV)
import Math.HiddenMarkovModel.Distribution (State(State))
import Math.HiddenMarkovModel.Utility
          (randomItemProp, normalizeProb, attachOnes)

import qualified Numeric.LinearAlgebra.Algorithms as Algo
import qualified Numeric.Container as NC
import qualified Data.Packed.Matrix as Matrix
import qualified Data.Packed.Vector as Vector
import Data.Packed.Matrix (Matrix)
import Data.Packed.Vector (Vector)

import qualified Text.CSV.Lazy.String as CSV

import qualified System.Random as Rnd

import qualified Control.Monad.Exception.Synchronous as ME
import qualified Control.Monad.Trans.State as MS
import qualified Control.Monad.HT as Monad

import qualified Data.NonEmpty as NonEmpty
import qualified Data.Array as Array
import Data.Traversable (Traversable, mapAccumL)
import Data.Foldable (Foldable)
import Data.Array (accumArray)

state :: Int -> State
state = State

type DiscreteTrained prob symbol = Trained (Distr.DiscreteTrained prob symbol) prob
type Discrete prob symbol = T (Distr.Discrete prob symbol) prob

type GaussianTrained a = Trained (Distr.GaussianTrained a) a
type Gaussian a = T (Distr.Gaussian a) a

{- |
Create a model with uniform probabilities
for initial vector and transition matrix
given a distribution for the emissions.
You can use this as a starting point for 'Normalized.trainUnsupervised'.
-}
uniform ::
   (Distr.Info distr, Distr.Probability distr ~ prob) =>
   distr -> T distr prob
uniform distr =
   let n = Distr.numberOfStates distr
       c = recip $ fromIntegral n
   in  Cons {
          initial = NC.constant c n,
          transition = NC.konst c (n,n),
          distribution = distr
       }

probabilitySequence ::
   (Traversable f, Distr.EmissionProb distr,
    Distr.Probability distr ~ prob, Distr.Emission distr ~ emission) =>
   T distr prob -> f (State, emission) -> f prob
probabilitySequence hmm =
   snd
   .
   mapAccumL
      (\index (State s, e) ->
         (NC.atIndex (transition hmm) . flip (,) s,
          index s * Distr.emissionStateProb (distribution hmm) e (State s)))
      (NC.atIndex (initial hmm))

generate ::
   (Rnd.RandomGen g, Ord prob, Rnd.Random prob,
    Distr.Generate distr, Distr.Probability distr ~ prob, Distr.Emission distr ~ emission) =>
   T distr prob -> g -> [emission]
generate hmm =
   MS.evalState $
   flip MS.evalStateT (initial hmm) $
   Monad.repeat $ MS.StateT $ \v0 -> do
      s <- randomItemProp $ zip [0..] (Vector.toList v0)
      x <- Distr.generate (distribution hmm) (State s)
      return (x, takeColumn s $ transition hmm)

takeColumn :: (Matrix.Element a) => Int -> Matrix a -> Vector a
takeColumn n  =  Matrix.flatten . Matrix.extractColumns [n]

{- |
Contribute a manually labeled emission sequence to a HMM training.
-}
trainSupervised ::
   (Distr.Estimate tdistr, Distr.Distribution tdistr ~ distr,
    Distr.Probability distr ~ prob, Distr.Emission distr ~ emission) =>
   Int -> NonEmpty.T [] (State, emission) -> Trained tdistr prob
trainSupervised n xs =
   let getState (State s, _x) = s
   in  Trained {
          trainedInitial = NC.assoc n 0 [(getState (NonEmpty.head xs), 1)],
          trainedTransition =
             Matrix.trans $ NC.accum (NC.konst 0 (n,n)) (+) $
             attachOnes $ NonEmpty.mapAdjacent (,) $ fmap getState xs,
          trainedDistribution =
             Distr.accumulateEmissions $ map attachOnes $ Array.elems $
             accumArray (flip (:)) [] (State 0, State (n-1)) $ NonEmpty.flatten xs
       }

finishTraining ::
   (Distr.Estimate tdistr, Distr.Distribution tdistr ~ distr,
    Distr.Probability distr ~ prob) =>
   Trained tdistr prob -> T distr prob
finishTraining hmm =
   Cons {
      initial = normalizeProb $ trainedInitial hmm,
      transition =
         Matrix.fromColumns $ map normalizeProb $
         Matrix.toColumns $ trainedTransition hmm,
      distribution = Distr.normalize $ trainedDistribution hmm
   }

trainMany ::
   (Distr.Estimate tdistr, Distr.Distribution tdistr ~ distr,
    Distr.Probability distr ~ prob,
    Foldable f) =>
   (trainingData -> Trained tdistr prob) ->
   NonEmpty.T f trainingData -> T distr prob
trainMany train =
   finishTraining . NonEmpty.foldl1Map mergeTrained train

{- |
Compute maximum deviation between initial and transition probabilities.
You can use this as abort criterion for unsupervised training.
We omit computation of differences between the emission probabilities.
This simplifies matters a lot and
should suffice for defining an abort criterion.
-}
deviation ::
   (Algo.Field prob, Ord prob) => T distr prob -> T distr prob -> prob
deviation hmm0 hmm1 =
   deviationVec (initial hmm0) (initial hmm1)
   `max`
   deviationVec (transition hmm0) (transition hmm1)

deviationVec ::
   (Ord a, NC.Container c a) =>
   c a -> c a -> a
deviationVec x y =
   let d = NC.sub x y
   in  NC.maxElement d `max` negate (NC.minElement d)

toCSV ::
   (Distr.CSV distr, Algo.Field prob, Show prob) =>
   T distr prob -> String
toCSV hmm =
   CSV.ppCSVTable $ snd $ CSV.toCSVTable $ HMMCSV.padTable "" $
   toCells hmm

fromCSV ::
   (Distr.CSV distr, Algo.Field prob, Read prob) =>
   String -> ME.Exceptional String (T distr prob)
fromCSV =
   MS.evalStateT parseCSV . map HMMCSV.fixShortRow . CSV.parseCSV


{-# LANGUAGE OverloadedLists #-}
{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE Strict #-}
{-# LANGUAGE TupleSections #-}
module Physics.Mechanics where

import           Numeric.LinearAlgebra (Matrix, Vector, scale, (#>), cross, cmap,dot, ident, tr, (<>), norm_2)
import qualified Number.Quaternion as Q
import           Control.Lens hiding ((<.>))
import qualified Data.Vector.Storable as VS
import           Control.Arrow hiding ((<+>))
import           Debug.Trace
import           Data.Maybe (fromMaybe)
import           Data.Function (on)

import Data.PhysicsData
import Data.GameObject
import Data.Mesh
import Physics.Collision
import qualified Data.Vector as V
import Data.List.Lens

import Utils(step,norm)
import Utils.Quaternions

infixr 1 <+>
(<+>) :: Vector Float -> Vector Float -> Vector Float
(<+>) = VS.zipWith (+)

-- (<*>) :: Vector Float -> Vector Float -> Vector Float
-- (<*>) = VS.zipWith (*)

(<.>) :: Vector Float -> Vector Float -> Float
(<.>) = (VS.sum .) . VS.zipWith (*)

addMaybeForce :: Maybe (Vector Float) -> Maybe PhysicsData -> Maybe PhysicsData
addMaybeForce mf mph = over acc . (<+>) <$> mf <*> mph

putForceLoc :: Vector Float -> Vector Float -> GameObject -> GameObject
putForceLoc point force obj =
          let m    = preview (physics._Just.massInv) obj
              lin  = flip scale (scale (force <.> point) force) <$> m
              ang  = cross force point
          in over physics (addMaybeForce lin)
           . over (physics._Just.torque) (<+>ang)
           $ obj

putForce :: Vector Float -> Vector Float -> GameObject -> GameObject
putForce point force obj =
          let p    = point - obj^.location
          in putForceLoc p force obj

putForceOnMid = over (physics._Just.acc) . (+)

putForceOnTip vec obj =
  let quat = obj^.mesh.quaternion
      tip  = rotateWithQuat quat [0,0,-0.1] + obj^.location
  in putForce tip vec obj

putForceOnTipLoc vec obj =
  let quat = obj^.mesh.quaternion
      f    = rotateWithQuat quat vec
  in putForceOnTip f obj

putForceOnSide vec obj =
  let quat = obj^.mesh.quaternion
      tip  = rotateWithQuat quat [1,0,0] <+> obj^.location
  in putForce tip vec obj

putForceOnSideLoc vec obj =
  let quat = obj^.mesh.quaternion
      f    = rotateWithQuat quat vec
  in putForceOnSide f obj

addScaled
  :: Setter GameObject GameObject (VS.Vector Float) (VS.Vector Float)
  -> Lens PhysicsData PhysicsData (VS.Vector Float) (VS.Vector Float)
  -> GameObject -> GameObject
addScaled lens1 lens2 = applyWithPhysics helper
  where helper ph = over lens1 (<+> VS.map (*step) (ph^.lens2))

applySpeed = addScaled location speed

gForce :: Vector Float
gForce = [0,9,0]

gravityF = over (physics._Just.acc) (subtract gForce)

gravCoeff = scale step gForce

getVec :: Lens PhysicsData PhysicsData (VS.Vector Float) (VS.Vector Float) -> GameObject -> VS.Vector Float
getVec l = fromMaybe [0,0,0] . preview (physics._Just.l)

getMat :: Lens PhysicsData PhysicsData (Matrix Float) (Matrix Float) -> GameObject -> Matrix Float
getMat l = fromMaybe (ident 3) . preview (physics._Just.l)

floatOnGround obj =
  let floorNormal = [0,1,0]
      loc       = obj^.location
      d         = dot floorNormal loc + 2
      f         = negate (2-d)^(2 :: Int)
  in if d < 1
  then putForceOnMid (scale f gForce) obj
  else obj

calculateRelPos contact obj = contact^.contactPoint - obj^.location

linearInertia obj =
  fromMaybe 0 $ obj^?physics._Just.massInv

adjustMove relPos p@(l, an) =
  let limit = norm relPos * 0.1
      newAn = signum an * limit
  in if abs an > limit
     then (l + an - newAn,newAn)
     else p

collisionMoves :: Contact -> (GameObject, Vector Float) -> (GameObject, Vector Float) -> ((Float, Vector Float), (Float, Vector Float))
collisionMoves contact p1 p2 =
  let ang1 = angularInertia contact p1
      lin1 = linearInertia $ fst p1
      ang2 = angularInertia contact p2
      lin2 = linearInertia $ fst p2
      coef = contact^.penetration / (ang1 + ang2 + lin1 + lin2)
      (fli1, fan1) = adjustMove (snd p1) (coef*lin1, coef*ang1)
      (fli2, fan2) = adjustMove (snd p2) (coef*lin2, coef*ang2)
      finang1 = getRotationPM contact p1 fan1 ang1
      finang2 = getRotationPM contact p1 fan2 ang2
  in ((fli1, finang1),(fli2, finang2))

getRotationPM :: Contact -> (GameObject, Vector Float) -> Float -> Float -> Vector Float
getRotationPM contact (obj, relPos) ang angIn =
  if angIn == 0
  then [0,0,0]
  else
    scale (ang) .
    scale (1/angIn) $
    getInertiaInv obj #> cross relPos (contact^.contactNormal)

applyMoves :: Contact -> (Float, Vector Float) -> GameObject -> GameObject
applyMoves contact (lin, ang) =
  over (mesh.quaternion) (qvMul ang) .
  over location (+ scale lin (contact^.contactNormal))

calculateDeltaVel :: Contact -> (GameObject, Vector Float) -> Float
calculateDeltaVel contact p@(obj, _) =
  case obj^.physics of
    Just pd ->
      (pd^.massInv +) $
      angularInertia contact p
      -- dot (contact^.contactNormal) .
      -- flip cross relPos .
      -- ((pd^.inertiaInv) #>) .
      -- flip cross (contact^.contactNormal) $ relPos
    Nothing -> 0

angCoeff = 1

angularInertia :: Contact ->  (GameObject, Vector Float) -> Float
angularInertia contact (obj, relPos) =
  case obj^.physics of
    Just pd ->
--      (*angCoeff) .
      dot (contact^.contactNormal) .        -- get component of speed pointing along the normal
      (\angS -> cross angS relPos) .        -- get linear speed of pt
      ((pd^.inertiaInv) #>) $               -- Get ang speed
      cross relPos (contact^.contactNormal) -- for unit impulse
    Nothing -> 0

calculateClosingVel :: Contact -> (GameObject, Vector Float) -> Float
calculateClosingVel contact (obj, relPos) =
  case obj^.physics of
    Just pd ->
      -- (makeContactBasis contact #>) .
      dot (contact^.contactNormal) .
      (pd^.speed +) . cross (pd^.angularS) $ relPos
    Nothing -> 0

calculateImpulse :: Contact -> (GameObject, Vector Float) -> (GameObject, Vector Float) -> Vector Float
calculateImpulse contact o1 o2 =
  let deltaV = (("Delta: "++) . show >>= trace) $
               ((+) `on` calculateDeltaVel contact) o1 o2
      closiV = (("Closing: "++) . show >>= trace) $
               ((+) `on` calculateClosingVel contact) o1 o2
      contV  = negate $ closiV
      restitution = if abs contV < 2
                    then 0.0
                    else 0.0
      impulse = -- if signum contV /= signum (contact^.sign)
                -- then
                  scale (contV * (1.0 + restitution)/deltaV) $ contact^.contactNormal
                -- else [0,0,0]
  in impulse

applyImpulse :: Vector Float -> (GameObject, Vector Float) -> GameObject
applyImpulse impulse (obj, relPos) =
  let vChange = -- impulse
        (("Linear impulse: "++) . show >>= trace) .
        fromMaybe [0,0,0] $
        flip scale impulse <$> (obj^?physics._Just.massInv)
      rotChange = (("Angular impulse: "++) . show >>= trace) .
--        scale angCoeff .
        fromMaybe [0,0,0] $
        (#> cross impulse relPos) <$> (obj^?physics._Just.inertiaInv)
  in over (physics._Just.speed) (+ vChange) .
     over (physics._Just.angularS) (+ rotChange) $ obj

resolveCollision :: (Int, Int) -> [Contact] -> V.Vector GameObject -> V.Vector GameObject
resolveCollision (ix1, ix2) (contact:cs) objs =
  let obj1 = objs V.! ix1
      obj2 = objs V.! ix2
      movement = scale (contact^.penetration*contact^.sign) $
                 contact^.contactNormal
      move = case obj1^.physics of
               Just _ -> over (ix 0._2.location) (+movement)
               Nothing -> over (ix 1._2.location) (subtract movement)
      relPos1 = (("RelPos1: "++) . show >>= trace) $ calculateRelPos contact obj1
      relPos2 = (("RelPos2: "++) . show >>= trace) $ calculateRelPos contact obj2
      impulse = (("Impulse: "++) . show >>= trace) .
                scale (contact^.sign) $
                calculateImpulse contact (obj1, relPos1) (obj2, relPos2)
      (m1, m2) = collisionMoves contact (obj1, relPos1) (obj2, relPos2)
      update i p =
        -- applyMoves contact p .
        applyImpulse i
      updates = move $
                [(ix1, update (impulse) m1 (obj1, relPos1))
                ,(ix2, update (-impulse) m2 (obj2, relPos2))]
      newObjs = V.unsafeUpd objs updates
      check   = V.unsafeIndex newObjs
      -- newColls =
      --   take (length cs) $
      --   detectCollision (prepareCPrim $ check ix1) (prepareCPrim $ check ix2)^.contacts
  in
    resolveCollision (ix1, ix2) cs newObjs
resolveCollision _ [] objs = objs

floorReaction obj =
  let floorHeight = [0,-2,0]
      floorNormal = [0,1,0] :: Vector Float
      vert      = obj^.mesh.vertices
      loc       = obj^.location
      r         = quatToMat3 $ obj^.mesh.quaternion
      m         = (1/) $ fromMaybe 1 $ obj^?physics._Just.massInv
      translate = (+loc) . subtract floorHeight
      collides = filter ((<0) . snd)
               . map (id &&& (dot floorNormal . translate) . (tr r #>))
               $ vert
      minD :: Float
      minD = negate . foldr min 0 . map snd $ collides
      response (point, depth) =
        let
          vrel = dot floorNormal $ getVec speed obj + cross (getVec angularS obj) point
          aobj = dot floorNormal $ cross (i #> cross point floorNormal) point
          i    = getMat inertiaInv obj
          j    = negate (0.8) * vrel/(m+aobj)
          force = scale j floorNormal
          torque = getMat inertiaInv obj #> cross point force
            -- getMat inertiaInv obj #> cross point force
        in (force, torque)
      (f,t) = foldr (\(a,b) (c,d) -> (a+c,b+d)) ([0,0,0],[0,0,0])
            . map response
            $ collides
      n = fromIntegral $ length collides
  in if n /= 0
  then over (physics._Just.angularS) (+ scale (1/n) t)
     . over (physics._Just.speed) (+ scale (1/n) f)
     . over (location) (+ scale minD floorNormal)
     $ obj
  else obj

applyAcc :: GameObject -> GameObject
applyAcc = addScaled (physics._Just.speed) acc

applyForce :: GameObject -> GameObject
applyForce = applyWithPhysics helper
  where helper ph obj =
         let tq    = ph^.torque
             om    = ph^.angularS
             inInv = ph^.inertiaInv
             newOm = om <+> (inInv #> scale step tq)
             quatD = qvMul (scale step newOm) $ obj^.mesh.quaternion
             newQ  = Q.normalize . addQuat quatD $ obj^.mesh.quaternion
             inInD = quatToLInv (ph^.inertiaD) newQ
         in set (physics._Just.torque) [0,0,0]
          . set (physics._Just.acc) [0,0,0]
          . applySpeed
--          . floorReaction
          . applyAcc
--          . floatOnGround
          . set (physics._Just.angularS) newOm
          . set (physics._Just.inertiaInv) inInD
          . set (mesh.quaternion) newQ
          $ obj

applyLinear :: GameObject -> GameObject
applyLinear = set (physics._Just.torque) [0,0,0]
              . set (physics._Just.acc) [0,0,0]
              . applySpeed
              -- . floorReaction
              . applyAcc
              -- . gravityF

applyWithPhysics :: (PhysicsData -> GameObject -> GameObject)
                   -> GameObject -> GameObject
applyWithPhysics f obj = case obj^.physics of
  Just phy -> f phy obj
  Nothing  -> obj

maxspeed = 50

yolo :: Float -> Vector Float -> Vector Float
yolo maxVal vec =
  let vecNorm = norm vec
      coeff   = (vecNorm/maxVal)**(1/2)
  in scale (0 - max 1 coeff) vec

friction = applyWithPhysics helper
  where helper phy =
          let maxspeed  = 1000000
              maxtorque = 100
              coeff1    = scale step . yolo maxspeed $ phy^.speed
              coeff2    = scale step . yolo maxtorque $ phy^.angularS
          in over (physics._Just.angularS) (<+> coeff2)
            . over (physics._Just.speed) (<+> coeff1)

  -- applyWithPhysics helper
  -- where helper phy obj =
  --         let quat   = obj^.mesh.quaternion
  --             rot    = rotateWithQuat quat [0,0,-1]
  --               -- (\(a,b,c) -> [a,b,c]) . imag . flip quatConcat (qInverse quat) $ quatConcat quat (0 +:: (0,0,-1))
  --             force  = 40 -- exp (maxspeed - sNorm) - 1
  --         in over (physics._Just.acc) (<+>cmap (*force) rot) obj


{-# LANGUAGE DataKinds             #-}
{-# LANGUAGE DeriveAnyClass        #-}
{-# LANGUAGE DeriveGeneric         #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE TypeFamilies          #-}
{-# LANGUAGE TypeOperators         #-}
{-|
Module      : Grenade.Layers.Logit
Description : Exponential linear unit layer
Copyright   : (c) Huw Campbell, 2016-2017
License     : BSD2
Stability   : experimental
-}
module Grenade.Layers.Elu (
    Elu (..)
  ) where

import           Control.DeepSeq              (NFData)
import           Data.Serialize
import           GHC.Generics                 (Generic)

import           GHC.TypeLits
import           Grenade.Core

import qualified Numeric.LinearAlgebra.Static as LAS

-- | An exponential linear unit.
--   A layer which can act between any shape of the same dimension, acting as a
--   diode on every neuron individually.
data Elu = Elu
  deriving (Generic, NFData, Show)

instance UpdateLayer Elu where
  type Gradient Elu = ()
  runUpdate _ _ _ = Elu

instance RandomLayer Elu where
  createRandomWith _ _ = return Elu

instance Serialize Elu where
  put _ = return ()
  get = return Elu

instance ( KnownNat i) => Layer Elu ('D1 i) ('D1 i) where
  type Tape Elu ('D1 i) ('D1 i) = LAS.R i

  runForwards _ (S1D y) = (y, S1D (elu y))
    where
      elu = LAS.dvmap (\a -> if a <= 0 then exp a - 1 else a)
  runBackwards _ y (S1D dEdy) = ((), S1D (elu' y * dEdy))
    where
      elu' = LAS.dvmap (\a -> if a <= 0 then exp a else 1)

instance (KnownNat i, KnownNat j) => Layer Elu ('D2 i j) ('D2 i j) where
  type Tape Elu ('D2 i j) ('D2 i j) = S ('D2 i j)

  runForwards _ (S2D y) = (S2D y, S2D (elu y))
    where
      elu = LAS.dmmap (\a -> if a <= 0 then exp a - 1 else a)
  runBackwards _ (S2D y) (S2D dEdy) = ((), S2D (elu' y * dEdy))
    where
      elu' = LAS.dmmap (\a -> if a <= 0 then exp a else 1)

instance (KnownNat i, KnownNat j, KnownNat k) => Layer Elu ('D3 i j k) ('D3 i j k) where

  type Tape Elu ('D3 i j k) ('D3 i j k) = S ('D3 i j k)

  runForwards _ (S3D y) = (S3D y, S3D (elu y))
    where
      elu = LAS.dmmap (\a -> if a <= 0 then exp a - 1 else a)
  runBackwards _ (S3D y) (S3D dEdy) = ((), S3D (elu' y * dEdy))
    where
      elu' = LAS.dmmap (\a -> if a <= 0 then exp a else 1)

-------------------- GNum instances --------------------

instance GNum Elu where
  _ |* Elu = Elu
  _ |+ Elu = Elu
  gFromRational _ = Elu


{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE CPP #-}
module Numeric.QpOases where

import Control.Monad
import Control.Monad.IO.Class
import Control.Monad.Trans.Except
import qualified Data.Vector.Storable as VS
import Data.Matrix.CSC
import Foreign hiding (void)
import Foreign.C.Types
import Numeric.LinearAlgebra
import Numeric.LinearAlgebra.Devel

data OasesError = SizeMismatch | OasesError Int deriving (Show)

data HessianType = HST_ZERO | HST_IDENTITY | HST_POSDEF | HST_POSDEF_NULLSPACE | HST_SEMIDEF | HST_INDEF | HST_UNKNOWN deriving (Enum, Show)

newtype SQProblem = SQProblem (ForeignPtr SQProblem) deriving (Eq, Show)

newtype SQProblemSchur = SQProblemSchur (ForeignPtr SQProblemSchur) deriving (Eq, Show)

setupSQP :: Int -> Int -> HessianType -> IO SQProblem
setupSQP nV nC h = do
    ptr2 <- alloca $ \ptr -> do
        !_ <- sqproblem_setup (fromIntegral nV) (fromIntegral nC) (fromIntegral $ fromEnum h) ptr
        peek ptr
    SQProblem <$> newForeignPtr sqproblem_cleanup ptr2

setupSQPSchur :: Int -> Int -> HessianType -> IO SQProblemSchur
setupSQPSchur nV nC h = do
    ptr2 <- alloca $ \ptr -> do
        !_ <- sqproblem_schur_setup (fromIntegral nV) (fromIntegral nC) (fromIntegral $ fromEnum h) ptr
        peek ptr
    SQProblemSchur <$> newForeignPtr sqproblem_schur_cleanup ptr2

withSQProblem :: SQProblem -> (Ptr SQProblem -> IO a) -> IO a
withSQProblem (SQProblem fPtr) f = withForeignPtr fPtr f

withSQProblemSchur :: SQProblemSchur -> (Ptr SQProblemSchur -> IO a) -> IO a
withSQProblemSchur (SQProblemSchur fPtr) f = withForeignPtr fPtr f

initSQP :: SQProblem -> Matrix Double -> Vector Double -> Matrix Double
        -> Vector Double -> Vector Double -> Maybe Double
        -> ExceptT OasesError IO (Vector Double, Double)
initSQP sqp h g a lbA ubA m'cpuTime =
  mat' h $ \hRow hCol hPtr ->
  vec' g $ \gSize gPtr ->
  mat' a $ \aRow aCol aPtr ->
  vec' lbA $ \lbASize lbAPtr ->
  vec' ubA $ \ubASize ubAPtr -> do
    let nVar = gSize
    let nConstr = aRow
    unless (nVar == hRow) $ throwE SizeMismatch
    unless (nVar == hCol) $ throwE SizeMismatch
    unless (nVar == aCol) $ throwE SizeMismatch
    unless (nConstr == lbASize) $ throwE SizeMismatch
    unless (nConstr == ubASize) $ throwE SizeMismatch
    primalFP <- liftIO $ mallocForeignPtrArray (fromIntegral nVar)
    (obj, status) <- liftIO $ withForeignPtr primalFP $ \xPtr ->
      allocaArray (fromIntegral $ nVar + nConstr) $ \yPtr ->
      alloca $ \objPtr ->
      alloca $ \statusPtr ->
      with (5 * fromIntegral (nVar + nConstr)) $ \nWSRPtr ->
      maybe ($ nullPtr) with m'cpuTime $ \cpuTimePtr -> do
        _ <- withSQProblem sqp $ \ptr ->
               sqproblem_init ptr hPtr gPtr aPtr nullPtr nullPtr lbAPtr ubAPtr
               nWSRPtr cpuTimePtr xPtr yPtr objPtr statusPtr
        !obj <- peek objPtr
        !status <- fromIntegral <$> peek statusPtr
        return $ (obj, status)
    let solutionVec = VS.unsafeFromForeignPtr0 primalFP (fromIntegral nVar)
    if status == 0
      then return (solutionVec, obj)
      else throwE (OasesError status)
  where
    mat' m f = ExceptT $ mat (cmat m) $ \church -> church $ \nrow ncol ptr -> runExceptT $ f nrow ncol ptr
    vec' v f = ExceptT $ vec v $ \church -> church $ \size ptr -> runExceptT $ f size ptr

initSparseSQP :: SQProblem -> CSC Double -> Vector Double -> Matrix Double
              -> Vector Double -> Vector Double -> Maybe Double
              -> ExceptT OasesError IO (Vector Double, Double)
initSparseSQP sqp h g a lbA ubA m'cpuTime =
  vec' (cscVals h) $ \_hValsSize hValsPtr ->
  vec' (cscCols h) $ \_hColsSize hColsPtr ->
  vec' (cscRows h) $ \_hRowsSize hRowsPtr ->
  vec' g $ \gSize gPtr ->
  mat' a $ \aRow aCol aPtr ->
  vec' lbA $ \lbASize lbAPtr ->
  vec' ubA $ \ubASize ubAPtr -> do
    let nVar = gSize
    let nConstr = aRow
    unless (nVar == aCol) $ throwE SizeMismatch
    unless (nConstr == lbASize) $ throwE SizeMismatch
    unless (nConstr == ubASize) $ throwE SizeMismatch
    primalFP <- liftIO $ mallocForeignPtrArray (fromIntegral nVar)
    (obj, status) <- liftIO $ withForeignPtr primalFP $ \xPtr ->
      allocaArray (fromIntegral $ nVar + nConstr) $ \yPtr ->
      alloca $ \objPtr ->
      alloca $ \statusPtr ->
      with (5 * fromIntegral (nVar + nConstr)) $ \nWSRPtr ->
      maybe ($ nullPtr) with m'cpuTime $ \cpuTimePtr -> do
        _ <- withSQProblem sqp $ \ptr ->
               sqproblem_sparse_init ptr hRowsPtr hColsPtr hValsPtr gPtr aPtr
               nullPtr nullPtr lbAPtr ubAPtr
               nWSRPtr cpuTimePtr xPtr yPtr objPtr statusPtr
        !obj <- peek objPtr
        !status <- fromIntegral <$> peek statusPtr
        return $ (obj, status)
    let solutionVec = VS.unsafeFromForeignPtr0 primalFP (fromIntegral nVar)
    if status == 0
      then return (solutionVec, obj)
      else throwE (OasesError status)
  where
    mat' m f = ExceptT $ mat (cmat m) $ \church -> church $ \nrow ncol ptr -> runExceptT $ f nrow ncol ptr
    vec' v f = ExceptT $ vec v $ \church -> church $ \size ptr -> runExceptT $ f size ptr

initSparseSQPSchur :: SQProblemSchur -> CSC Double -> Vector Double -> Matrix Double
                   -> Vector Double -> Vector Double -> Maybe Double
                   -> ExceptT OasesError IO (Vector Double, Double)
initSparseSQPSchur sqp h g a lbA ubA m'cpuTime =
  vec' (cscVals h) $ \_hValsSize hValsPtr ->
  vec' (cscCols h) $ \_hColsSize hColsPtr ->
  vec' (cscRows h) $ \_hRowsSize hRowsPtr ->
  vec' g $ \gSize gPtr ->
  mat' a $ \aRow aCol aPtr ->
  vec' lbA $ \lbASize lbAPtr ->
  vec' ubA $ \ubASize ubAPtr -> do
    let nVar = gSize
    let nConstr = aRow
    unless (nVar == aCol) $ throwE SizeMismatch
    unless (nConstr == lbASize) $ throwE SizeMismatch
    unless (nConstr == ubASize) $ throwE SizeMismatch
    primalFP <- liftIO $ mallocForeignPtrArray (fromIntegral nVar)
    (obj, status) <- liftIO $ withForeignPtr primalFP $ \xPtr ->
      allocaArray (fromIntegral $ nVar + nConstr) $ \yPtr ->
      alloca $ \objPtr ->
      alloca $ \statusPtr ->
      with (5 * fromIntegral (nVar + nConstr)) $ \nWSRPtr ->
      maybe ($ nullPtr) with m'cpuTime $ \cpuTimePtr -> do
        _ <- withSQProblemSchur sqp $ \ptr ->
               sqproblem_sparse_schur_init ptr hRowsPtr hColsPtr hValsPtr gPtr aPtr
               nullPtr nullPtr lbAPtr ubAPtr
               nWSRPtr cpuTimePtr xPtr yPtr objPtr statusPtr
        !obj <- peek objPtr
        !status <- fromIntegral <$> peek statusPtr
        return $ (obj, status)
    let solutionVec = VS.unsafeFromForeignPtr0 primalFP (fromIntegral nVar)
    if status == 0
      then return (solutionVec, obj)
      else throwE (OasesError status)
  where
    mat' m f = ExceptT $ mat (cmat m) $ \church -> church $ \nrow ncol ptr -> runExceptT $ f nrow ncol ptr
    vec' v f = ExceptT $ vec v $ \church -> church $ \size ptr -> runExceptT $ f size ptr

initSparseSparseSQPSchur :: SQProblemSchur -> CSC Double -> Vector Double -> CSC Double
                   -> Maybe (Vector Double) -> Maybe (Vector Double)
                   -> Vector Double -> Vector Double -> Maybe Double
                   -> ExceptT OasesError IO (Vector Double, Double)
initSparseSparseSQPSchur sqp h g a m'lb m'ub lbA ubA m'cpuTime =
  vec' (cscVals h) $ \_hValsSize hValsPtr ->
  vec' (cscCols h) $ \_hColsSize hColsPtr ->
  vec' (cscRows h) $ \_hRowsSize hRowsPtr ->
  vec' g $ \gSize gPtr ->
  vec' (cscVals a) $ \_aValsSize aValsPtr ->
  vec' (cscCols a) $ \_aColsSize aColsPtr ->
  vec' (cscRows a) $ \_aRowsSize aRowsPtr ->
  maybe (\f -> f undefined nullPtr) vec' m'lb $ \_lbSize lbPtr ->
  maybe (\f -> f undefined nullPtr) vec' m'ub $ \_ubSize ubPtr ->
  vec' lbA $ \lbASize lbAPtr ->
  vec' ubA $ \ubASize ubAPtr -> do
    let nVar = gSize
    let nConstr = cscNRows a
    unless (nVar == cscNCols a) $ throwE SizeMismatch
    unless (nConstr == lbASize) $ throwE SizeMismatch
    unless (nConstr == ubASize) $ throwE SizeMismatch
    primalFP <- liftIO $ mallocForeignPtrArray (fromIntegral nVar)
    (obj, status) <- liftIO $ withForeignPtr primalFP $ \xPtr ->
      allocaArray (fromIntegral $ nVar + nConstr) $ \yPtr ->
      alloca $ \objPtr ->
      alloca $ \statusPtr ->
      with (5 * fromIntegral (nVar + nConstr)) $ \nWSRPtr ->
      maybe ($ nullPtr) with m'cpuTime $ \cpuTimePtr -> do
        _ <- withSQProblemSchur sqp $ \ptr ->
               sqproblem_sparse_sparse_schur_init ptr hRowsPtr hColsPtr hValsPtr gPtr
               aRowsPtr aColsPtr aValsPtr
               lbPtr ubPtr lbAPtr ubAPtr
               nWSRPtr cpuTimePtr xPtr yPtr objPtr statusPtr
        !obj <- peek objPtr
        !status <- fromIntegral <$> peek statusPtr
        return $ (obj, status)
    let solutionVec = VS.unsafeFromForeignPtr0 primalFP (fromIntegral nVar)
    if status == 0
      then return (solutionVec, obj)
      else throwE (OasesError status)
  where
    mat' m f = ExceptT $ mat (cmat m) $ \church -> church $ \nrow ncol ptr -> runExceptT $ f nrow ncol ptr
    vec' v f = ExceptT $ vec v $ \church -> church $ \size ptr -> runExceptT $ f size ptr

hotstartSQP :: SQProblem -> Matrix Double -> Vector Double -> Matrix Double
            -> Vector Double -> Vector Double -> Maybe Double
            -> ExceptT OasesError IO (Vector Double, Double)
hotstartSQP sqp h g a lbA ubA m'cpuTime =
  mat' h $ \hRow hCol hPtr ->
  vec' g $ \gSize gPtr ->
  mat' a $ \aRow aCol aPtr ->
  vec' lbA $ \lbASize lbAPtr ->
  vec' ubA $ \ubASize ubAPtr -> do
    let nVar = gSize
    let nConstr = aRow
    unless (nVar == hRow) $ throwE SizeMismatch
    unless (nVar == hRow) $ throwE SizeMismatch
    unless (nVar == hCol) $ throwE SizeMismatch
    unless (nVar == aCol) $ throwE SizeMismatch
    unless (nConstr == lbASize) $ throwE SizeMismatch
    unless (nConstr == ubASize) $ throwE SizeMismatch
    primalFP <- liftIO $ mallocForeignPtrArray (fromIntegral nVar)
    (obj, status) <- liftIO $ withForeignPtr primalFP $ \xPtr ->
      allocaArray (fromIntegral $ nVar + nConstr) $ \yPtr ->
      alloca $ \objPtr ->
      alloca $ \statusPtr ->
      with (5 * fromIntegral (nVar + nConstr)) $ \nWSRPtr ->
      maybe ($ nullPtr) with m'cpuTime $ \cpuTimePtr -> do
        _ <- withSQProblem sqp $ \ptr ->
               sqproblem_hotstart ptr hPtr gPtr aPtr nullPtr nullPtr lbAPtr ubAPtr
               nWSRPtr cpuTimePtr xPtr yPtr objPtr statusPtr
        !obj <- peek objPtr
        !status <- fromIntegral <$> peek statusPtr
        return $ (obj, status)
    let solutionVec = VS.unsafeFromForeignPtr0 primalFP (fromIntegral nVar)
    if status == 0
      then return (solutionVec, obj)
      else throwE (OasesError status)
  where
    mat' m f = ExceptT $ mat (cmat m) $ \church -> church $ \nrow ncol ptr -> runExceptT $ f nrow ncol ptr
    vec' v f = ExceptT $ vec v $ \church -> church $ \size ptr -> runExceptT $ f size ptr

hotstartSparseSQP :: SQProblem -> CSC Double -> Vector Double -> Matrix Double
                  -> Vector Double -> Vector Double -> Maybe Double
                  -> ExceptT OasesError IO (Vector Double, Double)
hotstartSparseSQP sqp h g a lbA ubA m'cpuTime =
  vec' (cscVals h) $ \_hValsSize hValsPtr ->
  vec' (cscCols h) $ \_hColsSize hColsPtr ->
  vec' (cscRows h) $ \_hRowsSize hRowsPtr ->
  vec' g $ \gSize gPtr ->
  mat' a $ \aRow aCol aPtr ->
  vec' lbA $ \lbASize lbAPtr ->
  vec' ubA $ \ubASize ubAPtr -> do
    let nVar = gSize
    let nConstr = aRow
    unless (nVar == aCol) $ throwE SizeMismatch
    unless (nConstr == lbASize) $ throwE SizeMismatch
    unless (nConstr == ubASize) $ throwE SizeMismatch
    primalFP <- liftIO $ mallocForeignPtrArray (fromIntegral nVar)
    (obj, status) <- liftIO $ withForeignPtr primalFP $ \xPtr ->
      allocaArray (fromIntegral $ nVar + nConstr) $ \yPtr ->
      alloca $ \objPtr ->
      alloca $ \statusPtr ->
      with (5 * fromIntegral (nVar + nConstr)) $ \nWSRPtr -> do
      maybe ($ nullPtr) with m'cpuTime $ \cpuTimePtr -> do
        _ <- withSQProblem sqp $ \ptr ->
               sqproblem_sparse_hotstart ptr hRowsPtr hColsPtr hValsPtr gPtr
               aPtr nullPtr nullPtr lbAPtr ubAPtr
               nWSRPtr cpuTimePtr xPtr yPtr objPtr statusPtr
        !obj <- peek objPtr
        !status <- fromIntegral <$> peek statusPtr
        return $ (obj, status)
    let solutionVec = VS.unsafeFromForeignPtr0 primalFP (fromIntegral nVar)
    if status == 0
      then return (solutionVec, obj)
      else throwE (OasesError status)
  where
    mat' m f = ExceptT $ mat (cmat m) $ \church -> church $ \nrow ncol ptr -> runExceptT $ f nrow ncol ptr
    vec' v f = ExceptT $ vec v $ \church -> church $ \size ptr -> runExceptT $ f size ptr

hotstartSparseSQPSchur :: SQProblemSchur -> CSC Double -> Vector Double -> Matrix Double
                  -> Vector Double -> Vector Double -> Maybe Double
                  -> ExceptT OasesError IO (Vector Double, Double)
hotstartSparseSQPSchur sqp h g a lbA ubA m'cpuTime =
  vec' (cscVals h) $ \_hValsSize hValsPtr ->
  vec' (cscCols h) $ \_hColsSize hColsPtr ->
  vec' (cscRows h) $ \_hRowsSize hRowsPtr ->
  vec' g $ \gSize gPtr ->
  mat' a $ \aRow aCol aPtr ->
  vec' lbA $ \lbASize lbAPtr ->
  vec' ubA $ \ubASize ubAPtr -> do
    let nVar = gSize
    let nConstr = aRow
    unless (nVar == aCol) $ throwE SizeMismatch
    unless (nConstr == lbASize) $ throwE SizeMismatch
    unless (nConstr == ubASize) $ throwE SizeMismatch
    primalFP <- liftIO $ mallocForeignPtrArray (fromIntegral nVar)
    (obj, status) <- liftIO $ withForeignPtr primalFP $ \xPtr ->
      allocaArray (fromIntegral $ nVar + nConstr) $ \yPtr ->
      alloca $ \objPtr ->
      alloca $ \statusPtr ->
      with (5 * fromIntegral (nVar + nConstr)) $ \nWSRPtr ->
      maybe ($ nullPtr) with m'cpuTime $ \cpuTimePtr -> do
        _ <- withSQProblemSchur sqp $ \ptr ->
               sqproblem_sparse_schur_hotstart ptr hRowsPtr hColsPtr hValsPtr gPtr
               aPtr nullPtr nullPtr lbAPtr ubAPtr
               nWSRPtr cpuTimePtr xPtr yPtr objPtr statusPtr
        !obj <- peek objPtr
        !status <- fromIntegral <$> peek statusPtr
        return $ (obj, status)
    let solutionVec = VS.unsafeFromForeignPtr0 primalFP (fromIntegral nVar)
    if status == 0
      then return (solutionVec, obj)
      else throwE (OasesError status)
  where
    mat' m f = ExceptT $ mat (cmat m) $ \church -> church $ \nrow ncol ptr -> runExceptT $ f nrow ncol ptr
    vec' v f = ExceptT $ vec v $ \church -> church $ \size ptr -> runExceptT $ f size ptr

hotstartSparseSparseSQPSchur :: SQProblemSchur -> CSC Double -> Vector Double -> CSC Double
                  -> Maybe (Vector Double) -> Maybe (Vector Double) -> Vector Double -> Vector Double -> Maybe Double
                  -> ExceptT OasesError IO (Vector Double, Double)
hotstartSparseSparseSQPSchur sqp h g a m'lb m'ub lbA ubA m'cpuTime =
  vec' (cscVals h) $ \_hValsSize hValsPtr ->
  vec' (cscCols h) $ \_hColsSize hColsPtr ->
  vec' (cscRows h) $ \_hRowsSize hRowsPtr ->
  vec' g $ \gSize gPtr ->
  vec' (cscVals a) $ \_aValsSize aValsPtr ->
  vec' (cscCols a) $ \_aColsSize aColsPtr ->
  vec' (cscRows a) $ \_aRowsSize aRowsPtr ->
  maybe (\f -> f undefined nullPtr) vec' m'lb $ \_lbSize lbPtr ->
  maybe (\f -> f undefined nullPtr) vec' m'ub $ \_ubSize ubPtr ->
  vec' lbA $ \lbASize lbAPtr ->
  vec' ubA $ \ubASize ubAPtr -> do
    let nVar = gSize
    let nConstr = cscNRows a
    unless (nVar == cscNCols a) $ throwE SizeMismatch
    unless (nConstr == lbASize) $ throwE SizeMismatch
    unless (nConstr == ubASize) $ throwE SizeMismatch
    primalFP <- liftIO $ mallocForeignPtrArray (fromIntegral nVar)
    (obj, status) <- liftIO $ withForeignPtr primalFP $ \xPtr ->
      allocaArray (fromIntegral $ nVar + nConstr) $ \yPtr ->
      alloca $ \objPtr ->
      alloca $ \statusPtr ->
      with (5 * fromIntegral (nVar + nConstr)) $ \nWSRPtr ->
      maybe ($ nullPtr) with m'cpuTime $ \cpuTimePtr -> do
        _ <- withSQProblemSchur sqp $ \ptr ->
               sqproblem_sparse_sparse_schur_hotstart ptr hRowsPtr hColsPtr hValsPtr gPtr
               aRowsPtr aColsPtr aValsPtr lbPtr ubPtr lbAPtr ubAPtr
               nWSRPtr cpuTimePtr xPtr yPtr objPtr statusPtr
        !obj <- peek objPtr
        !status <- fromIntegral <$> peek statusPtr
        return $ (obj, status)
    let solutionVec = VS.unsafeFromForeignPtr0 primalFP (fromIntegral nVar)
    if status == 0
      then return (solutionVec, obj)
      else throwE (OasesError status)
  where
    mat' m f = ExceptT $ mat (cmat m) $ \church -> church $ \nrow ncol ptr -> runExceptT $ f nrow ncol ptr
    vec' v f = ExceptT $ vec v $ \church -> church $ \size ptr -> runExceptT $ f size ptr

data Options

foreign import ccall "sqproblem_init"
    sqproblem_init
        :: Ptr SQProblem -> Ptr Double -> Ptr Double -> Ptr Double -> Ptr Double
        -> Ptr Double -> Ptr Double -> Ptr Double -> Ptr CInt -> Ptr Double
        -> Ptr Double -> Ptr Double -> Ptr Double -> Ptr CInt -> IO CInt

foreign import ccall "sqproblem_sparse_init"
    sqproblem_sparse_init
        :: Ptr SQProblem -> Ptr CInt -> Ptr CInt -> Ptr Double
        -> Ptr Double -> Ptr Double -> Ptr Double
        -> Ptr Double -> Ptr Double -> Ptr Double -> Ptr CInt -> Ptr Double
        -> Ptr Double -> Ptr Double -> Ptr Double -> Ptr CInt -> IO CInt

foreign import ccall "sqproblem_sparse_schur_init"
    sqproblem_sparse_schur_init
        :: Ptr SQProblemSchur -> Ptr CInt -> Ptr CInt -> Ptr Double
        -> Ptr Double -> Ptr Double -> Ptr Double
        -> Ptr Double -> Ptr Double -> Ptr Double -> Ptr CInt -> Ptr Double
        -> Ptr Double -> Ptr Double -> Ptr Double -> Ptr CInt -> IO CInt

foreign import ccall "sqproblem_sparse_sparse_schur_init"
    sqproblem_sparse_sparse_schur_init
        :: Ptr SQProblemSchur -> Ptr CInt -> Ptr CInt -> Ptr Double
        -> Ptr Double -> Ptr CInt -> Ptr CInt -> Ptr Double -> Ptr Double
        -> Ptr Double -> Ptr Double -> Ptr Double -> Ptr CInt -> Ptr Double
        -> Ptr Double -> Ptr Double -> Ptr Double -> Ptr CInt -> IO CInt

foreign import ccall "sqproblem_hotstart"
    sqproblem_hotstart
        :: Ptr SQProblem -> Ptr Double -> Ptr Double -> Ptr Double -> Ptr Double
        -> Ptr Double -> Ptr Double -> Ptr Double -> Ptr CInt -> Ptr Double
        -> Ptr Double -> Ptr Double -> Ptr Double -> Ptr CInt -> IO CInt

foreign import ccall "sqproblem_sparse_hotstart"
    sqproblem_sparse_hotstart
        :: Ptr SQProblem -> Ptr CInt -> Ptr CInt -> Ptr Double
        -> Ptr Double -> Ptr Double -> Ptr Double
        -> Ptr Double -> Ptr Double -> Ptr Double -> Ptr CInt -> Ptr Double
        -> Ptr Double -> Ptr Double -> Ptr Double -> Ptr CInt -> IO CInt

foreign import ccall "sqproblem_sparse_schur_hotstart"
    sqproblem_sparse_schur_hotstart
        :: Ptr SQProblemSchur -> Ptr CInt -> Ptr CInt -> Ptr Double
        -> Ptr Double -> Ptr Double -> Ptr Double
        -> Ptr Double -> Ptr Double -> Ptr Double -> Ptr CInt -> Ptr Double
        -> Ptr Double -> Ptr Double -> Ptr Double -> Ptr CInt -> IO CInt

foreign import ccall "sqproblem_sparse_sparse_schur_hotstart"
    sqproblem_sparse_sparse_schur_hotstart
        :: Ptr SQProblemSchur -> Ptr CInt -> Ptr CInt -> Ptr Double
        -> Ptr Double -> Ptr CInt -> Ptr CInt -> Ptr Double -> Ptr Double
        -> Ptr Double -> Ptr Double -> Ptr Double -> Ptr CInt -> Ptr Double
        -> Ptr Double -> Ptr Double -> Ptr Double -> Ptr CInt -> IO CInt

foreign import ccall "sqproblem_setup"
    sqproblem_setup
        :: CInt -> CInt -> CInt -> Ptr (Ptr SQProblem)-> IO CInt

foreign import ccall "sqproblem_schur_setup"
    sqproblem_schur_setup
        :: CInt -> CInt -> CInt -> Ptr (Ptr SQProblemSchur)-> IO CInt

foreign import ccall "&sqproblem_cleanup"
    sqproblem_cleanup
        :: FunPtr (Ptr SQProblem -> IO ())

foreign import ccall "&sqproblem_schur_cleanup"
    sqproblem_schur_cleanup
        :: FunPtr (Ptr SQProblemSchur -> IO ())