#include "fileio.h"

////////////////////////////////////////////////////////////////////////////////////

////////////////////////////////////////////////////////////////////////////////////

//Gets the recombination fractions between loci
int read_recom_fracs(std::vector<locus> *loci, int *n_loci, std::string loc_filename, GeneDrop* ui)
{
    std::ifstream input;
    std::string line;
    int line_num=0;

    boost::tokenizer<boost::char_separator<char> >::iterator field;
    boost::char_separator<char> sep(",");

    input.clear();
    input.open(loc_filename.c_str());

    if(!input.good())
    {
        std::stringstream message;

        input.close();
        message<<loc_filename<<" was not found.";
        throw std::runtime_error(message.str());
    }

    getline(input, line);

    *n_loci=0;

    while(getline(input, line))
    {
        (*n_loci)++;
    }

    loci->resize(*n_loci);

    input.clear();
    input.seekg(0, std::ios::beg);

    getline(input, line);	// To throw away header line
    getline(input, line);

    while(!input.eof())
    {
        boost::tokenizer<boost::char_separator<char> > tok(line, sep);
        field=tok.begin();

        loci->at(line_num).change_ID(*field);
        loci->at(line_num).change_pos(line_num);
        field++;
        loci->at(line_num).change_dist(atof(field->c_str()));
        field++;
        loci->at(line_num).change_chrom(*field);

        line_num++;
        getline(input, line);
    }

    input.close();

    return 0;
}

////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////

////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////

//Gets the genotypes of the founder lines
int read_genotypes(int *current_pop1, std::vector<plant> *pop1, int n_loci, int *n_founders, boost::uniform_01<boost::mt19937>* rng, std::vector<locus> loci, GeneDrop* ui, std::string ped_filename, std::string fou_filename)
{
    int chromo_choose;
    int i, j;
    int offset;
    int p1_found, p2_found;
    int p;
    int value;
    std::ifstream input;
    std::string line;
    std::string p1_check, p2_check;
    std::vector<plant> founders;	//Holds founder generation
    std::vector<std::vector<int> > hetero;	//Holds the position of any heterozygous loci with coordinates [founder][locus]

    boost::tokenizer<boost::char_separator<char> >::iterator field;
    boost::char_separator<char> sep(",");

    input.clear();
    input.open(fou_filename.c_str());

    if(!input.good())
    {
        std::stringstream message;

        input.close();
        message<<fou_filename<<" was not found.";
        ui->print(message);
        return 1;
    }

    founders.resize(*n_founders, n_loci);
    hetero.resize(*n_founders);

    getline(input, line);	// To throw away header line

    for(i=0; i<*n_founders; i++)
    {
        getline(input, line);
        boost::tokenizer<boost::char_separator<char> > tok(line, sep);
        field=tok.begin();

        founders[i].change_ID(*field);
        field++;

        for(j=0; j<n_loci; j++)
        {
            //Check if heterozygote or homozygote (by checking for / notation)
            if(field->find('/')==std::string::npos)
            {
                value=atoi(field->c_str());
                founders[i].change_genotype(j, value);
                founders[i].change_genotype(j+n_loci, value);

            }
            else
            {
                hetero.at(i).push_back(j);
                founders[i].change_genotype(j, atoi(field->substr(0, 1).c_str()));
                founders[i].change_genotype(j+n_loci, atoi(field->substr(2, 3).c_str()));
            }
            field++;
        }
    }

    input.close();

    input.clear();
    input.open(ped_filename.c_str());

    if(!input.good())
    {
        std::stringstream message;

        input.close();
        message<<ped_filename<<" was not found.";
        ui->print(message);
        return 1;
    }

    getline(input, line);	//Throw away header

    *current_pop1=0;

    for(i=0; getline(input, line); i++)
    {
        boost::tokenizer<boost::char_separator<char> > tok(line, sep);
        field=tok.begin();

        //Resize if about to be too small
        if(*current_pop1==(int)(pop1->size()-1))
        {
            pop1->resize(2*(pop1->size()), n_loci);
        }

        if(atoi(field->c_str())!=0)
        {
            break;
        }

        field++;
        pop1->at(i).change_ID(*field);
        field++;
        pop1->at(i).change_p1(*field);
        field++;
        pop1->at(i).change_p2(*field);

        (*current_pop1)++;
    }

    if(*current_pop1==0)
    {
        ui->print("Failed to count initial population size. Did you forget to label the first generation as 0 in the pedigree file?");
        return 1;
    }

    input.close();

    for(i=0; i<*current_pop1; i++)
    {
        p1_check=pop1->at(i).get_p1();
        p2_check=pop1->at(i).get_p2();

        p1_found=-1;
        p2_found=-1;

        for(j=0; j<*n_founders; j++)
        {
            if(p1_check==founders[j].get_ID())
            {
                p1_found=j;
            }
            if(p2_check==founders[j].get_ID())
            {
                p2_found=j;
            }
            if(p1_found>=0 && p2_found>=0)
            {
                break;
            }
        }

        if(p1_found==-1 || p2_found==-1)
        {
            ui->print("There was an error in determining which founders give rise to which lines.");
            return 1;
        }

        //Check if there are any heterozygous loci for this parental pair
        if(hetero.at(p1_found).empty() && hetero.at(p2_found).empty())
        {
            for(j=0; j<n_loci; j++)
            {
                pop1->at(i).change_genotype(j, founders[p1_found].get_genotype(j));
                pop1->at(i).change_genotype(j+n_loci, founders[p2_found].get_genotype(j));
            }
        }
        else	//Do more complex breeding taking recombination into account if heterozygous loci, as detailed in crossing()
        {
            //Produce both chromosomes by moving input index by offset of 0 or n_loci
            for(offset=0; offset<n_loci-1; offset+=n_loci)
            {
                //Depending on offset, draw alleles from different parent
                if(offset==0)
                {
                    p=p1_found;
                }
                else
                {
                    p=p2_found;
                }

                if((*rng)()<0.5)
                {
                    chromo_choose=0;
                    pop1->at(i).change_genotype(offset, founders.at(p).get_genotype(0));
                }
                else
                {
                    chromo_choose=1;
                    pop1->at(i).change_genotype(offset, founders.at(p).get_genotype(n_loci));
                }

                for(j=0; j<(n_loci-1); j++)
                {
                    if((*rng)()<get_recom(loci[j], loci[j+1]))
                    {
                        chromo_switch(&chromo_choose);
                    }

                    if(chromo_choose==1)
                    {
                        pop1->at(i).change_genotype(offset+j+1, founders.at(p).get_genotype(n_loci+j+1));
                    }
                    else if(chromo_choose==0)
                    {
                        pop1->at(i).change_genotype(offset+j+1, founders.at(p).get_genotype(j+1));
                    }
                    else
                    {
                        std::stringstream message;

                        message<<"There was an error in determining which chromosome to choose an allele from at the "<<offset+j+1<<"th position in individual number "<<i+1<<" in the initial founder breeding round.";
                        ui->print(message);
                        return 1;
                    }
                }
            }
        }
    }

    return 0;
}

