from __future__ import annotations

import json
import os
import time
import gc
from pathlib import Path

import gradio as gr
import numpy as np
import spaces
import torch
import torch.nn.functional as F
from transformers import AutoModel, AutoTokenizer, BitsAndBytesConfig


MODEL_PATH = Path("/models/qwen3-embedding-0.6b")
DATA_ROOT = Path(os.environ.get("REPRO_DATA_ROOT", "/data"))
INPUT_PATH = DATA_ROOT / "inputs/retrieval_pairs.json"
MLX_REFERENCE = DATA_ROOT / "local-reference/qwen3-embedding-0.6b/bf16.npz"
OUTPUT_DIR = DATA_ROOT / "cloud-results/qwen3-embedding-0.6b"
CUDA_BF16_REFERENCE = OUTPUT_DIR / "cuda-bf16.npz"
TASK = (
    "Given a natural-language search query, retrieve the single passage "
    "that best answers it"
)


tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH, padding_side="left")
model = AutoModel.from_pretrained(
    MODEL_PATH,
    torch_dtype=torch.bfloat16,
    trust_remote_code=True,
).to("cuda").eval()


def detailed_instruction(query: str) -> str:
    return f"Instruct: {TASK}\nQuery:{query}"


def encode_with(active_model, text: str) -> np.ndarray:
    batch = tokenizer(text, return_tensors="pt", truncation=True, max_length=32768)
    batch = {key: value.to("cuda") for key, value in batch.items()}
    with torch.inference_mode():
        output = active_model(**batch)
        vector = F.normalize(output.last_hidden_state[:, -1, :].float(), p=2, dim=-1)
    return vector[0].cpu().numpy()


def encode(text: str) -> np.ndarray:
    return encode_with(model, text)


def retrieval_metrics(queries: np.ndarray, documents: np.ndarray) -> tuple[dict, np.ndarray, np.ndarray]:
    scores = queries @ documents.T
    order = np.argsort(-scores, axis=1)
    ranks = np.array([
        int(np.where(order[index] == index)[0][0]) + 1
        for index in range(len(queries))
    ])
    positive = np.diag(scores)
    negative = scores.copy()
    np.fill_diagonal(negative, -np.inf)
    margins = positive - negative.max(axis=1)
    return {
        "pair_count": len(queries),
        "top1": float(np.mean(ranks == 1)),
        "recall_at_5": float(np.mean(ranks <= 5)),
        "mrr": float(np.mean(1.0 / ranks)),
        "mean_margin": float(margins.mean()),
        "minimum_margin": float(margins.min()),
        "mean_rank": float(ranks.mean()),
        "worst_rank": int(ranks.max()),
    }, scores, ranks


@spaces.GPU(duration=300)
def run_bf16_control() -> dict:
    pairs = json.loads(INPUT_PATH.read_text())
    query_texts = [detailed_instruction(item["query"]) for item in pairs]
    document_texts = [item["document"] for item in pairs]

    if torch.cuda.is_available():
        torch.cuda.reset_peak_memory_stats()
        torch.cuda.synchronize()
    started = time.perf_counter()
    queries = np.stack([encode(text) for text in query_texts])
    documents = np.stack([encode(text) for text in document_texts])
    if torch.cuda.is_available():
        torch.cuda.synchronize()
    elapsed = time.perf_counter() - started
    metrics, scores, ranks = retrieval_metrics(queries, documents)
    metrics.update({
        "lane": "cuda-zerogpu-bf16",
        "model": "Qwen/Qwen3-Embedding-0.6B",
        "source_revision": "97b0c614be4d77ee51c0cef4e5f07c00f9eb65b3",
        "elapsed_seconds": elapsed,
        "texts_per_second": len(query_texts + document_texts) / elapsed,
        "torch_version": torch.__version__,
        "cuda_device": torch.cuda.get_device_name(0) if torch.cuda.is_available() else None,
        "cuda_peak_bytes": torch.cuda.max_memory_allocated() if torch.cuda.is_available() else None,
    })

    comparison = None
    if MLX_REFERENCE.exists():
        reference = np.load(MLX_REFERENCE)
        ref_all = np.concatenate([reference["queries"], reference["documents"]])
        cuda_all = np.concatenate([queries, documents])
        aligned = np.sum(ref_all * cuda_all, axis=1)
        score_delta = scores - reference["scores"]
        comparison = {
            "mean_aligned_cosine_cuda_vs_mlx_bf16": float(aligned.mean()),
            "minimum_aligned_cosine_cuda_vs_mlx_bf16": float(aligned.min()),
            "score_rmse_cuda_vs_mlx_bf16": float(np.sqrt(np.mean(score_delta ** 2))),
            "queries_with_rank_change": int(np.count_nonzero(ranks - reference["ranks"])),
        }

    result = {"metrics": metrics, "mlx_bf16_comparison": comparison}
    OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
    np.savez_compressed(
        OUTPUT_DIR / "cuda-bf16.npz",
        queries=queries,
        documents=documents,
        scores=scores,
        ranks=ranks,
        metrics=np.array(json.dumps(metrics)),
    )
    (OUTPUT_DIR / "cuda-bf16.json").write_text(json.dumps(result, indent=2) + "\n")
    return result


