#!/usr/bin/env python3
"""
Identity tagging for user comments (rule-based + embeddings).

Usage:
    python identify_identities.py budgetbakers_pain_points.csv
"""

import sys
import pandas as pd
import numpy as np
import re
from openai import OpenAI
from sklearn.metrics.pairwise import cosine_similarity

# ---------------------------
# Config
# ---------------------------
MODEL = "text-embedding-3-small"
client = OpenAI()

# Expanded prototypes for embedding similarity
IDENTITY_PROTOTYPES = {
    "Parent": [
        "My kids are expensive to feed.",
        "As a parent, it’s hard to budget.",
        "My family expenses keep piling up.",
        "Feeding my children costs a lot.",
        "My daughter’s school lunch is expensive.",
        "Childcare takes up most of my budget.",
        "Raising kids makes saving money difficult."
    ],
    "Student": [
        "In college, I struggle to manage money.",
        "Studying makes it hard to cook at home.",
        "Campus food is expensive.",
        "As a student, I can’t afford much.",
        "Living in a dorm means limited cooking.",
        "Exam stress makes me eat out more.",
        "On a student budget, it’s hard to plan meals."
    ],
    "Worker": [
        "At the office cafeteria, I overspend.",
        "My job makes meal planning difficult.",
        "During work hours, I just grab fast food.",
        "In the office, I can’t manage healthy eating.",
        "My boss expects long hours, so I skip meals.",
        "At work, I eat quickly between meetings.",
        "My career leaves me little time for cooking."
    ]
}

# Rule-based keyword sets
IDENTITY_KEYWORDS = {
    "Parent": ["kid", "kids", "child", "children", "daughter", "son", "family", "mom", "dad", "parent"],
    "Student": ["college", "university", "school", "class", "exam", "study", "studying", "semester", "dorm", "campus"],
    "Worker": ["work", "job", "office", "boss", "coworker", "shift", "career", "freelance", "workplace", "salary"]
}


# ---------------------------
# Helpers
# ---------------------------
def get_embeddings(texts, model=MODEL, batch_size=100):
    """
    Get embeddings from OpenAI in batches.
    Returns a list of embedding vectors.
    """
    embeddings = []
    for i in range(0, len(texts), batch_size):
        batch = texts[i:i+batch_size]
        response = client.embeddings.create(
            model=model,
            input=batch
        )
        embeddings.extend([d.embedding for d in response.data])
    return embeddings


def detect_identity_rule(text):
    """
    Simple rule-based identity detection.
    """
    text = text.lower()
    for identity, keywords in IDENTITY_KEYWORDS.items():
        for kw in keywords:
            if re.search(rf"\b{kw}\b", text):
                return identity
    return "Unclear"


# ---------------------------
# Main pipeline
# ---------------------------
def main(input_file):
    # Load data
    df = pd.read_csv(input_file)
    df["comment"] = df["comment"].astype(str).fillna("")

    # Rule-based detection
    df["identity_rule"] = df["comment"].apply(detect_identity_rule)

    # Build prototype embeddings
    proto_sentences = []
    proto_labels = []
    for label, examples in IDENTITY_PROTOTYPES.items():
        for ex in examples:
            proto_sentences.append(ex)
            proto_labels.append(label)

    proto_embeddings = get_embeddings(proto_sentences)

    # Compute comment embeddings
    comment_embeddings = get_embeddings(df["comment"].tolist())

    # Compute similarity
    sims = cosine_similarity(comment_embeddings, proto_embeddings)

    # Assign identity by embedding similarity
    best_idx = sims.argmax(axis=1)
    best_scores = sims.max(axis=1)
    df["identity_embedding"] = [proto_labels[i] for i in best_idx]

    # Apply confidence threshold
    threshold = 0.75
    df.loc[best_scores < threshold, "identity_embedding"] = "Unclear"
    

    # Save results
    output_file = input_file.replace(".csv", "_with_identity.csv")
    df.to_csv(output_file, index=False)
    print(f"✅ Saved with identities → {output_file}")

    # Show quick distribution
    print("\nRule-based distribution:")
    print(df["identity_rule"].value_counts())
    print("\nEmbedding-based distribution:")
    print(df["identity_embedding"].value_counts())


if __name__ == "__main__":
    if len(sys.argv) < 2:
        print("Usage: python identify_identities.py <input_file.csv>")
        sys.exit(1)
    main(sys.argv[1])