-- TODO: better naming conventions

{-# LANGUAGE TypeFamilies, TypeOperators #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE DataKinds, PolyKinds, ConstraintKinds #-}
{-# LANGUAGE FlexibleInstances, FlexibleContexts, UndecidableInstances #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE PartialTypeSignatures #-}

module GrammarNNDAG where

import Control.Monad (foldM)
import Data.Function (fix)

import Generics.SOP hiding (SingI)
import Generics.SOP.NP
import Generics.SOP.NS
import Generics.SOP.Constraint
import Generics.SOP.Dict
import Data.Vinyl hiding (Dict)
import Data.Singletons.Prelude hiding (All, And)
import Data.Singletons.Prelude.List (Length)

import Utils
import List
import Grammar
import TensorHMatrix
import DAGIO
import TensorDAG
import VinylUtils
import Random

import Data.Default
import DefaultM

import Numeric.LinearAlgebra (Numeric)
import qualified Data.Vector.Storable as Vector

type family Size (s :: k -> *) t :: k
type IntegralS (s :: k -> *) = IntegralK ('KProxy :: KProxy k)

class SingI (Size s t) => KnownSize s t
instance SingI (Size s t) => KnownSize s t

class All2 (KnownSize s) (Code t) => KnownSizes s t
instance All2 (KnownSize s) (Code t) => KnownSizes s t

class SingI (FromInteger (Length (Code t)) :: k) => KnownCode (s :: k -> *) t
instance SingI (FromInteger (Length (Code t)) :: k) => KnownCode (s :: k -> *) t

class HasNodes t where
  getNodes :: t -> [Some Node]

newtype Repr a dim = Repr { runRepr :: Node (Tensor '[dim] a) }

instance HasNodes (Repr a dim) where
  getNodes (Repr node) = [Some node]

newtype ReprT a s t = ReprT { runReprT :: Node (Tensor '[Size s t] a) }

instance HasNodes (ReprT a s t) where
  getNodes (ReprT node) = [Some node]

instance (Default a, Usable a, IntegralS s, KnownSize s t) => DefaultM IO (ReprT a s t) where
  defM = ReprT <$> defM

--newtype Linear a outDim inDim = Linear (Node (Tensor '[outDim, inDim] a))
newtype LinearT a s tOut tIn = LinearT (Node (Tensor '[Size s tOut, Size s tIn] a))

instance HasNodes (LinearT a s tOut tIn) where
  getNodes (LinearT node) = [Some node]

--instance (Default a, Usable a, IntegralN outDim, SingI outDim, SingI inDim) => DefaultM IO (Linear a outDim inDim) where
--  defM = Linear <$> defM

instance (Default a, Usable a, IntegralS s, KnownSize s tIn, KnownSize s tOut) => DefaultM IO (LinearT a s tOut tIn) where
  defM = LinearT <$> defM

--linear :: Numeric a => Linear a outDim inDim -> Repr a inDim -> IO (Node (Tensor a '[outDim]))
--linear (Linear m) (Repr v) = makeMV m v

linearT :: (SingI (Size s tOut), IntegralS s, Usable a) => LinearT a s tOut tIn -> ReprT a s tIn -> IO (ReprT a s tOut)
linearT (LinearT m) (ReprT v) = ReprT <$> makeMV m v

data EncodeParams' a s parent children = EncodeParams' (ReprT a s parent) (NP (LinearT a s parent) children)

instance HasNodes (EncodeParams' a s parent children) where
  getNodes (EncodeParams' bias weights) = getNodes bias ++ concat (collapse_NP $ liftA_NP' (K . getNodes) weights)

instance (Default a, Usable a, IntegralS s, KnownSize s parent, SListI children, All (KnownSize s) children) => DefaultM IO (EncodeParams' a s parent children) where
  defM = EncodeParams' <$> defM <*> sequence'_NP (cpure_NP (Proxy::Proxy (KnownSize s)) (Comp defM))

encodeParent' :: (Usable a, Floating a, Floating (Vector.Vector a), IntegralS s, KnownSize s parent) =>
  EncodeParams' a s parent children -> NP (ReprT a s) children -> IO (ReprT a s parent)
encodeParent' (EncodeParams' (ReprT bias) weights) inputs = do
  let linearT' m v = Comp $ K <$> linearT m v
  children <- map runReprT . collapse_NP <$> (sequence'_NP $ liftA2_NP' linearT' weights inputs)
  parent <- foldM (makeBinary (+)) bias children
  ReprT <$> makeUnary tanh parent

newtype EncodeParams a s t = EncodeParams { runEncodeParams :: NP (EncodeParams' a s t) (Code t) }

instance HasNodes (EncodeParams a s t) where
  getNodes (EncodeParams params) = concat . collapse_NP $ liftA_NP' (K . getNodes) params

encodeParent :: (Floating a, Floating (Vector.Vector a), Usable a, IntegralS s, KnownSize s t) =>
  EncodeParams a s t -> SOP (ReprT a s) (Code t) -> IO (ReprT a s t)
encodeParent (EncodeParams params) (SOP sop) = collapse_NS $ liftA2_NS' encode params sop
  where encode p cs = K $ encodeParent' p cs

data Encoding a s t where
  Primitive :: ReprT a s t -> Encoding a s t
  Generic :: ReprT a s t -> SOP (Encoding a s) (Code t) -> Encoding a s t

getRepr :: Encoding a s t -> ReprT a s t
getRepr (Primitive repr) = repr
getRepr (Generic repr _) = repr

instance (Default a, Usable a, IntegralS s, KnownSize s t, KnownSizes s t) => DefaultM IO (EncodeParams a s t) where
  defM = EncodeParams <$> sequence'_NP (cpure_NP (Proxy::Proxy (All (KnownSize s))) (Comp defM))

newtype Encoder a s t = Encoder { runEncoder :: t -> IO (Encoding a s t) }

encodeRec :: forall ts a s t. (Floating a, Floating (Vector.Vector a), Usable a, KnownSize s t, IntegralS s) =>
  Rec (Encoder a s) ts -> Dict (Contained ts) t -> EncodeParams a s t -> Encoder a s t

encodeRec encoders Dict params = Encoder f where
  encoders' = cpure_POP (Proxy::Proxy (Find ts)) (Fn $ Comp . runEncoder (index find encoders) . unI)
  f t = do
    children <- sequence_SOP' (ap_SOP encoders' (from t))
    let childReprs = liftA_SOP' getRepr children
    parent <- encodeParent params childReprs
    return $ Generic parent children

makeEncoders :: forall p g a s. (Floating a, Floating (Vector.Vector a), Usable a, IntegralS s, All (KnownSize s) g) =>
  Complete p g -> NP (EncodeParams a s) g -> Rec (Encoder a s) p -> Rec (Encoder a s) (p :++ g)
makeEncoders complete params prim = fix f where
  f encoders = rAppend prim $ np2rec (cliftA2_NP (Proxy::Proxy (KnownSize s)) (encodeRec encoders) (unAll_NP complete) params)

newtype AnyEncoder a s ts = AnyEncoder { runAnyEncoder :: forall t. Find ts t => t -> IO (Encoding a s t) }

makeEncoder :: (Floating a, Floating (Vector.Vector a), Usable a, IntegralS s, All (KnownSize s) g) =>
  Complete p g -> NP (EncodeParams a s) g -> Rec (Encoder a s) p -> AnyEncoder a s (p :++ g)
makeEncoder complete params prim = AnyEncoder $ runEncoder (index find encoders)
  where encoders = makeEncoders complete params prim

newtype LinearIn a s t t' = LinearIn (Node (Tensor '[Size s t', Size s t] a))

instance HasNodes (LinearIn a s t t') where
  getNodes (LinearIn node) = [Some node]

instance (Default a, Usable a, IntegralS s, KnownSize s t, KnownSize s t') => DefaultM IO (LinearIn a s t t') where
  defM = LinearIn <$> defM

linearIn :: (IntegralS s, KnownSize s t', Usable a) => LinearIn a s t t' -> ReprT a s t -> IO (ReprT a s t')
linearIn (LinearIn m) (ReprT v) = ReprT <$> makeMV m v

data AffineIn a s t t' = AffineIn (ReprT a s t') (LinearIn a s t t')

instance HasNodes (AffineIn a s t t') where
  getNodes (AffineIn bias linear) = getNodes bias ++ getNodes linear

instance (Default a, Usable a, IntegralS s, KnownSize s t, KnownSize s t') => DefaultM IO (AffineIn a s t t') where
  defM = AffineIn <$> defM <*> defM

affineIn :: (Floating a, Floating (Vector.Vector a), Usable a, IntegralS s, KnownSize s t') => AffineIn a s t t' -> ReprT a s t -> IO (ReprT a s t')
affineIn (AffineIn (ReprT bias) weights) input = do
  ReprT l <- linearIn weights input
  b <- makeBinary (+) bias l
  t <- makeUnary tanh b
  return $ ReprT t

data Affine' a outDim inDim =
  Affine' (Node (Tensor '[outDim] a)) (Node (Tensor '[outDim, inDim] a))

instance HasNodes (Affine' a outDim inDim) where
  getNodes (Affine' bias linear) = [Some bias, Some linear]

instance (Default a, Usable a, IntegralN outDim, SingI outDim, SingI inDim) => DefaultM IO (Affine' a outDim inDim) where
  defM = Affine' <$> defM <*> defM

sigmoid x = 1 / (1 + exp (-x))

affine' :: (Floating a, Floating (Vector.Vector a), Usable a, IntegralN outDim, SingI outDim) => Affine' a outDim inDim -> Repr a inDim -> IO (Repr a outDim)
affine' (Affine' bias weight) (Repr input) = do
  l <- makeMV weight input
  b <- makeBinary (+) bias l
  Repr <$> makeUnary sigmoid b

data DecodeParams a s t =
  DecodeParams (Affine' a (Length (Code t)) (Size s t)) (POP (AffineIn a s t) (Code t))

instance HasNodes (DecodeParams a s t) where
  getNodes (DecodeParams aff params) = getNodes aff ++ concat (collapse_POP' $ liftA_POP' (K . getNodes) params)

instance (Default a, Usable a, KnownCode s t, KnownSize s t, KnownSizes s t) => DefaultM IO (DecodeParams a s t) where
  defM = DecodeParams <$> defM <*> sequence'_POP (cpure_POP (Proxy::Proxy (KnownSize s)) (Comp defM))

decodeParent :: forall a s t. (Real a, Floating a, Floating (Vector.Vector a), Usable a, All2 (KnownSize s) (Code t), KnownCode s t) =>
  DecodeParams a s t -> ReprT a s t -> IO (SOP (ReprT a s) (Code t))
decodeParent (DecodeParams aff params) parent = do
  Repr node <- affine' aff (Repr $ runReprT parent)
  Vector v <- evalNode node
  let weights = map toRational $ Vector.toList v

  let children = np2ns . unPOP $ cliftA_POP (Proxy::Proxy (KnownSize s)) (Comp . flip affineIn parent) params

  child <- sample $ zip children weights
  sequence'_SOP $ SOP child

newtype Decoder a s t = Decoder { runDecoder :: ReprT a s t -> IO t }

decodeRec :: forall ts a s t. (Real a, Floating a, Floating (Vector.Vector a), Usable a, KnownSizes s t, KnownCode s t) =>
  Rec (Decoder a s) ts -> Dict (Contained ts) t -> DecodeParams a s t -> Decoder a s t

decodeRec decoders Dict params = Decoder f where
  decoders' = cpure_POP (Proxy::Proxy (Find ts)) (Fn $ Comp . (I <$>) . runDecoder (index find decoders))
  f repr = do
    childReprs <- decodeParent params repr
    children <- sequence_SOP' (ap_SOP decoders' childReprs)
    return $ to children

makeDecoders :: forall p g a s. (Real a, Floating a, Floating (Vector.Vector a), Usable a, All (And (KnownSizes s) (KnownCode s)) g) =>
  Complete p g -> NP (DecodeParams a s) g -> Rec (Decoder a s) p -> Rec (Decoder a s) (p :++ g)
makeDecoders complete params prim = fix f where
  f decoders = rAppend prim $ np2rec (cliftA2_NP (Proxy::Proxy (And (KnownSizes s) (KnownCode s))) (decodeRec decoders) (unAll_NP complete) params)

newtype AnyDecoder a s ts = AnyDecoder { runAnyDecoder :: forall t. Find ts t => ReprT a s t -> IO t }

makeDecoder :: (Real a, Floating a, Floating (Vector.Vector a), Usable a, All (And (KnownSizes s) (KnownCode s)) g) =>
  Complete p g -> NP (DecodeParams a s) g -> Rec (Decoder a s) p -> AnyDecoder a s (p :++ g)
makeDecoder complete params prim = AnyDecoder $ runDecoder (index find decoders)
  where decoders = makeDecoders complete params prim

newtype AutoDecoder a s t = AutoDecoder { runAutoDecoder :: Encoding a s t -> IO (Node a) }

autoDecodeRec :: forall ts a s t. (Floating a, Floating (Vector.Vector a), Usable a, KnownCode s t, All2 (KnownSize s) (Code t)) =>
  Rec (AutoDecoder a s) ts -> Dict (Contained ts) t -> DecodeParams a s t -> AutoDecoder a s t

autoDecodeRec autoDecoders Dict (DecodeParams aff params) = AutoDecoder f where
  autoDecoders' = cpure_POP (Proxy::Proxy (Find ts)) (Fn $ Comp . (K <$>) . runAutoDecoder (index find autoDecoders))
  f :: Encoding a s t -> IO (Node a)
  f (Generic parent children) = do
    Repr node <- affine' aff (Repr $ runReprT parent)
    let i = ns2int (unSOP children)

    prob <- makeSelect i node
    log_score <- makeUnary (negate . log) prob

    let childReprs = liftA_SOP getRepr children
    let decodings = cliftA_POP (Proxy::Proxy (KnownSize s)) (Comp . flip affineIn parent) params

    let dist :: forall t'. KnownSize s t' => (IO :.: ReprT a s) t' -> ReprT a s t' -> (IO :.: K (Node a)) t'
        dist dIO (ReprT c) = Comp $ do
          ReprT d <- unComp dIO
          diff <- makeBinary (-) d c
          norm <- makeDot diff diff
          return $ K norm

    norms <- collapse_SOP <$> (sequence_SOP' $ cliftA2_SOP (Proxy::Proxy (KnownSize s)) dist decodings childReprs)

    childNodes <- collapse_SOP <$> sequence_SOP' (ap_SOP autoDecoders' children)

    foldM (makeBinary (+)) log_score (norms ++ childNodes)

primAutoDecoder :: Num a => AutoDecoder a s t
primAutoDecoder = AutoDecoder f where
  f (Primitive _) = makeSource 0

makeAutoDecoders :: forall proxy p g a s. (Real a, Floating a, Floating (Vector.Vector a), Usable a, All (And (KnownSizes s) (KnownCode s)) g) =>
  Complete p g -> NP (DecodeParams a s) g -> Rec (AutoDecoder a s) p -> Rec (AutoDecoder a s) (p :++ g)
makeAutoDecoders complete params prim = fix f where
  --prim = np2rec $ pure_NP primitiveAutoDecoder :: Rec _ p
  f autoDecoders = rAppend prim $ np2rec (cliftA2_NP (Proxy::Proxy (And (KnownSizes s) (KnownCode s))) (autoDecodeRec autoDecoders) (unAll_NP complete) params)

newtype AnyAutoDecoder a s ts = AnyAutoDecoder { runAnyAutoDecoder :: forall t. Find ts t => Encoding a s t -> IO (Node a) }

makeAutoDecoder :: (Real a, Floating a, Floating (Vector.Vector a), Usable a, All (And (KnownSizes s) (KnownCode s)) g) =>
  Complete p g -> NP (DecodeParams a s) g -> Rec (AutoDecoder a s) p -> AnyAutoDecoder a s (p :++ g)
makeAutoDecoder complete params prim = AnyAutoDecoder $ runAutoDecoder (index find autoDecoders)
  where autoDecoders = makeAutoDecoders complete params prim

newtype AutoEncoder a s t = AutoEncoder { runAutoEncoder :: t -> IO (Node a) }

makeAutoEncoders :: forall p g a s. (Real a, Floating a, Floating (Vector.Vector a), Usable a, All (KnownSize s) g, All (And (KnownSizes s) (KnownCode s)) g) =>
  Complete p g -> NP (EncodeParams a s) g -> Rec (Encoder a s) p -> NP (DecodeParams a s) g -> Rec (AutoDecoder a s) p -> Rec (AutoEncoder a s) (p :++ g)

makeAutoEncoders complete encodeParams primEncoders decodeParams primAutoDecoders = autoEncoders where
  encoders = makeEncoders complete encodeParams primEncoders
  autoDecoders = makeAutoDecoders complete decodeParams primAutoDecoders

  makeAutoEncoder (Encoder e) (AutoDecoder d) = AutoEncoder f where
    f t = e t >>= d

  autoEncoders = rZipWith makeAutoEncoder encoders autoDecoders

newtype AnyAutoEncoder a s ts = AnyAutoEncoder { runAnyAutoEncoder :: forall t. Find ts t => t -> IO (Node a) }

makeAutoEncoder :: (Real a, Floating a, Floating (Vector.Vector a), Usable a, All (KnownSize s) g, All (And (KnownSizes s) (KnownCode s)) g) =>
  Complete p g -> NP (EncodeParams a s) g -> Rec (Encoder a s) p -> NP (DecodeParams a s) g -> Rec (AutoDecoder a s) p -> AnyAutoEncoder a s (p :++ g)
makeAutoEncoder complete encodeParams primEncoders decodeParams primAutoDecoders = AnyAutoEncoder $ runAutoEncoder (index find autoEncoders)
  where autoEncoders = makeAutoEncoders complete encodeParams primEncoders decodeParams primAutoDecoders


{-# LANGUAGE ForeignFunctionInterface #-}

module Numeric.LinearAlgebra.HArpack
       ( eigs
       , eigsN
       , eigsS
       , linearOperator
       , SparseMatrix (..)
       , Problem (..)
       , SWhich (..)
       , NWhich (..)
       , Result (..)
       , Dimension
       , HowMany
       , Tolerance
       , Iterations ) where

import           Control.Monad
import           Control.Monad.Loops
import           Data.Complex
import           Data.Vector.Storable         (Vector)
import qualified Data.Vector.Storable         as V
import           Data.Vector.Storable.Mutable (IOVector)
import qualified Data.Vector.Storable.Mutable as MV
import           Foreign
import           Foreign.C.String
import           Foreign.C.Types
import           Foreign.ForeignPtr.Unsafe    as FPUnsafe
import           System.IO.Unsafe

-- | Functions of type 'LinearOperator' represent a matrix vector product.
-- The values in the second vector should be overwritten by operator applied to
-- the first.

data SparseMatrix =
  SparseMatrix { dim ::Int
               , indexes :: [((Int,Int),Double)] }
  deriving Show

data LinearOperator =
  LinearOperator
    { applyOperator :: IOVector CDouble -> IOVector CDouble -> IO () }

linearOperator :: [((Int,Int),Double)] -> LinearOperator
linearOperator l = LinearOperator $ \x y ->
  do MV.set y 0
     forM_ l $ \((i,j),v) -> do
       v' <- (* realToFrac v) <$> MV.read x j
       MV.modify y (+ v') i

--------------------------------------------------------------------------------
-- | Which eigenvalues to calculate:
-- - The first character refers to 'S'ymmetric or 'N'on-Symmetric;
-- - The second refers to 'L'argest or 'S'mallest;
-- - The third refers to 'M'agnitude, 'A'lgebraic, 'R'eal, 'I'mmaginary;
-- - The option 'BE' means 'BE'tween, for central eigenvalues.

data SWhich = SLM | SSM | SLA | SSA | SBE

instance Show SWhich where
  show SLM = "LM"
  show SSM = "SM"
  show SLA = "LA"
  show SSA = "SA"
  show SBE = "BE"

----------

data NWhich = NLM | NSM | NLR | NSR | NLI | NSI

instance Show NWhich where
  show NLM = "LM"
  show NSM = "SM"
  show NLR = "LR"
  show NSR = "SR"
  show NLI = "LI"
  show NSI = "SI"

----------

data Problem = Symmetric SWhich | NonSymmetric NWhich
     deriving (Show)

data Eigenpairs = Error String
                | EReal [(Double,Vector Double)]
                | EComplex [(Complex Double, Vector (Complex Double))]

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

-- | Encapsulates the arrays and data needed to invoke _naupd routines in ARPACK

data ArpackSetup =
  ArpackSetup { aIDO    :: Ptr CInt
              , aBMAT   :: Ptr CChar
              , aN      :: Ptr CInt
              , aWHICH  :: Ptr CChar
              , aNEV    :: Ptr CInt
              , aTOL    :: Ptr CDouble
              , aRESID  :: Ptr CDouble
              , aNCV    :: Ptr CInt
              , aV      :: Ptr CDouble
              , aLDV    :: Ptr CInt
              , aIPARAM :: Ptr CInt
              , aIPNTR  :: Ptr CInt
              , aWORKD  :: Ptr CDouble
              , aWORKL  :: Ptr CDouble
              , aLWORKL :: Ptr CInt
              , aINFO   :: Ptr CInt }
  deriving (Show)

type Dimension = Int
type HowMany = Int
type Tolerance = Double
type Iterations = Int

-- | The 'setupArpack' function allocates the arrays and data
-- needed by ARPACK to solve an eigenvalue problem

setupArpack :: Problem
            -> Dimension
            -> HowMany
            -> Tolerance
            -> Iterations
            -> IO ArpackSetup

setupArpack (NonSymmetric which') n nev tol maxIter =
  do let ncv = min n (2 * nev + 1)
         lWorkl = 3 * ncv * ncv + 6 * ncv
         ldv = n

     idoPtr <- ptrInt 0
     nPtr   <- ptrInt n
     nevPtr <- ptrInt nev
     ncvPtr <- ptrInt ncv
     ldvPtr <- ptrInt ldv

     bmatPtr  <- ptrCharArray "I"
     whichPtr <- ptrCharArray (show which')

     tolPtr   <- ptrDouble tol

     residPtr <- mallocDoubleArray n
     vPtr     <- mallocDoubleArray (n * ncv)
     workdPtr <- mallocDoubleArray (3 * n)
     worklPtr <- mallocDoubleArray lWorkl

     iParamPtr <- mallocIntArray 11

     pokeElemOff iParamPtr 0 1
     pokeElemOff iParamPtr 2 (fromIntegral maxIter)
     pokeElemOff iParamPtr 3 1
     pokeElemOff iParamPtr 6 1

     iPntrPtr  <- mallocIntArray 14

     lworklPtr <- ptrInt lWorkl
     infoPtr   <- ptrInt 0

     setup <- return $
              ArpackSetup idoPtr bmatPtr nPtr whichPtr nevPtr tolPtr
                          residPtr ncvPtr vPtr ldvPtr iParamPtr iPntrPtr
                          workdPtr worklPtr lworklPtr infoPtr

     dnaupd setup
     return setup

setupArpack (Symmetric which) n nev tol maxIter =
  do let ncv = min n (2 * nev + 1)
         lWorkl = ncv * ncv + 8 * ncv
         ldv = n

     idoPtr <- ptrInt 0
     nPtr   <- ptrInt n
     nevPtr <- ptrInt nev
     ncvPtr <- ptrInt ncv
     ldvPtr <- ptrInt ldv

     bmatPtr  <- ptrCharArray "I"
     whichPtr <- ptrCharArray (show which)

     tolPtr   <- ptrDouble tol

     residPtr <- mallocDoubleArray n
     vPtr     <- mallocDoubleArray (n * ncv)
     workdPtr <- mallocDoubleArray (3 * n)
     worklPtr <- mallocDoubleArray lWorkl

     iParamPtr <- mallocIntArray 11

     pokeElemOff iParamPtr 0 1
     pokeElemOff iParamPtr 2 (fromIntegral maxIter)
     pokeElemOff iParamPtr 3 1
     pokeElemOff iParamPtr 6 1

     iPntrPtr  <- mallocIntArray 11

     lworklPtr <- ptrInt lWorkl
     infoPtr   <- ptrInt 0

     setup <- return $
              ArpackSetup idoPtr bmatPtr nPtr whichPtr nevPtr tolPtr
                          residPtr ncvPtr vPtr ldvPtr iParamPtr iPntrPtr
                          workdPtr worklPtr lworklPtr infoPtr

     dsaupd setup

     return setup

-- |Invoke ARPACK's iteration through the reverse communication interface.
iterateArpack :: Problem -> LinearOperator -> ArpackSetup -> IO ()
iterateArpack (NonSymmetric _) f setup = do
  workdXElem <- peekElemOff (aIPNTR setup) 0
  workdYElem <- peekElemOff (aIPNTR setup) 1

  n <- peek (aN setup)

  xPtr <- newForeignPtr_ $ advancePtr (aWORKD setup) (fromIntegral workdXElem - 1)
  yPtr <- newForeignPtr_ $ advancePtr (aWORKD setup) (fromIntegral workdYElem - 1)

  let x = MV.unsafeFromForeignPtr xPtr 0 (fromIntegral n)
      y = MV.unsafeFromForeignPtr yPtr 0 (fromIntegral n)

  applyOperator f x y

  dnaupd setup

iterateArpack (Symmetric _) f setup = do
  workdXElem <- peekElemOff (aIPNTR setup) 0
  workdYElem <- peekElemOff (aIPNTR setup) 1

  n <- peek (aN setup)

  xPtr <- newForeignPtr_ $ advancePtr (aWORKD setup) (fromIntegral workdXElem - 1)
  yPtr <- newForeignPtr_ $ advancePtr (aWORKD setup) (fromIntegral workdYElem - 1)

  let x = MV.unsafeFromForeignPtr xPtr 0 (fromIntegral n)
      y = MV.unsafeFromForeignPtr yPtr 0 (fromIntegral n)

  applyOperator f x y

  dsaupd setup

arpack :: LinearOperator -> Problem -> Dimension -> HowMany -> Tolerance ->
          Iterations -> IO ArpackSetup
arpack op problem n nev tol mxItr = do
  ar <- setupArpack problem n nev tol mxItr
  iterateArpack problem op ar
  whileM_ ((== 1) <$> peek (aIDO ar)) (iterateArpack problem op ar)
  return ar

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

data D = DReal (IOVector CDouble)
       | DComplex (IOVector CDouble) (IOVector CDouble)

data Sigma = SReal (Ptr CDouble)
           | SComplex (Ptr CDouble) (Ptr CDouble)

data ArpackResults =
  ArpackResults { eRVEC    :: Ptr CInt
                , eHOWMNY  :: Ptr CChar
                , eSELECT  :: Ptr CInt
                , eD       :: D
                , eZ       :: IOVector CDouble
                , eLDZ     :: Ptr CInt
                , eSIGMA   :: Sigma
                , eWORKEV  :: Maybe (Ptr CDouble)
                , eSetup   :: ArpackSetup }

errCheck :: ArpackResults -> IO Bool
errCheck ar = do
  info <- (peek . aINFO . eSetup) ar
  print info
  return (info < 0)

parseArpackOutput :: Problem -> ArpackSetup -> IO ArpackResults
parseArpackOutput (NonSymmetric _) setup = do
  ncv <- peek (aNCV setup)
  nev <- peek (aNEV setup)
  n <- peek (aN setup)

  rvecPtr <- ptrInt 1
  howmnyPtr <- ptrCharArray "A"
  drVec <- newVectorDouble (nev + 1)
  diVec <- newVectorDouble (nev + 1)
  selectPtr <- mallocIntArray (fromIntegral ncv)
  zVec <- newVectorDouble (n*(nev+1))
  ldzPtr <- ptrInt (fromIntegral n)
  sigmarPtr <- mallocDouble
  sigmaiPtr <- mallocDouble
  workevPtr <- mallocDoubleArray (fromIntegral (3*ncv))

  let (drfPtr,_,_) = MV.unsafeToForeignPtr drVec
      drPtr = FPUnsafe.unsafeForeignPtrToPtr drfPtr
      (difPtr,_,_) = MV.unsafeToForeignPtr diVec
      diPtr = FPUnsafe.unsafeForeignPtrToPtr difPtr
      (zfPtr,_,_) = MV.unsafeToForeignPtr zVec
      zPtr = FPUnsafe.unsafeForeignPtrToPtr zfPtr

  c_dneupd rvecPtr howmnyPtr selectPtr drPtr diPtr zPtr ldzPtr sigmarPtr
    sigmaiPtr workevPtr (aBMAT setup) (aN setup) (aWHICH setup) (aNEV setup)
    (aTOL setup) (aRESID setup) (aNCV setup) (aV setup) (aLDV setup)
    (aIPARAM setup) (aIPNTR setup) (aWORKD setup) (aWORKL setup) (aLWORKL setup)
    (aINFO setup)

  return $ ArpackResults rvecPtr howmnyPtr selectPtr (DComplex drVec diVec) zVec
               ldzPtr (SComplex sigmarPtr sigmaiPtr) (Just workevPtr) setup

parseArpackOutput (Symmetric _) setup = do
  ncv <- peek (aNCV setup)
  nev <- peek (aNEV setup)
  n <- peek (aN setup)

  rvecPtr <- ptrInt 1
  howmnyPtr <- ptrCharArray "A"
  dVec <- newVectorDouble nev
  selectPtr <- mallocIntArray (fromIntegral ncv)
  zVec <- newVectorDouble (n*(nev+1))
  ldzPtr <- ptrInt (fromIntegral n)
  sigmaPtr <- mallocDouble

  let (drfPtr,_,_) = MV.unsafeToForeignPtr dVec
      dPtr = FPUnsafe.unsafeForeignPtrToPtr drfPtr
      (zfPtr,_,_) = MV.unsafeToForeignPtr zVec
      zPtr = FPUnsafe.unsafeForeignPtrToPtr zfPtr

  c_dseupd rvecPtr howmnyPtr selectPtr dPtr zPtr ldzPtr
    sigmaPtr (aBMAT setup) (aN setup) (aWHICH setup) (aNEV setup)
    (aTOL setup) (aRESID setup) (aNCV setup) (aV setup) (aLDV setup)
    (aIPARAM setup) (aIPNTR setup) (aWORKD setup) (aWORKL setup) (aLWORKL setup)
    (aINFO setup)

  return $ ArpackResults rvecPtr howmnyPtr selectPtr (DReal dVec) zVec
               ldzPtr (SReal sigmaPtr) Nothing setup

getComplexEigenvalue :: ArpackResults -> Int -> IO (Complex Double)
getComplexEigenvalue results index = do
  let DComplex reEigs imEigs = eD results
  re <- MV.read reEigs index
  im <- MV.read imEigs index
  return (realToFrac re :+ realToFrac im)

getRealEigenvector :: ArpackResults -> Int -> IO (Vector (Complex Double))
getRealEigenvector results index = do
  n <- (peek . aN . eSetup) results
  ldz <- peek $ eLDZ results
  let z = eZ results
  evec <- V.freeze $ MV.slice (index * fromIntegral ldz) (fromIntegral n) z
  return $ V.map ((:+0) . realToFrac) evec

getComplexEigenvector
  :: ArpackResults -> Int -> Bool -> IO (Vector (Complex Double))
getComplexEigenvector results index conjugate = do
  n <- (peek . aN . eSetup) results
  ldz <- peek $ eLDZ results
  let z = eZ results
  reEvec <- V.freeze $ MV.slice (index * fromIntegral ldz) (fromIntegral n) z
  imEvec <- V.freeze $ MV.slice ((index+1) * fromIntegral ldz) (fromIntegral n) z
  if conjugate
     then return $
          V.zipWith (:+) (V.map realToFrac reEvec)
                         (V.map (negate . realToFrac) imEvec)
     else return $
          V.zipWith (:+) (V.map realToFrac reEvec)
                         (V.map realToFrac imEvec)

getEigenpair
  :: ArpackResults -> Int -> IO (Complex Double, V.Vector (Complex Double))
getEigenpair results index = do
  eig <- getComplexEigenvalue results index
  let evec = case compare (imagPart eig) 0 of
        EQ -> getRealEigenvector results index
        LT -> getComplexEigenvector results (index - 1) True
        GT -> getComplexEigenvector results index False
  evec' <- evec
  return (eig,evec')

-- | The function 'eigs' computes a few eigenvalues of a real valued
-- unsymmetric matrix by invoking the ARPACK library. First element is true if
-- ARPACK converged. If converged the second element is a list of eigenpairs.

eigsN :: LinearOperator
      -> NWhich
      -> Dimension
      -> HowMany
      -> Tolerance
      -> Iterations
      -> Maybe [(Complex Double, Vector (Complex Double))]
eigsN f which' n' nev' tol' iters = unsafePerformIO $ do
  setup <- arpack f (NonSymmetric which') n' nev' tol' iters
  results <- parseArpackOutput (NonSymmetric which') setup
  err <- errCheck results
  if err
    then return Nothing
    else do
      pairs <- mapM (getEigenpair results) [0..(nev'-1)]
      return (Just pairs)

data Result = RError
            | RReal [(Double,Vector Double)]
            | RComplex [(Complex Double, Vector (Complex Double))]
              deriving Show

eigs :: SparseMatrix
     -> Problem
     -> HowMany
     -> Result
eigs m (NonSymmetric w) h =
  case eigsN (linearOperator (indexes m)) w (dim m) h (1e-10) 6000 of
  Nothing -> RError
  Just x -> RComplex x
eigs m (Symmetric w) h =
  case eigsS (linearOperator (indexes m)) w (dim m) h (1e-10) 6000 of
  Nothing -> RError
  Just x -> RReal x

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

getSymmetricEigenvalue :: ArpackResults -> Int -> IO Double
getSymmetricEigenvalue results index = do
  let DReal eigV = eD results
  eig <- MV.read eigV index
  return (realToFrac eig)

getSymmetricEigenvector :: ArpackResults -> Int -> IO (Vector Double)
getSymmetricEigenvector results index = do
  n <- fromIntegral <$> (peek . aN . eSetup) results
  ldz <- peek $ eLDZ results
  let z = eZ results
  evec <- V.freeze $ MV.slice (index * fromIntegral ldz) n z
  return $ V.map realToFrac evec

getSymmetricEigenpair :: ArpackResults -> Int -> IO (Double, Vector Double)
getSymmetricEigenpair results index = do
  eig <- getSymmetricEigenvalue results index
  evec <- getSymmetricEigenvector results index
  return (eig,evec)

eigsS :: LinearOperator
      -> SWhich
      -> Dimension
      -> HowMany
      -> Tolerance
      -> Iterations
      -> Maybe [(Double,Vector Double)]
eigsS f which n nev tol' iters = unsafePerformIO $ do
  setup <- arpack f (Symmetric which) n nev tol' iters
  results <- parseArpackOutput (Symmetric which) setup
  err <- errCheck results
  if err
    then return Nothing
    else do
      pairs <- mapM (getSymmetricEigenpair results) [0..(nev-1)]
      return (Just pairs)

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

mallocInt :: IO (Ptr CInt)
mallocInt = malloc

ptrInt :: Int -> IO (Ptr CInt)
ptrInt k =
  do ptr <- malloc
     poke ptr (fromIntegral k)
     return ptr

mallocIntArray :: Int -> IO (Ptr CInt)
mallocIntArray = mallocArray

----------

mallocDouble :: IO (Ptr CDouble)
mallocDouble = malloc

ptrDouble :: Double -> IO (Ptr CDouble)
ptrDouble x =
  do ptr <- malloc
     poke ptr (realToFrac x)
     return ptr

mallocDoubleArray :: Int -> IO (Ptr CDouble)
mallocDoubleArray = mallocArray

newVectorDouble :: CInt -> IO (MV.IOVector CDouble)
newVectorDouble = MV.new . fromIntegral

----------

mallocChar :: IO (Ptr CChar)
mallocChar = malloc

mallocCharArray :: Int -> IO (Ptr CChar)
mallocCharArray = mallocArray

ptrCharArray :: String -> IO (Ptr CChar)
ptrCharArray s =
  do ptr <- mallocCharArray (length s + 1)
     pokeArray ptr (map castCharToCChar s)
     return ptr

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

foreign import ccall unsafe "arpack.h dnaupd_"
  c_dnaupd :: Ptr CInt    -- ^ reverse communication flag, to be started with 0
           -> Ptr CChar   -- ^ 'I' for standard problem, 'G' for generalized
           -> Ptr CInt    -- ^ dimension of the eigenproblem
           -> Ptr CChar   -- ^ which eigenpairs to compute
           -> Ptr CInt    -- ^ N : number of eigenpairs to compute
           -> Ptr CDouble -- ^ tolerance
           -> Ptr CDouble -- ^ residual vector
           -> Ptr CInt    -- ^ NCV : number of columns of V matrix
           -> Ptr CDouble -- ^ N x NCV output array
           -> Ptr CInt    -- ^ leading dimension of V
           -> Ptr CInt    -- ^ array for internal parameters, length 11
           -> Ptr CInt    -- ^ pointer of starting locations, length 14
           -> Ptr CDouble -- ^ workd : work array
           -> Ptr CDouble -- ^ workl : work array
           -> Ptr CInt    -- ^ length of workl
           -> Ptr CInt    -- ^ info
           -> IO ()

dnaupd :: ArpackSetup -> IO ()
dnaupd setup =
  c_dnaupd (aIDO setup) (aBMAT setup) (aN setup) (aWHICH setup) (aNEV setup)
           (aTOL setup) (aRESID setup) (aNCV setup) (aV setup) (aLDV setup)
           (aIPARAM setup) (aIPNTR setup) (aWORKD setup) (aWORKL setup)
           (aLWORKL setup) (aINFO setup)

foreign import ccall unsafe "arpack.h dneupd_"
  c_dneupd :: Ptr CInt    -- ^ whether to calculate eigenvectors
           -> Ptr CChar   -- ^ whether to calculate Ritz or Schur vectors
           -> Ptr CInt    -- ^ selects the vectors to be computed
           -> Ptr CDouble -- ^ real part of eigenvalues
           -> Ptr CDouble -- ^ imaginary part of eigenvalues
           -> Ptr CDouble -- ^ approximate eigenvectors
           -> Ptr CInt    -- ^ leading dimension of former array
           -> Ptr CDouble -- ^ real part of the shift
           -> Ptr CDouble -- ^ imaginary part of the shift
           -> Ptr CDouble -- ^ work array
           -> Ptr CChar -> Ptr CInt -> Ptr CChar -> Ptr CInt -> Ptr CDouble
           -> Ptr CDouble -> Ptr CInt -> Ptr CDouble -> Ptr CInt -> Ptr CInt
           -> Ptr CInt -> Ptr CDouble -> Ptr CDouble -> Ptr CInt -> Ptr CInt
           -> IO ()

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

foreign import ccall unsafe "arpack.h dsaupd_"
  c_dsaupd :: Ptr CInt    -- ^ reverse communication flag, to be started with 0
           -> Ptr CChar   -- ^ 'I' for standard problem, 'G' for generalized
           -> Ptr CInt    -- ^ dimension of the eigenproblem
           -> Ptr CChar   -- ^ which eigenpairs to compute
           -> Ptr CInt    -- ^ N : number of eigenpairs to compute
           -> Ptr CDouble -- ^ tolerance
           -> Ptr CDouble -- ^ residual vector
           -> Ptr CInt    -- ^ NCV : number of columns of V matrix
           -> Ptr CDouble -- ^ N x NCV output array
           -> Ptr CInt    -- ^ leading dimension of V
           -> Ptr CInt    -- ^ array for internal parameters, length 11
           -> Ptr CInt    -- ^ pointer of starting locations, length 14
           -> Ptr CDouble -- ^ workd : work array
           -> Ptr CDouble -- ^ workl : work array
           -> Ptr CInt    -- ^ length of workl
           -> Ptr CInt    -- ^ info
           -> IO ()

dsaupd :: ArpackSetup -> IO ()
dsaupd setup =
  c_dsaupd (aIDO setup) (aBMAT setup) (aN setup) (aWHICH setup) (aNEV setup)
           (aTOL setup) (aRESID setup) (aNCV setup) (aV setup) (aLDV setup)
           (aIPARAM setup) (aIPNTR setup) (aWORKD setup) (aWORKL setup)
           (aLWORKL setup) (aINFO setup)

foreign import ccall unsafe "arpack.h dseupd_"
  c_dseupd :: Ptr CInt    -- ^ whether to calculate eigenvectors
           -> Ptr CChar   -- ^ whether to calculate Ritz or Schur vectors
           -> Ptr CInt    -- ^ selects the vectors to be computed
           -> Ptr CDouble -- ^ eigenvalues
           -> Ptr CDouble -- ^ approximate eigenvectors
           -> Ptr CInt    -- ^ leading dimension of former array
           -> Ptr CDouble -- ^ shift
           -> Ptr CChar -> Ptr CInt -> Ptr CChar -> Ptr CInt -> Ptr CDouble
           -> Ptr CDouble -> Ptr CInt -> Ptr CDouble -> Ptr CInt -> Ptr CInt
           -> Ptr CInt -> Ptr CDouble -> Ptr CDouble -> Ptr CInt -> Ptr CInt
           -> IO ()


{-# LANGUAGE DataKinds             #-}
{-# LANGUAGE ScopedTypeVariables   #-}
{-# LANGUAGE StandaloneDeriving    #-}
{-# LANGUAGE GADTs                 #-}
{-# LANGUAGE TypeOperators         #-}
{-# LANGUAGE TypeFamilies          #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE FlexibleContexts      #-}
{-# LANGUAGE DeriveGeneric #-}
{-# LANGUAGE DeriveAnyClass #-}
{-|
Module      : Grenade.Core.Pooling
Description : Max Pooling layer for 2D and 3D images
Copyright   : (c) Huw Campbell, 2016-2017
License     : BSD2
Stability   : experimental
-}
module Grenade.Layers.Pooling (
    Pooling (..)
  ) where

import           Data.Maybe
import           Data.Proxy
import           Data.Serialize
import           Data.Singletons.TypeLits hiding (natVal)
import           GHC.TypeLits
import           Control.DeepSeq (NFData)
import           GHC.Generics    (Generic)

import           Grenade.Core
import           Grenade.Layers.Internal.Pooling

import           Numeric.LinearAlgebra.Static as LAS hiding ((|||), build, toRows)

-- | A pooling layer for a neural network.
--
--   Does a max pooling, looking over a kernel similarly to the convolution network, but returning
--   maxarg only. This layer is often used to provide minor amounts of translational invariance.
--
--   The kernel size dictates which input and output sizes will "fit". Fitting the equation:
--   `out = (in - kernel) / stride + 1` for both dimensions.
--
data Pooling :: Nat -> Nat -> Nat -> Nat -> * where
  Pooling :: Pooling kernelRows kernelColumns strideRows strideColumns
  deriving (NFData, Generic)

instance Show (Pooling k k' s s') where
  show Pooling = "Pooling"

instance UpdateLayer (Pooling kernelRows kernelColumns strideRows strideColumns) where
  type Gradient (Pooling kernelRows kernelColumns strideRows strideColumns) = ()
  runUpdate _ Pooling _ = Pooling

instance RandomLayer (Pooling k k' s s') where
  createRandomWith _ _ = return Pooling

instance Serialize (Pooling kernelRows kernelColumns strideRows strideColumns) where
  put _ = return ()
  get = return Pooling

-- | A two dimentional image can be pooled.
instance ( KnownNat kernelRows
         , KnownNat kernelColumns
         , KnownNat strideRows
         , KnownNat strideColumns
         , KnownNat inputRows
         , KnownNat inputColumns
         , KnownNat outputRows
         , KnownNat outputColumns
         , ((outputRows - 1) * strideRows) ~ (inputRows - kernelRows)
         , ((outputColumns - 1) * strideColumns) ~ (inputColumns - kernelColumns)
         ) => Layer (Pooling kernelRows kernelColumns strideRows strideColumns) ('D2 inputRows inputColumns) ('D2 outputRows outputColumns) where
  type Tape (Pooling kernelRows kernelColumns strideRows strideColumns) ('D2 inputRows inputColumns) ('D2 outputRows outputColumns) = S ('D2 inputRows inputColumns)
  runForwards Pooling (S2D input) =
    let height = fromIntegral $ natVal (Proxy :: Proxy inputRows)
        width  = fromIntegral $ natVal (Proxy :: Proxy inputColumns)
        kx = fromIntegral $ natVal (Proxy :: Proxy kernelRows)
        ky = fromIntegral $ natVal (Proxy :: Proxy kernelColumns)
        sx = fromIntegral $ natVal (Proxy :: Proxy strideRows)
        sy = fromIntegral $ natVal (Proxy :: Proxy strideColumns)
        ex = extract input
        r  = poolForward 1 height width kx ky sx sy ex
        rs = fromJust . create $ r
    in  (S2D input, S2D rs)
  runBackwards Pooling (S2D input) (S2D dEdy) =
    let height = fromIntegral $ natVal (Proxy :: Proxy inputRows)
        width  = fromIntegral $ natVal (Proxy :: Proxy inputColumns)
        kx = fromIntegral $ natVal (Proxy :: Proxy kernelRows)
        ky = fromIntegral $ natVal (Proxy :: Proxy kernelColumns)
        sx = fromIntegral $ natVal (Proxy :: Proxy strideRows)
        sy = fromIntegral $ natVal (Proxy :: Proxy strideColumns)
        ex = extract input
        eo = extract dEdy
        vs = poolBackward 1 height width kx ky sx sy ex eo
    in  ((), S2D . fromJust . create $ vs)

-- | A three dimensional image can be pooled on each layer.
instance ( KnownNat kernelRows
         , KnownNat kernelColumns
         , KnownNat strideRows
         , KnownNat strideColumns
         , KnownNat inputRows
         , KnownNat inputColumns
         , KnownNat outputRows
         , KnownNat outputColumns
         , KnownNat channels
         , KnownNat (outputRows * channels)
         , ((outputRows - 1) * strideRows) ~ (inputRows - kernelRows)
         , ((outputColumns - 1) * strideColumns) ~ (inputColumns - kernelColumns)
         ) => Layer (Pooling kernelRows kernelColumns strideRows strideColumns) ('D3 inputRows inputColumns channels) ('D3 outputRows outputColumns channels) where
  type Tape (Pooling kernelRows kernelColumns strideRows strideColumns) ('D3 inputRows inputColumns channels) ('D3 outputRows outputColumns channels) = S ('D3 inputRows inputColumns channels)
  runForwards Pooling (S3D input) =
    let ix = fromIntegral $ natVal (Proxy :: Proxy inputRows)
        iy = fromIntegral $ natVal (Proxy :: Proxy inputColumns)
        kx = fromIntegral $ natVal (Proxy :: Proxy kernelRows)
        ky = fromIntegral $ natVal (Proxy :: Proxy kernelColumns)
        sx = fromIntegral $ natVal (Proxy :: Proxy strideRows)
        sy = fromIntegral $ natVal (Proxy :: Proxy strideColumns)
        ch = fromIntegral $ natVal (Proxy :: Proxy channels)
        ex = extract input
        r  = poolForward ch ix iy kx ky sx sy ex
        rs = fromJust . create $ r
    in  (S3D input, S3D rs)
  runBackwards Pooling (S3D input) (S3D dEdy) =
    let ix = fromIntegral $ natVal (Proxy :: Proxy inputRows)
        iy = fromIntegral $ natVal (Proxy :: Proxy inputColumns)
        kx = fromIntegral $ natVal (Proxy :: Proxy kernelRows)
        ky = fromIntegral $ natVal (Proxy :: Proxy kernelColumns)
        sx = fromIntegral $ natVal (Proxy :: Proxy strideRows)
        sy = fromIntegral $ natVal (Proxy :: Proxy strideColumns)
        ch = fromIntegral $ natVal (Proxy :: Proxy channels)
        ex = extract input
        eo = extract dEdy
        vs = poolBackward ch ix iy kx ky sx sy ex eo
    in  ((), S3D . fromJust . create $ vs)

-------------------- GNum instance --------------------

instance GNum (Pooling k k' s s') where
  _ |* _ = Pooling
  _ |+ _ = Pooling
  gFromRational _ = Pooling


{-# LANGUAGE DataKinds #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE FlexibleContexts #-}

module HVX.Primitives
-- TODO(mh): Currently primitives that should generate implicit constraints are
-- unsupported. The should be supported shortily. (2014-06-04)
  ( apply
  , hadd
  , (+~)
  , hmul
  , (*~)
  , habs
  , neg
  , hlog
  , hexp
  , logsumexp
  , hmax
  , hmin
  , norm
  , berhu
  , huber
  , quadform
  , powBaseP0
  , powBaseP01
  , powBaseP1
  , powBaseP1InfEven
  , powBaseP1InfNotInt
  ) where

import Numeric.LinearAlgebra
import Numeric.LinearAlgebra.HMatrix (eigSH')

import HVX.Internal.DCP
import HVX.Internal.Matrix
import HVX.Internal.Primitives
import HVX.Internal.Util

apply :: (ApplyVex vf mf ve me ~ vr, ValidVex vr)
  => Fun vf mf -> Expr ve me -> Expr vr (ApplyMon mf me)
apply = EFun

hadd :: (AddVex v1 v2 ~ v3, ValidVex v3)
  => Expr v1 m1 -> Expr v2 m2 -> Expr v3 (AddMon m1 m2)
hadd = EAdd

infixl 6 +~
(+~) :: (AddVex v1 v2 ~ v3, ValidVex v3)
  => Expr v1 m1 -> Expr v2 m2 -> Expr v3 (AddMon m1 m2)
(+~) = hadd

-- Constructors that enforce DCP constraints.
hmul :: (ApplyVex 'Affine 'Nonmon v1 m1 ~ v2, ValidVex v2)
  => Expr 'Affine 'Const -> Expr v1 m1 -> Expr v2 (ApplyMon 'Nonmon m1)
hmul (EConst a) e = apply (Mul a) e
hmul _ _ = error "the left argument of a multiply must be a constant"

infixl 7 *~
(*~) :: (ApplyVex 'Affine 'Nonmon v1 m1 ~ v2, ValidVex v2)
  => Expr 'Affine 'Const -> Expr v1 m1 -> Expr v2 (ApplyMon 'Nonmon m1)
(*~) (EConst a) e = apply (Mul a) e
(*~) _ _ = error "the left argument of a multiply must be a constant"

habs :: (ApplyVex 'Convex 'Nonmon v1 m1 ~ v2, ValidVex v2)
  => Expr v1 m1 -> Expr v2 (ApplyMon 'Nonmon m1)
habs = apply Abs

neg :: (ApplyVex 'Affine 'Noninc v1 m1 ~ v2, ValidVex v2)
  => Expr v1 m1 -> Expr v2 (ApplyMon 'Noninc m1)
neg = apply Neg

hlog :: (ApplyVex 'Concave 'Nondec v1 m1 ~ v2, ValidVex v2)
  => Expr v1 m1 -> Expr v2 (ApplyMon 'Nondec m1)
hlog = apply Log

hexp :: (ApplyVex 'Convex 'Nondec v1 m1 ~ v2, ValidVex v2)
  => Expr v1 m1 -> Expr v2 (ApplyMon 'Nondec m1)
hexp e = apply Exp e

logsumexp :: (ApplyVex 'Convex 'Nondec v1 m1 ~ v2, ValidVex v2)
  => Expr v1 m1 -> Expr v2 (ApplyMon 'Nondec m1)
logsumexp = apply LogSumExp

hmax :: (ApplyVex 'Convex 'Nondec v1 m1 ~ v2, ValidVex v2)
  => Expr v1 m1 -> Expr v2 (ApplyMon 'Nondec m1)
hmax = apply Max

hmin :: (ApplyVex 'Concave 'Nondec v1 m1 ~ v2, ValidVex v2)
  => Expr v1 m1 -> Expr v2 (ApplyMon 'Nondec m1)
hmin = apply Min

norm :: (ApplyVex 'Convex 'Nonmon v1 m1 ~ v2, ValidVex v2)
  => Double -> Expr v1 m1 -> Expr v2 (ApplyMon 'Nonmon m1)
norm p
  | p == infinity = error "Internal: Infinity norm should become max . abs."
  | 1 <= p = apply (Norm p)
  | otherwise = error "Internal: Norm only supports p >= 1."

berhu :: (ApplyVex 'Convex 'Nonmon v1 m1 ~ v2, ValidVex v2)
  => Double -> Expr v1 m1 -> Expr v2 (ApplyMon 'Nonmon m1)
berhu m
  | 0 < m = apply (Berhu m)
  | otherwise = error "Internal: Berhu only supports m >= 0."

huber :: (ApplyVex 'Convex 'Nonmon v1 m1 ~ v2, ValidVex v2)
  => Double -> Expr v1 m1 -> Expr v2 (ApplyMon 'Nonmon m1)
huber m
  | 0 < m = apply (Huber m)
  | otherwise = error "Internal: Huber only supports m >= 0."

quadform :: (ApplyVex 'Convex 'Nonmon 'Affine m1 ~ v2, ValidVex v2)
  => Expr 'Affine 'Const -> Expr 'Affine m1 -> Expr v2 (ApplyMon 'Nonmon m1)
quadform (EConst a) e
  | rows a == cols a
    && fpequalsMat a (tr a)
    && 0 <= (maxElement.fst.eigSH' $ a) = apply (Quadform a) e -- we have checked that the matrix being passed in is Hermitian, so it's safe to use eigSH'.
  | otherwise = error "Matrices in quadratic forms must be positive semidefinite."
quadform _ _ = error "The matrix sandwitched by the quadratic form must be constant."

powBaseP0 :: (ApplyVex 'Affine 'Const v1 m1 ~ v2, ValidVex v2)
  => Double -> Expr v1 m1 -> Expr v2 (ApplyMon 'Const m1)
powBaseP0 p
  | p `fpequals` 0 = apply PowBaseP0
  | otherwise = error "Internal: PowBaseP0 only supports p == 0."

powBaseP01 :: (ApplyVex 'Concave 'Nondec v1 m1 ~ v2, ValidVex v2)
  => Double -> Expr v1 m1 -> Expr v2 (ApplyMon 'Nondec m1)
powBaseP01 p
  | 0 < p && p < 1 = apply (PowBaseP01 p)
  | otherwise = error "Internal: PowBaseP01 only supports 0 < p < 1."

powBaseP1 :: (ApplyVex 'Affine 'Nondec v1 m1 ~ v2, ValidVex v2)
  => Double -> Expr v1 m1 -> Expr v2 (ApplyMon 'Nondec m1)
powBaseP1 p
  | p `fpequals` 1 = apply PowBaseP1
  | otherwise = error "Internal: PowBaseP1 only supports p == 1."

powBaseP1InfEven :: (ApplyVex 'Convex 'Nonmon v1 m1 ~ v2, ValidVex v2)
  => Double -> Expr v1 m1 -> Expr v2 (ApplyMon 'Nonmon m1)
powBaseP1InfEven p
  | 1 < p && isInteger p && even intP = apply (PowBaseP1InfEven intP)
  | otherwise = error "Internal: PowBaseP1InfEven only supports even p > 1."
  where intP = round p :: Integer

powBaseP1InfNotInt :: (ApplyVex 'Convex 'Nondec v1 m1 ~ v2, ValidVex v2)
  => Double -> Expr v1 m1 -> Expr v2 (ApplyMon 'Nondec m1)
powBaseP1InfNotInt p
  | 1 < p && not (isInteger p) = apply (PowBaseP1InfNotInt p)
  | otherwise = error "Internal: PowBaseP1InfNotInt only supports non integral p > 1."


{-# LANGUAGE CPP                  #-}
{-# LANGUAGE DeriveDataTypeable   #-}
{-# LANGUAGE FlexibleContexts     #-}
{-# LANGUAGE NoImplicitPrelude    #-}
{-# LANGUAGE ScopedTypeVariables  #-}
{-# LANGUAGE StandaloneDeriving   #-}
{-# LANGUAGE UndecidableInstances #-}

--------------------------------------------------------------------------------
-- |
-- Module:     Numeric.LevMar
-- Copyright:  (c) 2009 - 2014 Roel van Dijk & Bas van Dijk
-- License:    BSD-style (see the file LICENSE)
-- Maintainer: Roel van Dijk <vandijk.roel@gmail.com>
--             Bas van Dijk <v.dijk.bas@gmail.com>
-- Stability:  Experimental
--
-- For additional documentation see the
-- <http://www.ics.forth.gr/~lourakis/levmar/ documentation of the levmar C>
-- library which this library is based on:
--
--------------------------------------------------------------------------------

module Numeric.LevMar
    ( -- * Model & Jacobian.
      Params, Samples
    , Model, Jacobian

      -- * Levenberg-Marquardt algorithm.
    , LevMarable(levmar)

      -- * Minimization options.
    , Options(..), defaultOpts

      -- * Constraints
    , Constraints(..), LinearConstraints

      -- * Output
    , Info(..), StopReason(..), LevMarError(..)
    ) where

-------------------------------------------------------------------------------
-- Imports
-------------------------------------------------------------------------------

-- from base:
import Control.Monad         ( return, mplus )
import Control.Exception     ( Exception )
import Data.Bool             ( (&&), (||), otherwise )
import Data.Data             ( Data )
import Data.Typeable         ( Typeable )
import Data.Either           ( Either(Left, Right) )
import Data.Eq               ( Eq, (==), (/=) )
import Data.Function         ( (.), ($) )
import Data.Functor          ( (<$>) )
import Data.Int              ( Int )
import Data.List             ( lookup, (++) )
import Data.Maybe            ( Maybe(Nothing, Just), isJust, fromJust, fromMaybe )
import Data.Monoid           ( Monoid, mempty, mappend )
import Data.Ord              ( Ord, (<) )
import Foreign.C.Types       ( CInt )
import Foreign.Marshal.Array ( allocaArray, withArray, peekArray, copyArray )
import Foreign.Ptr           ( Ptr, nullPtr )
import Foreign.ForeignPtr    ( ForeignPtr, newForeignPtr_, withForeignPtr )
import Foreign.Storable      ( Storable )
import Prelude               ( Num, Enum, Fractional, RealFrac, Float, Double
                             , fromIntegral, toEnum, (-), (*), error, floor
                             )
import System.IO             ( IO )
import System.IO.Unsafe      ( unsafePerformIO )
import Text.Read             ( Read )
import Text.Show             ( Show, show )

#if __GLASGOW_HASKELL__ >= 605
import GHC.ForeignPtr        ( mallocPlainForeignPtrBytes )
import Prelude               ( undefined )
import Foreign.Storable      ( sizeOf )
#else
import Foreign.ForeignPtr    ( mallocForeignPtrArray )
#endif

#if __GLASGOW_HASKELL__ < 700
import Prelude               ( fromInteger, (>>=), (>>), fail )
#endif

-- from hmatrix:
#if MIN_VERSION_hmatrix(0,17,0)
import Numeric.LinearAlgebra.Data ( Matrix, flatten, rows, reshape )
import Numeric.LinearAlgebra      ( Container, Element )
#else
import Data.Packed.Matrix         ( Matrix, Element, flatten, rows, reshape )
import Numeric.Container          ( Container )
import Numeric.LinearAlgebra      ( {- Instances for Matrix -} )
#endif

-- from vector:
import           Data.Vector.Storable       ( Vector )
import qualified Data.Vector.Storable as VS ( unsafeWith, length
                                            , unsafeFromForeignPtr
                                            , length
                                            )

-- from bindings-levmar:
import Bindings.LevMar ( c'LM_INFO_SZ

                       , withModel
                       , withJacobian

                       , c'LM_ERROR
                       , c'LM_ERROR_LAPACK_ERROR
                       , c'LM_ERROR_FAILED_BOX_CHECK
                       , c'LM_ERROR_MEMORY_ALLOCATION_FAILURE
                       , c'LM_ERROR_CONSTRAINT_MATRIX_ROWS_GT_COLS
                       , c'LM_ERROR_CONSTRAINT_MATRIX_NOT_FULL_ROW_RANK
                       , c'LM_ERROR_TOO_FEW_MEASUREMENTS
                       , c'LM_ERROR_SINGULAR_MATRIX
                       , c'LM_ERROR_SUM_OF_SQUARES_NOT_FINITE

                       , c'LM_INIT_MU
                       , c'LM_STOP_THRESH
                       , c'LM_DIFF_DELTA
                       )
import qualified Bindings.LevMar ( Model, Jacobian )

-- from levmar (this package):
import Bindings.LevMar.CurryFriendly ( LevMarDer,     LevMarDif
                                     , LevMarBCDer,   LevMarBCDif
                                     , LevMarLecDer,  LevMarLecDif
                                     , LevMarBLecDer, LevMarBLecDif

                                     , dlevmar_der,      slevmar_der
                                     , dlevmar_dif,      slevmar_dif
                                     , dlevmar_bc_der,   slevmar_bc_der
                                     , dlevmar_bc_dif,   slevmar_bc_dif
                                     , dlevmar_lec_der,  slevmar_lec_der
                                     , dlevmar_lec_dif,  slevmar_lec_dif
                                     , dlevmar_blec_der, slevmar_blec_der
                                     , dlevmar_blec_dif, slevmar_blec_dif
                                     )

--------------------------------------------------------------------------------
-- Model & Jacobian.
--------------------------------------------------------------------------------

-- | Parameter vector of length @m@.
--
-- Ensure that @m <= n@ where @n@ is the length of the 'Samples' vector.
type Params r = Vector r

-- | Sample vector of length @n@.
--
-- Ensure that @n >= m@ where @m@ is the length of the 'Params' vector.
type Samples r = Vector r

{-| A functional relation describing measurements represented as a function
from a vector of parameters to a vector of expected samples.

 * Ensure that the length @m@ of the parameter vector equals the length of the
   initial parameter vector in 'levmar'.

 * Ensure that the length @n@ of the output sample vector equals the length of
   the sample vector in 'levmar'.

 * Ensure that the length @n@ of the output sample vector vector is bigger than or
   equal to the length @m@ of the parameter vector.
-}
type Model r = Params r -> Samples r

{-| The <http://en.wikipedia.org/wiki/Jacobian_matrix_and_determinant jacobian>
of the 'Model' function. Expressed as a function from a vector of
parameters to a matrix which for each expected sample describes the
partial derivatives of the parameters.

 * Ensure that the length @m@ of the parameter vector equals the length of the
   initial parameter vector in 'levmar'.

 * Ensure that the output matrix has the dimension @n><m@ where @n@ is the
   number of samples and @m@ is the number of parameters.
-}
type Jacobian r = Params r -> Matrix r

--------------------------------------------------------------------------------
-- Levenberg-Marquardt algorithm.
--------------------------------------------------------------------------------

-- | The Levenberg-Marquardt algorithm is overloaded to work on 'Double' and 'Float'.
class LevMarable r where

    -- | The Levenberg-Marquardt algorithm.
    --
    -- Returns a triple of the found parameters, a structure containing
    -- information about the minimization and the covariance matrix
    -- corresponding to LS solution.
    --
    -- Ensure that @n >= m@.
    levmar :: Model r            -- ^ Model
           -> Maybe (Jacobian r) -- ^ Optional jacobian
           -> Params r           -- ^ Initial parameters of length @m@
           -> Samples r          -- ^ Sample vector of length @n@
           -> Int                -- ^ Maximum iterations
           -> Options r          -- ^ Minimization options
           -> Constraints r      -- ^ Constraints
           -> Either LevMarError (Params r, Info r, Matrix r)

instance LevMarable Float where
    levmar = gen_levmar slevmar_der
                        slevmar_dif
                        slevmar_bc_der
                        slevmar_bc_dif
                        slevmar_lec_der
                        slevmar_lec_dif
                        slevmar_blec_der
                        slevmar_blec_dif

instance LevMarable Double where
    levmar = gen_levmar dlevmar_der
                        dlevmar_dif
                        dlevmar_bc_der
                        dlevmar_bc_dif
                        dlevmar_lec_der
                        dlevmar_lec_dif
                        dlevmar_blec_der
                        dlevmar_blec_dif

{-| @gen_levmar@ takes the low-level C functions as arguments and
executes one of them depending on the optional jacobian and constraints.

Preconditions:

@
  length ys >= length ps

     isJust mLowBs && length (fromJust mLowBs) == length ps
  && isJust mUpBs  && length (fromJust mUpBs)  == length ps

  boxConstrained && (all $ zipWith (<=) (fromJust mLowBs) (fromJust mUpBs))
@
-}
gen_levmar :: forall r. (RealFrac r, Element r)
           => LevMarDer r
           -> LevMarDif r
           -> LevMarBCDer r
           -> LevMarBCDif r
           -> LevMarLecDer r
           -> LevMarLecDif r
           -> LevMarBLecDer r
           -> LevMarBLecDif r

           -> Model r            -- ^ Model
           -> Maybe (Jacobian r) -- ^ Optional jacobian
           -> Params r           -- ^ Initial parameters
           -> Samples r          -- ^ Samples
           -> Int                -- ^ Maximum iterations
           -> Options r          -- ^ Options
           -> Constraints r      -- ^ Constraints
           -> Either LevMarError (Params r, Info r, Matrix r)
gen_levmar f_der
           f_dif
           f_bc_der
           f_bc_dif
           f_lec_der
           f_lec_dif
           f_blec_der
           f_blec_dif
           model mJac ps ys itMax opts (Constraints mLowBs mUpBs mWeights mLinC)
               | m == 0    = Left LevMarError -- LAPACK will crash otherwise!
               | otherwise =
  -- All effects are contained, so we can safely perform:
  unsafePerformIO $ do

    -- We need to pass the initial parameters 'ps' to the C function.
    -- However, we can't just pass a pointer to them because the C function
    -- will modify the parameters during execution which will violate
    -- referential transparanency. Instead we allocate new space
    -- and copy the parameters to it.
    --
    -- Note that, in the end, the array is returned from this function.
    -- This means that the only way to guarantee its finalisation
    -- is to allocate it using a ForeignPtr:
    psFP <- fastMallocForeignPtrArray m
    withForeignPtr psFP $ \psPtr -> do
      VS.unsafeWith ps $ \psPtrInp ->
        copyArray psPtr psPtrInp m

      -- Retrieve the (read-only) pointer 'ysPtr' to the samples vector 'ys'
      -- so we can pass it to the C function:
      VS.unsafeWith ys $ \ysPtr ->

        -- Convert the Options 'opts' to a list and then to an array
        -- so we can pass the (read-only) pointer 'optsPtr' to the C function:
        withArray (optsToList opts) $ \optsPtr ->

          -- Allocate space for the info array
          -- so we can pass it to the C function.
          -- Note that, in the end, this array is converted to an Info value
          -- and returned from this function.
          allocaArray c'LM_INFO_SZ $ \infoPtr -> do

            -- Allocate space for the covariance matrix
            -- so we can pass it to the C function.
            -- Like the parameters array the matrix
            -- needs to be returned from this function.
            -- So we also allocate it using a ForeignPtr:
            covarFP <- fastMallocForeignPtrArray mm
            withForeignPtr covarFP $ \covarPtr ->

              -- 'cmodel' is the low-level model function which is converted
              -- to the FunPtr 'modelFunPtr' and passed to the C function.
              -- 'cmodel' will first convert the parameters pointer 'parPtr'
              -- into a Vector after converting it into a ForeignPtr
              -- (without a finalizer).
              -- Then it will apply the high-level 'model' function
              -- to this parameter vector. The resulting vector is then copied
              -- to the output buffer 'hxPtr':
              let cmodel :: Bindings.LevMar.Model r
                  cmodel parPtr hxPtr _ _ _ = do
                    parFP <- newForeignPtr_ parPtr
                    let psV = VS.unsafeFromForeignPtr parFP 0 m
                        vector = model psV
                    VS.unsafeWith vector $ \p -> copyArray hxPtr p (VS.length vector)
              in withModel cmodel $ \modelFunPtr -> do

                 -- All the low-level C functions share a common set of arguments.
                 -- 'runDif' applies these arguments to the given C function 'f':
                 let runDif :: LevMarDif r -> IO CInt
                     runDif f = f modelFunPtr
                                  psPtr
                                  ysPtr
                                  (fromIntegral m)
                                  (fromIntegral n)
                                  (fromIntegral itMax)
                                  optsPtr
                                  infoPtr
                                  nullPtr
                                  covarPtr
                                  nullPtr

                 err <- case mJac of
                   Nothing -> if boxConstrained
                              then if linConstrained
                                   then withBoxConstraints
                                          (withLinConstraints $ withWeights runDif)
                                          f_blec_dif
                                   else withBoxConstraints runDif f_bc_dif
                              else if linConstrained
                                   then withLinConstraints runDif f_lec_dif
                                   else runDif f_dif

                   Just jac ->
                     let cjacobian :: Bindings.LevMar.Jacobian r
                         cjacobian parPtr jPtr _ _ _ = do
                           parFP <- newForeignPtr_ parPtr
                           let psV    = VS.unsafeFromForeignPtr parFP 0 m
                               matrix = jac psV
                               vector = flatten matrix
                           VS.unsafeWith vector $ \p ->
                             copyArray jPtr p (VS.length vector)
                     in withJacobian cjacobian $ \jacobPtr ->

                       let runDer :: LevMarDer r -> IO CInt
                           runDer f = runDif $ f jacobPtr
                       in if boxConstrained
                          then if linConstrained
                               then withBoxConstraints
                                      (withLinConstraints $ withWeights runDer)
                                      f_blec_der
                               else withBoxConstraints runDer f_bc_der
                          else if linConstrained
                               then withLinConstraints runDer f_lec_der
                               else runDer f_der

                 -- Handling errors:
                 if err < 0
                    -- we don't treat the following two as an error:
                    && err /= c'LM_ERROR_SINGULAR_MATRIX
                    && err /= c'LM_ERROR_SUM_OF_SQUARES_NOT_FINITE
                   then return $ Left $ convertLevMarError err

                   else do -- Converting results:
                     info <- listToInfo <$> peekArray c'LM_INFO_SZ infoPtr
                     let psV = VS.unsafeFromForeignPtr psFP 0 m
                     let covarM = reshape m $ VS.unsafeFromForeignPtr covarFP 0 mm

                     return $ Right (psV, info, covarM)
  where
    m  = VS.length ps
    n  = VS.length ys
    mm = m*m

    -- Whether the parameters are constrained by a linear equation.
    linConstrained = isJust   mLinC
    (cMat, rhcVec) = fromJust mLinC

    -- Whether the parameters are constrained by a bounding box.
    boxConstrained = isJust mLowBs || isJust mUpBs

    withBoxConstraints f g =
        maybeWithArray mLowBs $ \lBsPtr ->
          maybeWithArray mUpBs $ \uBsPtr ->
            f $ g lBsPtr uBsPtr

    withLinConstraints f g =
        VS.unsafeWith (flatten cMat) $ \cMatPtr ->
          VS.unsafeWith rhcVec $ \rhcVecPtr ->
            f . g cMatPtr rhcVecPtr $ fromIntegral $ rows cMat

    withWeights f g = maybeWithArray mWeights $ f . g

maybeWithArray :: (Storable a) => Maybe (Vector a) -> (Ptr a -> IO β) -> IO β
maybeWithArray Nothing  f = f nullPtr
maybeWithArray (Just v) f = VS.unsafeWith v f

#if __GLASGOW_HASKELL__ >= 605
{-# INLINE fastMallocForeignPtrArray #-}
fastMallocForeignPtrArray :: forall a. Storable a => Int -> IO (ForeignPtr a)
fastMallocForeignPtrArray n = mallocPlainForeignPtrBytes
                                (n * sizeOf (undefined :: a))
#else
fastMallocForeignPtrArray :: forall a. Storable a => Int -> IO (ForeignPtr a)
fastMallocForeignPtrArray = mallocForeignPtrArray
#endif

--------------------------------------------------------------------------------
-- Minimization options.
--------------------------------------------------------------------------------

-- | Minimization options
data Options r =
    Opts { optScaleInitMu      :: !r -- ^ Scale factor for initial /mu/.
         , optStopNormInfJacTe :: !r -- ^ Stopping thresholds for @||J^T e||_inf@.
         , optStopNorm2Dp      :: !r -- ^ Stopping thresholds for @||Dp||_2@.
         , optStopNorm2E       :: !r -- ^ Stopping thresholds for @||e||_2@.
         , optDelta            :: !r -- ^ Step used in the difference
                                     -- approximation to the Jacobian. If
                                     -- @optDelta<0@, the Jacobian is approximated
                                     -- with central differences which are more
                                     -- accurate (but slower!)  compared to the
                                     -- forward differences employed by default.
         } deriving (Eq, Ord, Read, Show, Data, Typeable)

-- | Default minimization options
defaultOpts :: Fractional r => Options r
defaultOpts = Opts { optScaleInitMu      = c'LM_INIT_MU
                   , optStopNormInfJacTe = c'LM_STOP_THRESH
                   , optStopNorm2Dp      = c'LM_STOP_THRESH
                   , optStopNorm2E       = c'LM_STOP_THRESH
                   , optDelta            = c'LM_DIFF_DELTA
                   }

optsToList :: Options r -> [r]
optsToList (Opts mu  eps1  eps2  eps3  delta) =
                [mu, eps1, eps2, eps3, delta]

--------------------------------------------------------------------------------
-- Constraints
--------------------------------------------------------------------------------

-- | Ensure that these vectors have the same length as the number of parameters.
data Constraints r = Constraints
    { lowerBounds       :: !(Maybe (Params r))            -- ^ Optional lower bounds
    , upperBounds       :: !(Maybe (Params r))            -- ^ Optional upper bounds
    , weights           :: !(Maybe (Params r))            -- ^ Optional weights
    , linearConstraints :: !(Maybe (LinearConstraints r)) -- ^ Optional linear constraints
    } deriving (Read, Show, Typeable)

deriving instance (Eq r, Container Vector r, Num r) => Eq (Constraints r)

-- | Linear constraints consisting of a constraints matrix, @k><m@ and
--   a right hand constraints vector, of length @k@ where @m@ is the number of
--   parameters and @k@ is the number of constraints.
type LinearConstraints r = (Matrix r, Vector r)

-- | * 'mempty' is defined as a 'Constraints' where all fields are 'Nothing'.
--
--   * 'mappend' merges two 'Constraints' by taking the first non-'Nothing' value
--     for each field.
instance Monoid (Constraints r) where
    mempty = Constraints Nothing Nothing Nothing Nothing
    mappend (Constraints lb1 ub1 w1 l1)
            (Constraints lb2 ub2 w2 l2) = Constraints (lb1 `mplus` lb2)
                                                      (ub1 `mplus` ub2)
                                                      (w1  `mplus` w2)
                                                      (l1  `mplus` l2)

--------------------------------------------------------------------------------
-- Output
--------------------------------------------------------------------------------

-- | Information regarding the minimization.
data Info r = Info
  { infNorm2initE      :: !r          -- ^ @||e||_2@             at initial parameters.
  , infNorm2E          :: !r          -- ^ @||e||_2@             at estimated parameters.
  , infNormInfJacTe    :: !r          -- ^ @||J^T e||_inf@       at estimated parameters.
  , infNorm2Dp         :: !r          -- ^ @||Dp||_2@            at estimated parameters.
  , infMuDivMax        :: !r          -- ^ @\mu/max[J^T J]_ii ]@ at estimated parameters.
  , infNumIter         :: !Int        -- ^ Number of iterations.
  , infStopReason      :: !StopReason -- ^ Reason for terminating.
  , infNumFuncEvals    :: !Int        -- ^ Number of function evaluations.
  , infNumJacobEvals   :: !Int        -- ^ Number of jacobian evaluations.
  , infNumLinSysSolved :: !Int        -- ^ Number of linear systems solved,
                                      --   i.e. attempts for reducing error.
  } deriving (Eq, Ord, Read, Show, Data, Typeable)

listToInfo :: (RealFrac r) => [r] -> Info r
listToInfo [a,b,c,d,e,f,g,h,i,j] =
    Info { infNorm2initE      = a
         , infNorm2E          = b
         , infNormInfJacTe    = c
         , infNorm2Dp         = d
         , infMuDivMax        = e
         , infNumIter         = floor f
         , infStopReason      = toEnum $ floor g - 1
         , infNumFuncEvals    = floor h
         , infNumJacobEvals   = floor i
         , infNumLinSysSolved = floor j
         }
listToInfo _ = error "liftToInfo: wrong list length"

-- | Reason for terminating.
data StopReason
  = SmallGradient  -- ^ Stopped because of small gradient @J^T e@.
  | SmallDp        -- ^ Stopped because of small Dp.
  | MaxIterations  -- ^ Stopped because maximum iterations was reached.
  | SingularMatrix -- ^ Stopped because of singular matrix. Restart from current
                   --   estimated parameters with increased 'optScaleInitMu'.
  | SmallestError  -- ^ Stopped because no further error reduction is
                   --   possible. Restart with increased 'optScaleInitMu'.
  | SmallNorm2E    -- ^ Stopped because of small @||e||_2@.
  | InvalidValues  -- ^ Stopped because model function returned invalid values
                   --   (i.e. NaN or Inf). This is a user error.
    deriving (Eq, Ord, Read, Show, Data, Typeable, Enum)

--------------------------------------------------------------------------------
-- Error
--------------------------------------------------------------------------------

data LevMarError
    = LevMarError                    -- ^ Generic error (not one of the others)
    | LapackError                    -- ^ A call to a lapack subroutine failed
                                     --   in the underlying C levmar library.
    | FailedBoxCheck                 -- ^ At least one lower bound exceeds the
                                     --   upper one.
    | MemoryAllocationFailure        -- ^ A call to @malloc@ failed in the
                                     --   underlying C levmar library.
    | ConstraintMatrixRowsGtCols     -- ^ The matrix of constraints cannot have
                                     --   more rows than columns.
    | ConstraintMatrixNotFullRowRank -- ^ Constraints matrix is not of full row
                                     --   rank.
    | TooFewMeasurements             -- ^ Cannot solve a problem with fewer
                                     --   measurements than unknowns.  In case
                                     --   linear constraints are provided, this
                                     --   error is also returned when the number
                                     --   of measurements is smaller than the
                                     --   number of unknowns minus the number of
                                     --   equality constraints.
      deriving (Eq, Ord, Read, Show, Data, Typeable)

-- Handy in case you want to thow a LevMarError as an exception:
instance Exception LevMarError

levmarCErrorToLevMarError :: [(CInt, LevMarError)]
levmarCErrorToLevMarError =
    [ (c'LM_ERROR,                                     LevMarError)
    , (c'LM_ERROR_LAPACK_ERROR,                        LapackError)
  --, (c'LM_ERROR_NO_JACOBIAN,                         can never happen)
  --, (c'LM_ERROR_NO_BOX_CONSTRAINTS,                  can never happen)
    , (c'LM_ERROR_FAILED_BOX_CHECK,                    FailedBoxCheck)
    , (c'LM_ERROR_MEMORY_ALLOCATION_FAILURE,           MemoryAllocationFailure)
    , (c'LM_ERROR_CONSTRAINT_MATRIX_ROWS_GT_COLS,      ConstraintMatrixRowsGtCols)
    , (c'LM_ERROR_CONSTRAINT_MATRIX_NOT_FULL_ROW_RANK, ConstraintMatrixNotFullRowRank)
    , (c'LM_ERROR_TOO_FEW_MEASUREMENTS,                TooFewMeasurements)
  --, (c'LM_ERROR_SINGULAR_MATRIX,                     we don't treat this as an error)
  --, (c'LM_ERROR_SUM_OF_SQUARES_NOT_FINITE,           we don't treat this as an error)
    ]

convertLevMarError :: CInt -> LevMarError
convertLevMarError err = fromMaybe (error $ "Unknown levmar error: " ++ show err)
                                   (lookup err levmarCErrorToLevMarError)


{-# LANGUAGE GeneralizedNewtypeDeriving, FlexibleInstances #-}

module SMC where

import PPL ( Sample(sample), Infer(score, observe, logscore), Model, runModel, sampleWeighted, WeightedT )
import Randomness ( Random, Weighted, Weight, parallel )
import Control.Monad ( ap, replicateM, zipWithM )
import Data.Either ( partitionEithers )
import Particle ( ParticleT(..), reweight )
import Control.Monad.Trans.Class (lift)
import Data.Monoid (Product(..))
import Control.Monad.Trans.Writer (WriterT, tell, runWriterT)
import Numeric.Log (Log(..))
import Distribution ( Dist (samplefrom), logpdf, Uniform (Uniform) )
import Statistics.Resampling (Bootstrap(resamples))
import Control.Monad.ST (ST)
import Control.Parallel (pseq)
import Data.List ( partition, sortOn, groupBy )
import Debug.Trace ( trace )
import Control.Monad.RWS (gets)

type Address = Int

instance (Monoid w, Sample m) => Sample (ParticleT l w m) where
  sample = Compute . fmap return . sample

instance Sample m => Infer (ParticleT l (Product Weight) m) where
  score w = Done (Product $ Exp $ log w) ()
  logscore w = Done (Product w) ()
  observe d x = Done (Product $ logpdf d x) ()

-- SMC with seperate score and resample statements
type SMCModel = ParticleT () (Product Weight) Random
type AddressSMC = ParticleT Address (Product Weight) Random

avg :: [Weighted Double] -> Double
avg xs = let total = sum (map snd xs) in
         exp $ ln $ sum [Exp (log a) * w / total | (a, w) <- xs]

stats :: Functor f => f [Weighted (Double, Double) ] -> f Double
stats = fmap (avg . map (\((a,b),c) -> (a,c)))

weight :: ParticleT l (Product Weight) m a -> Weight
weight (Done w _) = getProduct w
weight (Resample l w _) = getProduct w
weight (Compute _) = error "Compute has no weight"

-- A stochastic monad is one that can perform some sort of stochastic compute step
class Monad m => Stochastic m  where
  compute :: m a -> Random (m a)

instance Stochastic (ParticleT l (Product Weight) Random) where
  -- run a computation until it reaches a resample or is done
  compute (Compute ma) = ma >>= compute
  compute p = return p

class Stochastic m => Resampleable m where
  resample :: m ()

instance Resampleable SMCModel where
  resample = Resample () mempty $ return ()

class Infer m => SMC m where
  resampleSMC :: Int -> [m a] -> Random [Weighted a]
  normalisingConstant :: Int -> [m a] -> Random Weight

class Infer m => SampleWeighted m where
  doResample :: [m a] -> Random [m a]

-- SMC where every resample is explicit
newtype ExplicitResample a = ExplicitResample { unexplicit :: SMCModel a } deriving (Functor, Applicative, Monad, Sample, Resampleable, Infer, SMC)

-- SMC where every score triggers a resample
newtype AlwaysResample a = AlwaysResample { unalways :: SMCModel a } deriving (Functor, Applicative, Monad, Sample, Resampleable, SMC)

-- SMC where weight zero particles are resampled
newtype AliveFilter a = AliveFilter { unalive :: SMCModel a } deriving (Functor, Applicative, Monad, Sample, Resampleable, SMC)

-- SMC where subsets of particles are resampled based on addresses
newtype AlphaSMC a = AlphaSMC { unalpha :: AddressSMC a } deriving (Functor, Applicative, Monad, Sample, Infer)

-- SMC where subsets of particles are resampled based on addresses
newtype Adaptive a = Adaptive { unadapt :: SMCModel a } deriving (Functor, Applicative, Monad, Sample, Resampleable, Show)

-- SMC with Startified Resampling
newtype Stratified a = Stratified { unstrat :: SMCModel a } deriving (Functor, Applicative, Monad, Sample, Resampleable)

-- SMC with Subset Resampling
newtype Labelled a = Labelled { unlab :: AddressSMC a } deriving (Functor, Applicative, Monad, Sample, Infer)

-- these all have implicit resamples:
instance Infer AliveFilter where
  score w = AliveFilter (score w >> resample)
  observe d x = AliveFilter (observe d x >> resample)

instance Infer Adaptive where
  score w = Adaptive (score w >> resample)
  observe d x = Adaptive (observe d x >> resample)

instance Infer AlwaysResample where
  score w = AlwaysResample (score w >> resample)
  observe d x = AlwaysResample (observe d x >> resample)

instance Infer Stratified where
  score w = Stratified (score w >> resample)
  observe d x = Stratified (observe d x >> resample)

-- Generalized Newtype Deriving isn't quite enough to do these
instance Stochastic ExplicitResample where
  compute (ExplicitResample a) = ExplicitResample <$> compute a

instance Stochastic Labelled where
  compute (Labelled a) = Labelled <$> compute a

instance Stochastic AlwaysResample where
  compute (AlwaysResample a) = AlwaysResample <$> compute a

instance Stochastic Adaptive where
  compute (Adaptive a) = Adaptive <$> compute a

instance Stochastic Stratified where
  compute (Stratified a) = Stratified <$> compute a

resampleLab :: [Labelled a] -> Random [Labelled a]
resampleLab ms = do
    -- do all the stochastic compuations so everything
    -- is either done or at a resample statement
    ms' <- parallel $ map (compute . unlab) ms
    -- check if all particles are done...
    if all isDone ms'
      -- ... in which case return
      then return $ map Labelled ms'
      -- otherwise we do resampling
      else do
        -- we need to prune the resample we dealt with and reset weights
        resampled <- labResample $ map getweight ms'
        -- recursive call to continue inference
        resampleLab $ map Labelled resampled

labResample :: [Weighted (AddressSMC a)] -> Random [AddressSMC a]
labResample ms = let grouped = groupBy (\(x,_) (y,_) -> address x == address y) ms in
    fmap concat $ parallel $ map stratResample grouped

resampleStrat :: [Stratified a] -> Random [Stratified a]
resampleStrat ms = do
    -- do all the stochastic compuations so everything
    -- is either done or at a resample statement
    ms' <- parallel $ map (compute . unstrat) ms
    -- check if all particles are done...
    if all isDone ms'
      -- ... in which case return
      then return $ map Stratified ms'
      -- otherwise we do resampling
      else do
        -- we need to prune the resample we dealt with and reset weights
        resampled <- stratResample $ map getweight ms'
        -- recursive call to continue inference
        resampleStrat $ map Stratified resampled

getSamples :: [Weight] -> [Weighted a] -> [a]
getSamples [] _ = []
getSamples (r:rs) values = let (y:ys) = dropWhile (\(x,w) -> w < r) values
                           in fst y : getSamples rs (y:ys)

residual :: [Weighted a] -> ([a], [Weighted a])
residual [] = ([], [])
residual ((a, w) : rest) = let (xs, wxs) = residual rest
                           in let w' = floor w
                           in (replicate w' a ++ xs, (a, w - fromIntegral w') : wxs)

stratResample :: Functor m => [Weighted (ParticleT l (Product Weight) m a)] -> Random [ParticleT l (Product Weight) m a]
stratResample xs = do
    -- get sum of weights
    let weight = sum $ map snd xs
    let n = length xs
    let avgweight = weight / Exp (log $ fromIntegral n)
    -- residuals
    let (guaranteed, rest) = residual [(a, w / avgweight) | (a,w) <- xs]
    let n' = n - length guaranteed
    -- rest
    let weighted = let (as, ws) = unzip rest in zip as (scanl1 (+) ws)
    let avgrest = snd (last weighted) / Exp (log $ fromIntegral n')

    -- list of n uniform[i-1/n,i/n] random values
    let between a b = do
          r <- samplefrom (Uniform (fromIntegral a) (fromIntegral b))
          return $ avgrest * Exp (log r)

    rs <- zipWithM between [0..n'-1] [1..n']
    return $ map (reweight $ Product avgweight) (getSamples rs weighted ++ guaranteed)
{-
stratResample :: [Weighted (SMCModel a)] -> Random [SMCModel a]
stratResample xs = do
    -- [(a1, w1), (a2, w2), ...] becomes [(a1, w1), (a2, w1+w2), ...]
    let weighted = let (as, ws) = unzip xs in zip as (scanl1 (+) ws)
    -- get sum of weights
    let weight = snd $ last weighted
    let n = length weighted
    let avgweight = weight / Exp (log $ fromIntegral n)

    -- list of n uniform[i-1/n,i/n] random values
    let between a b = do
          r <- samplefrom (Uniform (fromIntegral a) (fromIntegral b))
          return $ avgweight * Exp (log r)
    rs <- zipWithM between [0..n-1] [1..n]

    -- list of k weighted samples
    let result = getSamples rs weighted
    return $ map (reweight $ Product avgweight) result
-}

class Infer m => Addressable m where
  resample' :: Address -> m ()

  score' :: Address -> Double -> m ()
  score' a w = score w >> resample' a

  observe' :: Dist d v => Address -> d -> v -> m ()
  observe' a d v = observe d v >> resample' a

-- subset SMC specific functions
instance Addressable Labelled where
  resample' a = Labelled $ Resample a mempty $ return ()

-- alpha-SMC specific functions
instance Addressable AlphaSMC where
  resample' a = AlphaSMC $ Resample a mempty $ return ()

-- get address from a Dome or Resample
address :: Bounded l => ParticleT l w m a -> l
address (Resample a _ _) = a
address (Done _ _) = maxBound
address (Compute _ ) = error "compute has no address"

-- helper function used by alpha-SMC
resampleMinSubset :: [AddressSMC a] -> Random [AddressSMC a]
resampleMinSubset ps = do
  let mini = minimum $ map address ps
  let (minima, rest) = partition ((== mini) . address) ps
  let w = sum $ map weight minima
  let n = length minima
  minima' <- sampleWeighted n $ map getweight minima
  let normalized = map (reweight $ Product (w / Exp (log $ fromIntegral n))) minima'
  return (normalized ++ rest)

resampleAlpha :: Int -> [AlphaSMC a] -> Random [AlphaSMC a]
resampleAlpha n ms = do
    -- do all the stochastic compuations so everything
    -- is either done or at a resample statement
    ms' <- parallel $ map (compute . unalpha) ms
    -- check if all particles are done...
    if all isDone ms'
      -- ... in which case return
      then return $ map AlphaSMC ms'
      -- otherwise we do resampling
      else do
        -- we need to prune the resample we dealt with and reset weights
        resampled <- resampleMinSubset ms'
        -- recursive call to continue inference
        resampleAlpha n $ map AlphaSMC resampled

resampleESS :: Int -> ([Adaptive a], Weight) -> Random ([Adaptive a], Weight)
resampleESS n (ms, w) = do
    -- do all the stochastic compuations so everything
    -- is either done or at a resample statement
    ms' <- parallel $ map (compute . unadapt) ms
    let w1 = sum (map weight ms')
    let w2 = sum (map ((^2) . weight) ms')
    let ess = w1 * w1 / w2
    let w' = w * (w1 / Exp (log $ fromIntegral n))

    -- check if all particles are done...
    if all isDone ms'
      -- ... in which case return
      then return (map Adaptive ms', w')
      -- otherwise we do resampling
      else do
        if ess < Exp (log (1 * fromIntegral n))
        then do
          resampled <- sampleWeighted n (map getweight ms')
          resampleESS n (map Adaptive resampled, w')
        else
          resampleESS n (map (Adaptive . ignoreResample) ms', w)

instance SMC AlphaSMC where
  resampleSMC n ms = do
    ms' <- resampleAlpha n ms
    return (map (\(AlphaSMC (Done w a)) -> (a, getProduct w)) ms')
  normalisingConstant n ms = do
    ms' <- resampleAlpha n ms
    let w = sum (map (\(AlphaSMC (Done w a)) -> getProduct w) ms')
    return $ w / Exp (log $ fromIntegral n)

instance SMC Labelled where
  resampleSMC n ms = do
    ms' <- resampleLab ms
    return (map (\(Labelled (Done w a)) -> (a, getProduct w)) ms')
  normalisingConstant n ms = do
    ms' <- resampleLab ms
    return $ getProduct (sum $ map (\(Labelled (Done w a)) -> w) ms') / Exp (log $ fromIntegral $ length ms')

instance SMC Adaptive where
  resampleSMC n ms = do
    (ms', w) <- resampleESS n (ms, Exp 0)
    return (map (\(Adaptive (Done w a)) -> (a, getProduct w)) ms')
  normalisingConstant n ms = snd <$> resampleESS n (ms, Exp 0)

instance SMC Stratified where
  resampleSMC n ms = do
    ms' <- resampleStrat ms
    return (map (\(Stratified (Done w a)) -> (a, getProduct w)) ms')
  normalisingConstant n ms = do
    ms' <- resampleStrat ms
    return $ getProduct (sum $ map (\(Stratified (Done w a)) -> w) ms') / Exp (log $ fromIntegral $ length ms')

ignoreResample :: ParticleT l (Product Weight) Random a -> ParticleT l (Product Weight) Random a
ignoreResample (Resample _ w a) = reweight w a
ignoreResample p@(Done _ _) = p
ignoreResample (Compute _) = error "unexpected compute"

instance Stochastic AliveFilter where
  -- keep trying until we get a nonzero weight
  compute p@(AliveFilter (Compute ma)) = do
    a <- compute (Compute ma)
    if weight a == Exp (log 0) -- note that log 0 = -Infinity
    then compute p
    else return $ AliveFilter a
  compute p = return p

instance Stochastic AlphaSMC where
  -- keep trying until we get a nonzero weight
  compute p@(AlphaSMC (Compute ma)) = do
    a <- compute (Compute ma)
    if weight a == Exp (log 0) -- note that log 0 = -Infinity
    then compute p
    else return $ AlphaSMC a
  compute p = return p

-- get weights, clear internal weights, remove resample
getweight :: ParticleT l (Product Weight) m a -> Weighted (ParticleT l (Product Weight) m a)
getweight (Done w a) = (Done mempty a, getProduct w)
getweight (Resample l w a) = (a, getProduct w)
getweight (Compute a) = error "unexpected compuation"

-- check if particle is done
isDone :: ParticleT l w m a -> Bool
isDone (Done _ _) = True
isDone _ = False

smcmodel :: Int -> ([SMCModel a], Weight) -> Random ([SMCModel a], Weight)
smcmodel n (ms, w) = do
  -- do all the stochastic compuations so everything
    -- is either done or at a resample statement
    ms' <- parallel $ map compute ms
    let w' = w * sum (map weight ms') / Exp (log $ fromIntegral n)
    -- check if all particles are done...
    if all isDone ms'
      -- ... in which case return
      then return (ms', w')
      -- otherwise we do resampling
      else do
        -- we need to prune the resample we dealt with and reset weights
        resampled <- sampleWeighted n (map getweight ms')
        -- recursive call to continue inference
        smcmodel n (resampled, w')

instance SMC SMCModel where
  resampleSMC n ms = do
    (ms', w) <- smcmodel n (ms, Exp 0)
    return (map (\(Done w a) -> (a, getProduct w)) ms')
  normalisingConstant n ms = snd <$> smcmodel n (ms, Exp 0)

inferSMC :: SMC m => Int -> m a -> Random [Weighted a]
inferSMC n m = resampleSMC n (replicate n m)

inferConstant :: SMC m => Int -> m a -> Random Weight
inferConstant n m = normalisingConstant n (replicate n m)