#!/usr/bin/env plams
from functools import reduce

# Make sure this script was called correctly
if not "resultsdir" in globals():
    print("\n USAGE: $AMSBIN/plams new-cross-link-density.py -v resultsdir=[path to .results folder]\n")
    quit()

# Load the job associated with that .results folder
bondboost = AMSJob.load_external(resultsdir)

# determine number of possible CN bonds from first frame
max_CN_Bonds = 0
pre_existing_CN_Bonds = 0
iSys = bondboost.results.get_history_molecule(1)
iSys.guess_bonds()
for bond in iSys.bonds:
    if bond.atom1.symbol == "N" and bond.atom2.symbol == "H":
        max_CN_Bonds += 1
    elif bond.atom1.symbol == "H" and bond.atom2.symbol == "N":
        max_CN_Bonds += 1
    elif bond.atom1.symbol == "C" and bond.atom2.symbol == "N":
        pre_existing_CN_Bonds += 1
    elif bond.atom1.symbol == "N" and bond.atom2.symbol == "C":
        pre_existing_CN_Bonds += 1

nEntries = bondboost.results.readrkf("History", "nEntries")
for iHistEntry in range(2, nEntries + 1):
    # Get the geometry for the accepted MC step
    iSys = bondboost.results.get_history_molecule(iHistEntry)
    iSys.guess_bonds()
    # Find N atoms
    CN_Bonds = 0
    for bond in iSys.bonds:
        if bond.atom1.symbol == "C" and bond.atom2.symbol == "N":
            CN_Bonds += 1
        elif bond.atom1.symbol == "N" and bond.atom2.symbol == "C":
            CN_Bonds += 1
    cross_link_density = (CN_Bonds - pre_existing_CN_Bonds) / max_CN_Bonds
    print("({:4d}/{:4d}) {:2.3f} {:2.3f}".format(iHistEntry, nEntries, CN_Bonds, cross_link_density))

# compute density from lattice vectors and total mass, only for last frame
latticevecs = bondboost.results.readrkf("Molecule", "eeLatticeVectors")
latticevecs = list(filter(lambda a: a != 0, latticevecs))
volume = reduce((lambda x, y: x * y), latticevecs) * 0.14818  # convert Bohr**3 to Ang**3
masses = bondboost.results.readrkf("Molecule", "AtomMasses")
mass = sum(masses)
print(
    "\n  Final density: {:2.3f} \n  Cross-linking ratio: {:2.3f} \n".format(mass / volume * 1.6601, cross_link_density)
)