{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE AllowAmbiguousTypes #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE FunctionalDependencies #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE InstanceSigs #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE StandaloneDeriving #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE TemplateHaskell #-}
{-# LANGUAGE DeriveGeneric #-}
{-# LANGUAGE UndecidableInstances #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE PolyKinds #-}
{-# LANGUAGE PartialTypeSignatures #-}
{-# OPTIONS_GHC -Wno-partial-type-signatures #-}
{-# LANGUAGE StandaloneKindSignatures #-}
module Inference.Conjugate where

import           Control.Monad.Primitive        ( PrimMonad
                                                , PrimState
                                                )
import           Control.Monad.Reader           ( ReaderT
                                                , runReaderT
                                                )
import           Control.Monad.Reader.Class     ( MonadReader(ask) )
import           Control.Monad.State            ( State
                                                , StateT
                                                , evalState
                                                , evalStateT
                                                , execStateT
                                                , runStateT
                                                )
import           Control.Monad.State.Class      ( get
                                                , modify
                                                , put
                                                )
import           Control.Monad.Trans.Maybe      ( MaybeT(MaybeT)
                                                , runMaybeT
                                                )
import           Control.Monad.Writer
import           Data.Dynamic                   ( Dynamic
                                                , Typeable
                                                , fromDynamic
                                                , toDyn
                                                )
import           Data.Kind
import           Data.Maybe                     ( fromMaybe )
import           Data.MultiSet                 as MS
import qualified Data.Sequence                 as S
import           Data.Typeable                  ( Proxy(Proxy)
                                                , typeRep
                                                )
import qualified Data.Vector                   as V
import qualified Debug.Trace                   as DT
import           GHC.Float                      ( int2Double )
import           GHC.Generics
import           GHC.TypeNats
import           Lens.Micro
import           Lens.Micro.Extras
import           Lens.Micro.TH                  ( makeLenses )
import           Numeric.SpecFunctions          ( logBeta
                                                , logChoose
                                                , logFactorial
                                                , logGamma
                                                )
import           System.Random.MWC.Probability
                                         hiding ( Uniform )

----------------------------------------------------------------
-- type classes and families for distributions and conjugates --
----------------------------------------------------------------

-- | Describes a family of distributions with a fixed form.
-- For example, a 'Bernoulli' distribution is parameterized by a probability @p@
-- and produces binary samples
-- (@True@ with probability @p@, @False@ with probability @1-p@).
--
-- Its 'Distribution' instance is:
-- > instance Distribution Bernoulli where
-- >   type Params Bernoulli = Double
-- >   type Support Bernoulli = Bool
-- >   distSample _ = uncurry bernoulli
-- >   distLogP _ p True = log p
-- >   distLogP _ p False = log (1 - p)
class Distribution a where
  type Params a :: Type
  type Support a :: Type
  distSample :: (PrimMonad m) => a -> Params a -> Prob m (Support a)
  distLogP :: a -> Params a -> Support a -> Double

-- -- | Used as a type-lifted kind for conjugate distribution pairs.
-- data Conj a b = Conj a bb

-- | A type-level marker for treating a distribution as a prior.
newtype AsPrior p = AsPrior p

-- -- | A type-level marker for treating a distribution as a likelihood.
-- newtype AsLk l = AsLk l

-- | Marks two distributions as a conjugate pair of prior and likelihood.
-- The property of such a pair is that the posterior has the same form as the prior
-- (including the same 'Params' and 'Support'),
-- and that its parameters can be obtained analytically from the parameters of the prior
-- and a set of observations.
--
-- The class method 'updatePrior' returns the parameters of the posterior
-- given the prior parameters after a single observation.
class (Distribution p, Distribution l, Support p ~ Params l) => Conjugate p l where
  priorSingleton :: p -- ^ provides a singleton instance of the prior distribution
                      -- in order to make sampling from priors easier.
  updatePrior :: l -> Params p -> Support l -> Params p
  predLogP :: l -> Params p -> Support l -> Double

type family Hyper (a :: k) :: Type
type instance Hyper (AsPrior p) = Params p

type family Probs (a :: k) :: Type
type instance Probs (AsPrior p) = Support p

-- type family Value (a :: k) :: Type
-- type instance Value (Conj p l) = Support l

-- helper types for instantiating hyperparameters, parameters, and values
-- ----------------------------------------------------------------------

newtype HyperRep p = HyperRep { runHyper :: Hyper (AsPrior p) }
deriving instance Show (Hyper (AsPrior p)) => Show (HyperRep p)

type instance Hyper (a :: (Type -> Type) -> Type) = a HyperRep

newtype ProbsRep p = ProbsRep { runProbs :: Probs (AsPrior p)}
deriving instance Show (Probs (AsPrior p)) => Show (ProbsRep p)

type instance Probs (a :: (Type -> Type) -> Type) = a ProbsRep

-- newtype ValueRep p l = ValueRep { runValue :: Value (Conj p l) }
-- deriving instance Show (Value (Conj p l)) => Show (ValueRep p l)

-- type instance Value (a :: (Type -> Type -> Type) -> Type) = a ValueRep

-----------------------------------------------------
-- generic magic for conjugates and parameter sets --
-----------------------------------------------------

-- Jeffreys prior
-- ---------------

class Jeffreys a where
  jeffreysPrior :: Hyper a

class GJeffreys t where
  gjeffreysPrior :: forall p. t p

instance GJeffreys V1 where
  gjeffreysPrior = undefined -- ok

instance GJeffreys U1 where
  gjeffreysPrior = U1

-- base case: k is a conjugate distribution
instance (Jeffreys (AsPrior p)) => GJeffreys (K1 i (HyperRep p)) where
  gjeffreysPrior = K1 $ HyperRep $ jeffreysPrior @(AsPrior p)

-- recursive case: k is another record
instance (Jeffreys k, k HyperRep ~ Hyper k) => GJeffreys (K1 i (k HyperRep)) where
  gjeffreysPrior = K1 $ jeffreysPrior @k

instance (GJeffreys t) => GJeffreys (M1 i c (t :: Type -> Type)) where
  gjeffreysPrior = M1 gjeffreysPrior

instance (GJeffreys ta, GJeffreys tb) => GJeffreys (ta :*: tb) where
  gjeffreysPrior = gjeffreysPrior @ta :*: gjeffreysPrior @tb

instance (Generic (t HyperRep), GJeffreys (Rep (t HyperRep))) => Jeffreys (t :: (Type -> Type) -> Type) where
  jeffreysPrior = GHC.Generics.to (gjeffreysPrior @(Rep (t HyperRep)))

-- uniform prior
-- -------------

class Uniform a where
  uniformPrior :: Hyper a

class GUniform t where
  guniformPrior :: forall p. t p

instance GUniform V1 where
  guniformPrior = undefined -- ok

instance GUniform U1 where
  guniformPrior = U1

-- base case: k is a conjugate distribution
instance (Uniform (AsPrior p)) => GUniform (K1 i (HyperRep p)) where
  guniformPrior = K1 $ HyperRep $ uniformPrior @(AsPrior p)

-- recursive case: k is another record
instance (Uniform k, k HyperRep ~ Hyper k) => GUniform (K1 i (k HyperRep)) where
  guniformPrior = K1 $ uniformPrior @k

instance (GUniform t) => GUniform (M1 i c (t :: Type -> Type)) where
  guniformPrior = M1 guniformPrior

instance (GUniform ta, GUniform tb) => GUniform (ta :*: tb) where
  guniformPrior = guniformPrior @ta :*: guniformPrior @tb

instance (Generic (t HyperRep), GUniform (Rep (t HyperRep))) => Uniform (t :: (Type -> Type) -> Type) where
  uniformPrior = GHC.Generics.to (guniformPrior @(Rep (t HyperRep)))

-- sampling from prior
-- ------------------

class Prior a where
  sampleProbs :: (PrimMonad m) => Hyper a -> Prob m (Probs a)

class GPrior i o where
  gsampleProbs :: forall m p. PrimMonad m => i p -> Prob m (o p)

instance GPrior V1 V1 where
  gsampleProbs = undefined -- ok

instance GPrior U1 U1 where
  gsampleProbs _ = pure U1

-- base case: k is a conjugate distribution
instance (Prior (AsPrior p)) => GPrior (K1 i (HyperRep p)) (K1 i (ProbsRep p)) where
  gsampleProbs (K1 (HyperRep hyper)) =
    K1 . ProbsRep <$> sampleProbs @(AsPrior p) hyper

-- recursive case: k is another record
instance (Prior k, k HyperRep ~ Hyper k, k ProbsRep ~ Probs k) =>
         GPrior (K1 i (k HyperRep)) (K1 i (k ProbsRep)) where
  gsampleProbs (K1 hyper) = K1 <$> sampleProbs @k hyper

instance (GPrior ti to) => GPrior (M1 i c ti) (M1 i' c' to) where
  gsampleProbs (M1 x) = M1 <$> gsampleProbs x

instance (GPrior ia oa, GPrior ib ob) => GPrior (ia :*: ib) (oa :*: ob) where
  gsampleProbs (a :*: b) = (:*:) <$> gsampleProbs a <*> gsampleProbs b

instance (GPrior ia oa, GPrior ib ob) => GPrior (ia :+: ib) (oa :+: ob) where
  gsampleProbs (L1 a) = L1 <$> gsampleProbs a
  gsampleProbs (R1 b) = R1 <$> gsampleProbs b

instance ( Generic (a HyperRep)
         , Generic (a ProbsRep)
         , GPrior (Rep (a HyperRep)) (Rep (a ProbsRep))
         ) => Prior (a :: (Type -> Type) -> Type) where
  sampleProbs hyper = GHC.Generics.to <$> gsampleProbs (from hyper)

-----------------------
-- likelihood monads --
-----------------------

type Accessor r p = forall f . Lens' (r f) (f p)

class Monad m => RandomInterpreter m r | m -> r where
  type SampleCtx m a :: Constraint
  sampleValue :: (Conjugate p l, SampleCtx m l) => String -> l -> Accessor r p -> m (Support l)
  sampleConst :: (Distribution d, SampleCtx m d) => String -> d -> Params d -> m (Support d)
  permutationPlate :: (Ord a) => Int -> m a -> m [a]

newtype Trace (r :: (Type -> Type) -> Type)  = Trace { runTrace :: S.Seq Dynamic }
  deriving (Show)

observeValue
  :: (Conjugate p l, Typeable (Support l), Monad m)
  => String
  -> l
  -> Accessor r p
  -> Support l
  -> StateT (Trace r) m ()
observeValue _ _ _ val = modify $ \(Trace st) -> Trace $ st S.|> toDyn val

observeConst
  :: (Distribution d, Typeable (Support d), Monad m)
  => String
  -> d
  -> Params d
  -> Support d
  -> StateT (Trace r) m ()
observeConst _ _ _ val = modify $ \(Trace st) -> Trace $ st S.|> toDyn val

takeTrace :: Typeable a => Trace r -> Maybe (a, Trace r)
takeTrace (Trace t) = do
  (valDyn, rest) <- case S.viewl t of
    S.EmptyL         -> Nothing
    valDyn S.:< rest -> Just (valDyn, rest)
  val <- fromDynamic valDyn
  pure (val, Trace rest)

-- just sample
-- -----------

newtype SampleI m r a = SampleI (ReaderT (r ProbsRep) (Prob m) a)
  deriving (Functor, Applicative, Monad)

instance (PrimMonad m) => RandomInterpreter (SampleI m r) r where
  type SampleCtx (SampleI m r) a = ()
  sampleValue
    :: forall p l
     . (Conjugate p l)
    => String
    -> l
    -> Accessor r p
    -> SampleI m r (Support l)
  sampleValue _ lk getProbs = SampleI $ do
    probs <- ask
    lift $ distSample lk $ runProbs $ view getProbs probs
  sampleConst
    :: forall d
     . (Distribution d)
    => String
    -> d
    -> Params d
    -> SampleI m r (Support d)
  sampleConst _ dist params = SampleI $ lift $ distSample dist params
  permutationPlate = replicateM

sampleResult :: p ProbsRep -> SampleI m p a -> Gen (PrimState m) -> m a
sampleResult probs (SampleI a) = sample (runReaderT a probs)

-- sample and trace the execution process
-- --------------------------------------

newtype TraceI m r a = TraceI (ReaderT (r ProbsRep) (StateT (Trace r) (Prob m)) a)
  deriving (Functor, Applicative, Monad)

instance (PrimMonad m) => RandomInterpreter (TraceI m r) r where
  type SampleCtx (TraceI m r) l = Typeable (Support l)
  sampleValue
    :: forall p l
     . (Conjugate p l, Typeable (Support l))
    => String
    -> l
    -> Accessor r p
    -> TraceI m r (Support l)
  sampleValue _ lk getProbs = TraceI $ do
    probs <- ask
    val   <- lift $ lift $ distSample lk $ runProbs $ view getProbs probs
    modify $ \(Trace obs) -> Trace $ obs S.|> toDyn val
    pure val
  sampleConst
    :: forall d
     . (Distribution d, Typeable (Support d))
    => String
    -> d
    -> Params d
    -> TraceI m r (Support d)
  sampleConst _ dist params = TraceI $ do
    val <- lift $ lift $ distSample dist params
    modify $ \(Trace obs) -> Trace $ obs S.|> toDyn val
    pure val
  permutationPlate = replicateM

sampleTrace
  :: r ProbsRep -> TraceI m r a -> Gen (PrimState m) -> m (a, Trace r)
sampleTrace probs (TraceI a) = do
  let st = runReaderT a probs
      pr = runStateT st (Trace mempty)
  sample pr

-- evaluate the probability of a trace
-- -----------------------------------

newtype EvalTraceI r a = EvalTraceI (ReaderT (r ProbsRep) (StateT (Trace r, Double) Maybe) a)
  deriving (Functor, Applicative, Monad)

instance RandomInterpreter (EvalTraceI r) r where
  type SampleCtx (EvalTraceI r) l = Typeable (Support l)
  sampleValue
    :: forall p l
     . (Conjugate p l, Typeable (Support l))
    => String
    -> l
    -> Accessor r p
    -> EvalTraceI r (Support l)
  sampleValue _ lk getProbs = EvalTraceI $ do
    probs              <- ask
    (trace, totalLogP) <- get
    (val  , trace'   ) <- lift $ lift $ takeTrace trace
    let logP = distLogP lk (runProbs $ view getProbs probs) val
    put (trace', totalLogP + logP)
    pure val
  sampleConst
    :: forall d
     . (Distribution d, Typeable (Support d))
    => String
    -> d
    -> Params d
    -> EvalTraceI r (Support d)
  sampleConst _ dist params = EvalTraceI $ do
    (trace, totalLogP) <- get
    (val  , trace'   ) <- lift $ lift $ takeTrace trace
    let logP = distLogP dist params val
    put (trace', totalLogP + logP)
    pure val
  permutationPlate n submodel = EvalTraceI $ do
    probs                     <- ask
    (trace  , totalLogP     ) <- get
    (results, (trace', logP)) <- lift $ lift $ runTraceLogP
      probs
      trace
      (replicateM n submodel)
    let unique = MS.fromList results
        permutations =
          logFactorial n - sum (logFactorial . snd <$> MS.toOccurList unique)
    put (trace', totalLogP + logP + permutations)
    pure results

runTraceLogP
  :: r ProbsRep -> Trace r -> EvalTraceI r a -> Maybe (a, (Trace r, Double))
runTraceLogP probs trace (EvalTraceI model) = do
  runStateT (runReaderT model probs) (trace, 0)

evalTraceLogP :: r ProbsRep -> Trace r -> EvalTraceI r a -> Maybe (a, Double)
evalTraceLogP probs trace model = do
  (val, (_trace, logp)) <- runTraceLogP probs trace model
  pure (val, logp)

-- evaluate the predictive probability of a trace
-- -----------------------------------

newtype EvalPredTraceI r a = EvalPredTraceI (ReaderT (r HyperRep) (StateT (Trace r, Double) Maybe) a)
  deriving (Functor, Applicative, Monad)

instance RandomInterpreter (EvalPredTraceI r) r where
  type SampleCtx (EvalPredTraceI r) l = Typeable (Support l)
  sampleValue
    :: forall p l
     . (Conjugate p l, Typeable (Support l))
    => String
    -> l
    -> Accessor r p
    -> EvalPredTraceI r (Support l)
  sampleValue _ lk getHyper = EvalPredTraceI $ do
    hyper              <- ask
    (trace, totalLogP) <- get
    (val  , trace'   ) <- lift $ lift $ takeTrace trace
    let logP = predLogP @p lk (runHyper $ view getHyper hyper) val
    put (trace', totalLogP + logP)
    pure val
  sampleConst
    :: forall d
     . (Distribution d, Typeable (Support d))
    => String
    -> d
    -> Params d
    -> EvalPredTraceI r (Support d)
  sampleConst _ dist params = EvalPredTraceI $ do
    (trace, totalLogP) <- get
    (val  , trace'   ) <- lift $ lift $ takeTrace trace
    let logP = distLogP dist params val
    put (trace', totalLogP + logP)
    pure val
  permutationPlate n submodel = EvalPredTraceI $ do
    probs                     <- ask
    (trace  , totalLogP     ) <- get
    (results, (trace', logP)) <- lift $ lift $ runTracePredLogP
      probs
      trace
      (replicateM n submodel)
    let unique = MS.fromList results
        permutations =
          logFactorial n - sum (logFactorial . snd <$> MS.toOccurList unique)
    put (trace', totalLogP + logP + permutations)
    pure results

runTracePredLogP
  :: r HyperRep -> Trace r -> EvalPredTraceI r a -> Maybe (a, (Trace r, Double))
runTracePredLogP hyper trace (EvalPredTraceI model) = do
  runStateT (runReaderT model hyper) (trace, 0)

evalTracePredLogP
  :: r HyperRep -> Trace r -> EvalPredTraceI r a -> Maybe (a, Double)
evalTracePredLogP hyper trace model = do
  (val, (_trace, logp)) <- runTracePredLogP hyper trace model
  pure (val, logp)

-- update priors
-- -------------

newtype UpdatePriorsI r a = UpdatePriorsI (StateT (Trace r, r HyperRep) Maybe a)
  deriving (Functor, Applicative, Monad)

instance RandomInterpreter (UpdatePriorsI r) r where
  type SampleCtx (UpdatePriorsI r) l = Typeable (Support l)
  sampleValue
    :: forall p l
     . (Conjugate p l, Typeable (Support l))
    => String
    -> l
    -> Accessor r p
    -> UpdatePriorsI r (Support l)
  sampleValue _ lk accessor = UpdatePriorsI $ do
    (trace, priors) <- get
    (val  , trace') <- lift $ takeTrace trace
    let priors' :: r HyperRep
        priors' = over
          accessor
          (\(HyperRep pr) -> HyperRep $ updatePrior @p @l lk pr val)
          priors
    put (trace', priors')
    pure val
  sampleConst _ _ _ = UpdatePriorsI $ do
    (trace, priors) <- get
    (val  , trace') <- lift $ takeTrace trace
    put (trace', priors)
    pure val
  permutationPlate = replicateM

getPosterior
  :: r HyperRep -> Trace r -> UpdatePriorsI r a -> Maybe (r HyperRep)
getPosterior priors trace (UpdatePriorsI model) = do
  (_trace, posteriors) <- execStateT model (trace, priors)
  pure posteriors

-- show how a trace is generated by a model
-- ----------------------------------------

newtype ShowTraceI r a = ShowTraceI (MaybeT (WriterT String (State (Trace r))) a)
  deriving (Functor, Applicative, Monad)

showTraceItem
  :: forall l r
   . (Show (Support l), Typeable l, Typeable (Support l))
  => String
  -> ShowTraceI r (Support l)
showTraceItem name = ShowTraceI $ do
  trace         <- get
  (val, trace') <- MaybeT $ pure $ takeTrace trace
  put trace'
  let distName = show (typeRep (Proxy :: Proxy l))
  tell
    $  "Sampled value "
    <> show val
    <> " from a "
    <> distName
    <> " at "
    <> name
    <> ".\n"
  pure val

instance RandomInterpreter (ShowTraceI r) r where
  type SampleCtx (ShowTraceI r) l
    = (Typeable (Support l), Typeable l, Show (Support l))
  sampleValue
    :: forall p l
     . (Conjugate p l, Typeable (Support l), Typeable l, Show (Support l))
    => String
    -> l
    -> Accessor r p
    -> ShowTraceI r (Support l)
  sampleValue name _ acc = showTraceItem @l name
  sampleConst
    :: forall d
     . (Distribution d, SampleCtx (ShowTraceI r) d)
    => String
    -> d
    -> Params d
    -> ShowTraceI r (Support d)
  sampleConst name _ _ = showTraceItem @d name
  permutationPlate = replicateM

showTrace :: Trace r -> ShowTraceI r a -> (Maybe a, String)
showTrace trace (ShowTraceI model) =
  evalState (runWriterT (runMaybeT model)) trace

printTrace :: Trace r -> ShowTraceI r a -> IO ()
printTrace trace model = do
  let (res, txt) = showTrace trace model
  putStrLn txt
  case res of
    Nothing -> do
      putStrLn "Trace does not match the model (stops here)"
    Just _ -> pure ()

-- tracing interpreter (for debugging)
-- -----------------------------------

newtype TraceTraceI r a = TraceTraceI (State (Trace r) a)
  deriving (Functor, Applicative, Monad)

traceTraceItem
  :: forall l r
   . (Show (Support l), Typeable l, Typeable (Support l))
  => String
  -> TraceTraceI r (Support l)
traceTraceItem name = TraceTraceI $ do
  trace <- get
  let loc = show (typeRep (Proxy :: Proxy l)) <> " at " <> name
  case takeTrace trace of
    Just (val, trace') -> do
      put trace'
      DT.traceM $ "Sampled value " <> show val <> " from a " <> loc <> "."
      pure val
    Nothing -> error $ "Incompatible trace at " <> loc <> "."

instance RandomInterpreter (TraceTraceI r) r where
  type SampleCtx (TraceTraceI r) l
    = (Typeable (Support l), Typeable l, Show (Support l))
  sampleValue
    :: forall p l
     . (Conjugate p l, Typeable (Support l), Typeable l, Show (Support l))
    => String
    -> l
    -> Accessor r p
    -> TraceTraceI r (Support l)
  sampleValue name _ acc = traceTraceItem @l name
  sampleConst
    :: forall d
     . (Distribution d, SampleCtx (TraceTraceI r) d)
    => String
    -> d
    -> Params d
    -> TraceTraceI r (Support d)
  sampleConst name _ _ = traceTraceItem @d name
  permutationPlate = replicateM

traceTrace :: Trace r -> TraceTraceI r a -> a
traceTrace trace (TraceTraceI model) = evalState model trace

-------------------
-- distributions --
-------------------

-- Beta
-- ----

data Beta = Beta
  deriving (Eq, Ord, Show, Generic)

instance Distribution Beta where
  type Params Beta = (Double, Double)
  type Support Beta = Double
  distSample _ = uncurry beta
  distLogP _ (a, b) p =
    log (p ** (a - 1)) + log ((1 - p) ** (b - 1)) - logBeta a b

instance Jeffreys (AsPrior Beta) where
  jeffreysPrior = (0.5, 0.5)

instance Uniform (AsPrior Beta) where
  uniformPrior = (1, 1)

instance Prior (AsPrior Beta) where
  sampleProbs = distSample Beta

-- Bernoulli
-- ---------

data Bernoulli = Bernoulli
  deriving (Eq, Ord, Show, Generic)

instance Distribution Bernoulli where
  type Params Bernoulli = Double
  type Support Bernoulli = Bool
  distSample _ = bernoulli
  distLogP _ p True  = log p
  distLogP _ p False = log (1 - p)

-- Binomial
-- --------

newtype Binomial = Binomial Int
  deriving (Eq, Ord, Show, Generic)

instance Distribution Binomial where
  type Params Binomial = Double
  type Support Binomial = Int
  distSample (Binomial n) = binomial n
  distLogP (Binomial n) p k =
    logChoose n k + (log p * k') + (log (1 - p) * (n' - k'))
   where
    k' = int2Double k
    n' = int2Double n

-- Categorical
-- -----------

type Categorical :: Nat -> Type
data Categorical n = Categorical
  deriving (Eq, Ord, Show, Generic)

instance Distribution (Categorical n) where
  type Params (Categorical n) = V.Vector Double
  type Support (Categorical n) = Int
  distSample _ = categorical
  distLogP _ ps cat = log $ fromMaybe 0 $ ps V.!? cat

-- Dirichlet
-- ---------

type Dirichlet :: Nat -> Type
data Dirichlet n = Dirichlet
  deriving (Eq, Ord, Show, Generic)

instance Distribution (Dirichlet n) where
  type Params (Dirichlet n) = V.Vector Double
  type Support (Dirichlet n) = V.Vector Double
  distSample _ = dirichlet
  distLogP _ counts probs = logp + logz
   where
    logp = sum (V.zipWith (\a x -> log x * (a - 1)) counts probs)
    logz = logGamma (sum counts) - sum (logGamma <$> counts)

instance KnownNat n => Jeffreys (AsPrior (Dirichlet n)) where
  jeffreysPrior = V.replicate (fromIntegral $ natVal (Proxy :: Proxy n)) 0.5

instance KnownNat n => Uniform (AsPrior (Dirichlet n)) where
  uniformPrior = V.replicate (fromIntegral $ natVal (Proxy :: Proxy n)) 1

instance Prior (AsPrior (Dirichlet n)) where
  sampleProbs = distSample Dirichlet

-- Geometric (from 0)
-- ------------------

data Geometric0 = Geometric0
  deriving (Eq, Ord, Show, Generic)

instance Distribution Geometric0 where
  type Params Geometric0 = Double
  type Support Geometric0 = Int
  distSample _ = geometric0
   where
    geometric0 p = do
      coin <- bernoulli p
      if coin then pure 0 else (1 +) <$> geometric0 p
  distLogP _ p val | val >= 0  = (log (1 - p) * int2Double val) + log p
                   | otherwise = log 0

-- Geometric (from 1)
-- ------------------

data Geometric1 = Geometric1
  deriving (Eq, Ord, Show, Generic)

instance Distribution Geometric1 where
  type Params Geometric1 = Double
  type Support Geometric1 = Int
  distSample _ = geometric1
   where
    geometric1 p = do
      coin <- bernoulli p
      if coin then pure 1 else (1 +) <$> geometric1 p
  distLogP _ p val | val >= 1  = (log (1 - p) * int2Double (val - 1)) + log p
                   | otherwise = log 0

-----------------------------
-- conjugate distributions --
-----------------------------

-- beta bernoulli
-- --------------

instance Conjugate Beta Bernoulli where
  priorSingleton = Beta
  updatePrior _ (a, b) False = (a, b + 1)
  updatePrior _ (a, b) True  = (a + 1, b)
  predLogP _ (a, b) False = log a - log (a + b)
  predLogP _ (a, b) True  = log b - log (a + b)

-- B(1+a, b)   G(1+a) G(b) G(a+b)   G(1+a) G(a+b)   a G(a) G(a+b)        a
-- --------- = ------------------ = ------------- = ----------------- = ---
-- B(a, b)     G(a) G(b) G(a+b+1)   G(a) G(a+b+1)   G(a) (a+b) G(a+b)   a+b

-- beta binomial
-- -------------

instance Conjugate Beta Binomial where
  priorSingleton = Beta
  updatePrior (Binomial n) (a, b) x = (a + x', b + (n' - x'))
   where
    x' = int2Double x
    n' = int2Double n
  predLogP (Binomial n) (a, b) k =
    logChoose n k + logBeta (k' + a) (n' - k' + a) - logBeta a b
   where
    n' = int2Double n
    k' = int2Double k

-- beta geometric0
-- ---------------

instance Conjugate Beta Geometric0 where
  priorSingleton = Beta
  updatePrior _ (a, b) k = (a + 1, b + int2Double k)
  predLogP _ (a, b) k =
    (log a + logGamma (a + b) + logGamma (k' + b))
      - (logGamma b + logGamma (a + b + k' + 1))
    where k' = int2Double k

-- beta geometric1
-- ---------------

instance Conjugate Beta Geometric1 where
  priorSingleton = Beta
  updatePrior _ (a, b) k = (a + 1, b + int2Double (k - 1))
  predLogP _ (a, b) k =
    (log a + logGamma (a + b) + logGamma (k' + b))
      - (logGamma b + logGamma (a + b + k' + 1))
    where k' = int2Double k - 1

-- dirichlet categorical
-- ---------------------

instance Conjugate (Dirichlet n) (Categorical n) where
  priorSingleton = Dirichlet
  updatePrior _ counts obs
    | obs >= 0, obs < V.length counts
    = counts V.// [(obs, (counts V.! obs) + 1)]
    | otherwise
    = counts
  predLogP _ counts obs =
    log (fromMaybe 0 $ counts V.!? obs) - log (V.sum counts)


{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE BangPatterns #-}

module Statistics.Quantile.Exact (
    external
  , external'
  , countOut
) where

import Control.Applicative
import Control.Monad.IO.Class

import Data.Monoid

import Debug.Trace

import Numeric

import Pipes
import qualified Pipes.Prelude as P

import Statistics.Quantile.Types
import Statistics.Quantile.Util

import System.Directory
import System.Exit
import System.IO
import System.IO.Temp
import System.FilePath
import System.Process

external :: Selector IO
external =  Selector external'

external' :: Quantile
          -> Stream IO
          -> IO Double
external' q src = do
  (fp, fh) <- openTempFile tmpDir "quantile.txt"
  n <- countOut fh src
  hClose fh
  fp' <- externalMergeSort fp
  fh' <- openFile fp' ReadMode
  quant <- select (qidx q n) $
                P.fromHandle fh'
            >-> P.map realOrDie
  removeFile fp >> removeFile fp'
  pure quant
  where realOrDie :: String -> Double
        realOrDie = read

countOut :: Handle
         -> Stream IO
         -> IO Int
countOut fh (Stream src) =
  P.foldM (incrementWrite fh) (pure 0) pure src
  where incrementWrite h n v = do
          -- Can't use the Show instance because scientific notation
          -- would break sorting.
          hPutStrLn h (showFFloat Nothing v "")
          pure $! (n + 1)

select :: Monad m
       => Int
       -> Producer Double m ()
       -> m Double
select n p0 = next p0 >>= \case
  (Left _)       -> error $ show EmptySample
  (Right (v, p)) -> go 1 v p
   where go i v p
           | i == n    = pure v
           | otherwise =
               next p >>= \case
                 (Left _) -> error "universe broke"
                 (Right (v', p')) -> go (i + 1) v' p'

externalMergeSort :: FilePath -> IO FilePath
externalMergeSort fp =
  let fp' = fp <> ".sorted" in do
  fh <- openFile fp' WriteMode
  ph <- runProcess "/usr/bin/sort" sortOpts Nothing Nothing Nothing (Just fh) Nothing
  st <- waitForProcess ph
  hClose fh
  case st of
    ExitSuccess -> pure fp'
    _ -> error $ "sort -n " <> fp <> " failed"
  where
    -- Numeric, TMPDIR=/var/tmp so we don't fill up a tmpfs.
    sortOpts = [ "-n", "-T", "/var/tmp", fp]


{-# LANGUAGE CPP #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE ConstraintKinds #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE PolyKinds #-}
{-# OPTIONS_GHC -fno-warn-deprecations #-}
module Data.Constraint.Lifting
  ( Lifting(..)
  , Lifting2(..)
  ) where

import Control.Applicative
import Control.Applicative.Backwards
import Control.Applicative.Lift
import Control.DeepSeq
import Control.Monad
import Control.Monad.Cont.Class
import Control.Monad.Error.Class
import Control.Monad.Fix
import Control.Monad.IO.Class
import Control.Monad.RWS.Class
import Control.Monad.Trans.Cont
import Control.Monad.Trans.Error
import Control.Monad.Trans.Except
import Control.Monad.Trans.Identity
import Control.Monad.Trans.List
import Control.Monad.Trans.Maybe
import Control.Monad.Trans.Reader
import Control.Monad.Trans.RWS.Lazy as Lazy
import Control.Monad.Trans.RWS.Strict as Strict
import Control.Monad.Trans.State.Lazy as Lazy
import Control.Monad.Trans.State.Strict as Strict
import Control.Monad.Trans.Writer.Lazy as Lazy
import Control.Monad.Trans.Writer.Strict as Strict
import Data.Binary
import Data.Complex
import Data.Constraint
import Data.Functor.Classes
import Data.Functor.Compose as Functor
import Data.Functor.Identity
import Data.Functor.Product as Functor
import Data.Functor.Reverse as Functor
import Data.Functor.Sum as Functor
import Data.Hashable
import Data.Ratio
import GHC.Arr

class Lifting p f where
  lifting :: p a :- p (f a)

instance Lifting Eq [] where lifting = Sub Dict
instance Lifting Ord [] where lifting = Sub Dict
instance Lifting Show [] where lifting = Sub Dict
instance Lifting Read [] where lifting = Sub Dict
instance Lifting Hashable [] where lifting = Sub Dict
instance Lifting Binary [] where lifting = Sub Dict
instance Lifting NFData [] where lifting = Sub Dict

instance Lifting Eq Maybe where lifting = Sub Dict
instance Lifting Ord Maybe where lifting = Sub Dict
instance Lifting Show Maybe where lifting = Sub Dict
instance Lifting Read Maybe where lifting = Sub Dict
instance Lifting Hashable Maybe where lifting = Sub Dict
instance Lifting Binary Maybe where lifting = Sub Dict
instance Lifting NFData Maybe where lifting = Sub Dict
instance Lifting Semigroup Maybe where lifting = Sub Dict
instance Lifting Monoid Maybe where lifting = Sub Dict

instance Lifting Eq Ratio where lifting = Sub Dict
-- instance Lifting Show Ratio where lifting = Sub Dict -- requires 7.10

instance Lifting Eq Complex where lifting = Sub Dict
instance Lifting Read Complex where lifting = Sub Dict
instance Lifting Show Complex where lifting = Sub Dict
instance Lifting Semigroup ((->) a) where lifting = Sub Dict
instance Lifting Monoid ((->) a) where lifting = Sub Dict

instance Eq a => Lifting Eq (Either a) where lifting = Sub Dict
instance Ord a => Lifting Ord (Either a) where lifting = Sub Dict
instance Show a => Lifting Show (Either a) where lifting = Sub Dict
instance Read a => Lifting Read (Either a) where lifting = Sub Dict
instance Hashable a => Lifting Hashable (Either a) where lifting = Sub Dict
instance Binary a => Lifting Binary (Either a) where lifting = Sub Dict
instance NFData a => Lifting NFData (Either a) where lifting = Sub Dict

instance Eq a => Lifting Eq ((,) a) where lifting = Sub Dict
instance Ord a => Lifting Ord ((,) a) where lifting = Sub Dict
instance Show a => Lifting Show ((,) a) where lifting = Sub Dict
instance Read a => Lifting Read ((,) a) where lifting = Sub Dict
instance Hashable a => Lifting Hashable ((,) a) where lifting = Sub Dict
instance Binary a => Lifting Binary ((,) a) where lifting = Sub Dict
instance NFData a => Lifting NFData ((,) a) where lifting = Sub Dict
instance Semigroup a => Lifting Semigroup ((,) a) where lifting = Sub Dict
instance Monoid a => Lifting Monoid ((,) a) where lifting = Sub Dict
instance Bounded a => Lifting Bounded ((,) a) where lifting = Sub Dict
instance Ix a => Lifting Ix ((,) a) where lifting = Sub Dict

instance Functor f => Lifting Functor (Compose f) where lifting = Sub Dict
instance Foldable f => Lifting Foldable (Compose f) where lifting = Sub Dict
instance Traversable f => Lifting Traversable (Compose f) where lifting = Sub Dict
instance Applicative f => Lifting Applicative (Compose f) where lifting = Sub Dict
instance Alternative f => Lifting Alternative (Compose f) where lifting = Sub Dict -- overconstrained

instance Show1 f => Lifting Show1 (Compose f) where lifting = Sub Dict
instance Eq1 f => Lifting Eq1 (Compose f) where lifting = Sub Dict
instance Ord1 f => Lifting Ord1 (Compose f) where lifting = Sub Dict
instance Read1 f => Lifting Read1 (Compose f) where lifting = Sub Dict
instance (Eq1 f, Eq1 g) => Lifting Eq (Compose f g) where lifting = Sub Dict
instance (Ord1 f, Ord1 g) => Lifting Ord (Compose f g) where lifting = Sub Dict
instance (Read1 f, Read1 g) => Lifting Read (Compose f g) where lifting = Sub Dict
instance (Show1 f, Show1 g) => Lifting Show (Compose f g) where lifting = Sub Dict

instance Functor f => Lifting Functor (Functor.Product f) where lifting = Sub Dict
instance Foldable f => Lifting Foldable (Functor.Product f) where lifting = Sub Dict
instance Traversable f => Lifting Traversable (Functor.Product f) where lifting = Sub Dict
instance Applicative f => Lifting Applicative (Functor.Product f) where lifting = Sub Dict
instance Alternative f => Lifting Alternative (Functor.Product f) where lifting = Sub Dict
instance Monad f => Lifting Monad (Functor.Product f) where lifting = Sub Dict
instance MonadFix f => Lifting MonadFix (Functor.Product f) where lifting = Sub Dict
instance MonadPlus f => Lifting MonadPlus (Functor.Product f) where lifting = Sub Dict
instance Show1 f => Lifting Show1 (Functor.Product f) where lifting = Sub Dict
instance Eq1 f => Lifting Eq1 (Functor.Product f) where lifting = Sub Dict
instance Ord1 f => Lifting Ord1 (Functor.Product f) where lifting = Sub Dict
instance Read1 f => Lifting Read1 (Functor.Product f) where lifting = Sub Dict
instance (Eq1 f, Eq1 g) => Lifting Eq (Functor.Product f g) where lifting = Sub Dict
instance (Ord1 f, Ord1 g) => Lifting Ord (Functor.Product f g) where lifting = Sub Dict
instance (Read1 f, Read1 g) => Lifting Read (Functor.Product f g) where lifting = Sub Dict
instance (Show1 f, Show1 g) => Lifting Show (Functor.Product f g) where lifting = Sub Dict

instance Functor f => Lifting Functor (Functor.Sum f) where lifting = Sub Dict
instance Foldable f => Lifting Foldable (Functor.Sum f) where lifting = Sub Dict
instance Traversable f => Lifting Traversable (Functor.Sum f) where lifting = Sub Dict
instance Show1 f => Lifting Show1 (Functor.Sum f) where lifting = Sub Dict
instance Eq1 f => Lifting Eq1 (Functor.Sum f) where lifting = Sub Dict
instance Ord1 f => Lifting Ord1 (Functor.Sum f) where lifting = Sub Dict
instance Read1 f => Lifting Read1 (Functor.Sum f) where lifting = Sub Dict
instance (Eq1 f, Eq1 g) => Lifting Eq (Functor.Sum f g) where lifting = Sub Dict
instance (Ord1 f, Ord1 g) => Lifting Ord (Functor.Sum f g) where lifting = Sub Dict
instance (Read1 f, Read1 g) => Lifting Read (Functor.Sum f g) where lifting = Sub Dict
instance (Show1 f, Show1 g) => Lifting Show (Functor.Sum f g) where lifting = Sub Dict

instance Lifting Functor (Strict.StateT s) where lifting = Sub Dict
instance Lifting Monad (Strict.StateT s) where lifting = Sub Dict
instance Lifting MonadFix (Strict.StateT s) where lifting = Sub Dict
instance Lifting MonadIO (Strict.StateT s) where lifting = Sub Dict
instance Lifting MonadPlus (Strict.StateT s) where lifting = Sub Dict

instance Lifting Functor (Lazy.StateT s) where lifting = Sub Dict
instance Lifting Monad (Lazy.StateT s) where lifting = Sub Dict
instance Lifting MonadFix (Lazy.StateT s) where lifting = Sub Dict
instance Lifting MonadIO (Lazy.StateT s) where lifting = Sub Dict
instance Lifting MonadPlus (Lazy.StateT s) where lifting = Sub Dict

instance Lifting Functor (Lazy.RWST r w s) where lifting = Sub Dict
instance Monoid w => Lifting Monad (Lazy.RWST r w s) where lifting = Sub Dict
instance Monoid w => Lifting MonadFix (Lazy.RWST r w s) where lifting = Sub Dict
instance Monoid w => Lifting MonadPlus (Lazy.RWST r w s) where lifting = Sub Dict
instance Monoid w => Lifting MonadIO (Lazy.RWST r w s) where lifting = Sub Dict

instance Lifting Functor (Strict.RWST r w s) where lifting = Sub Dict
instance Monoid w => Lifting Monad (Strict.RWST r w s) where lifting = Sub Dict
instance Monoid w => Lifting MonadFix (Strict.RWST r w s) where lifting = Sub Dict
instance Monoid w => Lifting MonadPlus (Strict.RWST r w s) where lifting = Sub Dict
instance Monoid w => Lifting MonadIO (Strict.RWST r w s) where lifting = Sub Dict

instance Lifting Functor (ReaderT e) where lifting = Sub Dict
instance Lifting Applicative (ReaderT e) where lifting = Sub Dict
instance Lifting Alternative (ReaderT e) where lifting = Sub Dict
instance Lifting Monad (ReaderT e) where lifting = Sub Dict
instance Lifting MonadPlus (ReaderT e) where lifting = Sub Dict
instance Lifting MonadFix (ReaderT e) where lifting = Sub Dict
instance Lifting MonadIO (ReaderT e) where lifting = Sub Dict

instance Lifting Functor (ErrorT e) where lifting = Sub Dict
instance Lifting Foldable (ErrorT e) where lifting = Sub Dict
instance Lifting Traversable (ErrorT e) where lifting = Sub Dict
instance Error e => Lifting Monad (ErrorT e) where lifting = Sub Dict
instance Error e => Lifting MonadFix (ErrorT e) where lifting = Sub Dict
instance Error e => Lifting MonadPlus (ErrorT e) where lifting = Sub Dict -- overconstrained!
instance Error e => Lifting MonadIO (ErrorT e) where lifting = Sub Dict
instance Show e => Lifting Show1 (ErrorT e) where lifting = Sub Dict
instance Eq e => Lifting Eq1 (ErrorT e) where lifting = Sub Dict
instance Ord e => Lifting Ord1 (ErrorT e) where lifting = Sub Dict
instance Read e => Lifting Read1 (ErrorT e) where lifting = Sub Dict
instance (Show e, Show1 m) => Lifting Show (ErrorT e m) where lifting = Sub Dict
instance (Eq e, Eq1 m) => Lifting Eq (ErrorT e m) where lifting = Sub Dict
instance (Ord e, Ord1 m) => Lifting Ord (ErrorT e m) where lifting = Sub Dict
instance (Read e, Read1 m) => Lifting Read (ErrorT e m) where lifting = Sub Dict

instance Lifting Functor (ExceptT e) where lifting = Sub Dict
instance Lifting Foldable (ExceptT e) where lifting = Sub Dict
instance Lifting Traversable (ExceptT e) where lifting = Sub Dict
instance Lifting Monad (ExceptT e) where lifting = Sub Dict
instance Lifting MonadFix (ExceptT e) where lifting = Sub Dict
instance Monoid e => Lifting MonadPlus (ExceptT e) where lifting = Sub Dict -- overconstrained!
instance Lifting MonadIO (ExceptT e) where lifting = Sub Dict
instance Show e => Lifting Show1 (ExceptT e) where lifting = Sub Dict
instance Eq e => Lifting Eq1 (ExceptT e) where lifting = Sub Dict
instance Ord e => Lifting Ord1 (ExceptT e) where lifting = Sub Dict
instance Read e => Lifting Read1 (ExceptT e) where lifting = Sub Dict
instance (Show e, Show1 m) => Lifting Show (ExceptT e m) where lifting = Sub Dict
instance (Eq e, Eq1 m) => Lifting Eq (ExceptT e m) where lifting = Sub Dict
instance (Ord e, Ord1 m) => Lifting Ord (ExceptT e m) where lifting = Sub Dict
instance (Read e, Read1 m) => Lifting Read (ExceptT e m) where lifting = Sub Dict

instance Lifting Functor (Strict.WriterT w) where lifting = Sub Dict
instance Monoid w => Lifting Applicative (Strict.WriterT w) where lifting = Sub Dict
instance Monoid w => Lifting Alternative (Strict.WriterT w) where lifting = Sub Dict
instance Monoid w => Lifting Monad (Strict.WriterT w) where lifting = Sub Dict
instance Monoid w => Lifting MonadFix (Strict.WriterT w) where lifting = Sub Dict
instance Monoid w => Lifting MonadPlus (Strict.WriterT w) where lifting = Sub Dict
instance Lifting Foldable (Strict.WriterT w) where lifting = Sub Dict
instance Lifting Traversable (Strict.WriterT w) where lifting = Sub Dict
instance Monoid w => Lifting MonadIO (Strict.WriterT w) where lifting = Sub Dict
instance Show w => Lifting Show1 (Strict.WriterT w) where lifting = Sub Dict
instance Eq w => Lifting Eq1 (Strict.WriterT w) where lifting = Sub Dict
instance Ord w => Lifting Ord1 (Strict.WriterT w) where lifting = Sub Dict
instance Read w => Lifting Read1 (Strict.WriterT w) where lifting = Sub Dict
instance (Show w, Show1 m) => Lifting Show (Strict.WriterT w m) where lifting = Sub Dict
instance (Eq w, Eq1 m) => Lifting Eq (Strict.WriterT w m) where lifting = Sub Dict
instance (Ord w, Ord1 m) => Lifting Ord (Strict.WriterT w m) where lifting = Sub Dict
instance (Read w, Read1 m) => Lifting Read (Strict.WriterT w m) where lifting = Sub Dict

instance Lifting Functor (Lazy.WriterT w) where lifting = Sub Dict
instance Monoid w => Lifting Applicative (Lazy.WriterT w) where lifting = Sub Dict
instance Monoid w => Lifting Alternative (Lazy.WriterT w) where lifting = Sub Dict
instance Monoid w => Lifting Monad (Lazy.WriterT w) where lifting = Sub Dict
instance Monoid w => Lifting MonadFix (Lazy.WriterT w) where lifting = Sub Dict
instance Monoid w => Lifting MonadPlus (Lazy.WriterT w) where lifting = Sub Dict
instance Lifting Foldable (Lazy.WriterT w) where lifting = Sub Dict
instance Lifting Traversable (Lazy.WriterT w) where lifting = Sub Dict
instance Monoid w => Lifting MonadIO (Lazy.WriterT w) where lifting = Sub Dict
instance Show w => Lifting Show1 (Lazy.WriterT w) where lifting = Sub Dict
instance Eq w => Lifting Eq1 (Lazy.WriterT w) where lifting = Sub Dict
instance Ord w => Lifting Ord1 (Lazy.WriterT w) where lifting = Sub Dict
instance Read w => Lifting Read1 (Lazy.WriterT w) where lifting = Sub Dict
instance (Show w, Show1 m) => Lifting Show (Lazy.WriterT w m) where lifting = Sub Dict
instance (Eq w, Eq1 m) => Lifting Eq (Lazy.WriterT w m) where lifting = Sub Dict
instance (Ord w, Ord1 m) => Lifting Ord (Lazy.WriterT w m) where lifting = Sub Dict
instance (Read w, Read1 m) => Lifting Read (Lazy.WriterT w m) where lifting = Sub Dict

instance Lifting Functor (ContT r) where lifting = Sub Dict -- overconstrained
instance Lifting Applicative (ContT r) where lifting = Sub Dict -- overconstrained
instance Lifting Monad (ContT r) where lifting = Sub Dict -- overconstrained
instance Lifting MonadIO (ContT r) where lifting = Sub Dict

instance Lifting Functor IdentityT where lifting = Sub Dict
instance Lifting Applicative IdentityT where lifting = Sub Dict
instance Lifting Alternative IdentityT where lifting = Sub Dict
instance Lifting Monad IdentityT where lifting = Sub Dict
instance Lifting MonadPlus IdentityT where lifting = Sub Dict
instance Lifting MonadFix IdentityT where lifting = Sub Dict
instance Lifting Foldable IdentityT where lifting = Sub Dict
instance Lifting Traversable IdentityT where lifting = Sub Dict
instance Lifting MonadIO IdentityT where lifting = Sub Dict
instance Lifting Show1 IdentityT where lifting = Sub Dict
instance Lifting Read1 IdentityT where lifting = Sub Dict
instance Lifting Ord1 IdentityT where lifting = Sub Dict
instance Lifting Eq1 IdentityT where lifting = Sub Dict
instance Show1 m => Lifting Show (IdentityT m) where lifting = Sub Dict
instance Read1 m => Lifting Read (IdentityT m) where lifting = Sub Dict
instance Ord1 m => Lifting Ord (IdentityT m) where lifting = Sub Dict
instance Eq1 m => Lifting Eq (IdentityT m) where lifting = Sub Dict

instance Lifting Functor ListT where lifting = Sub Dict
instance Lifting Applicative ListT where lifting = Sub Dict
instance Lifting Alternative ListT where lifting = Sub Dict -- overconstrained
instance Lifting Monad ListT where lifting = Sub Dict
-- instance Lifting MonadFix ListT where lifting = Sub Dict
instance Lifting MonadPlus ListT where lifting = Sub Dict -- overconstrained
instance Lifting Foldable ListT where lifting = Sub Dict
instance Lifting Traversable ListT where lifting = Sub Dict
instance Lifting MonadIO ListT where lifting = Sub Dict
instance Lifting Show1 ListT where lifting = Sub Dict
instance Lifting Read1 ListT where lifting = Sub Dict
instance Lifting Ord1 ListT where lifting = Sub Dict
instance Lifting Eq1 ListT where lifting = Sub Dict
instance Show1 m => Lifting Show (ListT m) where lifting = Sub Dict
instance Read1 m => Lifting Read (ListT m) where lifting = Sub Dict
instance Ord1 m => Lifting Ord (ListT m) where lifting = Sub Dict
instance Eq1 m => Lifting Eq (ListT m) where lifting = Sub Dict

instance Lifting Functor MaybeT where lifting = Sub Dict
instance Lifting Monad MaybeT where lifting = Sub Dict
-- instance Lifting MonadFix MaybeT where lifting = Sub Dict
instance Lifting MonadPlus MaybeT where lifting = Sub Dict -- overconstrained
instance Lifting Foldable MaybeT where lifting = Sub Dict
instance Lifting Traversable MaybeT where lifting = Sub Dict
instance Lifting MonadIO MaybeT where lifting = Sub Dict
instance Lifting Show1 MaybeT where lifting = Sub Dict
instance Lifting Read1 MaybeT where lifting = Sub Dict
instance Lifting Ord1 MaybeT where lifting = Sub Dict
instance Lifting Eq1 MaybeT where lifting = Sub Dict
instance Show1 m => Lifting Show (MaybeT m) where lifting = Sub Dict
instance Read1 m => Lifting Read (MaybeT m) where lifting = Sub Dict
instance Ord1 m => Lifting Ord (MaybeT m) where lifting = Sub Dict
instance Eq1 m => Lifting Eq (MaybeT m) where lifting = Sub Dict

instance Lifting Functor Reverse where lifting = Sub Dict
instance Lifting Applicative Reverse where lifting = Sub Dict
instance Lifting Alternative Reverse where lifting = Sub Dict
instance Lifting Foldable Reverse where lifting = Sub Dict
instance Lifting Traversable Reverse where lifting = Sub Dict
instance Lifting Show1 Reverse where lifting = Sub Dict
instance Lifting Read1 Reverse where lifting = Sub Dict
instance Lifting Ord1 Reverse where lifting = Sub Dict
instance Lifting Eq1 Reverse where lifting = Sub Dict
instance Show1 f => Lifting Show (Reverse f) where lifting = Sub Dict
instance Read1 f => Lifting Read (Reverse f) where lifting = Sub Dict
instance Ord1 f => Lifting Ord (Reverse f) where lifting = Sub Dict
instance Eq1 f => Lifting Eq (Reverse f) where lifting = Sub Dict

instance Lifting Functor Backwards where lifting = Sub Dict
instance Lifting Foldable Backwards where lifting = Sub Dict
instance Lifting Traversable Backwards where lifting = Sub Dict
instance Lifting Applicative Backwards where lifting = Sub Dict
instance Lifting Alternative Backwards where lifting = Sub Dict
instance Lifting Show1 Backwards where lifting = Sub Dict
instance Lifting Read1 Backwards where lifting = Sub Dict
instance Lifting Ord1 Backwards where lifting = Sub Dict
instance Lifting Eq1 Backwards where lifting = Sub Dict
instance Show1 f => Lifting Show (Backwards f) where lifting = Sub Dict
instance Read1 f => Lifting Read (Backwards f) where lifting = Sub Dict
instance Ord1 f => Lifting Ord (Backwards f) where lifting = Sub Dict
instance Eq1 f => Lifting Eq (Backwards f) where lifting = Sub Dict

instance Lifting Functor Lift where lifting = Sub Dict
instance Lifting Foldable Lift where lifting = Sub Dict
instance Lifting Traversable Lift where lifting = Sub Dict
instance Lifting Applicative Lift where lifting = Sub Dict
instance Lifting Alternative Lift where lifting = Sub Dict
instance Lifting Show1 Lift where lifting = Sub Dict
instance Lifting Read1 Lift where lifting = Sub Dict
instance Lifting Ord1 Lift where lifting = Sub Dict
instance Lifting Eq1 Lift where lifting = Sub Dict
instance Show1 f => Lifting Show (Lift f) where lifting = Sub Dict
instance Read1 f => Lifting Read (Lift f) where lifting = Sub Dict
instance Ord1 f => Lifting Ord (Lift f) where lifting = Sub Dict
instance Eq1 f => Lifting Eq (Lift f) where lifting = Sub Dict

instance Lifting Eq Identity where lifting = Sub Dict
instance Lifting Ord Identity where lifting = Sub Dict
instance Lifting Show Identity where lifting = Sub Dict
instance Lifting Read Identity where lifting = Sub Dict

instance Lifting MonadCont MaybeT where lifting = Sub Dict
instance Lifting MonadCont ListT where lifting = Sub Dict
instance Lifting MonadCont IdentityT where lifting = Sub Dict
instance Monoid w => Lifting MonadCont (Strict.WriterT w) where lifting = Sub Dict
instance Monoid w => Lifting MonadCont (Lazy.WriterT w) where lifting = Sub Dict
instance Error e => Lifting MonadCont (ErrorT e) where lifting = Sub Dict
instance Lifting MonadCont (ExceptT w) where lifting = Sub Dict
instance Lifting MonadCont (Strict.StateT s) where lifting = Sub Dict
instance Lifting MonadCont (Lazy.StateT s) where lifting = Sub Dict
instance Lifting MonadCont (ReaderT e) where lifting = Sub Dict
instance Monoid w => Lifting MonadCont (Strict.RWST r w s) where lifting = Sub Dict
instance Monoid w => Lifting MonadCont (Lazy.RWST r w s) where lifting = Sub Dict

instance Lifting (MonadError e) MaybeT where lifting = Sub Dict
instance Lifting (MonadError e) ListT where lifting = Sub Dict
instance Lifting (MonadError e) IdentityT where lifting = Sub Dict
instance Monoid w => Lifting (MonadError e) (Strict.WriterT w) where lifting = Sub Dict
instance Monoid w => Lifting (MonadError e) (Lazy.WriterT w) where lifting = Sub Dict
instance Lifting (MonadError e) (Strict.StateT s) where lifting = Sub Dict
instance Lifting (MonadError e) (Lazy.StateT s) where lifting = Sub Dict
instance Lifting (MonadError e) (ReaderT r) where lifting = Sub Dict
instance Monoid w => Lifting (MonadError e) (Strict.RWST r w s) where lifting = Sub Dict
instance Monoid w => Lifting (MonadError e) (Lazy.RWST r w s) where lifting = Sub Dict

instance Lifting (MonadRWS r w s) MaybeT where lifting = Sub Dict
instance Lifting (MonadRWS r w s) IdentityT where lifting = Sub Dict
instance Error e => Lifting (MonadRWS r w s) (ErrorT e) where lifting = Sub Dict
instance Lifting (MonadRWS r w s) (ExceptT e) where lifting = Sub Dict

instance Lifting (MonadReader r) MaybeT where lifting = Sub Dict
instance Lifting (MonadReader r) ListT where lifting = Sub Dict
instance Lifting (MonadReader r) IdentityT where lifting = Sub Dict
instance Monoid w => Lifting (MonadReader r) (Strict.WriterT w) where lifting = Sub Dict
instance Monoid w => Lifting (MonadReader r) (Lazy.WriterT w) where lifting = Sub Dict
instance Lifting (MonadReader r) (Strict.StateT s) where lifting = Sub Dict
instance Lifting (MonadReader r) (Lazy.StateT s) where lifting = Sub Dict
instance Lifting (MonadReader r) (ExceptT e) where lifting = Sub Dict
instance Error e => Lifting (MonadReader r) (ErrorT e) where lifting = Sub Dict
instance Lifting (MonadReader r) (ContT r') where lifting = Sub Dict

instance Lifting (MonadState s) MaybeT where lifting = Sub Dict
instance Lifting (MonadState s) ListT where lifting = Sub Dict
instance Lifting (MonadState s) IdentityT where lifting = Sub Dict
instance Monoid w => Lifting (MonadState s) (Strict.WriterT w) where lifting = Sub Dict
instance Monoid w => Lifting (MonadState s) (Lazy.WriterT w) where lifting = Sub Dict
instance Lifting (MonadState s) (ReaderT r) where lifting = Sub Dict
instance Lifting (MonadState s) (ExceptT e) where lifting = Sub Dict
instance Error e => Lifting (MonadState s) (ErrorT e) where lifting = Sub Dict
instance Lifting (MonadState s) (ContT r') where lifting = Sub Dict

class Lifting2 p f where
  lifting2 :: p a :- Lifting p (f a) -- (p a, p b) :- p (f a b)

instance Lifting2 Eq Either where lifting2 = Sub Dict
instance Lifting2 Ord Either where lifting2 = Sub Dict
instance Lifting2 Show Either where lifting2 = Sub Dict
instance Lifting2 Read Either where lifting2 = Sub Dict
instance Lifting2 Hashable Either where lifting2 = Sub Dict
instance Lifting2 Binary Either where lifting2 = Sub Dict
instance Lifting2 NFData Either where lifting2 = Sub Dict

instance Lifting2 Eq (,) where lifting2 = Sub Dict
instance Lifting2 Ord (,) where lifting2 = Sub Dict
instance Lifting2 Show (,) where lifting2 = Sub Dict
instance Lifting2 Read (,) where lifting2 = Sub Dict
instance Lifting2 Hashable (,) where lifting2 = Sub Dict
instance Lifting2 Binary (,) where lifting2 = Sub Dict
instance Lifting2 NFData (,) where lifting2 = Sub Dict
instance Lifting2 Semigroup (,) where lifting2 = Sub Dict
instance Lifting2 Monoid (,) where lifting2 = Sub Dict
instance Lifting2 Bounded (,) where lifting2 = Sub Dict
instance Lifting2 Ix (,) where lifting2 = Sub Dict

instance Lifting2 Functor Compose where lifting2 = Sub Dict
instance Lifting2 Foldable Compose where lifting2 = Sub Dict
instance Lifting2 Traversable Compose where lifting2 = Sub Dict
instance Lifting2 Applicative Compose where lifting2 = Sub Dict
instance Lifting2 Alternative Compose where lifting2 = Sub Dict -- overconstrained

instance Lifting2 Functor Functor.Product where lifting2 = Sub Dict
instance Lifting2 Foldable Functor.Product where lifting2 = Sub Dict
instance Lifting2 Traversable Functor.Product where lifting2 = Sub Dict
instance Lifting2 Applicative Functor.Product where lifting2 = Sub Dict
instance Lifting2 Alternative Functor.Product where lifting2 = Sub Dict
instance Lifting2 Monad Functor.Product where lifting2 = Sub Dict
instance Lifting2 MonadPlus Functor.Product where lifting2 = Sub Dict
instance Lifting2 MonadFix Functor.Product where lifting2 = Sub Dict
instance Lifting2 Show1 Functor.Product where lifting2 = Sub Dict
instance Lifting2 Eq1 Functor.Product where lifting2 = Sub Dict
instance Lifting2 Ord1 Functor.Product where lifting2 = Sub Dict
instance Lifting2 Read1 Functor.Product where lifting2 = Sub Dict

instance Lifting2 Functor Functor.Sum where lifting2 = Sub Dict
instance Lifting2 Foldable Functor.Sum where lifting2 = Sub Dict
instance Lifting2 Traversable Functor.Sum where lifting2 = Sub Dict
instance Lifting2 Show1 Functor.Sum where lifting2 = Sub Dict
instance Lifting2 Eq1 Functor.Sum where lifting2 = Sub Dict
instance Lifting2 Ord1 Functor.Sum where lifting2 = Sub Dict
instance Lifting2 Read1 Functor.Sum where lifting2 = Sub Dict


{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE FunctionalDependencies #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE TemplateHaskell #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE UndecidableInstances #-}
{-# OPTIONS_HADDOCK show-extensions #-}

-- {-# LANGUAGE OverlappingInstances   #-}
-- |
-- Module      : Energy
-- Description : Particle Swarm Optimisation
-- Copyright   : (c) Tom Westerhout, 2017
-- License     : BSD3
-- Maintainer  : t.westerhout@student.ru.nl
-- Stability   : experimental
module PSO.Energy
  ( writeEnergies2TSV
  , writeVariances2TSV
  ) where

import Lens.Micro
import Lens.Micro.Extras
import Data.Monoid
import Data.List(intersperse)
import qualified Data.List as List
import qualified Numeric.LinearAlgebra as LA
import qualified Data.ByteString.Lazy as LB
import Data.ByteString.Builder

import PSO.Swarm

class EncodableAsDecimal a where
  encodeAsDecimal :: a -> Builder

instance EncodableAsDecimal Float where
  encodeAsDecimal = floatDec

instance EncodableAsDecimal Double where
  encodeAsDecimal = doubleDec

energies2TSV :: (EncodableAsDecimal t)
             => (Bee α β r -> t) -> [Swarm m γ β α r] -> LB.ByteString
energies2TSV func =
  let encodeSwarm = mconcat . intersperse (charUtf8 '\t')
        . map (encodeAsDecimal . func) . view bees
   in toLazyByteString . mconcat . intersperse (charUtf8 '\n')
        . map encodeSwarm

variances2TSV :: (EncodableAsDecimal r, HasVar γ r)
              => [Swarm m γ β α r] -> LB.ByteString
variances2TSV = toLazyByteString . mconcat . intersperse (charUtf8 '\n')
                . map (encodeAsDecimal . view (guide . var))

writeEnergies2TSV :: (EncodableAsDecimal t)
                  => FilePath -> (Bee α β r -> t) -> [Swarm m γ β α r] -> IO ()
writeEnergies2TSV file func = LB.writeFile file . energies2TSV func

writeVariances2TSV :: (EncodableAsDecimal r, HasVar γ r)
                   => FilePath -> [Swarm m γ β α r] -> IO ()
writeVariances2TSV file = LB.writeFile file . variances2TSV


{-# LANGUAGE TypeOperators, GADTs, CPP #-}
#ifndef NO_SAFE_HASKELL
{-# LANGUAGE Safe #-}
#endif
#if defined(__GLASGOW_HASKELL__) && __GLASGOW_HASKELL__ >= 708
{-# LANGUAGE PatternSynonyms, ViewPatterns #-}
#endif

#ifndef NO_GENERICS
{-# LANGUAGE DefaultSignatures, FlexibleContexts #-}
#endif

#ifndef NO_POLYKINDS
{-# LANGUAGE PolyKinds #-}
#endif

-- | Generation of random shrinkable, showable functions.
-- See the paper \"Shrinking and showing functions\" by Koen Claessen.
--
-- __Note__: most of the contents of this module are re-exported by
-- "Test.QuickCheck". You probably do not need to import it directly.
--
-- Example of use:
--
-- >>> :{
-- >>> let prop :: Fun String Integer -> Bool
-- >>>     prop (Fun _ f) = f "monkey" == f "banana" || f "banana" == f "elephant"
-- >>> :}
-- >>> quickCheck prop
-- *** Failed! Falsified (after 3 tests and 134 shrinks):
-- {"elephant"->1, "monkey"->1, _->0}
--
-- To generate random values of type @'Fun' a b@,
-- you must have an instance @'Function' a@.
-- If your type has a 'Show' instance, you can use 'functionShow' to write the instance; otherwise,
-- use 'functionMap' to give a bijection between your type and a type that is already an instance of 'Function'.
-- See the @'Function' [a]@ instance for an example of the latter.
module Test.QuickCheck.Function
  ( Fun(..)
  , applyFun
  , apply
  , applyFun2
  , applyFun3
  , (:->)
  , Function(..)
  , functionMap
  , functionShow
  , functionIntegral
  , functionRealFrac
  , functionBoundedEnum
#if defined(__GLASGOW_HASKELL__) && __GLASGOW_HASKELL__ >= 708
  , pattern Fn
  , pattern Fn2
  , pattern Fn3
#endif
  )
 where

--------------------------------------------------------------------------
-- imports

import Test.QuickCheck.Arbitrary
import Test.QuickCheck.Poly

import Control.Applicative
import Data.Char
import Data.Word
import Data.List( intersperse )
import Data.Ratio
import qualified Data.IntMap as IntMap
import qualified Data.IntSet as IntSet
import qualified Data.Map as Map
import qualified Data.Set as Set
import qualified Data.Sequence as Sequence
import Data.Int
import Data.Complex
import Data.Foldable(toList)
import Data.Functor.Identity
import qualified Data.Monoid as Monoid

#ifndef NO_FIXED
import Data.Fixed
#endif

#ifndef NO_GENERICS
import GHC.Generics hiding (C)
#endif

--------------------------------------------------------------------------
-- concrete functions

-- | The type of possibly partial concrete functions
data a :-> c where
  Pair  :: (a :-> (b :-> c)) -> ((a,b) :-> c)
  (:+:) :: (a :-> c) -> (b :-> c) -> (Either a b :-> c)
  Unit  :: c -> (() :-> c)
  Nil   :: a :-> c
  Table :: Eq a => [(a,c)] -> (a :-> c)
  Map   :: (a -> b) -> (b -> a) -> (b :-> c) -> (a :-> c)

instance Functor ((:->) a) where
  fmap f (Pair p)    = Pair (fmap (fmap f) p)
  fmap f (p:+:q)     = fmap f p :+: fmap f q
  fmap f (Unit c)    = Unit (f c)
  fmap f Nil         = Nil
  fmap f (Table xys) = Table [ (x,f y) | (x,y) <- xys ]
  fmap f (Map g h p) = Map g h (fmap f p)

instance (Show a, Show b) => Show (a:->b) where
  show p = showFunction p Nothing

-- only use this on finite functions
showFunction :: (Show a, Show b) => (a :-> b) -> Maybe b -> String
showFunction p md =
  "{" ++ concat (intersperse ", " ( [ show x ++ "->" ++ show c
                                    | (x,c) <- table p
                                    ]
                                 ++ [ "_->" ++ show d
                                    | Just d <- [md]
                                    ] )) ++ "}"

-- turning a concrete function into an abstract function (with a default result)
abstract :: (a :-> c) -> c -> (a -> c)
abstract (Pair p)    d (x,y) = abstract (fmap (\q -> abstract q d y) p) d x
abstract (p :+: q)   d exy   = either (abstract p d) (abstract q d) exy
abstract (Unit c)    _ _     = c
abstract Nil         d _     = d
abstract (Table xys) d x     = head ([y | (x',y) <- xys, x == x'] ++ [d])
abstract (Map g _ p) d x     = abstract p d (g x)

-- generating a table from a concrete function
table :: (a :-> c) -> [(a,c)]
table (Pair p)    = [ ((x,y),c) | (x,q) <- table p, (y,c) <- table q ]
table (p :+: q)   = [ (Left x, c) | (x,c) <- table p ]
                 ++ [ (Right y,c) | (y,c) <- table q ]
table (Unit c)    = [ ((), c) ]
table Nil         = []
table (Table xys) = xys
table (Map _ h p) = [ (h x, c) | (x,c) <- table p ]

--------------------------------------------------------------------------
-- Function

-- | The class @Function a@ is used for random generation of showable
-- functions of type @a -> b@.
--
-- There is a default implementation for 'function', which you can use
-- if your type has structural equality. Otherwise, you can normally
-- use 'functionMap' or 'functionShow'.
class Function a where
  function :: (a->b) -> (a:->b)
#ifndef NO_GENERICS
  default function :: (Generic a, GFunction (Rep a)) => (a->b) -> (a:->b)
  function = genericFunction
#endif

-- basic instances

-- | Provides a 'Function' instance for types with 'Bounded' and 'Enum'.
-- Use only for small types (i.e. not integers): creates
-- the list @['minBound'..'maxBound']@!
functionBoundedEnum :: (Eq a, Bounded a, Enum a) => (a->b) -> (a:->b)
functionBoundedEnum f = Table [(x,f x) | x <- [minBound..maxBound]]

-- | Provides a 'Function' instance for types with 'RealFrac'.
functionRealFrac :: RealFrac a => (a->b) -> (a:->b)
functionRealFrac = functionMap toRational fromRational

-- | Provides a 'Function' instance for types with 'Integral'.
functionIntegral :: Integral a => (a->b) -> (a:->b)
functionIntegral = functionMap fromIntegral fromInteger

-- | Provides a 'Function' instance for types with 'Show' and 'Read'.
functionShow :: (Show a, Read a) => (a->c) -> (a:->c)
functionShow f = functionMap show read f

-- | The basic building block for 'Function' instances.
-- Provides a 'Function' instance by mapping to and from a type that
-- already has a 'Function' instance.
functionMap :: Function b => (a->b) -> (b->a) -> (a->c) -> (a:->c)
functionMap = functionMapWith function

functionMapWith :: ((b->c) -> (b:->c)) -> (a->b) -> (b->a) -> (a->c) -> (a:->c)
functionMapWith function g h f = Map g h (function (\b -> f (h b)))

instance Function () where
  function f = Unit (f ())

instance Function a => Function (Const a b) where
  function = functionMap getConst Const

instance Function a => Function (Identity a) where
  function = functionMap runIdentity Identity

instance (Function a, Function b) => Function (a,b) where
  function = functionPairWith function function

functionPairWith :: ((a->b->c) -> (a:->(b->c))) -> ((b->c) -> (b:->c)) -> ((a,b)->c) -> ((a,b):->c)
functionPairWith func1 func2 f = Pair (func2 `fmap` func1 (curry f))

instance (Function a, Function b) => Function (Either a b) where
  function = functionEitherWith function function

functionEitherWith :: ((a->c) -> (a:->c)) -> ((b->c) -> (b:->c)) -> (Either a b->c) -> (Either a b:->c)
functionEitherWith func1 func2 f = func1 (f . Left) :+: func2 (f . Right)

-- tuple convenience instances

instance (Function a, Function b, Function c) => Function (a,b,c) where
  function = functionMap (\(a,b,c) -> (a,(b,c))) (\(a,(b,c)) -> (a,b,c))

instance (Function a, Function b, Function c, Function d) => Function (a,b,c,d) where
  function = functionMap (\(a,b,c,d) -> (a,(b,c,d))) (\(a,(b,c,d)) -> (a,b,c,d))

instance (Function a, Function b, Function c, Function d, Function e) => Function (a,b,c,d,e) where
  function = functionMap (\(a,b,c,d,e) -> (a,(b,c,d,e))) (\(a,(b,c,d,e)) -> (a,b,c,d,e))

instance (Function a, Function b, Function c, Function d, Function e, Function f) => Function (a,b,c,d,e,f) where
  function = functionMap (\(a,b,c,d,e,f) -> (a,(b,c,d,e,f))) (\(a,(b,c,d,e,f)) -> (a,b,c,d,e,f))

instance (Function a, Function b, Function c, Function d, Function e, Function f, Function g) => Function (a,b,c,d,e,f,g) where
  function = functionMap (\(a,b,c,d,e,f,g) -> (a,(b,c,d,e,f,g))) (\(a,(b,c,d,e,f,g)) -> (a,b,c,d,e,f,g))

-- other instances

instance Function a => Function [a] where
  function = functionMap g h
   where
    g []     = Left ()
    g (x:xs) = Right (x,xs)

    h (Left _)       = []
    h (Right (x,xs)) = x:xs

instance Function a => Function (Maybe a) where
  function = functionMap g h
   where
    g Nothing  = Left ()
    g (Just x) = Right x

    h (Left _)  = Nothing
    h (Right x) = Just x

instance Function Bool where
  function = functionMap g h
   where
    g False = Left ()
    g True  = Right ()

    h (Left _)  = False
    h (Right _) = True

instance Function Integer where
  function = functionMap gInteger hInteger
   where
    gInteger n | n < 0     = Left (gNatural (abs n - 1))
               | otherwise = Right (gNatural n)

    hInteger (Left ws)  = -(hNatural ws + 1)
    hInteger (Right ws) = hNatural ws

    gNatural 0 = []
    gNatural n = (fromIntegral (n `mod` 256) :: Word8) : gNatural (n `div` 256)

    hNatural []     = 0
    hNatural (w:ws) = fromIntegral w + 256 * hNatural ws

instance Function Int where
  function = functionIntegral

instance Function Word where
  function = functionIntegral

instance Function Char where
  function = functionMap ord chr

instance Function Float where
  function = functionRealFrac

instance Function Double where
  function = functionRealFrac

-- instances for assorted types in the base package

instance Function Ordering where
  function = functionMap g h
    where
      g LT = Left False
      g EQ = Left True
      g GT = Right ()

      h (Left False) = LT
      h (Left True)  = EQ
      h (Right _)    = GT

instance (Integral a, Function a) => Function (Ratio a) where
  function = functionMap g h
   where
     g r = (numerator r, denominator r)
     h (n, d) = n % d

#ifndef NO_FIXED
instance HasResolution a => Function (Fixed a) where
  function = functionRealFrac
#endif

instance (RealFloat a, Function a) => Function (Complex a) where
  function = functionMap g h
   where
     g (x :+ y) = (x,   y)
     h (x,   y) =  x :+ y

instance (Ord a, Function a) => Function (Set.Set a) where
  function = functionMap Set.toList Set.fromList

instance (Ord a, Function a, Function b) => Function (Map.Map a b) where
  function = functionMap Map.toList Map.fromList

instance Function IntSet.IntSet where
  function = functionMap IntSet.toList IntSet.fromList

instance Function a => Function (IntMap.IntMap a) where
  function = functionMap IntMap.toList IntMap.fromList

instance Function a => Function (Sequence.Seq a) where
  function = functionMap toList Sequence.fromList

instance Function Int8 where
  function = functionBoundedEnum

instance Function Int16 where
  function = functionIntegral

instance Function Int32 where
  function = functionIntegral

instance Function Int64 where
  function = functionIntegral

instance Function Word8 where
  function = functionBoundedEnum

instance Function Word16 where
  function = functionIntegral

instance Function Word32 where
  function = functionIntegral

instance Function Word64 where
  function = functionIntegral

-- instances for Data.Monoid newtypes

instance Function a => Function (Monoid.Dual a) where
  function = functionMap Monoid.getDual Monoid.Dual

instance Function Monoid.All where
  function = functionMap Monoid.getAll Monoid.All

instance Function Monoid.Any where
  function = functionMap Monoid.getAny Monoid.Any

instance Function a => Function (Monoid.Sum a) where
  function = functionMap Monoid.getSum Monoid.Sum

instance Function a => Function (Monoid.Product a) where
  function = functionMap Monoid.getProduct Monoid.Product

instance Function a => Function (Monoid.First a) where
  function = functionMap Monoid.getFirst Monoid.First

instance Function a => Function (Monoid.Last a) where
  function = functionMap Monoid.getLast Monoid.Last

#if MIN_VERSION_base(4,8,0)
instance Function (f a) => Function (Monoid.Alt f a) where
  function = functionMap Monoid.getAlt Monoid.Alt
#endif

-- poly instances

instance Function A where
  function = functionMap unA A

instance Function B where
  function = functionMap unB B

instance Function C where
  function = functionMap unC C

instance Function OrdA where
  function = functionMap unOrdA OrdA

instance Function OrdB where
  function = functionMap unOrdB OrdB

instance Function OrdC where
  function = functionMap unOrdC OrdC

-- instance Arbitrary

instance (Function a, CoArbitrary a, Arbitrary b) => Arbitrary (a:->b) where
  arbitrary = function `fmap` arbitrary
  shrink    = shrinkFun shrink

--------------------------------------------------------------------------
-- generic function instances

#ifndef NO_GENERICS
-- | Generic 'Function' implementation.
genericFunction :: (Generic a, GFunction (Rep a)) => (a->b) -> (a:->b)
genericFunction = functionMapWith gFunction from to

class GFunction f where
  gFunction :: (f a -> b) -> (f a :-> b)

instance GFunction U1 where
  gFunction = functionMap (\U1 -> ()) (\() -> U1)

instance (GFunction f, GFunction g) => GFunction (f :*: g) where
  gFunction = functionMapWith (functionPairWith gFunction gFunction) g h
   where
     g (x :*: y) = (x, y)
     h (x, y) = x :*: y

instance (GFunction f, GFunction g) => GFunction (f :+: g) where
  gFunction = functionMapWith (functionEitherWith gFunction gFunction) g h
   where
     g (L1 x) = Left x
     g (R1 x) = Right x
     h (Left x) = L1 x
     h (Right x) = R1 x

instance GFunction f => GFunction (M1 i c f) where
  gFunction = functionMapWith gFunction (\(M1 x) -> x) M1

instance Function a => GFunction (K1 i a) where
  gFunction = functionMap (\(K1 x) -> x) K1
#endif

--------------------------------------------------------------------------
-- shrinking

shrinkFun :: (c -> [c]) -> (a :-> c) -> [a :-> c]
shrinkFun shr (Pair p) =
  [ pair p' | p' <- shrinkFun (\q -> shrinkFun shr q) p ]
 where
  pair Nil = Nil
  pair p   = Pair p

shrinkFun shr (p :+: q) =
  [ p .+. Nil | not (isNil q) ] ++
  [ Nil .+. q | not (isNil p) ] ++
  [ p  .+. q' | q' <- shrinkFun shr q ] ++
  [ p' .+. q  | p' <- shrinkFun shr p ]
 where
  isNil :: (a :-> b) -> Bool
  isNil Nil = True
  isNil _   = False

  Nil .+. Nil = Nil
  p   .+. q   = p :+: q

shrinkFun shr (Unit c) =
  [ Nil ] ++
  [ Unit c' | c' <- shr c ]

shrinkFun shr (Table xys) =
  [ table xys' | xys' <- shrinkList shrXy xys ]
 where
  shrXy (x,y) = [(x,y') | y' <- shr y]

  table []  = Nil
  table xys = Table xys

shrinkFun shr Nil =
  []

shrinkFun shr (Map g h p) =
  [ mapp g h p' | p' <- shrinkFun shr p ]
 where
  mapp g h Nil = Nil
  mapp g h p   = Map g h p

--------------------------------------------------------------------------
-- the Fun modifier

-- | Generation of random shrinkable, showable functions.
--
-- To generate random values of type @'Fun' a b@,
-- you must have an instance @'Function' a@.
--
-- See also 'applyFun', and 'Fn' with GHC >= 7.8.
data Fun a b = Fun (a :-> b, b, Shrunk) (a -> b)
data Shrunk = Shrunk | NotShrunk deriving Eq

instance Functor (Fun a) where
  fmap f (Fun (p, d, s) g) = Fun (fmap f p, f d, s) (f . g)

#if defined(__GLASGOW_HASKELL__) && __GLASGOW_HASKELL__ >= 708
-- | A modifier for testing functions.
--
-- > prop :: Fun String Integer -> Bool
-- > prop (Fn f) = f "banana" == f "monkey"
-- >            || f "banana" == f "elephant"
#if __GLASGOW_HASKELL__ >= 800
pattern Fn :: (a -> b) -> Fun a b
#endif
pattern Fn f <- (applyFun -> f)

-- | A modifier for testing binary functions.
--
-- > prop_zipWith :: Fun (Int, Bool) Char -> [Int] -> [Bool] -> Bool
-- > prop_zipWith (Fn2 f) xs ys = zipWith f xs ys == [ f x y | (x, y) <- zip xs ys]
#if __GLASGOW_HASKELL__ >= 800
pattern Fn2 :: (a -> b -> c) -> Fun (a, b) c
#endif
pattern Fn2 f <- (applyFun2 -> f)

-- | A modifier for testing ternary functions.
#if __GLASGOW_HASKELL__ >= 800
pattern Fn3 :: (a -> b -> c -> d) -> Fun (a, b, c) d
#endif
pattern Fn3 f <- (applyFun3 -> f)
#endif

mkFun :: (a :-> b) -> b -> Fun a b
mkFun p d = Fun (p, d, NotShrunk) (abstract p d)

-- | Alias to 'applyFun'.
apply :: Fun a b -> (a -> b)
apply = applyFun

-- | Extracts the value of a function.
--
-- 'Fn' is the pattern equivalent of this function.
--
-- > prop :: Fun String Integer -> Bool
-- > prop f = applyFun f "banana" == applyFun f "monkey"
-- >       || applyFun f "banana" == applyFun f "elephant"
applyFun :: Fun a b -> (a -> b)
applyFun (Fun _ f) = f

-- | Extracts the value of a binary function.
--
-- 'Fn2' is the pattern equivalent of this function.
--
--  > prop_zipWith :: Fun (Int, Bool) Char -> [Int] -> [Bool] -> Bool
--  > prop_zipWith f xs ys = zipWith (applyFun2 f) xs ys == [ applyFun2 f x y | (x, y) <- zip xs ys]
--
applyFun2 :: Fun (a, b) c -> (a -> b -> c)
applyFun2 (Fun _ f) a b = f (a, b)

-- | Extracts the value of a ternary function. 'Fn3' is the
-- pattern equivalent of this function.
applyFun3 :: Fun (a, b, c) d -> (a -> b -> c -> d)
applyFun3 (Fun _ f) a b c = f (a, b, c)

instance (Show a, Show b) => Show (Fun a b) where
  show (Fun (_, _, NotShrunk) _) = "<fun>"
  show (Fun (p, d, Shrunk) _)    = showFunction p (Just d)

instance (Function a, CoArbitrary a, Arbitrary b) => Arbitrary (Fun a b) where
  arbitrary =
    do p <- arbitrary
       d <- arbitrary
       return (mkFun p d)

  shrink (Fun (p, d, s) f) =
    [ mkFun p' d' | (p', d') <- shrink (p, d) ] ++
    [ Fun (p, d, Shrunk) f | s == NotShrunk ]

--------------------------------------------------------------------------
-- the end.


{-# LANGUAGE NumericUnderscores, OverloadedStrings #-}

module Main where

import Test.QuickCheck

import qualified Data.Vector.Unboxed as U
import qualified Data.Vector as V
-- import Data.Vector.Unboxed  ((!))
import Data.Time (fromGregorian, UTCTime)
import Data.Time.Format (parseTimeM, defaultTimeLocale)
import qualified Data.HashMap.Strict as Hm
import Data.Text (Text)
import qualified Data.Complex as Cx
import Data.Complex (Complex ( (:+) ) )

import Data.Approx

main :: IO ()
main = do

  -- print "Float equality test"
  quickCheck $ (0.0 :: Float) =~ (-1e-9)
  quickCheck $ (0.0 :: Double) =~ (-0.0)
  quickCheck $ (1.0e+7 :: Double) =~ 10_000_000.05

  quickCheck $ Just (10.0000000000000007 :: Double) =~ Just 10.0
  quickCheck $ Just (10 :: Double) /~ Just 1.0
  quickCheck $ Just (10 :: Double) /~ Nothing
  quickCheck $ (Nothing :: Maybe Double) /~ Just 1.0
  quickCheck $ (Nothing :: Maybe Double) =~ Nothing

  quickCheck $ ([1.2, 3.4, 5.6] :: [Double]) /~ [1.2, 3.4, 5.65]
  quickCheck $ ([2.3, 1,2, 1000] :: [Double]) =~ [2.3, 1,2, 1e+3]
  quickCheck $ ([1.0,1.0,1.0,1.0] :: [Double]) /~ [1.0,1.0,1.0,1.0,2.0]

  quickCheck $ ((1.2, 3.4) :: (Double, Double)) =~ (1.20000008, 3.399999999)
  quickCheck $ ((1.2, 3.5) :: (Double, Double)) /~ (1.20000008, 3.399999999)

  quickCheck $ ((1.2, 3.4, 2.5) :: (Double, Double, Double)) =~ (1.20000008, 3.399999999, 2.5000001)
  quickCheck $ ((1.2, 3.4, 2.5) :: (Double, Double, Double)) /~ (1.2, 3.399999999, 2.4)

  quickCheck $ ((1.2, 3.4, 2.5) :: (Float, Float, Float)) =~ (1.20000008, 3.399999999, 2.5000001)
  quickCheck $ ((1.2, 3.4, 2.5) :: (Float, Float, Float)) /~ (1.2, 3.399999999, 2.4)

  quickCheck $ ((1, 3, 2) :: (Int, Int, Int)) =~ (1, 3, 2)
  quickCheck $ ((1, 3, 2) :: (Int, Int, Int)) /~ (1, 3, 5)

  quickCheck $ ((2,5.35) :: (Int,Double)) =~ (2,5.35)

  quickCheck $ ((3.0 :+ 5.0)::(Cx.Complex Float)) =~ (3.0 :+ 5.0)

  quickCheck $ ((2,5.35,"happ",[1,2],["ret","we"],[6.78,3.5])
    :: (Int,Double,Text,[Int],[[Char]],[Double]))

    =~ (2,5.35,"happ",[1,2],["ret","we"],[6.78,3.5])

  quickCheck $ ((1.0,2.0,3.0,4.0,5.0,6.0,7.0,8.0,9.0,10.0,11.0) :: (Double,Double,Double,Double,Double,Double,Double,Double,Double,Double,Double)) =~ (1.0,2.0,3.0,4.0,5.0,6.0,7.0,8.0,9.0,10.0,11.0)

  quickCheck $ ((1.0,2.0,3.0,4.0,5.0,6.0,7.0,8.0,9.0,10.0,11.1) :: (Double,Double,Double,Double,Double,Double,Double,Double,Double,Double,Double))   /~ (1.0,2.0,3.0,4.0,5.0,6.0,7.0,8.0,9.0,10.0,11.0)

  quickCheck $ ((1.0,2.0,3.0,4.0,5.0,6.0,7.0,8.0,9.0,10.0) :: (Double,Double,Double,Double,Double,Double,Double,Double,Double,Double)) =~ (1.0,2.0,3.0,4.0,5.0,6.0,7.0,8.0,9.0,10.0)

  quickCheck $ ((1.0,2.0,3.0,4.0,5.0,6.0,7.0,8.0,9.0,10.0) :: (Double,Double,Double,Double,Double,Double,Double,Double,Double,Double))   /~ (1.0,2.0,3.0,4.0,5.0,6.0,7.0,8.0,9.0,10.1)

  quickCheck $ False /~ True
  quickCheck $ True =~ True
  quickCheck $ False =~ False

  quickCheck $ ("star"::Text) =~ ("star"::Text)
  quickCheck $ ("star"::Text) /~ ("start"::Text)
  quickCheck $ ("start"::Text) /~ ("star"::Text)

  quickCheck $ ("star"::String) =~ ("star"::String)
  quickCheck $ ("star"::String) /~ ("start"::String)
  quickCheck $ ("start"::String) /~ ("star"::String)

  quickCheck $ fromGregorian 2018 3 30 =~ fromGregorian 2018 3 30
  quickCheck $ fromGregorian 2018 3 31 /~ fromGregorian 2018 3 30

  quickCheck $ (parseTimeM True defaultTimeLocale "%Y-%m-%d %H:%M:%S" "2020-01-15 15:02:15" :: Maybe UTCTime) /~ parseTimeM True defaultTimeLocale "%Y-%m-%d %H:%M:%S" "2020-01-15 14:02:15"

  quickCheck $ (parseTimeM True defaultTimeLocale "%Y-%m-%d %H:%M:%S" "2020-01-15 15:02:15" :: Maybe UTCTime) =~ parseTimeM True defaultTimeLocale "%Y-%m-%d %H:%M:%S" "2020-01-15 15:01:50"

  quickCheck $ (parseTimeM True defaultTimeLocale "%Y-%m-%d %H:%M:%S" "2020-01-15 15:01:50" :: Maybe UTCTime) =~ parseTimeM True defaultTimeLocale "%Y-%m-%d %H:%M:%S" "2020-01-15 15:02:15"
  quickCheck $ (parseTimeM True defaultTimeLocale "%Y-%m-%d %H:%M:%S" "2020-01-16 15:02:15" :: Maybe UTCTime) /~ parseTimeM True defaultTimeLocale "%Y-%m-%d %H:%M:%S" "2020-01-15 15:02:15"

  quickCheck $ (parseTimeM True defaultTimeLocale "%Y-%m-%d %H:%M:%S" "2019-01-15 15:02:15" :: Maybe UTCTime) /~ parseTimeM True defaultTimeLocale "%Y-%m-%d %H:%M:%S" "2020-01-15 15:02:15"

  quickCheck $ (parseTimeM True defaultTimeLocale "%Y-%m-%d %H:%M:%S" "2020-01-15 15:02:15" :: Maybe UTCTime) =~ parseTimeM True defaultTimeLocale "%Y-%m-%d %H:%M:%S" "2020-01-15 15:02:15"

  let rb = Hm.fromList [("Cash", 10.23), ("CurrentAdvances", 56.25), ("AccountPayables", 0.0)] :: Hm.HashMap String Double
  let tb = Hm.fromList [("Cash", 10.23), ("CurrentAdvances", 56.25)] :: Hm.HashMap String Double

  quickCheck $ rb /~ tb
  quickCheck $ tb /~ rb

  let rb = Hm.fromList [("Cash", 10.23), ("CurrentAdvances", 56.25)] :: Hm.HashMap String Double

  quickCheck $ rb =~ tb
  quickCheck $ tb =~ rb

  let rb = Hm.fromList [("Cash", 10.21), ("CurrentAdvances", 56.25)] :: Hm.HashMap String Double

  quickCheck $ rb /~ tb
  quickCheck $ tb /~ rb

  let rb = Hm.fromList [("Cash", Just 10.23), ("CurrentAdvances", Just 56.25), ("AccountPayables", Just 0.0)] :: Hm.HashMap String (Maybe Double)
  let tb = Hm.fromList [("Cash", Just 10.23), ("CurrentAdvances", Just 56.25)] :: Hm.HashMap String (Maybe Double)

  quickCheck $ rb /~ tb
  quickCheck $ tb /~ rb

  let rb = Hm.fromList [("Cash", Just 10.23), ("CurrentAdvances", Just 56.25)] :: Hm.HashMap String (Maybe Double)

  quickCheck $ rb =~ tb
  quickCheck $ tb =~ rb

  let rb = V.fromList [("Cash", Just 10.23), ("CurrentAdvances", Just 56.25), ("AccountPayables", Just 0.0)] :: V.Vector (String, Maybe Double)
  let tb = V.fromList [("Cash", Just 10.23), ("CurrentAdvances", Just 56.25)] :: V.Vector (String, Maybe Double)

  quickCheck $ rb /~ tb

  let rb = V.fromList [("Cash", Just 10.24), ("CurrentAdvances", Just 56.25)] :: V.Vector (String, Maybe Double)

  quickCheck $ rb /~ tb

  quickCheck $ U.fromList ([2.5, 1.5, 3.56] :: [Double]) =~ U.fromList [2.500000001, 1.5, 3.56]

  quickCheck $ (2,4) `inRange` (3::Int)
  quickCheck $ not $ (2,4) `inRange` (5::Int)

  quickCheck $ (2.5, 2.75) `inRange` (2.62::Double)
  quickCheck $ not $ (2.5, 2.75) `inRange` (2.8::Double)

  quickCheck $ (4,2) `safeInRange` (3::Int)
  quickCheck $ not $ (2,4) `safeInRange` (5::Int)

  quickCheck $ (2.5, 2.75) `safeInRange` (2.62::Double)
  quickCheck $ not $ (2.75, 2.5) `safeInRange` (2.8::Double)

  quickCheck $ inTol 0.25 6.25 (6.1::Double)
  quickCheck $ not $ inTol 0.125 6.25 (6.1::Double)