import { useMemo } from 'react';
import Plot from 'react-plotly.js';

interface TrialData {
  trial_number: number;
  values: number[];
  source?: 'FEA' | 'NN' | 'V10_FEA';
  user_attrs?: Record<string, any>;
}

interface PlotlySurrogateQualityProps {
  trials: TrialData[];
  height?: number;
}

export function PlotlySurrogateQuality({
  trials,
  height = 400
}: PlotlySurrogateQualityProps) {
  const { feaTrials, nnTrials, timeline } = useMemo(() => {
    const fea = trials.filter(t => t.source === 'FEA' || t.source === 'V10_FEA');
    const nn = trials.filter(t => t.source === 'NN');

    // Sort by trial number for timeline
    const sorted = [...trials].sort((a, b) => a.trial_number - b.trial_number);

    // Calculate source distribution over time
    const timeline: { trial: number; feaCount: number; nnCount: number }[] = [];
    let feaCount = 0;
    let nnCount = 0;

    sorted.forEach(t => {
      if (t.source === 'NN') nnCount++;
      else feaCount++;

      timeline.push({
        trial: t.trial_number,
        feaCount,
        nnCount
      });
    });

    return {
      feaTrials: fea,
      nnTrials: nn,
      timeline
    };
  }, [trials]);

  if (nnTrials.length === 0) {
    return (
      <div className="h-64 flex items-center justify-center text-dark-400">
        <p>No neural network evaluations in this study</p>
      </div>
    );
  }

  // Objective distribution by source
  const feaObjectives = feaTrials.map(t => t.values[0]).filter(v => v !== undefined && !isNaN(v));
  const nnObjectives = nnTrials.map(t => t.values[0]).filter(v => v !== undefined && !isNaN(v));

  return (
    <div className="space-y-6">
      {/* Source Distribution Over Time */}
      <Plot
        data={[
          {
            x: timeline.map(t => t.trial),
            y: timeline.map(t => t.feaCount),
            type: 'scatter',
            mode: 'lines',
            name: 'FEA Cumulative',
            line: { color: '#3b82f6', width: 2 },
            fill: 'tozeroy',
            fillcolor: 'rgba(59, 130, 246, 0.2)'
          },
          {
            x: timeline.map(t => t.trial),
            y: timeline.map(t => t.nnCount),
            type: 'scatter',
            mode: 'lines',
            name: 'NN Cumulative',
            line: { color: '#a855f7', width: 2 },
            fill: 'tozeroy',
            fillcolor: 'rgba(168, 85, 247, 0.2)'
          }
        ]}
        layout={{
          title: {
            text: 'Evaluation Source Over Time',
            font: { color: '#fff', size: 14 }
          },
          height: height * 0.6,
          margin: { l: 60, r: 30, t: 50, b: 50 },
          paper_bgcolor: 'transparent',
          plot_bgcolor: 'transparent',
          xaxis: {
            title: { text: 'Trial Number', font: { color: '#888' } },
            tickfont: { color: '#888' },
            gridcolor: 'rgba(255,255,255,0.05)'
          },
          yaxis: {
            title: { text: 'Cumulative Count', font: { color: '#888' } },
            tickfont: { color: '#888' },
            gridcolor: 'rgba(255,255,255,0.1)'
          },
          legend: {
            font: { color: '#888' },
            bgcolor: 'rgba(0,0,0,0.5)',
            orientation: 'h',
            y: 1.1
          },
          showlegend: true
        }}
        config={{
          displayModeBar: true,
          modeBarButtonsToRemove: ['lasso2d', 'select2d'],
          displaylogo: false
        }}
        style={{ width: '100%' }}
      />

      {/* Objective Distribution by Source */}
      <Plot
        data={[
          {
            x: feaObjectives,
            type: 'histogram',
            name: 'FEA',
            marker: { color: 'rgba(59, 130, 246, 0.7)' },
            opacity: 0.8
          } as any,
          {
            x: nnObjectives,
            type: 'histogram',
            name: 'NN',
            marker: { color: 'rgba(168, 85, 247, 0.7)' },
            opacity: 0.8
          } as any
        ]}
        layout={{
          title: {
            text: 'Objective Distribution by Source',
            font: { color: '#fff', size: 14 }
          },
          height: height * 0.5,
          margin: { l: 60, r: 30, t: 50, b: 50 },
          paper_bgcolor: 'transparent',
          plot_bgcolor: 'transparent',
          xaxis: {
            title: { text: 'Objective Value', font: { color: '#888' } },
            tickfont: { color: '#888' },
            gridcolor: 'rgba(255,255,255,0.05)'
          },
          yaxis: {
            title: { text: 'Count', font: { color: '#888' } },
            tickfont: { color: '#888' },
            gridcolor: 'rgba(255,255,255,0.1)'
          },
          barmode: 'overlay',
          legend: {
            font: { color: '#888' },
            bgcolor: 'rgba(0,0,0,0.5)',
            orientation: 'h',
            y: 1.1
          },
          showlegend: true
        }}
        config={{
          displayModeBar: true,
          modeBarButtonsToRemove: ['lasso2d', 'select2d'],
          displaylogo: false
        }}
        style={{ width: '100%' }}
      />

      {/* FEA vs NN Best Values Comparison */}
      {feaObjectives.length > 0 && nnObjectives.length > 0 && (
        <div className="grid grid-cols-2 gap-4 mt-4">
          <div className="bg-dark-750 rounded-lg p-4 border border-dark-600">
            <div className="text-xs text-dark-400 uppercase mb-2">FEA Best</div>
            <div className="text-xl font-mono text-blue-400">
              {Math.min(...feaObjectives).toExponential(4)}
            </div>
            <div className="text-xs text-dark-500 mt-1">
              from {feaObjectives.length} evaluations
            </div>
          </div>
          <div className="bg-dark-750 rounded-lg p-4 border border-dark-600">
            <div className="text-xs text-dark-400 uppercase mb-2">NN Best</div>
            <div className="text-xl font-mono text-purple-400">
              {Math.min(...nnObjectives).toExponential(4)}
            </div>
            <div className="text-xs text-dark-500 mt-1">
              from {nnObjectives.length} predictions
            </div>
          </div>
        </div>
      )}
    </div>
  );
}