import unittest

import onnx_light.onnx_core.expressions as onnx_expressions
from onnx_light.ext_test_case import ExtTestCase
from onnx_light.onnx_core.expressions import (
    simplify_expression,
    simplify_two_expressions,
    compare_expressions,
    CompareResult,
    DimRange,
    evaluate_expression,
    parse_expression_tokens,
    rename_expression,
    rename_dynamic_expression,
    dim_add,
    dim_sub,
    dim_mul,
    dim_multi_mul,
    dim_div,
    dim_exact_div,
    dim_mod,
    dim_max,
    dim_min,
    dim_ranges_from_expressions,
)

# ─────────────────────────────────────────────────────────────────────────────
# test_simplify_expressions.py (adapted)
# ─────────────────────────────────────────────────────────────────────────────


class TestSimplifyExpressions(ExtTestCase):
    def test_import_new_module_path(self):
        self.assertEqual(onnx_expressions.simplify_expression("x + x"), "2*x")

    def test_simplify_expression(self):
        self.assertEqual(simplify_expression("x - y + y"), "x")
        self.assertEqual(simplify_expression("2*x + 3*x - x"), "4*x")
        self.assertEqual(simplify_expression("a + b - a"), "b")
        self.assertEqual(simplify_expression("5 + x - 2 + 3"), "x+6")
        self.assertEqual(simplify_expression("x - x"), "0")
        self.assertEqual(simplify_expression("Max(x,y)"), "x^y")
        self.assertEqual(str(simplify_expression("2*batch//batch")), "2")

    def test_simplify_expression_to_int(self):
        self.assertEqual(simplify_expression("2*batch//batch"), 2)
        self.assertEqual(simplify_expression("2^2"), 2)
        self.assertEqual(simplify_expression("2*batch//batch^2"), 2)

    def test_simplify_expression2(self):
        self.assertEqual(simplify_expression("5 + x - (2 + 3)"), "x")

    def test_simplify_expression3(self):
        self.assertEqual(simplify_expression("x - 1"), "x-1")
        self.assertEqual(simplify_expression("1 - x"), "-x+1")

    def test_simplify_expression4(self):
        self.assertEqual(simplify_expression("-1+1+length"), "length")
        self.assertEqual(simplify_expression("1+length-1"), "length")
        self.assertEqual(simplify_expression("-1+length+1"), "length")

    def test_simplify_two_expressions(self):
        self.assertEqual(
            simplify_two_expressions("s52+seq_length", "s52+s70"), {"s70": -1, "seq_length": 1}
        )
        self.assertEqual(simplify_two_expressions("e*2", "e+e"), {})

    def test_compare_expressions(self):
        def check(expr1, expr2, result, difference):
            c = compare_expressions(expr1, expr2)
            self.assertEqual(c.result, result)
            self.assertEqual(c.difference, difference)

        check("a+b", "b+a", CompareResult.Equal, "0")
        check("5", "3", CompareResult.Greater, -2)
        check("3", "5", CompareResult.Smaller, 2)
        check("a+1", "a", CompareResult.Greater, "-1")
        check("a", "a+1", CompareResult.Smaller, "1")
        check("a+b+1", "a", CompareResult.Greater, "-b-1")
        check("a", "b", CompareResult.Unknown, "b-a")
        check("2*a", "a", CompareResult.Unknown, "-a")

    def test_compare_expressions_syntax_error(self):
        with self.assertRaises(RuntimeError):
            compare_expressions("a", "b +")

    def test_simplify_expression_bracket(self):
        self.assertEqual("x", simplify_expression("2*x//2"))
        self.assertEqual("x", simplify_expression("(2*x)//2"))
        self.assertEqual("x", simplify_expression("(x*y)//y"))
        self.assertEqual("x", simplify_expression("(x*(y+1))//(y+1)"))
        self.assertEqual("c//2", simplify_expression("((c)//(2))"))

    def test_simplify_expression_bracket_max(self):
        self.assertEqual("x^1+y", simplify_expression("(x)^(y+1)"))
        self.assertEqual("1+x^y", simplify_expression("(x+1)^(y)"))

    def test_simplify_add_sub(self):
        self.assertEqual("b+c", simplify_expression("b+c-CeilToInt(b+c,2)+CeilToInt(b+c,2)"))

    def test_simplify_floordiv_add_ring(self):
        # sum_{i=0..n-1} floor((y + i) / n) == y (for integer y)
        self.assertEqual("b+c", simplify_expression("(1+b+c)//2+(b+c)//2"))
        self.assertEqual("a", simplify_expression("a//2+(a+1)//2"))
        self.assertEqual("x", simplify_expression("x//3+(x+1)//3+(x+2)//3"))
        self.assertEqual("x+5", simplify_expression("(x+5)//3+(x+6)//3+(x+7)//3"))
        self.assertEqual("a+b", simplify_expression("a//2+(a+1)//2+b"))
        self.assertEqual("b+c", simplify_expression("CeilToInt(b+c, 2)+(b+c)//2"))
        # Insufficient terms / non-contiguous offsets must be preserved.
        self.assertEqual("(1+x)//3+x//3", simplify_expression("x//3+(x+1)//3"))
        self.assertEqual("(2+a)//2+a//2", simplify_expression("a//2+(a+2)//2"))

    def test_simplify_function(self):
        self.assertEqual("(1+b+c)//2", simplify_expression("CeilToInt(b+c,2)"))

    def test_simplify_function_order(self):
        self.assertEqual("a+b", simplify_expression("b+a"))

    def test_simplify_function_order3(self):
        self.assertEqual("a+b+c", simplify_expression("c+b+a"))
        self.assertEqual("a+b+c", simplify_expression("b+c+a"))
        self.assertEqual("a+b+c", simplify_expression("a+c+b"))

    def test_simplify_function_floordiv_int(self):
        self.assertEqual("512*a", simplify_expression("1024*a//2"))
        self.assertEqual("a", simplify_expression("1024*a//1024"))
        self.assertEqual("a+b", simplify_expression("1024*(a+b)//1024"))
        self.assertEqual("2*a+2*b", simplify_expression("1024*(a+b)//1024*2"))

    def test_simplify_floordiv_not_exact(self):
        # `//` is floor division, so a factor cannot cross the division
        # boundary: a*(x//a) == x only when x is a multiple of a, whereas
        # (a*x)//a == x always holds.
        self.assertEqual("2*(H//2)", simplify_expression("2*(H//2)"))
        self.assertEqual("H", simplify_expression("(2*H)//2"))
        self.assertEqual("3*(H//3)", simplify_expression("3*(H//3)"))
        self.assertEqual("H", simplify_expression("(3*H)//3"))
        self.assertEqual("2*(n//2)+1", simplify_expression("2*(n//2)+1"))
        self.assertEqual("(a-1)//2", simplify_expression("(a-3)//2+1"))
        # Concrete counter-example: 2*(3//2) == 2, not 3.
        self.assertEqual(2, evaluate_expression("2*(H//2)", {"H": 3}))
        self.assertEqual(3, evaluate_expression("(2*H)//2", {"H": 3}))
        self.assertEqual(-1, evaluate_expression("-1//2", {}))
        self.assertEqual(0, evaluate_expression("(a-3)//2+1", {"a": 2}))
        self.assertEqual(0, evaluate_expression("(a-1)//2", {"a": 2}))

    def test_simplify_function_floordiv_distribute(self):
        self.assertEqual("b+c", simplify_expression("(2*b+2*c)//2"))
        self.assertEqual("2*b+c", simplify_expression("(4*b+2*c)//2"))
        self.assertEqual("a-b", simplify_expression("(2*a-2*b)//2"))
        # Constant residual smaller than the divisor: 1//2 == 0, so the
        # residual vanishes and the result is just the distributed quotient.
        self.assertEqual("b+c", simplify_expression("(2*b+2*c+1)//2"))
        self.assertEqual("b+c", simplify_expression("(1+2*b+2*c)//2"))
        # Constant residual equal to or greater than the divisor folds into
        # an additive integer constant.
        self.assertEqual("b+c+1", simplify_expression("(2*b+2*c+3)//2"))
        self.assertEqual("sequence//5-2", simplify_expression("(sequence-10)//5"))
        self.assertEqual("sequence//5+2", simplify_expression("(sequence+10)//5"))

    def test_simplify_nested_floordiv(self):
        # (x // a) // b  →  x // (a * b)
        self.assertEqual("x//10", simplify_expression("x//5//2"))
        self.assertEqual("x//6", simplify_expression("x//2//3"))
        self.assertEqual("x//16", simplify_expression("x//4//4"))
        # Verify numerically for positive and negative values.
        self.assertEqual(
            evaluate_expression("x//10", {"x": 37}), evaluate_expression("x//5//2", {"x": 37})
        )
        self.assertEqual(
            evaluate_expression("x//10", {"x": -37}), evaluate_expression("x//5//2", {"x": -37})
        )
        # Should not simplify when a divisor is not a constant.
        # // is left-associative so (x//a)//2 prints without redundant parens.
        self.assertEqual("x//a//2", simplify_expression("(x//a)//2"))
        self.assertEqual("x//2//b", simplify_expression("(x//2)//b"))

    def test_simplify_expression_negation(self):
        self.assertEqual("length", simplify_expression("-1+1+length"))
        self.assertEqual("x-1", simplify_expression("-1+x"))
        self.assertEqual("x", simplify_expression("-x+2*x"))

    def test_simplify_like_terms_multi_symbolic_factor(self):
        # When an integer constant is not a direct child of the Mult node
        # (e.g. 4096*a*b parses as Mult(Mult(4096,a),b)), the like-term
        # collector must still extract the coefficient and combine terms
        # with the same symbolic product.
        self.assertEqual("4104*a*b", simplify_expression("4096*a*b+8*a*b"))
        self.assertEqual("8*x*y", simplify_expression("3*x*y+5*x*y"))
        self.assertEqual("5*a*b*c", simplify_expression("2*a*b*c+3*a*b*c"))
        # Full expression from the issue.
        expr = (
            "(past_sequence_length+sequence_length)"
            "+16384*batch_size*past_sequence_length"
            "+4096*batch_size*sequence_length"
            "+8*batch_size*sequence_length"
            "+8*batch_size*total_sequence_length"
            "+748180122"
        )
        result = simplify_expression(expr)
        self.assertEqual(
            "16384*batch_size*past_sequence_length"
            "+4104*batch_size*sequence_length"
            "+8*batch_size*total_sequence_length"
            "+past_sequence_length+sequence_length+748180122",
            result,
        )


