"""Step 4/5. Assemble the three AI codings of the pilot, align them, and draw the human validation sample.

Inputs : data/pilot/v2/coder{1,2,3}_batch{1,2,3}.json   three independent AI coders x three batches
         data/pilot_articles.parquet                      the 150 pilot articles
Outputs: site/files/dataset_ai_coded.xlsx          Coder 1's dataset (the reference dataset used on the Analysis page)
         site/files/dataset_ai_three_coders.xlsx   one sheet per coder, instance_id aligned across coders, plus README
         site/files/human_sample_blank.xlsx        10 rows, coded columns emptied, for human coders
         data/human_sample_key.xlsx                the same 10 rows with Coder 1's codes (instructor key; not published)

Alignment. Each coder numbered its own rows. Rows are matched across coders on (doc_id, enterprise_name without
spaces) for relevant rows and on doc_id alone for relevance-0 rows. A matched row keeps Coder 1's instance_id; a row
only one or two coders produced receives an id ending in -u<n> and is listed on the README sheet as a unit-of-analysis
disagreement. validate.py joins the sheets on the aligned instance_id.
"""
import json, re
from pathlib import Path
import pandas as pd

ROOT = Path(__file__).resolve().parents[1]
DATA, FILES = ROOT / "data", ROOT / "site" / "files"
SAMPLE_N = 10                     # workshop default; see Step 5
SAMPLE_SPLIT = (6, 3, 1)          # confidence-3 rows, confidence<=2 rows, relevance-0 rows
CODE_COLS = ["relevance", "enterprise_name", "sector", "event_type", "leader_visit", "province", "evidence_text", "coder_confidence", "notes"]
ART_COLS = ["doc_id", "article_id", "article_date", "year", "article_category", "article_long_title", "article_url", "article_cleaned_body"]
CODEBOOK = "codebook_nk_enterprises_v2.xlsx"


def load_coder(k):
    rows = []
    for b in (1, 2, 3):
        rows += json.loads((DATA / "pilot" / "v2" / f"coder{k}_batch{b}.json").read_text(encoding="utf-8"))
    df = pd.DataFrame(rows)
    df["own_instance_id"] = df.instance_id
    df["_key"] = [match_key(r) for r in df.itertuples()]
    return df


def match_key(r):
    name = r.enterprise_name if isinstance(r.enterprise_name, str) else ""
    name = re.sub(r"\s+", "", name)
    # strip a parent-combine prefix so "덕천지구탄광련합기업소 제남탄광" matches "제남탄광"
    name = re.sub(r"^.*?련합기업소(?=.+)", "", name)
    return (r.doc_id, name if r.relevance == 1 else "")


