"""Tests: physics sanity, symmetry, and that ES actually improves fitness."""

import os
import sys

import numpy as np
import pytest

sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from policies import ChargerPolicy, MLPPolicy, RandomPolicy, init_params, run_matches  # noqa: E402
from sumo_env import (  # noqa: E402
    K_CONTACT,
    MAX_STEPS,
    OBS_DIM,
    ROBOT_RADIUS,
    PhysParams,
    SumoEnv,
    random_starts,
)
from train import Adam, centered_ranks, es_gradient, run_stage, shaped_fitness  # noqa: E402


class ZeroPolicy:
    """Does nothing (a 'sitting duck')."""

    def reset(self, n, rng=None):
        pass

    def act(self, obs):
        return np.zeros((obs.shape[0], 2))


class ForwardPolicy:
    def reset(self, n, rng=None):
        pass

    def act(self, obs):
        return np.tile([1.0, 0.0], (obs.shape[0], 1))


def make_starts(p0, p1, h0, h1):
    return {"pos": np.array([[p0, p1]], float), "heading": np.array([[h0, h1]], float)}


# ---------------------------------------------------------------- physics
def test_pushed_out_robot_loses():
    env = SumoEnv(1)
    starts = make_starts([0.5, 0.0], [0.5 + 2 * ROBOT_RADIUS + 0.01, 0.0], 0.0, np.pi / 2)
    res = run_matches(env, ForwardPolicy(), ZeroPolicy(), starts=starts)
    assert res["outcome"][0] == 1, "robot 0 should push the passive robot 1 out"
    assert res["end_step"][0] < MAX_STEPS
    # and with roles swapped, robot 1 wins
    starts = make_starts([0.5 + 2 * ROBOT_RADIUS + 0.01, 0.0], [0.5, 0.0], np.pi / 2, 0.0)
    res = run_matches(SumoEnv(1), ZeroPolicy(), ForwardPolicy(), starts=starts)
    assert res["outcome"][0] == -1


def test_leaving_ring_alone_loses():
    env = SumoEnv(1)
    env.reset(make_starts([0.95, 0.0], [-0.3, 0.0], 0.0, 0.0))
    for _ in range(20):
        _, _, done, outcome = env.step(np.array([[1.0, 0.0]]), np.zeros((1, 2)))
    assert done[0] and outcome[0] == -1


def test_passive_robots_tie_at_time_cap():
    env = SumoEnv(4, seed=0)
    res = run_matches(env, ZeroPolicy(), ZeroPolicy())
    assert np.all(res["outcome"] == 0) and np.all(res["end_step"] == MAX_STEPS)


def test_symmetric_start_symmetric_outcome():
    """Point-reflected starts + identical deterministic policy => mirrored trajectories, tie."""
    rng = np.random.default_rng(3)
    n = 32
    p0 = random_starts(n, rng)["pos"][:, 0]
    h0 = rng.uniform(-np.pi, np.pi, n)
    starts = {"pos": np.stack([p0, -p0], 1), "heading": np.stack([h0, h0 + np.pi], 1)}
    env = SumoEnv(n)
    res = run_matches(env, ChargerPolicy(), ChargerPolicy(), starts=starts, record=True)
    for fr in res["frames"]:
        np.testing.assert_allclose(fr["pos"][:, 0], -fr["pos"][:, 1], atol=1e-6)
    assert np.all(res["outcome"] == 0)


def test_contact_dissipates_energy_and_is_stable():
    """Frictionless head-on collision: total mechanical energy (kinetic + contact
    spring) never grows, discs don't pass through each other, robots bounce apart."""
    p = PhysParams.nominal(1)
    p.mu[:] = 0.0  # no wheel forces at all -> only contact + air drag
    env = SumoEnv(1, p)
    env.reset({"pos": np.array([[[-0.3, 0.0], [0.3, 0.0]]]), "heading": np.zeros((1, 2)),
               "vel": np.array([[[1.5, 0.0], [-1.5, 0.0]]])})

    def total_energy():
        overlap = max(0.0, 2 * ROBOT_RADIUS - np.linalg.norm(env.pos[0, 1] - env.pos[0, 0]))
        return env.kinetic_energy()[0] + 0.5 * K_CONTACT * overlap**2

    e0 = total_energy()
    energies, min_gap = [], np.inf
    for _ in range(30):
        env.step(np.zeros((1, 2)), np.zeros((1, 2)))
        energies.append(total_energy())
        min_gap = min(min_gap, np.linalg.norm(env.pos[0, 1] - env.pos[0, 0]))
    # explicit integration may wiggle the energy by a hair; allow 1% of the initial energy
    assert np.all(np.diff([e0] + energies) <= 0.01 * e0)
    assert energies[-1] < e0
    assert min_gap > ROBOT_RADIUS  # discs compress a little but do not pass through each other
    assert env.vel[0, 0, 0] < 0 < env.vel[0, 1, 0]  # bounced apart


