solver_combinatorics.py

from __future__ import annotations

import math
import re
from math import comb, factorial, perm
from typing import Optional

from models import SolverResult


# ----------------------------
# Core helpers
# ----------------------------

def _clean(text: str) -> str:
    t = (text or "").strip()
    t = t.replace("×", "x")
    t = t.replace("–", "-").replace("—", "-")
    t = t.replace("“", '"').replace("”", '"')
    t = t.replace("’", "'")
    t = re.sub(r"\s+", " ", t)
    return t


def _lower(text: str) -> str:
    return _clean(text).lower()


def _safe_comb(n: int, r: int) -> Optional[int]:
    if n < 0 or r < 0 or r > n:
        return None
    return comb(n, r)


def _safe_perm(n: int, r: int) -> Optional[int]:
    if n < 0 or r < 0 or r > n:
        return None
    return perm(n, r)


def _fact(n: int) -> Optional[int]:
    if n < 0:
        return None
    return factorial(n)


def _make_result(topic: str, answer: int, steps: list[str]) -> SolverResult:
    # Keep numeric answer internal / structured.
    # The outer reply layer can decide whether to reveal it.
    return SolverResult(
        domain="quant",
        solved=True,
        topic=topic,
        answer_value=str(answer),
        internal_answer=str(answer),
        steps=steps,
    )


def _numbers(text: str) -> list[int]:
    return [int(x) for x in re.findall(r"\d+", text)]


def _has_any(text: str, words: list[str]) -> bool:
    return any(w in text for w in words)


def _word_frequency_from_quoted_or_caps(raw: str) -> Optional[dict[str, int]]:
    """
    Tries to detect repeated-letter arrangement prompts:
    - letters of "BALLOON"
    - word MISSISSIPPI
    - arrange the letters in BOOKKEEPER
    """
    # quoted word
    m = re.search(r'letters?\s+of\s+"([A-Za-z]+)"', raw, re.I)
    if not m:
        m = re.search(r"word\s+([A-Za-z]+)", raw, re.I)
    if not m:
        m = re.search(r'arrange\s+the\s+letters?\s+(?:in|of)\s+"?([A-Za-z]+)"?', raw, re.I)

    if not m:
        return None

    word = m.group(1).strip()
    if not word.isalpha() or len(word) < 2:
        return None

    freq: dict[str, int] = {}
    for ch in word.upper():
        freq[ch] = freq.get(ch, 0) + 1
    return freq


def _multiset_permutations(freq: dict[str, int]) -> int:
    total = sum(freq.values())
    out = factorial(total)
    for c in freq.values():
        out //= factorial(c)
    return out


def _extract_n_r_from_choose_style(lower: str) -> Optional[tuple[int, int]]:
    patterns = [
        r"(?:choose|select|pick|form|make)\s+(\d+)\s+(?:from|out of)\s+(\d+)",
        r"(?:choose|select|pick)\s+(\d+)\s+of\s+(\d+)",
        r"(?:committee|team|group|delegation|panel)\s+of\s+(\d+).*?(?:from|out of)\s+(\d+)",
        r"(\d+)\s+(?:chosen|selected|picked)\s+from\s+(\d+)",
        r"how many combinations.*?(\d+).*?(?:from|out of)\s+(\d+)",
        r"n\s*=\s*(\d+)\s*,?\s*r\s*=\s*(\d+)",  # only if combo cues present elsewhere
    ]

    for pat in patterns:
        m = re.search(pat, lower)
        if m:
            a, b = int(m.group(1)), int(m.group(2))
            # for n=,r= pattern we want n,r order
            if "n=" in pat:
                return a, b
            return b, a
    return None


def _extract_n_r_from_perm_style(lower: str) -> Optional[tuple[int, int]]:
    patterns = [
        r"(?:arrange|order|rank|assign|seat)\s+(\d+)\s+(?:from|out of)\s+(\d+)",
        r"(?:choose|select|pick)\s+(\d+)\s+(?:from|out of)\s+(\d+)\s+(?:and|then)\s+(?:arrange|order|rank)",
        r"permutations?\s+of\s+(\d+)\s+(?:from|out of)\s+(\d+)",
        r"(?:how many ways|number of ways).*?(?:arrange|order|rank).*?(\d+).*?(?:from|out of)\s+(\d+)",
        r"(\d+)\s+(?:positions|places|slots).*?(?:filled|chosen).*?(?:from|out of)\s+(\d+)",
    ]

    for pat in patterns:
        m = re.search(pat, lower)
        if m:
            r, n = int(m.group(1)), int(m.group(2))
            return n, r
    return None


