#!/usr/bin/env python
from KDataPy.database import kdatadb
import KDataPy.util
import ROOT
import scipy.signal as sig
import numpy as np
import matplotlib.pyplot as plt
plt.ion()

#acquire data location from the database
db = kdatadb()
viewResults = db.view('proc/raw', key = 'mf10a002')  #returns the location of all RAW .root files for this run

#create a KPulseAnalysisChain to remove the baseline and apply a highpass filter
bas = ROOT.KBaselineRemoval() #create a new baseline removal object
nyquist = 0.5/(2.016e-3)  #the nyquist frequency for typical heat channels
(b,a) = sig.iirfilter(2, 2.5/nyquist, btype='highpass') #b/a parameters for a second order butterworth highpass filter @ 2.5 Hz
myFilter = ROOT.KIIRSecondOrder(a[1], a[2], b[0], b[1], b[2])
(b,a) = sig.iirfilter(2, 40.0/nyquist, btype='lowpass') #b/a parameters for a second order butterworth lowpass filter @ 40.0 Hz
myFilterLowPass = ROOT.KIIRSecondOrder(a[1], a[2], b[0], b[1], b[2])
chain = ROOT.KPulseAnalysisChain()
chain.AddProcessor(bas)
chain.AddProcessor(myFilter)
chain.AddProcessor(myFilterLowPass)

pulseWidths = [] #i will store my results in this array. i could use a ROOT Tree, if I want to save the results for later
minVals = []
minPos = []
#my analysis function is defined here - which will be called by the 'looppulse' function below.
def myAnalysis(pulse,pta,**kwargs):
  global pulseWidths # type this in order to access the results array initiliazed above
  global minVals
  outputPulse = KDataPy.util.get_out(pta) #this gives you the output of the KPulseAnalysisChain defined above

  #assume a negative going pulse
  minval = np.amin(outputPulse)
  minpos = np.argmin(outputPulse)
  minPos.append(minpos)

  #find the 90% position
  ninPos = 0
  for i in range(minpos+1):
    if outputPulse[i] <= 0.9 * minval:  #remember, i'm assuming that min is a negative value
      ninPos = i
      break

  #find the 10% position
  tenPos = ninPos
  for i in range(ninPos, len(outputPulse)):
    if outputPulse[i] >= 0.1 * minval:
      tenPos = i
      break

  pulseWidths.append(tenPos - ninPos)
  minVals.append(minval)
  #print 'position', minpos, 'width', tenPos-ninPos, 'amplitude', minval


#now loop through the files returned by the database
for row in viewResults:
  doc = db[row['id']]
  filename = doc['proc1']['file']  # full path to each file of interest

  #find each H1 ZM1 pulse in the file, apply the "chain" and then pass the chain to your myAnalysis function
  #KDataPy.util.plotpulse(filename, name="H1 ZM1", match=True, pta=chain, analysisFunction = myAnalysis)
  KDataPy.util.looppulse(filename, name="H1 ZM1", match=True, pta=chain, analysisFunction = myAnalysis)


#analyze the results
plt.scatter(np.array(minVals), np.array(pulseWidths))

raw_input('hit the enter key to quit...')