"""Generate instruction-conditioned embeddings using Qwen3-Embedding-8B.

This script generates embeddings with a router-specific instruction that guides
the model to encode task type, difficulty, and capability requirements rather
than semantic similarity for retrieval.

The key insight is that Qwen3-Embedding-8B is instruction-aware, and the default
"retrieve relevant passages" instruction optimizes for the wrong objective.
For LLM routing, we want embeddings that cluster by:
- Task type (math, coding, reasoning, factual, creative)
- Difficulty level (easy, medium, hard)
- Required LLM capabilities (tool use, long context, domain expertise)

Usage:
    cd ~/embergraphrouter
    .venv-qwen/bin/python scripts/generate_qwen3_instruct_embeddings.py

    # With specific dimensions:
    .venv-qwen/bin/python scripts/generate_qwen3_instruct_embeddings.py --dims 256,512,1024

    # Dry run to test instruction format:
    .venv-qwen/bin/python scripts/generate_qwen3_instruct_embeddings.py --dry-run
"""

import argparse
import json
import pickle
import time
from pathlib import Path
from typing import Optional

import numpy as np
import pandas as pd
import torch
from sentence_transformers import SentenceTransformer
from tqdm import tqdm

# Router-specific instruction for query embeddings
# This instruction guides the model to encode routing-relevant features
ROUTER_INSTRUCTION = (
    "Classify this query by task type (reasoning, mathematics, coding, "
    "factual recall, creative writing), cognitive difficulty level, and "
    "the language model capabilities required to answer it accurately."
)

# Alternative instructions for ablation study
ALTERNATIVE_INSTRUCTIONS = {
    "router_v1": ROUTER_INSTRUCTION,
    "router_v2": (
        "Analyze this query to determine: (1) the type of cognitive task, "
        "(2) the difficulty level from trivial to expert, and (3) what "
        "specialized capabilities an AI model would need to answer well."
    ),
    "retrieval_default": (
        "Given a web search query, retrieve relevant passages that answer the query"
    ),
    "classification": (
        "Classify this text by topic, complexity, and required expertise level."
    ),
}


def generate_embeddings_with_instruction(
    model: SentenceTransformer,
    texts: list[str],
    instruction: str,
    batch_size: int = 32,
    show_progress: bool = True,
    target_dim: Optional[int] = None,
) -> np.ndarray:
    """Generate embeddings with a custom instruction prefix.

    Args:
        model: SentenceTransformer model.
        texts: List of texts to embed.
        instruction: Instruction to prepend to each text.
        batch_size: Batch size for inference.
        show_progress: Whether to show progress bar.
        target_dim: Target dimension for Matryoshka truncation (optional).

    Returns:
        Embeddings array of shape (len(texts), embedding_dim).
    """
    # Format texts with instruction prefix (Qwen3 format)
    formatted_texts = [f"Instruct: {instruction}\nQuery: {text}" for text in texts]

    embeddings = []
    iterator = range(0, len(formatted_texts), batch_size)
    if show_progress:
        iterator = tqdm(iterator, desc="Embedding")

    for i in iterator:
        batch = formatted_texts[i:i + batch_size]
        # Use raw encode without prompt_name to avoid double-prefixing
        batch_embeddings = model.encode(
            batch,
            convert_to_numpy=True,
            normalize_embeddings=True,
        )
        embeddings.append(batch_embeddings)

    result = np.vstack(embeddings)

    # Apply Matryoshka truncation if target dimension specified
    if target_dim is not None and target_dim < result.shape[1]:
        result = result[:, :target_dim]
        # Re-normalize after truncation
        norms = np.linalg.norm(result, axis=1, keepdims=True)
        result = result / np.maximum(norms, 1e-12)

    return result


def generate_document_embeddings(
    model: SentenceTransformer,
    texts: list[str],
    batch_size: int = 32,
    show_progress: bool = True,
    target_dim: Optional[int] = None,
) -> np.ndarray:
    """Generate document embeddings (no instruction prefix for documents)."""
    embeddings = []
    iterator = range(0, len(texts), batch_size)
    if show_progress:
        iterator = tqdm(iterator, desc="Document Embedding")

    for i in iterator:
        batch = texts[i:i + batch_size]
        batch_embeddings = model.encode(
            batch,
            convert_to_numpy=True,
            normalize_embeddings=True,
        )
        embeddings.append(batch_embeddings)

    result = np.vstack(embeddings)

    if target_dim is not None and target_dim < result.shape[1]:
        result = result[:, :target_dim]
        norms = np.linalg.norm(result, axis=1, keepdims=True)
        result = result / np.maximum(norms, 1e-12)

    return result


