import csv
import numpy as np
import torch


def convert_pc_to_box(obj_pc):
    xmin = np.min(obj_pc[:,0])
    ymin = np.min(obj_pc[:,1])
    zmin = np.min(obj_pc[:,2])
    xmax = np.max(obj_pc[:,0])
    ymax = np.max(obj_pc[:,1])
    zmax = np.max(obj_pc[:,2])
    center = [(xmin+xmax)/2, (ymin+ymax)/2, (zmin+zmax)/2]
    box_size = [xmax-xmin, ymax-ymin, zmax-zmin]
    return center, box_size

class LabelConverter(object):
    def __init__(self, file_path):
        self.raw_name_to_id = {}
        self.nyu40id_to_id = {}
        self.nyu40_name_to_id = {}
        self.scannet_name_to_scannet_id = {'cabinet':0, 'bed':1, 'chair':2, 'sofa':3, 'table':4,
            'door':5, 'window':6,'bookshelf':7,'picture':8, 'counter':9, 'desk':10, 'curtain':11,
            'refrigerator':12, 'shower curtain':13, 'toilet':14, 'sink':15, 'bathtub':16, 'others':17}  
        self.id_to_scannetid = {}
        self.scannet_raw_id_to_raw_name = {}
        self.raw_name_to_scannet_raw_id = {}

        with open(file_path, encoding='utf-8') as fd:
            rd = list(csv.reader(fd, delimiter="\t", quotechar='"'))
            for i in range(1, len(rd)):
                raw_id = i - 1
                scannet_raw_id = int(rd[i][0])
                raw_name = rd[i][1]
                nyu40_id = int(rd[i][4])
                nyu40_name = rd[i][7]
                self.raw_name_to_id[raw_name] = raw_id
                self.scannet_raw_id_to_raw_name[scannet_raw_id] = raw_name
                self.raw_name_to_scannet_raw_id[raw_name] = scannet_raw_id
                self.nyu40id_to_id[nyu40_id] = raw_id
                self.nyu40_name_to_id[nyu40_name] = raw_id
                if nyu40_name not in self.scannet_name_to_scannet_id:
                    self.id_to_scannetid[raw_id] = self.scannet_name_to_scannet_id['others']
                else:
                    self.id_to_scannetid[raw_id] = self.scannet_name_to_scannet_id[nyu40_name]

def build_rotate_mat(split, rot_aug=True, rand_angle='axis'):
    if rand_angle == 'random':
        theta = np.random.rand() * np.pi * 2
    else:
        ROTATE_ANGLES = [0, np.pi/2, np.pi, np.pi*3/2]
        theta_idx = np.random.randint(len(ROTATE_ANGLES))
        theta = ROTATE_ANGLES[theta_idx]
    if (theta is not None) and (theta != 0) and (split == 'train') and rot_aug:
        rot_matrix = np.array([
            [np.cos(theta), -np.sin(theta), 0],
            [np.sin(theta), np.cos(theta), 0],
            [0, 0, 1]
        ], dtype=np.float32)
    else:
        rot_matrix = None
    return rot_matrix

def obj_processing_post(obj_pcds, rot_aug=True):
        obj_pcds = torch.from_numpy(obj_pcds)
        rot_matrix = build_rotate_mat('val', rot_aug)
        if rot_matrix is not None:
            rot_matrix = torch.from_numpy(rot_matrix.transpose())
            obj_pcds[:, :, :3] @= rot_matrix
        
        xyz = obj_pcds[:, :, :3]
        center = xyz.mean(1)
        xyz_min = xyz.min(1).values
        xyz_max = xyz.max(1).values
        box_center = (xyz_min + xyz_max) / 2
        size = xyz_max - xyz_min
        obj_locs = torch.cat([center, size], dim=1)
        obj_boxes = torch.cat([box_center, size], dim=1)

        # centering
        obj_pcds[:, :, :3].sub_(obj_pcds[:, :, :3].mean(1, keepdim=True))

        # normalization
        max_dist = (obj_pcds[:, :, :3]**2).sum(2).sqrt().max(1).values
        max_dist.clamp_(min=1e-6)
        obj_pcds[:, :, :3].div_(max_dist[:, None, None])
        
        return obj_pcds, obj_locs, obj_boxes, rot_matrix


def pad_sequence(sequence_list, max_len=None, pad=0, return_mask=False):
    lens = [x.shape[0] for x in sequence_list]
    if max_len is None:
        max_len = max(lens)
        
    shape = list(sequence_list[0].shape)
    shape[0] = max_len
    shape = [len(sequence_list)] + shape
    dtype = sequence_list[0].dtype
    device = sequence_list[0].device
    padded_sequence = torch.ones(shape, dtype=dtype, device=device) * pad
    for i, tensor in enumerate(sequence_list):
        padded_sequence[i, :tensor.shape[0]] = tensor
    padded_sequence = padded_sequence.to(dtype)

    if return_mask:
        mask = torch.arange(max_len).to(device)[None, :] >= torch.LongTensor(lens).to(device)[:, None] # True as masked.
        return padded_sequence, mask
    else:
        return padded_sequence