PIG-Demo/src/components/demo/RewardEditor.tsx

import { useMemo, useState } from 'react';
import { ArrowDown, ArrowUp, Minus, RotateCcw } from 'lucide-react';
import {
  isEdited,
  pruneOverrides,
  reweight,
  weightsAreDegenerate,
} from '@/lib/demo-kit/reward';
import type { WeightOverrides } from '@/lib/demo-kit/reward';
import { rewardTotal } from '@/lib/demo-kit/episode';
import type { RewardSpec, RewardValues } from '@/lib/demo-kit/types';
import { Button } from '@/components/ui/button';
import { Slider } from '@/components/ui/slider';
import { cn } from '@/lib/utils';
import { formatNumber, formatOrDash } from './format';
import { EditedChip } from './StatStrip';

/** One recorded arm — a run, or a group of runs already reduced to one score. */
export interface RewardArm {
  id: string;
  label: string;
  /** The environment's per-component scores. Re-weighted, never re-run. */
  values: RewardValues;
  note?: string;
}

export interface RewardPreset {
  id: string;
  label: string;
  description?: string;
  /** Raw weights, before normalisation. Keyed by component. */
  weights: Record<string, number>;
}

export interface RewardEditorProps {
  spec: RewardSpec;
  arms: RewardArm[];
  /** Two is the right number. More and the visitor reads instead of playing. */
  presets?: RewardPreset[];
  onWeightsChange?: (overrides: WeightOverrides, edited: boolean) => void;
  className?: string;
}

const STEP = 0.05;

/**
 * Presets built from the spec's own labels, for a demo that does not supply its
 * own. Each is phrased as a position someone would actually argue for in a
 * meeting, not as "preset A" and "preset B".
 */
function derivePresets(spec: RewardSpec): RewardPreset[] {
  const shipped: Record<string, number> = {};
  for (const component of spec.components) shipped[component.key] = component.weight;

  const counterweights = spec.components.filter((c) => c.role === 'counterweight');
  const objectives = spec.components.filter((c) => c.role === 'objective');
  const presets: RewardPreset[] = [
    {
      id: 'shipped',
      label: 'What we ship',
      description: 'The weights in the environment as committed.',
      weights: shipped,
    },
  ];

  const firstObjective = objectives[0];
  const firstCounterweight = counterweights[0];
  if (firstCounterweight && firstObjective) {
    const objectiveOnly = { ...shipped };
    for (const component of counterweights) objectiveOnly[component.key] = 0;
    presets.push({
      id: 'objective-only',
      label: `${firstObjective.label} at any cost`,
      description: `Drops ${counterweights.map((c) => c.label.toLowerCase()).join(' and ')} to zero.`,
      weights: objectiveOnly,
    });

    const doubled = { ...shipped };
    for (const component of counterweights) doubled[component.key] = component.weight * 2;
    presets.push({
      id: 'counterweight-heavy',
      // Quoted, because a component label is a phrase written for a table cell
      // ("Found it early") and reads as gibberish spliced into a sentence.
      label: `Twice as much "${firstCounterweight.label}"`,
      description: 'What a risk-averse buyer would ask for.',
      weights: doubled,
    });
  }
  return presets;
}

function sameWeights(spec: RewardSpec, overrides: WeightOverrides, preset: RewardPreset): boolean {
  return spec.components.every((component) => {
    const current = overrides[component.key] ?? component.weight;
    const target = preset.weights[component.key] ?? component.weight;
    return Math.abs(current - target) < 1e-9;
  });
}

/**
 * Change what "good" means and watch the ranking move.
 *
 * The honesty problem this component has to solve: re-weighting recorded scores
 * is NOT training. It shows the ranking a different reward would have produced
 * over these exact attempts; it cannot show the different attempts a model
 * trained on that reward would have made. That distinction is the permanent
 * caption at the bottom, and it is not collapsible.
 *
 * Weights are normalised to sum to 1 before scoring — `reweight` does it — so
 * dragging one slider up trades weight away from the others instead of lifting
 * every arm at once. Without that, the totals all rise together and the ranking
 * appears to move when only the scale did.
 */