def main():
    parser = argparse.ArgumentParser(
        description="Generate router-instruction-conditioned Qwen3 embeddings"
    )
    parser.add_argument(
        "--instruction",
        type=str,
        default="router_v1",
        choices=list(ALTERNATIVE_INSTRUCTIONS.keys()),
        help="Instruction variant to use (default: router_v1)",
    )
    parser.add_argument(
        "--dims",
        type=str,
        default="256,512,1024,2048,4096",
        help="Comma-separated list of target dimensions (default: 256,512,1024,2048,4096)",
    )
    parser.add_argument(
        "--batch-size",
        type=int,
        default=64,
        help="Batch size for embedding generation (default: 64)",
    )
    parser.add_argument(
        "--dry-run",
        action="store_true",
        help="Test instruction format without full generation",
    )
    parser.add_argument(
        "--device",
        type=str,
        default="cuda:0",
        help="Device to use (default: cuda:0)",
    )
    parser.add_argument(
        "--input",
        type=str,
        default="data/router_data.csv",
        help="Input data file (default: data/router_data.csv)",
    )
    args = parser.parse_args()

    target_dims = [int(d.strip()) for d in args.dims.split(",")]
    instruction = ALTERNATIVE_INSTRUCTIONS[args.instruction]
    instruction_tag = args.instruction

    print("=" * 70)
    print("Qwen3-Embedding-8B: Instruction-Conditioned Embedding Generation")
    print("=" * 70)
    print(f"\nInstruction variant: {instruction_tag}")
    print(f"Instruction text:\n  \"{instruction}\"")
    print(f"\nTarget dimensions: {target_dims}")
    print(f"Device: {args.device}")

    if args.dry_run:
        print("\n" + "=" * 70)
        print("DRY RUN: Testing instruction format")
        print("=" * 70)

        sample_queries = [
            "What is 2 + 2?",
            "Explain the theory of relativity in simple terms",
            "Write a Python function to sort a list",
            "What are the ethical implications of artificial general intelligence?",
        ]

        print("\nSample formatted inputs:")
        for q in sample_queries:
            formatted = f"Instruct: {instruction}\nQuery: {q}"
            print(f"\n  Input: {q[:50]}...")
            print(f"  Formatted: {formatted[:100]}...")

        # Load model and test a few embeddings
        print("\nLoading model for test...")
        model = SentenceTransformer(
            "Qwen/Qwen3-Embedding-8B",
            device=args.device,
            trust_remote_code=True,
        )

        test_embs = generate_embeddings_with_instruction(
            model, sample_queries, instruction, batch_size=4, show_progress=False
        )
        print(f"\nTest embeddings shape: {test_embs.shape}")
        print(f"Embedding norm range: [{np.linalg.norm(test_embs, axis=1).min():.4f}, {np.linalg.norm(test_embs, axis=1).max():.4f}]")

        # Compare with default instruction
        default_embs = generate_embeddings_with_instruction(
            model, sample_queries, ALTERNATIVE_INSTRUCTIONS["retrieval_default"],
            batch_size=4, show_progress=False
        )

        # Cosine similarity between router and retrieval embeddings
        similarities = np.sum(test_embs * default_embs, axis=1)
        print(f"\nCosine similarity between router vs retrieval embeddings:")
        for q, sim in zip(sample_queries, similarities):
            print(f"  {q[:40]}...: {sim:.4f}")

        print("\nDry run complete. Use without --dry-run to generate full embeddings.")
        return

    # Full embedding generation
    output_dir = Path("configs")
    data_path = Path(args.input)

    print("\nLoading model...")
    start = time.time()
    model = SentenceTransformer(
        "Qwen/Qwen3-Embedding-8B",
        device=args.device,
        trust_remote_code=True,
    )
    print(f"Model loaded in {time.time() - start:.1f}s")
    print(f"Native embedding dimension: {model.get_sentence_embedding_dimension()}")

    # Load data
    print(f"\nLoading data from {data_path}...")
    df = pd.read_csv(data_path)
    print(f"Total rows: {len(df):,}")

    # Get unique queries and task descriptions
    unique_queries = df["query"].unique().tolist()
    unique_tasks = df[["task_id", "task_description"]].drop_duplicates()

    print(f"Unique queries: {len(unique_queries):,}")
    print(f"Unique tasks: {len(unique_tasks):,}")

    # Generate query embeddings with router instruction (full dim first)
    print("\n" + "=" * 70)
    print(f"Generating QUERY embeddings with instruction: {instruction_tag}")
    print("=" * 70)
    start = time.time()
    query_embeddings_full = generate_embeddings_with_instruction(
        model, unique_queries, instruction, batch_size=args.batch_size
    )
    print(f"Full query embeddings shape: {query_embeddings_full.shape}")
    print(f"Time: {time.time() - start:.1f}s")

    # Generate task description embeddings (no instruction - documents)
    print("\n" + "=" * 70)
    print("Generating TASK DESCRIPTION embeddings (document mode, no instruction)")
    print("=" * 70)
    task_descriptions = unique_tasks["task_description"].tolist()
    start = time.time()
    task_embeddings_full = generate_document_embeddings(
        model, task_descriptions, batch_size=8
    )
    print(f"Task embeddings shape: {task_embeddings_full.shape}")
    print(f"Time: {time.time() - start:.1f}s")

    # Generate LLM description embeddings
    print("\n" + "=" * 70)
    print("Generating LLM DESCRIPTION embeddings (document mode)")
    print("=" * 70)

    with open("configs/LLM_Descriptions.json") as f:
        llm_descriptions = json.load(f)

    llm_names = df["llm"].unique().tolist()
    llm_texts = []
    for name in llm_names:
        desc = llm_descriptions.get(name, {})
        if isinstance(desc, dict):
            llm_texts.append(desc.get("feature", name))
        else:
            llm_texts.append(str(desc))

    print(f"LLMs: {llm_names}")
    start = time.time()
    llm_embeddings_full = generate_document_embeddings(
        model, llm_texts, batch_size=8
    )
    print(f"LLM embeddings shape: {llm_embeddings_full.shape}")
    print(f"Time: {time.time() - start:.1f}s")

    # Save embeddings at each target dimension
    print("\n" + "=" * 70)
    print("Saving embeddings at each target dimension")
    print("=" * 70)

    for dim in target_dims:
        print(f"\nDimension: {dim}")

        # Truncate and re-normalize (Matryoshka)
        if dim < query_embeddings_full.shape[1]:
            query_emb = query_embeddings_full[:, :dim]
            task_emb = task_embeddings_full[:, :dim]
            llm_emb = llm_embeddings_full[:, :dim]

            # Re-normalize
            for arr in [query_emb, task_emb, llm_emb]:
                norms = np.linalg.norm(arr, axis=1, keepdims=True)
                arr /= np.maximum(norms, 1e-12)
        else:
            query_emb = query_embeddings_full
            task_emb = task_embeddings_full
            llm_emb = llm_embeddings_full

        # Save with instruction tag in filename
        suffix = f"qwen3_instruct_{instruction_tag}_{dim}"

        # Query embeddings
        query_map = {q: emb for q, emb in zip(unique_queries, query_emb)}
        query_path = output_dir / f"query_embeddings_{suffix}.pkl"
        with open(query_path, "wb") as f:
            pickle.dump(query_map, f)
        print(f"  Saved: {query_path.name}")

        # Task embeddings
        task_map = {
            row["task_id"]: task_emb[i]
            for i, (_, row) in enumerate(unique_tasks.iterrows())
        }
        task_path = output_dir / f"task_embeddings_{suffix}.pkl"
        with open(task_path, "wb") as f:
            pickle.dump(task_map, f)
        print(f"  Saved: {task_path.name}")

        # LLM embeddings
        llm_path = output_dir / f"llm_embedding_{suffix}.pkl"
        with open(llm_path, "wb") as f:
            pickle.dump({
                "embeddings": np.array(llm_emb, dtype=np.float32),
                "names": llm_names,
            }, f)
        print(f"  Saved: {llm_path.name}")

    # Also save the parquet with new columns
    print("\n" + "=" * 70)
    print("Adding embeddings to parquet")
    print("=" * 70)

    # Create query embedding lookup
    query_embedding_lookup = {q: query_embeddings_full[i] for i, q in enumerate(unique_queries)}

    # Add embedding columns for each dimension
    for dim in target_dims:
        col_name = f"query_embedding_qwen3_instruct_{instruction_tag}_{dim}"

        def get_truncated_embedding(query, d=dim):
            emb = query_embedding_lookup.get(query)
            if emb is None:
                return None
            truncated = emb[:d]
            norm = np.linalg.norm(truncated)
            if norm > 1e-12:
                truncated = truncated / norm
            return truncated.tolist()

        df[col_name] = df["query"].apply(get_truncated_embedding)
        print(f"  Added column: {col_name}")

    # Save enriched parquet
    output_parquet = Path("data") / f"router_data_qwen3_instruct_{instruction_tag}.parquet"
    df.to_parquet(output_parquet, index=False)
    print(f"\nSaved: {output_parquet}")

    print("\n" + "=" * 70)
    print("COMPLETE!")
    print("=" * 70)
    print(f"\nGenerated embeddings with instruction: {instruction_tag}")
    print(f"Dimensions: {target_dims}")
    print(f"\nTo run comparison sweep, use the new parquet file:")
    print(f"  {output_parquet}")


if __name__ == "__main__":
    main()
