"""

    This file include some functions for converting raw Apnea-ECG database to many txt files, each txt file including

     a 60s ECG segment corresponding with labels came from raw Apnea-ECG database.

    Before run this file, you first set path information.

    If you want to know more information about Apnea-ECG database, please see https://physionet.org/physiobank/database/apnea-ecg/.

"""

__version__ = '0.2'

__time__ = "2019.06.22"

__author__ = "zzklove3344"

import wfdb

import os

import numpy as np

# path information

# You need to set these file path before you run this file.

# Raw apnea-ecg database. You must download firstly.

APNEA_ECG_DATABASE_PATH = "G:/Apnea-ecg/raw records/"

# Folder for writing apnea-ecg 60s segments

SEGMENTS_BASE_PATH = "F:/Apnea-ecg/ecg segments/"

# The number of segments in train set

SEGMENTS_NUMBER_TRAIN = 17045

# The number of segments in test set

SEGMENTS_NUMBER_TEST = 17268

APNEA_ECG_TRAIN_FILENAME = [

    "a01", "a02", "a03", "a04", "a05", "a06", "a07", "a08", "a09", "a10",

    "a11", "a12", "a13", "a14", "a15", "a16", "a17", "a18", "a19", "a20",

    "b01", "b02", "b03", "b04", "b05",

    "c01", "c02", "c03", "c04", "c05", "c06", "c07", "c08", "c09", "c10"

]

# The number of 60s segments for every subject in a01-a20, b01-b05, c01-c10

TRAIN_LABEL_AMOUNT = [489, 528, 519, 492, 454,

              510, 511, 501, 495, 517,

              466, 577, 495, 509, 510,

              482, 485, 489, 502, 510,

              487, 517, 441, 429, 433,

              484, 502, 454, 482, 466,

              468, 429, 513, 468, 431]

APNEA_ECG_TEST_FILENAME = [

    "x01", "x02", "x03", "x04", "x05", "x06", "x07", "x08", "x09", "x10",

    "x11", "x12", "x13", "x14", "x15", "x16", "x17", "x18", "x19", "x20",

    "x21", "x22", "x23", "x24", "x25", "x26", "x27", "x28", "x29", "x30",

    "x31", "x32", "x33", "x34", "x35"

]

# The number of 60s segments for every subject in x01-x35

TEST_LABEL_AMOUNT = [523, 469, 465, 482, 505,

              450, 509, 517, 508, 510,

              457, 527, 506, 490, 498,

              515, 400, 459, 487, 513,

              510, 482, 527, 429, 510,

              520, 498, 495, 470, 511,

              557, 538, 473, 475, 483]

ECG_RAW_FREQUENCY = 100

