//------------------------------------------------------------------------------

lfunction ComputeOnGrid (lf_id, grid, handler, callback) {

    jobs = mpi.PartitionIntoBlocks(grid);

    scores = {};

    queue  = mpi.CreateQueue ({^"terms.mpi.LikelihoodFunctions": {{lf_id}},
                               ^"terms.mpi.Headers" : utility.GetListOfLoadedModules ("libv3/")});

    for (i = 1; i < Abs (jobs); i += 1) {
        mpi.QueueJob (queue, handler, {"0" : lf_id,
                                       "1" : jobs [i],
                                       "2" : &scores}, callback);
    }

    Call (callback, -1, Call (handler, lf_id, jobs[0], &scores), {"0" : lf_id, "1" : jobs [0], "2" : &scores});

    mpi.QueueComplete (queue);

    return scores;

}

//------------------------------------------------------------------------------


lfunction pass1.result_handler (node, result, arguments) {
    utility.Extend (^(arguments[2]), result);
}

//------------------------------------------------------------------------------

lfunction pass1.evaluator (lf_id, tasks, scores) {
    LFCompute (^lf_id, LF_START_COMPUTE);

    results = {};
    task_ids = utility.Keys (tasks);
    task_count = Abs (tasks);
    for (i = 0; i < task_count; i+=1) {
        
        parameters.SetValues (tasks[task_ids[i]]);
        LFCompute (^lf_id, ll);
        results [task_ids[i]] = ll;
        
    }

     LFCompute (^lf_id, LF_DONE_COMPUTE);

    return results;
}

//------------------------------------------------------------------------------

lfunction pass2.evaluator (lf_id, tasks, scores) {

    results = {};
    task_ids = utility.Keys (tasks);
    task_count = Abs (tasks);
    for (i = 0; i < task_count; i+=1) {
        parameters.SetValues (tasks[task_ids[i]]);
        ConstructCategoryMatrix(site_likelihoods,^lf_id,SITE_LOG_LIKELIHOODS);
        /*
        if (( (tasks[task_ids[i]]) ["FADE bias"])["MLE"] == 0.0) {
            console.log (tasks[task_ids[i]]);
            console.log (site_likelihoods);
        }
        */
        results [task_ids[i]] = site_likelihoods ["Max(_MATRIX_ELEMENT_VALUE_,-1e200)"];

        // to avoid returning -inf
    }


    return results;
}

//----------------------------------------------------------------------------

lfunction LogDrichletDensity (dir_weights, alpha) {
     if (Min(dir_weights, 0) <= 1e-10) {
        return -1e10;
     }
     if (alpha == 1) {
        return 0;
     }
     dim = Columns (dir_weights);
     return  (+dir_weights["Log(_MATRIX_ELEMENT_VALUE_)*(alpha-1)"]+LnGamma(alpha*dim)-dim*LnGamma(alpha));
}


//------------------------------------------------------------------------------

