plot_latest_loss.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
可视化最近一次生成的 lossX.csv：

新增：
    - 适配新的表头：step,eval_mse_nmlz,eval_d2_nmlz,eval_combined,...
    - 同时绘制多条曲线（至少 MSE_nmlz 与 D2_nmlz；若存在 eval_combined 也可附加）
    - 向后兼容旧格式（只有 step,value 两列时仍然绘制单条）

特性：
    - 自动在当前目录查找匹配 loss*.csv 的最新文件（按修改时间）
    - x 轴线性；y 轴线性，从0开始
    - 使用 Pillow 绘制折线图，保存为 latest_loss_plot.png 或 latest_loss_plot_multi.png
    - 若用户指定 --out 则使用指定输出路径

依赖：仅需 Pillow（requirements.txt 已包含 pillow）

使用示例：
    python plot_latest_loss.py
    python plot_latest_loss.py --file loss3.csv --open
    python plot_latest_loss.py --out out.png --width 1600 --height 900
"""
from __future__ import annotations

import argparse
import csv
import math
import os
import re
import sys
from dataclasses import dataclass
from typing import List, Tuple, Optional, Dict

from PIL import Image, ImageDraw, ImageFont


HERE = os.path.abspath(os.path.dirname(__file__))


@dataclass
class Series:
    name: str
    points: List[Tuple[float, float]]  # (x=step, y=value)


def is_float(s: str) -> bool:
    try:
        float(s)
        return True
    except Exception:
        return False


def find_latest_loss_csv(folder: str) -> Optional[str]:
    """在目录中查找最近修改的标准命名 loss*.csv 文件。
    仅匹配正则：^loss(\d*)\.csv$（例如 loss.csv, loss1.csv, loss4.csv），排除“副本”等拷贝文件。
    """
    pat = re.compile(r"^loss(\d*)\.csv$", re.IGNORECASE)
    candidates = []
    for name in os.listdir(folder):
        if not pat.match(name):
            continue
        path = os.path.join(folder, name)
        if not os.path.isfile(path):
            continue
        try:
            mtime = os.path.getmtime(path)
        except OSError:
            continue
        candidates.append((mtime, path))
    if not candidates:
        return None
    candidates.sort(key=lambda x: x[0], reverse=True)
    return candidates[0][1]


def parse_loss_csv(path: str) -> List[Series]:
    """解析 loss*.csv，多列支持：
    优先匹配列名：
        step 列："step"
        MSE：依次尝试 eval_mse_nmlz, mse_nmlz, eval_mse, mse
        D2：  eval_d2_nmlz, d2_nmlz, eval_d2, d2
        Combined：eval_combined, combined
    若未检测到 header 或只有两列数字，则退化为旧格式 (step,value)。
    返回 Series 列表（过滤 y<=0 的点）。
    """
    with open(path, 'r', encoding='utf-8', newline='') as f:
        rdr = csv.reader(f)
        rows = [row for row in rdr if row]
    if not rows:
        return []
    header = [c.strip() for c in rows[0]]
    lower_header = [h.lower() for h in header]

    def find_idx(candidates: List[str]) -> Optional[int]:
        for name in candidates:
            if name in lower_header:
                return lower_header.index(name)
        return None

    # 判断 header 是否真的包含 step（否则视为无header旧格式）
    step_idx = find_idx(["step"]) if any(h == "step" for h in lower_header) else None

    series_list: List[Series] = []
    if step_idx is not None:
        mse_idx = find_idx(["eval_mse_nmlz", "mse_nmlz", "eval_mse", "mse"])
        d2_idx = find_idx(["eval_d2_nmlz", "d2_nmlz", "eval_d2", "d2"])
        comb_idx = find_idx(["eval_combined", "combined"])

        # 解析数据行
        for name, col_idx, title in [
            ("mse_nmlz", mse_idx, "MSE_nmlz"),
            ("d2_nmlz", d2_idx, "D2_nmlz"),
            ("combined", comb_idx, "Combined")
        ]:
            if col_idx is None:
                continue
            pts: List[Tuple[float, float]] = []
            for row in rows[1:]:
                if len(row) <= max(step_idx, col_idx):
                    continue
                xs = row[step_idx].strip()
                ys = row[col_idx].strip()
                if not (is_float(xs) and is_float(ys)):
                    continue
                x = float(xs); y = float(ys)
                if y > 0:
                    pts.append((x, y))
            if pts:
                series_list.append(Series(name=title, points=pts))
        return series_list
    else:
        # 旧格式：尝试按两列数字读取
        points: List[Tuple[float, float]] = []
        for row in rows:
            if len(row) < 2:
                continue
            xs, ys = row[0].strip(), row[1].strip()
            if is_float(xs) and is_float(ys):
                x = float(xs); y = float(ys)
                if y > 0:
                    points.append((x, y))
        if points:
            return [Series(name=os.path.basename(path), points=points)]
        return []


def _nice_step(span: float, max_ticks: int) -> float:
    """计算一个漂亮的步长，偏好 1/2/5×10^k。"""
    if span <= 0 or max_ticks <= 1:
        return span if span > 0 else 1.0
    raw = span / (max_ticks - 1)
    power = 10 ** math.floor(math.log10(raw))
    base = raw / power
    if base <= 1:
        nice = 1
    elif base <= 2:
        nice = 2
    elif base <= 5:
        nice = 5
    else:
        nice = 10
    return nice * power


def nice_ticks_linear(vmin: float, vmax: float, max_ticks: int = 8) -> List[float]:
    """在线性空间内生成漂亮刻度（尽可能整数，避免科学计数法）。"""
    if vmax == vmin:
        return [vmin]
    if vmin > vmax:
        vmin, vmax = vmax, vmin
    span = vmax - vmin
    step = _nice_step(span, max_ticks)
    start = math.floor(vmin / step) * step
    end = math.ceil(vmax / step) * step
    ticks = []
    x = start
    # 防止浮点误差导致无限循环
    guard = 0
    while x <= end + step * 1e-9 and guard < 1000:
        ticks.append(x)
        x += step
        guard += 1
    # 保证至少一个刻度
    if not ticks:
        ticks = [vmin]
    return ticks


def _is_int_like(x: float, tol: float = 1e-9) -> bool:
    return abs(x - round(x)) <= tol


def format_tick(x: float, step: Optional[float] = None) -> str:
    """格式化刻度：优先整数，禁止科学计数法，按步长决定小数位。"""
    if step is None:
        step = 0.0
    # 如果接近整数或步长>=1，则用整数
    if _is_int_like(x, max(1e-9, 0.01 * step)) or step >= 1:
        return f"{int(round(x))}"
    # 根据步长决定小数位
    if step >= 0.5:
        return f"{x:.0f}"
    elif step >= 0.1:
        return f"{x:.1f}"
    elif step >= 0.01:
        return f"{x:.2f}"
    else:
        return f"{x:.3f}"


def draw_plot_multi(series_list: List[Series], width: int = 1200, height: int = 600, out_path: str = None,
                    open_after: bool = False) -> str:
    if not series_list:
        raise RuntimeError("没有可绘制的数据序列")
    # 过滤掉空序列
    series_list = [s for s in series_list if s.points]
    if not series_list:
        raise RuntimeError("序列中没有有效数据点")

    # 合并范围
    all_x = [x for s in series_list for (x, _) in s.points]
    all_y = [y for s in series_list for (_, y) in s.points if y > 0]
    if not all_y:
        raise RuntimeError("所有 y 均 <=0，无法绘制")
    x_min, x_max = min(all_x), max(all_x)
    
    # 检查是否有 D2_Loss 序列
    has_d2_loss = any(s.name.lower() in ["d2_nmlz", "d2_loss", "d2"] for s in series_list)
    
    # 计算主要轴范围（排除D2_Loss）
    if has_d2_loss:
        main_y = [y for s in series_list for (_, y) in s.points if y > 0 and s.name.lower() not in ["d2_nmlz", "d2_loss", "d2"]]
        y_min, y_max = 0, max(main_y) if main_y else max(all_y)
        
        # 计算D2_Loss的范围
        d2_series = None
        d2_y = []
        for s in series_list:
            if s.name.lower() in ["d2_nmlz", "d2_loss", "d2"]:
                d2_series = s
                d2_y = [y for (_, y) in s.points if y > 0]
                break
        d2_y_min, d2_y_max = 0, max(d2_y) if d2_y else 0
    else:
        y_min, y_max = 0, max(all_y)

    # 画布与边距 - 如果有双轴，右侧边距增加
    w, h = width, height
    left, right, top, bottom = 90, 40, 60, 80
    if has_d2_loss:
        right = 90  # 增加右侧边距用于第二个 y轴
    plot_w = max(1, w - left - right)
    plot_h = max(1, h - top - bottom)

    img = Image.new('RGB', (w, h), color=(255, 255, 255))
    draw = ImageDraw.Draw(img)

    # 字体（尽量使用默认字体）
    try:
        font = ImageFont.truetype("arial.ttf", 14)
        font_bold = ImageFont.truetype("arial.ttf", 18)
    except Exception:
        font = ImageFont.load_default()
        font_bold = font

    # 坐标映射
    def x_to_px(x: float) -> int:
        if x_max == x_min:
            return left
        return int(left + (x - x_min) / (x_max - x_min) * plot_w)

    def y_to_px(y: float) -> int:
        if y_max == y_min:
            return h - bottom
        return int(h - bottom - (y - y_min) / (y_max - y_min) * plot_h)
    
    # D2_Loss 的 y 轴映射（如果有双轴）
    def d2_y_to_px(y: float) -> int:
        if not has_d2_loss or d2_y_max == d2_y_min:
            return h - bottom
        return int(h - bottom - (y - d2_y_min) / (d2_y_max - d2_y_min) * plot_h)

    # 绘制坐标轴
    axis_color = (0, 0, 0)
    grid_color = (230, 230, 230)
    draw.line([(left, h - bottom), (w - right, h - bottom)], fill=axis_color, width=2)
    draw.line([(left, h - bottom), (left, top)], fill=axis_color, width=2)
    if has_d2_loss:
        draw.line([(w - right, h - bottom), (w - right, top)], fill=axis_color, width=2)

    # 网格与刻度
    # x 轴：线性空间漂亮刻度（直接使用 step 值）
    xticks = nice_ticks_linear(x_min, x_max, 8)
    x_step = xticks[1] - xticks[0] if len(xticks) >= 2 else (x_max - x_min if x_max != x_min else 1)

    # y 轴：在 log2 域上生成刻度，当范围小时使用小数刻度
    def nice_ticks_log2(y_log_min: float, y_log_max: float, max_ticks: int = 8):
        if y_log_min > y_log_max:
            y_log_min, y_log_max = y_log_max, y_log_min
        if y_log_min == y_log_max:
            return [y_log_min]
        
        span = y_log_max - y_log_min
        
        # 如果范围很小（<2），使用更细的刻度
        if span < 2:
            # 使用 0.2 或 0.5 作为步长
            if span < 0.5:
                step = 0.1
            elif span < 1:
                step = 0.2
            else:
                step = 0.5
            
            start = math.ceil(y_log_min / step) * step
            end = math.floor(y_log_max / step) * step
            ticks = []
            current = start
            while current <= end + step * 0.1:  # 小的误差容忍
                ticks.append(current)
                current += step
                if len(ticks) > max_ticks:  # 防止太多刻度
                    break
            
            # 确保至少有端点
            if not ticks or ticks[0] > y_log_min + 0.01:
                ticks.insert(0, round(y_log_min, 1))
            if not ticks or ticks[-1] < y_log_max - 0.01:
                ticks.append(round(y_log_max, 1))
        else:
            # 范围较大时，使用整数刻度
            imin = math.floor(y_log_min)
            imax = math.ceil(y_log_max)
            span_int = imax - imin
            if span_int + 1 <= max_ticks:
                step = 1
            else:
                step = max(1, math.ceil(span_int / (max_ticks - 1)))
            start = math.ceil(imin / step) * step
            ticks = list(range(start, imax + 1, step))
            if not ticks:
                ticks = [imin]
        
        return ticks

    yticks = nice_ticks_linear(y_min, y_max, 8)
    y_step = yticks[1] - yticks[0] if len(yticks) >= 2 else (y_max - y_min if y_max != y_min else 1)
    
    # D2_Loss 的刻度（如果有双轴）
    d2_yticks = []
    if has_d2_loss:
        d2_yticks = nice_ticks_linear(d2_y_min, d2_y_max, 6)  # 右侧轴使用较少的刻度

    for xv in xticks:
        xpx = x_to_px(xv)
        draw.line([(xpx, h - bottom), (xpx, top)], fill=grid_color, width=1)
        draw.line([(xpx, h - bottom), (xpx, h - bottom + 5)], fill=axis_color, width=1)
        label = format_tick(xv, x_step)
        bbox = draw.textbbox((0, 0), label, font=font)
        tw, th = bbox[2] - bbox[0], bbox[3] - bbox[1]
        draw.text((xpx - tw // 2, h - bottom + 8), label, fill=axis_color, font=font)

    for yv in yticks:
        # 直接在线性空间绘制
        ypx = y_to_px(yv)
        draw.line([(left, ypx), (w - right, ypx)], fill=grid_color, width=1)
        draw.line([(left - 5, ypx), (left, ypx)], fill=axis_color, width=1)
        
        # 显示数值
        label = format_tick(yv, y_step)
        
        bbox = draw.textbbox((0, 0), label, font=font)
        tw, th = bbox[2] - bbox[0], bbox[3] - bbox[1]
        draw.text((left - tw - 8, ypx - th // 2), label, fill=axis_color, font=font)
    
    # 右侧 y 轴刻度（D2_Loss）
    if has_d2_loss:
        for yv in d2_yticks:
            ypx = d2_y_to_px(yv)
            # 右侧轴只有局部网格线
            draw.line([(w - right, ypx), (w - right + 5, ypx)], fill=axis_color, width=1)
            
            # 显示数值
            label = format_tick(yv, 0)  # D2_Loss通常不需要特殊的步长格式化
            
            bbox = draw.textbbox((0, 0), label, font=font)
            tw, th = bbox[2] - bbox[0], bbox[3] - bbox[1]
            draw.text((w - right + 8, ypx - th // 2), label, fill=axis_color, font=font)
    title = "Latest loss plot (linear y)"
    if has_d2_loss:
        title += " - Dual axis for D2_Loss"
    tb = draw.textbbox((0, 0), title, font=font_bold)
    tw, th = tb[2] - tb[0], tb[3] - tb[1]
    draw.text(((w - tw) // 2, 12), title, fill=(10, 10, 10), font=font_bold)

    x_label = "step (linear)"
    # 纵轴显示数值（线性刻度）
    y_label = "value (linear scale)"
    # x 轴标签
    xb = draw.textbbox((0, 0), x_label, font=font)
    xtw, xth = xb[2] - xb[0], xb[3] - xb[1]
    draw.text(((w - xtw) // 2, h - bottom + 35), x_label, fill=axis_color, font=font)
    # y 轴标签（竖排）
    # 简单起见，逐字竖排
    y_chars = list(y_label)
    y_total_h = 0
    for c in y_chars:
        cb = draw.textbbox((0, 0), c, font=font)
        y_total_h += cb[3] - cb[1]
    y_start = top + (plot_h - y_total_h) // 2
    cy = y_start
    for c in y_chars:
        cb = draw.textbbox((0, 0), c, font=font)
        cw, ch = cb[2] - cb[0], cb[3] - cb[1]
        draw.text((15, cy), c, fill=axis_color, font=font)
        cy += ch
    
    # y 轴标签（右侧，如果有双轴）
    if has_d2_loss:
        y_chars_right = list("D2_Loss (linear scale)")
        y_total_h_right = 0
        for c in y_chars_right:
            cb = draw.textbbox((0, 0), c, font=font)
            y_total_h_right += cb[3] - cb[1]
        y_start_right = top + (plot_h - y_total_h_right) // 2
        cy = y_start_right
        for c in y_chars_right:
            cb = draw.textbbox((0, 0), c, font=font)
            cw, ch = cb[2] - cb[0], cb[3] - cb[1]
            draw.text((w - right + 15, cy), c, fill=axis_color, font=font)
            cy += ch

    # 折线（多序列）
    palette = [
        (31,119,180),   # 蓝 - mse
        (214,39,40),    # 红 - d2
        (44,160,44),    # 绿 - combined
        (255,127,14),   # 橙
        (148,103,189)   # 紫
    ]
    legend_items = []
    for idx, s in enumerate(series_list):
        color = palette[idx % len(palette)]
        
        # 选择使用哪个 y 轴
        use_right_axis = has_d2_loss and s.name.lower() in ["d2_nmlz", "d2_loss", "d2"]
        
        if use_right_axis:
            pts = [(x_to_px(x), d2_y_to_px(y)) for (x, y) in s.points]
        else:
            pts = [(x_to_px(x), y_to_px(y)) for (x, y) in s.points]
            
        if len(pts) == 1:
            xpx, ypx = pts[0]
            r = 3
            draw.ellipse((xpx - r, ypx - r, xpx + r, ypx + r), fill=color, outline=color)
        else:
            draw.line(pts, fill=color, width=3)
        # 末尾标记当前值
        last_x, last_y = s.points[-1]
        lx, ly = x_to_px(last_x), (d2_y_to_px(last_y) if use_right_axis else y_to_px(last_y))
        draw.text((lx + 6, ly - 8), f"{s.name}:{last_y:.4g}", fill=color, font=font)
        legend_items.append((s.name, color))

    # 图例（左上角）
    lg_x, lg_y = left + 6, top + 6
    for (name, color) in legend_items:
        box_h = 14
        draw.rectangle((lg_x, lg_y + 2, lg_x + 18, lg_y + 12), fill=color)
        draw.text((lg_x + 24, lg_y), name, fill=(0,0,0), font=font)
        lg_y += box_h + 4

    # 边框
    draw.rectangle((left, top, w - right, h - bottom), outline=(180, 180, 180), width=1)

    # 输出
    default_name = "latest_loss_plot_multi.png" if len(series_list) > 1 else "latest_loss_plot.png"
    out_path = out_path or os.path.join(HERE, default_name)
    img.save(out_path)

    print(f"已生成图像: {out_path}")
    if open_after:
        try:
            if sys.platform.startswith('win'):
                os.startfile(out_path)  # type: ignore[attr-defined]
            elif sys.platform == 'darwin':
                os.system(f"open '{out_path}'")
            else:
                os.system(f"xdg-open '{out_path}'")
        except Exception as e:
            print(f"无法自动打开图像: {e}")
    return out_path


def main():
    parser = argparse.ArgumentParser(description="绘制最近的 loss*.csv（y=线性，从0开始）")
    parser.add_argument("--file", dest="file", type=str, default=None, help="指定 CSV 文件（默认自动查找最新 loss*.csv）")
    parser.add_argument("--out", dest="out", type=str, default=None, help="输出 PNG 路径（默认 latest_loss_plot.png）")
    parser.add_argument("--width", dest="width", type=int, default=1200, help="图像宽度像素")
    parser.add_argument("--height", dest="height", type=int, default=600, help="图像高度像素")
    parser.add_argument("--open", dest="open_after", action="store_true", help="生成后尝试打开图像（默认已打开）")
    parser.add_argument("--no-open", dest="no_open", action="store_true", help="生成后不打开图像")
    args = parser.parse_args()

    folder = HERE
    csv_path = args.file
    if not csv_path:
        csv_path = find_latest_loss_csv(folder)
        if not csv_path:
            print("未找到任何 loss*.csv 文件。")
            sys.exit(1)
    if not os.path.isabs(csv_path):
        csv_path = os.path.join(folder, csv_path)

    if not os.path.exists(csv_path):
        print(f"CSV 文件不存在: {csv_path}")
        sys.exit(1)

    series_list = parse_loss_csv(csv_path)
    if not series_list:
        print("未解析到任何有效数据点。")
        sys.exit(2)

    # 默认自动打开，除非显式 --no-open；如果用户传了 --open 也强制打开
    open_after = True
    if args.no_open:
        open_after = False
    if args.open_after:
        open_after = True

    try:
        draw_plot_multi(series_list, width=args.width, height=args.height, out_path=args.out, open_after=open_after)
    except Exception as e:
        print(f"绘图失败: {e}")
        sys.exit(3)


if __name__ == "__main__":
    main()