"""
Visualization utilities for BADAS model validation.

Provides plotting functions for:
- Prediction time-series with gradient fills
- Event markers and annotations
- Model comparison plots
- PR/ROC curves
"""

import json
import numpy as np
import pandas as pd
import matplotlib
# Only use Agg backend if not in interactive environment (notebook/IPython)
try:
    from IPython import get_ipython
    if get_ipython() is None:
        matplotlib.use('Agg')
except (ImportError, AttributeError):
    matplotlib.use('Agg')
import matplotlib.pyplot as plt
from matplotlib.colors import LinearSegmentedColormap
from typing import Optional, List, Dict, Any, Union


# =============================================================================
# Constants
# =============================================================================

VIDEO_FPS_DEFAULT = 30  # Default video frame rate

# Risk-based colormap (green -> yellow -> orange -> red)
RISK_COLORS = ['#27ae60', '#f1c40f', '#e67e22', '#e74c3c']
RISK_CMAP = LinearSegmentedColormap.from_list('risk', RISK_COLORS)

# Standard colors
COLORS = {
    'smoothed': '#2980b9',      # Blue for smoothed predictions
    'raw': '#e74c3c',           # Red for raw predictions
    'caution': '#f39c12',       # Orange for caution threshold
    'danger': '#e74c3c',        # Red for danger threshold
    'event': '#8e44ad',         # Purple for event marker
}

# Default thresholds
DEFAULT_CAUTION_THRESHOLD = 0.5
DEFAULT_DANGER_THRESHOLD = 0.8


# =============================================================================
# Core Plotting Functions
# =============================================================================

