import os
import sys
sys.path.append(os.path.join(os.path.dirname(__file__), 'vigs'))
import time
import torch
import cv2
import argparse
import numpy as np
import lietorch
import resource
import socket
import struct
import threading
import queue as stdlib_queue
from collections import deque

rlimit = resource.getrlimit(resource.RLIMIT_NOFILE)
resource.setrlimit(resource.RLIMIT_NOFILE, (100000, rlimit[1]))

from vigs import VIGS
from line_profiler import profile

HOST = "0.0.0.0"
PORT = 9999
RES = 341 * 640

# Maximum frames to buffer. When full, oldest is dropped so SLAM always
# processes near-real-time frames and latency stays bounded.
IMAGE_QUEUE_DEPTH = 3


# ---------------------------------------------------------------------------
# Thread-safe IMU buffer
# ---------------------------------------------------------------------------

class IMUStreamBuffer:
    """
    Accumulates synced IMU rows as they arrive from the socket receiver.
    Rows are stored pre-converted to VIGS internal format:
        [t_s, gx, gy, gz, ax_ms2, ay_ms2, az_ms2]
    Thread-safe: append from receiver thread, snapshot from SLAM thread.
    """
    def __init__(self):
        self._rows = []
        self._lock = threading.Lock()
        self._last_t = -1.0          # last received IMU timestamp (seconds)
        self._cond = threading.Condition(self._lock)

    def append(self, row):
        with self._cond:
            self._rows.append(row)
            self._last_t = row[0]
            self._cond.notify_all()

    def wait_until(self, t_s, timeout=2.0):
        """Block until IMU buffer has at least one sample at or past t_s."""
        with self._cond:
            return self._cond.wait_for(lambda: self._last_t >= t_s, timeout=timeout)

    def snapshot(self):
        """Return a numpy (N, 7) copy of all rows so far."""
        with self._lock:
            if len(self._rows) == 0:
                return np.zeros((0, 7), dtype=np.float64)
            return np.array(self._rows, dtype=np.float64)

    def __len__(self):
        with self._lock:
            return len(self._rows)


# ---------------------------------------------------------------------------
# Socket helpers (identical to stream_data.py)
# ---------------------------------------------------------------------------

def recv_exact(conn, n):
    data = b""
    while len(data) < n:
        chunk = conn.recv(n - len(data))
        if not chunk:
            raise ConnectionError("socket closed")
        data += chunk
    return data


def interp_sample(buf, t_ns):
    if len(buf) < 2:
        return None
    for i in range(len(buf) - 1):
        t0, x0, y0, z0 = buf[i]
        t1, x1, y1, z1 = buf[i + 1]
        if t0 <= t_ns <= t1:
            if t1 == t0:
                return (x0, y0, z0)
            alpha = (t_ns - t0) / float(t1 - t0)
            return (
                (1.0 - alpha) * x0 + alpha * x1,
                (1.0 - alpha) * y0 + alpha * y1,
                (1.0 - alpha) * z0 + alpha * z1,
            )
    return None


def prune_gyro_buffer(gyro_buf, oldest_needed_t):
    while len(gyro_buf) >= 2 and gyro_buf[1][0] < oldest_needed_t:
        gyro_buf.popleft()


# ---------------------------------------------------------------------------
# Socket receiver (runs in a background daemon thread)
# ---------------------------------------------------------------------------

