Atomizer/atomizer-dashboard/frontend/src/components/ConvergencePlot.tsx

/**
 * Convergence Plot
 * Shows optimization progress over time with running best and improvement rate
 * Features log scale toggle to better visualize early improvements
 */

import { useMemo, useState } from 'react';
import {
  Line,
  XAxis,
  YAxis,
  Tooltip,
  ResponsiveContainer,
  ReferenceLine,
  Area,
  ComposedChart,
  Legend
} from 'recharts';

interface Trial {
  trial_number: number;
  values: number[];
  state?: string;
  constraint_satisfied?: boolean;
  user_attrs?: Record<string, any>;
}

// Penalty threshold - objectives above this are considered failed/penalty trials
const PENALTY_THRESHOLD = 100000;

interface ConvergencePlotProps {
  trials: Trial[];
  objectiveIndex?: number;
  objectiveName?: string;
  direction?: 'minimize' | 'maximize';
  showLogScaleToggle?: boolean;
}

export function ConvergencePlot({
  trials,
  objectiveIndex = 0,
  objectiveName = 'Objective',
  direction = 'minimize',
  showLogScaleToggle = true
}: ConvergencePlotProps) {
  const [useLogScale, setUseLogScale] = useState(false);
  const [zoomToTail, setZoomToTail] = useState(false);

  const convergenceData = useMemo(() => {
    if (!trials || trials.length === 0) return [];

    // Sort by trial number, filtering out failed/penalty trials
    const sortedTrials = [...trials]
      .filter(t => {
        // Must have valid values
        if (!t.values || t.values.length <= objectiveIndex) return false;
        // Filter out failed state
        if (t.state === 'FAIL') return false;
        // Filter out penalty values (e.g., 1000000 = solver failure)
        const val = t.values[objectiveIndex];
        if (val >= PENALTY_THRESHOLD) return false;
        // Filter out constraint violations
        if (t.constraint_satisfied === false) return false;
        // Filter out pruned trials
        if (t.user_attrs?.pruned === true || t.user_attrs?.fail_reason) return false;
        return true;
      })
      .sort((a, b) => a.trial_number - b.trial_number);

    if (sortedTrials.length === 0) return [];

    let runningBest = direction === 'minimize' ? Infinity : -Infinity;
    const data = sortedTrials.map((trial, idx) => {
      const value = trial.values[objectiveIndex];

      // Update running best
      if (direction === 'minimize') {
        runningBest = Math.min(runningBest, value);
      } else {
        runningBest = Math.max(runningBest, value);
      }

      // Calculate improvement from first trial
      const firstValue = sortedTrials[0].values[objectiveIndex];
      const improvement = direction === 'minimize'
        ? ((firstValue - runningBest) / firstValue) * 100
        : ((runningBest - firstValue) / firstValue) * 100;

      return {
        trial: trial.trial_number,
        value,
        best: runningBest,
        improvement: Math.max(0, improvement),
        index: idx + 1
      };
    });

    return data;
  }, [trials, objectiveIndex, direction]);

  // Calculate statistics
  const stats = useMemo(() => {
    if (convergenceData.length === 0) return null;

    const values = convergenceData.map(d => d.value);
    const best = convergenceData[convergenceData.length - 1]?.best ?? 0;
    const first = convergenceData[0]?.value ?? 0;
    const improvement = direction === 'minimize'
      ? ((first - best) / first) * 100
      : ((best - first) / first) * 100;

    // Calculate when we found 90% of the improvement
    const targetImprovement = 0.9 * improvement;
    let trialsTo90 = convergenceData.length;
    for (let i = 0; i < convergenceData.length; i++) {
      if (convergenceData[i].improvement >= targetImprovement) {
        trialsTo90 = i + 1;
        break;
      }
    }

    return {
      best,
      first,
      improvement: Math.max(0, improvement),
      totalTrials: convergenceData.length,
      trialsTo90,
      mean: values.reduce((a, b) => a + b, 0) / values.length,
      std: Math.sqrt(values.map(v => Math.pow(v - values.reduce((a, b) => a + b, 0) / values.length, 2)).reduce((a, b) => a + b, 0) / values.length)
    };
  }, [convergenceData, direction]);

  // Transform data for log scale display
  const displayData = useMemo(() => {
    if (!useLogScale) {
      return zoomToTail && convergenceData.length > 20
        ? convergenceData.slice(-Math.ceil(convergenceData.length * 0.2))
        : convergenceData;
    }

    // For log scale, we need to handle values <= 0
    // Shift all values to be positive if needed
    const minVal = Math.min(...convergenceData.map(d => d.value), ...convergenceData.map(d => d.best));
    const offset = minVal <= 0 ? Math.abs(minVal) + 1 : 0;

    const transformed = convergenceData.map(d => ({
      ...d,
      value: Math.log10(d.value + offset),
      best: Math.log10(d.best + offset),
      originalValue: d.value,
      originalBest: d.best
    }));

    return zoomToTail && transformed.length > 20
      ? transformed.slice(-Math.ceil(transformed.length * 0.2))
      : transformed;
  }, [convergenceData, useLogScale, zoomToTail]);

  if (convergenceData.length === 0) {
    return (
      <div className="bg-dark-700 rounded-lg p-6 border border-dark-500 shadow-sm">
        <h3 className="text-lg font-semibold mb-4 text-dark-100">Convergence Plot</h3>
        <div className="h-64 flex items-center justify-center text-dark-400">
          No trial data available
        </div>
      </div>
    );
  }

  return (
    <div className="bg-dark-700 rounded-lg p-6 border border-dark-500 shadow-sm">
      <div className="flex items-start justify-between mb-4">
        <div>
          <h3 className="text-lg font-semibold text-dark-100">Convergence Plot</h3>
          <p className="text-sm text-dark-400 mt-1">
            {objectiveName} over {convergenceData.length} trials
            {useLogScale && <span className="ml-2 text-primary-400">(log₁₀ scale)</span>}
            {zoomToTail && <span className="ml-2 text-yellow-400">(last 20%)</span>}
          </p>
        </div>

        <div className="flex items-center gap-3">
          {/* Scale toggle buttons */}
          {showLogScaleToggle && (
            <div className="flex gap-1">
              <button
                onClick={() => setUseLogScale(!useLogScale)}
                className={`px-2 py-1 text-xs rounded transition-colors ${
                  useLogScale
                    ? 'bg-primary-600 text-white'
                    : 'bg-dark-600 text-dark-300 hover:bg-dark-500'
                }`}
                title="Toggle logarithmic scale - better for viewing early improvements"
              >
                Log
              </button>
              <button
                onClick={() => setZoomToTail(!zoomToTail)}
                className={`px-2 py-1 text-xs rounded transition-colors ${
                  zoomToTail
                    ? 'bg-yellow-600 text-white'
                    : 'bg-dark-600 text-dark-300 hover:bg-dark-500'
                }`}
                title="Zoom to last 20% of trials"
              >
                Tail
              </button>
            </div>
          )}

          {stats && (
            <div className="flex gap-4 text-sm border-l border-dark-500 pl-3">
              <div className="text-right">
                <div className="text-dark-400 text-xs">Best</div>
                <div className="font-semibold text-green-400">{stats.best.toFixed(4)}</div>
              </div>
              <div className="text-right">
                <div className="text-dark-400 text-xs">Improvement</div>
                <div className="font-semibold text-blue-400">{stats.improvement.toFixed(1)}%</div>
              </div>
              <div className="text-right">
                <div className="text-dark-400 text-xs">90% at</div>
                <div className="font-semibold text-purple-400">#{stats.trialsTo90}</div>
              </div>
            </div>
          )}
        </div>
      </div>

      <div className="h-72">
        <ResponsiveContainer width="100%" height="100%">
          <ComposedChart data={displayData} margin={{ top: 10, right: 30, left: 10, bottom: 10 }}>
            <defs>
              <linearGradient id="colorValue" x1="0" y1="0" x2="0" y2="1">
                <stop offset="5%" stopColor="#60a5fa" stopOpacity={0.3} />
                <stop offset="95%" stopColor="#60a5fa" stopOpacity={0} />
              </linearGradient>
            </defs>
            <XAxis
              dataKey="trial"
              tick={{ fill: '#94a3b8', fontSize: 11 }}
              stroke="#334155"
              label={{ value: 'Trial', position: 'bottom', offset: -5, fill: '#94a3b8', fontSize: 12 }}
            />
            <YAxis
              tick={{ fill: '#94a3b8', fontSize: 11 }}
              stroke="#334155"
              tickFormatter={(v) => useLogScale ? `10^${v.toFixed(1)}` : v.toFixed(2)}
              label={{
                value: useLogScale ? `log₁₀(${objectiveName})` : objectiveName,
                angle: -90,
                position: 'insideLeft',
                fill: '#94a3b8',
                fontSize: 12
              }}
            />
            <Tooltip
              contentStyle={{
                backgroundColor: '#1e293b',
                border: '1px solid #334155',
                borderRadius: '8px',
                boxShadow: '0 4px 6px rgba(0,0,0,0.3)'
              }}
              labelStyle={{ color: '#e2e8f0' }}
              formatter={(value: number, name: string, props: any) => {
                const label = name === 'value' ? 'Trial Value' :
                              name === 'best' ? 'Running Best' : name;
                // Show original values in tooltip when using log scale
                if (useLogScale) {
                  const original = name === 'value'
                    ? props.payload.originalValue
                    : props.payload.originalBest;
                  return [original?.toFixed(4) ?? value.toFixed(4), label];
                }
                return [value.toFixed(4), label];
              }}
              labelFormatter={(label) => `Trial #${label}`}
            />
            <Legend
              verticalAlign="top"
              height={36}
              wrapperStyle={{ color: '#94a3b8' }}
              formatter={(value) => value === 'value' ? 'Trial Value' : 'Running Best'}
            />

            {/* Area under trial values */}
            <Area
              type="monotone"
              dataKey="value"
              stroke="#60a5fa"
              fill="url(#colorValue)"
              strokeWidth={0}
            />

            {/* Trial values as scatter points */}
            <Line
              type="monotone"
              dataKey="value"
              stroke="#60a5fa"
              strokeWidth={1}
              dot={{ fill: '#60a5fa', r: 3 }}
              activeDot={{ fill: '#93c5fd', r: 5 }}
            />

            {/* Running best line */}
            <Line
              type="stepAfter"
              dataKey="best"
              stroke="#10b981"
              strokeWidth={2.5}
              dot={false}
            />

            {/* Reference line at best value */}
            {stats && (
              <ReferenceLine
                y={stats.best}
                stroke="#10b981"
                strokeDasharray="5 5"
                strokeOpacity={0.5}
              />
            )}
          </ComposedChart>
        </ResponsiveContainer>
      </div>

      {/* Summary stats bar */}
      {stats && (
        <div className="grid grid-cols-4 gap-4 mt-4 pt-4 border-t border-dark-500 text-center text-sm">
          <div>
            <div className="text-dark-400">First Value</div>
            <div className="font-medium text-dark-100">{stats.first.toFixed(4)}</div>
          </div>
          <div>
            <div className="text-dark-400">Mean</div>
            <div className="font-medium text-dark-100">{stats.mean.toFixed(4)}</div>
          </div>
          <div>
            <div className="text-dark-400">Std Dev</div>
            <div className="font-medium text-dark-100">{stats.std.toFixed(4)}</div>
          </div>
          <div>
            <div className="text-dark-400">Total Trials</div>
            <div className="font-medium text-dark-100">{stats.totalTrials}</div>
          </div>
        </div>
      )}
    </div>
  );
}