"""
Heatmap video generation for BADAS.

Renders per-window spatial attribution maps as a video overlay on the
original footage.  Uses ``attention_map`` (always_compute background maps)
and ``hazard_map`` (triggered hazard-event maps) from prediction dicts.
Typically produced with ``attach_attribution(always_compute=True)``.

Typical usage::

    predictor.attach_attribution(block_range=(8, 12), always_compute=True)
    predictions = predictor.predict_video("video.mp4", stride=1)
    generate_heatmap_video("video.mp4", predictions, "video_heatmap.mp4")
"""

import numpy as np
import cv2
from pathlib import Path
from typing import List, Dict, Any, Optional

from badas.visualization.video_base import BaseVideoProcessor


class HeatmapVideoProcessor(BaseVideoProcessor):
    """Renders attribution heatmaps onto a video frame-by-frame.

    Each output frame shows the original video frame with the heatmap from
    the nearest inference window blended on top, plus an optional risk panel
    at the bottom (reusing OverlayRenderer for visual consistency with the
    regular overlay video).
    """

    def __init__(self):
        super().__init__()
        from badas.visualization.overlay import OverlayRenderer
        self.renderer = OverlayRenderer()

    # -------------------------------------------------------------------------
    # Public API
    # -------------------------------------------------------------------------

    def create_heatmap_video(
        self,
        video_path: str,
        predictions: List[Dict[str, Any]],
        output_path: Optional[str] = None,
        heatmap_alpha: float = 0.5,
        show_bbox: bool = False,
        show_risk: bool = True,
        quiet: bool = False,
    ) -> str:
        """Write a heatmap overlay video to *output_path* and return the path.

        Args:
            video_path:    Path to the source video.
            predictions:   Output of ``BADAS.predict_video()``.  Predictions
                           without an ``attention_map`` or ``hazard_map`` key are
                           skipped (the raw frame is written unchanged).
            output_path:   Destination path.  Defaults to
                           ``<stem>_heatmap.mp4`` next to the source video.
            heatmap_alpha: Blend weight of the heatmap layer (0 = invisible,
                           1 = opaque).  Default 0.5.
            show_bbox:     If True, draw bounding boxes from ``bboxes`` scaled
                           to the original video resolution.
            show_risk:     If True, render the standard risk panel at the
                           bottom of every frame (reuses OverlayRenderer).
            quiet:         Suppress progress messages.
        """
        attr_preds = [p for p in predictions if 'attention_map' in p or 'hazard_map' in p]
        if not attr_preds:
            raise ValueError(
                "No predictions contain an attention_map or hazard_map.  "
                "Call attach_attribution(always_compute=True) before predict_video()."
            )

        video_path_obj = Path(video_path)
        if output_path is None:
            output_path = str(video_path_obj.parent / f"{video_path_obj.stem}_heatmap.mp4")

        if self._has_ffmpeg:
            return self._create_with_ffmpeg(
                video_path, predictions, str(output_path),
                heatmap_alpha, show_bbox, show_risk, quiet,
            )
        return self._create_with_opencv(
            video_path, predictions, str(output_path),
            heatmap_alpha, show_bbox, show_risk, quiet,
        )

    # -------------------------------------------------------------------------
    # Frame rendering
    # -------------------------------------------------------------------------

    def _render_frame(
        self,
        frame: np.ndarray,
        pred: Dict[str, Any],
        heatmap_alpha: float,
        show_bbox: bool,
        show_risk: bool,
    ) -> np.ndarray:
        """Composite heatmap + optional overlays onto *frame* (BGR, in-place copy)."""
        height, width = frame.shape[:2]
        attr_map = pred.get('attention_map') or pred.get('hazard_map')

        # Upscale attr_map to frame dimensions and colourise
        heat_u8 = (cv2.resize(attr_map.astype(np.float32), (width, height),
                              interpolation=cv2.INTER_LINEAR) * 255).astype(np.uint8)
        heat_bgr = cv2.applyColorMap(heat_u8, cv2.COLORMAP_INFERNO)

        # Blend heatmap onto frame
        result = cv2.addWeighted(frame, 1.0 - heatmap_alpha, heat_bgr, heatmap_alpha, 0)

        # Bounding boxes — scaled from 256×256 model space to video resolution
        if show_bbox and pred.get('bboxes'):
            bx = width / 256.0
            by = height / 256.0
            for rank, box in enumerate(pred['bboxes']):
                x1, y1, x2, y2 = (int(box[0] * bx), int(box[1] * by),
                                   int(box[2] * bx), int(box[3] * by))
                # White border + coloured fill; primary box is brighter
                cv2.rectangle(result, (x1 - 1, y1 - 1), (x2 + 1, y2 + 1), (255, 255, 255), 1)
                alpha = max(80, int(220 * (1.0 - rank * 0.3)))
                cv2.rectangle(result, (x1, y1), (x2, y2),
                               (30, alpha // 3, alpha), max(1, 3 - rank))

        # Standard risk panel at the bottom
        if show_risk:
            result = self.renderer.draw_prediction_overlay(
                result,
                probability=pred['probability'],
                raw_probability=pred.get('raw_probability', pred['probability']),
                current_time=pred['timestamp'],
                latency_ms=pred.get('latency_ms', 0.0),
            )

        return result

    # -------------------------------------------------------------------------
    # Video writers
    # -------------------------------------------------------------------------

    def _create_with_ffmpeg(
        self,
        video_path: str,
        predictions: List[Dict[str, Any]],
        output_path: str,
        heatmap_alpha: float,
        show_bbox: bool,
        show_risk: bool,
        quiet: bool,
    ) -> str:
        cap, fps, width, height, total_frames = self._open_video(video_path)

        if not quiet:
            print(f"Creating heatmap video: {total_frames} frames @ {fps:.1f} fps...")

        frame_lookup = self._create_frame_lookup(predictions, fps, total_frames)

        proc = self._open_ffmpeg_pipe(width, height, fps, output_path)

        frame_idx = 0
        try:
            while True:
                ret, frame = cap.read()
                if not ret:
                    break
                pred = frame_lookup.get(frame_idx)
                if pred is not None and ('attention_map' in pred or 'hazard_map' in pred):
                    frame = self._render_frame(frame, pred, heatmap_alpha, show_bbox, show_risk)
                proc.stdin.write(frame.tobytes())
                frame_idx += 1
        finally:
            cap.release()
            proc.stdin.close()
            proc.wait()

        if not quiet:
            print(f"Saved to: {output_path}")
        return output_path

    def _create_with_opencv(
        self,
        video_path: str,
        predictions: List[Dict[str, Any]],
        output_path: str,
        heatmap_alpha: float,
        show_bbox: bool,
        show_risk: bool,
        quiet: bool,
    ) -> str:
        cap, fps, width, height, total_frames = self._open_video(video_path)

        if not quiet:
            print(f"Creating heatmap video: {total_frames} frames @ {fps:.1f} fps...")

        frame_lookup = self._create_frame_lookup(predictions, fps, total_frames)

        fourcc = cv2.VideoWriter_fourcc(*'mp4v')
        out = cv2.VideoWriter(output_path, fourcc, fps, (width, height))
        if not out.isOpened():
            out = cv2.VideoWriter(output_path,
                                  cv2.VideoWriter_fourcc(*'avc1'), fps, (width, height))

        frame_idx = 0
        try:
            while True:
                ret, frame = cap.read()
                if not ret:
                    break
                pred = frame_lookup.get(frame_idx)
                if pred is not None and ('attention_map' in pred or 'hazard_map' in pred):
                    frame = self._render_frame(frame, pred, heatmap_alpha, show_bbox, show_risk)
                out.write(frame)
                frame_idx += 1
        finally:
            cap.release()
            out.release()

        if not quiet:
            print(f"Saved to: {output_path}")
        return output_path