lfunction ExecuteMCMC (grid, conditionals, ignored, settings, arguments){

    points             = Rows(grid); // # of grid points (settings["chain-length"])
    sites              = Columns(conditionals["conditionals"]); // settings["chain-length"] # of sites
    normalize_by_site  = ({1,points}["1"])*(conditionals["conditionals"]);  // site log L sums
    normalized_weights = (conditionals["conditionals"])*({sites,sites}["1/normalize_by_site[_MATRIX_ELEMENT_ROW_]*(_MATRIX_ELEMENT_ROW_==_MATRIX_ELEMENT_COLUMN_)"]);
                         // rescale each site so that its conditional log L add up to 1
    sum_by_site        = normalized_weights * ({sites,1}["1"]);
                         // settings["chain-length"] initial "loading" of each grid point


    chain_block        = {};

    chain_names        = utility.Keys (arguments);

    for (chain_id = 0; chain_id < Abs (arguments); chain_id += 1) {
        if (mpi.IsMasterNode ()) {
            if (mpi.NodeCount () > 1) {
                io.ReportProgressMessageMD                  ("fubar", "mcmc", "* Running MCMC chain " + (1+chain_names[chain_id]) + " on the master node; other chains will be run on compute MPI nodes without visible feedback");

            } else {
                io.ReportProgressMessageMD                  ("fubar", "mcmc", "* Running MCMC chain " + (1+chain_names[chain_id]));
            }
        }
         weights = {1,points}["Random(sum_by_site[_MATRIX_ELEMENT_COLUMN_]*0.8,sum_by_site[_MATRIX_ELEMENT_COLUMN_]*1.2)"];
        // jiggle initial loading by 20% in each direction, randomly
        weights = weights * (1/(+weights));
            // normalize initial grid point loading


        grid_sampled = {points, 3};
        for (k = 0; k < points; k+=1) {
            grid_sampled[k][0] = grid[k][0];
            grid_sampled[k][1] = grid[k][1];
            grid_sampled[k][2] = weights[k];
        }


        default_step = Max(Min(0.001,1/sites),(weights%0)[points*50$100]);
        // default MCMC move is either the 25-th percentile loading or the smaller of 1/sites or 0.001

        current_site_likelihoods                = weights*(conditionals["conditionals"]);
        current_site_log_sum                    = +(current_site_likelihoods["Log(_MATRIX_ELEMENT_VALUE_)"]);
        current_log_likelihood = {1,2};
        current_log_likelihood [0] = +(((weights *(conditionals["conditionals"]))["Log(_MATRIX_ELEMENT_VALUE_)"])+conditionals["scalers"]);
        current_log_likelihood [1] = LogDrichletDensity (weights, settings["concentration"]);


        individual_grid_point_contribs = {};

        for (k = 0; k < points; k+=1) {
            individual_grid_point_contribs [k] = (conditionals["conditionals"])[k][-1];
        }


        contracting            = settings["chain-length"]*50;
        sample                 = (settings["chain-length"]-settings["burn-in"])$settings["samples"];
        sample_count           = settings["samples"];
        sampled_weights        = {sample_count,points};
        sampled_likelihoods    = {1,sample_count};

        time0                = Time(1);
        sample_index         = 0;

        baseline_step         = default_step;
        reduction_factor      = 1;
        accepted_steps        = 0;

        total_step_sum        = 0; // keeps track of the total change in point weights
        mean_samples_log_likelihood      = 0;
        initialLogL         = current_log_likelihood[0];

        for (steps = 0; steps < settings["chain-length"]; steps += 1) {

            // select two points to switch
            idx    = Random (0, points-1e-10)$1;
            idx2   = Random (0, points-1e-10)$1;
            while (idx == idx2) {
                idx2   = Random (0, points-1e-10)$1;
            }

            // adjust proposal step to keep acceptance rate OK
            if ((steps+1) % contracting == 0) {
                acc_rate = accepted_steps/steps;
                if (acc_rate < 0.25) {
                    baseline_step = baseline_step/1.6;
                } else if (acc_rate > 0.5) {
                    baseline_step = baseline_step*1.6;
                }
           }

            change = Random (0,baseline_step);
            total_step_sum += change;

            if (weights[idx] > change) {
                diffVector          = (individual_grid_point_contribs[idx2]-individual_grid_point_contribs[idx])*change;
                logLDiff            = +((current_site_likelihoods + diffVector)["Log(_MATRIX_ELEMENT_VALUE_)"]) - current_site_log_sum;
                diffPrior           = (settings["concentration"]-1)*(Log((weights[idx]-change)/weights[idx])+Log((weights[idx2]+change)/weights[idx2]));
                costOfMove          = logLDiff+diffPrior;

                if (Random (0,1) <= Exp (costOfMove)) {

                    current_log_likelihood[0] += logLDiff;
                    current_log_likelihood[1] += diffPrior;

                    current_site_likelihoods += diffVector;
                    current_site_log_sum += logLDiff;

                    weights[idx]  += (-change);
                    weights[idx2] += (+change);
                    accepted_steps += 1;
                }
            }

            if (steps > settings["burn-in"]) {
                if ((steps - settings["burn-in"] + 1) % sample == 0) {
                    for (dd = 0; dd < points; dd += 1) {
                        sampled_weights[sample_index][dd] = weights[dd];
                    }
                    sampled_likelihoods[sample_index] = current_log_likelihood[0];
                    mean_samples_log_likelihood += current_log_likelihood[0];
                    sample_index += 1;
                    logLString = Format (mean_samples_log_likelihood/sample_index, 12, 4);
                }
            } else {
                logLString = "(burning in)";
            }

            if ((1+steps) % sample == 0) {
                    io.ReportProgressBar ("MCMC", "Running MCMC chain ID "+ (1+chain_names[chain_id]) + ". Current step: " + (1+steps) + "/" + settings["chain-length"] + ". Mean sampled log(L) = " + logLString
                    + ". Acceptance rate = " + Format (accepted_steps/steps, 6, 3));
            }

        }

        io.ClearProgressBar ();

        chain_block[chain_names[chain_id]] = {
                                                "likelihoods" : sampled_likelihoods,
                                                "weights"     : sampled_weights
                                            };

    }



    return chain_block;
}

//----------------------------------------------------------------------------

