import random
import time
from datetime import datetime

import cv2
import numpy as np
import os


class GenerateVideoWithImage():
    def __myResizeImage(self, image, resizeHeight, resizeWidth, BaseHeight=True):

        h, w, c = image.shape
        if BaseHeight:
            my_height = int(resizeHeight)
            my_width = int(my_height * (w / h))
            my_image = cv2.resize(image, (my_width, my_height))
            if my_width > resizeWidth:
                xo = int((my_width - resizeWidth) / 2)
                x1 = xo + resizeWidth
                result_image = my_image[0:resizeHeight, xo:x1]

            else:
                result_image = np.zeros((resizeHeight, resizeWidth, 3), dtype=np.uint8)
                result_image[:, :, 0] = 255
                result_image[:, :, 1] = 255
                result_image[:, :, 2] = 255

                xo = int((resizeWidth - my_width) / 2)
                x1 = xo + my_width
                result_image[0:resizeHeight, xo:x1] = my_image

        else:
            my_width = int(resizeWidth)
            my_height = int(my_width * (h / w))
            my_image = cv2.resize(image, (my_width, my_height))
            if my_height > resizeHeight:
                result_image = my_image[0:resizeHeight, 0:resizeWidth]
            else:
                result_image = np.zeros((resizeHeight, resizeWidth, 3), dtype=np.uint8)
                result_image[:, :, 0] = 255
                result_image[:, :, 1] = 255
                result_image[:, :, 2] = 255
                result_image[0:my_height, 0:my_width] = my_image

        return result_image

    def __myMergeImage(self, image_path_array, mergeImageHeight, mergeImageWidth):
        mergeImage = np.zeros((mergeImageHeight, mergeImageWidth, 3), dtype=np.uint8)
        # mergeImage[:, :, 0] = 255
        # mergeImage[:, :, 1] = 255
        # mergeImage[:, :, 2] = 255
        image_count = len(image_path_array)

        if image_count == 4:
            for i, image_path in enumerate(image_path_array):

                image = self.__myResizeImage(cv2.imread(image_path_array[i]), int(mergeImageHeight / 2),
                                             int(mergeImageWidth / 2))
                yo = 0
                y1 = int(mergeImageHeight / 2)
                x0 = int((mergeImageWidth / 2) * i)
                if i > 1:
                    yo = int(mergeImageHeight / 2)
                    y1 = mergeImageHeight
                    x0 = int((mergeImageWidth / 2) * (i - 2))
                x1 = x0 + int(mergeImageWidth / 2)

                mergeImage[yo:y1, x0:x1] = image

            # image1 = self.__myResizeImage(cv2.imread(image_path_array[0]), int(mergeImageHeight / 2),
            #                               int(mergeImageWidth / 2))
            # image2 = self.__myResizeImage(cv2.imread(image_path_array[1]), int(mergeImageHeight / 2),
            #                               int(mergeImageWidth / 2))
            # image3 = self.__myResizeImage(cv2.imread(image_path_array[2]), int(mergeImageHeight / 2),
            #                               int(mergeImageWidth / 2))
            # image4 = self.__myResizeImage(cv2.imread(image_path_array[3]), int(mergeImageHeight / 2),
            #                               int(mergeImageWidth / 2))
            # mergeImage[0:int(mergeImageHeight / 2), 0:int(mergeImageWidth / 2)] = image1
            # mergeImage[0:int(mergeImageHeight / 2), int(mergeImageWidth / 2):mergeImageWidth] = image2
            # mergeImage[int(mergeImageHeight / 2):mergeImageHeight, 0:int(mergeImageWidth / 2)] = image3
            # mergeImage[int(mergeImageHeight / 2):mergeImageHeight, int(mergeImageWidth / 2):mergeImageWidth] = image4
        elif image_count == 5:
            image1 = None
            for i, image_path in enumerate(image_path_array):
                if i < 4:
                    image = self.__myResizeImage(cv2.imread(image_path_array[i]), int(mergeImageHeight / 2),
                                                 int(mergeImageWidth / 2))
                    if i == 0:
                        image1 = image
                    yo = 0
                    y1 = int(mergeImageHeight / 2)
                    x0 = int((mergeImageWidth / 2) * i)
                    if i > 1:
                        yo = int(mergeImageHeight / 2)
                        y1 = mergeImageHeight
                        x0 = int((mergeImageWidth / 2) * (i - 2))
                    x1 = x0 + int(mergeImageWidth / 2)

                    mergeImage[yo:y1, x0:x1] = image
                else:
                    mergeImage[int(mergeImageHeight / 4):int(mergeImageHeight / 4) * 3,
                    int(mergeImageWidth / 4):int(mergeImageWidth / 4) * 3] = image1

        elif image_count == 8:
            for i, image_path in enumerate(image_path_array):

                image = self.__myResizeImage(cv2.imread(image_path_array[i]), int(mergeImageHeight / 2),
                                             int(mergeImageWidth / 4))
                yo = 0
                y1 = int(mergeImageHeight / 2)
                x0 = int((mergeImageWidth / 4) * i)
                if i > 3:
                    yo = int(mergeImageHeight / 2)
                    y1 = mergeImageHeight
                    x0 = int((mergeImageWidth / 4) * (i - 4))
                x1 = x0 + int(mergeImageWidth / 4)

                mergeImage[yo:y1, x0:x1] = image

        return mergeImage

    def __putImageInBg(self, image, bg_h, bg_w):

        bg = np.zeros((bg_h, bg_w, 3), dtype=np.uint8)

        bg[:, :, 0] = 255
        bg[:, :, 1] = 255
        bg[:, :, 2] = 255

        h, w, _ = image.shape

        change_h = bg_h - h
        change_w = bg_w - w

        if change_h >= 0 and change_w >= 0:
            y0 = random.randint(0, change_h)
            x0 = random.randint(0, change_w)

            bg[y0:y0 + h, x0:x0 + w] = image

        return bg

    def generateVideo_LongWidth(self, image_dir, save_video_dir, save_video_name):
        image_names = os.listdir(image_dir)
        print("image_dir=%s,len(image_names)=%d" % (image_dir, len(image_names)))

        videoWidth = 1920
        videoHeight = 1080
        videoFps = 25
        frameCycleNum = 25  # 一帧合成后的图片循环展示的次数
        frameFragmentCount = 4  # 一帧图片分成的片段数量
        frameFragmentWidth = int(videoWidth / frameFragmentCount)
        frameFragmentHeight = int(videoHeight)

        save_video_path = "%s/%d_%d_%s" % (save_video_dir, frameFragmentCount, frameCycleNum, save_video_name)

        videoSize = (videoWidth, videoHeight)
        fourcc = cv2.VideoWriter_fourcc(*'mp4v')
        video = cv2.VideoWriter(save_video_path, fourcc, videoFps, videoSize, True)

        fragments = []

        for image_name in image_names:

            if len(fragments) > frameFragmentCount:
                for _ in range(frameFragmentCount):
                    del fragments[0]

            image_path = os.path.join(image_dir, image_name)
            image = cv2.imread(image_path)
            image = cv2.resize(image, (videoWidth, videoHeight))

            for i in range(frameFragmentCount):
                top = 0
                bottom = frameFragmentHeight
                left = int(i * frameFragmentWidth)
                right = int((i + 1) * frameFragmentWidth)

                fragment = image[top:bottom, left:right]
                fragments.append(fragment)

            if len(fragments) > frameFragmentCount:
                for index in range(len(fragments) - frameFragmentCount):
                    new_image_fragments = fragments[index:index + frameFragmentCount]
                    new_image = np.hstack(new_image_fragments)

                    # cv2.imshow('new_image', new_image)
                    # cv2.waitKey(0)

                    t1 = time.time()
                    for _ in range(frameCycleNum):
                        video.write(new_image)
                    t2 = time.time()

                    print("video.write 1 frame spend %.3f" % ((t2 - t1)))

        video.release()
        cv2.destroyAllWindows()

    def generateVideo_OneByOne(self, image_dir, save_video_dir, save_video_name):
        image_names = os.listdir(image_dir)
        print("image_dir=%s,len(image_names)=%d" % (image_dir, len(image_names)))

        if not os.path.exists(save_video_dir):
            os.makedirs(save_video_dir)

        save_video_path = os.path.join(save_video_dir, save_video_name)

        videoWidth = 1920
        videoHeight = 1080
        videoFps = 12
        frameCycleNum = 25  # 一帧合成后的图片循环展示的次数

        videoSize = (videoWidth, videoHeight)
        fourcc = cv2.VideoWriter_fourcc(*'mp4v')
        video = cv2.VideoWriter(save_video_path, fourcc, videoFps, videoSize, True)

        print(image_names)
        i = 0
        for image_name in image_names:
            if image_name.endswith(".jpg"):
                i += 1
                if i > 200:
                    break
                image_path = os.path.join(image_dir, image_name)
                image = cv2.imread(image_path)

                image = self.__myResizeImage(image, videoHeight, videoWidth)
                # 默认使用双线性插值法

                t1 = time.time()
                for _ in range(frameCycleNum):
                    video.write(image)
                t2 = time.time()
                print("video.write 1 frame spend %.3f" % ((t2 - t1)))

                # cv2.imshow('image', image)
                # cv2.waitKey(0)

        video.release()

        # cv2.destroyAllWindows()

    def generateVideo_RandomMergePicture(self, image_dir, save_video_dir, save_video_name):

        image_names = os.listdir(image_dir)
        print("image_dir=%s,len(image_names)=%d" % (image_dir, len(image_names)))

        save_video_path = os.path.join(save_video_dir, save_video_name)

        videoWidth = 1920
        videoHeight = 1080
        videoFps = 30
        frameCycleNum = 20  # 一帧合成后的图片循环展示的次数

        videoSize = (videoWidth, videoHeight)
        fourcc = cv2.VideoWriter_fourcc(*'mp4v')
        video = cv2.VideoWriter(save_video_path, fourcc, videoFps, videoSize, True)

        for i, image_name in enumerate(image_names):
            image_path1 = os.path.join(image_dir, image_names[i])

            for j in range(4):
                image_path_array = []
                index = i + j
                if index < len(image_names):
                    pass
                else:
                    index = i - j
                image_path2 = os.path.join(image_dir, image_names[index])

                image_path_array.append(image_path1)
                image_path_array.append(image_path2)
                image_path_array.append(os.path.join(image_dir, random.choice(image_names)))
                image_path_array.append(os.path.join(image_dir, random.choice(image_names)))
                # image_path_array.append(None)
                # image_path_array.append(os.path.join(image_dir,random.choice(image_names)))
                # image_path_array.append(os.path.join(image_dir,random.choice(image_names)))
                # image_path_array.append(os.path.join(image_dir,random.choice(image_names)))
                # image_path_array.append(os.path.join(image_dir,random.choice(image_names)))

                mergeImage = self.__myMergeImage(image_path_array, mergeImageHeight=videoHeight,
                                                 mergeImageWidth=videoWidth)
                try:
                    t1 = time.time()
                    for _ in range(random.randint(15, 60)):
                        video.write(mergeImage)
                    t2 = time.time()
                    print("video.write 1 frame spend %.3f" % ((t2 - t1)))

                    # cv2.imshow('image', image)
                    # cv2.waitKey(0)
                except Exception as e:
                    print("mergeImage error:%s" % str(e))

        video.release()


if __name__ == '__main__':
    # 指定合成视频的图片文件夹绝对路径
    image_dir = "E:\\project\\ai\\data\\body\\FallDataset\\img"
    save_video_dir = "E:\\project\\ai\\data\\body\\FallDataset\\img_video"
    save_video_name = "img_video.mp4"

    # 实例化一个合成视频类的对象
    gen = GenerateVideoWithImage()
    # 图片按序一张一张的合并成视频
    gen.generateVideo_OneByOne(image_dir=image_dir, save_video_dir=save_video_dir, save_video_name=save_video_name)

    # 图片按照从左往右滑动的的方式合成视频
    # gen.generateVideo_LongWidth(image_dir=image_dir, save_video_dir=save_video_dir, save_video_name=save_video_name)

    # 按照内部结构随机合成视频，具体需要参考一下内部的代码实现。当下支持4,5,8张图片合成视频的方式
    # gen.generateVideo_RandomMergePicture(image_dir=image_dir, save_video_dir=save_video_dir,save_video_name=save_video_name)