frontend/src/components/KnowledgeGraph.tsx

import { useEffect, useMemo, useState } from 'react';
import { Link, useNavigate } from 'react-router-dom';
import { ArrowRight } from 'lucide-react';
import { api, type InnerGraph, type OuterGraph } from '../lib/api';
import { prefetch, useApi } from '../lib/hooks';
import { FIDELITY_INFO, fidelityColor } from '../lib/fidelity';
import { chainLabel, chainPath, cleanFunction, fmtInt, fmtNum, methodShort, sequencePath } from '../lib/format';
import { ForceGraph, type FGGroup, type FGLink, type FGNode } from './ForceGraph';
import { ErrorState, Spinner, Swatch } from './ui';

/* Categorical slots (validated order) for the largest states; the rest fold into "other". */
const STATE_COLORS = ['#2a78d6', '#eb6834', '#1baf7a', '#eda100', '#e87ba4', '#008300', '#4a3aa7'];
const OTHER = '#98a2b3';

/* Ordinal ramp for "number of states" on outer-graph nodes. */
const STATE_BINS: Array<{ max: number; label: string; color: string }> = [
  { max: 1, label: '1 state', color: 'var(--nst-1)' },
  { max: 2, label: '2', color: 'var(--nst-2)' },
  { max: 5, label: '3–5', color: 'var(--nst-3)' },
  { max: 20, label: '6–20', color: 'var(--nst-4)' },
  { max: Infinity, label: '> 20', color: 'var(--nst-5)' },
];
const statesColor = (n: number) => STATE_BINS.find((b) => n <= b.max)!.color;

/**
 * Thin a dense complete graph without hiding transitions: keep each node's k most
 * similar same-state partners, and the `perPair` most similar edges between every
 * pair of states (so every state pair stays visibly connected).
 */
function backbone(edges: InnerGraph['edges'], state: number[], k = 2, perPair = 3): Set<number> {
  const within = new Map<number, number[]>();
  const across = new Map<string, number[]>();
  edges.forEach((e, i) => {
    const a = state[e.s];
    const b = state[e.t];
    if (a === b) {
      for (const n of [e.s, e.t]) {
        const arr = within.get(n) ?? [];
        arr.push(i);
        within.set(n, arr);
      }
    } else {
      const key = a < b ? `${a}|${b}` : `${b}|${a}`;
      const arr = across.get(key) ?? [];
      arr.push(i);
      across.set(key, arr);
    }
  });
  const bySim = (x: number, y: number) => (edges[y].sim ?? -1) - (edges[x].sim ?? -1);
  const keep = new Set<number>();
  for (const arr of within.values()) arr.sort(bySim).slice(0, k).forEach((i) => keep.add(i));
  for (const arr of across.values()) arr.sort(bySim).slice(0, perPair).forEach((i) => keep.add(i));
  return keep;
}

/* ───────────────────────────── outer layer: one ECOD homology group */

function OuterGraphView({
  data,
  selected,
  onSelect,
  names,
}: {
  data: OuterGraph;
  selected: string;
  onSelect: (id: string) => void;
  names: Record<string, string>;
}) {
  const group = data.groups[0];
  const { nodes, links } = useMemo(() => {
    const seqs = group?.sequences ?? [];
    const nodes: FGNode[] = seqs.map((s) => ({
      id: s.sequence_id,
      r: 5 + 2.1 * Math.log2(1 + s.n_obs),
      fill: statesColor(s.n_states),
    }));
    const links: FGLink[] = [];
    for (let i = 0; i < seqs.length; i++)
      for (let j = i + 1; j < seqs.length; j++)
        links.push({ source: seqs[i].sequence_id, target: seqs[j].sequence_id, distance: 140, color: 'var(--faint)', opacity: 0.3 });
    return { nodes, links };
  }, [group]);
  const withRing = useMemo(() => nodes.map((n) => ({ ...n, ring: n.id === selected })), [nodes, selected]);
  const bySeq = useMemo(() => new Map((group?.sequences ?? []).map((s) => [s.sequence_id, s])), [group]);
  if (!group) return <p className="muted">This sequence has no ECOD homology assignment.</p>;

  return (
    <ForceGraph
      ariaLabel={`Outer graph: ${group.label}`}
      nodes={withRing}
      links={links}
      charge={-220}
      linkStrength={0.035}
      onNodeClick={onSelect}
      nodeTip={(id) => {
        const s = bySeq.get(id)!;
        return (
          <>
            <div className="chart-tip__title">
              {names[id] ?? <span className="mono">{id}</span>}
              {names[id] && <span className="muted mono"> · {id}</span>}
            </div>
            {s.top_function && <div className="chart-tip__wrap">{cleanFunction(s.top_function).slice(0, 90)}</div>}
            <div className="chart-tip__row">
              <span>Observations</span>
              <b className="tabular">{fmtInt(s.n_obs)}</b>
            </div>
            <div className="chart-tip__row">
              <span>States</span>
              <b className="tabular">{fmtInt(s.n_states)}</b>
            </div>
            {s.uniprot_id && (
              <div className="chart-tip__row">
                <span>UniProt</span>
                <b className="mono">{s.uniprot_id}</b>
              </div>
            )}
            <div className="chart-tip__foot">Click to open its inner graph · drag to move</div>
          </>
        );
      }}
    />
  );
}