def plot_prediction_timeseries(
    timestamps: Union[List[float], np.ndarray],
    probabilities: Union[List[float], np.ndarray],
    raw_probabilities: Optional[Union[List[float], np.ndarray]] = None,
    event_time: Optional[float] = None,
    event_frame: Optional[int] = None,
    event_type: Optional[str] = None,
    video_fps: int = VIDEO_FPS_DEFAULT,
    threshold: float = DEFAULT_CAUTION_THRESHOLD,
    danger_threshold: float = DEFAULT_DANGER_THRESHOLD,
    title: Optional[str] = None,
    ax: Optional[plt.Axes] = None,
    show_legend: bool = True,
    show_stats: bool = True,
    show_gradient_fill: bool = True,
    figsize: tuple = (12, 5),
) -> plt.Axes:
    """
    Plot prediction time series with professional styling.

    Args:
        timestamps: Array of timestamps in seconds
        probabilities: Array of (smoothed) probability values
        raw_probabilities: Optional array of raw probability values
        event_time: Time of event in seconds
        event_frame: Frame number of event (converted using video_fps)
        event_type: Type of event for annotation
        video_fps: Video frame rate for frame-to-time conversion
        threshold: Caution threshold (default: 0.5)
        danger_threshold: Danger threshold (default: 0.8)
        title: Plot title
        ax: Matplotlib axes (created if None)
        show_legend: Whether to show legend
        show_stats: Whether to show statistics annotation
        show_gradient_fill: Whether to show gradient fill under curve
        figsize: Figure size if creating new figure

    Returns:
        Matplotlib axes
    """
    if ax is None:
        fig, ax = plt.subplots(figsize=figsize)

    timestamps = np.array(timestamps)
    probabilities = np.array(probabilities)

    if len(timestamps) == 0 or len(probabilities) == 0:
        ax.text(0.5, 0.5, 'No predictions', ha='center', va='center',
                transform=ax.transAxes, fontsize=12)
        if title:
            ax.set_title(title, fontsize=11, fontweight='bold')
        return ax

    if raw_probabilities is None:
        raw_probabilities = probabilities
    else:
        raw_probabilities = np.array(raw_probabilities)

    # Calculate event time from frame if needed
    actual_event_time = event_time
    if actual_event_time is None and event_frame is not None and event_frame > 0:
        actual_event_time = event_frame / video_fps

    # Gradient fill under curve
    if show_gradient_fill and len(timestamps) > 1:
        for i in range(len(timestamps) - 1):
            color = RISK_CMAP(probabilities[i])
            ax.fill_between(
                [timestamps[i], timestamps[i+1]],
                [0, 0],
                [probabilities[i], probabilities[i+1]],
                alpha=0.5, color=color, linewidth=0
            )

    # Main smoothed line
    ax.plot(timestamps, probabilities, color=COLORS['smoothed'],
            linewidth=2, label='Prediction', zorder=5)

    # Raw probability if different from smoothed
    if raw_probabilities is not None:
        diff = np.abs(probabilities - raw_probabilities)
        if np.any(diff > 0.01):
            ax.plot(timestamps, raw_probabilities, color=COLORS['raw'],
                    linestyle='--', linewidth=1, alpha=0.5, label='Raw')

    # Threshold lines
    ax.axhline(y=threshold, color=COLORS['caution'], linestyle='--',
               linewidth=1.5, alpha=0.7, label=f'Caution ({threshold:.0%})')
    ax.axhline(y=danger_threshold, color=COLORS['danger'], linestyle='--',
               linewidth=1.5, alpha=0.7, label=f'Danger ({danger_threshold:.0%})')

    # Event marker
    if actual_event_time is not None:
        ax.axvline(x=actual_event_time, color=COLORS['event'], linestyle='-',
                   linewidth=2.5, alpha=0.9, zorder=10)

        event_label = f"{event_type.replace('_', ' ').title()}" if event_type else "Event"
        ax.annotate(
            event_label,
            xy=(actual_event_time, 0.95),
            xytext=(5, 0),
            textcoords='offset points',
            fontsize=9,
            fontweight='bold',
            color=COLORS['event'],
            verticalalignment='top',
        )

    # Styling
    ax.set_xlabel('Time (seconds)', fontsize=11)
    ax.set_ylabel('Collision Probability', fontsize=11)
    if title:
        ax.set_title(title, fontsize=12, fontweight='bold')

    ax.set_ylim([0, 1.05])
    ax.set_xlim([timestamps[0], timestamps[-1]])
    ax.grid(True, alpha=0.3)

    if show_legend:
        ax.legend(loc='upper right', fontsize=9, framealpha=0.9)

    # Statistics annotation
    if show_stats:
        max_risk = float(np.max(probabilities))
        max_idx = int(np.argmax(probabilities))
        max_time = timestamps[max_idx]
        safe_pct = np.sum(probabilities < threshold) / len(probabilities) * 100

        lead_time_str = ""
        if actual_event_time is not None:
            above_threshold = probabilities >= threshold
            times_above = timestamps[above_threshold]
            times_before_event = times_above[times_above < actual_event_time]
            if len(times_before_event) > 0:
                first_alert = times_before_event[0]
                lead_time = actual_event_time - first_alert
                lead_time_str = f" | Lead: {lead_time:.1f}s"

        stats_text = f"Max: {max_risk:.0%} @ {max_time:.1f}s | Safe: {safe_pct:.0f}%{lead_time_str}"
        ax.text(
            0.02, 0.98, stats_text,
            transform=ax.transAxes, fontsize=9,
            verticalalignment='top', fontfamily='monospace',
            bbox=dict(boxstyle='round,pad=0.3', facecolor='white', alpha=0.9, edgecolor='gray')
        )

    return ax


