#!/usr/bin/env python3
"""AI くささ検証 段階1: 文章の「目に見える癖」を数える。

入力: JSONL（1 行 1 文章。text、または --normalize-html 指定時は htmlPath が必須）
出力: 文章ごとの癖の数（1,000 字あたり）の CSV と、条件別の平均の表（Markdown）。
      --classify を付けると、癖の数だけで「どの会社の文章か」を当てられるかを
      お題を 1 本ずつ抜く交差検証（最近傍重心）で測る。

  python3 -I ai_style_count.py texts.jsonl --out out/ [--classify]
  python3 -I ai_style_count.py --from-headless <ヘッドレスの結果フォルダ> --out out/   # 既存の単体回答で試す

癖の定義は MARKERS と count_text() が正本。足す・直すときは PLAN.md も同時に直す。
"""
import argparse
import csv
import json
import math
import os
import re
import statistics
import sys
from collections import defaultdict

# (id, 区分, 説明, 正規表現, 数え方)  数え方: "match"=一致数 / "line"=当てはまる行数
MARKERS = [
    # 定型句
    ("P01", "定型句", "結論から言うと／結論として", r"結論から(言う|申し上げる)と|結論として", "match"),
    ("P02", "定型句", "いかがでしたか", r"いかがでしたか", "match"),
    ("P03", "定型句", "〜と言えるでしょう／〜と言えます", r"と言えるでしょう|と言えます", "match"),
    ("P04", "定型句", "〜ではないでしょうか", r"ではないでしょうか", "match"),
    ("P05", "定型句", "〜が重要です／大切です／鍵です", r"(が|は)(非常に|とても)?(重要|大切|鍵|カギ)(です|となります|になります)", "match"),
    ("P06", "定型句", "さまざまな／多様な", r"さまざまな|様々な|多様な", "match"),
    ("P07", "定型句", "まとめると／最後に／いずれにせよ", r"まとめると|最後に|いずれにせよ|総じて", "match"),
    ("P08", "定型句", "〜することができます", r"することができ(ます|る)", "match"),
    ("P09", "定型句", "〜において", r"において", "match"),
    ("P10", "定型句", "ぜひ〜してみてください", r"ぜひ.{0,20}(てみてください|ください)", "match"),
    # 修辞
    ("R01", "修辞", "AだけでなくBも", r"だけでなく|のみならず", "match"),
    ("R02", "修辞", "AではなくB（対比の言い切り）", r"ではなく、", "match"),
    ("R03", "修辞", "問いかけで始める（行頭の〜か？）", r"^[^\n]{0,40}(でしょうか|ますか|ですか)[？?]", "line"),
    ("R04", "修辞", "〜のです／〜なのです", r"(の|な)のです", "match"),
    # ぼかし
    ("H01", "ぼかし", "〜かもしれません／可能性があります", r"かもしれません|可能性があります", "match"),
    ("H02", "ぼかし", "〜と考えられます／〜と思われます", r"と考えられます|と思われます", "match"),
    # 書式
    ("F01", "書式", "見出し（# または【】行）", r"^\s*(#{1,6}\s|【[^】]+】\s*$)", "line"),
    ("F02", "書式", "箇条書きの行", r"^\s*([-*・•]|\d+[.)．、])\s*", "line"),
    ("F03", "書式", "太字（**）", r"\*\*[^*\n]+\*\*", "match"),
    ("F04", "書式", "絵文字", "[\U0001F300-\U0001FAFF☀-➿]", "match"),
    ("F05", "書式", "ダッシュ（——・―）", r"——|―|—", "match"),
    ("F06", "書式", "行頭の見出し風コロン（〜：）", r"^[^\n。]{1,20}[：:]\s*\S", "line"),
    ("F07", "書式", "鉤括弧での強調（「」）", r"「[^」\n]{1,15}」", "match"),
]

# v1 は生成前に固定した正本として保持。v1.1 は誤検出確認後の修正版。
SEPARATOR_PATTERN = r"^\s*[-‐‑‒–—―−_]+\s*$"
MARKERS_V11 = []
for marker in MARKERS:
    mid, cat, desc, pat, how = marker
    if mid == "F02":
        pat = r"^\s*(?:[-*•]\s+|・\s*|\d+[.)．、]\s+)"
    elif mid == "F05":
        MARKERS_V11.extend([
            ("F05a", cat, "区切り線（ダッシュ類・---・___ だけの行）", SEPARATOR_PATTERN, "line"),
            ("F05b", cat, "文中のダッシュ（区切り線以外・連続を1回）", r"[—―]+", "inline_dash"),
        ])
        continue
    elif mid == "F06":
        # 行頭の一語のラベルのみ。時刻・URL・文中や括弧内のラベルを除く。
        pat = r"^\s*(?!(?i:https?|ftp|mailto|file):)[^\W\d_][\w・／/ー-]{0,19}[：:]\s*(?!//)\S"
    elif mid == "F07":
        desc = "鉤括弧の数（強調に限らない）"
    MARKERS_V11.append((mid, cat, desc, pat, how))


