#!/usr/bin/env python3
"""
code_corpus.py -- automated (script-computed) content-coding variables for the
WP29 mandate corpus (research/corpus/*.md).

Standard library only.  Run from anywhere:

    python3 -I research/code_corpus.py

Writes research/corpus_automated.csv (one row per file, including the
byte-identical WP15 duplicate, which is flagged in `duplicate_of`).

Every variable here is AUTOMATED (computed from the bytes of the file).  The
definitions below are the authoritative ones referenced by
research/corpus_coding_protocol.md.  Keyword counts are *mentions*, not
judgements: a document that says "no budget is specified" would still count
one mention of the `budget` family.
"""

import csv
import hashlib
import os
import re
import sys

HERE = os.path.dirname(os.path.abspath(__file__))
CORPUS_DIR = os.path.join(HERE, "corpus")
MANIFEST = os.path.join(HERE, "corpus_manifest.csv")
OUT = os.path.join(HERE, "corpus_automated.csv")

# --------------------------------------------------------------------------
# Regular expressions (documented for the codebook)
# --------------------------------------------------------------------------

# Markdown ATX headings: 1-6 '#' at line start followed by a space.
RE_HEADING = re.compile(r"^#{1,6}\s+\S", re.MULTILINE)

# Explicit phase labels anywhere in the text, e.g. "Phase 0", "PHASE 3 —",
# "phase_XX" is NOT matched (requires digits).  Distinct numbers are counted.
RE_PHASE_LABEL = re.compile(r"\bphase\s*(\d+)", re.IGNORECASE)

# Lettered stage labels ("Stage A", "Stage G") -- WP15 uses stages, not phases.
# Case-insensitive on the word "stage", but the label itself must be a capital
# letter.  NB: WP28 also uses "Stage A..G" for its coordinator/worker message
# protocol; because WP28 has explicit "Phase N" labels those stages are not
# used as its phase count.
RE_STAGE_LABEL = re.compile(r"(?i:\bstage)\s+([A-Z])\b")

# Gate labels ("Gate 1", "Gate 12") -- distinct numbers counted.
RE_GATE_LABEL = re.compile(r"\bgate\s*(\d+)", re.IGNORECASE)

# Numbered list items at line start ("1. ", "  12. ").
RE_NUMBERED_ITEM = re.compile(r"^\s*\d+\.\s", re.MULTILINE)

# Headings whose text indicates a phase/sequence/execution-order section.  The
# numbered items *inside* such sections are counted as declared phases when the
# document has no explicit "Phase N" labels.
RE_PHASE_SECTION_HEADING = re.compile(
    r"^#{1,6}\s+.*\b(phase|phases|stage|stages|sequence|execution order|"
    r"work sequence|execution protocol|execution phases|milestones)\b",
    re.IGNORECASE | re.MULTILINE,
)

# Markdown table separator rows: |---|---:| etc.  One separator row == one table.
RE_TABLE_SEP = re.compile(r"^\s*\|?\s*:?-{3,}:?\s*(\|\s*:?-{3,}:?\s*)+\|?\s*$", re.MULTILINE)

# Fenced code blocks: pairs of ``` fences.
RE_FENCE = re.compile(r"^\s*```", re.MULTILINE)

# Keyword families (case-insensitive).  Each is a mention count + boolean.
KEYWORD_FAMILIES = {
    "preregist": r"pre-?regist\w*",
    "estimand": r"\bestimands?\b",
    "hypothes": r"\bhypothes\w*",
    "falsif": r"\bfalsif\w*",
    "baseline": r"\bbaselines?\b",
    "leakage": r"\bleakage\b|\blook-?ahead\b|\blook ahead\b|\bpoint-in-time\b|\bpoint in time\b",
    "provenance": r"\bprovenance\b|\bhash(?:es|ed|ing)?\b|\bsha-?256\b|\bchecksums?\b",
    "reproduc": r"\breproduc\w*",
    "stopping": r"\bstopping\b|\bstop rules?\b|\bstop conditions?\b|\btermination\b|\bstop and\b|\bstop if\b|\bstop only\b",
    "recovery": r"\brecover\w*|\bresum\w*|\bcheckpoint\w*",
    "gate": r"\bgates?\b|\bacceptance\b|\bcriteria\b",
    "exhibit": r"\bscreenshots?\b|\bfigures?\b|\btables?\b",
    "disclos": r"\bdisclos\w*|\bAI-use\b|\bAI use\b|\bAI contribution\b",
    "budget": r"\bbudgets?\b|\bcosts?\b|\bUSD\b|\$",
    "license": r"\blicen[cs]e[sd]?\b|\blicensing\b",
    "seed": r"\bseeds?\b|\bseeded\b",
}
KEYWORD_RE = {k: re.compile(v, re.IGNORECASE) for k, v in KEYWORD_FAMILIES.items()}

