# =================================================================================================
# This script is written to check the CRC values on the TX and RX sides,
# as well as check the CRC for captured video frames.
# The main check is performed on standard patterns, but you can also use custom images.
# In order to check the DSC of an image, you need to change the functions to select the CRC values,
# as well as take a value from a stream variable link.status.stream(stream_number).dsc_crc.
# =================================================================================================


# Import UniTAP and others module.
import UniTAP
import time
from UniTAP.utils import calculate_crc

max_streams = 1
log_fid = open("rb_testing.log", 'w')


# =================================================================================================
#   Function. Setting and applying Pattern Generator (PG)
# =================================================================================================
def apply_PG(role, max_streams, timing, bpc, color_format, pg_progress_max, pg_progress_cur, pattern):
    color_mode = UniTAP.ColorInfo()
    if color_format == color_mode.ColorFormat.CF_RGB:
        color_mode.bpc = bpc
        color_mode.dynamic_range = color_mode.DynamicRange.DR_VESA
    else:
        color_mode.dynamic_range = color_mode.DynamicRange.DR_CTA
        # if bpc == 6:
        #    return 0
        color_mode.bpc = bpc
    color_mode.color_format = color_format
    color_mode.colorimetry = UniTAP.ColorInfo.Colorimetry.CM_ITUR_BT601
    if not color_mode.is_valid():
        print("Color mode " + color_mode.__str__() + " is not valid")
        exit(1)
    video_mode = UniTAP.VideoMode(timing=timing, color_info=color_mode)
    for stream in range(max_streams):
        role.pg[stream].set_vm(vm=video_mode)
        role.pg[stream].set_pattern(pattern=pattern)
        role.pg[stream].apply()
    # logger.debug("Set video mode: " + video_mode.__str__())
    # print("Set video mode: " + video_mode.__str__())
    progress(pg_progress_max, pg_progress_cur, "Timing: {}".format(video_mode), log_fid)
    time.sleep(0.5)


# =================================================================================================
#   Function. TX, RX CRCs comparison
# =================================================================================================
def check_CRCs_the_same(roles, max_streams, log_fid):
    for stream in range(max_streams):
        tx_crc = list(map(hex, roles.dptx.link.status.stream(stream).crc))
        rx_crc = list(map(hex, roles.dprx.link.status.stream(stream).crc))
        if tx_crc != rx_crc:
            msg = "@Error: stream " + stream.__str__() + ": CRCs are different: TX: " + tx_crc.__str__() + ";   RX: " + rx_crc.__str__()
            print(msg)
            log_fid.write(msg + '\n')


# =================================================================================================
#   Function. Check CRCs on captured frames: stability and equality with dprx.stream.crc
# =================================================================================================
def check_capturing(log_fid):
    buffers = 3
    role_rx.video_capturer.start(frames_count=buffers)
    role_rx.video_capturer.stop()
    result = role_rx.video_capturer.capture_result
    stream_crc = list(map(hex, role_rx.link.status.stream(0).crc))
    bin_fid = open('buf.dat', 'bw')
    bin_fid.write(result.buffer[0].data)
    bin_fid.close()
    for b in range(buffers):
        capt_crc = list(map(hex, calculate_crc(result.buffer[b])))
        if stream_crc != capt_crc:
            msg = "@Error: stream 0 CRCs are different: RX: " + stream_crc.__str__() + ";   Captured: " + capt_crc.__str__()
            print(msg)
            log_fid.write(msg + '\n')


# =================================================================================================
#   Function. Progress
# =================================================================================================
def progress(max, cur, msg='', log_fid=0):
    print("Progress: %3.2f%%: %s" % (cur * 100 / max, msg), end='\n')
    log_fid.write(msg + '\n')


# =================================================================================================
#   MAIN
# =================================================================================================
# -------------------------------------------------------------------------------------------------
#   Open roles and others init routines
# -------------------------------------------------------------------------------------------------
if __name__ == '__main__':
    # init UniTAP library wrapper
    lUniTAP = UniTAP.TsiLib()
    # open device and set the role
    serial_number = "2150C474"
    dev = lUniTAP.open(serial_number)
    roles = dev.select_role(UniTAP.dev.UCD500.USBCSourceUSBCSink)
    role_tx = roles.dptx
    role_rx = roles.dprx
    # -------------------------------------------------------------------------------------------------
    #   Setting up links
    # -------------------------------------------------------------------------------------------------
    # make LT
    role_tx.link.start_link_training()

    # wait some time after LT
    time.sleep(2)

    # -------------------------------------------------------------------------------------------------
    #   Get all video modes
    # -------------------------------------------------------------------------------------------------
    timing_manager = role_tx.pg.timing_manager
    color_mode = UniTAP.ColorInfo()
    # set color_formats list
    color_formats = [
        # color_mode.ColorFormat.CF_RGB,
        # color_mode.ColorFormat.CF_YCbCr_444,
        color_mode.ColorFormat.CF_YCbCr_422,
        # color_mode.ColorFormat.CF_YCbCr_420
    ]

    # set pattern (can be value from defaults, like UniTAP.VideoPattern.SolidWhite) or name of image
    pattern = UniTAP.VideoPattern.SolidWhite

    # pattern = "image.png"

    # set bpc list
    color_bpcs = [6, 8, 10, 12, 16]

    # get RB timings only (it was a test case of Maxim Zaitsev)
    rb_timings = []
    for t in timing_manager.get_all():
        if t.reduce_blanking > 0:
            rb_timings.append(t)

    pg_progress_max = len(rb_timings) * len(color_formats) * len(color_bpcs)
    pg_progress_cur = 0
    for t in rb_timings:
        for cf in color_formats:
            for bpc in color_bpcs:
                if cf != color_mode.ColorFormat.CF_RGB and bpc == 6:
                    pass
                else:
                    apply_PG(role_tx, max_streams, t, bpc, cf, pg_progress_max, pg_progress_cur, pattern)
                    check_CRCs_the_same(roles, max_streams, log_fid)
                    check_capturing(log_fid)
                    pg_progress_cur += 1
    log_fid.close()