"""Classify ONNX outputs into HDF5 layout rows and stack batches."""

from __future__ import annotations

from typing import Any, Dict, List, Optional, Sequence

import numpy as np
import onnxruntime as ort

from chess_tutor.lc0_onnx.session import resolve_lc0_onnx_heads
from chess_tutor.training.tensor_names import tensor_dataset_name


def classify_outputs(
    sess: ort.InferenceSession,
    onnx_out: Dict[str, np.ndarray],
    probe_names: Sequence[str],
) -> Dict[str, Any]:
    """
    Build layout description and one-row numpy copies for each stored tensor.

    Probes listed in ``probe_names`` are only included if present in ``onnx_out``.
    """
    policy_n, value_n, wdl_n = resolve_lc0_onnx_heads(sess)
    layout: Dict[str, Any] = {"probes": []}

    if policy_n and policy_n in onnx_out:
        pol = np.asarray(onnx_out[policy_n], dtype=np.float32)
        layout["policy_key"] = tensor_dataset_name(policy_n)
        layout["policy_vec"] = pol.reshape(1, -1)

    if wdl_n and wdl_n in onnx_out:
        w = np.asarray(onnx_out[wdl_n], dtype=np.float32).reshape(1, -1)
        layout["wdl_key"] = tensor_dataset_name(wdl_n)
        layout["wdl_vec"] = w
    elif value_n and value_n in onnx_out:
        v = np.asarray(onnx_out[value_n], dtype=np.float32).reshape(1, -1)
        layout["value_key"] = tensor_dataset_name(value_n)
        layout["value_vec"] = v

    for name in probe_names:
        if name not in onnx_out:
            continue
        arr = np.asarray(onnx_out[name], dtype=np.float32)
        if arr.ndim == 0:
            arr = arr.reshape(1, 1)
        elif arr.shape[0] != 1:
            arr = np.reshape(arr, (1,) + arr.shape)
        dk = tensor_dataset_name(name)
        layout["probes"].append({"onnx_name": name, "dataset_key": dk, "sample": arr})

    if "policy_vec" not in layout:
        raise RuntimeError("Could not resolve policy output from ONNX session.")
    return layout


def merge_layout_batch(
    acc: Optional[Dict[str, Any]], one: Dict[str, Any]
) -> Dict[str, Any]:
    """Ensure batch of rows matches same tensor layout as first row."""
    if acc is None:
        return one
    if one["policy_vec"].shape != acc["policy_vec"].shape:
        raise RuntimeError(
            f"Policy shape mismatch: {one['policy_vec'].shape} vs {acc['policy_vec'].shape}"
        )
    for k in ("wdl_vec", "value_vec"):
        if k in acc and k in one:
            if one[k].shape != acc[k].shape:
                raise RuntimeError(f"{k} shape mismatch across rows")
    if len(one["probes"]) != len(acc["probes"]):
        raise RuntimeError("Probe count mismatch across rows")
    for a, b in zip(acc["probes"], one["probes"]):
        if a["dataset_key"] != b["dataset_key"] or a["onnx_name"] != b["onnx_name"]:
            raise RuntimeError("Probe ordering/name mismatch across rows")
        if a["sample"].shape != b["sample"].shape:
            raise RuntimeError(
                f"Probe {a['onnx_name']!r} shape mismatch: {b['sample'].shape} vs {a['sample'].shape}"
            )
    return acc


def stack_batch(layout_rows: List[Dict[str, Any]]) -> Dict[str, np.ndarray]:
    """Stack single-row layout dicts into batch arrays (B, ...)."""
    pol = np.concatenate([x["policy_vec"] for x in layout_rows], axis=0)
    out: Dict[str, np.ndarray] = {"policy": pol}
    if layout_rows and "wdl_vec" in layout_rows[0]:
        out["wdl"] = np.concatenate([x["wdl_vec"] for x in layout_rows], axis=0)
    if layout_rows and "value_vec" in layout_rows[0]:
        out["value_head"] = np.concatenate([x["value_vec"] for x in layout_rows], axis=0)
    probes0 = layout_rows[0]["probes"]
    for i, p in enumerate(probes0):
        key = p["dataset_key"]
        out[key] = np.concatenate([x["probes"][i]["sample"] for x in layout_rows], axis=0)
    return out