# -----------------------------------------------------------
# Copyright (C) 2020 NVIDIA Corporation. All rights reserved.
# Nvidia Source Code License-NC
# Code written by Xueting Li.
# -----------------------------------------------------------

# Script to compute a semantic template given trained reconstruction network.
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
from absl import app
from absl import flags

from ..nnutils import loss_utils
from ..nnutils import test_utils
from ..nnutils import geom_utils
from ..nnutils.smr import SoftRenderer
from ..nnutils import cub_mesh_s1 as mesh_net

from ..data import cub as cub_data
from ..utils import tf_visualizer
from ..utils import transformations
from ..utils import image as image_utils

import os
import cv2
import random
import scipy.misc
import torchvision
import numpy as np
from tqdm import tqdm
import os.path as osp
from PIL import Image
from collections import OrderedDict

import soft_renderer as sr
import neural_renderer as nr
import soft_renderer.functional as srf

import torch
import torch.nn as nn
import torchvision.utils as vutils

flags.DEFINE_integer('image_size', 256, 'training image size')
flags.DEFINE_boolean('pred_cam', True, 'If true predicts camera')
flags.DEFINE_integer('axis', 1, 'symmetric axis')
flags.DEFINE_integer('num_parts', 4, 'number of semantic parts')
flags.DEFINE_boolean('use_scops', True, 'If true load SCOPS in loader')
flags.DEFINE_string('model_path', 'unsup-mesh/cachedir/snapshots/cub_net/pred_net_latest.pth', 'model path')
flags.DEFINE_string('out_dir', 'unsup-mesh/cachedir/snapshots/cub_net/', 'output directory')

opts = flags.FLAGS