# ─────────────────────────────────────────────────────────────────────────────
# test_evaluate_expressions.py (adapted)
# ─────────────────────────────────────────────────────────────────────────────


class TestEvaluateExpressions(ExtTestCase):
    def test_evaluate_expression(self):
        self.assertEqual(-1, evaluate_expression("x - y", {"x": 5, "y": 6}))
        self.assertEqual(-5, evaluate_expression("- x", {"x": 5}))

    def test_evaluate_expression_syntax_error(self):
        with self.assertRaises(RuntimeError):
            evaluate_expression("x +", {"x": 5})


# ─────────────────────────────────────────────────────────────────────────────
# test_rename_expressions.py (adapted)
# ─────────────────────────────────────────────────────────────────────────────


class TestRenameExpressions(ExtTestCase):
    def test_rename_expression2(self):
        self.assertEqual("B+seq_length", rename_expression("s52+seq_length", {"s52": "B"}))

    def test_rename_expression_max(self):
        cst = {
            "A": "s77",
            "s77": "A",
            "B": "s27",
            "s27": "B",
            "D": "s3",
            "s3": "D",
            "E": "s10",
            "s10": "E",
            "E^D": "Max(s10,s3)",
            "Max(s10,s3)": "E^D",
        }
        self.assertEqual("E^D", rename_expression("s10^s3", cst))
        self.assertEqual("E^D", rename_expression("Max(s10,s3)", cst))
        self.assertEqual("E^D", rename_dynamic_expression("Max(s10,s3)", cst))

    def test_rename_expression(self):
        self.assertEqual("B+seq_length", rename_expression("s52+seq_length", {"s52": "B"}))

    def test_parse_expression_tokens_syntax_error(self):
        invalid_expr = "a +"
        result = parse_expression_tokens(invalid_expr)
        self.assertEqual({invalid_expr}, result)

    def test_rename_dynamic_expression_syntax_error(self):
        invalid_expr = "a +"
        result = rename_dynamic_expression(invalid_expr, {"a": "b"})
        self.assertEqual(invalid_expr, result)

    def test_rename_dynamic_expression(self):
        replacements = {
            "DYN0": "DYN0",
            "batch": "batch",
            "cache_length": "cache_length",
            "s0": "batch",
            "s1": "seq_length",
            "s10": "batch",
            "s12": "batch",
            "s14": "batch",
            "s2": "batch",
            "s3": "DYN0",
            "s8": "batch",
            "seq_length": "seq_length",
            "s11": "cache_length",
            "s15": "cache_length",
            "s9": "cache_length",
            "s13": "cache_length",
        }
        expression = "s9+seq_length"
        renamed = rename_dynamic_expression(expression, replacements)
        self.assertEqual(renamed, "cache_length+seq_length")


