"""Observation normalization — exact replication of legacy norm procedures.

Legacy norm computation (from FIFO data):
  - force_norm_fact = factor * max(|forces|)
  - sens_deviation[i] = mean(sensor_i)
  - sens_norm_fact[i] = factor * max(|sensor_i - sens_deviation[i]|)

Normalised observation:
  - forces = raw_forces / force_norm_fact
  - sens = (raw_sens - sens_deviation) / sens_norm_fact
  - observation = clip(hstack([forces, sens]), -1, 1)
"""
from __future__ import annotations

from typing import Dict, Optional, Tuple

import numpy as np


def compute_norm(
    fifo_states: np.ndarray,
    force_norm_factor: float = 6.0,
    sens_norm_factor: float = 5.0,
    force_slice: Tuple[int, int] = (6, 12),
    sens_slice: Tuple[int, int] = (0, 6),
) -> Dict[str, np.ndarray]:
    """Compute norm values from a FIFO of observations.

    Exact replica of legacy env norm computation.

    Args:
        fifo_states: shape ``(FIFO_LEN, N_obs)`` array of raw observations.
        force_norm_factor: Multiplier for force norm (6 for cloak, 100 for erase, etc.)
        sens_norm_factor: Multiplier for sensor norm (5 for cloak, 10 for erase, etc.)
        force_slice: Slice indices for forces within the obs array.
        sens_slice: Slice indices for sensors within the obs array.

    Returns:
        dict with keys:
            force_norm_fact: scalar float32
            sens_deviation: (6,) float32
            sens_norm_fact: (6,) float32
    """
    arr = np.asarray(fifo_states, dtype=np.float64)

    s_start, s_end = sens_slice
    f_start, f_end = force_slice

    forces = arr[:, f_start:f_end]
    sensors = arr[:, s_start:s_end]

    n_sens = sensors.shape[1]
    force_norm_fact = np.float32(force_norm_factor * np.max(np.abs(forces)))

    sens_deviation = np.mean(sensors, axis=0).astype(np.float32)
    sens_norm_fact = np.zeros(n_sens, dtype=np.float32)
    for i in range(n_sens):
        deviation = np.max(np.abs(sensors[:, i] - sens_deviation[i]))
        sens_norm_fact[i] = np.float32(sens_norm_factor * deviation)

    return {
        "force_norm_fact": force_norm_fact,
        "sens_deviation": sens_deviation,
        "sens_norm_fact": sens_norm_fact,
    }


def normalize_observation(
    raw_obs_slice: np.ndarray,
    norm: Dict[str, np.ndarray],
    force_slice: Tuple[int, int] = (6, 12),
    sens_slice: Tuple[int, int] = (0, 6),
) -> np.ndarray:
    """Normalize a raw observation slice using pre-computed norm values.

    Args:
        raw_obs_slice: shape ``(N_obs,)`` raw observation values.
        norm: dict with keys 'force_norm_fact', 'sens_deviation', 'sens_norm_fact'.
        force_slice: Slice for forces within the obs array.
        sens_slice: Slice for sensors within the obs array.

    Returns:
        shape ``(S_DIM,)`` normalized observation, clipped to [-1, 1].
    """
    obs = np.asarray(raw_obs_slice, dtype=np.float32)
    s_start, s_end = sens_slice
    f_start, f_end = force_slice

    forces = obs[f_start:f_end] / norm["force_norm_fact"]
    sens = (obs[s_start:s_end] - norm["sens_deviation"]) / norm["sens_norm_fact"]

    return np.clip(np.hstack([forces, sens]), -1.0, 1.0).astype(np.float32)


def load_norm(path: str) -> Dict[str, np.ndarray]:
    """Load norm values from a .npz file."""
    data = np.load(path)
    return {
        "force_norm_fact": data["force_norm_fact"],
        "sens_deviation": data["sens_deviation"],
        "sens_norm_fact": data["sens_norm_fact"],
    }


def save_norm(path: str, norm: Dict[str, np.ndarray]) -> None:
    """Save norm values to a .npz file."""
    np.savez_compressed(
        path,
        force_norm_fact=norm["force_norm_fact"],
        sens_deviation=norm["sens_deviation"],
        sens_norm_fact=norm["sens_norm_fact"],
    )
