import json
from pathlib import Path
import subprocess
import tempfile

import matplotlib.pyplot as plt
import numpy as np
import numpy.typing as npt
from tqdm import tqdm

from flandre.utils.RfFrame import b2t
from flandre.utils.RfMeta import RfFrameMeta
from flandre.utils.archive import to_zip
from flandre.utils.robot import rtsi_float2int


def match(a: int, b: npt.NDArray[np.int_]):
    arr = b.copy()
    arr -= a
    arr = np.abs(arr)
    argmin = np.argmin(arr)
    m: int = b[argmin]
    return m, argmin, a - m


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

robot_folder = rec_folder / "robot"
device_folder = rec_folder / "device"
camera_folder = rec_folder / "camera"

device_ts_map = {int(file.stem): file for file in device_folder.glob("*.bin")}
camera_ts_map = {int(file.stem): file for file in camera_folder.glob("*.jpg")}

device_host_arr = np.array(sorted(device_ts_map.keys()))[100:-100]
device_host_zero = device_host_arr[0]

robot_ts_map = {
    int(k): v for k, v in json.loads((rec_folder / "robot.json").read_text()).items()
}

device_device_list: list[int] = []
device_seq_list: list[int] = []
for k in device_host_arr:
    seq, encoder, host_ts, device_deivce, buffer = b2t(
        device_ts_map[k].open("rb").read(100)
    )
    device_device_list.append(device_deivce * 1000)
    device_seq_list.append(seq)
device_device_arr = np.array(device_device_list)
device_seq_arr = np.array(device_seq_list)
device_device_zero = device_device_arr[0]
robot_host_arr = np.array(sorted(robot_ts_map.keys()))
robot_device_arr = np.array([robot_ts_map[k]["device_ts"] for k in robot_host_arr])
camera_host_arr = np.array(sorted(camera_ts_map.keys()))

device_host0 = device_host_arr[0]
device_host1 = device_host_arr[-1]
m0, robot_host_id0, robot_d0 = match(device_host0, robot_host_arr)
robot_device0 = robot_ts_map[m0]["device_ts"]
m1, robot_host_id1, robot_d1 = match(device_host1, robot_host_arr)
robot_device1 = robot_ts_map[m1]["device_ts"]

camera_host0, camera_host_id0, camera_d0 = match(device_host0, camera_host_arr)
camera_host1, camera_host_id1, camera_d1 = match(device_host1, camera_host_arr)

seq, encoder, host_ts, device_deivce0, buffer = b2t(
    device_ts_map[device_host0].read_bytes()
)
device_deivce0 *= 10**3
seq, encoder, host_ts, device_deivce1, buffer = b2t(
    device_ts_map[device_host1].read_bytes()
)
device_deivce1 *= 10**3

device_host_diff = device_host1 - device_host0
device_deivce_diff = device_deivce1 - device_deivce0
robot_device_diff = robot_device1 - robot_device0
camera_host_diff = camera_host1 - camera_host0

print(robot_d0 / 10**9, robot_d1 / 10**9)
print(camera_d0 / 10**9, camera_d1 / 10**9)

print("device_host_diff", device_host_diff / 10**9)
print("device_deivce_diff", device_deivce_diff / 10**9)
print("robot_device_diff", robot_device_diff / 10**9)
print("camera_host_diff", camera_host_diff / 10**9)

print("device_host - device_deivce", (device_host_diff - device_deivce_diff) / 10**9)
print("device_host - robot_device", (device_host_diff - robot_device_diff) / 10**9)
print("device_host - camera_host", (device_host_diff - camera_host_diff) / 10**9)
print("device_deivce - robot_device", (device_deivce_diff - robot_device_diff) / 10**9)

device_device_cut = device_device_arr - device_device_zero

robot_device_cut = (
    robot_device_arr[robot_host_id0 : robot_host_id1 + 1]
    - robot_device_arr[robot_host_id0]
    - robot_d0
)

camera_host_cut = (
    camera_host_arr[camera_host_id0 : camera_host_id1 + 1] - device_host_zero
)

# x = (device_host_arr - device_host_zero) / 10**9
# y = np.zeros_like(x)
# plt.scatter(x, y)

# x = (robot_host_arr[robot_host_id0 : robot_host_id1 + 1] - device_host_zero) / 10**9
# y = np.zeros_like(x) + 1
# plt.scatter(x, y)

# x = camera_host_cut / 10**9
# y = np.zeros_like(x) + 2
# plt.scatter(x, y)

# x = device_device_cut / 10**9
# y = np.zeros_like(x) + 3
# plt.scatter(x, y)

# x = robot_device_cut / 10**9
# y = np.zeros_like(x) + 4
# plt.scatter(x, y)

# plt.show()

avif_folder = Path(tempfile.gettempdir()) / "avif_temp"
avif_folder.mkdir(exist_ok=True)

zip_li: list[tuple[bytes, RfFrameMeta, list[tuple[Path, str]]]] = []
for i, k in enumerate(tqdm(device_device_cut)):
    res1, idx1, offset1 = match(k, robot_device_cut)
    res2, idx2, offset2 = match(k, camera_host_cut)

    r = robot_ts_map[robot_host_arr[idx1 + robot_host_id0]]
    c = camera_ts_map[camera_host_arr[idx2 + camera_host_id0]]

    target = avif_folder / c.with_suffix(".avif").name
    if not target.exists():
        subprocess.run(
            [
                "avifenc",
                "-y",
                "420",
                "-q",
                "50",
                c,
                target,
            ],
            stdout=subprocess.DEVNULL,
        )
    file = device_ts_map[device_host_arr[i]]
    seq, encoder, host_ts, device_ts, buffer = b2t(file.read_bytes())

    robot_rtsi = rtsi_float2int(r)

    zip_li.append(
        (
            buffer,
            RfFrameMeta(
                sequence_id=seq,
                encoder=encoder,
                robot_x=robot_rtsi["x"],
                robot_y=robot_rtsi["y"],
                robot_z=robot_rtsi["z"],
                robot_rx=robot_rtsi["rx"],
                robot_ry=robot_rtsi["ry"],
                robot_rz=robot_rtsi["rz"],
                robot_force_x=robot_rtsi["fx"],
                robot_force_y=robot_rtsi["fy"],
                robot_force_z=robot_rtsi["fz"],
                robot_force_rx=robot_rtsi["frx"],
                robot_force_ry=robot_rtsi["fry"],
                robot_force_rz=robot_rtsi["frz"],
            ),
            [(c, ".avif")],
        )
    )
to_zip(zip_li, Path("/tmp/f1"), Path("/tmp/f2"))
# print(
#     offset1 * 10**-9,
#     # (device_host_arr[i] - r["host_ts"]) * 10**-9,
#     offset2 * 10**-9,
#     # (device_host_arr[i] - camera_host_arr[idx2 + camera_host_id0]) * 10**-9,
# )