import os, sys, shutil
from pathlib import Path
import os
import matplotlib.pyplot as plt
from pymatgen.analysis.transition_state import NEBAnalysis

plt.figure(figsize=(8, 6))

neb_directory = "LPSCl_bulk_diffusion"
neb_analysis = NEBAnalysis.from_dir(neb_directory)
neb_analysis.setup_spline({"saddle_point": "zero_slope"})

# Extract reaction coordinates and energies
reaction_coords = neb_analysis.r
energies = neb_analysis.energies
rxn_coord = neb_analysis.get_extrema(normalize_rxn_coordinate=False)

print(rxn_coord)
# {“saddle_point”: “zero_slope”} 
ax = neb_analysis.get_plot(normalize_rxn_coordinate = False, label_barrier = True)

plt.title("NEB Energy Profile")
plt.xlabel("Reaction Coordinate")
plt.ylabel("Energy (eV)")
plt.grid(False)
plt.tight_layout()

plt.savefig("neb_energy_profile.pdf", bbox_inches="tight", dpi=300)