def plot_prediction_from_inference(
    predictions: List[Dict],
    event_time: Optional[float] = None,
    event_frame: Optional[int] = None,
    event_type: Optional[str] = None,
    video_fps: int = VIDEO_FPS_DEFAULT,
    ax: Optional[plt.Axes] = None,
    **kwargs
) -> plt.Axes:
    """
    Plot prediction time series from BADAS inference output.

    Args:
        predictions: List of prediction dicts from BADAS.predict_video()
                    Each dict has: timestamp, probability, raw_probability
        event_time: Time of event in seconds
        event_frame: Frame number of event
        event_type: Type of event
        video_fps: Video frame rate
        ax: Matplotlib axes
        **kwargs: Additional arguments passed to plot_prediction_timeseries

    Returns:
        Matplotlib axes
    """
    if not predictions:
        if ax is None:
            fig, ax = plt.subplots(figsize=kwargs.get('figsize', (12, 5)))
        ax.text(0.5, 0.5, 'No predictions', ha='center', va='center',
                transform=ax.transAxes, fontsize=12)
        return ax

    timestamps = [p.get('timestamp', i/8) for i, p in enumerate(predictions)]
    probabilities = [p['probability'] for p in predictions]
    raw_probs = [p.get('raw_probability', p['probability']) for p in predictions]

    return plot_prediction_timeseries(
        timestamps=timestamps,
        probabilities=probabilities,
        raw_probabilities=raw_probs,
        event_time=event_time,
        event_frame=event_frame,
        event_type=event_type,
        video_fps=video_fps,
        ax=ax,
        **kwargs
    )


# =============================================================================
# Multi-Panel Figures
# =============================================================================

def plot_validation_results(
    results: Dict[str, List[Dict]],
    event_info: Optional[Dict[str, Dict]] = None,
    output_path: Optional[str] = None,
    max_videos: int = 9,
    video_fps: int = VIDEO_FPS_DEFAULT,
    threshold: float = DEFAULT_CAUTION_THRESHOLD,
    danger_threshold: float = DEFAULT_DANGER_THRESHOLD,
    figsize_per_plot: tuple = (6, 4),
) -> plt.Figure:
    """
    Create a multi-panel figure showing validation results.

    Args:
        results: Dict mapping video_id to list of prediction dicts
        event_info: Dict mapping video_id to event info
                   (keys: event_time/time_of_event, event_frame, event_type)
        output_path: Path to save figure
        max_videos: Maximum number of videos to plot
        video_fps: Video frame rate
        threshold: Decision threshold
        danger_threshold: Danger threshold
        figsize_per_plot: Size of each subplot

    Returns:
        Matplotlib figure
    """
    video_names = list(results.keys())
    n_plots = min(len(video_names), max_videos)

    if n_plots == 0:
        print("No videos to plot")
        return None

    cols = min(3, n_plots)
    rows = (n_plots + cols - 1) // cols

    fig_width = figsize_per_plot[0] * cols
    fig_height = figsize_per_plot[1] * rows

    fig, axes = plt.subplots(rows, cols, figsize=(fig_width, fig_height), squeeze=False)
    axes = axes.flatten()

    for idx, video_name in enumerate(video_names[:n_plots]):
        predictions = results[video_name]

        video_event = event_info.get(video_name, {}) if event_info else {}
        event_time = video_event.get('event_time') or video_event.get('time_of_event')
        event_frame = video_event.get('event_frame')
        event_type = video_event.get('event_type')

        title = f"{video_name[:30]}..." if len(video_name) > 30 else video_name
        if event_type:
            title += f" ({event_type})"

        plot_prediction_from_inference(
            predictions,
            event_time=event_time,
            event_frame=event_frame,
            event_type=event_type,
            video_fps=video_fps,
            ax=axes[idx],
            title=title,
            threshold=threshold,
            danger_threshold=danger_threshold,
            show_legend=(idx == 0),
        )

    for idx in range(n_plots, len(axes)):
        axes[idx].set_visible(False)

    plt.tight_layout()

    if output_path:
        plt.savefig(output_path, dpi=150, bbox_inches='tight', facecolor='white')
        print(f"Saved validation figure to: {output_path}")

    return fig


