frontend/src/components/StateViews.tsx

import { useMemo, useState } from 'react';
import type { Member, StateInfo, StatePair } from '../lib/api';
import { FIDELITY_INFO, fidelityKey, similarityColor } from '../lib/fidelity';
import { chainLabel, fmtNum } from '../lib/format';
import { Tooltip, useWidth, type TipState } from './charts';
import { Swatch } from './ui';

const MAX_TILES = 40;
const MAX_DOTS = 120;

/**
 * One tile per conformational state; each dot is one observation, coloured
 * by ligand state. Clicking a tile selects the state.
 */
export function StateLandscape({
  states,
  members,
  selected,
  onSelect,
  highlight,
}: {
  states: StateInfo[];
  members: Member[];
  selected: number | null;
  onSelect: (s: number | null) => void;
  highlight?: { pdb_id: string; auth_asym_id: string } | null;
}) {
  const byState = useMemo(() => {
    const m = new Map<number, Member[]>();
    for (const mem of members) {
      const arr = m.get(mem.state_id) ?? [];
      arr.push(mem);
      m.set(mem.state_id, arr);
    }
    return m;
  }, [members]);
  const shown = states.slice(0, MAX_TILES);
  const rest = states.slice(MAX_TILES);
  const restObs = rest.reduce((s, x) => s + x.n_members, 0);

  return (
    <div className="landscape">
      <div className="landscape__grid">
        {shown.map((st) => {
          const mem = byState.get(st.state_id) ?? [];
          const holo = mem.filter((m) => m.binding_status === 'holo').length;
          const isSel = selected === st.state_id;
          const hasHi = highlight && mem.some((m) => m.pdb_id === highlight.pdb_id && m.auth_asym_id === highlight.auth_asym_id);
          return (
            <button
              key={st.state_id}
              type="button"
              className={`state-tile ${isSel ? 'is-selected' : ''} ${hasHi ? 'is-home' : ''}`}
              onClick={() => onSelect(isSel ? null : st.state_id)}
              aria-pressed={isSel}
              title={`State ${st.state_id}: ${st.n_members} observations (${holo} holo)`}
            >
              <span className="state-tile__head">
                <span className="state-tile__name">State {st.state_id}</span>
                <span className="state-tile__n tabular">{st.n_members.toLocaleString()}</span>
              </span>
              <span className="state-tile__dots">
                {mem.slice(0, MAX_DOTS).map((m) => {
                  const isHi = highlight && m.pdb_id === highlight.pdb_id && m.auth_asym_id === highlight.auth_asym_id;
                  return (
                    <span
                      key={`${m.pdb_id}.${m.auth_asym_id}`}
                      className={`state-dot ${m.binding_status === 'holo' ? 'is-holo' : ''} ${isHi ? 'is-hi' : ''}`}
                    />
                  );
                })}
                {mem.length > MAX_DOTS && <span className="state-tile__more">+{mem.length - MAX_DOTS}</span>}
              </span>
            </button>
          );
        })}
      </div>
      <div className="landscape__foot">
        <ul className="legend">
          <li>
            <Swatch color="var(--apo)" shape="dot" /> Apo observation
          </li>
          <li>
            <Swatch color="var(--holo)" shape="dot" /> Holo observation
          </li>
          {highlight && (
            <li>
              <span className="state-dot is-hi legend-dot" /> This chain
            </li>
          )}
        </ul>
        {rest.length > 0 && (
          <span className="muted">
            + {rest.length.toLocaleString()} smaller states ({restObs.toLocaleString()} observations) not shown
          </span>
        )}
      </div>
    </div>
  );
}

