"""第 2—4、22 课：材料选择、去重、分卷和关键词练习题。"""
import json
import re
from collections import Counter
from pathlib import Path

from opencc import OpenCC

from .common import digest, read_json, write_json, read_jsonl

SPECIAL = ["<pad>", "<bos>", "<eos>", "<五言>", "<七言>", "<关键词>", "<正文>", "<隔>", "<unk>"]
HAN = re.compile(r"[\u3400-\u9fff]+\Z")
CONVERTER = OpenCC("t2s")
WORDS = "春 夏 秋 冬 山 水 江 河 海 风 雨 雪 月 日 云 花 鸟 松 竹 梅 柳 莲 舟 酒 夜 客 归 梦 愁 春风 春雨 秋风 秋雨 明月 清风 白云 青山 流水 山水 江南 故乡 归乡 落花 桃花 梅花 荷花 杨柳 松竹 山川 日暮 黄昏 夕阳 斜阳 孤舟 渔舟 扁舟 寒山 白雪 秋月 春山 江水 青松 竹林 故人 归客 离别 相思 清明 中秋 田园 山寺 归舟 夜雨 远山 烟雨 天涯 浮云 夕照 落日 长江 黄河".split()
THEMES = {"思乡": "乡归客", "送别": "送别离", "山居": "山居林", "咏梅": "梅", "边塞": "塞戍边", "田园": "田农稼"}


def normalize(text):
    return CONVERTER.convert(text).strip()


def keywords_for(text):
    words = [word for word in WORDS if word in text]
    words += [theme for theme, clues in THEMES.items() if any(ch in text for ch in clues)]
    # 主题为规则推测，记录来源；不伪装成人工标注。
    return sorted(set(words), key=lambda word: (-len(word), word))


def clean_record(record):
    paragraphs = record.get("paragraphs")
    if not isinstance(paragraphs, list) or not all(isinstance(x, str) for x in paragraphs):
        return None, "正文结构异常"
    body = normalize("".join(paragraphs))
    body = re.sub(r"\s+", "", body)
    if not body:
        return None, "正文为空"
    lines = [x for x in re.split(r"[，。！？；,!?;]", body) if x]
    if len(lines) != 4:
        return None, "不是四句"
    if not all(HAN.fullmatch(x) for x in lines):
        return None, "缺字或夹注等非汉字"
    width = len(lines[0])
    if width not in (5, 7) or any(len(x) != width for x in lines):
        return None, "不是整齐五言七言"
    text = "".join(line + ("，" if i % 2 == 0 else "。") for i, line in enumerate(lines))
    plain = "".join(lines)
    return {"id": digest(plain), "title": normalize(record.get("title", "无题")),
            "author": normalize(record.get("author", "佚名")), "form": width,
            "text": text, "lines": lines, "keywords": keywords_for(plain)}, None


def group_variants(records):
    """共享两句或以上的版本放同组，避免改一字的同诗跨卷；整首精确重复已去除。"""
    parent = list(range(len(records)))
    sizes = [1] * len(records)
    signatures = {}

    def find(x):
        while x != parent[x]:
            parent[x] = parent[parent[x]]
            x = parent[x]
        return x

    for i, record in enumerate(records):
        lines = record["lines"]
        for a in range(4):
            for b in range(a + 1, 4):
                key = (a, lines[a], b, lines[b])
                if key in signatures:
                    x, y = find(i), find(signatures[key])
                    if x != y:
                        if sizes[x] < sizes[y]: x, y = y, x
                        parent[y] = x
                        sizes[x] += sizes[y]
                else:
                    signatures[key] = i
    ids = {}
    for i, record in enumerate(records):
        root = find(i)
        ids[root] = min(ids.get(root, record["id"]), record["id"])
    return [ids[find(i)] for i in range(len(records))]


