"""Heuristic epitope/neoantigen scoring backend."""

from __future__ import annotations

from peptidegym.peptide.properties import HYDROPHOBICITY
from peptidegym.peptide.selfpeptide import max_kmer_overlap

# HLA-A*02:01 anchor preferences
_P2_PREFERENCES: dict[str, float] = {
    "L": 0.9, "M": 0.7, "I": 0.5, "V": 0.4, "A": 0.3, "T": 0.3,
}
_CTERMINAL_PREFERENCES: dict[str, float] = {
    "V": 0.9, "L": 0.8, "I": 0.7, "A": 0.5, "T": 0.4,
}
_DEFAULT_ANCHOR_SCORE = 0.1


class HeuristicEpitopeScorer:
    """Scores peptide sequences for MHC-I epitope fitness (HLA-A*02:01 default)."""

    def score(self, sequence: str, hla_allele: str = "HLA-A*02:01", **kwargs) -> dict[str, float]:
        length = len(sequence)

        # ── binding_score (0-1) ──
        if length < 2:
            binding_score = 0.0
        else:
            # P2 = index 1 (second residue)
            p2_aa = sequence[1]
            p2_score = _P2_PREFERENCES.get(p2_aa, _DEFAULT_ANCHOR_SCORE)

            # C-terminal = last residue
            ct_aa = sequence[-1]
            ct_score = _CTERMINAL_PREFERENCES.get(ct_aa, _DEFAULT_ANCHOR_SCORE)

            binding_score = (p2_score + ct_score) / 2.0

        # ── self_dissimilarity (0-1) ──
        overlap = max_kmer_overlap(sequence)
        self_dissimilarity = 1.0 - overlap

        # ── immunogenicity_proxy (0-1) ──
        # TCR-facing residues: positions 3-6 (0-indexed), i.e. 4th through 7th residues
        if length >= 7:
            tcr_residues = sequence[3:7]
        elif length > 3:
            tcr_residues = sequence[3:length]
        else:
            tcr_residues = ""

        if tcr_residues:
            mean_h = sum(HYDROPHOBICITY[aa] for aa in tcr_residues) / len(tcr_residues)
            # Kyte-Doolittle range is roughly -4.5 to 4.5; rescale to 0-1
            immunogenicity_proxy = (mean_h + 4.5) / 9.0
            immunogenicity_proxy = max(0.0, min(1.0, immunogenicity_proxy))
        else:
            immunogenicity_proxy = 0.0

        # ── length_fitness (0-1) ──
        if length == 9:
            length_fitness = 1.0
        elif length in (8, 10):
            length_fitness = 0.8
        elif length == 11:
            length_fitness = 0.5
        else:
            length_fitness = 0.0

        return {
            "binding_score": round(binding_score, 6),
            "self_dissimilarity": round(self_dissimilarity, 6),
            "immunogenicity_proxy": round(immunogenicity_proxy, 6),
            "length_fitness": round(length_fitness, 6),
        }
