// XYZbins.cc
// Robert Baertsch 
// 24 March 2002
//
//	Heavily modified by Kevin Karplus
//	8 April 2002

#include <assert.h>
#include "XYZbins.h"
#include "FixStartLength.h"

    const double MaxCoord= 9.e6;

// This constructor creates the bins and fills them with the
// all the points in the subarray of P.
// As in Pointlist, fix_start_length() is used to interpret
// negative values of start and length:
// start<0 is interpreted as
//    P.num_points()+start (so the last element of the list is -1).
// Length<0 is interpreted in a similar way, so that
//    (start,-1) means from start to the end of the list,
//    (start,-2) omits the last element, and so on.
//
// If atom_present is provided, then only those points whose
//	atom_present[p] is true will be included.
// The subscripting is identical in Atoms and atom_present
XYZbins::XYZbins(const Pointlist& Atoms, 
    double BinWidth,
    const bool *atom_present,
    int start, int length
    ) : Points(Atoms)
{
    fix_start_length(start,length, Atoms.num_points());
    int after= start+length;
    
    assert(BinWidth>0);
    
    // find range of points, and count points
    minp.X=minp.Y=minp.Z=  1.e35;
    maxp.X=maxp.Y=maxp.Z= -1.e35;
    int num_points=0;
    for (int a=start; a<after; a++)
    {	if (atom_present and !atom_present[a]) continue;
        const XYZpoint& p = Atoms[a];
	if (isnan(p.X) || isnan(p.Y) || isnan(p.Z)) continue;
	if (p.X >= MaxCoord) continue; // skip far points
        if (p.Y >= MaxCoord) continue; // skip far points
        if (p.Z >= MaxCoord) continue; // skip far points
    	if (p.X < minp.X) minp.X = p.X;
        if (p.X > maxp.X) maxp.X = p.X;
        if (p.Y < minp.Y) minp.Y = p.Y;
        if (p.Y > maxp.Y) maxp.Y = p.Y;
        if (p.Z < minp.Z) minp.Z = p.Z;
        if (p.Z > maxp.Z) maxp.Z = p.Z;
	num_points ++;
    }
    if (num_points==0)
    {	// we've been asked to create a empty set of bins
    	NumBinsX=NumBinsY=NumBinsZ = 0;
	ScaleX=ScaleY=ScaleZ=0;
	Bins = new IntArray* [1];	// avoid 0 pointer
	Bins[0] = 0;
	return;
    }
    
    assert (num_points > 0);
    NumBinsX = static_cast<int>(ceil((maxp.X - minp.X +1.e-5 ) / BinWidth));
    NumBinsY = static_cast<int>(ceil((maxp.Y - minp.Y +1.e-5) / BinWidth));
    NumBinsZ = static_cast<int>(ceil((maxp.Z - minp.Z +1.e-5) / BinWidth));
    
    // Since there must be at least 1 point, maxp>=minp,
    // and the slight extra 1.e-5 added should guarantee that NumBins
    // are positive.
    assert(NumBinsX >0);
    assert(NumBinsY >0);
    assert(NumBinsZ >0);
    
    int bin_count = NumBinsX * NumBinsY * NumBinsZ;
    Bins = new IntArray* [bin_count];
//    cerr << "XYZbins range: " << minp << " to " << maxp << "\n"
//    	<< "num_points = " << num_points << " out of " << Atoms.num_points()
//	<< "\n"
//	<< " binX=" << NumBinsX << " binY=" << NumBinsY 
//	<< " binZ=" << NumBinsZ << "\n" << flush; 
    
    ScaleX = NumBinsX/(maxp.X-minp.X+1.e-6);
    ScaleY = NumBinsY/(maxp.Y-minp.Y+1.e-6);
    ScaleZ = NumBinsZ/(maxp.Z-minp.Z+1.e-6);

    for (int i=bin_count-1; i>=0; i--)
        Bins[i]= 0;
    
    // How many points should be allowed for in each bin?
    // Typical maximum density of proteins is 0.045 atoms/cubic Angstrom
    int DefaultBinLength = static_cast<int>(ceil((BinWidth+1)*(BinWidth+1)*(BinWidth+1)*0.045));
    
    // distribute atoms to the appropriate bins
    for (int a=start; a<start+length; a++)
    {	if (atom_present and !atom_present[a]) continue;
    	const XYZpoint& p = Atoms[a];
	if (isnan(p.X) || isnan(p.Y) || isnan(p.Z)) continue;
        if (p.X >= MaxCoord) continue; // skip far points
        if (p.Y >= MaxCoord) continue; // skip far points
        if (p.Z >= MaxCoord) continue; // skip far points
        IntArray*& pl = bin_element(p);
        if (pl == 0)
        {   pl = new IntArray(DefaultBinLength);
        }
        pl->append_int(a);
    }

}

