"""
VCD (Video Content Description) library v5.0.1.

Project website: http://vcd.vicomtech.org

Copyright (C) 2023, Vicomtech (http://www.vicomtech.es/),
(Spain) all rights reserved.

VCD is a Python library to create and manage OpenLABEL content.
VCD is distributed under MIT License. See LICENSE.
"""


import json
import os
import time

import converter
import cv2 as cv
import numpy as np

from vcd import core, types, utils

use_profiler = False
show_images = True
quick_test = True
root_name = "openlabel"

with open("config.json", encoding="utf-8") as f:
    mapillary_config = json.load(f)

labels = mapillary_config["labels"]
path = "labels"


def convert_labels_to_poly(img_name, vcd_poly_mode):
    print("VCD Poly2d mode " + vcd_poly_mode.name)
    start = time.time()
    vcd_class = converter.mapillary_classes_to_vcd(
        img_name=img_name,
        mapillary_config=mapillary_config,
        vcd_poly_mode=vcd_poly_mode,
        vcd_root_name=root_name,
    )
    end = time.time()
    print(f"\tConverted class labels to VCD in {(end - start) * 1000:.2f} ms")

    start = time.time()
    vcd_class.save(
        "vcd_files/"
        + root_name
        + img_name[0 : len(img_name) - 4]
        + "_class_"
        + vcd_poly_mode.name
        + ".json"
    )
    end = time.time()

    print(f"\tSaved to VCD JSON file in {(end - start) * 1000:.2f} ms")

    start = time.time()
    img_class_out = converter.vcd_mapillary_classes_to_png(
        vcd=vcd_class, mapillary_config=mapillary_config
    )
    end = time.time()
    print(f"\tClass reconstructed from VCD in {(end - start) * 1000:.2f} ms")

    # Prepare mosaic & compute differences
    img_label = cv.imread("labels/" + img_name)
    difference_class = cv.subtract(img_label, img_class_out)
    # Check equals
    difference_class_sum = np.sum(difference_class)
    print("\tClass differences: " + str(difference_class_sum))

    if show_images:
        stack = np.hstack((img_label, img_class_out, difference_class))
        cv.namedWindow(img_name + vcd_poly_mode.name, cv.WINDOW_NORMAL)
        cv.resizeWindow(
            img_name + vcd_poly_mode.name, img_label.shape[0], img_label.shape[1]
        )
        cv.imshow(img_name + vcd_poly_mode.name, stack)
        cv.waitKey(0)


def convert_instances_to_poly(img_name, vcd_poly_mode):
    print("VCD Poly2d mode " + vcd_poly_mode.name)
    start = time.time()
    vcd_ins = converter.mapillary_instances_to_vcd(
        img_name=img_name,
        mapillary_config=mapillary_config,
        vcd_poly_mode=vcd_poly_mode,
        vcd_root_name=root_name,
    )
    end = time.time()
    print(f"\tConverted instances to VCD in {(end - start) * 1000:.2f} ms")

    start = time.time()
    vcd_ins.save(
        "vcd_files/"
        + root_name
        + img_name[0 : len(img_name) - 4]
        + "_instances_"
        + vcd_poly_mode.name
        + ".json"
    )
    end = time.time()

    print(f"\tSaved to VCD JSON file in {(end - start) * 1000:.2f} ms")

    start = time.time()
    img_instances_out = converter.vcd_mapillary_instances_to_png(
        vcd=vcd_ins, mapillary_config=mapillary_config
    )
    end = time.time()
    print(f"\tInstances reconstructed from VCD in {(end - start) * 1000:.2f} ms")

    # Prepare mosaic & compute differences
    img_instances = cv.imread("instances/" + img_name, cv.IMREAD_ANYDEPTH)

    # Check equals
    difference_instances = cv.subtract(img_instances, img_instances_out)
    difference_instances_sum = np.sum(difference_instances)
    print("\tInstances differences: " + str(difference_instances_sum))

    height = img_instances.shape[0]
    width = img_instances.shape[1]
    _temp = np.zeros((height, width), np.uint8)

    if show_images:
        stack = np.hstack((img_instances, img_instances_out, difference_instances))
        cv.namedWindow(img_name + "_instances", cv.WINDOW_NORMAL)
        cv.resizeWindow(
            img_name + "_instances", img_instances.shape[0], img_instances.shape[1]
        )
        cv.imshow(img_name + "_instances", stack)
        cv.waitKey(0)


