import pandas as pd
from sentence_transformers import SentenceTransformer, util

# Load model (can be replaced with any other from sentence-transformers)
model = SentenceTransformer('all-MiniLM-L6-v2')

# Load your comments
comments_df = pd.read_csv("reddit_smb_marketing_painpoints_search_v1_results_comments_jtbd_comments.csv")
comment_col = next((c for c in comments_df.columns if 'comment' in c.lower()), comments_df.columns[0])
comments = comments_df[comment_col].fillna("").tolist()

# Load the NAICS taxonomy
naics_df = pd.read_csv("naics_2022.csv")
naics_texts = (naics_df["Title"] + ". " + naics_df["Description"]).fillna("").tolist()

# Compute embeddings
comment_embeddings = model.encode(comments, convert_to_tensor=True)
naics_embeddings = model.encode(naics_texts, convert_to_tensor=True)

# Perform semantic search (top 1 best match)
matches = util.semantic_search(comment_embeddings, naics_embeddings, top_k=1)

# Attach best match and similarity score
matched_titles = []
matched_codes = []
similarities = []

for m in matches:
    best = m[0]
    idx = best['corpus_id']
    matched_titles.append(naics_df.iloc[idx]["Title"])
    matched_codes.append(naics_df.iloc[idx]["Code"])
    similarities.append(best['score'])

comments_df["matched_industry"] = matched_titles
comments_df["naics_code"] = matched_codes
comments_df["similarity"] = similarities

comments_df.to_csv("comments_with_naics_embeddings.csv", index=False)
print("✅ Done! Saved to comments_with_naics_embeddings.csv")