--- a/dpmix_plot.py	Fri Jul 26 12:51:13 2013 -0400
+++ b/dpmix_plot.py	Fri Sep 20 13:25:27 2013 -0400
@@ -3,8 +3,10 @@
 import os
 import sys
 import math
+
 import matplotlib as mpl
 mpl.use('PDF')
+from matplotlib.backends.backend_pdf import PdfPages
 import matplotlib.pyplot as plt
 from matplotlib.path import Path
 import matplotlib.patches as patches
@@ -226,18 +228,49 @@
     return vals, labels
 
 ################################################################################
+################################################################################
+################################################################################
+################################################################################
+
+def space_for_legend(plot_params):
+    space = 0.0
+
+    legend_states = plot_params['legend_states']
+    if legend_states:
+        ind_space = plot_params['ind_space']
+        ind_height = plot_params['ind_height']
+        space += len(legend_states) * (ind_space + ind_height) - ind_space
+
+    return space
+
+################################################################################
+
+def space_for_chroms(plot_params, chroms, individuals, data):
+    space_dict = {}
+
+    chrom_height = plot_params['chrom_height']
+    ind_space = plot_params['ind_space']
+    ind_height = plot_params['ind_height']
+
+    for chrom in chroms:
+        space_dict[chrom] = chrom_height
+
+        individual_count = 0
+        for individual in individuals:
+            if individual in data[chrom]:
+                individual_count += 1
+
+        space_dict[chrom] += individual_count * (ind_space + ind_height)
+
+    return space_dict
+
+################################################################################
 
 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, 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
-
+    ## populate chrom_len
     for chrom in chrom_len.keys():
         if chrom in fs_chrom_len:
             chrom_len[chrom] = fs_chrom_len[chrom]
@@ -245,135 +278,177 @@
     #check_chroms(chroms, chrom_len, input_dbkey)
     check_data(data, chrom_len, input_dbkey)
 
-    ## units below are inches
-    top_space = 0.10
-    chrom_space = 0.25
-    chrom_height = 0.25
-    ind_space = 0.10
-    ind_height = 0.25
+    ## plot parameters
+    plot_params = {
+        'plot_dpi':        300,
+        'page_width':     8.50,
+        'page_height':   11.00,
+        'top_margin':     0.10,
+        'bottom_margin':  0.10,
+        'chrom_space':    0.25,
+        'chrom_height':   0.25,
+        'ind_space':      0.10,
+        'ind_height':     0.25,
+        'legend_space':   0.10
+    }
 
-    total_height = 0.0
-
-    ## make a legend
-    ## only print out states that are
+    ## in the 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
+    plot_params['legend_states'] = legend_states
+
+    ## choose the correct make_state_rectangle method
+    if populations == 3:
+        plot_params['rectangle_method'] = make_state_rectangle_3pop
+    elif populations == 2:
+        plot_params['rectangle_method'] = make_state_rectangle_2pop
+
+    pdf_pages = PdfPages(output_file)
+
+	## generate a list of chroms for each page
+
+    needed_for_legend = space_for_legend(plot_params)
+    needed_for_chroms = space_for_chroms(plot_params, chroms, individuals, data)
 
-    for chrom in chroms:
-        if at_top:
-            total_height += (top_space + chrom_height)
-            at_top = False
+    chrom_space_per_page = plot_params['page_height']
+    chrom_space_per_page -= plot_params['top_margin'] + plot_params['bottom_margin']
+    chrom_space_per_page -= needed_for_legend + plot_params['legend_space']
+    chrom_space_per_page -= plot_params['chrom_space']
+
+    chroms_left = chroms[:]
+    pages = []
+
+    space_left = chrom_space_per_page
+    chrom_list = []
+
+    while chroms_left:
+        chrom = chroms_left.pop(0)
+        space_needed = needed_for_chroms[chrom] + plot_params['chrom_space']
+        if (space_needed > chrom_space_per_page):
+            print >> sys.stderr, 'Multipage chroms not yet supported'
+            sys.exit(1)
+
+		## sometimes 1.9 - 1.9 < 0 (-4.4408920985e-16)
+		## so, we make sure it's not more than a millimeter over
+        if space_left - space_needed > -0.04:
+            chrom_list.append(chrom)
+            space_left -= space_needed
         else:
-            total_height += (top_space + chrom_space + chrom_height)
-    
-        individual_count = 0
-        for individual in individuals:
-            if individual in data[chrom]:
-                individual_count += 1
-        total_height += individual_count * (ind_space + ind_height)
+            pages.append(chrom_list[:])
+            chrom_list = []
+            chroms_left.insert(0, chrom)
+            space_left = chrom_space_per_page
 
-    width = 7.5
-    height = math.ceil(total_height)
-
-    bottom = 1.0
-
-    fig = plt.figure(figsize=(width, height))
+    ############################################################################
 
