"""第 26 章：把短句语言模型微调为颜色/物品提取助手，并检查旧任务退化。"""

import argparse
from copy import deepcopy
from hashlib import sha256
import json
from pathlib import Path
import torch
from torch import nn
from tiny_transformer import TinyTransformer
from small_language_model import corpus, make_windows, evaluate

PAD, BOS, EOS, USER, ASSISTANT = range(5)
IGNORE = -100
INSTRUCTIONS = {"颜色": "只回答颜色。", "物品": "只回答物品。"}
SPECIAL = ["<PAD>", "<BOS>", "<EOS>", "<USER>", "<ASSISTANT>"]


def setup():
    torch.set_num_threads(1)
    torch.use_deterministic_algorithms(True)


def vocabulary(texts):
    chars = sorted(set("".join(texts)+"".join(INSTRUCTIONS.values())))
    return SPECIAL+chars


def instruction_records(texts):
    records = []
    for text in texts:
        for task in INSTRUCTIONS:
            answer = text[4] if task == "颜色" else text[6:8]
            records.append(dict(prompt=INSTRUCTIONS[task]+text[:14], answer=answer,
                                task=task, source=text))
    return records


def encode_prefix(prompt, vocab):
    ids = {token: i for i, token in enumerate(vocab)}
    return [BOS, USER]+[ids[c] for c in prompt]+[ASSISTANT]


def supervised_data(records, vocab, assistant_only=True):
    ids = {token: i for i, token in enumerate(vocab)}
    inputs, targets = [], []
    for row in records:
        prefix = encode_prefix(row["prompt"], vocab)
        sequence = prefix+[ids[c] for c in row["answer"]]+[EOS]
        x, y = sequence[:-1], sequence[1:]
        if assistant_only:
            y[:len(prefix)-1] = [IGNORE]*(len(prefix)-1)
        if len(x) > 32:
            raise ValueError("示例不截断对话；输入超出 32 个位置")
        inputs.append(x+[PAD]*(32-len(x)))
        targets.append(y+[IGNORE]*(32-len(y)))
    return torch.tensor(inputs), torch.tensor(targets)


def loss_value(model, data, ignore=IGNORE):
    x, y = data[:2]
    return nn.functional.cross_entropy(model(x).flatten(0, 1), y.flatten(), ignore_index=ignore)


def persistent_bytes(model, optimizer):
    # 当前实际存在的张量；不包括激活、临时工作区、Python 对象等。
    weights = sum(p.numel()*p.element_size() for p in model.parameters())
    gradients = sum(p.grad.numel()*p.grad.element_size() for p in model.parameters() if p.grad is not None)
    states = sum(t.numel()*t.element_size() for state in optimizer.state.values()
                 for t in state.values() if torch.is_tensor(t))
    return dict(weights=weights, gradients=gradients, optimizer=states)


def pretrain(train_texts, vocab, steps=600):
    torch.manual_seed(2026)
    model = TinyTransformer(len(vocab), width=32).double()
    optimizer = torch.optim.AdamW(model.parameters(), lr=.003, weight_decay=.01)
    sampler = torch.Generator().manual_seed(2027)
    x, y, _ = make_windows(train_texts, vocab, 32)
    for _ in range(steps):
        selected = torch.randint(len(x), (16,), generator=sampler)
        optimizer.zero_grad(set_to_none=True)
        loss = loss_value(model, (x[selected], y[selected]), ignore=PAD)
        loss.backward()
        nn.utils.clip_grad_norm_(model.parameters(), 1, error_if_nonfinite=True)
        optimizer.step()
    model.zero_grad(set_to_none=True)
    return model


def finetune(model, records, old_texts, vocab, steps=160, replay=0, learning_rate=.001):
    model = deepcopy(model)
    model.zero_grad(set_to_none=True)
    trainable = [p for p in model.parameters() if p.requires_grad]
    optimizer = torch.optim.AdamW(trainable, lr=learning_rate, weight_decay=.01)
    sampler = torch.Generator().manual_seed(2028)
    old_sampler = torch.Generator().manual_seed(2029)
    x, y = supervised_data(records, vocab)
    old_x, old_y, _ = make_windows(old_texts, vocab, 32)
    targets, old_targets, history = 0, 0, []
    for step in range(1, steps+1):
        model.train()
        selected = torch.randint(len(x), (16,), generator=sampler)
        optimizer.zero_grad(set_to_none=True)
        instruction_loss = loss_value(model, (x[selected], y[selected]))
        loss = (1-replay)*instruction_loss
        targets += (y[selected] != IGNORE).sum().item()
        if replay:
            old_selected = torch.randint(len(old_x), (16,), generator=old_sampler)
            old_loss = loss_value(model, (old_x[old_selected], old_y[old_selected]), ignore=PAD)
            loss = loss+replay*old_loss
            old_targets += (old_y[old_selected] != PAD).sum().item()
        if not torch.isfinite(loss):
            raise RuntimeError("非有限训练损失")
        loss.backward()
        nn.utils.clip_grad_norm_(trainable, 1, error_if_nonfinite=True)
        optimizer.step()
        if step % 40 == 0:
            history.append(dict(step=step, loss=loss.item()))
    stats = dict(trainable=sum(p.numel() for p in trainable), steps=steps,
                 targets=targets, old_targets=old_targets, state_bytes=persistent_bytes(model, optimizer),
                 history=history)
    model.zero_grad(set_to_none=True)
    return model, stats


