from os import makedirs
from pathlib import Path

import matplotlib.pyplot as plt
import numpy as np
from scipy.stats import linregress
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_phase_transition_alt(indir, outdir):
    if not (path := Path(outdir)).exists():
        makedirs(path)

    files = [
        "size_20.txt",
        "size_40.txt",
        "size_60.txt",
        "size_80.txt",
        "size_100.txt",
        "size_500.txt",
    ]
    labels = ["L = 20", "L = 40", "L = 60", "L = 80", "L = 100", "L = 500"]

    figure1, ax1 = plt.subplots()
    figure2, ax2 = plt.subplots()
    figure3, ax3 = plt.subplots()
    figure4, ax4 = plt.subplots()
    figure5, ax5 = plt.subplots()

    # For linear regression
    L = []
    Tc = []
    size = 20

    for file, label in zip(files, labels):
        t = []
        e = []
        m = []
        CV = []
        X = []

        # Append the lattice size
        L.append(size)
        size += 20

        with open(Path(indir, file)) as f:
            lines = f.readlines()
            for line in lines:
                l = line.strip().split(",")
                t.append(float(l[0]))
                e.append(float(l[1]))
                m.append(float(l[2]))
                CV.append(float(l[3]))
                X.append(float(l[4]))

        # Append the critical temp for the current lattice size
        Tc.append(t[X.index(max(X))])

        ax1.plot(t, e, label=label)
        ax2.plot(t, m, label=label)
        ax3.plot(t, CV, label=label)
        ax4.plot(t, X, label=label)

    inv_L = list(map(lambda x: 1 / x, L))
    # Attempt linear regression
    x = np.linspace(0, 1 / 20, 1001)
    regression = linregress(inv_L, Tc)
    f = lambda x: regression[0] * x + regression[1]
    stats = (
        f"$\\beta_{0}$ = {regression[1]:.4f}\n"
        f"$\\beta_{1}$ = {regression[0]:.4f}"
    )
    bbox = dict(boxstyle="round", pad=0.3, fc="white", ec="white", alpha=0.5)
    ax5.text(
        0.6, 0.6, stats, bbox=bbox, transform=ax1.transAxes, ha="right", va="center"
    )
    ax5.scatter(inv_L, Tc)
    ax5.plot(x, f(x))

    ax1.set_xlabel(r"T $(J/k_{B})$")
    ax1.set_ylabel(r"$\langle \epsilon \rangle$ $(J)$")
    ax1.legend(title="Lattice size", loc="upper right")

    ax2.set_xlabel(r"T $(J/k_{B})$")
    ax2.set_ylabel(r"$\langle |m| \rangle$ $(unitless)$")
    ax2.legend(title="Lattice size", loc="upper right")

    ax3.set_xlabel(r"T $(J/k_{B})$")
    ax3.set_ylabel(r"$C_{V}$ $(k_{B})$")
    ax3.legend(title="Lattice size", loc="upper right")

    ax4.set_xlabel(r"T $(J/k_{B})$")
    ax4.set_ylabel(r"$\chi$ $(1/J)$")
    ax4.legend(title="Lattice size", loc="upper right")

    # ax5.legend()

    figure1.savefig(Path(outdir, "energy.pdf"), bbox_inches="tight")
    figure2.savefig(Path(outdir, "magnetization.pdf"), bbox_inches="tight")
    figure3.savefig(Path(outdir, "heat_capacity.pdf"), bbox_inches="tight")
    figure4.savefig(Path(outdir, "susceptibility.pdf"), bbox_inches="tight")
    figure5.savefig(Path(outdir, "linreg.pdf"), bbox_inches="tight")

    plt.close(figure1)
    plt.close(figure2)
    plt.close(figure3)
    plt.close(figure4)
    plt.close(figure5)


