<commit_before>"""
This module defines the building blocks of a CP2K input file. The cp2k input structure is essentially a collection
of "sections" which are similar to dictionary objects that activate modules of the cp2k executable, and then
"keywords" which adjust variables inside of those modules. For example, FORCE_EVAL section will activate CP2K's ability
to calculate forces, and inside FORCE_EVAL, the Keyword "METHOD can be set to "QS" to set the method of force evaluation
to be the quickstep (DFT) module.

A quick overview of the module:

-- Section class defines the basis of Cp2k input and contains methods for manipulating these objects similarly to Dicts.
-- Keyword class defines the keywords used inside of Section objects that changes variables in Cp2k program
-- Cp2kInput class is special instantiation of Section that is used to represent the full cp2k calculation input.
-- The rest of the classes are children of Section intended to make initialization of common sections easier.
"""

import re
import os
import copy
import textwrap
import warnings
from typing import Dict, List, Tuple, Union, Sequence
from monty.json import MSONable
from monty.io import zopen
from pymatgen.io.cp2k.utils import _postprocessor
from pymatgen import Lattice, Structure, Molecule

__author__ = "<NAME>"
__version__ = "0.2"
__email__ = "<EMAIL>"
__date__ = "January 2019"


class Keyword(MSONable):

    """
    Class representing a keyword argument in CP2K. Within CP2K Secitons, which activate features
    of the CP2K code, the keywords are arguments that control the functionality of that feature.
    For example, the section "FORCE_EVAL" activates the evaluation of forces/energies, but within
    "FORCE_EVAL" the keyword "METHOD" controls whether or not this will be done with, say,
    "Quickstep" (DFT) or "EIP" (empirical interatomic potential).
    """

    def __init__(
        self,
        name: str,
        *values,
        description: str = None,
        units: str = None,
        verbose: bool = True,
        repeats: bool = False,
    ):
        """
        Initializes a keyword. These Keywords and the value passed to them are sometimes as simple as KEYWORD VALUE,
        but can also be more elaborate such as KEYWORD [UNITS] VALUE1 VALUE2, which is why this class exists:
        to handle many values and control easy printing to an input file.

        Args:
            name (str): The name of this keyword. Must match an acceptable keyword from CP2K
            args: All non-keyword arguments after 'name' are interpreted as the values to set for
                this keyword. i.e: KEYWORD ARG1 ARG2 would provide two values to the keyword.
            description (str): The description for this keyword. This can make readability of
                input files easier for some. Default=None.
            units (str): The units for this keyword. If not specified, CP2K default units will be
                used. Consult manual for default units. Default=None.
            repeats (bool): Whether or not this keyword may be repeated. Default=False.
        """
        self.name = name
        self.values = values
        self.description = description
        self.repeats = repeats
        self.units = units
        self.verbose = verbose

    def __str__(self):
        return (
            self.name.__str__()
            + " "
            + ("[{}] ".format(self.units) if self.units else "")
            + " ".join(map(str, self.values))
            + (
                " ! " + self.description
                if (self.description and self.verbose)
                else ""
            )
        )

    def __eq__(self, other):
        if self.name.upper() == other.name.upper():
            v1 = [_.upper() if isinstance(_, str) else _ for _ in self.values]
            v2 = [_.upper() if isinstance(_, str) else _ for _ in other.values]
            if v1 == v2:
                if self.units == self.units:
                    return True
        return False

    def __add__(self, other):
        return KeywordList(keywords=[self, other])

    def __getitem__(self, item):
        return self.values[item]

    def as_dict(self):
        """
        Get a dictionary representation of the Keyword
        """
        d = {}
        d["@module"] = self.__class__.__module__
        d["@class"] = self.__class__.__name__
        d["name"] = self.name
        d["values"] = self.values
        d["description"] = self.description
        d["repeats"] = self.repeats
        d["units"] = self.units
        d["verbose"] = self.verbose
        return d

    def get_string(self):
        """
        String representation of Keyword
        """
        return self.__str__()

    @classmethod
    def from_dict(cls, d):
        """
        Initialise from dictonary
        """
        return Keyword(
            d["name"],
            *d["values"],
            description=d["description"],
            repeats=d["repeats"],
            units=d["units"],
            verbose=d["verbose"]
        )

    @staticmethod
    def from_string(s):
        """
        Initialise from a string
        """
        s = s.strip()
        if s.__contains__('!') or s.__contains__('#'):
            s, description = re.split('(?:!|#)', s)
            description = description.strip()
        else:
            description = None
        units = re.findall(r"\[(.*)\]", s) or [None]
        s = re.sub(r"\[(.*)\]", "", s)
        return Keyword(*map(_postprocessor, s.split()),
                       units=units[0],
                       description=description)

    def silence(self):
        """
        Deactivate the printing of this keyword's description.
        """
        self.verbose = False


