#
# Import UniTAP module.
#

import UniTAP
from UniTAP.utils import encode_video_frame, calculate_dsc_slice_size
from UniTAP.common import get_vf_from_image, CompressionInfo

#
# To initialize UniTAP library wrapper user should create UniTAP.TsiLib() object.
#
lUniTAP = UniTAP.TsiLib()

#
# For opening device, please, put serial number of the device as 8 symbol str or put index of device.
#
# dev = lUniTAP.open("NNNNNNNN")
dev = lUniTAP.open(0)

# After opening device as in UCD Console device role should be selected.
role = dev.select_role(UniTAP.dev.UCD400.DPSourceDPSink)
# role = dev.select_role(UniTAP.dev.UCD500.DPSourceDPSink)

#
# For using our predefined device specific timings, you can use 'timing manager' with which you can get them.
#
timing_manager = role.dptx.pg.timing_manager

color_mode = UniTAP.ColorInfo()
color_mode.color_format = UniTAP.ColorInfo.ColorFormat.CF_RGB
color_mode.bpc = 8
color_mode.colorimetry = UniTAP.ColorInfo.Colorimetry.CM_sRGB

target_width = 1920
target_height = 1080
path_custom_image = "Your_Path_To_Image.jpg"

vf = get_vf_from_image(path_custom_image, target_width, target_height)

params = CompressionInfo()
params.color_format = CompressionInfo.DscColorFormat.CF_RGB
params.bpp = 128
params.version = (1, 2)
params.v_slice_size = calculate_dsc_slice_size(target_width, 4)
params.h_slice_size = calculate_dsc_slice_size(target_height, 4)
params.buffer_bit_depth = vf.color_info.bpc + 1

dsc_vf = encode_video_frame(vf, params)

video_mode = UniTAP.VideoMode(timing=timing_manager.get_cta(76), color_info=color_mode)

caps = role.dprx.link.capabilities.link_caps_status()
caps.dsc = True
caps.fec = True
role.dprx.link.capabilities.set(caps)

role.dptx.link.start_link_training()

role.dptx.pg.set_vm(video_mode)

# Also you can put image path (dsc image path also) to function 'set_pattern' as argument.
# path_custom_image = "Your_Path_To_Image.jpg"
# role.dptx.pg.set_pattern(path_custom_image)
role.dptx.pg.set_pattern(dsc_vf)
role.dptx.pg.apply()
res_app = role.dptx.pg.status().error
print(f"Stream {0} - Apply {res_app.__str__()}")

role.dprx.video_capturer.start()
dsc_captured_frame = role.dprx.video_capturer.pop_element()
role.dprx.video_capturer.stop()

print(dsc_captured_frame.is_compressed())
print(dsc_captured_frame)

role.dptx.pg.set_pattern(UniTAP.VideoPattern.ColorSquares)
role.dptx.pg.apply()
res_app = role.dptx.pg.status().error
print(f"Stream {0} - Apply {res_app.__str__()}")

role.dprx.video_capturer.start()
captured_frame = role.dprx.video_capturer.pop_element()
role.dprx.video_capturer.stop()
print(captured_frame.is_compressed())
print(captured_frame)

#
# Since the 3.5 version, TsiLib and TSIDevice objects have the option to be closed earlier.
# TSIDevice can be closed with the TsiLib method close(). TsiLib can be closed with cleanup().
# Clean up will close all opened devices and block ability to open any devices
# with same TsiLib object.
#
lUniTAP.close(dev)

lUniTAP.cleanup()