import React, { useState, useMemo, useEffect, useRef } from 'react';
import { format } from 'date-fns';
import {
  LineChart, Line, XAxis, YAxis, CartesianGrid, Tooltip,
  Legend, ResponsiveContainer,
} from 'recharts';
import { SubjectSidebar } from './SubjectSidebar';
import { sortPayloadByValueDesc, toggleInSet, setIdsHidden, serializePlotConfig, readPlotConfig } from '../lib/crossSubjectChart';

// ── Color palettes ─────────────────────────────────────────────────────────────

const CAT_COLORS   = { Success: '#22c55e', Failure: '#ef4444', Other: '#9ca3af' };
const EXTRA_CAT    = ['#f59e0b', '#a855f7', '#ec4899', '#14b8a6', '#f97316'];
const GROUP_COLORS = ['#6366f1', '#f59e0b', '#22c55e', '#ef4444', '#a855f7', '#14b8a6', '#ec4899', '#f97316', '#84cc16', '#06b6d4'];

function catColor(cat, idx = 0) {
  return CAT_COLORS[cat] ?? EXTRA_CAT[idx % EXTRA_CAT.length];
}

// ── Shared tooltip ─────────────────────────────────────────────────────────────

function ChartTooltip({ active, payload, label, labelKey = 'dateStr', unit = '' }) {
  if (!active || !payload?.length) return null;
  const header = payload[0]?.payload?.[labelKey] ?? `Day ${label}`;
  return (
    <div className="bg-white border border-gray-200 rounded shadow-sm px-3 py-2 text-xs space-y-0.5">
      <p className="font-semibold text-gray-700 mb-1">{header}</p>
      {payload.map((p) => (
        <p key={p.dataKey} style={{ color: p.stroke ?? p.color }}>
          {p.name}: {p.value != null ? `${p.value}${unit}` : '—'}
        </p>
      ))}
    </div>
  );
}

// ── Shared helpers ─────────────────────────────────────────────────────────────

const DAYS_REACH_RE = /#?\s*days?\s*reach/i;
const SELECT_CLS    = 'border border-gray-200 rounded px-1.5 py-0.5 text-xs bg-white focus:outline-none focus:ring-1 focus:ring-indigo-400';
const INPUT_CLS     = 'w-16 border border-gray-200 rounded px-1.5 py-0.5 text-xs bg-white focus:outline-none focus:ring-1 focus:ring-indigo-400';
const STEP_OPTIONS  = [1, 2, 5, 10];

function defaultXField(fields) {
  return fields.find((f) => DAYS_REACH_RE.test(f.label))?.fieldId ?? '__days__';
}

function numericFieldVal(status, field) {
  if (!field) return null;
  const raw = field.builtin ? status[field.key] : status.custom_fields?.[field.fieldId];
  const n = parseFloat(raw);
  return isNaN(n) ? null : n;
}

function buildTicks(data, step) {
  const xs = data.map((d) => d.x).filter((x) => x != null && isFinite(x));
  if (!xs.length) return undefined;
  const lo = Math.floor(Math.min(...xs));
  const hi = Math.ceil(Math.max(...xs));
  const ticks = [];
  for (let t = Math.ceil(lo / step) * step; t <= hi; t += step) ticks.push(t);
  if (!ticks.length || ticks[0] > lo) ticks.unshift(lo);
  if (ticks[ticks.length - 1] < hi) ticks.push(hi);
  return ticks;
}

function TickStepControl({ value, onChange }) {
  return (
    <div className="flex items-center gap-1.5 text-xs">
      <span className="text-gray-400">Tick every</span>
      <div className="flex rounded border border-gray-200 overflow-hidden">
        {STEP_OPTIONS.map((s) => (
          <button key={s} type="button" onClick={() => onChange(s)}
            className={`px-2 py-0.5 transition-colors border-l border-gray-200 first:border-l-0
              ${value === s ? 'bg-gray-700 text-white' : 'bg-white text-gray-500 hover:bg-gray-50'}`}>
            {s}
          </button>
        ))}
      </div>
    </div>
  );
}

// ── Per-chart customisation hook & UI ──────────────────────────────────────────

const HEIGHT_OPTIONS = [
  { label: 'S', value: 140 },
  { label: 'M', value: 220 },
  { label: 'L', value: 340 },
];

