#!/usr/bin/env amspython
import sys

from scm.plams import AMSJob, Units


def main():
    if len(sys.argv) != 2:
        print(f"\nUsage: {sys.argv[0]} <path to .results folder>\n")
        sys.exit(1)

    resultsdir = sys.argv[1]

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

    # Get energy and number of atoms of the initial state (before adding any Li)
    initEnergy = gcmc.results.readrkf("GCMC", "InitialEnergy")
    initNumAtoms = len(gcmc.results.get_input_molecule())

    # Get chemical potential of Li
    muLi = gcmc.results.readrkf("GCMC", "Mol1.chemPot")
    # Get Faraday's constant in atomic units
    ha_to_v = Units.convert(1, "Hartree", "kcal/mol") / 23.06
    # (Note: 23.06 is Faraday's constant in kcal per volt gram equivalent)

    # Loop over entries in the History and make dictionaries for the voltages and
    # volumes, using the number of Li atoms as the keys to the dictionaries.
    # (Note: We start at history step 2, because the first one is the initial
    # system without any Li.)
    voltages = {}
    volumes = {}
    for iHistEntry in range(2, gcmc.results.readrkf("History", "nEntries") + 1):
        # Get the geometry for the accepted MC step
        iSys = gcmc.results.get_history_molecule(iHistEntry)

        # Find out how many Li atoms it has
        numLiAtoms = sum(atom.symbol == "Li" for atom in iSys.atoms)
        if numLiAtoms == 0:
            continue

        totalEnergy = gcmc.results.readrkf("History", f"Energy({iHistEntry:d})")
        voltage = -ha_to_v * (totalEnergy - initEnergy - numLiAtoms * muLi) / numLiAtoms
        if numLiAtoms not in voltages:
            voltages[numLiAtoms] = []
        voltages[numLiAtoms].append(voltage)

        volume = iSys.unit_cell_volume()
        if numLiAtoms not in volumes:
            volumes[numLiAtoms] = []
        volumes[numLiAtoms].append(volume)

    # Calculate and print averages for voltages and volumes.
    print("# Nr. Li atoms, fraction of Li-atoms, Voltage in V, Volume in Angstrom**3")
    for numLiAtoms in voltages:
        voltage = sum(voltages[numLiAtoms]) / len(voltages[numLiAtoms])
        volume = sum(volumes[numLiAtoms]) / len(volumes[numLiAtoms])
        print(
            "{:3d} {:2.3f} {:2.3f} {:2.2f}".format(
                numLiAtoms, numLiAtoms / initNumAtoms, voltage, volume
            )
        )


if __name__ == "__main__":
    main()