solver_work_rate.py

from __future__ import annotations

import math
import re
from fractions import Fraction
from typing import Callable, Iterable, Optional

from models import SolverResult


# ============================================================
# Core helpers
# ============================================================

NUMBER = r"-?\d+(?:\.\d+)?"
FRACTION = r"\d+\s*/\s*\d+"
NUMTOKEN = rf"(?:{FRACTION}|{NUMBER})"


def _norm(text: str) -> str:
    s = (text or "").strip()
    s = s.replace("’", "'").replace("“", '"').replace("”", '"')
    s = s.replace("–", "-").replace("—", "-")
    s = s.replace("per cent", "percent")
    s = re.sub(r"\s+", " ", s)
    return s


def _to_float(token: str) -> float:
    token = token.strip()
    if "/" in token and re.fullmatch(r"\d+\s*/\s*\d+", token):
        a, b = token.split("/")
        return float(Fraction(int(a.strip()), int(b.strip())))
    return float(token)


def _safe_div(a: float, b: float) -> Optional[float]:
    if abs(b) < 1e-12:
        return None
    return a / b


def _approx_int(x: float, tol: float = 1e-9) -> bool:
    return abs(x - round(x)) <= tol


def _clean_num(x: Optional[float]) -> Optional[str]:
    if x is None:
        return None
    if not math.isfinite(x):
        return None
    if _approx_int(x):
        return str(int(round(x)))
    return f"{x:.6g}"


def _make_result(
    topic: str,
    internal_answer: Optional[float],
    steps: list[str],
    method: Optional[str] = None,
    what_is_asked: Optional[str] = None,
) -> SolverResult:
    extra_steps = []
    if what_is_asked:
        extra_steps.append(f"What the question is asking: {what_is_asked}")
    if method:
        extra_steps.append(f"Method: {method}")

    # Do not reveal the final value in visible fields.
    return SolverResult(
        domain="quant",
        solved=internal_answer is not None,
        topic=topic,
        answer_value=None,
        internal_answer=_clean_num(internal_answer),
        steps=extra_steps + steps,
    )


def _contains_any(lower: str, phrases: Iterable[str]) -> bool:
    return any(p in lower for p in phrases)


def _extract_all_numbers(text: str) -> list[float]:
    return [_to_float(m.group(0)) for m in re.finditer(NUMTOKEN, text)]


def _extract_time_units(lower: str) -> bool:
    return bool(re.search(r"\b(hour|hours|hr|hrs|minute|minutes|min|day|days)\b", lower))


def _workish(lower: str) -> bool:
    keywords = [
        "work", "together", "alone", "finish", "complete", "completed",
        "job", "task", "paint", "printer", "prints", "machine", "machines",
        "crew", "worker", "workers", "pipe", "pipes", "tap", "faucet",
        "fill", "filled", "drain", "drains", "empty", "empties",
        "leak", "leaks", "produce", "produces", "manufacture", "manufactures",
        "pages per minute", "sprockets per hour", "rate"
    ]
    return _contains_any(lower, keywords)


def _distanceish(lower: str) -> bool:
    keywords = [
        "distance", "mile", "miles", "km", "kilometer", "kilometers",
        "speed", "mph", "kph", "miles per hour", "kilometers per hour",
        "travel", "travels", "moving", "train", "car", "bike", "walk",
        "upstream", "downstream", "current", "stream"
    ]
    return _contains_any(lower, keywords)


def _find_times_complete_job(text: str) -> list[float]:
    """
    Finds patterns like:
    - A can do the job in 3 hours
    - finishes a task in 24 minutes
    - alone would finish in 60 minutes
    """
    pats = [
        rf"\bin\s+({NUMTOKEN})\s*(?:hour|hours|hr|hrs|minute|minutes|min|day|days)\b",
        rf"\btakes?\s+({NUMTOKEN})\s*(?:hour|hours|hr|hrs|minute|minutes|min|day|days)\b",
        rf"\bfinish(?:es)?(?:\s+the\s+\w+)?\s+in\s+({NUMTOKEN})\s*(?:hour|hours|hr|hrs|minute|minutes|min|day|days)\b",
        rf"\bcomplete(?:s|d)?(?:\s+the\s+\w+)?\s+in\s+({NUMTOKEN})\s*(?:hour|hours|hr|hrs|minute|minutes|min|day|days)\b",
    ]
    vals: list[float] = []
    for pat in pats:
        for m in re.finditer(pat, text, flags=re.I):
            vals.append(_to_float(m.group(1)))
    return vals


