# jtbd_cluster_evaluator.py

import pandas as pd
import numpy as np
from textblob import TextBlob
from sentence_transformers import SentenceTransformer, util

def compute_scores(file_path, output_path="scored_jtbd_clusters.csv"):
    # Load the CSV
    df = pd.read_csv(file_path)

    # Check required columns
    required_cols = {'situation', 'struggle', 'outcome', 'comment'}
    if not required_cols.issubset(df.columns):
        raise ValueError(f"Input CSV must contain columns: {required_cols}")

    # Load sentence transformer model
    print("Loading embedding model...")
    model = SentenceTransformer('all-MiniLM-L6-v2')

    # Encode situation, struggle, outcome
    print("Embedding text...")
    s_emb = model.encode(df['situation'].tolist(), convert_to_tensor=True)
    str_emb = model.encode(df['struggle'].tolist(), convert_to_tensor=True)
    o_emb = model.encode(df['outcome'].tolist(), convert_to_tensor=True)

    # Coherence = mean cosine similarity between S, Str, O
    print("Computing coherence scores...")
    coherence_scores = []
    for s, strg, o in zip(s_emb, str_emb, o_emb):
        sim1 = util.cos_sim(s, strg).item()
        sim2 = util.cos_sim(s, o).item()
        sim3 = util.cos_sim(strg, o).item()
        coherence_scores.append(np.mean([sim1, sim2, sim3]))

    # Specificity = mean token length of situation + struggle + outcome
    print("Computing specificity scores...")
    specificity_scores = df[['situation', 'struggle', 'outcome']].applymap(lambda x: len(str(x).split()))
    specificity_scores = specificity_scores.mean(axis=1)

    # Emotional intensity = sentiment polarity magnitude from comment
    print("Computing emotional intensity...")
    intensity_scores = df['comment'].apply(lambda x: abs(TextBlob(str(x)).sentiment.polarity))

    # Completeness = 1 if all fields are non-empty
    completeness_scores = df[['situation', 'struggle', 'outcome']].notnull().all(axis=1).astype(int)

    # Normalize features to [0, 1]
    def normalize(series):
        return (series - series.min()) / (series.max() - series.min())

    coherence_norm = normalize(pd.Series(coherence_scores))
    specificity_norm = normalize(specificity_scores)
    intensity_norm = normalize(intensity_scores)

    # Final JTBD Value Score
    jtbd_score = (
        0.3 * coherence_norm +
        0.2 * specificity_norm +
        0.2 * intensity_norm +
        0.3 * completeness_scores
    )

    # Add results to DataFrame
    df['coherence'] = coherence_scores
    df['specificity'] = specificity_scores
    df['intensity'] = intensity_scores
    df['completeness'] = completeness_scores
    df['jtbd_value_score'] = jtbd_score

    # Save to CSV
    df.to_csv(output_path, index=False)
    print(f"\n✅ Scored clusters saved to: {output_path}")

if __name__ == "__main__":
    import argparse
    parser = argparse.ArgumentParser()
    parser.add_argument("input", help="Path to your JTBD cluster CSV file")
    parser.add_argument("--output", help="Where to save the scored output", default="scored_jtbd_clusters.csv")
    args = parser.parse_args()

    compute_scores(args.input, args.output)