"""Measure Compact-M K209 sweep acquisition throughput through VISA HiSLIP."""

from __future__ import annotations

from dataclasses import dataclass
import statistics
import time

import numpy as np

from python_app.hardware_full.compact_m_k209_service import CompactMK209Service
from python_app.models.run_config_model import RadarSweepModel


RESOURCE = "TCPIP0::127.0.0.1::hislip0,4880::INSTR"
VISA_LIBRARY = "@ivi"

START_HZ = 100_000_000.0
STOP_HZ = 6_000_000_000.0
POINTS = 1501
IFBW_HZ = 10_000.0
POWER_DBM = -20.0

TIMEOUT_MS = 20_000
WARMUP_SWEEPS = 3
TIMED_SWEEPS = 20
PRESET_ON_OPEN = False

# False measures the device/transport hot path: one synchronized TRIG:SING/*OPC?
# message plus S11/S21 REAL32 reads.
# True also includes public SweepResult construction and complex array conversion.
INCLUDE_RESULT_CONVERSION = False


@dataclass(frozen=True, slots=True)
class BenchmarkResult:
    """Timing summary for repeated sweep acquisition."""

    durations_s: list[float]
    points: int

    @property
    def total_s(self) -> float:
        """Return total timed acquisition duration."""
        return sum(self.durations_s)

    @property
    def sweeps_per_s(self) -> float:
        """Return completed sweeps per second."""
        return len(self.durations_s) / self.total_s

    @property
    def points_per_s(self) -> float:
        """Return measured sweep points per second."""
        return (len(self.durations_s) * self.points) / self.total_s

    @property
    def binary_payload_mb_per_s(self) -> float:
        """Return S11+S21 binary payload throughput, excluding SCPI headers."""
        payload_bytes = len(self.durations_s) * self.points * 2 * 2 * np.dtype(np.float32).itemsize
        return payload_bytes / self.total_s / 1_000_000.0


def _validate_config() -> None:
    if POINTS < 2:
        raise ValueError("POINTS must be >= 2")
    if WARMUP_SWEEPS < 0:
        raise ValueError("WARMUP_SWEEPS must be >= 0")
    if TIMED_SWEEPS <= 0:
        raise ValueError("TIMED_SWEEPS must be > 0")
    if IFBW_HZ <= 0.0:
        raise ValueError("IFBW_HZ must be > 0")
    if STOP_HZ < START_HZ:
        raise ValueError("STOP_HZ must be >= START_HZ")
    if TIMEOUT_MS <= 0:
        raise ValueError("TIMEOUT_MS must be > 0")


def _validate_interleaved(raw, expected_points: int) -> None:
    if raw.frequency_hz.shape != (expected_points,):
        raise RuntimeError(f"Unexpected frequency shape: {raw.frequency_hz.shape}")
    if raw.s11_values.shape != (expected_points * 2,):
        raise RuntimeError(f"Unexpected S11 shape: {raw.s11_values.shape}")
    if raw.s21_values.shape != (expected_points * 2,):
        raise RuntimeError(f"Unexpected S21 shape: {raw.s21_values.shape}")
    if not np.all(np.isfinite(raw.frequency_hz)):
        raise RuntimeError("Frequency axis contains non-finite values")
    if np.any(np.diff(raw.frequency_hz) < 0.0):
        raise RuntimeError("Frequency axis is not monotonic")
    if not np.all(np.isfinite(raw.s11_values)):
        raise RuntimeError("S11 contains non-finite values")
    if not np.all(np.isfinite(raw.s21_values)):
        raise RuntimeError("S21 contains non-finite values")


def _validate_result(result, expected_points: int) -> None:
    if result.x.shape != (expected_points,):
        raise RuntimeError(f"Unexpected frequency shape: {result.x.shape}")
    for name in ("s11", "s21"):
        values = result.trace(name)
        if values.shape != (expected_points,):
            raise RuntimeError(f"Unexpected {name.upper()} shape: {values.shape}")
        if not np.all(np.isfinite(values.real)) or not np.all(np.isfinite(values.imag)):
            raise RuntimeError(f"{name.upper()} contains non-finite values")


