#!/usr/bin/env amspython
import argparse
import os
import scm.plams as plams
from collections import Counter


def main():
    parser = argparse.ArgumentParser(description="Process a path to an existing directory.")

    parser.add_argument("directory", type=str, help="Path to an existing directory", metavar="AMS_RESULTS_DIRECTORY")

    args = parser.parse_args()

    if not os.path.isdir(args.directory):
        print(f"The path {args.directory} is not a valid directory.")
        exit(1)

    job = plams.AMSJob.load_external(args.directory)

    # determine number of possible CN bonds from first frame
    input_system = job.results.get_input_molecule()

    counter = Counter((bond.atom1.symbol, bond.atom2.symbol) for bond in input_system.bonds)
    max_new_CN_bonds = counter[("N", "H")] + counter[("H", "N")]
    pre_existing_CN_bonds = counter[("N", "C")] + counter[("C", "N")]
    assert max_new_CN_bonds > 0, "No N-H bonds found in input system"

    traj = plams.RKFTrajectoryFile(job.results.rkfpath())
    mol = traj.get_plamsmol()
    n_entries = len(traj)

    # loop over trajectory
    for i, (_, _) in enumerate(traj(mol)):
        counter = Counter((bond.atom1.symbol, bond.atom2.symbol) for bond in mol.bonds)
        CN_bonds = counter[("N", "C")] + counter[("C", "N")]
        cross_link_density = (CN_bonds - pre_existing_CN_bonds) / max_new_CN_bonds
        print("({:4d}/{:4d}) {:3d} {:.3f}".format(i, n_entries, CN_bonds, cross_link_density))

    density = job.results.get_main_molecule().get_density()
    print(f"Final density: {density:.3f} kg/m^3")
    print(f"Cross linking ratio: {cross_link_density}")


if __name__ == "__main__":
    main()