/* nag_zggqrf (f08zsc) Example Program.
 *
 * Copyright 2017 Numerical Algorithms Group.
 *
 * Mark 26.1, 2017.
 */

#include <stdio.h>
#include <nag.h>
#include <nag_stdlib.h>
#include <naga02.h>
#include <nagf07.h>
#include <nagf08.h>
#include <nagf16.h>

int main(void)
{
  /* Scalars */
  Complex alpha, beta;
  Complex zero = { 0.0, 0.0 };
  double rnorm;
  Integer i, j, m, n, nm, p, pda, pdb, pdd, pnm, zrow;
  Integer exit_status = 0;

  /* Arrays */
  Complex *a = 0, *b = 0, *d = 0, *taua = 0, *taub = 0, *y = 0;

  /* Nag Types */
  NagError fail;
  Nag_OrderType order;

#ifdef NAG_COLUMN_MAJOR
#define A(I, J) a[(J-1)*pda + I - 1]
#define B(I, J) b[(J-1)*pdb + I - 1]
  order = Nag_ColMajor;
#else
#define A(I, J) a[(I-1)*pda + J - 1]
#define B(I, J) b[(I-1)*pdb + J - 1]
  order = Nag_RowMajor;
#endif

  INIT_FAIL(fail);

  printf("nag_zggqrf (f08zsc) Example Program Results\n\n");

  /* Skip heading in data file */
  scanf("%*[^\n]");
  scanf("%" NAG_IFMT "%" NAG_IFMT "%" NAG_IFMT "%*[^\n]", &n, &m, &p);
  if (n < 0 || m < 0 || p < 0) {
    printf("Invalid n, m or p\n");
    exit_status = 1;
    goto END;
  }

#ifdef NAG_COLUMN_MAJOR
  pda = n;
  pdb = n;
  pdd = n;
#else
  pda = m;
  pdb = p;
  pdd = 1;
#endif

  /* Allocate memory */
  if (!(a = NAG_ALLOC(n * m, Complex)) ||
      !(b = NAG_ALLOC(n * p, Complex)) ||
      !(d = NAG_ALLOC(n, Complex)) ||
      !(taua = NAG_ALLOC(MIN(n, m), Complex)) ||
      !(taub = NAG_ALLOC(MIN(n, p), Complex)) || !(y = NAG_ALLOC(p, Complex)))
  {
    printf("Allocation failure\n");
    exit_status = -1;
    goto END;
  }

  /* Read A, B and d from data file */
  for (i = 1; i <= n; ++i)
    for (j = 1; j <= m; ++j)
      scanf(" ( %lf , %lf )", &A(i, j).re, &A(i, j).im);
  scanf("%*[^\n]");
  for (i = 1; i <= n; ++i)
    for (j = 1; j <= p; ++j)
      scanf(" ( %lf , %lf )", &B(i, j).re, &B(i, j).im);
  scanf("%*[^\n]");
  for (i = 0; i < n; ++i)
    scanf(" ( %lf , %lf )", &d[i].re, &d[i].im);
  scanf("%*[^\n]");

  /* Compute the generalized QR factorization of (A,B) as
   * A = Q*(R),   B = Q*(T11 T12)*Z
   *       (0)          ( 0  T22)
   * using  nag_dggqrf (f08zec).
   */
  nag_zggqrf(order, n, m, p, a, pda, taua, b, pdb, taub, &fail);
  if (fail.code != NE_NOERROR) {
    printf("Error from nag_zggqrf (f08zsc).\n%s\n", fail.message);
    exit_status = 1;
    goto END;
  }

  /* Solve weighted least squares problem for case n > m */
  if (n <= m)
    goto END;

  nm = n - m;
  pnm = p - nm;
  /* Multiply Q^H through d = Ax + By to get 
   *     (c1) = Q^H * d = (R) * x + (T11 T12) * Z * (y1)
   *     (c2)             (0)       ( 0  T22)       (y2)
   * Compute C using nag_zunmqr (f08auc).
   */
  nag_zunmqr(order, Nag_LeftSide, Nag_ConjTrans, n, 1, m, a, pda, taua, d,
             pdd, &fail);
  if (fail.code != NE_NOERROR) {
    printf("Error from nag_zunmqr (f08auc).\n%s\n", fail.message);
    exit_status = 1;
    goto END;
  }

  /* Let Z*(y1) = (w1) and solving for w2 we have to solve the triangular sytem
   *       (y2) = (w2)
   *                     T22 * w2 = c2
   * This is done by putting c2 in y2 and backsolving to get w2 in y2.
   *
   * Copy c2 (at d[m]) into y2 using nag_zge_copy (f16tfc).
   */
  nag_zge_copy(Nag_ColMajor, Nag_NoTrans, nm, 1, &d[m], n - m, &y[pnm], nm,
               &fail);
  if (fail.code != NE_NOERROR) {
    printf("Error from nag_zge_copy (f16tfc).\n%s\n", fail.message);
    exit_status = 1;
    goto END;
  }
  /* Solve T22*w2 = c2 using nag_ztrtrs (f07tsc).
   * T22 is stored in a submatrix of matrix B of dimension n-m by n-m
   * with first element at B(m+1,p-(n-m)+1). y2 is stored from y[p-(n-m)].
   */
  nag_ztrtrs(order, Nag_Upper, Nag_NoTrans, Nag_NonUnitDiag, nm, 1,
             &B(m + 1, pnm + 1), pdb, &y[pnm], nm, &fail);
  if (fail.code != NE_NOERROR) {
    printf("Error from nag_ztrtrs (f07tsc).\n%s\n", fail.message);
    exit_status = 1;
    goto END;
  }

  /* set w1 = 0 for minimum norm y. */
  nag_zload(pnm, zero, y, 1, &fail);
  if (fail.code != NE_NOERROR) {
    printf("Error from nag_zload (f16hbc).\n%s\n", fail.message);
    exit_status = 1;
    goto END;
  }

  /* Compute estimate of the square root of the residual sum of squares 
   * norm(y) = norm(w2) with y1 = 0 using  nag_dge_norm (f16uac).
   */
  nag_zge_norm(Nag_ColMajor, Nag_FrobeniusNorm, nm, 1, &y[pnm], nm, &rnorm,
               &fail);
  if (fail.code != NE_NOERROR) {
    printf("Error from nag_zge_norm (f16uac).\n%s\n", fail.message);
    exit_status = 1;
    goto END;
  }

  /* The top half of the system remains:
   *     (c1) = Q^H * d = (R) * x + (T11 T12) * ( 0)
   *                                            (w2)
   * =>      c1 = R * x + T12 * w2
   * =>   R * x = c1    - T12 * w2;
   * 
   * first form d = c1 - T12*w2 where c1 is stored in d
   * using nag_zgemv (f16sac).
   */
  alpha = nag_complex(-1.0, 0.0);
  beta = nag_complex(1.0, 0.0);
  nag_zgemv(order, Nag_NoTrans, m, nm, alpha, &B(1, pnm + 1), pdb, &y[pnm], 1,
            beta, d, 1, &fail);
  if (fail.code != NE_NOERROR) {
    printf("Error from nag_zgemv (f16sac).\n%s\n", fail.message);
    exit_status = 1;
    goto END;
  }

  /* Next, solve R * x = d for x (in d) where R is stored in leading submatrix
   * of A in a. This gives the least squares solution x in d.
   * Using nag_dtrtrs (f07tec).
   */
  nag_ztrtrs(order, Nag_Upper, Nag_NoTrans, Nag_NonUnitDiag, m, 1, a, pda, d,
             pdd, &fail);
  if (fail.code != NE_NOERROR) {
    printf("Error from nag_ztrtrs (f07tsc).\n%s\n", fail.message);
    exit_status = 1;
    goto END;
  }

  /* Compute the minimum norm residual vector y = (Z^T)*w 
   * using nag_dzunmrq (f08cxc).
   */
  zrow = MAX(1, n - p + 1);
  nag_zunmrq(order, Nag_LeftSide, Nag_ConjTrans, p, 1, MIN(n, p), &B(zrow, 1),
             pdb, taub, y, pdd, &fail);
  if (fail.code != NE_NOERROR) {
    printf("Error from nag_zunmrq (f08cxc).\n%s\n", fail.message);
    exit_status = 1;
    goto END;
  }

  /* Print least squares solution x */
  printf("Generalized least squares solution\n");
  for (i = 0; i < m; ++i)
    printf(" (%9.4f, %9.4f)%s", d[i].re, d[i].im, i % 3 == 2 ? "\n" : "");

  /* Print residual vector y */
  printf("\n");
  printf("\nResidual vector\n");
  for (i = 0; i < p; ++i)
    printf(" (%9.2e, %9.2e)%s", y[i].re, y[i].im, i % 3 == 2 ? "\n" : "");

  /* Print estimate of the square root of the residual sum of  squares. */
  printf("\n\nSquare root of the residual sum of squares\n");
  printf("%11.2e\n", rnorm);

END:
  NAG_FREE(a);
  NAG_FREE(b);
  NAG_FREE(d);
  NAG_FREE(taua);
  NAG_FREE(taub);
  NAG_FREE(y);

  return exit_status;
}