/**
 * @file structs.h
 * @author Govind M. Chari <govindchari1@gmail.com>
 *
 * @section LICENSE
 *
 * Copyright (c) 2024, Govind M. Chari
 * This source code is licensed under the BSD 3-Clause License
 *
 * @section DESCRIPTION
 *
 * Defines all structs used by QOCO.
 */

#ifndef QOCO_STRUCTS_H
#define QOCO_STRUCTS_H

#include "definitions.h"
#include "qoco_linalg.h"
#include "timer.h"

/**
 * @brief SOCP problem data.
 *
 */
typedef struct {
  /** Quadratic cost term. */
  QOCOMatrix* P; //

  /** Linear cost term. */
  QOCOVectorf* c;

  /** Affine equality constraint matrix. */
  QOCOMatrix* A;

  /** Transpose of A (used in Ruiz for fast row norm calculations of A). */
  QOCOMatrix* At;

  /** Affine equality constraint offset. */
  QOCOVectorf* b;

  /** Conic constraint matrix. */
  QOCOMatrix* G;

  /** Transpose of G (used in Ruiz for fast row norm calculations of G). */
  QOCOMatrix* Gt;

  /** Mapping from A to At. */
  QOCOInt* AtoAt;

  /** Mapping from G to Gt. */
  QOCOInt* GtoGt;

  /** Conic constraint offset. */
  QOCOVectorf* h;

  /** Dimension of non-negative orthant in cone C. */
  QOCOInt l;

  /** Number of second-order cones in C */
  QOCOInt nsoc;

  /** Dimension of each second-order cone (length of nsoc)*/
  QOCOVectori* q;

  /** Number of primal variables. */
  QOCOInt n;

  /** Number of conic constraints. */
  QOCOInt m;

  /** Number of affine equality constraints. */
  QOCOInt p;

  /** Indices of P->x that were added due to regularization. */
  QOCOInt* Pnzadded_idx;

  /** Number of elements of P->x that were added due to regularization. */
  QOCOInt Pnum_nzadded;

  /** Scaling statistics: minimum absolute value in P (prior to regularization)
   * and c. */
  QOCOFloat obj_range_min;

  /** Scaling statistics: maximum absolute value in P (prior to regularization)
   * and c. */
  QOCOFloat obj_range_max;

  /** Scaling statistics: minimum absolute value in A and G. */
  QOCOFloat constraint_range_min;

  /** Scaling statistics: maximum absolute value in A and G. */
  QOCOFloat constraint_range_max;

  /** Scaling statistics: minimum absolute value in b and h. */
  QOCOFloat rhs_range_min;

  /** Scaling statistics: maximum absolute value in b and h. */
  QOCOFloat rhs_range_max;

} QOCOProblemData;

/**
 * @brief QOCO solver settings
 *
 */
typedef struct {
  /** Maximum number of IPM iterations. */
  QOCOInt max_iters;

  /** Number of bisection iterations for linesearch. */
  QOCOInt bisect_iters;

  /** Number of Ruiz equilibration iterations. */
  QOCOInt ruiz_iters;

  /** Number of iterative refinement iterations performed. */
  QOCOInt iter_ref_iters;

  /** Static regularization parameter for KKT system. */
  QOCOFloat kkt_static_reg;

  /** Dynamic regularization parameter for KKT system. */
  QOCOFloat kkt_dynamic_reg;

  /** Absolute tolerance. */
  QOCOFloat abstol;

  /** Relative tolerance. */
  QOCOFloat reltol;

  /** Low tolerance stopping criteria. */
  QOCOFloat abstol_inacc;

  /** Low tolerance stopping criteria. */
  QOCOFloat reltol_inacc;

  /** 0 for quiet anything else for verbose. */
  unsigned char verbose;
} QOCOSettings;

/**
 * @brief Contains all data related to problem scaling and equilibration.
 *
 */
typedef struct {

  /** Diagonal of scaling matrix. */
  QOCOVectorf* delta;

  /** Diagonal of scaling matrix. */
  QOCOVectorf* Druiz;

  /** Diagonal of scaling matrix. */
  QOCOVectorf* Eruiz;

  /** Diagonal of scaling matrix. */
  QOCOVectorf* Fruiz;

  /** Inverse of Druiz. */
  QOCOVectorf* Dinvruiz;

  /** Inverse of Eruiz. */
  QOCOVectorf* Einvruiz;

  /** Inverse of Fruiz. */
  QOCOVectorf* Finvruiz;

  /** Cost scaling factor. */
  QOCOFloat k;

  /** Inverse of cost scaling factor. */
  QOCOFloat kinv;

} QOCOScaling;

/**
 * @brief QOCO Workspace
 */
