solver_distance_rate_time.py

from __future__ import annotations

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

from models import SolverResult


_NUMBER = r"(-?\d+(?:\.\d+)?)"


def _clean(text: str) -> str:
    t = text or ""
    t = t.replace("×", "x").replace("–", "-").replace("—", "-")
    t = t.replace("hrs", "hours").replace("hr", "hour")
    t = t.replace("mins", "minutes").replace("min", "minute")
    t = t.replace("secs", "seconds").replace("sec", "second")
    t = t.replace("kms", "km").replace("kilometers", "km").replace("kilometres", "km")
    t = t.replace("miles per hour", "mph")
    t = t.replace("kilometers per hour", "kmph")
    t = t.replace("kilometres per hour", "kmph")
    return t.strip()


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


def _fmt_num(x: float) -> str:
    if abs(x - round(x)) < 1e-9:
        return str(int(round(x)))
    return f"{x:.6g}"


def _safe_positive(x: float) -> bool:
    return math.isfinite(x) and x > 0


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


def _extract_number_before_unit(text: str, unit_pattern: str) -> Optional[float]:
    m = re.search(rf"{_NUMBER}\s*{unit_pattern}", text)
    if m:
        return _to_float(m.group(1))
    return None


def _extract_all_number_unit_pairs(text: str, unit_pattern: str) -> List[float]:
    return [_to_float(x) for x in re.findall(rf"{_NUMBER}\s*{unit_pattern}", text)]


def _extract_speeds(text: str) -> List[float]:
    speeds = []
    speeds += _extract_all_number_unit_pairs(text, r"mph\b")
    speeds += _extract_all_number_unit_pairs(text, r"kmph\b")
    speeds += _extract_all_number_unit_pairs(text, r"m/s\b")
    speeds += _extract_all_number_unit_pairs(text, r"ft/s\b")
    return speeds


def _extract_times_in_hours(text: str) -> List[float]:
    times = []
    times += _extract_all_number_unit_pairs(text, r"hours?\b")
    times += [x / 60.0 for x in _extract_all_number_unit_pairs(text, r"minutes?\b")]
    times += [x / 3600.0 for x in _extract_all_number_unit_pairs(text, r"seconds?\b")]
    return times


def _extract_distances(text: str) -> List[Tuple[float, str]]:
    out: List[Tuple[float, str]] = []
    for val in re.findall(rf"{_NUMBER}\s*(miles?|mi\b)", text):
        out.append((_to_float(val[0] if isinstance(val, tuple) else val), "mile"))
    for val in re.findall(rf"{_NUMBER}\s*(km\b)", text):
        out.append((_to_float(val[0] if isinstance(val, tuple) else val), "km"))
    for val in re.findall(rf"{_NUMBER}\s*(meters?|m\b)", text):
        out.append((_to_float(val[0] if isinstance(val, tuple) else val), "m"))
    for val in re.findall(rf"{_NUMBER}\s*(feet|foot|ft\b)", text):
        out.append((_to_float(val[0] if isinstance(val, tuple) else val), "ft"))
    return out


def _all_numbers(text: str) -> List[float]:
    return [_to_float(x) for x in re.findall(_NUMBER, text)]


def _question_asks_time(text: str) -> bool:
    return _contains_any(
        text,
        [
            "how long", "what is the time", "find the time", "time taken",
            "how many hours", "how many minutes", "when will", "time will"
        ],
    )


def _question_asks_speed(text: str) -> bool:
    return _contains_any(
        text,
        [
            "what is the speed", "find the speed", "how fast",
            "average speed", "rate of", "what speed"
        ],
    )


def _question_asks_distance(text: str) -> bool:
    return _contains_any(
        text,
        [
            "what is the distance", "find the distance", "how far",
            "distance traveled", "distance travelled"
        ],
    )


