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

# Which population to plot
species_names = parse_midas_data.parse_good_species_list()

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

####################################################################
#
#  Function controlling which mutations are colored
#
#      Returns: true for colored
#               false for not colored
#
####################################################################
def color_condition(population_idx, gene_name, times, gene_copynums, marker_coverages):
        
    # don't use interpolated freqs yet.    
    
    start_idxs = (times==sample_time_map[parse_timecourse_data.highcoverage_start_2])
    end_idxs = (times==sample_time_map[parse_timecourse_data.highcoverage_end])
    
    if (start_idxs.sum() > 0) and (end_idxs.sum() > 0):
        
        one_present = ((gene_copynums>=0.5)*(gene_copynums<=2))[ end_idxs].any()
        one_absent =  (gene_copynums<=0.05)[start_idxs].any()   
        
        condition = one_present and one_absent
        
    else:
        condition = False
    
    # First two conditions require mutation to start at low frequency
    # i.e., to focus on "de novo" mutations
    #if gene_name in antibiotic_resistance_genes:
    #    condition=True
    #else:
    #    condition=False
    

    return condition
