#!/usr/bin/env python3
"""One-time script to convert LPIPS linear weights to numpy format.

Run this once to generate the lpips_lins.npz file:
    uv run python scripts/convert_lpips_weights.py

The output file should be committed to src/lpips_nnx/data/
"""

import hashlib
from pathlib import Path

import numpy as np
import requests
import torch

LPIPS_WEIGHTS_URL = "https://heibox.uni-heidelberg.de/f/607503859c864bc1b30b/?dl=1"
LPIPS_WEIGHTS_MD5 = "d507d7349b931f0638a25a48a722f98a"


def download_file(url: str, dest: Path, expected_md5: str) -> Path:
    """Download a file with MD5 verification."""
    if dest.exists():
        with open(dest, "rb") as f:
            if hashlib.md5(f.read()).hexdigest() == expected_md5:
                return dest
        dest.unlink()

    print(f"Downloading {url}...")
    response = requests.get(url, stream=True)
    response.raise_for_status()

    dest.parent.mkdir(parents=True, exist_ok=True)
    with open(dest, "wb") as f:
        for chunk in response.iter_content(chunk_size=8192):
            f.write(chunk)

    with open(dest, "rb") as f:
        actual_md5 = hashlib.md5(f.read()).hexdigest()
    if actual_md5 != expected_md5:
        raise ValueError(f"MD5 mismatch: expected {expected_md5}, got {actual_md5}")

    return dest


def convert_conv_weight(pt_weight: np.ndarray) -> np.ndarray:
    """Convert PyTorch conv weight (OIHW) to JAX format (HWIO)."""
    return np.transpose(pt_weight, (2, 3, 1, 0))


def main():
    # Download original LPIPS weights
    cache_dir = Path.home() / ".cache" / "lpips_nnx"
    cache_dir.mkdir(parents=True, exist_ok=True)
    pth_path = download_file(LPIPS_WEIGHTS_URL, cache_dir / "vgg_lpips.pth", LPIPS_WEIGHTS_MD5)

    # Load and convert
    state_dict = torch.load(pth_path, map_location="cpu", weights_only=True)

    weights = {}
    for i in range(5):
        key = f"lin{i}.model.1.weight"
        pt_weight = state_dict[key].numpy()
        weights[f"lin{i}"] = convert_conv_weight(pt_weight)
        print(f"lin{i}: {pt_weight.shape} -> {weights[f'lin{i}'].shape}")

    # Save to package data directory
    output_dir = Path(__file__).parent.parent / "src" / "lpips_nnx" / "data"
    output_dir.mkdir(parents=True, exist_ok=True)
    output_path = output_dir / "lpips_lins.npz"

    np.savez(output_path, **weights)
    print(f"\nSaved to {output_path}")
    print(f"File size: {output_path.stat().st_size} bytes")


if __name__ == "__main__":
    main()
