"""Drinkware reviews: what makes buyers unhappy, and can an AI sort reviews well enough to watch them every month?

Data:  2,000 Amazon reviews of insulated drinkware from four brands (YETI, Stanley, Hydro Flask, Owala),
       500 per brand at random, January 2022 to September 2023 (Amazon Reviews 2023, McAuley Lab, UC San Diego).
AI:    Laya, an open "System One" decision model: it answers typed questions (yes/no, pick one) with probabilities
       instead of writing text. Its answers are in data/ai_answers.csv (made by run_model.py).
       Claude, an LLM, named the main complaint of each one- and two-star review: data/llm_complaints.csv (llm_code.py).
Human: an optional check against hand labels on 80 one- and two-star reviews (data/my_labels.csv, made with
       label.html); Step 7 runs only when that file exists.

Run:   python3 analysis.py      ->  results.json (every number on the page comes from here)
Needs: pandas and numpy.
"""
import json
import os

import numpy as np
import pandas as pd

# Shared constants: complaint types, display names, brand order
TYPES = ["leaks_lid", "temperature", "damage", "wrong_item", "size_design", "cleaning_taste", "price"]
NAMES = {"leaks_lid": "Leaks or lid problems", "temperature": "Doesn't keep cold or hot", "damage": "Damaged, dented or broken",
         "wrong_item": "Wrong or missing parts", "size_design": "Size, weight or design", "cleaning_taste": "Cleaning, smell or taste",
         "price": "Too expensive", "other": "Something else"}
BRANDS = ["YETI", "Stanley", "Hydro Flask", "Owala"]

# Results container and display settings
out = {}                                        # every number for the page goes in here, then into results.json
pd.set_option("display.width", 150)             # print wide tables on one line
pd.set_option("display.max_columns", 20)

# ---------- Step 1. Load the reviews and the AI's answers ----------
reviews = pd.read_csv("data/reviews.csv")
ai = pd.read_csv("data/ai_answers.csv")
df = reviews.merge(ai, on="review_id", how="left")
print("Reviews:", len(reviews), "| AI answers:", len(ai), "| Together:", len(df))

# ---------- Step 2. Check the data before using it ----------
# Data-quality checks before any modeling
words = df["title"].fillna("") + df["text"].fillna("")
checks = {
    # Volume, brand balance, duplicates
    "rows": len(df),
    "per_brand": df["brand"].value_counts().reindex(BRANDS).tolist(),
    "duplicate_ids": int(df["review_id"].duplicated().sum()),
    # Content: empty text, non-English reviews
    "empty_text": int(words.str.strip().eq("").sum()),
    "spanish": int((df["language"] == "es").sum()),
    # Coverage: date range, rating distribution
    "first_day": df["date"].min(),
    "last_day": df["date"].max(),
    "stars": df["rating"].value_counts().sort_index().tolist(),
}
print(checks)

# Reference label from star ratings (the reviewer's own verdict)
df["happy_by_stars"] = np.nan                    # 3 stars stay empty: a mixed review, left out of the mood check
df.loc[df["rating"] >= 4, "happy_by_stars"] = 1
df.loc[df["rating"] <= 2, "happy_by_stars"] = 0
out["checks"] = checks

# ---------- Step 3. Does the AI read the mood the way the reviewer rated it? ----------
# Evaluation set: star label and model score both present; threshold at 0.5
both = df[df["happy_by_stars"].notna() & df["p_happy"].notna()].copy()
both["ai_happy"] = (both["p_happy"] >= 0.5).astype(int)

# Agreement with the reference: confusion table and accuracy
table = pd.crosstab(both["happy_by_stars"], both["ai_happy"])
print(table)
accuracy = (both["ai_happy"] == both["happy_by_stars"]).mean()


# Chance-corrected agreement (Cohen's kappa)
def kappa(a, b):
    """Cohen's kappa: how much two raters agree beyond what they would agree on by luck."""
    a, b = pd.Series(list(a)), pd.Series(list(b))
    observed = (a == b).mean()
    chance = 0
    for label in set(a) | set(b):
        chance += (a == label).mean() * (b == label).mean()
    return (observed - chance) / (1 - chance)


# Mood metrics: kappa and unhappy reviews scored as happy
k_mood = kappa(both["ai_happy"], both["happy_by_stars"])
missed = int(((both["happy_by_stars"] == 0) & (both["ai_happy"] == 1)).sum())
print(f"Mood agrees with the stars on {accuracy:.1%} of {len(both)} reviews; kappa {k_mood:.2f}; unhappy reviews it missed: {missed}")

# Persist results for the report
out["mood"] = {"n": len(both), "accuracy": round(accuracy, 3), "kappa": round(k_mood, 2),
               "unhappy_n": int((both["happy_by_stars"] == 0).sum()), "unhappy_missed": missed,
               "table": {"stars_happy_ai_happy": int(table.loc[1, 1]), "stars_happy_ai_unhappy": int(table.loc[1, 0]),
                         "stars_unhappy_ai_happy": int(table.loc[0, 1]), "stars_unhappy_ai_unhappy": int(table.loc[0, 0])}}

