# ab-Initio Configuration Sampling tool kit (abICS)
# Copyright (C) 2019- The University of Tokyo
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program. If not, see http://www.gnu.org/licenses/.

"""
To deal with VASP
"""

from __future__ import annotations

import os, os.path
import time
from collections import namedtuple

import numpy as np
from pymatgen.io.vasp.inputs import Poscar, VaspInput
from pymatgen.io.vasp.outputs import Oszicar
from pymatgen.core import Structure

from .params import ALParams, DFTParams
from .base_solver import SolverBase


class VASPSolver(SolverBase):
    """
    VASP solver

    Attributes
    ----------
    path_to_solver : str
        Path to the solver
    input : VASPSolver.Input
        Input manager
    output : VASPSolver.Output
        Output manager
    """

    def __init__(self, path_to_solver, ignore_species=None):
        """
        Initialize the solver.

        Parameters
        ----------
        path_to_solver : str
            Path to the solver.
        """
        super(VASPSolver, self).__init__(path_to_solver)
        self.path_to_solver = path_to_solver
        self.input = VASPSolver.Input(ignore_species)
        self.output = VASPSolver.Output()

    def name(self):
        return "VASP"

    class Input(object):
        """
        Input manager for VASP

        Attributes
        ----------
        base_info : dict
            Common parameters
        pos_info : pymatgen.Structure
            Atom positions
        """

        def __init__(self, ignore_species):
            self.base_info = None
            self.pos_info = None
            self.ignore_species = ignore_species

        def from_directory(self, base_input_dir):
            """

            Parameters
            ----------
            base_input_dir : str
                Path to the directory including base input files.
            Returns
            -------
            base_vasp_input : VaspInput (defined in pymatgen)
                vasp input object
            """
            self.base_vasp_input = VaspInput.from_directory(base_input_dir)
            self.base_info = self.base_vasp_input.get("INCAR")
            return self.base_vasp_input

        def update_info_by_structure(self, structure):
            """
            Update information by structure file

            Parameters
            ----------
            structure : pymatgen.Structure
                Atomic structure
            """
            if self.ignore_species is not None:
                structure = structure.copy()
                structure.remove_species(self.ignore_species)
                structure.sort(key=lambda site: site.species_string)
            if "seldyn" in structure.site_properties.keys():
                self.pos_info = Poscar(
                    structure=structure,
                    selective_dynamics=structure.site_properties["seldyn"],
                )
            else:
                self.pos_info = Poscar(structure=structure)
            self.base_vasp_input.update({"POSCAR": self.pos_info})

        def update_info_from_files(self, workdir, rerun):
            """
            Update information from output files.

            Parameters
            ----------
            workdir: str
                Path to working directory.
            rerun: int
                Mode for rerunning (0: not rerun, 1: rerun (base file is replaced by new one))
            """
            if rerun == 1:
                info_dict = ["BASE", "POS"]
            elif rerun > 0:
                info_dict = ["POS"]

            # Add update procedure
            # 1. Check kind of the update information
            # 2. Replace the information
            for info in info_dict:
                if info == "BASE":
                    # Modify parameters
                    self.base_info = self.base_vasp_input.get("INCAR")
                    self.base_info.update({"IBRION": 3, "POTIM": 0.2})
                    self.base_vasp_input.update({"INCAR": self.base_info})
                elif info == "POS":
                    # Update positions
                    self.pos_info = Poscar.from_file(os.path.join(workdir, "CONTCAR"))
                    self.base_vasp_input.update({"POSCAR": self.pos_info})

        def write_input(self, output_dir):
            """
            Generate input files of the solver program.

            Parameters
            ----------
            workdir : str
                Path to working directory.
            """
            if self.base_info is None:
                raise AttributeError("Fail to set base_info.")
            self.base_vasp_input.write_input(output_dir=output_dir)

        def cl_args(self, nprocs, nthreads, workdir):
            """
            Generate command line argument of the solver program.

            Parameters
            ----------
            nprocs : int
                The number of processes.
            nthreads : int
                The number of threads.
            workdir : str
                Path to the working directory.

            Returns
            -------
            args : list[str]
                Arguments of command
            """
            return [workdir]

    class Output(object):
        """
        Output manager.
        """

        # def __init__(self):
        #    self.drone = SimpleVaspToComputedEntryDrone(inc_structure=True)
        #    self.queen = BorgQueen(self.drone)

        def get_results(self, workdir):
            """
            Get energy and structure obtained by the solver program.

            Parameters
            ----------
            workdir : str
                Path to the working directory.

            Returns
            -------
            phys : named_tuple("energy", "structure")
                Total energy and atomic structure.
                The energy is measured in the units of eV
                and coodinates is measured in the units of Angstrom.
            """
            # Read results from files in output_dir and calculate values
            Phys = namedtuple("PhysVaules", ("energy", "structure"))

            # Make sure output is newer than input

            counter = 0
            while True:
                st_in = os.stat(os.path.join(workdir, "INCAR"))
                st_out = os.stat(os.path.join(workdir, "OSZICAR"))
                if st_in.st_mtime - st_out.st_mtime < 0:
                    break
                elif counter > 10:
                    raise TimeoutError(
                        "VASP OSZICAR is older than INCAR and no update was seen for 1 minute. "
                        "If VASP was actually run, then this might be due to a very slow file system."
                    )
                else:
                    print(
                        "VASP OSZICAR is older than INCAR; waiting for file system update"
                    )
                    counter += 1
                    time.sleep(10)

            energy = Oszicar(os.path.join(workdir, "OSZICAR")).final_energy
            structure = Structure.from_file(os.path.join(workdir, "CONTCAR"))

            return Phys(np.float64(energy), structure)

    def solver_run_schemes(self):
        return ("mpi_spawn_ready",)

    # -- factory

    @classmethod
    def create(cls, params: ALParams | DFTParams):
        path = params.path
        ignore_species = params.ignore_species
        return cls(path, ignore_species)
