#!/usr/bin/env python3
"""Extract a low-memory ComfyUI LoRA from two MiniMax-H3 checkpoints.

The intended use is to approximate::

    ref2va = fl2va + lora

Only one source tensor is materialized at a time. Matrix deltas are compressed
with a deterministic randomized SVD; vectors and biases are stored as exact
ComfyUI ``.diff``/``.diff_b`` patches. Work is checkpointed per source tensor
so an interrupted extraction can be resumed by running the same command.
"""

from __future__ import annotations

import argparse
import hashlib
import json
import math
import os
import shutil
import struct
import time
from pathlib import Path
from typing import Any

import torch
from safetensors import safe_open
from safetensors.torch import save_file


HEADER_LIMIT = 100 * 1024 * 1024
COPY_BUFFER_SIZE = 16 * 1024 * 1024
DEFAULT_SAMPLES = (
    "condition_proj.weight",
    "blocks.0.adaln_proj.linear.weight",
    "blocks.0.attn.qkv_proj.weight",
    "blocks.0.mlp.fc1.weight",
    "blocks.0.mlp.fc2.weight",
    "blocks.24.attn.qkv_proj.weight",
    "blocks.49.attn.qkv_proj.weight",
    "token_refiner.blocks.0.attn.qkv_proj.weight",
)


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Stream FL2VA and REF2VA checkpoints into a ComfyUI LoRA."
    )
    parser.add_argument("--base", type=Path, required=True, help="FL2VA checkpoint")
    parser.add_argument("--target", type=Path, required=True, help="REF2VA checkpoint")
    parser.add_argument("--output", type=Path, help="Output .safetensors LoRA")
    parser.add_argument("--rank", type=int, default=64)
    parser.add_argument("--oversample", type=int, default=16)
    parser.add_argument("--power-iters", type=int, default=2)
    parser.add_argument("--device", default="cuda")
    parser.add_argument(
        "--factor-dtype", choices=("bf16", "fp16", "fp32"), default="bf16"
    )
    parser.add_argument(
        "--analyze-only",
        action="store_true",
        help="Measure representative delta spectra without writing a LoRA",
    )
    parser.add_argument(
        "--sample", action="append", default=[], help="Tensor key to analyze"
    )
    parser.add_argument("--analysis-rank", type=int, default=128)
    parser.add_argument(
        "--keep-parts", action="store_true", help="Keep resumable part files"
    )
    return parser.parse_args()


def tensor_info(handle, key: str) -> tuple[tuple[int, ...], str]:
    view = handle.get_slice(key)
    return tuple(view.get_shape()), str(view.get_dtype())


def validate_sources(base, target) -> list[str]:
    base_keys = set(base.keys())
    target_keys = set(target.keys())
    if base_keys != target_keys:
        only_base = sorted(base_keys - target_keys)[:5]
        only_target = sorted(target_keys - base_keys)[:5]
        raise ValueError(
            "Checkpoint keys differ: "
            f"only_base={only_base}, only_target={only_target}"
        )

    for key in sorted(base_keys):
        base_info = tensor_info(base, key)
        target_info = tensor_info(target, key)
        if base_info != target_info:
            raise ValueError(
                f"Tensor metadata differs for {key}: {base_info} != {target_info}"
            )
    return sorted(base_keys)


def load_delta(base, target, key: str, device: torch.device) -> torch.Tensor:
    """Load one target-base delta directly onto the compute device in FP32."""

    target_cpu = target.get_tensor(key)
    base_cpu = base.get_tensor(key)
    delta = target_cpu.to(device=device, dtype=torch.float32)
    delta.sub_(base_cpu.to(device=device, dtype=torch.float32))
    del target_cpu, base_cpu
    return delta


