# Copyright 2021 Google LLC
# 
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
# 
#     https://www.apache.org/licenses/LICENSE-2.0
# 
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

from absl import app
from absl import flags
import os
import os.path as osp
import sys
sys.path.insert(0,'third_party')

import numpy as np
import configparser
import time
import torch
import torchvision
import torch.nn as nn
from torch.autograd import Variable
import torch.nn.functional as F
import trimesh
import pytorch3d
import pytorch3d.loss

import pdb
from ext_utils import mesh
from ext_utils.quatlib import q_rnd_m, q_scale_m
from ext_utils.util_rot import compute_geodesic_distance_from_two_matrices
from ext_utils import geometry as geom_utils
from ext_nnutils import net_blocks as nb
from ext_nnutils.mesh_net import MeshNet
import kornia
import configparser
import soft_renderer as sr
from nnutils.geom_utils import pinhole_cam, obj_to_cam, render_flow_soft_3
from nnutils.geom_utils import label_colormap
citylabs = label_colormap()

#-------------- flags -------------#
#----------------------------------#
flags.DEFINE_boolean('noise', True, 'Add random noise to pose')
flags.DEFINE_boolean('symmetric', True, 'Use symmetric mesh or not')
flags.DEFINE_boolean('symmetric_loss', True, 'Use symmetric loss or not')
flags.DEFINE_integer('nz_feat', 200, 'Encoded feature size')

flags.DEFINE_boolean('texture', True, 'if true uses texture!')
flags.DEFINE_boolean('symmetric_texture', True, 'if true texture is symmetric!')

flags.DEFINE_integer('subdivide', 3, '# to subdivide icosahedron, 3=642verts, 4=2562 verts')
flags.DEFINE_integer('symidx', 0, 'symmetry index: 0-x 1-y 2-z')
flags.DEFINE_integer('n_bones', 1, 'num of meshes')
flags.DEFINE_string('n_faces', '1280','number of faces for remeshing')
flags.DEFINE_integer('n_hypo', 1, 'num of hypothesis cameras')

flags.DEFINE_boolean('only_mean_sym', False, 'If true, only the meanshape is symmetric')
flags.DEFINE_string('dataname', 'fashion', 'name of the test data')
flags.DEFINE_string('opt_tex', 'no', 'optimize texture')
flags.DEFINE_float('rscale', 1.0, 'scale random variance')
flags.DEFINE_float('l1tex_wt', 1.0, 'weight of l1 texture')
flags.DEFINE_float('sigval', 1e-4, 'blur radius of soft renderer')

def render_flow_soft_2(renderer_soft, verts, faces, verts_pos0, verts_pos1, pp0, pp1, proj_cam0,proj_cam1):
    # flow (no splat): 1) get mask; 2) render 3D coords for 1st/2nd frame 
    n_hypo = verts.shape[0] // faces.shape[0]
    faces = faces[:,None].repeat(1,n_hypo,1,1).view(-1,faces.shape[1],3)
    verts_pos0 = verts_pos0.clone()
    verts_pos1 = verts_pos1.clone()
    offset = torch.Tensor( renderer_soft.transform.transformer._eye ).cuda()[np.newaxis,np.newaxis]
    verts_pre = verts[:,:,:3]+offset; verts_pre[:,:,1] = -1*verts_pre[:,:,1]

    nb = verts.shape[0]
    verts_pos_px = renderer_soft.render_mesh(sr.Mesh(torch.cat([verts_pre         ,verts_pre],0),
                                                     torch.cat([faces             ,faces],0), 
                                            textures=torch.cat([verts_pos0[:,:,:3],verts_pos1[:,:,:3]],0),texture_type='vertex')).clone()
    fgmask = verts_pos_px[:nb,-1]
    verts_pos_px = verts_pos_px[:,:3]
    verts_pos0_px = verts_pos_px[:nb].permute(0,2,3,1)
    verts_pos1_px = verts_pos_px[nb:].permute(0,2,3,1)

    bgmask = (verts_pos0_px[:,:,:,2]<1e-9) | (verts_pos1_px[:,:,:,2]<1e-9)
    verts_pos0_px[bgmask]=10
    verts_pos1_px[bgmask]=10

    # projet 3D verts with different intrinsics
    verts_pos0_px[:,:,:,1] = pp0[:,1:2,np.newaxis]+verts_pos0_px[:,:,:,1].clone()*proj_cam0[:,:1,np.newaxis] / verts_pos0_px[:,:,:,2].clone()
    verts_pos0_px[:,:,:,0] = pp0[:,0:1,np.newaxis]+verts_pos0_px[:,:,:,0].clone()*proj_cam0[:,:1,np.newaxis] / verts_pos0_px[:,:,:,2].clone()
    verts_pos1_px[:,:,:,1] = pp1[:,1:2,np.newaxis]+verts_pos1_px[:,:,:,1].clone()*proj_cam1[:,:1,np.newaxis] / verts_pos1_px[:,:,:,2].clone()
    verts_pos1_px[:,:,:,0] = pp1[:,0:1,np.newaxis]+verts_pos1_px[:,:,:,0].clone()*proj_cam1[:,:1,np.newaxis] / verts_pos1_px[:,:,:,2].clone()
    flow_fw = (verts_pos1_px - verts_pos0_px.detach())[:,:,:,:2]
    flow_fw[bgmask] = flow_fw[bgmask].detach()
    return flow_fw, bgmask, fgmask

