"""Phase 4 — measured writing profiles (train split only) → profiles/*.yaml, figures, data/derived/features.parquet."""
from __future__ import annotations
import json, re
from pathlib import Path
import numpy as np, pandas as pd, yaml
from . import config, segmentation
from .provenance import stamp, now_iso

PROFILE_DEFS = {  # name -> (filter description, predicate on the merged feature row)
    "corpus_all": ("all included train documents", lambda r: True),
    "empirical_asset_pricing": ("annotation.subfield == asset_pricing and genre in {empirical, mixed_theory_empirical}", lambda r: r.get("subfield") == "asset_pricing" and r.get("genre") in ("empirical", "mixed_theory_empirical")),
    "corporate_finance": ("annotation.subfield == corporate_finance and genre in {empirical, mixed_theory_empirical}", lambda r: r.get("subfield") == "corporate_finance" and r.get("genre") in ("empirical", "mixed_theory_empirical")),
    "financial_intermediation_banking": ("annotation.subfield == financial_intermediation_banking (empirical/mixed)", lambda r: r.get("subfield") == "financial_intermediation_banking" and r.get("genre") in ("empirical", "mixed_theory_empirical")),
    "household_behavioral_finance": ("subfield in {household_finance, behavioral_finance} (empirical/mixed)", lambda r: r.get("subfield") in ("household_finance", "behavioral_finance") and r.get("genre") in ("empirical", "mixed_theory_empirical")),
    "financial_econometrics": ("subfield == financial_econometrics or genre == methodological_econometric", lambda r: r.get("subfield") == "financial_econometrics" or r.get("genre") == "methodological_econometric"),
    "theoretical_finance": ("genre == theoretical", lambda r: r.get("genre") == "theoretical"),
    "empirical_all": ("genre in {empirical, mixed_theory_empirical}", lambda r: r.get("genre") in ("empirical", "mixed_theory_empirical")),
}
NUMERIC = ["pages", "n_words_body", "n_words_appendix", "abstract_words", "n_sections", "n_subsections", "n_paragraphs", "n_equations", "n_tables", "n_figures", "n_footnotes", "n_references", "n_citations_in_text",
           "intro_words", "intro_paragraphs", "intro_share", "conclusion_words", "conclusion_share", "theory_share", "data_share", "results_share", "robustness_share", "words_per_table", "refs_per_1000_words", "equations_per_1000_words"]

def doc_features(doc_id: str) -> dict:
    d = segmentation.load_parsed(doc_id); st = d["stats"]
    roles = {}
    for s in d["sections"]:
        role = segmentation.section_role(s["heading"], s.get("is_appendix", False), s.get("is_intro", False)); roles[role] = roles.get(role, 0) + s["n_words"]
    body = max(st["n_words_body"], 1); intro = next((s for s in d["sections"] if s.get("is_intro")), None)
    f = {"doc_id": doc_id, "year": d["metadata"]["year"], "article_type": d["metadata"]["article_type"], "pages": d["pages"], **{k: st[k] for k in ["n_words_body", "n_words_appendix", "abstract_words", "n_sections", "n_subsections", "n_paragraphs", "n_equations", "n_tables", "n_figures", "n_footnotes", "n_references", "n_citations_in_text"]},
         "intro_words": roles.get("introduction", 0), "intro_paragraphs": len(segmentation.paragraphs(intro)) if intro else 0, "intro_share": roles.get("introduction", 0) / body,
         "conclusion_words": roles.get("conclusion", 0), "conclusion_share": roles.get("conclusion", 0) / body, "theory_share": roles.get("theory", 0) / body, "data_share": roles.get("data", 0) / body,
         "results_share": roles.get("results", 0) / body, "robustness_share": roles.get("robustness", 0) / body, "has_robustness_section": roles.get("robustness", 0) > 0, "has_literature_section": roles.get("literature", 0) > 0,
         "has_appendix": st["n_words_appendix"] > 0, "has_internet_appendix": d["markers"]["has_internet_appendix"], "words_per_table": body / max(st["n_tables"], 1), "refs_per_1000_words": 1000 * st["n_references"] / body, "equations_per_1000_words": 1000 * st["n_equations"] / body,
         "section_headings": " | ".join(f"{s.get('number') or ''} {s['heading']}".strip() for s in d["sections"]), "parse_confidence": d["confidence"]["overall"]}
    return f

