# main.py

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import math
import sys

import Bar
from case1 import solve, get_analytical_datapoints, get_analytical_heat_convection
from common import calc_extrapolated, calc_error, calc_beta, fdm_heat_extraction, fem_heat_extraction, heat_extraction


bar1 = Bar.Bar(
    radius=0.1,
    k=0.5,
    h=0.25
)
bar2 = Bar.Bar(
    radius=0.1,
    k=2,
    h=0.25
)

data_dict = {}
data_dict["k_ratios"] = np.array([1/16, 1/8, 1/4, 1/2, 1, 2, 4, 8, 16])
data_dict["x_bar_values"] = np.array([1/2, 2/np.pi])
data_dict["length"] = 1

def analytical():
    '''Analytical solutions with varying k_ratio'''
    print("Section 1: Analytical")
    file_path = "images/analytical/"

    l = data_dict["length"]
    
    # Calculate the results for every k ratio for x_bar_1
    x_bar_1 = data_dict["x_bar_values"][0]

    data_dict["analytical"] = {}
    data_dict["analytical"]["x_bar_1"] = {}
    data_dict["analytical"]["x_bar_1"]["x_values_fine"] = np.linspace(0, l, 50)
    data_dict["analytical"]["x_bar_1"]["x_values_fine"] = np.sort(np.append(data_dict["analytical"]["x_bar_1"]["x_values_fine"], x_bar_1))
    data_dict["analytical"]["x_bar_1"]["x_values_course"] = np.linspace(0, l, 9)
    data_dict["analytical"]["x_bar_1"]["x_values_course"] = np.sort(np.append(data_dict["analytical"]["x_bar_1"]["x_values_course"], x_bar_1))

    results = []
    results_course = []
    for k_ratio in data_dict["k_ratios"]:
        Bar.set_k_ratio(k_ratio, bar1, bar2) # Set the k value of bar2 = k*bar1
        result = get_analytical_datapoints(data_dict["analytical"]["x_bar_1"]["x_values_fine"], 
                                                  bar1=bar1, 
                                                  bar2=bar2, 
                                                  x_bar=x_bar_1, 
                                                  l=l)
        result_course = get_analytical_datapoints(data_dict["analytical"]["x_bar_1"]["x_values_course"], 
                                                  bar1=bar1, 
                                                  bar2=bar2, 
                                                  x_bar=x_bar_1, 
                                                  l=l)
        results.append(result)
        results_course.append(result_course)

    data_dict["analytical"]["x_bar_1"]["temp_results"] = np.array(results)

    df_x_bar_1 = pd.DataFrame(results_course, index=[f'k2/k1 = {a}' for a in data_dict["k_ratios"]], columns=[f'x = {x:.3f}' for x in data_dict["analytical"]["x_bar_1"]["x_values_course"]])

    print("Analytical Temperature Results x_bar = 0.5:")
    print(df_x_bar_1.to_string())
    print("\n" * 2)

    # Plotting x_bar_1
    plt.figure(figsize=(10, 6))
    for idx, k_ratio in enumerate(data_dict["k_ratios"]):
        plt.plot(data_dict["analytical"]["x_bar_1"]["x_values_fine"], data_dict["analytical"]["x_bar_1"]["temp_results"][idx], label=f'k2/k1 = {k_ratio}')
    plt.axvline(x=x_bar_1, color='r', linestyle='--', label='x_bar')
    plt.xlabel('Position (cm)')
    plt.ylabel('Temperature (°C)')
    plt.legend()
    plt.grid(True)
    plt.savefig(f'{file_path}x_bar_1/temperature_distribution.png', dpi=300)
    #plt.show()
    plt.clf()

    # Calculate the temperature results for every k ratio for x_bar_1
    x_bar_2 = data_dict["x_bar_values"][1]

    data_dict["analytical"]["x_bar_2"] = {}
    data_dict["analytical"]["x_bar_2"]["x_values_fine"] = np.linspace(0, l, 50)
    data_dict["analytical"]["x_bar_2"]["x_values_fine"] = np.sort(np.append(data_dict["analytical"]["x_bar_2"]["x_values_fine"], x_bar_2))
    data_dict["analytical"]["x_bar_2"]["x_values_course"] = np.linspace(0, l, 9)
    data_dict["analytical"]["x_bar_2"]["x_values_course"] = np.sort(np.append(data_dict["analytical"]["x_bar_2"]["x_values_course"], x_bar_2))

    results = []
    results_course = []
    for k_ratio in data_dict["k_ratios"]:
        Bar.set_k_ratio(k_ratio, bar1, bar2)
        result = get_analytical_datapoints(data_dict["analytical"]["x_bar_2"]["x_values_fine"], 
                                                  bar1=bar1, 
                                                  bar2=bar2, 
                                                  x_bar=x_bar_2, 
                                                  l=l)
        result_course = get_analytical_datapoints(data_dict["analytical"]["x_bar_2"]["x_values_course"], 
                                                  bar1=bar1, 
                                                  bar2=bar2, 
                                                  x_bar=x_bar_2, 
                                                  l=l)
        results.append(result)
        results_course.append(result_course)

    data_dict["analytical"]["x_bar_2"]["temp_results"] = np.array(results)

    df_x_bar_2 = pd.DataFrame(results_course, index=[f'k2/k1 = {a}' for a in data_dict["k_ratios"]], columns=[f'x = {x:.3f}' for x in data_dict["analytical"]["x_bar_2"]["x_values_course"]])

    print("Analytical Temperature Results x_bar = 2/pi:")
    print(df_x_bar_2.to_string())
    print("\n" * 2)
    
    # Plotting x_bar_2
    plt.figure(figsize=(10, 6))
    for idx, k_ratio in enumerate(data_dict["k_ratios"]):
        plt.plot(data_dict["analytical"]["x_bar_2"]["x_values_fine"], data_dict["analytical"]["x_bar_2"]["temp_results"][idx], label=f'k2/k1 = {k_ratio}')
    plt.axvline(x=x_bar_2, color='r', linestyle='--', label='x_bar')
    plt.xlabel('Position (cm)')
    plt.ylabel('Temperature (°C)')
    plt.legend()
    plt.grid(True)
    plt.savefig(f'{file_path}x_bar_2/temperature_distribution.png', dpi=300)
    #plt.show()
    plt.clf()



    # Calculate the analytical heat convection leaving the rod at x=l for varying k ratio

    results_1 = []
    results_2 = []
    for k_ratio in data_dict["k_ratios"]:
        Bar.set_k_ratio(k_ratio, bar1, bar2)
        result_1 = get_analytical_heat_convection(x=l, 
                                                      bar1=bar1, 
                                                      bar2=bar2, 
                                                      x_bar=x_bar_1, 
                                                      l=l)
        result_2 = get_analytical_heat_convection(x=l, 
                                                      bar1=bar1, 
                                                      bar2=bar2, 
                                                      x_bar=x_bar_2, 
                                                      l=l)
        results_1.append([result_1])
        results_2.append([result_2])

    data_dict["analytical"]["x_bar_1"]["heat_results"] = np.array(results_1)
    data_dict["analytical"]["x_bar_2"]["heat_results"] = np.array(results_2)

    df_x_bar_1 = pd.DataFrame(results_1, index=[f'k2/k1 = {a}' for a in data_dict["k_ratios"]], columns=["Heat Convection"])
    df_x_bar_2 = pd.DataFrame(results_2, index=[f'k2/k1 = {a}' for a in data_dict["k_ratios"]], columns=["Heat Convection"])

    print("Analytical Heat Convection Results x_bar = 0.5:")
    print(df_x_bar_1.to_string())
    print("\n" * 2)

    print("Analytical Heat Convection Results x_bar = 2/pi:")
    print(df_x_bar_2.to_string())
    print("\n" * 2)

    # Plotting x_bar_1 heat convection
    plt.figure(figsize=(10, 6))
    plt.plot(data_dict["k_ratios"], data_dict["analytical"]["x_bar_1"]["heat_results"], label=f'Heat Convection')
    plt.xlabel('k2/k1')
    plt.ylabel('Q_Dot (W)')
    plt.legend()
    plt.grid(True)
    plt.savefig(f'{file_path}x_bar_1/heat_convection.png', dpi=300)
    #plt.show()
    plt.clf()

    # Plotting x_bar_2 heat convection
    plt.figure(figsize=(10, 6))
    plt.plot(data_dict["k_ratios"], data_dict["analytical"]["x_bar_2"]["heat_results"], label=f'Heat Convection')
    plt.xlabel('k2/k1')
    plt.ylabel('Q_Dot (W)')
    plt.legend()
    plt.grid(True)
    plt.savefig(f'{file_path}x_bar_2/heat_convection.png', dpi=300)
    #plt.show()
    plt.clf()


