#include <list>
#include <cmath>
#include <vector>
#include <fstream>
#include <iostream>
#include <cstdlib>
#include <algorithm>

using namespace std;

#define ILOUSESTL

#include <ilcplex/ilocplex.h>
#include <time.h>

ILOSTLBEGIN

// problem parameters
int   number_of_groups;
int   number_of_timeperiods;
int   number_of_states;
int   number_of_bandit_in_total;   // total number of bandits
int   maximum_number_per_group; // (int)(number_of_bandit_in_total * 6 / number_of_groups);

int* size_of_groups;
double** init_prob;
double**** prob;
double*** reward;

double      pulled_ratio;  // the maximum percentage of bandits that are allowed to pulled 
int         maximum_number_of_pulled_bandits; // = (int)(pulled_ratio * number_of_bandit_in_total);

void generate_dist(int dimension, double dist[], double lower_bound) {
    double prob_sum = 0.0;
    for (int i = 0; i < dimension; i++) {
        dist[i] = (double)rand() / (RAND_MAX)+lower_bound;
        prob_sum += dist[i];
    }
    for (int i = 0; i < dimension; i++) dist[i] /= prob_sum;
}

int generate_number_in_dist(int dimension, double dist[]) {
    double random_num = (double)rand() / (RAND_MAX);
    for (int k = 0; k < dimension; k++) {
        if (random_num <= dist[k]) return k;
        else random_num -= dist[k];
    }
    return dimension - 1;
}

int generate_number_uniformly(int dimension) {
    int number = (int)(dimension * ((double)rand() / (RAND_MAX)));
    if (number == dimension) return dimension - 1;
    else return number;
}

void generate_feasible_action(vector <int>& action) {
    std::fill(action.begin(), action.end(), 0);
    for (int i = 0; i < pulled_ratio * number_of_bandit_in_total; i++) {
        action[generate_number_uniformly(number_of_bandit_in_total)] = 1;
    }
}

int sort_trancated(vector <double>& weight, vector <int>& action) {
    int sum_pulled_bandits = 0;
    vector <double> weight_copy = weight;
    std::fill(action.begin(), action.end(), 0);
    std::sort(weight.begin(), weight.end());
    if (weight[number_of_bandit_in_total - maximum_number_of_pulled_bandits - 1] >= 0) {
        for (int i = 0; i < number_of_bandit_in_total; i++) {
            if (weight_copy[i] > weight[number_of_bandit_in_total - maximum_number_of_pulled_bandits - 1]) {
                action[i] = 1;
                sum_pulled_bandits += 1;
            }
        }
        for (int i = 0; i < number_of_bandit_in_total; i++) {
            if (sum_pulled_bandits >= maximum_number_of_pulled_bandits) break;
            if (weight_copy[i] == weight[number_of_bandit_in_total - maximum_number_of_pulled_bandits - 1]) {
                action[i] = 1;
                sum_pulled_bandits += 1;
            }
        }
    }
    else {
        for (int i = 0; i < number_of_bandit_in_total; i++) {
            if (weight_copy[i] > 0) {
                action[i] = 1;
                sum_pulled_bandits += 1;
            }
        }
    }
    return sum_pulled_bandits;
}

int random_trancated(vector <double>& weight, vector <int>& action) {
    int number, number_old = -1, sum_pulled_bandits = 0;
    vector <double> weight_copy = weight;
    std::fill(action.begin(), action.end(), 0);
    int index[number_of_bandit_in_total];
    for (int i = 0; i < number_of_bandit_in_total; i++) {
        if (weight_copy[i] > 0) {
            action[i] = 1;
            index[sum_pulled_bandits] = i;
            sum_pulled_bandits += 1;
        }
    }
    if (sum_pulled_bandits > maximum_number_of_pulled_bandits) {
        int overnumber = sum_pulled_bandits - maximum_number_of_pulled_bandits;
        for (int i = 0; i < overnumber; i++) {
            number = (int)(sum_pulled_bandits / overnumber * (i + (double)rand() / (RAND_MAX)));
            if (number > number_old) {
                action[index[number]] = 0;
                number_old = number;
            }
            else {
                action[index[number + 1]] = 0;
                number_old = number + 1;
            }
        }
        sum_pulled_bandits = 0;
        for (int i = 0; i < number_of_bandit_in_total; i++) {
            if (action[i] == 1) {
                sum_pulled_bandits += 1;
            }
        }
    }
    return sum_pulled_bandits;


}

void initilizaton() {
    // each group has same size
    for (int j = 0; j < number_of_groups; ++j)
        size_of_groups[j] = number_of_bandit_in_total / number_of_groups;

    /*// randomly partition the size for each group
    double size_percentage[number_of_groups];
    generate_dist(number_of_groups, size_percentage, 0.2);
    int size_sum = 0;
    for (int j = 0; j < number_of_groups - 1; ++j) {
        size_of_groups[j] = (int)(number_of_bandit_in_total * size_percentage[j]);
        size_sum += size_of_groups[j];
    }
    size_of_groups[number_of_groups - 1] = number_of_bandit_in_total - size_sum;*/

    // randomly generate the initial probability
    for (int j = 0; j < number_of_groups; ++j)
        generate_dist(number_of_states, init_prob[j], 0);

    // transition probability
    for (int t = 0; t < number_of_timeperiods; ++t) {
        for (int j = 0; j < number_of_groups; ++j) {
            for (int k = 0; k < number_of_states; k++)
                generate_dist(number_of_states, prob[t][j][k], 0);
        }
    }

    // reward function
    for (int t = 0; t < number_of_timeperiods; ++t) {
        for (int j = 0; j < number_of_groups; ++j) {
            for (int k = 0; k < number_of_states; k++)
	        reward[t][j][k] = 2.0 * ((double)rand() / (RAND_MAX)) - 1.0;
        }
    }

}

