#!/usr/bin/env python3

# authors M.Appleby + some code donated by J.Beale

"""
Martin needs to tidy this
# aim
Calculate FWHM of x-ray beam from an edge scan

# protocol
give run directory

To run and edge scan in Slic - may be a limit on how far and fast motor can move
# to execute directly in slic:

from devices.knife_edge import KnifeEdge
kn = KnifeEdge

scan = Scanner(default_acquisitions=[daq], condition=None)
scan.scan1D(kn.x, 4, 5, 0.1, 10, 'test_knife_edge_evr_only', detectors=[], channels=["SARES30-LSCP1-FNS:CH5:VAL_GET"], pvs=[])
"""

import os, glob
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from scipy import stats
from scipy.optimize import curve_fit
from scipy.signal import peak_widths, find_peaks
from scipy import asarray as ar,exp
from math import sqrt, pi
from scipy.signal import savgol_filter
import argparse

from sfdata import SFDataFiles, SFScanInfo


def loadData(directory):
    num_of_acqs=len([name for name in os.listdir(os.path.join(directory,'data'))])
    print(num_of_acqs)
    diode_data = []
    for acquisition in range(1,num_of_acqs+1):
        with SFDataFiles(os.path.join(directory,'data','acq{:04}.BSDATA.h5'.format(acquisition))) as data:
            diode = data["SARES30-LSCP1-FNS:CH5:VAL_GET"].data
            diode = diode.mean()
            diode_data.append(diode)
    scan = SFScanInfo(os.path.join(directory,'meta','scan.json'))
    motor_positions = scan.readbacks
    print(motor_positions)
    print(len(diode_data), len(motor_positions))
    return motor_positions, diode_data

def gauss(x, *p):
    
    A, mu, sigma = p
    return A * np.exp(-(x-mu)**2/(2.*sigma**2))

#def gauss(x, C, A, x0, sigma):
#	return C + A*np.exp(-(x-x0)**2/(2.*sigma**2))

def getFWHM(motor_positions, diode_data):
	#print(df.columns)

	x_label='x motor'
	y_label='diode'
	x_vals = motor_positions*1000
	y_vals = [i for i in diode_data]
	y_vals = savgol_filter(y_vals, 15, 3)	

	fig, ax1 = plt.subplots()

	ax1.scatter(x_vals, y_vals, label='edge_scan', color='blue')
	plt.ylabel('Intensity (counts)')
	plt.xlabel('Motor position (um)')
	#plt.show()
	
	dydx = np.gradient(y_vals, x_vals)
	#dydx = dydx/max(dydx)

	mean = sum((x_vals)*dydx)/sum(dydx)
	sigma = sum(dydx*(mean)**2)/sum(dydx)

	A= 1/(sigma*2*sqrt(pi))
	
	dydx[np.isnan(dydx)] = 0

	p0 = [100, mean, 5 ]
	parameters, covariance = curve_fit( gauss, x_vals, dydx, p0=p0 )
	x0_op = parameters[1]
	# gaussian fit
	gauss_y = gauss(x_vals, *parameters)
	FWHM_x = np.abs(2*np.sqrt(2*np.log(2))*parameters[2])

    # find peak centre
	peaks = find_peaks( gauss_y )
	
    # find fwhm peak width
	fwhm_peak = peak_widths( gauss_y, peaks[0], rel_height=0.5 )

	#fwhm_str = x_vals[int( round( fwhm_peak[2][0], 0 ))]
	#fwhm_end = x_vals[int( round( fwhm_peak[3][0], 0 ))]

	#fwhm = fwhm_str - fwhm_end
	#fwhm_height = gauss_y[int( round( fwhm_peak[2][0], 0 ))]
	
    # find 1/e2 peak width
	#full_peak = peak_widths( gauss_y, peaks[0], rel_height=0.865 )

	#full_str = x_vals[int( round( full_peak[2][0], 0 ))]
	#full_end = x_vals[int( round( full_peak[3][0], 0 ))]

	#full = full_str - full_end
	#full_height = gauss_y[int( round( full_peak[2][0], 0 ))]
	
	#print( "FWHM = {0}".format( fwhm ) )
	print(FWHM_x)
	#print( "1/e2 = {0}".format( full ) )


	# plot
	ax2 = ax1.twinx()
	ax2.plot(x_vals, dydx,label ='derivative',color='red')
	ax2.plot(x_vals,gauss_y,label='Gaussian fit',color ='orange')
	
	plt.fill_between(x_vals,gauss_y,color='orange',alpha=0.5)
	plt.axvspan(x0_op+FWHM_x/2,x0_op-FWHM_x/2, color='blue', alpha=0.75, lw=0, label='FWHM = {0}'.format(FWHM_x))
	#plt.hlines(fwhm_height, fwhm_str, fwhm_end, color='green', label='FWHM (find_peaks) = {0}'.format(fwhm))
	#plt.hlines(full_height, full_str, full_end, color='yellow', label='1/e2 = {0}'.format(full))
	fig.legend(loc="upper left")
	fig.tight_layout()
	plt.show()
	#plt.savefig(output_name+filename.split('/')[-1]+'FWHM_{0}.png'.format(sciFWHM_x))
	return

if __name__ == '__main__':
	parser = argparse.ArgumentParser()
	parser.add_argument(
		"-i",
		"--input",
		help="location of input file",
		type=str,
		default="/sf/cristallina/data/p21736/raw/run0002/"
	)
	parser.add_argument(
		"-t",
		"--type",
		help="laser or xray",
		type=str,
		#default="/sf/cristallina/data/p19150/raw/run0371/"
	)
	args = parser.parse_args()
	x_data, y_data = loadData(args.input)
	getFWHM(x_data, y_data)