lfunction ComputeRandNeff (samples) {

    chain_count  = Abs (samples);
    within_var   = {chain_count,1};
    within_means = {chain_count, 1};


    chain_length = utility.Array1D (samples[0]);

    for (id = 0;  id < chain_count; id += 1) {
        chain_mean = (+(samples[id]))/chain_length;
        chain_var  = +((samples[id])["(_MATRIX_ELEMENT_VALUE_-chain_mean)^2"]);
        within_var   [id] = chain_var/(chain_length-1);
        within_means [id] = chain_mean;
    }

    overall_mean = (+within_means)/chain_count;

    if (chain_count > 1) {
        B           = (+within_means["(_MATRIX_ELEMENT_VALUE_-overall_mean)^2"])*chain_length / (chain_count-1);
        W           = (+within_var)/chain_count;
        VarEst      = (chain_length-1)/chain_length *W + B/chain_length;
        return      {{overall_mean, Sqrt(VarEst/W), VarEst/B*chain_count*chain_length, B, W, VarEst}};
    } else {
        return      {{overall_mean, 0, 0, 0, 0, 0}};
    }
}

//----------------------------------------------------------------------------


lfunction     RunCollapsedGibbs (settings, grid, conditionals, methodname) {

    if (None != methodname) {
        io.ReportProgressMessageMD                  (methodname, "mcmc", "Running a collapsed Gibbs sampler to sample the posterior of rate weights");
        io.ReportProgressMessageMD                  (methodname, "mcmc", "* Using the following settings");
        io.ReportProgressMessageMD                  (methodname, "mcmc", "\t* Steps            : " + settings["chain-length"]);
        io.ReportProgressMessageMD                  (methodname, "mcmc", "\t* Burn-in steps    : " + settings["burn-in"]);
        io.ReportProgressMessageMD                  (methodname, "mcmc", "\t* Samples.         : " + settings["samples"]);
        io.ReportProgressMessageMD                  (methodname, "mcmc", "\t* Dirichlet alpha  : " + settings["concentration"]);
    }


    points             = Rows(grid); // # of grid points
    sites              = Columns(conditionals["conditionals"]);  // GRID_POINTS / SITES
    unit_row           = {1,points} ["1"];
    unit_column        = {sites,1} ["1"];


    current_sampling_weights    = {points, sites};  // each COLUMN vector is a 1 sum vector of probabilities for drawing category J for this site
                                                    // initially just uninformed prior weighted by likelihoods
                                                    // "conditionals" are already normalized to sum to 1 for each site


    grid_weights    = conditionals["conditionals"] * unit_column;

    tolerance                   = 1e-8; // convergence tolerance
    step                        = 0;
    max_steps                   = 100000;
    diagonal_points             = {points, points};
    diagonal_sites              = {sites, sites};
    chain_sample                = {};

    sample                 = (settings["chain-length"]-settings["burn-in"])$settings["samples"];
    sample_count           = settings["samples"];
    sampled_weights        = {sample_count,points};
    current_sample         = random.dirichlet ({points, 1} ['settings["concentration"]']);
    sample_index           = 0;
    current_sampling_weights_norm   = {sites, 1};

    do {
        step += 1;


        current_sampling_weights = current_sample $ conditionals["conditionals"];
            // [rates x sites] element (i,j) => Prob (site j | category i) x Prob (category i), not normalized
        current_sampling_weights_norm = unit_row * current_sampling_weights;
            // normalization weights for each _site_
        current_sample           = random.dirichlet ((current_sampling_weights / current_sampling_weights_norm) * unit_column + (settings["concentration"]));


        if (step > settings["burn-in"]) {
            if ((step - settings["burn-in"] + 1) % sample == 0) {
                for (dd = 0; dd < points; dd += 1) {
                    sampled_weights[sample_index][dd] = current_sample[dd];
                }
                sample_index += 1;
                logLString = " / " + sample_index + " samples drawn";
            }
        } else {
            logLString = " (burning in)";
        }

        if ((step) % (Max (1000, 5*sample)) == 0) {
             io.ReportProgressBar ("mcmc", "Grid weight sampling step " + Format (step, 10, 0) + logLString);
        }


    } while (step < settings["chain-length"]);

    io.ClearProgressBar                   ();

    return {"0" : {"weights" : sampled_weights}};

}

//----------------------------------------------------------------------------