function useChartSettings(defaultHeight, initial = {}) {
  const [expanded,       setExpanded]       = useState(false);
  const [height,         setHeight]         = useState(initial.height ?? defaultHeight);
  const [yMin,           setYMin]           = useState(initial.yMin ?? '');
  const [yMax,           setYMax]           = useState(initial.yMax ?? '');
  const [xZoomMin,       setXZoomMin]       = useState(initial.xZoomMin ?? '');
  const [xZoomMax,       setXZoomMax]       = useState(initial.xZoomMax ?? '');
  const [xFieldOverride, setXFieldOverride] = useState(initial.xFieldOverride ?? null); // null → use parent's xField

  const yDomain = useMemo(() => [
    yMin !== '' ? parseFloat(yMin) : 'auto',
    yMax !== '' ? parseFloat(yMax) : 'auto',
  ], [yMin, yMax]);

  const xDomain = useMemo(() => {
    const hasMin = xZoomMin !== '';
    const hasMax = xZoomMax !== '';
    if (!hasMin && !hasMax) return null;
    return [hasMin ? parseFloat(xZoomMin) : 'dataMin', hasMax ? parseFloat(xZoomMax) : 'dataMax'];
  }, [xZoomMin, xZoomMax]);

  const isDirty = yMin !== '' || yMax !== '' || xZoomMin !== '' || xZoomMax !== ''
    || xFieldOverride !== null || height !== defaultHeight;

  const snapshot = useMemo(
    () => ({ height, yMin, yMax, xZoomMin, xZoomMax, xFieldOverride }),
    [height, yMin, yMax, xZoomMin, xZoomMax, xFieldOverride],
  );

  function reset() {
    setHeight(defaultHeight);
    setYMin(''); setYMax('');
    setXZoomMin(''); setXZoomMax('');
    setXFieldOverride(null);
  }

  return {
    expanded, setExpanded,
    height, setHeight,
    yMin, setYMin, yMax, setYMax,
    xZoomMin, setXZoomMin, xZoomMax, setXZoomMax,
    xFieldOverride, setXFieldOverride,
    yDomain, xDomain, isDirty, reset, snapshot,
  };
}

function ChartHeader({ title, isDirty, expanded, onToggle }) {
  return (
    <div className="flex items-center justify-between mb-2">
      <p className="text-xs text-gray-500 font-medium">{title}</p>
      <button type="button" onClick={onToggle}
        className={[
          'text-xs px-2 py-0.5 rounded border transition-colors',
          expanded
            ? 'bg-gray-100 border-gray-300 text-gray-700'
            : isDirty
              ? 'bg-indigo-50 border-indigo-200 text-indigo-600'
              : 'bg-white border-gray-200 text-gray-400 hover:text-gray-700 hover:border-gray-300',
        ].join(' ')}
      >
        {expanded ? '▲ collapse' : '▾ customize'}
      </button>
    </div>
  );
}

// Shown only when the chart card is expanded — renders per-chart controls.
function ChartSettingsBar({ s, xAxisFields, globalXField, extra }) {
  return (
    <div className="mb-3 p-3 bg-gray-50 border border-gray-100 rounded-lg flex flex-wrap gap-x-5 gap-y-2 items-center text-xs">

      {/* Size */}
      <div className="flex items-center gap-1.5">
        <span className="text-gray-400 shrink-0">Size:</span>
        <div className="flex rounded border border-gray-200 overflow-hidden">
          {HEIGHT_OPTIONS.map((opt) => (
            <button key={opt.value} type="button" onClick={() => s.setHeight(opt.value)}
              className={`px-2 py-0.5 transition-colors border-l border-gray-200 first:border-l-0
                ${s.height === opt.value ? 'bg-gray-700 text-white' : 'bg-white text-gray-500 hover:bg-gray-50'}`}>
              {opt.label}
            </button>
          ))}
        </div>
      </div>

      {/* Y range */}
      <div className="flex items-center gap-1.5">
        <span className="text-gray-400 shrink-0">Y range:</span>
        <input type="number" placeholder="min" value={s.yMin} onChange={(e) => s.setYMin(e.target.value)} className={INPUT_CLS} />
        <span className="text-gray-300">–</span>
        <input type="number" placeholder="max" value={s.yMax} onChange={(e) => s.setYMax(e.target.value)} className={INPUT_CLS} />
      </div>

      {/* X zoom */}
      <div className="flex items-center gap-1.5">
        <span className="text-gray-400 shrink-0">X zoom:</span>
        <input type="number" placeholder="min" value={s.xZoomMin} onChange={(e) => s.setXZoomMin(e.target.value)} className={INPUT_CLS} />
        <span className="text-gray-300">–</span>
        <input type="number" placeholder="max" value={s.xZoomMax} onChange={(e) => s.setXZoomMax(e.target.value)} className={INPUT_CLS} />
      </div>

      {/* Per-chart X axis field override */}
      {xAxisFields.length > 0 && (
        <div className="flex items-center gap-1.5">
          <span className="text-gray-400 shrink-0">X data:</span>
          <select
            value={s.xFieldOverride ?? globalXField}
            onChange={(e) => {
              const v = e.target.value;
              s.setXFieldOverride(v === globalXField ? null : v);
            }}
            className={SELECT_CLS}
          >
            <option value="__days__"># Days Reach (ordinal)</option>
            {xAxisFields.map((f) => <option key={f.fieldId} value={f.fieldId}>{f.label}</option>)}
          </select>
        </div>
      )}

      {/* Slot for chart-specific controls (e.g. Y line visibility) */}
      {extra}

      {s.isDirty && (
        <button type="button" onClick={s.reset}
          className="text-xs text-red-400 hover:text-red-600 transition-colors ml-auto">
          ↺ Reset
        </button>
      )}
    </div>
  );
}

