"""Pinned official Diffusion Policy components, adapted to the SO101 data contract.

The upstream encoder, ResNet-18 getter, crop randomizer, Mish U-Net and EMA are
vendored without edits. This wrapper deliberately keeps Issue 18's *current-t*
action convention, train-only min/max statistics and full-window dataset. It is
not an unmodified official task/runner or a pretrained real-robot policy.
"""
from __future__ import annotations

import copy
import hashlib
import importlib.metadata
import inspect
import json
import os
from pathlib import Path
import sys
import types

import numpy as np
import torch
from torch import Tensor, nn

HERE = Path(__file__).resolve().parent
REPO_ROOT = HERE.parents[1]
VENDOR = HERE / "vendor"
UPSTREAM_COMMIT = "5ba07ac6661db573af695b419a7947ecb704690f"
FORMAT = "so101_official_dp.visual_policy.v1"
NORMALIZER_KEYS = ("action_center", "action_span", "state_center", "state_span")
STATS_KEYS = ("action_min", "action_max", "state_min", "state_max")
DEFAULTS = {
    "history": 2, "horizon": 16, "image_size": 96,
    "crop_shape": None, "resize_shape": None, "random_crop": False, "imagenet_norm": True,
    "cameras": ["front", "gripper"], "context_dim": 2060,
    "execution_start": 0, "execute_steps": 8, "control_hz": 10,
    "down_dims": [512, 1024, 2048], "diffusion_step_embed_dim": 128,
    "kernel_size": 5, "n_groups": 8, "cond_predict_scale": True,
    "num_train_timesteps": 100, "beta_schedule": "squaredcos_cap_v2",
    "prediction_type": "epsilon", "clip_sample": True,
    "variance_type": "fixed_small", "default_sampler": "ddpm",
    "default_inference_steps": 100,
    "total_updates": 100000, "batch_size": 64, "microbatch_size": 8,
    "learning_rate": 1e-4, "adam_betas": [0.95, 0.999],
    "weight_decay": 1e-6, "warmup_steps": 500, "gradient_clip": None,
    "ema_power": 0.75, "ema_max_value": 0.9999,
    "checkpoint_every": 10000, "log_every": 100, "num_workers": 0,
    "validation_batch_size": 8, "validation_sampler": "ddim",
    "validation_inference_steps": 16, "validation_noise_seed": 7041,
    "no_real_robot": True, "test_only": False,
}


def canonical_config(config: dict) -> dict:
    cfg = copy.deepcopy(DEFAULTS)
    cfg.update(copy.deepcopy(config))
    required = {"history": 2, "horizon": 16, "image_size": 96,
                "crop_shape": None, "resize_shape": None, "random_crop": False, "imagenet_norm": True,
                "cameras": ["front", "gripper"], "context_dim": 2060,
                "execution_start": 0, "execute_steps": 8, "control_hz": 10,
                "no_real_robot": True, "num_train_timesteps": 100,
                "beta_schedule": "squaredcos_cap_v2", "prediction_type": "epsilon",
                "clip_sample": True, "variance_type": "fixed_small"}
    for key, value in required.items():
        if cfg[key] != value:
            raise ValueError(f"SO101 official-component contract requires {key}={value!r}")
    if not cfg["test_only"]:
        for key in ("down_dims", "diffusion_step_embed_dim", "kernel_size", "n_groups",
                    "cond_predict_scale"):
            if cfg[key] != DEFAULTS[key]:
                raise ValueError(f"Production architecture requires {key}={DEFAULTS[key]!r}")
    if len(cfg["down_dims"]) != 3 or any(int(x) <= 0 or int(x) % cfg["n_groups"] for x in cfg["down_dims"]):
        raise ValueError("Use three positive group-divisible U-Net channel widths")
    for key in ("batch_size", "microbatch_size", "total_updates", "checkpoint_every",
                "log_every", "validation_batch_size", "validation_inference_steps"):
        if isinstance(cfg[key], bool) or int(cfg[key]) != cfg[key] or cfg[key] <= 0:
            raise ValueError(f"{key} must be a positive integer")
    if cfg["batch_size"] % cfg["microbatch_size"]:
        raise ValueError("microbatch_size must divide effective batch_size")
    if cfg["num_workers"] != 0:
        raise ValueError("Exact-resume batch stream uses num_workers=0")
    if cfg["warmup_steps"] < 0 or cfg["learning_rate"] <= 0:
        raise ValueError("Invalid optimizer schedule")
    if cfg["gradient_clip"] is not None and cfg["gradient_clip"] <= 0:
        raise ValueError("gradient_clip must be positive or None")
    return cfg


