"""Narrow a case's variants the way a clinical scientist does, and say why.

The rank is a weighted mean of four lines of evidence a reviewer can audit. ClinVar is deliberately
not one of them: it is shown beside the result as independent confirmation, so a variant never
ranks highly merely because ClinVar already called it pathogenic.

Rarity and consequence *filter* (the usual first pass); phenotype only *ranks*, because a real
diagnosis can sit in a gene nobody has annotated yet and filtering on it would hide exactly that.

Two rules keep the number honest:

**A line of evidence that was never looked up abstains.** It does not score zero, and it certainly
does not score full marks. Treating "no gnomAD frequency in the annotation run" as "absent from
gnomAD, therefore maximally rare" awarded every variant a free 0.25, which is a guess wearing the
costume of a measurement. `Evidence` says what the run actually produced, and the weights
renormalise over whatever is left, so the score stays on a 0-1 scale and means the same thing.

**Each line of evidence is counted once.** The model used to take allele frequency as a feature
while `rarity` scored the same frequency again, so roughly 45% of the rank was one measurement
double-counted. The model no longer sees frequency (see rarelens_ml.features); it earns its weight
only when it has something the other three do not already say, which means CADD or AlphaMissense.
"""
from collections.abc import Mapping, Sequence
from dataclasses import dataclass, field

from app.models import Variant

WEIGHTS = {"phenotype": 0.35, "rarity": 0.25, "consequence": 0.20, "model": 0.20}

RARE_AF = 0.001
CANDIDATE_IMPACTS = frozenset({"HIGH", "MODERATE"})
IMPACT_SEVERITY = {"HIGH": 1.0, "MODERATE": 0.6, "LOW": 0.2, "MODIFIER": 0.0}
# Allele frequency ceiling -> score, rarest first.
RARITY_STEPS = ((0.0, 1.0), (0.0001, 0.8), (0.001, 0.5), (0.01, 0.2))
# A term HPO has never annotated to any gene cannot match anything, so its information content is
# unknown. Treating it as maximally specific keeps it in the denominator and depresses every gene
# equally, which is the neutral choice.
DEFAULT_IC = 10.0


@dataclass(frozen=True)
class Evidence:
    """What the annotation run actually produced, and therefore which components may score.

    Decided once per job rather than per variant: components must be in play for every variant in
    a case, or two variants would be scored against different denominators and their ranks would
    not be comparable.
    """

    frequencies: bool = False       # did the run look up allele frequencies at all?
    effect_scores: bool = False     # CADD / AlphaMissense, the only features the model adds

    @property
    def missing(self) -> list[str]:
        absent = []
        if not self.frequencies:
            absent.append("rarity")
        if not self.effect_scores:
            absent.append("model")
        return absent


@dataclass(frozen=True)
class Ontology:
    """HPO reference data: what each gene is annotated with, and how specific each term is.

    `gene_terms` is expected to be propagated up the ontology by scripts/load-hpo.py, so a case
    term matches a gene annotated with any of its descendants. `ic` is information content,
    -ln(fraction of genes carrying the term): "Bifid uvula" is worth many times "Abnormality of
    the head", which nearly every gene in the corpus carries.
    """

    gene_terms: Mapping[str, set[str]] = field(default_factory=dict)
    ic: Mapping[str, float] = field(default_factory=dict)

    def weight(self, term: str) -> float:
        return self.ic.get(term, DEFAULT_IC)


@dataclass(frozen=True)
class Funnel:
    """How many variants survive each narrowing step; the headline of the case page."""

    total: int
    rare: int
    candidates: int
    phenotype_matched: int
    frequencies: bool = False  # False means the "rare" step filtered nothing, because it could not


@dataclass(frozen=True)
class Candidate:
    variant: Variant
    score: float
    components: dict[str, float | None]  # None: this evidence was not available
    matched_terms: list[str]
    scored: bool


def rarity_score(af: float | None) -> float:
    """Only meaningful when frequencies were annotated; None then means absent from gnomAD."""
    if af is None:
        return 1.0
    for ceiling, score in RARITY_STEPS:
        if af <= ceiling:
            return score
    return 0.0


def consequence_score(impact: str | None) -> float:
    return IMPACT_SEVERITY.get(impact or "", 0.0)


def phenotype_score(
    gene: str | None, case_terms: Sequence[str], ontology: Ontology
) -> tuple[float, list[str]]:
    """How much of the patient's phenotype HPO associates with this gene, weighted by specificity.

    Information-content-weighted recall: the share of the *total specificity* of the patient's
    terms that this gene accounts for. Plain term counting let a common term like global
    developmental delay count as much as a near-pathognomonic one.
    """
    if not gene or not case_terms:
        return 0.0, []
    annotated = ontology.gene_terms.get(gene, set())
    matched = [term for term in case_terms if term in annotated]
    total = sum(ontology.weight(term) for term in case_terms)
    if total <= 0:
        return 0.0, matched
    return sum(ontology.weight(term) for term in matched) / total, matched


def is_rare(variant: Variant) -> bool:
    return variant.gnomad_af is None or variant.gnomad_af < RARE_AF


def is_candidate(variant: Variant) -> bool:
    return is_rare(variant) and variant.impact in CANDIDATE_IMPACTS


def combine(components: Mapping[str, float | None]) -> float:
    """Weighted mean over the components that have evidence, renormalised to 0-1."""
    weight = sum(WEIGHTS[name] for name, value in components.items() if value is not None)
    if weight <= 0:
        return 0.0
    return sum(WEIGHTS[name] * value for name, value in components.items() if value is not None) / weight


def weights_in_use(evidence: Evidence) -> dict[str, float]:
    """The weights as actually applied, so the UI never shows a bar the score did not use."""
    live = {name: w for name, w in WEIGHTS.items() if name not in evidence.missing}
    total = sum(live.values())
    return {name: round(w / total, 4) for name, w in live.items()} if total else {}


def funnel(variants: Sequence[Variant], case_terms: Sequence[str], ontology: Ontology,
           evidence: Evidence) -> Funnel:
    rare = [v for v in variants if is_rare(v)]
    candidates = [v for v in rare if v.impact in CANDIDATE_IMPACTS]
    matched = sum(1 for v in candidates if phenotype_score(v.gene, case_terms, ontology)[1])
    return Funnel(len(variants), len(rare), len(candidates), matched, evidence.frequencies)


def evaluate(
    variant: Variant, case_terms: Sequence[str], ontology: Ontology, evidence: Evidence
) -> Candidate:
    """Score one variant, whether or not it survived the filters."""
    phenotype, matched = phenotype_score(variant.gene, case_terms, ontology)
    prediction = variant.prediction
    model: float | None = None
    if evidence.effect_scores and prediction is not None:
        model = float(prediction.score)
    components: dict[str, float | None] = {
        "phenotype": phenotype,
        "rarity": rarity_score(variant.gnomad_af) if evidence.frequencies else None,
        "consequence": consequence_score(variant.impact),
        "model": model,
    }
    return Candidate(variant, combine(components), components, matched, prediction is not None)


def rank(
    variants: Sequence[Variant], case_terms: Sequence[str], ontology: Ontology, evidence: Evidence
) -> list[Candidate]:
    candidates = [evaluate(v, case_terms, ontology, evidence) for v in variants if is_candidate(v)]
    # id breaks ties, so equal scores do not shuffle between requests.
    candidates.sort(key=lambda c: (-c.score, c.variant.id))
    return candidates