def reg_decay(curr_steps, max_steps, min_wt,max_wt):
    """
    max weight to min weight
    """
    if curr_steps>max_steps:current = min_wt
    else:
        current = np.exp(curr_steps/float(max_steps)*(np.log(min_wt)-np.log(max_wt))) * max_wt 
    return current

class LASR(MeshNet):
    def __init__(self, input_shape, opts, nz_feat=100):
        super(LASR, self).__init__(input_shape, opts, nz_feat)
        self.rest_rs = torch.Tensor([[0,0,0,1]]).repeat(opts.n_bones-1,1).cuda()
        self.transg =  torch.Tensor([[0,0,0]]).cuda().repeat(opts.n_bones-1,1)  # not including the body-to-world transform
        self.ctl_rs =  torch.Tensor([[0,0,0,1]]).repeat(opts.n_hypo*(opts.n_bones-1),1).cuda()
        self.rest_ts = torch.zeros(opts.n_hypo*(opts.n_bones-1),3).cuda()
        self.ctl_ts =  torch.zeros(opts.n_hypo*(opts.n_bones-1),3).cuda()
        self.log_ctl = torch.Tensor([[0,0,0]]).cuda().repeat(opts.n_hypo*(opts.n_bones-1),1)  # control point varuance

        if self.opts.n_bones>1:
            self.ctl_rs  = nn.Parameter(self.ctl_rs) 
            self.rest_ts = nn.Parameter(self.rest_ts)
            self.ctl_ts  = nn.Parameter(self.ctl_ts) 
            self.log_ctl = nn.Parameter(self.log_ctl)

        # For renderering.
        self.renderer_soft = sr.SoftRenderer(image_size=opts.img_size, sigma_val=1e-4,
                       camera_mode='look_at',perspective=False, aggr_func_rgb='hard',
                       light_mode='vertex', light_intensity_ambient=1.,light_intensity_directionals=0.)
        
        self.renderer_softflf = sr.SoftRenderer(image_size=opts.img_size, sigma_val=1e-4, gamma_val=1e-2,
                       camera_mode='look_at',perspective=False, aggr_func_rgb='softmax',
                       light_mode='vertex', light_intensity_ambient=1.,light_intensity_directionals=0.)
        self.renderer_softflb = sr.SoftRenderer(image_size=opts.img_size, sigma_val=1e-4, gamma_val=1e-2,
                       camera_mode='look_at',perspective=False, aggr_func_rgb='softmax',
                       light_mode='vertex', light_intensity_ambient=1.,light_intensity_directionals=0.)
        
        self.renderer_softtex = sr.SoftRenderer(image_size=opts.img_size, sigma_val=1e-4,gamma_val=1e-2,
                       camera_mode='look_at',perspective=False, aggr_func_rgb='softmax',
                       light_mode='vertex', light_intensity_ambient=1.,light_intensity_directionals=0.)
        
        self.renderer_softpart = sr.SoftRenderer(image_size=opts.img_size, sigma_val=1e-4,gamma_val=1e-4, 
                       camera_mode='look_at',perspective=False, aggr_func_rgb='softmax',
                       light_mode='vertex', light_intensity_ambient=1.,light_intensity_directionals=0.)
        

    def forward(self, batch_input):
        if self.training:
            local_batch_size = batch_input['input_imgs  '].shape[0]//2
            for k,v in batch_input.items():
                batch_input[k] = v.view(local_batch_size,2,-1).permute(1,0,2).reshape(v.shape)
            self.input_imgs   = batch_input['input_imgs  ']
            self.imgs         = batch_input['imgs        ']
            self.masks        = batch_input['masks       ']
            self.cams         = batch_input['cams        ']
            self.depth_gt     = batch_input['depth_gt    ']
            self.flow         = batch_input['flow        ']
            self.dts_barrier  = batch_input['dts_barrier ']
            self.ddts_barrier = batch_input['ddts_barrier']
            self.mask_contour = batch_input['mask_contour']
            self.pp           = batch_input['pp          ']
            self.occ          = batch_input['occ         ']
            self.oriimg_shape = batch_input['oriimg_shape']
            self.frameid      = batch_input['frameid']
            self.dataid      = batch_input['dataid']
            self.is_canonical = batch_input['is_canonical']
        else:
            local_batch_size = 1
            self.input_imgs = batch_input
        img = self.input_imgs

