/* refine.c++
These routines use the "Correct - Diamond" approach for finding the force
  on an atoms
NEW global variables used:
geo.lmax = maximum distance from atom center that extend atoms density
          = 3.0?.
1=ca,2=c,3=o;4=n,5=cb; 6=cb2 in pos
*/
#include "drawwindow.h"
#include "math.h"
#include "trace.h"
#include "peptide.h"
#include "tordyn.h"
extern struct maiddatafile maidfile;
extern struct pepfit fit[NFIT]; // defined in peptide.h, see trace.c++ for definitin
extern int fitnum[NFIT];// see trace.c++ for definition

#define okhead struct geometry& geo,int ivec[3],float diffv[3],float rad,float bfact,struct grid& gzd // header for functions called by oneatomgrid
#define okvar geo,ivec,diffv,rad,bfact,gzd  // func variable for okfunc

#define gzdhead struct griddata& den1,struct geometry& geo,float cenpos[3],float force[3]

#define NDIV 5  // Number points divide 1 grid step into where atom center is placed (i.e. if = 5,then 1.0,1.2,1.4,1.6,1.8)
int maxgzd;
int scalechar =100;  // factor that mult den. by to convert to unsidhed character

float maxden = 0.0;
struct grid
{
   int num;  // number of grid points with radial dist from grid point to cenpos >geo.lmax
   int **dvec; // the x,y,z, value of the grid point  (centered at 0,0,0)
   int **forv;  // x,y,z value of force vecfor each grid point (int - to agree with char qden
   int *den; // den for each grid point
}; 

struct grid gzd[2+(NDIV+1)*(NDIV+1)*(NDIV+1)];


unsigned char  ***qden;

int igrid(int i,int j,int k)
/* convert 3d grid pos to 1D   1,2,....max
  Save 0 pos for max possible value
*/
{
   return 1+i+ (NDIV+1)*j + (NDIV+1)*(NDIV+1)*k;
}

void allocategzd(int npts)
{
   int i,j,k,ii;
   for(i=0;i<=NDIV;i++)
     for(j=0;j<=NDIV;j++)
       for(k=0;k<=NDIV;k++){
          ii = igrid(i,j,k);
          imat2(gzd[ii].dvec,npts,3);
          gzd[ii].den = new int[npts];
          imat2(gzd[ii].forv,npts,3);
       }
}

int gridint(int nn,float x)
/* returns the fractioal interval position
For example, if nn = 5, then round interval
  to 4.0, 4.2, 4.4., 4.6, 4.8, 5.0 corresponds
to (0,1,2,3,4, and 5)
*/
{
   int ix,rnd;
   ix = floor(x);
   rnd = nn*(x-ix) + 0.5;
   printf("nn = %d x = %5.3f  rnd = %d  Grid point %5.3f\n",nn,x,rnd,ix+(float)rnd/nn);
   return rnd;
}
   
void gridintvec(int nn,float fvec[3],int ivec[3],int ival[3])
/* 
  Takes the float fvec and returns the floor of (fvec) and 
    the rounded inteval number 
*/
{
   int k;

   for(k=0;k<=2;k++){
      ivec[k] = floor(fvec[k]);
      ival[k] = nn*(fvec[k]-ivec[k]) + 0.5;
   }
}   

float atomden1(float rad,float bfact)
/* Simplest function - 1 paramter, assume same for Carbon, nitrogen, etc.
   rad is distance from center; bfact = tempffact/(4Pi^2)
*/
{
   return exp(-rad*rad/bfact);
}

void countfunc(okhead)
/* This just counts number of grid points
*/
{
   gzd.num = gzd.num+1;
}

void gzdfunc(okhead)
/* This assigns values to gzd
*/
{  
   int k;
 
   gzd.num = gzd.num+1;
   if(gzd.num>maxgzd-1){
      printf("ERROR: gzd.num = %d is too big (maxgzd = %d)\n",gzd.num,maxgzd-2);
      exit(0);
   }
   for(k=0;k<=2;k++)
       gzd.dvec[gzd.num][k] =ivec[k];
   gzd.den[gzd.num] = scalechar* exp(-rad*rad/bfact) +0.5; // mult by 100 and round
   for(k=0;k<=2;k++)
       gzd.forv[gzd.num][k] = 4.0*rad*gzd.den[gzd.num]*diffv[k]/bfact;   
}

