"""
ComprehensiveMathVerifier — Programmatic answer verification for all math domains.
===================================================================================
Verifies LLM math answers without calling another LLM. Uses SymPy, networkx,
and Python stdlib to check correctness across:
  - Algebra (equations, inequalities, calculus)
  - Number theory (primes, modular arithmetic, divisibility)
  - Combinatorics (counting, permutations, probability)
  - Geometry (area, perimeter, Pythagorean, coordinate)
  - Graph theory (structural properties via networkx)
  - Proof verification (sample-testing claimed identities)
  - Finance (compound interest, NPV, annuity formulas)

Risk output:
  0.10  verified correct  (answer satisfies problem constraints)
  0.50  unverifiable      (problem type not extractable — neutral)
  0.90  verified wrong    (answer does NOT satisfy constraints)

Same score() interface as MathVerifier — but RISK_WRONG=0.90 (not 0.80). NOT a drop-in replacement.
"""
from __future__ import annotations

import math
import re
from typing import List, Optional, Tuple


_CLASSIFY_PATTERNS = [
    ("number_theory",  re.compile(
        r"\b(prime|divisib|factor|gcd|lcm|modulo|\bmod\b|totient|congruent|"
        r"greatest common|least common multiple|euler)\b", re.I)),
    ("combinatorics",  re.compile(
        r"\b(choose|combinations?|permutations?|ways to|arrangements?|"
        r"probability|P\(|C\(|nCr|nPr|how many ways)\b", re.I)),
    ("geometry",       re.compile(
        r"\b(area|perimeter|circumference|triangle|circle|rectangle|square|"
        r"volume|angle|distance|coordinate|slope|midpoint|radius|diameter|"
        r"polygon|trapezoid|parallelogram|sphere|cylinder|cone|cube)\b", re.I)),
    ("graph_theory",   re.compile(
        r"\b(vertex|vertices|edge|edges|\bgraph\b|path|cycle|connected|"
        r"chromatic|degree|spanning|eulerian|hamiltonian)\b", re.I)),
    ("finance",        re.compile(
        r"\b(compound interest|principal|annual rate|present value|future value|"
        r"npv|annuity|amortize|mortgage|investment|return on|roi|apr|"
        r"percentage change|percent change|simple interest)\b", re.I)),
    ("proof",          re.compile(
        r"\b(prove|show that|for all|for every|for any integer|for any n|"
        r"for all n|for all positive)\b", re.I)),
    ("algebra",        re.compile(
        r"\b(equation|solve for|roots|zeros|inequality|polynomial|quadratic|"
        r"absolute value|system of|simplify|expand|factor)\b", re.I)),
]

_SQRT_RE = re.compile(r"^\\sqrt\{([^}]+)\}$")
# Note: nested braces not supported (e.g. \frac{\sqrt{3}}{2} → None). This is intentional.
_FRAC_RE = re.compile(r"^(-?)\\frac\{([^}]+)\}\{([^}]+)\}$")
_THOU_SEP_RE = re.compile(r"^-?\d{1,3}(,\d{3})*(\.\d+)?$")


