# -*- coding: UTF-8 -*-
'''=================================================
@Project -> File   pram -> refframe
@IDE    PyCharm
@Author fx221@cam.ac.uk
@Date   04/03/2024 10:06
=================================================='''
import numpy as np
from localization.camera import Camera
from colmap_utils.camera_intrinsics import intrinsics_from_camera
from colmap_utils.read_write_model import qvec2rotmat


class RefFrame:
    def __init__(self, camera: Camera, id: int, qvec: np.ndarray, tvec: np.ndarray,
                 point3D_ids: np.ndarray = None, keypoints: np.ndarray = None,
                 name: str = None, scene_name: str = None):
        self.camera = camera
        self.id = id
        self.qvec = qvec
        self.tvec = tvec
        self.name = name
        self.scene_name = scene_name
        self.width = camera.width
        self.height = camera.height
        self.image_size = np.array([self.height, self.width])

        self.point3D_ids = point3D_ids
        self.keypoints = keypoints
        self.descriptors = None
        self.keypoint_segs = None
        self.xyzs = None

    def get_keypoints_by_sid(self, sid: int):
        mask = (self.keypoint_segs == sid)
        return {
            'point3D_ids': self.point3D_ids[mask],
            'keypoints': self.keypoints[mask][:, :2],
            'descriptors': self.descriptors[mask],
            'scores': self.keypoints[mask][:, 2],
            'xyzs': self.xyzs[mask],
            'camera': self.camera,
        }

        valid_p3d_ids = []
        valid_kpts = []
        valid_descs = []
        valid_scores = []
        valid_xyzs = []
        for i, v in enumerate(self.point3D_ids):
            if v in point3Ds.keys():
                p3d = point3Ds[v]
                if p3d.seg_id == sid:
                    valid_kpts.append(self.keypoints[i])
                    valid_p3d_ids.append(v)
                    valid_xyzs.append(p3d.xyz)
                    valid_descs.append(p3d.descriptor)
                    valid_scores.append(p3d.error)
        return {
            'point3D_ids': np.array(valid_p3d_ids),
            'keypoints': np.array(valid_kpts),
            'descriptors': np.array(valid_descs),
            'scores': np.array(valid_scores),
            'xyzs': np.array(valid_xyzs),
        }

    def get_keypoints(self):
        return {
            'point3D_ids': self.point3D_ids,
            'keypoints': self.keypoints[:, :2],
            'descriptors': self.descriptors,
            'scores': self.keypoints[:, 2],
            'xyzs': self.xyzs,
            'camera': self.camera,
        }

        valid_p3d_ids = []
        valid_kpts = []
        valid_descs = []
        valid_scores = []
        valid_xyzs = []
        for i, v in enumerate(self.point3D_ids):
            if v in point3Ds.keys():
                p3d = point3Ds[v]
                valid_kpts.append(self.keypoints[i])
                valid_p3d_ids.append(v)
                valid_xyzs.append(p3d.xyz)
                valid_descs.append(p3d.descriptor)
                valid_scores.append(p3d.error)
        return {
            'points3D_ids': np.array(valid_p3d_ids),
            'keypoints': np.array(valid_kpts),
            'descriptors': np.array(valid_descs),
            'scores': 1 / np.clip(np.array(valid_scores) * 5, a_min=1., a_max=20.),
            'xyzs': np.array(valid_xyzs),
            'camera': self.camera,
        }

    def associate_keypoints_with_point3Ds(self, point3Ds: dict):
        xyzs = []
        descs = []
        scores = []
        p3d_ids = []
        kpt_sids = []
        for i, v in enumerate(self.point3D_ids):
            if v in point3Ds.keys():
                p3d = point3Ds[v]
                p3d_ids.append(v)
                xyzs.append(p3d.xyz)
                descs.append(p3d.descriptor)
                scores.append(p3d.error)

                kpt_sids.append(p3d.seg_id)

        xyzs = np.array(xyzs)
        if xyzs.shape[0] == 0:
            return False

        descs = np.array(descs)
        scores = 1 / np.clip(np.array(scores) * 5, a_min=1., a_max=20.)
        p3d_ids = np.array(p3d_ids)
        uvs = self.project(xyzs=xyzs)
        self.keypoints = np.hstack([uvs, scores.reshape(-1, 1)])
        self.descriptors = descs
        self.point3D_ids = p3d_ids
        self.xyzs = xyzs
        self.keypoint_segs = np.array(kpt_sids)

        return True

    def project(self, xyzs):
        '''
        :param xyzs: [N, 3]
        :return:
        '''
        K = intrinsics_from_camera(camera_model=self.camera.model, params=self.camera.params)  # [3, 3]
        Rcw = qvec2rotmat(self.qvec)
        tcw = self.tvec.reshape(3, 1)
        Tcw = np.eye(4, dtype=float)
        Tcw[:3, :3] = Rcw
        Tcw[:3, 3:] = tcw
        xyzs_homo = np.hstack([xyzs, np.ones(shape=(xyzs.shape[0], 1))])  # [N 4]

        xyzs_cam = Tcw @ xyzs_homo.transpose()  # [4, N]
        uvs = K @ xyzs_cam[:3, :]  # [3, N]
        uvs[:2, :] = uvs[:2, :] / uvs[2, :]
        return uvs[:2, :].transpose()