#!/usr/bin/env python3
from __future__ import annotations

import json
import sys
from dataclasses import dataclass
from pathlib import Path
from typing import Literal

import numpy as np
import torch
import typer
from datasets import load_dataset
from rich.progress import track
from torch.utils.data import DataLoader

PROJECT_ROOT = Path(__file__).resolve().parents[1]
MINI_NAV_ROOT = PROJECT_ROOT / "mini-nav"
if str(MINI_NAV_ROOT) not in sys.path:
    sys.path.insert(0, str(MINI_NAV_ROOT))

from compressors.model_loader import get_dino_dim, load_dino_model, load_hash_compressor  # noqa: E402
from utils import get_device  # noqa: E402

DEFAULT_COMPRESSOR_PATH = Path("outputs/hash_compressor.pt")
DEFAULT_OUTPUT_ROOT = Path("outputs/cam_retrieval_benchmark/datasets")
DEFAULT_MAX_QUERIES = 8192


@dataclass(frozen=True)
class PreparedArtifact:
    npz_path: Path
    metadata_path: Path
    num_rows: int
    num_queries: int


def dataset_config(dataset: str) -> tuple[str, str]:
    if dataset == "cifar10":
        return "uoft-cs/cifar10", "label"
    if dataset == "cifar100":
        return "uoft-cs/cifar100", "fine_label"
    raise ValueError("dataset must be cifar10 or cifar100")


def stratified_indices(labels: list[int], *, total: int, seed: int) -> list[int]:
    if total <= 0:
        raise ValueError("total must be greater than 0")
    if total > len(labels):
        raise ValueError("total cannot exceed number of labels")

    rng = np.random.default_rng(seed)
    by_label: dict[int, list[int]] = {}
    for idx, label in enumerate(labels):
        by_label.setdefault(int(label), []).append(idx)

    label_ids = sorted(by_label)
    base = total // len(label_ids)
    remainder = total % len(label_ids)
    selected: list[int] = []

    for offset, label in enumerate(label_ids):
        want = base + (1 if offset < remainder else 0)
        candidates = list(by_label[label])
        if want > len(candidates):
            raise ValueError(f"not enough examples for label {label}: need {want}, have {len(candidates)}")
        chosen = rng.choice(candidates, size=want, replace=False).tolist()
        selected.extend(int(i) for i in chosen)

    selected.sort()
    return selected


