import torch
import torch.nn as nn
import render_util
import geo_transform
import numpy as np


def compute_tri_normal(geometry, tris):
    geometry = geometry.permute(0, 2, 1)
    tri_1 = tris[:, 0]
    tri_2 = tris[:, 1]
    tri_3 = tris[:, 2]

    vert_1 = torch.index_select(geometry, 2, tri_1)
    vert_2 = torch.index_select(geometry, 2, tri_2)
    vert_3 = torch.index_select(geometry, 2, tri_3)

    nnorm = torch.cross(vert_2 - vert_1, vert_3 - vert_1, 1)
    normal = nn.functional.normalize(nnorm).permute(0, 2, 1)
    return normal


class Compute_normal_base(torch.autograd.Function):
    @staticmethod
    def forward(ctx, normal):
        (normal_b,) = render_util.normal_base_forward(normal)
        ctx.save_for_backward(normal)
        return normal_b

    @staticmethod
    def backward(ctx, grad_normal_b):
        (normal,) = ctx.saved_tensors
        (grad_normal,) = render_util.normal_base_backward(grad_normal_b, normal)
        return grad_normal


class Normal_Base(torch.nn.Module):
    def __init__(self):
        super(Normal_Base, self).__init__()

    def forward(self, normal):
        return Compute_normal_base.apply(normal)


def preprocess_render(geometry, euler, trans, cam, tris, vert_tris, ori_img):
    point_num = geometry.shape[1]
    rott_geo = geo_transform.euler_trans_geo(geometry, euler, trans)
    proj_geo = geo_transform.proj_geo(rott_geo, cam)
    rot_tri_normal = compute_tri_normal(rott_geo, tris)
    rot_vert_normal = torch.index_select(rot_tri_normal, 1, vert_tris)
    is_visible = -torch.bmm(
        rot_vert_normal.reshape(-1, 1, 3),
        nn.functional.normalize(rott_geo.reshape(-1, 3, 1)),
    ).reshape(-1, point_num)
    is_visible[is_visible < 0.01] = -1
    pixel_valid = torch.zeros(
        (ori_img.shape[0], ori_img.shape[1] * ori_img.shape[2]),
        dtype=torch.float32,
        device=ori_img.device,
    )
    return rott_geo, proj_geo, rot_tri_normal, is_visible, pixel_valid


class Render_Face(torch.autograd.Function):
    @staticmethod
    def forward(
        ctx, proj_geo, texture, nbl, ori_img, is_visible, tri_inds, pixel_valid
    ):
        batch_size, h, w, _ = ori_img.shape
        ori_img = ori_img.view(batch_size, -1, 3)
        ori_size = torch.cat(
            (
                torch.ones((batch_size, 1), dtype=torch.int32, device=ori_img.device)
                * h,
                torch.ones((batch_size, 1), dtype=torch.int32, device=ori_img.device)
                * w,
            ),
            dim=1,
        ).view(-1)
        tri_index, tri_coord, render, real = render_util.render_face_forward(
            proj_geo, ori_img, ori_size, texture, nbl, is_visible, tri_inds, pixel_valid
        )
        ctx.save_for_backward(
            ori_img, ori_size, proj_geo, texture, nbl, tri_inds, tri_index, tri_coord
        )
        return render, real

    @staticmethod
    def backward(ctx, grad_render, grad_real):
        (
            ori_img,
            ori_size,
            proj_geo,
            texture,
            nbl,
            tri_inds,
            tri_index,
            tri_coord,
        ) = ctx.saved_tensors
        grad_proj_geo, grad_texture, grad_nbl = render_util.render_face_backward(
            grad_render,
            grad_real,
            ori_img,
            ori_size,
            proj_geo,
            texture,
            nbl,
            tri_inds,
            tri_index,
            tri_coord,
        )
        return grad_proj_geo, grad_texture, grad_nbl, None, None, None, None


class Render_RGB(nn.Module):
    def __init__(self):
        super(Render_RGB, self).__init__()

    def forward(
        self, proj_geo, texture, nbl, ori_img, is_visible, tri_inds, pixel_valid
    ):
        return Render_Face.apply(
            proj_geo, texture, nbl, ori_img, is_visible, tri_inds, pixel_valid
        )


def cal_land(proj_geo, is_visible, lands_info, land_num):
    (land_index,) = render_util.update_contour(lands_info, is_visible, land_num)
    proj_land = torch.index_select(proj_geo.reshape(-1, 3), 0, land_index)[
        :, :2
    ].reshape(-1, land_num, 2)
    return proj_land


class Render_Land(nn.Module):
    def __init__(self):
        super(Render_Land, self).__init__()
        lands_info = np.loadtxt("../data/3DMM/lands_info.txt", dtype=np.int32)
        self.lands_info = torch.as_tensor(lands_info).cuda()
        tris = np.loadtxt("../data/3DMM/tris.txt", dtype=np.int64)
        self.tris = torch.as_tensor(tris).cuda() - 1
        vert_tris = np.loadtxt("../data/3DMM/vert_tris.txt", dtype=np.int64)
        self.vert_tris = torch.as_tensor(vert_tris).cuda()
        self.normal_baser = Normal_Base().cuda()
        self.renderer = Render_RGB().cuda()

    def render_mesh(self, geometry, euler, trans, cam, ori_img, light):
        batch_size, h, w, _ = ori_img.shape
        ori_img = ori_img.view(batch_size, -1, 3)
        ori_size = torch.cat(
            (
                torch.ones((batch_size, 1), dtype=torch.int32, device=ori_img.device)
                * h,
                torch.ones((batch_size, 1), dtype=torch.int32, device=ori_img.device)
                * w,
            ),
            dim=1,
        ).view(-1)
        rott_geo, proj_geo, rot_tri_normal, _, _ = preprocess_render(
            geometry, euler, trans, cam, self.tris, self.vert_tris, ori_img
        )
        tri_nb = self.normal_baser(rot_tri_normal.contiguous())
        nbl = torch.bmm(
            tri_nb, (light.reshape(-1, 9, 3))[:, :, 0].unsqueeze(-1).repeat(1, 1, 3)
        )
        texture = torch.ones_like(geometry) * 200
        (render,) = render_util.render_mesh(
            proj_geo, ori_img, ori_size, texture, nbl, self.tris
        )
        return render.view(batch_size, h, w, 3).byte()

    def cal_loss_rgb(self, geometry, euler, trans, cam, ori_img, light, texture, lands):
        rott_geo, proj_geo, rot_tri_normal, is_visible, pixel_valid = preprocess_render(
            geometry, euler, trans, cam, self.tris, self.vert_tris, ori_img
        )
        tri_nb = self.normal_baser(rot_tri_normal.contiguous())
        nbl = torch.bmm(tri_nb, light.reshape(-1, 9, 3))
        render, real = self.renderer(
            proj_geo, texture, nbl, ori_img, is_visible, self.tris, pixel_valid
        )
        proj_land = cal_land(proj_geo, is_visible, self.lands_info, lands.shape[1])
        col_minus = torch.norm((render - real).reshape(-1, 3), dim=1).reshape(
            ori_img.shape[0], -1
        )
        col_dis = torch.mean(col_minus * pixel_valid) / (
            torch.mean(pixel_valid) + 0.00001
        )
        land_dists = torch.norm((proj_land - lands).reshape(-1, 2), dim=1).reshape(
            ori_img.shape[0], -1
        )
        lan_dis = torch.mean(land_dists)
        return col_dis, lan_dis