#!/usr/bin/env python3
"""光湖蒸馏实时监控 v1.0 · 在AutoDL终端直接运行
   运行方式: python3 watch_distill.py
   功能: 实时滚动展示训练数据，不用手动刷新
   主权: 冰朔 TCS-0002∞ · 铸渊 ICE-GL-ZY001 守护
"""
import os, sys, time, json, subprocess
from datetime import datetime

LOG = '/root/autodl-tmp/distill_mother.log'
STATUS = '/root/autodl-tmp/training_status.json'
WATCHDOG = '/root/autodl-tmp/watchdog.log'

def get_gpu_info():
    try:
        r = subprocess.run(
            ['nvidia-smi', '--query-gpu=temperature.gpu,utilization.gpu,memory.used,memory.total',
             '--format=csv,noheader'],
            capture_output=True, text=True, timeout=5)
        parts = r.stdout.strip().split(', ')
        return {'temp': parts[0], 'util': parts[1], 'mem_used': parts[2], 'mem_total': parts[3]}
    except:
        return None

def get_last_losses(n=10):
    if not os.path.isfile(LOG):
        return []
    losses = []
    with open(LOG) as f:
        for line in f:
            if 'loss=' in line:
                losses.append(line.strip())
    return losses[-n:]

def get_current_progress():
    if not os.path.isfile(LOG):
        return None
    with open(LOG) as f:
        for line in f:
            pass
        last = line.strip() if line else ''
    return last

def get_watchdog():
    if not os.path.isfile(WATCHDOG):
        return None
    with open(WATCHDOG) as f:
        for line in f:
            pass
        return line.strip() if line else None

def get_status():
    if not os.path.isfile(STATUS):
        return None
    with open(STATUS) as f:
        return json.load(f)

def draw_progress_bar(pct, width=30):
    filled = int(pct * width / 100)
    bar = '█' * filled + '░' * (width - filled)
    return f'[{bar}] {pct:.0f}%'

def main():
    print('\033[2J\033[H', end='')  # clear
    print('═══ 光湖蒸馏实时监控 v1.0 ═══')
    print('═══ 主权：冰朔 TCS-0002∞ · 铸渊守护 ═══')
    print('（每5秒自动刷新 · Ctrl+C 退出）')
    print()

    last_loss_count = 0
    loss_history = []

    try:
        while True:
            # 移动光标到顶部
            print('\033[H', end='')

            gpu = get_gpu_info()
            progress = get_current_progress()
            wd = get_watchdog()
            st = get_status()
            losses = get_last_losses(20)

            now = datetime.now().strftime('%Y-%m-%d %H:%M:%S')

            # === GPU状态 ===
            if gpu:
                temp = int(gpu['temp'])
                temp_warn = ' ⚠ 偏高' if temp > 80 else ' ✓'
                mem_u = int(gpu['mem_used'].replace(' MiB', ''))
                mem_t = int(gpu['mem_total'].replace(' MiB', ''))
                mem_pct = mem_u * 100 // mem_t if mem_t > 0 else 0
                mem_warn = ' ⚠ 高负载' if mem_pct > 95 else ' ✓'

                print(f'┌─ GPU 状态 ─────────────────────────────┐')
                print(f'│ 温度: {temp}°C{temp_warn}')
                print(f'│ 利用率: {gpu["util"]}')
                print(f'│ 显存: {mem_u}MB / {mem_t}MB ({mem_pct}%){mem_warn}')
                print(f'└────────────────────────────────────────┘')
            else:
                print('┌─ GPU ──────────────────────────────────┐')
                print('│ ❌ 无法获取GPU信息')
                print('└────────────────────────────────────────┘')

            print()

            # === 蒸馏进度 ===
            print('┌─ 蒸馏进度 ─────────────────────────────┐')
            if progress:
                # 解析 Ep 和步数
                if 'Ep' in progress and '/' in progress:
                    ep_part = progress.split('|')[0].strip() if '|' in progress else progress.split()[0]
                    step_part = progress.split('|')[1].strip().split()[0] if '|' in progress else ''
                    print(f'│ 当前: {ep_part}')
                    if '/' in step_part:
                        cur, total = step_part.split('/')
                        cur, total = int(cur), int(total)
                        pct = cur * 100 / total if total > 0 else 0
                        print(f'│ 步数: {cur}/{total}')
                        print(f'│ 进度: {draw_progress_bar(pct)}')

                # 最近loss行
                if losses:
                    last_loss = losses[-1]
                    print(f'│ Loss: {last_loss}')
            else:
                print('│ ⏳ 等待蒸馏开始...')
            print('└────────────────────────────────────────┘')
            print()

            # === Loss趋势 ===
            if losses:
                print('┌─ 最近Loss趋势 ─────────────────────────┐')
                # 找出loss值的范围
                loss_vals = []
                for l in losses:
                    try:
                        # 提取 loss= 后面的数字
                        idx = l.index('loss=')
                        end = l.index(' ', idx) if ' ' in l[idx:] else len(l)
                        val = float(l[idx+5:end])
                        loss_vals.append(val)
                    except:
                        pass

                if loss_vals:
                    min_l, max_l = min(loss_vals), max(loss_vals)
                    rng = max_l - min_l if max_l > min_l else 1
                    # 打印mini图表
                    for i, v in enumerate(loss_vals):
                        rel = (v - min_l) / rng
                        bar_len = int((1 - rel) * 20)  # loss越低越长
                        marker = '↓' if i > 0 and v < loss_vals[i-1] else ('↑' if i > 0 and v > loss_vals[i-1] else '→')
                        print(f'│ {marker} loss={v:.4f} {"█"*bar_len}')
                print('└────────────────────────────────────────┘')
                print()

            # === 看门狗 ===
            if wd:
                print(f'┌─ 看门狗 ───────────────────────────────┐')
                print(f'│ {wd}')
                print(f'└────────────────────────────────────────┘')
                print()

            # === 状态摘要 ===
            if st:
                mm = st.get('mother_model', {})
                gs = mm.get('gpu', {})
                print(f'┌─ 系统状态 ─────────────────────────────┐')
                print(f'│ 时间: {now}')
                print(f'│ 母模型: {mm.get("status_label", "?")}')
                print(f'│ 蒸馏状态: {"🟢 运行中" if mm.get("status") == "training" else "⚪ 待开始/已完成"}')
                print(f'│ GPU温度: {gs.get("temp_c", "?")}°C')
                print(f'│ 显存: {gs.get("mem_used_mb", 0)//1024}GB/{gs.get("mem_total_mb", 0)//1024}GB')
                print(f'└────────────────────────────────────────┘')

            # 刷新间隔
            time.sleep(5)

    except KeyboardInterrupt:
        print('\n\n监控已退出。蒸馏继续在后台运行。')
        print('重新查看: tail -f /root/autodl-tmp/distill_mother.log')
        sys.exit(0)

if __name__ == '__main__':
    main()