def plot_phase_transition(indir, outdir):
    if not (path := Path(outdir)).exists():
        makedirs(path)

    files = [
        "size_20.txt",
        "size_40.txt",
        "size_60.txt",
        "size_80.txt",
        "size_100.txt",
    ]
    labels = [
        "20",
        "40",
        "60",
        "80",
        "100",
    ]

    figure1, ax1 = plt.subplots()
    figure2, ax2 = plt.subplots()
    figure3, ax3 = plt.subplots()
    figure4, ax4 = plt.subplots()
    figure5, ax5 = plt.subplots()

    # For linear regression
    L = []
    Tc = []
    size = 20

    for file, label in zip(files, labels):
        t = []
        e = []
        m = []
        CV = []
        X = []

        # Append the lattice size
        L.append(size)
        size += 20

        with open(Path(indir, file)) as f:
            lines = f.readlines()
            for line in lines:
                l = line.strip().split(",")
                t.append(float(l[0]))
                e.append(float(l[1]))
                m.append(float(l[2]))
                CV.append(float(l[3]))
                X.append(float(l[4]))

        # Append the critical temp for the current lattice size
        Tc.append(t[X.index(max(X))])

        ax1.plot(t, e, label=label)
        ax2.plot(t, m, label=label)
        ax3.plot(t, CV, label=label)
        ax4.plot(t, X, label=label)

    inv_L = list(map(lambda x: 1 / x, L))
    # Attempt linear regression
    x = np.linspace(0, 1 / 20, 1001)
    regression = linregress(inv_L, Tc)
    f = lambda x: regression[0] * x + regression[1]
    stats = (
        f"$\\beta_{0}$ = {regression[1]:.4f}\n"
        f"$\\beta_{1}$ = {regression[0]:.4f}"
    )
    bbox = dict(boxstyle="round", pad=0.3, fc="white", ec="white", alpha=0.5)
    ax5.text(
        0.6, 0.6, stats, bbox=bbox, transform=ax1.transAxes, ha="right", va="center"
    )
    ax5.scatter(inv_L, Tc)
    ax5.plot(x, f(x))

    ax1.set_xlabel(r"T $(J/k_{B})$")
    ax1.set_ylabel(r"$\langle \epsilon \rangle$ $(J)$")
    ax1.legend(title="Lattice size", loc="upper right")

    ax2.set_xlabel(r"T $(J/k_{B})$")
    ax2.set_ylabel(r"$\langle |m| \rangle$ $(unitless)$")
    ax2.legend(title="Lattice size", loc="upper right")

    ax3.set_xlabel(r"T $(J/k_{B})$")
    ax3.set_ylabel(r"$C_{V}$ $(k_{B})$")
    ax3.legend(title="Lattice size", loc="upper right")

    ax4.set_xlabel(r"T $(J/k_{B})$")
    ax4.set_ylabel(r"$\chi$ $(1/J)$")
    ax4.legend(title="Lattice size", loc="upper right")

    # print(Tc)

    figure1.savefig(Path(outdir, "energy.pdf"), bbox_inches="tight")
    figure2.savefig(Path(outdir, "magnetization.pdf"), bbox_inches="tight")
    figure3.savefig(Path(outdir, "heat_capacity.pdf"), bbox_inches="tight")
    figure4.savefig(Path(outdir, "susceptibility.pdf"), bbox_inches="tight")
    figure5.savefig(Path(outdir, "linreg.pdf"), bbox_inches="tight")

    plt.close(figure1)
    plt.close(figure2)
    plt.close(figure3)
    plt.close(figure4)
    plt.close(figure5)


if __name__ == "__main__":
    plot_phase_transition(
        "data/fox/phase_transition/wide/10M/",
        "latex/images/phase_transition/fox/wide/10M/",
    )
    plot_phase_transition(
        "data/fox/phase_transition/wide/1M/",
        "latex/images/phase_transition/fox/wide/1M/",
    )
    plot_phase_transition(
        "data/fox/phase_transition/narrow/10M/",
        "latex/images/phase_transition/fox/narrow/10M/",
    )
    plot_phase_transition(
        "data/hp/phase_transition/", "latex/images/phase_transition/hp/"
    )
    plot_phase_transition(
        "data/hp/phase_transition/", "latex/images/phase_transition/hp/"
    )
    plot_phase_transition(
        "data/hp/phase_transition/",
        "latex/images/phase_transition/hp/",
    )