#!/usr/bin/env python3
"""
embedding_engine.py

Pipeline for generating vector embeddings from PubMed abstracts using the Ollama embedding model
and building a FAISS similarity index for retrieval in RAG (Retrieval-Augmented Generation) workflows.

This module supports query-specific subfolders, enabling multiple case studies to coexist without 
overwriting existing data. It is designed to work seamlessly with the `rag_data_loader.py` pipeline 
for fetching PubMed articles.

Features:
1. Load abstracts from JSON files generated by `rag_data_loader.py`.
   - Skips empty abstracts.
   - Supports query-specific subfolders for reproducibility.

2. Generate embeddings for each abstract using Ollama:
   - Handles API failures gracefully by inserting zero vectors.
   - Checkpoints progress to allow resuming embedding generation.
   - Supports batch processing and optional delay between calls to avoid API overload.

3. Build a FAISS index:
   - Inner product (IP) index for cosine similarity.
   - Normalizes embeddings to unit length.
   - Supports GPU acceleration if configured; falls back to CPU if unavailable.

4. Save FAISS index and PMID mapping:
   - Stores FAISS index and a JSON mapping of vector indices to PubMed IDs.
   - Query-specific folders are used for organized storage.

Configuration:
- Uses `ragbio.config` for paths, model names, and GPU usage:
    - ABSTRACT_FOLDER : Base folder containing abstract JSON files.
    - INDEX_FOLDER    : Base folder to store FAISS indices and mapping files.
    - MODEL_NAME      : Ollama embedding model name.
    - USE_GPU         : Flag to enable GPU acceleration for FAISS.

Checkpointing:
- Saves embedding checkpoint in JSON format under the query-specific index folder.
- Checkpoints are updated every 500 abstracts to reduce API strain and allow resumption.

Command-line Arguments:
- --study : Optional string to specify a subfolder for abstracts and index files.
                 Useful for running multiple RAG case studies in parallel.

Dependencies:
- Python >=3.8
- numpy
- faiss (CPU or GPU)
- ollama
- tqdm
- ragbio.config
- json, os, time, argparse

Usage Example:
    python -m ragbio.embeddings.embedding_engine --study "Alzheimer_CaseStudy"

Author:
    Manish Kumar
"""


import os
import json
import time
import faiss
import numpy as np
import ollama
from tqdm import tqdm
import argparse
from ragbio import config

# ==========================================
# Argument parser
# ==========================================
def parse_args():
    """
    Parse command-line arguments.

    Returns:
        argparse.Namespace: Contains optional study.
    """
    parser = argparse.ArgumentParser(description="Generate embeddings and FAISS index for PubMed abstracts")
    parser.add_argument(
        "--study",
        type=str,
        default=None,
        help="Optional query-specific folder name under ABSTRACT_FOLDER and INDEX_FOLDER"
    )
    return parser.parse_args()

# ==========================================
# Load abstracts
# ==========================================
def load_abstracts(study=None):
    """
    Load abstracts from JSON files, skipping empty ones.

    Args:
        study (str, optional): Name of query-specific subfolder.

    Returns:
        abstracts (list[str]): List of abstract texts.
        pmids (list[str]): Corresponding PMIDs.
    """
    folder = config.ABSTRACT_FOLDER
    if study:
        folder = os.path.join(folder, study)

    abstracts, pmids = [], []
    for file in os.listdir(folder):
        if file.endswith(".json"):
            with open(os.path.join(folder, file), "r", encoding="utf-8") as f:
                data = json.load(f)
                text = data.get("abstract", "").strip()
                if text:
                    abstracts.append(text)
                    pmids.append(data["pmid"])

    print(f"[INFO] Loaded {len(abstracts)} abstracts from {folder}.")
    return abstracts, pmids

