import numpy as np
from scipy.signal import resample
from baseline_wander_removal import bw_remover
import pywt
# import matplotlib.pyplot as plt
import codecs

def normalize(sig):
    return 2*((sig-np.min(sig))/(np.max(sig)-np.min(sig)))

def prepare_all_leads(path):
    if path.endswith(".txt"):
        filecp = codecs.open(path, "r", encoding="base64")
        sig = np.loadtxt(filecp, delimiter=',', unpack=True, encoding=None)
    elif path.endswith(".npy"):
        sig = np.load(path, allow_pickle=True)
        x = pywt.wavedec(sig[0], 'db6', level=2)[0]
        y = pywt.wavedec(sig[1], 'db6', level=2)[0]
        z = pywt.wavedec(sig[2], 'db6', level=2)[0]
        # visualize_sig([x, y, z])
        return x[None, :], y[None, :], z[None, :]
    # return sig
    freq = sig.shape[1] // 60
    sig[0] = bw_remover(freq, sig[0])
    sig[1] = bw_remover(freq, sig[1])
    sig[2] = bw_remover(freq, sig[2])

    # visualize_sig(sig)
    sig_length = freq*2
    total_samples = sig[0].shape[0] // sig_length

    lead_1 = []
    lead_2 = []
    lead_3 = []
    for i in range(total_samples):
        x = sig[0][i*sig_length:(i+1)*sig_length]
        x = pywt.wavedec(x, 'db6', level=1)[0]
        x = resample(x, 259)
        x = normalize(x)
        lead_1.append(x)

        y = sig[1][i*sig_length:(i+1)*sig_length]
        y = pywt.wavedec(y, 'db6', level=1)[0]
        y = resample(y, 259)
        y = normalize(y)
        lead_2.append(y)

        z = sig[2][i*sig_length:(i+1)*sig_length]
        z = pywt.wavedec(z, 'db6', level=1)[0]
        z = resample(z, 259)
        z = normalize(z)
        lead_3.append(z)
    
    return np.asarray(lead_1), np.asarray(lead_2), np.asarray(lead_3)