#!/usr/bin/env python3
"""PraMD live demo backend. The same file runs on a Hugging Face Space and on a laptop.

Loads the frozen OMG-100M generator with the Run 3 MuRIL adapter exactly as tools/render_run3.py
does -- build("100m", ..., norm_mode="corpus"), the Run 3 calibration, adapter_step036000.pt -- and
serves text -> Unitree G1 qpos_36 as JSON for the project website.

Weights    local files under PRAMD_MODELS if present, otherwise downloaded from the Hugging Face Hub.
           The adapter repo is private: give the server a read-only HF_TOKEN (a Space secret).
           Never commit the token.
Model      exactly the one requested: mode="pramd" sends every instruction, in any script, through
           MuRIL + the adapter; mode="stock" sends every instruction through the backbone's own T5.
Noise      drawn on the CPU from the request seed and handed to the sampler, so one seed gives the
           same motion on CPU, Apple GPU or CUDA. Examples verified on one machine hold on another.
Start pose one real standing frame held still for the whole history, so the instruction alone
           decides the motion.
"""
from __future__ import annotations

import os

os.environ.setdefault("PYTORCH_ENABLE_MPS_FALLBACK", "1")   # must precede the torch import

import hashlib
import json
import math
import sys
import threading
import time
import traceback
from contextlib import asynccontextmanager
from pathlib import Path

import numpy as np
import torch
from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from fastapi.staticfiles import StaticFiles
from pydantic import BaseModel, Field

HERE = Path(__file__).resolve().parent
REPO = Path(os.environ.get("PRAMD_REPO", str(HERE / "omg"))).expanduser().resolve()
MODELS = Path(os.environ.get("PRAMD_MODELS", str(HERE / "models"))).expanduser().resolve()
SITE = Path(os.environ["PRAMD_SITE"]).expanduser().resolve() if os.environ.get("PRAMD_SITE") else None
HOST = os.environ.get("PRAMD_HOST", "0.0.0.0")
PORT = int(os.environ.get("PRAMD_PORT", "7860"))
SEED_KIND = os.environ.get("PRAMD_SEED_KIND", "static")          # "static" | "window"
THREADS = int(os.environ.get("PRAMD_THREADS", "0"))               # 0 keeps torch's default
ALLOWED_ORIGINS = [o.strip() for o in os.environ.get("PRAMD_ALLOWED_ORIGINS", "*").split(",") if o.strip()]

# file under MODELS  ->  (Hub repo, repo type, path in repo)
FILES = {
    "omg": ("omg/checkpoints/updated/100m/sstep=170000.ckpt",
            "THU-MARS/OMG", "model", "checkpoints/updated/100m/sstep=170000.ckpt"),
    "adapter": ("adapter/maila-run3-20260909-204625/adapter_step036000.pt",
                "CodeSushh/maila-run3-motionpivot", "model", "maila-run3-20260909-204625/adapter_step036000.pt"),
    "calibration": ("calib/calibration/adapter_calibration_run3.pt",
                    "CodeSushh/maila-run3-motionpivot", "dataset", "calibration/adapter_calibration_run3.pt"),
}
SNAPSHOTS = {  # directory under MODELS -> (Hub repo, files to fetch)
    "t5": ("t5-base-local", "t5-base", ["*.json", "*.model", "model.safetensors"]),
    "muril": ("muril-base-cased", "google/muril-base-cased", ["*.json", "*.txt", "pytorch_model.bin"]),
}
PATHS = {
    **{k: MODELS / v[0] for k, v in FILES.items()},
    **{k: MODELS / v[0] for k, v in SNAPSHOTS.items()},
    "seed": HERE / "seed_qpos10.npy",
}
# sha256[:16] recorded in RUN3_SSOT.md section 10 for the evaluated artifacts
EXPECTED_SHA16 = {"adapter": "2a3ea341891d9609", "omg": "7ec1afc66786ef27"}

L, H, FPS = 10, 60, 30.0

SCRIPTS = [  # (code, label, first, last) -- Unicode blocks
    ("hi", "Hindi", 0x0900, 0x097F), ("bn", "Bengali", 0x0980, 0x09FF),
    ("pa", "Punjabi", 0x0A00, 0x0A7F), ("gu", "Gujarati", 0x0A80, 0x0AFF),
    ("or", "Odia", 0x0B00, 0x0B7F), ("ta", "Tamil", 0x0B80, 0x0BFF),
    ("te", "Telugu", 0x0C00, 0x0C7F), ("kn", "Kannada", 0x0C80, 0x0CFF),
    ("ml", "Malayalam", 0x0D00, 0x0D7F),
]
TRAINED = {"hi", "bn", "ta", "te"}

