import os
import numpy as np
import faiss

EXP = "josias_v1"
ROOT = os.path.dirname(os.path.abspath(__file__))
LOGDIR = os.path.join(ROOT, "logs", EXP)

fea_path = os.path.join(LOGDIR, "total_fea.npy")
assert os.path.exists(fea_path), f"Missing: {fea_path}"

x = np.load(fea_path)
# Asegurar float32 y 2D
x = np.asarray(x, dtype=np.float32)
if x.ndim == 1:
    x = x.reshape(-1, 1)

n, d = x.shape
print("Loaded features:", (n, d), x.dtype)

if n == 0:
    raise RuntimeError("total_fea.npy is empty (0 rows). Cannot build index.")

# Parámetros FAISS (compatibles con nombres típicos)
nlist = 1058          # igual a tu nombre IVF1058
nprobe = 1            # igual a nprobe_1

# Entrenamiento del cuantizador IVF
quantizer = faiss.IndexFlatL2(d)
index = faiss.IndexIVFFlat(quantizer, d, nlist, faiss.METRIC_L2)

print("Training IVF index...")
index.train(x)

print("Adding vectors...")
index.add(x)

index.nprobe = nprobe

out_index = os.path.join(LOGDIR, f"added_IVF{nlist}_Flat_nprobe_{nprobe}_{EXP}_v2.index")
faiss.write_index(index, out_index)

print("Wrote index:", out_index)
print("ntotal =", index.ntotal)