# geney/tis/metrics.py — TIS shift analysis and metrics.
from __future__ import annotations

from dataclasses import dataclass
from typing import Optional

import numpy as np
import pandas as pd

from .model import _load_titer_model, get_codon_score
from .prediction import (
    predict_context_score,
    predict_context_scores_batch,
    extract_tis_context,
    find_all_atg_positions,
)


# ============================================================================
# TIS Shift Analysis and Metrics
# ============================================================================

@dataclass
class TISShiftResult:
    """
    Results of TIS shift analysis comparing reference to variant.

    Key metrics for understanding translation initiation changes.
    """
    # Canonical TIS metrics
    canonical_ref_prob: float          # Canonical TIS probability in reference
    canonical_var_prob: float          # Canonical TIS probability in variant
    canonical_prob_change: float       # Absolute change (var - ref)
    canonical_prob_fold_change: float  # Fold change (var / ref)

    # Usage fractions (normalized probabilities)
    canonical_ref_usage: float         # Fraction of initiation at canonical in ref
    canonical_var_usage: float         # Fraction of initiation at canonical in var
    canonical_usage_change: float      # Change in usage fraction

    # Alternative TIS metrics
    num_alternative_tis: int           # Number of alternative TIS candidates
    top_alternative_position: Optional[int]     # Position of strongest alternative
    top_alternative_prob: float        # Probability of strongest alternative
    top_alternative_location: Optional[str]     # upstream/downstream
    top_alternative_in_frame: Optional[bool]    # In-frame with canonical?

    # Summary scores
    tis_shift_score: float             # 0-1 score indicating shift away from canonical
    initiation_intact: bool            # Is canonical initiation preserved?

    # Raw data
    analysis_df: pd.DataFrame          # Full analysis DataFrame


