#!/usr/bin/env amspython
"""
Usage: 
$AMSBIN/amspython viscosity_extractor.py /path/to/ams.rkf

Feel free to change the max_dt_fs to change the maximum correlation time!

The script will produce a file green_kubo_viscosity.txt and viscosity.pdf in the same directory as ams.rkf.
"""

from scm.plams import *
import numpy as np
from scipy.optimize import curve_fit
import matplotlib.pyplot as plt
import os
import sys


def main():
    if len(sys.argv) > 1:
        path_to_rkf = sys.argv[1]
    else:
        path_to_rkf = "ams.rkf"

    assert os.path.exists(path_to_rkf), f"{path_to_rkf} does not seem to exist. It should be an ams.rkf file."
    max_dt_fs = 200000
    outfile = os.path.join(os.path.dirname(os.path.abspath(path_to_rkf)), "green_kubo_viscosity.txt")
    compute_pressure_tensor_acf(path_to_rkf, outfile=outfile, max_dt_fs=max_dt_fs)

    assert os.path.exists(outfile), f"{outfile} does not seem to exist!"
    imagefile = os.path.join(os.path.dirname(outfile), "viscosity.pdf")
    viscosity = do_fit(outfile, imagefile=imagefile)

    print(f"Final viscosity: {viscosity:.3f} mPa s")


# The double exponential function
def viscosity_double_exponential(x, A, lam, tau1, tau2):
    return A * (lam * (1 - np.exp(-x / tau1)) + (1 - lam) * (1 - np.exp(-x / tau2)))


# The main function
def compute_pressure_tensor_acf(path_to_rkf, outfile, max_dt_fs=100000, discard_first_fs=0):

    # Change this to the path of your ams.rkf file
    job = AMSJob.load_external(path_to_rkf)

    print(f"Computing viscosity for {path_to_rkf}")
    print("Note: this may take several hours if the trajectory is long!")

    # Here we get the off-diagonal pressure tensor from binlog
    xy = job.results.get_history_property("PressureTensor_xy", history_section="BinLog")
    xz = job.results.get_history_property("PressureTensor_xz", history_section="BinLog")
    yz = job.results.get_history_property("PressureTensor_yz", history_section="BinLog")

    minN = min(len(xy), len(xz), len(yz))
    pressuretensor = np.stack((np.zeros(minN), np.zeros(minN), np.zeros(minN), yz[:minN], xz[:minN], xy[:minN]), axis=1)

    # Volume, temperature, time
    volume = job.results.get_main_molecule().unit_cell_volume()
    temperature = np.mean(job.results.get_history_property("Temperature", history_section="MDHistory"))
    times = job.results.get_history_property("Time", history_section="BinLog")
    timestep = times[1] - times[0]

    # We skip first 10000 frames
    discard_first_frames = int(discard_first_fs / timestep)
    pressuretensor = pressuretensor[discard_first_frames:, :]

    # Time interval to compute the ACF
    # We found that max_dt_fs = 100000 fs is a good value
    max_dt = int(max_dt_fs / timestep)

    # Some info to print
    print(f"Averaged temperature: {temperature}")
    print(f"Timestep: {timestep}")
    print(f"Total simulation time: {times[-1]} fs")
    print(f"Total #frames: {minN}")
    print(f"Discarding first {discard_first_fs:.2f} fs as equilibration period")
    print(f"Discarding first {discard_first_frames} frames as equilibration period")
    print(f"Maximum correlation time: {max_dt_fs:.2f} fs")
    print(f"Maximum correlation time: {max_dt} frames")
    print(f"Unit cell volume: {volume:.3f} ang^3")

    print(f"Calculating the pressure tensor autocorrelation function (this may take a while....)")

    # Compute the viscosity based on GK relation
    x, y = AMSResults._get_green_kubo_viscosity(
        pressuretensor=pressuretensor, max_dt=max_dt, time_step=timestep, volume=volume, temperature=temperature
    )

    # We save the data
    print(f"Writing cumulative integral of pressure tensor autocorrelation function to {outfile}")
    A = np.stack((x, y), axis=1)
    np.savetxt(outfile, A, header="Time(fs) Viscosity(mPa*s)")


def do_fit(fname, imagefile):
    A = np.loadtxt(fname)
    x = A[:, 0]
    y = A[:, 1]

    # The fit
    f, p0 = viscosity_double_exponential, (20.0, 0.5, 1000, 10000)
    popt, _ = curve_fit(f, x, y, p0=p0)
    prediction = f(x, *popt)

    # We plot the results
    plt.title(f"Limiting value: {popt[0]:.2f} mPa s")
    plt.plot(x, y, color="tab:blue")
    plt.plot(x, prediction, color="tab:red")
    plt.xlabel("Correlation Time (fs)")
    plt.ylabel("η (mPa s)")
    plt.savefig(imagefile)
    plt.show()

    return popt[0]


if __name__ == "__main__":
    main()