import numpy as np
import os, shutil

from pymatgen.io.vasp.inputs import Poscar
from pymatgen.io.vasp.outputs import Vasprun
from pymatgen.io.vasp.sets import VaspInputSet
from pymatgen.core.periodic_table import Element

from subprocess import call
from flare.struc import Structure
from typing import List, Union
from json import dump, load
from flare.util import NumpyEncoder

name="VASP"

def check_vasprun(vasprun: Union[str, Vasprun], vasprun_kwargs: dict = {}) -> \
        Vasprun:
    """
    Helper utility to take a vasprun file name or Vasprun object
    and return a vasprun object.
    :param vasprun: vasprun filename or object
    """

    if type(vasprun) == str:
        return Vasprun(vasprun, **vasprun_kwargs)
    elif type(vasprun) == Vasprun:
        return vasprun
    else:
        raise ValueError('Vasprun argument is not a string or '
                         'Vasprun instance!')


def parse_dft_input(input: str):
    """
    Returns the positions, species, and cell of a POSCAR file.
    Outputs are specced for OTF module.

    :param input: POSCAR file input
    :return:
    """

    pmg_structure = Poscar.from_file(input).structure
    flare_structure = Structure.from_pmg_structure(pmg_structure)

    positions = flare_structure.positions
    species = flare_structure.species_labels
    cell = flare_structure.cell
    #TODO Allow for custom masses in POSCAR

    elements = set(species)

    # conversion from amu to md units
    mass_dict = {elt: Element(elt).atomic_mass * 0.000103642695727 for
                 elt in elements}

    return positions, species, cell, mass_dict

def run_dft(calc_dir: str, dft_loc: str,
            structure: Structure = None, en: bool = False, vasp_cmd="{}"):
    """
    Run a VASP calculation.
    :param calc_dir: Name of directory to perform calculation in
    :param dft_loc: Name of VASP command (i.e. executable location)
    :param structure: flare structure object
    :param en: whether to return the final energy along with the forces
    :param vasp_cmd: Command to run VASP (leave a formatter "{}" to insert dft_loc);
        this can be used to specify mpirun, srun, etc.
    Returns:
        forces on each atom (and energy if en=True)
    """

    currdir = os.getcwd()
    if structure:
        edit_dft_input_positions("POSCAR", structure)
    try:
        if currdir != calc_dir:
            os.chdir(calc_dir)
        call(vasp_cmd.format(dft_loc).split())

        if en:
            parse_func = parse_dft_forces_and_energy
        else:
            parse_func = parse_dft_forces

        try:
            forces = parse_func("vasprun.xml")
        except FileNotFoundError:
            raise FileNotFoundError("""Could not load vasprun.xml.
            The calculation may not have finished.
            Current directory is %s""" % os.getcwd())

        os.chdir(currdir)
        return forces

    except Exception as e:
        os.chdir(currdir)
        raise e


def run_dft_par(dft_input: str, structure: Structure,
                dft_command:str= None, n_cpus=1,
                dft_out="vasprun.xml",
                parallel_prefix="mpi",
                mpi = None, npool = None,
                **dft_kwargs):
    # TODO Incorporate Custodian.
    edit_dft_input_positions(dft_input, structure)

    if dft_command is None and not os.environ.get('VASP_COMMAND'):
        raise FileNotFoundError\
            ("Warning: No VASP Command passed, or stored in "
            "environment as VASP_COMMAND. ")

    if n_cpus > 1:
        # why is parallel prefix needed?
        if (parallel_prefix == "mpi") and (mpi == 'mpi'):
            dft_command = f'mpirun -np {n_cpus} {dft_command}'
        elif mpi == 'srun':
            dft_command = f'srun -n {n_cpus} {dft_command}'

    else:
        serial_prefix = dft_kwargs.get('serial_prefix', '')
        dft_command = f'{serial_prefix} {dft_command}'

    call(dft_command, shell=True)

    return parse_dft_forces(dft_out)


def dft_input_to_structure(poscar: str):
    """
    Parse the DFT input in a directory.
    :param vasp_input: directory of vasp input
    """
    return Structure.from_pmg_structure(Poscar.from_file(poscar).structure)


def edit_dft_input_positions(output_name: str, structure: Structure):
    """
    Writes a VASP POSCAR file from structure with the name poscar .
    WARNING: Destructively replaces the file with the name specified by poscar
    :param poscar: Name of file
    :param structure: structure to write to file
    """
    if os.path.isfile(output_name):
        shutil.copyfile(output_name, output_name + '.bak')
    poscar = Poscar(structure.to_pmg_structure())
    poscar.write_file(output_name)
    return output_name


def parse_dft_forces(vasprun: Union[str, Vasprun]):
    """
    Parses the DFT forces from the last ionic step of a VASP vasprun.xml file
    :param vasprun: pymatgen Vasprun object or vasprun filename
    """
    vasprun = check_vasprun(vasprun)
    istep = vasprun.ionic_steps[-1]
    return np.array(istep['forces'])


def parse_dft_forces_and_energy(vasprun: Union[str, Vasprun]):
    """
    Parses the DFT forces and energy from a VASP vasprun.xml file
    :param vasprun: pymatgen Vasprun object or vasprun filename
    """
    vasprun = check_vasprun(vasprun)
    istep = vasprun.ionic_steps[-1]
    return np.array(istep['forces']), \
           istep["electronic_steps"][-1]["e_0_energy"]


def md_trajectory_from_vasprun(vasprun: Union[str, Vasprun],
                               ionic_step_skips=1, vasprun_kwargs: dict={}):
    """
    Returns a list of flare Structure objects decorated with forces, stress,
    and total energy from a MD trajectory performed in VASP.
    :param vasprun: pymatgen Vasprun object or vasprun filename
    :param ionic_step_skips: if True, only samples the configuration every
        ionic_skip_steps steps.
    """
    vasprun = check_vasprun(vasprun, vasprun_kwargs)

    struc_lst = []
    for step in vasprun.ionic_steps[::ionic_step_skips]:
        structure = Structure.from_pmg_structure(step["structure"])
        structure.energy = step["electronic_steps"][-1]["e_0_energy"]
        # TODO should choose e_wo_entrp or e_fr_energy?
        structure.forces = np.array(step["forces"])
        structure.stress = np.array(step["stress"])
        struc_lst.append(structure)

    return struc_lst
