/* Test program to read in matrix and pass to subroutine HICHOL
   for Cholesky decomposition */

#include <stdio.h>
#include <stdlib.h>
#include <math.h>

#define FITS 34         /* Total number of nights */
#define STARS 12        /* Total number of stars */
#define NSTARS 12       /* Number of stars/night */
#define FITPARMS 4      /* Number of non-global parameters per night */ 
#define STARPARMS 2     /* Number of parameters per star */

#define NV FITS*(FITPARMS*(FITPARMS+1)/2+FITPARMS*NSTARS*STARPARMS)
#define NT STARS*STARPARMS*(STARS*STARPARMS+1)/2
#define NP FITS*FITPARMS+STARS*STARPARMS

#define NVD FITS*(FITPARMS*(FITPARMS+1)/2+FITPARMS*NSTARS*STARPARMS) \
          + STARS*STARPARMS*(STARPARMS+1)/2

int hichol(float v[],unsigned long iv[],unsigned short jv[],
           unsigned short ni[],unsigned short ns,unsigned short nm, 
           float t[]);
int hi_chol(float v[],unsigned long iv[],unsigned short jv[],
            unsigned short ir,unsigned short nr,unsigned short mb,
            float t[],unsigned short nm,unsigned short nb);
int cholesky(float t[],unsigned short n);
int hi_solve(float v[],unsigned long iv[],unsigned short jv[],
             unsigned short nm,float t[],unsigned short nb,float r[]);

main()
        {
        char filename[80];
        FILE *fp;

        unsigned long i,j,ii,jj;
        unsigned short nm,ns,nb,np;
        unsigned long nv,nt;
        double sum;

        /* Input data */
        double vd[NVD+1];
        double rs[NP+1],rn[NP+1];
        unsigned short jv[NVD+1];
        unsigned short ni[FITS+1];
        unsigned long iv[NP+2];

        float v[NV+1],t[NT+1],r1[NP+1],r2[NP+1];

        puts("\nThis program reads a matrix and ");
        puts("calculates the Cholesky decomposition.");

        printf("%s","\nPlease enter filename of sparse matrix: ");
        scanf("%s",filename);
        if((fp=fopen(filename,"r"))==NULL){
                fprintf(stderr,"Error opening file.");
                exit(1);       
                }
        i=1;
        while(i<=NVD){
                if(fscanf(fp,"%lf %hu",&vd[i],&jv[i])==EOF)break;
                i+=1;
                }
        fclose(fp);

        printf("%s","Please enter filename for array iv: ");
        scanf("%s",filename);
        if((fp=fopen(filename,"r"))==NULL){
                fprintf(stderr, "Error opening file.");
                exit(1);       
                }
        i=1;
        while(i<=NP+2){
                if(fscanf(fp,"%lu",&iv[i])==EOF)break;
                i+=1;
                }
        nm=i-2;
        fclose(fp);

        printf("%s","Please enter filename for right hand side: ");
        scanf("%s",filename);
        if((fp=fopen(filename,"r"))==NULL){
                fprintf(stderr, "Error opening file.");
                exit(1);       
                }
        i=1;
        while(i<=NP+1){
                if(fscanf(fp,"%lf",&rs[i])==EOF)break;
                i+=1;
                }
        fclose(fp);

        /* Copy right hand side to r1 */
        for(i=1;i<=nm;i++){r1[i]=rs[i];}

        printf("Dimension of matrix is: %d",nm);
        printf("%s","\nEnter number of fits: ");
        scanf("%hu",&ns);

        /* Set array of sub-matrix dimensions */
        for(i=1;i<=ns;i++){ni[i]=4;}

        /* Prepare matrices v and t */
        nb=0;
        for(i=1;i<=ns;i++){nb+=ni[i];}
        nb=nm-nb;  /* Number of border columns */
        np=nm-nb;  /* Number of columns in block diagonal part */
        nv=iv[np+1];
        for(i=1;i<nv;i++){v[i]=vd[i];}
        for(i=1;i<NT+1;i++){t[i]=0.0;}
        for(i=np+1;i<nm+1;i++){
                ii=i-np;
                for(j=iv[i];j<iv[i+1];j++){
                        jj=jv[j]-np;
                        t[(ii-1)*nb-(ii-1)*(ii-2)/2+1+jj-ii]=vd[j];
                        }
                }

        puts("Call hichol now...");
        hichol(v,iv,jv,ni,ns,nm,t);
        puts("Finished hichol.");

        hi_solve(v,iv,jv,nm,t,nb,r1);

        if((fp=fopen("r1.out","w"))==NULL){
                fprintf(stderr, "Error opening file.");
                exit(1);
                }
        for(i=1;i<=nm;i++){
                fprintf(fp,"%4.13lf \n ",r1[i]);
                }
        fclose(fp);

        /* Try to improve solution */
        for(i=1;i<=nm;i++){
                sum=-rs[i];
                for(j=iv[i];j<iv[i+1];j++){sum+=vd[j]*r1[jv[j]];}
                rn[i]=sum;
                }
        for(i=1;i<=nm;i++){
                for(j=iv[i]+1;j<iv[i+1];j++){rn[jv[j]]+=vd[j]*r1[i];}
                r2[i]=rn[i];
                }
        hi_solve(v,iv,jv,nm,t,nb,r2);
        for(i=1;i<=nm;i++){r1[i]-=r2[i];} 

        if((fp=fopen("r2.out","w"))==NULL){
                fprintf(stderr, "Error opening file.");
                exit(1);
                }
        for(i=1;i<=nm;i++){
                fprintf(fp,"%4.13lf \n ",r1[i]);
                }
        fclose(fp);

        /* Try to improve solution even further */
        for(i=1;i<=nm;i++){
                sum=-rs[i];
                for(j=iv[i];j<iv[i+1];j++){sum+=vd[j]*r1[jv[j]];}
                rn[i]=sum;
                }
        for(i=1;i<=nm;i++){
                for(j=iv[i]+1;j<iv[i+1];j++){rn[jv[j]]+=vd[j]*r1[i];}
                r2[i]=rn[i];
                }
        hi_solve(v,iv,jv,nm,t,nb,r2);
        for(i=1;i<=nm;i++){r1[i]-=r2[i];} 

        if((fp=fopen("r3.out","w"))==NULL){
                fprintf(stderr, "Error opening file.");
                exit(1);
                }
        for(i=1;i<=nm;i++){
                fprintf(fp,"%4.13lf \n ",r1[i]);
                }
        fclose(fp);

/*      puts("\nNow start calculating inverse...");
        for(i=1;i<=nm;i++){
                for(j=1;j<=nm;j++){r2[j]=0.0;}
                r2[i]=1.0;
                hi_solve(v,iv,jv,nm,t,nb,r2);
                }
        puts("Finished calculating inverse");
*/
        return 0;
        }

