Source code for scm.pisa.block

from __future__ import annotations

import copy
import re
import textwrap
import warnings
from dataclasses import dataclass
from typing import TYPE_CHECKING, Any, Literal, Sequence, TypeVar

from scm.pisa.entry import Entry, _check_unique
from scm.pisa.enums import BlockType, InputCategory
from scm.pisa.key import FloatListKey, IntListKey, Key
from scm.pisa.utils import str_to_enum, to_valid_identifier

if TYPE_CHECKING:
    from scm.plams.core.settings import Settings

INDENT = " " * 2

T_Block = TypeVar("T_Block", bound="Block")


@dataclass(frozen=True)
class InputParserEngineBlock:
    """Simple container for data created by the InputParser for EngineBlocks. Only used in the :meth:`Block.from_text`
    functionality
    """

    _h: str
    _1: Sequence[str]


@dataclass(frozen=True)
class InputParserFreeBlock:
    """Simple container for data created by the InputParser for FreeBlocks. Only used in the :meth:`Block.from_text`
    functionality"""

    _1: Sequence[str]

    def __str__(self) -> str:
        return "\n".join(self._1) + "\n"


class Block(Entry):
    """Base class representing an input block for AMS."""

    _header = ""

    def __post_init__(self):
        """Recursively sets key and block attributes present in ``self._keyword_dict``"""
        if self._keyword_dict is not None:
            for key, val in self._keyword_dict.items():
                if key.startswith("_") or not isinstance(val, dict):
                    continue
                category = str_to_enum(val["_category"], InputCategory)
                if category == InputCategory.BLOCK:
                    setattr(self, key, Block.from_def_dict(name=key, def_dict=val))
                elif category == InputCategory.KEY:
                    setattr(self, key, Key.from_def_dict(name=key, def_dict=val))

    @property
    def header(self):
        return self._header

    @header.setter
    def header(self, val: str) -> None:
        """Set the header of this block.

        :param val: String header
        :type val: str
        :raises ValueError: Will raise exceptions if block is not an EngineBlock or does not allow headers
        """
        # adjusting the default value the allow_header attribute for EngineBlocks would require me to repeat
        # the entirety of the Entry.__init__ signature, which would take ~40 lines
        if self.allow_header or isinstance(self, EngineBlock):
            self._header = str(val)
        else:
            raise ValueError(f"Block {self.name} does not allow setting of header")

    @staticmethod
    def from_def_dict(name: str, def_dict: dict[str, Any]) -> Block:
        """Instantiate the relevant ``Block``  object from a definition dictionary from the json input definitions. Only
        used in the autogeneration of ``scm.input_classes``

        :param name: Name of the block to create
        :type name: str
        :param def_dict: Definition dictionary from the json input definitions.
        :type def_dict: dict[str, Any]
        :return: The created ``Block`` instance
        :rtype: Block
        """
        def_dict = dict(def_dict)
        assert str_to_enum(def_dict["_category"], InputCategory) == InputCategory.BLOCK
        block_type_lookup = {
            BlockType.FIXED: FixedBlock,
            BlockType.ENGINE: EngineBlock,
            BlockType.FREE: FreeBlock,
            BlockType.DRIVER: DriverBlock,
            BlockType.INPUT: InputBlock,
        }
        block_type = str_to_enum(def_dict["_type"], BlockType)
        assert block_type in block_type_lookup
        # category and type information already present in class
        def_dict.pop("_category")
        def_dict.pop("_type")
        # remove leading underscore to match constructor signature
        normal_kwargs = {k[1:]: v for k, v in def_dict.items() if k.startswith("_")}
        if len(normal_kwargs) == len(def_dict):
            keyword_dict = None
        else:
            keyword_dict = {k: v for k, v in def_dict.items() if k not in normal_kwargs}
        return block_type_lookup[block_type](name=name, **normal_kwargs, _keyword_dict=keyword_dict)

    @classmethod
    def from_settings(cls: type[T_Block], settings: Settings) -> T_Block:
        """Generate a block object from a settings object. Should only be used for backwards compatibility, e.g. old
        job classes or recipes that will break on receiving a PISA object. Not guaranteed to always work, since it does
        a to text input and from text input conversion. Examples:

        .. code:: python3

            from scm.input_classes import AMS, DFTB
            from scm.plams.interfaces.adfsuite.ams import AMSJob
            from scm.plams.core.settings import Settings

            settings = Settings()
            settings.input
            settings.input.ams.Task = 'GeometryOptimization'
            settings.input.ams.Properties.NormalModes = 'Yes'
            settings.input.DFTB.Model = 'SCC-DFTB'

            ams = AMS.from_settings(settings)

            # this is supported, but fragile:
            dftb = DFTB.from_settings(settings)
            # this is safer:
            ams = AMS.from_settings(settings)
            dftb = ams.engine


        :param cls: Block class you want to generate from settings. Usually should be AMS
        :type cls: type[T_Block]
        :param settings: Settings object to convert from. Should contain 'input' attribute.
        :type settings: Settings
        :raises ValueError: If settings input does not contain 'input' attribute
        :return: Block object with keys and blocks set according to information in settings object.
        :rtype: T_Block
        """
        # prevent circular import
        from scm.plams.interfaces.adfsuite.ams import AMSJob

        if "input" not in settings:
            raise ValueError(
                f"Settings object {settings} passed to {cls.from_settings} does not have 'input' attribute!"
            )

        job = AMSJob(settings=settings)
        text_input = job.get_input()
        return cls.from_text(text_input)

    @classmethod
    def from_text(cls: type[T_Block], text: str) -> T_Block:
        """Class method to instantiate a block object from text input. Need to call this on the relevant subclass to
        work correctly. E.g. for an AMS input string one needs to call AMS.from_text, not Block.from_text

        :param cls: Relevant block class
        :type cls: type[T_Block]
        :param text: String representing valid AMS input
        :type text: str
        :return: Block object representing the text input
        :rtype: T_Block
        """
        # prevent circular imports
        from scm.libbase import InputError, InputParser
        from scm.pisa.input_def import DRIVER_BLOCK_FILES, ENGINE_BLOCK_FILES

        try:
            # input_dict = InputParser().to_dict(program=cls.__name__.lower(), text_input=text, string_leafs=False)
            PROGRAMS = {value: key for key, value in ENGINE_BLOCK_FILES.items()}
            PROGRAMS.update({value: key for key, value in DRIVER_BLOCK_FILES.items()})
            program = PROGRAMS[cls.__name__].name.split(".")[0]
            input_dict = InputParser().to_dict(program=program, text_input=text, string_leafs=False)
        except InputError as e:
            # Try to filter out the relevant block. Probably breaks on headers, repeated blocks etc?
            re_match = re.search(f"(?<={cls.__name__.lower()}).*(?=End)", text, flags=re.I | re.M | re.S)
            if re_match:
                new_text = re_match.group()
                msg = f"Could not generate {cls} from textinput:\n{text} \nRetrying (fragile!) with subset:\n{new_text}"
                warnings.warn(msg)
                return cls.from_text(new_text)
            else:
                raise ValueError(f"Could not generate {cls} object from text input: \n\n {text}") from e
        input_object = cls()
        input_object.set_from_input_dict(input_dict)
        return input_object

    def set_from_input_dict(self, input_dict: dict[str, Any]) -> None:
        """Called from ``Block.from_text`` after instantiation of the Block to correctly set the values of the key/block
        attributes as defined in the text input.

        :param input_dict: dictionary returned by ``scm.libbase.InputParser``
        :type input_dict: dict[str, Any]
        :raises ValueError: When the text input references an Entry not present in the JSON input definitions, would be caused by faulty text input.
        """
        # only to be used by the dictionary returned by scm.libbase.InputParser
        for key, val in input_dict.items():
            try:
                attr = self._lowercase_dict_entry[to_valid_identifier(key, assert_success=False).lower()]
            except KeyError:
                raise ValueError(
                    f"Can not set {self.__class__.__name__} block from dict {input_dict}, unrecognized entry {key}"
                )
            if isinstance(attr, Key):
                multiple_entries = False
                if not attr.unique:
                    # special case, these are lists, so could be list of lists
                    if isinstance(attr, (FloatListKey, IntListKey)):
                        if isinstance(val[0], list):
                            multiple_entries = True
                    else:
                        if isinstance(val, list):
                            multiple_entries = True
                if multiple_entries:
                    for i, v in enumerate(val):
                        attr[i].val = v
                else:
                    attr.val = val
            elif isinstance(attr, Block):
                assert isinstance(val, (dict, list))
                # for engineblocks, instead of setting the values of it's keys, we need to overwrite it with an
                # instance of a specific engine.
                if isinstance(attr, EngineBlock):
                    if isinstance(val, list):
                        assert not attr.unique
                        for i, input_dict in enumerate(val):
                            attr[i] = EngineBlock._from_input_parser(InputParserEngineBlock(**input_dict))
                    elif isinstance(val, dict):
                        setattr(self, key.capitalize(), EngineBlock._from_input_parser(InputParserEngineBlock(**val)))
                else:
                    if isinstance(val, list):
                        assert not attr.unique
                        for i, input_dict in enumerate(val):
                            header = input_dict.pop("_h", None)
                            attr[i].set_from_input_dict(input_dict)
                            if header:
                                attr[i].header = header
                    elif isinstance(val, dict):
                        header = val.pop("_h", None)
                        attr.set_from_input_dict(val)
                        if header:
                            attr.header = header

    @property
    def blocks(self) -> tuple[Block, ...]:
        """Property returning all the blocks in this block

        :return: All the blocks in this block
        :rtype: tuple[Block, ...]
        """
        return tuple(
            sorted(
                (b for b in self.__dict__.values() if isinstance(b, Block)),
                key=lambda b: b.name,
            )
        )

    @property
    def keys(self) -> tuple[Key, ...]:
        """Property returning all the keys in this block

        :return: All the keys in this block
        :rtype: tuple[Key, ...]
        """
        return tuple(
            sorted(
                (b for b in self.__dict__.values() if isinstance(b, Key)),
                key=lambda b: b.name,
            )
        )

    @property
    def entries(self) -> tuple[Entry, ...]:
        """Property returning all the entries (blocks and keys) in this block. Return will be sorted with possible
        engine blocks at the end.

        :return: Tuple of all entries in this block
        :rtype: tuple[Entry, ...]
        """
        input_attrs = sorted(
            (attr for attr in self.__dict__.values() if isinstance(attr, Entry)),
            key=lambda attr: attr.name,
        )
        # if we have an engineblock, move it to the end of the list
        # because that seems to be the convention in the inputfile
        to_end = None
        for attr in input_attrs:
            if isinstance(attr, EngineBlock):
                to_end = attr
                break
        if to_end is not None:
            input_attrs.remove(to_end)
            input_attrs.append(to_end)
        return tuple(input_attrs)

    @property
    def _lowercase_dict_entry(self) -> dict[str, Entry]:
        """Lowercase dictionary of all Entries and their lowercase names, used to check if accessed ``Entry`` attributes
        in :meth:`__setattr__` are an incorrect case variation of an existing attributes.

        :return: Dictionary with lowercase entry names as keys, and entry attributes as values
        :rtype: dict[str, Entry]
        """
        return {k.lower(): v for k, v in self.__dict__.items() if isinstance(v, Entry)}

    @property
    def _dict_entry(self) -> dict[str, Entry]:
        """Dictionary of all Entries and their names, to find the proper case names of misspelled Entries
        in :meth:`__setattr__` .

        :return: Dictionary with entry names as keys, and entry attributes as values
        :rtype: dict[str, Entry]
        """
        return {k: v for k, v in self.__dict__.items() if isinstance(v, Entry)}

    def __setattr__(self, __name: str, __value: Any) -> None:
        """Overrides the default setattribute functionality for multiple added features:

        **Checking for case mistakes**
        The lookup of attributes is case sensitive, like the python standard, but will provide a helpful error message
        when you try to do ``block.temperature = 50`` instead of ``block.Temperature = 50``, by providing the user with
        the proper case name.

        **Allows overriding of attributes of type ``EngineBlock`` with an ``EngineBlock``**
        Normally, key and block attributes can not be overridden, since this would break the input system, but for
        engine blocks this is intended behavior.

        **Easier syntax for setting key values**
        Allows ``some_block.some_key = 10`` as an alias for ``some_block.some_key.val = 10``

        **Allows overwriting of ``FreeBlock`` attributes with multiline strings or sequences of strings**

        :param __name: Name of instance attribute to set the value of
        :type __name: str
        :param __value: Value to set for th attribute
        :type __value: Any
        :raises Exception: When trying to access a case variation of a key/block attribute
        """
        # check if name is case variation of existing entry
        if __name not in self.__dict__ and __name.lower() in self._lowercase_dict_entry:
            proper_case_name = None
            for name in self._dict_entry:
                if name.lower() == __name.lower():
                    proper_case_name = name
            raise Exception(f"Attribute '{__name}' not found, did you mean '{proper_case_name}'?")
        if not hasattr(self, __name):
            super().__setattr__(__name, __value)
            return
        else:
            attr = getattr(self, __name)
        # should be allowed to override engine block attributes with engine blocks
        if isinstance(attr, EngineBlock):
            if isinstance(__value, EngineBlock):
                __value._value_changed = True
                super().__setattr__(__name, __value)
            else:
                raise TypeError(f"Expect type 'EngineBlock', received type {type(__value)}: {repr(__value)}")
        elif isinstance(attr, FreeBlock):
            is_iterable = False
            try:
                iter(__value)
                is_iterable = True
            except TypeError:
                pass
            if isinstance(__value, str):
                attr._multiline_string = __value + "\n"
            elif is_iterable:
                attr._multiline_string = "\n".join(str(line) for line in __value) + "\n"
            else:
                raise TypeError(
                    f"Expected a string or an iterable of strings to assign to FreeBlock attribute {attr}. "
                    f"Received {__value}"
                )
            attr._value_changed = True
        elif isinstance(attr, Block):
            raise AttributeError(f"Not allowed to overwrite block attribute {attr}")
        # allow syntax like AMS().Task = 'SinglePoint'
        elif isinstance(attr, Key) and not isinstance(__value, Key):
            attr.val = __value
        else:
            super().__setattr__(__name, __value)

    @_check_unique
    def __setitem__(self, index: int, value: Block) -> None:
        """Adds support for setting heterogenous repeated blocks (currently only ``EngineBlock`` attributes):

        .. code:: python3

            driver = HYBRID()
            driver.Engine[0] = ADF()
            driver.Engine[1] = BAND()

        :param index: Index of the item of the repeated block. Must be used in sequential order, e.g. can't use [1] before [0] is used.
        :type index: int
        :param value: Block object for the repeated block.
        :type value: Block
        """
        assert isinstance(value, Block)
        try:
            self._entry_list[index] = value
        except IndexError:
            if index == len(self._entry_list):
                self._entry_list = [*self._entry_list, value]
            else:
                self[index]

    def _get_single_input_string(self, indent_lvl: int = 0) -> str:
        """A wrapper that concatenates the outputs of  ``self._get_block_start``, ``self._get_block_body``, and
        ``self._get_block_end``. Those are the methods to override in the subclasses.

        :param indent_lvl: Indentation level, defaults to 0
        :type indent_lvl: int, optional
        :return: Input string for a unique block.
        :rtype: str
        """
        input_str = ""
        input_str += self._get_block_start(indent_lvl)
        if len(self.header):
            input_str += " " + self.header
        input_str += "\n"
        input_str += self._get_block_body(indent_lvl + 1)
        input_str += self._get_block_end(indent_lvl)
        return input_str

    def _get_block_start(self, indent_lvl: int) -> str:
        """Inserts the block name

        :param indent_lvl: Indentation level
        :type indent_lvl: int
        :return: Properly indented block name
        :rtype: str
        """
        return f"{INDENT * indent_lvl}{self.name}"

    def _get_block_body(self, indent_lvl: int) -> str:
        """Recursively searches all the key/block attributes build the body by concatenating the outputs of
        ``get_input_string`` for attributes that were changed/touched by the user.

        :param indent_lvl: Indentation level
        :type indent_lvl: int
        :return: block body
        :rtype: str
        """
        body = ""
        for attr in self.entries:
            if attr.value_changed:
                body += attr.get_input_string(indent_lvl=indent_lvl)
        return body

    def _get_block_end(self, indent_lvl: int) -> str:
        """Simply returns the string `end` with required indentation.

        :param indent_lvl: Indentation level
        :type indent_lvl: int
        :return: `End`
        :rtype: str
        """
        return f"{INDENT * indent_lvl}End\n"

    @property
    def value_changed(self) -> bool:
        """Recursively searches all ``Entry`` attributes to see if any of them has ``value_changed=True``. We consider
        a key changed if it's value is changed by the user, even if it is set to the default value. A block is changed
        when any of its keys or block have their value changed. Used as a condition to include a block in the input
        string

        :return: Whether the value of this block is different from the default or touched by the user.
        :rtype: bool
        """
        if super().value_changed:
            return True
        for attr in self.entries:
            if attr.value_changed:
                return True
        return False

    def __add__(self: T_Block, other: T_Block) -> T_Block:
        """Implements a soft update of this block with another block via the '+' operator.

        :param self: Left hand object in the addition
        :type self: T_Block
        :param other: Right hand object in the addition
        :type other: T_Block
        :raises TypeError: When both blocks are not the same type
        :return: Returns a deepcopy of the lefthand block, where all the attributes that were set in the right hand block, but not in the left hand block, are copied over.
        :rtype: T_Block
        """

        if not type(self) is type(other):
            raise TypeError(f"Can only add blocks of the same type, not {self} and {other}")
        self_copy = copy.deepcopy(self)
        for name, attr in self_copy.__dict__.items():
            if isinstance(attr, Block):
                self_copy.__dict__[name] = attr + other.__dict__[name]
            elif isinstance(attr, Key):
                other_key = other.__dict__[name]
                if not attr.value_changed and attr.val != other_key.val:
                    attr.val = other_key.val
        return self_copy


