#!/usr/bin/env python3
"""
持续训练: 20k局/轮, 无限循环
每轮: dotnet自对弈 → V训练(Neural-only+replay) → eval → ONNX导出 → 存档
改进: dropout正则化, 只训Neural数据, 经验回放, 每轮eval验证回滚
"""
import subprocess, os, sys, time, shutil, glob, pickle, re
from datetime import datetime

HJHA_DIR = '/home/xiaoou/projects/hjha-server'
TRAINER_DIR = '/home/xiaoou/projects/paodekuai-trainer'
LOG = os.path.join(HJHA_DIR, 'training_loop.log')
REPLAY_PATH = os.path.join(TRAINER_DIR, 'data/replay_buffer.pkl')
REPLAY_MAX_ROUNDS = 3  # keep last N rounds in buffer
EVAL_GAMES = 200  # eval games per round

def now(): return datetime.now().strftime('%m%d %H:%M:%S')
def log(msg):
    line = f"[{now()}] {msg}"
    print(line, flush=True)
    with open(LOG, 'a') as f: f.write(line + '\n')

def run(cmd, cwd=HJHA_DIR, to=86400):
    try:
        r = subprocess.run(cmd, cwd=cwd, capture_output=True, text=True, timeout=to)
        ok = r.returncode == 0
        if not ok:
            log(f"  FAIL rc={r.returncode}: {(r.stderr or r.stdout)[:200]}")
        return ok, r.stdout + r.stderr
    except Exception as e:
        log(f"  EXCEPTION: {e}")
        return False, str(e)

def selfplay(games):
    log(f"── 自对弈 {games}局 ──")
    td = os.path.join(HJHA_DIR, 'training_data')
    # 先备份旧数据, 崩了也不丢
    if os.path.exists(td):
        backup = f"{td}_prev"
        if os.path.exists(backup): shutil.rmtree(backup)
        shutil.move(td, backup)
    os.makedirs(td)
    
    env = os.environ.copy()
    env['PATH'] = f"/usr/lib/dotnet:{env.get('PATH','')}"
    
    ok, out = run(['/usr/lib/dotnet/dotnet', 'run', '--project', 'hjha-console', '-c', 'Release',
                   '--', '--mix-rotate', str(games)], to=max(36000, games*2))
    
    csv = len(glob.glob(os.path.join(td, 'game_*.csv')))
    sf = os.path.join(HJHA_DIR, 'stress_summary.txt')
    if os.path.exists(sf):
        with open(sf) as f: log(f"  {f.read().strip()}")
    log(f"  {csv} CSV")
    return csv

def load_csv_states(csv_dir):
    """Load states from CSV.
    New format (q_value): prefer ISMCTS states (player!=2) with ISMCTS evaluation as target.
    Old format (reward): fall back to Neural-only with game result."""
    import numpy as np
    states, results = [], []
    filtered = 0
    ismcts_samples = 0
    for fp in sorted(glob.glob(f'{csv_dir}/game_*.csv')):
        try:
            with open(fp) as f:
                header = f.readline().strip()
                has_qvalue = 'q_value' in header
                has_player = header.endswith(',player')
                for line in f:
                    if line.startswith('#'): continue
                    p = line.strip().split('"')
                    if len(p) < 3: continue
                    sf = [float(x) for x in p[1].split(',')]
                    rest = p[2].strip(',').split(',')
                    if len(sf) != 364: continue
                    if len(rest) < 2: continue
                    label = float(rest[1])  # q_value or reward
                    if has_qvalue and has_player and len(rest) >= 3:
                        player = int(rest[2])
                        if player == 2 and abs(label) < 0.01:
                            # Neural state with no q_value (0.0) → skip
                            filtered += 1
                            continue
                        if player != 2:
                            ismcts_samples += 1
                    elif not has_qvalue and has_player and len(rest) >= 3:
                        # Old format: Neural-only filter
                        player = int(rest[2])
                        if player != 2:
                            filtered += 1
                            continue
                    states.append(sf)
                    results.append(label)
        except: pass
    if ismcts_samples:
        log(f"  {ismcts_samples} ISMCTS q-value samples")
    if filtered:
        log(f"  filtered {filtered} unusable states")
    return np.array(states, dtype=np.float32) if states else None, np.array(results, dtype=np.float32) if results else None

