solver_ratio.py

from __future__ import annotations

import re
from fractions import Fraction
from math import gcd, isclose
from typing import Dict, List, Optional, Tuple

from models import SolverResult


_NUMBER = r"-?\d+(?:\.\d+)?"
_INT = r"-?\d+"
_NAME = r"[a-zA-Z][a-zA-Z0-9_]*"


# =========================
# core helpers
# =========================

def _clean(text: str) -> str:
    s = (text or "").strip()
    s = s.replace("–", "-").replace("—", "-").replace("−", "-")
    s = s.replace("∶", ":")
    s = re.sub(r"\s+", " ", s)
    return s


def _to_float(x: str) -> float:
    return float(x)


def _to_int_if_whole(x: float):
    if isclose(x, round(x), rel_tol=1e-9, abs_tol=1e-9):
        return int(round(x))
    return x


def _fmt_num(x: float) -> str:
    xi = _to_int_if_whole(x)
    if isinstance(xi, int):
        return str(xi)
    return f"{x:.6g}"


def _reduce_ratio(a: int, b: int) -> Tuple[int, int]:
    if a == 0 and b == 0:
        return 0, 0
    if a == 0:
        return 0, 1 if b > 0 else -1
    if b == 0:
        return 1 if a > 0 else -1, 0
    g = gcd(abs(a), abs(b))
    a //= g
    b //= g
    if b < 0:
        a, b = -a, -b
    return a, b


def _safe_ratio_answer(a: float, b: float) -> str:
    if isclose(a, round(a), abs_tol=1e-9) and isclose(b, round(b), abs_tol=1e-9):
        ra, rb = _reduce_ratio(int(round(a)), int(round(b)))
        return f"{ra}:{rb}"
    return f"{_fmt_num(a)}:{_fmt_num(b)}"


def _fraction_from_text(s: str) -> Fraction:
    s = s.strip()
    if "/" in s:
        a, b = s.split("/", 1)
        return Fraction(a.strip()) / Fraction(b.strip())
    return Fraction(s)


def _make_result(
    *,
    topic: str,
    internal_answer: Optional[str],
    steps: List[str],
    solved: bool = True,
    answer_value: Optional[str] = None,
) -> SolverResult:
    return SolverResult(
        domain="quant",
        solved=solved,
        topic=topic,
        answer_value=answer_value if answer_value is not None else internal_answer,
        internal_answer=internal_answer,
        steps=steps,
    )


# =========================
# detection helpers
# =========================

def _has_ratio_intent(lower: str) -> bool:
    triggers = [
        "ratio",
        "proportion",
        ":",
        " to ",
        "for every",
        "out of",
        "divide",
        "split",
        "share",
        "distributed",
        "mixture",
        "combined",
        "boys",
        "girls",
        "men",
        "women",
        "students",
        "teachers",
    ]
    return any(t in lower for t in triggers)


def _extract_ratio_pair(lower: str) -> Optional[Tuple[float, float]]:
    patterns = [
        rf"ratio.*?is\s+({_NUMBER})\s*:\s*({_NUMBER})",
        rf"ratio.*?is\s+({_NUMBER})\s+to\s+({_NUMBER})",
        rf"in\s+the\s+ratio\s+({_NUMBER})\s*:\s*({_NUMBER})",
        rf"in\s+the\s+ratio\s+({_NUMBER})\s+to\s+({_NUMBER})",
        rf"({_NUMBER})\s*:\s*({_NUMBER})",
        rf"({_NUMBER})\s+to\s+({_NUMBER})",
    ]
    for pat in patterns:
        m = re.search(pat, lower)
        if m:
            return _to_float(m.group(1)), _to_float(m.group(2))
    return None


def _extract_three_part_ratio(lower: str) -> Optional[Tuple[float, float, float]]:
    patterns = [
        rf"ratio.*?is\s+({_NUMBER})\s*:\s*({_NUMBER})\s*:\s*({_NUMBER})",
        rf"in\s+the\s+ratio\s+({_NUMBER})\s*:\s*({_NUMBER})\s*:\s*({_NUMBER})",
        rf"({_NUMBER})\s*:\s*({_NUMBER})\s*:\s*({_NUMBER})",
    ]
    for pat in patterns:
        m = re.search(pat, lower)
        if m:
            return _to_float(m.group(1)), _to_float(m.group(2)), _to_float(m.group(3))
    return None


