#!python

import os
import argparse
import numpy as np
import shutil
from vaspgibbs import __version__, read_outcar, read_poscar, prepare_incar, prepare_poscar, run_vasp, compute_thermo, reposition
from vaspgibbs.thermo import get_inertia, Itol, get_masses

def read_options():

    parser = argparse.ArgumentParser()
    parser.add_argument("-n","--ncores",dest="ncores",type=int, default=1, help="Number of cores for vasp run")
    parser.add_argument("-c","--mpi-command",dest="command",type=str, default="srun", help="Mpi execution command e.g. mpirun, mpiexec, srun")
    parser.add_argument("-v","--vasp",dest="vasp",type=str, default="vasp_std", help="Which vasp executable to run")
    parser.add_argument("-t","--top",dest="top",type=int, default=0, help="Number of atoms to be considered from the top of the c axis for vibrationals modes")
    parser.add_argument("-o","--only",dest="list_atoms",type=str, nargs='+', default=[], help="Specific atoms to consider fro vibrational modes")
    parser.add_argument("-b","--ibrion",dest="ibrion",type=int, default=5, help="IBRION value", choices=[5, 6, 7, 8])
    parser.add_argument("-T","--temperature",dest="T",type=float, default=298.15, help="Temperature of the system")
    parser.add_argument("-P","--pressure",dest="P",type=float, default=101.3, help="Pressure of the system")
    parser.add_argument("-m","--molecule",dest="mol",action="store_true", default=False, help="Use this flag if the sytem is a molecule")
    parser.add_argument("-g","--gen-only",dest="gen_only",action="store_true", default=False, help="Do not run VASP only generate input files")
    parser.add_argument("--version",dest="version",action="store_true", default=False, help="Display the version of VaspGibbs and stop")

    args = parser.parse_args()

    return args

def print_results(vgout, linear, T, P, E_dft, G, H, S, E_zpe, elec, vib, rot=None, trans=None):
    print("\n# Output\n", file=vgout)

    print("## System properties", file=vgout)
    if rot is not None and trans is not None:
        print("*This system is a {:}molecule*".format("linear " if linear else ""), file=vgout)
    s_char="| {:^16} | {:^23} |"
    f_char="| {:^16} | {:^ 14.7g} {:^8} |"
    print(s_char.format("Property", "Value"), file=vgout)
    print(s_char.format(":--------------:", ":---------------------:"), file=vgout)
    print(f_char.format("DFT Total Energy", E_dft, "eV"), file=vgout)
    print(f_char.format("Temperature", T, "K"), file=vgout)
    print(f_char.format("Pressure", P, "KPa"), file=vgout)
    if rot is not None:
        print(f_char.format("Sigma", rot.sigma, ""), file=vgout)
        print(s_char.format("**P. axes**", ""), file=vgout)
        print(f_char.format("I~1", rot.I[0], "eV/THz^2"), file=vgout)
        print(f_char.format("I~2", rot.I[1], "eV/THz^2"), file=vgout)
        print(f_char.format("I~3", rot.I[2], "eV/THz^2"), file=vgout)
    print(file=vgout)
    
    print("## Energy corrections", file=vgout)
    s_char="| {:^14} | {:^14} | {:^14} | {:^14} | {:^14} |"
    f_char="| {:^14} | {:^ 14.7g} | {:^ 14.7g} | {:^ 14.7g} | {:^ 14.7g} |"
    print(s_char.format("Type", "Z", "E (eV)", "S (eV/K)", "F (eV)"), file=vgout)
    print(s_char.format(*[":------------:"]*5), file=vgout)
    print(s_char.format("ZPE", "N/A", "{:^ 14.7g}".format(E_zpe), "N/A", "N/A"), file=vgout)
    print(f_char.format("Electronic", elec.Z, elec.E, elec.S, elec.E - T*elec.S), file=vgout)
    print(f_char.format("Vibrational", vib.Z, vib.E, vib.S, vib.E - T*vib.S), file=vgout)
    if rot is not None and trans is not None:
        print(f_char.format("Rotational", rot.Z, rot.E, rot.S, rot.E - T*rot.S), file=vgout)
        print(f_char.format("Translational", trans.Z, trans.E, trans.S, trans.E - T*trans.S), file=vgout)
    print(file=vgout)
    
    print("## Thermodynamic Quantities", file=vgout)
    
    s_char="| {:^17} | {:^19} |"
    f_char="| {:^17} | {:^ 14.7g} {:^4} |"
    print(s_char.format("Quantity", "Value"), file=vgout)
    print(s_char.format(":---------------:", ":-----------------:"), file=vgout)
    print(f_char.format("Enthalpy", H, "eV"), file=vgout)
    print(f_char.format("Entropy", S, "eV/K"), file=vgout)
    print(f_char.format("Gibbs Free Energy", G, "eV"), file=vgout)
    print(f_char.format("G - E_dft", G - E_dft, "eV"), file=vgout)
    print(f_char.format("TS", T*S, "eV"), file=vgout)

