"""
Multi-class video overlay with voting-based temporal smoothing.

For each frame, aggregates predictions from all windows that contain it
using soft voting (probability averaging), then displays the smoothed result.
"""

import cv2
import numpy as np
import subprocess
from pathlib import Path
from typing import List, Dict, Any, Optional, Tuple
from collections import defaultdict

from badas.visualization.renderer import _check_ffmpeg
from badas.visualization.video_base import BaseVideoProcessor


# Class colors (BGR format for OpenCV)
CLASS_COLORS = {
    # Collision events - Red tones
    'collision_front': (0, 0, 220),
    'collision_rear_or_side': (0, 50, 200),
    'collision_curb': (0, 100, 180),
    'collision_other': (50, 50, 180),

    # Near collision / Hard brake - Orange tones
    'near_collision': (0, 140, 255),
    'hard_brake_with_3rd_party': (0, 165, 255),
    'hard_brake_without_3rd_party': (0, 200, 255),

    # Active driving - Blue tones
    'sharp_cornering': (255, 140, 0),
    'harsh_acceleration': (255, 100, 50),
    'passing_or_overtaking': (255, 180, 0),
    'intersection': (200, 150, 50),

    # Normal - Green
    'normal_driving': (0, 200, 0),
}

# Risk categories
HIGH_RISK_CLASSES = {
    'collision_front', 'collision_rear_or_side', 'collision_curb',
    'collision_other', 'near_collision'
}

MEDIUM_RISK_CLASSES = {
    'hard_brake_with_3rd_party', 'hard_brake_without_3rd_party'
}

# Display names (cleaner formatting)
DISPLAY_NAMES = {
    'collision_front': 'Front Collision',
    'collision_rear_or_side': 'Rear/Side Collision',
    'collision_curb': 'Curb Collision',
    'collision_other': 'Collision',
    'near_collision': 'Near Collision',
    'hard_brake_with_3rd_party': 'Hard Brake (3rd Party)',
    'hard_brake_without_3rd_party': 'Hard Brake',
    'harsh_acceleration': 'Harsh Acceleration',
    'sharp_cornering': 'Sharp Cornering',
    'passing_or_overtaking': 'Passing/Overtaking',
    'intersection': 'Intersection',
    'normal_driving': 'Normal Driving',
}


def get_class_color(class_name: str) -> Tuple[int, int, int]:
    """Get BGR color for a class."""
    return CLASS_COLORS.get(class_name, (200, 200, 200))


def get_display_name(class_name: str) -> str:
    """Get human-readable display name for a class."""
    return DISPLAY_NAMES.get(class_name, class_name.replace('_', ' ').title())


def get_risk_level(class_name: str) -> int:
    """Get risk level: 2=high, 1=medium, 0=low."""
    if class_name in HIGH_RISK_CLASSES:
        return 2
    elif class_name in MEDIUM_RISK_CLASSES:
        return 1
    return 0


