import os
from pathlib import Path

import matplotlib.pyplot as plt
import numpy as np
import seaborn as sns
from scipy.stats import norm

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(infile, outfile):
    if not (outdir := Path(outfile).parent).exists():
        os.makedirs(outdir)

    figure1, ax1 = plt.subplots()
    eps = []
    with open(infile) as f:
        lines = f.readlines()
        for line in lines:
            vals = line.strip().split(",")
            eps.append(float(vals[0]))

    ax1.hist(eps, np.arange(min(eps), max(eps) + 0.02, 0.02), color="steelblue", density=True, ec="white", lw=0.2)
    ax1.set_xlabel(r"$\epsilon$")
    ax1.set_ylabel("Density")

    mu, sigma = np.mean(eps), np.std(eps) 
    x = np.arange(min(eps), max(eps) + 0.02, 0.02)

    stats = (f"$\\langle \\epsilon \\rangle$ = {np.mean(eps):.4f}\n"
             f"Var$(\\epsilon)$ = {np.var(eps):.4f}")
    bbox = dict(boxstyle="round", pad=0.3, fc="white", ec="white", alpha=0.5)
    ax1.text(0.95, 0.9, stats, bbox=bbox, transform=ax1.transAxes, ha="right", va="center")
    ax1.plot(x, norm.pdf(x, mu, sigma), color="mediumseagreen")
    figure1.savefig(outfile, bbox_inches="tight")


if __name__ == "__main__":
    plot("output/pd_estimate/estimate_1_0.txt", "../latex/images/pd_estimate_1_0.pdf")
    plot("output/pd_estimate/estimate_2_4.txt", "../latex/images/pd_estimate_2_4.pdf")
