import numpy as np
import torch
import sys
from collections import deque
# prob_thres = 0.5
# # spatial_thres = 3
# norm_thres = np.cos(np.deg2rad(30))
# RADIUS =

import torch.nn.functional as F
# from vPlaneRecover.util import coordinates
import trimesh
from skimage import measure

# grouping the 3 sigma region defined in prepare_data
GROUPING_RADIUS =  {16: 1, 8: 2, 4: 3}

#  window, blind, pillow, mirror, clothes, shower curtain, person, toilet, sink, lamp, bag
NONE_PLANE_ID = [0, -1, 9, 13, 18, 19, 21, 28, 31, 33, 34,  35, 37 ]

# for generating
STD =  {16: 1 ,8: 2, 4: 3 }
ADOPT_THRES = 0.03
PI = torch.tensor(3.1415926)
SQRT_2PI = 2.506628

ORTHOGONAL_THRES = np.cos(np.deg2rad(60))
# for inference
PLANE_MIN_N_VOXELS = {16: 25 ,8: 100, 4: 400 } #0.04**3*1600=0.1024 m3 -- marchine cube need at least 2

# SEM_INS_MAP = {1:0, 2:1, 22:2, 30:3, 3:4, 4:5, 5:6, 6:7, 7:8, 8:9, 10:10,
#                11:11, 12:12, 14:13, 16:14, 24:15, 28:16, 33:17, 36:18, 39:19}
#
# CLASS_LABELS = ['wall', 'floor', 'ceiling', 'whiteboard', 'cabinet', 'bed', 'chair', 'sofa', 'table', 'door',
#                 'bookshelf', 'picture',  'counter', 'desk', 'curtain', 'refrigerator', 'shower curtain', 'toilet', 'bathtub', #'sink', 'window',
#                 'otherfurniture']

# SEM_INS_MAP = {1:0, 2:1, 3:2, 4:3, 5:4, 6:5, 7:6, 8:7, 9:8, 10:9,
#                11:10, 12:11, 14:12, 16:13, 24:14,   33:15, 36:16, 39:17, 22:18, 30:19}
#
# CLASS_LABELS = ['wall', 'floor', 'cabinet', 'bed', 'chair', 'sofa', 'table', 'door', 'window',
#                 'bookshelf', 'picture',  'counter', 'desk', 'curtain', 'refrigerator',   'toilet', 'bathtub', #'sink', 'window',
#                 'otherfurniture', 'ceiling', 'whiteboard']

SEM_INS_MAP = {1:0, 2:1, 22:2} # 8:2,

CLASS_LABELS = ['wall', 'floor','ceiling'] #'door',

LAYOUT_SEM = {1, 2,22, 30, 3, 7, 12, 14, 36}
# ============= winner take all assignment (deprecated) ===========================
def get_centers(voxel_coord, voxel_norm, semsegs, center_prob_in, surface_mask, center_thres, radius, norm_thres, topk=5000):
    # center_prob = center_prob.clamp(0, 1)

    # pick top k, to save memory
    center_prob = center_prob_in.reshape([-1])
    _, indices = torch.sort(center_prob[surface_mask.reshape([-1])], descending=True)
    ori_indx = surface_mask.reshape([-1]).nonzero(as_tuple=False)
    _idx = indices[:topk]
    idx = ori_indx[_idx]
    valid_mask = torch.zeros_like(center_prob).type(torch.bool)
    valid_mask[idx] = (center_prob[idx] > center_thres)
    valid_mask = valid_mask.view(center_prob_in.size())

    if valid_mask.sum() == 0:
        return [], [], []

    # get all inliners
    inlier_coord, inlier_norm= voxel_coord[:, valid_mask.reshape([-1])], voxel_norm[:, valid_mask].reshape([3,-1])
    inlier_semseg, inlier_prob = semsegs[valid_mask].reshape([-1]), center_prob[valid_mask.reshape([-1])]

    # sort prob
    desc_prob, indices = torch.sort(inlier_prob, descending=True)
    tmp_indices = indices.clone()
    norm_ignore_mask =  (torch.sum(inlier_norm[:, indices], dim=0) == 0)
    sem_ignore_mask = (inlier_semseg[indices] == 0)
    centers = []
    center_segm_lst = []
    center_norm_lst = []
    idx = 0
    while idx < tmp_indices.shape[0]:
        if tmp_indices[idx] != -1:
            cur_center = inlier_coord[:, tmp_indices[idx]].reshape([3,1])
            center_norm = inlier_norm[:, tmp_indices[idx]].reshape([3,1])
            center_semseg = inlier_semseg[tmp_indices[idx]].reshape([1])

            # group the surrounding ones if they share same semantic label and the normal dist. < cos(30),
            sp_dist = torch.sum((inlier_coord[:, indices] - cur_center) ** 2, dim=0)

            # if no valid norm, we ingore the norm dist constrain
            if (torch.sum(center_norm.abs(), dim=0) == 0):
                norm_mask = torch.ones_like(center_semseg).type(torch.bool)
            else:
                norm_mask = (torch.sum((inlier_norm[:, indices] * center_norm), dim=0).abs()> norm_thres) | norm_ignore_mask

            # if no valid semseg, we ingore the semseg dist
            if center_semseg == 0:
                semseg_mask = torch.ones_like(center_semseg).type(torch.bool)
            else:
                semseg_mask = ((center_semseg == inlier_semseg[indices])  | sem_ignore_mask)

            center_mask = (sp_dist < radius * radius) & semseg_mask  & norm_mask
            tmp_indices[center_mask] = -1      # nms

            # set the center as the weighted sum of the highest one's surrounding --- this will cause center fell outside
            # of the semantic area, just pick the highest on
            # new_center =((inlier_coord[:, indices[center_mask]] * inlier_prob[indices[center_mask]]).sum(dim=-1) / \
            #              (inlier_prob[indices[center_mask]]).sum(dim=-1)).round().type(torch.int)

            # find the nearest voxel has semantic and normal
            # new_center_norm = voxel_norm[:, new_center[0], new_center[1], new_center[2]]
            # new_center_semg = semsegs[new_center[0], new_center[1], new_center[2]]
            # dist = torch.sum((inlier_coord[:, indices] - new_center.unsqueeze(1)) ** 2, dim=0)
            _, new_idx = torch.sort(sp_dist) #acedend
            new_center = cur_center.clone()
            new_center_norm = center_norm.clone()
            new_center_semg = center_semseg.clone()
            cnt = 0
            while ((torch.sum(new_center_norm.abs(), dim=0) == 0) or (new_center_semg == 0)) and (cnt < new_idx.shape[0]):
                cur_idx = new_idx[cnt]
                new_center = inlier_coord[:, indices][:, cur_idx].type(torch.long).view(cur_center.shape)
                new_center_norm = voxel_norm[:, new_center[0], new_center[1], new_center[2]].view(center_norm.shape)
                new_center_semg = semsegs[new_center[0], new_center[1], new_center[2]].view(center_semseg.shape)
                cnt+=1

            # if the whole ball is in empty space
            if ((torch.sum(new_center_norm.abs(), dim=0) == 0) or (new_center_semg == 0)):
                continue
            else:
                centers.append(new_center)
                center_segm_lst.append(new_center_semg)
                center_norm_lst.append(new_center_norm)

            # debug
            # if idx == 0:
            #     vert_viz = np.concatenate([vert, inlier_coord.T.numpy()], axis=0)
            #     vert_color, voxel_color = np.zeros([vert.shape[0], 3]), np.ones([inlier_coord.shape[1], 3])*255
            #     voxel_color[:,2] = 0
            # voxel_color[indices.numpy()[center_mask.numpy()]] =  np.array([0,0,255])
            # print(center_mask.sum(), new_center.round().type(torch.int))
            # color_viz = np.concatenate([vert_color, voxel_color], axis=0).astype(np.int)
            # pld = trimesh.points.PointCloud(vertices=vert_viz,  colors=color_viz, process=False)
            # pld.show()

        idx += 1

    return centers, center_segm_lst, center_norm_lst

def get_planeIns(tsdf, cfg):

    voxel_size =  tsdf.voxel_size
    normals =   tsdf.attribute_vols['plane_norm']
    semIns = tsdf.attribute_vols['semseg']
    center_prob = tsdf.attribute_vols['centroid_prob']

    mask_surface = tsdf.tsdf_vol.abs() < 1

    radius =  GROUPING_RADIUS[int(voxel_size * 100)]
    norm_thres =  np.cos(np.deg2rad(cfg.MODEL.GROUPING.NORM_THRES))
    prob_thres =  cfg.MODEL.GROUPING.PROB_THRES
    topk_prob = cfg.MODEL.GROUPING.TOPK_PROB

    coords = coordinates(center_prob.shape,device=tsdf.device)
    centers, center_segms, center_norms = get_centers( coords, normals, semIns, center_prob, mask_surface,
                                                     prob_thres, radius,norm_thres, topk_prob)

    normals, semIns = normals.reshape([3, -1]), semIns.reshape([-1])
    planeIns = torch.zeros_like(semIns)

    for i, (center_coord, center_seg, center_norm) in enumerate(zip(centers, center_segms, center_norms)):
        # semantic should be same
        semseg_mask = (semIns == center_seg)

        # normal should be similiar
        norm_mask = (torch.sum((normals * center_norm), dim=0).abs() > norm_thres)

        # distance to the plane should under threshold
        planeD = (center_norm * center_coord).sum()
        cluster_plane_dist = ((center_norm * coords).sum(dim=0) - planeD).abs()
        spatial_mask = cluster_plane_dist <= radius

        # only assign once with highest prob center
        # note we do not prevent a center to be assigned to another center with higher prob, so planeIns.max()
        # is not equal to the plane_num
        available_mask = planeIns == 0

        # assign plane instance to the picked voxels
        cluster_mask = semseg_mask & norm_mask & spatial_mask & available_mask

        # some center may do not have ins because of the limited number of inliner
        if cluster_mask.sum() > PLANE_MIN_N_VOXELS[int(voxel_size * 100)]:
            planeIns[cluster_mask] = (i + 1)
            normals[:, cluster_mask] = center_norm
        # print(cluster_mask.sum())

    return planeIns, normals,  centers, center_segms, center_norms