STATE: dict = {"ready": False, "stage": "starting", "error": None, "device": None, "queue": 0}
MODEL: dict = {}
LOCK = threading.Lock()
QUEUE_LOCK = threading.Lock()


def sha16(path: Path) -> str:
    h = hashlib.sha256()
    with open(path, "rb") as f:
        for block in iter(lambda: f.read(1 << 20), b""):
            h.update(block)
    return h.hexdigest()[:16]


def pick_device() -> torch.device:
    if os.environ.get("PRAMD_DEVICE"):
        return torch.device(os.environ["PRAMD_DEVICE"])
    if torch.cuda.is_available():
        return torch.device("cuda")
    if torch.backends.mps.is_available():
        return torch.device("mps")
    return torch.device("cpu")


def ensure_weights() -> None:
    """Fetch anything missing from the Hub. Local files are used as they are."""
    from huggingface_hub import hf_hub_download, snapshot_download

    for key, (rel, repo, rtype, name) in FILES.items():
        if (MODELS / rel).is_file():
            continue
        STATE["stage"] = f"downloading {key}"
        target_dir = MODELS / rel[: -len(name)].rstrip("/")
        try:
            hf_hub_download(repo_id=repo, filename=name, repo_type=rtype, local_dir=str(target_dir))
        except Exception as exc:
            if repo.startswith("CodeSushh/"):
                raise RuntimeError(f"could not download the private {key} from {repo}. Set a read-only "
                                   "HF_TOKEN secret with access to that repo.") from exc
            raise
    for key, (rel, repo, patterns) in SNAPSHOTS.items():
        if (MODELS / rel / "config.json").is_file():
            continue
        STATE["stage"] = f"downloading {key}"
        snapshot_download(repo_id=repo, allow_patterns=patterns, local_dir=str(MODELS / rel))


def detect_script(text: str) -> tuple[str, str]:
    counts: dict[str, int] = {}
    latin = 0
    for ch in text:
        if "a" <= ch.lower() <= "z":
            latin += 1
            continue
        o = ord(ch)
        for code, _, lo, hi in SCRIPTS:
            if lo <= o <= hi:
                counts[code] = counts.get(code, 0) + 1
                break
    if counts:
        code = max(counts, key=counts.get)
        return code, next(label for c, label, *_ in SCRIPTS if c == code)
    return ("en", "English") if latin else ("unknown", "Unrecognised script")


def _qmul(a: np.ndarray, b: np.ndarray) -> np.ndarray:
    aw, ax, ay, az = np.moveaxis(a, -1, 0)
    bw, bx, by, bz = np.moveaxis(b, -1, 0)
    return np.stack([aw * bw - ax * bx - ay * by - az * bz,
                     aw * bx + ax * bw + ay * bz - az * by,
                     aw * by - ax * bz + ay * bw + az * bx,
                     aw * bz + ax * by - ay * bx + az * bw], axis=-1)


def origin_facing_forward(q: np.ndarray) -> np.ndarray:
    """Move the seed to the origin and remove its heading, so every generation starts at the
    same place facing +x. Heading is a free gauge: the codec canonicalises about the anchor
    frame's yaw, so this changes where the robot stands, not what it does."""
    q = q.astype(np.float64).copy()
    w, x, y, z = q[0, 3:7]
    yaw = math.atan2(2 * (w * z + x * y), 1 - 2 * (y * y + z * z))
    rz = np.array([math.cos(-yaw / 2), 0.0, 0.0, math.sin(-yaw / 2)])
    q[:, 3:7] = _qmul(np.broadcast_to(rz, q[:, 3:7].shape), q[:, 3:7])
    c, s = math.cos(-yaw), math.sin(-yaw)
    xy = q[:, 0:2] - q[0, 0:2]
    q[:, 0], q[:, 1] = c * xy[:, 0] - s * xy[:, 1], s * xy[:, 0] + c * xy[:, 1]
    return q.astype(np.float32)


