#!/usr/bin/env python3
"""
JspaceAI —— 对话版（只有语言，控制台交互）

基于字符级语言模型 + 自主进化（EWC + 经验回放 + 专家可塑性）。
交互式对话：用户输入文本，模型生成回复，同时持续学习。

模式：
    --mode chat:     交互对话（默认）
    --mode train:    先在 Shakespeare 语料上预训练若干步，再进入对话
    --mode generate: 给定提示词一次性生成

运行：
    python main_chat.py
    python main_chat.py --mode train --steps 100
    python main_chat.py --mode generate --prompt "To be"
"""
from __future__ import annotations
import argparse
import torch
import numpy as np
from pathlib import Path

from jspaceai import (
    LanguageConfig, JSpaceLanguageModel, EvolutionTrainer,
    CharTokenizer, load_chinese_corpus,
)
from train_chat import clean_corpus


def get_config(vocab_size: int) -> LanguageConfig:
    return LanguageConfig(
        vocab_size=vocab_size,
        embed_dim=64, input_dim=32,
        workspace_dim=128, expert_dim=64,
        num_experts=12, num_wells=8,
        ode_steps=4, dt=0.1, tau_w=0.6,
        jacobian_sparsity=32, noise_std=0.001,
        use_rk4=True, use_layer_norm=True,
    )


def load_or_init_model(device: str):
    """加载已保存的模型或初始化新模型"""
    text = clean_corpus()  # 清洗后语料（繁简统一+过滤）

    mp = Path('outputs/chat_model.pt')
    if mp.exists():
        try:
            ckpt = torch.load(mp, map_location=device, weights_only=False)
            # 用保存的 tokenizer chars 确保一致
            saved_chars = ckpt.get('tokenizer_chars', None)
            if saved_chars:
                tokenizer = CharTokenizer(
                    chars=saved_chars,
                    char_to_idx={c: i for i, c in enumerate(saved_chars)},
                    idx_to_char={i: c for i, c in enumerate(saved_chars)},
                )
            else:
                tokenizer = CharTokenizer.from_text(text)
            config = get_config(tokenizer.vocab_size)
            model = JSpaceLanguageModel(config).to(device)
            model.load_state_dict(ckpt['model'])
            print(f"已加载模型: {mp}（vocab={tokenizer.vocab_size}）")
            return model, config, tokenizer, text
        except Exception as e:
            print(f"模型加载失败: {e}，全新初始化")

    tokenizer = CharTokenizer.from_text(text)
    config = get_config(tokenizer.vocab_size)
    model = JSpaceLanguageModel(config).to(device)
    print("全新初始化（首次对话）")
    return model, config, tokenizer, text


def save_model(model, config, tokenizer):
    Path('outputs').mkdir(exist_ok=True)
    torch.save({
        'model': model.state_dict(),
        'config': config,
        'tokenizer_chars': tokenizer.chars,
    }, 'outputs/chat_model.pt')


def train(model, tokenizer, text, n_steps: int, device: str):
    """在 Shakespeare 语料上预训练

    训练时临时关闭 RK4 用 Euler 加速（快 4 倍），训练完恢复 RK4。
    """
    print("\n" + "=" * 60)
    print(f"预训练 {n_steps} 步（Shakespeare 语料）")
    print("=" * 60)

    config = model.config

    # 训练时临时关 RK4 加速（Euler 快 4 倍）
    original_rk4 = config.use_rk4
    for expert in model.experts:
        expert.use_rk4 = False
    print(f"训练模式: Euler（加速），训练后恢复 RK4")

    trainer = EvolutionTrainer(
        model, config, lr=5e-3, ewc_lambda=0.05, device=device,
    )
    chunks = [text[i:i+200] for i in range(0, len(text), 200)]
    trainer.evolve(
        chunks, tokenizer,
        seq_len=64, batch_size=8,
        consolidate_every=50, generate_every=50,
        max_steps=n_steps, prompt_text="学而时习之",
    )

    # 恢复 RK4
    for expert in model.experts:
        expert.use_rk4 = original_rk4

    save_model(model, config, tokenizer)
    print(f"\n模型已保存: outputs/chat_model.pt")


def generate_response(model, tokenizer, prompt: str, n_new: int = 100,
                      temperature: float = 0.8, top_k: int = 5) -> str:
    """生成回复"""
    # 把用户输入编码（未知字符用 0）
    prompt_ids = tokenizer.encode(prompt)
    if not prompt_ids:
        prompt_ids = [0]
    generated = model.generate(
        prompt_ids, n_new=n_new, temperature=temperature, top_k=top_k,
    )
    return tokenizer.decode(generated)