double fluid_LP(vector< vector< vector<double> > >& sigma_LP, vector< vector< vector<double> > >& pi_LP, ofstream& fout) {

    IloEnv env;
    IloModel model(env);

    typedef IloArray<IloNumVarArray> NumVarMatrix;
    typedef IloArray<NumVarMatrix>   NumVarThreeMatrix;

    NumVarThreeMatrix sigma(env, number_of_timeperiods);
    NumVarThreeMatrix pi(env, number_of_timeperiods);
    for (int t = 0; t < number_of_timeperiods; ++t) {
        sigma[t] = NumVarMatrix(env, number_of_groups);
        pi[t] = NumVarMatrix(env, number_of_groups);
        for (int j = 0; j < number_of_groups; ++j) {
            sigma[t][j] = IloNumVarArray(env, number_of_states, 0, +IloInfinity);
            pi[t][j] = IloNumVarArray(env, number_of_states, 0, +IloInfinity);
        }
    }

    IloExpr obj_func(env);

    for (int t = 0; t < number_of_timeperiods; ++t) {
        for (int j = 0; j < number_of_groups; ++j) {
            for (int k = 0; k < number_of_states; ++k) {
                obj_func += reward[t][j][k] * pi[t][j][k];
            }
        }
    }
    IloObjective objc = IloMaximize(env, obj_func);
    model.add(objc);
    obj_func.end();

    for (int j = 0; j < number_of_groups; j++)
        for (int k = 0; k < number_of_states; k++)
            model.add(sigma[0][j][k] == init_prob[j][k] * size_of_groups[j]);

    for (int t = 0; t < number_of_timeperiods; t++) {
        IloExpr curr_LHS(env);
        for (int j = 0; j < number_of_groups; j++) {
            for (int k = 0; k < number_of_states; k++) {
                model.add(pi[t][j][k] <= sigma[t][j][k]);
                curr_LHS += pi[t][j][k];
            }
        }
        model.add(curr_LHS <= pulled_ratio * number_of_bandit_in_total);
        curr_LHS.end();

        if (t < number_of_timeperiods - 1) {
            for (int j = 0; j < number_of_groups; j++) {
                for (int k = 0; k < number_of_states; k++) {
                    IloExpr curr_RHS(env);
                    curr_RHS += sigma[t][j][k] - pi[t][j][k];
                    for (int l = 0; l < number_of_states; l++) {
                        curr_RHS += pi[t][j][l] * prob[t][j][l][k];
                    }
                    model.add(sigma[t + 1][j][k] == curr_RHS);
                    curr_RHS.end();
                }
            }
        }

    }

    IloCplex cplex(model);
    cplex.setOut(env.getNullStream());
    cplex.setParam(IloCplex::Param::ClockType, 2);
    cplex.setParam(IloCplex::RootAlg, IloCplex::Dual);
    cplex.setParam(IloCplex::Threads, 1);
    cplex.solve();

    double objval = cplex.getObjValue();

    for (int t = 0; t < number_of_timeperiods; t++) {
        for (int j = 0; j < number_of_groups; j++) {
            IloNumArray values_sigma(env), values_pi(env);
            cplex.getValues(values_sigma, sigma[t][j]);
            cplex.getValues(values_pi, pi[t][j]);
            for (int k = 0; k < number_of_states; k++) {
                sigma_LP[t][j][k] = values_sigma[k];
                pi_LP[t][j][k] = values_pi[k];
            }
        }
    }

    env.end();

    fout << objval << endl;

    return objval;
}

double fluid_LP_reduced(vector< vector< vector<double> > >& sigma_LP, vector< vector< vector<double> > >& pi_LP, double phi, ofstream& fout) {

    IloEnv env;
    IloModel model(env);

    typedef IloArray<IloNumVarArray> NumVarMatrix;
    typedef IloArray<NumVarMatrix>   NumVarThreeMatrix;

    NumVarThreeMatrix sigma(env, number_of_timeperiods);
    NumVarThreeMatrix pi(env, number_of_timeperiods);
    for (int t = 0; t < number_of_timeperiods; ++t) {
        sigma[t] = NumVarMatrix(env, number_of_groups);
        pi[t] = NumVarMatrix(env, number_of_groups);
        for (int j = 0; j < number_of_groups; ++j) {
            sigma[t][j] = IloNumVarArray(env, number_of_states, 0, +IloInfinity);
            pi[t][j] = IloNumVarArray(env, number_of_states, 0, +IloInfinity);
        }
    }

    IloExpr obj_func(env);

    for (int t = 0; t < number_of_timeperiods; ++t) {
        for (int j = 0; j < number_of_groups; ++j) {
            for (int k = 0; k < number_of_states; ++k) {
                obj_func += reward[t][j][k] * pi[t][j][k];
            }
        }
    }
    IloObjective objc = IloMaximize(env, obj_func);
    model.add(objc);
    obj_func.end();

    for (int j = 0; j < number_of_groups; j++)
        for (int k = 0; k < number_of_states; k++)
            model.add(sigma[0][j][k] == init_prob[j][k] * size_of_groups[j]);

    double epsilon = 0;
    for (int i = 0; i < number_of_groups; i++) {
        if (epsilon < sqrt(log(size_of_groups[i]) / (2 * size_of_groups[i])))
            epsilon = sqrt(log(size_of_groups[i]) / (2 * size_of_groups[i]));
    }

    //    double maximum_number_of_pulled_bandits_reduced = pulled_ratio * number_of_bandit_in_total - \
    //        phi * epsilon * number_of_bandit_in_total * number_of_groups;
    //    if (maximum_number_of_pulled_bandits_reduced < 0) 
    //        maximum_number_of_pulled_bandits_reduced = 0;

    // ***** CHANGE IMPLEMENTED BY WW *****
    double maximum_number_of_pulled_bandits_reduced = (1.0 - phi) * pulled_ratio * number_of_bandit_in_total;
    // ************************************


    for (int t = 0; t < number_of_timeperiods; t++) {
        IloExpr curr_LHS(env);
        for (int j = 0; j < number_of_groups; j++) {
            for (int k = 0; k < number_of_states; k++) {
                model.add(pi[t][j][k] <= sigma[t][j][k]);
                curr_LHS += pi[t][j][k];
            }
        }
        model.add(curr_LHS <= maximum_number_of_pulled_bandits_reduced);
        curr_LHS.end();

        if (t < number_of_timeperiods - 1) {
            for (int j = 0; j < number_of_groups; j++) {
                for (int k = 0; k < number_of_states; k++) {
                    IloExpr curr_RHS(env);
                    curr_RHS += sigma[t][j][k] - pi[t][j][k];
                    for (int l = 0; l < number_of_states; l++) {
                        curr_RHS += pi[t][j][l] * prob[t][j][l][k];
                    }
                    model.add(sigma[t + 1][j][k] == curr_RHS);
                    curr_RHS.end();
                }
            }
        }

    }

    IloCplex cplex(model);
    cplex.setOut(env.getNullStream());
    cplex.setParam(IloCplex::Param::ClockType, 2);
    cplex.setParam(IloCplex::RootAlg, IloCplex::Dual);
    cplex.solve();

    double objval = cplex.getObjValue();

    for (int t = 0; t < number_of_timeperiods; t++) {
        for (int j = 0; j < number_of_groups; j++) {
            IloNumArray values_sigma(env), values_pi(env);
            cplex.getValues(values_sigma, sigma[t][j]);
            cplex.getValues(values_pi, pi[t][j]);
            for (int k = 0; k < number_of_states; k++) {
                sigma_LP[t][j][k] = values_sigma[k];
                pi_LP[t][j][k] = values_pi[k];
            }
        }
    }

    env.end();

    fout << objval << endl;

    return objval;
}

