import os

import matplotlib.pyplot as plt
import pandas as pd
import seaborn as sns
from algorithmContract.confBusiness import *
from behavior.analyst import Analyst
from utils.directoryUtil import DirectoryUtil as dir


class PowerOscillationAnalyst(Analyst):
    """
    风电机组功率震荡分析
    """

    def typeAnalyst(self):
        return "power_diff"

    def turbineAnalysis(self,
                 dataFrame,
                 outputAnalysisDir,
                 outputFilePath,
                 confData: ConfBusiness,
                 turbineName):
        self.power_diff(dataFrame, outputFilePath,
                        confData.field_power,confData.field_rotor_speed)

    def power_diff(self, dataFrame, outputFilePath,  field_Active_Power, field_Rotor_Speed):
        # Floor the power column to the nearest 10
        dataFrame['power_col_floor'] = (
            dataFrame[field_Active_Power] // 10 * 10).astype('int32')

        # Group by the floored power column
        grouped = dataFrame.groupby('power_col_floor')

        # Calculate max, min, and diff of generator speed within each group
        agg_df = grouped[field_Rotor_Speed].agg(
            speed_max='max',
            speed_min='min'
        )
        agg_df['speed_diff'] = agg_df['speed_max'] - agg_df['speed_min']

        # Sort by the floored power column
        agg_df = agg_df.sort_index()

        # Write the result to a CSV file
        agg_df.to_csv(outputFilePath)

    def turbinesAnalysis(self,dataFrameMerge,outputAnalysisDir, confData: ConfBusiness):
        self.plot_power_oscillations(outputAnalysisDir,confData.farm_name)

    def plot_power_oscillations(self, csvFileDirOfCp, farm_name, encoding='utf-8'):
        """
        Plot TSR trend from CSV files in a given input path and save the plots to an output path.

        Parameters:
        - csvFileDirOfCp: str, path to the directory containing input CSV files.
        - farm_name: str, name of the wind farm.
        - encoding: str, encoding of the input CSV files. Defaults to 'utf-8'.
        """
        sns.set_palette('deep')
        field_Name_Turbine = "turbine_name"
        x_name = 'power_col_floor'
        y_name = 'speed_diff'

        # 初始化结果DataFrame
        res = pd.DataFrame()

        # 遍历输入路径下的所有文件
        for root, dir_names, file_names in dir.list_directory(csvFileDirOfCp):
            for file_name in file_names:

                if not file_name.endswith(CSVSuffix):
                    continue

                file_path = os.path.join(root, file_name)
                frame = pd.read_csv(file_path, encoding=encoding)

                # 获取输出文件名前缀
                turbine_name = file_name.split(CSVSuffix)[0]
                # 添加设备名作为新列
                frame[field_Name_Turbine] = turbine_name

                selected_data = frame.loc[:, [field_Name_Turbine, x_name, y_name]]
                res = pd.concat([res, selected_data], axis=0)

        # 重置索引
        ress = res.reset_index(drop=True)

        # 绘制所有设备的功率震荡图
        fig, ax = plt.subplots(figsize=(16, 8))
        ax = sns.lineplot(x=x_name, y=y_name, data=ress, hue=field_Name_Turbine)
        ax.set_title(f'功率震荡-{self.turbineModelInfo[Field_MachineTypeCode].iloc[0]}')
        plt.legend(ncol=4)
        plt.savefig(csvFileDirOfCp+ r'/{}-Power-Oscillation.png'.format(farm_name), bbox_inches='tight', dpi=300)
        plt.close()

        # 分组绘制每个设备的功率震荡图
        grouped = ress.groupby(field_Name_Turbine)
        for name, group in grouped:
            color = ["lightgrey"] * len(ress[field_Name_Turbine].unique())
            fig, ax = plt.subplots(figsize=(8, 8))
            ax = sns.lineplot(x=x_name, y=y_name, data=ress, hue=field_Name_Turbine, palette=sns.color_palette(color), legend=False)
            ax = sns.lineplot(x=x_name, y=y_name, data=group, color='darkblue', legend=False)
            ax.set_title('turbine_name={}'.format(name))
            plt.savefig(csvFileDirOfCp+ r'/{}.png'.format(name), bbox_inches='tight', dpi=120)
            plt.close()