"""
示例 1: 基础用法 - 单传感器同步处理

本示例演示 orisys 的基本使用流程：
1. 初始化传感器（视频文件或摄像头）
2. 循环获取图像并计算形变和触碰点
3. 读取各类传感器数据
4. 可视化显示


"""
import os
os.environ["OPENCV_VIDEOIO_MSMF_ENABLE_HW_TRANSFORMS"] = "0" 
import cv2
import orisys
import numpy as np
import argparse
import time


def _tile_to_bgr(img):
    if img.ndim == 2:
        return cv2.cvtColor(img, cv2.COLOR_GRAY2BGR)
    return img


def _tile_to_gray_bgr(bgr):
    if bgr.ndim == 2:
        return bgr
    return cv2.cvtColor(bgr, cv2.COLOR_BGR2GRAY)


def _ensure_u8_gray(gray):
    if gray.dtype == np.uint8:
        return gray
    return cv2.normalize(gray, None, 0, 255, cv2.NORM_MINMAX).astype(np.uint8)


def _preprocess_gray_for_orb(g0, g1):
    g0 = cv2.GaussianBlur(g0, (5, 5), 0)
    g1 = cv2.GaussianBlur(g1, (5, 5), 0)
    g0 = cv2.equalizeHist(g0)
    g1 = cv2.equalizeHist(g1)
    return g0, g1


def _translation_from_H_to_ref(H):
    """
    Use only the translation part of H (tile -> reference): where tile origin (0,0) lands in ref coords.
    Ignores rotation/scale/shear for fast compositing (slice paste).
    """
    if H is None:
        return 0.0, 0.0
    H = np.asarray(H, dtype=np.float64)
    if abs(H[2, 2]) > 1e-12:
        H = H / H[2, 2]
    pt = H @ np.array([[0.0], [0.0], [1.0]], dtype=np.float64)
    return float(pt[0, 0]), float(pt[1, 0])


def _paste_bgr(dst, src, y0, x0):
    """Paste src into dst at top-left (y0, x0), clipping to dst bounds."""
    h, w = src.shape[:2]
    H, W = dst.shape[:2]
    y0d, x0d = max(0, y0), max(0, x0)
    y1d, x1d = min(H, y0 + h), min(W, x0 + w)
    if y0d >= y1d or x0d >= x1d:
        return
    sy0 = y0d - y0
    sx0 = x0d - x0
    sh = y1d - y0d
    sw = x1d - x0d
    dst[y0d:y1d, x0d:x1d] = src[sy0 : sy0 + sh, sx0 : sx0 + sw]


def _accumulate_tile_bgr(acc, wsum, src, y0, x0):
    """
    Add tile BGR into float accumulators for mean blending where tiles overlap.
    acc: (H,W,3) float32, wsum: (H,W) float32 counts per pixel.
    """
    h, w = src.shape[:2]
    H, W = acc.shape[:2]
    y0d, x0d = max(0, y0), max(0, x0)
    y1d, x1d = min(H, y0 + h), min(W, x0 + w)
    if y0d >= y1d or x0d >= x1d:
        return
    sy0 = y0d - y0
    sx0 = x0d - x0
    sh = y1d - y0d
    sw = x1d - x0d
    patch = src[sy0 : sy0 + sh, sx0 : sx0 + sw].astype(np.float32, copy=False)
    acc[y0d:y1d, x0d:x1d] += patch
    wsum[y0d:y1d, x0d:x1d] += 1.0


def _finalize_mean_blend(acc, wsum):
    """Average accumulated colors; pixels with no coverage stay 0."""
    out_h, out_w = wsum.shape
    stitched = np.zeros((out_h, out_w, 3), dtype=np.uint8)
    valid = wsum > 1e-6
    stitched[valid] = np.clip(acc[valid] / wsum[valid, None], 0.0, 255.0).astype(np.uint8)
    return stitched


def _h_list_signature(H_use_list):
    """Stable signature to detect when cached canvas geometry must be rebuilt."""
    parts = []
    for H in H_use_list:
        if H is None:
            parts.append("None")
        else:
            parts.append(np.asarray(H, dtype=np.float64).tobytes())
    return tuple(parts)


