#PRODUCTO INTEGRADOR DE APRENDIZAJE
# FÍSICA COMPUTACIONAL (2:00-3:00 PM)
#INTEGRANTES: IRIDIAN ÁVILA, NORALEE COVARRUBIAS, CARLOS RODRÍGUEZ

import numpy as np
import matplotlib.pyplot as plt
from scipy.integrate import romberg

# Parámetros del sistema
m = 10 # Masa de un electrón en kg
k = 1  # Constante elástica en N/m (se ha multiplicado por 1e6 para ajustar la escala)
h = 6  # Constante de Planck reducida en J·s (simbolizada con la letra "h")
x_min = -6  # Valor mínimo de x en metros
x_max = 6  # Valor máximo de x en metros
num_p = 103 # Número de puntos en el espacio de coordenadas (donde el numero tiene que ser impar para que el metodo funcione)
niveles_E = 5  # Número de niveles de energía a calcular

# Método de Runge-Kutta
# derivas la funcion de ondas
def resolver_runge_kutta():
    def potencial(x):
        return k * x**2 / 2

    def derivada_funcion_onda(x, psi, E): # se deriva la funcion de onda
        psi_val, psi_der = psi  # Desempaquetar valores de psi
        return np.array([psi_der, (2 * m / h**2) * (potencial(x) - E) * psi_val])

    x = np.linspace(x_min, x_max, num_p) # se declaran los valores de x
    dx = (x_max - x_min) / (num_p- 1) # se deriva los valores de x
    energias = np.zeros(niveles_E)
    funciones_onda = np.zeros((num_p, niveles_E))

    for i in range(niveles_E):
        E = (i + 0.5) *h* np.sqrt(k / m)
        psi = np.zeros((num_p, 2))
        psi[0, 0] = 1.0

        for j in range(1, num_p):
            k1 = dx * derivada_funcion_onda(x[j-1], psi[j-1], E)
            k2 = dx * derivada_funcion_onda(x[j-1] + dx/2, [psi[j-1][0] + k1[0]/2, psi[j-1][1] + k1[1]/2], E)
            k3 = dx * derivada_funcion_onda(x[j-1] + dx/2, [psi[j-1][0] + k2[0]/2, psi[j-1][1] + k2[1]/2], E)
            k4 = dx * derivada_funcion_onda(x[j-1] + dx, [psi[j-1][0] + k3[0], psi[j-1][1] + k3[1]], E)
            psi[j] = [psi[j-1][0] + (k1[0] + 2*k2[0] + 2*k3[0] + k4[0]) / 6, psi[j-1][1] + (k1[1] + 2*k2[1] + 2*k3[1] + k4[1]) / 6]

        factor_normalizacion = np.sqrt(np.trapz(np.abs(psi[:, 0])**2, x=x))
        psi[:, 0] /= factor_normalizacion

        energias[i] = E
        funciones_onda[:, i] = psi[:, 0]

    return x, energias, funciones_onda

# Resolución utilizando el método de Runge-Kutta
#se calcula la densidad
x_rk, energias_rk, funciones_onda_rk = resolver_runge_kutta()

# Método de Romberg
def calcular_densidad(x, funcion_onda):
    densidad = np.abs(funcion_onda)**2
    integral = romberg(lambda x: np.abs(np.interp(x, x_rk, densidad))**2, x_min, x_max, divmax=20)
    return densidad / integral

# Cálculo de las densidades de probabilidad
densidades_rk = []
for i in range(niveles_E):
    densidad = calcular_densidad(x_rk, funciones_onda_rk[:, i])
    densidad_normalizada = densidad / np.max(densidad)
    densidades_rk.append(densidad_normalizada)

# Método del trapecio para integración numérica
def metodo_trapecio(x, y):
    dx = x[1] - x[0]
    integral = 0.5 * dx * (y[0] + 2 * np.sum(y[1:-1]) + y[-1])
    return integral

# Calcular valores esperados de posición y momento
def calcular_valores_esperados(x, funcion_onda):
    densidad = np.abs(funcion_onda)**2
    x_esperado = metodo_trapecio(x, x * densidad)
    p_esperado = metodo_trapecio(x, np.conj(funcion_onda) * (h/ 1) * np.gradient(funcion_onda, x))
    return x_esperado, p_esperado

# Calcular valores esperados para cada nivel de energía
valores_esperados_x = []
valores_esperados_p = []
for i in range(niveles_E):
    x_esperado, p_esperado = calcular_valores_esperados(x_rk, funciones_onda_rk[:, i])
    valores_esperados_x.append(x_esperado)
    valores_esperados_p.append(p_esperado)

# Mostrar resultados
print("Energías permitidas (Método de Runge-Kutta):")

# Imprimir los valores de energía
for i in range(niveles_E):
    print(f"Energía del estado {i}: {energias_rk[i]}",) #estado de Energia (E0)

print("Valores esperados de posición:")
for i, x_esperado in enumerate(valores_esperados_x): # estado de Energia de (E1)
    print(f"Estado {i}: {x_esperado}")

print("Valores esperados de momento:")
for i, p_esperado in enumerate(valores_esperados_p): #estado de Energia de (E2)
    print(f"Estado {i}: {p_esperado}")
# Visualización de los resultados
plt.figure()
# Gráficas de las funciones de onda para los primeros tres niveles de energía
for i in range(niveles_E):
    plt.plot(x_rk,funciones_onda_rk[:, i],label=f'Estado {i}')   # se imprimen en un ciclo para todos los valores de E

plt.xlabel('Posición',color="purple")
plt.ylabel('Función de onda',color="purple")
plt.title('Función de onda(Método de Runge-Kutta)',color="blue")
plt.legend()
plt.tight_layout()
plt.show()
#imprimimos otra grafica para los valores de el metodo de romberg
plt.figure()
# Visualización de las densidades de probabilidad
for i in range(niveles_E):
    plt.plot(x_rk, densidades_rk[i], label=f'Estado {i}')

plt.xlabel('Posición',color="purple")
plt.ylabel('Densidad de probabilidad (normalizada)',color="purple")
plt.title('Densidades de probabilidad (Método de Romberg)',color="blue")
plt.legend()
plt.show()