"""各个实验脚本共用的部分：语料、批次、训练循环、计算量统计、结果存档。"""
import datetime
import glob
import os
import sys
import sysconfig
import time

import numpy as np

HERE = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, os.path.join(HERE, ".."))
import zerocg as zc  # noqa: E402
from models import ZeroX  # noqa: E402
from zerocg import functional as F  # noqa: E402

ITERS = int(os.environ.get("ITERS", 2000))
T, B = 64, 32
CFG = dict(n_embd=64, n_head=4, n_layer=2)


def load_corpus(path=None):
    """默认语料：Python 标准库源码（纯 ASCII 的 .py 文件，约 150 万字符）。"""
    if path:
        return open(path, encoding="utf-8").read()
    parts, total = [], 0
    for p in sorted(glob.glob(os.path.join(sysconfig.get_paths()["stdlib"], "*.py"))):
        try:
            s = open(p, encoding="utf-8").read()
        except (UnicodeDecodeError, OSError):
            continue
        if s.isascii():
            parts.append(s)
            total += len(s)
        if total > 1_500_000:
            break
    return "\n".join(parts)


def prepare_data(path=None):
    text = load_corpus(path)
    chars = sorted(set(text))
    stoi = {c: i for i, c in enumerate(chars)}
    data = np.array([stoi[c] for c in text], dtype=np.int64)
    split = int(0.9 * len(data))
    return text, len(chars), data[:split], data[split:]


def batches(data, n, length, seed, batch=B):
    r = np.random.default_rng(seed)
    for _ in range(n):
        ix = r.integers(0, len(data) - length - 1, batch)
        yield np.stack([data[i:i + length] for i in ix]), np.stack([data[i + 1:i + length + 1] for i in ix])


def loss_of(logits, y):
    return F.cross_entropy(logits.reshape(-1, logits.shape[-1]), y.reshape(-1)).item()


def train(name, model, train_data, val_data, seed, lr=3e-3, iters=None, seq_len=None, batch=None):
    """返回验证 loss（固定在 64 字上评估，便于各实验比较）。
    学习率：100 步线性预热 + 余弦退火到 10%；AdamW，权重衰减 0.01，梯度裁剪 1.0。
    seq_len / batch 默认 T / B；改训练长度时请同时改 batch，保持每步的字数不变。"""
    iters = iters or ITERS
    opt = zc.optim.AdamW(model.parameters(), lr=lr, weight_decay=0.01)
    t0 = time.time()
    for it, (x, y) in enumerate(batches(train_data, iters, seq_len or T, seed, batch or B)):
        opt.lr = lr * min(1, (it + 1) / 100) * (0.1 + 0.9 * 0.5 * (1 + np.cos(np.pi * it / iters)))
        logits = model(x)
        loss = F.cross_entropy(logits.reshape(-1, logits.shape[-1]), y.reshape(-1))
        opt.zero_grad()
        loss.backward()
        zc.optim.clip_grad_norm(model.parameters(), 1.0)
        opt.step()
        if it % 500 == 0:
            print(f"  [{name} seed={seed} lr={lr:g}] step {it:5d}  train {loss.item():.3f}  "
                  f"{time.time() - t0:.0f}s", flush=True)
    return evaluate(model, val_data)


def evaluate(model, val_data, n=40):
    model.eval()
    with zc.no_grad():
        val = float(np.mean([loss_of(model(x), y) for x, y in batches(val_data, n, T, 999)]))
    model.train()
    return val


def forward_flops_per_char(model, V):
    """前向一次、平均每个字多少次浮点运算（只算矩阵乘）。训练一步约是它的 3 倍。"""
    x = np.zeros((1, T), dtype=np.int64)
    with zc.no_grad(), zc.count_flops() as c:
        model(x)
    return c.flops / T


def slot_usage(model):
    """每层"有效槽数"：最近一次 forward 的平均读/写地址分布的 exp(熵)，对头取平均。
    16 个槽都用得一样多 = 16；全挤在一个槽 = 1。"""
    out = []
    for b in model.blocks:
        read, write = b.mix.usage
        eff = lambda p: float(np.mean(np.exp(-(p * np.log(p + 1e-12)).sum(-1))))  # noqa: E731
        out.append((eff(read), eff(write)))
    return out


def save_result(name, lines):
    """结果写两份：experiments/results/<时间>-<name>.md（存档，不覆盖）和 experiments/<name>.md（最新）。"""
    text = "\n".join(lines) + "\n"
    stamp = datetime.datetime.now().strftime("%Y%m%d-%H%M")
    os.makedirs(os.path.join(HERE, "results"), exist_ok=True)
    for path in (os.path.join(HERE, "results", f"{stamp}-{name}.md"), os.path.join(HERE, f"{name}.md")):
        open(path, "w").write(text)
    print("\n" + text)


def is_zeropoint(model):
    return isinstance(model, ZeroX)