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

### ========== USER PARAMETERS ========== ###
photonEnergy = [180, 380, 650, 930] # Photon energies given in eV
inelasticMeanFreePath = [7.79, 12.5, 19.07, 25.51] # IMFP in Carbon at each photon energy given in A
intensityBeforeDissolution = [154124, 388258, 211069, 54654] # SRPES peak intensity before dissolution at each photon energy
intensityAfterDissolution = [93216, 281611, 167056, 47772] # SRPES peak intensity after dissolution at each photon energy
depositedLayerThickness = 3.6 # Average layer thickness given in A
depositedMetalDensity = 21450 # Bulk Density given in kg/m3
depositedMetalAtomMass = 3.23945e-25 # Atomic Mass given in kg

### ========== FUNCTION DECLARATION ========== ###
def calculateAtomsPerArea():
    """Calculate the number of atoms per square meter from global variables"""
    return depositedMetalDensity * depositedLayerThickness * 1e-10 / depositedMetalAtomMass

def calculateDissolvedAtoms(d):
    """Calculate the number of dissolved atoms from global variables and the thickness d"""
    dissolved = np.zeros(len(photonEnergy))
    for i in range(len(photonEnergy)):
        attIntensity = intensityBeforeDissolution[i] * np.exp(-d / inelasticMeanFreePath[i])
        dissolved[i] = atomsPerArea * (attIntensity - intensityAfterDissolution[i]) / attIntensity
    return dissolved

### ========== CALCULATION AND MINIMIZATION ========== ###
atomsPerArea = calculateAtomsPerArea()
result = minimize(lambda d: np.std(calculateDissolvedAtoms(d)), x0=10, tol=1e12)
optimizedD = result.x[0]
dissolvedAtomsPercentage = calculateDissolvedAtoms(optimizedD) / atomsPerArea * 100
dissolvedAtomsMeanPercentage = np.mean(dissolvedAtomsPercentage)

### ========== OUTPUT ========== ###
print(result)
print(f"Dissolved Atoms for each Photon Energy:\n{', '.join([f'{v:.5g}' for v in calculateDissolvedAtoms(optimizedD)])}")
print(f"Average Dissolved Atoms: {np.mean(calculateDissolvedAtoms(optimizedD)):.5g}\nTotal Atoms: {atomsPerArea:.5g}")
print(f"Calculated Thickness: d = {optimizedD:.5g} A\nDissolution Amount: {dissolvedAtomsMeanPercentage:.4g} %")
plt.figure(figsize = (9,6), dpi = 240)
plt.plot(photonEnergy, dissolvedAtomsPercentage, "oc-", label = "Calculated")
plt.plot(photonEnergy, [dissolvedAtomsMeanPercentage]*len(photonEnergy), color = "k", label = "Average")
plt.title("Dissolution Fit")
plt.xlabel("Photon Energy [eV]")
plt.ylabel("Dissolution [%]")
plt.legend()
plt.show()