/* -------------------------------------------------------------------------- *
 *                                   OpenMM                                   *
 * -------------------------------------------------------------------------- *
 * This is part of the OpenMM molecular simulation toolkit originating from   *
 * Simbios, the NIH National Center for Physics-Based Simulation of           *
 * Biological Structures at Stanford, funded under the NIH Roadmap for        *
 * Medical Research, grant U54 GM072970. See https://simtk.org.               *
 *                                                                            *
 * Portions copyright (c) 2009 Stanford University and the Authors.           *
 * Authors: Mark Friedrichs, Mike Houston                                     *
 * Contributors:                                                              *
 *                                                                            *
 * This program is free software: you can redistribute it and/or modify       *
 * it under the terms of the GNU Lesser General Public License as published   *
 * by the Free Software Foundation, either version 3 of the License, or       *
 * (at your option) any later version.                                        *
 *                                                                            *
 * This program is distributed in the hope that it will be useful,            *
 * but WITHOUT ANY WARRANTY; without even the implied warranty of             *
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the              *
 * GNU Lesser General Public License for more details.                        *
 *                                                                            *
 * You should have received a copy of the GNU Lesser General Public License   *
 * along with this program.  If not, see <http://www.gnu.org/licenses/>.      *
 * -------------------------------------------------------------------------- */

/* After forces above, we have the forces for even numbered particles
 * in one stream, odd numbered particles in another.
 * In each stream, the forces are in several parts depending on how
 * many times we replicated the input stream.
 *
 * This kernel should not be a bottle neck, but if it turns out to
 * be so, we can try some work arounds.
 *
 * To avoid an extra kernel to zero forces, this sets the forces
 * rather than adding to it.
 * */

kernel static void kmerge_partial_forces( 
		float repfac, 
		float atomStrWidth, 
		float pforceStrWidth,
		float natoms,
		float3 pforce1[][], 
		float3 pforce2[][], 
		out float3 force<> )
{
	float linind;
	float2 pindex;
	float odd;
	float i;
	
	float2 adjustcount = (indexof force);
	
	//convert to linear atom index
	linind = adjustcount.x + adjustcount.y * atomStrWidth;

	//If odd or even, we pick from diferent streams.
	odd = linind - floor( linind / 2.0f ) * 2.0f;

	//Now linear index is the index into partial_forces
	linind = floor( linind / 2.0f );
	force = float3( 0.0f, 0.0f, 0.0f );

	//If we have predicated conditionals, we should 
	//keep the conditional inside the loop
	for ( i = 0; i < repfac; i+=1.0f ) {
		pindex.y = floor( linind / pforceStrWidth );		
		pindex.x = linind - pindex.y * pforceStrWidth;
		//pindex.x = fmod( linind, pforceStrWidth );
		
		if ( odd > 0.5f ) { //is odd
			force += pforce2[ pindex ];
		}
		else {
			force += pforce1[ pindex ];
		}
		linind += natoms/2.0f; 
	}
}

kernel void kMergeFloat3_4( 
      float repfac, 
      float atomStreamWidth, 
      float pStreamWidth,
      float natoms,
      float roundNatoms,
      float iUnroll,
      float3 pstream1[][], 
      float3 pstream2[][], 
      float3 pstream3[][], 
      float3 pstream4[][], 
      out float3 outstream<> )
{
   float atomIndex, forceIndex, qIndex, qOff;
   float2 pindex;
   float i;
   
   // given atom index find force indices and streams

   pindex      = indexof( outstream );
   atomIndex   = pindex.x + pindex.y*atomStreamWidth;
   forceIndex  = atomIndex;

   outstream   = float3( 0.0f, 0.0f, 0.0f );
   
   for( i = 0.0f; i < repfac; i += 1.0f ){

      qIndex            = round( (forceIndex - fmod( forceIndex, iUnroll))/iUnroll );
      qOff              = forceIndex - iUnroll*qIndex;
      pindex.y          = round( (qIndex - fmod( qIndex, pStreamWidth ))/pStreamWidth );
      pindex.x          = qIndex - pindex.y*pStreamWidth;

      if ( qOff < 0.5f ){ 
         outstream += pstream1[ pindex ];
      } else if( qOff < 1.5f ){
         outstream += pstream2[ pindex ];
      } else if( qOff < 2.5f ){
         outstream += pstream3[ pindex ];
      } else {
         outstream += pstream4[ pindex ];
      }

      forceIndex += roundNatoms;
   }
}


kernel void kMergeFloat3_4_nobranch( 
      float repfac, 
      float atomStreamWidth, 
      float pStreamWidth,
      float natoms,
      float roundNatoms,
      float iUnroll,
      float3 pstream1[][], 
      float3 pstream2[][], 
      float3 pstream3[][], 
      float3 pstream4[][], 
      out float3 outstream<> )
{
   float atomIndex, forceIndex, qIndex, qOff;
   float2 pindex;
   float i;
   float3 o1,o2,o3,o4;
   float3 tmp;
   
   // given atom index find force indices and streams

   pindex      = indexof( outstream );
   atomIndex   = pindex.x + pindex.y*atomStreamWidth;
   forceIndex  = atomIndex;

   outstream   = float3( 0.0f, 0.0f, 0.0f );
   
   for( i = 0.0f; i < repfac; i += 1.0f ){

      qIndex            = round( (forceIndex - fmod( forceIndex, iUnroll))/iUnroll );
      qOff              = forceIndex - iUnroll*qIndex;
      pindex.y          = round( (qIndex - fmod( qIndex, pStreamWidth ))/pStreamWidth );
      pindex.x          = qIndex - pindex.y*pStreamWidth;

      o1 = pstream1[ pindex ];
      o2 = pstream2[ pindex ];
      o3 = pstream3[ pindex ];
      o4 = pstream4[ pindex ];
      
      tmp = qOff < 0.5f ? o1 : o2;
      tmp = qOff < 1.5f ? tmp : o3;
      tmp = qOff < 2.5f ? tmp : o4;

      outstream += tmp;
      
      forceIndex += roundNatoms;
   }
}

