/* **********************************************************
 * Purpose::
 * heuristic multi-user waterfilling algorithm implementation
 * example invocation included; compile and run with
 * gcc -Wall (thisFile.c); ./a.out
 *
 * PLEASE NOTE:
 * This is a quick rewrite that hasn't been extensively tested.
 * Please verify independently that the code does what it says.
 */

/* **********************************************************
 * Header file
 * **********************************************************/

#ifndef WATERFILL_H
#define WATERFILL_H
#include <stdio.h>
#include <assert.h>
#include <math.h>
#include <stdlib.h>
#include <string.h>

/* **********************************************************
 * Waterfilling algorithm data structure
 * In C++ this would become an object.
 * Usage:
 * wfalg_new(...);
 * while (need to invoke it multiple times on changing data){
 *   wfalg_run(...);
 *   (process results)
 * }
 * wfalg_free(...); now the results become invalid  
 * **********************************************************/
typedef struct wfalg_t_ wfalg_t;

/* Purpose:
 * Create data structure for multiuser waterfiller "object"
 * Use for multiple invocations on varying data
 *
 * Parameters:
 * nRB: Number of "resource blocks" (frequency bins)
 *
 * nUser: number of simultaneous transmissions to allocate
 *
 * pmode: 0 conventional multi-user waterfilling: Per-resource power based on
 * channel quality estimate
 *        otherwise: select resources based on channel quality, but distribute the
 * per-user power evenly
 */
wfalg_t* wfalg_new(int /*nRB*/, int /*nUser*/, int /*mode*/);

/* Purpose:
 * Invoke multi-user waterfilling for the following parameters:
 * Rmax: maximum number of RBs that may be allocated to each individual user.
 *    in "ideal" waterfilling, a single user with a very good channel hogs up
 *    most of the resources. While this (ideally) maximizes sum throughput, there are
 *    usually other criteria that are equally or more important.
 *    The caller can limit the number of resources for a particular user, and thus
 *    control the "fairness" of allocation.
 *    Set Rmax >= nRB for each user to disable this limitation.
 *
 * SNR_threshold: Do not allocate a resource if the signal-to-noise ratio falls below the
 *    threshold. For example, even heavily coded QPSK doesn't lead to any meaningful
 *    throughput when the signal-to-noise ratio is below -3 dB
 *    if so, choose SNR_threshold=10*log10(-3/10) = 0.5
 *
 * Tinv: inverse of channel quality, noise-to-signal ratio N/S where "signal" includes
 *    the channel |H(f)| with f the frequency of a resource.
 *    Matrix: first user resource 0..nRB-1, second user resources 0..nRB-1, last user
 *
 * porder: Process users in the given order (array of user indices 0..nUser-1 with length nUser)
 *    The order is important in pathological cases, such as "more users than RBs."
 *    If so, the scheduling decision is effectively made by the caller, via the processing order:
 *    first come, first served, and the later ones get nothing.
 *
 * After execution,
 *   - resultAlloc points to a size-nRB array containing the index of the allocated user for each resource block.
     - resultPower points to a size-nRB array with the power allocated to a resource.
 * All powers for user j sum up to pUser[j]
 * The number of resources where resultAlloc==j is <= Rmax[j].
 * There is only one storage space per wfalg_t "object". Earlier results become invalid with
 * the next invocation.
 */
void wfalg_run(wfalg_t* /*obj*/, const double* /*pUser*/, const int* /*Rmax*/, const int* /*porder*/, const double* /*TinvMatrix*/, double /*SNR_threshold*/, int** /*resultAlloc*/, double** /*resultPower*/);

/* Purpose:
 * Release allocated memory.
 * obj becomes invalid.
 * any space provided by wfalg_run becomes invalid.
 */
void wfalg_free(wfalg_t* /* obj */);
#endif

/* **********************************************************
 * .c file
 * **********************************************************/

/* Data structure for multiple invocations without unnecessary malloc / free */
struct wfalg_t_ {
  /* parameters */
  int nRB;
  int nUser;
 
  /* Storage space for results */
  int* RB_alloc;
  double* RB_power;
 
  /* Internal temporary storage space
   * all Tinv values for the single user that is
   * processed by waterfill() */
  double* temp_singleUserTinv;
 
  /* internal temporary storage: How many resources are allocated per user? */
  int* temp_userNResources;
  int* order;
 
  /* Mode switch parameter */
  int mode;  
};

