import array, math
import numpy as np
import matplotlib.pyplot as plt

###########################################
# Global user-set parameters
###########################################
f1 = 16		# How long the first (smoothing) filter is
f2 = 1024	# How long the second (baselining) filter is
f3 = 100	# How many samples per segment

###########################################
# Global user-set per-channel parameters
###########################################
ch1_Schmidt0=100	# Call the signal a 0 when it dips lower than this.
ch1_Schmidt1=400	#   "   "     "   " 1   "  "   "   higher "     ".

# Our final "RMS," doesn't take the sqrt or divide by f3. So compensate.
# sqrt(N/f3)=S; N/f3=S**2; N=f3*S**2
ch1_Schmidt0=f3 * ch1_Schmidt0 * ch1_Schmidt0
ch1_Schmidt1=f3 * ch1_Schmidt1 * ch1_Schmidt1

###########################################
# Global variables
###########################################
f1_ptr=0
f2_ptr=0
pos_in_current_segment = 0	# I.e., we're at the beginning of a segment.
n_reads=0			# Number of reads since forever.

f1_mask = f1-1			# If f1=16, then 0xF
f2_mask = f2-1			# If f2=64, then 0x3F

###########################################
# Global per-channel variables
###########################################
# 2B/entry, holds the last 'f2' ADC values.
ch1_circ_buffer1=array.array('H', (0 for _ in range(f1)))
ch1_circ_buffer2=array.array('H', (0 for _ in range(f2)))
ch1_f1_run_sum=0	# running sum for smoothing filter
ch1_f2_run_sum=0	# running sum for baselining filter
ch1_run_sumsqr=0	# running sum of baselined**2 on current segment
ch1_On = False		# Current state for the Schmidt trigger.

###########################################
# Arrays just to save up numbers for plotting
###########################################
plot_f1=[]
plot_baseline=[]
plot_baselined=[]
plot_rms=[]
plot_schmidt=[]

###########################################
# Library routines.
###########################################

# Find i such that 1<<i = numb
def shiftR_amount (numb):
    for i in range(15):
        if (1<<i == numb):
            return(i)
    assert False, str(numb) + " is not a power of 2"

# Returns On, changed
def Schmidt (On, run_sumsqr, Schmidt0, Schmidt1):
    if ((not On) and (run_sumsqr > Schmidt1)):
        return(True, True)
    if (On and (run_sumsqr < Schmidt0)):
        return(False, True)
    return (On, False)

def mavg_filters (val):
    global n_reads, f1_ptr, f2_ptr, pos_in_current_segment, ch1_f1_run_sum, \
	   ch1_f2_run_sum, ch1_run_sumsqr, ch1_On

    f1_ptr = n_reads & f1_mask
    ch1_f1_run_sum += (val-ch1_circ_buffer1[f1_ptr]) #Sum of last 'f1' readings
    f1_out = ch1_f1_run_sum >> divide_f1	# avg over last 'f1' readings
    ch1_circ_buffer1[f1_ptr] = val

    f2_ptr = n_reads & f2_mask
    ch1_f2_run_sum += (val-ch1_circ_buffer2[f2_ptr]) #Sum of last 'f2' readings
    ch1_circ_buffer2[f2_ptr] = val
    baseline  = ch1_f2_run_sum >> divide_f2	# avg over last 'f2' readings

    baselined = f1_out - baseline
    ch1_run_sumsqr += (baselined * baselined)

    # Pos_in_current_segment counts cyclically in [0,f3) to flag when we're
    # done with a segment.
    pos_in_current_segment += 1
    if (pos_in_current_segment==f3):
        pos_in_current_segment=0
        ch1_On,changed=Schmidt(ch1_On,ch1_run_sumsqr,ch1_Schmidt0,ch1_Schmidt1)
        #muscle_fun (ch1_On, changed)	# User's function to call.
        for i in range(f3):
            plot_rms.append(math.sqrt(ch1_run_sumsqr/f3))
            plot_schmidt.append(ch1_On*1000)
        ch1_run_sumsqr=0

    # Save the signals into lists to plot later.
    plot_f1.append(f1_out)
    plot_baseline.append(baseline)
    plot_baselined.append(baselined)

    n_reads += 1

def main():
    global divide_f1, divide_f2
    divide_f1 = shiftR_amount (f1)	# If f1=16, then 4 (since >>4 is /16)
    divide_f2 = shiftR_amount (f2)	# If f2=64, then 6 (since >>6 is /64)

    list = []
    dir = '..\\..\\data\\csv_files\\8_13_21'
    file='ll_la_cleaner'
    with open(dir + '\\' + file) as f:
        for line in f:
            list.append (int(float(line.strip())))
    raw = np.array (list)
    plt.plot (raw[1000:], label="raw")
    for i in range(len(raw)):
        mavg_filters (raw[i])
    #plt.plot (plot_f1, label=f"moving-avg{f1}")
    #plt.plot (plot_baseline, label=f"baseline{f2}")
    #plt.plot (plot_baselined, label=f"baselined{f2}")
    #plt.plot (plot_rms, label=f"RMS{f3}")
    #plt.plot (plot_schmidt, label="Schmidt")

    plt.legend (loc="upper left")
    plt.show()

main()

#ecg_8232_ll_la_1_cleaned: my normal ECG. F2=128 yields a really low T wave
#	that's as big (negatively) as QRS. But F2=1024 is only slightly
#	negative.
#ecg_8232_ll_ra_1: big T. Even with F2=1024, we get a big peak on T