"""Which residues evolution has refused to change — computed, not predicted.

WHY THIS EXISTS ALONGSIDE ESM-2
-------------------------------
ESM-2 gives a learned opinion about whether a substitution looks plausible.
Conservation across real homologs gives an observed fact: in 40 orthologs,
this position is serine 40 times. Those are different kinds of evidence and
they fail differently — the model is weakest exactly where the audit says it
is (membrane proteins, disordered regions, multi-domain assemblies), and a
frequency count is unaffected by any of that.

So this complements the scorer rather than duplicating it, which is why the
audit rated it High: it is the cheapest real signal for "don't mutate this
residue", and unlike a prediction it can be checked by counting.

TWO PIECES
----------
`align_many` builds a progressive multiple alignment on top of the pairwise
aligner that already exists (`align.py`), anchored on the longest sequence.
This is the classic progressive approach and it is approximate: a true
simultaneous MSA is exponential, and every practical tool approximates. It is
honest about being an approximation rather than presenting itself as ground
truth.

`score` then counts. Shannon entropy per column, the most common residue, and
the fraction of sequences that agree. No model, no training, no weights.

THE HONEST LIMIT, WHICH IS ABOUT INPUT NOT ALGORITHM
----------------------------------------------------
Conservation is only as meaningful as the homolog set. Forty sequences that
are 99% identical to each other say nothing — they are one sequence counted
forty times. The result therefore reports the diversity of the input, and
refuses to present a confident conservation call on a set with no spread.
"""
from __future__ import annotations

import math
import re
from collections import Counter
from typing import Any, Dict, List, Optional, Sequence

from dee.core.align import align

GAP = "-"
# Below this many sequences a "conserved" call means almost nothing.
MIN_SEQS = 3
# Mean pairwise identity above which the set is effectively one sequence.
REDUNDANT_ABOVE = 95.0


def _clean(s: str) -> str:
    return re.sub(r"[^A-Za-z*]", "", s or "").upper()


def align_many(sequences: Sequence[str],
               names: Optional[Sequence[str]] = None) -> Dict[str, Any]:
    """Progressive multiple alignment, anchored on the longest sequence.

    Each sequence is aligned pairwise to the anchor and its gaps merged into a
    common frame. Approximate by construction — stated in the result so no
    caller mistakes it for a simultaneous optimum.
    """
    seqs = [_clean(s) for s in (sequences or [])]
    keep = [(i, s) for i, s in enumerate(seqs) if s]
    if len(keep) < 2:
        return {"ok": False, "error": "Need at least two non-empty sequences."}

    labels = list(names or [])
    def label(i: int) -> str:
        return str(labels[i]) if i < len(labels) and labels[i] else f"seq{i + 1}"

    anchor_i, anchor = max(keep, key=lambda kv: len(kv[1]))

    # Columns the anchor must gain, keyed by the anchor index they precede.
    inserts: Dict[int, int] = {}
    pairs: List[Dict[str, Any]] = []
    for i, s in keep:
        if i == anchor_i:
            continue
        try:
            r = align(anchor, s, mode="global")
        except ValueError as exc:
            return {"ok": False, "kind": "too_large", "error": str(exc),
                    "next": "Align shorter sequences, or fewer of them."}
        pairs.append({"i": i, "a": r["aligned_a"], "b": r["aligned_b"],
                      "identity": r["identity"]})
        pos = 0
        run = 0
        for ca in r["aligned_a"]:
            if ca == GAP:
                run += 1
            else:
                if run:
                    inserts[pos] = max(inserts.get(pos, 0), run)
                    run = 0
                pos += 1
        if run:
            inserts[pos] = max(inserts.get(pos, 0), run)

    def expand(aligned_a: str, aligned_b: str) -> str:
        """Re-lay one pairwise result into the common frame."""
        out: List[str] = []
        pos = 0
        run: List[str] = []
        for ca, cb in zip(aligned_a, aligned_b):
            if ca == GAP:
                run.append(cb)
                continue
            need = inserts.get(pos, 0)
            out.append("".join(run).ljust(need, GAP)[:need] if need else "")
            run = []
            out.append(cb)
            pos += 1
        need = inserts.get(pos, 0)
        out.append("".join(run).ljust(need, GAP)[:need] if need else "")
        return "".join(out)

    rows: List[Dict[str, Any]] = []
    frame_anchor = []
    for pos, ch in enumerate(anchor):
        frame_anchor.append(GAP * inserts.get(pos, 0) + ch)
    frame_anchor.append(GAP * inserts.get(len(anchor), 0))
    anchor_row = "".join(frame_anchor)
    rows.append({"name": label(anchor_i), "aligned": anchor_row,
                 "identity_to_anchor": 100.0, "is_anchor": True})
    for p in pairs:
        rows.append({"name": label(p["i"]), "aligned": expand(p["a"], p["b"]),
                     "identity_to_anchor": p["identity"], "is_anchor": False})

    width = max(len(r["aligned"]) for r in rows)
    for r in rows:
        r["aligned"] = r["aligned"].ljust(width, GAP)

    ids = [r["identity_to_anchor"] for r in rows if not r["is_anchor"]]
    return {
        "ok": True,
        "rows": rows,
        "columns": width,
        "sequences": len(rows),
        "anchor": label(anchor_i),
        "mean_identity_to_anchor": round(sum(ids) / len(ids), 1) if ids else 100.0,
        "method": ("Progressive alignment onto the longest sequence, using the "
                   "engine's Needleman-Wunsch. Approximate: a simultaneous "
                   "optimum is exponential and every practical tool "
                   "approximates. Treat column boundaries in gappy regions as "
                   "indicative."),
    }