# ============= RANSAC based assignment =====================
def check_connection_bfs(mask, seed_idx, vol_shape):
    # bfs to build connection mask --- deprecated too slow !!!
    # vol_shape: h, w, d, where h w in 2D, d is the height dim
    h, w, d = vol_shape
    vol_mask = mask.view(vol_shape)

    # set the seed
    candidate_mask = torch.zeros_like(mask).bool()
    candidate_mask[seed_idx] = True
    candidate_mask = candidate_mask.view(vol_shape)
    idx = (candidate_mask == True).nonzero(as_tuple=False).squeeze().cpu().numpy() #shape (1,3), x y z

    # 6 direction search
    x, y, z = idx[0], idx[1], idx[2]
    queue = deque([(x, y, z)])
    while len(queue) > 0:
        cur_x, cur_y, cur_z = queue.popleft()
        sys.stdout.write('\r connection: {}/{}'.format(candidate_mask.sum(), mask.sum()))
        sys.stdout.flush()

        if cur_x - 1 >= 0 and vol_mask[cur_x-1, cur_y, cur_z] == True and candidate_mask[cur_x-1, cur_y, cur_z] == False:
            candidate_mask[cur_x-1, cur_y, cur_z] = True
            queue.append((cur_x-1, cur_y, cur_z))

        if cur_x + 1 < h and vol_mask[cur_x +1, cur_y, cur_z] == True and candidate_mask[cur_x+1, cur_y, cur_z] == False:
            candidate_mask[cur_x+1, cur_y, cur_z] = True
            queue.append((cur_x +1, cur_y, cur_z))

        if cur_y - 1 >= 0 and vol_mask[cur_x, cur_y - 1, cur_z] == True and candidate_mask[cur_x, cur_y-1, cur_z] == False:
            candidate_mask[cur_x , cur_y -1, cur_z] = True
            queue.append((cur_x, cur_y - 1, cur_z))

        if cur_y + 1 < h and vol_mask[cur_x, cur_y + 1, cur_z] == True and candidate_mask[cur_x, cur_y+1, cur_z] == False:
            candidate_mask[cur_x, cur_y + 1, cur_z] = True
            queue.append((cur_x, cur_y + 1, cur_z))

        if cur_z - 1 >= 0 and vol_mask[cur_x, cur_y, cur_z - 1] == True and candidate_mask[cur_x, cur_y, cur_z-1] == False:
            candidate_mask[cur_x, cur_y, cur_z - 1] = True
            queue.append((cur_x, cur_y, cur_z - 1))

        if cur_z + 1 < h and vol_mask[cur_x, cur_y, cur_z + 1] == True and candidate_mask[cur_x, cur_y, cur_z+1] == False:
            candidate_mask[cur_x , cur_y, cur_z + 1] = True
            queue.append((cur_x, cur_y, cur_z + 1))

    return candidate_mask.view(-1)

def check_connection(mask, seed_idx, vol_shape):
    # use make pooling to keep update until it reach
    vol_mask = mask.view(vol_shape).unsqueeze(0).unsqueeze(0)

    # set the seed
    candidate_mask = torch.zeros_like(mask).bool()
    candidate_mask[seed_idx] = True
    candidate_mask = candidate_mask.view(vol_shape).unsqueeze(0).unsqueeze(0)

    memo_mat = candidate_mask.clone()
    pre_mask = candidate_mask.clone()

    # use 3D max pooling to propgate seed
    for cnt in range(max(vol_shape)): # longest dist to flood fill equals to the largest dim
        candidate_mask = F.max_pool3d(candidate_mask.float(), kernel_size=3, stride=1, padding=1).bool() & vol_mask
        memo_mat = F.max_pool3d(memo_mat.float(), kernel_size=3, stride=1, padding=1)
        if memo_mat.sum() >= vol_shape[0] * vol_shape[1] * vol_shape[2] or (candidate_mask == vol_mask).all()\
                or (pre_mask == candidate_mask).all():
            break
        pre_mask = candidate_mask.clone()
        # sys.stdout.write('\rflood fill step: {}'.format(cnt))
        # sys.stdout.flush()
    return candidate_mask.view(-1)

def seq_ransac(coords, normals, semLab, prob, planeIns, valid_mask, mask_surface,
               norm_thres, radius, area_thres, cur_planeID,  vol_shape, n_iter=100):
    # sequential one-point plane ransac, the principle is the same as onePoint_ransan in fit_plane/util,
    # but the code is slightly different  because we do not have the instance level label,
    # so we need consider semantic label consistency here

    # sample the seeds w.r.t. their prob
    prob[~valid_mask] = 0

    # in case replacement == False, and n_iter > (prob > 0).sum(), it will return idx whose weight == 0
    # therefore, we should make sure the n_iter is always < (prob>0).sum()
    idxs = torch.multinomial(prob, min(n_iter, (prob > 0).sum()), replacement=False)

    resume = False
    n_inliers = 0
    cnt = 0
    best_mask = torch.zeros_like(semLab).type(torch.bool)
    for i in idxs:
        cnt += 1
        sample_pnt = coords[:, i].unsqueeze(1)
        sample_norm = normals[:, i].unsqueeze(1)
        sample_semg = semLab[i]

        sys.stdout.write('\rprocessing: {}, assigning plane: {}, iter {}'.format(sample_semg.item(), cur_planeID, cnt))
        sys.stdout.flush()

        # semantic should be same
        semseg_mask = (semLab == sample_semg)

        # normal should be similiar
        norm_mask = (torch.sum((normals * sample_norm), dim=0).abs() > norm_thres)

        # distance to the plane should under threshold
        planeD = (sample_norm * sample_pnt).sum()
        cluster_plane_dist = ((sample_norm * coords).sum(dim=0) - planeD).abs()
        spatial_mask = cluster_plane_dist <= radius

        # only sign once
        available_mask = planeIns == 0

        # connection mask -- only if the voxel have path connected to the sample pnt can be assigned
        cluster_mask = semseg_mask & norm_mask & spatial_mask & available_mask & mask_surface

        # put all the masks together --- Ideally we should put connection mask here, but it is too slow, we do it in post process instead
        # cluster_mask = cluster_mask & connection_mask

        n =  cluster_mask.sum()
        if n > n_inliers:
            best_mask = cluster_mask.clone()
            n_inliers = n
            center_idx, center_coord, center_semseg, center_norm = i, sample_pnt, sample_semg, sample_norm

    # best_mask = best_mask | (planeIns == cur_planeID) # for NN case
    # ransac will stop if the best plane_area < area_thres
    if n_inliers >= area_thres:
        # break the mask if contain 2 seperate part
        connection_mask = check_connection(best_mask, center_idx, vol_shape)
        best_mask = connection_mask & best_mask

        # only update label if the area is sufficent large
        if best_mask.sum() < area_thres:
            return planeIns, valid_mask, resume, normals, sample_pnt, sample_semg, sample_norm
        # ======= debug========
        # voxels = coords[:, mask_surface].T.cpu().numpy()
        # voxel_color = np.zeros([voxels.shape[0], 3])
        # voxel_color[best_mask.cpu().numpy()[mask_surface.cpu().numpy()]] =  np.array([0,0,255])
        #
        # pld = trimesh.points.PointCloud(vertices=voxels, colors=voxel_color, process=False)
        # pld.show()
        #
        # best_mask = connection_mask & best_mask
        # voxel_color[best_mask[mask_surface].cpu().numpy()] = np.array([255, 0, 0])
        # pld = trimesh.points.PointCloud(vertices=voxels, colors=voxel_color, process=False)
        # pld.show()

        # -------- connect
        # if sample_semg ==3:
        #     pld = trimesh.points.PointCloud(vertices=voxels, colors=voxel_color, process=False)
        #     pld.show()

            # vol_mask = best_mask.view(vol_shape).unsqueeze(0).unsqueeze(0)
            #
            # # set the seed
            # candidate_mask = torch.zeros_like(best_mask).bool()
            # candidate_mask[center_idx] = True
            # candidate_mask = candidate_mask.view(vol_shape).unsqueeze(0).unsqueeze(0)
            #
            # memo_mat = candidate_mask.clone()
            # pre_mask = candidate_mask.clone()
            # # use 3D max pooling to propgate seed
            # for cnt in range(max(vol_shape)):  # longest dist to flood fill equals to the largest dim
            #     candidate_mask = F.max_pool3d(candidate_mask.float(), kernel_size=3, stride=1, padding=1).bool() & vol_mask
            #
            #     connection_mask = candidate_mask.view(-1)
            #     voxel_color[connection_mask[mask_surface].cpu().numpy()] = np.array([255, 0, 0])
            #     pld = trimesh.points.PointCloud(vertices=voxels, colors=voxel_color, process=False)
            #     pld.show()
            #
            #     memo_mat = F.max_pool3d(memo_mat.float(), kernel_size=3, stride=1, padding=1)
            #     if memo_mat.sum() >= vol_shape[0] * vol_shape[1] * vol_shape[2] or (candidate_mask == vol_mask).all():
            #         break

        planeIns[best_mask] = cur_planeID
        valid_mask[best_mask] = False # seed will only be assigned once as well
        resume = True

        # get weird result -- perform least square fit in the ininlier -- we should not change normal, because the voxel norm if from tsdf_grad
        # it can be very different in the boundary

        mean_center = torch.mean(coords[:,best_mask].float(), dim=1, keepdim=True)
        dist = torch.sum((coords[:, best_mask].float() - mean_center).abs(), dim=0) # use l1 dist find nearest inliers to mean_center
        new_idx = dist.argmin(dim=0)  # acedend
        center_coord = coords[:,best_mask][:, new_idx].unsqueeze(1)
        normals[:, best_mask] = center_norm
    else:
        center_coord, center_semseg, center_norm = sample_pnt, sample_semg, sample_norm # just return sth, will not be used

    return planeIns, valid_mask, resume, normals,  center_coord, center_semseg, center_norm

def get_planeIns_RANSAC(tsdf, cfg):
    # init necessary variables
    voxel_size =  tsdf.voxel_size
    normals =   tsdf.attribute_vols['plane_norm']
    semLab = tsdf.attribute_vols['semseg']
    center_prob = tsdf.attribute_vols['centroid_prob']

    radius =  GROUPING_RADIUS[int(voxel_size * 100)]
    norm_thres =  np.cos(np.deg2rad(cfg.MODEL.GROUPING.NORM_THRES))
    prob_thres =  cfg.MODEL.GROUPING.PROB_THRES
    area_thres =  PLANE_MIN_N_VOXELS[int(voxel_size * 100)]

    coords = coordinates(center_prob.shape,device=tsdf.device)

    normals, semLab = normals.reshape([3, -1]), semLab.reshape([-1])
    planeIns = torch.zeros_like(semLab)

    # pick valid center
    mask_surface = tsdf.tsdf_vol.abs() < 1
    center_prob_flat = center_prob.clone().reshape([-1])
    seed_mask = (center_prob_flat > prob_thres) & mask_surface.reshape([-1])

    # _, indices = torch.sort(center_prob_flat[seed_mask], descending=True)
    # ori_indx = seed_mask.nonzero(as_tuple=False)
    # idx_in_ori_idx = ori_indx[indices] # convert the valid set idx to the whole set idx

    if seed_mask.sum() == 0:
        return planeIns, normals, [], [], []

    # start sequential RANSAC for each pred_semantic label
    cur_planeId = 1
    centers, center_segms, center_norms = [], [], []
    for semid in torch.unique(semLab):
        if semid <= 0: continue # ignore invalid semantic label
        resume_ransac = True
        tmp_valid_mask = seed_mask & (semLab == semid)
        # if semid == 22:
        #     print(123)
        # Start seq_Ransac,
        while resume_ransac:
            if tmp_valid_mask.sum() == 0: #quit if no seeds exist
                resume_ransac = False
            else:
                planeIns, tmp_valid_mask, resume_ransac,  normals, center_coord, center_semseg, center_norm  =\
                    seq_ransac(coords, normals, semLab, center_prob_flat.clone(), planeIns, tmp_valid_mask , mask_surface.reshape([-1]),
                           norm_thres, radius,  area_thres, cur_planeId, vol_shape=center_prob.squeeze().shape, n_iter=500)

                if resume_ransac:
                    cur_planeId += 1
                    centers.append(center_coord)
                    center_segms.append(center_semseg)
                    center_norms.append(center_norm)

    return planeIns, normals,  centers, center_segms, center_norms

