"""Unitree G1 rough-terrain locomotion training, on the Space's GPU.

Training runs in a background thread so the Gradio server stays responsive; the
UI is a log tail plus start/stop. Brax PPO cannot be interrupted from outside,
so the stop button sets a flag that the progress callback checks and raises on.
"""

from __future__ import annotations

import os
import pickle
import threading
import traceback
from collections import deque
from datetime import datetime
from pathlib import Path

import gradio as gr

ENV_NAME = os.environ.get("ENV_NAME", "G1JoystickRoughTerrain")
NUM_TIMESTEPS = int(os.environ.get("NUM_TIMESTEPS", 200_000_000))
SEED = int(os.environ.get("SEED", 0))
HF_REPO = os.environ.get("HF_REPO", "").strip()
AUTO_START = os.environ.get("AUTO_START", "1") == "1"

# /data exists only when persistent storage is attached; fall back to /tmp.
OUT = Path("/data/ckpt") if Path("/data").is_dir() else Path("/tmp/ckpt")

LOG: deque[str] = deque(maxlen=800)
STATE = {"thread": None, "status": "idle", "stop": False, "step": 0}


def log(msg: str) -> None:
    LOG.append(f"[{datetime.now():%H:%M:%S}] {msg}")
    print(msg, flush=True)


class Stopped(Exception):
    pass


def _upload(path: Path) -> None:
    if not HF_REPO:
        return
    try:
        from huggingface_hub import HfApi
        api = HfApi()
        api.create_repo(HF_REPO, repo_type="model", exist_ok=True)
        api.upload_file(path_or_fileobj=str(path), path_in_repo=path.name,
                        repo_id=HF_REPO, repo_type="model")
        log(f"uploaded {path.name} -> {HF_REPO}")
    except Exception as e:                       # never let upload kill training
        log(f"upload failed ({type(e).__name__}): {e}")


def install_jax_pmap_shims() -> None:
    """Re-add jax.device_put_replicated / device_put_sharded if this JAX removed them.

    JAX deprecated both in 0.8.1 and removed them in 0.10.0 (April 2026), but the
    current brax *release* (0.14.2, which playground 0.2.0 requires) still calls
    device_put_replicated; only brax main has the sharding-based replacement.
    Shimming the two functions is a smaller intervention than pinning JAX down,
    which would risk the MJX/Warp stack that already compiles cleanly here.
    """
    import jax
    import jax.numpy as jnp
    import numpy as np

    def _sharding(devices):
        mesh = jax.sharding.Mesh(np.array(list(devices)), axis_names=("i",))
        return jax.sharding.NamedSharding(mesh, jax.sharding.PartitionSpec("i"))

    def device_put_replicated(x, devices):
        sharding, n = _sharding(devices), len(devices)

        def rep(leaf):
            stack = jnp.stack if isinstance(leaf, jax.Array) else np.stack
            return jax.device_put(stack([leaf] * n), sharding)

        return jax.tree_util.tree_map(rep, x)

    def device_put_sharded(shards, devices):
        sharding = _sharding(devices)

        def put(*leaves):
            stack = jnp.stack if isinstance(leaves[0], jax.Array) else np.stack
            return jax.device_put(stack(list(leaves)), sharding)

        return jax.tree_util.tree_map(put, *shards)

    for name, fn in (("device_put_replicated", device_put_replicated),
                     ("device_put_sharded", device_put_sharded)):
        try:
            getattr(jax, name)
        except AttributeError:
            setattr(jax, name, fn)
            log(f"shimmed jax.{name} (removed in this JAX version)")