def markers_for(version="v1"):
    if version not in ("v1", "v1.1"):
        raise ValueError("markers は v1 または v1.1")
    return MARKERS if version == "v1" else MARKERS_V11


SENT_SPLIT = re.compile(r"(?<=[。！？!?])|\n+")


def strip_ws(text):
    return re.sub(r"\s", "", text)


def sentences(text):
    out = []
    for s in SENT_SPLIT.split(text):
        s = s.strip()
        if s and not re.fullmatch(r"[#\-*・\d.)．、\s]*", s):
            out.append(s)
    return out


def ending_of(sent):
    """文末の型（です/ます/た/その他の文末…）。同じ型の連続を見るため。"""
    s = re.sub(r"[。！？!?」』）)\s]+$", "", sent)
    for pat, tag in [
        (r"でしょう$", "でしょう"), (r"ません$", "ません"), (r"ました$", "ました"),
        (r"ます$", "ます"), (r"でした$", "でした"), (r"です$", "です"),
        (r"(だ|である)$", "だ"), (r"た$", "た"), (r"[るうくすつぬむぶぐ]$", "る"),
    ]:
        if re.search(pat, s):
            return tag
    return "その他の文末"


def paragraphs(text):
    """空行で分け、見出し・箇条書きの行だけのブロックを除く。"""
    out = []
    for block in re.split(r"\n\s*\n", text):
        lines = [line for line in block.splitlines() if line.strip()]
        if lines and not all(re.match(r"^\s*(?:#|[-*・]|\d+\.\s)", line)
                             for line in lines):
            out.append(block.strip())
    return out


def count_text(text, markers="v1"):
    n_chars = max(1, len(strip_ws(text)))
    lines = text.splitlines()
    row = {"chars": n_chars}
    for mid, _cat, _desc, pat, how in markers_for(markers):
        if how == "line":
            rx = re.compile(pat)
            c = sum(1 for ln in lines if rx.search(ln))
        elif how == "inline_dash":
            c = sum(len(re.findall(pat, ln)) for ln in lines
                    if not re.fullmatch(SEPARATOR_PATTERN, ln))
        else:
            c = len(re.findall(pat, text, flags=re.M))
        row[mid] = round(c * 1000 / n_chars, 3)
    sents = sentences(text)
    lens = [len(strip_ws(s)) for s in sents] or [0]
    row["S01_mean_sent_len"] = round(statistics.mean(lens), 2)
    row["S02_sd_sent_len"] = round(statistics.pstdev(lens), 2)
    ends = [ending_of(s) for s in sents]
    same = sum(1 for a, b in zip(ends, ends[1:]) if a == b)
    row["S03_same_ending_rate"] = round(same / max(1, len(ends) - 1), 3)
    if markers == "v1":
        row["S04_taigen_rate"] = round(ends.count("その他の文末") / max(1, len(ends)), 3)
    mean_len = statistics.mean(lens)
    row["S05_cv_sent_len"] = round(statistics.pstdev(lens) / mean_len, 3) if len(sents) >= 2 else 0
    row["S06_adj_len_change"] = round(
        statistics.mean(abs(a - b) for a, b in zip(lens, lens[1:])) / mean_len, 3
    ) if len(sents) >= 2 else 0
    para_counts = [len(sentences(p)) for p in paragraphs(text)]
    mean_para_sents = statistics.mean(para_counts) if para_counts else 0
    row["S07_para_sent_cv"] = round(
        statistics.pstdev(para_counts) / mean_para_sents, 3
    ) if len(para_counts) >= 2 and mean_para_sents else 0
    row["S08_mean_para_sents"] = round(mean_para_sents, 3)
    # ending_of() と同じ文末記号を除いて判定。最後の文は後続がないため除く。
    neg_then_pos = sum(bool(re.search(
        r"(?:ではない|じゃない|ではありません|わけではない|わけではありません|だけではない|だけではありません)$",
        re.sub(r"[。！？!?」』）)\s]+$", "", s),
    )) for s in sents[:-1])
    row["R05_neg_then_pos"] = round(neg_then_pos * 1000 / n_chars, 3)
    return row


