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

linked_strs = ["1120974.3.peg.1084:147, 1120974.3.peg.1083:112, 1120974.3.peg.1085:67, 1120974.3.peg.1086:47, 1120974.3.peg.1088:24, 1120974.3.peg.1089:22, 1120974.3.peg.1080:17, 1499682.3.peg.2843:16, 1120974.3.peg.1090:16, 1120974.3.peg.1079:16, 1120974.3.peg.1087:14, 1120974.3.peg.1104:13, 1120974.3.peg.1091:12, 1120974.3.peg.1105:10, 1499682.3.peg.2823:8, 1120974.3.peg.1081:8, 1120974.3.peg.1026:8, 1499682.3.peg.2818:7, 1499682.3.peg.2847:6, 1120974.3.peg.1025:6, 1120974.3.peg.1021:6, 1499682.3.peg.2850:5, 1120974.3.peg.1112:5, 1120974.3.peg.1106:5, 1499682.3.peg.2851:4, 1499682.3.peg.2816:4, 1120974.3.peg.1023:4, 1499682.3.peg.2848:3, 1499682.3.peg.2824:3, 1499682.3.peg.1783:3, 1120974.3.peg.1109:3, 1120974.3.peg.1019:3, 1499682.3.peg.364:2, 1499682.3.peg.2853:2, 1499682.3.peg.2814:2, 1499682.3.peg.2813:2, 1499682.3.peg.1929:2, 1499682.3.peg.1782:2, 1120974.3.peg.705:2, 1120974.3.peg.2744:2, 1120974.3.peg.2741:2, 1120974.3.peg.229:2, 1120974.3.peg.1607:2, 1120974.3.peg.1539:2, 1120974.3.peg.1280:2", 
    "1120974.3.peg.19:132, 1120974.3.peg.91:32, 1120974.3.peg.2376:23, 1120974.3.peg.45:18, 1120974.3.peg.116:18, 1120974.3.peg.92:10, 1120974.3.peg.2377:10, 1120974.3.peg.2375:10, 1120974.3.peg.24:9, 1120974.3.peg.2378:6, 1120974.3.peg.7:5, 1120974.3.peg.8:4, 1120974.3.peg.2379:4, 1120974.3.peg.176:4, 1499682.3.peg.897:3, 1120974.3.peg.746:3, 1120974.3.peg.6:3, 1120974.3.peg.47:3, 1499682.3.peg.896:2, 1499682.3.peg.2289:2, 1499682.3.peg.1238:2, 1120974.3.peg.504:2, 1120974.3.peg.502:2, 1120974.3.peg.360:2, 1120974.3.peg.3267:2, 1120974.3.peg.3:2, 1120974.3.peg.2428:2, 1120974.3.peg.1746:2, 1120974.3.peg.1406:2", 
    "1499682.3.peg.802:289, 1120974.3.peg.1039:23, 1499682.3.peg.1543:14, 1120974.3.peg.2947:9, 1120974.3.peg.2945:7, 1499682.3.peg.365:6, 1499682.3.peg.1544:6, 1499682.3.peg.1546:5, 1499682.3.peg.373:4, 1499682.3.peg.1548:4, 1499682.3.peg.367:3, 1499682.3.peg.364:3, 1499682.3.peg.1551:3, 1499682.3.peg.1550:3, 1120974.3.peg.507:3, 1120974.3.peg.495:3, 1120974.3.peg.341:3, 1120974.3.peg.2951:3, 1120974.3.peg.2943:3, 1120974.3.peg.2866:3, 1499682.3.peg.785:2, 1499682.3.peg.587:2, 1499682.3.peg.453:2, 1499682.3.peg.2466:2, 1499682.3.peg.2243:2, 1499682.3.peg.2238:2, 1499682.3.peg.2234:2, 1499682.3.peg.2084:2, 1499682.3.peg.1607:2, 1499682.3.peg.1606:2, 1499682.3.peg.1555:2, 1499682.3.peg.1540:2, 1120974.3.peg.340:2, 1120974.3.peg.338:2, 1120974.3.peg.3200:2, 1120974.3.peg.2940:2, 1120974.3.peg.2919:2, 1120974.3.peg.2918:2, 1120974.3.peg.2421:2, 1120974.3.peg.1955:2, 1120974.3.peg.1797:2, 1120974.3.peg.1440:2, 1120974.3.peg.1087:2, 1120974.3.peg.1086:2"]

# Which population to plot
species_names = ['Bacteroides_caccae_53434']

target_genes = ['997873.3.peg.3761','997873.3.peg.3760','997873.3.peg.3793']

linked_genes = []
for linked_str in linked_strs:
    linked_genes.append([])
    for item in linked_str.split(","):
        subitems = item.split(":")
        gene_name = subitems[0].strip()
        count = long(subitems[1])
        if count>=10:
            linked_genes[-1].append(gene_name)
    linked_genes[-1] = set(linked_genes[-1])

# Output filename (alter as needed)
filename = parse_midas_data.analysis_directory+'Bacteroides_caccae_gene_gain_timecourse.png'

####################################################################
#
#  Function controlling which mutations are colored
#
#      Returns: true for colored
#               false for not colored
#
####################################################################
def color_condition(population_idx, gene_name, times, gene_copynums, marker_coverages):
        
    # don't use interpolated freqs yet.    
    
    #condition = (gene_name in linked_genes[population_idx])
    condition = (gene_name in target_genes)
 
    initial_idx = numpy.fabs(times-47).argmin()
    final_idx = numpy.fabs(times-73).argmin()
 
    #condition = ((gene_copynums>0.5)*(gene_copynums<2))[final_idx] #*((gene_copynums<0.1))[initial_idx]
 
    return condition
