import os
import cv2
import shutil
import numpy as np
import torch


def build_fold(path):
    if os.path.exists(path):
        return True
    os.makedirs(path)
    return False
    

def visualize_frame_with_mask(grid0, grid1, grid2, a0_mask, a1_mask, a2_mask, point_coords, resolution, name, args=None):
    os.makedirs(name + 'frames_seg', exist_ok=True)
    a0_dir, a1_dir, a2_dir = name + 'frames_seg/x', name + 'frames_seg/y', name + 'frames_seg/z'
    # build_fold(a1_dir)
    # build_fold(a2_dir)
    
    a0_mask, a1_mask, a2_mask = a0_mask.repeat(1, 3, 1, 1), a1_mask.repeat(1, 3, 1, 1), a2_mask.repeat(1, 3, 1, 1)
    a0_mask[:, 1:], a1_mask[:, 1:], a2_mask[:, 1:] = a0_mask[:, 1:] * 0, a1_mask[:, 1:] * 0, a2_mask[:, 1:] * 0
    
    grid0, grid1, grid2 = grid0 * 0.7 + a0_mask * 0.3, grid1 * 0.7 + a1_mask * 0.3, grid2 * 0.7 + a2_mask * 0.3
    
    grid0[point_coords[0], :, point_coords[1], point_coords[2]] = torch.Tensor([0., 1., 0.])
    grid1[point_coords[1], :, point_coords[0], point_coords[2]] = torch.Tensor([0., 1., 0.])
    grid2[point_coords[2], :, point_coords[0], point_coords[1]] = torch.Tensor([0., 1., 0.])
    if not build_fold(a0_dir):
        visualize_per_frame(grid0, a0_dir, resolution, args)
    if not build_fold(a1_dir):
        visualize_per_frame(grid1, a1_dir, resolution, args)
    if not build_fold(a2_dir):
        visualize_per_frame(grid2, a2_dir, resolution, args)
    

def visualize_per_frame(grid, foldpath, resolution, args=None):
    grid = torch.nn.functional.interpolate(grid, size=(resolution, resolution), mode=args.mode)
    
    imgs = grid.cpu().numpy()
    n, _, _, _ = grid.shape
    for ii in range(n):
        r = np.uint8(imgs[ii, 0, :, :]*255)
        g = np.uint8(imgs[ii, 1, :, :]*255)
        b = np.uint8(imgs[ii, 2, :, :]*255)
        img = cv2.merge([b, g, r])
        
        cv2.imwrite('{}/{}.png'.format(foldpath, ii), img)
    return

def cal(input, points):
    reference_point_3d = np.array(input)  
    distances = np.linalg.norm(points - reference_point_3d, axis=1)
    closest_index = np.argmin(distances)
    closest_point = points[closest_index]
    return [closest_point[0], closest_point[1], closest_point[2]]