import matplotlib.pyplot as plt
import seaborn as sns

sns.set_theme()
params = {
    "font.family": "Serif",
    "font.serif": "Roman",
    "text.usetex": True,
    "axes.titlesize": "large",
    "axes.labelsize": "large",
    "xtick.labelsize": "large",
    "ytick.labelsize": "large",
    "legend.fontsize": "medium",
}
plt.rcParams.update(params)


def plot_1_0():
    files = [
        # "data/hp/burn_in_time/unordered_1_0_1421110368.txt",
        # "data/hp/burn_in_time/ordered_1_0_611577739.txt",
        "data/hp/burn_in_time/unordered_1_0.txt",
        "data/hp/burn_in_time/ordered_1_0.txt",
    ]
    labels = [
        "Unordered",
        "Ordered",
    ]
    colors = [
        "darkred",
        "seagreen",
    ]

    fig1, ax1 = plt.subplots()
    fig2, ax2 = plt.subplots()

    for infile, label, color in zip(files, labels, colors):
        with open(infile) as f:
            lines = f.readlines()

            t = []
            energy = []
            magnetization = []
            for line in lines:
                items = line.strip().split(",")
                t.append(int(items[0]))
                energy.append(float(items[1]))
                magnetization.append(float(items[5]))

            ax1.plot(t, energy, label=f"{label}", color=color)
            ax2.plot(t, magnetization, label=f"{label}", color=color)

    ax1.set_xlabel("t (MC cycles)")
    ax1.set_ylabel(r"$\langle \epsilon \rangle$ $(J)$")
    ax1.legend(title="Initial state", loc="upper right")
    fig1.savefig("./latex/images/burn_in_time_energy_1_0.pdf", bbox_inches="tight")

    ax2.set_ylabel(r"$\langle |m| \rangle$ (unitless)")
    ax2.set_xlabel("t (MC cycles)")
    ax2.legend(title="Initial state", loc="upper right")
    fig2.savefig(
        "./latex/images/burn_in_time_magnetization_1_0.pdf", bbox_inches="tight"
    )


def plot_2_4():
    files = [
        # "data/hp/burn_in_time/unordered_2_4_1212892317.txt",
        # "data/hp/burn_in_time/ordered_2_4_2408603856.txt",
        "data/hp/burn_in_time/unordered_2_4.txt",
        "data/hp/burn_in_time/ordered_2_4.txt",
    ]
    labels = [
        "Unordered",
        "Ordered",
    ]
    colors = [
        "darkred",
        "seagreen",
    ]

    fig1, ax1 = plt.subplots()
    fig2, ax2 = plt.subplots()

    for infile, label, color in zip(files, labels, colors):
        with open(infile) as f:
            lines = f.readlines()

            t = []
            energy = []
            magnetization = []
            for line in lines:
                items = line.strip().split(",")
                t.append(int(items[0]))
                energy.append(float(items[1]))
                magnetization.append(float(items[5]))

            ax1.plot(t, energy, label=f"{label}", color=color)
            ax2.plot(t, magnetization, label=f"{label}", color=color)

    ax1.set_xlabel("t (MC cycles)")
    ax1.set_ylabel(r"$\langle \epsilon \rangle$ $(J)$")
    ax1.legend(title="Initial state", loc="upper right")
    fig1.savefig("./latex/images/burn_in_time_energy_2_4.pdf", bbox_inches="tight")

    ax2.set_ylabel(r"$\langle |m| \rangle$ (unitless)")
    ax2.set_xlabel("t (MC cycles)")
    ax2.legend(title="Initial state", loc="upper right")
    fig2.savefig(
        "./latex/images/burn_in_time_magnetization_2_4.pdf", bbox_inches="tight"
    )


if __name__ == "__main__":
    plot_1_0()
    plot_2_4()