def _extract_total_distinct_arrangement_n(lower: str) -> Optional[int]:
    patterns = [
        r"(?:arrange|order|permute|seat)\s+(\d+)\s+(?:different|distinct)?\s*(?:objects|items|people|books|letters|marbles|students|digits)?",
        r"permutations?\s+of\s+(\d+)\s+(?:different|distinct)?",
        r"arrangements?\s+of\s+(\d+)\s+(?:different|distinct)?",
    ]
    for pat in patterns:
        m = re.search(pat, lower)
        if m:
            return int(m.group(1))
    return None


def _extract_circular_n(lower: str) -> Optional[int]:
    patterns = [
        r"(?:around|in)\s+a\s+circle.*?(\d+)",
        r"circular arrangements?.*?(\d+)",
        r"seat\s+(\d+).*?(?:around|in)\s+a\s+(?:round\s+)?table",
        r"arrange\s+(\d+).*?(?:around|in)\s+a\s+circle",
    ]
    for pat in patterns:
        m = re.search(pat, lower)
        if m:
            return int(m.group(1))
    return None


def _extract_required_excluded_committee(lower: str) -> Optional[tuple[int, int, int, int]]:
    """
    Returns (n_total, r_choose, required_count, excluded_count)
    for patterns like:
    - committee of 4 from 10 if A must be included
    - choose 3 from 8 if 2 specific people are excluded
    """
    base = _extract_n_r_from_choose_style(lower)
    if not base:
        return None

    n, r = base
    required = 0
    excluded = 0

    # specific people forced in
    req_patterns = [
        r"(?:must|has to|have to|should)\s+be\s+included",
        r"including\s+\w+",
        r"include\s+\w+",
        r"with\s+\w+\s+included",
    ]
    exc_patterns = [
        r"(?:must|has to|have to)\s+not\s+be\s+included",
        r"excluding\s+\w+",
        r"exclude\s+\w+",
        r"without\s+\w+",
    ]

    if any(re.search(p, lower) for p in req_patterns):
        required = 1
    if any(re.search(p, lower) for p in exc_patterns):
        excluded = 1

    m = re.search(r"(\d+)\s+specific\s+(?:people|members|students|persons)\s+(?:must|have to)\s+be\s+included", lower)
    if m:
        required = int(m.group(1))

    m = re.search(r"(\d+)\s+specific\s+(?:people|members|students|persons)\s+(?:must|have to)\s+be\s+excluded", lower)
    if m:
        excluded = int(m.group(1))

    if required == 0 and excluded == 0:
        return None

    return n, r, required, excluded


def _extract_group_selection(lower: str) -> Optional[tuple[int, int, int, int, str]]:
    """
    Handles common grouped committee/team cases:
    - 3 men and 2 women from 7 men and 5 women
    - at least 2 women from 6 men and 4 women for a committee of 4
    Returns:
      (a_total, b_total, choose_total, threshold, mode)
    mode in {"exact_a", "at_least_a"}
    Here group A is first-mentioned category.
    """
    # exact split: x men and y women from M men and W women
    m = re.search(
        r"(\d+)\s+\w+\s+and\s+(\d+)\s+\w+.*?(?:from|out of)\s+(\d+)\s+\w+\s+and\s+(\d+)\s+\w+",
        lower,
    )
    if m:
        a_choose = int(m.group(1))
        b_choose = int(m.group(2))
        a_total = int(m.group(3))
        b_total = int(m.group(4))
        # pack choose_total as first component sum, threshold as a_choose,
        # mode exact_a means choose threshold from first group and remainder from second
        return a_total, b_total, a_choose + b_choose, a_choose, "exact_a"

    # at least k from first group
    m = re.search(
        r"at least\s+(\d+)\s+\w+.*?(?:committee|team|group|delegation|selection)\s+of\s+(\d+).*?(?:from|out of)\s+(\d+)\s+\w+\s+and\s+(\d+)\s+\w+",
        lower,
    )
    if m:
        threshold = int(m.group(1))
        choose_total = int(m.group(2))
        a_total = int(m.group(3))
        b_total = int(m.group(4))
        return a_total, b_total, choose_total, threshold, "at_least_a"

    return None


