"""Scores a predicted GFF3 against a reference the way the NTv3 benchmark did.

Tiberius / BRAKER-style (Gabriel et al. 2024): a CDS is correct if chromosome,
start, end, strand and phase all match; a gene is correct if its whole set of CDS
matches a reference gene. Reference genes are restricted to those the NTv3
training labels keep, predicted genes below 201 nt of CDS are dropped, and
predicted CDS in pseudogene regions away from reference genes are ignored.

Standard library only, so it runs anywhere the notebook does.
"""

import bisect
import gzip
import re
from collections import defaultdict
from collections.abc import Callable, Iterable, Iterator, Mapping
from dataclasses import dataclass, field
from pathlib import Path
from typing import IO, NamedTuple

CdsKey = tuple[str, int, int, str, str]  # (chrom, start, end, strand, phase), 1-based
Intervals = dict[str, list[tuple[int, int]]]

# GENCODE keeps pseudogenes as generic rows with a pseudogene ``gene_type``.
_PSEUDOGENE_GENE_TYPE = re.compile(r"(?:^|;)gene_type=[^;]*pseudogene")

LEVELS = {
    "CDS": "Exon",
    "Gene": "Gene",
    "CDS_no_phase": "Exon (no phase)",
    "Gene_no_phase": "Gene (no phase)",
}


# ---------------------------------------------------------------------------
# GFF reading
# ---------------------------------------------------------------------------


class GffRow(NamedTuple):
    """One GFF3 feature row (1-based closed coordinates)."""

    seqid: str
    type: str
    start: int
    end: int
    strand: str
    phase: str
    attribute: str

    def attributes(self) -> dict[str, str]:
        """Column 9 as ``{tag: value}``."""
        out = {}
        for item in self.attribute.split(";"):
            if "=" in item:
                key, value = item.strip().split("=", 1)
                out[key.strip()] = value.strip()
        return out

    def get_id(self) -> str | None:
        """First ``ID=`` value."""
        for item in self.attribute.split(";"):
            item = item.strip()
            if item.startswith("ID="):
                return item[3:].split(",")[0].strip() or None
        return None

    def get_parents(self) -> list[str]:
        """All ``Parent=`` values."""
        for item in self.attribute.split(";"):
            item = item.strip()
            if item.startswith("Parent="):
                return [v.strip() for v in item[7:].split(",") if v.strip()]
        return []


def open_text(path: str | Path) -> IO[str]:
    """Opens a plain or gzipped text file."""
    if str(path).endswith(".gz"):
        return gzip.open(path, "rt", encoding="utf-8")
    return open(path, encoding="utf-8")


def iter_gff(
    path: str | Path, chromosomes: Iterable[str] | None = None
) -> Iterator[GffRow]:
    """Yields feature rows, skipping comments and malformed lines."""
    keep = set(chromosomes) if chromosomes is not None else None
    with open_text(path) as handle:
        for line in handle:
            if not line.strip() or line.startswith("#"):
                continue
            cols = line.rstrip("\n").split("\t")
            if len(cols) < 9 or (keep is not None and cols[0] not in keep):
                continue
            try:
                yield GffRow(cols[0], cols[2], int(cols[3]), int(cols[4]), *cols[6:9])
            except ValueError:
                continue


def merge_intervals(intervals: Iterable[tuple[int, int]]) -> list[tuple[int, int]]:
    """Merges overlapping or adjacent 1-based closed intervals."""
    merged: list[tuple[int, int]] = []
    for start, end in sorted(intervals):
        if merged and start <= merged[-1][1] + 1:
            merged[-1] = (merged[-1][0], max(merged[-1][1], end))
        else:
            merged.append((start, end))
    return merged


# ---------------------------------------------------------------------------
# Reference preparation
# ---------------------------------------------------------------------------