# ============== NN association ===========

def seq_ransac_NN(coords, normals, semLab, planeIns, valid_mask, mask_surface,
               norm_thres, radius, area_thres, cur_planeID,  vol_shape, n_iter=100):
    # sequential one-point plane ransac, the principle is the same as onePoint_ransan in fit_plane/util,
    # but the code is slightly different  because we do not have the instance level label,
    # so we need consider semantic label consistency here

    # sample the seeds w.r.t. their prob
    prob = valid_mask.float().clone()

    # in case replacement == False, and n_iter > (prob > 0).sum(), it will return idx whose weight == 0
    # therefore, we should make sure the n_iter is always < (prob>0).sum()
    idxs = torch.multinomial(prob, min(n_iter, (prob > 0).sum()), replacement=False)

    resume = False
    n_inliers = 0
    cnt_iter = 0
    best_mask = torch.zeros_like(semLab).type(torch.bool)
    # ransac iter
    for i in idxs:
        cnt_iter += 1
        sample_pnt = coords[:, i].unsqueeze(1)
        sample_norm = normals[:, i].unsqueeze(1)
        sample_semg = semLab[i]

        sys.stdout.write('\rprocessing: {}, assigning plane: {}, iter {}'.format(sample_semg.item(), cur_planeID, cnt_iter))
        sys.stdout.flush()

        # semantic should be same
        semseg_mask = (semLab == sample_semg)

        # normal should be similiar
        norm_mask = (torch.sum((normals * sample_norm), dim=0).abs() > norm_thres)

        # distance to the plane should under threshold
        planeD = (sample_norm * sample_pnt).sum()
        cluster_plane_dist = ((sample_norm * coords).sum(dim=0) - planeD).abs()
        spatial_mask = cluster_plane_dist <= radius

        # only sign once
        available_mask = (planeIns == 0) #| cur_planeID

        # connection mask -- only if the voxel have path connected to the sample pnt can be assigned
        cluster_mask = semseg_mask & norm_mask & spatial_mask & available_mask & mask_surface

        # put all the masks together --- Ideally we should put connection mask here, but it is too slow, we do it in post process instead
        # cluster_mask = cluster_mask & connection_mask

        n =  cluster_mask.sum()
        if n > n_inliers:
            best_mask = cluster_mask.clone()
            n_inliers = n
            center_idx, center_coord, center_semseg, center_norm = i, sample_pnt, sample_semg, sample_norm

    # best_mask = best_mask | (planeIns == cur_planeID) # for NN case
    # ransac will stop if the best plane_area < area_thres
    if n_inliers >= area_thres:
        print("{} assinged".format(cur_planeID))
        # break the mask if contain 2 seperate part
        connection_mask = check_connection(best_mask | valid_mask, center_idx, vol_shape)
        final_best_mask = connection_mask & best_mask

        # only update label if the area is sufficent large
        if final_best_mask.sum() < area_thres:
            return planeIns, valid_mask, resume, normals, sample_pnt, sample_semg, sample_norm, final_best_mask
        # ======= debug========
        # voxels = coords[:, mask_surface].T.cpu().numpy()
        # voxel_color = np.zeros([voxels.shape[0], 3])
        # voxel_color[best_mask.cpu().numpy()[mask_surface.cpu().numpy()]] =  np.array([0,0,255])
        #
        # pld = trimesh.points.PointCloud(vertices=voxels, colors=voxel_color, process=False)
        # pld.show()
        #
        # best_mask = connection_mask & best_mask
        # voxel_color[best_mask[mask_surface].cpu().numpy()] = np.array([255, 0, 0])
        # pld = trimesh.points.PointCloud(vertices=voxels, colors=voxel_color, process=False)
        # pld.show()

        planeIns[final_best_mask] = cur_planeID
        # planeIns[torch.logical_xor(final_best_mask, best_mask)] = 0 # the discard note connected (if itone set to be 0
        # valid_mask[final_best_mask] = True # add new added area for
        resume = True

        # get weird result -- perform least square fit in the ininlier -- we should not change normal, because the voxel norm if from tsdf_grad
        # it can be very different in the boundary
        mean_center = torch.mean(coords[:,best_mask].float(), dim=1, keepdim=True)
        dist = torch.sum((coords[:, best_mask].float() - mean_center).abs(), dim=0) # use l1 dist find nearest inliers to mean_center
        new_idx = dist.argmin(dim=0)  # acedend
        center_coord = coords[:,best_mask][:, new_idx].unsqueeze(1)
        normals[:, final_best_mask] = center_norm
    else:
        print("{} is too small, iliner {} not assinged".format(cur_planeID, n_inliers))
        center_coord, center_semseg, center_norm, final_best_mask = sample_pnt, sample_semg, sample_norm, best_mask # just return sth, will not be used

    return planeIns, valid_mask, resume, normals,  center_coord, center_semseg, center_norm, final_best_mask

def FloodFill(mask, normals, semLab, norm_thres, seed_idx, vol_shape):

    # set the seed
    candidate_mask = torch.zeros_like(mask).bool()
    candidate_mask[seed_idx] = True
    ctr_semg = semLab[seed_idx]
    # ctr_norm = normals[:, seed_idx]

    if ctr_semg in NONE_PLANE_ID:
        return candidate_mask.view(-1)

    candidate_mask = candidate_mask.view(vol_shape).unsqueeze(0).unsqueeze(0)

    pre_mask = candidate_mask.clone()

    # use 3D max pooling to propgate seed
    for cnt in range(max(vol_shape)): # longest dist to flood fill equals to the largest dim
        candidate_mask = F.max_pool3d(candidate_mask.float(), kernel_size=3, stride=1, padding=1).bool()

        # semantic should be same
        semseg_mask = (semLab == ctr_semg)
        # normal should be similiar
        # norm_mask = (torch.sum((normals * ctr_norm), dim=0).abs() > ORTHOGONAL_THRES) # same as gt generate
        # distance to the plane should under threshold -- emperical no difference

        candidate_mask = (candidate_mask.view(-1) & semseg_mask &  mask).view(vol_shape).unsqueeze(0).unsqueeze(0)

        if  (pre_mask == candidate_mask).all():
            break

        # new_mask = torch.logical_xor(pre_mask, candidate_mask) # the filled region in this iter
        pre_mask = candidate_mask.clone()

        # sys.stdout.write('\rflood fill step: {}'.format(cnt))
        # sys.stdout.flush()
    return candidate_mask.view(-1)

def get_center_FloodFill(coords, normals, semLab, seed_mask, norm_thres, idx_in_ori_idx, vol_shape):
    # center_prob is in vol_shape
    vaild_mask = seed_mask.clone()
    centers, center_segm_lst, center_norm_lst, center_ids, center_mask = [], [], [], [], []

    for id in idx_in_ori_idx:
        if vaild_mask[id]:
            cur_seed_mask = FloodFill(seed_mask, normals, semLab, norm_thres, id, vol_shape)
            if cur_seed_mask.sum() <= 4: continue # a instance center should have at least other 3 support

            mean_center = torch.mean(coords[:, cur_seed_mask].float(), dim=1, keepdim=True)
            dist = torch.sum((coords[:, cur_seed_mask].float() - mean_center).abs(), dim=0)

            new_idx = dist.argmin(dim=0)  # acedend

            # this normal update lead to wrong normal sometimes , but on average it leads to a bette results
            norm = normals[:, cur_seed_mask]
            if norm.shape[1] > 10000:   # to control the memory usage
                mask =  torch.randperm(norm.shape[1])[:10000]
                norm = norm[:, mask]
            u, s, v = torch.svd((norm.T))
            new_normal = v[:, 0].unsqueeze(1)
            seed_norm = normals[:, cur_seed_mask][:, new_idx].unsqueeze(1)
            new_normal = -new_normal if (new_normal*seed_norm).sum() < 0 else new_normal

            # we cannot use least square, some slice mask will turn to change their normal direction
            # new_normal, _ = torch.lstsq(torch.ones_like(coords[:1, cur_seed_mask]).float().T, coords[:, cur_seed_mask].float().T)
            # new_normal = new_normal[:coords.shape[0]]/new_normal[:coords.shape[0]].norm()

            centers.append(coords[:, cur_seed_mask][:, new_idx].unsqueeze(1))
            center_norm_lst.append(new_normal)
            # center_norm_lst.append(normals[:, cur_seed_mask][:, new_idx].unsqueeze(1))
            center_segm_lst.append(semLab[cur_seed_mask][new_idx])
            center_ids.append(id)
            center_mask.append(cur_seed_mask)

            vaild_mask[cur_seed_mask] = False

    return centers, center_norm_lst, center_segm_lst, center_ids, center_mask