class Mit2Segment:

    """

    Mit to 60s segments.

    """

    def __init__(self):

        self.raw_ecg_data = None  # list, raw ecg data

        self.denoised_ecg_data = None   # list, raw ecg data

        r"""basic attributes"""

        self.label = None  # int, 0 or 1

        self.database_name = None  # string, "apnea ecg"

        self.filename = None  # string, like "a01", "x02"

        self.local_id = None  # int, the ID in filename, like 101 in "a01"

        self.global_id = None  # int, global ID in database(train set or test set)

        self.samplefrom = None  # int, sample from where

        self.sampleto = None  # int, sample to where

        self.base_file_path = None      # string

    def write_ecg_segment(self, rdf):

        """

        Write minute-by-minute ECG segment to txt file.

        :param int rdf: 0 means to write to raw ecg file, 1 means to write to denoised ecg file.

        :return: None

        """

        # a01-a10, b01-b05 and c01-c10 belong to train set;

        # x01-x35 belong to test set.

        # if self.filename.find('x') >= 0:

        #   file_path = SEGMENTS_BASE_PATH + "test/" + str(self.global_id) + "/"

        # else:

        #   file_path = SEGMENTS_BASE_PATH + "train/" + str(self.global_id) + "/"

        if not os.path.exists(self.base_file_path):

            os.makedirs(self.base_file_path)

        if rdf == 0:

            filename = "raw_ecg_segment_data.txt"

            ecg_data = self.raw_ecg_data

        elif rdf == 1:

            filename = "denosing_ecg_segment_data.txt"

            ecg_data = self.denoised_ecg_data

        else:

            raise Exception("Error rdf value.")

        attr_name = "database_name file_name local_id samplefrom sampleto global_id label\n"

        # 将标签转化为数字

        if self.label == 'A':

            self.label = 1

        elif self.label == 'N':

            self.label = 0

        with open(self.base_file_path + filename, "w") as f:

            r"""attributes name """

            f.write(attr_name)

            r"""attributes value"""

            f.write(

                self.database_name[0] + " " + self.database_name[1] + " "

                + self.filename + " " + str(self.local_id) + " "

                + str(self.global_id) + " " + str(self.samplefrom) + " "

                + str(self.sampleto) + " " + str(self.label) + "\n")

            r"""data"""

            for value in ecg_data:

                f.write(str(value[0]) + "\n")

    def read_ecg_segment(self, rdf, database_name_or_path):

        """

        Read Minute-by-minute ECG segment from TXT file

        :param string or list database_name_or_path: the database or the file path you want to read

        :param int rdf: 0 means to read to raw ecg file, 1 means to read to denoised ecg file.

        :return: None

        """

        if rdf == 0:

            filename = "raw_ecg_segment_data.txt"

        elif rdf == 1:

            filename = "denosing_ecg_segment_data.txt"

        else:

            raise Exception("Error rdf value.")

        if database_name_or_path == ["apnea-ecg", "train"]:

            file_path = SEGMENTS_BASE_PATH + "train/" + str(self.global_id) + "/" + filename

        elif database_name_or_path == ["apnea-ecg", "test"]:

            file_path = SEGMENTS_BASE_PATH + "test/" + str(self.global_id) + "/" + filename

        else:

            file_path = database_name_or_path

        with open(file_path) as f:

            _ = f.readline()

            # attribute values

            attrs_value = f.readline().replace("\n", "").split(" ")

            self.database_name = [attrs_value[0], attrs_value[1]]

            self.filename = attrs_value[2]

            self.local_id = int(attrs_value[3])

            self.global_id = int(attrs_value[4])

            self.samplefrom = int(attrs_value[5])

            self.sampleto = int(attrs_value[6])

            self.label = int(attrs_value[7])

            self.base_file_path = SEGMENTS_BASE_PATH + self.database_name[1] + "/" + str(self.global_id) + "/"

            # ECG segment data

            ecg_data = []

            data_value = f.readline().replace("\n", "")

            while data_value != "":

                ecg_data.append(float(data_value))

                data_value = f.readline().replace("\n", "")

            if rdf == 0:

                self.raw_ecg_data = ecg_data

            elif rdf == 1:

                self.denoised_ecg_data = ecg_data

    def read_edr(self, flag):

        """

        flag为0时读取原始edr信号,为1时读取下采样之后的edr信号.

        :return: None

        """

        edr = []

        if self.filename.find('x') >= 0:

            if flag == 0:

                file_path = SEGMENTS_BASE_PATH + "test/" + str(self.global_id) + "/edr.txt"

            elif flag == 1:

                file_path = SEGMENTS_BASE_PATH + "test/" + str(self.global_id) + "/downsampling_EDR.txt"

            else:

                file_path = ""

                print("edr file path error....")

        else:

            if flag == 0:

                file_path = SEGMENTS_BASE_PATH + "train/" + str(self.global_id) + "/edr.txt"

            elif flag == 1:

                file_path = SEGMENTS_BASE_PATH + "train/" + str(self.global_id) + "/downsampling_EDR.txt"

            else:

                file_path = ""

                print("edr file path error....")

        with open(file_path) as f:

            data_value = f.readline().replace("\n", "")

            while data_value != "":

                edr.append(float(data_value))

                data_value = f.readline().replace("\n", "")

            edr = np.array(edr)

        return edr