def randomized_svd(
    matrix: torch.Tensor,
    rank: int,
    oversample: int,
    power_iters: int,
    seed: int,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    """Deterministic randomized SVD suitable for very large dense deltas."""

    if matrix.ndim != 2:
        raise ValueError(f"Expected a matrix, got shape {tuple(matrix.shape)}")
    rows, cols = matrix.shape
    max_rank = min(rows, cols)
    requested_rank = min(rank, max_rank)
    sketch_rank = min(requested_rank + max(0, oversample), max_rank)

    generator = torch.Generator(device=matrix.device)
    generator.manual_seed(seed)
    omega = torch.randn(
        (cols, sketch_rank),
        generator=generator,
        device=matrix.device,
        dtype=torch.float32,
    )
    q, _ = torch.linalg.qr(matrix @ omega, mode="reduced")
    del omega

    for _ in range(max(0, power_iters)):
        z, _ = torch.linalg.qr(matrix.T @ q, mode="reduced")
        q, _ = torch.linalg.qr(matrix @ z, mode="reduced")
        del z

    small = q.T @ matrix
    u_small, singular_values, vh = torch.linalg.svd(small, full_matrices=False)
    u = q @ u_small
    return (
        u[:, :requested_rank],
        singular_values[:requested_rank],
        vh[:requested_rank, :],
    )


def key_seed(key: str) -> int:
    return int.from_bytes(hashlib.sha256(key.encode("utf-8")).digest()[:8], "little")


def factor_dtype(name: str) -> torch.dtype:
    return {
        "bf16": torch.bfloat16,
        "fp16": torch.float16,
        "fp32": torch.float32,
    }[name]


def lora_prefix(key: str) -> str:
    if not key.endswith(".weight"):
        raise ValueError(f"Matrix key is not a weight: {key}")
    return f"diffusion_model.{key[:-len('.weight')]}"


def exact_patch_key(key: str, all_keys: set[str]) -> str:
    if key.endswith(".bias") and f"{key[:-len('.bias')]}.weight" in all_keys:
        return f"diffusion_model.{key[:-len('.bias')]}.diff_b"
    if key.endswith(".weight"):
        return f"diffusion_model.{key[:-len('.weight')]}.diff"
    return f"diffusion_model.{key}.diff"


def decompose_delta(
    delta: torch.Tensor,
    key: str,
    rank: int,
    oversample: int,
    power_iters: int,
    output_dtype: torch.dtype,
) -> tuple[dict[str, torch.Tensor], dict[str, Any]]:
    total_energy = float(torch.sum(delta * delta).item())
    delta_norm = math.sqrt(total_energy)
    if total_energy == 0.0:
        return {}, {
            "key": key,
            "kind": "unchanged",
            "delta_norm": 0.0,
            "seconds": 0.0,
        }

    started = time.monotonic()
    effective_rank = min(rank, delta.shape[0], delta.shape[1])
    u, s, vh = randomized_svd(
        delta,
        rank=effective_rank,
        oversample=oversample,
        power_iters=power_iters,
        seed=key_seed(key),
    )
    captured_energy = float(torch.sum(s * s).item())
    capture = min(1.0, captured_energy / total_energy)
    sqrt_s = torch.sqrt(s)
    up = (u * sqrt_s.unsqueeze(0)).to(dtype=output_dtype, device="cpu")
    down = (sqrt_s.unsqueeze(1) * vh).to(dtype=output_dtype, device="cpu")
    prefix = lora_prefix(key)
    tensors = {
        f"{prefix}.lora_up.weight": up.contiguous(),
        f"{prefix}.lora_down.weight": down.contiguous(),
        f"{prefix}.alpha": torch.tensor(float(effective_rank), dtype=torch.float32),
    }
    stats = {
        "key": key,
        "kind": "lora",
        "shape": list(delta.shape),
        "rank": effective_rank,
        "delta_norm": delta_norm,
        "total_energy": total_energy,
        "captured_energy": captured_energy,
        "capture": capture,
        "relative_frobenius_error": math.sqrt(max(0.0, 1.0 - capture)),
        "seconds": time.monotonic() - started,
    }
    return tensors, stats


def analyze_samples(
    base,
    target,
    keys: list[str],
    device: torch.device,
    analysis_rank: int,
    oversample: int,
    power_iters: int,
) -> None:
    candidates = [key for key in (keys or list(DEFAULT_SAMPLES)) if key in set(base.keys())]
    if not candidates:
        raise ValueError("None of the requested sample keys exist")

    ranks = sorted({r for r in (8, 16, 32, 64, 128, analysis_rank) if r <= analysis_rank})
    print("key,shape,delta_norm," + ",".join(f"capture_r{r}" for r in ranks), flush=True)
    for key in candidates:
        shape, _ = tensor_info(base, key)
        if len(shape) != 2:
            print(f"Skipping non-matrix sample {key}: {shape}", flush=True)
            continue
        delta = load_delta(base, target, key, device)
        total_energy = float(torch.sum(delta * delta).item())
        max_rank = min(analysis_rank, *delta.shape)
        _, s, _ = randomized_svd(
            delta,
            rank=max_rank,
            oversample=oversample,
            power_iters=power_iters,
            seed=key_seed(key),
        )
        values = []
        for rank in ranks:
            actual = min(rank, len(s))
            values.append(float(torch.sum(s[:actual] ** 2).item()) / total_energy)
        print(
            f"{key},{'x'.join(map(str, shape))},{math.sqrt(total_energy):.8g},"
            + ",".join(f"{value:.8f}" for value in values),
            flush=True,
        )
        del delta, s
        if device.type == "cuda":
            torch.cuda.empty_cache()


def read_header(path: Path) -> tuple[dict[str, Any], int, int]:
    with path.open("rb") as handle:
        raw_length = handle.read(8)
        if len(raw_length) != 8:
            raise ValueError(f"Truncated safetensors length in {path}")
        header_length = struct.unpack("<Q", raw_length)[0]
        if header_length == 0 or header_length > HEADER_LIMIT:
            raise ValueError(f"Invalid safetensors header length in {path}")
        raw_header = handle.read(header_length)
        if len(raw_header) != header_length:
            raise ValueError(f"Truncated safetensors header in {path}")
    header = json.loads(raw_header)
    data_start = 8 + header_length
    data_length = path.stat().st_size - data_start
    return header, data_start, data_length


def merge_parts(parts: list[Path], output: Path, metadata: dict[str, str]) -> int:
    merged_header: dict[str, Any] = {"__metadata__": metadata}
    regions: list[tuple[Path, int]] = []
    offset = 0
    tensor_count = 0

    for part in parts:
        header, data_start, data_length = read_header(part)
        max_end = 0
        for key, descriptor in header.items():
            if key == "__metadata__":
                continue
            if key in merged_header:
                raise KeyError(f"Duplicate output tensor {key} in {part}")
            start, end = descriptor["data_offsets"]
            if not (0 <= start <= end <= data_length):
                raise ValueError(f"Bad offsets for {key} in {part}")
            merged_header[key] = {
                **descriptor,
                "data_offsets": [offset + start, offset + end],
            }
            max_end = max(max_end, end)
            tensor_count += 1
        if max_end != data_length:
            raise ValueError(f"Unmapped bytes in {part}: {max_end} != {data_length}")
        regions.append((part, data_start))
        offset += data_length

    encoded = json.dumps(merged_header, separators=(",", ":")).encode("utf-8")
    encoded += b" " * ((-len(encoded)) % 8)
    output.parent.mkdir(parents=True, exist_ok=True)
    temporary = output.with_name(f".{output.name}.tmp")
    if temporary.exists():
        temporary.unlink()
    with temporary.open("xb") as destination:
        destination.write(struct.pack("<Q", len(encoded)))
        destination.write(encoded)
        for part, data_start in regions:
            with part.open("rb") as source:
                source.seek(data_start)
                shutil.copyfileobj(source, destination, length=COPY_BUFFER_SIZE)
        destination.flush()
        os.fsync(destination.fileno())
    os.replace(temporary, output)
    return tensor_count


def write_json_atomic(path: Path, value: Any) -> None:
    temporary = path.with_name(f".{path.name}.tmp")
    temporary.write_text(json.dumps(value, indent=2, sort_keys=True) + "\n")
    os.replace(temporary, path)


def part_path(parts_dir: Path, index: int, key: str) -> Path:
    digest = hashlib.sha1(key.encode("utf-8")).hexdigest()[:12]
    return parts_dir / f"{index:04d}-{digest}.safetensors"


def run_extraction(args: argparse.Namespace, base, target, keys: list[str]) -> None:
    if args.output is None:
        raise ValueError("--output is required unless --analyze-only is used")
    if args.output.exists():
        raise FileExistsError(f"Refusing to overwrite {args.output}")
    if args.rank < 1:
        raise ValueError("--rank must be positive")

    output_dtype = factor_dtype(args.factor_dtype)
    device = torch.device(args.device)
    all_keys = set(keys)
    parts_dir = Path(f"{args.output}.parts")
    parts_dir.mkdir(parents=True, exist_ok=True)
    progress_path = parts_dir / "progress.json"
    progress = {"completed": {}, "settings": {}}
    if progress_path.exists():
        progress = json.loads(progress_path.read_text())

    settings = {
        "base": str(args.base.resolve()),
        "target": str(args.target.resolve()),
        "rank": args.rank,
        "oversample": args.oversample,
        "power_iters": args.power_iters,
        "factor_dtype": args.factor_dtype,
    }
    if progress.get("settings") and progress["settings"] != settings:
        raise ValueError(
            f"Existing parts use different settings: {progress['settings']} != {settings}"
        )
    progress["settings"] = settings
    completed: dict[str, Any] = progress.setdefault("completed", {})

    if device.type == "cuda":
        if not torch.cuda.is_available():
            raise RuntimeError("CUDA was requested but is unavailable")
        torch.cuda.reset_peak_memory_stats(device)

    started_all = time.monotonic()
    for index, key in enumerate(keys):
        expected_part = part_path(parts_dir, index, key)
        previous = completed.get(key)
        if previous is not None and (
            previous.get("kind") == "unchanged" or expected_part.exists()
        ):
            print(f"[{index + 1:03d}/{len(keys)}] resume {key}", flush=True)
            continue

        shape, _ = tensor_info(base, key)
        delta = load_delta(base, target, key, device)
        started = time.monotonic()
        if len(shape) >= 2 and key.endswith(".weight"):
            tensors, stats = decompose_delta(
                delta,
                key=key,
                rank=args.rank,
                oversample=args.oversample,
                power_iters=args.power_iters,
                output_dtype=output_dtype,
            )
        else:
            delta_norm = float(torch.linalg.vector_norm(delta).item())
            if delta_norm == 0.0:
                tensors = {}
                stats = {"key": key, "kind": "unchanged", "delta_norm": 0.0}
            else:
                patch_key = exact_patch_key(key, all_keys)
                tensors = {patch_key: delta.to(dtype=torch.float32, device="cpu").contiguous()}
                stats = {
                    "key": key,
                    "kind": "exact",
                    "shape": list(shape),
                    "delta_norm": delta_norm,
                    "output_key": patch_key,
                }
            stats["seconds"] = time.monotonic() - started

        if tensors:
            temporary_part = expected_part.with_name(f".{expected_part.name}.tmp")
            if temporary_part.exists():
                temporary_part.unlink()
            save_file(
                tensors,
                str(temporary_part),
                metadata={"format": "pt", "source_key": key},
            )
            os.replace(temporary_part, expected_part)
            stats["part"] = expected_part.name
            stats["output_bytes"] = expected_part.stat().st_size

        completed[key] = stats
        write_json_atomic(progress_path, progress)
        detail = stats["kind"]
        if detail == "lora":
            detail += (
                f" r={stats['rank']} capture={stats['capture']:.4%} "
                f"err={stats['relative_frobenius_error']:.4%}"
            )
        print(
            f"[{index + 1:03d}/{len(keys)}] {key} {shape}: {detail} "
            f"({stats.get('seconds', 0.0):.2f}s)",
            flush=True,
        )
        del delta, tensors
        if device.type == "cuda":
            torch.cuda.empty_cache()

    stats_list = [completed[key] for key in keys]
    matrix_stats = [item for item in stats_list if item["kind"] == "lora"]
    exact_stats = [item for item in stats_list if item["kind"] == "exact"]
    unchanged_stats = [item for item in stats_list if item["kind"] == "unchanged"]
    total_energy = sum(item["total_energy"] for item in matrix_stats)
    captured_energy = sum(item["captured_energy"] for item in matrix_stats)
    weighted_capture = captured_energy / total_energy if total_energy else 1.0
    ranks = [item["rank"] for item in matrix_stats]

    summary = {
        "source_tensors": len(keys),
        "lora_matrices": len(matrix_stats),
        "exact_patches": len(exact_stats),
        "unchanged_tensors": len(unchanged_stats),
        "rank_requested": args.rank,
        "rank_min": min(ranks) if ranks else 0,
        "rank_max": max(ranks) if ranks else 0,
        "rank_mean": sum(ranks) / len(ranks) if ranks else 0.0,
        "matrix_delta_energy_capture": weighted_capture,
        "matrix_delta_relative_frobenius_error": math.sqrt(
            max(0.0, 1.0 - weighted_capture)
        ),
        "elapsed_seconds": time.monotonic() - started_all,
        "peak_cuda_bytes": (
            torch.cuda.max_memory_allocated(device) if device.type == "cuda" else 0
        ),
    }
    summary_path = args.output.with_suffix(args.output.suffix + ".json")
    write_json_atomic(summary_path, {"summary": summary, "layers": stats_list})

    metadata = {
        "format": "pt",
        "modelspec.architecture": "minimax_h3",
        "modelspec.title": "MiniMax-H3 FL2VA to REF2VA extracted LoRA",
        "modelspec.description": (
            "Approximate REF2VA-FL2VA matrix deltas with randomized SVD; "
            "vectors and biases use exact ComfyUI diff patches."
        ),
        "ss_network_module": "networks.lora",
        "ss_network_dim": str(args.rank),
        "ss_network_alpha": str(args.rank),
        "source_base": args.base.name,
        "source_target": args.target.name,
        "rank": str(args.rank),
        "oversample": str(args.oversample),
        "power_iters": str(args.power_iters),
        "factor_dtype": args.factor_dtype,
        "matrix_delta_energy_capture": f"{weighted_capture:.10f}",
    }
    parts = sorted(parts_dir.glob("[0-9][0-9][0-9][0-9]-*.safetensors"))
    tensor_count = merge_parts(parts, args.output, metadata)

    with safe_open(args.output, framework="pt", device="cpu") as final:
        final_keys = list(final.keys())
        if len(final_keys) != tensor_count:
            raise ValueError(
                f"Final tensor count mismatch: {len(final_keys)} != {tensor_count}"
            )

    summary["output_tensors"] = tensor_count
    summary["output_bytes"] = args.output.stat().st_size
    write_json_atomic(summary_path, {"summary": summary, "layers": stats_list})
    print(json.dumps(summary, indent=2), flush=True)

    if not args.keep_parts:
        for part in parts:
            part.unlink()
        progress_path.unlink(missing_ok=True)
        parts_dir.rmdir()


def main() -> None:
    args = parse_args()
    if not args.base.is_file() or not args.target.is_file():
        raise FileNotFoundError("Both --base and --target must exist")

    device = torch.device(args.device)
    with safe_open(args.base, framework="pt", device="cpu") as base, safe_open(
        args.target, framework="pt", device="cpu"
    ) as target:
        keys = validate_sources(base, target)
        print(f"Validated {len(keys)} matching source tensors", flush=True)
        if args.analyze_only:
            analyze_samples(
                base,
                target,
                keys=args.sample,
                device=device,
                analysis_rank=args.analysis_rank,
                oversample=args.oversample,
                power_iters=args.power_iters,
            )
        else:
            run_extraction(args, base, target, keys)


if __name__ == "__main__":
    main()