void assignforce(gzdhead)
/* Goes through all the grid points corresponding to center pos of atom
    and 
*/
{
   int i,k,ii,ivec[3],ivec0[3],ival[3];
   float pvec[3];
   float denscale;

   for(k=0;k<=2;k++)
     force[k] = 0.0;
   xtof(den1,cenpos,pvec); // conver to floating point grid pos.
   gridintvec(NDIV,pvec,ivec0,ival); // approx. pvec[k] by ivec[k] + ival[k]*1/NDIM
   ii = igrid(ival[0],ival[1],ival[2]); // convert to gzd point of cen pos.
   for(i=1;i<=gzd[ii].num;i++){
      for(k=0;k<=2;k++){ // the position of ith grid point
         ivec[k] = ivec0[k]+gzd[ii].dvec[i][k];
         if( (ivec[k]<0)||(ivec[k]>=den1.griddim[k]) )
           goto skipi;
      }
      denscale =(int)den1.den[ivec[0]][ivec[1]][ivec[2]];
      denscale = geo.denmult*denscale + geo.dentrans;
      for(k=0;k<=2;k++)
         force[k] = force[k]+(denscale - 
               qden[ivec[0]][ivec[1]][ivec[2]])*gzd[ii].forv[i][k];
      skipi:;
   } // i
}   

void zeroden(gzdhead)
/* Goes through all the grid points corresponding to center pos of atom
    and set qden = 0
*/
{
   int i,k,ii,ivec[3],ivec0[3],ival[3];
   float pvec[3];
   xtof(den1,cenpos,pvec); // conver to floating point grid pos.
   gridintvec(NDIV,pvec,ivec0,ival); // approx. pvec[k] by ivec[k] + ival[k]*1/NDIM
   ii = igrid(ival[0],ival[1],ival[2]); // convert to gzd point of cen pos.
   for(i=1;i<=gzd[ii].num;i++){
      for(k=0;k<=2;k++){ // the position of ith grid point
         ivec[k] = ivec0[k]+gzd[ii].dvec[i][k];
         if( (ivec[k]<0)||(ivec[k]>=den1.griddim[k]) )
           goto skipi;
      }
      qden[ ivec[0] ][ ivec[1] ][ ivec[2] ] = 0;
      skipi:;
   } // i
}   


void assignqden(gzdhead)
/* Goes through all the grid points corresponding to center pos of atom
    and den to qden
*/
{
   int i,k,ii,ivec[3],ivec0[3],ival[3];
   float pvec[3];
   xtof(den1,cenpos,pvec); // conver to floating point grid pos.
   gridintvec(NDIV,pvec,ivec0,ival); // approx. pvec[k] by ivec[k] + ival[k]*1/NDIM
   ii = igrid(ival[0],ival[1],ival[2]); // convert to gzd point of cen pos.
   for(i=1;i<=gzd[ii].num;i++){
      for(k=0;k<=2;k++){ // the position of ith grid point
         ivec[k] = ivec0[k]+gzd[ii].dvec[i][k];
         if( (ivec[k]<0)||(ivec[k]>=den1.griddim[k]) )
           goto skipi;
      }
float fden;
  fden = qden[ ivec[0] ][ ivec[1] ][ ivec[2] ] +gzd[ii].den[i];
if(fden> maxden)
  maxden = fden;  
      qden[ ivec[0] ][ ivec[1] ][ ivec[2] ] = qden[ ivec[0] ][ ivec[1] ][ ivec[2] ]
                              +gzd[ii].den[i];
      skipi:;
   } // i
}   