# ==========================================
# Generate embeddings
# ==========================================
def generate_embeddings(texts, model_name, batch_size=8, checkpoint_path=None):
    """
    Generate embeddings using Ollama with checkpoint support.

    Args:
        texts (list[str]): List of abstracts.
        model_name (str): Ollama embedding model.
        batch_size (int): Number of abstracts per batch (unused here but placeholder).
        checkpoint_path (str): Optional JSON checkpoint file.

    Returns:
        np.ndarray: Embeddings as float32 array.
    """
    embeddings = []
    start_idx = 0

    if checkpoint_path and os.path.exists(checkpoint_path):
        with open(checkpoint_path, "r") as f:
            saved = json.load(f)
        start_idx = len(saved.get("embeddings", []))
        embeddings = saved["embeddings"]
        print(f"[RESUME] Resuming from index {start_idx}")

    for i in tqdm(range(start_idx, len(texts)), desc="Generating embeddings"):
        text = texts[i]
        try:
            result = ollama.embeddings(model=model_name, prompt=text)
            emb = result["embedding"]
        except Exception as e:
            print(f"[WARN] Embedding failed for index {i}: {e}")
            emb = [0] * 768
        embeddings.append(emb)

        # Save checkpoint every 500 embeddings
        if checkpoint_path and (i + 1) % 500 == 0:
            os.makedirs(os.path.dirname(checkpoint_path), exist_ok=True)  # <-- create folder
            with open(checkpoint_path, "w") as f:
                json.dump({"embeddings": embeddings}, f)
            print(f"[CHECKPOINT] Saved {i + 1} embeddings")

        time.sleep(0.1)

    if checkpoint_path:
        os.makedirs(os.path.dirname(checkpoint_path), exist_ok=True)  # <-- ensure folder exists
        with open(checkpoint_path, "w") as f:
            json.dump({"embeddings": embeddings}, f)
    return np.array(embeddings, dtype="float32")

# ==========================================
# Build FAISS index
# ==========================================
def build_faiss_index(embeddings):
    """Build FAISS index (GPU if available) and normalize embeddings."""
    faiss.normalize_L2(embeddings)
    dim = embeddings.shape[1]

    if config.USE_GPU:
        try:
            print("[INFO] Using GPU FAISS index")
            res = faiss.StandardGpuResources()
            index_flat = faiss.IndexFlatIP(dim)
            index = faiss.index_cpu_to_gpu(res, 0, index_flat)
        except Exception as e:
            print(f"[WARN] GPU FAISS failed, using CPU. Reason: {e}")
            index = faiss.IndexFlatIP(dim)
    else:
        print("[INFO] Using CPU FAISS index")
        index = faiss.IndexFlatIP(dim)

    index.add(embeddings)
    print(f"[INFO] FAISS index built with {index.ntotal} vectors.")
    return index

# ==========================================
# Save index and PMID map
# ==========================================
def save_index(index, pmids, study=None):
    """
    Save FAISS index and corresponding PMID map.

    Args:
        index (faiss.Index): FAISS index object.
        pmids (list[str]): List of PMIDs.
        study (str, optional): Query-specific subfolder.
    """
    folder = config.INDEX_FOLDER
    if study:
        folder = os.path.join(folder, study)
    os.makedirs(folder, exist_ok=True)

    faiss.write_index(index, os.path.join(folder, "pubmed_index.faiss"))
    with open(os.path.join(folder, "pmid_map.json"), "w", encoding="utf-8") as f:
        json.dump(pmids, f, indent=2)

    print(f"[OK] Index saved → {folder}/pubmed_index.faiss")
    print(f"[OK] PMID map saved → {folder}/pmid_map.json")

# ==========================================
# Main
# ==========================================
def main():
    args = parse_args()
    abstracts, pmids = load_abstracts(args.study)
    if not abstracts:
        print("[ERROR] No abstracts found. Run rag_data_loader.py first.")
        return

    checkpoint_file = os.path.join(
        config.INDEX_FOLDER,
        args.study if args.study else "",
        "embedding_checkpoint.json"
    )
    embeddings = generate_embeddings(abstracts, config.MODEL_NAME, checkpoint_path=checkpoint_file)
    index = build_faiss_index(embeddings)
    save_index(index, pmids, args.study)
    print("[DONE] Embedding pipeline complete.")

if __name__ == "__main__":
    main()
