"""Collator emitting the AC-3 frozen batch tensor contract.

Consumes per-triple samples in the public `csi_path_*` schema and batches
them into the AC-3 contract:
    scene_idx         [Q]      int64
    tx_scene_idx      [Q]      int64
    rx_xyz_norm       [Q, 3]   float32
    tx_xyz_norm       [Q, 3]   float32
    scene_extent      [Q, 3]   float32
    gt_num_paths      [Q]      int64
    gt_delay_ns       [Q, M_max]  float32
    gt_peak_db        [Q, M_max]  float32
    gt_path_mask      [Q, M_max]  uint8
    gt_truncated      [Q]      int64

Zero-path triples pass through (mask=0; num_paths=0).
"""

from __future__ import annotations

from typing import Any, Sequence

import numpy as np
import torch


def _to_float(x: Any) -> torch.Tensor:
    return torch.as_tensor(x, dtype=torch.float32)


def triple_collate(samples: Sequence[dict[str, Any]]) -> dict[str, Any]:
    if not samples:
        raise ValueError("triple_collate received empty sample list")

    Q = len(samples)
    # Public csi_path_* schema: each sample's csi_path_valid has shape [M_max]
    M_max = max(int(np.asarray(s["csi_path_valid"]).shape[0]) for s in samples)

    scene_seen: dict[Any, int] = {}
    tx_seen: dict[tuple[Any, Any], int] = {}
    scene_voxel_levels: list[Any] = []

    scene_idx = torch.zeros((Q,), dtype=torch.long)
    tx_scene_idx = torch.zeros((Q,), dtype=torch.long)
    tx_scene_voxel_idx = torch.zeros((Q,), dtype=torch.long)
    tx_xyz_norm = torch.zeros((Q, 3), dtype=torch.float32)
    rx_xyz_norm = torch.zeros((Q, 3), dtype=torch.float32)
    scene_extent = torch.zeros((Q, 3), dtype=torch.float32)
    gt_num_paths = torch.zeros((Q,), dtype=torch.long)
    gt_delay_ns = torch.zeros((Q, M_max), dtype=torch.float32)
    gt_peak_db = torch.zeros((Q, M_max), dtype=torch.float32)
    gt_path_mask = torch.zeros((Q, M_max), dtype=torch.uint8)
    gt_truncated = torch.zeros((Q,), dtype=torch.long)

    for q, s in enumerate(samples):
        sid = s["scene_id"]
        if sid not in scene_seen:
            scene_seen[sid] = len(scene_seen)
            scene_voxel_levels.append(s.get("voxel_level"))
        scene_idx[q] = scene_seen[sid]
        tx_scene_voxel_idx[q] = scene_idx[q]

        key = (sid, s["tx_id"])
        if key not in tx_seen:
            tx_seen[key] = len(tx_seen)
        tx_scene_idx[q] = tx_seen[key]

        tx_xyz_norm[q] = _to_float(s["tx_xyz_norm"])
        rx_xyz_norm[q] = _to_float(s["rx_xyz_norm"])
        scene_extent[q] = _to_float(s["scene_extent"])

        count = int(s["csi_path_count"])
        gt_num_paths[q] = count
        gt_truncated[q] = int(s.get("csi_path_truncation_count", 0))

        m = int(np.asarray(s["csi_path_valid"]).shape[0])
        gt_delay_ns[q, :m] = _to_float(s["csi_path_delay_ns"])
        gt_peak_db[q, :m] = _to_float(s["csi_path_peak_db"])
        gt_path_mask[q, :m] = torch.as_tensor(s["csi_path_valid"], dtype=torch.uint8)

    return {
        "scene_idx": scene_idx,
        "tx_scene_idx": tx_scene_idx,
        "tx_scene_voxel_idx": tx_scene_voxel_idx,
        "tx_xyz_norm": tx_xyz_norm,
        "rx_xyz_norm": rx_xyz_norm,
        "scene_extent": scene_extent,
        "gt_num_paths": gt_num_paths,
        "gt_delay_ns": gt_delay_ns,
        "gt_peak_db": gt_peak_db,
        "gt_path_mask": gt_path_mask,
        "gt_truncated": gt_truncated,
        "scene_voxel_levels": scene_voxel_levels,
        "unique_scene_count": len(scene_seen),
        "unique_scene_tx_count": len(tx_seen),
    }


__all__ = ["triple_collate"]