from __future__ import annotations

import json
import os
import time
import gc
import traceback
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")
# The completed 0.6B phase is intentionally not resident while the larger GTE
# startup path is validated.  Packing both models together exceeded the
# current Space startup envelope at 6.16 GB.
model = None

# ZeroGPU optimizes CUDA placements made during module startup.  The previous
# larger-family path loaded this model inside the decorated call and exhausted
# that reservation before producing an artifact.
GTE_MODEL_PATH = Path("/models/gte-qwen2-1.5b")
gte_tokenizer = None
gte_model = None

QWEN8_MODEL_PATH = Path("/models/qwen3-embedding-8b")
qwen8_tokenizer = AutoTokenizer.from_pretrained(QWEN8_MODEL_PATH, padding_side="left")
qwen8_model = None


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


def encode_with(active_model, text: str, active_tokenizer=None) -> np.ndarray:
    active_tokenizer = active_tokenizer or tokenizer
    batch = active_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, use_cache=False)
        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:
    if model is None:
        raise RuntimeError("0.6B BF16 control is offline during the GTE phase")
    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


BF16_FAMILIES = {
    "gte-qwen2-1.5b": {
        "path": Path("/models/gte-qwen2-1.5b"),
        "source": "Alibaba-NLP/gte-Qwen2-1.5B-instruct",
        "revision": "a9af15a6372d7d6b25e9fb07c2ccb9e1fe645644",
    },
    "qwen3-embedding-8b": {
        "path": Path("/models/qwen3-embedding-8b"),
        "source": "Qwen/Qwen3-Embedding-8B",
        "revision": "1d8ad4ca9b3dd8059ad90a75d4983776a23d44af",
    },
}


def _run_family_bf16_control(family: str) -> dict:
    if family not in BF16_FAMILIES:
        raise ValueError(f"unsupported family: {family}")
    spec = BF16_FAMILIES[family]
    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]

    torch.cuda.reset_peak_memory_stats()
    torch.cuda.synchronize()
    allocation_before = int(torch.cuda.memory_allocated())
    family_tokenizer = AutoTokenizer.from_pretrained(
        spec["path"], padding_side="left", trust_remote_code=True
    )
    load_started = time.perf_counter()
    family_model = AutoModel.from_pretrained(
        spec["path"],
        dtype=torch.bfloat16,
        trust_remote_code=True,
        low_cpu_mem_usage=True,
        device_map={"": "cuda"},
    ).eval()
    torch.cuda.synchronize()
    load_seconds = time.perf_counter() - load_started
    loading_strategy = "in-call-direct-cuda-device-map"
    owns_model = True
    allocation_after_load = int(torch.cuda.memory_allocated())

    encode_started = time.perf_counter()
    queries = np.stack([
        encode_with(family_model, text, family_tokenizer) for text in query_texts
    ])
    documents = np.stack([
        encode_with(family_model, text, family_tokenizer) 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": "cuda-zerogpu-bf16",
        "family": family,
        "model": spec["source"],
        "source_revision": spec["revision"],
        "load_seconds": load_seconds,
        "loading_strategy": loading_strategy,
        "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()),
    })
    local_reference = (
        DATA_ROOT / "local-results/full-q4-q6-q8" / family / "quality/bf16.npz"
    )
    result = {
        "metrics": metrics,
        "mlx_bf16_comparison": compare_vectors(
            queries, documents, scores, ranks, local_reference, "cuda_vs_mlx_bf16"
        ),
        "previous_cuda_bf16_comparison": compare_vectors(
            queries,
            documents,
            scores,
            ranks,
            DATA_ROOT / "cloud-results" / family / "cuda-bf16-root-pack.npz",
            "direct_vs_root_pack_bf16",
        ),
    }
    output_dir = DATA_ROOT / "cloud-results" / family
    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")

    if owns_model:
        del family_model, family_tokenizer
        gc.collect()
        torch.cuda.empty_cache()
    return result


@spaces.GPU(duration=300)
def run_family_bf16_control(family: str) -> dict:
    try:
        return _run_family_bf16_control(family)
    except Exception as error:
        return {
            "diagnostic_error": type(error).__name__,
            "diagnostic_message": str(error),
            "diagnostic_traceback": traceback.format_exc(),
        }


def _run_family_quantized_control(family: str, variant: str) -> dict:
    if family not in BF16_FAMILIES:
        raise ValueError(f"unsupported family: {family}")
    if variant not in QUANTIZERS:
        raise ValueError(f"unsupported quantizer: {variant}")

    spec = BF16_FAMILIES[family]
    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]
    family_tokenizer = AutoTokenizer.from_pretrained(
        spec["path"], padding_side="left", trust_remote_code=True
    )

    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()
    family_model = AutoModel.from_pretrained(
        spec["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(family_model, text, family_tokenizer) for text in query_texts
    ])
    documents = np.stack([
        encode_with(family_model, text, family_tokenizer) 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}",
        "family": family,
        "model": spec["source"],
        "source_revision": spec["revision"],
        "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()),
    })
    output_dir = DATA_ROOT / "cloud-results" / family
    result = {
        "metrics": metrics,
        "cuda_bf16_comparison": compare_vectors(
            queries,
            documents,
            scores,
            ranks,
            output_dir / "cuda-bf16.npz",
            "vs_cuda_bf16",
        ),
        "mlx_bf16_comparison": compare_vectors(
            queries,
            documents,
            scores,
            ranks,
            DATA_ROOT / "local-results/full-q4-q6-q8" / family / "quality/bf16.npz",
            "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 family_model, family_tokenizer
    gc.collect()
    torch.cuda.empty_cache()
    return result


@spaces.GPU(duration=300)
def run_family_quantized_control(family: str, variant: str) -> dict:
    try:
        return _run_family_quantized_control(family, variant)
    except Exception as error:
        return {
            "diagnostic_error": type(error).__name__,
            "diagnostic_message": str(error),
            "diagnostic_traceback": traceback.format_exc(),
        }


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",
    )
    gr.Markdown("## Larger-family BF16 cross-runtime controls")
    family = gr.Dropdown(
        choices=list(BF16_FAMILIES),
        value="gte-qwen2-1.5b",
        label="Model family",
    )
    family_button = gr.Button("Run bounded family BF16 control")
    family_output = gr.JSON(label="Family result")
    family_button.click(
        fn=run_family_bf16_control,
        inputs=family,
        outputs=family_output,
        concurrency_limit=1,
        api_name="run_family_bf16_control",
    )
    gr.Markdown("## Larger-family CUDA-native quantizer controls")
    quant_family = gr.Dropdown(
        choices=list(BF16_FAMILIES),
        value="qwen3-embedding-8b",
        label="Model family",
    )
    family_quantizer = gr.Dropdown(
        choices=list(QUANTIZERS), value="bnb-int8", label="Quantizer"
    )
    family_quant_button = gr.Button("Run bounded family quantizer control")
    family_quant_output = gr.JSON(label="Family quantized result")
    family_quant_button.click(
        fn=run_family_quantized_control,
        inputs=[quant_family, family_quantizer],
        outputs=family_quant_output,
        concurrency_limit=1,
        api_name="run_family_quantized_control",
    )


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