#!/usr/bin/env python3
import sys
import pandas as pd
import spacy
import json
from collections import Counter

# load spaCy
nlp = spacy.load("en_core_web_sm")

# Simple emotion lexicon
EMOTIONS = {
    "anxiety", "anxious", "stress", "stressed", "sad", "sadness", "hopeless", "hopelessness",
    "guilt", "guilty", "regret", "regretting", "fear", "afraid", "angry", "anger",
    "frustrated", "frustration", "depression", "lonely", "loneliness", "useless", "uselessness"
}

def extract_contexts(doc):
    contexts = []
    for token in doc:
        if token.dep_ == "pobj" and token.head.dep_ == "prep":
            if token.head.text.lower() in {"in", "at", "on", "before", "after", "during", "while"}:
                phrase = f"{token.head.text} {token.text}"
                contexts.append(phrase.lower())
    return list(set(contexts))

def extract_feelings(doc):
    words = {t.lemma_.lower() for t in doc}
    return [w for w in words if w in EMOTIONS]

def extract_objects(doc):
    stop_nouns = {"thing", "way", "medium", "day", "time"}
    return list({t.lemma_.lower() for t in doc if t.pos_ == "NOUN" and t.lemma_.lower() not in stop_nouns})

def extract_behaviors(doc):
    stop_verbs = {"have", "do", "get", "go", "make", "be", "say"}
    return list({t.lemma_.lower() for t in doc if t.pos_ == "VERB" and t.lemma_.lower() not in stop_verbs})

def process_comment(text):
    doc = nlp(text)
    return {
        "comment": text,
        "contexts": extract_contexts(doc),
        "feelings": extract_feelings(doc),
        "objects": extract_objects(doc),
        "behaviors": extract_behaviors(doc)
    }

def main(input_file, output_json, output_csv):
    df = pd.read_csv(input_file)
    if "comment" not in df.columns:
        raise ValueError("CSV must contain a 'comment' column")

    unique_comments = df["comment"].dropna().drop_duplicates().tolist()
    print(f"Processing {len(unique_comments)} unique comments...")

    results = []
    ctx_counter, feel_counter, obj_counter, beh_counter = Counter(), Counter(), Counter(), Counter()
    co_obj_feel, co_ctx_feel, co_beh_obj = Counter(), Counter(), Counter()

    for comment in unique_comments:
        parsed = process_comment(comment)
        results.append(parsed)

        ctx_counter.update(parsed["contexts"])
        feel_counter.update(parsed["feelings"])
        obj_counter.update(parsed["objects"])
        beh_counter.update(parsed["behaviors"])

        for o in parsed["objects"]:
            for f in parsed["feelings"]:
                co_obj_feel[(o, f)] += 1
        for c in parsed["contexts"]:
            for f in parsed["feelings"]:
                co_ctx_feel[(c, f)] += 1
        for b in parsed["behaviors"]:
            for o in parsed["objects"]:
                co_beh_obj[(b, o)] += 1

    # save structured JSON
    with open(output_json, "w", encoding="utf-8") as f:
        json.dump(results, f, indent=2, ensure_ascii=False)
    print(f"Saved structured extraction → {output_json}")

    # build flat table for CSV
    rows = []
    for label, counter in [
        ("Context", ctx_counter),
        ("Feeling", feel_counter),
        ("Object", obj_counter),
        ("Behavior", beh_counter),
    ]:
        for item, count in counter.items():
            rows.append({"category": label, "item1": item, "item2": "", "count": count})

    for (o, f), count in co_obj_feel.items():
        rows.append({"category": "Object×Feeling", "item1": o, "item2": f, "count": count})
    for (c, f), count in co_ctx_feel.items():
        rows.append({"category": "Context×Feeling", "item1": c, "item2": f, "count": count})
    for (b, o), count in co_beh_obj.items():
        rows.append({"category": "Behavior×Object", "item1": b, "item2": o, "count": count})

    pd.DataFrame(rows).to_csv(output_csv, index=False)
    print(f"Saved analysis summary → {output_csv}")

if __name__ == "__main__":
    if len(sys.argv) < 4:
        print("Usage: python doomscroll_extractor.py <input.csv> <output.json> <output.csv>")
        sys.exit(1)

    main(sys.argv[1], sys.argv[2], sys.argv[3])