def get_ecg_data_annotations(database_name, is_debug=False):

    """

    Read files in specified database.

    :param list database_name: Database you want to read.

                                Reserved paras, it must be ["apnea-ecg", "train"] or ["apnea-ecg", "test"] now.

    :param bool is_debug: whether is debug mode.

    :return list: ecg data and annotations.

    example: data_set = get_ecg_data_annotations("train", True)

    """

    data_annotations_set = []

    file_name_set = None

    no_apn = None

    if database_name[0] == "apnea-ecg":

        root_file_path = APNEA_ECG_DATABASE_PATH

        if database_name[1] == "train":

            file_name_set = APNEA_ECG_TRAIN_FILENAME

            no_apn = False

        elif database_name[1] == "test":

            file_name_set = APNEA_ECG_TEST_FILENAME

            no_apn = True

    # if database name is test, we first read label file

    test_label_set = []

    if no_apn is True:

        # read event-2.txt, which is test label downloading from PhysioNet

        test_annotation_path = root_file_path + "event-2.txt"

        with open(test_annotation_path) as f:

            lines = f.readlines()

            for line in lines:

                line = line.replace("\n", "")

                for index_str in range(len(line)):

                    if line[index_str] == "A" or line[index_str] == "N":

                        test_label_set.append(line[index_str])

    file_count = 0  # use when the database name is test.

    test_label_index = 0  # use when the database name is test.

    for name in file_name_set:

        if is_debug:

            print("process file " + name + "...")

        file_path = root_file_path + name

        ecg_data = wfdb.rdrecord(file_path)  # use wfdb.rdrecord to read data

        if no_apn is False:

            # use wfdb.rdann to read annotation

            annotation = wfdb.rdann(file_path, "apn")

            # annotation range

            annotation_range_list = annotation.sample

            # annotation

            annotation_list = annotation.symbol

        else:

            annotation_range_list = []

            annotation_list = []

            for index_label in range(TEST_LABEL_AMOUNT[file_count]):

                annotation_range_list.append(np.array(index_label * 6000))

                annotation_list.append(test_label_set[test_label_index])

                test_label_index += 1

            file_count += 1

            annotation_range_list = np.array(annotation_range_list)

        data_annotations_set.append([ecg_data, annotation_range_list, annotation_list, name])

    return data_annotations_set

def process_ecg_data_segments(database_name, data_annotations_set, is_debug=False):

    """

    Divide ECG data to minute-by-minute ECG segment.

    :param list database_name: name of database.

                               Reserved paras, it must be ["apnea-ecg", "train"] or ["apnea-ecg", "test"] now.

    :param list data_annotations_set: output of function get_ecg_data_annotations.

    :param bool is_debug: whether is debug mode.

    :return: None

    """

    data_set = []

    global_counter = 0  # use for global id

    base_floder_path = None

    if database_name[0] == "apnea-ecg":

        if database_name[1] == "train":

            base_floder_path = SEGMENTS_BASE_PATH + "/train"

        elif database_name[1] == "test":

            base_floder_path = SEGMENTS_BASE_PATH + "/test"

    # ecg data segments divide

    for data_annotation in data_annotations_set:

        segment_amount = len(data_annotation[2])

        for index_segment in range(segment_amount):

            eds = Mit2Segment()

            eds.database_name = database_name

            eds.samplefrom = data_annotation[1][index_segment]

            if (data_annotation[1][index_segment] + 6000) > len(data_annotation[0].p_signal):

                eds.sampleto = len(data_annotation[0].p_signal)

            else:

                eds.sampleto = data_annotation[1][index_segment] + 6000

            eds.raw_ecg_data = data_annotation[0].p_signal[eds.samplefrom:eds.sampleto]

            eds.label = data_annotation[2][index_segment]

            eds.filename = data_annotation[3]

            eds.local_id = index_segment

            eds.global_id = global_counter

            eds.base_file_path = SEGMENTS_BASE_PATH + "/" + database_name[1] + "/" + str(eds.global_id) + "/"

            eds.write_ecg_segment(rdf=0)

            global_counter += 1

            data_set.append(eds)

            if is_debug:

                print("---------------------------------------------------")

                print(("global id: %s,  file name: %s, local id: %s") % (

                    str(eds.global_id), eds.filename, str(eds.local_id)))

                print("---------------------------------------------------")

    if not os.path.exists(base_floder_path):

        os.makedirs(base_floder_path)

    # extra_info, this file store number of all ECG segments.

    with open(base_floder_path + "/extra_info.txt", "w") as f:

        f.write("Number of ECG segments\n")

        f.write(str(global_counter))

    return data_set

def produce_database(database_name, is_debug):

    """

    Produce database. It will write many txt files in SEGMENTS_BASE_PATH.

    :param list database_name: name of database.

                               Reserved paras, it must be ["apnea-ecg", "train"] or ["apnea-ecg", "test"] now.

    :param bool is_debug: whether is debug mode.

    :return: None

    """

    # read files from a01-a35, every file including whole ecg data and the corresponding annotation

    data_annotations_set = get_ecg_data_annotations(database_name, is_debug)

    # divide ECG data to minute-by-minute ECG segments

    _ = process_ecg_data_segments(database_name, data_annotations_set, is_debug)

def produce_all_database(is_debug):

    """

    Produce train database and test database.

    :param bool is_debug: whether is debug mode.

    :return: None

    """

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

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

if __name__ == '__main__':

    print("fileIO test statements")

    # if you want to generate train database, you can run follow statement.

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

    # if you want to generate test database, you can run follow statement.

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

    # if you want to generate train and test database, you can run follow statement.

    produce_all_database(True)