# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License.  You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied.  See the License for the
# specific language governing permissions and limitations
# under the License.
# pylint: disable=invalid-name
"""Arithmetic data structure and utility"""

import enum
from contextlib import contextmanager

import tvm_ffi

from tvm import ir, tirx
from tvm.arith import IntSet
from tvm.runtime import Object

from . import _ffi_api


class ProofStrength(enum.IntEnum):
    """Proof strength of the analysis"""

    DEFAULT = 0
    SYMBOLIC_BOUND = 1


class CompareResult(enum.IntEnum):
    """Result of a transitive comparison.

    Values must match the C++ ``arith::CompareResult`` enum.
    """

    INCONSISTENT = 0
    EQ = 1
    LT = 2
    LE = 3
    GT = 4
    GE = 5
    NE = 6
    UNKNOWN = 7


class Extension(enum.Flag):
    """Extensions enabled for RewriteSimplifier

    Values should match `RewriteSimplifier::Extensions`
    """

    NoExtensions = 0
    TransitivelyProveInequalities = 1 << 0
    ConvertBooleanToAndOfOrs = 1 << 1
    ApplyConstraintsToBooleanBranches = 1 << 2
    ComparisonOfProductAndSum = 1 << 3


@tvm_ffi.register_object("arith.ModularSet")
class ModularSet(Object):
    """Represent range of (coeff * x + base) for x in Z"""

    def __init__(self, coeff, base):
        self.__init_handle_by_constructor__(_ffi_api.ModularSet, coeff, base)


@tvm_ffi.register_object("arith.ConstIntBound")
class ConstIntBound(Object):
    """Represent constant integer bound

    Parameters
    ----------
    min_value : int
        The minimum value of the bound.

    max_value : int
        The maximum value of the bound.
    """

    POS_INF = (1 << 63) - 1
    NEG_INF = -POS_INF

    def __init__(self, min_value, max_value):
        self.__init_handle_by_constructor__(_ffi_api.ConstIntBound, min_value, max_value)


class ConstraintScope:
    """Constraint scope.

    Parameters
    ----------
    fenter : function
        A function that will be called to create an enter context.

    Note
    ----
    Do not create object directly, use Analyzer.constraint_scope
    """

    def __init__(self, fenter):
        self._fenter = fenter
        self._fexit = None

    def __enter__(self):
        self._fexit = self._fenter()

    def __exit__(self, ptype, value, trace):
        self._fexit()


class Z3ContextScope:
    """Share one fresh Z3 context across Analyzers created in this scope.

    The outermost scope creates the context. Nested scopes on the same thread
    reuse it, and exiting the outermost scope releases its ownership.
    """

    def __enter__(self):
        _ffi_api.EnterZ3ContextScope()
        return self

    def __exit__(self, ptype, value, trace):
        _ffi_api.ExitZ3ContextScope()