def _extract_total(lower: str) -> Optional[float]:
    patterns = [
        rf"total(?:\s+number)?(?:\s+is|\s+of|\s*=)?\s*({_NUMBER})",
        rf"sum(?:\s+is|\s+of|\s*=)?\s*({_NUMBER})",
        rf"({_NUMBER})\s+(?:in total|altogether|total)",
        rf"add up to\s+({_NUMBER})",
    ]
    for pat in patterns:
        m = re.search(pat, lower)
        if m:
            return _to_float(m.group(1))
    return None


def _extract_difference(lower: str) -> Optional[float]:
    patterns = [
        rf"difference(?:\s+is|\s+of|\s*=)?\s*({_NUMBER})",
        rf"({_NUMBER})\s+more\s+than",
        rf"({_NUMBER})\s+less\s+than",
        rf"exceeds.*?by\s+({_NUMBER})",
    ]
    for pat in patterns:
        m = re.search(pat, lower)
        if m:
            return _to_float(m.group(1))
    return None


def _extract_out_of(lower: str) -> Optional[Tuple[float, float]]:
    m = re.search(rf"({_NUMBER})\s+out\s+of\s+({_NUMBER})", lower)
    if m:
        return _to_float(m.group(1)), _to_float(m.group(2))
    return None


def _extract_fraction_ratio(lower: str) -> Optional[Tuple[int, int]]:
    m = re.search(rf"ratio.*?is\s+({_NUMBER}\s*/\s*{_NUMBER})", lower)
    if not m:
        return None
    frac = _fraction_from_text(m.group(1).replace(" ", ""))
    return frac.numerator, frac.denominator


def _extract_missing_term_proportion(lower: str):
    patterns = [
        rf"(x|{_NUMBER})\s*:\s*(x|{_NUMBER})\s*=\s*(x|{_NUMBER})\s*:\s*(x|{_NUMBER})",
        rf"(x|{_NUMBER})\s+to\s+(x|{_NUMBER})\s*=\s*(x|{_NUMBER})\s+to\s+(x|{_NUMBER})",
    ]
    for pat in patterns:
        m = re.search(pat, lower)
        if m:
            vals = [m.group(i) for i in range(1, 5)]
            if sum(v == "x" for v in vals) == 1:
                return vals
    return None


def _extract_named_initial_ratio(lower: str):
    # e.g. ratio of boys to girls is 3:5
    m = re.search(
        rf"ratio\s+of\s+({_NAME})\s+to\s+({_NAME})\s+is\s+({_NUMBER})\s*(?::|to)\s*({_NUMBER})",
        lower,
    )
    if m:
        return {
            "name1": m.group(1),
            "name2": m.group(2),
            "a": _to_float(m.group(3)),
            "b": _to_float(m.group(4)),
        }
    return None


def _extract_named_change(lower: str):
    # e.g. after 6 girls join, the ratio becomes 2:3
    m = re.search(
        rf"(?:after|when)\s+({_NUMBER})\s+(more\s+|additional\s+|extra\s+)?({_NAME})\s+"
        rf"(join|are added|added|enter|left|leave|are removed|removed).*?"
        rf"(?:ratio.*?becomes|ratio.*?changes?\s+to|is\s+now)\s+({_NUMBER})\s*(?::|to)\s*({_NUMBER})",
        lower,
    )
    if not m:
        return None

    delta = _to_float(m.group(1))
    changed_name = m.group(3)
    verb = m.group(4)
    new_a = _to_float(m.group(5))
    new_b = _to_float(m.group(6))

    sign = -1 if verb in ["left", "leave", "are removed", "removed"] else 1

    return {
        "delta": sign * delta,
        "changed_name": changed_name,
        "new_a": new_a,
        "new_b": new_b,
    }


def _extract_transfer_between_groups(lower: str):
    # e.g. if 4 boys move to girls, ratio becomes ...
    m = re.search(
        rf"(?:if|when|after)\s+({_NUMBER})\s+({_NAME})\s+(?:move|transfer|are transferred)\s+to\s+({_NAME}).*?"
        rf"(?:ratio.*?becomes|ratio.*?changes?\s+to|is\s+now)\s+({_NUMBER})\s*(?::|to)\s*({_NUMBER})",
        lower,
    )
    if not m:
        return None
    return {
        "delta": _to_float(m.group(1)),
        "from_name": m.group(2),
        "to_name": m.group(3),
        "new_a": _to_float(m.group(4)),
        "new_b": _to_float(m.group(5)),
    }


