"""Robustness sweep: the same scene many times with randomized physics, in parallel on the GPU.

Each copy ("world") gets its own
    - ground/object friction scale      (geom_friction)
    - mass scale of every free object   (body_mass + body_inertia, objects only, not the robot)
    - small joint noise on the robot's starting pose
    - one horizontal push on the robot's root body, in a random direction, at `push_at` seconds

and runs `duration_s` with the same controller as /simulate/run. The result is how many worlds end with the
robot fallen, and which parameters those were. Engine: MuJoCo Warp (NVIDIA Warp, one GPU kernel launch per
step for all worlds) when CUDA is available, else plain MuJoCo on the CPU, one world after another. Both apply
exactly the same randomization, so results agree up to floating point.
"""

from __future__ import annotations

import json
import math
import time
from dataclasses import asdict, dataclass
from pathlib import Path

import mujoco
import numpy as np

from api.robots import ROBOTS, reset_to_robot_keyframe
from worker.sim_runner import _hold_controller, config_for

MAX_WORLDS = 256
MAX_DURATION_S = 10.0


@dataclass
class SweepConfig:
    n_worlds: int = 64
    duration_s: float = 5.0
    controller: str = "hold"
    friction: tuple = (0.5, 1.5)  # scale range
    object_mass: tuple = (0.5, 2.0)  # scale range for free objects
    push_n: tuple = (0.0, 80.0)  # push force range (N), applied for push_s
    push_at: float = 1.0
    push_s: float = 0.2
    joint_noise: float = 0.02  # rad, uniform
    seed: int = 0


@dataclass
class World:
    friction: float
    object_mass: float
    push_n: float
    push_dir_deg: float
    fell: bool = False
    fell_at: float | None = None
    min_height: float | None = None
    final_height: float | None = None


def sample_worlds(cfg: SweepConfig) -> list[World]:
    rng = np.random.default_rng(cfg.seed)
    return [
        World(
            friction=round(float(rng.uniform(*cfg.friction)), 3),
            object_mass=round(float(rng.uniform(*cfg.object_mass)), 3),
            push_n=round(float(rng.uniform(*cfg.push_n)), 1),
            push_dir_deg=round(float(rng.uniform(0, 360)), 1),
        )
        for _ in range(cfg.n_worlds)
    ]


def _robot_bodies(model: mujoco.MjModel, root: int) -> np.ndarray:
    inside = np.zeros(model.nbody, bool)
    for b in range(1, model.nbody):
        a = b
        while a > 0 and a != root:
            a = model.body_parentid[a]
        inside[b] = a == root and root > 0
    return inside


def _free_objects(model: mujoco.MjModel, robot_bodies: np.ndarray) -> np.ndarray:
    """Bodies (and their children) hanging off a free joint that are not part of the robot."""
    free = np.zeros(model.nbody, bool)
    for j in range(model.njnt):
        b = model.jnt_bodyid[j]
        if model.jnt_type[j] == mujoco.mjtJoint.mjJNT_FREE and not robot_bodies[b]:
            free[b] = True
    for b in range(1, model.nbody):  # children of free objects
        if free[model.body_parentid[b]]:
            free[b] = True
    return free


def _joint_noise_mask(model: mujoco.MjModel, robot_bodies: np.ndarray) -> np.ndarray:
    """qpos entries of the robot's hinge/slide joints (the free joint is left alone)."""
    mask = np.zeros(model.nq, bool)
    for j in range(model.njnt):
        if robot_bodies[model.jnt_bodyid[j]] and model.jnt_type[j] in (
            mujoco.mjtJoint.mjJNT_HINGE,
            mujoco.mjtJoint.mjJNT_SLIDE,
        ):
            mask[model.jnt_qposadr[j]] = True
    return mask


def _setup(scene: Path, robot: str | None, cfg: SweepConfig):
    model = mujoco.MjModel.from_xml_path(str(scene))
    data = mujoco.MjData(model)
    run_cfg = config_for(robot, cfg.duration_s, cfg.controller, render=False)
    reset_to_robot_keyframe(model, data, run_cfg.root_body, run_cfg.home_qpos)
    root = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_BODY, run_cfg.root_body) if run_cfg.root_body else -1
    robot_bodies = _robot_bodies(model, root)
    return model, data, run_cfg, root, robot_bodies


