"""
NAS example using DPO_Universal with a lightweight mock estimator.
"""

import numpy as np

from dpo.core.problem import NASProblem
from dpo.core.universal import DPO_Presets, DPO_Universal


class MockNASEstimator:
    """Simple deterministic+noise estimator for demonstration."""

    def estimate(self, arch_dict, **kwargs):
        rng = np.random.default_rng(kwargs.get("iteration", 0) + 17)

        operations = arch_dict.get("operations", [])
        kernels = arch_dict.get("kernels", [])
        depth = float(arch_dict.get("depth_multiplier", 1.0))
        channels = float(arch_dict.get("channel_multiplier", 1.0))

        op_complexity = sum((int(o) if isinstance(o, (int, np.integer)) else 1) for o in operations)
        kernel_cost = sum((int(k) if isinstance(k, (int, np.integer)) else 3) for k in kernels)

        latency = 20.0 + 0.5 * op_complexity + 0.8 * kernel_cost + 8.0 * depth
        memory = 10.0 + 0.2 * op_complexity + 5.0 * channels
        flops = 80.0 + 0.9 * kernel_cost + 12.0 * depth * channels

        base_acc = 0.80 + 0.08 * np.tanh((depth + channels) / 3.0)
        penalty = 0.0006 * latency + 0.0005 * flops + 0.0008 * memory
        noise = float(rng.normal(0.0, 0.002))
        accuracy = float(np.clip(base_acc - penalty + noise, 0.0, 1.0))

        fitness = 1.0 - accuracy
        return fitness, {
            "accuracy": accuracy,
            "latency_ms": float(latency),
            "memory_mb": float(memory),
            "flops_m": float(flops),
        }


def run_nas_example():
    config = DPO_Presets.NAS_Config(population_size=24, max_iterations=30)
    config.verbose = False

    problem = NASProblem(
        evaluator=MockNASEstimator(),
        constraints={"latency": 100.0, "memory": 50.0, "flops": 300.0},
    )

    optimizer = DPO_Universal(problem=problem, config=config)
    result = optimizer.optimize()

    print("NAS Example")
    print("-" * 40)
    print(f"Best fitness : {result['best_fitness']:.6f}")
    print(f"Best accuracy: {result.get('best_accuracy', 0.0):.6f}")
    print(f"Evaluations  : {result.get('total_evaluations', 'N/A')}")
    return result


if __name__ == "__main__":
    run_nas_example()
