import open3d as o3d
import numpy as np
from math import pi
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.theta = 2*pi
        self.accum_rotate = np.array([[1, 0, 0],[0, 1, 0],[0,0,1]])
        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]
        self.cloudpoints['colors'] = self.cloudpoints['colors'][:, ::-1]

        # 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 rotate_left(self):
        self.theta += pi/4
        self.cal_position()
    
    def rotate_right(self):
        self.theta -= pi/4
        self.cal_position()

    

    def cal_position(self):
        transform_matrix = np.array([
            [np.cos(self.theta)*self.scalar,  -np.sin(self.theta)*self.scalar,0, 0],
            [0,0 , self.scalar, 0],
            [np.sin(self.theta)*self.scalar, np.cos(self.theta)*self.scalar, 0, 0],
            [self.x, self.y, self.z, 1]])                              # rotate (x, y, z, theta)
        pokemon_points = self.cloudpoints['points']

        '''
        transform_matrix[3, :3] = transform_matrix[3, :3] @ self.accum_rotate
        tmp  =  transform_matrix[[0,2, 1], :3]
        self.accum_rotate = self.accum_rotate @ tmp
'''

        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']