import os
from os.path import isfile, join

import imutils as imutils
import numpy as np
import cv2
import constant
import mimetypes

try:
    from cv2 import cv2
finally:
    pass

proto_file = "./model/colorization_deploy_v2.prototxt"
model = "./model/colorization_release_v2.caffemodel"
points_file = "./model/pts_in_hull.npy"
network = cv2.dnn.readNetFromCaffe(proto_file, model)
points = np.load(points_file)

points = points.transpose().reshape(2, 313, 1, 1)

class_layer = network.getLayerId("class8_ab")
convolution_layer = network.getLayerId("conv8_313_rh")

network.getLayer(class_layer).blobs = [points.astype(constant.FLOAT32)]
network.getLayer(convolution_layer).blobs = [np.full([1, 313], 2.606, dtype=constant.FLOAT32)]


# Convert from grayscale to gray to remove all little noises, and then again to rgb
def clear_image(image):
    cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
    cv2.cvtColor(image, cv2.COLOR_GRAY2RGB)

    return image


def convert_frames_to_video(file_name):
    path_in = 'videos/colorized_frames/'
    path_out = 'videos/output/' + file_name.replace('/', '_')
    frame_array = []
    files = [f for f in os.listdir(path_in) if isfile(join(path_in, f))]
    files.sort(key=lambda x: int(x[5:-4]))

    for i in range(len(files)):
        filename = path_in + files[i]
        img = cv2.imread(filename)
        height, width, layers = img.shape
        size = (width, height)
        frame_array.append(img)

    out = cv2.VideoWriter(path_out, cv2.VideoWriter_fourcc(*'MJPG'), 30.0, size)

    for i in range(len(frame_array)):
        out.write(frame_array[i])

    out.release()


def colorize_image(image):
    image_scaled = image.astype(constant.FLOAT32) / 255.0
    lab = cv2.cvtColor(image_scaled, cv2.COLOR_RGB2LAB)
    resized = cv2.resize(lab, (224, 224))
    image_l_channel = cv2.split(resized)[0] - 50

    network.setInput(cv2.dnn.blobFromImage(image_l_channel))
    ab = network.forward()[0, :, :, :].transpose((1, 2, 0))
    ab = cv2.resize(ab, (image.shape[1], image.shape[0]))

    image_l_channel = cv2.split(lab)[0]
    colorized = np.concatenate((image_l_channel[:, :, np.newaxis], ab), axis=2)

    colorized = cv2.cvtColor(colorized, cv2.COLOR_LAB2RGB)
    colorized = np.clip(colorized, 0, 1)

    return (255 * colorized).astype("uint8")


def colorize(file_path: str) -> None:
    file_type = mimetypes.guess_type(file_path)[0].split('/')

    if file_type[0] == 'image':
        image = cv2.imread(file_path)
        colorized = colorize_image(image)

        cv2.imwrite("images/output/" + file_path.replace('/', '_'), cv2.cvtColor(colorized, cv2.COLOR_RGB2BGR))

        print('Success')
    elif file_type[0] == 'video':
        vs = cv2.VideoCapture(file_path)
        count = 0
        success = True

        while success:
            success, frame = vs.read()

            if frame is None:
                break

            frame = imutils.resize(frame, 500)

            colorized = colorize_image(frame)

            cv2.imwrite("./videos/colorized_frames/frame%d.jpg" % count, colorized)
            count += 1
            key = cv2.waitKey(1) & 0xFF
            if key == ord("q"):
                break

        vs.release()
        convert_frames_to_video(file_path)

        print('Success')
    else:
        print('Wrong media format')