class EngineBlock(Block):
    """Represents engine blocks. Input string will always start with 'Engine <engine_name>' and end with 'EndEngine'"""

    def _get_block_start(self, indent_lvl: int) -> str:
        """Prefixes the word `Engine` to the block start line containing the block name

        :param indent_lvl: Indentation level
        :type indent_lvl: int
        :return: <INDENT>Engine <block_name>
        :rtype: str
        """
        return f"\n{INDENT * indent_lvl}Engine {self.name}"

    def _get_block_end(self, indent_lvl: int) -> str:
        """Return `EndEngine` instead of `End`

        :param indent_lvl: Indentation level
        :type indent_lvl: int
        :return: <INDENT>EndEngine
        :rtype: str
        """
        return f"{INDENT * indent_lvl}EndEngine\n"

    @_check_unique
    def __getitem__(self, index: int) -> Any:
        return super().__getitem__(index)

    @staticmethod
    def _get_engine_class_and_header(header_str: str) -> tuple[type[EngineBlock], str | None]:
        """Only used when generating block objects from text input. Separates the engine name from the possible header
        and looks it up in ``scm.input_classes.engine``

        :param header_str: Line like ``BAND test_header`` that started the engine block in the text input
        :type header_str: str
        :raises ValueError: When the corresponding engine class can not be found
        :return: The Engine class and the possible header
        :rtype: tuple[type[EngineBlock], str | None]
        """
        engine_name = header_str.split()[0]

        # import here to prevent circular imports
        from scm import input_classes as ic

        lower_case_dir = [attr.lower() for attr in dir(ic)]
        try:
            engine_class = getattr(ic, dir(ic)[lower_case_dir.index(engine_name.lower())])
        except ValueError:
            raise ValueError(f"Unknown engine type {engine_name}, not present in {dir(ic)}")
        if len(engine_name) == len(header_str.strip()):
            return engine_class, None
        else:
            return engine_class, header_str[len(engine_name) :].lstrip()  # noqa: E203

    @staticmethod
    def _from_input_parser(d: InputParserEngineBlock) -> EngineBlock:
        """Create an ``EngineBlock`` object from the output of the ``scm.libbase.InputParser``

        :param d: Dataclass holding the engine block info
        :type d: InputParserEngineBlock
        :return: EngineBlock object
        :rtype: EngineBlock
        """
        engine_class, header = EngineBlock._get_engine_class_and_header(d._h)
        text_input = "\n".join(d._1)
        engine_block = engine_class.from_text(text_input)
        if header:
            engine_block.header = header
        return engine_block