double lagrangian_LP(vector< vector< vector<double> > >& value_Lagrangian, vector<double>& lambda_LG, ofstream& fout) {

    IloEnv env;
    IloModel model(env);

    typedef IloArray<IloNumVarArray> NumVarMatrix;
    typedef IloArray<NumVarMatrix>   NumVarThreeMatrix;

    NumVarThreeMatrix value(env, number_of_timeperiods + 1);
    for (int t = 0; t < number_of_timeperiods + 1; ++t) {
        value[t] = NumVarMatrix(env, number_of_groups);
        for (int j = 0; j < number_of_groups; ++j) {
            value[t][j] = IloNumVarArray(env, number_of_states, -IloInfinity, IloInfinity);
        }
    }

    IloNumVarArray lambda(env, number_of_timeperiods, 0, +IloInfinity);

    IloExpr obj_func(env);

    for (int t = 0; t < number_of_timeperiods; ++t)
        obj_func += lambda[t] * pulled_ratio * number_of_bandit_in_total;

    for (int j = 0; j < number_of_groups; ++j) {
        for (int k = 0; k < number_of_states; ++k)
            obj_func += value[0][j][k] * init_prob[j][k] * size_of_groups[j];
    }

    IloObjective objc = IloMinimize(env, obj_func);
    model.add(objc);
    obj_func.end();

    for (int j = 0; j < number_of_groups; j++)
        for (int k = 0; k < number_of_states; k++)
            model.add(value[number_of_timeperiods][j][k] == 0.0);

    for (int t = 0; t < number_of_timeperiods; t++) {
        for (int j = 0; j < number_of_groups; j++) {
            for (int k = 0; k < number_of_states; k++) {
                // action: pull the bandit
                IloExpr curr_RHS(env);
                curr_RHS += reward[t][j][k] - lambda[t];
                for (int l = 0; l < number_of_states; l++)
                    curr_RHS += prob[t][j][k][l] * value[t + 1][j][l];
                model.add(value[t][j][k] >= curr_RHS);
                curr_RHS.end();
                // action: do not pull the bandit
                model.add(value[t][j][k] >= value[t + 1][j][k]);
            }
        }

    }

    IloCplex cplex(model);
    cplex.setOut(env.getNullStream());
    cplex.setParam(IloCplex::Param::ClockType, 2);
    cplex.setParam(IloCplex::RootAlg, IloCplex::Dual);
    cplex.setParam(IloCplex::Threads, 1);
    cplex.solve();

    double objval = cplex.getObjValue();

    IloNumArray value_lambda(env);
    cplex.getValues(value_lambda, lambda);
    for (int t = 0; t < number_of_timeperiods; t++) {
        lambda_LG[t] = value_lambda[t];
        for (int j = 0; j < number_of_groups; j++) {
            IloNumArray values_LG(env);
            cplex.getValues(values_LG, value[t][j]);
            for (int k = 0; k < number_of_states; k++) {
                value_Lagrangian[t][j][k] = values_LG[k];
            }
        }
    }

    env.end();

    fout << objval << endl;

    return objval;
}