def write_longest_isoform(
    source: str | Path, output: str | Path, keep_pseudogenes: bool = False
) -> None:
    """Keeps each protein-coding gene with its longest-CDS transcript.

    Turns a raw multi-isoform annotation (e.g. an NCBI ``*_genomic.gff.gz``) into
    the single-isoform reference the benchmark scores against. Rows must be in
    the usual gene -> transcript -> exon/CDS order.

    Args:
        source: Raw GFF3, plain or gzipped.
        output: Destination GFF3, gzipped if it ends in ``.gz``.
        keep_pseudogenes: Also copy ``pseudogene`` rows, which turns on
            pseudogene filtering. The benchmark references carried none outside
            hg38, so leave off to reproduce it.
    """
    opener = gzip.open if str(output).endswith(".gz") else open
    transcripts: dict[str, list[str]] = {}
    lengths: dict[str, int] = {}
    gene_line: str | None = None
    gene_id: str | None = None
    transcript_id: str | None = None
    keep_gene = False

    def flush(out: IO[str]) -> None:
        if transcripts and gene_line is not None:
            longest = max(transcripts, key=lambda tid: lengths[tid])
            out.write(gene_line)
            out.writelines(transcripts[longest])

    with open_text(source) as src, opener(output, "wt", encoding="utf-8") as out:
        out.write("##gff-version 3\n")
        for line in src:
            if not line.strip() or line.startswith("#"):
                continue
            cols = line.rstrip("\n").split("\t")
            if len(cols) < 9:
                continue
            row = GffRow(cols[0], cols[2], int(cols[3]), int(cols[4]), *cols[6:9])
            if row.type in ("gene", "pseudogene"):
                flush(out)
                if row.type == "pseudogene" and keep_pseudogenes:
                    out.write(line)
                attrs = dict(
                    item.split("=", 1)
                    for item in row.attribute.split(";")
                    if "=" in item
                )
                biotype = attrs.get("gene_type") or attrs.get("gene_biotype")
                keep_gene = biotype == "protein_coding" or not biotype
                transcripts, lengths = {}, {}
                gene_id, gene_line = row.get_id(), line
            if not keep_gene:
                continue
            if row.type in ("mRNA", "transcript"):
                parents = row.get_parents()
                if parents and gene_id is not None and parents[0] == gene_id:
                    transcript_id = row.get_id()
                    if transcript_id is not None:
                        transcripts[transcript_id] = [line]
                        lengths[transcript_id] = 0
            elif (
                row.type in ("exon", "CDS")
                and transcript_id is not None
                and transcript_id in transcripts
            ):
                # Children join the last transcript seen, as in the pipeline.
                transcripts[transcript_id].append(line)
                if row.type == "CDS":
                    lengths[transcript_id] += row.end - row.start + 1
        flush(out)


@dataclass
class _CodingGene:
    gene_id: str
    name: str
    chrom: str
    strand: str
    cds: list[tuple[int, int]] = field(default_factory=list)


def _parse_coding_genes(
    path: str | Path, chromosomes: Iterable[str] | None = None
) -> dict[str, _CodingGene]:
    genes: dict[str, _CodingGene] = {}
    transcript_to_gene: dict[str, str] = {}
    cds_by_transcript: dict[str, list[tuple[int, int]]] = defaultdict(list)
    for row in iter_gff(path, chromosomes):
        if row.type == "gene":
            attrs = row.attributes()
            biotype = attrs.get("gene_type") or attrs.get("gene_biotype")
            gene_id = row.get_id()
            if (biotype == "protein_coding" or not biotype) and gene_id:
                if row.strand in ("+", "-"):
                    name = attrs.get("Name") or gene_id.removeprefix("gene-")
                    genes[gene_id] = _CodingGene(gene_id, name, row.seqid, row.strand)
        elif row.type in ("mRNA", "transcript"):
            parents, tid = row.get_parents(), row.get_id()
            if parents and tid:
                transcript_to_gene[tid] = parents[0]
        elif row.type == "CDS" and row.phase != ".":
            parents = row.get_parents()
            if parents:
                cds_by_transcript[parents[0]].append((row.start, row.end))
    for tid, gid in transcript_to_gene.items():
        if gid in genes:
            genes[gid].cds.extend(cds_by_transcript.get(tid, []))
    for gene in genes.values():
        gene.cds.sort()
    return genes


