class FaceFinder:
    def __init__(self):
        self.img = None
        self.img_orig = None
        self.patches = []
        self.preds = None

    def process_photo(self, img_name, patch_size=64):
        self.img = plt.imread(img_name)
        self.img = scipy.misc.imresize(self.img, (256, 256, 3))
        self.img_orig = self.img.copy()

        self._split_pic_into_frames(size=patch_size)
        self._run_predictions()
        # self._rectangify_patches()
        return self._gray_out()

    def _split_pic_into_frames(self, size, stride=1):
        sh_x = self.img.shape[0]
        sh_y = self.img.shape[1]
        len_x = size
        len_y = size
        stride_x = np.ceil(len_x*stride).astype(int)
        stride_y = np.ceil(len_y*stride).astype(int)

        for y in range(0, sh_y, stride_y):
            for x in range(0, sh_x, stride_x):
                patch = self.img[y:y+len_y, x:x+len_x]
                if patch.shape[0]*patch.shape[1]!=0:
                    self.patches.append([patch, (x, y), (x+len_x, y+len_y)])

    def _rectangify_patches(self):
        columns = int(np.sqrt(len(self.preds)))

        cont = True
        while cont:
            cont = False
            for r in range(columns-1):
                for c in range(columns-1):
                    if self.preds[r*columns+c] == 0:
                        LD = self.preds[(r+1)*columns+c-1]
                        SD = self.preds[(r+1)*columns+c]
                        RD = self.preds[(r+1)*columns+c+1]
                        L = self.preds[r*columns+c-1]
                        R = self.preds[r*columns+c+1]
                        LG = self.preds[(r-1)*columns+c-1]
                        SG = self.preds[(r-1)*columns+c]
                        RG = self.preds[(r-1)*columns+c+1]

                        answ = L*LD*SD + SD*RD*R + R*RG*SG + SG*LG*L

                        if answ >= 1:
                            cont = True
                            self.preds[r*columns+c] = 1

    def _gray_out(self):
        coord_start = np.array(self.patches)[:, 1]
        coord_end = np.array(self.patches)[:, 2]

        self.img.setflags(write=1)

        for start, end, pred in zip(coord_start, coord_end, self.preds):
            if pred < 0.5:
                self.img[max(0, start[0]-3):end[0]+3, max(0, start[1]-3):end[1]+3] = [255, 0, 0]

        for start, end, pred in zip(coord_start, coord_end, self.preds):
            R = self.img_orig[start[0]:end[0], start[1]:end[1]][:, :, 0]
            G = self.img_orig[start[0]:end[0], start[1]:end[1]][:, :, 1]
            B = self.img_orig[start[0]:end[0], start[1]:end[1]][:, :, 2]
            if pred < 0.5:
                self.img[start[0]:end[0], start[1]:end[1]][:, :, 0] = R
                self.img[start[0]:end[0], start[1]:end[1]][:, :, 1] = G
                self.img[start[0]:end[0], start[1]:end[1]][:, :, 2] = B

        return self.img, self.img_orig

    def _run_predictions(self):
        a = scipy.misc.imresize(self.patches[0][0], (64, 64, 3)).reshape((1, 64, 64, 3))
        for i in self.patches[1:]:
            a = np.vstack((a, scipy.misc.imresize(i[0], (64, 64, 3)).reshape((1, 64, 64, 3))))
        self.preds = predict(a, sess)