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

def interp2_normalized(f00, f10, f01, f11, x, y):
    a00 = f00
    a10 = f10 - f00
    a01 = f01 - f00
    a11 = f11 - f10 - f01 + f00
    return a00 + a10*x + a01*y + a11*x*y

def interp2(x0, y0, x1, y1, f00, f10, f01, f11, x, y):
    x_n = (x - x0)/(x1 - x0)
    y_n = (y - y0)/(y1 - y0)
    print(x_n, y_n)# DEBUG
    return interp2_normalized(f00, f10, f01, f11, x_n, y_n)

def update_pixel(surf, i, j, cam, img, interp_type='nearest'):
    ''' Updates a pixel in a surface by projecting a ray from the camera to the sphere centered on the camera.
        interp_type: 'linear' or 'nearest'
    '''
    # 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])
    if interp_type == 'linear':
        i0 = int(np.floor(i_img)) % img.shape[0]
        j0 = int(np.floor(j_img)) % img.shape[1]
        i1 = int(np.ceil(i_img)) % img.shape[0]
        j1 = int(np.ceil(j_img)) % img.shape[1]
        col = interp2(i0, j0, i1, j1, img[i0,j0,:], img[i1,j0,:], img[i0,j1,:], img[i1,j1,:], i_img, j_img)
        # print(i0, j0, i1, j1, i_img, j_img, img[i0,j0,:], img[i1,j0,:], img[i0,j1,:], img[i1,j1,:], col)# debug
        surf[i,j,:] = [int(col[0]), int(col[1]), int(col[2])]
    else:
        surf[i,j,:]  = img[int(i_img), int(j_img), :]
    return surf

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

    az = -10*deg
    cam = Camera(FOV=10*deg).set_target([np.cos(az),np.sin(az),0.05])
    # 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, interp_type)
    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()