import numpy as np
from Camera import *
import matplotlib.pyplot as plt

def update_pixel(surf, i, j, cam, img):
    ''' Updates a pixel in a surface by projecting a ray from the camera to the sphere centered on the camera. '''
    # Get the ray direction in inertial frame by projecting the pixel to the sphere
    p_sensor     = pixelToNormalizedCoordinates(i, j, surf.shape[1], surf.shape[0])
    p_sensor[0] *= -1
    p_sensor[1] *= surf.shape[0]/surf.shape[1]
    ray_dir      = cam.compute_ray_dir_inertial_frame(p_sensor[0], p_sensor[1])
    ray_dir_sph  = cart2sph(ray_dir)
    i_img, j_img = sph2equirectangular(ray_dir_sph, img.shape[1], img.shape[0])
    surf[i,j,:]  = img[i_img, j_img, :]
    return surf

if __name__ == '__main__':
    deg = np.pi/180.
    width = 60
    height = int(width/2)
    surf = np.zeros((height, width, 3), dtype=int)
    
    # load 360 image
    img = plt.imread('venise.jpg')
    plt.figure()
    plt.imshow(img)

    cam = Camera(FOV=90*deg).set_target([0.,1.,.8])
    # cam = Camera(FOV=90*deg,up=np.array([1.,0.,1.])).set_target([0.,0.,-1.])# straight down
    
    # update surface
    for i in range(height):
        for j in range(width):
            surf = update_pixel(surf, i, j, cam, img)
    plt.figure()
    plt.imshow(surf)

    # plot points from spherical coordinates for debug
    if 0:
        width2 = 21
        height2 = 11
        fig = plt.figure()
        ax = fig.add_subplot(111, projection='3d')
        points = np.zeros((width2*height2, 3))
        colors = np.zeros((width2*height2, 3), dtype=int)
        for i in range(height2):
            for j in range(width2):
                p = sph2cart(np.array([1.0, np.pi*(i/(height2-1)), 2*np.pi*(j/(width2-1))]))
                points[i*height2+j,:] = p
                colors[i*height2+j,:] = img[int((i/(height2-1))*(img.shape[0]-1)), int((j/(width2-1))*(img.shape[1]-1)), :]# get color from image
        ax.scatter(points[:,0], points[:,1], points[:,2], s=10, c=colors/255, marker='o')
    
    plt.show()