#!/usr/bin/perl 

#############
#usage stuff#
#############

if (@ARGV != 3) {
    die "\nUsage: perl WrightFisher.pl Pop_size #Generations Fitness_dominant\n\n";
}


############################################################################################
#Get the input variables from the first line.
############################################################################################
$N = shift;
$Gen = shift;
$fit_dom = shift;
#Intialize number of times allel 2 is fixed and the average of generations
$num_2_fix = 0;
$time_2_fix = 0;

# Do 100 runs...
for($num_runs=1;$num_runs<=100;$num_runs++){

#################################################################################
#Initialize Population.  Allel 2 is the rare allel (1 out of 2N). The Population#
#is set up in a 2D array: $Pop[Individual][1st or 2nd allel]                    #
#################################################################################
    @Pop = ();
    $Pop[0][0]=2; 
    $Pop[0][1]=1;
    for($i=1;$i<$N;$i++){
	$IntPop[$i][0] = $Pop[$i][1] = 1;
    }
    
    #Assign Fitness Values ($fit_dom comes from command line)
    $fit_1_1 = 1;
    $fit_1_2 = $fit_dom;
    $fit_2_2 = $fit_dom;

    @Prob_Mat = ();    
    #Assign Random Picking Matrix (Prob_Mat)
    $j=0;
    for($i=0;$i<$N;$i++){
	if($Pop[$i][0]==2){ 
	   if($Pop[$i][1]==2){
	       #If ind. is homo. for rare allele add fitness number of times
	       for($q=0;$q<$fit_2_2;$q++){
		   $Prob_Mat[$j] = $i;
		   $j++;
	       }
	   }
	   else{
	       #Ind. het
	       for($q=0;$q<$fit_1_2;$q++){
                   $Prob_Mat[$j] = $i;
                   $j++;
	       }
	   }
       }
	else{
	    #Ind. het
	    if($Pop[$i][1]==2){
		for($q=0;$q<$fit_1_2;$q++){
		    $Prob_Mat[$j] = $i;
		    $j++;
		}
	    }
	    #Ind. homo. for wt allele
	    else{
		for($q=0;$q<$fit_1_1;$q++){
		    $Prob_Mat[$j] = $i;
		    $j++;
		}
	    }
	}
    }

    #Do $Gen number of generations
    for($num_gen=1;$num_gen<=$Gen;$num_gen++){

	#Randomly make a new Population, using the Prob_Mat to make the picks
	for($i=0;$i<$N;$i++){
	    $ran = int rand @Prob_Mat; #Random number between 0 size of picking matrix
	    $ran2 = int rand 2;
	    $New_Pop[$i][0]=$Pop[$Prob_Mat[$ran]][$ran2];
	    $ran = int rand @Prob_Mat;
            $ran2 = int rand 2;
            $New_Pop[$i][1]=$Pop[$Prob_Mat[$ran]][$ran2];
	}

	#Make New_Pop the current Pop ($Pop)
	@Pop = @New_Pop;
	@New_Pop = ();
	@Prob_Mat = ();
        #Assign Random Picking Matrix (Prob_Mat)
	$count_2 = 0;
	$j=0;
	for($i=0;$i<$N;$i++){
	    if($Pop[$i][0]==2){
		if($Pop[$i][1]==2){
		    #If ind. is homo. for rare allele add fitness number of times
		    $count_2 = $count_2 + 2;
		    for($q=0;$q<$fit_2_2;$q++){
			$Prob_Mat[$j] = $i;
			$j++;
		    }
		}
		else{
		    #Ind. Het
		    $count_2++;
		    for($q=0;$q<$fit_1_2;$q++){
			$Prob_Mat[$j] = $i;
			$j++;
		    }
		}
	    }
	    else{
		if($Pop[$i][1]==2){
		    #Ind. Het
		    $count_2++;
		    for($q=0;$q<$fit_1_2;$q++){
			$Prob_Mat[$j] = $i;
			$j++;
		    }
		}
		else{
		    #Ind. homo. for wt allele
		    for($q=0;$q<$fit_1_1;$q++){
			$Prob_Mat[$j] = $i;
			$j++;
		    }
		}
	    }
	}
	
	######################################################
	#Calculate allel 2 frequency and check for fixation  #
	#If there's fixation, make num_gen greater Gen = stop#
	######################################################

	if($count_2 == 0){
	    print "$num_runs: Allel 1 fixed\t$num_gen\n";
	    $num_gen = $Gen+1;
	}
	elsif($count_2 == 2*$N){
	    print "$num_runs: Allel 2 fixed\t$num_gen\n";
            $num_2_fix++;
	    $time_2_fix = $time_2_fix + $num_gen;
	    $num_gen = $Gen+1;
	}
	if($num_gen == $Gen){
	    $freq_2 = $count_2/(2*$N);
	    print "$num_runs: No Fixation: frequency of allel 2 = $freq_2\n";
	}
    }
}

#Print out overall frequency
$freq = $num_2_fix / 100;
$avg_gen = $time_2_fix/$num_2_fix;
print "Allel 2 fixed $num_2_fix / 100 = $freq\tAverage # of generations: $avg_gen\n";