class ShapenetTester(test_utils.Tester):
    def define_model(self):
        opts = self.opts

        # define model
        img_size = (opts.image_size, opts.image_size)
        self.model = mesh_net.MeshNet(
            img_size, opts, nz_feat=opts.nz_feat,
            axis = opts.axis)

        self.load_my_state_dict(opts.model_path)
        self.model = self.model.cuda(device=opts.gpu_id)
        self.mean_shape = self.model.get_mean_shape()
        self.faces = self.model.faces.view(1, -1, 3)

        # define differentiable renderer
        self.renderer = SoftRenderer(opts.image_size, 'softmax')

        # define colorization tools
        self.colorize = image_utils.Colorize(opts.num_parts + 1)
        self.batch_colorize = image_utils.BatchColorize(opts.num_parts + 1)

        # define criterion functions
        self.texture_loss_fn = loss_utils.PerceptualTextureLoss()
        os.makedirs(opts.out_dir, exist_ok=True)

        return

    def init_dataset(self):
        opts = self.opts
        self.data_module = cub_data
        self.dataloader = self.data_module.data_loader(opts)
        self.resnet_transform = torchvision.transforms.Normalize(
            mean=[0.485, 0.456, 0.406],
            std=[0.229, 0.224, 0.225])

    def load_my_state_dict(self, resume_dir):
        saved_state_dict = torch.load(resume_dir)
        # registered buff tensors may have different batch size in training
        # so we do not load them from the pretrained models
        unwanted_keys = {"noise", "uv_sampler"}

        new_params = self.model.state_dict().copy()
        for name, param in new_params.items():
            if name not in unwanted_keys:
                new_params[name].copy_(saved_state_dict[name])
        self.model.load_state_dict(new_params)
        print(tf_visualizer.green("Loaded model from {}.".format(resume_dir)))

    def set_input(self, batch):
        opts = self.opts

        input_img_tensor = batch['img'].type(torch.FloatTensor)
        for b in range(input_img_tensor.size(0)):
            input_img_tensor[b] = self.resnet_transform(input_img_tensor[b])
        img_tensor = batch['img'].type(torch.FloatTensor)
        mask_tensor = batch['mask'].type(torch.FloatTensor)
        part_prob_tensor = batch['part_prob'].type(torch.FloatTensor)

        self.input_imgs = input_img_tensor.cuda()
        self.imgs = img_tensor.cuda()
        self.masks = mask_tensor.cuda()
        img_path = batch['img_path']
        self.part_segs = part_prob_tensor.permute(0, 3, 1, 2)

    def most_freq(self, arr):
        """
        Given an arr of N * D, return a N dimension array,
        indicating the most frequent element in 2nd dimension.
        """
        n,d = arr.size()
        k = torch.max(torch.unique(arr)) + 1
        arr_one_hot = torch.zeros(n * d, k).cuda()
        arr_flat = arr.view(-1, 1)
        arr_one_hot.scatter_(1, arr_flat, 1)
        arr_one_hot = arr_one_hot.view(n, d, k)
        arr_ = torch.sum(arr_one_hot, dim = 1)
        return torch.argmax(arr_, dim = 1)

    def compute_uv(self):
        dataloader = iter(self.dataloader)

        self.best_shape = None
        best_idx = 0
        best_mask_loss = 100000.0
        self.best_uv = None

        print(tf_visualizer.green("Search for the examplar instance."))
        for i, batch in tqdm(enumerate(self.dataloader)):
            self.set_input(batch)

            with torch.no_grad():
                outputs = self.model(self.input_imgs)
                # shape
                delta_v = outputs['delta_v']

                if(opts.symmetric):
                    delta_v = self.model.symmetrize(delta_v)
                    self.mean_shape = self.model.get_mean_shape()

                pred_vs = self.mean_shape + delta_v

                # camera
                proj_cam = outputs['cam']

                faces = self.faces.repeat(delta_v.size(0), 1, 1)
                ori_flip = torch.FloatTensor([1, -1, 1, 1, 1, -1, -1]).view(1,-1).cuda()
                proj_cam = proj_cam * ori_flip
                pred_seen, _, _ = self.renderer.forward(pred_vs, faces, proj_cam)
                rgb_pred_seen = pred_seen[:, 0:3, :, :]
                mask_pred_seen = pred_seen[:, 3, :, :]

                # flip images
                flip_imgs = torch.flip(self.imgs, [3])
                flip_masks = torch.flip(self.masks, [2])
                texture_loss = self.texture_loss_fn(rgb_pred_seen, flip_imgs, flip_masks, mask_pred_seen, avg = False)

                # the new best shape should also be close to the old one,
                # because we don't want the template to change all the time.
                mean_shape = self.mean_shape.unsqueeze(0).repeat(pred_vs.size(0), 1, 1)
                dist = torch.nn.MSELoss(reduction='none')(pred_vs, mean_shape)
                # scale is used to make sure these two losses are comparable
                scale = texture_loss.mean() / torch.sum(dist, dim = (1, 2)).mean()
                texture_loss += (torch.sum(dist, dim = (1, 2))) * scale

                # select a shape that has both low mask loss and close to the current best_shape
                min_mask, min_idx = torch.min(texture_loss, 0)
                if(best_mask_loss > min_mask):
                    best_idx = min_idx
                    best_mask_loss = min_mask
                    self.best_shape = pred_vs[best_idx].unsqueeze(0)
                    uvimage_pred = outputs['uvimage_pred']
                    uv_parts = torch.nn.functional.grid_sample(self.part_segs.cuda(), uvimage_pred.permute(0, 2, 3, 1))
                    self.best_uv = uv_parts[best_idx].unsqueeze(0)

                    # visulize semantic texture
                    tex_flow = outputs['tex_flow'][best_idx].unsqueeze(0)
                    tex = geom_utils.sample_textures(tex_flow, self.part_segs[best_idx].unsqueeze(0).cuda())
                    best_tex = tex.contiguous()

                    bs, fs, ts, _, cs = best_tex.size()
                    best_tex = best_tex.view(bs, fs, -1, cs)
                    best_tex = torch.argmax(best_tex, dim = -1)
                    self.best_tex = self.batch_colorize(best_tex.cpu())
                    self.best_tex = self.best_tex.permute(0, 2, 3, 1)

        self.best_shape = self.best_shape.repeat(opts.batch_size, 1, 1)
        self.best_uv = self.best_uv.repeat(opts.batch_size, 1, 1, 1)

        print(tf_visualizer.green("Start to compute semantic template."))
        counter = 0
        avg_uv_parts = None
        for i, batch in tqdm(enumerate(self.dataloader)):
            self.set_input(batch)

            with torch.no_grad():
                outputs = self.model(self.input_imgs)

                self.uvimage_pred = outputs['uvimage_pred']
                uv_parts = torch.nn.functional.grid_sample(self.part_segs.cuda(), self.uvimage_pred.permute(0, 2, 3, 1))

                uv_parts_ch = uv_parts.clone()
                best_uv_ch = self.best_uv.clone()
                dist = torch.nn.MSELoss(reduction='none')(uv_parts_ch, best_uv_ch)
                dist = torch.sum(dist, dim = (1, 2, 3))
                _, idx = torch.topk(dist, k = 5, largest = False)

                if(avg_uv_parts is None):
                    avg_uv_parts = torch.sum(uv_parts[idx, :, :, :], dim = 0).unsqueeze(0)
                else:
                    avg_uv_parts += torch.sum(uv_parts[idx, :, :, :], dim = 0).unsqueeze(0)
                counter += idx.size(0)

        avg_prob = avg_uv_parts / counter

        avg_prob = avg_prob.cpu().squeeze().numpy()
        avg_prob = avg_prob.transpose(1,2,0)
        uv_path = osp.join(opts.out_dir, "semantic_prob.npy")
        np.save(uv_path, avg_prob)
        avg_prob = np.asarray(np.argmax(avg_prob, axis=2), dtype=np.int)

        pil_image = Image.fromarray(avg_prob.astype(dtype=np.uint8))
        pil_image.save(osp.join(opts.out_dir, "semantic_seg.png"), 'PNG')

        color_vis = self.colorize(avg_prob)
        color_vis = torch.from_numpy(color_vis).float()

        # wrap the uv map onto template
        uv_sampler = self.model.uv_sampler[0].unsqueeze(0)
        tex = torch.nn.functional.grid_sample(color_vis.unsqueeze(0).cuda().float(), uv_sampler)
        tex = tex.view(tex.size(0), -1, tex.size(2), opts.tex_size, opts.tex_size).permute(0, 2, 3, 4, 1)
        tex_left = tex[:, -self.model.texture_predictor.num_sym_faces:]
        tex = torch.cat([tex, tex_left], 1)
        tex = tex.view(tex.size(0), tex.size(1), -1, 3)

        mean_v = self.model.get_mean_shape()
        mesh_ = sr.Mesh(mean_v, self.faces, tex)
        mesh_path = osp.join(opts.out_dir,"mean_template.obj")
        mesh_.save_obj(mesh_path, save_texture=True)

        # compute vertices/faces belong to each part
        uv_label = np.load(uv_path)
        uv_label = torch.from_numpy(uv_label).float().unsqueeze(0).permute(0, 3, 1, 2)
        uv_label = uv_label.cuda()
        tex_seg = torch.nn.functional.grid_sample(uv_label, uv_sampler)
        tex_seg = tex_seg.view(tex_seg.size(0), -1, tex_seg.size(2), opts.tex_size, opts.tex_size).permute(0, 2, 3, 4, 1)
        tex_left = tex_seg[:, -self.model.texture_predictor.num_sym_faces:]
        tex_seg = torch.cat([tex_seg, tex_left], 1)
        tex_seg = tex_seg.view(tex_seg.size(0), tex_seg.size(1), -1, (self.opts.num_parts + 1))
        tex_seg = torch.argmax(tex_seg, dim = -1)

        # obtain vertex label through face label
        tex_seg = self.most_freq(tex_seg.squeeze())
        tex_seg = tex_seg.float()
        face = self.faces[0]
        parts = []
        for cnt in range(opts.num_parts):
            parts.append([])

        # go through all vertices and compute their label
        # for sanity check
        vert_tex = [] 
        for cnt in range(face.max() + 1):
            v0 = (face[:, 0] == cnt) * 1
            v1 = (face[:, 1] == cnt) * 1
            v2 = (face[:, 2] == cnt) * 1
            v = v0 + v1 + v2
            # which faces relate to this vertex
            idxes = torch.nonzero(v).squeeze()
            labels = tex_seg[idxes].long().view(1, idxes.size(0))
            label = self.most_freq(labels)

            if(label > 0):
                parts[label-1].append(cnt)
            vert_tex.append(label)

        np.save(osp.join(opts.out_dir, "head_vertices.npy"), parts[0])
        np.save(osp.join(opts.out_dir, "neck_vertices.npy"), parts[1])
        np.save(osp.join(opts.out_dir, "back_vertices.npy"), parts[2])
        np.save(osp.join(opts.out_dir, "belly_vertices.npy"), parts[3])

        # visualize part label for each vertex
        vert_tex = torch.stack(vert_tex)
        vert_tex = self.colorize(vert_tex.view(642, 1).cpu().numpy())
        vert_tex = torch.from_numpy(vert_tex).float().squeeze()
        vert_tex = vert_tex.permute(1, 0)

        mesh_ = sr.Mesh(mean_v, self.faces, vert_tex.view(1, 642, 3), texture_type='vertex')
        mesh_path = osp.join(opts.out_dir,"vertex_label.obj")
        mesh_.save_obj(mesh_path, save_texture=True)
        
        torch.save(self.model.mean_v.detach().cpu(), osp.join(opts.out_dir, "mean_v.pth"))

        print(tf_visualizer.green("Semantic template saved at {}.".format(opts.out_dir)))

def set_seed(seed):
    torch.backends.cudnn.deterministic = True
    torch.backends.cudnn.benchmark = False
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)
    np.random.seed(seed)
    random.seed(seed)

def main(_):
    set_seed(0)
    tester = ShapenetTester(opts)
    tester.init_testing()
    tester.compute_uv()

if __name__ == '__main__':
    app.run(main)
