from __future__ import annotations

import json
import re
from pathlib import Path

import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import statsmodels.formula.api as smf
from scipy import stats


ROOT = Path(__file__).resolve().parent
FIG = ROOT / "figures"
FIG.mkdir(exist_ok=True)
plt.rcParams["font.family"] = ["Microsoft JhengHei", "Noto Sans CJK TC", "DejaVu Sans"]
plt.rcParams["axes.unicode_minus"] = False


def infer_brand(name: str) -> str:
    bracket = re.search(r"[【\[]([^】\]]+)", str(name))
    if bracket:
        return bracket.group(1).strip()[:20]
    cleaned = re.sub(r"^[★☆◆◇\s]+", "", str(name))
    return re.split(r"[\s_｜|（(]", cleaned, maxsplit=1)[0][:20]


products = pd.read_csv(ROOT / "products-raw.csv")
products["brand_proxy"] = products["name"].map(infer_brand)
products["log_unit_price"] = np.log(products["unit_price_ntd_per_kg"])
products["log_unit_weight"] = np.log(products["unit_weight_kg"])
products["log_bundle_count"] = np.log(products["bundle_count"].fillna(1))

# The confirmatory sample prioritizes weights found in product titles and excludes
# implausible unit prices. Sensitivity results use a broader parsed sample.
main = products[
    (products["weight_confidence"] == "title")
    & products["unit_weight_kg"].between(0.5, 12)
    & products["unit_price_ntd_per_kg"].between(50, 1500)
].copy()
main["size_group"] = pd.Categorical(
    main["package_size_group"],
    categories=["large_3_to_5kg", "bulk_over_5kg", "medium_2_to_3kg", "small_le_2kg"],
)
main["small_package"] = (main["unit_weight_kg"] <= 2).astype(int)

controls = "rice_type + organic + traceable + award + vacuum + gift + promoted + variety_named + origin_named + log_bundle_count"
elasticity_model = smf.ols(f"log_unit_price ~ log_unit_weight + {controls}", data=main).fit(cov_type="HC3")
category_model = smf.ols(f"log_unit_price ~ C(size_group, Treatment(reference='large_3_to_5kg')) + {controls}", data=main).fit(cov_type="HC3")

brand_counts = main["brand_proxy"].value_counts()
brand_sample = main[main["brand_proxy"].isin(brand_counts[brand_counts >= 3].index)].copy()
brand_model = smf.ols(f"log_unit_price ~ log_unit_weight + {controls} + C(brand_proxy)", data=brand_sample).fit(cov_type="HC3")

broader = products[
    products["weight_confidence"].isin(["title", "description"])
    & products["unit_weight_kg"].between(0.5, 12)
    & products["unit_price_ntd_per_kg"].between(50, 2000)
].copy()
broader_model = smf.ols(f"log_unit_price ~ log_unit_weight + {controls}", data=broader).fit(cov_type="HC3")

group_order = ["small_le_2kg", "medium_2_to_3kg", "large_3_to_5kg", "bulk_over_5kg"]
group_labels = {"small_le_2kg": "≤2 kg", "medium_2_to_3kg": ">2–3 kg", "large_3_to_5kg": ">3–5 kg", "bulk_over_5kg": ">5 kg"}
group_summary = (
    main.groupby("package_size_group", observed=True)["unit_price_ntd_per_kg"]
    .agg(n="size", median="median", mean="mean", q1=lambda s: s.quantile(.25), q3=lambda s: s.quantile(.75))
    .reindex(group_order)
)
group_summary.to_csv(ROOT / "package-size-summary.csv", encoding="utf-8-sig")

# Household structure: normalize English and legacy Chinese field names.
h = pd.read_csv(ROOT / "household-structure-2017-2025.csv", low_memory=False)
def coalesce(names):
    existing = [h[name] for name in names if name in h.columns]
    out = existing[0]
    for series in existing[1:]:
        out = out.fillna(series)
    return pd.to_numeric(out, errors="coerce").fillna(0)

h["roc_year"] = coalesce(["statistic_yyy", "\ufeffstatistic_yyy", "統計年"])
h["year"] = h["roc_year"] + 1911
h["one"] = coalesce(["household_single_total", "1人家戶"])
for number in range(2, 10):
    h[f"n{number}"] = coalesce([f"home_group_{number:02d}", f"{number}人家戶"])