def get_planeIns_NN(tsdf, cfg, ransac=False):
    # flood fill
    voxel_size = tsdf.voxel_size
    normals = tsdf.attribute_vols['plane_norm']
    semLab = tsdf.attribute_vols['semseg']
    plane_cls = tsdf.attribute_vols['plane_cls']
    center_prob = tsdf.attribute_vols['centroid_prob']

    radius = GROUPING_RADIUS[int(voxel_size * 100)]
    norm_thres = np.cos(np.deg2rad(cfg.MODEL.GROUPING.NORM_THRES))
    prob_thres = cfg.MODEL.GROUPING.PROB_THRES
    area_thres = PLANE_MIN_N_VOXELS[int(voxel_size * 100)]

    coords = coordinates(center_prob.shape, device=tsdf.device).type(torch.int16)

    normals, semLab, plane_cls = normals.reshape([3, -1]).type(torch.float), semLab.reshape([-1]), plane_cls.reshape([4, -1])#.type(coords.dtype)
    planeIns = torch.zeros_like(semLab).type(torch.int16)

    # pick valid center
    mask_surface = (tsdf.tsdf_vol.abs() < 1).reshape([-1]) & (plane_cls[0] > 0.5) #must be plane to be considered
    center_prob_flat = center_prob.clone().reshape([-1])
    seed_mask = (center_prob_flat > prob_thres) & mask_surface

    center_pred = (coords + plane_cls[1:].round().type(coords.dtype))

    if seed_mask.sum() == 0:
        return planeIns, normals, [], [], []

    # collect potential cluster center idx
    _, indices = torch.sort(center_prob_flat[seed_mask], descending=True)
    ori_indx = seed_mask.nonzero(as_tuple=False)
    idx_in_ori_idx = ori_indx[indices] # convert the valid set idx to the whole set idx

    # get center list
    centers, center_norms, center_segms, center_ids, center_masks = \
        get_center_FloodFill(coords, normals, semLab, seed_mask, norm_thres,  idx_in_ori_idx, center_prob.shape)

    n_ins = len(centers)
    # planeIns_distance = torch.ones([semLab.shape[0], n_ins], dtype=torch.int16).to(semLab.device).cpu() * 10000 #no gpu memory if use int64, if the center is hugh amount like scene 500_00 still not enough

    for i, (ctr_pnt, ctr_norm, ctr_semg, ctr_id, ctr_mask) in enumerate(zip(centers, center_norms, center_segms, center_ids, center_masks)):
        # semantic should be same
        cluster_mask = (semLab == ctr_semg) & mask_surface

        # normal should be similiar
        norm_mask = ((normals[:, cluster_mask] * ctr_norm).sum().abs() > norm_thres)
        cluster_mask[cluster_mask] &= norm_mask

        # distance to the plane should under threshold
        # planeD = (ctr_norm * ctr_pnt).sum()
        # cluster_plane_dist = ((coords * ctr_norm).sum(0) - planeD).abs()
        # spatial_mask = (cluster_plane_dist <= radius)
        # cluster_mask &= spatial_mask

       # the predicted center of the current candidates voxel is within the center support mask
        # https://stackoverflow.com/questions/41234161/check-common-elements-of-two-2d-numpy-arrays-either-row-or-column-wise
        cur_center_pred = center_pred[:, cluster_mask].T
        ctr_support_area = coords[:, ctr_mask].T
        # print(cur_center_pred.shape, ctr_support_area.shape, prob_thres)
        b_voted = (cur_center_pred[:, None] == ctr_support_area).all(-1).any(-1)
        cluster_mask[cluster_mask] &= b_voted

        # cluster_mask = check_connection(cluster_mask, ctr_id, center_prob.shape)
        planeIns[cluster_mask] = i +1

        # should associate to the cloest center instead of plane to preserve every instance
        # planeIns_distance[cluster_mask, i] =(coords[:,cluster_mask] - ctr_pnt).abs().sum(dim=0).type(torch.int16).cpu() # cluster_plane_dist[cluster_mask].type(torch.int16) #

    # discard very small instance
    if ransac:
        additional_id = 2
        for x in range(n_ins+1):
            cur_planeId = x + 1

            # merge unassigned voxels to exising instance first, and then generate plane in other unassigned region
            potentail_area_mask = (planeIns == cur_planeId) if x<n_ins else ((planeIns == 0) & mask_surface)  #center_mask[x] & (planeIns == 0)
            # print(n_ins)
            for semid in torch.unique(semLab[potentail_area_mask]): # for the planeIns case
                if semid.item() in NONE_PLANE_ID:continue
                resume_ransac = True
                tmp_valid_mask = potentail_area_mask & (semLab == semid)
                while resume_ransac:
                    if tmp_valid_mask.sum() == 0:  # quit if no seeds exist
                        resume_ransac = False
                    else:

                        planeIns, tmp_valid_mask, resume_ransac, normals, center_coord, center_semseg, center_norm, new_add_mask = \
                                seq_ransac_NN(coords, normals, semLab,  planeIns, tmp_valid_mask,
                                           mask_surface,
                                           norm_thres, radius, area_thres, cur_planeId, vol_shape=center_prob.squeeze().shape,
                                           n_iter=500)
                        # -------- debug -----------------
                        # if x == n_ins :
                        #     voxel_color[tmp_valid_mask[mask_surface.reshape([-1])].cpu().numpy()] = np.array([255, 0, 0])
                        #     voxel_color[new_add_mask[mask_surface.reshape([-1])].cpu().numpy()] = np.array([0, 255, 0])
                        #     print(resume_ransac)
                        #     pld = trimesh.points.PointCloud(vertices=voxels, colors=voxel_color, process=False)
                        #     pld.show()

                        # offer more id for the unassigned region
                        if cur_planeId >= n_ins + 1:
                            if resume_ransac:
                                centers.append(center_coord)
                                center_norms.append(center_norm)
                                center_segms.append(center_semseg)
                                tmp_valid_mask ^= new_add_mask

                            cur_planeId = n_ins + additional_id
                            additional_id += 1
                        else:
                            # enlarge merging seed
                            tmp_valid_mask |= new_add_mask

    for cur_planeId in torch.unique(planeIns):
        planeIns[planeIns == cur_planeId] = torch.where(planeIns[planeIns == cur_planeId].sum() < area_thres,
                                                        torch.zeros_like(planeIns[planeIns == cur_planeId]),
                                                        planeIns[planeIns == cur_planeId])

    return planeIns, normals, centers, center_segms, center_norms

def expl(val, ord=8, verbose = False):
    # approx the exp to speed up the process, as this part is done on cpu
    # idea from  https://codingforspeed.com/using-faster-exponential-approximation/
    # there is another brillant idea, but torch doest not support bit move >> or <<
    #  w.r.t. http://perso.citi-lab.fr/fdedinec/recherche/publis/2005-FPT.pdf
    x = 1. + val / (2**ord)
    for _ in range(ord):
        x *= x
        if verbose:
            print(x)
    return  x

def get_heatmap( planeIns_vol, voxel_size, tsdf_vol_sz):
    ## we expect flatten planeIns here
    # todo: This will limited only 1 batch in one GPU, change this to enable multiple batch
    coords_voxel = coordinates(tsdf_vol_sz[1:], device=planeIns_vol.device)
    centroid_vol = torch.zeros(planeIns_vol.view(-1).shape, device=planeIns_vol.device)

    # ====== build centroid heatmap =====
    unique_id = torch.unique(planeIns_vol)
    # centers = torch.zeros([3, len(unique_id)-1]).type(torch.long)
    for k, id in enumerate(unique_id):
        if id <= 0: continue  # none plane
        heatmap_mask = (planeIns_vol == id)  # ensure only update once
        cur_plane_voxels = coords_voxel[:, heatmap_mask]
        # cur_plane_norm = normal_vol[planeIns_vol == id, :]
        cur_center = torch.mean((cur_plane_voxels).type(torch.float), 1, keepdims=True)

        # if we use the mean center directly, it can fall into some empty voxel, then we cannot use its normal and depth
        # to generate plane proposal, therefore, we have to pick the one near to it and have the same plane normal
        dist = torch.sum((cur_plane_voxels - cur_center).abs(), dim=0)
        new_idx = torch.argmin(dist)
        new_center = cur_plane_voxels[:, new_idx].type(torch.long)

        dist = torch.sum((coords_voxel - new_center.unsqueeze(1).type(torch.float)) ** 2, dim=0)

        # plane ins id start from 1
        # only voxel belong to the instance can be update.
        # For a voxel does not have instance id, it normal will be different from plane norm, and so cannot be used to
        # propose the plane in the groupping stage

        std = (heatmap_mask.sum() * ADOPT_THRES + STD[voxel_size]) / 3
        # Gaussian_prob = torch.exp(-(dist / (2 * std))) / (torch.sqrt(std) * torch.sqrt(2 * PI))
        Gaussian_prob = expl(-(dist / (2 * std)).clamp(0, 500)) / ( torch.sqrt(std) * SQRT_2PI)  # clamp the dist for numerical stable
        Gaussian_prob[~heatmap_mask] = 0
        # change one of the max value to ensure only one center after approx
        max_val = Gaussian_prob.max()
        top_mask = (Gaussian_prob == max_val)

        if top_mask.sum() > 1:
            idx = top_mask.nonzero(as_tuple=False)
            Gaussian_prob[idx[0]] += 1e-4
        assert (Gaussian_prob == Gaussian_prob.max()).sum() == 1, "only one center for one instance"
        Gaussian_prob /= max_val  # normalize

        # only update the voxels have the same plane ins or those not belong to any plane once
        centroid_vol += Gaussian_prob
    return  centroid_vol.view(tsdf_vol_sz)

def coordinates(voxel_dim, device=torch.device('cuda'), b_flat=True):
    """ 3d meshgrid of given size.

    Args:
        voxel_dim: tuple of 3 ints (nx,ny,nz) specifying the size of the volume

    Returns:
        torch long tensor of size (3,nx*ny*nz)
    """

    nx, ny, nz = voxel_dim
    x = torch.arange(nx, dtype=torch.long, device=device)
    y = torch.arange(ny, dtype=torch.long, device=device)
    z = torch.arange(nz, dtype=torch.long, device=device)
    x, y, z = torch.meshgrid(x, y, z)
    if b_flat:
        return torch.stack((x.flatten(), y.flatten(), z.flatten()))
    else:
        return torch.stack((x, y, z))

def find_closet(src_arr, tgt_arr):
    # given 1D tgt_arr find closet val idx in src_arr
    # https://stackoverflow.com/questions/20780017/vectorize-finding-closest-value-in-an-array-for-each-element-in-another-array
    idx1 = torch.searchsorted(src_arr, tgt_arr).clamp(0, len(src_arr) -1)
    idx2 = (idx1 - 1).clamp(0, len(src_arr) - 1)

    diff1 = src_arr[idx1] - tgt_arr
    diff2 = tgt_arr - src_arr[idx2]

    final_idx = torch.where(diff1 <= diff2, idx1, idx2)
    return final_idx

def Eud2sphere(norm, d, rhos, thetas, phis):
    _rhos = d #param[3:, :]
    rhoIdxs = find_closet(rhos, _rhos)

    _thetas = torch.acos(norm[2]) #np.arccos(param[2:3, :])
    thetaIdxs = find_closet(thetas, _thetas)

    _phis = torch.acos(norm[0] / torch.sin(_thetas))
    phiIdxs = find_closet(phis, _phis)

    # ensure theta==0 only bring one activation
    if thetaIdxs == 0:
        phiIdxs = 0

    return rhoIdxs, thetaIdxs, phiIdxs

def sphere2norm(theta, phi):
    p_norm = torch.tensor([np.sin(theta) * np.cos(phi),
                         np.sin(theta) * np.sin(phi),
                         np.cos(theta)])
    return p_norm / p_norm.norm()
# ========================
# from vPlaneRecover.tsdf import SEM_INS_MAP

