// -*-mode:c++; c-style:k&r; c-basic-offset:4;-*-
//
// Copyright 2014-2021, Julian Catchen <jcatchen@illinois.edu>
//
// This file is part of Stacks.
//
// Stacks is free software: you can redistribute it and/or modify
// it under the terms of the GNU General Public License as published by
// the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// Stacks is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU General Public License for more details.
//
// You should have received a copy of the GNU General Public License
// along with Stacks.  If not, see <http://www.gnu.org/licenses/>.
//

#ifndef __BOOTSTRAP_H__
#define __BOOTSTRAP_H__

#include "constants.h"
#include "MetaPopInfo.h"
#include "PopSum.h"
#include "input.h"
#include "smoothing_utils.h"

extern double   sigma;
extern int      bootstrap_reps;
extern bool     bootstrap_wl;
extern set<int> bootstraplist;

enum class BootstrapT {
    nuc_popstat, hap_popstat,
    nuc_divstat, hap_divstat
};

//
// Bootstrap archive structure. Store all of the point values for various population statistics
// for use in resampling during bootstrapping.
//
class BSArchive {
    static const uint8_t _nuc_popstat_size = 2;
    static const uint8_t _hap_popstat_size = 3;
    static const uint8_t _nuc_divstat_size = 1;
    static const uint8_t _hap_divstat_size = 4;

    const MetaPopInfo *_mpopi;

    //
    // A key to map pairs of populations to their index in the stats vector.
    //   _div_key[pop_1][pop_2] will return the index into the vector for the population pair.
    //
    vector<vector<int>> _div_key;

    // One collection of double arrays for each population.
    vector<vector<array<double, _nuc_popstat_size>>> _nuc_pop_stats; // nucleotide-level per-population stats.
    vector<vector<array<double, _hap_popstat_size>>> _hap_pop_stats; // haplotye-level (locus) per-population stats.
    vector<vector<array<double, _nuc_divstat_size>>> _nuc_div_stats; // nucleotide-level population-pari divergence stats.
    vector<vector<array<double, _hap_divstat_size>>> _hap_div_stats; // haplotye-level (locus) population-pair divergence stats.

    // Set of references to vectors set to the current population / pair of populations; for efficiency.
    size_t _cur_pop_index;

public:
    BSArchive(const MetaPopInfo *);
    int    parse_archive();
    size_t num_stats(BootstrapT type);
    void   set_cur_pop(size_t pop_1);
    void   set_cur_pop(size_t pop_1, size_t pop_2);
    const double *sample(BootstrapT type);
};

//
// Bootstrap resamplign structure.
//
class BSample {
public:
    int    bp;
    int    alleles;
    bool   fixed;
    double stat[PopStatSize];

    BSample() {
        this->bp      = 0;
        this->alleles = 0;
        this->fixed   = false;
        for (uint i = 0; i < PopStatSize; i++)
            this->stat[i] = 0.0;
    }
};

template<class StatT=PopStat>
class Bootstrap {
    BootstrapT  _bootstrap_type;
    BSArchive  *_bs_archive; // Archive holding statistical values for resampling.
    double     *_weights;    // Weight matrix to apply while smoothing.

public:
    Bootstrap(const MetaPopInfo *mpopi, BSArchive *bsarchive, BootstrapT btype)  {
	this->_bootstrap_type = btype;
	//
 	// Class storing statistical values for resampling during bootstrapping
	// (must be generated via --bootstrap-archive prior to a bootstrapping run.)
	//
	this->_bs_archive = bsarchive;
	//
	// Prepare the weights for use in the sliding window average calculation.
	//
        this->_weights = calc_weights();
    }
    ~Bootstrap() {
        delete [] this->_weights;
    }

    void   set_cur_pop(size_t pop_1) { this->_bs_archive->set_cur_pop(pop_1); }
    void   set_cur_pop(size_t pop_1, size_t pop_2) { this->_bs_archive->set_cur_pop(pop_1, pop_2); }

    int    execute(vector<StatT *> &);
    double pval(double, vector<double> &);
};

