"""
Utility Functions for WiLoR Processing Pipeline

Helper functions for loading outputs, validation, and benchmarking.
"""

import time
from dataclasses import dataclass

import h5py
import numpy as np
import torch


@dataclass
class BenchmarkResult:
    """Results from benchmark run."""

    name: str
    total_time: float
    iterations: int
    items_processed: int

    @property
    def time_per_iter(self) -> float:
        return self.total_time / self.iterations

    @property
    def throughput(self) -> float:
        return self.items_processed / self.total_time


def load_wilor_output(h5_path: str) -> dict[str, np.ndarray | dict]:
    """
    Load WiLoR output from HDF5 file.

    Args:
        h5_path: Path to .h5 output file

    Returns:
        Dictionary with all stored data and metadata
    """
    with h5py.File(h5_path, "r") as hf:
        # Load metadata
        meta = dict(hf.attrs)

        # Load arrays
        data = {
            "meta": meta,
            "right": hf["right"][:].astype(bool),
            "img_idx": hf["img_idx"][:].astype(np.int32),
            "frame_idx": hf["frame_idx"][:].astype(np.int32),
            "cam": hf["cam"][:].astype(np.float32),
            "kpts_2d": hf["kpts_2d"][:].astype(np.float32),
            "kpts_3d": hf["kpts_3d"][:].astype(np.float32),
            "bboxes": hf["bboxes"][:].astype(np.float32),
        }

        # Load confidence if present (schema v1.1+)
        if "confidence" in hf:
            data["confidence"] = hf["confidence"][:].astype(np.float32)

        # Load MANO parameters
        data["mano"] = {
            "global_orient": hf["mano/global_orient"][:].astype(np.float32),
            "hand_pose": hf["mano/hand_pose"][:].astype(np.float32),
            "betas": hf["mano/betas"][:].astype(np.float32),
        }

    return data


def get_frame_detections(
    data: dict,
    frame_idx: int,
) -> dict[str, np.ndarray]:
    """
    Extract detections for a specific frame.

    Args:
        data: Loaded WiLoR output dictionary
        frame_idx: Frame index to query

    Returns:
        Dictionary with detection data for that frame only
    """
    # Use frame_idx lookup for O(1) access
    start_idx, count = data["frame_idx"][frame_idx]

    if count == 0:
        return {
            "right": np.array([], dtype=bool),
            "cam": np.array([]).reshape(0, 3),
            "kpts_2d": np.array([]).reshape(0, 21, 2),
            "kpts_3d": np.array([]).reshape(0, 21, 3),
            "bboxes": np.array([]).reshape(0, 4),
            "confidence": np.array([]),
        }

    end_idx = start_idx + count

    result = {
        "right": data["right"][start_idx:end_idx],
        "cam": data["cam"][start_idx:end_idx],
        "kpts_2d": data["kpts_2d"][start_idx:end_idx],
        "kpts_3d": data["kpts_3d"][start_idx:end_idx],
        "bboxes": data["bboxes"][start_idx:end_idx],
    }

    if "confidence" in data:
        result["confidence"] = data["confidence"][start_idx:end_idx]

    return result


def filter_by_confidence(
    data: dict,
    min_confidence: float = 0.5,
) -> dict[str, np.ndarray | dict]:
    """
    Filter detections by confidence threshold.

    Args:
        data: Loaded WiLoR output dictionary
        min_confidence: Minimum confidence to keep

    Returns:
        Filtered dictionary with only high-confidence detections
    """
    if "confidence" not in data:
        raise ValueError("No confidence data in output (schema < 1.1)")

    mask = data["confidence"] >= min_confidence

    filtered = {
        "meta": data["meta"],
        "right": data["right"][mask],
        "img_idx": data["img_idx"][mask],
        "cam": data["cam"][mask],
        "kpts_2d": data["kpts_2d"][mask],
        "kpts_3d": data["kpts_3d"][mask],
        "bboxes": data["bboxes"][mask],
        "confidence": data["confidence"][mask],
        "mano": {
            "global_orient": data["mano"]["global_orient"][mask],
            "hand_pose": data["mano"]["hand_pose"][mask],
            "betas": data["mano"]["betas"][mask],
        },
    }

    # Rebuild frame index
    num_frames = data["meta"]["num_frames"]
    filtered["frame_idx"] = _build_frame_index(filtered["img_idx"], num_frames)

    return filtered


def _build_frame_index(img_idx: np.ndarray, num_frames: int) -> np.ndarray:
    """Build frame index from img_idx array."""
    frame_idx = np.zeros((num_frames, 2), dtype=np.int32)

    for f in range(num_frames):
        detection_indices = np.where(img_idx == f)[0]
        if len(detection_indices) > 0:
            frame_idx[f, 0] = detection_indices.min()
            frame_idx[f, 1] = len(detection_indices)
        else:
            frame_idx[f, 0] = -1
            frame_idx[f, 1] = 0

    return frame_idx