// This constructor fills the bins only with selected points---
// those whose subscripts are in
// selected_atoms[0]..selected_atoms[num_selected-1] 
XYZbins::XYZbins(const Pointlist& Atoms, 
    const int *selected_atoms, 		// indexes of atoms to put in bins
    int num_selected,			// how long is selected_atoms?
    double BinWidth
    ) : Points(Atoms)
{
    assert(BinWidth>0);
    if (num_selected==0)
    {	// we've been asked to create a empty set of bins
    	NumBinsX=NumBinsY=NumBinsZ = 0;
	ScaleX=ScaleY=ScaleZ=0;
	Bins = new IntArray* [1];	// avoid 0 pointer
	Bins[0] = 0;
	return;
    }
    
    // find range of points, and count points
    minp.X=minp.Y=minp.Z=  1.e35;
    maxp.X=maxp.Y=maxp.Z= -1.e35;
    for (int i=0; i<num_selected; i++)
    {	const XYZpoint& p = Atoms[selected_atoms[i]];
	if (isnan(p.X) || isnan(p.Y) || isnan(p.Z)) continue;
        if (p.X >= MaxCoord) continue; // skip far points
        if (p.Y >= MaxCoord) continue; // skip far points
        if (p.Z >= MaxCoord) continue; // skip far points
    	if (p.X < minp.X) minp.X = p.X;
        if (p.X > maxp.X) maxp.X = p.X;
        if (p.Y < minp.Y) minp.Y = p.Y;
        if (p.Y > maxp.Y) maxp.Y = p.Y;
        if (p.Z < minp.Z) minp.Z = p.Z;
        if (p.Z > maxp.Z) maxp.Z = p.Z;
    }
    
    NumBinsX = static_cast<int>(ceil((maxp.X - minp.X +1.e-5 ) / BinWidth));
    NumBinsY = static_cast<int>(ceil((maxp.Y - minp.Y +1.e-5) / BinWidth));
    NumBinsZ = static_cast<int>(ceil((maxp.Z - minp.Z +1.e-5) / BinWidth));
    
    // Since there must be at least 1 point, maxp>=minp,
    // and the slight extra 1.e-5 added should guarantee that NumBins
    // are positive.
    assert(NumBinsX >0);
    assert(NumBinsY >0);
    assert(NumBinsZ >0);
    
    int bin_count = NumBinsX * NumBinsY * NumBinsZ;
    Bins = new IntArray* [bin_count];
//    cerr << "XYZbins range: " << minp << " to " << maxp << "\n"
//    	<< "num_selected = " << num_selected << " out of " << Atoms.num_points()
//	<< "\n"
//	<< " binX=" << NumBinsX << " binY=" << NumBinsY 
//	<< " binZ=" << NumBinsZ << "\n" << flush; 
    
    ScaleX = NumBinsX/(maxp.X-minp.X+1.e-6);
    ScaleY = NumBinsY/(maxp.Y-minp.Y+1.e-6);
    ScaleZ = NumBinsZ/(maxp.Z-minp.Z+1.e-6);

    for (int i=bin_count-1; i>=0; i--)
        Bins[i]= 0;
    
    // How many points should be allowed for in each bin?
    // Typical maximum density of proteins is 0.045 atoms/cubic Angstrom
    int DefaultBinLength = static_cast<int>(ceil((BinWidth+1)*(BinWidth+1)*(BinWidth+1)*0.045));
    
    // distribute atoms to the appropriate bins
    for (int i=0; i<num_selected; i++)
    {	int a = selected_atoms[i];
    	const XYZpoint& p = Atoms[a];
	if (isnan(p.X) || isnan(p.Y) || isnan(p.Z)) continue;
        if (p.X >= MaxCoord) continue; // skip far points
        if (p.Y >= MaxCoord) continue; // skip far points
        if (p.Z >= MaxCoord) continue; // skip far points
        IntArray*& pl = bin_element(p);
        if (pl == 0)
        {   pl = new IntArray(DefaultBinLength);
        }
        pl->append_int(a);
    }
}

int XYZbins::num_points_near(const XYZpoint& spot, 
    double spot_radius
    )
{	
    if (NumBinsX==0) return 0;	// empty bins
    
    // figure out what range of bins to check
    XYZpoint minspot = spot - (spot_radius +0.01);
    int min_i, min_j, min_k;
    get_subscripts(min_i, min_j, min_k, minspot);

    XYZpoint maxspot = spot + (spot_radius + 0.01);
    int max_i, max_j, max_k;
    get_subscripts(max_i, max_j, max_k, maxspot);

    int Count = 0;
    double max_dist2 = spot_radius*spot_radius;
    for (int i=min_i; i<=max_i; i++)
    {   for (int j=min_j; j<=max_j; j++)
	{   for (int k=min_k; k<=max_k; k++)
	    {   IntArray* pl = bin_element(i,j,k);
		for (int m=0; pl!=NULL && m<pl->num_ints(); m++)
		{   int ai = (*pl)[m];	// atom number to check against
		    const XYZpoint& p = Points[ai];	// get point
		    if (sq_dist(spot,p) <=max_dist2)
		    {    // spot found
			Count++;       // increment burial count
		    }
		}
	    }
	}
    }
    return Count;
}