# Language heuristic: tokens that are common in French prose and rare in
# English technical prose ("et" is excluded because of "et al.").
# Tokens shared with English usage in this corpus ("conventions", "journal",
# "du" as in "Universite du Quebec", "et" as in "et al.") are deliberately excluded.
FR_TOKENS = {
    "le", "la", "les", "des", "est", "une", "ce", "cette", "dans", "pour",
    "sur", "avec", "pas", "aux", "qui", "que", "ci-dessous", "dossier",
    "appelle", "correspond", "auteur",
}
EN_TOKENS = {
    "the", "and", "of", "to", "in", "is", "for", "with", "that", "not", "must",
    "do", "never", "each", "every", "all", "or", "as", "from", "if", "use",
}
RE_WORD = re.compile(r"[A-Za-zÀ-ÿ][A-Za-zÀ-ÿ'\-]*")


def sha256_of(path):
    h = hashlib.sha256()
    with open(path, "rb") as f:
        for chunk in iter(lambda: f.read(1 << 20), b""):
            h.update(chunk)
    return h.hexdigest()


def language_of(text):
    """Return (label, fr_share).  fr_share = FR hits / (FR hits + EN hits).

    en    : fr_share < 0.015
    mixed : 0.015 <= fr_share < 0.50  (e.g. a French preamble on an English body)
    fr    : fr_share >= 0.50
    (threshold chosen after inspection: WP18 = 0.043 and WP20 = 0.020 carry a
    French prefatory paragraph; all other files score 0.000)
    """
    toks = [t.lower() for t in RE_WORD.findall(text)]
    fr = sum(1 for t in toks if t in FR_TOKENS)
    en = sum(1 for t in toks if t in EN_TOKENS)
    total = fr + en
    share = (fr / total) if total else 0.0
    if share >= 0.50:
        label = "fr"
    elif share >= 0.015:
        label = "mixed"
    else:
        label = "en"
    return label, round(share, 4)


def numbered_items_in_phase_sections(text):
    """Count numbered list items inside sections whose heading mentions
    phase/stage/sequence/execution order.  A section runs from its heading to
    the next heading of the same or higher level."""
    lines = text.splitlines()
    heading_idx = []
    for i, line in enumerate(lines):
        m = re.match(r"^(#{1,6})\s+(.*)$", line)
        if m:
            heading_idx.append((i, len(m.group(1)), m.group(2)))
    count = 0
    for n, (i, level, title) in enumerate(heading_idx):
        if not RE_PHASE_SECTION_HEADING.match("#" * level + " " + title):
            continue
        end = len(lines)
        for j, lvl, _ in heading_idx[n + 1:]:
            if lvl <= level:
                end = j
                break
        block = "\n".join(lines[i + 1:end])
        count += len(RE_NUMBERED_ITEM.findall(block))
    return count


