module kkt;

nothrow @nogc extern(C):

import glob_opts;
import cs;
import types; // CSC matrix type
import constants; // for OSQP_NULL

version(EMBEDDED){}
else {
csc* form_KKT(const csc  *P,
              const  csc *A,
              c_int       format,
              c_float     param1,
              c_float    *param2,
              c_int      *PtoKKT,
              c_int      *AtoKKT,
              c_int     **Pdiag_idx,
              c_int      *Pdiag_n,
              c_int      *param2toKKT) {
  c_int  nKKT;            // Size, number of nonzeros and max number of nonzeros in KKT matrix
  c_int  nnzKKTmax;       // Size, number of nonzeros and max number of nonzeros in KKT matrix
  csc   *KKT_trip;        // KKT matrix in triplet format and CSC format
  csc   *KKT;             // KKT matrix in triplet format and CSC format
  c_int  ptr, i, j;       // Counters for elements (i,j) and index pointer
  c_int  zKKT = 0;        // Counter for total number of elements in P and in
                          // KKT
  c_int *KKT_TtoC;        // Pointer to vector mapping from KKT in triplet form
                          // to CSC

  // Get matrix dimensions
  nKKT = P.m + A.m;

  // Get maximum number of nonzero elements (only upper triangular part)
  nnzKKTmax = P.p[P.n] + // Number of elements in P
              P.m +       // Number of elements in param1 * I
              A.p[A.n] + // Number of nonzeros in A
              A.m;        // Number of elements in - diag(param2)

  // Preallocate KKT matrix in triplet format
  KKT_trip = csc_spalloc(nKKT, nKKT, nnzKKTmax, 1, 1);

  if (!KKT_trip) return OSQP_NULL;  // Failed to preallocate matrix

  // Allocate vector of indices on the diagonal. Worst case it has m elements
  if (Pdiag_idx != OSQP_NULL) {
    (*Pdiag_idx) = cast(c_int*)c_malloc(P.m * c_int.sizeof);
    *Pdiag_n     = 0; // Set 0 diagonal elements to start
  }

  // Allocate Triplet matrices
  // P + param1 I
  for (j = 0; j < P.n; j++) { // cycle over columns
    // No elements in column j => add diagonal element param1
    if (P.p[j] == P.p[j + 1]) {
      KKT_trip.i[zKKT] = j;
      KKT_trip.p[zKKT] = j;
      KKT_trip.x[zKKT] = param1;
      zKKT++;
    }

    for (ptr = P.p[j]; ptr < P.p[j + 1]; ptr++) { // cycle over rows
      // Get current row
      i = P.i[ptr];

      // Add element of P
      KKT_trip.i[zKKT] = i;
      KKT_trip.p[zKKT] = j;
      KKT_trip.x[zKKT] = P.x[ptr];

      if (PtoKKT != OSQP_NULL) PtoKKT[ptr] = zKKT;  // Update index from P to
                                                    // KKTtrip

      if (i == j) {                                 // P has a diagonal element,
                                                    // add param1
        KKT_trip.x[zKKT] += param1;

        // If index vector pointer supplied . Store the index
        if (Pdiag_idx != OSQP_NULL) {
          (*Pdiag_idx)[*Pdiag_n] = ptr;
          (*Pdiag_n)++;
        }
      }
      zKKT++;

      // Add diagonal param1 in case
      if ((i < j) &&                  // Diagonal element not reached
          (ptr + 1 == P.p[j + 1])) { // last element of column j
        // Add diagonal element param1
        KKT_trip.i[zKKT] = j;
        KKT_trip.p[zKKT] = j;
        KKT_trip.x[zKKT] = param1;
        zKKT++;
      }
    }
  }

  if (Pdiag_idx != OSQP_NULL) {
    // Realloc Pdiag_idx so that it contains exactly *Pdiag_n diagonal elements
    (*Pdiag_idx) = cast(c_int*)c_realloc((*Pdiag_idx), (*Pdiag_n) * c_int.sizeof);
  }


  // A' at top right
  for (j = 0; j < A.n; j++) {                      // Cycle over columns of A
    for (ptr = A.p[j]; ptr < A.p[j + 1]; ptr++) {
      KKT_trip.p[zKKT] = P.m + A.i[ptr];         // Assign column index from
                                                    // row index of A
      KKT_trip.i[zKKT] = j;                        // Assign row index from
                                                    // column index of A
      KKT_trip.x[zKKT] = A.x[ptr];                // Assign A value element

      if (AtoKKT != OSQP_NULL) AtoKKT[ptr] = zKKT;  // Update index from A to
                                                    // KKTtrip
      zKKT++;
    }
  }

  // - diag(param2) at bottom right
  for (j = 0; j < A.m; j++) {
    KKT_trip.i[zKKT] = j + P.n;
    KKT_trip.p[zKKT] = j + P.n;
    KKT_trip.x[zKKT] = -param2[j];

    if (param2toKKT != OSQP_NULL) param2toKKT[j] = zKKT;  // Update index from
                                                          // param2 to KKTtrip
    zKKT++;
  }

  // Allocate number of nonzeros
  KKT_trip.nz = zKKT;

  // Convert triplet matrix to csc format
  if (!PtoKKT && !AtoKKT && !param2toKKT) {
    // If no index vectors passed, do not store KKT mapping from Trip to CSC/CSR
    if (format == 0) KKT = triplet_to_csc(KKT_trip, OSQP_NULL);
    else KKT = triplet_to_csr(KKT_trip, OSQP_NULL);
  }
  else {
    // Allocate vector of indices from triplet to csc
    KKT_TtoC = cast(c_int*)c_malloc((zKKT) * c_int.sizeof);

    if (!KKT_TtoC) {
      // Error in allocating KKT_TtoC vector
      csc_spfree(KKT_trip);
      c_free(*Pdiag_idx);
      return OSQP_NULL;
    }

    // Store KKT mapping from Trip to CSC/CSR
    if (format == 0)
      KKT = triplet_to_csc(KKT_trip, KKT_TtoC);
    else
      KKT = triplet_to_csr(KKT_trip, KKT_TtoC);

    // Update vectors of indices from P, A, param2 to KKT (now in CSC format)
    if (PtoKKT != OSQP_NULL) {
      for (i = 0; i < P.p[P.n]; i++) {
        PtoKKT[i] = KKT_TtoC[PtoKKT[i]];
      }
    }

    if (AtoKKT != OSQP_NULL) {
      for (i = 0; i < A.p[A.n]; i++) {
        AtoKKT[i] = KKT_TtoC[AtoKKT[i]];
      }
    }

    if (param2toKKT != OSQP_NULL) {
      for (i = 0; i < A.m; i++) {
        param2toKKT[i] = KKT_TtoC[param2toKKT[i]];
      }
    }

    // Free mapping
    c_free(KKT_TtoC);
  }

  // Clean matrix in triplet format and return result
  csc_spfree(KKT_trip);

  return KKT;
}

} /* ifndef EMBEDDED */


