from PIL import Image
import glob
import os

from plot_helper import PlotHelper


class HistogramAnalysis:

    def __init__(self, image_dir, fname_stress, centroids=None):
        self.fname_stress = fname_stress
        self.image_dir = image_dir
        self.avg_hists = self.find_avg_histograms()
        self.plot_hists()

    def load_images(self):
        os.chdir(self.image_dir)

        for infile in glob.glob("*.jpg"):
            yield (infile, Image.open(infile))

    def bin_stress(self, stress, centroid=None):
        if not centroid is None:
            return "centroid" + str(centroid)
        if stress < 0.5:
            return "LOW"
        else:
            return "HIGH"

    def find_avg_histograms(self):

        res = dict()

        for (fname, img) in self.load_images():
            if fname in self.fname_stress:
                stress = self.fname_stress[fname]["stress"]
                centroid = self.fname_stress[fname]["centroid"]
                bin = self.bin_stress(stress, centroid)
                hist = img.histogram()

                if bin in res:
                    res[bin].append(hist)
                else:
                    res[bin] = [hist]

        comp = dict()

        def avg_hists(hists):
            sum_hist = None

            for hist in hists:
                if sum_hist:
                    for i in range(len(hist)):
                        sum_hist[i] += hist[i]
                else:
                    sum_hist = hist

            n = len(hists)

            sum_hist = [x / n for x in sum_hist]
            return sum_hist

        for bin in res.keys():
            comp[bin] = avg_hists(res[bin])

        print comp.keys()
        return comp

    def get_avg_hists(self):
        return self.avg_hists

    def plot_hists(self):

        ph = PlotHelper()

        i = 0
        mx = 0
        for k in self.avg_hists.keys():
            hist = self.avg_hists[k]
            reds = hist[0:255]
            greens = hist[255:511]
            blues = hist[511:]

            reds = reds[25:-25]
            greens = greens[25:-25]
            blues = blues[25:-25]

            mx = max([mx, max(reds), max(greens), max(blues)])

            ph.addSeries(range(25, 25 + len(reds)), reds, plt=i, label="reds "+k)
            ph.addSeries(range(25, 25 + len(greens)), greens, plt=i, label="greens "+k)
            ph.addSeries(range(25, 25 + len(blues)), blues, plt=i, label="blues "+k)
            i += 1

        ph.set_xy_labels("Color Channel Values", "Avg. Color Value")
        ph.set_y_range(mx)
        ph.plot()