"""This module defines an ASE interface to CP2K.

https://www.cp2k.org/
Author: Ole Schuett <ole.schuett@mat.ethz.ch>
"""

import os
import os.path
import subprocess
from contextlib import AbstractContextManager
from warnings import warn

import numpy as np

import ase.io
from ase.calculators.calculator import (
    Calculator,
    CalculatorSetupError,
    Parameters,
    all_changes,
)
from ase.config import cfg
from ase.units import Rydberg


class CP2K(Calculator, AbstractContextManager):
    """ASE-Calculator for CP2K.

    CP2K is a program to perform atomistic and molecular simulations of solid
    state, liquid, molecular, and biological systems. It provides a general
    framework for different methods such as e.g., density functional theory
    (DFT) using a mixed Gaussian and plane waves approach (GPW) and classical
    pair and many-body potentials.

    CP2K is freely available under the GPL license.
    It is written in Fortran 2003 and can be run efficiently in parallel.

    Check https://www.cp2k.org about how to obtain and install CP2K.
    Make sure that you also have the CP2K-shell available, since it is required
    by the CP2K-calulator.

    The CP2K-calculator relies on the CP2K-shell. The CP2K-shell was originally
    designed for interactive sessions. When a calculator object is
    instantiated, it launches a CP2K-shell as a subprocess in the background
    and communications with it through stdin/stdout pipes. This has the
    advantage that the CP2K process is kept alive for the whole lifetime of
    the calculator object, i.e. there is no startup overhead for a sequence
    of energy evaluations. Furthermore, the usage of pipes avoids slow file-
    system I/O. This mechanism even works for MPI-parallelized runs, because
    stdin/stdout of the first rank are forwarded by the MPI-environment to the
    mpiexec-process.

    The command used by the calculator to launch the CP2K-shell is
    ``cp2k_shell``. To run a parallelized simulation use something like this::

        CP2K.command="env OMP_NUM_THREADS=2 mpiexec -np 4 cp2k_shell.psmp"

    The CP2K-shell can be shut down by calling :meth:`close`.
    The close method will be called automatically when using the calculator as
    part of a with statement::

        with CP2K() as calc:
            calc.get_potential_energy(atoms)

    The shell will be restarted if you call the calculator object again.

    Arguments:

    auto_write: bool
        Flag to enable the auto-write mode. If enabled the
        ``write()`` routine is called after every
        calculation, which mimics the behavior of the
        ``FileIOCalculator``. Default is ``False``.
    basis_set: str
        Name of the basis set to be use.
        The default is ``DZVP-MOLOPT-SR-GTH``.
    basis_set_file: str
        Filename of the basis set file.
        Default is ``BASIS_MOLOPT``.
        Set the environment variable $CP2K_DATA_DIR
        to enabled automatic file discovered.
    charge: float
        The total charge of the system.  Default is ``0``.
    command: str
        The command used to launch the CP2K-shell.
        If ``command`` is not passed as an argument to the
        constructor, the class-variable ``CP2K.command``,
        and then the environment variable
        ``$ASE_CP2K_COMMAND`` are checked.
        Eventually, ``cp2k_shell`` is used as default.
    cutoff: float
        The cutoff of the finest grid level.  Default is ``400 * Rydberg``.
    debug: bool
        Flag to enable debug mode. This will print all
        communication between the CP2K-shell and the
        CP2K-calculator. Default is ``False``.
    force_eval_method: str
        The method CP2K uses to evaluate energies and forces.
        The default is ``Quickstep``, which is CP2K's
        module for electronic structure methods like DFT.
    inp: str
        CP2K input template. If present, the calculator will
        augment the template, e.g. with coordinates, and use
        it to launch CP2K. Hence, this generic mechanism
        gives access to all features of CP2K.
        Note, that most keywords accept ``None`` to disable the generation
        of the corresponding input section.

        This input template is important for advanced CP2K
        inputs, but is also needed for e.g. controlling the Brillouin
        zone integration. The example below illustrates some common
        options::

            inp = '''&FORCE_EVAL
               &DFT
                 &KPOINTS
                   SCHEME MONKHORST-PACK 12 12 8
                 &END KPOINTS
                 &SCF
                   ADDED_MOS 10
                   &SMEAR
                     METHOD FERMI_DIRAC
                     ELECTRONIC_TEMPERATURE [K] 500.0
                   &END SMEAR
                 &END SCF
               &END DFT
             &END FORCE_EVAL
            '''

    max_scf: int
        Maximum number of SCF iteration to be performed for
        one optimization. Default is ``50``.
    multiplicity: int, default=None
        Select the multiplicity of the system
        (two times the total spin plus one).
        If None, multiplicity is not explicitly given in the input file.
    poisson_solver: str
        The poisson solver to be used. Currently, the only supported
        values are ``auto`` and ``None``. Default is ``auto``.
    potential_file: str
        Filename of the pseudo-potential file.
        Default is ``POTENTIAL``.
        Set the environment variable $CP2K_DATA_DIR
        to enabled automatic file discovered.
    pseudo_potential: str
        Name of the pseudo-potential to be use.
        Default is ``auto``. This tries to infer the
        potential from the employed XC-functional,
        otherwise it falls back to ``GTH-PBE``.
    stress_tensor: bool
        Indicates whether the analytic stress-tensor should be calculated.
        Default is ``True``.
    uks: bool
        Requests an unrestricted Kohn-Sham calculations.
        This is need for spin-polarized systems, ie. with an
        odd number of electrons. Default is ``False``.
    xc: str
        Name of exchange and correlation functional.
        Accepts all functions supported by CP2K itself or libxc.
        Default is ``LDA``.
    print_level: str
        PRINT_LEVEL of global output.
        Possible options are:
        DEBUG Everything is written out, useful for debugging purposes only
        HIGH Lots of output
        LOW Little output
        MEDIUM Quite some output
        SILENT Almost no output
        Default is 'LOW'
    set_pos_file: bool
        Send updated positions to the CP2K shell via file instead of
        via stdin, which can bypass limitations for sending large
        structures via stdin for CP2K built with some MPI libraries.
        Requires CP2K 2024.2
    """

    implemented_properties = ['energy', 'free_energy', 'forces', 'stress']
    command = None

    default_parameters = dict(
        auto_write=False,
        basis_set='DZVP-MOLOPT-SR-GTH',
        basis_set_file='BASIS_MOLOPT',
        charge=0,
        cutoff=400 * Rydberg,
        force_eval_method="Quickstep",
        inp='',
        max_scf=50,
        multiplicity=None,
        potential_file='POTENTIAL',
        pseudo_potential='auto',
        stress_tensor=True,
        uks=False,
        poisson_solver='auto',
        xc='LDA',
        print_level='LOW',
        set_pos_file=False,
    )

    def __init__(self, restart=None,
                 ignore_bad_restart_file=Calculator._deprecated,
                 label='cp2k', atoms=None, command=None,
                 debug=False, **kwargs):
        """Construct CP2K-calculator object."""

        self._debug = debug
        self._force_env_id = None
        self._shell = None
        self.label = None
        self.parameters = None
        self.results = None
        self.atoms = None

        # Several places are check to determine self.command
        if command is not None:
            self.command = command
        elif CP2K.command is not None:
            self.command = CP2K.command
        else:
            self.command = cfg.get('ASE_CP2K_COMMAND', 'cp2k_shell')

        super().__init__(restart=restart,
                         ignore_bad_restart_file=ignore_bad_restart_file,
                         label=label, atoms=atoms, **kwargs)
        if restart is not None:
            self.read(restart)

        # Start the shell by default, which is how SocketIOCalculator
        self._shell = Cp2kShell(self.command, self._debug)

    def __del__(self):
        """Terminate cp2k_shell child process"""
        self.close()

    def __exit__(self, __exc_type, __exc_value, __traceback):
        self.close()

    def close(self):
        """Close the attached shell"""
        if self._shell is not None:
            self._shell.close()
            self._shell = None
            self._force_env_id = None  # Force env must be recreated

    def set(self, **kwargs):
        """Set parameters like set(key1=value1, key2=value2, ...)."""
        msg = '"%s" is not a known keyword for the CP2K calculator. ' \
              'To access all features of CP2K by means of an input ' \
              'template, consider using the "inp" keyword instead.'
        for key in kwargs:
            if key not in self.default_parameters:
                raise CalculatorSetupError(msg % key)

        changed_parameters = Calculator.set(self, **kwargs)
        if changed_parameters:
            self.reset()

    def write(self, label):
        'Write atoms, parameters and calculated results into restart files.'
        if self._debug:
            print("Writing restart to: ", label)
        self.atoms.write(label + '_restart.traj')
        self.parameters.write(label + '_params.ase')
        from ase.io.jsonio import write_json
        with open(label + '_results.json', 'w') as fd:
            write_json(fd, self.results)

    def read(self, label):
        'Read atoms, parameters and calculated results from restart files.'
        self.atoms = ase.io.read(label + '_restart.traj')
        self.parameters = Parameters.read(label + '_params.ase')
        from ase.io.jsonio import read_json
        with open(label + '_results.json') as fd:
            self.results = read_json(fd)

    def calculate(self, atoms=None, properties=None,
                  system_changes=all_changes):
        """Do the calculation."""

        if not properties:
            properties = ['energy']
        Calculator.calculate(self, atoms, properties, system_changes)

        # Start the shell if needed
        if self._shell is None:
            self._shell = Cp2kShell(self.command, self._debug)

        if self._debug:
            print("system_changes:", system_changes)

        if 'numbers' in system_changes:
            self._release_force_env()

        if self._force_env_id is None:
            self._create_force_env()

        # enable eV and Angstrom as units
        self._shell.send('UNITS_EV_A')
        self._shell.expect('* READY')

        n_atoms = len(self.atoms)
        if 'cell' in system_changes:
            cell = self.atoms.get_cell()
            self._shell.send('SET_CELL %d' % self._force_env_id)
            for i in range(3):
                self._shell.send('%.18e %.18e %.18e' % tuple(cell[i, :]))
            self._shell.expect('* READY')

        if 'positions' in system_changes:
            if self.parameters.set_pos_file:
                # TODO: Update version number when released
                if self._shell.version < 7:
                    raise ValueError('SET_POS_FILE requires > CP2K 2024.2')
                pos: np.ndarray = self.atoms.get_positions()
                fn = self.label + '.pos'
                with open(fn, 'w') as fp:
                    print(3 * n_atoms, file=fp)
                    for pos in self.atoms.get_positions():
                        print('%.18e %.18e %.18e' % tuple(pos), file=fp)
                self._shell.send(f'SET_POS_FILE {fn} {self._force_env_id}')
            else:
                if len(atoms) > 100 and 'psmp' in self.command:
                    warn('ASE may stall when passing large structures'
                         ' to MPI versions of CP2K.'
                         ' Consider using `set_pos_file=True`.')
                self._shell.send('SET_POS %d' % self._force_env_id)
                self._shell.send('%d' % (3 * n_atoms))
                for pos in self.atoms.get_positions():
                    self._shell.send('%.18e %.18e %.18e' % tuple(pos))
                self._shell.send('*END')
            max_change = float(self._shell.recv())
            assert max_change >= 0  # sanity check
            self._shell.expect('* READY')

        self._shell.send('EVAL_EF %d' % self._force_env_id)
        self._shell.expect('* READY')

        self._shell.send('GET_E %d' % self._force_env_id)
        self.results['energy'] = float(self._shell.recv())
        self.results['free_energy'] = self.results['energy']
        self._shell.expect('* READY')

        forces = np.zeros(shape=(n_atoms, 3))
        self._shell.send('GET_F %d' % self._force_env_id)
        nvals = int(self._shell.recv())
        assert nvals == 3 * n_atoms  # sanity check
        for i in range(n_atoms):
            line = self._shell.recv()
            forces[i, :] = [float(x) for x in line.split()]
        self._shell.expect('* END')
        self._shell.expect('* READY')
        self.results['forces'] = forces

        self._shell.send('GET_STRESS %d' % self._force_env_id)
        line = self._shell.recv()
        self._shell.expect('* READY')

        stress = np.array([float(x) for x in line.split()]).reshape(3, 3)
        assert np.all(stress == np.transpose(stress))  # should be symmetric
        # Convert 3x3 stress tensor to Voigt form as required by ASE
        stress = np.array([stress[0, 0], stress[1, 1], stress[2, 2],
                           stress[1, 2], stress[0, 2], stress[0, 1]])
        self.results['stress'] = -1.0 * stress  # cp2k uses the opposite sign

        if self.parameters.auto_write:
            self.write(self.label)

    def _create_force_env(self):
        """Instantiates a new force-environment"""
        assert self._force_env_id is None
        label_dir = os.path.dirname(self.label)
        if len(label_dir) > 0 and not os.path.exists(label_dir):
            print('Creating directory: ' + label_dir)
            os.makedirs(label_dir)  # cp2k expects dirs to exist

        inp = self._generate_input()
        inp_fn = self.label + '.inp'
        out_fn = self.label + '.out'
        self._write_file(inp_fn, inp)
        self._shell.send(f'LOAD {inp_fn} {out_fn}')
        self._force_env_id = int(self._shell.recv())
        assert self._force_env_id > 0
        self._shell.expect('* READY')

    def _write_file(self, fn, content):
        """Write content to a file"""
        if self._debug:
            print('Writting to file: ' + fn)
            print(content)
        if self._shell.version < 2.0:
            with open(fn, 'w') as fd:
                fd.write(content)
        else:
            lines = content.split('\n')
            if self._shell.version < 2.1:
                lines = [l.strip() for l in lines]  # save chars
            self._shell.send('WRITE_FILE')
            self._shell.send(fn)
            self._shell.send('%d' % len(lines))
            for line in lines:
                self._shell.send(line)
            self._shell.send('*END')
            self._shell.expect('* READY')

    def _release_force_env(self):
        """Destroys the current force-environment"""
        if self._force_env_id:
            if self._shell.isready:
                self._shell.send('DESTROY %d' % self._force_env_id)
                self._shell.expect('* READY')
            else:
                msg = "CP2K-shell not ready, could not release force_env."
                warn(msg, RuntimeWarning)
            self._force_env_id = None

    def _generate_input(self):
        """Generates a CP2K input file"""
        p = self.parameters
        root = parse_input(p.inp)
        root.add_keyword('GLOBAL', 'PROJECT ' + self.label)
        if p.print_level:
            root.add_keyword('GLOBAL', 'PRINT_LEVEL ' + p.print_level)
        if p.force_eval_method:
            root.add_keyword('FORCE_EVAL', 'METHOD ' + p.force_eval_method)
        if p.stress_tensor:
            root.add_keyword('FORCE_EVAL', 'STRESS_TENSOR ANALYTICAL')
            root.add_keyword('FORCE_EVAL/PRINT/STRESS_TENSOR',
                             '_SECTION_PARAMETERS_ ON')
        if p.basis_set_file:
            root.add_keyword('FORCE_EVAL/DFT',
                             'BASIS_SET_FILE_NAME ' + p.basis_set_file)
        if p.potential_file:
            root.add_keyword('FORCE_EVAL/DFT',
                             'POTENTIAL_FILE_NAME ' + p.potential_file)
        if p.cutoff:
            root.add_keyword('FORCE_EVAL/DFT/MGRID',
                             'CUTOFF [eV] %.18e' % p.cutoff)
        if p.max_scf:
            root.add_keyword('FORCE_EVAL/DFT/SCF', 'MAX_SCF %d' % p.max_scf)
            root.add_keyword('FORCE_EVAL/DFT/LS_SCF', 'MAX_SCF %d' % p.max_scf)

        if p.xc:
            legacy_libxc = ""
            for functional in p.xc.split():
                functional = functional.replace("LDA", "PADE")  # resolve alias
                xc_sec = root.get_subsection('FORCE_EVAL/DFT/XC/XC_FUNCTIONAL')
                # libxc input section changed over time
                if functional.startswith("XC_") and self._shell.version < 3.0:
                    legacy_libxc += " " + functional  # handled later
                elif functional.startswith("XC_") and self._shell.version < 5.0:
                    s = InputSection(name='LIBXC')
                    s.keywords.append('FUNCTIONAL ' + functional)
                    xc_sec.subsections.append(s)
                elif functional.startswith("XC_"):
                    s = InputSection(name=functional[3:])
                    xc_sec.subsections.append(s)
                else:
                    s = InputSection(name=functional.upper())
                    xc_sec.subsections.append(s)
            if legacy_libxc:
                root.add_keyword('FORCE_EVAL/DFT/XC/XC_FUNCTIONAL/LIBXC',
                                 'FUNCTIONAL ' + legacy_libxc)

        if p.uks:
            root.add_keyword('FORCE_EVAL/DFT', 'UNRESTRICTED_KOHN_SHAM ON')

        if p.multiplicity:
            root.add_keyword('FORCE_EVAL/DFT',
                             'MULTIPLICITY %d' % p.multiplicity)

        if p.charge and p.charge != 0:
            root.add_keyword('FORCE_EVAL/DFT', 'CHARGE %d' % p.charge)

        # add Poisson solver if needed
        if p.poisson_solver == 'auto' and not any(self.atoms.get_pbc()):
            root.add_keyword('FORCE_EVAL/DFT/POISSON', 'PERIODIC NONE')
            root.add_keyword('FORCE_EVAL/DFT/POISSON', 'PSOLVER  MT')

        # write coords
        syms = self.atoms.get_chemical_symbols()
        atoms = self.atoms.get_positions()
        for elm, pos in zip(syms, atoms):
            line = f'{elm} {pos[0]:.18e} {pos[1]:.18e} {pos[2]:.18e}'
            root.add_keyword('FORCE_EVAL/SUBSYS/COORD', line, unique=False)

        # write cell
        pbc = ''.join([a for a, b in zip('XYZ', self.atoms.get_pbc()) if b])
        if len(pbc) == 0:
            pbc = 'NONE'
        root.add_keyword('FORCE_EVAL/SUBSYS/CELL', 'PERIODIC ' + pbc)
        c = self.atoms.get_cell()
        for i, a in enumerate('ABC'):
            line = f'{a} {c[i, 0]:.18e} {c[i, 1]:.18e} {c[i, 2]:.18e}'
            root.add_keyword('FORCE_EVAL/SUBSYS/CELL', line)

        # determine pseudo-potential
        potential = p.pseudo_potential
        if p.pseudo_potential == 'auto':
            if p.xc and p.xc.upper() in ('LDA', 'PADE', 'BP', 'BLYP', 'PBE',):
                potential = 'GTH-' + p.xc.upper()
            else:
                msg = 'No matching pseudo potential found, using GTH-PBE'
                warn(msg, RuntimeWarning)
                potential = 'GTH-PBE'  # fall back

        # write atomic kinds
        subsys = root.get_subsection('FORCE_EVAL/SUBSYS').subsections
        kinds = {s.params: s for s in subsys if s.name == "KIND"}
        for elem in set(self.atoms.get_chemical_symbols()):
            if elem not in kinds.keys():
                s = InputSection(name='KIND', params=elem)
                subsys.append(s)
                kinds[elem] = s
            if p.basis_set:
                kinds[elem].keywords.append('BASIS_SET ' + p.basis_set)
            if potential:
                kinds[elem].keywords.append('POTENTIAL ' + potential)

        output_lines = ['!!! Generated by ASE !!!'] + root.write()
        return '\n'.join(output_lines)


