Atomizer/atomizer-dashboard/frontend/src/components/plotly/PlotlyParetoPlot.tsx

/**
 * PlotlyParetoPlot - Interactive Pareto front visualization using Plotly
 *
 * Features:
 * - 2D scatter with Pareto front highlighted
 * - 3D scatter for 3-objective problems
 * - Hover tooltips with trial details
 * - Pareto front connection line
 * - FEA vs NN differentiation
 * - Constraint satisfaction highlighting
 * - Dark mode styling
 * - Zoom, pan, and export
 */

import { useMemo, useState } from 'react';
import Plot from 'react-plotly.js';

interface Trial {
  trial_number: number;
  values: number[];
  params: Record<string, number>;
  user_attrs?: Record<string, any>;
  source?: 'FEA' | 'NN' | 'V10_FEA';
  constraint_satisfied?: boolean;
}

interface Objective {
  name: string;
  direction?: 'minimize' | 'maximize';
  unit?: string;
}

interface PlotlyParetoPlotProps {
  trials: Trial[];
  paretoFront: Trial[];
  objectives: Objective[];
  height?: number;
  showParetoLine?: boolean;
  showInfeasible?: boolean;
}

export function PlotlyParetoPlot({
  trials,
  paretoFront,
  objectives,
  height = 500,
  showParetoLine = true,
  showInfeasible = true
}: PlotlyParetoPlotProps) {
  const [viewMode, setViewMode] = useState<'2d' | '3d'>(objectives.length >= 3 ? '3d' : '2d');
  const [selectedObjectives, setSelectedObjectives] = useState<[number, number, number]>([0, 1, 2]);

  const paretoSet = useMemo(() => new Set(paretoFront.map(t => t.trial_number)), [paretoFront]);

  // Separate trials by source, Pareto status, and constraint satisfaction
  const { feaTrials, nnTrials, paretoTrials, infeasibleTrials, stats } = useMemo(() => {
    const fea: Trial[] = [];
    const nn: Trial[] = [];
    const pareto: Trial[] = [];
    const infeasible: Trial[] = [];

    trials.forEach(t => {
      const source = t.source || t.user_attrs?.source || 'FEA';
      const isFeasible = t.constraint_satisfied !== false && t.user_attrs?.constraint_satisfied !== false;

      if (!isFeasible && showInfeasible) {
        infeasible.push(t);
      } else if (paretoSet.has(t.trial_number)) {
        pareto.push(t);
      } else if (source === 'NN') {
        nn.push(t);
      } else {
        fea.push(t);
      }
    });

    // Calculate statistics
    const stats = {
      totalTrials: trials.length,
      paretoCount: pareto.length,
      feaCount: fea.length + pareto.filter(t => (t.source || 'FEA') !== 'NN').length,
      nnCount: nn.length + pareto.filter(t => t.source === 'NN').length,
      infeasibleCount: infeasible.length,
      hypervolume: 0 // Could calculate if needed
    };

    return { feaTrials: fea, nnTrials: nn, paretoTrials: pareto, infeasibleTrials: infeasible, stats };
  }, [trials, paretoSet, showInfeasible]);

  // Helper to get objective value
  const getObjValue = (trial: Trial, idx: number): number => {
    if (trial.values && trial.values[idx] !== undefined) {
      return trial.values[idx];
    }
    const objName = objectives[idx]?.name;
    return trial.user_attrs?.[objName] ?? 0;
  };

  // Build hover text
  const buildHoverText = (trial: Trial): string => {
    const lines = [`Trial #${trial.trial_number}`];
    objectives.forEach((obj, i) => {
      const val = getObjValue(trial, i);
      lines.push(`${obj.name}: ${val.toFixed(4)}${obj.unit ? ` ${obj.unit}` : ''}`);
    });
    const source = trial.source || trial.user_attrs?.source || 'FEA';
    lines.push(`Source: ${source}`);
    return lines.join('<br>');
  };

  // Create trace data
  const createTrace = (
    trialList: Trial[],
    name: string,
    color: string,
    symbol: string,
    size: number,
    opacity: number
  ) => {
    const [i, j, k] = selectedObjectives;

    if (viewMode === '3d' && objectives.length >= 3) {
      return {
        type: 'scatter3d' as const,
        mode: 'markers' as const,
        name,
        x: trialList.map(t => getObjValue(t, i)),
        y: trialList.map(t => getObjValue(t, j)),
        z: trialList.map(t => getObjValue(t, k)),
        text: trialList.map(buildHoverText),
        hoverinfo: 'text' as const,
        marker: {
          color,
          size,
          symbol,
          opacity,
          line: { color: '#fff', width: 1 }
        }
      };
    } else {
      return {
        type: 'scatter' as const,
        mode: 'markers' as const,
        name,
        x: trialList.map(t => getObjValue(t, i)),
        y: trialList.map(t => getObjValue(t, j)),
        text: trialList.map(buildHoverText),
        hoverinfo: 'text' as const,
        marker: {
          color,
          size,
          symbol,
          opacity,
          line: { color: '#fff', width: 1 }
        }
      };
    }
  };

  // Sort Pareto trials by first objective for line connection
  const sortedParetoTrials = useMemo(() => {
    const [i] = selectedObjectives;
    return [...paretoTrials].sort((a, b) => getObjValue(a, i) - getObjValue(b, i));
  }, [paretoTrials, selectedObjectives]);

  // Create Pareto front line trace (2D only)
  const createParetoLine = () => {
    if (!showParetoLine || viewMode === '3d' || sortedParetoTrials.length < 2) return null;
    const [i, j] = selectedObjectives;
    return {
      type: 'scatter' as const,
      mode: 'lines' as const,
      name: 'Pareto Front',
      x: sortedParetoTrials.map(t => getObjValue(t, i)),
      y: sortedParetoTrials.map(t => getObjValue(t, j)),
      line: {
        color: '#10B981',
        width: 2,
        dash: 'dot'
      },
      hoverinfo: 'skip' as const,
      showlegend: false
    };
  };

  const traces = [
    // Infeasible trials (background, red X)
    ...(showInfeasible && infeasibleTrials.length > 0 ? [
      createTrace(infeasibleTrials, `Infeasible (${infeasibleTrials.length})`, '#EF4444', 'x', 7, 0.4)
    ] : []),
    // FEA trials (blue circles)
    createTrace(feaTrials, `FEA (${feaTrials.length})`, '#3B82F6', 'circle', 8, 0.6),
    // NN trials (purple diamonds)
    createTrace(nnTrials, `NN (${nnTrials.length})`, '#A855F7', 'diamond', 8, 0.5),
    // Pareto front line (2D only)
    createParetoLine(),
    // Pareto front points (highlighted)
    createTrace(sortedParetoTrials, `Pareto (${sortedParetoTrials.length})`, '#10B981', 'star', 14, 1.0)
  ].filter(trace => trace && (trace.x as number[]).length > 0);

  const [i, j, k] = selectedObjectives;

  // Dark mode color scheme
  const colors = {
    text: '#E5E7EB',
    textMuted: '#9CA3AF',
    grid: 'rgba(255,255,255,0.1)',
    zeroline: 'rgba(255,255,255,0.2)',
    legendBg: 'rgba(30,30,30,0.9)',
    legendBorder: 'rgba(255,255,255,0.1)'
  };

  const layout: any = viewMode === '3d' && objectives.length >= 3
    ? {
        height,
        margin: { l: 50, r: 50, t: 30, b: 50 },
        paper_bgcolor: 'transparent',
        plot_bgcolor: 'transparent',
        scene: {
          xaxis: {
            title: { text: objectives[i]?.name || 'Objective 1', font: { color: colors.text } },
            gridcolor: colors.grid,
            zerolinecolor: colors.zeroline,
            tickfont: { color: colors.textMuted }
          },
          yaxis: {
            title: { text: objectives[j]?.name || 'Objective 2', font: { color: colors.text } },
            gridcolor: colors.grid,
            zerolinecolor: colors.zeroline,
            tickfont: { color: colors.textMuted }
          },
          zaxis: {
            title: { text: objectives[k]?.name || 'Objective 3', font: { color: colors.text } },
            gridcolor: colors.grid,
            zerolinecolor: colors.zeroline,
            tickfont: { color: colors.textMuted }
          },
          bgcolor: 'transparent'
        },
        legend: {
          x: 1,
          y: 1,
          font: { color: colors.text },
          bgcolor: colors.legendBg,
          bordercolor: colors.legendBorder,
          borderwidth: 1
        },
        font: { family: 'Inter, system-ui, sans-serif', color: colors.text }
      }
    : {
        height,
        margin: { l: 60, r: 30, t: 30, b: 60 },
        paper_bgcolor: 'transparent',
        plot_bgcolor: 'transparent',
        xaxis: {
          title: { text: objectives[i]?.name || 'Objective 1', font: { color: colors.text } },
          gridcolor: colors.grid,
          zerolinecolor: colors.zeroline,
          tickfont: { color: colors.textMuted }
        },
        yaxis: {
          title: { text: objectives[j]?.name || 'Objective 2', font: { color: colors.text } },
          gridcolor: colors.grid,
          zerolinecolor: colors.zeroline,
          tickfont: { color: colors.textMuted }
        },
        legend: {
          x: 1,
          y: 1,
          xanchor: 'right',
          font: { color: colors.text },
          bgcolor: colors.legendBg,
          bordercolor: colors.legendBorder,
          borderwidth: 1
        },
        font: { family: 'Inter, system-ui, sans-serif', color: colors.text },
        hovermode: 'closest' as const
      };

  if (!trials.length) {
    return (
      <div className="flex items-center justify-center h-64 text-dark-400">
        No trial data available
      </div>
    );
  }

  return (
    <div className="w-full">
      {/* Stats Bar */}
      <div className="flex gap-4 mb-4 text-sm">
        <div className="flex items-center gap-2 px-3 py-1.5 bg-dark-700 rounded-lg">
          <div className="w-3 h-3 bg-green-500 rounded-full" />
          <span className="text-dark-300">Pareto:</span>
          <span className="text-green-400 font-medium">{stats.paretoCount}</span>
        </div>
        <div className="flex items-center gap-2 px-3 py-1.5 bg-dark-700 rounded-lg">
          <div className="w-3 h-3 bg-blue-500 rounded-full" />
          <span className="text-dark-300">FEA:</span>
          <span className="text-blue-400 font-medium">{stats.feaCount}</span>
        </div>
        <div className="flex items-center gap-2 px-3 py-1.5 bg-dark-700 rounded-lg">
          <div className="w-3 h-3 bg-purple-500 rounded-full" />
          <span className="text-dark-300">NN:</span>
          <span className="text-purple-400 font-medium">{stats.nnCount}</span>
        </div>
        {stats.infeasibleCount > 0 && (
          <div className="flex items-center gap-2 px-3 py-1.5 bg-dark-700 rounded-lg">
            <div className="w-3 h-3 bg-red-500 rounded-full" />
            <span className="text-dark-300">Infeasible:</span>
            <span className="text-red-400 font-medium">{stats.infeasibleCount}</span>
          </div>
        )}
      </div>

      {/* Controls */}
      <div className="flex gap-4 items-center justify-between mb-3">
        <div className="flex gap-2 items-center">
          {objectives.length >= 3 && (
            <div className="flex rounded-lg overflow-hidden border border-dark-600">
              <button
                onClick={() => setViewMode('2d')}
                className={`px-3 py-1.5 text-sm font-medium transition-colors ${
                  viewMode === '2d'
                    ? 'bg-primary-600 text-white'
                    : 'bg-dark-700 text-dark-300 hover:bg-dark-600 hover:text-white'
                }`}
              >
                2D
              </button>
              <button
                onClick={() => setViewMode('3d')}
                className={`px-3 py-1.5 text-sm font-medium transition-colors ${
                  viewMode === '3d'
                    ? 'bg-primary-600 text-white'
                    : 'bg-dark-700 text-dark-300 hover:bg-dark-600 hover:text-white'
                }`}
              >
                3D
              </button>
            </div>
          )}
        </div>

        {/* Objective selectors */}
        <div className="flex gap-2 items-center text-sm">
          <label className="text-dark-400">X:</label>
          <select
            value={selectedObjectives[0]}
            onChange={(e) => setSelectedObjectives([parseInt(e.target.value), selectedObjectives[1], selectedObjectives[2]])}
            className="px-2 py-1.5 bg-dark-700 border border-dark-600 rounded text-white text-sm"
          >
            {objectives.map((obj, idx) => (
              <option key={idx} value={idx}>{obj.name}</option>
            ))}
          </select>

          <label className="text-dark-400 ml-2">Y:</label>
          <select
            value={selectedObjectives[1]}
            onChange={(e) => setSelectedObjectives([selectedObjectives[0], parseInt(e.target.value), selectedObjectives[2]])}
            className="px-2 py-1.5 bg-dark-700 border border-dark-600 rounded text-white text-sm"
          >
            {objectives.map((obj, idx) => (
              <option key={idx} value={idx}>{obj.name}</option>
            ))}
          </select>

          {viewMode === '3d' && objectives.length >= 3 && (
            <>
              <label className="text-dark-400 ml-2">Z:</label>
              <select
                value={selectedObjectives[2]}
                onChange={(e) => setSelectedObjectives([selectedObjectives[0], selectedObjectives[1], parseInt(e.target.value)])}
                className="px-2 py-1.5 bg-dark-700 border border-dark-600 rounded text-white text-sm"
              >
                {objectives.map((obj, idx) => (
                  <option key={idx} value={idx}>{obj.name}</option>
                ))}
              </select>
            </>
          )}
        </div>
      </div>

      <Plot
        data={traces as any}
        layout={layout}
        config={{
          displayModeBar: true,
          displaylogo: false,
          modeBarButtonsToRemove: ['lasso2d', 'select2d'],
          toImageButtonOptions: {
            format: 'png',
            filename: 'pareto_front',
            height: 800,
            width: 1200,
            scale: 2
          }
        }}
        style={{ width: '100%' }}
      />

      {/* Pareto Front Table for 2D view */}
      {viewMode === '2d' && sortedParetoTrials.length > 0 && (
        <div className="mt-4 max-h-48 overflow-auto">
          <table className="w-full text-sm">
            <thead className="sticky top-0 bg-dark-800">
              <tr className="border-b border-dark-600">
                <th className="text-left py-2 px-3 text-dark-400 font-medium">Trial</th>
                <th className="text-left py-2 px-3 text-dark-400 font-medium">{objectives[i]?.name || 'Obj 1'}</th>
                <th className="text-left py-2 px-3 text-dark-400 font-medium">{objectives[j]?.name || 'Obj 2'}</th>
                <th className="text-left py-2 px-3 text-dark-400 font-medium">Source</th>
              </tr>
            </thead>
            <tbody>
              {sortedParetoTrials.slice(0, 10).map(trial => (
                <tr key={trial.trial_number} className="border-b border-dark-700 hover:bg-dark-750">
                  <td className="py-2 px-3 font-mono text-white">#{trial.trial_number}</td>
                  <td className="py-2 px-3 font-mono text-green-400">
                    {getObjValue(trial, i).toExponential(4)}
                  </td>
                  <td className="py-2 px-3 font-mono text-green-400">
                    {getObjValue(trial, j).toExponential(4)}
                  </td>
                  <td className="py-2 px-3">
                    <span className={`px-2 py-0.5 rounded text-xs ${
                      (trial.source || trial.user_attrs?.source) === 'NN'
                        ? 'bg-purple-500/20 text-purple-400'
                        : 'bg-blue-500/20 text-blue-400'
                    }`}>
                      {trial.source || trial.user_attrs?.source || 'FEA'}
                    </span>
                  </td>
                </tr>
              ))}
            </tbody>
          </table>
          {sortedParetoTrials.length > 10 && (
            <div className="text-center py-2 text-dark-500 text-xs">
              Showing 10 of {sortedParetoTrials.length} Pareto-optimal solutions
            </div>
          )}
        </div>
      )}
    </div>
  );
}