#!/usr/bin/env python3
"""Symmetry augmentation wrapper for Karman Cloak (2000x600).

Applies the G-operator (up/down mirror) with 50% probability to encourage
the DRL policy to learn a symmetric control law.

Obs layout (env output, physical norm): [forces(6), sensors(6)]
  forces = [f_front_fx, f_front_fy, f_top_fx, f_top_fy, f_bot_fx, f_bot_fy]
  sensors = [s0_ux, s0_uy, s1_ux, s1_uy, s2_ux, s2_uy]
  (s0=top/+y, s1=center, s2=bot/-y; top=+y, bot=-y)

Action layout: [a_front, a_top, a_bot]

G-operator (y -> -y mirror):
  obs forces:  front_fx stays, front_fy negated
               top <-> bot swap, fy negated
  obs sensors: s0(top) <-> s2(bot) swap, uy negated
               s1(center) ux stays, uy negated
  action:      [aF, aT, aB] -> [-aF, -aB, -aT]

Wrapper placement: between KarmanCloakEnv and DummyVecEnv (inside VecNormalize).
The wrapper intercepts step(): with 50% prob, applies G to obs seen by policy,
then applies G_inv to action before passing to env. Reward and done are
unchanged (G is a symmetry of the physics).

Note: This wrapper must be applied BEFORE VecNormalize so that VecNormalize
sees the (possibly G-transformed) obs. Since G is an exact symmetry, the
running mean/std statistics remain valid (G just swaps/sign-flips dimensions,
preserving the overall distribution shape).

Usage in train script:
    env = KarmanCloakEnv(...)
    env = SymmetryAugmentWrapper(env, prob=0.5)
    vec_env = DummyVecEnv([lambda: env])
    vec_env = VecNormalize(vec_env, ...)
"""
from __future__ import annotations

from typing import Tuple

import numpy as np
import gymnasium as gym
from gymnasium import spaces


def apply_G_obs(obs: np.ndarray) -> np.ndarray:
    """Apply G-operator to obs [forces(6), sensors(6)].

    Returns G(obs) with same shape. G is its own inverse (G^2 = identity).
    """
    obs = np.asarray(obs)
    if obs.ndim != 1 or obs.shape[0] not in (12, 14):
        raise ValueError(
            f"Symmetry observation must be 12D Karman or 14D Illusion; got {obs.shape}."
        )
    out = np.empty_like(obs)
    # Forces: [f_front_fx, f_front_fy, f_top_fx, f_top_fy, f_bot_fx, f_bot_fy]
    out[0] = obs[0]      # front_fx stays
    out[1] = -obs[1]     # front_fy negated
    out[2] = obs[4]      # top_fx <- bot_fx
    out[3] = -obs[5]     # top_fy <- -bot_fy
    out[4] = obs[2]      # bot_fx <- top_fx
    out[5] = -obs[3]     # bot_fy <- -top_fy
    # Sensors: [s0_ux, s0_uy, s1_ux, s1_uy, s2_ux, s2_uy]
    out[6] = obs[10]     # s0_ux <- s2_ux
    out[7] = -obs[11]    # s0_uy <- -s2_uy
    out[8] = obs[8]      # s1_ux stays
    out[9] = -obs[9]     # s1_uy negated
    out[10] = obs[6]     # s2_ux <- s0_ux
    out[11] = -obs[7]    # s2_uy <- -s0_uy
    if obs.shape[0] == 14:
        out[12] = obs[12]    # target_cd stays
        out[13] = -obs[13]   # target_cl negated
    return out


def apply_G_action(action: np.ndarray) -> np.ndarray:
    """Apply G-operator to action [a_front, a_top, a_bot].

    [aF, aT, aB] -> [-aF, -aB, -aT]
    G is its own inverse (G^2 = identity).
    """
    out = np.empty_like(action)
    out[0] = -action[0]  # front reversed
    out[1] = -action[2]  # top <- -bot
    out[2] = -action[1]  # bot <- -top
    return out


