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

# Which population to plot
#species_names = ["Alistipes_onderdonkii_55464"]

species_names = parse_midas_data.parse_good_species_list()
 
# Output filename
filename = parse_midas_data.analysis_directory+'new_antibiotic_snps_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):
    
    condition = False
    
    pre_antibiotic_idxs = (times<sample_time_map[parse_timecourse_data.highcoverage_antibiotic])
    
    post_antibiotic_idxs = (times>=sample_time_map[parse_timecourse_data.highcoverage_antibiotic])
    
    if (pre_antibiotic_idxs.sum() > 1) and (post_antibiotic_idxs.sum() > 0):
        
        condition = (freqs[pre_antibiotic_idxs]<0.05).all()
        condition = condition and (freqs[post_antibiotic_idxs]>0.5).any()
            
    if condition:
        
        items = [chromosome, location, gene_name, variant_type]
        print_str = ", ".join([str(item) for item in items])
        sys.stderr.write("%s\n" % print_str)
    
    return condition
