"""Postprocessing helpers to normalize prediction outputs."""

from __future__ import annotations

from collections.abc import Mapping, Sequence
from typing import TYPE_CHECKING, Any

import numpy as np
import torch
import torch.nn.functional as nn_functional


if TYPE_CHECKING:
    from .schemas import (
        ActionPrediction,
        ClassificationPrediction,
        DetectionPrediction,
        FramewiseActionPredictions,
    )


def classification_top1_prediction(
    logits: torch.Tensor,
    *,
    confidence_digits: int = 2,
) -> ClassificationPrediction:
    """Return top-1 class prediction with confidence from logits.

    Args:
        logits: Model logits tensor with shape ``[N, C]``.
        confidence_digits: Number of decimal places for confidence rounding.

    Returns:
        Classification prediction dictionary with category and confidence.
    """
    if logits.ndim != 2:
        raise ValueError(f"Expected logits shape [N, C], got {tuple(logits.shape)}")
    probabilities = nn_functional.softmax(logits, dim=1)
    predicted_class = torch.argmax(logits, dim=1).item()
    confidence = round(float(probabilities[0, predicted_class].item()), confidence_digits)
    return {"category": str(predicted_class), "confidence": confidence}


def classification_top1_from_raw_output(
    raw_output: Any,
    *,
    confidence_digits: int = 2,
) -> ClassificationPrediction:
    """Normalize common runtime outputs and return top-1 prediction.

    Args:
        raw_output: Runtime-specific raw model output object.
        confidence_digits: Number of decimal places for confidence rounding.

    Returns:
        Classification prediction dictionary with category and confidence.
    """
    logits = _coerce_logits(raw_output)
    return classification_top1_prediction(logits, confidence_digits=confidence_digits)


def _coerce_logits(raw_output: Any) -> torch.Tensor:
    """Convert runtime outputs to ``[1, C]`` tensor logits.

    Args:
        raw_output: Runtime-specific raw model output object.

    Returns:
        Float tensor logits with shape ``[N, C]``.
    """
    if isinstance(raw_output, Mapping):
        if not raw_output:
            raise ValueError("Output mapping is empty.")
        raw_output = next(iter(raw_output.values()))

    if isinstance(raw_output, (list, tuple)):
        if not raw_output:
            raise ValueError("Output sequence is empty.")
        raw_output = raw_output[0]

    if isinstance(raw_output, torch.Tensor):
        logits = raw_output.detach().to("cpu")
    elif isinstance(raw_output, np.ndarray):
        logits = torch.from_numpy(raw_output)
    else:
        logits = torch.as_tensor(raw_output)

    logits = logits.float()
    if logits.ndim == 1:
        logits = logits.unsqueeze(0)
    if logits.ndim > 2:
        logits = logits.reshape(logits.shape[0], -1)
    return logits


def format_detection_from_normalized_tensors(
    labels: torch.Tensor,
    boxes_xyxy_normalized: torch.Tensor,
    scores: torch.Tensor,
    *,
    image_width: int,
    image_height: int,
    confidence_threshold: float = 0.5,
) -> list[DetectionPrediction]:
    """Convert normalized detection tensors into API-ready records.

    Args:
        labels: Tensor of class ids for detections.
        boxes_xyxy_normalized: Normalized ``[N, 4]`` boxes in ``xyxy`` format.
        scores: Confidence tensor aligned with boxes.
        image_width: Source image width in pixels.
        image_height: Source image height in pixels.
        confidence_threshold: Minimum confidence required to keep a detection.

    Returns:
        List of detection prediction dictionaries.
    """
    if boxes_xyxy_normalized.ndim != 2 or boxes_xyxy_normalized.size(-1) != 4:
        raise ValueError("Expected boxes with shape [N, 4].")

    device = boxes_xyxy_normalized.device
    scale = torch.tensor(
        [image_width, image_height, image_width, image_height],
        dtype=boxes_xyxy_normalized.dtype,
        device=device,
    )
    boxes = boxes_xyxy_normalized * scale

    output: list[DetectionPrediction] = []
    for idx in range(boxes.size(0)):
        score = float(scores[idx].item())
        if score < confidence_threshold:
            continue

        x_min, y_min, x_max, y_max = boxes[idx]
        output.append(
            {
                "category": str(int(labels[idx].item())),
                "confidence": score,
                "bounding_box": {
                    "xmin": int(x_min.item()),
                    "ymin": int(y_min.item()),
                    "xmax": int(x_max.item()),
                    "ymax": int(y_max.item()),
                },
            }
        )
    return output


def _timestamp_to_seconds(timestamp: Any) -> float:
    """Normalize timestamp payload to float seconds.

    Args:
        timestamp: Timestamp value such as ``1.2`` or ``"1.2s"``.

    Returns:
        Timestamp converted to seconds as float.
    """
    if isinstance(timestamp, str) and timestamp.endswith("s"):
        return float(timestamp[:-1])
    return float(timestamp)


def framewise_action_results(
    inference_results: Sequence[Mapping[str, Any]],
    *,
    fps: float,
    confidence_digits: int = 2,
) -> FramewiseActionPredictions:
    """Group frame-level action predictions keyed by frame index string.

    Args:
        inference_results: Iterable of per-event action inference records.
        fps: Frames-per-second used to convert timestamps to frame indices.
        confidence_digits: Number of decimal places for confidence rounding.

    Returns:
        Mapping of frame index string to action prediction list.
    """
    frame_results: dict[int, list[ActionPrediction]] = {}
    fps_int = int(fps)

    for result in inference_results:
        timestamp = result.get("timestamp", 0.0)
        frame_number = int(_timestamp_to_seconds(timestamp) * fps)
        frame_results.setdefault(frame_number, []).append(
            {
                "category": str(result.get("label", "")),
                "confidence": round(float(result.get("confidence", 0.0)), confidence_digits),
                "timestamp": result.get("timestamp", 0.0),
                "fps": fps_int,
            }
        )

    return {str(frame): detections for frame, detections in frame_results.items() if detections}
