#!/usr/bin/env python3
"""
Profiling script for CyborgDB Service using wiki_all_1M dataset.
This script loads data, creates an index via the REST API, upserts, trains, and queries.
Uses the cyborgdb Python SDK to communicate with cyborgdb-service.

Measures recall@k and QPS (queries per second).
"""

import sys
import os

# Disable threading in BLAS/LAPACK/OpenMP libraries for clearer profiling
os.environ["OPENBLAS_NUM_THREADS"] = "1"
os.environ["MKL_NUM_THREADS"] = "1"
os.environ["OMP_NUM_THREADS"] = "1"
os.environ["VECLIB_MAXIMUM_THREADS"] = "1"
os.environ["NUMEXPR_NUM_THREADS"] = "1"

import numpy as np
import h5py
import time
import argparse

# Import the CyborgDB Python SDK
import cyborgdb

# ========================= CONFIGURATION =========================
# Dataset configuration
DATASET_PATH = os.path.join(
    os.path.dirname(__file__),
    "..",
    "..",
    "cyborgdb-core",
    "datasets",
    "wiki_all_1M.hdf5",
)
NUM_VECTORS_TO_LOAD = 1000000  # Load 100K vectors by default (adjustable)
NUM_QUERY_VECTORS = 100  # Number of queries to run

# Index configuration
INDEX_TYPE = "ivfflat"  # Options: "ivf", "ivfflat", "ivfpq"
DIMENSION = 768  # Wiki embeddings dimension
N_LISTS = 0  # Number of clusters for IVF (0 = auto)
PQ_DIM = 768  # For IVFPQ only
PQ_BITS = 8  # For IVFPQ only

# Query configuration
TOP_K = 100  # Number of results to return
N_PROBES = 0  # Number of clusters to search (0 = auto)

# Training configuration
BATCH_SIZE = 256
MAX_ITERS = 50
TOLERANCE = 1e-6

# Service configuration
SERVICE_URL = os.getenv("CYBORGDB_SERVICE_URL", "http://localhost:8000")
API_KEY = os.getenv("CYBORGDB_API_KEY", "test-api-key")

# ================================================================


def load_data(num_vectors=NUM_VECTORS_TO_LOAD, num_queries=NUM_QUERY_VECTORS):
    """Load data from HDF5 file."""
    print(f"Loading data from {DATASET_PATH}...")

    if not os.path.exists(DATASET_PATH):
        raise FileNotFoundError(f"Dataset not found at {DATASET_PATH}")

    with h5py.File(DATASET_PATH, "r") as f:
        # Load train data (for upserting)
        train_data = np.array(f["train"][:num_vectors], dtype=np.float32)

        # Load test data (for querying)
        test_data = np.array(f["test"][:num_queries], dtype=np.float32)

        # Load neighbors (ground truth for recall calculation)
        neighbors = np.array(f["neighbors"][:num_queries, :TOP_K], dtype=np.int32)

        print(
            f"Loaded {train_data.shape[0]} training vectors with dimension {train_data.shape[1]}"
        )
        print(f"Loaded {test_data.shape[0]} query vectors")
        print(
            f"Loaded {neighbors.shape[0]} neighbor sets (top {neighbors.shape[1]} per query)"
        )

    return train_data, test_data, neighbors


def create_index(client, index_name):
    """Create an index based on configuration."""
    print(f"Creating {INDEX_TYPE.upper()} index with dimension {DIMENSION}...")

    if INDEX_TYPE == "ivf":
        index_config = cyborgdb.IndexIVF(dimension=DIMENSION)
    elif INDEX_TYPE == "ivfflat":
        index_config = cyborgdb.IndexIVFFlat(dimension=DIMENSION)
    elif INDEX_TYPE == "ivfpq":
        index_config = cyborgdb.IndexIVFPQ(
            dimension=DIMENSION, pq_dim=PQ_DIM, pq_bits=PQ_BITS
        )
    else:
        raise ValueError(f"Unknown index type: {INDEX_TYPE}")

    index_key = cyborgdb.Client.generate_key()
    index = client.create_index(index_name, index_key, index_config, metric="euclidean")

    print(f"Index '{index_name}' created successfully")
    return index, index_key


