import type { MethodDefinition } from "./methods/types";

/**
 * Wall-time estimate for grid-based methods. This is a calibration, not
 * physics: the plan (cell count, steps) comes from the method's plan function
 * running in numbl; only the seconds-per-cell-step constant lives here, seeded
 * from the method definition and refined by each finished run. The GPU path of
 * a method is calibrated separately: its cost is nearly independent of the
 * cell count at the grids used here, so its per-cell-step rate is a stand-in
 * that the measurements correct for the grid in use.
 */
export class CostModel {
  private us = new Map<string, number>();
  private readonly overheadMs = 150;

  private key(method: MethodDefinition, gpu: boolean): string {
    return gpu ? `${method.id}:gpu` : method.id;
  }

  private rate(method: MethodDefinition, gpu: boolean): number {
    return (
      this.us.get(this.key(method, gpu)) ??
      (gpu ? (method.gpu?.usPerCellStep ?? method.usPerCellStep) : method.usPerCellStep)
    );
  }

  estimateMs(method: MethodDefinition, cellSteps: number, gpu = false): number {
    return (cellSteps * this.rate(method, gpu)) / 1000 + this.overheadMs;
  }

  /** Blend a measurement in, weighting recent runs more. */
  observe(method: MethodDefinition, cellSteps: number, wallMs: number, gpu = false): void {
    if (cellSteps <= 0 || wallMs <= 0) return;
    const measured = ((wallMs - this.overheadMs) * 1000) / cellSteps;
    if (!Number.isFinite(measured) || measured <= 0) return;
    this.us.set(this.key(method, gpu), 0.5 * this.rate(method, gpu) + 0.5 * measured);
  }
}