"""Object-calibration sets and the acquisition state that makes them applicable."""

from __future__ import annotations

import hashlib
import struct
from typing import Any

import numpy as np

from python_app.models.dataset_model import ComboKey, SweepCollection
from python_app.storage.npz.serialize import validate_trace_collection


def geometry_signature(config: dict[str, Any]) -> dict[str, list[dict[str, int | float]]]:
    """Normalize physical antenna coordinates independently of table row order."""
    return {
        key: sorted(({
            identity: int(entry[identity]),
            'x_m': float(entry['x_m']), 'y_m': float(entry.get('y_m', 0.)), 'z_m': float(entry.get('z_m', 0.)),
        } for entry in config.get('gpr', {}).get(key, [])), key=lambda record: record[identity])
        for key, identity in (("tx_geometry", "output_pos"), ("rx_geometry", "input_pos"))
    }


def through_fingerprint(collection: SweepCollection) -> str:
    """Hash physical pairs, float32 frequencies and complex64 S21, excluding timing."""
    validate_trace_collection(collection, require_finite=True)
    digest = hashlib.sha256()
    for trace in sorted(collection.traces, key=lambda item: (item.combo.input, item.combo.output)):
        frequency = np.asarray(trace.frequency_hz, dtype="<f4")
        samples = np.asarray(trace.s21, dtype="<c8")
        digest.update(struct.pack("<III", trace.combo.input, trace.combo.output, frequency.size))
        digest.update(frequency.tobytes())
        digest.update(samples.tobytes())
    return digest.hexdigest()


def capture_provenance(
    config: dict[str, Any], radar_key: str, through: SweepCollection,
) -> dict[str, Any]:
    """Freeze the hardware, geometry and through calibration used for a capture."""
    return {
        "radar_key": radar_key,
        "geometry": geometry_signature(config),
        "through_name": config["preprocess"]["s21"]["calibration"]["set_name"].strip(),
        "through_hash": through_fingerprint(through),
    }


def validate_object_set(
    collection: SweepCollection,
    metadata: dict[str, Any],
    config: dict[str, Any],
    radar_key: str,
    through: SweepCollection,
) -> None:
    """Reject corrections whose acquisition state or selected pair grids differ."""
    validate_trace_collection(collection, require_finite=True)
    captured = metadata.get("provenance")
    if not isinstance(captured, dict):
        raise ValueError("Object calibration has no capture provenance; capture it again")
    current = capture_provenance(config, radar_key, through)
    for field, label in (
        ("radar_key", "radar configuration"),
        ("geometry", "antenna geometry"),
        ("through_name", "through calibration name"),
        ("through_hash", "through calibration content"),
    ):
        if captured.get(field) != current[field]:
            raise ValueError(f"Object calibration {label} differs from the captured state")
    corrections = {trace.combo: trace for trace in collection.traces}
    through_traces = {trace.combo: trace for trace in through.traces}
    for combo in config["run"]["combos"]:
        key = ComboKey(input=combo["input"], output=combo["output"])
        if key not in corrections or key not in through_traces:
            raise ValueError(f"Object calibration and through must cover input={key.input}, output={key.output}")
        correction = corrections[key]
        if np.any(correction.s21 == 0):
            raise ValueError("Object calibration corrections must be nonzero")
        if not np.array_equal(
            np.asarray(correction.frequency_hz, dtype=np.float32),
            np.asarray(through_traces[key].frequency_hz, dtype=np.float32),
        ):
            raise ValueError(f"Object calibration frequency grid differs for input={key.input}, output={key.output}")