def initial_noise(seed: int, frames: int, feat_dim: int) -> list[torch.Tensor]:
    """One CPU-drawn noise tensor per 60-frame chunk: identical on every device."""
    g = torch.Generator(device="cpu").manual_seed(int(seed))
    chunks, left = [], int(frames)
    while left > 0:
        n = min(H, left)
        chunks.append(torch.randn(1, n, feat_dim, generator=g))
        left -= n
    return chunks


def _generate(text: str, frames: int, seed: int, cfg_scale: float, use_t5: bool) -> np.ndarray:
    model, enc, dev, w = MODEL["model"], MODEL["enc"], MODEL["dev"], MODEL["window"]
    batch = {  # same batch as render_run3.make_batch
        "history_features": w[:, :L], "prev_state_features": w[:, :L], "motion_features": w[:, L:],
        "canon_root_pos": MODEL["anchor_pos"], "canon_root_quat": MODEL["anchor_quat"],
        "mask": {"valid": torch.ones(1, H, dtype=torch.bool, device=dev)},
        "has_text": torch.ones(1, dtype=torch.bool, device=dev),
        "fps": torch.tensor([FPS], device=dev), "caption": [text],
    }
    noise = [n.to(dev) for n in initial_noise(seed, frames, model.representation.feat_dim)]
    with QUEUE_LOCK:
        STATE["queue"] += 1
    try:
        with LOCK, torch.no_grad():
            enc.passthrough_t5 = bool(use_t5)
            try:
                out = model.generate(batch, num_frames=frames, cfg_scale=cfg_scale, initial_noise_chunks=noise)
            finally:
                enc.passthrough_t5 = False
    finally:
        with QUEUE_LOCK:
            STATE["queue"] -= 1
    return out["qpos_36"][0].float().cpu().numpy()


def load() -> None:
    t0 = time.time()
    try:
        if THREADS > 0:
            torch.set_num_threads(THREADS)
        MODELS.mkdir(parents=True, exist_ok=True)
        ensure_weights()
        missing = [f"{k}: {p}" for k, p in PATHS.items() if not p.exists()]
        if missing:
            raise FileNotFoundError("missing files -> " + "; ".join(missing))

        STATE["stage"] = "verifying checksums"
        STATE["sha16"] = {k: sha16(PATHS[k]) for k in EXPECTED_SHA16}
        wrong = {k: v for k, v in STATE["sha16"].items() if v != EXPECTED_SHA16[k]}
        if wrong:
            raise RuntimeError(f"checkpoint checksum mismatch {wrong}; expected {EXPECTED_SHA16}")

        sys.path[:0] = [str(REPO / "src"), str(REPO / "tools")]
        os.chdir(REPO)                                  # hydra + relative asset paths
        dev = pick_device()
        STATE["device"] = dev.type

        STATE["stage"] = "building OMG-100M with the MuRIL adapter"
        from maila_train import build
        from omg.generation.conditions.t5 import FrozenT5TextEncoder
        import convert_g1_npz_to_omg125 as conv

        model = build("100m", str(PATHS["omg"]), str(PATHS["t5"]), str(PATHS["muril"]), dev,
                      True, True, 0.0, t5_calibration=str(PATHS["calibration"]), norm_mode="corpus")
        enc = model.text_encoder
        STATE["stage"] = "loading the T5 path"
        enc._t5 = FrozenT5TextEncoder(model_name=str(PATHS["t5"]), max_length=50,
                                      output_dim=768).to(dev).eval()
        STATE["stage"] = "loading adapter step 36,000"
        state = torch.load(PATHS["adapter"], map_location="cpu", weights_only=False)["adapter"]
        enc.adapter.load_state_dict(state)
        model.eval()
        enc.adapter.eval()          # build() leaves the adapter in train() -- dropout would be live

        STATE["stage"] = "encoding the start pose"
        rep = model.representation
        seed = np.load(PATHS["seed"]).astype(np.float32)
        if SEED_KIND == "static":                        # zero-velocity history: one real frame, held
            seed = np.repeat(seed[:1], len(seed), axis=0)
        seed = origin_facing_forward(seed)
        q70 = np.concatenate([seed, np.repeat(seed[-1:], L + H - len(seed), axis=0)], axis=0)
        feats, apk, aqk = conv.encode_windows(rep, q70, [0], L, H, str(dev))

        STATE["stage"] = "geometry check"
        probe = ["एक व्यक्ति आगे चलता है।", "একজন ব্যক্তি সামনে হাঁটে।",
                 "ஒரு நபர் முன்னோக்கி நடக்கிறார்.", "ఒక వ్యక్తి ముందుకు నడుస్తాడు."]
        with torch.no_grad():
            o = enc(probe, device=dev)
            mk = o["mask"].float()
            tn = float((o["context"].float().norm(dim=-1) * mk).sum() / mk.sum().clamp_min(1))
        if not 4.0 < tn < 12.0:                          # render_run3 gate V5
            raise RuntimeError(f"adapter output geometry out of band: token norm {tn:.3f}")
        STATE["token_norm"] = round(tn, 3)

        names = getattr(rep.kinematics, "joint_order", None)
        if not names:
            names = json.loads((REPO / "assets/robots/g1/g1_kinematics.json").read_text())["joint_order"]

        MODEL.update(model=model, enc=enc, dev=dev, window=torch.from_numpy(feats).to(dev),
                     anchor_pos=torch.from_numpy(apk).to(dev), anchor_quat=torch.from_numpy(aqk).to(dev),
                     joint_names=list(names), seed_frame=np.round(q70[L - 1], 5).tolist())

        STATE["stage"] = "warm-up"
        _generate("a person stands still.", H, 0, 2.5, use_t5=True)
        STATE.update(ready=True, stage="ready", seed_kind=SEED_KIND, threads=torch.get_num_threads(),
                     load_seconds=round(time.time() - t0, 1))
        print(f"[pramd] ready on {dev.type} ({torch.get_num_threads()} threads) in {STATE['load_seconds']}s  "
              f"token_norm={tn:.3f}", flush=True)
    except Exception as exc:                            # surfaced through /api/health
        traceback.print_exc()
        STATE.update(stage="failed", error=f"{type(exc).__name__}: {exc}")


