// -*-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/>.
//

#include "bootstrap.h"

extern string out_path;
extern string out_prefix;
extern bool   bootstrap_pifis;
extern bool   bootstrap_div;
extern bool   bootstrap_fst;
extern bool   bootstrap_phist;

BSArchive::BSArchive(const MetaPopInfo *mpopi)
{
    this->_mpopi         = mpopi;
    this->_cur_pop_index = 0;
    
    //
    // Fix the internal vectors to their final size.
    //
    size_t pop_cnt = this->_mpopi->pops().size();
    this->_nuc_pop_stats.resize(pop_cnt, vector<array<double, _nuc_popstat_size>>());
    this->_hap_pop_stats.resize(pop_cnt, vector<array<double, _hap_popstat_size>>());

    if (pop_cnt == 1)
	return;

    this->_div_key.resize(pop_cnt, vector<int>(pop_cnt, -1));
    size_t pop_pair_cnt = 0;
    for (uint i = 0; i < pop_cnt; i++)
	for (uint j = i + 1; j < pop_cnt; j++) {
	    pop_pair_cnt++;
	}
    this->_nuc_div_stats.resize(pop_pair_cnt, vector<array<double, _nuc_divstat_size>>());
    this->_hap_div_stats.resize(pop_pair_cnt, vector<array<double, _hap_divstat_size>>());

}

void
BSArchive::set_cur_pop(size_t pop_1)
{
    if (pop_1 > this->_nuc_pop_stats.size() && pop_1 > this->_hap_pop_stats.size())
	throw exception();

    this->_cur_pop_index = pop_1;
}

void
BSArchive::set_cur_pop(size_t pop_1, size_t pop_2)
{
    if (pop_1 > this->_div_key.size() || pop_2 > this->_div_key[pop_1].size())
	throw exception();

    if (this->_div_key[pop_1][pop_2] == -1)
	throw exception();

    this->_cur_pop_index = this->_div_key[pop_1][pop_2];
}

size_t
BSArchive::num_stats(BootstrapT bs_type)
{
    size_t num_stats;
    switch(bs_type) {
    case BootstrapT::nuc_popstat:
	num_stats = this->_nuc_popstat_size;
	break;
    case BootstrapT::hap_popstat:
	num_stats = this->_hap_popstat_size;
	break;
    case BootstrapT::nuc_divstat:
	num_stats = this->_nuc_divstat_size;
	break;
    case BootstrapT::hap_divstat:
    default:
	num_stats = this->_hap_divstat_size;
	break;
    }
    return num_stats;
}

const double *
BSArchive::sample(BootstrapT bs_type)
{
    double *sample = NULL;
    size_t  size   = 1;
    int     index  = 0;
    
    switch(bs_type) {
    case BootstrapT::nuc_popstat:
        size   = this->_nuc_pop_stats[this->_cur_pop_index].size();
	index  = (int) size * (random() / (RAND_MAX + 1.0));
	sample = (this->_nuc_pop_stats[this->_cur_pop_index][index]).data();
	break;
    case BootstrapT::hap_popstat:
        size   = this->_hap_pop_stats[this->_cur_pop_index].size();
	index  = (int) size * (random() / (RAND_MAX + 1.0));
	sample = (this->_hap_pop_stats[this->_cur_pop_index][index]).data();
	break;
    case BootstrapT::nuc_divstat:
        size   = this->_nuc_div_stats[this->_cur_pop_index].size();
	index  = (int) size * (random() / (RAND_MAX + 1.0));
	sample = (this->_nuc_div_stats[this->_cur_pop_index][index]).data();
	break;
    case BootstrapT::hap_divstat:
	size   = this->_hap_div_stats[this->_cur_pop_index].size();
	index  = (int) size * (random() / (RAND_MAX + 1.0));
	sample = (this->_hap_div_stats[this->_cur_pop_index][index]).data();
	break;
    default:
	break;
    }

    // cerr << "    Randomly selected index: " << index << "; resulting in sample[0]: " << sample[0] << "\n";
    
    return sample;
}