/* ───────────────────────────── inner layer: observations, states as shaded circles */

function InnerGraphView({ g }: { g: InnerGraph }) {
  const navigate = useNavigate();
  const { nodes, links, groups, colorOf } = useMemo(() => {
    const top = g.states.slice(0, STATE_COLORS.length);
    const colorMap = new Map(top.map((s, i) => [s.state_id, STATE_COLORS[i]]));
    const colorOf = (st: number) => colorMap.get(st) ?? OTHER;
    const nodes: FGNode[] = g.nodes.map((n, i) => ({
      id: String(i),
      r: 6,
      fill: colorOf(n.state_id),
      group: n.state_id,
    }));
    const keep = g.nodes.length > 40 ? backbone(g.edges, g.nodes.map((n) => n.state_id)) : null;
    const links: FGLink[] = g.edges.map((e, i) => {
      const same = g.nodes[e.s].state_id === g.nodes[e.t].state_id;
      return {
        source: String(e.s),
        target: String(e.t),
        // spring length grows with structural dissimilarity
        distance: Math.min(230, 22 + 560 * (1 - (e.sim ?? 0.3))),
        color: same ? 'var(--line-strong)' : fidelityColor(e.fid),
        opacity: same ? 0.45 : keep ? 0.75 : 0.3,
        width: same ? 0.9 : keep ? 1.4 : 0.9,
        hidden: keep ? !keep.has(i) : false,
      };
    });
    const present = new Set(g.nodes.map((n) => n.state_id));
    const groups: FGGroup[] = g.states
      .filter((s) => present.has(s.state_id))
      .map((s) => ({
        key: s.state_id,
        color: colorOf(s.state_id),
        label: g.nodes.filter((n) => n.state_id === s.state_id).length > 1 ? `State ${s.state_id}` : undefined,
      }));
    return { nodes, links, groups, colorOf };
  }, [g]);

  return (
    <>
      <ForceGraph
        ariaLabel="Inner graph: observations of one sequence grouped by conformational state"
        nodes={nodes}
        links={links}
        groups={groups}
        charge={-55}
        linkStrength={0.03}
        onNodeClick={(id) => {
          const n = g.nodes[Number(id)];
          navigate(chainPath(n.pdb_id, n.auth_asym_id));
        }}
        nodeTip={(id) => {
          const n = g.nodes[Number(id)];
          return (
            <>
              <div className="chart-tip__title mono">{chainLabel(n.pdb_id, n.auth_asym_id)}</div>
              <div className="chart-tip__row">
                <Swatch color={colorOf(n.state_id)} shape="dot" />
                <span>State</span>
                <b>{n.state_id}</b>
              </div>
              <div className="chart-tip__row">
                <span>Method</span>
                <b>
                  {methodShort(n.experimental_method)}
                  {n.resolution != null ? ` · ${fmtNum(n.resolution, 2)} Å` : ''}
                </b>
              </div>
              <div className="chart-tip__row">
                <span>Ligand</span>
                <b>{n.binding_status ?? '—'}</b>
              </div>
              <div className="chart-tip__foot">Click to open this chain · drag to move</div>
            </>
          );
        }}
      />
      <ul className="legend">
        {g.states.slice(0, STATE_COLORS.length).map((s) => (
          <li key={s.state_id}>
            <Swatch color={colorOf(s.state_id)} shape="dot" /> State {s.state_id}
            <span className="muted">({fmtInt(s.n_members)})</span>
          </li>
        ))}

      </ul>
      <ul className="legend">
        <li>
          <span className="legend-line" style={{ background: 'var(--line-strong)' }} /> Same state
        </li>
        {['low', 'medium', 'high'].map((f) => (
          <li key={f}>
            <span className="legend-line" style={{ background: fidelityColor(f) }} /> {FIDELITY_INFO[f].label} fidelity
          </li>
        ))}
      </ul>
    </>
  );
}

/* ───────────────────────────── container */

