
"""
Auto-generated by CVXPYgen on October 12, 2023 at 11:22:36.
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 import cpg_module


standard_settings_names = {"max_iters": "maxit"}


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", "A"]
    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"))
    n = param_dict["A"].shape[0]
    A_coordinates = np.unique([coord[0]+coord[1]*n for coord in param_dict["A"].attributes["sparsity"]])
    A_value = []
    A_flat = param_dict["A"].value.flatten(order="F")
    for coord in A_coordinates:
        A_value.append(A_flat[coord])
        A_flat[coord] = 0
    if np.sum(np.abs(A_flat)) > 0:
        warnings.warn('Ignoring nonzero value outside of sparsity pattern for parameter A!')
    par.A = list(A_value)

    # 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(2))
    prob.constraints[0].save_dual_value(np.array(res.cpg_dual.d0).reshape(3))

    # 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