def _extract_adjacent_block_n_k(lower: str) -> Optional[tuple[int, int]]:
    """
    Tries to detect:
    - among n distinct objects, k specific objects must be together
    - A and B must sit together among n people
    Returns (n_total_distinct, block_size)
    """
    # count named items joined by and
    m = re.search(
        r"(\d+)\s+(?:distinct|different)?\s*(?:people|students|books|letters|objects|items|marbles).*?(?:must|have to)\s+be\s+together",
        lower,
    )
    if m:
        n = int(m.group(1))
        # if text explicitly lists "a and b", treat as 2
        pair = re.search(r"\b\w+\s+and\s+\w+\b.*?(?:must|have to)\s+be\s+together", lower)
        if pair:
            return n, 2

    # more direct "2 specific people must sit together among 7"
    m = re.search(
        r"(\d+)\s+specific\s+(?:people|students|books|letters|objects|items).*?(?:must|have to)\s+be\s+together.*?(?:among|from|out of)\s+(\d+)",
        lower,
    )
    if m:
        k = int(m.group(1))
        n = int(m.group(2))
        return n, k

    # "A and B together among 6 people"
    m = re.search(
        r"\b\w+\s+and\s+\w+\b.*?(?:together|next to each other|adjacent).*?(?:among|out of|from)\s+(\d+)",
        lower,
    )
    if m:
        n = int(m.group(1))
        return n, 2

    # fallback: if exactly one overall n and pair-together cue exists
    nums = _numbers(lower)
    if len(nums) == 1 and _has_any(lower, ["together", "next to each other", "adjacent"]):
        return nums[0], 2

    return None


def _extract_not_together_pair(lower: str) -> Optional[int]:
    if not _has_any(lower, ["not together", "not adjacent", "not next to each other", "cannot sit together"]):
        return None
    nums = _numbers(lower)
    if nums:
        return nums[-1] if len(nums) == 1 else max(nums)
    return None


def _extract_relative_order_n(lower: str) -> Optional[int]:
    """
    Detect simple relative-order cases:
    - A left of B
    - A before B
    - A ahead of B
    among n distinct objects/people
    """
    if not _has_any(lower, ["left of", "before", "ahead of", "to the left of"]):
        return None
    nums = _numbers(lower)
    if nums:
        return nums[-1] if len(nums) == 1 else max(nums)
    return None


def _extract_stars_bars(lower: str) -> Optional[tuple[int, int, bool]]:
    """
    Detect nonnegative / positive integer solution counts:
    - number of nonnegative integer solutions to x+y+z=10
    - positive integer solutions to a+b+c=12
    Returns (n_sum, k_vars, positive_required)
    """
    m = re.search(
        r"(nonnegative|positive)\s+integer\s+solutions?.*?to\s+([a-z](?:\s*\+\s*[a-z])+)\s*=\s*(\d+)",
        lower,
    )
    if not m:
        return None

    positivity = m.group(1)
    vars_expr = m.group(2)
    total = int(m.group(3))
    vars_count = len(re.findall(r"[a-z]", vars_expr))
    return total, vars_count, positivity == "positive"


def _extract_digit_codes(lower: str) -> Optional[tuple[int, bool, bool]]:
    """
    Handles password/code/number formation style:
    Returns (length, repetition_allowed, starts_with_zero_allowed_guess)
    """
    if not _has_any(lower, ["digit number", "digits", "code", "password", "license plate", "pin"]):
        return None

    m = re.search(r"(\d+)[-\s]digit", lower)
    if not m:
        m = re.search(r"code\s+of\s+length\s+(\d+)", lower)
    if not m:
        return None

    length = int(m.group(1))
    repetition_allowed = _has_any(lower, ["repetition allowed", "can repeat", "may repeat", "with repetition"])
    starts_zero_allowed = not _has_any(lower, ["first digit cannot be 0", "leading zero not allowed", "cannot start with 0"])

    return length, repetition_allowed, starts_zero_allowed