def compare_vectors(
    queries: np.ndarray,
    documents: np.ndarray,
    scores: np.ndarray,
    ranks: np.ndarray,
    reference_path: Path,
    prefix: str,
) -> dict | None:
    if not reference_path.exists():
        return None
    reference = np.load(reference_path)
    ref_all = np.concatenate([reference["queries"], reference["documents"]])
    candidate_all = np.concatenate([queries, documents])
    aligned = np.sum(ref_all * candidate_all, axis=1)
    score_delta = scores - reference["scores"]
    return {
        f"mean_aligned_cosine_{prefix}": float(aligned.mean()),
        f"minimum_aligned_cosine_{prefix}": float(aligned.min()),
        f"score_rmse_{prefix}": float(np.sqrt(np.mean(score_delta ** 2))),
        f"queries_with_rank_change_{prefix}": int(
            np.count_nonzero(ranks - reference["ranks"])
        ),
    }


QUANTIZERS = ("bnb-int8", "bnb-nf4")


def quantization_config(variant: str) -> BitsAndBytesConfig:
    if variant == "bnb-int8":
        return BitsAndBytesConfig(load_in_8bit=True)
    if variant == "bnb-nf4":
        return BitsAndBytesConfig(
            load_in_4bit=True,
            bnb_4bit_quant_type="nf4",
            bnb_4bit_compute_dtype=torch.bfloat16,
            bnb_4bit_use_double_quant=False,
        )
    raise ValueError(f"unsupported quantizer: {variant}")


@spaces.GPU(duration=300)
def run_quantized_control(variant: str) -> dict:
    if variant not in QUANTIZERS:
        raise ValueError(f"unsupported quantizer: {variant}")
    config = quantization_config(variant)
    pairs = json.loads(INPUT_PATH.read_text())
    query_texts = [detailed_instruction(item["query"]) for item in pairs]
    document_texts = [item["document"] for item in pairs]

    gc.collect()
    torch.cuda.empty_cache()
    torch.cuda.reset_peak_memory_stats()
    torch.cuda.synchronize()
    allocation_before = int(torch.cuda.memory_allocated())
    load_started = time.perf_counter()
    quantized_model = AutoModel.from_pretrained(
        MODEL_PATH,
        quantization_config=config,
        device_map={"": 0},
        trust_remote_code=True,
    ).eval()
    torch.cuda.synchronize()
    load_seconds = time.perf_counter() - load_started
    allocation_after_load = int(torch.cuda.memory_allocated())

    encode_started = time.perf_counter()
    queries = np.stack([encode_with(quantized_model, text) for text in query_texts])
    documents = np.stack([
        encode_with(quantized_model, text) for text in document_texts
    ])
    torch.cuda.synchronize()
    encode_seconds = time.perf_counter() - encode_started
    metrics, scores, ranks = retrieval_metrics(queries, documents)
    metrics.update({
        "lane": f"cuda-zerogpu-{variant}",
        "model": "Qwen/Qwen3-Embedding-0.6B",
        "source_revision": "97b0c614be4d77ee51c0cef4e5f07c00f9eb65b3",
        "quantizer": variant,
        "quantization_config": config.to_dict(),
        "load_seconds": load_seconds,
        "encode_seconds": encode_seconds,
        "texts_per_second": len(query_texts + document_texts) / encode_seconds,
        "torch_version": torch.__version__,
        "cuda_device": torch.cuda.get_device_name(0),
        "cuda_allocation_before_load": allocation_before,
        "cuda_allocation_after_load": allocation_after_load,
        "cuda_incremental_model_allocation": allocation_after_load - allocation_before,
        "cuda_peak_bytes": int(torch.cuda.max_memory_allocated()),
    })
    result = {
        "metrics": metrics,
        "cuda_bf16_comparison": compare_vectors(
            queries, documents, scores, ranks, CUDA_BF16_REFERENCE, "vs_cuda_bf16"
        ),
        "mlx_bf16_comparison": compare_vectors(
            queries, documents, scores, ranks, MLX_REFERENCE, "vs_mlx_bf16"
        ),
    }
    OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
    stem = f"cuda-{variant}"
    np.savez_compressed(
        OUTPUT_DIR / f"{stem}.npz",
        queries=queries,
        documents=documents,
        scores=scores,
        ranks=ranks,
        metrics=np.array(json.dumps(metrics)),
    )
    (OUTPUT_DIR / f"{stem}.json").write_text(json.dumps(result, indent=2) + "\n")

    del quantized_model
    gc.collect()
    torch.cuda.empty_cache()
    return result


with gr.Blocks() as demo:
    gr.Markdown(
        "# Embedding Quantization CUDA Control\n"
        "Runs one bounded 48-text BF16 control and writes the result to the attached private bucket. "
        "The requested GPU duration is capped at five minutes."
    )
    run_button = gr.Button("Run 0.6B BF16 CUDA control", variant="primary")
    output = gr.JSON(label="Result")
    run_button.click(
        fn=run_bf16_control,
        outputs=output,
        concurrency_limit=1,
        api_name="run_bf16_control",
    )
    gr.Markdown(
        "## CUDA-native quantizer controls\n"
        "These are bitsandbytes INT8/NF4 controls, not MLX Q/oQ/oQe replicas."
    )
    quantizer = gr.Dropdown(
        choices=list(QUANTIZERS), value="bnb-int8", label="Quantizer"
    )
    quant_button = gr.Button("Run bounded CUDA quantizer control")
    quant_output = gr.JSON(label="Quantized result")
    quant_button.click(
        fn=run_quantized_control,
        inputs=quantizer,
        outputs=quant_output,
        concurrency_limit=1,
        api_name="run_quantized_control",
    )


if __name__ == "__main__":
    demo.queue(default_concurrency_limit=1).launch()