import os
import time
import torch
import numpy as np
import rerun as rr
from lietorch import SE3
import vigs_backends

# --- headless safety (no Qt/GUI on cluster) ---
os.environ.setdefault("QT_QPA_PLATFORM", "offscreen")
os.environ.pop("DISPLAY", None)

# ---- helpers to draw a camera frustum as 3D line segments ----
_CAM_POINTS = np.array([
    [0, 0, 0],
    [-1, -1, 1.5],
    [ 1, -1, 1.5],
    [ 1,  1, 1.5],
    [-1,  1, 1.5],
    [-0.5, 1, 1.5],
    [ 0.5, 1, 1.5],
    [ 0, 1.2, 1.5]
], dtype=np.float32)

_CAM_LINES = np.array([
    [1,2], [2,3], [3,4], [4,1], [1,0], [0,2], [3,0], [0,4], [5,7], [7,6]
], dtype=np.int32)

def _camera_lines(scale=0.05):
    pts = (_CAM_POINTS * scale).astype(np.float32)
    return [np.stack([pts[a], pts[b]], axis=0) for a, b in _CAM_LINES]  # list[(2,3)]

def _to_rr_transform(mat4):
    mat4 = mat4.astype(np.float32)
    try:
        return rr.Transform3D(mat4x4=mat4)   # newer API
    except TypeError:
        return rr.Transform3D(mat4)          # older positional form

def _rr_init(app_id: str):
    # Be compatible across rerun-sdk versions
    try:
        rr.init(app_id)  # older positional
        return
    except TypeError:
        pass
    try:
        rr.init(application_id=app_id)  # some versions use 'application_id'
        return
    except TypeError:
        pass
    try:
        rr.init(app_id=app_id)  # newer keyword
    except TypeError:
        rr.init()  # last resort: no app id

def vigs_visualization_rerun(
    video,
    device="cuda:0",
    app_id="droid_viz",
    web_port=9876,                # HTTP port for the web viewer
    grpc_port=9877,               # gRPC port for the data stream
    enable_viewer=True,           # start live web viewer (--rerunvis)
    record_path=None,             # save .rrd file for replay (--rerun_record)
    warmup=8,
    filter_thresh=0.5,
    scale=1.0,
):
    torch.cuda.set_device(device)
    _rr_init(app_id)

    # Start gRPC data server + web viewer (no GUI/X11 required)
    started_server = False
    if enable_viewer:
        try:
            from urllib.parse import quote
            grpc_uri = rr.serve_grpc(grpc_port=grpc_port)
            rr.serve_web_viewer(web_port=web_port, open_browser=False, connect_to=grpc_uri)
            started_server = True
            web_url = f"http://127.0.0.1:{web_port}?url={quote(grpc_uri, safe='')}"
            print(f"[Rerun] Open in browser: {web_url}")
        except Exception as e:
            print(f"[Rerun] Could not start web server: {e}")

    # Record to file for later playback.
    # rr.save() overwrites the gRPC sink, so only use it when the live viewer is not active.
    if record_path and not started_server:
        try:
            rr.save(record_path)
            print(f"[Rerun] Recording to {record_path}")
        except Exception as e:
            print(f"[Rerun] Recording not started: {e}")

    # Global scene axes (best-effort across versions)
    try:
        rr.log("world", rr.ViewCoordinates.RDF, timeless=True)  # OpenCV-ish axes
    except Exception:
        pass

    filter_thresh = float(filter_thresh)
    cam_segs_local = _camera_lines(scale=0.1)  # larger for current cam

    try:
        while True:
            # Check which frames need updating
            with video.get_lock():
                dirty_index, = torch.where(video.dirty.clone())

            if len(dirty_index) == 0:
                time.sleep(0.01)
                continue

            # Mark processed
            video.dirty[dirty_index] = False

            high_res = False
            # Fetch tensors
            poses = torch.index_select(video.poses, 0, dirty_index)
            Ps = SE3(poses).inv().matrix().cpu().numpy()  # [N, 4, 4] world_T_cam
            images = torch.index_select(video.images, 0, dirty_index.cpu())
            intrinsics = video.intrinsics[0].clone()
            if high_res:
                disps = torch.index_select(video.disps_up, 0, dirty_index.cpu())
                images = images.cpu()[..., :, :].permute(0, 2, 3, 1) / 255.0  # [N, H, W, 3]
                intrinsics*=8
            else:
                disps = torch.index_select(video.disps, 0, dirty_index)
                images = images.cpu()[..., 3::8, 3::8].permute(0, 2, 3, 1) / 255.0  # [N, H, W, 3]
            # Consistency filter
            thresh = filter_thresh * torch.ones_like(disps.mean(dim=[1, 2]))
            if high_res:
                count = vigs_backends.depth_filter(video.poses, video.disps_up, intrinsics, dirty_index, thresh)
            else:
                count = vigs_backends.depth_filter(video.poses, video.disps, intrinsics, dirty_index, thresh)
            count = count.cpu()
            disps = disps.cpu()
            
            masks = ((count >= 2) & (disps > .5 * disps.mean(dim=[1, 2], keepdim=True)))
            # masks = ((count >= 2) & (disps > .1 * disps.mean(dim=[1, 2], keepdim=True))) # or use .1
            # masks = torch.ones_like(masks) # no filter
            
            points = vigs_backends.iproj(SE3(poses).inv().data, disps.cuda(), intrinsics).cpu()  # [N,H,W,3]

            
            # Log per keyframe
            for i in range(len(dirty_index)):
                kf_idx = int(dirty_index[i].item())
                world_from_cam = Ps[i]
                cam_path = f"world/cameras/{kf_idx:06d}"

                # Pose
                try:
                    rr.log(cam_path, _to_rr_transform(world_from_cam))
                except Exception:
                    pass

                # Frustum lines
                try:
                    segs_world = []
                    for seg in cam_segs_local:
                        seg_h = np.concatenate([seg, np.ones((2, 1), dtype=np.float32)], axis=1)  # (2,4)
                        seg_w = (world_from_cam @ seg_h.T).T[:, :3]
                        segs_world.append(seg_w)
                    rr.log(f"{cam_path}/frustum", rr.LineStrips3D(np.stack(segs_world, axis=0)))  # (S,2,3)
                except Exception:
                    pass

                # Colored point cloud
                mask = masks[i].reshape(-1).numpy()
                pts = points[i].reshape(-1, 3)[mask].numpy()
                clr = images[i].reshape(-1, 3)[mask].numpy()
                # still visualize high_res, but downsample, otherwise it's too dense, rerun very slow
                # even down sample it is different than the low_res one, should have more details
                if high_res:
                    downsample_factor = 512
                    pts=pts[::downsample_factor, :]
                    clr=clr[::downsample_factor, :]
                if pts.size > 0:
                    try:
                        rr.log(f"world/points/{kf_idx:06d}", rr.Points3D(pts, colors=(clr * 255).astype(np.uint8)))
                    except Exception:
                        rr.log(f"world/points/{kf_idx:06d}", rr.Points3D(pts))

    except KeyboardInterrupt:
        pass