#        torch.cuda.synchronize()
#        start_time = time.time()

        opts = self.opts
        pred_v, tex,faces= self.get_mean_shape(local_batch_size)
        if self.training:
            tex = tex.view(2*local_batch_size,-1,opts.n_hypo,tex.shape[-2],3)
            texnew = torch.zeros_like(tex[:,0])
            for i in range(local_batch_size*2):
                texnew[i] = tex[i,self.dataid[i]]
            tex = texnew.reshape(-1,tex.shape[-2],3)
        faces = faces.cuda()
        
        def set_bn_eval(m):
            classname = m.__class__.__name__
            if classname.find('BatchNorm') != -1:
                m.eval()
        #if opts.n_hypo==1:
        self.apply(set_bn_eval)

        img_feat = self.encoder.forward(img)
        scale, trans, quat, depth, ppoint = self.code_predictor.forward(img_feat)
        if not self.training:
           return scale,trans,quat, depth, ppoint
        
        # transform the CNN-predicted focal length 
        # (wrt original image without cropping) to the cropped image
        scale = self.cams[:,:1]*scale # here assumes intrinsic may change
        # transform depth as well, to ensure the rendered silhouette 
        # occupies enough pixels of the cropped image at initialization 
        depth[:,:1] = self.cams[:,:1]* depth[:,:1]; depth = depth.view(-1,1) 
       
        # transform the CNN-predicted principal points 
        # (wrt original image, before cropping) to the cropped image
        ppb1 = self.cams[:local_batch_size,:1]*self.pp[:local_batch_size]/(opts.img_size/2.)
        ppb2 = self.cams[local_batch_size:,:1]*self.pp[local_batch_size:]/(opts.img_size/2.)
        # represent pp of croped frame 2 as transformed pp of cropped frame 1
        # to reduce ambiguity caused by inconsistent pp over time
        ppa1 = ppoint[:local_batch_size] + ppb1 + 1
        ppa2 = ppa1 * (self.cams[local_batch_size:,:1] / self.cams[:local_batch_size,:1]) 
        ppoint[local_batch_size:]= ppa2 - ppb2 -1
        
        quat = quat.view(-1,9)
        if opts.noise and self.epoch>0 and self.iters<100 and self.iters>1:
            # add noise
            decay_factor = 0.2*(1e-4)**(self.iters/100)
            decay_factor_r = decay_factor * np.ones(quat.shape[0])
            ### smaller noise for bones
            decay_factor_r = decay_factor_r.reshape((-1,opts.n_bones))
            decay_factor_r[:,1:] *= 1
            decay_factor_r[:,0] *=  1
            decay_factor_r = decay_factor_r.flatten()
            noise = q_scale_m(q_rnd_m(b=quat.shape[0]), decay_factor_r)  # wxyz
            noise = torch.Tensor(noise).cuda()
            noise = torch.cat([noise[:,1:], noise[:,:1]],-1)
            quat = quat.view(-1,3,3).matmul(kornia.quaternion_to_rotation_matrix(noise)).view(-1,9)

            decay_factor_s = decay_factor
            scale = scale * (decay_factor_s*torch.normal(torch.zeros(scale.shape).cuda(),opts.rscale)).exp()
            
        depth = depth.view(local_batch_size*2,1,opts.n_bones,1).repeat(1,opts.n_hypo,1,1).view(-1,1)
        trans = trans.view(local_batch_size*2,1,opts.n_bones,2).repeat(1,opts.n_hypo,1,1).view(-1,2)

        if opts.use_gtpose:
            # w/ gt cam
            quat_pred = quat.clone()
            scale_pred = scale.clone()
            trans_pred = trans.clone()
            ppoint_pred = ppoint.clone()
            depth_pred = depth.clone()

            scale = 10*self.cams[:,:1]
            trans = self.cams[:,1:3]
            quat = kornia.quaternion_to_rotation_matrix( torch.cat( (self.cams[:,4:],self.cams[:,3:4] ) ,-1) ).view(-1,9)
            depth    = self.depth_gt[:]
            halforisize = 0.5*opts.img_size / self.cams[:,:1]  # half size before scaling 
            ppoint = (0.5*self.oriimg_shape - self.pp[:]) / halforisize - 1
    
        ## rendering
        # obj-cam rigid transform;  proj_cam: [focal, tx,ty,qw,qx,qy,qz]; 
        # 1st/2nd frame stored as 1st:[0:local_batch_size], 2nd: [local_batch_size:-1]
        # transforms [body-to-cam, part1-to-body, ...]
        Rmat = quat.view(-1,3,3).permute(0,2,1)
        Tmat = torch.cat([trans, depth],1)
        if opts.n_bones>1:
            # skin computation
            # GMM
            dis_norm = (self.ctl_ts.view(opts.n_hypo,-1,1,3) - pred_v.view(2*local_batch_size,opts.n_hypo,-1,3)[0,:,None].detach()) # p-v, H,J,1,3 - H,1,N,3
            dis_norm = dis_norm.matmul(kornia.quaternion_to_rotation_matrix(self.ctl_rs).view(opts.n_hypo,-1,3,3)) # h,j,n,3
            dis_norm = self.log_ctl.exp().view(opts.n_hypo,-1,1,3) * dis_norm.pow(2) # (p-v)^TS(p-v) 
            skin = (-10 * dis_norm.sum(3)).softmax(1)[:,:,:,None] # h,j,n,1
            # color palette
            colormap = torch.Tensor(citylabs[:skin.shape[1]]).cuda() # 5x3
            skin_colors = (skin[self.optim_idx] * colormap[:,None]).sum(0)/256.
            skin = skin.repeat(local_batch_size*2,1,1,1)
          
            # rest_ts: joint center
            # ctl_ts: control points
            rest_ts = self.rest_ts[:,None,:,None].repeat(local_batch_size*2,1,1,1).view(-1,opts.n_bones-1,3,1) # bh,m,3,1
            ctl_ts = self.ctl_ts[:,None,:,None].repeat(local_batch_size*2,1,1,1).view(-1,opts.n_bones-1,3,1) # bh,m,3,1
            
            Rmat = Rmat.view(-1,opts.n_bones,3,3)
            Tmat = Tmat.view(-1,opts.n_bones,3,1)
            Tmat[:,1:] = -Rmat[:,1:].matmul(rest_ts)+Tmat[:,1:]+rest_ts
            Rmat[:,1:] = Rmat[:,1:].permute(0,1,3,2)
            Rmat = Rmat.view(-1,3,3)
            Tmat = Tmat.view(-1,3)

            self.joints_proj = obj_to_cam(rest_ts[:,:,:,0], Rmat.detach(), Tmat[:,np.newaxis].detach(), opts.n_bones, opts.n_hypo,torch.eye(opts.n_bones-1)[None,:,:,None].cuda())
            self.joints_proj = pinhole_cam(self.joints_proj, ppoint.detach(), scale.detach())
            self.ctl_proj =    obj_to_cam(ctl_ts[:,:,:,0],  Rmat.detach(), Tmat[:,np.newaxis].detach(), opts.n_bones, opts.n_hypo, torch.eye(opts.n_bones-1)[None,:,:,None].cuda())
            self.ctl_proj = pinhole_cam(self.ctl_proj, ppoint.detach(), scale.detach())
        else:skin=None                

        self.deform_v = obj_to_cam(pred_v, Rmat.view(-1,3,3), Tmat[:,np.newaxis,:],opts.n_bones, opts.n_hypo,skin,tocam=False)

