import numpy as np
from matplotlib import pyplot as plt
# Se define una lista llamada coefi en la cual cada elemento se le define una posición de la siguiente forma [0,1,2,3,...,n] y el elemento en la posición 0,1,2,3,n es un diccionario
# Nota: Un diccionario es una estructura de datos parecida a una lista con la diferencia en que los elementos no se acceden por una posición si no por una clave, es decir que un elemento está definida por una clave y el elemento
# para acceder a los elementos del diccionario se tiene que especificar la clave de la siguiente forma: diccionario[clave]
# si se quiere agregar un elemento al diccionario entonces se tiene que hacer de la siguiente forma: diccionario{elementos_anteriores,clave:dato_a_agregar}
# dentro de cada elemento de la lista coefi existe un diccionario con la función externa y los coeficientes de la ecuación diferencial; Los coeficientes A(x),B(X),C(X) Y la función externa f(x) tienen las claves 1,2,3 y 4 respectivamente, hay que tener eso encuenta a la hora de querer acceder a los coeficientes de la EDO elegida.
# Finalmente para acceder a los coeficientes de un edo en particular tenemos que hacerlo de la siguiente manera coefi[posición_de_mi_ecuacion_en_la_lista][clave_de_mi_coeficiente_dentro_del_diccionario](parámetro_a_evaluar_en_el_coeficiente_de_la_edo)
# dentro de la funciones lambda se puede modificar las funciones al gusto de quien edite el codigo
# Nota 2: Si se va a ingresar una nueva ecuación diferencial entonces debe de ingresarse un nuevo elemento a la lista coefi
# El elemento que se vaya a agregar debe de ser un diccionario que obligatoriamente debe contener 4 datos
# Esos cuatro datos corresponde a los coeficiente A(X),B(X),C(X) y la función externa f(x)
# Cada dato del diccionario debe de contener una clave de 1 a 4, en donde la clave 1 corresponde al coeficiente A(x), la 2 al coeficiente B(x)
# Así sucesivamente hasta llegar a la clave 4 que corresponde a la función externa f(x)
coefi=[{1:lambda x: 1,2:lambda x: 0,3:lambda x: -0.01,4:lambda x: -0.2},{1:lambda x: 7,2:lambda x: -2,3:lambda x: -1,4:lambda x:-x},{1:lambda x: x**2,2:lambda x: -x,3:lambda x: -1,4:lambda x:2*x}]
# cfr es una lista que contiene en cada posición un diccionario en donde hay dos elementos con una clave, el primer elemento siempre corresponderá a la frontera izquierda y el segundo elemento corresponderá a la frontera derecha
# Nota si se quiere ingresar unas nuevas condiciones de frontera para una ecuación diferencial en específico, se debe ingresar una nueva posición en la lista
# el dato de esa posición debe de tener la misma estructura que la de la primera posición de la lista
# recordar que en una lista los datos se separan por comas
cfr=[{1:(0,40),2:(10,200)},{1:(0,5),2:(20,8)},{1:(1.962,0),2:(9.265,0)}]
funcanali=[lambda x:73.4523*np.exp(0.1*x)-53.4523*np.exp(-0.1*x)+20,lambda x:(-1.783e-4)*np.exp(((1+2*np.sqrt(2))/7)*x)+7.0001783*np.exp(((1-2*np.sqrt(2))/7)*x)+x-2,lambda x:0.03862212*np.float_power(x,(1+np.sqrt(2)))+2.33412631*np.float_power(x,(1-np.sqrt(2)))-x]
# cantidad de puntos, modificar si se necesita
n=10
def Graphs(id,ls,h):
    x=np.arange(cfr[id][1][0],cfr[id][2][0]+h,h) # se llena un vector con todos los valores de x
    y=funcanali[id](x) # se calculan los valores de la función en los x calculados
    plt.plot(x,y,label="Solución analítica")
    plt.plot(x,ls,'r--',label="Solución numérica")
    plt.xlabel("x")
    plt.ylabel("y")
    plt.legend()
    plt.grid()
    plt.title("Solución numérica vs analítica")
    plt.show()
