# main.py

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

import case1
import case2
import common

data_dict = {}
data_dict["a_values"] = np.array([0.29, 0.5, 1.35, 2.75, 4.75, 9.15])

def prob1():
    '''Analytical solutions'''
    print("Problem 1")
    # Problem 1 - Analytical solution
    data_dict["prob1"] = {}
    data_dict["prob1"]["x_values"] = np.array([0, 0.125, 0.25, 0.375, 0.5, 0.625, 0.75, 0.875, 1.0])
    data_dict["prob1"]["x_values_fine"] = np.linspace(0, 1)

    # Case 1 temperature calculations
    results = []
    results_fine = []
    for a_value in data_dict["a_values"]:
        result = case1.analytical(a=a_value, x=data_dict["prob1"]["x_values"])
        result_fine = case1.analytical(a=a_value, x=data_dict["prob1"]["x_values_fine"])
        results.append(result)
        results_fine.append(result_fine)

    data_dict["prob1"]["case1_temp_results"] = np.array(results_fine)

    df_case1 = pd.DataFrame(results, index=[f'a = {a}' for a in data_dict["a_values"]], columns=[f'x = {x:.3f}' for x in data_dict["prob1"]["x_values"]])

    # Pretty print DataFrame for Case 1 temperature
    print("Analytical Case 1 Temperature Results:")
    print(df_case1)
    print("\n" * 2)

    # Plotting Case 1 temperature
    plt.figure(figsize=(10, 6))
    for idx, a_value in enumerate(data_dict["a_values"]):
        plt.plot(data_dict["prob1"]["x_values_fine"], data_dict["prob1"]["case1_temp_results"][idx], label=f'α = {a_value}')
    plt.xlabel('Position (cm)')
    plt.ylabel('Temperature (°C)')
    plt.legend()
    plt.grid(True)
    plt.savefig('prob1_case1_analytical_temperature_vs_position.png', dpi=300)
    #plt.show()
    plt.close()

    # Case 2 temperature calculations
    results = []
    results_fine = []
    for a_value in data_dict["a_values"]:
        result = case2.analytical(a=a_value, x=data_dict["prob1"]["x_values"])
        result_fine = case2.analytical(a=a_value, x=data_dict["prob1"]["x_values_fine"])
        results.append(result)
        results_fine.append(result_fine)

    data_dict["prob1"]["case2_temp_results"] = np.array(results_fine)

    df_case2 = pd.DataFrame(results, index=[f'a = {a}' for a in data_dict["a_values"]], columns=[f'x = {x:.3f}' for x in data_dict["prob1"]["x_values"]])

    # Pretty print DataFrame for Case 2
    print("Analytical Case 2 Temperature Results:")
    print(df_case2)
    print("\n" * 2)

    # Plotting Case 2
    plt.figure(figsize=(10, 6))
    for idx, a_value in enumerate(data_dict["a_values"]):
        plt.plot(data_dict["prob1"]["x_values_fine"], data_dict["prob1"]["case2_temp_results"][idx], label=f'α = {a_value}')
    plt.xlabel('Position (cm)')
    plt.ylabel('Temperature (°C)')
    plt.legend()
    plt.grid(True)
    plt.savefig('prob1_case2_analytical_temperature_vs_position.png', dpi=300)
    #plt.show()
    plt.close()

    # Comparison plots between cases for temperature distribution
    fig, axes = plt.subplots(3, 2, figsize=(12, 12), sharex=True, sharey=True)

    for idx, a_value in enumerate(data_dict["a_values"]):
        ax = axes[idx // 2, idx % 2]
        ax.plot(data_dict["prob1"]["x_values_fine"], data_dict["prob1"]["case1_temp_results"][idx], label='Case 1')
        ax.plot(data_dict["prob1"]["x_values_fine"], data_dict["prob1"]["case2_temp_results"][idx], label='Case 2')
        ax.set_title(f'α = {a_value}')
        ax.set_xlabel('Position (cm)')
        ax.set_ylabel('Temperature (°C)')
        ax.legend()
        ax.grid(True)

    plt.tight_layout(rect=[0, 0, 1, 0.95])  # Adjust layout to make room for the main title
    plt.savefig('prob1_comparison_cases.png', dpi=300)
    #plt.show()
    plt.close()

    # Case 1 heat flux calculations
    results = []
    for a_value in data_dict["a_values"]:
        result = case1.analytical_heat_loss(a=a_value)
        results.append(result)

    data_dict["prob1"]["case1_heat_flux_results"] = np.array(results)

    # Create a DataFrame with alpha as the index and heat flux as the data
    df_case1_heat_flux = pd.DataFrame(results, index=[f'a = {a}' for a in data_dict["a_values"]], columns=['Heat Flux'])

    # Pretty print DataFrame for Case 1 heat flux
    print("Analytical Case 1 Heat Flux Results:")
    print(df_case1_heat_flux)
    print("\n" * 2)

    # Case 2 heat flux calculations
    results = []
    for a_value in data_dict["a_values"]:
        result = case2.analytical_heat_loss(a=a_value)
        results.append(result)

    data_dict["prob1"]["case2_heat_flux_results"] = np.array(results)

    # Create a DataFrame with alpha as the index and heat flux as the data
    df_case2_heat_flux = pd.DataFrame(results, index=[f'a = {a}' for a in data_dict["a_values"]], columns=['Heat Flux'])

    # Pretty print DataFrame for Case 2 heat flux
    print("Analytical Case 2 Heat Flux Results:")
    print(df_case2_heat_flux)
    print("\n" * 2)

    # Plotting heat flux vs alpha for both cases
    plt.figure(figsize=(10, 6))
    plt.plot(data_dict["a_values"], data_dict["prob1"]["case1_heat_flux_results"], label='Case 1')
    plt.plot(data_dict["a_values"], data_dict["prob1"]["case2_heat_flux_results"], label='Case 2')
    plt.xlabel('Alpha ($cm^{-1}$)')
    plt.ylabel('Heat Flux (W)')
    plt.legend()
    plt.grid(True)
    plt.savefig('prob1_heat_flux_comparison.png', dpi=300)
    #plt.show()
    plt.close()

def prob2():
    '''FDM with dx = 0.125 and alpha varying'''
    print("Problem 2")
    data_dict["prob2"] = {}
    data_dict["prob2"]["num_point"] = 8 # This does not consider the final T(L) point. So actually we have 9 points
    data_dict["prob2"]["dx"] = 0.125

    # Case 1
    results = []
    # num_point does not consider T(L). So "4" points is T0, T1, T2, T3, and then we add back T(L)
    for a_value in data_dict["a_values"]:
        fdm_array = case1.fdm_array(data_dict["prob2"]["num_point"], a = a_value)
        fdm_array_inverse = np.linalg.inv(fdm_array)

        right_array = case1.right_array(data_dict["prob2"]["num_point"])

        unkown_temps_array = fdm_array_inverse @ right_array
        solution = [0] # First point is 0 from BC
        solution[1:] = unkown_temps_array
        solution.append(100) # T(L) BC

        results.append(solution)

    data_dict["prob2"]["case1"] = {}
    data_dict["prob2"]["case1"]["temp_results"] = np.array(results)

    data_dict["prob2"]["x_values"] = np.array([0, 0.125, 0.25, 0.375, 0.5, 0.625, 0.75, 0.875, 1.0])
    df_case1 = pd.DataFrame(results, index=[f'a = {a}' for a in data_dict["a_values"]], columns=[f'x = {x:.3f}' for x in data_dict["prob2"]["x_values"]])

    # Pretty print DataFrame for Case 1 temperature
    print("FDM Case 1 Temperature Results:")
    print(df_case1)
    print("\n" * 2)

    # Plotting Case 1 temperature
    plt.figure(figsize=(10, 6))
    for idx, a_value in enumerate(data_dict["a_values"]):
        plt.plot(data_dict["prob2"]["x_values"], data_dict["prob2"]["case1"]["temp_results"][idx], label=f'α = {a_value}')
    plt.xlabel('Position (cm)')
    plt.ylabel('Temperature (°C)')
    plt.legend()
    plt.grid(True)
    plt.savefig('prob2_case1_fdm_temperature_vs_position.png', dpi=300)
    #plt.show()
    plt.close()

    # Case 2
    results = []
    # num_point does not consider T(L). So "4" points is T0, T1, T2, T3, and then we add back T(L)
    for a_value in data_dict["a_values"]:
        fdm_array = case2.fdm_array(data_dict["prob2"]["num_point"], a = a_value)
        fdm_array_inverse = np.linalg.inv(fdm_array)

        right_array = case2.right_array(data_dict["prob2"]["num_point"])  

        unkown_temps_array = fdm_array_inverse @ right_array

        solution = unkown_temps_array.tolist()
        solution.append(100) # T(L) BC

        results.append(solution)

    data_dict["prob2"]["case2"] = {}
    data_dict["prob2"]["case2"]["temp_results"] = np.array(results)

    data_dict["prob2"]["x_values"] = np.array([0, 0.125, 0.25, 0.375, 0.5, 0.625, 0.75, 0.875, 1.0])
    df_case2 = pd.DataFrame(results, index=[f'a = {a}' for a in data_dict["a_values"]], columns=[f'x = {x:.3f}' for x in data_dict["prob2"]["x_values"]])

    # Pretty print DataFrame for Case 2 temperature
    print("FDM Case 2 Temperature Results:")
    print(df_case2)
    print("\n" * 2)

    # Plotting Case 2 temperature
    plt.figure(figsize=(10, 6))
    for idx, a_value in enumerate(data_dict["a_values"]):
        plt.plot(data_dict["prob2"]["x_values"], data_dict["prob2"]["case2"]["temp_results"][idx], label=f'α = {a_value}')
    plt.xlabel('Position (cm)')
    plt.ylabel('Temperature (°C)')
    plt.legend()
    plt.grid(True)
    plt.savefig('prob2_case2_fdm_temperature_vs_position.png', dpi=300)
    #plt.show()
    plt.close()


    # Plot the comparisons similar to problem 1
    # Comparison plots between cases for temperature distribution
    fig, axes = plt.subplots(3, 2, figsize=(12, 12), sharex=True, sharey=True)

    for idx, a_value in enumerate(data_dict["a_values"]):
        ax = axes[idx // 2, idx % 2]
        ax.plot(data_dict["prob2"]["x_values"], data_dict["prob2"]["case1"]["temp_results"][idx], label='Case 1')
        ax.plot(data_dict["prob2"]["x_values"], data_dict["prob2"]["case2"]["temp_results"][idx], label='Case 2')
        ax.set_title(f'α = {a_value}')
        ax.set_xlabel('Position (cm)')
        ax.set_ylabel('Temperature (°C)')
        ax.legend()
        ax.grid(True)

    plt.tight_layout(rect=[0, 0, 1, 0.95])  # Adjust layout to make room for the main title
    plt.savefig('prob2_comparison_cases.png', dpi=300)
    #plt.show()
    plt.close()



    # Heat flux extraction
    # Case 1
    results = []
    for i in range(len(data_dict["a_values"])):
        a = data_dict["a_values"][i]
        (t_0, t_1) = data_dict["prob2"]["case1"]["temp_results"][i][-2:]
        dx = data_dict["prob2"]["dx"]
        first_order = common.taylor_extraction(t_0, t_1, dx, a, 1)
        second_order = common.taylor_extraction(t_0, t_1, dx, a, 2)
        results.append([first_order, second_order])

    data_dict["prob2"]["case1"]["heat_results"] = np.array(results)

    # Create a DataFrame with alpha as the index and heat flux as the data
    df_case1_heat_flux = pd.DataFrame(results, index=[f'a = {a}' for a in data_dict["a_values"]], columns=['1st Order', '2nd Order'])

    # Pretty print DataFrame for Case 1 heat flux
    print("FDM Case 1 Heat Flux Results:")
    print(df_case1_heat_flux)
    print("\n" * 2)

    # Plotting heat flux vs alpha for both extractions
    plt.figure(figsize=(10, 6))
    plt.plot(data_dict["a_values"], [i[0] for i in data_dict["prob2"]["case1"]["heat_results"]], label='1st Order Extraction')
    plt.plot(data_dict["a_values"], [i[1] for i in data_dict["prob2"]["case1"]["heat_results"]], label='2nd Order Extraction')
    plt.xlabel('Alpha ($cm^{-1}$)')
    plt.ylabel('Heat Flux (W)')
    plt.legend()
    plt.grid(True)
    plt.savefig('prob2_case1_heat_flux.png', dpi=300)
    #plt.show()
    plt.close()



    # Case 2
    results = []
    for i in range(len(data_dict["a_values"])):
        a = data_dict["a_values"][i]
        (t_0, t_1) = data_dict["prob2"]["case2"]["temp_results"][i][-2:]
        dx = data_dict["prob2"]["dx"]
        first_order = common.taylor_extraction(t_0, t_1, dx, a, 1)
        second_order = common.taylor_extraction(t_0, t_1, dx, a, 2)
        results.append([first_order, second_order])

    data_dict["prob2"]["case2"]["heat_results"] = np.array(results)

    # Create a DataFrame with alpha as the index and heat flux as the data
    df_case2_heat_flux = pd.DataFrame(results, index=[f'a = {a}' for a in data_dict["a_values"]], columns=['1st Order', '2nd Order'])

    # Pretty print DataFrame for Case 1 heat flux
    print("FDM Case 2 Heat Flux Results:")
    print(df_case2_heat_flux)
    print("\n" * 2)

    # Plotting heat flux vs alpha for both extractions
    plt.figure(figsize=(10, 6))
    plt.plot(data_dict["a_values"], [i[0] for i in data_dict["prob2"]["case2"]["heat_results"]], label='1st Order Extraction')
    plt.plot(data_dict["a_values"], [i[1] for i in data_dict["prob2"]["case2"]["heat_results"]], label='2nd Order Extraction')
    plt.xlabel('Alpha ($cm^{-1}$)')
    plt.ylabel('Heat Flux (W)')
    plt.legend()
    plt.grid(True)
    plt.savefig('prob2_case2_heat_flux.png', dpi=300)
    #plt.show()
    plt.close()

def prob3():
    '''FEM with dx = 0.125 and alpha varying'''
    print("Problem 3")
    data_dict["prob3"] = {}
    data_dict["prob3"]["num_point"] = 8 # This does not consider the final T(L) point. So actually we have 9 points
    data_dict["prob3"]["dx"] = 0.125

    # Case 1
    results = []
    # num_point does not consider T(L). So "4" points is T0, T1, T2, T3, and then we add back T(L)
    for a_value in data_dict["a_values"]:
        fdm_array = case1.fem_array(data_dict["prob3"]["num_point"], a = a_value)
        fdm_array_inverse = np.linalg.inv(fdm_array)

        right_array = case1.right_array(data_dict["prob3"]["num_point"])

        unkown_temps_array = fdm_array_inverse @ right_array
        solution = [0] # First point is 0 from BC
        solution[1:] = unkown_temps_array
        solution.append(100) # T(L) BC

        results.append(solution)

    data_dict["prob3"]["case1"] = {}
    data_dict["prob3"]["case1"]["temp_results"] = np.array(results)

    data_dict["prob3"]["x_values"] = np.array([0, 0.125, 0.25, 0.375, 0.5, 0.625, 0.75, 0.875, 1.0])
    df_case1 = pd.DataFrame(results, index=[f'a = {a}' for a in data_dict["a_values"]], columns=[f'x = {x:.3f}' for x in data_dict["prob3"]["x_values"]])

    # Pretty print DataFrame for Case 1 temperature
    print("FEM Case 1 Temperature Results:")
    print(df_case1)
    print("\n" * 2)

    # Plotting Case 1 temperature
    plt.figure(figsize=(10, 6))
    for idx, a_value in enumerate(data_dict["a_values"]):
        plt.plot(data_dict["prob3"]["x_values"], data_dict["prob3"]["case1"]["temp_results"][idx], label=f'α = {a_value}')
    plt.xlabel('Position (cm)')
    plt.ylabel('Temperature (°C)')
    plt.legend()
    plt.grid(True)
    plt.savefig('prob3_case1_fem_temperature_vs_position.png', dpi=300)
    #plt.show()
    plt.close()

    # Case 2
    results = []
    # num_point does not consider T(L). So "4" points is T0, T1, T2, T3, and then we add back T(L)
    for a_value in data_dict["a_values"]:
        fdm_array = case2.fem_array(data_dict["prob3"]["num_point"], a = a_value)
        fdm_array_inverse = np.linalg.inv(fdm_array)

        right_array = case2.right_array(data_dict["prob3"]["num_point"])  

        unkown_temps_array = fdm_array_inverse @ right_array

        solution = unkown_temps_array.tolist()
        solution.append(100) # T(L) BC

        results.append(solution)

    data_dict["prob3"]["case2"] = {}
    data_dict["prob3"]["case2"]["temp_results"] = np.array(results)

    data_dict["prob3"]["x_values"] = np.array([0, 0.125, 0.25, 0.375, 0.5, 0.625, 0.75, 0.875, 1.0])
    df_case2 = pd.DataFrame(results, index=[f'a = {a}' for a in data_dict["a_values"]], columns=[f'x = {x:.3f}' for x in data_dict["prob3"]["x_values"]])

    # Pretty print DataFrame for Case 2 temperature
    print("FEM Case 2 Temperature Results:")
    print(df_case2)
    print("\n" * 2)

    # Plotting Case 2 temperature
    plt.figure(figsize=(10, 6))
    for idx, a_value in enumerate(data_dict["a_values"]):
        plt.plot(data_dict["prob3"]["x_values"], data_dict["prob3"]["case2"]["temp_results"][idx], label=f'α = {a_value}')
    plt.xlabel('Position (cm)')
    plt.ylabel('Temperature (°C)')
    plt.legend()
    plt.grid(True)
    plt.savefig('prob3_case2_fem_temperature_vs_position.png', dpi=300)
    #plt.show()
    plt.close()


    # Plot the comparisons similar to problem 1
    # Comparison plots between cases for temperature distribution
    fig, axes = plt.subplots(3, 2, figsize=(12, 12), sharex=True, sharey=True)

    for idx, a_value in enumerate(data_dict["a_values"]):
        ax = axes[idx // 2, idx % 2]
        ax.plot(data_dict["prob3"]["x_values"], data_dict["prob3"]["case1"]["temp_results"][idx], label='Case 1')
        ax.plot(data_dict["prob3"]["x_values"], data_dict["prob3"]["case2"]["temp_results"][idx], label='Case 2')
        ax.set_title(f'α = {a_value}')
        ax.set_xlabel('Position (cm)')
        ax.set_ylabel('Temperature (°C)')
        ax.legend()
        ax.grid(True)

    plt.tight_layout(rect=[0, 0, 1, 0.95])  # Adjust layout to make room for the main title
    plt.savefig('prob3_comparison_cases.png', dpi=300)
    #plt.show()
    plt.close()



    # Heat flux extraction
    # Case 1
    results = []
    for i in range(len(data_dict["a_values"])):
        a = data_dict["a_values"][i]
        (t_0, t_1) = data_dict["prob3"]["case1"]["temp_results"][i][-2:]
        dx = data_dict["prob3"]["dx"]
        result = common.fem_heat_extraction(t_0, t_1, dx, a)
        results.append([result])

    data_dict["prob3"]["case1"]["heat_results"] = np.array(results)

    # Create a DataFrame with alpha as the index and heat flux as the data
    df_case1_heat_flux = pd.DataFrame(results, index=[f'a = {a}' for a in data_dict["a_values"]], columns=['Heat Extraction'])

    # Pretty print DataFrame for Case 1 heat flux
    print("FEM Case 1 Heat Flux Results:")
    print(df_case1_heat_flux)
    print("\n" * 2)

    # Plotting heat flux vs alpha for both extractions
    plt.figure(figsize=(10, 6))
    plt.plot(data_dict["a_values"], [i[0] for i in data_dict["prob3"]["case1"]["heat_results"]], label='Heat Extraction')
    plt.xlabel('Alpha ($cm^{-1}$)')
    plt.ylabel('Heat Flux (W)')
    plt.legend()
    plt.grid(True)
    plt.savefig('prob3_case1_heat_flux.png', dpi=300)
    #plt.show()
    plt.close()



    # Case 2
    results = []
    for i in range(len(data_dict["a_values"])):
        a = data_dict["a_values"][i]
        (t_0, t_1) = data_dict["prob3"]["case2"]["temp_results"][i][-2:]
        dx = data_dict["prob3"]["dx"]
        result = common.fem_heat_extraction(t_0, t_1, dx, a)
        results.append([result])

    data_dict["prob3"]["case2"]["heat_results"] = np.array(results)

    # Create a DataFrame with alpha as the index and heat flux as the data
    df_case2_heat_flux = pd.DataFrame(results, index=[f'a = {a}' for a in data_dict["a_values"]], columns=['Heat Extraction'])

    # Pretty print DataFrame for Case 1 heat flux
    print("FEM Case 2 Heat Flux Results:")
    print(df_case2_heat_flux)
    print("\n" * 2)

    # Plotting heat flux vs alpha for both extractions
    plt.figure(figsize=(10, 6))
    plt.plot(data_dict["a_values"], [i[0] for i in data_dict["prob3"]["case2"]["heat_results"]], label='Heat Extraction')
    plt.xlabel('Alpha ($cm^{-1}$)')
    plt.ylabel('Heat Flux (W)')
    plt.legend()
    plt.grid(True)
    plt.savefig('prob3_case2_heat_flux.png', dpi=300)
    #plt.show()
    plt.close()

def prob4():
    '''Convergence analysis of the cases'''
    print("Problem 4")
    data_dict["prob4"] = {}
    data_dict["prob4"]["num_points"] = [2**i for i in range(2, 8)] # This does not consider the final T(L) point. So actually we have n + 1 points
    data_dict["prob4"]["dx_values"] = [1 / i for i in data_dict["prob4"]["num_points"]]

    # Case 1
    # Get fdm temp results
    data_dict["prob4"]["case1"] = {}
    data_dict["prob4"]["case1"]["fdm"] = {}
    data_dict["prob4"]["case1"]["fdm"]["temp_results"] = []
    # num_point does not consider T(L). So "4" points is T0, T1, T2, T3, and then we add back T(L)
    for num_point in data_dict["prob4"]["num_points"]:
        fdm_array = case1.fdm_array(num_point)
        fdm_array_inverse = np.linalg.inv(fdm_array)

        right_array = case1.right_array(num_point)

        unkown_temps_array = fdm_array_inverse @ right_array
        solution = [0] # First point is 0 from BC
        solution[1:] = unkown_temps_array
        solution.append(100) # T(L) BC

        data_dict["prob4"]["case1"]["fdm"]["temp_results"].append(solution)

    # Get fem temp results
    data_dict["prob4"]["case1"]["fem"] = {}
    data_dict["prob4"]["case1"]["fem"]["temp_results"] = []
    # num_point does not consider T(L). So "4" points is T0, T1, T2, T3, and then we add back T(L)
    for num_point in data_dict["prob4"]["num_points"]:
        fdm_array = case1.fem_array(num_point)
        fdm_array_inverse = np.linalg.inv(fdm_array)

        right_array = case1.right_array(num_point)

        unkown_temps_array = fdm_array_inverse @ right_array
        solution = [0] # First point is 0 from BC
        solution[1:] = unkown_temps_array
        solution.append(100) # T(L) BC

        data_dict["prob4"]["case1"]["fem"]["temp_results"].append(solution)
    
    # We have all of the temperature results for varying dx for both fdm and fem
    # For each dx, get the midpoint temperature from fdm and fem
    results = []
    for i in range(len(data_dict["prob4"]["num_points"])):
        a = 2.75
        num_points = data_dict["prob4"]["num_points"][i]
        midpoint_index = int(num_points / 2)
        t_fdm = data_dict["prob4"]["case1"]["fdm"]["temp_results"][i][midpoint_index]
        t_fem = data_dict["prob4"]["case1"]["fem"]["temp_results"][i][midpoint_index]

        results.append([t_fdm, t_fem])

    data_dict["prob4"]["case1"]["midpoint_temp_results"] = np.array(results)
    data_dict["prob4"]["case1"]["midpoint_temp_true"] = case1.analytical(a = 2.75, x = 0.5)

    # Calculate % error and beta values for midpoint temperature
    errors = []
    betas = []
    t_exact = data_dict["prob4"]["case1"]["midpoint_temp_true"]
    t_results = data_dict["prob4"]["case1"]["midpoint_temp_results"]

    for i in range(0, len(t_results)):
        # % Error calculation
        t_fdm, t_fem = results[i]
        fdm_error = common.calc_error(t_exact, t_fdm)
        fem_error = common.calc_error(t_exact, t_fem)

        # Beta (convergence rate) calculation
        if i != 0:
            t_fdm_prev, t_fem_prev = results[i-1]
            dx_1 = data_dict["prob4"]["dx_values"][i-1]
            dx_2 = data_dict["prob4"]["dx_values"][i]
            fdm_beta = common.calc_beta(t_exact, t_fdm, t_fdm_prev, dx_2, dx_1)
            fem_beta = common.calc_beta(t_exact, t_fem, t_fem_prev, dx_2, dx_1)
        else:
            fdm_beta = 0
            fem_beta = 0
        
        errors.append([fdm_error, fem_error])
        betas.append([fdm_beta, fem_beta])

    data_dict["prob4"]["case1"]["temp_errors"] = errors
    data_dict["prob4"]["case1"]["temp_betas"] = betas

    # Combine dx values, results, errors, and betas into a DataFrame for display
    table_data = []
    for i in range(len(data_dict["prob4"]["dx_values"])):
        table_data.append([
            results[i][0],  # fdm result
            results[i][1],  # fem result
            data_dict["prob4"]["case1"]["midpoint_temp_true"],  # True T
            errors[i][0],  # fdm % error
            errors[i][1],  # fem % error
            betas[i][0],  # fdm beta
            betas[i][1]   # fem beta
        ])

    columns = ['FDM T', 'FEM T', 'True T', 'FDM % Error', 'FEM % Error', 'FDM β', 'FEM β']
    df_case1 = pd.DataFrame(table_data, index=[f'dx = {i}' for i in data_dict["prob4"]["dx_values"]], columns=columns)

    print("FDM and FEM Case 1 Midpoint Temperature Convergence Results:")
    print(df_case1)
    print("\n" * 2)
    # Plotting
    plt.figure(figsize=(8, 6))
    plt.plot(data_dict["prob4"]["dx_values"], [err[0] for err in data_dict["prob4"]["case1"]["temp_errors"]], label='FDM Temp Error')
    plt.plot(data_dict["prob4"]["dx_values"], [err[1] for err in data_dict["prob4"]["case1"]["temp_errors"]], label='FEM Temp Error')
    plt.xscale('log')
    plt.yscale('log')
    plt.xlabel('Δx')
    plt.ylabel('% Error')
    plt.legend()
    plt.grid(True)
    plt.savefig("prob4_case1_temp_convergence.png", dpi=300)
    #plt.show()
    plt.close()



    # For each dx, calculate the 2nd order heat flux for fdm and the heat flux for fem
    results = []
    for i in range(len(data_dict["prob4"]["dx_values"])):
        a = 2.75
        dx = data_dict["prob4"]["dx_values"][i]
        (t_0_fdm, t_1_fdm) = data_dict["prob4"]["case1"]["fdm"]["temp_results"][i][-2:]
        (t_0_fem, t_1_fem) = data_dict["prob4"]["case1"]["fem"]["temp_results"][i][-2:]

        fdm = common.taylor_extraction(t_0_fdm, t_1_fdm, dx, a, 2)
        fem = common.fem_heat_extraction(t_0_fem, t_1_fem, dx, a)
        results.append([fdm, fem])

    data_dict["prob4"]["case1"]["heat_results"] = np.array(results)
    data_dict["prob4"]["case1"]["heat_results_true"] = case1.analytical_heat_loss(a=2.75)

    # Calculate % error and beta values for heat loss
    errors = []
    betas = []
    q_exact = data_dict["prob4"]["case1"]["heat_results_true"]
    q_results = data_dict["prob4"]["case1"]["heat_results"]

    for i in range(0, len(q_results)):
        # % Error calculation
        q_fdm, q_fem = results[i]
        fdm_error = common.calc_error(q_exact, q_fdm)
        fem_error = common.calc_error(q_exact, q_fem)

        # Beta (convergence rate) calculation
        if i != 0:
            q_fdm_prev, q_fem_prev = results[i-1]
            dx_1 = data_dict["prob4"]["dx_values"][i-1]
            dx_2 = data_dict["prob4"]["dx_values"][i]
            fdm_beta = common.calc_beta(q_exact, q_fdm, q_fdm_prev, dx_2, dx_1)
            fem_beta = common.calc_beta(q_exact, q_fem, q_fem_prev, dx_2, dx_1)
        else:
            fdm_beta = 0
            fem_beta = 0
        
        errors.append([fdm_error, fem_error])
        betas.append([fdm_beta, fem_beta])

    data_dict["prob4"]["case1"]["heat_errors"] = errors
    data_dict["prob4"]["case1"]["heat_betas"] = betas

    # Combine dx values, results, errors, and betas into a DataFrame for display
    table_data = []
    for i in range(len(data_dict["prob4"]["dx_values"])):
        table_data.append([
            results[i][0],  # fdm result
            results[i][1],  # fem result
            q_exact,  # True Q
            errors[i][0],  # fdm % error
            errors[i][1],  # fem % error
            betas[i][0],  # fdm beta
            betas[i][1]   # fem beta
        ])

    columns = ['FDM Q', 'FEM Q', 'True Q', 'FDM % Error', 'FEM % Error', 'FDM β', 'FEM β']
    df_case1 = pd.DataFrame(table_data, index=[f'dx = {i}' for i in data_dict["prob4"]["dx_values"]], columns=columns)

    print("FDM and FEM Case 1 Heat Loss Convergence Results:")
    print(df_case1)
    print("\n" * 2)

    # Plotting
    plt.figure(figsize=(8, 6))
    plt.plot(data_dict["prob4"]["dx_values"], [err[0] for err in data_dict["prob4"]["case1"]["heat_errors"]], label='FDM Heat Flux Error')
    plt.plot(data_dict["prob4"]["dx_values"], [err[1] for err in data_dict["prob4"]["case1"]["heat_errors"]], label='FEM Heat Flux Error')
    plt.xscale('log')
    plt.yscale('log')
    plt.xlabel('Δx')
    plt.ylabel('% Error')
    plt.legend()
    plt.grid(True)
    plt.savefig("prob4_case1_heat_convergence.png", dpi=300)
    #plt.show()
    plt.close()




    # Case 2
    # Get fdm temp results
    data_dict["prob4"]["case2"] = {}
    data_dict["prob4"]["case2"]["fdm"] = {}
    data_dict["prob4"]["case2"]["fdm"]["temp_results"] = []
    # num_point does not consider T(L). So "4" points is T0, T1, T2, T3, and then we add back T(L)
    for num_point in data_dict["prob4"]["num_points"]:
        fdm_array = case2.fdm_array(num_point)
        fdm_array_inverse = np.linalg.inv(fdm_array)

        right_array = case2.right_array(num_point)

        unkown_temps_array = fdm_array_inverse @ right_array
        solution = []
        solution[0:] = unkown_temps_array
        solution.append(100) # T(L) BC

        data_dict["prob4"]["case2"]["fdm"]["temp_results"].append(solution)

    # Get fem temp results
    data_dict["prob4"]["case2"]["fem"] = {}
    data_dict["prob4"]["case2"]["fem"]["temp_results"] = []
    # num_point does not consider T(L). So "4" points is T0, T1, T2, T3, and then we add back T(L)
    for num_point in data_dict["prob4"]["num_points"]:
        fdm_array = case2.fem_array(num_point)
        fdm_array_inverse = np.linalg.inv(fdm_array)

        right_array = case2.right_array(num_point)

        unkown_temps_array = fdm_array_inverse @ right_array
        solution = []
        solution[0:] = unkown_temps_array
        solution.append(100) # T(L) BC

        data_dict["prob4"]["case2"]["fem"]["temp_results"].append(solution)
    
    # We have all of the temperature results for varying dx for both fdm and fem
    # For each dx, get the midpoint temperature from fdm and fem
    results = []
    for i in range(len(data_dict["prob4"]["num_points"])):
        a = 2.75
        num_points = data_dict["prob4"]["num_points"][i]
        midpoint_index = int(num_points / 2)
        t_fdm = data_dict["prob4"]["case2"]["fdm"]["temp_results"][i][midpoint_index]
        t_fem = data_dict["prob4"]["case2"]["fem"]["temp_results"][i][midpoint_index]

        results.append([t_fdm, t_fem])

    data_dict["prob4"]["case2"]["midpoint_temp_results"] = np.array(results)
    data_dict["prob4"]["case2"]["midpoint_temp_true"] = case2.analytical(a = 2.75, x = 0.5)

    # Calculate % error and beta values for midpoint temperature
    errors = []
    betas = []
    t_exact = data_dict["prob4"]["case2"]["midpoint_temp_true"]
    t_results = data_dict["prob4"]["case2"]["midpoint_temp_results"]

    for i in range(0, len(t_results)):
        # % Error calculation
        t_fdm, t_fem = results[i]
        fdm_error = common.calc_error(t_exact, t_fdm)
        fem_error = common.calc_error(t_exact, t_fem)

        # Beta (convergence rate) calculation
        if i != 0:
            t_fdm_prev, t_fem_prev = results[i-1]
            dx_1 = data_dict["prob4"]["dx_values"][i-1]
            dx_2 = data_dict["prob4"]["dx_values"][i]
            fdm_beta = common.calc_beta(t_exact, t_fdm, t_fdm_prev, dx_2, dx_1)
            fem_beta = common.calc_beta(t_exact, t_fem, t_fem_prev, dx_2, dx_1)
        else:
            fdm_beta = 0
            fem_beta = 0
        
        errors.append([fdm_error, fem_error])
        betas.append([fdm_beta, fem_beta])

    data_dict["prob4"]["case2"]["temp_errors"] = errors
    data_dict["prob4"]["case2"]["temp_betas"] = betas

    # Combine dx values, results, errors, and betas into a DataFrame for display
    table_data = []
    for i in range(len(data_dict["prob4"]["dx_values"])):
        table_data.append([
            results[i][0],  # fdm result
            results[i][1],  # fem result
            data_dict["prob4"]["case2"]["midpoint_temp_true"],  # True T
            errors[i][0],  # fdm % error
            errors[i][1],  # fem % error
            betas[i][0],  # fdm beta
            betas[i][1]   # fem beta
        ])

    columns = ['FDM T', 'FEM T', 'True T', 'FDM % Error', 'FEM % Error', 'FDM β', 'FEM β']
    df_case2 = pd.DataFrame(table_data, index=[f'dx = {i}' for i in data_dict["prob4"]["dx_values"]], columns=columns)

    print("FDM and FEM Case 2 Midpoint Temperature Convergence Results:")
    print(df_case2)
    print("\n" * 2)
    # Plotting
    plt.figure(figsize=(8, 6))
    plt.plot(data_dict["prob4"]["dx_values"], [err[0] for err in data_dict["prob4"]["case2"]["temp_errors"]], label='FDM Temp Error')
    plt.plot(data_dict["prob4"]["dx_values"], [err[1] for err in data_dict["prob4"]["case2"]["temp_errors"]], label='FEM Temp Error')
    plt.xscale('log')
    plt.yscale('log')
    plt.xlabel('Δx')
    plt.ylabel('% Error')
    plt.legend()
    plt.grid(True)
    plt.savefig("prob4_case2_temp_convergence.png", dpi=300)
    #plt.show()
    plt.close()



    # For each dx, calculate the 2nd order heat flux for fdm and the heat flux for fem
    results = []
    for i in range(len(data_dict["prob4"]["dx_values"])):
        a = 2.75
        dx = data_dict["prob4"]["dx_values"][i]
        (t_0_fdm, t_1_fdm) = data_dict["prob4"]["case2"]["fdm"]["temp_results"][i][-2:]
        (t_0_fem, t_1_fem) = data_dict["prob4"]["case2"]["fem"]["temp_results"][i][-2:]

        fdm = common.taylor_extraction(t_0_fdm, t_1_fdm, dx, a, 2)
        fem = common.fem_heat_extraction(t_0_fem, t_1_fem, dx, a)
        results.append([fdm, fem])

    data_dict["prob4"]["case2"]["heat_results"] = np.array(results)
    data_dict["prob4"]["case2"]["heat_results_true"] = case2.analytical_heat_loss(a=2.75)

    # Calculate % error and beta values for heat loss
    errors = []
    betas = []
    q_exact = data_dict["prob4"]["case2"]["heat_results_true"]
    q_results = data_dict["prob4"]["case2"]["heat_results"]

    for i in range(0, len(q_results)):
        # % Error calculation
        q_fdm, q_fem = results[i]
        fdm_error = common.calc_error(q_exact, q_fdm)
        fem_error = common.calc_error(q_exact, q_fem)

        # Beta (convergence rate) calculation
        if i != 0:
            q_fdm_prev, q_fem_prev = results[i-1]
            dx_1 = data_dict["prob4"]["dx_values"][i-1]
            dx_2 = data_dict["prob4"]["dx_values"][i]
            fdm_beta = common.calc_beta(q_exact, q_fdm, q_fdm_prev, dx_2, dx_1)
            fem_beta = common.calc_beta(q_exact, q_fem, q_fem_prev, dx_2, dx_1)
        else:
            fdm_beta = 0
            fem_beta = 0
        
        errors.append([fdm_error, fem_error])
        betas.append([fdm_beta, fem_beta])

    data_dict["prob4"]["case2"]["heat_errors"] = errors
    data_dict["prob4"]["case2"]["heat_betas"] = betas

    # Combine dx values, results, errors, and betas into a DataFrame for display
    table_data = []
    for i in range(len(data_dict["prob4"]["dx_values"])):
        table_data.append([
            results[i][0],  # fdm result
            results[i][1],  # fem result
            q_exact,  # True Q
            errors[i][0],  # fdm % error
            errors[i][1],  # fem % error
            betas[i][0],  # fdm beta
            betas[i][1]   # fem beta
        ])

    columns = ['FDM Q', 'FEM Q', 'True Q', 'FDM % Error', 'FEM % Error', 'FDM β', 'FEM β']
    df_case2 = pd.DataFrame(table_data, index=[f'dx = {i}' for i in data_dict["prob4"]["dx_values"]], columns=columns)

    print("FDM and FEM Case 2 Heat Loss Convergence Results:")
    print(df_case2)
    print("\n" * 2)

    # Plotting
    plt.figure(figsize=(8, 6))
    plt.plot(data_dict["prob4"]["dx_values"], [err[0] for err in data_dict["prob4"]["case2"]["heat_errors"]], label='FDM Heat Flux Error')
    plt.plot(data_dict["prob4"]["dx_values"], [err[1] for err in data_dict["prob4"]["case2"]["heat_errors"]], label='FEM Heat Flux Error')
    plt.xscale('log')
    plt.yscale('log')
    plt.xlabel('Δx')
    plt.ylabel('% Error')
    plt.legend()
    plt.grid(True)
    plt.savefig("prob4_case2_heat_convergence.png", dpi=300)
    #plt.show()
    plt.close()
    


if __name__ == "__main__":
    print("Homework 1 by Judson Upchurch")
    print("Please note that some of the tables are transposed compared to the LaTeX equivalent")
    print("")

    prob1() # Analyitical temperature distribution

    prob2() # FDM with dx = 0.125 and varying alpha

    prob3() # FEM with dx = 0.125 and varying alpha

    prob4() # Convergence analysis