#!/usr/bin/env python
# -*- coding:utf-8 -*-
# @FileName  :material.py
# @Time      :2025/4/29 11:07
# @Author    :David
# @Company: shenyang JY
import pandas as pd
from pathlib import Path
from concurrent.futures import ThreadPoolExecutor
from functools import partial
from app.model.config import myargparse

args = myargparse(discription="算法批量修模", add_help=False)

class MaterialLoader:
    def __init__(self, base_path, lazy_load=True):
        self.base_path = Path(base_path)
        self.lazy_load = lazy_load
        self._data_cache = {}
        self.opt = args.parse_args_and_yaml()

    def _load_material(self, station_path):
        """核心数据加载方法"""
        input_file = self.base_path / station_path
        # 根据您的原始代码逻辑简化的加载流程
        try:
            basic = pd.read_csv(input_file / self.opt.doc_mapping['basic'], sep=r'\s+', header=0)
            power = pd.read_csv(input_file / self.opt.doc_mapping['power'], sep=r'\s+', header=0)
            plant_type = int(basic.loc[basic['PropertyID'].tolist().index('PlantType'), 'Value'])
            assert plant_type == 0 or plant_type == 1
            # 根据电站类型加载数据
            nwp_v = pd.read_csv(input_file / '0' / self.opt.doc_mapping['nwp_v'], sep=r'\s+', header=0)
            nwp_v_h = pd.read_csv(input_file / '0' / self.opt.doc_mapping['nwp_v_h'], sep=r'\s+', header=0)
            nwp_own = pd.read_csv(input_file / '1' / self.opt.doc_mapping['nwp_own'], sep=r'\s+', header=0)
            nwp_own_h = pd.read_csv(input_file / '1' / self.opt.doc_mapping['nwp_own_h'], sep=r'\s+', header=0)
            if self.opt.switch_nwp_owner:
                nwp_v, nwp_v_h = nwp_own, nwp_own_h
            # 如果是风电
            if plant_type == 0:
                station_info = pd.read_csv(input_file / self.opt.doc_mapping['station_info_w'], sep=r'\s+', header=0)
                station_info_d = pd.read_csv(input_file / self.opt.doc_mapping['station_info_d_w'], sep=r'\s+', header=0)
                nwp = pd.read_csv(input_file / self.opt.doc_mapping['nwp_w'], sep=r'\s+', header=0)
                nwp_h = pd.read_csv(input_file / self.opt.doc_mapping['nwp_w_h'], sep=r'\s+', header=0)
                cap = float(station_info.loc[0, 'PlantCap'])
                if (input_file / self.opt.doc_mapping['env_wf']).exists():
                    env = pd.read_csv(input_file / self.opt.doc_mapping['env_wf'], sep=r'\s+', header=0)
                else:
                    env = None
            # 如果是光伏
            else:
                station_info = pd.read_csv(input_file / self.opt.doc_mapping['station_info_s'], sep=r'\s+', header=0)
                station_info_d = pd.read_csv(input_file / self.opt.doc_mapping['station_info_d_s'], sep=r'\s+', header=0)
                nwp = pd.read_csv(input_file / self.opt.doc_mapping['nwp_s'], sep=r'\s+', header=0)
                nwp_h = pd.read_csv(input_file / self.opt.doc_mapping['nwp_s_h'], sep=r'\s+', header=0)
                cap = float(station_info.loc[0, 'PlantCap'])
                if (input_file / self.opt.doc_mapping['env_sf']).exists():
                    env = pd.read_csv(input_file / self.opt.doc_mapping['env_sf'], sep=r'\s+', header=0)
                else:
                    env = None

            return {
                'station_info': station_info,
                'station_info_d': station_info_d,
                'nwp': nwp,
                'nwp_h': nwp_h,
                'power': power,
                'nwp_v': nwp_v,
                'nwp_v_h': nwp_v_h,
                'env': env,
                'cap': cap
            }
        except Exception as e:
            print(f"Error loading {station_path}: {str(e)}")
            return None

    def get_material(self, station_path):
        if self.lazy_load:
            if station_path not in self._data_cache:
                self._data_cache[station_path] = self._load_material(station_path)
            return self._data_cache[station_path]
        else:
            return self._load_material(station_path)



if __name__ == "__main__":
    run_code = 0