"""Auto-generated example: Bencher self-introspection: overhead vs problem size."""

import bencher as bn


class TrivialWorkload(bn.ParametrizedSweep):
    """A near-zero-cost worker so we measure framework overhead, not compute."""

    x = bn.FloatSweep(default=0, bounds=[0, 1], samples=2)
    result = bn.ResultFloat(units="v", doc="trivial output")

    def benchmark(self):
        self.result = self.x * 2


class BencherSelfBenchmark(bn.ParametrizedSweep):
    """Sweep over problem sizes and measure bencher's own timing phases."""

    num_samples = bn.IntSweep(
        default=10,
        bounds=[2, 100],
        samples=6,
        doc="Number of parameter samples in the inner sweep",
    )
    use_cache = bn.BoolSweep(default=False, doc="Whether sample caching is enabled")

    # Result variables — one per timing phase
    total_ms = bn.ResultFloat(units="ms", doc="Total sweep wall-clock time")
    dataset_setup_ms = bn.ResultFloat(units="ms", doc="Dataset initialization time")
    job_submission_ms = bn.ResultFloat(units="ms", doc="Job creation and submission time")
    job_execution_ms = bn.ResultFloat(units="ms", doc="Worker execution and result storage time")
    cache_check_ms = bn.ResultFloat(units="ms", doc="Benchmark cache lookup time")
    sample_cache_init_ms = bn.ResultFloat(units="ms", doc="Sample cache initialization time")
    throughput = bn.ResultFloat(units="samples/s", doc="Samples processed per second")

    def benchmark(self):
        workload = TrivialWorkload()
        x_sweep = bn.FloatSweep(default=0, bounds=[0, 1], samples=self.num_samples, doc="input")
        x_sweep.name = "x"

        inner_cfg = bn.BenchRunCfg()
        inner_cfg.repeats = 1
        inner_cfg.cache_samples = self.use_cache
        inner_cfg.cache_results = False
        inner_cfg.auto_plot = False

        bench = bn.Bench("inner_bench", workload, run_cfg=inner_cfg)
        bench.plot_sweep(input_vars=[x_sweep], result_vars=["result"])
        res = bench.results[-1]

        t = res.timings
        self.total_ms = t.total_ms
        self.dataset_setup_ms = t.dataset_setup_ms
        self.job_submission_ms = t.job_submission_ms
        self.job_execution_ms = t.job_execution_ms
        self.cache_check_ms = t.cache_check_ms
        self.sample_cache_init_ms = t.sample_cache_init_ms
        self.throughput = (self.num_samples / t.total_ms * 1000) if t.total_ms > 0 else 0


def example_perf_self_benchmark(run_cfg: bn.BenchRunCfg | None = None) -> bn.Bench:
    """Bencher self-introspection: overhead vs problem size."""
    bench = BencherSelfBenchmark().to_bench(run_cfg)
    bench.plot_sweep(
        input_vars=["num_samples"],
        result_vars=["total_ms", "dataset_setup_ms", "job_submission_ms", "job_execution_ms"],
        title="Phase Timing vs Problem Size",
    )
    bench.plot_sweep(
        input_vars=["num_samples"],
        result_vars=["throughput"],
        title="Throughput vs Problem Size",
    )
    bench.plot_sweep(
        input_vars=["num_samples", "use_cache"],
        result_vars=["total_ms"],
        title="Cache Impact on Total Time",
    )

    return bench


if __name__ == "__main__":
    bn.run(example_perf_self_benchmark)
