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

# open cluster files
cluster_file = 'Bacteroides_vulgatus_clusters.txt'
file = open(cluster_file,"r")
clusters = []
for line in file:
    items = line.split()
    cluster = set()
    for item in items:
        subitems = item.split("|")
        contig = subitems[0]
        location = long(subitems[1])
        cluster.add((contig,location))
    clusters.append(cluster)

if len(clusters)>10:
    clusters = clusters[0:10]
    
# Which population to plot
species_names = ['Bacteroides_vulgatus_57955']*len(clusters)

# 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_clusters_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
    
    cluster = clusters[population_idx]
    
    condition = ((chromosome,location) in cluster)
    
    if population_idx==2 and condition:
        
        if freqs[-1]>0.4:
            print "*%s|%d" % (chromosome, location)
        else:
            print "%s|%d" % (chromosome, location)
        
    return condition
