solver_factorial.py

from __future__ import annotations

import math
import re
from typing import Optional, List

from models import SolverResult


# ----------------------------
# Helpers
# ----------------------------

def _clean(text: str) -> str:
    return " ".join((text or "").strip().split())


def _has_factorial_signal(raw: str, lower: str) -> bool:
    return any(
        [
            re.search(r"\b\d+\s*!", raw) is not None,
            "factorial" in lower,
            "trailing zero" in lower,
            "trailing zeros" in lower,
            "zeros at the end" in lower,
            "power of" in lower and "!" in raw,
            "highest power" in lower and "!" in raw,
            "divide" in lower and "!" in raw,
        ]
    )


def _is_prime(n: int) -> bool:
    if n < 2:
        return False
    if n in (2, 3):
        return True
    if n % 2 == 0 or n % 3 == 0:
        return False
    f = 5
    while f * f <= n:
        if n % f == 0 or n % (f + 2) == 0:
            return False
        f += 6
    return True


def _valuation_in_factorial(n: int, p: int) -> int:
    """
    Exponent of prime p in n! using Legendre's formula:
    floor(n/p) + floor(n/p^2) + ...
    """
    total = 0
    power = p
    while power <= n:
        total += n // power
        power *= p
    return total


def _factorial_steps() -> List[str]:
    return [
        "A factorial means multiply consecutive positive integers down to 1.",
        "Rewrite the factorial in expanded form if needed.",
        "Then simplify carefully, often by canceling common factors first.",
    ]


def _trailing_zero_steps() -> List[str]:
    return [
        "Trailing zeros come from factors of 10.",
        "Each factor of 10 is made from one 2 and one 5.",
        "In a factorial, there are usually more 2s than 5s, so the limiting factor is the number of 5s.",
        "Count factors of 5 using multiples of 5, then add extra 5s from 25, 125, and so on.",
    ]


def _prime_power_steps(p: int) -> List[str]:
    return [
        f"To find the power of {p} in n!, count how many times {p} appears in the factorial product.",
        f"Use repeated floor division: floor(n/{p}) + floor(n/{p**2}) + floor(n/{p**3}) + ...",
        "Stop when the next power of the prime is larger than n.",
        "That total gives the exponent of the prime in the factorial.",
    ]


def _ratio_steps() -> List[str]:
    return [
        "Expand only the larger factorial until cancellation is possible.",
        "Do not expand both factorials fully unless the numbers are very small.",
        "Cancel the common descending product.",
        "Then simplify the remaining factors.",
    ]


def _equation_steps() -> List[str]:
    return [
        "Use factorial growth: 1!, 2!, 3!, 4!, ... gets large quickly.",
        "Compare nearby factorial values rather than expanding everything at once.",
        "Match the factorial expression to the target value or simplified form.",
    ]


def _normalize_factorial_expr(expr: str) -> str:
    expr = expr.replace("^", "**")
    expr = expr.replace("×", "*")
    expr = expr.replace("÷", "/")
    expr = expr.replace("–", "-")
    expr = expr.replace("—", "-")
    expr = expr.strip()
    return expr


def _factorial_to_python(expr: str) -> str:
    """
    Convert occurrences like 7! into math.factorial(7).
    Only supports numeric factorial arguments.
    """
    expr = _normalize_factorial_expr(expr)

    def repl(match: re.Match) -> str:
        num = match.group(1)
        return f"math.factorial({num})"

    return re.sub(r"(\d+)\s*!", repl, expr)


def _safe_eval_numeric_factorial_expr(expr: str) -> Optional[int]:
    """
    Evaluate a numeric factorial expression like:
    8!/5!
    (7!*3!)/(6!)
    Only for controlled numeric inputs after conversion.
    """
    try:
        py_expr = _factorial_to_python(expr)
        if not re.fullmatch(r"[0-9\(\)\s\+\-\*/,.a-zA-Z_]+", py_expr):
            return None
        value = eval(py_expr, {"__builtins__": {}}, {"math": math})
        if isinstance(value, (int, float)) and abs(value - round(value)) < 1e-9:
            return int(round(value))
        return None
    except Exception:
        return None


def _extract_smallest_positive_solution_for_factorial_value(target: int) -> Optional[int]:
    """
    Solve n! = target for integer n if possible.
    """
    if target < 1:
        return None
    n = 1
    fact = 1
    while fact < target and n <= 50:
        n += 1
        fact *= n
    if fact == target:
        return n
    return None


def _extract_numbers_from_text(lower: str) -> List[int]:
    return [int(x) for x in re.findall(r"\b\d+\b", lower)]


