#!/usr/bin/env python3
"""蒸馏训练实时监控 — 在GPU服务器终端直接运行
不刷新不覆盖，只追加新行，保持完整的输出历史。

用法（GPU服务器上）：
  cd /root/autodl-tmp
  python3 scripts/watch_distill.py

或者直接：
  python3 /root/autodl-tmp/scripts/watch_distill.py
"""

import time, re, os, sys
from datetime import datetime

LOG = "/root/autodl-tmp/distill_mother.log"
OUT = "/root/autodl-tmp/output/qwen25-15b-shuangyan-distill"

def fmt_time():
    return datetime.now().strftime("%H:%M:%S")

def get_gpu():
    """解析nvidia-smi输出"""
    try:
        r = os.popen(
            "nvidia-smi --query-gpu=memory.used,memory.total,utilization.gpu,temperature.gpu "
            "--format=csv,noheader,nounits 2>/dev/null"
        ).read().strip()
        if not r:
            return "N/A"
        parts = [p.strip() for p in r.split(", ")]
        if len(parts) >= 2:
            used, total = parts[0], parts[1]
            pct = int(used) / int(total) * 100 if int(total) > 0 else 0
            gpu_util = parts[2] if len(parts) >= 3 else "?"
            temp = parts[3] if len(parts) >= 4 else "?"
            return f"显存: {used}/{total} MiB ({pct:.0f}%) | GPU: {gpu_util}% | 温度: {temp}°C"
        return r
    except:
        return "N/A"

def parse_loss(line):
    """从训练日志行解析loss"""
    m = re.search(r"'loss':\s*'?([\d.]+)'?", line)
    if m:
        return float(m.group(1))
    m = re.search(r"loss[=:]\s*([\d.]+)", line)
    if m:
        return float(m.group(1))
    m = re.search(r"loss=([\d.]+)", line)
    if m:
        return float(m.group(1))
    return None

def parse_step(line):
    """解析训练步数和epoch"""
    m = re.search(r"(\d+)/(\d+)\s+\[", line)
    if m:
        return int(m.group(1)), int(m.group(2))
    m = re.search(r"step=(\d+)", line)
    if m:
        return int(m.group(1)), None
    return None, None

def parse_progress(line):
    """解析进度条百分比"""
    m = re.search(r"(\d+)%\|", line)
    if m:
        return int(m.group(1))
    return None

def parse_eta(line):
    """解析剩余时间"""
    m = re.search(r"<(\d+:\d+)", line)
    if m:
        return m.group(1)
    return None

def parse_epoch(line):
    m = re.search(r"'epoch':\s*'?([\d.]+)'?", line)
    if m:
        return float(m.group(1))
    return None

def main():
    print("=" * 60)
    print(f"  铸渊蒸馏监控 · {fmt_time()}")
    print(f"  Watch: {LOG}")
    print(f"  不刷新不覆盖，只追加新行")
    print("=" * 60)
    print()

    # 先读已有日志
    last_size = 0
    if os.path.exists(LOG):
        last_size = os.path.getsize(LOG)
        with open(LOG) as f:
            for line in f:
                line = line.rstrip()
                if line:
                    print(f"  [{fmt_time()}] {line}")

    gpu_interval = 15  # 每15秒打一次GPU状态
    last_gpu = 0
    last_loss = None
    last_step = None
    total_steps = None
    last_epoch = None
    progress = None

    print()
    print("-" * 40)
    print(f"  [{fmt_time()}] 🔄 进入实时监控模式，每2秒刷新")
    print("-" * 40)
    sys.stdout.flush()

    try:
        while True:
            now = time.time()
            
            # 读新增日志行
            if os.path.exists(LOG):
                new_size = os.path.getsize(LOG)
                if new_size > last_size:
                    with open(LOG) as f:
                        f.seek(last_size)
                        for line in f:
                            line = line.rstrip()
                            if not line:
                                continue
                            print(f"  [{fmt_time()}] {line}")
                            
                            # 解析关键指标
                            loss = parse_loss(line)
                            if loss is not None:
                                last_loss = loss
                            
                            step, total = parse_step(line)
                            if step is not None:
                                last_step = step
                                if total is not None:
                                    total_steps = total
                            
                            pct = parse_progress(line)
                            if pct is not None:
                                progress = pct
                            
                            epoch = parse_epoch(line)
                            if epoch is not None:
                                last_epoch = epoch
                    
                    last_size = new_size
                    sys.stdout.flush()
            
            # 每15秒打一次GPU和进度摘要
            if now - last_gpu >= gpu_interval:
                last_gpu = now
                gpu_info = get_gpu()
                
                summary_parts = [f"[{fmt_time()}] 📊 {gpu_info}"]
                if last_loss is not None:
                    summary_parts.append(f"loss={last_loss:.4f}")
                if last_step is not None and total_steps is not None:
                    summary_parts.append(f"step={last_step}/{total_steps}")
                if last_epoch is not None:
                    summary_parts.append(f"epoch={last_epoch:.2f}")
                if progress is not None:
                    summary_parts.append(f"progress={progress}%")
                
                print(f"  {' | '.join(summary_parts)}")
                sys.stdout.flush()
            
            time.sleep(2)
    
    except KeyboardInterrupt:
        print()
        print(f"  [{fmt_time()}] 👋 监控已退出")
        print(f"  最后状态: loss={last_loss}, step={last_step}/{total_steps}, epoch={last_epoch}")
        sys.exit(0)

if __name__ == "__main__":
    main()