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

# Which population to plot
species_names = ['Faecalibacterium_prausnitzii_57453','Faecalibacterium_prausnitzii_57453','Faecalibacterium_prausnitzii_57453','Faecalibacterium_prausnitzii_57453','Faecalibacterium_prausnitzii_57453','Faecalibacterium_prausnitzii_57453'] 

# 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+'Faecalibacterium_prausnitzii_57453_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
    
    initial_idxs = numpy.logical_or((times==sample_time_map[parse_timecourse_data.highcoverage_start_2]),(times==sample_time_map[parse_timecourse_data.highcoverage_start_2]))
    
    final_idxs = (times==sample_time_map[parse_timecourse_data.highcoverage_end])
    antibiotic_idxs = (times==sample_time_map[parse_timecourse_data.highcoverage_antibiotic])
    

    # don't use interpolated freqs yet.    

    condition = False
    
    if initial_idxs.sum()<0.5 or antibiotic_idxs.sum()<0.5 or final_idxs.sum()<0.5:
        return condition
    
    
    if population_idx==0: # get the main clade. hits high freq during pre-abx usage, then goes back down. 
    
        condition = (freqs[initial_idxs]>0.7).any() and (freqs[antibiotic_idxs]<0.2).any() and (freqs[final_idxs]<0.2)
    
    elif population_idx==1: 
    
        condition = (freqs[initial_idxs]>0.7).any() and (freqs[antibiotic_idxs]<0.2).any() and (freqs[final_idxs]>0.5)
    
    elif population_idx==2:
         
         #condition = (freqs[initial_idxs]<0.2).any() and (freqs[antibiotic_idxs]<0.2).any() and (freqs[final_idxs]>0.5)
        
         condition = (freqs[initial_idxs]>0.7).any() and (freqs[antibiotic_idxs]>0.5).any()
         
    elif population_idx==3:
         
         condition = (freqs[0]<0.1) and (freqs[initial_idxs]<0.1).any() and (freqs[antibiotic_idxs]>0.7).any() and (freqs[final_idxs]>0.5)
         
         if condition:
             print depths
         
    elif population_idx==4:
         
         gene_number = long(gene_name.split(".")[-1])
    
         # things that should go low after abx
         #condition = (gene_number>600)*(gene_number<=605)
         
         condition = (gene_number>=1461)*(gene_number<=1465)
         
    elif population_idx==5:
         
         gene_number = long(gene_name.split(".")[-1])
    
         #condition = (gene_number>606)*(gene_number<=610)
         
         #condition = condition and (freqs[antibiotic_idxs]<0.3).any()
         condition = (chromosome=='NZ_DS483491') and (location==29134 or location==29191 or location==29029)
    
    else:
        pass       
            

    if condition:
        print chromosome, location, gene_name, variant_type
        
    return condition