def _extract_chain_ratio(lower: str):
    m = re.search(
        rf"({_NAME})\s*:\s*({_NAME})\s*=\s*({_NUMBER})\s*:\s*({_NUMBER}).*?"
        rf"({_NAME})\s*:\s*({_NAME})\s*=\s*({_NUMBER})\s*:\s*({_NUMBER}).*?"
        rf"find\s+({_NAME})\s*:\s*({_NAME})",
        lower,
    )
    if m:
        return {
            "l1": m.group(1), "r1": m.group(2), "a": _to_float(m.group(3)), "b": _to_float(m.group(4)),
            "l2": m.group(5), "r2": m.group(6), "c": _to_float(m.group(7)), "d": _to_float(m.group(8)),
            "t1": m.group(9), "t2": m.group(10),
        }
    return None


def _extract_part_of_total_question(lower: str) -> Optional[str]:
    patterns = [
        r"how many\s+([a-zA-Z][a-zA-Z0-9_]*)",
        r"number of\s+([a-zA-Z][a-zA-Z0-9_]*)",
    ]
    for pat in patterns:
        m = re.search(pat, lower)
        if m:
            return m.group(1)
    return None


# =========================
# solver blocks
# =========================

def _solve_simplify_ratio(lower: str) -> Optional[SolverResult]:
    if not any(k in lower for k in ["simplify", "reduce", "lowest terms", "ratio"]):
        return None

    pair = _extract_ratio_pair(lower)
    if not pair:
        return None

    a, b = pair
    if not isclose(a, round(a), abs_tol=1e-9) or not isclose(b, round(b), abs_tol=1e-9):
        return _make_result(
            topic="ratio",
            internal_answer=f"{_fmt_num(a)}:{_fmt_num(b)}",
            steps=[
                "The question is asking for the comparison in the stated order.",
                "Write the two quantities as a ratio.",
                "Reduce only if both parts share a common factor cleanly.",
            ],
        )

    ra, rb = _reduce_ratio(int(round(a)), int(round(b)))
    return _make_result(
        topic="ratio",
        internal_answer=f"{ra}:{rb}",
        steps=[
            "The question is asking for the ratio in lowest terms.",
            "Write the ratio in the correct order.",
            "Divide both parts by their greatest common factor.",
        ],
    )


def _solve_out_of_ratio(lower: str) -> Optional[SolverResult]:
    vals = _extract_out_of(lower)
    if not vals:
        return None
    a, b = vals
    return _make_result(
        topic="ratio",
        internal_answer=_safe_ratio_answer(a, b),
        steps=[
            "The phrase 'out of' means part-to-whole.",
            "Write the first quantity over the total in the same order.",
            "Then simplify the ratio if possible.",
        ],
    )


def _solve_fraction_ratio(lower: str) -> Optional[SolverResult]:
    vals = _extract_fraction_ratio(lower)
    if not vals:
        return None
    a, b = vals
    return _make_result(
        topic="ratio",
        internal_answer=_safe_ratio_answer(a, b),
        steps=[
            "A ratio written as a fraction compares numerator to denominator.",
            "Rewrite the fraction as a:b.",
            "Then simplify if needed.",
        ],
    )


def _solve_partition_total_2(lower: str) -> Optional[SolverResult]:
    pair = _extract_ratio_pair(lower)
    total = _extract_total(lower)
    if not pair or total is None:
        return None

    a, b = pair
    unit = total / (a + b)
    x1 = a * unit
    x2 = b * unit

    ask = _extract_part_of_total_question(lower)

    if ask:
        named = _extract_named_initial_ratio(lower)
        if named:
            if ask == named["name1"]:
                ans = x1
            elif ask == named["name2"]:
                ans = x2
            else:
                ans = None
            if ans is not None:
                return _make_result(
                    topic="ratio_partition",
                    internal_answer=_fmt_num(ans),
                    steps=[
                        "The question is asking for one share from a total split in a given ratio.",
                        "Add the ratio parts to find the total number of equal units.",
                        "Divide the total by that number to get one unit.",
                        "Multiply by the requested part of the ratio.",
                    ],
                )

    return _make_result(
        topic="ratio_partition",
        internal_answer=f"{_fmt_num(x1)}, {_fmt_num(x2)}",
        steps=[
            "The question is asking you to split a total in a given ratio.",
            "Add the ratio parts to get the number of equal units.",
            "Find one unit by dividing the total by the total number of ratio parts.",
            "Multiply each ratio part by that unit value.",
        ],
    )