def _make_result(
    internal_answer: str,
    steps: List[str],
    solved: bool = True,
    topic: str = "distance_rate_time",
) -> SolverResult:
    return SolverResult(
        domain="quant",
        solved=solved,
        topic=topic,
        answer_value=None,          # keep final answer hidden from user-facing reply
        internal_answer=internal_answer,
        steps=steps,
    )


def _basic_formula_solver(text: str) -> Optional[SolverResult]:
    speeds = _extract_speeds(text)
    times = _extract_times_in_hours(text)
    distances = _extract_distances(text)

    # distance = speed * time
    if _question_asks_distance(text) and speeds and times:
        s = speeds[0]
        t = times[0]
        d = s * t
        return _make_result(
            internal_answer=_fmt_num(d),
            steps=[
                "Identify the two known quantities: speed and time.",
                "Use the relationship distance = speed × time.",
                f"Substitute the known values: distance = { _fmt_num(s) } × { _fmt_num(t) }.",
                "Multiply to get the travel distance.",
            ],
            topic="distance_rate_time_basic",
        )

    # time = distance / speed
    if _question_asks_time(text) and distances and speeds:
        d = distances[0][0]
        s = speeds[0]
        if not _safe_positive(s):
            return None
        t = d / s
        return _make_result(
            internal_answer=_fmt_num(t),
            steps=[
                "Identify the known distance and speed.",
                "Use the relationship time = distance ÷ speed.",
                f"Substitute: time = { _fmt_num(d) } ÷ { _fmt_num(s) }.",
                "Divide to get the travel time.",
            ],
            topic="distance_rate_time_basic",
        )

    # speed = distance / time
    if _question_asks_speed(text) and distances and times:
        d = distances[0][0]
        t = times[0]
        if not _safe_positive(t):
            return None
        s = d / t
        return _make_result(
            internal_answer=_fmt_num(s),
            steps=[
                "Identify the known distance and time.",
                "Use the relationship speed = distance ÷ time.",
                f"Substitute: speed = { _fmt_num(d) } ÷ { _fmt_num(t) }.",
                "Divide to get the speed.",
            ],
            topic="distance_rate_time_basic",
        )

    return None


def _average_speed_total_trip_solver(text: str) -> Optional[SolverResult]:
    if "average speed" not in text and "averaged" not in text:
        return None

    distances = _extract_distances(text)
    times = _extract_times_in_hours(text)

    # If total distance and total time are explicitly present
    if _question_asks_speed(text) and distances and times:
        d = sum(x[0] for x in distances)
        t = sum(times)
        if not _safe_positive(t):
            return None
        avg = d / t
        return _make_result(
            internal_answer=_fmt_num(avg),
            steps=[
                "For average speed, do not average the speeds directly.",
                "First find total distance and total time.",
                "Use average speed = total distance ÷ total time.",
                f"Here that means average speed = { _fmt_num(d) } ÷ { _fmt_num(t) }.",
            ],
            topic="distance_rate_time_average_speed",
        )

    # equal-distance round trip average speed from two speeds
    speeds = _extract_speeds(text)
    if len(speeds) >= 2 and _contains_any(text, ["round trip", "same distance", "back", "there and back"]):
        s1, s2 = speeds[0], speeds[1]
        if _safe_positive(s1) and _safe_positive(s2):
            avg = 2 * s1 * s2 / (s1 + s2)
            return _make_result(
                internal_answer=_fmt_num(avg),
                steps=[
                    "This is an equal-distance average-speed setup.",
                    "When the distances are the same, average speed is not the arithmetic mean.",
                    "Use harmonic-mean logic: average speed = 2ab ÷ (a + b).",
                    f"Set a = { _fmt_num(s1) } and b = { _fmt_num(s2) }.",
                ],
                topic="distance_rate_time_average_speed",
            )

    return None


