# src/utils/index_utils.py

import faiss
import json
import numpy as np
from typing import Union
from pathlib import Path

def normalize(embeddings):
    """Normalize the embeddings to unit vectors."""
    norms = np.linalg.norm(embeddings, axis=1, keepdims=True)
    return embeddings / np.clip(norms, a_min=1e-10, a_max=None)

def save_faiss_index(features: np.ndarray, file_path):
    """
    Builds a FAISS index from the given feature matrix and saves it to disk.
    """
    # Ensure features is float32
    features = np.asarray(features).astype(np.float32)
    
    # Create the index
    dim = features.shape[1]
    index = faiss.IndexFlatL2(dim)
    index.add(features)

    # Save the index
    faiss.write_index(index, str(file_path))
    print(f"[✓] FAISS index saved to {file_path}")
    
def load_features(path: Path) -> np.ndarray:
    assert path.exists(), f"Feature file not found at: {path}"
    return np.load(path, allow_pickle=True)

def save_default_config(dataset: str, weights: dict, filename: str = "default_weights.json") -> None:
    """
    Saves the default fusion weights to a JSON file.
    """
    path = Path(f"defaults/{dataset}")
    path.mkdir(parents=True, exist_ok=True)
    filepath = path / filename
    with open(filepath, "w") as f:
        json.dump(weights, f, indent=2)
    print(f"[✓] Saved default weights to {filepath}")