#
# Import UniTAP module.
#

import UniTAP
import time

#
# 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)

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

timings = [timing_manager.get_cta(16),
           timing_manager.get_cta(12),
           timing_manager.get_cta(1),
           timing_manager.get_cta(19)]

# Also, possible to select timing with using functions: 'get_dmt', 'get_cvt', 'get_by_list_index', 'get_all'

# Print all available timings
print(timing_manager.print_all())

#
# If you want to select pattern, you need to use enum 'VideoPattern' from UniTAP.
# Like this: UniTAP.VideoPattern.ColorRamp
#

# Iterate for each available streams
for i in range(role.dptx.pg.max_stream_count):
    video_mode = UniTAP.VideoMode(timing=timings[i], color_info=color_mode)
    role.dptx.pg[i].set_vm(vm=video_mode)
    role.dptx.pg[i].set_pattern(pattern=UniTAP.VideoPattern.ColorRamp if i == 0 else UniTAP.VideoPattern.SolidGreen)
    role.dptx.pg[i].apply()
    res_app = role.dptx.pg[i].status().error
    print(f"Stream {i} - Apply {res_app.__str__()}")

#
# After configure Pattern Generators, MST should be enabled in link module as well as selected
# stream count
#

cfg = role.dptx.link.config.get()
cfg.mst = True
cfg.mst_stream_count = 4
role.dptx.link.config.set(cfg)
role.dptx.link.start_link_training()

cfg = role.dprx.link.capabilities.link_caps_status()
cfg.mst = True
cfg.dp_128_132_bitrates = None
cfg.override_10g = None
role.dprx.link.capabilities.set(cfg)
role.dprx.link.hpd_pulse()
#
# Just random sleep, to be sure that Sink get video
#
time.sleep(2)

for i in range(role.dptx.pg.max_stream_count):
    res = role.dprx.link.status.stream(i).video_mode.timing == timings[i]
    print(f'Stream {i}\n'
          f'Reference timing: {timings[i]}. '
          f'Received timing RX: {role.dprx.link.status.stream(i).video_mode.timing}.\n'
          f'Result: {res}')

#
# 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()