# geney/protein/analysis.py — Protein change classification and scoring.
"""
Takes reference and variant protein sequences (from direct translation of
the mature mRNA) and classifies the mutation type, scores functional impact,
and predicts NMD eligibility.

This pipeline answers: "What happened to the protein?"
  - synonymous: no AA change
  - missense: one or more AA substitutions, same length
  - in_frame_insertion: extra AAs inserted, no frameshift
  - in_frame_deletion: AAs deleted, no frameshift
  - frameshift: reading frame disrupted → different C-terminal sequence
  - nonsense: premature stop codon introduced
  - start_loss: start codon destroyed
  - stop_loss: stop codon destroyed → read-through
  - no_protein: variant produces no protein
"""
from __future__ import annotations

from dataclasses import dataclass, field
from typing import Optional

import numpy as np


@dataclass
class ProteinChangeResult:
    """Result of protein-level mutation analysis."""
    status: str

    # Mutation classification
    mutation_type: str = "unknown"          # missense, frameshift, nonsense, etc.
    mutation_description: str = ""          # Human-readable summary

    # Protein metrics
    ref_length: int = 0
    var_length: int = 0
    length_change: int = 0                  # var - ref (negative = shorter)
    truncation_fraction: float = 0.0        # 1 - var/ref (positive = truncated)

    # Missense details (for missense/in-frame changes)
    n_mismatches: int = 0
    mismatch_positions: list[int] = field(default_factory=list)
    first_mismatch: Optional[int] = None

    # Oncosplice conservation scoring
    oncosplice_score: float = 0.0
    oncosplice_percentile: float = 0.0
    conservation_at_changes: list[float] = field(default_factory=list)
    mean_conservation_at_change: float = 0.0
    has_conservation_data: bool = False

    # NMD prediction
    has_ptc: bool = False
    nmd_eligible: bool = False
    nmd_reason: str = ""

    # Score (0-10 normalized)
    score: float = 0.0

    # Metadata
    mut_id: str = ""
    gene: str = ""
    transcript_id: Optional[str] = None
    reason: Optional[str] = None


def _classify_mutation(ref_prot: str, var_prot: str) -> tuple[str, str, int]:
    """Classify mutation type from protein sequences.

    Returns (mutation_type, description, first_mismatch_position).
    """
    if not var_prot:
        return "no_protein", "variant produces no protein", 0

    # Start loss
    if ref_prot and var_prot and ref_prot[0] == "M" and var_prot[0] != "M":
        return "start_loss", "start codon destroyed", 0

    ref_has_stop = ref_prot.endswith("*")
    var_has_stop = var_prot.endswith("*")

    ref_clean = ref_prot.rstrip("*")
    var_clean = var_prot.rstrip("*")

    if ref_clean == var_clean:
        if ref_has_stop and not var_has_stop:
            return "stop_loss", "stop codon destroyed, read-through expected", len(ref_clean)
        return "synonymous", "no amino acid change", -1

    # Find first mismatch
    min_len = min(len(ref_clean), len(var_clean))
    first_mm = -1
    for i in range(min_len):
        if ref_clean[i] != var_clean[i]:
            first_mm = i
            break

    if first_mm == -1:
        # One is a prefix of the other
        if len(var_clean) < len(ref_clean):
            # Check if variant has a premature stop
            if "*" in var_prot[:-1] if var_has_stop else "*" in var_prot:
                return "nonsense", f"premature stop at position {len(var_clean) + 1}", len(var_clean)
            return "in_frame_deletion", f"C-terminal deletion of {len(ref_clean) - len(var_clean)} residues", len(var_clean)
        else:
            return "in_frame_insertion", f"C-terminal extension of {len(var_clean) - len(ref_clean)} residues", len(ref_clean)

    # Check for premature stop in variant
    if "*" in var_clean:
        stop_pos = var_clean.index("*")
        return "nonsense", f"premature stop at position {stop_pos + 1}", first_mm

    # Same length → missense or multi-missense
    if len(ref_clean) == len(var_clean):
        n_mm = sum(1 for a, b in zip(ref_clean, var_clean) if a != b)
        if n_mm == 1:
            return "missense", f"{ref_clean[first_mm]}{first_mm + 1}{var_clean[first_mm]}", first_mm
        return "missense", f"{n_mm} amino acid substitutions, first at position {first_mm + 1}", first_mm

    # Different lengths: check if reading frame is preserved
    # Frameshift: the sequence diverges and lengths differ by non-multiple-of-1
    # Actually, at the protein level we can't distinguish frameshift from
    # in-frame indel purely by length — we check if the tail is completely different
    tail_ref = ref_clean[first_mm:]
    tail_var = var_clean[first_mm:]

    # If the tails share <20% identity, it's likely a frameshift
    min_tail = min(len(tail_ref), len(tail_var))
    if min_tail > 0:
        matches = sum(1 for a, b in zip(tail_ref[:min_tail], tail_var[:min_tail]) if a == b)
        tail_identity = matches / min_tail
    else:
        tail_identity = 0.0

    if tail_identity < 0.2 and abs(len(ref_clean) - len(var_clean)) > 0:
        return "frameshift", f"reading frame disruption at position {first_mm + 1}", first_mm

    diff = len(var_clean) - len(ref_clean)
    if diff > 0:
        return "in_frame_insertion", f"{diff} residue insertion near position {first_mm + 1}", first_mm
    elif diff < 0:
        return "in_frame_deletion", f"{-diff} residue deletion near position {first_mm + 1}", first_mm
    else:
        n_mm = sum(1 for a, b in zip(ref_clean, var_clean) if a != b)
        return "missense", f"{n_mm} substitutions starting at position {first_mm + 1}", first_mm