/////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////

/////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////

//Works out which crosses need to be simulated for the current generation
int read_crosses(int gen, int *current_pop1, int *current_pop2, std::vector<plant> *pop1, std::vector<plant> *pop2, int n_loci, GeneDrop *ui, std::string ped_filename)
{
    std::ifstream input;
    std::string line;
    int line_num=0;
    int start_pop=500;
    boost::tokenizer<boost::char_separator<char> >::iterator field;
    std::stringstream message;

    boost::char_separator<char> sep(",");

    input.clear();
    input.open(ped_filename.c_str());

    if(!input.good())
    {
        message<<"Reading in pedigree information failed(at line <<"<<line_num+1<<").";
        ui->print(message);
        input.close();
        return 1;
    }

    getline(input, line);		//Ignore header

    if(gen%2==0)
    {
        pop2->clear();
        pop2->resize(start_pop, n_loci);
        *current_pop2=0;
    }
    else if(gen%2!=0)
    {
        pop1->clear();
        pop1->resize(start_pop, n_loci);
        *current_pop1=0;
    }

    while(getline(input, line))
    {
        boost::tokenizer<boost::char_separator<char> > tok(line, sep);
        field=tok.begin();

        if(gen+1==atoi(field->c_str()))
        {
            if(gen%2==0)
            {
                //Resize if about to be too small
                if(*current_pop2==(int)(pop2->size()-1))
                {
                    pop2->resize(2*(pop2->size()), n_loci);
                }

                field++;
                pop2->at(line_num).change_ID(*field);
                field++;
                pop2->at(line_num).change_p1(*field);
                field++;
                pop2->at(line_num).change_p2(*field);

                (*current_pop2)++;
            }
            else
            {
                if(*current_pop1==(int)(pop1->size()-1))
                {
                    pop1->resize(2*(pop1->size()), n_loci);
                }

                field++;
                pop1->at(line_num).change_ID(*field);
                field++;
                pop1->at(line_num).change_p1(*field);
                field++;
                pop1->at(line_num).change_p2(*field);

                (*current_pop1)++;
            }
            line_num++;
        }
    }

    input.close();

    return 0;
}

/////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////

/////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////

