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

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

# 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_61481_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_1]),(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()<1.5 or antibiotic_idxs.sum()<0.5 or final_idxs.sum()<0.5:
        return condition
    
    if not (freqs[initial_idxs]<0.5).all():
        return condition       
    
    if population_idx==0: # get the main clade. hits high freq during ABX usage, then goes back down. 
    
        condition = (freqs[antibiotic_idxs]>0.8).any() and (freqs[final_idxs]<0.5).any()
    
    elif population_idx==1: 
    
        condition = (freqs[antibiotic_idxs]>0.8).any() and (freqs[final_idxs]>0.6).any() and (freqs[final_idxs]<0.9).any()
    
    elif population_idx==2:
         #condition = (chromosome=='NZ_GG697157')
         gene_name_items = gene_name.split(".")
         gene_number = long(gene_name_items[-1])
         condition = (gene_number>=2660)*(gene_number<=2702)
         #condition = (freqs[antibiotic_idxs]<0.6).any() and (freqs[antibiotic_idxs]>0.05).any() and (freqs[final_idxs]>0.6).any()
         #condition = condition and (len(gene_name)==len('411483.3.peg.26XX')) and (gene_name.startswith('411483.3.peg.266') or gene_name.startswith('411483.3.peg.267') or gene_name.startswith('411483.3.peg.268'))
         #condition = (len(gene_name)==len('411483.3.peg.26XX')) and (gene_name.startswith('411483.3.peg.266') or gene_name.startswith('411483.3.peg.267') or gene_name.startswith('411483.3.peg.268'))
         
    #elif population_idx==3:
    #     
         #condition = (freqs[antibiotic_idxs]<0.05).any() and (freqs[final_idxs]>0.8).any()
    
    else:
        pass       
            

    if condition and population_idx!=0:
        print chromosome, location, gene_name, variant_type, depths[-1], freqs[-1]
        
    return condition