def finitediferencies(id):
    ai, di, bi, ci, yi = [], [], [], [], [] # se definen las listas que van a contener los coeficientes de las ecuaciones lineales
    xi = cfr[id][1][0] #se trae la coordenada x de la frontera izquierda
    h=(cfr[id][2][0]-cfr[id][1][0])/n # se calculan el número de puntos a econctrar la solución
    for i in range(0,n-1): # empieza el bucle iterador para la variable i
        xi=xi+h # se incrementa el valor de xi para evaluar en el punto xi
        for j in range(0,n-1): # empieza el bucle iterador para la variable j
            if i==j: # si en la diagonal principal
                c2 = (-2 / h ** 2) * coefi[id][1](xi) + coefi[id][3](xi) # entonces calcule el coeficiente 2
                di.append(c2) # y haga una lista con los elementos de la diganonal principal
            if i==j-1: # si se está en la diagonal que está por encima de la principal
                c3 = (1 / h ** 2) * coefi[id][1](xi) + (1 / (2 * h)) * coefi[id][2](xi) #calcula el coeficiente 3
                ci.append(c3) #y haga una lista de valores  de la diagonal que está por encima de la diagonal principal
            if i==j+1: #si se está en la diagonal que está por debajo de la principal
                c1 = (1 / h ** 2) * coefi[id][1](xi) - (1 / (2 * h)) * coefi[id][2](xi) #calcule el cofiente 1
                ai.append(c1) # y haga una lista de valores con los elementos de la diagonal que está por debajo de la diagonal pricipal
        ti=coefi[id][4](xi) #calculo de los bi
        if i==0: # se verifica si se está en la ecuación 1
            ti=ti-((1 / h ** 2) * coefi[id][1](xi) - (1 / (2 * h)) * coefi[id][2](xi))*cfr[id][1][1] #si está en la ecuación 1 entonces se despeja el termino para suma cn bi
        elif i==n-2: # se verifica si se está la última ecuación
            ti=ti-((1 / h ** 2) * coefi[id][1](xi) + (1 / (2 * h)) * coefi[id][2](xi))*cfr[id][2][1] # si están en la última ecuación despeja el termino para sumar con bi
        bi.append(ti) #Se ingresa bi a la lista
    matriz = np.diag(di) + np.diag(ci, 1) + np.diag(ai, -1) # se construye la matriz tridiagonal
    sol=np.linalg.solve(matriz,bi) # se soluciona la matriz tridiagonal
    print("Resultados con la aproximación centrada, para un error de h^2 y un tamaño de paso de h="+str(h))
    xi=cfr[id][1][0]
    print("x="+str(xi)+"  y="+str(cfr[id][2][1])) # se muestra el resultado inicial
    ynum=[cfr[id][1][1]]
    for r in sol: # se recorre el vector solución
        xi=xi+h # se calcula el x para el cual se obtiene una solución de y
        print("x="+str(xi)+"  y="+str(r))
        ynum.append(r) # se guardan los valores del vector solución en otra lista para poder gráficar
    print("x=" + str(cfr[id][2][0]) + "  y=" + str(cfr[id][2][1])) #se muestra el resultado final
    ynum.append(cfr[id][2][1])
    Graphs(id,ynum,h) # se llama a una función que grafica las dos soluciones, si no se tiene función analítica para la función escogida, comentar esta línea con #

try:
    print("Actualmente hay "+str(len(coefi))+" ecuaciones diferenciales ingresadas")
    print("Las funciones están en una lista que van desde las posición 0 hasta "+str(len(coefi)-1)+", el id será esa posición en la lista")
    id=int(input("Ingrese el id de la ecuación diferencial para la cual se solucionará sobre un intervalo"))
    finitediferencies(id)
except:
    print("""Ha ocurrido un error al ejecutar el código, revisa si la información que ingresaste es válida""")