######################################################
#
# Make a copy of this file and edit it to your liking
#
######################################################

# Which population to plot
species_name='Eubacterium_eligens_61678'
species_names = [species_name]*5
# Alternative: can access entire list using function:
#species_names = parse_midas_data.parse_good_species_list()

# Output filename (alter as needed)
filename = parse_midas_data.analysis_directory+species_name+'_timecourse.png'

####################################################################
#
#  Function controlling which mutations are colored
#
#      Returns: true for colored
#               false for not colored
#
####################################################################
def color_condition(population_idx, chromosome, location, gene_name, variant_type, times, freqs, depths):

    # don't use interpolated freqs yet.    

    # First two conditions require mutation to start at low frequency
    # i.e., to focus on "de novo" mutations
    idx_130 = (numpy.fabs(times-75)).argmin() 
    idx_82 = (numpy.fabs(times-82)).argmin()
    idx_70 = (numpy.fabs(times-70)).argmin() 
    idx_39 = (numpy.fabs(times-39)).argmin()
    idx_60 = (numpy.fabs(times-60)).argmin()
    idx_95 = (numpy.fabs(times-95)).argmin()
    if population_idx==0:
	condition = (freqs[-1]<0.1)*(freqs[idx_130]<0.1)*(freqs[idx_82]<0.5)
    elif population_idx==1:
	condition =(freqs[0]<0.50)*(freqs[-1]>0.5)*(freqs[idx_70]>0.7)
    elif population_idx==2:
	condition = (freqs[0]<0.50)*(freqs[idx_39]<0.3)*(freqs[2]>0.4)*(freqs[idx_60]<0.50)
    elif population_idx==3:
        condition = (freqs[0]<0.1)*(freqs[idx_95]<0.4)
    elif population_idx==4:
        condition = (freqs[0]<0.1)*(freqs[idx_95]<0.4)*(freqs[2]<0.1)

        #print chromosome, location, gene_name, variant_type
    else:
        condition = False
        
    return condition
