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

module Jikka.Core.Language.ArithmeticExpr
  ( -- * Basic functions
    ArithmeticExpr,
    parseArithmeticExpr,
    formatArithmeticExpr,
    integerArithmeticExpr,
    negateArithmeticExpr,
    plusArithmeticExpr,
    minusArithmeticExpr,
    multArithmeticExpr,
    incrArithmeticExpr,
    decrArithmeticExpr,
    sumArithmeticExpr,
    isZeroArithmeticExpr,
    isOneArithmeticExpr,
    isIntegerArithmeticExpr,
    integerFromArithmeticExpr,

    -- * Advanced functions
    unNPlusKPattern,
    makeVectorFromArithmeticExpr,
    makeAffineFunctionFromArithmeticExpr,
    splitConstantFactorArithmeticExpr,
    splitToSumArithmeticExpr,
  )
where

import Control.Arrow
import Control.Monad
import Control.Monad.ST
import Control.Monad.Trans
import Control.Monad.Trans.Maybe
import Data.List (findIndices, groupBy, sort, sortBy)
import Data.STRef
import qualified Data.Vector as V
import qualified Data.Vector.Mutable as MV
import Jikka.Core.Language.BuiltinPatterns
import Jikka.Core.Language.Expr
import Jikka.Core.Language.FreeVars

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

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

newtype ArithmeticExpr = ArithmeticExpr {unArithmeticExpr :: SumExpr}
  deriving (Show)

instance Eq ArithmeticExpr where
  e1 == e2 = unArithmeticExpr (normalizeArithmeticExpr e1) == unArithmeticExpr (normalizeArithmeticExpr e2)

instance Ord ArithmeticExpr where
  e1 `compare` e2 = unArithmeticExpr (normalizeArithmeticExpr e1) `compare` unArithmeticExpr (normalizeArithmeticExpr e2)

integerProductExpr :: Integer -> ProductExpr
integerProductExpr n =
  ProductExpr
    { productExprConst = n,
      productExprList = []
    }

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
    }

iterateN :: Integer -> (a -> a) -> a -> a
iterateN n _ _ | n < 0 = error $ "iterateN: negative number: " ++ show n
iterateN 0 _ x = x
iterateN n f x = iterateN (n - 1) f (f x)

parseProductExpr :: Expr -> ProductExpr
parseProductExpr = \case
  LitInt' n -> ProductExpr {productExprConst = n, productExprList = []}
  Negate' e -> negateProductExpr (parseProductExpr e)
  Mult' e1 e2 -> multProductExpr (parseProductExpr e1) (parseProductExpr e2)
  Pow' e1 (LitInt' k) | 0 <= k && k < 10 -> iterateN k (multProductExpr (parseProductExpr e1)) (integerProductExpr 1)
  e -> ProductExpr {productExprConst = 1, productExprList = [e]}

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

arithmeticalExprFromProductExpr :: ProductExpr -> ArithmeticExpr
arithmeticalExprFromProductExpr = ArithmeticExpr . sumExprFromProductExpr

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

integerArithmeticExpr :: Integer -> ArithmeticExpr
integerArithmeticExpr = ArithmeticExpr . integerSumExpr

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

negateArithmeticExpr :: ArithmeticExpr -> ArithmeticExpr
negateArithmeticExpr (ArithmeticExpr e) = ArithmeticExpr $ negateSumExpr e

plusArithmeticExpr :: ArithmeticExpr -> ArithmeticExpr -> ArithmeticExpr
plusArithmeticExpr (ArithmeticExpr e1) (ArithmeticExpr e2) = ArithmeticExpr $ plusSumExpr e1 e2

minusArithmeticExpr :: ArithmeticExpr -> ArithmeticExpr -> ArithmeticExpr
minusArithmeticExpr (ArithmeticExpr e1) (ArithmeticExpr e2) = ArithmeticExpr $ plusSumExpr e1 (negateSumExpr e2)

multArithmeticExpr :: ArithmeticExpr -> ArithmeticExpr -> ArithmeticExpr
multArithmeticExpr (ArithmeticExpr e1) (ArithmeticExpr e2) = ArithmeticExpr $ multSumExpr e1 e2

