#!/usr/bin/env python3
"""Generate an official PyTorch reference for the browser-controlled flow loop."""

from __future__ import annotations

import argparse
import json
import os
import time
from pathlib import Path

from dit_common import (
    CAMERA_SHAPE,
    FEATURE1_SHAPE,
    FEATURE2_SHAPE,
    LATENT_SHAPE,
    choose_torch_device,
    load_official_flow_model,
    resolved_file,
    sha256_file,
    source_revision,
    synchronize_torch,
)


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--triposplat-repo", type=Path, required=True)
    parser.add_argument("--weights", type=Path, required=True)
    parser.add_argument("--input-fixture-dir", type=Path, required=True)
    parser.add_argument("--output-fixture-dir", type=Path, required=True)
    parser.add_argument("--device", choices=("cpu", "mps", "cuda", "auto"), default="cpu")
    parser.add_argument(
        "--internal-precision",
        choices=("fp16", "fp32"),
        default="fp16",
        help="Official flow-model precision used for the reference (default: %(default)s).",
    )
    parser.add_argument("--steps", type=int, default=4)
    parser.add_argument("--guidance-scale", type=float, default=3.0)
    parser.add_argument("--shift", type=float, default=3.0)
    parser.add_argument(
        "--record-trajectory",
        action="store_true",
        help="Record official per-invocation sample, timestep, and raw predictions.",
    )
    args = parser.parse_args()
    if args.steps <= 0:
        parser.error("--steps must be positive")
    if args.guidance_scale <= 1:
        parser.error("--guidance-scale must be greater than one to exercise CFG")
    if args.shift <= 0:
        parser.error("--shift must be positive")
    return args


def hard_link(source: Path, destination: Path) -> None:
    destination.parent.mkdir(parents=True, exist_ok=True)
    if destination.exists():
        if os.path.samefile(source, destination):
            return
        destination.unlink()
    os.link(source, destination)


