"""Analysis. Univariate, bivariate and multivariate description of the validated dataset.

    python scripts/analyze.py                       # default workshop dataset
    python scripts/analyze.py --data my_dataset.xlsx --min-confidence 2

Writes site/assets/figs/*.png and data/analysis_tables.json (used by analysis.html),
plus site/files/analysis_tables.xlsx for participants.
"""
import argparse, json
from pathlib import Path
import numpy as np
import pandas as pd
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt

ROOT = Path(__file__).resolve().parents[1]
FIGS = ROOT / "site" / "assets" / "figs"
NAVY, MUTED, LINK, SIGNAL, BORDER = "#10233e", "#586b82", "#1555aa", "#59ddff", "#9eb2c8"
PALETTE = ["#10233e", "#1555aa", "#59ddff", "#9ac3df", "#586b82", "#b7791f", "#9eb2c8", "#dfeaf3"]

plt.rcParams.update({"font.family": ["Malgun Gothic", "sans-serif"], "font.size": 10,
                     "axes.spines.top": False, "axes.spines.right": False, "axes.edgecolor": BORDER,
                     "axes.labelcolor": MUTED, "xtick.color": MUTED, "ytick.color": MUTED, "figure.dpi": 150})


def bar(series, title, fname, horizontal=True, color=NAVY):
    fig, ax = plt.subplots(figsize=(7, 0.42 * len(series) + 1.2) if horizontal else (7, 3.6))
    if horizontal:
        s = series.sort_values()
        ax.barh(s.index.astype(str), s.values, color=color)
        for i, v in enumerate(s.values):
            ax.text(v, i, f" {v}", va="center", fontsize=9, color=MUTED)
        ax.set_xlabel("instances")
    else:
        ax.bar(series.index.astype(str), series.values, color=color)
        ax.set_ylabel("instances")
    ax.set_title(title, loc="left", fontsize=12, color="#10213b", pad=10)
    fig.tight_layout(); fig.savefig(FIGS / fname, bbox_inches="tight"); plt.close(fig)


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--data", default=str(ROOT / "site" / "files" / "dataset_ai_coded.xlsx"))
    ap.add_argument("--min-confidence", type=int, default=2)
    args = ap.parse_args()
    FIGS.mkdir(parents=True, exist_ok=True)

    raw = pd.read_excel(args.data)
    n_rows, n_docs = len(raw), raw.doc_id.nunique()
    df = raw[(raw.relevance == 1) & (raw.coder_confidence >= args.min_confidence)].copy()
    df["year"] = pd.to_datetime(df.article_date).dt.year
    df["period"] = np.where(df.year <= 2017, "2012-2017", np.where(df.year <= 2021, "2018-2021", "2022-2025"))
    tables = {"n_rows": int(n_rows), "n_docs": int(n_docs), "n_relevant": int((raw.relevance == 1).sum()),
              "n_analysis": int(len(df)), "min_confidence": args.min_confidence}

    # ---------------- univariate
    uni = {}
    for var in ["sector", "event_type", "leader_visit", "province", "coder_confidence", "period"]:
        vc = df[var].value_counts()
        uni[var] = {"counts": vc.to_dict(), "pct": (vc / vc.sum() * 100).round(1).to_dict()}
    tables["univariate"] = uni
    bar(df.sector.value_counts(), "Instances by industrial sector", "uni_sector.png")
    bar(df.event_type.value_counts(), "Instances by event type", "uni_event.png", color=LINK)
    bar(df.year.value_counts().sort_index(), "Instances by year", "uni_year.png", horizontal=False, color=LINK)
    top = df.enterprise_name.value_counts().head(12)
    bar(top, "Most frequently reported enterprises", "uni_enterprise.png", color=MUTED)
    tables["top_enterprises"] = top.to_dict()

    # ---------------- bivariate
    biv = {}
    ct1 = pd.crosstab(df.sector, df.event_type)
    ct2 = pd.crosstab(df.sector, df.leader_visit)
    ct3 = pd.crosstab(df.period, df.event_type)
    ct4 = pd.crosstab(df.period, df.leader_visit, normalize="index").round(3) * 100
    for name, ct in {"sector_x_event": ct1, "sector_x_leader": ct2, "period_x_event": ct3, "period_x_leader_pct": ct4}.items():
        biv[name] = {"index": list(map(str, ct.index)), "columns": list(map(str, ct.columns)), "values": ct.values.tolist()}
    tables["bivariate"] = biv

    fig, ax = plt.subplots(figsize=(7, 3.8))
    share = pd.crosstab(df.period, df.event_type, normalize="index")
    bottom = np.zeros(len(share))
    for i, col in enumerate(share.columns):
        ax.bar(share.index, share[col].values, bottom=bottom, label=col, color=PALETTE[i % len(PALETTE)])
        bottom += share[col].values
    ax.set_ylabel("share of instances"); ax.set_ylim(0, 1)
    ax.set_title("Event type by period", loc="left", fontsize=12, color="#10213b", pad=10)
    ax.legend(bbox_to_anchor=(1.02, 1), loc="upper left", frameon=False, fontsize=8)
    fig.tight_layout(); fig.savefig(FIGS / "biv_period_event.png", bbox_inches="tight"); plt.close(fig)

    fig, ax = plt.subplots(figsize=(7, 3.8))
    hm = pd.crosstab(df.sector, df.leader_visit)
    im = ax.imshow(hm.values, cmap="Blues", aspect="auto")
    ax.set_xticks(range(hm.shape[1])); ax.set_xticklabels(hm.columns, rotation=0)
    ax.set_yticks(range(hm.shape[0])); ax.set_yticklabels(hm.index)
    for i in range(hm.shape[0]):
        for j in range(hm.shape[1]):
            ax.text(j, i, hm.values[i, j], ha="center", va="center", fontsize=9, color="#10213b" if hm.values[i, j] < hm.values.max() * .6 else "white")
    ax.set_title("Sector by visiting official", loc="left", fontsize=12, color="#10213b", pad=10)
    fig.tight_layout(); fig.savefig(FIGS / "biv_sector_leader.png", bbox_inches="tight"); plt.close(fig)

    # ---------------- multivariate: logit of a leadership visit on sector and period
    multi = {}
    try:
        import statsmodels.formula.api as smf
        d = df.copy()
        d["visit"] = (d.leader_visit != "none").astype(int)
        d = d[d.sector.isin(d.sector.value_counts()[lambda s: s >= 5].index)]
        base = d.sector.value_counts().idxmax()
        d["sector"] = pd.Categorical(d.sector, categories=[base] + [c for c in d.sector.unique() if c != base])
        d["period"] = pd.Categorical(d.period, categories=["2012-2017", "2018-2021", "2022-2025"])
        model = smf.logit("visit ~ C(sector) + C(period)", data=d).fit(disp=0, maxiter=200)
        tab = pd.DataFrame({"coef": model.params, "odds_ratio": np.exp(model.params), "se": model.bse, "p": model.pvalues}).round(3)
        multi = {"formula": "visit ~ C(sector) + C(period)", "n": int(model.nobs), "pseudo_r2": round(float(model.prsquared), 3),
                 "reference_sector": base, "table": tab.reset_index().rename(columns={"index": "term"}).to_dict("records")}
    except Exception as e:  # small pilot samples can fail to converge
        multi = {"error": str(e)}
    tables["multivariate"] = multi

    (ROOT / "data" / "analysis_tables.json").write_text(json.dumps(tables, ensure_ascii=False, indent=1, default=str), encoding="utf-8")
    with pd.ExcelWriter(ROOT / "site" / "files" / "analysis_tables.xlsx", engine="openpyxl") as w:
        for var in ["sector", "event_type", "leader_visit", "province", "period"]:
            df[var].value_counts().rename("instances").to_frame().to_excel(w, sheet_name=f"uni_{var}"[:31])
        ct1.to_excel(w, sheet_name="sector_x_event"); ct2.to_excel(w, sheet_name="sector_x_leader"); ct3.to_excel(w, sheet_name="period_x_event")
        if "table" in multi:
            pd.DataFrame(multi["table"]).to_excel(w, sheet_name="logit_visit", index=False)
    print(json.dumps({k: tables[k] for k in ["n_rows", "n_docs", "n_relevant", "n_analysis"]}))
    print("multivariate:", "ok" if "table" in multi else multi.get("error"))


if __name__ == "__main__":
    main()