"""Рабочий процесс замкнутого контура: слушает заявки от Isaac, отвечает классом товара.

Зачем отдельный процесс. torch внутри Isaac роняет процесс, поэтому CV не может жить в
одном процессе со сценой. Обмен идёт через каталог: Isaac кладёт заявку с путями к шести
кадрам, работник отвечает файлом с габаритами, k и классом.

Времени на это хватает с запасом. От ворот осмотра (x = -0.750) до пушера (x = -3.900)
товар при 1 м/с едет 3.15 с, до плуга (x = -7.85) - 7.1 с, а инференс на зафиксированном
бейзлайне занимает 469 мс. То есть класс успевает к обоим механизмам, и контур замыкается
по-настоящему, а не постфактум.

БЕЙЗЛАЙН ЗАФИКСИРОВАН: DEFOM-Stereo vitl, вход сети 480, iters 24, кроп зоны осмотра,
без сегментации (товар отделяется превышением над плоскостью полотна + отсев по плотности).
На потоке 700 мм он дал классы 8/9 = 89 % и габариты MAE 32.8 мм.
"""
import os, sys, json, time, glob, traceback

os.environ["IMPORT_ONLY"] = "1"          # measure_plane импортируется как библиотека
CT = "/home/dasha/robozon-sorter/control_test"
sys.path.insert(0, CT)

import numpy as np
import cv2
import measure_plane as MP
import measure_flow as MF
import classify as CL

RT = f"{CT}/runtime"
REQ, RES = f"{RT}/req", f"{RT}/res"
for d in (REQ, RES):
    os.makedirs(d, exist_ok=True)

calib = MF.calib
TARGET, BELT_Z, RIGS = MF.TARGET, MF.BELT_Z, MF.RIGS


def measure(files):
    """шесть кадров -> габариты, k, класс. Тот же тракт, что в зафиксированном бейзлайне."""
    pairs, metas = [], []
    for rig in RIGS:
        cam = calib[f"{rig}_Left"]
        IL = cv2.imread(files[f"{rig}_Left"])
        IR = cv2.imread(files[f"{rig}_Right"])
        if IL is None or IR is None:
            continue
        x0, y0, x1, y1 = MP.gate_crop_px(cam)
        maxd = int(np.ceil(MF.DPAD * cam["fx"] * cam["baseline"] / MF.ZMIN))
        x0 = max(0, x0 - maxd)
        pairs.append((IL[y0:y1, x0:x1].astype(np.float32), IR[y0:y1, x0:x1].astype(np.float32)))
        metas.append(((x0, y0, x1, y1), cam))
    if not pairs:
        return None
    disps = MP.cre_scaled(pairs, MP.SW)          # движок выбирается внутри (бейзлайн: DEFOM)
    clouds = []
    for d, (win, cam) in zip(disps, metas):
        # cloud_from_roi отдаёт ПАРУ: облако товара и облако полотна. В первой версии я
        # складывал кортеж целиком, и vstack падал на разнородных формах.
        obj, _belt = MP.cloud_from_roi(d, win, cam)
        if len(obj):
            clouds.append(obj)
    if not clouds:
        return None
    P = np.vstack(clouds)
    sel = MP.dense_only(P)
    if sel.sum() >= 60:
        P = P[sel]
    P = P[MP.biggest_blob_idx(P)]
    out = MF.dims_and_k(P)
    if out is None:
        return None
    dims, _ = out
    k, is_round, sec = MP.circular_section_K(P)
    cls = CL.classify(dims, 0.0 if not k else k)
    return dict(dims=[round(v, 1) for v in dims], k=round(float(k), 3),
                cls=cls, n=len(P))


print(f"работник запущен: движок {MP.STEREO} {MP.DEFOM_CKPT}, вход {MP.SW}, "
      f"iters {MP.DEFOM_ITERS}, кроп {MP.CROP}")
print(f"заявки: {REQ}   ответы: {RES}")

# прогрев, чтобы первый настоящий товар не ждал загрузку энкодера
try:
    z = np.zeros((240, 480, 3), np.float32)
    MP.cre_scaled([(z, z)], MP.SW)
    print("прогрев выполнен")
except Exception as ex:
    print("прогрев не удался:", ex)

open(f"{RT}/worker_ready", "w").write(str(time.time()))
seen = set()
idle = 0.0
while True:
    reqs = sorted(glob.glob(f"{REQ}/*.json"))
    if not reqs:
        time.sleep(0.05); idle += 0.05
        if idle > 1800:
            print("1800 с без заявок - выхожу"); break
        continue
    idle = 0.0
    for rq in reqs:
        try:
            j = json.load(open(rq))
        except Exception:
            continue
        name = j["name"]
        t0 = time.time()
        try:
            r = measure(j["files"])
        except Exception:
            traceback.print_exc(); r = None
        dt = (time.time() - t0) * 1000
        ans = dict(name=name, ms=round(dt), ok=r is not None)
        if r:
            ans.update(r)
        json.dump(ans, open(f"{RES}/{name}.json", "w"), ensure_ascii=False)
        os.remove(rq)
        print(f"  {name}: {ans.get('cls','-')} dims={ans.get('dims')} k={ans.get('k')} "
              f"за {dt:.0f} мс", flush=True)