template<class StatT>
int
Bootstrap<StatT>::execute(vector<StatT *> &sites)
{
    size_t num_stats = this->_bs_archive->num_stats(this->_bootstrap_type);

    #pragma omp parallel
    {
        StatT *c;
        double final_weight, sum, weighted_stat[num_stats];
        int  dist;
        uint pos_l = 0;
        uint pos_u = 0;

        #pragma omp for schedule(dynamic, 1)
        for (uint pos_c = 0; pos_c < sites.size(); pos_c++) {
            c = sites[pos_c];

            if (c == NULL || c->fixed == true)
                continue;

            if (bootstrap_wl && bootstraplist.count(c->loc_id) == 0)
                continue;

            determine_window_limits(sites, c->bp, pos_l, pos_u);

            int non_null_sitecnt = 0;
            for (uint i = pos_l; i < pos_u;  i++)
                if (sites[i] != NULL) non_null_sitecnt++;

            // cerr << "Bootstrapping locus: " << c->loc_id << "; pos_c: " << pos_c << "; bp: " << c->bp << "; "
            //      << "window start: " << pos_l << " (" << sites[pos_l]->bp << "), window end: " << pos_u << " (" << sites[pos_u]->bp << "); non-null site cnt: " << non_null_sitecnt << "\n";

            //
            // Allocate an array of bootstrap resampling objects.
            //
            BSample *bs = new BSample[non_null_sitecnt];

            //
            // Populate the BSample objects.
            //
            int j = 0;
            for (uint i = pos_l; i < pos_u;  i++) {
                if (sites[i] == NULL)
                    continue;
                bs[j].bp      = sites[i]->bp;
                bs[j].alleles = sites[i]->alleles;
                bs[j].fixed   = sites[i]->fixed;
                for (uint k = 0; k < num_stats; k++)
                    bs[j].stat[k] = sites[i]->stat(k);
                j++;
            }

            //
            // For some population statistics, if the site is fixed it does not make sense to calculate a value
            // for that site. So, Fis is not defined on fixed sites, but Pi is. So, we cannot mark the site fully
            // fixed, i.e. bs[j].fixed != TRUE. Instead, we rely on the value for that specific statistic to
            // be special, i.e. FSNC, defined in constants.h == Fixed Site Not Calculated. We want to exclude
            // these sites from all aspects of bootstrapping.
            // JC - Oct 21, 2021
            //
            
            //
            // Precompute the fraction of the window that will not change during resampling (the fixed sites).
            //
            double partial_weighted_stat[num_stats];
            double partial_sum = 0.0;
            memset(partial_weighted_stat, 0, num_stats);

            for (j = 0; j < non_null_sitecnt; j++) {
                if (bs[j].fixed == false) continue;

                dist = bs[j].bp > c->bp ? bs[j].bp - c->bp : c->bp - bs[j].bp;

                final_weight  = (bs[j].alleles - 1.0) * this->_weights[dist];
                partial_sum  += final_weight;
                for (uint k = 0; k < num_stats; k++)
                    if (bs[j].stat[k] > FSNC)
                        partial_weighted_stat[k] += bs[j].stat[k] * final_weight;
            }

            vector<vector<double>> resampled_stats(num_stats, vector<double>());
            for (uint i = 0; i < num_stats; i++)
                resampled_stats[i].reserve(bootstrap_reps);

            //
            // Bootstrap this chromosome.
            //
            for (int i = 0; i < bootstrap_reps; i++) {
                // cerr << "Bootsrap rep " << i << "\n";

                for (uint k = 0; k < num_stats; k++)
                    weighted_stat[k] = partial_weighted_stat[k];
                sum = partial_sum;

                for (j = 0; j < non_null_sitecnt; j++) {
                    if (bs[j].fixed == true) continue;

                    dist = bs[j].bp > c->bp ? bs[j].bp - c->bp : c->bp - bs[j].bp;

                    //
                    // Resample for this round of bootstrapping.
                    //
		    const double *sample = this->_bs_archive->sample(this->_bootstrap_type);
                    // cerr << "    Resampling position " << j << " (" << bs[j].bp << ") from " << bs[j].stat[1] << " to " << sample[1] << "\n";
                    for (uint k = 0; k < num_stats; k++)
                        if (bs[j].stat[k] > FSNC)
                            bs[j].stat[k] = sample[k];

                    final_weight = (bs[j].alleles - 1) * this->_weights[dist];
                    for (uint k = 0; k < num_stats; k++)
                        if (bs[j].stat[k] > FSNC)
                            weighted_stat[k] += bs[j].stat[k] * final_weight;
                    sum += final_weight;
                }

                // cerr << "    New weighted value: " << (weighted_stat[1] / sum) << "\n";
                for (uint k = 0; k < num_stats; k++)
                    resampled_stats[k].push_back(weighted_stat[k] / sum);
            }

            //
            // Cacluate the p-value for this window based on the empirical Fst distribution.
            //
            for (uint k = 0; k < num_stats; k++) {
                sort(resampled_stats[k].begin(), resampled_stats[k].end());
                c->bs(k, this->pval(c->smoothed(k), resampled_stats[k]));
            }

            delete [] bs;
        }
    }

    return 0;
}

template<class StatT>
double
Bootstrap<StatT>::pval(double stat, vector<double> &dist)
{
    // Rank `stat` relative to the bootstrapped distribution.
    // `first_greater` is between 0 (if `stat` is smaller than all values in the
    // distribution) and `dist.size()` (if `stat` is greater than all of them).
    size_t first_greater = upper_bound(dist.begin(), dist.end(), stat) - dist.begin();

    // cerr << "Generated Smoothed Fst Distribution:\n";
    // for (uint n = 0; n < dist.size(); n++)
    //         cerr << "  n: " << n << "; Fst: " << dist[n] << "\n";

    // cerr << "Comparing Fst value: " << stat
    //          << " at position " << (up - dist.begin()) << " out of "
    //          << dist.size() << " positions (converted position: " << pos << "); pvalue: " << res << ".\n";

    return 1.0 - (double) first_greater / dist.size();
}

#endif // __BOOTSTRAP_H__