# ----------------------------
# Main solver
# ----------------------------

def solve_factorial(text: str) -> Optional[SolverResult]:
    raw = text or ""
    lower = _clean(raw).lower()

    if not raw.strip():
        return None

    if not _has_factorial_signal(raw, lower):
        return None

    # --------------------------------------------------
    # 1) Trailing zeros in n!
    # Covers:
    # - trailing zeros in/of 32!
    # - zeros at the end of 100!
    # - how many zeros are at the end of 50 factorial
    # --------------------------------------------------
    trailing_patterns = [
        r"trailing zeros?.*?(?:in|of)?\s*(\d+)\s*!",
        r"zeros? at the end.*?(?:in|of)?\s*(\d+)\s*!",
        r"how many zeros?.*?(?:end|trailing).*?(\d+)\s*!",
        r"(\d+)\s*!\s*.*?trailing zeros?",
        r"(\d+)\s*factorial.*?trailing zeros?",
        r"zeros? at the end of\s*(\d+)\s*factorial",
    ]
    for pattern in trailing_patterns:
        m = re.search(pattern, lower)
        if m:
            n = int(m.group(1))
            if n < 0:
                return None

            count = _valuation_in_factorial(n, 5)

            return SolverResult(
                domain="quant",
                solved=True,
                topic="factorial_trailing_zeros",
                answer_value=None,
                internal_answer=str(count),
                steps=_trailing_zero_steps(),
            )

    # --------------------------------------------------
    # 2) Power of a prime p in n!
    # Covers:
    # - power of 2 in 25!
    # - exponent of 3 in 100!
    # - highest power of 5 dividing 80!
    # - how many times does 7 divide 50!
    # --------------------------------------------------
    prime_power_patterns = [
        r"power of\s+(\d+)\s+in\s+(\d+)\s*!",
        r"exponent of\s+(\d+)\s+in\s+(\d+)\s*!",
        r"highest power of\s+(\d+)\s+(?:that\s+)?divides\s+(\d+)\s*!",
        r"how many times does\s+(\d+)\s+divide\s+(\d+)\s*!",
        r"factor of\s+(\d+)\s+in\s+(\d+)\s*!",
    ]
    for pattern in prime_power_patterns:
        m = re.search(pattern, lower)
        if m:
            p = int(m.group(1))
            n = int(m.group(2))

            if n < 0 or not _is_prime(p):
                return None

            count = _valuation_in_factorial(n, p)

            return SolverResult(
                domain="quant",
                solved=True,
                topic="factorial_prime_power",
                answer_value=None,
                internal_answer=str(count),
                steps=_prime_power_steps(p),
            )

    # --------------------------------------------------
    # 3) Numeric factorial ratio / product evaluation
    # Covers:
    # - 8!/5!
    # - (9!)/(7!)
    # - 6! / (3! 2!)  [if user types with *]
    # - (10! * 3!) / 8!
    #
    # We only trigger if the expression contains at least
    # two factorials or one factorial plus arithmetic signs.
    # --------------------------------------------------
    numeric_factorials = re.findall(r"\d+\s*!", raw)
    if numeric_factorials:
        factorial_count = len(numeric_factorials)
        has_arithmetic = any(sym in raw for sym in ["/", "*", "+", "-", "(", ")"])

        if factorial_count >= 2 or (factorial_count >= 1 and has_arithmetic):
            expr_candidate = raw.strip()

            # Try to isolate a math-looking substring if wrapped in words
            expr_match = re.search(r"([\d\s!\(\)\+\-\*/×÷^]+)", raw)
            if expr_match:
                maybe_expr = expr_match.group(1).strip()
                if "!" in maybe_expr:
                    expr_candidate = maybe_expr

            value = _safe_eval_numeric_factorial_expr(expr_candidate)
            if value is not None:
                return SolverResult(
                    domain="quant",
                    solved=True,
                    topic="factorial_expression",
                    answer_value=None,
                    internal_answer=str(value),
                    steps=_ratio_steps(),
                )

    # --------------------------------------------------
    # 4) Factorial equations: n! = k
    # Covers:
    # - n! = 120
    # - which value of n satisfies n! = 720
    # --------------------------------------------------
    eq_match = re.search(
        r"(?:n|x)\s*!\s*=\s*(\d+)|(?:which value of|find)\s+(?:n|x).*?(?:n|x)\s*!\s*=\s*(\d+)",
        lower
    )
    if eq_match:
        target = eq_match.group(1) or eq_match.group(2)
        if target is not None:
            target_int = int(target)
            sol = _extract_smallest_positive_solution_for_factorial_value(target_int)
            if sol is not None:
                return SolverResult(
                    domain="quant",
                    solved=True,
                    topic="factorial_equation",
                    answer_value=None,
                    internal_answer=str(sol),
                    steps=_equation_steps(),
                )

    # --------------------------------------------------
    # 5) Direct small factorial value
    # Covers:
    # - 6!
    # - value of 7!
    # - factorial of 5
    #
    # Keep this after richer patterns so it does not steal
    # trailing-zero / prime-power questions.
    # --------------------------------------------------
    direct_patterns = [
        r"^\s*(\d+)\s*!\s*$",
        r"value of\s+(\d+)\s*!",
        r"what is\s+(\d+)\s*!",
        r"factorial of\s+(\d+)",
        r"value of\s+(\d+)\s+factorial",
    ]
    for pattern in direct_patterns:
        m = re.search(pattern, lower)
        if m:
            n = int(m.group(1))
            if n < 0:
                return None

            result = math.factorial(n)

            return SolverResult(
                domain="quant",
                solved=True,
                topic="factorial",
                answer_value=None,
                internal_answer=str(result),
                steps=_factorial_steps(),
            )

    # --------------------------------------------------
    # 6) “Last non-zero digit” / “units digit” in n!
    # GMAT-style edge coverage:
    # - last non-zero digit of 10!
    # - units digit of 7!
    #
    # Note: for n! with n >= 5, the units digit is 0.
    # For last non-zero digit, strip zeros.
    # --------------------------------------------------
    units_match = re.search(r"units digit.*?(\d+)\s*!|(\d+)\s*!\s*.*?units digit", lower)
    if units_match:
        n_str = units_match.group(1) or units_match.group(2)
        if n_str is not None:
            n = int(n_str)
            value = math.factorial(n)
            units = value % 10
            return SolverResult(
                domain="quant",
                solved=True,
                topic="factorial_units_digit",
                answer_value=None,
                internal_answer=str(units),
                steps=[
                    "Compute or reason about the final digit of the factorial product.",
                    "For factorials with n ≥ 5, a factor of 10 appears, so the units digit becomes 0.",
                    "For smaller factorials, expand directly.",
                ],
            )

    last_nonzero_match = re.search(
        r"last non-?zero digit.*?(\d+)\s*!|(\d+)\s*!\s*.*?last non-?zero digit",
        lower,
    )
    if last_nonzero_match:
        n_str = last_nonzero_match.group(1) or last_nonzero_match.group(2)
        if n_str is not None:
            n = int(n_str)
            value = math.factorial(n)
            while value % 10 == 0:
                value //= 10
            last_nonzero = value % 10

            return SolverResult(
                domain="quant",
                solved=True,
                topic="factorial_last_nonzero_digit",
                answer_value=None,
                internal_answer=str(last_nonzero),
                steps=[
                    "Find the factorial value or reason from its prime-factor structure.",
                    "Remove trailing zeros first.",
                    "Then take the final remaining digit.",
                ],
            )

    # --------------------------------------------------
    # 7) Divisibility-style basic recognition without a full parse
    # Example:
    # - Is 6! divisible by 9?
    # - greatest integer k such that 2^k divides 25!
    #
    # We only solve simple p^k-divides-n! patterns when p is prime.
    # --------------------------------------------------
    power_divides_match = re.search(
        r"greatest integer\s+(?:k|n).*?(\d+)\s*\^\s*(?:k|n)\s+divides\s+(\d+)\s*!",
        lower
    )
    if power_divides_match:
        base = int(power_divides_match.group(1))
        n = int(power_divides_match.group(2))
        if _is_prime(base):
            count = _valuation_in_factorial(n, base)
            return SolverResult(
                domain="quant",
                solved=True,
                topic="factorial_prime_power",
                answer_value=None,
                internal_answer=str(count),
                steps=_prime_power_steps(base),
            )

    # --------------------------------------------------
    # 8) Fallback recognition for factorial questions that are real
    # but not yet fully parsed. This is still useful because it gives
    # structured guidance without pretending to solve incorrectly.
    # --------------------------------------------------
    if "!" in raw or "factorial" in lower:
        return SolverResult(
            domain="quant",
            solved=False,
            topic="factorial",
            answer_value=None,
            internal_answer=None,
            steps=[
                "Identify whether the question is asking for direct evaluation, cancellation, trailing zeros, or prime-factor counting.",
                "If it is a ratio, expand only enough terms to cancel.",
                "If it is about zeros or divisibility, convert the question into counting prime factors inside the factorial.",
                "Then simplify step by step without expanding more than necessary.",
            ],
        )

    return None