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

-- |
-- Module      : Jikka.Core.Convert.SortAbs
-- Description : remove abs with sorting. / sort によって abs を除去します。
-- Copyright   : (c) Kimiyuki Onaka, 2021
-- License     : Apache License 2.0
-- Maintainer  : kimiyuki95@gmail.com
-- Stability   : experimental
-- Portability : portable
--
-- \[
--     \newcommand\int{\mathbf{int}}
--     \newcommand\bool{\mathbf{bool}}
--     \newcommand\list{\mathbf{list}}
-- \]
module Jikka.Core.Convert.SortAbs
  ( run,

    -- * internal rules
    rule,
  )
where

import Control.Monad.Trans.Maybe
import Jikka.Common.Alpha
import Jikka.Common.Error
import qualified Jikka.Core.Convert.Alpha as Alpha
import Jikka.Core.Language.ArithmeticExpr
import Jikka.Core.Language.Beta
import Jikka.Core.Language.BuiltinPatterns
import Jikka.Core.Language.Expr
import Jikka.Core.Language.Lint
import Jikka.Core.Language.QuasiRules
import Jikka.Core.Language.RewriteRules
import Jikka.Core.Language.Util

-- | @replaceAbsDelta x y z e@ replaces \(\levert x - y \rvert\) in \(e\) with \(z\).
replaceAbsDelta :: VarName -> VarName -> VarName -> Expr -> Expr
replaceAbsDelta x y z e = mapSubExpr go [] e
  where
    go _ = \case
      Abs' e | isZeroArithmeticExpr (parseArithmeticExpr (Minus' e (Minus' (Var x) (Var y)))) -> Var z
      Abs' e | isZeroArithmeticExpr (parseArithmeticExpr (Minus' e (Minus' (Var y) (Var x)))) -> Var z
      e -> e

swapTwoVars :: MonadAlpha m => VarName -> VarName -> Expr -> m Expr
swapTwoVars x y e = do
  x' <- genVarName x
  y' <- genVarName y
  e <- substitute x (Var x') e
  e <- substitute y (Var y') e
  e <- substitute x' (Var y) e
  substitute y' (Var x) e

-- | TODO: accept more functions
isSymmetric :: MonadAlpha m => VarName -> VarName -> Expr -> m Bool
isSymmetric x y f = do
  g <- swapTwoVars x y f
  return $ parseArithmeticExpr g == parseArithmeticExpr f

rule :: (MonadAlpha m, MonadError Error m) => RewriteRule m
rule = makeRewriteRule "sum/sum/abs/symmetric" $ \env -> \case
  Sum' (Map' IntTy _ (Lam x _ (Sum' (Map' _ _ (Lam y _ f) xs'))) xs) | xs' == xs -> runMaybeT $ do
    delta <- lift genVarName'
    let f' = replaceAbsDelta x y delta f
    guard $ f' /= f -- f has |x - y|
    guard =<< lift (isSymmetric x y f') -- symmetric
    ys <- lift $ genVarName'' xs
    i <- lift genVarName'
    j <- lift genVarName'
    lt <- lift $ substitute delta (Minus' (Var x) (Var y)) f'
    eq <- lift $ substitute delta (LitInt' 0) f'
    gt <- lift $ substitute delta (Minus' (Var y) (Var x)) f'
    let ctx = Let y IntTy (At' IntTy (Var ys) (Var j))
    let lt' = Sum' (Map' IntTy IntTy (Lam j IntTy (ctx lt)) (Range1' (Var i)))
    let eq' = Let j IntTy (Var i) (ctx eq)
    let gt' = Sum' (Map' IntTy IntTy (Lam j IntTy (ctx gt)) (Range2' (Plus' (Var i) (LitInt' 1)) (Len' IntTy (Var ys))))
    let e =
          Let ys (ListTy IntTy) (Sorted' IntTy xs) $
            Sum'
              ( Map'
                  IntTy
                  IntTy
                  ( Lam
                      i
                      IntTy
                      ( Let
                          x
                          IntTy
                          (At' IntTy (Var ys) (Var i))
                          (Plus' (Plus' lt' eq') gt')
                      )
                  )
                  (Range1' (Len' IntTy (Var ys)))
              )
    lift $ Alpha.runExpr (typeEnv env) e
  _ -> return Nothing

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

-- | `run` reduces \(\lvert \sum _ {a_i \in a} \sum _ {a_j \in a} f(a, a_i, a_j) \rvert\) to \(\mathbf{let}~ b = \mathrm{sort}(a) ~\mathbf{in}~ \sum \sum f'(a, a_i, a_j)\) when \(f\) contains \(\lvert a_i - a_j \rvert\) and \(f(a, a_i, a_j) = f(a, a_j, a_i)\) holds.
--
-- == Example
--
-- Before:
--
-- > sum (map (fun (a_i: int) ->
-- >     sum (map (fun (a_j: int) ->
-- >         abs (a_i - a_j)
-- >     ) a)
-- > ) a)
--
-- After:
--
-- > let b = sort a
-- > in sum (map (fun (i: int) ->
-- >     (sum (map (fun (b_j: int) ->
-- >         b_i - b_j
-- >     ) b[:i])
-- >     + 0
-- >     + sum (map (fun (b_j: int) ->
-- >         b_j - b_i
-- >     ) b[i + 1:]))
-- > ) (range (length b)))
run :: (MonadAlpha m, MonadError Error m) => Program -> m Program
run prog = wrapError' "Jikka.Core.Convert.SortAbs" $ do
  precondition $ do
    lint prog
  prog <- runProgram prog
  postcondition $ do
    lint prog
  return prog