def _find_percent_more_less(lower: str) -> Optional[tuple[str, float]]:
    """
    Returns ("more"/"less", percent_as_decimal)
    """
    m = re.search(rf"({NUMTOKEN})\s*%\s*(more|less)", lower)
    if m:
        return m.group(2), _to_float(m.group(1)) / 100.0
    return None


def _find_times_faster_slower(lower: str) -> Optional[tuple[str, float]]:
    """
    Examples:
    - B is 2 times as fast as A
    - Y's rate is 1/10 of X
    """
    m = re.search(rf"\b({NUMTOKEN})\s+times?\s+as\s+(fast|slow)\b", lower)
    if m:
        return m.group(2), _to_float(m.group(1))

    m2 = re.search(rf"\brate\s+is\s+({NUMTOKEN})\s+of\b", lower)
    if m2:
        return "multiplier", _to_float(m2.group(1))
    return None


# ============================================================
# Work-rate solving blocks
# ============================================================

def _solve_basic_together(text: str, lower: str) -> Optional[SolverResult]:
    """
    Basic:
    A in t1, B in t2, how long together?
    Also handles 3+ workers if multiple independent completion times appear.
    """
    if not _contains_any(lower, ["together", "work together", "working together"]):
        return None

    times = _find_times_complete_job(text)
    if len(times) < 2:
        # fallback: use first two numeric tokens only when wording is clearly work-like
        nums = _extract_all_numbers(text)
        if len(nums) >= 2 and _workish(lower) and _extract_time_units(lower):
            times = nums[:2]

    if len(times) < 2:
        return None

    rates = []
    for t in times:
        if t <= 0:
            return None
        rates.append(1.0 / t)

    total_rate = sum(rates)
    total_time = _safe_div(1.0, total_rate)
    if total_time is None:
        return None

    return _make_result(
        topic="work_rate",
        internal_answer=total_time,
        what_is_asked="how to combine individual completion times into one combined completion time",
        method="Convert each worker's time into rate = 1/job per time unit, add the rates, then invert.",
        steps=[
            "Identify each worker/machine's individual time for one full job.",
            "Convert each time into a per-unit work rate using rate = 1 ÷ time.",
            "Add the rates because work done in the same unit of time combines.",
            "Take the reciprocal of the combined rate to get total time.",
        ],
    )


def _solve_combined_and_one_alone(text: str, lower: str) -> Optional[SolverResult]:
    """
    Together time + one alone time => other alone time.
    Example:
    A and B together finish in 24 min. A alone in 60 min. How long for B alone?
    """
    if not _contains_any(lower, ["together", "alone"]):
        return None

    nums = _extract_all_numbers(text)
    if len(nums) < 2:
        return None

    # Heuristic: first number = together time, second = one-alone time
    t_together = nums[0]
    t_one = nums[1]
    if t_together <= 0 or t_one <= 0:
        return None

    rate_other = (1.0 / t_together) - (1.0 / t_one)
    t_other = _safe_div(1.0, rate_other)
    if t_other is None or t_other <= 0:
        return None

    ask_other_alone = _contains_any(
        lower,
        ["how long will", "how long would", "how much time", "alone", "by himself", "by herself", "by itself"]
    )
    if not ask_other_alone:
        return None

    return _make_result(
        topic="work_rate",
        internal_answer=t_other,
        what_is_asked="the missing individual completion time when the combined time and one person's time are known",
        method="Use combined rate minus known individual rate to isolate the unknown rate, then invert.",
        steps=[
            "Write the combined rate as 1 ÷ together-time.",
            "Write the known person's rate as 1 ÷ known-alone-time.",
            "Subtract to get the unknown person's rate.",
            "Invert that unknown rate to get the missing alone time.",
        ],
    )