def analyse(path):
    raw = open(path, "rb").read()
    text = raw.decode("utf-8", errors="replace")
    name = os.path.basename(path)
    wp = re.match(r"(WP\d+)", name).group(1)

    phase_numbers = {int(x) for x in RE_PHASE_LABEL.findall(text)}
    stage_letters = set(RE_STAGE_LABEL.findall(text))
    gate_numbers = {int(x) for x in RE_GATE_LABEL.findall(text)}
    items_in_phase_sections = numbered_items_in_phase_sections(text)

    # n_numbered_phases: distinct "Phase N" labels if any; else distinct
    # "Stage X" labels if any; else distinct "Gate N" labels if any; else
    # numbered items within phase/sequence sections.  (Analyst-coded
    # n_phases_declared in corpus_coding.csv is the authoritative value; this
    # is the mechanical proxy.)
    if phase_numbers:
        n_numbered_phases = len(phase_numbers)
        phase_rule = "distinct 'Phase N' labels"
    elif stage_letters:
        n_numbered_phases = len(stage_letters)
        phase_rule = "distinct 'Stage X' labels"
    elif gate_numbers:
        n_numbered_phases = len(gate_numbers)
        phase_rule = "distinct 'Gate N' labels"
    else:
        n_numbered_phases = items_in_phase_sections
        phase_rule = "numbered items inside phase/sequence sections"

    lang, fr_share = language_of(text)
    words = len(text.split())
    row = {
        "file": name,
        "wp": wp,
        "sha256": hashlib.sha256(raw).hexdigest(),
        "bytes": len(raw),
        "lines": text.count("\n") + (0 if text.endswith("\n") else 1),
        "words": words,
        "n_headings": len(RE_HEADING.findall(text)),
        "n_numbered_phases": n_numbered_phases,
        "n_numbered_phases_rule": phase_rule,
        "n_phase_labels_distinct": len(phase_numbers),
        "n_stage_labels_distinct": len(stage_letters),
        "n_gate_labels_distinct": len(gate_numbers),
        "n_numbered_items_in_phase_sections": items_in_phase_sections,
        "n_numbered_items_total": len(RE_NUMBERED_ITEM.findall(text)),
        "n_tables": len(RE_TABLE_SEP.findall(text)),
        "n_code_blocks": len(RE_FENCE.findall(text)) // 2,
        "language_of_document": lang,
        "fr_token_share": fr_share,
    }
    for fam, rx in KEYWORD_RE.items():
        n = len(rx.findall(text))
        row["mentions_%s_n" % fam] = n
        row["mentions_%s" % fam] = "yes" if n else "no"
    return row


def main():
    files = sorted(
        os.path.join(CORPUS_DIR, f) for f in os.listdir(CORPUS_DIR) if f.endswith(".md")
    )
    if not files:
        sys.exit("no .md files in %s" % CORPUS_DIR)

    # manifest lookup (url, retrieved_utc, expected sha)
    manifest = {}
    if os.path.exists(MANIFEST):
        with open(MANIFEST, newline="", encoding="utf-8") as f:
            for r in csv.DictReader(f):
                if r.get("local_file"):
                    manifest[os.path.basename(r["local_file"])] = r

    rows = [analyse(p) for p in files]

    # duplicates by sha256 (first file in sorted order is the canonical one)
    seen = {}
    for r in rows:
        r["duplicate_of"] = seen.get(r["sha256"], "")
        seen.setdefault(r["sha256"], r["file"])
        m = manifest.get(r["file"], {})
        r["manifest_url"] = m.get("url", "")
        r["manifest_retrieved_utc"] = m.get("retrieved_utc", "")
        r["sha256_matches_manifest"] = (
            "yes" if m and m.get("sha256") == r["sha256"] else ("no" if m else "not_in_manifest")
        )
        r["document_type"] = (
            "design_brief" if "BRIEF" in r["file"]
            else ("duplicate_claude_md" if r["duplicate_of"] else "claude_md")
        )

    fieldnames = (
        ["file", "wp", "document_type", "duplicate_of", "sha256", "sha256_matches_manifest",
         "manifest_url", "manifest_retrieved_utc", "bytes", "lines", "words", "n_headings",
         "n_numbered_phases", "n_numbered_phases_rule", "n_phase_labels_distinct",
         "n_stage_labels_distinct", "n_gate_labels_distinct",
         "n_numbered_items_in_phase_sections", "n_numbered_items_total", "n_tables",
         "n_code_blocks", "language_of_document", "fr_token_share"]
        + ["mentions_%s" % k for k in KEYWORD_FAMILIES]
        + ["mentions_%s_n" % k for k in KEYWORD_FAMILIES]
    )
    with open(OUT, "w", newline="", encoding="utf-8") as f:
        w = csv.DictWriter(f, fieldnames=fieldnames, quoting=csv.QUOTE_MINIMAL)
        w.writeheader()
        for r in rows:
            w.writerow({k: r.get(k, "") for k in fieldnames})

    # console summary
    print("wrote %s (%d rows)" % (OUT, len(rows)))
    for r in rows:
        print("%-38s %6d w  %3d h  phases=%-2s tables=%d code=%d lang=%s dup=%s" % (
            r["file"], r["words"], r["n_headings"], r["n_numbered_phases"], r["n_tables"],
            r["n_code_blocks"], r["language_of_document"], r["duplicate_of"] or "-"))


if __name__ == "__main__":
    main()