Return after RSI

// What came after an RSI reading on this market: for each band of RSI five points wide, the average
// return over the next bars, up where it rose and down where it fell. RSI is on the graph's own
// axis rather than the chart's time, so it shows in a Script graphs panel, which the legend's
// Graphs button opens. Points mark how many readings each band had, on a second graph.
indicator('Return after RSI', { overlay: true, shortName: 'RSI → return' });

const length = input.int('RSI length', 14, { min: 2 });
const ahead = input.int('Bars ahead', 12, { min: 1, max: 500 });

const chart = bars();
const rsi = ta.rsi(length);
const pastRsi = series(501);
const pastClose = series(501);
// A sum and a count per band.
const sums = state.map<number>(20);
const counts = state.map<number>(20);

const returns = graph(`Average return ${ahead} bars after RSI`, { format: 'percent', precision: 2 });
const average = returns.bars('Return');
const samples = graph('Readings per RSI band', { precision: 0 });
const count = samples.points('Readings', { size: 4 });

chart.onBar(() => {
  pastRsi.set(rsi.update(chart.close));
  pastClose.set(chart.close);
  // The reading `ahead` bars ago, and what price did since.
  const then = pastRsi.ago(ahead);
  const from = pastClose.ago(ahead);
  if (isNa(then) || isNa(from)) return;
  const band = Math.min(19, Math.floor(then / 5));
  const key = String(band);
  const change = (chart.close / from - 1) * 100;
  const n = (counts.get(key) ?? 0) + 1;
  counts.set(key, n);
  sums.set(key, (sums.get(key) ?? 0) + change);
  const mean = sums.get(key)! / n;
  // Each band at its middle: 32.5 for RSI 30 to 35.
  const middle = band * 5 + 2.5;
  average.set(middle, mean, mean >= 0 ? color.up : color.down);
  count.set(middle, n);
});