#!/usr/bin/env python3
"""
Estimate Importance, Satisfaction, and Opportunity scores (VOC + JTBD)
using emotion and sentiment analysis on archetype data.
"""

import os
import argparse
import pandas as pd
import numpy as np
from pathlib import Path
from transformers import pipeline
from sklearn.preprocessing import MinMaxScaler
import plotly.express as px

os.environ["TOKENIZERS_PARALLELISM"] = "false"
os.environ["OMP_NUM_THREADS"] = "1"
os.environ["MKL_NUM_THREADS"] = "1"
os.environ["NUMEXPR_NUM_THREADS"] = "1"
os.environ["OMP_THREAD_LIMIT"] = "1"

BASE_DIR = Path(__file__).resolve().parent
emotion_path = BASE_DIR / "go_emotions_model"
sentiment_path = BASE_DIR / "bertweet_sentiment_model"

print("🔍 Emotion model path:", emotion_path.exists(), emotion_path)
print("🔍 Sentiment model path:", sentiment_path.exists(), sentiment_path)

# =========================================================
# ARGUMENTS
# =========================================================
parser = argparse.ArgumentParser(description="Value–Opportunity Mapping from JTBD Archetypes")
parser.add_argument("--input", required=True, help="Path to archetype CSV file (output from JTBD pipeline)")
args = parser.parse_args()

INPUT_PATH = Path(args.input)
OUTPUT_DIR = INPUT_PATH.parent

input_path = Path(args.input)
base_name = input_path.stem

# =========================================================
# SAFETENSORS ENVIRONMENT
# =========================================================
os.environ["SAFETENSORS_ALWAYS_USE"] = "1"

# =========================================================
# LOAD DATA
# =========================================================
print(f"📂 Loading archetype data: {INPUT_PATH}")
df = pd.read_csv(INPUT_PATH)

# Ensure 'archetype_cluster' exists
if "archetype_cluster" not in df.columns:
    raise ValueError("Input file must include an 'archetype_cluster' column.")

# Keep only rows where relevance is '✅ Directly relevant'
# df = df[df["relevance"] == "⚙️ Partially relevant"].reset_index(drop=True)

from transformers import AutoTokenizer, AutoModelForSequenceClassification, pipeline
import os
from pathlib import Path

# ---- force offline ----
os.environ["TRANSFORMERS_OFFLINE"] = "1"
os.environ["HF_HUB_OFFLINE"] = "1"

from sentence_transformers import SentenceTransformer
import torch
torch.set_num_threads(1)
torch.set_num_interop_threads(1)
embedder = SentenceTransformer("all-MiniLM-L6-v2")

BASE_DIR = Path(__file__).resolve().parent
emotion_path = BASE_DIR / "go_emotions_model"
sentiment_path = BASE_DIR / "bertweet_sentiment_model"

print("🔍 Emotion model path:", emotion_path.exists(), emotion_path)
print("🔍 Sentiment model path:", sentiment_path.exists(), sentiment_path)

# ---- load tokenizer + model manually to guarantee local use ----
emotion_tok = AutoTokenizer.from_pretrained(str(emotion_path), local_files_only=True)
emotion_mod = AutoModelForSequenceClassification.from_pretrained(
    str(emotion_path), local_files_only=True
)

sent_tok = AutoTokenizer.from_pretrained(str(sentiment_path), local_files_only=True)
sent_mod = AutoModelForSequenceClassification.from_pretrained(
    str(sentiment_path), local_files_only=True
)

# ---- build pipelines from those objects ----
emotion_model = pipeline(
    "text-classification",
    model=emotion_mod,
    tokenizer=emotion_tok,
    device=-1,
    return_all_scores=True,
)

sentiment_model = pipeline(
    "sentiment-analysis",
    model=sent_mod,
    tokenizer=sent_tok,
    device=-1,
    return_all_scores=True,
)