def _solve_fraction_done_by_x_then_y(text: str, lower: str) -> Optional[SolverResult]:
    """
    Example:
    X takes 12 hours. He finishes 2/3 of the work.
    Rest is finished by Y whose rate is 1/10 of X.
    How much time does Y take for his portion?
    """
    if "rest of the work" not in lower and "remaining" not in lower and "finishes" not in lower:
        return None

    base_time_match = re.search(
        rf"\btakes?\s+({NUMTOKEN})\s*(hour|hours|hr|hrs|minute|minutes|min|day|days)\b",
        lower
    )
    frac_done_match = re.search(rf"\bfinishes?\s+({NUMTOKEN})\s+of\s+the\s+work\b", lower)
    rate_ratio_match = re.search(rf"\brate\s+is\s+({NUMTOKEN})\s+of\b", lower)

    if not (base_time_match and frac_done_match and rate_ratio_match):
        return None

    base_time = _to_float(base_time_match.group(1))
    frac_done = _to_float(frac_done_match.group(1))
    ratio = _to_float(rate_ratio_match.group(1))

    if base_time <= 0 or ratio <= 0 or frac_done <= 0 or frac_done >= 1:
        return None

    remaining = 1.0 - frac_done
    x_rate = 1.0 / base_time
    y_rate = ratio * x_rate
    y_time = _safe_div(remaining, y_rate)
    if y_time is None or y_time <= 0:
        return None

    return _make_result(
        topic="work_rate",
        internal_answer=y_time,
        what_is_asked="the time needed for a second worker to complete only the remaining fraction of the job",
        method="Convert the first worker's time to rate, scale the second worker's rate from the ratio, then divide remaining work by that rate.",
        steps=[
            "Convert the known worker's full-job time into a unit rate.",
            "Translate the completed fraction into remaining work: remaining = 1 - completed fraction.",
            "Use the stated rate ratio to get the second worker's rate.",
            "Use time = remaining work ÷ rate.",
        ],
    )


def _solve_work_after_some_time(text: str, lower: str) -> Optional[SolverResult]:
    """
    General remaining-work setup:
    A can do in a hours, B in b hours. They work together for x hours.
    How much remains / what fraction remains / how much completed?
    """
    if not _contains_any(lower, ["remain", "remains", "remaining", "left", "completed", "after"]):
        return None

    nums = _extract_all_numbers(text)
    if len(nums) < 3:
        return None

    # common GMAT-style heuristic: first two are full-job times, third is elapsed time
    a, b, x = nums[0], nums[1], nums[2]
    if a <= 0 or b <= 0 or x < 0:
        return None

    combined = (1.0 / a) + (1.0 / b)
    completed = combined * x
    remaining = 1.0 - completed

    if "remain" in lower or "left" in lower or "remaining" in lower:
        internal = remaining
        asked = "the fraction of work still left after some amount of joint work"
    else:
        internal = completed
        asked = "the fraction of work completed after some amount of joint work"

    return _make_result(
        topic="work_rate",
        internal_answer=internal,
        what_is_asked=asked,
        method="Find the combined rate, multiply by elapsed time to get completed work, then subtract from 1 if the question asks for what remains.",
        steps=[
            "Convert each worker's full-job time into a rate.",
            "Add the rates if they are working simultaneously.",
            "Multiply the combined rate by the elapsed time to get the completed fraction.",
            "If the question asks what remains, subtract that completed fraction from 1.",
        ],
    )


