#!/usr/bin/env amspython

import sys
from functools import reduce
from typing import List

from scm.plams import KFHistory, KFReader

# read command line arguments
if len(sys.argv) != 3:
    print("\n USAGE: $AMSBIN/amspython strain.py [path to .results folder] [startstep]\n")
    sys.exit(1)
resultsdir = str(sys.argv[1])
startstep = int(sys.argv[2])

# use the PLAMS KF reader to read from the ams.rkf binary output file
kf = KFReader(f"{resultsdir}/ams.rkf")
mdhist = KFHistory(kf, "MDHistory")
hist = KFHistory(kf, "History")

# iterate through the frames of the trajectory
print("# T[K]   a[Å]   b[Å]   c[Å]    V[Å**3]")
for counter, T in enumerate(mdhist.iter("Temperature"), 1):
    step: int = kf.read("History", f"Step({counter})")
    latticevecs: List[float] = kf.read("History", f"LatticeVectors({counter})")
    latticevecs = list(filter(lambda a: a != 0, latticevecs))
    volume = reduce((lambda x, y: x * y), latticevecs)
    if step >= int(startstep):
        print(f"{T:6.1f} {latticevecs[0]:2.3f} {latticevecs[1]:2.3f} {latticevecs[2]:2.3f} {volume:5.3f}")