# =========================================================
# SCORING FUNCTIONS
# =========================================================
def emotion_intensity(text):
    """Estimate emotional energy (importance) from outcome text."""
    if not isinstance(text, str) or not text.strip():
        return 0.0

    result = emotion_model(text[:512], return_all_scores=True)
    # result can be list-of-lists, list-of-dicts, or dict
    if isinstance(result, list):
        if isinstance(result[0], list):
            scores = result[0]
        elif isinstance(result[0], dict):
            scores = result
        else:
            return 0.0
    elif isinstance(result, dict):
        scores = [result]
    else:
        return 0.0

    non_neutral = [
        s["score"] for s in scores
        if isinstance(s, dict) and s.get("label", "").lower() != "neutral"
    ]
    return float(np.mean(non_neutral)) if non_neutral else 0.0



def satisfaction_from_sentiment(text):
    """Estimate satisfaction (positive vs negative sentiment) from struggle text."""
    if not isinstance(text, str) or not text.strip():
        return 0.0

    result = sentiment_model(text[:512], return_all_scores=True)
    if isinstance(result, list):
        if isinstance(result[0], list):
            scores = result[0]
        elif isinstance(result[0], dict):
            scores = result
        else:
            return 0.0
    elif isinstance(result, dict):
        scores = [result]
    else:
        return 0.0

    # normalize labels to lowercase
    normalized = {s["label"].lower(): s["score"] for s in scores if "score" in s}

    # handle multiple naming schemes
    pos_label = next((k for k in normalized if k in ["positive", "pos", "1"]), None)
    neg_label = next((k for k in normalized if k in ["negative", "neg", "0"]), None)

    pos = normalized.get(pos_label, 0.0)
    neg = normalized.get(neg_label, 0.0)

    return float(pos - neg)



# =========================================================
# APPLY SCORING
# =========================================================
print("🧮 Scoring emotional intensity (importance) and sentiment (satisfaction)...")
df["importance_score"] = df["outcome"].apply(emotion_intensity)
df["satisfaction_score"] = df["struggle"].apply(satisfaction_from_sentiment)

scaler = MinMaxScaler(feature_range=(1, 10))
df[["importance", "satisfaction"]] = scaler.fit_transform(
    df[["importance_score", "satisfaction_score"]]
)
df["opportunity"] = df["importance"] - df["satisfaction"]

# =========================================================
# AGGREGATE BY ARCHETYPE
# =========================================================
voc = (
    df.groupby("archetype_cluster")[["importance", "satisfaction", "opportunity"]]
    .mean()
    .reset_index()
    .sort_values("opportunity", ascending=False)
)

csv_out = OUTPUT_DIR / f"{base_name}_value_opportunity_map.csv"
voc.to_csv(csv_out, index=False)
print(f"✅ Saved table → {csv_out}")

# =========================================================
# INTERACTIVE VISUALIZATION
# =========================================================

print("\n📊 Top Opportunities:\n")
print(voc.to_string(index=False, float_format=lambda x: f"{x:.2f}"))

# === assume df has columns ===
# "struggle", "outcome", "importance_score", "satisfaction_score"

# ---- 1️⃣ Value Opportunity per Struggle ----
df_struggle = (
    df.groupby("struggle", dropna=True)[["importance_score", "satisfaction_score"]]
    .mean()
    .reset_index()
)
df_struggle["opportunity"] = (
    df_struggle["importance_score"] - df_struggle["satisfaction_score"]
)
df_struggle = df_struggle.sort_values("opportunity", ascending=False)

df_struggle.to_csv(f"{base_name}_value_opportunity_struggles.csv", index=False)
print("\n💥 Saved: value_opportunity_struggles.csv")
print(df_struggle.head(10))

# ---- 2️⃣ Value Opportunity per Outcome ----
df_outcome = (
    df.groupby("outcome", dropna=True)[["importance_score", "satisfaction_score"]]
    .mean()
    .reset_index()
)
df_outcome["opportunity"] = (
    df_outcome["importance_score"] - df_outcome["satisfaction_score"]
)
df_outcome = df_outcome.sort_values("opportunity", ascending=False)

df_outcome.to_csv(f"{base_name}_value_opportunity_outcomes.csv", index=False)
print("\n💥 Saved: value_opportunity_outcomes.csv")
print(df_outcome.head(10))