def _solve_join_leave_case(text: str, lower: str) -> Optional[SolverResult]:
    """
    Handles:
    - A starts alone, B joins later
    - A and B start together, one leaves
    Requires 3 times/numbers in common patterns.
    """
    join_words = ["joins", "join", "joined", "after", "left", "leaves", "leave"]
    if not _contains_any(lower, join_words):
        return None

    nums = _extract_all_numbers(text)
    if len(nums) < 3:
        return None

    a, b, x = nums[0], nums[1], nums[2]
    if a <= 0 or b <= 0 or x < 0:
        return None

    a_rate = 1.0 / a
    b_rate = 1.0 / b

    if _contains_any(lower, ["joins", "join", "joined"]):
        completed_before_join = a_rate * x
        remaining = 1.0 - completed_before_join
        together_rate = a_rate + b_rate
        rest_time = _safe_div(remaining, together_rate)
        total_time = None if rest_time is None else x + rest_time

        if total_time is None or total_time <= 0:
            return None

        return _make_result(
            topic="work_rate",
            internal_answer=total_time,
            what_is_asked="the total completion time when one worker starts first and another joins later",
            method="Split the job into stages: first-stage work by one worker, then remaining work by the combined rate.",
            steps=[
                "Convert each person's full-job time into a rate.",
                "Compute how much work the first worker finishes before the second joins.",
                "Subtract from 1 to get the remaining work.",
                "Use the combined rate for the remaining stage, then add the elapsed starting time.",
            ],
        )

    if _contains_any(lower, ["left", "leaves", "leave"]):
        completed_before_leave = (a_rate + b_rate) * x
        remaining = 1.0 - completed_before_leave
        # heuristic: after leaving, first worker continues
        rest_time = _safe_div(remaining, a_rate)
        total_time = None if rest_time is None else x + rest_time

        if total_time is None or total_time <= 0:
            return None

        return _make_result(
            topic="work_rate",
            internal_answer=total_time,
            what_is_asked="the total completion time when workers begin together and one stops after a known amount of time",
            method="Split the timeline into stages: together first, then remaining work by the continuing worker.",
            steps=[
                "Convert each worker's time into a unit work rate.",
                "Find the completed fraction during the stage when both are active.",
                "Subtract from 1 to get the remaining fraction.",
                "Finish the remainder using the continuing worker's rate and add stage times.",
            ],
        )

    return None


def _solve_fill_drain(text: str, lower: str) -> Optional[SolverResult]:
    """
    Pipes/tanks:
    fill in a hours, drain/empty in b hours => net rate = 1/a - 1/b
    """
    if not _contains_any(lower, ["pipe", "pipes", "fill", "filled", "drain", "drains", "empty", "empties", "leak", "leaks", "tank", "cistern"]):
        return None

    nums = _extract_all_numbers(text)
    if len(nums) < 2:
        return None

    a, b = nums[0], nums[1]
    if a <= 0 or b <= 0:
        return None

    fill_rate = 1.0 / a
    drain_rate = 1.0 / b
    net = fill_rate - drain_rate

    if net <= 0:
        return None

    t = _safe_div(1.0, net)
    if t is None:
        return None

    return _make_result(
        topic="work_rate",
        internal_answer=t,
        what_is_asked="the net completion time when one process adds work and another removes it",
        method="Treat filling as positive rate and draining/leaking as negative rate, then invert the net rate.",
        steps=[
            "Convert the filling time into a positive rate.",
            "Convert the draining/leaking time into a negative rate.",
            "Subtract the drain rate from the fill rate to get the net rate.",
            "Take the reciprocal of the net rate to get total time.",
        ],
    )


def _solve_page_rate_difference(text: str, lower: str) -> Optional[SolverResult]:
    """
    Example:
    Together finish in 24 min. A alone in 60 min.
    B prints 5 pages/min more than A.
    How many pages in task?
    """
    if "pages" not in lower and "pages a minute" not in lower and "prints" not in lower:
        return None

    nums = _extract_all_numbers(text)
    if len(nums) < 3:
        return None

    together_time, a_time, diff_pages = nums[0], nums[1], nums[2]
    if together_time <= 0 or a_time <= 0 or diff_pages <= 0:
        return None

    # difference in job/minute between B and A
    b_job_rate = (1.0 / together_time) - (1.0 / a_time)
    a_job_rate = 1.0 / a_time
    job_rate_difference = abs(b_job_rate - a_job_rate)

    if job_rate_difference <= 0:
        return None

    pages_total = _safe_div(diff_pages, job_rate_difference)
    if pages_total is None or pages_total <= 0:
        return None

    return _make_result(
        topic="work_rate",
        internal_answer=pages_total,
        what_is_asked="the total size of the job when job-rate differences correspond to a page-per-minute difference",
        method="Translate both workers into job/minute, compare those job-rates, then scale that difference up to the stated page difference.",
        steps=[
            "Write the combined job rate using the together time.",
            "Write the known worker's job rate using the alone time.",
            "Subtract to get the other worker's job rate.",
            "Find the difference in their job-rates per minute.",
            "Use the fact that this difference corresponds to the stated pages-per-minute gap to scale up to the full job size.",
        ],
    )


