"""

    Produce train set and test set based on Apnea-ECG database.

"""

import numpy as np

from ECGSegment import ECGSegment

import os

def produce_database(database_name):

    """

    :param database_name:

    :return None:

    """

    if database_name == ["apnea-ecg", "train"] or database_name == ["apnea-ecg", "test"]:

        clear_id_set = np.load(database_name[0] + "_" + database_name[1] + "_clear_id.npy")

    else:

        raise Exception("Error database name.")

    dataset = []

    RRI_set = []

    RAMP_set = []

    EDR_set = []

    label_set = []

    for id in clear_id_set:

        eds = ECGSegment()

        eds.global_id = id

        eds.read_ecg_segment(1, database_name)

        eds.read_rri_ramp_edr()

        label_set.append(eds.label)

        RRI_set.append(eds.RR_intervals)

        RAMP_set.append(eds.R_peaks_amplitude)

        EDR_set.append(eds.EDR)

        dataset.append(eds)

    # substract mean of RRI,RAMP and EDR, RAMP * 10， EDR * 10000

    mean = np.mean(RRI_set, axis=1)

    mean = np.reshape(mean, (mean.shape[0], 1))

    rri_set = RRI_set - mean

    rri_set = np.reshape(rri_set, (rri_set.shape[0], rri_set.shape[1], 1))

    mean = np.mean(RAMP_set, axis=1)

    mean = np.reshape(mean, (mean.shape[0], 1))

    ramp_set = RAMP_set - mean

    ramp_set = np.reshape(ramp_set, (ramp_set.shape[0], ramp_set.shape[1], 1))

    ramp_set = ramp_set * 100

    mean = np.mean(EDR_set, axis=1)

    mean = np.reshape(mean, (mean.shape[0], 1))

    edr_set = EDR_set - mean

    edr_set = np.reshape(edr_set, (edr_set.shape[0], edr_set.shape[1], 1))

    edr_set = edr_set * 10000

    rri_amp_edr_set = np.concatenate([rri_set, ramp_set, edr_set], axis=2)

    if not os.path.exists("data/"):

        os.makedirs("data/")

    np.save("data/" + database_name[0] + "_" + database_name[1] + "_clear_dataset.npy", np.array(dataset))

    np.save("data/" + database_name[0] + "_" + database_name[1] + "_clear_rri.npy", np.array(rri_set))

    np.save("data/" + database_name[0] + "_" + database_name[1] + "_clear_ramp.npy", np.array(ramp_set))

    np.save("data/" + database_name[0] + "_" + database_name[1] + "_clear_edr.npy", np.array(edr_set))

    np.save("data/" + database_name[0] + "_" + database_name[1] + "_clear_rri_ramp_edr.npy", np.array(rri_amp_edr_set))

    np.save("data/" + database_name[0] + "_" + database_name[1] + "_clear_label.npy", np.array(label_set))

if __name__ == '__main__':

    # produce_database(["apnea-ecg", "train"])

    produce_database(["apnea-ecg", "test"])