# geney/tis/visualization.py — TIS visualization functions.
from __future__ import annotations

from typing import Optional, Tuple

import numpy as np

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


# ============================================================================
# Visualization
# ============================================================================

def plot_tis_comparison(
    result: TISShiftResult,
    title: str = "Translation Initiation Site Usage",
    figsize: Tuple[int, int] = (12, 5),
    max_sites: int = 10,
    show_threshold: float = 0.05,
):
    """
    Create a visualization comparing TIS usage between reference and variant.

    Args:
        result: TISShiftResult from compute_tis_shift_metrics
        title: Plot title
        figsize: Figure size (width, height)
        max_sites: Maximum number of sites to show
        show_threshold: Minimum usage to display a site

    Returns:
        matplotlib Figure object
    """
    import matplotlib.pyplot as plt
    import matplotlib.patches as mpatches

    df = result.analysis_df.copy()

    # Filter to significant sites
    df = df[(df['ref_usage'] >= show_threshold) | (df['var_usage'] >= show_threshold)]
    df = df.head(max_sites)

    if df.empty:
        fig, ax = plt.subplots(figsize=figsize)
        ax.text(0.5, 0.5, 'No significant TIS detected', ha='center', va='center', fontsize=14)
        ax.set_xlim(0, 1)
        ax.set_ylim(0, 1)
        ax.axis('off')
        return fig

    # Create figure with two subplots
    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=figsize)

    # Color scheme
    colors = {
        'upstream': '#E74C3C',      # Red
        'canonical': '#2ECC71',     # Green
        'downstream': '#3498DB',    # Blue
    }

    # === Left plot: Bar comparison ===
    x = np.arange(len(df))
    width = 0.35

    bars1 = ax1.bar(x - width/2, df['ref_usage'], width, label='Reference',
                    color='lightgray', edgecolor='black', linewidth=1)

    bar_colors = [colors.get(loc, 'gray') for loc in df['location']]
    bars2 = ax1.bar(x + width/2, df['var_usage'], width, label='Variant',
                    color=bar_colors, edgecolor='black', linewidth=1)

    # Labels
    labels = []
    for _, row in df.iterrows():
        pos_label = f"{int(row['position'])}"
        if row['location'] == 'canonical':
            pos_label += '\n(canonical)'
        elif row['in_frame']:
            pos_label += '\n(in-frame)'
        else:
            pos_label += '\n(out-of-frame)'
        labels.append(pos_label)

    ax1.set_xlabel('TIS Position', fontsize=11)
    ax1.set_ylabel('Usage Fraction', fontsize=11)
    ax1.set_title('TIS Usage: Reference vs Variant', fontsize=12, fontweight='bold')
    ax1.set_xticks(x)
    ax1.set_xticklabels(labels, fontsize=9)
    ax1.legend(loc='upper right')
    ax1.set_ylim(0, max(df['ref_usage'].max(), df['var_usage'].max()) * 1.2)
    ax1.axhline(y=0, color='black', linewidth=0.5)

    # Add value labels on bars
    for bar, val in zip(bars2, df['var_usage']):
        if val > 0.02:
            ax1.text(bar.get_x() + bar.get_width()/2, bar.get_height() + 0.01,
                    f'{val:.0%}', ha='center', va='bottom', fontsize=8)

    # === Right plot: Shift summary ===
    ax2.set_xlim(0, 10)
    ax2.set_ylim(0, 10)
    ax2.axis('off')

    # Title
    ax2.text(5, 9.5, 'TIS Shift Summary', ha='center', va='top',
             fontsize=14, fontweight='bold')

    # Metrics
    y_pos = 8.0
    line_height = 0.9

    metrics = [
        ('Canonical Usage (Ref)', f"{result.canonical_ref_usage:.1%}"),
        ('Canonical Usage (Var)', f"{result.canonical_var_usage:.1%}"),
        ('Usage Change', f"{result.canonical_usage_change:+.1%}"),
        ('', ''),  # Spacer
        ('TIS Shift Score', f"{result.tis_shift_score:.2f}"),
        ('Initiation Intact', '\u2713 Yes' if result.initiation_intact else '\u2717 No'),
        ('', ''),  # Spacer
        ('Alternative TIS Count', str(result.num_alternative_tis)),
    ]

    if result.top_alternative_position is not None:
        metrics.extend([
            ('Top Alternative Position', str(result.top_alternative_position)),
            ('Top Alternative Prob', f"{result.top_alternative_prob:.2f}"),
            ('Top Alternative Location', result.top_alternative_location or 'N/A'),
        ])

    for label, value in metrics:
        if label:
            ax2.text(1, y_pos, label + ':', ha='left', va='center', fontsize=10)

            # Color the value based on content
            color = 'black'
            if 'Change' in label and result.canonical_usage_change < -0.1:
                color = '#E74C3C'  # Red for significant decrease
            elif 'Intact' in label:
                color = '#2ECC71' if result.initiation_intact else '#E74C3C'

            ax2.text(9, y_pos, value, ha='right', va='center', fontsize=10,
                    fontweight='bold', color=color)
        y_pos -= line_height

    # Legend for colors
    legend_patches = [
        mpatches.Patch(color=colors['upstream'], label='Upstream'),
        mpatches.Patch(color=colors['canonical'], label='Canonical'),
        mpatches.Patch(color=colors['downstream'], label='Downstream'),
    ]
    ax2.legend(handles=legend_patches, loc='lower center', ncol=3, fontsize=9)

    plt.suptitle(title, fontsize=14, fontweight='bold', y=1.02)
    plt.tight_layout()

    return fig


