"""Weight download and conversion utilities for LPIPS."""

from pathlib import Path

import jax.numpy as jnp
import numpy as np
import requests
from safetensors import safe_open

# VGG16 weights from timm (torchvision ImageNet weights)
VGG16_SAFETENSORS_URL = "https://huggingface.co/timm/vgg16.tv_in1k/resolve/main/model.safetensors"

# Cache directory
CACHE_DIR = Path.home() / ".cache" / "lpips_nnx"

# Bundled LPIPS linear weights
LPIPS_WEIGHTS_PATH = Path(__file__).parent / "data" / "lpips_lins.npz"


def get_cache_dir() -> Path:
    """Get or create the cache directory."""
    CACHE_DIR.mkdir(parents=True, exist_ok=True)
    return CACHE_DIR


def convert_conv_weight(pt_weight: np.ndarray) -> jnp.ndarray:
    """Convert PyTorch conv weight to JAX format.

    PyTorch: [out_ch, in_ch, kH, kW] (OIHW)
    JAX:     [kH, kW, in_ch, out_ch] (HWIO)
    """
    return jnp.array(np.transpose(pt_weight, (2, 3, 1, 0)))


def download_file(url: str, dest: Path) -> Path:
    """Download a file if it doesn't exist."""
    if dest.exists():
        return dest

    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)

    return dest


def load_lpips_weights() -> dict[str, jnp.ndarray]:
    """Load LPIPS linear layer weights from bundled npz file.

    Returns a dict with keys 'lin0' through 'lin4', each containing
    the 1x1 conv weights in HWIO format.
    """
    data = np.load(LPIPS_WEIGHTS_PATH)
    return {f"lin{i}": jnp.array(data[f"lin{i}"]) for i in range(5)}


def load_vgg16_weights() -> dict[str, dict[str, jnp.ndarray]]:
    """Load VGG16 pretrained weights from HuggingFace.

    Downloads the safetensors file from timm/vgg16.tv_in1k and converts
    weights to JAX format.

    Returns a dict mapping layer names to {'kernel': ..., 'bias': ...}.
    """
    cache_path = get_cache_dir() / "vgg16.safetensors"
    download_file(VGG16_SAFETENSORS_URL, cache_path)

    # VGG16 feature layer indices for conv layers
    conv_indices = [0, 2, 5, 7, 10, 12, 14, 17, 19, 21, 24, 26, 28]

    weights = {}
    with safe_open(str(cache_path), framework="numpy") as f:
        for idx in conv_indices:
            name = f"conv_{idx}"
            pt_weight = f.get_tensor(f"features.{idx}.weight")
            pt_bias = f.get_tensor(f"features.{idx}.bias")
            weights[name] = {
                "kernel": convert_conv_weight(pt_weight),
                "bias": jnp.array(pt_bias),
            }

    return weights