# ---------- Step 4. When the AI says "90% sure", is it right 90% of the time? (calibration) ----------
# Per-review confidence in the predicted class, and correctness
both["confidence"] = np.maximum(both["p_happy"], 1 - both["p_happy"])
both["right"] = (both["ai_happy"] == both["happy_by_stars"]).astype(int)
bands = [0.5, 0.6, 0.7, 0.8, 0.9, 1.0001]


# Reliability by confidence band and expected calibration error (ECE)
def reliability(confidence, right):
    """For each confidence band: how many reviews, how sure the AI said it was, and how often it was right."""
    t = pd.DataFrame({"said": confidence, "right": right})
    t["band"] = pd.cut(t["said"], bands, right=False)
    g = t.groupby("band", observed=True).agg(n=("right", "size"), said=("said", "mean"), right=("right", "mean"))
    error = (g["n"] * (g["said"] - g["right"]).abs()).sum() / g["n"].sum()
    return g, error


# Baseline calibration on all reviews
g_all, error_all = reliability(both["confidence"], both["right"])
print(g_all.round(3))
print(f"Calibration error (all reviews, before any fix): {error_all:.3f}")

# Temperature scaling: fit T on a random learn split, evaluate on the held-out test split
rng = np.random.default_rng(7)
both["half"] = np.where(rng.random(len(both)) < 0.5, "learn", "test")


def adjust(p, T):
    """Temperature: T above 1 pulls probabilities toward 50%, T below 1 pushes them toward 0% or 100%."""
    a, b = p ** (1 / T), (1 - p) ** (1 / T)
    return a / (a + b)


def log_loss(p, y):
    """The penalty for a probability: small when the AI is sure and right, large when it is sure and wrong."""
    p = np.clip(p, 1e-6, 1 - 1e-6)
    return -np.mean(y * np.log(p) + (1 - y) * np.log(1 - p))


# Grid search for T on the learn split (minimum log loss)
learn = both[both["half"] == "learn"]
test = both[both["half"] == "test"].copy()
best_T, best_loss = None, None
for T in np.round(np.arange(0.5, 5.01, 0.05), 2):
    loss = log_loss(adjust(learn["p_happy"], T), learn["happy_by_stars"])
    if best_loss is None or loss < best_loss:
        best_T, best_loss = float(T), loss

# Apply T to the test split; compare ECE before and after
test["p_fixed"] = adjust(test["p_happy"], best_T)
test["confidence_fixed"] = np.maximum(test["p_fixed"], 1 - test["p_fixed"])
g_before, error_before = reliability(test["confidence"], test["right"])
g_after, error_after = reliability(test["confidence_fixed"], test["right"])
print(f"Best T = {best_T}; on the test half the calibration error went from {error_before:.3f} to {error_after:.3f}")

# Persist results for the report
out["calibration"] = {"T": best_T, "ece_all_before": round(error_all, 3), "ece_test_before": round(error_before, 3),
                      "ece_test_after": round(error_after, 3), "test_n": len(test),
                      "bands_before": [{"band": str(b), "n": int(r.n), "said": round(r.said, 3), "right": round(r.right, 3)} for b, r in g_before.iterrows()],
                      "bands_after": [{"band": str(b), "n": int(r.n), "said": round(r.said, 3), "right": round(r.right, 3)} for b, r in g_after.iterrows()]}

# ---------- Step 5. A routing rule: the AI handles the sure ones, a person reads the rest ----------
# Confidence-threshold routing: automation rate vs. accuracy at each cut-off
routes = []
for cut in [0.6, 0.7, 0.8, 0.9, 0.95]:
    sure = test[test["confidence_fixed"] >= cut]
    routes.append({"cut": cut,
                   "auto_share": round(len(sure) / len(test), 3),
                   "auto_right": round(sure["right"].mean(), 3) if len(sure) else None,
                   "to_person": len(test) - len(sure)})
print(pd.DataFrame(routes))
out["routing"] = routes

# ---------- Step 6. What do unhappy buyers complain about? Two AI coders ----------
# Coder 1 (decision model): highest-probability complaint type, 'other' below 0.5
unhappy = df[(df["rating"] <= 2) & df["c_leaks_lid"].notna()].copy()
cols = ["c_" + t for t in TYPES]
unhappy["dm_main"] = unhappy[cols].idxmax(axis=1).str.replace("c_", "", regex=False)
unhappy["dm_main_p"] = unhappy[cols].max(axis=1)
unhappy.loc[unhappy["dm_main_p"] < 0.5, "dm_main"] = "other"