# =========================================================
# MOTIVATION POLARITY + SEMANTIC POLARITY (Text Summary)
# =========================================================
print("\n🧭 Computing Motivation and Semantic Polarity ...")

from sklearn.metrics.pairwise import cosine_similarity

# -----------------------------
# 1️⃣ Motivation Polarity
# -----------------------------
APPROACH = {"joy", "love", "excitement", "admiration", "gratitude", "pride", "hope"}
AVOIDANCE = {"fear", "anger", "disgust", "sadness", "embarrassment", "guilt", "shame", "confusion"}

def motivation_polarity(text):
    if not text or not isinstance(text, str):
        return np.nan
    try:
        preds = emotion_model(text[:512])[0]
    except Exception:
        return np.nan
    approach_sum = sum(p["score"] for p in preds if p["label"].lower() in APPROACH)
    avoidance_sum = sum(p["score"] for p in preds if p["label"].lower() in AVOIDANCE)
    return approach_sum - avoidance_sum  # positive = approach, negative = avoidance

df["motivation_polarity_struggle"] = df["struggle"].apply(motivation_polarity)
df["motivation_polarity_outcome"] = df["outcome"].apply(motivation_polarity)
df["motivation_polarity_avg"] = df[
    ["motivation_polarity_struggle", "motivation_polarity_outcome"]
].mean(axis=1)

# -----------------------------
# 2️⃣ Semantic Polarity
# -----------------------------
struggle_embs = embedder.encode(df["struggle"].fillna(""), normalize_embeddings=True)
outcome_embs = embedder.encode(df["outcome"].fillna(""), normalize_embeddings=True)
df["semantic_polarity"] = np.diag(cosine_similarity(struggle_embs, outcome_embs))

# -----------------------------
# 3️⃣ Aggregate by Archetype
# -----------------------------
behavioral_map = (
    df.groupby("archetype_cluster")[["motivation_polarity_avg", "semantic_polarity"]]
    .mean()
    .reset_index()
)
behavioral_csv = OUTPUT_DIR / f"{base_name}_behavioral_segmentation_map.csv"
behavioral_map.to_csv(behavioral_csv, index=False)

print(f"\n✅ Saved Behavioral Segmentation Map → {behavioral_csv}\n")

# -----------------------------
# 4️⃣ Print Insightful Summary
# -----------------------------
print("🧭 Behavioral Segmentation Summary")
print("────────────────────────────────────────────")

for _, row in behavioral_map.iterrows():
    cluster = int(row["archetype_cluster"])
    mot = row["motivation_polarity_avg"]
    sem = row["semantic_polarity"]

    # Interpret motivation
    if mot > 0.2:
        mot_desc = f"Approach-driven (+{mot:.2f})"
    elif mot < -0.2:
        mot_desc = f"Avoidance-driven ({mot:.2f})"
    else:
        mot_desc = f"Mixed ({mot:.2f})"

    # Interpret semantic polarity
    if sem < 0.3:
        sem_desc = "seeks Big Change"
    elif sem < 0.6:
        sem_desc = "seeks Moderate Change"
    else:
        sem_desc = "seeks Optimization"

    print(f"Archetype {cluster} → {mot_desc}, {sem_desc} (semantic={sem:.2f})")

print("\n🔗 Combining Value Opportunity with Motivation Polarity...")

# Load value opportunity data (if not already in memory)
vom = pd.read_csv("value_opportunity_map.csv")
pol = behavioral_map.copy()
pol = pol.rename(columns={"motivation_polarity_avg": "motivation_polarity"})

# Merge on archetype_cluster
merged = pd.merge(vom, pol, on="archetype_cluster", how="inner")

# Define polarity category
def polarity_label(val):
    if val < 0.45:
        return "Refinement"
    elif val < 0.60:
        return "Moderate Change"
    else:
        return "Transformation"

merged["change_type"] = merged["semantic_polarity"].apply(polarity_label)

# Sort by opportunity descending
merged = merged.sort_values("opportunity", ascending=False)

