#!/usr/bin/env amspython
from scm.plams import AMSJob, Units, init, finish
import numpy as np
import sys


def get_mean_property(job, prop, window_size, history_section):
    original_list = job.results.get_history_property(prop, history_section)  # a 1d list
    max_index = window_size * (len(original_list) // window_size)
    trimmed_list = original_list[:max_index]  # trim list to make sure the length is divisible by window_size
    two_d = np.array(trimmed_list).reshape(-1, window_size)  # an nXwindow_size 2D array
    ret = np.mean(two_d, axis=1)
    return ret


def main():
    # Make sure this script was called correctly
    args = sys.argv
    if len(args) != 2:
        print("\n USAGE: $AMSBIN/amspython densities.py [path to .results folder]\n")
        quit()

    init()

    # Load the job associated with that .results folder
    job = AMSJob.load_external(args[1])

    # Average how many frames
    window_size = 10

    mean_T = get_mean_property(job, "Temperature", window_size, "MDHistory")
    mean_density = get_mean_property(job, "Density", window_size, "MDHistory")
    mean_density *= (
        1e-3 * Units.convert(1.0, "au", "kg") * Units.convert(1.0, "bohr", "m") ** -3
    )  # from au/bohr^3 to g/cm^3

    print("Temperature_K Density_g/cm3")
    for t, dens in zip(mean_T, mean_density):
        print("{:4.1f} {:2.3f}".format(t, dens))


if __name__ == "__main__":
    main()