#!/usr/bin/env python3
"""Adaptive Gemma 4 12B recipe loop helpers.

This module is intentionally side-effect free on import so tests and other
automation can reuse the recipe policy without loading a model.
"""

from __future__ import annotations

from dataclasses import dataclass, replace


@dataclass(frozen=True)
class Recipe:
    name: str
    method: str = "informed"
    n_directions: int = 4
    min_layer_fraction: float = 0.35
    max_layer_fraction: float = 0.65
    regularization: float = 0.0
    refinement_passes: int = 2


BASE_RECIPES = [
    Recipe(name="gemma4_12b_balanced"),
    Recipe(
        name="gemma4_12b_surgical",
        method="surgical",
        n_directions=6,
        min_layer_fraction=0.30,
        max_layer_fraction=0.70,
    ),
]


def next_round_from(recipe: Recipe, metrics: dict[str, float]) -> list[Recipe]:
    """Derive the next candidate recipes from aggregate benchmark metrics."""
    refusal_rate = metrics.get("refusal_rate", 0.0)
    repetition_rate = metrics.get("repetition_rate", 0.0)
    short_rate = metrics.get("short_rate", 0.0)
    collapse_rate = max(repetition_rate, short_rate)

    next_recipes: list[Recipe] = []
    if refusal_rate > 0.0:
        next_recipes.append(
            replace(
                recipe,
                name=f"{recipe.name}_more_dirs",
                n_directions=min(recipe.n_directions + 2, 12),
            )
        )
        next_recipes.append(
            replace(
                recipe,
                name=f"{recipe.name}_wider_layers",
                max_layer_fraction=min(recipe.max_layer_fraction + 0.10, 0.95),
            )
        )

    if collapse_rate > 0.0:
        next_recipes.append(
            replace(
                recipe,
                name=f"{recipe.name}_more_conservative",
                n_directions=max(recipe.n_directions - 1, 1),
                max_layer_fraction=max(recipe.max_layer_fraction - 0.05, recipe.min_layer_fraction),
                regularization=min(recipe.regularization + 0.05, 0.5),
            )
        )

    return next_recipes or [recipe]


if __name__ == "__main__":
    for item in BASE_RECIPES:
        print(item)