#include "pdb.h"
#define MAYBE_WATER(ar) ((int) ar == (int) &maybe_cavity_cell || (int) ar == (int) &maybe_water_cell )


double shell_mulecule_radius;
/*                 dummy records          */
atom_record not_contact_cell ;   /*  cells not in contact with protein */
atom_record maybe_water_cell ;   /* these will become water */
atom_record maybe_cavity_cell ;   /*  this will become buried water */
atom_record modified_cavity_cell ;   /*  maybe  buried water that are not buried*/
atom_record void_cell ;   /* water that are too close to protein atom*/


/* fill the borders of the grid with not_contact_cells, so the rest of the functions
   can start from the next line in. 
*/
  
void shell_with_not_contact (grid_descriptor * gd)
{
	int xdim = gd->xdim;
	int ydim = gd->ydim;
	int zdim = gd->zdim;
	int yzdim = ydim * zdim ;
	int xyzdim = xdim * ydim * zdim ;
       atom_record **grid =  gd->grid ;	
	int ii,jj,kk,ll ;

	jj  = (xdim -1 ) * yzdim ;
	kk = (ydim -1) * zdim ;
	ll = zdim -1 ;
	for (ii = 0 ; ii < yzdim  ; ii++)
	   {
		grid[ii] = &not_contact_cell ;
		grid[ii + jj] = &not_contact_cell ;
  	   }

	for (ii = 0 ; ii < xyzdim  ; ii += yzdim) 
	   {

		for (jj = ii ; jj < ii + zdim ; jj ++)
		   {
			grid[jj] = &not_contact_cell ;
			grid[jj + kk] = &not_contact_cell ;
		   }
		for (jj = ii ; jj < ii + yzdim ; jj += zdim)
		   {
			grid[jj] = &not_contact_cell ;
			grid[jj + ll] = &not_contact_cell ;
		   }
	   }
}


/*    fill the cell with a not_contact_cell pointer if it is not in contact
      a protein cell
*/

void fill_not_contact (grid_descriptor *gd, int xx, int yy, int zz)
{
   GRID_SLOTS(gd)
	int ii,jj,kk;

	if (GRID(xx,yy,zz) != NULL) return ;
 	for(ii = (xx -1 ) * yzdim ; ii < (xx + 2) * yzdim ; ii += yzdim)
	  for (jj = ii + (yy -1 ) * zdim ; jj < ii + (yy + 2) * zdim ; jj += zdim)
	    for (kk = jj + zz -1 ; kk < jj + zz + 2 ; kk ++)
		{
		   if ((int) grid[kk] != (int) NULL && 
				(int) grid[kk] != (int) &not_contact_cell  &&
				(int) grid[kk] != (int) & maybe_water_cell)
		     {
			GRID(xx , yy, zz) = &maybe_water_cell ;

			return ;
		     }
                }
	GRID(xx , yy, zz) = &not_contact_cell ;
}


/*    Check potential cavities if they are in contact with a normal water,
      if they are convert them into modified_cavity_water. It could
      be normal water but it is like for comparison with richards
*/

void check_cavities_again(grid_descriptor *gd, int xx, int yy, int zz)
{
   GRID_SLOTS(gd)
	int ii,jj,kk;
	if (GRID(xx,yy,zz) != &maybe_cavity_cell ) return ;
 	for(ii = (xx -1 ) * yzdim ; ii < (xx + 2) * yzdim ; ii += yzdim)
	  for (jj = ii + (yy -1 ) * zdim ; jj < ii + (yy + 2) * zdim ; jj += zdim)
	    for (kk = jj + zz -1 ; kk < jj + zz + 2 ; kk ++)
		{
		   if ((int) grid[kk] == (int) &maybe_water_cell)
		     {
			GRID(xx , yy, zz) = & modified_cavity_cell ;
			return ;
		     }
                }
}

/*   change the cell to maybe cavity if it is in contact with proteins from two 
     two opposing sides 
*/


void check_for_cavities (grid_descriptor *gd, int xx, int yy, int zz)
{
   GRID_SLOTS(gd)

	atom_record *ar[27]  ;
	int ar_ind = 0 ,ii,jj,kk;
	if (GRID(xx,yy,zz) != &maybe_water_cell) return ;
 	for(ii = (xx -1 ) * yzdim ; ii < (xx + 2) * yzdim ; ii += yzdim)
	  for (jj = ii + (yy -1 ) * zdim ; jj < ii + (yy + 2) * zdim ; jj += zdim)
	    for (kk = jj + zz -1 ; kk < jj + zz + 2 ; kk ++)
		  ar[ar_ind++] = grid[kk]  ;
	for(ar_ind = 0 ; ar_ind < 13 ; ar_ind++)
	  if (ar[ar_ind]->atomtype > 0 && ar[26 - ar_ind]->atomtype > 0 )
		{
		  GRID(xx,yy,zz) = &maybe_cavity_cell ;
		  return;
                }
}



