"""Matplotlib-based prediction charts for BADAS."""

import textwrap
import numpy as np
from dataclasses import dataclass
from pathlib import Path
from typing import List, Dict, Any, Optional


@dataclass
class VisualizationConfig:
    """Configuration for visualizations."""
    # Colors (RGB for matplotlib, BGR for OpenCV)
    safe_color: tuple = (39, 174, 96)       # Green
    caution_color: tuple = (243, 156, 18)   # Orange
    danger_color: tuple = (231, 76, 60)     # Red

    # Thresholds — aligned with RiskConfig defaults
    caution_threshold: float = 0.5
    danger_threshold: float = 0.75

    # Figure settings
    figure_width: int = 14
    figure_height: int = 6


def visualize_predictions(
    predictions: List[Dict[str, Any]],
    title: str = "Collision Risk Analysis",
    show: bool = True,
    save_path: Optional[str] = None,
    config: VisualizationConfig = None,
    video_path: Optional[str] = None,
    show_preprocessing_latency: bool = False,
    show_heatmap_latency: bool = False,
    show_reasoning_latency: bool = False,
):
    """
    Create a clean visualization of collision predictions.

    Shows:
    - Risk timeline with color gradient
    - Latency distribution (if available)
    - For each attribution event: an arrow on the timeline pointing to the peak,
      with the frame+heatmap overlay embedded as an inset image above it.
      Requires video_path and predictions that contain 'hazard_map'.

    Args:
        predictions: List of prediction dicts with 'probability', 'timestamp', etc.
        title: Plot title
        show: Whether to display the plot
        save_path: Optional path to save the figure
        config: Visualization configuration
        video_path: Optional path to the source video. When provided alongside
                    predictions that contain 'hazard_map', the last frame of each
                    attribution window is rendered with the heatmap overlay and
                    embedded as an inset directly on the risk timeline.

    Returns:
        matplotlib Figure object
    """
    import matplotlib.pyplot as plt
    import matplotlib.gridspec as gridspec
    from matplotlib.colors import LinearSegmentedColormap

    if not predictions:
        print("No predictions to visualize")
        return None

    if config is None:
        config = VisualizationConfig()

    # Extract data
    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]
    preprocess_ms_vals  = [p.get('preprocess_ms', 0.0) for p in predictions]
    inference_ms_vals   = [p.get('inference_ms',
                                  p.get('latency_ms', 0.0) - p.get('preprocess_ms', 0.0))
                           for p in predictions]
    attribution_ms_vals = [p.get('attribution_ms', 0.0) for p in predictions]
    description_ms_vals = [p.get('description_ms', 0.0) for p in predictions]

    has_latencies      = any(v > 0 for v in inference_ms_vals) or any(v > 0 for v in preprocess_ms_vals)
    has_preprocess_ms  = any(v > 0 for v in preprocess_ms_vals)
    has_attribution_ms = any(v > 0 for v in attribution_ms_vals)
    has_description_ms = any(v > 0 for v in description_ms_vals)
    attr_preds = [p for p in predictions if 'hazard_map' in p]
    embed_images = video_path is not None and len(attr_preds) > 0

    # Layout: timeline always, latency optionally below.
    # When images are embedded on the timeline, the figure needs to be taller
    # to give the y-axis room above 1.0 for the inset images.
    height_ratios = [2, 1] if has_latencies else [2]
    base_height = config.figure_height * sum(height_ratios) / 3
    total_height = base_height * 1.6 if embed_images else base_height

    fig = plt.figure(figsize=(config.figure_width, total_height))
    gs = gridspec.GridSpec(len(height_ratios), 1,
                           height_ratios=height_ratios, hspace=0.35)

    ax1 = fig.add_subplot(gs[0])
    ax2 = fig.add_subplot(gs[1]) if has_latencies else None

    # Custom colormap: green → yellow → orange → red
    cmap = LinearSegmentedColormap.from_list('risk', ['#27ae60', '#f1c40f', '#e67e22', '#e74c3c'])

    # ============ Risk Timeline ============
    # When images are embedded, expand the y axis above 1.0 so the insets
    # (centred at y≈1.38) have space and the probability ticks stay clean.
    ylim_top = 1.80 if embed_images else 1.05
    ax1.set_ylim([0, ylim_top])

    for i in range(len(timestamps) - 1):
        ax1.fill_between([timestamps[i], timestamps[i+1]],
                         [probabilities[i], probabilities[i+1]],
                         alpha=0.6, color=cmap(probabilities[i]), linewidth=0)

    ax1.plot(timestamps, probabilities, 'b-', linewidth=2, label='Risk', zorder=5)

    if any(abs(p - r) > 0.01 for p, r in zip(probabilities, raw_probs)):
        ax1.plot(timestamps, raw_probs, 'r--', linewidth=1, alpha=0.4, label='Raw')

    ax1.axhline(y=config.caution_threshold, color='#f39c12', linestyle='--',
                linewidth=1.5, alpha=0.7, label=f'Caution ({int(config.caution_threshold*100)}%)')
    ax1.axhline(y=config.danger_threshold, color='#e74c3c', linestyle='--',
                linewidth=1.5, alpha=0.7, label=f'Danger ({int(config.danger_threshold*100)}%)')

    ax1.set_xlabel('Time (seconds)', fontsize=11)
    ax1.set_ylabel('Collision Probability', fontsize=11)
    ax1.set_title(title, fontsize=13, fontweight='bold')
    ax1.set_xlim([min(timestamps), max(timestamps)])
    ax1.set_yticks([0.0, 0.25, 0.5, 0.75, 1.0])
    ax1.grid(True, alpha=0.3)

    max_risk = max(probabilities)
    max_time = timestamps[probabilities.index(max_risk)]
    safe_pct = sum(1 for p in probabilities if p < config.caution_threshold) / len(probabilities) * 100
    stats_text = f"Max: {max_risk:.0%} @ {max_time:.1f}s | Safe: {safe_pct:.0f}%"
    ax1.text(0.02, 0.02, stats_text, transform=ax1.transAxes, fontsize=9,
             verticalalignment='bottom', fontfamily='monospace',
             bbox=dict(boxstyle='round,pad=0.3', facecolor='white', alpha=0.8))

    # ============ Attribution Event Insets ============
    if embed_images:
        import cv2
        from scipy.ndimage import zoom
        from matplotlib.offsetbox import OffsetImage, AnnotationBbox, TextArea, VPacker

        for pred in attr_preds:
            if pred.get('focus_confidence', 1.0) == 0.0:
                continue  # scattered heatmap — skip inset
            ts_ev = pred['timestamp']
            p_ev = pred['probability']
            attr_map = pred['hazard_map']
            bboxes    = pred.get('bboxes') or []
            risk = pred['probability']
            attr_ms = pred.get('attribution_ms', '?')
            frame_idx = pred.get('frame_index', 0)

            cap = cv2.VideoCapture(video_path)
            cap.set(cv2.CAP_PROP_POS_FRAMES, frame_idx)
            ret, frame_bgr = cap.read()
            cap.release()

            if not ret:
                # Fallback: plain arrow if frame unavailable
                ax1.annotate(
                    f't={ts_ev:.1f}s',
                    xy=(ts_ev, p_ev), xytext=(ts_ev, p_ev + 0.15),
                    ha='center', va='bottom', fontsize=8, color='#8e44ad',
                    arrowprops=dict(arrowstyle='->', color='#8e44ad', lw=1.5),
                )
                continue

            frame_256 = cv2.resize(frame_bgr, (256, 256))
            frame_rgb = cv2.cvtColor(frame_256, cv2.COLOR_BGR2RGB)

            # Subtle heatmap: 25% heat over 75% frame
            scale = 256 / attr_map.shape[0]
            heat = zoom(attr_map, scale, order=1).clip(0, 1)
            heat_rgb = (plt.cm.inferno(heat)[:, :, :3] * 255).astype(np.uint8)
            overlay = (frame_rgb * 0.75 + heat_rgb * 0.25).astype(np.uint8)

            # Clearly visible bboxes: white border + red fill, dimmer for secondary boxes
            for rank, box in enumerate(bboxes):
                x1, y1, x2, y2 = box[:4]
                alpha = max(80, int(220 * (1.0 - rank * 0.3)))
                thickness = max(1, 3 - rank)
                cv2.rectangle(overlay, (x1 - 1, y1 - 1), (x2 + 1, y2 + 1), (255, 255, 255), 1)
                cv2.rectangle(overlay, (x1, y1), (x2, y2), (alpha, 30, 30), thickness)

            # Pack label + image (+ optional description) into one annotation box
            title_color = '#e74c3c' if risk >= config.danger_threshold else '#8e44ad'
            label = f't={ts_ev:.1f}s   P={risk:.2f}   {attr_ms}ms'
            text_area = TextArea(label, textprops=dict(
                fontsize=7.5, color=title_color, fontweight='bold', ha='center'
            ))
            img_area = OffsetImage(overlay, zoom=0.35)
            children = [text_area, img_area]
            description = pred.get('description')
            if description:
                lines = textwrap.fill(description, 32).splitlines()
                if len(lines) > 2:
                    lines = lines[:2]
                    lines[-1] = lines[-1].rstrip() + '…'
                wrapped = '\n'.join(lines)
                desc_area = TextArea(wrapped, textprops=dict(
                    fontsize=7, color='#555555', ha='center'
                ))
                children.append(desc_area)
            packed = VPacker(children=children, pad=3, sep=3, align='center')

            # Arrow from inset box (xybox) down to the peak on the curve (xy).
            # clip_on=False lets the box render even if it sits near the axes edge.
            ab = AnnotationBbox(
                packed,
                xy=(ts_ev, p_ev),
                xybox=(ts_ev, 1.38),
                xycoords='data',
                boxcoords='data',
                arrowprops=dict(arrowstyle='->', color='#8e44ad', lw=1.5),
                bboxprops=dict(edgecolor='#8e44ad', linewidth=1.5, boxstyle='round,pad=0.15'),
                frameon=True,
                zorder=10,
            )
            ab.set_clip_on(False)
            ax1.add_artist(ab)

    elif attr_preds:
        # No video: plain arrow + dot for each attribution event (skip scattered)
        for pred in attr_preds:
            if pred.get('focus_confidence', 1.0) == 0.0:
                continue
            ts_ev = pred['timestamp']
            p_ev = pred['probability']
            ax1.annotate(
                f't={ts_ev:.1f}s',
                xy=(ts_ev, p_ev), xytext=(ts_ev, min(p_ev + 0.18, 0.98)),
                ha='center', va='bottom', fontsize=8, color='#8e44ad', fontweight='bold',
                arrowprops=dict(arrowstyle='->', color='#8e44ad', lw=1.5),
                zorder=10,
            )
            ax1.plot(ts_ev, p_ev, 'o', color='#8e44ad', markersize=7, zorder=10)

    ax1.legend(loc='upper right', fontsize=9)

    # ============ Latency Breakdown ============
    if ax2 is not None:
        stack_arrays = []
        stack_labels = []
        stack_colors = []

        if has_preprocess_ms and show_preprocessing_latency:
            stack_arrays.append(preprocess_ms_vals)
            stack_labels.append('Preprocessing')
            stack_colors.append('#27ae60')   # green

        stack_arrays.append(inference_ms_vals)
        stack_labels.append('Inference')
        stack_colors.append('#3498db')       # blue

        if has_attribution_ms and show_heatmap_latency:
            stack_arrays.append(attribution_ms_vals)
            stack_labels.append('Heatmap')
            stack_colors.append('#e67e22')   # orange

        if has_description_ms and show_reasoning_latency:
            stack_arrays.append(description_ms_vals)
            stack_labels.append('BADAS-Reasoning')
            stack_colors.append('#9b59b6')   # purple

        ax2.stackplot(timestamps, *stack_arrays,
                      labels=stack_labels, colors=stack_colors, alpha=0.85)

        total_ms_vals = [sum(v[i] for v in stack_arrays) for i in range(len(timestamps))]
        mean_total = np.mean(total_ms_vals)
        ax2.axhline(y=mean_total, color='#e74c3c', linestyle='--',
                    linewidth=1.5, label=f'Mean total: {mean_total:.1f}ms')

        ax2.set_xlabel('Time (seconds)', fontsize=11)
        ax2.set_ylabel('Latency (ms)', fontsize=11)
        ax2.set_title('Latency Breakdown', fontsize=12, fontweight='bold')
        ax2.legend(loc='upper right', fontsize=9)
        ax2.set_xlim([min(timestamps), max(timestamps)])
        ax2.grid(True, alpha=0.3)

    import warnings
    with warnings.catch_warnings():
        warnings.simplefilter("ignore")
        plt.tight_layout()

    if save_path:
        plt.savefig(save_path, dpi=150, bbox_inches='tight')
        print(f"Figure saved to: {save_path}")

    if show:
        plt.show()

    return fig


