#!/usr/bin/env amspython
import sys
import os
from scm.plams import *
import numpy as np

"""
Call this script using
$AMSBIN/amspython stress_strain_curve.py job_name

IMPORTANT: This only works for cubic or orthorhombic cells!

This script accompanies the tutorial at 
https://www.scm.com/doc/Tutorials/MolecularDynamicsAndMonteCarlo/PolymersMechanicalProperties.html
"""


def main():
    if len(sys.argv) != 2 or not os.path.isdir(sys.argv[1] + ".results"):
        print("\n USAGE: " + str(sys.argv[0]) + " job_name\n")
        exit(1)

    job_dir = sys.argv[1] + ".results"
    job = AMSJob.load_external(job_dir)

    lattice_vectors = job.results.get_history_property("LatticeVectors")
    lattice_vectors = np.array(lattice_vectors)
    lattice_vectors = lattice_vectors[:, (0, 4, 8)]  # assume orthorhombic, extract diagonal components
    # lattice_vectors now has shape (N, 3)

    strain = (lattice_vectors - lattice_vectors[0]) / lattice_vectors[0]

    stress = job.results.get_history_property("PressureTensor", "MDHistory")
    au_to_GPa = 29421.015209244444
    stress = -np.array(stress) * au_to_GPa  # change sign convert unit
    stress = stress[:, 0:3]  # first three columns (xx, yy, zz)
    # stress now has shape (N, 3)

    strain_stress = np.concatenate((strain, stress), axis=1)

    np.savetxt(
        "stress-strain-curve.csv",
        strain_stress,
        fmt="%.8f",
        header="strain_x strain_y strain_z stress_xx[GPa] stress_yy[GPa] stress_zz[GPa]",
        comments="",
    )

    print("Saved stress-strain-curve.csv")


if __name__ == "__main__":
    main()