//Prints out results to output file
int print_results(int run_num, int gen, int n_gens, int seed, int n_runs, int n_loci, int n_founders, int current_pop1, int current_pop2, std::vector<locus> loci, std::vector<plant> pop1, std::vector<plant> pop2, bool print_all, GeneDrop* ui)
{
    std::ofstream output;
    int i, j;
    static int gen_flag=0;	//Takes account of what files have been created for which generations
    static bool flag=false;	//For when printing out at end only
    static std::vector<std::string> filenames;
    static std::string stamp;
    static std::string stampForFilename;
    std::stringstream message;
    std::stringstream temp;
    time_t time_val;
    struct tm *time_holder;

    //Reset filenames on first use of print_results
    if(run_num==0 && (gen==0 || (!print_all && gen==n_gens)))
    {
        filenames.resize(0);
        flag=false;
        gen_flag=0;

        time_val=time(NULL);
        time_holder=localtime(&time_val);
        stamp=boost::lexical_cast<std::string>((int)time_holder->tm_hour)+":"+boost::lexical_cast<std::string>((int)time_holder->tm_min)+":"+boost::lexical_cast<std::string>((int)time_holder->tm_sec)+"@"+boost::lexical_cast<std::string>((int)time_holder->tm_mday)+"-"+boost::lexical_cast<std::string>(((int)time_holder->tm_mon)+1)+"-"+boost::lexical_cast<std::string>(time_holder->tm_year+1900);
        stampForFilename=boost::lexical_cast<std::string>((int)time_holder->tm_hour)+"-"+boost::lexical_cast<std::string>((int)time_holder->tm_min)+"-"+boost::lexical_cast<std::string>((int)time_holder->tm_sec)+"@"+boost::lexical_cast<std::string>((int)time_holder->tm_mday)+"-"+boost::lexical_cast<std::string>(((int)time_holder->tm_mon)+1)+"-"+boost::lexical_cast<std::string>(time_holder->tm_year+1900);
    }

    //Produce filenames vector of appropriate size on first run through function
    if(!flag && print_all)
    {
        filenames.resize(n_gens+1);
    }

    //Produce filename(s)
    if(!print_all && !flag)
    {
        temp<<"Output("<<stampForFilename<<").csv";
        filenames.push_back(temp.str());

        message<<"New output file produced - "<<filenames[0];
        ui->print(message);
    }
    else if(gen_flag<=gen && print_all)
    {
        temp<<"Output("<<stampForFilename<<")-"<<gen<<".csv";
        filenames[gen]=temp.str();

        message<<"New output file produced - "<<filenames[gen];
        ui->print(message);
    }

    if(print_all)
    {
        output.open(filenames[gen].c_str(), std::ios::out | std::ios::app);
    }
    else
    {
        output.open(filenames[0].c_str(), std::ios::out | std::ios::app);
    }

    if(!output.good())
    {
        ui->print("Opening the output file failed.");
        return 1;
    }

    if((!print_all && !flag) || (print_all && gen_flag<=gen))	//Label the file
    {
        output<<stamp<<std::endl;
        output<<"Number of loci,"<<n_loci<<std::endl;
        output<<"Number of founders,"<<n_founders<<std::endl;
        output<<"Number of generations,"<<n_gens+1<<std::endl;
        output<<"Seed,"<<seed<<std::endl;
        output<<"Number of runs,"<<n_runs<<std::endl;
        output<<std::endl;
        output<<std::endl;

        output<<"Run Number";
        output<<",ID";

        for(j=0; j<n_loci; j++)
        {
            output<<","<<loci[j].get_ID()<<" -1,"<<loci[j].get_ID()<<" -2";
        }

        output<<std::endl;

        flag=true;
        gen_flag++;
    }

    if(gen%2==0)	//Determines whether results are held in pop1 or pop2
    {
        for(i=0; i<current_pop1; i++)
        {
            output<<run_num<<","<<pop1.at(i).get_ID();

            for(j=0; j<n_loci; j++)
            {
                output<<","<<pop1.at(i).get_genotype(j);		//Print loci on the same chromosome next to each other
                output<<","<<pop1.at(i).get_genotype(j+n_loci);
            }
            output<<std::endl;
        }
    }
    else
    {
        for(i=0; i<current_pop2; i++)
        {
            output<<run_num<<","<<pop2.at(i).get_ID();

            for(j=0; j<n_loci; j++)
            {
                output<<","<<pop2.at(i).get_genotype(j);
                output<<","<<pop2.at(i).get_genotype(j+n_loci);
            }
            output<<std::endl;
        }
    }

    output.close();

    return 0;
}

/////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////

/////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////

//Works out how many generations there are in the pedigree
int get_num_gens(int *n_gens, std::string ped_filename, GeneDrop* ui)
{
    std::ifstream input;
    std::string line;
    int current_count;	//Holds the current number of generations counted
    boost::tokenizer<boost::char_separator<char> >::iterator field;
    boost::char_separator<char> sep(",");

    input.open(ped_filename.c_str());

    if(!input.good())
    {
        std::stringstream message;

        input.close();
        message<<ped_filename<<" was not found.";
        throw std::runtime_error(message.str());
    }

    getline(input, line);	//Chuck out header

    current_count=0;

    while(getline(input, line))
    {
        boost::tokenizer<boost::char_separator<char> > tok(line, sep);
        field=tok.begin();

        current_count=atoi(field->c_str());
    }

    *n_gens=current_count;

    input.close();

    return 0;
}

///////////////////////////////////////////////////////////////////////////////////////

///////////////////////////////////////////////////////////////////////////////////////

//Works out how many founder lines there are
int get_num_founders(int *n_founders, std::string fou_filename, GeneDrop* ui)
{
    std::ifstream input;
    std::string line;

    input.open(fou_filename.c_str());

    if(!input.good())
    {
        std::stringstream message;

        input.close();
        message<<fou_filename<<" was not found.";
        throw std::runtime_error(message.str());
    }

    getline(input, line);

    *n_founders=0;

    while(getline(input, line))
    {
        (*n_founders)++;
    }

    input.close();

    return 0;
}