def merged_features(split: str = "train") -> pd.DataFrame:
    sp = json.loads((config.ROOT / "data" / "splits" / "heldout_v1.json").read_text())
    ids = sp[split] if split in ("train", "heldout") else sp["train"] + sp["heldout"]
    df = pd.DataFrame([doc_features(i) for i in ids])
    ann_p = config.ROOT / "data" / "annotations" / "annotations_table.parquet"
    if ann_p.exists():
        ann = pd.read_parquet(ann_p); df = df.merge(ann.drop(columns=[c for c in ["model", "cost_usd"] if c in ann]), on="doc_id", how="left")
    df["split"] = split
    return df

def q(s: pd.Series) -> dict:
    s = pd.to_numeric(s, errors="coerce").dropna()
    if len(s) == 0: return {"n": 0}
    rng = np.random.default_rng(45); boots = [float(np.median(rng.choice(s.values, len(s)))) for _ in range(500)] if len(s) >= 5 else [float(s.median())]
    return {"n": int(len(s)), "median": round(float(s.median()), 3), "q25": round(float(s.quantile(.25)), 3), "q75": round(float(s.quantile(.75)), 3), "mean": round(float(s.mean()), 3), "min": round(float(s.min()), 3), "max": round(float(s.max()), 3),
            "median_ci95_bootstrap": [round(float(np.percentile(boots, 2.5)), 3), round(float(np.percentile(boots, 97.5)), 3)]}

def share(s: pd.Series) -> dict:
    s = s.dropna()
    return {"n": int(len(s)), "share": round(float(s.astype(bool).mean()), 3)} if len(s) else {"n": 0}

def cat_dist(s: pd.Series, top=8) -> dict:
    s = s.dropna(); n = len(s)
    return {"n": int(n), **{str(k): round(float(v) / n, 3) for k, v in s.value_counts().head(top).items()}} if n else {"n": 0}

def seq_dist(s: pd.Series, top=6) -> dict:
    """Most common introduction move orders (deduplicated consecutive moves)."""
    seqs = []
    for x in s.dropna():
        mv = [m for i, m in enumerate(x.split("|")) if i == 0 or m != x.split("|")[i - 1]]
        seqs.append(" > ".join(mv))
    s2 = pd.Series(seqs); return cat_dist(s2, top)

def move_prevalence(s: pd.Series) -> dict:
    """Share of introductions containing each move at least once, and median relative position of its first occurrence."""
    s = s.dropna(); n = len(s); out = {"n": int(n)}
    if not n: return out
    from collections import Counter
    c = Counter(); pos = {}
    for x in s:
        mv = x.split("|")
        for m in set(mv): c[m] += 1; pos.setdefault(m, []).append(mv.index(m) / max(len(mv) - 1, 1))
    for m, k in c.most_common(): out[m] = {"share": round(k / n, 3), "median_relative_position": round(float(np.median(pos[m])), 2)}
    return out

def first_move_dist(s: pd.Series) -> dict:
    return cat_dist(s.dropna().map(lambda x: x.split("|")[0] if x else None))