@torch.no_grad()
def generate_from_prefixes(model, prefixes, vocab, max_new):
    model.eval()
    if len({len(p) for p in prefixes}) != 1:
        raise ValueError("本例批量生成要求等长提示")
    if len(prefixes[0])+max_new-1 > model.config["max_context"]:
        raise ValueError("提示与生成上限超出上下文")
    ids = torch.tensor(prefixes)
    answers = [[] for _ in prefixes]
    finished = [False]*len(prefixes)
    for _ in range(max_new):
        logits = model(ids)[:, -1].clone()
        logits[:, [PAD, BOS, USER, ASSISTANT]] = -torch.inf
        chosen = logits.argmax(-1)
        for i, token in enumerate(chosen.tolist()):
            if not finished[i]:
                if token == EOS:
                    finished[i] = True
                else:
                    answers[i].append(vocab[token])
        ids = torch.cat([ids, chosen[:, None]], dim=1)
        if all(finished):
            break
    return [("".join(answer), stop) for answer, stop in zip(answers, finished)]


def generate_answers(model, records, vocab, max_new=6):
    return generate_from_prefixes(model, [encode_prefix(row["prompt"], vocab) for row in records],
                                  vocab, max_new)


@torch.no_grad()
def metrics(model, records, old_texts, vocab):
    model.eval()
    answer_loss = loss_value(model, supervised_data(records, vocab)).item()
    generated = generate_answers(model, records, vocab)
    correct, formatted = 0, 0
    failures = []
    for row, (answer, stopped) in zip(records, generated):
        valid = answer in (["红", "蓝", "绿", "白"] if row["task"] == "颜色"
                           else ["杯子", "盒子", "本子", "袋子"])
        formatted += valid and stopped
        correct += answer == row["answer"] and stopped
        if answer != row["answer"] or not stopped:
            failures.append(dict(prompt=row["prompt"], expected=row["answer"], generated=answer, stopped=stopped))
    old_loss, old_accuracy = evaluate(model, make_windows(old_texts, vocab, 32))
    ids = {token: i for i, token in enumerate(vocab)}
    old_prefixes = [[BOS]+[ids[c] for c in text[:14]] for text in old_texts]
    old_generated = generate_from_prefixes(model, old_prefixes, vocab, 16)
    old_exact = sum(answer == text[14:] and stopped for text, (answer, stopped) in zip(old_texts, old_generated))
    return dict(answer_loss=answer_loss, correct=correct, formatted=formatted, count=len(records),
                old_loss=old_loss, old_accuracy=old_accuracy, old_exact=old_exact, failures=failures,
                old_example=dict(expected=old_texts[0][14:], generated=old_generated[0][0],
                                 stopped=old_generated[0][1]))


def prepare():
    setup()
    train, validation, test = corpus()
    vocab = vocabulary(train)
    model = pretrain(train, vocab)
    return model, (train, validation, test), vocab


def run():
    base, (train, validation, test), vocab = prepare()
    records = instruction_records(train[:96])
    sft, sft_cost = finetune(base, records, train, vocab)
    replay, replay_cost = finetune(base, records, train, vocab, replay=.3)
    results = []
    for name, model, cost in [("基础模型", base, None), ("仅指令微调", sft, sft_cost),
                              ("混合旧文本", replay, replay_cost)]:
        results.append(dict(name=name, validation=metrics(model, instruction_records(validation), validation, vocab),
                            test=metrics(model, instruction_records(test), test, vocab), cost=cost))
    report = dict(vocab_size=len(vocab), parameters=sum(p.numel() for p in base.parameters()),
                  base_steps=600, instruction_records=len(records), results=results)
    return report, base, vocab


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--out", type=Path)
    parser.add_argument("--save-base", type=Path)
    args = parser.parse_args()
    report, base, vocab = run()
    print(f"词表={report['vocab_size']}；参数={report['parameters']:,}；"
          f"基础预训练=600 步；指令训练记录={report['instruction_records']}")
    for row in report["results"]:
        val, test = row["validation"], row["test"]
        print(f"{row['name']}：验证回答损失={val['answer_loss']:.3f}；"
              f"测试回答={test['correct']}/{test['count']}，格式={test['formatted']}/{test['count']}；"
              f"旧文本损失={test['old_loss']:.3f}，旧续写={test['old_exact']}/32")
        if row["cost"]:
            cost = row["cost"]
            print(f"  更新={cost['steps']}；回答目标={cost['targets']}；旧文本目标={cost['old_targets']}；"
                  f"状态张量={sum(cost['state_bytes'].values())/1024:.1f} KiB")
        if test["failures"]:
            failure = test["failures"][0]
            print(f"  首个失败：应答={failure['expected']}；实际={failure['generated']}；"
                  f"正常终止={failure['stopped']}")
        print(f"  旧任务续写示例：{test['old_example']['generated']}")
    if args.out:
        args.out.parent.mkdir(parents=True, exist_ok=True)
        args.out.write_text(json.dumps(report, ensure_ascii=False, indent=2))
        print(f"实验记录：{args.out}")
    if args.save_base:
        args.save_base.parent.mkdir(parents=True, exist_ok=True)
        train, _, _ = corpus()
        torch.save(dict(state=base.state_dict(), config=base.config, vocab=vocab,
                        experiment="instruction-base-v1", torch_version=str(torch.__version__),
                        train_sha256=sha256("\n".join(train).encode()).hexdigest()), args.save_base)
        print(f"基础模型：{args.save_base}")


if __name__ == "__main__":
    main()