{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE UndecidableInstances #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE FunctionalDependencies #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE NamedFieldPuns #-}
{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE ViewPatterns #-}

module Language.Grappa.Interp.StandardHORepr where

import Data.Vector (Vector)
import qualified Data.Vector as V
import Control.Monad

import Numeric.LinearAlgebra hiding (R, Uniform, (<>), Vector)

import Language.Grappa.Distribution
import Language.Grappa.Interp
import Language.Grappa.GrappaInternals
import Language.Grappa.Frontend.DataSource

import qualified Numeric.Log as Log

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

----------------------------------------------------------------------
-- * The StandardHORepr Representation
----------------------------------------------------------------------

-- | The type tag for the standard higher-order representation, which is
-- parameterized by a monad and by the representation types for reals and ints
data StandardHORepr (m :: * -> *) (r :: *) (i :: *) :: *

-- | The type family for 'StandardHORepr' expressions
type family StandardHOReprF m r i a :: * where
  StandardHOReprF m r i (a -> b) =
    (StandardHOReprF m r i a -> StandardHOReprF m r i b)
  StandardHOReprF m r i (Dist a) =
    (DistVar (StandardHORepr m r i) a -> m (StandardHOReprF m r i a))
  StandardHOReprF m r i (ADT adt) =
    adt (GExpr (StandardHORepr m r i)) (ADT adt)
  StandardHOReprF m r i (Vector a) = Vector (StandardHOReprF m r i a)
  StandardHOReprF m r i Bool    = Bool
  StandardHOReprF m r i Int     = i
  StandardHOReprF m r i Prob    = Log.Log r
  StandardHOReprF m r i R       = r
  StandardHOReprF m r i RMatrix = Matrix r
  StandardHOReprF m r i a       = a

instance ValidExprRepr (StandardHORepr m r i) where
  type GExprRepr (StandardHORepr m r i) a = StandardHOReprF m r i a
  interp__'bottom = error "StandardHORepr: unexpected bottom!"
  interp__'injTuple !tup = GExpr tup
  interp__'projTuple (GExpr !tup) k = k tup
  interp__'app (GExpr !f) (GExpr x) = GExpr (f x)
  interp__'lam f = GExpr (unGExpr . f . GExpr)
  interp__'fix f = f (interp__'fix f)

instance StrongTupleRepr (StandardHORepr m r i) where
  interp__'strongProjTuple (GExpr tup) = tup

-- | Helper to match on v-expressions of atomic type in the 'StandardHORepr': if
-- the 'DistVar' is not a 'VParam', then destructure it and pass it to the
-- continuation (the 2nd argument); otherwise, return the failure continuation
-- (the 3rd argument)
matchHOReprAtomicDistVar ::
  (IsAtomic a ~ 'True, EmbedRepr (StandardHORepr m r i) a) =>
  DistVar (StandardHORepr m r i) a ->
  (GExpr (StandardHORepr m r i) a -> ret) -> ret -> ret
matchHOReprAtomicDistVar VParam _ ret = ret
matchHOReprAtomicDistVar (VData (GData a)) k _ = k $ embedRepr a
matchHOReprAtomicDistVar (VData GNoData) _ ret = ret
matchHOReprAtomicDistVar (VExpr e) k _ = k e

-- | Test if a v-expression of atomic type is a "missing value"
isMissingHOReprAtomicDistVar ::
  (IsAtomic a ~ 'True, EmbedRepr (StandardHORepr m r i) a) =>
  DistVar (StandardHORepr m r i) a -> Bool
isMissingHOReprAtomicDistVar dv =
  matchHOReprAtomicDistVar dv (\_ -> False) True

-- | Helper to match on v-expressions in the 'StandardHORepr': if the 'DistVar'
-- is not a 'VParam', then destructure it and pass it to the continuation (the
-- 2nd argument); otherwise, return the failure continuation (the 3rd argument)
matchHOReprADTDistVar ::
  TraversableADT adt =>
  DistVar (StandardHORepr m r i) (ADT adt) ->
  (adt (DistVar (StandardHORepr m r i)) (ADT adt) -> ret) -> ret -> ret
matchHOReprADTDistVar VParam _ ret = ret
matchHOReprADTDistVar (VData (GData (ADT adt))) k _ =
  k $ mapADT (VData . GData . unId) adt
matchHOReprADTDistVar (VData GNoData) _ ret = ret
matchHOReprADTDistVar (VData (GADTData adt)) k _ = k $ mapADT VData adt
matchHOReprADTDistVar (VExpr (GExpr adt)) k _ = k $ mapADT VExpr adt
matchHOReprADTDistVar (VADT adt) k _ = k adt

-- | Recursively match a v-expression with list type, returning a list of
-- 'DistVar's in the list along with a 'Bool' flag indicating whether the list
-- ends with a "missing list", i.e., a 'VParam' or @'VData' 'GNoData'@
matchHOReprListDistVar :: DistVar (StandardHORepr m r i) (GList a) ->
                          ([DistVar (StandardHORepr m r i) a], Bool)
matchHOReprListDistVar dv =
  matchHOReprADTDistVar dv
  (\adt -> case adt of
      Nil -> ([], False)
      Cons hdv tlv ->
        let (l, flag) = matchHOReprListDistVar tlv in
        (hdv:l, flag))
  ([], True)

matchHOReprVectorDistVar ::
  DistVar (StandardHORepr m r i) (Vector a) ->
  (Vector (DistVar (StandardHORepr m r i) a) -> ret) -> ret -> ret
matchHOReprVectorDistVar VParam _ ret = ret
matchHOReprVectorDistVar (VData (GData a)) k _ = k $ V.map (VData . GData) a
matchHOReprVectorDistVar (VData GNoData) _ ret = ret
matchHOReprVectorDistVar (VExpr (GExpr dvs)) k _ = k $ V.map (VExpr . GExpr) dvs

instance Monad m => ValidRepr (StandardHORepr m r i) where
  type GVExprRepr (StandardHORepr m r i) a =
    DistVar (StandardHORepr m r i) a
  type GStmtRepr (StandardHORepr m r i) a = m (StandardHOReprF m r i a)

  interp__'projTupleStmt (GExpr !tup) k = k tup

  interp__'vInjTuple !tup = GVExpr (VADT $ mapADT unGVExpr tup)
  interp__'vProjTuple (GVExpr ve) k =
    matchHOReprADTDistVar ve (k . mapADT GVExpr)
    (k $ buildTuple $ GVExpr VParam)

  interp__'vwild k = k (GVExpr VParam)
  interp__'vlift e k = k (GVExpr $ VExpr e)

  interp__'return (GExpr !x) = GStmt (return x)
  interp__'let rhs body = body rhs
  interp__'sample (GExpr !d) (GVExpr !dv) k = GStmt $ do
    !x <- d dv
    unGStmt $ k (GExpr x)

  interp__'mkDist f = GExpr (\ dv -> unGStmt $ f $ GVExpr dv)

instance TraversableADT adt =>
         Interp__ADT__Expr (StandardHORepr m r i) adt where
  interp__'injADT adt = GExpr adt
  interp__'projADT (GExpr adt) k = k adt
  interp__'projMatchADT (GExpr adt) _ matcher k_succ k_fail =
    if applyCtorMatcher matcher adt then k_succ adt else k_fail

instance (Monad m, TraversableADT adt) =>
         Interp__ADT (StandardHORepr m r i) adt where
  interp__'vInjADT adt =
    GVExpr (VADT $ mapADT unGVExpr adt)
  interp__'vProjMatchADT (GVExpr ve) ctor matcher k_succ k_fail =
    matchHOReprADTDistVar ve
    (\adt ->
      if applyCtorMatcher matcher adt then
        k_succ (mapADT GVExpr adt)
      else k_fail)
    (k_succ $ mapADT (const $ GVExpr VParam) ctor)

instance Interp__'source (StandardHORepr m Double Int) a where
  interp__'source src =
    GVExpr . VData <$> interpSource src

instance EmbedRepr (StandardHORepr m Double i) R where
  embedRepr = GExpr

instance EmbedRepr (StandardHORepr m Double i) Prob where
  embedRepr = GExpr . fromProb

instance ADR.Reifies s ADR.Tape =>
         EmbedRepr (StandardHORepr m (ADR.Reverse s Double) i) R where
  embedRepr = GExpr . ADR.auto

instance ADR.Reifies s ADR.Tape =>
         EmbedRepr (StandardHORepr m (ADR.Reverse s Double) i) Prob where
  embedRepr = GExpr . fmap ADR.auto . fromProb

instance Num i => EmbedRepr (StandardHORepr m r i) Int where
  embedRepr = GExpr . fromIntegral

instance EmbedRepr (StandardHORepr m r i) Bool where
  embedRepr = GExpr

instance MapC (EmbedRepr (StandardHORepr m r i)) ts =>
         EmbedRepr (StandardHORepr m r i) (ADT (TupleF ts)) where
  embedRepr = GExpr . helper . unADT where
    helper :: MapC (EmbedRepr (StandardHORepr m r i)) ts' =>
              TupleF ts' Id any -> TupleF ts' (GExpr (StandardHORepr m r i)) any
    helper (Tuple0) = Tuple0
    helper (Tuple1 (Id a)) = Tuple1 (embedRepr a)
    helper (Tuple2 (Id a) (Id b)) =
      Tuple2 (embedRepr a) (embedRepr b)
    helper (Tuple3 (Id a) (Id b) (Id c)) =
      Tuple3 (embedRepr a) (embedRepr b) (embedRepr c)
    helper (Tuple4 (Id a) (Id b) (Id c) (Id d)) =
      Tuple4 (embedRepr a) (embedRepr b) (embedRepr c) (embedRepr d)
    helper (TupleN (Id a) (Id b) (Id c) (Id d) (Id e) tup) =
      TupleN (embedRepr a) (embedRepr b) (embedRepr c) (embedRepr d)
      (embedRepr e) (helper tup)

instance EmbedRepr (StandardHORepr m Double i) a =>
         EmbedRepr (StandardHORepr m Double i) (Vector a) where
  embedRepr xs = GExpr $ V.map (unGExpr . helper) xs where
    helper :: EmbedRepr (StandardHORepr m Double i) a => a ->
              GExpr (StandardHORepr m Double i) a
    helper = embedRepr

----------------------------------------------------------------------
-- Boolean and comparison operations
----------------------------------------------------------------------

instance Interp__'ifThenElse (StandardHORepr m r i) where
  interp__'ifThenElse (GExpr c) t e = if c then t else e

instance Monad m => Interp__'vmatchSwitch (StandardHORepr m r Int) where
  interp__'vmatchSwitch (GExpr i) stmts = stmts !! i

instance Interp__not (StandardHORepr m r i) where
  interp__not = GExpr not

instance Interp__'amp'amp (StandardHORepr m r i) where
  interp__'amp'amp = GExpr (&&)

instance Interp__'bar'bar (StandardHORepr m r i) where
  interp__'bar'bar = GExpr (||)

instance (Eq a, Eq (StandardHOReprF m r i a)) =>
         Interp__'eq'eq (StandardHORepr m r i) a where
  interp__'eq'eq = GExpr (==)

instance (Ord a, Ord (StandardHOReprF m r i a)) =>
         Interp__'lt (StandardHORepr m r i) a where
  interp__'lt = GExpr (<)

instance (Ord a, Ord (StandardHOReprF m r i a)) =>
         Interp__'gt (StandardHORepr m r i) a where
  interp__'gt = GExpr (>)

instance (Ord a, Ord (StandardHOReprF m r i a)) =>
         Interp__'lt'eq (StandardHORepr m r i) a where
  interp__'lt'eq = GExpr (<=)

instance (Ord a, Ord (StandardHOReprF m r i a)) =>
         Interp__'gt'eq (StandardHORepr m r i) a where
  interp__'gt'eq = GExpr (>=)

instance (Ord a, Ord (StandardHOReprF m r i a)) =>
         Interp__min (StandardHORepr m r i) a where
  interp__min = GExpr min

instance (Ord a, Ord (StandardHOReprF m r i a)) =>
         Interp__max (StandardHORepr m r i) a where
  interp__max = GExpr max

----------------------------------------------------------------------
-- Numeric Operations
----------------------------------------------------------------------

instance (Num a, Num (StandardHOReprF m r i a)) =>
         Interp__'plus (StandardHORepr m r i) a where
  interp__'plus = GExpr (+)

instance (Num a, Num (StandardHOReprF m r i a)) =>
         Interp__'minus (StandardHORepr m r i) a where
  interp__'minus = GExpr (-)

instance (Num a, Num (StandardHOReprF m r i a)) =>
         Interp__'times (StandardHORepr m r i) a where
  interp__'times = GExpr (*)

instance (Num a, Num (StandardHOReprF m r i a)) =>
         Interp__negate (StandardHORepr m r i) a where
  interp__negate = GExpr negate

instance (Num a, Num (StandardHOReprF m r i a)) =>
         Interp__abs (StandardHORepr m r i) a where
  interp__abs = GExpr abs

instance (Num a, Num (StandardHOReprF m r i a)) =>
         Interp__signum (StandardHORepr m r i) a where
  interp__signum = GExpr signum

instance (Num a, Num (StandardHOReprF m r i a)) =>
         Interp__fromInteger (StandardHORepr m r i) a where
  interp__fromInteger = GExpr fromInteger

instance (Num a, Num (StandardHOReprF m r i a)) =>
         Interp__'integer (StandardHORepr m r i) a where
  interp__'integer n = GExpr (fromInteger n)

instance (Interp__'integer (StandardHORepr m r i) a,
          Eq (StandardHOReprF m r i a))
         => Interp__'eqInteger (StandardHORepr m r i) a where
  interp__'eqInteger (GExpr x) (GExpr y) = GExpr (x == y)

instance (Fractional a, Fractional (StandardHOReprF m r i a)) =>
         Interp__'div (StandardHORepr m r i) a where
  interp__'div = GExpr (/)

instance (Fractional a, Fractional (StandardHOReprF m r i a)) =>
         Interp__recip (StandardHORepr m r i) a where
  interp__recip = GExpr recip

instance (Fractional a, Fractional (StandardHOReprF m r i a)) =>
         Interp__fromRational (StandardHORepr m r i) a where
  interp__fromRational = GExpr fromRational

instance (Fractional a, Fractional (StandardHOReprF m r i a)) =>
         Interp__'rational (StandardHORepr m r i) a where
  interp__'rational n = GExpr (fromRational n)

instance (Interp__'rational (StandardHORepr m r i) a,
          Eq (StandardHOReprF m r i a))
         => Interp__'eqRational (StandardHORepr m r i) a where
  interp__'eqRational (GExpr x) (GExpr y) = GExpr (x == y)

instance (Floating a, Floating (StandardHOReprF m r i a)) =>
         Interp__pi (StandardHORepr m r i) a where
  interp__pi = GExpr pi

instance (Floating a, Floating (StandardHOReprF m r i a)) =>
         Interp__exp (StandardHORepr m r i) a where
  interp__exp = GExpr exp

instance (Floating a, Floating (StandardHOReprF m r i a)) =>
         Interp__log (StandardHORepr m r i) a where
  interp__log = GExpr log

instance (Floating a, Floating (StandardHOReprF m r i a)) =>
         Interp__sqrt (StandardHORepr m r i) a where
  interp__sqrt = GExpr sqrt

instance (Floating a, Floating (StandardHOReprF m r i a)) =>
         Interp__'times'times (StandardHORepr m r i) a where
  interp__'times'times = GExpr (**)

instance (Floating a, Floating (StandardHOReprF m r i a)) =>
         Interp__logBase (StandardHORepr m r i) a where
  interp__logBase = GExpr logBase

instance (Floating a, Floating (StandardHOReprF m r i a)) =>
         Interp__sin (StandardHORepr m r i) a where
  interp__sin = GExpr sin

instance (Floating a, Floating (StandardHOReprF m r i a)) =>
         Interp__cos (StandardHORepr m r i) a where
  interp__cos = GExpr cos

instance (Floating a, Floating (StandardHOReprF m r i a)) =>
         Interp__tan (StandardHORepr m r i) a where
  interp__tan = GExpr tan

instance (Floating a, Floating (StandardHOReprF m r i a)) =>
         Interp__asin (StandardHORepr m r i) a where
  interp__asin = GExpr asin

instance (Floating a, Floating (StandardHOReprF m r i a)) =>
         Interp__acos (StandardHORepr m r i) a where
  interp__acos = GExpr acos

instance (Floating a, Floating (StandardHOReprF m r i a)) =>
         Interp__atan (StandardHORepr m r i) a where
  interp__atan = GExpr atan

instance (Floating a, Floating (StandardHOReprF m r i a)) =>
         Interp__sinh (StandardHORepr m r i) a where
  interp__sinh = GExpr sinh

instance (Floating a, Floating (StandardHOReprF m r i a)) =>
         Interp__cosh (StandardHORepr m r i) a where
  interp__cosh = GExpr cosh

instance (Floating a, Floating (StandardHOReprF m r i a)) =>
         Interp__tanh (StandardHORepr m r i) a where
  interp__tanh = GExpr tanh

instance (Floating a, Floating (StandardHOReprF m r i a)) =>
         Interp__asinh (StandardHORepr m r i) a where
  interp__asinh = GExpr asinh

instance (Floating a, Floating (StandardHOReprF m r i a)) =>
         Interp__acosh (StandardHORepr m r i) a where
  interp__acosh = GExpr acosh

instance (Floating a, Floating (StandardHOReprF m r i a)) =>
         Interp__atanh (StandardHORepr m r i) a where
  interp__atanh = GExpr atanh

----------------------------------------------------------------------
-- Probability expression operations
----------------------------------------------------------------------

instance (Ord r, Floating r) => Interp__realToProb (StandardHORepr m r i) where
  interp__realToProb = GExpr (Log.Exp . log . toNonNeg) where
    toNonNeg :: (Ord r, Floating r) => r -> r
    toNonNeg r = if r < 0 then 0 else r

instance Interp__logRealToProb (StandardHORepr m r i) where
  interp__logRealToProb = GExpr Log.Exp

instance Floating r => Interp__probToReal (StandardHORepr m r i) where
  interp__probToReal = GExpr (exp . Log.ln)

instance Interp__probToLogReal (StandardHORepr m r i) where
  interp__probToLogReal = GExpr Log.ln

instance HasGamma r => Interp__gammaProb (StandardHORepr m r i) where
  interp__gammaProb = GExpr (Log.Exp . logGamma)

instance HasGamma r => Interp__digamma (StandardHORepr m r i) where
  interp__digamma = GExpr digamma

----------------------------------------------------------------------
-- Misc operations
----------------------------------------------------------------------

instance (Show a, Show (StandardHOReprF m r i a), i ~ Int) =>
         Interp__gtrace (StandardHORepr m r i) a b where
  interp__gtrace = GExpr gtrace

instance i ~ Int => Interp__gerror (StandardHORepr m r i) a where
  interp__gerror = GExpr gerror

----------------------------------------------------------------------
-- * Distributions Supported by All 'StandardHORepr' Representations
----------------------------------------------------------------------

instance Monad m => Interp__ctorDist__ListF (StandardHORepr m r i) where
  interp__ctorDist__Nil = GExpr $ \ mkNil dv ->
    do _ <-
         matchHOReprADTDistVar dv
         (\adt -> case adt of
             Nil -> mkNil (VADT Tuple0)
             Cons _ _ -> error "Unexpected Cons")
         (mkNil VParam)
       return Nil

  interp__ctorDist__Cons = GExpr $ \ mkCons dv ->
    do (Tuple2 hd tl) <-
         matchHOReprADTDistVar dv
         (\adt -> case adt of
             Nil -> error "Unexpected Nil"
             Cons hdv tlv -> mkCons (VADT (Tuple2 hdv tlv)))
         (mkCons (VADT (Tuple2 VParam VParam)))
       return (Cons hd tl)

-- If a repr can do a categorical, it can do an ADT distribution on lists
instance (Monad m, Num i, Eq i, Show i,
          Interp__categorical (StandardHORepr m r i)) =>
         Interp__adtDist__ListF (StandardHORepr m r i) where
  interp__adtDist__ListF =
    GExpr $ \ probNil mkNil probCons mkCons dvList ->
    let
      -- Build a categorical distribution for the constructor, where 0 -> Nil
      -- and 1 -> Cons
      ctor_dist :: DistVar (StandardHORepr m r i) Int ->
                   m (StandardHOReprF m r i Int)
      ctor_dist =
        unGExpr (interp__categorical
                 :: GExpr (StandardHORepr m r i) (GList Prob -> Dist Int)) $
        fromHaskellListF GExpr [GExpr probNil, GExpr probCons]
      -- Helper wrapper around mkNil
      mkNilH = mkNil (VADT Tuple0)
      -- Helper wrapper around mkCons, that takes vars for the head and tail
      mkConsH hdv tlv =
        do Tuple2 hd tl <- mkCons (VADT (Tuple2 hdv tlv))
           return (Cons hd tl) in

    matchHOReprADTDistVar dvList
    (\adt -> case adt of
        Nil ->
          void (ctor_dist $ VData $ GData 0) >> void mkNilH >> return Nil
        Cons hdv tlv ->
          void (ctor_dist $ VData $ GData 1) >> mkConsH hdv tlv)
    (do ctor_choice <- ctor_dist VParam
        case ctor_choice of
          0 -> void mkNilH >> return Nil
          1 -> mkConsH VParam VParam
          _ -> error ("ListF: Invalid constructor choice: "
                      ++ show ctor_choice))

--
-- * Vector Operations
--

instance (Monad m, i ~ Int, EmbedRepr (StandardHORepr m r Int) a) =>
         Interp__vec_iid (StandardHORepr m r i) a where
  interp__vec_iid = GExpr $ \len d dv ->
    matchHOReprVectorDistVar dv
    (\dvs ->
      if V.length dvs == len then V.mapM d dvs else
        error "vec_iid: incorrect vector length on input")
    (V.replicateM len (d VParam))

{-
-- | Build a distribution on 'Vector's from one on lists
instance (Monad m, EmbedRepr (StandardHORepr m r Int) a) =>
         Interp__vec_dist (StandardHORepr m r Int) a where
  interp__vec_dist = GExpr $ \d dv ->
    matchHOReprVectorDistVar dv
    (\dvs ->
      V.fromList <$> toHaskellListF unGExpr <$> unGExpr <$>
      d _)
    _
-}

instance Monad m => Interp__vec_nil_dist (StandardHORepr m r i) a where
  interp__vec_nil_dist = GExpr $ \dv ->
    matchHOReprVectorDistVar dv
    (\dvs ->
      if V.length dvs == 0 then return V.empty else
        error "vec_nil_dist: non-empty vector!")
    (return V.empty)

instance Monad m => Interp__vec_cons_dist (StandardHORepr m r i) a where
  interp__vec_cons_dist = GExpr $ \d dv ->
    do (Tuple2 (GExpr hd) (GExpr tl)) <-
         case dv of
           VParam -> d VParam
           VData (GData xs) ->
             if V.length xs == 0 then
               error "vec_cons_dist: empty vector!"
             else
               d $ VData $ GData $ ADT $ Tuple2 (Id $ V.head xs) (Id $ V.tail xs)
           VData GNoData -> d (VData GNoData)
           VExpr (GExpr es) ->
             if V.length es == 0 then
               error "vec_cons_dist: empty vector!"
             else
               d $ VExpr $ GExpr $ Tuple2 (GExpr $ V.head es) (GExpr $ V.tail es)
       return (V.cons hd tl)

instance Interp__vec_head (StandardHORepr m r i) a where
  interp__vec_head = GExpr V.head

instance Interp__vec_tail (StandardHORepr m r i) a where
  interp__vec_tail = GExpr V.tail

instance i ~ Int => Interp__vec_length (StandardHORepr m r i) a where
  interp__vec_length = GExpr V.length

--
-- * Matrix Operations
--

-- FIXME: matrix operations!


{-# LANGUAGE DataKinds              #-}
{-# LANGUAGE DeriveGeneric          #-}
{-# LANGUAGE FlexibleContexts       #-}
{-# LANGUAGE FlexibleInstances      #-}
{-# LANGUAGE FunctionalDependencies #-}
{-# LANGUAGE GADTs                  #-}
{-# LANGUAGE LambdaCase             #-}
{-# LANGUAGE MultiParamTypeClasses  #-}
{-# LANGUAGE OverloadedStrings      #-}
{-# LANGUAGE PolyKinds              #-}
{-# LANGUAGE RankNTypes             #-}
{-# LANGUAGE RecordWildCards        #-}
{-# LANGUAGE ScopedTypeVariables    #-}
{-# LANGUAGE TemplateHaskell        #-}
{-# LANGUAGE TupleSections          #-}
{-# LANGUAGE TypeApplications       #-}
{-# LANGUAGE TypeFamilies           #-}
{-# LANGUAGE TypeOperators          #-}
{-# LANGUAGE TypeSynonymInstances   #-}

module Egg.Vehicle (
    Vehicle(..), vName, vBaseCapacity, vCosts
  , VehicleData(..), HasVehicleData(..), SomeVehicleData
  , DepotStatus(..), _DepotStatus, SomeDepotStatus, _SomeDepotStatus
  , SomeVehicleUpgradeError(..), _SVUERegression, _SVUENoSlot
  , withSomeDepotStatus
  , initDepotStatus
  , initSomeDepotStatus
  , baseDepotCapacity
  , totalDepotCapacity
  , vehicleHistory
  , vehicleCostAt
  , vehiclePrice
  , upgradeVehicle
  , upgradeSomeVehicle
  , vehicleUpgrades
  , someVehicleUpgrades
  ) where

import           Control.Applicative
import           Control.Lens hiding         ((.=))
import           Control.Monad
import           Data.Aeson
import           Data.Dependent.Sum
import           Data.Finite
import           Data.Finite.Internal
import           Data.Finite.Util            ()
import           Data.Functor
import           Data.Maybe
import           Data.Singletons
import           Data.Singletons.Prelude.Num
import           Data.Singletons.TypeLits
import           Data.Type.Combinator
import           Data.Vector.Sized.Util
import           Egg.Commodity
import           Egg.Research
import           GHC.Generics                (Generic)
import           Numeric.Lens
import           Numeric.Natural
import           Statistics.LinearRegression
import           Text.Printf
import qualified Data.Map                    as M
import qualified Data.Text                   as T
import qualified Data.Vector                 as V
import qualified Data.Vector.Sized           as SV
import qualified GHC.TypeLits                as TL

data Vehicle = Vehicle
        { _vName         :: T.Text
        , _vBaseCapacity :: Natural         -- ^ eggs per minute
        , _vCosts        :: V.Vector Bock
        , _vRegression   :: Maybe (Double, Double)
        }
  deriving (Show, Eq, Ord, Generic)

makeLenses ''Vehicle

newtype VehicleData vs = VehicleData { _vdVehicles :: SV.Vector vs Vehicle }
  deriving (Show, Eq, Ord, Generic)

makeClassy ''VehicleData
makeWrapped ''VehicleData

type SomeVehicleData = DSum Sing VehicleData

data DepotStatus vs slots
    = DepotStatus { _dsSlots :: M.Map (Finite slots) (Finite vs) }
  deriving (Show, Eq, Ord, Generic)

makePrisms ''DepotStatus
makeWrapped ''DepotStatus

data SomeVehicleUpgradeError vs
        = SVUERegression { _svuePrevious :: Finite vs }
        | SVUENoSlot
  deriving (Show, Eq, Ord)

makePrisms ''SomeVehicleUpgradeError

data SomeDepotStatus vs
    = SomeDepotStatus { _sdsSlots :: M.Map Natural (Finite vs) }
  deriving (Show, Eq, Ord, Generic)

_SomeDepotStatus
    :: Functor f
    => Bonuses
    -> (forall slots. KnownNat slots => LensLike' f (DepotStatus vs slots) b)
    -> LensLike' f (SomeDepotStatus vs) b
_SomeDepotStatus bs f g = traverseSomeDepotStatus bs (f g)

withSomeDepotStatus
    :: Bonuses
    -> SomeDepotStatus vs
    -> (forall slots. KnownNat slots => DepotStatus vs slots -> r)
    -> r
withSomeDepotStatus bs sd f = sd ^. _SomeDepotStatus bs (to f)

traverseSomeDepotStatus
    :: Functor f
    => Bonuses
    -> (forall slots. KnownNat slots => DepotStatus vs slots -> f (DepotStatus vs slots))
    -> SomeDepotStatus vs
    -> f (SomeDepotStatus vs)
traverseSomeDepotStatus bs f sds0 = withSomeSing fleetSize $ \(SNat :: Sing slots) ->
    let ds0 = DepotStatus
            . M.mapKeysMonotonic (Finite . fromIntegral)
            . M.filterWithKey (\k _ -> k < fleetSize)
            $ _sdsSlots sds0
    in  SomeDepotStatus . M.mapKeysMonotonic (fromIntegral . getFinite) . _dsSlots
          <$> f @slots ds0
  where
    fleetSize :: Natural
    fleetSize = 4 ^. bonusingFor @Double bs BTFleetSize . to round

vehicleParseOptions :: Options
vehicleParseOptions = defaultOptions
    { fieldLabelModifier = camelTo2 '-' . drop 2
    }

instance FromJSON Vehicle where
    parseJSON  = fmap fillRegression . genericParseJSON  vehicleParseOptions
instance ToJSON Vehicle where
    toJSON     = genericToJSON     vehicleParseOptions
    toEncoding = genericToEncoding vehicleParseOptions

fillRegression :: Vehicle -> Vehicle
fillRegression v = case v ^. vRegression of
    Just _  -> v
    Nothing -> v & vRegression .~ reg
  where
    reg = guard (V.length c > 1) $>
                 linearRegression (V.imap (\i _ -> fromIntegral i) c)
                                  (log . view bock <$> c)
      where
        c = v ^. vCosts

instance FromJSON SomeVehicleData where
    parseJSON o = do
      vs <- parseJSON o
      SV.withSized vs $ \vsV ->
        return $ SNat :=> VehicleData vsV
instance ToJSON SomeVehicleData where
    toJSON = \case
        _ :=> r -> toJSON r
    toEncoding = \case
        _ :=> r -> toEncoding r
instance KnownNat vs => FromJSON (VehicleData vs) where
    parseJSON o = do
      vs <- parseJSON o
      case SV.toSized vs of
        Nothing  -> fail $ printf "Bad number of items in list. (Expected %d, got %d)"
                         (fromSing (SNat @vs)) (length vs)
        Just vsV -> return $ VehicleData vsV
instance ToJSON (VehicleData habs) where
    toJSON     = toJSON     . SV.fromSized . _vdVehicles
    toEncoding = toEncoding . SV.fromSized . _vdVehicles

instance (KnownNat vs, KnownNat slots) => FromJSON (DepotStatus vs slots) where
    parseJSON o = DepotStatus <$> parseJSON o
instance ToJSON (DepotStatus vs slots) where
    toJSON = toJSON . _dsSlots
    toEncoding = toEncoding . _dsSlots
instance KnownNat vs => FromJSON (SomeDepotStatus vs) where
    parseJSON o = SomeDepotStatus <$> parseJSON o
instance ToJSON (SomeDepotStatus vs) where
    toJSON = toJSON . _sdsSlots
    toEncoding = toEncoding . _sdsSlots

-- | Initial 'DepotStatus' to start off the game.
initDepotStatus :: (KnownNat vs, KnownNat slots, 1 TL.<= vs, 1 TL.<= slots) => DepotStatus vs slots
initDepotStatus = DepotStatus $ M.singleton 0 0

-- | Initial 'SomeDepotStatus' to start off the game.
initSomeDepotStatus :: (KnownNat vs, 1 TL.<= vs) => SomeDepotStatus vs
initSomeDepotStatus = SomeDepotStatus $ M.singleton 0 0

-- | Total base capacity of all slots, in eggs per second.
baseDepotCapacity
    :: forall vs slots. KnownNat vs
    => VehicleData vs
    -> DepotStatus vs slots
    -> Double
baseDepotCapacity VehicleData{..} =
    sumOf $ _DepotStatus
          . folded
          . to (SV.index _vdVehicles)
          . vBaseCapacity
          . to fromIntegral
          . dividing 60

-- | Total capacity of all vehicles, factoring in bonuses.
totalDepotCapacity
    :: forall vs slots. KnownNat vs
    => VehicleData vs
    -> Bonuses
    -> DepotStatus vs slots
    -> Double
totalDepotCapacity vd@VehicleData{..} bs =
    sumOf $ to (baseDepotCapacity vd)
          . bonusingFor bs BTVehicleCapacity
          . bonusingFor bs BTVehicleSpeed

-- | How many of each vehicle has been purchased so far.  If key is not found,
-- zero purchases is implied.
vehicleHistory
    :: KnownNat (slots TL.+ 1)
    => DepotStatus vs slots
    -> M.Map (Finite vs) (Finite (slots TL.+ 1))
vehicleHistory = M.fromListWith (+)
               . toListOf (_DepotStatus . folded . to (, 1))

-- | Get the BASE price of a given vehicle, if a purchase were to be made.
-- Does not check if purchase is legal (see 'upgradeVehicle').
--
-- Returns Nothing if slots are all full?
vehiclePrice
    :: forall vs slots. (KnownNat vs, KnownNat slots)
    => VehicleData vs
    -> DepotStatus vs slots
    -> Finite vs
    -> Maybe Bock
vehiclePrice vd ds v = withKnownNat (SNat @slots %+ SNat @1) $
                           fmap priceOf
                         . strengthen
                         . fromMaybe 0
                         . M.lookup v
                         . vehicleHistory
                         $ ds
  where
    priceOf :: Finite slots -> Bock
    priceOf i = vd ^. vdVehicles
                    . ixSV v
                    . to (vehicleCostAt (fromIntegral i))

vehicleCostAt :: Int -> Vehicle -> Bock
vehicleCostAt i v = fromMaybe 100 $
        (v ^? vCosts . ix i) <|> (r <$> v ^. vRegression)
  where
    r (α, β) = realToFrac . exp $ α + β * fromIntegral i

-- | Purchase a vehicle upgrade.  Returns cost and new depot status,
-- if purchase is valid.
--
-- Purchase is invalid if purchasing a vehicle in a slot where a greater
-- vehicle is already purchased.
upgradeVehicle
    :: (KnownNat vs, KnownNat slots)
    => VehicleData vs
    -> Bonuses
    -> Finite slots
    -> Finite vs
    -> DepotStatus vs slots
    -> Either (Finite vs) (Bock, DepotStatus vs slots)
upgradeVehicle vd bs slot v ds0 =
    getComp . flip (_DepotStatus . at slot) ds0 $ \s0 -> Comp $ do
      case s0 of
        Just h | h >= v -> Left h
        _               -> Right ()
      -- should always be valid of the previous condition is true
      let price = fromJust $ vehiclePrice vd ds0 v
      return (price ^. bonusingFor bs BTVehicleCosts, Just v)

-- | Purchase a vehicle upgrade.  Returns cost and new depot status,
-- if purchase is valid.
--
-- Purchase is invalid if purchasing a vehicle in a slot where a greater
-- vehicle is already purchased, and also if there is no such slot.
upgradeSomeVehicle
    :: forall vs. KnownNat vs
    => VehicleData vs
    -> Bonuses
    -> Integer
    -> Finite vs
    -> SomeDepotStatus vs
    -> Either (SomeVehicleUpgradeError vs) (Bock, SomeDepotStatus vs)
upgradeSomeVehicle vd bs slot v sds0 =
    getComp . traverseSomeDepotStatus bs (Comp . go) $ sds0
  where
    go  :: KnownNat slots
        => DepotStatus vs slots
        -> Either (SomeVehicleUpgradeError vs) (Bock, DepotStatus vs slots)
    go ds0 = case packFinite slot of
      Nothing    -> Left SVUENoSlot
      Just slot' -> either (Left . SVUERegression) Right $
                      upgradeVehicle vd bs slot' v ds0

-- | List all possible vehicle upgrades.
--
-- Is a Left if the vehicle upgrade would be a regression.
vehicleUpgrades
    :: forall vs slots. KnownNat slots
    => VehicleData vs
    -> Bonuses
    -> DepotStatus vs slots
    -> (SV.Vector slots :.: SV.Vector vs :.: Either (Finite vs)) Bock
vehicleUpgrades vd bs ds = Comp . Comp $ SV.generate go
  where
    slots1 = SNat @slots %+ SNat @1
    hist :: M.Map (Finite vs) (Finite (slots TL.+ 1))
    hist = withKnownNat slots1 $
             vehicleHistory ds
    go :: Finite slots -> SV.Vector vs (Either (Finite vs) Bock)
    go s = vd ^. vdVehicles
              & SV.imap
                  (\j h -> withKnownNat slots1 $ do
                      case currVeh of
                        Just v | v >= j -> Left v
                        _               -> Right ()
                      let n = M.findWithDefault 0 j hist
                      return $ h ^. to (vehicleCostAt (fromIntegral n))
                                  . bonusingFor bs BTVehicleCosts
                  )
      where
        currVeh = ds ^? _DepotStatus . ix s

-- | List all possible vehicle upgrades for a 'SomeDepotStatus'.
--
-- Is a Left if the vehicle upgrade would be a regression.
someVehicleUpgrades
    :: forall vs. ()
    => VehicleData vs
    -> Bonuses
    -> SomeDepotStatus vs
    -> (V.Vector :.: SV.Vector vs :.: Either (Finite vs)) Bock
someVehicleUpgrades vd bs = view $ _SomeDepotStatus bs go
                                 . _Unwrapped
                                 . _Unwrapped
  where
    go  :: KnownNat slots
        => Getter (DepotStatus vs slots)
                  (V.Vector (SV.Vector vs (Either (Finite vs) Bock)))
    go = to (vehicleUpgrades vd bs)
       . _Wrapped
       . _Wrapped
       . to SV.fromSized


{-# LANGUAGE GADTs #-}
module Format.FastTime
  ( FastTime()
  , fastTime
  , ignoreHigh
  ) where

import Data.Time.Clock
import Statistics.Distribution
import Statistics.Distribution.Normal
import Text.Printf

newtype FastTime where
  FastTime :: Double -> FastTime

instance Show FastTime where
  showsPrec _ = showFastTime

fastTime :: NominalDiffTime -> FastTime
fastTime = FastTime . realToFrac

showFastTime :: FastTime -> ShowS
showFastTime (FastTime p) = showString go
  where go | p <= (-1000) = replicate 4 '\8593'
           | p >=      0  = replicate 4 '-'
           | p <=   (-10) = printf "-%3.0f" $ abs p
           | otherwise    = printf "%4.1f" p

ignoreHigh :: FastTime -> String
ignoreHigh a = go $ calc a
  where go p | p <= (-10) = "   "
             | otherwise  = show a

calc :: FastTime -> Double
calc (FastTime t) = undefined


{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE DeriveGeneric #-}
{-# LANGUAGE RecordWildCards #-}

module Main where
import Debug.Trace (traceShow)
import Data.Either.Combinators (fromRight, maybeToRight)
import BDSCOD.Llhd (llhdAndNB,initLlhdState, intervalLlhd)
import BDSCOD.Types
import BDSCOD.Utility
import Control.Monad (zipWithM)
import Control.Monad.Except (ExceptT, runExceptT, throwError)
import Control.Monad.Reader (ReaderT, asks, liftIO, runReaderT)
import qualified Data.Aeson as Json
import qualified Data.ByteString.Builder as BBuilder
import qualified Data.ByteString.Lazy as L
import qualified Data.Csv as Csv
import Data.List (intercalate, intersperse,nub)
import Data.Maybe (fromMaybe)
import qualified Data.Vector.Unboxed as Unboxed
import qualified Epidemic.BDSCOD as SimBDSCOD
import Epidemic.Types.Events
  ( EpidemicEvent(..)
  , asNewickString
  , eventTime
  , maybeEpidemicTree
  , maybeReconstructedTree
  )
import Epidemic.Types.Parameter
import Epidemic.Types.Population (Person(..), Identifier(..))
import qualified Epidemic.Utility as SimUtil
import GHC.Generics
import Numeric.Minimisation.Powell (minimise)
import Numeric.LinearAlgebra.Data (linspace, toList)
import Numeric.LinearAlgebra.HMatrix
import System.Environment (getArgs)
import System.Random.MWC
import Data.Word (Word32)
import System.IO.Unsafe (unsafePerformIO)

-- | Values of this type are used to specify an analysis of the data.
data InferenceConfiguration =
  InferenceConfiguration
    {
    -- | The time in the simulation at which the inference is carried out.
    inferenceTime :: AbsoluteTime
    ,
    -- | A couple of files to write the Newick representation to and the node labels.
    reconstructedTreeOutputFiles :: (FilePath, FilePath)
    ,
    -- | Where to write the sequence of observed events.
    observationsOutputCsv :: FilePath
    ,
    -- | Where to write the likelihood evaluations.
    llhdOutputCsv :: FilePath
    ,
    -- | Where to write the point estimate of prevalence.
    pointEstimatesCsv :: FilePath
    }
  deriving (Show, Generic)

-- | These objects describe the evaluation mesh when looking at the likelihood
-- profiles. This is useful because it allows us to provide these details at
-- runtime rather than hardcoding them here.
data LlhdProfileMesh =
  LlhdProfileMesh
  { lpmMeshSize :: Int
  , lpmLambdaBounds :: (Rate,Rate)
  , lpmMuBounds :: (Rate,Rate)
  , lpmPsiBounds :: (Rate,Rate)
  , lpmRhoBounds :: (Probability,Probability)
  , lpmOmegaBounds :: (Rate,Rate)
  , lpmNuBounds :: (Probability,Probability)
  } deriving (Show, Generic)

-- | This object configures the whole evaluation of this program and is to be
-- read in from a suitable JSON file.
data AppConfiguration =
  AppConfiguration
    { simulatedEventsOutputCsv :: FilePath
    , simulationParameters :: Parameters
    , simulationDuration :: TimeDelta
    , simulationSizeBounds :: (Int,Int)
    , simulationSeed :: Word32
    , inferenceConfigurations :: [InferenceConfiguration]
    , acLlhdProfileMesh :: LlhdProfileMesh
    }
  deriving (Show, Generic)

instance Json.FromJSON AppConfiguration

instance Json.FromJSON InferenceConfiguration

instance Json.FromJSON LlhdProfileMesh

type Simulation x = ReaderT AppConfiguration (ExceptT String IO) x

-- | This type is used to indicate if parameters are the true ones used in the
-- simulation or estimates parameters.
data ParameterKind
  = SimulationParameters
  | EstimatedParameters
  deriving (Show, Eq)

-- | A BDSCOD simulation configuration based on the parameters in the
-- environment.
bdscodConfiguration = do
  simParams <- asks simulationParameters
  (TimeDelta simDurDouble) <- asks simulationDuration
  let bdscodConfig =
        SimBDSCOD.configuration (AbsoluteTime simDurDouble) (unpackParameters simParams)
  case bdscodConfig of
    Just x -> do liftIO $ putStrLn (show simParams)
                 return x
    Nothing -> throwError "Could not construct BDSCOD configuration"

-- | Simulate the transmission process part of the epidemic making sure that the
-- results are acceptable in terms of the number of observed events before
-- returning a filtration of the events, i.e., the data that was availble at
-- several points in time.
--
-- __NOTE__ the filteration must happen before these are processed into
-- observations since the observations do not accumulate chronologically due to
-- birth events which can occur in the past due to new observations in the
-- present.
partialSimulatedEpidemic seedWord bdscodConfig =
  do
    gen <- liftIO $ initialize (Unboxed.fromList [seedWord])
    simEvents <- liftIO $ SimUtil.simulation' bdscodConfig SimBDSCOD.allEvents gen
    (sizeLowerBound,sizeUpperBound) <- asks simulationSizeBounds
    if length simEvents > sizeLowerBound && length simEvents < sizeUpperBound
      then do infTimes <- map inferenceTime <$> asks inferenceConfigurations
              simEventsCsv <- asks simulatedEventsOutputCsv
              liftIO $ L.writeFile simEventsCsv (Csv.encode simEvents)
              return [filter (\e -> eventTime e <= infTime) simEvents | infTime <- infTimes]
      else do liftIO $ putStrLn "Repeating epidemic simulation..."
              partialSimulatedEpidemic (seedWord + 1) bdscodConfig

-- | Run the actual observation of the simulation and record the results before
-- returning the dataset of observations generated by this epidemic.
simulatedObservations :: InferenceConfiguration
                      -> [EpidemicEvent]
                      -> Simulation (InferenceConfiguration,[Observation])
simulatedObservations infConfig@InferenceConfiguration{..} simEvents = do
  let Just (newickBuilder,newickMetaData) =
        do eTree <- maybeEpidemicTree simEvents
           rTree <- maybeReconstructedTree eTree
           asNewickString (AbsoluteTime 0, Person (Identifier 1)) rTree
      maybeObs = eventsAsObservations <$> SimBDSCOD.observedEvents simEvents
      (reconNewickTxt,reconNewickCsv) = reconstructedTreeOutputFiles
  case maybeObs of
    (Just obs) ->
      do
        liftIO $ L.writeFile reconNewickTxt (BBuilder.toLazyByteString newickBuilder)
        liftIO $ L.writeFile reconNewickCsv (Csv.encode newickMetaData)
        liftIO $ L.writeFile observationsOutputCsv (Csv.encode obs)
        return (infConfig,obs)
    _ -> throwError "Failed to simulate observations."

-- | If there is a unique timed value return that. This is used to make it
-- easier to extract the values of timed parameters so that you don't need to
-- store all of them.
uniqueTimedValue :: Eq x => Timed x -> Maybe x
uniqueTimedValue (Timed txs) = case txs of
  [] -> Nothing
  txs' -> let xs = nub [snd tx | tx <- txs']
              isUnique = 1 == length xs
            in if isUnique then Just (head xs) else Nothing

-- | Evaluate the NB posterior approximation of the prevalence for a single
-- point in parameter space and the LLHD over a list of points and write all of
-- the results to CSV.
recordLlhdCrossSections :: InferenceConfiguration
                          -> [Observation]
                          -> (Parameters,ParameterKind,[Parameters])
                          -> Simulation ()
recordLlhdCrossSections InferenceConfiguration {..} obs (singleParams, paramKind, evalParams) =
  let comma = BBuilder.charUtf8 ','
      parametersUsed = show paramKind
      parametersUsed' = BBuilder.stringUtf8 parametersUsed
      llhdVals = safeLlhdFromInit obs <$> evalParams
      nBValAndParams =
        Csv.encode . (:[]) $
        ( parametersUsed
        , show $ length obs
        , fromRight Zero $ snd <$> llhdAndNB obs singleParams initLlhdState
        , show $ getLambda singleParams
        , show $ getMu singleParams
        , show $ getPsi singleParams
        , show <$> uniqueTimedValue $ getRhos singleParams
        , show $ getOmega singleParams
        , show <$> uniqueTimedValue $ getNus singleParams )
      doublesAsString =
        BBuilder.toLazyByteString .
        mconcat .
        intersperse comma . (parametersUsed' :) . map BBuilder.doubleDec
   in liftIO $
      do
        L.appendFile llhdOutputCsv (doublesAsString llhdVals)
        L.appendFile pointEstimatesCsv nBValAndParams

-- | Evaluate the LLHD function on cross-sections centred at the true values
evaluateLLHD :: InferenceConfiguration -> [Observation] -> Simulation ()
evaluateLLHD infConfig obs = do
  simParams <- asks simulationParameters -- get the actual parameters used to simulate the observations
  llhdProfMesh <- asks acLlhdProfileMesh -- this specifies the cross section to use.
  let evalParams = crossSectionParameters llhdProfMesh simParams
  recordLlhdCrossSections infConfig obs (simParams,SimulationParameters,evalParams)

-- | Estimate the parameters of the of the model and then evaluate the LLHD
-- profiles and prevalence and append this to the file. The first value of the
-- CSV output now describes which parameters where used to to evaluate these
-- things.
estimateLLHD :: InferenceConfiguration -> [Observation] -> Simulation ()
estimateLLHD infConfig obs = do
  simParams@(Parameters (_,deathRate,_,_,_,_)) <- asks simulationParameters
  llhdProfMesh <- asks acLlhdProfileMesh
  liftIO $ putStrLn $ "\t\tEstimating parameters using " ++ show (length obs) ++ " observations..."
  let schedTimes = scheduledTimes simParams
      mleParams = estimateParameters deathRate schedTimes obs -- get the MLE estimate of the parameters
      evalParams = crossSectionParameters llhdProfMesh mleParams
  recordLlhdCrossSections infConfig obs (mleParams,EstimatedParameters,evalParams)

-- | A safe version of the log-likelihood assuming that you want to start from a
-- simple initial condition.
--
-- __NOTE__ This conditions the process against extinction.
--
safeLlhdFromInit :: [Observation] -> Parameters -> LogLikelihood
safeLlhdFromInit obs p@(Parameters (λ, μ, ψ, _, ω, _)) =
    let
      -- log-likelihood
      llhd = fromRight (-1e6) $ fst <$> llhdAndNB obs p initLlhdState
      -- log of the probability of not going extinct.
      lpne = log (1 - (μ + ψ + ω) / λ)
    in llhd - lpne

-- | Estimate of the MLE. This uses a simplex method.
--
-- __NOTE__ the ugly case statement means that this should work both with and
-- without scheduled observations.
--
-- __NOTE__ we fix the death rate to the true value because
-- this is assumed to be known a priori.
--
estimateParameters :: Rate -> ([AbsoluteTime],[AbsoluteTime]) -> [Observation] -> Parameters
estimateParameters deathRate sched obs =
  let energyFunc v2p x =
        negate $ safeLlhdFromInit obs (v2p x)
      mini rI v2p = minimise (energyFunc v2p) rI
  in case sched of
       -- there are no scheduled observations
       ([],[]) ->
         let -- randInit = [-1.5,-2.2,-2.2] -- initial point to start
           randInit = log <$> [0.228,4.8e-2,2.6e-2] -- _GROSS_ initial point as true value
           vec2Param vec =
             let [lnR1, lnR2, lnR3] = vec
             in packParameters ( exp lnR1
                               , deathRate
                               , exp lnR2
                               , []
                               , exp lnR3
                               , [])
         in case mini randInit vec2Param of
              Right (est,_,grad) -> traceShow grad $ vec2Param est
              Left msg -> error msg
       -- there is at least one scheduled observation
       (rhoTs, nuTs) ->
         let randInit = [-1.5,-2.5,-1,-2.5,-1]
             vec2Param vec =
               let [lnR1, lnR2, logitP1, lnR3, logitP2] = vec
                   rhoVal = invLogit logitP1
                   nuVal = invLogit logitP2
                   timed v ts = zip ts (repeat v)
               in packParameters ( exp lnR1
                                 , deathRate
                                 , exp lnR2
                                 , timed rhoVal rhoTs
                                 , exp lnR3
                                 , timed nuVal nuTs)
         in case mini randInit vec2Param of
              Right (est,_,grad) -> traceShow grad $ vec2Param est
              Left msg -> error msg

-- | List of parameters required to plot the cross sections.
crossSectionParameters :: LlhdProfileMesh -> Parameters -> [Parameters]
crossSectionParameters LlhdProfileMesh{..} ps =
  let mesh = toList . linspace lpmMeshSize
      lambdaMesh = mesh lpmLambdaBounds
      muMesh = mesh lpmMuBounds
      psiMesh = mesh lpmPsiBounds
      rhoProbMesh = mesh lpmRhoBounds
      omegaMesh = mesh lpmOmegaBounds
      nuProbMesh = mesh lpmNuBounds
      (rhoTimes,nuTimes) = scheduledTimes ps
      rhoMesh = [Timed [(t,r) | t <- rhoTimes] | r <- rhoProbMesh]
      nuMesh = [Timed [(t,n) | t <- nuTimes] | n <- nuProbMesh]
      -- for each dimension, construct a list of parameter values with the
      -- values from the mesh
      apply f = map (f ps)
      [lPs,mPs,pPs,oPs] =
        zipWith apply [putLambda,putMu,putPsi,putOmega] [lambdaMesh,muMesh,psiMesh,omegaMesh]
      [rPs,nPs] = zipWith apply [putRhos,putNus] [rhoMesh,nuMesh]
  in concat $
     if (rhoTimes,nuTimes) /= ([],[])
     then [lPs,mPs,pPs,rPs,oPs,nPs]
     else [lPs,mPs,pPs,oPs]

-- | Definition of the complete simulation study.
simulationStudy :: Simulation ()
simulationStudy = do
  bdscodConfig <- bdscodConfiguration
  liftIO $ putStrLn "\tRunning epidemic simulation"
  simSeed <- asks simulationSeed
  pEpi <- partialSimulatedEpidemic simSeed bdscodConfig
  infConfigs <- asks inferenceConfigurations
  liftIO $ putStrLn "\tExtracting observations from full simulation"
  pObs <- zipWithM simulatedObservations infConfigs pEpi
  liftIO $ putStrLn "\tEvaluating LLHD on cross-sections about true parameters"
  mapM_ (uncurry evaluateLLHD) pObs
  liftIO $ putStrLn "\tEvaluating LLHD on cross-sections about estimated parameters"
  mapM_ (uncurry estimateLLHD) pObs

main' :: IO ()
main' = do
  let configFilePath = "./examples/simulation-study-time-series/ts-config.json"
  maybeConfig <- getConfiguration configFilePath
  case maybeConfig of
    Nothing ->
      putStrLn $ "Could not get configuration from file: " ++ configFilePath
    Just config -> do
      putStrLn $ "Succeeded in reading configuration from file: " ++ configFilePath
      result <- runExceptT (runReaderT simulationStudy config)
      case result of
        Left errMsg -> putStrLn errMsg
        Right _ -> return ()

main :: IO ()
main = do
  configFilePath <- head <$> getArgs
  maybeConfig <- getConfiguration configFilePath
  case maybeConfig of
    Nothing ->
      putStrLn $ "Could not get configuration from file: " ++ configFilePath
    Just config -> do
      putStrLn $ "Succeeded in reading configuration from file: " ++ configFilePath
      result <- runExceptT (runReaderT simulationStudy config)
      case result of
        Left errMsg -> putStrLn errMsg
        Right _ -> return ()

-- | Attempt to read a configuration object from the given filepath.
getConfiguration :: FilePath -> IO (Maybe AppConfiguration)
getConfiguration fp = Json.decode <$> L.readFile fp


{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE FlexibleContexts #-}

module Taiji.Pipeline.SC.ATACSeq.Functions.Feature.Peak
    ( mkFeatMat
    , mkPeakMat
    , findPeaks
    , mergePeaks
    , getPeakEnrichment
    , mkCellClusterBed
    , computePeakRAS
    ) where

import Control.Arrow (first)
import           Bio.Pipeline
import qualified Data.HashSet as S
import Bio.Data.Bed hiding (NarrowPeak)
import qualified Bio.Data.Bed as Bed
import Data.Conduit.Internal (zipSinks, zipSources)
import qualified Data.Text as T
import qualified Data.ByteString.Char8 as B
import qualified Data.Matrix as Mat
import Data.Singletons.Prelude (SingI)
import Shelly hiding (FilePath)
import qualified Data.Vector as V
import qualified Data.Vector.Unboxed as U
import Control.DeepSeq (force)
import Data.Conduit.Zlib (gzip)

import Statistics.Distribution (quantile)
import Statistics.Distribution.ChiSquared (chiSquared)

import Taiji.Prelude
import Taiji.Pipeline.SC.ATACSeq.Types
import Taiji.Pipeline.SC.ATACSeq.Functions.Utils
import Taiji.Pipeline.SC.ATACSeq.Functions.QC (streamTN5Insertion)
import qualified Taiji.Utils.DataFrame as DF
import Taiji.Utils

-- | Make the read count matrix.
mkPeakMat :: SCATACSeqConfig config
          => FilePath
          -> SCATACSeq S (TagsAligned, File '[Gzip] 'NarrowPeak)
          -> ReaderT config IO (SCATACSeq S (File '[Gzip] 'Matrix))
mkPeakMat dir input = do
    passedQC <- getQCFunction
    let output = printf "%s/%s_rep%d_peak.mat.gz" dir (T.unpack $ input^.eid)
            (input^.replicates._1)
    input & replicates.traverse.files %%~ liftIO . (\(tagFl, regionFl) -> do
        regions <- runResourceT $ runConduit $
            streamBedGzip (regionFl^.location) .| sinkList :: IO [BED3]
        runResourceT $ runConduit $ streamTN5Insertion passedQC tagFl .|
            mapC (\xs -> (fromJust $ head xs^.name, xs)) .| mkCountMat regions .|
            sinkRows' (length regions) (fromJust . packDecimal) output
        return $ emptyFile & location .~ output )

-- | Make the read count matrix.
mkFeatMat :: SCATACSeqConfig config
          => FilePath
          -> SCATACSeq S (TagsAligned, File '[Gzip] 'NarrowPeak)
          -> ReaderT config IO ( SCATACSeq S
              ( File '[RowName, Gzip] 'Tsv
              , File '[ColumnName, Gzip] 'Tsv
              , File '[Gzip] 'Matrix ))
mkFeatMat dir input = do
    passedQC <- getQCFunction
    let output = printf "%s/%s_rep%d_peak.mat.gz" dir (T.unpack $ input^.eid)
            (input^.replicates._1)
        rownames = printf "%s/%s_rep%d_rownames.txt.gz" dir (T.unpack $ input^.eid)
            (input^.replicates._1)
        features = printf "%s/%s_rep%d_features.txt.gz" dir (T.unpack $ input^.eid)
            (input^.replicates._1)
        bcPrefix = B.pack $ T.unpack (input^.eid) <> "_" <> show (input^.replicates._1) <> "+"
    input & replicates.traverse.files %%~ liftIO . (\(tagFl, regionFl) -> do
        regions <- runResourceT $ runConduit $
            streamBedGzip (regionFl^.location) .| sinkList :: IO [BED3]
        let rowSink = mapC f .| unlinesAsciiC .| gzip .| sinkFile rownames
              where
                f (nm, xs) = force $ nm <> "\t" <> fromJust (packDecimal $ foldl1' (+) $ map snd xs)
            colSink = do
                vec <- colSumC $ length regions
                let bs = B.unlines $ zipWith showBed regions $ U.toList vec
                yield bs .| gzip .| sinkFile features
            showBed (BED3 chr s e) x = B.concat
                [ chr, ":", fromJust $ packDecimal s, "-"
                , fromJust $ packDecimal e, "\t", fromJust $ packDecimal x]
            sink = (,,) <$> ZipSink (sinkRows' (length regions) (fromJust . packDecimal) output)
                <*> ZipSink rowSink <*> ZipSink colSink
        _ <- runResourceT $ runConduit $ streamTN5Insertion passedQC tagFl .|
            mapC (\xs -> (fromJust $ head xs^.name, xs)) .|
            mkCountMat regions .| mapC (first (bcPrefix <>)) .|
            getZipSink sink
        return ( location .~ rownames $ emptyFile
               , location .~ features $ emptyFile
               , location .~ output $ emptyFile )
        )

-- | Call Peaks for aggregated files
findPeaks :: (SingI tags, SCATACSeqConfig config)
          => FilePath
          -> CallPeakOpts
          -> (B.ByteString, File tags 'Bed)
          -> ReaderT config IO (B.ByteString, File '[Gzip] 'NarrowPeak)
findPeaks prefix opts (cName, bedFl) = do
    tmpdir <- asks _scatacseq_tmp_dir
    dir <- asks _scatacseq_output_dir >>= getPath . (<> asDir prefix)
    let output = dir ++ "/" ++ B.unpack cName ++ ".narrowPeak.gz"
    asks _scatacseq_blacklist >>= \case
        Nothing -> liftIO $ do
            r <- callPeaks output bedFl Nothing opts
            return (cName, r)
        Just blacklist -> liftIO $ withTemp tmpdir $ \tmp -> do
            _ <- callPeaks tmp bedFl Nothing opts
            blackRegions <- readBed blacklist :: IO [BED3]
            let bedTree = bedToTree const $ map (\x -> (x, ())) blackRegions
            runResourceT $ runConduit $
                (streamBedGzip tmp :: ConduitT () Bed.NarrowPeak (ResourceT IO) ()) .|
                filterC (not . isIntersected bedTree) .| sinkFileBedGzip output
            return (cName, location .~ output $ emptyFile)

-- | Merge peaks
mergePeaks :: SCATACSeqConfig config
           => FilePath
           -> [(B.ByteString, File '[Gzip] 'NarrowPeak)]
           -> ReaderT config IO (Maybe (File '[Gzip] 'NarrowPeak))
mergePeaks _ [] = return Nothing
mergePeaks dir input = do
    tmpdir <- asks _scatacseq_tmp_dir
    let output = dir <> "/merged.narrowPeak.gz"
    liftIO $ withTemp tmpdir $ \tmp1 -> withTemp tmpdir $ \tmp2 -> do
        runResourceT $ runConduit $ mapM_ (streamBedGzip . (^._2.location)) input .|
            mapC resize .| sinkFileBed tmp1
        shelly $ escaping False $ bashPipeFail bash_ "cat" $
            [T.pack tmp1, "|", "sort", "-k1,1", "-k2,2n", "-k3,3n", ">", T.pack tmp2]
        let source = streamBed tmp2 .|
                mergeSortedBedWith iterativeMerge .| concatC
        runResourceT $ runConduit $ zipSources (iterateC succ (0 :: Int)) source .|
            mapC (\(i, p) -> name .~ Just ("p" <> B.pack (show i)) $ p) .|
            sinkFileBedGzip output
    return $ Just $ location .~ output $ emptyFile
  where
    iterativeMerge [] = []
    iterativeMerge peaks = bestPeak : iterativeMerge rest
      where
        rest = filter (\x -> sizeOverlapped x bestPeak == 0) peaks
        bestPeak = maximumBy (comparing (fromJust . (^.npPvalue))) peaks
    resize pk = chromStart .~ max 0 (summit - halfWindowSize) $
        chromEnd .~ summit + halfWindowSize$
        npPeak .~ Just halfWindowSize $ pk
      where
        summit = pk^.chromStart + fromJust (pk^.npPeak)
    halfWindowSize = 200

-- | Extract BEDs for each cluster.
mkCellClusterBed :: SCATACSeqConfig config
                 => SCATACSeq S ( File '[NameSorted, Gzip] 'Bed
                                , [CellCluster] )  -- ^ clusters
                 -> ReaderT config IO
                    (SCATACSeq S [(B.ByteString, File '[Gzip] 'Bed, Int)])
mkCellClusterBed input = do
    let idRep = asDir $ "/Bed/" <> T.unpack (input^.eid) <>
            "_rep" <> show (input^.replicates._1)
    dir <- asks _scatacseq_output_dir >>= getPath . (<> idRep)
    input & replicates.traverse.files %%~ liftIO . ( \(bed, cs) -> do
        let sinks = sequenceConduits $ flip map cs $ \CellCluster{..} -> do
                let output = dir ++ "/" ++ B.unpack _cluster_name ++ ".bed.gz"
                    cells = S.fromList $ map _cell_barcode _cluster_member
                    fl = location .~ output $ emptyFile
                (_, depth) <- filterC (f cells) .|
                    zipSinks (sinkFileBedGzip output) lengthC
                return (_cluster_name, fl, depth)
        runResourceT $ runConduit $
            streamBedGzip (bed^.location) .| sinks )
  where
    f :: S.HashSet B.ByteString -> BED -> Bool
    f cells x = fromJust (x^.name) `S.member` cells

-- | Get signal enrichment for each peak
getPeakEnrichment :: SCATACSeqConfig config
                  => ( [(B.ByteString, File '[Gzip] 'NarrowPeak)]
                     , Maybe (File '[Gzip] 'NarrowPeak) )
                  -> ReaderT config IO (File '[] 'Tsv, File '[] 'Tsv)
getPeakEnrichment (peaks, Just refPeak) = do
    dir <- asks _scatacseq_output_dir >>= getPath . (<> "/Feature/Peak/")
    let output1 = dir <> "peak_signal.tsv"
        output2 = dir <> "peak_pvalue.tsv"
    liftIO $ do
        list <- runResourceT $ runConduit $
            streamBedGzip (refPeak^.location) .| sinkList
        let col = map toString list
        (row, val) <- fmap unzip $ forM peaks $ \(nm, p) -> do
            peak <- runResourceT $ runConduit $ streamBedGzip (p^.location) .| sinkList
            return (T.pack $ B.unpack nm, getValue peak list)
        let (signal, pval) = DF.unzip $ DF.transpose $ DF.mkDataFrame row col val
        DF.writeTable output1 (T.pack . show) signal
        DF.writeTable output2 (T.pack . show) pval
        return (location .~ output1 $ emptyFile, location .~ output2 $ emptyFile)
  where
    toString x = T.pack (B.unpack $ x^.chrom) <> ":" <>
        T.pack (show $ x^.chromStart) <> "-" <> T.pack (show $ x^.chromEnd)
    getValue :: [Bed.NarrowPeak]  -- ^ query
            -> [BED3]  -- ^ Peak list
            -> [(Double, Double)]
    getValue query peakList = runIdentity $ runConduit $
        yieldMany peakList .| intersectBedWith f query .| sinkList
      where
        f _ [] = (0, 0)
        f _ [x] = (x^.npSignal, quantile (chiSquared 1) $ 10**(negate $ fromJust $ x^.npPvalue))
        f _ _ = undefined
getPeakEnrichment _ = undefined

-- | Compute the relative accessibility score.
computePeakRAS :: SCATACSeqConfig config
               => FilePath
               -> ( Maybe (File '[Gzip] 'NarrowPeak)
                  , [SCATACSeq S (File '[Gzip] 'Matrix)] )
               -> ReaderT config IO (Maybe (FilePath, FilePath, FilePath))
computePeakRAS _ (Nothing, _) = return Nothing
computePeakRAS prefix (peakFl, inputs) = do
    dir <- asks ((<> asDir prefix) . _scatacseq_output_dir) >>= getPath
    let output1 = dir <> "relative_accessibility_scores.tsv"
        output2 = dir <> "cell_specificity_score.tsv"
        output3 = dir <> "/cell_specificity_pvalue.tsv"
    liftIO $ do
        peaks <- fmap (map mkName) $ runResourceT $ runConduit $
            streamBedGzip (fromJust peakFl^.location) .| sinkList

        (names, cols) <- fmap unzip $ forM inputs $ \input -> do
            vec <- computeRAS (input^.replicates._2.files.location)
            return (input^.eid, V.convert vec)

        let ras = DF.map (*2.5) $ DF.fromMatrix peaks names $ Mat.fromColumns cols
            ss = computeSS $ DF.map (logBase 2 . (+1)) ras
        DF.writeTable output1 (T.pack . show) ras
        DF.writeTable output2 (T.pack . show) ss
        cdf <- computeCDF $ DF.map (logBase 2 . (+1)) ras
        DF.writeTable output3 (T.pack . show) $ DF.map (lookupP cdf) ss
        return $ Just (output1, output2, output3)
  where
    mkName :: BED3 -> T.Text
    mkName p = T.pack $ B.unpack (p^.chrom) <> ":" <> show (p^.chromStart) <>
        "-" <> show (p^.chromEnd)
    lookupP (vec, res, n) x | p == 0 = 1 / n
                            | otherwise = p
      where
        p = vec V.! i
        i = min (V.length vec - 1) $ truncate $ x / res


{-# LANGUAGE AllowAmbiguousTypes       #-}
{-# LANGUAGE DataKinds                 #-}
{-# LANGUAGE FlexibleContexts          #-}
{-# LANGUAGE GADTs                     #-}
{-# LANGUAGE NoMonomorphismRestriction #-}
{-# LANGUAGE OverloadedStrings         #-}
{-# LANGUAGE ScopedTypeVariables       #-}
{-# LANGUAGE TypeApplications          #-}
{-# LANGUAGE TypeOperators             #-}
{-# LANGUAGE QuasiQuotes               #-}
{-# LANGUAGE TypeFamilies              #-}
module Main where

import qualified Control.Foldl                   as FL
import           Control.Monad.IO.Class           (MonadIO (..))
import qualified Data.List                       as List
--import           Data.Maybe                       (fromMaybe)
import qualified Data.Map                        as M
import qualified Data.Profunctor                 as P
import qualified Data.Text                       as T
import qualified Data.Text.IO                    as T
import qualified Data.Text.Lazy                  as TL
import qualified Data.Vector.Storable            as V
import qualified Data.Vector                     as VB
import qualified Data.Vinyl                      as V
import qualified Frames                          as F
import qualified Frames.InCore                   as FI

import qualified Numeric.LinearAlgebra           as LA
import           Numeric.LinearAlgebra            (R
                                                  , Matrix
                                                  )

import qualified Text.Blaze.Html.Renderer.Text   as BH

import qualified Knit.Report                   as K

import qualified Frames.MapReduce as MR
import qualified Frames.Folds                    as FF
import qualified Frames.Transform                as FT
import           Frames.Table                    (blazeTable
                                                 , RecordColonnade
                                                 )

import           Data.String.Here

templateVars = M.fromList
  [
    ("lang", "English")
  , ("author", "Adam Conner-Sax")
  , ("pagetitle", "Map Reduce Examples")
--  , ("tufte","True")
  ]

-- this is annoying.  Where should this instance live?
type instance FI.VectorFor (Maybe a) = VB.Vector

mapReduceNotesMD
  = [here|
## Map Reduce Examples
|]

main :: IO ()
main = asPandoc

asPandoc :: IO ()
asPandoc = do
  let pandocWriterConfig = K.PandocWriterConfig (Just "pandoc-templates/minWithVega-pandoc.html")  templateVars K.mindocOptionsF
  htmlAsTextE <- K.knitHtml (Just "MapReduce.Main") K.logAll pandocWriterConfig $ do
    let rows :: Int = 20000
        vars = unweighted rows
        yNoise = 1.0
    K.logLE K.Info "Creating data"
    frame <- liftIO (noisyData vars yNoise coeffs >>= makeFrame vars)
    K.addMarkDown mapReduceNotesMD
    testMapReduce frame (showText rows <> " rows, average of X and Y by label.") mrAvgXYByLabel
    testMapReduce frame (showText rows <> " rows, filtered label, max of X and Y by label.") mrMaxXYByLabelABC
  case htmlAsTextE of
    Right htmlAsText -> T.writeFile "examples/html/mapReduce.html" $ TL.toStrict  $ htmlAsText
    Left err -> putStrLn $ "pandoc error: " ++ show err

type Label = "label" F.:-> T.Text
type Y = "y" F.:-> Double
type X = "x" F.:-> Double
type ZM = "zMaybe" F.:-> (Maybe Double)
type Weight = "weight" F.:-> Double
type IsDup = "is_duplicate" F.:-> Bool
type AllCols = [Label,Y,X,Weight]

-- let's make some map-reductions on Frames of AllCols

-- First some unpackings
noUnpack = MR.noUnpack
filterLabel ls = MR.filterUnpack (\r -> (F.rgetField @Label r) `List.elem` ls)
filterMinX minX = MR.filterUnpack ((>= minX) . F.rgetField @X)
editLabel f r = F.rputField @Label (f (F.rgetField @Label r)) r -- this would be better with lenses!!
unpackDup = MR.Unpack $ \r -> [r, editLabel (<> "2") r]

-- some assignings
assignToLabels = MR.assignKeys @'[Label]
assignDups = MR.assign @(F.Record '[IsDup]) (\r -> (T.length (F.rgetField @Label r) > 1) F.&: V.RNil) (F.rcast @[Y,X,Weight])

-- some reductions
--averageF :: FL.Fold (F.FrameRec '[X,Y]) F.Record '[X,Y]
averageF = FF.foldAllConstrained @RealFloat FL.mean

--maxX :: FL.Fold (F.Record '[X]) (F.Record '[ZM])
maxX = P.dimap (F.rgetField @X) (FT.recordSingleton @ZM) FL.maximum

--maxXY :: FL.Fold (F.Record '[X,Y]) (F.Record '[MX])
maxXY = P.dimap (\r -> Prelude.max (F.rgetField @X r) (F.rgetField @Y r)) (FT.recordSingleton @ZM) FL.maximum

-- put them together
--mrAvgXYByLabel :: FL.Fold (F.Record AllCols) (F.FrameRec AllCols)
mrAvgXYByLabel = MR.concatFold $ MR.mapReduceFold noUnpack (MR.splitOnKeys @'[Label]) (MR.foldAndAddKey averageF)

--mrAvgXYByLabelP :: FL.Fold (F.Record AllCols) (F.FrameRec AllCols)
--mrAvgXYByLabelP = MR.parBasicListHashableFold 1000 6 noUnpack (MR.splitOnKeys @'[Label]) (MR.foldAndAddKey averageF)
--  MR.MR.mapReduceGF (MRP.defaultParReduceGatherer pure)

mrMaxXYByLabelABC :: FL.Fold (F.Record AllCols) (F.FrameRec '[Label,ZM])
mrMaxXYByLabelABC = MR.concatFold $ MR.mapReduceFold (filterLabel ["A","B","C"]) assignToLabels (MR.foldAndAddKey maxXY)

noisyData :: [Double] -> Double -> LA.Vector R -> IO (LA.Vector R, LA.Matrix R)
noisyData variances noiseObs coeffs = do
  -- generate random measurements
  let d = LA.size coeffs
      nObs = List.length variances
  let xs0 :: Matrix R = LA.asColumn $ LA.fromList $ [-0.5 + (realToFrac i/realToFrac nObs) | i <- [0..(nObs-1)]]
  let xsC = 1 LA.||| xs0
      ys0 = xsC LA.<> LA.asColumn coeffs
  yNoise <- fmap (List.head . LA.toColumns) (LA.randn nObs 1) -- 0 centered normally (sigma=1) distributed random numbers
  let ys = ys0 + LA.asColumn (LA.scale noiseObs (V.zipWith (*) yNoise (LA.cmap sqrt (LA.fromList variances))))
  return (List.head (LA.toColumns ys), xsC)

makeFrame :: [Double] -> (LA.Vector R, LA.Matrix R) -> IO (F.FrameRec AllCols)
makeFrame vars (ys, xs) = do
  if snd (LA.size xs) /= 2 then error ("Matrix of xs wrong size (" ++ show (LA.size xs) ++ ") for makeFrame") else return ()
  let rows = LA.size ys
      rowIndex = [0..(rows - 1)]
      makeLabel :: Int -> Char
      makeLabel n = toEnum (fromEnum 'A' + n `mod` 26)
      makeRecord :: Int -> F.Record AllCols
      makeRecord n = T.pack [makeLabel n] F.&: ys `LA.atIndex` n F.&: xs `LA.atIndex` (n,1) F.&: vars List.!! n F.&: V.RNil -- skip bias column
  return $ F.toFrame $ makeRecord <$> rowIndex

unweighted :: Int -> [Double]
unweighted n = List.replicate n (1.0)

coeffs :: LA.Vector R = LA.fromList [1.0, 2.2]

testMapReduce :: ( RecordColonnade as
                 , K.Member K.ToPandoc effs
                 , K.PandocEffects effs
                 , MonadIO (K.Sem effs)
                 , Show (F.Record as))
              => F.FrameRec AllCols
              -> T.Text
              -> FL.Fold (F.Record AllCols) (F.FrameRec as)
              -> K.Sem effs ()
testMapReduce dataFrame title mrFold = do
  K.logLE K.Info $ "Doing map-reduce fold: " <> title
  let resFrame = FL.fold mrFold dataFrame
      header _ _ = title
  K.addMarkDown $ "\n## " <> title
  K.addBlaze $ blazeTable resFrame

showText :: Show a => a -> T.Text
showText = T.pack . show


{-# LANGUAGE AllowAmbiguousTypes #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE RecordWildCards #-}
{-# OPTIONS_GHC -Wincomplete-patterns #-}

--{-# OPTIONS_GHC -W #-}

module Symphony.Symphony where

import AbsHashedLang
import Codec.Picture
import Control.Monad (when)
import Control.Monad.Except
import qualified Data.Array as Array
import Data.Char (toLower)
import Data.Complex
import Data.List (intercalate)
import Data.List.Extra (firstJust)
import Data.Map (Map)
import qualified Data.Map.Strict as Map
import Data.Maybe (fromMaybe, mapMaybe)
import qualified Data.Set as Set
import Data.Tuple.HT (fst3)
import ErrM
import HashedExpression.Codegen
import qualified HashedExpression.Codegen.CSimple as CSimple
import HashedExpression.Internal
import HashedExpression.Internal.Expression (ExpressionMap, NodeID, Op (..))
import qualified HashedExpression.Internal.Expression as HE
import qualified HashedExpression.Internal.Node as HN
import HashedExpression.Internal.Utils
import qualified HashedExpression.Operation as HO
import HashedExpression.Prettify
import qualified HashedExpression.Problem as HS
import qualified HashedExpression.Value as HV
import LayoutHashedLang
import LexHashedLang
import ParHashedLang
import Symphony.Common
import Symphony.Exp
import Symphony.Solver
import Text.Regex.Posix

myLLexer :: String -> [Token]
myLLexer = resolveLayout True . myLexer

-------------------------------------------------------------------------------
-- Parse content of symphony file
-- Thanks to BNFC, we only need a simple call to `pProblem`
parse :: String -> Result Problem
parse fileContent =
  case pProblem . myLLexer $ fileContent of
    Ok problem -> return problem
    Bad errStr ->
      case (errStr =~ "line ([0-9]+), column ([0-9]+)" :: (String, String, String, [String])) of
        (_, _, _, [rs, cs]) ->
          -- because layout parsing add dummy ';'
          let r = min (read rs) (length r2c)
              c = min (read cs) (getNumColumn r)
           in throwError $ SyntaxError (r, c)
        _ -> throwError $ SyntaxError (1, 1)
  where
    r2c = Map.fromList . zip [1 ..] . map length . lines $ fileContent
    getNumColumn r = fromMaybe 0 $ Map.lookup r r2c

data ValidSymphony = ValidSymphony (ExpressionMap, NodeID) Vars Consts [HS.ConstraintStatement] AvailableSolver

-------------------------------------------------------------------------------
-- 1. Check if there is variables block
-- 2. Check if variable block is valid (no name clash)
-- 3. If there is a constant, check if it is valid (no name clash)
-- 4. Check if all expressions (in constraints and objective) is valid (should be scalar, all operation should be valid),
--        if yes construct Expression.
-- 5. Gen code
checkSemantic :: Problem -> Result ValidSymphony
checkSemantic problem = do
  let toVar name (shape, _) = varWithShape shape name
  let toParam name (shape, _) = paramWithShape shape name
  varDecls <- getVarDecls problem -- 1
  vars <- foldM checkVariableDecl Map.empty varDecls -- 2
  constDecls <- getConstDecls problem
  consts <- foldM (checkConstantDecl vars) Map.empty constDecls
  let initContext =
        Context
          { declarations = Map.mapWithKey toVar vars `Map.union` Map.mapWithKey toParam consts,
            vars = vars,
            consts = consts
          }
  letDecls <- getLetDecls problem
  finalContext <- foldM checkLetDecl initContext letDecls
  -- Complete processing context (variables, constants, declared values)
  -- Constraint
  constraintDecls <- getConstraintDecls problem
  css <- foldM (checkConstraintDecl finalContext) [] constraintDecls
  -- Objective
  parseObjectiveExp <- getMinimizeBlock problem
  objectiveExp <- constructExp finalContext (Just []) parseObjectiveExp
  -- Solver
  let shape = getShape objectiveExp
      nt = getNT objectiveExp
  solver <- getSolver problem
  when (shape /= [] || nt /= HE.R) $
    throwError $
      ErrorWithPosition
        ( "Objective should be a real scalar, here it is ("
            ++ toReadable shape
            ++ ", "
            ++ show nt
            ++ ")"
        )
        (getBeginningPosition parseObjectiveExp)
  liftIO $ putStrLn "Syntax & semantic is correct"
  return $ ValidSymphony objectiveExp vars consts css solver

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

-- | Generate code
generateCode :: String -> ValidSymphony -> Result ()
generateCode outputPath (ValidSymphony objectiveExp vars consts css solver) = do
  let problemGen =
        HS.constructProblem
          -- Objective
          (wrap objectiveExp)
          -- Constraints
          (HS.Constraint css)
  heProblem <- liftEitherString problemGen
--  case problemGen of
--    HS.ProblemValid heProblem -> do
  let valMap = Map.mapMaybeWithKey varVal vars `Map.union` Map.mapMaybeWithKey constVal consts
  res <- liftEitherString $ generateProblemCode CSimple.CSimpleConfig {output = CSimple.OutputHDF5} heProblem valMap
--  case generateProblemCode CSimple.CSimpleConfig {output = CSimple.OutputHDF5} heProblem valMap of
--    Invalid reason -> throwError $ GeneralError reason
--    Success res -> do
  liftIO $ putStrLn "Problem detail:"
  liftIO $ print $ heProblem
  liftIO $ res outputPath
  liftIO $ putStrLn "Download solver & adapter......"
  liftIO $ downloadSolver outputPath solver
--    HS.ProblemInvalid reason -> throwError $ GeneralError reason
  where
    varVal _ (_, Just val) = Just val
    varVal _ _ = Nothing
    constVal _ (_, val) = Just val

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

-- | Get variable declarations, there must be exact 1 variables block
getVarDecls :: Problem -> Result [VariableDecl]
getVarDecls (Problem blocks) =
  case filter isVariableBlock blocks of
    [] -> throwError $ GeneralError "No variale block"
    [BlockVariable declss] -> return $ concat declss
    _ -> throwError $ GeneralError "There are more than 1 variables block"
  where
    isVariableBlock (BlockVariable declss) = True
    isVariableBlock _ = False

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

-- | Get constants declarations, there must be at most 1 constants block
getConstDecls :: Problem -> Result [ConstantDecl]
getConstDecls (Problem blocks) =
  case filter isConstantBlock blocks of
    [] -> return []
    [BlockConstant declss] -> return $ concat declss
    _ -> throwError $ GeneralError "There are more than 1 constant block"
  where
    isConstantBlock (BlockConstant declss) = True
    isConstantBlock _ = False

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

-- | Get constraints declarations, there must be at most 1 constraints block
getConstraintDecls :: Problem -> Result [ConstraintDecl]
getConstraintDecls (Problem blocks) =
  case filter isConstraintBlock blocks of
    [] -> return []
    [BlockConstraint declss] -> return $ concat declss
    _ -> throwError $ GeneralError "There are more than 1 constraint block"
  where
    isConstraintBlock (BlockConstraint declss) = True
    isConstraintBlock _ = False

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

-- | Get immediate declarations, at most 1 block
getLetDecls :: Problem -> Result [LetDecl]
getLetDecls (Problem blocks) =
  case filter isLetBlock blocks of
    [] -> return []
    [BlockLet declss] -> return $ concat declss
    _ -> throwError $ GeneralError "There are more than 1 let block"
  where
    isLetBlock (BlockLet declss) = True
    isLetBlock _ = False

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

-- | Get objective block, must be 1 block
getMinimizeBlock :: Problem -> Result Exp
getMinimizeBlock (Problem blocks) =
  case filter isMinimizeBlock blocks of
    [] -> throwError $ GeneralError "No minimize block"
    [BlockMinimize exp] -> return exp
    _ -> throwError $ GeneralError "There are more than 1 minimize block"
  where
    isMinimizeBlock (BlockMinimize declss) = True
    isMinimizeBlock _ = False

-- | Get the solver
getSolver :: Problem -> Result AvailableSolver
getSolver (Problem blocks) =
  case filter isSolverBlock blocks of
    [] -> do
      lift $ putStrLn "No solver specified, default to L-BFGS-B (https://github.com/stephenbeckr/L-BFGS-B-C)"
      return LBFGSB
    [BlockSolver (SolverName name)]
      | map toLower name `elem` ["lbfgs-b", "lbfgsb", "l-bfgs-b"] -> return LBFGSB
      | map toLower name `elem` ["ipopt"] -> return Ipopt
      | otherwise -> throwError $ GeneralError "Unknown solver, should be (ipopt | lbfgs-b)"
    _ -> throwError $ GeneralError "More than one solver specified"
  where
    isSolverBlock (BlockSolver name) = True
    isSolverBlock _ = False

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

-- | Check variable declarations
checkVariableDecl :: Vars -> VariableDecl -> Result Vars
checkVariableDecl accRes decl = do
  let checkName name pos =
        when (name `Map.member` accRes) $
          throwError $
            ErrorWithPosition ("Duplicate declaration of " ++ name) pos
  case decl of
    VariableNoInit (PIdent (pos, name)) shape -> do
      checkName name pos
      return $ Map.insert name (toHEShape shape, Nothing) accRes
    VariableWithInit (PIdent (pos, name)) shape val -> do
      checkName name pos
      checkVal (toHEShape shape) val
      let heShape = toHEShape shape
      heVal <- toHEVal heShape val
      return $ Map.insert name (heShape, Just heVal) accRes

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

-- | Check constant declarations
checkConstantDecl :: Vars -> Consts -> ConstantDecl -> Result Consts
checkConstantDecl vars accRes decl = do
  let (ConstantDecl (PIdent (pos, name)) shape val) = decl
  when (name `Map.member` accRes) $
    throwError $
      ErrorWithPosition ("Duplicate declaration of " ++ name) pos
  when (name `Map.member` vars) $
    throwError $
      ErrorWithPosition (name ++ " already defined as variables") pos
  let heShape = toHEShape shape
  checkVal heShape val
  heVal <- toHEVal heShape val
  return $ Map.insert name (toHEShape shape, heVal) accRes

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

-- | Check constraints declarations
checkConstraintDecl ::
  Context ->
  [HS.ConstraintStatement] ->
  ConstraintDecl ->
  Result [HS.ConstraintStatement]
checkConstraintDecl context@Context {..} acc decl = do
  (constructedExp, boundVal) <-
    case (isVariable exp, bound) of
      (Just varExp, ConstantBound (PIdent (pos, name)))
        | Just (shape, val) <- Map.lookup name consts -> do
          when (shape /= getShape varExp) $
            throwError $
              ErrorWithPosition
                "Shape mismatched: the bound doesn't have same shape as the variable"
                pos
          return (varExp, val)
        | otherwise ->
          throwError $ ErrorWithPosition (name ++ " not found") pos
      (Just varExp, NumberBound num) ->
        return (varExp, HV.VNum (numToDouble num))
      _ -> do
        scalarExp <- constructExp context Nothing exp
        when (getShape scalarExp /= []) $
          throwError $
            ErrorWithPosition
              ( "Higher-dimension (in)equality is not supported yet, here the expression has shape "
                  ++ toReadable (getShape scalarExp)
              )
              (getBeginningPosition exp)
        case bound of
          ConstantBound (PIdent (pos, name))
            | Just (shape, val) <- Map.lookup name consts -> do
              when (shape /= getShape scalarExp) $
                throwError $
                  ErrorWithPosition "The bound must be scalar" pos
              return (scalarExp, val)
            | otherwise ->
              throwError $
                ErrorWithPosition (name ++ " not found") pos
          NumberBound num ->
            return (scalarExp, HV.VNum (numToDouble num))
  let newEntry =
        case decl of
          ConstraintLower {} -> HS.Lower constructedExp boundVal
          ConstraintUpper {} -> HS.Upper constructedExp boundVal
          ConstraintEqual {} ->
            HS.Between constructedExp (boundVal, boundVal)
  return $ acc ++ [newEntry]
  where
    (exp, bound) =
      case decl of
        ConstraintLower exp bound -> (exp, bound)
        ConstraintUpper exp bound -> (exp, bound)
        ConstraintEqual exp bound -> (exp, bound)
    isVariable exp =
      case exp of
        EIdent (PIdent (_, name))
          | Just constructedExp <- Map.lookup name declarations,
            name `Map.member` vars ->
            Just constructedExp
        _ -> Nothing

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

-- | Check immediate declarations
checkLetDecl :: Context -> LetDecl -> Result Context
checkLetDecl context@Context {..} (LetDecl (PIdent (pos, name)) exp) = do
  when (name `Map.member` consts) $
    throwError $
      ErrorWithPosition (name ++ " already defined as a constant") pos
  when (name `Map.member` vars) $
    throwError $
      ErrorWithPosition (name ++ " already defined as a variable") pos
  when (name `Map.member` declarations) $
    throwError $
      ErrorWithPosition (name ++ " already taken") pos
  constructedExp <- constructExp context Nothing exp
  let newDeclarations = Map.insert name constructedExp declarations
  return $ context {declarations = newDeclarations}

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

-- | To HashedExpression's value
toHEVal :: HE.Shape -> Val -> Result HV.Val
toHEVal shape v =
  case v of
    ValFile filePath -> return $ HV.VFile $ HV.TXT filePath
    ValDataset filePath dataset ->
      return $ HV.VFile $ HV.HDF5 filePath dataset
    ValPattern (KWDataPattern pattern) ->
      case (pattern, shape) of
        ("FIRST_ROW_1", [size1, size2]) ->
          return
            . HV.V2D
            . Array.listArray ((0, 0), (size1 - 1, size2 - 1))
            $ replicate size2 1 ++ repeat 0
        ("LAST_ROW_1", [size1, size2]) ->
          return
            . HV.V2D
            . Array.listArray ((0, 0), (size1 - 1, size2 - 1))
            $ replicate (size2 * (size1 - 1)) 0 ++ repeat 1
        ("FIRST_COLUMN_1", [size1, size2]) ->
          return
            . HV.V2D
            . Array.listArray ((0, 0), (size1 - 1, size2 - 1))
            $ concat $
              replicate size1 $
                1 : replicate (size2 - 1) 0
        ("LAST_COLUMN_1", [size1, size2]) ->
          return
            . HV.V2D
            . Array.listArray ((0, 0), (size1 - 1, size2 - 1))
            $ concat $
              replicate size1 $
                replicate (size2 - 1) 0 ++ [1]
        ("FIRST_ROW_0", [size1, size2]) ->
          return
            . HV.V2D
            . Array.listArray ((0, 0), (size1 - 1, size2 - 1))
            $ replicate size2 0 ++ repeat 1
        ("LAST_ROW_0", [size1, size2]) ->
          return
            . HV.V2D
            . Array.listArray ((0, 0), (size1 - 1, size2 - 1))
            $ replicate (size2 * (size1 - 1)) 1 ++ repeat 0
        ("FIRST_COLUMN_0", [size1, size2]) ->
          return
            . HV.V2D
            . Array.listArray ((0, 0), (size1 - 1, size2 - 1))
            $ concat $
              replicate size1 $
                0 : replicate (size2 - 1) 1
        ("LAST_COLUMN_0", [size1, size2]) ->
          return
            . HV.V2D
            . Array.listArray ((0, 0), (size1 - 1, size2 - 1))
            $ concat $
              replicate size1 $
                replicate (size2 - 1) 1 ++ [0]
        _ ->
          throwError $
            GeneralError $
              "Pattern "
                ++ pattern
                ++ " is incompatible with the shape or not supported yet"
    ValRandom -> return $ HV.VNum 3
    ValLiteral num -> return $ HV.VNum (numToDouble num)
    ValImage imgPath -> do
      a <- liftIO $ readImage imgPath
      case a of
        Left err -> throwError $ GeneralError $ "Error reading image at " <> imgPath <> ":" <> err
        Right v -> do
          -- TODO : only support grayscale yet, and this is very slow
          let img = convertRGB8 v
          let col = imageWidth img
              row = imageHeight img
              toGrayscale :: Pixel8 -> Pixel8 -> Pixel8 -> Double
              toGrayscale r g b = (0.2126 * (fromIntegral r) + 0.7152 * (fromIntegral g) + 0.0722 * (fromIntegral b)) / 256
          if (shape /= [row, col])
            then throwError $ GeneralError $ "image size and variable shape don't match, image size is " ++ show row ++ "x" ++ show col
            else
              return $
                HV.V2D $
                  Array.listArray
                    ((0, 0), (row - 1, col - 1))
                    [ toGrayscale r g b
                      | i <- [0 .. row - 1],
                        j <- [0 .. col - 1],
                        let (PixelRGB8 r g b) = pixelAt img j i
                    ]