def _meeting_opposite_direction_solver(text: str) -> Optional[SolverResult]:
    if not _contains_any(text, ["opposite", "toward each other", "towards each other", "meet", "meeting"]):
        return None

    distances = _extract_distances(text)
    speeds = _extract_speeds(text)

    # Basic meet: total distance / sum of speeds
    if distances and len(speeds) >= 2 and _question_asks_time(text):
        total_distance = distances[0][0]
        s1, s2 = speeds[0], speeds[1]
        if _safe_positive(s1 + s2):
            t = total_distance / (s1 + s2)
            return _make_result(
                internal_answer=_fmt_num(t),
                steps=[
                    "For motion toward each other, the distances add but the meeting time is the same.",
                    "So use relative speed = speed₁ + speed₂.",
                    f"Relative speed = { _fmt_num(s1) } + { _fmt_num(s2) }.",
                    "Then use time = total distance ÷ relative speed.",
                ],
                topic="distance_rate_time_meeting",
            )

    # If asking distance to meeting point from one side
    if distances and len(speeds) >= 2 and _question_asks_distance(text):
        total_distance = distances[0][0]
        s1, s2 = speeds[0], speeds[1]
        if _safe_positive(s1 + s2):
            t = total_distance / (s1 + s2)
            d1 = s1 * t
            return _make_result(
                internal_answer=_fmt_num(d1),
                steps=[
                    "At the meeting point, both travelers have moved for the same amount of time.",
                    "First find the shared meeting time using total distance ÷ (sum of speeds).",
                    "Then multiply that time by the speed of the traveler asked about.",
                ],
                topic="distance_rate_time_meeting",
            )

    return None


def _overtake_same_direction_solver(text: str) -> Optional[SolverResult]:
    if not _contains_any(text, ["overtake", "overtakes", "catch up", "catches up", "same direction"]):
        return None

    speeds = _extract_speeds(text)
    times = _extract_times_in_hours(text)
    nums = _all_numbers(text)

    # Explicit delayed start and catch-up time with speed difference
    # Example: freight is 20 mph slower, passenger overtakes in 3 hours, freight had 2-hour head start.
    slower_by = None
    m = re.search(rf"{_NUMBER}\s*(?:mph|kmph|m/s|ft/s)?\s*slower", text)
    if m:
        slower_by = _to_float(m.group(1))

    if slower_by is not None and len(times) >= 2 and _question_asks_speed(text):
        # Usually the smaller time is catch-up time; larger is head start or total.
        t_small = min(times)
        t_large = max(times)
        # Need head start and catch-up time; if text says "2 hours after" and "in 3 hours"
        head_start = None
        catch_time = None

        after_match = re.search(rf"{_NUMBER}\s*hours?\s*after", text)
        in_match = re.search(rf"in\s*{_NUMBER}\s*hours?", text)
        if after_match:
            head_start = _to_float(after_match.group(1))
        if in_match:
            catch_time = _to_float(in_match.group(1))

        if head_start is None or catch_time is None:
            if len(times) >= 2:
                head_start = min(times)
                catch_time = max(times)

        if head_start is not None and catch_time is not None:
            # x * catch_time = (x - slower_by) * (catch_time + head_start)
            denom = head_start
            if abs(denom) < 1e-12:
                return None
            x = slower_by * (catch_time + head_start) / head_start
            if _safe_positive(x):
                return _make_result(
                    internal_answer=_fmt_num(x),
                    steps=[
                        "In an overtaking problem, the distances are equal at the catch-up point.",
                        "Let the faster speed be the unknown.",
                        "Express the slower speed using the stated difference.",
                        "Use time carefully: the earlier starter travels longer because of the head start.",
                        "Set distance of faster traveler equal to distance of slower traveler and solve.",
                    ],
                    topic="distance_rate_time_overtake",
                )

    # Generic relative-speed catch-up: time = lead distance / (faster - slower)
    if len(speeds) >= 2 and distances and _question_asks_time(text):
        faster = max(speeds[0], speeds[1])
        slower = min(speeds[0], speeds[1])
        lead_distance = distances[0][0]
        if faster <= slower:
            return None
        t = lead_distance / (faster - slower)
        return _make_result(
            internal_answer=_fmt_num(t),
            steps=[
                "For same-direction catch-up, use relative speed = faster speed - slower speed.",
                f"Relative speed = { _fmt_num(faster) } - { _fmt_num(slower) }.",
                "Then use time = lead distance ÷ relative speed.",
            ],
            topic="distance_rate_time_overtake",
        )

    return None