def plot_model_comparison(
    results: Dict[str, Dict[str, List[Dict]]],
    model_names: List[str],
    event_info: Optional[Dict[str, Dict]] = None,
    output_path: Optional[str] = None,
    video_fps: int = VIDEO_FPS_DEFAULT,
) -> plt.Figure:
    """
    Create a multi-panel figure comparing models on multiple videos.

    Layout: Each row = one video, each column = one model

    Args:
        results: Nested dict {video_name: {model_name: predictions}}
        model_names: List of model names (column order)
        event_info: Dict {video_name: {'event_time': float, 'event_type': str, ...}}
        output_path: Path to save figure
        video_fps: Video frame rate

    Returns:
        Matplotlib figure
    """
    video_names = list(results.keys())
    n_videos = len(video_names)
    n_models = len(model_names)

    if n_videos == 0:
        print("No videos to plot")
        return None

    fig_height = 4 * n_videos
    fig_width = 6 * n_models

    fig, axes = plt.subplots(n_videos, n_models, figsize=(fig_width, fig_height), squeeze=False)

    for i, video_name in enumerate(video_names):
        video_event = event_info.get(video_name, {}) if event_info else {}
        event_time = video_event.get('event_time') or video_event.get('time_of_event')
        event_frame = video_event.get('event_frame')
        event_type = video_event.get('event_type')

        for j, model_name in enumerate(model_names):
            predictions = results[video_name].get(model_name, [])

            if i == 0:
                title = f"{model_name}\n{video_name[:40]}"
            else:
                title = video_name[:40]

            plot_prediction_from_inference(
                predictions,
                event_time=event_time,
                event_frame=event_frame,
                event_type=event_type,
                video_fps=video_fps,
                ax=axes[i, j],
                title=title,
                show_legend=(i == 0 and j == n_models - 1),
                show_stats=True,
            )

    fig.suptitle('Collision Risk Comparison', fontsize=14, fontweight='bold', y=1.01)
    plt.tight_layout()

    if output_path:
        plt.savefig(output_path, dpi=150, bbox_inches='tight', facecolor='white')
        print(f"Saved comparison figure to: {output_path}")

    return fig


# =============================================================================
# Curve Plots (PR, ROC)
# =============================================================================

def plot_pr_curve(
    precision: np.ndarray,
    recall: np.ndarray,
    auc_pr: float,
    ax: Optional[plt.Axes] = None,
    label: Optional[str] = None,
    color: Optional[str] = None,
) -> plt.Axes:
    """Plot precision-recall curve."""
    if ax is None:
        fig, ax = plt.subplots(figsize=(8, 8))

    curve_label = f'{label} (AUC = {auc_pr:.3f})' if label else f'AUC = {auc_pr:.3f}'
    ax.plot(recall, precision, linewidth=2, label=curve_label, color=color)

    ax.set_xlabel('Recall', fontsize=12)
    ax.set_ylabel('Precision', fontsize=12)
    ax.set_title('Precision-Recall Curve', fontsize=14, fontweight='bold')
    ax.legend(loc='best', fontsize=11)
    ax.grid(True, alpha=0.3)
    ax.set_xlim(0, 1.05)
    ax.set_ylim(0, 1.05)

    return ax


def plot_roc_curve(
    fpr: np.ndarray,
    tpr: np.ndarray,
    auc_roc: float,
    ax: Optional[plt.Axes] = None,
    label: Optional[str] = None,
    color: Optional[str] = None,
) -> plt.Axes:
    """Plot ROC curve."""
    if ax is None:
        fig, ax = plt.subplots(figsize=(8, 8))

    curve_label = f'{label} (AUC = {auc_roc:.3f})' if label else f'AUC = {auc_roc:.3f}'
    ax.plot(fpr, tpr, linewidth=2, label=curve_label, color=color)
    ax.plot([0, 1], [0, 1], 'k--', linewidth=1, label='Random', alpha=0.5)

    ax.set_xlabel('False Positive Rate', fontsize=12)
    ax.set_ylabel('True Positive Rate', fontsize=12)
    ax.set_title('ROC Curve', fontsize=14, fontweight='bold')
    ax.legend(loc='lower right', fontsize=11)
    ax.grid(True, alpha=0.3)
    ax.set_xlim(0, 1.05)
    ax.set_ylim(0, 1.05)

    return ax