def canonical_stats(stats: dict) -> dict:
    out = copy.deepcopy(stats)
    for key in STATS_KEYS:
        value = np.asarray(stats[key], dtype=np.float32)
        if value.shape != (6,) or not np.isfinite(value).all():
            raise ValueError(f"{key} requires six finite values")
        out[key] = value.tolist()
    for prefix in ("action", "state"):
        if np.any(np.asarray(out[prefix + "_max"]) < np.asarray(out[prefix + "_min"])):
            raise ValueError(f"Invalid {prefix} min/max")
    return out


def sha256(path: str | Path) -> str:
    h = hashlib.sha256()
    with Path(path).open("rb") as stream:
        for part in iter(lambda: stream.read(1024 * 1024), b""):
            h.update(part)
    return h.hexdigest()


def _official_imports():
    # A regular installed diffusion_policy package would otherwise take priority
    # over a namespace-only vendored snapshot. Do not silently mix the two.
    package_path = VENDOR / "diffusion_policy"
    existing = sys.modules.get("diffusion_policy")
    if existing is None:
        package = types.ModuleType("diffusion_policy")
        package.__path__ = [str(package_path)]
        package.__package__ = "diffusion_policy"
        sys.modules["diffusion_policy"] = package
    elif [Path(p).resolve() for p in getattr(existing, "__path__", [])] != [package_path.resolve()]:
        raise ImportError("Another diffusion_policy package is loaded; use an isolated process")
    from diffusion_policy.model.vision.multi_image_obs_encoder import MultiImageObsEncoder
    from diffusion_policy.model.vision.model_getter import get_resnet
    from diffusion_policy.model.diffusion.conditional_unet1d import ConditionalUnet1D
    from diffusion_policy.model.diffusion.ema_model import EMAModel
    for component in (MultiImageObsEncoder, get_resnet, ConditionalUnet1D, EMAModel):
        if not Path(inspect.getfile(component)).resolve().is_relative_to(VENDOR.resolve()):
            raise ImportError("An official component resolved outside the pinned vendor directory")
    return MultiImageObsEncoder, get_resnet, ConditionalUnet1D, EMAModel


def official_ema(model: nn.Module, cfg: dict):
    """Upstream EMA owns its input; deepcopy is REQUIRED, not done upstream."""
    *_, EMA = _official_imports()
    return EMA(copy.deepcopy(model), update_after_step=0, inv_gamma=1.0,
               power=float(cfg["ema_power"]), min_value=0.0,
               max_value=float(cfg["ema_max_value"]))


