frontend/src/components/StructureDiff.tsx

import { useMemo, useState } from 'react';
import { api, type CompareClass, type StructureCompare } from '../lib/api';
import { useApi } from '../lib/hooks';
import { chainLabel, fmtInt, fmtNum } from '../lib/format';
import { CustomStructureViewer, type ResidueColorRun, type ResidueTooltip } from './StructureViewer';
import { ErrorState, Segmented, Spinner, Swatch } from './ui';

/* Cα deviation grades and the "resolved in one structure only" grey. */
export const DIFF_COLORS: Record<CompareClass, string> = {
  low: '#2fa84f',
  medium: '#f2b705',
  high: '#d93b3b',
  only_a: '#b6bcc6',
  only_b: '#b6bcc6',
  missing: 'transparent',
};

/* In the superposed "both" view the alternative keeps the same colours but is drawn semi-transparent. */
export const ALT_OPACITY = 0.35;

type Show = 'both' | 'A' | 'B';

function runsFor(residues: StructureCompare['residues'], chain: 'A' | 'B', opacity?: number): ResidueColorRun[] {
  const absent: CompareClass = chain === 'A' ? 'only_b' : 'only_a';
  const runs: ResidueColorRun[] = [];
  for (const r of residues) {
    if (r.class === absent || r.class === 'missing') continue;
    const color = DIFF_COLORS[r.class];
    const last = runs[runs.length - 1];
    if (last && last.color === color && last.end === r.pos - 1) last.end = r.pos;
    else runs.push({ chain, start: r.pos, end: r.pos, color, opacity });
  }
  return runs;
}

function describe(r: StructureCompare['residues'][number], la: string, lb: string): string {
  const head = `${r.res} ${r.pos}`;
  if (r.class === 'only_a') return `${head}: resolved only in ${la}`;
  if (r.class === 'only_b') return `${head}: resolved only in ${lb}`;
  return `${head}: Cα deviation ${fmtNum(r.dev, 2)} Å`;
}

/**
 * Superposed view of two observations of one sequence, coloured per residue by
 * Cα deviation (green / yellow / red) with order changes in grey, plus a
 * sequence track of the same classes.
 */
export function StructureDiff({
  a,
  b,
  length,
}: {
  a: { pdb_id: string; auth_asym_id: string };
  b: { pdb_id: string; auth_asym_id: string };
  length: number;
}) {
  const [show, setShow] = useState<Show>('both');
  const la = chainLabel(a.pdb_id, a.auth_asym_id);
  const lb = chainLabel(b.pdb_id, b.auth_asym_id);
  const res = useApi(`compare:${la}:${lb}`, (sig) =>
    api.compare(a.pdb_id, a.auth_asym_id, b.pdb_id, b.auth_asym_id, sig),
  );
  const d = res.data;

  const runs = useMemo(() => {
    if (!d) return [];
    if (show === 'both') return [...runsFor(d.residues, 'A'), ...runsFor(d.residues, 'B', ALT_OPACITY)];
    return runsFor(d.residues, show);
  }, [d, show]);

  // hover text names the real chains (the file calls them A and B)
  const tooltips = useMemo<ResidueTooltip[]>(() => {
    if (!d) return [];
    const out: ResidueTooltip[] = [];
    for (const r of d.residues) {
      const what =
        r.class === 'only_a' || r.class === 'only_b' ? 'ordered in one structure only' : r.dev != null ? `Cα Δ ${fmtNum(r.dev, 2)} Å` : '';
      if (show !== 'B' && r.class !== 'only_b' && r.class !== 'missing')
        out.push({ chain: 'A', pos: r.pos, text: `<b>${la}</b> (this) · ${r.res} ${r.pos} · ${what}` });
      if (show !== 'A' && r.class !== 'only_a' && r.class !== 'missing')
        out.push({ chain: 'B', pos: r.pos, text: `<b>${lb}</b> (alternative) · ${r.res} ${r.pos} · ${what}` });
    }
    return out;
  }, [d, show, la, lb]);

  const url = api.compareStructureUrl(a.pdb_id, a.auth_asym_id, b.pdb_id, b.auth_asym_id, show === 'both' ? undefined : show);
  const L = Math.max(length, d?.residues[d.residues.length - 1]?.pos ?? 0);
  const byPos = useMemo(() => new Map((d?.residues ?? []).map((r) => [r.pos, r])), [d]);

  if (res.loading) return <Spinner label="Fetching coordinates and superposing…" />;
  if (res.error) return <ErrorState title="Could not compare the structures" message={res.error.message} />;
  if (!d) return null;

  const t = d.thresholds;
  return (
    <div className="sdiff">
      <div className="sdiff__bar">
        <Segmented<Show>
          label="Chains shown"
          value={show}
          onChange={setShow}
          options={[
            { value: 'both', label: 'Both' },
            { value: 'A', label: `${la} (this)` },
            { value: 'B', label: `${lb} (alternative)` },
          ]}
        />
        <p className="muted small sdiff__stats">
          {show === 'both' ? <>{la} is solid and {lb} semi-transparent, both coloured by deviation. </> : null}
          {lb} superposed onto {la} on their largest rigid core: {fmtInt(d.counts.low)} of {fmtInt(d.n_common)} shared
          residues within {t.low} Å (core RMSD {fmtNum(d.rmsd_core, 2)} Å), maximum deviation {fmtNum(d.max_dev, 1)} Å.
        </p>
      </div>
      <CustomStructureViewer
        url={url}
        runs={runs}
        tooltips={tooltips}
        height={460}
        label={show === 'both' ? `chain A = ${la} · chain B = ${lb}` : show === 'A' ? la : lb}
      />
      <ul className="legend sdiff__legend">
        {show === 'both' && (
          <li>
            <Swatch color={DIFF_COLORS.low} /> solid: {la}
            <span style={{ opacity: ALT_OPACITY + 0.15 }}>
              <Swatch color={DIFF_COLORS.low} />
            </span>{' '}
            faded: {lb}
          </li>
        )}
        <li>
          <Swatch color={DIFF_COLORS.low} /> &lt; {t.low} Å <b className="tabular">{fmtInt(d.counts.low)}</b>
        </li>
        <li>
          <Swatch color={DIFF_COLORS.medium} /> {t.low}–{t.high} Å <b className="tabular">{fmtInt(d.counts.medium)}</b>
        </li>
        <li>
          <Swatch color={DIFF_COLORS.high} /> ≥ {t.high} Å <b className="tabular">{fmtInt(d.counts.high)}</b>
        </li>
        <li>
          <Swatch color={DIFF_COLORS.only_a} /> Ordered in only one structure{' '}
          <b className="tabular">
            {fmtInt(d.counts.only_a)} / {fmtInt(d.counts.only_b)}
          </b>
        </li>
      </ul>
      <div className="sdiff__track" aria-label="Per-residue comparison along the sequence">
        {Array.from({ length: L }, (_, i) => {
          const r = byPos.get(i + 1);
          return (
            <span
              key={i}
              className={`sdiff__cell ${r ? '' : 'is-empty'}`}
              style={r ? { background: DIFF_COLORS[r.class] } : undefined}
              title={r ? describe(r, la, lb) : `${i + 1}: not resolved in either structure`}
            />
          );
        })}
      </div>
      <div className="sdiff__axis">
        <span>1</span>
        <span className="muted small">
          grey: {fmtInt(d.counts.only_a)} residues ordered only in {la}, {fmtInt(d.counts.only_b)} only in {lb} · blank:
          disordered in both
        </span>
        <span>{L}</span>
      </div>
    </div>
  );
}