import numpy as np

class Camera:
    def __init__(self, fwd=np.array([1.,0.,0.]), up=np.array([0.,0.,1.]), FOV=90.*np.pi/180.):
        self.fwd = fwd / np.linalg.norm(fwd)
        self.up  = up  / np.linalg.norm(up)
        self.FOV = FOV
        self.R   = np.eye(3)

        self.build_rotation_matrix()
    
    def build_rotation_matrix(self):
        ''' Builds the rotation matrix from the fwd and up vectors. '''
        left   = np.cross(self.up, self.fwd)
        left   = left/np.linalg.norm(left)
        cam_up = np.cross(self.fwd, left)
        self.R[:,0] = -left             # x-axis is to the right
        self.R[:,1] = cam_up            # y-axis is up
        self.R[:,2] = -self.fwd         # z-axis is backwards
        return self
    
    def set_target(self, target):
        ''' Orients the camera so that the forward vector points at the target. The position of the camera is (0,0,0). '''
        self.fwd = target/np.linalg.norm(target)
        self.build_rotation_matrix()
        return self

    def compute_ray_dir_camera_frame(self, x, y):
        ''' Computes the ray direction in camera frame for a point (x,y) on the surface of the camera's sensor.
            x and y range from -1 to 1.
        '''
        # Compute the ray direction in the camera's frame
        ray_dir = np.array([x, y, -1.0/np.tan(self.FOV/2.0)])
        ray_dir = ray_dir/np.linalg.norm(ray_dir)
        return ray_dir

    def compute_ray_dir_inertial_frame(self, x, y):
        ''' Computes the ray direction for a point (x,y) on the surface of the camera's sensor.
            x and y range from -1 to 1.
        '''
        return self.R@self.compute_ray_dir_camera_frame(x, y)

    def project_point_camera_frame_to_sensor(self, p_camera):
        ''' Projects a point in camera frame to the sensor. '''
        # Project the point to the sensor
        tanFOV2 = 1./np.tan(self.FOV/2.0)
        p_sensor = p_camera
        if p_camera[2] != 0.0:
           p_sensor *= tanFOV2/np.abs(p_camera[2])
        return p_sensor

    def project_inertial_point_to_sensor(self, p_inertial):
        ''' Projects a point in inertial frame (cartesian coordinates) to the sensor. '''
        return self.project_point_camera_frame_to_sensor(self.R.T@p_inertial)

    def is_point_visible_inertial_frame(self, p_inertial):
        ''' Returns true if the point (cartesian coordinates in inertial frame) is visible to the camera. '''
        ray_sensor = self.project_inertial_point_to_sensor(p_inertial)
        return (ray_sensor[0] >= -1.0) and (ray_sensor[0] <= 1.0) and (ray_sensor[1] >= -1.0) and (ray_sensor[1] <= 1.0) and (ray_sensor[2] < 0.0)

def sph2cart(p):
    ''' Converts a spherical vector [r, theta, phi] in ISO convention (theta=0 is Z+) to a cartesian vector. '''
    r = p[0]
    theta = p[1]
    phi = p[2]
    return np.array([r*np.sin(theta)*np.cos(phi), r*np.sin(theta)*np.sin(phi), r*np.cos(theta)])

def cart2sph(p):
    ''' Converts a cartesian vector to a spherical vector [r, theta, phi] in ISO convention (theta=0 is Z+). '''
    r = np.sqrt(p[0]**2 + p[1]**2 + p[2]**2)
    theta = np.arctan2(np.sqrt(p[0]**2 + p[1]**2), p[2])
    phi = np.arctan2(p[1], p[0])
    return np.array([r, theta, phi])

def pixelToNormalizedCoordinates(i, j, width, height):
    ''' Converts pixel coordinates to normalized coordinates. i is the row (y), j is the column (x). '''
    return np.array([2.0*j/(width - 1) - 1.0, -2.0*i/(height - 1) + 1.0])

def sph2equirectangular(p_sph, width, height):
    ''' Converts spherical coordinates to pixel coordinates of an equirectangular image. '''
    theta = p_sph[1]
    phi = p_sph[2]
    i = np.mod(((height - 1)*theta)/np.pi, height)
    j = np.mod((width*(np.pi + phi))/(2.0*np.pi), width)
    return [i,j]

def sph2equirectangular_int(p_sph, width, height):
    ''' Converts spherical coordinates to pixel coordinates of an equirectangular image. '''
    i, j = sph2equirectangular(p_sph, width, height)
    return [int(i), int(j)]