double approximate_LP(vector< vector< vector<double> > >& value_AproximateLP, int number_of_sample, ofstream& fout) {

    IloEnv env;
    IloModel model(env);

    typedef IloArray<IloNumVarArray> NumVarMatrix;
    typedef IloArray<NumVarMatrix>   NumVarThreeMatrix;

    NumVarThreeMatrix value(env, number_of_timeperiods + 1);
    for (int t = 0; t < number_of_timeperiods + 1; ++t) {
        value[t] = NumVarMatrix(env, number_of_groups);
        for (int j = 0; j < number_of_groups; ++j) {
            value[t][j] = IloNumVarArray(env, number_of_states, 0.0, IloInfinity);
        }
    }

    IloExpr obj_func(env);

    for (int j = 0; j < number_of_groups; ++j)
        for (int k = 0; k < number_of_states; ++k)
            obj_func += value[0][j][k] * init_prob[j][k] * size_of_groups[j];

    IloObjective objc = IloMinimize(env, obj_func);
    model.add(objc);
    obj_func.end();

    for (int j = 0; j < number_of_groups; j++) {
        for (int k = 0; k < number_of_states; k++) {
            model.add(value[number_of_timeperiods][j][k] == 0.0);
        }
    }

    for (int t = 0; t < number_of_timeperiods; t++) {
        for (int j = 0; j < number_of_groups; j++) {
            for (int k = 0; k < number_of_states; k++) {
                model.add(value[t][j][k] >= 0.0);
            }
        }
    }

    int t, k, index;
    vector<int> action(number_of_bandit_in_total, 0);

    double*** coeffs_LHS = new double** [number_of_timeperiods + 1], *** coeffs_RHS = new double** [number_of_timeperiods + 1];
    for (int t = 0; t < number_of_timeperiods + 1; ++t) {
        coeffs_LHS[t] = new double* [number_of_groups]; coeffs_RHS[t] = new double* [number_of_groups];
        for (int j = 0; j < number_of_groups; ++j) {
            coeffs_LHS[t][j] = new double[number_of_states]; coeffs_RHS[t][j] = new double[number_of_states];
        }
    }

    for (int sample_num = 0; sample_num < number_of_sample; sample_num++) {
        for (int t = 0; t < number_of_timeperiods; ++t)
            for (int j = 0; j < number_of_groups; ++j)
                for (int k = 0; k < number_of_states; ++k) {
                    coeffs_LHS[t][j][k] = 0.0;
                    coeffs_RHS[t][j][k] = 0.0;
                }

        IloExpr curr_LHS(env), curr_RHS(env);

        t = generate_number_uniformly(number_of_timeperiods);
        generate_feasible_action(action);
        index = 0;
        for (int j = 0; j < number_of_groups; j++) {
            for (int i = 0; i < size_of_groups[j]; i++) {
                k = generate_number_uniformly(number_of_states);
                coeffs_LHS[t][j][k] += 1.0;
                if (action[index] > 0) {
                    curr_RHS += reward[t][j][k];
                    for (int l = 0; l < number_of_states; l++) {
                        coeffs_RHS[t + 1][j][l] += prob[t][j][k][l];
                    }
                }
                else coeffs_RHS[t + 1][j][k] += 1.0;

                index += 1;
            }
        }

        for (int t = 0; t < number_of_timeperiods; ++t)
            for (int j = 0; j < number_of_groups; ++j)
                for (int k = 0; k < number_of_states; ++k) {
                    curr_LHS += coeffs_LHS[t][j][k] * value[t][j][k];
                    curr_RHS += coeffs_RHS[t][j][k] * value[t][j][k];
                }
        model.add(curr_LHS >= curr_RHS);
        curr_LHS.end();
        curr_RHS.end();
    }


    IloCplex cplex(model);
    cplex.setOut(env.getNullStream());
    cplex.setParam(IloCplex::Param::ClockType, 2);
    cplex.setParam(IloCplex::RootAlg, IloCplex::Dual);
    cplex.setParam(IloCplex::Threads, 1);
    cplex.solve();

    double objval = cplex.getObjValue();

    for (int t = 0; t < number_of_timeperiods; t++) {
        for (int j = 0; j < number_of_groups; j++) {
            IloNumArray values_ALP(env);
            cplex.getValues(values_ALP, value[t][j]);
            for (int k = 0; k < number_of_states; k++) {
                value_AproximateLP[t][j][k] = values_ALP[k];
            }
        }
    }

    env.end();

    fout << objval << endl;

    return objval;
}

double monte_carlo_randomized_policy_infeasible(const vector< vector< vector<double> > >& sigma_LP, \
    const vector < vector< vector<double> > >& pi_LP, ofstream& fout) {

  double result = 0.0;
  
  for (int run = 0; run < 10000; ++run) {
    double total_reward = 0.0;
    int number_of_active[number_of_timeperiods];
    int arm_state[number_of_groups][maximum_number_per_group];

    // initialization of state for each arm based on the initial probability
    for (int j = 0; j < number_of_groups; j++) {
      for (int i = 0; i < size_of_groups[j]; i++)
	arm_state[j][i] = generate_number_in_dist(number_of_states, init_prob[j]);
    }

    // simulations for each time period
    double active_prob;
    int k;
    for (int t = 0; t < number_of_timeperiods; t++) {
      number_of_active[t] = 0;
      for (int j = 0; j < number_of_groups; j++) {
	for (int i = 0; i < size_of_groups[j]; i++) {
	  k = arm_state[j][i];
	  active_prob = pi_LP[t][j][k] / sigma_LP[t][j][k];
	  if ((double)rand() / (RAND_MAX) < active_prob) {
	    total_reward += reward[t][j][k];
	    number_of_active[t] += 1;
	    arm_state[j][i] = generate_number_in_dist(number_of_states, prob[t][j][k]);
	  }
	}
      }
    }

    //cout << "Allowed bandit number: " << (int)(pulled_ratio * number_of_bandit_in_total) << endl;
    //for (int t = 0; t < number_of_timeperiods; t++) {
    //    cout << "number of active bandits: " << number_of_active[t] << endl;
    //    fout << number_of_active[t] << endl;
    //}

    result += total_reward;
  }

  result /= 10000.0;
  fout << result << endl;
  
  return result;
}