export function KnowledgeGraph({ featured }: { featured: Array<{ id: string; name: string }> }) {
  const names = Object.fromEntries(featured.map((f) => [f.id, f.name]));
  const [example, setExample] = useState(featured[0].id);
  const [selected, setSelected] = useState(featured[0].id);
  const outer = useApi(`outer1:${example}`, (sig) => api.outerGraph([example], sig));
  const inner = useApi(`inner:${selected}`, (sig) => api.innerGraph(selected, 60, sig));
  const group = outer.data?.groups[0];
  const g = inner.data;

  // once the first graph is up, warm the other examples so switching is instant
  const firstReady = !!inner.data;
  useEffect(() => {
    if (!firstReady) return;
    const t = setTimeout(() => {
      for (const f of featured) {
        prefetch(`inner:${f.id}`, () => api.innerGraph(f.id, 60));
        prefetch(`outer1:${f.id}`, () => api.outerGraph([f.id]));
      }
    }, 800);
    return () => clearTimeout(t);
    // eslint-disable-next-line react-hooks/exhaustive-deps
  }, [firstReady]);
  const selectedSeq = group?.sequences.find((s) => s.sequence_id === selected) ?? g?.sequence;

  return (
    <div className="kg">
      <div className="kg__picks" role="group" aria-label="Examples">
        <span className="muted small">Example</span>
        {featured.map((f) => (
          <button
            key={f.id}
            type="button"
            className={`chip ${f.id === example ? 'is-on' : ''}`}
            onClick={() => {
              setExample(f.id);
              setSelected(f.id);
            }}
          >
            {f.name}
          </button>
        ))}
      </div>
      <div className="kg__panels">
        <div className="card kg__panel">
          <header className="kg__head">
            <p className="eyebrow">Outer layer · sequences</p>
            <h3 className="card__title">{group ? group.label || group.hset.join(' + ') : 'ECOD homology group'}</h3>
            <p className="card__kicker">
              {group
                ? `${fmtInt(group.n_sequences)} sequences share this set of ECOD homology groups and form one clique; the ${group.sequences.length} most multistate are shown.`
                : 'Sequences sharing the same ECOD homology groups are linked.'}{' '}
              Node size: observations · colour: states.
            </p>
          </header>
          {outer.loading && <Spinner />}
          {outer.error && <ErrorState title="Could not load the sequence graph" message={outer.error.message} />}
          {outer.data && <OuterGraphView data={outer.data} selected={selected} onSelect={setSelected} names={names} />}
          <ul className="legend">
            {STATE_BINS.map((b) => (
              <li key={b.label}>
                <Swatch color={b.color} shape="dot" /> {b.label}
              </li>
            ))}
            <li>
              <span className="swatch swatch--dot swatch--sel" /> Selected
            </li>
          </ul>
          {group && (
            <div className="kg__foot">
              <Link to={`/browse?ecod=${encodeURIComponent(group.hset.join(';'))}`} className="btn btn--ghost btn--sm">
                All {fmtInt(group.n_sequences)} sequences <ArrowRight aria-hidden />
              </Link>
            </div>
          )}
        </div>

        <div className="card kg__panel">
          <header className="kg__head">
            <p className="eyebrow">Inner layer · observations</p>
            <h3 className="card__title">
              {names[selected] ?? (g?.sequence.uniprot_id ? `UniProt ${g.sequence.uniprot_id}` : 'Sequence')}{' '}
              <span className="muted mono">{selected}</span>
            </h3>
            <p className="card__kicker">
              Every observation is linked to every other; spring length follows structural dissimilarity. Shaded circles
              are conformational states; cross-state edges are coloured by fidelity.
            </p>
          </header>
          {inner.loading && <Spinner label="Loading inner graph…" />}
          {inner.error && <ErrorState title="Could not load the observation graph" message={inner.error.message} />}
          {g && <InnerGraphView g={g} />}
          {g && (
            <div className="kg__foot">
              <span className="muted small">
                {g.nodes.length} of {fmtInt(g.n_total)} observations
                {g.states.length > STATE_COLORS.length
                  ? ` · ${STATE_COLORS.length} largest of ${fmtInt(g.states.length)} states`
                  : ` · ${fmtInt(g.states.length)} states`}
                {g.nodes.length > 40 && ' · showing the 3 closest links between each pair of states; hover a node for all its triplets'}
              </span>
              <Link to={sequencePath(selected)} className="btn btn--ghost btn--sm">
                Open {selectedSeq ? `${fmtInt(selectedSeq.n_obs)} observations` : 'sequence'} <ArrowRight aria-hidden />
              </Link>
            </div>
          )}
        </div>
      </div>
    </div>
  );
}