wfalg_t* wfalg_new(int nRB, int nUser, int mode){

  wfalg_t* obj=(wfalg_t*)malloc(sizeof(wfalg_t)); assert(obj);  
  obj->nUser=nUser;
  obj->nRB=nRB;
  obj->mode=mode;

  obj->RB_alloc=(int*)malloc(sizeof(int) * nRB); assert(obj->RB_alloc);
  obj->RB_power=(double*)malloc(sizeof(double) * nRB); assert(obj->RB_power);
  obj->temp_singleUserTinv=(double*)malloc(sizeof(double) * nRB); assert(obj->temp_singleUserTinv);
  obj->temp_userNResources=(int*)malloc(sizeof(int) * nUser); assert(obj->temp_userNResources);
  memset((void*)obj->RB_power, 0, nRB * sizeof(double));
 
  return obj;
}

void wfalg_free(wfalg_t* obj){
  free(obj->RB_alloc);
  free(obj->RB_power);
  free(obj->temp_singleUserTinv);
  free(obj->temp_userNResources);
  free(obj);
}

/* ************************************************************************************
 * Conventional multi-user waterfilling on the set of resources where
 * RB_alloc[]==userindex
 * expects temp_singleUserTinv to be filled with the Tinv for the allocating user!
 * ************************************************************************************/
static void waterfill(wfalg_t* obj, int userindex, double pUser){
  const int nRB=obj->nRB;

  const double E0=pUser;
  const double threshold=1e-12* E0;
 
  /* ************************************************************************************
   * Count resource allocated to this user
   * ************************************************************************************/
  int nRB_user=0;
  int j_RB;
  for(j_RB=0; j_RB < nRB; ++j_RB){
    if (*(obj->RB_alloc+j_RB)==userindex){
      ++nRB_user;
    }
  }
  if (nRB_user == 0){
    /* no resources allocated - can't waterfill */
    return;
  }
 
  double E=E0;
  double waterlevel=0;      
  while (E > threshold){
    /* ************************************************************************************
     * Distribute remaining energy evenly over all RBs
     * since some of them have greater-than-zero Tinv, this approximation will always stay
     * below target (standard waterfilling algorithm)
     * ************************************************************************************/
    E *= 0.999;
    waterlevel += E/(double)nRB_user;
   
    /* ************************************************************************************
     * Determine energy that remains with current allocation
     * ************************************************************************************/
    E=E0;

    /* Note: temp_singleUserTinv contains the Tinv for the user who allocates the RB.
     * This avoids many repeated table lookups from the nUser*nRB Tinv matrix. */
    double* pTinv=obj->temp_singleUserTinv;

    double* pRB_power=obj->RB_power;
    int* palloc=obj->RB_alloc;

    int j_RB;
    for(j_RB=0; j_RB < nRB; ++j_RB){    
      int alloc=*(palloc++);
      double Tinv=*(pTinv++);
     
      /* resource is allocated to current user */
      if (alloc==userindex){
       
        /* determine actual per-resource energy */
        double d=waterlevel - Tinv;
        d=(d > 0.0 ? d : 0.0);
        E -= d;
        *(pRB_power)=d;
      } /* if allocated */
     
      ++pRB_power;

    } /* for RB */
  } /* while energy remains */

  assert(E >= 0.0);
}

/* picks the best unallocated RB for the given user */
static int find_best_RB(wfalg_t* obj, int userindex, const double* thisUserTinv){
  /* find RB with lowest Tinv for this user */
  int i_RB;
  int bestRB_index=-1;
  double bestRB_Tinv=9e99;
  int valid=0;
  const double* pTinv=thisUserTinv;
  for (i_RB=0; i_RB < obj->nRB; ++i_RB){
    double Tinv=*(pTinv++);
    int alloc=*(obj->RB_alloc+i_RB);
    if ((alloc < 0) && (Tinv < bestRB_Tinv)){
      bestRB_index=i_RB;
      bestRB_Tinv=Tinv;
      valid=1;
    }    
  }
 
  if (!valid){
    bestRB_index=-1;
  }

  return bestRB_index; /* -1 means: none */
}

static int try_alloc_one_RB(wfalg_t* obj, int userindex, double pThisUser, const double* thisUserTinv, double SNR_threshold){
  int RBindex=find_best_RB(obj, userindex, thisUserTinv);
  if (RBindex >= 0){
   
    /* channel quality on resource RBindex */
    double Tinv=*(thisUserTinv+RBindex);
   
    /* allocate RB, at least temporarily */
    *(obj->RB_alloc+RBindex)=userindex;

    double Tinv2=Tinv;
    if (obj->mode){
      /* Equal power allocation - disregard "floor" level in waterfilling
       * We'll use the original Tinv for the SNR check later */
      Tinv2=0.0;
    }
   
    *(obj->temp_singleUserTinv+RBindex) = Tinv2;

    waterfill(obj, userindex, pThisUser);

    /* check SNR of last RB */
    double P=*(obj->RB_power + RBindex);
    double SNR=P/Tinv;
    if (SNR >= SNR_threshold){
      return 1; /* success */
    } else {
      /* too low power to be useful.
       * undo allocation. */
      *(obj->RB_alloc+RBindex)=-1;      
    }
  }
  return 0; /* failure */
}

