import numpy as np
from dotenv import load_dotenv
from openai import OpenAI

load_dotenv()
client = OpenAI()


def get_embedding(text):
    response = client.embeddings.create(input=text, model="text-embedding-ada-002")
    return np.array(response.data[0].embedding)


documents = [
    "Python은 강력한 프로그래밍 언어입니다.",
    "OpenAI는 AI 연구를 선도하는 기업입니다.",
    "FAISS는 벡터 검색을 위한 라이브러리입니다.",
]
doc_embeddings = np.array([get_embedding(doc) for doc in documents], dtype=np.float32)


# Cosine Similarity(A,B) = A·B / (∥A∥ × ∥B∥)
def cosine_similarity(a, b):
    dot = np.dot(a, b)
    norm = np.linalg.norm(a) * np.linalg.norm(b)
    return dot / norm


def top_k_search(query_vec, doc_vecs, k=2):
    sims = [cosine_similarity(query_vec, d) for d in doc_vecs]
    ranked = np.argsort(sims)[::-1][:k]
    return [(i, sims[i]) for i in ranked]


# query = "프로그래밍이란?"
# query = "오늘 저녁에 뭐 먹을까?"   # 무관한 질문 → 유사도가 전부 낮게 나온다
query = "검색 라이브러리가 뭐야?"
query_vec = get_embedding(query)

results = top_k_search(query_vec, doc_embeddings, k=2)  # k=1, k=2, k=3
for idx, score in results:
    print(f"문서 {idx}: 유사도 {score:.4f} -> {documents[idx]}")
