import os
import pandas as pd
import numpy as np
np.random.seed(42)
from cache.data_cleaning import cleaning
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.weather_power = self.weather_power[self.weather_power[self.opt.usable_power["env"]] >= 0]
        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
        self.saveVar(os.path.dirname(current_path)+'/var/glob_rp.pickle', glob_rp)
        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, signal=False):
        glob_rp = self.segment_statis()
        new_weather_power = []
        # fig, ax = plt.subplots()
        for index, row in self.weather_power.iterrows():
            zfs = row[self.opt.usable_power["env"]]
            rp = row['C_REAL_VALUE']
            s = int(row['signal']) if signal is True else 2
            if self.filter_unlimited_power(zfs, rp, glob_rp):
                row['c'] = 'red' if s == 0 or s == 2 else 'cornflowerblue'
                new_weather_power.append(row)
            else:
                row['c'] = 'blue' if s == 1 or s == 2 else 'pink'
                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)))
        self.args.save_args_yml(self.opt)
        return new_weather_power.loc[:, ['C_TIME', 'C_REAL_VALUE', 'C_ABLE_VALUE']]

    def clean_limited_power_by_signal(self):
        # powers = self.weather_power.copy()
        self.weather_power["signal"] = self.weather_power.apply(lambda x: self.signal_result(x["C_IS_RATIONING_BY_MANUAL_CONTROL"], x["C_IS_RATIONING_BY_AUTO_CONTROL"]), axis=1)
        # powers['c'] = self.weather_power.apply(lambda x: 'blue' if bool(x["signal"]) is True else 'red', axis=1)
        # self.weather_power.plot.scatter(x=self.opt.usable_power["env"], y='C_REAL_VALUE', c='c')
        # plt.savefig(current_path + '/figs/限电.png')
        new_weather_power = self.weather_power[self.weather_power['signal'] == False]
        return new_weather_power.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

def plt_point(glob_rp, coe):
    """
    对最后的可用功范围进行画图
    :param glob_rp: 总辐射分段
    :param width: 偏移量
    :return:
    """
    fig, ax = plt.subplots()
    for x, row in glob_rp.items():
        ax.scatter(x, row[0]-row[1]*coe, s=10, alpha=0.5)
        ax.scatter(x, row[0]+row[1]*coe, s=10, alpha=0.5)
        ax.scatter(x, row[0], s=12, alpha=0.75, marker='x')
    plt.savefig('./7月份-优化1.png')


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='时间')  # 联立数据