def get_planeInsVert_frmHT(tsdf, voxel_sz, origin, semLab_in, param_htmap_in, vote_idx,  rhos, thetas, phis, cfg):

    tsdf_vol = -tsdf.squeeze()

    # don't close surfaces using unknown-empty boundry
    tsdf_vol[tsdf_vol == -1] = 1

    tsdf_vol = tsdf_vol.clamp(-1, 1).cpu().numpy()

    semLab = semLab_in.squeeze()
    param_htmap = param_htmap_in.squeeze()

    verts_mc, faces, _, _ = measure.marching_cubes(tsdf_vol, level=0)
    verts_ind = np.round(verts_mc).astype(int)
    n_verts = verts_mc.shape[0]

    # in sone weird case ind will exceede the range
    d, h, w = semLab.shape
    if np.any(verts_ind[:, 0] > d - 1) or np.any(verts_ind[:, 1] > h - 1) or np.any(verts_ind[:, 2] > w - 1) or \
            np.any(verts_ind[:, 0] < 0) or np.any(verts_ind[:, 1] < 0) or np.any(verts_ind[:, 2] < 0):
        return trimesh.Trimesh(vertices=np.zeros((1, 3)))

    verts = verts_mc * voxel_sz + origin.cpu().numpy() # n*3
    verts_ind = np.round(verts_mc).astype(int)
    tmp_mesh = trimesh.Trimesh(vertices=verts, faces=faces, process=False)
    verts_norm = torch.from_numpy(tmp_mesh.vertex_normals).float().to(semLab.device) # n*3

    semseg_vol = semLab_in.detach().cpu().numpy()
    semseg_verts = semseg_vol[verts_ind[:, 0], verts_ind[:, 1], verts_ind[:, 2]]

    voxel_coord_np = coordinates(semLab.shape, torch.device('cpu'), False).float().numpy() #semLab.device
    # voxel_coord_np = (voxel_coord_np - np.array([semLab.shape]).T / 2.).reshape([3, d, h, w])
    voxel_coord = torch.from_numpy(voxel_coord_np[:, verts_ind[:, 0], verts_ind[:, 1], verts_ind[:, 2]]).float()
    voxel_coord = (voxel_coord -  torch.tensor([semLab.shape]).T / 2.).to(semLab.device)

    prob_thres = cfg.MODEL.GROUPING.PROB_THRES
    nms_r = 7
    planeIns = np.zeros_like(semseg_verts)
    ins_id = 1

    uniq_sem = (torch.unique(semLab))
    _param_htmap = F.threshold(param_htmap, prob_thres, 0)
    for sem_id in uniq_sem:
        # sys.stdout.write('\r process:{}'.format(CLASS_LABELS[sem_id.item()]))
        # sys.stdout.flush()
        if sem_id.item() not in SEM_INS_MAP:
            continue

        # load htmap and voxels under current semantic
        param_channel = SEM_INS_MAP[sem_id.item()]
        tmp = _param_htmap[param_channel]
        semseg_mask_np = (semseg_verts == sem_id.item())  # .view(-1)

        # nms
        nms_padding = (nms_r) // 2
        tmp_pool = F.max_pool3d(tmp.unsqueeze(0).unsqueeze(0), kernel_size=nms_r, stride=1,
                                padding=nms_padding).squeeze()
        tmp[tmp != tmp_pool] = 0
        params = tmp.nonzero(as_tuple=False)

        # assign
        h, w, d = semLab.shape
        rho_tol = (cfg.MODEL.BACKBONE3D.RHO_STEP[-1] +1) #if sem_id.item() in LAYOUT_SEM else (cfg.MODEL.BACKBONE3D.RHO_STEP[-1] +1)
        # componets = torch.arange(h*w*d).reshape([h,w,d]).to(semLab.device).float()
        norm_score_glb = torch.zeros([n_verts]).float().to(semLab.device)
        # tmp = torch.zeros_like(cur_param_htmap)
        for param_id in params:
            # tmp[param_id[0], param_id[1], param_id[2]] = 1.
            cur_norm = sphere2norm(thetas[param_id[1]], phis[param_id[2]]).float().to(semLab.device)
            cur_rho = torch.tensor(rhos[param_id[0]] * 2).float().to(semLab.device)

            norm_score = (cur_norm.view(1, 3) * verts_norm).sum(dim=1).abs()
            verts_norm_mask = norm_score > norm_score_glb

            verts_in_mask = ((voxel_coord.T @ cur_norm - cur_rho).abs()) <= rho_tol

            valid_pln_ins_mask = np.logical_and( (verts_in_mask & verts_norm_mask).detach().cpu().numpy(), semseg_mask_np )

            if valid_pln_ins_mask.sum() < 4:
                continue

            planeIns[valid_pln_ins_mask] = ins_id
            ins_id += 1
            norm_score_glb[valid_pln_ins_mask] = torch.where((norm_score_glb < norm_score)[valid_pln_ins_mask], norm_score[valid_pln_ins_mask],
                                                   norm_score_glb[valid_pln_ins_mask])

    return planeIns

def get_planeIns_htmap(tsdf, voxel_sz, origin, semLab_in, param_htmap_in, plane_norm, vote_idx,  rhos, thetas, phis, cfg, upscale=2):
    # voxel_size = tsdf.voxel_size
    valid_tsdf= tsdf.abs().squeeze() < 1
    semLab = semLab_in.squeeze()
    param_htmap = param_htmap_in.squeeze()

    verts_mc, faces, _, _ = measure.marching_cubes(tsdf, level=0)
    verts_ind = np.round(verts_mc).astype(int)

    verts = verts_mc * voxel_sz + origin.cpu().numpy()

    prob_thres = cfg.MODEL.GROUPING.PROB_THRES
    nms_r = 3
    voxel_coord = coordinates(semLab.shape, semLab.device).float()
    voxel_coord = voxel_coord - torch.tensor([semLab.shape]).to(semLab.device).T / 2.

    planeIns = torch.zeros_like(semLab).float()
    ins_id = 1
    uniq_sem = (torch.unique(semLab))

    _param_htmap = F.threshold(param_htmap, prob_thres, 0)

    for sem_id in uniq_sem:

        if sem_id.item() not in SEM_INS_MAP:
            continue

        # load htmap and voxels under current semantic
        param_channel = SEM_INS_MAP[sem_id.item()]
        tmp = _param_htmap[param_channel]
        cur_mask = (semLab == sem_id) #.view(-1)

        # nms
        nms_padding = (nms_r - 1) // 2
        tmp_pool = F.max_pool3d(tmp.unsqueeze(0).unsqueeze(0), kernel_size=nms_r, stride=1, padding=nms_padding).squeeze()
        tmp[tmp != tmp_pool] = 0
        params = tmp.nonzero(as_tuple=False)

        # assign
        valid_vlxs = valid_tsdf & cur_mask
        h, w, d = semLab.shape
        rho_tol = (cfg.MODEL.BACKBONE3D.RHO_STEP[-1] * 2 +1) if sem_id.item() in LAYOUT_SEM else (cfg.MODEL.BACKBONE3D.RHO_STEP[-1] +1)
        componets = torch.arange(h*w*d).reshape([h,w,d]).to(semLab.device).float()
        norm_score_glb = torch.zeros_like(semLab).float()

        for param_id in params:
            # tmp[param_id[0], param_id[1], param_id[2]] = 1.
            cur_norm = sphere2norm(thetas[param_id[1]], phis[param_id[2]]).float().to(semLab.device)
            cur_rho = torch.tensor(rhos[param_id[0]] * 2).float().to(semLab.device)

            planes_in_vol = ((voxel_coord.T @ cur_norm - cur_rho).abs()).reshape(semLab.shape)<= rho_tol

            norm_score = (cur_norm.view(3,1,1,1) * plane_norm).sum(dim=0).abs()
            norm_mask =  norm_score > norm_score_glb

            valid_planes_vlxs = norm_mask  &  planes_in_vol & valid_vlxs

            if valid_planes_vlxs.sum()<8:
                continue

            # connected components with floodfill
            comp = componets.clone()
            comp[~valid_planes_vlxs] = 0
            pre_comp = comp.clone()
            for cnt in range(max(semLab.shape)):  # longest dist to flood fill equals to the largest dim
                comp[valid_planes_vlxs] = F.max_pool3d(comp.unsqueeze(0).unsqueeze(0),
                                              kernel_size=3, stride=1, padding=1).squeeze()[valid_planes_vlxs]
                if (pre_comp == comp).all():
                    break
                else:
                    pre_comp = comp.clone()

            # assign label
            uniq_ins = torch.unique(comp)
            tmp_mask_sum = torch.zeros(len(uniq_ins))
            for i, tmp_id in enumerate(uniq_ins):
                if tmp_id < 1: continue
                tmp_mask = (comp == tmp_id)
                tmp_mask_sum[i] = tmp_mask.sum()

            tmp_thres = tmp_mask_sum.max() * 0.1
            for i, tmp_id in enumerate(uniq_ins):
                if tmp_id < 1: continue
                tmp_mask = (comp == tmp_id)
                if tmp_mask_sum[i] < tmp_thres: continue # filter extremely small ones
                planeIns[tmp_mask] = ins_id

                ins_id += 1

                # cancel the score if the voxel is not assigned
                norm_score_glb[tmp_mask] = torch.where((norm_score_glb < norm_score)[tmp_mask], norm_score[tmp_mask],
                                                         norm_score_glb[tmp_mask])

    return planeIns.long()


# Copyright (c) 2020, Soohwan Kim. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np

from speech_transformer.modules import Linear
from torch import Tensor
from typing import Optional, Tuple

class ScaledDotProductAttention(nn.Module):
    """
    Scaled Dot-Product Attention proposed in "Attention Is All You Need"
    Compute the dot products of the query with all keys, divide each by sqrt(dim),
    and apply a softmax function to obtain the weights on the values

    Args: dim, mask
        dim (int): dimension of attention
        mask (torch.Tensor): tensor containing indices to be masked

    Inputs: query, key, value, mask
        - **query** (batch, q_len, d_model): tensor containing projection vector for decoder.
        - **key** (batch, k_len, d_model): tensor containing projection vector for encoder.
        - **value** (batch, v_len, d_model): tensor containing features of the encoded input sequence.
        - **mask** (-): tensor containing indices to be masked

    Returns: context, attn
        - **context**: tensor containing the context vector from attention mechanism.
        - **attn**: tensor containing the attention (alignment) from the encoder outputs.
    """
    def __init__(self, dim: int) -> None:
        super(ScaledDotProductAttention, self).__init__()
        self.sqrt_dim = np.sqrt(dim)

    def forward(self, query: Tensor, key: Tensor, value: Tensor, mask: Optional[Tensor] = None) -> Tuple[Tensor, Tensor]:
        score = torch.bmm(query, key.transpose(1, 2)) / self.sqrt_dim

        if mask is not None:
            score.masked_fill_(mask, -1e9)

        attn = F.softmax(score, -1)
        context = torch.bmm(attn, value)
        return context, attn