def _round_trip_solver(text: str) -> Optional[SolverResult]:
    if not _contains_any(text, ["round trip", "returns", "back", "same distance", "there and back"]):
        return None

    speeds = _extract_speeds(text)
    times = _extract_times_in_hours(text)
    distances = _extract_distances(text)

    # Boat/current round trip where current speed given and times given; ask calm-water speed
    if _contains_any(text, ["current", "downstream", "upstream", "against the current", "with the current"]):
        current = None
        m = re.search(rf"current(?: of)?\s*{_NUMBER}\s*(?:mph|kmph|m/s|ft/s)?", text)
        if m:
            current = _to_float(m.group(1))
        else:
            m = re.search(rf"{_NUMBER}\s*(?:mph|kmph|m/s|ft/s)\s*(?:current)", text)
            if m:
                current = _to_float(m.group(1))

        if current is not None and len(times) >= 2 and _question_asks_speed(text):
            t1, t2 = times[0], times[1]
            # equal distances: (b + c)t1 = (b - c)t2
            denom = t2 - t1
            if abs(denom) < 1e-12:
                return None
            b = current * (t1 + t2) / denom
            if _safe_positive(b) and b > current:
                return _make_result(
                    internal_answer=_fmt_num(b),
                    steps=[
                        "This is a same-distance out-and-back problem.",
                        "So the downstream distance equals the upstream distance.",
                        "Let the calm-water speed be the unknown.",
                        "Then downstream speed = boat speed + current and upstream speed = boat speed - current.",
                        "Set the two distances equal and solve.",
                    ],
                    topic="distance_rate_time_stream_current",
                )

    # Same-distance, two different speeds, total time, ask one leg or distance
    if len(speeds) >= 2 and times and _question_asks_distance(text):
        total_time = sum(times)
        s1, s2 = speeds[0], speeds[1]
        if _safe_positive(s1) and _safe_positive(s2):
            one_way_distance = total_time / (1 / s1 + 1 / s2)
            return _make_result(
                internal_answer=_fmt_num(one_way_distance),
                steps=[
                    "For a round trip, the outbound and return distances are the same.",
                    "Write time for each leg as distance ÷ speed.",
                    "Add the two times and match that to the total time.",
                    "Then solve for the one-way distance.",
                ],
                topic="distance_rate_time_round_trip",
            )

    # If total distance of round trip and total time given, ask average speed
    if distances and times and _question_asks_speed(text):
        total_distance = sum(x[0] for x in distances)
        total_time = sum(times)
        if _safe_positive(total_time):
            avg = total_distance / total_time
            return _make_result(
                internal_answer=_fmt_num(avg),
                steps=[
                    "For average speed on a round trip, use total distance ÷ total time.",
                    "Do not average the two speeds directly unless the times are equal, which is not guaranteed.",
                    "Compute the total distance and total time first.",
                ],
                topic="distance_rate_time_round_trip",
            )

    return None


def _stream_current_direct_solver(text: str) -> Optional[SolverResult]:
    if not _contains_any(text, ["current", "stream", "downstream", "upstream"]):
        return None

    speeds = _extract_speeds(text)

    # direct downstream/upstream speeds given -> ask calm water or current
    if len(speeds) >= 2 and _question_asks_speed(text):
        s1, s2 = speeds[0], speeds[1]

        if _contains_any(text, ["calm water", "still water", "boat's speed"]):
            calm = (s1 + s2) / 2
            return _make_result(
                internal_answer=_fmt_num(calm),
                steps=[
                    "For boat/current setups:",
                    "downstream speed = boat speed + current",
                    "upstream speed = boat speed - current",
                    "Add the two equations to isolate twice the calm-water speed.",
                ],
                topic="distance_rate_time_stream_current",
            )

        if _contains_any(text, ["speed of the current", "current speed", "stream speed"]):
            current = abs(s1 - s2) / 2
            return _make_result(
                internal_answer=_fmt_num(current),
                steps=[
                    "For boat/current setups:",
                    "downstream speed = boat speed + current",
                    "upstream speed = boat speed - current",
                    "Subtract the two equations to isolate twice the current speed.",
                ],
                topic="distance_rate_time_stream_current",
            )

    return None


