from mpl_toolkits.mplot3d import Axes3D
import matplotlib.pyplot as plt
from matplotlib import cm
import numpy as np
import os
import contorno
from constantes import INTERVALOS, PASSOS, TAMANHO_BARRA, DELTA_T, DELTA_X

z_temp = contorno.p_3
TAMANHO_BARRA = 2

x = np.linspace(0.0, TAMANHO_BARRA, INTERVALOS+1)
y = np.linspace(0.0, DELTA_T, PASSOS+1)

z = []
for k in range(PASSOS+1):
    z_k = np.copy(z_temp)
    z.append(z_k)
    for i in range(1, INTERVALOS):
        z_temp[i] = z_k[i] + (DELTA_T/(DELTA_X**2)) * (z_k[i+1]-2*z_k[i]+z_k[i-1])

z = np.asarray(z)
x, y = np.meshgrid(x, y)

fig = plt.figure()
ax = fig.gca(projection='3d')
surf = ax.plot_surface(x, y, z, cmap=cm.coolwarm, antialiased=False)

ax.set_xlabel('x')
ax.set_ylabel('t')
ax.set_zlabel('T(x,t)')

fig.colorbar(surf, shrink=0.5, aspect=5)

plt.show()

[(12, 'numpy.linspace', 'np.linspace', (['(0.0)', 'TAMANHO_BARRA', '(INTERVALOS + 1)'], {}), True, 'import numpy as np\n'), (13, 'numpy.linspace', 'np.linspace', (['(0.0)', 'DELTA_T', '(PASSOS + 1)'], {}), True, 'import numpy as np\n'), (22, 'numpy.asarray', 'np.asarray', (['z'], {}), True, 'import numpy as np\n'), (23, 'numpy.meshgrid', 'np.meshgrid', (['x', 'y'], {}), True, 'import numpy as np\n'), (25, 'matplotlib.pyplot.figure', 'plt.figure', ([], {}), True, 'import matplotlib.pyplot as plt\n'), (35, 'matplotlib.pyplot.show', 'plt.show', ([], {}), True, 'import matplotlib.pyplot as plt\n'), (17, 'numpy.copy', 'np.copy', (['z_temp'], {}), True, 'import numpy as np\n')]