"""NAS-Bench-201 Benchmark Wrapper for DPO-NAS"""
import logging
import numpy as np
from typing import Dict, List, Optional

logger = logging.getLogger(__name__)


class NASBench201Benchmark:
    """Wrapper for NAS-Bench-201 evaluation"""
    
    def __init__(self, data_dir: Optional[str] = None):
        """
        Initialize NAS-Bench-201 API
        
        Args:
            data_dir: Directory containing NAS-Bench-201 data
        """
        self.api = None
        self.cache = {}
        self.data_dir = data_dir
        
        try:
            from nas_201_api import NASBench201API as API
            file_path = f"{data_dir}/NAS-Bench-201-v1_1-096897.pth" if data_dir else None
            if file_path:
                self.api = API(file_path)
                logger.info("NAS-Bench-201 API initialized successfully")
            else:
                logger.warning("NAS-Bench-201 data_dir not provided, using simulation mode")
        except ImportError:
            logger.warning("nas_201_api library not installed, using simulation mode")
        except Exception as e:
            logger.warning(f"Failed to initialize NAS-Bench-201 API: {e}, using simulation mode")
    
    def _convert_to_nasbench201_format(self, arch_dict: Dict) -> str:
        """
        Convert DPO-NAS architecture to NAS-Bench-201 6-edge encoding string
        
        Args:
            arch_dict: DPO-NAS architecture dictionary
            
        Returns:
            NAS-Bench-201 architecture string (e.g., '|none~0|+|none~0|skip_connect~1|+|...')
        """
        # Map DPO operations to NAS-Bench-201 operations
        op_mapping = {
            'conv_3x3': 'nor_conv_3x3',
            'conv_5x5': 'nor_conv_3x3',  # NB201 doesn't have 5x5
            'conv_1x1': 'nor_conv_1x1',
            'dw_conv': 'nor_conv_3x3',
            'sep_conv': 'nor_conv_3x3',
            'max_pool': 'avg_pool_3x3',
            'avg_pool': 'avg_pool_3x3',
            'skip_connect': 'skip_connect',
        }
        
        operations = arch_dict.get('operations', [])
        skip_connections = arch_dict.get('skip_connections', [])
        
        # NAS-Bench-201 has 6 edges in the cell
        edges = []
        op_idx = 0
        
        # Edge pattern: (0->1), (0->2, 1->2), (0->3, 1->3, 2->3)
        for i in range(1, 4):  # nodes 1, 2, 3
            for j in range(i):  # connections from previous nodes
                if op_idx < len(operations):
                    op = operations[op_idx]
                    # Use skip if enabled in config
                    if op_idx < len(skip_connections) and skip_connections[op_idx]:
                        mapped_op = 'skip_connect'
                    else:
                        mapped_op = op_mapping.get(op, 'nor_conv_3x3')
                    edges.append(f"{mapped_op}~{j}")
                    op_idx += 1
        
        # Format: |edge0|+|edge1|edge2|+|edge3|edge4|edge5|
        arch_str = f"|{edges[0]}|+|{edges[1]}|{edges[2]}|+|{edges[3]}|{edges[4]}|{edges[5]}|"
        return arch_str
    
    def evaluate_multi_dataset(self, arch_dict: Dict) -> Dict:
        """
        Evaluate architecture on multiple datasets (CIFAR-10, CIFAR-100, ImageNet16-120)
        
        Args:
            arch_dict: DPO-NAS architecture dictionary
            
        Returns:
            Dictionary with results for each dataset
        """
        arch_hash = str(arch_dict)
        
        if arch_hash in self.cache:
            logger.info("Returning cached NAS-Bench-201 result")
            return self.cache[arch_hash]
        
        try:
            if self.api:
                arch_str = self._convert_to_nasbench201_format(arch_dict)
                arch_index = self.api.query_index_by_arch(arch_str)
                
                results = {}
                for dataset in ['cifar10', 'cifar100', 'ImageNet16-120']:
                    info = self.api.get_more_info(arch_index, dataset, hp='200', is_random=False)
                    results[dataset] = {
                        'accuracy': info['test-accuracy'],
                        'latency': info.get('latency', 0),
                        'parameters': info.get('params', 0),
                    }
            else:
                results = self._simulate_multi_dataset_evaluation(arch_dict)
            
            self.cache[arch_hash] = results
            return results
            
        except Exception as e:
            logger.error(f"NAS-Bench-201 evaluation failed: {e}, using simulation")
            return self._simulate_multi_dataset_evaluation(arch_dict)
    
    def _simulate_multi_dataset_evaluation(self, arch_dict: Dict) -> Dict:
        """Simulate multi-dataset evaluation"""
        num_layers = len(arch_dict.get('operations', []))
        num_skips = sum(arch_dict.get('skip_connections', []))
        
        base_accuracies = {
            'cifar10': 0.91,
            'cifar100': 0.71,
            'ImageNet16-120': 0.43,
        }
        
        results = {}
        for dataset, base_acc in base_accuracies.items():
            acc_bonus = 0.02 * (num_skips / max(1, len(arch_dict.get('skip_connections', [1]))))
            accuracy = base_acc + acc_bonus + np.random.normal(0, 0.01)
            
            results[dataset] = {
                'accuracy': float(np.clip(accuracy, 0, 1)),
                'latency': float(10 + num_layers * 2 + np.random.uniform(-1, 1)),
                'parameters': float(num_layers * 100000),
            }
        
        return results
    
    def robustness_analysis(self, arch_dict: Dict) -> Dict:
        """
        Analyze architecture robustness across datasets
        
        Args:
            arch_dict: DPO-NAS architecture dictionary
            
        Returns:
            Dictionary with robustness metrics
        """
        results = self.evaluate_multi_dataset(arch_dict)
        
        accuracies = [res['accuracy'] for res in results.values()]
        
        return {
            'mean_accuracy': float(np.mean(accuracies)),
            'std_accuracy': float(np.std(accuracies)),
            'min_accuracy': float(np.min(accuracies)),
            'max_accuracy': float(np.max(accuracies)),
            'dataset_results': results,
        }
