# geney/tis/analysis.py — TIS candidate analysis.
from __future__ import annotations

from dataclasses import dataclass
from typing import Dict, List, Optional

import pandas as pd

from .model import _load_titer_model
from .prediction import (
    predict_tis_probability,
    extract_tis_context,
    find_all_atg_positions,
)


# ============================================================================
# TIS Candidate Analysis
# ============================================================================

@dataclass
class TISCandidate:
    """Represents a potential translation initiation site."""
    position: int           # 0-based position in mRNA
    codon: str              # The codon (e.g., 'ATG')
    probability: float      # TITER probability
    ref_probability: float  # Reference probability (if comparing to ref)
    relative_score: float   # (prob - canonical_prob) / canonical_prob
    location: str           # 'upstream', 'canonical', or 'downstream'
    in_frame: bool          # True if in-frame with canonical start
    context_changed: bool   # True if surrounding context differs from ref


def analyze_tis_candidates(
    ref_mrna: str,
    var_mrna: str,
    canonical_start: int,
    threshold: float = 0.5,
    context_window: int = 10,
) -> pd.DataFrame:
    """
    Analyze translation initiation site candidates in variant vs reference mRNA.

    This implements the logic from the pseudocode:
    - Find all ATG codons
    - Compare mutation vs reference probabilities
    - Filter based on position and probability thresholds
    - Identify alternative TIS candidates

    Args:
        ref_mrna: Reference mRNA sequence
        var_mrna: Variant mRNA sequence
        canonical_start: 0-based position of canonical ATG in reference
        threshold: Minimum probability relative to canonical to consider as TIS (default 0.5)
        context_window: Window size to check for context changes (default 10)

    Returns:
        DataFrame with TIS candidates and their properties
    """
    model = _load_titer_model()

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

    # Get canonical probability in reference
    if canonical_start in ref_atg_positions:
        canonical_context = extract_tis_context(ref_mrna, canonical_start)
        canonical_prob = predict_tis_probability(canonical_context, model)
    else:
        canonical_prob = 0.0

    candidates = []

    # Analyze each ATG in the variant
    for pos in var_atg_positions:
        # Get context and probability
        var_context = extract_tis_context(var_mrna, pos)
        var_prob = predict_tis_probability(var_context, model)

        # Get reference probability if position exists in reference
        if pos in ref_atg_positions:
            ref_context = extract_tis_context(ref_mrna, pos)
            ref_prob = predict_tis_probability(ref_context, model)

            # Check if context changed
            context_start = max(0, pos - context_window)
            context_end = min(len(ref_mrna), pos + 3 + context_window)

            if context_end <= len(var_mrna):
                ref_local = ref_mrna[context_start:context_end]
                var_local = var_mrna[context_start:context_end]
                context_changed = ref_local != var_local
            else:
                context_changed = True
        else:
            ref_prob = 0.0
            context_changed = True  # New ATG

        # Determine location relative to canonical
        if pos < canonical_start:
            location = 'upstream'
        elif pos == canonical_start:
            location = 'canonical'
        else:
            location = 'downstream'

        # Check in-frame status
        distance_from_canonical = pos - canonical_start
        in_frame = (distance_from_canonical % 3) == 0

        # Calculate relative score
        if canonical_prob > 0:
            relative_score = (var_prob - canonical_prob) / canonical_prob
        else:
            relative_score = float('inf') if var_prob > 0 else 0.0

        # Apply filtering logic from pseudocode:
        # - Upstream with unchanged context and no probability increase: skip
        # - Upstream with mut_prob <= ref_prob + 15%: skip
        # - Keep if mut_prob > threshold * canonical_prob

        keep = True
        if location == 'upstream':
            if not context_changed:
                keep = False  # Unchanged context upstream - discard
            elif var_prob <= ref_prob * 1.15:
                keep = False  # No significant increase - discard

        if var_prob <= threshold * canonical_prob and location != 'canonical':
            keep = False  # Below threshold - discard

        if keep or location == 'canonical':
            candidates.append(TISCandidate(
                position=pos,
                codon='ATG',
                probability=var_prob,
                ref_probability=ref_prob,
                relative_score=relative_score,
                location=location,
                in_frame=in_frame,
                context_changed=context_changed,
            ))

    # Convert to DataFrame
    if not candidates:
        return pd.DataFrame(columns=[
            'position', 'codon', 'probability', 'ref_probability',
            'relative_score', 'location', 'in_frame', 'context_changed'
        ])

    df = pd.DataFrame([
        {
            'position': c.position,
            'codon': c.codon,
            'probability': round(c.probability, 4),
            'ref_probability': round(c.ref_probability, 4),
            'relative_score': round(c.relative_score, 4) if c.relative_score != float('inf') else float('inf'),
            'location': c.location,
            'in_frame': c.in_frame,
            'context_changed': c.context_changed,
        }
        for c in candidates
    ])

    # Sort by position
    df = df.sort_values('position').reset_index(drop=True)

    return df


def predict_alternative_tis(
    ref_mrna: str,
    var_mrna: str,
    canonical_start: Optional[int] = None,
    threshold: float = 0.5,
) -> Dict:
    """
    Predict alternative translation initiation sites for a variant.

    This is the main entry point for TIS analysis.

    Args:
        ref_mrna: Reference mRNA sequence
        var_mrna: Variant mRNA sequence
        canonical_start: Position of canonical start codon (auto-detected if None)
        threshold: Minimum probability relative to canonical (default 0.5)

    Returns:
        Dictionary with:
            - 'canonical_tis': canonical TIS info
            - 'alternative_tis': list of alternative TIS candidates
            - 'canonical_intact': bool, whether canonical start is preserved
            - 'strongest_alternative': strongest non-canonical TIS (if any)
            - 'analysis_df': full DataFrame of analysis
    """
    # 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]  # First ATG

    # Run analysis
    df = analyze_tis_candidates(
        ref_mrna=ref_mrna,
        var_mrna=var_mrna,
        canonical_start=canonical_start,
        threshold=threshold,
    )

    # Extract results
    canonical_row = df[df['location'] == 'canonical']
    alternative_rows = df[df['location'] != 'canonical']

    canonical_tis = None
    canonical_intact = False

    if not canonical_row.empty:
        row = canonical_row.iloc[0]
        canonical_tis = {
            'position': int(row['position']),
            'probability': float(row['probability']),
            'ref_probability': float(row['ref_probability']),
            'relative_change': float(row['probability'] - row['ref_probability']),
        }
        # Canonical is intact if probability is within 5% of reference
        canonical_intact = row['probability'] >= row['ref_probability'] * 0.95

    alternative_tis = []
    for _, row in alternative_rows.iterrows():
        alternative_tis.append({
            'position': int(row['position']),
            'probability': float(row['probability']),
            'location': row['location'],
            'in_frame': bool(row['in_frame']),
            'relative_score': float(row['relative_score']) if row['relative_score'] != float('inf') else None,
        })

    # Find strongest alternative
    strongest = None
    if alternative_tis:
        strongest = max(alternative_tis, key=lambda x: x['probability'])

    return {
        'canonical_tis': canonical_tis,
        'alternative_tis': alternative_tis,
        'canonical_intact': canonical_intact,
        'strongest_alternative': strongest,
        'analysis_df': df,
    }
