import numpy as np
import scipy.special as spl
import matplotlib.pyplot as plt

# Définition des paramètres
n_medium = 1
n_sphere = 0.05 + 1j * 4.152  # Indice de réfraction de la sphère (partie réelle et imaginaire)

# Longueur d'onde fixée à 635 nm
lam = 635 

# Tailles des particules
a_min = 0.1  # 0.1 micron 
a_max = 100  # 100 micron 
num_a = 100  # Nombre de valeurs de taille de particules
a = np.linspace(a_min, a_max, num_a)

# Calcul des paramètres
k = 2 * np.pi / (lam * 1e-9) * n_medium

mu1 = mu = 1  # Ratio de la perméabilité magnétique de la sphère au milieu

# Initialisation des tableaux pour stocker les valeurs de x et z
x_values = np.zeros(num_a, dtype=complex)
z_values = np.zeros(num_a, dtype=complex)

###############################################################################
#======================= Function to calculate mie coefficients ===============
###############################################################################

def mie_coeff(n, x, z):
        
    # Calculating spherical bessel & henkel function for n and n-1 order
    jnx = spl.spherical_jn(n, x)
    jnx_1 = spl.spherical_jn(n-1, x)
    ynx = spl.spherical_yn(n, x)
    ynx_1 = spl.spherical_yn(n-1, x)
    jnz = spl.spherical_jn(n, z)
    jnz_1 = spl.spherical_jn(n-1, z)
    hnx = jnx + 1j * ynx
    hnx_1 = jnx_1 + 1j * ynx_1

    # Recurrence Relationship to calculate derivative of product of spherical bessel function & x
    x_jnxp = x * jnx_1 - n * jnx
    z_jnzp = z * jnz_1 - n * jnz
    x_hnxp = x * hnx_1 - n * hnx
    
    # Mie scattering coefficients
    an = (mu * (m**2) * jnz * x_jnxp - mu1 * jnx * z_jnzp) / (mu * (m**2) * jnz * x_hnxp - mu1 * hnx * z_jnzp)
    bn = (mu1 * jnz * x_jnxp - mu * jnx * z_jnzp) / (mu1 * jnz * x_hnxp - mu * hnx * z_jnzp)
    
    return an, bn

###############################################################################
#========================= Calculating cross sections for given particle ======
###############################################################################

# Calcul des sections efficaces pour chaque taille de particule
Qsca = np.zeros(num_a)
Qext = np.zeros(num_a)

m = n_sphere / n_medium

for i in range(num_a):
    x = (k * a[i]).astype(complex) 
    n_max = int(np.real(2 + np.max(x) + (4 * np.max(x) ** (1 / 3))))
    z = m * x
    
    # Enregistrement des valeurs de x et z
    x_values[i] = x
    z_values[i] = z
    
    Csca = 0
    Cext = 0
    
    for n in range(1, n_max + 1):
        an, bn = mie_coeff(n, x, z)

        Csca += (2 * np.pi / (k ** 2)) * (((2 * n + 1) * (abs(an) ** 2)) + ((2 * n + 1) * (abs(bn) ** 2)))
        Cext += (2 * np.pi / (k ** 2)) * ((2 * n + 1) * np.real(an + bn))

    Cgeom = np.pi * (a[i] ** 2)
    Qsca[i] = Csca / Cgeom
    Qext[i] = Cext / Cgeom
    Qabs = Qext - Qsca

# Sélection d'une valeur spécifique de i pour les coefficients d'excitation
i_selected = 0

# Calcul des coefficients d'excitation
an1, bn1 = mie_coeff(1, x_values[i_selected], z_values[i_selected])
an2, bn2 = mie_coeff(2, x_values[i_selected], z_values[i_selected])
an3, bn3 = mie_coeff(3, x_values[i_selected], z_values[i_selected])

# Coefficients d'excitation dipolaires électriques et magnétiques
a1 = (2 / (x_values ** 2)) * 3 * np.abs(an1 ** 2)
b1 = (2 / (x_values ** 2)) * 3 * np.abs(bn1 ** 2)

# Coefficients d'excitation quadrupolaires électriques et magnétiques
a2 = (2 / (x_values ** 2)) * 5 * np.abs(an2 ** 2)
b2 = (2 / (x_values ** 2)) * 5 * np.abs(bn2 ** 2)

# Coefficients d'excitation octopolaire électrique et magnétique
a3 = (2 / (x_values ** 2)) * 7 * np.abs(an3 ** 2)
b3 = (2 / (x_values ** 2)) * 7 * np.abs(bn3 ** 2)


###############################################################################
#=================================== Plotting =================================
###############################################################################

fig1 = plt.figure(figsize=(8, 4))
fig1.subplots_adjust(left=0.15, bottom=0.15, right=0.85, top=0.95, wspace=0.3, hspace=0.35)

plt.subplot(121)
plt.plot(a, Qsca, label=r'$Q_{sca}$', color='k', linestyle='-', marker='', markersize=6)
plt.plot(a, np.real(a1), label=r'$ED$', color='b', linestyle='--', marker='', markersize=6)
plt.plot(a, np.real(a2), label=r'$EQ$', color='m', linestyle='--', marker='', markersize=6)
plt.plot(a, np.real(a3), label=r'$EO$', color='c', linestyle='--', marker='', markersize=6)

plt.xlabel(r'Taille (μm)', fontsize=12)
plt.ylabel(r'$Q_{scat}$', fontsize=12)
plt.legend()
plt.tight_layout()

plt.subplot(122)
plt.plot(a, Qext, label=r'$Q_{ext}$', color='k', linestyle='-', marker='', markersize=6)
plt.plot(a, Qsca, label=r'$Q_{sca}$', color='b', linestyle='-', marker='', markersize=6)
plt.plot(a, Qabs, label=r'$Q_{abs}$', color='r', linestyle='-', marker='', markersize=6)

plt.xlabel(r'Taille (m)', fontsize=12)
plt.ylabel(r'$Q_{ext}$, $Q_{scat}$, $Q_{abs}$', fontsize=12)
plt.legend()
plt.tight_layout()

plt.show()