def _world_params(model, worlds, robot_bodies, cfg: SweepConfig, qpos0: np.ndarray):
    """Per-world arrays: friction, masses, inertias, starting qpos, push force."""
    n = len(worlds)
    rng = np.random.default_rng(cfg.seed + 1)
    objects = _free_objects(model, robot_bodies)
    noise_mask = _joint_noise_mask(model, robot_bodies)
    friction = np.repeat(model.geom_friction[None], n, 0).copy()
    mass = np.repeat(model.body_mass[None], n, 0).copy()
    inertia = np.repeat(model.body_inertia[None], n, 0).copy()
    qpos = np.repeat(qpos0[None], n, 0).copy()
    force = np.zeros((n, 3))
    for i, w in enumerate(worlds):
        friction[i, :, 0] *= w.friction  # sliding friction of every geom
        mass[i, objects] *= w.object_mass
        inertia[i, objects] *= w.object_mass
        qpos[i, noise_mask] += rng.uniform(-cfg.joint_noise, cfg.joint_noise, noise_mask.sum())
        a = math.radians(w.push_dir_deg)
        force[i] = (w.push_n * math.cos(a), w.push_n * math.sin(a), 0.0)
    return friction, mass, inertia, qpos, force


def _finish(worlds, heights, fell_at, final, fall_height) -> None:
    for i, w in enumerate(worlds):
        w.min_height = round(float(heights[i]), 3)
        w.final_height = round(float(final[i]), 3)
        if fall_height is not None:
            w.fell = fell_at[i] is not None
            w.fell_at = fell_at[i]


# --------------------------------------------------------------------------- engines


def run_cpu(scene: Path, robot: str | None, cfg: SweepConfig, worlds: list[World]) -> None:
    model, data, run_cfg, root, robot_bodies = _setup(scene, robot, cfg)
    friction, mass, inertia, qpos, force = _world_params(model, worlds, robot_bodies, cfg, data.qpos.copy())
    base = (model.geom_friction.copy(), model.body_mass.copy(), model.body_inertia.copy())
    steps = int(math.ceil(cfg.duration_s / model.opt.timestep))
    push = (int(cfg.push_at / model.opt.timestep), int((cfg.push_at + cfg.push_s) / model.opt.timestep))
    n = len(worlds)
    heights, final, fell_at = np.full(n, np.inf), np.zeros(n), [None] * n
    for i in range(n):
        model.geom_friction[:], model.body_mass[:], model.body_inertia[:] = friction[i], mass[i], inertia[i]
        mujoco.mj_resetData(model, data)
        data.qpos[:] = qpos[i]
        mujoco.mj_forward(model, data)
        hold = _hold_controller(model, data, run_cfg) if cfg.controller == "hold" else None
        for s in range(steps):
            if hold is not None:
                hold.apply(data)
            if root > 0:
                data.xfrc_applied[root, :3] = force[i] if push[0] <= s < push[1] else 0.0
            mujoco.mj_step(model, data)
            if root > 0:
                z = float(data.xpos[root, 2])
                heights[i] = min(heights[i], z)
                if fell_at[i] is None and run_cfg.fall_height is not None and z < run_cfg.fall_height:
                    fell_at[i] = round(float(data.time), 2)
        final[i] = float(data.xpos[root, 2]) if root > 0 else 0.0
    model.geom_friction[:], model.body_mass[:], model.body_inertia[:] = base
    _finish(worlds, heights, fell_at, final, run_cfg.fall_height)


