import cv2
import numpy as np
from PIL import Image, ImageDraw
import random

from .color_utils import map_colors
from .transform import Transform


class LidarCamera:
    world2cam = Transform.from_euler([-90, 0, -90], degrees=True)

    def __init__(self, camera_id):
        self.id = camera_id
        self.pose = Transform(self.world2cam)
        self.K = np.eye(3, dtype=np.float32)
        self.D = np.zeros(6, dtype=np.float32)
        self.model = 'brown_conrady'
        self.model = None
        self.scale_factor = 1
        self.skew = 0

    @property
    def position(self) -> np.ndarray:
        return self.pose.position

    @property
    def rotation(self) -> np.ndarray:
        return self.pose.rotation

    @property
    def world_transform(self) -> Transform:
        return self.pose @ self.world2cam.T

    @property
    def fx(self):
        return self.K[0, 0]

    @property
    def fy(self):
        return self.K[1, 1]

    @property
    def cx(self):
        return self.K[0, 2]

    @property
    def cy(self):
        return self.K[1, 2]

    @property
    def intrinsic_matrix(self):
        return self.K

    @property
    def extrinsic_matrix(self):
        return self.pose.inverse[:3, :4]

    @property
    def projection_matrix(self):
        return self.K @ self.extrinsic_matrix

    @position.setter
    def position(self, position: np.ndarray):
        self.pose.position = position

    @rotation.setter
    def rotation(self, rotation):
        self.pose.rotation = Transform(rotation).rotation

    @world_transform.setter
    def world_transform(self, transform: Transform):
        self.pose = transform @ self.world2cam

    @extrinsic_matrix.setter
    def extrinsic_matrix(self, matrix):
        self.pose = Transform(matrix).inverse

    @projection_matrix.setter
    def projection_matrix(self, P):
        assert P.shape == (3, 4), 'Projection matrix should be 3x4'
        K, R, t, _, _, _, _ = cv2.decomposeProjectionMatrix(P)
        self.pose = Transform.from_Rt(R.T, t[:3, 0] / t[3, 0])
        self.K = K

    def calibrate(self, position=None, rotation=None, pose=None, extrinsic_matrix=None, projection_matrix=None, K=None, D=None, model=None, scale_factor=None, skew=None, **kwargs):
        """
        Helper for camera calibration
        :param position: Camera position [x, y, z]
        :param rotation: Camera rotation/heading
        :param pose: Camera pose (position + rotation)
        :param extrinsic_matrix: extrinsic 4x4 matrix (world to camera transform)
        :param projection_matrix: 3x4 projection matrix
        :param K: intrinsic 3x3 matrix
        :param D: distortion values
        :param model: camera model
        :param scale_factor: image scale_factor
        :param skew: camera skew coefficient
        """
        if position is not None:
            self.position = position
        if rotation is not None:
            self.rotation = rotation
        if pose is not None:
            self.pose = Transform(pose)
        if extrinsic_matrix is not None:
            self.extrinsic_matrix = extrinsic_matrix
        if projection_matrix is not None:
            self.projection_matrix = projection_matrix
        if K is not None:
            self.K = np.array(K[:3, :3])
        if D is not None:
            self.D = np.pad(D, (0, 6 - len(D)), 'constant')
        if model is not None:
            self.model = model
        if scale_factor is not None:
            self.scale_factor = scale_factor
        if skew is not None:
            self.skew = skew
        if 'world_transform' in kwargs:
            self.world_transform = kwargs['world_transform']
        if 'fx' in kwargs:
            self.K[0, 0] = kwargs['fx']
        if 'fy' in kwargs:
            self.K[1, 1] = kwargs['fy']
        if 'cx' in kwargs:
            self.K[0, 2] = kwargs['cx']
        if 'cy' in kwargs:
            self.K[1, 2] = kwargs['cy']

    def apply_transform(self, transform: Transform):
        self.pose = transform @ self.pose

    def rotate(self, angles, degrees=True):
        self.apply_transform(Transform.from_euler(angles, degrees=degrees))

    def translate(self, vector):
        self.apply_transform(Transform(vector))

    def project_points(self, points: np.ndarray, use_distortion=False):
        projected = Transform(self.projection_matrix) @ points[:, :3]

        # projected = ((points[:, :3] - self.position) @ self.rotation) @ self.intrinsic_matrix[:3, :3].T
        projected[:, 0] /= np.where(projected[:, 2] == 0, np.inf, projected[:, 2])
        projected[:, 1] /= np.where(projected[:, 2] == 0, np.inf, projected[:, 2])

        if use_distortion:
            projected[:, :2] = cv2.fisheye.projectPoints(
                objectPoints=np.array([points[:, :3]], dtype=np.float32),
                rvec=cv2.Rodrigues(self.extrinsic_matrix[:3, :3])[0],
                tvec=self.extrinsic_matrix[:3, 3],
                D=np.array([self.D[0], self.D[1], self.D[4], self.D[5]], dtype=np.float32),
                K=np.array(self.K, dtype=np.float32),
                alpha= self.skew
            )[0].reshape((-1, 2))

        return np.hstack([
            projected[:, :3],
            points[:, 3:]
        ])

    def get_projected_image(self, image, points, frame_transform, color_mode='default', oversample=3):
        assert image, 'No image loaded.'
        def crop_points(points, bounding_box):
            conditions = np.logical_and(points[:, :3] >= bounding_box[0], points[:, :3] < bounding_box[1])
            mask = np.all(conditions, axis=1)
            return points[mask]
        im = image.get_image().convert('RGBA')
        radius = 3
        points = np.array(random.sample(points.tolist(), int(len(points) / 2))) #reduce the number of points projected, no need to project everything
        # Project points image
        points_im = Image.new('RGBA', (im.size[0] * oversample, im.size[1] * oversample))

        draw = ImageDraw.Draw(points_im)

        if self.model == 'cylindrical':
            # handle cylindrical cameras
            epsilon = 0.0000001
            fisheye = frame_transform.inverse @ np.array([points[:,0],points[:,1], points[:,2] ], dtype=np.float)  # 3D point in camera coordinates
            fisheye = fisheye.T  # 3D point in camera coordinates
            fisheye = self.pose.inverse @ fisheye
            fisheye[:, 1] *= -1  # invert y because cylindrical y is up and cartesian y is down
            fisheye = Transform(self.extrinsic_matrix[:3, :3]) @ fisheye.T  # lift cylinder to stand up straight
            cylindrical = np.array([np.arctan2(fisheye[0, :], fisheye[2, :]),
                                    fisheye[1, :] / np.sqrt(fisheye[0, :] ** 2 + fisheye[2, :] ** 2),
                                    np.ones(fisheye.shape[1])])
            cylindrical[1, :] *= -1   # invert y because cylindrical y is up and cartesian y is down
            q = self.K @ cylindrical
            img_coords = q[[0, 1], :]  # pixels on image
            img_coords = img_coords.T

            for point in img_coords[:, :2]:
                draw.ellipse([tuple(point - radius), tuple(point + radius)], fill=tuple([255, 10, 10]))
            points_im = points_im.resize(im.size, Image.CUBIC)

            # Merge images
            projected_im = Image.composite(points_im, im, points_im)
            return projected_im
        if self.model == 'fisheye':
            wct = self.pose @ frame_transform.T
            projected = self.project_points(points, use_distortion=True)
            #  projected = crop_points(projected,
                                    #  np.array([[0, 0, 0.1], [im.size[0], im.size[1], np.inf]]))
            fisheye = self.world_transform @ np.array([projected[:,0],projected[:,1], projected[:,2] ], dtype=np.float)  # 3D point in camera coordinates
            fisheye = fisheye.T  # 3D point in camera coordinates
            fisheye = np.concatenate((np.array(fisheye),np.array(projected[:,3])[:,None]),axis=1)

            if not len(fisheye):
                return im

            colors = map_colors(fisheye, color_mode)

            for point, color in zip(fisheye[:, :2] * oversample, colors):
                draw.ellipse([tuple(point - radius), tuple(point + radius)], fill=tuple(color))
            points_im = points_im.resize(im.size, Image.CUBIC)

            # Merge images
            projected_im = Image.composite(points_im, im, points_im)
            return projected_im
        else:
            projected = self.project_points(points, use_distortion=False)
            projected = crop_points(projected,
                                    np.array([[0, 0, 0.1], [im.size[0], im.size[1], np.inf]]))

            # Returns original image if not projected points on image
            if not len(projected):
                return im

            colors = map_colors(projected, color_mode)

            for point, color in zip(projected[:, :2] * oversample, colors):
                draw.ellipse([tuple(point - radius), tuple(point + radius)], fill=tuple(color))
            points_im = points_im.resize(im.size, Image.CUBIC)

            # Merge images
            projected_im = Image.composite(points_im, im, points_im)
            return projected_im

    def __repr__(self):
        return 'LidarCamera({0}) {1}'.format(self.id, self.pose)
