
"""
Auto-generated by CVXPYgen on January 14, 2025 at 10:11:50.
Content: Custom solve method for CVXPY interface.
"""

import time
import warnings
import numpy as np
from cvxpy.reductions import Solution
from cvxpy.problems.problem import SolverStats
from nonneg_LS_OSQP import cpg_module


standard_settings_names = {'warm_start': 'warm_starting'}


def squeeze_scalar(val):
    if isinstance(val, np.ndarray):
        val = val.squeeze()
        if val.shape == ():
            return val.item()
    return val


def cpg_solve(prob, updated_params=None, **kwargs):

    # set flags for updated parameters
    upd = cpg_module.cpg_updated()
    if updated_params is None:
        updated_params = ["b"]
    for p in updated_params:
        try:
            setattr(upd, p, True)
        except AttributeError:
            raise AttributeError(f"{p} is not a parameter.")

    # set solver settings
    cpg_module.set_solver_default_settings()
    for key, value in kwargs.items():
        try:
            eval(f'cpg_module.set_solver_{standard_settings_names.get(key, key)}(value)')
        except AttributeError:
            raise AttributeError(f'Solver setting "{key}" not available.')

    # set parameter values
    par = cpg_module.cpg_params()
    param_dict = prob.param_dict
    par.b = list(param_dict["b"].value.flatten(order="F"))

    # solve
    t0 = time.time()
    res = cpg_module.solve(upd, par)
    t1 = time.time()

    # store solution in problem object
    prob._clear_solution()
    prob.var_dict['x'].save_value(np.array(res.cpg_prim.x).reshape(4))
    prob.constraints[0].save_dual_value(np.array(res.cpg_dual.d0).reshape(4))
    prob.constraints[1].save_dual_value(np.array(res.cpg_dual.d1).reshape(4))
    prob.constraints[2].save_dual_value(np.array(res.cpg_dual.d2).reshape(4))

    # store additional solver information in problem object
    prob._status = res.cpg_info.status
    if abs(res.cpg_info.obj_val) == 1e30:
        prob._value = np.sign(res.cpg_info.obj_val) * np.inf
    else:
        prob._value = res.cpg_info.obj_val
    primal_vars = {var.id: var.value for var in prob.variables()}
    dual_vars = {c.id: c.dual_value for c in prob.constraints}
    solver_specific_stats = {'obj_val': res.cpg_info.obj_val,
                             'status': prob._status,
                             'iter': res.cpg_info.iter,
                             'pri_res': res.cpg_info.pri_res,
                             'dua_res': res.cpg_info.dua_res,
                             'time': res.cpg_info.time}
    attr = {'solve_time': t1 - t0, 'solver_specific_stats': solver_specific_stats, 'num_iters': res.cpg_info.iter}
    prob._solution = Solution(prob.status, prob.value, primal_vars, dual_vars, attr)
    results_dict = {'solver_specific_stats': solver_specific_stats,
                    'num_iters': res.cpg_info.iter,
                    'solve_time': t1 - t0}
    prob._solver_stats = SolverStats(results_dict, 'OSQP')

    return prob.value


def cpg_gradient(prob, gradient_sol_primal=None, gradient_sol_dual=None):

    # set gradient primal and dual solutions if provided
    gradient_sol = cpg_module.cpg_gsol()
    if gradient_sol_primal is not None and gradient_sol_dual is not None:
        gradient_sol.primal = list(gradient_sol_primal)
        gradient_sol.dual = list(gradient_sol_dual)
        use_sol = True
    else:
        gradient_sol.primal = [0] * 8
        gradient_sol.dual = [0] * 16
        use_sol = False

    # set variable deltas
    vdelta = cpg_module.cpg_vdelta()
    vdelta.x = list(prob.var_dict["x"].gradient.flatten(order="F"))
    pdelta = cpg_module.gradient(vdelta, gradient_sol, use_sol)
    prob.param_dict['b'].gradient = np.array(pdelta.b).reshape(4)


def forward(params, context):

    info = {}
    kwargs = context.solver_args.copy()
    prob = kwargs.pop("problem")
    parameters = prob.parameters()
    for pid, val in zip(context.param_ids, params):
        next(p for p in parameters if p.id == pid).value = val
    updated_params = kwargs.pop("updated_params", None)
    _, info["gradient_primal"], info["gradient_dual"] = cpg_solve_and_gradient_info(prob, updated_params, **kwargs)
    info["prob"] = prob

    vars = prob.variables()
    return [next(v for v in vars if v.id == variable.id).value for variable in context.variables], info


def backward(dvars, context):

    prob = context.info["prob"]
    vars = prob.variables()
    for variable, dv in zip(context.variables, dvars):
        next(v for v in vars if v.id == variable.id).gradient = dv
    gradient_primal = context.info["gradient_primal"]
    gradient_dual = context.info["gradient_dual"]
    cpg_gradient(prob, gradient_primal, gradient_dual)

    params = prob.parameters()
    return [next(p for p in params if p.id == pid).gradient for pid in context.param_ids], {}
