#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import numpy as np
import cv2
from shapely.geometry import Polygon
from PIL import Image
import random
from os import listdir
from os import path
from os import makedirs
from random import choice
import argparse

max_w, max_h = 512, 512
test_count, train_count = 10, 20


def in_rectangle(width, height, x, y):
    if 0 < x < width and height > y > 0:
        return True
    return False


def poly_in_rectangle(points, width, height):
    for p in points:
        if not in_rectangle(width, height, p[0], p[1]):
            return False
    return True


def valid_intersection_polygons(i0, i1, i2, i3, max_x, max_y):
    inner_polygon = Polygon([(i0[0], i0[1]), (i1[0], i1[1]), (i2[0], i2[1]), (i3[0], i3[1])])
    outer_polygon = max_x * max_y
    p = inner_polygon.area / outer_polygon * 100
    if p > 0.5:
        return True
    return False


def point_rotare(point, matrix):
    r = matrix.dot(np.float32([[point[0]], [point[1]], [1]]))
    return r[0][0], r[1][0]


def find_centre(points):
    left_x, top_y, right_x, bottom_y = bounding_box(points)
    return (left_x + right_x) / 2, (top_y + bottom_y) / 2


def rotate_and_scale_points(points, angle_min, angle_max,
                            scale_min, scale_max,
                            img_width, img_height):
    angle = random.uniform(angle_min, angle_max)
    while True:
        scale = random.uniform(scale_min, scale_max)
        centre = find_centre(points)
        M = cv2.getRotationMatrix2D(centre, angle, scale)
        # M = cv2.getRotationMatrix2D((img_width/2,img_height/2),90,1)
        new_points = list(map(lambda x: point_rotare(x, M), points))
        if poly_in_rectangle(new_points, img_width, img_height):
            return new_points
        scale_max -= 0.01 * abs(scale_max - scale_min)  # Уменьшаем Scale пока не впишется


def bounding_box(points):
    left_x = min(points, key=lambda p: p[0])[0]
    right_x = max(points, key=lambda p: p[0])[0]
    bottom_y = min(points, key=lambda p: p[1])[1]
    top_y = max(points, key=lambda p: p[1])[1]
    return (left_x, top_y, right_x, bottom_y)


def move_poly(points, img_width, img_height):
    left_x, top_y, right_x, bottom_y = bounding_box(points)
    dx = random.uniform(-left_x, img_width - right_x)
    dy = random.uniform(-bottom_y, img_height - top_y)
    new_points = list(map(lambda p: [p[0] + dx, p[1] + dy], points))
    return new_points


def random_shift_point(point, d, img_width, img_height):
    x = random.uniform(max(0, point[0] - d), min(img_width, point[0] + d))
    y = random.uniform(max(0, point[1] - d), min(img_height, point[1] + d))
    return [x, y]


def random_perspective_points_transform(points, value, img_width, img_height):
    left_x, top_y, right_x, bottom_y = bounding_box(points)
    d = (right_x - left_x + top_y - bottom_y) / 2 * value
    new_points = list(
        map(lambda p: random_shift_point(p, d, img_width, img_height), points))
    return new_points


def gen_points():
    x0, y0 = 0, 0
    x1, y1 = x0 + max_w, y0
    x2, y2 = x0 + max_w, y0 + max_h
    x3, y3 = x0, y0 + max_h
    return x0, y0, x1, y1, x2, y2, x3, y3


def genM(img):
    x0, y0, x1, y1, x2, y2, x3, y3 = gen_points()
    rows, cols, ch = img.shape
    orig_points = np.float32(
        [[x0, y0],
         [x1, y1],
         [x2, y2],
         [x3, y3]])
    new_points = orig_points
    new_points = rotate_and_scale_points(new_points, 0, 360, 0.5, 2, cols, rows)
    new_points = move_poly(new_points, cols, rows)
    new_points = random_perspective_points_transform(new_points, 0.1, cols, rows)
    new_points = np.float32(new_points)
    M = cv2.getPerspectiveTransform(new_points, orig_points)
    return M, new_points


def save_tile(img, M, save_img_path, color=True):
    if color:
        transformed_img = cv2.cvtColor(cv2.warpPerspective(img, M, (max_w, max_h)), cv2.COLOR_BGR2RGB)
    else:
        transformed_img = cv2.cvtColor(cv2.warpPerspective(img, M, (max_w, max_h)), cv2.COLOR_GRAY2RGB)
    im_pil = Image.fromarray(transformed_img)
    im_pil.save(save_img_path)


def count_colours(img):
    unique, counts = np.unique(img, return_counts=True)
    return dict(zip(unique, counts))


def threshold_black(img, threshold):
    result = count_colours(img)
    try:
        if result[255] < (result[0] + result[255]) * threshold:
            return True
        return False
    except:
        return False


def crop(img_path, img_bg_path, img_mask_path, save_img_path, save_bg_path, threshold=0.5):
    img = cv2.imread(img_path)
    M, new_points = genM(img)
    if path.exists(img_mask_path):
        img_mask = cv2.imread(img_mask_path, cv2.IMREAD_UNCHANGED)[:, :, 3]
        transformed_img = cv2.cvtColor(cv2.warpPerspective(img_mask, M, (max_w, max_h)), cv2.COLOR_GRAY2RGB)
        # Тут проверка на точки через пороговое значение
        if threshold_black(transformed_img, threshold):
            return False
    img_bg = cv2.imread(img_bg_path, cv2.IMREAD_UNCHANGED)[:, :, 3]

    save_tile(img, M, save_img_path, True)
    save_tile(img_bg, M, save_bg_path, False)
    return True


def get_jpg_files_whithout_background(folder_path):
    files = []
    for file in listdir(folder_path):
        if file.endswith('.png') and '_bg' not in file and '_mask' not in file:
            files.append(path.join(folder_path, file))
    return files


def cropping(root_folder, save_folder, max_files=100):
    if not path.exists(save_folder):
        makedirs(save_folder)

    without_background = get_jpg_files_whithout_background(root_folder)
    count = 0
    while count < max_files:
        image = choice(without_background)
        base_name = path.basename(image)

        mask_fat_name = base_name[:-4] + '_mask.png'
        mask_fat_name_full = path.join(image[:-len(base_name)], mask_fat_name)

        bg_name = base_name[:-4] + '_bg.png'
        bg_name_full = path.join(image[:-len(base_name)], bg_name)

        save_img_path = path.join(save_folder, str(count) + '_' + base_name)
        save_bg_path = path.join(save_folder, str(count) + '_' + bg_name)

        print("Paths:", image, bg_name_full, mask_fat_name_full, save_img_path, save_bg_path)
        if crop(image, bg_name_full, mask_fat_name_full, save_img_path, save_bg_path, threshold=0.5):
            count += 1


def smart_crop():
    parser = argparse.ArgumentParser()
    parser.add_argument('-i', '--input', help='Input folder name', required=True)
    args = parser.parse_args()

    train_folder = path.join(args.input, 'train')
    if not path.exists(train_folder):
        makedirs(train_folder)

    test_folder = path.join(args.input, 'test')
    if not path.exists(test_folder):
        makedirs(test_folder)

    cropping('test_bad', path.join(test_folder, 'bad'), test_count)
    cropping('test_good', path.join(test_folder, 'good'), test_count)

    cropping('train_bad', path.join(train_folder, 'bad'), train_count)
    cropping('train_good', path.join(train_folder, 'good'), train_count)


if __name__ == '__main__':
    smart_crop()