# -*- coding: utf-8 -*-
"""
@author: wenger
4 sept 2019 - Modified to calculate first bimolecular recombination and then Saha's equation
"""

import argparse

import numpy as np
import matplotlib.pyplot as plt
from scipy.integrate import odeint
from scipy.constants import h, k, m_e, pi, e
from scipy.optimize import root

# Notes
# FWHM was remeasured with a fast photodiode (35ps rise time). It is 3.47 ns (FWHM)

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

# define the arguments to be parsed
parser.add_argument(
    't_sample',
    metavar='t_sample',
    type=float,
    help='t_sample (nm)')
	
parser.add_argument(
    't_ref',
    metavar='t_ref',
    type=float,
    help='t_ref (nm)')

parser.add_argument(
    'excdensity1',
    metavar='excdensity1',
    type=float,
    help='excdensity1 (cm-3)')
	
parser.add_argument(
    'k1',
    metavar='k1',
    type=float,
    help='k1 (s-1)')

parser.add_argument(
    'k2',
    metavar='k2',
    type=float,
    help='k2 (cm3/s)')
	
parser.add_argument(
    'k3',
    metavar='k3',
    type=float,
    help='k3 (cm6/s-2)')

parser.add_argument(
    'E_B',
    metavar='E_B',
    type=float,
    help='E_B (eV)')

# binding energy args added by Jongchul
# store all pared arguments in a list
args = parser.parse_args()

# assign arguments to variables
t_sample = args.t_sample
t_ref = args.t_ref
excdensity1 = args.excdensity1
k1 = args.k1
k2 = args.k2
k3 = args.k3
E_B = args.E_B

# biding energy variables added by Jongchul
# Parameters for Recombination
I0 = 0   # initial carrier density cm-3
ntot = excdensity1 # total photons absorbed calculated by Jongchul in cm-3

ntot = ntot * t_ref/t_sample

# Kinetic model parameters
# k1 = 1.8e6	    # s-1
# k2 = 6.8e-10  # 6.8e-10 cm3/s from Tim's paper
# k3 = 1.6e-28    # cm6/s-2 from Laura's paper
mu = 0      # time offset of gaussian
fwhm = 3.74e-9 # measured FWHM
sigma = fwhm / (2*np.sqrt(2*np.log(2)))

p = [k1, k2, k3, ntot]

# ODE definition
def df(c0,t,p):
    ne = c0
    k1, k2, k3, ntot = p
    G = gauss(t, mu, sigma)
    dnedt = ntot*G -k1*ne - k2*ne**2 - k3*ne**3
    return dnedt

def gauss(t, mu, sig):
    return 1/(sig*np.sqrt(2*np.pi)) * np.exp(-(1/2)*((t-mu)/sig)**2)

# Vary I0 and solve equation

[plt.close(i) for i in plt.get_fignums()] # close all existing figures

fig1, (ax1, ax2, ax3)= plt.subplots(3,1, figsize=(6, 8), sharex=False)

num_pnts = 1000 # number of points in the simulation
res = np.empty((num_pnts,0))
q = p.copy()
q[0] = q[1] = q[2] = 0  # This is for the curve without recombination

p[3] = ntot
q[3] = ntot
t = np.linspace(-1e-8, 0.5e-6, num_pnts) # in ns
sol = odeint(df, I0, t, args=(p,))
solg = odeint(df, I0, t, args=(q,))
res= np.append(res, sol, axis=1)
res= np.append(res, solg, axis=1)
ax1.semilogy(t*1e9, sol, label=("{:.2e}".format(ntot)))
ax1.semilogy(t*1e9, solg, '-')
ax2.plot(t*1e9, sol)
ax2.plot(t*1e9, solg, '-')



ax2.set_xlabel("Time [ns]")
for ax in [ax1, ax2]:
    ax.set_ylabel("n(t) [a.u.]")
    ax.set_xlim(right=100)
ax1.set_ylim(bottom=1e15)
# fig1.legend()

sol = sol.reshape(-1) # required for Saha's function

# Calculate free carriers using Saha's equation
# ---------------------------------------------

# Set physical variables
T = 300 # temp in K
#E_B = 0.011 # exciton binding energy in eV
# biding energy changed to args by Jongchul
mu_ex = 0.104 * m_e


# Define Saha as a root finding problem
def saha(n_ex, *args):
    """Function for numerically solving the Saha equation with root.
    
    Parameters
    ---------
    n_ex : float
        exciton density in cm^-3
    *args : list or tuple of float
        n_total (total carrier density in cm^-3), E_B (exciton binding energy in eV)
        
    Returns
    -------
    sol : float
       solution of the Saha equation to root of with root 
    
    """
    
    # Unpack *args
    n_total, E_B, mu_ex = args
    
    # Calculate free carrier density
    n_free = n_total - n_ex
    
    # Calculate free-carrier fraction
    a = n_free / n_total if n_total > 0 else 0
    
    # Solution to minimise
    res = (a**2 / (1 - a)) - ((1 / n_total) * (2 * pi * mu_ex * k * T / h**2)**(3/2) * np.exp(-(e * E_B) / (k * T))) if n_total > 0 else 0
    return res

# Define function for calculating free-carrier fraction
def a(n, E_B, mu_ex):
    # Solve the Saha equation to find the exciton density
    n_ex = []
    for m in n:
        ex0 = np.array([m / 1000]) # starts with a guess (lower end)
        solution = root(saha, ex0, args=(m, E_B, mu_ex))
        n_ex.append(solution.x[0])
    n_ex = np.array(n_ex)
    
    # # Convert to free carrier fraction
    a = []
    for i in range(len(n)):
        a.append((n[i] - n_ex[i]) / n[i] if n[i] > 0 else 0)
    
    return a

# Calculate free-carrier fraction as a function of charge density
# and exciton binding energy
a1 = a(sol*1e6, E_B, mu_ex) # Saha's equation is calculated in m3 not cm3, hence 1e6

sol_saha = np.nan_to_num(sol*a1)
#print(f'Time at max: {t[sol_saha.argmax()]*1e9:.2f} ns') # Find the time corresponding to n_max
#print(f'Carrier density at max: {sol_saha.max():.3e} cm3')
#revised by Jongchul for labviews

print(t[sol_saha.argmax()]*1e9) #revised by Jongchul for labviews
print(sol_saha.max()) #revised by Jongchul for labviews
print(a1[sol_saha.argmax()]) #revised by Jongchul for labviews

#ax2.plot(t*1e9, sol_saha)
#ax1.semilogy(t*1e9, sol_saha)
#ax3.semilogx(sol, a1)
#ax3.set_ylim(0)
#ax3.set_xlabel("n (cm^-3)")
#ax3.set_ylabel("n_free / n_total")
#plt.tight_layout()
#plt.show(block=True)
# 
# Format data for saving
res = np.c_[t, res, sol_saha, sol_saha, a1]
np.savetxt('res.txt', res, delimiter='\t')