{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE LambdaCase #-}

-- |
-- Module      : Jikka.Core.Convert.PropagateMod
-- Description : propagates modulo operations, and replaces integer functions with corresponding functions with modulo. / 剰余演算を伝播させ、整数の関数を対応する modulo 付きの関数で置き換えます。
-- Copyright   : (c) Kimiyuki Onaka, 2021
-- License     : Apache License 2.0
-- Maintainer  : kimiyuki95@gmail.com
-- Stability   : experimental
-- Portability : portable
module Jikka.Core.Convert.PropagateMod
  ( run,
  )
where

import Control.Monad.Trans.Maybe
import Data.List
import Data.Maybe
import Jikka.Common.Alpha
import Jikka.Common.Error
import Jikka.Core.Format (formatType)
import Jikka.Core.Language.Beta
import Jikka.Core.Language.BuiltinPatterns
import Jikka.Core.Language.Expr
import Jikka.Core.Language.Lint
import Jikka.Core.Language.ModuloExpr
import Jikka.Core.Language.RewriteRules
import Jikka.Core.Language.TypeCheck
import Jikka.Core.Language.Util

isModulo' :: Expr -> Expr -> Bool
isModulo' e m = e `isModulo` Modulo m

putFloorMod :: MonadAlpha m => Modulo -> Expr -> m (Maybe Expr)
putFloorMod (Modulo m) =
  runMaybeT . \case
    Negate' e -> return $ ModNegate' e m
    Plus' e1 e2 -> return $ ModPlus' e1 e2 m
    Minus' e1 e2 -> return $ ModMinus' e1 e2 m
    Mult' e1 e2 -> return $ ModMult' e1 e2 m
    JustDiv' e1 e2 -> return $ ModMult' e1 (ModInv' e2 m) m
    Pow' e1 e2 -> return $ ModPow' e1 e2 m
    MatAp' h w e1 e2 -> return $ ModMatAp' h w e1 e2 m
    MatAdd' h w e1 e2 -> return $ ModMatAdd' h w e1 e2 m
    MatMul' h n w e1 e2 -> return $ ModMatMul' h n w e1 e2 m
    MatPow' n e1 e2 -> return $ ModMatPow' n e1 e2 m
    Sum' e -> return $ ModSum' e m
    Product' e -> return $ ModProduct' e m
    LitInt' n -> case m of
      LitInt' m -> return $ LitInt' (n `mod` m)
      _ -> MaybeT $ return Nothing
    Proj' ts i e | isVectorTy' ts -> return $ Proj' ts i (VecFloorMod' (genericLength ts) e m)
    Proj' ts i e
      | isMatrixTy' ts ->
        let (h, w) = fromJust (sizeOfMatrixTy (TupleTy ts))
         in return $ Proj' ts i (MatFloorMod' (toInteger h) (toInteger w) e m)
    Map' t1 t2 f xs -> do
      f <- MaybeT $ putFloorMod (Modulo m) f
      return $ Map' t1 t2 f xs
    Foldl' t1 t2 f init xs -> do
      f <- MaybeT $ putFloorMod (Modulo m) f
      return $ Foldl' t1 t2 f init xs
    Lam x t body -> do
      -- TODO: rename only if required
      y <- lift $ genVarName x
      body <- lift $ substitute x (Var y) body
      body <- MaybeT $ putFloorMod (Modulo m) body
      return $ Lam y t body
    e@(App _ _) -> case curryApp e of
      (f@(Lam _ _ _), args) -> do
        f <- MaybeT $ putFloorMod (Modulo m) f
        return $ uncurryApp f args
      (Tuple' ts, es) | isVectorTy' ts -> do
        es' <- lift $ mapM (putFloorMod (Modulo m)) es
        if all isNothing es'
          then MaybeT $ return Nothing
          else return $ uncurryApp (Tuple' ts) (zipWith fromMaybe es es')
      (Tuple' ts, es) | isMatrixTy (TupleTy ts) -> do
        es' <- lift $ mapM (putFloorMod (Modulo m)) es
        if all isNothing es'
          then MaybeT $ return Nothing
          else return $ uncurryApp (Tuple' ts) (zipWith fromMaybe es es')
      _ -> MaybeT $ return Nothing
    _ -> MaybeT $ return Nothing

putFloorModGeneric :: MonadAlpha m => (Expr -> Modulo -> m Expr) -> Modulo -> Expr -> m Expr
putFloorModGeneric fallback m e =
  if e `isModulo` m
    then return e
    else do
      e' <- putFloorMod m e
      case e' of
        Just e' -> return e'
        Nothing -> fallback e m

putMapFloorMod :: MonadAlpha m => Modulo -> Expr -> m Expr
putMapFloorMod = putFloorModGeneric fallback
  where
    fallback e (Modulo m) = do
      x <- genVarName'
      return $ Map' IntTy IntTy (Lam x IntTy (FloorMod' (Var x) m)) e