# Focus on high-opportunity, moderate-change
ripe_segments = merged[
    (merged["opportunity"] > merged["opportunity"].mean()) &
    (merged["change_type"] == "Moderate Change")
]

print("\n💡 Ripe-for-Marketing Segments (High Opportunity + Moderate Polarity):")
print(ripe_segments[[
    "archetype_cluster", "opportunity", "motivation_polarity", "semantic_polarity", "change_type"
]])

# Save merged table
merged.to_csv(f"{base_name}_value_opportunity_polarity_map.csv", index=False)
print("\n💾 Saved combined map → value_opportunity_polarity_map.csv")

# ---- 🧭 Opportunity × Change Readiness Summary ----
print("\n🧭 Interpreting Opportunity × Change Readiness...")

# Define thresholds
opp_mean = merged["opportunity"].mean()

def opportunity_level(v):
    return "High" if v > opp_mean else "Low"

def readiness_level(v):
    if v < 0.45:
        return "Low"
    elif v < 0.60:
        return "Moderate"
    else:
        return "High"

merged["opportunity_level"] = merged["opportunity"].apply(opportunity_level)
merged["change_readiness"] = merged["semantic_polarity"].apply(readiness_level)

# Quadrant label
merged["segment_label"] = merged.apply(
    lambda r: f"{r['opportunity_level']} Opportunity × {r['change_readiness']} Readiness", axis=1
)

# Summarize by segment
summary = (
    merged.groupby("segment_label")
    .agg(
        archetypes=("archetype_cluster", list),
        avg_opportunity=("opportunity", "mean"),
        avg_readiness=("semantic_polarity", "mean"),
        avg_emotion=("motivation_polarity", "mean"),
    )
    .reset_index()
)

print("\n📋 Opportunity × Change Readiness Segments:")
print(summary)

# Save to CSV
summary.to_csv(f"{base_name}_opportunity_change_readiness_summary.csv", index=False)
print("💾 Saved → opportunity_change_readiness_summary.csv")

# Optional: highlight prime marketing zones
ripe = summary[
    summary["segment_label"].str.contains("High Opportunity × Moderate Readiness")
]
if not ripe.empty:
    print("\n💡 Prime-for-Marketing Segments:")
    for _, row in ripe.iterrows():
        print(
            f"  Archetypes {row['archetypes']} → want improvement, not reinvention "
            f"(avg_opportunity={row['avg_opportunity']:.2f}, readiness={row['avg_readiness']:.2f})"
        )

# ---- 6️⃣ Hybrid Behavioral Clustering: Data-driven + Interpretable ----
print("\n🤖 Discovering Optimal Behavioral Clusters (Opportunity × Motivation × Readiness)...")

from sklearn.preprocessing import StandardScaler
from sklearn.cluster import KMeans
from sklearn.metrics import silhouette_score

# Select behavioral features
features = ["opportunity", "motivation_polarity", "semantic_polarity"]
X = merged[features].fillna(0)

# Normalize for equal weighting
scaler = StandardScaler()
X_scaled = scaler.fit_transform(X)

# 🔍 1️⃣ Find best K automatically via silhouette
best_k, best_score = None, -1
# Prevent KMeans from exceeding available samples
max_k = min(10, len(X_scaled) - 1)
for k in range(2, max_k):
    model = KMeans(n_clusters=k, random_state=42, n_init=10)
    labels = model.fit_predict(X_scaled)
    score = silhouette_score(X_scaled, labels)
    if score > best_score:
        best_k, best_score = k, score

print(f"✅ Optimal cluster count: K = {best_k} (Silhouette = {best_score:.3f})")

# 🧩 2️⃣ Fit final KMeans model
final_model = KMeans(n_clusters=best_k, random_state=42, n_init=10)
merged["behavior_cluster"] = final_model.fit_predict(X_scaled)

# Get cluster centroids in original value space
centroids = pd.DataFrame(
    scaler.inverse_transform(final_model.cluster_centers_),
    columns=features
)
centroids["cluster"] = centroids.index