def equirectangular2sph(i, j, width, height):
    return np.array([1.0, np.pi*(i/(height-1)), 2.0*np.pi*(j/width - 0.5)])

if __name__ == '__main__':
    import matplotlib.pyplot as plt

    deg = np.pi/180.
    cam = Camera()
    print(cam.R)

    p = np.array([1., -2., 3.])
    s = cart2sph(p)
    print(p, s, sph2cart(s))

    # ray direction from camera
    print('Rays in camera frame')
    cam.FOV = 20.*deg
    print('initial R =\n', cam.R)
    cam.set_target([1.,.5,.3])
    print('final R   =\n', cam.R)
    print(cam.compute_ray_dir_camera_frame(0., 0.))
    print(cam.compute_ray_dir_camera_frame(1., 0.))
    print(cam.compute_ray_dir_camera_frame(1., 1.))
    print(cam.compute_ray_dir_camera_frame(1., -1.))
    print(cam.compute_ray_dir_camera_frame(-1., -1.))
    print('Rays in inertial frame')
    print(cam.compute_ray_dir_inertial_frame(0., 0.))
    print(cam.compute_ray_dir_inertial_frame(1., 0.))
    print(cam.compute_ray_dir_inertial_frame(1., 1.))
    print(cam.compute_ray_dir_inertial_frame(1., -1.))
    print(cam.compute_ray_dir_inertial_frame(-1., -1.))

    # projection of point on sensor
    print('Project (0,0) from camera frame to sensor frame')
    print(cam.project_point_camera_frame_to_sensor(cam.compute_ray_dir_camera_frame(0., 0.)))
    print('Project (1,1) from camera frame to sensor frame')
    print(cam.project_point_camera_frame_to_sensor(cam.compute_ray_dir_camera_frame(1., 1.)))
    print('Project (-1,0.5) from camera frame to sensor frame')
    print(cam.project_point_camera_frame_to_sensor(cam.compute_ray_dir_camera_frame(-1., .5)))

    print('Project (0,0) from inertial frame to sensor frame')
    print(cam.project_inertial_point_to_sensor(cam.compute_ray_dir_inertial_frame(0., 0.)))
    print('Project (1,1) from inertial frame to sensor frame')
    print(cam.project_inertial_point_to_sensor(cam.compute_ray_dir_inertial_frame(1., 1.)))
    print('Project (-1,0.5) from inertial frame to sensor frame')
    print(cam.project_inertial_point_to_sensor(cam.compute_ray_dir_inertial_frame(-1., .5)))

    # projection of many points on sensor
    if 1:
        Npts = 11
        xy = np.linspace(-1.,1.,Npts)
        is_point_in_FOV = np.zeros((Npts,Npts), dtype=bool)
        cam.FOV = 60.*deg
        cam.set_target([1.,0.,0.])
        for i in range(Npts):
            for j in range(Npts):
                x, y = xy[i], xy[j]
                p_inertial = [1., x, y]
                p_sensor   = cam.project_inertial_point_to_sensor(p_inertial)
                is_point_in_FOV[i,j] = (p_sensor[0] >= -1.) and (p_sensor[0] <= 1.) and (p_sensor[1] >= -1.) and (p_sensor[1] <= 1.) and (p_sensor[2] < 0.0)
                print(p_inertial, p_sensor, is_point_in_FOV[i,j], cam.is_point_visible_inertial_frame(p_inertial))

    if 1:
        width, height = 100, 50
        normalized_coordinates = np.zeros((height, width, 3))
        for i in range(height):
            for j in range(width):
                p_n = pixelToNormalizedCoordinates(i, j, width, height)
                normalized_coordinates[i,j,:] = [0.5, p_n[0], p_n[1]]
        plt.imshow(normalized_coordinates)
        plt.show()

    if 0: # 3D plot to check validity of transform
        tgt = [1.,-1.,.5]
        cam = Camera(FOV=30*deg).set_target(tgt)

        fig = plt.figure()
        ax = fig.add_subplot(111, projection='3d')
        ax.scatter(0., 0., 0., s=10, c='k', marker='o')
        ax.scatter(tgt[0], tgt[1], tgt[2], s=10, c='r', marker='x')
        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')
        for i in np.linspace(-1., 1., 11):
            for j in np.linspace(-1., 1., 11):
                p = cam.compute_ray_dir_inertial_frame(i, j)
                ax.scatter(p[0], p[1], p[2], s=1, c='k', marker='.')
        plt.show()