def _solve_partition_difference(lower: str) -> Optional[SolverResult]:
    pair = _extract_ratio_pair(lower)
    diff = _extract_difference(lower)
    if not pair or diff is None:
        return None

    a, b = pair
    part_diff = abs(a - b)
    if isclose(part_diff, 0.0):
        return None

    unit = diff / part_diff
    x1 = a * unit
    x2 = b * unit

    return _make_result(
        topic="ratio_difference",
        internal_answer=f"{_fmt_num(x1)}, {_fmt_num(x2)}",
        steps=[
            "The question gives the ratio and the actual difference between the quantities.",
            "Subtract the ratio parts to find the difference in ratio units.",
            "Match the real difference to those units to find one unit.",
            "Scale each ratio part by that unit value.",
        ],
    )


def _solve_missing_term_proportion(lower: str) -> Optional[SolverResult]:
    vals = _extract_missing_term_proportion(lower)
    if not vals:
        return None

    parsed: List[Optional[float]] = []
    x_index = None
    for i, v in enumerate(vals):
        if v == "x":
            parsed.append(None)
            x_index = i
        else:
            parsed.append(_to_float(v))

    a, b, c, d = parsed
    x = None

    if x_index == 0 and b not in (None, 0) and c is not None and d not in (None, 0):
        x = (b * c) / d
    elif x_index == 1 and a is not None and c is not None and c != 0 and d is not None:
        x = (a * d) / c
    elif x_index == 2 and a is not None and b not in (None, 0) and d is not None:
        x = (a * d) / b
    elif x_index == 3 and a not in (None, 0) and b is not None and c is not None:
        x = (b * c) / a

    if x is None:
        return None

    return _make_result(
        topic="proportion",
        internal_answer=_fmt_num(x),
        steps=[
            "The question is asking for a missing term in equivalent ratios.",
            "Keep the positions of corresponding quantities aligned.",
            "Use the fact that equal ratios represent the same multiplicative relationship.",
            "Solve for the unknown term without changing the order of the ratio.",
        ],
    )


def _solve_chain_ratio(lower: str) -> Optional[SolverResult]:
    info = _extract_chain_ratio(lower)
    if not info:
        return None

    if info["r1"] != info["l2"]:
        return None

    a, b, c, d = info["a"], info["b"], info["c"], info["d"]
    if not all(isclose(v, round(v), abs_tol=1e-9) for v in [a, b, c, d]):
        return None

    a, b, c, d = map(lambda z: int(round(z)), [a, b, c, d])

    lcm_mid = b * c // gcd(b, c)
    scale1 = lcm_mid // b
    scale2 = lcm_mid // c

    mapping = {
        info["l1"]: a * scale1,
        info["r1"]: lcm_mid,
        info["r2"]: d * scale2,
    }

    if info["t1"] not in mapping or info["t2"] not in mapping:
        return None

    ans = _safe_ratio_answer(mapping[info["t1"]], mapping[info["t2"]])
    return _make_result(
        topic="ratio_chain",
        internal_answer=ans,
        steps=[
            "The question is asking you to combine two linked ratios with a shared middle term.",
            "Make the shared term equal in both ratios.",
            "Then compare the two end terms and simplify.",
        ],
    )


def _solve_three_part_total(lower: str) -> Optional[SolverResult]:
    triple = _extract_three_part_ratio(lower)
    total = _extract_total(lower)
    if not triple or total is None:
        return None

    a, b, c = triple
    units = a + b + c
    if isclose(units, 0.0):
        return None

    u = total / units
    x1, x2, x3 = a * u, b * u, c * u

    return _make_result(
        topic="ratio_three_part",
        internal_answer=f"{_fmt_num(x1)}, {_fmt_num(x2)}, {_fmt_num(x3)}",
        steps=[
            "The question is asking you to divide a total among three quantities in a given ratio.",
            "Add all ratio parts to get the total number of equal units.",
            "Divide the total by that number to find one unit.",
            "Multiply each ratio part by the unit value.",
        ],
    )


