#!/usr/bin/env python3
"""Capture safe Kimodo PyTorch fixtures for GGML layer-parity tests.

The script is run inside the official upstream environment.  It keeps upstream
code read-only and writes only an NPZ/JSON pair.  The generated fixture covers
unconstrained single-prompt inference; constraints and post-processing are
captured separately so their behaviour cannot mask a denoiser error.
"""

from __future__ import annotations

import argparse
import hashlib
import json
import platform
import sys
import gc
from pathlib import Path
from typing import Any

import numpy as np
import torch


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--upstream", required=True, type=Path, help="Kimodo checkout")
    parser.add_argument("--model", default="kimodo-smplx-rp")
    parser.add_argument("--prompt", required=True)
    parser.add_argument("--frames", required=True, type=int)
    parser.add_argument("--steps", default=100, type=int)
    parser.add_argument("--seed", required=True, type=int)
    parser.add_argument("--device", default="cuda")
    parser.add_argument("--checkpoint-dir", type=Path,
                        help="local directory containing Kimodo-SMPLX-RP-v1; never download at capture time")
    parser.add_argument("--zero-embedding", action="store_true",
                        help="use a deterministic [1,1,4096] zero embedding; enables motion-only fixtures")
    parser.add_argument("--text-base", type=Path,
                        help="local Llama-3 base model; together with both adapters, encodes before loading diffusion")
    parser.add_argument("--text-mntp-adapter", type=Path,
                        help="local LLM2Vec MNTP adapter")
    parser.add_argument("--text-supervised-adapter", type=Path,
                        help="local LLM2Vec supervised adapter")
    parser.add_argument("--output", required=True, type=Path)
    return parser.parse_args()


def as_numpy(value: Any) -> np.ndarray:
    if not isinstance(value, torch.Tensor):
        raise TypeError(f"expected tensor, got {type(value)!r}")
    return value.detach().to(device="cpu", dtype=torch.float32).contiguous().numpy()


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


class CachedEmbedding:
    """A one-prompt text encoder that keeps the diffusion capture GPU-only.

    The real LLM2Vec model is deliberately released before Kimodo is loaded.
    This avoids an otherwise unnecessary combined 8B-LLM + diffusion VRAM peak
    while retaining the exact F32 embedding used by the upstream denoiser.
    """

    def __init__(self, embedding: torch.Tensor, prompt: str):
        self.embedding = embedding.detach().to(device="cpu", dtype=torch.float32).contiguous()
        self.prompt = prompt

    def __call__(self, texts: list[str] | str):
        values = [texts] if isinstance(texts, str) else texts
        if values != [self.prompt]:
            raise RuntimeError("capture cache only supports its recorded single prompt")
        return self.embedding.clone(), [1]


def capture_real_embedding(args: argparse.Namespace) -> CachedEmbedding:
    paths = [args.text_base, args.text_mntp_adapter, args.text_supervised_adapter]
    if any(path is None for path in paths):
        raise SystemExit("real text capture requires --text-base, --text-mntp-adapter, and --text-supervised-adapter")
    if any(not path.is_dir() for path in paths):
        raise SystemExit("each local text model path must be a directory")

    from kimodo.model.llm2vec import LLM2Vec  # pylint: disable=import-outside-toplevel
    from peft import PeftModel  # pylint: disable=import-outside-toplevel

    # LLM2Vec's upstream preset is base Llama + MNTP LoRA + supervised LoRA.
    # Merge only the MNTP stage; the second adapter remains active exactly as
    # it is in the upstream Kimodo LLM2VecEncoder.
    encoder = LLM2Vec.from_pretrained(
        str(args.text_base), peft_model_name_or_path=str(args.text_mntp_adapter),
        merge_peft=True, torch_dtype=torch.bfloat16,
    )
    encoder.model = PeftModel.from_pretrained(encoder.model, str(args.text_supervised_adapter))
    embedding = encoder.encode([args.prompt], batch_size=1, show_progress_bar=False, device=args.device)
    if tuple(embedding.shape) != (1, 4096):
        raise RuntimeError(f"LLM2Vec returned {tuple(embedding.shape)}, expected (1, 4096)")
    return CachedEmbedding(embedding[:, None, :], args.prompt)


