"""Vectorized 2D robot-sumo simulator (numpy only).

Two disc-shaped, differential-drive-style robots start inside a circular ring.
A robot loses as soon as its *center* leaves the ring. N independent matches
("envs") are simulated at once with array operations, which is what makes
evolution strategies on a laptop CPU feasible.

Units are SI-ish: metres, seconds, kilograms, newtons.

Physics model (per robot, per substep)
--------------------------------------
* Motor command ``a = (thrust, turn)`` in [-1, 1]^2 passes through a first-order
  lag with time constant ``tau`` (motors do not respond instantly).
* Longitudinal force  = F_MOTOR * thrust_state - DRAG * mass * v_long
  (the drag term models back-EMF / rolling losses and sets the top speed).
* Lateral force tries to cancel sideways velocity (wheels do not slide
  sideways) -- this is what lets a robot "dig in" against side pushes.
* The vector sum of longitudinal + lateral wheel forces is capped by the
  friction circle |F| <= mu * mass * g. On an icy floor (small mu) robots skid,
  accelerate slowly and get shoved around easily.
* Turn rate follows the lagged turn command; turning authority is also reduced
  on low-grip floors (a simplification of wheel slip during spin turns).
* Robots interact through a soft spring-damper contact (frictionless discs).
* A small uncapped air drag keeps everything numerically tame.

Integration is semi-implicit (symplectic) Euler with a fixed number of
substeps per control step for stability of the stiff contact spring.
"""

from __future__ import annotations

from dataclasses import dataclass

import numpy as np

# ---------------------------------------------------------------- constants
RING_RADIUS = 1.0  # m
ROBOT_RADIUS = 0.1  # m
G = 9.81  # m/s^2
F_MOTOR = 12.0  # N, peak motor force (more than a 1 kg robot's traction on a normal floor)
DRAG = 4.0  # 1/s, longitudinal drag per unit mass
AIR_DRAG = 0.1  # 1/s, uncapped velocity damping
OMEGA_MAX = 6.0  # rad/s, max turn rate
MU_TURN_REF = 0.5  # below this friction, turning authority scales down linearly
K_CONTACT = 3000.0  # N/m, contact spring stiffness
C_CONTACT = 15.0  # N s/m, contact damping
LATERAL_GAIN = 0.8  # fraction of sideways velocity the wheels try to cancel per substep

CONTROL_DT = 0.05  # s, policy acts at 20 Hz
SUBSTEPS = 5  # physics substeps per control step
EPISODE_SECONDS = 12.0
MAX_STEPS = int(round(EPISODE_SECONDS / CONTROL_DT))  # 240
MAX_DELAY = 4  # observation-delay buffer length (control steps)

VEL_SCALE = 2.0  # m/s, used to normalise velocities in observations
OBS_DIM = 16


@dataclass
class PhysParams:
    """Per-env physics parameters. Arrays have shape (n,) or (n, 2) (per robot)."""

    mu: np.ndarray  # (n,)   floor friction coefficient
    mass: np.ndarray  # (n, 2) robot masses
    tau: np.ndarray  # (n, 2) motor lag time constants (s)
    noise: np.ndarray  # (n, 2) observation noise std (in normalised obs units)
    delay: np.ndarray  # (n, 2) observation delay in control steps (int)

    @staticmethod
    def nominal(n: int) -> "PhysParams":
        """Default 'normal floor' parameters."""
        return PhysParams(
            mu=np.full(n, 0.8),
            mass=np.ones((n, 2)),
            tau=np.full((n, 2), 0.05),
            noise=np.zeros((n, 2)),
            delay=np.zeros((n, 2), dtype=int),
        )

    @staticmethod
    def randomized(n: int, rng: np.random.Generator) -> "PhysParams":
        """Domain randomization ranges used for training level L4."""
        return PhysParams(
            mu=np.exp(rng.uniform(np.log(0.15), np.log(1.0), n)),
            mass=rng.uniform(0.7, 2.2, (n, 2)),
            tau=np.exp(rng.uniform(np.log(0.02), np.log(0.35), (n, 2))),
            noise=rng.uniform(0.0, 0.2, (n, 2)) * (rng.random((n, 2)) < 0.6),
            delay=rng.integers(0, 3, (n, 2)),
        )

    def tile(self, reps: int) -> "PhysParams":
        """Repeat the whole batch ``reps`` times (used for common random numbers in ES)."""
        t = lambda a: np.tile(a, (reps,) + (1,) * (a.ndim - 1))  # noqa: E731
        return PhysParams(t(self.mu), t(self.mass), t(self.tau), t(self.noise), t(self.delay))

    def subset(self, idx) -> "PhysParams":
        return PhysParams(self.mu[idx], self.mass[idx], self.tau[idx], self.noise[idx], self.delay[idx])


