import os
import pickle
from PIL import Image
from stack3d.study.components import StudyData, MouseData, TimepointData


def merge_plots(work_dir, study_collection):

    for eachStudy in study_collection:
        print "Merging Study: ", eachStudy.raw_data_path
        study_dir = work_dir + "/" + eachStudy.raw_data_path
        for eachMouse in eachStudy.mice:
            print "Merging Mouse: ", eachMouse.raw_data_path
            mouse_dir = study_dir + "/" + eachMouse.raw_data_path

            image_paths = []
            for eachTimePoint in eachMouse.timepoints:
                time_dir = mouse_dir + "/" + eachTimePoint.raw_data_path
                base_path = "/analysis_analysis/plots/vessels/diameterPlot.png"
                if os.path.isfile(time_dir + base_path):
                    image_paths.append(time_dir + base_path)

            images = map(Image.open, image_paths)
            merge_path = mouse_dir + "/timeline_diameterPlot.jpg"
            merge_axis = 0
            print "Merging plots to: ", merge_path
            if len(images) > 0:
                widths, heights = zip(*(i.size for i in images))
                if merge_axis == 0:
                    total_width = sum(widths)
                    max_height = max(heights)
                    new_im = Image.new('RGB', (total_width, max_height))
                    x_offset = 0
                    for im in images:
                        new_im.paste(im, (x_offset, 0))
                        x_offset += im.size[0]
                    new_im.save(merge_path, optimize=True, quality=5)
                else:
                    max_width = max(widths)
                    total_height = sum(heights)
                    new_im = Image.new('RGB', (max_width, total_height))
                    y_offset = 0
                    for im in images:
                        new_im.paste(im, (0, y_offset))
                        y_offset += im.size[1]
                    new_im.save(merge_path, optimize=True, quality=5)


work_dir = "/scratch/jgrogan/stack-working/study/"
f = open(work_dir + "/study_collection.p", 'r')
study_collection = pickle.load(f)
merge_plots(work_dir, study_collection)