# Tuning (if stitch position is wrong or unstable):
# 1) ORB: nfeatures, fastThreshold, edgeThreshold in ORB_create.
# 2) Lowe ratio: lower => stricter matches; higher => more matches, more outliers.
# 3) ransac_reproj_thresh: affine/homography RANSAC inlier threshold (pixels).
# 4) Temporal: smooth_alpha, min_inliers, max_mean_reproj_err.
def estimate_bfmatcher_stitch_transform(
    t0,
    t1,
    *,
    H_prev=None,
    smooth_alpha=0.2,
    min_inliers=30,
    max_mean_reproj_err=8.0,
    ransac_reproj_thresh=5.0,
    ratio=0.9,
    orb_nfeatures=5000,
    verbose=True,
):
    """
    Estimate 3x3 transform H_use mapping t1 coordinates -> t0 coordinates (t0 is reference).

    Parameters
    ----------
    t0, t1 : ndarray
        Tile images (BGR or grayscale).
    H_prev : ndarray or None
        Previous frame's H_use for temporal smoothing; pass None on first frame.

    Returns
    -------
    H_use : ndarray (3,3) or None
        Normalized homography/affine-as-homography. None if not enough features or no valid model.
    H_prev_out : ndarray or None
        Value to pass as H_prev on the next call (same as H_use when successful, else unchanged H_prev).
    info : dict
        Keys may include: ok, method, inliers, mean_err, good_matches, k0, k1, t0_bgr, t1_bgr (for viz).
    """
    t0_bgr = _tile_to_bgr(t0)
    t1_bgr = _tile_to_bgr(t1)
    g0 = _ensure_u8_gray(_tile_to_gray_bgr(t0_bgr))
    g1 = _ensure_u8_gray(_tile_to_gray_bgr(t1_bgr))

    if verbose:
        print(
            f"BF stitcher debug: t0 dtype={t0.dtype}, t1 dtype={t1.dtype}, "
            f"g0 dtype={g0.dtype}, g0 range=({np.min(g0)}, {np.max(g0)}), "
            f"g1 dtype={g1.dtype}, g1 range=({np.min(g1)}, {np.max(g1)})"
        )

    g0, g1 = _preprocess_gray_for_orb(g0, g1)

    orb = cv2.ORB_create(
        nfeatures=orb_nfeatures, scaleFactor=1.2, nlevels=8, edgeThreshold=15, fastThreshold=5
    )
    k0, d0 = orb.detectAndCompute(g0, None)
    k1, d1 = orb.detectAndCompute(g1, None)

    info = {
        "ok": False,
        "method": None,
        "inliers": 0,
        "mean_err": None,
        "good_matches": 0,
        "k0": k0,
        "k1": k1,
        "t0_bgr": t0_bgr,
        "t1_bgr": t1_bgr,
    }

    if d0 is None or d1 is None or len(k0) < 4 or len(k1) < 4:
        if verbose:
            print(
                "OpenCV BFMatcher stitcher: not enough features.",
                f"k0={0 if k0 is None else len(k0)}",
                f"k1={0 if k1 is None else len(k1)}",
                f"d0_is_none={d0 is None}",
                f"d1_is_none={d1 is None}",
            )
        return None, H_prev, info

    bf = cv2.BFMatcher(cv2.NORM_HAMMING, crossCheck=False)
    matches = bf.knnMatch(d0, d1, k=2)

    good = []
    for m_n in matches:
        if len(m_n) != 2:
            continue
        m, n = m_n
        if m.distance < ratio * n.distance:
            good.append(m)
    info["good_matches"] = len(good)
    info["good"] = good
    if verbose:
        print(f"BFMatcher good matches: {len(good)}")

    if len(good) < 4:
        return None, H_prev, info

    pts0 = np.float32([k0[m.queryIdx].pt for m in good]).reshape(-1, 1, 2)
    pts1 = np.float32([k1[m.trainIdx].pt for m in good]).reshape(-1, 1, 2)

    method = "affine_partial2d"
    M, inliers_aff = cv2.estimateAffinePartial2D(
        pts1, pts0, method=cv2.RANSAC, ransacReprojThreshold=ransac_reproj_thresh
    )
    if M is not None:
        H = np.vstack([M, [0, 0, 1]]).astype(np.float64)
        inlier_mask = inliers_aff
    else:
        method = "homography"
        H, inlier_mask = cv2.findHomography(pts1, pts0, cv2.RANSAC, ransac_reproj_thresh)

    inliers = int(inlier_mask.sum()) if inlier_mask is not None else 0
    info["method"] = method
    info["inliers"] = inliers
    if verbose:
        print(f"Estimated {method} inliers: {inliers}")

    if H is None:
        info["transform_none"] = True
        return None, H_prev, info

    if abs(H[2, 2]) > 1e-12:
        H = H / H[2, 2]

    mean_err = 1e9
    if inlier_mask is not None:
        inlier_bool = inlier_mask.ravel().astype(bool)
        if np.any(inlier_bool):
            projected = cv2.perspectiveTransform(pts1[inlier_bool], H)
            err = np.linalg.norm(
                projected.reshape(-1, 2) - pts0[inlier_bool].reshape(-1, 2), axis=1
            )
            mean_err = float(err.mean()) if err.size else 1e9
    info["mean_err"] = mean_err

    H_use = H
    if H_prev is not None:
        if inliers < min_inliers or mean_err > max_mean_reproj_err:
            H_use = H_prev
        else:
            H_use = smooth_alpha * H + (1.0 - smooth_alpha) * H_prev
            if abs(H_use[2, 2]) > 1e-12:
                H_use = H_use / H_use[2, 2]

    info["ok"] = True
    return H_use, H_use.copy(), info


