#include "biofunctions.h"

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

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

//Runs simulation
int run_simulation(bool print_all, boost::uniform_01<boost::mt19937>* rng, int run_num, int n_gens, std::vector<locus> loci, int seed, int n_runs, int n_loci, int *n_founders, GeneDrop *ui, std::string ped_filename, std::string fou_filename)
{
    int current_pop1=0;	//Current size of pop1
    int current_pop2=0;	//Current size of pop2
    int i;
    int start_pop=500;	//Initial guesstimate of maximum population size
    std::vector<plant> pop1, pop2;	//Generations of plants

    //Zeroes out populations for next run
    pop1.clear();
    pop1.resize(start_pop, n_loci);
    pop2.clear();
    pop2.resize(start_pop, n_loci);

    //Read in the founder genotypes and produce F1
    try
    {
        read_genotypes(&current_pop1, &pop1, n_loci, n_founders, rng, loci, ui, ped_filename, fou_filename);
    }
    catch(std::exception &e)
    {
        std::stringstream message;

        message<<"Reading in the founder genotypes failed: "<<e.what();
        ui->print(message);
        throw;
    }
    catch(...)
    {
        ui->print("Reading in the founder genotypes failed.");
        throw;
    }
    boost::this_thread::interruption_point();

    //If printing at each generation, print out the data
    if(print_all)
    {
        try
        {
            print_results(run_num, 0, n_gens, seed, n_runs, n_loci, *n_founders, current_pop1, current_pop2, loci, pop1, pop2, print_all, ui);
        }
        catch(std::exception &e)
        {
            std::stringstream message;

            message<<"Printing out the results failed: "<<e.what();
            ui->print(message);
            throw;
        }
        catch(...)
        {
            ui->print("Printing out the results failed.");
            throw;
        }
    }

    //Run through each generation of crossing
    for(i=0; i<n_gens; i++)
    {
        try
        {
            read_crosses(i, &current_pop1, &current_pop2, &pop1, &pop2, n_loci, ui, ped_filename);
        }
        catch(std::exception &e)
        {
            std::stringstream message;

            message<<"Reading in the crosses failed: "<<e.what();
            ui->print(message);
            throw;
        }
        catch(...)
        {
            ui->print("Reading in the crosses failed.");
            throw;
        }
        boost::this_thread::interruption_point();

        try
        {
            crossing(i, rng, current_pop1, current_pop2, &pop1, &pop2, loci, n_loci, ui);
        }
        catch(std::exception &e)
        {
            std::stringstream message;

            message<<"Carrying out the crosses failed: "<<e.what();
            ui->print(message);
            throw;
        }
        catch(...)
        {
            ui->print("Carrying out the crosses failed.");
            throw;
        }
        boost::this_thread::interruption_point();

        if(print_all || i==n_gens-1)
        {
            try
            {
                print_results(run_num, i+1, n_gens, seed, n_runs, n_loci, *n_founders, current_pop1, current_pop2, loci, pop1, pop2, print_all, ui);
            }
            catch(std::exception &e)
            {
                std::stringstream message;

                message<<"Printing out the results failed: "<<e.what();
                ui->print(message);
                throw;
            }
            catch(...)
            {
                ui->print("Printing out the results failed.");
                throw;
            }
        }
    }

    return 0;
}

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

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

//Carries out the current generation's crosses
int crossing(int round, boost::uniform_01<boost::mt19937>* rng, int current_pop1, int current_pop2, std::vector<plant> *pop1, std::vector<plant> *pop2, std::vector<locus> loci, int n_loci, GeneDrop* ui)
{
    int chromo_choose;		//Determines which chromosome the next allele is chosen from
    int i, j;
    int offset;
    int num_read, num_change;	//Either current_pop1 or current_pop2 depending on the generation being simulated
    int p1_found;	//Holds position of the first parent in their population std::vector
    int p2_found;	//Holds position of the second parent in their population std::vector
    int p;
    std::string p1_check;	//Holds name of first parent one for comparisons
    std::string p2_check;	//Holds name of second parent one for comparisons
    std::vector<plant> *pop_read, *pop_change;	//Pointers to either pop1 or pop2, depending on which generation is being simulated
    void chromo_switch(int*);	//Switches chromosome to draw alleles from

    //Work out if generation number is odd or even and then change which population to draw alleles from appropriately
    if(round%2==0)
    {
        pop_read=pop1;	//Since pops are passed as pointers, this works
        pop_change=pop2;
        num_read=current_pop1;
        num_change=current_pop2;
    }
    else
    {
        pop_read=pop2;
        pop_change=pop1;
        num_read=current_pop2;
        num_change=current_pop1;
    }

    //Cycle through each individual in the population to be altered
    for(i=0; i<num_change; i++)
    {
        p1_found=-1;	//Initialise to -1 to show that position of parent is unknown
        p2_found=-1;

        p1_check=pop_change->at(i).get_p1();
        p2_check=pop_change->at(i).get_p2();

        for(j=0; j<num_read; j++)
        {
            if(pop_read->at(j).get_ID()==p1_check)
                    p1_found=j;
            if(pop_read->at(j).get_ID()==p2_check)
                    p2_found=j;

            if(p1_found>=0 && p2_found>=0)
            {
                break;
            }
        }

        if(p1_found==-1 || p2_found==-1)
        {
            throw std::runtime_error("There was an error in determining which individuals to cross.");
        }

        //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;
                pop_change->at(i).change_genotype(offset, pop_read->at(p).get_genotype(0));
            }
            else
            {
                chromo_choose=1;
                pop_change->at(i).change_genotype(offset, pop_read->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)
                {
                    pop_change->at(i).change_genotype(offset+j+1, pop_read->at(p).get_genotype(n_loci+j+1));
                }
                else if(chromo_choose==0)
                {
                    pop_change->at(i).change_genotype(offset+j+1, pop_read->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 round number "<<round<<".";
                    throw std::runtime_error(message.str());
                }
            }
        }
    }

    //Clear everything just in case
    pop_change=NULL;
    pop_read=NULL;

    return 0;
}

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

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

//Work out the recombination fraction between two loci
double get_recom(locus locus1, locus locus2)
{
    double dist;
    double recom_frac;

    if(locus1.get_chrom().compare(locus2.get_chrom())!=0)	//Check if loci are not on same chromosome
    {
        recom_frac=0.5;
    }
    else
    {
        //Get the mapping distance between the two loci in Morgans
        dist=(locus2.get_dist()-locus1.get_dist())/100;

        //Use the Haldane Equation to get recombination fraction
        recom_frac=0.5*(1-exp(-2*dist));
    }

    return recom_frac;
}

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

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

//Switches between chromosomes due to recombination
void chromo_switch(int* chromo_choose)
{
    if(*chromo_choose==0)
        *chromo_choose=1;
    else
        *chromo_choose=0;
}
