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

The rank is a weighted sum of four parts 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.
"""
from collections.abc import Mapping, Sequence
from dataclasses import dataclass

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))


@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


@dataclass(frozen=True)
class Candidate:
    variant: Variant
    score: float
    components: dict[str, float]
    matched_terms: list[str]
    scored: bool


def rarity_score(af: float | None) -> float:
    if af is None:  # absent from gnomAD
        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], gene_terms: Mapping[str, set[str]]
) -> tuple[float, list[str]]:
    """What fraction of the patient's terms HPO associates with this gene, and which ones."""
    if not gene or not case_terms:
        return 0.0, []
    annotated = gene_terms.get(gene, set())
    matched = [term for term in case_terms if term in annotated]
    return len(matched) / len(case_terms), 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 funnel(
    variants: Sequence[Variant], case_terms: Sequence[str], gene_terms: Mapping[str, set[str]]
) -> 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, gene_terms)[1])
    return Funnel(len(variants), len(rare), len(candidates), matched)


def evaluate(
    variant: Variant, case_terms: Sequence[str], gene_terms: Mapping[str, set[str]]
) -> Candidate:
    """Score one variant, whether or not it survived the filters."""
    phenotype, matched = phenotype_score(variant.gene, case_terms, gene_terms)
    prediction = variant.prediction
    components = {
        "phenotype": phenotype,
        "rarity": rarity_score(variant.gnomad_af),
        "consequence": consequence_score(variant.impact),
        "model": float(prediction.score) if prediction is not None else 0.0,
    }
    score = sum(WEIGHTS[name] * value for name, value in components.items())
    return Candidate(variant, score, components, matched, prediction is not None)


def rank(
    variants: Sequence[Variant], case_terms: Sequence[str], gene_terms: Mapping[str, set[str]]
) -> list[Candidate]:
    candidates = [evaluate(v, case_terms, gene_terms) 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