def main():
    arts = pd.read_parquet(DATA / "pilot_articles.parquet")
    coders = {k: load_coder(k) for k in (1, 2, 3)}

    # ---- align instance ids across coders
    ids = {}                                   # match key -> aligned instance_id
    for k, c1 in coders[1].set_index("_key").instance_id.items():
        ids[k] = c1
    extra = {}
    for k in (2, 3):
        for key in coders[k]._key:
            if key not in ids:
                extra.setdefault(key[0], [])
                if key not in extra[key[0]]:
                    extra[key[0]].append(key)
    for doc, keys in extra.items():
        for i, key in enumerate(keys, 1):
            ids[key] = f"{doc}-u{i}"
    present = {key: [k for k in (1, 2, 3) if key in set(coders[k]._key)] for key in ids}
    n_all = sum(1 for v in present.values() if len(v) == 3)

    order = ["instance_id", "own_instance_id"] + ART_COLS + CODE_COLS
    out = {}
    for k, df in coders.items():
        df["instance_id"] = df._key.map(ids)
        df = df.drop(columns=["_key"]).merge(arts[ART_COLS], on="doc_id", how="left")
        df["article_date"] = df.article_date.astype(str)
        missing = set(arts.doc_id) - set(df.doc_id)
        if missing:
            print(f"WARNING coder {k}: articles without a row:", sorted(missing))
        out[k] = df[order].sort_values(["doc_id", "instance_id"]).reset_index(drop=True)

    # ---- reference dataset = coder 1
    df1 = out[1].drop(columns=["own_instance_id"])
    readme = pd.DataFrame({"field": ["generated", "source corpus", "pilot sample", "coder", "codebook", "unit", "rows", "articles", "relevant rows", "confidence 3 / 2 / 1", "other coders"],
        "value": [pd.Timestamp.now().strftime("%Y-%m-%d %H:%M"), "corpus_nk_enterprises_2012-2025.xlsx (946 articles)",
                  f"{arts.doc_id.nunique()} articles stratified by year (random_state 2026)",
                  f"AI coder 1 of 3: Claude (Sonnet) sub-agent in Claude Code reading {CODEBOOK}; three batches of 50 articles",
                  CODEBOOK, "one row per named industrial enterprise per article; relevance 0 rows carry no codes",
                  len(df1), df1.doc_id.nunique(), int((df1.relevance == 1).sum()),
                  " / ".join(str(int((df1.coder_confidence == c).sum())) for c in (3, 2, 1)),
                  "dataset_ai_three_coders.xlsx holds coders 1-3 with aligned instance_id for the Step 5 comparison"]})
    with pd.ExcelWriter(FILES / "dataset_ai_coded.xlsx", engine="openpyxl") as w:
        df1.to_excel(w, sheet_name="instances", index=False)
        readme.to_excel(w, sheet_name="README", index=False)

    # ---- three-coder workbook
    unmatched = [{"instance_id": ids[key], "doc_id": key[0], "enterprise_name": key[1] or "(relevance 0 row)",
                  "coders_with_row": ", ".join(f"coder{k}" for k in present[key])} for key in ids if len(present[key]) < 3]
    readme3 = pd.DataFrame({"field": ["generated", "design", "codebook", "rows coder 1 / 2 / 3", "instances produced by all three coders", "instances produced by one or two coders", "alignment rule", "use"],
        "value": [pd.Timestamp.now().strftime("%Y-%m-%d %H:%M"),
                  "three independent Claude (Sonnet) sub-agents, each coding all 150 pilot articles with the same prompt and codebook, without seeing one another's output",
                  CODEBOOK, " / ".join(str(len(out[k])) for k in (1, 2, 3)), n_all, len(unmatched),
                  "matched on doc_id + enterprise name (spaces and parent-combine prefix removed); relevance-0 rows matched on doc_id; unmatched rows get ids ending -u<n>",
                  "python scripts/validate.py --coders dataset_ai_three_coders.xlsx  (one coder per sheet)"]})
    with pd.ExcelWriter(FILES / "dataset_ai_three_coders.xlsx", engine="openpyxl") as w:
        readme3.to_excel(w, sheet_name="README", index=False)
        for k in (1, 2, 3):
            out[k].drop(columns=["article_cleaned_body"]).to_excel(w, sheet_name=f"coder{k}", index=False)
        pd.DataFrame(unmatched).to_excel(w, sheet_name="unmatched_rows", index=False)

    # ---- human validation sample from coder 1
    n_hi, n_md, n_nr = SAMPLE_SPLIT
    rel = df1[df1.relevance == 1]
    hi = rel[rel.coder_confidence == 3].sample(n=n_hi, random_state=7)
    md = rel[rel.coder_confidence <= 2].sample(n=n_md, random_state=7)
    nr = df1[df1.relevance == 0].sample(n=n_nr, random_state=7)
    sample = pd.concat([hi, md, nr]).sort_values("instance_id")
    sample.to_excel(DATA / "human_sample_key.xlsx", index=False)
    blank = sample.copy()
    for c in CODE_COLS:
        blank[c] = None
    blank["coder_name"] = None
    blank = blank[["coder_name"] + [c for c in df1.columns]]
    instr = pd.DataFrame({"instructions": [
        "1. Enter your name in coder_name on every row.",
        "2. Open each article (article_url) or read article_cleaned_body. Code relevance first (1 = a named industrial enterprise is the subject; 0 = not).",
        f"3. If relevance = 1, fill enterprise_name, sector, event_type, leader_visit, province using only the codes in {CODEBOOK} (Dropdown_Values sheet).",
        "4. Paste the sentence(s) that justify your codes into evidence_text; rate coder_confidence 3/2/1.",
        "5. Do not change instance_id or delete rows. Save as human_coded_<yourname>.xlsx and upload it.",
    ]})
    with pd.ExcelWriter(FILES / "human_sample_blank.xlsx", engine="openpyxl") as w:
        blank.to_excel(w, sheet_name="sample", index=False)
        instr.to_excel(w, sheet_name="README", index=False)

    print(f"coder1 rows {len(df1)} articles {df1.doc_id.nunique()} relevant {(df1.relevance==1).sum()} | aligned instances {len(ids)}, all-three {n_all}, unmatched {len(unmatched)} | sample {len(sample)}")
    for k in (1, 2, 3):
        d = out[k][out[k].relevance == 1]
        print(k, d.sector.value_counts().to_dict(), d.event_type.value_counts().to_dict(), d.leader_visit.value_counts().to_dict(), out[k].coder_confidence.value_counts().to_dict())


if __name__ == "__main__":
    main()