# geney/results.py — Typed result dataclasses for pipeline outputs.
from __future__ import annotations

from dataclasses import dataclass, field, fields
from typing import Optional

import numpy as np
import pandas as pd


class _DictCompatMixin:
    """Mixin providing dict-like access for backward compatibility."""

    def __getitem__(self, key):
        try:
            return getattr(self, key)
        except AttributeError:
            raise KeyError(key)

    def get(self, key, default=None):
        return getattr(self, key, default)

    def keys(self):
        return [f.name for f in fields(self)]

    def __contains__(self, key):
        return hasattr(self, key)


# ============================================================================
# Splicing Result
# ============================================================================

@dataclass
class SplicingResult:
    """Result of splicing / Oncosplice analysis.

    Stores shared reference data once and per-isoform results in a DataFrame.

    Attributes:
        score: Max oncosplice_score across isoforms (0 = silent, higher = more disruptive).
        missplicing: Max splice-site probability delta.
        isoforms: DataFrame with per-isoform columns only (variant_protein,
            variant_mrna, oncosplice_score, percentile, isoform_prevalence, etc.).
        reference_protein: Reference protein sequence (shared across isoforms).
        reference_mrna: Reference mature mRNA sequence (shared across isoforms).
        conservation_vector: Per-residue conservation (shared across isoforms).
        splice_events: Splice site change metadata (donor/acceptor events JSON,
            proximity info).
    """
    # Summary
    score: float = 0.0                      # max splice-site probability delta
    missplicing: float = 0.0                # same as score (explicit name)
    max_oncosplice_score: float = 0.0       # max oncosplice_score across isoforms

    # Per-isoform data (variant-specific columns only)
    isoforms: pd.DataFrame = field(default_factory=pd.DataFrame)

    # Shared reference sequences (stored once, not per row)
    reference_protein: str = ""
    reference_mrna: str = ""
    conservation_vector: Optional[np.ndarray] = None

    # Shared metadata
    mut_id: str = ""
    gene: str = ""
    transcript_id: Optional[str] = None
    splicing_engine: str = ""
    central_position: int = 0

    # Splice site data (full per-position probabilities for visualization)
    donor_df: Optional[pd.DataFrame] = None     # all donors: ref_prob, event_prob, deltas
    acceptor_df: Optional[pd.DataFrame] = None  # all acceptors: ref_prob, event_prob, deltas

    # Splice event metadata (stored once)
    splice_events: Optional[pd.Series] = None

    @property
    def empty(self) -> bool:
        return self.isoforms.empty

    @property
    def columns(self) -> pd.Index:
        """Backward-compat: columns of the isoforms DataFrame."""
        return self.isoforms.columns

    def __len__(self) -> int:
        return len(self.isoforms)

    def to_dataframe(self) -> pd.DataFrame:
        """Reconstruct the full flat DataFrame (for backward compatibility)."""
        if self.isoforms.empty:
            return pd.DataFrame()
        df = self.isoforms.copy()
        df["mut_id"] = self.mut_id
        df["gene"] = self.gene
        df["transcript_id"] = self.transcript_id
        df["splicing_engine"] = self.splicing_engine
        df["central_position"] = self.central_position
        df["reference_protein"] = self.reference_protein
        df["reference_mrna"] = self.reference_mrna
        df["conservation_vector"] = [self.conservation_vector] * len(df)
        if self.splice_events is not None:
            for k, v in self.splice_events.items():
                df[k] = v
        return df


# ============================================================================
# TIS Result
# ============================================================================

@dataclass
class TISSite:
    """A single translation initiation site."""
    position: int
    location: str                # 'canonical', 'upstream', 'downstream'
    in_frame: bool
    ref_probability: float
    var_probability: float
    prob_change: float
    ref_usage: float = 0.0
    var_usage: float = 0.0


@dataclass
class TISResult(_DictCompatMixin):
    """Result of TIS (translation initiation site) shift analysis.

    Attributes:
        score: TIS shift score (0 = canonical intact, 1 = fully shifted away).
        status: 'completed', 'skipped', or 'error'.
        canonical: The canonical TIS site with ref/var probabilities.
        gained_alternatives: Sites that gained significant probability in the variant.
        all_sites: Every ATG site with ref/var probabilities.
        initiation_intact: Whether canonical initiation is preserved (within 20%).
    """
    status: str
    score: float                            # tis_shift_score

    # Canonical TIS
    canonical: Optional[TISSite] = None
    canonical_prob_ref: float = 0.0
    canonical_prob_var: float = 0.0
    canonical_prob_change: float = 0.0
    canonical_usage_ref: float = 0.0
    canonical_usage_var: float = 0.0
    canonical_usage_change: float = 0.0
    initiation_intact: bool = True

    # Alternatives
    gained_alternative_tis: list[dict] = field(default_factory=list)
    num_alternative_tis: int = 0
    top_alternative_position: Optional[int] = None
    top_alternative_prob: float = 0.0
    top_alternative_location: Optional[str] = None
    top_alternative_in_frame: Optional[bool] = None

    # All ATG sites
    atg_sites: list[dict] = field(default_factory=list)

    # Sequences
    ref_mrna_length: int = 0
    var_mrna_length: int = 0
    ref_mrna_window: str = ""
    var_mrna_window: str = ""
    mutation_mrna_positions: list[int] = field(default_factory=list)

    # Metadata
    mut_id: str = ""
    gene: str = ""
    transcript_id: Optional[str] = None
    summary: Optional[pd.Series] = None

    # Raw result for advanced use
    tis_result: object = None               # TISShiftResult

    # Skip/error metadata
    reason: Optional[str] = None
    proximity_window: Optional[int] = None

    # Legacy aliases
    @property
    def tis_shift_score(self) -> float:
        return self.score


