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

/**
 * Zernike Viewer Component
 * Displays interactive Zernike wavefront analysis for mirror optimization trials
 *
 * Features:
 * - 3D surface residual plots (Plotly)
 * - RMS metrics tables
 * - Zernike coefficient bar charts
 * - Tab navigation for different angle comparisons (40°, 60°, 90° vs 20°)
 */

import { useState, useEffect } from 'react';
import { X, RefreshCw, Activity, ChevronLeft, ChevronRight, ExternalLink } from 'lucide-react';

interface ZernikeComparison {
  html: string;
  rms_global: number;
  rms_filtered: number;
  title: string;
}

interface ZernikeData {
  study_id: string;
  trial_number: number;
  comparisons: Record<string, ZernikeComparison>;
  available_comparisons: string[];
}

interface ZernikeViewerProps {
  studyId: string;
  trialNumber: number;
  onClose: () => void;
}

export function ZernikeViewer({ studyId, trialNumber, onClose }: ZernikeViewerProps) {
  const [data, setData] = useState<ZernikeData | null>(null);
  const [loading, setLoading] = useState(true);
  const [error, setError] = useState<string | null>(null);
  const [selectedTab, setSelectedTab] = useState<string>('40_vs_20');

  const fetchZernikeData = async () => {
    setLoading(true);
    setError(null);
    try {
      const response = await fetch(
        `/api/optimization/studies/${studyId}/trials/${trialNumber}/zernike`
      );
      if (!response.ok) {
        const errorData = await response.json().catch(() => ({}));
        throw new Error(errorData.detail || `HTTP ${response.status}`);
      }
      const result = await response.json();
      setData(result);
      // Select first available tab
      if (result.available_comparisons?.length > 0) {
        setSelectedTab(result.available_comparisons[0]);
      }
    } catch (err) {
      setError(err instanceof Error ? err.message : 'Failed to load Zernike analysis');
    } finally {
      setLoading(false);
    }
  };

  useEffect(() => {
    fetchZernikeData();
  }, [studyId, trialNumber]);

  // Tab labels for display
  const tabLabels: Record<string, string> = {
    '40_vs_20': '40° vs 20°',
    '60_vs_20': '60° vs 20°',
    '90_vs_20': '90° vs 20° (Mfg)',
  };

  // Get current comparison data
  const currentComparison = data?.comparisons[selectedTab];

  // Navigate between tabs
  const navigateTab = (direction: 'prev' | 'next') => {
    if (!data?.available_comparisons) return;
    const currentIndex = data.available_comparisons.indexOf(selectedTab);
    if (direction === 'prev' && currentIndex > 0) {
      setSelectedTab(data.available_comparisons[currentIndex - 1]);
    } else if (direction === 'next' && currentIndex < data.available_comparisons.length - 1) {
      setSelectedTab(data.available_comparisons[currentIndex + 1]);
    }
  };

  // Open in new window
  const openInNewWindow = () => {
    if (!currentComparison?.html) return;
    const newWindow = window.open('', '_blank');
    if (newWindow) {
      newWindow.document.write(currentComparison.html);
      newWindow.document.close();
    }
  };

  return (
    <div
      className="fixed inset-0 z-50 flex items-center justify-center bg-black/80"
      onClick={onClose}
    >
      <div
        className="bg-dark-800 rounded-xl shadow-2xl w-[98vw] max-w-[1800px] h-[95vh] flex flex-col border border-dark-600"
        onClick={(e) => e.stopPropagation()}
      >
        {/* Header */}
        <div className="flex items-center justify-between px-6 py-4 border-b border-dark-600 bg-dark-700/50">
          <div className="flex items-center gap-3">
            <Activity className="text-primary-400" size={24} />
            <div>
              <h2 className="text-lg font-semibold text-dark-100">
                Zernike Analysis - Trial #{trialNumber}
              </h2>
              <p className="text-sm text-dark-400">{studyId}</p>
            </div>
          </div>
          <div className="flex items-center gap-2">
            {/* RMS Quick Summary */}
            {currentComparison && (
              <div className="flex items-center gap-4 mr-4 px-4 py-2 bg-dark-600 rounded-lg">
                <div className="text-center">
                  <div className="text-xs text-dark-400">Global RMS</div>
                  <div className="text-sm font-mono text-primary-300">
                    {currentComparison.rms_global.toFixed(2)} nm
                  </div>
                </div>
                <div className="text-center border-l border-dark-500 pl-4">
                  <div className="text-xs text-dark-400">Filtered RMS</div>
                  <div className="text-sm font-mono text-green-400">
                    {currentComparison.rms_filtered.toFixed(2)} nm
                  </div>
                </div>
              </div>
            )}
            <button
              onClick={fetchZernikeData}
              className="p-2 hover:bg-dark-600 rounded-lg transition-colors"
              title="Refresh"
            >
              <RefreshCw size={18} className={`text-dark-300 ${loading ? 'animate-spin' : ''}`} />
            </button>
            <button
              onClick={openInNewWindow}
              className="p-2 hover:bg-dark-600 rounded-lg transition-colors"
              title="Open in new window"
              disabled={!currentComparison}
            >
              <ExternalLink size={18} className="text-dark-300" />
            </button>
            <button
              onClick={onClose}
              className="p-2 hover:bg-dark-600 rounded-lg transition-colors"
            >
              <X size={20} className="text-dark-300" />
            </button>
          </div>
        </div>

        {/* Tabs */}
        {data?.available_comparisons && data.available_comparisons.length > 0 && (
          <div className="flex items-center justify-between px-6 py-2 border-b border-dark-600 bg-dark-700/30">
            <button
              onClick={(e) => {
                e.stopPropagation();
                navigateTab('prev');
              }}
              disabled={data.available_comparisons.indexOf(selectedTab) === 0}
              className="p-2 hover:bg-dark-600 rounded-lg transition-colors disabled:opacity-30"
            >
              <ChevronLeft size={20} className="text-dark-300" />
            </button>

            <div className="flex gap-2">
              {data.available_comparisons.map((tab) => (
                <button
                  key={tab}
                  onClick={(e) => {
                    e.stopPropagation();
                    e.preventDefault();
                    setSelectedTab(tab);
                  }}
                  className={`px-4 py-2 rounded-lg text-sm font-medium transition-colors ${
                    selectedTab === tab
                      ? 'bg-primary-600 text-white'
                      : 'bg-dark-600 text-dark-200 hover:bg-dark-500'
                  }`}
                >
                  {tabLabels[tab] || tab}
                </button>
              ))}
            </div>

            <button
              onClick={(e) => {
                e.stopPropagation();
                navigateTab('next');
              }}
              disabled={data.available_comparisons.indexOf(selectedTab) === data.available_comparisons.length - 1}
              className="p-2 hover:bg-dark-600 rounded-lg transition-colors disabled:opacity-30"
            >
              <ChevronRight size={20} className="text-dark-300" />
            </button>
          </div>
        )}

        {/* Content */}
        <div className="flex-1 overflow-hidden bg-white">
          {loading && (
            <div className="flex flex-col items-center justify-center h-full bg-dark-700">
              <div className="animate-spin rounded-full h-12 w-12 border-b-2 border-primary-400 mb-4"></div>
              <p className="text-dark-300">Generating Zernike analysis...</p>
              <p className="text-dark-500 text-sm mt-2">This may take a few seconds for large meshes</p>
            </div>
          )}

          {error && (
            <div className="flex flex-col items-center justify-center h-full text-dark-400 bg-dark-700">
              <Activity size={48} className="mb-4 opacity-50" />
              <p className="text-lg font-medium text-yellow-400 mb-2">
                {error.includes('surrogate') || error.includes('NN') ? 'No FEA Results' : 'Analysis Failed'}
              </p>
              <p className="text-dark-300 mb-4 max-w-lg text-center">{error}</p>
              {error.includes('surrogate') || error.includes('NN') ? (
                <p className="text-dark-500 text-sm mb-4">
                  Try selecting a trial with "FEA" source tag instead of "NN"
                </p>
              ) : (
                <button
                  onClick={fetchZernikeData}
                  className="px-4 py-2 text-sm bg-dark-600 hover:bg-dark-500 text-dark-200 rounded-lg transition-colors"
                >
                  Try Again
                </button>
              )}
              <button
                onClick={onClose}
                className="mt-4 px-4 py-2 text-sm bg-dark-700 hover:bg-dark-600 text-dark-300 rounded-lg transition-colors border border-dark-500"
              >
                Close
              </button>
            </div>
          )}

          {!loading && !error && currentComparison && (
            <iframe
              srcDoc={currentComparison.html}
              className="w-full h-full border-0"
              title={`Zernike Analysis - ${tabLabels[selectedTab] || selectedTab}`}
              sandbox="allow-scripts allow-same-origin"
            />
          )}
        </div>
      </div>
    </div>
  );
}

/**
 * Button component to trigger Zernike viewer for a trial
 * Used in Dashboard trial list
 */
interface ZernikeButtonProps {
  studyId: string;
  trialNumber: number;
  compact?: boolean;
}

export function ZernikeButton({ studyId, trialNumber, compact = false }: ZernikeButtonProps) {
  const [isOpen, setIsOpen] = useState(false);

  return (
    <>
      <button
        onClick={(e) => {
          e.stopPropagation();  // Prevent trial row expansion
          setIsOpen(true);
        }}
        className={`flex items-center gap-1 ${
          compact
            ? 'px-2 py-1 text-xs'
            : 'px-3 py-1.5 text-sm'
        } bg-indigo-600 hover:bg-indigo-700 text-white rounded transition-colors font-medium`}
        title="View Zernike wavefront analysis"
      >
        <Activity size={compact ? 12 : 14} />
        <span>{compact ? 'Zernike' : 'View Zernike'}</span>
      </button>

      {isOpen && (
        <ZernikeViewer
          studyId={studyId}
          trialNumber={trialNumber}
          onClose={() => setIsOpen(false)}
        />
      )}
    </>
  );
}