#
# Copyright (C) 2023, Inria
# GRAPHDECO research group, https://team.inria.fr/graphdeco
# All rights reserved.
#
# This software is free for non-commercial, research and evaluation use 
# under the terms of the LICENSE.md file.
#
# For inquiries contact  george.drettakis@inria.fr
#

import torch
from scene import Scene
import os
import sys
from tqdm import tqdm
from os import makedirs
from scene import Scene, GaussianModel
from gaussian_renderer import render
import torchvision
from utils.general_utils import safe_state
from argparse import ArgumentParser
from arguments import ModelParams, PipelineParams, get_combined_args, Namespace
from utils.sh_utils import RGB2SH, SH2RGB
from predictor import FaceMeshPredictor
import torch.nn.functional as F
# os.environ["CUDA_VISIBLE_DEVICES"] = "3"


def rot_mat_from_6dof(v: torch.Tensor) -> torch.Tensor:
    assert v.shape[-1] == 6
    v = v.view(-1, 6)
    vx, vy = v[..., :3].clone(), v[..., 3:].clone()

    b1 = F.normalize(vx, dim=-1)
    b3 = F.normalize(torch.cross(b1, vy), dim=-1)
    b2 = - torch.cross(b1, b3)

    return torch.stack((b1, b2, b3), dim=-1)
def render_set(use_sh, model_path, name, iteration, views, gaussians, pipeline, background):
    render_path = os.path.join(model_path, name, "ours_{}".format(iteration), "renders")
    gts_path = os.path.join(model_path, name, "ours_{}".format(iteration), "gt")
    flame_predictor = FaceMeshPredictor.dad_3dnet()
    makedirs(render_path, exist_ok=True)
    makedirs(gts_path, exist_ok=True)
    verts_noaling = gaussians.get_xyz.unsqueeze(0)
    for idx, view in enumerate(tqdm(views, desc="Rendering progress")):
        gt = view.original_image[0:3, :, :]
        gt_input = gt.permute(1, 2, 0).cpu().numpy()[::2, ::2, :] * 255
        gt_parameters = flame_predictor(gt_input)
        rotation_mat = rot_mat_from_6dof(gt_parameters['flame_params'].rotation).type(verts_noaling.dtype).cuda()
        # vertices = torch.matmul(rotation_mat.unsqueeze(1), verts_noaling.unsqueeze(-1))
        # vertices = vertices[..., 0]
        # if idx == 0:
        #     scale = torch.clamp(gt_parameters['flame_params'].scale[:, None] + 1.0, 1e-8).cuda()
        scale = torch.clamp(gt_parameters['flame_params'].scale[:, None] + 1.0, 1e-8).cuda()
        gaussians._xyz = verts_noaling.squeeze(0) * scale.squeeze()
        translation = gt_parameters['flame_params'].translation.cuda()
        translation[..., 2] = 2183.5264
        view.mk_new_cam(R=rotation_mat.squeeze(), T=translation.squeeze())
        if not use_sh:
            rendering = RGB2SH(render(view, gaussians, pipeline, background)["render"]) # use rgb
        else:
            rendering = render(view, gaussians, pipeline, background)["render"] # use sh
        # print(torch.max(rendering))
        torchvision.utils.save_image(rendering, os.path.join(render_path, '{0:05d}'.format(idx) + ".png"))
        torchvision.utils.save_image(gt, os.path.join(gts_path, '{0:05d}'.format(idx) + ".png"))

def render_sets(dataset : ModelParams, iteration : int, pipeline : PipelineParams, skip_train : bool, skip_test : bool):
    with torch.no_grad():
        gaussians = GaussianModel(dataset.sh_degree, dataset.data_device, dataset.xyz_grad, dataset.align, dataset.free_flame)
        # print(gaussians._xyz_scale)
        scene = Scene(dataset, gaussians, load_iteration=iteration, shuffle=False)

        bg_color = [1,1,1] if dataset.white_background else [0, 0, 0]
        background = torch.tensor(bg_color, dtype=torch.float32, device=dataset.data_device)

        if not skip_train:
             render_set(dataset.use_sh, dataset.model_path, "train", scene.loaded_iter, scene.getTrainCameras(), gaussians, pipeline, background)

        if not skip_test:
             render_set(dataset.use_sh, dataset.model_path, "test", scene.loaded_iter, scene.getTestCameras(), gaussians, pipeline, background)

if __name__ == "__main__":
    # Set up command line argument parser
    parser = ArgumentParser(description="Testing script parameters")
    model = ModelParams(parser, sentinel=True)
    pipeline = PipelineParams(parser)
    parser.add_argument("--iteration", default=-1, type=int)
    parser.add_argument("--skip_train", action="store_true")
    parser.add_argument("--skip_test", action="store_true")
    parser.add_argument("--quiet", action="store_true")
    args = get_combined_args(parser)
    print("Rendering " + args.model_path)

    # Initialize system state (RNG)
    safe_state(args.quiet)

    render_sets(model.extract(args), args.iteration, pipeline.extract(args), args.skip_train, args.skip_test)