Atomizer/atomizer-dashboard/frontend/src/pages/Dashboard.tsx

import { useState, useEffect } from 'react';
import {
  LineChart, Line, ScatterChart, Scatter,
  XAxis, YAxis, CartesianGrid, Tooltip, Legend, ResponsiveContainer, Cell
} from 'recharts';
import { useOptimizationWebSocket } from '../hooks/useWebSocket';
import { apiClient } from '../api/client';
import { Card } from '../components/common/Card';
import { MetricCard } from '../components/dashboard/MetricCard';
import { StudyCard } from '../components/dashboard/StudyCard';
import { OptimizerPanel } from '../components/OptimizerPanel';
import { ParetoPlot } from '../components/ParetoPlot';
import { ParallelCoordinatesPlot } from '../components/ParallelCoordinatesPlot';
import type { Study, Trial, ConvergenceDataPoint, ParameterSpaceDataPoint } from '../types';

export default function Dashboard() {
  const [studies, setStudies] = useState<Study[]>([]);
  const [selectedStudyId, setSelectedStudyId] = useState<string | null>(null);
  const [allTrials, setAllTrials] = useState<Trial[]>([]);
  const [displayedTrials, setDisplayedTrials] = useState<Trial[]>([]);
  const [bestValue, setBestValue] = useState<number>(Infinity);
  const [prunedCount, setPrunedCount] = useState<number>(0);
  const [alerts, setAlerts] = useState<Array<{ id: number; type: 'success' | 'warning'; message: string }>>([]);
  const [alertIdCounter, setAlertIdCounter] = useState(0);
  const [expandedTrials, setExpandedTrials] = useState<Set<number>>(new Set());
  const [sortBy, setSortBy] = useState<'performance' | 'chronological'>('performance');

  // Protocol 13: New state for metadata and Pareto front
  const [studyMetadata, setStudyMetadata] = useState<any>(null);
  const [paretoFront, setParetoFront] = useState<any[]>([]);

  // Load studies on mount
  useEffect(() => {
    apiClient.getStudies()
      .then(data => {
        setStudies(data.studies);
        if (data.studies.length > 0) {
          // Check LocalStorage for last selected study
          const savedStudyId = localStorage.getItem('lastSelectedStudyId');
          const studyExists = data.studies.find(s => s.id === savedStudyId);

          if (savedStudyId && studyExists) {
            setSelectedStudyId(savedStudyId);
          } else {
            const running = data.studies.find(s => s.status === 'running');
            setSelectedStudyId(running?.id || data.studies[0].id);
          }
        }
      })
      .catch(console.error);
  }, []);

  const showAlert = (type: 'success' | 'warning', message: string) => {
    const id = alertIdCounter;
    setAlertIdCounter(prev => prev + 1);
    setAlerts(prev => [...prev, { id, type, message }]);
    setTimeout(() => {
      setAlerts(prev => prev.filter(a => a.id !== id));
    }, 5000);
  };

  // WebSocket connection
  const { connectionStatus } = useOptimizationWebSocket({
    studyId: selectedStudyId,
    onMessage: (msg) => {
      if (msg.type === 'trial_completed') {
        const trial = msg.data as Trial;
        setAllTrials(prev => [...prev, trial]);
        if (trial.objective !== null && trial.objective !== undefined && trial.objective < bestValue) {
          setBestValue(trial.objective);
          showAlert('success', `New best: ${trial.objective.toFixed(4)} (Trial #${trial.trial_number})`);
        }
      } else if (msg.type === 'trial_pruned') {
        setPrunedCount(prev => prev + 1);
        showAlert('warning', `Trial pruned: ${msg.data.pruning_cause}`);
      }
    }
  });

  // Load initial trial history when study changes
  useEffect(() => {
    if (selectedStudyId) {
      setAllTrials([]);
      setBestValue(Infinity);
      setPrunedCount(0);
      setExpandedTrials(new Set());

      // Save to LocalStorage
      localStorage.setItem('lastSelectedStudyId', selectedStudyId);

      apiClient.getStudyHistory(selectedStudyId)
        .then(data => {
          const validTrials = data.trials.filter(t => t.objective !== null && t.objective !== undefined);
          setAllTrials(validTrials);
          if (validTrials.length > 0) {
            const minObj = Math.min(...validTrials.map(t => t.objective));
            setBestValue(minObj);
          }
        })
        .catch(console.error);

      apiClient.getStudyPruning(selectedStudyId)
        .then(data => {
          setPrunedCount(data.pruned_trials?.length || 0);
        })
        .catch(console.error);

      // Protocol 13: Fetch metadata
      fetch(`/api/optimization/studies/${selectedStudyId}/metadata`)
        .then(res => res.json())
        .then(data => {
          setStudyMetadata(data);
        })
        .catch(err => console.error('Failed to load metadata:', err));

      // Protocol 13: Fetch Pareto front (raw format for Protocol 13 components)
      fetch(`/api/optimization/studies/${selectedStudyId}/pareto-front`)
        .then(res => res.json())
        .then(paretoData => {
          if (paretoData.is_multi_objective && paretoData.pareto_front) {
            setParetoFront(paretoData.pareto_front);
          } else {
            setParetoFront([]);
          }
        })
        .catch(err => console.error('Failed to load Pareto front:', err));
    }
  }, [selectedStudyId]);

  // Sort trials based on selected sort order
  useEffect(() => {
    let sorted = [...allTrials];
    if (sortBy === 'performance') {
      // Sort by objective (best first)
      sorted.sort((a, b) => {
        const aObj = a.objective ?? Infinity;
        const bObj = b.objective ?? Infinity;
        return aObj - bObj;
      });
    } else {
      // Chronological (newest first)
      sorted.sort((a, b) => b.trial_number - a.trial_number);
    }
    setDisplayedTrials(sorted);
  }, [allTrials, sortBy]);

  // Auto-refresh polling (every 3 seconds) for trial history
  useEffect(() => {
    if (!selectedStudyId) return;

    const refreshInterval = setInterval(() => {
      apiClient.getStudyHistory(selectedStudyId)
        .then(data => {
          const validTrials = data.trials.filter(t => t.objective !== null && t.objective !== undefined);
          setAllTrials(validTrials);
          if (validTrials.length > 0) {
            const minObj = Math.min(...validTrials.map(t => t.objective));
            setBestValue(minObj);
          }
        })
        .catch(err => console.error('Auto-refresh failed:', err));
    }, 3000); // Poll every 3 seconds

    return () => clearInterval(refreshInterval);
  }, [selectedStudyId]);

  // Prepare chart data with proper null/undefined handling
  const convergenceData: ConvergenceDataPoint[] = allTrials
    .filter(t => t.objective !== null && t.objective !== undefined)
    .sort((a, b) => a.trial_number - b.trial_number)
    .map((trial, idx, arr) => {
      const previousTrials = arr.slice(0, idx + 1);
      const validObjectives = previousTrials.map(t => t.objective).filter(o => o !== null && o !== undefined);
      return {
        trial_number: trial.trial_number,
        objective: trial.objective,
        best_so_far: validObjectives.length > 0 ? Math.min(...validObjectives) : trial.objective,
      };
    });

  const parameterSpaceData: ParameterSpaceDataPoint[] = allTrials
    .filter(t => t.objective !== null && t.objective !== undefined && t.design_variables)
    .map(trial => {
      const params = Object.values(trial.design_variables);
      return {
        trial_number: trial.trial_number,
        x: params[0] || 0,
        y: params[1] || 0,
        objective: trial.objective,
        isBest: trial.objective === bestValue,
      };
    });

  // Calculate average objective
  const validObjectives = allTrials.filter(t => t.objective !== null && t.objective !== undefined).map(t => t.objective);
  const avgObjective = validObjectives.length > 0
    ? validObjectives.reduce((sum, obj) => sum + obj, 0) / validObjectives.length
    : 0;

  // Get parameter names
  const paramNames = allTrials.length > 0 && allTrials[0].design_variables
    ? Object.keys(allTrials[0].design_variables)
    : [];

  // Toggle trial expansion
  const toggleTrialExpansion = (trialNumber: number) => {
    setExpandedTrials(prev => {
      const newSet = new Set(prev);
      if (newSet.has(trialNumber)) {
        newSet.delete(trialNumber);
      } else {
        newSet.add(trialNumber);
      }
      return newSet;
    });
  };

  // Export functions
  const exportJSON = () => {
    if (allTrials.length === 0) return;
    const data = JSON.stringify(allTrials, null, 2);
    const blob = new Blob([data], { type: 'application/json' });
    const url = URL.createObjectURL(blob);
    const a = document.createElement('a');
    a.href = url;
    a.download = `${selectedStudyId}_trials.json`;
    a.click();
    URL.revokeObjectURL(url);
    showAlert('success', 'JSON exported successfully!');
  };

  const exportCSV = () => {
    if (allTrials.length === 0) return;
    const headers = ['trial_number', 'objective', ...paramNames].join(',');
    const rows = allTrials.map(t => [
      t.trial_number,
      t.objective,
      ...paramNames.map(k => t.design_variables?.[k] ?? '')
    ].join(','));
    const csv = [headers, ...rows].join('\n');
    const blob = new Blob([csv], { type: 'text/csv' });
    const url = URL.createObjectURL(blob);
    const a = document.createElement('a');
    a.href = url;
    a.download = `${selectedStudyId}_trials.csv`;
    a.click();
    URL.revokeObjectURL(url);
    showAlert('success', 'CSV exported successfully!');
  };

  return (
    <div className="container mx-auto">
      {/* Alerts */}
      <div className="fixed top-4 right-4 z-50 space-y-2">
        {alerts.map(alert => (
          <div
            key={alert.id}
            className={`px-4 py-3 rounded-lg shadow-lg transition-all duration-300 ${
              alert.type === 'success'
                ? 'bg-green-900 border-l-4 border-green-400 text-green-100'
                : 'bg-yellow-900 border-l-4 border-yellow-400 text-yellow-100'
            }`}
          >
            {alert.message}
          </div>
        ))}
      </div>

      {/* Header */}
      <header className="mb-8 flex items-center justify-between border-b border-dark-600 pb-4">
        <div>
          <h1 className="text-3xl font-bold text-primary-400">Live Dashboard</h1>
          <p className="text-dark-300 mt-1">Real-time optimization monitoring</p>
        </div>
        <div className="flex gap-2">
          <button
            onClick={() => {
              if (selectedStudyId) {
                window.open(`http://localhost:8080?study=${selectedStudyId}`, '_blank');
              }
            }}
            className="btn-secondary"
            disabled={!selectedStudyId}
            title="Open Optuna Dashboard (make sure it's running on port 8080)"
          >
            Optuna Dashboard
          </button>
          <button onClick={exportJSON} className="btn-secondary" disabled={allTrials.length === 0}>
            Export JSON
          </button>
          <button onClick={exportCSV} className="btn-secondary" disabled={allTrials.length === 0}>
            Export CSV
          </button>
        </div>
      </header>

      <div className="grid grid-cols-12 gap-6">
        {/* Sidebar - Study List */}
        <aside className="col-span-3">
          <Card title="Active Studies">
            <div className="space-y-3 max-h-[calc(100vh-200px)] overflow-y-auto">
              {studies.map(study => (
                <StudyCard
                  key={study.id}
                  study={study}
                  isActive={study.id === selectedStudyId}
                  onClick={() => setSelectedStudyId(study.id)}
                />
              ))}
            </div>
          </Card>
        </aside>

        {/* Main Content */}
        <main className="col-span-9">
          {/* Metrics Grid */}
          <div className="grid grid-cols-4 gap-4 mb-6">
            <MetricCard label="Total Trials" value={allTrials.length} />
            <MetricCard
              label="Best Value"
              value={bestValue === Infinity ? '-' : bestValue.toFixed(4)}
              valueColor="text-green-400"
            />
            <MetricCard
              label="Avg Objective"
              value={avgObjective > 0 ? avgObjective.toFixed(4) : '-'}
              valueColor="text-blue-400"
            />
            <MetricCard
              label="Pruned"
              value={prunedCount}
              valueColor={prunedCount > 0 ? 'text-red-400' : 'text-green-400'}
            />
          </div>

          {/* Protocol 13: Intelligent Optimizer & Pareto Front */}
          {selectedStudyId && paretoFront.length > 0 && studyMetadata && studyMetadata.objectives && (
            <div className="grid grid-cols-2 gap-6 mb-6">
              <Card title="Optimizer Strategy">
                <div className="space-y-2">
                  <div className="text-sm text-dark-300">
                    <span className="font-semibold text-dark-100">Algorithm:</span> {studyMetadata.sampler || 'NSGA-II'}
                  </div>
                  <div className="text-sm text-dark-300">
                    <span className="font-semibold text-dark-100">Type:</span> Multi-objective
                  </div>
                  <div className="text-sm text-dark-300">
                    <span className="font-semibold text-dark-100">Objectives:</span> {studyMetadata.objectives?.length || 2}
                  </div>
                  <div className="text-sm text-dark-300">
                    <span className="font-semibold text-dark-100">Design Variables:</span> {studyMetadata.design_variables?.length || 0}
                  </div>
                </div>
              </Card>
              <ParetoPlot
                paretoData={paretoFront}
                objectives={studyMetadata.objectives}
              />
            </div>
          )}

          {/* Parallel Coordinates (full width for multi-objective) */}
          {paretoFront.length > 0 && studyMetadata && studyMetadata.objectives && studyMetadata.design_variables && (
            <div className="mb-6">
              <ParallelCoordinatesPlot
                paretoData={paretoFront}
                objectives={studyMetadata.objectives}
                designVariables={studyMetadata.design_variables}
              />
            </div>
          )}

          {/* Charts */}
          <div className="grid grid-cols-2 gap-6 mb-6">
            {/* Convergence Chart */}
            <Card title="Convergence Plot">
              {convergenceData.length > 0 ? (
                <ResponsiveContainer width="100%" height={300}>
                  <LineChart data={convergenceData}>
                    <CartesianGrid strokeDasharray="3 3" stroke="#334155" />
                    <XAxis
                      dataKey="trial_number"
                      stroke="#94a3b8"
                      label={{ value: 'Trial Number', position: 'insideBottom', offset: -5, fill: '#94a3b8' }}
                    />
                    <YAxis
                      stroke="#94a3b8"
                      label={{ value: 'Objective', angle: -90, position: 'insideLeft', fill: '#94a3b8' }}
                    />
                    <Tooltip
                      contentStyle={{ backgroundColor: '#1e293b', border: 'none', borderRadius: '8px' }}
                      labelStyle={{ color: '#e2e8f0' }}
                    />
                    <Legend />
                    <Line
                      type="monotone"
                      dataKey="objective"
                      stroke="#60a5fa"
                      name="Objective"
                      dot={{ r: 3 }}
                    />
                    <Line
                      type="monotone"
                      dataKey="best_so_far"
                      stroke="#10b981"
                      name="Best So Far"
                      strokeWidth={2}
                      dot={{ r: 4 }}
                    />
                  </LineChart>
                </ResponsiveContainer>
              ) : (
                <div className="h-64 flex items-center justify-center text-dark-300">
                  No trial data yet
                </div>
              )}
            </Card>

            {/* Parameter Space Chart */}
            <Card title={`Parameter Space (${paramNames[0] || 'X'} vs ${paramNames[1] || 'Y'})`}>
              {parameterSpaceData.length > 0 ? (
                <ResponsiveContainer width="100%" height={300}>
                  <ScatterChart>
                    <CartesianGrid strokeDasharray="3 3" stroke="#334155" />
                    <XAxis
                      type="number"
                      dataKey="x"
                      stroke="#94a3b8"
                      name={paramNames[0] || 'X'}
                      label={{ value: paramNames[0] || 'Parameter 1', position: 'insideBottom', offset: -5, fill: '#94a3b8' }}
                    />
                    <YAxis
                      type="number"
                      dataKey="y"
                      stroke="#94a3b8"
                      name={paramNames[1] || 'Y'}
                      label={{ value: paramNames[1] || 'Parameter 2', angle: -90, position: 'insideLeft', fill: '#94a3b8' }}
                    />
                    <Tooltip
                      cursor={{ strokeDasharray: '3 3' }}
                      contentStyle={{ backgroundColor: '#1e293b', border: 'none', borderRadius: '8px' }}
                      labelStyle={{ color: '#e2e8f0' }}
                      formatter={(value: any, name: string) => {
                        if (name === 'objective') return [value.toFixed(4), 'Objective'];
                        return [value.toFixed(3), name];
                      }}
                    />
                    <Scatter name="Trials" data={parameterSpaceData}>
                      {parameterSpaceData.map((entry, index) => (
                        <Cell
                          key={`cell-${index}`}
                          fill={entry.isBest ? '#10b981' : '#60a5fa'}
                          r={entry.isBest ? 8 : 5}
                        />
                      ))}
                    </Scatter>
                  </ScatterChart>
                </ResponsiveContainer>
              ) : (
                <div className="h-64 flex items-center justify-center text-dark-300">
                  No trial data yet
                </div>
              )}
            </Card>
          </div>

          {/* Trial History with Sort Controls */}
          <Card
            title={
              <div className="flex items-center justify-between w-full">
                <span>Trial History ({displayedTrials.length} trials)</span>
                <div className="flex gap-2">
                  <button
                    onClick={() => setSortBy('performance')}
                    className={`px-3 py-1 rounded text-sm ${
                      sortBy === 'performance'
                        ? 'bg-primary-500 text-white'
                        : 'bg-dark-500 text-dark-200 hover:bg-dark-400'
                    }`}
                  >
                    Best First
                  </button>
                  <button
                    onClick={() => setSortBy('chronological')}
                    className={`px-3 py-1 rounded text-sm ${
                      sortBy === 'chronological'
                        ? 'bg-primary-500 text-white'
                        : 'bg-dark-500 text-dark-200 hover:bg-dark-400'
                    }`}
                  >
                    Newest First
                  </button>
                </div>
              </div>
            }
          >
            <div className="space-y-2 max-h-[600px] overflow-y-auto">
              {displayedTrials.length > 0 ? (
                displayedTrials.map(trial => {
                  const isExpanded = expandedTrials.has(trial.trial_number);
                  const isBest = trial.objective === bestValue;

                  return (
                    <div
                      key={trial.trial_number}
                      className={`rounded-lg transition-all duration-200 cursor-pointer ${
                        isBest
                          ? 'bg-green-900 border-l-4 border-green-400'
                          : 'bg-dark-500 hover:bg-dark-400'
                      }`}
                      onClick={() => toggleTrialExpansion(trial.trial_number)}
                    >
                      {/* Collapsed View */}
                      <div className="p-3">
                        <div className="flex justify-between items-center">
                          <span className="font-semibold text-primary-400">
                            Trial #{trial.trial_number}
                            {isBest && <span className="ml-2 text-xs bg-green-700 text-green-100 px-2 py-1 rounded">BEST</span>}
                          </span>
                          <div className="flex items-center gap-3">
                            <span className={`font-mono text-lg ${
                              isBest ? 'text-green-400 font-bold' : 'text-dark-100'
                            }`}>
                              {trial.objective !== null && trial.objective !== undefined
                                ? trial.objective.toFixed(4)
                                : 'N/A'}
                            </span>
                            <span className="text-dark-400 text-sm">
                              {isExpanded ? '▼' : '▶'}
                            </span>
                          </div>
                        </div>

                        {/* Quick Preview */}
                        {!isExpanded && trial.results && Object.keys(trial.results).length > 0 && (
                          <div className="text-xs text-primary-300 flex flex-wrap gap-3 mt-2">
                            {trial.results.mass && (
                              <span>Mass: {trial.results.mass.toFixed(2)}g</span>
                            )}
                            {trial.results.frequency && (
                              <span>Freq: {trial.results.frequency.toFixed(2)}Hz</span>
                            )}
                          </div>
                        )}
                      </div>

                      {/* Expanded View */}
                      {isExpanded && (
                        <div className="px-3 pb-3 space-y-3">
                          {/* Design Variables */}
                          {trial.design_variables && Object.keys(trial.design_variables).length > 0 && (
                            <div className="border-t border-dark-400 pt-3">
                              <h4 className="text-sm font-semibold text-dark-200 mb-2">Design Variables</h4>
                              <div className="grid grid-cols-2 gap-2 text-xs">
                                {Object.entries(trial.design_variables).map(([key, val]) => (
                                  <div key={key} className="flex justify-between bg-dark-600 px-2 py-1 rounded">
                                    <span className="text-dark-300">{key}:</span>
                                    <span className="text-dark-100 font-mono">{val.toFixed(4)}</span>
                                  </div>
                                ))}
                              </div>
                            </div>
                          )}

                          {/* Results */}
                          {trial.results && Object.keys(trial.results).length > 0 && (
                            <div className="border-t border-dark-400 pt-3">
                              <h4 className="text-sm font-semibold text-dark-200 mb-2">Extracted Results</h4>
                              <div className="grid grid-cols-2 gap-2 text-xs">
                                {Object.entries(trial.results).map(([key, val]) => (
                                  <div key={key} className="flex justify-between bg-dark-600 px-2 py-1 rounded">
                                    <span className="text-dark-300">{key}:</span>
                                    <span className="text-primary-300 font-mono">
                                      {typeof val === 'number' ? val.toFixed(4) : String(val)}
                                    </span>
                                  </div>
                                ))}
                              </div>
                            </div>
                          )}

                          {/* All User Attributes */}
                          {trial.user_attrs && Object.keys(trial.user_attrs).length > 0 && (
                            <div className="border-t border-dark-400 pt-3">
                              <h4 className="text-sm font-semibold text-dark-200 mb-2">All Attributes</h4>
                              <div className="max-h-48 overflow-y-auto">
                                <pre className="text-xs text-dark-300 bg-dark-700 p-2 rounded">
                                  {JSON.stringify(trial.user_attrs, null, 2)}
                                </pre>
                              </div>
                            </div>
                          )}

                          {/* Timestamps */}
                          {trial.start_time && trial.end_time && (
                            <div className="border-t border-dark-400 pt-3 text-xs text-dark-400">
                              <div className="flex justify-between">
                                <span>Duration:</span>
                                <span>
                                  {((new Date(trial.end_time).getTime() - new Date(trial.start_time).getTime()) / 1000).toFixed(1)}s
                                </span>
                              </div>
                            </div>
                          )}
                        </div>
                      )}
                    </div>
                  );
                })
              ) : (
                <div className="text-center py-8 text-dark-300">
                  No trials yet. Waiting for optimization to start...
                </div>
              )}
            </div>
          </Card>
        </main>
      </div>
    </div>
  );
}