import numpy as np
import cv2
from time import clock

# generate some 3d points arranged into two layers
x, y = np.mgrid[-5:5, -5:5]
z = np.ones_like(x)
p3d_1 = np.float32(np.dstack([x, y, z*40]).reshape(-1, 3))
p3d_2 = np.float32(np.dstack([x, y, z*50]).reshape(-1, 3))
p3d = np.vstack([p3d_1, p3d_2])

# setup some camera matrix 
w, h = 640, 480
f = 1280.0
K = np.diag([f, f, 1.0])
K[:2, 2] = 0.5*w, 0.5*h


def render(p, color):
    img = np.zeros((h, w, 3), np.uint8)
    n = len(p)/2
    for (x1, y1), (x2, y2) in np.int32(zip(p[:n], p[n:])):
        cv2.line(img, (x1, y1), (x2, y2), color, 2, cv2.CV_AA)
    return img

# render first (stationary) view
z3 = np.zeros(3)
p1 = cv2.projectPoints(p3d, z3, z3, K, np.zeros(4))[0].reshape(-1, 2)
cv2.imshow("camera1", render(p1, (0, 0, 255)))

while True:
    # render second view
    t = clock()*0.5
    tvec = np.float64([np.cos(t), np.sin(t), 0.0])*2
    p2 = cv2.projectPoints(p3d, z3, tvec, K, np.zeros(4))[0].reshape(-1, 2)
    img2 = render(p2, (128, 128, 0))

    # try to rectify two views
    F, _ = cv2.findFundamentalMat(p1, p2, cv2.FM_8POINT) # synthetic test -> no outliers -> not using RANSAC
    ret, H1, H2 = cv2.stereoRectifyUncalibrated(p1.reshape(1, -1, 2), p2.reshape(1, -1, 2), F, (w, h))

    # render rectified image
    p1t = cv2.perspectiveTransform(p1.reshape(1, -1, 2), H1).reshape(-1, 2)
    p2t = cv2.perspectiveTransform(p2.reshape(1, -1, 2), H2).reshape(-1, 2)
    img1t = render(p1t, (0, 0, 255))
    img2t = render(p2t, (128, 128, 0))
    rectified = cv2.addWeighted(img1t, 1.0, img2t, 1.0, 0.0)

    cv2.imshow("rectified", rectified)
    cv2.imshow("camera2", img2)
    ch = cv2.waitKey(10)
    if ch == 27:
        break