// ── Pure data-computation helpers ──────────────────────────────────────────────

function computeSubjectChartData(xField, statuses, xAxisFields) {
  const xAxisField = xAxisFields.find((f) => f.fieldId === xField);
  const xLabel     = xAxisField ? xAxisField.label : '# Days Reach';

  const withMetrics = statuses
    .filter((s) => s.analysis_summary?.total != null)
    .sort((a, b) => String(a.date).localeCompare(String(b.date)));

  const rows = withMetrics.map((s, i) => {
    const xVal    = xField === '__days__' || !xAxisField ? i + 1 : numericFieldVal(s, xAxisField);
    const dateStr = format(new Date(String(s.date).slice(0, 10) + 'T12:00:00'), 'MMM d');
    const { counts = {}, total, success_rate } = s.analysis_summary;
    return { _x: xVal, dateStr, total, rate: success_rate != null ? +(success_rate * 100).toFixed(1) : null, ...counts };
  }).filter((r) => r._x != null).sort((a, b) => a._x - b._x);

  const data = rows.map((r) => ({ ...r, x: r._x }));

  const catSet = new Set();
  for (const d of data) {
    for (const k of Object.keys(d)) {
      if (!['x', '_x', 'dateStr', 'total', 'rate'].includes(k)) catSet.add(k);
    }
  }
  return { data, cats: [...catSet], xLabel };
}

function computeExpChartData(xField, chartType, metric, animalRows, lines, xAxisFields) {
  // chartType: 'count' | 'rate'
  const xAxisField = xAxisFields.find((f) => f.fieldId === xField);
  const xLabel     = xAxisField ? xAxisField.label : '# Days Reach';

  function getX(row, dayIndex) {
    if (xField === '__days__') return dayIndex + 1;
    return numericFieldVal(row, xAxisField);
  }

  if (xField === '__days__') {
    const maxDay = Math.max(0, ...lines.map((l) => animalRows[l.id].length));
    if (maxDay === 0) return { data: [], xLabel };
    const data = Array.from({ length: maxDay }, (_, i) => ({ x: i + 1 }));
    for (const line of lines) {
      animalRows[line.id].forEach((row, i) => {
        const s = row.analysis_summary;
        if (chartType === 'count') {
          data[i][line.id] = metric === 'total' ? s.total : (s.counts?.[metric] ?? null);
        } else {
          data[i][line.id] = s.success_rate != null ? +(s.success_rate * 100).toFixed(1) : null;
        }
      });
    }
    return { data, xLabel };
  }

  // Field-based x axis
  const allPts = [];
  for (const line of lines) {
    animalRows[line.id].forEach((row, i) => {
      const x = getX(row, i);
      if (x !== null) allPts.push({ x, animalId: line.id, summary: row.analysis_summary });
    });
  }
  const sortedX = [...new Set(allPts.map((p) => p.x))].sort((a, b) => a - b);
  if (sortedX.length === 0) return { data: [], xLabel };

  const data = sortedX.map((x) => {
    const pt = { x };
    for (const p of allPts.filter((pp) => pp.x === x)) {
      const s = p.summary;
      if (chartType === 'count') {
        pt[p.animalId] = metric === 'total' ? s.total : (s.counts?.[metric] ?? null);
      } else {
        pt[p.animalId] = s.success_rate != null ? +(s.success_rate * 100).toFixed(1) : null;
      }
    }
    return pt;
  });
  return { data, xLabel };
}

