content: module Language.Futhark.Interpreter.AD
    ( Op (..),
      ADVariable (..),
      ADValue (..),
      Tape (..),
      VJPValue (..),
      JVPValue (..),
      Counter (..),
      Depth (..),
      doOp,
      addFor,
      tapePrimal,
      primitive,
      varPrimal,
      deriveTape,
      unionWithM,
      unionsWithM,
    )
  where
  
  import Control.Monad (foldM, zipWithM)
  import Control.Monad.Trans.Class (lift)
  import Control.Monad.Trans.Except (ExceptT, catchE, runExceptT, throwE)
  import Control.Monad.Trans.State (State, get, modify, runState)
  import Data.Either (isRight)
  import Data.Foldable (find, foldlM)
  import Data.Functor ((<&>))
  import Data.Map qualified as M
  import Data.Maybe (fromJust, fromMaybe)
  import Data.Text qualified as T
  import Futhark.AD.Derivatives (pdBinOp, pdBuiltin, pdUnOp)
  import Futhark.Analysis.PrimExp (PrimExp (..))
  import Language.Futhark.Core (VName (..), nameFromString, nameFromText)
  import Language.Futhark.Primitive
    ( BinOp (Add, FAdd, FMul, LogAnd, LogOr, Mul),
      CmpOp,
      ConvOp,
      Overflow (OverflowWrap),
      PrimType (Bool, FloatType, IntType),
      PrimValue (BoolValue),
      UnOp,
      binOpType,
      blankPrimValue,
      cmpOpType,
      convOpType,
      doBinOp,
      doCmpOp,
      doConvOp,
      doUnOp,
      flipConvOp,
      primFuns,
      primValueType,
      unOpType,
    )
  
  -- | Used to uniquely identify values.
  newtype Counter = Counter Int
    deriving (Eq, Ord, Num, Show)
  
  type ADMonad = ExceptT String (State Counter)
  
  incCounter :: ADMonad ()
  incCounter = lift $ modify $ \i -> i + 1
  
  -- Mathematical operations subject to AD.
  data Op
    = OpBin BinOp
    | OpCmp CmpOp
    | OpUn UnOp
    | OpFn T.Text
    | OpConv ConvOp
    deriving (Show)
  
  -- Checks if an operation matches the types of its operands
  opTypeMatch :: Op -> [PrimType] -> Bool
  opTypeMatch (OpBin op) p = all (\x -> binOpType op == x) p
  opTypeMatch (OpCmp op) p = all (\x -> cmpOpType op == x) p
  opTypeMatch (OpUn op) p = all (\x -> unOpType op == x) p
  opTypeMatch (OpConv op) p = all (\x -> fst (convOpType op) == x) p
  opTypeMatch (OpFn fn) p = case M.lookup fn primFuns of
    Just (t, _, _) -> and $ zipWith (==) t p
    Nothing -> error "opTypeMatch" -- It is assumed that the function exists
  
  -- Gets the return type of an operation
  opReturnType :: Op -> PrimType
  opReturnType (OpBin op) = binOpType op
  opReturnType (OpCmp op) = cmpOpType op
  opReturnType (OpUn op) = unOpType op
  opReturnType (OpConv op) = snd $ convOpType op
  opReturnType (OpFn fn) = case M.lookup fn primFuns of
    Just (_, t, _) -> t
    Nothing -> error "opReturnType" -- It is assumed that the function exists
  
  -- Returns the operation which performs addition (or an
  -- equivalent operation) on the given type
  addFor :: PrimType -> BinOp
  addFor (IntType t) = Add t OverflowWrap
  addFor (FloatType t) = FAdd t
  addFor Bool = LogOr
  addFor t = error $ "addFor: " ++ show t
  
  -- Returns the function which performs multiplication
  -- (or an equivalent operation) on the given type
  mulFor :: PrimType -> BinOp
  mulFor (IntType t) = Mul t OverflowWrap
  mulFor (FloatType t) = FMul t
  mulFor Bool = LogAnd
  mulFor t = error $ "mulFor: " ++ show t
  
  -- | An indication of the nesting depth of AD. This is used to avoid
  -- pertubation confusion.
  newtype Depth = Depth Int
    deriving (Ord, Eq, Show)
  
  -- Types and utility functions--
  -- When taking the partial derivative of a function, we
  -- must differentiate between the values which are kept
  -- constant, and those which are not
  data ADValue
    = Variable Depth ADVariable
    | Constant PrimValue
    deriving (Show)
  
  -- When performing automatic differentiation, each derived
  -- variable must be augmented with additional data. This
  -- value holds the primitive value of the variable, as well
  -- as its data
  data ADVariable
    = VJP VJPValue
    | JVP JVPValue
    deriving (Show)
  
  depth :: ADValue -> Depth
  depth (Variable d _) = d
  depth (Constant _) = Depth 0
  
  primal :: ADValue -> ADValue
  primal (Variable _ (VJP (VJPValue t))) = tapePrimal t
  primal (Variable _ (JVP (JVPValue v _))) = primal v
  primal (Constant v) = Constant v
  
  primalFor :: Depth -> ADValue -> ADValue
  primalFor cur v@(Variable tag _) | cur /= tag = v
  primalFor _ (Variable _ (VJP (VJPValue t))) = tapePrimal t
  primalFor cur (Variable _ (JVP (JVPValue v _))) = primalFor cur v
  primalFor _ (Constant v) = Constant v
  
  primitive :: ADValue -> PrimValue
  primitive (Variable _ v) = varPrimal v
  primitive (Constant v) = v
  
  varPrimal :: ADVariable -> PrimValue
  varPrimal (VJP (VJPValue t)) = primitive $ tapePrimal t
  varPrimal (JVP (JVPValue v _)) = primitive $ primal v
  
  -- Evaluates a PrimExp using doOp'
  evalPrimExp :: M.Map VName ADValue -> PrimExp VName -> ADMonad ADValue
  evalPrimExp m (LeafExp n _) =
    maybe (throwE $ "Unknown variable " <> show n) pure $ M.lookup n m
  evalPrimExp _ (ValueExp pv) =
    pure $ Constant pv
  evalPrimExp m (BinOpExp op x y) = do
    x' <- evalPrimExp m x
    y' <- evalPrimExp m y
    doOp' (OpBin op) [x', y']
  evalPrimExp m (CmpOpExp op x y) = do
    x' <- evalPrimExp m x
    y' <- evalPrimExp m y
    doOp' (OpCmp op) [x', y']
  evalPrimExp m (UnOpExp op x) = do
    x' <- evalPrimExp m x
    doOp' (OpUn op) [x']
  evalPrimExp m (ConvOpExp op x) = do
    x' <- evalPrimExp m x
    doOp' (OpConv op) [x']
  evalPrimExp m (FunExp fn p _) = do
    p' <- mapM (evalPrimExp m) p
    doOp' (OpFn fn) p'
  
  -- Returns a list of PrimExps calculating the partial
  -- derivative of each operands of a given operation
  lookupPDs :: Op -> [PrimExp VName] -> Maybe [PrimExp VName]
  lookupPDs (OpBin op) [x, y] = Just $ do
    let (a, b) = pdBinOp op x y
    [a, b]
  lookupPDs (OpUn op) [x] = Just [pdUnOp op x]
  lookupPDs (OpFn fn) p = pdBuiltin (nameFromText fn) p
  lookupPDs _ _ = Nothing
  
  -- Shared AD logic--
  -- This function performs a mathematical operation on a
  -- list of operands, performing automatic differentiation
  -- if one or more operands is a Variable (of depth > 0)
  doOp :: Op -> [ADValue] -> Counter -> Either String (ADValue, Counter)
  doOp op o uid = case runState (runExceptT $ doOp' op o) uid of
    (Left s, _) -> Left s
    (Right v, uid') -> Right (v, uid')
  
  doOp' :: Op -> [ADValue] -> ADMonad ADValue
  doOp' op o
    | not $ opTypeMatch op (map primValueType pv) =
        -- This function may be called with arguments of invalid types,
        -- because it is used as part of an overloaded operator.
        throwE $ unwords ["invalid types for op", show op, "and operands", show o]
    | otherwise = do
        let dep = case op of
              OpCmp _ -> Depth 0 -- AD is not well-defined for comparason operations
              -- There are no derivatives for those written in
              -- PrimExp (check lookupPDs)
              _ -> maximum (map depth o)
        if dep == Depth 0
          then maybe (throwE "failed to evaluate const") pure constCase <* incCounter
          else nonconstCase dep
    where
      pv = map primitive o
  
      divideDepths :: Depth -> ADValue -> Either ADValue ADVariable
      divideDepths _ v@(Constant {}) = Left v
      divideDepths d v@(Variable d' v') = if d' < d then Left v else Right v'
  
      -- TODO: There may be a more graceful way of
      -- doing this
      extractVJP :: Either ADValue ADVariable -> Either ADValue VJPValue
      extractVJP (Right (VJP v)) = Right v
      extractVJP (Left v) = Left v
      extractVJP _ =
        -- This will never be called when the maximum depth layer is JVP
        error "extractVJP"
  
      -- TODO: There may be a more graceful way of
      -- doing this
      extractJVP :: Either ADValue ADVariable -> Either ADValue JVPValue
      extractJVP (Right (JVP v)) = Right v
      extractJVP (Left v) = Left v
      extractJVP _ =
        -- This will never be called when the maximum depth layer is VJP
        error "extractJVP"
  
      -- In this case, every operand is a constant, and the
      -- mathematical operation can be applied as it would be
      -- otherwise
      constCase =
        Constant <$> case (op, pv) of
          (OpBin op', [x, y]) -> doBinOp op' x y
          (OpCmp op', [x, y]) -> BoolValue <$> doCmpOp op' x y
          (OpUn op', [x]) -> doUnOp op' x
          (OpConv op', [x]) -> doConvOp op' x
          (OpFn fn, _) -> do
            (_, _, f) <- M.lookup fn primFuns
            f pv
          _ -> error "doOp': opTypeMatch"
  
      nonconstCase dep = do
        -- In this case, some values are variables. We therefore
        -- have to perform the necessary steps for AD
  
        -- First, we calculate the value for the previous depth
        let oprev = map (primalFor dep) o
        vprev <- doOp' op oprev
  
        -- Then we separate the values of the maximum depth from
        -- those of a lower depth
        let o' = map (divideDepths dep) o
        -- Then we find out what type of AD is being performed
        case find isRight o' of
          -- Finally, we perform the necessary steps for the given
          -- type of AD
          Just (Right (VJP {})) ->
            Variable dep . VJP . VJPValue
              <$> vjpHandleOp op (map extractVJP o') vprev
          Just (Right (JVP {})) ->
            Variable dep . JVP . JVPValue vprev
              <$> jvpHandleOp op (map extractJVP o')
          _ ->
            -- Since the maximum depth is non-zero, there must be at
            -- least one variable of depth > 0
            error "find isRight"
  
  calculatePDs :: Op -> [ADValue] -> ADMonad [ADValue]
  calculatePDs op args =
    -- Create a unique VName for each operand
    let n = map (\i -> VName (nameFromString $ "x" ++ show i) i) [1 .. length args]
        -- Put the operands in the environment
        m = M.fromList $ zip n args
  
        -- Look up, and calculate the partial derivative
        -- of the operation with respect to each operand
        pde =
          fromMaybe (error "lookupPDs failed") $
            lookupPDs op $
              zipWith (\v val -> LeafExp v $ primValueType $ primitive val) n args
        res = mapM (\x -> catchE (evalPrimExp m x) $ error . ("evalPrimExp failed: " <>)) pde
     in res
  
  -- VJP / Reverse mode automatic differentiation--
  -- In reverse mode AD, the entire computation
  -- leading up to a variable must be saved
  -- This is represented as a Tape
  newtype VJPValue = VJPValue Tape
    deriving (Show)
  
  -- | Represents a computation tree, as well as every intermediate
  -- value in its evaluation.
  data Tape
    = -- | This represents a variable. Each variable is given a unique ID,
      -- and has an initial value
      TapeID Counter ADValue
    | -- | This represents a constant.
      TapeConst ADValue
    | -- | This represents the application of a mathematical operation.
      -- Each parameter is given by its Tape, and the return value of
      -- the operation is saved
      TapeOp Op [Tape] Counter ADValue
    deriving (Show)
  
  -- | Returns the primal value of a Tape.
  tapePrimal :: Tape -> ADValue
  tapePrimal (TapeID _ v) = v
  tapePrimal (TapeConst v) = v
  tapePrimal (TapeOp _ _ _ v) = v
  
  -- This updates Tape of a VJPValue with a new operation,
  -- treating all operands of a lower depth as constants
  vjpHandleOp :: Op -> [Either ADValue VJPValue] -> ADValue -> ADMonad Tape
  vjpHandleOp op p v = do
    i <- lift get
    pure $ TapeOp op (map toTape p) i v
    where
      toTape (Left v') = TapeConst v'
      toTape (Right (VJPValue t)) = t
  
  unionWithM :: (Monad m, Ord k) => (a -> a -> m a) -> M.Map k a -> M.Map k a -> m (M.Map k a)
  unionWithM f m1 m2 = do
    let m = M.union (M.difference m1 m2) (M.difference m2 m1)
    let k = M.keys $ M.intersection m1 m2
    v <- mapM (\k' -> f (fromJust $ M.lookup k' m1) (fromJust $ M.lookup k' m2)) k
    pure $ foldl (\m' (k', v') -> M.insert k' v' m') m (zip k v)
  
  unionsWithM :: (Foldable f, Monad m, Ord k) => (a -> a -> m a) -> f (M.Map k a) -> m (M.Map k a)
  unionsWithM f = foldM (unionWithM f) M.empty
  
  -- | This calculates every partial derivative of a 'Tape'. The result
  -- is a map of the partial derivatives, each key corresponding to the
  -- ID of a free variable (see TapeID).
  deriveTape :: Tape -> ADValue -> Counter -> Either String (M.Map Counter ADValue, Counter)
  deriveTape tp s uid = case runState (runExceptT $ deriveTape' tp s) uid of
    (Left e, _) -> Left e
    (Right v, uid') -> Right (v, uid')
  
  deriveTape' :: Tape -> ADValue -> ADMonad (M.Map Counter ADValue)
  deriveTape' (TapeID i _) s = pure $ M.singleton i s
  deriveTape' (TapeConst _) _ = pure M.empty
  deriveTape' tp@(TapeOp op p uid _) s =
    fst <$> derive tp s M.empty (countReferences p $ M.singleton (-uid - 1) 1)
    where
      add x y = doOp' (OpBin $ addFor $ opReturnType op) [x, y]
      mul x y = doOp' (OpBin $ mulFor $ opReturnType op) [x, y]
      madd :: Counter -> ADValue -> M.Map Counter ADValue -> ADMonad (M.Map Counter ADValue)
      madd i a m = case M.lookup i m of
        Just b -> add a b <&> (\x -> M.insert i x m)
        Nothing -> pure $ M.insert i a m
      derive ::
        Tape ->
        ADValue ->
        M.Map Counter ADValue ->
        M.Map Counter Int ->
        ADMonad (M.Map Counter ADValue, M.Map Counter Int)
      derive (TapeID i _) s' ss rs = madd i s' ss <&> (,rs)
      derive (TapeConst _) _ ss rs = pure (ss, rs)
      derive (TapeOp op' p' uid' _) s' ss rs = do
        -- Decrease the reference counter
        let r = fromJust (M.lookup (-uid' - 1) rs) - 1
            rs' = M.insert (-uid' - 1) r rs
        -- Add the sensitivity
        ss' <- madd (-uid' - 1) s' ss
        -- If there are still more references left, do nothing
        if r > 0
          then pure (ss', rs')
          else -- Otherwise, derive the tape
            if r == 0
              then do
                let s'' = fromJust (M.lookup (-uid' - 1) ss')
  
                -- Calculate the new sensitivities
                s''' <- case op' of
                  OpConv op'' ->
                    -- In case of type conversion, simply convert the sensitivity
                    sequence [doOp' (OpConv $ flipConvOp op'') [s'']]
                  _ -> calculatePDs op' (map tapePrimal p') >>= mapM (mul s'')
  
                -- Propagate the new sensitivities
                foldlM (\(ss'', rs'') (p'', s'''') -> derive p'' s'''' ss'' rs'') (ss', rs') $ zip p' s'''
              else error "TODO: This branch is unreachable unless `countReferences` undercounts"
      countReferences :: [Tape] -> M.Map Counter Int -> M.Map Counter Int
      countReferences p' d' = foldl f d' p'
      f d'' x =
        case x of
          (TapeOp _ p'' uid'' _) -> case M.lookup (-uid'' - 1) d'' of
            Just v -> M.insert (-uid'' - 1) (v + 1) d''
            Nothing -> countReferences p'' $ M.insert (-uid'' - 1) 1 d''
          _ -> d''
  
  -- JVP / Forward mode automatic differentiation--
  
  -- | In JVP, the derivative of the variable must be saved. This is
  -- represented as a second value.
  data JVPValue = JVPValue ADValue ADValue
    deriving (Show)
  
  -- | This calculates the tangent part of the JVPValue resulting
  -- from the application of a mathematical operation on one or more
  -- JVPValues.
  jvpHandleOp :: Op -> [Either ADValue JVPValue] -> ADMonad ADValue
  jvpHandleOp op p = do
    case op of
      OpConv _ ->
        -- In case of type conversion, simply convert
        -- the old tangent
        doOp' op [tangent $ head p]
      _ -> do
        -- Calculate the new tangent using the chain rule
        pds <- calculatePDs op $ map primal' p
        vs <- zipWithM mul pds $ map tangent p
        foldM add (Constant $ blankPrimValue op_t) vs
    where
      op_t = opReturnType op
      primal' (Left v) = v
      primal' (Right (JVPValue v _)) = v
      tangent (Left _) = Constant $ blankPrimValue $ opReturnType op
      tangent (Right (JVPValue _ d)) = d
      add x y = doOp' (OpBin $ addFor $ opReturnType op) [x, y]
      mul x y = doOp' (OpBin $ mulFor $ opReturnType op) [x, y]
  