class DriverBlock(Block):
    """Represents the highest level block."""

    def _get_block_start(self, indent_lvl: int) -> Literal[""]:
        """Driver blocks don't have start line or an end line, simply returns an empty string

        :param indent_lvl: Indentation level, ignored.
        :type indent_lvl: int
        :return: Empty string
        :rtype: Literal[""]
        """
        return ""

    def _get_block_end(self, indent_lvl: int) -> Literal[""]:
        """Driver blocks don't have start line or an end line, simply returns an empty string

        :param indent_lvl: Indentation level, ignored.
        :type indent_lvl: int
        :return: Empty string
        :rtype: Literal[""]
        """
        return ""

    def get_input_string(self, indent_lvl: int = -1) -> str:
        return super().get_input_string(indent_lvl).lstrip("\n")

    def to_settings(self) -> Settings:
        """Convert the Block to a PLAMS Settings object. Can be used for Job classes/recipes that don't support PISA
        objects. Example:

        .. code:: python3

            from scm.input_classes import AMS, DFTB
            from scm.plams.interfaces.adfsuite.ams import AMSJob
            from scm.plams.core.settings import Settings

            ams = AMS()
            ams.Task = 'SinglePoint'
            ams.Engine = DFTB()

            s = Settings()
            s.input = ams.to_settings()
            AMSJob(settings=s)


        This conversion works via serializing to text and using the libbase InputParser to create a settings object
        from text. Since this is a bit fragile, only recommended to use this method for backwards compatibility.

        :return: Settings object representation of the Block
        :rtype: Settings
        """
        # prevent circular imports
        from scm.libbase import InputParser
        from scm.pisa.input_def import DRIVER_BLOCK_FILES, ENGINE_BLOCK_FILES

        PROGRAMS = {value: key for key, value in ENGINE_BLOCK_FILES.items()}
        PROGRAMS.update({value: key for key, value in DRIVER_BLOCK_FILES.items()})
        program = PROGRAMS[self.__class__.__name__].name.split(".")[0]
        return InputParser().to_settings(program=program, text_input=self.get_input_string())


