import { useEffect, useMemo, useRef } from '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 }: { series: ScalarSeries[]; smoothing: number }) {
  const host = useRef<HTMLDivElement>(null);
  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 chart = new uPlot(
      {
        width: Math.max(320, element.clientWidth),
        height: 360,
        title: '训练与评估 Scalars',
        cursor: { drag: { x: true, y: true, setScale: true } },
        scales: { x: { time: false } },
        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,
    );
    const observer = new ResizeObserver((entries) => {
      const width = entries[0]?.contentRect.width;
      if (width) chart.setSize({ width: Math.max(320, Math.floor(width)), height: 360 });
    });
    observer.observe(element);
    return () => {
      observer.disconnect();
      chart.destroy();
    };
  }, [prepared, series]);

  if (series.length === 0)
    return (
      <div className="grid h-[360px] place-items-center rounded-lg border border-border bg-app text-xs text-text-tertiary">
        当前 trial 尚无 scalar 数据
      </div>
    );
  return <div ref={host} className="min-w-0 overflow-hidden rounded-lg bg-app p-2" />;
}