void assignqdenpos(gzdhead)
/* Same as assignqden, but only add densityh if qden>0
*/
{
   int i,ii,k,ivec[3],ivec0[3],ival[3];
   float pvec[3];
   xtof(den1,cenpos,pvec); // conver to floating point grid pos.
   gridintvec(NDIV,pvec,ivec0,ival); // approx. pvec[k] by ivec[k] + ival[k]*1/NDIM
   ii = igrid(ival[0],ival[1],ival[2]); // convert to gzd point of cen pos.
   for(i=1;i<=gzd[ii].num;i++){
      for(k=0;k<=2;k++){ // the position of ith grid point
         ivec[k] = ivec0[k]+gzd[ii].dvec[i][k];
         if( (ivec[k]<0)||(ivec[k]>=den1.griddim[k]) )
           goto skipi;
      }
      if((int)qden[ ivec[0] ][ ivec[1] ][ ivec[2] ]>0)
          qden[ ivec[0] ][ ivec[1] ][ ivec[2] ] = qden[ ivec[0] ][ ivec[1] ][ ivec[2] ]
                              +gzd[ii].den[i];
      skipi:;
   } // i
}   
    
void oneatomgrid(struct griddata& den1,struct geometry& geo,int ivec0[3],
          float cenpos[3],struct grid& gzd,float bfact,void (*okfunc)(okhead))
{
   int k,ig,istep,foundj;
   int jg,jstep,foundk,kg,kstep,ivec[3];
   float rad,fvec[3],diffv[3];

         ig= 1; /* step # in units of del, converted to proper fractional co-ord below */
         istep=0;
         foundj=1;
         while(foundj==1){
              foundj=0; 
              ivec[0]=ivec0[0]+ig*istep;
              jg= 1;
              jstep=0;
              foundk=1;
              while(foundk==1){
                 foundk=0;
                 ivec[1]=ivec0[1]+jg*jstep;
                 kg= 1;
                 kstep=0;
                 while(1){// exit this while by goto nextj
                    ivec[2]=ivec0[2]+kg*kstep;
                    imcv(den1,ivec,fvec);  //convert from  integer xi,xj (grid point)to world co-ord
                    for(k=0;k<=2;k++)
                      diffv[k] = fvec[k]-cenpos[k];
                    rad= sqrt( dot(diffv,diffv) );
                    if(rad>geo.lmax){
                       if(kg==1){/* reverse sign and repeat k search*/
                          kg= -1;
                          kstep=1;
                          goto reversek;
                       }
                       else { /* go to next j value */
                          goto nextj;
                       }
                     }
                     else{ // this point satisfies distance condition
                       okfunc(okvar);                        
                       foundk=1;
                       foundj=1;
                       kstep = kstep+1;
                     } // else                     
                     reversek:;
                 } //while
                 nextj:
                 if(foundk==1){ /* increase j and repeat */
                    jstep=jstep+1;
                 }
                 else{ /* foundk = 0: if jg = 1 , revrese and repeat, otherwise advance to next i*/
                    if(jg==1){ /* reverse step and repeat*/
                       foundk=1;
                       jg = -1;
                       jstep = 1;
                    }
                 }                  
              } /* while  foundk =1*/
              if(foundj==1){ /* increase i and repeat */
                    istep=istep+1;
              }
              else{ /* foundj = 0   - if ig =1, reverse and repeat otherwise quit*/
                    if(ig==1){ /* reverse step and repeat*/
                       foundj=1;
                       ig = -1;
                       istep = 1;
                    }
              } 
          } /*while foundj =1 */ 
}

#if 1
void tordataatoms(struct griddata& den1,struct geometry& geo,
      struct tordyn& tordata,float bfact,void (*okfunc)(gzdhead))
