"""
Generic Problem interface for DPO - define any optimization problem.
"""

from abc import ABC, abstractmethod
from typing import Dict, Any, Tuple, Optional, List
import numpy as np
from .solution import Solution


class Problem(ABC):
    """
    Abstract base class for optimization problems.
    DPO adapts to ANY problem by implementing this interface.
    
    Examples:
    - NAS: Evaluate architecture, return loss + metrics
    - TSP: Calculate tour length
    - Pathfinding: Find path cost
    - ML Tuning: Train model, return validation loss
    """
    
    @abstractmethod
    def evaluate(self, solution: Solution, **kwargs) -> Tuple[float, Dict[str, float]]:
        """
        Evaluate a solution.
        
        Args:
            solution: Solution object to evaluate
            **kwargs: Additional context (iteration, is_initial, etc.)
        
        Returns:
            Tuple of (fitness_value, metrics_dict)
            - fitness_value: float, lower is better
            - metrics_dict: dict of secondary metrics
        
        Examples:
            >>> fitness, metrics = problem.evaluate(solution)
            >>> # NAS: fitness = loss, metrics = {accuracy, latency, memory}
            >>> # TSP: fitness = tour_length, metrics = {gap_to_opt, time}
        """
        pass
    
    @abstractmethod
    def create_solution(self, **kwargs) -> Solution:
        """Create a random solution for this problem."""
        pass
    
    def get_constraints(self) -> Dict[str, Tuple[float, str]]:
        """
        Define problem constraints.
        
        Returns:
            Dict mapping constraint_name -> (limit, metric_name)
            E.g.: {"latency": (100.0, "latency_ms"), "memory": (50.0, "memory_mb")}
        """
        return {}
    
    def get_problem_info(self) -> Dict[str, Any]:
        """Return metadata about the problem."""
        return {
            'name': self.__class__.__name__,
            'constraints': self.get_constraints(),
        }


class ContinuousOptimizationProblem(Problem):
    """
    For continuous optimization problems (HPO, weights tuning, etc.)
    """
    
    def __init__(self, 
                 objective_fn,
                 param_bounds: List[Tuple[float, float]],
                 param_names: List[str] = None,
                 constraint_fn=None):
        """
        Args:
            objective_fn: Function that takes param dict, returns (fitness, metrics_dict)
            param_bounds: List of (min, max) tuples
            param_names: Names of parameters
            constraint_fn: Optional function for constraint checking
        """
        self.objective_fn = objective_fn
        self.param_bounds = param_bounds
        self.param_names = param_names or [f"param_{i}" for i in range(len(param_bounds))]
        self.constraint_fn = constraint_fn
        
        assert len(self.param_bounds) == len(self.param_names)
    
    def evaluate(self, solution: Solution, **kwargs) -> Tuple[float, Dict]:
        param_dict = solution.to_dict()
        return self.objective_fn(param_dict)
    
    def create_solution(self, **kwargs):
        from .solution import NumericSolution
        initial_values = np.array([
            np.random.uniform(vmin, vmax)
            for vmin, vmax in self.param_bounds
        ])
        return NumericSolution(initial_values, self.param_bounds, self.param_names)
    
    def get_constraints(self) -> Dict[str, Tuple[float, str]]:
        if self.constraint_fn:
            return self.constraint_fn()
        return {}


class CombinatoricOptimizationProblem(Problem):
    """
    For combinatoric problems (TSP, job scheduling, etc.)
    """
    
    def __init__(self,
                 objective_fn,
                 problem_size: int,
                 constraint_fn=None):
        """
        Args:
            objective_fn: Function(sequence) -> (fitness, metrics_dict)
            problem_size: Size of the problem (n for n-city TSP)
            constraint_fn: Optional constraint function
        """
        self.objective_fn = objective_fn
        self.problem_size = problem_size
        self.constraint_fn = constraint_fn
    
    def evaluate(self, solution: Solution, **kwargs) -> Tuple[float, Dict]:
        sequence_dict = solution.to_dict()
        return self.objective_fn(sequence_dict)
    
    def create_solution(self, **kwargs):
        from .solution import CombinatoricSolution
        sequence = list(np.random.permutation(self.problem_size))
        return CombinatoricSolution(sequence, self.problem_size)
    
    def get_constraints(self) -> Dict[str, Tuple[float, str]]:
        if self.constraint_fn:
            return self.constraint_fn()
        return {}