def excerpt_at_sentence_end(text, limit):
    """空白を除いて limit 字を超えた後、最初の文末または改行まで残す。"""
    n = 0
    for i, char in enumerate(text):
        if not char.isspace():
            n += 1
        if n > limit and (char in "。！？!?" or char == "\n"):
            return text[:i + 1].rstrip()
    return text


RHYTHM_DESCRIPTIONS = {
    "S05_cv_sent_len": "文長（空白を除いた字数）の変動係数（母標準偏差 ÷ 平均）。文が2未満なら0。",
    "S06_adj_len_change": "隣り合う文長の差の絶対値の平均 ÷ 文長の平均。文が2未満なら0。",
    "S07_para_sent_cv": "空行区切りの段落ごとの文数の変動係数。見出し（#）・箇条書き（- / * / ・ / 数字. ）だけのブロックは除外。段落が2未満なら0。",
    "S08_mean_para_sents": "S07と同じ段落の文数の平均。段落が0なら0。",
    "R05_neg_then_pos": "否定（ではない／じゃない／ではありません／わけではない／わけではありません／だけではない／だけではありません）で終わる文に後続文がある組の数（1,000字あたり）。後続文の肯否は判定しない。",
}

FEATURES = [m[0] for m in MARKERS] + ["S01_mean_sent_len", "S02_sd_sent_len", "S03_same_ending_rate", "S04_taigen_rate"] + list(RHYTHM_DESCRIPTIONS)


def features_for(markers="v1"):
    return [m[0] for m in markers_for(markers)] + [
        f for f in FEATURES[len(MARKERS):]
        if markers == "v1" or f != "S04_taigen_rate"
    ]


def rhythm_descriptions_for(markers="v1"):
    desc = dict(RHYTHM_DESCRIPTIONS)
    if markers == "v1.1":
        desc["R05_neg_then_pos"] = (
            "否定で終わる文の数（直後に文があるもの）。対象の否定語尾と計数はv1と同じ。"
            "1,000字あたり。後続文の肯否・言い直しかどうかは判定しない。")
    return desc


def load_jsonl(path):
    with open(path, encoding="utf-8") as f:
        return [json.loads(l) for l in f if l.strip()]


def load_headless(dirpath):
    """既存のヘッドレス討論の結果から単体回答（baseline）だけを取り出す（数え方の試運転用）。"""
    out = []
    for name in sorted(os.listdir(dirpath)):
        if not re.match(r"\d{3}-.*\.json$", name):
            continue
        d = json.load(open(os.path.join(dirpath, name), encoding="utf-8"))
        for b in d.get("baseline") or []:
            if b.get("content"):
                out.append({"text": b["content"], "model": b["model"], "questionId": d.get("questionId"),
                            "condition": b.get("variant", "plain"), "runIndex": d.get("runIndex", 0)})
    return out


def zscore(rows, markers="v1", *, fit_rows=None):
    """fit_rows の平均・母標準偏差で rows を変換（省略時は従来どおり全行）。"""
    features = features_for(markers)
    fit_rows = rows if fit_rows is None else fit_rows
    stats = {}
    for f in features:
        vals = [r[f] for r in fit_rows]
        mu, sd = statistics.mean(vals), statistics.pstdev(vals)
        stats[f] = (mu, sd if sd > 0 else 1.0)
    return [{f: (r[f] - stats[f][0]) / stats[f][1] for f in features} for r in rows]


def classify_loto_predictions(rows, markers="v1", *, train_only_scaling=False):
    """LOTO予測。train_only_scaling=Trueなら伏せたお題を標準化の当てはめからも除く。"""
    features = features_for(markers)
    z = None if train_only_scaling else zscore(rows, markers)
    qids = sorted({r["questionId"] for r in rows})
    predictions = [None] * len(rows)
    for q in qids:
        if train_only_scaling:
            z = zscore(rows, markers, fit_rows=[r for r in rows if r["questionId"] != q])
        train = [(r, zz) for r, zz in zip(rows, z) if r["questionId"] != q]
        test = [(i, zz) for i, (r, zz) in enumerate(zip(rows, z)) if r["questionId"] == q]
        cents = {}
        for m in sorted({r["model"] for r, _ in train}):
            vs = [zz for r, zz in train if r["model"] == m]
            cents[m] = {f: statistics.mean(v[f] for v in vs) for f in features}
        for i, zz in test:
            pred = min(cents, key=lambda m: math.dist([zz[f] for f in features], [cents[m][f] for f in features]))
            predictions[i] = pred
    return predictions


