src/benchmark/openhands-aggregate.ts

import fs from "fs";
import path from "path";
import type { OpenHandsBenchReport, TraceBenchRow } from "./openhands-run.ts";
import { ROI_PATTERN_TYPES } from "./openhands-run.ts";

export type MetricsCsvRow = {
  instance_id: string;
  arm: string;
  resolved: string;
  total_tokens: string;
  prompt_tokens: string;
  completion_tokens: string;
  usd: string;
  tcg_savings_pct: string;
  repeat_read_rate: string;
  explore_fix_ratio: string;
  pattern_hit_count: string;
  pattern_types: string;
  optimize_latency_ms: string;
  log_count: string;
};

export type ParetoPoint = {
  instance_id: string;
  arm: string;
  savings_rate: number;
  resolved: number;
  usd_c0: number | null;
  usd_arm: number | null;
};

export type P1AggregateSummary = {
  generated_at: string;
  instance_count: number;
  arms: string[];
  resolve_rate_by_arm: Record<string, number>;
  resolve_drop_pp: number | null;
  avg_savings_pct_c1: number;
  avg_repeat_read_rate_c1: number;
  two_plus_pattern_rate_c1: number;
  pareto_points: ParetoPoint[];
};

const CSV_HEADER: (keyof MetricsCsvRow)[] = [
  "instance_id",
  "arm",
  "resolved",
  "total_tokens",
  "prompt_tokens",
  "completion_tokens",
  "usd",
  "tcg_savings_pct",
  "repeat_read_rate",
  "explore_fix_ratio",
  "pattern_hit_count",
  "pattern_types",
  "optimize_latency_ms",
  "log_count",
];

function csvEscape(value: string): string {
  if (value.includes(",") || value.includes('"') || value.includes("\n")) {
    return `"${value.replace(/"/g, '""')}"`;
  }
  return value;
}

export function traceToCsvRow(row: TraceBenchRow): MetricsCsvRow {
  const swe = row.swe;
  const resolved = swe?.resolved != null ? String(swe.resolved) : "";
  return {
    instance_id: row.instance_id,
    arm: row.arm,
    resolved,
    total_tokens: swe?.total_tokens != null ? String(swe.total_tokens) : String(row.tokens_from_logs),
    prompt_tokens: swe?.prompt_tokens != null ? String(swe.prompt_tokens) : "",
    completion_tokens: swe?.completion_tokens != null ? String(swe.completion_tokens) : "",
    usd: swe?.usd != null ? String(swe.usd) : String(row.usd_est_from_logs),
    tcg_savings_pct: row.savings_percentage != null ? String(row.savings_percentage) : "",
    repeat_read_rate: String(row.metrics.repeat_read_rate),
    explore_fix_ratio: String(row.metrics.explore_fix_ratio),
    pattern_hit_count: String(row.pattern_hit_count),
    pattern_types: row.pattern_types.filter((p) => ROI_PATTERN_TYPES.has(p)).join("|"),
    optimize_latency_ms: String(row.optimize_latency_ms),
    log_count: String(row.log_count),
  };
}

export function resolveRate(rows: TraceBenchRow[]): number {
  const withOutcome = rows.filter((r) => r.swe?.resolved != null);
  if (withOutcome.length === 0) return 0;
  return withOutcome.filter((r) => r.swe?.resolved).length / withOutcome.length;
}

export function loadBenchReports(reportsDir: string): OpenHandsBenchReport[] {
  if (!fs.existsSync(reportsDir)) return [];
  return fs
    .readdirSync(reportsDir)
    .filter((f) => f.startsWith("bench-") && f.endsWith(".json"))
    .map((f) => JSON.parse(fs.readFileSync(path.join(reportsDir, f), "utf-8")) as OpenHandsBenchReport)
    .sort((a, b) => a.arm.localeCompare(b.arm));
}

