pi-ext

import { writeFile } from "node:fs/promises";
import { Session } from "node:inspector/promises";
import { arch, platform } from "node:os";
import { performance } from "node:perf_hooks";
import { mock } from "node:test";
import { pathToFileURL } from "node:url";
import type { Theme } from "@earendil-works/pi-coding-agent";
import type { TUI } from "@earendil-works/pi-tui";
import { type AfkWidget, createAfkWidget } from "../src/widget.ts";

const TICKS = 1_024;
const SAMPLES = 31;
const WARMUPS = 3;
const BOOTSTRAPS = 10_000;
type Factory = typeof createAfkWidget;
type Sample = {
	wallMs: number;
	cpuMs: number;
	requests: number;
	redundant: number;
};
type Pair = { baseline: Sample; candidate: Sample };
let sink = 0;

const theme = {
	fg: (_color: string, text: string) => `\u001b[36m${text}\u001b[0m`,
} as unknown as Theme;

function simulate(factory: Factory, napping: boolean): Sample {
	mock.timers.enable({ apis: ["setInterval", "setTimeout"] });
	let widget: AfkWidget | undefined;
	let requests = 0;
	let redundant = 0;
	let previous = "";
	try {
		widget = factory(
			{
				requestRender() {
					requests += 1;
					if (!widget)
						throw new Error("Widget requested render before mounting");
					const rendered = widget.render(80).join("\n");
					if (rendered === previous) redundant += 1;
					previous = rendered;
					sink ^= rendered.length;
				},
			} as unknown as TUI,
			theme,
			() => napping,
		);
		previous = widget.render(80).join("\n");
		widget.invalidate();
		requests = 0;
		redundant = 0;
		const cpu = process.cpuUsage();
		const start = performance.now();
		for (let tick = 0; tick < TICKS; tick++) mock.timers.tick(700);
		const wallMs = performance.now() - start;
		const usage = process.cpuUsage(cpu);
		return {
			wallMs,
			cpuMs: (usage.user + usage.system) / 1_000,
			requests,
			redundant,
		};
	} finally {
		widget?.dispose();
		mock.timers.reset();
	}
}

function quantile(values: number[], fraction: number): number {
	const sorted = [...values].sort((a, b) => a - b);
	const index = (sorted.length - 1) * fraction;
	const lower = Math.floor(index);
	const low = sorted[lower];
	const high = sorted[Math.ceil(index)];
	if (low === undefined || high === undefined)
		throw new Error("Empty quantile sample");
	return low + (high - low) * (index - lower);
}
function mean(values: number[]): number {
	return values.reduce((sum, value) => sum + value, 0) / values.length;
}
function summarize(pairs: Pair[]) {
	const bounds = (key: "baseline" | "candidate"): [number, number] => {
		const times = pairs.map((pair) => pair[key].wallMs);
		const q1 = quantile(times, 0.25);
		const q3 = quantile(times, 0.75);
		return [q1 - 1.5 * (q3 - q1), q3 + 1.5 * (q3 - q1)];
	};
	const left = bounds("baseline");
	const right = bounds("candidate");
	const retained = pairs.filter(
		(pair) =>
			pair.baseline.wallMs >= left[0] &&
			pair.baseline.wallMs <= left[1] &&
			pair.candidate.wallMs >= right[0] &&
			pair.candidate.wallMs <= right[1],
	);
	if (retained.length < 10)
		throw new Error("Too few retained pairs for comparison");
	let seed = 0xaff;
	const randomIndex = () => {
		seed = (Math.imul(seed, 1_664_525) + 1_013_904_223) >>> 0;
		return Math.floor((seed / 2 ** 32) * retained.length);
	};
	const savings: number[] = [];
	const differences: number[] = [];
	for (let iteration = 0; iteration < BOOTSTRAPS; iteration++) {
		let baseline = 0;
		let candidate = 0;
		for (let sample = 0; sample < retained.length; sample++) {
			const pair = retained[randomIndex()];
			if (!pair) throw new Error("Invalid bootstrap index");
			baseline += pair.baseline.wallMs;
			candidate += pair.candidate.wallMs;
		}
		savings.push(100 * (1 - candidate / baseline));
		differences.push((baseline - candidate) / retained.length);
	}
	const first = retained[0];
	if (!first) throw new Error("Empty retained sample");
	const describe = (key: "baseline" | "candidate") => {
		const times = retained.map((pair) => pair[key].wallMs);
		const average = mean(times);
		const sdMs = Math.sqrt(
			times.reduce((sum, value) => sum + (value - average) ** 2, 0) /
				(times.length - 1),
		);
		return {
			meanMs: average,
			sdMs,
			meanCpuMs: mean(retained.map((pair) => pair[key].cpuMs)),
			requests: first[key].requests,
			redundant: first[key].redundant,
		};
	};
	const rawBaselineMs = mean(pairs.map((pair) => pair.baseline.wallMs));
	const rawCandidateMs = mean(pairs.map((pair) => pair.candidate.wallMs));
	return {
		pairs: pairs.length,
		removedPairs: pairs.length - retained.length,
		unfilteredSavingPercent: 100 * (1 - rawCandidateMs / rawBaselineMs),
		baseline: describe("baseline"),
		candidate: describe("candidate"),
		savingPercent95CI: [quantile(savings, 0.025), quantile(savings, 0.975)],
		savedMs95CI: [quantile(differences, 0.025), quantile(differences, 0.975)],
	};
}