h["n10"] = coalesce(["home_group_10up", "10人以上家戶"])
household = h.groupby("year")[["one"] + [f"n{i}" for i in range(2, 11)]].sum()
household["total_households"] = household.sum(axis=1)
household["one_person_share"] = household["one"] / household["total_households"]
household["one_two_person_share"] = (household["one"] + household["n2"]) / household["total_households"]
household = household.loc[2017:2025]
household.to_csv(ROOT / "household-trend-summary.csv", encoding="utf-8-sig")

# Descriptive text coding for the full analytic sample and a 30-item purposive close-reading set.
text = (main["name"].fillna("") + " " + main["description"].fillna(""))
themes = {
    "freshness_storage": r"真空|保鮮|新鮮|夾鏈|分裝|保存",
    "convenience_portion": r"小包裝|小家庭|單身|一人|方便|少量|份量",
    "quality_safety": r"有機|產銷履歷|SGS|檢驗|無農藥|認證",
    "origin_story": r"池上|關山|花蓮|富里|宜蘭|台東|臺東|農會|小農|產地",
    "taste_variety": r"越光|香米|台[梗稉]|口感|香Q|黏性|壽司",
}
for theme, pattern in themes.items():
    main[theme] = text.str.contains(pattern, regex=True).astype(int)
theme_summary = main.groupby("small_package")[[*themes]].mean().T
theme_summary.columns = ["over_2kg_share", "le_2kg_share"]
theme_summary["difference_pp"] = (theme_summary["le_2kg_share"] - theme_summary["over_2kg_share"]) * 100
theme_summary.to_csv(ROOT / "text-theme-summary.csv", encoding="utf-8-sig")
close_reading = main[main["small_package"] == 1].copy()
close_reading["theme_count"] = close_reading[[*themes]].sum(axis=1)
close_reading = close_reading.sort_values(["theme_count", "text_small_household", "unit_price_ntd_per_kg"], ascending=False).head(30)
close_reading[["product_id", "name", "description", "unit_weight_kg", "unit_price_ntd_per_kg", *themes]].to_csv(
    ROOT / "qualitative-subsample-30.csv", index=False, encoding="utf-8-sig"
)

# Figure 1: unit price distributions by package size.
fig, ax = plt.subplots(figsize=(9, 5.4))
data = [main.loc[main["package_size_group"] == group, "unit_price_ntd_per_kg"] for group in group_order]
ax.boxplot(data, labels=[group_labels[g] for g in group_order], showfliers=False, patch_artist=True,
           boxprops={"facecolor": "#dce9df", "edgecolor": "#315c4c"}, medianprops={"color": "#8a3b24", "linewidth": 2})
ax.set_ylabel("Unit price (NT$/kg)")
ax.set_xlabel("Weight of each rice package")
ax.set_title("Smaller rice packages have higher observed unit prices")
ax.grid(axis="y", alpha=.2)
fig.tight_layout(); fig.savefig(FIG / "unit-price-by-size.png", dpi=180); plt.close(fig)

# Figure 2: log-log relation and fitted line.
fig, ax = plt.subplots(figsize=(9, 5.4))
ax.scatter(main["unit_weight_kg"], main["unit_price_ntd_per_kg"], s=20, alpha=.45, color="#315c4c")
xgrid = np.geomspace(main["unit_weight_kg"].min(), main["unit_weight_kg"].max(), 100)
simple = smf.ols("log_unit_price ~ log_unit_weight", data=main).fit()
ygrid = np.exp(simple.params["Intercept"] + simple.params["log_unit_weight"] * np.log(xgrid))
ax.plot(xgrid, ygrid, color="#a84b2f", linewidth=2.2)
ax.set_xscale("log"); ax.set_yscale("log")
ax.set_xlabel("Package weight (kg, log scale)"); ax.set_ylabel("Unit price (NT$/kg, log scale)")
ax.set_title("Package weight and listed unit price")
ax.grid(alpha=.2, which="both")
fig.tight_layout(); fig.savefig(FIG / "weight-price-scatter.png", dpi=180); plt.close(fig)