# ─────────────────────────────────────────────────────────────────────────────
# test_operations.py (adapted)
# ─────────────────────────────────────────────────────────────────────────────


class TestDimOperations(ExtTestCase):
    # ------------------------------------------------------------------ dim_mul
    def test_dim_mul_int_int(self):
        self.assertEqual(dim_mul(3, 4), 12)

    def test_dim_mul_zero(self):
        self.assertEqual(dim_mul(0, 5), 0)

    def test_dim_mul_symbolic(self):
        result = dim_mul("n", 2)
        self.assertIsInstance(result, str)
        self.assertIn("n", result)
        self.assertIn("2", result)

    def test_dim_mul_both_symbolic(self):
        result = dim_mul("a", "b")
        self.assertIsInstance(result, str)
        self.assertIn("a", result)
        self.assertIn("b", result)

    # --------------------------------------------------------------- dim_multi_mul
    def test_dim_multi_mul_all_int(self):
        self.assertEqual(dim_multi_mul(2, 3, 4), 24)

    def test_dim_multi_mul_single_int(self):
        self.assertEqual(dim_multi_mul(7), 7)

    def test_dim_multi_mul_with_symbolic(self):
        result = dim_multi_mul(2, "n", 3)
        self.assertIsInstance(result, str)
        self.assertIn("n", result)

    # ------------------------------------------------------------------ dim_add
    def test_dim_add_int_int(self):
        self.assertEqual(dim_add(3, 4), 7)

    def test_dim_add_symbolic(self):
        result = dim_add("n", 1)
        self.assertIsInstance(result, str)
        self.assertIn("n", result)
        self.assertIn("1", result)

    # ------------------------------------------------------------------ dim_sub
    def test_dim_sub_int_int(self):
        self.assertEqual(dim_sub(10, 3), 7)

    def test_dim_sub_same_symbol(self):
        result = dim_sub("n", "n")
        self.assertEqual(str(result), "0")

    # ------------------------------------------------------------------ dim_div
    def test_dim_div_int_int_exact(self):
        self.assertEqual(dim_div(12, 4), 3)

    def test_dim_div_int_int_floor(self):
        self.assertEqual(dim_div(7, 2), 3)

    def test_dim_div_symbolic(self):
        result = dim_div("2*n", 2)
        self.assertEqual(str(result), "n")

    # ------------------------------------------------------------------ dim_mod
    def test_dim_mod_int_int(self):
        self.assertEqual(dim_mod(10, 3), 1)

    def test_dim_mod_exact(self):
        self.assertEqual(dim_mod(12, 4), 0)

    # ------------------------------------------------------------------ dim_max
    def test_dim_max_int_int(self):
        self.assertEqual(dim_max(7, 3), 7)
        self.assertEqual(dim_max(2, 9), 9)
        self.assertEqual(dim_max(5, 5), 5)

    def test_dim_max_same_symbol(self):
        result = dim_max("n", "n")
        self.assertEqual(str(result), "n")

    # ------------------------------------------------------------------ dim_min
    def test_dim_min_int_int(self):
        self.assertEqual(dim_min(2, 9), 2)
        self.assertEqual(dim_min(8, 3), 3)
        self.assertEqual(dim_min(4, 4), 4)

    # ----------------------------------------------------------- dim_exact_div
    def test_dim_exact_div_int_int(self):
        self.assertEqual(dim_exact_div(12, 4), 3)

    def test_dim_exact_div_int_int_not_exact_raises(self):
        with self.assertRaises(RuntimeError):
            dim_exact_div(7, 2)

    def test_dim_exact_div_symbolic_simplifies(self):
        result = dim_exact_div("2*n", 2)
        self.assertEqual(str(result), "n")

    def test_dim_exact_div_symbolic_mult_outside(self):
        result = dim_exact_div("batch*4", 2)
        self.assertEqual(str(result), "2*batch")


