/****************************************************************************
 *
 * BlkTriAux.c
 *
 * Abstract
 *   Define auxiliar functions for block Lanczos tridiagonalization.
 *   
 *
 * Student:  Guohong Liu
 *
 * ID:       0385117
 *
 ****************************************************************************/

#include <stdio.h>
#include <stdlib.h>
#include <malloc.h>
#include <memory.h>
#include <math.h>
#include "blklan.h"

/****************************************************************************
 *
 * Global variables
 *
 ****************************************************************************/

int n;                          /* Order of given matrix A.         */
int bs;                         /* Block size.                      */
double eps;       
double mideps;
int steps;

DoubleComplexMat *A = NULL;     /* Matrix A. n * n                  */
DoubleComplexMat *Q = NULL;     /* Unitary. one block = n * bs      */   
DoubleComplexMat *P = NULL;     /* Unitary. n * n                   */
DoubleComplexMat *M = NULL;     /* Diagonal blocks. bs * bs.        */                                     
DoubleComplexMat *B = NULL;     /* Off-Diagonal blocks. bs * bs     */                                  
DoubleComplexMat *WCur = NULL;  /* Detectors of current iter. bs*bs */
DoubleComplexMat *WOld = NULL;  /* Detectors of previous iter.bs*bs */
DoubleComplexMat *R = NULL;     /* Work matrix. n * bs              */
DoubleComplexMat *Work1 = NULL; /* Work matrix. bs * bs             */
DoubleComplexMat *Work2 = NULL; /* Work matrix. bs * bs             */

/* Memories for corresponding matrix or vector.                     */
DoubleComplex *Q_m = NULL;      
DoubleComplex *P_m = NULL;
DoubleComplex **M_m = NULL;                       
DoubleComplex **B_m = NULL;
DoubleComplex **WCur_m = NULL;   
DoubleComplex **WOld_m = NULL;

/* Memories for main diagonal and subdiagonal of tridiagonal        */ 
DoubleComplex *a_m = NULL;      
DoubleReal *b_m = NULL;

char separator[22] = "=====================\0";

/***************************************************************************
 *
 * Abstract: 
 *   Perform initialization for block tridiagonalization operation.
 *   Allocate necessary memories.
 *
 * Input:    
 *   na     Order of matrix A.
 *   A_p    Pointer to the memory of matrix A.
 *   blk    Block size.
 *   S_p    Pointer to the starting matrix S.
 *
 * Output:     
 *   a_p    Pointer to the memory of vector a.
 *   b_p    Pointer to the memory of vector b.
 *   Q_p    Pointer to the memory of matrix Q.
 *   P_p    Pointer to the memory of matrix P.
 *
 * Return:  
 *   0      Successful exit
 *   1      Memory allocation exception occurs
 *
 * NOTE:    
 *   There is a little difference in the definition of MIDEPS in
 *   C implementation from the paper and MATLAB code. Matlab code:
 *      SQRTEPS = sqrt(eps) 
 *      MIDEPS = sqrt(sqrt(SQRTEPS))^7
 *
 *   Later in the detection of loss of orthogonality,
 *      max(abs(W(:, colW, k, old))) >= SQRTEPS
 *   Supporse w is maximum element of matrix W, then 
 *   max(abs(W(:, colW, k, old))) > = SQRTEPS means
 *      sqrt(w.r^2 + w.i^2) >= sqrt(eps)
 *
 *   For the sake of performance, C implementation will be,
 *       MIDEPS = sqrt(sqrt(eps))^7
 *   Supporse w is one element of matrix W, then use the following code 
 *   to test if this element loses orthogonality,
 *       w.r^2 + w.i^2 >= eps
 *
 ***************************************************************************/