class NASProblem(Problem):
    """
    For Neural Architecture Search (backward compatible with old code).
    """
    
    def __init__(self,
                 evaluator,
                 architecture_fn=None,
                 constraints: Dict = None):
        """
        Args:
            evaluator: Estimator with estimate() method
            architecture_fn: Function to convert solution to arch dict
            constraints: Dict of constraint limits
        """
        self.evaluator = evaluator
        self.architecture_fn = architecture_fn
        self._constraints = constraints or {
            'latency': 100.0,
            'memory': 50.0,
            'flops': 300.0,
        }
    
    def evaluate(self, solution: Solution, **kwargs) -> Tuple[float, Dict]:
        # Convert solution to architecture dict
        if self.architecture_fn:
            arch_dict = self.architecture_fn(solution)
        else:
            arch_dict = solution.to_dict()
        
        fitness, metrics = self.evaluator.estimate(
            arch_dict, 
            search_mode=kwargs.get('search_mode', True),
            iteration=kwargs.get('iteration', 0)
        )
        return fitness, metrics
    
    def create_solution(self, **kwargs):
        """Create random NAS architecture."""
        from ..architecture.gene import ArchitectureGene
        from .solution import HybridSolution
        
        # Create old-style gene for compatibility
        gene = ArchitectureGene()
        
        # Convert to hybrid solution
        operations = gene.gene[:gene.num_layers].astype(int)
        kernels = gene.gene[gene.num_layers:2*gene.num_layers].astype(int)
        skips = gene.gene[2*gene.num_layers:2*gene.num_layers+gene.num_cells].astype(int)
        multipliers = gene.gene[-2:]
        
        numeric_part = np.array(multipliers, dtype=np.float32)
        combinatoric_part = list(operations) + list(kernels) + list(skips)
        
        return HybridSolution(
            numeric_part=numeric_part,
            combinatoric_part=combinatoric_part,
            numeric_bounds=[(0.5, 1.5), (0.5, 1.5)],
            numeric_names=['depth_mult', 'channel_mult']
        )
    
    def get_constraints(self) -> Dict[str, Tuple[float, str]]:
        return {
            'latency': (self._constraints.get('latency', 100.0), 'latency_ms'),
            'memory': (self._constraints.get('memory', 50.0), 'memory_mb'),
            'flops': (self._constraints.get('flops', 300.0), 'flops_m'),
        }


class HybridProblem(Problem):
    """
    For problems with both continuous and discrete variables.
    E.g., robot control with continuous actions + discrete strategy selection.
    """
    
    def __init__(self,
                 objective_fn,
                 numeric_bounds: List[Tuple[float, float]],
                 discrete_options: Dict[str, List[Any]],
                 param_names: List[str] = None,
                 constraint_fn=None):
        """
        Args:
            objective_fn: Function(params_dict) -> (fitness, metrics)
            numeric_bounds: List of (min, max) for continuous params
            discrete_options: Dict of {param_name: [options]}
            param_names: Names of parameters
            constraint_fn: Optional constraints function
        """
        self.objective_fn = objective_fn
        self.numeric_bounds = numeric_bounds
        self.discrete_options = discrete_options
        self.constraint_fn = constraint_fn
        
        n_numeric = len(numeric_bounds)
        n_discrete = len(discrete_options)
        self.param_names = param_names or \
            [f"num_{i}" for i in range(n_numeric)] + \
            [f"disc_{k}" for k in discrete_options.keys()]
    
    def evaluate(self, solution: Solution, **kwargs) -> Tuple[float, Dict]:
        param_dict = solution.to_dict()
        return self.objective_fn(param_dict)
    
    def create_solution(self, **kwargs):
        from .solution import HybridSolution
        
        numeric_part = np.array([
            np.random.uniform(vmin, vmax)
            for vmin, vmax in self.numeric_bounds
        ], dtype=np.float32)
        
        combinatoric_part = [
            np.random.choice(options)
            for options in self.discrete_options.values()
        ]
        
        return HybridSolution(
            numeric_part=numeric_part,
            combinatoric_part=combinatoric_part,
            numeric_bounds=self.numeric_bounds,
            numeric_names=list(self.discrete_options.keys())
        )
    
    def get_constraints(self) -> Dict[str, Tuple[float, str]]:
        if self.constraint_fn:
            return self.constraint_fn()
        return {}