double monte_carlo_randomized_policy_feasible(const vector< vector< vector<double> > >& sigma_LP, \
    const vector < vector< vector<double> > >& pi_LP, ofstream& fout) {

  double result = 0.0;

  for (int run = 0; run < 10000; ++run) {  
    double total_reward = 0.0;
    int number_of_active[number_of_timeperiods];
    int arm_state[number_of_groups][maximum_number_per_group];

    // initialization of state for each arm based on the initial probability
    for (int j = 0; j < number_of_groups; j++) {
      for (int i = 0; i < size_of_groups[j]; i++)
	arm_state[j][i] = generate_number_in_dist(number_of_states, init_prob[j]);
    }

    // simulations for each time period
    double active_prob;
    int k, index;
    vector<int> action(number_of_bandit_in_total, 0);
    vector<double> weight(number_of_bandit_in_total, 0);
    for (int t = 0; t < number_of_timeperiods; t++) {
      number_of_active[t] = 0;
      index = 0;
      for (int j = 0; j < number_of_groups; j++) {
	for (int i = 0; i < size_of_groups[j]; i++) {
	  k = arm_state[j][i];
	  active_prob = pi_LP[t][j][k] / sigma_LP[t][j][k];
	  if (((double)rand() / (RAND_MAX)) < active_prob) weight[index] = 1 + active_prob;
	  else weight[index] = active_prob;
	  index++;
	}
      }
      number_of_active[t] = sort_trancated(weight, action);
      index = 0;
      for (int j = 0; j < number_of_groups; j++) {
	for (int i = 0; i < size_of_groups[j]; i++) {
	  k = arm_state[j][i];
	  if (action[index] == 1) {
	    total_reward += reward[t][j][k];
	    arm_state[j][i] = generate_number_in_dist(number_of_states, prob[t][j][k]);
	  }
	  index++;
	}
      }
    }

    //cout << "Allowed bandit number: " << (int)(pulled_ratio * number_of_bandit_in_total) << endl;
    //for (int t = 0; t < number_of_timeperiods; t++) {
    //    cout << "number of active bandits: " << number_of_active[t] << endl;
    //    fout << number_of_active[t] << endl;
    //}

    result += total_reward;
  }

  result /= 10000.0;

  fout << result << endl;

  return result;
}

double monte_carlo_Lagrangian_policy_infeasible(const vector< vector< vector<double> > >& \
    value_Lagrangian, const vector<double>& lambda, ofstream& fout) {

  double result = 0.0;

  for (int run = 0; run < 10000; ++run) {  
    double total_reward = 0.0;
    int number_of_active[number_of_timeperiods];
    int arm_state[number_of_groups][maximum_number_per_group];

    // initialization of state for each arm based on the initial probability
    for (int j = 0; j < number_of_groups; j++) {
      for (int i = 0; i < size_of_groups[j]; i++)
	arm_state[j][i] = generate_number_in_dist(number_of_states, init_prob[j]);
    }

    double active_value;
    int k;
    for (int t = 0; t < number_of_timeperiods; t++) {
      number_of_active[t] = 0;
      for (int j = 0; j < number_of_groups; j++) {
	for (int i = 0; i < size_of_groups[j]; i++) {
	  k = arm_state[j][i];
	  //active_value = reward[t][j][k] - lambda[t];
	  active_value = reward[t][j][k];
	  for (int l = 0; l < number_of_states; l++)
	    active_value += prob[t][j][k][l] * value_Lagrangian[t + 1][j][l];
	  if (active_value > value_Lagrangian[t + 1][j][k]) {
	    total_reward += reward[t][j][k];
	    number_of_active[t] += 1;
	    arm_state[j][i] = generate_number_in_dist(number_of_states, prob[t][j][k]);
	  }
	}
      }
    }

    //cout << "Allowed bandit number: " << (int)(pulled_ratio * number_of_bandit_in_total) << endl;
    //for (int t = 0; t < number_of_timeperiods; t++) {
    //    cout << "number of active bandits: " << number_of_active[t] << endl;
    //    fout << number_of_active[t] << endl;
    //}

    result += total_reward;
  }

  result /= 10000.0;

  fout << result << endl;

  return result;
}

double monte_carlo_Lagrangian_policy_feasible(const vector< vector< vector<double> > >& \
    value_Lagrangian, const vector<double>& lambda, ofstream& fout) {

  double result = 0.0;
  
  for (int run = 0; run < 10000; ++run) {
    double total_reward = 0.0;
    int number_of_active[number_of_timeperiods];
    int arm_state[number_of_groups][maximum_number_per_group];

    // initialization of state for each arm based on the initial probability
    for (int j = 0; j < number_of_groups; j++) {
      for (int i = 0; i < size_of_groups[j]; i++)
	arm_state[j][i] = generate_number_in_dist(number_of_states, init_prob[j]);
    }

    double active_value;
    int k, index;
    vector<int> action(number_of_bandit_in_total, 0);
    vector<double> weight(number_of_bandit_in_total, 0);
    for (int t = 0; t < number_of_timeperiods; t++) {
      number_of_active[t] = 0;
      index = 0;
      for (int j = 0; j < number_of_groups; j++) {
	for (int i = 0; i < size_of_groups[j]; i++) {
	  k = arm_state[j][i];
	  //active_value = reward[t][j][k] - lambda[t];
	  active_value = reward[t][j][k];
	  for (int l = 0; l < number_of_states; l++)
	    active_value += prob[t][j][k][l] * value_Lagrangian[t + 1][j][l];
	  weight[index] = active_value - value_Lagrangian[t + 1][j][k];
	  index++;
	}
      }
      number_of_active[t] = sort_trancated(weight, action);
      index = 0;
      for (int j = 0; j < number_of_groups; j++) {
	for (int i = 0; i < size_of_groups[j]; i++) {
	  k = arm_state[j][i];
	  if (action[index] == 1) {
	    total_reward += reward[t][j][k];
	    arm_state[j][i] = generate_number_in_dist(number_of_states, prob[t][j][k]);
	  }
	  index++;
	}
      }
    }

    //cout << "Allowed bandit number: " << (int)(pulled_ratio * number_of_bandit_in_total) << endl;
    //for (int t = 0; t < number_of_timeperiods; t++) {
    //    cout << "number of active bandits: " << number_of_active[t] << endl;
    //    fout << number_of_active[t] << endl;
    //}

    result += total_reward;
  }

  result /= 10000.0;

  fout << result << endl;

  return result;
}

