'''
ECE 8527
Anna K. Chau
Intro to Machine Learning

Final Project -- Signal File Generation
--> script to generate .npy files of ECG signal data

usage: 
python signals.py <path to input .list containing records path>  <path for output .npy file>

note: .list needs to be like '00001_N1.dat 00001_N1.hea' to be able to use wfdb tools

'''
import os
import pandas as pd
import argparse
import tqdm
import sys
import numpy as np
import wfdb

def read_ecg(path, format='wfdb'):
    if format == 'wfdb':
        record = wfdb.rdrecord(path)
    return record.p_signal.T, record.fs, record.sig_name

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Generate CSV from files')
    parser.add_argument('input_file', type=str, help='path to RECORDS file (.list)')
    parser.add_argument('out_file', type=str, help='output .npy file containing the data')  # Update output file extension
    parser.add_argument('--root_dir', type=str, help='The root directory which the relative paths refer to'
                                                     'If None just use the folder where `input_file` is')
    args = parser.parse_args()

    files = pd.read_csv(args.input_file, header=None).values.flatten()
    
    if args.root_dir is None:
        folder = os.path.dirname(args.input_file)
    else:
        folder = args.root_dir

    all_data = []

    for i, f in enumerate(tqdm.tqdm(files, leave=False)):
        try:
            ecg, sample_rate, leads = read_ecg(os.path.join(folder, f), format='wfdb')
            #print(ecg)
            ecg_flat = np.ravel(ecg)
            all_data.append(ecg_flat) 
            #print(f"Data type of the signal in file {f}: {ecg_flat.dtype}")
        except FileNotFoundError as e:
            print(f"File not found: {f}")
            continue

    np_array = np.array(all_data)
    print(np_array)
    print(np_array.shape)

    np.save(args.out_file, np_array)  

    print("NumPy array saved to file successfully.")