from autogluon.common.utils.resource_utils import ResourceManager

# 获取系统资源信息
cpu_count = ResourceManager.get_cpu_count()
gpu_count = ResourceManager.get_gpu_count()
total_memory = ResourceManager.get_memory_size("GB")
available_memory = ResourceManager.get_available_virtual_mem("GB")

print(f"CPU核心数: {cpu_count}")
print(f"GPU数量: {gpu_count}") 
print(f"总内存: {total_memory:.1f} GB")
print(f"可用内存: {available_memory:.1f} GB")

from autogluon.tabular import TabularPredictor

predictor = TabularPredictor(label="target").fit(
    train_data,
    num_cpus=4,           # 指定CPU核心数
    num_gpus=1,           # 指定GPU数量
    memory_limit="16GB",  # 内存限制
    time_limit=3600       # 时间限制（秒）
)

def predict_memory_requirements(n_samples, n_features, task_type):
    """
    预测AutoGluon训练所需内存
    
    参数:
    n_samples: 样本数量
    n_features: 特征数量  
    task_type: 任务类型（'tabular', 'multimodal', 'timeseries'）
    
    返回:
    预估内存需求（GB）
    """
    base_memory = 2.0  # 基础内存开销
    
    if task_type == 'tabular':
        memory_per_sample = 0.001 * n_features
        return base_memory + n_samples * memory_per_sample / 1024
    
    elif task_type == 'multimodal':
        memory_per_sample = 0.01 * n_features  # 多媒体数据内存需求更高
        return base_memory + n_samples * memory_per_sample / 1024
    
    elif task_type == 'timeseries':
        memory_per_sample = 0.0005 * n_features
        return base_memory + n_samples * memory_per_sample / 1024
    
    return base_memory

def predict_cpu_requirements(n_samples, n_features):
    """
    预测所需的CPU核心数
    """
    if n_samples < 10000:
        return min(4, ResourceManager.get_cpu_count())
    elif n_samples < 100000:
        return min(8, ResourceManager.get_cpu_count())
    else:
        # 大数据集使用更多核心，但不超过物理核心数
        physical_cores = ResourceManager.get_cpu_count(only_physical_cores=True)
        return min(16, physical_cores)

# SageMaker训练实例选择指南
instance_mapping = {
    'small': {'cpu': 4, 'memory': 16, 'gpu': 0},
    'medium': {'cpu': 8, 'memory': 32, 'gpu': 0},
    'large': {'cpu': 16, 'memory': 64, 'gpu': 1},
    'xlarge': {'cpu': 32, 'memory': 128, 'gpu': 2},
    '2xlarge': {'cpu': 64, 'memory': 256, 'gpu': 4}
}

def recommend_instance_type(n_samples, n_features, task_type):
    memory_needed = predict_memory_requirements(n_samples, n_features, task_type)
    
    for size, specs in instance_mapping.items():
        if (specs['memory'] >= memory_needed and 
            specs['cpu'] >= predict_cpu_requirements(n_samples, n_features)):
            return size, specs
    
    return 'custom', {'memory': max(512, memory_needed), 'cpu': 64, 'gpu': 8}

class DynamicResourceManager:
    def __init__(self, initial_config):
        self.current_config = initial_config
        self.performance_metrics = []
    
    def adjust_resources(self, current_usage, training_progress):
        """
        根据实时使用情况动态调整资源
        """
        cpu_usage = current_usage['cpu']
        memory_usage = current_usage['memory']
        gpu_usage = current_usage['gpu']
        
        # 内存使用超过90%时增加内存
        if memory_usage > 0.9:
            new_memory = self.current_config['memory'] * 1.2
            logger.info(f"内存使用率高({memory_usage:.1%})，增加内存到 {new_memory:.1f}GB")
            self.current_config['memory'] = new_memory
        
        # CPU使用率低时减少核心数
        if cpu_usage < 0.3 and training_progress > 0.5:
            new_cpus = max(2, int(self.current_config['cpu'] * 0.8))
            logger.info(f"CPU使用率低({cpu_usage:.1%})，减少核心到 {new_cpus}")
            self.current_config['cpu'] = new_cpus
        
        return self.current_config

def cost_optimized_training(train_data, label_column, budget_per_hour):
    """
    成本优化的训练策略
    """
    # 第一阶段：小规模试训练
    sample_data = train_data.sample(frac=0.1)
    predictor = TabularPredictor(label=label_column).fit(
        sample_data,
        num_cpus=4,
        num_gpus=0,  # 第一阶段不使用GPU
        time_limit=1800,
        presets='medium_quality'
    )
    
    # 分析资源需求
    resource_needs = analyze_resource_needs(predictor, train_data)
    
    # 选择成本最优的实例类型
    optimal_instance = select_cost_effective_instance(
        resource_needs, 
        budget_per_hour
    )
    
    # 全量训练
    final_predictor = TabularPredictor(label=label_column).fit(
        train_data,
        num_cpus=optimal_instance['cpu'],
        num_gpus=optimal_instance['gpu'],
        memory_limit=f"{optimal_instance['memory']}GB",
        time_limit=optimal_instance['time_limit']
    )
    
    return final_predictor

class ResourceMonitor:
    def __init__(self, predictor, check_interval=60):
        self.predictor = predictor
        self.check_interval = check_interval
        self.usage_history = []
    
    def start_monitoring(self):
        while training_in_progress():
            current_usage = self.get_current_usage()
            self.usage_history.append({
                'timestamp': time.time(),
                'cpu_usage': current_usage['cpu'],
                'memory_usage': current_usage['memory'],
                'gpu_usage': current_usage['gpu']
            })
            
            if self.check_resource_constraints(current_usage):
                self.trigger_alert(current_usage)
            
            time.sleep(self.check_interval)
    
    def check_resource_constraints(self, usage):
        # 检查是否接近资源限制
        if usage['memory'] > 0.85:  # 内存使用超过85%
            return True
        if usage['cpu'] > 0.9:     # CPU使用超过90%
            return True
        return False
    
    def generate_resource_report(self):
        """生成资源使用报告"""
        report = {
            'peak_memory': max([u['memory'] for u in self.usage_history]),
            'avg_cpu_usage': np.mean([u['cpu'] for u in self.usage_history]),
            'total_training_time': self.usage_history[-1]['timestamp'] - self.usage_history[0]['timestamp']
        }
        return report