import numpy as np
import pandas as pd


def fix_ep_zero(ep, dif_quantile, dif_mul):
    f_dif = ep.diff().fillna(0).abs()
    b_dif = ep.diff(periods=-1).fillna(0).abs()
    zeros = ep.loc[(f_dif > f_dif.quantile(dif_quantile) * dif_mul) &
                   (b_dif > b_dif.quantile(dif_quantile) * dif_mul)]
    ep.loc[zeros.index] = np.nan
    ep = ep.interpolate("linear")
    return ep


def fix_ep_minus(ep, dif_minus):
    dif = ep.diff().dropna()
    minus_dif = dif.loc[dif < dif_minus]
    for idx, val in minus_dif.iteritems():
        ep.loc[ep.index >= idx] += abs(val)
    return ep


def fix_ep_jump(ep, dif_quantile, dif_mul):
    dif = ep.diff().dropna()
    jump_dif = dif.loc[dif > dif.quantile(dif_quantile) * dif_mul]
    for idx, val in jump_dif.iteritems():
        ep.loc[ep.index >= idx] -= abs(val)
    return ep


def ep_minus_error(ep, dif_minus, max_minus_thres):
    dif = ep.diff().dropna()
    minus_dif = dif.loc[dif < dif_minus]
    if minus_dif.shape[0] > max_minus_thres:
        return True
    else:
        return False


def fix_ep(ep_raw: pd.Series):
    """
    修正电能曲线值：1. 零值突变点，2.负向跃变点，3.正向跃变点
    :param ep_raw: pd.Series
            原始电能数据时间序列
    :return ep: pd.Series
            修正后的电能数据时间序列
    """
    dif_quantile = 0.99  # 异常电能阈值
    dif_mul = 30  # 异常电能倍率阈值
    dif_minus = 0  # 负向电能阈值
    max_minus_thres = 1000  # 负向电能错误量最大值
    ep = ep_raw.copy()
    ep = fix_ep_zero(ep, dif_quantile, dif_mul)  # 修正零值突变点

    if ep_minus_error(ep, dif_minus, max_minus_thres):
        print("ep DATA ERROR! TOO MANY NEGATIVE ep TREND!")
        return
    else:
        ep = fix_ep_minus(ep, dif_minus)  # 修正负向跃变点
        ep = fix_ep_jump(ep, dif_quantile, dif_mul)  # 修正正向跃变点
        return ep


