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

# Which population to plot
species_names = ["Alistipes_finegoldii_56071", "Alistipes_sp_60764", "Eubacterium_rectale_56927", "Bacteroides_vulgatus_57955",  "Bacteroides_massiliensis_44749", "Bacteroides_coprocola_61586"]
 
# Output filename
filename = parse_midas_data.analysis_directory+'example_fourgamete_figure.pdf'
stats_filename = parse_midas_data.analysis_directory+'example_fourgamete_stats_figure.pdf'
output_filename = parse_midas_data.analysis_directory+'example_fourgametefigure_snps.txt'

# Load fourgamete files:
file = open('new_cross_species_snp_barcode_statistics_nonshared.txt',"r")

fourgamete_location_map = {}

for line in file:    
    items = line.split()
    species_name = items[0]
    lmin = long(items[1])
    testable = long(items[2])
    onegametes = long(items[3])
    twogametes = long(items[4])
    threegametes = long(items[5])
    testable_fourgametes = long(items[6])
    fourgametes = long(items[7])
      
    if lmin<100:
    	continue
    	  
    if species_name not in fourgamete_location_map:
    	fourgamete_location_map[species_name] = {}
        
    if len(items)>7:
    	
    	fourgamete_items = items[8:]
        for item in fourgamete_items:
        
            subitems = item.split(",")
            focal_snp_items = subitems[0].strip().split("|")
            target_snp_items = subitems[1].strip().split("|")
                
            ns = numpy.array([long(subitem) for subitem in subitems[2:]])
            min_n = ns.min()
    
            focal_contig = focal_snp_items[0]
            focal_location = long(focal_snp_items[1])
        
            target_contig = target_snp_items[0]
            target_location = long(target_snp_items[1])
        
            if True: #min_n>=8 and fabs(target_location-focal_location)>100:
            	if (focal_contig, focal_location) not in fourgamete_location_map[species_name]:
            		
                    fourgamete_location_map[species_name][(focal_contig, focal_location)] = set()
        	    fourgamete_location_map[species_name][(focal_contig, focal_location)].add((target_contig, target_location))
        	    
        	    if (target_contig, target_location) not in fourgamete_location_map[species_name]:
        	        fourgamete_location_map[species_name][(target_contig, target_location)] = set()   
        	    
        	    fourgamete_location_map[species_name][(target_contig, target_location)].add((focal_contig, focal_location))
        		
####################################################################
#
#  Function controlling which mutations are colored
#
#      Returns: true for colored
#               false for not colored
#
####################################################################
def color_rule_in_condition(population_idx, chromosome, location, gene_name, variant_type, times, freqs, depths):
    
    condition = (chromosome,location) in fourgamete_location_map[species_names[population_idx]]
        
    return condition
    
def color_rule_out_condition(population_idx, chromosome, location, gene_name, variant_type, times, freqs, depths):
    
    rule_in_condition = ((chromosome,location) in fourgamete_location_map[species_names[population_idx]])
    
    if rule_in_condition:
        print (chromosome, location), "->",  sorted(fourgamete_location_map[species_names[population_idx]][(chromosome,location)])
    
    condition = not rule_in_condition
    
    return condition
