{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE PatternSynonyms #-}
{-# LANGUAGE ViewPatterns #-}

module Jikka.Core.Language.AssertedHint
  ( AssertedHint (..),
    pattern EqualHint,
    parseHints,

    -- * Functions using hints
    lowerBoundWithHints,
    upperBoundWithHints,
    isZeroWithHints,
    nullWithHints,
    lengthWithHints,
  )
where

import Control.Monad
import Data.Semigroup
import Jikka.Core.Language.ArithmeticExpr
import Jikka.Core.Language.BuiltinPatterns
import Jikka.Core.Language.Expr

data AssertedHint
  = -- | @NotEqualHint n@ for a integer variable @x@ means @x /= n@
    NotEqualHint ArithmeticExpr
  | -- | @BoundHint l r@ for a integer variable @x@ means @l <= x < r@
    BoundHint (Maybe ArithmeticExpr) (Maybe ArithmeticExpr)
  | -- | @LengthHint t hint@ for a list variable @xs@ means @length xs@ satisfies @hint@
    LengthHint Type AssertedHint
  | -- | @AllHint t hint@ for a list variable @xs@ means that @x@ satisfies @hint@ for all @x@ in @xs@
    AllHint Type AssertedHint
  deriving (Eq, Ord, Show)

pattern EqualHint :: ArithmeticExpr -> AssertedHint
pattern EqualHint l <-
  (\case BoundHint (Just l) (Just r) | incrArithmeticExpr l == r -> Just l; _ -> Nothing -> Just l)
  where
    EqualHint l = BoundHint (Just l) (Just (incrArithmeticExpr l))

parseNilHint :: Expr -> [(VarName, AssertedHint)]
parseNilHint = \case
  -- Jikka.Core.Language.EqualitiySolving makes rhs nil
  Equal' _ (Var xs) (Nil' t) -> [(xs, LengthHint t (BoundHint (Just (integerArithmeticExpr 0)) (Just (integerArithmeticExpr 1))))]
  NotEqual' _ (Var xs) (Nil' t) -> [(xs, LengthHint t (BoundHint (Just (integerArithmeticExpr 1)) Nothing))]
  _ -> []

data Term
  = VarTerm VarName
  | LenTerm Type VarName
  | NegatedTerm Term
  deriving (Eq, Ord, Show, Read)

-- `decomposeArithmeticExpr` makes @a + b + c@ to @[(a, b + c), (b, a + c), (c, a + b)]@ in a convenient way, because Jikka.Core.Language.EqualitiySolving makes @a == b@ to @a - b == 0@.
decomposeArithmeticExpr :: Expr -> [(Term, ArithmeticExpr)]
decomposeArithmeticExpr e =
  let es = splitToSumArithmeticExpr $ parseArithmeticExpr e
   in (`concatMap` [0 .. length es - 1]) $ \i ->
        let e = es !! i
            e' = sumArithmeticExpr $ take i es ++ drop (i + 1) es
         in case formatArithmeticExpr e of
              Var x -> [(VarTerm x, e')]
              Negate' (Var x) -> [(NegatedTerm (VarTerm x), e')]
              Len' t (Var x) -> [(LenTerm t x, e')]
              Negate' (Len' t (Var x)) -> [(NegatedTerm (LenTerm t x), e')]
              _ -> []

parseBoundHint :: Expr -> [(VarName, AssertedHint)]
parseBoundHint = \case
  Equal' IntTy e (LitInt' 0) -> do
    (x, e) <- decomposeArithmeticExpr e
    case x of
      VarTerm x -> return (x, EqualHint (negateArithmeticExpr e))
      NegatedTerm (VarTerm x) -> return (x, EqualHint e)
      LenTerm t x -> return (x, LengthHint t (EqualHint (negateArithmeticExpr e)))
      NegatedTerm (LenTerm t x) -> return (x, LengthHint t (EqualHint e))
      _ -> []
  NotEqual' IntTy e (LitInt' 0) -> do
    (x, e) <- decomposeArithmeticExpr e
    case x of
      VarTerm x -> return (x, NotEqualHint (negateArithmeticExpr e))
      NegatedTerm (VarTerm x) -> return (x, NotEqualHint e)
      LenTerm t x | isZeroArithmeticExpr e -> return (x, LengthHint t (BoundHint (Just (integerArithmeticExpr 1)) Nothing))
      NegatedTerm (LenTerm t x) | isZeroArithmeticExpr e -> return (x, LengthHint t (BoundHint (Just (integerArithmeticExpr 1)) Nothing))
      _ -> []
  LessEqual' IntTy e (LitInt' 0) -> do
    (x, e) <- decomposeArithmeticExpr e
    case x of
      VarTerm x -> return (x, BoundHint Nothing (Just (incrArithmeticExpr (negateArithmeticExpr e))))
      NegatedTerm (VarTerm x) -> return (x, BoundHint (Just e) Nothing)
      LenTerm t x -> return (x, LengthHint t (BoundHint Nothing (Just (incrArithmeticExpr (negateArithmeticExpr e)))))
      NegatedTerm (LenTerm t x) -> return (x, LengthHint t (BoundHint (Just e) Nothing))
      _ -> []
  LessThan' IntTy e (LitInt' 0) -> do
    (x, e) <- decomposeArithmeticExpr e
    case x of
      VarTerm x -> return (x, BoundHint Nothing (Just (negateArithmeticExpr e)))
      NegatedTerm (VarTerm x) -> return (x, BoundHint (Just (incrArithmeticExpr e)) Nothing)
      LenTerm t x -> return (x, LengthHint t (BoundHint Nothing (Just (negateArithmeticExpr e))))
      NegatedTerm (LenTerm t x) -> return (x, LengthHint t (BoundHint (Just (incrArithmeticExpr e)) Nothing))
      _ -> []
  -- Jikka.Core.Language.EqualitiySolving makes rhs 0 and removes GreaterEqual and GreaterThan
  _ -> []

