

import numpy as np
import logging
from typing import Dict, Optional
from collections import deque

from .config import DPO_Config
from .agent import SearchAgent
from ..architecture.gene import ArchitectureGene
from ..evaluation.ensemble import EnsembleEstimator
from ..constraints.handler import AdvancedConstraintHandler
from ..utils.logger import get_logger

class DPO_NAS:
    def __init__(self,
                 config: DPO_Config,
                 estimator: Optional[EnsembleEstimator] = None,
                 constraint_handler: Optional[AdvancedConstraintHandler] = None,
                 logger: Optional[logging.Logger] = None):
        self.config = config
        self.estimator = estimator or EnsembleEstimator()
        self.constraint_handler = constraint_handler or AdvancedConstraintHandler(config)
        self.logger = logger or get_logger('DPO-NAS', self.config.verbose)
        self.islands = [[] for _ in range(config.num_islands if config.island_model else 1)]
        self.best_agent: Optional[SearchAgent] = None
        self.best_fitness: float = float('inf')
        self.best_per_island = [None] * len(self.islands)
        self.history = {
            'iterations': [],
            'best_fitness': [],
            'avg_fitness': [],
            'worst_fitness': [],
            'improvement_rate': [],
            'best_architecture': None,
        }
        self.improvement_tracker = deque(maxlen=30)
        
        # Pre-calculate bounds for enforcement
        self._ops_len = len(ArchitectureGene.OPERATIONS) - 0.01
        self._kers_len = len(ArchitectureGene.KERNELS) - 0.01

    def initialize_population(self) -> None:
        agents_per_island = max(1, self.config.population_size // len(self.islands))
        self.logger.info(f"Initializing population: {self.config.population_size} agents")
        
        for island_idx, island in enumerate(self.islands):
            for _ in range(agents_per_island):
                gene = ArchitectureGene()
                # Create agent
                agent = SearchAgent(gene=gene, island_id=island_idx)
                # Evaluate
                self._evaluate_agent(agent, iteration=0)
                island.append(agent)
                
                # Update bests
                if agent.fitness < self.best_fitness:
                    self.best_fitness = agent.fitness
                    self.best_agent = agent
                
                if self.best_per_island[island_idx] is None or agent.fitness < self.best_per_island[island_idx].fitness:
                    self.best_per_island[island_idx] = agent

    def _evaluate_agent(self, agent: SearchAgent, iteration: int) -> None:
        arch_dict = agent.gene.to_architecture_dict()
        loss, metrics = self.estimator.estimate(arch_dict)
        penalty = self.constraint_handler.compute_adaptive_penalty(metrics, iteration, self.config.max_iterations)
        
        # Fast arithmetic
        lat_norm = metrics['latency_ms'] / self.config.latency_constraint
        mem_norm = metrics['memory_mb'] / self.config.memory_constraint
        flop_norm = metrics['flops_m'] / self.config.flops_constraint
        
        agent.fitness = (
            self.config.w_loss * loss +
            self.config.w_latency * lat_norm +
            self.config.w_memory * mem_norm +
            self.config.w_flops * flop_norm +
            penalty
        )
        agent.metrics = metrics

    def _get_adaptive_parameters(self, iteration: int) -> Dict[str, float]:
        t = iteration / max(1, self.config.max_iterations)
        
        # Vectorized-style scalar operations
        decay = (1 - t) ** self.config.decay_power
        alpha = self.config.alpha_0 * decay
        beta = self.config.beta_0 + 0.5 * t
        gamma = self.config.gamma_0 * (1 - 0.5 * t)
        delta = self.config.delta_0 * (1 - 0.3 * t)
        
        if self.config.adaptive_alpha and len(self.improvement_tracker) > 10:
            # sum / len is faster than np.mean for deque
            it = self.improvement_tracker
            recent_avg = sum(list(it)[-10:]) / 10.0
            alpha *= 1.2 if recent_avg < 0.001 else 0.9
            
        return {
            'alpha': float(np.clip(alpha, 0.01, 0.3)),
            'beta': float(np.clip(beta, 0.8, 1.5)),
            'gamma': float(np.clip(gamma, 0.5, 1.5)),
            'delta': float(np.clip(delta, 0.05, 0.4)),
        }

    def _clip_gene_array(self, gene_arr: np.ndarray, num_layers: int, num_cells: int) -> None:
        """Helper to enforce bounds in-place on numpy array"""
        # Operations
        np.clip(gene_arr[:num_layers], 0, self._ops_len, out=gene_arr[:num_layers])
        # Kernels
        np.clip(gene_arr[num_layers:2*num_layers], 0, self._kers_len, out=gene_arr[num_layers:2*num_layers])
        # Skip connections
        np.clip(gene_arr[2*num_layers:2*num_layers+num_cells], 0, 1, out=gene_arr[2*num_layers:2*num_layers+num_cells])
        # Multipliers
        np.clip(gene_arr[-2:], 0.3, 2.0, out=gene_arr[-2:])

    def _dpo_step(self, agent: SearchAgent, params: Dict, iteration: int) -> SearchAgent:
        # OPTIMIZED: Perform vector math on arrays without creating intermediate ArchitectureGene objects
        # This reduces object overhead significantly
        
        current_vec = agent.gene.gene
        
        # 1. Debt Gene Logic (Agent + Perturbation)
        perturbation = np.random.randn(agent.gene.D) * params['alpha']
        debt_vec = current_vec + perturbation
        self._clip_gene_array(debt_vec, agent.gene.num_layers, agent.gene.num_cells)
        
        # Debt Vector
        debt_delta = debt_vec - current_vec
        
        # 2. Repay Gene Logic (Agent - Beta * Debt)
        repay_vec = current_vec - (params['beta'] * debt_delta)
        self._clip_gene_array(repay_vec, agent.gene.num_layers, agent.gene.num_cells)
        
        # 3. Double Gene Logic (Repay - Gamma * Debt)
        double_vec = repay_vec - (params['gamma'] * debt_delta)
        self._clip_gene_array(double_vec, agent.gene.num_layers, agent.gene.num_cells)
        
        # 4. Smart Gene Logic (Double + Delta * (GlobalBest - Current))
        best_global = self.best_agent or agent
        smart_vec = double_vec + params['delta'] * (best_global.gene.gene - current_vec)
        self._clip_gene_array(smart_vec, agent.gene.num_layers, agent.gene.num_cells)
        
        # Create Final Object
        smart_gene = ArchitectureGene(agent.gene.num_layers, agent.gene.num_cells)
        smart_gene.gene = smart_vec
        
        candidate = SearchAgent(gene=smart_gene, island_id=agent.island_id)
        self._evaluate_agent(candidate, iteration)
        
        if candidate.fitness < agent.fitness:
            candidate.improvements = agent.improvements + 1
            return candidate
            
        agent.age += 1
        return agent

    def _migrate_between_islands(self) -> None:
        if not self.config.island_model or len(self.islands) < 2:
            return
            
        num_migrate = max(1, self.config.population_size // 20)
        
        for island_idx in range(len(self.islands)):
            # Find best neighbors
            best_neighbors = [
                self.best_per_island[o] 
                for o in range(len(self.islands)) 
                if o != island_idx and self.best_per_island[o]
            ]
            
            if best_neighbors:
                best_neighbor = min(best_neighbors, key=lambda a: a.fitness)
                # Sort current island
                self.islands[island_idx].sort(key=lambda a: a.fitness, reverse=True) # Worst first
                
                # Replace worst
                for i in range(num_migrate):
                    if i < len(self.islands[island_idx]):
                        # Deep copy gene
                        new_gene = best_neighbor.gene.copy()
                        new_agent = SearchAgent(gene=new_gene, island_id=island_idx)
                        self._evaluate_agent(new_agent, iteration=0)
                        self.islands[island_idx][i] = new_agent

    def _inject_diversity(self) -> None:
        for island in self.islands:
            n = len(island)
            elite_count = max(1, int(n * self.config.elite_ratio))
            island.sort(key=lambda a: a.fitness)
            elite = island[:elite_count]
            
            # Replace non-elites
            for i in range(elite_count, n):
                if elite:
                    parent = elite[np.random.randint(0, len(elite))]
                    mutant_gene = parent.gene.mutate()
                    new_agent = SearchAgent(gene=mutant_gene, island_id=parent.island_id)
                    self._evaluate_agent(new_agent, iteration=0)
                    island[i] = new_agent

    def optimize(self) -> Dict:
        self.logger.info("=" * 70)
        self.logger.info("Starting DPO-NAS Optimization")
        self.logger.info(f"Strategy: {self.config.eval_strategy}")
        self.logger.info(f"Population: {self.config.population_size} | Islands: {len(self.islands)}")
        self.logger.info("=" * 70)
        
        self.initialize_population()
        
        for iteration in range(self.config.max_iterations):
            params = self._get_adaptive_parameters(iteration)
            
            for island in self.islands:
                for idx, agent in enumerate(island):
                    new_agent = self._dpo_step(agent, params, iteration)
                    island[idx] = new_agent
                    
                    if new_agent.fitness < self.best_fitness:
                        self.best_fitness = new_agent.fitness
                        self.best_agent = new_agent
                        self.logger.info(f"[Iter {iteration}] New best: {self.best_fitness:.4f}")

            # Periodic tasks
            if (iteration + 1) % self.config.migration_freq == 0 and self.config.island_model:
                self._migrate_between_islands()
                
            if (iteration + 1) % self.config.diversity_inject_freq == 0:
                self._inject_diversity()

            # Tracking
            # Optim: generator expression for fitness sum/stats
            all_fitness = np.array([a.fitness for island in self.islands for a in island])
            
            improvement = (self.history['best_fitness'][-1] - self.best_fitness) if self.history['best_fitness'] else 0.0
            self.improvement_tracker.append(improvement)
            
            self.history['iterations'].append(iteration)
            self.history['best_fitness'].append(self.best_fitness)
            self.history['avg_fitness'].append(float(np.mean(all_fitness)))
            self.history['worst_fitness'].append(float(np.max(all_fitness)))
            self.history['improvement_rate'].append(float(improvement))
            
            if iteration % 10 == 0:
                self.logger.info(f"[Iter {iteration:3d}] Best: {self.best_fitness:.4f} | Avg: {np.mean(all_fitness):.4f} | Improvement: {improvement:.6f}")

        if self.best_agent:
            self.history['best_architecture'] = self.best_agent.gene.to_architecture_dict()

        self.logger.info("=" * 70)
        self.logger.info("Optimization Complete!")
        self.logger.info(f"Best Fitness: {self.best_fitness:.4f}")
        self.logger.info("=" * 70)

        return {
            'best_fitness': self.best_fitness,
            'best_architecture': self.history.get('best_architecture'),
            'best_metrics': getattr(self.best_agent, 'metrics', {}),
            'history': self.history,
            'config': {
                'population_size': self.config.population_size,
                'islands': len(self.islands),
                'eval_strategy': self.config.eval_strategy,
                'max_iterations': len(self.history['iterations']),
            }
        }