def main():

    args = read_options()

    if args.version:
       print("VaspGibbs", __version__)
       return

    vgout = open("VaspGibbs.md","w")  

    print("# VaspGibbs BETA\n", file=vgout)
    print("## Parameters:", file=vgout)
    for arg in vars(args):
        print(" * ", arg, ":", vars(args)[arg], file=vgout)
    print(file=vgout, flush=True)

    # Read the old outcar
    success, ibrion, freq, E_dft = read_outcar()

    cell, atoms = read_poscar()

    # If frequencies have been calculated, use them
    if ibrion in [5,6,7,8]:
        print("*Vibrational frequencies have been calculated, using existing files.*", file=vgout, flush=True)

    elif args.mol and len(atoms) == 1:
        freq = np.array([])

    else:
        if not success:
            print("> **Warning**: VaspGibbs must start from a completed calculation. From your OUTCAR, it seems like the calculation did not properly finish.", file=vgout, flush=True)
    
        prepare_incar(args.ibrion)

        if not os.path.isfile("POSCAR.save"):
            shutil.copyfile("POSCAR", "POSCAR.save")
        if ibrion != -1 and ibrion is not None:
            print("> **Warning**: starting from non-static run (IBRION = %2d)"%(ibrion), file=vgout, flush=True)
            print("> VaspGibbs will copy the CONTCAR to POSCAR", file=vgout, flush=True)
            shutil.copyfile("CONTCAR", "POSCAR")
            reposition()
            cell, atoms = read_poscar()
    
        cell, atoms = prepare_poscar(cell, atoms, args.list_atoms, args.top)

        if not os.path.isfile("OUTCAR.save"):
            shutil.copyfile("OUTCAR", "OUTCAR.save")

        if args.gen_only:
            print(file=vgout)
            print("VASP inputs generated.", file=vgout, flush=True)
            return

        run_vasp(args.command, args.ncores, args.vasp)

        # Reading the new outcar
        success, ibrion, freq, E_dft = read_outcar()

    I_mat = get_inertia(cell, atoms, get_masses(atoms))

    linear = False
    if args.mol:
        if len(atoms) == 1:
            freq = []
        if np.linalg.det(I_mat) < Itol*np.max(I_mat)**3:
            freq = freq[:3*len(atoms)-5]
            linear = True
        else:
            freq = freq[:3*len(atoms)-6]

    if freq is None:
        raise RuntimeError("Could not find frequencies in OUTCAR")

    for i,f in enumerate(freq):
        if not np.isreal(f):
            print("> **Warning**: frequency %d is imaginary (i*%s). Your structure may not be properly relaxed. Ignoring this frequency."%(i, np.imag(f)), file=vgout, flush=True)
            

    results = compute_thermo(args.T, args.P, freq, E_dft, cell, atoms, args.mol)

    print_results(vgout, linear, args.T, args.P, E_dft, *results)

main()