// ── Subject charts ─────────────────────────────────────────────────────────────

export function SubjectAnalysisCharts({ statuses, dailyTemplate }) {
  const xAxisFields = useMemo(
    () => (dailyTemplate ?? []).filter((f) => f.active),
    [dailyTemplate],
  );
  const [xField,    setXField]    = useState(() => defaultXField(xAxisFields));
  const [tickStep,  setTickStep]  = useState(1);

  const countsS = useChartSettings(200);
  const rateS   = useChartSettings(160);

  // Shared base computation — used by both charts when no override is set
  const sharedData = useMemo(
    () => computeSubjectChartData(xField, statuses, xAxisFields),
    [xField, statuses, xAxisFields],
  );

  // Per-chart overrides: only recompute when the override differs from the shared field
  const countsData = useMemo(() => {
    const ef = countsS.xFieldOverride;
    if (!ef || ef === xField) return sharedData;
    return computeSubjectChartData(ef, statuses, xAxisFields);
  }, [countsS.xFieldOverride, xField, sharedData, statuses, xAxisFields]);

  const rateData = useMemo(() => {
    const ef = rateS.xFieldOverride;
    if (!ef || ef === xField) return sharedData;
    return computeSubjectChartData(ef, statuses, xAxisFields);
  }, [rateS.xFieldOverride, xField, sharedData, statuses, xAxisFields]);

  const countsTicks = useMemo(() => buildTicks(countsData.data, tickStep), [countsData.data, tickStep]);
  const rateTicks   = useMemo(() => buildTicks(rateData.data,   tickStep), [rateData.data,   tickStep]);

  // Y-line visibility for the counts chart (null = show all)
  const [visibleCats, setVisibleCats] = useState(null);
  useEffect(() => { setVisibleCats(null); }, [countsData.cats.join(',')]); // reset when cats change

  if (sharedData.data.length === 0) return null;

  function toggleCat(cat) {
    setVisibleCats((prev) => {
      if (prev === null) return new Set([cat]);
      const next = new Set(prev);
      next.has(cat) ? next.delete(cat) : next.add(cat);
      return next.size === 0 ? null : next;
    });
  }

  function xAxisProps(data, ticks, s) {
    const ef         = s.xFieldOverride ?? xField;
    const xAxisField = xAxisFields.find((f) => f.fieldId === ef);
    const xLabel     = xAxisField ? xAxisField.label : '# Days Reach';
    const domain     = s.xDomain ?? (ticks?.length ? [ticks[0], ticks[ticks.length - 1]] : ['auto', 'auto']);
    return { dataKey: 'x', type: 'number', tick: { fontSize: 10 }, ticks, domain, allowDecimals: false,
      label: { value: xLabel, position: 'insideBottomRight', offset: -4, fontSize: 10 } };
  }

  return (
    <div className="bg-white border border-gray-200 rounded-xl p-5 mb-6 space-y-5">
      {/* Global controls */}
      <div className="flex items-center flex-wrap gap-3">
        <h2 className="text-sm font-semibold text-gray-700 uppercase tracking-wide">Session metrics over time</h2>
        <div className="flex items-center gap-1.5 text-xs">
          <span className="text-gray-400">X axis:</span>
          <select value={xField} onChange={(e) => setXField(e.target.value)} className={SELECT_CLS}>
            <option value="__days__"># Days Reach (ordinal)</option>
            {xAxisFields.map((f) => <option key={f.fieldId} value={f.fieldId}>{f.label}</option>)}
          </select>
        </div>
        <TickStepControl value={tickStep} onChange={setTickStep} />
      </div>

      {/* Chart 1 — counts */}
      <div>
        <ChartHeader title="Attempts per session" isDirty={countsS.isDirty} expanded={countsS.expanded} onToggle={() => countsS.setExpanded((e) => !e)} />
        {countsS.expanded && (
          <ChartSettingsBar s={countsS} xAxisFields={xAxisFields} globalXField={xField}
            extra={
              countsData.cats.length > 0 && (
                <div className="flex items-center gap-2 flex-wrap">
                  <span className="text-gray-400 shrink-0">Y lines:</span>
                  {countsData.cats.map((cat, i) => {
                    const on = visibleCats === null || visibleCats.has(cat);
                    return (
                      <button key={cat} type="button" onClick={() => toggleCat(cat)}
                        style={{ color: catColor(cat, i), borderColor: on ? catColor(cat, i) : '#e5e7eb' }}
                        className={`px-2 py-0.5 rounded border text-xs transition-opacity ${on ? 'opacity-100' : 'opacity-35'}`}>
                        {cat}
                      </button>
                    );
                  })}
                </div>
              )
            }
          />
        )}
        <ResponsiveContainer width="100%" height={countsS.height}>
          <LineChart data={countsData.data} margin={{ top: 4, right: 16, left: 0, bottom: 4 }}>
            <CartesianGrid strokeDasharray="3 3" stroke="#f3f4f6" />
            <XAxis {...xAxisProps(countsData.data, countsTicks, countsS)} />
            <YAxis tick={{ fontSize: 10 }} width={30} allowDecimals={false} domain={countsS.yDomain} />
            <Tooltip content={<ChartTooltip labelKey="dateStr" />} />
            <Legend wrapperStyle={{ fontSize: 11 }} />
            {countsData.cats
              .filter((cat) => visibleCats === null || visibleCats.has(cat))
              .map((cat, i) => (
                <Line key={cat} type="monotone" dataKey={cat} stroke={catColor(cat, i)} strokeWidth={2}
                  dot={{ r: 3 }} activeDot={{ r: 4 }} connectNulls isAnimationActive={false} />
              ))}
            <Line type="monotone" dataKey="total" name="Total" stroke="#3b82f6" strokeWidth={1.5}
              strokeDasharray="5 3" dot={{ r: 3 }} connectNulls isAnimationActive={false} />
          </LineChart>
        </ResponsiveContainer>
      </div>

      {/* Chart 2 — success rate */}
      <div>
        <ChartHeader title="Success rate (%)" isDirty={rateS.isDirty} expanded={rateS.expanded} onToggle={() => rateS.setExpanded((e) => !e)} />
        {rateS.expanded && (
          <ChartSettingsBar s={rateS} xAxisFields={xAxisFields} globalXField={xField} />
        )}
        <ResponsiveContainer width="100%" height={rateS.height}>
          <LineChart data={rateData.data} margin={{ top: 4, right: 16, left: 0, bottom: 4 }}>
            <CartesianGrid strokeDasharray="3 3" stroke="#f3f4f6" />
            <XAxis {...xAxisProps(rateData.data, rateTicks, rateS)} />
            <YAxis tick={{ fontSize: 10 }} width={36}
              domain={rateS.yDomain[0] !== 'auto' || rateS.yDomain[1] !== 'auto' ? rateS.yDomain : [0, 100]}
              tickFormatter={(v) => `${v}%`} />
            <Tooltip content={<ChartTooltip unit="%" labelKey="dateStr" />} />
            <Line type="monotone" dataKey="rate" name="Success rate" stroke="#6366f1" strokeWidth={2}
              dot={{ r: 3 }} connectNulls isAnimationActive={false} />
          </LineChart>
        </ResponsiveContainer>
      </div>
    </div>
  );
}

