// Copyright (c) 2021, salesforce.com, inc.
// All rights reserved.
// SPDX-License-Identifier: BSD-3-Clause
// For full license text, see the LICENSE file in the repo root
// or https://opensource.org/licenses/BSD-3-Clause


// Real Business Chain implementation in CUDA C

#include <curand_kernel.h>
#include <math.h>

typedef enum {
  kConsumerType,
  kFirmType,
  kGovernmentType,
} AgentType;

const size_t kBatchSize = M_BATCHSIZE;
const size_t kNumConsumers = M_NUMCONSUMERS;
const bool kCountFirmReward = M_COUNTFIRMREWARD;
const size_t kNumFirms = M_NUMFIRMS;
const size_t kNumGovts = M_NUMGOVERNMENTS;
const float kMaxTime = M_MAXTIME;
const float kCrraParam = M_CRRA_PARAM;
const float kInterestRate = M_INTERESTRATE;
const size_t kNumAgents = kNumConsumers + kNumFirms + kNumGovts;
const bool kIncentivizeFirmActivity = M_SHOULDBOOSTFIRMREWARD;
const float kFirmBoostRewardFactor = M_BOOSTFIRMREWARDFACTOR;
const bool kUseImporter = M_USEIMPORTER;
const float kImporterPrice = M_IMPORTERPRICE;
const float kImporterQuantity = M_IMPORTERQUANTITY;
const float kLaborFloor = M_LABORFLOOR;

// Global state =
const size_t kNumPrices = kNumFirms;  // - prices,
const size_t kNumWages = kNumFirms;  // - wages,
const size_t kNumInventories = kNumFirms;  // - stocks,
const size_t kNumOverdemandFlags = kNumFirms;  // - good overdemanded flag,
const size_t kNumCorporateTaxes = kNumGovts;  // - corporate tax rate
const size_t kNumIncomeTaxes = kNumGovts;  // - income tax rate
const size_t kNumTimeDimensions = 1;  // - time step
const size_t kGlobalStateSize = kNumPrices + kNumWages + kNumInventories + kNumOverdemandFlags + kNumCorporateTaxes + kNumIncomeTaxes + kNumTimeDimensions;

const size_t kIdxPricesOffset = 0;
const size_t kIdxWagesOffset = kNumPrices;
const size_t kIdxStockOffset = kIdxWagesOffset + kNumInventories;
const size_t kIdxOverdemandOffset = kIdxStockOffset + kNumOverdemandFlags;
const size_t kIdxIncomeTaxOffset = kGlobalStateSize - 3;
const size_t kIdxCorporateTaxOffset = kGlobalStateSize - 2;
const size_t kIdxTimeOffset = kGlobalStateSize - 1;

// Consumer actions: consume, work, choose which firm to work for
const size_t kActionSizeConsumer = kNumFirms + 1 + 1;

// add budget and theta
const size_t kStateSizeConsumer = kGlobalStateSize + 1 + 1;
const size_t kIdxConsumerBudgetOffset = 0;

// offset from agent-specific state part of array
const size_t kIdxConsumerThetaOffset = 1;

// UNUSED for consumer. Actions are floats.
// __constant__ float cs_index_to_action[num_actions_consumer *
// kActionSizeConsumer]; const size_t kActionSizeConsumer = kNumFirms +
// kNumFirms; // consume + work
/*const size_t num_actions_consumer =
    NUMACTIONSkConsumerType; // depends on discretization*/

// Firm actions: set wage, set price, invest in capital
const size_t kActionSizeFirm = 3;

// Number of actions depends on discretization of continuous action space.
const size_t kNumActionsFirm = M_NUMACTIONSFIRM;

// budget, capital, production alpha, and one-hot firm ID
const size_t kStateSizeFirm = kGlobalStateSize + 1 + 1 + 1 + kNumFirms;

// offset from agent-specific state part of array
const size_t kIdxFirmBudgetOffset = 0;
const size_t kIdxFirmCapitalOffset = 1;
const size_t kIdxFirmAlphaOffset = 2;
const size_t kIdxFirmOnehotOffset = 3;

// Constant memory available from ALL threads.
// See https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html#constant
__constant__ float kFirmIndexToAction[kNumActionsFirm * kActionSizeFirm];

// Corporate + income tax rates
const size_t kGovtActionSize = 2;
const size_t kNumActionsGovernment = M_NUMACTIONSGOVERNMENT;
const size_t kGovtStateSize = kGlobalStateSize;
__constant__ float
    kGovernmentIndexToAction[kNumActionsGovernment * kGovtActionSize];

// One RNG state for each thread. Each thread is assigned to an agent in an env.
__device__ curandState_t
    *rng_state_arr[kBatchSize * kNumAgents]; // not sure best way to do this

