# coding: utf-8
# Copyright (c) 2018-2019, Taku MURAKAMI. All rights reserved.
# Distributed under the terms of the BSD 3-clause License.

import logging
import os
import csv
import subprocess
import xml.etree.cElementTree as ET
import pymatgen
from pymatgen.io.vasp.inputs import Poscar
from pymatgen.io.vasp.inputs import Incar
# from pymatgen.io.vasp.inputs import Kpoints
from pymatgen.io.vasp.outputs import Vasprun

"""
Class to performe high-throughput first-principles calculation with VASP.
"""

logger = logging.getLogger(__name__)


class Calculation_vasp(object):
    """
    Class to perform first-principles calculation with VASP.
    """
    
    def __init__(self, struct_name, struct, calculator, potential_path):
        """
        Arguments
        ---------
        struct_name: str
            Name of structure data.
        struct: pymatgen.Structure
            Atomic structure itself.
        calculator: dict
            Calculation configulation as ASE format.
        potential_path: str
            Path to pseudo-potential database using in VASP calculation.
        """
        self._struct_name = struct_name
        self._struct = struct
        try:
            self._output_path = self._set_output_path(calculator["txt"])
            self._write_input_files(struct, calculator, potential_path)
        except KeyError:
            pass
        except TypeError:
            pass
    
    def _set_output_path(self, output_path):
        """
        Sets output file path of calculation.
        
        Arguments
        ---------
        output_path: str or None
            Output file path, which taking the form
            "(output file path)/(structure name).txt".
        
        Returns
        -------
        path: str
            Output file path, if it is None in arguments,
            it is "./outputs/(structure name)".
        """
        if output_path is not None:
            path = os.path.splitext(output_path)[0]
        else:
            return None
        if not os.path.exists(path):
            os.makedirs(path, exist_ok=True)
        return path
    
    def _write_input_files(self, struct, calculator, potential_path):
        """
        Writes input files using in VASP calculation.
        
        Arguments
        ---------
        struct: pymatgen.Structure
            Atomic structure itself.
        calculator: dict
            Calculation configulation as ASE format.
        potential_path: str
            Path to pseudo-potential database using in VASP calculation.
        """
        self._write_poscar(struct)
        self._write_incar(calculator)
        self._write_kpoints(calculator["kpts"])
        self._write_potcar(struct, potential_path)
    
    def _write_poscar(self, struct):
        """
        Writes POSCAR file.
        
        Arguments
        ---------
        struct: pymatgen.Strucuture
            Atomic structure itself.
        """
        with open("POSCAR", mode="w") as file:
            file.writelines(str(Poscar(struct)))
    
    def _write_incar(self, calculator):
        """
        Writes INCAR file.
        
        Arguments
        ---------
        calculator: dict
            Calculation configulation as ASE format.
        
        Parameters
        ----------
        incar_dict: dict
            Dictionary of incar setting.
        """
        incar_dict = {"SYSTEM": self._struct_name,
                      "ALGO": "VeryFast",
                      "LWAVE": False,
                      "ISTART": 0,
                      "ICHARG": 2,
                      "ISPIN": 1,
                      "IBRION": 1,
                      "ISYM": 0,
                      "ISIF": 3}
        
        if calculator.get("maxiter") is not None:
            incar_dict["NELM"] = calculator["maxiter"]
        
        if calculator.get("isif") is not None:
            incar_dict["isif"] = calculator["isif"]
        
        with open("INCAR", mode="w") as file:
            file.writelines(str(Incar(incar_dict)))
    
    def _write_kpoints(self, kpts):
        """
        Writes KPOINTS file.
        
        Arguments
        ---------
        kpts: dict
            Calculation configulation of KPOINTS as ASE format.
        
        Parameters
        ----------
        size: tuple
            The number of KPOINTS aligned to a, b and c direction.
            Default: (1, 1, 1)
        shift: bool
            Gamma centered or not (Monkhorst-Pack).
            Default: "Gamma"
        """
        
        if kpts.get("size") is not None:
            size = kpts["size"]
        else:
            size = (1, 1, 1)
        
        if kpts.get("gamma") is not None:
            shift = "Gamma" if kpts["gamma"] is True else "Monkhorst-Pack"
        else:
            shift = "Gamma"
        
        with open("KPOINTS", mode="w") as file:
            file.write("Automatic mesh generated by pythroughput\n")
            file.write("  0\n")
            file.write(shift + "\n")
            file.write("  " + str(size[0]) + "  " + str(size[1]) + "  " + str(size[2]) + "\n")
            file.write("  0. 0. 0.\n")
    
    def _write_potcar(self, struct, potential_path):
        """
        Writes POTCAR file.
        
        Arguments
        ---------
        struct: pymatgen.Structure
            Atomic structure itself.
        potential_path: str
            Path to pseudo-potential database using in VASP calculation.
        """
        if os.path.exists("POTCAR") is True:
            os.remove("POTCAR")
        
        recommended_potential = self._read_potential(potential_path)
        
        with open("POTCAR", mode="a") as file:
            for symbol in struct.symbol_set:
                file.writelines(open(
                    potential_path + recommended_potential[symbol] + "/POTCAR"
                ).read())
    
    def _read_potential(self, potential_path):
        """
        Reads recommended potential dict from database.
        
        Arguments
        ---------
        potential_path: str
            Path of the recommended pseudo-potential database which used in VASP calculation.
        
        Returns
        -------
        potential: dict
            Recommended potential dict which has specie as key and potential name as value.
        """
        potential = []
        
        with open(potential_path+"db_recommended_paw.csv", "r") as file:
            reader = csv.reader(file, delimiter="\t")
            for row in reader:
                potential.append(row)
        
        return dict(potential)
    
    def get_results(self, steps=1, n_jobs=4,
                    results_list=["struct_name",
                                  "initial_energy",
                                  "total_energy",
                                  "initial_forces",
                                  "final_forces"],
                    backup_file_list=["POSCAR", "vasprun.xml", "INCAR"]):
        """
        Gets calculation results.
        
        Arguments
        ---------
        steps: int
            Number of relaxation steps
        n_jobs: int
            Number of CPU using in parallel calculation.
        results_list: list
            List of required results.
        backup_file_list: list
            List of backuped files.
        
        Parameters
        ----------
        results: dict
            dictionary of calculation results.
        vasprun: pymatgen.io.vasp.outputs.Vasprun
            vasprun.xml file, which includes all calculation results.
        
        Returns
        -------
        results: dict
            dictionary of caluclation results.
        """
        results = {}
        
        if not steps is 1:
            with open("INCAR", mode="a") as file:
                file.write("NSW="+str(steps))
            backup_file_list.append("CONTCAR")
        
        self._run_vasp(n_jobs)
        if self._output_path is not None:
            self._mv_output_files(backup_file_list)
        
        results = self.read_results(results_list)
        
        return results
    
    def read_results(self, results_list=["struct_name",
                                         "initial_energy",
                                         "total_energy",
                                         "initial_forces",
                                         "final_forces"]):
        """
        Reads results from calculated files, vasprun.xml.
        
        Arguments
        ---------
        vasprun: pymatgen.io.vasp.outputs.Vasprun
            vasprun.xml file, which includes all calculation results.
        results_list
            List of required results.
        
        Returns
        -------
        results: dict
            dictionary of caluclation results.
        """
        results = {}
        try:
            if self._output_path is not None:
                vasprun = Vasprun(self._output_path + "/vasprun.xml")
            else:
                vasprun = Vasprun("vasprun.xml")
        except (ET.ParseError, ValueError) as error:
            results["error"] = error
            return results
        
        for term in results_list:
            if term == "struct_name":
                results[term] = self._struct_name
            elif term == "initial_energy":
                results[term] = self._get_initial_energy(vasprun)
            elif term == "total_energy":
                results[term] = self._get_total_energy(vasprun)
            elif term == "initial_forces":
                results[term] = self._get_initial_forces(vasprun)
            elif term == "final_forces":
                results[term] = self._get_final_forces(vasprun)
            elif term == "formula":
                results[term] = self._struct.formula
            else:
                results[term] = "Undefined parameter"
        return results
    
    def _run_vasp(self, n_jobs):
        """
        Runs VASP calculation.
        
        Arguments
        ---------
        n_jobs: int
            Number of CPU using in parallel calculation.
        """
        cmd = "mpirun -np " + str(n_jobs) + " vasp &"
        devnull = open("/dev/null", "w")
        subprocess.run(cmd.split(), stdout=devnull)
    
    def _get_total_energy(self, vasprun):
        """
        Gets total energy.
        
        Arguments
        ---------
        vasprun: pymatgen.io.vasp.outputs.Vasprun
            vasprun.xml file, which includes all calculation results.
        
        Returns
        -------
        float
            Calculated total energy [eV].
        """
        if vasprun.converged_electronic is True:
            return float(vasprun.final_energy)
        else:
            return "Unconverged"
    
    def _get_initial_energy(self, vasprun):
        """
        Gets initial energy of first step ().
        
        Arguments
        ---------
        vasprun: pymatgen.io.vasp.outputs.Vasprun
            vasprun.xml file, which includes all calculation results.
        
        Returns
        -------
        initial_energy: float
            Calculated total energy of initial structure [eV].
        """
        
        try:
            initial_istep = vasprun.ionic_steps[0]
            if initial_istep["electronic_steps"] < vasprun.parameters["NELM"]:
                if initial_istep["e_wo_entrp"] != initial_istep[
                        'electronic_steps'][-1]["e_0_energy"]:
                    return float(initial_istep["e_wo_entrp"])
                return float(initial_istep['electronic_steps'][-1]["e_0_energy"])
            else:
                return "Unconverged"
        except (IndexError, KeyError):
            return float('inf')
    
    def _get_initial_forces(self, vasprun):
        """
        Gets initial force.
        
        Arguments
        ---------
        vasprun: pytmatgen.io.vasp.outputs.Vasprun
            vasprun.xml file, which includes all calculation results.
        
        Returns
        -------
        list
            Calculated force for every elements [eV/A].
        """
        return vasprun.ionic_steps[0]["forces"]
    
    def _get_final_forces(self, vasprun):
        """
        Gets final force.
        
        Arguments
        ---------
        vasprun: pymatgen.io.vasp.outputs.Vasprun
            vasprun.xml file, which includes all calculation results.
        
        Returns
        -------
        list
            Calculated force for every elements [eV/A].
        """
        return vasprun.ionic_steps[-1]["forces"]
    
    def _is_optimize_converged(self, vasprun):
        """
        Is optimization converged.
        
        Arguments
        ---------
        vasprun: pymatgen.io.vasp.outputs.Vasprun
            vasprun.xml file, which includes all calculation results.
        
        Returns
        -------
        bool
            Is optimization converged.
        """
        return vasprun.converged_ionic
    
    def _mv_output_files(self, backup_file_list):
        """
        Moves output files to output path.
        
        Arguments
        ---------
        backup_file_list: tuple
            File list moving to output path.
        """
        for file in backup_file_list:
            if os.path.exists(self._output_path):
                os.makedirs(self._output_path, exist_ok=True)
            cmd = "mv " + file + " " + self._output_path
            subprocess.run(cmd.split())
    
