{-# 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.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 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
  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 e2 e3 -> do
    idx <- interpret' env e1
    val <- interpret' env e2
    accum <- interpret' env e3
    accumAddSparse t p accum idx val
  EZero _ t ezi -> do
    zi <- interpret' env ezi
    return $ zeroM 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

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 (AcIdx 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)

newAcZero :: SMTy t -> Rep (ZeroInfo t) -> IO (RepAc t)
newAcZero typ zi = case typ of
  SMTNil -> return ()
  SMTPair t1 t2 -> bitraverse (newAcZero t1) (newAcZero t2) zi
  SMTLEither{} -> newIORef Nothing
  SMTMaybe _ -> newIORef Nothing
  SMTArr _ t -> arrayMapM (newAcZero t) zi
  SMTScal sty -> numericIsNum sty $ newIORef 0

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

newAcSparse :: SMTy a -> SAcPrj p a b -> Rep (AcIdx p a) -> Rep b -> IO (RepAc a)
newAcSparse typ prj idx val = case (typ, prj) of
  (_, SAPHere) -> newAcDense typ val

  (SMTPair t1 t2, SAPFst prj') ->
    (,) <$> newAcSparse t1 prj' (fst idx) val <*> newAcZero t2 (snd idx)
  (SMTPair t1 t2, SAPSnd prj') ->
    (,) <$> newAcZero t1 (fst idx) <*> newAcSparse t2 prj' (snd idx) val

  (SMTLEither t1 _, SAPLeft prj') -> newIORef . Just . Left =<< newAcSparse t1 prj' idx val
  (SMTLEither _ t2, SAPRight prj') -> newIORef . Just . Right =<< newAcSparse t2 prj' idx val

  (SMTMaybe t1, SAPJust prj') -> newIORef . Just =<< newAcSparse t1 prj' idx val

  (SMTArr n t, SAPArrIdx prj') -> onehotArray (\idx' -> newAcSparse t prj' idx' val) (newAcZero t) n prj' idx

onehotArray :: Monad m
            => (Rep (AcIdx p a) -> m v)  -- ^ the "one"
            -> (Rep (ZeroInfo a) -> m v)  -- ^ the "zero"
            -> SNat n -> SAcPrj p a b -> Rep (AcIdx (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

accumAddDense :: SMTy a -> RepAc a -> Rep a -> AcM s ()
accumAddDense typ ref val = case typ of
  SMTNil -> return ()
  SMTPair t1 t2 -> do
    accumAddDense t1 (fst ref) (fst val)
    accumAddDense t2 (snd ref) (snd val)
  SMTLEither{} ->
    case val of
      Nothing -> return ()
      Just (Left val1) -> accumAddSparse typ (SAPLeft SAPHere) ref () val1
      Just (Right val2) -> accumAddSparse typ (SAPRight SAPHere) ref () val2
  SMTMaybe{} ->
    case val of
      Nothing -> return ()
      Just val' -> accumAddSparse typ (SAPJust SAPHere) ref () val'
  SMTArr _ t1 ->
    forM_ [0 .. arraySize ref - 1] $ \i ->
      accumAddDense t1 (arrayIndexLinear ref i) (arrayIndexLinear val i)
  SMTScal sty -> numericIsNum sty $ AcM $ atomicModifyIORef' ref (\x -> (x + val, ()))

accumAddSparse :: SMTy a -> SAcPrj p a b -> RepAc a -> Rep (AcIdx p a) -> Rep b -> AcM s ()
accumAddSparse typ prj ref idx val = case (typ, prj) of
  (_, SAPHere) -> accumAddDense typ ref val

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

  (SMTLEither t1 _, SAPLeft prj') ->
    realiseMaybeSparse ref (Left <$> newAcSparse t1 prj' idx val)
                           (\case Left ac1 -> accumAddSparse t1 prj' ac1 idx val
                                  Right{} -> error "Mismatched Either in accumAddSparse (r +l)")
  (SMTLEither _ t2, SAPRight prj') ->
    realiseMaybeSparse ref (Right <$> newAcSparse t2 prj' idx val)
                           (\case Right ac2 -> accumAddSparse t2 prj' ac2 idx val
                                  Left{} -> error "Mismatched Either in accumAddSparse (l +r)")

  (SMTMaybe t1, SAPJust prj') ->
    realiseMaybeSparse ref (newAcSparse t1 prj' idx val)
                           (\ac -> accumAddSparse t1 prj' ac idx val)

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


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)