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

# Which population to plot
species_name='Odoribacter_splanchnicus_62174'
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()
    if population_idx==0:
	condition = (freqs[-1]<0.1)*(freqs[idx_130]<0.1)
    elif population_idx==1:
	condition =(freqs[0]<0.50)*(freqs[idx_39]<0.3)*(freqs[2]>0.4)*(freqs[-1]>0.5)
    elif population_idx==2:
	condition = (freqs[idx_130]>0.75)*(freqs[1]>0.75)  
    elif population_idx==3:
	condition = (freqs[idx_39]>0.5)*(freqs[2]<0.4)*(freqs[-1]>0.5)
    elif population_idx==4:
	condition = (freqs[idx_39]>0.5)*(freqs[2]<0.4)*(freqs[-1]<0.5)

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