const args = process.argv.slice(2);
function option(name: string): string | undefined {
	const index = args.indexOf(name);
	if (index < 0) return undefined;
	const value = args[index + 1];
	if (!value || value.startsWith("--"))
		throw new Error(`Missing ${name} value`);
	return value;
}
const baselinePath = option("--baseline");
const baseline: Factory = baselinePath
	? (await import(pathToFileURL(baselinePath).href)).createAfkWidget
	: createAfkWidget;
const profilePath = option("--profile");
if (profilePath) {
	for (let warmup = 0; warmup < WARMUPS; warmup++) simulate(baseline, false);
	const profiler = new Session();
	profiler.connect();
	try {
		await profiler.post("Profiler.enable");
		await profiler.post("Profiler.start");
		for (const napping of [false, true]) {
			for (let repetition = 0; repetition < 8; repetition++)
				simulate(baseline, napping);
		}
		const { profile } = await profiler.post("Profiler.stop");
		await writeFile(profilePath, JSON.stringify(profile));
	} finally {
		profiler.disconnect();
	}
}
const cases: Record<string, ReturnType<typeof summarize>> = {};
for (const [name, napping] of [
	["awake", false],
	["asleep", true],
] as const) {
	for (let warmup = 0; warmup < WARMUPS; warmup++) {
		simulate(baseline, napping);
		simulate(createAfkWidget, napping);
	}
	const pairs: Pair[] = [];
	for (let sample = 0; sample < SAMPLES; sample++) {
		const candidateFirst = sample % 2 === 1;
		const first = simulate(
			candidateFirst ? createAfkWidget : baseline,
			napping,
		);
		const second = simulate(
			candidateFirst ? baseline : createAfkWidget,
			napping,
		);
		pairs.push(
			candidateFirst
				? { baseline: second, candidate: first }
				: { baseline: first, candidate: second },
		);
	}
	cases[name] = summarize(pairs);
}
const report = {
	runtime: process.version,
	platform: platform(),
	arch: arch(),
	ticks: TICKS,
	tickMs: 700,
	warmups: WARMUPS,
	bootstrapResamples: BOOTSTRAPS,
	cases,
	sink,
};
const output = JSON.stringify(report, null, 2);
const outputPath = option("--output");
if (outputPath) await writeFile(outputPath, `${output}\n`);
console.log(output);