"""
Copyright 2025, the CVXPY developers

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
"""

from __future__ import annotations

from typing import TYPE_CHECKING

import numpy as np

from cvxpy.constraints import (
    Equality,
    Inequality,
)
from cvxpy.reductions.inverse_data import InverseData
from cvxpy.reductions.solvers.solver import Solver
from cvxpy.reductions.utilities import (
    lower_equality,
    lower_ineq_to_nonneg,
)

if TYPE_CHECKING:
    from cvxpy.problems.problem import Problem


class NLPsolver(Solver):
    """
    A non-linear programming (NLP) solver.
    """
    REQUIRES_CONSTR = False
    MIP_CAPABLE = False
    BOUNDED_VARIABLES = True

    def accepts(self, problem: Problem) -> bool:
        """
        Only accepts disciplined nonlinear programs.
        """
        return problem.is_dnlp()

    def apply(self, problem: Problem) -> tuple[dict, InverseData]:
        """
        Construct NLP problem data stored in a dictionary.
        The NLP has the following form

            minimize      f(x)
            subject to    g^l <= g(x) <= g^u
                          x^l <= x <= x^u
        where f and g are non-linear (and possibly non-convex) functions
        """
        problem, data, inv_data = self._prepare_data_and_inv_data(problem)

        return data, inv_data

    def _prepare_data_and_inv_data(
        self, problem: Problem
    ) -> tuple[Problem, dict, InverseData]:
        data = dict()
        bounds = Bounds(problem)
        inverse_data = InverseData(bounds.new_problem)
        inverse_data.offset = 0.0
        data["problem"] = bounds.new_problem
        data["cl"], data["cu"] = bounds.cl, bounds.cu
        data["lb"], data["ub"] = bounds.lb, bounds.ub
        data["x0"] = bounds.x0
        data["_bounds"] = bounds  # Store for deferred Oracles creation in solve_via_data
        return problem, data, inverse_data

class Bounds:
    """Extracts variable and constraint bounds from a CVXPY problem.

    Converts the problem into the standard NLP form::

        g^l <= g(x) <= g^u,   x^l <= x <= x^u

    Inequalities are lowered to nonneg form and equalities to zero
    constraints.  The resulting ``new_problem`` attribute holds the
    canonicalized problem used by the solver oracles.
    """

    def __init__(self, problem: Problem) -> None:
        self.problem = problem
        self.main_var = problem.variables()
        self.get_constraint_bounds()
        self.get_variable_bounds()
        self.construct_initial_point()

    def get_constraint_bounds(self) -> None:
        """
        Get constraint bounds for all constraints.
        Also converts inequalities to nonneg form,
        as well as equalities to zero constraints and forms
        a new problem from the canonicalized problem.
        """
        lower, upper = [], []
        new_constr = []
        for constraint in self.problem.constraints:
            if isinstance(constraint, Equality):
                lower.extend([0.0] * constraint.size)
                upper.extend([0.0] * constraint.size)
                new_constr.append(lower_equality(constraint))
            elif isinstance(constraint, Inequality):
                lower.extend([0.0] * constraint.size)
                upper.extend([np.inf] * constraint.size)
                new_constr.append(lower_ineq_to_nonneg(constraint))
        canonicalized_prob = self.problem.copy([self.problem.objective, new_constr])
        self.new_problem = canonicalized_prob
        self.cl = np.array(lower)
        self.cu = np.array(upper)

    def get_variable_bounds(self) -> None:
        """
        Get variable bounds for all variables.
        Uses the variable's get_bounds() method which handles bounds attributes,
        nonneg/nonpos attributes, and properly broadcasts scalar bounds.
        """
        var_lower, var_upper = [], []
        for var in self.main_var:
            # get_bounds() returns arrays broadcastable to var.shape
            # and handles all edge cases (scalar bounds, sign attributes, etc.)
            lb, ub = var.get_bounds()
            # Flatten in column-major (Fortran) order and convert to contiguous array
            # (broadcast_to creates read-only views that need to be copied)
            lb_flat = np.asarray(lb).flatten(order='F')
            ub_flat = np.asarray(ub).flatten(order='F')
            var_lower.extend(lb_flat)
            var_upper.extend(ub_flat)
        self.lb = np.array(var_lower)
        self.ub = np.array(var_upper)

    def construct_initial_point(self) -> None:
        """ Loop through all variables and collect the intial point."""
        x0 = []
        for var in self.main_var:
            if var.value is None:
                raise ValueError("Variable %s has no value. This is a bug and should be reported."
                                  % var.name())

            x0.append(np.atleast_1d(var.value).flatten(order='F'))
        self.x0 = np.concatenate(x0, axis=0)

