import numpy as np
import torch
import torch.nn.functional as F
import argparse
import tqdm
import json
import cv2 as cv
import os, glob
import math


from render_utils.lib.utils.graphics_utils import focal2fov, getProjectionMatrix
from diff_gaussian_rasterization import GaussianRasterizationSettings, GaussianRasterizer


def render3(
    gaussian_vals: dict,
    bg_color: torch.Tensor,
    extr: torch.Tensor,
    intr: torch.Tensor,
    img_w: int,
    img_h: int,
    scaling_modifier = 1.0,
    override_color = None,
    compute_cov3D_python = False
):
    means3D = gaussian_vals['positions']
    # Create zero tensor. We will use it to make pytorch return gradients of the 2D (screen-space) means
    screenspace_points = torch.zeros_like(means3D, dtype = means3D.dtype, requires_grad = True, device = "cuda") + 0
    try:
        screenspace_points.retain_grad()
    except:
        pass
    means2D = screenspace_points
    opacity = gaussian_vals['opacity']

    # If precomputed 3d covariance is provided, use it. If not, then it will be computed from
    # scaling / rotation by the rasterizer.
    scales = None
    rotations = None
    cov3D_precomp = None
    scales = gaussian_vals['scales']
    rotations = gaussian_vals['rotations']

    # If precomputed colors are provided, use them. Otherwise, if it is desired to precompute colors
    # from SHs in Python, do it. If not, then SH -> RGB conversion will be done by rasterizer.
    shs = None
    # colors_precomp = None
    # if override_color is None:
    #     shs = gaussian_vals['shs']
    # else:
    #     colors_precomp = override_color
    if 'colors' in gaussian_vals:
        colors_precomp = gaussian_vals['colors']
    else:
        colors_precomp = None

    # Set up rasterization configuration
    FoVx = focal2fov(intr[0, 0].item(), img_w)
    FoVy = focal2fov(intr[1, 1].item(), img_h)
    tanfovx = math.tan(FoVx * 0.5)
    tanfovy = math.tan(FoVy * 0.5)
    world_view_transform = extr.transpose(1, 0).cuda()
    projection_matrix = getProjectionMatrix(znear = 0.1, zfar = 100, fovX = FoVx, fovY = FoVy, K = intr, img_w = img_w, img_h = img_h).transpose(0, 1).cuda()
    full_proj_transform = (world_view_transform.unsqueeze(0).bmm(projection_matrix.unsqueeze(0))).squeeze(0)
    camera_center = torch.linalg.inv(extr)[:3, 3]

    raster_settings = GaussianRasterizationSettings(
        image_height = img_h,
        image_width = img_w,
        tanfovx = tanfovx,
        tanfovy = tanfovy,
        bg = bg_color,
        scale_modifier = scaling_modifier,
        viewmatrix = world_view_transform,
        projmatrix = full_proj_transform,
        sh_degree = gaussian_vals['max_sh_degree'],
        campos = camera_center,
        prefiltered = False,
        debug = False
    )

    rasterizer = GaussianRasterizer(raster_settings = raster_settings)

    # Rasterize visible Gaussians to image, obtain their radii (on screen).
    rendered_image, radii = rasterizer(
        means3D = means3D,
        means2D = means2D,
        shs = shs,
        colors_precomp = colors_precomp,
        opacities = opacity,
        scales = scales,
        rotations = rotations,
        cov3D_precomp = cov3D_precomp)

    # Those Gaussians that were frustum culled or had a radius of 0 were not visible.
    # They will be excluded from value updates used in the splitting criteria.
    return {
        "render": rendered_image,
        "viewspace_points": screenspace_points,
        "visibility_filter": radii > 0,
        "radii": radii
    }


def blend_color(head_facial_color, body_facial_color, blend_weight):
    blend_weight = blend_weight.reshape([len(blend_weight)] + [1]*(len(head_facial_color.shape)-1))
    result = head_facial_color * blend_weight + body_facial_color * (1-blend_weight)
    return result


@torch.no_grad()
def paste_back_with_linear_interp(pasteback_scale, pasteback_center, src, tgt_size):
    pasteback_topleft = [pasteback_center[0] - src.shape[1]/2/pasteback_scale, 
                         pasteback_center[1] - src.shape[0]/2/pasteback_scale]

    h, w = src.shape[0], src.shape[1]
    grayscale = False
    if len(src.shape) == 2:
        src = src.reshape([h, w, 1])
        grayscale = True
    src = torch.from_numpy(src)
    src = src.permute(2, 0, 1).unsqueeze(0)
    grid = torch.meshgrid(torch.arange(0, tgt_size[0]), torch.arange(0, tgt_size[1]), indexing='xy')
    grid = torch.stack(grid, dim = -1).float().to(src.device).unsqueeze(0)
    grid[..., 0] = (grid[..., 0] - pasteback_topleft[0]) * pasteback_scale
    grid[..., 1] = (grid[..., 1] - pasteback_topleft[1]) * pasteback_scale

    grid[..., 0] = grid[..., 0] / (src.shape[-1] / 2.0) - 1.0
    grid[..., 1] = grid[..., 1] / (src.shape[-2] / 2.0) - 1.0
    out = F.grid_sample(src, grid, align_corners = True)
    out = out[0].detach().permute(1, 2, 0).cpu().numpy()
    if grayscale:
        out = out[:, :, 0]
    return out


def soften_blending_mask(blending_mask, valid_mask):
    blending_mask = np.clip(blending_mask*2.0, 0.0, 1.0)
    blending_mask = cv.erode(blending_mask, np.ones((5, 5))) * valid_mask
    blending_mask_bk = np.copy(blending_mask)
    blending_mask = cv.blur(blending_mask*valid_mask, (25, 25))
    valid_mask = cv.blur(valid_mask, (25, 25))
    blending_mask = blending_mask / (valid_mask + 1e-6) * blending_mask_bk
    return blending_mask