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

module Jikka.Core.Language.ModuloExpr
  ( -- * Basic functions
    ModuloExpr,
    parseModuloExpr,
    formatTopModuloExpr,
    formatBottomModuloExpr,
    integerModuloExpr,
    negateModuloExpr,
    plusModuloExpr,
    minusModuloExpr,
    multModuloExpr,
    isZeroModuloExpr,
    isOneModuloExpr,
    moduloOfModuloExpr,
    arithmeticExprFromModuloExpr,

    -- * Utilities
    Modulo (..),
    formatModulo,
    isModulo,
  )
where

import Data.List
import Jikka.Common.Error
import Jikka.Core.Format (formatExpr)
import Jikka.Core.Language.ArithmeticExpr
import Jikka.Core.Language.BuiltinPatterns
import Jikka.Core.Language.Expr
import Jikka.Core.Language.Runtime (modinv, modpow)
import Jikka.Core.Language.Util

-- | `Modulo` is just a newtype to avoid mistakes that swapping left and right of mod-op.
newtype Modulo = Modulo {unModulo :: Expr}
  deriving (Eq, Ord, Show, Read)

formatModulo :: Modulo -> String
formatModulo = formatExpr . unModulo

isModulo :: Expr -> Modulo -> Bool
isModulo e (Modulo m) = case e of
  FloorMod' _ m' -> m' == m
  ModNegate' _ m' -> m' == m
  ModPlus' _ _ m' -> m' == m
  ModMinus' _ _ m' -> m' == m
  ModMult' _ _ m' -> m' == m
  ModInv' _ m' -> m' == m
  ModPow' _ _ m' -> m' == m
  VecFloorMod' _ _ m' -> m' == m
  MatFloorMod' _ _ _ m' -> m' == m
  ModMatAp' _ _ _ _ m' -> m' == m
  ModMatAdd' _ _ _ _ m' -> m' == m
  ModMatMul' _ _ _ _ _ m' -> m' == m
  ModMatPow' _ _ _ m' -> m' == m
  ModSum' _ m' -> m' == m
  ModProduct' _ m' -> m' == m
  LitInt' n -> case m of
    LitInt' m -> 0 <= n && n < m
    _ -> False
  Proj' ts _ e | isVectorTy' ts -> e `isModulo` Modulo m
  Proj' ts _ e | isMatrixTy' ts -> e `isModulo` Modulo m
  Map' _ _ f _ -> f `isModulo` Modulo m
  Lam _ _ body -> body `isModulo` Modulo m
  e@(App _ _) -> case curryApp e of
    (e@(Lam _ _ _), _) -> e `isModulo` Modulo m
    (Tuple' ts, es) | isVectorTy' ts -> all (`isModulo` Modulo m) es
    (Tuple' ts, es) | isMatrixTy' ts -> all (`isModulo` Modulo m) es
    _ -> False
  _ -> False

data ProductExpr = ProductExpr
  { productExprConst :: Integer,
    productExprList :: [Expr],
    productExprInvList :: [Expr]
  }
  deriving (Eq, Ord, Show, Read)

data SumExpr = SumExpr
  { sumExprList :: [ProductExpr],
    sumExprConst :: Integer
  }
  deriving (Eq, Ord, Show, Read)

data ModuloExpr = ModuloExpr
  { unModuloExpr :: SumExpr,
    modulo :: Modulo
  }
  deriving (Show)

instance Eq ModuloExpr where
  e1 == e2 = unModuloExpr (normalizeModuloExpr e1) == unModuloExpr (normalizeModuloExpr e2) && modulo e1 == modulo e2

instance Ord ModuloExpr where
  e1 `compare` e2 = (unModuloExpr (normalizeModuloExpr e1), modulo e1) `compare` (unModuloExpr (normalizeModuloExpr e2), modulo e2)

negateProductExpr :: ProductExpr -> ProductExpr
negateProductExpr e = e {productExprConst = negate (productExprConst e)}

multProductExpr :: ProductExpr -> ProductExpr -> ProductExpr
multProductExpr e1 e2 =
  ProductExpr
    { productExprConst = productExprConst e1 * productExprConst e2,
      productExprList = productExprList e1 ++ productExprList e2,
      productExprInvList = productExprInvList e1 ++ productExprInvList e2
    }

powProductExpr :: Integer -> ProductExpr -> Integer -> ProductExpr
powProductExpr m e k =
  ProductExpr
    { productExprConst = fromSuccess $ modpow (productExprConst e) k m,
      productExprList = map (\e -> ModPow' e (LitInt' k) (LitInt' m)) (productExprList e),
      productExprInvList = map (\e -> ModPow' e (LitInt' k) (LitInt' m)) (productExprInvList e)
    }

partitionMaybe :: (a -> Maybe b) -> [a] -> ([b], [a])
partitionMaybe f = \case
  [] -> ([], [])
  x : xs ->
    let (ys, xs') = partitionMaybe f xs
     in case f x of
          Just y -> (y : ys, xs')
          Nothing -> (ys, x : xs')

fromLitInt :: Expr -> Maybe Integer
fromLitInt (LitInt' k) = Just k
fromLitInt _ = Nothing

invProductExpr :: Modulo -> ProductExpr -> ProductExpr
invProductExpr m e =
  let (invKs, invList) = partitionMaybe fromLitInt (productExprInvList e)
      e' =
        ProductExpr
          { productExprConst = product invKs,
            productExprList = LitInt' (productExprConst e) : invList,
            productExprInvList = productExprList e
          }
   in case m of
        Modulo (LitInt' m) -> case modinv (productExprConst e) m of
          Right k ->
            ProductExpr
              { productExprConst = (k * product invKs) `mod` m,
                productExprList = invList,
                productExprInvList = productExprList e
              }
          _ -> e'
        _ -> e'

isInteger :: Modulo -> Bool
isInteger (Modulo (LitInt' _)) = True
isInteger _ = False

moduloToInteger :: Modulo -> Integer
moduloToInteger (Modulo (LitInt' m)) = m
moduloToInteger m = error $ "Jikka.Core.Language.ModuloExpr.moduloToInteger: modulo is not an integer" ++ formatModulo m

parseProductExpr :: Modulo -> Expr -> ProductExpr
parseProductExpr m = \case
  LitInt' n -> ProductExpr {productExprConst = n, productExprList = [], productExprInvList = []}
  Negate' e -> negateProductExpr (parseProductExpr m e)
  Mult' e1 e2 -> multProductExpr (parseProductExpr m e1) (parseProductExpr m e2)
  JustDiv' e1 e2 -> multProductExpr (parseProductExpr m e1) (invProductExpr m (parseProductExpr m e2))
  Pow' e1 (LitInt' k) | isInteger m -> powProductExpr (moduloToInteger m) (parseProductExpr m e1) k
  ModNegate' e m' | Modulo m' == m -> negateProductExpr (parseProductExpr m e)
  ModMult' e1 e2 m' | Modulo m' == m -> multProductExpr (parseProductExpr m e1) (parseProductExpr m e2)
  ModPow' e1 (LitInt' k) (LitInt' m') | Modulo (LitInt' m') == m -> powProductExpr m' (parseProductExpr m e1) k
  ModInv' e1 m' | Modulo m' == m -> invProductExpr m (parseProductExpr m e1)
  e -> ProductExpr {productExprConst = 1, productExprList = [e], productExprInvList = []}

sumExprFromProductExpr :: ProductExpr -> SumExpr
sumExprFromProductExpr e =
  SumExpr
    { sumExprList = [e],
      sumExprConst = 0
    }

integerSumExpr :: Integer -> SumExpr
integerSumExpr n =
  SumExpr
    { sumExprConst = n,
      sumExprList = []
    }

integerModuloExpr :: Modulo -> Integer -> ModuloExpr
integerModuloExpr m k = ModuloExpr (integerSumExpr k) m

negateSumExpr :: SumExpr -> SumExpr
negateSumExpr e =
  SumExpr
    { sumExprList = map negateProductExpr (sumExprList e),
      sumExprConst = negate (sumExprConst e)
    }

plusSumExpr :: SumExpr -> SumExpr -> SumExpr
plusSumExpr e1 e2 =
  SumExpr
    { sumExprList = sumExprList e1 ++ sumExprList e2,
      sumExprConst = sumExprConst e1 + sumExprConst e2
    }

multSumExpr :: Modulo -> SumExpr -> SumExpr -> SumExpr
multSumExpr m e1 e2 =
  SumExpr
    { sumExprList =
        let es1 = parseProductExpr m (LitInt' (sumExprConst e1)) : sumExprList e1
            es2 = parseProductExpr m (LitInt' (sumExprConst e2)) : sumExprList e2
         in tail $ map (uncurry multProductExpr) ((,) <$> es1 <*> es2),
      sumExprConst = sumExprConst e1 * sumExprConst e2
    }

negateModuloExpr :: ModuloExpr -> ModuloExpr
negateModuloExpr (ModuloExpr e m) = ModuloExpr (negateSumExpr e) m

plusModuloExpr :: ModuloExpr -> ModuloExpr -> Maybe ModuloExpr
plusModuloExpr (ModuloExpr e1 m) (ModuloExpr e2 m') | m == m' = Just $ ModuloExpr (plusSumExpr e1 e2) m
plusModuloExpr _ _ = Nothing

minusModuloExpr :: ModuloExpr -> ModuloExpr -> Maybe ModuloExpr
minusModuloExpr (ModuloExpr e1 m) (ModuloExpr e2 m') | m == m' = Just $ ModuloExpr (plusSumExpr e1 (negateSumExpr e2)) m
minusModuloExpr _ _ = Nothing

multModuloExpr :: ModuloExpr -> ModuloExpr -> Maybe ModuloExpr
multModuloExpr (ModuloExpr e1 m) (ModuloExpr e2 m') | m == m' = Just $ ModuloExpr (multSumExpr m e1 e2) m
multModuloExpr _ _ = Nothing

parseSumExpr :: Modulo -> Expr -> SumExpr
parseSumExpr m = \case
  LitInt' n -> SumExpr {sumExprList = [], sumExprConst = n}
  Negate' e -> negateSumExpr (parseSumExpr m e)
  Plus' e1 e2 -> plusSumExpr (parseSumExpr m e1) (parseSumExpr m e2)
  Minus' e1 e2 -> plusSumExpr (parseSumExpr m e1) (negateSumExpr (parseSumExpr m e2))
  Mult' e1 e2 -> multSumExpr m (parseSumExpr m e1) (parseSumExpr m e2)
  FloorMod' e m' | Modulo m' == m -> parseSumExpr m e
  ModNegate' e m' | Modulo m' == m -> negateSumExpr (parseSumExpr m e)
  ModPlus' e1 e2 m' | Modulo m' == m -> plusSumExpr (parseSumExpr m e1) (parseSumExpr m e2)
  ModMinus' e1 e2 m' | Modulo m' == m -> plusSumExpr (parseSumExpr m e1) (negateSumExpr (parseSumExpr m e2))
  ModMult' e1 e2 m' | Modulo m' == m -> multSumExpr m (parseSumExpr m e1) (parseSumExpr m e2)
  e -> sumExprFromProductExpr (parseProductExpr m e)

getModuloFromExpr :: Expr -> Maybe Modulo
getModuloFromExpr = \case
  FloorMod' _ m -> Just $ Modulo m
  ModNegate' _ m -> Just $ Modulo m
  ModPlus' _ _ m -> Just $ Modulo m
  ModMinus' _ _ m -> Just $ Modulo m
  ModMult' _ _ m -> Just $ Modulo m
  _ -> Nothing

-- | `parseModuloExpr` converts a given expr to a normal form \(\sum_i \prod_j e _ {i,j}) \bmod m\).
-- This assumes given exprs have the type \(\mathbf{int}\).
parseModuloExpr :: Expr -> Maybe ModuloExpr
parseModuloExpr e = do
  m <- getModuloFromExpr e
  return $ ModuloExpr (parseSumExpr m e) m

formatTopProductExpr :: Modulo -> ProductExpr -> Expr
formatTopProductExpr m e =
  let k = LitInt' (productExprConst e)
      k' e' = case productExprConst e of
        0 -> LitInt' 0
        1 -> e'
        -1 -> Negate' e'
        _ -> Mult' e' k
      invList = map (`FloorMod'` unModulo m) (productExprInvList e)
   in case productExprList e ++ invList of
        [] -> k
        eHead : esTail -> k' (foldl Mult' eHead esTail)

formatTopSumExpr :: Modulo -> SumExpr -> Expr
formatTopSumExpr m e = case sumExprList e of
  [] -> LitInt' (sumExprConst e)
  eHead : esTail ->
    let op e'
          | productExprConst e' > 0 = Plus'
          | productExprConst e' < 0 = Minus'
          | otherwise = const
        go e1 e2 = op e2 e1 (formatTopProductExpr m (e2 {productExprConst = abs (productExprConst e2)}))
        k' e'
          | sumExprConst e > 0 = Plus' e' (LitInt' (sumExprConst e))
          | sumExprConst e < 0 = Minus' e' (LitInt' (abs (sumExprConst e)))
          | otherwise = e'
     in k' (foldl go (formatTopProductExpr m eHead) esTail)

-- | `formatTopModuloExpr` convert `ModuloExpr` to `Expr` with adding `FloorMod` at only the root.
formatTopModuloExpr :: ModuloExpr -> Expr
formatTopModuloExpr e = FloorMod' (formatTopSumExpr (modulo e) . unModuloExpr $ normalizeModuloExpr e) (unModulo $ modulo e)

formatBottomInteger :: Modulo -> Integer -> Expr
formatBottomInteger m k = case unModulo m of
  LitInt' m -> LitInt' (k `mod` m)
  m -> FloorMod' (LitInt' k) m

formatBottomProductExpr :: Modulo -> ProductExpr -> Expr
formatBottomProductExpr m e =
  let k = formatBottomInteger m (productExprConst e)
      k' e' = case productExprConst e of
        0 -> LitInt' 0
        1 -> e'
        -1 -> ModNegate' e' (unModulo m)
        _ -> ModMult' e' k (unModulo m)
      invList = map (`ModInv'` unModulo m) (productExprInvList e)
      list = map (\e -> if e `isModulo` m then e else FloorMod' e (unModulo m)) (productExprList e)
   in case list ++ invList of
        [] -> k
        eHead : esTail -> k' (foldl (\e1 e2 -> ModMult' e1 e2 (unModulo m)) eHead esTail)

formatBottomSumExpr :: Modulo -> SumExpr -> Expr
formatBottomSumExpr m e = case sumExprList e of
  [] -> formatBottomInteger m (sumExprConst e)
  eHead : esTail ->
    let go e1 e2
          | productExprConst e2 == 0 = e1
          | productExprConst e2 == -1 = ModMinus' e1 (formatBottomProductExpr m (e2 {productExprConst = 1})) (unModulo m)
          | otherwise = ModPlus' e1 (formatBottomProductExpr m e2) (unModulo m)
        plusConst e'
          | sumExprConst e == 0 = e'
          | otherwise = ModPlus' e' (formatBottomInteger m (sumExprConst e)) (unModulo m)
     in plusConst (foldl go (formatBottomProductExpr m eHead) esTail)

-- | `formatBottomModuloExpr` convert `ModuloExpr` to `Expr` with adding `FloorMod` at every nodes.
formatBottomModuloExpr :: ModuloExpr -> Expr
formatBottomModuloExpr e = formatBottomSumExpr (modulo e) . unModuloExpr $ normalizeModuloExpr e

normalizeProductExpr :: Modulo -> ProductExpr -> ProductExpr
normalizeProductExpr m e =
  let k = case m of
        Modulo (LitInt' m) -> productExprConst e `mod` m
        _ -> productExprConst e
      es =
        if k == 0
          then []
          else sort (productExprList e)
      es' =
        if k == 0
          then []
          else sort (productExprInvList e)
   in e
        { productExprList = es,
          productExprInvList = es',
          productExprConst = k
        }

normalizeSumExpr :: Modulo -> SumExpr -> SumExpr
normalizeSumExpr m e =
  let f e = (productExprList e, productExprInvList e)
      cmp e1 e2 = f e1 `compare` f e2
      cmp' e1 e2 = cmp e1 e2 == EQ
      es = sortBy cmp (map (normalizeProductExpr m) (sumExprList e))
      es' = groupBy cmp' es
      es'' =
        map
          ( \group ->
              ProductExpr
                { productExprConst = sum (map productExprConst group),
                  productExprList = productExprList (head group),
                  productExprInvList = productExprInvList (head group)
                }
          )
          es'
      es''' = filter (\e -> productExprConst e /= 0 && not (null (productExprList e))) es''
      k = sum (map (\e -> if null (productExprList e) then productExprConst e else 0) es'')
      k' = sumExprConst e + k
      k'' = case m of
        Modulo (LitInt' m) -> k' `mod` m
        _ -> k'
   in SumExpr
        { sumExprList = es''',
          sumExprConst = k''
        }

normalizeModuloExpr :: ModuloExpr -> ModuloExpr
normalizeModuloExpr e = ModuloExpr (normalizeSumExpr (modulo e) (unModuloExpr e)) (modulo e)

isZeroModuloExpr :: ModuloExpr -> Bool
isZeroModuloExpr e = normalizeModuloExpr e == integerModuloExpr (modulo e) 0

isOneModuloExpr :: ModuloExpr -> Bool
isOneModuloExpr e = normalizeModuloExpr e == integerModuloExpr (modulo e) 1

moduloOfModuloExpr :: ModuloExpr -> Expr
moduloOfModuloExpr = unModulo . modulo

arithmeticExprFromModuloExpr :: ModuloExpr -> ArithmeticExpr
arithmeticExprFromModuloExpr e = parseArithmeticExpr . formatTopSumExpr (modulo e) $ unModuloExpr e