#!/usr/bin/env python3
# ============================================================
# extract_audience_psychology.py
# ============================================================
# Extracts psychological profiles of the target audience from
# social media comments, segmented by problem cluster.
#
# The idea: Different problems attract different psychologies.
# Someone struggling with "phone addiction affecting sleep" has
# different fears/motivations than "missing real life moments".
#
# This script:
#   1. Loads problem clusters (from cluster_canonical_problems_and_match_comments.py)
#   2. For each cluster, samples matched comments
#   3. Extracts psychological dimensions from those comments
#   4. Synthesizes a profile per problem cluster
#
# Output: Per-cluster profiles showing what psychological levers
# resonate with each problem segment.
#
# Usage:
#   python extract_audience_psychology.py canonical_problem_clusters.json \
#       --comments_path comments.csv --samples_per_cluster 30
# ============================================================

import os
import sys
import json
import re
import pandas as pd
import numpy as np
from pathlib import Path
from collections import Counter, defaultdict
from concurrent.futures import ThreadPoolExecutor, as_completed
from openai import OpenAI

client = OpenAI()

# ============================================================
# COMMENT PSYCHOLOGY EXTRACTION
# ============================================================

def extract_comment_psychology(comment_text: str) -> dict:
    """
    Extract psychological dimensions from a single comment.
    Returns traits about the commenter's psychology, not the comment's persuasion.
    """
    prompt = f"""Analyze this social media comment to extract the commenter's psychological profile.

Comment: "{comment_text}"

Extract the following dimensions (return "None" if not present):

1. **Emotional State**: What emotions is this person experiencing or expressing? (e.g., frustration, hope, skepticism, excitement, fear, resignation)

2. **Core Motivation**: What underlying need or desire drives this comment? (e.g., seeking validation, looking for solutions, expressing identity, social connection)

3. **Pain Points**: What struggles or problems does this comment reveal? Be specific about the user's experienced difficulty.

4. **Decision Barriers**: What objections, skepticism, or resistance does this comment show? What would stop them from taking action?

5. **Aspirations**: What does this person want to achieve or become? What positive outcome are they seeking?

6. **Identity Signals**: How does this person see themselves or want to be seen? What group do they identify with?

7. **Information Seeking**: What questions or gaps in understanding does this comment reveal?

8. **Social Dynamics**: Does this comment show social proof seeking, peer influence, or conformity pressure?

Respond with JSON only:
{{
  "emotional_state": "...",
  "core_motivation": "...",
  "pain_points": "...",
  "decision_barriers": "...",
  "aspirations": "...",
  "identity_signals": "...",
  "information_seeking": "...",
  "social_dynamics": "..."
}}

Be specific and grounded in the actual comment text. If a dimension isn't clearly present, return "None"."""

    try:
        response = client.chat.completions.create(
            model="gpt-4.1-mini",
            messages=[
                {"role": "system", "content": "You are a consumer psychologist analyzing social media comments to understand audience psychology."},
                {"role": "user", "content": prompt}
            ],
            temperature=0.2,
        )
        content = response.choices[0].message.content.strip()
        content = re.sub(r"```(?:json)?|```", "", content).strip()
        return json.loads(content)
    except Exception as e:
        print(f"Error extracting psychology: {e}")
        return {k: None for k in [
            "emotional_state", "core_motivation", "pain_points",
            "decision_barriers", "aspirations", "identity_signals",
            "information_seeking", "social_dynamics"
        ]}


# ============================================================
# PROFILE AGGREGATION
# ============================================================

def aggregate_to_profile(psychology_list: list[dict]) -> dict:
    """
    Aggregate individual comment psychology into an audience profile.
    """
    dimensions = [
        "emotional_state", "core_motivation", "pain_points",
        "decision_barriers", "aspirations", "identity_signals",
        "information_seeking", "social_dynamics"
    ]

    aggregated = {}
    for dim in dimensions:
        values = [p.get(dim, "") for p in psychology_list if p.get(dim) and p.get(dim).lower() != "none"]
        aggregated[dim] = values

    return aggregated


