#!/usr/bin/env python3
"""Sparrow core pipeline (seed v1 — WORKS, but crude: yours to improve).

The seed is a functioning end-to-end baseline: it detects corners, tracks
them frame to frame, and keeps a naive sparse map. It is deliberately
unimpressive so the company has real work to do (see backlog).

Conventions the whole repo follows:
  - extract(frame)      -> (points Nx2 float32, scores)
  - step(prev, frame, prev_pts) -> (tracked Nx2, lost_fraction, pair idx)
  - study of the fallback modules: features.py / tracker.py / map.py
    may be created by the team and are imported if present (duck-typed
    to these same APIs — the rest of the repo won't know or care).
"""
import cv2
import numpy as np

try:
    from features import extract as feature_extract
except Exception:
    feature_extract = None

try:
    from tracker import step as tracker_step
except Exception:
    tracker_step = None

try:
    from map import SparseMap
except Exception:
    SparseMap = None

TIME_STEP = 1.0 / 30.0


def extract_builtin(frame):
    import cv2
    gray = frame if frame.ndim == 2 else cvt_gray(frame)
    pts = cv2.goodFeaturesToTrack(gray, 200, qualityLevel=0.01, minDistance=8)
    return (pts.reshape(-1, 2), np.ones(len(pts))) if pts is not None else (np.zeros((0, 2)), np.zeros(0))


def flow_builtin(gray_prev, gray_cur, pts):
    p1, st, err = cv2.calcOpticalFlowPyrLK(gray_prev, gray_cur, pts, None, winSize=(21, 21))
    st = st.ravel().astype(bool)
    missing = int((~st).sum())
    return p1[st].reshape(-1, 2), pts[st].reshape(-1, 2), missing / max(1, len(st))


def cvt_gray(f):
    import cv2
    return cv2.cvtColor(f, cv2.COLOR_BGR2GRAY)


def extract(frame):
    """Extract features (seed = plain goodFeaturesToTrack)."""
    return feature_extract(frame) if feature_extract else extract_builtin(frame)


def step(prev_gray, gray, prev_pts):
    """Track points into the frame; returns (cur_pts, prev_pts_stage, lost_frac)."""
    if tracker_step:
        return tracker_step(prev_gray, gray, prev_pts)
    return flow_builtin(prev_gray, gray, prev_pts)


class WeakMap:
    """Seed map: just the current points + snapshot(). Replace via map.py."""

    def __init__(self):
        self.points = {}          # idx -> [x, z, age]
        self.nid = 0

    def update(self, pts, prev_pts):
        n = len(pts)
        for i in range(n):
            self.nid += 1
            self.points[self.nid] = [float(pts[i][0]), float(prev_pts[i][1]), 0.0]

    def snapshot(self):
        return {
            "image_pts": [(v[0], v[1], v[2]) for v in list(self.points.values())],
            "count": len(self.points),
        }


if SparseMap:
    WeakMap = SparseMap


class Pipeline:
    def __init__(self):
        self.prev_gray = None
        self.prev_pts = None
        self.map = WeakMap()

    def on_frame(self, frame):
        import cv2
        gray = frame if frame.ndim == 2 else cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
        if self.prev_gray is None or self.prev_pts is None or len(self.prev_pts) < 4:
            pts, scores = extract(frame)
            self.prev_gray, self.prev_pts = gray, pts.astype(np.float32)
            return {"features": len(pts), "lost": 0.0, "new": len(pts)}
        cur, prev_matched, lost = step(self.prev_gray, gray, self.prev_pts)
        if len(cur) < 4:
            pts, _ = extract(frame)
            self.prev_gray, self.prev_pts = gray, pts.astype(np.float32)
            return {"features": len(pts), "lost": 1.0, "new": len(pts)}
        self.map.update(cur, prev_matched)
        # reseed occasionally so the stream never starves
        if len(cur) < 60:
            pts, _ = extract(frame)
            self.prev_pts = np.vstack([cur, pts.astype(np.float32)])
        else:
            self.prev_pts = cur
        self.prev_gray = gray
        return {"features": len(cur), "lost": float(lost), "new": len(cur) - len(prev_matched)}


if __name__ == "__main__":
    import cv2
    import scene
    p = Pipeline()
    t = 0.0
    for _ in range(10):
        f = scene.render(t)
        r = p.on_frame(f)
        print({k: round(v, 2) for k, v in r.items()})
        t += TIME_STEP