def online_learn(model, config, tokenizer, text: str, user_input: str, device: str):
    """在线学习用户输入 + 回放一段 Shakespeare 防遗忘"""
    import torch.nn.functional as F
    # 把用户输入作为新语料学习
    user_tokens = tokenizer.encode(user_input)
    if len(user_tokens) < 4:
        return  # 太短不学

    # 重复用户输入凑够 seq_len
    seq_len = 48
    if len(user_tokens) < seq_len:
        user_tokens = user_tokens * (seq_len // len(user_tokens) + 1)
    user_seq = user_tokens[:seq_len]
    token_tensor = torch.tensor([user_seq], dtype=torch.long).to(device)

    # forward + next-token loss
    model.train()
    model.zero_grad()
    logits, _ = model(token_tensor)
    pred = logits[:, :-1]
    target = token_tensor[:, 1:]
    loss = F.cross_entropy(
        pred.reshape(-1, config.vocab_size),
        target.reshape(-1),
    )
    loss.backward()
    torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)

    # 手动 SGD step（EvolutionTrainer 内部有 EWC，这里简化用直接 step）
    with torch.no_grad():
        for p in model.parameters():
            if p.grad is not None:
                p -= 5e-3 * p.grad
    model.eval()
    return loss.item()


def chat(model, config, tokenizer, text, device: str):
    """交互对话循环"""
    print("\n" + "=" * 60)
    print("JspaceAI 对话模式")
    print("=" * 60)
    print("输入文本与模型对话，模型会持续学习你的输入。")
    print("命令:  /quit 退出  /save 保存  /reset 重置  /train N 预训练N步")
    print("=" * 60 + "\n")

    model.eval()
    while True:
        try:
            user_input = input("你: ").strip()
        except (EOFError, KeyboardInterrupt):
            print("\n再见")
            break

        if not user_input:
            continue

        if user_input.startswith('/'):
            cmd = user_input.lower()
            if cmd in ('/quit', '/exit', '/q'):
                save_model(model, config, tokenizer)
                print(f"(状态已保存到 outputs/chat_model.pt)")
                print("再见")
                break
            elif cmd == '/save':
                save_model(model, config, tokenizer)
                print(f"(已保存到 outputs/chat_model.pt)")
                continue
            elif cmd == '/reset':
                model = JSpaceLanguageModel(config).to(device)
                print("(模型已重置为随机初始化)")
                continue
            elif cmd.startswith('/train'):
                parts = cmd.split()
                n = int(parts[1]) if len(parts) > 1 else 50
                train(model, tokenizer, text, n, device)
                model.eval()
                continue
            else:
                print("未知命令。可用: /quit /save /reset /train N")
                continue

        # 在线学习用户输入
        loss = online_learn(model, config, tokenizer, text, user_input, device)

        # 生成回复
        response = generate_response(
            model, tokenizer, user_input,
            n_new=80, temperature=0.8, top_k=5,
        )
        print(f"AI: {response}")
        if loss is not None:
            print(f"   (学习 loss={loss:.3f})")


def generate_once(model, tokenizer, prompt: str, n_new: int = 200):
    """一次性生成"""
    model.eval()
    response = generate_response(model, tokenizer, prompt, n_new=n_new,
                                  temperature=0.7, top_k=5)
    print(f"提示: {prompt}")
    print(f"生成: {response}")


def main():
    p = argparse.ArgumentParser(description='JspaceAI 对话版')
    p.add_argument('--mode', default='chat',
                   choices=['chat', 'train', 'generate'],
                   help='运行模式: chat=交互, train=预训练, generate=一次性生成')
    p.add_argument('--steps', type=int, default=600, help='train 模式步数')
    p.add_argument('--prompt', default='To be', help='generate 模式提示词')
    p.add_argument('--device', default='cpu', help='设备 (cpu/cuda/mps/auto)')
    args = p.parse_args()

    dev = args.device
    if dev == 'auto':
        dev = 'cuda' if torch.cuda.is_available() else (
            'mps' if torch.backends.mps.is_available() else 'cpu')

    model, config, tokenizer, text = load_or_init_model(dev)

    if args.mode == 'chat':
        chat(model, config, tokenizer, text, dev)
    elif args.mode == 'train':
        train(model, tokenizer, text, args.steps, dev)
    elif args.mode == 'generate':
        generate_once(model, tokenizer, args.prompt)


if __name__ == '__main__':
    main()