class MultiHeadAttention(nn.Module):
    """
    Multi-Head Attention proposed in "Attention Is All You Need"
    Instead of performing a single attention function with d_model-dimensional keys, values, and queries,
    project the queries, keys and values h times with different, learned linear projections to d_head dimensions.
    These are concatenated and once again projected, resulting in the final values.
    Multi-head attention allows the model to jointly attend to information from different representation
    subspaces at different positions.

    MultiHead(Q, K, V) = Concat(head_1, ..., head_h) · W_o
        where head_i = Attention(Q · W_q, K · W_k, V · W_v)

    Args:
        d_model (int): The dimension of keys / values / quries (default: 512)
        num_heads (int): The number of attention heads. (default: 8)

    Inputs: query, key, value, mask
        - **query** (batch, q_len, d_model): tensor containing projection vector for decoder.
        - **key** (batch, k_len, d_model): tensor containing projection vector for encoder.
        - **value** (batch, v_len, d_model): tensor containing features of the encoded input sequence.
        - **mask** (-): tensor containing indices to be masked

    Returns: output, attn
        - **output** (batch, output_len, dimensions): tensor containing the attended output features.
        - **attn** (batch * num_heads, v_len): tensor containing the attention (alignment) from the encoder outputs.
    """
    def __init__(self, d_model: int = 512, num_heads: int = 8) -> None:
        super(MultiHeadAttention, self).__init__()

        assert d_model % num_heads == 0, "hidden_dim % num_heads should be zero."

        self.d_head = int(d_model / num_heads)
        self.num_heads = num_heads
        self.query_proj = Linear(d_model, self.d_head * num_heads)
        self.key_proj = Linear(d_model, self.d_head * num_heads)
        self.value_proj = Linear(d_model, self.d_head * num_heads)
        self.sqrt_dim = np.sqrt(d_model)
        self.scaled_dot_attn = ScaledDotProductAttention(self.d_head)

    def forward(self, query: Tensor, key: Tensor, value: Tensor, mask: Optional[Tensor] = None) -> Tuple[Tensor, Tensor]:
        batch_size = value.size(0)

        query = self.query_proj(query).view(batch_size, -1, self.num_heads, self.d_head)  # BxQ_LENxNxD
        key = self.key_proj(key).view(batch_size, -1, self.num_heads, self.d_head)        # BxK_LENxNxD
        value = self.value_proj(value).view(batch_size, -1, self.num_heads, self.d_head)  # BxV_LENxNxD

        query = query.permute(2, 0, 1, 3).contiguous().view(batch_size * self.num_heads, -1, self.d_head)  # BNxQ_LENxD
        key = key.permute(2, 0, 1, 3).contiguous().view(batch_size * self.num_heads, -1, self.d_head)      # BNxK_LENxD
        value = value.permute(2, 0, 1, 3).contiguous().view(batch_size * self.num_heads, -1, self.d_head)  # BNxV_LENxD

        if mask is not None:
            mask = mask.repeat(self.num_heads, 1, 1)

        context, attn = self.scaled_dot_attn(query, key, value, mask)
        context = context.view(self.num_heads, batch_size, -1, self.d_head)
        context = context.permute(1, 2, 0, 3).contiguous().view(batch_size, -1, self.num_heads * self.d_head)  # BxTxND

        return context, attn


import numpy as np
import pandas as pd
import networkx as nx

def m_time(pos1,pos2): #计算rgv在两点间运动的最短时间
    d=abs(pos1-pos2)
    if d==1:
        return m1
    elif d==2:
        return m2
    elif d==3:
        return m3
    elif d==0:
        return 0
#m1,m2,m3,p1,p21,p22,to,te,w=20,33,46,560,400,378,28,31,25
#m1,m2,m3,p1,p21,p22,to,te,w=23,41,59,580,280,500,30,35,30
m1,m2,m3,p1,p21,p22,to,te,w=18,32,46,545,455,182,27,32,25
t=0
v_pos = 1  # 初始位置为1，右侧依次为2、3、4
v_sit = 0  # 0停止，1移动，2上料，3下料、上料、清洗
rgv=np.array([[t,v_pos,v_sit]])

c_sit=0 #表示cnc的作业状态，0为空闲且无孰料，1为空闲且有孰料，2为工作中
cnc=[]
cnc_inf=[]
c_pos=0 #表示初始位置，无实际意义
for num in range(1,9):
    if num==1 or num==2:
        c_pos=1
    elif num==3 or num==4:
        c_pos=2
    elif num==5 or num==6:
        c_pos=3
    elif num==7 or num==8:
        c_pos=4
    if num==1 or num==3 or num==5 or num==7:
        tc=to
    else:
        tc=te
    cnc.append(np.array([[t, num,c_pos,c_sit]]))
    cnc_inf.append(np.array([num,p1,p21,p22,tc])) #储存了cnc的加工、上下料时间，是固定的信息

judge_cnc=np.zeros((1,8),dtype=int) #初始cnc加工→休息判断矩阵
judge_move=0
judge_load=np.zeros((1,8),dtype=int)
judge_layoff=np.zeros((1,8),dtype=int)
load_num=0
layoff_num=0
result=np.array([[0,0,0,0]])

print("stop")

for t in range(1,28800): #从第1秒到第5秒
    #print(t)
    #################################
    rgv_new =np.array([t, rgv[t-1,1], rgv[t-1,2]])
    rgv=np.vstack((rgv, rgv_new))#t改变，其他rgv参数复制上一秒的信息
    for num in range(1, 9):
        cnc_new=np.array([t, cnc[num-1][t-1,1],cnc[num-1][t-1,2],cnc[num-1][t-1,3]])
        cnc[num-1]=np.vstack((cnc[num-1], cnc_new))#t改变，其他cnc参数复制上一秒的信息
    #################################
    for num in range(1, 9):
        if cnc[num-1][t,3]==2 and t>=judge_cnc[0,num-1] and t<judge_cnc[0,num-1]+1: #cnc加工完成，给rgv发信号
            cnc[num - 1][t, 3] = 1   #cnc从2（工作状态）到1（空闲，上有孰料状态）
    #################################
    if rgv[t,2]>0:    #如果rgv上一秒在运动
        if t>=judge_move and t <judge_move+1:     #判断rgb移动是否完成
            rgv[t,2]=0     #rgv瞬间进入移动后的等待状态，按照算法，这里应当立即对cnc进行操作
            for num in range(1, 9):
                if cnc[num - 1][t, 3] == 1:  # 如果cnc是空闲,上有孰料状态
                    rgv[t, 2] = 3
                    judge_layoff[0, num - 1] = t + cnc_inf[num - 1][4] + w
                    load_num = load_num + 1
                    layoff_num = layoff_num + 1
                    print("第%d个CNC在第%d秒开始下第%d个料" % (num, t, layoff_num))
                    print("第%d个CNC在第%d秒开始上第%d个料" % (num, t + cnc_inf[num - 1][4] / 2, load_num))
                    result_new = np.array([load_num, num, t + cnc_inf[num - 1][4] / 2, 0])
                    result = np.vstack((result, result_new))
                    result[layoff_num, 3] = t
                    break  # 一个时间rgv只能进入一种状态，这里进入了下料、上料、清洗状态
                if cnc[num - 1][t, 3] == 0:  # 如果cnc也是空闲，无料状态（初始状态）
                    rgv[t, 2] = 2  # rgv进入上料状态
                    judge_load[0, num - 1] = t + cnc_inf[num - 1][4] / 2  # 假设空cnc上料时间为上下料时间的一半
                    load_num = load_num + 1
                    print("第%d个CNC在第%d秒开始上第%d个料" % (num, t, load_num))
                    result_new = np.array([load_num, num, t, 0])
                    result = np.vstack((result, result_new))
                    break  # 一个时间rgv只能进入一种状态，这里进入了上料状态
        for num in range(1, 9):
            if t>=judge_load[0,num-1] and t<judge_load[0,num-1]+1:       #这个情况只有在最开始的时候情况
                cnc[num-1][t,3]=2   #cnc从0（空闲，上无料状态）到2（工作状态）
                judge_cnc[0, num - 1] = t + cnc_inf[num - 1][1]  #当前时间加cnc工作时间，为cnc的预测空闲时间
                rgv[t,2]=0       #rgv瞬间进入等待状态
            if t>=judge_layoff[0,num-1] and t<judge_layoff[0,num-1]+1:     #上下料同时进行
                cnc[num-1][t,3]=2 #cnc从1（空闲，上有孰料状态）到2（工作状态）
                judge_cnc[0, num - 1] = t + cnc_inf[num - 1][1]  #当前时间加cnc工作时间，为cnc的预测空闲时间
                rgv[t,2] = 0  # rgv瞬间进入等待状态,因为我们判断下料后一定会上料，这里要不要直接跳过搜索空闲cnc的过程？
###################################################################
    if rgv[t,2]==0:  #如果rgv空闲
            cnc_stop_num=[]
            for num in range(1, 9):
                if cnc[num-1][t,3]==0 or cnc[num-1][t,3]==1:   #如果cnc也是空闲状态
                    cnc_stop_num.append(cnc[num-1][t,1])   #调出空闲状态的cnc机号
