"""Heuristic antimicrobial peptide (AMP) scoring backend."""

from __future__ import annotations

import math

from peptidegym.peptide.properties import (
    net_charge, mean_hydrophobicity, amphipathicity_score,
)


class HeuristicAMPScorer:
    """Scores peptide sequences for antimicrobial peptide fitness using biophysical heuristics."""

    def score(self, sequence: str, **kwargs) -> dict[str, float]:
        charge = net_charge(sequence)
        hydro = mean_hydrophobicity(sequence)
        amphi = amphipathicity_score(sequence)
        length = len(sequence)

        # ── activity_score (0-1) ──
        # Cationic preference: ideal charge 2-7
        if 2 <= charge <= 7:
            charge_score = 1.0
        elif charge > 7:
            charge_score = max(0.0, 1.0 - (charge - 7) / 10.0)
        elif charge > 0:
            charge_score = charge / 2.0
        else:
            charge_score = 0.0

        # Hydrophobicity: ideal range -0.5 to 0.5 for AMPs
        if -0.5 <= hydro <= 0.5:
            hydro_score = 1.0
        else:
            dist = min(abs(hydro - (-0.5)), abs(hydro - 0.5)) if hydro < -0.5 or hydro > 0.5 else 0.0
            hydro_score = max(0.0, 1.0 - dist / 3.0)

        # Amphipathicity contribution (normalize to 0-1, typical range 0-3)
        amphi_norm = min(1.0, amphi / 2.5)

        # Motif presence
        motif_patterns = ["KK", "KR", "RR", "LK"]
        motif_hits = sum(1 for pat in motif_patterns if pat in sequence)
        motif_score = min(1.0, motif_hits / 2.0)

        activity_score = min(1.0, max(0.0,
            0.35 * charge_score + 0.25 * hydro_score + 0.25 * amphi_norm + 0.15 * motif_score
        ))

        # ── toxicity_penalty (0-1) ──
        charge_penalty = max(0.0, min(1.0, (charge - 10) / 5.0)) if charge > 10 else 0.0
        trp_density = sequence.count("W") / max(1, length)
        trp_penalty = max(0.0, min(1.0, (trp_density - 0.20) / 0.10)) if trp_density > 0.20 else 0.0
        toxicity_penalty = min(1.0, 0.6 * charge_penalty + 0.4 * trp_penalty)

        # ── amphipathicity_bonus (0-1) ──
        amphipathicity_bonus = min(1.0, amphi / 2.5)

        # ── length_fitness (0-1) ──
        length_fitness = math.exp(-0.5 * ((length - 20) / 10.0) ** 2)

        return {
            "activity_score": round(activity_score, 6),
            "toxicity_penalty": round(toxicity_penalty, 6),
            "amphipathicity_bonus": round(amphipathicity_bonus, 6),
            "length_fitness": round(length_fitness, 6),
        }
