nDays=3650
times = np.arange(0, nDays, 1)
TSurf = 0.5*(TSurf_min+TSurf_max) + 0.5 * (TSurf_max-TSurf_min)* np.sin((times - 91.25)/365 * 2* np.pi)
plt.plot(times, TSurf, '-k')
uBt = []
for i in range(nDays):
    uBt.append([[1, TSurf[i]], [3, TDepth]])

dof = mesh.nodeCount()
u = np.zeros((len(times), dof))
u[0, :]= u0

dt = times[1] - times[0]
a = pg.solver.parseArgToArray(alpha,ndof=mesh.cellCount(),mesh=mesh)
A = pg.solver.createStiffnessMatrix(mesh, a)
M = pg.solver.createMassMatrix(mesh)

ut = pg.RVector(dof, 0.0)
rhs = pg.RVector(dof, 0.0)
b = pg.RVector(dof, 0.0)

theta = 0.5

boundUNeu = pg.solver.parseArgToBoundaries([[1, 0], [3, 0]], mesh)

for n in range(1, len(times)):
    b = (M + A * (dt*(theta - 1.0))) * u[n-1] + rhs * (dt*(1.0 - theta)) + rhs * dt * theta

    S = M + A * dt * theta

    boundUdir = pg.solver.parseArgToBoundaries([[1, TSurf[n]], [3, TDepth]], mesh)
    pg.solver.assembleDirichletBC(S, boundUdir, time=times[n], rhs=b)
    pg.solver.assembleNeumannBC(S, boundUNeu, time=times[n])

    solve = pg.LinSolver(S)
    solve.solve(b, ut)

    u[n, :] = ut

for t in range(3650-365,nDays,100):
    fig, ax = plt.subplots(figsize=(8,10))
    patches = pg.mplviewer.drawField(ax, mesh, data=u[t], colorBar=True, showLater=True, omitLines=True, nLevs=64)
    pg.mplviewer.createColorbar(patches, orientation = 'vertical', label='Temperature [°C] at day %s\n Surf. Temp. %2.2f °C' %(t, TSurf[t]))
    pg.mplviewer.setCbarLevels(cbar, cMin=TSurf_min, cMax=TSurf_max, nLevs=16)