class SocketReceiver(threading.Thread):
    """
    Receives iPhone packets in a background thread and feeds:
      - images into `image_queue`  (item: (t_ns, jpeg_bytes, w, h))
      - synced IMU into `imu_buffer` (pre-converted to VIGS format)
      - camera intrinsics into `self.intrinsics` + sets `intrinsics_event`

    IMU sync: gyro is interpolated to accel timestamps (same as stream_data.py).
    Unit/axis conversion to match existing iphone handling in demo.py:
        raw accel (assumed in g) → × 9.81 → m/s²
        column order: [t_s, gx, gy, gz, ax_ms2, ay_ms2, az_ms2]
    """
    def __init__(self, conn, image_queue, imu_buffer, intrinsics_event, imu_time_offset=0.0):
        super().__init__(daemon=True)
        self.conn = conn
        self.image_queue = image_queue
        self.imu_buffer = imu_buffer
        self.intrinsics_event = intrinsics_event
        self.imu_time_offset = imu_time_offset

        self.intrinsics = None   # (fx, fy, cx, cy, w, h)
        self.running = True
        self._frame_count = 0    # for stride skipping
        self.stop_event = threading.Event()

    def _flush_synced_imu(self, accel_buf, gyro_buf):
        if len(accel_buf) == 0 or len(gyro_buf) < 2:
            return
        latest_gyro_t = gyro_buf[-1][0]
        while len(accel_buf) > 0:
            t_a, ax, ay, az = accel_buf[0]
            if t_a > latest_gyro_t:
                break
            g = interp_sample(gyro_buf, t_a)
            if g is None:
                if t_a < gyro_buf[0][0]:
                    accel_buf.popleft()
                    continue
                break
            # Convert to VIGS format: [t_s, gx, gy, gz, ax_ms2, ay_ms2, az_ms2]
            t_s = t_a / 1e9 + self.imu_time_offset
            self.imu_buffer.append([t_s, g[0], g[1], g[2], ax * -9.81, ay * -9.81, az * -9.81])
            accel_buf.popleft()

    def run(self):
        accel_buf = deque()
        gyro_buf = deque()
        try:
            while not self.stop_event.is_set():
                header = recv_exact(self.conn, 4)
                msg_len = struct.unpack(">I", header)[0]
                msg = recv_exact(self.conn, msg_len)
                msg_type = msg[0]
                payload = msg[1:]

                if msg_type == 1:
                    # Image — only queue after intrinsics are known to avoid
                    # building a backlog that floods SLAM on startup.
                    if len(payload) < 20:
                        continue
                    t_ns, w, h, jpeg_size = struct.unpack(">QIII", payload[:20])
                    jpeg_bytes = payload[20:20 + jpeg_size]
                    if len(jpeg_bytes) != jpeg_size:
                        continue
                    self._frame_count += 1
                    if self.intrinsics is not None and self._frame_count % 3 == 0:
                        # Drop-oldest: if queue is full, discard the stale frame
                        # so SLAM always sees near-real-time images.
                        if self.image_queue.full():
                            try:
                                self.image_queue.get_nowait()
                            except stdlib_queue.Empty:
                                pass
                        self.image_queue.put_nowait((t_ns, jpeg_bytes, w, h))

                elif msg_type == 2:
                    # IMU
                    if len(payload) != 8 + 1 + 8 * 3:
                        continue
                    t_ns = struct.unpack(">Q", payload[:8])[0]
                    sensor_type = payload[8]
                    x, y, z = struct.unpack(">ddd", payload[9:])
                    if sensor_type == 0:
                        accel_buf.append((t_ns, x, y, z))
                    elif sensor_type == 1:
                        gyro_buf.append((t_ns, x, y, z))
                    else:
                        continue
                    self._flush_synced_imu(accel_buf, gyro_buf)
                    # Prune gyro buffer
                    if len(accel_buf) > 0:
                        prune_gyro_buffer(gyro_buf, accel_buf[0][0])
                    elif len(gyro_buf) > 2:
                        while len(gyro_buf) > 2:
                            gyro_buf.popleft()

                elif msg_type == 3:
                    # Camera intrinsics
                    if len(payload) != 4 + 4 + 4 * 5:
                        continue
                    width, height = struct.unpack(">II", payload[:8])
                    fx, fy, cx, cy, lens_position = struct.unpack(">fffff", payload[8:])
                    if self.intrinsics is None:
                        self.intrinsics = (fx, fy, cx, cy, width, height)
                        print(f"[INFO] Received camera intrinsics: fx={fx:.2f} fy={fy:.2f} cx={cx:.2f} cy={cy:.2f} ({width}x{height})")
                        self.intrinsics_event.set()

        except ConnectionError:
            print("\n[INFO] iPhone disconnected.")
        except Exception as e:
            print(f"[ERROR] Receiver: {e}")
        finally:
            self.running = False
            # Unblock the main thread if it's waiting on the queue
            self.image_queue.put(None)


# ---------------------------------------------------------------------------
# Trajectory saving (keyframe-only, for streaming mode)
# ---------------------------------------------------------------------------

def save_trajectory_stream(vigs, tstamps_full, output, suffix='', final=False):
    """Save keyframe trajectory. tstamps_full is a list of per-frame timestamps."""
    t = vigs.video.counter.value
    if t == 0:
        return
    tstamps = vigs.video.tstamp[:t]
    poses_wc = lietorch.SE3(vigs.video.poses[:t]).inv().data
    if final:
        np.save(f"{output}/intrinsics.npy", vigs.video.intrinsics[0].cpu().numpy() * 8)
    kf_indices = tstamps.cpu().numpy().astype(int)
    ts_arr = np.array(tstamps_full)
    valid = kf_indices < len(ts_arr)
    kf_indices = kf_indices[valid]
    poses_wc_np = poses_wc.cpu().numpy()[valid]
    tstamps_kf = ts_arr[kf_indices][..., np.newaxis]
    ttraj_kf = np.concatenate([tstamps_kf, poses_wc_np], axis=1)
    os.makedirs(output, exist_ok=True)
    np.savetxt(f"{output}/traj_kf{suffix}.txt", ttraj_kf)
    print(f"[INFO] Saved keyframe trajectory ({len(ttraj_kf)} KFs) → {output}/traj_kf{suffix}.txt")