/** Mean pair similarity between states (diagonal = within-state). */
export function StateHeatmap({
  states,
  pairs,
  selected,
  onSelect,
}: {
  states: StateInfo[];
  pairs: StatePair[];
  selected: number | null;
  onSelect: (s: number | null) => void;
}) {
  const [ref, width] = useWidth<HTMLDivElement>();
  const [tip, setTip] = useState<TipState | null>(null);
  const [hover, setHover] = useState<[number, number] | null>(null);
  const order = states.slice(0, MAX_TILES).map((s) => s.state_id);
  const lookup = useMemo(() => {
    const m = new Map<string, StatePair>();
    for (const p of pairs) {
      m.set(`${p.a}|${p.b}`, p);
      m.set(`${p.b}|${p.a}`, p);
    }
    return m;
  }, [pairs]);

  const n = order.length;
  const label = 34;
  const cell = n ? Math.max(8, Math.min(34, Math.floor((Math.min(width, 560) - label - 4) / n))) : 0;
  const size = label + n * cell;
  const showLabels = cell >= 14;

  if (n < 2) return null;
  return (
    <div className="heatmap" ref={ref}>
      {width > 0 && (
        <svg width={size} height={size} role="img" aria-label="State similarity matrix">
          {order.map((s, i) => (
            <g key={s}>
              {(showLabels || i % 5 === 0) && (
                <>
                  <text x={label - 6} y={label + i * cell + cell / 2} dy="0.32em" textAnchor="end" className="tick">
                    {s}
                  </text>
                  <text x={label + i * cell + cell / 2} y={label - 8} textAnchor="middle" className="tick">
                    {s}
                  </text>
                </>
              )}
            </g>
          ))}
          {order.map((a, i) =>
            order.map((b, j) => {
              const p = lookup.get(`${a}|${b}`);
              const dim =
                (selected != null && a !== selected && b !== selected) ||
                (hover && hover[0] !== i && hover[1] !== j);
              return (
                <rect
                  key={`${a}-${b}`}
                  x={label + j * cell + 1}
                  y={label + i * cell + 1}
                  width={cell - 2}
                  height={cell - 2}
                  rx={Math.min(3, cell / 6)}
                  fill={similarityColor(p?.similarity)}
                  opacity={dim ? 0.35 : 1}
                  className="heatmap__cell"
                  onClick={() => onSelect(selected === b ? null : b)}
                  onMouseMove={(e) => {
                    const r = (e.currentTarget.ownerSVGElement as SVGSVGElement).getBoundingClientRect();
                    setHover([i, j]);
                    setTip({
                      x: e.clientX - r.left,
                      y: e.clientY - r.top,
                      content: (
                        <>
                          <div className="chart-tip__title">
                            {a === b ? `Within state ${a}` : `State ${a} ↔ state ${b}`}
                          </div>
                          <div className="chart-tip__row">
                            <span>Mean similarity</span>
                            <b className="tabular">{fmtNum(p?.similarity, 3)}</b>
                          </div>
                          {a !== b && p && (
                            <div className="chart-tip__row">
                              <span>Fidelity</span>
                              <b>{FIDELITY_INFO[fidelityKey(p.fidelity)].label}</b>
                            </div>
                          )}
                          <div className="chart-tip__foot">{p ? `${p.n_pairs.toLocaleString()} observation pairs` : 'No pairs'}</div>
                        </>
                      ),
                    });
                  }}
                  onMouseLeave={() => {
                    setTip(null);
                    setHover(null);
                  }}
                />
              );
            }),
          )}
        </svg>
      )}
      <Tooltip tip={tip} />
      <div className="heatmap__scale">
        <span>Mean similarity</span>
        {[
          ['< 0.50', 0.4],
          ['0.50', 0.6],
          ['0.70', 0.8],
          ['0.85', 0.9],
          ['0.95', 0.96],
          ['0.98+', 0.99],
        ].map(([l, v]) => (
          <span key={l as string} className="heatmap__step">
            <Swatch color={similarityColor(v as number)} />
            {l}
          </span>
        ))}
      </div>
    </div>
  );
}

export function memberKey(m: { pdb_id: string; auth_asym_id: string }) {
  return chainLabel(m.pdb_id, m.auth_asym_id);
}