"""Protein-level checks of a predicted annotation (needs the genome FASTA).

Translates every gene, checks predicted ORFs (start with M, no internal stop), and
scores each reference protein by the identity of its best overlapping prediction.
"""

import dataclasses
from collections.abc import Callable, Iterable, Mapping
from dataclasses import dataclass, field

from Bio.Align import PairwiseAligner
from Bio.Seq import Seq

from .analysis import CdsIndex
from .evaluation import Gene

IDENTITY_THRESHOLDS = (100.0, 95.0, 90.0, 50.0)
IDENTITY_BINS = ("100", "95-100", "90-95", "50-90", "<50")
# Shorter proteins are fragments, not scored for ORF shape.
MIN_PROTEIN_LENGTH = 20


@dataclass
class ProteinCounts:
    """Protein-level tallies; ORF fields cover predictions of 20+ residues."""

    n_reference: int = 0
    recovered: dict[str, int] = field(
        default_factory=lambda: {f"{t:g}": 0 for t in IDENTITY_THRESHOLDS}
    )
    no_counterpart: int = 0
    n_predicted: int = 0
    starts_with_methionine: int = 0
    no_internal_stop: int = 0
    well_formed: int = 0
    identity_histogram: dict[str, int] = field(
        default_factory=lambda: dict.fromkeys(IDENTITY_BINS, 0)
    )


def translate_gene(gene: Gene, sequence: str) -> str:
    """Protein encoded by a gene's CDS on its (uppercase) chromosome sequence.

    CDS are joined in ascending genomic order and reverse-complemented once on the
    minus strand; the phase of the transcript's first CDS is honoured, a trailing
    partial codon and one terminal stop are dropped.
    """
    ascending = sorted(gene.cds_keys, key=lambda k: k[1])
    coding = "".join(sequence[start - 1 : end] for _, start, end, _, _ in ascending)
    if gene.strand == "-":
        coding = str(Seq(coding).reverse_complement())
    phase = (ascending[-1] if gene.strand == "-" else ascending[0])[4]
    coding = coding[int(phase) if phase.isdigit() else 0 :]
    coding = coding[: len(coding) - len(coding) % 3]
    if not coding:
        return ""
    protein = str(Seq(coding).translate())
    return protein[:-1] if protein.endswith("*") else protein


def identity(a: str, b: str, aligner: PairwiseAligner) -> float:
    """Percent identity of a global alignment, over the longer protein."""
    if a == b:
        return 100.0
    if not a or not b:
        return 0.0
    alignment = aligner.align(a, b)[0]
    matches = sum(
        x == y and x != "-" for x, y in zip(alignment[0], alignment[1], strict=True)
    )
    return 100.0 * matches / max(len(a), len(b))


def _identity_bin(value: float) -> str:
    for threshold, label in zip(IDENTITY_THRESHOLDS, IDENTITY_BINS, strict=False):
        if value >= threshold:
            return label
    return "<50"


def measure_proteins(
    ref_genes: dict[str, Gene],
    pred_genes: dict[str, Gene],
    genome: Mapping[str, str] | Callable[[str], str],
    chromosomes: Iterable[str] | None = None,
) -> dict[str, object]:
    """ORF shape of the predictions and identity recovery of the reference.

    Args:
        ref_genes: Reference genes, e.g. ``ScoringSets.ref_genes``.
        pred_genes: Predicted genes.
        genome: Chromosome name to sequence, or a function returning it (e.g. over a
            ``pyfaidx.Fasta``); read one chromosome at a time.
        chromosomes: Chromosomes to measure; defaults to every one with a gene.

    Returns:
        ``ProteinCounts`` as a dict.
    """
    fetch = genome if callable(genome) else genome.__getitem__
    if chromosomes is None:
        chromosomes = sorted(
            {g.chrom for g in ref_genes.values()}
            | {g.chrom for g in pred_genes.values()}
        )
    aligner = PairwiseAligner(scoring="blastp")
    aligner.mode = "global"
    counts = ProteinCounts()

    for chrom in chromosomes:
        chrom_ref = {
            g: x for g, x in ref_genes.items() if x.chrom == chrom and x.cds_keys
        }
        chrom_pred = {
            g: x for g, x in pred_genes.items() if x.chrom == chrom and x.cds_keys
        }
        if not chrom_ref and not chrom_pred:
            continue
        sequence = str(fetch(chrom)).upper()
        ref_proteins = {g: translate_gene(x, sequence) for g, x in chrom_ref.items()}
        pred_proteins = {g: translate_gene(x, sequence) for g, x in chrom_pred.items()}

        for protein in pred_proteins.values():
            if len(protein) < MIN_PROTEIN_LENGTH:
                continue
            starts, clean = protein.startswith("M"), "*" not in protein
            counts.n_predicted += 1
            counts.starts_with_methionine += starts
            counts.no_internal_stop += clean
            counts.well_formed += starts and clean

        index = CdsIndex(chrom_pred)
        for gid, gene in chrom_ref.items():
            reference = ref_proteins[gid]
            if not reference:
                continue
            counts.n_reference += 1
            candidates = [
                pred_proteins[other]
                for other in index.overlapping_gene(gene)
                if pred_proteins.get(other)
            ]
            if not candidates:
                counts.no_counterpart += 1
                counts.identity_histogram["<50"] += 1
                continue
            best = max(identity(reference, c, aligner) for c in candidates)
            for threshold in IDENTITY_THRESHOLDS:
                counts.recovered[f"{threshold:g}"] += best >= threshold
            counts.identity_histogram[_identity_bin(best)] += 1
    return dataclasses.asdict(counts)