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

# Which population to plot
species_names = ['Bacteroides_vulgatus_57955']*2

# 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+'Bacteroides_vulgatus_57955_linkage_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):

    #if gene_name not in core_genes:
        #return False

    # don't use interpolated freqs yet.    
    
    # First two conditions require mutation to start at low frequency
    # i.e., to focus on "de novo" mutations
    
    condition = False
    
    if population_idx==0:
       
        if not ((154 in times) and (130 in times)):
            return False
            
        # otherwise continue
        
        idx_154 = numpy.fabs(times-154).argmin()
        idx_130 = numpy.fabs(times-130).argmin()
       
        # Highest freq guys at last timepoint
        condition = (freqs[idx_154] > 0.6)*(freqs[idx_130]<0.5)
        
    elif population_idx==1:
    
        if not ((154 in times) and (130 in times)):
            return False
            
        # otherwise continue
        
        idx_154 = numpy.fabs(times-154).argmin()
        idx_130 = numpy.fabs(times-130).argmin()
       
        # Highest guys at second-to-last timepoint
        condition = (freqs[idx_130]>0.5)*(freqs[idx_154]>0.4)*(freqs[idx_154]<0.6)
        
        
    return condition