version (EMBEDDED_1){}
else {

void update_KKT_P(csc          *KKT,
                  const csc    *P,
                  const c_int  *PtoKKT,
                  const c_float param1,
                  const c_int  *Pdiag_idx,
                  const c_int   Pdiag_n) {
  c_int i, j; // Iterations

  // Update elements of KKT using P
  for (i = 0; i < P.p[P.n]; i++) {
    KKT.x[PtoKKT[i]] = P.x[i];
  }

  // Update diagonal elements of KKT by adding sigma
  for (i = 0; i < Pdiag_n; i++) {
    j                  = Pdiag_idx[i]; // Extract index of the element on the
                                       // diagonal
    KKT.x[PtoKKT[j]] += param1;
  }
}

void update_KKT_A(csc *KKT, const csc *A, const c_int *AtoKKT) {
  c_int i; // Iterations

  // Update elements of KKT using A
  for (i = 0; i < A.p[A.n]; i++) {
    KKT.x[AtoKKT[i]] = A.x[i];
  }
}

void update_KKT_param2(csc *KKT, const c_float *param2,
                       const c_int *param2toKKT, const c_int m) {
  c_int i; // Iterations

  // Update elements of KKT using param2
  for (i = 0; i < m; i++) {
    KKT.x[param2toKKT[i]] = -param2[i];
  }
}

} // EMBEDDED != 1