# ____________________________________________________________________________________
#
# 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 logging
import glob

from pyomo.common.tempfiles import TempfileManager
from pyomo.solvers.plugins.solvers.ASL import ASL

logger = logging.getLogger('pyomo.contrib.gjh')


def readgjh(fname=None):
    """
    Build objective gradient and constraint Jacobian
    from gjh file written by the ASL gjh 'solver'.

    gjh solver may be called using pyomo and 'keepfiles=True'.
    Enable 'symbolic_solver_labels' as well to write
    .row and col file to get variable mappings.

    Parameters
    ----------
    fname : string, optional
        gjh file name. The default is None.

    Returns
    -------
    g : list
        Current objective gradient.
    J : list
        Current objective Jacobian.
    H : list
        Current objective Hessian.
    variableList : list
        Variables as defined by `*.col` file.
    constraintList : list
        Constraints as defined by `*.row` file.

    """
    if fname is None:
        files = list(glob.glob("*.gjh"))
        fname = files.pop(0)
        if len(files) > 1:
            print("WARNING: More than one gjh file in current directory")
            print("  Processing: %s\nIgnoring: %s" % (fname, '\n\t\t'.join(files)))

    with open(fname, "r") as f:
        line = "dummy_str"
        # Skip top lines until g value is reached
        while line != "param g :=\n":
            line = f.readline()

        line = f.readline()
        g = []
        # ; is the escape character
        while line[0] != ';':
            """
            When printed via ampl interface:
            ampl: display g;
            g [*] :=
                         1   0.204082
                         2   0.367347
                         3   0.44898
                         4   0.44898
                         5   0.244898
                         6  -0.173133
                         7  -0.173133
                         8  -0.0692532
                         9   0.0692532
                        10   0.346266
                        ;
            """
            index = int(line.split()[0]) - 1
            value = float(line.split()[1])
            g.append([index, value])
            line = f.readline()

        # Skip lines until J value is reached
        while line != "param J :=\n":
            line = f.readline()

        line = f.readline()
        J = []
        while line[0] != ';':
            """
            When printed via ampl interface:
            ampl: display J;
            J [*,*]
                :         1             2           3          4           5            6
                 :=
                1    -0.434327       0.784302     .          .           .          -0.399833
                2     2.22045e-16     .          1.46939     .           .          -0.831038
                3     0.979592        .           .         1.95918      .          -0.9596
                4     1.79592         .           .          .          2.12245     -0.692532
                5     0.979592        .           .          .           .           0
                6      .            -0.0640498   0.545265    .           .            .
                7      .             0.653061     .         1.14286      .            .
                8      .             1.63265      .          .          1.63265       .
                9      .             1.63265      .          .           .            .
                10     .              .          0.262481   0.262481     .            .
                11     .              .          1.14286     .          0.653061      .
                12     .              .          1.95918     .           .            .
                13     .              .           .         0.545265   -0.0640498     .
                14     .              .           .         1.95918      .            .
                15     .              .           .          .          1.63265       .
                16     .              .           .          .           .          -1

                :        7           8           9          10       :=
                1     0.399833     .           .          .
                2      .          0.831038     .          .
                3      .           .          0.9596      .
                4      .           .           .         0.692532
                6    -0.799667    0.799667     .          .
                7    -1.38506      .          1.38506     .
                8    -1.33278      .           .         1.33278
                9     0            .           .          .
                10     .         -0.9596      0.9596      .
                11     .         -1.38506      .         1.38506
                12     .          0            .          .
                13     .           .         -0.799667   0.799667
                14     .           .          0           .
                15     .           .           .         0
                16    1            .           .          .
                17   -1           1            .          .
                18     .         -1           1           .
                19     .           .         -1          1
            ;
            """
            if line[0] == '[':
                # Jacobian row index
                row = int(''.join(filter(str.isdigit, line))) - 1
                line = f.readline()

            column = int(line.split()[0]) - 1
            value = float(line.split()[1])
            J.append([row, column, value])
            line = f.readline()

        while line != "param H :=\n":
            line = f.readline()

        line = f.readline()
        H = []
        while line[0] != ';':
            """
            When printed via ampl interface:
            ampl: display H;
                H [*,*]
                :       1           2           3           4           5           6         :=
                1      .         0.25         .           .           .         -0.35348
                2     0.25        .          0.25         .           .         -0.212088
                3      .         0.25         .          0.25         .           .
                4      .          .          0.25         .          0.25         .
                5      .          .           .          0.25         .           .
                6    -0.35348   -0.212088     .           .           .         -0.0999584
                7     0.35348   -0.212088   -0.35348      .           .          0.0999584
                8      .         0.424176   -0.070696   -0.424176     .           .
                9      .          .          0.424176    0.070696   -0.424176     .
                10     .          .           .          0.35348     0.424176     .

                :        7            8           9          10        :=
                1     0.35348       .           .           .
                2    -0.212088     0.424176     .           .
                3    -0.35348     -0.070696    0.424176     .
                4      .          -0.424176    0.070696    0.35348
                5      .            .         -0.424176    0.424176
                6     0.0999584     .           .           .
                7    -0.299875     0.199917     .           .
                8     0.199917    -0.439817    0.2399       .
                9      .           0.2399     -0.439817    0.199917
                10     .            .          0.199917   -0.199917
                ;
            """
            if line[0] == '[':
                # Hessian row index
                row = int(''.join(filter(str.isdigit, line))) - 1
                line = f.readline()

            column = int(line.split()[0]) - 1
            value = float(line.split()[1])
            H.append([row, column, value])
            line = f.readline()

    with open(fname[:-3] + 'col', 'r') as f:
        data = f.read()
        variableList = data.split()

    with open(fname[:-3] + 'row', 'r') as f:
        data = f.read()
        constraintList = data.split()

    return g, J, H, variableList, constraintList


class GJHSolver(ASL):
    """
    An interface to the AMPL GJH "solver" for evaluating a model at a
    point.
    """

    def __init__(self, **kwds):
        kwds['type'] = 'gjh'
        kwds['symbolic_solver_labels'] = True
        super().__init__(**kwds)
        self.options.solver = 'gjh'
        self._metasolver = False

    # A hackish way to hold on to the model so that we can parse the
    # results.
    def _initialize_callbacks(self, model):
        self._model = model
        self._model._gjh_info = None
        super()._initialize_callbacks(model)

    def _presolve(self, *args, **kwds):
        super()._presolve(*args, **kwds)
        self._gjh_file = self._soln_file[:-3] + 'gjh'
        TempfileManager.add_tempfile(self._gjh_file, exists=False)

    def _postsolve(self):
        #
        # TODO: We should return the information using a better data
        # structure (ComponentMap?) so that the GJH solver does not need
        # to be called with symbolic_solver_labels=True
        #
        self._model._gjh_info = readgjh(self._gjh_file)
        self._model = None
        return super()._postsolve()
