/*
Copyright 2025 Haihao Lu
Copyright 2026 Hongpei Li

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.
*/

#include "internal_types.h"
#include "pdhcg.h"
#include "preconditioner.h"
#include "solver.h"
#include "solver_state.h"
#include "utils.h"
#include <cublas_v2.h>
#include <cuda_runtime.h>
#include <cusparse.h>
#include <math.h>
#include <stdbool.h>
#include <stdio.h>
#include <time.h>

static void initialize_sparse_component_obj(pdhg_solver_state_t *state,
                                            const qp_problem_t *problem) {
  state->quadratic_objective_term->objective_sparse_matrix =
      (cu_sparse_matrix_csr_t *)safe_malloc(sizeof(cu_sparse_matrix_csr_t));

  memset(state->quadratic_objective_term->objective_sparse_matrix, 0,
         sizeof(cu_sparse_matrix_csr_t));
  state->quadratic_objective_term->objective_sparse_matrix->num_rows =
      problem->num_variables;
  state->quadratic_objective_term->objective_sparse_matrix->num_cols =
      problem->num_variables;
  state->quadratic_objective_term->objective_sparse_matrix->num_nonzeros =
      problem->objective_sparse_matrix_num_nonzeros;

  ALLOC_AND_COPY_CSR(state->quadratic_objective_term->objective_sparse_matrix,
                     problem->objective_sparse_matrix, problem->num_variables,
                     problem->objective_sparse_matrix_num_nonzeros);

  CUSPARSE_CHECK(cusparseCreateCsr(
      &state->quadratic_objective_term->matQ, state->num_variables,
      state->num_variables,
      state->quadratic_objective_term->objective_sparse_matrix->num_nonzeros,
      state->quadratic_objective_term->objective_sparse_matrix->row_ptr,
      state->quadratic_objective_term->objective_sparse_matrix->col_ind,
      state->quadratic_objective_term->objective_sparse_matrix->val,
      CUSPARSE_INDEX_32I, CUSPARSE_INDEX_32I, CUSPARSE_INDEX_BASE_ZERO,
      CUDA_R_64F));

  size_t primal_spmv_buffer_size;
  CUSPARSE_CHECK(cusparseSpMV_bufferSize(
      state->sparse_handle, CUSPARSE_OPERATION_NON_TRANSPOSE, &HOST_ONE,
      state->quadratic_objective_term->matQ, state->vec_primal_sol, &HOST_ZERO,
      state->quadratic_objective_term->vec_primal_obj_prod, CUDA_R_64F,
      CUSPARSE_SPMV_CSR_ALG2, &primal_spmv_buffer_size));
  CUDA_CHECK(
      cudaMalloc(&state->quadratic_objective_term->primal_obj_spmv_buffer,
                 primal_spmv_buffer_size));
  CUSPARSE_CHECK(cusparseSpMV_preprocess(
      state->sparse_handle, CUSPARSE_OPERATION_NON_TRANSPOSE, &HOST_ONE,
      state->quadratic_objective_term->matQ, state->vec_primal_sol, &HOST_ZERO,
      state->quadratic_objective_term->vec_primal_obj_prod, CUDA_R_64F,
      CUSPARSE_SPMV_CSR_ALG2,
      state->quadratic_objective_term->primal_obj_spmv_buffer));
}