def _extract_available_digits(lower: str) -> Optional[int]:
    m = re.search(r"from\s+(\d+)\s+(?:digits|numbers)", lower)
    if m:
        return int(m.group(1))
    # decimal digit default
    if _has_any(lower, ["digit number", "digits", "code", "password", "pin"]):
        return 10
    return None


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

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

    if not raw:
        return None

    combinatorics_cues = [
        "combination", "combinations", "permutation", "permutations",
        "arrange", "arrangement", "arrangements", "order", "ordered",
        "choose", "select", "pick", "committee", "team", "group",
        "delegation", "panel", "seat", "seating", "circular", "circle",
        "round table", "together", "adjacent", "left of", "before",
        "anagram", "letters of", "word", "integer solutions", "password",
        "license plate", "pin", "code", "nonnegative integer", "positive integer",
    ]
    if not _has_any(lower, combinatorics_cues):
        return None

    # ----------------------------------------
    # 1) Repeated-letter arrangements / anagrams
    # ----------------------------------------
    freq = _word_frequency_from_quoted_or_caps(raw)
    if freq and _has_any(lower, ["letters", "word", "arrange", "anagram", "distinct arrangements"]):
        result = _multiset_permutations(freq)
        return _make_result(
            topic="combinatorics",
            answer=result,
            steps=[
                "Treat this as an arrangement of letters with repetition.",
                "Start with the factorial of the total number of letters.",
                "Then divide by factorials of each repeated-letter count so identical rearrangements are not overcounted.",
                "Evaluate that expression at the end.",
            ],
        )

    # ----------------------------------------
    # 2) Circular arrangements
    # ----------------------------------------
    n_circle = _extract_circular_n(lower)
    if n_circle is not None and n_circle >= 1:
        result = 1 if n_circle == 1 else factorial(n_circle - 1)
        return _make_result(
            topic="combinatorics",
            answer=result,
            steps=[
                "For circular arrangements of distinct objects, rotations count as the same arrangement.",
                "Fix one object to remove rotational duplicates.",
                "Then arrange the remaining objects in the remaining positions.",
                "Evaluate the factorial expression at the end.",
            ],
        )

    # ----------------------------------------
    # 3) Integer-solution counting (stars and bars)
    # ----------------------------------------
    stars = _extract_stars_bars(lower)
    if stars:
        total, vars_count, positive_required = stars
        if positive_required:
            adjusted = total - vars_count
            if adjusted < 0:
                return _make_result(
                    topic="combinatorics",
                    answer=0,
                    steps=[
                        "For positive integer solutions, first give each variable 1.",
                        "That reduces the remaining amount to distribute.",
                        "If the remaining amount is negative, no valid solutions exist.",
                    ],
                )
            result = comb(adjusted + vars_count - 1, vars_count - 1)
            return _make_result(
                topic="combinatorics",
                answer=result,
                steps=[
                    "This is a stars-and-bars counting problem with positive integers.",
                    "Give each variable 1 first so the positivity condition is satisfied.",
                    "Then count the nonnegative distributions of the remaining total among the variables.",
                    "Evaluate the resulting combination.",
                ],
            )
        else:
            result = comb(total + vars_count - 1, vars_count - 1)
            return _make_result(
                topic="combinatorics",
                answer=result,
                steps=[
                    "This is a stars-and-bars counting problem with nonnegative integers.",
                    "Interpret the total as stars and the separators between variables as bars.",
                    "Count the placements of the bars among the stars.",
                    "Evaluate the resulting combination.",
                ],
            )

    # ----------------------------------------
    # 4) Exact grouped selections
    # ----------------------------------------
    grouped = _extract_group_selection(lower)
    if grouped:
        a_total, b_total, choose_total, threshold, mode = grouped
        if mode == "exact_a":
            a_choose = threshold
            b_choose = choose_total - a_choose
            left = _safe_comb(a_total, a_choose)
            right = _safe_comb(b_total, b_choose)
            if left is not None and right is not None:
                result = left * right
                return _make_result(
                    topic="combinations",
                    answer=result,
                    steps=[
                        "Split the selection by category.",
                        "Choose the required number from the first group.",
                        "Choose the remaining number from the second group.",
                        "Multiply the independent counts.",
                    ],
                )

        if mode == "at_least_a":
            total_count = 0
            valid = False
            for a_choose in range(threshold, choose_total + 1):
                b_choose = choose_total - a_choose
                left = _safe_comb(a_total, a_choose)
                right = _safe_comb(b_total, b_choose)
                if left is None or right is None:
                    continue
                total_count += left * right
                valid = True
            if valid:
                return _make_result(
                    topic="combinations",
                    answer=total_count,
                    steps=[
                        "Break the problem into valid cases based on how many can come from the first group.",
                        "For each valid case, choose from the first group and choose the remainder from the second group.",
                        "Add the case counts together.",
                    ],
                )

    # ----------------------------------------
    # 5) Committee / selection with required or excluded members
    # ----------------------------------------
    req_exc = _extract_required_excluded_committee(lower)
    if req_exc:
        n, r, required, excluded = req_exc

        # required members first
        if required > 0 and excluded == 0:
            remaining_n = n - required
            remaining_r = r - required
            result = _safe_comb(remaining_n, remaining_r)
            if result is not None:
                return _make_result(
                    topic="combinations",
                    answer=result,
                    steps=[
                        "Treat the required members as already chosen.",
                        "Reduce both the total pool and the number still to be selected.",
                        "Then count the remaining selection with combinations.",
                    ],
                )

        # excluded members first
        if excluded > 0 and required == 0:
            remaining_n = n - excluded
            result = _safe_comb(remaining_n, r)
            if result is not None:
                return _make_result(
                    topic="combinations",
                    answer=result,
                    steps=[
                        "Remove the excluded members from the pool first.",
                        "Then choose the full committee from the reduced pool.",
                    ],
                )

        # both required and excluded
        remaining_n = n - required - excluded
        remaining_r = r - required
        result = _safe_comb(remaining_n, remaining_r)
        if result is not None:
            return _make_result(
                topic="combinations",
                answer=result,
                steps=[
                    "Force the required members into the selection.",
                    "Remove the excluded members from the pool.",
                    "Then choose the remaining spots from the remaining eligible people.",
                ],
            )

    # ----------------------------------------
    # 6) Pair not together / not adjacent
    # ----------------------------------------
    n_not_together = _extract_not_together_pair(lower)
    if n_not_together is not None and n_not_together >= 2:
        total = factorial(n_not_together)
        together = 2 * factorial(n_not_together - 1)
        result = total - together
        return _make_result(
            topic="permutations",
            answer=result,
            steps=[
                "Count all unrestricted arrangements first.",
                "Then count the arrangements where the two specified objects stay together by treating them as one block.",
                "Because the two objects can switch places inside the block, include that internal ordering factor.",
                "Subtract the together-case count from the total.",
            ],
        )

    # ----------------------------------------
    # 7) Adjacent/together block arrangements
    # ----------------------------------------
    block = _extract_adjacent_block_n_k(lower)
    if block:
        n, k = block
        if n >= k >= 2:
            result = factorial(n - k + 1) * factorial(k)
            return _make_result(
                topic="permutations",
                answer=result,
                steps=[
                    "Treat the required adjacent objects as one block.",
                    "Count the arrangements of that block together with the remaining distinct objects.",
                    "Then multiply by the internal arrangements of the objects inside the block.",
                ],
            )

    # ----------------------------------------
    # 8) Relative order: A before B / left of B
    # ----------------------------------------
    n_order = _extract_relative_order_n(lower)
    if n_order is not None and n_order >= 2:
        result = factorial(n_order) // 2
        return _make_result(
            topic="permutations",
            answer=result,
            steps=[
                "Start from all arrangements of the distinct objects.",
                "For any arrangement, swapping the two named objects reverses their relative order.",
                "So exactly half of all arrangements satisfy the required order condition.",
            ],
        )

    # ----------------------------------------
    # 9) Exact combination selection nCr
    # ----------------------------------------
    choose_nr = _extract_n_r_from_choose_style(lower)
    if choose_nr and _has_any(lower, ["choose", "select", "pick", "committee", "team", "group", "delegation", "panel"]):
        n, r = choose_nr
        result = _safe_comb(n, r)
        if result is not None:
            return _make_result(
                topic="combinations",
                answer=result,
                steps=[
                    "This is a selection problem where order does not matter.",
                    "Use combinations: choose the required number from the total pool.",
                    "Set up the combination expression and evaluate it at the end.",
                ],
            )

    # ----------------------------------------
    # 10) Exact permutation / ordered selection nPr
    # ----------------------------------------
    perm_nr = _extract_n_r_from_perm_style(lower)
    if perm_nr:
        n, r = perm_nr
        result = _safe_perm(n, r)
        if result is not None:
            return _make_result(
                topic="permutations",
                answer=result,
                steps=[
                    "This is an ordered selection problem, so order matters.",
                    "Choose and arrange the required number of positions from the total available objects.",
                    "Set up the permutation expression and evaluate it at the end.",
                ],
            )

    # ----------------------------------------
    # 11) Full arrangement of n distinct objects
    # ----------------------------------------
    if _has_any(lower, ["arrange", "arrangement", "arrangements", "permutation", "permutations", "seat", "seating", "order"]):
        n_all = _extract_total_distinct_arrangement_n(lower)
        if n_all is not None:
            result = _fact(n_all)
            if result is not None:
                return _make_result(
                    topic="permutations",
                    answer=result,
                    steps=[
                        "All distinct objects are being arranged.",
                        "Fill positions one by one: first position, second position, and so on.",
                        "That leads to the factorial count for all distinct arrangements.",
                    ],
                )

    # ----------------------------------------
    # 12) Digits / codes / passwords
    # ----------------------------------------
    digit_case = _extract_digit_codes(lower)
    if digit_case:
        length, repetition_allowed, starts_zero_allowed = digit_case
        available = _extract_available_digits(lower) or 10

        if repetition_allowed:
            if starts_zero_allowed:
                result = available ** length
            else:
                result = (available - 1) * (available ** (length - 1))
            return _make_result(
                topic="combinatorics",
                answer=result,
                steps=[
                    "Treat each position independently.",
                    "Use the allowed number of choices for the first position, paying attention to any leading-zero restriction.",
                    "Then multiply by the allowed choices for each remaining position.",
                ],
            )

        # no repetition allowed
        if length > available:
            return _make_result(
                topic="combinatorics",
                answer=0,
                steps=[
                    "Without repetition, you cannot fill more positions than the number of available distinct digits.",
                    "So this setup has no valid arrangements.",
                ],
            )

        if starts_zero_allowed:
            result = perm(available, length)
        else:
            result = (available - 1) * perm(available - 1, length - 1)

        return _make_result(
            topic="combinatorics",
            answer=result,
            steps=[
                "This is an ordered arrangement of digits.",
                "Because repetition is not allowed, the number of choices decreases from position to position.",
                "Handle the first digit separately if leading zero is not allowed.",
                "Then multiply the sequential choices.",
            ],
        )

    # ----------------------------------------
    # 13) Direct nCr / nPr notation
    # ----------------------------------------
    m = re.search(r"(\d+)\s*c\s*(\d+)", lower)
    if m:
        n, r = int(m.group(1)), int(m.group(2))
        result = _safe_comb(n, r)
        if result is not None:
            return _make_result(
                topic="combinations",
                answer=result,
                steps=[
                    "Interpret the notation as a combination count.",
                    "Evaluate the combination expression carefully.",
                ],
            )

    m = re.search(r"(\d+)\s*p\s*(\d+)", lower)
    if m:
        n, r = int(m.group(1)), int(m.group(2))
        result = _safe_perm(n, r)
        if result is not None:
            return _make_result(
                topic="permutations",
                answer=result,
                steps=[
                    "Interpret the notation as a permutation count.",
                    "Evaluate the permutation expression carefully.",
                ],
            )

    return None