class FreeBlock(Block):
    """A block that can be assigned a multiline string of an iterable of strings, instead of directly assigning
    attributes. Will simply result in a series of lines in the input text
    """

    _multiline_string: str = ""

    def __post_init__(self) -> None:
        assert self._keyword_dict is None

    def _get_block_body(self, indent_lvl: int = 0) -> str:
        """Simply returns an indented version of the multiline string in :attr:`FreeBlock._multiline_string`

        :param indent_lvl: Indentation level, defaults to 0
        :type indent_lvl: int, optional
        :return: Indented version of multiline string.
        :rtype: str
        """
        return textwrap.indent(self._multiline_string, str(indent_lvl * INDENT))

    def set_from_input_dict(self, input_dict: dict[str, Any]) -> None:
        self._multiline_string = str(InputParserFreeBlock(**input_dict))
        self._value_changed = True


class FixedBlock(Block):
    """Standard block with a predefined set of possible keys and blocks."""

    pass


class InputBlock(FreeBlock):
    """Identical to the :class:`FreeBlock`, but end with `EndInput` instead of `End`"""

    def _get_block_end(self, indent_lvl: int) -> str:
        return f"{INDENT * indent_lvl}EndInput\n"


def type_hint(b: Block) -> str | None:
    """Produces the proper type hint for all ``Block`` classes. Only necessary for ``Block`` classes that are allowed to
    overwritten as an attribute, such as an ``EngineBlock`` or a ``FreeBlock`` with a sequence of strings.

    :param b: Block instance to produce type hint for
    :type b: Block
    :return: Proper type hint.
    :rtype: str | None
    """
    if isinstance(b, EngineBlock):
        return EngineBlock.__name__
    if isinstance(b, FreeBlock):
        return f"str | Sequence[str] | {FreeBlock.__name__}"
    if isinstance(b, InputBlock):
        return f"str | Sequence[str] | {InputBlock.__name__}"
    else:
        return None


ALL_BLOCK_CLASSES = [Block, FixedBlock, EngineBlock, InputBlock, FreeBlock, DriverBlock]