def compute_tis_shift_metrics(
    ref_mrna: str,
    var_mrna: str,
    canonical_start: Optional[int] = None,
    probability_threshold: float = 0.1,
    use_prior: bool = False,
) -> TISShiftResult:
    """
    Compute comprehensive TIS shift metrics comparing reference to variant.

    This is the main function for analyzing how translation initiation changes
    between reference and variant sequences.

    Implements the logic from the pseudocode:
    - Score = relative change from canonical: (prob - canonical_prob) / canonical_prob
    - Tracks upstream/downstream/canonical status
    - Tracks in-frame/out-of-frame status
    - Filters based on context changes and probability thresholds

    Args:
        ref_mrna: Reference mRNA sequence
        var_mrna: Variant mRNA sequence
        canonical_start: Position of canonical start codon (auto-detected if None)
        probability_threshold: Minimum ContextScore to consider a site (default 0.1)
        use_prior: Whether to use CodonScore prior (default False, uses ContextScore only)

    Returns:
        TISShiftResult with all metrics
    """
    model = _load_titer_model()

    # Auto-detect canonical start if not provided
    if canonical_start is None:
        ref_atgs = find_all_atg_positions(ref_mrna)
        if not ref_atgs:
            raise ValueError("No ATG codons found in reference mRNA")
        canonical_start = ref_atgs[0]

    # Find all ATG positions
    ref_atg_positions = find_all_atg_positions(ref_mrna)
    var_atg_positions = find_all_atg_positions(var_mrna)

    # Compute ContextScores for reference (batched)
    ref_contexts = [extract_tis_context(ref_mrna, pos) for pos in ref_atg_positions]
    ref_scores = predict_context_scores_batch(ref_contexts) if ref_contexts else np.array([])
    ref_probs = dict(zip(ref_atg_positions, ref_scores.tolist())) if len(ref_scores) else {}

    # Compute ContextScores for variant (batched)
    var_contexts = [extract_tis_context(var_mrna, pos) for pos in var_atg_positions]
    var_scores = predict_context_scores_batch(var_contexts) if var_contexts else np.array([])
    var_probs = dict(zip(var_atg_positions, var_scores.tolist())) if len(var_scores) else {}

    # Optionally apply codon prior to get full TISScore
    if use_prior:
        for pos in ref_probs:
            codon = ref_mrna[pos:pos+3].upper()
            ref_probs[pos] *= get_codon_score(codon)
        for pos in var_probs:
            codon = var_mrna[pos:pos+3].upper()
            var_probs[pos] *= get_codon_score(codon)

    # Canonical TIS metrics
    canonical_ref_prob = ref_probs.get(canonical_start, 0.0)
    canonical_var_prob = var_probs.get(canonical_start, 0.0)
    canonical_prob_change = canonical_var_prob - canonical_ref_prob

    if canonical_ref_prob > 0:
        canonical_prob_fold_change = canonical_var_prob / canonical_ref_prob
    else:
        canonical_prob_fold_change = float('inf') if canonical_var_prob > 0 else 1.0

    # Compute usage fractions (normalized probabilities)
    ref_total = sum(p for p in ref_probs.values() if p >= probability_threshold)
    var_total = sum(p for p in var_probs.values() if p >= probability_threshold)

    if ref_total > 0:
        canonical_ref_usage = canonical_ref_prob / ref_total
    else:
        canonical_ref_usage = 1.0 if canonical_ref_prob > 0 else 0.0

    if var_total > 0:
        canonical_var_usage = canonical_var_prob / var_total
    else:
        canonical_var_usage = 1.0 if canonical_var_prob > 0 else 0.0

    canonical_usage_change = canonical_var_usage - canonical_ref_usage

    # Find alternative TIS candidates in variant
    alternatives = []
    for pos, prob in var_probs.items():
        if pos == canonical_start:
            continue
        if prob < probability_threshold:
            continue

        location = 'upstream' if pos < canonical_start else 'downstream'
        distance = pos - canonical_start
        in_frame = (distance % 3) == 0

        alternatives.append({
            'position': pos,
            'probability': prob,
            'location': location,
            'in_frame': in_frame,
            'ref_probability': ref_probs.get(pos, 0.0),
        })

    # Sort by probability
    alternatives.sort(key=lambda x: x['probability'], reverse=True)

    # Top alternative
    if alternatives:
        top_alt = alternatives[0]
        top_alternative_position = top_alt['position']
        top_alternative_prob = top_alt['probability']
        top_alternative_location = top_alt['location']
        top_alternative_in_frame = top_alt['in_frame']
    else:
        top_alternative_position = None
        top_alternative_prob = 0.0
        top_alternative_location = None
        top_alternative_in_frame = None

    # TIS shift score: how much initiation shifted away from canonical
    # Score of 0 = all initiation at canonical, 1 = no initiation at canonical
    if var_total > 0:
        tis_shift_score = 1.0 - canonical_var_usage
    else:
        tis_shift_score = 0.0 if canonical_var_prob > 0 else 1.0

    # Is canonical initiation preserved? (within 20% of reference usage)
    initiation_intact = (
        canonical_var_prob >= canonical_ref_prob * 0.8 and
        canonical_var_usage >= canonical_ref_usage * 0.8
    )

    # Build analysis DataFrame
    all_positions = sorted(set(ref_atg_positions) | set(var_atg_positions))
    rows = []
    for pos in all_positions:
        ref_p = ref_probs.get(pos, 0.0)
        var_p = var_probs.get(pos, 0.0)

        if pos == canonical_start:
            location = 'canonical'
        elif pos < canonical_start:
            location = 'upstream'
        else:
            location = 'downstream'

        distance = pos - canonical_start
        in_frame = (distance % 3) == 0

        rows.append({
            'position': pos,
            'location': location,
            'in_frame': in_frame,
            'ref_probability': round(ref_p, 4),
            'var_probability': round(var_p, 4),
            'prob_change': round(var_p - ref_p, 4),
            'ref_usage': round(ref_p / ref_total, 4) if ref_total > 0 else 0.0,
            'var_usage': round(var_p / var_total, 4) if var_total > 0 else 0.0,
        })

    analysis_df = pd.DataFrame(rows)
    if not analysis_df.empty:
        analysis_df = analysis_df.sort_values('position').reset_index(drop=True)

    return TISShiftResult(
        canonical_ref_prob=round(canonical_ref_prob, 4),
        canonical_var_prob=round(canonical_var_prob, 4),
        canonical_prob_change=round(canonical_prob_change, 4),
        canonical_prob_fold_change=round(canonical_prob_fold_change, 4),
        canonical_ref_usage=round(canonical_ref_usage, 4),
        canonical_var_usage=round(canonical_var_usage, 4),
        canonical_usage_change=round(canonical_usage_change, 4),
        num_alternative_tis=len(alternatives),
        top_alternative_position=top_alternative_position,
        top_alternative_prob=round(top_alternative_prob, 4),
        top_alternative_location=top_alternative_location,
        top_alternative_in_frame=top_alternative_in_frame,
        tis_shift_score=round(tis_shift_score, 4),
        initiation_intact=initiation_intact,
        analysis_df=analysis_df,
    )