def test_random_play_stays_finite():
    p = PhysParams.randomized(256, np.random.default_rng(0))
    env = SumoEnv(256, p, seed=0)
    res = run_matches(env, RandomPolicy(hold=3), RandomPolicy(hold=3), np.random.default_rng(0))
    assert np.all(np.isfinite(env.pos)) and np.all(np.isfinite(env.vel))
    assert np.all(np.abs(env.vel) < 20)
    assert set(np.unique(res["outcome"])) <= {-1, 0, 1}


def test_icy_floor_is_slower():
    """Same command, lower friction -> less acceleration (traction-limited)."""
    speeds = []
    for mu in (0.8, 0.2):
        p = PhysParams.nominal(1)
        p.mu[:] = mu
        env = SumoEnv(1, p)
        env.reset(make_starts([-0.5, 0.0], [0.5, 0.5], 0.0, 0.0))
        for _ in range(5):
            env.step(np.array([[1.0, 0.0]]), np.zeros((1, 2)))
        speeds.append(np.linalg.norm(env.vel[0, 0]))
    assert speeds[1] < 0.5 * speeds[0]


# ---------------------------------------------------------------- observations
def test_observation_swap_symmetry():
    """Swapping the two robots swaps their observations exactly."""
    rng = np.random.default_rng(1)
    n = 16
    s = random_starts(n, rng)
    s["vel"] = rng.normal(size=(n, 2, 2))
    env = SumoEnv(n)
    o0, o1 = env.reset(s)
    sw = {k: v[:, ::-1].copy() for k, v in s.items()}
    q0, q1 = SumoEnv(n).reset(sw)
    np.testing.assert_allclose(o0, q1, atol=1e-12)
    np.testing.assert_allclose(o1, q0, atol=1e-12)
    assert o0.shape == (n, OBS_DIM)


def test_observation_rotation_invariance():
    rng = np.random.default_rng(2)
    n = 16
    s = random_starts(n, rng)
    s["vel"] = rng.normal(size=(n, 2, 2))
    a = 0.7
    R = np.array([[np.cos(a), -np.sin(a)], [np.sin(a), np.cos(a)]])
    rs = {"pos": s["pos"] @ R.T, "vel": s["vel"] @ R.T, "heading": s["heading"] + a}
    o = SumoEnv(n).reset(s)
    q = SumoEnv(n).reset(rs)
    for x, y in zip(o, q):
        np.testing.assert_allclose(x, y, atol=1e-9)


# ---------------------------------------------------------------- ES
def test_centered_ranks():
    r = centered_ranks(np.array([10.0, -3.0, 2.0]))
    np.testing.assert_allclose(r, [0.5, -0.5, 0.0])


def test_es_on_quadratic():
    rng = np.random.default_rng(0)
    target = rng.normal(size=20)
    theta = np.zeros(20)
    opt = Adam(20, lr=0.05)
    for _ in range(300):
        eps = rng.standard_normal((16, 20))
        members = np.concatenate([theta + 0.1 * eps, theta - 0.1 * eps])
        fit = -np.sum((members - target) ** 2, 1)
        theta = theta + opt.step(es_gradient(fit, eps, 0.1))
    assert np.sum((theta - target) ** 2) < 0.05 * np.sum(target**2)


def test_es_improves_sumo_fitness(tmp_path):
    """A tiny ES run against the hand-coded charger raises the policy's fitness."""
    rng = np.random.default_rng(0)
    theta0 = init_params(rng)

    def score(theta):
        env = SumoEnv(256, seed=99)
        res = run_matches(env, MLPPolicy(theta), ChargerPolicy(), np.random.default_rng(99))
        return shaped_fitness(res, 0.3).mean()

    cfg = dict(gens=12, pairs=8, episodes=16, sigma=0.05, lr=0.03, l2=0.0, shape=0.3,
               eval_every=10**9, snap_every=10**9, eval_n=16)
    theta, _, log = run_stage("T", theta0, cfg, str(tmp_path), opp_mode="charger", seed=1, verbose=False)
    before, after = score(theta0), score(theta)
    assert after > before + 0.2, (before, after)
    assert os.path.exists(tmp_path / "T.npz")


# ---------------------------------------------------------------- tools
def test_watch_renders_gif(tmp_path):
    import watch

    out = tmp_path / "m.gif"
    watch.main(["--p0", "L1", "--p1", "L0", "--seed", "1", "--save", str(out), "--dpi", "30"])
    assert out.exists() and out.stat().st_size > 1000


def test_evaluate_play_rates_sum_to_one():
    import evaluate

    r = evaluate.play("L1", "L0", 64, perturbation="icy")
    assert abs(r["win"] + r["loss"] + r["tie"] - 1) < 1e-9
    assert r["win"] > r["loss"]  # the charger beats random wandering, even on ice