def train_worker() -> None:
    try:
        STATE["status"] = "importing"
        log("importing jax / playground / brax ...")
        import functools
        import jax

        log(f"jax {jax.__version__}  devices: {jax.devices()}")
        install_jax_pmap_shims()
        import brax
        log(f"brax {brax.__version__}")
        if not any(d.platform == "gpu" for d in jax.devices()):
            log("WARNING: no GPU visible to JAX -- this will be very slow")

        from brax.training.agents.ppo import networks as ppo_networks
        from brax.training.agents.ppo import train as ppo
        from mujoco_playground import registry, wrapper
        from mujoco_playground.config import locomotion_params

        STATE["status"] = "building env"
        log(f"loading {ENV_NAME} ...")
        env = registry.load(ENV_NAME)
        env_cfg = registry.get_default_config(ENV_NAME)
        ppo_params = locomotion_params.brax_ppo_config(ENV_NAME)
        ppo_params.num_timesteps = NUM_TIMESTEPS
        log(f"num_envs={ppo_params.get('num_envs')} "
            f"batch_size={ppo_params.get('batch_size')} "
            f"timesteps={NUM_TIMESTEPS:,}")

        OUT.mkdir(parents=True, exist_ok=True)

        def progress(step, metrics):
            if STATE["stop"]:
                raise Stopped()
            STATE["step"] = int(step)
            rew = metrics.get("eval/episode_reward", float("nan"))
            ln = metrics.get("eval/avg_episode_length", float("nan"))
            log(f"step {int(step):>12,}  reward {rew:8.2f}  ep_len {ln:7.1f}")

        def policy_params_fn(step, make_policy, params):
            p = OUT / f"ckpt_{int(step)}.pkl"
            with open(p, "wb") as f:
                pickle.dump(params, f)
            _upload(p)

        ppo_kwargs = dict(ppo_params)
        net_cfg = ppo_kwargs.pop("network_factory", None)
        if net_cfg is not None:
            ppo_kwargs["network_factory"] = functools.partial(
                ppo_networks.make_ppo_networks, **dict(net_cfg))

        STATE["status"] = "training"
        log("compiling (first step takes several minutes) ...")
        _, params, _ = ppo.train(
            **ppo_kwargs,
            environment=env,
            eval_env=registry.load(ENV_NAME, config=env_cfg),
            wrap_env_fn=wrapper.wrap_for_brax_training,
            randomization_fn=registry.get_domain_randomizer(ENV_NAME),
            progress_fn=progress,
            policy_params_fn=policy_params_fn,
            seed=SEED,
        )
        final = OUT / "final.pkl"
        with open(final, "wb") as f:
            pickle.dump(params, f)
        _upload(final)
        STATE["status"] = "done"
        log("training complete")

    except Stopped:
        STATE["status"] = "stopped"
        log("stopped by user")
    except Exception as e:
        STATE["status"] = f"error: {type(e).__name__}"
        log(f"FAILED: {type(e).__name__}: {e}")
        log(traceback.format_exc()[-2000:])


def start():
    t = STATE["thread"]
    if t is not None and t.is_alive():
        return status_md()
    STATE["stop"] = False
    STATE["thread"] = threading.Thread(target=train_worker, daemon=True)
    STATE["thread"].start()
    return status_md()


def stop():
    STATE["stop"] = True
    log("stop requested -- will halt at the next eval boundary")
    return status_md()


def status_md() -> str:
    alive = STATE["thread"] is not None and STATE["thread"].is_alive()
    pct = 100.0 * STATE["step"] / max(NUM_TIMESTEPS, 1)
    return (f"**{ENV_NAME}** — status: `{STATE['status']}`"
            f"{' (running)' if alive else ''}  \n"
            f"step {STATE['step']:,} / {NUM_TIMESTEPS:,}  ({pct:.1f}%)  \n"
            f"checkpoints: `{OUT}`"
            + (f" → pushing to `{HF_REPO}`" if HF_REPO else
               "  \n⚠️ `HF_REPO` unset — checkpoints are not pushed to the Hub"))


def logs() -> str:
    return "\n".join(LOG) or "(no output yet)"


with gr.Blocks(title="G1 rough-terrain training") as demo:
    gr.Markdown("# Unitree G1 — rough-terrain locomotion training")
    st = gr.Markdown(status_md())
    with gr.Row():
        gr.Button("Start", variant="primary").click(start, outputs=st)
        gr.Button("Stop", variant="stop").click(stop, outputs=st)
    out = gr.Textbox(label="training log", lines=26, max_lines=26,
                     autoscroll=True, value=logs())
    timer = gr.Timer(3.0)
    timer.tick(logs, outputs=out)
    timer.tick(status_md, outputs=st)

if AUTO_START:
    start()

demo.queue().launch()