// Query the bins data structure repeatedly
// storing the result of QuerySpot for each spot in pre-allocated
// array burial.
// As in Pointlist, fix_start_length() is used to interpret
// negative values of start and length:
// The burial[0] value is for spot (*spots)[start].
void XYZbins::num_points_near(unsigned short* burial,
    const Pointlist *spots,
    double spot_radius,
    int start, int length
    )
{
    if (spots==NULL)	return;
    fix_start_length(start,length, spots->num_points());
    
    // for each spot, check distance to atoms
    for (long int a=start+length-1; a>=start; a--)
    {	burial[a-start] = num_points_near((*spots)[a], spot_radius);
    }
}

// Query repeatedly, interleaving the queries for spots1 and spots2.
// Preallocated burial array should be TWICE the length.
// NOTE: use of the 2-Pointlist version of num_spots_near is
// deprecated---it is provided only for compatibility with old
// Kestrel-based code. 
//
// If spots1==NULL or spots2==NULL, this code is equivalent
// to 1-Pointlist num_points_near---no interleaving is done.
void XYZbins::num_points_near(unsigned short* burial,
    const Pointlist *spots1,
    double spot_radius1,
    const Pointlist *spots2,
    double spot_radius2,
    int start, int length
    )
{
    if (spots1==NULL)
    {   num_points_near(burial, spots2, spot_radius2, start, length);
    	return;
    }
    if (spots2==NULL)
    {   num_points_near(burial, spots1, spot_radius1, start, length);
    	return;
    }
    
    assert(spots1->num_points() == spots2->num_points() );
    fix_start_length(start,length, spots1->num_points());
    
    // for each spot, check distance to atoms
    int burial_index=0;
    for (long int a=start; a<start+length; a++)
    {	burial[burial_index++] = num_points_near((*spots1)[a], spot_radius1);
	burial[burial_index++] = num_points_near((*spots2)[a], spot_radius2);
    }
}

void XYZbins::apply_to_points_near(
    ClosePairFcn *f,
    const XYZpoint& spot, 
    double spot_radius,
    int spot_index
    )
{	
    if (NumBinsX==0) return;	// empty bins
    
    // figure out what range of bins to check
    XYZpoint minspot = spot - (spot_radius +0.01);
    int min_i, min_j, min_k;
    get_subscripts(min_i, min_j, min_k, minspot);

    XYZpoint maxspot = spot + (spot_radius +0.01);
    int max_i, max_j, max_k;
    get_subscripts(max_i, max_j, max_k, maxspot);

    double max_dist2 = spot_radius*spot_radius;
    for (int i=min_i; i<=max_i; i++)
    {   for (int j=min_j; j<=max_j; j++)
	{   for (int k=min_k; k<=max_k; k++)
	    {   IntArray* pl = bin_element(i,j,k);
		for (int m=0; pl!=NULL && m<pl->num_ints(); m++)
		{   int ai = (*pl)[m];	// atom number to check against
		    const XYZpoint& p = Points[ai];	// get point
		    double dist2 = sq_dist(spot,p);
		    if (dist2 <=max_dist2)
		    {    // spot found
			(*f)(ai, p, spot_index, spot, dist2);
		    }
		}
	    }
	}
    }
}

// Query the bins data structure repeatedly
// applying the function f to any close pairs found.
// As in Pointlist, fix_start_length() is used to interpret
// negative values of start and length:
// The burial[0] value is for spot (*spots)[start].
void XYZbins::apply_to_points_near(ClosePairFcn *f,
    const Pointlist *spots,
    double spot_radius,
    int start, int length
    )
{
    if (spots==NULL)	return;
    fix_start_length(start,length, spots->num_points());
    
    // for each spot, check distance to atoms
    for (long int a=start+length-1; a>=start; a--)
    {	apply_to_points_near(f, (*spots)[a], spot_radius, a);
    }
}

// Report the subscript of the closest point to spot
// counting only points in the same bin as spot.
// Returns -1 if no point in same bin.
//
// dist2 is set to sq_dist(spot, Points[return value]), if point found.
int XYZbins::approx_closest(double &dist2, const XYZpoint& spot)
{
    dist2 = MaxCoord;
    IntArray* bin= bin_element(spot);
    if (bin==NULL)	return -1;
    int closest_subscr=-1;
    
    for (int m=bin->num_ints()-1; m>=0; m--)
    {   int ai = (*bin)[m];	// point number to check against
        const XYZpoint& p = Points[ai];	// get point
	double this_dist2 = sq_dist(spot, p);
	if (this_dist2 <dist2)
	{   // new closest found
	    dist2 = this_dist2;
	    closest_subscr =ai;
	}
    }
    return closest_subscr;
}