class KeywordList(MSONable, Sequence):

    """
    Some keywords can be repeated, which makes accessing them via the normal dictionary
    methods a little unnatural. This class deals with this by defining a collection
    of same-named keywords that are accessed by one name.
    """

    def __init__(
            self,
            keywords: Sequence[Keyword]
    ):
        """
        Initializes a keyword. These Keywords and the value passed to them are sometimes as simple as KEYWORD VALUE,
        but can also be more elaborate such as KEYWORD [UNITS] VALUE1 VALUE2, which is why this class exists:
        to handle many values and control easy printing to an input file.

        Args:
            keywords: A list of keywords. Must all have the same name (case-insensitive)
        """
        assert all(k.name.upper() == keywords[0].name.upper() for k in keywords) if keywords else True
        self.name = keywords[0].name if keywords else None
        self.keywords = keywords

    def __str__(self):
        return self.get_string()

    def __eq__(self, other):
        return all(k == o for k, o in zip(self.keywords, other.keywords))

    def __add__(self, other):
        return self.extend(other)

    def __len__(self):
        return len(self.keywords)

    def __getitem__(self, item):
        return self.keywords[item]

    def append(self, item):
        """
        append the keyword list
        """
        self.keywords.append(item)

    def extend(self, l):
        """
        extend the keyword list
        """
        self.keywords.extend(l)

    def get_string(self, indent=0):
        """
        String representation of Keyword
        """
        return " \n".join(["\t" * indent + k.__str__() for k in self.keywords])

    def silence(self):
        """
        Silence all keywords in keyword list
        """
        for k in self.keywords:
            k.silence()


class Section(MSONable):

    """
    Basic input representation of input to Cp2k. Activates functionality inside of the Cp2k executable.
    """

    required_sections: tuple = ()
    required_keywords: tuple = ()

    def __init__(
        self,
        name: str,
        subsections: dict = None,
        repeats: bool = False,
        description: Union[str, None] = None,
        keywords: Dict = {},
        section_parameters: Union[List, Tuple] = [],
        location: str = None,
        verbose: bool = True,
        alias: str = None,
        **kwargs
    ):
        """
        Basic object representing a CP2K Section. Sections activate different parts of the calculation.
        For example, FORCE_EVAL section will activate CP2K's ability to calculate forces. Sections are
        described with the following:

        Args:
            name (str): The name of the section (must match name in CP2K)
            subsections (dict): A dictionary of subsections that are nested in this section. Format is
                {'NAME': Section(*args, **kwargs). The name you chose for 'NAME' to index that subsection
                does not *have* to be the same as the section's true name, but we recommend matching
                them. You can specify a blank dictionary if there are no subsections, or if you want to
                insert the subsections later.
            repeats (bool): Whether or not this section can be repeated. Most sections cannot. Default=False.
            description (str): Description of this section for easier readability/more descriptiveness
            keywords (list): the keywords to be set for this section. Each element should be a Keyword
                object. This can be more cumbersome than simply using kwargs for building a class in a script,
                but is more convenient for the class instantiations of CP2K sections (see below).
            section_parameters (list): the section parameters for this section. Section parameters are
                specialized keywords that modify the behavior of the section overall. Most
                sections do not have section parameters, but some do. Unlike normal Keywords, these are
                specified as strings and not as Keyword objects.
            location (str): the path to the section in the form 'SECTION/SUBSECTION1/SUBSECTION3', example for
                QS module: 'FORCE_EVAL/DFT/QS'. This location is used to automatically determine if a subsection
                requires a supersection to be activated.
            verbose (str): Controls how much is printed to Cp2k input files (Also see Keyword). If True, then
                a description of the section will be printed with it as a comment (if description is set).
                Default=True.

            kwargs are interpreted as keyword, value pairs and added to the keywords array as Keyword objects
        """

        self.name = name
        self.subsections = subsections if subsections is not None else {}
        self.repeats = repeats
        self.description = description
        self.keywords = keywords
        self.section_parameters = section_parameters
        self.location = location
        self.verbose = verbose
        self.alias = alias
        self.kwargs = kwargs
        for k, v in self.kwargs.items():
            self.keywords[k] = Keyword(k, v)

        for k in self.required_sections:
            if not self.check(k):
                raise UserWarning(
                    "WARNING: REQUIRED SECTION {} HAS NOT BEEN INITIALIZED".format(
                        k
                    )
                )
        for k in self.required_keywords:
            if k not in self.keywords:
                raise UserWarning(
                    "WARNING: REQUIRED KEYWORD {} HAS NOT BEEN PROVIDED".format(
                        k
                    )
                )

    def __str__(self):
        return self.get_string()

    def __eq__(self, d):
        d2 = copy.deepcopy(d)
        s2 = copy.deepcopy(self)
        d2.silence()
        s2.silence()
        return d2.as_dict() == s2.as_dict()

    def __deepcopy__(self, memodict={}):
        c = copy.deepcopy(self.as_dict())
        return getattr(
            __import__(c['@module'], globals(),
                       locals(), c['@class'], 0),
            c['@class']
        ).from_dict(copy.deepcopy(self.as_dict()))

    def __getitem__(self, d):
        for k in self.keywords:
            if k.upper() == d.upper():
                return self.keywords[k]
        for k in self.subsections:
            if k.upper() == d.upper():
                return self.subsections[k]
        raise KeyError

    def get(self, d, default=None):
        """
        Similar to get for dictionaries. This will attempt to retrieve the
        section or keyword matching d. Will not raise an error if d does not
        exist.

        Args:
             d: the key to retrieve, if present
             default: what to return if d is not found
        """
        for k in self.keywords:
            if k.upper() == d.upper():
                return self.keywords[k]
        for k in self.subsections:
            if k.upper() == d.upper():
                return self.subsections[k]
        return default

    def __setitem__(self, key, value):
        if isinstance(value, Section):
            if key in self.subsections:
                self.subsections[key] = value.__deepcopy__()
            else:
                self.insert(value)
        if not isinstance(value, (Keyword, KeywordList)):
            value = Keyword(key, value)
        if isinstance(value, (Keyword, KeywordList)):
            match = [k for k in self.keywords if key.upper() == k.upper()]
            if match:
                del self.keywords[match[0]]
            self.keywords[key] = value

    def __delitem__(self, key):
        """
        Delete section with name matching key OR delete all keywords<commit_msg>Conditional error<commit_after>"""
