#ifndef __PARMARGINALS_CU__
#define __PARMARGINALS_CU__

/* PARMARGINALS.CU calculate marginals from energy
 *
 * Augustus Lidaka, Shitanshu Aggarwal, and Jerod Weinman
 *
 * Created July 15, 2009
 * Modified June 2, 2010
 *
 * $Id$
 */
#include <cuda.h>
#include "cuda_value.h"

/* Create some variables shared across each block.
 * Note that Z and uarray must have 2^n elements for some n.
 * uarray holds uint representations of the float elements in B for finding
 *      the max.
 * Z holds actual float values calculated from B used as intermediate values
 *      for the marginal probabilities--used in summing those elements.
 * fmx holds simply the maximum value from B.
 */
__shared__ __device__ unsigned int uarray[MAX_THREADS_PER_BLOCK];
__shared__ __device__ float Z[MAX_THREADS_PER_BLOCK];
__shared__ __device__ float fmx;


/* Notes: float_as_int is necessary to do an apparent bit-interpretation 
 * preserving cast. However, for negative numbers, the result seems to be
 * opposite of what we expect. Shifting down by 31 should leave us 
 * with the sign bit in the one's place, and zeros elsewhere. However, after 
 * negating and making the cast to uint32, we end up with one, rather than 
 * the expected "all ones." Need to see what the result of the shift is to
 * determine whether the shift is repeating the leading one or the uint32 cast
 * is giving unexpected results.
 * Leftmost bit is copied on a right shift.
 */
__device__ unsigned int floatFlip(float theFloat)
{
	unsigned int mask = (__float_as_int(theFloat) >> 31) | 0x80000000;
	return __float_as_int(theFloat) ^ mask;
}

__device__ float invFloatFlip(unsigned int theUint)
{
	unsigned int mask = ((theUint >> 31) - 1) | 0x80000000;
	return __int_as_float((int)(theUint ^ mask));
}


/* See value_cu_kernels.h for detailed documentation.
 * Recall the preconditions that gridDim.x >= N (preferably ==) and that
 * blockDim.x <= MAX_THREADS_PER_BLOCK and blockDim.x <= L (preferably ==).
 */