def _solve_percent_more_output_fixed_total(text: str, lower: str) -> Optional[SolverResult]:
    """
    Example:
    A takes 10 hours longer than B to produce 660 sprockets.
    B produces 10% more per hour than A.
    Find A's rate.
    """
    if not _contains_any(lower, ["more per hour", "less per hour", "produces", "produce", "manufacture", "sprockets", "pages"]):
        return None

    nums = _extract_all_numbers(text)
    pm = _find_percent_more_less(lower)

    if len(nums) < 3 or pm is None:
        return None

    total_units = nums[0]
    time_difference = nums[1]
    kind, pct = pm

    if total_units <= 0 or time_difference <= 0 or pct <= 0:
        return None

    # Let A = r units/time. B = r(1+p) if "more", else r(1-p)
    multiplier = (1.0 + pct) if kind == "more" else (1.0 - pct)
    if multiplier <= 0:
        return None

    # total/r - total/(multiplier*r) = time_difference
    # total * (1 - 1/multiplier) / r = time_difference
    numerator = total_units * (1.0 - 1.0 / multiplier)
    r = _safe_div(numerator, time_difference)
    if r is None or r <= 0:
        return None

    return _make_result(
        topic="work_rate",
        internal_answer=r,
        what_is_asked="the base production rate when total output, time difference, and relative efficiency are given",
        method="Let the slower rate be r, express the faster rate using the percent relationship, write time = total ÷ rate for each, and subtract the times.",
        steps=[
            "Assign a variable to the unknown base production rate.",
            "Translate the percent-more/less statement into the other machine's rate.",
            "Write each completion time as total output ÷ rate.",
            "Use the stated time difference to form one equation and solve for the base rate.",
        ],
    )


def _solve_direct_rate_from_units_and_time(text: str, lower: str) -> Optional[SolverResult]:
    """
    Direct unit-rate:
    660 sprockets in 12 hours -> 55 per hour
    distance/rate/time analogue included.
    """
    if not _contains_any(lower, ["per hour", "per minute", "per day", "rate", "speed", "mph", "kph"]):
        return None

    nums = _extract_all_numbers(text)
    if len(nums) < 2:
        return None

    a, b = nums[0], nums[1]
    if a <= 0 or b <= 0:
        return None

    # Heuristic: if asks "per hour" or "speed", compute units/time
    if _contains_any(lower, ["per hour", "per minute", "per day", "speed", "mph", "kph"]):
        rate = _safe_div(a, b)
        if rate is None or rate <= 0:
            return None

        topic = "distance_rate" if _distanceish(lower) else "work_rate"
        asked = (
            "the unit rate or speed from a total amount and a time"
            if topic == "work_rate"
            else "the speed from distance and time"
        )
        method = (
            "Use rate = total work ÷ time."
            if topic == "work_rate"
            else "Use speed = distance ÷ time."
        )

        return _make_result(
            topic=topic,
            internal_answer=rate,
            what_is_asked=asked,
            method=method,
            steps=[
                "Identify the total amount completed or traveled.",
                "Identify the time taken.",
                "Use rate = amount ÷ time.",
            ],
        )

    return None


# ============================================================
# Distance/rate sub-blocks
# ============================================================

def _solve_distance_speed_time(text: str, lower: str) -> Optional[SolverResult]:
    """
    Basic distance = rate * time family.
    Supports solving for one missing quantity in simple 2-number problems.
    """
    if not _distanceish(lower):
        return None

    nums = _extract_all_numbers(text)
    if len(nums) < 2:
        return None

    a, b = nums[0], nums[1]
    if a <= 0 or b <= 0:
        return None

    if "how far" in lower or "distance" in lower:
        d = a * b
        return _make_result(
            topic="distance_rate",
            internal_answer=d,
            what_is_asked="the distance traveled from speed and time",
            method="Use distance = rate × time.",
            steps=[
                "Identify the rate/speed.",
                "Identify the travel time.",
                "Multiply rate by time to get distance.",
            ],
        )

    if "how long" in lower or "how much time" in lower:
        t = _safe_div(a, b)
        if t is None or t <= 0:
            return None
        return _make_result(
            topic="distance_rate",
            internal_answer=t,
            what_is_asked="the travel time from distance and speed",
            method="Use time = distance ÷ rate.",
            steps=[
                "Identify the total distance.",
                "Identify the speed.",
                "Divide distance by speed to get time.",
            ],
        )

    if "what speed" in lower or "what is the speed" in lower or "mph" in lower or "kph" in lower:
        s = _safe_div(a, b)
        if s is None or s <= 0:
            return None
        return _make_result(
            topic="distance_rate",
            internal_answer=s,
            what_is_asked="the speed from distance and time",
            method="Use speed = distance ÷ time.",
            steps=[
                "Identify the distance.",
                "Identify the time.",
                "Divide distance by time to get speed.",
            ],
        )

    return None