def plot_attribution_events(
    predictions: List[Dict[str, Any]],
    video_path: str,
    figsize: tuple = (15, 5),
    show: bool = True,
) -> List[Any]:
    """
    Visualise every attribution event in *predictions* as a 3-panel figure:
    original frame · heatmap · overlay with bounding box.

    Args:
        predictions: Output of ``BADAS.predict_video()``.
        video_path:  Path to the source video (used to read frames).
        figsize:     Matplotlib figure size per event. Default (15, 5).
        show:        Call ``plt.show()`` after each figure. Default True.

    Returns:
        List of ``matplotlib.figure.Figure`` objects, one per event.

    Example::

        figs = plot_attribution_events(video_predictions, VIDEO_PATH)
    """
    import cv2
    import matplotlib.pyplot as plt
    from scipy.ndimage import zoom

    attr_preds = [p for p in predictions if "hazard_map" in p]
    print(f"{len(attr_preds)} attribution event(s) to plot")

    figs = []
    for pred in attr_preds:
        frame_idx = pred["frame_index"]
        attr_map  = pred["hazard_map"]
        bboxes    = pred.get("bboxes") or []
        risk      = pred["probability"]
        ts        = pred["timestamp"]
        attr_ms   = pred.get("attribution_ms", "?")

        cap = cv2.VideoCapture(video_path)
        cap.set(cv2.CAP_PROP_POS_FRAMES, frame_idx)
        ret, frame_bgr = cap.read()
        cap.release()
        if not ret:
            print(f"  Could not read frame {frame_idx}, skipping")
            continue

        frame_256 = cv2.resize(frame_bgr, (256, 256))
        frame_rgb = cv2.cvtColor(frame_256, cv2.COLOR_BGR2RGB)

        heat = zoom(attr_map, 256 / attr_map.shape[0], order=1).clip(0, 1)

        heat_rgb = (plt.cm.inferno(heat)[:, :, :3] * 255).astype(np.uint8)
        overlay  = (frame_rgb * 0.5 + heat_rgb * 0.5).astype(np.uint8)

        for rank, box in enumerate(bboxes):
            x1, y1, x2, y2 = box[:4]
            alpha = max(80, int(220 * (1.0 - rank * 0.3)))
            cv2.rectangle(overlay, (x1, y1), (x2, y2),
                          (255, alpha // 3, alpha // 3), max(1, 2 - rank))

        fig, axes = plt.subplots(1, 3, figsize=figsize)

        axes[0].imshow(frame_rgb)
        axes[0].set_title(f"Frame {frame_idx}  t={ts:.2f}s", fontsize=10)
        axes[0].axis("off")

        axes[1].imshow(heat, cmap="inferno", vmin=0, vmax=1)
        axes[1].set_title(f"Attribution map  ({attr_ms}ms)", fontsize=10)
        axes[1].axis("off")

        axes[2].imshow(overlay)
        axes[2].set_title(
            f"Overlay + bbox  risk={risk:.3f}", fontsize=10,
            color="red" if risk >= 0.5 else "black",
        )
        axes[2].axis("off")

        plt.suptitle(f"t={ts:.2f}s  |  risk={risk:.3f}  |  {attr_ms}ms", fontsize=11)
        plt.tight_layout()

        if show:
            plt.show()

        figs.append(fig)

    return figs