// Offsets into action vectors
const size_t kIdxConsumerDemandedOffset = 0;
const size_t kIdxConsumerWorkedOffset = kNumFirms;
const size_t kIdxConsumerWhichFirmOffset = kNumFirms + 1;

// currently 1 govt
const size_t kIdxThisThreadGovtId = 0;

extern "C" {

// ------------------
// CUDA C Utilities
// ------------------
__device__ void CopyFloatArraySlice(float *start_point, int num_elems,
                                    float *destination) {
  for (int i = 0; i < num_elems; i++) {
    destination[i] = start_point[i];
  }
}

__device__ void CopyIntArraySlice(int *start_point, int num_elems,
                                  float *destination) {
  for (int i = 0; i < num_elems; i++) {
    destination[i] = start_point[i];
  }
}

// unfortunately, you can't do templates with extern "C" linkage required for
// CUDA, so we have to define different functions for each case.
__device__ int *GetPointerFromMultiIndexFor3DIntTensor(int *array,
                                                       const dim3 &sizes,
                                                       const dim3 &index) {
  unsigned int flat_index =
      index.z + index.y * (sizes.z) + index.x * (sizes.z * sizes.y);
  return &(array[flat_index]);
}

__device__ float *GetPointerFromMultiIndexFor3DFloatTensor(float *array,
                                                           const dim3 &sizes,
                                                           const dim3 &index) {
  unsigned int flat_index =
      index.z + index.y * (sizes.z) + index.x * (sizes.z * sizes.y);
  return &(array[flat_index]);
}

__device__ float *
GetPointerFromMultiIndexFor4DTensor(float *array, const size_t *sizes,
                                    const size_t *multi_index) {
  // don't use this for arrays that arne't exactly size 4!!!
  unsigned int flat_index = multi_index[3] + multi_index[2] * sizes[3] +
                            multi_index[1] * sizes[3] * sizes[2] +
                            multi_index[0] * sizes[3] * sizes[2] * sizes[1];
  return &(array[flat_index]);
}

__global__ void CudaInitKernel(int seed) {
  // we want to reset random seeds for all firms and consumers
  int tidx = threadIdx.x;
  const int kThisThreadGlobalArrayIdx = blockIdx.x * kNumAgents + threadIdx.x;

  if (tidx < kNumAgents) {
    curandState_t *s = new curandState_t;
    if (s != 0) {
      curand_init(seed, kThisThreadGlobalArrayIdx, 0, s);
    }
    rng_state_arr[kThisThreadGlobalArrayIdx] = s;
  }
}

__global__ void CudaFreeRand() {
  int tidx = threadIdx.x;
  const int kThisThreadGlobalArrayIdx = blockIdx.x * kNumAgents + threadIdx.x;

  if (tidx < kNumAgents) {
    curandState_t *s = rng_state_arr[kThisThreadGlobalArrayIdx];
    delete s;
  }
}

__device__ int SearchIndex(float *distr, float p, int l, int r) {
  int mid;
  int left = l;
  int right = r;

  while (left <= right) {
    mid = left + (right - left) / 2;
    if (distr[mid] == p) {
      return mid;
    } else if (distr[mid] < p) {
      left = mid + 1;
    } else {
      right = mid - 1;
    }
  }
  return left > r ? r : left;
}

// --------------------
// Simulation Utilities
// --------------------
__device__ AgentType GetAgentType(const int agent_id) {
  if (agent_id < kNumConsumers) {
    return kConsumerType;
  } else if (agent_id < (kNumConsumers + kNumFirms)) {
    return kFirmType;
  } else {
    return kGovernmentType;
  }
}

__device__ float GetCRRAUtil(float consumption, float crra_param) {
  return (powf(consumption + 1, 1.0 - crra_param) - 1.0) / (1.0 - crra_param);
}

__global__ void CudaResetEnv(float *cs_state_arr, float *fm_state_arr,
                             float *govt_state_arr, float *cs_state_ckpt_arr,
                             float *fm_state_ckpt_arr,
                             float *govt_state_ckpt_arr,
                             float theta_anneal_factor) {
  /*
  Resets the environment by writing the initial state (checkpoint) into the
  state array for this agent (thread).
  */

  const int kBlockId = blockIdx.x;
  const int kWithinBlockAgentId = threadIdx.x;

  if (kWithinBlockAgentId >= kNumAgents) {
    return;
  }

  AgentType ThisThreadAgentType = GetAgentType(kWithinBlockAgentId);

  float *state_arr;
  float *ckpt_arr;
  dim3 my_state_shape, my_state_idx;
  size_t my_state_size;

  if (ThisThreadAgentType == kConsumerType) {
    // This thread/agent is a consumer.
    my_state_size = kStateSizeConsumer;
    my_state_shape = {kBatchSize, kNumConsumers, kStateSizeConsumer};
    my_state_idx = {kBlockId, kWithinBlockAgentId, 0};
    state_arr = cs_state_arr;
    ckpt_arr = cs_state_ckpt_arr;
  } else if (ThisThreadAgentType == kFirmType) {
    // This thread/agent is a firm.
    my_state_size = kStateSizeFirm;
    my_state_shape = {kBatchSize, kNumFirms, kStateSizeFirm};
    my_state_idx = {kBlockId,
                    (unsigned int)(kWithinBlockAgentId - kNumConsumers), 0};
    state_arr = fm_state_arr;
    ckpt_arr = fm_state_ckpt_arr;
  } else {
    // This thread/agent is government.
    my_state_size = kGovtStateSize;
    my_state_shape = {kBatchSize, kNumGovts, kGovtStateSize};
    my_state_idx = {
        kBlockId,
        (unsigned int)(kWithinBlockAgentId - kNumConsumers - kNumFirms), 0};
    state_arr = govt_state_arr;
    ckpt_arr = govt_state_ckpt_arr;
  }

  float *my_state_arr = GetPointerFromMultiIndexFor3DFloatTensor(
      state_arr, my_state_shape, my_state_idx);

  float *my_ckpt_ptr = GetPointerFromMultiIndexFor3DFloatTensor(
      ckpt_arr, my_state_shape, my_state_idx);

  CopyFloatArraySlice(my_ckpt_ptr, my_state_size, my_state_arr);

  // anneal theta
  if (ThisThreadAgentType == kConsumerType) {
    my_state_arr[kGlobalStateSize + kIdxConsumerThetaOffset] *=
        theta_anneal_factor;
  }
}

__device__ void GetAction(float *action_arr, float *index_to_action_arr,
                          int index, int agent_idx, int agent_action_size) {

  // it needs to be possible to call this for either agents or firms
  // Note: each thread is an agent.
  for (int i = 0; i < agent_action_size; i++) {
    action_arr[agent_idx * agent_action_size + i] =
        index_to_action_arr[index * agent_action_size + i];
  }
}

__global__ void CudaSampleFirmAndGovernmentActions(
    float *fm_distr, int *fm_action_indices_arr, float *fm_actions_arr,
    float *govt_distr, int *govt_action_indices_arr, float *govt_actions_arr) {
  // Samples actions for firms and governments. Consumer actions are sampled in
  // Pytorch...
  const int kWithinBlockAgentId = threadIdx.x;

  // Unused threads should not do anything.
  if (threadIdx.x >= kNumAgents) {
    return;
  }

  AgentType ThisThreadAgentType = GetAgentType(kWithinBlockAgentId);

  // Index into rand states array
  int kThisThreadGlobalArrayIdx = blockIdx.x * kNumAgents + threadIdx.x;
  curandState_t rng_state = *rng_state_arr[kThisThreadGlobalArrayIdx];
  *rng_state_arr[kThisThreadGlobalArrayIdx] = rng_state;

  // float cs_cum_dist[num_actions_consumer];
  float fm_cum_dist[kNumActionsFirm];
  float govt_cum_dist[kNumActionsGovernment];

  float *my_cumul_dist;
  float *my_dist;
  int *my_indices;
  float *my_actions;
  float *index_to_action;
  int this_thread_global_array_idx;
  size_t my_num_actions;
  int my_action_size;

  // Consumers have multiple action heads, hence sampling is more complicated.
  if (ThisThreadAgentType == kConsumerType) {
    return;
  } else if (ThisThreadAgentType == kFirmType) {
    // on firm thread
    my_cumul_dist = fm_cum_dist;
    my_dist = fm_distr;
    my_indices = fm_action_indices_arr;
    my_actions = fm_actions_arr;
    my_num_actions = kNumActionsFirm;
    index_to_action = kFirmIndexToAction;
    my_action_size = kActionSizeFirm;
    this_thread_global_array_idx =
        (blockIdx.x * kNumFirms) + (threadIdx.x - kNumConsumers);
  } else {
    my_cumul_dist = govt_cum_dist;
    my_dist = govt_distr;
    my_indices = govt_action_indices_arr;
    my_actions = govt_actions_arr;
    my_num_actions = kNumActionsGovernment;
    index_to_action = kGovernmentIndexToAction;
    my_action_size = kGovtActionSize;
    this_thread_global_array_idx =
        (blockIdx.x * kNumGovts) + (threadIdx.x - kNumConsumers - kNumFirms);
  }

  // Compute CDF
  my_cumul_dist[0] = my_dist[this_thread_global_array_idx * my_num_actions];
  for (int i = 1; i < my_num_actions; i++) {
    my_cumul_dist[i] =
        my_dist[this_thread_global_array_idx * my_num_actions + i] +
        my_cumul_dist[i - 1];
  }

  // Given sampled action which is a float in [0, 1], find the corresponding
  // discrete action.
  float sampled_float = curand_uniform(&rng_state);
  const int index =
      SearchIndex(my_cumul_dist, sampled_float, 0, (int)(my_num_actions - 1));
  my_indices[this_thread_global_array_idx] = index;
  GetAction(my_actions, index_to_action, index, this_thread_global_array_idx,
            my_action_size);
}

__device__ float GetFirmProduction(float technology, float capital, float hours,
                                   float alpha) {
  if (hours < kLaborFloor) {
      hours = 0.0;
  }
  return technology * powf(capital, 1.0 - alpha) * powf(hours, alpha);
}

// --------------------
// Simulation Logic
// --------------------
__global__ void
CudaStep(float *cs_state_arr, float *cs_actions_arr, float *cs_rewards_arr,
         float *cs_state_arr_batch, float *cs_rewards_arr_batch,

         float *fm_state_arr, int *fm_action_indices_arr, float *fm_actions_arr,
         float *fm_rewards_arr, float *fm_state_arr_batch,
         int *fm_actions_arr_batch, float *fm_rewards_arr_batch,

         float *govt_state_arr, int *govt_action_indices_arr,
         float *govt_actions_arr, float *govt_rewards_arr,
         float *govt_state_arr_batch, int *govt_actions_arr_batch,
         float *govt_rewards_arr_batch,
         float *consumer_aux_batch,
         float *firm_aux_batch,
         int iter) {
  // This function should be called with 1 block per copy of the environment.
  // Within a block, each agent should have a thread.
  const int kWithinBlockAgentId = threadIdx.x;

  // return if we're on an extra thread not corresponding to an agent
  if (kWithinBlockAgentId >= kNumAgents) {
    return;
  }

  // -------------------------------------
  // Start of variables and pointers defs.
  // -------------------------------------

  // __shared__ variables are block-local: can be seen by each thread ** in the
  // block **
  __shared__ float gross_demand_arr[kNumFirms];
  __shared__ int num_consumer_demand_arr[kNumFirms];
  __shared__ float hours_worked_arr[kNumFirms];
  __shared__ float total_actually_consumer_arr[kNumFirms];
  __shared__ float bought_by_importer_arr[kNumFirms];
  __shared__ float next_global_state_arr[kGlobalStateSize];
  __shared__ float tax_revenue_arr[kNumGovts];
  __shared__ float total_utility_arr[kNumGovts];
  __shared__ bool need_to_ration_this_good_arr[kNumFirms]; // whether or not to
                                                           // ration good i

  float net_demand_arr[kNumFirms]; // amount demanded after budget constraints
                                   // by a consumer (ignore for non-consumers)

  int num_iter = (int)kMaxTime;
  AgentType ThisThreadAgentType = GetAgentType(kWithinBlockAgentId);
  float this_agent_reward = 0.0;

  // pointer to start of state vector
  // state vector consists of global state, then
  // agent-specific state global part is of same size for
  // all agents, but needs to be sliced out of different
  // arrays depending on agent type
  float *my_global_state_ptr;

  float *my_action_arr;

  // pointer to start of state vector in batch history
  float *batch_state_ptr;

  // sizes and indices for strided array access
  dim3 my_state_shape, my_state_idx, action_shape;

  // shape for batched array of scalars (action ind and reward)
  dim3 batch_single_shape, single_idx;

  // pointer to action index
  int *batch_action_value_ptr;

  // pointer to batch reward
  float *batch_reward_value_ptr;

  // pointers to action index for current arrays
  int *my_action_value_ptr;

  // pointers to reward index for current arrays
  float *my_reward_value_ptr;

  float *my_aux_batch_ptr;

  // -----------------------------------
  // End of variables and pointers defs.
  // -----------------------------------

  if (ThisThreadAgentType == kConsumerType) {
    // get current state
    my_state_shape = {kBatchSize, kNumConsumers, kStateSizeConsumer};
    my_state_idx = {blockIdx.x, threadIdx.x, 0};
    my_global_state_ptr = GetPointerFromMultiIndexFor3DFloatTensor(
        cs_state_arr, my_state_shape, my_state_idx); // index)

    // get current action
    action_shape = {kBatchSize, kNumConsumers, kActionSizeConsumer};
    my_action_arr = GetPointerFromMultiIndexFor3DFloatTensor(
        cs_actions_arr, action_shape, my_state_idx);

    // index into the episode history and save prev state into it
    size_t my_batch_state_shape[] = {kBatchSize, num_iter, kNumConsumers,
                                     kStateSizeConsumer};
    size_t my_batch_state_idx[] = {blockIdx.x, iter, threadIdx.x, 0};
    batch_state_ptr = GetPointerFromMultiIndexFor4DTensor(
        cs_state_arr_batch, my_batch_state_shape, my_batch_state_idx);
    CopyFloatArraySlice(my_global_state_ptr, kStateSizeConsumer,
                        batch_state_ptr);

    size_t my_aux_batch_shape[] = {kBatchSize, num_iter, kNumConsumers, kNumFirms};
    my_aux_batch_ptr = GetPointerFromMultiIndexFor4DTensor(
        consumer_aux_batch, my_aux_batch_shape, my_batch_state_idx
    );


    // Extract pointers to rewards, batch and current
    batch_single_shape = {kBatchSize, (unsigned int)num_iter, kNumConsumers};
    single_idx = {blockIdx.x, (unsigned int)iter, threadIdx.x};

    batch_reward_value_ptr = GetPointerFromMultiIndexFor3DFloatTensor(
        cs_rewards_arr_batch, batch_single_shape, single_idx);

    my_reward_value_ptr =
        &(cs_rewards_arr[blockIdx.x * kNumConsumers + threadIdx.x]);
  }

  if (ThisThreadAgentType == kFirmType) {
    // get current state
    size_t this_thread_firm_id = (threadIdx.x - kNumConsumers);
    my_state_shape = {kBatchSize, kNumFirms, kStateSizeFirm};
    my_state_idx = {blockIdx.x, (unsigned int)this_thread_firm_id, 0};
    my_global_state_ptr = GetPointerFromMultiIndexFor3DFloatTensor(
        fm_state_arr, my_state_shape, my_state_idx);

    // get current action
    action_shape = {kBatchSize, kNumFirms, kActionSizeFirm};
    my_action_arr = GetPointerFromMultiIndexFor3DFloatTensor(
        fm_actions_arr, action_shape, my_state_idx);

    // index into the episode history and save prev state into it
    size_t my_batch_state_shape[] = {kBatchSize, num_iter, kNumFirms,
                                     kStateSizeFirm};
    size_t my_batch_state_idx[] = {blockIdx.x, iter, this_thread_firm_id, 0};
    batch_state_ptr = GetPointerFromMultiIndexFor4DTensor(
        fm_state_arr_batch, my_batch_state_shape, my_batch_state_idx);
    CopyFloatArraySlice(my_global_state_ptr, kStateSizeFirm, batch_state_ptr);

    dim3 my_aux_batch_shape = {kBatchSize, num_iter, kNumFirms};
    dim3 aux_batch_idx = {blockIdx.x, iter, this_thread_firm_id};
    my_aux_batch_ptr = GetPointerFromMultiIndexFor3DFloatTensor(
        firm_aux_batch, my_aux_batch_shape, aux_batch_idx
    );
    // extract pointers to action indices and rewards, batch and current
    batch_single_shape = {kBatchSize, (unsigned int)num_iter, kNumFirms};
    single_idx = {blockIdx.x, (unsigned int)iter,
                  (unsigned int)this_thread_firm_id};
    batch_action_value_ptr = GetPointerFromMultiIndexFor3DIntTensor(
        fm_actions_arr_batch, batch_single_shape, single_idx);
    batch_reward_value_ptr = GetPointerFromMultiIndexFor3DFloatTensor(
        fm_rewards_arr_batch, batch_single_shape, single_idx);

    const int kThisThreadFirmIdx = blockIdx.x * kNumFirms + this_thread_firm_id;
    my_action_value_ptr = &(fm_action_indices_arr[kThisThreadFirmIdx]);
    my_reward_value_ptr = &(fm_rewards_arr[kThisThreadFirmIdx]);
  }

  if (ThisThreadAgentType == kGovernmentType) {

    int this_thread_govt_id = (threadIdx.x - kNumConsumers - kNumFirms);

    my_state_shape = {kBatchSize, kNumGovts, kGovtStateSize};
    my_state_idx = {blockIdx.x, (unsigned int)this_thread_govt_id, 0};
    my_global_state_ptr = GetPointerFromMultiIndexFor3DFloatTensor(
        govt_state_arr, my_state_shape, my_state_idx); // index)

    // get current action
    action_shape = {kBatchSize, kNumGovts, kGovtActionSize};
    my_action_arr = GetPointerFromMultiIndexFor3DFloatTensor(
        govt_actions_arr, action_shape, my_state_idx);

    // index into the episode history and save prev state into it
    size_t my_batch_state_shape[] = {kBatchSize, num_iter, kNumGovts,
                                     kGovtStateSize};
    size_t my_batch_state_idx[] = {blockIdx.x, iter, this_thread_govt_id, 0};
    batch_state_ptr = GetPointerFromMultiIndexFor4DTensor(
        govt_state_arr_batch, my_batch_state_shape, my_batch_state_idx);
    CopyFloatArraySlice(my_global_state_ptr, kGovtStateSize, batch_state_ptr);

    // extract pointers to action indices and rewards, batch and current
    batch_single_shape = {kBatchSize, (unsigned int)num_iter, kNumGovts};
    single_idx = {blockIdx.x, (unsigned int)iter,
                  (unsigned int)this_thread_govt_id};
    batch_action_value_ptr = GetPointerFromMultiIndexFor3DIntTensor(
        govt_actions_arr_batch, batch_single_shape, single_idx);
    batch_reward_value_ptr = GetPointerFromMultiIndexFor3DFloatTensor(
        govt_rewards_arr_batch, batch_single_shape, single_idx);

    const int kThisThreadGovtIdx = blockIdx.x * kNumGovts + this_thread_govt_id;
    my_action_value_ptr = &(govt_action_indices_arr[kThisThreadGovtIdx]);
    my_reward_value_ptr = &(govt_rewards_arr[kThisThreadGovtIdx]);
  }

  // ----------------------------------------------
  // State pointers and variables
  // Create pointers to agent-specific state that will be updated by the
  // simulation logic.
  // ----------------------------------------------
  float *my_state_arr = &(my_global_state_ptr[kGlobalStateSize]);
  float *prices_arr = &(my_global_state_ptr[kIdxPricesOffset]);
  float *wages_arr = &(my_global_state_ptr[kIdxWagesOffset]);
  float *available_stock_arr = &(my_global_state_ptr[kIdxStockOffset]);
  float time = my_global_state_ptr[kIdxTimeOffset];
  float income_tax_rate = my_global_state_ptr[kIdxIncomeTaxOffset];
  float corporate_tax_rate = my_global_state_ptr[kIdxCorporateTaxOffset];

  // -------------------------------
  // Safely initialize shared memory
  // -------------------------------
  if (ThisThreadAgentType == kFirmType) {
    int this_thread_firm_id = threadIdx.x - kNumConsumers;
    gross_demand_arr[this_thread_firm_id] = 0.0;
    num_consumer_demand_arr[this_thread_firm_id] = 0;
    hours_worked_arr[this_thread_firm_id] = 0.0;
    total_actually_consumer_arr[this_thread_firm_id] = 0.0;
    need_to_ration_this_good_arr[this_thread_firm_id] = false;
  }

  if (ThisThreadAgentType == kGovernmentType) {
    tax_revenue_arr[0] = 0.0;
  }

  __syncthreads();
  // -------------------------------------
  // End - Safely initialize shared memory
  // -------------------------------------

  // -------------------------------------
  // Process actions
  // -------------------------------------
  if (ThisThreadAgentType == kConsumerType) {
    // amount demanded is just the first part of the action vector
    float *this_agent_gross_demand_arr =
        &(my_action_arr[kIdxConsumerDemandedOffset]);
    const float this_agent_hours_worked =
        my_action_arr[kIdxConsumerWorkedOffset];
    const int worked_for_this_firm_id =
        (int)my_action_arr[kIdxConsumerWhichFirmOffset];

    // here, need to scale demands to meet the budget. put them in a local array
    // *budgetDemanded logic should be: compute total expenditure given prices.
    // if less than budget, copy existing demands
    float __cost_of_demand = 0.0;
    for (int i = 0; i < kNumFirms; i++) {
      __cost_of_demand += this_agent_gross_demand_arr[i] * prices_arr[i];
    }

    // Scale demand to ensure that total demand at most equals total supply
    // we want: my_state_arr being 0 always sends __scale_factor to 0
    float __scale_factor = 1.0;

    if ((__cost_of_demand > 0.0) && (__cost_of_demand > my_state_arr[0])) {
      __scale_factor = my_state_arr[0] / __cost_of_demand;
    }

    // otherwise scale all demands down to meet budget
    // copy them into a demanded array
    for (int i = 0; i < kNumFirms; i++) {
      net_demand_arr[i] = __scale_factor * this_agent_gross_demand_arr[i];
    }

    // adding up demand across threads **in the block**
    // somehow store amount demanded per firm in an array demanded (copy from
    // action) also store amount worked per firm in array worked

    // Every thread executes atomicAdd_block in a memory-safe way.
    for (int i = 0; i < kNumFirms; i++) {
      // sum across threads in block
      atomicAdd_block(&(gross_demand_arr[i]), net_demand_arr[i]);

      // increment count of consumers who want good i
      if (net_demand_arr[i] > 0) {
        atomicAdd_block(&(num_consumer_demand_arr[i]), 1);
      }
    }

    // increment total hours worked for firm i
    atomicAdd_block(&(hours_worked_arr[worked_for_this_firm_id]),
                    this_agent_hours_worked);
  }

  // wait for everyone to finish tallying up their adding
  __syncthreads();

  if (ThisThreadAgentType == kFirmType) {
    // check each firm if rationing needed
    int this_thread_firm_id = threadIdx.x - kNumConsumers;
    need_to_ration_this_good_arr[this_thread_firm_id] =
        ((gross_demand_arr[this_thread_firm_id] > 0.0) && (gross_demand_arr[this_thread_firm_id] >
         available_stock_arr[this_thread_firm_id]));
  }

  // wait for single thread to finish checking demands
  __syncthreads();

  // ----------------------------------------
  // Consumers: Rationing demand + Utility
  // ----------------------------------------
  // Logic:
  // case 1: no overdemand
  // case 2: overdemand, but some want less than 1/N -- fill everyone up to
  // max(theirs, 1/N) case 3: overdemand, everyone wants more -- fill everyone
  // up to max(theirs, 1/N)
  float net_consumed_arr[kNumFirms]; // per consumer thread
  // always add negligible positive money to avoid budgets becoming small
  // negative numbers otherwise, when computing proportions, one may end up with
  // negative stocks.
  float cs_budget_delta = 0.01;
  float fm_budget_delta = 0.01;
  float capital_delta = 0.0;

  if (ThisThreadAgentType == kConsumerType) {
    // find out how much consumed
    for (int i = 0; i < kNumFirms; i++) {
      float __ration_factor = 1.0;

      if (need_to_ration_this_good_arr[i]) {
        // overdemanded
        __ration_factor = available_stock_arr[i] / gross_demand_arr[i];
      }
      
      net_consumed_arr[i] = __ration_factor * net_demand_arr[i];

      atomicAdd_block(&(total_actually_consumer_arr[i]), net_consumed_arr[i]);
    }

    // store amount actually consumed for this consumer
    CopyFloatArraySlice(net_consumed_arr, kNumFirms, my_aux_batch_ptr);

    // ----------------------------------------
    // Compute consumer utility
    // ----------------------------------------
    float hours_worked = my_action_arr[kIdxConsumerWorkedOffset];
    int worked_for_this_firm_id =
        (int)my_action_arr[kIdxConsumerWhichFirmOffset];

    // budget is first elem of consumer state, theta second
    float __theta = my_state_arr[1];

    float __this_consumer_util = 0.0;
    float __total_hours_worked = 0.0;
    float __gross_income = 0.0;

    // Compute expenses
    // Each consumer can consume from each firm, so loop over them.
    for (int i = 0; i < kNumFirms; i++) {
      __this_consumer_util += GetCRRAUtil(net_consumed_arr[i], kCrraParam);
      cs_budget_delta -= prices_arr[i] * net_consumed_arr[i];
    }

    // Compute income
    __total_hours_worked += hours_worked;
    __gross_income += wages_arr[worked_for_this_firm_id] * hours_worked;
    float __income_tax_paid = income_tax_rate * __gross_income;
    cs_budget_delta += (__gross_income - __income_tax_paid);

    // Update tax revenue (government)
    atomicAdd_block(&(tax_revenue_arr[kIdxThisThreadGovtId]),
                    __income_tax_paid);

    // Compute reward
    this_agent_reward +=
        __this_consumer_util - (__theta / 2.0) * (__total_hours_worked);
  }

  __syncthreads();


  // ----------------------------------------
  // Firms Exports: Add external consumption.
  // ----------------------------------------
  if (ThisThreadAgentType == kFirmType ) {
      const int this_thread_firm_id = threadIdx.x - kNumConsumers;
      if (kUseImporter) {
          // sell remaining goods, if any, to importer, if price is high enough.
          float __this_firm_price = prices_arr[this_thread_firm_id];
          float __stock_after_consumers = available_stock_arr[this_thread_firm_id] - total_actually_consumer_arr[this_thread_firm_id];

          if (__this_firm_price >= kImporterPrice) {
              bought_by_importer_arr[this_thread_firm_id] = fmaxf(fminf(__stock_after_consumers, kImporterQuantity), 0.0); // floor to zero to avoid small negative floats
          }
          else {
              bought_by_importer_arr[this_thread_firm_id] = 0.0;
          }
      }
      else {
          bought_by_importer_arr[this_thread_firm_id] = 0.0;
      }
  }

  // ----------------------------------------
  // Firms: Rationing demand + Utility
  // ----------------------------------------
  if (ThisThreadAgentType == kFirmType) {
    const int this_thread_firm_id = threadIdx.x - kNumConsumers;

    float __this_firm_revenue =
        (total_actually_consumer_arr[this_thread_firm_id] + bought_by_importer_arr[this_thread_firm_id]) *
        prices_arr[this_thread_firm_id];
    float __wages_paid =
        hours_worked_arr[this_thread_firm_id] * wages_arr[this_thread_firm_id];

    // Firms can invest in new capital. This increases their production factor
    // (see GetFirmProduction).
    // here, after consumers consume, if price is >= than importer price, importer consumes up to their maximum of the goods, at the importer price

    float __gross_income = __this_firm_revenue - __wages_paid;
    capital_delta = fmaxf(my_action_arr[2] * __gross_income, 0.0);
    float __gross_profit = __gross_income - capital_delta;
    float __corp_tax_paid = corporate_tax_rate * fmaxf(__gross_profit, 0.0);
    fm_budget_delta = (__gross_profit - __corp_tax_paid);
    if (kIncentivizeFirmActivity) {
       if ((fm_budget_delta + my_state_arr[0]) > 0.0) {  // if positive budget
            this_agent_reward += (kFirmBoostRewardFactor * __this_firm_revenue);
       }
    }
    this_agent_reward += (__gross_profit - __corp_tax_paid);

    atomicAdd_block(&(tax_revenue_arr[0]), __corp_tax_paid);

    float __production = GetFirmProduction(0.01, my_state_arr[kIdxFirmCapitalOffset],
                          hours_worked_arr[this_thread_firm_id],  my_state_arr[kIdxFirmAlphaOffset]);

    // -------------------
    // Update global state
    // -------------------
    // update prices in global state
    next_global_state_arr[kIdxPricesOffset + this_thread_firm_id] =
        my_action_arr[0];
    // update wages in global state
    next_global_state_arr[kIdxWagesOffset + this_thread_firm_id] =
        my_action_arr[1];
    // update stocks in global state
    next_global_state_arr[kIdxStockOffset + this_thread_firm_id] =
        available_stock_arr[this_thread_firm_id] -
        total_actually_consumer_arr[this_thread_firm_id] - bought_by_importer_arr[this_thread_firm_id] + __production;

    *my_aux_batch_ptr = bought_by_importer_arr[this_thread_firm_id];

    // update overdemanded in global state
    next_global_state_arr[kIdxOverdemandOffset + this_thread_firm_id] =
        need_to_ration_this_good_arr[this_thread_firm_id] ? 1.0 : 0.0;
  }

  // -----------------
  // Move time forward
  // -----------------
  // Let first firm tick time
  if (ThisThreadAgentType == kFirmType) {
      const int this_thread_firm_id = threadIdx.x - kNumConsumers;
    if (this_thread_firm_id == 0) {
      next_global_state_arr[kIdxTimeOffset] =
          my_global_state_ptr[kIdxTimeOffset] + 1.0;
    }
  }

  __syncthreads();

  // ----------------------------------------
  // Subsidies
  // ----------------------------------------
  // need to redistribute tax revenues
  if (ThisThreadAgentType == kConsumerType) {
    float __redistribution = tax_revenue_arr[0] / kNumConsumers;
    cs_budget_delta += __redistribution;
  }

  // ----------------------------------------
  // Social welfare
  // ----------------------------------------
  // After this point consumers and firms know their final reward, so can inform
  // the government thread via shared memory

  __syncthreads();

  // ----------------------------------------
  // Government sets taxes for the next round
  // ----------------------------------------
  if (ThisThreadAgentType == kGovernmentType) {
    next_global_state_arr[kIdxIncomeTaxOffset] = my_action_arr[0];
    next_global_state_arr[kIdxCorporateTaxOffset] = my_action_arr[1];
  }

  __syncthreads();

  // -----------------------------------------------
  // Copy next_global_state_arr into my global state
  // -----------------------------------------------
  // All agents need to see the updated global state
  CopyFloatArraySlice(next_global_state_arr, kGlobalStateSize,
                      my_global_state_ptr);

  // -----------------------------------------------
  // Update budgets
  // -----------------------------------------------
  // Update budget (same for all agents)
  if (ThisThreadAgentType == kConsumerType) {
    my_state_arr[0] += cs_budget_delta;
  }
  if (ThisThreadAgentType == kFirmType) {
    my_state_arr[0] += fm_budget_delta;
  }

  // Add interest rate on savings
  if ((ThisThreadAgentType == kConsumerType) ||
      (ThisThreadAgentType == kFirmType)) {
    if (my_state_arr[0] > 0.0) {
      my_state_arr[0] += my_state_arr[0] * kInterestRate;
    }
  }

  // Add new capital
  if (ThisThreadAgentType == kFirmType) {
    my_state_arr[kIdxFirmCapitalOffset] += capital_delta;
  }

  // Add new capital
  if ((ThisThreadAgentType == kFirmType) ||
      (ThisThreadAgentType == kGovernmentType)) {
    *batch_action_value_ptr = *my_action_value_ptr;
  }

  // Update rewards in global state
  *my_reward_value_ptr = this_agent_reward;
  *batch_reward_value_ptr = this_agent_reward;
}

// ************************
// End of extern "C" block.
// ************************
}