double monte_carlo_Lagrangian_policy_feasible_BrownSmith(const vector< vector< vector<double> > >& \
    value_Lagrangian, const vector<double>& lambda, const vector< vector< vector<double> > >& sigma_LP, \
    const vector < vector< vector<double> > >& pi_LP, ofstream& fout) {

  double result = 0.0;

  for (int run = 0; run < 10000; ++run) {
    double total_reward = 0.0;
    int number_of_active[number_of_timeperiods];
    int arm_state[number_of_groups][maximum_number_per_group];

    // initialization of state for each arm based on the initial probability
    for (int j = 0; j < number_of_groups; j++) {
      for (int i = 0; i < size_of_groups[j]; i++)
	arm_state[j][i] = generate_number_in_dist(number_of_states, init_prob[j]);
    }

    double active_value, epsilon = 1e-5;
    int k, index;
    vector<int> action(number_of_bandit_in_total, 0);
    vector<double> weight(number_of_bandit_in_total, 0);
    for (int t = 0; t < number_of_timeperiods; t++) {
      number_of_active[t] = 0;
      index = 0;
      for (int j = 0; j < number_of_groups; j++) {
	for (int i = 0; i < size_of_groups[j]; i++) {
	  k = arm_state[j][i];
        active_value = reward[t][j][k];
                        if ( ((double)rand() / (RAND_MAX)) < pi_LP[t][j][k] / sigma_LP[t][j][k])
                            active_value += epsilon;
	  for (int l = 0; l < number_of_states; l++)
	    active_value += prob[t][j][k][l] * value_Lagrangian[t + 1][j][l];
	  weight[index] = active_value - value_Lagrangian[t + 1][j][k];
	  index++;
	}
      }
      number_of_active[t] = sort_trancated(weight, action);
      index = 0;
      for (int j = 0; j < number_of_groups; j++) {
	for (int i = 0; i < size_of_groups[j]; i++) {
	  k = arm_state[j][i];
	  if (action[index] == 1) {
	    total_reward += reward[t][j][k];
	    arm_state[j][i] = generate_number_in_dist(number_of_states, prob[t][j][k]);
	  }
	  index++;
	}
      }
    }

    //cout << "Allowed bandit number: " << (int)(pulled_ratio * number_of_bandit_in_total) << endl;
    //for (int t = 0; t < number_of_timeperiods; t++) {
    //    cout << "number of active bandits: " << number_of_active[t] << endl;
    //    fout << number_of_active[t] << endl;
    //}
    
    result += total_reward;
  }

  result /= 10000.0;

  fout << result << endl;

  return result;
}

double monte_carlo_ApproximateLP_policy_infeasible(const vector< vector< vector<double> > >& \
    value_ApproximateLP, ofstream& fout) {

  double result = 0.0;

  for (int run = 0; run < 10000; ++run) {
    double total_reward = 0.0;
    int number_of_active[number_of_timeperiods];
    int arm_state[number_of_groups][maximum_number_per_group];

    // initialization of state for each arm based on the initial probability
    for (int j = 0; j < number_of_groups; j++) {
      for (int i = 0; i < size_of_groups[j]; i++)
	arm_state[j][i] = generate_number_in_dist(number_of_states, init_prob[j]);
    }

    // simulations for each time period
    double active_value;
    int k, index;

    for (int t = 0; t < number_of_timeperiods; t++) {
      number_of_active[t] = 0;
      for (int j = 0; j < number_of_groups; j++) {
	for (int i = 0; i < size_of_groups[j]; i++) {
	  k = arm_state[j][i];
	  active_value = reward[t][j][k];
	  for (int l = 0; l < number_of_states; l++)
	    active_value += prob[t][j][k][l] * value_ApproximateLP[t + 1][j][l];
	  if (active_value > value_ApproximateLP[t + 1][j][k]) {
	    total_reward += reward[t][j][k];
	    number_of_active[t] += 1;
	    arm_state[j][i] = generate_number_in_dist(number_of_states, prob[t][j][k]);
	  }
	}
      }
    }

    //cout << "Allowed bandit number: " << (int)(pulled_ratio * number_of_bandit_in_total) << endl;
    //for (int t = 0; t < number_of_timeperiods; t++) {
    //    cout << "number of active bandits: " << number_of_active[t] << endl;
    //    fout << number_of_active[t] << endl;
    //}

    result += total_reward;
  }

  result /= 10000.0;

  fout << result << endl;

  return result;
}