def section_2(method="FDM"):
    '''FDM and FEM with varying k_ratio'''
    print(f"Section 2: {method} with Varying k_ratio")
    file_path = f"images/{method}/"

    l = data_dict["length"]

    data_dict["section2"] = {}
    data_dict["section2"]["num_sections"] = 8


    # Calculate the results for every k ratio for x_bar_1
    x_bar_1 = data_dict["x_bar_values"][0]

    # For every k_ratio, calculate the FDM and FEM temperature results
    fdm_results = []
    fem_results = []
    x_fdm = []
    x_fem = []
    for k_ratio in data_dict["k_ratios"]:
        Bar.set_k_ratio(k_ratio, bar1, bar2) # Set the k value of bar2 = k*bar1
        fdm_temps, x_fdm = solve(bar1=bar1,
                                       bar2=bar2,
                                       num_sections=data_dict["section2"]["num_sections"],
                                       x_bar=x_bar_1,
                                       l=l,
                                       method=method,
                                       order=2)
        fem_temps, x_fem = solve(bar1=bar1,
                                       bar2=bar2, 
                                       num_sections=data_dict["section2"]["num_sections"],
                                       x_bar=x_bar_1,
                                       l=l,
                                       method=method,
                                       order=4)
        
        fdm_results.append(fdm_temps)
        fem_results.append(fem_temps)

    data_dict["section2"]["x_bar_1"] = {}
    data_dict["section2"]["x_bar_1"]["x_values"] = x_fdm # fdm and fem use same x_values
    data_dict["section2"]["x_bar_1"]["fdm_results"] = np.array(fdm_results)
    data_dict["section2"]["x_bar_1"]["fem_results"] = np.array(fem_results)

    df_fdm = pd.DataFrame(fdm_results, index=[f'k2/k1 = {a}' for a in data_dict["k_ratios"]], columns=[f'x = {x:.3f}' for x in x_fdm])
    df_fem = pd.DataFrame(fem_results, index=[f'k2/k1 = {a}' for a in data_dict["k_ratios"]], columns=[f'x = {x:.3f}' for x in x_fem])


    foo = "2nd Order" if method=="FDM" else "p=1"
    bar = "4th Order" if method=="FDM" else "p=2"

    print(f"{method} {foo} Temperature Results x_bar = 0.5:")
    print(df_fdm.to_string())
    print("\n" * 2)
    print(f"{method} {bar} Temperature Results x_bar = 0.5:")
    print(df_fem.to_string())
    print("\n" * 2)

    # Now that the data is gathered for FDM and FEM, plot it
    # Plotting x_bar_1, FDM
    plt.figure(figsize=(10, 6))
    for idx, k_ratio in enumerate(data_dict["k_ratios"]):
        plt.plot(data_dict["section2"]["x_bar_1"]["x_values"], data_dict["section2"]["x_bar_1"]["fdm_results"][idx], label=f'k2/k1 = {k_ratio}')
    plt.axvline(x=x_bar_1, color='r', linestyle='--', label='x_bar')
    plt.xlabel('Position (cm)')
    plt.ylabel('Temperature (°C)')
    plt.legend()
    plt.grid(True)
    plt.savefig(f'{file_path}x_bar_1/temperature_distribution_1.png', dpi=300)
    #plt.show()
    plt.clf()

    # Plotting x_bar_1
    plt.figure(figsize=(10, 6))
    for idx, k_ratio in enumerate(data_dict["k_ratios"]):
        plt.plot(data_dict["section2"]["x_bar_1"]["x_values"], data_dict["section2"]["x_bar_1"]["fem_results"][idx], label=f'k2/k1 = {k_ratio}')
    plt.axvline(x=x_bar_1, color='r', linestyle='--', label='x_bar')
    plt.xlabel('Position (cm)')
    plt.ylabel('Temperature (°C)')
    plt.legend()
    plt.grid(True)
    plt.savefig(f'{file_path}x_bar_1/temperature_distribution_2.png', dpi=300)
    #plt.show()
    plt.clf()


    # Calculate the results for every k ratio for x_bar_2
    x_bar_2 = data_dict["x_bar_values"][1]

    # For every k_ratio, calculate the FDM and FEM temperature results
    fdm_results = []
    fem_results = []
    x_fdm = []
    x_fem = []
    for k_ratio in data_dict["k_ratios"]:
        Bar.set_k_ratio(k_ratio, bar1, bar2) # Set the k value of bar2 = k*bar1
        fdm_temps, x_fdm = solve(bar1=bar1,
                                       bar2=bar2,
                                       num_sections=data_dict["section2"]["num_sections"],
                                       x_bar=x_bar_2,
                                       l=l,
                                       method=method, order=2)
        fem_temps, x_fem = solve(bar1=bar1,
                                       bar2=bar2, 
                                       num_sections=data_dict["section2"]["num_sections"],
                                       x_bar=x_bar_2,
                                       l=l,
                                       method=method, order=4)
        
        fdm_results.append(fdm_temps)
        fem_results.append(fem_temps)

    data_dict["section2"]["x_bar_2"] = {}
    data_dict["section2"]["x_bar_2"]["x_values"] = x_fdm # fdm and fem use same x_values
    data_dict["section2"]["x_bar_2"]["fdm_results"] = np.array(fdm_results)
    data_dict["section2"]["x_bar_2"]["fem_results"] = np.array(fem_results)

    df_fdm = pd.DataFrame(fdm_results, index=[f'k2/k1 = {a}' for a in data_dict["k_ratios"]], columns=[f'x = {x:.3f}' for x in x_fdm])
    df_fem = pd.DataFrame(fem_results, index=[f'k2/k1 = {a}' for a in data_dict["k_ratios"]], columns=[f'x = {x:.3f}' for x in x_fem])

    print(f"{method} {foo} Temperature Results x_bar = 2/pi:")
    print(df_fdm.to_string())
    print("\n" * 2)
    print(f"{method} {bar} Temperature Results x_bar = 2/pi:")
    print(df_fem.to_string())
    print("\n" * 2)

    # Plotting x_bar_2, FDM
    plt.figure(figsize=(10, 6))
    for idx, k_ratio in enumerate(data_dict["k_ratios"]):
        plt.plot(data_dict["section2"]["x_bar_2"]["x_values"], data_dict["section2"]["x_bar_2"]["fdm_results"][idx], label=f'k2/k1 = {k_ratio}')
    plt.axvline(x=x_bar_2, color='r', linestyle='--', label='x_bar')
    plt.xlabel('Position (cm)')
    plt.ylabel('Temperature (°C)')
    plt.legend()
    plt.grid(True)
    plt.savefig(f'{file_path}x_bar_2/temperature_distribution_1.png', dpi=300)
    #plt.show()
    plt.clf()

    # Plotting x_bar_2, FEM
    plt.figure(figsize=(10, 6))
    for idx, k_ratio in enumerate(data_dict["k_ratios"]):
        plt.plot(data_dict["section2"]["x_bar_2"]["x_values"], data_dict["section2"]["x_bar_2"]["fem_results"][idx], label=f'k2/k1 = {k_ratio}')
    plt.axvline(x=x_bar_2, color='r', linestyle='--', label='x_bar')
    plt.xlabel('Position (cm)')
    plt.ylabel('Temperature (°C)')
    plt.legend()
    plt.grid(True)
    plt.savefig(f'{file_path}x_bar_2/temperature_distribution_2.png', dpi=300)
    #plt.show()
    plt.clf()



    # After calculating temperature values, extract the heat convection at x=l

    # x_bar_1
    # for every k ratio, calculate the heat convection from fdm and fem temp results
    fdm_heat_results = []
    fem_heat_results = []
    for i, k_ratio in enumerate(data_dict["k_ratios"]):
        Bar.set_k_ratio(k_ratio, bar1, bar2) # Set the k value of bar2 = k*bar1
        dx = data_dict["section2"]["x_bar_1"]["x_values"][-1] - data_dict["section2"]["x_bar_1"]["x_values"][-2]
        (fdm_t0, fdm_t1) = data_dict["section2"]["x_bar_1"]["fdm_results"][i][-2:]
        fdm_heat_result = heat_extraction(fdm_t0, fdm_t1, dx, bar2, order=2, method=method)
        fdm_heat_results.append(fdm_heat_result)

        (fem_t0, fem_t1) = data_dict["section2"]["x_bar_1"]["fem_results"][i][-2:]
        fem_heat_result = heat_extraction(fem_t0, fem_t1, dx, bar2, order=2, method=method)
        fem_heat_results.append(fem_heat_result)
        

    data_dict["section2"]["x_bar_1"]["fdm_heat_results"] = np.array(fdm_heat_results)
    data_dict["section2"]["x_bar_1"]["fem_heat_results"] = np.array(fem_heat_results)

    # Create DataFrame with two columns for FDM and FEM heat results
    df_heat = pd.DataFrame(
        {
            f"{foo} Heat Convection": fdm_heat_results,
            f"{bar} Heat Convection": fem_heat_results
        },
        index=[f'k2/k1 = {a}' for a in data_dict["k_ratios"]]
    )

    # Print the DataFrame with labeled columns and indices
    print(f"{method} Heat Convection Results for x_bar = 0.5:")
    print(df_heat.to_string())
    print("\n")

    # x_bar_2
    # for every k ratio, calculate the heat convection from fdm and fem temp results
    fdm_heat_results = []
    fem_heat_results = []
    for i, k_ratio in enumerate(data_dict["k_ratios"]):
        Bar.set_k_ratio(k_ratio, bar1, bar2) # Set the k value of bar2 = k*bar1
        dx = data_dict["section2"]["x_bar_2"]["x_values"][-1] - data_dict["section2"]["x_bar_2"]["x_values"][-2]
        (fdm_t0, fdm_t1) = data_dict["section2"]["x_bar_2"]["fdm_results"][i][-2:]
        fdm_heat_result = heat_extraction(fdm_t0, fdm_t1, dx, bar2, order=2, method=method)
        fdm_heat_results.append(fdm_heat_result)

        (fem_t0, fem_t1) = data_dict["section2"]["x_bar_2"]["fem_results"][i][-2:]
        fem_heat_result = heat_extraction(fem_t0, fem_t1, dx, bar2, order=4, method=method)
        fem_heat_results.append(fem_heat_result)
        

    data_dict["section2"]["x_bar_2"]["fdm_heat_results"] = np.array(fdm_heat_results)
    data_dict["section2"]["x_bar_2"]["fem_heat_results"] = np.array(fem_heat_results)

    # Create DataFrame with two columns for FDM and FEM heat results
    df_heat = pd.DataFrame(
        {
            f"{foo} Heat Convection": fdm_heat_results,
            f"{bar} Heat Convection": fem_heat_results
        },
        index=[f'k2/k1 = {a}' for a in data_dict["k_ratios"]]
    )

    # Print the DataFrame with labeled columns and indices
    print(f"{method} Heat Convection Results for x_bar = 2/pi:")
    print(df_heat.to_string())
    print("\n")

    # Plotting x_bar_1 heat convection
    plt.figure(figsize=(10, 6))
    plt.plot(data_dict["k_ratios"], data_dict["section2"]["x_bar_1"]["fdm_heat_results"], label=f'2nd Order')
    plt.plot(data_dict["k_ratios"], data_dict["section2"]["x_bar_1"]["fem_heat_results"], label=f'4th Order')
    #plt.plot(data_dict["k_ratios"], data_dict["analytical"]["x_bar_1"]["heat_results"], label=f'Analytical Heat Convection')
    plt.xlabel('k2/k1')
    plt.ylabel('Q_Dot (W)')
    plt.legend()
    plt.grid(True)
    plt.savefig(f'{file_path}x_bar_1/heat_convection.png', dpi=300)
    #plt.show()
    plt.clf()

    # Plotting x_bar_2 heat convection
    plt.figure(figsize=(10, 6))
    plt.plot(data_dict["k_ratios"], data_dict["section2"]["x_bar_2"]["fdm_heat_results"], label=f'2nd Order')
    plt.plot(data_dict["k_ratios"], data_dict["section2"]["x_bar_2"]["fem_heat_results"], label=f'4th Order')
    #plt.plot(data_dict["k_ratios"], data_dict["analytical"]["x_bar_2"]["heat_results"], label=f'Analytical Heat Convection')
    plt.xlabel('k2/k1')
    plt.ylabel('Q_Dot (W)')
    plt.legend()
    plt.grid(True)
    plt.savefig(f'{file_path}x_bar_2/heat_convection.png', dpi=300)
    #plt.show()
    plt.clf()