def build_profile(name: str, df: pd.DataFrame) -> dict:
    desc, pred = PROFILE_DEFS[name]; sub = df[[pred(r) for r in df.to_dict("records")]]
    has_ann = "genre" in sub and sub["genre"].notna().any()
    prof = {"profile": name, "ontology_version": "1.0", "built": stamp(), "selection_rule": desc, "n_documents": int(len(sub)), "split": "train (heldout_v1 excluded)",
            "caveat": "Observed frequencies describe what published JoF articles did in 2017–2020; they are not evidence that these choices are editorially superior or required.",
            "structure": {k: q(sub[k]) for k in NUMERIC if k in sub},
            "presence": {k: share(sub[k]) for k in ["has_robustness_section", "has_literature_section", "has_appendix", "has_internet_appendix"] if k in sub},
            "by_year": {int(y): {"n": int(len(g)), "n_words_body_median": float(g["n_words_body"].median()), "n_tables_median": float(g["n_tables"].median())} for y, g in sub.groupby("year")}}
    if has_ann:
        a = sub[sub["schema_valid"] == True]
        prof["rhetoric"] = {"n_annotated": int(len(a)), "genre": cat_dist(a["genre"]), "subfield": cat_dist(a["subfield"]), "primary_method": cat_dist(a["primary_method"]),
            "abstract_moves_present": {k: share(a[f"abs_{k}"]) for k in ["question", "gap", "method", "results", "magnitude", "contribution"]},
            "abstract_move_order_median_sentence": {k: q(a[a[f"abs_{k}_idx"] >= 0][f"abs_{k}_idx"]) for k in ["question", "gap", "method", "results", "magnitude", "contribution"]},
            "abstract_sentences": q(a["abstract_n_sent"]), "magnitude_in_abstract_regex": share(a["magnitude_regex_abstract"]),
            "introduction": {"paragraphs": q(a["intro_n_paragraphs"]), "words": q(a["intro_n_words"]), "contribution_paragraph_index": q(a[a["contribution_para_idx"] >= 0]["contribution_para_idx"]),
                             "contribution_relative_position": q((a[a["contribution_para_idx"] >= 0]["contribution_para_idx"] / a[a["contribution_para_idx"] >= 0]["intro_n_paragraphs"].clip(lower=1))),
                             "first_move": first_move_dist(a["intro_moves_seq"]), "move_prevalence": move_prevalence(a["intro_moves_seq"]), "moves_per_intro": q(a["intro_n_moves"]), "common_move_sequences": seq_dist(a["intro_moves_seq"]), "roadmap_present_llm": share(a["roadmap_llm"]), "roadmap_present_regex": share(a["roadmap_regex"]),
                             "literature_positioning_location": cat_dist(a["lit_positioning"]), "contribution_types": cat_dist(a["contribution_types"].dropna().str.split("|").explode())},
            "identification": {"causal_language_strength": q(a["causal_strength"]), "uses_causal_claim": share(a["uses_causal_claim"])},
            "results": {"economic_significance_discussed": share(a["econ_signif"]), "statistical_language": cat_dist(a["stat_lang"]), "magnitude_statements_in_intro": share(a["magnitude_in_intro"])},
            "robustness": {"location": cat_dist(a["robust_location"]), "types": cat_dist(a["robust_types"].dropna().str.split("|").explode())},
            "tables": {"standard_errors": cat_dist(a["se_desc"]), "units_in_notes": share(a["units_in_notes"]), "clustered_regex": share(a["clustered_se_regex"])},
            "figures": {"uncertainty_shown": cat_dist(a["fig_uncertainty"])},
            "conclusion": {k: share(a[c]) for k, c in [("bounded_claims", "concl_bounded"), ("policy_implications", "concl_policy"), ("open_questions", "concl_open")]},
            "hedging": cat_dist(a["hedging"]), "claims_with_stated_evidence_share": q(a["claims_with_evidence"] / a["n_claims"].clip(lower=1)),
            "llm_vs_rule_section_role_agreement": q(a["role_agreement"]), "llm_confidence": q(a["llm_confidence"])}
    return prof

def build_profiles() -> dict:
    df = merged_features("train"); allf = pd.concat([df, merged_features("heldout")]); allf.to_parquet(config.derived_dir() / "features.parquet", index=False); allf.to_csv(config.derived_dir() / "features.csv", index=False)
    out = {}; pdir = config.ROOT / "profiles"; pdir.mkdir(exist_ok=True)
    for name in PROFILE_DEFS:
        prof = build_profile(name, df); out[name] = prof["n_documents"]
        (pdir / f"{name}.yaml").write_text("# Generated by `jof250 profiles build` — do not edit by hand (see docs/DATA_DICTIONARY.md)\n" + yaml.safe_dump(prof, sort_keys=False, allow_unicode=True, width=120))
    make_figures(df)
    return {"profiles": out, "n_train": int(len(df)), "annotated_in_train": int(df["schema_valid"].fillna(False).sum()) if "schema_valid" in df else 0}