##########################################################################################################################################################
            if cnc_stop_num == []:
                continue  # 跳出大循环，表示rgv处于等候状态
            elif len(cnc_stop_num)==1:  #如果只有一个机器发指令，直接接受指令
                mov_pro=cnc_stop_num[0]
            else:
                G = nx.MultiDiGraph()
                for s_num in cnc_stop_num:##########开始遍历第一层
                    t_1=t
                    G.add_edge(0,cnc[s_num - 1][t_1, 1], key=0, weight=m_time(cnc[s_num - 1][t_1, 2], rgv[t_1,1]) + cnc_inf[s_num - 1][4]+w)
                    rgv_1 = np.array([[t_1, rgv[t_1,1], rgv[t_1,2]]])
                    cnc_1 = []
                    for num in range(1, 9):
                        if num == 1 or num == 2:
                            c_pos = 1
                        elif num == 3 or num == 4:
                            c_pos = 2
                        elif num == 5 or num == 6:
                            c_pos = 3
                        elif num == 7 or num == 8:
                            c_pos = 4
                        cnc_1.append(np.array([[t_1, num, cnc[num-1][t_1,2], cnc[num-1][t_1,3]]]))
                    judge_cnc_1 = np.zeros((1, 8), dtype=int)  # 初始cnc加工→休息判断矩阵
                    judge_move_1 = 0
                    judge_load_1 = np.zeros((1, 8), dtype=int)
                    judge_layoff_1 = np.zeros((1, 8), dtype=int)
                    load_num_1 = 0
                    layoff_num_1 = 0
                    #################################################
                    mov_pro_1 = s_num
                    mov_pos_1 = 0
                    if mov_pro_1 == 1 or mov_pro_1 == 2:
                        mov_pos_1 = 1
                    elif mov_pro_1 == 3 or mov_pro_1 == 4:
                        mov_pos_1 = 2
                    elif mov_pro_1 == 5 or mov_pro_1 == 6:
                        mov_pos_1 = 3
                    elif mov_pro_1 == 7 or mov_pro_1 == 8:
                        mov_pos_1 = 4  # 计算目标机位（mov_pro 1-8）所在的位置（mov_pos 1-4）
                    if mov_pos_1 != rgv_1[0, 1]:  # 如果判断的移动位置和rgv所在位置不同，则立即进入移动状态
                        rgv_1[0, 2] = 1  # 0停止，1移动，2上料，3下料、上料、清洗
                        judge_move_1 = t_1 + m_time(mov_pos_1, rgv_1[0, 1])
                        rgv_1[0, 1] = mov_pos_1  # 移动开始，未知参数就变为目标位置
                    elif mov_pos_1 == rgv_1[0, 1]:  # 如果判断的移动位置和rgv所在位置相同，则立即进入 上料 或者 下料、上料、清洗 状态
                        for num in range(1, 9):
                            if cnc_1[num - 1][0, 3] == 1:  # 如果cnc是空闲,上有孰料状态
                                rgv_1[0, 2] = 3
                                judge_layoff_1[0, num - 1] = t_1 + cnc_inf[num - 1][4] + w
                                load_num_1 = load_num_1 + 1
                                layoff_num_1 = layoff_num_1 + 1
                                break # 一个时间rgv只能进入一种状态，这里进入了下料、上料、清洗状态
                            if cnc_1[num - 1][0, 3] == 0:  # 如果cnc也是空闲，无料状态（初始状态）
                                rgv_1[0, 2] = 2  # rgv进入上料状态
                                judge_load_1[0, num - 1] = t_1 + cnc_inf[num - 1][4] / 2  # 假设空cnc上料时间为上下料时间的一半
                                load_num_1 = load_num_1 + 1
                                break# 一个时间rgv只能进入一种状态，这里进入了上料状态
                    a_1 = 0
                    ###########################################################################
                    while True:
                        t_1=t_1+1
                        a_1 = a_1 + 1
                        #################################
                        rgv_new_1 = np.array([t_1, rgv_1[a_1 - 1, 1], rgv_1[a_1 - 1, 2]])
                        rgv_1 = np.vstack((rgv_1, rgv_new_1))  # t改变，其他rgv参数复制上一秒的信息
                        for num in range(1, 9):
                            cnc_new_1 = np.array([t_1, cnc_1[num - 1][a_1 - 1, 1], cnc_1[num - 1][a_1 - 1, 2], cnc_1[num - 1][a_1 - 1, 3]])
                            cnc_1[num - 1] = np.vstack((cnc_1[num - 1], cnc_new_1))  # t改变，其他cnc参数复制上一秒的信息
                        #################################
                        for num in range(1, 9):
                            if cnc_1[num - 1][a_1, 3] == 2 and t_1 >= judge_cnc_1[0, num - 1] and t_1 < judge_cnc_1[0, num - 1] + 1:  # cnc加工完成，给rgv发信号
                                cnc_1[num - 1][a_1, 3] = 1  # cnc从2（工作状态）到1（空闲，上有孰料状态）
                        #################################
                        if rgv_1[a_1, 2] > 0:  # 如果rgv上一秒在运动
                            if t_1 >= judge_move_1 and t_1 < judge_move_1 + 1:  # 判断rgb移动是否完成
                                rgv_1[a_1, 2] = 0  # rgv瞬间进入移动后的等待状态，按照算法，这里应当立即对cnc进行操作
                                for num in range(1, 9):
                                    if cnc_1[num - 1][a_1, 3] == 1:  # 如果cnc是空闲,上有孰料状态
                                        rgv_1[a_1, 2] = 3
                                        judge_layoff_1[0, num - 1] = t_1 + cnc_inf[num - 1][4] + w
                                        load_num_1 = load_num_1 + 1
                                        layoff_num_1 = layoff_num_1 + 1
                                        break  # 一个时间rgv只能进入一种状态，这里进入了下料、上料、清洗状态
                                    if cnc_1[num - 1][a_1, 3] == 0:  # 如果cnc也是空闲，无料状态（初始状态）
                                        rgv_1[a_1, 2] = 2  # rgv进入上料状态
                                        judge_load_1[0, num - 1] = t_1 + cnc_inf[num - 1][4] / 2  # 假设空cnc上料时间为上下料时间的一半
                                        load_num_1 = load_num_1 + 1
                                        break  # 一个时间rgv只能进入一种状态，这里进入了上料状态
                            for num in range(1, 9):
                                if t_1 >= judge_load_1[0, num - 1] and t_1 < judge_load_1[0, num - 1] + 1:  # 这个情况只有在最开始的时候情况
                                    cnc_1[num - 1][a_1, 3] = 2  # cnc从0（空闲，上无料状态）到2（工作状态）
                                    judge_cnc_1[0, num - 1] = t_1 + cnc_inf[num - 1][1]  # 当前时间加cnc工作时间，为cnc的预测空闲时间
                                    rgv_1[a_1, 2] = 0  # rgv瞬间进入等待状态
                                if t_1 >= judge_layoff_1[0, num - 1] and t_1 < judge_layoff_1[0, num - 1] + 1:  # 上下料同时进行
                                    cnc_1[num - 1][a_1, 3] = 2  # cnc从1（空闲，上有孰料状态）到2（工作状态）
                                    judge_cnc_1[0, num - 1] = t_1 + cnc_inf[num - 1][1]  # 当前时间加cnc工作时间，为cnc的预测空闲时间
                                    rgv_1[a_1, 2] = 0  # rgv瞬间进入等待状态,因为我们判断下料后一定会上料，这里要不要直接跳过搜索空闲cnc的过程？
                        ###################################################################
                        if rgv_1[a_1, 2] == 0:  # 如果rgv空闲 #直到运行到某时刻空闲，且有分叉，while才会停止
                                cnc_stop_num_1 = []
                                for num in range(1, 9):
                                    if cnc_1[num - 1][a_1, 3] == 0 or cnc_1[num - 1][a_1, 3] == 1:  # 如果cnc也是空闲状态
                                            cnc_stop_num_1.append(cnc_1[num - 1][a_1, 1])  # 调出空闲状态的cnc机号
                                    ##################################################################################################################################\
                                if cnc_stop_num_1 == []:
                                    continue  # 跳出大循环，表示rgv处于等候状态
                                else:
                                    for s_num_1 in cnc_stop_num_1:  ##########开始遍历第二层
                                        t_2=t_1
                                        G.add_edge(cnc[s_num - 1][t, 1], cnc_1[s_num_1 - 1][a_1, 1], key=1,weight=m_time(cnc_1[s_num_1 - 1][a_1, 2], rgv_1[a_1, 1]) + cnc_inf[s_num_1 - 1][4] + w)
                                    break
                GA_weight = []
                GA_choice = []
                for edge_0 in G.edges:
                    if edge_0[2] == 0:  # 遍历所有边，取key为0的边：edge_0
                        for edge_1 in G.edges:
                            if edge_1[2] == 1 and edge_0[1] == edge_1[0] and edge_1[0]!=edge_1[1]:  # 遍历所有边，取key为1的边,且首尾相连的边,且非环的边（两个节点重复）：edge_1
                                GA_weight.append(
                                    G[edge_0[0]][edge_0[1]][edge_0[2]]['weight'] + G[edge_1[0]][edge_1[1]][edge_1[2]]['weight'])
                                GA_choice.append(edge_0[1])  # 我们只关心下一步的点是哪个,这个点在1-8中
                mov_pro=GA_choice[GA_weight.index(min(GA_weight))]
                ##########结束遍历第一层
                # 获得较好的空闲状态机号mov_pro，属于(1-8)

            #返回一个1-4的数，是rgv到cnc距离和上料时间和最短的位置，进行移动和上料
            mov_pos=0
            if  mov_pro==1 or mov_pro==2:
                mov_pos=1
            elif mov_pro==3 or mov_pro==4:
                mov_pos=2
            elif mov_pro==5 or mov_pro==6:
                mov_pos=3
            elif mov_pro==7 or mov_pro==8:
                mov_pos=4             #计算目标机位（mov_pro 1-8）所在的位置（mov_pos 1-4）
            if mov_pos!=rgv[t,1]:  #如果判断的移动位置和rgv所在位置不同，则立即进入移动状态
                rgv[t,2] = 1      # 0停止，1移动，2上料，3下料、上料、清洗
                judge_move=t+m_time(mov_pos,rgv[t,1])
                rgv[t, 1] = mov_pos  # 移动开始，未知参数就变为目标位置
            elif mov_pos==rgv[t,1]:#如果判断的移动位置和rgv所在位置相同，则立即进入 上料 或者 下料、上料、清洗 状态
                for num in range(1, 9):
                    if cnc[num - 1][t,3] == 1:  # 如果cnc是空闲,上有孰料状态
                        rgv[t, 2] = 3
                        judge_layoff[0, num - 1] = t + cnc_inf[num - 1][4]+w
                        load_num = load_num + 1
                        layoff_num=layoff_num+1
                        print("第%d个CNC在第%d秒开始下第%d个料" % (num, t, layoff_num))
                        print("第%d个CNC在第%d秒开始上第%d个料" % (num, t+cnc_inf[num-1][4]/2, load_num))
                        result_new=np.array([load_num,num,t+cnc_inf[num-1][4]/2,0])
                        result = np.vstack((result, result_new))
                        result[layoff_num,3]=t

                        break  # 一个时间rgv只能进入一种状态，这里进入了下料、上料、清洗状态
                    if cnc[num - 1][t,3] == 0:  # 如果cnc也是空闲，无料状态（初始状态）
                        rgv[t, 2] = 2  #rgv进入上料状态
                        judge_load[0,num-1]=t+cnc_inf[num-1][4]/2  #假设空cnc上料时间为上下料时间的一半
                        load_num=load_num+1
                        print("第%d个CNC在第%d秒开始上第%d个料"%(num,t,load_num))
                        result_new = np.array([load_num, num, t , 0])
                        result = np.vstack((result, result_new))
                        break #一个时间rgv只能进入一种状态，这里进入了上料状态

result_df = pd.DataFrame(result)
writer = pd.ExcelWriter('test.xlsx')
result_df.to_excel(writer,'page_1',float_format='%.5f') # float_format 控制精度
writer.save()


# -*- coding: UTF-8 -*-
__author__ = 'Jens-Kristian Krogager'

import numpy as np
import VoigtFit

def print_T_model_pars(dataset, filename=None):
    """Print the turbulence and T parameters for physical model."""
    N_comp = len(dataset.components.values()[0])
    print("")
    print(u"  No:     Temperature [K]       Turbulence [km/s]")
    if filename:
        out_file = open(filename, 'w')
        out_file.write(u"# No:     Temperature [K]       Turbulence [km/s] \n")

    for comp_num in range(N_comp):
        T_name = 'T_%i' % comp_num
        turb_name = 'turb_%i' % comp_num
        T_fit = dataset.best_fit[T_name]
        turb_fit = dataset.best_fit[turb_name]
        par_tuple = (comp_num, T_fit.value, T_fit.stderr,
                     turb_fit.value, turb_fit.stderr)
        print(u"  %-3i   %.2e ± %.2e    %.2e ± %.2e" % par_tuple)
        if filename:
            out_file.write(u"  %-3i   %.2e ± %.2e    %.2e ± %.2e \n" % par_tuple)

    print("")
    if filename:
        out_file.close()

# -- Fit DLA towards quasar Q1313+1441
#    Observed in X-shooter P089.A-0068

z_DLA = 0.00345

# If log(NHI) is not known use:
logNHI = None

# -- Load X-SHOOTER UVB and VIS data in ASCII format:
fname = 'thermal_model_2comp.dat'
res = 6.6

wl, spec, err = np.loadtxt(fname, unpack=True)

# -- Here you can load your data in any way you wish
#    Only requirement is that wl, spec, and err have the same dimensions.

# -- A dataset which has already been defined can be loaded like this:
# dataset = VoigtFit.LoadDataSet('test_data.hdf5')

dataset = VoigtFit.DataSet(z_DLA)
dataset.set_name('test_2comp')
dataset.verbose = True
dataset.velspan = 150.
dataset.cheb_order = -1

# -- Add the data loaded from the
dataset.add_data(wl, spec, res, err=err, normalized=True)