def _average_speed_trick_solver(text: str) -> Optional[SolverResult]:
    # famous impossible target-average setup
    if "average" not in text:
        return None

    lap_match = re.search(rf"{_NUMBER}\s*(?:mile|mi|km|m)\b.*?(?:track|lap)", text)
    target_match = re.search(rf"average\s*{_NUMBER}\s*(?:mph|kmph|m/s|ft/s)", text)
    first_leg_match = re.search(rf"first\s*(?:lap|mile|part).*?{_NUMBER}\s*(?:mph|kmph|m/s|ft/s)", text)

    if not (lap_match and target_match and first_leg_match):
        return None

    dist_each = _to_float(lap_match.group(1))
    target = _to_float(target_match.group(1))
    first_speed = _to_float(first_leg_match.group(1))

    if not (_safe_positive(dist_each) and _safe_positive(target) and _safe_positive(first_speed)):
        return None

    total_distance = 2 * dist_each
    allowed_total_time = total_distance / target
    first_leg_time = dist_each / first_speed

    if first_leg_time >= allowed_total_time - 1e-12:
        return _make_result(
            internal_answer="impossible",
            steps=[
                "This is a target-average-speed trap.",
                "Convert the target average into the maximum total time allowed for the full trip.",
                "Then compare that allowed total time with the time already used on the first leg.",
                "If the first leg already uses all available time (or more), no finite second-leg speed can fix it.",
            ],
            topic="distance_rate_time_average_speed_trick",
        )

    return None


def _component_sum_solver(text: str) -> Optional[SolverResult]:
    # Two-part journey where total distance and total time given, one segment distance asked
    if not (_question_asks_distance(text) and _contains_any(text, ["entire distance", "entire trip", "total distance", "total trip"])):
        return None

    speeds = _extract_speeds(text)
    distances = _extract_distances(text)
    times = _extract_times_in_hours(text)

    if len(speeds) >= 2 and distances and times:
        total_distance = distances[0][0]
        total_time = times[0] if len(times) == 1 else sum(times)

        # Let x be the second segment distance:
        # (total_distance - x)/s1 + x/s2 = total_time
        s1, s2 = speeds[0], speeds[1]
        denom = (1 / s2) - (1 / s1)
        rhs = total_time - (total_distance / s1)
        if abs(denom) < 1e-12:
            return None
        x = rhs / denom
        if math.isfinite(x):
            return _make_result(
                internal_answer=_fmt_num(x),
                steps=[
                    "Break the trip into components.",
                    "Let the unknown segment distance be x, so the other segment is total distance - x.",
                    "Write time for each segment as distance ÷ speed.",
                    "Add the segment times and set that equal to the total trip time.",
                    "Solve the resulting linear equation.",
                ],
                topic="distance_rate_time_components",
            )

    return None


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

    if not _contains_any(
        lower,
        [
            "distance", "speed", "time", "rate", "mph", "kmph", "m/s", "ft/s",
            "meet", "meeting", "overtake", "catch up", "current", "stream",
            "upstream", "downstream", "round trip", "average speed", "lap"
        ],
    ):
        return None

    solvers = [
        _average_speed_trick_solver,
        _stream_current_direct_solver,
        _round_trip_solver,
        _meeting_opposite_direction_solver,
        _overtake_same_direction_solver,
        _component_sum_solver,
        _average_speed_total_trip_solver,
        _basic_formula_solver,
    ]

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

    return None