// ── Experiment charts ──────────────────────────────────────────────────────────

const METRIC_OPTIONS = [
  { value: 'total',   label: 'Total attempts' },
  { value: 'Success', label: 'Success count'  },
  { value: 'Failure', label: 'Failure count'  },
];

export function ExperimentAnalysisCharts({ animals, experimentStatuses, subjectTemplate, dailyTemplate, plotConfig, onPersistConfig }) {
  const groupableFields = useMemo(() => (subjectTemplate ?? []).filter((f) => f.active), [subjectTemplate]);
  const xAxisFields     = useMemo(() => (dailyTemplate    ?? []).filter((f) => f.active), [dailyTemplate]);

  const defaultGroup = (groupableFields.find((f) => /^group$/i.test(f.label)) ?? groupableFields[0])?.fieldId ?? '';
  // Read the saved config exactly once (component only mounts after the experiment loads).
  const initialConfigRef = useRef(null);
  if (initialConfigRef.current === null) initialConfigRef.current = readPlotConfig(plotConfig);
  const initialConfig = initialConfigRef.current;
  const [groupBy,  setGroupBy]  = useState(initialConfig.groupBy ?? defaultGroup);
  const [metric,   setMetric]   = useState(initialConfig.metric ?? 'total');
  const [xField,   setXField]   = useState(() => initialConfig.xField ?? defaultXField(xAxisFields));
  const [tickStep, setTickStep] = useState(initialConfig.tickStep ?? 1);

  // Cross-subject interaction state (client-only; global across both charts)
  const [sidebarOpen,        setSidebarOpen]        = useState(initialConfig.sidebarOpen);
  const [hiddenSubjects,     setHiddenSubjects]     = useState(() => new Set(initialConfig.hiddenSubjects));
  const [highlightedSubject, setHighlightedSubject] = useState(null);

  const toggleHidden = (id)          => setHiddenSubjects((prev) => toggleInSet(prev, id));
  const toggleGroup  = (ids, hidden) => setHiddenSubjects((prev) => setIdsHidden(prev, ids, hidden));
  const showAll      = ()            => setHiddenSubjects(new Set());

  const countsS = useChartSettings(200, initialConfig.counts);
  const rateS   = useChartSettings(160, initialConfig.rate);

  // ── Persist display settings (debounced auto-save) ────────────────────────────
  const [saveState, setSaveState] = useState('idle'); // 'idle' | 'saving' | 'saved' | 'error'

  const currentConfig = useMemo(
    () => serializePlotConfig({
      xField, groupBy, tickStep, metric, hiddenSubjects, sidebarOpen,
      counts: countsS.snapshot, rate: rateS.snapshot,
    }),
    [xField, groupBy, tickStep, metric, hiddenSubjects, sidebarOpen, countsS.snapshot, rateS.snapshot],
  );

  // Keep the latest persist callback in a ref so the save effect does not re-run
  // when the parent passes a new inline function each render.
  const persistRef = useRef(onPersistConfig);
  persistRef.current = onPersistConfig;

  const lastSavedRef = useRef(null);
  useEffect(() => {
    const serialized = JSON.stringify(currentConfig);
    if (lastSavedRef.current === null) { lastSavedRef.current = serialized; return; } // skip initial (hydration)
    // Reverted to the last-saved value before the debounce fired → nothing pending.
    if (serialized === lastSavedRef.current) { setSaveState('idle'); return; }
    if (!persistRef.current) { lastSavedRef.current = serialized; setSaveState('idle'); return; }
    setSaveState('saving');
    const t = setTimeout(() => {
      Promise.resolve(persistRef.current(currentConfig))
        .then(() => { lastSavedRef.current = serialized; setSaveState('saved'); })
        .catch(() => setSaveState('error'));
    }, 1000);
    return () => clearTimeout(t);
  }, [currentConfig]);

  // Step 1: animalRows — does NOT depend on xField; shared by both charts
  const { animalRows, lines } = useMemo(() => {
    const animalGroup = {};
    for (const animal of animals) {
      const raw = groupBy ? animal.subject_info?.[groupBy] : null;
      animalGroup[animal.id] = raw != null && raw !== '' ? String(raw) : '—';
    }
    const groups = [...new Set(Object.values(animalGroup))].sort();

    const statusesByAnimal = {};
    for (const s of experimentStatuses) {
      if (!statusesByAnimal[s.animal_id]) statusesByAnimal[s.animal_id] = [];
      statusesByAnimal[s.animal_id].push(s);
    }
    const animalRows = {};
    for (const animal of animals) {
      animalRows[animal.id] = (statusesByAnimal[animal.id] ?? [])
        .filter((s) => s.analysis_summary?.total != null)
        .sort((a, b) => String(a.date).localeCompare(String(b.date)));
    }

    const lines = animals
      .filter((a) => animalRows[a.id]?.length > 0)
      .map((a) => ({
        id:    a.id,
        name:  a.animal_name,
        group: animalGroup[a.id],
        color: GROUP_COLORS[groups.indexOf(animalGroup[a.id]) % GROUP_COLORS.length],
      }));

    return { animalRows, lines };
  }, [animals, experimentStatuses, groupBy]);

  // Step 2: count chart data — depends on xField for counts + metric
  const { data: countData, xLabel: countsXLabel } = useMemo(() => {
    if (lines.length === 0) return { data: [], xLabel: '' };
    const ef = countsS.xFieldOverride ?? xField;
    return computeExpChartData(ef, 'count', metric, animalRows, lines, xAxisFields);
  }, [animalRows, lines, countsS.xFieldOverride, xField, metric, xAxisFields]);

  // Step 3: rate chart data — depends on xField for rate
  const { data: rateData, xLabel: rateXLabel } = useMemo(() => {
    if (lines.length === 0) return { data: [], xLabel: '' };
    const ef = rateS.xFieldOverride ?? xField;
    return computeExpChartData(ef, 'rate', null, animalRows, lines, xAxisFields);
  }, [animalRows, lines, rateS.xFieldOverride, xField, xAxisFields]);

  const countsTicks = useMemo(() => buildTicks(countData, tickStep), [countData, tickStep]);
  const rateTicks   = useMemo(() => buildTicks(rateData,  tickStep), [rateData,  tickStep]);

  const hasData = lines.length > 0;

  function xAxisProps(data, ticks, s, fallbackLabel) {
    const domain = s.xDomain ?? (ticks?.length ? [ticks[0], ticks[ticks.length - 1]] : ['auto', 'auto']);
    const label  = (s.xFieldOverride ? xAxisFields.find((f) => f.fieldId === s.xFieldOverride)?.label : null) ?? fallbackLabel;
    return { dataKey: 'x', type: 'number', tick: { fontSize: 10 }, ticks, domain, allowDecimals: false,
      label: { value: label, position: 'insideBottomRight', offset: -4, fontSize: 10 } };
  }

  function ExpTooltip({ active, payload, label, unit = '' }) {
    if (!active || !payload?.length) return null;
    const rows = sortPayloadByValueDesc(payload);
    return (
      <div className="bg-white border border-gray-200 rounded shadow-sm px-3 py-2 text-xs space-y-0.5 max-h-48 overflow-y-auto">
        <p className="font-semibold text-gray-700 mb-1">{label}</p>
        {rows.map((p) => {
          const line = lines.find((l) => l.id === p.dataKey);
          return (
            <p key={p.dataKey} style={{ color: p.stroke }}>
              {line?.name ?? p.dataKey}{line?.group && line.group !== '—' ? ` (${line.group})` : ''}: {p.value}{unit}
            </p>
          );
        })}
      </div>
    );
  }

  return (
    <div className="bg-white border border-gray-200 rounded-xl p-5 mb-6 space-y-5">
      {/* Global controls */}
      <div className="flex items-center flex-wrap gap-3">
        <h2 className="text-sm font-semibold text-gray-700 uppercase tracking-wide">Cross-subject metrics</h2>
        <div className="flex items-center gap-1.5 text-xs">
          <span className="text-gray-400">X axis:</span>
          <select value={xField} onChange={(e) => setXField(e.target.value)} className={SELECT_CLS}>
            <option value="__days__"># Days Reach (ordinal)</option>
            {xAxisFields.map((f) => <option key={f.fieldId} value={f.fieldId}>{f.label}</option>)}
          </select>
        </div>
        {groupableFields.length > 0 && (
          <div className="flex items-center gap-1.5 text-xs">
            <span className="text-gray-400">Color by:</span>
            <select value={groupBy} onChange={(e) => setGroupBy(e.target.value)} className={SELECT_CLS}>
              <option value="">— none —</option>
              {groupableFields.map((f) => <option key={f.fieldId} value={f.fieldId}>{f.label}</option>)}
            </select>
          </div>
        )}
        <TickStepControl value={tickStep} onChange={setTickStep} />
        {hasData && (
          <div className="ml-auto flex items-center gap-2">
            {saveState === 'saving' && <span className="text-xs text-gray-400">Saving…</span>}
            {saveState === 'saved'  && <span className="text-xs text-emerald-500">Saved ✓</span>}
            {saveState === 'error'  && <span className="text-xs text-amber-500">Couldn't save layout</span>}
            <button type="button" onClick={() => setSidebarOpen((o) => !o)}
              className={[
                'text-xs px-2 py-0.5 rounded border transition-colors',
                sidebarOpen
                  ? 'bg-gray-100 border-gray-300 text-gray-700'
                  : 'bg-white border-gray-200 text-gray-500 hover:text-gray-700 hover:border-gray-300',
              ].join(' ')}>
              {sidebarOpen ? '◂ Subjects' : '▸ Subjects'}
            </button>
          </div>
        )}
      </div>

      {!hasData && (
        <p className="text-sm text-gray-400 italic">No sessions with saved metrics yet.</p>
      )}

      {hasData && (
        <div className="flex gap-3 items-start">
          <div className="flex-1 min-w-0 space-y-5">
          {/* Chart 1 — counts */}
          <div>
            <div className="flex items-center justify-between mb-2 flex-wrap gap-2">
              <div className="flex items-center gap-3">
                <p className="text-xs text-gray-500 font-medium">Count per session</p>
                {/* Y-metric selector stays at chart level — primary control */}
                <div className="flex rounded border border-gray-200 overflow-hidden text-[10px]">
                  {METRIC_OPTIONS.map((opt) => (
                    <button key={opt.value} type="button" onClick={() => setMetric(opt.value)}
                      className={`px-2 py-0.5 transition-colors border-l border-gray-200 first:border-l-0
                        ${metric === opt.value ? 'bg-gray-700 text-white' : 'bg-white text-gray-500 hover:bg-gray-50'}`}>
                      {opt.label}
                    </button>
                  ))}
                </div>
              </div>
              <button type="button" onClick={() => countsS.setExpanded((e) => !e)}
                className={[
                  'text-xs px-2 py-0.5 rounded border transition-colors',
                  countsS.expanded
                    ? 'bg-gray-100 border-gray-300 text-gray-700'
                    : countsS.isDirty
                      ? 'bg-indigo-50 border-indigo-200 text-indigo-600'
                      : 'bg-white border-gray-200 text-gray-400 hover:text-gray-700 hover:border-gray-300',
                ].join(' ')}>
                {countsS.expanded ? '▲ collapse' : '▾ customize'}
              </button>
            </div>
            {countsS.expanded && (
              <ChartSettingsBar s={countsS} xAxisFields={xAxisFields} globalXField={xField} />
            )}
            <ResponsiveContainer width="100%" height={countsS.height}>
              <LineChart data={countData} margin={{ top: 4, right: 16, left: 0, bottom: 4 }}>
                <CartesianGrid strokeDasharray="3 3" stroke="#f3f4f6" />
                <XAxis {...xAxisProps(countData, countsTicks, countsS, countsXLabel)} />
                <YAxis tick={{ fontSize: 10 }} width={30} allowDecimals={false} domain={countsS.yDomain} />
                <Tooltip content={<ExpTooltip />} />
                <Legend wrapperStyle={{ fontSize: 11 }}
                  formatter={(value) => lines.find((l) => l.id === value)?.name ?? value}
                  onClick={(o) => toggleHidden(o.dataKey)}
                  onMouseEnter={(o) => setHighlightedSubject(o.dataKey)}
                  onMouseLeave={() => setHighlightedSubject(null)} />
                {lines.map((l) => (
                  <Line key={l.id} type="monotone" dataKey={l.id} name={l.id}
                    stroke={l.color}
                    strokeWidth={highlightedSubject === l.id ? 3.5 : 2}
                    hide={hiddenSubjects.has(l.id)}
                    dot={{ r: 3 }} activeDot={{ r: 5 }}
                    connectNulls isAnimationActive={false} />
                ))}
              </LineChart>
            </ResponsiveContainer>
          </div>

          {/* Chart 2 — success rate */}
          <div>
            <ChartHeader title="Success rate (%) — each line is one subject" isDirty={rateS.isDirty} expanded={rateS.expanded} onToggle={() => rateS.setExpanded((e) => !e)} />
            {rateS.expanded && (
              <ChartSettingsBar s={rateS} xAxisFields={xAxisFields} globalXField={xField} />
            )}
            <ResponsiveContainer width="100%" height={rateS.height}>
              <LineChart data={rateData} margin={{ top: 4, right: 16, left: 0, bottom: 4 }}>
                <CartesianGrid strokeDasharray="3 3" stroke="#f3f4f6" />
                <XAxis {...xAxisProps(rateData, rateTicks, rateS, rateXLabel)} />
                <YAxis tick={{ fontSize: 10 }} width={36}
                  domain={rateS.yDomain[0] !== 'auto' || rateS.yDomain[1] !== 'auto' ? rateS.yDomain : [0, 100]}
                  tickFormatter={(v) => `${v}%`} />
                <Tooltip content={<ExpTooltip unit="%" />} />
                <Legend wrapperStyle={{ fontSize: 11 }}
                  formatter={(value) => lines.find((l) => l.id === value)?.name ?? value}
                  onClick={(o) => toggleHidden(o.dataKey)}
                  onMouseEnter={(o) => setHighlightedSubject(o.dataKey)}
                  onMouseLeave={() => setHighlightedSubject(null)} />
                {lines.map((l) => (
                  <Line key={l.id} type="monotone" dataKey={l.id} name={l.id}
                    stroke={l.color}
                    strokeWidth={highlightedSubject === l.id ? 3.5 : 2}
                    hide={hiddenSubjects.has(l.id)}
                    dot={{ r: 3 }} activeDot={{ r: 5 }}
                    connectNulls isAnimationActive={false} />
                ))}
              </LineChart>
            </ResponsiveContainer>
          </div>
          </div>
          {sidebarOpen && (
            <SubjectSidebar
              lines={lines}
              hiddenSubjects={hiddenSubjects}
              onToggleSubject={toggleHidden}
              onToggleGroup={toggleGroup}
              onShowAll={showAll}
              onHighlight={setHighlightedSubject}
            />
          )}
        </div>
      )}
    </div>
  );
}