int blkTriInit (int na, DoubleComplex *A_p,
                 int blk, DoubleComplex *S_p, 
                 DoubleComplex *a_p, DoubleReal *b_p, 
                 DoubleComplex *Q_p, DoubleComplex *P_p)
{
    int i;
#ifdef INFO
    DoubleComplexMat S;
#endif

    n = na;
    bs = blk;
    steps = n / bs;
#ifdef MATLAB_COMP
    eps = 2.22044604925031e-016;
#else
    eps = getEps ();
#endif       
    mideps = pow (sqrt (sqrt (eps)), 7);

    /* Accept memory parameters, A, S, a, b, Q, from the caller.*/
    if (!(A = (DoubleComplexMat *)malloc (sizeof (DoubleComplexMat)))) {
        blkTriExcept ("blkTriInit", "Malloc for *A failed");
        return 1;
    }
    A->mat = A_p;
    A->m = n;
    A->n = n;

    Q_m = Q_p;
    P_m = P_p;
    if (!(Q = (DoubleComplexMat *)malloc (sizeof (DoubleComplexMat)))) {
        blkTriExcept ("blkTriInit", "Malloc for *Q failed");
        return 1;
    }
    /* Set the address of the first Q block.                    */ 
    Q->mat = Q_m;
    Q->m = n;
    Q->n = bs;
    /* The first block of Q is the starting matrix S.           */
    memcpy (Q->mat, S_p, n * bs * sizeof (DoubleComplex));

    /* Allocate other memory used in block Lanczos algorithm.   */
    if (!(M = (DoubleComplexMat *)malloc (sizeof (DoubleComplexMat)))) {
        blkTriExcept ("blkTriInit", "Malloc for *M failed");
        return 1;
    }
    if (!(M_m = (DoubleComplex **)
                malloc (steps * sizeof (DoubleComplex *)))) {
        blkTriExcept ("blkTriInit", "Malloc for **M_m failed");
        return 1;
    }
    for (i = 0; i < steps; i++) {
        if (!(M_m[i] = (DoubleComplex *)
                        malloc (bs * bs * sizeof (DoubleComplex)))) {
            blkTriExcept ("blkTriInit", "Malloc for M_m[i]->mat failed");
            return 1;
        }
    }

    if (!(B = (DoubleComplexMat *)malloc (sizeof (DoubleComplexMat)))) {
        blkTriExcept ("blkTriInit", "Malloc for *B failed");
        return 1;
    }
    if (!(B_m = (DoubleComplex **)
                malloc ((steps - 1) * sizeof (DoubleComplex *)))) {
        blkTriExcept ("blkTriInit", "Malloc for **B_m failed");
        return 1;
    }
    for (i = 0; i < steps - 1; i++) {
        if (!(B_m[i] = (DoubleComplex *)
                        malloc (bs * bs * sizeof (DoubleComplex)))) {
            blkTriExcept ("blkTriInit", "Malloc for B_m[i] failed");
            return 1;
        }
    }

    if (!(WCur = (DoubleComplexMat *)
                    malloc (sizeof (DoubleComplexMat)))) {
        blkTriExcept ("blkTriInit", "Malloc for *WCur failed");
        return 1;
    }
    if (!(WCur_m = (DoubleComplex **)
                    malloc (steps * sizeof (DoubleComplex *)))) {
        blkTriExcept ("blkTriInit", "Malloc for **WCur_m failed");
        return 1;
    }
    for (i = 0; i < steps; i++) {
        if (!(WCur_m[i] = (DoubleComplex *)
                    malloc (bs * bs * sizeof (DoubleComplex)))) {
            blkTriExcept ("blkTriInit", "Malloc for WCur_m[i] failed");
            return 1;
        }
        /* Set initial value of W1[i]   */
        zeros (WCur_m[i], bs * bs);
    }

    if (!(WOld = (DoubleComplexMat *)
                    malloc (sizeof (DoubleComplexMat)))) {
        blkTriExcept ("blkTriInit", "Malloc for *WOld failed");
        return 1;
    }
    if (!(WOld_m = (DoubleComplex **)
                    malloc (steps * sizeof (DoubleComplex *)))) {
        blkTriExcept ("blkTriInit", "Malloc for **WOld_m failed");
        return 1;
    }
    for (i = 0; i < steps; i++) {
        if (!(WOld_m[i] = (DoubleComplex *)
                            malloc (bs * bs * sizeof (DoubleComplex)))) {
            blkTriExcept ("blkTriInit", "Malloc for WOld_m[i] failed");
            return -1;
        }
        /* Set initial value of W1[i]   */
        zeros (WOld_m[i], bs * bs);
    }

    a_m = a_p;
    b_m = b_p;

    if (!(R = (DoubleComplexMat *)
                malloc (steps * sizeof (DoubleComplexMat)))) {
        blkTriExcept ("blkTriInit", "Malloc for *R failed");
        return 1;
    }
    if (!(R->mat = (DoubleComplex *)
                    malloc (n * bs * sizeof (DoubleComplex)))) {
        blkTriExcept ("blkTriInit", "Malloc for R->mat failed");
        return 1;
    }
    R->m = n;
    R->n = bs;

    if (!(Work1 = (DoubleComplexMat *)
                    malloc (steps * sizeof (DoubleComplexMat)))) {
        blkTriExcept ("blkTriInit", "Malloc for *Work1 failed");
        return 1;
    }
    if (!(Work1->mat = (DoubleComplex *)
                        malloc (n * bs * sizeof (DoubleComplex)))) {
        blkTriExcept ("blkTriInit", "Malloc for Work1->mat failed");
        return 1;
    }
    Work1->m = bs;
    Work1->n = bs;

    if (!(Work2 = (DoubleComplexMat *)
                    malloc (steps * sizeof (DoubleComplexMat)))) {
        blkTriExcept ("blkTriInit", "Malloc for *Work2 failed");
        return 1;
    }
    if (!(Work2->mat = (DoubleComplex *)
                        malloc (n * bs * sizeof (DoubleComplex)))) {
        blkTriExcept ("blkTriInit", "Malloc for Work2->mat failed");
        return 1;
    }
    Work2->m = bs;
    Work2->n = bs;

    return 0;
}