static void zero_unallocated_RBs(const wfalg_t* obj){
  int* ap=obj->RB_alloc;  
  double* pp=obj->RB_power;
  int iRB;
  for (iRB=0; iRB < obj->nRB; ++iRB){
    *pp=(*ap >= 0 ? *pp : 0.0);
    ++pp;
    ++ap;
  }
}

void wfalg_run(wfalg_t* obj, const double* pUser, const int* Rmax, const int* porder, const double* TinvMatrix, double SNR_threshold, int** resultAlloc, double** resultPower){

  /* Deallocate all RBs */
  int* p=obj->RB_alloc;
  int i_RB;  
  for (i_RB=0; i_RB < obj->nRB; ++i_RB){
    *(p++) = -1;
  }
  memset((void*)obj->temp_userNResources, 0, obj->nUser * sizeof(int));

  int nAllocResourcesThisRound=-1; /* any nonzero value to enter the loop */

  /* Repeat until we fail to allocate another resource */
  while (nAllocResourcesThisRound != 0){
    nAllocResourcesThisRound=0;

    /* Iterate over all users ... */
    int jj_user;
    for (jj_user=0; jj_user < obj->nUser; ++jj_user){
      /* ... in the specified order */
      int j_user=*(porder+jj_user);

      /* power budget of current user */
      double p_jUser=*(pUser+j_user);
     
      /* Are we allowed to add another resource to this user? (Rmax constraint) */
      if (*(obj->temp_userNResources+j_user) < *(Rmax+j_user)){
       
        /* Look up the channel quality array for j_user */
        const double* thisUserTinv=TinvMatrix + j_user * obj->nRB;
       
        /* Try to allocate one more resource to this user */
        if (try_alloc_one_RB(obj, j_user, p_jUser, thisUserTinv, SNR_threshold)){
          /* Successfully allocated */
          ++ *(obj->temp_userNResources+j_user); /* count resources allocated to this user */
          ++nAllocResourcesThisRound; /* continue iterating */
        } else {
          /* this user can't allocate any more resources
           * There is no need to try again in future iterations,
           * disregard this user by making the Rmax test fail
           *
           * note, power allocation is not valid at this point */
          *(obj->temp_userNResources+j_user) = *(Rmax+j_user);
        }
      }      
    } /* for j_user */
  } /* while allocations succeed */
 
  /* The previous step fails exactly once per user, making the
   * power allocation invalid.
   *
   * Recalculate.
   * This is single-user waterfilling without any interaction, therefore order does not matter
   * Note that waterfill() requires a valid temp_singleUserTinv (which is the case at this point):
   * For each resource, it contains the Tinv of the allocating user
   */
  int j_user;
  for (j_user=0; j_user < obj->nUser; ++j_user){
    double p_jUser=*(pUser+j_user);    
    waterfill(obj, j_user, p_jUser);
  }
 
  /* Set zero power allocation for resources that aren't allocated to any user */
  zero_unallocated_RBs(obj);
 
  *resultAlloc=obj->RB_alloc;
  *resultPower=obj->RB_power;
}

/* **********************************************************
 * Example use (still .c file)
 * **********************************************************/
#if 1
int main(void){
  const int nUser=5;
  const int nRB=40;

  /* allocate and create random multi-user channel data */
  double* Tinv=(double*)malloc(nUser*nRB*sizeof(double));
  int i; int j;
  double* p=Tinv;
  for (i=0; i < nUser; ++i){
    for (j=0; j < nRB; ++j){
      *(p++)=(double)random() / (double)RAND_MAX;
    }
  }

  /* power per user */
  double pUser[]={1, 2, 3, 4, 5};
   
  /* processing order */
  int order[]={0, 1, 2, 3, 4};
 
  /* maximum number of resources per user */
  const int RMax[]={nRB, nRB, nRB, nRB, nRB};

  /* create waterfiller "object" */
  wfalg_t* wfalg=wfalg_new(nRB, nUser, /* mode */0);

  /* Invoke waterfilling */
  int* RB_alloc;
  double* RB_power;  
  wfalg_run(wfalg, pUser, RMax, order, Tinv, /* SNR_threshold */1.0, &RB_alloc, &RB_power);
 
  /* print result */
  int j_user; int i_RB;
  for (j_user=0; j_user  < nUser; ++j_user){
    printf("user %i ", j_user);
    double sum=0;
    for (i_RB=0; i_RB < nRB; ++i_RB){
      if (*(RB_alloc+i_RB)==j_user){
        double p=*(RB_power+i_RB);
        printf("%i=%lf ", i_RB, p);
        sum += p;
      }
    }
    printf("sum=%lf\n", sum);
  }
 
  /* clean up */
  wfalg_free(wfalg);
  free(Tinv);

  return EXIT_SUCCESS;
}
#endif