src/solver.py

"""Deterministic answer extraction and response normalization.

Bypasses the model for questions where the answer can be computed
directly from retrieved evidence (exon counts, domain matching, NMD
prediction, variant-specific therapy lookups).
"""

import re
import time
from difflib import SequenceMatcher

from src.equivalence import canonical_drug_name, find_equivalent_in_choices
from src.dags import clinical_trials as ct_dag
from src.dags import variant_assessment as va_dag


def parse_answer_choices(question_text: str) -> list[str]:
    m = re.search(r"answer choices:\s*(.+)$", question_text, flags=re.IGNORECASE)
    if not m:
        return []
    raw = m.group(1).strip().rstrip(".")
    return [c.strip() for c in raw.split(",") if c.strip()]


def _best_choice_match(text: str, choices: list[str]) -> str | None:
    if not choices:
        return None

    text_norm = text.strip().lower()
    for c in choices:
        if text_norm == c.strip().lower():
            return c
    for c in choices:
        c_norm = c.strip().lower()
        if c_norm in text_norm or text_norm in c_norm:
            return c

    scored = [(SequenceMatcher(None, text_norm, c.strip().lower()).ratio(), c) for c in choices]
    scored.sort(reverse=True, key=lambda x: x[0])
    if scored and scored[0][0] >= 0.55:
        return scored[0][1]
    return None


def _extract_numeric(text: str) -> str | None:
    m = re.search(r"(-?\d+(?:\.\d+)?)", text)
    return m.group(1) if m else None


def normalize_response(response: str, input_data: dict, evidence: list[dict]) -> str:
    answer_format = input_data["question"]["answer_format"]
    question_text = input_data["question"]["prompt"]
    text = response.strip().strip('"').strip("'").strip()

    if answer_format == "binary":
        low = text.lower()
        if low in {"yes", "y", "true", "eligible"}:
            return "Yes"
        if low in {"no", "n", "false", "ineligible"}:
            return "No"
        if "yes" in low:
            return "Yes"
        if "no" in low:
            return "No"
        return text

    if answer_format == "numeric_match":
        n = _extract_numeric(text)
        return n if n is not None else text

    if answer_format == "string_match":
        if "clinical trial" in question_text.lower() or "nct" in question_text.lower():
            trial = re.search(r"\bNCT\d{6,}\b", text, flags=re.IGNORECASE)
            if trial:
                return trial.group(0).upper()
            for e in evidence:
                m = re.search(r"\bNCT\d{6,}\b", e.get("snippet", ""), flags=re.IGNORECASE)
                if m:
                    return m.group(0).upper()
        return text

    if answer_format == "multiple_choice":
        choices = parse_answer_choices(question_text)
        matched = _best_choice_match(text, choices)
        if matched:
            return matched
        # Equivalence fallback: did the model produce a clinically equivalent drug?
        eq = find_equivalent_in_choices(text, choices)
        if eq:
            return eq
        return text

    return text


def _extract_variant_position(input_data: dict) -> int | None:
    prompt = input_data["question"]["prompt"]
    m = re.search(r"position\s+(\d+)", prompt, flags=re.IGNORECASE)
    if m:
        return int(m.group(1))
    protein = input_data["patient"]["genotype"][0].get("variant_protein", "")
    m = re.search(r"p\.\(?[A-Za-z*]+(\d+)", protein)
    if m:
        return int(m.group(1))
    return None


def _parse_uniprot_domains(snippet: str) -> list[tuple[str, int, int]]:
    domains = []
    for m in re.finditer(r"(?:Domain|Region):\s*(.*?)\s*\(aa\s*(\d+)-(\d+)\)", snippet):
        domains.append((m.group(1).strip(), int(m.group(2)), int(m.group(3))))
    return domains