def classify_loto(rows, markers="v1"):
    """お題を 1 本ずつ抜き、残りの会社ごとの重心で抜いたお題の文章の会社を当てる。"""
    predictions = classify_loto_predictions(rows, markers)
    confusion = defaultdict(lambda: defaultdict(int))
    for r, pred in zip(rows, predictions):
        confusion[r["model"]][pred] += 1
    return sum(pred == r["model"] for r, pred in zip(rows, predictions)), len(rows), confusion


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("input", nargs="?")
    ap.add_argument("--markers", choices=("v1", "v1.1"), default="v1")
    ap.add_argument("--from-headless")
    ap.add_argument("--out", required=True)
    ap.add_argument("--classify", action="store_true")
    ap.add_argument("--normalize-html", action="store_true", help="text のない行の htmlPath を本文へ変換")
    ap.add_argument("--excerpt", type=int, metavar="N", help="冒頭 N 字超の最初の文末までを集計")
    a = ap.parse_args()
    features = features_for(a.markers)
    rhythm_descriptions = rhythm_descriptions_for(a.markers)
    if a.excerpt is not None and a.excerpt < 1:
        ap.error("--excerpt は 1 以上の整数を指定してください")
    if a.normalize_html:
        # -I ではスクリプトのディレクトリが sys.path に入らない。
        sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
        from html_to_text import html_to_markdownish, read_html, slice_by_heading
    items = load_headless(a.from_headless) if a.from_headless else load_jsonl(a.input)
    if not items:
        sys.exit("文章が 0 件")
    os.makedirs(a.out, exist_ok=True)
    rows = []
    for it in items:
        if "text" in it:
            content = it["text"]
        elif a.normalize_html and "htmlPath" in it:
            html_path = it["htmlPath"]
            if not os.path.isabs(html_path):
                html_path = os.path.join(os.path.dirname(os.path.abspath(a.input)), html_path)
            content = html_to_markdownish(read_html(html_path))
            if it.get("sectionStart"):
                content = slice_by_heading(content, it["sectionStart"], it.get("sectionEnd") or None)
            elif it.get("sectionEnd"):
                ap.error("sectionEnd には sectionStart が必要です")
        else:
            ap.error("各行に text が必要です。htmlPath を使う場合は --normalize-html を指定してください")
        r = {k: it.get(k) for k in ("model", "questionId", "condition", "runIndex")}
        if a.excerpt is not None:
            r["original_chars"] = len(strip_ws(content))
            content = excerpt_at_sentence_end(content, a.excerpt)
        r.update(count_text(content, a.markers))
        rows.append(r)
    with open(os.path.join(a.out, "per-text.csv"), "w", newline="", encoding="utf-8") as f:
        w = csv.DictWriter(f, fieldnames=list(rows[0].keys()))
        w.writeheader()
        w.writerows(rows)

    groups = defaultdict(list)
    for r in rows:
        groups[(r["model"], r["condition"])].append(r)
    keys = sorted(groups)
    md = ["| 癖 | 区分 | " + " | ".join(f"{m.split('/')[-1]}・{c}" for m, c in keys) + " |",
          "|---|---|" + "---:|" * len(keys)]
    desc = {m[0]: (m[1], m[2]) for m in markers_for(a.markers)}
    desc.update({f: ("修辞" if f.startswith("R") else "文", d)
                 for f, d in rhythm_descriptions.items()})
    for f in ["chars"] + features:
        cat, d = desc.get(f, ("文", f))
        md.append(f"| {f} {d if f in desc else ''} | {cat} | " +
                  " | ".join(f"{statistics.mean(r[f] for r in groups[k]):.2f}" for k in keys) + " |")
    md.append("")
    md.append(f"件数: " + "、".join(f"{m.split('/')[-1]}・{c} {len(groups[(m, c)])}" for m, c in keys))
    md.append("癖の値は 1,000 字あたりの数（S01〜S08 を除く）。chars は空白を除いた字数の平均。")
    md.extend(f"{f}: {d}" for f, d in rhythm_descriptions.items())
    if a.classify:
        hit, total, conf = classify_loto(rows, a.markers)
        n_models = len({r["model"] for r in rows})
        md += ["", f"## 癖の数だけで会社を当てる（お題を 1 本ずつ抜く交差検証・最近傍重心）", "",
               f"正解率 {hit}/{total} = {hit / total:.1%}（当て推量は {1 / n_models:.1%}）", ""]
        ms = sorted(conf)
        md.append("| 実際＼予測 | " + " | ".join(m.split("/")[-1] for m in ms) + " |")
        md.append("|---|" + "---:|" * len(ms))
        for m in ms:
            md.append(f"| {m.split('/')[-1]} | " + " | ".join(str(conf[m][p]) for p in ms) + " |")
    with open(os.path.join(a.out, "summary.md"), "w", encoding="utf-8") as f:
        f.write("\n".join(md) + "\n")
    print("\n".join(md))


if __name__ == "__main__":
    main()