double monte_carlo_ApproximateLP_policy_feasible(const vector< vector< vector<double> > >& \
    value_ApproximateLP, ofstream& fout) {

  double result = 0.0;

  for (int run = 0; run < 10000; ++run) {  
    double total_reward = 0.0;
    int number_of_active[number_of_timeperiods];
    int arm_state[number_of_groups][maximum_number_per_group];

    // initialization of state for each arm based on the initial probability
    for (int j = 0; j < number_of_groups; j++) {
      for (int i = 0; i < size_of_groups[j]; i++)
	arm_state[j][i] = generate_number_in_dist(number_of_states, init_prob[j]);
    }

    // simulations for each time period
    double active_value;
    int k, index;

    vector<int> action(number_of_bandit_in_total, 0);
    vector<double> weight(number_of_bandit_in_total, 0);
    for (int t = 0; t < number_of_timeperiods; t++) {
      number_of_active[t] = 0;
      index = 0;
      for (int j = 0; j < number_of_groups; j++) {
	for (int i = 0; i < size_of_groups[j]; i++) {
	  k = arm_state[j][i];
	  active_value = reward[t][j][k];
	  for (int l = 0; l < number_of_states; l++)
	    active_value += prob[t][j][k][l] * value_ApproximateLP[t + 1][j][l];
	  weight[index] = active_value - value_ApproximateLP[t + 1][j][k];
	  index++;
	}
      }
      number_of_active[t] = sort_trancated(weight, action);
      index = 0;
      for (int j = 0; j < number_of_groups; j++) {
	for (int i = 0; i < size_of_groups[j]; i++) {
	  k = arm_state[j][i];
	  active_value = reward[t][j][k];
	  if (action[index] == 1) {
	    total_reward += reward[t][j][k];
	    arm_state[j][i] = generate_number_in_dist(number_of_states, prob[t][j][k]);
	  }
	  index++;
	}
      }
    }

    //cout << "Allowed bandit number: " << (int)(pulled_ratio * number_of_bandit_in_total) << endl;
    //for (int t = 0; t < number_of_timeperiods; t++) {
    //    cout << "number of active bandits: " << number_of_active[t] << endl;
    //    fout << number_of_active[t] << endl;
    //}

    result += total_reward;
  }

  result /= 10000.0;

  fout << result << endl;

  return result;
}

double monte_carlo_uniform_policy_infeasible(ofstream& fout) {
  double result = 0.0;

  for (int run = 0; run < 10000; ++run) {  
    double total_reward = 0.0;
    int number_of_active[number_of_timeperiods];
    int arm_state[number_of_groups][maximum_number_per_group];

    // initialization of state for each arm based on the initial probability
    for (int j = 0; j < number_of_groups; j++) {
      for (int i = 0; i < size_of_groups[j]; i++)
	arm_state[j][i] = generate_number_in_dist(number_of_states, init_prob[j]);
    }

    // simulations for each time period
    int k;
    for (int t = 0; t < number_of_timeperiods; t++) {
      number_of_active[t] = 0;
      for (int j = 0; j < number_of_groups; j++) {
	for (int i = 0; i < size_of_groups[j]; i++) {
	  k = arm_state[j][i];
	  if ((double)rand() / (RAND_MAX) < pulled_ratio) {
	    total_reward += reward[t][j][k];
	    number_of_active[t] += 1;
	    arm_state[j][i] = generate_number_in_dist(number_of_states, prob[t][j][k]);
	  }
	}
      }
    }

    //cout << "Allowed bandit number: " << (int)(pulled_ratio * number_of_bandit_in_total) << endl;
    //for (int t = 0; t < number_of_timeperiods; t++) {
    //    cout << "number of active bandits: " << number_of_active[t] << endl;
    //    fout << number_of_active[t] << endl;
    //}

    result += total_reward;
  }

  result /= 10000.0;

  fout << result << endl;

  return result;
}

double monte_carlo_uniform_policy_feasible(ofstream& fout) {
  double result = 0.0;
  
  for (int run = 0; run < 10000; ++run) {
    double total_reward = 0.0;
    int number_of_active[number_of_timeperiods];
    int arm_state[number_of_groups][maximum_number_per_group];

    // initialization of state for each arm based on the initial probability
    for (int j = 0; j < number_of_groups; j++) {
      for (int i = 0; i < size_of_groups[j]; i++)
	arm_state[j][i] = generate_number_in_dist(number_of_states, init_prob[j]);
    }

    // simulations for each time period
    int k, index;
    vector<int> action(number_of_bandit_in_total, 0);
    vector<double> weight(number_of_bandit_in_total, 0);
    for (int t = 0; t < number_of_timeperiods; t++) {
      number_of_active[t] = 0;
      index = 0;
      for (int j = 0; j < number_of_groups; j++) {
	for (int i = 0; i < size_of_groups[j]; i++) {
	  if (((double)rand() / (RAND_MAX)) < pulled_ratio) weight[index] = 1;
	  else weight[index] = 0;
	  index++;
	}
      }
      number_of_active[t] = random_trancated(weight, action);
      index = 0;
      for (int j = 0; j < number_of_groups; j++) {
	for (int i = 0; i < size_of_groups[j]; i++) {
	  k = arm_state[j][i];
	  if (action[index] == 1) {
	    total_reward += reward[t][j][k];
	    arm_state[j][i] = generate_number_in_dist(number_of_states, prob[t][j][k]);
	  }
	  index++;
	}
      }
    }

    //cout << "Allowed bandit number: " << (int)(pulled_ratio * number_of_bandit_in_total) << endl;
    //for (int t = 0; t < number_of_timeperiods; t++) {
    //    cout << "number of active bandits: " << number_of_active[t] << endl;
    //    fout << number_of_active[t] << endl;
    //}

    result += total_reward;
  }

  result /= 10000.0;

  fout << result << endl;

  return result;
}

