# ____________________________________________________________________________________
#
# Pyomo: Python Optimization Modeling Objects
# Copyright (c) 2008-2026 National Technology and Engineering Solutions of Sandia, LLC
# Under the terms of Contract DE-NA0003525 with National Technology and Engineering
# Solutions of Sandia, LLC, the U.S. Government retains certain rights in this
# software.  This software is distributed under the 3-clause BSD License.
# ____________________________________________________________________________________

import math
import enum
from pyomo.opt.results.container import MapContainer, ListContainer, ignore
from pyomo.common.collections import Bunch, OrderedDict

default_print_options = Bunch(
    schema=False,
    sparse=True,
    num_solutions=None,
    ignore_time=False,
    ignore_defaults=False,
)


class SolutionStatus(str, enum.Enum):
    bestSoFar = 'bestSoFar'
    error = 'error'
    feasible = 'feasible'
    globallyOptimal = 'globallyOptimal'
    infeasible = 'infeasible'
    locallyOptimal = 'locallyOptimal'
    optimal = 'optimal'
    other = 'other'
    stoppedByLimit = 'stoppedByLimit'
    unbounded = 'unbounded'
    unknown = 'unknown'
    unsure = 'unsure'

    # Overloading __str__ is needed to match the behavior of the old
    # pyutilib.enum class (removed June 2020). There are spots in the
    # code base that expect the string representation for items in the
    # enum to not include the class name. New uses of enum shouldn't
    # need to do this.
    def __str__(self):
        return self.value


intlist = (int,)
numlist = (float, int)


class Solution(MapContainer):
    def __init__(self):
        MapContainer.__init__(self)

        self.declare('gap')
        self.declare('status', value=SolutionStatus.unknown)
        self.declare('message')

        self.declare('problem', value={})
        self.declare('objective', value={})
        self.declare('variable', value={})
        self.declare('constraint', value={})

        self._option = default_print_options

    def load(self, repn):
        # delete key from dictionary, call base class load, handle variable loading.
        for field in ("Variable", "Constraint", "Problem", "Objective"):
            if field in repn:
                tmp_ = repn[field]
                del repn[field]
            if tmp_ == "No values":
                tmp_ = {}
            setattr(self, field.lower(), tmp_)
        MapContainer.load(self, repn)

    def pprint(self, ostream, option, from_list=False, prefix="", repn=None):
        #
        # the following is specialized logic for handling variable and
        # constraint maps - which are dictionaries of dictionaries, with
        # at a minimum an "id" element per sub-dictionary.
        #
        first = True
        for key, item in self.items():
            if not key in repn or key == 'Problem':
                continue
            if type(item.value) is not dict:
                #
                # Do a normal print
                #
                if first:
                    ostream.write(key + ": ")
                    first = False
                else:
                    ostream.write(prefix + key + ": ")
                item.pprint(ostream, option, prefix=prefix + "  ", repn=repn[key])
            elif len(item.value) == 0:
                #
                # The dictionary is empty
                #
                ostream.write(prefix + key + ": No values\n")
            else:
                print_zeros = key in ['Objective']
                first = True
                ostream.write(prefix + key + ":")
                prefix_ = prefix
                prefix = prefix + "  "
                #
                # Print values in the dictionary
                #
                value = item.value
                id_ctr = 0
                id_dict_map = {}
                id_name_map = (
                    {}
                )  # the name could be an integer or float - so convert prior to       printing (see code below)
                id_nonzeros_map = {}  # are any of the non-id entries are non-zero?
                entries_to_print = False

                for entry_name, entry_dict in value.items():
                    entry_id = id_ctr
                    id_ctr += 1
                    id_name_map[entry_id] = entry_name
                    id_dict_map[entry_id] = entry_dict
                    id_nonzeros_map[entry_id] = False  # until proven otherwise
                    for attr_name, attr_value in entry_dict.items():
                        if print_zeros or math.fabs(attr_value) > 1e-16:
                            id_nonzeros_map[entry_id] = True
                            entries_to_print = True

                if entries_to_print:
                    for entry_id in sorted(
                        id_dict_map.keys(), key=lambda id: id_name_map[id]
                    ):
                        if id_nonzeros_map[entry_id]:
                            if first:
                                ostream.write("\n")
                                first = False
                            ostream.write(prefix + str(id_name_map[entry_id]) + ":\n")
                            entry_dict = id_dict_map[entry_id]
                            for attr_name in sorted(entry_dict.keys()):
                                attr_value = entry_dict[attr_name]
                                if isinstance(attr_value, float) and (
                                    math.floor(attr_value) == attr_value
                                ):
                                    attr_value = int(attr_value)
                                ostream.write(
                                    prefix
                                    + "  "
                                    + attr_name.capitalize()
                                    + ": "
                                    + str(attr_value)
                                    + '\n'
                                )
                else:
                    ostream.write(" No nonzero values\n")
                prefix = prefix_


class SolutionSet(ListContainer):
    def __init__(self):
        ListContainer.__init__(self, Solution)
        self._option = default_print_options

    def _repn_(self, option):
        if not option.schema and not self._active and not self._required:
            return ignore
        if option.schema and len(self) == 0:
            self.add()
            self.add()
        if option.num_solutions is None:
            num = len(self)
        else:
            num = min(option.num_solutions, len(self))
        i = 0
        tmp = []
        for item in self._list:
            tmp.append(item._repn_(option))
            i = i + 1
            if i == num:
                break
        return [
            OrderedDict(
                [
                    ('number of solutions', len(self)),
                    ('number of solutions displayed', num),
                ]
            )
        ] + tmp

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

    def __call__(self, i=1):
        return self._list[i - 1]

    def pprint(self, ostream, option, prefix="", repn=None):
        if not option.schema and not self._active and not self._required:
            return ignore
        ostream.write("\n")
        ostream.write(prefix + "- ")
        spaces = ""
        for key in repn[0]:
            ostream.write(prefix + spaces + key + ": " + str(repn[0][key]) + '\n')
            spaces = "  "
        i = 0
        for i in range(len(self._list)):
            item = self._list[i]
            ostream.write(prefix + '- ')
            item.pprint(
                ostream, option, from_list=True, prefix=prefix + "  ", repn=repn[i + 1]
            )

    def load(self, repn):
        #
        # Note: we ignore the first element of the repn list, since
        # it was generated on the fly by the SolutionSet object.
        #
        for data in repn[1:]:  # repn items 1 through N are individual solutions.
            item = self.add()
            item.load(data)
