import os
import matplotlib.pyplot as plt
import numpy as np

plt.rcParams.update({
    'font.size': 15,              # Default font size for all text
    'axes.labelsize': 15,         # Font size for axis labels
    'xtick.labelsize': 15,        # Font size for x-axis tick labels
    'ytick.labelsize': 15,        # Font size for y-axis tick labels
    'legend.fontsize': 10         # Font size for legend text
    })

def get_current_toughness_dict(filepaths_dict : dict):
    output_dict = {}
    for grid_type, path in filepaths_dict.items():
        output_dict[grid_type] = {}
        for crack_lenght in range(10,60,10):
            filepath = path + f"/a_{crack_lenght}/strength"
            with open(filepath, 'r') as file:
                data_dict = {}
                data_list = []
                data_list = [line.rstrip() for line in file]
                data_list = [line.split(" ") for line in data_list]
                data_dict["ipavg"] = float(data_list[0][0])
                data_dict["ipsig"] = float(data_list[0][1])
                data_dict["tavg"] = float(data_list[1][0])
                data_dict["tsig"] = float(data_list[1][1])
                output_dict[grid_type][crack_lenght] = data_dict
    return output_dict

def plot_current_and_toughness(filepaths_dict : dict):
    data_dict = get_current_toughness_dict(filepaths_dict)             
    fig, axs = plt.subplots(1,2) 
    fig.set_figwidth(10)
    fig.tight_layout(pad=2.5)
    
    hierarchical_list = []
    random_reference_list = []
    for crack_length in range(10,60,10):
        hierarchical_list.append([crack_length, data_dict["hierarchical"][crack_length]["ipavg"], data_dict["hierarchical"][crack_length]["ipsig"], 
                                 data_dict["hierarchical"][crack_length]["tavg"], data_dict["hierarchical"][crack_length]["tsig"]])
        random_reference_list.append([crack_length, data_dict["random equal density per layer"][crack_length]["ipavg"], data_dict["random equal density per layer"][crack_length]["ipsig"], 
                                 data_dict["random equal density per layer"][crack_length]["tavg"], data_dict["random equal density per layer"][crack_length]["tsig"]])
    hierarchical_array = np.array(hierarchical_list)
    random_reference_array = np.array(random_reference_list)
    print(hierarchical_array[:,0])

    axs[0].errorbar(hierarchical_array[:,0], hierarchical_array[:,1], yerr = hierarchical_array[:,2],
                    fmt='x', color = 'black')
    axs[0].errorbar(random_reference_array[:,0], random_reference_array[:,1], yerr = random_reference_array[:,2],
                   fmt='o', color = 'grey')
    axs[1].errorbar(hierarchical_array[:,0], hierarchical_array[:,3], yerr = hierarchical_array[:,4],
                    fmt='x', color = 'black', label = "hierarchical")
    axs[1].errorbar(random_reference_array[:,0], random_reference_array[:,3], yerr = random_reference_array[:,4],
                   fmt='o', color = 'grey', label = "random same density per layer")
    
    axs[0].set_xlabel("Crack length")
    axs[1].set_xlabel("Crack length")
    axs[0].set_ylabel("Peak current")
    axs[1].set_ylabel("Toughness")
    axs[1].legend()
    plt.show()
    fig.savefig("current_toughness_crack_lenghts.png")

if __name__ == "__main__":
    base_path = "/FASTTEMP/p7/lpyka/hierarchical_interface"
    os.chdir(base_path)
    filepaths_dict = {"hierarchical": base_path + "/5_precracked", "random equal density per layer": base_path + "/6_precrack_density_offset"}
    plot_current_and_toughness(filepaths_dict=filepaths_dict)