@tvm_ffi.register_object("arith.Analyzer")
class Analyzer(Object):
    """Integer arithmetic analyzer

    This is a stateful analyzer class that can be used to perform
    various symbolic integer analysis. The same analyzer instance can
    be passed to FFI APIs to share accumulated facts across calls.
    """

    def __init__(self):
        self.__init_handle_by_constructor__(_ffi_api.Analyzer)

    @property
    def is_z3_enabled(self) -> bool:
        """Whether this build includes the Z3 backend (``USE_Z3=ON``).

        The Z3-specific methods (:py:meth:`get_smtlib2`, :py:meth:`get_z3_stats`,
        :py:meth:`set_z3_timeout_ms`, :py:meth:`set_z3_rlimit`) only work when
        this is ``True``.
        """
        return bool(_ffi_api.AnalyzerIsZ3Enabled(self))

    def _check_z3_enabled(self) -> None:
        if not self.is_z3_enabled:
            raise RuntimeError(
                "The Z3 backend is not available in this build. "
                "Rebuild TVM with USE_Z3=ON to use Z3-specific Analyzer APIs."
            )

    def get_smtlib2(self, expr: tirx.Expr | None = None) -> str:
        """Get the current Z3 problem in SMT-LIB2 format.

        Raises
        ------
        RuntimeError
            If TVM was built without Z3 (``USE_Z3=OFF``), since there is no
            solver state to export. Use :py:attr:`is_z3_enabled` to check first.

        Parameters
        ----------
        expr : Optional[Expr]
            The expression to prove. If provided, its negation is added to the problem.
        """
        self._check_z3_enabled()
        return _ffi_api.AnalyzerGetSMTLIB2(self, expr)

    def set_z3_timeout_ms(self, timeout_ms: int) -> None:
        """Set Z3 timeout in milliseconds.

        Raises
        ------
        RuntimeError
            If TVM was built without Z3 (``USE_Z3=OFF``).

        Parameters
        ----------
        timeout_ms : int
            The timeout in milliseconds.
        """
        self._check_z3_enabled()
        _ffi_api.AnalyzerSetZ3TimeoutMs(self, timeout_ms)

    def set_z3_rlimit(self, rlimit: int) -> None:
        """Set Z3 resource limit.

        The resource limit gives deterministic solver budgeting (unlike a wall
        clock timeout). A value of ``0`` disables the limit.

        Raises
        ------
        RuntimeError
            If TVM was built without Z3 (``USE_Z3=OFF``).

        Parameters
        ----------
        rlimit : int
            The resource limit.
        """
        self._check_z3_enabled()
        _ffi_api.AnalyzerSetZ3RLimit(self, rlimit)

    def get_z3_stats(self) -> str:
        """Get Z3 solver statistics.

        Raises
        ------
        RuntimeError
            If TVM was built without Z3 (``USE_Z3=OFF``).

        Returns
        -------
        stats : str
            The Z3 statistics.
        """
        self._check_z3_enabled()
        return _ffi_api.AnalyzerGetZ3Stats(self)

    def const_int_bound(self, expr: tirx.Expr) -> ConstIntBound:
        """Find constant integer bound for expr.

        Parameters
        ----------
        expr : Expr
            The expression.

        Returns
        -------
        bound : ConstIntBound
            The result bound
        """
        return _ffi_api.AnalyzerConstIntBound(self, expr)

    def const_int_bound_is_bound(self, var: tirx.Var) -> bool:
        """Check if a variable is bound to a range.

        Parameters
        ----------
        var : tvm.tirx.Var
            The variable.

        Returns
        -------
        result : bool
            Whether the variable is bound to a range.
        """
        return _ffi_api.AnalyzerConstIntBoundIsBound(self, var)

    def modular_set(self, expr: tirx.Expr) -> ModularSet:
        """Find a modular set that expr belongs to.

        Parameters
        ----------
        expr : Expr
            The expression.

        Returns
        -------
        result : ModularSet
            The result.
        """
        return _ffi_api.AnalyzerModularSet(self, expr)

    def simplify(self, expr: tirx.Expr, steps: int = 2) -> tirx.Expr:
        """Simplify expression via both rewrite and canonicalization.

        Parameters
        ----------
        expr : Expr
            The expression.
        steps : The simplification runs in the order of
                rewrite_simplify (step 1) -> canonical_simplify (step 2) ->
                rewrite_simplify (step 3) -> canonical_simplify (step 4) -> ...
                param steps controls how many steps to run.
                Default is 2, i.e., rewrite_simplify + canonical_simplify.

        Returns
        -------
        result : Expr
            The result.
        """
        return _ffi_api.AnalyzerSimplify(self, expr, steps)

    def clone(self) -> "Analyzer":
        """Return a deep copy of this analyzer with independent state.

        The returned analyzer carries the same accumulated facts (variable
        bounds, modular sets, bindings, integer-set domains, literal
        constraints and transitive comparisons) as this one, but owns its own
        state: binding or simplifying on either analyzer afterwards does not
        affect the other. Unlike copying the handle, this is a true deep copy.

        Do not call this while a constraint scope is active on this analyzer.

        Returns
        -------
        result : Analyzer
            A new analyzer holding an independent copy of the facts.
        """
        return _ffi_api.AnalyzerClone(self)

    def rewrite_simplify(self, expr: tirx.Expr) -> tirx.Expr:
        """Simplify expression via rewriting rules.

        Parameters
        ----------
        expr : Expr
            The expression.

        Returns
        -------
        result : Expr
            The result.
        """
        return _ffi_api.AnalyzerRewriteSimplify(self, expr)

    @property
    def rewrite_simplify_stats(self):
        return _ffi_api.AnalyzerGetRewriteSimplifyStats(self)

    def reset_rewrite_simplify_stats(self):
        _ffi_api.AnalyzerResetRewriteSimplifyStats(self)

    def canonical_simplify(self, expr: tirx.Expr) -> tirx.Expr:
        """Simplify expression via canonicalization.

        Parameters
        ----------
        expr : Expr
            The expression.

        Returns
        -------
        result : Expr
            The result.
        """
        return _ffi_api.AnalyzerCanonicalSimplify(self, expr)

    def int_set(self, expr: tirx.Expr, dom_map: dict[tirx.Var, IntSet] | None = None) -> IntSet:
        """Compute a symbolic IntSet that covers expr for all values in dom_map.

        Parameters
        ----------
        expr : Expr
            The expression.

        dom_map : Optional[Dict[tvm.tirx.Var, tvm.arith.IntSet]]
            The domain for variables to be relaxed.  When omitted, the analyzer
            uses the domains of the variables already bound to it.

        Returns
        -------
        result : IntSet
            The result.
        """
        return _ffi_api.AnalyzerIntSet(self, expr, dom_map)

    def can_prove(self, expr: tirx.Expr, strength: ProofStrength = ProofStrength.DEFAULT) -> bool:
        """Check whether we can prove expr to be true.

        Parameters
        ----------
        expr : Expr
            The expression.

        strength: ProofStrength
            The proof strength. When TVM is built with Z3 (``USE_Z3=ON``), the
            optional Z3 fallback is only consulted at ``SYMBOLIC_BOUND`` or
            higher, after the native analyzers fail to prove the predicate.

        Returns
        -------
        result : Expr
            The result.
        """
        return _ffi_api.AnalyzerCanProve(self, expr, strength)

    def set_maximum_rewrite_steps(self, maximum: int) -> None:
        """Set the maximum allowed number of rewrite-simplify steps.

        When a positive limit is set, the simplifier raises an exception once
        it exceeds that number of rewrite steps.  This is useful for guarding
        against performance regressions in tests.

        Parameters
        ----------
        maximum : int
            The maximum number of rewrite steps, or a non-positive value to
            allow an unlimited number of steps.
        """
        _ffi_api.AnalyzerSetMaximumRewriteSteps(self, maximum)

    def bind(
        self,
        var: tirx.Var,
        expr: tirx.Expr | ir.Range,
        allow_override: bool = False,
    ) -> None:
        """Bind a variable to the expression.

        Parameters
        ----------
        var : tvm.tirx.Var
            The variable.

        expr : Union[tirx.Expr, ir.Range]
            The expression or the range to bind to.

        allow_override : bool
            Whether to allow overriding an existing binding for the variable.
        """
        return _ffi_api.AnalyzerBind(self, var, expr, allow_override)

    def constraint_scope(self, constraint: tirx.Expr) -> ConstraintScope:
        """Create a constraint scope.

        Parameters
        ----------
        constraint : Expr
            The constraint expression.

        returns
        -------
        scope : ConstraintScope
            The constraint scope

        Examples
        --------
        .. code-block:: python

          x = te.var("x")
          analyzer = tvm.arith.Analyzer()
          with analyzer.constraint_scope(x % 3 == 0):
              # constraint in effect
              assert analyzer.modular_set(x).coeff == 3
          # constraint no longer in effect
          assert analyzer.modular_set(x).coeff != 3
        """

        def _fenter():
            return _ffi_api.AnalyzerEnterConstraintContext(self, constraint)

        return ConstraintScope(_fenter)

    def update(
        self, var: tirx.Var, info: ConstIntBound | ModularSet | IntSet, override: bool = False
    ) -> None:
        """Update information about var.

        Parameters
        ----------
        var : tvm.tirx.Var
            The variable.

        info : Union[ConstIntBound, ModularSet, IntSet]
            Related information.  A ``ConstIntBound`` updates the constant
            integer bound, a ``ModularSet`` updates the modular set, and an
            ``IntSet`` updates the integer-set domain of ``var``.

        override : bool
            Whether allow override.
        """
        if isinstance(info, ConstIntBound):
            _ffi_api.AnalyzerConstIntBoundUpdate(self, var, info, override)
        elif isinstance(info, ModularSet):
            _ffi_api.AnalyzerModularSetUpdate(self, var, info, override)
        elif isinstance(info, IntSet):
            _ffi_api.AnalyzerIntSetUpdate(self, var, info, override)
        else:
            raise TypeError(f"Do not know how to handle type {type(info)}")

    def can_prove_equal(self, lhs: tirx.Expr, rhs: tirx.Expr) -> bool:
        """Whether we can prove that lhs == rhs

        Parameters
        ----------
        lhs: Expr
            The left-hand side of the comparison

        rhs: Expr
            The right-hand side of the comparison

        Returns
        -------
        result: bool
            Whether we can prove that lhs == rhs
        """
        return _ffi_api.AnalyzerCanProveEqual(self, lhs, rhs)

    def try_compare(
        self, lhs: tirx.Expr, rhs: tirx.Expr, propagate_inequalities: bool = True
    ) -> CompareResult:
        """Compare lhs and rhs using previously provided known comparisons.

        Parameters
        ----------
        lhs : Expr
            The left-hand side of the comparison.

        rhs : Expr
            The right-hand side of the comparison.

        propagate_inequalities : bool
            If true, attempt to find a sequence of transitive inequalities that
            allow lhs and rhs to be compared.

        Returns
        -------
        result : CompareResult
            The most specific result that can be proven about the comparison.
            Returns ``CompareResult.UNKNOWN`` when nothing can be proven.
        """
        return CompareResult(_ffi_api.AnalyzerTryCompare(self, lhs, rhs, propagate_inequalities))

    @property
    def enabled_extensions(self) -> Extension:
        """Return the currently enabled extensions"""
        value = _ffi_api.AnalyzerGetEnabledExtensions(self)
        return Extension(value)

    @enabled_extensions.setter
    def enabled_extensions(self, flags: int | Extension):
        """Enable extensions for the analyzer

        Parameters
        ----------
        flags: Union[int,Extension]

            The extensions to enable.
        """
        flags = Extension(flags).value
        _ffi_api.AnalyzerSetEnabledExtensions(self, flags)


@contextmanager
def allow_uint_as_index():
    """Opt-in no-overflow domain for unsigned index arithmetic.

    Within the scope, the analyzer treats uint32/uint64 scalars as index
    types: it assumes their values never approach the type limit (no
    wraparound), so the signed rewrite rule set (floordiv/floormod
    decomposition, modular analysis) applies to them. Every unsigned
    expression admitted under the assumption is logged once per process
    (WARNING) so the assumption stays auditable.

    Off by default. Intended for compile-time layout proofs (offsets that
    provably stay small), NOT for runtime codegen semantics where unsigned
    wraparound is real behavior.
    """
    _ffi_api.EnterAllowUintAsIndex()
    try:
        yield
    finally:
        _ffi_api.ExitAllowUintAsIndex()
