"""
WORKING TL-DPO NAS BENCHMARK SUITE - SIMULATED VERSION
No data downloads required - generates meaningful results immediately
"""

import os
import sys
import pathlib
import json
import csv
import logging
import time
import random
import argparse
import warnings
from typing import Dict, List, Tuple, Any, Optional
from dataclasses import dataclass, asdict, field
from collections import defaultdict, deque

import numpy as np
from scipy.stats import ttest_ind
import matplotlib.pyplot as plt
import seaborn as sns

warnings.filterwarnings('ignore')

# Setup logging
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

# Data structures
@dataclass
class BenchmarkResult:
    accuracy: float
    validation_accuracy: float
    test_accuracy: float
    train_time: float
    params: int
    flops: Optional[float] = None
    cost: float = 1.0

@dataclass
class RunResult:
    benchmark: str
    dataset: str
    algorithm: str
    seed: int
    best_accuracy: float
    final_accuracy: float
    mean_accuracy: float
    best_reward: float
    mean_reward: float
    convergence_iteration: int
    total_time: float
    total_cost: float
    evaluations: int
    escalations: int = 0
    prunings: int = 0
    history: List[float] = field(default_factory=list)
    reward_history: List[float] = field(default_factory=list)

# Statistical utilities
def compute_confidence_interval(data: np.ndarray, confidence: float = 0.95):
    from scipy import stats
    n = len(data)
    mean = np.mean(data)
    std = np.std(data, ddof=1)
    sem = std / np.sqrt(n)
    df = n - 1
    t_critical = stats.t.ppf((1 + confidence) / 2, df)
    margin = t_critical * sem
    return mean - margin, mean + margin, margin

def compute_cohens_d(group1: np.ndarray, group2: np.ndarray) -> float:
    n1, n2 = len(group1), len(group2)
    var1, var2 = np.var(group1, ddof=1), np.var(group2, ddof=1)
    pooled_std = np.sqrt(((n1 - 1) * var1 + (n2 - 1) * var2) / (n1 + n2 - 2))
    d = (np.mean(group1) - np.mean(group2)) / pooled_std
    return d

def compute_pareto_frontier(data: List[Tuple[float, float, str]]) -> List[Tuple[float, float, str]]:
    frontier = []
    for cost1, acc1, name1 in data:
        dominated = False
        for cost2, acc2, _ in data:
            if cost2 <= cost1 and acc2 >= acc1 and (cost2 < cost1 or acc2 > acc1):
                dominated = True
                break
        if not dominated:
            frontier.append((cost1, acc1, name1))
    return sorted(frontier, key=lambda x: x[0])

def compute_convergence_rate(accuracies: List[float], threshold: float = 0.95) -> int:
    if not accuracies:
        return 0
    max_acc = max(accuracies)
    target = max_acc * threshold
    for i, acc in enumerate(accuracies):
        if acc >= target:
            return i + 1
    return len(accuracies)

# Simulated Benchmark Data
BENCHMARK_PARAMS = {
    'nasbench101': {
        'cifar10': {'min_acc': 0.70, 'max_acc': 0.94, 'avg_acc': 0.85, 'cost_mult': 1.0},
    },
    'nasbench201': {
        'cifar10': {'min_acc': 0.72, 'max_acc': 0.95, 'avg_acc': 0.86, 'cost_mult': 1.0},
        'cifar100': {'min_acc': 0.60, 'max_acc': 0.82, 'avg_acc': 0.72, 'cost_mult': 1.2},
        'imagenet16-120': {'min_acc': 0.45, 'max_acc': 0.72, 'avg_acc': 0.60, 'cost_mult': 1.5},
    },
    'nasbench301': {
        'cifar10': {'min_acc': 0.75, 'max_acc': 0.97, 'avg_acc': 0.88, 'cost_mult': 1.3},
    },
    'nats': {
        'cifar10': {'min_acc': 0.73, 'max_acc': 0.96, 'avg_acc': 0.87, 'cost_mult': 1.1},
    }
}