/* Before call diamondforce, must find the sum of the density from all atoms
    that are being moved to the grid points within geo.lmax of each atoms
It returns: qden[i]; i =1,2, 3 .... total number of grid points covered
  and jden where jden[j] = i where j is the 1D grid point and 
               jden[j] = i is the count number corresponding to this 1D point
Note: jden was allocated for entire map by scaleden - (different definition)
    qden was also allocated.  
    They are both treated as global variables - however, they are only visible 
       in this file.
 Since uses tordata.pos - this routine must sum over all atoms in tordata -
  The following code is excerpted from runtordyn()
*/
{
  int i,ii,k,j,jj;
  int icnt[2],tcount,cunit1,count,itype,natoms,cstep; 
  float force[3]; // not used here but passed to okfunc

   itype = 1; // indicates first unit is unit2
   icnt[0] = 1; // number of first unit1
   icnt[1] = 1;// number of first unit2
   tcount = tordata.nunit1+tordata.nunit2; 
   cunit1=icnt[0]-1; /*The unit 1 number minus 1 */
   for(count=1;count<=tcount;count++){ /* find inertia, force ietc for main chain atoms */
     if(itype==0){ /* unit 1 type*/
         cunit1=cunit1+1; /* counts unit1 type */
         natoms=2; //Always use Cb, even  if not in densith
         if( (tordata.sidedata[cunit1][0]>0)||(tordata.sidedata[cunit1][2]>=0)  ){// if side atoms or if is a residue type
            cstep=2; /* skip Cb because it is considered as part of side chain */
/*cstep=1;*/
            jj=tordata.sidedata[cunit1][1]; /*Assign Ca to uside:  jj= # of first side on this unit1 (cunit1)*/
            for(k=0;k<=2;k++)
              tordata.uside[jj][0][k]=tordata.unit[0][cunit1][0][k]; /* The pos. of Ca */
         }
         else{
            cstep=1;
         }
     } // if type 1
     else{ /* type 2 */
         if( (count==1)||(count==tcount) )/* at end - add atoms */
           natoms=4;
         else
           natoms=3; 
         cstep=1;
     }
     for(ii=0;ii<natoms;ii=ii+cstep){// sum over atoms in unit - set jden = 1 if density point
         assignqden(den1,geo,tordata.unit[itype][ icnt[itype] ][ii],force); 
     }//ii
     icnt[itype]=icnt[itype]+1;
     itype=(itype+1)%2;
   }// count

   for(i=1;i<=tordata.nunit1;i++){ /* Repeat above for sidecha. atoms*/
      if(tordata.sidedata[i][0]>0){
         for(j=tordata.sidedata[i][1]+1;j<tordata.sidedata[i][1]+tordata.sidedata[i][0];j++){/*Do not need for Ca */
            if( j==(tordata.sidedata[i][1]+tordata.sidedata[i][0]-1)){  /* last 2 atoms - find inertia etc.*/
               for(ii=0;ii<2;ii++){/* sum over atoms in unit - generate force and torge vectores ten.*/        
                    assignqden(den1,geo,tordata.uside[j][ii],force);
//                  findforden(den1,geo,tordata.uside[j][ii],forv);
             } /* ii */
            } /* if j = last unit*/
            else{ /* j not last unit, just one atom */
               assignqden(den1,geo,tordata.uside[j][0],force);
//               findforden(den1,geo,tordata.uside[j][0],forv);
           } /* if j != last unit*/
        } //j 
      } // sidedata
   } //i 
}
#endif   

void setupdiamondforce(struct griddata& den1,struct geometry& geo,
      struct tordyn& tordata)
/* Before call diamondforce, must find the sum of the density from all atoms
    that are being moved to the grid points within geo.lmax of each atoms
1) Calls tordataatoms with zeroden func -
    this sets qden[][][] = 0 for all values. 
   Note, this function call is inefficeint since many values overlap and
      are set = 0 several times.  Would be faster to store values
      which are set = 0 in some array and then just reset them here.
Runs through all tordata atoms and sums the density from all atoms
   and assigns to qden[][][]
Note this density is the atom density multiplied by 100 and rounded
*/
{
   float bfact;

   bfact = geo.tempfact/(4.0*PI*PI);
   tordataatoms(den1,geo,tordata,bfact,zeroden);
   tordataatoms(den1,geo,tordata,bfact,assignqden);
}

void diamondfor(struct griddata& den1,struct geometry& geo,float atompos[3],
       float force[3])