class VotingSmoothedOverlay(BaseVideoProcessor):
    """
    Creates video overlays with voting-based temporal smoothing.

    For each video frame:
    1. Find all prediction windows that contain this frame
    2. Average their probability distributions (soft voting)
    3. Display the class with highest averaged probability
    """

    INIT_PERIOD = 1.0  # Initialization period in seconds

    def __init__(
        self,
        window_size: int = 16,
        model_fps: int = 8,
        stride: int = 4
    ):
        """
        Args:
            window_size: Number of frames per prediction window
            model_fps: FPS at which model processes video
            stride: Stride between prediction windows (in model frames)
        """
        super().__init__()
        self.window_size = window_size
        self.model_fps = model_fps
        self.stride = stride
        self.window_duration = window_size / model_fps  # seconds

    def compute_smoothed_predictions(
        self,
        predictions: List[Dict[str, Any]],
        video_fps: float,
        total_frames: int,
        class_names: Optional[List[str]] = None
    ) -> Dict[int, Dict[str, Any]]:
        """
        Compute smoothed predictions for each video frame using voting.

        For each frame, finds all windows that overlap with it and averages
        their probability distributions.

        Args:
            predictions: List of window predictions with 'timestamp' and 'all_probs'
            video_fps: FPS of the video
            total_frames: Total frames in the video
            class_names: List of class names (inferred from predictions if not provided)

        Returns:
            Dict mapping frame_idx to smoothed prediction dict
        """
        if not predictions:
            return {}

        # Infer class names from first prediction
        if class_names is None:
            class_names = list(predictions[0].get('all_probs', {}).keys())

        num_classes = len(class_names)
        if num_classes == 0:
            return {}

        # For each window, compute the frame range it covers (in video frames)
        # Window is centered at timestamp, spans window_duration seconds
        window_ranges = []
        for pred in predictions:
            center_time = pred['timestamp']
            start_time = center_time - self.window_duration / 2
            end_time = center_time + self.window_duration / 2

            start_frame = max(0, int(start_time * video_fps))
            end_frame = min(total_frames, int(end_time * video_fps))

            window_ranges.append({
                'start_frame': start_frame,
                'end_frame': end_frame,
                'prediction': pred
            })

        # For each video frame, aggregate predictions from overlapping windows
        frame_predictions = {}

        for frame_idx in range(total_frames):
            current_time = frame_idx / video_fps

            # Skip init period
            if current_time < self.INIT_PERIOD:
                continue

            # Find all windows that contain this frame
            overlapping_preds = []
            for wr in window_ranges:
                if wr['start_frame'] <= frame_idx < wr['end_frame']:
                    overlapping_preds.append(wr['prediction'])

            if not overlapping_preds:
                # Use nearest window if no overlap
                min_dist = float('inf')
                nearest = None
                for wr in window_ranges:
                    center = (wr['start_frame'] + wr['end_frame']) / 2
                    dist = abs(frame_idx - center)
                    if dist < min_dist:
                        min_dist = dist
                        nearest = wr['prediction']
                if nearest:
                    overlapping_preds = [nearest]

            if not overlapping_preds:
                continue

            # Soft voting: average probability distributions
            avg_probs = np.zeros(num_classes)
            for pred in overlapping_preds:
                all_probs = pred.get('all_probs', {})
                for i, cls in enumerate(class_names):
                    avg_probs[i] += all_probs.get(cls, 0)
            avg_probs /= len(overlapping_preds)

            # Get winning class
            pred_idx = int(np.argmax(avg_probs))
            pred_class = class_names[pred_idx]
            confidence = float(avg_probs[pred_idx])

            frame_predictions[frame_idx] = {
                'predicted_name': pred_class,
                'confidence': confidence,
                'num_votes': len(overlapping_preds),
                'timestamp': current_time
            }

        return frame_predictions

    def draw_overlay(
        self,
        frame: np.ndarray,
        predicted_class: str,
        confidence: float,
        current_time: float,
        num_votes: int = 1,
        timeline_history: Optional[List[Tuple[str, float]]] = None,
        total_duration: float = 0
    ) -> np.ndarray:
        """
        Draw a clean, single-action overlay on the frame.

        Args:
            frame: Input BGR frame
            predicted_class: Name of predicted class
            confidence: Confidence (0-1)
            current_time: Current video time in seconds
            num_votes: Number of windows that voted
            timeline_history: List of (class_name, timestamp) for timeline
            total_duration: Total video duration for timeline scaling

        Returns:
            Frame with overlay
        """
        height, width = frame.shape[:2]

        # Get class info
        display_name = get_display_name(predicted_class)
        color = get_class_color(predicted_class)
        risk_level = get_risk_level(predicted_class)

        # Format text with probability
        label_text = f"{display_name}  {confidence*100:.0f}%"

        # --- Timeline at bottom ---
        timeline_height = 12
        timeline_margin = 15
        timeline_y = height - timeline_height - timeline_margin
        timeline_x = timeline_margin
        timeline_width = width - 2 * timeline_margin

        # Timeline background
        cv2.rectangle(frame, (timeline_x, timeline_y),
                     (timeline_x + timeline_width, timeline_y + timeline_height),
                     (30, 30, 30), -1)

        # Draw timeline segments from history
        if timeline_history and total_duration > 0:
            for i, (cls, ts) in enumerate(timeline_history):
                # Calculate segment position
                start_x = timeline_x + int((ts / total_duration) * timeline_width)

                # Get end position from next entry or current time
                if i + 1 < len(timeline_history):
                    end_x = timeline_x + int((timeline_history[i + 1][1] / total_duration) * timeline_width)
                else:
                    end_x = timeline_x + int((current_time / total_duration) * timeline_width)

                # Draw colored segment
                seg_color = get_class_color(cls)
                if end_x > start_x:
                    cv2.rectangle(frame, (start_x, timeline_y),
                                 (end_x, timeline_y + timeline_height),
                                 seg_color, -1)

        # Timeline border
        cv2.rectangle(frame, (timeline_x, timeline_y),
                     (timeline_x + timeline_width, timeline_y + timeline_height),
                     (80, 80, 80), 1)

        # Current position marker (white line)
        if total_duration > 0:
            marker_x = timeline_x + int((current_time / total_duration) * timeline_width)
            cv2.line(frame, (marker_x, timeline_y - 3),
                    (marker_x, timeline_y + timeline_height + 3),
                    (255, 255, 255), 2)

        # --- Main label above timeline ---
        font = cv2.FONT_HERSHEY_SIMPLEX
        font_scale = 1.0
        thickness = 2

        text_size = cv2.getTextSize(label_text, font, font_scale, thickness)[0]
        padding_x = 30
        padding_y = 15

        label_width = text_size[0] + padding_x * 2
        label_height = text_size[1] + padding_y * 2
        label_x = (width - label_width) // 2
        label_y = timeline_y - label_height - 10

        # Opaque black background
        cv2.rectangle(frame, (label_x, label_y),
                     (label_x + label_width, label_y + label_height),
                     (0, 0, 0), -1)

        # Colored border based on class
        border_thickness = 4 if risk_level == 2 else 3 if risk_level == 1 else 2
        cv2.rectangle(frame, (label_x, label_y),
                     (label_x + label_width, label_y + label_height),
                     color, border_thickness)

        # Class name + probability (centered)
        text_x = label_x + padding_x
        text_y = label_y + padding_y + text_size[1]

        cv2.putText(frame, label_text, (text_x, text_y),
                   font, font_scale, color, thickness)

        # --- Timestamp in top-left corner ---
        time_text = f"{current_time:.1f}s"
        cv2.putText(frame, time_text, (15, 30),
                   cv2.FONT_HERSHEY_SIMPLEX, 0.7, (200, 200, 200), 2)

        # --- Color legend (compact, top-right) ---
        # Groups: Collision (red), Near Miss (orange), Hard Brake (yellow),
        #         Maneuver (blue), Normal (green)
        legend_items = [
            ('Collision', (0, 0, 220)),           # Red - front/rear/curb/other
            ('Near Collision', (0, 140, 255)),    # Orange
            ('Hard Brake', (0, 200, 255)),        # Yellow-orange
            ('Acceleration', (255, 100, 50)),     # Blue-ish
            ('Cornering', (255, 140, 0)),         # Blue
            ('Passing', (255, 180, 0)),           # Light blue
            ('Intersection', (200, 150, 50)),     # Teal
            ('Normal', (0, 200, 0)),              # Green
        ]

        legend_x = width - 130
        legend_y = 12
        item_height = 16
        dot_radius = 4

        # Opaque background for legend
        legend_bg_height = len(legend_items) * item_height + 8
        cv2.rectangle(frame, (legend_x - 8, legend_y - 4),
                     (width - 8, legend_y + legend_bg_height),
                     (0, 0, 0), -1)
        cv2.rectangle(frame, (legend_x - 8, legend_y - 4),
                     (width - 8, legend_y + legend_bg_height),
                     (60, 60, 60), 1)

        for i, (label, clr) in enumerate(legend_items):
            y_pos = legend_y + i * item_height + 10
            # Color dot
            cv2.circle(frame, (legend_x, y_pos - 3), dot_radius, clr, -1)
            # Label
            cv2.putText(frame, label, (legend_x + 10, y_pos),
                       cv2.FONT_HERSHEY_SIMPLEX, 0.35, (180, 180, 180), 1)

        # --- Warning indicator for high risk ---
        if risk_level == 2 and confidence > 0.4:
            # Red border around entire frame
            cv2.rectangle(frame, (0, 0), (width-1, height-1), color, 6)

            # Warning icon in top-right
            warn_x = width - 50
            warn_y = 40
            cv2.circle(frame, (warn_x, warn_y), 20, color, -1)
            cv2.putText(frame, "!", (warn_x - 6, warn_y + 8),
                       cv2.FONT_HERSHEY_SIMPLEX, 1.0, (255, 255, 255), 3)

        return frame

    def draw_init_overlay(
        self,
        frame: np.ndarray,
        current_time: float
    ) -> np.ndarray:
        """Draw initialization overlay."""
        height, width = frame.shape[:2]

        # Darken frame
        overlay = frame.copy()
        cv2.rectangle(overlay, (0, 0), (width, height), (0, 0, 0), -1)
        frame = cv2.addWeighted(frame, 0.4, overlay, 0.6, 0)

        # Center box
        box_width = 300
        box_height = 80
        box_x = (width - box_width) // 2
        box_y = (height - box_height) // 2

        cv2.rectangle(frame, (box_x, box_y),
                     (box_x + box_width, box_y + box_height),
                     (40, 40, 40), -1)
        cv2.rectangle(frame, (box_x, box_y),
                     (box_x + box_width, box_y + box_height),
                     (100, 100, 100), 2)

        # Text
        text = "Initializing..."
        text_size = cv2.getTextSize(text, cv2.FONT_HERSHEY_SIMPLEX, 0.8, 2)[0]
        text_x = box_x + (box_width - text_size[0]) // 2
        cv2.putText(frame, text, (text_x, box_y + 35),
                   cv2.FONT_HERSHEY_SIMPLEX, 0.8, (200, 200, 200), 2)

        # Progress bar
        progress = min(1.0, current_time / self.INIT_PERIOD)
        bar_width = box_width - 40
        bar_x = box_x + 20
        bar_y = box_y + 55

        cv2.rectangle(frame, (bar_x, bar_y),
                     (bar_x + bar_width, bar_y + 8),
                     (60, 60, 60), -1)
        fill_width = int(bar_width * progress)
        if fill_width > 0:
            cv2.rectangle(frame, (bar_x, bar_y),
                         (bar_x + fill_width, bar_y + 8),
                         (0, 200, 0), -1)

        return frame

    def create_overlay_video(
        self,
        video_path: str,
        predictions: List[Dict[str, Any]],
        output_path: Optional[str] = None,
        quiet: bool = False
    ) -> str:
        """
        Create a video with smoothed prediction overlays.

        Args:
            video_path: Path to input video
            predictions: List of prediction dicts from inference
            output_path: Path for output video (auto-generated if None)
            quiet: Suppress progress output

        Returns:
            Path to output video
        """
        if output_path is None:
            p = Path(video_path)
            output_path = str(p.parent / f"{p.stem}_classified.mp4")

        # Open video
        cap, fps, width, height, total_frames = self._open_video(video_path)

        if not quiet:
            print(f"Video: {total_frames} frames @ {fps:.1f}fps")
            print(f"Computing smoothed predictions with voting...")

        # Compute smoothed predictions using voting
        frame_predictions = self.compute_smoothed_predictions(
            predictions, fps, total_frames
        )

        if not quiet:
            print(f"Smoothed {len(frame_predictions)} frames from {len(predictions)} windows")

        # Create output video
        if self._has_ffmpeg:
            return self._create_with_ffmpeg(
                cap, frame_predictions, output_path,
                fps, width, height, total_frames, quiet
            )
        else:
            return self._create_with_opencv(
                cap, frame_predictions, output_path,
                fps, width, height, total_frames, quiet
            )

    def _build_timeline_history(
        self,
        frame_predictions: Dict[int, Dict],
        fps: float,
        total_frames: int
    ) -> List[Tuple[str, float]]:
        """Build timeline history from frame predictions (class changes only)."""
        timeline = []
        last_class = None

        for frame_idx in range(total_frames):
            pred = frame_predictions.get(frame_idx)
            if pred is not None:
                cls = pred['predicted_name']
                ts = frame_idx / fps
                if cls != last_class:
                    timeline.append((cls, ts))
                    last_class = cls

        return timeline

    def _create_with_ffmpeg(
        self,
        cap: cv2.VideoCapture,
        frame_predictions: Dict[int, Dict],
        output_path: str,
        fps: float,
        width: int,
        height: int,
        total_frames: int,
        quiet: bool
    ) -> str:
        """Create video using FFmpeg pipe."""
        ffmpeg_cmd = [
            'ffmpeg', '-y',
            '-f', 'rawvideo',
            '-vcodec', 'rawvideo',
            '-s', f'{width}x{height}',
            '-pix_fmt', 'bgr24',
            '-r', str(fps),
            '-i', '-',
            '-c:v', 'libx264', '-preset', 'fast',
            '-pix_fmt', 'yuv420p',
            '-movflags', '+faststart',
            output_path
        ]

        proc = subprocess.Popen(
            ffmpeg_cmd,
            stdin=subprocess.PIPE,
            stdout=subprocess.DEVNULL,
            stderr=subprocess.DEVNULL if quiet else subprocess.PIPE
        )

        # Build full timeline history upfront
        full_timeline = self._build_timeline_history(frame_predictions, fps, total_frames)
        total_duration = total_frames / fps

        frame_idx = 0
        last_pred = None

        try:
            while True:
                ret, frame = cap.read()
                if not ret:
                    break

                current_time = frame_idx / fps

                if current_time < self.INIT_PERIOD:
                    frame = self.draw_init_overlay(frame, current_time)
                else:
                    pred = frame_predictions.get(frame_idx)
                    if pred is not None:
                        last_pred = pred

                    if last_pred is not None:
                        # Get timeline history up to current time
                        timeline_so_far = [(c, t) for c, t in full_timeline if t <= current_time]

                        frame = self.draw_overlay(
                            frame,
                            predicted_class=last_pred['predicted_name'],
                            confidence=last_pred['confidence'],
                            current_time=current_time,
                            num_votes=last_pred.get('num_votes', 1),
                            timeline_history=timeline_so_far,
                            total_duration=total_duration
                        )

                proc.stdin.write(frame.tobytes())
                frame_idx += 1

                if not quiet and frame_idx % 100 == 0:
                    pct = 100 * frame_idx / total_frames
                    print(f"  Processing: {frame_idx}/{total_frames} ({pct:.0f}%)", end='\r')

        except BrokenPipeError:
            pass
        finally:
            cap.release()
            proc.stdin.close()
            proc.wait()

        if not quiet:
            print(f"\nSaved: {output_path}")

        return output_path

    def _create_with_opencv(
        self,
        cap: cv2.VideoCapture,
        frame_predictions: Dict[int, Dict],
        output_path: str,
        fps: float,
        width: int,
        height: int,
        total_frames: int,
        quiet: bool
    ) -> str:
        """Fallback using OpenCV."""
        fourcc = cv2.VideoWriter_fourcc(*'mp4v')
        out = cv2.VideoWriter(output_path, fourcc, fps, (width, height))

        # Build full timeline history upfront
        full_timeline = self._build_timeline_history(frame_predictions, fps, total_frames)
        total_duration = total_frames / fps

        frame_idx = 0
        last_pred = None

        while True:
            ret, frame = cap.read()
            if not ret:
                break

            current_time = frame_idx / fps

            if current_time < self.INIT_PERIOD:
                frame = self.draw_init_overlay(frame, current_time)
            else:
                pred = frame_predictions.get(frame_idx)
                if pred is not None:
                    last_pred = pred

                if last_pred is not None:
                    # Get timeline history up to current time
                    timeline_so_far = [(c, t) for c, t in full_timeline if t <= current_time]

                    frame = self.draw_overlay(
                        frame,
                        predicted_class=last_pred['predicted_name'],
                        confidence=last_pred['confidence'],
                        current_time=current_time,
                        num_votes=last_pred.get('num_votes', 1),
                        timeline_history=timeline_so_far,
                        total_duration=total_duration
                    )

            out.write(frame)
            frame_idx += 1

        cap.release()
        out.release()

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

        return output_path


def create_multiclass_overlay_video(
    video_path: str,
    predictions: List[Dict[str, Any]],
    output_path: Optional[str] = None,
    quiet: bool = False,
    window_size: int = 16,
    model_fps: int = 8,
    stride: int = 4
) -> str:
    """
    Create a video with voting-smoothed classification overlays.

    Each frame shows the majority-voted prediction from all overlapping
    prediction windows, providing temporal smoothing.

    Args:
        video_path: Path to input video
        predictions: List of prediction dicts with 'timestamp', 'predicted_name',
                    'confidence', and 'all_probs'
        output_path: Output path (auto-generated if None)
        quiet: Suppress progress output
        window_size: Frames per prediction window
        model_fps: Model's processing FPS
        stride: Stride between windows

    Returns:
        Path to output video
    """
    processor = VotingSmoothedOverlay(
        window_size=window_size,
        model_fps=model_fps,
        stride=stride
    )
    return processor.create_overlay_video(video_path, predictions, output_path, quiet)
