--- a/dpmix_plot.py	Mon Jun 03 12:29:29 2013 -0400
+++ b/dpmix_plot.py	Mon Jul 15 10:47:35 2013 -0400
@@ -37,6 +37,7 @@
     individuals = []
     data = {}
     chrom_len = {}
+    used_states = []
 
     with open(input_file) as fh:
         for line in fh:
@@ -47,6 +48,9 @@
                 p1, p2, state = map(int, elems[1:4])
                 id = elems[4]
 
+                if state not in used_states:
+                    used_states.append(state)
+
                 if chrom not in chroms:
                     chroms.append(chrom)
 
@@ -60,7 +64,7 @@
                 if p2 > chrom_len.setdefault(chrom, 0):
                     chrom_len[chrom] = p2
 
-    return chroms, individuals, data, chrom_len
+    return chroms, individuals, data, chrom_len, used_states
 
 def check_chroms(chroms, chrom_len, dbkey):
     error = 0
@@ -112,15 +116,43 @@
     patch2 = make_rectangle(p1, p2, top_color, bottom=0.5)
     return [patch1, patch2]
 
-def make_state_rectangle(p1, p2, state, chrom, individual):
+def make_state_rectangle_2pop(p1, p2, state, chrom, individual):
+    p1_color = 'r'
+    p2_color = 'g'
+    heterochromatin_color = '#c7c7c7'
+
     if state == 0:
-        return [ make_rectangle(p1, p2, 'r') ]
+        return [ make_rectangle(p1, p2, heterochromatin_color) ]
     elif state == 1:
-        return make_split_rectangle(p1, p2, 'r', 'g')
+        return [ make_rectangle(p1, p2, p1_color) ]
     elif state == 2:
-        return [ make_rectangle(p1, p2, 'g') ]
+        return [ make_rectangle(p1, p2, p2_color) ]
     elif state == 3:
-        return [ make_rectangle(p1, p2, '#c7c7c7') ]
+        return make_split_rectangle(p1, p2, p1_color, p2_color)
+    else:
+        print >> sys.stderr, "Unknown state: {0}: {1} {2} {3} {4}".format(state, chrom, p1, p2, state, individual)
+        sys.exit(1)
+
+def make_state_rectangle_3pop(p1, p2, state, chrom, individual):
+    p1_color = 'r'
+    p2_color = 'g'
+    p3_color = 'b'
+    heterochromatin_color = '#c7c7c7'
+
+    if state == 0:
+        return [ make_rectangle(p1, p2, heterochromatin_color) ]
+    if state == 1:
+        return [ make_rectangle(p1, p2, p1_color) ]
+    if state == 2:
+        return [ make_rectangle(p1, p2, p2_color) ]
+    if state == 3:
+        return [ make_rectangle(p1, p2, p3_color) ]
+    if state == 4:
+        return make_split_rectangle(p1, p2, p1_color, p2_color)
+    if state == 5:
+        return make_split_rectangle(p1, p2, p1_color, p3_color)
+    if state == 6:
+        return make_split_rectangle(p1, p2, p2_color, p3_color)
     else:
         print >> sys.stderr, "Unknown state: {0}: {1} {2} {3} {4}".format(state, chrom, p1, p2, state, individual)
         sys.exit(1)
@@ -195,9 +227,16 @@
 
 ################################################################################
 
-def make_dpmix_plot(input_dbkey, input_file, output_file, galaxy_data_index_dir):
+def make_dpmix_plot(input_dbkey, input_file, output_file, galaxy_data_index_dir, state2name=None, populations=3):
     fs_chrom_len = build_chrom_len_dict(input_dbkey, galaxy_data_index_dir)
-    chroms, individuals, data, chrom_len = parse_input_file(input_file)
+    chroms, individuals, data, chrom_len, used_states = parse_input_file(input_file)
+
+    if populations == 3:
+        make_state_rectangle = make_state_rectangle_3pop
+    elif populations == 2:
+        make_state_rectangle = make_state_rectangle_2pop
+    else:
+        pass
 
     for chrom in chrom_len.keys():
         if chrom in fs_chrom_len:
@@ -214,7 +253,26 @@
     ind_height = 0.25
 
     total_height = 0.0
-    at_top = True
+
+    ## make a legend
+    ## only print out states that are
+    ##   1) in the data
+    ##    - AND -
+    ##   2) in the state2name map
+    ## here, we only calculate the space needed
+    legend_states = []
+    if state2name is not None:
+        for state in used_states:
+            if state in state2name:
+                legend_states.append(state)
+
+    if legend_states:
+        total_height += len(legend_states) * (ind_space + ind_height)
+        total_height += (top_space - ind_space)
+        at_top = False
+    else:
+        at_top = True
+
     for chrom in chroms:
         if at_top:
             total_height += (top_space + chrom_height)
@@ -235,7 +293,29 @@
 
     fig = plt.figure(figsize=(width, height))
 
-    at_top = True
+    if legend_states:
+        at_top = True
+        for state in sorted(legend_states):
+            if at_top:
+                bottom -= (top_space + ind_height)/height
+                at_top = False
+            else:
+                bottom -= (ind_space + ind_height)/height
+            ## add code here to draw legend
+            # [left, bottom, width, height]
+            ax1 = fig.add_axes([0.0, bottom, 0.09, ind_height/height])
+            plt.axis('off')
+            ax1.set_xlim(0, 1)
+            ax1.set_ylim(0, 1)
+            for patch in make_state_rectangle(0, 1, state, 'legend', state2name[state]):
+                ax1.add_patch(patch)
+
+            ax2 = fig.add_axes([0.10, bottom, 0.88, ind_height/height], frame_on=False)
+            plt.axis('off')
+            plt.text(0.0, 0.5, state2name[state], fontsize=10, ha='left', va='center')
+    else:
+        at_top = True
+
     for_webb = False
 
     for chrom in chroms:
@@ -283,6 +363,7 @@
                     else:
                         plt.axis('off')
                 for p1, p2, state in sorted(data[chrom][individual]):
+                    #for patch in make_state_rectangle(p1, p2, state, chrom, individual):
                     for patch in make_state_rectangle(p1, p2, state, chrom, individual):
                         ax2.add_patch(patch)
 
@@ -295,3 +376,6 @@
     make_dpmix_plot(input_dbkey, input_file, output_file, galaxy_data_index_dir)
     sys.exit(0)
 
+## notes
+# 1) pass in a state to name mapping
+# 2) only print out names for states which exist in the data, and are in the state to name mapping