static void initialize_lowrank_component_obj(pdhg_solver_state_t *state,
                                             const qp_problem_t *problem) {
  state->quadratic_objective_term->num_rank_lowrank_obj =
      problem->num_rank_lowrank_obj;
  ALLOC_ZERO(state->quadratic_objective_term->Rx_product,
             problem->num_rank_lowrank_obj * sizeof(double));

  CUSPARSE_CHECK(cusparseCreateDnVec(
      &state->quadratic_objective_term->vec_primal_obj_prod,
      state->num_variables, state->quadratic_objective_term->primal_obj_product,
      CUDA_R_64F));
  CUSPARSE_CHECK(cusparseCreateDnVec(
      &state->quadratic_objective_term->vec_Rx_prod,
      state->quadratic_objective_term->num_rank_lowrank_obj,
      state->quadratic_objective_term->Rx_product, CUDA_R_64F));

  state->quadratic_objective_term->objective_lowrank_matrix =
      (cu_sparse_matrix_csr_t *)safe_malloc(sizeof(cu_sparse_matrix_csr_t));

  memset(state->quadratic_objective_term->objective_lowrank_matrix, 0,
         sizeof(cu_sparse_matrix_csr_t));
  state->quadratic_objective_term->objective_lowrank_matrix->num_rows =
      problem->num_rank_lowrank_obj;
  state->quadratic_objective_term->objective_lowrank_matrix->num_cols =
      problem->num_variables;
  state->quadratic_objective_term->objective_lowrank_matrix->num_nonzeros =
      problem->objective_lowrank_matrix_num_nonzeros;

  ALLOC_AND_COPY_CSR(state->quadratic_objective_term->objective_lowrank_matrix,
                     problem->objective_lowrank_matrix,
                     problem->num_rank_lowrank_obj,
                     problem->objective_lowrank_matrix_num_nonzeros);

  state->quadratic_objective_term->objective_lowrank_matrix_t =
      (cu_sparse_matrix_csr_t *)safe_malloc(sizeof(cu_sparse_matrix_csr_t));

  memset(state->quadratic_objective_term->objective_lowrank_matrix_t, 0,
         sizeof(cu_sparse_matrix_csr_t));

  state->quadratic_objective_term->objective_lowrank_matrix_t->num_rows =
      problem->num_variables;
  state->quadratic_objective_term->objective_lowrank_matrix_t->num_cols =
      problem->num_rank_lowrank_obj;
  state->quadratic_objective_term->objective_lowrank_matrix_t->num_nonzeros =
      problem->objective_lowrank_matrix_num_nonzeros;

  CUDA_CHECK(cudaMalloc(
      &state->quadratic_objective_term->objective_lowrank_matrix_t->row_ptr,
      (problem->num_variables + 1) * sizeof(int)));
  CUDA_CHECK(cudaMalloc(
      &state->quadratic_objective_term->objective_lowrank_matrix_t->col_ind,
      problem->objective_lowrank_matrix_num_nonzeros * sizeof(int)));
  CUDA_CHECK(cudaMalloc(
      &state->quadratic_objective_term->objective_lowrank_matrix_t->val,
      problem->objective_lowrank_matrix_num_nonzeros * sizeof(double)));

  size_t buffer_size = 0;
  void *buffer = nullptr;
  CUSPARSE_CHECK(cusparseCsr2cscEx2_bufferSize(
      state->sparse_handle,
      state->quadratic_objective_term->objective_lowrank_matrix->num_rows,
      state->quadratic_objective_term->objective_lowrank_matrix->num_cols,
      state->quadratic_objective_term->objective_lowrank_matrix->num_nonzeros,
      state->quadratic_objective_term->objective_lowrank_matrix->val,
      state->quadratic_objective_term->objective_lowrank_matrix->row_ptr,
      state->quadratic_objective_term->objective_lowrank_matrix->col_ind,
      state->quadratic_objective_term->objective_lowrank_matrix_t->val,
      state->quadratic_objective_term->objective_lowrank_matrix_t->row_ptr,
      state->quadratic_objective_term->objective_lowrank_matrix_t->col_ind,
      CUDA_R_64F, CUSPARSE_ACTION_NUMERIC, CUSPARSE_INDEX_BASE_ZERO,
      CUSPARSE_CSR2CSC_ALG_DEFAULT, &buffer_size));
  CUDA_CHECK(cudaMalloc(&buffer, buffer_size));

  CUSPARSE_CHECK(cusparseCsr2cscEx2(
      state->sparse_handle,
      state->quadratic_objective_term->objective_lowrank_matrix->num_rows,
      state->quadratic_objective_term->objective_lowrank_matrix->num_cols,
      state->quadratic_objective_term->objective_lowrank_matrix->num_nonzeros,
      state->quadratic_objective_term->objective_lowrank_matrix->val,
      state->quadratic_objective_term->objective_lowrank_matrix->row_ptr,
      state->quadratic_objective_term->objective_lowrank_matrix->col_ind,
      state->quadratic_objective_term->objective_lowrank_matrix_t->val,
      state->quadratic_objective_term->objective_lowrank_matrix_t->row_ptr,
      state->quadratic_objective_term->objective_lowrank_matrix_t->col_ind,
      CUDA_R_64F, CUSPARSE_ACTION_NUMERIC, CUSPARSE_INDEX_BASE_ZERO,
      CUSPARSE_CSR2CSC_ALG_DEFAULT, buffer));

  CUDA_CHECK(cudaFree(buffer));

  CUSPARSE_CHECK(cusparseCreateCsr(
      &state->quadratic_objective_term->matR,
      state->quadratic_objective_term->num_rank_lowrank_obj,
      state->num_variables,
      state->quadratic_objective_term->objective_lowrank_matrix->num_nonzeros,
      state->quadratic_objective_term->objective_lowrank_matrix->row_ptr,
      state->quadratic_objective_term->objective_lowrank_matrix->col_ind,
      state->quadratic_objective_term->objective_lowrank_matrix->val,
      CUSPARSE_INDEX_32I, CUSPARSE_INDEX_32I, CUSPARSE_INDEX_BASE_ZERO,
      CUDA_R_64F));

  CUSPARSE_CHECK(cusparseCreateCsr(
      &state->quadratic_objective_term->matRt, state->num_variables,
      state->quadratic_objective_term->num_rank_lowrank_obj,
      state->quadratic_objective_term->objective_lowrank_matrix->num_nonzeros,
      state->quadratic_objective_term->objective_lowrank_matrix_t->row_ptr,
      state->quadratic_objective_term->objective_lowrank_matrix_t->col_ind,
      state->quadratic_objective_term->objective_lowrank_matrix_t->val,
      CUSPARSE_INDEX_32I, CUSPARSE_INDEX_32I, CUSPARSE_INDEX_BASE_ZERO,
      CUDA_R_64F));

  size_t primal_Rx_spmv_buffer_size;
  CUSPARSE_CHECK(cusparseSpMV_bufferSize(
      state->sparse_handle, CUSPARSE_OPERATION_NON_TRANSPOSE, &HOST_ONE,
      state->quadratic_objective_term->matR, state->vec_primal_sol, &HOST_ZERO,
      state->quadratic_objective_term->vec_Rx_prod, CUDA_R_64F,
      CUSPARSE_SPMV_CSR_ALG2, &primal_Rx_spmv_buffer_size));
  CUDA_CHECK(cudaMalloc(&state->quadratic_objective_term->Rx_spmv_buffer,
                        primal_Rx_spmv_buffer_size));
  CUSPARSE_CHECK(cusparseSpMV_preprocess(
      state->sparse_handle, CUSPARSE_OPERATION_NON_TRANSPOSE, &HOST_ONE,
      state->quadratic_objective_term->matR, state->vec_primal_sol, &HOST_ZERO,
      state->quadratic_objective_term->vec_Rx_prod, CUDA_R_64F,
      CUSPARSE_SPMV_CSR_ALG2, state->quadratic_objective_term->Rx_spmv_buffer));

  size_t primal_RRx_spmv_buffer_size;
  CUSPARSE_CHECK(cusparseSpMV_bufferSize(
      state->sparse_handle, CUSPARSE_OPERATION_NON_TRANSPOSE, &HOST_ONE,
      state->quadratic_objective_term->matRt,
      state->quadratic_objective_term->vec_Rx_prod, &HOST_ZERO,
      state->quadratic_objective_term->vec_primal_obj_prod, CUDA_R_64F,
      CUSPARSE_SPMV_CSR_ALG2, &primal_RRx_spmv_buffer_size));
  CUDA_CHECK(cudaMalloc(&state->quadratic_objective_term->RRx_spmv_buffer,
                        primal_RRx_spmv_buffer_size));
  CUSPARSE_CHECK(cusparseSpMV_preprocess(
      state->sparse_handle, CUSPARSE_OPERATION_NON_TRANSPOSE, &HOST_ONE,
      state->quadratic_objective_term->matRt,
      state->quadratic_objective_term->vec_Rx_prod, &HOST_ZERO,
      state->quadratic_objective_term->vec_primal_obj_prod, CUDA_R_64F,
      CUSPARSE_SPMV_CSR_ALG2,
      state->quadratic_objective_term->RRx_spmv_buffer));
}

