import numpy as np
import matplotlib.pyplot as plt
import ast
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():
    files = [
        "data/screen/single_slit.txt",
        "data/screen/double_slit.txt",
        "data/screen/triple_slit.txt",
    ]
    outputs = [
        "latex/images/single_slit_detector.pdf",
        "latex/images/double_slit_detector.pdf",
        "latex/images/triple_slit_detector.pdf",
    ]
    for file, output in zip(files, outputs):
        with open(file) as f:
            lines = f.readlines();
            size = int(lines[0])
            column = int(.8 * size)
            fig, ax = plt.subplots()
            arr = lines[1].strip().split("\t")
            arr = np.asarray(list(map(lambda x: complex(*ast.literal_eval(x)), arr)))

            arr = arr.reshape(size,size)
            x = np.linspace(0,1, size)
            slice = arr[column, :]
            norm = 1. / np.sqrt(sum([(i*i.conjugate()).real for i in slice]))
            slice *= norm
            slice = np.asarray([i * i.conjugate() for i in slice])

            ax.plot(x, slice)
            plt.savefig(output)
            plt.close(fig)


if __name__ == "__main__":
    plot()