def try_direct_answer(input_data: dict, evidence: list[dict]) -> dict | None:
    """Return a deterministic answer if one can be computed from evidence."""
    category = input_data["question"]["category"]
    answer_format = input_data["question"]["answer_format"]
    prompt = input_data["question"]["prompt"].lower()

    # ----------------------------------------------------------------
    # Tier 1.5: Reasoning DAGs for Clinical_Trials and Variant_Assessment —
    # explicit sub-question decomposition and composition.
    # ----------------------------------------------------------------

    # Clinical_Trials DAG
    if category == "Clinical_Trials":
        result = ct_dag.evaluate_with_meta(input_data, evidence)
        if result.resolved:
            return _make_payload(result.value, _stub_evidence(result),
                                 result.justification[:500])

    # Variant_Assessment: try DAGs in priority order
    if category == "Variant_Assessment":
        # Numeric questions (amino-acid count / fraction)
        if answer_format == "numeric_match":
            r = va_dag.NUMERIC_DAG.evaluate(input_data, evidence)
            if r.resolved:
                return _make_payload(r.value, _stub_evidence(r),
                                     r.justification[:500])
        # NMD binary
        if answer_format == "binary" and "nonsense mediated decay" in prompt:
            r = va_dag.NMD_DAG.evaluate(input_data, evidence)
            if r.resolved:
                return _make_payload(r.value, _stub_evidence(r),
                                     r.justification[:500])
        # Domain matching
        if answer_format == "multiple_choice":
            r = va_dag.DOMAIN_DAG.evaluate(input_data, evidence)
            if r.resolved:
                return _make_payload(r.value, _stub_evidence(r),
                                     r.justification[:500])

    # Pre-computed variant-specific therapy lookups
    for e in evidence:
        if e.get("source_name") == "Variant_Therapy_Lookup":
            answer = e.get("_precomputed_answer", "")
            if answer:
                if answer_format == "multiple_choice":
                    choices = parse_answer_choices(input_data["question"]["prompt"])
                    if choices:
                        matched = _best_choice_match(answer, choices)
                        if matched:
                            answer = matched
                return _make_payload(answer, e, e.get("_mechanism", "")[:500])

    # Pre-computed supportive-care lookups (EDS-overlap connective tissue disorders)
    for e in evidence:
        if e.get("source_name") == "Supportive_Care_Lookup":
            answer = e.get("_precomputed_answer", "")
            if answer:
                return _make_payload(answer, e, e.get("_mechanism", "")[:500])

    # Clinical Trials: deterministic eligibility filter.
    # If every retrieved trial fails at least one eligibility gate (not testing a
    # new therapeutic, observational only, or patient age out of range), the
    # answer is "None".
    if (category == "Clinical_Trials"
            and answer_format == "string_match"):
        trial_evidence = [e for e in evidence if e.get("source_name") == "ClinicalTrials.gov"]
        asks_new_therapeutic = any(kw in prompt for kw in (
            "new therapeutic", "new treatment", "novel therap",
            "new therapy", "investigational",
        ))

        def _eligible(e: dict) -> bool:
            if asks_new_therapeutic and not e.get("_is_new_therapeutic", True):
                return False
            if asks_new_therapeutic and e.get("_is_observational", False):
                return False
            if not e.get("_age_eligible", True):
                return False
            return True

        if trial_evidence:
            eligible_trials = [e for e in trial_evidence if _eligible(e)]
            if not eligible_trials:
                reasons = []
                if asks_new_therapeutic and all(not e.get("_is_new_therapeutic", True) for e in trial_evidence):
                    reasons.append("all retrieved trials test only already-FDA-approved drugs")
                if all(e.get("_is_observational", False) for e in trial_evidence) and asks_new_therapeutic:
                    reasons.append("all retrieved trials are observational")
                if all(not e.get("_age_eligible", True) for e in trial_evidence):
                    reasons.append("patient age outside trial enrollment window")
                why = "; ".join(reasons) if reasons else "no retrieved trial meets eligibility criteria"
                return _make_payload(
                    "None", trial_evidence[0],
                    f"No eligible trial: {why}.",
                )

    # Pre-computed DMD exon-skipping lookup
    for e in evidence:
        if e.get("source_name") == "DMD_ExonSkip_Lookup":
            dmd = e.get("_dmd_result", {})
            drug = dmd.get("amenable_drug", "")
            if drug:
                if answer_format == "multiple_choice":
                    choices = parse_answer_choices(input_data["question"]["prompt"])
                    if choices:
                        matched = _best_choice_match(drug, choices)
                        if matched:
                            drug = matched
                return _make_payload(drug, e, dmd.get("mechanism", "")[:500])

    if category != "Variant_Assessment":
        return None

    # Extract answers from Ensembl evidence
    for e in evidence:
        snippet = e.get("snippet", "")
        low = snippet.lower()

        if answer_format == "numeric_match":
            if "amino acids" in prompt:
                m = re.search(r"coding\s+(\d+)\s+amino acids", low)
                if m:
                    return _make_payload(m.group(1), e, "Amino acid count from Ensembl exon mapping.")
            if "percentage" in prompt or "decimal" in prompt or "fraction" in prompt:
                m = re.search(r"fraction of total protein:\s*\d+/\d+\s*=\s*([0-9]*\.?[0-9]+)", low)
                if m:
                    return _make_payload(m.group(1), e, "Exon fraction from Ensembl evidence.")

        if answer_format == "binary" and "nonsense mediated decay" in prompt:
            if "trigger nmd" in low:
                return _make_payload("Yes", e, "Variant in early exon predicted to trigger NMD.")
            if "escape nmd" in low:
                return _make_payload("No", e, "Variant in last/penultimate exon predicted to escape NMD.")

    # Domain matching from UniProt evidence.
    #
    # When multiple UniProt annotations contain the variant position and several
    # appear among the answer choices, select with this preference order:
    #   1. Among choice-matching candidates (score >= 0.85), prefer the earliest
    #      start position. UniProt returns features sorted by start, and the
    #      earliest annotation containing a position is typically its most
    #      specific canonical functional unit.
    #   2. Tiebreaker: prefer typed `Domain` over typed `Region` (typed Domain
    #      annotations are the most authoritative functional units).
    #   3. Tiebreaker: prefer the broader span (larger end-start) — broader
    #      structural classifications often supersede narrower sub-features.
    if answer_format == "multiple_choice":
        choices = parse_answer_choices(input_data["question"]["prompt"])
        if choices:
            position = _extract_variant_position(input_data)

            for e in evidence:
                if e.get("source_name") != "UniProt":
                    continue
                snippet = e.get("snippet", "")

                # Extract ALL domains and their type from the snippet.
                # _parse_uniprot_domains returns (name, start, end). We need
                # type too — re-parse with the typed regex.
                typed = re.findall(
                    r"(Domain|Region):\s*(.*?)\s*\(aa\s*(\d+)-(\d+)\)",
                    snippet,
                )
                if not typed and position is None:
                    continue

                # Build the typed list and filter to those containing the position
                if position is not None:
                    typed_containing = [
                        {"type": t, "name": n.strip(), "start": int(s), "end": int(e_)}
                        for t, n, s, e_ in typed
                        if int(s) <= position <= int(e_)
                    ]
                else:
                    typed_containing = []

                if typed_containing:
                    # Score each containing annotation against every answer choice.
                    # Keep only "candidates" with a strong match (>= 0.85).
                    candidates = []
                    for ann in typed_containing:
                        for choice in choices:
                            score = SequenceMatcher(
                                None, ann["name"].lower(), choice.lower()
                            ).ratio()
                            if score >= 0.85:
                                candidates.append({
                                    "ann": ann,
                                    "choice": choice,
                                    "score": score,
                                })

                    if candidates:
                        # When both a typed Domain AND a Region match the answer
                        # choices, AND the Domain is much broader than the Region
                        # (>2x span), prefer the Region — it's a more specific
                        # functional sub-feature within the broader Domain.
                        # (Handles LMNA-like cases: IF rod 356bp Domain vs Coil 1B
                        # 137bp Region; the Region is the right answer.)
                        domain_cands = [c for c in candidates if c["ann"]["type"] == "Domain"]
                        region_cands = [c for c in candidates if c["ann"]["type"] == "Region"]
                        if domain_cands and region_cands:
                            d_span = domain_cands[0]["ann"]["end"] - domain_cands[0]["ann"]["start"]
                            r_span = region_cands[0]["ann"]["end"] - region_cands[0]["ann"]["start"]
                            if d_span > 2 * r_span:
                                # Prefer the more-specific Region
                                region_cands.sort(key=lambda c: (
                                    c["ann"]["start"],
                                    -(c["ann"]["end"] - c["ann"]["start"]),
                                    -c["score"],
                                ))
                                best = region_cands[0]
                                return _make_payload(
                                    best["choice"], e,
                                    f"Position {position} maps to Region {best['ann']['name']} "
                                    f"(aa {best['ann']['start']}-{best['ann']['end']}). "
                                    f"Region selected as more specific than the broader containing Domain.",
                                )

                        # Default ordering:
                        #   1. Earliest start position
                        #   2. Domain type: typed `Domain` before `Region`
                        #   3. Larger span (broader category)
                        #   4. Higher string-match score
                        type_rank = {"Domain": 0, "Region": 1}
                        candidates.sort(key=lambda c: (
                            c["ann"]["start"],
                            type_rank.get(c["ann"]["type"], 9),
                            -(c["ann"]["end"] - c["ann"]["start"]),
                            -c["score"],
                        ))
                        best = candidates[0]
                        return _make_payload(
                            best["choice"], e,
                            f"Position {position} maps to {best['ann']['type']}: "
                            f"{best['ann']['name']} "
                            f"(aa {best['ann']['start']}-{best['ann']['end']}). "
                            f"Selected by earliest-start preference among answer-choice candidates.",
                        )

                    # No high-confidence candidate — fall back to looser matching
                    # over containing annotations only (not the entire domain list).
                    best_choice = None
                    best_score = -1.0
                    best_ann = None
                    for ann in typed_containing:
                        for choice in choices:
                            score = SequenceMatcher(
                                None, ann["name"].lower(), choice.lower()
                            ).ratio()
                            if score > best_score:
                                best_score = score
                                best_choice = choice
                                best_ann = ann

                    if best_choice and best_score >= 0.4:
                        return _make_payload(
                            best_choice, e,
                            f"Position {position} maps to {best_ann['name']}. "
                            f"Matched answer choice: {best_choice} (score={best_score:.2f}).",
                        )

                # Fallback: use the primary DOMAIN MATCH line
                dm = re.search(
                    r"DOMAIN MATCH: amino acid position \d+ falls in (?:Domain|Region): (.+?) \(aa \d+-\d+\)",
                    snippet,
                )
                if dm:
                    best = _best_choice_match(dm.group(1).strip(), choices)
                    if best:
                        return _make_payload(best, e, f"Pre-computed domain match: {dm.group(1).strip()}.")

    return None


def _make_payload(response: str, evidence_item: dict, justification: str) -> dict:
    return {
        "response": response,
        "evidence": [
            {
                "source": evidence_item.get("url", ""),
                "time_accessed": int(time.time()),
                "justification": justification[:500],
            }
        ],
    }


def _stub_evidence(node_result) -> dict:
    """Convert a reasoning_dag.NodeResult into the evidence-dict shape
    `_make_payload` expects."""
    return {"url": getattr(node_result, "evidence_url", "")}