# -- Define absorption lines:
dataset.add_many_lines(['FeII_2344', 'FeII_2374', 'FeII_2382'])
dataset.add_many_lines(['FeII_1608', 'FeII_1611'])
dataset.add_line('FeII_2260')
dataset.add_line('FeII_2249')
dataset.add_line('CrII_2056')
dataset.add_line('CrII_2066')
dataset.add_line('CrII_2026')
dataset.add_line('ZnII_2026')
dataset.add_line('CrII_2062')
dataset.add_line('ZnII_2062')
# -- dataset.add_many_lines is equivalent to dataset.add_lines:
dataset.add_many_lines(['CII_1036', 'CII_1334'])
dataset.add_many_lines(['OI_1302', 'OI_1039', 'OI_1355'])
dataset.add_many_lines(['SiII_1526', 'SiII_1808', 'SiII_1304'])
dataset.add_many_lines(['SiII_1260', 'FeII_1260', 'SII_1259'])
dataset.add_many_lines(['SII_1250', 'SII_1253'])

# -- If a line has been defined, and you don't want to fit it
#    it can either be removed from the dataset completely:
# dataset.remove_line('CrII_2056')

# -- or deactivated:
# dataset.deactivate_line('FeII_2374')

# -- Deactivated lines will not be included in the fit, but their line definitions
#    and components remain in the dataset for future reference.

# -- To use the physical model, make sure that all components are cleared:
dataset.reset_components()

# -- Add components for each ion:
#                      ion    z         b   logN
dataset.add_component('FeII', 0.003290, 5., 15.0, var_z=1, var_b=1)
dataset.add_component('FeII', 0.003620, 5., 14.5, var_z=1, var_b=1)

# -- The physical model requires that all ions have the same velocity structure:
#    The default order is 'to' , 'from' :
# dataset.copy_components('CrII', 'FeII')
# -- But the ions can be specified using keywords to ease the call:
dataset.copy_components(from_ion='FeII', to_ion='CrII', tie_b=False)
dataset.copy_components(from_ion='FeII', to_ion='ZnII', tie_b=False)
dataset.copy_components(from_ion='FeII', to_ion='SiII', tie_b=False)
dataset.copy_components(from_ion='FeII', to_ion='SII', tie_b=False)
dataset.copy_components(from_ion='FeII', to_ion='CII', tie_b=False)
dataset.copy_components(from_ion='FeII', to_ion='OI', tie_b=False)

# -- This copies the two components defined for FeII to the other ions and
#    keeps the same pattern of initial guesses for column density scaled
#    to the Solar abundance ratio.

# -- Individual components which are not observed for weaker lines can be removed:
# dataset.delete_component('ZnII', 1)
# dataset.delete_component('ZnII', 0)
#
# NOTE - components should be deleted from last component to first component
#        not the other way around as that messes up the component numbering.
#        Components are zero-indexed!

# -- Prepare the dataset: This will prompt the user for interactive
#    masking and normalization, as well as initiating the Parameters:
dataset.prepare_dataset(norm=False, mask=False)

# -- Define masks for individual lines:
# dataset.mask_line('ZnII_2026')

# --- This is where the magic happens ----------------------------------------
# Set up the thermal and turbulence parameters for each component:
dataset.pars.add('turb_0', value=5., vary=True, min=0.)
dataset.pars.add('turb_1', value=5., vary=True, min=0.)

dataset.pars.add('T_0', value=5000., vary=1, min=0.)
dataset.pars.add('T_1', value=5000., vary=1, min=0.)
# -- This can be defined in a for loop assuming the same intial guess for T:
# T_init = 1.e4
# for comp_num in range(len(dataset.components.values()[0])):
#     dataset.pars.add('T_%i'%comp_num, value=T_init, vary=True, min=0.)

# -- Now set up the links for the 'b'-parameter of each component of each ion:
# 2k_B/m_u in (km/s)^2 units:
K = 0.0166287
for ion, comp in dataset.components.items():
    N_comp = len(comp)
    for comp_num in range(N_comp):
        par_name = 'b%i_%s' % (comp_num, ion)
        lines_for_ion = dataset.get_lines_for_ion(ion)
        m_ion = lines_for_ion[0].mass
        const = K/m_ion
        T_num = dataset.pars['T_%i' % comp_num].value
        b_eff = np.sqrt(5.**2 + K*T_num/m_ion)
        model_constraint = 'sqrt((turb_%i)**2 + %.6f*T_%i)' % (comp_num,
                                                               const,
                                                               comp_num)
        dataset.pars[par_name].set(expr=model_constraint, value=b_eff)

# ---------------------------------------------------------------------------

# -- Fit the dataset:
popt, chi2 = dataset.fit(verbose=True, plot=False, factor=10.)

dataset.plot_fit(filename=dataset.name)

# -- Print total column densities
dataset.print_total()

if logNHI:
    dataset.print_metallicity(*logNHI)

print_T_model_pars(dataset)

# -- Save the dataset to file: taken from the dataset.name
dataset.save()


from functools import reduce
import jax.numpy as np
import jax.linear_util as lu
from jax.util import unzip2, unzip3, safe_zip, safe_map, partial, WrapHashably
from jax.abstract_arrays import ShapedArray
from jax.experimental import stax
from jax.interpreters import partial_eval as pe
from jax.interpreters.batching import get_aval
from jax.interpreters import batching
from jax.api_util import (
    wraps, pytree_to_jaxtupletree, pytree_fun_to_jaxtupletree_fun)
import jax.core as jc
from jax import random
from jax import lax

zip = safe_zip
map = safe_map

def merge_params(params):
    if len(params) > 0:
        p = params[0]
        for param in params[1:]:
            p.update(param)
        return p
    else:
        return {}

class Layer(jc.Primitive):
    def __init__(self, name, init_fun, apply_fun):
        self.init_fun = init_fun
        self.apply_fun = apply_fun
        super(Layer, self).__init__(name)
        def layer_abstract_eval(*avals):
            akey = ShapedArray((2,), 'uint32')
            def init_and_apply(key, *inputs):
                params = init_fun(key, *inputs)
                return apply_fun(params, *inputs)
            return pe.abstract_eval_fun(init_and_apply, akey, *avals)
        self.def_abstract_eval(layer_abstract_eval)
        def layer_batch(batched_args, batch_dims, **params):
            # TODO: figure out batching/debatching for init_fun
            batched_apply_fun = (
                lambda params, *batch_inputs:
                batching.batch(lu.wrap_init(partial(self.apply_fun, params)),
                               batch_inputs, batch_dims, 0))
            batched_layer = Layer(name, init_fun, batched_apply_fun)
            return batched_layer.bind(*batched_args, **params), 0
        batching.primitive_batchers[self] = layer_batch

def init_interpreter(rng, jaxpr, consts, freevar_vals, net_params, *args):
    def read(v):
        if type(v) is jc.Literal:
            return v.val
        else:
            return env[v]

    def write(v, val):
        env[v] = val

    env = {}
    write(jc.unitvar, jc.unit)
    jc.pat_fmap(write, jaxpr.constvars, consts)
    jc.pat_fmap(write, jaxpr.invars, args)
    jc.pat_fmap(write, jaxpr.freevars, freevar_vals)
    for eqn in jaxpr.eqns:
        rng, prim_rng = random.split(rng)
        if not eqn.restructure:
            in_vals = map(read, eqn.invars)
        else:
            in_vals = [pack(map(read, invars)) if type(invars) is tuple
                       else read(invars) for invars in eqn.invars]
        if eqn.bound_subjaxprs:
            subjaxprs, sub_consts, sub_freevar_vals = unzip3([
                (subjaxpr,
                 map(read, const_vars),
                 map(read, bound_vars))
                for subjaxpr, const_vars, bound_vars in eqn.bound_subjaxprs])
            ans, net_params = get_primitive_init(eqn.primitive)(
                prim_rng, eqn.params, sub_consts, sub_freevar_vals, in_vals,
                net_params)
        else:
            ans, net_params = get_primitive_init(eqn.primitive)(
                prim_rng, net_params, *in_vals, **eqn.params)
        outvals = list(ans) if eqn.destructure else [ans]
        map(write, eqn.outvars, outvals)
    return net_params

init_rules = {}

def layer_init(layer, rng, net_params, *inputs):
    if layer.name not in net_params:
        layer_params = layer.init_fun(rng, *inputs)
        net_params[layer.name] = layer_params
    return layer.apply_fun(net_params[layer.name], *inputs), net_params

def get_primitive_init(primitive):
    if primitive in init_rules:
        return primitive
    elif isinstance(primitive, Layer):
        return partial(layer_init, primitive)
    else:
        return (lambda _, net_params, *in_vals, **params:
                (primitive.bind(*in_vals, **params), net_params))

def init_fun(net_fun, rng, *example_inputs, **kwargs):
    net_fun = lu.wrap_init(net_fun)
    def pv_like(x):
        return pe.PartialVal((get_aval(x), jc.unit))
    pvals = map(pv_like, example_inputs)
    jaxpr, _, consts = pe.trace_to_jaxpr(net_fun, pvals, **kwargs)
    return init_interpreter(rng, jaxpr, consts, [], {}, *example_inputs)

class ApplyTracer(jc.Tracer):
    __slots__ = ['val', 'net_params']

    def __init__(self, trace, net_params, val):
        self.trace = trace
        self.val = val
        self.net_params = net_params

    @property
    def aval(self):
        return jc.get_aval(self.val)

    def unpack(self):
        return tuple(self.val)

    def full_lower(self):
        return self

class ApplyTrace(jc.Trace):
    def pure(self, val):
        return ApplyTracer(self, {}, val)

    def lift(self, val):
        return ApplyTracer(self, {}, val)

    def sublift(self, val):
        return ApplyTracer(self, {}, val.val)

    def process_primitive(self, primitive, tracers, params):
        vals_in, net_params = unzip2((t.val, t.net_params) for t in tracers)
        net_params = merge_params(net_params)
        if isinstance(primitive, Layer):
            apply_fun = primitive.apply_fun
            layer_params = net_params[primitive.name]
            return ApplyTracer(
                self, net_params, apply_fun(layer_params, *vals_in))
        else:
            return ApplyTracer(
                self, net_params, primitive.bind(*vals_in, **params))

    def process_call(self, call_primitive, f, tracers, params):
        if call_primitive in pe.map_primitives:
            raise NotImplementedError
        vals, net_params = unzip2((t.val, t.net_params) for t in tracers)
        if any(net_params):
            net_params = merge_params(net_params)
            f = apply_subtrace(f, self.master, WrapHashably(net_params))
            val_out = call_primitive.bind(f, *vals, **params)
            return ApplyTracer(self, net_params, val_out)
        else:
            return call_primitive.bind(f, *vals, **params)

@lu.transformation
def apply_transform(net_params, inputs):
    with jc.new_master(ApplyTrace) as master:
        trace = ApplyTrace(master, jc.cur_sublevel())
        ans = yield map(partial(ApplyTracer, trace, net_params), inputs), {}
        out_tracer = trace.full_raise(ans)
        out_val = out_tracer.val
        del master, out_tracer
    yield out_val

@lu.transformation
def apply_subtrace(master, net_params, *vals):
    net_params = net_params.val
    trace = ApplyTrace(master, jc.cur_sublevel())
    ans = yield map(partial(ApplyTracer, trace, net_params), vals), {}
    out_tracer = trace.full_raise(ans)
    yield out_tracer.val

def apply_fun(net_fun, params, *inputs):
    return apply_transform(lu.wrap_init(net_fun), params).call_wrapped(inputs)