def synthesize_profile(aggregated: dict, sample_size: int) -> dict:
    """
    Use LLM to synthesize aggregated data into a coherent psychological profile.
    """
    # Format the aggregated data for the prompt
    dim_summaries = []
    for dim, values in aggregated.items():
        if values:
            # Take a sample if too many
            sample = values[:50] if len(values) > 50 else values
            dim_summaries.append(f"**{dim.replace('_', ' ').title()}** ({len(values)} mentions):\n" +
                               "\n".join(f"- {v}" for v in sample[:20]))

    aggregated_text = "\n\n".join(dim_summaries)

    prompt = f"""Based on analyzing {sample_size} social media comments, here are the extracted psychological dimensions:

{aggregated_text}

Synthesize this into a unified PSYCHOLOGICAL PROFILE of the target audience. Structure your response as:

1. **Dominant Emotional Landscape**: What emotional states characterize this audience? What are they feeling?

2. **Core Motivations**: What fundamentally drives this audience? What do they want at a deep level?

3. **Primary Pain Points**: What are the 3-5 most common struggles this audience experiences?

4. **Key Decision Barriers**: What objections or resistance patterns would prevent action? What makes them skeptical?

5. **Aspirational Identity**: Who does this audience want to become? What transformation do they seek?

6. **Persuasion Vulnerabilities**: Based on this profile, what psychological levers would be most effective for persuasion? What would resonate?

7. **Messaging Landmines**: What should persuasion AVOID? What would trigger resistance?

Be specific and actionable. This profile will be used to craft persuasion text."""

    try:
        response = client.chat.completions.create(
            model="gpt-4.1",  # Use stronger model for synthesis
            messages=[
                {"role": "system", "content": "You are a consumer psychologist creating an actionable audience profile for marketing strategy."},
                {"role": "user", "content": prompt}
            ],
            temperature=0.3,
        )
        return {
            "profile_text": response.choices[0].message.content.strip(),
            "sample_size": sample_size,
            "dimension_counts": {k: len(v) for k, v in aggregated.items()}
        }
    except Exception as e:
        print(f"Error synthesizing profile: {e}")
        return {"error": str(e)}


# ============================================================
# MAIN
# ============================================================

def run_extraction_by_cluster(
    clusters_path: str,
    comments_path: str,
    text_col: str = "text",
    samples_per_cluster: int = 30,
    min_cluster_comments: int = 10,
    output_path: str = None,
    max_workers: int = 10
):
    print("=" * 60)
    print("AUDIENCE PSYCHOLOGICAL PROFILE BY PROBLEM CLUSTER")
    print("=" * 60)

    # Load clusters
    print(f"\nLoading clusters from {clusters_path}...")
    with open(clusters_path) as f:
        clusters_data = json.load(f)

    # Load comments for lookup
    print(f"Loading comments from {comments_path}...")
    ext = Path(comments_path).suffix.lower()
    if ext == ".parquet":
        comments_df = pd.read_parquet(comments_path)
    elif ext == ".csv":
        comments_df = pd.read_csv(comments_path)
    else:
        raise ValueError(f"Unsupported file type: {ext}")

    # Build comment text lookup
    comments_df[text_col] = comments_df[text_col].astype(str).str.strip()
    comment_texts = set(comments_df[text_col].tolist())

    print(f"Loaded {len(clusters_data)} clusters, {len(comment_texts)} unique comments")

    # Process each cluster
    all_profiles = {}

    for cluster_info in clusters_data:
        cluster_id = cluster_info.get("cluster_id", 0)
        # Use first canonical member as label, or generate one
        canonical_members = cluster_info.get("canonical_members", [])
        cluster_label = canonical_members[0][:80] if canonical_members else f"Cluster {cluster_id}"
        matched_comments = cluster_info.get("example_comments", [])

        # Filter to comments that exist in our dataset
        valid_comments = [c for c in matched_comments if c.get("text", "") in comment_texts]

        if len(valid_comments) < min_cluster_comments:
            print(f"\n[{cluster_id}] {cluster_label}: Only {len(valid_comments)} comments, skipping")
            continue

        print(f"\n{'=' * 60}")
        print(f"[{cluster_id}] {cluster_label}")
        print(f"  Matched comments: {len(valid_comments)}")

        # Sample from matched comments (prefer higher similarity)
        if len(valid_comments) > samples_per_cluster:
            # Sort by similarity score, take top N
            sorted_comments = sorted(valid_comments, key=lambda x: x.get("similarity", 0), reverse=True)
            sampled = sorted_comments[:samples_per_cluster]
        else:
            sampled = valid_comments

        print(f"  Sampling: {len(sampled)} comments")

        # Filter valid comments
        valid_sampled = [
            c for c in sampled
            if c.get("text", "") and len(c.get("text", "")) >= 15
        ]

        # Extract psychology in parallel
        psychology_results = []

        def process_comment(comment_info):
            comment_text = comment_info.get("text", "")
            psych = extract_comment_psychology(comment_text)
            psych["_comment"] = comment_text[:100]
            psych["_similarity"] = comment_info.get("similarity", 0)
            return psych

        with ThreadPoolExecutor(max_workers=max_workers) as executor:
            futures = {executor.submit(process_comment, c): c for c in valid_sampled}
            done_count = 0
            for future in as_completed(futures):
                try:
                    psych = future.result()
                    psychology_results.append(psych)
                    done_count += 1
                    if done_count % 10 == 0:
                        print(f"    Processed {done_count}/{len(valid_sampled)}")
                except Exception as e:
                    print(f"    Error: {e}")

        if len(psychology_results) < 5:
            print(f"  Too few valid extractions ({len(psychology_results)}), skipping")
            continue

        # Aggregate and synthesize
        aggregated = aggregate_to_profile(psychology_results)
        profile = synthesize_profile(aggregated, len(psychology_results))

        # Add problem context to profile
        profile["problem_cluster"] = cluster_label
        profile["canonical_problems"] = cluster_info.get("canonical_members", [])[:5]
        profile["n_comments_analyzed"] = len(psychology_results)

        all_profiles[str(cluster_id)] = profile

        print(f"\n  Profile for '{cluster_label}':")
        print("-" * 40)
        # Print abbreviated profile
        profile_text = profile.get("profile_text", "")
        lines = profile_text.split("\n")[:15]
        for line in lines:
            print(f"  {line}")
        if len(profile_text.split("\n")) > 15:
            print("  ...")

        # Save after each cluster (incremental)
        out_file = output_path or "psychological_profiles_by_cluster.json"
        with open(out_file, "w") as f:
            json.dump(all_profiles, f, indent=2, default=str)
        print(f"  💾 Saved ({len(all_profiles)} profiles so far)")

    print(f"\n{'=' * 60}")
    print(f"✅ Done! {len(all_profiles)} cluster profiles saved to: {output_path or 'psychological_profiles_by_cluster.json'}")

    return all_profiles