def _intervals_overlap(a: list[tuple[int, int]], b: list[tuple[int, int]]) -> bool:
    i = j = 0
    while i < len(a) and j < len(b):
        if a[i][0] <= b[j][1] and b[j][0] <= a[i][1]:
            return True
        if a[i][1] < b[j][1]:
            i += 1
        else:
            j += 1
    return False


def _cluster_by_span(
    gene_ids: list[str], genes: dict[str, _CodingGene]
) -> list[list[str]]:
    ordered = sorted(gene_ids, key=lambda g: (genes[g].cds[0][0], genes[g].cds[-1][1]))
    clusters: list[list[str]] = []
    end = -1
    for gid in ordered:
        if clusters and genes[gid].cds[0][0] <= end:
            clusters[-1].append(gid)
            end = max(end, genes[gid].cds[-1][1])
        else:
            clusters.append([gid])
            end = genes[gid].cds[-1][1]
    return clusters


def _is_readthrough(name: str, other_names: set[str]) -> bool:
    parts = name.split("-")
    return any(
        "-".join(parts[:i]) in other_names and "-".join(parts[i:]) in other_names
        for i in range(1, len(parts))
    )


def _cds_length(gene: _CodingGene) -> int:
    return sum(end - start + 1 for start, end in gene.cds)


def _conflicting_gene_ids(genes: dict[str, _CodingGene]) -> set[str]:
    """Readthrough fusions, then the shorter of two genes sharing CDS bases."""
    by_strand: dict[tuple[str, str], list[str]] = defaultdict(list)
    for gid, gene in genes.items():
        by_strand[(gene.chrom, gene.strand)].append(gid)
    conflicts: set[str] = set()
    for gids in by_strand.values():
        for cluster in _cluster_by_span(gids, genes):
            if len(cluster) < 2:
                continue
            names = {gid: genes[gid].name for gid in cluster}
            all_names = set(names.values())
            remaining = []
            for gid in cluster:
                if _is_readthrough(names[gid], all_names - {names[gid]}):
                    conflicts.add(gid)
                else:
                    remaining.append(gid)
            kept: list[str] = []
            for gid in remaining:
                clash = next(
                    (
                        other
                        for other in kept
                        if genes[gid].cds[0][0] <= genes[other].cds[-1][1]
                        and _intervals_overlap(genes[other].cds, genes[gid].cds)
                    ),
                    None,
                )
                if clash is None:
                    kept.append(gid)
                elif _cds_length(genes[gid]) > _cds_length(genes[clash]):
                    conflicts.add(clash)
                    kept.remove(clash)
                    kept.append(gid)
                else:
                    conflicts.add(gid)
    return conflicts


def scorable_reference_genes(
    path: str | Path, chromosomes: Iterable[str] | None = None
) -> set[str]:
    """IDs of the reference genes the benchmark scores.

    Drops genes with no CDS, overlapping CDS, readthrough fusions and CDS shared
    with a longer gene, then keeps the longest gene of every same-strand cluster
    of overlapping CDS spans (so a gene nested in another's intron goes too).
    """
    genes = {
        gid: gene
        for gid, gene in _parse_coding_genes(path, chromosomes).items()
        if gene.cds
        and all(b[0] > a[1] for a, b in zip(gene.cds, gene.cds[1:], strict=False))
    }
    conflicts = _conflicting_gene_ids(genes)
    genes = {gid: gene for gid, gene in genes.items() if gid not in conflicts}
    by_strand: dict[tuple[str, str], list[str]] = defaultdict(list)
    for gid, gene in genes.items():
        by_strand[(gene.chrom, gene.strand)].append(gid)
    return {
        max(cluster, key=lambda g: (_cds_length(genes[g]), g))
        for gids in by_strand.values()
        for cluster in _cluster_by_span(gids, genes)
    }


# ---------------------------------------------------------------------------
# Scoring
# ---------------------------------------------------------------------------


@dataclass
class Gene:
    """A gene reduced to what the scorer compares."""

    gene_id: str
    chrom: str
    strand: str
    cds_keys: list[CdsKey] = field(default_factory=list)