-    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)
+    plot_dpi = plot_params['plot_dpi']
+    page_width = plot_params['page_width']
+    page_height = plot_params['page_height']
+    top_margin = plot_params['top_margin']
+    ind_space = plot_params['ind_space']
+    ind_height = plot_params['ind_height']
+    make_state_rectangle = plot_params['rectangle_method']
+    legend_space = plot_params['legend_space']
+    chrom_space = plot_params['chrom_space']
+    chrom_height = plot_params['chrom_height']
+
+    for page in pages:
+        fig = plt.figure(figsize=(page_width, page_height), dpi=plot_dpi)
+        bottom = 1.0 - (top_margin/page_height)
+
+        # print legend
+        if legend_states:
+            top = True
+            for state in sorted(legend_states):
+                if top:
+                    bottom -= ind_height/page_height
+                    top = False
+                else:
+                    bottom -= (ind_space + ind_height)/page_height
 
-            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
+                ax1 = fig.add_axes([0.0, bottom, 0.09, ind_height/page_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)
 
-    for_webb = False
+                ax2 = fig.add_axes([0.10, bottom, 0.88, ind_height/page_height], frame_on=False)
+                plt.axis('off')
+                plt.text(0.0, 0.5, state2name[state], fontsize=10, ha='left', va='center')
+
+            bottom -= legend_space/page_height
 
-    for chrom in chroms:
-        length = chrom_len[chrom]
-        vals, labels = tick_foo(0, length)
+        # print chroms
+        top = True
+        for chrom in page:
+            length = chrom_len[chrom]
+            vals, labels = tick_foo(0, length)
 
-        if at_top:
-            bottom -= (top_space + chrom_height)/height
-            at_top = False
-        else:
-            bottom -= (top_space + chrom_space + chrom_height)/height
+            if top:
+                bottom -= chrom_height/page_height
+                top = False
+            else:
+                bottom -= (chrom_space + chrom_height)/page_height
 
-        if not for_webb:
-            ax = fig.add_axes([0.0, bottom, 1.0, chrom_height/height])
+            ax = fig.add_axes([0.0, bottom, 1.0, chrom_height/page_height])
             plt.axis('off')
             plt.text(0.5, 0.5, chrom, fontsize=14, ha='center')
 
-        individual_count = 0
-        for individual in individuals:
-            if individual in data[chrom]:
-                individual_count += 1
+            individual_count = 0
+            for individual in individuals:
+                if individual in data[chrom]:
+                    individual_count += 1
 
-        i = 0
-        for individual in individuals:
-            if individual in data[chrom]:
-                i += 1
+            i = 0
+            for individual in individuals:
+                if individual in data[chrom]:
+                    i += 1
+                    bottom -= (ind_space + ind_height)/page_height
 
-                bottom -= (ind_space + ind_height)/height
-                if not for_webb:
-                    # [left, bottom, width, height]
-                    ax1 = fig.add_axes([0.0, bottom, 0.09, ind_height/height])
+                    ax1 = fig.add_axes([0.0, bottom, 0.09, ind_height/page_height])
                     plt.axis('off')
                     plt.text(1.0, 0.5, individual, fontsize=10, ha='right', va='center')
-                # [left, bottom, width, height]
-                ax2 = fig.add_axes([0.10, bottom, 0.88, ind_height/height], frame_on=False)
-                ax2.set_xlim(0, length)
-                ax2.set_ylim(0, 1)
-                if i != individual_count:
-                    plt.axis('off')
-                else:
-                    if not for_webb:
+
+                    ax2 = fig.add_axes([0.10, bottom, 0.88, ind_height/page_height], frame_on=False)
+                    ax2.set_xlim(0, length)
+                    ax2.set_ylim(0, 1)
+
+                    if i != individual_count:
+                        plt.axis('off')
+                    else:
                         ax2.tick_params(top=False, left=False, right=False, labelleft=False)
                         ax2.set_xticks(vals)
                         ax2.set_xticklabels(labels)
-                    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)
+
+                    for p1, p2, state in sorted(data[chrom][individual]):
+                        for patch in make_state_rectangle(p1, p2, state, chrom, individual):
+                            ax2.add_patch(patch)
 
-    plt.savefig(output_file)
+                    # extend last state to end of chrom
+                    if p2 < length:
+                        for patch in make_state_rectangle(p2, length, state, chrom, individual):
+                            ax2.add_patch(patch)
+
+
+        pdf_pages.savefig(fig)
+        plt.close(fig)
+
+    pdf_pages.close()
 
 ################################################################################
 
 if __name__ == '__main__':
-    input_dbkey, input_file, output_file, galaxy_data_index_dir = sys.argv[1:5]
-    make_dpmix_plot(input_dbkey, input_file, output_file, galaxy_data_index_dir)
+    make_dpmix_plot('loxAfr3', 'output.dat', 'output2_files/picture.pdf', '/scratch/galaxy/home/oocyte/galaxy_oocyte/tool-data', state2name={0: 'heterochromatin', 1: 'reference', 2: 'asian'}, populations=2)
+#    input_dbkey, input_file, output_file, galaxy_data_index_dir = sys.argv[1:5]
+#    make_dpmix_plot(input_dbkey, input_file, output_file, galaxy_data_index_dir)
     sys.exit(0)
 
 ## notes