int main (int argc, const char* argv[], const char* env[]) {
    // get job number
    int job_number = -1;
    int repetition_number = -1;

    for (int i = 0; env[i] != NULL; ++i) {
        string str = env[i];
        if (str.find("PBS_ARRAY_INDEX") != string::npos) {
            job_number = atoi(str.substr(str.find_last_of("=") + 1).c_str());
            break;
        }
    }
    cout << "job number: " << job_number << endl;
    if (job_number == -1) exit(-1);

    // extract parameter values
    /*int counter = 0;
    for (int index_number_of_bandits = 0; index_number_of_bandits < 3; ++index_number_of_bandits)
        for (int index_number_of_states = 0; index_number_of_states < 4; ++index_number_of_states)
            for (int index_number_of_time_stages = 0; index_number_of_time_stages < 3; ++index_number_of_time_stages)
                for (int index_percentage_pulled = 0; index_percentage_pulled < 4; ++index_percentage_pulled)
                    for (int index_repetition_number = 0; index_repetition_number < 100; ++index_repetition_number) {
                        if (job_number == counter) {
                            number_of_groups = 10;

                            switch (index_number_of_time_stages) {
                            case 0: number_of_timeperiods = 5; break;
                            case 1: number_of_timeperiods = 25; break;
                            case 2: number_of_timeperiods = 50; break;
                            default: exit(-1); break;
                            }

                            switch (index_number_of_states) {
                            case 0: number_of_states = 2; break;
                            case 1: number_of_states = 3; break;
                            case 2: number_of_states = 4; break;
                            case 3: number_of_states = 5; break;
                            default: exit(-1); break;
                            }

                            switch (index_number_of_bandits) {
                            case 0: number_of_bandit_in_total = number_of_groups * 1000; maximum_number_per_group = 10000; break;
                            case 1: number_of_bandit_in_total = number_of_groups * 10000; maximum_number_per_group = 100000; break;
                            case 2: number_of_bandit_in_total = number_of_groups * 100000; maximum_number_per_group = 1000000; break;
                            default: exit(-1); break;
                            }

                            switch (index_percentage_pulled) {
                            case 0: pulled_ratio = 0.01; break;
                            case 1: pulled_ratio = 0.05; break;
                            case 2: pulled_ratio = 0.10; break;
                            case 3: pulled_ratio = 0.25; break;
                            default: exit(-1); break;
                            }
                            maximum_number_of_pulled_bandits = (int)(pulled_ratio * number_of_bandit_in_total);

                            repetition_number = index_repetition_number;

                            goto done;
                        }
                        ++counter;
                    }
		    done:*/

    // extract parameter values
    ifstream fin ("/rds/general/user/wwiesema/home/HKL/list.txt");
    assert (fin.is_open());
    
    if (job_number > 0)
      for (int i = 0; i < job_number; ++i) {
	string str;
	getline (fin, str);
      }

    number_of_groups = 10;
    fin >> number_of_timeperiods;
    fin >> number_of_states;
    fin >> number_of_bandit_in_total;
    maximum_number_per_group = number_of_bandit_in_total / number_of_groups;
    fin >> pulled_ratio;
    pulled_ratio /= 100.0;
    maximum_number_of_pulled_bandits = (int)(pulled_ratio * number_of_bandit_in_total);
    fin >> repetition_number;
 
    // allocate space for data structures
    size_of_groups = new int[number_of_groups];

    // init_prob = new double[number_of_groups][number_of_states];
    init_prob = new double* [number_of_groups];
    for (int i = 0; i < number_of_groups; ++i)
        init_prob[i] = new double[number_of_states];

    // prob = new double[number_of_timeperiods][number_of_groups][number_of_states][number_of_states];
    prob = new double*** [number_of_timeperiods];
    for (int i = 0; i < number_of_timeperiods; ++i) {
        prob[i] = new double** [number_of_groups];
        for (int j = 0; j < number_of_groups; ++j) {
            prob[i][j] = new double* [number_of_states];
            for (int k = 0; k < number_of_states; ++k)
                prob[i][j][k] = new double[number_of_states];
        }
    }

    // reward = new double[number_of_timeperiods][number_of_groups][number_of_states];
    reward = new double** [number_of_timeperiods];
    for (int i = 0; i < number_of_timeperiods; ++i) {
        reward[i] = new double* [number_of_groups];
        for (int j = 0; j < number_of_groups; ++j)
            reward[i][j] = new double[number_of_states];
    }

    // print parameter values
    cout << "number of groups:           " << number_of_groups << endl;
    cout << "number of time periods:     " << number_of_timeperiods << endl;
    cout << "number of states:           " << number_of_states << endl;
    cout << "number of bandits in total: " << number_of_bandit_in_total << endl;
    cout << "pulled ratio:               " << pulled_ratio << endl;
    cout << endl;
    cout << "repetition number:          " << repetition_number << endl;

    // initialization
    srand(job_number);
    initilizaton();

    // create output file
    stringstream output_name;
    output_name << "/rds/general/user/wwiesema/home/HKL/" << number_of_timeperiods << "_" << number_of_states << "_" << number_of_bandit_in_total << "_" << (int)(100.0 * pulled_ratio) << "_" << repetition_number << ".txt";
    ofstream fout(output_name.str().c_str());

    vector < vector< vector<double> > > sigma_LP(number_of_timeperiods, vector< vector<double> > (number_of_groups, vector<double>(number_of_states, 0)));
    vector < vector< vector<double> > > pi_LP(number_of_timeperiods, vector< vector<double> > (number_of_groups, vector<double>(number_of_states, 0)));

    double total_reward = fluid_LP(sigma_LP, pi_LP, fout);
    cout << "objective value of LP: " << total_reward << endl << endl;

    cout << endl << "****************Lagrangian LP***************************" << endl;
    vector < vector< vector<double> > > value_Lagrangian(number_of_timeperiods + 1, vector< vector<double> > \
        (number_of_groups, vector<double>(number_of_states, 0)));
    vector<double> lambda(number_of_timeperiods, 0);
    total_reward = lagrangian_LP(value_Lagrangian, lambda, fout);
    cout << endl << "objective value of Lagrangian LP: " << total_reward << endl << endl;
    
    total_reward = monte_carlo_Lagrangian_policy_feasible_BrownSmith(value_Lagrangian, lambda, sigma_LP, pi_LP, fout);
    cout << endl << "total reward of Brown&Smith feasible Lagrangian policy: " << total_reward << endl << endl;


    return 0;
}