def analyze_protein_change(
    reference_protein: str,
    variant_protein: str,
    conservation_vector: np.ndarray | None = None,
    exon_boundaries: list[int] | None = None,
    window_length: int = 13,
    mut_id: str = "",
    gene: str = "",
    transcript_id: str | None = None,
) -> ProteinChangeResult:
    """Analyze the protein-level consequence of a mutation.

    Takes the directly-translated reference and variant proteins and:
    1. Classifies the mutation type
    2. Scores functional impact via Oncosplice (conservation-weighted)
    3. Predicts NMD eligibility
    4. Produces a 0-10 normalized score

    Args:
        reference_protein: Reference protein sequence.
        variant_protein: Variant protein sequence (from direct translation).
        conservation_vector: Per-residue conservation scores. If None,
            Oncosplice scoring is skipped (but classification still runs).
        exon_boundaries: AA positions of exon-exon junctions (for NMD).
        window_length: Oncosplice smoothing window (default 13).
        mut_id: Mutation identifier.
        gene: Gene name.
        transcript_id: Transcript identifier.

    Returns:
        ProteinChangeResult with full classification and scoring.
    """
    ref = reference_protein or ""
    var = variant_protein or ""
    ref_clean = ref.rstrip("*")
    var_clean = var.rstrip("*")

    if not ref_clean:
        return ProteinChangeResult(
            status="error", reason="no_reference_protein",
            mut_id=mut_id, gene=gene, transcript_id=transcript_id)

    if not var_clean:
        return ProteinChangeResult(
            status="completed", mutation_type="no_protein",
            mutation_description="variant produces no protein",
            ref_length=len(ref_clean), score=10.0,
            mut_id=mut_id, gene=gene, transcript_id=transcript_id)

    # Classify
    mut_type, description, first_mm = _classify_mutation(ref, var)

    ref_len = len(ref_clean)
    var_len = len(var_clean)
    trunc = 1.0 - (var_len / ref_len) if ref_len > 0 else 0.0

    # Find all mismatches
    min_len = min(ref_len, var_len)
    mismatches = [i for i in range(min_len) if ref_clean[i] != var_clean[i]]

    # Oncosplice scoring
    onco_score = 0.0
    onco_pct = 0.0
    cons_at_changes: list[float] = []
    has_cons = False

    if conservation_vector is not None and hasattr(conservation_vector, '__len__') and len(conservation_vector) > 0:
        has_cons = not np.allclose(conservation_vector, 1.0)

        if ref_clean != var_clean:
            try:
                from ..oncosplice import Oncosplice
                onco = Oncosplice(ref_clean, var_clean, conservation_vector,
                                  window_length=window_length)
                onco_score = onco.score
                onco_pct = onco.percentile
            except Exception:
                pass

        # Conservation at change positions
        for pos in mismatches:
            if pos < len(conservation_vector):
                cons_at_changes.append(float(conservation_vector[pos]))

    # NMD prediction
    has_ptc = var_len < ref_len * 0.95 and mut_type in ("nonsense", "frameshift")
    nmd_eligible = False
    nmd_reason = "no_ptc"

    if has_ptc:
        from ..enrichment import predict_nmd_eligibility
        nmd = predict_nmd_eligibility(var, ref, exon_boundaries=exon_boundaries)
        nmd_eligible = nmd["nmd_eligible"]
        nmd_reason = nmd["reason"]

    # Score (0-10)
    score = _compute_score(mut_type, onco_score, trunc, has_ptc,
                           nmd_eligible, has_cons, len(mismatches))

    return ProteinChangeResult(
        status="completed",
        mutation_type=mut_type,
        mutation_description=description,
        ref_length=ref_len,
        var_length=var_len,
        length_change=var_len - ref_len,
        truncation_fraction=round(trunc, 4),
        n_mismatches=len(mismatches),
        mismatch_positions=mismatches,
        first_mismatch=first_mm if first_mm >= 0 else None,
        oncosplice_score=round(onco_score, 4),
        oncosplice_percentile=round(onco_pct, 4),
        conservation_at_changes=cons_at_changes,
        mean_conservation_at_change=(
            round(sum(cons_at_changes) / len(cons_at_changes), 4)
            if cons_at_changes else 0.0
        ),
        has_conservation_data=has_cons,
        has_ptc=has_ptc,
        nmd_eligible=nmd_eligible,
        nmd_reason=nmd_reason,
        score=round(score, 1),
        mut_id=mut_id,
        gene=gene,
        transcript_id=transcript_id,
    )


def _compute_score(
    mut_type: str,
    onco_score: float,
    truncation: float,
    has_ptc: bool,
    nmd_eligible: bool,
    has_conservation: bool,
    n_mismatches: int,
) -> float:
    """Compute 0-10 normalized score from protein change characteristics."""
    if mut_type == "synonymous":
        return 0.0
    if mut_type == "no_protein":
        return 10.0

    score = 0.0

    # Base score from mutation type severity
    type_base = {
        "missense": 2.0,
        "in_frame_insertion": 3.0,
        "in_frame_deletion": 3.5,
        "frameshift": 7.0,
        "nonsense": 7.0,
        "start_loss": 8.0,
        "stop_loss": 4.0,
    }.get(mut_type, 1.0)

    score = type_base

    # Conservation boost for missense (the Oncosplice value IS the score here)
    if mut_type == "missense" and has_conservation and onco_score > 0:
        score = max(score, min(10.0, onco_score))

    # NMD boost
    if has_ptc and nmd_eligible:
        score = max(score, 8.0)

    # Truncation severity
    if truncation > 0.5:
        score = max(score, 8.0)
    elif truncation > 0.2:
        score = max(score, 6.0)

    return min(10.0, score)