# ============================================================
# CLI
# ============================================================

if __name__ == "__main__":
    if len(sys.argv) < 2:
        print("Usage: python extract_audience_psychology.py canonical_problem_clusters.json \\")
        print("           --comments_path comments.csv [--samples_per_cluster N]")
        print("")
        print("Arguments:")
        print("  clusters.json          Path to canonical_problem_clusters.json")
        print("  --comments_path PATH   Path to comments CSV/parquet (required)")
        print("  --samples_per_cluster  Comments to analyze per cluster (default: 30)")
        print("  --text_col COL         Column containing comment text (default: text)")
        print("  --workers N            Parallel workers for LLM calls (default: 10)")
        print("")
        print("Output:")
        print("  Creates psychological_profiles_by_cluster.json with:")
        print("    - One psychological profile per problem cluster")
        print("    - Persuasion vulnerabilities specific to each problem")
        sys.exit(1)

    clusters_path = sys.argv[1]
    comments_path = None
    samples_per_cluster = 30
    text_col = "text"
    max_workers = 10

    # Parse optional args
    i = 2
    while i < len(sys.argv):
        if sys.argv[i] == "--comments_path" and i + 1 < len(sys.argv):
            comments_path = sys.argv[i + 1]
            i += 2
        elif sys.argv[i] == "--samples_per_cluster" and i + 1 < len(sys.argv):
            samples_per_cluster = int(sys.argv[i + 1])
            i += 2
        elif sys.argv[i] == "--text_col" and i + 1 < len(sys.argv):
            text_col = sys.argv[i + 1]
            i += 2
        elif sys.argv[i] == "--workers" and i + 1 < len(sys.argv):
            max_workers = int(sys.argv[i + 1])
            i += 2
        else:
            i += 1

    if comments_path is None:
        print("Error: --comments_path is required")
        sys.exit(1)

    run_extraction_by_cluster(
        clusters_path=clusters_path,
        comments_path=comments_path,
        text_col=text_col,
        samples_per_cluster=samples_per_cluster,
        max_workers=max_workers
    )