"""Sensor-unit schema helpers for calibration, policy observations, and DTW."""
from __future__ import annotations

from typing import Any, Mapping

import numpy as np

CALIBRATION_SCHEMA_NATIVE_V2 = "drl-pinball-calibration-v2"
SENSOR_UNITS_NATIVE = "celeris-area-time-average"
POLICY_SENSOR_UNITS_LEGACY = "legacy-policy-v1"
LEGACY_SENSOR_FACTOR = 78.0


def is_native_sensor_schema(calibration: Mapping[str, Any]) -> bool:
    """Return whether sensor-valued fields use CelerisLab native units."""
    return (
        calibration.get("schema_version") == CALIBRATION_SCHEMA_NATIVE_V2
        and calibration.get("sensor_units") == SENSOR_UNITS_NATIVE
    )


def dtw_sensor_factor(calibration: Mapping[str, Any]) -> float:
    """Scale a native solver sensor sample into the calibration's DTW units."""
    if is_native_sensor_schema(calibration):
        return 1.0
    return float(calibration.get("SENSOR_CC", LEGACY_SENSOR_FACTOR))


def policy_sensor_scale(
    calibration: Mapping[str, Any], policy_sensor_units: str | None = None
) -> float:
    """Return the divisor that produces the policy's historical raw observation.

    Native-v2 policies default to native sensor units. ``legacy-policy-v1`` is an
    explicit compatibility mode for an existing model/VecNormalize pair.
    Schema-less calibrations already store the legacy policy divisor and remain
    unchanged.
    """
    scale = float(calibration["SENS_SCALE"])
    if is_native_sensor_schema(calibration) and policy_sensor_units == POLICY_SENSOR_UNITS_LEGACY:
        return scale * LEGACY_SENSOR_FACTOR
    return scale


def convert_legacy_calibration(calibration: Mapping[str, Any]) -> dict[str, Any]:
    """Convert sensor-valued legacy calibration fields to native-v2 units."""
    if is_native_sensor_schema(calibration):
        return dict(calibration)
    factor = float(calibration.get("SENSOR_CC", LEGACY_SENSOR_FACTOR))
    converted = dict(calibration)
    converted["schema_version"] = CALIBRATION_SCHEMA_NATIVE_V2
    converted["sensor_units"] = SENSOR_UNITS_NATIVE
    converted["SENS_SCALE"] = float(calibration["SENS_SCALE"]) / factor
    converted["dtw_norm_scale"] = float(calibration["dtw_norm_scale"]) / factor
    converted.pop("SENSOR_CC", None)
    return converted


def convert_legacy_target(target: np.ndarray, factor: float = LEGACY_SENSOR_FACTOR) -> np.ndarray:
    """Convert a sensor target array from legacy-equivalent to native units."""
    converted = np.asarray(target).copy()
    if converted.ndim < 2 or converted.shape[-1] <= 6:
        return converted / float(factor)
    converted[..., :6] /= float(factor)
    return converted


def convert_legacy_harmonics(document: Any, factor: float = LEGACY_SENSOR_FACTOR) -> Any:
    """Convert the first six sensor channels in an illusion harmonics document."""
    wrapped = isinstance(document, dict)
    harmonics = document["harmonics"] if wrapped else document
    converted = []
    for index, source in enumerate(harmonics):
        item = dict(source)
        if index < 6:
            item["dc"] = float(item["dc"]) / factor
            item["amps"] = [float(value) / factor for value in item["amps"]]
        converted.append(item)
    if not wrapped:
        return converted
    result = dict(document)
    result["harmonics"] = converted
    return result
