#!/usr/bin/env python3

# authors M.Appleby

"""
# 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.signal import peak_widths, find_peaks
from scipy import asarray as ar,exp
from math import sqrt, pi
import argparse

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

def sigmoid(x, *p):
    L, x0, k, b = p
    return L / (1 + np.exp(-k * (x -x0))) + b


def getFWHM(filename, output_name, motor):
	motor = motor.capitalize()
	df = pd.read_csv(filename, sep='\t')
	#df = df[ [0,1] ]
	df = df.set_index( f"SAR-EXPMX:MOT_F{motor}.VAL", drop=True )
	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)')
        
	L_guess = np.max(y_vals)
	b_guess = np.min(y_vals)
	x0_guess = x_vals[np.argmin(np.abs(y_vals - (L_guess + b_guess) / 2))]
	k_guess = 1 / (x_vals[-1] - x_vals[0])

	s0 = [L_guess, x0_guess, k_guess, b_guess]
	params, params_covariance = curve_fit(sigmoid, x_vals, y_vals, p0=s0)
        #L_fit, x0_fit, k_fit, b_fit = params
	plt.plot(x_vals, y_vals)
	plt.plot(x_vals, sigmoid(x_vals, *params))
	plt.show()
	y_fit = sigmoid(x_vals, *params)

	dydx = np.gradient(y_fit, x_vals)
	if motor == 'Y':
		dydx=-dydx
	mean = x_vals[np.argmax(dydx)]
	sigma = (np.max(x_vals) - np.min(x_vals))/4
	print(mean, sigma)

	A= np.max(dydx)
	print(x_vals, '\\', y_vals)
	print(dydx)
	dydx[np.isnan(dydx)] = 0
	print(dydx)	

	p0 = [A, mean, sigma]
	parameters, covariance = curve_fit( gaus, x_vals, dydx, p0=p0 )
	x0_op = parameters[1]
	sigma_op = parameters[2]
	print(parameters)

	gauss_y = gaus(x_vals,*parameters)
	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.legend()
	plt.show()


	# find peak centre
	peaks = find_peaks( gauss_y )

	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 ) )




	#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/"
	)
	parser.add_argument(
		"-m",
		"--motor",
		help="X or Y",
		type=str,
		default="X"
	)

	args = parser.parse_args()
	getFWHM(args.input, args.output, args.motor)