int
BSArchive::parse_archive()
{
    string path = out_path + out_prefix + ".bsarchive.tsv";

    ifstream fh(path.c_str(), ifstream::in);

    if (fh.fail()) {
        cerr << "Error: Failed to open the boostrap archive file '" << path << "'.\n";
        throw exception();
    }

    const vector<Pop>& pops = this->_mpopi->pops();

    cerr << "Parsing bootstrap archive file to obtain values for resampling...";

    char line[max_len];
    memset(line, '\0', max_len);
    vector<string> parts;

    //
    // Parse the population list first, make sure the same popmap is being used.
    // We expect the number of populations on the first line, followed by one
    // line for each of those pops.
    //
    size_t lineno = 0;
    fh.getline(line, max_len);
    lineno++;
    int pop_cnt = is_integer(line);
    if (pop_cnt < 1) {
	cerr << "\nError: malformed bootstrap archive file '" << path << "', line 1 should contain number of populations.\n";
	throw exception();
    }

    for (int i = 0; i < pop_cnt; i++) {
	fh.getline(line, max_len);
	lineno++;

	//
	// Check that the population name matches the expected numerical index.
	//
	parse_tsv(line, parts);
	if (parts.size() != 2) {
	    cerr << "\nError: malformed bootstrap archive file '" << path << "', "
		 << "line " << lineno << " should contain two columns describing a population and its members.\n";
	    throw exception();
	}

	if (pops[i].name != parts[0]) {
	    cerr << "\nError: malformed bootstrap archive file '" << path << "', "
		 << "line " << lineno << "; population listed does not match popmap file, has popmap changed since the archive was written?\n";
	    throw exception();
	}
    }

    //
    // Generate the list of population pairs that we expect to find in the archive file.
    //
    size_t index = 0;
    for (int i = 0; i < pop_cnt; i++)
	for (int j = i+1; j < pop_cnt; j++) {
	    this->_div_key[i][j] = index;
	    index++;
	}
    
    double stat_1, stat_2, stat_3, stat_4;
    int    pop_index_1, pop_index_2;

    //
    // Parse the body of the file. 
    //
    while (fh.getline(line, max_len)) {
	lineno++;
	
        size_t len = strlen(line);

        // Skip empty lines and comments
        if (len == 0 || line[0] == '#')
            continue;

        // Check for Windows line endings
        if (line[len - 1] == '\r') {
            line[len - 1] = '\0';
            len -= 1;
        }

        //
        // Parse the contents, we expect one of four different types of lines:
        // type<tab>population index 1[<tab>population index 2]<tab>stat[<tab>stat...]
        // 'type' can be:
	//    a haplotype-level ('P') population statistic or a nucleotide-level ('p') pop statistic.
	//    a haplotype-level ('D') diversity statistic or a nucleotide-level ('d') diversity statistic.
	//  Population statistics are keyed to a single population, diversity statistics to a pair of populations.
	//
        parse_tsv(line, parts);
        for (string& part : parts)
            strip_whitespace(part);

        if (parts.size() < 5 || parts.size() > 7) {
            cerr << "\nError: Malformed bootstrap archive file '" << path << "', line: " << lineno << "\n";
            throw exception();
        }

	if (bootstrap_pifis && parts[0] == "p") {
	    try {
		pop_index_1 = stoi(parts[1], NULL);
		stat_1      = stod(parts[3], NULL);
		stat_2      = stod(parts[4], NULL);

		this->_nuc_pop_stats[pop_index_1].push_back(array<double, _nuc_popstat_size>{stat_1, stat_2});

	    } catch (exception &e) {
		cerr << "\nError: Malformed bootstrap archive file '" << path << "', line: " << lineno << "\n";
	    }		
	} else if (bootstrap_div && parts[0] == "P") {
	    try {
		pop_index_1 = stoi(parts[1], NULL);
		stat_1      = stod(parts[2], NULL);
		stat_2      = stod(parts[3], NULL);
		stat_3      = stod(parts[4], NULL);

		this->_hap_pop_stats[pop_index_1].push_back(array<double, _hap_popstat_size>{stat_1, stat_2, stat_3});

	    } catch (exception &e) {
		cerr << "\nError: Malformed bootstrap archive file '" << path << "', line: " << lineno << "\n";
	    }
	} else if (bootstrap_fst && parts[0] == "d") {
	    try {
		pop_index_1 = stoi(parts[1], NULL);
		pop_index_2 = stoi(parts[2], NULL);
		stat_1      = stod(parts[4], NULL);

		if (pop_index_1 < 0 || pop_index_2 < 0)
		    throw exception();

		index = this->_div_key[pop_index_1][pop_index_2];
		this->_nuc_div_stats[index].push_back(array<double, _nuc_divstat_size>{stat_1});

	    } catch (exception &e) {
		cerr << "\nError: Malformed bootstrap archive file '" << path << "', line: " << lineno << "\n";
	    }
	} else if (bootstrap_phist && parts[0] == "D") {
	    try {
		pop_index_1 = stoi(parts[1], NULL);
		pop_index_2 = stoi(parts[2], NULL);
		stat_1      = stod(parts[3], NULL);
		stat_2      = stod(parts[4], NULL);
		stat_3      = stod(parts[5], NULL);
		stat_4      = stod(parts[6], NULL);

		if (pop_index_1 < 0 || pop_index_2 < 0)
		    throw exception();

		index = this->_div_key[pop_index_1][pop_index_2];
		this->_hap_div_stats[index].push_back(array<double, _hap_divstat_size>{stat_1, stat_2, stat_3, stat_4});

	    } catch (exception &e) {
		cerr << "\nError: Malformed bootstrap archive file '" << path << "', line: " << lineno << "\n";
	    }
	}
    }

    cerr << "done.\n";

    if (bootstrap_pifis) {
	for (int i = 0; i < pop_cnt; i++)
	    if (this->_nuc_pop_stats[i].size() == 0) {
		cerr << "Error: Bootstrap resampling of smoothed nucleotide statistics requested (--bootstrap-pifis), "
		     << "but no data found in bootstrap archive for population " << this->_mpopi->pops()[i].name << ".\n";
		exit(1);
	    }	
    }

    if (bootstrap_div) {
	for (int i = 0; i < pop_cnt; i++)
	    if (this->_hap_pop_stats[i].size() == 0) {
		cerr << "Error: Bootstrap resampling of smoothed haplotype statistics requested (--bootstrap-div), "
		     << "but no data found in bootstrap archive for population " << this->_mpopi->pops()[i].name << ".\n";
		exit(1);
	    }	
    }

    if (bootstrap_fst) {
    	for (uint i = 0; i < this->_div_key.size(); i++)
    	    for (uint j = i + 1; j < this->_div_key[i].size(); j++)
    		if (this->_nuc_div_stats[this->_div_key[i][j]].size() == 0) {
    		    cerr << "Error: Bootstrap resampling of smoothed Fst nucleotide statistics requested (--bootstrap-fst), "
    			 << "but no data found in bootstrap archive for populations "
    			 << "'" << this->_mpopi->pops()[i].name << "' and '" << this->_mpopi->pops()[j].name << "'"
    			 << " (did you turn on --fstats when generating the archive?).\n";
    		    exit(1);
    		}
    }

    if (bootstrap_phist) {
    	for (uint i = 0; i < this->_div_key.size(); i++)
    	    for (uint j = i + 1; j < this->_div_key[i].size(); j++)
    		if (this->_hap_div_stats[this->_div_key[i][j]].size() == 0) {
    		    cerr << "Error: Bootstrap resampling of smoothed Fst haplotype statistics requested (--bootstrap-phist), "
    			 << "but no data found in bootstrap archive for populations "
    			 << "'" << this->_mpopi->pops()[i].name << "' and '" << this->_mpopi->pops()[j].name << "'"
    			 << " (did you turn on --fstats when generating the archive?)..\n";
    		    exit(1);
    		}
    }

    for (int i = 0; i < pop_cnt; i++) {
	cerr << "  " << pops[i].name << ": "
	     << this->_hap_pop_stats[i].size() << " sets of haplotype-based per-population statistics read; "
	     << this->_nuc_pop_stats[i].size() << " sets of nucleotide-based per-population statistics read.\n";
    }
    for (uint i = 0; i < this->_div_key.size(); i++)
	for (uint j = i + 1; j < this->_div_key[i].size(); j++) {
	    cerr << "  " << pops[i].name << "-" << pops[j].name << ": "
		 << this->_hap_div_stats[this->_div_key[i][j]].size() << " sets of haplotype-based divergence statistics read; "
		 << this->_nuc_div_stats[this->_div_key[i][j]].size() << " sets of nucleotide-based divergence statistics read.\n";
	}
    cerr << "\n";
    
    return 0;
}
