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

# Constantes físicas
hbar = 1.0545718e-34  # J·s
m = 9.10938356e-31  # kg (masa del electrón)
eV_to_J = 1.60218e-19  # eV a J
Å_to_m = 1e-10  # Å a metros

# Parámetros del pozo
a = 2 * Å_to_m  # Ancho del pozo en metros (2 Å)
V0_eV = 14  # Altura del pozo en eV
V0_J = V0_eV * eV_to_J  # Convertir a Joules

# Definir 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

# Eigenvalores de energía (de la pregunta 4)
E_eigenvalues_eV = [1.41, 5.52, 11.45]  # Primeros tres eigenvalores (dos pares, uno impar)
E_eigenvalues_J = [E * eV_to_J for E in E_eigenvalues_eV]

# Funciones de onda para estados pares e impares
def psi_even(x, E_J, A):
    k2_val = k2(E_J)
    k1_val = k1(E_J)
    if x < -a:
        return A * np.exp(k1_val * (x + a))
    elif -a <= x <= a:
        return A * np.cos(k2_val * x)
    else:
        return A * np.exp(-k1_val * (x - a))

def psi_odd(x, E_J, B):
    k2_val = k2(E_J)
    k1_val = k1(E_J)
    if x < -a:
        return B * np.exp(k1_val * (x + a))
    elif -a <= x <= a:
        return B * np.sin(k2_val * x)
    else:
        return B * np.exp(-k1_val * (x - a))

# Normalización de las funciones de onda
def normalize_even(E_J):
    A = 1.0  # Inicialmente asumimos A = 1
    norm_factor = np.sqrt(quad(lambda x: psi_even(x, E_J, A)**2, -10*a, 10*a)[0])
    return 1.0 / norm_factor

def normalize_odd(E_J):
    B = 1.0  # Inicialmente asumimos B = 1
    norm_factor = np.sqrt(quad(lambda x: psi_odd(x, E_J, B)**2, -10*a, 10*a)[0])
    return 1.0 / norm_factor

# Graficar las funciones de onda normalizadas
x_vals = np.linspace(-5 * a, 5 * a, 1000)

plt.figure(figsize=(10, 6))

# Estado par 1
A1 = normalize_even(E_eigenvalues_J[0])
psi_vals_even_1 = [psi_even(x, E_eigenvalues_J[0], A1) for x in x_vals]
plt.plot(x_vals / Å_to_m, psi_vals_even_1, label=f'Estado par 1 (E = {E_eigenvalues_eV[0]:.2f} eV)', color='blue')

# Estado impar 1
B1 = normalize_odd(E_eigenvalues_J[1])
psi_vals_odd_1 = [psi_odd(x, E_eigenvalues_J[1], B1) for x in x_vals]
plt.plot(x_vals / Å_to_m, psi_vals_odd_1, label=f'Estado impar 1 (E = {E_eigenvalues_eV[1]:.2f} eV)', color='green')

# Estado par 2
A2 = normalize_even(E_eigenvalues_J[2])
psi_vals_even_2 = [psi_even(x, E_eigenvalues_J[2], A2) for x in x_vals]
plt.plot(x_vals / Å_to_m, psi_vals_even_2, label=f'Estado par 2 (E = {E_eigenvalues_eV[2]:.2f} eV)', color='red')

# Configuración de la gráfica
plt.title('Funciones de onda normalizadas para los primeros estados ligados')
plt.xlabel('Posición $x$ (Å)')
plt.ylabel('Función de onda $\psi(x)$')
plt.axvline(-a / Å_to_m, color='black', linestyle='--')
plt.axvline(a / Å_to_m, color='black', linestyle='--')
plt.legend()
plt.grid(True)
plt.show()