
import numpy as np
import os
import sys
from tensorflow.keras.models import model_from_json
import wfdb
from utils import qrs_detect, comp_cosEn, save_dict
import tensorflow as tf
from matplotlib import pyplot as plt
from sklearn import preprocessing
"""
Written by:  <NAME>, <NAME>
             School of Instrument Science and Engineering
             Southeast University, China
             <EMAIL>

Save answers to '.json' files, the format is as {‘predict_endpoints’: [[s0, e0], [s1, e1], …, [sm-1, em-2]]}.
"""


def load_model(path_weight , path_json):  # 导入数据和模型，对数据进行归一化处理，输出数据和模型
    with open(path_json, "r") as f:
        json_string = f.read()  # 读取本地模型的json文件
    model = model_from_json(json_string)  # 创建一个模型
    #model = fixnet.mix_layer((600,1))
    model.load_weights(path_weight)  # 导入模型参数

    return model


def load_data(sample_path):
    sig, fields = wfdb.rdsamp(sample_path)
    length = len(sig)
    fs = fields['fs']

    return sig, length, fs


def ngrams_rr(data, length):
    grams = []
    for i in range(0, length - 12, 12):
        grams.append(data[i: i + 12])
    return grams


def challenge_entry(sample_path):
    """
    This is a baseline method.
    """
    path_weight = r'the101171800th training record  of 0.h5'
    path_json = r'fixresnet.json'
    model = load_model(path_weight, path_json)

    # model = unet((256, 1), 2, 0.001, maxpool=True, weights=None)
    # model.load_weights(path_weight)
    sig, _, fs = load_data(sample_path)
    min_max_scaler = preprocessing.MinMaxScaler()  # 数据归一化处理
    sig = min_max_scaler.fit_transform(sig)
    # sig = sig[:, 1]
    outdata = np.zeros(2048)
    num_class = []
    num_class.append(1)
    # r_peaks = qrs_detect(sig, fs=200)
    co = 0
    x = [1,0]
    r_peaks = wfdb.rdann(sample_path, extension="atr").sample
    former = 0
    for i, sample in enumerate(r_peaks):
        if int(sample - 1024) < 0 :
            former +=1
            data_2leads = np.hstack((sig[int(0):int(1024), 0], sig[int(0):int(1024), 1]))
            outdata = np.vstack((outdata, data_2leads))
            continue
        if int(sample + 1024) > len(sig):
            former += 1
            data_2leads = np.hstack(
                (sig[int(len(sig) - 1024):int(len(sig)), 0], sig[int(len(sig) - 1024):int(len(sig)), 1]))
            outdata = np.vstack((outdata, data_2leads))
            continue
        data_2leads = np.hstack((sig[int(sample - 512):int(sample + 512),0],sig[int(sample - 512):int(sample + 512),1]))
        outdata = np.vstack((outdata, data_2leads))
    outdata = np.delete(outdata, 0, axis=0)
    val_dataset = tf.data.Dataset.from_tensor_slices((outdata))
    val_dataset = val_dataset.batch(len(r_peaks))
    for x_batch_val in val_dataset:
        val_logits = model(x_batch_val, training=False)
        num_class = np.argmax(tf.nn.softmax(val_logits), 1)

    class_af = (np.sum(np.array(num_class) == 1) / len(num_class))
    endpoint = []
    flag = 0
    window_len = 8
    if class_af > 0.8:
        window_len = 7
    elif class_af > 0.6:
        window_len = 6
    elif class_af > 0.4:
        window_len = 5
    elif class_af > 0.2:
        window_len = 6
    else:
        window_len = 7
    #window_len = 5
    for i in range(len(num_class) - window_len):
        average_class = np.mean(num_class[i:i + window_len])  # 检测从一到零的变化
        if average_class == 1 and flag == 0:
            flag = 1
            startindex = r_peaks[i]
        if flag == 1 and average_class == 0:
            endpoint.append(startindex)
            endpoint.append(r_peaks[i])
            flag = 0
        elif i == len(num_class) - (window_len + 1) and flag == 1:
            endpoint.append(startindex)
            endpoint.append(len(sig) - 1)
    endpoint = np.array(endpoint).reshape(-1, 2).tolist()  # 如果没有检测出又房颤段那么为空，如果全是一则返回起始心拍和最终心拍，房颤段的检测标准为有连续大于五个的房颤心拍（1）或者全部是1的时候


    print(class_af)
    print(endpoint)

    pred_dcit = {'predict_endpoints': endpoint}

    return pred_dcit


if __name__ == '__main__':
    # DATA_PATH = sys.argv[1]
    # RESULT_PATH = sys.argv[2]

    DATA_PATH = r'F:\2019级\魏树鸿\21暑假\挑战赛\orgin data\training_II'
    RESULT_PATH = r'F:\2019级\魏树鸿\21暑假\挑战赛\orgin data\21'



    if not os.path.exists(RESULT_PATH):
        os.makedirs(RESULT_PATH)

    test_set = open(os.path.join(DATA_PATH, 'RECORDS'), 'r').read().splitlines()
    for i, sample in enumerate(test_set):
        #print(sample)
        sample_path = os.path.join(DATA_PATH, sample)
        pred_dict = challenge_entry(sample_path)
        #print(wfdb.rdheader(sample_path).comments)
        save_dict(os.path.join(RESULT_PATH, sample + '.json'), pred_dict)