class Cp2kShell:
    """Wrapper for CP2K-shell child-process"""

    def __init__(self, command, debug):
        """Construct CP2K-shell object"""

        self.isready = False
        self.version = 1.0  # assume oldest possible version until verified
        self._debug = debug

        # launch cp2k_shell child process
        assert 'cp2k_shell' in command
        if self._debug:
            print(command)
        self._child = subprocess.Popen(
            command, shell=True, universal_newlines=True,
            stdin=subprocess.PIPE, stdout=subprocess.PIPE, bufsize=1)
        self.expect('* READY')

        # check version of shell
        self.send('VERSION')
        line = self.recv()
        if not line.startswith('CP2K Shell Version:'):
            raise RuntimeError('Cannot determine version of CP2K shell.  '
                               'Probably the shell version is too old.  '
                               'Please update to CP2K 3.0 or newer.  '
                               f'Received: {line}')

        shell_version = line.rsplit(":", 1)[1]
        self.version = float(shell_version)
        assert self.version >= 1.0

        self.expect('* READY')

        # enable harsh mode, stops on any error
        self.send('HARSH')
        self.expect('* READY')

    def __del__(self):
        """Terminate cp2k_shell child process"""
        self.close()

    def close(self):
        """Terminate cp2k_shell child process"""
        if self.isready:
            self.send('EXIT')
            self._child.communicate()
            rtncode = self._child.wait()
            assert rtncode == 0  # child process exited properly?
        elif getattr(self, '_child', None) is not None:
            warn('CP2K-shell not ready, sending SIGTERM.', RuntimeWarning)
            self._child.terminate()
            self._child.communicate()
        self._child = None
        self.version = None
        self.isready = False

    def send(self, line):
        """Send a line to the cp2k_shell"""
        assert self._child.poll() is None  # child process still alive?
        if self._debug:
            print('Sending: ' + line)
        if self.version < 2.1 and len(line) >= 80:
            raise Exception('Buffer overflow, upgrade CP2K to r16779 or later')
        assert len(line) < 800  # new input buffer size
        self.isready = False
        self._child.stdin.write(line + '\n')

    def recv(self):
        """Receive a line from the cp2k_shell"""
        assert self._child.poll() is None  # child process still alive?
        line = self._child.stdout.readline().strip()
        if self._debug:
            print('Received: ' + line)
        self.isready = line == '* READY'
        return line

    def expect(self, line):
        """Receive a line and asserts that it matches the expected one"""
        received = self.recv()
        assert received == line