/* Calculates the diamond force on atom given its center position
Procedure:
  Applies oneatomgrid to atom with cen pos, finds force on atom which is 
the integral over all grid positions within geo.lmax of center,
and each posiiton  finds difference between the scaleddensity and
  the density from all the atoms in the moving structure
All the details in forcefunc() called by oneatomgrid
This force is scaled by 10^4 because of factor 100 in char den. This
   factor interes into both the (den-qden) and the gzd.for term
*/
{ 
   int k;
   assignforce(den1,geo,atompos,force);
   for(k=0;k<=2;k++){      
//       force[k] = force[k]/60000.0;
       force[k] = force[k]/6000.0;
   } 
//printf("force = %5.3f %5.3f %5.3f\n",force[0],force[1],force[2]);
}

void findgzd(struct griddata& den1,struct geometry& geo,float bfact)
/*
Assigns the grid values, etc. to gzd
First go through oneatomgrid for (ivec = 0,0,0) and counts points, set gzd.maxnum
Then allocates gzd.
Then finds values for gzd
*/
{
   int ii,i,j,k,ivec0[3];
   float cenpos[3],fvec[3],del;

   ivec0[0]=ivec0[1]=ivec0[2]=0;
   imcv(den1,ivec0,cenpos);  //convert from  integer xi,xj (grid point)to world co-ord
   gzd[0].num = 0; 
   oneatomgrid(den1,geo,ivec0,cenpos,gzd[0],bfact,countfunc);
   printf("totnum = %d\n",gzd[0].num);
   maxgzd = gzd[0].num = 1.5*gzd[0].num; // increase by 15%
   allocategzd(gzd[0].num);
   del = 1.0/NDIV;
   for(i=0;i<=NDIV;i++){
     fvec[0] = ivec0[0]+ i*del;
     for(j=0;j<=NDIV;j++){
       fvec[1] = ivec0[1]+ j*del;
       for(k=0;k<=NDIV;k++){
         fvec[2] = ivec0[2]+ k*del;
         mcv(den1,fvec,cenpos);
         ii = igrid(i,j,k);
         gzd[ii].num = 0;    
         oneatomgrid(den1,geo,ivec0,cenpos,gzd[ii],bfact,gzdfunc);
       }
     }
   }// i      
}

void scaleden0(struct griddata& den1,struct geometry& geo)
/* called if fit not known - used previously determined scale factors
*/
{
    int i,j,k;
    float bfact;

    printf(" USING DEFAULT SCALING FOR DIAMOND FORCE\n");
    bfact = geo.tempfact/(4.0*PI*PI);
    if(geo.denmult < -1000.0) // then qden not allocated
       unscmat3(qden,den1.griddim[0],den1.griddim[1],den1.griddim[2]);// allocate qden as unsigned char (0 to 255)
    for(i=0;i<den1.griddim[0];i++)
      for(j=0;j<den1.griddim[1];j++)
       for(k=0;k<den1.griddim[2];k++)
           qden[i][j][k] = 0;
    findgzd(den1,geo,bfact);
#if 1 // ctd
    geo.denmult=24.186;
    geo.dentrans = 46.479;
#endif
#if 0 // shad
    geo.denmult=23.298;
    geo.dentrans = 54.353;
#endif
}



