# 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():
    '''Convergence analysis of the cases'''
    print("Problem 3")
    data_dict["prob3"] = {}
    data_dict["prob3"]["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["prob3"]["dx_values"] = [1 / i for i in data_dict["prob3"]["num_points"]]

    # Case 1
    data_dict["prob3"]["case1"] = {}
    data_dict["prob3"]["case1"]["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["prob3"]["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["prob3"]["case1"]["temp_results"].append(solution)

    #data_dict["prob3"]["x_values_fine"] = np.linspace(0, 1)
    #data_dict["prob3"]["case1"]["true_temp_results"] = case1.analytical(a=2.75, x=data_dict["prob3"]["x_values_fine"])
    
    # We have all of the temperature results for varying dx
    # For each dx, calculate the 1st and 2nd order heat flux
    results = []
    for i in range(len(data_dict["prob3"]["dx_values"])):
        a = 2.75
        dx = data_dict["prob3"]["dx_values"][i]
        (t_0, t_1) = data_dict["prob3"]["case1"]["temp_results"][i][-2:]

        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["prob3"]["case1"]["heat_results"] = np.array(results)

    data_dict["prob3"]["case1"]["heat_results_true"] = case1.analytical_heat_loss(a=2.75)

    # Calculate % error and beta values
    errors = []
    betas = []
    q_exact = data_dict["prob3"]["case1"]["heat_results_true"]
    for i in range(0, len(results)):
        
        # % Error calculation
        q_1_curr, q_2_curr = results[i]
        first_order_error = common.calc_error(q_exact, q_1_curr)
        second_order_error = common.calc_error(q_exact, q_2_curr)
        
        # Beta (convergence rate) calculation
        if i != 0:
            q_1_prev, q_2_prev = results[i-1]
            dx_1 = data_dict["prob3"]["dx_values"][i-1]
            dx_2 = data_dict["prob3"]["dx_values"][i]
            first_order_beta = common.calc_beta(q_exact, q_1_curr, q_1_prev, dx_2, dx_1)
            second_order_beta = common.calc_beta(q_exact, q_2_curr, q_2_prev, dx_2, dx_1)
        else:
            first_order_beta = 0
            second_order_beta = 0
        
        errors.append([first_order_error, second_order_error])
        betas.append([first_order_beta, second_order_beta])

    data_dict["prob3"]["case1"]["errors"] = errors
    data_dict["prob3"]["case1"]["betas"] = betas

    # Combine dx values, results, errors, and betas into a DataFrame for display
    table_data = []

    for i in range(len(data_dict["prob3"]["dx_values"])):
        table_data.append([
            results[i][0],  # 1st order result
            results[i][1],  # 2nd order result
            data_dict["prob3"]["case1"]["heat_results_true"],  # True Q
            errors[i][0],  # 1st order % error
            errors[i][1],  # 2nd order % error
            betas[i][0],  # 1st order beta
            betas[i][1]   # 2nd order beta
        ])

    columns = ['1st Q', '2nd Q', 'True Q', '1st % Error', '2nd % Error', '1st β', '2nd β']
    df_case1 = pd.DataFrame(table_data, index=[f'dx = {i}' for i in data_dict["prob3"]["dx_values"]], columns=columns)

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

    # Plotting
    plt.figure(figsize=(8, 6))
    plt.plot(data_dict["prob3"]["dx_values"], [err[0] for err in data_dict["prob3"]["case1"]["errors"]], label='1st Order Error')
    plt.plot(data_dict["prob3"]["dx_values"], [err[1] for err in data_dict["prob3"]["case1"]["errors"]], label='2nd Order Error')
    plt.xscale('log')
    plt.yscale('log')
    plt.xlabel('Δx')
    plt.ylabel('% Error')
    plt.legend()
    plt.grid(True)
    plt.savefig("prob3_case1_convergence.png", dpi=300)
    #plt.show()
    plt.close()




    # Case 2
    data_dict["prob3"]["case2"] = {}
    data_dict["prob3"]["case2"]["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["prob3"]["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["prob3"]["case2"]["temp_results"].append(solution)
    
    # We have all of the temperature results for varying dx
    # For each dx, calculate the 1st and 2nd order heat flux
    results = []
    for i in range(len(data_dict["prob3"]["dx_values"])):
        a = 2.75
        dx = data_dict["prob3"]["dx_values"][i]
        (t_0, t_1) = data_dict["prob3"]["case2"]["temp_results"][i][-2:]

        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["prob3"]["case2"]["heat_results"] = np.array(results)

    data_dict["prob3"]["case2"]["heat_results_true"] = case2.analytical_heat_loss(a=2.75)

    # Calculate % error and beta values
    errors = []
    betas = []
    q_exact = data_dict["prob3"]["case2"]["heat_results_true"]
    for i in range(0, len(results)):
        
        # % Error calculation
        q_1_curr, q_2_curr = results[i]
        first_order_error = common.calc_error(q_exact, q_1_curr)
        second_order_error = common.calc_error(q_exact, q_2_curr)
        
        # Beta (convergence rate) calculation
        if i != 0:
            q_1_prev, q_2_prev = results[i-1]
            dx_1 = data_dict["prob3"]["dx_values"][i-1]
            dx_2 = data_dict["prob3"]["dx_values"][i]
            first_order_beta = common.calc_beta(q_exact, q_1_curr, q_1_prev, dx_2, dx_1)
            second_order_beta = common.calc_beta(q_exact, q_2_curr, q_2_prev, dx_2, dx_1)
        else:
            first_order_beta = 0
            second_order_beta = 0
        
        errors.append([first_order_error, second_order_error])
        betas.append([first_order_beta, second_order_beta])
    
    data_dict["prob3"]["case2"]["errors"] = errors
    data_dict["prob3"]["case2"]["betas"] = betas

    # Combine dx values, results, errors, and betas into a DataFrame for display
    table_data = []

    for i in range(len(data_dict["prob3"]["dx_values"])):
        table_data.append([
            results[i][0],  # 1st order result
            results[i][1],  # 2nd order result
            data_dict["prob3"]["case2"]["heat_results_true"],  # True Q
            errors[i][0],  # 1st order % error
            errors[i][1],  # 2nd order % error
            betas[i][0],  # 1st order beta
            betas[i][1]   # 2nd order beta
        ])

    columns = ['1st Q', '2nd Q', 'True Q', '1st % Error', '2nd % Error', '1st β', '2nd β']
    df_case2 = pd.DataFrame(table_data, index=[f'dx = {i}' for i in data_dict["prob3"]["dx_values"]], columns=columns)

    print("FDM Case 2 Heat Flux Convergence Results:")
    print(df_case2)
    print("\n" * 2)

    # Plotting
    plt.figure(figsize=(8, 6))
    plt.plot(data_dict["prob3"]["dx_values"], [err[0] for err in data_dict["prob3"]["case2"]["errors"]], label='1st Order Error')
    plt.plot(data_dict["prob3"]["dx_values"], [err[1] for err in data_dict["prob3"]["case2"]["errors"]], label='2nd Order Error')
    plt.xscale('log')
    plt.yscale('log')
    plt.xlabel('Δx')
    plt.ylabel('% Error')
    plt.legend()
    plt.grid(True)
    plt.savefig("prob3_case2_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() # Convergence analysis