/*   change every potentail water cell where the center of the cell is inside
     a vdw radius of a protein atom to void cell
*/


void check_in_vdw (atom_record *atom, grid_descriptor *gd,int index)
{
	if (atom->flags & IGNORE_ATOM) return ;
  {
   GRID_SLOTS(gd)


	  int xx = (int)floor((atom->x - origx)/edge) ;
	  int yy = (int)floor((atom->y - origy)/edge) ;
	  int zz = (int)floor((atom->z - origz)/edge) ;
	  double edge2 = edge / 2;
	  double xoff = atom->x - ( (xx - 0.5) * edge + origx) ;/* offsets from */
	  double yoff = atom->y - ( (yy - 0.5) * edge + origy) ;/* the atom     */	
	  double zoff = atom->z - ( (zz - 0.5) * edge + origz) ;
	
	int ii, jj, kk ;

	for (ii = xx -1; ii < xx +2 ; ii++)
           {
		double xsq = xoff * xoff ; 
		double ayoff = yoff ;
 	  	for (jj = yy -1; jj < yy +2 ; jj++)
		  {
			double ysq = ayoff * ayoff ; 
			double azoff = zoff ; 
			for (kk = zz -1; kk < zz +2 ; kk++ )
			   {
				atom_record *ar  = GRID(ii,jj,kk) ;

				if (MAYBE_WATER(ar) &&
					atom->definition->eradsq > xsq + ysq + azoff * azoff )
				  {
				    GRID(ii,jj,kk) = &void_cell;


				  }
				azoff -= edge;
			   }
			ayoff -= edge ;
								
		   }		
		xoff -= edge ;	  
	   }
  }
}




grid_descriptor * generate_shell_grid (file_records * file,double  probe)
{

	setup_file_records(file, 0.0,1) ;
	setup_minmax_xyzr(file,1) ;
     {
	grid_descriptor * gd = gridify(file,probe + probe, 6 * probe) ;


	shell_with_not_contact (gd) ;
	
	traverse_grid(gd, fill_not_contact, 1) ;

	traverse_grid(gd, check_for_cavities, 1) ;

	traverse_file_atom_records (file, check_in_vdw, gd) ;

	traverse_grid(gd, check_cavities_again, 1) ;
	
	return gd ;
     }
}


static int potential_water_count ;
file_records *water_shell ; 


/* this just counts the water mulecules (atoms) we are going to generate */
void increment_water_counter(grid_descriptor *gd, int xx , int yy, int zz)
{
	GRID_SLOTS(gd)
	int ar  =(int) GRID(xx,yy,zz) ;
	
	if(ar == (int) & maybe_water_cell  || ar == (int) & maybe_cavity_cell ||
				ar == (int) & modified_cavity_cell)
		potential_water_count++ ;
}

/*   generate the water mulecules. water_shell if the file_record to put it in,
     and potential_water_count is the index of the next place.
*/

void generate_dumm_water(grid_descriptor *gd, int xx , int yy, int zz)
{
	GRID_SLOTS (gd)
	int ar  =(int) GRID(xx,yy,zz) ; 
	int resnum ;
	
	if(ar == (int) & maybe_water_cell )
	      resnum = 0;
	else if ( ar == (int) & maybe_cavity_cell)
	    resnum = -2 ;
	else if (ar == (int) & modified_cavity_cell)
	    resnum = -5 ;
	else return ;

   {  	atom_record * new = (atom_record *) malloc (sizeof (atom_record)) ;
	new->atomtype = RECORD_DUMMY ;
	strcpy(new->ident,"DUMMY ") ;	
	new->number = potential_water_count ;
	strcpy(new->name, "O  ") ;
	new->extraflag = ' ' ;
	strcpy(new->resname,"HOH") ; 
	new->resnum = resnum ;
	new->x = origx + edge * (xx + 0.5) ;
	new->y = origy + edge * (yy + 0.5) ;
	new->z = origz + edge * (zz + 0.5) ;
	new->occ = 1.0 ;
	new->chain = ' '; 
	new->b = 0.0 ;  
	new->flags = DONT_CALCULATE_ATOM ;
	new->restype	 = find_residue_type(new->resname) ; 
        new->definition = find_atom_definition(new->restype,new->name) ;
	water_shell->atoms[ potential_water_count] = new ;
    }
   potential_water_count++ ;
}

	

	
file_records * generate_shell(file_records *file, double probe)
{
	grid_descriptor *gd = generate_shell_grid(file, probe) ;
	
	potential_water_count = 0 ;

	traverse_grid(gd, increment_water_counter, 1);
	water_shell = (file_records *) malloc_and_zero (sizeof(file_records)) ;
	
	water_shell->atomnum =  potential_water_count;
	water_shell-> atoms = (atom_record ** )malloc (sizeof(atom_record *) * potential_water_count);
	potential_water_count = 0 ;
	traverse_grid(gd, generate_dumm_water, 1);

	free_grid(gd) ;
	return water_shell ;
}