/****************************************************************************
 *
 * LanTriAux.c
 *
 * Abstract
 *   Define auxiliar functions for 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
 *
 ****************************************************************************/

DoubleComplexVec *p = NULL;     /* One column of vector P, size n.      */
extern DoubleComplexMat *Q;     /* Unitary. one block = n * bs          */
extern DoubleComplexMat *M;     /* Diagonal blocks. bs * bs.            */                                     
extern DoubleComplexMat *B;     /* Off-Diagonal blocks. bs * bs         */
DoubleComplexVec *a = NULL;     /* Main diagonal of the tridiagonal. n  */
DoubleRealVec *b = NULL;        /* subdiagonal of the tridiagonal. n-1  */
DoubleComplexVec *wCur = NULL;  /* Dectector of current iter. size n.   */
DoubleComplexVec *wOld = NULL;  /* Dectector of previous iter.          */
DoubleComplexVec *r = NULL;     /* Starting vector and work vector. n   */
DoubleComplexVec *workZ = NULL; /* Work complex vector. size = n        */
DoubleRealVec *workD = NULL;    /* Work real vector. size = n           */

/* Memories for corresponding matrix or vector.                         */
extern DoubleComplex *P_m;      
extern DoubleComplex **M_m;                       
extern DoubleComplex **B_m;
extern DoubleComplex *a_m;
extern DoubleReal *b_m;
DoubleComplex *wCur_m = NULL;
DoubleComplex *wOld_m = NULL;   
DoubleComplex *r_m = NULL;
DoubleComplex *workZ_m = NULL;
DoubleReal *workD_m = NULL;


extern int n;
extern int bs;
int blks = 0;
int *up2 = NULL, *low2 = NULL;  /* Orthogonalization intervals.         */

extern char separator[35];

/***************************************************************************
 *
 * Abstract
 *   Perform initialization for tridiagonalization operation. Allocate
 *   necessary memories for computation.
 *
 * Input    
 *   None
 *
 * Output
 *   None
 *
 * Return  
 *   0      Successful exit
 *   1      Memory allocation exception occurs
 *
 ***************************************************************************/

int lanTriInit ()
{
    int i;
    int *array;

    blks = n / bs;

    if (!(p = (DoubleComplexVec *)malloc (sizeof (DoubleComplexVec)))) {
        lanTriExcept ("lanTriInit", "Malloc for *p failed");
        return 1;
    }
    p->n = n;

    if (!(a = (DoubleComplexVec *)malloc (sizeof (DoubleComplexVec)))) {
        lanTriExcept ("lanTriInit", "Malloc for *a failed");
        return 1;
    }
    a->n = n;
    a->vec = a_m;

    if (!(b = (DoubleRealVec *)malloc (sizeof (DoubleRealVec)))) {
        lanTriExcept ("lanTriInit", "Malloc for *a failed");
        return 1;
    }
    b->n = n - 1;
    b->vec = b_m;

    if (!(wOld_m = (DoubleComplex *)malloc (n * sizeof (DoubleComplex)))) {
        lanTriExcept ("lanTriInit", "Malloc for *wOld_m failed");
        return 1;
    }
    zeros (wOld_m, n);
    if (!(wOld = (DoubleComplexVec *)malloc (sizeof (DoubleComplexVec)))) {
        lanTriExcept ("lanTriInit", "Malloc for *wOld failed");
        return 1;
    }

    if (!(wCur_m = (DoubleComplex *)malloc (n * sizeof (DoubleComplex)))) {
        lanTriExcept ("lanTriInit", "Malloc for *wCur_m failed");
        return 1;
    }
    zeros (wCur_m, n);
    if (!(wCur = (DoubleComplexVec *)malloc (sizeof (DoubleComplexVec)))) {
        lanTriExcept ("lanTriInit", "Malloc for *wCur failed");
        return 1;
    }

    if (!(r = (DoubleComplexVec *)malloc (sizeof (DoubleComplexVec)))) {
        lanTriExcept ("lanTriInit", "Malloc for *r failed");
        return 1;
    }
    if (!(r_m = (DoubleComplex *)malloc (n * sizeof (DoubleComplex)))) {
        lanTriExcept ("lanTriInit", "Malloc for r_m failed");
        return 1;
    }
    /* Dynamically set r->n and r->vec later.   */

    if (!(workZ = (DoubleComplexVec *)malloc (sizeof (DoubleComplexVec)))) {
        lanTriExcept ("lanTriInit", "Malloc for *workZ failed");
        return 1;
    }
    if (!(workZ_m = (DoubleComplex *)malloc (n * sizeof (DoubleComplex)))) {
        lanTriExcept ("lanTriInit", "Malloc for *workZ_m failed");
        return 1;
    }
    /* Dynamically set workZ->n and workZ->vec  */

    if (!(workD = (DoubleRealVec *)malloc (sizeof (DoubleRealVec)))) {
        lanTriExcept ("lanTriInit", "Malloc for *workD failed");
        return 1;
    }
    if (!(workD_m = (DoubleReal *)malloc (n * sizeof (DoubleReal)))) {
        lanTriExcept ("lanTriInit", "Malloc for *workD_m failed");
        return 1;
    }

    if (!(up2 = (int *)malloc (n * sizeof (int)))) {
        lanTriExcept ("lanTriInit", "Malloc for *up failed");
        return 1;
    }
    array = up2;
    for (i = 0; i < n; i++) {
        array[i] = 1;
    }

    if (!(low2 = (int *)malloc (n * sizeof (int)))) {
        lanTriExcept ("lanTriInit", "Malloc for *low failed");
        return 1;
    }
    array = low2;
    for (i = 0; i < n; i++) {
        array[i] = 1;
    }

    return 0;
}


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