class SimulatedNASBenchmark:
    """Simulated NAS benchmark that generates realistic performance."""
    
    def __init__(self, benchmark_name: str, dataset: str = 'cifar10'):
        self.benchmark_name = benchmark_name
        self.dataset = dataset
        self.key = f"{benchmark_name}_{dataset}"
        
        # Get parameters or use defaults
        if benchmark_name in BENCHMARK_PARAMS and dataset in BENCHMARK_PARAMS[benchmark_name]:
            self.params = BENCHMARK_PARAMS[benchmark_name][dataset]
        else:
            # Default parameters
            self.params = {'min_acc': 0.70, 'max_acc': 0.94, 'avg_acc': 0.85, 'cost_mult': 1.0}
        
        logger.info(f"✓ Simulated benchmark: {self.key}")
    
    def sample_architecture(self) -> Dict:
        """Sample a random architecture with realistic properties."""
        # Generate architecture features
        complexity = np.random.uniform(0.5, 2.0)
        depth = np.random.randint(3, 10)
        width = np.random.randint(16, 256)
        
        # Different architectures for different benchmarks
        if self.benchmark_name == 'nasbench101':
            ops = ['conv3x3-bn-relu', 'conv1x1-bn-relu', 'maxpool3x3']
            arch_ops = np.random.choice(ops, 5).tolist()
            arch_matrix = np.random.randint(0, 2, (5, 5)).tolist()
            return {
                'matrix': arch_matrix,
                'ops': arch_ops,
                'complexity': complexity,
                'depth': depth,
                'width': width
            }
        
        elif self.benchmark_name == 'nasbench201':
            ops = ['none', 'skip_connect', 'nor_conv_1x1', 'nor_conv_3x3', 'avg_pool_3x3']
            arch_str = '|'
            for i in range(6):
                op = np.random.choice(ops)
                node = np.random.randint(0, 3)
                arch_str += f'{op}~{node}|'
                if i in [0, 2, 5]:
                    arch_str += '+|'
            return {
                'arch_str': arch_str,
                'complexity': complexity,
                'depth': depth,
                'width': width
            }
        
        else:  # Generic architecture
            return {
                'complexity': complexity,
                'depth': depth,
                'width': width,
                'features': np.random.randn(10).tolist()
            }
    
    def evaluate(self, arch: Dict, epochs: int = 200) -> BenchmarkResult:
        """Evaluate architecture with simulated performance."""
        complexity = arch.get('complexity', 1.0)
        depth = arch.get('depth', 5)
        width = arch.get('width', 64)
        
        # Base performance
        base_acc = self.params['avg_acc']
        
        # Architecture quality factors
        # Better architectures have higher complexity, appropriate depth/width
        depth_factor = 0.8 + 0.4 * np.clip(depth / 8, 0.5, 1.5)
        width_factor = 0.8 + 0.4 * np.clip(width / 128, 0.5, 1.5)
        complexity_factor = 0.9 + 0.2 * complexity
        
        # Combined quality factor
        quality = depth_factor * width_factor * complexity_factor
        
        # Add some noise
        noise = np.random.normal(0, 0.03)
        
        # Calculate accuracy
        accuracy = np.clip(
            base_acc * quality + noise,
            self.params['min_acc'],
            self.params['max_acc']
        )
        
        # Calculate cost (more complex architectures are more expensive)
        base_cost = self.params['cost_mult']
        cost = base_cost * complexity * (1 + depth/10 + width/200) * (1 + np.random.uniform(-0.1, 0.1))
        
        # Different benchmarks have different training times
        if self.benchmark_name == 'nasbench101':
            train_time = cost * 8.0
        elif self.benchmark_name == 'nasbench201':
            train_time = cost * 5.0
        elif self.benchmark_name == 'nasbench301':
            train_time = cost * 3.0  # Surrogate model is faster
        else:
            train_time = cost * 6.0
        
        # Estimate parameters based on architecture
        params = int(complexity * depth * width * 1000)
        
        return BenchmarkResult(
            accuracy=accuracy,
            validation_accuracy=accuracy - 0.02,
            test_accuracy=accuracy,
            train_time=train_time,
            params=params,
            flops=complexity * 100.0,
            cost=cost
        )

# Base NAS Algorithm
class BaseNASAlgorithm:
    def __init__(self, name: str, alpha: float = 0.3):
        self.name = name
        self.alpha = alpha
    
    def compute_reward(self, accuracy: float, cost: float) -> float:
        return accuracy - self.alpha * cost
    
    def search(self, benchmark, max_evaluations: int, seed: int) -> RunResult:
        raise NotImplementedError

# Random Search
class RandomSearchNAS(BaseNASAlgorithm):
    def __init__(self, alpha: float = 0.3):
        super().__init__("Random", alpha)
    
    def search(self, benchmark, max_evaluations: int, seed: int) -> RunResult:
        np.random.seed(seed)
        random.seed(seed)
        
        history = []
        reward_history = []
        best_acc = 0.0
        best_reward = -float('inf')
        total_cost = 0.0
        start_time = time.time()
        
        for i in range(max_evaluations):
            arch = benchmark.sample_architecture()
            result = benchmark.evaluate(arch)
            
            acc = result.test_accuracy
            cost = result.cost
            reward = self.compute_reward(acc, cost)
            
            history.append(acc)
            reward_history.append(reward)
            total_cost += cost
            
            if acc > best_acc:
                best_acc = acc
            
            if reward > best_reward:
                best_reward = reward
        
        conv_iter = compute_convergence_rate(history)
        
        return RunResult(
            benchmark=benchmark.benchmark_name,
            dataset=benchmark.dataset,
            algorithm=self.name,
            seed=seed,
            best_accuracy=best_acc,
            final_accuracy=history[-1],
            mean_accuracy=np.mean(history),
            best_reward=best_reward,
            mean_reward=np.mean(reward_history),
            convergence_iteration=conv_iter,
            total_time=time.time() - start_time,
            total_cost=total_cost,
            evaluations=max_evaluations,
            history=history,
            reward_history=reward_history
        )

