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

/**
 * Parallel Coordinates Plot - Enhanced Multi-Objective Visualization
 * Shows design variables → objectives → constraints in proper research format
 * Light theme with high visibility
 */

import { useState, useMemo } from 'react';

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

interface Objective {
  name: string;
  type?: 'minimize' | 'maximize';
  direction?: 'minimize' | 'maximize';  // Alternative name used in some configs
  unit?: string;
}

interface DesignVariable {
  name: string;
  parameter?: string;  // Optional: the actual parameter name if different from name
  unit?: string;
  min?: number;
  max?: number;
}

interface Constraint {
  name: string;
  threshold: number;
  type: 'less_than' | 'greater_than';
  unit?: string;
}

interface ParallelCoordinatesPlotProps {
  paretoData: ParetoTrial[];
  objectives: Objective[];
  designVariables: DesignVariable[];
  constraints?: Constraint[];
  paretoFront?: ParetoTrial[];
  maxTrials?: number; // Limit displayed trials for performance (default: 1000)
}

export function ParallelCoordinatesPlot({
  paretoData,
  objectives,
  designVariables,
  constraints: _constraints = [],
  paretoFront = [],
  maxTrials = 500
}: ParallelCoordinatesPlotProps) {
  const [hoveredTrial, setHoveredTrial] = useState<number | null>(null);
  const [selectedTrials, setSelectedTrials] = useState<Set<number>>(new Set());
  const [showNN, setShowNN] = useState<boolean>(true);  // Toggle for NN trials visibility

  // Create set of Pareto front trial numbers for easy lookup
  const paretoTrialNumbers = new Set(paretoFront.map(t => t.trial_number));

  // Check if we have any NN trials (to show/hide toggle)
  const hasNNTrials = useMemo(() => {
    return paretoData.some(t => t.source === 'NN' || t.user_attrs?.source === 'NN');
  }, [paretoData]);

  // Filter to best N trials for performance when dataset is large
  // Also filter by source type based on toggle
  const filteredParetoData = useMemo(() => {
    // First filter by source if NN is hidden
    let data = paretoData;
    if (!showNN) {
      data = data.filter(t => {
        const source = t.source || t.user_attrs?.source || 'FEA';
        return source !== 'NN';
      });
    }

    if (data.length <= maxTrials) return data;

    // Sort by first objective value and take best N
    const sorted = [...data].sort((a, b) => {
      const aVal = a.values[0] ?? Infinity;
      const bVal = b.values[0] ?? Infinity;
      return aVal - bVal;
    });
    return sorted.slice(0, maxTrials);
  }, [paretoData, maxTrials, showNN]);

  const totalTrialsCount = paretoData.length;
  const displayedTrialsCount = filteredParetoData.length;
  const isFiltered = totalTrialsCount > maxTrials;

  // Safety checks
  if (!paretoData || paretoData.length === 0) {
    return (
      <div className="bg-white rounded-lg p-6 border border-gray-300 shadow-sm">
        <h3 className="text-lg font-semibold mb-4 text-gray-900">Parallel Coordinates Plot</h3>
        <div className="h-96 flex items-center justify-center text-gray-500">
          No Pareto front data available
        </div>
      </div>
    );
  }

  if (!objectives || objectives.length === 0 || !designVariables || designVariables.length === 0) {
    return (
      <div className="bg-white rounded-lg p-6 border border-gray-300 shadow-sm">
        <h3 className="text-lg font-semibold mb-4 text-gray-900">Parallel Coordinates Plot</h3>
        <div className="h-96 flex items-center justify-center text-gray-500">
          Missing objectives or design variables metadata
        </div>
      </div>
    );
  }

  // Structure axes: Design Variables → Objectives → Constraints
  const axes: Array<{
    name: string;
    label: string;
    type: 'design_var' | 'objective' | 'constraint';
    unit?: string;
  }> = [];

  // Add design variables
  designVariables.forEach(dv => {
    const paramName = dv.parameter || dv.name;  // Support both formats
    axes.push({
      name: paramName,
      label: dv.unit ? `${paramName}\n(${dv.unit})` : paramName,
      type: 'design_var',
      unit: dv.unit
    });
  });

  // Add objectives
  objectives.forEach((obj, i) => {
    axes.push({
      name: `objective_${i}`,
      label: obj.unit ? `${obj.name}\n(${obj.unit})` : obj.name,
      type: 'objective',
      unit: obj.unit
    });
  });

  // Add constraints (extract from user_attrs)
  const constraintNames = new Set<string>();
  filteredParetoData.forEach(trial => {
    if (trial.user_attrs) {
      // Common constraint metrics
      if (trial.user_attrs.max_stress !== undefined) constraintNames.add('max_stress');
      if (trial.user_attrs.max_displacement !== undefined) constraintNames.add('max_displacement');
      if (trial.user_attrs.frequency !== undefined && objectives.findIndex(obj => obj.name.toLowerCase().includes('freq')) === -1) {
        constraintNames.add('frequency');
      }
      if (trial.user_attrs.mass !== undefined && objectives.findIndex(obj => obj.name.toLowerCase().includes('mass')) === -1) {
        constraintNames.add('mass');
      }
    }
  });

  constraintNames.forEach(name => {
    const unit = name.includes('stress') ? 'MPa' :
                 name.includes('displacement') ? 'mm' :
                 name.includes('frequency') ? 'Hz' :
                 name.includes('mass') ? 'g' : '';
    axes.push({
      name: `constraint_${name}`,
      label: unit ? `${name}\n(${unit})` : name,
      type: 'constraint',
      unit
    });
  });

  // Extract values for each axis
  const trialData = filteredParetoData.map(trial => {
    const values: number[] = [];

    // Design variables - use .parameter field from metadata
    designVariables.forEach(dv => {
      const paramName = dv.parameter || dv.name;  // Support both formats
      values.push(trial.params[paramName] ?? 0);
    });

    // Objectives
    trial.values.forEach(val => {
      values.push(val);
    });

    // Constraints
    constraintNames.forEach(name => {
      values.push(trial.user_attrs?.[name] ?? 0);
    });

    return {
      trial_number: trial.trial_number,
      values,
      feasible: trial.constraint_satisfied !== false,
      objectiveValues: trial.values || [],
      source: trial.source || trial.user_attrs?.source || 'FEA'  // Default to FEA
    };
  });

  // Rank trials by their first objective (for multi-objective, this is just one metric)
  // For proper multi-objective ranking, we use Pareto dominance
  const rankedTrials = [...trialData].sort((a, b) => {
    // Primary: Pareto front members come first
    const aIsPareto = paretoTrialNumbers.has(a.trial_number);
    const bIsPareto = paretoTrialNumbers.has(b.trial_number);
    if (aIsPareto && !bIsPareto) return -1;
    if (!aIsPareto && bIsPareto) return 1;

    // Secondary: Sort by first objective value (minimize assumed)
    const aObj = a.objectiveValues[0] ?? Infinity;
    const bObj = b.objectiveValues[0] ?? Infinity;
    return aObj - bObj;
  });

  // Create ranking map: trial_number -> rank (0-indexed)
  const trialRanks = new Map<number, number>();
  rankedTrials.forEach((trial, index) => {
    trialRanks.set(trial.trial_number, index);
  });

  // Calculate min/max for normalization
  const ranges = axes.map((_, axisIdx) => {
    const values = trialData.map(d => d.values[axisIdx]);
    return {
      min: Math.min(...values),
      max: Math.max(...values)
    };
  });

  // Normalize to [0, 1]
  const normalize = (value: number, axisIdx: number): number => {
    const range = ranges[axisIdx];
    if (range.max === range.min) return 0.5;
    return (value - range.min) / (range.max - range.min);
  };

  // Chart dimensions
  const width = 1400;
  const height = 500;
  const margin = { top: 100, right: 50, bottom: 60, left: 50 };
  const plotWidth = width - margin.left - margin.right;
  const plotHeight = height - margin.top - margin.bottom;

  const axisSpacing = plotWidth / (axes.length - 1);

  // Toggle trial selection
  const toggleTrial = (trialNum: number) => {
    const newSelected = new Set(selectedTrials);
    if (newSelected.has(trialNum)) {
      newSelected.delete(trialNum);
    } else {
      newSelected.add(trialNum);
    }
    setSelectedTrials(newSelected);
  };

  // Color scheme - FEA trials: blue gradient, NN trials: orange gradient
  const totalTrials = trialData.length;

  const getLineColor = (trial: typeof trialData[0], isHovered: boolean, isSelected: boolean) => {
    if (isSelected) return '#FF6B00'; // Bright orange for selected
    if (isHovered) return '#2563EB'; // Blue for hover
    if (!trial.feasible) return '#DC2626'; // Red for infeasible

    const rank = trialRanks.get(trial.trial_number) ?? totalTrials - 1;
    const t = totalTrials > 1 ? rank / (totalTrials - 1) : 0;

    // FEA trials: dark blue to light blue gradient (validated, trustworthy)
    // NN trials: dark orange to light orange gradient (predictions, less reliable)
    const isNN = trial.source === 'NN';

    if (isNN) {
      // Orange gradient: rgb(194, 65, 12) to rgb(254, 215, 170)
      const r = Math.round(194 + t * (254 - 194));
      const g = Math.round(65 + t * (215 - 65));
      const b = Math.round(12 + t * (170 - 12));
      return `rgb(${r}, ${g}, ${b})`;
    } else {
      // Blue gradient: rgb(30, 64, 175) to rgb(191, 219, 254)
      const r = Math.round(30 + t * (191 - 30));
      const g = Math.round(64 + t * (219 - 64));
      const b = Math.round(175 + t * (254 - 175));
      return `rgb(${r}, ${g}, ${b})`;
    }
  };

  // Opacity gradient: best trials are more opaque, worst are more transparent
  const getLineOpacity = (trial: typeof trialData[0], isHighlighted: boolean, hasSelection: boolean, isSelected: boolean) => {
    if (hasSelection) {
      return isSelected ? 0.95 : 0.08;
    }
    if (isHighlighted) return 0.95;

    const rank = trialRanks.get(trial.trial_number) ?? totalTrials - 1;
    const t = totalTrials > 1 ? rank / (totalTrials - 1) : 0;

    // Best trials: 0.9 opacity, worst trials: 0.2 opacity
    return 0.9 - t * 0.7;
  };

  return (
    <div className="bg-white rounded-lg p-6 border border-gray-300 shadow-sm">
      <div className="flex items-center justify-between mb-4">
        <div>
          <h3 className="text-lg font-semibold text-gray-900">
            Parallel Coordinates Plot ({displayedTrialsCount} solutions{isFiltered ? ` of ${totalTrialsCount}` : ''})
          </h3>
          <p className="text-sm text-gray-600 mt-1">
            Design Variables → Objectives → Constraints
            {isFiltered && <span className="ml-2 text-amber-600">(showing best {maxTrials} for performance)</span>}
          </p>
        </div>
        <div className="flex items-center gap-3">
          {/* NN Toggle - only show if there are NN trials */}
          {hasNNTrials && (
            <button
              onClick={() => setShowNN(!showNN)}
              className={`text-xs px-3 py-1.5 rounded font-medium transition-colors flex items-center gap-1.5 ${
                showNN
                  ? 'bg-orange-100 text-orange-700 hover:bg-orange-200 border border-orange-300'
                  : 'bg-gray-100 text-gray-500 hover:bg-gray-200 border border-gray-300'
              }`}
              title={showNN ? 'Click to hide NN predictions' : 'Click to show NN predictions'}
            >
              <span className={`w-2 h-2 rounded-full ${showNN ? 'bg-orange-500' : 'bg-gray-400'}`} />
              {showNN ? 'NN Visible' : 'NN Hidden'}
            </button>
          )}

          {selectedTrials.size > 0 && (
            <button
              onClick={() => setSelectedTrials(new Set())}
              className="text-xs px-4 py-2 bg-gray-200 hover:bg-gray-300 rounded text-gray-800 font-medium transition-colors"
            >
              Clear Selection ({selectedTrials.size})
            </button>
          )}
        </div>
      </div>

      <div className="overflow-x-auto">
        <svg width={width} height={height} className="overflow-visible">
          <g transform={`translate(${margin.left}, ${margin.top})`}>
            {/* Draw axes */}
            {axes.map((axis, i) => {
              const x = i * axisSpacing;
              const bgColor = axis.type === 'design_var' ? '#EFF6FF' :
                             axis.type === 'objective' ? '#F0FDF4' : '#FEF3C7';

              return (
                <g key={axis.name} transform={`translate(${x}, 0)`}>
                  {/* Background highlight */}
                  <rect
                    x={-15}
                    y={-20}
                    width={30}
                    height={plotHeight + 40}
                    fill={bgColor}
                    opacity={0.3}
                  />

                  {/* Axis line */}
                  <line
                    y1={0}
                    y2={plotHeight}
                    stroke="#374151"
                    strokeWidth={2.5}
                  />

                  {/* Axis label */}
                  <text
                    y={-30}
                    textAnchor="middle"
                    fill="#111827"
                    fontSize={12}
                    fontWeight="600"
                    className="select-none"
                  >
                    {(axis.label || '').split('\n').map((line, idx) => (
                      <tspan key={idx} x={0} dy={idx === 0 ? 0 : 14}>{line}</tspan>
                    ))}
                  </text>

                  {/* Type badge */}
                  <text
                    y={-55}
                    textAnchor="middle"
                    fill="#6B7280"
                    fontSize={9}
                    fontWeight="500"
                    className="select-none"
                  >
                    {axis.type === 'design_var' ? 'DESIGN VAR' :
                     axis.type === 'objective' ? 'OBJECTIVE' : 'CONSTRAINT'}
                  </text>

                  {/* Min/max labels */}
                  <text
                    y={plotHeight + 20}
                    textAnchor="middle"
                    fill="#374151"
                    fontSize={10}
                    fontWeight="500"
                  >
                    {ranges[i].min.toFixed(2)}
                  </text>
                  <text
                    y={-10}
                    textAnchor="middle"
                    fill="#374151"
                    fontSize={10}
                    fontWeight="500"
                  >
                    {ranges[i].max.toFixed(2)}
                  </text>
                </g>
              );
            })}

            {/* Draw polylines for each trial - sorted so best trials render on top */}
            {[...trialData]
              .sort((a, b) => {
                // Render worst trials first (bottom), best trials last (top)
                const rankA = trialRanks.get(a.trial_number) ?? totalTrials;
                const rankB = trialRanks.get(b.trial_number) ?? totalTrials;
                return rankB - rankA; // Higher rank (worse) renders first
              })
              .map(trial => {
              const isHovered = hoveredTrial === trial.trial_number;
              const isSelected = selectedTrials.has(trial.trial_number);
              const isHighlighted = isHovered || isSelected;

              // Build path
              const pathData = axes.map((_, i) => {
                const x = i * axisSpacing;
                const normalizedY = normalize(trial.values[i], i);
                const y = plotHeight * (1 - normalizedY);
                return i === 0 ? `M ${x} ${y}` : `L ${x} ${y}`;
              }).join(' ');

              return (
                <path
                  key={trial.trial_number}
                  d={pathData}
                  fill="none"
                  stroke={getLineColor(trial, isHovered, isSelected)}
                  strokeWidth={isHighlighted ? 3 : 1.5}
                  opacity={getLineOpacity(trial, isHighlighted, selectedTrials.size > 0, isSelected)}
                  strokeLinecap="round"
                  strokeLinejoin="round"
                  className="transition-all duration-150 cursor-pointer"
                  style={{
                    filter: isHighlighted ? 'drop-shadow(0 2px 4px rgba(0,0,0,0.2))' : 'none'
                  }}
                  onMouseEnter={() => setHoveredTrial(trial.trial_number)}
                  onMouseLeave={() => setHoveredTrial(null)}
                  onClick={() => toggleTrial(trial.trial_number)}
                />
              );
            })}

            {/* Hover tooltip */}
            {hoveredTrial !== null && (
              <g transform={`translate(${plotWidth + 20}, 20)`}>
                <rect
                  x={0}
                  y={0}
                  width={140}
                  height={70}
                  fill="white"
                  stroke="#D1D5DB"
                  strokeWidth={2}
                  rx={6}
                  style={{ filter: 'drop-shadow(0 4px 6px rgba(0,0,0,0.1))' }}
                />
                <text x={12} y={24} fill="#111827" fontSize={13} fontWeight="700">
                  Trial #{hoveredTrial}
                </text>
                <text x={12} y={44} fill="#6B7280" fontSize={11}>
                  Click to select/deselect
                </text>
                <text x={12} y={60} fill="#374151" fontSize={10} fontWeight="600">
                  {selectedTrials.has(hoveredTrial) ? '✓ Selected' : '○ Not selected'}
                </text>
              </g>
            )}
          </g>
        </svg>
      </div>

      {/* Legend */}
      <div className="flex gap-6 justify-center mt-6 text-sm border-t border-gray-200 pt-4">
        <div className="flex items-center gap-2">
          <div className="w-16 h-2 rounded" style={{ background: 'linear-gradient(to right, #1E40AF, #60A5FA, #BFDBFE)' }} />
          <span className="text-gray-700 font-medium">FEA (validated)</span>
        </div>
        <div className="flex items-center gap-2">
          <div className="w-16 h-2 rounded" style={{ background: 'linear-gradient(to right, #C2410C, #FB923C, #FED7AA)' }} />
          <span className="text-gray-700 font-medium">NN (predicted)</span>
        </div>
        <div className="flex items-center gap-2">
          <div className="w-8 h-2 rounded" style={{ backgroundColor: '#DC2626' }} />
          <span className="text-gray-700 font-medium">Infeasible</span>
        </div>
        <div className="flex items-center gap-2">
          <div className="w-8 h-2 rounded" style={{ backgroundColor: '#FF6B00' }} />
          <span className="text-gray-700 font-medium">Selected</span>
        </div>
        <div className="flex items-center gap-2">
          <div className="w-8 h-2 rounded" style={{ backgroundColor: '#2563EB' }} />
          <span className="text-gray-700 font-medium">Hover</span>
        </div>
      </div>

      {/* Axis type legend */}
      <div className="flex gap-6 justify-center mt-3 text-xs">
        <div className="flex items-center gap-2">
          <div className="w-4 h-4 rounded" style={{ backgroundColor: '#EFF6FF' }} />
          <span className="text-gray-600">Design Variables</span>
        </div>
        <div className="flex items-center gap-2">
          <div className="w-4 h-4 rounded" style={{ backgroundColor: '#F0FDF4' }} />
          <span className="text-gray-600">Objectives</span>
        </div>
        <div className="flex items-center gap-2">
          <div className="w-4 h-4 rounded" style={{ backgroundColor: '#FEF3C7' }} />
          <span className="text-gray-600">Constraints</span>
        </div>
      </div>
    </div>
  );
}