import scipy.signal as signal
import numpy as np
import matplotlib.pyplot as plt

# Read a file that has one integer per line, and was sampled at 500 Hz.
# Subtract off the mean, so it's a zero-mean signal (this makes the FFT much
# more reasonable).
# Then return the resulting numpy array.
# The 'beg' and 'end' arguments allow you to subset just a portion of the file;
# e.g., to use only one heartbeat so that the resulting file will just repeat
# the one heartbeat over and over, making it very repeatable.
def read_data_file (fname, beg=0, end=None):
    data = read_data_file_raw (fname, beg, end)
    return (data - np.mean(data))

# Read an input file just like read_data_file(). This time, just print it out
# again, but with 10 values per line, and comma-separated.
# This makes it easy for an .c file to include it and make it part
# of an array definition.
def convert_csv_to_disco (fname, beg=0, end=None):
    import math, re
    data = read_data_file_raw (fname, beg, end)

    # Change the output file name so we don't overwrite the original file.
    outfile_name = fname
    if (re.search ("\..?.?.?$", fname)):
        outfile_name = re.sub ("\..?.?.?$", ".c_data", outfile_name, 1)
    outfile = open (outfile_name, "w")

    # Now output it, 10 values/line.
    n_rows = math.floor (data.size/10)
    for i in range (n_rows):
        s = ",".join ([str(int(n)) for n in data[i*10:(i+1)*10]])
        if (i<n_rows-1):
            s += ","
        outfile.write (s + "\n")
    outfile.close()

# Common routine used by both read_data_file() and convert_csv_to_disco().
def read_data_file_raw (fname, beg=0, end=None):
    infile = open (fname, "r")
    numbs = []
    for line in infile:
        # Should be a single number
        line.rstrip ('\r\n')
        numbs.append (int(line))
    data = np.array (numbs)
    infile.close()
    if (end != None):
        data = data[0:end]	# If the user only wants part of the file,
    data = data[beg:]		# then chop it to size.
    return (data)

# Assume we've already called plt.plot (data) one or more times.
# Just issue the final calls to get the plot onto the screen.
def show_time_series(title):
    plt.xlabel ("time (ms)")
    plt.legend()
    plt.title (title)
    plt.show()

# Given a time series of data, plot out its frequency components.
# Use np.fft.rfft() to do the work (that's just an FFT routine, but specialized
# to only work on real-valued inputs).
def plot_freq_spectrum (time_ser, Fsamp, title):
    xformed = np.fft.rfft(time_ser)	# coefficients are complex numbers
    xformed = np.abs (xformed)		# so get the magnitude

    # This tells us what frequency each FFT coefficient stands for.
    # It's fft_freqs[i] = .1Hz * i. That's because, from DFT theory,
    # exponential #k has frequency k*Fsamp/N. And we have Fsamp=500Hz, N=5000
    # samples. So each DFT coefficient has freq=500/5000=.1Hz higher than the
    # previous one. The rightmost frequency is 250Hz, which is half the sampling
    # frequency (if we had used .fft() rather .rfft(), we would have seen
    # a frequency spectrum from 0 up to 500 Hz).
    fft_freqs = np.fft.rfftfreq (len(time_ser), 1/Fsamp)

    plt.plot (fft_freqs, xformed)
    plt.xlabel('freq(Hz)'); plt.ylabel('DFT coef')
    plt.title('Frequency spectrum of ' + title)
    plt.show()

# Given a filter in sos form, plot its frequency response.
# We don't call plt.show() -- you must do that yourself after calling
# plot_filter_response(); this lets you plot two filters on the same graph.
def plot_filter_response (label, filt, Fsamp):
    # 'W' is an array of frequencies (the x axis of the plot). It runs from
    # 0 to fs/2, with the same units as fs (so in our case, 0 to 500 Hz).
    # 'H' is a corresponding array, for the gain (a complex number).
    (w,h) = signal.sosfreqz (filt,fs=Fsamp)
    gain = np.abs(h)		# gain is the magnitude of the complex number.
    #phase = np.angle(h)	# we can plot the phase if we want to.

    plt.plot (w, gain,  label=label)
    plt.ylabel ("gain")
    #plt.plot (w, phase, label=label);     plt.ylabel ("phase")
    plt.xlabel ("freq")
    plt.legend()

def part1(data_file_name):
    # Read the data file & plot it.
    data = read_data_file (data_file_name)
    time = np.arange (data.size) * .002	# 2ms per sample
    plt.plot (data, label="Original ECG")
    show_time_series ("Original ECG")

    # Plot the frequency spectrum.
    plot_freq_spectrum (data, 500, "original ECG")
    
def part2(data_file_name):
    data = read_data_file (data_file_name)
    time = np.arange (data.size) * .002	# 2ms per sample
    filt = signal.iirfilter (N=2, Wn=[59,61], btype='bandstop',
                             output='sos',fs=500)
    plot_filter_response ("simple filter", filt, 500)
    plt.show()

    data_filtered = signal.sosfilt (filt, data)
    plt.plot (time, data,  label="original data")
    plt.plot (time, data_filtered,  label="filtered data")
    show_time_series ('ECG before/after filtering')


#part1 ("../../ecg_data/ecg_normal_board_calm1.csv")
#part1 ("../../ecg_data/ecg_data_500samp_per_sec/ll_ra_warren1.csv")
#part2 ("../../ecg_data/ecg_normal_board_calm1.csv")
convert_csv_to_disco ("../../ecg_data/ecg_data_500samp_per_sec/ecg_normal_board_calm1_redone.csv", 1263*2, 1263*2 + 992)