import torch
import torch.nn.functional as F
from tqdm import tqdm
import kaolin
import lpips
from einops import rearrange
from pytorch3d.transforms import so3_exponential_map

from kaolin.ops.mesh import index_vertices_by_faces
from kaolin.metrics.trianglemesh import point_to_mesh_distance


def laplace_regularizer_const(mesh_verts, mesh_faces):
    term = torch.zeros_like(mesh_verts)
    norm = torch.zeros_like(mesh_verts[..., 0:1])

    v0 = mesh_verts[mesh_faces[:, 0], :]
    v1 = mesh_verts[mesh_faces[:, 1], :]
    v2 = mesh_verts[mesh_faces[:, 2], :]

    term.scatter_add_(0, mesh_faces[:, 0:1].repeat(1,3), (v1 - v0) + (v2 - v0))
    term.scatter_add_(0, mesh_faces[:, 1:2].repeat(1,3), (v0 - v1) + (v2 - v1))
    term.scatter_add_(0, mesh_faces[:, 2:3].repeat(1,3), (v0 - v2) + (v1 - v2))

    two = torch.ones_like(v0) * 2.0
    norm.scatter_add_(0, mesh_faces[:, 0:1], two)
    norm.scatter_add_(0, mesh_faces[:, 1:2], two)
    norm.scatter_add_(0, mesh_faces[:, 2:3], two)

    term = term / torch.clamp(norm, min=1.0)

    return torch.mean(term**2)




class MeshHeadTrainer():
    def __init__(self, dataloader, meshhead, camera, optimizer, recorder, gpu_id):
        self.dataloader = dataloader
        self.meshhead = meshhead
        self.camera = camera
        self.optimizer = optimizer
        self.recorder = recorder
        self.device = torch.device('cuda:%d' % gpu_id)

    def train(self, start_epoch=0, epochs=1):
        for epoch in range(start_epoch, epochs):
            for idx, data in tqdm(enumerate(self.dataloader)):
                
                # prepare data
                to_cuda = ['images', 'masks', 'visibles', 'intrinsics', 'extrinsics', 'pose', 'scale', 'exp_coeff', 'landmarks_3d', 'exp_id']
                for data_item in to_cuda:
                    data[data_item] = data[data_item].to(device=self.device)

                images = data['images'].permute(0, 1, 3, 4, 2)
                masks = data['masks'].permute(0, 1, 3, 4, 2)
                visibles = data['visibles'].permute(0, 1, 3, 4, 2)
                resolution = images.shape[2]

                R = so3_exponential_map(data['pose'][:, :3])
                T = data['pose'][:, 3:, None]
                S = data['scale'][:, :, None]
                landmarks_3d_can = (torch.bmm(R.permute(0,2,1), (data['landmarks_3d'].permute(0, 2, 1) - T)) / S).permute(0, 2, 1)
                landmarks_3d_neutral = self.meshhead.get_landmarks()[None].repeat(data['landmarks_3d'].shape[0], 1, 1)
                data['landmarks_3d_neutral'] = landmarks_3d_neutral

                deform_data = {
                    'exp_coeff': data['exp_coeff'],
                    'query_pts': landmarks_3d_neutral
                }
                deform_data = self.meshhead.deform(deform_data)
                pred_landmarks_3d_can = deform_data['deformed_pts']
                loss_def = F.mse_loss(pred_landmarks_3d_can, landmarks_3d_can)

                deform_data = self.meshhead.query_sdf(deform_data)
                sdf_landmarks_3d = deform_data['sdf']
                loss_lmk = torch.abs(sdf_landmarks_3d[:, :, 0]).mean()

                data = self.meshhead.reconstruct(data)
                data = self.camera.render(data, resolution)
                render_images = data['render_images']
                render_soft_masks = data['render_soft_masks']
                exp_deform = data['exp_deform']
                pose_deform = data['pose_deform']
                verts_list = data['verts_list']
                faces_list = data['faces_list']

                loss_rgb = F.l1_loss(render_images[:, :, :, :, 0:3] * visibles, images * visibles)
                loss_sil = kaolin.metrics.render.mask_iou((render_soft_masks * visibles[:, :, :, :, 0]).view(-1, resolution, resolution), (masks * visibles).squeeze().view(-1, resolution, resolution))
                loss_offset = (exp_deform ** 2).sum(-1).mean() + (pose_deform ** 2).sum(-1).mean()

                loss_lap = 0.0
                for b in range(len(verts_list)):
                    loss_lap += laplace_regularizer_const(verts_list[b], faces_list[b])
                
                loss = loss_rgb * 1e-1 + loss_sil * 1e-1 + loss_def * 1e0 + loss_offset * 1e-2 + loss_lmk * 1e-1 + loss_lap * 1e2

                self.optimizer.zero_grad()
                loss.backward()
                self.optimizer.step()

                log = {
                    'data': data,
                    'meshhead' : self.meshhead,
                    'loss_rgb' : loss_rgb,
                    'loss_sil' : loss_sil,
                    'loss_def' : loss_def,
                    'loss_offset' : loss_offset,
                    'loss_lmk' : loss_lmk,
                    'loss_lap' : loss_lap,
                    'epoch' : epoch,
                    'iter' : idx + epoch * len(self.dataloader)
                }
                self.recorder.log(log)
                if idx + epoch * len(self.dataloader) >= 10000:
                    return