"""Fast FoundationStereo вместо CREStereo, на том же прогоне потока 700 мм.

Модель взята ONNX-экспортом (23_36_37, iters 4, 320x736) и гоняется через onnxruntime-gpu:
это снимает зависимость от окружения репозитория (torch 2.6 + xformers), которое с
питоном Isaac не совпадает. Веса лежат в fast-foundationstereo/weights.

ВХОД У МОДЕЛИ ФИКСИРОВАННЫЙ - 320x736. Кроп приводится к этому размеру, а полученная
диспаратность масштабируется ОБРАТНО по горизонтали: диспаратность измеряется в пикселях,
и при сжатии кадра в k раз по ширине она сжимается во столько же. Без этого множителя
глубина уехала бы ровно в k раз - самая частая ошибка при подмене стереосети.

Сравнение честное: сегментация, отделение полотна, слияние и метрики те же, что в
варианте без сегментации (лучшем на сегодня). Меняется ТОЛЬКО стереодвижок.
"""
import os, sys, json, time
import numpy as np
import cv2

CV = "/home/dasha/isaac_assets/cv"
CT = "/home/dasha/robozon-sorter/control_test"
FFS = f"{CV}/fast-foundationstereo/weights/ffs_23_36_37_iters4_320x736.onnx"
OUT = f"{CT}/diag/ffs"
os.makedirs(OUT, exist_ok=True)
sys.path.insert(0, CT)
os.environ.setdefault("MASK_MODE", "gate")
import measure_flow as MF
import classify as CL

import onnxruntime as ort
so = ort.SessionOptions()
so.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
sess = ort.InferenceSession(FFS, so, providers=["CUDAExecutionProvider", "CPUExecutionProvider"])
IN_H, IN_W = 320, 736
print(f"Fast FoundationStereo: {os.path.basename(FFS)}, вход {IN_W}x{IN_H}, "
      f"провайдер {sess.get_providers()[0]}")

man, calib = MF.man, MF.calib
TARGET, BELT_Z, RIGS = MF.TARGET, MF.BELT_Z, MF.RIGS
FLOOR, GATE_R = 0.020, 0.28
DENS_CELL, DENS_MIN = 0.010, 12
PITCH_S = 0.70


def ffs_disp(L, R):
    """диспаратность одной пары. Кроп -> 320x736 -> обратный масштаб по горизонтали."""
    # cv2.imread даёт BGR, а модель обучена на RGB (run_demo.py читает через imageio).
    # Нормализацию (ImageNet mean/std после деления на 255) модель делает ВНУТРИ forward,
    # см. core/foundation_stereo.py:32 - значит она экспортирована в ONNX, и подавать надо
    # именно СЫРЫЕ 0-255. Ошибкой был только порядок каналов.
    h, w = L.shape[:2]
    sx = IN_W / float(w)
    lr = cv2.resize(cv2.cvtColor(L, cv2.COLOR_BGR2RGB), (IN_W, IN_H)).astype(np.float32)
    rr = cv2.resize(cv2.cvtColor(R, cv2.COLOR_BGR2RGB), (IN_W, IN_H)).astype(np.float32)
    a = lr.transpose(2, 0, 1)[None]
    b = rr.transpose(2, 0, 1)[None]
    d = sess.run(["disparity"], {"left_image": a, "right_image": b})[0][0, 0]
    d = cv2.resize(d, (w, h), interpolation=cv2.INTER_LINEAR)
    return np.abs(d) / sx          # пиксели исходного кропа, а не приведённого


def dense_only(P, cell=DENS_CELL, need=DENS_MIN):
    if len(P) < 60:
        return np.ones(len(P), bool)
    key = np.floor(P[:, :2] / cell).astype(np.int64)
    uniq, inv, cnt = np.unique(key, axis=0, return_inverse=True, return_counts=True)
    return cnt[inv] >= need


def biggest_blob_idx(P, grid=0.012):
    if len(P) < 40:
        return np.ones(len(P), bool)
    cx = np.round(P[:, 0] / grid).astype(int); cy = np.round(P[:, 1] / grid).astype(int)
    im = np.zeros((cy.max() - cy.min() + 3, cx.max() - cx.min() + 3), np.uint8)
    im[cy - cy.min() + 1, cx - cx.min() + 1] = 255
    im = cv2.morphologyEx(im, cv2.MORPH_CLOSE, np.ones((3, 3), np.uint8))
    n, lab = cv2.connectedComponents(im)
    li = lab[cy - cy.min() + 1, cx - cx.min() + 1]
    best, bn = None, 0
    for k in range(1, n):
        m = li == k
        if m.sum() > bn:
            best, bn = m, m.sum()
    return np.ones(len(P), bool) if best is None else best


