#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Tue Apr  8 16:32:02 2025

@author: christian

Finite Difference Method for 1D Heat Equation
This code simulates the temperature distribution in a 1D rod over time using the finite difference method (FDM).
The rod is divided into nodes, and the temperature at each node is calculated iteratively.
The code uses the Forward Time Centered Space (FTCS) scheme to solve the heat equation.
Input:
- alpha: thermal diffusivity (m^2/s)
- T_i: initial temperature (°C)
- dt: time step (s)
- T_0: boundary temperature (°C)
- F_0: Fourier number (dimensionless)
- thickness: length of the rod (m)
- t_final: final time (s)
Output:
- Temp: 2D array of temperatures at each node over time
- Plot: temperature distribution at all nodes over time
- Animation: visualizes the temperature distribution over time
"""

import numpy as np
import math
from matplotlib.animation import FuncAnimation
import matplotlib.pyplot as plt

# Input parameters
alpha = 1.5e-6 # m^2/s
T_i = 85 # degC
dt = 300 # s
T_0 = 20 # degC
F_0 = 0.5
thickness = 0.12 # m
t_final = 45 * 60 # s

dx = np.sqrt(alpha*dt/F_0)
t = np.arange(0,t_final+1,dt)

no_internal_nodes = int(thickness/dx)
no_phantom_nodes = math.ceil((len(t) - no_internal_nodes)/2)
total_nodes = 1+no_internal_nodes+no_phantom_nodes

Temp = np.zeros((len(t),total_nodes))

# Imposed temperature T_0
Temp[:,0] = T_0

# Initialize temperatures of upper triangle
for i in range(len(t)):
    for j in range(i,total_nodes-1):
        Temp[i,j+1]=T_i

# Calculates averages temperatures of lower triangle
for i in range(1,len(t)):
    for j in range(1,total_nodes-1):
        Temp[i,j]=round((Temp[i-1,j-1]+Temp[i-1,j+1])/2, 3)

# Remove phantom nodes
Temp = np.delete(Temp, np.s_[-no_phantom_nodes:], axis=1)
print(Temp)


# Plotting time vs temperature at all nodes
plt.figure(figsize=(10, 6))
plt.title('Temperature at all nodes over time')
plt.xlabel('Time (s)')
plt.ylabel('Temperature (°C)')
plt.xlim(0, t_final)
plt.ylim(0, T_i)
plt.grid()
for j in range(Temp.shape[1]):
    plt.plot(t, Temp[:, j], label=f'Node {j}')
plt.legend(loc='upper right')
plt.show()


# Animation
fig, ax = plt.subplots(figsize=(10, 6))
ax.set_title('Temperature distribution over time')
ax.set_xlabel('Distance (m)')
ax.set_ylabel('Temperature (°C)')
ax.set_xlim(0, thickness)
ax.set_ylim(0, T_i)
ax.grid()

line, = ax.plot([], [], label='Time = 0 s')
legend = ax.legend(loc='upper right')

def init():
    line.set_data([], [])
    return line,

def update(frame):
    time = frame * dt
    line.set_data(np.linspace(0, thickness, Temp.shape[1]), Temp[frame])
    line.set_label(f'Time = {time} s')
    legend = ax.legend(loc='upper right')
    return line, legend

ani = FuncAnimation(fig, update, frames=len(t), init_func=init, blit=True, interval=100)

plt.show()