"""G-operator and equivariance tools.

Provides G-operator transformations, dimensionless conversion,
and equivariance diagnostics for PPO control laws.
"""
from __future__ import annotations

from typing import Any, Dict, Optional, Tuple

import numpy as np

from .feature_builder import compute_dimensionless as _compute_dimless


def apply_G_alpha(alpha: np.ndarray) -> np.ndarray:
    """Apply mirror G to action: [aF, aT, aB] -> [-aF, -aB, -aT]."""
    return np.array([-alpha[0], -alpha[2], -alpha[1]], dtype=alpha.dtype)


def apply_G_raw(obs_slice: np.ndarray,
                a_prev: np.ndarray,
                a_prev2: np.ndarray) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
    """Apply G to raw obs slice [sensor(6)+force(6)] and action arrays.

    Parameters
    ----------
    obs_slice : (12,) raw obs [s0_ux,uy, s1_ux,uy, s2_ux,uy, f0_fx,fy, f1_fx,fy, f2_fx,fy]
    a_prev : (3,) physical omega at t-1
    a_prev2 : (3,) physical omega at t-2

    Returns
    -------
    G_obs : (12,) transformed obs slice
    G_a_prev : (3,) transformed a_prev
    G_a_prev2 : (3,) transformed a_prev2
    """
    G_obs = np.zeros(12, dtype=np.float64)
    # sensors: swap top(0,1) <-> bottom(4,5), negate v
    G_obs[0] = obs_slice[4]
    G_obs[1] = -obs_slice[5]
    G_obs[2] = obs_slice[2]
    G_obs[3] = -obs_slice[3]
    G_obs[4] = obs_slice[0]
    G_obs[5] = -obs_slice[1]
    # forces: swap bottom(2,3) <-> top(4,5), negate fy
    G_obs[6] = obs_slice[6]
    G_obs[7] = -obs_slice[7]
    G_obs[8] = obs_slice[10]
    G_obs[9] = -obs_slice[11]
    G_obs[10] = obs_slice[8]
    G_obs[11] = -obs_slice[9]

    G_a_prev = np.array([-a_prev[0], -a_prev[2], -a_prev[1]], dtype=np.float64)
    G_a_prev2 = np.array([-a_prev2[0], -a_prev2[2], -a_prev2[1]], dtype=np.float64)
    return G_obs, G_a_prev, G_a_prev2


def check_equivariance(
    model: Any,
    obs_slice_series: np.ndarray,  # (T, 12) raw obs
    actions_phys: np.ndarray,      # (T, 3) physical omega
    norm: dict,
    action_scale: float = 8.0,
    action_bias: Tuple[float, float, float] = (0.0, -4.0, 4.0),
    u0: float = 0.01,
) -> Dict[str, float]:
    """Check G-equivariance of a PPO model over a time series.

    Returns dict with front/rear equivariance errors.
    """
    from .cfd_interface import build_observation, action_to_physical

    T = min(obs_slice_series.shape[0], actions_phys.shape[0])
    ef, eb, et = [], [], []

    for t in range(2, T):
        # Get current obs
        osl = obs_slice_series[t]
        a_prev = actions_phys[t - 1] if t > 0 else actions_phys[0]
        a_prev2 = actions_phys[t - 2] if t > 1 else actions_phys[0]

        # Predict action for current state
        obs = build_observation(osl, norm)
        act, _ = model.predict(obs, deterministic=True)
        act = act.astype(np.float32).flatten()
        alpha = action_to_physical(act.reshape(1, 3),
                                   scale=action_scale, bias=action_bias, u0=u0).flatten()

        # Apply G to state
        G_obs, _, _ = apply_G_raw(osl, a_prev, a_prev2)
        obs_G = build_observation(G_obs, norm)
        act_G, _ = model.predict(obs_G, deterministic=True)
        act_G = act_G.astype(np.float32).flatten()
        alpha_G = action_to_physical(act_G.reshape(1, 3),
                                     scale=action_scale, bias=action_bias, u0=u0).flatten()

        # Expected: G(alpha) = [-aF, -aB, -aT]
        expected = apply_G_alpha(alpha)

        ef.append(abs(float(alpha_G[0]) - float(expected[0])))
        eb.append(abs(float(alpha_G[1]) - float(expected[1])))
        et.append(abs(float(alpha_G[2]) - float(expected[2])))

    ef_arr = np.array(ef)
    eb_arr = np.array(eb)
    et_arr = np.array(et)
    alpha_range = float(np.max(np.abs(actions_phys[2:])))

    return {
        "front_mean_abs_error": float(np.mean(ef_arr)),
        "front_rel_error": float(np.mean(ef_arr) / (alpha_range + 1e-12)),
        "rear_bottom_rel_error": float(np.mean(eb_arr) / (alpha_range + 1e-12)),
        "rear_top_rel_error": float(np.mean(et_arr) / (alpha_range + 1e-12)),
        "alpha_range": alpha_range,
    }


def diagnose_one_re(model, ff, target_states, norm, config, n_steps=150) -> dict:
    """Run PPO inference and check equivariance.

    Parameters
    ----------
    model : loaded PPO model
    ff : FlowField instance (must be at saved checkpoint state)
    target_states : (FIFO_LEN, 6) target sensor signals
    norm : norm dict
    config : scene config dict with action_scale, action_bias, u0, etc.

    Returns
    -------
    dict with equivariance metrics.
    """
    from collections import deque
    from .cfd_interface import (build_observation, scale_action,
                                 action_to_physical, compute_similarity)

    action_scale = config.get("action_scale", 8.0)
    action_bias = config.get("action_bias", (0.0, -4.0, 4.0))
    u0 = config.get("u0", 0.01)
    sample_interval = config.get("sample_interval", 800)
    fifo_len = config.get("fifo_len", 150)
    n_obj_total = config.get("n_objects_total", 7)

    ff.restore_ddf()
    ff.apply_ddf()

    # Bias FIFO init
    fifo = deque(maxlen=fifo_len)
    bias_arr = scale_action(np.zeros(3, dtype=np.float32),
                            scale=action_scale, bias=action_bias,
                            u0=u0, n_total_bodies=n_obj_total)
    for _ in range(fifo_len):
        ff.run(sample_interval, bias_arr)
        fifo.append(ff.obs.copy()[2:14])

    # Inference
    obs_array = []
    action_array = []
    obs = np.zeros(12, dtype=np.float32)

    for _ in range(n_steps):
        act, _ = model.predict(obs, deterministic=True)
        act = act.astype(np.float32).flatten()
        action_array.append(act.copy())

        action_arr = scale_action(act, scale=action_scale, bias=action_bias,
                                  u0=u0, n_total_bodies=n_obj_total)
        ff.context.push()
        ff.run(sample_interval, action_arr)
        ff.context.pop()

        obs_slice = ff.obs.copy()[2:14]
        fifo.append(obs_slice)
        obs_array.append(obs_slice)
        obs = build_observation(obs_slice, norm)

    obs_series = np.array(obs_array, dtype=np.float64)
    actions_phys = action_to_physical(np.array(action_array),
                                      scale=action_scale, bias=action_bias, u0=u0)
    states_arr = np.array(list(fifo), dtype=np.float32)
    sim = compute_similarity(target_states, states_arr[:, 0:6],
                             config.get("conv_len", 30))

    # Equivariance check
    eq = check_equivariance(model, obs_series, actions_phys, norm,
                            action_scale, action_bias, u0)

    return {
        "similarity": sim,
        "equivariance": eq,
    }