def main() -> None:
    args = parse_args()
    upstream = args.upstream.resolve()
    if not (upstream / "kimodo").is_dir():
        raise SystemExit(f"not a Kimodo checkout: {upstream}")
    if args.frames <= 0 or args.steps <= 0:
        raise SystemExit("--frames and --steps must be positive")
    if args.zero_embedding and any((args.text_base, args.text_mntp_adapter, args.text_supervised_adapter)):
        raise SystemExit("--zero-embedding cannot be combined with real text model paths")
    if args.checkpoint_dir:
        checkpoint = args.checkpoint_dir.resolve()
        if not (checkpoint / "Kimodo-SMPLX-RP-v1" / "config.yaml").is_file():
            raise SystemExit("--checkpoint-dir must contain Kimodo-SMPLX-RP-v1/config.yaml")
        # This is deliberately set only for the reference subprocess.  It
        # prevents a missing local model from silently falling back to HF.
        import os
        os.environ["CHECKPOINT_DIR"] = str(checkpoint)

    sys.path.insert(0, str(upstream))
    from kimodo import load_model  # pylint: disable=import-outside-toplevel

    torch.manual_seed(args.seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed_all(args.seed)

    text_encoder = None
    if args.zero_embedding:
        class ZeroEmbedding:
            def __call__(self, text: list[str] | str):
                batch = len(text) if isinstance(text, list) else 1
                return torch.zeros((batch, 1, 4096), dtype=torch.float32), [1] * batch
        text_encoder = ZeroEmbedding()
    elif any((args.text_base, args.text_mntp_adapter, args.text_supervised_adapter)):
        text_encoder = capture_real_embedding(args)
        # The 8B encoder owns the GPU while encoding.  The cached output is
        # CPU F32, so release all allocator-owned memory before diffusion.
        gc.collect()
        if torch.cuda.is_available():
            torch.cuda.empty_cache()
    model, resolved_name = load_model(
        args.model, device=args.device, return_resolved_name=True, text_encoder=text_encoder
    )
    root_calls: list[tuple[tuple[torch.Tensor, ...], torch.Tensor]] = []
    body_calls: list[tuple[tuple[torch.Tensor, ...], torch.Tensor]] = []
    root_layer0: list[torch.Tensor] = []
    root_layers: list[list[torch.Tensor]] = [[] for _ in range(16)]
    root_layer0_inputs: list[torch.Tensor] = []
    root_attention0: list[torch.Tensor] = []
    root_norm10: list[torch.Tensor] = []
    root_motion_projection: list[torch.Tensor] = []
    root_text_projection: list[torch.Tensor] = []
    root_timestep_projection: list[torch.Tensor] = []
    root_heading_projection: list[torch.Tensor] = []
    sampling_inputs: list[torch.Tensor] = []
    sampling_outputs: list[torch.Tensor] = []

    def capture(calls: list[tuple[tuple[torch.Tensor, ...], torch.Tensor]]):
        def hook(_module: torch.nn.Module, inputs: tuple[Any, ...], output: torch.Tensor) -> None:
            tensors = tuple(item for item in inputs if isinstance(item, torch.Tensor))
            calls.append((tensors, output))
        return hook

    root_hook = model.denoiser.model.root_model.register_forward_hook(capture(root_calls))
    body_hook = model.denoiser.model.body_model.register_forward_hook(capture(body_calls))
    root_layer_hook = model.denoiser.model.root_model.seqTransEncoder.layers[0].register_forward_hook(
        lambda _module, _inputs, output: root_layer0.append(output)
    )
    root_layer_hooks = [layer.register_forward_hook(
        lambda _module, _inputs, output, index=index: root_layers[index].append(output)
    ) for index, layer in enumerate(model.denoiser.model.root_model.seqTransEncoder.layers)]
    root_layer_input_hook = model.denoiser.model.root_model.seqTransEncoder.layers[0].register_forward_pre_hook(
        lambda _module, inputs: root_layer0_inputs.append(inputs[0])
    )
    root_attention_hook = model.denoiser.model.root_model.seqTransEncoder.layers[0].self_attn.register_forward_hook(
        lambda _module, _inputs, output: root_attention0.append(output[0])
    )
    root_norm1_hook = model.denoiser.model.root_model.seqTransEncoder.layers[0].norm1.register_forward_hook(
        lambda _module, _inputs, output: root_norm10.append(output)
    )
    root_motion_hook = model.denoiser.model.root_model.input_linear.register_forward_hook(
        lambda _module, _inputs, output: root_motion_projection.append(output)
    )
    root_text_hook = model.denoiser.model.root_model.embed_text.register_forward_hook(
        lambda _module, _inputs, output: root_text_projection.append(output)
    )
    root_timestep_hook = model.denoiser.model.root_model.embed_timestep.register_forward_hook(
        lambda _module, _inputs, output: root_timestep_projection.append(output)
    )
    root_heading_hook = model.denoiser.model.root_model.linear_first_heading_angle.register_forward_hook(
        lambda _module, _inputs, output: root_heading_projection.append(output)
    )
    original_denoising_step = model.denoising_step

    def capture_denoising_step(*args: Any, **kwargs: Any) -> torch.Tensor:
        sampling_inputs.append(args[0].detach().clone())
        value = original_denoising_step(*args, **kwargs)
        sampling_outputs.append(value.detach().clone())
        return value

    model.denoising_step = capture_denoising_step
    try:
        text_features, text_lengths = model.text_encoder([args.prompt])
        output = model(
            args.prompt,
            num_frames=args.frames,
            num_denoising_steps=args.steps,
            num_samples=1,
            cfg_weight=[2.0, 2.0],
            cfg_type="separated",
            post_processing=False,
            return_numpy=True,
            progress_bar=lambda values: values,
        )
    finally:
        root_hook.remove()
        body_hook.remove()
        root_layer_hook.remove()
        for hook in root_layer_hooks:
            hook.remove()
        root_layer_input_hook.remove()
        root_attention_hook.remove()
        root_norm1_hook.remove()
        root_motion_hook.remove()
        root_text_hook.remove()
        root_timestep_hook.remove()
        root_heading_hook.remove()
        model.denoising_step = original_denoising_step

    if not root_calls or not body_calls or not root_layer0 or any(not values for values in root_layers) or not root_layer0_inputs or not root_attention0 or not root_norm10 or not root_motion_projection or not root_text_projection or not root_timestep_projection or not root_heading_projection:
        raise RuntimeError("the denoiser hooks did not observe inference")

    # CFG may invoke a transformer more than once per diffusion step.  Capturing
    # call zero is a stable, fully specified first boundary; later calls remain
    # reproducible from the final output and are added as C++ reaches CFG.
    root_inputs, root_output = root_calls[0]
    body_inputs, body_output = body_calls[0]
    # This is the exact stage boundary used to construct body_input_0: it is
    # deliberately captured separately so C++ can validate the representation
    # conversion independently of either Transformer graph.
    root_lengths = root_inputs[1].sum(-1)
    root_local = model.denoiser.model.motion_rep.global_root_to_local_root(
        root_output, normalized=True, lengths=root_lengths
    )
    if not sampling_inputs or len(sampling_inputs) != len(sampling_outputs):
        raise RuntimeError("the sampler trajectory hook did not observe inference")
    tensors: dict[str, np.ndarray] = {
        "text_features": as_numpy(text_features),
        "text_lengths": np.asarray(text_lengths, dtype=np.int64),
        "root_output": as_numpy(root_output),
        "root_local": as_numpy(root_local),
        "body_output": as_numpy(body_output),
        "root_layer0_output": as_numpy(root_layer0[0]),
        "root_layer0_input": as_numpy(root_layer0_inputs[0]),
        "root_attention0_output": as_numpy(root_attention0[0]),
        "root_norm10_output": as_numpy(root_norm10[0]),
        "root_motion_projection": as_numpy(root_motion_projection[0]),
        "root_text_projection": as_numpy(root_text_projection[0]),
        "root_timestep_projection": as_numpy(root_timestep_projection[0]),
        "root_heading_projection": as_numpy(root_heading_projection[0]),
        "sampling_initial_noise": as_numpy(sampling_inputs[0]),
        "sampling_final_state": as_numpy(sampling_outputs[-1]),
    }
    for index, values in enumerate(root_layers):
        tensors[f"root_layer{index}_output"] = as_numpy(values[0])
    for prefix, inputs in (("root", root_inputs), ("body", body_inputs)):
        for index, value in enumerate(inputs):
            tensors[f"{prefix}_input_{index}"] = as_numpy(value)
    for key, value in output.items():
        if isinstance(value, np.ndarray):
            tensors[f"motion_{key}"] = value
    for index, (sample_in, sample_out) in enumerate(zip(sampling_inputs, sampling_outputs)):
        tensors[f"sampling_input_{index}"] = as_numpy(sample_in)
        tensors[f"sampling_output_{index}"] = as_numpy(sample_out)

    args.output.parent.mkdir(parents=True, exist_ok=True)
    np.savez_compressed(args.output.with_suffix(".npz"), **tensors)
    metadata = {
        "fixture_format": 1,
        "upstream": str(upstream),
        "resolved_model": resolved_name,
        "prompt": args.prompt,
        "frames": args.frames,
        "diffusion_steps": args.steps,
        "seed": args.seed,
        "cfg_type": "separated",
        "cfg_weight": [2.0, 2.0],
        "post_processing": False,
        "device": args.device,
        "checkpoint_dir": str(args.checkpoint_dir) if args.checkpoint_dir else None,
        "zero_embedding": args.zero_embedding,
        "torch": torch.__version__,
        "cuda": torch.version.cuda,
        "python": platform.python_version(),
        "root_call_count": len(root_calls),
        "body_call_count": len(body_calls),
        "sampling_step_count": len(sampling_inputs),
        "npz_sha256": sha256(args.output.with_suffix(".npz")),
    }
    args.output.with_suffix(".json").write_text(json.dumps(metadata, indent=2) + "\n", encoding="utf-8")


if __name__ == "__main__":
    main()