#!/usr/bin/env python
# -*- coding:utf-8 -*-
# @FileName  :task_worker.py
# @Time      :2025/4/29 11:05
# @Author    :David
# @Company: shenyang JY

import logging
import pandas as pd
from app.model.tf_model_train import ModelTrainer
from app.model.material import MaterialLoader


def station_task(station_id, config):
    """单个场站训练任务"""
    try:
        mate = MaterialLoader(base_path=config['input_file'])
        # 动态生成场站数据路径
        input_file = config['data_template'].format(station_id=station_id)

        # 加载数据
        data_objects = mate.get_material(input_file)

        # 数据合并
        train_data = pd.merge(data_objects.nwp_v_h, data_objects.power, on=config['time_column'])

        # 模型训练
        model = ModelTrainer(input_file, train_data, capacity=config['capacities'][station_id], gpu_id=config.get('gpu_assignment'))
        model.train()

        return {'status': 'success', 'station_id': station_id}
    except Exception as e:
        logging.error(f"Station {station_id} failed: {str(e)}")
        return {'status': 'failed', 'station_id': station_id}


def area_task():
    """区域级训练任务"""
    pass