import numpy as np

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
    if not 'resultsdir' in globals():
        print("\n USAGE: $AMSBIN/plams densities.py -v resultsdir=[path to .results folder]\n")
        quit()

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

    # 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()