class InputSection:
    """Represents a section of a CP2K input file"""

    def __init__(self, name, params=None):
        self.name = name.upper()
        self.params = params
        self.keywords = []
        self.subsections = []

    def write(self):
        """Outputs input section as string"""
        output = []
        for k in self.keywords:
            output.append(k)
        for s in self.subsections:
            if s.params:
                output.append(f'&{s.name} {s.params}')
            else:
                output.append(f'&{s.name}')
            for l in s.write():
                output.append(f'   {l}')
            output.append(f'&END {s.name}')
        return output

    def add_keyword(self, path, line, unique=True):
        """Adds a keyword to section."""
        parts = path.upper().split('/', 1)
        candidates = [s for s in self.subsections if s.name == parts[0]]
        if len(candidates) == 0:
            s = InputSection(name=parts[0])
            self.subsections.append(s)
            candidates = [s]
        elif len(candidates) != 1:
            raise Exception(f'Multiple {parts[0]} sections found ')

        key = line.split()[0].upper()
        if len(parts) > 1:
            candidates[0].add_keyword(parts[1], line, unique)
        elif key == '_SECTION_PARAMETERS_':
            if candidates[0].params is not None:
                msg = f'Section parameter of section {parts[0]} already set'
                raise Exception(msg)
            candidates[0].params = line.split(' ', 1)[1].strip()
        else:
            old_keys = [k.split()[0].upper() for k in candidates[0].keywords]
            if unique and key in old_keys:
                msg = 'Keyword %s already present in section %s'
                raise Exception(msg % (key, parts[0]))
            candidates[0].keywords.append(line)

    def get_subsection(self, path):
        """Finds a subsection"""
        parts = path.upper().split('/', 1)
        candidates = [s for s in self.subsections if s.name == parts[0]]
        if len(candidates) > 1:
            raise Exception(f'Multiple {parts[0]} sections found ')
        if len(candidates) == 0:
            s = InputSection(name=parts[0])
            self.subsections.append(s)
            candidates = [s]
        if len(parts) == 1:
            return candidates[0]
        return candidates[0].get_subsection(parts[1])


def parse_input(inp):
    """Parses the given CP2K input string"""
    root_section = InputSection('CP2K_INPUT')
    section_stack = [root_section]

    for line in inp.split('\n'):
        line = line.split('!', 1)[0].strip()
        if len(line) == 0:
            continue

        if line.upper().startswith('&END'):
            s = section_stack.pop()
        elif line[0] == '&':
            parts = line.split(' ', 1)
            name = parts[0][1:]
            if len(parts) > 1:
                s = InputSection(name=name, params=parts[1].strip())
            else:
                s = InputSection(name=name)
            section_stack[-1].subsections.append(s)
            section_stack.append(s)
        else:
            section_stack[-1].keywords.append(line)

    return root_section