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


def get_time(path):
    
    no_space = path.replace(" ", "")
    day_index = no_space.find("Day")
    day_string = no_space[day_index+3:day_index+5]
    day_string = day_string.replace("_", "")
    return int(day_string)

def merge_plots(work_dir, study_collection):

    plot_name = "diameterPlot"
    plot_name = "histogramDiameter"
    plot_name = "rendering"
    reduction_factor = 1
    reduction_factor2 = 1

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

            image_paths = []
            # first sort time points
            days = []
            for eachTimePoint in eachMouse.timepoints:
                days.append(get_time(eachTimePoint.raw_data_path))
            if len(days) > 0:
                days, eachMouse.timepoints = (list(t) for t in zip(*sorted(zip(days, eachMouse.timepoints))))

            for eachTimePoint in eachMouse.timepoints:
                time_dir = mouse_dir + "/" + eachTimePoint.raw_data_path
                #base_path = "/analysis_analysis/plots/vessels/" + plot_name + ".png"
                base_path = "/analysis_analysis/histograms/" + plot_name + ".png"
                base_path = "/analysis_analysis//" + plot_name + ".tiff"
                if os.path.isfile(time_dir + base_path):
                    image_paths.append(time_dir + base_path)

            merge_path = mouse_dir + "/timeline_" + plot_name + ".jpg"
            merge_axis = 0
            print "Merging plots to: ", merge_path

            if len(image_paths) > 0:

                widths = []
                heights = []
                for path in image_paths:
                    im = Image.open(path)
                    widths.append(im.size[0])
                    heights.append(im.size[1])
                    im.close()
                total_width = int(0.5*sum(widths))
                max_height = int(0.8*max(heights))
                new_im = Image.new('RGB', (total_width, max_height))
                x_offset = 0
                for path in image_paths:
                    im = Image.open(path)
                    img2 = im.crop((int(0.25*im.size[0]), int(0.2*im.size[1]), int(0.75*im.size[0]), im.size[1]))
                    new_im.paste(img2, (x_offset, 0))
                    x_offset += int(0.5*im.size[0])
                    im.close()
                new_im.thumbnail((total_width/reduction_factor, max_height/reduction_factor))
                merge_paths.append(merge_path)
                new_im.save(merge_path)
        if len(merge_paths) > 0:
            widths = []
            heights = []
            for path in merge_paths:
                im = Image.open(path)
                widths.append(im.size[0])
                heights.append(im.size[1])
                im.close()
            max_width = max(widths)
            total_height = sum(heights)
            new_im = Image.new('RGB', (max_width, total_height))
            y_offset = 0
            for path in merge_paths:
                im = Image.open(path)
                new_im.paste(im, (0, y_offset))
                y_offset += im.size[1]
                im.close()
            new_im.thumbnail((max_width/reduction_factor2, total_height/reduction_factor2))
            new_im.save(study_dir + "/timeline_" + plot_name + ".jpg")


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)