// Report the subscript of the closest point to spot
// counting only points with squared distance <= max_dist2.
// (Unless max_dist2< 0, then put no constraints on search.)
// Returns -1 if no point meeting constraints can be found.
int XYZbins::closest(double &dist2, const XYZpoint& spot, double max_dist2)
{
    // distance to closest spot found (or upper bound on such distance)
    dist2= max_dist2>=0? max_dist2: sq_dist(minp,maxp);
    
    // first search just the bin of spot
    int closest_subscr = approx_closest(dist2, spot);
    
    // figure out what range of bins to check
    double radius = sqrt(dist2);
    XYZpoint minspot = spot - (radius +0.01);
    int min_i, min_j, min_k;
    get_subscripts(min_i, min_j, min_k, minspot);

    XYZpoint maxspot = spot + (radius +0.01);
    int max_i, max_j, max_k;
    get_subscripts(max_i, max_j, max_k, maxspot);

    if (min_i==max_i && min_j==max_j && min_k==max_k)
    {   // no possibility of closer point outside already checked bin 
    	return closest_subscr;
    }

    // BAD CODE: it would probably be more efficient to check gradually
    // increasing shells around the central bin, rather than
    // sweeping the entire volume, but that would be messier code.
    for (int i=min_i; i<=max_i; i++)
    {   for (int j=min_j; j<=max_j; j++)
	{   for (int k=min_k; k<=max_k; k++)
	    {   IntArray* bin = bin_element(i,j,k);
		for (int m=0; bin!=NULL && m<bin->num_ints(); m++)
		{   int ai = (*bin)[m];	// atom number to check against
		    const XYZpoint& p = Points[ai];	// get point
		    double this_dist2 = sq_dist(spot, p);
		    if (this_dist2 <dist2)
		    {   // new closest found
    		    	dist2 = this_dist2;
			closest_subscr =ai;
		    }
		}
	    }
	}
    }
    
    return closest_subscr;        
}

// Return the subscript of the closest point to any points in spots
// counting only pairs with squared distance <= max_dist2.
// (Unless max_dist2< 0, then put no constraints on search.)
// Returns -1 if no pair meeting constraints can be found.
// spot_subscript is set to the index in spots of the point closest
// to the returned point.
// (Won't report spots for which spot_present[s]=false, if
//    spot_present provided.)
int XYZbins::closest(
	double &dist2, 
	int & spot_subscript,
	const Pointlist* spots,
	double max_dist2,
	const bool* spot_present,
	int start, int length)
{
    dist2= max_dist2>=0? max_dist2: sq_dist(minp,maxp);
    spot_subscript=-1;
    if (spots==NULL) return -1;
    
    int closest_subscr = -1;
    
    fix_start_length(start,length, spots->num_points());
    
    // first check all pairs in same bins
    for (int s=start+length-1; s>=start; s--)
    {   if (spot_present && !spot_present[s]) continue;
    	double this_dist2;
	int closest_to_this= approx_closest(this_dist2, (*spots)[s]);
	if (closest_to_this >=0 && this_dist2 < dist2)
	{   spot_subscript = s;
	    closest_subscr = closest_to_this;
	    dist2 = this_dist2;
	}
    }
    
    // Now, rescan to find real closest (shrink spot size as needed)
    for (int s=start+length-1; s>=start; s--)
    {   if (spot_present && !spot_present[s]) continue;
    	double this_dist2;
	int closest_to_this= closest(this_dist2, (*spots)[s], dist2);
	if (closest_to_this >=0 && this_dist2 < dist2)
	{   spot_subscript = s;
	    closest_subscr = closest_to_this;
	    dist2 = this_dist2;
	}
    }
    
    return closest_subscr;
    
}

// CHANGE LOG:
//	22 August 2002	Kevin Karplus
//	    Added 0.01 to radius when finding minspot and maxspot,
//		to avoid rounding errors for points near boundaries.
// 15 March 2004 Kevin Karplus
//	Fixed old-style casts
// Tue Jul 26 12:42:05 PDT 2005 Kevin Karplus
//	Added atom_present argument to basic constructor.
// Tue Jul 26 13:25:42 PDT 2005 Kevin Karplus
//	Added spot_present to closest 
// Mon Mar 27 12:55:55 PST 2006 Kevin Karplus
//	Added isnan test to constructors, so that points in list
//	could be skipped by making them NaN, keeping numbering intact.