def get_tis_summary(result: TISShiftResult) -> pd.Series:
    """
    Get a pandas Series with key TIS metrics for display/export.

    Args:
        result: TISShiftResult from compute_tis_shift_metrics

    Returns:
        Series with displayable metrics
    """
    return pd.Series({
        'canonical_prob_ref': result.canonical_ref_prob,
        'canonical_prob_var': result.canonical_var_prob,
        'canonical_prob_change': result.canonical_prob_change,
        'canonical_usage_ref': f"{result.canonical_ref_usage:.1%}",
        'canonical_usage_var': f"{result.canonical_var_usage:.1%}",
        'canonical_usage_change': f"{result.canonical_usage_change:+.1%}",
        'tis_shift_score': result.tis_shift_score,
        'initiation_intact': result.initiation_intact,
        'num_alternative_tis': result.num_alternative_tis,
        'top_alternative_position': result.top_alternative_position,
        'top_alternative_prob': result.top_alternative_prob,
        'top_alternative_location': result.top_alternative_location,
        'top_alternative_in_frame': result.top_alternative_in_frame,
    })


# ============================================================================
# Cross-mutation batched TIS analysis
# ============================================================================

def _build_tis_result_from_probs(
    ref_mrna: str,
    var_mrna: str,
    ref_atg_positions: list[int],
    var_atg_positions: list[int],
    ref_probs: dict[int, float],
    var_probs: dict[int, float],
    canonical_start: int,
    probability_threshold: float,
    use_prior: bool,
) -> TISShiftResult:
    """Build a TISShiftResult from pre-computed probability dicts.

    Shared logic extracted from :func:`compute_tis_shift_metrics` so both
    single-mutation and batch paths produce identical results.
    """
    if use_prior:
        for pos in ref_probs:
            codon = ref_mrna[pos:pos + 3].upper()
            ref_probs[pos] *= get_codon_score(codon)
        for pos in var_probs:
            codon = var_mrna[pos:pos + 3].upper()
            var_probs[pos] *= get_codon_score(codon)

    canonical_ref_prob = ref_probs.get(canonical_start, 0.0)
    canonical_var_prob = var_probs.get(canonical_start, 0.0)
    canonical_prob_change = canonical_var_prob - canonical_ref_prob

    if canonical_ref_prob > 0:
        canonical_prob_fold_change = canonical_var_prob / canonical_ref_prob
    else:
        canonical_prob_fold_change = float('inf') if canonical_var_prob > 0 else 1.0

    ref_total = sum(p for p in ref_probs.values() if p >= probability_threshold)
    var_total = sum(p for p in var_probs.values() if p >= probability_threshold)

    if ref_total > 0:
        canonical_ref_usage = canonical_ref_prob / ref_total
    else:
        canonical_ref_usage = 1.0 if canonical_ref_prob > 0 else 0.0

    if var_total > 0:
        canonical_var_usage = canonical_var_prob / var_total
    else:
        canonical_var_usage = 1.0 if canonical_var_prob > 0 else 0.0

    canonical_usage_change = canonical_var_usage - canonical_ref_usage

    alternatives = []
    for pos, prob in var_probs.items():
        if pos == canonical_start or prob < probability_threshold:
            continue
        location = 'upstream' if pos < canonical_start else 'downstream'
        distance = pos - canonical_start
        in_frame = (distance % 3) == 0
        alternatives.append({
            'position': pos,
            'probability': prob,
            'location': location,
            'in_frame': in_frame,
            'ref_probability': ref_probs.get(pos, 0.0),
        })
    alternatives.sort(key=lambda x: x['probability'], reverse=True)

    if alternatives:
        top_alt = alternatives[0]
        top_alternative_position = top_alt['position']
        top_alternative_prob = top_alt['probability']
        top_alternative_location = top_alt['location']
        top_alternative_in_frame = top_alt['in_frame']
    else:
        top_alternative_position = None
        top_alternative_prob = 0.0
        top_alternative_location = None
        top_alternative_in_frame = None

    if var_total > 0:
        tis_shift_score = 1.0 - canonical_var_usage
    else:
        tis_shift_score = 0.0 if canonical_var_prob > 0 else 1.0

    initiation_intact = (
        canonical_var_prob >= canonical_ref_prob * 0.8
        and canonical_var_usage >= canonical_ref_usage * 0.8
    )

    all_positions = sorted(set(ref_atg_positions) | set(var_atg_positions))
    rows = []
    for pos in all_positions:
        ref_p = ref_probs.get(pos, 0.0)
        var_p = var_probs.get(pos, 0.0)
        if pos == canonical_start:
            location = 'canonical'
        elif pos < canonical_start:
            location = 'upstream'
        else:
            location = 'downstream'
        distance = pos - canonical_start
        in_frame = (distance % 3) == 0
        rows.append({
            'position': pos,
            'location': location,
            'in_frame': in_frame,
            'ref_probability': round(ref_p, 4),
            'var_probability': round(var_p, 4),
            'prob_change': round(var_p - ref_p, 4),
            'ref_usage': round(ref_p / ref_total, 4) if ref_total > 0 else 0.0,
            'var_usage': round(var_p / var_total, 4) if var_total > 0 else 0.0,
        })

    analysis_df = pd.DataFrame(rows)
    if not analysis_df.empty:
        analysis_df = analysis_df.sort_values('position').reset_index(drop=True)

    return TISShiftResult(
        canonical_ref_prob=round(canonical_ref_prob, 4),
        canonical_var_prob=round(canonical_var_prob, 4),
        canonical_prob_change=round(canonical_prob_change, 4),
        canonical_prob_fold_change=round(canonical_prob_fold_change, 4),
        canonical_ref_usage=round(canonical_ref_usage, 4),
        canonical_var_usage=round(canonical_var_usage, 4),
        canonical_usage_change=round(canonical_usage_change, 4),
        num_alternative_tis=len(alternatives),
        top_alternative_position=top_alternative_position,
        top_alternative_prob=round(top_alternative_prob, 4),
        top_alternative_location=top_alternative_location,
        top_alternative_in_frame=top_alternative_in_frame,
        tis_shift_score=round(tis_shift_score, 4),
        initiation_intact=initiation_intact,
        analysis_df=analysis_df,
    )