def parse_cds_and_genes(
    path: str | Path, chromosomes: Iterable[str] | None = None
) -> tuple[set[CdsKey], dict[str, Gene]]:
    """All CDS keys, and each gene's CDS keys (CDS reach genes via transcripts)."""
    cds: set[CdsKey] = set()
    genes: dict[str, Gene] = {}
    transcript_to_gene: dict[str, str] = {}
    for row in iter_gff(path, chromosomes):
        kind = row.type.lower()
        if kind in ("gene", "pseudogene"):
            gid = row.get_id()
            if gid and gid not in genes:
                genes[gid] = Gene(gid, row.seqid, row.strand)
        elif kind in ("mrna", "transcript"):
            tid, parents = row.get_id(), row.get_parents()
            if tid and parents:
                transcript_to_gene[tid] = parents[0]
        elif kind == "cds":
            parents = row.get_parents()
            parent = parents[0] if parents else ""
            gid = transcript_to_gene.get(parent, parent)
            key = (row.seqid, row.start, row.end, row.strand, row.phase)
            cds.add(key)
            if gid:
                genes.setdefault(gid, Gene(gid, row.seqid, row.strand))
                genes[gid].cds_keys.append(key)
    return cds, genes


def parse_pseudogene_regions(
    path: str | Path, chromosomes: Iterable[str] | None = None
) -> list[tuple[str, int, int]]:
    """``pseudogene`` rows, or rows whose ``gene_type`` is a pseudogene (GENCODE)."""
    return [
        (row.seqid, row.start, row.end)
        for row in iter_gff(path, chromosomes)
        if row.type == "pseudogene" or _PSEUDOGENE_GENE_TYPE.search(row.attribute)
    ]


def _interval_index(regions: Iterable[tuple[str, int, int]]) -> Intervals:
    by_chrom: dict[str, list[tuple[int, int]]] = defaultdict(list)
    for chrom, start, end in regions:
        by_chrom[chrom].append((start, end))
    return {chrom: merge_intervals(ivs) for chrom, ivs in by_chrom.items()}


def _overlaps_any(chrom: str, start: int, end: int, index: Intervals) -> bool:
    intervals = index.get(chrom)
    if not intervals:
        return False
    i = bisect.bisect_left(intervals, (end + 1,)) - 1
    return i >= 0 and intervals[i][1] >= start


@dataclass
class ScoringSets:
    """Reference and prediction after every benchmark filter."""

    ref_cds: set[CdsKey]
    ref_genes: dict[str, Gene]
    pred_cds: set[CdsKey]
    pred_genes: dict[str, Gene]


def prepare_scoring_sets(
    reference: str | Path,
    prediction: str | Path,
    pseudogenes: str | Path | None = None,
    chromosomes: Iterable[str] | None = None,
    min_pred_gene_cds_length: int = 201,
    full_reference: bool = False,
) -> ScoringSets:
    """Parses both files and applies the benchmark filters.

    Args:
        reference: Single-isoform reference GFF3.
        prediction: Predicted GFF3.
        pseudogenes: GFF3 to read pseudogene regions from; the benchmark passed
            the reference itself.
        chromosomes: Seqids to score; None scores the whole reference, so a
            contig the prediction skips counts as missed.
        min_pred_gene_cds_length: Predicted genes with less CDS (nt) are dropped.
        full_reference: Score every reference gene with a CDS instead of the
            ones NTv3 training labels keep.

    Returns:
        The filtered CDS and gene sets.
    """
    chroms = sorted(set(chromosomes)) if chromosomes is not None else None
    ref_cds, ref_genes = parse_cds_and_genes(reference, chroms)
    if full_reference:
        ref_genes = {gid: gene for gid, gene in ref_genes.items() if gene.cds_keys}
    else:
        keep = scorable_reference_genes(reference, chroms)
        ref_genes = {gid: gene for gid, gene in ref_genes.items() if gid in keep}
    ref_cds = {key for gene in ref_genes.values() for key in gene.cds_keys}

    pred_cds, pred_genes = parse_cds_and_genes(prediction, chroms)
    if min_pred_gene_cds_length > 0:
        short = {
            gid
            for gid, gene in pred_genes.items()
            if sum(end - start + 1 for _, start, end, _, _ in gene.cds_keys)
            < min_pred_gene_cds_length
        }
        pred_cds -= {key for gid in short for key in pred_genes[gid].cds_keys}
        pred_genes = {gid: g for gid, g in pred_genes.items() if gid not in short}

    if pseudogenes is not None:
        ref_index = _interval_index((c, s, e) for c, s, e, _, _ in ref_cds)
        away_from_genes = _interval_index(
            region
            for region in parse_pseudogene_regions(pseudogenes, chroms)
            if not _overlaps_any(*region, ref_index)
        )
        pred_cds = {k for k in pred_cds if not _overlaps_any(*k[:3], away_from_genes)}
        pred_genes = {
            gid: Gene(gid, g.chrom, g.strand, kept)
            for gid, g in pred_genes.items()
            if (kept := [k for k in g.cds_keys if k in pred_cds])
        }
    return ScoringSets(ref_cds, ref_genes, pred_cds, pred_genes)