#        torch.cuda.synchronize()
#        print('before rend time:%.2f'%(time.time()-start_time))

        
        # 1) flow rendering 
        verts_fl = obj_to_cam(pred_v, Rmat, Tmat[:,np.newaxis,:],opts.n_bones, opts.n_hypo,skin)
        verts_fl = torch.cat([verts_fl,torch.ones_like(verts_fl[:, :, 0:1])], dim=-1)
        verts_pos0 = verts_fl.view(2*local_batch_size,opts.n_hypo,-1,4)[:local_batch_size].clone().view(local_batch_size*opts.n_hypo,-1,4)
        verts_pos1 = verts_fl.view(2*local_batch_size,opts.n_hypo,-1,4)[local_batch_size:].clone().view(local_batch_size*opts.n_hypo,-1,4)
        verts_fl = pinhole_cam(verts_fl, ppoint, scale)

        dmax=verts_fl[:,:,-2].max()
        dmin=verts_fl[:,:,-2].min()
        self.renderer_softflf.rasterizer.near=dmin-(dmax-dmin)/2
        self.renderer_softflf.rasterizer.far= dmax+(dmax-dmin)/2
        self.renderer_softflb.rasterizer.near=dmin-(dmax-dmin)/2
        self.renderer_softflb.rasterizer.far= dmax+(dmax-dmin)/2
        self.renderer_softtex.rasterizer.near=dmin-(dmax-dmin)/2
        self.renderer_softtex.rasterizer.far= dmax+(dmax-dmin)/2
        if opts.sigval!=1e-4:
            self.renderer_soft.rasterizer.sigma_val=   opts.sigval 
            self.renderer_softflf.rasterizer.sigma_val=opts.sigval 
            self.renderer_softflb.rasterizer.sigma_val=opts.sigval 
            self.renderer_softtex.rasterizer.sigma_val=opts.sigval 

        self.flow_fw, self.bgmask_fw, self.fgmask_flowf = render_flow_soft_2(self.renderer_softflf, verts_fl.view(2*local_batch_size,opts.n_hypo,-1,4)[:local_batch_size].view(-1,verts_fl.shape[1],4), 
                                              faces[:local_batch_size], 
                                              verts_pos0, verts_pos1, 
                                       ppoint[:,None][:local_batch_size].repeat(1,opts.n_hypo,1).view(-1,2),
                                       ppoint[:,None][local_batch_size:].repeat(1,opts.n_hypo,1).view(-1,2), 
                                        scale[:,None][:local_batch_size].view(-1,1),
                                        scale[:,None][local_batch_size:].view(-1,1))