export function RewardEditor({
  spec,
  arms,
  presets,
  onWeightsChange,
  className,
}: RewardEditorProps) {
  const [overrides, setOverrides] = useState<WeightOverrides>({});
  const effectivePresets = useMemo(() => presets ?? derivePresets(spec), [presets, spec]);

  const edited = isEdited(overrides, spec.components);
  const inForce = useMemo(() => reweight(spec.components, overrides), [spec, overrides]);
  const shippedNormalised = useMemo(() => reweight(spec.components, {}), [spec]);
  const degenerate = weightsAreDegenerate(inForce);

  const max = useMemo(
    () => Math.max(2, ...spec.components.map((c) => Math.ceil(Math.abs(c.weight) * 2))),
    [spec],
  );

  const apply = (next: WeightOverrides) => {
    // Pruned before it goes into state: an override that equals the shipped
    // weight is not an edit, and leaving it in makes the "edited" chip stick
    // after the visitor drags a slider back where it started.
    const pruned = pruneOverrides(next, spec.components);
    setOverrides(pruned);
    onWeightsChange?.(pruned, isEdited(pruned, spec.components));
  };

  const ranked = useMemo(() => {
    const rows = arms.map((arm) => ({
      arm,
      total: degenerate ? null : rewardTotal(arm.values, inForce),
      shippedTotal: rewardTotal(arm.values, shippedNormalised),
    }));
    // Nulls sort last: an unscored arm is not a zero-scoring arm.
    const byTotal = (a: { total: number | null }, b: { total: number | null }) => {
      if (a.total === null && b.total === null) return 0;
      if (a.total === null) return 1;
      if (b.total === null) return -1;
      return b.total - a.total;
    };
    const shippedOrder = [...rows]
      .sort((a, b) => byTotal({ total: a.shippedTotal }, { total: b.shippedTotal }))
      .map((row) => row.arm.id);
    return [...rows].sort(byTotal).map((row, index) => ({
      ...row,
      rank: index + 1,
      shippedRank: shippedOrder.indexOf(row.arm.id) + 1,
    }));
  }, [arms, inForce, shippedNormalised, degenerate]);

  const span = useMemo(() => {
    const totals = ranked.map((row) => row.total).filter((total): total is number => total !== null);
    if (totals.length === 0) return { lo: 0, hi: 1 };
    const lo = Math.min(0, ...totals);
    const hi = Math.max(...totals);
    return { lo, hi: hi === lo ? lo + 1 : hi };
  }, [ranked]);

  const leader = ranked[0];

  return (
    <div className={cn('grid grid-cols-1 gap-4 lg:grid-cols-2', className)}>
      <section aria-label="Reward weights" className="card p-3">
        <div className="flex flex-wrap items-center gap-2">
          <h3 className="text-sm font-semibold">Change what good means</h3>
          {edited ? <EditedChip /> : null}
          <Button
            variant="outline"
            size="sm"
            className="tap ml-auto"
            onClick={() => apply({})}
            disabled={!edited}
          >
            <RotateCcw aria-hidden="true" />
            Reset to shipped
          </Button>
        </div>

        <div className="mt-2 flex flex-wrap gap-2">
          {effectivePresets.map((preset) => {
            const active = sameWeights(spec, overrides, preset);
            return (
              <Button
                key={preset.id}
                variant={active ? 'subtle' : 'outline'}
                size="touch"
                aria-pressed={active}
                title={preset.description ?? preset.label}
                onClick={() => apply({ ...preset.weights })}
                className="text-xs"
              >
                {preset.label}
              </Button>
            );
          })}
        </div>

        <div className="mt-4 space-y-4">
          {spec.components.map((component, index) => {
            const raw = overrides[component.key] ?? component.weight;
            const changed = Math.abs(raw - component.weight) > 1e-9;
            const share = inForce[index]?.weight ?? 0;
            return (
              <div key={component.key}>
                <div className="flex items-baseline justify-between gap-2">
                  {/* A <label htmlFor> would point at Radix's root <span>, which
                      is not a labelable element — the association would silently
                      do nothing. The thumb takes its name from this text. */}
                  <span id={`weight-label-${component.key}`} className="text-sm font-medium text-fg">
                    {component.label}
                  </span>
                  <span
                    className={cn(
                      'nums font-mono text-sm',
                      changed ? 'font-semibold text-accent-fg' : 'text-muted',
                    )}
                  >
                    {formatNumber(raw, 2)}
                    <span className="ml-1.5 text-xs font-normal text-muted">
                      {degenerate ? '—' : `${Math.round(share * 100)}%`}
                    </span>
                  </span>
                </div>
                <p className="mb-1 text-xs leading-snug text-muted">{component.description}</p>
                <Slider
                  min={0}
                  max={max}
                  step={STEP}
                  value={[raw]}
                  onValueChange={(next) =>
                    apply({ ...overrides, [component.key]: next[0] ?? component.weight })
                  }
                  aria-labelledby={`weight-label-${component.key}`}
                />
              </div>
            );
          })}
        </div>

        <p className="mt-3 text-xs leading-relaxed text-muted">
          The percentage is the share of the reward each term carries once the weights are
          normalised. Raising one lowers the others — that is the trade a reward designer
          actually makes.
        </p>
      </section>

      <section aria-label="Ranking under this reward" className="card flex flex-col p-3">
        <h3 className="flex items-center gap-2 text-sm font-semibold">
          Ranking under this reward
          {edited ? <EditedChip /> : null}
        </h3>

        {degenerate ? (
          <p className="mt-3 text-sm leading-relaxed text-muted">
            Every weight is zero, so there is no reward left to rank by. That is not a score of
            nought — it is a reward that expresses no preference at all.
          </p>
        ) : (
          <ol className="mt-3 space-y-2">
            {ranked.map((row) => {
              const moved = row.rank - row.shippedRank;
              const width =
                row.total === null
                  ? 0
                  : Math.max(2, ((row.total - span.lo) / (span.hi - span.lo)) * 100);
              return (
                <li
                  key={row.arm.id}
                  className={cn(
                    'rounded-lg border p-2.5 transition-colors duration-3 ease-enter',
                    row.rank === 1 ? 'border-brand bg-accent-subtle/50' : 'border-border',
                  )}
                >
                  <div className="flex items-baseline gap-2">
                    <span className="nums text-sm font-semibold text-muted">{row.rank}</span>
                    <span className="min-w-0 flex-1 truncate text-sm font-medium">
                      {row.arm.label}
                    </span>
                    <RankMove moved={moved} />
                    <span className="nums font-mono text-sm font-semibold">
                      {formatOrDash(row.total)}
                    </span>
                  </div>
                  <div
                    aria-hidden="true"
                    className="mt-1.5 h-1.5 w-full overflow-hidden rounded-full bg-surface-2"
                  >
                    <div
                      className="h-full rounded-full bg-brand transition-[width] duration-3 ease-enter"
                      style={{ width: `${width}%` }}
                    />
                  </div>
                  {row.arm.note ? <p className="mt-1 text-xs text-muted">{row.arm.note}</p> : null}
                </li>
              );
            })}
          </ol>
        )}

        {/* Safe to announce here without fighting the shell's step-change
            region: a slider drag and a step advance never fire from one action. */}
        <p role="status" className="sr-only">
          {leader && leader.total !== null
            ? `Leading under this reward: ${leader.arm.label}, ${formatOrDash(leader.total)}.`
            : ''}
        </p>

        <p className="mt-auto pt-3 text-xs leading-relaxed text-muted">
          We re-scored the same recorded attempts under your reward. Training on it would change
          the behaviour, not just the ranking.
        </p>
      </section>
    </div>
  );
}

function RankMove({ moved }: { moved: number }) {
  if (moved === 0) {
    return (
      <span className="inline-flex items-center text-muted" title="Same rank as the shipped reward">
        <Minus className="size-3.5" aria-hidden="true" />
        <span className="sr-only">rank unchanged</span>
      </span>
    );
  }
  const up = moved < 0;
  const places = Math.abs(moved);
  return (
    <span
      className={cn('nums inline-flex items-center text-xs', up ? 'text-positive' : 'text-danger')}
      title={`${places} place${places === 1 ? '' : 's'} ${up ? 'up' : 'down'} from the shipped reward`}
    >
      {up ? <ArrowUp className="size-3.5" aria-hidden="true" /> : <ArrowDown className="size-3.5" aria-hidden="true" />}
      {places}
      <span className="sr-only">{up ? ' places up' : ' places down'}</span>
    </span>
  );
}