import os

import re

import wfdb

import numpy as np

from wfdb import processing

import sys

from tqdm import tqdm

from scipy.signal import butter, iirnotch, filtfilt

from sklearn.decomposition import PCA

from ecgdetectors import Detectors #https://pypi.org/project/py-ecg-detectors/

def interpolate_nans(record):

    '''

    interpolate over NaN values in each lead of a given signal

    Parameters:

    - record (numpy.ndarray): 2D numpy array representing the ECG signal with leads in rows

                                shape (channels x samples)

    Returns:

    -numpy.ndarray: the signal with NaNs interpolated, leads with all NaNs remain unchanged.

    '''

    for i in range(record.shape[0]):  ###for future: instead of looping through each lead, applying interpolation in a vectorized manner

        lead_signal = record[i, :]

        if np.isnan(lead_signal).all():

            print(f"Warning: Lead {i} contains only NaNs.")

            continue

        if np.isnan(lead_signal).any():

            valid_mask = ~np.isnan(lead_signal)

            record[i, :] = np.interp(

                np.arange(len(lead_signal)),

                valid_mask.nonzero()[0],

                lead_signal[valid_mask]

            )

    return record

def filter_signal(signal, fs):

    """

    apply a high-pass and notch filter to an ECG signal

    Parameters:

    signal (array): the ECG signal to filter with shape (channels x samples)

                    can be 1D (single lead) or 2D (multichannel)

    fs (int): sampling frequency of the signal

    Returns:

    array or None: The filtered signal or None if an error occurs.

    """

    try:

        if np.isnan(signal).any():  # if the interpolation didn't work this will catch it

            print("NaN values detected during filtering, interpolating...")

            signal = np.nan_to_num(signal)  # Replace NaNs with zeros

        [b, a] = butter(3, (0.5, 40), btype='bandpass', fs=fs)

        signal = filtfilt(b, a, signal, axis=1)

        [bn, an] = iirnotch(50, 3, fs=fs)

        signal = filtfilt(bn, an, signal, axis=1)

        if signal.size == 0:

            print("Warning: Filtered signal is empty.")

        return signal

    except Exception as e:

        print(f"An error occurred during filtering: {e}")

        return None

#I'm using regular expression for this one coz info.txt is messy

def extract_patient_ids(info_path):

    '''

    extracts patient IDs from info.txt file

    returns:

    - list of IDs

    '''

    patient_ids = []

    id_pattern = re.compile(r'^\d{4}\b')

    with open(info_path, 'r') as file:

        for line in file:

            match = id_pattern.match(line)

            if match:

                patient_id = match.group()

                patient_ids.append(patient_id)

    return patient_ids

def preprocess_and_save(data_path, save_path):

    '''

    Preprocesses ECG signal data and saves the filtered signals and QRS complex indices.

    steps:

    - itaration over ECG records specified in 'info.txt'

    - NaN interpolation and filtering of the signals

    - QRS detection in a two-step process: 

        - an initial detection using the Pan-Tompkins algorithm

        - refinement with WFDB's `correct_peaks`

    parameters:

    - data_path (str): path to database

    - save_path (str): where to store preprocessed signals

    signals are saved in separate files with shape (number of channels x number of samples)

    '''

    os.makedirs(save_path, exist_ok=True)

    info_path = os.path.join(data_path, 'info.txt')

    patient_ids = extract_patient_ids(info_path)

    patient_ids_path = os.path.join(save_path, 'patient_ids.txt')

    with open(patient_ids_path, 'w') as f:

        for id in patient_ids:

            f.write(f"{id}\n")

    for patient_id in tqdm(patient_ids, desc='Processing Patients'):

        record_path = os.path.join(data_path, f"0{patient_id}") #I'm adding a zero before id coz thats how the files are saved

        save_signals_file = os.path.join(save_path, f"{patient_id}_signal.npy")

        save_qrs_file = os.path.join(save_path, f"{patient_id}_qrs.npy")

        record = wfdb.rdrecord(record_path)

        fs = record.fs

        signals = record.p_signal[(fs*60*5):, :].T #getting the shape (number of channels x number of samples)

        interpolated_signals = interpolate_nans(signals)

        filtered_signals = filter_signal(interpolated_signals, record.fs)

        if filtered_signals is None:

            print(f"Filtering failed for patient ID {patient_id}. Skipping QRS detection.")

            continue

        ### for future: fs is constant so filter coefficients don't have to be recalculated for every call

        ### move filtering from seperate function to preprocess_and_save

        #its way faster to find approximate peaks and correct them with wfdb rather than doing the whole search with wfdb

        detectors = Detectors(record.fs) 

        qrs_inds = detectors.pan_tompkins_detector(interpolated_signals[0])

        corrected_peak_inds = processing.correct_peaks(filtered_signals[0],

                                                    peak_inds=qrs_inds,

                                                    search_radius=int(record.fs*0.2),

                                                    smooth_window_size=int(record.fs*0.1))

        np.save(save_signals_file, filtered_signals)

        np.save(save_qrs_file, corrected_peak_inds)

if __name__ == "__main__":

    if len(sys.argv) != 3:

        print("Usage: python preprocess.py <data_path> <save_path>")

        sys.exit(1)

    data_path = sys.argv[1]

    save_path = sys.argv[2]

    preprocess_and_save(data_path, save_path)