#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Script to view a plot of a data file from a DIY / low cost spectrometer via
audio.

Based on software from the DIY particle detector by Oliver Keller.
https://github.com/ozel/DIY_particle_detector

@author: Jeff Moe
@date: January 2022
"""

import pandas as pd
import numpy as np
import matplotlib as mpl
import matplotlib.pyplot as plt
from scipy.optimize import curve_fit
import matplotlib.gridspec as gridspec
import decimal as D
import msgpack
import argparse

mpl.rcParams['font.size']=12

parser = argparse.ArgumentParser()
parser.add_argument("file_name", help="filename to import, e.g. pulses.pkl")
args = parser.parse_args()

def poly1(x, a, b):
    return a*x + b

try:
    df = pd.read_pickle(args.file_name)
except:
    print('No such file')
    exit()

print(df)

THL = -300

lp = []
for i,p in df.iterrows():
    if p.pulse.min() < (THL):
        lp.append(p.pulse)

lp = np.asarray(lp[:])

count=0
loopcnt =0
areas = []
peaks = []
min_alpha_peak = 1243

min_g = -20
max_g = -3300
min_length = 44
max_length = 120
min_skip = min_length

xpulses = []
ypulses = []

fig = plt.figure()
wf = fig.add_subplot(111)
wf.set_xlim(0,120)
wf.set_xticks(np.arange(0,120,0.5e-3*48e3), minor=False)
wf.set_xticks(np.arange(0,120,0.1e-3*48e3), minor=True)
wf.set_xticklabels(list(map(str, np.arange(0,2.5,0.5))), minor=False)
#wf.set_xticklabels(list(map(str, np.arange(0.1,0.9,0.1))), minor=True)
wf.set_xlabel("Time [ms]", fontsize=14)
wf.xaxis.set_label_coords(0.96, -0.025) #right
wf.set_ylabel('Audio signal amplitude [arb. unit]', fontsize = 14)
#wf.set_ylim(-1300,400) # beta pulse range
wf.set_ylim(-17000,5000) # complete alpha pulse range
fig.tight_layout() #rect=(0.05,-0.010,1.01,1.02))

gmax=0
for i,y in enumerate(lp[:]):
    dydx=np.gradient(y)
    gy= dydx.min()
    gx= dydx.argmin()

    while True:
        if y[0] > THL and y[0] < np.abs(THL) and y.min() < THL and y[gx] <= np.abs(THL) and gy < min_g and gy > max_g:
            trigx=np.where(dydx < min_g)[0]
            trigy=dydx[dydx < min_g]
            peakx1=trigx[0]
            crossing_x1 = (y[peakx1+min_length:] > y[peakx1]).argmax() if (y[peakx1+min_length:] > y[peakx1]).any() else -1
            if crossing_x1 > 0:
                peakx2=peakx1+crossing_x1+min_length
                diff = np.abs(peakx2 - peakx1)
                peak = y[peakx1:peakx2].min()
                if diff >= min_length and diff <= max_length and peak <= THL :
                     area = np.absolute(y[peakx1:peakx2]).sum()
                     areas.append(area)
                     peak = np.absolute(peak) + y[peakx1]
                     peaks.append(peak)
                     ypulse = np.roll(y[peakx1-100:peakx2+100],50-y[peakx1:peakx2].argmin())[130:]
                     if ypulse.size > min_length:
                         xpulses.append(list(range(len(ypulse))))
                         ypulses.append(ypulse)





                     x = range(len(y))
                     # wf.plot(y, "black", alpha=0.1)
#                     if OVERLAY_PULSES:
                     wf.plot(ypulse, "green", alpha=(1/255)) #alpha=(1/255) for smallest setting, 0.1 otherwise
                         # uncomment below to print index labels next to max pulse amplitude
                         #wf.text(x[y[peakx1:peakx2].argmin()],y[peakx1:peakx2].min(), i) #lable pulse with id
#                     else:
                     if peak > min_alpha_peak:
                         wf.plot(x[peakx1:peakx2],y[peakx1:peakx2], "red") #, alpha=0.1)
                         print('peaky')
                     else:
                         wf.plot(x[peakx1:peakx2],y[peakx1:peakx2], "blue") #, alpha=0.1)
                         print('nonpeaky')





                     if gy < gmax:
                         gmax=gy
                     count+=1
                else:
                    peakx2=peakx1+min_skip
            else:
                peakx2 = peakx1 + min_skip
        else:
            peakx1=0
            remaining = len(y)
            if remaining <  min_length:
                break
            else:
                peakx2=min_skip
        y = np.delete(y, slice(0,peakx2))
        if len(y) < min_length:
            break
        else:
            dydx=np.delete(dydx, slice(0,peakx2))
            gy= dydx.min()
            gx= dydx.argmin()
        loopcnt+=1

time_diff = df.iloc[-1,0] - df.iloc[0,0]
time_diff = time_diff.total_seconds() / 60.0
cpm = count/time_diff
print("detected pulses:",count, "in", round(time_diff), "minutes ->", round(cpm,3), "CPM")
peaks=np.asarray(peaks)