from core.OperationBase import OperationBase
import random
import cv2
import numpy as np


class Folding(OperationBase):

    def __init__(self, image):
        OperationBase.__init__(self, image)
        self.no_folding = 0
        self.slices = []
        self.sections = []
        self.moves = []
        self.src_points = []
        self.moved_points = []
        self.max_width = 0

    def do_operation(self, kwargs):
        if 'folding_no' in kwargs and 'min_move' in kwargs and 'max_move' in kwargs:
            self.no_folding = kwargs['folding_no']
            self.generate_random_sections()
            self.get_slices()
            self.get_moves(kwargs['min_move'], kwargs['max_move'])
            self.generate_points()
            self.move_points()
        return self.apply_folding()

    def generate_random_sections(self):
        self.sections = []
        start = self.rows / 10
        end = self.rows / self.no_folding

        for i in range(self.no_folding):
            section = random.randint(start + start / 10, end - end / 10)
            start = section
            end += self.rows / self.no_folding
            self.sections.append(section)

    def get_slices(self):
        start = 0
        for i in self.sections:
            end = i
            self.slices.append(self.image[start:end, :])
            start = end
        self.slices.append(self.image[start:self.rows, :])

    def get_moves(self, min_move, max_move):
        self.max_width = int(self.cols + self.cols * (max_move / 100.0))
        self.moves = []
        for i in range(len(self.slices)):
            self.moves.append(random.randint(min_move, max_move) / 100.0)

    def generate_points(self):
        self.src_points = []
        for i in range(len(self.slices)):
            self.src_points.append([[0, 0], [0, self.slices[i].shape[0]], [self.cols, self.slices[i].shape[0]]])

    def move_points(self):
        self.moved_points = self.src_points[:]
        if self.no_folding == 1:
            self.moved_points[0] = [[self.cols * self.moves[0], 0], [0, self.slices[0].shape[0]],
                                    [self.cols, self.slices[0].shape[0]]]
            self.moved_points[1] = [[0, 0], [self.cols * self.moves[1], self.slices[1].shape[0]],
                                    [self.cols + self.cols * self.moves[1], self.slices[1].shape[0]]]
        elif self.no_folding == 2:
            self.moved_points[0] = [[self.cols * self.moves[0], 0], [0, self.slices[0].shape[0]],
                                    [self.cols, self.slices[0].shape[0]]]
            self.moved_points[1] = [[0, 0], [self.cols * self.moves[1], self.slices[1].shape[0]],
                                    [self.cols + self.cols * self.moves[1], self.slices[1].shape[0]]]
            self.moved_points[2] = [[self.cols * self.moves[1], 0], [self.cols * self.moves[2], self.slices[2].shape[0]],
                                    [self.cols + self.cols * self.moves[2], self.slices[2].shape[0]]]
        elif self.no_folding == 3:
            self.moved_points[0] = [[self.cols * self.moves[0], 0], [0, self.slices[0].shape[0]],
                                    [self.cols, self.slices[0].shape[0]]]
            self.moved_points[1] = [[0, 0], [self.cols * self.moves[1], self.slices[1].shape[0]],
                                    [self.cols + self.cols * self.moves[1], self.slices[1].shape[0]]]
            self.moved_points[2] = [[self.cols * self.moves[1], 0], [self.cols * self.moves[2], self.slices[2].shape[0]],
                                    [self.cols + self.cols * self.moves[2], self.slices[2].shape[0]]]
            self.moved_points[3] = [[self.cols * self.moves[2], 0], [self.cols * self.moves[3], self.slices[3].shape[0]],
                                    [self.cols + self.cols * self.moves[3], self.slices[3].shape[0]]]

    def apply_folding(self):
        images_finale = []
        for i in range(len(self.slices)):
            aux = cv2.getAffineTransform(np.float32(self.src_points[i]), np.float32(self.moved_points[i]))
            #FLAGS=CV2.INTER_NEAREST FOR RECTS if needed
            images_finale.append(cv2.warpAffine(self.slices[i], aux, (self.max_width, self.slices[i].shape[0]),
                                                flags=cv2.INTER_NEAREST))

        image = images_finale[0]
        for i in range(1, len(images_finale), 1):
            image = np.vstack((image, images_finale[i]))
        return image