"""Phase C — build the analysis dataset from extraction JSONs: data/final/studies.csv, data/final/coefficients_raw.csv,
data/final/effects.csv (harmonised, with moderators). Primary synthesis uses confidence in {high, medium} and family != other.
Usage: python scripts/harmonization/build_effects.py"""
import json, re, sys
from pathlib import Path
import numpy as np, pandas as pd, yaml

ROOT = Path(__file__).resolve().parents[2]; sys.path.insert(0, str(ROOT / "scripts" / "harmonization"))
from harmonize import harmonise_row  # noqa: E402

JS = ROOT / "data" / "extracted" / "json"; FINAL = ROOT / "data" / "final"; FINAL.mkdir(parents=True, exist_ok=True)
countries = yaml.safe_load((ROOT / "scripts" / "harmonization" / "countries.yaml").read_text())
tax = yaml.safe_load((ROOT / "scripts" / "harmonization" / "taxonomy.yaml").read_text())
TAX_GROUP = {leaf: grp for grp, leaves in tax.items() for leaf in leaves}
ALIAS = {"usa": "United States", "us": "United States", "u.s.": "United States", "united states of america": "United States", "uk": "United Kingdom", "england": "United Kingdom",
         "scotland": "United Kingdom", "wales": "United Kingdom", "great britain": "United Kingdom", "korea": "South Korea", "republic of korea": "South Korea", "prc": "China",
         "people's republic of china": "China", "hong kong sar": "Hong Kong", "the netherlands": "Netherlands", "holland": "Netherlands", "czechia": "Czech Republic", "türkiye": "Turkey", "turkiye": "Turkey"}
EST_CLASS = {"ols": "ols", "wls": "ols", "gls": "ols", "fe_panel": "panel", "panel": "panel", "iv": "iv", "2sls": "iv", "quantile": "quantile", "spatial_lag": "spatial", "spatial_error": "spatial", "sac": "spatial",
             "sdm": "spatial", "gwr": "spatial", "spatial": "spatial", "ml_marginal": "ml", "random_forest": "ml", "xgboost": "ml", "gbm": "ml", "neural": "ml", "shap": "ml", "did": "quasi_experimental", "rdd": "quasi_experimental", "box-cox": "ols", "boxcox": "ols"}


def canon_country(c):
    if not c: return "unknown"
    c0 = str(c).split(",")[0].split(";")[0].split("/")[0].strip(); key = c0.lower()
    if key in ALIAS: return ALIAS[key]
    for name in countries:
        if name.lower() == key or name.lower() in key: return name
    return c0.title() if c0 else "unknown"


def est_class(e):
    e = (e or "").lower()
    for k, v in EST_CLASS.items():
        if k in e: return v
    return "other" if e else "unknown"


def build_ai():
    rows = []
    for f in sorted((ROOT / "data" / "extracted" / "ai").glob("*.json")):
        try: j = json.loads(f.read_text())
        except Exception: continue
        if not j.get("eligible", True): continue
        j["record_id"] = j.get("record_id", f.stem); j["country"] = canon_country(j.get("country")); j["methods_compared"] = ";".join(j.get("methods_compared") or [])
        rows.append({k: v for k, v in j.items() if not isinstance(v, (dict, list))})
    if rows:
        pd.DataFrame(rows).to_csv(FINAL / "ai_studies.csv", index=False); print("ai studies", len(rows))