def score(sequences: Sequence[str], names: Optional[Sequence[str]] = None,
          *, positions: Optional[Sequence[int]] = None) -> Dict[str, Any]:
    """Per-column conservation across an alignment of homologs.

    Positions, when given, are numbered along the ANCHOR (the longest input),
    because that is the sequence a user is designing against.
    """
    seqs = [_clean(s) for s in (sequences or [])]
    if len([s for s in seqs if s]) < MIN_SEQS:
        return {"ok": False, "kind": "too_few",
                "error": f"Conservation needs at least {MIN_SEQS} sequences; "
                         f"got {len([s for s in seqs if s])}.",
                "next": "Add orthologs — BLAST the sequence and take the hits."}

    msa = align_many(seqs, names)
    if not msa.get("ok"):
        return msa

    rows = msa["rows"]
    anchor_row = next(r for r in rows if r["is_anchor"])["aligned"]
    n = len(rows)

    cols: List[Dict[str, Any]] = []
    anchor_pos = 0
    for c in range(msa["columns"]):
        column = [r["aligned"][c] for r in rows]
        anchor_ch = anchor_row[c]
        if anchor_ch != GAP:
            anchor_pos += 1
        residues = [ch for ch in column if ch != GAP]
        if not residues:
            continue
        counts = Counter(residues)
        top, top_n = counts.most_common(1)[0]
        # Shannon entropy over observed residues. 0 = every sequence agrees.
        total = len(residues)
        H = -sum((k / total) * math.log2(k / total) for k in counts.values())
        cols.append({
            "column": c,
            "anchor_position": anchor_pos if anchor_ch != GAP else None,
            "anchor_residue": None if anchor_ch == GAP else anchor_ch,
            "consensus": top,
            "agreement_pct": round(100.0 * top_n / total, 1),
            # abs(): -sum(...) of an all-agreeing column yields IEEE -0.0, which
            # prints as "-0.0 bits" and reads like a bug. Entropy is never
            # negative.
            "entropy_bits": abs(round(H, 3)),
            "gaps": n - total,
            # A plain word, because "0.0 bits" is not what a bench scientist
            # reads. Thresholds are stated rather than hidden.
            "call": ("invariant" if H == 0.0 else
                     "highly conserved" if H < 0.5 else
                     "conserved" if H < 1.0 else
                     "variable"),
        })

    wanted = None
    if positions:
        want = {int(p) for p in positions}
        wanted = [c for c in cols if c["anchor_position"] in want]

    invariant = [c for c in cols if c["call"] == "invariant"]
    mean_id = msa["mean_identity_to_anchor"]
    redundant = mean_id > REDUNDANT_ABOVE
    return {
        "ok": True,
        "sequences": n,
        "anchor": msa["anchor"],
        "columns": len(cols),
        "invariant_count": len(invariant),
        "conservation": wanted if wanted is not None else cols,
        "mean_identity": mean_id,
        "alignment_method": msa["method"],
        # The limit that is about the INPUT, not the algorithm. Forty
        # sequences at 99% identity are one sequence counted forty times, and
        # every column will look invariant.
        "diversity_warning": (
            f"The homologs are {mean_id}% identical to each other on average. "
            f"At that redundancy nearly every column looks conserved because "
            f"the set carries little independent evidence. Use more divergent "
            f"orthologs." if redundant else None),
        "trustworthy": (not redundant) and n >= MIN_SEQS,
        "caveat": ("Conservation is an OBSERVATION about the sequences given, "
                   "not a property of the protein. It is only as meaningful as "
                   "the homolog set: too few or too similar and it says "
                   "nothing. It complements ESM-2 rather than confirming it — "
                   "agreement between the two is genuine corroboration, "
                   "disagreement is worth investigating, not averaging."),
    }