class OfficialDP(nn.Module):
    kind = "teacher"

    def __init__(self, cfg: dict, stats: dict):
        super().__init__()
        self.cfg = canonical_config(cfg)
        self.stats = canonical_stats(stats)
        for prefix in ("action", "state"):
            low = torch.tensor(self.stats[prefix + "_min"], dtype=torch.float32)
            high = torch.tensor(self.stats[prefix + "_max"], dtype=torch.float32)
            self.register_buffer(prefix + "_center", (high + low) / 2)
            self.register_buffer(prefix + "_span", (high - low).clamp_min(1e-3))
        Encoder, get_resnet, Unet, _ = _official_imports()
        shape_meta = {"obs": {
            "front": {"shape": [3, 96, 96], "type": "rgb"},
            "gripper": {"shape": [3, 96, 96], "type": "rgb"},
            "q": {"shape": [6], "type": "low_dim"}},
            "action": {"shape": [6]}}
        self.encoder = Encoder(
            shape_meta=shape_meta, rgb_model=get_resnet("resnet18", weights=None),
            resize_shape=None, crop_shape=None, random_crop=False,
            use_group_norm=True, share_rgb_model=False, imagenet_norm=True)
        self.denoiser = Unet(
            input_dim=6, local_cond_dim=None, global_cond_dim=2060,
            diffusion_step_embed_dim=self.cfg["diffusion_step_embed_dim"],
            down_dims=self.cfg["down_dims"], kernel_size=self.cfg["kernel_size"],
            n_groups=self.cfg["n_groups"], cond_predict_scale=self.cfg["cond_predict_scale"])
        from diffusers.schedulers.scheduling_ddpm import DDPMScheduler
        from diffusers.schedulers.scheduling_ddim import DDIMScheduler
        schedule = dict(num_train_timesteps=100, beta_start=0.0001, beta_end=0.02,
                        beta_schedule="squaredcos_cap_v2", clip_sample=True,
                        prediction_type="epsilon", timestep_spacing="leading",
                        steps_offset=0)
        self.noise_scheduler = DDPMScheduler(**schedule, variance_type="fixed_small")
        self.ddim_scheduler = DDIMScheduler(**schedule, set_alpha_to_one=True)

    def normalize_action(self, action: Tensor) -> Tensor:
        return 2 * (action - self.action_center) / self.action_span

    def normalize_state(self, q: Tensor) -> Tensor:
        return 2 * (q - self.state_center) / self.state_span

    def decode(self, normalized_action: Tensor) -> Tensor:
        """Linear inverse only; the evaluator logs both explicit clamp stages."""
        return normalized_action * self.action_span / 2 + self.action_center

    def encode(self, images: Tensor, q: Tensor) -> Tensor:
        if images.ndim != 6 or tuple(images.shape[1:]) != (2, 2, 3, 96, 96):
            raise ValueError("Expected [B,previous/current,front/gripper,RGB,96,96]")
        if images.dtype != torch.uint8 or tuple(q.shape) != (images.shape[0], 2, 6):
            raise ValueError("Images must be uint8 and q must be [B,2,6] environment coordinates")
        if not q.is_floating_point():
            raise ValueError("q must contain floating-point environment-normalized joint states")
        b = images.shape[0]
        rgb = images.float() / 255.0
        state = self.normalize_state(q)
        obs = {"front": rgb[:, :, 0].reshape(b * 2, 3, 96, 96),
               "gripper": rgb[:, :, 1].reshape(b * 2, 3, 96, 96),
               "q": state.reshape(b * 2, 6)}
        return self.encoder(obs).reshape(b, 2060)

    def head_predict(self, context: Tensor, noise: Tensor | None = None, *,
                     sampler: str = "ddpm", inference_steps: int = 100,
                     generator: torch.Generator | None = None) -> Tensor:
        """Return all 16 actions starting at current t, BEFORE deployment clamps.

        DDPM draws new Gaussian noise at every nonterminal reverse step. Supplying
        only initial noise does not make DDPM deterministic: pass a seeded generator
        as well. DDIM uses eta=0. Both use diffusers' official update equations.
        """
        if context.ndim != 2 or context.shape[1] != 2060:
            raise ValueError("Expected context [B,2060]")
        if sampler not in ("ddpm", "ddim") or not 1 <= int(inference_steps) <= 100:
            raise ValueError("Use ddpm/ddim with 1..100 inference steps")
        shape = (context.shape[0], 16, 6)
        if noise is None:
            sample = torch.randn(shape, dtype=context.dtype, device=context.device, generator=generator)
        else:
            if tuple(noise.shape) != shape or noise.device != context.device or noise.dtype != context.dtype:
                raise ValueError("Initial noise must match prediction shape/device/dtype")
            sample = noise.clone()
        scheduler = self.noise_scheduler if sampler == "ddpm" else self.ddim_scheduler
        scheduler.set_timesteps(int(inference_steps), device=context.device)
        finite = torch.isfinite(sample).all() & torch.isfinite(context).all()
        for t in scheduler.timesteps:
            epsilon = self.denoiser(sample, t, global_cond=context)
            finite = finite & torch.isfinite(epsilon).all()
            if sampler == "ddpm":
                sample = scheduler.step(epsilon, t, sample, generator=generator).prev_sample
            else:
                sample = scheduler.step(epsilon, t, sample, eta=0.0, generator=generator).prev_sample
            finite = finite & torch.isfinite(sample).all()
        if not bool(finite):
            raise FloatingPointError("Non-finite intermediate diffusion sample")
        return sample

    def predict(self, images: Tensor, q: Tensor, noise: Tensor | None = None, *,
                sampler: str = "ddpm", inference_steps: int = 100,
                generator: torch.Generator | None = None) -> Tensor:
        raw = self.head_predict(self.encode(images, q), noise, sampler=sampler,
                                inference_steps=inference_steps, generator=generator)
        return self.decode(raw.clamp(-1, 1)).clamp(-1, 1)

    def training_loss(self, batch: dict, *, generator: torch.Generator | None = None) -> Tensor:
        target = self.normalize_action(batch["action"])
        if target.shape[1:] != (16, 6) or batch["mask"].shape != target.shape[:2]:
            raise ValueError("Training target must be a masked [B,16,6] current-time chunk")
        noise = torch.randn(target.shape, device=target.device, dtype=target.dtype, generator=generator)
        timestep = torch.randint(0, 100, (target.shape[0],), device=target.device, generator=generator)
        noisy = self.noise_scheduler.add_noise(target, noise, timestep)
        prediction = self.denoiser(noisy, timestep, global_cond=self.encode(batch["images"], batch["q"]))
        mask = batch["mask"].unsqueeze(-1).expand_as(target)
        count = mask.sum()
        if count <= 0:
            raise ValueError("Empty action mask")
        return ((prediction - noise).square() * mask).sum() / count