__global__ void parmarginals(float *B, const unsigned int *Y, const float
    *I, unsigned int N, int L, float eps, float *vv, int *error)
{

  unsigned int i;

  
  /* There are several reasons for a thread (or its entire thread block) to
   * exit immediately:
   *    1. The error flag is set.
   *    2. There are more threads per block than elements per row on which a
   *       block is operating (threadIdx.x > L).
   *    3. There are more blocks than rows on which to operate (blockIdx.x >
   *       N). 
   */
  if ( (*error) || (threadIdx.x > L) || (blockIdx.x > N) )
  {
    return;
  }

  /* Initialize some useful variables:
   *    B_block_start is this block's starting index in B
   *    B_block_end is the index of the first element in the next block's
   *        section of B
   *    B_thread_start is this thread's starting index in B
   *    B_thread_end is the index of the first element in the next thread's
   *        section of B
   * I.e., this block will be responsible for the blockIdx.x instance in B,
   * so indices (blockIdx.x * L) : ((blockIdx.x + 1) * L - 1), or
   * [B_block_start, B_block_end). Each thread will be operating on
   * [B_thread_start, B_thread_end). */
  unsigned int B_block_start = blockIdx.x * L;
  unsigned int B_block_end = (blockIdx.x + 1) * L;
  unsigned int B_thread_start;
  unsigned int B_thread_end;
  if (blockDim.x < L)
  {
    B_thread_start = B_block_start + ((L / blockDim.x) * threadIdx.x);
    B_thread_end = B_block_start + ((L / blockDim.x) * (threadIdx.x + 1));
    // The last thread will also handle the remainder due to integer division
    if (threadIdx.x == (blockDim.x - 1)) { B_thread_end = B_block_end; }
  }
  else
  {
    B_thread_start = B_block_start + threadIdx.x;
    B_thread_end = B_block_start + (threadIdx.x + 1);
  }



  /* MATLAB: fmx = max(B)
   * Calculate local maxima within a thread's chunk of B, then put the maximum
   * over all of this part of B (this instance) in fmx using the float/uint
   * bijection and shared array uarray. Synchronize all threads at a barrier
   * immediately afterward.
   */
  float local_fmx = B[B_thread_start];
  for (i = B_thread_start+1; i < B_thread_end; i++)
      local_fmx = fmaxf(local_fmx, B[i]);
  uarray[threadIdx.x] = floatFlip(local_fmx);

  /* Make sure each thread has calculated and uploaded its value before
   * continuing */
  __syncthreads();

  /* See cudamaxent/trunk/floatFlip.cu for explanation of float/uint bijection
   * Z is all zeroed, so we can treat it as if it had a power-of-two number
   * of elements. 2^(floor(log2(blockDim.x-1))) should be the highest power
   * of two less than L, which is at most MAX_THREADS_PER_BLOCK.
   */
  for (i = powf(2,floorf(log2f(blockDim.x-1))); i > 0; i >>= 1)
  {
      if (threadIdx.x < i && threadIdx.x + i < blockDim.x)
          atomicMax(&(uarray[threadIdx.x]), uarray[threadIdx.x + i]);

      __syncthreads();
  }

  // Make sure this is a valid value? Can you even do this for UINTs?
  if (isnan(uarray[0]))
  {
    *error = UINT_MAX_IS_NAN;
    return;
  }
  


  // Probably faster to have each thread change uint back to float than just
  // one thread and calling syncthreads() again.
  fmx = invFloatFlip(uarray[0]);



  // Make sure inverted (uint->float) value is valid
  if (isnan(fmx))
  {
    *error = FLOAT_MAX_IS_NAN;
    return;
  }



  /* MATLAB: Z(threadIdx.x) =
   *            exp(B(B_thread_start:B_thread_end) - fmx(ones(...)))
   * Normalize B by finding the sum of each row in log space and subtracting
   * that from B.
   */
  float local_Z = 0;
  for (i = B_thread_start; i < B_thread_end; i++)
      local_Z += expf(B[i] - fmx);
  Z[threadIdx.x] = local_Z;


  // Check that Z is non-negative
  if (local_Z < 0)
  {
    *error = THREAD_Z_SANS_MAX_IS_NEG;
    return;
  }

  if (isinf(local_Z))
  {
    *error = THREAD_Z_SANS_MAX_IS_INF;
    return;
  }


  /* Make sure each thread has calculated and uploaded its value before
   * continuing */
  __syncthreads();



  /* MATLAB: Z(1) = sum(Z)
   * Although Z is no longer zeroed, each thread uses its own respective
   * place--i.e., only the first blockDim.x cells have been used and will be
   * used again. The rest is still zero.
   */
  for (i = powf(2,floorf(log2f(blockDim.x-1))); i > 0; i >>= 1)
  {
      // Don't need read/write protection since only threads 0:i-1 will be
      // running and each thread is accessing, respectively, indices i:2i-1.
      if (threadIdx.x < i && threadIdx.x + i < blockDim.x)
          Z[threadIdx.x] += Z[threadIdx.x + i];

      __syncthreads();
  }

  float expZ, logZ;

  expZ = Z[0];

  /* Check:  sum(exp(B-mx)) */
  if (isnan(expZ))
  {
    *error = EXP_Z_SANS_MAX_IS_NAN;
    return;
  }

  /* Check:  log(sum(exp(B-mx))) */

  if (0 == expZ)
  {
    *error = EXP_Z_SANS_MAX_IS_ZERO;
    return;
  }


  if (isinf(expZ))
  {
    *error = EXP_Z_SANS_MAX_IS_INF;
    return;
  }

  if (expZ < 0)
  {
    *error = EXP_Z_SANS_MAX_IS_NEG;
    return;
  }

//  Removed. This seems redundant with the next check (and log is not cheap)
//  if (isnan(logf(expZ)))
//  {
//    *error = LOG_Z_SANS_MAX_IS_NAN;
//    return;
//  }


  /* MATLAB: log Z = fmx + log(sum(exp(Z-fmx)))
   * Probably faster to have all threads do this than just one and force
   * another barrier call. Note that we can't put the value back in fmx or
   * Z[0] since that would introduce a race condition. Fortunately, Z[1] is
   * free for use.
   */
  logZ = fmx + logf(expZ);

  /* check that value of Z is valid */
  if (isnan(logZ))
  {
    *error = LOG_Z_IS_NAN;
    return;
  }


  /* If the error flag is set, exit. */
  if (*error)
    return;

  /* Let leader thread calculate the value (log probability) for the instance */
  if (threadIdx.x == 0)
  {
    i = B_block_start + Y[blockIdx.x] - 1; // Index of correct label
    vv[blockIdx.x] = B[i] - logZ;  // Normalized log probability

    if (NULL != I)
      vv[blockIdx.x] *= I[blockIdx.x];
  }

  __syncthreads();


  /* Now we calculate the marginals
   * Start by subtracting log Z from B. If there are any NaN
   * values at that point, set the error flag and exit. Then scale by I if
   * it is populated (I != NULL), and exponentiate B. */

  float bval;
  for (i = B_thread_start; i < B_thread_end; i++)
  {

    bval = B[i] - logZ;

    if (isnan(bval))
    {
      *error = LOG_B_IS_NAN;
      return;
    }

    bval = expf(bval);
    
    if (NULL != I)
      bval *= I[blockIdx.x];

    B[i] = bval;


  }

  return;

}

#endif



/*

  Copyright 2009,2010 Augustus Lidaka, Shitanshu Aggarwal, Jerod Weinman

  This file is part of CUDAMAXENT.

    CUDAMAXENT is free software: you can redistribute it and/or modify
    it under the terms of the GNU General Public License as published by
    the Free Software Foundation, either version 3 of the License, or
    (at your option) any later version.

    CUDAMAXENT 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
    General Public License for more details.

    You should have received a copy of the GNU General Public License
    along with CUDAMAXENT.  If not, see <http://www.gnu.org/licenses/>.

*/
