# -*- coding: utf-8 -*-
"""
Analysis of Photoconductivity measurements
20/01/2018
BWe
"""
import argparse

import numpy as np
import matplotlib.pyplot as plt
from scipy.optimize import least_squares

# create an argument parser object
parser = argparse.ArgumentParser(description='Test argument parser')

# define the arguments to be parsed
parser.add_argument(
    'Wch',
    metavar='Wch',
    type=float,
    help='width (mm)')
	
parser.add_argument(
    'thick',
    metavar='thick',
    type=float,
    help='thickenss (nm)')

parser.add_argument(
    'Rterm',
    metavar='Rterm',
    type=float,
    help='Rterm (ohm)')	

parser.add_argument(
    'Lch',
    metavar='Lch',
    type=float,
    help='Lch (mm)')	
	
parser.add_argument(
    'Vapp',
    metavar='Vapp',
    type=float,
    help='voltage (V)')	

parser.add_argument(
    'Sample',
    metavar='Sample',
    type=str,
    help='Sample (name)')	
	
parser.add_argument(
    'fitstart',
    metavar='fitstart',
    type=float,
    help='fitstart (ns)')	
	
parser.add_argument(
    'fitend',
    metavar='fitend',
    type=float,
    help='fitend (ns)')	
	
#parser.add_argument(
    #'folderpath',
    #metavar='folderpath',
    #type=str,
    #help='folderpath')	
	
	
	# store all pared arguments in a list
args = parser.parse_args()

# assign arguments to variables
Wch = args.Wch
thick = args.thick
Rterm = args.Rterm
Lch = args.Lch
Vapp = args.Vapp
Sample = args.Sample
fitstart = args.fitstart
fitend = args.fitend
#folderpath = args.folderpath

# ///////////////////////////////
# USER INPUT
#Sample = 'A11'
#Wch    = 4e-3 # Channel width [m]
fitrange = [fitstart, fitend] # fit range [ns]
p_1exp = [0.1, 10, 1e-5] # guess parameters [A1, tau, t0, cst]
saveTraces = True # True or False 
directory_path = './'
# ///////////////////////////////


# Experimental conditions
# --------------------------------------------------------- 
#thick  = 400e-9 # Film thickness [m]
#Rterm  = 999 # Terminator resistance [ohm] 
Rscope = 1e6 # Scope impedance [ohm]
#Lch    = 4.5e-3 # Channel length [m]
#Vapp   = 26.52 # Applied voltage [V]

R = 1 / (1/Rterm + 1/Rscope)

# Fit functions 
# --------------------------------------------------------- 
def singleExp(t, p):
    A1, tau1, C = p
    # Create f = (ne',np'):
    f = A1*np.exp(-t/tau1) + C
    return f

def doubleExp(t, p):
    A1, tau1, A2, tau2, C = p
    # Create f = (ne',np'):
    f = A1*np.exp(-t/tau1) + A2*np.exp(-t/tau2)  + C
    return f

def score(p, t, d, fitfunc):
	# function to minimise
	err = d - fitfunc(t, p)
	return err


# File loading
# ---------------------------------------------------------
file_type = '.txt'

data = np.loadtxt(directory_path + Sample + file_type, delimiter='\t')
OD = data[0,:]
data = data[1:, :]

# Data processing
# ---------------------------------------------------------
plt.clf()
smax = [] # max conductivity
fitparams = [] # fit parameters
export = []
hd = [] # header
for col in range(1, data.shape[1]):
	tosc = data[:,0]
	Vosc = data[:,col]

	# Calculate conductivity
	sigma = Vosc / (R * (Vapp-Vosc)) * Lch / (thick * Wch)
	sigma /= 100 # convert in S/cm

	# Fit
	lowLimit = np.nanargmin(np.abs(1e9*tosc-fitrange[0]))
	highLimit = np.nanargmin(np.abs(1e9*tosc-fitrange[1]))
	tfit = tosc[lowLimit:highLimit].copy()
	sfit = sigma[lowLimit:highLimit].copy()
	# p_2exp = [0.1, 10, 1e-2, 100, 1e-5]

	plsq = least_squares(score, p_1exp, args=(tfit*1e9, sfit, singleExp), method='lm')

	# Plot
	plt.semilogy(tosc*1e9, sigma)
	plt.semilogy(tosc*1e9, singleExp(tosc*1e9, plsq.x))

	# Record results
	smax.append(singleExp(0,plsq.x))
	fitparams.append(plsq.x)
	hd.extend([str(OD[col]) + '_time', str(OD[col]) + '_sigma', str(OD[col]) + '_fit' ])
	export.extend([tosc, sigma, singleExp(tosc*1e9, plsq.x)])
	#print(OD[col], '\t', smax[-1], '\n')
	#print(OD[col], '\t', fitparams[-1], '\n')


plt.xlabel('Time [ns]')
plt.ylabel('Cond [S/cm]')
plt.show(0)

export = np.asarray(export)

# Save files
# ---------------------------------------------------------
fileName = directory_path + Sample + '_res.txt'
fn = directory_path + Sample + '_traces.txt'
pName = directory_path + Sample + '_fit.txt'
toSave = np.c_[OD[1:].tolist(), smax]
#np.savetxt(fileName, toSave[toSave[:,0].argsort()], header=("OD"+"\t"+"sigma0"), newline='\r\n', fmt="%i"+"\t"+"%f", comments='')
np.savetxt(pName, fitparams,  delimiter='\t'), #header=("A1"+"\t"+"tau"+"\t"+"C"))
if saveTraces == True:
	np.savetxt(fn, export.T, header=','.join(hd).replace(',','\t'), delimiter='\t', comments='')
tracefile = Sample + '_traces.txt'

print(tracefile)