# Named perturbations used by evaluate.py / watch.py. Robot 0 is "us", robot 1
# the opponent. Each entry modifies a nominal PhysParams in place.
def apply_perturbation(p: PhysParams, name: str) -> PhysParams:
    if name == "normal":
        pass
    elif name == "icy":
        p.mu[:] = 0.2
    elif name == "heavy_opp":
        p.mass[:, 1] = 2.0
    elif name == "laggy":
        p.tau[:, 0] = 0.3
        p.delay[:, 0] = 2
    elif name == "noisy":
        p.noise[:, 0] = 0.15
        p.delay[:, 0] = 1
    else:
        raise ValueError(f"unknown perturbation {name!r}")
    return p


PERTURBATIONS = ["normal", "icy", "heavy_opp", "laggy", "noisy"]


def random_starts(n: int, rng: np.random.Generator) -> dict:
    """Random, statistically symmetric start states for n matches.

    Robots sit on roughly opposite sides of the center at 0.25-0.6 R with random
    headings, so neither side is systematically favoured.
    """
    phi = rng.uniform(-np.pi, np.pi, n)
    r0 = rng.uniform(0.25, 0.6, n) * RING_RADIUS
    r1 = rng.uniform(0.25, 0.6, n) * RING_RADIUS
    phi1 = phi + np.pi + rng.uniform(-0.6, 0.6, n)
    pos = np.zeros((n, 2, 2))
    pos[:, 0] = np.stack([r0 * np.cos(phi), r0 * np.sin(phi)], -1)
    pos[:, 1] = np.stack([r1 * np.cos(phi1), r1 * np.sin(phi1)], -1)
    heading = rng.uniform(-np.pi, np.pi, (n, 2))
    return {"pos": pos, "heading": heading}


def tile_starts(starts: dict, reps: int) -> dict:
    return {k: np.tile(v, (reps,) + (1,) * (v.ndim - 1)) for k, v in starts.items()}


