/* least-squares  inversion */

#include <stdio.h>
#include <math.h>
#include "../iolib/message.h"

void leastinv(dim,ata,atb,cov,xi,step_damp,flag)
int dim;
float **ata;  /*upper triangle filled up*/
float *atb;
float **cov;  /*inverse model covariance matrix, upper triangule*/
float *xi;    /* model from previous iteration minus a priori model*/
float step_damp;
int flag;     /* 0 --> no inverse ata is calculated */
{
float rms,trace;
double add;
int i,j,k;
void invmat();

  rms=0.0;
  for(i=0;i<dim;i++)
  { add=0.0;
    for(j=0;j<i;j++) add+=cov[j][i]*xi[j];
    for(j=i;j<dim;j++) 
    { add+=cov[i][j]*xi[j];
/*    if(i==338) printf("%d %g %g %g\n",j,add,cov[i][j],xi[j]); */
    }
    atb[i]-=add;
    
    rms+=add*xi[i];
    
    
    for(j=i;j<dim;j++) ata[i][j]+=cov[i][j];
  }
  rms=sqrt((double)(rms/(double)dim));
  printf("RMS of the starting model is %e\n",rms);
  fflush(stdout);

  k=-100;
  for(i=0;i<dim;i++)
  { j=log10(abs((double)ata[i][i]));
    if(j>k) k=j;
  }
  trace=1.0;
  if(k>0) for(i=0;i<k;i++) trace*=10.0;
  if(k<0) for(i=0;i>k;i--) trace*=0.1;
  
  for(i=0;i<dim;i++)
  { atb[i]/=trace;
    for(j=i;j<dim;j++) ata[i][j]/=trace;
  }
  
  invmat(dim,ata,atb,flag);
  
  if(flag)
  { for(i=0;i<dim;i++)
    for(j=i;j<dim;j++) ata[i][j]/=trace;
  }
  
  for(i=0;i<dim;i++) atb[i]=xi[i]+step_damp*atb[i];
  
  rms=0.0;
  for(i=0;i<dim;i++)
  { add=0.0;
    for(j=0;j<i;j++) add+=cov[j][i]*atb[j];
    for(j=i;j<dim;j++) add+=cov[i][j]*atb[j];
    rms+=add*atb[i];
  }
  rms=sqrt((double)(rms/(double)dim));
  printf("RMS of the new model is %e\n",rms);
  
}

void invmat(dim,ata,atb,flag)
int dim,flag;
float **ata,*atb;
{
int i,j,counter;

  if(flag)
  { float **a,**b,**farray2();
    a=farray2(1,dim,1,dim);  /*RECIPES package*/
    b=farray2(1,dim,1,1);    /*RECIPES package*/
    for(i=1;i<=dim;i++)
    { b[i][1]=atb[i-1];
      for(j=1;j<i;j++) a[i][j]=ata[j-1][i-1];
      for(j=i;j<=dim;j++) a[i][j]=ata[i-1][j-1];
    }
    gaussj(a,dim,b,1);
    for(i=0;i<dim;i++)
    { atb[i]=b[i+1][1];
      for(j=0;j<dim;j++) ata[i][j]=a[i+1][j+1];
    }
    free_farray2(a,1,dim,1);
    free_farray2(b,1,dim,1);
  }
  
  else
  { float *a,*b,*wk,*farray1();
    int err;
    a=farray1(1,dim*(dim+1)/2);
    b=farray1(1,dim);
    wk=farray1(1,dim);
    counter=0;
    for(j=1;j<=dim;j++)
    { b[j]=atb[j-1];
      for(i=1;i<=j;i++) a[++counter]=ata[i-1][j-1];
    }
    if(counter!=dim*(dim+1)/2) STOP(leastinv: counter error); 
    choles(a,b,wk,dim,&err);
    if(err) STOP(leastinv: error returned from Choles);
    for(i=0;i<dim;i++) atb[i]=wk[i+1];
    free_farray1(a,1);
    free_farray1(b,1);
    free_farray1(wk,1);
  }
  
  
}