putVecFloorMod :: (MonadError Error m, MonadAlpha m) => [(VarName, Type)] -> Modulo -> Expr -> m Expr
putVecFloorMod env = putFloorModGeneric fallback
  where
    fallback e (Modulo m) = do
      t <- typecheckExpr env e
      case t of
        TupleTy ts -> return $ VecFloorMod' (genericLength ts) e m
        _ -> throwInternalError $ "not a vector: " ++ formatType t

putMatFloorMod :: (MonadError Error m, MonadAlpha m) => [(VarName, Type)] -> Modulo -> Expr -> m Expr
putMatFloorMod env = putFloorModGeneric fallback
  where
    fallback e (Modulo m) = do
      t <- typecheckExpr env e
      case t of
        TupleTy ts@(TupleTy ts' : _) -> return $ MatFloorMod' (genericLength ts) (genericLength ts') e m
        _ -> throwInternalError $ "not a matrix: " ++ formatType t

rule :: (MonadAlpha m, MonadError Error m) => RewriteRule m
rule =
  let go0 :: Expr -> Maybe Expr
      go0 e = do
        e' <- formatBottomModuloExpr <$> parseModuloExpr e
        guard $ e' /= e
        return e'
      go1 m f (t1, e1) = Just <$> (f <$> t1 (Modulo m) e1 <*> pure m)
      go2 m f (t1, e1) (t2, e2) = Just <$> (f <$> t1 (Modulo m) e1 <*> t2 (Modulo m) e2 <*> pure m)
   in makeRewriteRule "Jikka.Core.Convert.PropagateMod" $ \env -> \case
        e@(ModNegate' _ _) -> return $ go0 e
        e@(ModPlus' _ _ _) -> return $ go0 e
        e@(ModMinus' _ _ _) -> return $ go0 e
        e@(ModMult' _ _ _) -> return $ go0 e
        e@(ModInv' _ _) -> return $ go0 e
        e@(ModPow' _ _ _) -> return $ go0 e
        ModMatAp' h w e1 e2 m | not (e1 `isModulo'` m) || not (e2 `isModulo'` m) -> go2 m (ModMatAp' h w) (putMatFloorMod (typeEnv env), e1) (putVecFloorMod (typeEnv env), e2)
        ModMatAdd' h w e1 e2 m | not (e1 `isModulo'` m) || not (e2 `isModulo'` m) -> go2 m (ModMatAdd' h w) (putMatFloorMod (typeEnv env), e1) (putMatFloorMod (typeEnv env), e2)
        ModMatMul' h n w e1 e2 m | not (e1 `isModulo'` m) || not (e2 `isModulo'` m) -> go2 m (ModMatMul' h n w) (putMatFloorMod (typeEnv env), e1) (putMatFloorMod (typeEnv env), e2)
        ModMatPow' n e1 e2 m | not (e1 `isModulo'` m) -> go2 m (ModMatPow' n) (putMatFloorMod (typeEnv env), e1) (\_ e -> return e, e2)
        ModSum' e m | not (e `isModulo'` m) -> go1 m ModSum' (putMapFloorMod, e)
        ModProduct' e m | not (e `isModulo'` m) -> go1 m ModProduct' (putMapFloorMod, e)
        FloorMod' e m ->
          if e `isModulo'` m
            then return $ Just e
            else putFloorMod (Modulo m) e
        VecFloorMod' _ e m ->
          if e `isModulo'` m
            then return $ Just e
            else putFloorMod (Modulo m) e
        MatFloorMod' _ _ e m ->
          if e `isModulo'` m
            then return $ Just e
            else putFloorMod (Modulo m) e
        _ -> return Nothing

runProgram :: (MonadAlpha m, MonadError Error m) => Program -> m Program
runProgram = applyRewriteRuleProgram' rule

-- | `run` propagates `FloorMod` to leaves of exprs.
-- For example, this converts the following:
--
-- > mod ((fun x -> x * x + x) y) 1000000007
--
-- to:
--
-- > (fun x -> mod (mod (x * x) 1000000007 + x) 1000000007) y
run :: (MonadAlpha m, MonadError Error m) => Program -> m Program
run prog = wrapError' "Jikka.Core.Convert.PropagateMod" $ do
  precondition $ do
    lint prog
  prog <- runProgram prog
  postcondition $ do
    lint prog
  return prog