"""噪声生成模块。"""

import copy
import random

import numpy as np
import pandas as pd

from .leak_simulator import simple_add_leak, simple_recover_wn, simple_simulation_pf


def add_noise_pd(data, noise_type, noise_para):
    output_data = copy.deepcopy(data)
    if type(output_data) == pd.core.frame.Series:
        if noise_type == "uni":
            for x in output_data.index:
                noise = (np.random.random() - 0.5) * 2
                output_data[x] = output_data[x] + noise * noise_para

        elif noise_type == "gauss":
            noise = np.random.normal(loc=0, scale=noise_para, size=output_data.shape)
            output_data = output_data + noise
    elif type(output_data) == pd.core.frame.DataFrame:
        if noise_type == "uni":
            noise = (np.random.random(size=output_data.shape) - 0.5) * 2
            output_data = output_data + noise * noise_para

        elif noise_type == "gauss":
            noise = np.random.normal(loc=0, scale=noise_para, size=output_data.shape)
            output_data = output_data + noise
    return output_data


def add_noise_number(data, noise_type, noise_para):
    output_data = copy.deepcopy(data)
    if noise_type == "uni":
        noise = (np.random.random() - 0.5) * 2
        output_data = output_data + noise * noise_para

    elif noise_type == "gauss":
        noise = random.gauss(0, noise_para)
        output_data = output_data + noise
    return output_data


def add_noise_number_flow(data, noise_para_mean, noise_para_std1, noise_para_std2):
    output_data = copy.deepcopy(data)
    noise_flag1 = np.random.random() - 0.5
    if noise_flag1 < 0:
        noise = noise_para_mean - abs(np.random.normal(loc=0, scale=noise_para_std1))
    else:
        noise = noise_para_mean + abs(np.random.normal(loc=0, scale=noise_para_std2))
    noise_flag2 = np.random.random() - 0.5
    if noise_flag2 < 0:
        noise_f = noise * (-1)
    else:
        noise_f = noise
    output_data = output_data + noise_f
    return output_data


def produce_noise_number(noise_type, noise_para):
    if noise_type == "uni":
        noise = (np.random.random() - 0.5) * 2
        noise = noise * noise_para
    elif noise_type == "gauss":
        noise = random.gauss(0, noise_para)
    else:
        noise = 0
    return noise


def add_noise_percentage_pd(data, noise_type, noise_para):
    output_data = copy.deepcopy(data)

    if type(output_data) == pd.core.frame.Series:
        if noise_type == "uni":
            for x in output_data.index:
                noise = (np.random.random() - 0.5) * 2
                output_data[x] = output_data[x] * (1 + noise * noise_para / 100)

        elif noise_type == "gauss":
            for x in output_data.index:
                noise = np.random.gauss(0, noise_para)
                output_data[x] = output_data[x] * (1 + noise / 100)
        # std_noise = noise.std()
    elif type(output_data) == pd.core.frame.DataFrame:
        if noise_type == "uni":
            noise = (np.random.random(size=output_data.shape) - 0.5) * 2
            output_data = output_data * (1 + noise * noise_para / 100)

        elif noise_type == "gauss":
            noise = np.random.normal(loc=0, scale=noise_para, size=output_data.shape)
            output_data = output_data * (1 + noise / 100)
        # std_noise = noise.std().mean()
    return output_data


def add_noise_in_wn_pf(
    wn,
    pipe_c_noise,
    timestep_list,
    pipe_coefficient,
    sensor_name,
    sensor_f_name,
    all_node,
    basic_demand_pd,
    noise_type,
    noise_para,
    leak_pipe,
    leak_flow,
):

    wn.options.time.duration = 0
    pipe_roughness_change = add_noise_pd(pipe_coefficient, noise_type, pipe_c_noise)
    wn = change_para_of_wn(wn, pipe_roughness_change)
    record_pressure = pd.DataFrame(index=timestep_list, columns=sensor_name)
    record_flow = pd.DataFrame(index=timestep_list, columns=sensor_f_name)
    record_noise_all = pd.DataFrame(
        index=pd.MultiIndex.from_product([timestep_list, all_node]),
        columns=basic_demand_pd.columns,
    )
    record_noise_all = record_noise_all.sort_index()

    # normal 获取添加噪声后的监测点数据
    for i in range(len(timestep_list)):
        wn, record_noise = change_node_demand(
            wn, basic_demand_pd, all_node, noise_type, noise_para
        )
        record_noise_all.loc[timestep_list[i]].loc[:, :] = record_noise
        pressure_temp, flow_temp = simple_simulation_pf(
            wn, sensor_name, sensor_f_name, [], []
        )
        record_pressure.iloc[i, :] = pressure_temp
        record_flow.iloc[i, :] = flow_temp

    # leak_simulation 获取添加漏损后的监测点数据
    record_pressure_leak = pd.DataFrame(index=timestep_list, columns=sensor_name)
    record_flow_leak = pd.DataFrame(index=timestep_list, columns=sensor_f_name)
    # 改_wz_________________________________________
    # add leak
    wn, whole_inf, add_pipe1 = simple_add_leak(wn, leak_flow, leak_pipe)
    # simulation
    for i in range(len(timestep_list)):
        record_noise = record_noise_all.loc[timestep_list[i]]
        wn = change_node_demand_leak(wn, record_noise, all_node)
        pressure_temp, flow_temp = simple_simulation_pf(
            wn, sensor_name, sensor_f_name, leak_pipe, add_pipe1
        )
        record_pressure_leak.iloc[i, :] = pressure_temp
        record_flow_leak.iloc[i, :] = flow_temp
    # delete leak
    wn = simple_recover_wn(wn, whole_inf)
    return wn, record_pressure, record_flow, record_pressure_leak, record_flow_leak


def change_node_demand(wn, basic_demand_pd, all_node, noise_type, noise_para):
    # 改_wz_____________________________________
    record_noise = pd.DataFrame(index=all_node, columns=basic_demand_pd.columns)
    for each_node in all_node:
        node = wn.get_node(each_node)
        num_columns = len(basic_demand_pd.columns)
        # 处理前N-1列（如果有）
        for i in range(num_columns - 1):
            # 获取原始值并添加噪声
            record_noise.loc[each_node].iloc[i] = (
                1 + produce_noise_number(noise_type, noise_para)
            ) * basic_demand_pd.loc[each_node].iloc[i]
            node.demand_timeseries_list[i].base_value = record_noise.loc[
                each_node
            ].iloc[i]
        # 处理最后一列（当列数>=1时）
        if num_columns >= 1:
            last_col = basic_demand_pd.columns[-1]
            original_last = basic_demand_pd.loc[each_node, last_col]
            record_noise.loc[each_node, last_col] = original_last
            node.demand_timeseries_list[-1].base_value = original_last

    return wn, record_noise


def change_node_demand_leak(wn, record_noise, all_node):
    sample_node = wn.get_node(all_node[0])
    # num_categories = len(sample_node.demand_timeseries_list)
    num_categories = 1
    for each in all_node:
        node = wn.get_node(each)
        for i in range(num_categories):
            node.demand_timeseries_list[i].base_value = record_noise.loc[each].iloc[i]
    return wn


def change_para_of_wn(wn, pipe_roughness_change):
    for pipe_name, pipe in wn.pipes():
        pipe.roughness = pipe_roughness_change[pipe_name]
    return wn


