from flandre.utils.RfFrame import b2t
import json
from pathlib import Path
import itertools

import cv2
import numpy as np
from nvidia import nvimgcodec
decoder = nvimgcodec.Decoder()
encoder = nvimgcodec.Encoder()


rec_folder = Path(
    '/run/media/lambda/b86dccdc-f134-464b-a310-6575ee9ae85c/record2/')

robot_folder = rec_folder / 'robot'
robot_folder = rec_folder / 'robot'

d: dict[int, dict] = dict()
for file in robot_folder.glob('*'):
    for line in file.read_text().splitlines():
        s = line.split(',')
        v = dict(
            host_ts=int(s[0]),
            device_ts=int(float(s[1])*10**9),
            x=float(s[2]),
            y=float(s[3]),
            z=float(s[4]),
            rx=float(s[5]),
            ry=float(s[6]),
            rz=float(s[7]),
            fx=float(s[8]),
            fy=float(s[9]),
            fz=float(s[10]),
            frx=float(s[11]),
            fry=float(s[12]),
            frz=float(s[13]),
        )
        d[v['host_ts']] = v
(rec_folder/'robot.json').write_text(json.dumps(d))

for file in Path('/run/media/lambda/b86dccdc-f134-464b-a310-6575ee9ae85c/record2/camera_raw').glob('*.bin'):
    fb = Path(file).read_bytes()
    b = np.frombuffer(fb, dtype=np.uint8).reshape(
        (1080, 1920, 3))
    b = cv2.cvtColor(b, cv2.COLOR_BGR2RGB)
    # cv2.imshow('',)
    # cv2.waitKey(0)
    # enc_params = nvimgcodec.EncodeParams(quality=5, chroma_subsampling=nvimgcodec.ChromaSubsampling.CSS_GRAY)
    # encoder.write('a.jpg', b, params=enc_params)
    encoder.write((file.parent.parent/'camera' /
                  file.with_suffix('.jpg').name).__str__(), b)