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

# Constantes
hbar = 1.0545718e-34  # J·s
m = 9.10938356e-31  # kg (masa del electrón)
eV_to_J = 1.60218e-19  # eV a J
a = 2e-10  # m (2 Å)
V0_eV = 14  # eV
V0_J = V0_eV * eV_to_J  # J

# Definir la función k1 y k2 en función de la energía
def k1(E_J):
    return np.sqrt(2 * m * (V0_J - E_J)) / hbar

def k2(E_J):
    return np.sqrt(2 * m * E_J) / hbar

# Ecuaciones trascendentales para pares e impares
def even_transcendental(E_eV):
    E_J = E_eV * eV_to_J
    k2_val = k2(E_J)
    k1_val = k1(E_J)
    return k2_val * np.tan(k2_val * a) - k1_val

def odd_transcendental(E_eV):
    E_J = E_eV * eV_to_J
    k2_val = k2(E_J)
    k1_val = k1(E_J)
    return k2_val / np.tan(k2_val * a) + k1_val

# Encontrar las raíces numéricas (intersecciones)
E_initial_guesses_even = [1, 5, 10]  # Guesses iniciales para estados pares
E_initial_guesses_odd = [3, 7, 12]  # Guesses iniciales para estados impares

E_eigenvalues_even = [fsolve(even_transcendental, E_guess)[0] for E_guess in E_initial_guesses_even]
E_eigenvalues_odd = [fsolve(odd_transcendental, E_guess)[0] for E_guess in E_initial_guesses_odd]

# Mostrar los eigenvalores
print("Eigenvalores para estados pares (eV):", E_eigenvalues_even)
print("Eigenvalores para estados impares (eV):", E_eigenvalues_odd)

# Graficar las ecuaciones trascendentales y los eigenvalores encontrados
E_eV = np.linspace(0, V0_eV, 1000)
E_J = E_eV * eV_to_J
k1_vals = k1(E_J) * a
k2_vals = k2(E_J) * a

# Ecuaciones trascendentales para graficar
lhs_even = k2_vals * np.tan(k2_vals)  # Para funciones pares
lhs_odd = -k2_vals / np.tan(k2_vals)  # Para funciones impares

# Graficar
plt.figure(figsize=(10, 6))
plt.plot(E_eV, lhs_even, label=r'$k_2 \tan(k_2 a)$ (pares)', color='blue')
plt.plot(E_eV, lhs_odd, label=r'$k_2 \cot(k_2 a)$ (impares)', color='green')
plt.plot(E_eV, k1_vals, label=r'$k_1$', color='red', linestyle='--')

# Añadir puntos para los eigenvalores
for E in E_eigenvalues_even:
    plt.plot(E, 0, 'bo')  # Estado par
for E in E_eigenvalues_odd:
    plt.plot(E, 0, 'go')  # Estado impar

plt.ylim(-20, 20)
plt.xlim(0, V0_eV)
plt.axhline(0, color='black', linestyle='--')
plt.title('Solución gráfica de las ecuaciones trascendentales y eigenvalores')
plt.xlabel('Energía (eV)')
plt.ylabel('Valores de las funciones')
plt.legend()
plt.grid(True)
plt.show()