# legacy_test/core/io_helpers.py
"""I/O utilities for legacy test scripts.

Saves controlled/target/uncontrolled output and visualisations.
"""

from __future__ import annotations

import json
import os
from typing import Any, Dict, Optional, Tuple

import numpy as np


def save_signals(
    out_dir: str,
    sensors: np.ndarray,
    forces: np.ndarray,
    actions: np.ndarray,
    name: str = "controlled",
) -> str:
    """Save sensor/force/action arrays as compressed .npz.

    Args:
        out_dir: Output directory.
        sensors: (N, 6) raw sensor velocities.
        forces: (N, 6) raw force values.
        actions: (N, 3) normalised PPO actions in [-1, 1].
        name: Base filename without extension.

    Returns:
        Full path to the saved file.
    """
    os.makedirs(out_dir, exist_ok=True)
    path = os.path.join(out_dir, f"{name}.npz")
    np.savez_compressed(
        path,
        sensors=np.asarray(sensors, dtype=np.float32),
        forces=np.asarray(forces, dtype=np.float32),
        actions=np.asarray(actions, dtype=np.float32),
    )
    return path


def save_target(out_dir: str, target_states: np.ndarray) -> str:
    """Save target sensor signals.

    Args:
        out_dir: Output directory.
        target_states: (FIFO_LEN, 6) target sensor data.

    Returns:
        Full path to the saved file.
    """
    os.makedirs(out_dir, exist_ok=True)
    path = os.path.join(out_dir, "target.npz")
    np.savez_compressed(path, target_states=np.asarray(target_states, dtype=np.float32))
    return path


def save_norm(out_dir: str, norm: Dict[str, Any]) -> str:
    """Save normalisation constants as JSON.

    Args:
        out_dir: Output directory.
        norm: dict with force_norm_fact, sens_deviation, sens_norm_fact.

    Returns:
        Full path to the saved file.
    """
    os.makedirs(out_dir, exist_ok=True)
    path = os.path.join(out_dir, "norm.json")
    out = {
        "force_norm_fact": float(norm["force_norm_fact"]),
        "sens_deviation": [float(x) for x in norm["sens_deviation"]],
        "sens_norm_fact": [float(x) for x in norm["sens_norm_fact"]],
    }
    with open(path, "w") as f:
        json.dump(out, f, indent=2)
    return path


def load_reference_signals(ref_dir: str) -> Dict[str, np.ndarray]:
    """Load SR_analysis reference controlled.npz.

    Args:
        ref_dir: Path to the scene directory, e.g.
            ``src/SR_analysis/data/karman/karman_re100/``.

    Returns:
        dict with keys: sensors (N,6), forces (N,6), actions (N,3).
    """
    path = os.path.join(ref_dir, "controlled.npz")
    if not os.path.isfile(path):
        raise FileNotFoundError(f"Reference file not found: {path}")
    data = np.load(path)
    return {
        "sensors": np.asarray(data["sensors"], dtype=np.float32),
        "forces": np.asarray(data["forces"], dtype=np.float32),
        "actions": np.asarray(data["actions"], dtype=np.float32),
    }


def load_reference_target(ref_dir: str) -> np.ndarray:
    """Load SR_analysis reference target.npz.

    Returns:
        target_states: (FIFO_LEN, N) reference target sensor data.
    """
    path = os.path.join(ref_dir, "target.npz")
    if not os.path.isfile(path):
        raise FileNotFoundError(f"Reference file not found: {path}")
    return np.asarray(np.load(path)["target_states"], dtype=np.float32)


def load_reference_norm(ref_dir: str) -> Dict[str, Any]:
    """Load SR_analysis reference norm.json.

    Returns:
        dict with force_norm_fact, sens_deviation, sens_norm_fact.
    """
    path = os.path.join(ref_dir, "norm.json")
    if not os.path.isfile(path):
        raise FileNotFoundError(f"Reference file not found: {path}")
    with open(path) as f:
        d = json.load(f)
    return {
        "force_norm_fact": np.float32(d["force_norm_fact"]),
        "sens_deviation": np.array(d["sens_deviation"], dtype=np.float32),
        "sens_norm_fact": np.array(d["sens_norm_fact"], dtype=np.float32),
    }