def _solve_named_changed_ratio(lower: str) -> Optional[SolverResult]:
    initial = _extract_named_initial_ratio(lower)
    change = _extract_named_change(lower)
    if not initial or not change:
        return None

    name1 = initial["name1"]
    name2 = initial["name2"]
    a = initial["a"]
    b = initial["b"]

    changed = change["changed_name"]
    delta = change["delta"]
    new_a = change["new_a"]
    new_b = change["new_b"]

    # original quantities = ax and bx
    # if name1 changed: (ax + delta)/(bx) = new_a/new_b
    # if name2 changed: (ax)/(bx + delta) = new_a/new_b
    if changed == name1:
        denom = new_b * a - new_a * b
        if isclose(denom, 0.0):
            return None
        x = -(new_b * delta) / denom
    elif changed == name2:
        denom = new_b * a - new_a * b
        if isclose(denom, 0.0):
            return None
        x = (new_a * delta) / denom
    else:
        return None

    if x < 0:
        return None

    q1 = a * x
    q2 = b * x

    return _make_result(
        topic="ratio_change",
        internal_answer=f"{_fmt_num(q1)}, {_fmt_num(q2)}",
        steps=[
            "This is a ratio problem with an absolute change to one named group.",
            "Represent the original groups as ratio parts multiplied by the same scale factor.",
            "Apply the increase or decrease only to the group that changed.",
            "Set the new expression equal to the new ratio and solve for the common scale factor.",
            "Use that factor to recover the original quantities.",
        ],
    )


def _solve_transfer_between_groups(lower: str) -> Optional[SolverResult]:
    initial = _extract_named_initial_ratio(lower)
    transfer = _extract_transfer_between_groups(lower)
    if not initial or not transfer:
        return None

    name1 = initial["name1"]
    name2 = initial["name2"]
    a = initial["a"]
    b = initial["b"]

    from_name = transfer["from_name"]
    to_name = transfer["to_name"]
    delta = transfer["delta"]
    new_a = transfer["new_a"]
    new_b = transfer["new_b"]

    # original ax, bx
    if from_name == name1 and to_name == name2:
        # (ax - d)/(bx + d) = new_a/new_b
        denom = new_b * a - new_a * b
        if isclose(denom, 0.0):
            return None
        x = delta * (new_a + new_b) / denom
    elif from_name == name2 and to_name == name1:
        # (ax + d)/(bx - d) = new_a/new_b
        denom = new_b * a - new_a * b
        if isclose(denom, 0.0):
            return None
        x = -delta * (new_a + new_b) / denom
    else:
        return None

    if x < 0:
        return None

    q1 = a * x
    q2 = b * x

    return _make_result(
        topic="ratio_transfer",
        internal_answer=f"{_fmt_num(q1)}, {_fmt_num(q2)}",
        steps=[
            "This is a transfer problem, so one group decreases while the other increases by the same amount.",
            "Represent the original groups using a common scale factor based on the initial ratio.",
            "Apply the transfer in opposite directions to the two groups.",
            "Match the resulting quantities to the new ratio and solve for the scale factor.",
            "Then recover the original group sizes.",
        ],
    )


def _solve_direct_proportion_text(lower: str) -> Optional[SolverResult]:
    if not any(k in lower for k in ["proportion", "proportional", "same rate", "spread evenly", "uniformly"]):
        return None

    nums = [float(x) for x in re.findall(_NUMBER, lower)]
    if len(nums) < 3:
        return None

    whole1, whole2, part1 = nums[0], nums[1], nums[2]
    if isclose(whole1, 0.0):
        return None

    x = whole2 * part1 / whole1

    return _make_result(
        topic="proportion",
        internal_answer=_fmt_num(x),
        steps=[
            "The question is asking for a missing amount under the same rate or same distribution pattern.",
            "Set up equivalent ratios so whole-to-whole matches part-to-part.",
            "Keep corresponding quantities aligned in the same order.",
            "Then solve the proportion for the missing term.",
        ],
    )


def _solve_ratio(text: str) -> Optional[SolverResult]:
    lower = _clean(text).lower()

    if not _has_ratio_intent(lower):
        return None

    solvers = [
        _solve_named_changed_ratio,
        _solve_transfer_between_groups,
        _solve_partition_total_2,
        _solve_three_part_total,
        _solve_partition_difference,
        _solve_missing_term_proportion,
        _solve_chain_ratio,
        _solve_direct_proportion_text,
        _solve_fraction_ratio,
        _solve_out_of_ratio,
        _solve_simplify_ratio,
    ]

    for solver in solvers:
        result = solver(lower)
        if result is not None:
            return result

    return _make_result(
        topic="ratio",
        internal_answer=None,
        answer_value=None,
        solved=False,
        steps=[
            "This looks like a ratio or proportion problem, but the exact structure is not fully clear yet.",
            "First identify which quantities are being compared and in what order.",
            "Then decide whether the problem is about simplifying, splitting a total, using a difference, handling a change, or setting up an equivalent proportion.",
            "Once that structure is clear, the relationship can be translated into ratio units or algebra cleanly.",
        ],
    )


def solve_ratio(text: str) -> Optional[SolverResult]:
    return _solve_ratio(text)