def _solve_relative_speed(text: str, lower: str) -> Optional[SolverResult]:
    """
    Relative speed:
    - towards each other => add speeds
    - same direction / catches up => subtract speeds
    """
    if not _distanceish(lower):
        return None

    if not _contains_any(lower, ["towards each other", "opposite directions", "same direction", "catch up", "catches up"]):
        return None

    nums = _extract_all_numbers(text)
    if len(nums) < 2:
        return None

    v1, v2 = nums[0], nums[1]
    if v1 <= 0 or v2 <= 0:
        return None

    if _contains_any(lower, ["towards each other", "opposite directions"]):
        rel = v1 + v2
        ask = "the relative speed when two objects move toward each other or in opposite directions"
        meth = "Add the speeds because separation changes by both rates together."
    else:
        rel = abs(v1 - v2)
        ask = "the relative speed when two objects move in the same direction"
        meth = "Subtract the speeds because only the speed difference closes the gap."

    return _make_result(
        topic="distance_rate",
        internal_answer=rel,
        what_is_asked=ask,
        method=meth,
        steps=[
            "Decide whether the motion uses addition or subtraction of speeds.",
            "Use addition for opposite directions/toward each other.",
            "Use subtraction for same-direction chase situations.",
        ],
    )


# ============================================================
# Question-explainer fallback
# ============================================================

def _explain_work_rate_without_solving(lower: str) -> Optional[SolverResult]:
    explain_triggers = [
        "what is this asking",
        "what does this mean",
        "how do i set this up",
        "how should i think about this",
        "explain the question",
        "decode the question",
    ]
    if not _contains_any(lower, explain_triggers):
        return None

    if _workish(lower):
        return _make_result(
            topic="work_rate",
            internal_answer=None,
            what_is_asked="how to translate a work/rate word problem into rate equations",
            method="Use a unit-job model: job = 1, rate = fraction of job per unit time, time = work ÷ rate.",
            steps=[
                "Treat the whole task as 1 complete job.",
                "Convert each person's completion time into a rate using 1 ÷ time.",
                "Add rates when people work together at the same time.",
                "Subtract rates when one process undoes another, like draining or leaking.",
                "For multi-stage problems, track completed work and remaining work separately.",
                "For percentage or ratio statements, convert the wording into algebra before solving.",
            ],
        )

    if _distanceish(lower):
        return _make_result(
            topic="distance_rate",
            internal_answer=None,
            what_is_asked="how to translate a distance/rate question into equations",
            method="Use the distance-rate-time relationship and choose the form that matches the missing quantity.",
            steps=[
                "Identify which quantity is missing: distance, speed, or time.",
                "Use distance = speed × time.",
                "Rearrange to speed = distance ÷ time or time = distance ÷ speed when needed.",
                "For relative motion, add speeds when moving toward each other and subtract when moving in the same direction.",
            ],
        )

    return None


# ============================================================
# Main entry point
# ============================================================

def solve_work_rate(text: str) -> Optional[SolverResult]:
    raw = _norm(text)
    lower = raw.lower()

    # Broad gateway: this solver now intentionally covers work/rate and
    # a useful distance/rate subset because the user wanted wider rate coverage.
    if not (_workish(lower) or _distanceish(lower) or "rate" in lower):
        return None

    solvers: list[Callable[[str, str], Optional[SolverResult]]] = [
        _solve_fraction_done_by_x_then_y,
        _solve_page_rate_difference,
        _solve_percent_more_output_fixed_total,
        _solve_join_leave_case,
        _solve_fill_drain,
        _solve_work_after_some_time,
        _solve_combined_and_one_alone,
        _solve_basic_together,
        _solve_relative_speed,
        _solve_distance_speed_time,
        _solve_direct_rate_from_units_and_time,
        _explain_work_rate_without_solving,
    ]

    for solver in solvers:
        try:
            result = solver(raw, lower)
            if result is not None:
                return result
        except Exception:
            continue

    return None