def cloud_from_roi(disp, win, cam):
    x0, y0, x1, y1 = win
    depth = np.where(disp > 0.5, cam["fx"] * cam["baseline"] / np.maximum(disp, 1e-6), np.nan)
    vs, us = np.mgrid[y0:y0 + depth.shape[0], x0:x0 + depth.shape[1]]
    ok = np.isfinite(depth) & (depth > 1e-3)
    u, v, z = us[ok], vs[ok], depth[ok]
    P = np.stack([(u - cam["cx"]) * z / cam["fx"],
                  -(v - cam["cy"]) * z / cam["fy"], -z, np.ones_like(z)], 1)
    P = (P @ np.array(cam["M"]))[:, :3]
    near = np.hypot(P[:, 0] - TARGET[0], P[:, 1] - TARGET[1]) < GATE_R
    return P[near & (P[:, 2] > BELT_Z + FLOOR) & (P[:, 2] < BELT_Z + 0.60)]


for _ in range(2):                                   # прогрев
    ffs_disp(np.zeros((240, 480, 3), np.uint8), np.zeros((240, 480, 3), np.uint8))
print("подача: RGB 0-255, нормализация внутри графа")

print(f"\n  {'товар':18s} {'эталон, мм':>18s} {'предсказано, мм':>20s} {'MAE':>6s} "
      f"{'k':>6s} {'класс':>12s} {'точек':>7s} {'мс':>6s}")
print("  " + "-" * 100)
rows, times = [], []
for name, e in man["items"].items():
    t0 = time.time()
    clouds = []
    for rig in RIGS:
        cam = calib[f"{rig}_Left"]
        IL = cv2.imread(e["files"][f"{rig}_Left"])
        IR = cv2.imread(e["files"][f"{rig}_Right"])
        if IL is None or IR is None:
            continue
        x0, y0, x1, y1 = MF.belt_roi_px(cam, MF.POLY3)
        maxd = int(np.ceil(MF.DPAD * cam["fx"] * cam["baseline"] / MF.ZMIN))
        x0 = max(0, x0 - maxd)
        d = ffs_disp(IL[y0:y1, x0:x1], IR[y0:y1, x0:x1])
        P = cloud_from_roi(d, (x0, y0, x1, y1), cam)
        if len(P):
            clouds.append(P)
        if name in ("bag", "bucket", "box_300x200x200"):
            dn = cv2.normalize(d, None, 0, 255, cv2.NORM_MINMAX).astype(np.uint8)
            cv2.imwrite(f"{OUT}/{name}_{rig}_disp.png", cv2.applyColorMap(dn, cv2.COLORMAP_TURBO))
    dt = (time.time() - t0) * 1000
    times.append(dt)
    gt = e["gt"]; gtd = sorted(gt["dims_mm"], reverse=True)
    if not clouds:
        print(f"  {name:18s} {str([round(v) for v in gtd]):>18s}   облако пустое"); continue
    P = np.vstack(clouds)
    s1 = dense_only(P)
    if s1.sum() >= 60:
        P = P[s1]
    P = P[biggest_blob_idx(P)]
    out = MF.dims_and_k(P)
    if out is None:
        print(f"  {name:18s} габариты не взялись"); continue
    dims, k = out
    mae = float(np.mean(np.abs(np.array(dims) - np.array(gtd))))
    cls = CL.classify(dims, 0.0 if np.isnan(k) else k)
    ok = "верно" if cls == gt["zone_scene"] else f"ОШ({gt['zone_scene']})"
    print(f"  {name:18s} {str([round(v) for v in gtd]):>18s} "
          f"{str([round(v) for v in dims]):>20s} {mae:6.1f} {k:6.2f} {cls+' '+ok:>12s} "
          f"{len(P):7d} {dt:6.0f}")
    rows.append(dict(name=name, gt=gtd, gt_cls=gt["zone_scene"], pred=[round(v, 1) for v in dims],
                     pred_k=round(float(k), 3), pred_cls=cls, mae=round(mae, 1),
                     n=len(P), ms=round(dt)))

print("\n  === МЕТРИКИ (Fast FoundationStereo, без сегментации) ===")
if rows:
    maes = [r["mae"] for r in rows]
    acc = sum(1 for r in rows if r["pred_cls"] == r["gt_cls"])
    print(f"  габариты: MAE медиана {np.median(maes):.1f} мм, среднее {np.mean(maes):.1f}, "
          f"худший {max(maes):.1f} ({max(rows, key=lambda r: r['mae'])['name']})")
    print(f"  классы: {acc}/{len(rows)} = {100.0*acc/len(rows):.0f}%")
    print(f"  k макс {max(r['pred_k'] for r in rows):.2f}")
    lab = ["B", "C", "D"]
    print("  матрица (строки истина, столбцы предсказание):")
    print("        " + "".join(f"{c:>5s}" for c in lab))
    for a in lab:
        print(f"     {a:3s}  " + "".join(
            f"{sum(1 for r in rows if r['gt_cls']==a and r['pred_cls']==b):5d}" for b in lab))
print(f"  время: медиана {np.median(times):.0f} мс, такт {PITCH_S*1000:.0f} мс -> "
      f"{'УКЛАДЫВАЕТСЯ' if np.median(times) < PITCH_S*1000 else 'НЕ УКЛАДЫВАЕТСЯ'} "
      f"(запас {PITCH_S*1000 - np.median(times):+.0f} мс)")
json.dump(rows, open(f"{OUT}/metrics.json", "w"), indent=1, ensure_ascii=False)
print(f"\n  -> {OUT}/metrics.json")