def prepare(source, output, min_count=2):
    source, output = Path(source), Path(output)
    if min_count < 1: raise ValueError("min-count 至少为 1。")
    if (output / "manifest.json").exists():
        raise ValueError("目标材料目录已存在，请使用新的 --output，避免覆盖已训练模型对应的材料。")
    selected = []
    seen = set()
    reasons, raw_counts = Counter(), Counter()
    sources = []
    for dynasty in ("tang", "song"):
        paths = sorted((source / "全唐诗").glob(f"poet.{dynasty}.*.json"))
        if not paths: raise ValueError(f"在 {source}/全唐诗 找不到 {dynasty} 正文文件。")
        for path in paths:
            raw = path.read_bytes()
            import hashlib
            sources.append({"file": str(path.relative_to(source)), "sha256": hashlib.sha256(raw).hexdigest()})
            for record in json.loads(raw):
                raw_counts[dynasty] += 1
                cleaned, reason = clean_record(record)
                if reason:
                    reasons[reason] += 1
                    continue
                if cleaned["id"] in seen:
                    reasons["繁简统一后正文重复"] += 1
                    continue
                seen.add(cleaned["id"])
                cleaned.update(source=str(path.relative_to(source)), dynasty=dynasty)
                selected.append(cleaned)
    selected.sort(key=lambda x: x["id"])
    groups = group_variants(selected)
    splits = {"train": [], "val": [], "test": []}
    for record, group in zip(selected, groups):
        bucket = int(digest("poetry-course-v1:" + group)[:8], 16) % 100
        split = "train" if bucket < 90 else "val" if bucket < 95 else "test"
        record["group"] = group
        splits[split].append(record)
    counts = Counter("".join(x["text"] for x in splits["train"]))
    alphabet = {char for char, count in counts.items() if count >= min_count}
    # 标签本身是公开格式定义，不读取验证集/测试集建立词表。
    alphabet.update("".join(WORDS) + "".join(THEMES))
    vocab = SPECIAL + sorted(alphabet)
    for name in splits:
        before = len(splits[name])
        splits[name] = [r for r in splits[name] if set(r["text"]) <= alphabet]
        reasons[f"{name} 含低频未收录字"] = before - len(splits[name])
    output.mkdir(parents=True, exist_ok=True)
    for name, records in splits.items():
        (output / f"{name}.jsonl").write_text("".join(json.dumps(r, ensure_ascii=False) + "\n" for r in records), encoding="utf-8")
    write_json(output / "vocab.json", vocab)
    counts_by_split = {name: dict(Counter(str(r["form"]) for r in rs)) for name, rs in splits.items()}
    manifest = {"version": 1, "source_counts": dict(raw_counts), "source_files": sources,
                "rejections": dict(reasons), "split_counts": {k: len(v) for k,v in splits.items()},
                "forms": counts_by_split, "vocab_size": len(vocab), "min_count": min_count,
                "grouping": "繁简统一后正文去重；任意相同位置两句相同的诗归组，再按组固定分卷",
                "keyword_method": "正文中出现的词 + 可审计的主题字表；自动弱标签，非人工标注",
                "keyword_themes": THEMES,
                "split_sha256": {name: digest((output / f"{name}.jsonl").read_text()) for name in splits},
                "vocab_sha256": digest(json.dumps(vocab, ensure_ascii=False)),
                "note": "筛选四句五七言外形，不宣称符合平水韵和平仄；繁简自动转换可能改变个别古字含义。"}
    write_json(output / "manifest.json", manifest)
    return manifest


class Tokenizer:
    def __init__(self, vocabulary):
        self.tokens = vocabulary
        self.ids = {token: i for i, token in enumerate(vocabulary)}

    def encode(self, text):
        missing = sorted(set(text) - self.ids.keys())
        if missing: raise ValueError("字表中没有这些字：" + "、".join(missing))
        return [self.ids[x] for x in text]

    def decode(self, ids):
        return "".join(self.tokens[int(i)] for i in ids)

    def prefix(self, form, keywords=()):
        out = [self.ids["<bos>"], self.ids["<五言>" if form == 5 else "<七言>"], self.ids["<关键词>"]]
        for i, word in enumerate(keywords):
            if i: out.append(self.ids["<隔>"])
            out.extend(self.encode(word))
        return out + [self.ids["<正文>"]]


class Batcher:
    def __init__(self, records, tokenizer, context, seed):
        import random
        self.records, self.tokenizer, self.context = records, tokenizer, context
        self.random = random.Random(seed)
        if not records: raise ValueError("材料为空。")

    def batch(self, batch_size, device):
        import torch
        x = torch.zeros((batch_size, self.context), dtype=torch.long)
        y = torch.full_like(x, -100)
        for i in range(batch_size):
            record = self.random.choice(self.records)
            words = record["keywords"]
            k = self.random.choices([0, 1, 2], weights=[3, 4, 3])[0]
            keywords = self.random.sample(words, min(k, len(words)))
            prefix = self.tokenizer.prefix(record["form"], keywords)
            sequence = prefix + self.tokenizer.encode(record["text"]) + [self.tokenizer.ids["<eos>"]]
            if len(sequence) > self.context + 1: raise ValueError("练习题超过模型前文长度。")
            n = len(sequence) - 1
            x[i, :n] = torch.tensor(sequence[:-1])
            y[i, len(prefix)-1:n] = torch.tensor(sequence[len(prefix):])
        return x.to(device), y.to(device)