# -*- coding: utf-8 -*-
"""
 @Time    : 2020/6/14 19:35
 @Author  : 杰森·家乐森
 @File    : sae_main.py
 @Software: PyCharm
"""
import json
import time
import torch
import itertools
import numpy as np
import pandas as pd
import torch.nn as nn
import torch.utils.data as Data
from sae_diagnosis import diagnosis
from torch.autograd import Variable
from sklearn.metrics import r2_score
from sklearn.preprocessing import MinMaxScaler


def get_diagnosis(model, fault_data, true_fault, spe, name):
    """
    诊断函数
    :param model: 模型信息
    :param fault_data: 故障数据
    :param true_fault: 真实故障
    :param spe: spe限值
    :param name: 诊出结果保存文件名
    :return:
    """
    dection_array, amp_array = diagnosis(model, fault_data[:100, :], spe)
    dection_array = dection_array * 10
    fault_array = true_fault[:100, :]
    detection = fault_array - dection_array
    FDR = 1 - np.sum(detection > 0) / np.sum(fault_array > 0)
    FAR = np.sum(detection == -10) / np.sum(fault_array == 0)
    detection_cfg = {
        "FDR": FDR,
        "FAR": FAR
    }
    with open(name, 'w') as f:
        f.write(json.dumps(detection_cfg))


def diag(hi, lo, name):
    # data = np.array(pd.read_excel('旧数据.xlsx'))
    # data1 = np.array(pd.read_excel('旧数据.xlsx'))
    with open('train_data.json', 'r') as f:
        data = np.array(json.load(f)['data'])
    with open('test_data.json', 'r') as f:
        data1 = np.array(json.load(f)['data'])[:500, :]
    mms = MinMaxScaler()
    mms.fit_transform(data)
    data = mms.transform(data[:100, :])
    fault_dir = np.zeros(data.shape)
    index_array = np.random.randint(0, 10, (data.shape[0], 1))
    index_array1 = np.random.randint(0, 10, (data.shape[0], 1))
    index_array2 = np.random.randint(0, 10, (data.shape[0], 1))
    for n in range(data.shape[1]):
        fault_dir[np.where(index_array[:, 0] == n), n] = 1
        fault_dir[np.where(index_array1[:, 0] == n), n] = 1
        fault_dir[np.where(index_array2[:, 0] == n), n] = 1
    amp_array = np.random.random_sample(data.shape) * 1
    fault_dir[np.where(amp_array == 0)] = 0
    fault_a = fault_dir * amp_array * (hi - lo) + fault_dir * lo
    fault_a = fault_dir * 0.5
    fault_data = data + fault_a
    data_fault = mms.inverse_transform(fault_data)
    bias_data = data_fault - data1[:100, :]
    true_fault = np.zeros(data.shape)
    true_fault[np.where(fault_a > 0)] = 10
    true_fault[np.where(fault_a == 0)] = 0
    with open('model.json', 'r') as f:
        model_cfg = json.load(f)
    model = model_cfg["model"]
    spe = model_cfg["spe"]
    # spe = 0.009
    get_diagnosis(model, fault_data, true_fault, spe, 'rbct%d.json' % name)
    print('OK')


if __name__ == '__main__':
    # hi = [0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 1.0, 2.0, 3.0, 5.0, 10.0]
    # lo = [0, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 1.0, 2.0, 3.0, 5.0]
    # for i in range(len(hi)):
    #     diag(hi[i], lo[i], i)

    diag(1, 0, 4)
    print('OK')