import argparse

import os

import os.path as osp

import cv2

import matplotlib.pyplot as plt

import numpy as np

import wfdb

from sklearn.preprocessing import scale

from wfdb import rdrecord

# Choose from peak to peak or centered

# mode = [20, 20]

mode = 128

image_size = 128

output_dir = "../data"

# dpi fix

fig = plt.figure(frameon=False)

dpi = fig.dpi

# fig size / image size

figsize = (image_size / dpi, image_size / dpi)

image_size = (image_size, image_size)

def plot(signal, filename):

    plt.figure(figsize=figsize, frameon=False)

    plt.axis("off")

    plt.subplots_adjust(top=1, bottom=0, right=1, left=0, hspace=0, wspace=0)

    # plt.margins(0, 0) # use for generation images with no margin

    plt.plot(signal)

    plt.savefig(filename)

    plt.close()

    im_gray = cv2.imread(filename, cv2.IMREAD_GRAYSCALE)

    im_gray = cv2.resize(im_gray, image_size, interpolation=cv2.INTER_LANCZOS4)

    cv2.imwrite(filename, im_gray)

if __name__ == "__main__":

    parser = argparse.ArgumentParser()

    parser.add_argument("--file", required=True)

    args = parser.parse_args()

    ecg = args.file

    name = osp.basename(ecg)

    record = rdrecord(ecg)

    ann = wfdb.rdann(ecg, extension="atr")

    for sig_name, signal in zip(record.sig_name, record.p_signal.T):

        if not np.all(np.isfinite(signal)):

            continue

        signal = scale(signal)

        for i, (label, peak) in enumerate(zip(ann.symbol, ann.sample)):

            if label == "/":

                label = "\\"

            print("\r{} [{}/{}]".format(sig_name, i + 1, len(ann.symbol)), end="")

            if isinstance(mode, list):

                if np.all([i > 0, i + 1 < len(ann.sample)]):

                    left = ann.sample[i - 1] + mode[0]

                    right = ann.sample[i + 1] - mode[1]

                else:

                    continue

            elif isinstance(mode, int):

                left, right = peak - mode // 2, peak + mode // 2

            else:

                raise Exception("Wrong mode in script beginning")

            if np.all([left > 0, right < len(signal)]):

                one_dim_data_dir = osp.join(output_dir, "1D", name, sig_name, label)

                two_dim_data_dir = osp.join(output_dir, "2D", name, sig_name, label)

                os.makedirs(one_dim_data_dir, exist_ok=True)

                os.makedirs(two_dim_data_dir, exist_ok=True)

                filename = osp.join(one_dim_data_dir, "{}.npy".format(peak))

                np.save(filename, signal[left:right])

                filename = osp.join(two_dim_data_dir, "{}.png".format(peak))

                plot(signal[left:right], filename)