def restrict_to_region(
    sets: ScoringSets, chrom: str, start: int, end: int
) -> ScoringSets:
    """Keeps the genes whose CDS lie inside [start, end) (1-based, closed ends).

    Genes cut by the region edges are dropped on both sides, so neither the
    reference nor the prediction is penalised for what lies outside.
    """

    def inside(gene: Gene) -> bool:
        return gene.chrom == chrom and all(
            start <= k[1] and k[2] <= end for k in gene.cds_keys
        )

    ref_genes = {g: x for g, x in sets.ref_genes.items() if inside(x)}
    pred_genes = {g: x for g, x in sets.pred_genes.items() if inside(x)}
    return ScoringSets(
        {k for x in ref_genes.values() for k in x.cds_keys},
        ref_genes,
        {k for x in pred_genes.values() for k in x.cds_keys},
        pred_genes,
    )


def _prf(ref_tp: int, n_ref: int, pred_tp: int, n_pred: int) -> dict[str, float]:
    recall = ref_tp / n_ref if n_ref else 0.0
    precision = pred_tp / n_pred if n_pred else 0.0
    f1 = 2 * recall * precision / (recall + precision) if recall + precision else 0.0
    return {
        "Ref_TP": ref_tp,
        "Ref_Total": n_ref,
        "Pred_TP": pred_tp,
        "Pred_Total": n_pred,
        "Recall": recall,
        "Precision": precision,
        "F1": f1,
    }


def _no_phase(keys: Iterable[CdsKey]) -> frozenset[tuple[str, int, int, str]]:
    return frozenset(key[:4] for key in keys)


def _cds_scores(ref: set, pred: set) -> dict[str, float]:
    tp = len(ref & pred)
    scores = _prf(tp, len(ref), tp, len(pred))
    return {
        "TP": tp,
        "FN": len(ref) - tp,
        "FP": len(pred) - tp,
        **{k: scores[k] for k in ("Recall", "Precision", "F1")},
    }


def _gene_scores(
    ref_genes: Iterable[Gene],
    pred_genes: Iterable[Gene],
    structure: Callable[[Gene], frozenset],
    all_ref: set[frozenset],
    all_pred: set[frozenset],
) -> dict[str, float]:
    ref = [structure(g) for g in ref_genes]
    pred = [structure(g) for g in pred_genes]
    return _prf(
        sum(s in all_pred for s in ref),
        len(ref),
        sum(s in all_ref for s in pred),
        len(pred),
    )


def score_sets(sets: ScoringSets) -> dict[str, dict[str, float]]:
    """CDS- and gene-level scores, with and without phase (fractions, not %)."""
    return _score(sets, sets.ref_cds, sets.pred_cds, sets.ref_genes, sets.pred_genes)