void lanTriEnd ()
{
    int i;

    /* Release the memory which is allocated at block tridiagonlization
     * stage, and still used at tridiagonalization stage. 
     * Q->mat is assigned by user, so don't free it 
     */
    free (Q);
    Q = NULL;

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

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

    /* Release memory allocated at tridiagonalization stage. P_m is 
     * assigned by user, so don't free it.
     */
    free (p);
    p = NULL;

    free (a);
    a = NULL;
    free (b);
    b = NULL;

    free (wOld_m);
    wOld_m = NULL;
    free (wOld);
    wOld = NULL;

    free (wCur_m);
    wCur_m = NULL;
    free (wCur);
    wCur = NULL;

    free (r_m);
    r_m = NULL;
    free (r);
    r = NULL;

    free (workZ);
    workZ = NULL;
    free (workZ_m);
    workZ_m = NULL;

    free (workD);
    workD = NULL;
    free (workD_m);
    workD_m = NULL;

    free (up2);
    up2 = NULL;

    free (low2);
    low2 = NULL;
}

/***************************************************************************
 *
 * Abstract
 *   Handle exceptions thrown in the tridiagonalization stage. Release
 *   memories one by one.
 *
 * Input  
 *   pos    String to describe where the exception was thrown
 *   msg    String about the exception
 *
 * Output
 *   None
 *
 ***************************************************************************/

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

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

    /* Release the memory which is allocated at block tridiagonlization
     * stage, and still used at tridiagonalization stage.      
     */
    if (Q != NULL) {
        free (Q);
        Q = NULL;
    }
    if (M != NULL) {
        free (M);
        M = NULL;
    }
    if (M_m != NULL) {
        for (i = 0; i < blks; 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 < blks - 1; i++) {
            if (B_m[i] != NULL) {
                free (B_m[i]);
                B_m[i] = NULL;
            }
        }
        free (B_m);
        B_m = NULL;
    }

    /* Release memory allocated at tridiagonalization stage.    */
    if (p != NULL) {
        free (p);
        p = NULL;
    }
    
    if (a != NULL) {
        free (a);
        a = NULL;
    }
    if (b != NULL) {
        free (b);
        b = NULL;
    }

    if (wOld_m != NULL) {
        free (wOld_m);
        wOld_m = NULL;
    }
    if (wOld != NULL) {
        free (wOld);
        wOld = NULL;
    }

    if (wCur_m != NULL) {
        free (wCur_m);
        wCur_m = NULL;
    }
    if (wCur != NULL) {
        free (wCur);
        wCur = NULL;
    }

    if (r_m != NULL) {
        free (r_m);
        r_m = NULL;
    }
    if (r != NULL) {
        free (r);
        r = NULL;
    }

    if (workZ != NULL) {
        free (workZ);
        workZ = NULL;
    }
    if (workZ_m != NULL) {
        free (workZ_m);
        workZ_m = NULL;
    }

    if (workD != NULL) {
        free (workD);
        workD = NULL;
    }
    if (workD_m != NULL) {
        free (workD_m);
        workD_m = NULL;
    }

    if (up2 != NULL) {
        free (up2);
        up2 = NULL;
    }
    if (low2 != NULL) {
        free (low2);
        low2 = NULL;
    }
}

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

void lanTriInfo (int iter)
{
    char tag[50];
    DoubleComplexVec p1, a1, wCur1;
    DoubleRealVec b1;

    printf ("\n%s LanTri Iteration #%d %s\n", separator, iter, separator);
    
    if (iter == 0) {
        sprintf (tag, "\nMatrix P column #%d", iter + 1);
        p1.vec = P_m + 1;
        p1.n = n;
        showComplexVec (&p1, tag);
    }
    else {
        if (iter != n) {
            sprintf (tag, "\nMatrix P column #%d", iter + 1);
            p1.vec = P_m + iter * n + 1;
            p1.n = n;
            showComplexVec (&p1, tag);
        }
        
        a1.vec = &a_m[1];
        a1.n = iter;
        showComplexVec (&a1, "\nVector a");
        
        if (iter != n) {
            b1.vec = &b_m[1];
            b1.n = iter;
            showRealVec (&b1, "\nVector b");
        }
    }
}