lfunction     RunVariationalBayes (settings, grid, conditionals, methodname) {

    if (None != methodname) {
        io.ReportProgressMessageMD                  (methodname, "mcmc", "Running an iterative zeroth order variational Bayes procedure to estimate the posterior mean of rate weights");
        io.ReportProgressMessageMD                  (methodname, "mcmc", "* Using the following settings");
        io.ReportProgressMessageMD                  (methodname, "mcmc", "\t* Dirichlet alpha  : " + settings["concentration"]);
    }

    points             = Rows(grid); // # of grid points
    sites              = Columns(conditionals["conditionals"]);  // GRID_POINTS / SITES
    unit_row           = {1,points} ["1"];
    unit_column        = {sites,1} ["1"];

    current_sampling_weights    = {points, sites};  // each COLUMN vector is a 1 sum vector of probabilities for drawing category J for this site
                                                    // initially just uninformed prior weighted by likelihoods
                                                    // "conditionals" are already normalized to sum to 1 for each site


    grid_weights    = conditionals["conditionals"] * unit_column;

    tolerance                   = 1e-8; // convergence tolerance
    step                        = 0;
    max_steps                   = 100000;
    diagonal_points             = {points, points};
    diagonal_sites              = {sites, sites};
    current_sampling_weights_norm = {sites, 1};
    for (k = 0; k < points; k += 1) {
        diagonal_points[k][k] := last_grid_weights [k__];
    }
    for (k = 0; k < sites; k += 1) {
        diagonal_sites[k][k] := 1/current_sampling_weights_norm [k__];
    }

    do {
        last_grid_weights = grid_weights;

        step         += 1;
        current_sampling_weights = last_grid_weights $ conditionals["conditionals"];
            // [rates x sites] element (i,j) => Prob (site j | category i) x Prob (category i), not normalized

        current_sampling_weights_norm = unit_row * current_sampling_weights;

            // normalization weights for each _site_
        grid_weights           = (current_sampling_weights / current_sampling_weights_norm) * unit_column + settings["concentration"];
        grid_weights           = grid_weights * (1/ (+grid_weights));

        current_error = Abs (grid_weights-last_grid_weights);

        io.ReportProgressBar ("mcmc", "Posterior mean estimation step " + Format (step, 10, 0) + ", error = " + Format (current_error, 10, 7));
    } while (current_error > tolerance && step < max_steps);

    io.ClearProgressBar                   ();

    return grid_weights;

}

//----------------------------------------------------------------------------

lfunction   ConvertToConditionals (raw_scores) {
    total_grid_points         = Abs (raw_scores);
    site_count                = utility.Array1D (raw_scores[0]);

    conditionals.matrix       = {total_grid_points, site_count};
    conditionals.scalers      = {1,site_count};

    utility.ForEachPair (raw_scores, "_point_", "_conditionals_",
        '
            `&index` = 0 + _point_;
            for (_r = 0; _r < `&site_count`; _r += 1) {
                `&conditionals.matrix` [`&index`][_r] = _conditionals_[_r];
            }
        '
    );

    for (s = 0; s < site_count; s += 1) {
        this_site = conditionals.matrix[-1][s];


        best_log_l = Min (this_site*(-1), 0);
        this_site = (this_site + best_log_l)["Exp(_MATRIX_ELEMENT_VALUE_)"];
        normalizer = +this_site;
        this_site  = (this_site)*(1/normalizer);
        conditionals.scalers[s] = -best_log_l + Log(normalizer);

        for (g = 0; g < total_grid_points  ; g += 1) {
            conditionals.matrix [g][s] = this_site [g];
        }
    }

    return {"conditionals" : conditionals.matrix,
            "scalers" : conditionals.scalers };
}

//----------------------------------------------------------------------------


lfunction     RunMCMC (settings, grid, conditionals, callback, methodname) {
    if (None != methodname) {
        io.ReportProgressMessageMD                  (methodname, "mcmc", "Running MCMC chains to obtain a posterior sample of rate weights");
        io.ReportProgressMessageMD                  (methodname, "mcmc", "* Using the following settings");
        io.ReportProgressMessageMD                  (methodname, "mcmc", "\t* Number of chains : " + settings["chains"]);
        io.ReportProgressMessageMD                  (methodname, "mcmc", "\t* Steps/chain      : " + settings["chain-length"]);
        io.ReportProgressMessageMD                  (methodname, "mcmc", "\t* Burn-in steps    : " + settings["burn-in"]);
        io.ReportProgressMessageMD                  (methodname, "mcmc", "\t* Samples/chain    : " + settings["samples"]);
        io.ReportProgressMessageMD                  (methodname, "mcmc", "\t* Dirichlet alpha  : " + settings["concentration"]);
    }

    jobs = mpi.PartitionIntoBlocks(utility.Range (settings["chains"], 0, 1));

    queue  = mpi.CreateQueue (None);
    chains = {};

    for (i = 1; i < Abs (jobs); i += 1) {
        mpi.QueueJob (queue, "ExecuteMCMC", {"0" : grid,
                                                   "1" : conditionals,
                                                   "3" : settings,
                                                   "4" : jobs[i],
                                                   "2" : &chains},
                                                   callback);
    }

    utility.Extend (chains, ExecuteMCMC (grid, conditionals, &chains, settings, jobs[0]));

    mpi.QueueComplete (queue);

    return chains;
}