export function aggregateP1(reports: OpenHandsBenchReport[]): {
  csvRows: MetricsCsvRow[];
  summary: P1AggregateSummary;
} {
  const csvRows: MetricsCsvRow[] = [];
  const byInstanceArm = new Map<string, TraceBenchRow>();

  for (const report of reports) {
    for (const trace of report.traces) {
      csvRows.push(traceToCsvRow(trace));
      byInstanceArm.set(`${trace.instance_id}::${trace.arm}`, trace);
    }
  }

  const instanceIds = [...new Set(reports.flatMap((r) => r.traces.map((t) => t.instance_id)))].sort();
  const arms = [...new Set(reports.map((r) => r.arm))].sort();
  const resolve_rate_by_arm: Record<string, number> = {};
  for (const arm of arms) {
    const armRows = reports.filter((r) => r.arm === arm).flatMap((r) => r.traces);
    resolve_rate_by_arm[arm] = Number(resolveRate(armRows).toFixed(4));
  }

  const c0Rate = resolve_rate_by_arm.C0 ?? null;
  const c1Rate = resolve_rate_by_arm.C1 ?? null;
  const resolve_drop_pp =
    c0Rate != null && c1Rate != null ? Number(((c0Rate - c1Rate) * 100).toFixed(2)) : null;

  const c1Traces = reports.filter((r) => r.arm === "C1").flatMap((r) => r.traces);
  const savingsVals = c1Traces.map((t) => t.savings_percentage).filter((v): v is number => v != null);
  const avg_savings_pct_c1 = savingsVals.length
    ? Number((savingsVals.reduce((a, b) => a + b, 0) / savingsVals.length).toFixed(4))
    : 0;

  const avg_repeat_read_rate_c1 = c1Traces.length
    ? Number((c1Traces.reduce((s, t) => s + t.metrics.repeat_read_rate, 0) / c1Traces.length).toFixed(4))
    : 0;

  const twoPlus = c1Traces.filter(
    (t) => new Set(t.pattern_types.filter((p) => ROI_PATTERN_TYPES.has(p))).size >= 2
  ).length;
  const two_plus_pattern_rate_c1 = c1Traces.length
    ? Number((twoPlus / c1Traces.length).toFixed(4))
    : 0;

  const pareto_points: ParetoPoint[] = [];
  for (const instance_id of instanceIds) {
    const c0 = byInstanceArm.get(`${instance_id}::C0`);
    const c1 = byInstanceArm.get(`${instance_id}::C1`);
    const target = c1 ?? byInstanceArm.get(`${instance_id}::C2`);
    if (!target) continue;
    const usdC0 = c0?.swe?.usd ?? c0?.usd_est_from_logs ?? null;
    const usdArm = target.swe?.usd ?? target.usd_est_from_logs ?? null;
    let savings_rate = target.savings_percentage ?? 0;
    if (usdC0 != null && usdArm != null && usdC0 > 0) {
      savings_rate = (usdC0 - usdArm) / usdC0;
    }
    pareto_points.push({
      instance_id,
      arm: target.arm,
      savings_rate: Number(savings_rate.toFixed(4)),
      resolved: target.swe?.resolved ? 1 : 0,
      usd_c0: usdC0,
      usd_arm: usdArm,
    });
  }

  return {
    csvRows,
    summary: {
      generated_at: new Date().toISOString(),
      instance_count: instanceIds.length,
      arms,
      resolve_rate_by_arm,
      resolve_drop_pp,
      avg_savings_pct_c1,
      avg_repeat_read_rate_c1,
      two_plus_pattern_rate_c1,
      pareto_points,
    },
  };
}

export function writeMetricsCsv(rows: MetricsCsvRow[], outPath: string): void {
  const lines = [
    CSV_HEADER.join(","),
    ...rows.map((row) => CSV_HEADER.map((k) => csvEscape(row[k])).join(",")),
  ];
  fs.mkdirSync(path.dirname(outPath), { recursive: true });
  fs.writeFileSync(outPath, `${lines.join("\n")}\n`, "utf-8");
}

export function writeParetoSvg(points: ParetoPoint[], outPath: string): void {
  const width = 640;
  const height = 400;
  const pad = 48;
  const plotW = width - pad * 2;
  const plotH = height - pad * 2;

  const dots = points
    .map((p, i) => {
      const x = pad + Math.max(0, Math.min(1, p.savings_rate)) * plotW;
      const y = pad + (1 - p.resolved) * plotH * 0.85 + (i % 5) * 3;
      const color = p.resolved ? "#16a34a" : "#dc2626";
      return `<circle cx="${x.toFixed(1)}" cy="${y.toFixed(1)}" r="5" fill="${color}" opacity="0.85"><title>${p.instance_id} saved=${(p.savings_rate * 100).toFixed(1)}%</title></circle>`;
    })
    .join("\n");

  const svg = `<?xml version="1.0" encoding="UTF-8"?>
<svg xmlns="http://www.w3.org/2000/svg" width="${width}" height="${height}" viewBox="0 0 ${width} ${height}">
  <rect width="100%" height="100%" fill="#fafafa"/>
  <text x="${pad}" y="24" font-family="sans-serif" font-size="14" fill="#111">Pareto: savings rate vs resolve (C1/C2 vs C0 USD)</text>
  <line x1="${pad}" y1="${height - pad}" x2="${width - pad}" y2="${height - pad}" stroke="#333"/>
  <line x1="${pad}" y1="${pad}" x2="${pad}" y2="${height - pad}" stroke="#333"/>
  <text x="${width / 2}" y="${height - 12}" text-anchor="middle" font-size="11" fill="#444">Estimated / realized savings rate →</text>
  <text x="14" y="${height / 2}" transform="rotate(-90 14 ${height / 2})" font-size="11" fill="#444">Resolve (jittered)</text>
  ${dots}
</svg>
`;
  fs.mkdirSync(path.dirname(outPath), { recursive: true });
  fs.writeFileSync(outPath, svg, "utf-8");
}