#
# @BEGIN LICENSE
#
# Psi4: an open-source quantum chemistry software package
#
# Copyright (c) 2007-2022 The Psi4 Developers.
#
# The copyrights for code used from other parties are included in
# the corresponding files.
#
# This file is part of Psi4.
#
# Psi4 is free software; you can redistribute it and/or modify
# it under the terms of the GNU Lesser General Public License as published by
# the Free Software Foundation, version 3.
#
# Psi4 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 Lesser General Public License for more details.
#
# You should have received a copy of the GNU Lesser General Public License along
# with Psi4; if not, write to the Free Software Foundation, Inc.,
# 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA.
#
# @END LICENSE
#
r"""Module to provide mechanism to store and restore option states in driver.

"""
import sys
from contextlib import contextmanager

from psi4 import core

from .exceptions import ValidationError


class OptionState(object):
    """Class to store the state of a single *option*. If *module* given, the *option*
    value and has_changed value is stored for global, local to *module*, and used by
    *module* scopes; otherwise (used for BASIS keywords), only global scope is stored.
    Class can store, print, and restore option values. ::

        >>> OptionState('E_CONVERGENCE', 'SCF')

        >>> print(OptionState('DF_BASIS_MP2'))

    """

    def __init__(self, option, module=None):
        self.option = option.upper()
        if module:
            self.module = module.upper()
        else:
            self.module = None

        self.value_global = core.get_global_option(option)
        self.haschanged_global = core.has_global_option_changed(option)
        if self.module:
            self.value_local = core.get_local_option(self.module, option)
            self.haschanged_local = core.has_local_option_changed(self.module, option)
            self.value_used = core.get_option(self.module, option)
            self.haschanged_used = core.has_option_changed(self.module, option)
        else:
            self.value_local = None
            self.haschanged_local = None
            self.value_used = None
            self.haschanged_used = None

    def __str__(self):
        text = ''
        if self.module:
            text += """  ==> %s Option in Module %s <==\n\n""" % (self.option, self.module)
            text += """  Global (has changed?) value: %7s %s\n""" % ('(' + str(self.haschanged_global) + ')',
                                                                     self.value_global)
            text += """  Local (has changed?) value:  %7s %s\n""" % ('(' + str(self.haschanged_local) + ')',
                                                                     self.value_local)
            text += """  Used (has changed?) value:   %7s %s\n""" % ('(' + str(self.haschanged_used) + ')',
                                                                     self.value_used)
        else:
            text += """  ==> %s Option in Global Scope <==\n\n""" % (self.option)
            text += """  Global (has changed?) value: %7s %s\n""" % ('(' + str(self.haschanged_global) + ')',
                                                                     self.value_global)
        text += """\n"""
        return text

    def restore(self):
        core.set_global_option(self.option, self.value_global)
        if not self.haschanged_global:
            core.revoke_global_option_changed(self.option)
        if self.module:
            core.set_local_option(self.module, self.option, self.value_local)
            if not self.haschanged_local:
                core.revoke_local_option_changed(self.module, self.option)


class OptionsState(object):
    """Class to contain multiple :py:func:`~psi4.driver.p4util.OptionState` objects.
    Used in python driver functions to collect several options before altering
    them, then restoring before function return. ::

        >>> optstash = OptionsState(
                ['DF_BASIS_SCF'],
                ['SCF_TYPE'],
                ['SCF', 'REFERENCE'])

        >>> print(optstash)

        >>> optstash.restore()

    """

    def __init__(self, *largs):
        self.data = {}
        for item in largs:
            self.add_option(item)

    def add_option(self, item):
        if len(item) == 2:
            key = (item[1], item[0])
        elif len(item) == 1:
            key = (item[0], )
        else:
            raise ValidationError(
                'Each argument to OptionsState should be an array, the first element of which is     the module scope and the second element of which is the module name. Bad argument: %s'
                % (item))

        if key in self.data:
            raise ValidationError(
                'Malformed options state, duplicate key adds of "{}". This should not happen, please raise a issue on github.com/psi4/psi4'.format(key))
        else:
            self.data[key] = OptionState(*key)

    def __str__(self):
        text = ''
        for key, item in self.data.items():
            text += str(item)
        return text

    def restore(self):
        for key, item in self.data.items():
            item.restore()


@contextmanager
def OptionsStateCM(osd):
    oso = OptionsState(osd)
    yield
    oso.restore()