print("\n📊 Cluster Centroids:")
print(centroids.round(3))

# 🧠 3️⃣ Interpret clusters using behavioral logic
def interpret_cluster(row):
    mot = row["motivation_polarity"]
    sem = row["semantic_polarity"]
    opp = row["opportunity"]

    # Heuristic logic for motivational archetypes
    if opp > 1 and mot < 0 and sem > 0.5:
        return "Frustrated Reformers (feel bad, want big change)"
    elif opp > 0.5 and mot > 0 and sem < 0.5:
        return "Optimistic Maintainers (feel good, want refinement)"
    elif opp < 0.3 and mot < 0 and sem < 0.45:
        return "Cautious Skeptics (low readiness, low reward)"
    else:
        return "Steady Improvers (moderate emotion, steady growth)"

# Label each data point
merged["cluster_label"] = merged.apply(interpret_cluster, axis=1)

# 🧮 4️⃣ Summarize clusters
cluster_summary = (
    merged.groupby(["behavior_cluster", "cluster_label"])[features]
    .mean()
    .reset_index()
    .sort_values("opportunity", ascending=False)
)

print("\n🧭 Behavioral Cluster Summary:")
print(cluster_summary.round(3))

# 💾 Save
cluster_summary.to_csv(f"{base_name}_behavioral_clusters_summary.csv", index=False)
print("💾 Saved → behavioral_clusters_summary.csv")

# Optional: show cluster membership counts
print("\n📈 Cluster Sizes:")
print(merged["cluster_label"].value_counts())

# ---- 7️⃣ Cross-Tab: Archetype × Behavioral Cluster ----
print("\n🔗 Mapping Archetypes ↔ Behavioral Motivation Clusters...")

# Keep only what we need
cross_tab = merged[[
    "archetype_cluster",
    "behavior_cluster",
    "cluster_label",
    "opportunity",
    "motivation_polarity",
    "semantic_polarity"
]].copy()

# Aggregate archetypes within each behavioral cluster
summary = (
    cross_tab.groupby(["cluster_label", "archetype_cluster"])
    .agg({
        "opportunity": "mean",
        "motivation_polarity": "mean",
        "semantic_polarity": "mean"
    })
    .reset_index()
    .sort_values(["cluster_label", "opportunity"], ascending=[True, False])
)

# Save results
summary.to_csv(f"{base_name}_archetype_behavioral_map.csv", index=False)
print("💾 Saved → archetype_behavioral_map.csv")

# Pretty print summary
print("\n📊 Archetype × Behavioral Motivation Map:")
for label, group in summary.groupby("cluster_label"):
    print(f"\n🧭 {label}:")
    for _, row in group.iterrows():
        print(
            f"  • Archetype {int(row['archetype_cluster'])} → "
            f"Opportunity={row['opportunity']:.2f}, "
            f"Motivation={row['motivation_polarity']:.2f}, "
            f"Readiness={row['semantic_polarity']:.2f}"
        )

# =========================================================
# FINAL EXPORT: Append behavioral metrics to main df
# =========================================================
# =========================================================
# FINAL EXPORT: Add behavioral metrics to main df (no columns removed)
# =========================================================
print("\n💾 Saving full archetype data with importance, opportunity, satisfaction, motivation, polarity, and readiness...")

# Map change_readiness from merged table back to df via archetype_cluster
readiness_map = merged.set_index("archetype_cluster")["change_readiness"].to_dict()
df["change_readiness"] = df["archetype_cluster"].map(readiness_map)

# Average motivation polarity across struggle + outcome
df["motivation_polarity"] = df[["motivation_polarity_struggle", "motivation_polarity_outcome"]].mean(axis=1)

# Ensure semantic_polarity exists (it’s computed earlier)
if "semantic_polarity" not in df.columns:
    df["semantic_polarity"] = np.nan

# Save the full dataframe with all existing columns + new metrics
final_out = OUTPUT_DIR / f"{base_name}_jtbd_value_opportunity_full.csv"
df.to_csv(final_out, index=False)
print(f"✅ Saved full dataset → {final_out}")