"""
Generic Solution interface for DPO - works with any problem domain.
This replaces domain-specific ArchitectureGene with a flexible solution representation.
"""

from abc import ABC, abstractmethod
import numpy as np
from typing import Dict, Any, List
from copy import deepcopy


class Solution(ABC):
    """
    Abstract base class for any optimization problem solution.
    
    Examples:
    - NAS: Neural network architecture
    - TSP: Tour/route
    - Pathfinding: Path
    - ML HPO: Hyperparameters
    """
    
    @abstractmethod
    def to_dict(self) -> Dict[str, Any]:
        """Convert solution to dictionary representation."""
        pass
    
    @abstractmethod
    def mutate(self, mutation_strength: float = 0.15) -> 'Solution':
        """Create a mutated copy of this solution."""
        pass
    
    @abstractmethod
    def copy(self) -> 'Solution':
        """Create a deep copy of this solution."""
        pass
    
    @abstractmethod
    def crossover(self, other: 'Solution') -> 'Solution':
        """Create offspring by combining this and other solution."""
        pass
    
    def get_gene_vector(self) -> np.ndarray:
        """
        Optional: Return normalized gene vector for diversity calculation.
        Default: Return empty array if not applicable.
        """
        return np.array([])


class NumericSolution(Solution):
    """
    Generic numeric solution for continuous/discrete problems.
    Suitable for hyperparameter tuning, weights, routes, etc.
    """
    
    def __init__(self, 
                 values: np.ndarray,
                 bounds: List[tuple] = None,
                 names: List[str] = None):
        """
        Args:
            values: numpy array of solution values
            bounds: list of (min, max) tuples for each dimension
            names: optional names for each dimension
        """
        self.values = np.array(values, dtype=np.float32)
        self.dim = len(self.values)
        
        # Set bounds: if not provided, derive from initial values
        if bounds is None:
            bounds = [(v * 0.5, v * 1.5) for v in self.values]
        self.bounds = bounds
        
        self.names = names or [f"param_{i}" for i in range(self.dim)]
        
        # Enforce bounds
        self._enforce_bounds()
    
    def _enforce_bounds(self):
        """Clip values to bounds."""
        for i, (vmin, vmax) in enumerate(self.bounds):
            self.values[i] = np.clip(self.values[i], vmin, vmax)
    
    def to_dict(self) -> Dict[str, Any]:
        return {name: float(val) for name, val in zip(self.names, self.values)}
    
    def mutate(self, mutation_strength: float = 0.15) -> 'NumericSolution':
        mutant = self.copy()
        num_mutations = max(1, int(self.dim * mutation_strength))
        indices = np.random.choice(self.dim, num_mutations, replace=False)
        
        for idx in indices:
            vmin, vmax = self.bounds[idx]
            range_val = vmax - vmin
            
            # Gaussian mutation
            delta = np.random.normal(0, range_val * 0.1)
            mutant.values[idx] += delta
        
        mutant._enforce_bounds()
        return mutant
    
    def copy(self) -> 'NumericSolution':
        return NumericSolution(
            values=self.values.copy(),
            bounds=deepcopy(self.bounds),
            names=self.names.copy()
        )
    
    def crossover(self, other: 'NumericSolution') -> 'NumericSolution':
        """Uniform crossover."""
        mask = np.random.rand(self.dim) < 0.5
        offspring_values = np.where(mask, self.values, other.values)
        return NumericSolution(
            values=offspring_values,
            bounds=deepcopy(self.bounds),
            names=self.names.copy()
        )
    
    def get_gene_vector(self) -> np.ndarray:
        if not self.bounds:
            return self.values.copy()
        normalized = []
        for value, (vmin, vmax) in zip(self.values, self.bounds):
            denom = max(1e-8, vmax - vmin)
            normalized.append((float(value) - vmin) / denom)
        return np.array(normalized, dtype=np.float32)


class CombinatoricSolution(Solution):
    """
    Generic combinatoric solution for TSP, job scheduling, etc.
    Represents solutions as sequences or permutations.
    """
    
    def __init__(self, sequence: List[Any], problem_size: int = None):
        """
        Args:
            sequence: List representing the solution (e.g., tour for TSP)
            problem_size: Total problem size (if different from sequence length)
        """
        self.sequence = list(sequence)
        self.problem_size = problem_size or len(sequence)
    
    def to_dict(self) -> Dict[str, Any]:
        return {
            'sequence': self.sequence,
            'length': len(self.sequence),
            'problem_size': self.problem_size
        }
    
    def mutate(self, mutation_strength: float = 0.15) -> 'CombinatoricSolution':
        """2-opt or swap mutation."""
        mutant = self.copy()
        num_mutations = max(1, int(len(mutant.sequence) * mutation_strength))
        
        for _ in range(num_mutations):
            if np.random.rand() < 0.5:
                # 2-opt: reverse a segment
                i, j = np.random.choice(len(mutant.sequence), 2, replace=False)
                if i > j:
                    i, j = j, i
                mutant.sequence[i:j+1] = mutant.sequence[i:j+1][::-1]
            else:
                # Swap mutation
                i, j = np.random.choice(len(mutant.sequence), 2, replace=False)
                mutant.sequence[i], mutant.sequence[j] = mutant.sequence[j], mutant.sequence[i]
        
        return mutant
    
    def copy(self) -> 'CombinatoricSolution':
        return CombinatoricSolution(
            sequence=self.sequence.copy(),
            problem_size=self.problem_size
        )
    
    def crossover(self, other: 'CombinatoricSolution') -> 'CombinatoricSolution':
        """Order crossover (OX)."""
        n = len(self.sequence)
        i, j = np.sort(np.random.choice(n, 2, replace=False))
        
        offspring = [None] * n
        offspring[i:j+1] = self.sequence[i:j+1]
        
        # Fill remaining positions from other, maintaining order
        other_idx = 0
        for pos in range(n):
            if offspring[pos] is None:
                while other_idx < n and other.sequence[other_idx] in offspring:
                    other_idx += 1
                if other_idx < n:
                    offspring[pos] = other.sequence[other_idx]
        
        return CombinatoricSolution(offspring, self.problem_size)

    def get_gene_vector(self) -> np.ndarray:
        if not self.sequence:
            return np.array([])
        denom = max(1, self.problem_size - 1)
        return np.array([val / denom for val in self.sequence], dtype=np.float32)