def train_v(prev_winrate=None):
    log(f"── V训练 ──")
    
    sys.path.insert(0, TRAINER_DIR)
    import numpy as np, torch
    from models.network import PdkNet
    
    csv_dir = os.path.join(HJHA_DIR, 'training_data')
    cur_states, cur_results = load_csv_states(csv_dir)
    
    if cur_states is None or len(cur_states) == 0:
        log("  0 samples!")
        return False, None
    
    cur_games = len(glob.glob(f'{csv_dir}/game_*.csv'))
    log(f"  current: {len(cur_states)} Neural samples from {cur_games} games")
    
    # --- experience replay: load past rounds ---
    replay_loaded = 0
    if os.path.exists(REPLAY_PATH):
        try:
            with open(REPLAY_PATH, 'rb') as f:
                replay_data = pickle.load(f)
            # replay_data is list of (states, results) tuples
            all_r = []
            for rs, rr in replay_data[-REPLAY_MAX_ROUNDS:]:
                if len(rs) > 0:
                    all_r.append((rs, rr))
                    replay_loaded += len(rs)
            if all_r:
                replay_states = np.concatenate([r[0] for r in all_r])
                replay_results = np.concatenate([r[1] for r in all_r])
                states = np.concatenate([cur_states, replay_states])
                results = np.concatenate([cur_results, replay_results])
                log(f"  replay: {replay_loaded} samples from {min(len(replay_data), REPLAY_MAX_ROUNDS)} past rounds → total {len(states)}")
            else:
                states, results = cur_states, cur_results
        except Exception as e:
            log(f"  replay load failed: {e}, using current only")
            states, results = cur_states, cur_results
    else:
        states, results = cur_states, cur_results
    
    model = PdkNet()
    opt = torch.optim.Adam(model.parameters(), lr=0.001, weight_decay=1e-5)
    scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(opt, T_max=50, eta_min=1e-5)
    loss_fn = torch.nn.MSELoss()
    n = len(states)
    t0 = datetime.now()
    best_loss = float('inf')
    no_improve = 0
    
    for ep in range(50):
        perm = np.random.permutation(n)
        losses = []
        for i in range(0, n, 128):
            idx = perm[i:i+128]
            x = torch.from_numpy(states[idx]).float()
            y = torch.from_numpy(results[idx]).float().unsqueeze(1)
            model.train(); pred = model(x)
            loss = loss_fn(pred, y)
            opt.zero_grad(); loss.backward(); opt.step()
            losses.append(loss.item())
        scheduler.step()
        avg_loss = np.mean(losses)
        if avg_loss < best_loss * 0.999:
            best_loss = avg_loss
            no_improve = 0
        else:
            no_improve += 1
        if ep % 5 == 0 or ep == 49:
            dt = (datetime.now() - t0).total_seconds()
            lr = scheduler.get_last_lr()[0]
            log(f"  ep{ep+1}/50 loss={avg_loss:.4f} lr={lr:.6f} ({dt:.0f}s)")
        if no_improve >= 10:
            log(f"  early stop ep{ep+1}, plateau 10 epochs")
            break
    
    elapsed = (datetime.now() - t0).total_seconds()
    log(f"  done {elapsed:.0f}s, final loss={avg_loss:.4f}")
    
    # --- eval before deploying ---
    log(f"── eval {EVAL_GAMES}局 ──")
    model.eval()
    mp = os.path.join(TRAINER_DIR, 'data/model.pt')
    model.save(mp)
    # export ONNX for eval
    onnx_tmp = os.path.join(TRAINER_DIR, 'data/model.onnx')
    dummy = torch.randn(1, 364)
    torch.onnx.export(model, dummy, onnx_tmp,
        input_names=['state'], output_names=['v'],
        dynamic_axes={'state': {0: 'batch'}, 'v': {0: 'batch'}},
        opset_version=15)
    # copy ONNX to hjha for eval
    eval_dest = os.path.join(HJHA_DIR, 'model.onnx')
    shutil.copy(onnx_tmp, eval_dest)
    data_src = onnx_tmp + '.data'
    if os.path.exists(data_src):
        shutil.copy(data_src, eval_dest + '.data')
    
    winrate = 0.0
    try:
        ok, out = run(['/usr/lib/dotnet/dotnet', 'run', '--project', 'hjha-console', '-c', 'Release',
                       '--', '--eval', str(EVAL_GAMES)], to=600)
        sf = os.path.join(HJHA_DIR, 'stress_summary.txt')
        if os.path.exists(sf):
            with open(sf) as f:
                summary = f.read()
            log(f"  {summary.strip()}")
            # parse Neural win rate from #BOT_SCORES (总分法: 每局 Score>0 算赢)
            in_bot_scores = False
            for line in summary.splitlines():
                line = line.strip()
                if line == "#BOT_SCORES":
                    in_bot_scores = True
                    continue
                if in_bot_scores and line.startswith("NEURAL="):
                    # Format: NEURAL=wins/games,wr=XX.X%,score=YYY
                    m = re.search(r'wr=(\d+\.?\d*)%', line)
                    if m:
                        winrate = float(m.group(1))
                    break
            if winrate == 0.0:
                log("  ⚠️ BOT_SCORES not found, fallback to legacy parse")
        log(f"  Neural winrate: {winrate:.1f}%")
    except Exception as e:
        log(f"  eval failed: {e}")
    
    # --- rollback if worse ---
    if prev_winrate is not None and winrate < prev_winrate * 0.9:
        log(f"  ⚠️ winrate dropped {prev_winrate:.1f}% → {winrate:.1f}%, rolling back model")
        # restore previous model.pt and ONNX
        prev_pt = os.path.join(TRAINER_DIR, 'data/model_prev.pt')
        if os.path.exists(prev_pt):
            shutil.copy(prev_pt, mp)
            log(f"  restored model.pt from backup")
        return False, prev_winrate
    
    # --- save and deploy ---
    # backup current as prev
    if os.path.exists(mp):
        shutil.copy(mp, os.path.join(TRAINER_DIR, 'data/model_prev.pt'))
    
    # final ONNX copy to hjha (already done above for eval, but ensure)
    dest = os.path.join(HJHA_DIR, 'model.onnx')
    shutil.copy(onnx_tmp, dest)
    if os.path.exists(data_src):
        shutil.copy(data_src, dest + '.data')
    
    # archive round
    tag = datetime.now().strftime('%m%d_%H%M')
    rd = os.path.join(HJHA_DIR, f'training_rounds/r{tag}')
    os.makedirs(rd, exist_ok=True)
    shutil.copy(onnx_tmp, os.path.join(rd, 'model.onnx'))
    if os.path.exists(data_src):
        shutil.copy(data_src, os.path.join(rd, 'model.onnx.data'))
    shutil.copy(mp, os.path.join(rd, 'model.pt'))
    if os.path.exists(csv_dir):
        shutil.move(csv_dir, os.path.join(rd, 'training_data'))
    os.makedirs(csv_dir, exist_ok=True)
    
    sz = os.path.getsize(dest) + (os.path.getsize(dest + '.data') if os.path.exists(dest + '.data') else 0)
    log(f"  ONNX → {dest} ({sz//1024}KB), 存档 → {rd}")
    
    # --- update replay buffer ---
    try:
        if os.path.exists(REPLAY_PATH):
            with open(REPLAY_PATH, 'rb') as f:
                replay_data = pickle.load(f)
        else:
            replay_data = []
        replay_data.append((cur_states, cur_results))
        if len(replay_data) > REPLAY_MAX_ROUNDS * 2:
            replay_data = replay_data[-REPLAY_MAX_ROUNDS * 2:]
        with open(REPLAY_PATH, 'wb') as f:
            pickle.dump(replay_data, f)
        log(f"  replay buffer: {len(replay_data)} rounds stored")
    except Exception as e:
        log(f"  replay save failed: {e}")
    
    return True, winrate

def main():
    log("=" * 50)
    log("持续训练: 20k局/轮 | Neural-only + replay + eval")
    log("=" * 50)
    
    rn, total = 0, 0
    winrate = None  # track best winrate for rollback
    
    while True:
        rn += 1
        games = 20000
        log(f"\n── 第{rn}轮: {games}局 ──")
        
        csv = selfplay(games)
        if csv == 0:
            log("自对弈失败! 跳过")
            continue
        total += games
        
        ok, wr = train_v(prev_winrate=winrate)
        if not ok:
            log("训练/eval失败! 跳过")
            continue
        # 训练成功, 清理备份
        backup = os.path.join(HJHA_DIR, 'training_data_prev')
        if os.path.exists(backup): shutil.rmtree(backup)
        if wr is not None:
            winrate = wr
        
        log(f"累计: {rn}轮, {total}局, best wr={winrate:.1f}%" if winrate else f"累计: {rn}轮, {total}局")

if __name__ == '__main__':
    main()