def render_stitched_preview(
    t0_bgr,
    t1_bgr,
    H_use,
    *,
    pad=None,
    frame_idx=0,
    print_every=30,
    verbose=True,
):
    """
    Warp t1 with H_use (t1->t0), paste t0 on a fixed padded canvas. Returns stitched BGR image.
    """
    if H_use is None:
        return None
    h0c, w0c = t0_bgr.shape[:2]
    if pad is None:
        pad = max(50, int(0.1 * w0c))
    out_w = w0c + 2 * pad
    out_h = h0c + 2 * pad
    T_pad = np.array([[1, 0, pad], [0, 1, pad], [0, 0, 1]], dtype=np.float64)
    H_total = T_pad @ H_use

    if verbose and frame_idx % print_every == 0:
        h1c, w1c = t1_bgr.shape[:2]
        corners1 = np.float32([[0, 0], [w1c, 0], [w1c, h1c], [0, h1c]]).reshape(-1, 1, 2)
        corners1_warped = cv2.perspectiveTransform(corners1, H_total).reshape(-1, 2)
        corners0_out = np.float32(
            [
                [pad, pad],
                [pad + w0c, pad],
                [pad + w0c, pad + h0c],
                [pad, pad + h0c],
            ]
        )
        print("=== Stitch transform debug ===")
        print("H_use (t1 -> t0):\n", H_use)
        print("H_total (output canvas):\n", H_total)
        print("t1 corners warped into output (x,y):", corners1_warped)
        print("t0 corners in output (x,y):", corners0_out)

    stitched = cv2.warpPerspective(t1_bgr, H_total, (out_w, out_h), flags=cv2.INTER_LINEAR)
    stitched[pad : pad + h0c, pad : pad + w0c] = t0_bgr
    return stitched


def _geom_cache_key(t_ref_bgr, tiles_bgr_c, H_use_list, pad):
    """Invalidate canvas layout when ref/tile sizes, transforms, or pad change."""
    ref_h, ref_w = t_ref_bgr.shape[:2]
    shapes = tuple(tb.shape[:2] for tb in tiles_bgr_c)
    return (ref_h, ref_w, shapes, pad, _h_list_signature(H_use_list))