class HybridSolution(Solution):
    """
    Hybrid solution combining numeric and combinatoric variables.
    E.g., NAS: operations (discrete) + multipliers (continuous)
    """
    
    def __init__(self, 
                 numeric_part: np.ndarray = None,
                 combinatoric_part: List = None,
                 numeric_bounds: List[tuple] = None,
                 numeric_names: List[str] = None):
        self.numeric_part = numeric_part if numeric_part is not None else np.array([])
        self.combinatoric_part = combinatoric_part if combinatoric_part is not None else []
        self.numeric_bounds = numeric_bounds or []
        self.numeric_names = numeric_names or [f"n_{i}" for i in range(len(self.numeric_part))]
    
    def to_dict(self) -> Dict[str, Any]:
        result = {}
        if len(self.numeric_part) > 0:
            for name, val in zip(self.numeric_names, self.numeric_part):
                result[name] = float(val)
        if self.combinatoric_part:
            result['sequence'] = self.combinatoric_part
        return result
    
    def mutate(self, mutation_strength: float = 0.15) -> 'HybridSolution':
        mutant = self.copy()
        
        # Mutate numeric part
        if len(mutant.numeric_part) > 0:
            num_mutations = max(1, int(len(mutant.numeric_part) * mutation_strength))
            indices = np.random.choice(len(mutant.numeric_part), num_mutations, replace=False)
            for idx in indices:
                if idx < len(self.numeric_bounds):
                    vmin, vmax = self.numeric_bounds[idx]
                    delta = np.random.normal(0, (vmax - vmin) * 0.1)
                    mutant.numeric_part[idx] = np.clip(
                        mutant.numeric_part[idx] + delta, vmin, vmax
                    )
        
        # Mutate combinatoric part
        if mutant.combinatoric_part and np.random.rand() < mutation_strength:
            i, j = np.random.choice(len(mutant.combinatoric_part), 2, replace=False)
            mutant.combinatoric_part[i], mutant.combinatoric_part[j] = \
                mutant.combinatoric_part[j], mutant.combinatoric_part[i]
        
        return mutant
    
    def copy(self) -> 'HybridSolution':
        return HybridSolution(
            numeric_part=self.numeric_part.copy() if len(self.numeric_part) > 0 else None,
            combinatoric_part=self.combinatoric_part.copy() if self.combinatoric_part else None,
            numeric_bounds=deepcopy(self.numeric_bounds),
            numeric_names=self.numeric_names.copy()
        )
    
    def crossover(self, other: 'HybridSolution') -> 'HybridSolution':
        """Hybrid crossover."""
        # Numeric: uniform crossover
        numeric_offspring = None
        if len(self.numeric_part) > 0:
            mask = np.random.rand(len(self.numeric_part)) < 0.5
            numeric_offspring = np.where(mask, self.numeric_part, other.numeric_part)
        
        # Combinatoric: order crossover
        combinatoric_offspring = None
        if self.combinatoric_part:
            n = len(self.combinatoric_part)
            i, j = np.sort(np.random.choice(n, 2, replace=False))
            offspring = [None] * n
            offspring[i:j+1] = self.combinatoric_part[i:j+1]
            
            other_idx = 0
            for pos in range(n):
                if offspring[pos] is None:
                    while other_idx < n and other.combinatoric_part[other_idx] in offspring:
                        other_idx += 1
                    if other_idx < n:
                        offspring[pos] = other.combinatoric_part[other_idx]
            combinatoric_offspring = offspring
        
        return HybridSolution(
            numeric_part=numeric_offspring,
            combinatoric_part=combinatoric_offspring,
            numeric_bounds=deepcopy(self.numeric_bounds),
            numeric_names=self.numeric_names.copy()
        )
    
    def get_gene_vector(self) -> np.ndarray:
        parts = []
        if len(self.numeric_part) > 0:
            if self.numeric_bounds:
                for value, (vmin, vmax) in zip(self.numeric_part, self.numeric_bounds):
                    denom = max(1e-8, vmax - vmin)
                    parts.append((float(value) - vmin) / denom)
            else:
                parts.extend([float(v) for v in self.numeric_part])
        if self.combinatoric_part:
            max_val = max(1, max(self.combinatoric_part))
            parts.extend([val / max_val for val in self.combinatoric_part])
        return np.array(parts, dtype=np.float32)