def run_warp(
    scene: Path, robot: str | None, cfg: SweepConfig, worlds: list[World], device: str | None = None
) -> str | None:
    """Returns a note when the model had to be adapted for MuJoCo Warp."""
    import mujoco_warp as mjw
    import warp as wp

    wp.config.quiet = True
    model, data, run_cfg, root, robot_bodies = _setup(scene, robot, cfg)
    note = None
    friction, mass, inertia, qpos, force = _world_params(model, worlds, robot_bodies, cfg, data.qpos.copy())
    n = len(worlds)
    hold = _hold_controller(model, data, run_cfg) if cfg.controller == "hold" else None
    with wp.ScopedDevice(device):
        batch = {"geom_friction": n, "body_mass": n, "body_inertia": n}
        try:
            m = mjw.put_model(model, batch_sizes=batch)
        except NotImplementedError as e:
            if "margin" not in str(e):
                raise
            # e.g. Menagerie's Go2 gives its feet 1 mm contact margins, which MuJoCo Warp's CCD does not support.
            note = f"contact margins up to {model.geom_margin.max() * 1000:g} mm set to 0 for MuJoCo Warp"
            model.geom_margin[:] = 0.0
            m = mjw.put_model(model, batch_sizes=batch)
        m.geom_friction.assign(friction.astype(np.float32))
        m.body_mass.assign(mass.astype(np.float32))
        m.body_inertia.assign(inertia.astype(np.float32))
        d = mjw.put_data(model, data, nworld=n, nconmax=max(64, 8 * model.ngeom), njmax=max(256, 40 * model.nv))
        d.qpos.assign(qpos.astype(np.float32))
        mjw.forward(m, d)

        ctrl_base = np.repeat((hold.ctrl if hold else np.zeros(model.nu))[None], n, 0).astype(np.float32)
        pd = hold is not None and len(hold.motors) > 0
        if pd:
            motors = wp.array(hold.motors.astype(np.int32), dtype=int)
            qadr = wp.array(hold.qadr.astype(np.int32), dtype=int)
            vadr = wp.array(hold.vadr.astype(np.int32), dtype=int)
            target, kp, kd, lo, hi = (
                wp.array(a.astype(np.float32), dtype=float) for a in (hold.target, hold.kp, hold.kd, hold.lo, hold.hi)
            )
        d.ctrl.assign(ctrl_base)

        xfrc = np.zeros((n, model.nbody, 6), np.float32)
        steps = int(math.ceil(cfg.duration_s / model.opt.timestep))
        push = (int(cfg.push_at / model.opt.timestep), int((cfg.push_at + cfg.push_s) / model.opt.timestep))
        sample_every = max(1, int(0.02 / model.opt.timestep))
        heights, fell_at = np.full(n, np.inf), [None] * n
        for s in range(steps):
            if root > 0 and s in push:
                xfrc[:, root, :3] = force if s == push[0] else 0.0
                d.xfrc_applied.assign(xfrc)
            if pd:
                wp.launch(
                    _pd_kernel(),
                    dim=(n, len(hold.motors)),
                    inputs=[d.qpos, d.qvel, motors, qadr, vadr, target, kp, kd, lo, hi],
                    outputs=[d.ctrl],
                )
            mjw.step(m, d)
            if root > 0 and (s % sample_every == 0 or s == steps - 1):
                z = d.xpos.numpy()[:, root, 2]
                heights = np.minimum(heights, z)
                if run_cfg.fall_height is not None:
                    t = round((s + 1) * model.opt.timestep, 2)
                    for i in np.nonzero(z < run_cfg.fall_height)[0]:
                        if fell_at[i] is None:
                            fell_at[i] = t
        final = d.xpos.numpy()[:, root, 2] if root > 0 else np.zeros(n)
    _finish(worlds, heights, fell_at, final, run_cfg.fall_height)
    return note


def _make_pd_kernel():
    import warp as wp

    @wp.kernel
    def pd(
        qpos: wp.array2d(dtype=float),
        qvel: wp.array2d(dtype=float),
        motors: wp.array(dtype=int),
        qadr: wp.array(dtype=int),
        vadr: wp.array(dtype=int),
        target: wp.array(dtype=float),
        kp: wp.array(dtype=float),
        kd: wp.array(dtype=float),
        lo: wp.array(dtype=float),
        hi: wp.array(dtype=float),
        ctrl: wp.array2d(dtype=float),
    ):
        w, k = wp.tid()
        tau = kp[k] * (target[k] - qpos[w, qadr[k]]) - kd[k] * qvel[w, vadr[k]]
        ctrl[w, motors[k]] = wp.clamp(tau, lo[k], hi[k])

    return pd


_PD_KERNEL = None


def _pd_kernel():
    """Compiled on first use, so importing this module never needs Warp."""
    global _PD_KERNEL
    if _PD_KERNEL is None:
        _PD_KERNEL = _make_pd_kernel()
    return _PD_KERNEL


def cuda_available() -> bool:
    try:
        import warp as wp

        wp.config.quiet = True
        wp.init()
        return wp.is_cuda_available()
    except Exception:
        return False


# --------------------------------------------------------------------------- job