def main():
    build_ai()
    studies, coefs = [], []
    for f in sorted(JS.glob("*.json")):
        try: j = json.loads(f.read_text())
        except Exception as e:
            studies.append({"record_id": f.stem, "json_error": repr(e)}); continue
        st = j.get("study") or {}; rid = j.get("record_id", f.stem)
        country = canon_country(st.get("country")); meta = countries.get(country, [None, None, None, None, None])
        srow = {"record_id": rid, "eligible": bool(j.get("eligible", False)), "exclusion_reason": j.get("exclusion_reason"), "title": st.get("title"), "year": st.get("year"), "doi": st.get("doi"),
                "country": country, "continent": meta[0], "income_group": meta[1], "iso3": meta[2], "lat": meta[3], "lon": meta[4], "region_city": st.get("region_city"),
                "sample_start": st.get("sample_start"), "sample_end": st.get("sample_end"), "property_type": st.get("property_type") or "unknown", "price_type": st.get("price_type") or "unknown",
                "document_type": st.get("document_type") or "unknown", "estimators": ";".join(st.get("estimators") or []), "spatial_method": st.get("spatial_method"), "time_fe": st.get("time_fe"),
                "location_fe": st.get("location_fe"), "se_type": st.get("se_type"), "n_obs_total": st.get("n_obs_total"), "ai_methods": ";".join(st.get("ai_methods") or []),
                "n_models": len(j.get("models") or []), "n_coefficients": len(j.get("coefficients") or [])}
        studies.append(srow)
        models = {m.get("model_id"): m for m in (j.get("models") or [])}
        for c in j.get("coefficients") or []:
            m = models.get(c.get("model_id"), {})
            coefs.append({"record_id": rid, "study_id": rid, "model_id": f"{rid}:{c.get('model_id')}", "table": m.get("table"), "column": m.get("column"), "page": c.get("page") or m.get("page"),
                          "var_name_printed": c.get("var_name_printed"), "taxonomy_id": c.get("taxonomy_id") or "other", "taxonomy_group": TAX_GROUP.get(c.get("taxonomy_id"), "other"),
                          "x_transform": c.get("x_transform"), "x_unit": c.get("x_unit"), "beta": c.get("beta"), "se": c.get("se"), "t": c.get("t"), "p": c.get("p"), "p_stars": c.get("p_stars"),
                          "ci_lo": c.get("ci_lo"), "ci_hi": c.get("ci_hi"), "confidence": c.get("confidence") or "low", "notes": c.get("notes"),
                          "dep_var": m.get("dep_var"), "dep_transform": m.get("dep_transform"), "estimator": m.get("estimator") or (st.get("estimators") or [""])[0], "n_obs": m.get("n_obs") or st.get("n_obs_total"),
                          "r2": m.get("r2"), "paren": m.get("se_or_t_in_parentheses"), "sample_subset": m.get("sample_subset")})
    S = pd.DataFrame(studies); C = pd.DataFrame(coefs)
    S.to_csv(FINAL / "studies.csv", index=False); C.to_csv(FINAL / "coefficients_raw.csv", index=False)
    if len(C) == 0:
        print("no coefficients yet"); return
    # harmonise
    H = []
    for r in C.to_dict("records"):
        for k in ("beta", "se", "t", "p", "ci_lo", "ci_hi", "n_obs"):
            try: r[k] = float(r[k]) if r[k] is not None and r[k] != "" else np.nan
            except Exception: r[k] = np.nan
        if np.isnan(r["beta"]): continue
        h = harmonise_row({**r, "k_params": None}); H.append({**r, **h})
    E = pd.DataFrame(H).merge(S[["record_id", "year", "country", "continent", "income_group", "sample_start", "sample_end", "property_type", "price_type", "document_type", "time_fe", "location_fe", "se_type", "ai_methods"]], on="record_id", how="left")
    for c in ("sample_start", "sample_end", "year"):
        E[c] = pd.to_numeric(E[c], errors="coerce"); E.loc[(E[c] < 1900) | (E[c] > 2027), c] = np.nan
    E["data_mid_year"] = E[["sample_start", "sample_end"]].mean(axis=1); E["data_mid_year"] = E.data_mid_year.fillna(E.year - 3)
    E["estimator_class"] = E.estimator.map(est_class); E["pub_type"] = E.document_type.map(lambda d: "journal" if d == "journal" else "other")
    E["usable"] = (E.es_se.notna() & (E.es_se > 0) & E.confidence.isin(["high", "medium"]) & E.family.isin(["pct_premium", "semi_elasticity", "elasticity"]) & (E.taxonomy_id != "other")
                   & (E.harm_flags.fillna("") == ""))
    E["usable_partial_r"] = E.partial_r.notna() & E.confidence.isin(["high", "medium"]) & (E.taxonomy_id != "other") & ~E.harm_flags.fillna("").str.contains("year_built")
    # winsorisation flag for extreme harmonised effects within family (robustness, not applied to primary)
    E["extreme"] = False
    for (fam, tx), g in E[E.usable].groupby(["family", "taxonomy_id"]):
        if len(g) >= 10:
            lo, hi = g.es.quantile([0.01, 0.99]); E.loc[g.index[(g.es < lo) | (g.es > hi)], "extreme"] = True
    E.to_csv(FINAL / "effects_all.csv", index=False); E[E.usable].to_csv(FINAL / "effects.csv", index=False)
    print("studies", len(S), "eligible", int(S.eligible.sum()) if "eligible" in S else 0, "coefficients", len(C), "harmonised usable", int(E.usable.sum()), "families", E[E.usable].family.value_counts().to_dict())


if __name__ == "__main__":
    main()