incrArithmeticExpr :: ArithmeticExpr -> ArithmeticExpr
incrArithmeticExpr = plusArithmeticExpr (integerArithmeticExpr 1)

decrArithmeticExpr :: ArithmeticExpr -> ArithmeticExpr
decrArithmeticExpr = plusArithmeticExpr (integerArithmeticExpr (-1))

sumArithmeticExpr :: [ArithmeticExpr] -> ArithmeticExpr
sumArithmeticExpr = foldl plusArithmeticExpr (integerArithmeticExpr 0)

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

-- | `parseArithmeticExpr` converts a given expr to a normal form \(\sum_i \prod_j e _ {i,j})\).
-- This assumes given exprs have the type \(\mathbf{int}\).
parseArithmeticExpr :: Expr -> ArithmeticExpr
parseArithmeticExpr = ArithmeticExpr . parseSumExpr

formatProductExpr :: ProductExpr -> Expr
formatProductExpr e =
  let k = LitInt' (productExprConst e)
      k' e' = case productExprConst e of
        0 -> LitInt' 0
        1 -> e'
        -1 -> Negate' e'
        _ -> Mult' e' k
   in case productExprList e of
        [] -> k
        eHead : esTail -> k' (foldl Mult' eHead esTail)

formatSumExpr :: SumExpr -> Expr
formatSumExpr 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 (formatProductExpr (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 (formatProductExpr eHead) esTail)

formatArithmeticExpr :: ArithmeticExpr -> Expr
formatArithmeticExpr = formatSumExpr . unArithmeticExpr . normalizeArithmeticExpr

normalizeProductExpr :: ProductExpr -> ProductExpr
normalizeProductExpr e =
  let es =
        if productExprConst e == 0
          then []
          else sort (productExprList e)
   in e {productExprList = es}

normalizeSumExpr :: SumExpr -> SumExpr
normalizeSumExpr e =
  let cmp e1 e2 = productExprList e1 `compare` productExprList e2
      cmp' e1 e2 = cmp e1 e2 == EQ
      es = sortBy cmp (map normalizeProductExpr (sumExprList e))
      es' = groupBy cmp' es
      es'' = map (\group -> ProductExpr {productExprConst = sum (map productExprConst group), productExprList = productExprList (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'')
   in SumExpr
        { sumExprList = es''',
          sumExprConst = sumExprConst e + k
        }

normalizeArithmeticExpr :: ArithmeticExpr -> ArithmeticExpr
normalizeArithmeticExpr = ArithmeticExpr . normalizeSumExpr . unArithmeticExpr

-- | `makeVectorFromArithmeticExpr` makes a vector \(f\) and a expr \(c\) from a given vector of variables \(x_0, x_1, \dots, x _ {n - 1}\) and a given expr \(e\) s.t. \(f\) and \(c\) don't have \(x_0, x_1, \dots, x _ {n - 1}\) as free variables and \(e = c + f \cdot (x_0, x_1, \dots, x _ {n - 1})\) holds.
-- This assumes given variables and exprs have the type \(\mathbf{int}\).
--
-- * The returned exprs are normalized with `normalizeArithmeticExpr`.
makeVectorFromArithmeticExpr :: V.Vector VarName -> ArithmeticExpr -> Maybe (V.Vector ArithmeticExpr, ArithmeticExpr)
makeVectorFromArithmeticExpr xs es = runST $ do
  runMaybeT $ do
    f <- lift $ MV.replicate (V.length xs) (integerArithmeticExpr 0)
    c <- lift $ newSTRef (integerArithmeticExpr (sumExprConst (unArithmeticExpr es)))
    forM_ (sumExprList (unArithmeticExpr es)) $ \e -> do
      let indices = V.imap (\i x -> map (i,) (findIndices (x `isFreeVar`) (productExprList e))) xs
      case concat (V.toList indices) of
        [] -> lift $ modifySTRef c (plusArithmeticExpr (arithmeticalExprFromProductExpr e))
        [(i, j)] ->
          if productExprList e !! j == Var (xs V.! i)
            then do
              let e' = e {productExprList = take j (productExprList e) ++ drop (j + 1) (productExprList e)}
              lift $ MV.modify f (plusArithmeticExpr (arithmeticalExprFromProductExpr e')) i -- k x_i
            else MaybeT $ return Nothing -- e.g. f(x_i)
        _ -> MaybeT $ return Nothing -- e.g. x_1 x_2
    f <- V.freeze f
    c <- lift $ readSTRef c
    return (V.map normalizeArithmeticExpr f, normalizeArithmeticExpr c)

isZeroArithmeticExpr :: ArithmeticExpr -> Bool
isZeroArithmeticExpr e = normalizeArithmeticExpr e == integerArithmeticExpr 0

isOneArithmeticExpr :: ArithmeticExpr -> Bool
isOneArithmeticExpr e = normalizeArithmeticExpr e == integerArithmeticExpr 1

isIntegerArithmeticExpr :: ArithmeticExpr -> Bool
isIntegerArithmeticExpr e = normalizeArithmeticExpr e == integerArithmeticExpr (sumExprConst (unArithmeticExpr e))

integerFromArithmeticExpr :: ArithmeticExpr -> Maybe Integer
integerFromArithmeticExpr e
  | isIntegerArithmeticExpr e = Just (sumExprConst (unArithmeticExpr e))
  | otherwise = Nothing

-- | `unNPlusKPattern` recognizes a pattern of \(x + k\) for a variable \(x\) and an integer constant \(k \in \mathbb{Z}\).
unNPlusKPattern :: ArithmeticExpr -> Maybe (VarName, Integer)
unNPlusKPattern e = case normalizeArithmeticExpr e of
  ArithmeticExpr
    SumExpr
      { sumExprList =
          [ ProductExpr
              { productExprConst = 1,
                productExprList = [Var x]
              }
            ],
        sumExprConst = k
      } -> Just (x, k)
  _ -> Nothing

-- | `makeAffineFunctionFromArithmeticExpr` is a specialized version of `makeVectorFromArithmeticExpr`.
-- This function returns \(a, b\) for a given variable \(x\) and a given expr \(e = a x + b\) where \(a, b\) which doesn't use \(x\) free.
makeAffineFunctionFromArithmeticExpr :: VarName -> ArithmeticExpr -> Maybe (ArithmeticExpr, ArithmeticExpr)
makeAffineFunctionFromArithmeticExpr x es = first V.head <$> makeVectorFromArithmeticExpr (V.singleton x) es

-- | `splitConstantFactorArithmeticExpr` finds \(k\) and \(e'\) for given \(e\) s.t. \(e = k e'\).
splitConstantFactorArithmeticExpr :: ArithmeticExpr -> (Integer, ArithmeticExpr)
splitConstantFactorArithmeticExpr e =
  let e' = unArithmeticExpr $ normalizeArithmeticExpr e
   in case (sumExprConst e', sumExprList e') of
        (0, []) -> (0, integerArithmeticExpr 0)
        (k, []) -> (k, integerArithmeticExpr 1)
        (0, [e]) -> second arithmeticalExprFromProductExpr $ splitConstantFactorProductExpr e
        (k, es) ->
          let kes = map splitConstantFactorProductExpr es
              d = foldl gcd k (map fst kes)
           in ( d,
                ArithmeticExpr
                  SumExpr
                    { sumExprConst = k `div` d,
                      sumExprList = map (\(k, e) -> e {productExprConst = (k * productExprConst e) `div` d}) kes
                    }
              )

splitConstantFactorProductExpr :: ProductExpr -> (Integer, ProductExpr)
splitConstantFactorProductExpr e = (productExprConst e, e {productExprConst = 1})

splitToSumArithmeticExpr :: ArithmeticExpr -> [ArithmeticExpr]
splitToSumArithmeticExpr e =
  let e' = unArithmeticExpr $ normalizeArithmeticExpr e
      es = map arithmeticalExprFromProductExpr (sumExprList e')
      k = if sumExprConst e' == 0 then [] else [integerArithmeticExpr (sumExprConst e')]
   in es ++ k