"""Comprehensive Benchmark Evaluator for DPO-NAS"""
import logging
import json
import numpy as np
from pathlib import Path
from typing import Dict, List, Optional
from datetime import datetime

from .nasbench101_wrapper import NASBench101Benchmark
from .nasbench201_wrapper import NASBench201Benchmark
from .nasbench301_wrapper import NASBench301Benchmark
from .hpobench_wrapper import HPOBenchBenchmark
from .nats_bench_wrapper import NATSBenchBenchmark

logger = logging.getLogger(__name__)


class ComprehensiveBenchmarkEvaluator:
    """Unified evaluator for all NAS benchmarks"""

    def __init__(self, data_dir: Optional[str] = None):
        self.data_dir = data_dir
        logger.info("Initializing comprehensive benchmark evaluator")

        self.benchmarks = {
            'nasbench101': NASBench101Benchmark(data_dir),
            'nasbench201': NASBench201Benchmark(data_dir),
            'nasbench301': NASBench301Benchmark(data_dir),
            'hpobench': HPOBenchBenchmark(),
            'nats_bench': NATSBenchBenchmark(data_dir),
        }

        self.evaluation_history = []
        logger.info("All benchmarks initialized")

    def evaluate_architecture(self, arch_dict: Dict) -> Dict:
        logger.info("Starting comprehensive architecture evaluation")

        results = {
            'architecture': arch_dict,
            'timestamp': datetime.now().isoformat(),
            'benchmarks': {},
        }

        try:
            nb101_result = self.benchmarks['nasbench101'].evaluate(arch_dict)
            results['benchmarks']['nasbench101'] = {
                'validation_accuracy': nb101_result.validation_accuracy,
                'test_accuracy': nb101_result.test_accuracy,
                'training_time': nb101_result.training_time,
                'computational_cost': nb101_result.computational_cost,
            }
        except Exception as e:
            results['benchmarks']['nasbench101'] = {'error': str(e)}

        try:
            nb201_result = self.benchmarks['nasbench201'].robustness_analysis(arch_dict)
            results['benchmarks']['nasbench201'] = nb201_result
        except Exception as e:
            results['benchmarks']['nasbench201'] = {'error': str(e)}

        try:
            nb301_result = self.benchmarks['nasbench301'].predict_performance(arch_dict)
            results['benchmarks']['nasbench301'] = nb301_result
        except Exception as e:
            results['benchmarks']['nasbench301'] = {'error': str(e)}

        try:
            hpo_result = self.benchmarks['hpobench'].evaluate_anytime(arch_dict)
            results['benchmarks']['hpobench'] = hpo_result
        except Exception as e:
            results['benchmarks']['hpobench'] = {'error': str(e)}

        try:
            nats_scaling = self.benchmarks['nats_bench'].evaluate_scaling(arch_dict)
            nats_transfer = self.benchmarks['nats_bench'].evaluate_transfer(arch_dict)
            results['benchmarks']['nats_bench'] = {
                'scaling': nats_scaling,
                'transfer': nats_transfer,
            }
        except Exception as e:
            results['benchmarks']['nats_bench'] = {'error': str(e)}

        self.evaluation_history.append(results)
        return results

    def evaluate_population(self, architectures: List[Dict]) -> Dict:
        individual_results = [self.evaluate_architecture(arch) for arch in architectures]
        summary = self._compute_summary(individual_results)
        return {
            'individual_results': individual_results,
            'summary': summary,
            'population_size': len(architectures),
        }

    def _compute_summary(self, evaluations: List[Dict]) -> Dict:
        summary = {}

        nb101_accs = [
            e['benchmarks'].get('nasbench101', {}).get('test_accuracy')
            for e in evaluations
            if 'test_accuracy' in e['benchmarks'].get('nasbench101', {})
        ]
        if nb101_accs:
            summary['nasbench101'] = {
                'mean_accuracy': float(np.mean(nb101_accs)),
                'std_accuracy': float(np.std(nb101_accs)),
                'max_accuracy': float(np.max(nb101_accs)),
                'min_accuracy': float(np.min(nb101_accs)),
            }

        nb201_accs = [
            e['benchmarks'].get('nasbench201', {}).get('mean_accuracy')
            for e in evaluations
            if 'mean_accuracy' in e['benchmarks'].get('nasbench201', {})
        ]
        if nb201_accs:
            summary['nasbench201'] = {
                'mean_accuracy': float(np.mean(nb201_accs)),
                'std_accuracy': float(np.std(nb201_accs)),
                'max_accuracy': float(np.max(nb201_accs)),
                'min_accuracy': float(np.min(nb201_accs)),
            }

        return summary

    def save_results(self, filepath: str) -> None:
        output_path = Path(filepath)
        output_path.parent.mkdir(parents=True, exist_ok=True)
        with open(output_path, 'w') as f:
            json.dump(self.evaluation_history, f, indent=2, default=str)

    def generate_report(self, output_dir: str) -> str:
        output_path = Path(output_dir)
        output_path.mkdir(parents=True, exist_ok=True)
        report_file = output_path / f"benchmark_report_{datetime.now().strftime('%Y%m%d_%H%M%S')}.md"

        with open(report_file, 'w') as f:
            f.write("# DPO-NAS Comprehensive Benchmark Report\n\n")
            f.write(f"Generated: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}\n\n")
            f.write(f"Total Evaluations: {len(self.evaluation_history)}\n\n")

        return str(report_file)
