#!/usr/bin/env python3

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

"""
# aim
Calculate FWHM of x-ray beam from an edge scan

# protocol
complete edge scan
## IMPORTANT ##
- save data as .txt file 

# usage
python convert-scan-for-pyfai.py -j <jugfrau-name> -s <path to scan file> -n <name of output file>

# output
creates a .png of fit including FWHM title
"""

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 import asarray as ar,exp
from math import sqrt, pi
import argparse

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


def getFWHM(filename, output_name):
	df = pd.read_csv(filename, sep='\t', index_col=0)
	print(df)
	#print(df.columns)

	x_label=df.index.name
	y_label=df.columns[0]
	x_vals = df.index.tolist()
	y_vals = df[y_label].tolist()

	plt.plot(x_vals, y_vals, label = 'edge_scan')
	plt.ylabel('Intensity (counts)')
	plt.xlabel('Motor position (mm)')

	dydx = np.gradient(y_vals, x_vals)

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

	A= 1/(sigma*2*sqrt(pi))
	print(x_vals, '\\', y_vals)
	print(dydx)
	dydx[np.isnan(dydx)] = 0
	print(dydx)	

	param_optimised,param_covariance_matrix = curve_fit(gaus,x_vals,dydx,p0=[min(y_vals), A, mean,sigma],maxfev=5000)
	offset_op, amp_op, x0_op, sigma_op = param_optimised[0], param_optimised[1], param_optimised[2], param_optimised[3]

	print(param_optimised)

	gauss_y = gaus(x_vals,*param_optimised)
	FWHM_x = np.abs(2*np.sqrt(2*np.log(2))*sigma_op)

	plt.plot(x_vals, dydx, label = 'derivative')
	plt.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='green', alpha=0.75, lw=0, label='FWHM = {0}'.format(FWHM_x))
	print(FWHM_x)
	#plt.plot(x_label, y_label, data=df)
	plt.show()
	plt.savefig(output_name+filename.split('/')[-1]+'FWHM_{0}.png'.format(FWHM_x))
	return

if __name__ == '__main__':
	parser = argparse.ArgumentParser()
	parser.add_argument(
		"-p",
		"--path",
		help="path of input and output file, not currently in use",
		type=str,
		default="/sf/cristallina/data/p21224/res/pshell/edge_scans/"
	)
	parser.add_argument(
		"-i",
		"--input",
		help="location of input file",
		type=str,
		default="/sf/cristallina/data/p21224/res/pshell/edge_scans/0.5_x/0.5_x"
	)
	parser.add_argument(
		"-o",
		"--output",
		help="output path to save figure",
		type=str,
		default="/sf/cristallina/data/p21224/res/pshell/edge_scans/"
	)
	args = parser.parse_args()
	getFWHM(args.input, args.output)