#!/usr/bin/env amspython

import argparse
from scm.plams import *
import matplotlib.pyplot as plt
import numpy as np

"""

Analyze a GCMC run and plot the number of guest molecule as a function of GCMC steps.

Usage:

$AMSBIN/amspython -u analyze_gcmc.py ams.rkf [ams2.rkf ...] 

"""

ap = argparse.ArgumentParser()
ap.add_argument("results", type=str, help="AMS results directory paths", nargs="+")
ap.add_argument("--nevery", type=int, default=1, help="Will parse rkf every nevery frame")
args = ap.parse_args()

every = args.nevery

print(args.results)
Nres = len(args.results)
print("{:d} files to parse..".format(Nres))

jobs = []
nMC, history_list, history_move = [], [], []
for i in range(Nres):

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

    # Get energy and number of atoms of the initial state (before adding any Li)
    NIterMCtried = job.results.readrkf("GCMC", "NIterMCtried")
    HistoryAccepted = job.results.readrkf("GCMC", "HistoryAccepted")
    HistoryMoveType = job.results.readrkf("GCMC", "HistoryMoveType")
    #    HistoryAccepted = [hi for hi in HistoryAccepted if hi != -1]
    nMC.append(NIterMCtried)
    history_list.append(HistoryAccepted)
    history_move.append(HistoryMoveType)

    nguest = job.results.readrkf("GCMC", "Mol1.nAtoms")

    initNumAtoms = len(job.results.get_input_molecule())
    masses = job.results.readrkf("InputMolecule", "AtomMasses")
    totalmass = sum(masses)

    jobs.append(job)

print(
    "Initial molecule {:6.3f} g/mol ({:d} atoms) and {:d} atoms in guest molecule".format(
        totalmass, initNumAtoms, nguest
    )
)

NIterMCtried = min(nMC)

istep = [0] * Nres
load = 0.0

load_dict = {}
accept_dict = {}
for j in range(Nres):
    load_dict[j] = []
    accept_dict[j] = []

for i in np.arange(0, NIterMCtried, every):

    line = "MC step {:^9d}".format(i)

    for j in range(Nres):

        # index for the geometry corresponds to the sum of all accepted
        # because mol history only included when mol is modified i.e. an accepted move
        istep[j] = sum(history_list[j][: i + 1])

        # load
        # get the geometry for the accepted MC step
        if istep[j] == 0:
            iSys = job.results.readrkf("InputMolecule", "AtomMasses")
            natoms = len(iSys)
        else:
            iSys = jobs[j].results.get_history_molecule(istep[j])
            natoms = len(iSys.atoms)
        n_mol_guest = (natoms - initNumAtoms) / nguest
        load = 1e3 * n_mol_guest / totalmass

        #        print(natoms,initNumAtoms,nguest,history_list[j][i],history_move[j][i])

        # acceptance ratio
        # number of accepted insert / total number of MC moves
        n_accept = 0
        for k in range(i + 1):
            if history_list[j][k] == 1 and history_move[j][k] == 0:
                n_accept += 1
        accept = n_accept / (i + 1)

        # printing
        line += " astep {:^9d} {:7.5f} {:7.3f} mol/kg | ".format(istep[j], accept, load)

        load_dict[j].append(load)
        accept_dict[j].append(accept)

    print(line)

for j in range(Nres):
    ngcmc2 = int(len(load_dict[j]) / 2)
    load_ave = np.mean(load_dict[j][ngcmc2:])
    load_std = np.std(load_dict[j][ngcmc2:])
    print("File {:s} {:6.3f} +/- {:6.3f}".format(args.results[j], load_ave, load_std))

    fig, ax1 = plt.subplots()
    ax2 = ax1.twinx()
    ax2.plot(np.array(range(len(load_dict[j]))) * every, load_dict[j], color="tab:blue")
    ax1.plot(np.array(range(len(accept_dict[j]))) * every, accept_dict[j], color="tab:red")
    ax2.set_ylabel("Load mol/kg", fontsize=16)
    ax1.set_ylabel("Acceptance ratio", fontsize=16)
    ax1.set_xlabel("MC step", fontsize=16)
    plt.tight_layout()
    plt.savefig(f"load_list{j}.png")
    fig.clear()

fig, ax1 = plt.subplots()
for j in range(Nres):
    ax1.plot(np.array(range(len(load_dict[j]))) * every, load_dict[j], label=args.results[j])
ax1.set_ylabel("Load mol/kg", fontsize=16)
ax1.set_xlabel("MC step", fontsize=16)
# ax1.legend()
plt.tight_layout()
plt.savefig(f"load_list.png")