def upsert_data(index, train_data, batch_size=100000):
    """Upsert training data into the index in batches."""
    print(f"Upserting {len(train_data)} vectors in batches of {batch_size}...")

    start_time = time.time()
    total_vectors = len(train_data)

    for batch_start in range(0, total_vectors, batch_size):
        batch_end = min(batch_start + batch_size, total_vectors)
        batch_data = train_data[batch_start:batch_end]

        # upsert auto-detects numpy arrays and uses binary format
        ids = [str(i) for i in range(batch_start, batch_end)]
        index.upsert(ids, batch_data)

        progress = (batch_end / total_vectors) * 100
        print(f"  Upserted {batch_end}/{total_vectors} vectors ({progress:.1f}%)")

    elapsed = time.time() - start_time
    print(
        f"Upsert completed in {elapsed:.2f} seconds ({len(train_data) / elapsed:.2f} vectors/sec)"
    )


def train_index(index, n_lists=N_LISTS):
    """Train the index."""
    effective_n_lists = n_lists if n_lists > 0 else None
    print(f"Training index with n_lists={effective_n_lists or 'auto'}...")

    start_time = time.time()

    index.train(
        n_lists=effective_n_lists,
        # batch_size=BATCH_SIZE,
        # max_iters=MAX_ITERS,
        # tolerance=TOLERANCE,
    )

    elapsed = time.time() - start_time
    print(f"Training completed in {elapsed:.2f} seconds")

    # Wait for training to complete (it may be async)
    print("Waiting for training to complete...")
    max_wait = 300  # 5 minutes max
    wait_start = time.time()
    while index.is_training():
        if time.time() - wait_start > max_wait:
            print("WARNING: Training is taking longer than expected")
            break
        time.sleep(2)
        print("  Still training...")

    if not index.is_trained():
        raise RuntimeError("Index training failed")

    print("Index is trained and ready for queries")


def calculate_recall(all_results, neighbors):
    """Calculate recall@k for all queries in batch."""
    recalls = []
    for i, result in enumerate(all_results):
        # Extract IDs from the result dictionaries
        if isinstance(result, list) and len(result) > 0:
            if isinstance(result[0], list):
                # Batch result format - take first query's results
                result_ids = set(int(item["id"]) for item in result[0])
            else:
                result_ids = set(int(item["id"]) for item in result)
        else:
            result_ids = set()

        # Get ground truth neighbors for this query
        ground_truth = set(neighbors[i])

        # Calculate recall for this query
        intersection = len(result_ids.intersection(ground_truth))
        recall = intersection / len(ground_truth) if len(ground_truth) > 0 else 0.0
        recalls.append(recall)

    return np.array(recalls)


def run_queries(index, test_data, neighbors, n_probes=N_PROBES):
    """Run queries on the index one at a time - THIS IS THE MAIN PROFILING TARGET."""
    effective_n_probes = n_probes if n_probes > 0 else None
    print(
        f"\nRunning {len(test_data)} queries (one at a time), n_probes={effective_n_probes or 'auto'}..."
    )
    print("=" * 70)

    query_times = []
    all_results = []
    start_time = time.time()

    # Execute queries one at a time (not batched)
    for i, query_vector in enumerate(test_data):
        query_start = time.time()

        # Single query at a time
        results = index.query(
            query_vectors=query_vector.reshape(1, -1),
            top_k=TOP_K,
            n_probes=effective_n_probes,
        )

        query_elapsed = time.time() - query_start
        query_times.append(query_elapsed)
        all_results.append(results)

        # Print progress every 10 queries
        if (i + 1) % 10 == 0:
            avg_time = np.mean(query_times[-10:])
            print(
                f"  Completed {i + 1}/{len(test_data)} queries (avg last 10: {avg_time * 1000:.3f}ms per query)"
            )

    total_elapsed = time.time() - start_time

    print("\n" + "=" * 70)
    print("Query Performance Summary:")
    print(f"  Total queries executed: {len(test_data)}")
    print(f"  Total time: {total_elapsed:.3f} seconds")
    print(f"  Average time per query: {np.mean(query_times) * 1000:.3f}ms")
    print(f"  Median time per query: {np.median(query_times) * 1000:.3f}ms")
    print(f"  Std dev: {np.std(query_times) * 1000:.3f}ms")
    print(f"  Min time: {np.min(query_times) * 1000:.3f}ms")
    print(f"  Max time: {np.max(query_times) * 1000:.3f}ms")
    print(f"  Queries per second (QPS): {len(test_data) / total_elapsed:.2f}")
    print("=" * 70)

    # Calculate recall
    recalls = calculate_recall(all_results, neighbors)

    print("\nRecall Summary:")
    print(f"  Average Recall@{TOP_K}: {np.mean(recalls) * 100:.2f}%")
    print(f"  Median Recall@{TOP_K}: {np.median(recalls) * 100:.2f}%")
    print(f"  Min Recall@{TOP_K}: {np.min(recalls) * 100:.2f}%")
    print(f"  Max Recall@{TOP_K}: {np.max(recalls) * 100:.2f}%")
    print("=" * 70)

    return {
        "qps": len(test_data) / total_elapsed,
        "avg_latency_ms": np.mean(query_times) * 1000,
        "median_latency_ms": np.median(query_times) * 1000,
        "recall": np.mean(recalls),
        "query_times": query_times,
        "recalls": recalls,
    }


