Atomizer/atomizer-dashboard/frontend/src/components/tracker/OptimizerStatePanel.tsx

import { useEffect, useState } from 'react';
import { Cpu, Layers, Target, Database, Brain, RefreshCw } from 'lucide-react';
import { apiClient } from '../../api/client';

interface OptimizerState {
  available: boolean;
  source?: string;
  phase?: string;
  phase_description?: string;
  phase_progress?: number;
  current_strategy?: string;
  sampler?: {
    name: string;
    description: string;
  };
  objectives?: Array<{
    name: string;
    direction: string;
    current_best?: number;
    unit?: string;
  }>;
  plan?: {
    total_phases: number;
    current_phase: number;
    phases: string[];
  };
  completed_trials?: number;
  total_trials?: number;
}

interface OptimizerStatePanelProps {
  studyId?: string;
  // Legacy props for backwards compatibility
  sampler?: string;
  nTrials?: number;
  completedTrials?: number;
  feaTrials?: number;
  nnTrials?: number;
  objectives?: Array<{ name: string; direction: string }>;
  isMultiObjective?: boolean;
  paretoSize?: number;
}

export function OptimizerStatePanel({
  studyId,
  sampler: legacySampler = 'TPESampler',
  nTrials: legacyNTrials = 100,
  completedTrials: legacyCompletedTrials = 0,
  feaTrials = 0,
  nnTrials = 0,
  objectives: legacyObjectives = [],
  isMultiObjective = false,
  paretoSize = 0
}: OptimizerStatePanelProps) {
  const [state, setState] = useState<OptimizerState | null>(null);
  const [loading, setLoading] = useState(false);

  // Fetch dynamic state if studyId is provided
  useEffect(() => {
    if (!studyId) return;

    const fetchState = async () => {
      setLoading(true);
      try {
        const result = await apiClient.getOptimizerState(studyId);
        setState(result);
      } catch {
        setState(null);
      } finally {
        setLoading(false);
      }
    };

    fetchState();
    const interval = setInterval(fetchState, 5000); // Refresh every 5 seconds
    return () => clearInterval(interval);
  }, [studyId]);

  // Use dynamic state if available, otherwise fall back to legacy props
  const sampler = state?.sampler?.name || legacySampler;
  const samplerDescription = state?.sampler?.description;
  const nTrials = state?.total_trials || legacyNTrials;
  const completedTrials = state?.completed_trials ?? legacyCompletedTrials;
  const objectives = state?.objectives?.map(o => ({ name: o.name, direction: o.direction })) || legacyObjectives;
  const phase = state?.phase || getPhase(completedTrials, nTrials);
  const phaseDescription = state?.phase_description;
  const phaseProgress = state?.phase_progress ?? getPhaseProgress(completedTrials, nTrials);
  const plan = state?.plan;

  // Determine optimizer phase based on progress (fallback)
  function getPhase(completed: number, total: number): string {
    if (completed === 0) return 'initializing';
    if (completed < 10) return 'exploration';
    if (completed < total * 0.5) return 'exploitation';
    if (completed < total * 0.9) return 'refinement';
    return 'convergence';
  }

  function getPhaseProgress(completed: number, total: number): number {
    if (completed === 0) return 0;
    if (completed < 10) return completed / 10;
    if (completed < total * 0.5) return (completed - 10) / (total * 0.5 - 10);
    if (completed < total * 0.9) return (completed - total * 0.5) / (total * 0.4);
    return (completed - total * 0.9) / (total * 0.1);
  }

  // Format sampler name for display
  const formatSampler = (s: string) => {
    const samplers: Record<string, string> = {
      'TPESampler': 'TPE (Bayesian)',
      'NSGAIISampler': 'NSGA-II',
      'NSGAIIISampler': 'NSGA-III',
      'CmaEsSampler': 'CMA-ES',
      'RandomSampler': 'Random',
      'GridSampler': 'Grid',
      'QMCSampler': 'Quasi-Monte Carlo'
    };
    return samplers[s] || s;
  };

  const phaseColors: Record<string, string> = {
    'initializing': 'text-dark-400',
    'exploration': 'text-primary-400',
    'exploitation': 'text-yellow-400',
    'refinement': 'text-blue-400',
    'convergence': 'text-green-400'
  };

  const phaseLabels: Record<string, string> = {
    'initializing': 'Initializing',
    'exploration': 'Exploration',
    'exploitation': 'Exploitation',
    'refinement': 'Refinement',
    'convergence': 'Convergence'
  };

  return (
    <div className="bg-dark-750 rounded-lg border border-dark-600 p-4">
      {/* Header */}
      <div className="flex items-center justify-between mb-4">
        <div className="flex items-center gap-2">
          <Cpu className="w-5 h-5 text-primary-400" />
          <span className="font-semibold text-white">Optimizer State</span>
        </div>
        {loading && <RefreshCw className="w-4 h-4 text-dark-400 animate-spin" />}
      </div>

      {/* Sampler with description */}
      <div className="bg-dark-700 rounded-lg p-3 mb-3">
        <div className="flex items-center gap-2 mb-1">
          <Target className="w-4 h-4 text-primary-400" />
          <span className="font-medium text-white">{formatSampler(sampler)}</span>
        </div>
        {samplerDescription && (
          <p className="text-xs text-dark-400 mt-1">{samplerDescription}</p>
        )}
      </div>

      {/* Phase progress */}
      <div className="bg-dark-700 rounded-lg p-3 mb-3">
        <div className="flex justify-between mb-2">
          <span className="text-dark-300 text-sm">
            Phase {plan ? `${plan.current_phase}/${plan.total_phases}` : ''}
          </span>
          <span className={`font-medium ${phaseColors[phase] || 'text-primary-400'}`}>
            {phaseLabels[phase] || phase}
          </span>
        </div>
        <div className="h-2 bg-dark-600 rounded-full overflow-hidden">
          <div
            className="h-full bg-primary-500 transition-all duration-500"
            style={{ width: `${Math.min(100, phaseProgress * 100)}%` }}
          />
        </div>
        {phaseDescription && (
          <p className="text-xs text-dark-400 mt-2">{phaseDescription}</p>
        )}
      </div>

      {/* FEA vs NN Trials (for hybrid optimizations) */}
      {(feaTrials > 0 || nnTrials > 0) && (
        <div className="mb-3">
          <div className="text-xs text-dark-400 uppercase mb-2">Trial Sources</div>
          <div className="flex gap-2">
            <div className="flex-1 bg-dark-700 rounded-lg p-2 text-center">
              <Database className="w-4 h-4 text-blue-400 mx-auto mb-1" />
              <div className="text-lg font-bold text-blue-400">{feaTrials}</div>
              <div className="text-xs text-dark-400">FEA</div>
            </div>
            <div className="flex-1 bg-dark-700 rounded-lg p-2 text-center">
              <Brain className="w-4 h-4 text-purple-400 mx-auto mb-1" />
              <div className="text-lg font-bold text-purple-400">{nnTrials}</div>
              <div className="text-xs text-dark-400">Neural Net</div>
            </div>
          </div>
          {nnTrials > 0 && (
            <div className="mt-2 text-xs text-dark-400 text-center">
              {((nnTrials / (feaTrials + nnTrials)) * 100).toFixed(0)}% acceleration from surrogate
            </div>
          )}
        </div>
      )}

      {/* Live objectives with current best */}
      {objectives.length > 0 && (
        <div className="mb-3">
          <div className="flex items-center gap-2 mb-2">
            <Layers className="w-4 h-4 text-dark-400" />
            <span className="text-xs text-dark-400 uppercase">
              {isMultiObjective || objectives.length > 1 ? 'Multi-Objective' : 'Single Objective'}
            </span>
          </div>
          <div className="space-y-1">
            {objectives.slice(0, 3).map((obj, idx) => {
              const stateObj = state?.objectives?.[idx];
              return (
                <div
                  key={idx}
                  className="flex items-center justify-between text-sm bg-dark-700 rounded px-2 py-1"
                >
                  <span className="text-dark-300 truncate" title={obj.name}>
                    {obj.name.length > 15 ? obj.name.slice(0, 13) + '...' : obj.name}
                  </span>
                  <div className="flex items-center gap-2">
                    {stateObj?.current_best !== undefined && stateObj.current_best !== null && (
                      <span className="font-mono text-primary-400 text-xs">
                        {stateObj.current_best.toFixed(2)}
                        {stateObj.unit && <span className="text-dark-500 ml-0.5">{stateObj.unit}</span>}
                      </span>
                    )}
                    <span className={`text-xs px-1.5 py-0.5 rounded ${
                      obj.direction === 'minimize' ? 'bg-green-900/50 text-green-400' : 'bg-blue-900/50 text-blue-400'
                    }`}>
                      {obj.direction === 'minimize' ? 'min' : 'max'}
                    </span>
                  </div>
                </div>
              );
            })}
            {objectives.length > 3 && (
              <div className="text-xs text-dark-500 text-center">
                +{objectives.length - 3} more
              </div>
            )}
          </div>
        </div>
      )}

      {/* Pareto Front Size (for multi-objective) */}
      {(isMultiObjective || objectives.length > 1) && paretoSize > 0 && (
        <div className="pt-3 border-t border-dark-600">
          <div className="flex items-center justify-between">
            <span className="text-xs text-dark-400">Pareto Front Size</span>
            <span className="text-sm font-medium text-primary-400">
              {paretoSize} solutions
            </span>
          </div>
        </div>
      )}

      {/* Source indicator (for debugging) */}
      {state?.source && (
        <div className="mt-3 pt-2 border-t border-dark-700">
          <span className="text-xs text-dark-500">
            Source: {state.source}
          </span>
        </div>
      )}
    </div>
  );
}