def sweep(scene: Path, robot: str | None, cfg: SweepConfig, engine: str = "auto") -> dict:
    if not 1 <= cfg.n_worlds <= MAX_WORLDS:
        raise ValueError(f"n_worlds must be 1..{MAX_WORLDS}")
    cfg.duration_s = min(cfg.duration_s, MAX_DURATION_S)
    if engine == "auto":
        engine = "warp" if cuda_available() else "mujoco"
    worlds = sample_worlds(cfg)
    started = time.monotonic()
    note = None
    if engine == "warp":
        try:
            note = run_warp(scene, robot, cfg, worlds, device="cuda:0" if cuda_available() else "cpu")
        except NotImplementedError as e:  # a model feature MuJoCo Warp lacks: same sweep on the CPU
            engine, note = "mujoco", f"MuJoCo Warp cannot run this scene ({e}); ran on the CPU instead"
            worlds = sample_worlds(cfg)
            run_cpu(scene, robot, cfg, worlds)
    else:
        run_cpu(scene, robot, cfg, worlds)
    label = "MuJoCo (CPU)"
    if engine == "warp":
        label = "MuJoCo Warp (GPU)" if cuda_available() else "MuJoCo Warp (CPU device)"
    result = summarize(worlds, cfg, robot, label, time.monotonic() - started)
    if note:
        result["summary"]["engine_note"] = note
    return result


def summarize(worlds: list[World], cfg: SweepConfig, robot: str | None, engine: str, wall: float) -> dict:
    spec = ROBOTS.get(robot) if robot else None
    detect = spec is not None and spec.fall_height is not None
    fell = [w for w in worlds if w.fell]
    n = len(worlds)
    summary = {
        "worlds": n,
        "engine": engine,
        "simulated_seconds_each": cfg.duration_s,
        "wall_seconds": round(wall, 1),
        "controller": cfg.controller,
        "ranges": {
            "friction_scale": list(cfg.friction),
            "object_mass_scale": list(cfg.object_mass),
            "push_newtons": list(cfg.push_n),
            "push_at_s": cfg.push_at,
            "push_duration_s": cfg.push_s,
            "joint_noise_rad": cfg.joint_noise,
        },
        "seed": cfg.seed,
    }
    if detect:
        summary["fell"] = len(fell)
        summary["fall_rate"] = round(len(fell) / n, 3)
        if fell:
            summary["mean_push_n_fell"] = round(float(np.mean([w.push_n for w in fell])), 1)
            summary["min_push_n_fell"] = min(w.push_n for w in fell)
        stood = [w for w in worlds if not w.fell]
        if stood:
            summary["max_push_n_survived"] = max(w.push_n for w in stood)
        summary["headline"] = f"{len(fell)}/{n} fell ({100 * len(fell) / n:.0f}%)"
        if fell:
            summary["headline"] += f" · lightest push that toppled it: {summary['min_push_n_fell']:g} N"
    else:
        summary["headline"] = f"{n} worlds simulated (no fall detection for this robot)"
    return {"summary": summary, "worlds": [asdict(w) for w in worlds]}


def run_sweep_job(job_input: dict, menagerie_dir: Path | None = None, engine: str = "auto") -> dict:
    """Worker entry point. Returns the summary and the per-world table as an inline artifact."""
    import os
    import shutil

    from api.mjcf_check import SCENE_FILE, prepare_scene_dir, static_check

    mjcf, robot = job_input["mjcf"], job_input.get("robot")
    failed = static_check(mjcf, robot)
    if failed:
        raise ValueError(f"scene rejected: {failed.error}")
    fields = {k: job_input[k] for k in ("n_worlds", "duration_s", "controller", "seed") if k in job_input}
    for k in ("friction", "object_mass", "push_n"):
        if k in job_input:
            fields[k] = tuple(float(x) for x in job_input[k])
    cfg = SweepConfig(**fields)
    menagerie = Path(menagerie_dir or os.environ.get("MENAGERIE_DIR", "/opt/mujoco_menagerie"))
    scene_dir = prepare_scene_dir(mjcf, robot, menagerie)
    try:
        result = sweep(scene_dir / SCENE_FILE, robot, cfg, engine)
    finally:
        shutil.rmtree(scene_dir, ignore_errors=True)
    worker_seconds = result["summary"]["wall_seconds"]
    return {
        "summary": result["summary"],
        "worker_seconds": worker_seconds,
        "artifacts": {"sweep.json": json.dumps(result, indent=1).encode()},
    }