def _percentile(values: list[float], percentile: float) -> float:
    sorted_values = sorted(values)
    index = round((len(sorted_values) - 1) * percentile)
    return sorted_values[index]


def _run_benchmark(service: CompactMK209Service, *, points: int, warmup: int, sweeps: int, convert: bool) -> BenchmarkResult:
    acquire = service.acquire if convert else service.acquire_interleaved

    first = acquire()
    if convert:
        _validate_result(first, points)
    else:
        _validate_interleaved(first, points)

    for _ in range(warmup):
        acquire()

    durations_s: list[float] = []
    for _ in range(sweeps):
        start_ns = time.perf_counter_ns()
        last = acquire()
        end_ns = time.perf_counter_ns()
        durations_s.append((end_ns - start_ns) / 1_000_000_000.0)

    if convert:
        _validate_result(last, points)
    else:
        _validate_interleaved(last, points)
    return BenchmarkResult(durations_s=durations_s, points=points)


def _print_limits(limits: dict[str, float | int]) -> None:
    print(
        "K209 limits: "
        f"frequency={limits['min_frequency_hz']:.0f}..{limits['max_frequency_hz']:.0f} Hz, "
        f"IFBW={limits['min_ifbw_hz']:.0f}..{limits['max_ifbw_hz']:.0f} Hz, "
        f"power={limits['min_power_dbm']:.1f}..{limits['max_power_dbm']:.1f} dBm, "
        f"max_points={limits['max_points']}"
    )


def main() -> int:
    _validate_config()

    sweep = RadarSweepModel(
        start_hz=START_HZ,
        stop_hz=STOP_HZ,
        points=POINTS,
        if_bandwidth_hz=IFBW_HZ,
        power_dbm=POWER_DBM,
    )
    service = CompactMK209Service(
        resource=RESOURCE,
        timeout_ms=TIMEOUT_MS,
        preset_on_open=PRESET_ON_OPEN,
        visa_library=VISA_LIBRARY,
    )

    try:
        service.open()
        print(f"K209 IDN: {service.query_identity()}")
        _print_limits(service.read_device_limits())
        service.configure(sweep)
        print(
            "Benchmark settings: "
            f"start_hz={START_HZ:.3f}, stop_hz={STOP_HZ:.3f}, "
            f"points={POINTS}, ifbw_hz={IFBW_HZ:.3f}, power_dbm={POWER_DBM:.3f}, "
            f"warmup={WARMUP_SWEEPS}, sweeps={TIMED_SWEEPS}, "
            f"mode={'SweepResult' if INCLUDE_RESULT_CONVERSION else 'raw interleaved REAL32'}"
        )
        result = _run_benchmark(
            service,
            points=POINTS,
            warmup=WARMUP_SWEEPS,
            sweeps=TIMED_SWEEPS,
            convert=INCLUDE_RESULT_CONVERSION,
        )
        system_error = service.query_system_error()
        if not system_error.startswith("0,"):
            raise RuntimeError(f"K209 SCPI error after benchmark: {system_error}")
    finally:
        service.close()

    durations_ms = [value * 1_000.0 for value in result.durations_s]
    print("Benchmark result:")
    print(f"  total_s={result.total_s:.6f}")
    print(f"  sweep_mean_ms={statistics.fmean(durations_ms):.3f}")
    print(f"  sweep_median_ms={statistics.median(durations_ms):.3f}")
    print(f"  sweep_min_ms={min(durations_ms):.3f}")
    print(f"  sweep_max_ms={max(durations_ms):.3f}")
    print(f"  sweep_p95_ms={_percentile(durations_ms, 0.95):.3f}")
    print(f"  sweeps_per_s={result.sweeps_per_s:.3f}")
    print(f"  points_per_s={result.points_per_s:.1f}")
    print(f"  s11_s21_payload_mb_per_s={result.binary_payload_mb_per_s:.3f}")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())