static void initialize_quadratic_obj_term(pdhg_solver_state_t *state,
                                          const qp_problem_t *problem) {
  state->quadratic_objective_term = (quadratic_objective_term_t *)safe_malloc(
      sizeof(quadratic_objective_term_t));

  if ((!problem->objective_sparse_matrix ||
       problem->objective_sparse_matrix_num_nonzeros == 0) &&
      (!problem->objective_lowrank_matrix ||
       problem->objective_lowrank_matrix_num_nonzeros == 0)) {
    state->quadratic_objective_term->quad_obj_type = PDHCG_NON_Q;
  } else {
    state->quadratic_objective_term->quad_obj_type = detect_q_type(
        problem->objective_sparse_matrix, problem->objective_lowrank_matrix,
        problem->num_variables, problem->num_rank_lowrank_obj);
  }

  int n = problem->num_variables;
  if (state->quadratic_objective_term->quad_obj_type == PDHCG_NON_Q)
    return;

  size_t var_bytes = n * sizeof(double);
  ALLOC_ZERO(state->quadratic_objective_term->primal_obj_product, var_bytes);
  CUSPARSE_CHECK(cusparseCreateDnVec(
      &state->quadratic_objective_term->vec_primal_obj_prod,
      state->num_variables, state->quadratic_objective_term->primal_obj_product,
      CUDA_R_64F));

  switch (state->quadratic_objective_term->quad_obj_type) {
  case PDHCG_NON_Q:
    break;
  case PDHCG_DIAG_Q: {
    double *h_diag = (double *)safe_calloc(n, sizeof(double));

    CsrComponent *csr = problem->objective_sparse_matrix;
    if (csr->row_ptr && csr->col_ind && csr->val) {
      for (int i = 0; i < n; ++i) {
        for (int k = csr->row_ptr[i]; k < csr->row_ptr[i + 1]; ++k) {
          int col = csr->col_ind[k];
          if (col < n) {
            h_diag[col] = csr->val[k] + 1e-12;
          }
        }
      }
    }

    ALLOC_AND_COPY(state->quadratic_objective_term->diagonal_objective_matrix,
                   h_diag, n * sizeof(double));

    free(h_diag);
    state->quadratic_objective_term->objective_sparse_matrix = NULL;
    state->quadratic_objective_term->objective_lowrank_matrix = NULL;
    break;
  }

  case PDHCG_SPARSE_Q: {

    initialize_sparse_component_obj(state, problem);
    CUDA_CHECK(cudaGetLastError());

    state->quadratic_objective_term->diagonal_objective_matrix = NULL;
    break;
  }

  case PDHCG_LOW_RANK_Q: {

    initialize_lowrank_component_obj(state, problem);
    CUDA_CHECK(cudaGetLastError());

    state->quadratic_objective_term->diagonal_objective_matrix = NULL;
    break;
  }

  case PDHCG_LOW_RANK_PLUS_SPARSE_Q: {
    initialize_sparse_component_obj(state, problem);
    CUDA_CHECK(cudaGetLastError());
    initialize_lowrank_component_obj(state, problem);
    CUDA_CHECK(cudaGetLastError());

    state->quadratic_objective_term->diagonal_objective_matrix = NULL;
    break;
  }

  default:
    fprintf(stderr, "Error: Unknown Quadratic Objective Type detected.\n");
    exit(EXIT_FAILURE);
  }
}

static void initialize_inner_solver(pdhg_solver_state_t *state) {
  state->inner_solver =
      (inner_solver_t *)safe_calloc(1, sizeof(inner_solver_t));
  state->inner_solver->has_inner_loop = false;
  if (!(state->quadratic_objective_term->quad_obj_type == PDHCG_NON_Q ||
        state->quadratic_objective_term->quad_obj_type == PDHCG_DIAG_Q)) {
    ALLOC_ZERO(state->inner_solver->primal_buffer,
               state->num_variables * sizeof(double));
    ALLOC_ZERO(state->inner_solver->dual_buffer,
               state->num_constraints * sizeof(double));
  }
  switch (state->quadratic_objective_term->quad_obj_type) {
  case PDHCG_NON_Q:
    break;
  case PDHCG_DIAG_Q:
    break;
  case PDHCG_SPARSE_Q:
  case PDHCG_LOW_RANK_Q:
  case PDHCG_LOW_RANK_PLUS_SPARSE_Q:
    state->inner_solver->has_inner_loop = true;
    state->inner_solver->bb_step_size =
        (bb_step_size_t *)safe_malloc(sizeof(bb_step_size_t));
    ALLOC_ZERO(state->inner_solver->bb_step_size->gradient,
               state->num_variables * sizeof(double));
    ALLOC_ZERO(state->inner_solver->bb_step_size->direction,
               state->num_variables * sizeof(double));
    state->inner_solver->iteration_limit = 1000;
    state->inner_solver->tol = 1e-3;
    break;
  default:
    fprintf(stderr, "Error: Unknown Quadratic Objective Type detected.\n");
    exit(EXIT_FAILURE);
  }
}

__global__ void element_wise_mul_kernel(const double *__restrict__ A,
                                        const double *__restrict__ B,
                                        double *__restrict__ C, int n) {
  for (int idx = blockDim.x * blockIdx.x + threadIdx.x; idx < n;
       idx += blockDim.x * gridDim.x) {
    C[idx] = A[idx] * B[idx];
  }
}

