{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE DerivingStrategies #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE ImplicitParams #-}
{-# LANGUAGE KindSignatures #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE MultiWayIf #-}
{-# LANGUAGE PatternSynonyms #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TupleSections #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE TypeOperators #-}
module Interpreter (
  interpret,
  interpretOpen,
  Value(..),
) where

import Control.Monad (foldM, join, when, forM_)
import Data.Bifunctor (bimap)
import Data.Bitraversable (bitraverse)
import Data.Char (isSpace)
import Data.Functor.Identity
import qualified Data.Functor.Product as Product
import Data.Int (Int64)
import Data.IORef
import System.IO (hPutStrLn, stderr)
import System.IO.Unsafe (unsafePerformIO)

import Debug.Trace

import Array
import AST
import AST.Pretty
import AST.Sparse.Types
import Data
import Interpreter.Rep


newtype AcM s a = AcM { unAcM :: IO a }
  deriving newtype (Functor, Applicative, Monad)

runAcM :: (forall s. AcM s a) -> a
runAcM (AcM m) = unsafePerformIO m

acmDebugLog :: String -> AcM s ()
acmDebugLog s = AcM (hPutStrLn stderr s)

data V t = V (STy t) (Rep t)

interpret :: Ex '[] t -> Rep t
interpret = interpretOpen False SNil SNil

-- | Bool: whether to trace execution with debug prints (very verbose)
interpretOpen :: Bool -> SList STy env -> SList Value env -> Ex env t -> Rep t
interpretOpen prints env venv e =
  runAcM $
    let ?depth = 0
        ?prints = prints
    in interpret' (slistMap (\(Product.Pair t (Value v)) -> V t v) (slistZip env venv)) e

interpret' :: forall env t s. (?prints :: Bool, ?depth :: Int)
           => SList V env -> Ex env t -> AcM s (Rep t)
interpret' env e = do
  let tenv = slistMap (\(V t _) -> t) env
  let dep = ?depth
  let lenlimit = max 20 (100 - dep)
  let replace a b = map (\c -> if c == a then b else c)
  let trunc s | length s > lenlimit = take (lenlimit - 3) (replace '\n' ' ' s) ++ "..."
              | otherwise           = replace '\n' ' ' s
  when ?prints $ acmDebugLog $ replicate dep ' ' ++ "ev: " ++ trunc (ppExpr tenv e)
  res <- let ?depth = dep + 1 in interpret'Rec env e
  when ?prints $ acmDebugLog $ replicate dep ' ' ++ "<- " ++ showValue 0 (typeOf e) res ""
  return res

interpret'Rec :: forall env t s. (?prints :: Bool, ?depth :: Int) => SList V env -> Ex env t -> AcM s (Rep t)
interpret'Rec env = \case
  EVar _ _ i -> case slistIdx env i of V _ x -> return x
  ELet _ a b -> do
    x <- interpret' env a
    let ?depth = ?depth - 1 in interpret' (V (typeOf a) x `SCons` env) b
  expr | False && trace ("<i> " ++ takeWhile (not . isSpace) (show expr)) False -> undefined
  EPair _ a b -> (,) <$> interpret' env a <*> interpret' env b
  EFst _ e -> fst <$> interpret' env e
  ESnd _ e -> snd <$> interpret' env e
  ENil _ -> return ()
  EInl _ _ e -> Left <$> interpret' env e
  EInr _ _ e -> Right <$> interpret' env e
  ECase _ e a b ->
    let STEither t1 t2 = typeOf e
    in interpret' env e >>= \case
         Left x -> interpret' (V t1 x `SCons` env) a
         Right y -> interpret' (V t2 y `SCons` env) b
  ENothing _ _ -> return Nothing
  EJust _ e -> Just <$> interpret' env e
  EMaybe _ a b e ->
    let STMaybe t1 = typeOf e
    in maybe (interpret' env a) (\x -> interpret' (V t1 x `SCons` env) b) =<< interpret' env e
  ELNil _ _ _ -> return Nothing
  ELInl _ _ e -> Just . Left <$> interpret' env e
  ELInr _ _ e -> Just . Right <$> interpret' env e
  ELCase _ e a b c ->
    let STLEither t1 t2 = typeOf e
    in interpret' env e >>= \case
         Nothing -> interpret' env a
         Just (Left x) -> interpret' (V t1 x `SCons` env) b
         Just (Right y) -> interpret' (V t2 y `SCons` env) c
  EConstArr _ _ _ v -> return v
  EBuild _ dim a b -> do
    sh <- unTupRepIdx ShNil ShCons dim <$> interpret' env a
    arrayGenerateM sh (\idx -> interpret' (V (tTup (sreplicate dim tIx)) (tupRepIdx ixUncons dim idx) `SCons` env) b)
  EFold1Inner _ _ a b c -> do
    let t = typeOf b
    let f = \x y -> interpret' (V t y `SCons` V t x `SCons` env) a
    x0 <- interpret' env b
    arr <- interpret' env c
    let sh `ShCons` n = arrayShape arr
    arrayGenerateM sh $ \idx -> foldM f x0 [arrayIndex arr (idx `IxCons` i) | i <- [0 .. n - 1]]
  ESum1Inner _ e -> do
    arr <- interpret' env e
    let STArr _ (STScal t) = typeOf e
        sh `ShCons` n = arrayShape arr
    numericIsNum t $ return $ arrayGenerate sh $ \idx -> sum [arrayIndex arr (idx `IxCons` i) | i <- [0 .. n - 1]]
  EUnit _ e -> arrayGenerateLinM ShNil (\_ -> interpret' env e)
  EReplicate1Inner _ a b -> do
    n <- fromIntegral @Int64 @Int <$> interpret' env a
    arr <- interpret' env b
    let sh = arrayShape arr
    return $ arrayGenerate (sh `ShCons` n) (\(idx `IxCons` _) -> arrayIndex arr idx)
  EMaximum1Inner _ e -> do
    arr <- interpret' env e
    let STArr _ (STScal t) = typeOf e
        sh `ShCons` n = arrayShape arr
    numericIsNum t $ return $
      arrayGenerate sh (\idx -> maximum [arrayIndex arr (idx `IxCons` i) | i <- [0 .. n-1]])
  EMinimum1Inner _ e -> do
    arr <- interpret' env e
    let STArr _ (STScal t) = typeOf e
        sh `ShCons` n = arrayShape arr
    numericIsNum t $ return $
      arrayGenerate sh (\idx -> minimum [arrayIndex arr (idx `IxCons` i) | i <- [0 .. n-1]])
  EConst _ _ v -> return v
  EIdx0 _ e -> (`arrayIndexLinear` 0) <$> interpret' env e
  EIdx1 _ a b -> arrayIndex1 <$> interpret' env a <*> (fromIntegral @Int64 @Int <$> interpret' env b)
  EIdx _ a b ->
    let STArr n _ = typeOf a
    in arrayIndex <$> interpret' env a <*> (unTupRepIdx IxNil IxCons n <$> interpret' env b)
  EShape _ e | STArr n _ <- typeOf e -> tupRepIdx shUncons n . arrayShape <$> interpret' env e
  EOp _ op e -> interpretOp op <$> interpret' env e
  ECustom _ t1 t2 _ pr _ _ e1 e2 -> do
    e1' <- interpret' env e1
    e2' <- interpret' env e2
    interpret' (V t2 e2' `SCons` V t1 e1' `SCons` SNil) pr
  ERecompute _ e -> interpret' env e
  EWith _ t e1 e2 -> do
    initval <- interpret' env e1
    withAccum t (typeOf e2) initval $ \accum ->
      interpret' (V (STAccum t) accum `SCons` env) e2
  EAccum _ t p e1 sp e2 e3 -> do
    idx <- interpret' env e1
    val <- interpret' env e2
    accum <- interpret' env e3
    accumAddSparseD t p accum idx sp val
  EZero _ t ezi -> do
    zi <- interpret' env ezi
    return $ zeroM t zi
  EDeepZero _ t ezi -> do
    zi <- interpret' env ezi
    return $ deepZeroM t zi
  EPlus _ t a b -> do
    a' <- interpret' env a
    b' <- interpret' env b
    return $ addM t a' b'
  EOneHot _ t p a b -> do
    a' <- interpret' env a
    b' <- interpret' env b
    return $ onehotM p t a' b'
  EError _ _ s -> error $ "Interpreter: Program threw error: " ++ s

interpretOp :: SOp a t -> Rep a -> Rep t
interpretOp op arg = case op of
  OAdd st -> numericIsNum st $ uncurry (+) arg
  OMul st -> numericIsNum st $ uncurry (*) arg
  ONeg st -> numericIsNum st $ negate arg
  OLt st -> numericIsNum st $ uncurry (<) arg
  OLe st -> numericIsNum st $ uncurry (<=) arg
  OEq st -> styIsEq st $ uncurry (==) arg
  ONot -> not arg
  OAnd -> uncurry (&&) arg
  OOr -> uncurry (||) arg
  OIf -> if arg then Left () else Right ()
  ORound64 -> round arg
  OToFl64 -> fromIntegral arg
  ORecip st -> floatingIsFractional st $ recip arg
  OExp st -> floatingIsFractional st $ exp arg
  OLog st -> floatingIsFractional st $ log arg
  OIDiv st -> integralIsIntegral st $ uncurry quot arg
  OMod st -> integralIsIntegral st $ uncurry rem arg
  where
    styIsEq :: SScalTy t -> (Eq (Rep (TScal t)) => r) -> r
    styIsEq STI32 = id
    styIsEq STI64 = id
    styIsEq STF32 = id
    styIsEq STF64 = id
    styIsEq STBool = id

zeroM :: SMTy t -> Rep (ZeroInfo t) -> Rep t
zeroM typ zi = case typ of
  SMTNil -> ()
  SMTPair t1 t2 -> (zeroM t1 (fst zi), zeroM t2 (snd zi))
  SMTLEither _ _ -> Nothing
  SMTMaybe _ -> Nothing
  SMTArr _ t -> arrayMap (zeroM t) zi
  SMTScal sty -> case sty of
                  STI32 -> 0
                  STI64 -> 0
                  STF32 -> 0.0
                  STF64 -> 0.0

deepZeroM :: SMTy t -> Rep (DeepZeroInfo t) -> Rep t
deepZeroM typ zi = case typ of
  SMTNil -> ()
  SMTPair t1 t2 -> (deepZeroM t1 (fst zi), deepZeroM t2 (snd zi))
  SMTLEither t1 t2 -> fmap (bimap (deepZeroM t1) (deepZeroM t2)) zi
  SMTMaybe t -> fmap (deepZeroM t) zi
  SMTArr _ t -> arrayMap (deepZeroM t) zi
  SMTScal sty -> case sty of
                  STI32 -> 0
                  STI64 -> 0
                  STF32 -> 0.0
                  STF64 -> 0.0

addM :: SMTy t -> Rep t -> Rep t -> Rep t
addM typ a b = case typ of
  SMTNil -> ()
  SMTPair t1 t2 -> (addM t1 (fst a) (fst b), addM t2 (snd a) (snd b))
  SMTLEither t1 t2 -> case (a, b) of
    (Nothing, _) -> b
    (_, Nothing) -> a
    (Just (Left x), Just (Left y)) -> Just (Left (addM t1 x y))
    (Just (Right x), Just (Right y)) -> Just (Right (addM t2 x y))
    _ -> error "Plus of inconsistent LEithers"
  SMTMaybe t -> case (a, b) of
    (Nothing, _) -> b
    (_, Nothing) -> a
    (Just x, Just y) -> Just (addM t x y)
  SMTArr _ t ->
    let sh1 = arrayShape a
        sh2 = arrayShape b
    in if | shapeSize sh1 == 0 -> b
          | shapeSize sh2 == 0 -> a
          | sh1 == sh2 -> arrayGenerateLin sh1 (\i -> addM t (arrayIndexLinear a i) (arrayIndexLinear b i))
          | otherwise -> error "Plus of inconsistently shaped arrays"
  SMTScal sty -> numericIsNum sty $ a + b

onehotM :: SAcPrj p a b -> SMTy a -> Rep (AcIdxS p a) -> Rep b -> Rep a
onehotM SAPHere _ _ val = val
onehotM (SAPFst prj) (SMTPair a b) idx val = (onehotM prj a (fst idx) val, zeroM b (snd idx))
onehotM (SAPSnd prj) (SMTPair a b) idx val = (zeroM a (fst idx), onehotM prj b (snd idx) val)
onehotM (SAPLeft prj) (SMTLEither a _) idx val = Just (Left (onehotM prj a idx val))
onehotM (SAPRight prj) (SMTLEither _ b) idx val = Just (Right (onehotM prj b idx val))
onehotM (SAPJust prj) (SMTMaybe a) idx val = Just (onehotM prj a idx val)
onehotM (SAPArrIdx prj) (SMTArr n a) idx val =
  runIdentity $ onehotArray (\idx' -> Identity (onehotM prj a idx' val)) (\zi -> Identity (zeroM a zi)) n prj idx

withAccum :: SMTy t -> STy a -> Rep t -> (RepAc t -> AcM s (Rep a)) -> AcM s (Rep a, Rep t)
withAccum t _ initval f = AcM $ do
  accum <- newAcDense t initval
  out <- unAcM $ f accum
  val <- readAc t accum
  return (out, val)

newAcDense :: SMTy a -> Rep a -> IO (RepAc a)
newAcDense typ val = case typ of
  SMTNil -> return ()
  SMTPair t1 t2 -> bitraverse (newAcDense t1) (newAcDense t2) val
  SMTLEither t1 t2 -> newIORef =<< traverse (bitraverse (newAcDense t1) (newAcDense t2)) val
  SMTMaybe t1 -> newIORef =<< traverse (newAcDense t1) val
  SMTArr _ t1 -> arrayMapM (newAcDense t1) val
  SMTScal _ -> newIORef val

onehotArray :: Monad m
            => (Rep (AcIdxS p a) -> m v)  -- ^ the "one"
            -> (Rep (ZeroInfo a) -> m v)  -- ^ the "zero"
            -> SNat n -> SAcPrj p a b -> Rep (AcIdxS (APArrIdx p) (TArr n a)) -> m (Array n v)
onehotArray mkone mkzero n _ ((arrindex', ziarr), idx) =
  let arrindex = unTupRepIdx IxNil IxCons n arrindex'
      arrsh = arrayShape ziarr
      !linindex = toLinearIndex arrsh arrindex
  in arrayGenerateLinM arrsh (\i -> if i == linindex then mkone idx else mkzero (ziarr `arrayIndexLinear` i))

readAc :: SMTy t -> RepAc t -> IO (Rep t)
readAc typ val = case typ of
  SMTNil -> return ()
  SMTPair t1 t2 -> bitraverse (readAc t1) (readAc t2) val
  SMTLEither t1 t2 -> traverse (bitraverse (readAc t1) (readAc t2)) =<< readIORef val
  SMTMaybe t -> traverse (readAc t) =<< readIORef val
  SMTArr _ t -> traverse (readAc t) val
  SMTScal _ -> readIORef val

accumAddSparseD :: SMTy a -> SAcPrj p a b -> RepAc a -> Rep (AcIdxD p a) -> Sparse b c -> Rep c -> AcM s ()
accumAddSparseD typ prj ref idx sp val = case (typ, prj) of
  (_, SAPHere) -> accumAddDense typ ref sp val

  (SMTPair t1 _, SAPFst prj') -> accumAddSparseD t1 prj' (fst ref) idx sp val
  (SMTPair _ t2, SAPSnd prj') -> accumAddSparseD t2 prj' (snd ref) idx sp val

  (SMTLEither t1 _, SAPLeft prj') ->
    realiseMaybeSparse ref (error "Accumulating Left into LNil (EWith requires EDeepZero)")
                           (\case Left ac1 -> accumAddSparseD t1 prj' ac1 idx sp val
                                  Right{} -> error "Mismatched Either in accumAddSparseD (r +l)")
  (SMTLEither _ t2, SAPRight prj') ->
    realiseMaybeSparse ref (error "Accumulating Right into LNil (EWith requires EDeepZero)")
                           (\case Right ac2 -> accumAddSparseD t2 prj' ac2 idx sp val
                                  Left{} -> error "Mismatched Either in accumAddSparseD (l +r)")

  (SMTMaybe t1, SAPJust prj') ->
    realiseMaybeSparse ref (error "Accumulating Just into Nothing (EWith requires EDeepZero)")
                           (\ac -> accumAddSparseD t1 prj' ac idx sp val)

  (SMTArr n t1, SAPArrIdx prj') ->
    let (arrindex', idx') = idx
        arrindex = unTupRepIdx IxNil IxCons n arrindex'
        arrsh = arrayShape ref
        linindex = toLinearIndex arrsh arrindex
    in accumAddSparseD t1 prj' (arrayIndexLinear ref linindex) idx' sp val

accumAddDense :: SMTy a -> RepAc a -> Sparse a b -> Rep b -> AcM s ()
accumAddDense typ ref sp val = case (typ, sp) of
  (_, _) | isAbsent sp -> return ()
  (_, SpAbsent) -> return ()
  (_, SpSparse s) ->
    case val of
      Nothing -> return ()
      Just val' -> accumAddDense typ ref s val'
  (SMTPair t1 t2, SpPair s1 s2) -> do
    accumAddDense t1 (fst ref) s1 (fst val)
    accumAddDense t2 (snd ref) s2 (snd val)
  (SMTLEither t1 t2, SpLEither s1 s2) ->
    case val of
      Nothing -> return ()
      Just (Left val1) ->
        realiseMaybeSparse ref (error "Accumulating Left into LNil (EWith requires EDeepZero)")
                               (\case Left ac1 -> accumAddDense t1 ac1 s1 val1
                                      Right{} -> error "Mismatched Either in accumAddSparse (r +l)")
      Just (Right val2) ->
        realiseMaybeSparse ref (error "Accumulating Right into LNil (EWith requires EDeepZero)")
                               (\case Right ac2 -> accumAddDense t2 ac2 s2 val2
                                      Left{} -> error "Mismatched Either in accumAddSparse (l +r)")
  (SMTMaybe t, SpMaybe s) ->
    case val of
      Nothing -> return ()
      Just val' ->
        realiseMaybeSparse ref (error "Accumulating Just into Nothing (EAccum requires EDeepZero)")
                               (\ac -> accumAddDense t ac s val')
  (SMTArr _ t1, SpArr s) ->
    forM_ [0 .. arraySize ref - 1] $ \i ->
      accumAddDense t1 (arrayIndexLinear ref i) s (arrayIndexLinear val i)
  (SMTScal sty, SpScal) -> numericIsNum sty $ AcM $ atomicModifyIORef' ref (\x -> (x + val, ()))

-- TODO: makeval is always 'error' now. Simplify?
realiseMaybeSparse :: IORef (Maybe a) -> IO a -> (a -> AcM s ()) -> AcM s ()
realiseMaybeSparse ref makeval modifyval =
  -- Try modifying what's already in ref. The 'join' makes the snd
  -- of the function's return value a _continuation_ that is run after
  -- the critical section ends.
  AcM $ join $ atomicModifyIORef' ref $ \ac -> case ac of
    -- Oops, ref's contents was still sparse. Have to initialise
    -- it first, then try again.
    Nothing -> (ac, do val <- makeval
                       join $ atomicModifyIORef' ref $ \ac' -> case ac' of
                                Nothing -> (Just val, return ())
                                Just val' -> (ac', unAcM $ modifyval val'))
    -- Yep, ref already had a value in there, so we can just add
    -- val' to it recursively.
    Just val -> (ac, unAcM $ modifyval val)


numericIsNum :: ScalIsNumeric st ~ True => SScalTy st -> ((Num (ScalRep st), Ord (ScalRep st)) => r) -> r
numericIsNum STI32 = id
numericIsNum STI64 = id
numericIsNum STF32 = id
numericIsNum STF64 = id

floatingIsFractional :: ScalIsFloating st ~ True => SScalTy st -> ((Floating (ScalRep st), Ord (ScalRep st), ScalIsNumeric st ~ True, ScalIsFloating st ~ True) => r) -> r
floatingIsFractional STF32 = id
floatingIsFractional STF64 = id

integralIsIntegral :: ScalIsIntegral st ~ True => SScalTy st -> ((Integral (ScalRep st), Ord (ScalRep st), ScalIsNumeric st ~ True, ScalIsIntegral st ~ True) => r) -> r
integralIsIntegral STI32 = id
integralIsIntegral STI64 = id

unTupRepIdx :: f Z -> (forall m. f m -> Int -> f (S m))
            -> SNat n -> Rep (Tup (Replicate n TIx)) -> f n
unTupRepIdx nil _    SZ     _        = nil
unTupRepIdx nil cons (SS n) (idx, i) = unTupRepIdx nil cons n idx `cons` fromIntegral @Int64 @Int i

tupRepIdx :: (forall m. f (S m) -> (f m, Int))
          -> SNat n -> f n -> Rep (Tup (Replicate n TIx))
tupRepIdx _      SZ     _   = ()
tupRepIdx uncons (SS n) tup =
  let (tup', i) = uncons tup
  in ((,) $! tupRepIdx uncons n tup') $! fromIntegral @Int @Int64 i

ixUncons :: Index (S n) -> (Index n, Int)
ixUncons (IxCons idx i) = (idx, i)

shUncons :: Shape (S n) -> (Shape n, Int)
shUncons (ShCons idx i) = (idx, i)