// new version
void scaleden(struct griddata& den1,struct geometry& geo,float *rr)
/* Finds the scale factors for the density using the main chain atoms
   of all of the fits.
  Minimizes:
    R= sum(denscale - denatom)^^2 where
    denscale = ss*den +tt   where den is density in sigma units
Define:
 qden[][][] = unsigned char  - = sum of un scaled density from all atoms.
   Since unscaled atom den  = exp(-r^2/bfact) - max density from one atom = 1;
      and thus max density from all atoms at a point should be less than 2
   Thus, since unsigned char goes from 0 to 255, mult den by 100 and round
       before assigning to qden;
Set up the gzd structure: this tabulates grid points, density, etc. for all
  possible positions of the atom.  It rounds pos. to 0.2 of a grid length 
  (i.e. NDIM = 5).
Then go through all fits and sum density from each atom and assign to qden.
Then go through entire map and where ever qden>0 compare qden and observed
   den and find scale factors.
Coefficient in least squares routine:
   aa = sum of observed density^2 over all points
   m = total num of points
   bb = sum of observed density
   dd = sum of product of atom density*obsevvedden
   ee = sum of atom density
 aa,bb,dd, m(same as above = tot. num) and ee are coeficient in least square routine   
*/
{
   int  i,j,k,m;
   float bfact; 
   float oden,denom,norm;
   int itype,count;
   float aa,bb,dd,ee;
   float force[3]; // this is not used but need it in header for assignqden

//geo.lmax = 3.0;
   printf("Using %d fits to scaleden FOR DIAMOND FORCE\n",fitnum[0]);
   bfact = geo.tempfact/(4.0*PI*PI);
   if(geo.denmult < -1000.0) // then qden not allocated
       unscmat3(qden,den1.griddim[0],den1.griddim[1],den1.griddim[2]);// allocate qden as unsigned char (0 to 255)
   for(i=0;i<den1.griddim[0];i++)
      for(j=0;j<den1.griddim[1];j++)
       for(k=0;k<den1.griddim[2];k++)
           qden[i][j][k] = 0;
   findgzd(den1,geo,bfact);
   for(i=1;i<=fitnum[0];i++){ // go through  loop  sum density
     for(j=1;j<fit[i].nres-1;j++){ // each residue with side chains
       for(itype = 1;itype<=4;itype++){ // main cain atoms Ca,C, O, N
         assignqden(den1,geo,fit[i].pos[j][itype],force);
       }// itype
     }//j
   }//i
   for(i=1;i<=fitnum[0];i++){ // go through  loop again and add contribution from Cb only add if qden >0
     for(j=1;j<fit[i].nres-1;j++){ // each residue with side chains
         itype = 5; // Cb
         assignqdenpos(den1,geo,fit[i].pos[j][itype],force);
      }//j
   }//i
   for(i=1;i<=fitnum[0];i++){ // add density from end atoms
      j=0;
      itype = 2;
      assignqden(den1,geo,fit[i].pos[j][itype],force);
      j=fit[i].nres-1 ;
      itype = 4;
      assignqdenpos(den1,geo,fit[i].pos[j][itype],force);
   } // i
   // at this point have determined qden - density from all atoms.
   aa = bb = dd = ee = 0.0;
   count = 0;
   for(i=0;i<den1.griddim[0];i++){
     for(j=0;j<den1.griddim[1];j++)
       for(k=0;k<den1.griddim[2];k++){ // sum over all grid points
         if((int)qden[i][j][k]>0){
            count = count+1;
            oden = (int)den1.den[i][j][k]; // experimental density, 
            oden = oden/maidfile.denscale;  // convert oden to stand. dev. scale
            aa = aa + oden*oden;
            bb = bb + oden;
            dd = dd + oden*(int)qden[i][j][k];
            ee = ee + (int)qden[i][j][k];
         } // if >0
       }// k
   }//i
   m=count;
   denom = aa*m - bb*bb;
   geo.denmult = (dd*m - bb*ee)/denom;
   geo.dentrans = (aa*ee - dd*bb)/denom; 
   *rr = 0;  // function to be minimized
   norm = 0; // denomiator of rfactor
   for(i=0;i<den1.griddim[0];i++){
     for(j=0;j<den1.griddim[1];j++)
       for(k=0;k<den1.griddim[2];k++){ // sum over all grid points
         if((int)qden[i][j][k]>0){
            oden = (int)den1.den[i][j][k]; // experimental density, 
            oden = oden/maidfile.denscale;  // convert oden to stand. dev. scale
            *rr = *rr + sqr( geo.denmult*oden + geo.dentrans - (int)qden[i][j][k] );
             norm  = norm + sqr( geo.denmult*oden + geo.dentrans + (int)qden[i][j][k] );
         } // if >0
       }// k
   }//i
   norm = norm/2.0;
   printf("total points = %d \n",count);
   *rr = *rr/norm;
   printf("geo.lmax = %5.3f tempfact = %5.3f Scale factors: ss = %5.3f, tt = %5.3f, error = %5.3f\n",
      geo.lmax,geo.tempfact,geo.denmult,geo.dentrans,*rr);
printf("maxden = %5.3f\n",maxden);
}   