def temp_convergence():
    '''Convergence of FDM and FEM of interface temperature'''
    print("Section 3 Temp: Convergence of FDM and FEM Temp with Varying k Ratio")
    file_path = "images/convergence/"
    l = data_dict["length"]
    
    data_dict["convergence"] = {}
    data_dict["convergence"]["num_sections"] = [4, 8, 16, 32, 64, 128, 256, 512]

    for x_bar_idx, x_bar in enumerate(data_dict["x_bar_values"]):
        num_k_ratios = len(data_dict["k_ratios"])
        num_rows = math.ceil(num_k_ratios / 2)

        fig_temp, axs_temp = plt.subplots(num_rows, 2, figsize=(15, 5 * num_rows))

        axs_temp = axs_temp.flatten()

        for idx, k_ratio in enumerate(data_dict["k_ratios"]):
            Bar.set_k_ratio(k_ratio, bar1, bar2)

            # Analytical temperature at the interface
            analy_interface_temp = get_analytical_datapoints([x_bar], bar1, bar2, x_bar, l)[0]
            convergence_data = {
                "fdm_2_values": [], "fdm_4_values": [], "fem_2_values": [], "fem_4_values": [],
                "fdm_2_errors": [], "fdm_4_errors": [], "fem_2_errors": [], "fem_4_errors": [],
                "fdm_2_extrap": [], "fdm_4_extrap": [], "fem_2_extrap": [], "fem_4_extrap": [],
                "fdm_2_b": [], "fdm_4_b": [], "fem_2_b": [], "fem_4_b": []
            }

            # Calculate temperature values and errors for different section counts
            for num_sections in data_dict["convergence"]["num_sections"]:
                temp_results = {
                    "FDM_2": solve(bar1, bar2, num_sections, x_bar, l, method="FDM", order=2)[0],
                    "FDM_4": solve(bar1, bar2, num_sections, x_bar, l, method="FDM", order=4)[0],
                    "FEM_2": solve(bar1, bar2, num_sections, x_bar, l, method="FEM", order=2)[0],
                    "FEM_4": solve(bar1, bar2, num_sections, x_bar, l, method="FEM", order=4)[0]
                }

                interface_idx = num_sections // 2
                for method, temps in temp_results.items():
                    interface_temp = temps[interface_idx]
                    convergence_data[f"{method.lower()}_values"].append(interface_temp)
                    error = calc_error(analy_interface_temp, interface_temp)
                    convergence_data[f"{method.lower()}_errors"].append(error)

            # Calculate convergence rates `B` for each method and apply extrapolation from the third entry onward
            for method in ["fdm_2", "fdm_4", "fem_2", "fem_4"]:
                values = convergence_data[f"{method}_values"]
                errors = convergence_data[f"{method}_errors"]

                # Initialize `extrap` with NaNs for the first two entries
                extrapolated_values = [float('NaN')] * 2
                b_values = [float('NaN')]  # Initialize `B` list with NaN for the first entry

                # Compute extrapolated values starting from the third entry if possible
                for i in range(2, len(values)):
                    if i >= 2:
                        extrap_value = calc_extrapolated(values[i - 2], values[i - 1], values[i])
                        extrapolated_values.append(extrap_value)

                        # Calculate extrapolated error for the current value
                        extrap_error = calc_error(extrap_value, values[i])
                        errors[i] = extrap_error
                    else:
                        extrapolated_values.append(float('NaN'))

                # Calculate convergence rates `B` for each consecutive pair
                for i in range(1, len(values)):
                    beta = calc_beta(analy_interface_temp, values[i], values[i - 1],
                                    data_dict["convergence"]["num_sections"][i],
                                    data_dict["convergence"]["num_sections"][i - 1])
                    b_values.append(beta)
                
                # Store the calculated `B` and extrapolated values in the convergence data
                convergence_data[f"{method}_b"] = b_values
                convergence_data[f"{method}_extrap"] = extrapolated_values



            # Prepare data for FDM and FEM tables
            rows = data_dict["convergence"]["num_sections"]
            fdm_table_data = []
            fem_table_data = []

            for i, num_sections in enumerate(rows):
                # Append data for FDM
                fdm_table_data.append([
                    analy_interface_temp,
                    convergence_data["fdm_2_values"][i],
                    convergence_data["fdm_2_extrap"][i],
                    convergence_data["fdm_2_errors"][i],
                    convergence_data["fdm_2_b"][i],
                    calc_error(convergence_data["fdm_2_extrap"][i], convergence_data["fdm_2_values"][i]) if i >= 2 else float('NaN'),
                    calc_beta(convergence_data["fdm_2_extrap"][i], convergence_data["fdm_2_values"][i], convergence_data["fdm_2_values"][i-1], num_sections, rows[i-1]) if i >= 3 else float('NaN'),

                    convergence_data["fdm_4_values"][i],
                    convergence_data["fdm_4_extrap"][i],
                    convergence_data["fdm_4_errors"][i],
                    convergence_data["fdm_4_b"][i],
                    calc_error(convergence_data["fdm_4_extrap"][i], convergence_data["fdm_4_values"][i]) if i >= 2 else float('NaN'),
                    calc_beta(convergence_data["fdm_4_extrap"][i], convergence_data["fdm_4_values"][i], convergence_data["fdm_4_values"][i-1], num_sections, rows[i-1]) if i >= 3 else float('NaN'),
                ])
                
                # Append data for FEM
                fem_table_data.append([
                    analy_interface_temp,
                    convergence_data["fem_2_values"][i],
                    convergence_data["fem_2_extrap"][i],
                    convergence_data["fem_2_errors"][i],
                    convergence_data["fem_2_b"][i],
                    calc_error(convergence_data["fem_2_extrap"][i], convergence_data["fem_2_values"][i]) if i >= 2 else float('NaN'),
                    calc_beta(convergence_data["fem_2_extrap"][i], convergence_data["fem_2_values"][i], convergence_data["fem_2_values"][i-1], num_sections, rows[i-1]) if i >= 3 else float('NaN'),

                    convergence_data["fem_4_values"][i],
                    convergence_data["fem_4_extrap"][i],
                    convergence_data["fem_4_errors"][i],
                    convergence_data["fem_4_b"][i],
                    calc_error(convergence_data["fem_4_extrap"][i], convergence_data["fem_4_values"][i]) if i >= 2 else float('NaN'),
                    calc_beta(convergence_data["fem_4_extrap"][i], convergence_data["fem_4_values"][i], convergence_data["fem_4_values"][i-1], num_sections, rows[i-1]) if i >= 3 else float('NaN'),
                ])
                
                # Store extrapolated % errors for plotting
                if i >= 2:
                    if "fdm_2_extrap_errors" not in convergence_data:
                        convergence_data["fdm_2_extrap_errors"] = [float('NaN')] * 2
                        convergence_data["fdm_4_extrap_errors"] = [float('NaN')] * 2
                        convergence_data["fem_2_extrap_errors"] = [float('NaN')] * 2
                        convergence_data["fem_4_extrap_errors"] = [float('NaN')] * 2
                    convergence_data["fdm_2_extrap_errors"].append(calc_error(convergence_data["fdm_2_extrap"][i], convergence_data["fdm_2_values"][i]))
                    convergence_data["fdm_4_extrap_errors"].append(calc_error(convergence_data["fdm_4_extrap"][i], convergence_data["fdm_4_values"][i]))
                    convergence_data["fem_2_extrap_errors"].append(calc_error(convergence_data["fem_2_extrap"][i], convergence_data["fem_2_values"][i]))
                    convergence_data["fem_4_extrap_errors"].append(calc_error(convergence_data["fem_4_extrap"][i], convergence_data["fem_4_values"][i]))
                else:
                    if "fdm_2_extrap_errors" in convergence_data:
                        convergence_data["fdm_2_extrap_errors"].append(float('NaN'))
                        convergence_data["fdm_4_extrap_errors"].append(float('NaN'))
                        convergence_data["fem_2_extrap_errors"].append(float('NaN'))
                        convergence_data["fem_4_extrap_errors"].append(float('NaN'))


            # Update columns for new fields
            fdm_columns = [
                'True T', 'FDM 2 T', 'FDM 2 Extrap T', 'FDM 2 % Error', 'FDM 2 Exact B',
                'FDM 2 Extrap % Error', 'FDM 2 Extrap B',
                'FDM 4 T', 'FDM 4 Extrap T', 'FDM 4 % Error', 'FDM 4 Exact B',
                'FDM 4 Extrap % Error', 'FDM 4 Extrap B'
            ]

            fem_columns = [
                'True T', 'FEM 2 T', 'FEM 2 Extrap T', 'FEM 2 % Error', 'FEM 2 Exact B',
                'FEM 2 Extrap % Error', 'FEM 2 Extrap B',
                'FEM 4 T', 'FEM 4 Extrap T', 'FEM 4 % Error', 'FEM 4 Exact B',
                'FEM 4 Extrap % Error', 'FEM 4 Extrap B'
            ]

            fdm_df = pd.DataFrame(fdm_table_data, index=[f'N = {n}' for n in rows], columns=fdm_columns)
            fem_df = pd.DataFrame(fem_table_data, index=[f'N = {n}' for n in rows], columns=fem_columns)

            print(f"Convergence of FDM Temperature at x_bar = {x_bar:.3f} and k_ratio = {k_ratio}")
            print(fdm_df.to_string())
            print("\n")

            print(f"Convergence of FEM Temperature at x_bar = {x_bar:.3f} and k_ratio = {k_ratio}")
            print(fem_df.to_string())
            print("\n")

            # Plot convergence for FDM and FEM errors and extrapolated errors
            ax_temp = axs_temp[idx]
            num_sections = data_dict["convergence"]["num_sections"]

            # Plot analytical error convergence for FDM and FEM, orders 2 and 4
            ax_temp.plot(num_sections, convergence_data["fdm_2_errors"], 'o-', label="FDM 2 % Error (Analytical)")
            ax_temp.plot(num_sections, convergence_data["fdm_4_errors"], 's-', label="FDM 4 % Error (Analytical)")
            ax_temp.plot(num_sections, convergence_data["fem_2_errors"], '^-', label="FEM 2 % Error (Analytical)")
            ax_temp.plot(num_sections, convergence_data["fem_4_errors"], 'd-', label="FEM 4 % Error (Analytical)")

            # Plot extrapolated error convergence for FDM and FEM, orders 2 and 4
            ax_temp.plot(num_sections, convergence_data["fdm_2_extrap_errors"], 'o--', label="FDM 2 % Error (Extrapolated)")
            ax_temp.plot(num_sections, convergence_data["fdm_4_extrap_errors"], 's--', label="FDM 4 % Error (Extrapolated)")
            ax_temp.plot(num_sections, convergence_data["fem_2_extrap_errors"], '^--', label="FEM 2 % Error (Extrapolated)")
            ax_temp.plot(num_sections, convergence_data["fem_4_extrap_errors"], 'd--', label="FEM 4 % Error (Extrapolated)")

            # Set log scales, labels, and legends
            ax_temp.set_xscale("log")
            ax_temp.set_yscale("log")
            ax_temp.set_xlabel("Number of Sections (N)")
            ax_temp.set_ylabel("% Error")
            ax_temp.set_title(f"Interface Temp Convergence (k_ratio={k_ratio})")
            ax_temp.legend()
            ax_temp.grid(True)


        fig_temp.tight_layout(rect=[0, 0.03, 1, 0.95])
        x_bar_path = "x_bar_1/" if x_bar_idx == 0 else "x_bar_2/"
        fig_temp.savefig(f"{file_path+x_bar_path}temp_convergence.png", dpi=300)