typedef struct {
  /** Contains SOCP problem data. */
  QOCOProblemData* data;

  /** Solve timer. */
  QOCOTimer solve_timer;

  /** Contains all data related problem scaling. */
  QOCOScaling* scaling;

  /** Iterate of primal variables. */
  QOCOVectorf* x;

  /** Iterate of slack variables associated with conic constraint. */
  QOCOVectorf* s;

  /** Iterate of dual variables associated with affine equality constraint. */
  QOCOVectorf* y;

  /** Iterate of dual variables associated with conic constraint. */
  QOCOVectorf* z;

  /** Gap (s'*z / m) */
  QOCOFloat mu;

  /** Newton Step-size */
  QOCOFloat a;

  /** Centering parameter */
  QOCOFloat sigma;

  /** Number of nonzeros in upper triangular part of Nesterov-Todd Scaling. */
  QOCOInt Wnnz;

  /** Number of nonzeros in full Nesterov-Todd Scaling. */
  QOCOInt Wnnzfull;

  /** Upper triangular part of Nesterov-Todd Scaling */
  QOCOVectorf* W;

  /** Full Nesterov-Todd Scaling */
  QOCOVectorf* Wfull;

  /** Upper triangular part of inverse of Nesterov-Todd Scaling */
  QOCOVectorf* Winv;

  /** Full inverse of Nesterov-Todd Scaling */
  QOCOVectorf* Winvfull;

  /** Nesterov-Todd Scaling squared */
  QOCOVectorf* WtW;

  /** Scaled variables. */
  QOCOVectorf* lambda;

  /** Temporary array needed in Nesterov-Todd scaling calculations. Length of
   * max(q). */
  QOCOVectorf* sbar;

  /** Temporary array needed in Nesterov-Todd scaling calculations. Length of
   * max(q). */
  QOCOVectorf* zbar;

  /** Temporary variable of length n. */
  QOCOVectorf* xbuff;

  /** Temporary variable of length p. */
  QOCOVectorf* ybuff;

  /** Temporary variable of length m. */
  QOCOVectorf* ubuff1;

  /** Temporary variable of length m. */
  QOCOVectorf* ubuff2;

  /** Temporary variable of length m. */
  QOCOVectorf* ubuff3;

  /** Search direction for slack variables. Length of m. */
  QOCOVectorf* Ds;

  /** RHS of KKT system. */
  QOCOVectorf* rhs;

  /** Solution of KKT system. */
  QOCOVectorf* xyz;

  /** Buffer of size n + m + p. */
  QOCOVectorf* xyzbuff1;

  /** Buffer of size n + m + p. */
  QOCOVectorf* xyzbuff2;

  /** Residual of KKT condition. */
  QOCOVectorf* kktres;

} QOCOWorkspace;

typedef struct {
  /**Primal solution. */
  QOCOFloat* x;

  /**Slack variable for conic constraints. */
  QOCOFloat* s;

  /**Dual variables for affine equality constraints. */
  QOCOFloat* y;

  /**Dual variables for conic constraints. */
  QOCOFloat* z;

  /**Number of iterations. */
  QOCOInt iters;

  /**Setup time. */
  QOCOFloat setup_time_sec;

  /**Solve time. */
  QOCOFloat solve_time_sec;

  /**Optimal objective value. */
  QOCOFloat obj;

  /** Primal residual. */
  QOCOFloat pres;

  /** Dual residual. */
  QOCOFloat dres;

  /** Duality gap. */
  QOCOFloat gap;

  /**Solve status. */
  QOCOInt status;

} QOCOSolution;

// Linear system data structs for backends to define.
typedef struct LinSysData LinSysData;

typedef struct {
  const char* (*linsys_name)();
  LinSysData* (*linsys_setup)(QOCOProblemData* data, QOCOSettings* settings,
                              QOCOInt Wnnz);
  void (*linsys_update_nt)(LinSysData* linsys_data, QOCOVectorf* WtW_vec,
                           QOCOFloat kkt_static_reg, QOCOInt m);
  void (*linsys_update_data)(LinSysData* linsys_data, QOCOProblemData* data);
  void (*linsys_factor)(LinSysData* linsys_data, QOCOInt n,
                        QOCOFloat kkt_dynamic_reg);
  void (*linsys_solve)(LinSysData* linsys_data, QOCOWorkspace* work,
                       QOCOVectorf* b_vec, QOCOVectorf* x_vec,
                       QOCOInt iter_ref_iters);
  void (*linsys_cleanup)(LinSysData* linsys_data);
} LinSysBackend;

/**
 * @brief QOCO Solver struct. Contains all information about the state of the
 * solver.
 *
 */
typedef struct {
  /** Solver settings. */
  QOCOSettings* settings;

  /** Solver workspace. */
  QOCOWorkspace* work;

  /** Linear system backend. */
  LinSysBackend* linsys;

  /** Linear system backend. */
  LinSysData* linsys_data;

  /** Solution struct. */
  QOCOSolution* sol;

} QOCOSolver;

#endif /* #ifndef QOCO_STRUCTS_H */