def run_batch_queries(index, test_data, neighbors, n_probes=N_PROBES, batch_size=10):
    """Run queries in batches for higher throughput measurement."""
    effective_n_probes = n_probes if n_probes > 0 else None
    print(
        f"\nRunning {len(test_data)} queries in batches of {batch_size}, n_probes={effective_n_probes or 'auto'}..."
    )
    print("=" * 70)

    batch_times = []
    all_results = []
    start_time = time.time()

    num_batches = (len(test_data) + batch_size - 1) // batch_size

    for batch_idx in range(num_batches):
        batch_start_idx = batch_idx * batch_size
        batch_end_idx = min(batch_start_idx + batch_size, len(test_data))
        batch_vectors = test_data[batch_start_idx:batch_end_idx]

        batch_start = time.time()

        # Batch query
        results = index.query(
            query_vectors=batch_vectors, top_k=TOP_K, n_probes=effective_n_probes
        )

        batch_elapsed = time.time() - batch_start
        batch_times.append(batch_elapsed)

        # Results is a list of lists for batch queries
        if isinstance(results, list) and len(results) > 0:
            if isinstance(results[0], list):
                all_results.extend(results)
            else:
                all_results.append(results)

    total_elapsed = time.time() - start_time
    total_queries = len(test_data)

    print("\n" + "=" * 70)
    print("Batch Query Performance Summary:")
    print(f"  Total queries executed: {total_queries}")
    print(f"  Batch size: {batch_size}")
    print(f"  Total time: {total_elapsed:.3f} seconds")
    print(f"  Average batch time: {np.mean(batch_times) * 1000:.3f}ms")
    print(f"  Queries per second (QPS): {total_queries / total_elapsed:.2f}")
    print("=" * 70)

    # Calculate recall
    recalls = calculate_recall(all_results, neighbors)

    print("\nRecall Summary:")
    print(f"  Average Recall@{TOP_K}: {np.mean(recalls) * 100:.2f}%")
    print(f"  Median Recall@{TOP_K}: {np.median(recalls) * 100:.2f}%")
    print("=" * 70)

    return {
        "qps": total_queries / total_elapsed,
        "avg_batch_latency_ms": np.mean(batch_times) * 1000,
        "recall": np.mean(recalls),
    }


