from huggingface_hub import hf_hub_download, InferenceClient
from FlagEmbedding import FlagICLModel
import pandas as pd
import faiss
import numpy as np
import os
import json

# Define HF tokens
HF_TOKEN_read = os.environ.get("HF_TOKEN_read")
HF_TOKEN_inference = os.environ.get("HF_TOKEN_inf")

# Setup InferenceClient for embeddings
client = InferenceClient(provider="nebius", api_key=HF_TOKEN_inference)

# Dataset repo (private)
DATASET_REPO = "luciagomez/MrPhil_vector"

# -------------------------------------------------------------------
# 1. Download files from Hugging Face dataset
# -------------------------------------------------------------------
parquet_path = hf_hub_download(
    repo_id=DATASET_REPO,
    filename="foundations.parquet",
    repo_type="dataset",
    token=HF_TOKEN_read
)

faiss_path = hf_hub_download(
    repo_id=DATASET_REPO,
    filename="faiss.index",
    repo_type="dataset",
    token=HF_TOKEN_read
)

meta_path = hf_hub_download(
    repo_id=DATASET_REPO,
    filename="meta.json",
    repo_type="dataset",
    token=HF_TOKEN_read
)


# -------------------------------------------------------------------
# 2. Load data
# -------------------------------------------------------------------
df = pd.read_parquet(parquet_path)
index = faiss.read_index(faiss_path)
with open(meta_path, "r") as f:
    meta = json.load(f)

dim = meta.get("embedding_dim", None)
n = meta.get("num_vectors", len(df))

print(f"Loaded FAISS index with {n} vectors of dimension {dim}")


# -------------------------------------------------------------------
# 3. Initialize BGE-ICL model for queries
# -------------------------------------------------------------------
examples = [
            {
                "instruct": "Retrieve foundations whose mission aligns with the given perspective.",
                "query": "Protect marine life while educating children about ocean conservation",
                "response": "Foundations working on marine conservation and youth education."
            },
            {
                "instruct": "Retrieve foundations whose mission aligns with the given perspective.",
                "query": "Promote renewable energy education and community awareness",
                "response": "Foundations focused on clean energy advocacy and public education."
            }
        ]


#model = FlagICLModel(
#    "BAAI/bge-en-icl",
#    query_instruction_for_retrieval="Given a mission statement, retrieve foundations with aligned purposes.",
#    examples_for_task=examples,
#    use_fp16=False
#)

def encode_query(perspective):
    payload = {"model": "BAAI/bge-en-icl", 
               "inputs": perspective, 
               "parameters": {"instruction": "Retrieve foundations aligned with perspective.", 
                              "examples": examples}}
    response = client.feature_extraction(**payload)
    return np.array(response)

# -------------------------------------------------------------------
# 4. Retrieval function
# -------------------------------------------------------------------
def find_similar_foundations(perspective, top_k=5):
    q_emb = encode_query(perspective).astype("float32")
    faiss.normalize_L2(q_emb)
    scores, idxs = index.search(q_emb, top_k)
    return [
        {"title": df.iloc[i]["Title"], "purpose": df.iloc[i]["Purpose"], "similarity": float(scores[0][j])}
        for j, i in enumerate(idxs[0])
    ]