# -*- coding: utf-8 -*-
import matplotlib
matplotlib.use('TkAgg')
from scipy.integrate import odeint
from scipy.optimize import fsolve
from pylab import figure, plot, show, xlabel, ylabel ,draw, ion,ylim,ioff, pause
import numpy as np

def FH(u,t,n,beta):
    '''
    równania na rozchodzenie się ciepła:
    
     u_t = beta * u_xx
    
     warunki pocz. u[x,0]=sin(2*np.pi*x) + 2*x + 1
     warunki brzegowe:
      z lewej: u[0,t] = g1(t) = 1
      z prawej: u_x[1,t] = g2(t) = 2

    n - ilość węzłów
    
    zamieniamy równanie różniczkowe cząstkowe na układ równań różniczkowych poprzez 
    dyskretyzację przestrzeni.
    Korzystamy z przybliżenia drugiej pochodnej:
    d2u/dx2 = [u(x+h, t) -2u(x,t) +u(x-h,t) ]/h^2
    
    aby przeprowadzić tą dyskretyzację wybieramy siatkę n punktow równomiernie 
    rozłożonych na odcinku 0,1
    '''
    # bierzemy siatkę o n węzłach
    h = 1./n # to jest odstęp między węzłami
    h2= h**2
    # wytwarzamy miejsce na równania (dostaniemy jedno równanie na węzeł):
    du = np.zeros(n)
    A= np.zeros((n,n))
    for wiersz in range(n):
        for kol in range(n):
            if wiersz==kol:
                A[wiersz,kol] =-2
            elif kol+1==wiersz:
                A[wiersz,kol] =1
            elif kol-1==wiersz:
                A[wiersz,kol] = 1
    A[-1,-1]+=1           
    A *= beta/h2
    
    b = np.zeros(n)
    b[0] = 1
    b[-1] = 2*h
    b *= beta/h2
    du = np.dot(A,u) + b

    return du

## całkowanie
m = 200 # ilość punktów w czasie w rozwiązaniu
time = np.linspace(0.0,2000.0,m) 
n=100#ilość węzłów
x = np.linspace(0,1,n+2)
# warunki początkowe

uinit = np.sin(2*np.pi*x) + 2*x + 1
uinit = uinit[1:-1]
beta = 1e-4
u=np.zeros((m,n+2))
u[:,1:n+1] = odeint(FH, uinit, time, args =(n,beta))
u[:,0] = 1 # lewy war. brzegowy
u[:,n+1] = u[:,n]+2*1./n 
ion()
l,= plot(x,u[0,:])
ylim([0,3.1])
for t in range(0, u.shape[0]):
    l.set_ydata(u[t,:]) # y[:,0] pierwsza kolumna y - zawiera zmienną u
    draw()
    pause(0.1)

xlabel('x') 
ylabel('u') 



show()