Mujoco_WASM/web_platform/src/tuning/ScalarChart.tsx

import { useEffect, useMemo, useRef } from 'react';
import { RotateCcw, ZoomIn, ZoomOut } from 'lucide-react';
import uPlot from 'uplot';
import type { ScalarSeries } from '../training/types';

const COLORS = ['#38d39f', '#60a5fa', '#f59e0b', '#f472b6', '#a78bfa', '#fb7185'];

function smooth(values: (number | null)[], factor: number): (number | null)[] {
  if (factor <= 0) return values;
  let previous: number | undefined;
  return values.map((value) => {
    if (value === null) return null;
    previous = previous === undefined ? value : factor * previous + (1 - factor) * value;
    return previous;
  });
}

export function ScalarChart({
  series,
  smoothing,
  title = '训练与评估 Scalars',
}: {
  series: ScalarSeries[];
  smoothing: number;
  title?: string;
}) {
  const host = useRef<HTMLDivElement>(null);
  const chartRef = useRef<uPlot | null>(null);
  const trackZoom = useRef(false);
  const zoomRanges = useRef<Partial<Record<'x' | 'y', { min: number; max: number }>>>({});
  const prepared = useMemo(() => {
    const steps = Array.from(
      new Set(series.flatMap((item) => item.points.map((point) => point.step))),
    ).sort((a, b) => a - b);
    const columns: uPlot.AlignedData = [steps];
    for (const item of series) {
      const byStep = new Map(item.points.map((point) => [point.step, point.value]));
      columns.push(
        smooth(
          steps.map((step) => byStep.get(step) ?? null),
          smoothing,
        ),
      );
    }
    return columns;
  }, [series, smoothing]);

  useEffect(() => {
    if (!host.current || series.length === 0 || prepared[0].length === 0) return;
    const element = host.current;
    const width = Math.max(280, Math.floor(element.getBoundingClientRect().width));
    trackZoom.current = false;
    const chart = new uPlot(
      {
        width,
        height: 280,
        cursor: { drag: { x: true, y: true, setScale: true } },
        scales: { x: { time: false } },
        hooks: {
          setScale: [
            (instance, key) => {
              if (!trackZoom.current || (key !== 'x' && key !== 'y')) return;
              const scale = instance.scales[key];
              if (typeof scale.min === 'number' && typeof scale.max === 'number')
                zoomRanges.current[key] = { min: scale.min, max: scale.max };
            },
          ],
        },
        axes: [
          { stroke: '#8fa0b5', grid: { stroke: '#213044' } },
          { stroke: '#8fa0b5', grid: { stroke: '#213044' } },
        ],
        series: [
          { label: 'Step' },
          ...series.map((item, index) => ({
            label: item.tag,
            stroke: COLORS[index % COLORS.length],
            width: 2,
            spanGaps: true,
          })),
        ],
      },
      prepared,
      element,
    );
    chartRef.current = chart;
    trackZoom.current = true;
    for (const key of ['x', 'y'] as const) {
      const range = zoomRanges.current[key];
      if (range) chart.setScale(key, range);
    }
    const wheelZoom = (event: WheelEvent) => {
      if (!event.deltaY) return;
      event.preventDefault();
      event.stopPropagation();
      const bounds = chart.over.getBoundingClientRect();
      if (!bounds.width || !bounds.height) return;
      const xRatio = Math.min(1, Math.max(0, (event.clientX - bounds.left) / bounds.width));
      const yRatio = Math.min(1, Math.max(0, (event.clientY - bounds.top) / bounds.height));
      const unit = event.deltaMode === 1 ? 16 : event.deltaMode === 2 ? window.innerHeight : 1;
      const factor = Math.min(2, Math.max(0.5, Math.exp(event.deltaY * unit * 0.002)));
      for (const key of ['x', 'y'] as const) {
        const scale = chart.scales[key];
        if (typeof scale.min !== 'number' || typeof scale.max !== 'number') continue;
        const ratio = key === 'x' ? xRatio : 1 - yRatio;
        const anchor = scale.min + (scale.max - scale.min) * ratio;
        chart.setScale(key, {
          min: anchor - (anchor - scale.min) * factor,
          max: anchor + (scale.max - anchor) * factor,
        });
      }
    };
    chart.over.addEventListener('wheel', wheelZoom, { passive: false });
    let frame = 0;
    let lastWidth = width;
    const observer = new ResizeObserver(() => {
      window.cancelAnimationFrame(frame);
      frame = window.requestAnimationFrame(() => {
        const nextWidth = Math.max(280, Math.floor(element.getBoundingClientRect().width));
        if (nextWidth !== lastWidth) {
          lastWidth = nextWidth;
          chart.setSize({ width: nextWidth, height: 280 });
        }
      });
    });
    observer.observe(element);
    return () => {
      observer.disconnect();
      chart.over.removeEventListener('wheel', wheelZoom);
      window.cancelAnimationFrame(frame);
      trackZoom.current = false;
      chartRef.current = null;
      chart.destroy();
    };
  }, [prepared, series]);

  const zoom = (factor: number) => {
    const chart = chartRef.current;
    if (!chart) return;
    for (const key of ['x', 'y']) {
      const scale = chart.scales[key];
      if (typeof scale.min !== 'number' || typeof scale.max !== 'number') continue;
      const center = (scale.min + scale.max) / 2;
      const radius = ((scale.max - scale.min) * factor) / 2 || 1;
      chart.setScale(key, { min: center - radius, max: center + radius });
    }
  };
  const resetZoom = () => {
    const chart = chartRef.current;
    if (!chart) return;
    zoomRanges.current = {};
    trackZoom.current = false;
    chart.setData(chart.data, true);
    trackZoom.current = true;
  };

  if (series.length === 0)
    return (
      <div className="grid h-[280px] place-items-center rounded-lg border border-border bg-app text-xs text-text-tertiary">
        当前 trial 尚无 scalar 数据
      </div>
    );
  return (
    <section className="min-w-0 overflow-hidden rounded-lg border border-border bg-app">
      <header className="flex items-center justify-between gap-2 border-b border-border px-3 py-2">
        <h3 className="min-w-0 truncate text-[10px] font-semibold" title={title}>
          {title}
        </h3>
        <div className="flex shrink-0 items-center gap-1">
          <button
            type="button"
            aria-label={`${title} 放大`}
            title="放大"
            className="rounded p-1 text-text-tertiary hover:bg-element-hover hover:text-text-primary"
            onClick={() => zoom(0.7)}
          >
            <ZoomIn className="h-3.5 w-3.5" />
          </button>
          <button
            type="button"
            aria-label={`${title} 缩小`}
            title="缩小"
            className="rounded p-1 text-text-tertiary hover:bg-element-hover hover:text-text-primary"
            onClick={() => zoom(1.4)}
          >
            <ZoomOut className="h-3.5 w-3.5" />
          </button>
          <button
            type="button"
            aria-label={`${title} 重置缩放`}
            title="重置缩放"
            className="rounded p-1 text-text-tertiary hover:bg-element-hover hover:text-text-primary"
            onClick={resetZoom}
          >
            <RotateCcw className="h-3.5 w-3.5" />
          </button>
        </div>
      </header>
      <div ref={host} className="min-w-0 w-full overflow-hidden" />
      <p className="border-t border-border px-3 py-1.5 text-[9px] text-text-tertiary">
        图表区域滚轮以指针位置缩放；也可拖拽框选，或使用右上角按钮缩放和复位。
      </p>
    </section>
  );
}