int hichol(float v[],unsigned long iv[],unsigned short jv[],
           unsigned short ni[],unsigned short ns,unsigned short nm, 
           float t[])
        {
        unsigned short i;
        unsigned short nb,ib,nr,mb;

        /*
        Performs Cholesky-decomposition of sparse matrices.
        Matrix must be positive definite, symmetric doubly-bordered
        block diagonal. Diagonal blocks can have different sizes.
        Border band may be sparse too, but must consist of blocks with
        vertical extent matching the corresponding diagonal block. 
        The upper triangle of the matrix is stored in row-wise format 
        in v[], except for the upper triangle of the lower right block
        formed by the overlap of both borders, which is stored row-wise
        in t[]. The storage of the former must account for zeros in 
        lower right part; the latter must be filled, i.e. non-sparse.
        The transpose of v[] is vt[]. iv[] indicates position in v[],
        where row i begins, jv[] gives column number of element v[i].
        Dimension of the matrix is nm. There are ns diagonal blocks,
        each with dimension ni[i]. (All arrays [0..N-1].) */


        /* Calculate number of columns in border band of matrix.
        This is the dimension of matrix t[] at the same time. */
        nb=0;
        for(i=1;i<=ns;i++){nb+=ni[i];}
        nb=nm-nb;

        /* Calculate decomposition down to lower right triangle.*/
        ib=1;              /* We start with first block */
        nr=ni[ib];         /* Number of rows in block to operate on */
        mb=iv[2]-iv[1]-nr; /* Number of non-zero columns in border block */
        for(i=1;i<=nm-nb;i++){
                hi_chol(v,iv,jv,i,nr,mb,t,nm,nb); 
                nr-=1;
                if(nr==0){
                        ib+=1;
                        nr=ni[ib];
                        mb=iv[i+2]-iv[i+1]-nr;
                        }
                }

        /* Calculate decomposition of lower right triangle */
        cholesky(t,nb);          

        return 0;
        }

