# -*- coding:utf-8 -*-
# Author: RubanSeven

# import cv2
import numpy as np
# from transform import get_perspective_transform, warp_perspective
from warp_mls import WarpMLS


def distort(src, segment):
    img_h, img_w = src.shape[:2]

    cut = img_w // segment
    thresh = cut // 3
    # thresh = img_h // segment // 3
    # thresh = img_h // 5

    src_pts = list()
    dst_pts = list()

    src_pts.append([0, 0])
    src_pts.append([img_w, 0])
    src_pts.append([img_w, img_h])
    src_pts.append([0, img_h])

    dst_pts.append([np.random.randint(thresh), np.random.randint(thresh)])
    dst_pts.append([img_w - np.random.randint(thresh), np.random.randint(thresh)])
    dst_pts.append([img_w - np.random.randint(thresh), img_h - np.random.randint(thresh)])
    dst_pts.append([np.random.randint(thresh), img_h - np.random.randint(thresh)])

    half_thresh = thresh * 0.5

    for cut_idx in np.arange(1, segment, 1):
        src_pts.append([cut * cut_idx, 0])
        src_pts.append([cut * cut_idx, img_h])
        dst_pts.append([cut * cut_idx + np.random.randint(thresh) - half_thresh,
                        np.random.randint(thresh) - half_thresh])
        dst_pts.append([cut * cut_idx + np.random.randint(thresh) - half_thresh,
                        img_h + np.random.randint(thresh) - half_thresh])

    trans = WarpMLS(src, src_pts, dst_pts, img_w, img_h)
    dst = trans.generate()

    return dst


def stretch(src, segment):
    img_h, img_w = src.shape[:2]

    cut = img_w // segment
    thresh = cut * 4 // 5
    # thresh = img_h // segment // 3
    # thresh = img_h // 5

    src_pts = list()
    dst_pts = list()

    src_pts.append([0, 0])
    src_pts.append([img_w, 0])
    src_pts.append([img_w, img_h])
    src_pts.append([0, img_h])

    dst_pts.append([0, 0])
    dst_pts.append([img_w, 0])
    dst_pts.append([img_w, img_h])
    dst_pts.append([0, img_h])

    half_thresh = thresh * 0.5

    for cut_idx in np.arange(1, segment, 1):
        move = np.random.randint(thresh) - half_thresh
        src_pts.append([cut * cut_idx, 0])
        src_pts.append([cut * cut_idx, img_h])
        dst_pts.append([cut * cut_idx + move, 0])
        dst_pts.append([cut * cut_idx + move, img_h])

    trans = WarpMLS(src, src_pts, dst_pts, img_w, img_h)
    dst = trans.generate()

    return dst


def perspective(src):
    img_h, img_w = src.shape[:2]

    thresh = img_h // 2

    src_pts = list()
    dst_pts = list()

    src_pts.append([0, 0])
    src_pts.append([img_w, 0])
    src_pts.append([img_w, img_h])
    src_pts.append([0, img_h])

    dst_pts.append([0, np.random.randint(thresh)])
    dst_pts.append([img_w, np.random.randint(thresh)])
    dst_pts.append([img_w, img_h - np.random.randint(thresh)])
    dst_pts.append([0, img_h - np.random.randint(thresh)])

    trans = WarpMLS(src, src_pts, dst_pts, img_w, img_h)
    dst = trans.generate()

    return dst

# def distort(src, segment):
#     img_h, img_w = src.shape[:2]
#     dst = np.zeros_like(src, dtype=np.uint8)
#
#     cut = img_w // segment
#     thresh = img_h // 8
#
#     src_pts = list()
#     # dst_pts = list()
#
#     src_pts.append([-np.random.randint(thresh), -np.random.randint(thresh)])
#     src_pts.append([-np.random.randint(thresh), img_h + np.random.randint(thresh)])
#
#     # dst_pts.append([0, 0])
#     # dst_pts.append([0, img_h])
#     dst_box = np.array([[0, 0], [0, img_h], [cut, 0], [cut, img_h]], dtype=np.float32)
#
#     half_thresh = thresh * 0.5
#
#     for cut_idx in np.arange(1, segment, 1):
#         src_pts.append([cut * cut_idx + np.random.randint(thresh) - half_thresh,
#                         np.random.randint(thresh) - half_thresh])
#         src_pts.append([cut * cut_idx + np.random.randint(thresh) - half_thresh,
#                         img_h + np.random.randint(thresh) - half_thresh])
#
#         # dst_pts.append([cut * i, 0])
#         # dst_pts.append([cut * i, img_h])
#
#         src_box = np.array(src_pts[-4:-2] + src_pts[-2:-1] + src_pts[-1:], dtype=np.float32)
#
#         # mat = cv2.getPerspectiveTransform(src_box, dst_box)
#         # print(mat)
#         # dst[:, cut * (cut_idx - 1):cut * cut_idx] = cv2.warpPerspective(src, mat, (cut, img_h))
#
#         mat = get_perspective_transform(dst_box, src_box)
#         dst[:, cut * (cut_idx - 1):cut * cut_idx] = warp_perspective(src, mat, (cut, img_h))
#         # print(mat)
#
#     src_pts.append([img_w + np.random.randint(thresh) - half_thresh,
#                     np.random.randint(thresh) - half_thresh])
#     src_pts.append([img_w + np.random.randint(thresh) - half_thresh,
#                     img_h + np.random.randint(thresh) - half_thresh])
#     src_box = np.array(src_pts[-4:-2] + src_pts[-2:-1] + src_pts[-1:], dtype=np.float32)
#
#     # mat = cv2.getPerspectiveTransform(src_box, dst_box)
#     # dst[:, cut * (segment - 1):] = cv2.warpPerspective(src, mat, (img_w - cut * (segment - 1), img_h))
#     mat = get_perspective_transform(dst_box, src_box)
#     dst[:, cut * (segment - 1):] = warp_perspective(src, mat, (img_w - cut * (segment - 1), img_h))
#
#     return dst

