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 { useWebSocket } from '../hooks/useWebSocket';
import { Card } from '../components/Card';
import { MetricCard } from '../components/MetricCard';
import { StudyCard } from '../components/StudyCard';
import { OptimizerPanel } from '../components/OptimizerPanel';
import { ParetoPlot } from '../components/ParetoPlot';
import { ParallelCoordinatesPlot } from '../components/ParallelCoordinatesPlot';
import type { Study, Trial, ConvergenceDataPoint, ParameterSpaceDataPoint } from '../types';

interface DashboardProps {
  studies: Study[];
  selectedStudyId: string | null;
  onStudySelect: (studyId: string) => void;
}

export default function Dashboard({ studies, selectedStudyId, onStudySelect }: DashboardProps) {
  const [trials, setTrials] = useState<Trial[]>([]);
  const [allTrials, setAllTrials] = 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);

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

  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 { isConnected } = useWebSocket({
    studyId: selectedStudyId,
    onTrialCompleted: (trial) => {
      setTrials(prev => [trial, ...prev].slice(0, 20));
      setAllTrials(prev => [...prev, trial]);
      if (trial.objective < bestValue) {
        setBestValue(trial.objective);
        showAlert('success', `New best: ${trial.objective.toFixed(4)} (Trial #${trial.trial_number})`);
      }
    },
    onNewBest: (trial) => {
      console.log('New best trial:', trial);
    },
    onTrialPruned: (pruned) => {
      setPrunedCount(prev => prev + 1);
      showAlert('warning', `Trial #${pruned.trial_number} pruned: ${pruned.pruning_cause}`);
    },
  });

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

      // Fetch full history
      fetch(`/api/optimization/studies/${selectedStudyId}/history`)
        .then(res => res.json())
        .then(data => {
          const sortedTrials = data.trials.sort((a: Trial, b: Trial) => a.trial_number - b.trial_number);
          setAllTrials(sortedTrials);
          setTrials(sortedTrials.slice(-20).reverse());
          if (sortedTrials.length > 0) {
            const minObj = Math.min(...sortedTrials.map((t: Trial) => t.objective));
            setBestValue(minObj);
          }
        })
        .catch(err => console.error('Failed to load history:', err));

      // Fetch pruning count
      fetch(`/api/optimization/studies/${selectedStudyId}/pruning`)
        .then(res => res.json())
        .then(data => {
          setPrunedCount(data.pruned_trials?.length || 0);
        })
        .catch(err => console.error('Failed to load pruning data:', err));

      // 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
      fetch(`/api/optimization/studies/${selectedStudyId}/pareto-front`)
        .then(res => res.json())
        .then(data => {
          if (data.is_multi_objective) {
            setParetoFront(data.pareto_front);
          } else {
            setParetoFront([]);
          }
        })
        .catch(err => console.error('Failed to load Pareto front:', err));
    }
  }, [selectedStudyId]);

  // Prepare chart data
  const convergenceData: ConvergenceDataPoint[] = allTrials.map((trial, idx) => ({
    trial_number: trial.trial_number,
    objective: trial.objective,
    best_so_far: Math.min(...allTrials.slice(0, idx + 1).map(t => t.objective)),
  }));

  const parameterSpaceData: ParameterSpaceDataPoint[] = allTrials.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 avgObjective = allTrials.length > 0
    ? allTrials.reduce((sum, t) => sum + t.objective, 0) / allTrials.length
    : 0;

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

  // Helper: Format parameter label with unit from metadata
  const getParamLabel = (paramName: string, index: number): string => {
    if (!studyMetadata?.design_variables) {
      return paramName || `Parameter ${index + 1}`;
    }
    const dv = studyMetadata.design_variables.find((v: any) => v.name === paramName);
    if (dv && dv.unit) {
      return `${paramName} (${dv.unit})`;
    }
    return paramName || `Parameter ${index + 1}`;
  };

  // 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 p-6">
      {/* 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-500 pb-4">
        <div>
          <h1 className="text-3xl font-bold text-primary-400">Atomizer Dashboard</h1>
          <p className="text-dark-200 mt-2">Real-time optimization monitoring</p>
        </div>
        <div className="flex gap-2">
          <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={() => onStudySelect(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="Connection"
              value={isConnected ? 'Connected' : 'Disconnected'}
              valueColor={isConnected ? 'text-green-400' : 'text-red-400'}
            />
          </div>

          <div className="grid grid-cols-4 gap-4 mb-6">
            <MetricCard
              label="Pruned"
              value={prunedCount}
              valueColor={prunedCount > 0 ? 'text-red-400' : 'text-green-400'}
            />
          </div>

          {/* Protocol 13: Intelligent Optimizer & Pareto Front */}
          {selectedStudyId && (
            <div className="grid grid-cols-2 gap-6 mb-6">
              <OptimizerPanel studyId={selectedStudyId} />
              {paretoFront.length > 0 && studyMetadata && (
                <ParetoPlot
                  paretoData={paretoFront}
                  objectives={studyMetadata.objectives || []}
                />
              )}
            </div>
          )}

          {/* Parallel Coordinates (full width for multi-objective) */}
          {paretoFront.length > 0 && studyMetadata && (
            <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: getParamLabel(paramNames[0], 0), position: 'insideBottom', offset: -5, fill: '#94a3b8' }}
                    />
                    <YAxis
                      type="number"
                      dataKey="y"
                      stroke="#94a3b8"
                      name={paramNames[1] || 'Y'}
                      label={{ value: getParamLabel(paramNames[1], 1), 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 Feed */}
          <Card title="Recent Trials">
            <div className="space-y-2 max-h-96 overflow-y-auto">
              {trials.length > 0 ? (
                trials.map(trial => (
                  <div
                    key={trial.trial_number}
                    className={`p-3 rounded-lg transition-all duration-200 ${
                      trial.objective === bestValue
                        ? 'bg-green-900 border-l-4 border-green-400'
                        : 'bg-dark-500 hover:bg-dark-400'
                    }`}
                  >
                    <div className="flex justify-between items-center mb-1">
                      <span className="font-semibold text-primary-400">
                        Trial #{trial.trial_number}
                      </span>
                      <span className={`font-mono text-lg ${
                        trial.objective === bestValue ? 'text-green-400 font-bold' : 'text-dark-100'
                      }`}>
                        {trial.objective.toFixed(4)}
                      </span>
                    </div>
                    <div className="text-xs text-dark-200 flex flex-wrap gap-3">
                      {Object.entries(trial.design_variables).map(([key, val]) => (
                        <span key={key}>
                          <span className="text-dark-400">{key}:</span> {val.toFixed(3)}
                        </span>
                      ))}
                    </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>
  );
}