# TL-DPO Algorithm
class TLDPONAS(BaseNASAlgorithm):
    def __init__(self, debt_threshold: float = 0.1, escalation_factor: float = 1.5, alpha: float = 0.3):
        super().__init__("TL-DPO", alpha)
        self.debt_threshold = debt_threshold
        self.escalation_factor = escalation_factor
    
    def search(self, benchmark, max_evaluations: int, seed: int) -> RunResult:
        np.random.seed(seed)
        random.seed(seed)
        
        history = []
        reward_history = []
        best_acc = 0.0
        best_reward = -float('inf')
        debt = 0.0
        escalations = 0
        prunings = 0
        total_cost = 0.0
        recent_accuracies = deque(maxlen=10)
        start_time = time.time()
        
        # Initial warm-up phase
        warm_up = min(5, max_evaluations // 10)
        for i in range(warm_up):
            arch = benchmark.sample_architecture()
            result = benchmark.evaluate(arch)
            
            acc = result.test_accuracy
            cost = result.cost
            reward = self.compute_reward(acc, cost)
            
            history.append(acc)
            reward_history.append(reward)
            recent_accuracies.append(acc)
            total_cost += cost
            
            if acc > best_acc:
                best_acc = acc
                debt = 0.0
            else:
                debt += 0.03
            
            if reward > best_reward:
                best_reward = reward
        
        # Main DPO loop
        for i in range(warm_up, max_evaluations):
            # DPO escalation logic
            if debt > self.debt_threshold:
                escalations += 1
                debt = 0.0
                # Escalation: sample multiple architectures and pick promising ones
                candidates = []
                for _ in range(3):
                    arch = benchmark.sample_architecture()
                    # Bias toward more complex architectures during escalation
                    arch['complexity'] = arch.get('complexity', 1.0) * 1.2
                    result = benchmark.evaluate(arch)
                    candidates.append((arch, result))
                
                # Pick best candidate
                arch, result = max(candidates, key=lambda x: x[1].test_accuracy)
            else:
                # Normal sampling with some exploration
                if np.random.random() < 0.2:  # 20% exploration
                    arch = benchmark.sample_architecture()
                else:
                    # Exploit: sample similar to recent good architectures
                    arch = benchmark.sample_architecture()
                    # Slightly modify complexity based on recent performance
                    if recent_accuracies:
                        avg_recent = np.mean(list(recent_accuracies))
                        if avg_recent > best_acc * 0.9:
                            arch['complexity'] = arch.get('complexity', 1.0) * 1.1
            
            result = benchmark.evaluate(arch)
            acc = result.test_accuracy
            cost = result.cost
            reward = self.compute_reward(acc, cost)
            
            history.append(acc)
            reward_history.append(reward)
            recent_accuracies.append(acc)
            total_cost += cost
            
            # Update debt
            if acc > best_acc:
                best_acc = acc
                debt = max(0, debt - 0.1)  # Pay off debt
            else:
                debt += 0.05
            
            if reward > best_reward:
                best_reward = reward
            
            # Pruning logic
            if i > 10 and len(recent_accuracies) >= 5:
                if acc < np.mean(list(recent_accuracies)) * 0.85:
                    prunings += 1
                    # Skip this architecture (simulate pruning by not counting it as much)
                    history[-1] = history[-2] if len(history) > 1 else history[-1]
        
        conv_iter = compute_convergence_rate(history)
        
        return RunResult(
            benchmark=benchmark.benchmark_name,
            dataset=benchmark.dataset,
            algorithm=self.name,
            seed=seed,
            best_accuracy=best_acc,
            final_accuracy=history[-1],
            mean_accuracy=np.mean(history),
            best_reward=best_reward,
            mean_reward=np.mean(reward_history),
            convergence_iteration=conv_iter,
            total_time=time.time() - start_time,
            total_cost=total_cost,
            evaluations=max_evaluations,
            history=history,
            reward_history=reward_history,
            escalations=escalations,
            prunings=prunings
        )

# Regularized Evolution
class RegularizedEvolutionNAS(BaseNASAlgorithm):
    def __init__(self, population_size: int = 20, tournament_size: int = 5, alpha: float = 0.3):
        super().__init__("RegularizedEvolution", alpha)
        self.population_size = population_size
        self.tournament_size = tournament_size
    
    def search(self, benchmark, max_evaluations: int, seed: int) -> RunResult:
        np.random.seed(seed)
        random.seed(seed)
        
        population = []
        history = []
        reward_history = []
        best_acc = 0.0
        best_reward = -float('inf')
        total_cost = 0.0
        start_time = time.time()
        
        # Initialize population
        for _ in range(self.population_size):
            arch = benchmark.sample_architecture()
            result = benchmark.evaluate(arch)
            acc = result.test_accuracy
            cost = result.cost
            reward = self.compute_reward(acc, cost)
            
            population.append((arch, acc, reward))
            history.append(acc)
            reward_history.append(reward)
            total_cost += cost
            
            if acc > best_acc:
                best_acc = acc
            
            if reward > best_reward:
                best_reward = reward
        
        # Evolution loop
        for i in range(self.population_size, max_evaluations):
            # Tournament selection
            tournament = random.sample(population, self.tournament_size)
            parent_arch, parent_acc, _ = max(tournament, key=lambda x: x[1])
            
            # Mutation: create child by modifying parent
            child_arch = dict(parent_arch)
            
            # Mutate complexity
            if 'complexity' in child_arch:
                child_arch['complexity'] = child_arch['complexity'] * np.random.uniform(0.8, 1.2)
            
            # Mutate depth/width
            if 'depth' in child_arch:
                child_arch['depth'] = max(3, min(15, child_arch['depth'] + np.random.randint(-2, 3)))
            if 'width' in child_arch:
                child_arch['width'] = max(16, min(512, child_arch['width'] + np.random.randint(-32, 33)))
            
            result = benchmark.evaluate(child_arch)
            acc = result.test_accuracy
            cost = result.cost
            reward = self.compute_reward(acc, cost)
            
            # Replace oldest individual
            population.append((child_arch, acc, reward))
            population.pop(0)
            
            history.append(acc)
            reward_history.append(reward)
            total_cost += cost
            
            current_best = max(p[1] for p in population)
            if current_best > best_acc:
                best_acc = current_best
            
            current_best_reward = max(p[2] for p in population)
            if current_best_reward > best_reward:
                best_reward = current_best_reward
        
        conv_iter = compute_convergence_rate(history)
        
        return RunResult(
            benchmark=benchmark.benchmark_name,
            dataset=benchmark.dataset,
            algorithm=self.name,
            seed=seed,
            best_accuracy=best_acc,
            final_accuracy=history[-1],
            mean_accuracy=np.mean(history),
            best_reward=best_reward,
            mean_reward=np.mean(reward_history),
            convergence_iteration=conv_iter,
            total_time=time.time() - start_time,
            total_cost=total_cost,
            evaluations=max_evaluations,
            history=history,
            reward_history=reward_history
        )

# CMA-ES (simplified)
class CMAESNAS(BaseNASAlgorithm):
    def __init__(self, alpha: float = 0.3):
        super().__init__("CMA-ES", alpha)
    
    def search(self, benchmark, max_evaluations: int, seed: int) -> RunResult:
        np.random.seed(seed)
        random.seed(seed)
        
        history = []
        reward_history = []
        best_acc = 0.0
        best_reward = -float('inf')
        total_cost = 0.0
        start_time = time.time()
        
        # Simple evolutionary strategy
        current_solution = {
            'complexity': 1.0,
            'depth': 5,
            'width': 64
        }
        sigma = 0.3  # Mutation strength
        
        for i in range(max_evaluations):
            # Generate candidate by adding noise
            candidate = dict(current_solution)
            
            # Add Gaussian noise to parameters
            if 'complexity' in candidate:
                candidate['complexity'] = max(0.5, min(2.0, 
                    candidate['complexity'] + np.random.normal(0, sigma)))
            
            if 'depth' in candidate:
                candidate['depth'] = max(3, min(15,
                    candidate['depth'] + np.random.normal(0, sigma * 2)))
            
            if 'width' in candidate:
                candidate['width'] = max(16, min(512,
                    candidate['width'] + np.random.normal(0, sigma * 20)))
            
            # Add other architecture features
            arch = benchmark.sample_architecture()
            for key in arch:
                if key not in candidate:
                    candidate[key] = arch[key]
            
            result = benchmark.evaluate(candidate)
            acc = result.test_accuracy
            cost = result.cost
            reward = self.compute_reward(acc, cost)
            
            history.append(acc)
            reward_history.append(reward)
            total_cost += cost
            
            # Update solution if better
            if reward > best_reward:
                best_reward = reward
                current_solution = candidate
                sigma = max(0.1, sigma * 0.95)  # Reduce exploration
            
            if acc > best_acc:
                best_acc = acc
        
        conv_iter = compute_convergence_rate(history)
        
        return RunResult(
            benchmark=benchmark.benchmark_name,
            dataset=benchmark.dataset,
            algorithm=self.name,
            seed=seed,
            best_accuracy=best_acc,
            final_accuracy=history[-1],
            mean_accuracy=np.mean(history),
            best_reward=best_reward,
            mean_reward=np.mean(reward_history),
            convergence_iteration=conv_iter,
            total_time=time.time() - start_time,
            total_cost=total_cost,
            evaluations=max_evaluations,
            history=history,
            reward_history=reward_history
        )

# Bayesian Optimization (simplified)
class BayesianOptimizationNAS(BaseNASAlgorithm):
    def __init__(self, alpha: float = 0.3):
        super().__init__("BayesianOptimization", alpha)
    
    def search(self, benchmark, max_evaluations: int, seed: int) -> RunResult:
        np.random.seed(seed)
        random.seed(seed)
        
        history = []
        reward_history = []
        explored_solutions = []
        best_acc = 0.0
        best_reward = -float('inf')
        total_cost = 0.0
        start_time = time.time()
        
        # Exploration vs exploitation
        exploration_phase = min(10, max_evaluations // 5)
        
        for i in range(max_evaluations):
            if i < exploration_phase or len(explored_solutions) < 3:
                # Exploration: random sampling
                arch = benchmark.sample_architecture()
            else:
                # Exploitation: use information from previous evaluations
                # Simple heuristic: pick parameters similar to best so far
                best_solution = max(explored_solutions, key=lambda x: x[1])[0]
                arch = dict(best_solution)
                
                # Add some exploration noise
                for key in ['complexity', 'depth', 'width']:
                    if key in arch:
                        arch[key] = arch[key] * np.random.uniform(0.9, 1.1)
            
            result = benchmark.evaluate(arch)
            acc = result.test_accuracy
            cost = result.cost
            reward = self.compute_reward(acc, cost)
            
            history.append(acc)
            reward_history.append(reward)
            explored_solutions.append((arch, reward))
            total_cost += cost
            
            if acc > best_acc:
                best_acc = acc
            
            if reward > best_reward:
                best_reward = reward
        
        conv_iter = compute_convergence_rate(history)
        
        return RunResult(
            benchmark=benchmark.benchmark_name,
            dataset=benchmark.dataset,
            algorithm=self.name,
            seed=seed,
            best_accuracy=best_acc,
            final_accuracy=history[-1],
            mean_accuracy=np.mean(history),
            best_reward=best_reward,
            mean_reward=np.mean(reward_history),
            convergence_iteration=conv_iter,
            total_time=time.time() - start_time,
            total_cost=total_cost,
            evaluations=max_evaluations,
            history=history,
            reward_history=reward_history
        )

# Benchmark Orchestrator
class BenchmarkOrchestrator:
    def __init__(self):
        self.algorithms = {
            'tl_dpo': TLDPONAS(),
            'random': RandomSearchNAS(),
            'evolution': RegularizedEvolutionNAS(),
            'cma_es': CMAESNAS(),
            'bayesian': BayesianOptimizationNAS(),
        }
        logger.info(f"✓ Loaded {len(self.algorithms)} algorithms")
    
    def run_single(self, benchmark_name: str, dataset: str, algorithm: str, 
                   max_evals: int, seed: int, alpha: float = 0.3) -> RunResult:
        benchmark = SimulatedNASBenchmark(benchmark_name, dataset)
        algo = self.algorithms[algorithm]
        algo.alpha = alpha
        
        logger.info(f"  Running {algorithm} on {benchmark_name}_{dataset} (seed={seed}, α={alpha})")
        return algo.search(benchmark, max_evals, seed)
    
    def run_full_benchmark(self, benchmark_names: List[str], datasets: List[str],
                          algorithms: List[str], max_evals: int, seeds: int, alpha: float = 0.3) -> Dict:
        all_results = defaultdict(list)
        
        total_runs = len(benchmark_names) * len(datasets) * len(algorithms) * seeds
        current = 0
        
        logger.info(f"\n{'='*80}")
        logger.info(f"Starting {total_runs} runs...")
        logger.info(f"Alpha (cost penalty): {alpha}")
        logger.info(f"{'='*80}\n")
        
        for benchmark_name in benchmark_names:
            for dataset in datasets:
                for algorithm in algorithms:
                    key = f"{benchmark_name}_{dataset}_{algorithm}"
                    
                    for seed in range(seeds):
                        current += 1
                        logger.info(f"[{current}/{total_runs}] {key} seed={seed}")
                        
                        try:
                            result = self.run_single(
                                benchmark_name, dataset, algorithm, max_evals, seed, alpha
                            )
                            all_results[key].append(asdict(result))
                        except Exception as e:
                            logger.error(f"Failed: {e}")
                            all_results[key].append({
                                'error': str(e),
                                'seed': seed
                            })
        
        return dict(all_results)

# Reporting Module
class ReportGenerator:
    @staticmethod
    def aggregate_results(results: Dict) -> Dict:
        aggregated = defaultdict(lambda: defaultdict(list))
        
        for key, runs in results.items():
            parts = key.rsplit('_', 1)
            if len(parts) == 2:
                benchmark_dataset, algorithm = parts
            else:
                continue
            
            for run in runs:
                if 'error' not in run:
                    aggregated[benchmark_dataset][algorithm].append(run)
        
        return dict(aggregated)
    
    @staticmethod
    def compute_statistics(runs: List[Dict]) -> Dict:
        if not runs:
            return {}
        
        best_accs = [r['best_accuracy'] for r in runs]
        mean_accs = [r['mean_accuracy'] for r in runs]
        best_rewards = [r['best_reward'] for r in runs]
        mean_rewards = [r['mean_reward'] for r in runs]
        times = [r['total_time'] for r in runs]
        costs = [r['total_cost'] for r in runs]
        conv_iters = [r['convergence_iteration'] for r in runs]
        escalations = [r.get('escalations', 0) for r in runs]
        prunings = [r.get('prunings', 0) for r in runs]
        
        stats = {
            'best_acc_mean': np.mean(best_accs),
            'best_acc_std': np.std(best_accs),
            'mean_acc_mean': np.mean(mean_accs),
            'mean_acc_std': np.std(mean_accs),
            'best_reward_mean': np.mean(best_rewards),
            'best_reward_std': np.std(best_rewards),
            'mean_reward_mean': np.mean(mean_rewards),
            'mean_reward_std': np.std(mean_rewards),
            'time_mean': np.mean(times),
            'time_std': np.std(times),
            'cost_mean': np.mean(costs),
            'cost_std': np.std(costs),
            'conv_iter_mean': np.mean(conv_iters),
            'conv_iter_std': np.std(conv_iters),
            'escalations_mean': np.mean(escalations),
            'escalations_std': np.std(escalations),
            'prunings_mean': np.mean(prunings),
            'prunings_std': np.std(prunings),
        }
        
        # Confidence intervals
        if len(best_accs) >= 2:
            lower, upper, _ = compute_confidence_interval(np.array(best_accs))
            stats['best_acc_ci_lower'] = lower
            stats['best_acc_ci_upper'] = upper
        
        return stats
    
    @staticmethod
    def print_summary_table(results: Dict, alpha: float = 0.3):
        aggregated = ReportGenerator.aggregate_results(results)
        
        print(f"\n{'='*120}")
        print(f"NAS BENCHMARK SUMMARY (α={alpha})")
        print(f"{'='*120}\n")
        
        print(f"{'Benchmark':<25} {'Algorithm':<20} {'Best Acc':<15} {'Mean Acc':<15} {'Time (s)':<12} {'Cost':<12}")
        print(f"{'-'*120}")
        
        for bench_dataset, algorithms in sorted(aggregated.items()):
            for algorithm, runs in sorted(algorithms.items()):
                stats = ReportGenerator.compute_statistics(runs)
                
                if stats:
                    print(f"{bench_dataset:<25} {algorithm:<20} "
                          f"{stats['best_acc_mean']:.4f}±{stats['best_acc_std']:.4f}  "
                          f"{stats['mean_acc_mean']:.4f}±{stats['mean_acc_std']:.4f}  "
                          f"{stats['time_mean']:.2f}±{stats['time_std']:.2f}  "
                          f"{stats['cost_mean']:.2f}±{stats['cost_std']:.2f}")
        
        print(f"{'='*120}\n")
    
    @staticmethod
    def print_benchmark_report(results: Dict, benchmark_name: str, alpha: float):
        """Print comprehensive benchmark report."""
        aggregated = ReportGenerator.aggregate_results(results)
        
        # Filter for current benchmark
        bench_results = {k: v for k, v in aggregated.items() if benchmark_name in k}
        
        if not bench_results:
            logger.warning(f"No results for benchmark: {benchmark_name}")
            return
        
        method_stats = {}
        for bench_key, algorithms in bench_results.items():
            for algo, runs in algorithms.items():
                stats = ReportGenerator.compute_statistics(runs)
                method_stats[algo] = stats
        
        print(f"\n{'='*160}")
        print(f"NAS BENCHMARK REPORT: {benchmark_name.upper()}")
        print(f"Alpha (cost penalty): {alpha}")
        print(f"{'='*160}\n")
        
        # 1. PRIMARY PERFORMANCE METRICS
        print("1. PRIMARY PERFORMANCE METRICS (Mean ± Std)")
        print(f"{'='*160}\n")
        print(f"{'Method':<20} | {'Accuracy':<18} | {'Best Acc':<18} | "
              f"{'Mean Reward':<18} | {'Total Cost':<18} | {'Time':<12}")
        print(f"{'-'*160}")
        
        for algo, stats in method_stats.items():
            if stats:
                print(f"{algo:<20} | {stats['mean_acc_mean']:.4f}±{stats['mean_acc_std']:.4f}      | "
                      f"{stats['best_acc_mean']:.4f}±{stats['best_acc_std']:.4f}      | "
                      f"{stats['mean_reward_mean']:.4f}±{stats['mean_reward_std']:.4f}      | "
                      f"{stats['cost_mean']:.4f}±{stats['cost_std']:.4f}      | "
                      f"{stats['time_mean']:.2f}±{stats['time_std']:.2f}")
        
        # 2. EFFICIENCY & ADAPTIVITY
        print(f"\n{'='*160}")
        print("2. EFFICIENCY & ADAPTIVITY")
        print(f"{'='*160}\n")
        print(f"{'Method':<20} | {'Conv. Iter':<15} | {'Escalations':<15} | {'Prunings':<15}")
        print(f"{'-'*80}")
        
        for algo, stats in method_stats.items():
            if stats:
                print(f"{algo:<20} | {stats['conv_iter_mean']:<15.1f} | "
                      f"{stats['escalations_mean']:<15.1f} | {stats['prunings_mean']:<15.1f}")
        
        # 3. PARETO FRONTIER
        print(f"\n{'='*160}")
        print("3. PARETO FRONTIER ANALYSIS (Cost vs Accuracy)")
        print(f"{'='*160}\n")
        
        data_points = [(stats['cost_mean'], stats['best_acc_mean'], algo) 
                      for algo, stats in method_stats.items() if stats]
        frontier = compute_pareto_frontier(data_points)
        
        print(f"{'Method':<20} | {'Total Cost':<15} | {'Best Acc':<15} | {'Pareto-Optimal':<20}")
        print(f"{'-'*75}")
        
        for cost_val, acc_val, algo in sorted(data_points):
            is_pareto = any(p[2] == algo for p in frontier)
            status = "✓ Yes" if is_pareto else "  No"
            print(f"{algo:<20} | {cost_val:<15.4f} | {acc_val:<15.4f} | {status:<20}")
        
        # 4. STATISTICAL SIGNIFICANCE
        print(f"\n{'='*160}")
        print("4. STATISTICAL SIGNIFICANCE (T-Test vs TL-DPO)")
        print(f"{'='*160}\n")
        
        if 'tl_dpo' in method_stats:
            print(f"{'Method':<20} | {'t-stat':<12} | {'p-value':<12} | {'Sig.':<10}")
            print(f"{'-'*60}")
            
            for algo, stats in method_stats.items():
                if algo != 'tl_dpo' and algo in aggregated.get(list(bench_results.keys())[0], {}):
                    tl_dpo_runs = aggregated[list(bench_results.keys())[0]]['tl_dpo']
                    algo_runs = aggregated[list(bench_results.keys())[0]][algo]
                    
                    if tl_dpo_runs and algo_runs:
                        tl_accs = np.array([r['best_accuracy'] for r in tl_dpo_runs])
                        algo_accs = np.array([r['best_accuracy'] for r in algo_runs])
                        
                        t_stat, p_val = ttest_ind(tl_accs, algo_accs)
                        cohen_d = compute_cohens_d(tl_accs, algo_accs)
                        
                        # Interpret effect size
                        if abs(cohen_d) < 0.2:
                            effect = "Negligible"
                        elif abs(cohen_d) < 0.5:
                            effect = "Small"
                        elif abs(cohen_d) < 0.8:
                            effect = "Medium"
                        else:
                            effect = "Large"
                        
                        sig = "Yes*" if p_val < 0.05 else "No"
                        print(f"{algo:<20} | {t_stat:<12.4f} | {p_val:<12.6f} | {sig:<10}")
                        print(f"  Effect size: d={cohen_d:.3f} ({effect})")
        
        print(f"\n{'='*160}\n")
    
    @staticmethod
    def plot_convergence_curves(results: Dict, save_path: str = 'nas_convergence.png'):
        """Plot convergence curves."""
        aggregated = ReportGenerator.aggregate_results(results)
        
        bench_with_history = {
            b: algs for b, algs in aggregated.items()
            if any(any('history' in r and r['history'] for r in runs) for runs in algs.values())
        }
        
        n_benchmarks = len(bench_with_history)
        if n_benchmarks == 0:
            logger.warning("No convergence data available. Skipping convergence plot.")
            return
        
        ncols = min(n_benchmarks, 3)
        fig, axes = plt.subplots(1, ncols, figsize=(6 * ncols, 5))
        if ncols == 1:
            axes = [axes]
        
        for idx, (bench_dataset, algorithms) in enumerate(sorted(bench_with_history.items())[:ncols]):
            ax = axes[idx]
            
            for algorithm, runs in sorted(algorithms.items()):
                histories = [r['history'] for r in runs if 'history' in r and r['history']]
                if not histories:
                    continue
                
                max_len = max(len(h) for h in histories)
                histories_padded = []
                for h in histories:
                    if len(h) == 0:
                        continue
                    padded = h + [h[-1]] * (max_len - len(h))
                    histories_padded.append(padded)
                
                if not histories_padded:
                    continue
                
                avg_history = np.mean(histories_padded, axis=0)
                ax.plot(avg_history, label=algorithm, linewidth=2)
            
            ax.set_xlabel('Evaluations', fontsize=12)
            ax.set_ylabel('Best Accuracy', fontsize=12)
            ax.set_title(bench_dataset.replace('_', ' ').title(), fontsize=14, fontweight='bold')
            ax.legend(fontsize=9)
            ax.grid(True, alpha=0.3)
        
        plt.tight_layout()
        plt.savefig(save_path, dpi=300, bbox_inches='tight')
        logger.info(f"✓ Convergence plot saved: {save_path}")
        plt.close()
    
    @staticmethod
    def plot_performance_comparison(results: Dict, save_path: str = 'nas_performance.png'):
        """Plot performance comparison."""
        aggregated = ReportGenerator.aggregate_results(results)
        
        all_data = defaultdict(lambda: defaultdict(list))
        for bench_dataset, algorithms in aggregated.items():
            for algorithm, runs in algorithms.items():
                stats = ReportGenerator.compute_statistics(runs)
                if stats:
                    all_data[algorithm]['best_acc'].append(stats['best_acc_mean'])
                    all_data[algorithm]['time'].append(stats['time_mean'])
                    all_data[algorithm]['conv_iter'].append(stats['conv_iter_mean'])
                    all_data[algorithm]['cost'].append(stats['cost_mean'])
        
        if not all_data:
            logger.warning("No performance data available. Skipping performance plots.")
            return
        
        fig, axes = plt.subplots(2, 2, figsize=(16, 12))
        
        # Plot 1: Best Accuracy
        ax = axes[0, 0]
        algorithms = sorted(all_data.keys())
        means = [np.mean(all_data[a]['best_acc']) for a in algorithms]
        ax.bar(algorithms, means, color='steelblue')
        ax.set_ylabel('Mean Best Accuracy', fontsize=12)
        ax.set_title('Average Performance Across Benchmarks', fontsize=14, fontweight='bold')
        ax.tick_params(axis='x', rotation=45)
        ax.grid(True, axis='y', alpha=0.3)
        
        # Plot 2: Time
        ax = axes[0, 1]
        times = [np.mean(all_data[a]['time']) for a in algorithms]
        ax.bar(algorithms, times, color='coral')
        ax.set_ylabel('Mean Time (s)', fontsize=12)
        ax.set_title('Computational Cost', fontsize=14, fontweight='bold')
        ax.tick_params(axis='x', rotation=45)
        ax.grid(True, axis='y', alpha=0.3)
        
        # Plot 3: Pareto (Cost vs Accuracy)
        ax = axes[1, 0]
        for algo in algorithms:
            avg_acc = np.mean(all_data[algo]['best_acc'])
            avg_cost = np.mean(all_data[algo]['cost'])
            ax.scatter(avg_cost, avg_acc, s=100, label=algo, alpha=0.7)
        ax.set_xlabel('Cost', fontsize=12)
        ax.set_ylabel('Best Accuracy', fontsize=12)
        ax.set_title('Pareto: Cost vs Accuracy', fontsize=14, fontweight='bold')
        ax.legend(fontsize=9)
        ax.grid(True, alpha=0.3)
        
        # Plot 4: Convergence Speed
        ax = axes[1, 1]
        conv_iters = [np.mean(all_data[a]['conv_iter']) for a in algorithms]
        ax.bar(algorithms, conv_iters, color='mediumseagreen')
        ax.set_ylabel('Mean Convergence Iteration', fontsize=12)
        ax.set_title('Convergence Speed', fontsize=14, fontweight='bold')
        ax.tick_params(axis='x', rotation=45)
        ax.grid(True, axis='y', alpha=0.3)
        
        plt.tight_layout()
        plt.savefig(save_path, dpi=300, bbox_inches='tight')
        logger.info(f"✓ Performance plot saved: {save_path}")
        plt.close()
    
    @staticmethod
    def save_results(results: Dict, filename: str = 'nas_benchmark_results.json'):
        filename = str(filename)
        with open(filename, 'w') as f:
            json.dump(results, f, indent=2)
        logger.info(f"✓ Results saved: {filename}")
    
    @staticmethod
    def save_csv(results: Dict, filename: str = 'nas_benchmark_results.csv'):
        rows = []
        for key, runs in results.items():
            for run in runs:
                if 'error' not in run:
                    row = {
                        'key': key,
                        **{k: v for k, v in run.items() if k != 'history' and k != 'reward_history'}
                    }
                    rows.append(row)
        
        if rows:
            with open(filename, 'w', newline='') as f:
                writer = csv.DictWriter(f, fieldnames=rows[0].keys())
                writer.writeheader()
                writer.writerows(rows)
            logger.info(f"✓ CSV saved: {filename}")

# Main function
def main():
    parser = argparse.ArgumentParser(
        description='TL-DPO NAS Benchmark Suite (Simulated)',
        formatter_class=argparse.RawDescriptionHelpFormatter
    )
    
    parser.add_argument('--benchmarks', nargs='+', 
                       default=['nasbench201'],
                       choices=['nasbench101', 'nasbench201', 'nasbench301', 'nats'],
                       help='Benchmarks to run')
    
    parser.add_argument('--datasets', nargs='+',
                       default=['cifar10'],
                       choices=['cifar10', 'cifar100', 'imagenet16-120'],
                       help='Datasets to use')
    
    parser.add_argument('--algorithms', nargs='+',
                       default=['tl_dpo', 'random', 'evolution'],
                       choices=['tl_dpo', 'random', 'evolution', 'cma_es', 'bayesian'],
                       help='Algorithms to compare')
    
    parser.add_argument('--max-evals', type=int, default=100,
                       help='Maximum evaluations per run')
    
    parser.add_argument('--seeds', type=int, default=5,
                       help='Number of random seeds')
    
    parser.add_argument('--alpha', type=float, default=0.3,
                       help='Cost penalty weight')
    
    parser.add_argument('--output-dir', type=str, default='./nas_results',
                       help='Output directory')
    
    args = parser.parse_args()
    
    # Create output directory
    output_dir = pathlib.Path(args.output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)
    
    print(f"\n{'='*80}")
    print("TL-DPO NAS BENCHMARK SUITE (SIMULATED)")
    print(f"{'='*80}")
    print(f"Benchmarks: {', '.join(args.benchmarks)}")
    print(f"Datasets: {', '.join(args.datasets)}")
    print(f"Algorithms: {', '.join(args.algorithms)}")
    print(f"Max Evaluations: {args.max_evals}")
    print(f"Seeds: {args.seeds}")
    print(f"Alpha (cost penalty): {args.alpha}")
    print(f"{'='*80}\n")
    
    # Run benchmarks
    orchestrator = BenchmarkOrchestrator()
    results = orchestrator.run_full_benchmark(
        args.benchmarks,
        args.datasets,
        args.algorithms,
        args.max_evals,
        args.seeds,
        args.alpha
    )
    
    # Generate reports
    print(f"\n{'='*80}")
    print("GENERATING REPORTS")
    print(f"{'='*80}\n")
    
    ReportGenerator.print_summary_table(results, args.alpha)
    
    for benchmark in args.benchmarks:
        for dataset in args.datasets:
            ReportGenerator.print_benchmark_report(
                results, 
                f"{benchmark}_{dataset}", 
                args.alpha
            )
    
    # Save outputs
    ReportGenerator.save_results(results, output_dir / 'nas_results.json')
    ReportGenerator.save_csv(results, output_dir / 'nas_results.csv')
    ReportGenerator.plot_convergence_curves(results, output_dir / 'nas_convergence.png')
    ReportGenerator.plot_performance_comparison(results, output_dir / 'nas_performance.png')
    
    # Print final summary
    print(f"\n{'='*80}")
    print("✅ NAS BENCHMARK COMPLETE")
    print(f"{'='*80}")
    print(f"Results saved to: {output_dir}")
    print(f"  - nas_results.json")
    print(f"  - nas_results.csv")
    print(f"  - nas_convergence.png")
    print(f"  - nas_performance.png")
    print(f"{'='*80}\n")
    
    # Print which algorithm works best on which benchmark
    print("\n📊 FINAL SUMMARY: WHICH ALGORITHM WORKS BEST WHERE")
    print("="*80)
    
    aggregated = ReportGenerator.aggregate_results(results)
    for bench_dataset in sorted(aggregated.keys()):
        best_algo = None
        best_acc = -1
        
        for algorithm, runs in aggregated[bench_dataset].items():
            stats = ReportGenerator.compute_statistics(runs)
            if stats and stats['best_acc_mean'] > best_acc:
                best_acc = stats['best_acc_mean']
                best_algo = algorithm
        
        if best_algo:
            print(f"{bench_dataset:<40} → {best_algo:<20} (Accuracy: {best_acc:.4f})")

if __name__ == '__main__':
    main()