# ---------------------------------------------------------------------------
# Main
# ---------------------------------------------------------------------------

@profile
def main():
    parser = argparse.ArgumentParser(description="VIGS-SLAM streaming demo (iPhone via socket)")
    parser.add_argument("--config", type=str, required=True, help="path to configuration yaml")
    parser.add_argument("--output", default="outputs/stream", help="path to save outputs")
    parser.add_argument("--weights", default=os.path.join(os.path.dirname(__file__), "pretrained_models/droid.pth"))
    parser.add_argument("--buffer", type=int, default=1200)
    parser.add_argument("--host", type=str, default=HOST)
    parser.add_argument("--port", type=int, default=PORT)
    parser.add_argument("--IMU_poseinit_after", type=int, default=25,
                        help="enable IMU init after this keyframe index") #20
    parser.add_argument("--gsmapping", action="store_true")
    parser.add_argument("--gsvis", action="store_true")
    parser.add_argument("--droidvis", action="store_true")
    parser.add_argument("--rerunvis", action="store_true")
    parser.add_argument("--final_ba_inertial", action="store_true")
    args = parser.parse_args()

    if torch.cuda.is_available():
        print("GPU Available:", torch.cuda.get_device_name(0))

    os.makedirs(args.output, exist_ok=True)
    os.system(f"cp {args.config} {args.output}/config.yaml")

    # Fields expected by VIGS / parse_extra_params
    args.imagedir = "iphone_stream"   # triggers iphone axis handling path
    args.imufile = None
    args.undistort = False
    args.cropborder = 0
    args.stride = 1
    args.start = 0
    args.length = 100000

    # Empty IMU array — will be replaced with live snapshots before each frame
    # parse_extra_params in vigs.py runs in-place ops on this; empty array → no-op
    args.imus = np.zeros((0, 7), dtype=np.float64)

    torch.multiprocessing.set_start_method('spawn')

    # ------------------------------------------------------------------
    # Socket server
    # ------------------------------------------------------------------
    server = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    server.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
    server.bind((args.host, args.port))
    server.listen(1)
    print(f"[INFO] Listening on {args.host}:{args.port} — waiting for iPhone...")
    conn, addr = server.accept()
    print(f"[INFO] Connected from {addr}")

    image_queue = stdlib_queue.Queue(maxsize=IMAGE_QUEUE_DEPTH)
    imu_buffer = IMUStreamBuffer()
    intrinsics_event = threading.Event()

    receiver = SocketReceiver(conn, image_queue, imu_buffer, intrinsics_event)
    receiver.start()

    # ------------------------------------------------------------------
    # Wait for camera intrinsics before starting SLAM
    # ------------------------------------------------------------------
    print("[INFO] Waiting for camera intrinsics from iPhone...")
    intrinsics_event.wait()
    fx, fy, cx, cy, src_w, src_h = receiver.intrinsics

    # ------------------------------------------------------------------
    # Wait for a minimal IMU backlog before init (avoids preint with no data)
    # ------------------------------------------------------------------
    print("[INFO] Waiting for initial IMU samples...")
    while len(imu_buffer) < 20 and receiver.running:
        time.sleep(0.01)

    # ------------------------------------------------------------------
    # Compute scaled image size once (from reported camera dimensions)
    # ------------------------------------------------------------------
    h0, w0 = src_h, src_w
    h1 = int(h0 * np.sqrt(RES / (h0 * w0)))
    w1 = int(w0 * np.sqrt(RES / (h0 * w0)))
    h1 = h1 - h1 % 8
    w1 = w1 - w1 % 8
    scale_x = w1 / w0
    scale_y = h1 / h0
    scaled_intrinsics = torch.tensor([fx * scale_x, fy * scale_y, cx * scale_x, cy * scale_y])
    args.image_size = [h1, w1]
    print(f"[INFO] Image size: {w0}×{h0} → {w1}×{h1}, intrinsics: {scaled_intrinsics.tolist()}")

    # ------------------------------------------------------------------
    # SLAM loop
    # ------------------------------------------------------------------
    vigs = None
    t_global = 0             # sequential frame counter (used as index into tstamps_full)
    tstamps_full = []        # per-frame timestamps (seconds), parallel to t_global

    # debug_img_dir = os.path.join(args.output, "debug_images")
    # os.makedirs(debug_img_dir, exist_ok=True)
    # print(f"[INFO] Saving debug images to {debug_img_dir}")
    print("[INFO] Starting SLAM ...")

    while True:
        try:
            item = image_queue.get(timeout=5.0)
        except stdlib_queue.Empty:
            if not receiver.running:
                print("[INFO] Receiver stopped and queue empty — finishing.")
                break
            continue

        if item is None:
            # Receiver sent sentinel → connection closed
            print("\n[INFO] Received end-of-stream sentinel.")
            break

        t_ns, jpeg_bytes, img_w, img_h = item
        timestamp = t_ns / 1e9
        tstamps_full.append(timestamp)

        # Decode JPEG
        arr = np.frombuffer(jpeg_bytes, dtype=np.uint8)
        image = cv2.imdecode(arr, cv2.IMREAD_COLOR)
        if image is None:
            t_global += 1
            continue

        # Save original JPEG for debugging (BGR, full resolution)
        # cv2.imwrite(os.path.join(debug_img_dir, f"{t_global:06d}_{t_ns}.jpg"), image)

        try:
            cv2.namedWindow("SLAM input", cv2.WINDOW_NORMAL)
            # cv2.resizeWindow("SLAM input", 640*2, 480*2-200) # width, height in pixels
            cv2.resizeWindow("SLAM input", int(900), int((480*2-200)//2*1.3)) # width, height in pixels
            cv2.moveWindow("SLAM input", 0, 0) # x, y from top-left of screen 
            cv2.imshow("SLAM input", image)
            cv2.waitKey(1)
        except cv2.error:
            pass  # headless

        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)

        # Resize to SLAM resolution
        image = cv2.resize(image, (w1, h1))
        image_tensor = torch.as_tensor(image).permute(2, 0, 1)

        # Wait for IMU buffer to have data past this frame's timestamp so that
        # preintegration between the previous and current keyframe has coverage.
        if not imu_buffer.wait_until(timestamp, timeout=2.0):
            print(f"[WARN] IMU timeout waiting for t={timestamp:.3f}, proceeding anyway")

        # Snapshot current IMU data and inject into VIGS before tracking
        imu_snapshot = imu_buffer.snapshot()

        if vigs is None:
            # First frame: init VIGS.
            # parse_extra_params in vigs.py modifies args.imus in-place (e.g. divides
            # timestamps by 1e9 when imu_in_nanoseconds=True, even though our streaming
            # rows are already in seconds).  Pass a *copy* so the original snapshot is
            # not corrupted, then immediately replace video.imus with a fresh snapshot
            # that has correct units.
            args.imus = imu_snapshot.copy()
            vigs = VIGS(args)
            vigs.video.imus = imu_buffer.snapshot()   # correct seconds-based data
            # print(f"[INFO] VIGS initialized. IMU rows so far: {len(vigs.video.imus)}")
            print(f"[INFO] VIGS Instance Created.")
        else:
            # Inject fresh IMU snapshot before tracking this frame
            vigs.video.imus = imu_snapshot

        is_last = (not receiver.running) and image_queue.empty()
        vigs.track(t_global, timestamp, image_tensor[None], intrinsics=scaled_intrinsics[None].clone(), is_last=is_last)

        if args.gsmapping:
            n_gs = vigs.gs.gaussians._xyz.shape[0] if hasattr(vigs, 'gs') else 0
            print(f"\r[SLAM] frame={t_global} kf={vigs.video.counter.value} gs={n_gs}", end="", flush=True)
        else:
            print(f"\r[SLAM] frame={t_global} kf={vigs.video.counter.value}", end="", flush=True)

        t_global += 1

        if is_last:
            break

    print()  # newline after \r status
    receiver.stop_event.set()
    conn.close()
    server.close()

    if vigs is None:
        print("[WARN] VIGS was never initialized — no data received.")
        return

    # ------------------------------------------------------------------
    # Finalize
    # ------------------------------------------------------------------
    if hasattr(vigs, 'mp_backend'):
        vigs.video.pgobuf.stop()
        vigs.mp_backend.join(timeout=1.0)

    # Inject final full IMU snapshot for termination BA
    vigs.video.imus = imu_buffer.snapshot()

    traj_full_beforeBA = vigs.traj_filler(vigs.images)
    traj_before = traj_full_beforeBA.inv().data.cpu().numpy() # full trajectory before BA

    # Build tstamps_full array for save_trajectory_stream
    save_trajectory_stream(vigs, tstamps_full, args.output, suffix="_beforeBA", final=True)
    # save_trajectory_stream(vigs, tstamps_full, args.output, suffix="", final=True)

    vigs.gs.gaussians.save_ply(f"{args.output}/3dgs_before_final.ply")

    traj = vigs.terminate(inertial=args.final_ba_inertial & vigs.video.IMU_initialized) # final trajectory after BA
    vigs.gs.gaussians.save_ply(f"{args.output}/3dgs_final.ply")
    save_trajectory_stream(vigs, tstamps_full, args.output, suffix="_afterBA", final=False)

    print(f"[INFO] Done. Outputs in {args.output}")
    time.sleep(1000)


if __name__ == "__main__":
    main()
