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

/**
 * Wall-time estimates for the two stages of a run. This is a calibration,
 * not physics: the sizes (the plan's `work` for the modes stage, samples
 * times modes for the strike) come from the method's plan function running in
 * numbl; only the milliseconds-per-unit constants live here, seeded from the
 * method definition and refined by each finished stage.
 */
export class CostModel {
  private rates = new Map<string, number>();
  private readonly overheadMs = 150;

  private rate(method: MethodDefinition, stage: "modes" | "strike"): number {
    return (
      this.rates.get(`${method.id}:${stage}`) ??
      (stage === "modes" ? method.cost.msPerWork : method.cost.msPerSampleMode)
    );
  }

  modesMs(method: MethodDefinition, work: number): number {
    return work * this.rate(method, "modes") + this.overheadMs;
  }

  strikeMs(method: MethodDefinition, samples: number, modes: number): number {
    return samples * Math.max(modes, 1) * this.rate(method, "strike") + this.overheadMs;
  }

  /** Blend a measurement in, weighting recent runs more. */
  observe(method: MethodDefinition, stage: "modes" | "strike", units: number, wallMs: number) {
    if (units <= 0 || wallMs <= 0) return;
    const measured = (wallMs - this.overheadMs) / units;
    if (!Number.isFinite(measured) || measured <= 0) return;
    this.rates.set(`${method.id}:${stage}`, 0.5 * this.rate(method, stage) + 0.5 * measured);
  }
}