import path_setup  # noqa: F401 — adds project root to sys.path

%load_ext autoreload
%autoreload 2

import pandas as pd

df_ground_truth = pd.read_csv("data/ground_truth.csv")
ground_truth = df_ground_truth.to_dict(orient="records")

from ingest import load_faq_data, build_index

documents = load_faq_data()

df_ground_truth = pd.read_csv("data/clinical_eval_set.csv")
clinical_eval_set = df_ground_truth.to_dict(orient="records")

import os
import time
import pandas as pd
from sqlitesearch import VectorSearchIndex
from embedder import Embedder
from tqdm.auto import tqdm
import numpy as np

# 1. DEFINE THE 4 TOURNAMENT CANDIDATES
model_info = {
        "name": "all-MiniLM-L6-v2",
        "path": "models/Xenova/all-MiniLM-L6-v2",
        "dim": 384,
        "notes": "Lightweight Baseline"
    }

sqlite_vector_modes = ['lsh', 'ivf', 'hnsw']

comparison_results = []

# 3. EXECUTE THE BENCHMARK LOOP
for mode in sqlite_vector_modes:
    print(f"\n--- Loading and evaluating mode: {mode} ---")

    # Initialize the model locally via sentence-transformers (easy local alternative for embedding generation)
    model = Embedder(model_info["path"])

    # Prepare database storage path
    db_path = f"sqlite/vector/evals/embedding-models/{model_info['name']}/eval_db.db"
    if os.path.exists(db_path):
        os.remove(db_path)

    # Text strings to index
    faq_texts = [f"Q: {item['question']} A: {item['answer']}" for item in documents]

    # Track indexing speed
    print("Generating Embeddings in Batches...")
    start_gen_time = time.time()

    batch_size = 50
    db_vectors = []

    for i in tqdm(range(0, len(faq_texts), batch_size)):
        batch = faq_texts[i:i + batch_size]
        batch_vectors = model.encode_batch(batch)
        db_vectors.extend(batch_vectors)

    db_vectors = np.array(db_vectors)

    gen_duration = time.time() - start_gen_time
    print(f"Embedding Generation Complete! Time taken: {gen_duration:.2f} seconds.")

    # Insert vectors into sqlitesearch DB
    print(f"Building {mode} Index in SQLite...")

    start_fit_time = time.time()
    vector_index = VectorSearchIndex(mode=mode, db_path=db_path)
    vector_index.fit(db_vectors, documents)
    fit_duration = time.time() - start_fit_time

    print(f"Database Indexing Complete! Time taken: {fit_duration:.2f} seconds.")

    # Run the query test evaluation
    hits_at_1 = 0
    total_queries = len(clinical_eval_set)
    total_latency_ms = 0

    for i in tqdm(range(0, len(clinical_eval_set))):
        test = clinical_eval_set[i]
        # Measure latency per retrieval
        start_query_time = time.time()

        # Generate query vector
        q_vector = model.encode(test["question"])

        # Query local sqlite vector database (limit=1 forces strict top rank scoring)
        results = vector_index.search(q_vector, num_results=1)

        query_duration = (time.time() - start_query_time) * 1000
        total_latency_ms += query_duration

        # Check if the top result is the correct clinical match
        if results and results[0].get("id") == test["document"]:
            hits_at_1 += 1

    # Calculate performance matrix stats
    accuracy_rate = (hits_at_1 / total_queries) * 100
    avg_latency = total_latency_ms / total_queries
    db_file_size_kb = os.path.getsize(db_path) / 1024

    # Append analytics payloads
    comparison_results.append({
        "Mode": mode,
        "Hit Rate @ 1 (Accuracy)": f"{accuracy_rate:.1f}%",
        "Avg Search Latency (ms)": f"{avg_latency:.2f}ms",
        "SQLite DB Size (KB)": f"{db_file_size_kb:.1f} KB"
    })

    # Clear the temporary DB file footprint from memory/disk
    if os.path.exists(db_path):
        vector_index.close()
        os.remove(db_path)

# 4. PRINT PERFORMANCE MATRIX TABLE
df_comparison = pd.DataFrame(comparison_results)
print("\n" + "="*80)
print("                  FIRST AID EMBEDDING MODEL TOURNAMENT MATRIX                  ")
print("="*80)
print(df_comparison.to_string(index=False))
print("="*80)

--- Loading and evaluating mode: lsh ---
Generating Embeddings in Batches...



  0%|          | 0/56 [00:00<?, ?it/s]


Embedding Generation Complete! Time taken: 98.93 seconds.
Building HNSW Index in SQLite...
Database Indexing Complete! Time taken: 1.08 seconds.



  0%|          | 0/5526 [00:00<?, ?it/s]



--- Loading and evaluating mode: ivf ---
Generating Embeddings in Batches...



  0%|          | 0/56 [00:00<?, ?it/s]


Embedding Generation Complete! Time taken: 90.34 seconds.
Building HNSW Index in SQLite...
Database Indexing Complete! Time taken: 1.47 seconds.



  0%|          | 0/5526 [00:00<?, ?it/s]



--- Loading and evaluating mode: hnsw ---
Generating Embeddings in Batches...



  0%|          | 0/56 [00:00<?, ?it/s]


Embedding Generation Complete! Time taken: 86.77 seconds.
Building HNSW Index in SQLite...
Database Indexing Complete! Time taken: 9.84 seconds.



  0%|          | 0/5526 [00:00<?, ?it/s]



================================================================================
                  FIRST AID EMBEDDING MODEL TOURNAMENT MATRIX                  
================================================================================
Mode Hit Rate @ 1 (Accuracy) Avg Search Latency (ms) SQLite DB Size (KB)
 lsh                   12.3%                 31.02ms           8120.0 KB
 ivf                   15.6%                 28.02ms           6244.0 KB
hnsw                   15.5%                 27.12ms           6356.0 KB
================================================================================

df_comparison.to_csv("data/sqlite-vector-modes-comparison.csv", index=False)