/***************************************************************************
 *
 * Abstract
 *   Release memories allocated for block tridiagonalization operation.
 *
 * Input
 *   None
 *
 * Output
 *   None
 *
 ***************************************************************************/

void blkTriEnd ()
{
    int i;

    free (A);
    A = NULL;

    for (i = 0; i < steps; i++) {
        free (WCur_m[i]);
        free (WOld_m[i]);
    }
    free (WCur_m);
    free (WOld_m);
    WCur_m = NULL;
    WOld_m = NULL;
    free (WCur);
    free (WOld);
    WCur = NULL;
    WOld = NULL;

    free (R->mat);
    free (R);
    R = NULL;
    free (Work1->mat);
    free (Work1);
    Work1 = NULL;
    free (Work2->mat);
    free (Work2);
    Work2 = NULL;
}

/***************************************************************************
 *
 * Abstract
 *   Handle exceptions thrown in the block tridiagonalization stage.
 *   Release memories. Usually there is no memory available during 
 *   initialization period, call this function.
 *
 * Input  
 *   pos    String to describe where the exception was thrown
 *   msg    String about the exception
 *
 * Output
 *   None
 *
 ***************************************************************************/

void blkTriExcept (char *pos, char *msg)
{
    int i;

    printf("%s: %s\n", pos, msg);

    if (A != NULL) {
        /* Don't free A->mat which user allocates.                  */ 
        free (A);
        A = NULL;
    }

    if (Q != NULL) {
        /* Q->mat is assigned by user instead, So don't free Q->mat */ 
        free (Q);
        Q = NULL;
    }

    if (M != NULL) {
        free (M);
        M = NULL;
    }
    if (M_m != NULL) {
        for (i = 0; i < steps; i++) {
            if (M_m[i] != NULL) {
                free (M_m[i]);
                M_m[i] = NULL;
            }
        }
        free (M_m);
        M_m = NULL;
    }

    if (B != NULL) {
        free (B);
        B = NULL;
    }
    if (B_m != NULL) {
        for (i = 0; i < steps - 1; i++) {
            if (B_m[i] != NULL) {
                free (B_m[i]);
                B_m[i] = NULL;
            }
        }
        free (B_m);
        B_m = NULL;
    }

    if (WCur != NULL) {
        free (WCur);
        WCur = NULL;
    }
    if (WCur_m != NULL) {
        for (i = 0; i < steps; i++) {
            if (WCur_m[i] != NULL) {
                free (WCur_m[i]);
                WCur_m[i] = NULL;
            }
        }
        free (WCur_m);
        WCur_m = NULL;
    }

    if (WOld != NULL) {
        free (WOld);
        WOld = NULL;
    }
    if (WOld_m != NULL) {
        for (i = 0; i < steps; i++) {
            if (WOld_m[i] != NULL) {
                free (WOld_m[i]);
                WOld_m[i] = NULL;
            }
        }
        free (WOld_m);
        WOld_m = NULL;
    }

    if (R != NULL) {
        if (R->mat != NULL) {
            free (R->mat);
            R->mat = NULL;
        }
        free (R);
        R = NULL;
    }

    if (Work1 != NULL) {
        if (Work1->mat != NULL) {
            free (Work1->mat);
            Work1->mat = NULL;
        }
        free (Work1);
        Work1 = NULL;
    }
    if (Work2 != NULL) {
        if (Work2->mat != NULL) {
            free (Work2->mat);
            Work2->mat = NULL;
        }
        free (Work2);
        Work2 = NULL;
    }
}

/***************************************************************************
 *
 * Abstract
 *   Print information about block Lanczos tridiagonalization.
 *
 * Input   
 *   iter    the current iteration. Starting from 1.   
 *
 * Output
 *   None
 *
 ***************************************************************************/

void blkTriInfo (int iter)
{
    char tag[50];
    DoubleComplexMat Q_next, M1, B1, WCur1;

    printf ("\n%s BlkTri Iteration #%d %s\n", separator, iter, separator);
    
    if (iter == 0) {
        showComplexMat (A, "\nMatrix A");
        showComplexMat (Q, "\nMatrix Q block #1");
    }
    else {
        sprintf (tag, "\nMatrix Q block #%d", iter + 1);
        Q_next.m = Q->m;
        Q_next.n = Q->n;
        Q_next.mat = Q->mat + Q->m * bs;
        showComplexMat (&Q_next, tag);
        
        sprintf (tag, "\nMatrix M block #%d", iter);
        M1.mat = M_m[iter];
        M1.m = bs;
        M1.n = bs;
        showComplexMat (&M1, tag);
        
        if (iter != steps) {
            sprintf (tag, "\nMatrix B block #%d", iter);
            B1.mat = B_m[iter];
            B1.m = bs;
            B1.n = bs;
            showComplexMat (&B1, tag);
        }        
    }
}