from os import makedirs
from pathlib import Path

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_timing(indir, outdir):
    if not (path := Path(outdir)).exists():
        makedirs(path)

    files = [
        "lattice_sizes.txt",
        "sample_sizes.txt",
    ]
    labels = [
        "Lattice sizes",
        "Sample sizes",
    ]
    xlabels = ["Lattice size", "Sampling size"]
    outfiles = ["lattice_size.pdf", "sample_sizes.pdf"]

    for file, label, xlabel, outfile in zip(files, labels, xlabels, outfiles):
        figure1, ax1 = plt.subplots()
        x = []
        t = []

        with open(Path(indir, file)) as f:
            lines = f.readlines()
            for line in lines:
                l = line.strip().split(",")
                x.append(float(l[0]))
                t.append(float(l[1]))

        ax1.plot(x, t, label=label)
        ax1.set_xlabel(xlabel)
        ax1.set_ylabel("time (seconds)")

        ax1.legend()

        figure1.savefig(Path(outdir, outfile), bbox_inches="tight")

        plt.close(figure1)


if __name__ == "__main__":
    plot_timing(
        "data/hp/timing/",
        "./latex/images/timing",
    )