int hi_chol(float v[],unsigned long iv[],unsigned short jv[],
            unsigned short ir,unsigned short nr,unsigned short mb,
            float t[],unsigned short nm,unsigned short nb)
        {
        float v1;
        unsigned long i,j,ii,jj,jjmax,i1,i2;

        /* Work on current row */
        v1=v[iv[ir]]=sqrt(v[iv[ir]]);
        for(i=iv[ir]+1;i<=iv[ir+1]-1;i++){v[i]/=v1;}

        /* Work on remaining rows of current block */
        jj=ii=iv[ir]+1;
        jjmax=iv[ir+1];
        for(i=iv[ir+1];i<=iv[ir+nr]-1;i++){
                v[i]-=v[ii]*v[jj];
                if((jj+=1)==jjmax){jj=++ii;}
                }

        /* Work on rows in lower right triangle */
        i2=iv[ir+1]-1;
        i1=i2-mb+1;
        for(i=i1;i<=i2;i++){
                ii=jv[i]-(nm-nb);         
                for(j=i;j<=i2;j++){
                        jj=jv[j]-(nm-nb); 
                        t[(ii-1)*nb-(ii-1)*(ii-2)/2+1+jj-ii]-=v[i]*v[j];
                        }
                }   
        return 0;
        }

int cholesky(float t[],unsigned short n)
        {
        float t1;
        unsigned short i;
        unsigned long ii,j,k,l;

        ii=1;
        for(i=n;i>0;i--){
                t1=t[ii]=sqrt(t[ii]);
                for(j=ii+1;j<ii+i;j++){t[j]/=t1;}
                k=l=ii+1;
                for(j=ii+i;j<=n*(n+1)/2;j++){
                        t[j]-=t[k]*t[l];
                        if(++l==ii+i){l=++k;}
                        }
                ii+=i;
                }
        return 0;
        }

int hi_solve(float v[],unsigned long iv[],unsigned short jv[],
             unsigned short nm,float t[],unsigned short nb,float r[])
        {
        /* Given the Cholesky decomposition of a matrix as returned
        by HICHOL, solve for right hand side r[] by forward substitution
        and backsubstitution.*/

        float sum;
        unsigned long ii,jj,i,j,k,l;

        /* Skip leading zeros in right hand side */
        for(ii=1;r[ii]==0;ii++);

        /* Forward substitution in sparse part...*/
        if(ii<=nm-nb){
                for(i=ii;i<=nm-nb;i++){
                        r[i]/=v[iv[i]];
                        for(j=iv[i]+1;j<iv[i+1];j++){
                                r[jv[j]]-=v[j]*r[i];
                                }
                        }
                }
        /* ...and in triangle */
        if(ii<nm-nb+1){ii=nm-nb+1;}
        ii-=nm-nb;
        for(i=ii;i<=nb;i++){
                r[i+nm-nb]/=t[(i-1)*nb-(i-1)*(i-2)/2+1];
                for(j=i+1;j<=nb;j++){
                        r[j+nm-nb]-=t[(i-1)*nb-(i-1)*(i-2)/2+1+j-i]*r[i+nm-nb];
                        }
                }

        /* Back substitution in sparse part...*/
        for(i=nm;i>nm-nb;i--){
                sum=r[i];
                k=i-(nm-nb);
                for(j=i+1;j<=nm;j++){
                        l=j-(nm-nb);
                        sum-=t[(k-1)*nb-(k-1)*(k-2)/2+1+l-k]*r[j];
                        }
                r[i]=sum/t[(k-1)*nb-(k-1)*(k-2)/2+1];
                }
        /* ...and in triangle */
        for(i=nm-nb;i>0;i--){
                sum=r[i];
                for(j=iv[i]+1;j<=iv[i+1]-1;j++){
                        sum-=v[j]*r[jv[j]];
                        }
                r[i]=sum/v[iv[i]];
                }

        return 0;
        }