# Coder 2 (LLM): join Claude's main-complaint codes
llm = pd.read_csv("data/llm_complaints.csv")
unhappy = unhappy.merge(llm, on="review_id", how="left")

# Inter-coder agreement
print("One- and two-star reviews coded:", len(unhappy))
print(pd.DataFrame({"decision model": unhappy["dm_main"].value_counts(),
                    "LLM": unhappy["llm_complaint"].value_counts()}).fillna(0).astype(int))
coders_agree = (unhappy["dm_main"] == unhappy["llm_complaint"]).mean()
print(f"The two AI coders agree on {coders_agree:.0%} of the unhappy reviews")
out["coders_agree"] = round(coders_agree, 3)

# ---------- Step 7. The human check: my labels against both AI coders ----------
human = None
if os.path.exists("data/my_labels.csv"):
    # Join human labels with both coders' outputs
    labels = pd.read_csv("data/my_labels.csv")
    labels = labels.merge(unhappy[["review_id", "brand", "dm_main", "llm_complaint"]], on="review_id")

    # Agreement with human labels: accuracy and Cohen's kappa per coder
    dm_agree = (labels["my_complaint"] == labels["dm_main"]).mean()
    llm_agree = (labels["my_complaint"] == labels["llm_complaint"]).mean()
    dm_kappa = kappa(labels["dm_main"], labels["my_complaint"])
    llm_kappa = kappa(labels["llm_complaint"], labels["my_complaint"])
    print(f"Decision model matches my labels on {dm_agree:.0%} of {len(labels)} reviews; kappa {dm_kappa:.2f}")
    print(f"LLM matches my labels on {llm_agree:.0%} of {len(labels)} reviews; kappa {llm_kappa:.2f}")
    print(pd.crosstab(labels["my_complaint"], labels["llm_complaint"]))

    # Per-type recall against human labels
    per_type = []
    for t in TYPES + ["other"]:
        mine = labels[labels["my_complaint"] == t]
        if len(mine) > 0:
            per_type.append({"type": t, "name": NAMES[t], "mine": len(mine),
                             "llm_matched": int((mine["llm_complaint"] == t).sum()),
                             "dm_matched": int((mine["dm_main"] == t).sum())})

    # Persist results for the report
    human = {"n": len(labels), "dm_agree": round(dm_agree, 3), "dm_kappa": round(dm_kappa, 2),
             "llm_agree": round(llm_agree, 3), "llm_kappa": round(llm_kappa, 2), "per_type": per_type}
else:
    print("No data/my_labels.csv yet: label the 80 reviews in label.html first, then run this again.")
out["human_check"] = human

# ---------- Step 8. The so-what: the main complaint in unhappy reviews, by brand (LLM coder) ----------
# Complaint mix by brand (LLM coder): shares of each brand's unhappy reviews
by_brand = []
for b in BRANDS:
    brand_unhappy = unhappy[unhappy["brand"] == b]
    brand_all = df[(df["brand"] == b) & (df["language"] == "en")]
    row = {"brand": b, "reviews": len(brand_all), "unhappy": len(brand_unhappy),
           "unhappy_share": round(len(brand_unhappy) / len(brand_all), 3)}
    for t in TYPES + ["other"]:
        row[t] = round((brand_unhappy["llm_complaint"] == t).mean(), 3)
    by_brand.append(row)
print(pd.DataFrame(by_brand).set_index("brand"))

# Persist results for the report
out["by_brand"] = by_brand
out["complaint_totals"] = unhappy["llm_complaint"].value_counts().to_dict()

# Drill-down: YETI 'wrong or missing parts' reviews that report a missing lid part (keyword rules)
yeti_parts = unhappy[(unhappy["brand"] == "YETI") & (unhappy["llm_complaint"] == "wrong_item")]
text = (yeti_parts["title"].fillna("") + ". " + yeti_parts["text"].fillna("")).str.lower()
said_missing = text.str.contains(r"missing|without|not included|no lid|didn.t (?:get|receive|come)|did not (?:receive|come|get|ship)|doesn.t come|not come with")
lid_part = text.str.contains(r"slider|slide|lid|cap|top|piece|part|insert")
missing_lid_part = said_missing & lid_part
slider = missing_lid_part & text.str.contains(r"slider|slide")
print(f"YETI: {missing_lid_part.sum()} of its {len(yeti_parts)} 'wrong or missing parts' reviews say a lid part was missing; {slider.sum()} name the slider")

# Persist results for the report
out["yeti_missing_lid_part"] = {"n": int(missing_lid_part.sum()), "of": len(yeti_parts), "slider": int(slider.sum()),
                                "unhappy": int((unhappy["brand"] == "YETI").sum())}
out["unhappy_total"] = len(unhappy)
out["names"] = NAMES

# Write results.json (the report reads only this file)
with open("results.json", "w") as f:
    json.dump(out, f, indent=1, default=str)
print("Saved results.json")