def _score(
    sets: ScoringSets,
    ref_cds: set[CdsKey],
    pred_cds: set[CdsKey],
    ref_genes: dict[str, Gene],
    pred_genes: dict[str, Gene],
) -> dict[str, dict[str, float]]:
    with_phase: Callable[[Gene], frozenset] = lambda g: frozenset(g.cds_keys)  # noqa: E731
    without_phase: Callable[[Gene], frozenset] = lambda g: _no_phase(g.cds_keys)  # noqa: E731
    results = {"CDS": _cds_scores(ref_cds, pred_cds)}
    results["Gene"] = _gene_scores(
        ref_genes.values(),
        pred_genes.values(),
        with_phase,
        {with_phase(g) for g in sets.ref_genes.values()},
        {with_phase(g) for g in sets.pred_genes.values()},
    )
    results["CDS_no_phase"] = _cds_scores(
        {key[:4] for key in ref_cds}, {key[:4] for key in pred_cds}
    )
    results["Gene_no_phase"] = _gene_scores(
        ref_genes.values(),
        pred_genes.values(),
        without_phase,
        {without_phase(g) for g in sets.ref_genes.values()},
        {without_phase(g) for g in sets.pred_genes.values()},
    )
    return results


def evaluate(
    reference: str | Path,
    prediction: str | Path,
    pseudogenes: str | Path | None = None,
    chromosomes: Iterable[str] | None = None,
    min_pred_gene_cds_length: int = 201,
    full_reference: bool = False,
) -> dict[str, dict[str, float]]:
    """Benchmark scores of one prediction; see ``prepare_scoring_sets`` for args."""
    return score_sets(
        prepare_scoring_sets(
            reference,
            prediction,
            pseudogenes,
            chromosomes,
            min_pred_gene_cds_length,
            full_reference,
        )
    )


def score_per_contig(sets: ScoringSets) -> dict[str, dict[str, dict[str, float]]]:
    """Scores per seqid, matched against the whole genome (counts add up)."""
    chroms = sorted(
        {g.chrom for g in sets.ref_genes.values()}
        | {g.chrom for g in sets.pred_genes.values()}
    )
    return {
        chrom: _score(
            sets,
            {k for k in sets.ref_cds if k[0] == chrom},
            {k for k in sets.pred_cds if k[0] == chrom},
            {gid: g for gid, g in sets.ref_genes.items() if g.chrom == chrom},
            {gid: g for gid, g in sets.pred_genes.items() if g.chrom == chrom},
        )
        for chrom in chroms
    }


def compare_seqids(reference: str | Path, prediction: str | Path) -> dict[str, list]:
    """Seqids with CDS in each file, to catch naming mismatches (``chr1`` vs ``1``)."""
    ref = {key[0] for key in parse_cds_and_genes(reference)[0]}
    pred = {key[0] for key in parse_cds_and_genes(prediction)[0]}
    return {
        "shared": sorted(ref & pred),
        "reference_only": sorted(ref - pred),
        "prediction_only": sorted(pred - ref),
    }


def rename_seqids(
    source: str | Path,
    output: str | Path,
    rename: Mapping[str, str] | Callable[[str], str],
) -> None:
    """Writes a copy of a GFF3 with column 1 (and ``##sequence-region``) renamed.

    Args:
        source: GFF3 to rename, plain or gzipped.
        output: Destination, gzipped if it ends in ``.gz``.
        rename: Old to new seqid (unlisted ones are kept), or a function.
    """
    if callable(rename):
        new = rename
    else:
        mapping = dict(rename)

        def new(seqid: str) -> str:
            return mapping.get(seqid, seqid)

    opener = gzip.open if str(output).endswith(".gz") else open
    with open_text(source) as src, opener(output, "wt", encoding="utf-8") as out:
        for line in src:
            if line.startswith("##sequence-region"):
                tag, seqid, *rest = line.split()
                line = " ".join([tag, str(new(seqid)), *rest]) + "\n"
            elif line.strip() and not line.startswith("#"):
                seqid, rest = line.split("\t", 1)
                line = f"{new(seqid)}\t{rest}"
            out.write(line)


def format_summary(results: dict[str, dict[str, float]]) -> list[str]:
    """One line per level, in %."""
    return [
        f"{LEVELS.get(level, level):>16s}:  Recall {100 * m['Recall']:6.2f}%  "
        f"Precision {100 * m['Precision']:6.2f}%  F1 {100 * m['F1']:6.2f}%"
        for level, m in results.items()
    ]