void update_obj_product(pdhg_solver_state_t *state, double *primal_solution) {
  switch (state->quadratic_objective_term->quad_obj_type) {
  case PDHCG_NON_Q:
    return;
  case PDHCG_SPARSE_Q:
    CUSPARSE_CHECK(
        cusparseDnVecSetValues(state->vec_primal_sol, primal_solution));
    CUSPARSE_CHECK(cusparseDnVecSetValues(
        state->quadratic_objective_term->vec_primal_obj_prod,
        state->quadratic_objective_term->primal_obj_product));
    CUSPARSE_CHECK(cusparseSpMV(
        state->sparse_handle, CUSPARSE_OPERATION_NON_TRANSPOSE, &HOST_ONE,
        state->quadratic_objective_term->matQ, state->vec_primal_sol,
        &HOST_ZERO, state->quadratic_objective_term->vec_primal_obj_prod,
        CUDA_R_64F, CUSPARSE_SPMV_CSR_ALG2,
        state->quadratic_objective_term->primal_obj_spmv_buffer));
    return;
  case PDHCG_DIAG_Q:
    element_wise_mul_kernel<<<state->num_blocks_primal, THREADS_PER_BLOCK>>>(
        state->quadratic_objective_term->diagonal_objective_matrix,
        primal_solution, state->quadratic_objective_term->primal_obj_product,
        state->num_variables);
    return;
  case PDHCG_LOW_RANK_Q:
    CUSPARSE_CHECK(
        cusparseDnVecSetValues(state->vec_primal_sol, primal_solution));
    CUSPARSE_CHECK(cusparseSpMV(
        state->sparse_handle, CUSPARSE_OPERATION_NON_TRANSPOSE, &HOST_ONE,
        state->quadratic_objective_term->matR, state->vec_primal_sol,
        &HOST_ZERO, state->quadratic_objective_term->vec_Rx_prod, CUDA_R_64F,
        CUSPARSE_SPMV_CSR_ALG2,
        state->quadratic_objective_term->Rx_spmv_buffer));
    CUSPARSE_CHECK(
        cusparseSpMV(state->sparse_handle, CUSPARSE_OPERATION_NON_TRANSPOSE,
                     &HOST_ONE, state->quadratic_objective_term->matRt,
                     state->quadratic_objective_term->vec_Rx_prod, &HOST_ZERO,
                     state->quadratic_objective_term->vec_primal_obj_prod,
                     CUDA_R_64F, CUSPARSE_SPMV_CSR_ALG2,
                     state->quadratic_objective_term->RRx_spmv_buffer));
    return;
  case PDHCG_LOW_RANK_PLUS_SPARSE_Q:
    CUSPARSE_CHECK(
        cusparseDnVecSetValues(state->vec_primal_sol, primal_solution));

    CUSPARSE_CHECK(cusparseSpMV(
        state->sparse_handle, CUSPARSE_OPERATION_NON_TRANSPOSE, &HOST_ONE,
        state->quadratic_objective_term->matQ, state->vec_primal_sol,
        &HOST_ZERO, state->quadratic_objective_term->vec_primal_obj_prod,
        CUDA_R_64F, CUSPARSE_SPMV_CSR_ALG2,
        state->quadratic_objective_term->primal_obj_spmv_buffer));

    CUSPARSE_CHECK(cusparseSpMV(
        state->sparse_handle, CUSPARSE_OPERATION_NON_TRANSPOSE, &HOST_ONE,
        state->quadratic_objective_term->matR, state->vec_primal_sol,
        &HOST_ZERO, state->quadratic_objective_term->vec_Rx_prod, CUDA_R_64F,
        CUSPARSE_SPMV_CSR_ALG2,
        state->quadratic_objective_term->Rx_spmv_buffer));

    CUSPARSE_CHECK(
        cusparseSpMV(state->sparse_handle, CUSPARSE_OPERATION_NON_TRANSPOSE,
                     &HOST_ONE, state->quadratic_objective_term->matRt,
                     state->quadratic_objective_term->vec_Rx_prod, &HOST_ONE,
                     state->quadratic_objective_term->vec_primal_obj_prod,
                     CUDA_R_64F, CUSPARSE_SPMV_CSR_ALG2,
                     state->quadratic_objective_term->RRx_spmv_buffer));
    return;

  default:
    fprintf(stderr, "Error: Unknown Quadratic Objective Type detected.\n");
    exit(EXIT_FAILURE);
  }
}

double compute_xQx(pdhg_solver_state_t *state, double *primal_sol,
                   double *primal_obj_product) {
  if (state->quadratic_objective_term->quad_obj_type == PDHCG_NON_Q)
    return 0.0;

  double xQx = 0.0;
  CUBLAS_CHECK(cublasDdot(state->blas_handle, state->num_variables, primal_sol,
                          1, primal_obj_product, 1, &xQx));
  return xQx;
}

static void decide_problem_type(pdhg_solver_state_t *state) {
  if (state->quadratic_objective_term->quad_obj_type == PDHCG_NON_Q)
    state->problem_type = LP;
  else
    state->problem_type = CONVEX_QP;
}

