"""Meta-generator: Optimization & Pareto.

Shows optimization direction, Optuna integration, and multi-objective Pareto.
"""

from typing import Any

import bencher as bch
from bencher.example.meta.meta_generator_base import MetaGeneratorBase

OUTPUT_DIR = "optimization"

CLASS_CODE = "\n".join(
    [
        "class ServerOptimizer(bch.ParametrizedSweep):",
        '    """Optimizes server configuration for performance vs cost tradeoff."""',
        "",
        '    cpu_cores = bch.FloatSweep(default=4, bounds=[1, 32], doc="Number of CPU cores")',
        '    memory_gb = bch.FloatSweep(default=8, bounds=[1, 64], doc="Memory in GB")',
        "",
        '    performance = bch.ResultVar("score", bch.OptDir.maximize, doc="Performance score (maximize)")',
        '    cost = bch.ResultVar("$/hr", bch.OptDir.minimize, doc="Hourly cost (minimize)")',
        "",
        '    noise_scale = bch.FloatSweep(default=0.0, bounds=[0.0, 1.0], doc="Noise scale")',
        "",
        "    def __call__(self, **kwargs):",
        "        self.update_params_from_kwargs(**kwargs)",
        "        self.performance = math.log2(self.cpu_cores + 1) * math.sqrt(self.memory_gb) * 10",
        "        self.cost = 0.05 * self.cpu_cores + 0.02 * self.memory_gb",
        "        if self.noise_scale > 0:",
        "            self.performance += random.gauss(0, self.noise_scale * 5)",
        "            self.cost += random.gauss(0, self.noise_scale * 0.1)",
        "        return super().__call__()",
    ]
)


class MetaOptimization(MetaGeneratorBase):
    """Generate Python examples demonstrating optimization features."""

    n_objectives = bch.IntSweep(default=1, bounds=(1, 2), doc="Number of objectives")
    input_dims = bch.IntSweep(default=1, bounds=(1, 2), doc="Number of input dimensions")

    def __call__(self, **kwargs: Any) -> Any:
        self.update_params_from_kwargs(**kwargs)

        filename = f"optim_{self.n_objectives}obj_{self.input_dims}d"
        function_name = f"example_optim_{self.n_objectives}obj_{self.input_dims}d"
        title = f"Optimization: {self.n_objectives} objective(s), {self.input_dims}D input"

        if self.n_objectives == 1:
            result_vars_code = '["performance"]'
        else:
            result_vars_code = '["performance", "cost"]'

        if self.input_dims == 1:
            input_vars_code = '["cpu_cores"]'
        else:
            input_vars_code = '["cpu_cores", "memory_gb"]'

        if self.n_objectives == 1:
            description = (
                f"Single-objective optimization over {self.input_dims}D input space using Optuna. "
                "The optimizer searches for the parameter combination that maximizes performance."
            )
            post_description = (
                "The Optuna importance plot shows which input parameters most affect the objective."
            )
        else:
            description = (
                f"Multi-objective optimization over {self.input_dims}D input space using Optuna. "
                "The optimizer finds the Pareto front trading off performance vs cost."
            )
            post_description = (
                "The Pareto front shows optimal trade-offs — no point can improve one "
                "objective without worsening the other."
            )

        level = 2 if self.input_dims >= 2 else 3
        self.generate_sweep_example(
            title=title,
            output_dir=OUTPUT_DIR,
            filename=filename,
            function_name=function_name,
            benchable_class="ServerOptimizer",
            benchable_module=None,
            class_code=CLASS_CODE,
            extra_imports=["import math", "import random"],
            input_vars=input_vars_code,
            result_vars=result_vars_code,
            const_vars="dict(noise_scale=0.1)",
            description=description,
            post_description=post_description,
            post_sweep_line="res.to_optuna_plots()",
            run_cfg_lines=["run_cfg.use_optuna = True"],
            run_kwargs={"level": level, "repeats": 3},
        )

        return super().__call__()


def example_meta_optimization(run_cfg: bch.BenchRunCfg | None = None) -> bch.Bench:
    bench = MetaOptimization().to_bench(run_cfg)

    bench.plot_sweep(
        title="Optimization",
        input_vars=[
            bch.p("n_objectives", [1, 2]),
            bch.p("input_dims", [1, 2]),
        ],
    )

    return bench


if __name__ == "__main__":
    bch.run(example_meta_optimization)