def make_figures(df: pd.DataFrame) -> None:
    from . import viz; viz.style(); import matplotlib.pyplot as plt
    fdir = config.ROOT / "figures"; fdir.mkdir(exist_ok=True); pf = config.ROOT / "paper" / "figures"; pf.mkdir(exist_ok=True)
    # 1. structural distributions
    fig, axes = plt.subplots(2, 3, figsize=(7.2, 4.2))
    for ax, (col, lab) in zip(axes.flat, [("n_words_body", "Body words"), ("n_tables", "Tables"), ("n_figures", "Figures"), ("n_equations", "Numbered equations"), ("n_references", "References"), ("intro_share", "Introduction share of body")]):
        ax.hist(df[col].dropna(), bins=20, color=viz.PALETTE[0], edgecolor=viz.SURFACE, linewidth=0.5); ax.set_title(lab, fontsize=9); ax.axvline(df[col].median(), color=viz.PALETTE[1], linewidth=1.2)
    fig.suptitle(f"Structural distributions, JoF train split (n = {len(df)}); orange line = median", fontsize=9); fig.tight_layout()
    for p in (fdir, pf): fig.savefig(p / "fig_structure_distributions.pdf"); fig.savefig(p / "fig_structure_distributions.png")
    plt.close(fig)
    if "genre" in df and df["genre"].notna().any():
        a = df[df["schema_valid"] == True]
        # 2. abstract moves prevalence by genre
        moves = ["question", "gap", "method", "results", "magnitude", "contribution"]; genres = [g for g in ["empirical", "mixed_theory_empirical", "theoretical"] if (a["genre"] == g).sum() >= 5]
        fig, ax = plt.subplots(figsize=(6.2, 2.8)); w = 0.8 / max(len(genres), 1)
        for i, g in enumerate(genres):
            sub = a[a["genre"] == g]; vals = [sub[f"abs_{m}"].mean() for m in moves]
            ax.bar(np.arange(len(moves)) + i * w, vals, w, label=f"{g} (n={len(sub)})", color=viz.PALETTE[i], edgecolor=viz.SURFACE, linewidth=0.8)
        ax.set_xticks(np.arange(len(moves)) + w * (len(genres) - 1) / 2); ax.set_xticklabels(moves); ax.set_ylabel("Share of abstracts"); ax.set_ylim(0, 1); ax.legend(fontsize=7); ax.set_title("Abstract moves present, by genre (LLM annotation, train split)", fontsize=9)
        fig.tight_layout(); [fig.savefig(p / "fig_abstract_moves.pdf") for p in (fdir, pf)]; plt.close(fig)
        # 3. contribution relative position
        rel = (a[a["contribution_para_idx"] >= 0]["contribution_para_idx"] / a[a["contribution_para_idx"] >= 0]["intro_n_paragraphs"].clip(lower=1))
        fig, ax = plt.subplots(figsize=(4.2, 2.6)); ax.hist(rel, bins=10, range=(0, 1), color=viz.PALETTE[0], edgecolor=viz.SURFACE); ax.set_xlabel("Relative position of the contribution paragraph in the introduction"); ax.set_ylabel("Documents")
        ax.set_title(f"Contribution placement (n = {len(rel)})", fontsize=9); fig.tight_layout(); [fig.savefig(p / "fig_contribution_position.pdf") for p in (fdir, pf)]; plt.close(fig)
        # 4. intro first moves + robustness location
        fig, axes = plt.subplots(1, 2, figsize=(7.2, 2.6))
        fm = a["intro_moves_seq"].dropna().map(lambda x: x.split("|")[0]).value_counts(normalize=True).head(6)
        axes[0].barh(fm.index[::-1], fm.values[::-1], color=viz.PALETTE[0]); axes[0].set_title("First rhetorical move of the introduction", fontsize=9); axes[0].set_xlabel("Share")
        rl = a["robust_location"].value_counts(normalize=True).head(6)
        axes[1].barh(rl.index[::-1], rl.values[::-1], color=viz.PALETTE[2]); axes[1].set_title("Where robustness checks are presented", fontsize=9); axes[1].set_xlabel("Share")
        fig.tight_layout(); [fig.savefig(p / "fig_moves_robustness.pdf") for p in (fdir, pf)]; plt.close(fig)