def pack_bits_to_words_le(bits: np.ndarray, *, hash_bits: int) -> np.ndarray:
    if bits.ndim != 2:
        raise ValueError("bits must have shape [N, hash_bits]")
    if bits.shape[1] != hash_bits:
        raise ValueError(f"expected {hash_bits} bits, got {bits.shape[1]}")
    if hash_bits % 64 != 0:
        raise ValueError("hash_bits must be divisible by 64")

    words = np.zeros((bits.shape[0], hash_bits // 64), dtype=np.uint64)
    bits_u8 = bits.astype(np.uint8, copy=False)
    for word_idx in range(hash_bits // 64):
        word = np.zeros(bits.shape[0], dtype=np.uint64)
        for bit in range(64):
            word |= bits_u8[:, word_idx * 64 + bit].astype(np.uint64) << np.uint64(bit)
        words[:, word_idx] = word
    return words


def words_le_to_int(words: np.ndarray) -> int:
    value = 0
    for idx, word in enumerate(words.tolist()):
        value |= int(word) << (64 * idx)
    return value


@torch.no_grad()
def encode_dataset_bits(
    dataset,
    *,
    img_column: str,
    batch_size: int,
    dino_model: str,
    compressor_path: Path,
    hash_bits: int,
) -> np.ndarray:
    if not compressor_path.exists():
        raise FileNotFoundError(f"missing HashCompressor weights: {compressor_path}")

    device = get_device()
    processor, dino = load_dino_model(dino_model)
    compressor = load_hash_compressor(
        input_dim=get_dino_dim(dino_model),
        hash_bits=hash_bits,
        compressor_path=str(compressor_path),
    ).to(device)
    compressor.eval()

    loader = DataLoader(dataset.with_format("torch"), batch_size=batch_size, shuffle=False, num_workers=4)
    all_bits: list[np.ndarray] = []
    for batch in track(loader, description="Encoding CIFAR hash bits"):
        images = batch[img_column]
        inputs = processor(images=images, return_tensors="pt").to(device)
        tokens = dino(**inputs).last_hidden_state
        bits = compressor.encode(tokens).detach().cpu().numpy().astype(np.uint8)
        all_bits.append(bits)
    return np.concatenate(all_bits, axis=0)


def prepare_artifact(
    *,
    dataset_name: Literal["cifar10", "cifar100"],
    num_rows: int,
    max_queries: int,
    compressor_path: Path,
    output_root: Path,
    dino_model: str = "facebook/dinov2-large",
    hash_bits: int = 512,
    batch_size: int = 64,
    seed: int = 20260522,
) -> PreparedArtifact:
    hf_id, label_column = dataset_config(dataset_name)
    loaded = load_dataset(hf_id)
    train = loaded["train"]
    test = loaded["test"] if "test" in loaded else loaded["validation"]

    row_indices = stratified_indices([int(x) for x in train[label_column]], total=num_rows, seed=seed)
    query_indices = stratified_indices([int(x) for x in test[label_column]], total=max_queries, seed=seed + 1)

    row_subset = train.select(row_indices)
    query_subset = test.select(query_indices)

    row_bits = encode_dataset_bits(
        row_subset,
        img_column="img",
        batch_size=batch_size,
        dino_model=dino_model,
        compressor_path=compressor_path,
        hash_bits=hash_bits,
    )
    query_bits = encode_dataset_bits(
        query_subset,
        img_column="img",
        batch_size=batch_size,
        dino_model=dino_model,
        compressor_path=compressor_path,
        hash_bits=hash_bits,
    )

    rows_words = pack_bits_to_words_le(row_bits, hash_bits=hash_bits)
    queries_words = pack_bits_to_words_le(query_bits, hash_bits=hash_bits)
    row_labels = np.array([int(x) for x in row_subset[label_column]], dtype=np.int64)
    query_labels = np.array([int(x) for x in query_subset[label_column]], dtype=np.int64)

    output_root.mkdir(parents=True, exist_ok=True)
    stem = f"{dataset_name}_hash{hash_bits}_rows{num_rows}_queries{max_queries}"
    npz_path = output_root / f"{stem}.npz"
    metadata_path = output_root / f"{stem}.json"

    np.savez_compressed(
        npz_path,
        rows_words=rows_words,
        row_labels=row_labels,
        queries_words=queries_words,
        query_labels=query_labels,
    )
    metadata = {
        "dataset": dataset_name,
        "hf_id": hf_id,
        "label_column": label_column,
        "compressor_path": str(compressor_path),
        "dino_model": dino_model,
        "hash_bits": hash_bits,
        "bit_source": "HashCompressor.encode(tokens)",
        "bit_values": "{0,1} int32",
        "num_rows": int(num_rows),
        "num_queries": int(max_queries),
        "rows_source_split": "train",
        "queries_source_split": "test",
        "seed": int(seed),
    }
    metadata_path.write_text(json.dumps(metadata, indent=2, sort_keys=True) + "\n", encoding="utf-8")
    return PreparedArtifact(npz_path=npz_path, metadata_path=metadata_path, num_rows=num_rows, num_queries=max_queries)


def main(
    dataset: Literal["cifar10", "cifar100"] = typer.Option("cifar100"),
    num_rows: int = typer.Option(512, min=5),
    max_queries: int = typer.Option(DEFAULT_MAX_QUERIES, min=1),
    compressor_path: Path = typer.Option(DEFAULT_COMPRESSOR_PATH),
    output_root: Path = typer.Option(DEFAULT_OUTPUT_ROOT),
    dino_model: str = typer.Option("facebook/dinov2-large"),
    hash_bits: int = typer.Option(512),
    batch_size: int = typer.Option(64, min=1),
    seed: int = typer.Option(20260522),
) -> None:
    artifact = prepare_artifact(
        dataset_name=dataset,
        num_rows=num_rows,
        max_queries=max_queries,
        compressor_path=compressor_path,
        output_root=output_root,
        dino_model=dino_model,
        hash_bits=hash_bits,
        batch_size=batch_size,
        seed=seed,
    )
    print(f"CAM_RETRIEVAL_DATASET={artifact.npz_path}")
    print(f"CAM_RETRIEVAL_METADATA={artifact.metadata_path}")


if __name__ == "__main__":
    typer.run(main)