void initialize_quadratic_term_information(pdhg_solver_state_t *state,
                                           const pdhg_parameters_t *params) {
  if (state->quadratic_objective_term->quad_obj_type == PDHCG_SPARSE_Q) {
    state->quadratic_objective_term->norm = estimate_maximum_eigenvalue(
        state->sparse_handle, state->blas_handle,
        state->quadratic_objective_term->objective_sparse_matrix,
        params->sv_max_iter, params->sv_tol);
    state->quadratic_objective_term->nonconvexity = estimate_minimum_eigenvalue(
        state->sparse_handle, state->blas_handle,
        state->quadratic_objective_term->objective_sparse_matrix,
        state->quadratic_objective_term->norm, params->sv_max_iter,
        params->sv_tol);
    return;
  }
  if (state->quadratic_objective_term->quad_obj_type == PDHCG_DIAG_Q) {
    double max_eigen = 0.0;
    double min_eigen = 0.0;
    double *temp_diag_host =
        (double *)malloc(state->num_variables * sizeof(double));
    cudaMemcpy(temp_diag_host,
               state->quadratic_objective_term->diagonal_objective_matrix,
               state->num_variables * sizeof(double), cudaMemcpyDeviceToHost);
    for (int i = 0; i < state->num_variables; i++) {
      double item = temp_diag_host[i];
      max_eigen = fmax(max_eigen, fabs(item));
      min_eigen = fmin(min_eigen, item);
    }
    state->quadratic_objective_term->norm = max_eigen;
    state->quadratic_objective_term->nonconvexity = min_eigen;
  }
}

