"""DPO: Debt-Paying Optimizer for optimization problems beyond NAS.

This package exposes both low-level optimizer classes and high-level convenience
helpers that route to :class:`dpo.core.universal.DPO_Universal`.
"""

from __future__ import annotations

from typing import Callable, Dict, List, Optional, Tuple, Union

__version__ = "2.6.0"
__author__ = "Arya H"
__email__ = "arya.h1718@gmail.com"

from .core.config import DPO_Config
from .core.optimizer import DPO_NAS
from .core.universal import DPO_Universal, DPO_Presets
from .evaluation.ensemble import EnsembleEstimator
from .constraints.handler import AdvancedConstraintHandler
from .core.problem import (
    Problem,
    ContinuousOptimizationProblem,
    CombinatoricOptimizationProblem,
    NASProblem,
    HybridProblem,
)
from .core.solution import (
    Solution,
    NumericSolution,
    CombinatoricSolution,
    HybridSolution,
)


def dpo(
    problem: Optional[Problem] = None,
    *,
    config: Optional[DPO_Config] = None,
    preset: Optional[str] = None,
    # Convenience kwargs – build a Problem automatically when *problem* is None
    problem_type: Optional[str] = None,
    objective_fn: Optional[Callable] = None,
    param_bounds: Optional[List[Tuple[float, float]]] = None,
    param_names: Optional[List[str]] = None,
    problem_size: Optional[int] = None,
    **config_overrides,
) -> Dict:
    """One-line optimization helper.

    Can be called in two ways:

    1. With an explicit *problem* instance::

        result = dpo(problem=my_problem, preset='balanced')

    2. With convenience kwargs (builds the problem for you)::

        result = dpo(
            problem_type='continuous',
            objective_fn=my_fn,
            param_bounds=[(0, 1), (0, 1)],
            param_names=['x', 'y'],
            preset='balanced',
            max_iterations=100,
            population_size=40,
        )

    Supported *problem_type* values: ``'continuous'``, ``'combinatoric'``.

    Any extra keyword arguments (``max_iterations``, ``population_size``, etc.)
    are forwarded as config overrides.
    """
    # ── Build problem from convenience args if not provided ───────────
    if problem is None and objective_fn is not None:
        ptype = (problem_type or '').lower()
        if ptype == 'continuous' or (param_bounds is not None and ptype != 'combinatoric'):
            problem = ContinuousOptimizationProblem(
                objective_fn=objective_fn,
                param_bounds=param_bounds,
                param_names=param_names,
            )
        elif ptype == 'combinatoric' and problem_size is not None:
            problem = CombinatoricOptimizationProblem(
                objective_fn=objective_fn,
                problem_size=problem_size,
            )

    # ── Build config from preset + overrides ─────────────────────────
    if config is None and (preset or config_overrides):
        config = DPO_Config.from_preset(preset or 'balanced')
        for key, value in config_overrides.items():
            if hasattr(config, key):
                setattr(config, key, value)

    optimizer = DPO_Universal(problem=problem, config=config, preset=None if config else preset)
    return optimizer.optimize()


def dpo_optimize(
    *,
    objective: Callable,
    bounds: List[Tuple[float, float]],
    names: Optional[List[str]] = None,
    preset: Union[str, DPO_Config] = "balanced",
    **config_overrides,
) -> Dict:
    """Convenience helper for continuous optimization problems."""
    problem = ContinuousOptimizationProblem(
        objective_fn=lambda params: objective(params),
        param_bounds=bounds,
        param_names=names,
    )

    if isinstance(preset, DPO_Config):
        config = preset
    else:
        config = DPO_Config.from_preset(preset)

    for key, value in config_overrides.items():
        if hasattr(config, key):
            setattr(config, key, value)

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


def dpo_solve_tsp(
    *,
    distance_matrix,
    preset: Union[str, DPO_Config] = "balanced",
    **config_overrides,
) -> Dict:
    """Convenience helper for Traveling Salesman Problem optimization."""
    def tsp_objective(seq_dict):
        tour = seq_dict["sequence"]
        total_dist = sum(
            float(distance_matrix[tour[i], tour[(i + 1) % len(tour)]])
            for i in range(len(tour))
        )
        return total_dist, {
            "accuracy": float(1.0 / (1.0 + total_dist)),
            "latency_ms": float(total_dist),
            "memory_mb": float(len(tour) * 0.01),
            "flops_m": float(len(tour)),
        }

    if isinstance(preset, DPO_Config):
        config = preset
    else:
        config = DPO_Config.from_preset(preset)

    for key, value in config_overrides.items():
        if hasattr(config, key):
            setattr(config, key, value)

    problem = CombinatoricOptimizationProblem(
        objective_fn=tsp_objective,
        problem_size=len(distance_matrix),
    )
    optimizer = DPO_Universal(problem=problem, config=config)
    return optimizer.optimize()


def dpo_solve_nas(
    *,
    estimator,
    architecture_fn: Optional[Callable] = None,
    constraints: Optional[Dict] = None,
    preset: Union[str, DPO_Config] = "balanced",
    **config_overrides,
) -> Dict:
    """Convenience helper for NAS optimization."""
    if isinstance(preset, DPO_Config):
        config = preset
    else:
        config = DPO_Config.from_preset(preset)

    for key, value in config_overrides.items():
        if hasattr(config, key):
            setattr(config, key, value)

    problem = NASProblem(
        evaluator=estimator,
        architecture_fn=architecture_fn,
        constraints=constraints,
    )
    optimizer = DPO_Universal(problem=problem, config=config)
    return optimizer.optimize()

__all__ = [
    # Version info
    '__version__',
    
    # Core classes
    'DPO_Config',
    'DPO_NAS',
    'DPO_Universal',
    'DPO_Presets',
    'EnsembleEstimator',
    'AdvancedConstraintHandler',
    
    # APIs
    'dpo',
    'dpo_optimize',
    'dpo_solve_tsp',
    'dpo_solve_nas',
    
    # Problem classes
    'Problem',
    'ContinuousOptimizationProblem',
    'CombinatoricOptimizationProblem',
    'NASProblem',
    'HybridProblem',
    
    # Solution classes
    'Solution',
    'NumericSolution',
    'CombinatoricSolution',
    'HybridSolution',
]