class SumoEnv:
    """N simultaneous sumo matches.

    State arrays (index 0/1 = robot):
        pos (n,2,2), vel (n,2,2), heading (n,2), omega (n,2),
        thrust_state (n,2), turn_state (n,2)
    """

    def __init__(self, n: int, params: PhysParams | None = None, seed: int | None = None):
        self.n = n
        self.params = params if params is not None else PhysParams.nominal(n)
        self.rng = np.random.default_rng(seed)
        self.reset()

    # ------------------------------------------------------------ reset
    def reset(self, starts: dict | None = None):
        """Reset all matches. Returns (obs0, obs1), each (n, OBS_DIM)."""
        n = self.n
        if starts is None:
            starts = random_starts(n, self.rng)
        self.pos = starts["pos"].astype(float).copy()
        self.heading = starts["heading"].astype(float).copy()
        self.vel = starts.get("vel", np.zeros((n, 2, 2))).astype(float).copy()
        self.omega = np.zeros((n, 2))
        self.thrust_state = np.zeros((n, 2))
        self.turn_state = np.zeros((n, 2))
        self.t = 0
        self.done = np.zeros(n, dtype=bool)
        self.outcome = np.zeros(n, dtype=int)  # +1 robot0 won, -1 robot1 won, 0 tie/ongoing
        self.end_step = np.full(n, MAX_STEPS)
        clean = self._observe_true()
        self._hist = np.repeat(clean[None], MAX_DELAY + 1, axis=0)  # (D+1, n, 2, OBS)
        return self._delayed_noisy_obs(clean)

    # ------------------------------------------------------------ observations
    def _observe_true(self) -> np.ndarray:
        """Egocentric observation for both robots, shape (n, 2, OBS_DIM).

        The same function is used for both robots with roles swapped, so the
        observation is symmetric: one policy can play either side. Everything
        is expressed in the robot's own body frame (x forward, y left), so the
        observation is also invariant to rotating the whole arena.
        """
        out = np.empty((self.n, 2, OBS_DIM))
        for i in (0, 1):
            j = 1 - i
            c, s = np.cos(self.heading[:, i]), np.sin(self.heading[:, i])

            def rot(v):  # world -> body frame of robot i
                return np.stack([c * v[:, 0] + s * v[:, 1], -s * v[:, 0] + c * v[:, 1]], -1)

            dh = self.heading[:, j] - self.heading[:, i]
            o = out[:, i]
            o[:, 0:2] = rot(self.pos[:, j] - self.pos[:, i]) / RING_RADIUS  # opponent position
            o[:, 2:4] = rot(self.vel[:, j] - self.vel[:, i]) / VEL_SCALE  # opponent rel. velocity
            o[:, 4] = np.cos(dh)  # opponent heading, relative
            o[:, 5] = np.sin(dh)
            o[:, 6:8] = rot(self.vel[:, i]) / VEL_SCALE  # own velocity (fwd, sideways)
            o[:, 8] = self.omega[:, i] / OMEGA_MAX  # own turn rate
            o[:, 9:11] = rot(-self.pos[:, i]) / RING_RADIUS  # vector to ring center
            o[:, 11] = np.linalg.norm(self.pos[:, i], axis=-1) / RING_RADIUS  # own edge proximity
            o[:, 12] = np.linalg.norm(self.pos[:, j], axis=-1) / RING_RADIUS  # opponent's
            o[:, 13] = self.thrust_state[:, i]  # own motor states (proprioception)
            o[:, 14] = self.turn_state[:, i]
            o[:, 15] = 1.0 - self.t / MAX_STEPS  # time left
        return out

    def _delayed_noisy_obs(self, clean: np.ndarray):
        """Apply per-robot observation delay and Gaussian sensor noise."""
        self._hist = np.roll(self._hist, 1, axis=0)
        self._hist[0] = clean
        idx = np.arange(self.n)
        obs = []
        for i in (0, 1):
            o = self._hist[self.params.delay[:, i], idx, i].copy()
            sd = self.params.noise[:, i]
            if np.any(sd > 0):
                noise = self.rng.standard_normal(o.shape) * sd[:, None]
                noise[:, 15] = 0.0  # the clock is not noisy
                o += noise
            obs.append(o)
        return obs[0], obs[1]

    # ------------------------------------------------------------ dynamics
    def step(self, a0: np.ndarray, a1: np.ndarray):
        """Advance one control step.

        Returns (obs0, obs1, done, outcome); ``done``/``outcome`` are sticky:
        once a match has ended its result never changes (the physics keeps
        running for that env but is ignored).
        """
        act = np.clip(np.stack([a0, a1], axis=1), -1.0, 1.0)  # (n,2,2)
        p = self.params
        m = p.mass  # (n,2)
        mu = p.mu[:, None]  # (n,1)
        dt = CONTROL_DT / SUBSTEPS
        lag = 1.0 - np.exp(-dt / p.tau)  # (n,2) per-substep motor-lag blend factor
        turn_authority = np.minimum(1.0, mu / MU_TURN_REF)
        grip = mu * m * G  # (n,2) friction-circle radius

        for _ in range(SUBSTEPS):
            # motors (first-order lag)
            self.thrust_state += lag * (act[..., 0] - self.thrust_state)
            self.turn_state += lag * (act[..., 1] - self.turn_state)
            self.omega = self.turn_state * OMEGA_MAX * turn_authority

            h = np.stack([np.cos(self.heading), np.sin(self.heading)], -1)  # forward (n,2,2)
            nrm = np.stack([-h[..., 1], h[..., 0]], -1)  # left
            v_long = np.sum(self.vel * h, -1)
            v_lat = np.sum(self.vel * nrm, -1)

            # wheel/floor forces, capped by the friction circle
            f_long = F_MOTOR * self.thrust_state - DRAG * m * v_long
            f_lat = -LATERAL_GAIN * m * v_lat / dt
            f_wheel = f_long[..., None] * h + f_lat[..., None] * nrm
            f_norm = np.linalg.norm(f_wheel, axis=-1)
            scale = np.minimum(1.0, grip / np.maximum(f_norm, 1e-9))
            force = f_wheel * scale[..., None] - AIR_DRAG * m[..., None] * self.vel

            # robot-robot spring-damper contact
            d = self.pos[:, 1] - self.pos[:, 0]  # (n,2), from robot 0 to robot 1
            dist = np.linalg.norm(d, axis=-1)
            overlap = 2 * ROBOT_RADIUS - dist
            touching = overlap > 0
            if np.any(touching):
                u = d / np.maximum(dist, 1e-9)[:, None]
                approach = -np.sum((self.vel[:, 1] - self.vel[:, 0]) * u, -1)  # >0 when closing
                fc = np.where(touching, np.maximum(0.0, K_CONTACT * overlap + C_CONTACT * approach), 0.0)
                force[:, 0] -= fc[:, None] * u
                force[:, 1] += fc[:, None] * u

            # semi-implicit Euler
            self.vel += force / m[..., None] * dt
            self.pos += self.vel * dt
            self.heading += self.omega * dt

        self.heading = (self.heading + np.pi) % (2 * np.pi) - np.pi
        self.t += 1

        # termination: a robot whose center leaves the ring loses
        out = np.linalg.norm(self.pos, axis=-1) > RING_RADIUS  # (n,2)
        new = ~self.done
        r0_won = out[:, 1] & ~out[:, 0]
        r1_won = out[:, 0] & ~out[:, 1]
        ended = (out[:, 0] | out[:, 1] | (self.t >= MAX_STEPS)) & new
        self.outcome[new & r0_won] = 1
        self.outcome[new & r1_won] = -1
        self.end_step[ended] = self.t
        self.done |= ended

        o0, o1 = self._delayed_noisy_obs(self._observe_true())
        return o0, o1, self.done.copy(), self.outcome.copy()

    # ------------------------------------------------------------ helpers
    def kinetic_energy(self) -> np.ndarray:
        """Translational kinetic energy per env (n,)."""
        return 0.5 * np.sum(self.params.mass * np.sum(self.vel**2, -1), -1)

    def snapshot(self) -> dict:
        """Copy of the state for rendering."""
        return {
            "pos": self.pos.copy(),
            "heading": self.heading.copy(),
            "vel": self.vel.copy(),
            "done": self.done.copy(),
            "outcome": self.outcome.copy(),
            "t": self.t,
        }