def validate_output(h5_path: str) -> tuple[bool, list[str]]:
    """
    Validate WiLoR output file integrity.

    Args:
        h5_path: Path to .h5 output file

    Returns:
        Tuple of (is_valid, list_of_issues)
    """
    issues = []

    try:
        with h5py.File(h5_path, "r") as hf:
            # Check required datasets
            required = ["right", "img_idx", "frame_idx", "cam", "kpts_2d", "kpts_3d", "bboxes"]

            for name in required:
                if name not in hf:
                    issues.append(f"Missing dataset: {name}")

            # Check MANO group
            if "mano" not in hf:
                issues.append("Missing mano group")
            else:
                for name in ["global_orient", "hand_pose", "betas"]:
                    if f"mano/{name}" not in hf:
                        issues.append(f"Missing mano/{name}")

            # Check required attributes
            required_attrs = ["video_name", "num_frames", "schema_version"]
            for attr in required_attrs:
                if attr not in hf.attrs:
                    issues.append(f"Missing attribute: {attr}")

            # Validate array shapes
            if "right" in hf and "img_idx" in hf:
                n_det = hf["right"].shape[0]

                if hf["img_idx"].shape[0] != n_det:
                    issues.append(f"img_idx shape mismatch: {hf['img_idx'].shape[0]} vs {n_det}")

                if hf["kpts_2d"].shape[0] != n_det:
                    issues.append("kpts_2d shape mismatch")

                if hf["kpts_2d"].shape[1:] != (21, 2):
                    issues.append(f"kpts_2d wrong shape: {hf['kpts_2d'].shape}")

                if hf["kpts_3d"].shape[1:] != (21, 3):
                    issues.append(f"kpts_3d wrong shape: {hf['kpts_3d'].shape}")

            # Validate frame_idx
            if "frame_idx" in hf and "num_frames" in hf.attrs:
                num_frames = hf.attrs["num_frames"]
                if hf["frame_idx"].shape[0] != num_frames:
                    issues.append(
                        f"frame_idx length mismatch: {hf['frame_idx'].shape[0]} vs {num_frames}"
                    )

    except Exception as e:
        issues.append(f"Failed to open file: {e}")

    return len(issues) == 0, issues


def benchmark_vitdet_gpu(
    processor,
    batch_sizes: list[int] = [64, 128, 256, 512],
    num_detections: int = 100,
    img_size: tuple[int, int] = (1080, 1920),
    num_warmup: int = 3,
    num_runs: int = 10,
) -> list[BenchmarkResult]:
    """
    Benchmark ViTDet GPU processor at various batch sizes.

    Args:
        processor: ViTDetGPUProcessor instance
        batch_sizes: List of batch sizes to test
        num_detections: Number of detections per batch
        img_size: (H, W) of test images
        num_warmup: Warmup iterations
        num_runs: Timed iterations

    Returns:
        List of BenchmarkResult for each batch size
    """
    device = processor.device
    results = []

    for batch_size in batch_sizes:
        # Create dummy data
        frames = torch.randint(
            0, 255, (batch_size, img_size[0], img_size[1], 3), dtype=torch.uint8, device=device
        )

        # Random boxes within image bounds
        boxes = torch.zeros(num_detections, 4, device=device)
        boxes[:, 0] = torch.randint(0, img_size[1] - 200, (num_detections,))
        boxes[:, 1] = torch.randint(0, img_size[0] - 200, (num_detections,))
        boxes[:, 2] = boxes[:, 0] + torch.randint(100, 200, (num_detections,))
        boxes[:, 3] = boxes[:, 1] + torch.randint(100, 200, (num_detections,))

        is_right = torch.randint(0, 2, (num_detections,), device=device).float()
        img_indices = torch.randint(0, batch_size, (num_detections,), device=device).long()
        confidences = torch.rand(num_detections, device=device)

        # Warmup
        for _ in range(num_warmup):
            _ = processor(frames, boxes, is_right, img_indices, confidences)
            torch.cuda.synchronize()

        # Timed runs
        torch.cuda.synchronize()
        start = time.perf_counter()

        for _ in range(num_runs):
            _ = processor(frames, boxes, is_right, img_indices, confidences)

        torch.cuda.synchronize()
        elapsed = time.perf_counter() - start

        results.append(
            BenchmarkResult(
                name=f"batch_{batch_size}",
                total_time=elapsed,
                iterations=num_runs,
                items_processed=num_detections * num_runs,
            )
        )

    return results


def print_benchmark_results(results: list[BenchmarkResult]):
    """Pretty print benchmark results."""
    print("\n" + "=" * 60)
    print("BENCHMARK RESULTS")
    print("=" * 60)
    print(f"{'Name':<20} {'Time/iter (ms)':<18} {'Throughput (det/s)':<20}")
    print("-" * 60)

    for r in results:
        print(f"{r.name:<20} {r.time_per_iter*1000:<18.2f} {r.throughput:<20.0f}")


def estimate_memory_usage(
    num_frames: int,
    avg_detections_per_frame: float,
    img_size: tuple[int, int] = (1080, 1920),
    patch_size: int = 256,
) -> dict[str, float]:
    """
    Estimate GPU memory usage for processing.

    Args:
        num_frames: Number of frames to process together
        avg_detections_per_frame: Average hands per frame
        img_size: (H, W) of input frames
        patch_size: Size of extracted patches

    Returns:
        Dictionary with memory estimates in GB
    """
    num_det = int(num_frames * avg_detections_per_frame)

    # Frame buffer: N x H x W x 3 (uint8) + float32 copy
    frame_mem = num_frames * img_size[0] * img_size[1] * 3 * (1 + 4)

    # Patch buffer: M x 3 x patch x patch (float32)
    patch_mem = num_det * 3 * patch_size * patch_size * 4

    # Model activations (rough estimate)
    model_mem = 4e9  # ~4GB for WiLoR + YOLO

    # Output buffers
    output_mem = num_det * (
        21 * 3 * 4  # kpts_3d
        + 21 * 2 * 4  # kpts_2d
        + 3 * 4  # cam
        + 4 * 4  # bbox
        + 1  # right
        + 4  # confidence
    )

    total = frame_mem + patch_mem + model_mem + output_mem

    return {
        "frames_gb": frame_mem / 1e9,
        "patches_gb": patch_mem / 1e9,
        "model_gb": model_mem / 1e9,
        "output_gb": output_mem / 1e9,
        "total_gb": total / 1e9,
    }