class Oracles:
    """Oracle interface for NLP solvers using the C-based diff engine.

    Provides function and derivative oracles (objective, gradient, constraints,
    Jacobian, Hessian) by wrapping the ``C_problem`` class from the diff engine.

    Forward passes are cached per solver iteration: calling ``objective`` or
    ``constraints`` sets a flag so that ``gradient``/``jacobian``/``hessian``
    can reuse the cached forward values.  The ``intermediate`` callback resets
    these flags at the start of each new solver iteration.

    Sparsity structures (Jacobian and Hessian) are computed once on first
    access and cached for the lifetime of the object.
    """

    def __init__(
        self,
        problem: Problem,
        verbose: bool = True,
        use_hessian: bool = True,
    ) -> None:
        from cvxpy.reductions.solvers.nlp_solvers.diff_engine import C_problem

        self.c_problem = C_problem(problem, verbose=verbose)
        self.use_hessian = use_hessian

        # Always initialize Jacobian
        self.c_problem.init_jacobian_coo()

        # Only initialize Hessian if needed (not for quasi-Newton methods)
        if use_hessian:
            self.c_problem.init_hessian_coo_lower_tri()

        # Cached sparsity structures
        self._jac_structure: tuple[np.ndarray, np.ndarray] | None = None
        self._hess_structure: tuple[np.ndarray, np.ndarray] | None = None

    def objective(self, x: np.ndarray) -> float:
        """Returns the scalar value of the objective given x."""
        return self.c_problem.objective_forward(x)

    def gradient(self, x: np.ndarray) -> np.ndarray:
        """Returns the gradient of the objective with respect to x."""
        return self.c_problem.gradient()

    def constraints(self, x: np.ndarray) -> np.ndarray:
        """Returns the constraint values."""
        return self.c_problem.constraint_forward(x)

    def jacobian(self, x: np.ndarray) -> np.ndarray:
        """Returns the Jacobian values in COO format at the sparsity structure. """
        return self.c_problem.eval_jacobian_vals()

    def jacobianstructure(self) -> tuple[np.ndarray, np.ndarray]:
        """Returns the sparsity structure of the Jacobian."""
        if self._jac_structure is not None:
            return self._jac_structure

        rows, cols = self.c_problem.get_jacobian_sparsity_coo()
        self._jac_structure = (rows, cols)
        return self._jac_structure

    def hessian(self, x: np.ndarray, duals: np.ndarray, obj_factor: float) -> np.ndarray:
        """Returns the lower triangular Hessian values in COO format. """
        if not self.use_hessian:
            raise ValueError("Hessian oracle called but use_hessian is False. "
                             "This is a bug and should be reported.")

        return self.c_problem.eval_hessian_vals_coo_lower_tri(obj_factor, duals)

    def hessianstructure(self) -> tuple[np.ndarray, np.ndarray]:
        """Returns the COO sparsity structure of the lower part of the Hessian.
           The returned rows are ascending, and within each row the columns are
           ascending."""
        if not self.use_hessian:
            # IPOPT calls this function even when hessian_approximation='limited-memory',
            # so return empty structure
            return (np.array([]), np.array([]))

        if self._hess_structure is not None:
            return self._hess_structure

        rows, cols = self.c_problem.get_problem_hessian_sparsity_coo()
        self._hess_structure = (rows, cols)
        return self._hess_structure

    def update_params(self, problem) -> None:
        """Update parameter values in the C DAG from the problem's parameters.

        Sparsity structures remain valid after this call.
        """
        if not problem.parameters():
            raise ValueError("update_params called but problem has no parameters. "
                             "This is a bug and should be reported.")

        theta = np.concatenate([
            np.asarray(p.value, dtype=np.float64).flatten(order='F')
            for p in problem.parameters()
        ])
        self.c_problem.update_params(theta)