This module defines the building blocks of a CP2K input file. The cp2k input structure is essentially a collection
of "sections" which are similar to dictionary objects that activate modules of the cp2k executable, and then
"keywords" which adjust variables inside of those modules. For example, FORCE_EVAL section will activate CP2K's ability
to calculate forces, and inside FORCE_EVAL, the Keyword "METHOD can be set to "QS" to set the method of force evaluation
to be the quickstep (DFT) module.

A quick overview of the module:

-- Section class defines the basis of Cp2k input and contains methods for manipulating these objects similarly to Dicts.
-- Keyword class defines the keywords used inside of Section objects that changes variables in Cp2k program
-- Cp2kInput class is special instantiation of Section that is used to represent the full cp2k calculation input.
-- The rest of the classes are children of Section intended to make initialization of common sections easier.
"""

import re
import os
import copy
import textwrap
import warnings
from typing import Dict, List, Tuple, Union, Sequence
from monty.json import MSONable
from monty.io import zopen
from pymatgen.io.cp2k.utils import _postprocessor
from pymatgen import Lattice, Structure, Molecule

__author__ = "<NAME>"
__version__ = "0.2"
__email__ = "<EMAIL>"
__date__ = "January 2019"


class Keyword(MSONable):

    """
    Class representing a keyword argument in CP2K. Within CP2K Secitons, which activate features
    of the CP2K code, the keywords are arguments that control the functionality of that feature.
    For example, the section "FORCE_EVAL" activates the evaluation of forces/energies, but within
    "FORCE_EVAL" the keyword "METHOD" controls whether or not this will be done with, say,
    "Quickstep" (DFT) or "EIP" (empirical interatomic potential).
    """

    def __init__(
        self,
        name: str,
        *values,
        description: str = None,
        units: str = None,
        verbose: bool = True,
        repeats: bool = False,
    ):
        """
        Initializes a keyword. These Keywords and the value passed to them are sometimes as simple as KEYWORD VALUE,
        but can also be more elaborate such as KEYWORD [UNITS] VALUE1 VALUE2, which is why this class exists:
        to handle many values and control easy printing to an input file.

        Args:
            name (str): The name of this keyword. Must match an acceptable keyword from CP2K
            args: All non-keyword arguments after 'name' are interpreted as the values to set for
                this keyword. i.e: KEYWORD ARG1 ARG2 would provide two values to the keyword.
            description (str): The description for this keyword. This can make readability of
                input files easier for some. Default=None.
            units (str): The units for this keyword. If not specified, CP2K default units will be
                used. Consult manual for default units. Default=None.
            repeats (bool): Whether or not this keyword may be repeated. Default=False.
        """
        self.name = name
        self.values = values
        self.description = description
        self.repeats = repeats
        self.units = units
        self.verbose = verbose

    def __str__(self):
        return (
            self.name.__str__()
            + " "
            + ("[{}] ".format(self.units) if self.units else "")
            + " ".join(map(str, self.values))
            + (
                " ! " + self.description
                if (self.description and self.verbose)
                else ""
            )
        )

    def __eq__(self, other):
        if self.name.upper() == other.name.upper():
            v1 = [_.upper() if isinstance(_, str) else _ for _ in self.values]
            v2 = [_.upper() if isinstance(_, str) else _ for _ in other.values]
            if v1 == v2:
                if self.units == self.units:
                    return True
        return False

    def __add__(self, other):
        return KeywordList(keywords=[self, other])

    def __getitem__(self, item):
        return self.values[item]

    def as_dict(self):
        """
        Get a dictionary representation of the Keyword
        """
        d = {}
        d["@module"] = self.__class__.__module__
        d["@class"] = self.__class__.__name__
        d["name"] = self.name
        d["values"] = self.values
        d["description"] = self.description
        d["repeats"] = self.repeats
        d["units"] = self.units
        d["verbose"] = self.verbose
        return d

    def get_string(self):
        """
        String representation of Keyword
        """
        return self.__str__()

    @classmethod
    def from_dict(cls, d):
        """
        Initialise from dictonary
        """
        return Keyword(
            d["name"],
            *d["values"],
            description=d["description"],
            repeats=d["repeats"],
            units=d["units"],
            verbose=d["verbose"]
        )

    @staticmethod
    def from_string(s):
        """
        Initialise from a string
        """
        s = s.strip()
        if s.__contains__('!') or s.__contains__('#'):
            s, description = re.split('(?:!|#)', s)
            description = description.strip()
        else:
            description = None
        units = re.findall(r"\[(.*)\]", s) or [None]
        s = re.sub(r"\[(.*)\]", "", s)
        return Keyword(*map(_postprocessor, s.split()),
                       units=units[0],
                       description=description)

    def silence(self):
        """
        Deactivate the printing of this keyword's description.
        """
        self.verbose = False


class KeywordList(MSONable, Sequence):

    """
    Some keywords can be repeated, which makes accessing them via the normal dictionary
    methods a little unnatural. This class deals with this by defining a collection
    of same-named keywords that are accessed by one name.
    """

    def __init__(
            self,
            keywords: Sequence[Keyword]
    ):
        """
        Initializes a keyword. These Keywords and the value passed to them are sometimes as simple as KEYWORD VALUE,
        but can also be more elaborate such as KEYWORD [UNITS] VALUE1 VALUE2, which is why this class exists:
        to handle many values and control easy printing to an input file.

        Args:
            keywords: A list of keywords. Must all have the same name (case-insensitive)
        """
        assert all(k.name.upper() == keywords[0].name.upper() for k in keywords) if keywords else True
        self.name = keywords[0].name if keywords else None
        self.keywords = keywords

    def __str__(self):
        return self.get_string()

    def __eq__(self, other):
        return all(k == o for k, o in zip(self.keywords, other.keywords))

    def __add__(self, other):
        return self.extend(other)

    def __len__(self):
        return len(self.keywords)

    def __getitem__(self, item):
        return self.keywords[item]

    def append(self, item):
        """
        append the keyword list
        """
        self.keywords.append(item)

    def extend(self, l):
        """
        extend the keyword list
        """
        self.keywords.extend(l)

    def get_string(self, indent=0):
        """
        String representation of Keyword
        """
        return " \n".join(["\t" * indent + k.__str__() for k in self.keywords])

    def silence(self):
        """
        Silence all keywords in keyword list
        """
        for k in self.keywords:
            k.silence()


class Section(MSONable):

    """
    Basic input representation of input to Cp2k. Activates functionality inside of the Cp2k executable.
    """

    required_sections: tuple = ()
    required_keywords: tuple = ()

    def __init__(
        self,
        name: str,
        subsections: dict = None,
        repeats: bool = False,
        description: Union[str, None] = None,
        keywords: Dict = {},
        section_parameters: Union[List, Tuple] = [],
        location: str = None,
        verbose: bool = True,
        alias: str = None,
        **kwargs
    ):
        """
        Basic object representing a CP2K Section. Sections activate different parts of the calculation.
        For example, FORCE_EVAL section will activate CP2K's ability to calculate forces. Sections are
        described with the following:

        Args:
            name (str): The name of the section (must match name in CP2K)
            subsections (dict): A dictionary of subsections that are nested in this section. Format is
                {'NAME': Section(*args, **kwargs). The name you chose for 'NAME' to index that subsection
                does not *have* to be the same as the section's true name, but we recommend matching
                them. You can specify a blank dictionary if there are no subsections, or if you want to
                insert the subsections later.
            repeats (bool): Whether or not this section can be repeated. Most sections cannot. Default=False.
            description (str): Description of this section for easier readability/more descriptiveness
            keywords (list): the keywords to be set for this section. Each element should be a Keyword
                object. This can be more cumbersome than simply using kwargs for building a class in a script,
                but is more convenient for the class instantiations of CP2K sections (see below).
            section_parameters (list): the section parameters for this section. Section parameters are
                specialized keywords that modify the behavior of the section overall. Most
                sections do not have section parameters, but some do. Unlike normal Keywords, these are
                specified as strings and not as Keyword objects.
            location (str): the path to the section in the form 'SECTION/SUBSECTION1/SUBSECTION3', example for
                QS module: 'FORCE_EVAL/DFT/QS'. This location is used to automatically determine if a subsection
                requires a supersection to be activated.
            verbose (str): Controls how much is printed to Cp2k input files (Also see Keyword). If True, then
                a description of the section will be printed with it as a comment (if description is set).
                Default=True.

            kwargs are interpreted as keyword, value pairs and added to the keywords array as Keyword objects
        """

        self.name = name
        self.subsections = subsections if subsections is not None else {}
        self.repeats = repeats
        self.description = description
        self.keywords = keywords
        self.section_parameters = section_parameters
        self.location = location
        self.verbose = verbose
        self.alias = alias
        self.kwargs = kwargs
        for k, v in self.kwargs.items():
            self.keywords[k] = Keyword(k, v)

        for k in self.required_sections:
            if not self.check(k):
                raise UserWarning(
                    "WARNING: REQUIRED SECTION {} HAS NOT BEEN INITIALIZED".format(
                        k
                    )
                )
        for k in self.required_keywords:
            if k not in self.keywords:
                raise UserWarning(
                    "WARNING: REQUIRED KEYWORD {} HAS NOT BEEN PROVIDED".format(
                        k
                    )
                )

    def __str__(self):
        return self.get_string()

    def __eq__(self, d):
        d2 = copy.deepcopy(d)
        s2 = copy.deepcopy(self)
        d2.silence()
        s2.silence()
        return d2.as_dict() == s2.as_dict()

    def __deepcopy__(self, memodict={}):
        c = copy.deepcopy(self.as_dict())
        return getattr(
            __import__(c['@module'], globals(),
                       locals(), c['@class'], 0),
            c['@class']
        ).from_dict(copy.deepcopy(self.as_dict()))

    def __getitem__(self, d):
        for k in self.keywords:
            if k.upper() == d.upper():
                return self.keywords[k]
        for k in self.subsections:
            if k.upper() == d.upper():
                return self.subsections[k]
        raise KeyError

    def get(self, d, default=None):
        """
        Similar to get for dictionaries. This will attempt to retrieve the
        section or keyword matching d. Will not raise an error if d does not
        exist.

        Args:
             d: the key to retrieve, if present
             default: what to return if d is not found
        """
        for k in self.keywords:
            if k.upper() == d.upper():
                return self.keywords[k]
        for k in self.subsections:
            if k.upper() == d.upper():
                return self.subsections[k]
        return default

    def __setitem__(self, key, value):
        if isinstance(value, Section):
            if key in self.subsections:
                self.subsections[key] = value.__deepcopy__()
            else:
                self.insert(value)
        else:
            if not isinstance(value, (Keyword, KeywordList)):
                value = Keyword(key, value)
            elif isinstance(value, (Keyword, KeywordList)):
                match = [k for k in self.keywords if key.upper() == k.upper()]
                if match:
                    del self.keywords[match[0]]
                self.keywords[key] = value

    def __delitem__(self, key):
        """
        Delete section with name matching key OR delete all keywords
{'fraction_non_alphanumeric': 0.08433338211620078, 'fraction_numerical': 0.0037318893604566077, 'mean_word_length': 2.010573851745787, 'pattern_counts': {'":': 0, '<': 7, '<?xml version=': 0, '>': 7, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '2457711', 'n_tokens_mistral': 7203, 'n_tokens_neox': 6904, 'n_words': 2777}