class ComprehensiveMathVerifier:
    """
    Programmatic math answer verifier covering all major problem categories.
    No LLM calls. Zero new dependencies beyond sympy + networkx (already installed).
    """

    RISK_CORRECT     = 0.10
    RISK_UNKNOWN     = 0.50
    RISK_WRONG       = 0.90

    def is_math_chain(
        self,
        question: str,
        steps: List[dict],
        final_answer: str,
    ) -> bool:
        """
        Return True when this chain is a math/computation task.

        Uses the same classify patterns as score() — if any known math domain
        matches the question, it is considered a math chain. Also returns True
        when the final answer is purely numeric and the question contains digits.

        This method exists so ComprehensiveMathVerifier is a drop-in replacement
        for MathVerifier in AgentGuard.score_chain() which calls is_math_chain()
        before score().
        """
        combined = f"{question} {final_answer}"
        # If any _CLASSIFY_PATTERNS match, it's math
        for _, pat in _CLASSIFY_PATTERNS:
            if pat.search(combined):
                return True
        # Final answer is a bare number AND question has a digit → numeric math
        try:
            float(final_answer.replace(",", "").strip().rstrip("%"))
            if re.search(r"\d", question):
                return True
        except ValueError:
            pass
        return False

    def score(
        self,
        question: str,
        steps: List[dict],
        final_answer: str,
    ) -> float:
        """
        Return risk score for a completed math chain.
        0.10 = verified correct, 0.50 = unverifiable, 0.90 = verified wrong.
        """
        proposed = self._parse_answer(final_answer)
        if proposed is None:
            return self.RISK_UNKNOWN

        problem_type = self._classify(question)
        result = self._dispatch(problem_type, question, proposed)
        return result

    # ── Answer parser ──────────────────────────────────────────────────────

    def _parse_answer(self, text: str) -> Optional[float]:
        """Parse proposed answer to float. Handles LaTeX, fractions, decimals."""
        if not text:
            return None
        s = text.strip()

        # \sqrt{x} — simple numeric sqrt only (e.g. \sqrt{3}); nested LaTeX not supported
        if s.strip().startswith("\\sqrt{"):
            m = _SQRT_RE.match(s.strip())
            if m:
                try:
                    return math.sqrt(float(m.group(1)))
                except (ValueError, TypeError):
                    pass

        # \frac{a}{b} or -\frac{a}{b}
        m = _FRAC_RE.match(s)
        if m:
            try:
                sign = -1.0 if m.group(1) == "-" else 1.0
                return sign * float(m.group(2)) / float(m.group(3))
            except (ValueError, ZeroDivisionError):
                return None

        # Reject non-numeric LaTeX
        if "\\" in s:
            return None

        # Reject tuple/list (4,6,14 — not a thousand separator)
        if "," in s and not _THOU_SEP_RE.match(s):
            return None

        # Plain integer or decimal (including thousand separators like 1,000)
        try:
            return float(s.replace(",", ""))
        except ValueError:
            pass

        return None

    # ── Classifier ─────────────────────────────────────────────────────────

    def _classify(self, question: str) -> str:
        """Return problem type string. Uses regex priority order."""
        for ptype, pattern in _CLASSIFY_PATTERNS:
            if pattern.search(question):
                return ptype
        return "algebra"

    # ── Dispatcher ─────────────────────────────────────────────────────────

    def _dispatch(self, problem_type: str, question: str, proposed: float) -> float:
        """Route to the correct sub-verifier. Return risk float."""
        verifiers = {
            "algebra":       self._verify_algebra,
            "number_theory": self._verify_number_theory,
            "combinatorics": self._verify_combinatorics,
            "geometry":      self._verify_geometry,
            "graph_theory":  self._verify_graph,
            "proof":         self._verify_proof_sampling,
            "finance":       self._verify_finance,
        }
        fn = verifiers.get(problem_type, self._verify_algebra)
        verified, correct = fn(question, proposed)
        if not verified:
            return self.RISK_UNKNOWN
        return self.RISK_CORRECT if correct else self.RISK_WRONG

    # ── Sub-verifiers (stubs — implemented in subsequent tasks) ────────────

    def _verify_algebra(self, q: str, proposed: float) -> Tuple[bool, bool]:
        """
        Extract equation from problem, solve with SymPy, check proposed answer.
        Returns (verified: bool, correct: bool).
        """
        try:
            import sympy
            from sympy import Abs, solve, sympify, symbols
            x = symbols("x", real=True)

            # Clean LaTeX dollars and backslashes
            clean_q = re.sub(r"\$", "", q)
            clean_q = re.sub(r"\\", " ", clean_q)

            # Pattern: "expr1 = expr2" where both sides have x
            eq_match = re.search(
                r"([0-9x\s\|\+\-\*\/\^\(\)\.]+)\s*=\s*([0-9x\s\|\+\-\*\/\^\(\)\.]+)",
                clean_q
            )
            if not eq_match:
                return False, False

            lhs_str = eq_match.group(1).strip()
            rhs_str = eq_match.group(2).strip()

            # Skip if neither side has 'x' (pure arithmetic, handled by MathVerifier)
            if "x" not in lhs_str and "x" not in rhs_str:
                return False, False

            # Replace |expr| with Abs(expr) — must pair opening and closing pipes
            def _pipe_to_abs(s: str) -> str:
                result, depth = "", 0
                i = 0
                while i < len(s):
                    if s[i] == "|":
                        if depth % 2 == 0:
                            result += "Abs("
                            depth += 1
                        else:
                            result += ")"
                            depth -= 1
                    else:
                        result += s[i]
                    i += 1
                return result

            def _prep_expr(s: str) -> str:
                # Insert * for implicit multiplication (e.g. 2x → 2*x)
                s = re.sub(r"(\d)(x)", r"\1*\2", s)
                s = re.sub(r"(x)(\d)", r"\1*\2", s)
                # Replace ^ with ** for exponentiation
                s = s.replace("^", "**")
                return s

            # Build SymPy equation: lhs - rhs = 0
            lhs = sympify(_prep_expr(_pipe_to_abs(lhs_str)), locals={"x": x, "Abs": Abs})
            rhs = sympify(_prep_expr(_pipe_to_abs(rhs_str)), locals={"x": x, "Abs": Abs})
            equation = sympy.Eq(lhs, rhs)

            solutions = solve(equation, x)
            if not solutions:
                return False, False

            # Convert solutions to floats
            sol_floats = []
            for sol in solutions:
                try:
                    sol_floats.append(float(sol.evalf()))
                except Exception:
                    pass
            if not sol_floats:
                return False, False

            # Check if proposed answer matches any solution
            correct = any(self._approx_eq(proposed, s) for s in sol_floats)
            return True, correct

        except Exception:
            return False, False

    def _verify_number_theory(self, q: str, proposed: float) -> Tuple[bool, bool]:
        try:
            import sympy

            # GCD
            m = re.search(r"gcd\s+of\s+(\d+)\s+and\s+(\d+)", q, re.I)
            if m:
                expected = math.gcd(int(m.group(1)), int(m.group(2)))
                return True, self._approx_eq(proposed, float(expected))

            # LCM
            m = re.search(r"lcm\s+of\s+(\d+)\s+and\s+(\d+)", q, re.I)
            if m:
                a, b = int(m.group(1)), int(m.group(2))
                expected = abs(a * b) // math.gcd(a, b)
                return True, self._approx_eq(proposed, float(expected))

            # N mod M
            m = re.search(r"(\d+)\s+mod(?:ulo)?\s+(\d+)", q, re.I)
            if m:
                expected = int(m.group(1)) % int(m.group(2))
                return True, self._approx_eq(proposed, float(expected))

            # Is N prime? → proposed should be 1 (yes) or 0 (no)
            m = re.search(r"is\s+(\d+)\s+(?:a\s+)?prime", q, re.I)
            if m:
                n = int(m.group(1))
                is_prime = sympy.isprime(n)
                expected = 1.0 if is_prime else 0.0
                return True, self._approx_eq(proposed, expected)

            # How many primes less than N?
            m = re.search(r"how many prime\w*\s+(?:numbers?\s+)?(?:are\s+)?less than\s+(\d+)", q, re.I)
            if m:
                n = int(m.group(1))
                expected = float(sympy.primepi(n))
                return True, self._approx_eq(proposed, expected)

            # Euler's totient φ(N)
            m = re.search(r"(?:totient|phi|φ)\s*(?:of\s*)?(?:\()?\s*(\d+)", q, re.I)
            if m:
                expected = float(sympy.totient(int(m.group(1))))
                return True, self._approx_eq(proposed, expected)

            return False, False
        except Exception:
            return False, False

    def _verify_combinatorics(self, q: str, proposed: float) -> Tuple[bool, bool]:
        try:
            # C(n, k): "choose k from n" → group(1)=k, group(2)=n
            m = re.search(r"choose\s+(\d+)(?:\s+items?)?\s+from\s+(\d+)", q, re.I)
            if m:
                k, n = int(m.group(1)), int(m.group(2))
                if k > n:
                    return False, False  # impossible — unverifiable, not wrong
                expected = float(math.comb(n, k))
                return True, self._approx_eq(proposed, expected)

            # "n choose k" → group(1)=n, group(2)=k
            m = re.search(r"(\d+)\s+choose\s+(\d+)", q, re.I)
            if m:
                n, k = int(m.group(1)), int(m.group(2))
                if k > n:
                    return False, False  # impossible — unverifiable, not wrong
                expected = float(math.comb(n, k))
                return True, self._approx_eq(proposed, expected)

            # P(n, k): "permutations of k from n"
            m = re.search(r"permutations?\s+of\s+(\d+)(?:\s+items?)?\s+from\s+(\d+)", q, re.I)
            if m:
                k, n = int(m.group(1)), int(m.group(2))
                expected = float(math.perm(n, k))
                return True, self._approx_eq(proposed, expected)

            # "N distinct objects arranged" → N!
            m = re.search(r"(\d+)\s+distinct\s+\w+\s+(?:be\s+)?arrang", q, re.I)
            if m:
                expected = float(math.factorial(int(m.group(1))))
                return True, self._approx_eq(proposed, expected)

            # Fair coin: P(heads) = 0.5, P(tails) = 0.5
            if re.search(r"fair coin", q, re.I) and re.search(r"probability of (heads|tails)", q, re.I):
                return True, self._approx_eq(proposed, 0.5)

            # Fair die: P(specific face) = 1/6
            if re.search(r"fair (?:six-sided )?die", q, re.I):
                m = re.search(r"probability of (?:rolling a? )?(?:a )?(\d)", q, re.I)
                if m:
                    return True, self._approx_eq(proposed, 1.0 / 6.0)

            return False, False
        except Exception:
            return False, False

    def _verify_geometry(self, q: str, proposed: float) -> Tuple[bool, bool]:
        try:
            # Circle area: π r²
            if re.search(r"area\b.*(circle|disk)", q, re.I) or \
               re.search(r"(circle|disk).*(area)", q, re.I):
                m = re.search(r"radius\s+(?:of\s+)?(\d+\.?\d*)", q, re.I)
                if m:
                    r = float(m.group(1))
                    expected = math.pi * r * r
                    return True, self._approx_eq(proposed, expected)

            # Circle circumference/perimeter: 2πr
            if re.search(r"(circumference|perimeter).*(circle)", q, re.I) or \
               re.search(r"(circle).*(circumference|perimeter)", q, re.I):
                m = re.search(r"radius\s+(?:of\s+)?(\d+\.?\d*)", q, re.I)
                if m:
                    r = float(m.group(1))
                    expected = 2 * math.pi * r
                    return True, self._approx_eq(proposed, expected)

            # Rectangle area: l × w — only for non-negative dimensions
            if re.search(r"area.*rectangle", q, re.I) or re.search(r"rectangle.*area", q, re.I):
                m = re.search(r"length\s+(-?\d+\.?\d*)\s+and\s+width\s+(-?\d+\.?\d*)", q, re.I)
                if not m:
                    m = re.search(r"(-?\d+\.?\d*)\s+by\s+(-?\d+\.?\d*)", q, re.I)
                if m:
                    l, w = float(m.group(1)), float(m.group(2))
                    if l < 0 or w < 0:
                        return False, False  # negative dimensions — unverifiable
                    expected = l * w
                    return True, self._approx_eq(proposed, expected)

            # Pythagorean: legs a, b → hypotenuse c = sqrt(a²+b²)
            if re.search(r"(hypotenuse|right triangle)", q, re.I):
                nums = re.findall(r"(\d+\.?\d*)", q)
                floats = [float(x) for x in nums]
                if len(floats) >= 2:
                    a, b = floats[0], floats[1]
                    c = math.sqrt(a*a + b*b)
                    return True, self._approx_eq(proposed, c)

            # Coordinate distance: points (x1,y1) and (x2,y2)
            m = re.search(
                r"\(\s*(-?\d+\.?\d*)\s*,\s*(-?\d+\.?\d*)\s*\)\s*and\s*\(\s*(-?\d+\.?\d*)\s*,\s*(-?\d+\.?\d*)\s*\)",
                q
            )
            if m:
                x1,y1,x2,y2 = float(m.group(1)),float(m.group(2)),float(m.group(3)),float(m.group(4))
                expected = math.sqrt((x2-x1)**2 + (y2-y1)**2)
                return True, self._approx_eq(proposed, expected)

            return False, False
        except Exception:
            return False, False

    def _verify_graph(self, q: str, proposed: float) -> Tuple[bool, bool]:
        try:
            # Handshaking lemma: sum of degrees = 2 × |E|
            # "vertices with degrees d1, d2, ..., dn. How many edges?"
            m = re.search(r"degrees?\s+([\d,\s\band\b]+?)\.?\s+how many edges", q, re.I)
            if m:
                degrees = [int(x) for x in re.findall(r"\d+", m.group(1))]
                if degrees and sum(degrees) % 2 == 0:
                    expected = sum(degrees) / 2
                    return True, self._approx_eq(proposed, expected)

            return False, False
        except Exception:
            return False, False

    def _verify_proof_sampling(self, q: str, proposed: float) -> Tuple[bool, bool]:
        """
        Sample-test "f(n) = g(n) for all n" claims using sympy.lambdify — NO eval().
        Only verifiable when the expression uses basic algebraic operations on n.
        """
        try:
            from sympy import symbols, sympify, lambdify

            n_sym = symbols("n", integer=True, positive=True)

            def _safe_sample(expr_str: str):
                """Parse expr_str with sympy (safe — no exec/eval) and evaluate at n=1..20."""
                try:
                    expr = sympify(expr_str, locals={"n": n_sym})
                    # Reject if expression contains free symbols other than n
                    if expr.free_symbols - {n_sym}:
                        return None
                    f = lambdify(n_sym, expr, modules="math")
                    return [float(f(k)) for k in range(1, 21)]
                except Exception:
                    return None

            # "f(n) is always even"
            m = re.search(r"(n[\^*+\-\s\d()]+)\s+is\s+always\s+even", q, re.I)
            if m:
                expr_str = re.sub(r"\^", "**", m.group(1).strip())
                vals = _safe_sample(expr_str)
                if vals is not None:
                    all_even = all(int(round(v_)) % 2 == 0 for v_ in vals)
                    return True, (proposed >= 0.5) == all_even

            # "f(n) is always odd"
            m = re.search(r"(n[\^*+\-\s\d()]+)\s+is\s+always\s+odd", q, re.I)
            if m:
                expr_str = re.sub(r"\^", "**", m.group(1).strip())
                vals = _safe_sample(expr_str)
                if vals is not None:
                    all_odd = all(int(round(v_)) % 2 != 0 for v_ in vals)
                    return True, (proposed >= 0.5) == all_odd

            return False, False
        except Exception:
            return False, False

    def _verify_finance(self, q: str, proposed: float) -> Tuple[bool, bool]:
        try:
            # Compound interest: A = P(1 + r/n)^(nt)
            p_m = re.search(r"\$?\s*(\d[\d,]*\.?\d*)\s*(?:dollars?|principal)?", q, re.I)
            r_m = re.search(r"(\d+\.?\d*)\s*%\s*(?:annual|per\s+year|per\s+annum)?", q, re.I)
            t_m = re.search(r"(\d+)\s+years?", q, re.I)

            if p_m and r_m and t_m:
                P = float(p_m.group(1).replace(",", ""))
                r = float(r_m.group(1)) / 100.0
                t = float(t_m.group(1))

                if re.search(r"compound\w*\s+monthly", q, re.I):
                    n = 12.0
                elif re.search(r"compound\w*\s+quarterly", q, re.I):
                    n = 4.0
                elif re.search(r"compound\w*\s+(?:semi-?annually|twice)", q, re.I):
                    n = 2.0
                elif re.search(r"simple interest", q, re.I):
                    n = None  # simple interest
                else:
                    n = 1.0  # annually by default

                if n is None:
                    expected = P * (1 + r * t)
                else:
                    expected = P * ((1 + r / n) ** (n * t))

                # Use absolute tolerance of $0.10 for dollar-denominated answers
                return True, abs(proposed - expected) <= 0.10

            # Percentage change: (new - old) / old * 100
            m = re.search(
                r"(?:from|increased from|decreased from)\s+\$?(\d[\d,\.]*)\s+to\s+\$?(\d[\d,\.]*)",
                q, re.I
            )
            if m:
                old = float(m.group(1).replace(",", ""))
                new = float(m.group(2).replace(",", ""))
                if old != 0:
                    expected = (new - old) / old * 100
                    return True, self._approx_eq(proposed, expected)

            return False, False
        except Exception:
            return False, False

    # ── Shared helpers ─────────────────────────────────────────────────────

    def _approx_eq(self, a: float, b: float, tol: float = 0.02) -> bool:
        if abs(b) < 1e-9:
            return abs(a - b) < 1e-3
        return abs(a - b) / abs(b) <= tol
