{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE UndecidableInstances #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE ConstraintKinds #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE TupleSections #-}
{-# LANGUAGE DeriveGeneric #-}

module Language.Grappa.Interp.PVIE where

import Data.List
import GHC.Generics hiding (R)
import Data.Functor.Const
import Data.Functor.Compose
-- import Data.Functor.Product
import Control.Lens
import Control.Monad.State
import Control.Monad.Reader
import Control.Monad.Primitive
import Data.IORef
import qualified Numeric.Log as Log

import qualified Data.ByteString.Lazy as BS
import Data.Aeson
import System.IO
import System.IO.Unsafe
import System.Directory (doesFileExist)

import Data.Vector (Vector)
import qualified Data.Vector as V
import qualified Numeric.LinearAlgebra.Data as SV (Vector)
import qualified Data.Vector.Generic as SV hiding (Vector)
import qualified Data.Vector.Generic as SVGen
import qualified Data.Vector.Storable.Mutable as SMV

import qualified Numeric.AD.Mode.Reverse as ADR
import qualified Numeric.AD.Internal.Reverse as ADR
import qualified Data.Reflection as ADR (Reifies)

import Numeric.GSL.Minimization

import Language.Grappa.Distribution
import Language.Grappa.GrappaInternals

import Language.Grappa.Rand.MWCRandM
import qualified System.Random.MWC as MWC
import qualified System.Random.MWC.Distributions as MWC

import qualified Text.PrettyPrint.ANSI.Leijen as PP
import Text.PrettyPrint.ANSI.Leijen ((<+>), (<>))

import Options.Applicative

import Debug.Trace

-- | Read JSON file @filename@ and parse it to an @a@, where @filename@ can be
-- @"-"@ to denote 'stdin'
readJSONfile :: FromJSON a => String -> IO a
readJSONfile f =
  withFileContents f $ \contents ->
  case eitherDecode' contents of
    Left err -> error ("Could not parse JSON file " ++ f ++ ": " ++ err)
    Right a -> return a
  where
    withFileContents :: String -> (BS.ByteString -> IO a) -> IO a
    withFileContents "-" k = BS.hGetContents stdin >>= k
    withFileContents filename k =
      withFile filename ReadMode $ \h -> BS.hGetContents h >>= k

-- | Write JSON file @filename@, which can be @"-"@ to indicate 'stdout'
writeJSONfile :: ToJSON a => String -> a -> IO ()
writeJSONfile "-" a = BS.putStr $ encode a
writeJSONfile filename a = encodeFile filename a

----------------------------------------------------------------------
-- * Dimensionality Expressions
----------------------------------------------------------------------

-- | A "dimensionality variable"
-- (FIXME: explain what this means: an index into a 'VIDimAsgn'?)
newtype VIDimVar = VIDimVar { viDimVarIx :: Int } deriving Eq

-- | Build a list of the first @n@ dimensionality variables
viDimVars :: Int -> [VIDimVar]
viDimVars n = take n $ map VIDimVar $ [0 ..]

-- | The "first" dimensionality variable
zeroVar :: VIDimVar
zeroVar = VIDimVar 0

-- | Get the next dimensionality variable after the given one
nextVar :: VIDimVar -> VIDimVar
nextVar (VIDimVar i) = (VIDimVar (i+1))

maxVar :: VIDimVar -> VIDimVar -> VIDimVar
maxVar (VIDimVar v1) (VIDimVar v2) = VIDimVar $ max v1 v2

-- | An assignment to dimensionality variables, which should be non-negative
newtype VIDimAsgn = VIDimAsgn [Int] deriving Generic

instance FromJSON VIDimAsgn where

instance ToJSON VIDimAsgn where
  toEncoding = genericToEncoding defaultOptions

-- | Create an assignment where all variables below the given one are set to 0
zeroAsgn :: VIDimVar -> VIDimAsgn
zeroAsgn (VIDimVar i) = VIDimAsgn $ take i $ repeat 1

-- | Increment the assignment to a particular dimensionality variable
incrAsgn :: VIDimVar -> VIDimAsgn -> VIDimAsgn
incrAsgn (VIDimVar i) (VIDimAsgn asgn) = VIDimAsgn (over (ix i) (+ 1) asgn)

-- | Increment the assignment to a particular dimensionality variable
incrAsgnAll :: VIDimAsgn -> VIDimAsgn
incrAsgnAll (VIDimAsgn asgn) = VIDimAsgn (map (+1) asgn)

-- | Look up the value of a dimensionality variable in an assignment
viDimVarVal :: VIDimVar -> VIDimAsgn -> Int
viDimVarVal (VIDimVar i) (VIDimAsgn ns) = ns !! i

-- | An expression for the dimensionality of variational inference parameters in
-- terms of some dimensionality variables, given as a function on the values of
-- those variables along with an indicator of how many variables are
-- needed. Note that this function should always return a non-negative value.
data VIDim = VIDim {
  evalVIDim :: VIDimAsgn -> Int,
  viDimFirstUnusedVar :: VIDimVar
  }

-- | Build a constant dimensionality
constVIDim :: Int -> VIDim
constVIDim k = VIDim (const k) zeroVar

-- | Build a variable dimensionality
varVIDim :: VIDimVar -> VIDim
varVIDim v = VIDim (viDimVarVal v) (nextVar v)

-- | Apply a binary operation to a 'VIDim', ensuring the resulting dimension
-- evaluation is non-negative
viDimBinOp :: (Int -> Int -> Int) -> VIDim -> VIDim -> VIDim
viDimBinOp binop (VIDim eval1 v1) (VIDim eval2 v2) =
  VIDim (\asgn -> max (eval1 asgn `binop` eval2 asgn) 0) (maxVar v1 v2)

instance Num VIDim where
  (+) = viDimBinOp (+)
  (-) = viDimBinOp (-)
  (*) = viDimBinOp (*)
  abs dim = dim -- dimensions are always non-negative
  signum _ = 1 -- satisfies abs dim * signum dim = dim (up to equivalence)
  fromInteger = constVIDim . fromInteger

----------------------------------------------------------------------
-- * Differentiable Functions
----------------------------------------------------------------------

type Params = SV.Vector Double
type MutParams = SMV.MVector (PrimState IO) Double
type ParamsGrad = SMV.MVector (PrimState IO) Double

-- | Create a normal 'Vector' for a (small number of) 'Params'. This is needed
-- in order to use the reverse-mode AD 'ADR.grad' function.
vectorOfParams :: Int -> Params -> Vector Double
vectorOfParams n params = V.generate n (params SV.!)

-- | Helper class to write "differentiable" functions on our parameters
class FromDouble a where
  fromDouble :: Double -> a

instance FromDouble Double where
  fromDouble = id

instance ADR.Reifies s ADR.Tape => FromDouble (ADR.Reverse s Double) where
  fromDouble = ADR.auto

-- | All the typeclasses we want for types used in differentiable functions
type DiffConstr a = (RealFloat a, Ord a, Show a, FromDouble a,
                     HasGamma a, Log.Precise a)

-- | The type of functions that are differentiable using AD
type DiffFun = forall r. DiffConstr r => Vector r -> r

-- | Take a differentiable function on @n@ (the first argument) variables and
-- turn it into a function that computes its gradient, scaled by the given
-- factor, and accumulates the result into the supplied mutable vector
scaledGrad :: Int -> Double -> DiffFun -> Params -> ParamsGrad -> IO ()
scaledGrad n scale f params params_grad =
  let raw_grad = ADR.grad f (vectorOfParams n params) in
  forM_ [0 .. n-1] $ \i ->
  SMV.modify params_grad ((scale * raw_grad V.! i) +) i

-- | Take a differentiable function on @n@ (the first argument) variables and
-- turn it into a function that computes its gradient, accumulating the result
-- into the supplied mutable vector, and also returns the result of the function
gradAndRet :: Int -> DiffFun -> Params -> ParamsGrad -> IO Double
gradAndRet n f params params_grad =
  do scaledGrad n 1 f params params_grad
     return (f $ vectorOfParams n params)

----------------------------------------------------------------------
-- * Computations Over Parameters
----------------------------------------------------------------------

newtype ParamsT m a = ParamsT (ReaderT (IORef Params) m a)
                    deriving (Functor, Applicative, Monad, MonadTrans)

runParamsT :: MonadIO m => ParamsT m a -> Params -> m a
runParamsT (ParamsT m) params =
  liftIO (newIORef params) >>= runReaderT m

instance MonadIO m => MonadIO (ParamsT m) where
  liftIO m = ParamsT $ liftIO m

class Monad m => MonadParams m where
  getParams :: Int -> m Params

instance MonadParams m => MonadParams (ReaderT r m) where
  getParams n = lift $ getParams n

instance MonadIO m => MonadParams (ParamsT m) where
  getParams n =
    ParamsT $ do r <- ask
                 ps <- liftIO $ readIORef r
                 liftIO $ writeIORef r (SV.drop n ps)
                 return $ SV.take n ps

newtype MutParamsT m a = MutParamsT (ReaderT (IORef MutParams) m a)
                       deriving (Functor, Applicative, Monad, MonadTrans)

runMutParamsT :: MonadIO m => MutParamsT m a -> MutParams -> m a
runMutParamsT (MutParamsT m) mut_params =
  liftIO (newIORef mut_params) >>= runReaderT m

instance MonadIO m => MonadIO (MutParamsT m) where
  liftIO m = MutParamsT $ liftIO m

instance MonadParams m => MonadParams (MutParamsT m) where
  getParams n = MutParamsT $ getParams n

class Monad m => MonadMutParams m where
  getMutParams :: Int -> m MutParams

instance MonadIO m => MonadMutParams (MutParamsT m) where
  getMutParams n =
    MutParamsT $ do r <- ask
                    ps <- liftIO $ readIORef r
                    liftIO $ writeIORef r (SMV.drop n ps)
                    return (SMV.take n ps)

instance MonadMutParams m => MonadMutParams (ReaderT r m) where
  getMutParams n = lift $ getMutParams n

-- | A computation for generating samples of distribution families
type SamplingM = ReaderT VIDimAsgn (ParamsT MWCRandM)

-- | Run a 'SamplingM' computation
runSamplingM :: SamplingM a -> MWC.GenIO -> VIDimAsgn -> Params -> IO a
runSamplingM m gen asgn params =
  runRand gen (runParamsT (runReaderT m asgn) params)

embedSamplingFun :: VIDim -> (Params -> MWCRandM a) -> SamplingM a
embedSamplingFun dim f =
  do asgn <- ask
     params <- getParams (evalVIDim dim asgn)
     lift $ lift $ f params

-- | A computation that consumes a vector of parameters and a vector of mutable
-- parameters, the latter being used for accumulating gradients
type ParamsGradM = ReaderT VIDimAsgn (MutParamsT (ParamsT IO))

runParamsGradM :: ParamsGradM a -> VIDimAsgn -> Params -> MutParams -> IO a
runParamsGradM m asgn params mut_params =
  runParamsT (runMutParamsT (runReaderT m asgn) mut_params) params

-- | Take the scaled gradient of a differentiable function with @n@ parameters
-- and accumulate it into the gradient of the current computation
embedScaledGrad :: VIDim -> DiffFun -> Double -> ParamsGradM ()
embedScaledGrad dim f scale =
  do n <- evalVIDim dim <$> ask
     params <- getParams n
     params_grad <- getMutParams n
     liftIO $ scaledGrad n scale f params params_grad

-- | Evaluate a differentiable function with @n@ parameters and its gradient,
-- returning the value and accumulating the gradient into the gradient of the
-- current computation
embedGradAndRet :: VIDim -> DiffFun -> ParamsGradM Double
embedGradAndRet dim f =
  do n <- evalVIDim dim <$> ask
     params <- getParams n
     params_grad <- getMutParams n
     liftIO $ gradAndRet n f params params_grad

-- | A computation for initializing parameters
type InitM = ReaderT VIDimAsgn (MutParamsT MWCRandM)

runInitM :: InitM a -> MWC.GenIO -> VIDimAsgn -> MutParams -> IO a
runInitM m gen asgn mut_params =
  runRand gen $ runMutParamsT (runReaderT m asgn) mut_params

simpleInitFun :: (VIDimAsgn -> MWCRandM [Double]) -> InitM ()
simpleInitFun initFun =
  do asgn <- ask
     params <- lift $ lift $ initFun asgn
     let params_v = SV.fromList params
     mut_params <- getMutParams (length params)
     liftIO $ SV.copy mut_params params_v

-- | A computation for growing sequences of parameters
type GrowM = ReaderT (VIDimAsgn, VIDimAsgn) (MutParamsT (ParamsT MWCRandM))

runGrowM :: GrowM a -> MWC.GenIO -> VIDimAsgn -> VIDimAsgn ->
            Params -> MutParams -> IO a
runGrowM m gen asgn new_asgn params new_params =
  runRand gen $
  runParamsT (runMutParamsT (runReaderT m (asgn, new_asgn)) new_params) params

simpleGrowFun :: VIDim -> (VIDimAsgn -> MWCRandM [Double]) -> GrowM ()
simpleGrowFun dim initFun =
  do (old_asgn, new_asgn) <- ask
     params <- getParams (evalVIDim dim old_asgn)
     old_mut_params <- getMutParams (evalVIDim dim old_asgn)
     liftIO $ SV.copy old_mut_params params
     new_mut_params <-
       getMutParams (evalVIDim dim new_asgn - evalVIDim dim old_asgn)
     inits_l <- lift $ lift $ lift $ initFun new_asgn
     let inits_v = SV.fromList $ drop (SV.length params) inits_l
     liftIO $ SV.copy new_mut_params inits_v

type PPFun = ReaderT VIDimAsgn (ParamsT IO) PP.Doc

simplePPFun :: VIDim -> String -> PPFun
simplePPFun dim nm =
  do n <- evalVIDim dim <$> ask
     params <- getParams n
     return (PP.text nm <+> PP.tupled (map PP.pretty $ SV.toList params))

applyPPFun :: PPFun -> VIDimAsgn -> Params -> IO PP.Doc
applyPPFun f asgn params = runParamsT (runReaderT f asgn) params

----------------------------------------------------------------------
-- * Distribution Approximators for Variational Inference
----------------------------------------------------------------------

type EntropyFun = ParamsGradM Double
type ScaledGradPDFFun a = Double -> a -> ParamsGradM ()
type InitFun = InitM ()
type GrowFun = GrowM ()

{-
-- | The trivial entropy function for the only distribution over ()
unitEntropyFun :: EntropyFun
unitEntropyFun _ _ _ = return 0

-- | Add two independent entropy functions, where the first has the given
-- dimensionality expression
addEntropyFuns :: VIDim -> EntropyFun -> EntropyFun -> EntropyFun
addEntropyFuns dim1 f1 f2 =
  \asgn ps ps_grad ->
  (+) <$>
  f1 asgn (SV.take (evalVIDim dim1 asgn) ps)
  (SMV.take (evalVIDim dim1 asgn) ps_grad)
  <*>
  f2 asgn (SV.drop (evalVIDim dim1 asgn) ps)
  (SMV.drop (evalVIDim dim1 asgn) ps_grad)

-- | The trivial log PDF function with 0 gradient
unitGradPDFFun :: GradPDFFun a
unitGradPDFFun _ _ _ _ _ = return ()

-- | Add two log PDF functions, assuming they are independent, given a
-- volume-preserving map from the codomain type to the two domain types
-- (represented as a pair of functions instead of a function to a pair) along
-- with a dimensionality expression for the number of params of the first
addGradPDFFuns :: (a -> b) -> (a -> c) -> VIDim ->
                  GradPDFFun b -> GradPDFFun c -> GradPDFFun a
addGradPDFFuns split_l split_r dim1 f1 f2 =
  \scale a asgn ps ps_grad ->
  f1 scale (split_l a) asgn (SV.take (evalVIDim dim1 asgn) ps)
  (SMV.take (evalVIDim dim1 asgn) ps_grad)
  >>
  f2 scale (split_r a) asgn (SV.drop (evalVIDim dim1 asgn) ps)
  (SMV.drop (evalVIDim dim1 asgn) ps_grad)

-- | The trivial 'viDistGrowParams' function for no parameters
unitGrowFun :: GrowFun
unitGrowFun _ _ _ _ = return ()

-- | Combine two 'viDistGrowParams' fields, given the dimensionality expression
-- for the first one
combineGrowFuns :: VIDim -> GrowFun -> GrowFun -> GrowFun
combineGrowFuns dim1 f1 f2 =
  \old_asgn new_asgn params mut_params ->
  f1 old_asgn new_asgn (SV.take (evalVIDim dim1 old_asgn) params)
  (SMV.take (evalVIDim dim1 new_asgn) mut_params)
  >>
  f2 old_asgn new_asgn (SV.drop (evalVIDim dim1 old_asgn) params)
  (SMV.drop (evalVIDim dim1 new_asgn) mut_params)
-}

-- | A family of distributions for variational inference
data VIDistFam a =
  VIDistFam
  {
    -- | The dimensionality of the parameters
    viDistDim :: VIDim,

    -- | Sampling function to randomly generate an @a@ given 'viDistDim'
    -- 'Double'-valued parameters
    viDistSample :: SamplingM a,

    -- | Evaluate the entropy at the given parameters and also add the gradient
    -- of the entropy with respect to those parameters into the mutable vector
    viDistEntropy :: EntropyFun,

    -- | Evaluate the gradient with respect to the parameters of the log PDF of
    -- the distribution, scaled by the supplied factor, and add the result to
    -- the supplied mutable vector
    viDistScaledGradPDF :: ScaledGradPDFFun a,

    -- | Randomly initialize the parameters of this distribution family
    viDistInit :: InitFun,

    -- | Grow a vector of parameters from the dimensionality implied by one
    -- assignment to that of another, larger one, storing the result in the
    -- given mutable vector (assuming it has the correct size) and initializing
    -- any new parameters
    viDistGrowParams :: GrowFun,

    -- | Pretty-print a distribution given its parameters
    viDistPP :: PPFun
  }

-- | Build a distribution family for a "simple" distribution, meaning it is not
-- a composite of multiple distributions on sub-components.  Such a distribution
-- is defined by a dimensionality expression, a sampling function,
-- differentiable functions for the entropy and log PDF, and an initialization
-- function for the parameters.
simpleVIFam :: String -> VIDim -> (Params -> MWCRandM a) -> DiffFun ->
               (a -> DiffFun) -> (VIDimAsgn -> MWCRandM [Double]) ->
               VIDistFam a
simpleVIFam nm dim sampleFun entropyFun pdfFun initFun =
  VIDistFam
  { viDistDim = dim
  , viDistSample = embedSamplingFun dim sampleFun
  , viDistEntropy = embedGradAndRet dim entropyFun
  , viDistScaledGradPDF =
      (\scale a -> embedScaledGrad dim (pdfFun a) scale)
  , viDistInit = simpleInitFun initFun
  , viDistGrowParams = simpleGrowFun dim initFun
  , viDistPP = simplePPFun dim nm
  }

-- | Build a distribution family for tuples from a tuple of families
tupleVIFam :: TupleF ts (Compose VIDistFam f) (ADT (TupleF ts)) ->
              VIDistFam (TupleF ts f (ADT (TupleF ts)))
tupleVIFam ds =
  VIDistFam
  { viDistDim = foldrADT (viDistDim . getCompose) (+) 0 ds
  , viDistSample = traverseADT (viDistSample . getCompose) ds
  , viDistEntropy =
      foldrADT (viDistEntropy . getCompose)
      (\m1 m2 -> (+) <$> m1 <*> m2) (return 0) ds
  , viDistScaledGradPDF =
      (\scale tup ->
        foldrADT getConst (>>) (return ()) $
        mapTuple2 (\(Compose d) a ->
                    Const $ viDistScaledGradPDF d scale a) ds tup)
  , viDistInit = foldrADT (viDistInit . getCompose) (>>) (return ()) ds
  , viDistGrowParams =
    foldrADT (viDistGrowParams . getCompose) (>>) (return ()) ds
  , viDistPP =
    PP.tupled <$> sequence (ctorArgsADT (viDistPP . getCompose) ds)
  }

-- | Transform a distribution family over one type to another type using a
-- volume-preserving isomorphism. A volume-preserving map is one that does not
-- squeeze any objects closer together or move them farther apart; the technical
-- definition is that it preserves the measures of all sets with respect to the
-- reference measure on the two spaces.
xformVIDistFam :: (a -> b) -> (b -> a) -> VIDistFam a -> VIDistFam b
xformVIDistFam f_to f_from (VIDistFam {..}) =
  VIDistFam
  {
    viDistSample = fmap f_to $ viDistSample,
    viDistScaledGradPDF = (\scale a -> viDistScaledGradPDF scale (f_from a)),
    ..
  }

-- | Map a vector of distribution families to a distribution family of vectors
vecDistVIFam :: (SVGen.Vector f a, SVGen.Vector f (VIDistFam a)) =>
                f (VIDistFam a) -> VIDistFam (f a)
vecDistVIFam ds =
  VIDistFam
  { viDistDim = SVGen.foldl' (\dim d -> dim + viDistDim d) 0 ds
  , viDistSample = SVGen.mapM viDistSample ds
  , viDistEntropy = SVGen.foldM' (\e d -> (e +) <$> viDistEntropy d) 0 ds
  , viDistScaledGradPDF =
    (\scale v ->
      if SVGen.length v == SVGen.length ds then
        SVGen.zipWithM_ (\d x -> viDistScaledGradPDF d scale x) ds v
      else
        error "vecDist distribution: wrong size vector!")
  , viDistInit = SVGen.mapM_ viDistInit ds
  , viDistGrowParams = SVGen.mapM_ viDistGrowParams ds
  , viDistPP =
    do pps <- mapM viDistPP $ SVGen.toList ds
       return (PP.text "vecDist" <+> PP.list pps)
  }

-- | The distribution family over vectors with a given length whose elements are
-- drawn IID from the supplied distribution family
vecIIDVIFam :: SVGen.Vector f a => VIDim -> VIDistFam a -> VIDistFam (f a)
vecIIDVIFam len d =
  VIDistFam
  { viDistDim = len * viDistDim d
  , viDistSample =
    do asgn <- ask
       SVGen.replicateM (evalVIDim len asgn) (viDistSample d)
  , viDistEntropy =
    do asgn <- ask
       entropies <- replicateM (evalVIDim len asgn) (viDistEntropy d)
       return $ sum entropies
  , viDistScaledGradPDF =
    (\scale v ->
      do asgn <- ask
         if SVGen.length v == evalVIDim len asgn then
           SVGen.forM_ v (viDistScaledGradPDF d scale)
           else
           error "IID distribution: wrong size vector!")
  , viDistInit =
    do asgn <- ask
       replicateM_ (evalVIDim len asgn) (viDistInit d)
  , viDistGrowParams =
    do (old_asgn, new_asgn) <- ask
       replicateM_ (evalVIDim len old_asgn) (viDistGrowParams d)
       g <- lift $ lift $ lift $ nthGenIO 0
       let num_new_elems = evalVIDim len new_asgn - evalVIDim len old_asgn
       replicateM_ num_new_elems $
         do new_params <- getMutParams (evalVIDim (viDistDim d) new_asgn)
            liftIO $ runInitM (viDistInit d) g new_asgn new_params
  , viDistPP =
    do asgn <- ask
       let n = evalVIDim len asgn
       pps <- replicateM n (viDistPP d)
       return (PP.text ("IID(" ++ show n ++ ")") <+> PP.list pps)
  }

----------------------------------------------------------------------
-- * Distribution Approximator Expressions
----------------------------------------------------------------------

-- | An "expression" for building up a family of distributions, which keeps
-- track of how many dimensionality variables we have used so far and also a
-- list of the data file names to be used by 'readJSONVIDistFamExpr'
newtype VIDistFamExpr a =
  VIDistFamExpr { runVIDistFamExpr ::
                    StateT (VIDimVar, [String]) IO (VIDistFam a) }

-- | Evaluate a distribution family expression into a distribution family
evalVIDistFamExpr :: [String] -> VIDistFamExpr a -> IO (VIDistFam a)
evalVIDistFamExpr files (VIDistFamExpr m) = evalStateT m (VIDimVar 0, files)

entropyNaNProtect :: String -> DiffFun -> DiffFun
entropyNaNProtect s f ps =
  let ret = f ps in
  if isNaN ret then
    error ("NaN entropy for distribution " ++ s ++ ", params = " ++ show ps)
  else
    ret

-- | Build a distribution family expression for a "simple" distribution, meaning
-- it is not a composite of multiple distributions on sub-components.  Such a
-- distribution is defined by a dimensionality expression, a sampling function,
-- and differentiable functions for the entropy and log PDF.
simpleVIFamExpr :: String -> VIDim -> (Params -> MWCRandM a) -> DiffFun ->
                   (a -> DiffFun) -> (VIDimAsgn -> MWCRandM [Double]) ->
                   VIDistFamExpr a
simpleVIFamExpr nm dim sampleFun entropyFun pdfFun initFun =
  VIDistFamExpr $ return $
  simpleVIFam nm dim sampleFun (entropyNaNProtect nm entropyFun) pdfFun initFun

-- | The constant distribution (also known as the delta distribution), that
-- returns a single value with unit probability
deltaVIFamExpr :: (Eq a, GrappaShow a) => a -> VIDistFamExpr a
deltaVIFamExpr a =
  simpleVIFamExpr ("Delta(" ++ grappaShow a ++ ")")
  0 (\_ -> return a) (\_ -> 0)
  (\x _ -> if x == a then 0 else log 0)
  (const $ return [])

-- | Build a distribution family expression for the normal distribution, where
-- we use absolute value of @sigma@ to map it to the non-negative reals
normalVIFamExpr :: VIDistFamExpr R
normalVIFamExpr =
  simpleVIFamExpr "Normal" 2 (\ps -> mwcNormal (ps SV.! 0) (abs (ps SV.! 1)))
  (\ps ->  0.5 * (1 + log (2 * pi) + log ((ps V.! 1) * (ps V.! 1))))
  (\x ps -> Log.ln $ normalDensityUnchecked (ps V.! 0) (abs (ps V.! 1)) (fromDouble x))
  (\_ ->
    -- Init mu from the standard normal and sigma from the inverse gamma
    do mu <- random MWC.standard
       inv_sigma_sq <- mwcGamma 1 1
       let sigma = 1 / sqrt inv_sigma_sq
       return [mu,sigma])

-- | Build a distribution family expression for the uniform distribution over
-- the range @(min a b, max a b]@
uniformVIFamExpr :: VIDistFamExpr R
uniformVIFamExpr =
  simpleVIFamExpr "Uniform" 2 (\ps -> mwcUniform (ps SV.! 0) (ps SV.! 1))
  (\ps -> log $ abs (ps V.! 1 - ps V.! 0))
  (\x ps ->
    if min (ps V.! 0) (ps V.! 1) < fromDouble x
       && fromDouble x <= max (ps V.! 0) (ps V.! 1) then
      log $ abs (ps V.! 1 - ps V.! 0)
    else log 0)
  (\_ ->
    -- Initialize to a sub-interval of the unit interval by setting lower =
    -- uniform(0,1) and upper = lower + uniform(0,1) * (1 - lower)
    do lower <- mwcUniform 0 1
       upper_nonnorm <- mwcUniform 0 1
       let upper = lower + upper_nonnorm * (1 - lower)
       return [lower, upper])

-- | Build a distribution family expression for the categorical distribution
-- over @[0,..,n-1]@ with relative probabilities @[a1/a,..,an/a]@, where each
-- @ai@ is the absolute value of the @i@th parameter and @a@ is the sum of the
-- @ai@s. The special case that @n=0@ is treated like @n=1@, meaning it always
-- returns @0@.
categoricalVIFamExpr :: VIDim -> VIDistFamExpr Int
categoricalVIFamExpr dim =
  simpleVIFamExpr "Categorical" dim
  (\ps ->
    -- NOTE: MWC.categorical normalizes the weights for us
    random $ MWC.categorical $ SV.map abs ps)
  (\ps ->
    let ps' = V.map abs ps
        p_sum = V.sum ps' in
    V.sum $ V.map (\p -> (p / p_sum) * log (p / p_sum)) ps')
  (\x ps ->
    let ps' = V.map abs ps
        p_sum = V.sum ps'
        n = V.length ps in
    if n == 0 && x == 0 then 1 else
      if x < 0 || x >= n then 0 else
        log ((ps' V.! x) / p_sum))
  (\asgn ->
    -- Initialize the probabilities from a Dirichlet distribution
    let n = evalVIDim dim asgn in
    if n == 0 then return [] else
      mwcDirichlet (replicate n 1))

-- | Build a distribution family expression for the gamma distribution, where
-- we use absolute value of @k@ and @theta@ to them to the non-negative reals
gammaVIFamExpr :: VIDistFamExpr R
gammaVIFamExpr =
  simpleVIFamExpr "Gamma" 2 (\ps -> mwcGamma (abs (ps SV.! 0)) (abs (ps SV.! 1)))
  (\ps ->
    let k = abs (ps SV.! 0)
        theta = abs (ps SV.! 1) in
    k + log theta + logGamma k + (1-k) * digamma k)
  (\x ps ->
    Log.ln $
    gammaDensityUnchecked (abs (ps V.! 0)) (abs (ps V.! 1)) (fromDouble x))
  (\_ ->
    -- Initialize k and theta from exponentials (because their actual conjugate
    -- priors are messy and I'm too busy for that)
    do k <- mwcExponential 1
       theta <- mwcExponential 1
       return [k,theta])

-- | Build a distribution family expression for the gamma distribution over
-- probabilities, i.e., in log space, where the @k@ and @theta@ parameters are
-- also in log space
gammaProbVIFamExpr :: VIDistFamExpr Prob
gammaProbVIFamExpr =
  simpleVIFamExpr "GammaProb" 2
  (\ps -> Prob <$> mwcGammaLogLog (Log.Exp (ps SV.! 0)) (Log.Exp (ps SV.! 1)))
  (\ps ->
    let k = exp (ps SV.! 0)
        log_theta = ps SV.! 1 in
    k + log_theta + logGamma k + (1-k) * digamma k)
  (\x ps ->
    Log.ln $
    gammaDensityUnchecked (exp (ps V.! 0)) (exp (ps V.! 1))
    (fromDouble $ probToLogR x))
  (\_ ->
    -- Initialize (exp k) and (exp theta) from the exponential distribution, in
    -- keeping with the initialization for gammaVIFamExpr, above. NOTE: this
    -- kind of sucks, because we are just taking log (log uniform(0,1)), and
    -- taking two logarithms in a row probably loses a lot of precision.
    do exp_k <- mwcExponential 1
       exp_theta <- mwcExponential 1
       return [log exp_k, log exp_theta])

-- | Build a distribution family expression for the beta distribution, where
-- we use absolute value of @alpha@ and @beta@ to keep them non-negative
betaVIFamExpr :: VIDistFamExpr R
betaVIFamExpr =
  simpleVIFamExpr "Beta" 2 (\ps -> mwcBeta (abs (ps SV.! 0)) (abs (ps SV.! 1)))
  (\ps ->
    let alpha = abs (ps SV.! 0)
        beta = abs (ps SV.! 1) in
    logGamma alpha + logGamma beta - logGamma (alpha + beta)
    + (alpha + beta - 2) * digamma (alpha + beta)
    - (alpha - 1) * digamma alpha - (beta - 1) * digamma beta)
  (\x ps ->
    Log.ln $
    betaDensityUnchecked (abs (ps V.! 0)) (abs (ps V.! 1)) (fromDouble x))
  (\_ ->
    -- Initialize alpha and beta from the exponential distribution, because why
    -- not? And because, as Wikipedia says (for the Dirichlet distribution,
    -- which is similar), "In the published literature there is no practical
    -- algorithm to efficiently generate samples from" the conjugate prior.
    do alpha <- mwcExponential 1
       beta <- mwcExponential 1
       return [alpha,beta])

-- | Build a distribution family expression for the beta distribution over
-- probabilities, i.e., in log space, where the @alpha@ and @beta@ parameters
-- are also in log space
betaProbVIFamExpr :: VIDistFamExpr Prob
betaProbVIFamExpr =
  simpleVIFamExpr "BetaProb" 2
  (\ps -> Prob <$> mwcBetaLogLog (Log.Exp (ps SV.! 0)) (Log.Exp (ps SV.! 1)))
  (\ps ->
    let alpha = exp (ps SV.! 0)
        beta = exp (ps SV.! 1) in
    logGamma alpha + logGamma beta - logGamma (alpha + beta)
    + (alpha + beta - 2) * digamma (alpha + beta)
    - (alpha - 1) * digamma alpha - (beta - 1) * digamma beta)
  (\x ps ->
    Log.ln $
    betaDensityLog (exp (ps V.! 0)) (exp (ps V.! 1))
    (fmap fromDouble $ fromProb x))
  (\_ ->
    -- Initialize (exp alpha) and (exp beta) from the exponential distribution,
    -- in keeping with the initialization for betaVIFamExpr, above. NOTE: this
    -- kind of sucks, because we are just taking log (log uniform(0,1)), and
    -- taking two logarithms in a row probably loses a lot of precision.
    do exp_alpha <- mwcExponential 1
       exp_beta <- mwcExponential 1
       return [log exp_alpha, log exp_beta])

-- | Build a distribution family for the dirichlet distribution, where the
-- alphas are in log space, so they can never be negative and so that the
-- optimization can do better at picking very small alphas
dirichletVIFamExpr :: VIDim -> VIDistFamExpr [R]
dirichletVIFamExpr dim =
  simpleVIFamExpr "Dirichlet" dim
  (\ps -> mwcDirichlet (SV.toList $ SV.map exp ps))
  (\ps ->
    let alphas = SV.toList $ SV.map exp ps
        k = fromIntegral $ SV.length ps
        alpha0 = sum alphas in
    sum (map logGamma alphas) - logGamma alpha0
    + (alpha0 - k) * digamma alpha0
    - sum (flip map alphas $ \alpha -> (alpha - 1) * digamma alpha))
  (\x ps -> Log.ln $
            dirichletDensity (SV.toList $ SV.map exp ps) (map fromDouble x))
  (\asgn ->
    -- Initialize alphas from the exponential distribution, because why not? And
    -- because, as Wikipedia says, "In the published literature there is no
    -- practical algorithm to efficiently generate samples from" the conjugate
    -- prior.
    replicateM (evalVIDim dim asgn) (mwcExponential 1))

-- | Build a distribution family for the dirichlet distribution over
-- probabilities, i.e., over lists of reals in log space. The alphas are in log
-- space, so they can never be negative and so that the optimization can do
-- better at picking very small alphas
dirichletProbVIFamExpr :: VIDim -> VIDistFamExpr [Prob]
dirichletProbVIFamExpr dim =
  simpleVIFamExpr "DirichletProb" dim
  (\ps -> map Prob <$> mwcDirichletLog (SV.toList $ SV.map exp ps))
  (\ps ->
    let alphas = SV.toList $ SV.map exp ps
        k = fromIntegral $ SV.length ps
        alpha0 = sum alphas in
    sum (map logGamma alphas) - logGamma alpha0
    + (alpha0 - k) * digamma alpha0
    - sum (flip map alphas $ \alpha -> (alpha - 1) * digamma alpha))
  (\x ps -> Log.ln $
            dirichletDensityLog (SV.toList $ SV.map exp ps)
            (map (fmap fromDouble . fromProb) x))
  (\asgn ->
    -- Initialize (exp alphas) from the exponential distribution, because why
    -- not? And because, as Wikipedia says, "In the published literature there
    -- is no practical algorithm to efficiently generate samples from" the
    -- conjugate prior. NOTE: this kind of sucks, because we are just taking log
    -- (log uniform(0,1)), and taking two logarithms in a row probably loses a
    -- lot of precision.
    map log <$> replicateM (evalVIDim dim asgn) (mwcExponential 1))

-- | Bind a fresh dimensionality variable in a distribution family expression
bindVIDimFamExpr :: (VIDim -> VIDistFamExpr a) -> VIDistFamExpr a
bindVIDimFamExpr f =
  VIDistFamExpr $ do
  (v,files) <- get
  put $ (nextVar v, files)
  runVIDistFamExpr $ f (varVIDim v)

-- | Build a distribution family expression for a tuple from a tuple of
-- distribution family expressions
tupleVIFamExpr :: TupleF ts (Compose VIDistFamExpr f) (ADT (TupleF ts)) ->
                  VIDistFamExpr (TupleF ts f (ADT (TupleF ts)))
tupleVIFamExpr d_exprs =
  VIDistFamExpr $ fmap tupleVIFam $
  traverseADT ((Compose <$>) . runVIDistFamExpr . getCompose) d_exprs

-- | Transform a distribution family expression over one type to another type
-- using a volume-preserving isomorphism; see 'xformVIDistFam'.
xformVIDistFamExpr :: (a -> b) -> (b -> a) -> VIDistFamExpr a -> VIDistFamExpr b
xformVIDistFamExpr f_to f_from expr =
  VIDistFamExpr (xformVIDistFam f_to f_from <$> runVIDistFamExpr expr)

-- | The distribution family expression over vectors with a given length whose
-- elements are drawn IID from the supplied distribution family expression
vecIIDVIFamExpr :: VIDim -> VIDistFamExpr a -> VIDistFamExpr (Vector a)
vecIIDVIFamExpr len d_expr =
  VIDistFamExpr (vecIIDVIFam len <$> runVIDistFamExpr d_expr)

-- | Map a vector of distribution family expressions to a distribution family
-- expression over vectors
vecDistVIFamExpr :: (SVGen.Vector f a, SVGen.Vector f (VIDistFam a),
                     SVGen.Vector f (VIDistFamExpr a)) =>
                    f (VIDistFamExpr a) -> VIDistFamExpr (f a)
vecDistVIFamExpr d_exprs =
  VIDistFamExpr (vecDistVIFam <$> SVGen.mapM runVIDistFamExpr d_exprs)

-- | The distribution family expression over lists with a given length whose
-- elements are drawn IID from the supplied distribution family expression
iidVIFamExpr :: VIDim -> VIDistFamExpr a -> VIDistFamExpr [a]
iidVIFamExpr len d_expr =
  xformVIDistFamExpr V.toList V.fromList $ vecIIDVIFamExpr len d_expr

-- | The distribution family expression over unboxed 'ProbVector's with a given
-- length whose elements are drawn IID from the supplied distribution family
-- expression
iidPVVIFamExpr :: VIDim -> VIDistFamExpr Prob -> VIDistFamExpr ProbVector
iidPVVIFamExpr len d_expr =
  xformVIDistFamExpr ProbVector unProbVector $
  VIDistFamExpr (vecIIDVIFam len <$> runVIDistFamExpr
                 (xformVIDistFamExpr probToLogR logRToProb d_expr))

-- | The distribution family expression over unboxed 'RVector's with a given
-- length whose elements are drawn IID from the supplied distribution family
-- expression
iidVVIFamExpr :: VIDim -> VIDistFamExpr R -> VIDistFamExpr RVector
iidVVIFamExpr len d_expr =
  xformVIDistFamExpr RVector unRVector $
  VIDistFamExpr (vecIIDVIFam len <$> runVIDistFamExpr d_expr)

-- | This distribution family reads an input of type @a@ from @stdin@ as a JSON
-- file and then passes that input to the supplied function to build a
-- distribution family of type @b@
mapJSONVIDistFamExpr :: FromJSON a => (a -> VIDistFamExpr b) -> VIDistFamExpr b
mapJSONVIDistFamExpr f =
  VIDistFamExpr $
  do (v,files) <- get
     let (file, files') =
           if length files == 0 then
             error "mapJSONVIDistFamExpr: not enough data files specified"
           else (head files, tail files)
     put (v, files')
     a <- liftIO $ readJSONfile file
     runVIDistFamExpr $ f a

-- | This distribution family is a delta distribution (i.e., one that always
-- returns the same value, like 'deltaVIFamExpr') that reads its input from
-- stdin as a JSON file
readJSONVIDistFamExpr :: (Eq a, GrappaShow a, FromJSON a) => VIDistFamExpr a
readJSONVIDistFamExpr = mapJSONVIDistFamExpr deltaVIFamExpr

----------------------------------------------------------------------
-- * Variational Inference, Yay!
----------------------------------------------------------------------

-- | A PVIE model is a dimensionaltiy variable assignment + model parameters
data PVIEModel =
  PVIEModel { pvieModelAsgn :: VIDimAsgn,
              pvieModelParams :: Params }
  deriving Generic

instance FromJSON PVIEModel where

instance ToJSON PVIEModel where
  toEncoding = genericToEncoding defaultOptions

readPVIEModel :: MWC.GenIO -> VIDistFam a -> String -> IO PVIEModel
readPVIEModel g d filename =
  doesFileExist filename >>= \ex ->
  if ex then
    -- If filename exists, then read it
    readJSONfile filename
  else
    -- Otherwise, start with 0 for all dimensionality variables, and
    -- initialize all params to 1
    do let asgn = zeroAsgn $ viDimFirstUnusedVar $ viDistDim d
       mut_params <- SMV.new (evalVIDim (viDistDim d) asgn)
       runInitM (viDistInit d) g asgn mut_params
       params <- SV.freeze mut_params
       return $ PVIEModel asgn params

-- | The PVIE modes: training and evaluation
data PVIEMode = PVIETrainMode | PVIEEvalMode

-- | The PVIE command-line options
data PVIEOpts =
  PVIEOpts {
  pvieModelFile :: String,
  pvieDataFiles :: [String],
  pvieMode :: PVIEMode,
  pvieVerbosity :: Int,
  pvieDebugTruncate :: Maybe Int,
  pvieNumSamples :: Int,
  pvieGD :: Bool
  }

-- | Parser for PVIE command-line options
pvieOptsParser :: Parser PVIEOpts
pvieOptsParser =
  PVIEOpts
  <$> (strOption (long "model"
                  <> short 'm'
                  <> metavar "MODEL"
                  <> help "Model file"
                  <> value "-"
                  <> showDefault))
  <*> many (strOption (long "data"
                       <> short 'd'
                       <> metavar "DATAFILE"
                       <> help "Data file"))
  <*> (flag' PVIETrainMode (long "train" <> short 't'
                            <> help "Specifies training mode")
       <|>
       flag' PVIEEvalMode (long "eval" <> short 'e'
                            <> help "Specifies evaluation mode"))
  <*> (option auto (long "verbosity" <> short 'v'
                    <> help "Verbosity level for debugging"
                    <> showDefault <> value 0 <> metavar "VERBOSITY"))
  <*> (option (Just <$> auto)
       (long "truncate" <> short 'T'
        <> help "Truncate length for debugging strings"
        <> showDefault <> value Nothing <> metavar "TRUNCATE"))
  <*> (option auto (long "numsamples" <> short 'N'
                    <> help "Number of samples for approximating the variational family"
                    <> showDefault <> value 1000 <> metavar "NUMSAMPLES"))
  <*> (flag False True
       (long "gd" <>
        help "Use gradient descent instead of BFGS to optimize (only meaningful in training mode)"))

-- | Parse the command-line options
parsePVIEOpts :: IO PVIEOpts
parsePVIEOpts = execParser (info (pvieOptsParser <**> helper)
                            (fullDesc <>
                             progDesc "FIXME: description of PVIE"))

debugTruncate :: PVIEOpts -> String -> String
debugTruncate opts s
  | Just l <- pvieDebugTruncate opts
  , length s >= l = take l s ++ " ..."
debugTruncate _ s = s

-- | Print debugging info if the verbosity level is @>= level@
debugM :: PVIEOpts -> Int -> String -> IO ()
debugM opts level s
  | pvieVerbosity opts >= level = traceM (debugTruncate opts s)
debugM _ _ _ = return ()

-- | Find the parameters that /minimize/ the given differentiable function
optimize :: PVIEOpts -> (Params -> MutParams -> IO Double) ->
            MutParams -> IO Double
optimize opts f mut_params =
  do let len = SMV.length mut_params
     let meth = if pvieGD opts then SteepestDescent else VectorBFGS2
     params <- SV.freeze mut_params
     let eval_f ps =
           unsafePerformIO $
           do debugM opts 2 ("neg_elbo: params = " ++ show ps)
              grad <- SMV.replicate len 0
              ret <- f ps grad
              debugM opts 2 ("surprisal = " ++ show ret)
              return ret
     let grad_f ps =
           unsafePerformIO $
           do debugM opts 2 ("neg_elbo_grad: params = " ++ show ps)
              grad <- SMV.replicate len 0
              _ <- f ps grad
              ret <- SV.unsafeFreeze grad
              debugM opts 2 ("grad = " ++ show ret)
              return ret
     let (opt_params,_) =
           minimizeVD meth 0.0001 1000 100 0.1 eval_f grad_f params
     SV.copy mut_params opt_params
     return $ eval_f opt_params

-- FIXME HERE: make pvie_epsilon into command-line options

-- | The amount that PVIE has to improve in an interation to seem "better"
pvie_epsilon :: Double
pvie_epsilon = 1.0e-6

-- | The negative infinity value
negInfinity :: Double
negInfinity = log 0

-- | Test if a 'Double' is (positive or negative) infinity or NaN
isInfiniteOrNaN :: Double -> Bool
isInfiniteOrNaN x = isInfinite x || isNaN x

-- | Compute the negative Evidence Lower BOund (or ELBO) and its gradient
neg_elbo_with_grad :: GrappaShow a => PVIEOpts -> MWC.GenIO ->
                      VIDistFam a -> (a -> Double) -> VIDimAsgn ->
                      (Params -> MutParams -> IO Double)
neg_elbo_with_grad opts g d log_p asgn params grad =
  do ret <- elbo_with_grad opts g d log_p asgn params grad
     forM_ [0 .. (SMV.length grad)-1] (SMV.modify grad negate)
     return (-ret)

-- | Compute the Evidence Lower BOund (or ELBO) and its gradient
elbo_with_grad :: GrappaShow a => PVIEOpts -> MWC.GenIO ->
                  VIDistFam a -> (a -> Double) -> VIDimAsgn ->
                  (Params -> MutParams -> IO Double)
elbo_with_grad opts g d log_p asgn params grad =
  (replicateM (pvieNumSamples opts) $
   do s <- runSamplingM (viDistSample d) g asgn params
      return (s, log_p s)) >>= \samples_log_ps ->
  case find (isInfiniteOrNaN . snd) samples_log_ps of
    Just (bad_samp, p) ->
      -- If any of our samples have 0 probability in our model (i.e., from
      -- log_p), the elbo is -infinity, and the gradient is undefined. However,
      -- knowing our optimization algorithm, we cheat, and tell it that the
      -- gradient is just the sum of the negatives of the gradients of d at
      -- those 0 probability samples. That way, our optimization algorithm will
      -- keep trying to reduce the probabilities of generating bad samples until
      -- it finally succeeds. Note that we still return -infinity as the value.
      do debugM opts 3 ((if isNaN p then "NaN" else
                           if p < 0 then "Zero-probability"
                           else "Infinite (?) probability") ++
                        " sample: " ++ grappaShow bad_samp)
         forM_ samples_log_ps $ \(samp, samp_p) ->
           when (isInfiniteOrNaN samp_p) $
           runParamsGradM (viDistScaledGradPDF d (-1e9) samp) asgn params grad
         return negInfinity
    _ ->
      do let n = fromIntegral (pvieNumSamples opts)
         entr <- runParamsGradM (viDistEntropy d) asgn params grad
         debugM opts 2 ("Entropy: " ++ show entr)
         entr_grad <- SV.freeze grad
         debugM opts 3 ("Entropy grad: " ++ show entr_grad)
         forM_ samples_log_ps $ \(samp, p) ->
           runParamsGradM (viDistScaledGradPDF d (p / n) samp) asgn params grad
         -- grad_const <- SV.unsafeFreeze grad
         -- traceM ("grad = " ++ show grad_const)
         let ret = entr + (1/n) * sum (map snd samples_log_ps)
         -- traceM ("surprisal = " ++ show ret)
         debugM opts 4 ("Sample 0: " ++ grappaShow (fst $ head samples_log_ps))
         return ret

-- | The main entrypoint for the PVIE engine
pvie_main :: GrappaShow a => VIDistFamExpr a -> (a -> Double) -> IO ()
pvie_main dist_expr log_p =
  do opts <- parsePVIEOpts
     dist_fam <- evalVIDistFamExpr (pvieDataFiles opts) dist_expr
     case pvieMode opts of
       PVIETrainMode ->
         do (model@(PVIEModel asgn params), val) <-
              pvie_train opts dist_fam log_p
            pp <- applyPPFun (viDistPP dist_fam) asgn params
            debugM opts 1 ("Final model:\n" ++ show pp)
            debugM opts 1 ("Surprisal score: " ++ show val)
            writeJSONfile (pvieModelFile opts) model
       PVIEEvalMode ->
         do val <- pvie_eval opts dist_fam log_p
            putStrLn ("Surprisal score: " ++ show val)

-- | The main entrypoint for the PVIE engine with anomaly score
pvie_anom_main :: GrappaShow a => VIDistFamExpr a -> (a -> Double) ->
                  (a -> Prob) -> IO ()
pvie_anom_main dist_expr log_p anom_score =
  do opts <- parsePVIEOpts
     dist_fam <- evalVIDistFamExpr (pvieDataFiles opts) dist_expr
     case pvieMode opts of
       PVIETrainMode ->
         do (model@(PVIEModel asgn params), val) <-
              pvie_train opts dist_fam log_p
            pp <- applyPPFun (viDistPP dist_fam) asgn params
            debugM opts 1 ("Final model:\n" ++ show pp)
            debugM opts 1 ("Surprisal score: " ++ show val)
            writeJSONfile (pvieModelFile opts) model
       PVIEEvalMode ->
         do score <- pvie_eval_anom opts dist_fam anom_score
            putStrLn $ show (1 - min 1 (probToR score))

-- | The PVIE eval mode
pvie_eval :: GrappaShow a => PVIEOpts -> VIDistFam a -> (a -> Double) ->
             IO Double
pvie_eval opts d log_p =
  do PVIEModel asgn params <- readJSONfile $ pvieModelFile opts
     g <- MWC.createSystemRandom
     grad <- SMV.new (evalVIDim (viDistDim d) asgn)
     val <- neg_elbo_with_grad opts g d log_p asgn params grad
     return val

-- | Evaluate the anomaly score, which is the average over @N@ samples of the
-- supplied anomaly score function applied to each sample
pvie_eval_anom :: GrappaShow a => PVIEOpts -> VIDistFam a -> (a -> Prob) ->
                  IO Prob
pvie_eval_anom opts d anom_score =
  do PVIEModel asgn params <- readJSONfile $ pvieModelFile opts
     g <- MWC.createSystemRandom
     anom_scores <-
       replicateM (pvieNumSamples opts) $
       do samp <- runSamplingM (viDistSample d) g asgn params
          return (probToLogR $ anom_score samp)
     return $ logRToProb (sum anom_scores / fromIntegral (pvieNumSamples opts))

-- | The PVIE training mode
pvie_train :: GrappaShow a => PVIEOpts -> VIDistFam a -> (a -> Double) ->
              IO (PVIEModel, Double)
pvie_train opts d log_p = init_pvie where

  -- Initialize PVIE and start it running
  init_pvie :: IO (PVIEModel, Double)
  init_pvie = do
    g <- MWC.createSystemRandom
    PVIEModel asgn params <- readPVIEModel g d (pvieModelFile opts)
    mut_params <- SV.unsafeThaw params
    -- Generate the initial value to try to beat
    val <- optimize opts (neg_elbo_with_grad opts g d log_p asgn) mut_params
    params' <- SV.unsafeFreeze mut_params
    -- If there are no dimensionality variables in our distribution family, then
    -- there is nothing to increment, so we are done
    if viDimFirstUnusedVar (viDistDim d) == zeroVar then
      outer g asgn params' val False
      else
      -- Perform the first iteration of optimizing the dimensionality variables
      outer g asgn params' val True

  -- The main outer loop
  outer :: MWC.GenIO -> VIDimAsgn -> Params -> Double -> Bool ->
           IO (PVIEModel, Double)
  outer _ asgn params last_val False =
    -- If our last iteration of the outer loop did not improve, we are done, and
    -- return the current assignment, parameters, and value
    return (PVIEModel asgn params, last_val)

  outer g asgn params last_val True =
    do
      -- Otherwise, start by allocating two new arrays, one for the current
      -- parameters and one for the scratch space. The size should be big enough
      -- for if we increment all dimensionality variables by one, which is the
      -- most we will do in the upcoming iteration of the outer loop.
      new_mut_params <- SMV.new (evalVIDim (viDistDim d) $ incrAsgnAll asgn)
      scratch <- SMV.new (evalVIDim (viDistDim d) $ incrAsgnAll asgn)
      -- Copy the old parameters into the new mutable params
      let new_mut_params_slice = SMV.slice 0 (SV.length params) new_mut_params
      SV.copy new_mut_params_slice params
      new_params <- SV.unsafeFreeze new_mut_params
      -- Start the inner loop at the first dimensionality variable, telling it
      -- that we have not yet improved
      inner g asgn new_params last_val zeroVar scratch False

  -- The main inner loop: takes in the current assignment to dimensionality
  -- variables, the current "best" params for the associated size, the value for
  -- those params, the next dimensionality variable to try to increment,
  -- pre-allocated scratch space for parameters, and whether we have already
  -- seen improvement in this iteration of the outer loop
  inner :: MWC.GenIO -> VIDimAsgn -> Params -> Double -> VIDimVar ->
           MutParams -> Bool -> IO (PVIEModel, Double)
  inner g asgn last_params last_val next_var _ improved
    | next_var == viDimFirstUnusedVar (viDistDim d)
    = outer g asgn last_params last_val improved

  inner g asgn last_params last_val next_var scratch improved = do
    -- The next assignment we will try = increment next_var
    let new_asgn = incrAsgn next_var asgn
    -- Copy our current best parameters into our scratch area
    runGrowM (viDistGrowParams d) g asgn new_asgn last_params scratch
    -- Optimize those parameters
    new_val <- optimize opts (neg_elbo_with_grad opts g d log_p asgn) scratch
    -- Test how much we improved
    if last_val - new_val >= pvie_epsilon then
      -- If we did improve, swap last_params and scratch, then iterate
      do new_params <- SV.unsafeFreeze scratch
         new_scratch <- SV.unsafeThaw last_params
         inner g new_asgn new_params new_val (nextVar next_var) new_scratch True
      else
      -- Otherwise, keep the same assignment, params, and scratch
      inner g asgn last_params last_val (nextVar next_var) scratch improved


{-# OPTIONS -cpp  #-}
-----------------------------------------------------------------------------
-- |
-- Module      :  Control.Parallel.Strategies
-- Copyright   :  (c) The University of Glasgow 2001
-- License     :  BSD-style (see the file libraries/base/LICENSE)
--
-- Maintainer  :  libraries@haskell.org
-- Stability   :  experimental
-- Portability :  non-portable
--
-- Parallel strategy combinators. See
-- <http://www.macs.hw.ac.uk/~dsg/gph/papers/html/Strategies/strategies.html>
-- for more information.
--
-- Original authors:
--	Phil Trinder, Hans-Wolfgang Loidl, Kevin Hammond et al.
--
-----------------------------------------------------------------------------
-- #ifdef __PARALLEL_HASKELL__
-- module Control.Parallel.Strategies (
-- #else
module Parallel.Strategies (
-- #endif
   -- *	Strategy Type, Application and Semantics
   Done, Strategy,
   (>|), (>||),
   using, demanding, sparking,
   -- *	Basic Strategies
   r0, rwhnf, NFData(..),
   -- * Strategic Function Application
   ($|), ($||),
   (.|), (.||),
   (-|), (-||),
   -- * Tuples
   seqPair, parPair,
   seqTriple, parTriple,
   -- * Lists: Parallel Strategies
   parList, parListN, parListNth, parListChunk,
   parMap, parFlatMap, parZipWith,
   -- * Lists: Sequential Strategies
   seqList, seqListN, seqListNth, parBuffer,
   -- *	Arrays
   seqArr, parArr,
   -- * Deprecated types and functions
   sPar, sSeq,
   Assoc(..),
   fstPairFstList, force, sforce
  ) where

-- based on hslibs/concurrent/Strategies.lhs; see it for more detailed
-- code comments.
#ifdef __PARALLEL_HASKELL__
#else
#warning "Using nonparallel GHC. The following line will fail."
#endif
import Control.Parallel as Parallel (pseq)
import Data.Array
import Data.Complex
import Data.Int
import qualified Data.IntMap (IntMap, toList)
import qualified Data.IntSet (IntSet, toList)
import qualified Data.Map (Map, toList)
import qualified Data.Set (Set, toList)
import qualified Data.Tree (Tree(..))
import Data.Word

import Prelude hiding (seq)
import qualified Prelude (seq)

-- not a terribly portable way of getting at Ratio rep.
#ifdef __GLASGOW_HASKELL__
import GHC.Real	(Ratio(..))	-- The basic defns for Ratio
#endif

#ifdef __HUGS__
import Hugs.Prelude(Ratio(..) )
#endif

#ifdef __NHC__
import Ratio (Ratio(..) )
#endif

infixl 0 `using`,`demanding`,`sparking`              -- weakest precedence!

infixr 2 >||                -- another name for par
infixr 3 >|                 -- another name for seq
infixl 6 $||, $|            -- strategic function application (seq and par)
infixl 9 .|, .||, -|, -||   -- strategic (inverse) function composition

-- We need 'pseq', not the Prelude 'seq' here.  See the documentation
-- with 'pseq' in Control.Parallel.
seq = Parallel.pseq
par x y = y

------------------------------------------------------------------------------
-- *			Strategy Type, Application and Semantics
------------------------------------------------------------------------------

{-
The basic combinators for strategies are 'par' and 'seq' but with types that
indicate that they only combine the results of a strategy application.

NB: This version can be used with Haskell 1.4 (GHC 2.05 and beyond), *but*
    you won't get strategy checking on seq (only on par)!

The operators >| and >|| are alternative names for `seq` and `par`.
With the introduction of a Prelude function `seq` separating the Prelude
function from the Strategy function becomes a pain. The notation also matches
the notation for strategic function application.
-}

type Done = ()

-- | A strategy takes a value and returns a 'Done' value to indicate that
--   the specifed evaluation has been performed.
type Strategy a = a -> Done

-- | Evaluates the first argument before the second.
(>|) :: Done -> Done -> Done
{-# INLINE (>|) #-}
(>|) = Prelude.seq

-- | Evaluates the first argument in parallel with the second.
(>||) :: Done -> Done -> Done
{-# INLINE (>||) #-}
(>||) = par -- Parallel.par

-- | Takes a value and a strategy, and applies the strategy to the
-- value before returning the value. Used to express data-oriented
-- parallelism. @x \`using\` s@ is a projection on @x@, i.e. both:
--
-- [a retraction] @x \`using\` s@ ⊑ @x@
--
-- [idempotent] @(x \`using\` s) \`using\` s@ = @x \`using\` s@
--
using :: a -> Strategy a -> a
using x s = s x `seq` x

-- | Evaluates the second argument before the first.
-- Used to express control-oriented parallelism. The second
-- argument is usually a strategy application.
demanding :: a -> Done -> a
demanding = flip seq

-- | Evaluates the second argument in parallel with the first.
-- Used to express control-oriented
-- parallelism. The second argument is usually a strategy application.
sparking :: a -> Done -> a
sparking  = flip par -- Parallel.par
-- Sparking should only be used
-- with a singleton sequence as it is not necessarily executed.

-- | A strategy corresponding to 'par':
-- @x \`par\` e@ = @e \`using\` sPar x@.
--
-- 'sPar' has been superceded by 'sparking'.
-- Replace @e \`using\` sPar x@ with @e \`sparking\` rwhnf x@.
{-# DEPRECATED sPar "Use sparking instead." #-}
sPar :: a -> Strategy b
sPar x y = x `par` ()

-- | A strategy corresponding to 'seq':
-- @x \`seq\` e@ = @e \`using\` sSeq x@.
--
-- 'sSeq' has been superceded by 'demanding'.
-- Replace @e \`using\` sSeq x@ with @e \`demanding\` rwhnf x@.
{-# DEPRECATED sSeq "Use demanding instead." #-}
sSeq :: a -> Strategy b
sSeq x y = x `seq` ()

-----------------------------------------------------------------------------
-- *			Basic Strategies
-----------------------------------------------------------------------------

-- | Performs /no/ evaluation of its argument.
r0 :: Strategy a
r0 x = ()

-- | Reduces its argument to weak head normal form.
rwhnf :: Strategy a
rwhnf x = x `seq` ()

class NFData a where
  -- | Reduces its argument to (head) normal form.
  rnf :: Strategy a
  -- Default method. Useful for base types. A specific method is necessay for
  -- constructed types
  rnf = rwhnf

class (NFData a, Integral a) => NFDataIntegral a
class (NFData a, Ord a) => NFDataOrd a

------------------------------------------------------------------------------
-- *                     Strategic Function Application
------------------------------------------------------------------------------

{-
These are very
handy when writing pipeline parallelism asa sequence of @$@, @$|@ and
@$||@'s. There is no need of naming intermediate values in this case. The
separation of algorithm from strategy is achieved by allowing strategies
only as second arguments to @$|@ and @$||@.
-}

-- | Sequential function application. The argument is evaluated using
--   the given strategy before it is given to the function.
($|) :: (a -> b) -> Strategy a -> a -> b
f $| s  = \ x -> f x `demanding` s x

-- | Parallel function application. The argument is evaluated using
-- the given strategy, in parallel with the function application.
($||) :: (a -> b) -> Strategy a -> a -> b
f $|| s = \ x -> f x `sparking` s x

-- | Sequential function composition. The result of
-- the second function is evaluated using the given strategy,
-- and then given to the first function.
(.|) :: (b -> c) -> Strategy b -> (a -> b) -> (a -> c)
(.|) f s g = \ x -> let  gx = g x
                    in   f gx `demanding` s gx

-- | Parallel function composition. The result of the second
-- function is evaluated using the given strategy,
-- in parallel with the application of the first function.
(.||) :: (b -> c) -> Strategy b -> (a -> b) -> (a -> c)
(.||) f s g = \ x -> let  gx = g x
                     in   f gx `sparking` s gx

-- | Sequential inverse function composition,
-- for those who read their programs from left to right.
-- The result of the first function is evaluated using the
-- given strategy, and then given to the second function.
(-|) :: (a -> b) -> Strategy b -> (b -> c) -> (a -> c)
(-|) f s g = \ x -> let  fx = f x
                    in   g fx `demanding` s fx

-- | Parallel inverse function composition,
-- for those who read their programs from left to right.
-- The result of the first function is evaluated using the
-- given strategy, in parallel with the application of the
-- second function.
(-||) :: (a -> b) -> Strategy b -> (b -> c) -> (a -> c)
(-||) f s g = \ x -> let  fx = f x
                     in   g fx `sparking` s fx

------------------------------------------------------------------------------
--			Marking a Strategy
------------------------------------------------------------------------------

{-
Marking a strategy.

Actually, @markStrat@  sticks a label @n@  into the sparkname  field of the
thread executing strategy @s@. Together with a runtime-system that supports
propagation of sparknames to the children this means that this strategy and
all its children have  the sparkname @n@ (if the  static sparkname field in
the @parGlobal@ annotation contains the value 1). Note, that the @SN@ field
of starting the marked strategy itself contains the sparkname of the parent
thread. The END event contains @n@ as sparkname.
-}

#if 0
markStrat :: Int -> Strategy a -> Strategy a
markStrat n s x = unsafePerformPrimIO (
     _casm_ ``%r = set_sparkname(CurrentTSO, %0);'' n `thenPrimIO` \ z ->
     returnPrimIO (s x))
#endif

-----------------------------------------------------------------------------
--			Strategy Instances and Functions
-----------------------------------------------------------------------------

-----------------------------------------------------------------------------
-- *	                Tuples
-----------------------------------------------------------------------------

{-
We currently support up to 9-tuples. If you need longer tuples you have to
add the instance explicitly to your program.
-}

instance (NFData a, NFData b) => NFData (a,b) where
  rnf (x,y) = rnf x `seq` rnf y

instance (NFData a, NFData b, NFData c) => NFData (a,b,c) where
  rnf (x,y,z) = rnf x `seq` rnf y `seq` rnf z

instance (NFData a, NFData b, NFData c, NFData d) => NFData (a,b,c,d) where
  rnf (x1,x2,x3,x4) = rnf x1 `seq`
		        rnf x2 `seq`
		        rnf x3 `seq`
		        rnf x4

instance (NFData a1, NFData a2, NFData a3, NFData a4, NFData a5) =>
         NFData (a1, a2, a3, a4, a5) where
  rnf (x1, x2, x3, x4, x5) =
                  rnf x1 `seq`
                  rnf x2 `seq`
                  rnf x3 `seq`
                  rnf x4 `seq`
                  rnf x5

instance (NFData a1, NFData a2, NFData a3, NFData a4, NFData a5, NFData a6) =>
         NFData (a1, a2, a3, a4, a5, a6) where
  rnf (x1, x2, x3, x4, x5, x6) =
                  rnf x1 `seq`
                  rnf x2 `seq`
                  rnf x3 `seq`
                  rnf x4 `seq`
                  rnf x5 `seq`
                  rnf x6

instance (NFData a1, NFData a2, NFData a3, NFData a4, NFData a5, NFData a6, NFData a7) =>
         NFData (a1, a2, a3, a4, a5, a6, a7) where
  rnf (x1, x2, x3, x4, x5, x6, x7) =
                  rnf x1 `seq`
                  rnf x2 `seq`
                  rnf x3 `seq`
                  rnf x4 `seq`
                  rnf x5 `seq`
                  rnf x6 `seq`
                  rnf x7

instance (NFData a1, NFData a2, NFData a3, NFData a4, NFData a5, NFData a6, NFData a7, NFData a8) =>
         NFData (a1, a2, a3, a4, a5, a6, a7, a8) where
  rnf (x1, x2, x3, x4, x5, x6, x7, x8) =
                  rnf x1 `seq`
                  rnf x2 `seq`
                  rnf x3 `seq`
                  rnf x4 `seq`
                  rnf x5 `seq`
                  rnf x6 `seq`
                  rnf x7 `seq`
                  rnf x8

instance (NFData a1, NFData a2, NFData a3, NFData a4, NFData a5, NFData a6, NFData a7, NFData a8, NFData a9) =>
         NFData (a1, a2, a3, a4, a5, a6, a7, a8, a9) where
  rnf (x1, x2, x3, x4, x5, x6, x7, x8, x9) =
                  rnf x1 `seq`
                  rnf x2 `seq`
                  rnf x3 `seq`
                  rnf x4 `seq`
                  rnf x5 `seq`
                  rnf x6 `seq`
                  rnf x7 `seq`
                  rnf x8 `seq`
                  rnf x9

-- | Apply two strategies to the elements of a pair sequentially
--   from left to right.
seqPair :: Strategy a -> Strategy b -> Strategy (a,b)
seqPair strata stratb (x,y) = strata x `seq` stratb y

-- | Apply two strategies to the elements of a pair in parallel.
parPair :: Strategy a -> Strategy b -> Strategy (a,b)
parPair strata stratb (x,y) = strata x `par` stratb y `par` ()
-- The reason for the last 'par' is so that the strategy terminates
-- quickly. This is important if the strategy is used as the 1st
-- argument of a seq

-- | Apply three strategies to the elements of a triple in sequentially
--   from left to right.
seqTriple :: Strategy a -> Strategy b -> Strategy c -> Strategy (a,b,c)
seqTriple strata stratb stratc p@(x,y,z) =
  strata x `seq`
  stratb y `seq`
  stratc z

-- | Apply three strategies to the elements of a triple in parallel.
parTriple :: Strategy a -> Strategy b -> Strategy c -> Strategy (a,b,c)
parTriple strata stratb stratc (x,y,z) =
  strata x `par`
  stratb y `par`
  stratc z `par`
  ()

-----------------------------------------------------------------------------
-- 			Atomic types
-----------------------------------------------------------------------------

{-
Weak head normal form and normal form are identical for integers, so the
default rnf is sufficient.
-}
instance NFData Int
instance NFData Integer
instance NFData Float
instance NFData Double

instance NFData Int8
instance NFData Int16
instance NFData Int32
instance NFData Int64

instance NFData Word8
instance NFData Word16
instance NFData Word32
instance NFData Word64

instance NFDataIntegral Int
instance NFDataOrd Int

--Rational and complex numbers.

instance (Integral a, NFData a) => NFData (Ratio a) where
  rnf (x:%y) = rnf x `seq`
               rnf y `seq`
               ()

instance (RealFloat a, NFData a) => NFData (Complex a) where
  rnf (x:+y) = rnf x `seq`
	         rnf y `seq`
               ()

instance NFData Char
instance NFData Bool
instance NFData ()

-----------------------------------------------------------------------------
-- 			Various library types
-----------------------------------------------------------------------------

instance NFData a => NFData (Maybe a) where
    rnf Nothing  = ()
    rnf (Just x) = rnf x

instance (NFData a, NFData b) => NFData (Either a b) where
    rnf (Left x)  = rnf x
    rnf (Right y) = rnf y

instance (NFData k, NFData a) => NFData (Data.Map.Map k a) where
    rnf = rnf . Data.Map.toList

instance NFData a => NFData (Data.Set.Set a) where
    rnf = rnf . Data.Set.toList

instance NFData a => NFData (Data.Tree.Tree a) where
    rnf (Data.Tree.Node r f) = rnf r `seq` rnf f

instance NFData a => NFData (Data.IntMap.IntMap a) where
    rnf = rnf . Data.IntMap.toList

instance NFData Data.IntSet.IntSet where
    rnf = rnf . Data.IntSet.toList

-----------------------------------------------------------------------------
-- 			Lists
-----------------------------------------------------------------------------

instance NFData a => NFData [a] where
  rnf [] = ()
  rnf (x:xs) = rnf x `seq` rnf xs

----------------------------------------------------------------------------
-- *                   Lists: Parallel Strategies
----------------------------------------------------------------------------

-- | Applies a strategy to every element of a list in parallel.
parList :: Strategy a -> Strategy [a]
parList strat []     = ()
parList strat (x:xs) = strat x `par` (parList strat xs)

-- | Applies a strategy to the first @n@ elements of a list in parallel.
parListN :: (Integral b) => b -> Strategy a -> Strategy [a]
parListN n strat []     = ()
parListN 0 strat xs     = ()
parListN n strat (x:xs) = strat x `par` (parListN (n-1) strat xs)

-- | Evaluates @n@ elements of the spine of the argument list and applies
-- the given strategy to the @n@th element (if there is one) in parallel with
-- the result. E.g. @parListNth 2 [e1, e2, e3]@ evaluates @e3@.
parListNth :: Int -> Strategy a -> Strategy [a]
parListNth n strat xs
  | null rest = ()
  | otherwise = strat (head rest) `par` ()
  where
    rest = drop n xs

-- | Splits a list into chunks (sub-sequences) of length @n@,
-- and applies a strategy sequentially to the elements in each
-- chunk. The chunks are evaluated in parallel.
-- This is useful for increasing the grain size.
parListChunk :: Int -> Strategy a -> Strategy [a]
parListChunk n strat [] = ()
parListChunk n strat xs = seqListN n strat xs `par`
			    parListChunk n strat (drop n xs)

-- | Applies a function to each element of a list and
-- and evaluates the result list in parallel,
-- using the given strategy for each element.
parMap :: Strategy b -> (a -> b) -> [a] -> [b]
parMap strat f xs = map f xs `using` parList strat

-- | Uses 'parMap' to apply a list-valued function to each
-- element of a list in parallel, and concatenates the results.
parFlatMap :: Strategy [b] -> (a -> [b]) -> [a] -> [b]
parFlatMap strat f xs = concat (parMap strat f xs)

-- | Zips together two lists using a function,
-- and evaluates the result list in parallel.
parZipWith :: Strategy c -> (a -> b -> c) -> [a] -> [b] -> [c]
parZipWith strat z as bs =
  zipWith z as bs `using` parList strat

----------------------------------------------------------------------------
-- *                     Lists: Sequential Strategies
----------------------------------------------------------------------------

-- | Sequentially applies a strategy to each element of a list.
seqList :: Strategy a -> Strategy [a]
seqList strat []     = ()
seqList strat (x:xs) = strat x `seq` (seqList strat xs)

-- | Sequentially applies a strategy to the first n elements of a list.
seqListN :: (Integral a) => a -> Strategy b -> Strategy [b]
seqListN n strat []     = ()
seqListN 0 strat xs     = ()
seqListN n strat (x:xs) = strat x `seq` (seqListN (n-1) strat xs)

-- | Applies a strategy to the @n@th element of a list
--  (if there is one) before returning the result.
--  E.g. @seqListNth 2 [e1, e2, e3]@ evaluates @e3@.
seqListNth :: Int -> Strategy b -> Strategy [b]
seqListNth n strat xs
  | null rest = ()
  | otherwise = strat (head rest)
  where
    rest = drop n xs

-- | Parallel n-buffer function added for the revised version of the strategies
-- paper. 'parBuffer' supersedes the older @fringeList@. It has the same
-- semantics.
parBuffer :: Int -> Strategy a -> [a] -> [a]
parBuffer n s xs =
  return xs (start n xs)
  where
    return (x:xs) (y:ys) = (x:return xs ys) `sparking` s y
    return xs     []     = xs

    start n []     = []
    start 0 ys     = ys
    start n (y:ys) = start (n-1) ys `sparking` s y

{-
 'fringeList' implements a `rolling buffer' of length n, i.e.applies a
 strategy to the nth element of list when the head is demanded. More
 precisely:

   semantics:         fringeList n s = id :: [b] -> [b]
   dynamic behaviour: evalutates the nth element of the list when the
		      head is demanded.

 The idea is to provide a `rolling buffer' of length n.
fringeList :: (Integral a) => a -> Strategy b -> [b] -> [b]
fringeList n strat [] = []
fringeList n strat (r:rs) =
  seqListNth n strat rs `par`
  r:fringeList n strat rs
-}

------------------------------------------------------------------------------
-- *			Arrays
------------------------------------------------------------------------------
instance (Ix a, NFData a, NFData b) => NFData (Array a b) where
  rnf x = rnf (bounds x) `seq` seqList rnf (elems x) `seq` ()

-- | Apply a strategy to all elements of an array sequentially.
seqArr :: (Ix b) => Strategy a -> Strategy (Array b a)
seqArr s arr = seqList s (elems arr)

-- | Apply a strategy to all elements of an array in parallel.
parArr :: (Ix b) => Strategy a -> Strategy (Array b a)
parArr s arr = parList s (elems arr)

{-# DEPRECATED Assoc "Does not belong in Control.Parallel.Strategies" #-}
data  Assoc a b =  a := b  deriving ()

instance (NFData a, NFData b) => NFData (Assoc a b) where
  rnf (x := y) = rnf x `seq` rnf y `seq` ()

------------------------------------------------------------------------------
-- *	                Some strategies specific for Lolita
------------------------------------------------------------------------------

{-# DEPRECATED fstPairFstList "This was just an example. Write your own." #-}
fstPairFstList :: (NFData a) => Strategy [(a,b)]
fstPairFstList = seqListN 1 (seqPair rwhnf r0)

-- Some HACKs for Lolita. AFAIK force is just another name for our rnf and
-- sforce is a shortcut (definition here is identical to the one in Force.lhs)

{-# DEPRECATED force, sforce "Lolita-specific hacks." #-}
force :: (NFData a) => a -> a
sforce :: (NFData a) => a -> b -> b

force = id $| rnf
sforce x y = force x `seq` y


{-# LANGUAGE CPP #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE Rank2Types #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE DefaultSignatures #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE UndecidableInstances #-}
{-# LANGUAGE TypeOperators #-}

#ifdef TRUSTWORTHY
{-# LANGUAGE Trustworthy #-}
#endif

#include "lens-common.h"

-----------------------------------------------------------------------------
-- |
-- Module      :  Control.Lens.At
-- Copyright   :  (C) 2012-16 Edward Kmett
-- License     :  BSD-style (see the file LICENSE)
-- Maintainer  :  Edward Kmett <ekmett@gmail.com>
-- Stability   :  experimental
-- Portability :  non-portable
--
----------------------------------------------------------------------------
module Control.Lens.At
  (
  -- * At
    At(at)
    , sans
    , iat
  -- * Ixed
  , Index
  , IxValue
  , Ixed(ix)
  , ixAt
  , iix
  -- * Contains
  , Contains(contains)
  , icontains
  ) where

import Prelude ()

import Control.Lens.Each
import Control.Lens.Internal.Prelude
import Control.Lens.Traversal
import Control.Lens.Lens
import Control.Lens.Setter
import Control.Lens.Indexed
import Control.Monad (guard)
import Data.Array.IArray as Array
import Data.Array.Unboxed
import qualified Data.ByteString as StrictB
import qualified Data.ByteString.Lazy as LazyB
import Data.Complex
import Data.Functor (($>))
import Data.Hashable
import qualified Data.HashMap.Lazy as HashMap
import Data.HashMap.Lazy (HashMap)
import qualified Data.HashSet as HashSet
import Data.HashSet (HashSet)
import Data.Int
import qualified Data.IntMap as IntMap
import Data.IntMap (IntMap)
import qualified Data.IntSet as IntSet
import Data.IntSet (IntSet)
import qualified Data.Map as Map
import Data.Map (Map)
import Data.Maybe (isJust)
import qualified Data.Set as Set
import Data.Set (Set)
import qualified Data.Sequence as Seq
import Data.Sequence (Seq)
import qualified Data.Text as StrictT
import qualified Data.Text.Lazy as LazyT
import Data.Tree
import qualified Data.Vector as Vector
import qualified Data.Vector.Primitive as Prim
import Data.Vector.Primitive (Prim)
import qualified Data.Vector.Storable as Storable
import qualified Data.Vector.Unboxed as Unboxed
import Data.Vector.Unboxed (Unbox)
import Data.Word
import Foreign.Storable (Storable)

type family Index (s :: *) :: *
type instance Index (e -> a) = e
type instance Index IntSet = Int
type instance Index (Set a) = a
type instance Index (HashSet a) = a
type instance Index [a] = Int
type instance Index (NonEmpty a) = Int
type instance Index (Seq a) = Int
type instance Index (a,b) = Int
type instance Index (a,b,c) = Int
type instance Index (a,b,c,d) = Int
type instance Index (a,b,c,d,e) = Int
type instance Index (a,b,c,d,e,f) = Int
type instance Index (a,b,c,d,e,f,g) = Int
type instance Index (a,b,c,d,e,f,g,h) = Int
type instance Index (a,b,c,d,e,f,g,h,i) = Int
type instance Index (IntMap a) = Int
type instance Index (Map k a) = k
type instance Index (HashMap k a) = k
type instance Index (Array.Array i e) = i
type instance Index (UArray i e) = i
type instance Index (Vector.Vector a) = Int
type instance Index (Prim.Vector a) = Int
type instance Index (Storable.Vector a) = Int
type instance Index (Unboxed.Vector a) = Int
type instance Index (Complex a) = Int
type instance Index (Identity a) = ()
type instance Index (Maybe a) = ()
type instance Index (Tree a) = [Int]
type instance Index StrictT.Text = Int
type instance Index LazyT.Text = Int64
type instance Index StrictB.ByteString = Int
type instance Index LazyB.ByteString = Int64

-- $setup
-- >>> :set -XNoOverloadedStrings
-- >>> import Control.Lens
-- >>> import qualified Data.IntSet as IntSet
-- >>> import qualified Data.Sequence as Seq
-- >>> import qualified Data.Map as Map
-- >>> import Debug.SimpleReflect.Expr
-- >>> import Debug.SimpleReflect.Vars as Vars hiding (f,g)
-- >>> let f  :: Expr -> Expr; f = Debug.SimpleReflect.Vars.f
-- >>> let g  :: Expr -> Expr; g = Debug.SimpleReflect.Vars.g
-- >>> let f' :: Int -> Expr -> Expr; f' = Debug.SimpleReflect.Vars.f'
-- >>> let h  :: Int -> Expr; h = Debug.SimpleReflect.Vars.h

-- |
-- This class provides a simple 'Lens' that lets you view (and modify)
-- information about whether or not a container contains a given 'Index'.
class Contains m where
  -- |
  -- >>> IntSet.fromList [1,2,3,4] ^. contains 3
  -- True
  --
  -- >>> IntSet.fromList [1,2,3,4] ^. contains 5
  -- False
  --
  -- >>> IntSet.fromList [1,2,3,4] & contains 3 .~ False
  -- fromList [1,2,4]
  contains :: Index m -> Lens' m Bool

-- | An indexed version of 'contains'.
--
-- >>> IntSet.fromList [1,2,3,4] ^@. icontains 3
-- (3,True)
--
-- >>> IntSet.fromList [1,2,3,4] ^@. icontains 5
-- (5,False)
--
-- >>> IntSet.fromList [1,2,3,4] & icontains 3 %@~ \i x -> if odd i then not x else x
-- fromList [1,2,4]
--
-- >>> IntSet.fromList [1,2,3,4] & icontains 3 %@~ \i x -> if even i then not x else x
-- fromList [1,2,3,4]
icontains :: Contains m => Index m -> IndexedLens' (Index m) m Bool
icontains i f = contains i (indexed f i)
{-# INLINE icontains #-}

instance Contains IntSet where
#if MIN_VERSION_containers(0,6,3)
  contains k f = IntSet.alterF f k
#else
  -- This is a flipped copy of the implementation of `IntSet.alterF`.  Unlike a
  -- `Set`, we don't have to worry about expensive comparisons from descending
  -- multiple times into an `IntSet`. We are careful to share the results of
  -- insertion or deletion across multiple positions in the `Functor`.
  contains k f s = fmap choose (f member_)
    where
      member_ = IntSet.member k s

      (inserted, deleted)
        | member_   = (s, IntSet.delete k s)
        | otherwise = (IntSet.insert k s, s)

      choose True  = inserted
      choose False = deleted
#endif
  {-# INLINE contains #-}

instance Ord a => Contains (Set a) where
#if MIN_VERSION_containers(0,6,3)
  contains k f = Set.alterF f k
#else
  contains k f s = f (Set.member k s) <&> \b ->
    if b then Set.insert k s else Set.delete k s
#endif
  {-# INLINE contains #-}

instance (Eq a, Hashable a) => Contains (HashSet a) where
  contains k f s = HashSet.fromMap <$>
    HashMap.alterF (fmap guard . f . isJust) k (HashSet.toMap s)
  {-# INLINE contains #-}

-- | This provides a common notion of a value at an index that is shared by both 'Ixed' and 'At'.
type family IxValue (m :: *) :: *

-- | Provides a simple 'Traversal' lets you 'traverse' the value at a given
-- key in a 'Map' or element at an ordinal position in a list or 'Seq'.
class Ixed m where
  -- |
  -- /NB:/ Setting the value of this 'Traversal' will only set the value in
  -- 'at' if it is already present.
  --
  -- If you want to be able to insert /missing/ values, you want 'at'.
  --
  -- >>> Seq.fromList [a,b,c,d] & ix 2 %~ f
  -- fromList [a,b,f c,d]
  --
  -- >>> Seq.fromList [a,b,c,d] & ix 2 .~ e
  -- fromList [a,b,e,d]
  --
  -- >>> Seq.fromList [a,b,c,d] ^? ix 2
  -- Just c
  --
  -- >>> Seq.fromList [] ^? ix 2
  -- Nothing
  ix :: Index m -> Traversal' m (IxValue m)
  default ix :: At m => Index m -> Traversal' m (IxValue m)
  ix = ixAt
  {-# INLINE ix #-}

-- | An indexed version of 'ix'.
--
-- >>> Seq.fromList [a,b,c,d] & iix 2 %@~ f'
-- fromList [a,b,f' 2 c,d]
--
-- >>> Seq.fromList [a,b,c,d] & iix 2 .@~ h
-- fromList [a,b,h 2,d]
--
-- >>> Seq.fromList [a,b,c,d] ^@? iix 2
-- Just (2,c)
--
-- >>> Seq.fromList [] ^@? iix 2
-- Nothing
iix :: Ixed m => Index m -> IndexedTraversal' (Index m) m (IxValue m)
iix i f = ix i (indexed f i)
{-# INLINE iix #-}

-- | A definition of 'ix' for types with an 'At' instance. This is the default
-- if you don't specify a definition for 'ix'.
ixAt :: At m => Index m -> Traversal' m (IxValue m)
ixAt i = at i . traverse
{-# INLINE ixAt #-}

type instance IxValue (e -> a) = a
instance Eq e => Ixed (e -> a) where
  ix e p f = p (f e) <&> \a e' -> if e == e' then a else f e'
  {-# INLINE ix #-}

type instance IxValue (Maybe a) = a
instance Ixed (Maybe a) where
  ix ~() f (Just a) = Just <$> f a
  ix ~() _ Nothing  = pure Nothing
  {-# INLINE ix #-}

type instance IxValue [a] = a
instance Ixed [a] where
  ix k f xs0 | k < 0     = pure xs0
             | otherwise = go xs0 k where
    go [] _ = pure []
    go (a:as) 0 = f a <&> (:as)
    go (a:as) i = (a:) <$> (go as $! i - 1)
  {-# INLINE ix #-}

type instance IxValue (NonEmpty a) = a
instance Ixed (NonEmpty a) where
  ix k f xs0 | k < 0 = pure xs0
             | otherwise = go xs0 k where
    go (a:|as) 0 = f a <&> (:|as)
    go (a:|as) i = (a:|) <$> ix (i - 1) f as
  {-# INLINE ix #-}

type instance IxValue (Identity a) = a
instance Ixed (Identity a) where
  ix ~() f (Identity a) = Identity <$> f a
  {-# INLINE ix #-}

type instance IxValue (Tree a) = a
instance Ixed (Tree a) where
  ix xs0 f = go xs0 where
    go [] (Node a as) = f a <&> \a' -> Node a' as
    go (i:is) t@(Node a as)
      | i < 0     = pure t
      | otherwise = Node a <$> ix i (go is) as
  {-# INLINE ix #-}

type instance IxValue (Seq a) = a
instance Ixed (Seq a) where
  ix i f m
    | 0 <= i && i < Seq.length m = f (Seq.index m i) <&> \a -> Seq.update i a m
    | otherwise                  = pure m
  {-# INLINE ix #-}

type instance IxValue (IntMap a) = a
instance Ixed (IntMap a) where
  ix k f m = case IntMap.lookup k m of
     Just v -> f v <&> \v' -> IntMap.insert k v' m
     Nothing -> pure m
  {-# INLINE ix #-}

type instance IxValue (Map k a) = a
instance Ord k => Ixed (Map k a) where
  ix k f m = case Map.lookup k m of
     Just v  -> f v <&> \v' -> Map.insert k v' m
     Nothing -> pure m
  {-# INLINE ix #-}

type instance IxValue (HashMap k a) = a
instance (Eq k, Hashable k) => Ixed (HashMap k a) where
  ix k f m = case HashMap.lookup k m of
     Just v  -> f v <&> \v' -> HashMap.insert k v' m
     Nothing -> pure m
  {-# INLINE ix #-}

type instance IxValue (Set k) = ()
instance Ord k => Ixed (Set k) where
  ix k f m = if Set.member k m
     then f () $> m
     else pure m
  {-# INLINE ix #-}

type instance IxValue IntSet = ()
instance Ixed IntSet where
  ix k f m = if IntSet.member k m
     then f () $> m
     else pure m
  {-# INLINE ix #-}

type instance IxValue (HashSet k) = ()
instance (Eq k, Hashable k) => Ixed (HashSet k) where
  ix k f m = if HashSet.member k m
     then f () $> m
     else pure m
  {-# INLINE ix #-}

type instance IxValue (Array.Array i e) = e
-- |
-- @
-- arr '!' i ≡ arr 'Control.Lens.Getter.^.' 'ix' i
-- arr '//' [(i,e)] ≡ 'ix' i 'Control.Lens.Setter..~' e '$' arr
-- @
instance Ix i => Ixed (Array.Array i e) where
  ix i f arr
    | inRange (bounds arr) i = f (arr Array.! i) <&> \e -> arr Array.// [(i,e)]
    | otherwise              = pure arr
  {-# INLINE ix #-}

type instance IxValue (UArray i e) = e
-- |
-- @
-- arr '!' i ≡ arr 'Control.Lens.Getter.^.' 'ix' i
-- arr '//' [(i,e)] ≡ 'ix' i 'Control.Lens.Setter..~' e '$' arr
-- @
instance (IArray UArray e, Ix i) => Ixed (UArray i e) where
  ix i f arr
    | inRange (bounds arr) i = f (arr Array.! i) <&> \e -> arr Array.// [(i,e)]
    | otherwise              = pure arr
  {-# INLINE ix #-}

type instance IxValue (Vector.Vector a) = a
instance Ixed (Vector.Vector a) where
  ix i f v
    | 0 <= i && i < Vector.length v = f (v Vector.! i) <&> \a -> v Vector.// [(i, a)]
    | otherwise                     = pure v
  {-# INLINE ix #-}

type instance IxValue (Prim.Vector a) = a
instance Prim a => Ixed (Prim.Vector a) where
  ix i f v
    | 0 <= i && i < Prim.length v = f (v Prim.! i) <&> \a -> v Prim.// [(i, a)]
    | otherwise                   = pure v
  {-# INLINE ix #-}

type instance IxValue (Storable.Vector a) = a
instance Storable a => Ixed (Storable.Vector a) where
  ix i f v
    | 0 <= i && i < Storable.length v = f (v Storable.! i) <&> \a -> v Storable.// [(i, a)]
    | otherwise                       = pure v
  {-# INLINE ix #-}

type instance IxValue (Unboxed.Vector a) = a
instance Unbox a => Ixed (Unboxed.Vector a) where
  ix i f v
    | 0 <= i && i < Unboxed.length v = f (v Unboxed.! i) <&> \a -> v Unboxed.// [(i, a)]
    | otherwise                      = pure v
  {-# INLINE ix #-}

type instance IxValue StrictT.Text = Char
instance Ixed StrictT.Text where
  ix e f s = case StrictT.splitAt e s of
     (l, mr) -> case StrictT.uncons mr of
       Nothing      -> pure s
       Just (c, xs) -> f c <&> \d -> StrictT.concat [l, StrictT.singleton d, xs]
  {-# INLINE ix #-}

type instance IxValue LazyT.Text = Char
instance Ixed LazyT.Text where
  ix e f s = case LazyT.splitAt e s of
     (l, mr) -> case LazyT.uncons mr of
       Nothing      -> pure s
       Just (c, xs) -> f c <&> \d -> LazyT.append l (LazyT.cons d xs)
  {-# INLINE ix #-}

type instance IxValue StrictB.ByteString = Word8
instance Ixed StrictB.ByteString where
  ix e f s = case StrictB.splitAt e s of
     (l, mr) -> case StrictB.uncons mr of
       Nothing      -> pure s
       Just (c, xs) -> f c <&> \d -> StrictB.concat [l, StrictB.singleton d, xs]
  {-# INLINE ix #-}

type instance IxValue LazyB.ByteString = Word8
instance Ixed LazyB.ByteString where
  -- TODO: we could be lazier, returning each chunk as it is passed
  ix e f s = case LazyB.splitAt e s of
     (l, mr) -> case LazyB.uncons mr of
       Nothing      -> pure s
       Just (c, xs) -> f c <&> \d -> LazyB.append l (LazyB.cons d xs)
  {-# INLINE ix #-}

-- | 'At' provides a 'Lens' that can be used to read,
-- write or delete the value associated with a key in a 'Map'-like
-- container on an ad hoc basis.
--
-- An instance of 'At' should satisfy:
--
-- @
-- 'ix' k ≡ 'at' k '.' 'traverse'
-- @
class Ixed m => At m where
  -- |
  -- >>> Map.fromList [(1,"world")] ^.at 1
  -- Just "world"
  --
  -- >>> at 1 ?~ "hello" $ Map.empty
  -- fromList [(1,"hello")]
  --
  -- /Note:/ 'Map'-like containers form a reasonable instance, but not 'Array'-like ones, where
  -- you cannot satisfy the 'Lens' laws.
  at :: Index m -> Lens' m (Maybe (IxValue m))

-- | Delete the value associated with a key in a 'Map'-like container
--
-- @
-- 'sans' k = 'at' k .~ Nothing
-- @
sans :: At m => Index m -> m -> m
sans k m = m & at k .~ Nothing
{-# INLINE sans #-}

-- | An indexed version of 'at'.
--
-- >>> Map.fromList [(1,"world")] ^@. iat 1
-- (1,Just "world")
--
-- >>> iat 1 %@~ (\i x -> if odd i then Just "hello" else Nothing) $ Map.empty
-- fromList [(1,"hello")]
--
-- >>> iat 2 %@~ (\i x -> if odd i then Just "hello" else Nothing) $ Map.empty
-- fromList []
--
iat :: At m => Index m -> IndexedLens' (Index m) m (Maybe (IxValue m))
iat i f = at i (indexed f i)
{-# INLINE iat #-}

instance At (Maybe a) where
  at ~() f = f
  {-# INLINE at #-}

instance At (IntMap a) where
#if MIN_VERSION_containers(0,5,8)
  at k f = IntMap.alterF f k
#else
  at k f m = f mv <&> \r -> case r of
    Nothing -> maybe m (const (IntMap.delete k m)) mv
    Just v' -> IntMap.insert k v' m
    where mv = IntMap.lookup k m
#endif
  {-# INLINE at #-}

instance Ord k => At (Map k a) where
#if MIN_VERSION_containers(0,5,8)
  at k f = Map.alterF f k
#else
  at k f m = f mv <&> \r -> case r of
    Nothing -> maybe m (const (Map.delete k m)) mv
    Just v' -> Map.insert k v' m
    where mv = Map.lookup k m
#endif
  {-# INLINE at #-}

instance (Eq k, Hashable k) => At (HashMap k a) where
  at k f = HashMap.alterF f k
  {-# INLINE at #-}

instance At IntSet where
  -- This is a gently modified copy of the implementation of `IntSet.alterF`.
  -- Unlike a `Set`, we don't have to worry about expensive comparisons from
  -- descending multiple times into an `IntSet`. We are careful to share the
  -- results of insertion or deletion across multiple positions in the
  -- `Functor`.
  at k f s = fmap choose (f (guard member_))
    where
      member_ = IntSet.member k s

      (inserted, deleted)
        | member_   = (s, IntSet.delete k s)
        | otherwise = (IntSet.insert k s, s)

      choose (Just ~()) = inserted
      choose Nothing = deleted
  {-# INLINE at #-}

instance Ord k => At (Set k) where
#if MIN_VERSION_containers(0,6,3)
  at k f = Set.alterF (fmap isJust . f . guard) k
#else
  at k f m = f mv <&> \r -> case r of
    Nothing -> maybe m (const (Set.delete k m)) mv
    Just ~() -> maybe (Set.insert k m) (const m) mv
    where mv = if Set.member k m then Just () else Nothing
#endif
  {-# INLINE at #-}

instance (Eq k, Hashable k) => At (HashSet k) where
  at k f s = HashSet.fromMap <$> HashMap.alterF f k (HashSet.toMap s)
  {-# INLINE at #-}

-- | @'ix' :: 'Int' -> 'Traversal'' (a,a) a@
type instance IxValue (a,a2) = a
instance (a~a2) => Ixed (a,a2) where
  ix p = elementOf each p

-- | @'ix' :: 'Int' -> 'Traversal'' (a,a,a) a@
type instance IxValue (a,a2,a3) = a
instance (a~a2, a~a3) => Ixed (a,a2,a3) where
  ix p = elementOf each p

-- | @'ix' :: 'Int' -> 'Traversal'' (a,a,a,a) a@
type instance IxValue (a,a2,a3,a4) = a
instance (a~a2, a~a3, a~a4) => Ixed (a,a2,a3,a4) where
  ix p = elementOf each p

-- | @'ix' :: 'Int' -> 'Traversal'' (a,a,a,a,a) a@
type instance IxValue (a,a2,a3,a4,a5) = a
instance (a~a2, a~a3, a~a4, a~a5) => Ixed (a,a2,a3,a4,a5) where
  ix p = elementOf each p

-- | @'ix' :: 'Int' -> 'Traversal'' (a,a,a,a,a,a) a@
type instance IxValue (a,a2,a3,a4,a5,a6) = a
instance (a~a2, a~a3, a~a4, a~a5, a~a6) => Ixed (a,a2,a3,a4,a5,a6) where
  ix p = elementOf each p

-- | @'ix' :: 'Int' -> 'Traversal'' (a,a,a,a,a,a,a) a@
type instance IxValue (a,a2,a3,a4,a5,a6,a7) = a
instance (a~a2, a~a3, a~a4, a~a5, a~a6, a~a7) => Ixed (a,a2,a3,a4,a5,a6,a7) where
  ix p = elementOf each p

-- | @'ix' :: 'Int' -> 'Traversal'' (a,a,a,a,a,a,a,a) a@
type instance IxValue (a,a2,a3,a4,a5,a6,a7,a8) = a
instance (a~a2, a~a3, a~a4, a~a5, a~a6, a~a7, a~a8) => Ixed (a,a2,a3,a4,a5,a6,a7,a8) where
  ix p = elementOf each p

-- | @'ix' :: 'Int' -> 'Traversal'' (a,a,a,a,a,a,a,a,a) a@
type instance IxValue (a,a2,a3,a4,a5,a6,a7,a8,a9) = a
instance (a~a2, a~a3, a~a4, a~a5, a~a6, a~a7, a~a8, a~a9) => Ixed (a,a2,a3,a4,a5,a6,a7,a8,a9) where
  ix p = elementOf each p


-- |
-- Module      : Occlusion.Window
-- Description :
-- Copyright   : (c) Jonatan H Sundqvist, 2015
-- License     : MIT
-- Maintainer  : Jonatan H Sundqvist
-- Stability   : experimental|stable
-- Portability : POSIX (not sure)
--

-- Created September 20 2015

-- TODO | -
--        -

-- SPEC | -
--        -

--------------------------------------------------------------------------------------------------------------------------------------------
-- GHC Pragmas
--------------------------------------------------------------------------------------------------------------------------------------------

--------------------------------------------------------------------------------------------------------------------------------------------
-- API
--------------------------------------------------------------------------------------------------------------------------------------------
module Occlusion.Window where

--------------------------------------------------------------------------------------------------------------------------------------------
-- We'll need these
--------------------------------------------------------------------------------------------------------------------------------------------
import Data.Complex

import Graphics.UI.Gtk
import qualified Graphics.Rendering.Cairo as Cairo

import Occlusion.Types
import qualified Occlusion.Core   as Core
import qualified Occlusion.Events as Events

--------------------------------------------------------------------------------------------------------------------------------------------
-- Functions
--------------------------------------------------------------------------------------------------------------------------------------------
-- |
create :: Complex Double -> IO (Window, DrawingArea)
create (winx:+winy) = do
  initGUI
  window <- windowNew
  frame  <- frameNew

  set window [windowTitle := "Occlusion"]

  canvas <- drawingAreaNew
  containerAdd frame canvas
  set window [ containerChild := frame ]
  windowSetDefaultSize window (round winx) (round winy)
  windowSetIconFromFile window "C:/Users/Jonatan/Desktop/Haskell/projects/Occlusion/assets/images/TreeAtDusk.png"

  widgetAddEvents canvas [PointerMotionMask, ButtonMotionMask, ButtonPressMask] -- MouseButton1Mask
  widgetShowAll window

  return (window, canvas)