# ============================================================================
# Folding Result
# ============================================================================

@dataclass
class FoldingWindow:
    """A single folding window around a mutation site."""
    ref_sequence: str
    var_sequence: str
    ref_structure: str              # dot-bracket notation
    var_structure: str              # dot-bracket notation
    ref_mfe: float                  # kcal/mol
    var_mfe: float                  # kcal/mol
    delta_mfe: float                # var_mfe - ref_mfe
    mutation_offset: int            # position within window


@dataclass
class FoldingMutation(_DictCompatMixin):
    """Folding result for a single mutation position."""
    genomic_position: int
    ref: str
    alt: str
    status: str                     # 'ok', 'not_in_mrna', 'ref_mismatch', 'fold_error'
    delta_mfe: Optional[float] = None
    score: Optional[float] = None
    mrna_position: Optional[int] = None
    coding_alt: Optional[str] = None
    window: Optional[FoldingWindow] = None
    # Error details
    expected_ref: Optional[str] = None
    error: Optional[str] = None

    # Legacy dict keys from _central_window_dict
    @property
    def ref_window(self) -> Optional[str]:
        return self.window.ref_sequence if self.window else None

    @property
    def var_window(self) -> Optional[str]:
        return self.window.var_sequence if self.window else None

    @property
    def ref_structure(self) -> Optional[str]:
        return self.window.ref_structure if self.window else None

    @property
    def var_structure(self) -> Optional[str]:
        return self.window.var_structure if self.window else None

    @property
    def mfe_ref(self) -> Optional[float]:
        return self.window.ref_mfe if self.window else None

    @property
    def mfe_var(self) -> Optional[float]:
        return self.window.var_mfe if self.window else None

    @property
    def mutation_offset(self) -> Optional[int]:
        return self.window.mutation_offset if self.window else None


@dataclass
class FoldingResult(_DictCompatMixin):
    """Result of mRNA folding analysis.

    Attributes:
        score: Folding score (from delta-MFE normalization).
        delta_mfe: Mean delta-MFE across windows (kcal/mol).
            Positive = destabilized, negative = stabilized.
        mutations: Per-position folding results.
        central_window: The best-centered folding window (if available).
    """
    status: str
    score: float = 0.0
    delta_mfe: Optional[float] = None
    mutations: list[dict] = field(default_factory=list)
    central_window: Optional[FoldingWindow] = None

    # Metadata
    mut_id: str = ""
    gene: str = ""
    transcript_id: Optional[str] = None
    window_size: int = 39

    # Error
    reason: Optional[str] = None


# ============================================================================
# Transcription Result
# ============================================================================

@dataclass
class TranscriptionResult(_DictCompatMixin):
    """Result of Borzoi transcription-effect prediction.

    Attributes:
        score: Regulatory score (0-1 normalized from max |log2FC|).
        tissue_fold_changes: Per-tissue log2 fold-changes.
        max_lfc_tissue: Tissue with largest absolute fold-change.
        max_lfc_value: log2FC of that tissue (signed).
        num_affected_tissues: Count of tissues exceeding threshold.
    """
    status: str
    score: float = 0.0                      # regulatory_score

    tissue_fold_changes: dict[str, float] = field(default_factory=dict)
    tissue_predicted_levels: dict[str, float] = field(default_factory=dict)
    max_lfc_tissue: str = ""
    max_lfc_value: float = 0.0
    max_abs_lfc: float = 0.0
    num_affected_tissues: int = 0

    # Metadata
    mut_id: str = ""
    gene: str = ""

    # Error
    reason: Optional[str] = None

    # Legacy alias
    @property
    def regulatory_score(self) -> float:
        return self.score


# ============================================================================
# Mutation Result (top-level pipeline output)
# ============================================================================

@dataclass
class MutationResult:
    """Result of the unified mutation pipeline.

    Attributes:
        splicing: Splicing / Oncosplice analysis result.
        tis: TIS shift analysis result.
        folding: mRNA folding analysis result.
        transcription: Borzoi transcription-effect result.
        transcript: The mutated transcript object.
        mutation_context: Pre-mRNA clone around the mutation site.
        timing: Per-stage wall-clock durations in seconds.
        errors: Pipeline name → error message for any pipeline that failed.
    """
    splicing: SplicingResult = field(default_factory=SplicingResult)
    tis: Optional[TISResult] = None
    folding: Optional[FoldingResult] = None
    transcription: Optional[TranscriptionResult] = None
    transcript: object = None
    mutation_context: object = None
    timing: dict = field(default_factory=dict)
    errors: dict = field(default_factory=dict)

    # Backward-compat: namedtuple-like interface
    _fields = ("splicing", "tis", "folding", "transcription", "transcript", "mutation_context")

    def __iter__(self):
        return iter((self.splicing, self.tis, self.folding,
                     self.transcription, self.transcript, self.mutation_context))

    def __getitem__(self, idx):
        return tuple(self)[idx]
