#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Multi-Category Convolution-based NAS Configuration
"""

import os

# LLM Configuration
LLM_CONFIG = {
    'api_key': os.getenv('OPENAI_API_KEY', ''),
    'model': 'gpt-4.1',
    'temperature': 0.7,
    'max_tokens': 2000,
    'max_retries': 3,
    'concurrent_requests': True  # Enable parallel requests
}

# Convolution-based Classification System
SEARCH_CONFIG = {
    'max_iterations': 20,
    'conv_categories': {
        'no_convolution': {
            'description': 'No convolution operations',
            'conv_3x3_count': 0,
            'conv_1x1_count': 0,
            'target_count': 10,
            'prompt_file': 'no_convolution.txt'
        },
        'pure_lightweight': {
            'description': 'Only 1x1 convolutions, no 3x3',
            'conv_3x3_count': 0,
            'conv_1x1_count': (1, 6),  # Range: 1-6
            'target_count': 10, 
            'prompt_file': 'pure_lightweight.txt'
        },
        'basic_mixed': {
            'description': 'Exactly 1 3x3 convolution',
            'conv_3x3_count': 1,
            'conv_1x1_count': (0, 5),  # Range: 0-5
            'target_count': 10, 
            'prompt_file': 'basic_mixed.txt'
        },
        'medium_mixed': {
            'description': 'Exactly 2 3x3 convolutions',
            'conv_3x3_count': 2,
            'conv_1x1_count': (0, 4),  # Range: 0-4
            'target_count': 10,
            'prompt_file': 'medium_mixed.txt'
        },
        'heavy_mixed': {
            'description': 'Exactly 3 3x3 convolutions',
            'conv_3x3_count': 3,
            'conv_1x1_count': (0, 3),  # Range: 0-3
            'target_count': 10,
            'prompt_file': 'heavy_mixed.txt'
        },
        'extreme_compute': {
            'description': '4+ 3x3 convolutions for maximum performance',
            'conv_3x3_count': (4, 6),  # Range: 4-6
            'conv_1x1_count': (0, 2),  # Range: 0-2
            'target_count': 10,
            'prompt_file': 'extreme_compute.txt'
        }
    },
    'total_architectures_per_iteration': 60,
    
    # 反馈循环配置
    'feedback_config': {
        'reference_architectures_count': 10,  # 每次反馈提供的参考架构数量
        'generation_strategy': {
            'mutation_count': 5,     # 基于现有架构变异的数量
            'crossover_count': 5,    # 基于两个架构交叉的数量
            'pure_generation_count': 0  # 纯生成的数量
        }
    }
}

# Prompt Configuration
PROMPT_CONFIG = {
    'prompt_dir': 'pel_nas/llm/prompts',
    'system_prompt': """You are an expert in Neural Architecture Search (NAS) for {hardware_device} hardware.
Generate architectures optimized for accuracy and latency on this specific device.

Operations: none, avg_pool_3x3, nor_conv_1x1, nor_conv_3x3, skip_connect
Format: |op0~0|+|op1~0|op2~1|+|op3~0|op4~1|op5~2|

Return ONLY architecture strings, one per line."""
}

# Data & Predictor Configuration
DATA_CONFIG = {
    'metrics_csv': 'nb201_hw_metrics.csv',
    'zero_cost_features': 'zc_predictor/data/nb201_zc_scores.json',
}

EVAL_CONFIG = {
    'primary_metrics': ['accuracy', 'latency'],
    'pareto_objectives': ['accuracy', 'latency'],
    'default_latency_device': 'edgegpu',
    'fallback_accuracy': 0.0,
}

PREDICTOR_CONFIG = {
    'enabled': True,
    'model_paths': {
        'cifar10': 'zc_predictor/models/cifar10.json',
        'cifar100': 'zc_predictor/models/cifar100.json',
        'ImageNet16-120': 'zc_predictor/models/ImageNet16-120.json',
    },
    'feature_order': None,
} 

# Backwards compatibility for legacy scripts referencing DATA_PATHS
DATA_PATHS = {
    'nas_bench_201': DATA_CONFIG.get('metrics_csv'),
    'hw_nas_bench': None,
}
