#!/usr/bin/env python3

import numpy as np
import os
import argparse
import matplotlib.pyplot as plt
import ase.units

plt.switch_backend('Agg')


if __name__=='__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument("-d", "--detail-file", default="./results",type=str, help="The file containing details of energy force and virial accuracy")
    args = parser.parse_args()

    folder=args.detail_file

    f=np.concatenate([np.loadtxt(os.path.join(folder,filename)) for filename in os.listdir(folder) if filename.endswith('f.out')])
    f*=ase.units.eV/(ase.units.kcal/ase.units.mol)

    #np.save("f.npy",f)
    #f=np.load("f.npy")
    #e=np.concatenate([np.loadtxt(os.path.join(folder,filename),ndmin=2) for filename in os.listdir(folder) if filename.endswith('e.out')])
    f_error=f[:,:3]-f[:,3:6]
    f_mae=np.mean(np.abs(f_error))
    f_rmse=np.sqrt(np.mean(np.square(f_error)))

    #e_rmse=np.sqrt(np.mean(np.square(e[:,:1]-e[:,1:2])))
    #print("No of energy: ",e.shape[0],"RMSE: ",e_rmse)
    print("No of force: ",f.shape[0],"MAE",f_mae,"RMSE: ",f_rmse)

    plt.scatter(f[:,:3].ravel(),f[:,3:6].ravel(),s=1)
    plt.xlabel('DFT forces (kcal/(mol·Å))')
    plt.ylabel('DL forces (kcal/(mol·Å))')

    lims = [
        np.min([plt.xlim(), plt.ylim()]),  # min of both axes
        np.max([plt.xlim(), plt.ylim()]),  # max of both axes
    ]

    # now plot both limits against eachother
    plt.plot(lims, lims, 'k-', alpha=0.75, color='r')
    plt.axes().set_aspect('equal')
    plt.xlim(lims)
    plt.ylim(lims)

    plt.text(lims[0]*0.1,lims[0]*0.9,f"MAE={f_mae:.4f} kcal/(mol·Å)\nRMSE={f_rmse:.4f} kcal/(mol·Å)")

    plt.axes([0.3, 0.6, 0.2, 0.2])
    plt.scatter(f[:,:3].ravel(),f[:,3:6].ravel(),s=1)
    plt.plot([-250,250],[-250,250],'k-', alpha=0.75, color='r')
    plt.xlim([-250,250])
    plt.ylim([-250,250])


    plt.savefig("f.png")