parseAndHint :: Expr -> [(VarName, AssertedHint)]
parseAndHint = \case
  And' e e' -> parseHints e ++ parseHints e'
  _ -> []

parseAllHint :: Expr -> [(VarName, AssertedHint)]
parseAllHint = \case
  All' (Map' t _ (Lam x _ pred) (Var xs)) -> do
    (x', hint) <- parseHints pred
    guard $ x' == x
    return (xs, AllHint t hint)
  _ -> []

parseHints :: Expr -> [(VarName, AssertedHint)]
parseHints e =
  concat
    [ parseNilHint e,
      parseBoundHint e,
      parseAndHint e,
      parseAllHint e
    ]

selectHints :: [(VarName, AssertedHint)] -> VarName -> [AssertedHint]
selectHints hints x = map snd (filter (\hint -> fst hint == x) hints)

maximum' :: Ord a => [Maybe a] -> Maybe a
maximum' = (getMax <$>) . mconcat . map (Max <$>)

minimum' :: Ord a => [Maybe a] -> Maybe a
minimum' = (getMin <$>) . mconcat . map (Min <$>)

lowerBoundWithHints :: [(VarName, AssertedHint)] -> Expr -> Maybe Integer
lowerBoundWithHints hints = \case
  LitInt' n -> Just n
  Negate' e -> do
    e <- upperBoundWithHints hints e
    return $ - (e - 1)
  Plus' e1 e2 -> do
    e1 <- lowerBoundWithHints hints e1
    e2 <- lowerBoundWithHints hints e2
    return $ e1 + e2
  Minus' e1 e2 -> do
    e1 <- lowerBoundWithHints hints e1
    e2 <- upperBoundWithHints hints e2
    return $ e1 - (e2 - 1)
  Mult' e1 e2 -> do
    e1 <- lowerBoundWithHints hints e1
    e2 <- lowerBoundWithHints hints e2
    guard $ e1 >= 0 && e2 >= 0
    return $ e1 * e2
  Var x ->
    let go :: AssertedHint -> Maybe Integer
        go = \case
          BoundHint (Just l) _ -> lowerBoundWithHints hints (formatArithmeticExpr l)
          _ -> Nothing
     in maximum' (map go (selectHints hints x))
  _ -> Nothing

upperBoundWithHints :: [(VarName, AssertedHint)] -> Expr -> Maybe Integer
upperBoundWithHints hints = \case
  LitInt' n -> Just (n + 1)
  Negate' e -> do
    e <- lowerBoundWithHints hints e
    return $ - e + 1
  Plus' e1 e2 -> do
    e1 <- upperBoundWithHints hints e1
    e2 <- upperBoundWithHints hints e2
    return $ e1 + e2 - 1
  Minus' e1 e2 -> do
    e1 <- upperBoundWithHints hints e1
    e2 <- lowerBoundWithHints hints e2
    return $ e1 - e2
  Mult' e1 e2 -> do
    l1 <- lowerBoundWithHints hints e1
    l2 <- lowerBoundWithHints hints e2
    r1 <- upperBoundWithHints hints e1
    r2 <- upperBoundWithHints hints e2
    guard $ l1 >= 0 && l2 >= 0
    return $ (r1 - 1) * (r2 - 1) + 1
  Var x ->
    let go :: AssertedHint -> Maybe Integer
        go = \case
          BoundHint _ (Just r) -> pred <$> upperBoundWithHints hints (formatArithmeticExpr r)
          _ -> Nothing
     in minimum' (map go (selectHints hints x))
  _ -> Nothing

isZeroWithHints :: [(VarName, AssertedHint)] -> Expr -> Maybe Bool
isZeroWithHints hints e =
  case (lowerBoundWithHints hints e, upperBoundWithHints hints e) of
    (Just 0, Just 1) -> Just True
    (Just l, _) | l >= 1 -> Just False
    (_, Just r) | r <= 0 -> Just False
    _ -> Nothing

integerWithHints :: [(VarName, AssertedHint)] -> Expr -> Maybe Integer
integerWithHints hints = \case
  LitInt' n -> Just n
  Negate' e -> do
    e <- integerWithHints hints e
    return $ - e
  Plus' e1 e2 -> do
    e1 <- integerWithHints hints e1
    e2 <- integerWithHints hints e2
    return $ e1 + e2
  Minus' e1 e2 -> do
    e1 <- integerWithHints hints e1
    e2 <- integerWithHints hints e2
    return $ e1 - e2
  Mult' e1 e2 -> do
    e1 <- integerWithHints hints e1
    e2 <- integerWithHints hints e2
    return $ e1 * e2
  e@(Var _) ->
    case (lowerBoundWithHints hints e, upperBoundWithHints hints e) of
      (Just l, Just r) | l + 1 == r -> Just l
      _ -> Nothing
  _ -> Nothing

nullWithHints :: [(VarName, AssertedHint)] -> Expr -> Maybe Bool
nullWithHints hints = \case
  Nil' _ -> Just True
  Cons' _ _ _ -> Just False
  Range1' n -> isZeroWithHints hints n
  Var xs -> do
    let go :: AssertedHint -> Maybe Integer
        go = \case
          LengthHint _ (BoundHint (Just l) _) -> lowerBoundWithHints hints (formatArithmeticExpr l)
          _ -> Nothing
    l <- maximum' (map go (selectHints hints xs))
    if l >= 1
      then Just False
      else do
        let go :: AssertedHint -> Maybe Integer
            go = \case
              LengthHint _ (NotEqualHint n) -> integerWithHints hints (formatArithmeticExpr n)
              _ -> Nothing
        let ns = map go (selectHints hints xs)
        if Just 0 `elem` ns
          then Just False
          else do
            let go :: AssertedHint -> Maybe Integer
                go = \case
                  LengthHint _ (BoundHint _ (Just r)) -> pred <$> upperBoundWithHints hints (formatArithmeticExpr r)
                  _ -> Nothing
            r <- minimum' (map go (selectHints hints xs))
            if r <= 1
              then Just True
              else Nothing
  _ -> Nothing

lengthWithHints :: [(VarName, AssertedHint)] -> Expr -> Maybe Integer
lengthWithHints hints = \case
  Nil' _ -> Just 0
  Cons' _ _ xs -> do
    n <- lengthWithHints hints xs
    return $ n + 1
  Range1' e -> case (lowerBoundWithHints hints e, upperBoundWithHints hints e) of
    (Just l, Just r) | l + 1 == r -> Just l
    _ -> Nothing
  Var xs -> do
    let go :: AssertedHint -> Maybe Integer
        go = \case
          LengthHint _ (BoundHint (Just l) _) -> lowerBoundWithHints hints (formatArithmeticExpr l)
          _ -> Nothing
    let minimum' :: Ord a => [Maybe a] -> Maybe a
        minimum' = (getMin <$>) . mconcat . map (Min <$>)
    minimum' (map go (selectHints hints xs))
  _ -> Nothing