class SymmetryAugmentWrapper(gym.Wrapper):
    """Augment training data with G-operator symmetry (per-rollout, 50% prob).

    Instead of randomly mirroring each step, this wrapper decides ONCE per
    rollout (every `rollout_len` steps) whether to mirror. If mirrored, ALL
    steps in that rollout use G-transform consistently. This avoids the noise
    of per-step random switching.

    reset() selects the frame for the first rollout; evaluation sets prob=0.
    """

    def __init__(self, env: gym.Env, prob: float = 0.5, seed: int = 42,
                 rollout_len: int = 2048):
        super().__init__(env)
        obs_shape = getattr(env.observation_space, "shape", None)
        if obs_shape not in ((12,), (14,)):
            raise ValueError(
                f"Symmetry observation must be 12D Karman or 14D Illusion; got {obs_shape}."
            )
        if not 0.0 <= prob <= 1.0:
            raise ValueError(f"prob must be in [0, 1], got {prob}")
        if rollout_len <= 0:
            raise ValueError(f"rollout_len must be positive, got {rollout_len}")
        self.prob = prob
        self._rollout_len = rollout_len
        self._rng = np.random.default_rng(seed)
        self._mirrored = False
        self._step_count = 0

    def reset(self, *, seed=None, options=None) -> Tuple[np.ndarray, dict]:
        obs, info = self.env.reset(seed=seed, options=options)
        # Select the frame before returning the first observation and retain it
        # for this rollout. Evaluation sets prob=0 explicitly.
        self._mirrored = self._rng.random() < self.prob
        self._step_count = 0
        if self._mirrored:
            obs = apply_G_obs(obs)
        return obs, info

    def step(self, action: np.ndarray) -> Tuple[np.ndarray, float, bool, bool, dict]:
        # If currently mirrored, policy's action is in mirrored frame
        if self._mirrored:
            real_action = apply_G_action(action)
        else:
            real_action = action

        obs, reward, terminated, truncated, info = self.env.step(real_action)

        # Decide mirror state for next rollout boundary
        self._step_count += 1
        if self._step_count % self._rollout_len == 0:
            # Start of new rollout — flip coin
            self._mirrored = self._rng.random() < self.prob

        # Apply mirror to obs if in mirrored mode
        if self._mirrored:
            obs = apply_G_obs(obs)

        return obs, reward, terminated, truncated, info


# ---------------------------------------------------------------------------
# Verification
# ---------------------------------------------------------------------------
def _verify_g_involution():
    """G(G(x)) should equal x for both obs and action."""
    rng = np.random.default_rng(0)
    # Test obs
    for _ in range(100):
        obs = rng.standard_normal(12)
        assert np.allclose(apply_G_obs(apply_G_obs(obs)), obs), "G_obs not involution"
    # Test action
    for _ in range(100):
        act = rng.standard_normal(3)
        assert np.allclose(apply_G_action(apply_G_action(act)), act), "G_action not involution"
    print("  G is involution (G^2 = identity): PASS")


def _verify_symmetry_semantics():
    """Verify G transforms match the physical mirror semantics."""
    # Construct a known obs: front has +fx, +fy; top has +fx, +fy; bot has +fx, +fy
    # sensors: s0(top) ux=1, uy=2; s1 ux=3, uy=4; s2(bot) ux=5, uy=6
    obs = np.array([1, 2, 3, 4, 5, 6,   # forces: front(1,2), top(3,4), bot(5,6)
                    10, 20, 30, 40, 50, 60], dtype=np.float64)  # sensors
    g = apply_G_obs(obs)
    # Expected:
    # front_fx=1 (stays), front_fy=-2 (negated)
    # top_fx <- bot_fx = 5, top_fy <- -bot_fy = -6
    # bot_fx <- top_fx = 3, bot_fy <- -top_fy = -4
    # s0_ux <- s2_ux = 50, s0_uy <- -s2_uy = -60
    # s1_ux = 30 (stays), s1_uy = -40 (negated)
    # s2_ux <- s0_ux = 10, s2_uy <- -s0_uy = -20
    expected = np.array([1, -2, 5, -6, 3, -4,
                         50, -60, 30, -40, 10, -20], dtype=np.float64)
    assert np.allclose(g, expected), f"G_obs mismatch:\n  got={g}\n  exp={expected}"
    print("  G_obs semantics: PASS")

    # Action: [aF, aT, aB] -> [-aF, -aB, -aT]
    act = np.array([1.0, 2.0, 3.0])
    g_act = apply_G_action(act)
    expected_act = np.array([-1.0, -3.0, -2.0])
    assert np.allclose(g_act, expected_act), f"G_action mismatch:\n  got={g_act}\n  exp={expected_act}"
    print("  G_action semantics: PASS")


if __name__ == "__main__":
    print("=== SymmetryAugmentWrapper verification ===")
    _verify_g_involution()
    _verify_symmetry_semantics()
    print("\nAll checks passed. Wrapper is ready for use.")
    print("\nUsage in train script:")
    print("    env = KarmanCloakEnv(...)")
    print("    env = SymmetryAugmentWrapper(env, prob=0.5)")
    print("    vec_env = DummyVecEnv([lambda: env])")
    print("    vec_env = VecNormalize(vec_env, ...)")