def plot_tis_probability_landscape(
    ref_mrna: str,
    var_mrna: str,
    canonical_start: Optional[int] = None,
    window: int = 500,
    figsize: Tuple[int, int] = (14, 6),
):
    """
    Plot the TIS probability landscape around the canonical start site.

    Shows probabilities for all ATG codons within a window around the canonical start.

    Args:
        ref_mrna: Reference mRNA sequence
        var_mrna: Variant mRNA sequence
        canonical_start: Position of canonical start (auto-detected if None)
        window: Window size around canonical start to display
        figsize: Figure size

    Returns:
        matplotlib Figure object
    """
    import matplotlib.pyplot as plt

    model = _load_titer_model()

    # Auto-detect canonical start
    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]

    # Define window
    start_pos = max(0, canonical_start - window)
    end_pos = min(len(ref_mrna), canonical_start + window)

    # Find ATGs in window
    ref_atgs = [p for p in find_all_atg_positions(ref_mrna) if start_pos <= p <= end_pos]
    var_atgs = [p for p in find_all_atg_positions(var_mrna) if start_pos <= p <= end_pos]

    # Compute probabilities
    ref_probs = {}
    for pos in ref_atgs:
        context = extract_tis_context(ref_mrna, pos)
        ref_probs[pos] = predict_tis_probability(context, model)

    var_probs = {}
    for pos in var_atgs:
        context = extract_tis_context(var_mrna, pos)
        var_probs[pos] = predict_tis_probability(context, model)

    # Create figure
    fig, (ax1, ax2) = plt.subplots(2, 1, figsize=figsize, sharex=True)

    # Plot reference
    for pos, prob in ref_probs.items():
        color = '#2ECC71' if pos == canonical_start else '#3498DB'
        ax1.bar(pos, prob, width=3, color=color, edgecolor='black', linewidth=0.5)

    ax1.axvline(x=canonical_start, color='green', linestyle='--', alpha=0.5, label='Canonical')
    ax1.set_ylabel('TIS Probability', fontsize=11)
    ax1.set_title('Reference', fontsize=12, fontweight='bold')
    ax1.set_ylim(0, 1.0)
    ax1.legend(loc='upper right')

    # Plot variant
    for pos, prob in var_probs.items():
        if pos == canonical_start:
            color = '#2ECC71'
        elif pos < canonical_start:
            color = '#E74C3C'  # Upstream - red
        else:
            color = '#3498DB'  # Downstream - blue
        ax2.bar(pos, prob, width=3, color=color, edgecolor='black', linewidth=0.5)

    ax2.axvline(x=canonical_start, color='green', linestyle='--', alpha=0.5)
    ax2.set_xlabel('Position in mRNA', fontsize=11)
    ax2.set_ylabel('TIS Probability', fontsize=11)
    ax2.set_title('Variant', fontsize=12, fontweight='bold')
    ax2.set_ylim(0, 1.0)

    plt.suptitle('TIS Probability Landscape', fontsize=14, fontweight='bold')
    plt.tight_layout()

    return fig