class TestExactDiv(ExtTestCase):
    """Tests for the /: (exact division) operator."""

    def test_parse_unparse(self):
        from onnx_light.onnx_core.expressions import simplify_expression

        # Round-trip through parse/simplify.
        self.assertEqual(simplify_expression("a/:b"), "a/:b")

    def test_simplify_concrete(self):
        self.assertEqual(simplify_expression("12/:4"), 3)
        self.assertEqual(simplify_expression("(a*b)/:(a*b)"), 1)

    def test_simplify_cancel_common_factor(self):
        # Exact division cancels common factors from numerator.
        self.assertEqual(simplify_expression("(2*H)/:2"), "H")
        self.assertEqual(simplify_expression("(3*H)/:3"), "H")
        self.assertEqual(simplify_expression("(batch*seq)/:seq"), "batch")
        self.assertEqual(simplify_expression("(1024*a)/:2"), "512*a")
        self.assertEqual(simplify_expression("(sequence-10)/:5"), "sequence/:5-2")
        self.assertEqual(simplify_expression("(sequence+10)/:5"), "sequence/:5+2")

    def test_simplify_exact_commutes_with_mult(self):
        # Unlike floor division (//), exact division (/: ) commutes with
        # multiplication: c*(a/:b) == (c*a)/:b is valid for exact division,
        # so the simplifier is allowed to cancel the factor even when it
        # appears outside the division.
        self.assertEqual(simplify_expression("2*(H/:2)"), "H")

    def test_simplify_floordiv_still_blocked(self):
        # Floor division still does NOT commute.
        self.assertEqual(simplify_expression("2*(H//2)"), "2*(H//2)")

    def test_evaluate(self):
        self.assertEqual(evaluate_expression("12/:4", {}), 3)
        self.assertEqual(evaluate_expression("(batch*4)/:4", {"batch": 5}), 5)

    def test_evaluate_error_on_remainder(self):
        # /: must be exact: evaluating 7/:2 should raise.
        with self.assertRaises(RuntimeError):
            evaluate_expression("7/:2", {})