def convert_labels_to_base64(img_name):
    img_label = cv.imread("labels/" + img_name)
    payload_b64_str = utils.image_to_base64(img_label)

    if root_name == "vcd":
        vcd = core.VCD()
    elif root_name == "openlabel":
        vcd = core.OpenLABEL()

    # Create VCD
    height = img_label.shape[0]
    width = img_label.shape[1]
    vcd.add_metadata_properties(properties={"labels": "class"})
    vcd.add_stream(
        stream_name="Camera",
        uri=img_name,
        description="Mapillary image",
        stream_type=core.StreamType.camera,
    )
    vcd.add_stream_properties(
        stream_name="Camera", properties={"height": height, "width": width}
    )
    vcd_image = types.image("labels", payload_b64_str, "image/png", "base64")
    uid = vcd.add_object("", "")
    vcd.add_object_data(uid, vcd_image)

    vcd.save(
        "vcd_files/" + root_name + img_name[0 : len(img_name) - 4] + "_class_b64.json"
    )

    # Decode
    od = vcd.get_object_data(uid, "labels")
    payload_b64_read = od["val"]
    img_class_out = utils.base64_to_image(payload_b64_read)

    # Check equals
    difference_class = cv.subtract(img_label, img_class_out)
    difference_class_sum = np.sum(difference_class)
    print("\tClass differences: " + str(difference_class_sum))

    if show_images:
        stack = np.hstack((img_label, img_class_out, difference_class))
        cv.namedWindow(img_name + "_b64", cv.WINDOW_NORMAL)
        cv.resizeWindow(img_name + "_b64", img_label.shape[0], img_label.shape[1])
        cv.imshow(img_name + "_b64", stack)
        cv.waitKey(0)


def convert_instances_to_base64(img_name):
    img_instances = cv.imread("instances/" + img_name, cv.IMREAD_ANYDEPTH)
    payload_b64_str = utils.image_to_base64(img_instances)

    if root_name == "vcd":
        vcd = core.VCD()
    elif root_name == "openlabel":
        vcd = core.OpenLABEL()

    # Create VCD
    height = img_instances.shape[0]
    width = img_instances.shape[1]
    vcd.add_metadata_properties(properties={"labels": "class"})
    vcd.add_stream(
        stream_name="Camera",
        uri=img_name,
        description="Mapillary image",
        stream_type=core.StreamType.camera,
    )
    vcd.add_stream_properties(
        stream_name="Camera", properties={"height": height, "width": width}
    )
    vcd_image = types.image("instances", payload_b64_str, "image/png", "base64")
    uid = vcd.add_object("", "")
    vcd.add_object_data(uid, vcd_image)

    vcd.save(
        "vcd_files/" + root_name + img_name[0 : len(img_name) - 4] + "_instances_b64.json"
    )

    # Decode
    od = vcd.get_object_data(uid, "instances")
    payload_b64_read = od["val"]
    img_instances_out = utils.base64_to_image(payload_b64_read, cv.IMREAD_ANYDEPTH)

    # Save image
    compr_params = [int(cv.IMWRITE_PNG_COMPRESSION), 9]
    cv.imwrite(
        filename="vcd_files/"
        + root_name
        + img_name[0 : len(img_name) - 4]
        + "_instances_reconstructed_from_b64.png",
        img=img_instances_out,
        params=compr_params,
    )

    # Check equals
    difference_instances = cv.subtract(img_instances, img_instances_out)
    difference_instances_sum = np.sum(difference_instances)
    print("\tInstances differences: " + str(difference_instances_sum))

    if show_images:
        stack = np.hstack((img_instances, img_instances_out, difference_instances))
        cv.namedWindow(img_name + "_b64", cv.WINDOW_NORMAL)
        cv.resizeWindow(img_name + "_b64", img_instances.shape[0], img_instances.shape[1])
        cv.imshow(img_name + "_b64", stack)
        cv.waitKey(0)


for img_name in os.listdir(path):
    # Class LABELS
    print(img_name)

    # Convert class images
    convert_labels_to_base64(img_name=img_name)
    convert_labels_to_poly(
        img_name=img_name, vcd_poly_mode=types.Poly2DType.MODE_POLY2D_SRF6DCC
    )
    convert_labels_to_poly(
        img_name=img_name, vcd_poly_mode=types.Poly2DType.MODE_POLY2D_ABSOLUTE
    )

    # Convert instance images
    convert_instances_to_base64(img_name=img_name)
    convert_instances_to_poly(
        img_name=img_name, vcd_poly_mode=types.Poly2DType.MODE_POLY2D_SRF6DCC
    )
    convert_instances_to_poly(
        img_name=img_name, vcd_poly_mode=types.Poly2DType.MODE_POLY2D_ABSOLUTE
    )

    if quick_test:
        break
