"""Raw-state G transformations and empirical policy symmetry diagnostics."""
from __future__ import annotations

from typing import Any, Mapping, Optional, Sequence

import numpy as np

from .data_contracts import BatchDataset, build_batch_dataset
from .g_operator import apply_G_alpha, apply_G_raw12, apply_G_target_force


def build_mirrored_batch_dataset(
    sensors: np.ndarray,
    forces: np.ndarray,
    actions: np.ndarray,
    feat_keys: Sequence[str],
    *,
    target_forces: Optional[np.ndarray] = None,
    **kwargs: Any,
) -> BatchDataset:
    """Apply G to raw physical state/actions, then rebuild derived features."""
    raw = np.column_stack((np.asarray(sensors), np.asarray(forces)))
    mirrored = apply_G_raw12(raw)
    mirrored_target = None if target_forces is None else apply_G_target_force(target_forces)
    return build_batch_dataset(
        mirrored[:, :6], mirrored[:, 6:], apply_G_alpha(actions), feat_keys,
        target_forces=mirrored_target, **kwargs,
    )


def empirical_equivariance_defect(
    actions: np.ndarray,
    mirrored_policy_actions: np.ndarray,
) -> Mapping[str, Any]:
    """Compare policy actions at Gx against the exact expected G(action(x))."""
    expected = apply_G_alpha(actions)
    observed = np.asarray(mirrored_policy_actions, dtype=np.float64)
    if observed.shape != expected.shape:
        raise ValueError(f"mirrored actions have shape {observed.shape}; expected {expected.shape}")
    residual = observed - expected
    scale = np.ptp(np.asarray(actions, dtype=np.float64), axis=0)
    rmse = np.sqrt(np.mean(residual**2, axis=0))
    return {
        "channels": ["front", "upper", "lower"],
        "rmse": rmse.tolist(),
        "range_normalized_rmse": (rmse / np.where(scale > 1e-12, scale, 1.0)).tolist(),
        "max_abs_error": np.max(np.abs(residual), axis=0).tolist(),
        "n_samples": int(len(observed)),
    }
