##  @file plot_3d.py
#   
#   @author Cory Alexander Balaton (coryab)
#   @author Janita Ovidie Sandtrøen Willumsen (janitaws)
#
#   @version 1.0
#
#   @brief Plot 2 particles with and without particle interactions in 3D.
#
#   @bug No known bugs

import matplotlib.pyplot as plt
import numpy as np
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",
    "figure.autolayout": True
}
plt.rcParams.update(params)

def main():
    files = [
        "output/simulate_2_particles/no_interaction/particle_0_r.txt",
        "output/simulate_2_particles/no_interaction/particle_1_r.txt",
        "output/simulate_2_particles/with_interaction/particle_0_r.txt",
        "output/simulate_2_particles/with_interaction/particle_1_r.txt"
    ]
    labels = [
        r"$p_1$",
        r"$p_2$",
        r"$\hat{p}_1$",
        r"$\hat{p}_2$",
    ]
    colors = [
        "mediumaquamarine",
        "salmon",
        "seagreen",
        "darkred"
    ]
    ax = plt.figure().add_subplot(projection="3d")
    for label, color, file in zip(labels, colors, files):
        with open(file) as f:
            lines = f.readlines()
            t = np.linspace(0, 50, len(lines))
            r = np.array([list(map(float, line.strip().split(","))) for line in lines])
            ax.plot(r[:,0], r[:,1], r[:,2], label=label, color=color)
    
    ax.set_xlabel(r"x $(\mu m)$")
    ax.set_ylabel(r"y $(\mu m)$")
    ax.set_zlabel(r"z $(\mu m)$")
    ax.view_init(10,-35)

    plt.legend()
    plt.savefig("../latex/images/3d_plot.pdf")


if __name__ == "__main__":
    main()