def compute_tis_shift_metrics_batch(
    pairs: list[tuple[str, str]],
    canonical_starts: list[int | None] | None = None,
    probability_threshold: float = 0.1,
    use_prior: bool = False,
) -> list[TISShiftResult]:
    """Compute TIS shift metrics for multiple ref/var mRNA pairs in one pass.

    Collects ALL ATG contexts from ALL mutations, runs a single mega-batch
    through ``predict_context_scores_batch()``, then scatters scores back to
    build per-mutation :class:`TISShiftResult` objects.

    Args:
        pairs: List of ``(ref_mrna, var_mrna)`` strings.
        canonical_starts: Optional list of canonical start positions (one per
            pair).  ``None`` entries are auto-detected.
        probability_threshold: Min probability for TIS consideration.
        use_prior: Whether to apply codon prior.

    Returns:
        List of :class:`TISShiftResult`, one per input pair.
    """
    if canonical_starts is None:
        canonical_starts = [None] * len(pairs)

    # -- Collect all ATG positions and contexts across mutations --
    per_mut_info: list[dict] = []
    all_contexts: list[str] = []  # mega-batch of contexts
    # Track where each mutation's contexts start/end in the flat list
    for i, (ref_mrna, var_mrna) in enumerate(pairs):
        cs = canonical_starts[i]
        if cs is None:
            ref_atgs = find_all_atg_positions(ref_mrna)
            if not ref_atgs:
                raise ValueError(f"No ATG codons found in reference mRNA (pair {i})")
            cs = ref_atgs[0]

        ref_atg_positions = find_all_atg_positions(ref_mrna)
        var_atg_positions = find_all_atg_positions(var_mrna)

        ref_ctxs = [extract_tis_context(ref_mrna, pos) for pos in ref_atg_positions]
        var_ctxs = [extract_tis_context(var_mrna, pos) for pos in var_atg_positions]

        ref_start = len(all_contexts)
        all_contexts.extend(ref_ctxs)
        ref_end = len(all_contexts)

        var_start = len(all_contexts)
        all_contexts.extend(var_ctxs)
        var_end = len(all_contexts)

        per_mut_info.append({
            'ref_mrna': ref_mrna,
            'var_mrna': var_mrna,
            'canonical_start': cs,
            'ref_atg_positions': ref_atg_positions,
            'var_atg_positions': var_atg_positions,
            'ref_slice': (ref_start, ref_end),
            'var_slice': (var_start, var_end),
        })

    # -- Single mega-batch prediction --
    if all_contexts:
        all_scores = predict_context_scores_batch(all_contexts)
    else:
        all_scores = np.array([], dtype=np.float32)

    # -- Scatter scores back and build results --
    results: list[TISShiftResult] = []
    for info in per_mut_info:
        rs, re = info['ref_slice']
        vs, ve = info['var_slice']

        ref_scores_slice = all_scores[rs:re] if re > rs else np.array([])
        var_scores_slice = all_scores[vs:ve] if ve > vs else np.array([])

        ref_probs = dict(zip(info['ref_atg_positions'], ref_scores_slice.tolist())) if len(ref_scores_slice) else {}
        var_probs = dict(zip(info['var_atg_positions'], var_scores_slice.tolist())) if len(var_scores_slice) else {}

        result = _build_tis_result_from_probs(
            info['ref_mrna'], info['var_mrna'],
            info['ref_atg_positions'], info['var_atg_positions'],
            ref_probs, var_probs,
            info['canonical_start'],
            probability_threshold, use_prior,
        )
        results.append(result)

    return results
