From 10456dcfadba3edfc07e42995835fdd7ee837d0d Mon Sep 17 00:00:00 2001
From: JavaZero <71128095+JavaZeroo@users.noreply.github.com>
Date: Tue, 19 Aug 2025 19:36:49 +0800
Subject: [PATCH] feat: support step keyword alignment

---
 src/App.jsx                                   |  40 ++++-
 src/components/ChartContainer.jsx             | 139 ++++++++++++------
 .../__tests__/ChartContainer.test.jsx         |  31 ++++
 3 files changed, 167 insertions(+), 43 deletions(-)

diff --git a/src/App.jsx b/src/App.jsx
index fce2a15..53c9a07 100644
--- a/src/App.jsx
+++ b/src/App.jsx
@@ -40,6 +40,8 @@ function App() {
   const [xRange, setXRange] = useState({ min: undefined, max: undefined });
   const [maxStep, setMaxStep] = useState(0);
   const [sidebarVisible, setSidebarVisible] = useState(true);
+  const [alignSteps, setAlignSteps] = useState(false);
+  const [stepKeyword, setStepKeyword] = useState('step:');
 
   const handleFilesUploaded = useCallback((files) => {
     const filesWithDefaults = files.map(file => ({
@@ -340,12 +342,44 @@ function App() {
                   <h4 className="text-xs font-medium text-gray-700 mb-2">📊 图表显示</h4>
                   <p className="text-xs text-gray-500">上传文件后自动展示所有已配置的指标图表</p>
                 </div>
-                
+
+                <div className="border-t pt-3">
+                  <h4 className="text-xs font-medium text-gray-700 mb-2">步骤对齐</h4>
+                  <div className="space-y-2">
+                    <label className="flex items-center text-xs text-gray-700">
+                      <input
+                        type="checkbox"
+                        className="mr-2"
+                        checked={alignSteps}
+                        onChange={e => setAlignSteps(e.target.checked)}
+                      />
+                      启用基于关键词的步骤对齐
+                    </label>
+                    {alignSteps && (
+                      <div>
+                        <label
+                          htmlFor="step-keyword"
+                          className="block text-xs font-medium text-gray-700 mb-1"
+                        >
+                          关键词
+                        </label>
+                        <input
+                          id="step-keyword"
+                          type="text"
+                          value={stepKeyword}
+                          onChange={e => setStepKeyword(e.target.value)}
+                          className="w-full px-2 py-1 text-xs border border-gray-300 rounded-md focus:ring-2 focus:ring-blue-500 focus:border-blue-500 focus:outline-none"
+                        />
+                      </div>
+                    )}
+                  </div>
+                </div>
+
                 <div className="border-t pt-3">
                   <h4 className="text-xs font-medium text-gray-700 mb-2">基准线设置</h4>
                   <div className="space-y-3">
                     <div>
-                      <label 
+                      <label
                         htmlFor="relative-baseline"
                         className="block text-xs font-medium text-gray-700 mb-1"
                       >
@@ -414,6 +448,8 @@ function App() {
               xRange={xRange}
               onXRangeChange={setXRange}
               onMaxStepChange={setMaxStep}
+              alignSteps={alignSteps}
+              stepKeyword={stepKeyword}
             />
           </section>
         </main>
diff --git a/src/components/ChartContainer.jsx b/src/components/ChartContainer.jsx
index b10f0f5..40e34a1 100644
--- a/src/components/ChartContainer.jsx
+++ b/src/components/ChartContainer.jsx
@@ -68,7 +68,9 @@ export default function ChartContainer({
   absoluteBaseline = 0.005,
   xRange = { min: undefined, max: undefined },
   onXRangeChange,
-  onMaxStepChange
+  onMaxStepChange,
+  alignSteps = false,
+  stepKeyword = 'step:'
 }) {
   const chartRefs = useRef(new Map());
   const registerChart = useCallback((id, inst) => {
@@ -103,40 +105,76 @@ export default function ChartContainer({
       const lines = file.content.split('\n');
       const metricsData = {};
 
-      const extractByKeyword = (content, keyword) => {
-        const results = [];
-        const numberRegex = /[+-]?\d+(?:\.\d+)?(?:[eE][+-]?\d+)?/;
-        content.split('\n').forEach(line => {
-          const idx = line.toLowerCase().indexOf(keyword.toLowerCase());
-          if (idx !== -1) {
-            const after = line.substring(idx + keyword.length);
-            const match = after.match(numberRegex);
-            if (match) {
-              const v = parseFloat(match[0]);
-              if (!isNaN(v)) results.push(v);
+      const numberRegex = /[+-]?\d+(?:\.\d+)?(?:[eE][+-]?\d+)?/;
+
+      if (alignSteps) {
+        const metricKeys = metrics.map((m, idx) => m.name || m.keyword || `metric${idx + 1}`);
+        metricKeys.forEach(k => { metricsData[k] = []; });
+
+        lines.forEach(line => {
+          const stepIdx = line.toLowerCase().indexOf(stepKeyword.toLowerCase());
+          if (stepIdx === -1) return;
+          const afterStep = line.substring(stepIdx + stepKeyword.length);
+          const stepMatch = afterStep.match(/\d+/);
+          if (!stepMatch) return;
+          const step = parseInt(stepMatch[0]);
+          metrics.forEach((metric, mi) => {
+            const key = metricKeys[mi];
+            let value;
+            if (metric.mode === 'keyword') {
+              const idx = line.toLowerCase().indexOf(metric.keyword.toLowerCase());
+              if (idx !== -1) {
+                const after = line.substring(idx + metric.keyword.length);
+                const match = after.match(numberRegex);
+                if (match) value = parseFloat(match[0]);
+              }
+            } else if (metric.regex) {
+              const reg = new RegExp(metric.regex);
+              reg.lastIndex = 0;
+              const m = reg.exec(line);
+              if (m && m[1]) value = parseFloat(m[1]);
             }
-          }
+            if (value !== undefined && !isNaN(value)) {
+              metricsData[key].push({ x: step, y: value });
+            }
+          });
         });
-        return results;
-      };
-
-      metrics.forEach(metric => {
-        let values = [];
-        if (metric.mode === 'keyword') {
-          values = extractByKeyword(file.content, metric.keyword);
-        } else if (metric.regex) {
-          const reg = new RegExp(metric.regex);
-          lines.forEach(line => {
-            reg.lastIndex = 0;
-            const m = reg.exec(line);
-            if (m && m[1]) {
-              const v = parseFloat(m[1]);
-              if (!isNaN(v)) values.push(v);
+      } else {
+        const extractByKeyword = (content, keyword) => {
+          const results = [];
+          content.split('\n').forEach(line => {
+            const idx = line.toLowerCase().indexOf(keyword.toLowerCase());
+            if (idx !== -1) {
+              const after = line.substring(idx + keyword.length);
+              const match = after.match(numberRegex);
+              if (match) {
+                const v = parseFloat(match[0]);
+                if (!isNaN(v)) results.push(v);
+              }
             }
           });
-        }
-        metricsData[metric.name || metric.keyword] = values.map((v, i) => ({ x: i, y: v }));
-      });
+          return results;
+        };
+
+        metrics.forEach((metric, idx) => {
+          let values = [];
+          if (metric.mode === 'keyword') {
+            values = extractByKeyword(file.content, metric.keyword);
+          } else if (metric.regex) {
+            const reg = new RegExp(metric.regex);
+            lines.forEach(line => {
+              reg.lastIndex = 0;
+              const m = reg.exec(line);
+              if (m && m[1]) {
+                const v = parseFloat(m[1]);
+                if (!isNaN(v)) values.push(v);
+              }
+            });
+          }
+          const key = metric.name || metric.keyword || `metric${idx + 1}`;
+          metricsData[key] = values.map((v, i) => ({ x: i, y: v }));
+        });
+      }
 
       const range = file.config?.dataRange;
       if (range && (range.start > 0 || range.end !== undefined)) {
@@ -147,7 +185,7 @@ export default function ChartContainer({
           const endIndex = Math.min(data.length, end);
           return data.slice(start, endIndex);
         };
-        const reindex = data => data.map((p, idx) => ({ x: idx, y: p.y }));
+        const reindex = alignSteps ? (data => data) : (data => data.map((p, idx) => ({ x: idx, y: p.y })));
         Object.keys(metricsData).forEach(k => {
           metricsData[k] = reindex(applyRange(metricsData[k]));
         });
@@ -155,7 +193,7 @@ export default function ChartContainer({
 
       return { ...file, metricsData };
     });
-  }, [files, metrics]);
+  }, [files, metrics, alignSteps, stepKeyword]);
 
   useEffect(() => {
     const maxStep = parsedData.reduce((m, f) => {
@@ -166,15 +204,34 @@ export default function ChartContainer({
   }, [parsedData, onMaxStepChange]);
 
   useEffect(() => {
-    const minSteps = getMinSteps(parsedData);
-    if (minSteps > 0) {
-      onXRangeChange(prev => {
-        const next = { min: 0, max: minSteps - 1 };
-        if (prev.min === next.min && prev.max === next.max) return prev;
-        return next;
-      });
+    if (alignSteps) {
+      const enabled = parsedData.filter(f => f.enabled !== false);
+      if (enabled.length > 0) {
+        let minStart = Infinity;
+        let maxEnd = 0;
+        enabled.forEach(f => {
+          Object.values(f.metricsData).forEach(d => {
+            if (d.length > 0) {
+              minStart = Math.min(minStart, d[0].x);
+              maxEnd = Math.max(maxEnd, d[d.length - 1].x);
+            }
+          });
+        });
+        if (minStart !== Infinity) {
+          onXRangeChange({ min: minStart, max: maxEnd });
+        }
+      }
+    } else {
+      const minSteps = getMinSteps(parsedData);
+      if (minSteps > 0) {
+        onXRangeChange(prev => {
+          const next = { min: 0, max: minSteps - 1 };
+          if (prev.min === next.min && prev.max === next.max) return prev;
+          return next;
+        });
+      }
     }
-  }, [parsedData, onXRangeChange]);
+  }, [parsedData, onXRangeChange, alignSteps]);
 
   const colors = ['#ef4444', '#3b82f6', '#10b981', '#f59e0b', '#8b5cf6', '#f97316'];
   const createChartData = dataArray => ({
diff --git a/src/components/__tests__/ChartContainer.test.jsx b/src/components/__tests__/ChartContainer.test.jsx
index 62a9315..c880e61 100644
--- a/src/components/__tests__/ChartContainer.test.jsx
+++ b/src/components/__tests__/ChartContainer.test.jsx
@@ -165,4 +165,35 @@ describe('ChartContainer', () => {
     opts.plugins.zoom.pan.onPanComplete({ chart: { scales: { x: { min: 0, max: 10 } } } });
     opts.plugins.zoom.zoom.onZoomComplete({ chart: { scales: { x: { min: 2, max: 4 } } } });
   });
+
+  it('aligns data points using step keyword when enabled', () => {
+    const onXRangeChange = vi.fn();
+    const onMaxStepChange = vi.fn();
+    const files = [
+      { name: 'a.log', enabled: true, content: 'step:1 loss: 1\nstep:2 loss: 2' },
+      { name: 'b.log', enabled: true, content: 'step:5 loss: 5\nstep:6 loss: 6' }
+    ];
+    const prevCount = __lineProps.length;
+    render(
+      <ChartContainer
+        files={files}
+        metrics={[{ name: 'loss', keyword: 'loss', mode: 'keyword' }]}
+        compareMode="normal"
+        alignSteps
+        stepKeyword="step:"
+        onXRangeChange={onXRangeChange}
+        onMaxStepChange={onMaxStepChange}
+      />
+    );
+
+    const props = __lineProps[prevCount];
+    const dsA = props.data.datasets[0].data;
+    const dsB = props.data.datasets[1].data;
+    expect(dsA[0].x).toBe(1);
+    expect(dsA[1].x).toBe(2);
+    expect(dsB[0].x).toBe(5);
+    expect(dsB[1].x).toBe(6);
+    expect(onXRangeChange).toHaveBeenCalledWith({ min: 1, max: 6 });
+    expect(onMaxStepChange).toHaveBeenCalledWith(6);
+  });
 });