class GenerateRequest(BaseModel):
    text: str = Field(..., min_length=1, max_length=400)
    seconds: float = Field(6.0, ge=2.0, le=10.0)
    mode: str = Field("pramd", pattern="^(pramd|stock)$")
    seed: int = Field(0, ge=0, le=2**31 - 1)
    cfg_scale: float = Field(2.5, ge=1.0, le=5.0)


@asynccontextmanager
async def lifespan(_: FastAPI):
    threading.Thread(target=load, daemon=True).start()
    yield


app = FastAPI(title="PraMD live demo", lifespan=lifespan)
app.add_middleware(CORSMiddleware, allow_origins=ALLOWED_ORIGINS, allow_methods=["GET", "POST"],
                   allow_headers=["Content-Type"])


@app.get("/api/health")
def health() -> dict:
    return {**STATE, "joint_names": MODEL.get("joint_names"), "seed_frame": MODEL.get("seed_frame")}


@app.post("/api/generate")
def generate(req: GenerateRequest) -> dict:
    if not STATE["ready"]:
        raise HTTPException(503, STATE.get("error") or f"model is still loading ({STATE['stage']})")
    text = req.text.strip()
    if not text:
        raise HTTPException(422, "instruction is empty")
    code, label = detect_script(text)
    use_t5 = req.mode == "stock"                  # the user's selection alone decides the model
    frames = int(round(req.seconds * FPS))
    t0 = time.time()
    try:
        q = _generate(text, frames, req.seed, req.cfg_scale, use_t5)
    except Exception as exc:
        traceback.print_exc()
        raise HTTPException(500, f"generation failed: {type(exc).__name__}: {exc}") from exc
    note = None
    if not use_t5 and code not in TRAINED:
        note = f"{label} is outside the adapter's training languages (Hindi, Bengali, Tamil, Telugu)."
    return {
        "frames": np.round(q, 4).tolist(), "fps": FPS, "lang": code, "language": label,
        "path": "Stock OMG (T5)" if use_t5 else "PraMD adapter (MuRIL)", "mode": req.mode,
        "gen_ms": int((time.time() - t0) * 1000), "device": STATE["device"],
        "joint_names": MODEL["joint_names"], "note": note,
    }


@app.get("/")
def root() -> dict:
    return {"service": "PraMD live demo", "health": "/api/health", "generate": "POST /api/generate"}


if SITE and SITE.exists():                         # local use: serve the project page too
    app.mount("/site", StaticFiles(directory=str(SITE), html=True), name="site")


if __name__ == "__main__":
    import uvicorn

    print(f"[pramd] repo {REPO}\n[pramd] models {MODELS}\n[pramd] listening on {HOST}:{PORT}", flush=True)
    uvicorn.run(app, host=HOST, port=PORT, log_level="warning")