def make_dataset(root: str | Path, cfg: dict, split: str):
    """Use the ORIGINAL manifest configuration; new optimizer keys are unrelated.

    No new demonstrations, boundary padding, statistics fit or test examples are
    introduced. Requested semantic fields must agree with the recorded dataset.
    """
    if str(REPO_ROOT) not in sys.path:
        sys.path.insert(0, str(REPO_ROOT))
    from experiments.so101_pbit_sim.data import SimDataset
    manifest = json.loads((Path(root) / "manifest.json").read_text(encoding="utf-8"))
    original = manifest["config"]
    shared = ("task_id", "cameras", "camera_layout", "image_size", "horizon", "history",
              "execution_start", "control_hz", "success_contract", "action", "no_real_robot",
              "train_seeds", "validation_seeds", "development_seeds")
    checked = canonical_config(cfg)
    for key in shared:
        if key in checked and checked[key] != original.get(key):
            raise ValueError(f"Original dataset semantic contract differs at {key}")
    if "test_seeds" in checked and checked["test_seeds"] != original["test_seeds"]:
        if checked.get("historical_test_seeds") != original["test_seeds"]:
            raise ValueError("Declare the complete historical test set before choosing a fresh holdout")
        used = set().union(*(set(original[k]) for k in
                             ("train_seeds", "validation_seeds", "test_seeds", "development_seeds")))
        if used.intersection(checked["test_seeds"]) or len(set(checked["test_seeds"])) != len(checked["test_seeds"]):
            raise ValueError("Fresh holdout overlaps a prior split or contains duplicate seeds")
    data = SimDataset(root, original, split)
    if data.stats.get("fit_seeds") != original["train_seeds"]:
        raise ValueError("Normalization was not fit exclusively on training episodes")
    return data


def source_record() -> dict:
    manifest_path = VENDOR / "source_manifest.json"
    manifest = json.loads(manifest_path.read_text(encoding="utf-8"))
    if manifest.get("commit") != UPSTREAM_COMMIT:
        raise ValueError("Unexpected upstream source revision")
    for row in manifest["files"]:
        path = VENDOR / row["path"]
        if sha256(path) != row["sha256"]:
            raise ValueError(f"Pinned vendor file changed: {row['path']}")
    return {"upstream_commit": UPSTREAM_COMMIT,
            "vendor_manifest_sha256": sha256(manifest_path),
            "git_commit_declared_by_launcher": os.environ.get("SOURCE_SHA"),
            "file_sha256": {p.relative_to(REPO_ROOT).as_posix(): sha256(p)
                            for p in (HERE / "model.py", HERE / "train.py") if p.is_file()},
            "versions": {name: importlib.metadata.version(name)
                         for name in ("torch", "torchvision", "diffusers", "einops")}}


def policy_payload(model: OfficialDP, *, step: int, manifest_sha: str,
                   extra: dict | None = None) -> dict:
    return {"format": FORMAT, "kind": "teacher", "config": model.cfg,
            "stats": model.stats, "state_dict": model.state_dict(), "step": int(step),
            "dataset_manifest_sha256": manifest_sha, "source": source_record(),
            "action_time_origin": "current_t", "execution_start": 0,
            "clamp_contract": "data_clamp_then_linear_decode_then_environment_clamp",
            "weights": "EMA", "extra": copy.deepcopy(extra or {})}


def atomic_torch_save(payload: dict, path: str | Path) -> dict:
    path = Path(path)
    path.parent.mkdir(parents=True, exist_ok=True)
    temporary = path.with_name(path.name + ".tmp")
    torch.save(payload, temporary)
    temporary.replace(path)
    return {"path": path.name, "sha256": sha256(path), "bytes": path.stat().st_size,
            "step": payload.get("step")}


def load_policy(path: str | Path, device: str | torch.device = "cpu") -> tuple[OfficialDP, dict]:
    ck = torch.load(Path(path), map_location="cpu", weights_only=True)
    if ck.get("format") != FORMAT or ck.get("action_time_origin") != "current_t":
        raise ValueError("Not an official-component SO101 EMA inference checkpoint")
    current_source = source_record()
    for key in ("upstream_commit", "vendor_manifest_sha256"):
        if ck.get("source", {}).get(key) != current_source[key]:
            raise ValueError(f"Checkpoint official-component identity differs: {key}")
    model = OfficialDP(ck["config"], ck["stats"])
    expected = {name: getattr(model, name).clone() for name in NORMALIZER_KEYS}
    model.load_state_dict(ck["state_dict"], strict=True)
    for name, value in expected.items():
        if not torch.equal(value, getattr(model, name)):
            raise ValueError(f"Checkpoint normalizer contradicts train statistics: {name}")
    for name, value in model.state_dict().items():
        if not torch.isfinite(value).all():
            raise ValueError(f"Non-finite checkpoint tensor: {name}")
    return model.to(device).eval(), ck