def heat_convergence():
    '''Convergence of FDM and FEM of interface heat convection'''
    print("Section 3 Temp: Convergence of FDM and FEM Temp with Varying k Ratio")
    file_path = "images/convergence/"
    l = data_dict["length"]
    
    data_dict["convergence"] = {}
    data_dict["convergence"]["num_sections"] = [4, 8, 16, 32, 64, 128, 256, 512]

    for x_bar_idx, x_bar in enumerate(data_dict["x_bar_values"]):
        num_k_ratios = len(data_dict["k_ratios"])
        num_rows = math.ceil(num_k_ratios / 2)

        fig_heat, axs_heat = plt.subplots(num_rows, 2, figsize=(15, 5 * num_rows))

        axs_heat = axs_heat.flatten()

        for idx, k_ratio in enumerate(data_dict["k_ratios"]):
            Bar.set_k_ratio(k_ratio, bar1, bar2)

            # Analytical temperature at the interface
            analy_heat = get_analytical_heat_convection(l, bar1, bar2, x_bar, l)
            convergence_data = {
                "fdm_2_values": [], "fdm_4_values": [], "fem_2_values": [], "fem_4_values": [],
                "fdm_2_errors": [], "fdm_4_errors": [], "fem_2_errors": [], "fem_4_errors": [],
                "fdm_2_extrap": [], "fdm_4_extrap": [], "fem_2_extrap": [], "fem_4_extrap": [],
                "fdm_2_b": [], "fdm_4_b": [], "fem_2_b": [], "fem_4_b": []
            }

            # Calculate temperature values and errors for different section counts
            for num_sections in data_dict["convergence"]["num_sections"]:
                temp_results = {
                    "FDM_2": solve(bar1, bar2, num_sections, x_bar, l, method="FDM", order=2)[0],
                    "FDM_4": solve(bar1, bar2, num_sections, x_bar, l, method="FDM", order=4)[0],
                    "FEM_2": solve(bar1, bar2, num_sections, x_bar, l, method="FEM", order=2)[0],
                    "FEM_4": solve(bar1, bar2, num_sections, x_bar, l, method="FEM", order=4)[0]
                }

                dx = (l - x_bar) / (num_sections // 2)
                for method, temps in temp_results.items():
                    heat = heat_extraction(temps[-2], temps[-1], dx, bar2, order=int(method[-1]), method=method[:3])
                    convergence_data[f"{method.lower()}_values"].append(heat)
                    error = calc_error(analy_heat, heat)
                    convergence_data[f"{method.lower()}_errors"].append(error)

            # Calculate convergence rates `B` for each method and apply extrapolation from the third entry onward
            for method in ["fdm_2", "fdm_4", "fem_2", "fem_4"]:
                values = convergence_data[f"{method}_values"]
                errors = convergence_data[f"{method}_errors"]

                # Initialize `extrap` with NaNs for the first two entries
                extrapolated_values = [float('NaN')] * 2
                b_values = [float('NaN')]  # Initialize `B` list with NaN for the first entry

                # Compute extrapolated values starting from the third entry if possible
                for i in range(2, len(values)):
                    if i >= 2:
                        extrap_value = calc_extrapolated(values[i - 2], values[i - 1], values[i])
                        extrapolated_values.append(extrap_value)

                        # Calculate extrapolated error for the current value
                        extrap_error = calc_error(extrap_value, values[i])
                        errors[i] = extrap_error
                    else:
                        extrapolated_values.append(float('NaN'))

                # Calculate convergence rates `B` for each consecutive pair
                for i in range(1, len(values)):
                    beta = calc_beta(analy_heat, values[i], values[i - 1],
                                    data_dict["convergence"]["num_sections"][i],
                                    data_dict["convergence"]["num_sections"][i - 1])
                    b_values.append(beta)
                
                # Store the calculated `B` and extrapolated values in the convergence data
                convergence_data[f"{method}_b"] = b_values
                convergence_data[f"{method}_extrap"] = extrapolated_values



            # Prepare data for FDM and FEM tables
            rows = data_dict["convergence"]["num_sections"]
            fdm_table_data = []
            fem_table_data = []

            for i, num_sections in enumerate(rows):
                # Append data for FDM
                fdm_table_data.append([
                    analy_heat,
                    convergence_data["fdm_2_values"][i],
                    convergence_data["fdm_2_extrap"][i],
                    convergence_data["fdm_2_errors"][i],
                    convergence_data["fdm_2_b"][i],
                    calc_error(convergence_data["fdm_2_extrap"][i], convergence_data["fdm_2_values"][i]) if i >= 2 else float('NaN'),
                    calc_beta(convergence_data["fdm_2_extrap"][i], convergence_data["fdm_2_values"][i], convergence_data["fdm_2_values"][i-1], num_sections, rows[i-1]) if i >= 3 else float('NaN'),

                    convergence_data["fdm_4_values"][i],
                    convergence_data["fdm_4_extrap"][i],
                    convergence_data["fdm_4_errors"][i],
                    convergence_data["fdm_4_b"][i],
                    calc_error(convergence_data["fdm_4_extrap"][i], convergence_data["fdm_4_values"][i]) if i >= 2 else float('NaN'),
                    calc_beta(convergence_data["fdm_4_extrap"][i], convergence_data["fdm_4_values"][i], convergence_data["fdm_4_values"][i-1], num_sections, rows[i-1]) if i >= 3 else float('NaN'),
                ])
                
                # Append data for FEM
                fem_table_data.append([
                    analy_heat,
                    convergence_data["fem_2_values"][i],
                    convergence_data["fem_2_extrap"][i],
                    convergence_data["fem_2_errors"][i],
                    convergence_data["fem_2_b"][i],
                    calc_error(convergence_data["fem_2_extrap"][i], convergence_data["fem_2_values"][i]) if i >= 2 else float('NaN'),
                    calc_beta(convergence_data["fem_2_extrap"][i], convergence_data["fem_2_values"][i], convergence_data["fem_2_values"][i-1], num_sections, rows[i-1]) if i >= 3 else float('NaN'),

                    convergence_data["fem_4_values"][i],
                    convergence_data["fem_4_extrap"][i],
                    convergence_data["fem_4_errors"][i],
                    convergence_data["fem_4_b"][i],
                    calc_error(convergence_data["fem_4_extrap"][i], convergence_data["fem_4_values"][i]) if i >= 2 else float('NaN'),
                    calc_beta(convergence_data["fem_4_extrap"][i], convergence_data["fem_4_values"][i], convergence_data["fem_4_values"][i-1], num_sections, rows[i-1]) if i >= 3 else float('NaN'),
                ])
                
                # Store extrapolated % errors for plotting
                if i >= 2:
                    if "fdm_2_extrap_errors" not in convergence_data:
                        convergence_data["fdm_2_extrap_errors"] = [float('NaN')] * 2
                        convergence_data["fdm_4_extrap_errors"] = [float('NaN')] * 2
                        convergence_data["fem_2_extrap_errors"] = [float('NaN')] * 2
                        convergence_data["fem_4_extrap_errors"] = [float('NaN')] * 2
                    convergence_data["fdm_2_extrap_errors"].append(calc_error(convergence_data["fdm_2_extrap"][i], convergence_data["fdm_2_values"][i]))
                    convergence_data["fdm_4_extrap_errors"].append(calc_error(convergence_data["fdm_4_extrap"][i], convergence_data["fdm_4_values"][i]))
                    convergence_data["fem_2_extrap_errors"].append(calc_error(convergence_data["fem_2_extrap"][i], convergence_data["fem_2_values"][i]))
                    convergence_data["fem_4_extrap_errors"].append(calc_error(convergence_data["fem_4_extrap"][i], convergence_data["fem_4_values"][i]))
                else:
                    if "fdm_2_extrap_errors" in convergence_data:
                        convergence_data["fdm_2_extrap_errors"].append(float('NaN'))
                        convergence_data["fdm_4_extrap_errors"].append(float('NaN'))
                        convergence_data["fem_2_extrap_errors"].append(float('NaN'))
                        convergence_data["fem_4_extrap_errors"].append(float('NaN'))


            # Update columns for new fields
            fdm_columns = [
                'True Q', 'FDM 2 Q', 'FDM 2 Extrap Q', 'FDM 2 % Error', 'FDM 2 Exact B',
                'FDM 2 Extrap % Error', 'FDM 2 Extrap B',
                'FDM 4 Q', 'FDM 4 Extrap Q', 'FDM 4 % Error', 'FDM 4 Exact B',
                'FDM 4 Extrap % Error', 'FDM 4 Extrap B'
            ]

            fem_columns = [
                'True Q', 'FEM 2 Q', 'FEM 2 Extrap Q', 'FEM 2 % Error', 'FEM 2 Exact B',
                'FEM 2 Extrap % Error', 'FEM 2 Extrap B',
                'FEM 4 Q', 'FEM 4 Extrap Q', 'FEM 4 % Error', 'FEM 4 Exact B',
                'FEM 4 Extrap % Error', 'FEM 4 Extrap B'
            ]

            fdm_df = pd.DataFrame(fdm_table_data, index=[f'N = {n}' for n in rows], columns=fdm_columns)
            fem_df = pd.DataFrame(fem_table_data, index=[f'N = {n}' for n in rows], columns=fem_columns)

            print(f"Convergence of FDM Heat with x_bar = {x_bar:.3f} and k_ratio = {k_ratio}")
            print(fdm_df.to_string())
            print("\n")

            print(f"Convergence of FEM Heat with x_bar = {x_bar:.3f} and k_ratio = {k_ratio}")
            print(fem_df.to_string())
            print("\n")

            # Plot convergence for FDM and FEM errors and extrapolated errors
            ax_heat = axs_heat[idx]
            num_sections = data_dict["convergence"]["num_sections"]

            # Plot analytical error convergence for FDM and FEM, orders 2 and 4
            ax_heat.plot(num_sections, convergence_data["fdm_2_errors"], 'o-', label="FDM 2 % Error (Analytical)")
            ax_heat.plot(num_sections, convergence_data["fdm_4_errors"], 's-', label="FDM 4 % Error (Analytical)")
            ax_heat.plot(num_sections, convergence_data["fem_2_errors"], '^-', label="FEM 2 % Error (Analytical)")
            ax_heat.plot(num_sections, convergence_data["fem_4_errors"], 'd-', label="FEM 4 % Error (Analytical)")

            # Plot extrapolated error convergence for FDM and FEM, orders 2 and 4
            ax_heat.plot(num_sections, convergence_data["fdm_2_extrap_errors"], 'o--', label="FDM 2 % Error (Extrapolated)")
            ax_heat.plot(num_sections, convergence_data["fdm_4_extrap_errors"], 's--', label="FDM 4 % Error (Extrapolated)")
            ax_heat.plot(num_sections, convergence_data["fem_2_extrap_errors"], '^--', label="FEM 2 % Error (Extrapolated)")
            ax_heat.plot(num_sections, convergence_data["fem_4_extrap_errors"], 'd--', label="FEM 4 % Error (Extrapolated)")

            # Set log scales, labels, and legends
            ax_heat.set_xscale("log")
            ax_heat.set_yscale("log")
            ax_heat.set_xlabel("Number of Sections (N)")
            ax_heat.set_ylabel("% Error")
            ax_heat.set_title(f"Interface Heat Convergence (k_ratio={k_ratio})")
            ax_heat.legend()
            ax_heat.grid(True)


        fig_heat.tight_layout(rect=[0, 0.03, 1, 0.95])
        x_bar_path = "x_bar_1/" if x_bar_idx == 0 else "x_bar_2/"
        fig_heat.savefig(f"{file_path+x_bar_path}heat_convergence.png", dpi=300)





def clear_directory():
    # Make directories for plots
    import os
    import shutil

    base_dir = "images"
    sub_dirs = ["analytical", "convergence", "FDM", "FEM"]
    nested_dirs = ["x_bar_1", "x_bar_2"]

    # Create the base directory if it doesn't exist
    if not os.path.exists(base_dir):
        os.mkdir(base_dir)

    # Create or clear subdirectories and their nested directories
    for sub_dir in sub_dirs:
        section_path = os.path.join(base_dir, sub_dir)
        try:
            if os.path.exists(section_path):
                shutil.rmtree(section_path)  # Remove all contents of the directory
            os.makedirs(section_path)  # Recreate the directory
        except PermissionError as e:
            pass
        except Exception as e:
            pass

        # Create nested directories within each section
        for nested_dir in nested_dirs:
            nested_path = os.path.join(section_path, nested_dir)
            try:
                os.makedirs(nested_path)  # Create the nested directory
            except PermissionError as e:
                pass
            except Exception as e:
                pass

def redirect_stdout():
    # Redirect all print statements to a file
    sys.stdout = open("output.txt", "w", encoding="utf-8")

def restore_stdout():
    # Restore original stdout and close the file
    sys.stdout.close()
    sys.stdout = sys.__stdout__


def main():
    # Delete previous images from other runs
    clear_directory()


    redirect_stdout()  # Redirects the stdout to output.txt

    analytical() # Analytical solutions
    section_2(method="FDM") # FDM temperature and heat convection for different k ratios
    section_2(method="FEM") # FEM temperature and heat convection for different k ratios
    temp_convergence() # Convergence of FDM and FEM temperature
    heat_convergence()  # Convergence of FDM and FEM temperature

    restore_stdout()
    


if __name__ == "__main__":
    main()