def main() -> None:
    try:
        import numpy as np
        import torch
    except ImportError as exc:
        raise SystemExit(f"PyTorch and NumPy are required: {exc}") from exc

    args = parse_args()
    repository = args.triposplat_repo.expanduser().resolve()
    weights = resolved_file(args.weights, "TripoSplat flow-model weights")
    input_dir = args.input_fixture_dir.expanduser().resolve()
    output_dir = args.output_fixture_dir.expanduser().resolve()
    shapes = {
        "latent": LATENT_SHAPE,
        "camera": CAMERA_SHAPE,
        "feature1": FEATURE1_SHAPE,
        "feature2": FEATURE2_SHAPE,
    }
    arrays = {}
    for name, shape in shapes.items():
        path = resolved_file(input_dir / f"{name}.f32", f"fixture {name}")
        array = np.fromfile(path, dtype="<f4")
        expected = int(np.prod(shape))
        if array.size != expected:
            raise ValueError(f"{path} has {array.size} floats; expected {expected}")
        arrays[name] = array.reshape(shape)

    device = choose_torch_device(torch, args.device)
    commit, dirty = source_revision(repository)
    print(f"Loading untouched official flow model at {commit} on {device}")
    model, source = load_official_flow_model(
        torch=torch,
        triposplat_repo=repository,
        weights=weights,
        device=device,
        internal_precision=args.internal_precision,
        low_memory_construction=True,
    )
    noise = {
        "latent": torch.from_numpy(arrays["latent"]).to(device=device, dtype=torch.float32),
        "camera": torch.from_numpy(arrays["camera"]).to(device=device, dtype=torch.float32),
    }
    condition = {
        "feature1": torch.from_numpy(arrays["feature1"]).to(device=device, dtype=torch.float32),
        "feature2": torch.from_numpy(arrays["feature2"]).to(device=device, dtype=torch.float32),
    }
    negative_condition = {name: torch.zeros_like(value) for name, value in condition.items()}
    sampler = source.pipeline_module.FlowEulerCfgSampler()
    trajectory = []

    class RecordingModel:
        def __call__(self, x_t, timestep, cond):
            prediction = model(x_t, timestep, cond)
            trajectory.append(
                {
                    "sample": {name: value.detach().clone() for name, value in x_t.items()},
                    "timestep": timestep.detach().clone(),
                    "prediction": {
                        name: value.detach().clone() for name, value in prediction.items()
                    },
                }
            )
            return prediction

    sampling_model = RecordingModel() if args.record_trajectory else model

    synchronize_torch(torch, device)
    started = time.perf_counter()
    with torch.inference_mode():
        result = sampler.sample(
            sampling_model,
            {name: value.clone() for name, value in noise.items()},
            cond=condition,
            neg_cond=negative_condition,
            steps=args.steps,
            guidance_scale=args.guidance_scale,
            shift=args.shift,
        )
    synchronize_torch(torch, device)
    duration_ms = (time.perf_counter() - started) * 1000

    output_dir.mkdir(parents=True, exist_ok=True)
    for name in shapes:
        hard_link(input_dir / f"{name}.f32", output_dir / f"{name}.f32")
    output_files = {}
    for name, tensor in result.items():
        path = output_dir / f"flow{args.steps}_{name}.f32"
        np.ascontiguousarray(tensor.detach().float().cpu().numpy(), dtype="<f4").tofile(path)
        output_files[name] = {
            "path": path.name,
            "shape": list(tensor.shape),
            "sha256": sha256_file(path),
        }

    trajectory_files = []
    if args.record_trajectory:
        expected_invocations = args.steps * 2
        if len(trajectory) != expected_invocations:
            raise RuntimeError(
                f"Recorded {len(trajectory)} model calls; expected {expected_invocations}"
            )
        trajectory_dir = output_dir / "trajectory"
        trajectory_dir.mkdir(parents=True, exist_ok=True)
        for index, record in enumerate(trajectory):
            invocation = index + 1
            prefix = f"invocation_{invocation:02d}"
            paths = {}
            tensors = {
                "sample_latent": record["sample"]["latent"],
                "sample_camera": record["sample"]["camera"],
                "t": record["timestep"],
                "pred_latent": record["prediction"]["latent"],
                "pred_camera": record["prediction"]["camera"],
            }
            for name, tensor in tensors.items():
                path = trajectory_dir / f"{prefix}_{name}.f32"
                np.ascontiguousarray(tensor.detach().float().cpu().numpy(), dtype="<f4").tofile(path)
                paths[name] = {
                    "path": str(path.relative_to(output_dir)),
                    "shape": list(tensor.shape),
                    "sha256": sha256_file(path),
                }
            trajectory_files.append(
                {
                    "invocation": invocation,
                    "step": index // 2 + 1,
                    "pass": "conditional" if index % 2 == 0 else "unconditional",
                    "timestep": float(record["timestep"].flatten()[0].detach().cpu()),
                    "tensors": paths,
                }
            )

    manifest = {
        "source": {
            "repository": str(repository),
            "commit": commit,
            "tracked_source_dirty": dirty,
            "loader": "official triposplat.load_flow_model",
            "sampler": "official triposplat.FlowEulerCfgSampler",
        },
        "settings": {
            "steps": args.steps,
            "guidance_scale": args.guidance_scale,
            "shift": args.shift,
            "internal_precision": args.internal_precision,
            "conditional_invocations": args.steps,
            "unconditional_invocations": args.steps,
        },
        "pytorch": {"device": str(device), "duration_ms": duration_ms},
        "outputs": output_files,
    }
    if trajectory_files:
        manifest["trajectory"] = trajectory_files
    manifest_path = output_dir / "flow.json"
    manifest_path.write_text(json.dumps(manifest, indent=2) + "\n", encoding="utf-8")
    print(f"Wrote {manifest_path} ({duration_ms:.1f} ms official PyTorch sampling)")


if __name__ == "__main__":
    main()