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

interface TrialData {
  trial_number: number;
  values: number[];
  params: Record<string, number>;
}

interface PlotlyCorrelationHeatmapProps {
  trials: TrialData[];
  objectiveName?: string;
  height?: number;
}

// Calculate Pearson correlation coefficient
function pearsonCorrelation(x: number[], y: number[]): number {
  const n = x.length;
  if (n === 0 || n !== y.length) return 0;

  const meanX = x.reduce((a, b) => a + b, 0) / n;
  const meanY = y.reduce((a, b) => a + b, 0) / n;

  let numerator = 0;
  let denomX = 0;
  let denomY = 0;

  for (let i = 0; i < n; i++) {
    const dx = x[i] - meanX;
    const dy = y[i] - meanY;
    numerator += dx * dy;
    denomX += dx * dx;
    denomY += dy * dy;
  }

  const denominator = Math.sqrt(denomX) * Math.sqrt(denomY);
  return denominator === 0 ? 0 : numerator / denominator;
}

export function PlotlyCorrelationHeatmap({
  trials,
  objectiveName = 'Objective',
  height = 500
}: PlotlyCorrelationHeatmapProps) {
  const { matrix, labels, annotations } = useMemo(() => {
    if (trials.length < 3) {
      return { matrix: [], labels: [], annotations: [] };
    }

    // Get parameter names
    const paramNames = Object.keys(trials[0].params);
    const allLabels = [...paramNames, objectiveName];

    // Extract data columns
    const columns: Record<string, number[]> = {};
    paramNames.forEach(name => {
      columns[name] = trials.map(t => t.params[name]).filter(v => v !== undefined && !isNaN(v));
    });
    columns[objectiveName] = trials.map(t => t.values[0]).filter(v => v !== undefined && !isNaN(v));

    // Calculate correlation matrix
    const n = allLabels.length;
    const correlationMatrix: number[][] = [];
    const annotationData: any[] = [];

    for (let i = 0; i < n; i++) {
      const row: number[] = [];
      for (let j = 0; j < n; j++) {
        const col1 = columns[allLabels[i]];
        const col2 = columns[allLabels[j]];

        // Ensure same length
        const minLen = Math.min(col1.length, col2.length);
        const corr = pearsonCorrelation(col1.slice(0, minLen), col2.slice(0, minLen));
        row.push(corr);

        // Add annotation
        annotationData.push({
          x: allLabels[j],
          y: allLabels[i],
          text: corr.toFixed(2),
          showarrow: false,
          font: {
            color: Math.abs(corr) > 0.5 ? '#fff' : '#888',
            size: 11
          }
        });
      }
      correlationMatrix.push(row);
    }

    return {
      matrix: correlationMatrix,
      labels: allLabels,
      annotations: annotationData
    };
  }, [trials, objectiveName]);

  if (trials.length < 3) {
    return (
      <div className="h-64 flex items-center justify-center text-dark-400">
        <p>Need at least 3 trials to compute correlations</p>
      </div>
    );
  }

  return (
    <Plot
      data={[
        {
          z: matrix,
          x: labels,
          y: labels,
          type: 'heatmap',
          colorscale: [
            [0, '#ef4444'],    // -1: strong negative (red)
            [0.25, '#f87171'], // -0.5: moderate negative
            [0.5, '#1a1b26'],  // 0: no correlation (dark)
            [0.75, '#60a5fa'], // 0.5: moderate positive
            [1, '#3b82f6']     // 1: strong positive (blue)
          ],
          zmin: -1,
          zmax: 1,
          showscale: true,
          colorbar: {
            title: { text: 'Correlation', font: { color: '#888' } },
            tickfont: { color: '#888' },
            len: 0.8
          },
          hovertemplate: '%{y} vs %{x}<br>Correlation: %{z:.3f}<extra></extra>'
        }
      ]}
      layout={{
        title: {
          text: 'Parameter-Objective Correlation Matrix',
          font: { color: '#fff', size: 14 }
        },
        height,
        margin: { l: 120, r: 60, t: 60, b: 120 },
        paper_bgcolor: 'transparent',
        plot_bgcolor: 'transparent',
        xaxis: {
          tickangle: 45,
          tickfont: { color: '#888', size: 10 },
          gridcolor: 'rgba(255,255,255,0.05)'
        },
        yaxis: {
          tickfont: { color: '#888', size: 10 },
          gridcolor: 'rgba(255,255,255,0.05)'
        },
        annotations: annotations
      }}
      config={{
        displayModeBar: true,
        modeBarButtonsToRemove: ['lasso2d', 'select2d'],
        displaylogo: false
      }}
      style={{ width: '100%' }}
    />
  );
}
