import os
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import pickle
current_path = os.path.dirname(__file__)


class LimitPower(object):
    def __init__(self, logger, args, weather_power):
        self.logger = logger
        self.args = args
        self.opt = self.args.parse_args_and_yaml()
        self.weather_power = weather_power
        self.step = self.opt.usable_power['step']
        self.segs = np.array([x * self.step for x in range(1, 50)])  # 对风速以50为间隔进行分段
        self.xs = np.array([self.segs[i - 1] + x if i > 0 else self.step / 2 for i, x in
              enumerate([self.step / 2 for _ in self.segs])])  # 分段的中间点

    def segment_statis(self):
        """
        对机头风速-实际功率进行分段处理，获取分度的中位点，四分位间距和斜率
        :return: glob_rp 总辐射分段
        """
        glob_rp = {}       # dict: key 辐照度分段中间点 value 分段内的实际功率
        for index, row in self.weather_power.iterrows():
            ws_ = row[self.opt.usable_power["env"]]
            rp = row['C_REAL_VALUE']
            for i, seg in enumerate(self.segs):
                if ws_ <= seg:
                    glob_rp.setdefault(self.xs[i], []).append(rp)
                    break
        for i, x in enumerate(self.xs):
            rps = glob_rp.get(x)
            if rps is None:
                continue
            mean = np.around(np.mean(rps), 3)     # 实际功率均值
            std = np.around(np.std(rps), 3)  # 实际功率标准差

            glob_rp[x] = [mean, std]   # 更新dict
        return glob_rp

    def saveVar(self, path, data):
        os.makedirs(os.path.dirname(path), exist_ok=True)
        with open(path, 'wb') as file:
            pickle.dump(data, file)

    def filter_unlimited_power(self, ws, real_power, glob_rp):
        """
        预测可用功主方法
        :param zfs: 要预测可用功率的总辐射
        :param k: 斜率
        :param b: 偏移量
        :return: 预测的可用功率
        """
        coe = self.opt.usable_power['outliers_threshold']
        seg = self.xs[np.argmax(self.segs >= ws)]
        if seg in glob_rp:
            mean, std = glob_rp[seg][0], glob_rp[seg][1]
            high = mean + std*coe if mean + std*coe < self.opt.cap else self.opt.cap
            low = mean - std*coe if mean - std*coe > 0 else 0
            if low <= real_power <= high:
                return True
            else:
                return False
        else:
            return True

    def clean_limited_power(self, name, cluster=False):
        glob_rp = self.segment_statis()
        if cluster is True:
            self.saveVar(os.path.dirname(current_path) + '/var/glob_rp.pickle', glob_rp)
        new_weather_power, number = [], 0
        # fig, ax = plt.subplots()
        for index, row in self.weather_power.iterrows():
            zfs = row[self.opt.usable_power["env"]]
            rp = row['C_REAL_VALUE']
            if zfs < 0 or rp < 0:
                continue
            if self.filter_unlimited_power(zfs, rp, glob_rp):
                row['c'] = 'red'
                new_weather_power.append(row)
            else:
                row['c'] = 'blue'
                new_weather_power.append(row)
        new_weather_power = pd.concat(new_weather_power, axis=1).T
        new_weather_power.plot.scatter(x=self.opt.usable_power["env"], y='C_REAL_VALUE', c='c')
        plt.savefig(current_path + '/figs/测风法{}.png'.format(name))
        new_weather_power = new_weather_power[new_weather_power['c'] == 'red']
        number = len(new_weather_power)
        self.logger.info("未清洗限电前，总共有：{}条数据".format(len(self.weather_power)))
        self.logger.info("清除限电后保留的点有：" + str(number) + " 占比：" + str(round(number / len(self.weather_power), 2)))
        return new_weather_power.loc[:, ['C_TIME', 'C_REAL_VALUE', 'C_ABLE_VALUE']]

    def clean_limited_power_by_signal(self, name):
        weather_power1 = self.weather_power.copy()
        weather_power1["signal"] = weather_power1.apply(
            lambda x: self.signal_result(x["C_IS_RATIONING_BY_MANUAL_CONTROL"], x["C_IS_RATIONING_BY_AUTO_CONTROL"]),
            axis=1)
        weather_power1['c'] = weather_power1.apply(lambda x: 'cornflowerblue' if bool(x["signal"]) is True else 'pink',
                                                   axis=1)
        weather_power1.plot.scatter(x=self.opt.usable_power["env"], y='C_REAL_VALUE', c='c')
        plt.savefig(current_path + '/figs/信号法{}.png'.format(name))
        weather_power1 = weather_power1[weather_power1['signal'] == False]
        self.logger.info("信号法-未清洗限电前，总共有：{}条数据".format(len(self.weather_power)))
        self.logger.info("信号法-清除限电后保留的点有：" + str(len(weather_power1)) + " 占比：" + str(
            round(len(weather_power1) / len(self.weather_power), 2)))
        return weather_power1.loc[:, ['C_TIME', 'C_REAL_VALUE', 'C_ABLE_VALUE']]

    def signal_result(self, manual, auto):
        if int(manual) == 0:
            if int(auto) == 0:
                return False
            else:
                return True
        else:
            if int(auto) == 1:
                return True
            else:
                return False



if __name__ == '__main__':
    power = pd.read_csv('2023-12-01至2023-12-23实际功率导出文件.csv', date_parser=['时间'])
    weather = pd.read_csv('2023-12-01至2023-12-23气象站数据导出文件.csv', date_parser=['时间'])
    weather_power = pd.merge(weather, power, on='时间')  # 联立数据