[((33, 19, 33, 43), 'numpy.arange', 'np.arange', ({(33, 29, 33, 30): '(1)', (33, 32, 33, 39): 'segment', (33, 41, 33, 42): '(1)'}, {}), '(1, segment, 1)', True, 'import numpy as np\n'), ((41, 12, 41, 56), 'warp_mls.WarpMLS', 'WarpMLS', ({(41, 20, 41, 23): 'src', (41, 25, 41, 32): 'src_pts', (41, 34, 41, 41): 'dst_pts', (41, 43, 41, 48): 'img_w', (41, 50, 41, 55): 'img_h'}, {}), '(src, src_pts, dst_pts, img_w, img_h)', False, 'from warp_mls import WarpMLS\n'), ((70, 19, 70, 43), 'numpy.arange', 'np.arange', ({(70, 29, 70, 30): '(1)', (70, 32, 70, 39): 'segment', (70, 41, 70, 42): '(1)'}, {}), '(1, segment, 1)', True, 'import numpy as np\n'), ((77, 12, 77, 56), 'warp_mls.WarpMLS', 'WarpMLS', ({(77, 20, 77, 23): 'src', (77, 25, 77, 32): 'src_pts', (77, 34, 77, 41): 'dst_pts', (77, 43, 77, 48): 'img_w', (77, 50, 77, 55): 'img_h'}, {}), '(src, src_pts, dst_pts, img_w, img_h)', False, 'from warp_mls import WarpMLS\n'), ((101, 12, 101, 56), 'warp_mls.WarpMLS', 'WarpMLS', ({(101, 20, 101, 23): 'src', (101, 25, 101, 32): 'src_pts', (101, 34, 101, 41): 'dst_pts', (101, 43, 101, 48): 'img_w', (101, 50, 101, 55): 'img_h'}, {}), '(src, src_pts, dst_pts, img_w, img_h)', False, 'from warp_mls import WarpMLS\n'), ((26, 20, 26, 45), 'numpy.random.randint', 'np.random.randint', ({(26, 38, 26, 44): 'thresh'}, {}), '(thresh)', True, 'import numpy as np\n'), ((26, 47, 26, 72), 'numpy.random.randint', 'np.random.randint', ({(26, 65, 26, 71): 'thresh'}, {}), '(thresh)', True, 'import numpy as np\n'), ((27, 55, 27, 80), 'numpy.random.randint', 'np.random.randint', ({(27, 73, 27, 79): 'thresh'}, {}), '(thresh)', True, 'import numpy as np\n'), ((29, 20, 29, 45), 'numpy.random.randint', 'np.random.randint', ({(29, 38, 29, 44): 'thresh'}, {}), '(thresh)', True, 'import numpy as np\n'), ((71, 15, 71, 40), 'numpy.random.randint', 'np.random.randint', ({(71, 33, 71, 39): 'thresh'}, {}), '(thresh)', True, 'import numpy as np\n'), ((96, 23, 96, 48), 'numpy.random.randint', 'np.random.randint', ({(96, 41, 96, 47): 'thresh'}, {}), '(thresh)', True, 'import numpy as np\n'), ((97, 27, 97, 52), 'numpy.random.randint', 'np.random.randint', ({(97, 45, 97, 51): 'thresh'}, {}), '(thresh)', True, 'import numpy as np\n'), ((27, 28, 27, 53), 'numpy.random.randint', 'np.random.randint', ({(27, 46, 27, 52): 'thresh'}, {}), '(thresh)', True, 'import numpy as np\n'), ((28, 28, 28, 53), 'numpy.random.randint', 'np.random.randint', ({(28, 46, 28, 52): 'thresh'}, {}), '(thresh)', True, 'import numpy as np\n'), ((28, 63, 28, 88), 'numpy.random.randint', 'np.random.randint', ({(28, 81, 28, 87): 'thresh'}, {}), '(thresh)', True, 'import numpy as np\n'), ((29, 55, 29, 80), 'numpy.random.randint', 'np.random.randint', ({(29, 73, 29, 79): 'thresh'}, {}), '(thresh)', True, 'import numpy as np\n'), ((98, 35, 98, 60), 'numpy.random.randint', 'np.random.randint', ({(98, 53, 98, 59): 'thresh'}, {}), '(thresh)', True, 'import numpy as np\n'), ((99, 31, 99, 56), 'numpy.random.randint', 'np.random.randint', ({(99, 49, 99, 55): 'thresh'}, {}), '(thresh)', True, 'import numpy as np\n'), ((37, 24, 37, 49), 'numpy.random.randint', 'np.random.randint', ({(37, 42, 37, 48): 'thresh'}, {}), '(thresh)', True, 'import numpy as np\n'), ((36, 40, 36, 65), 'numpy.random.randint', 'np.random.randint', ({(36, 58, 36, 64): 'thresh'}, {}), '(thresh)', True, 'import numpy as np\n'), ((38, 40, 38, 65), 'numpy.random.randint', 'np.random.randint', ({(38, 58, 38, 64): 'thresh'}, {}), '(thresh)', True, 'import numpy as np\n'), ((39, 32, 39, 57), 'numpy.random.randint', 'np.random.randint', ({(39, 50, 39, 56): 'thresh'}, {}), '(thresh)', True, 'import numpy as np\n')]