import open3d as o3d
import numpy as np

class Pokemon():
    def __init__(self, pokemon_type, animal):
        self.type = pokemon_type
        self.animal = animal

        self.scalar = 8
        self.x = 7
        self.y = 10
        self.z = 5
        
        self.cloudpoints = {
            'points': None,
            'colors': None,
        }

        self.load_model()

    def load_model(self):
        self.cloudpoints['points'] = np.load('models/{}_{}/0/points.npy'.format(self.type, self.animal))
        self.cloudpoints['colors'] = np.load('models/{}_{}/0/colors.npy'.format(self.type, self.animal))

        self.cloudpoints['colors'] = self.cloudpoints['colors'] / 255
        self.cloudpoints['colors'] = self.cloudpoints['colors'][:, :3]

        # down sample
        point_cloud = o3d.geometry.PointCloud()
        point_cloud.points = o3d.utility.Vector3dVector(self.cloudpoints['points'])
        point_cloud.colors = o3d.utility.Vector3dVector(self.cloudpoints['colors'])
        
        downsampled_cloud = point_cloud.uniform_down_sample(every_k_points=5)
        downsampled_points = np.asarray(downsampled_cloud.points)
        downsampled_colors = np.asarray(downsampled_cloud.colors)

        self.cloudpoints['points'] = downsampled_points
        self.cloudpoints['colors'] = downsampled_colors

        self.cal_position()

    def walk_forward(self):
        self.x += 2
        self.cal_position()

    def walk_backward(self):
        self.x -= 2
        self.cal_position()

    def walk_left(self):
        self.y -= 2
        self.cal_position()

    def walk_right(self):
        self.y += 2
        self.cal_position()

    def cal_position(self):
        transform_matrix = np.array([
            [self.scalar, 0, 0, 0],
            [0, 0, self.scalar, 0],
            [0, self.scalar, 0, 0],
            [self.x, self.y, self.z, 1]])             # shift (前後, 左右, 高度)

        pokemon_points = self.cloudpoints['points']

        ones = np.ones((pokemon_points.shape[0], 1))
        pokemon_points = np.concatenate([pokemon_points, ones], axis=1)
        pokemon_points = pokemon_points @ transform_matrix
        pokemon_points[:, 2] = -pokemon_points[:, 2]
        self.wcs_point = pokemon_points

    def get_position(self):
        # return self.cloudpoints['points']
        return self.wcs_point
    
    def get_colors(self):
        return self.cloudpoints['colors']