# ─────────────────────────────────────────────────────────────────────────────
# test_dim_ranges_from_expressions.py (adapted)
# ─────────────────────────────────────────────────────────────────────────────


class TestDimRangesFromExpressions(ExtTestCase):
    """Tests for :func:`dim_ranges_from_expressions`."""

    # ── Direct equality ────────────────────────────────────────────────────
    def test_direct_equality(self):
        # a == d//5  →  a ∈ [d//5, d//5],  d ∈ [5*a, 5*a+4]
        result = dim_ranges_from_expressions([("a", "d//5")])
        self.assertIsInstance(result["a"], DimRange)
        self.assertEqual(result["a"].lower, "d//5")
        self.assertEqual(result["a"].upper, "d//5")
        self.assertIsInstance(result["d"], DimRange)
        self.assertEqual(simplify_expression(result["d"].lower), "5*a")
        self.assertIn(simplify_expression(result["d"].upper), ("4+5*a", "5*a+4"))

    def test_direct_variable_equality(self):
        # a == b  →  a ∈ [b, b],  b ∈ [a, a]
        result = dim_ranges_from_expressions([("a", "b")])
        self.assertIsInstance(result["a"], DimRange)
        self.assertEqual(result["a"].lower, "b")
        self.assertEqual(result["a"].upper, "b")
        self.assertIsInstance(result["b"], DimRange)
        self.assertEqual(result["b"].lower, "a")
        self.assertEqual(result["b"].upper, "a")

    def test_direct_constant_equality(self):
        # a == 3  →  a ∈ [3, 3]
        result = dim_ranges_from_expressions([("a", "3")])
        self.assertIsInstance(result["a"], DimRange)
        self.assertEqual(result["a"].lower, 3)
        self.assertEqual(result["a"].upper, 3)

    # ── Double floor-division chain ────────────────────────────────────────
    def test_double_floordiv(self):
        # a == d//5//2  ⟹  d//10 == a  →  d ∈ [10*a, 10*a+9],  a ∈ [d//10, d//10]
        result = dim_ranges_from_expressions([("a", "d//5//2")])
        self.assertIsInstance(result["d"], DimRange)
        self.assertEqual(simplify_expression(result["d"].lower), "10*a")
        self.assertIn(simplify_expression(result["d"].upper), ("9+10*a", "10*a+9"))
        self.assertIsInstance(result["a"], DimRange)
        self.assertEqual(simplify_expression(result["a"].lower), "d//10")
        self.assertEqual(simplify_expression(result["a"].upper), "d//10")

    # ── LHS is the floor-div chain ─────────────────────────────────────────
    def test_lhs_floordiv_chain(self):
        # d//5 == a  →  same as a == d//5
        result = dim_ranges_from_expressions([("d//5", "a")])
        self.assertIn("d", result)
        self.assertIsInstance(result["d"], DimRange)
        self.assertEqual(simplify_expression(result["d"].lower), "5*a")

    # ── Multiple equalities ────────────────────────────────────────────────
    def test_multiple_equalities(self):
        # a == d//5  AND  b == 1  →  ranges for a, b, d
        result = dim_ranges_from_expressions([("a", "d//5"), ("b", "1")])
        self.assertIn("a", result)
        self.assertIsInstance(result["b"], DimRange)
        self.assertEqual(result["b"].lower, 1)
        self.assertEqual(result["b"].upper, 1)
        self.assertIn("d", result)

    # ── Token filtering ───────────────────────────────────────────────────
    def test_tokens_filter(self):
        result = dim_ranges_from_expressions([("a", "d//5")], tokens=["d"])
        self.assertIn("d", result)
        self.assertNotIn("a", result)

    def test_tokens_filter_empty_list(self):
        # Empty tokens list means "all tokens".
        result = dim_ranges_from_expressions([("a", "d//5")], tokens=[])
        self.assertIn("a", result)
        self.assertIn("d", result)

    def test_tokens_filter_none(self):
        # None means "all tokens".
        result = dim_ranges_from_expressions([("a", "d//5")], tokens=None)
        self.assertIn("a", result)
        self.assertIn("d", result)

    # ── DimRange repr and equality ─────────────────────────────────────────
    def test_dim_range_repr_symbolic(self):
        result = dim_ranges_from_expressions([("a", "d//5")])
        r = result["a"]
        self.assertEqual(repr(r), "DimRange(lower='d//5', upper='d//5')")

    def test_dim_range_repr_numeric(self):
        result = dim_ranges_from_expressions([("a", "3")])
        r = result["a"]
        self.assertEqual(repr(r), "DimRange(lower=3, upper=3)")

    # ── Unsupported patterns return empty (variable absent) ────────────────
    def test_unsupported_two_variables_on_one_side(self):
        # x + y == 5 — x+y is not a chain of a single variable; skip both sides.
        result = dim_ranges_from_expressions([("x+y", "5")])
        self.assertNotIn("x", result)
        self.assertNotIn("y", result)

    def test_unsupported_symbolic_divisor(self):
        # x//(a*b) == rhs — non-integer divisor; cannot invert.
        result = dim_ranges_from_expressions([("x//(a*b)", "rhs")])
        self.assertNotIn("x", result)

    def test_empty_equalities(self):
        self.assertEqual(dim_ranges_from_expressions([]), {})

    # ── INFINITY constant ──────────────────────────────────────────────────
    def test_infinity_constant_exists(self):
        self.assertEqual(onnx_expressions.INFINITY, "+inf")

    def test_infinity_is_string(self):
        self.assertIsInstance(onnx_expressions.INFINITY, str)


if __name__ == "__main__":
    unittest.main(verbosity=2)
