import numpy as np
import matplotlib.patches as pa


class Sphere(object):

    def __init__(self, xy, z=0., radius=0.1, color=None):
        self.xy = np.asarray(xy)
        self.z = z
        self.radius = radius
        self.color = color

    def draw_top_view(self, ax):
        circle = pa.Circle(tuple(self.xy), radius=self.radius, facecolor=self.color)
        ax.add_artist(circle)

    def calculate_image_view_shape(self, camera):
        """
        :param camera:
        :return:
            pos: 2, | [u, v] position of element in image
            - not in image -> nans
            - may be outside of image bounds, use image size to filter these
            radius: | in pixels
        """
        xyz1 = np.hstack([self.xy, self.z, 1.])
        top_bot_xyz1 = np.array([xyz1, xyz1]).T
        top_bot_xyz1[2, :] += [self.radius, -self.radius]
        top_bot_xy1 = camera.project(top_bot_xyz1)
        if np.any(top_bot_xy1[-1] <= 0):
            return np.array([np.nan] * 2), np.nan
        top_bot_xy1 /= top_bot_xy1[-1]
        xy = top_bot_xy1[:-1].mean(axis=1)
        r = np.linalg.norm(top_bot_xy1[:-1, 0] - xy)/2
        return xy, r

    def calculate_image_view_center(self, camera):
        """
        :param camera:
        :return:
            pos: 2, | [u, v] position of element in image
            - not in image -> nans
            - may be outside of image bounds, use image size to filter these
        """
        xy = self.calculate_image_view_shape(camera)[0]
        return xy

    def draw_image_view(self, ax, camera, zorder=0):
        """
        Do not draw if behind camera.
        (Out of frame objects can be drawn, they will just be cut off)
        :param ax:
        :param camera:
        :param zorder: ordering for fake depth
        :return:
        """
        xy, r = self.calculate_image_view_shape(camera)
        if np.isnan(xy).any():
            return
        # reverse xy to place into plot coordinates, since cam x is vertical
        circle = pa.Circle(
            tuple(xy[::-1]), radius=r,
            facecolor=self.color, zorder=zorder
        )
        ax.add_artist(circle)