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

def update_pixel(surf, i, j, cam):
    ''' Updates a pixel in a surface by projecting a ray from the sphere centered on the camera to the camera's sensor and checking whether it is in the field of view or not.
        The surface must be 2x1 aspect ratio (width=2, height=1).
    '''
    # Get the ray direction in inertial frame using spherical coordinates
    ray_dir_sph = equirectangular2sph(i, j, surf.shape[1], surf.shape[0])
    ray_dir_cart = sph2cart(ray_dir_sph)
    # surf[i,j] += 1000*max(0., np.dot(ray_dir_cart, cam.fwd))# dot product with forward vector
    if cam.is_point_visible_inertial_frame(ray_dir_cart):
        surf[i,j] += 1
    return surf

def add_camera_coverage(surf, cam):
    ''' Adds the camera coverage to a surface. '''
    for i in range(surf.shape[0]):
        for j in range(surf.shape[1]):
            surf = update_pixel(surf, i, j, cam)
    return surf

if __name__ == '__main__':
    deg = np.pi/180.
    width = 300
    height = int(width/2)
    surf = np.zeros((height, width), dtype=int)
    plt.figure()
    
    if 1:# arbitrary test
        surf = add_camera_coverage(surf, Camera(FOV=90*deg).set_target([1.,1.,1.]))
        surf = add_camera_coverage(surf, Camera(FOV=50*deg).set_target([-1.,0.5,-.5]))
        surf = add_camera_coverage(surf, Camera(FOV=20*deg).set_target([.2,-0.3,1.]))
        surf = add_camera_coverage(surf, Camera(FOV=10*deg).set_target([0.,1.,-.5]))
        plt.imshow(surf)
    
    if 0:# Regularly spaced cameras in azimuth
        surf = add_camera_coverage(surf, Camera(FOV=96*deg).set_target([1.,0.,0.]))
        surf = add_camera_coverage(surf, Camera(FOV=96*deg).set_target([0.,1.,0.]))
        surf = add_camera_coverage(surf, Camera(FOV=96*deg).set_target([-1.,0.,0.]))
        surf = add_camera_coverage(surf, Camera(FOV=96*deg).set_target([0.,-1.,0.]))
        plt.imshow(surf)
    
    if 0:# Camera 3D display DEBUG
        FOV = 30*deg
        cam1 = Camera(FOV=FOV).set_target([1.,0.,0.])
        cam2 = Camera(FOV=FOV).set_target([0.,1.,0.])
        cam3 = Camera(FOV=FOV).set_target([-1.,0.,0.])
        cam4 = Camera(FOV=FOV).set_target([0.,-1.,0.])
        cam5 = Camera(FOV=FOV).set_target([1.,-1.,1.])
        cams = [cam1, cam2, cam3, cam4, cam5]

        fig = plt.figure()
        ax = fig.add_subplot(111, projection='3d')
        ax.scatter(0., 0., 0., s=10, c='k', marker='o')
        for cam in cams:
            ax.scatter(cam.R[0,0], cam.R[1,0], cam.R[2,0], s=10, c='r', marker='o')
            ax.scatter(cam.R[0,1], cam.R[1,1], cam.R[2,1], s=10, c='g', marker='o')
            ax.scatter(cam.R[0,2], cam.R[1,2], cam.R[2,2], s=10, c='b', marker='o')
            NptsGrid = 11
            for i in np.linspace(-1., 1., NptsGrid):
                for j in np.linspace(-1., 1., NptsGrid):
                    p = cam.compute_ray_dir_inertial_frame(i, j)
                    ax.scatter(p[0], p[1], p[2], s=1, color=[.5, (i+1.)/2., (j+1.)/2.], marker='.')

    if 0:# azimuth
        surf = add_camera_coverage(surf, Camera(FOV=60*deg).set_target([1.,0.,0.]))
        surf = add_camera_coverage(surf, Camera(FOV=60*deg).set_target([0.,1.,0.]))
        plt.imshow(surf)
    
    if 0:# elevation
        surf = add_camera_coverage(surf, Camera(FOV=60*deg).set_target([1.,0.,0.]))
        surf = add_camera_coverage(surf, Camera(FOV=60*deg).set_target([1.,0.,1.]))
        surf = add_camera_coverage(surf, Camera(FOV=60*deg).set_target([-1.,0.,-1.]))
        surf = add_camera_coverage(surf, Camera(FOV=60*deg).set_target([0.,1.,-1000.]))
        plt.imshow(surf)

    if 0: # show coordinates in the image for debug
        surf_coords = np.zeros((height, width, 3))
        for i in range(surf_coords.shape[0]):
            for j in range(surf_coords.shape[1]):
                surf_coords[i,j,:] = equirectangular2sph(i, j, surf_coords.shape[1], surf_coords.shape[0])
        print(surf_coords)
        plt.figure()
        plt.imshow(surf_coords/np.pi)
    
    if 0: # plot points from spherical coordinates for debug
        width2 = 21
        height2 = 11
        fig = plt.figure()
        ax = fig.add_subplot(111, projection='3d')
        points = np.zeros((width2*height2, 3))
        for i in range(width2):
            for j in range(height2):
                p = sph2cart(np.array([1.0, np.pi*(i/(width2-1)), 2.0*np.pi*(j/(height2-1) - 0.5)]))
                points[i*height2+j,:] = p
        ax.scatter(points[:,0], points[:,1], points[:,2], s=10, c='k', marker='o')
    
    plt.show()