pdhg_solver_state_t *
initialize_solver_state(const pdhg_parameters_t *params,
                        const qp_problem_t *working_problem,
                        const rescale_info_t *rescale_info) {
  pdhg_solver_state_t *state =
      (pdhg_solver_state_t *)safe_calloc(1, sizeof(pdhg_solver_state_t));

  int n_vars = working_problem->num_variables;
  int n_cons = working_problem->num_constraints;
  size_t var_bytes = n_vars * sizeof(double);
  size_t con_bytes = n_cons * sizeof(double);

  state->num_variables = n_vars;
  state->num_constraints = n_cons;
  state->objective_constant = working_problem->objective_constant;

  state->constraint_matrix =
      (cu_sparse_matrix_csr_t *)safe_malloc(sizeof(cu_sparse_matrix_csr_t));
  state->constraint_matrix_t =
      (cu_sparse_matrix_csr_t *)safe_malloc(sizeof(cu_sparse_matrix_csr_t));

  state->constraint_matrix->num_rows = n_cons;
  state->constraint_matrix->num_cols = n_vars;
  state->constraint_matrix->num_nonzeros =
      working_problem->constraint_matrix_num_nonzeros;

  state->constraint_matrix_t->num_rows = n_vars;
  state->constraint_matrix_t->num_cols = n_cons;
  state->constraint_matrix_t->num_nonzeros =
      working_problem->constraint_matrix_num_nonzeros;

  state->termination_reason = TERMINATION_REASON_UNSPECIFIED;

  state->rescaling_time_sec = rescale_info->rescaling_time_sec;

  ALLOC_AND_COPY_CSR(
      state->constraint_matrix, rescale_info->scaled_problem->constraint_matrix,
      rescale_info->scaled_problem->num_constraints,
      rescale_info->scaled_problem->constraint_matrix_num_nonzeros);

  CUDA_CHECK(cudaMalloc(&state->constraint_matrix_t->row_ptr,
                        (n_vars + 1) * sizeof(int)));
  CUDA_CHECK(
      cudaMalloc(&state->constraint_matrix_t->col_ind,
                 rescale_info->scaled_problem->constraint_matrix_num_nonzeros *
                     sizeof(int)));
  CUDA_CHECK(
      cudaMalloc(&state->constraint_matrix_t->val,
                 rescale_info->scaled_problem->constraint_matrix_num_nonzeros *
                     sizeof(double)));

  CUSPARSE_CHECK(cusparseCreate(&state->sparse_handle));
  CUBLAS_CHECK(cublasCreate(&state->blas_handle));
  CUBLAS_CHECK(
      cublasSetPointerMode(state->blas_handle, CUBLAS_POINTER_MODE_HOST));
  if (state->constraint_matrix->num_nonzeros > 0) {
    size_t buffer_size = 0;
    void *buffer = nullptr;
    CUSPARSE_CHECK(cusparseCsr2cscEx2_bufferSize(
        state->sparse_handle, state->constraint_matrix->num_rows,
        state->constraint_matrix->num_cols,
        state->constraint_matrix->num_nonzeros, state->constraint_matrix->val,
        state->constraint_matrix->row_ptr, state->constraint_matrix->col_ind,
        state->constraint_matrix_t->val, state->constraint_matrix_t->row_ptr,
        state->constraint_matrix_t->col_ind, CUDA_R_64F,
        CUSPARSE_ACTION_NUMERIC, CUSPARSE_INDEX_BASE_ZERO,
        CUSPARSE_CSR2CSC_ALG_DEFAULT, &buffer_size));
    CUDA_CHECK(cudaMalloc(&buffer, buffer_size));

    CUSPARSE_CHECK(cusparseCsr2cscEx2(
        state->sparse_handle, state->constraint_matrix->num_rows,
        state->constraint_matrix->num_cols,
        state->constraint_matrix->num_nonzeros, state->constraint_matrix->val,
        state->constraint_matrix->row_ptr, state->constraint_matrix->col_ind,
        state->constraint_matrix_t->val, state->constraint_matrix_t->row_ptr,
        state->constraint_matrix_t->col_ind, CUDA_R_64F,
        CUSPARSE_ACTION_NUMERIC, CUSPARSE_INDEX_BASE_ZERO,
        CUSPARSE_CSR2CSC_ALG_DEFAULT, buffer));

    CUDA_CHECK(cudaFree(buffer));
  } else {
    CUDA_CHECK(cudaMemset(state->constraint_matrix_t->row_ptr, 0,
                          (state->num_variables + 1) * sizeof(int)));
  }
  CUDA_CHECK(cudaGetLastError());
  ALLOC_AND_COPY(state->variable_lower_bound,
                 rescale_info->scaled_problem->variable_lower_bound, var_bytes);
  ALLOC_AND_COPY(state->variable_upper_bound,
                 rescale_info->scaled_problem->variable_upper_bound, var_bytes);
  ALLOC_AND_COPY(state->objective_vector,
                 rescale_info->scaled_problem->objective_vector, var_bytes);
  ALLOC_AND_COPY(state->constraint_lower_bound,
                 rescale_info->scaled_problem->constraint_lower_bound,
                 con_bytes);
  ALLOC_AND_COPY(state->constraint_upper_bound,
                 rescale_info->scaled_problem->constraint_upper_bound,
                 con_bytes);
  ALLOC_AND_COPY(state->constraint_rescaling, rescale_info->con_rescale,
                 con_bytes);
  ALLOC_AND_COPY(state->variable_rescaling, rescale_info->var_rescale,
                 var_bytes);

  state->constraint_bound_rescaling = rescale_info->con_bound_rescale;
  state->objective_vector_rescaling = rescale_info->obj_vec_rescale;

  ALLOC_ZERO(state->initial_primal_solution, var_bytes);
  ALLOC_ZERO(state->current_primal_solution, var_bytes);
  ALLOC_ZERO(state->pdhg_primal_solution, var_bytes);
  ALLOC_ZERO(state->reflected_primal_solution, var_bytes);
  ALLOC_ZERO(state->dual_product, var_bytes);
  ALLOC_ZERO(state->dual_slack, var_bytes);
  ALLOC_ZERO(state->dual_residual, var_bytes);
  ALLOC_ZERO(state->delta_primal_solution, var_bytes);

  ALLOC_ZERO(state->initial_dual_solution, con_bytes);
  ALLOC_ZERO(state->current_dual_solution, con_bytes);
  ALLOC_ZERO(state->pdhg_dual_solution, con_bytes);
  ALLOC_ZERO(state->reflected_dual_solution, con_bytes);
  ALLOC_ZERO(state->primal_product, con_bytes);
  ALLOC_ZERO(state->primal_slack, con_bytes);
  ALLOC_ZERO(state->primal_residual, con_bytes);
  ALLOC_ZERO(state->delta_dual_solution, con_bytes);

  if (working_problem->primal_start) {
    double *rescaled = (double *)safe_malloc(var_bytes);
    for (int i = 0; i < n_vars; ++i)
      rescaled[i] = working_problem->primal_start[i] *
                    rescale_info->var_rescale[i] *
                    rescale_info->con_bound_rescale;
    CUDA_CHECK(cudaMemcpy(state->initial_primal_solution, rescaled, var_bytes,
                          cudaMemcpyHostToDevice));
    CUDA_CHECK(cudaMemcpy(state->current_primal_solution, rescaled, var_bytes,
                          cudaMemcpyHostToDevice));
    CUDA_CHECK(cudaMemcpy(state->pdhg_primal_solution, rescaled, var_bytes,
                          cudaMemcpyHostToDevice));
    free(rescaled);
  }
  if (working_problem->dual_start) {
    double *rescaled = (double *)safe_malloc(con_bytes);
    for (int i = 0; i < n_cons; ++i)
      rescaled[i] = working_problem->dual_start[i] *
                    rescale_info->con_rescale[i] *
                    rescale_info->obj_vec_rescale;
    CUDA_CHECK(cudaMemcpy(state->initial_dual_solution, rescaled, con_bytes,
                          cudaMemcpyHostToDevice));
    CUDA_CHECK(cudaMemcpy(state->current_dual_solution, rescaled, con_bytes,
                          cudaMemcpyHostToDevice));
    CUDA_CHECK(cudaMemcpy(state->pdhg_dual_solution, rescaled, con_bytes,
                          cudaMemcpyHostToDevice));
    free(rescaled);
  }
  CUDA_CHECK(cudaGetLastError());
  double *temp_host = (double *)safe_malloc(fmax(var_bytes, con_bytes));
  for (int i = 0; i < n_cons; ++i)
    temp_host[i] =
        isfinite(rescale_info->scaled_problem->constraint_lower_bound[i])
            ? rescale_info->scaled_problem->constraint_lower_bound[i]
            : 0.0;
  ALLOC_AND_COPY(state->constraint_lower_bound_finite_val, temp_host,
                 con_bytes);
  for (int i = 0; i < n_cons; ++i)
    temp_host[i] =
        isfinite(rescale_info->scaled_problem->constraint_upper_bound[i])
            ? rescale_info->scaled_problem->constraint_upper_bound[i]
            : 0.0;
  ALLOC_AND_COPY(state->constraint_upper_bound_finite_val, temp_host,
                 con_bytes);
  for (int i = 0; i < n_vars; ++i)
    temp_host[i] =
        isfinite(rescale_info->scaled_problem->variable_lower_bound[i])
            ? rescale_info->scaled_problem->variable_lower_bound[i]
            : 0.0;
  ALLOC_AND_COPY(state->variable_lower_bound_finite_val, temp_host, var_bytes);
  for (int i = 0; i < n_vars; ++i)
    temp_host[i] =
        isfinite(rescale_info->scaled_problem->variable_upper_bound[i])
            ? rescale_info->scaled_problem->variable_upper_bound[i]
            : 0.0;
  ALLOC_AND_COPY(state->variable_upper_bound_finite_val, temp_host, var_bytes);
  free(temp_host);

  double sum_of_squares = 0.0;
  double max_val = 0.0;
  double val = 0.0;

  for (int i = 0; i < n_vars; ++i) {
    if (params->optimality_norm == NORM_TYPE_L_INF) {
      val = fabs(working_problem->objective_vector[i]);
      if (val > max_val)
        max_val = val;
    } else {
      sum_of_squares += working_problem->objective_vector[i] *
                        working_problem->objective_vector[i];
    }
  }

  if (params->optimality_norm == NORM_TYPE_L_INF) {
    state->objective_vector_norm = max_val;
  } else {
    state->objective_vector_norm = sqrt(sum_of_squares);
  }

  sum_of_squares = 0.0;
  max_val = 0.0;
  val = 0.0;

  for (int i = 0; i < n_cons; ++i) {
    double lower = working_problem->constraint_lower_bound[i];
    double upper = working_problem->constraint_upper_bound[i];

    if (params->optimality_norm == NORM_TYPE_L_INF) {
      if (isfinite(lower) && (lower != upper)) {
        val = fabs(lower);
        if (val > max_val)
          max_val = val;
      }
      if (isfinite(upper)) {
        val = fabs(upper);
        if (val > max_val)
          max_val = val;
      }
    } else {
      if (isfinite(lower) && (lower != upper)) {
        sum_of_squares += lower * lower;
      }
      if (isfinite(upper)) {
        sum_of_squares += upper * upper;
      }
    }
  }

  if (params->optimality_norm == NORM_TYPE_L_INF) {
    state->constraint_bound_norm = max_val;
  } else {
    state->constraint_bound_norm = sqrt(sum_of_squares);
  }

  state->num_blocks_primal =
      (state->num_variables + THREADS_PER_BLOCK - 1) / THREADS_PER_BLOCK;
  state->num_blocks_dual =
      (state->num_constraints + THREADS_PER_BLOCK - 1) / THREADS_PER_BLOCK;
  state->num_blocks_primal_dual =
      (state->num_variables + state->num_constraints + THREADS_PER_BLOCK - 1) /
      THREADS_PER_BLOCK;

  state->best_primal_dual_residual_gap = INFINITY;
  state->last_trial_fixed_point_error = INFINITY;
  state->step_size = 0.0;
  state->is_this_major_iteration = false;

  size_t primal_spmv_buffer_size;
  size_t dual_spmv_buffer_size;

  CUSPARSE_CHECK(cusparseCreateCsr(
      &state->matA, state->num_constraints, state->num_variables,
      state->constraint_matrix->num_nonzeros, state->constraint_matrix->row_ptr,
      state->constraint_matrix->col_ind, state->constraint_matrix->val,
      CUSPARSE_INDEX_32I, CUSPARSE_INDEX_32I, CUSPARSE_INDEX_BASE_ZERO,
      CUDA_R_64F));

  CUDA_CHECK(cudaGetLastError());

  CUSPARSE_CHECK(cusparseCreateCsr(
      &state->matAt, state->num_variables, state->num_constraints,
      state->constraint_matrix_t->num_nonzeros,
      state->constraint_matrix_t->row_ptr, state->constraint_matrix_t->col_ind,
      state->constraint_matrix_t->val, CUSPARSE_INDEX_32I, CUSPARSE_INDEX_32I,
      CUSPARSE_INDEX_BASE_ZERO, CUDA_R_64F));
  CUDA_CHECK(cudaGetLastError());

  CUSPARSE_CHECK(cusparseCreateDnVec(&state->vec_primal_sol,
                                     state->num_variables,
                                     state->pdhg_primal_solution, CUDA_R_64F));
  CUSPARSE_CHECK(cusparseCreateDnVec(&state->vec_dual_sol,
                                     state->num_constraints,
                                     state->pdhg_dual_solution, CUDA_R_64F));
  CUSPARSE_CHECK(cusparseCreateDnVec(&state->vec_primal_prod,
                                     state->num_constraints,
                                     state->primal_product, CUDA_R_64F));
  CUSPARSE_CHECK(cusparseCreateDnVec(&state->vec_dual_prod,
                                     state->num_variables, state->dual_product,
                                     CUDA_R_64F));
  CUSPARSE_CHECK(cusparseSpMV_bufferSize(
      state->sparse_handle, CUSPARSE_OPERATION_NON_TRANSPOSE, &HOST_ONE,
      state->matA, state->vec_primal_sol, &HOST_ZERO, state->vec_primal_prod,
      CUDA_R_64F, CUSPARSE_SPMV_CSR_ALG2, &primal_spmv_buffer_size));

  CUSPARSE_CHECK(cusparseSpMV_bufferSize(
      state->sparse_handle, CUSPARSE_OPERATION_NON_TRANSPOSE, &HOST_ONE,
      state->matAt, state->vec_dual_sol, &HOST_ZERO, state->vec_dual_prod,
      CUDA_R_64F, CUSPARSE_SPMV_CSR_ALG2, &dual_spmv_buffer_size));
  CUDA_CHECK(cudaMalloc(&state->primal_spmv_buffer, primal_spmv_buffer_size));
  CUSPARSE_CHECK(cusparseSpMV_preprocess(
      state->sparse_handle, CUSPARSE_OPERATION_NON_TRANSPOSE, &HOST_ONE,
      state->matA, state->vec_primal_sol, &HOST_ZERO, state->vec_primal_prod,
      CUDA_R_64F, CUSPARSE_SPMV_CSR_ALG2, state->primal_spmv_buffer));

  CUDA_CHECK(cudaMalloc(&state->dual_spmv_buffer, dual_spmv_buffer_size));
  CUSPARSE_CHECK(cusparseSpMV_preprocess(
      state->sparse_handle, CUSPARSE_OPERATION_NON_TRANSPOSE, &HOST_ONE,
      state->matAt, state->vec_dual_sol, &HOST_ZERO, state->vec_dual_prod,
      CUDA_R_64F, CUSPARSE_SPMV_CSR_ALG2, state->dual_spmv_buffer));

  initialize_quadratic_obj_term(state, rescale_info->scaled_problem);
  initialize_quadratic_term_information(state, params);
  initialize_inner_solver(state);

  CUDA_CHECK(
      cudaMalloc(&state->ones_primal_d, state->num_variables * sizeof(double)));
  CUDA_CHECK(
      cudaMalloc(&state->ones_dual_d, state->num_constraints * sizeof(double)));

  double *ones_primal_h =
      (double *)safe_malloc(state->num_variables * sizeof(double));
  for (int i = 0; i < state->num_variables; ++i)
    ones_primal_h[i] = 1.0;
  CUDA_CHECK(cudaMemcpy(state->ones_primal_d, ones_primal_h,
                        state->num_variables * sizeof(double),
                        cudaMemcpyHostToDevice));
  free(ones_primal_h);

  double *ones_dual_h =
      (double *)safe_malloc(state->num_constraints * sizeof(double));
  for (int i = 0; i < state->num_constraints; ++i)
    ones_dual_h[i] = 1.0;
  CUDA_CHECK(cudaMemcpy(state->ones_dual_d, ones_dual_h,
                        state->num_constraints * sizeof(double),
                        cudaMemcpyHostToDevice));
  decide_problem_type(state);
  free(ones_dual_h);
  if (params->verbose >= 2) {
    printf(
        "---------------------------------------------------------------------"
        "------------------\n");
    printf("Problem Type: %s\n", problem_type_to_string(state->problem_type));
    printf("Quadratic Objective Matrix Type: %s\n",
           quad_obj_type_to_string(
               state->quadratic_objective_term->quad_obj_type));
    if (state->quadratic_objective_term->quad_obj_type != PDHCG_NON_Q) {
      printf("L2 Norm of Quadratic Objective Matrix: %.3e\n",
             state->quadratic_objective_term->norm);
      printf("Non-Convexity of Quadratic Objective Matrix: %.3e\n",
             state->quadratic_objective_term->nonconvexity);
    }
    printf(
        "---------------------------------------------------------------------"
        "------------------\n");
    printf("%s | %s | %s | %s \n", "   runtime    ", "    objective     ",
           "  absolute residuals   ", "  relative residuals   ");
    printf("%s %s | %s %s | %s %s %s | %s %s %s \n", "  iter", "  time ",
           " pr obj ", "  du obj ", " pr res", " du res", "  gap  ", " pr res",
           " du res", "  gap  ");
    printf(
        "---------------------------------------------------------------------"
        "------------------\n");
  }

  return state;
}
void pdhg_solver_state_free(pdhg_solver_state_t *state) {
  if (state == NULL) {
    return;
  }

  if (state->variable_lower_bound)
    CUDA_CHECK(cudaFree(state->variable_lower_bound));
  if (state->variable_upper_bound)
    CUDA_CHECK(cudaFree(state->variable_upper_bound));
  if (state->objective_vector)
    CUDA_CHECK(cudaFree(state->objective_vector));
  if (state->constraint_matrix->row_ptr)
    CUDA_CHECK(cudaFree(state->constraint_matrix->row_ptr));
  if (state->constraint_matrix->col_ind)
    CUDA_CHECK(cudaFree(state->constraint_matrix->col_ind));
  if (state->constraint_matrix->val)
    CUDA_CHECK(cudaFree(state->constraint_matrix->val));
  if (state->constraint_matrix_t->row_ptr)
    CUDA_CHECK(cudaFree(state->constraint_matrix_t->row_ptr));
  if (state->constraint_matrix_t->col_ind)
    CUDA_CHECK(cudaFree(state->constraint_matrix_t->col_ind));
  if (state->constraint_matrix_t->val)
    CUDA_CHECK(cudaFree(state->constraint_matrix_t->val));
  if (state->constraint_lower_bound)
    CUDA_CHECK(cudaFree(state->constraint_lower_bound));
  if (state->constraint_upper_bound)
    CUDA_CHECK(cudaFree(state->constraint_upper_bound));
  if (state->constraint_lower_bound_finite_val)
    CUDA_CHECK(cudaFree(state->constraint_lower_bound_finite_val));
  if (state->constraint_upper_bound_finite_val)
    CUDA_CHECK(cudaFree(state->constraint_upper_bound_finite_val));
  if (state->variable_lower_bound_finite_val)
    CUDA_CHECK(cudaFree(state->variable_lower_bound_finite_val));
  if (state->variable_upper_bound_finite_val)
    CUDA_CHECK(cudaFree(state->variable_upper_bound_finite_val));
  if (state->initial_primal_solution)
    CUDA_CHECK(cudaFree(state->initial_primal_solution));
  if (state->current_primal_solution)
    CUDA_CHECK(cudaFree(state->current_primal_solution));
  if (state->pdhg_primal_solution)
    CUDA_CHECK(cudaFree(state->pdhg_primal_solution));
  if (state->reflected_primal_solution)
    CUDA_CHECK(cudaFree(state->reflected_primal_solution));
  if (state->dual_product)
    CUDA_CHECK(cudaFree(state->dual_product));
  if (state->initial_dual_solution)
    CUDA_CHECK(cudaFree(state->initial_dual_solution));
  if (state->current_dual_solution)
    CUDA_CHECK(cudaFree(state->current_dual_solution));
  if (state->pdhg_dual_solution)
    CUDA_CHECK(cudaFree(state->pdhg_dual_solution));
  if (state->reflected_dual_solution)
    CUDA_CHECK(cudaFree(state->reflected_dual_solution));
  if (state->primal_product)
    CUDA_CHECK(cudaFree(state->primal_product));
  if (state->constraint_rescaling)
    CUDA_CHECK(cudaFree(state->constraint_rescaling));
  if (state->variable_rescaling)
    CUDA_CHECK(cudaFree(state->variable_rescaling));
  if (state->primal_slack)
    CUDA_CHECK(cudaFree(state->primal_slack));
  if (state->dual_slack)
    CUDA_CHECK(cudaFree(state->dual_slack));
  if (state->primal_residual)
    CUDA_CHECK(cudaFree(state->primal_residual));
  if (state->dual_residual)
    CUDA_CHECK(cudaFree(state->dual_residual));
  if (state->delta_primal_solution)
    CUDA_CHECK(cudaFree(state->delta_primal_solution));
  if (state->delta_dual_solution)
    CUDA_CHECK(cudaFree(state->delta_dual_solution));
  if (state->ones_primal_d)
    CUDA_CHECK(cudaFree(state->ones_primal_d));
  if (state->ones_dual_d)
    CUDA_CHECK(cudaFree(state->ones_dual_d));

  free(state);
}

void rescale_info_free(rescale_info_t *info) {
  if (info == NULL) {
    return;
  }

  qp_problem_free(info->scaled_problem);
  free(info->con_rescale);
  free(info->var_rescale);

  free(info);
}
