"""Shared pieces for the RL policy on the real SO101: robot/camera setup, blue-candy detection,
and the phone-pixel -> table homography."""

import json
import os
import time

import cv2
import numpy as np

HERE = os.path.dirname(os.path.abspath(__file__))
CALIB = os.path.join(HERE, "table_calib.json")
FOLLOWER_PORT = "/dev/serial/by-id/usb-1a86_USB_Single_Serial_5AE6057813-if00"
FOLLOWER_ID = "my_follower_arm"
PHONE_URL = os.environ.get("PHONE_URL", "http://127.0.0.1:4747/video")  # DroidCam on localhost; PHONE_URL=... to override
IMG_W, IMG_H = 640, 360  # the phone view is used at this size (same as the half-res datasets)
JOINTS = ["shoulder_pan", "shoulder_lift", "elbow_flex", "wrist_flex", "wrist_roll", "gripper"]


def connect_robot():
    from lerobot.robots.so_follower.config_so_follower import SOFollowerRobotConfig
    from lerobot.robots.so_follower.so_follower import SOFollower

    robot = SOFollower(SOFollowerRobotConfig(port=FOLLOWER_PORT, id=FOLLOWER_ID))
    robot.connect(calibrate=False)
    return robot


def read_joints(robot):
    obs = robot.get_observation()
    return np.array([obs[f"{n}.pos"] for n in JOINTS], dtype=np.float64)


class Phone:
    """DroidCam stream, resized to 640x360 RGB.

    A background thread reads the stream continuously and keeps only the newest frame, so a photo
    is always current (OpenCV otherwise hands back frames buffered seconds earlier, which made the
    "before" and "after" photos identical). Reconnects if the HTTP stream drops."""

    def __init__(self, url=PHONE_URL):
        import threading

        self.url = url
        self.cap = None
        self.latest, self.stamp, self.count = None, 0.0, 0
        self.lock = threading.Lock()
        self.running = True
        self._open()
        self.thread = threading.Thread(target=self._reader, daemon=True)
        self.thread.start()
        t0 = time.time()
        while self.latest is None and time.time() - t0 < 10:
            time.sleep(0.05)
        if self.latest is None:
            raise RuntimeError(f"no frames from {self.url} within 10 s")

    def _open(self):
        if self.cap is not None:
            self.cap.release()
        self.cap = cv2.VideoCapture(self.url)
        if not self.cap.isOpened():
            raise RuntimeError(f"cannot open phone stream {self.url} (is DroidCam running? close any browser tab showing it)")

    def _reader(self):
        fails = 0
        while self.running:
            ok, img = self.cap.read()
            if ok and img is not None:
                fails = 0
                with self.lock:
                    self.latest, self.stamp = img, time.time()
                    self.count += 1
                continue
            fails += 1
            print(f"  (phone stream dropped, reconnecting: attempt {fails})")
            time.sleep(1.0)
            try:
                self._open()
            except RuntimeError as e:
                print("  ", e)

    def frame(self, max_age=0.5, timeout=10.0):
        """Newest frame (taken after this call started), 640x360 RGB."""
        t_call = time.time()
        while time.time() - t_call < timeout:
            with self.lock:
                img, stamp = self.latest, self.stamp
            if img is not None and stamp > t_call and time.time() - stamp < max_age:
                return cv2.cvtColor(cv2.resize(img, (IMG_W, IMG_H), interpolation=cv2.INTER_AREA), cv2.COLOR_BGR2RGB)
            time.sleep(0.02)
        raise RuntimeError("no fresh phone frame for %.0f s; check the DroidCam app and connection" % timeout)

    def close(self):
        self.running = False
        self.thread.join(timeout=2)
        self.cap.release()


def blue_mask(rgb):
    hsv = cv2.cvtColor(rgb, cv2.COLOR_RGB2HSV)
    m = cv2.inRange(hsv, (95, 90, 70), (125, 255, 255))  # OpenCV hue 0..180: blue ~ 100-125
    return cv2.morphologyEx(m, cv2.MORPH_OPEN, np.ones((3, 3), np.uint8))


def blue_report(rgb):
    """Why a photo has (no) blue blobs: blue pixel count and the largest blue component."""
    m = blue_mask(rgb)
    n, _, stats, cent = cv2.connectedComponentsWithStats(m)
    big = max(range(1, n), key=lambda i: stats[i, cv2.CC_STAT_AREA], default=None)
    hsv = cv2.cvtColor(rgb, cv2.COLOR_RGB2HSV)
    return {"blue_px": int((m > 0).sum()), "largest_area": int(stats[big, cv2.CC_STAT_AREA]) if big else 0,
            "largest_at": cent[big].round(0).tolist() if big else None, "mean_brightness": float(hsv[..., 2].mean())}


def blue_blobs(rgb, min_px=25, ignore=None):
    """Blue wrapper blobs: list of dicts with pixel centroid, area and the blob's long-axis angle.
    ``ignore``: a reference image; blue that was already there (dilated a little) is left out."""
    m = blue_mask(rgb)
    if ignore is not None:
        m = cv2.bitwise_and(m, cv2.bitwise_not(cv2.dilate(blue_mask(ignore), np.ones((9, 9), np.uint8))))
    n, lab, stats, cent = cv2.connectedComponentsWithStats(m)
    out = []
    for i in range(1, n):
        if stats[i, cv2.CC_STAT_AREA] < min_px:
            continue
        ys, xs = np.nonzero(lab == i)
        pts = np.c_[xs, ys].astype(np.float64)
        c = pts.mean(0)
        evals, evecs = np.linalg.eigh(np.cov((pts - c).T))
        major = evecs[:, np.argmax(evals)]
        out.append({"uv": c, "area": int(stats[i, cv2.CC_STAT_AREA]), "axis_uv": major})
    return out


def fit_homography(uv, xy):
    """Least-squares DLT homography (u, v) -> (x, y)."""
    A = []
    for (u, v), (x, y) in zip(uv, xy):
        A.append([-u, -v, -1, 0, 0, 0, x * u, x * v, x])
        A.append([0, 0, 0, -u, -v, -1, y * u, y * v, y])
    _, _, Vt = np.linalg.svd(np.array(A))
    return Vt[-1].reshape(3, 3) / Vt[-1, -1]


def apply_h(H, uv):
    uv = np.atleast_2d(uv)
    q = np.c_[uv, np.ones(len(uv))] @ H.T
    return q[:, :2] / q[:, 2:3]


def load_calib():
    if not os.path.exists(CALIB):
        raise SystemExit(f"No {CALIB}. Run calibrate_table.py first.")
    return np.array(json.load(open(CALIB))["H"])


def candy_on_table(H, blob):
    """Table x, y [m] and yaw [rad] of a blob (yaw from the blob's long axis mapped onto the table)."""
    c = blob["uv"]
    p0, p1 = apply_h(H, [c - 8 * blob["axis_uv"], c + 8 * blob["axis_uv"]])
    d = p1 - p0
    xy = apply_h(H, c)[0]
    return xy, float(np.arctan2(d[1], d[0]))