import json
import numpy as np
import faiss
import os

def create_faiss_index_from_json():
    # مسیر فایل‌ها
    json_file_path = './nahj_data/nahj_vector_bge_m3.json'
    faiss_index_path = './data-faiss/faiss_index_nahj.index'
    metadata_file_path = './data-faiss/faiss_index_nahj_metadata.json'

    # --- 1. بارگذاری داده‌ها از JSON ---
    with open(json_file_path, 'r', encoding='utf-8') as f:
        data = json.load(f)

    ids = []
    part_ids = []
    context_ids = []
    large_titles = []
    norm_sentences = []
    titles = []
    urls = []
    embeddings_list = []
    arabic_texts = []
    Interpretation_links = []
    types = []

    for item in data:
        # print(k)
        # print(item)
        ids.append(item['id'])
        part_ids.append(item['part_id'])
        context_ids.append(item['context_id'])
        large_titles.append(item['large_title'])
        norm_sentences.append(item['sentence'])
        urls.append(item['url'])
        titles.append(item['title'])
        embeddings_list.append(item['embeddings'])
        arabic_texts.append(item['arabic_text'])
        Interpretation_links.append(item['Interpretation_link'])
        types.append(item['type'])

    embeddings = np.array(embeddings_list).astype('float32')  # ابعاد: (n, d)
    # نرمال‌سازی برای cosine similarity
    faiss.normalize_L2(embeddings)
    dimension = embeddings.shape[1]
    print(f"Loaded {len(embeddings)} embeddings with dimension {dimension}")

    # --- 2. ایجاد ایندکس FAISS برای GPU ---
    # اگر فقط CPU دارید، از faiss.IndexFlatL2 استفاده کنید.
    # اگر GPU دارید، ابتدا ایندکس را روی CPU ایجاد و سپس به GPU انتقال دهید.
    # cpu_index = faiss.IndexFlatL2(dimension)  # معیار فاصله L2 (Euclidean)
    cpu_index = faiss.IndexFlatIP(dimension)
    
    # انتقال ایندکس به GPU
    if faiss.get_num_gpus() > 0:
        print("Using GPU for FAISS index...")
        res = faiss.StandardGpuResources()
        gpu_index = faiss.index_cpu_to_gpu(res, 0, cpu_index)
    else:
        print("GPU not available, using CPU.")
        gpu_index = cpu_index

    # --- 3. افزودن داده‌ها به ایندکس ---
    gpu_index.add(embeddings)
    print(f"Total vectors indexed: {gpu_index.ntotal}")

    # --- 4. ذخیره ایندکس به فایل ---
    # برای ذخیره باید به CPU منتقل شود
    final_index = faiss.index_gpu_to_cpu(gpu_index) if isinstance(gpu_index, faiss.Index) and faiss.get_num_gpus() > 0 else gpu_index
    os.makedirs(os.path.dirname(faiss_index_path), exist_ok=True)
    faiss.write_index(final_index, faiss_index_path)
    print(f"FAISS index saved to {faiss_index_path}")

    # --- 5. ذخیره متادیتا (برای نگاشت نتایج جستجو) ---
    metadata = [{
        "id": i,
        "context_id" : c,
        "part_id" : p,
        "title": t,
        "large_title": l,
        "normalized_sentence": ns,
        "url": u,
        "types" : ty,
        "arabic_text" : at,
        "Interpretation_links" : il
        } for i, c, p, l, t, ns, u, ty, at, il in zip(ids, context_ids, part_ids, large_titles, titles, norm_sentences, urls, types, arabic_texts, Interpretation_links)]
    # metadata = [{ "sentence": s} for s in zip(titles, sentences)]
    with open(metadata_file_path, 'w', encoding='utf-8') as f:
        json.dump(metadata, f, ensure_ascii=False, indent=2)
    print(f"Metadata saved to {metadata_file_path}")

if __name__ == '__main__':
    # استفاده از متد
    # json_file_path = './output-speechs/embedding_FastText_khamenei.json'  # مسیر فایل JSON
    # index_file_path = './data-faiss/faiss_index_khamenei.index'  # مسیر فایل ایندکس FAISS
    # index_file_path = './data-faiss/faiss_index_khamenei_metadata.json'  # مسیر فایل متادیتا FAISS

    create_faiss_index_from_json()