# Figure 3: household composition trends.
fig, ax = plt.subplots(figsize=(9, 5.4))
ax.plot(household.index, household["one_person_share"] * 100, marker="o", label="One-person households")
ax.plot(household.index, household["one_two_person_share"] * 100, marker="o", label="One- or two-person households")
ax.set_ylabel("Share of registered households (%)"); ax.set_xlabel("Year")
ax.set_title("Registered household structure in Taiwan, 2017–2025")
ax.legend(frameon=False); ax.grid(alpha=.2)
fig.tight_layout(); fig.savefig(FIG / "household-structure-trend.png", dpi=180); plt.close(fig)

# Figure 4: text themes by size.
fig, ax = plt.subplots(figsize=(9, 5.4))
labels = ["Freshness/storage", "Convenience/portion", "Quality/safety", "Origin story", "Taste/variety"]
y = np.arange(len(labels)); width = .36
ax.barh(y - width/2, theme_summary["le_2kg_share"] * 100, height=width, label="≤2 kg", color="#315c4c")
ax.barh(y + width/2, theme_summary["over_2kg_share"] * 100, height=width, label=">2 kg", color="#c59d61")
ax.set_yticks(y, labels); ax.invert_yaxis(); ax.set_xlabel("Listings mentioning theme (%)")
ax.set_title("Marketing themes in public product text"); ax.legend(frameon=False); ax.grid(axis="x", alpha=.2)
fig.tight_layout(); fig.savefig(FIG / "text-themes-by-size.png", dpi=180); plt.close(fig)

def model_record(model, key):
    beta = float(model.params[key]); se = float(model.bse[key])
    return {
        "coefficient": beta,
        "robust_se": se,
        "p_value": float(model.pvalues[key]),
        "ci95_low": beta - 1.96 * se,
        "ci95_high": beta + 1.96 * se,
        "percent_change": (np.exp(beta) - 1) * 100,
    }

cat_key = "C(size_group, Treatment(reference='large_3_to_5kg'))[T.small_le_2kg]"
results = {
    "collection_date": str(products["collected_at"].iloc[0]),
    "raw_listings": int(len(products)),
    "parsed_listings": int(products["unit_price_ntd_per_kg"].notna().sum()),
    "main_analytic_n": int(len(main)),
    "broader_analytic_n": int(len(broader)),
    "main_sample_rules": "title-parsed weight 0.5–12 kg; listed unit price NT$50–1,500/kg",
    "group_summary": group_summary.reset_index().where(pd.notna(group_summary.reset_index()), None).to_dict("records"),
    "elasticity_model": {**model_record(elasticity_model, "log_unit_weight"), "n": int(elasticity_model.nobs), "r_squared": float(elasticity_model.rsquared)},
    "small_vs_3_to_5kg_model": {**model_record(category_model, cat_key), "n": int(category_model.nobs), "r_squared": float(category_model.rsquared)},
    "brand_fixed_effect_model": {**model_record(brand_model, "log_unit_weight"), "n": int(brand_model.nobs), "r_squared": float(brand_model.rsquared)},
    "broader_sample_model": {**model_record(broader_model, "log_unit_weight"), "n": int(broader_model.nobs), "r_squared": float(broader_model.rsquared)},
    "household_2017": {
        "total": int(household.loc[2017, "total_households"]),
        "one_person_share": float(household.loc[2017, "one_person_share"]),
        "one_two_person_share": float(household.loc[2017, "one_two_person_share"]),
    },
    "household_2025": {
        "total": int(household.loc[2025, "total_households"]),
        "one_person_share": float(household.loc[2025, "one_person_share"]),
        "one_two_person_share": float(household.loc[2025, "one_two_person_share"]),
    },
    "theme_summary": theme_summary.reset_index(names="theme").to_dict("records"),
    "normality_note": "Inference uses HC3 heteroskedasticity-robust standard errors; observational listings do not identify consumer demand or causality.",
}
(ROOT / "analysis-results.json").write_text(json.dumps(results, ensure_ascii=False, indent=2), encoding="utf-8")
main.to_csv(ROOT / "products-analysis.csv", index=False, encoding="utf-8-sig")
print(json.dumps(results, ensure_ascii=False, indent=2))