def main():
    """Main profiling workflow."""
    parser = argparse.ArgumentParser(
        description="Profile CyborgDB Service with wiki dataset"
    )
    parser.add_argument(
        "--vectors",
        type=int,
        default=NUM_VECTORS_TO_LOAD,
        help=f"Number of vectors to load (default: {NUM_VECTORS_TO_LOAD})",
    )
    parser.add_argument(
        "--queries",
        type=int,
        default=NUM_QUERY_VECTORS,
        help=f"Number of query vectors (default: {NUM_QUERY_VECTORS})",
    )
    parser.add_argument(
        "--n-lists",
        type=int,
        default=N_LISTS,
        help=f"Number of IVF lists (default: {N_LISTS}, 0=auto)",
    )
    parser.add_argument(
        "--n-probes",
        type=int,
        default=N_PROBES,
        help=f"Number of probes for query (default: {N_PROBES}, 0=auto)",
    )
    parser.add_argument(
        "--service-url",
        type=str,
        default=SERVICE_URL,
        help=f"CyborgDB service URL (default: {SERVICE_URL})",
    )
    parser.add_argument(
        "--api-key", type=str, default=API_KEY, help="API key for the service"
    )
    parser.add_argument(
        "--skip-train",
        action="store_true",
        help="Skip training (run untrained queries only)",
    )
    parser.add_argument(
        "--batch-queries",
        action="store_true",
        help="Also run batch queries for throughput measurement",
    )
    parser.add_argument(
        "--batch-size",
        type=int,
        default=10,
        help="Batch size for batch queries (default: 10)",
    )
    args = parser.parse_args()

    print("=" * 70)
    print("CyborgDB Service Profiling Script - Wiki Dataset")
    print("=" * 70)
    print("\nConfiguration:")
    print(f"  Dataset: {DATASET_PATH}")
    print(f"  Vectors to load: {args.vectors:,}")
    print(f"  Index type: {INDEX_TYPE.upper()}")
    print(f"  Dimension: {DIMENSION}")
    print(f"  N_lists: {args.n_lists or 'auto'}")
    print(f"  Query vectors: {args.queries}")
    print(f"  Top K: {TOP_K}")
    print(f"  N_probes: {args.n_probes or 'auto'}")
    print(f"  Service URL: {args.service_url}")
    print("=" * 70 + "\n")

    # Step 1: Load data
    train_data, test_data, neighbors = load_data(args.vectors, args.queries)
    print()

    # Step 2: Create client and verify connectivity
    print("Creating CyborgDB client...")
    client = cyborgdb.Client(base_url=args.service_url, api_key=args.api_key)

    # Check health
    try:
        health = client.get_health()
        print(f"Service health: {health}")
    except Exception as e:
        print(f"WARNING: Could not get service health: {e}")
    print()

    # Step 3: Create index
    index_name = f"profile_wiki_{int(time.time())}"
    index, index_key = create_index(client, index_name)
    print()

    try:
        # Step 4: Upsert data
        upsert_data(index, train_data)
        print()

        # Step 5: Train index (optional)
        if not args.skip_train:
            train_index(index, args.n_lists)
            print()

        # Step 6: Run queries
        single_results = run_queries(index, test_data, neighbors, args.n_probes)
        print()

        # Step 7: Run batch queries (optional)
        if args.batch_queries:
            batch_results = run_batch_queries(
                index, test_data, neighbors, args.n_probes, args.batch_size
            )
            print()

        # Summary
        print("=" * 70)
        print("FINAL SUMMARY")
        print("=" * 70)
        print("Single Query Performance:")
        print(f"  QPS: {single_results['qps']:.2f}")
        print(f"  Avg Latency: {single_results['avg_latency_ms']:.3f}ms")
        print(f"  Recall@{TOP_K}: {single_results['recall'] * 100:.2f}%")
        if args.batch_queries:
            print(f"\nBatch Query Performance (batch_size={args.batch_size}):")
            print(f"  QPS: {batch_results['qps']:.2f}")
            print(f"  Recall@{TOP_K}: {batch_results['recall'] * 100:.2f}%")
        print("=" * 70)

    finally:
        # Cleanup
        print("\nCleaning up...")
        try:
            index.delete_index()
            print("Index deleted successfully")
        except Exception as e:
            print(f"Error deleting index: {e}")

    print("Done!")


if __name__ == "__main__":
    try:
        main()
    except KeyboardInterrupt:
        print("\n\nInterrupted by user")
        sys.exit(1)
    except Exception as e:
        print(f"\n\nError: {e}")
        import traceback

        traceback.print_exc()
        sys.exit(1)