def render_multi_tiles_preview(
    t_ref_bgr,
    tiles_bgr,
    H_use_list,
    *,
    pad=60,
    frame_idx=0,
    print_every=30,
    verbose=False,
    canvas_cache=None,
):
    """
    Composite multiple tiles into one canvas using translation only (where each tile's
    origin maps in ref coords). No per-tile warpPerspective.

    - t_ref_bgr: reference tile; output is in a shifted ref-aligned frame.
    - each H_use maps tile_i -> reference; only the translation of origin (0,0) is used.
    - tiles_bgr and H_use_list must have the same length.
    - canvas_cache: optional dict, updated in place to reuse out_w/out_h and paste origins
      when geometry is unchanged.
    Overlapping regions use mean blending (equal weight per contributing tile).
    """
    if H_use_list is None:
        return None

    if canvas_cache is None:
        canvas_cache = {}

    t_ref_bgr = _tile_to_bgr(t_ref_bgr)
    h0c, w0c = t_ref_bgr.shape[:2]
    tiles_bgr_c = [_tile_to_bgr(tb) for tb in tiles_bgr]

    geom_key = _geom_cache_key(t_ref_bgr, tiles_bgr_c, H_use_list, pad)

    if canvas_cache.get("geom_key") != geom_key:
        ref_corners = np.array([[0, 0], [w0c, 0], [w0c, h0c], [0, h0c]], dtype=np.float64)
        pts = [ref_corners]
        for tile_bgr, H_use in zip(tiles_bgr_c, H_use_list):
            if H_use is None:
                continue
            hi, wi = tile_bgr.shape[:2]
            tx, ty = _translation_from_H_to_ref(H_use)
            pts.append(
                np.array(
                    [[tx, ty], [tx + wi, ty], [tx + wi, ty + hi], [tx, ty + hi]],
                    dtype=np.float64,
                )
            )
        all_xy = np.concatenate(pts, axis=0)
        x_min_f = float(all_xy[:, 0].min()) - pad
        y_min_f = float(all_xy[:, 1].min()) - pad
        x_max_f = float(all_xy[:, 0].max()) + pad
        y_max_f = float(all_xy[:, 1].max()) + pad

        out_w = max(1, int(np.ceil(x_max_f - x_min_f)))
        out_h = max(1, int(np.ceil(y_max_f - y_min_f)))
        ref_y0 = int(np.floor(0.0 - y_min_f))
        ref_x0 = int(np.floor(0.0 - x_min_f))

        pastes = []
        for tile_bgr, H_use in zip(tiles_bgr_c, H_use_list):
            if H_use is None:
                pastes.append(None)
                continue
            tx, ty = _translation_from_H_to_ref(H_use)
            pastes.append((int(np.floor(ty - y_min_f)), int(np.floor(tx - x_min_f))))

        canvas_cache["geom_key"] = geom_key
        canvas_cache["out_w"] = out_w
        canvas_cache["out_h"] = out_h
        canvas_cache["ref_y0"] = ref_y0
        canvas_cache["ref_x0"] = ref_x0
        canvas_cache["pastes"] = pastes
        canvas_cache["x_min_f"] = x_min_f
        canvas_cache["y_min_f"] = y_min_f
    else:
        out_w = canvas_cache["out_w"]
        out_h = canvas_cache["out_h"]
        ref_y0 = canvas_cache["ref_y0"]
        ref_x0 = canvas_cache["ref_x0"]
        pastes = canvas_cache["pastes"]

    acc = np.zeros((out_h, out_w, 3), dtype=np.float32)
    wsum = np.zeros((out_h, out_w), dtype=np.float32)
    _accumulate_tile_bgr(acc, wsum, t_ref_bgr, ref_y0, ref_x0)
    for tile_bgr, H_use, paste in zip(tiles_bgr_c, H_use_list, pastes):
        if H_use is None or paste is None:
            continue
        y0, x0 = paste
        _accumulate_tile_bgr(acc, wsum, tile_bgr, y0, x0)
    stitched = _finalize_mean_blend(acc, wsum)
    
    # Crop away any uncovered regions (which correspond to black padding at borders).
    # Use `wsum` for robustness even if a tile's actual content is black.
    valid = wsum > 1e-6
    if np.any(valid):
        ys, xs = np.where(valid)
        y0c, y1c = int(ys.min()), int(ys.max()) + 1
        x0c, x1c = int(xs.min()), int(xs.max()) + 1
        stitched = stitched[y0c:y1c, x0c:x1c]

    stitched = cv2.GaussianBlur(stitched, (5, 5), 0)


    if verbose and frame_idx % print_every == 0:
        xmf = canvas_cache.get("x_min_f", 0.0)
        ymf = canvas_cache.get("y_min_f", 0.0)
        print(
            f"render_multi_tiles_preview: canvas out_w={out_w}, out_h={out_h}, "
            f"x_min={xmf:.2f}, y_min={ymf:.2f}"
        )

    return stitched


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--video","-v", type=str, default="0", help="Camera ID / Video file path")
    parser.add_argument("--config","-c", type=str, default="./config/ddjx01.json", help="Configuration file path")
    parser.add_argument("--cal","-cal", type=str, default="./config/ddjx01.npy", help="Calibration file path")
    parser.add_argument("--verbose","-verbose", type=bool, default=True, help="Print detailed information")
    args = parser.parse_args()

    # 步骤1：创建传感器对象 (必要步骤)，
    # config_name为配置文件名（内置配置文件名）或路径（自定义配置文件路径）
    # cal_path为校准文件路径, 每一个设备需要创建/使用一个校准文件

    # 方式1: 从视频文件读取
    try:
        input = int(args.video)
        is_camera = True
        sensor = orisys.Sensor(input, config_name=args.config, cal_path= args.cal, verbose=args.verbose)
    except:
        is_camera = False
        sensor = orisys.Sensor(args.video, config_name=args.config, cal_path= args.cal, verbose=args.verbose)

    print("\n按 'q' 退出程序，按 'r' 重置传感器追踪器")

    # For stabilizing frame-to-frame stitching (reduce oscillation).
    # Stores the last good transform (mapping t1 -> t0) and exponentially smooths it.
    smooth_alpha = 0.2  # 0..1; higher => less smoothing, more jitter
    min_inliers = 30    # if inliers are too low, keep previous transform
    max_mean_reproj_err = 8.0  # if reprojection error is too high, keep previous transform
    frame_idx = 0
    print_every = 30  # reduce console spam; set to 1 for per-frame output

    # Stitch a 3x3 (9 tiles) block from the row-major `tile_list`.
    # For your 3x5 layout (nrow=3, ncol=5), this stitches the "middle 3x3"
    # using the center tile of that 3x3 as the reference frame.
    H_prev_tiles = {}  # key: tile_list index, value: H_prev
    # Cache BFMatcher-estimated transforms from the first frame.
    H_use_list_cached = None
    moving_indices_cached = None
    crop_pad = 5  # tighter crop in stitched canvas
    canvas_cache = {}

    while True:
        # 步骤2：信息提取
        # 步骤 2.1: 获取并拼接图像 (必要步骤)
        img = sensor.get_img()
        if img is None:
            print("Failed to get image (end of stream or capture error). Exiting.")
            break
        # boxes = sensor.stitcher.boxes
        # draw boxes for checking
        # for box in boxes:
        #     x0, y0, x1, y1 = box
        #     cv2.rectangle(sensor.frame, (x0, y0), (x1, y1), (0, 255, 0), 2)
        # cv2.imshow("boxes", sensor.frame)
        # cv2.waitKey(0)
        
        tiles = sensor.stitcher.tile_list
        frame_idx += 1

        ncol = getattr(sensor.stitcher, "ncol", 0) or 0
        nrow = getattr(sensor.stitcher, "nrow", 0) or 0

        grid_r, grid_c = 3, 3
        needed = grid_r * grid_c
        if len(tiles) < needed or ncol < grid_c or nrow < grid_r:
            print(f"tile_list too small for 9-tile stitch: len={len(tiles)} ncol={ncol} nrow={nrow} needed={needed}")
            cv2.waitKey(1)
            continue

        # Center the 3x3 block in the full (nrow x ncol) layout.
        # Use the center tile of that 3x3 as reference frame.
        start_r = (nrow - grid_r) // 2
        start_c = (ncol - grid_c) // 2

        idx_ref = (start_r + 1) * ncol + (start_c + 1)
        if idx_ref < 0 or idx_ref >= len(tiles):
            print("Computed reference tile index out of range:", idx_ref)
            cv2.waitKey(1)
            continue

        t_ref = tiles[idx_ref]

        # Stitcher: warp each moving tile into the reference tile coordinate system.
        moving_tiles = []
        moving_indices = []
        for r in range(grid_r):
            for c in range(grid_c):
                idx = (start_r + r) * ncol + (start_c + c)
                if idx == idx_ref:
                    continue
                if idx < 0 or idx >= len(tiles):
                    continue
                moving_tiles.append(tiles[idx])
                moving_indices.append(idx)

        # Only run BFMatcher transform estimation on the first frame.
        if H_use_list_cached is None or moving_indices_cached != moving_indices:
            H_use_list = []
            do_verbose_first = (frame_idx % print_every == 0)
            for j, (tile, tile_idx) in enumerate(zip(moving_tiles, moving_indices)):
                H_prev_j = H_prev_tiles.get(tile_idx)
                H_use_j, H_prev_j_out, info_j = estimate_bfmatcher_stitch_transform(
                    t_ref,
                    tile,
                    H_prev=H_prev_j,
                    smooth_alpha=smooth_alpha,
                    min_inliers=min_inliers,
                    max_mean_reproj_err=max_mean_reproj_err,
                    verbose=(do_verbose_first and j == 0),
                )
                H_prev_tiles[tile_idx] = H_prev_j_out
                H_use_list.append(H_use_j)

                if (
                    j == 0
                    and do_verbose_first
                    and info_j.get("good") is not None
                    and len(info_j["good"]) >= 4
                ):
                    draw = cv2.drawMatches(
                        info_j["t0_bgr"],
                        info_j["k0"],
                        info_j["t1_bgr"],
                        info_j["k1"],
                        info_j["good"],
                        None,
                        matchColor=(0, 255, 0),
                        singlePointColor=(255, 0, 0),
                        flags=cv2.DrawMatchesFlags_NOT_DRAW_SINGLE_POINTS,
                    )
                    cv2.imshow("BFMatcher Matches (t_ref vs t_moving0)", draw)

            H_use_list_cached = H_use_list
            moving_indices_cached = list(moving_indices)
        else:
            H_use_list = H_use_list_cached

        start_time = time.perf_counter()
        stitched = render_multi_tiles_preview(
            t_ref,
            moving_tiles,
            H_use_list,
            pad=crop_pad,
            frame_idx=frame_idx,
            print_every=print_every,
            verbose=False,
            canvas_cache=canvas_cache,
        )
        end_time = time.perf_counter()
        print(f"render_multi_tiles_preview: {end_time - start_time} seconds")
        
        if stitched is not None:
            cv2.imshow("OpenCV Stitch (9 tiles, middle-ref stabilized)", stitched)
        # cv2.imshow("tiles_1_2 (two rows, 2nd shifted)", canvas)
        # cv2.waitKey(0)
        # 步骤 2.2: 计算形变场 (必要步骤)
        # is_touched = sensor.compute_deformation(check_motion=True, threshold=6)
        # 步骤 2.3: 计算触碰点 
       
        #步骤 3.5: 键盘控制
        key = cv2.waitKey(1) & 0xFF
        if key == ord("q"):
            print("\n退出程序...")
            break
        if key == ord("r"):
            # 重置光流追踪器（当追踪漂移时使用）
            sensor.reset()
            print("传感器已重置")
        
    # =========================================================================
    # 4. 清理资源
    # =========================================================================
    sensor.disconnect()  # 释放视频源和相关资源


if __name__ == '__main__':
    main()