import os
from time import time

import numpy as np
import matplotlib.pyplot as plt

import torch

# My libraries
import datasets.data_loader as data_loader
import utils.data_utils as data_utils
# os.environ['CUDA_VISIBLE_DEVICES'] = "0" # TODO: Change this if you have more than 1 GPU
from models.two_stream_attention_lang_fusion import TwoStreamAttentionLangFusionLat
import torch.nn as nn
import argparse
from tqdm import tqdm
import json
from tensorboardX import SummaryWriter
import datetime
from models.loss import BCEWithLogitsLossWeighted

def mkdir_if_missing(dst_dir):
    if not os.path.exists(dst_dir):
        os.makedirs(dst_dir)

parser = argparse.ArgumentParser()
parser.add_argument('--batch_size', type=int, default=2)
parser.add_argument('--num_workers', type=int, default=4)
parser.add_argument('--epoch', type=int, default=1000)
parser.add_argument('--save_rate', type=int, default=1000)
parser.add_argument('--lr', type=float, default=1e-3)
parser.add_argument('--model_suffix', type=str, default='')
parser.add_argument('--config_name', type=str, required=True)
parser.add_argument('--load_pretrained', action='store_true')

args = parser.parse_args()

args.model_suffix += args.config_name
with open(os.path.join('configs', args.config_name + '.json')) as f:
    data_loading_params = json.load(f)

if 'OCID' in data_loading_params['dataset']:
    dl = data_loader.get_OCID_train_dataloader(data_loading_params['dataset'], data_loading_params, batch_size=args.batch_size, num_workers=args.num_workers, shuffle=True)
else:
    dl = data_loader.get_TOD_train_dataloader(data_loading_params['dataset'], data_loading_params, batch_size=args.batch_size, num_workers=args.num_workers, shuffle=True)
loss_func = BCEWithLogitsLossWeighted(weighted=True) # torch.nn.BCEWithLogitsLoss(reduction='none')
net = TwoStreamAttentionLangFusionLat().cuda()
params_exclude_clip = net.params_exclude_clip
net = torch.nn.DataParallel(net)
optimizer = torch.optim.Adam(params_exclude_clip, lr=args.lr)
total_losses = data_utils.AverageMeter()
scheduler = torch.optim.lr_scheduler.MultiStepLR(optimizer, milestones=[1, 3, 5, 7], gamma=0.5)
total_iter = 0
logdir = 'tensorboard/{}_{}'.format(args.model_suffix, datetime.datetime.now().strftime("%d_%m_%Y_%H:%M:%S"))
mkdir_if_missing(logdir)
print("logdir:", logdir)

tb_logger = SummaryWriter(logdir=logdir)
mkdir_if_missing('output')

# train
epoch_num = args.epoch
for n in range(epoch_num):
    for data in tqdm(dl):
        foreground_labels = data['target_labels'] # Shape: [N x H x W]
        N, H, W = foreground_labels.shape
        x = torch.cat((data['rgb'], data['xyz']), dim=1)
        x = torch.nn.functional.interpolate(x, (240, 320))
        fg_logits  = net(x.cuda(), data['tokenize_prompt'].cuda())
        ### Foreground Loss ###

        foreground_labels = torch.nn.functional.interpolate(foreground_labels[:,None], (240, 320))
        fg_masks = foreground_labels.clamp(0,2).cuda()
        fg_loss = loss_func(fg_logits, fg_masks).mean()
        optimizer.zero_grad()
        loss = fg_loss
        loss.backward()
        optimizer.step()
        total_losses.update(loss.item())
        print('epoch: {} iter: {} loss: {}'.format(n, total_iter, total_losses.avg))
        total_iter += 1

        if total_iter % 10 == 0:
            tb_logger.add_scalar('bce_loss', total_losses.avg, total_iter)
            # also write image
        if total_iter % 100 == 0:
            tb_logger.add_image("gt_image", data_utils.torch_to_numpy(data['rgb'], is_standardized_image=True)[0], total_iter, dataformats='HWC')
            tb_logger.add_image("gt_depth", data_utils.torch_to_numpy(data['xyz'])[0], total_iter, dataformats='HWC')
            tb_logger.add_image("gt_mask", fg_masks[0], total_iter, dataformats='CHW')
            tb_logger.add_image("pred_mask", torch.sigmoid(fg_logits[0]), total_iter, dataformats='CHW')

        if total_iter % args.save_rate == 0:
            filename = "output/model_{}.pth".format(args.model_suffix)
            checkpoint = {'model' : net.state_dict()}
            torch.save(checkpoint, filename)
    scheduler.step()
    print("epoch: {}".format(n))

# test


import sys
sys.path.append(".")
import torch.nn as nn
import torch
import torch.nn.functional as F
import torchvision.models as models
import cv2
import numpy as np
import math
from functools import reduce
import matplotlib.pyplot as plt
import imageio
from skimage.transform import resize
from PIL import Image
from skimage.io import imread, imsave
import neural_renderer
import torchvision.transforms as transforms
from torchvision.utils import save_image
from .StyleLoss import get_style_model_and_losses, image_loader, run_style_transfer

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

class StyleTransferModel(nn.Module):
    def __init__(
            self,
            filename_mesh,
            filename_style,
            texture_size=4,
            camera_distance=2.732,
            camera_distance_noise=0.1,
            elevation_min=20,
            elevation_max=40,
            lr_vertices=0.01,
            lr_textures=1.0,
            lambda_style=1,
            lambda_content=1e5,
            lambda_tv=1e2,
            image_size=256,
    ):
        super(StyleTransferModel, self).__init__()
        self.image_size = image_size
        self.camera_distance = camera_distance
        self.camera_distance_noise = camera_distance_noise
        self.elevation_min = elevation_min
        self.elevation_max = elevation_max
        self.lambda_style = lambda_style
        self.lambda_content = lambda_content
        self.lambda_tv = lambda_tv
        self.xp = np

        # load reference image
        self.reference_image = imread(filename_style).astype('float32') / 255.
        #self.reference_image = cv2.resize(self.reference_image,(image_size,image_size))
        self.reference_image = torch.from_numpy(self.reference_image).permute(2,0,1)[None, ::].to(device)
        self.cnn = models.vgg19(pretrained=True).features.to(device).eval()
        for param in self.cnn.parameters():
            param.requires_grad = False
        self.cnn_normalization_mean = torch.tensor([0.485, 0.456, 0.406]).to(device)
        self.cnn_normalization_std = torch.tensor([0.229, 0.224, 0.225]).to(device)
        self.style_layers = ['conv_2', 'conv_4', 'conv_7', 'conv_10']

        self.background_color = self.reference_image.mean((0, 2, 3))

        # load .obj
        vertices,faces = neural_renderer.load_obj(filename_mesh)
        self.vertices_original = torch.clone(vertices[None,:,:]).detach()
        self.vertices = nn.Parameter(vertices[None,:,:])

        self.register_buffer('faces', faces[None, :, :])

        texture_size = 4
        textures = torch.zeros(1, self.faces.shape[1], texture_size, texture_size, texture_size, 3, dtype=torch.float32)
        self.textures = nn.Parameter(textures)

        # setup renderer
        renderer = neural_renderer.Renderer(camera_mode='look_at')
        renderer.perspective = False
        renderer.light_intensity_directional = 0.0
        renderer.light_intensity_ambient = 1.0
        self.renderer = renderer

    def image_loader(image_name,imsize,device):
        loader = transforms.Compose([transforms.Resize((imsize,imsize)),transforms.ToTensor()])
        image = Image.open(image_name)
        # fake batch dimension required to fit network's input dimensions
        image = loader(image).unsqueeze(0)
        return image.to(device, torch.float)

    def compute_style_loss(self, images):
        style_loss = 0
        # print('Reference Image',self.reference_image)
        # print('Content Image',images[0])
        # plt.imshow(self.reference_image[0].permute(1,2,0).cpu().detach().numpy())
        # plt.show()
        # plt.imshow(images[0].permute(1,2,0).cpu().detach().numpy())
        # plt.show()
        for image in images:
            image = image.unsqueeze(0)
            style_loss += (run_style_transfer(self.cnn,self.cnn_normalization_mean,self.cnn_normalization_std,image,self.reference_image,self.style_layers))*100000
            style_loss += torch.sum((image - self.reference_image) ** 2)
        return style_loss

    def compute_content_loss(self):
        loss = torch.sum(torch.square(self.vertices - self.vertices_original))
        return loss

    def compute_tv_loss(self, images, masks):
        s1 = torch.square(images[:, :, 1:, :-1] - images[:, :, :-1, :-1])
        s2 = torch.square(images[:, :, :-1, 1:] - images[:, :, :-1, :-1])
        masks = masks[:, None, :-1, :-1].expand(s1.shape)
        masks = masks.data == 1
        return torch.sum(masks * (s1 + s2))

    def __call__(self, batch_size):
        xp = self.xp

        # set random viewpoints
        self.renderer.eye = neural_renderer.get_points_from_angles(2.732, 0, np.random.uniform(0, 360))

        # compute loss

        x = self.vertices.repeat(batch_size,1,1)
        x = x.to(device)

        y = self.faces.repeat(batch_size,1,1)
        y = y.to(device)

        z = self.textures.repeat(batch_size,1,1,1,1,1)
        z = z.to(device)

        images = self.renderer.render_rgb(x,y,torch.tanh(z))
        #images = torch.flip(images,dims=[1])
        # for image in images:
        #     plt.imshow(image.permute(1,2,0).cpu().detach().numpy())
        #     plt.show()
        masks = self.renderer.render_silhouettes(x,y)
        # for image in masks:
        #     plt.imshow(image.cpu().detach().numpy())
        #     plt.show()
        #     break
        # import IPython
        # IPython.embed()
        loss_style = self.compute_style_loss(images)
        loss_content = self.compute_content_loss()
        loss_tv = self.compute_tv_loss(images, masks)
        # print('Content weight=',self.lambda_content)
        # print('Content Loss=',loss_content)
        # print('Style weight=',self.lambda_style)
        # print('Style Loss=',loss_style)
        # print('TV Loss=',loss_tv)

        # print('Content Loss (UPGRADED)=',self.lambda_content*loss_content)
        # print('Style weight=',self.lambda_style)
        # print('Style Loss (UPGRADED)=',self.lambda_style*loss_style)
        # print('TV Loss=(UPGRADED)',self.lambda_tv*loss_tv)

        loss = self.lambda_style * loss_style + self.lambda_content * loss_content + self.lambda_tv * loss_tv
        # set default lighting direction

        return loss


import torch
from model.base_model import BaseModel
from model.networks import base_function, external_function
import model.networks as network
from util import task, util
import itertools
import data as Dataset
import numpy as np
from itertools import islice
import random
import os
import matplotlib.pyplot as plt
from collections import OrderedDict
import glob
import cv2

class Face(BaseModel):
    """
    Face Image Animation using edge image
    """
    def name(self):
        return "Face Image Animation"

    @staticmethod
    def modify_options(parser, is_train=True):
        """Add new options and rewrite default values for existing options"""
        parser.add_argument('--attn_layer', action=util.StoreList, metavar="VAL1,VAL2...")
        parser.add_argument('--kernel_size', action=util.StoreDictKeyPair, metavar="KEY1=VAL1,KEY2=VAL2...")

        parser.add_argument('--layers', type=int, default=3, help='number of layers in G')
        parser.add_argument('--netG', type=str, default='face', help='The name of net Generator')
        parser.add_argument('--netD', type=str, default='res', help='The name of net Discriminator')
        parser.add_argument('--netD_V', type=str, default='res', help='The name of net Discriminator')
        parser.add_argument('--init_type', type=str, default='orthogonal', help='Initial type')

        # if is_train:
        parser.add_argument('--ratio_g2d', type=float, default=0.1, help='learning rate ratio G to D')
        parser.add_argument('--lambda_rec', type=float, default=5.0, help='weight for image reconstruction loss')
        parser.add_argument('--lambda_g', type=float, default=2.0, help='weight for generation loss')
        parser.add_argument('--lambda_correct', type=float, default=5.0, help='weight for Sampling Correctness loss')
        parser.add_argument('--lambda_style', type=float, default=500.0, help='weight for the VGG19 style loss')
        parser.add_argument('--lambda_content', type=float, default=0.5, help='weight for the VGG19 content loss')
        parser.add_argument('--lambda_regularization', type=float, default=0.0025, help='weight for the affine regularization loss')
        parser.add_argument('--frames_D_V', type=int, default=3, help='number of frames of D_V')

        parser.add_argument('--use_spect_g', action='store_false')
        parser.add_argument('--use_spect_d', action='store_false')
        parser.set_defaults(use_spect_g=False)
        parser.set_defaults(use_spect_d=True)
        parser.set_defaults(display_freq=100)
        parser.set_defaults(eval_iters_freq=1000)
        parser.set_defaults(print_freq=100)
        parser.set_defaults(save_latest_freq=1000)
        parser.set_defaults(save_iters_freq=10000)

        return parser

    def __init__(self, opt):
        BaseModel.__init__(self, opt)
        self.loss_names = ['app_gen','correctness_p', 'correctness_r','content_gen','style_gen',
                            'regularization_p', 'regularization_r',
                            'ad_gen','dis_img_gen',
                            'ad_gen_v', 'dis_img_gen_v']

        self.visual_names = ['P_reference','BP_reference', 'P_frame_step','BP_frame_step','img_gen', 'flow_fields', 'masks']
        self.model_names = ['G','D','D_V']

        self.FloatTensor = torch.cuda.FloatTensor if len(self.gpu_ids)>0 \
            else torch.FloatTensor
        self.ByteTensor = torch.cuda.ByteTensor if len(self.gpu_ids)>0 \
            else torch.ByteTensor

        # define the Animation model
        self.net_G = network.define_g(opt, image_nc=opt.image_nc, structure_nc=opt.structure_nc, ngf=64, img_f=512,
                                      layers=opt.layers, num_blocks=2, use_spect=opt.use_spect_g, attn_layer=opt.attn_layer,
                                      norm='instance', activation='LeakyReLU', extractor_kz=opt.kernel_size)
        if len(opt.gpu_ids) > 1:
            self.net_G = torch.nn.DataParallel(self.net_G, device_ids=self.gpu_ids)

        self.flow2color = util.flow2color()

        self.net_D = network.define_d(opt, ndf=32, img_f=128, layers=4, use_spect=opt.use_spect_d)
        if len(opt.gpu_ids) > 1:
            self.net_D = torch.nn.DataParallel(self.net_D, device_ids=self.gpu_ids)

        input_nc = (opt.frames_D_V-1) * opt.image_nc
        self.net_D_V = network.define_d(opt, input_nc=input_nc, ndf=32, img_f=128, layers=4, use_spect=opt.use_spect_d)
        if len(opt.gpu_ids) > 1:
            self.net_D_V = torch.nn.DataParallel(self.net_D_V, device_ids=self.gpu_ids)

        if self.isTrain:
            # define the loss functions
            self.GANloss = external_function.AdversarialLoss(opt.gan_mode).to(opt.device)
            self.L1loss = torch.nn.L1Loss()
            self.L2loss = torch.nn.MSELoss()
            self.Correctness = external_function.PerceptualCorrectness().to(opt.device)
            self.Regularization = external_function.MultiAffineRegularizationLoss(kz_dic=opt.kernel_size).to(opt.device)
            self.Vggloss = external_function.VGGLoss().to(opt.device)

            # define the optimizer
            self.optimizer_G = torch.optim.Adam(itertools.chain(
                                               filter(lambda p: p.requires_grad, self.net_G.parameters())),
                                               lr=opt.lr, betas=(0.0, 0.999))
            self.optimizers.append(self.optimizer_G)

            self.optimizer_D = torch.optim.Adam(itertools.chain(
                                filter(lambda p: p.requires_grad, self.net_D.parameters()),
                                filter(lambda p: p.requires_grad, self.net_D_V.parameters())),
                                lr=opt.lr*opt.ratio_g2d, betas=(0.0, 0.999))
            self.optimizers.append(self.optimizer_D)
        else:
            self.results_dir_base = self.opt.results_dir
        self.setup(opt)

    def set_input(self, data):
        # move to GPU and change data types
        opt = self.opt
        _, n_frames_total, self.height, self.width = data['P'].size()  #
        self.n_frames_total = n_frames_total // opt.image_nc
        n_frames_load = opt.max_frames_per_gpu                        # number of total frames loaded into GPU at a time for each batch
        self.n_frames_load = min(n_frames_load, self.n_frames_total)

        if self.isTrain:
            self.P_reference  = data['P'][:,:opt.image_nc, ...].cuda()
            self.BP_reference = data['BP'][:, :opt.structure_nc, ...].cuda()
            self.P_previous = None
            self.BP_previous = None

        self.P_images = data['P']
        self.BP_structures = data['BP']
        self.image_paths = [path[0] for path in data['P_path']]
        self.change_seq=data['change_seq']

        if not self.isTrain:
            assert self.opt.batchSize == 1
            if data['frame_idx'] == self.opt.start_frame + self.opt.n_frames_pre_load_test:
                self.P_previous = None
                self.BP_previous = None
                self.P_reference  = data['P'][:,:opt.image_nc, ...].cuda()
                self.BP_reference = data['BP'][:, :opt.structure_nc, ...].cuda()

            # else:
            #     self.P_previous = self.test_generated
            #     self.BP_previous = self.test_BP_previous
            self.opt.results_dir = os.path.join(self.results_dir_base,
                                                self.image_paths[0].split('/')[-2])

    def write2video(self, name_list):
        images=[]
        for name in name_list:
            images.append(sorted(glob.glob(self.opt.results_dir+'/*_'+name+'.png')))

        image_array=[]
        for i in range(len(images[0])):
            cat_im=None
            for image_list in images:
                im = cv2.imread(image_list[i])
                if cat_im is not None:
                    cat_im = np.concatenate((cat_im, im), axis=1)
                else:
                    cat_im = im
            image_array.append(cat_im)

        res=''
        for name in name_list:
            res += (name +'_')
        out_name = self.opt.results_dir+'_'+res+'.mp4'
        print('write video %s'%out_name)
        height, width, layers = cat_im.shape
        size = (width,height)
        out = cv2.VideoWriter(out_name, cv2.VideoWriter_fourcc(*'mp4v'), 15, size)

        for i in range(len(image_array)):
            out.write(image_array[i])
        out.release()

    def get_current_visuals(self):
        """Return visualization images"""
        visual_ret = OrderedDict()
        for name in self.visual_names:
            if isinstance(name, str):
                value = getattr(self, name)
                if 'frame_step' in name:
                    value = value[0]
                    list_value=[]
                    for i in range(value.size(0)):
                        list_value.append(value[i].unsqueeze(0))
                    value=list_value
                if 'flow_field' in name or 'masks' in name:
                    list_value = [item for sub_list in value for item in sub_list]
                    value = list_value

                if isinstance(value, list):
                    # visual multi-scale ouputs
                    for i in range(len(value)):
                        visual_ret[name + str(i)] = self.convert2im(value[i], name)
                    # visual_ret[name] = util.tensor2im(value[-1].data)
                else:
                    visual_ret[name] =self.convert2im(value, name)
        return visual_ret

    def test(self, save_features=False, save_all=False, generate_edge=True):
        """Forward function used in test time"""
        # img_gen, flow_fields, masks = self.net_G(self.input_P1, self.input_BP1, self.input_BP2)
        height, width = self.height, self.width
        image_nc, structure_nc = self.opt.image_nc, self.opt.structure_nc
        n_frames_pre_load = self.opt.n_frames_pre_load_test

        self.BP_frame_step = self.BP_structures.view(-1, n_frames_pre_load, structure_nc, height, width).cuda()
        self.test_generated, self.flow_fields, self.masks, _ = self.net_G(self.BP_frame_step,
                                                                self.P_reference,
                                                                self.BP_reference,
                                                                self.P_previous,
                                                                self.BP_previous)
        self.P_previous = self.test_generated[-1]
        self.BP_previous = self.BP_frame_step[:,-1,... ]

        self.test_generated = torch.cat(self.test_generated, 0)
        self.save_results(self.test_generated, data_name='vis', data_ext='png')

        if generate_edge:
            value = self.BP_frame_step[:,:,0:1,...][0]
            value = (1-value)*2-1
            self.save_results(value, data_name='edge', data_ext='png')

        if self.change_seq:
            name_list=[] if not generate_edge else ['edge']
            name_list.append('vis')
            print(self.opt.results_dir)
            self.write2video(name_list)

    def update(self):
        """Run forward processing to get the inputs"""
        image_nc, structure_nc = self.opt.image_nc, self.opt.structure_nc
        n_frames_total, n_frames_load = self.n_frames_total, self.n_frames_load
        height, width = self.height, self.width
        for i in range(0, n_frames_total, n_frames_load):
            self.P_frame_step  = self.P_images[:, i*image_nc:(i+n_frames_load)*image_nc].cuda()
            self.BP_frame_step = self.BP_structures[:, i*structure_nc:(i+n_frames_load)*structure_nc].cuda()
            self.P_frame_step = self.P_frame_step.view(-1, n_frames_load, image_nc, height, width)
            self.BP_frame_step = self.BP_frame_step.view(-1, n_frames_load, structure_nc, height, width)

            self.img_gen, self.flow_fields, self.masks, self.P_previous_recoder = self.net_G(self.BP_frame_step,
                                                                    self.P_reference,
                                                                    self.BP_reference,
                                                                    self.P_previous,
                                                                    self.BP_previous)
            self.optimizer_D.zero_grad()
            self.backward_D()
            self.optimizer_D.step()

            self.optimizer_G.zero_grad()
            self.backward_G()
            self.optimizer_G.step()

            self.P_previous  = self.img_gen[-1].detach()
            self.BP_previous = self.BP_frame_step[:,-1,...].detach()

    def backward_D_basic(self, netD, real, fake):
        """Calculate GAN loss for the discriminator"""
        # Real
        D_real = netD(real)
        D_real_loss = self.GANloss(D_real, True, True)
        # fake
        D_fake = netD(fake.detach())
        D_fake_loss = self.GANloss(D_fake, False, True)
        # loss for discriminator
        D_loss = (D_real_loss + D_fake_loss) * 0.5
        # if print_loss:
            # print(D_real_loss)
            # print(D_fake_loss)
        # gradient penalty for wgan-gp
        if self.opt.gan_mode == 'wgangp':
            gradient_penalty, gradients = external_function.cal_gradient_penalty(netD, real, fake.detach())
            D_loss += gradient_penalty

        D_loss.backward()

        return D_loss

    def backward_D(self):
        """Calculate the GAN loss for the discriminators"""
        base_function._unfreeze(self.net_D)
        i = np.random.randint(len(self.img_gen))
        fake = self.img_gen[i]
        real = self.P_frame_step[:,i,...]
        self.loss_dis_img_gen = self.backward_D_basic(self.net_D, real, fake)

        base_function._unfreeze(self.net_D_V)
        i = np.random.randint(len(self.img_gen)-self.opt.frames_D_V+1)
        # fake = [self.img_gen[i]]
        # real = [self.P_frame_step[:,i,...]]
        fake = []
        real = []
        for frame in range(self.opt.frames_D_V-1):
            fake.append(self.img_gen[i+frame]-self.img_gen[i+frame+1])
            real.append(self.P_frame_step[:,i+frame,...]
                       -self.P_frame_step[:,i+frame+1,...])
        fake = torch.cat(fake, dim=1)
        real = torch.cat(real, dim=1)
        self.loss_dis_img_gen_v = self.backward_D_basic(self.net_D_V, real, fake)

    def backward_G(self):
        """Calculate training loss for the generator"""
        # gen_tensor = torch.cat([v.unsqueeze(1) for v in self.img_gen], 1)
        loss_style_gen, loss_content_gen, loss_app_gen=0,0,0

        for i in range(len(self.img_gen)):
            gen = self.img_gen[i]
            gt = self.P_frame_step[:,i,...]
            loss_app_gen += self.L1loss(gen, gt)

            content_gen, style_gen = self.Vggloss(gen, gt)
            loss_style_gen += style_gen
            loss_content_gen += content_gen

        self.loss_style_gen = loss_style_gen * self.opt.lambda_style
        self.loss_content_gen = loss_content_gen * self.opt.lambda_content
        self.loss_app_gen = loss_app_gen * self.opt.lambda_rec

        loss_correctness_p, loss_regularization_p=0, 0
        loss_correctness_r, loss_regularization_r=0, 0

        for i in range(len(self.flow_fields)):
            flow_field_i = self.flow_fields[i]
            flow_p, flow_r=[],[]
            for j in range(0, len(flow_field_i), 2):
                flow_p.append(flow_field_i[j])
                flow_r.append(flow_field_i[j+1])

            correctness_r = self.Correctness(self.P_frame_step[:,i,...], self.P_reference,
                                                    flow_r, self.opt.attn_layer)
            correctness_p = self.Correctness(self.P_frame_step[:,i,...], self.P_previous_recoder[i].detach(),
                                                    flow_p, self.opt.attn_layer)
            loss_correctness_p += correctness_p
            loss_correctness_r += correctness_r
            loss_regularization_p += self.Regularization(flow_p)
            loss_regularization_r += self.Regularization(flow_r)

        self.loss_correctness_p = loss_correctness_p * self.opt.lambda_correct
        self.loss_correctness_r = loss_correctness_r * self.opt.lambda_correct
        self.loss_regularization_p = loss_regularization_p * self.opt.lambda_regularization
        self.loss_regularization_r = loss_regularization_r * self.opt.lambda_regularization

        # rec loss fake
        base_function._freeze(self.net_D)
        i = np.random.randint(len(self.img_gen))
        fake = self.img_gen[i]
        D_fake = self.net_D(fake)
        self.loss_ad_gen = self.GANloss(D_fake, True, False) * self.opt.lambda_g

        ##########################################################################
        base_function._freeze(self.net_D_V)
        i = np.random.randint(len(self.img_gen)-self.opt.frames_D_V+1)
        # fake = [self.img_gen[i]]
        fake = []
        for frame in range(self.opt.frames_D_V-1):
            fake.append(self.img_gen[i+frame]-self.img_gen[i+frame+1])
        fake = torch.cat(fake, dim=1)
        D_fake = self.net_D_V(fake)
        self.loss_ad_gen_v = self.GANloss(D_fake, True, False) * self.opt.lambda_g
        ##########################################################################

        total_loss = 0
        for name in self.loss_names:
            if name != 'dis_img_gen_v' and name != 'dis_img_gen':
                total_loss += getattr(self, "loss_" + name)
        total_loss.backward()

    def optimize_parameters(self):
        self.update()


#!/usr/bin/env python

import torch
from torch.utils.data import DataLoader
from gradslam.datasets.icl import ICL
from gradslam.datasets.tum import TUM
from gradslam.slam.pointfusion import PointFusion
from gradslam import Pointclouds, RGBDImages
from gradslam.datasets.datautils import normalize_image
from time import time
from tqdm import tqdm
import numpy as np
# ROS
import rospy
from cv_bridge import CvBridge, CvBridgeError
from sensor_msgs.msg import Image, CameraInfo, PointCloud2
from nav_msgs.msg import Path
from geometry_msgs.msg import PoseStamped, TransformStamped
import message_filters
import tf2_ros
from ros_numpy import msgify, numpify
from tf.transformations import quaternion_from_matrix
import rospkg

def publish_tf_pose(pose, quat, child_frame_id, frame_id="world"):
    assert len(pose) == 3
    assert len(quat) == 4
    br = tf2_ros.TransformBroadcaster()
    t = TransformStamped()
    t.header.stamp = rospy.Time.now()
    t.header.frame_id = frame_id
    t.child_frame_id = child_frame_id
    t.transform.translation.x = pose[0]
    t.transform.translation.y = pose[1]
    t.transform.translation.z = pose[2]
    t.transform.rotation.x = quat[0]
    t.transform.rotation.y = quat[1]
    t.transform.rotation.z = quat[2]
    t.transform.rotation.w = quat[3]
    br.sendTransform(t)

def get_camera_info_msg(image_width, image_height, K,
                        frame_id="camera_frame",
                        distortion_model="plumb_bob"):
    camera_info_msg = CameraInfo()
    camera_info_msg.header.frame_id = frame_id
    camera_info_msg.header.stamp = rospy.Time.now()
    camera_info_msg.width = image_width
    camera_info_msg.height = image_height
    camera_info_msg.K = K
    camera_info_msg.distortion_model = distortion_model
    return camera_info_msg

class GradslamROS:
    def __init__(self,
                 dataset_path: str,
                 seqlen: float = 100,
                 odometry: str = 'gt',
                 width: int = 320,
                 height: int = 240):
        # select device
        self.width, self.height = width, height
        self.device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
        dataset = TUM(dataset_path, seqlen=seqlen, height=height, width=width)
        loader = DataLoader(dataset=dataset, batch_size=1)
        colors, depths, intrinsics, poses, *_ = next(iter(loader))
        self.rgbdimages = RGBDImages(colors, depths, intrinsics, poses, channels_first=False, device=self.device)
        self.slam = PointFusion(odom=odometry, dsratio=1, device=self.device)

        self.world_frame = 'world'
        self.route = Path()
        self.route.header.frame_id = self.world_frame
        self.route_pub = rospy.Publisher('~route', Path, queue_size=2)
        self.pc_pub = rospy.Publisher('~cloud', PointCloud2, queue_size=1)
        self.caminfo_pub = rospy.Publisher('~frustum', CameraInfo, queue_size=1)

    def run(self):
        # SLAM: step by step
        seq_len = self.rgbdimages.shape[1]
        initial_poses = torch.eye(4, device=self.device).view(1, 1, 4, 4)
        prev_frame = None
        pointclouds = Pointclouds(device=self.device)
        for s in tqdm(range(seq_len)):
            if rospy.is_shutdown():
                break
            t1 = time()
            live_frame = self.rgbdimages[:, s]
            # print(torch.min(live_frame.depth_image), torch.max(live_frame.depth_image), live_frame.depth_image.dtype)
            # if s == 0 and live_frame.poses is None:
            #     live_frame.poses = initial_poses

            pointclouds, live_frame.poses = self.slam.step(pointclouds, live_frame, prev_frame)
            rospy.logdebug(f'SLAM step took {(time() - t1):.3f} sec')
            prev_frame = live_frame if self.slam.odom != 'gt' else None

            # publish odometry / path
            assert live_frame.poses.shape == (1, 1, 4, 4)
            pose = PoseStamped()
            pose.header.frame_id = self.world_frame
            pose.header.stamp = rospy.Time.now()
            p = live_frame.poses[..., :3, 3].squeeze()
            pose.pose.position.x = p[0]
            pose.pose.position.y = p[1]
            pose.pose.position.z = p[2]
            q = quaternion_from_matrix(live_frame.poses[0, 0].cpu().numpy())
            pose.pose.orientation.x = q[0]
            pose.pose.orientation.y = q[1]
            pose.pose.orientation.z = q[2]
            pose.pose.orientation.w = q[3]
            # publish camera tf frame
            publish_tf_pose(pose=p, quat=q, child_frame_id='camera_frame')
            caminfo_msg = get_camera_info_msg(image_width=self.width, image_height=self.height,
                                              K=live_frame.intrinsics.squeeze(),
                                              frame_id='camera_frame')
            self.caminfo_pub.publish(caminfo_msg)

            self.route.poses.append(pose)
            self.route.header.stamp = rospy.Time.now()
            self.route_pub.publish(self.route)

            # publish point cloud
            n_pts = pointclouds.points_padded.shape[1]
            cloud = np.zeros((n_pts,), dtype=[('x', 'f4'), ('y', 'f4'), ('z', 'f4'),
                                              ('r', 'f4'), ('g', 'f4'), ('b', 'f4')])
            for i, f in enumerate(['x', 'y', 'z']):
                cloud[f] = pointclouds.points_padded[..., i].squeeze().cpu().numpy()
            for i, f in enumerate(['r', 'g', 'b']):
                cloud[f] = pointclouds.colors_padded[..., i].squeeze().cpu().numpy() / 255.
            pc_msg = msgify(PointCloud2, cloud)
            pc_msg.header.stamp = rospy.Time.now()
            pc_msg.header.frame_id = self.world_frame
            self.pc_pub.publish(pc_msg)

if __name__ == "__main__":
    rospy.init_node('gradslam_ros', log_level=rospy.INFO)
    dataset_path = '/home/ruslan/subt/thirdparty/gradslam/examples/tutorials/TUM/'
    odometry = rospy.get_param('~odometry')  # gt, icp, gradicp
    proc = GradslamROS(dataset_path=dataset_path, odometry=odometry, seqlen=500)
    try:
        proc.run()
    except KeyboardInterrupt:
        print("Shutting down")


import numpy as np
import os
import torch
from PIL import Image, ImageDraw, ImageFont
import h5py
import re
import scipy

from .proj_utils.local_utils import mkdirs, imresize_shape
from .proj_utils.torch_utils import to_torch, to_binary, to_numpy, roll

def test_gan(dataloader, save_root, model_folder, model_marker, netG, netEs, netEb, args):
    highest_res  = 64

    h5_filename = model_marker+'_align_%s' % args.align

    if args.shape_noise:
        h5_filename += '_shape_noise'
    if args.background_noise:
        h5_filename += '_background_noise'

    h5_filename += '.h5'

    save_h5      = os.path.join(save_root, h5_filename)
    org_h5path   = os.path.join(save_root, 'original.h5')

    ''' load model '''
    assert args.load_from_epoch != '', 'args.load_from_epoch is empty'
    G_weightspath  = os.path.join(model_folder, 'G_epoch{}.pth'.format(args.load_from_epoch))
    Es_weightspath = os.path.join(model_folder, 'Es_epoch{}.pth'.format(args.load_from_epoch))
    Eb_weightspath = os.path.join(model_folder, 'Eb_epoch{}.pth'.format(args.load_from_epoch))
    print('reload weights from {}'.format(G_weightspath))
    print('reload weights from {}'.format(Es_weightspath))
    print('reload weights from {}'.format(Eb_weightspath))

    netG.load_state_dict(torch.load(G_weightspath))
    netEs.load_state_dict(torch.load(Es_weightspath))
    netEb.load_state_dict(torch.load(Eb_weightspath))

    num_examples = len(dataloader.dataset)
    total_number = num_examples * args.test_sample_num

    all_choosen_caption = []
    org_file_not_exists = not os.path.exists(org_h5path)

    if org_file_not_exists:
        org_h5   = h5py.File(org_h5path,'w')
        org_dset = org_h5.create_dataset('output_{}'.format(highest_res),
                                                shape=(num_examples, 3, highest_res, highest_res),
                                                dtype=np.uint8)
        org_emb_dset = org_h5.create_dataset('embedding', shape=(num_examples, 50, 300), dtype=np.float) # embeddings
    else:
        org_dset = None
        org_emb_dset = None

    with h5py.File(save_h5,'w') as h5file:

        start_count = 0
        data_count = {}
        dset = {}
        gen_samples = []
        img_samples = []
        vis_samples = {}
        tmp_samples = {}
        init_flag = True

        netG.eval()
        netEs.eval()
        netEb.eval()

        for data in dataloader:
            test_images, _, segs, txt_data, txt_len, chosen_captions, saveIDs, classIDs = data

            segs = segs.cuda()
            test_images = test_images.cuda()
            txt_data = txt_data.cuda()

            test_bimages = test_images

            np_test_images = to_numpy(test_images)

            this_batch_size = np_test_images.shape[0]

            all_choosen_caption.extend(chosen_captions)
            if org_dset is not None:
                org_dset[start_count:start_count+this_batch_size] = ((np_test_images + 1) * 127.5 ).astype(np.uint8)
                org_emb_dset[start_count:start_count+this_batch_size] = txt_data.detach().cpu()

            start_count += this_batch_size

            for t in range(args.test_sample_num):

                # alignment
                if args.align == 'shape':
                    test_bimages = roll(test_images, 2, dim=0) # for text and seg mismatched backgrounds
                elif args.align == 'background':
                    segs = roll(segs, 1, dim=0) # for text mismatched segmentations
                elif args.align == 'all':
                    pass
                elif args.align == 'none':
                    test_bimages = roll(test_images, 2, dim=0) # for text and seg mismatched backgrounds
                    segs = roll(segs, 1, dim=0) # for text mismatched segmentations

                ''' Encode Segmentation'''
                segs_code = netEs(segs)
                bkgs_code = netEb(test_bimages)

                test_outputs = {}
                _, _, _, fake_images, _ = netG(txt_data, txt_len, segs_code, bkgs_code,
                                                                    shape_noise=args.shape_noise,
                                                                    background_noise=args.background_noise)

                test_outputs['output_64'] = fake_images

                if  t == 0:
                    if init_flag is True:
                        dset['saveIDs']   = h5file.create_dataset('saveIDs', shape=(total_number,), dtype=np.int64)
                        dset['classIDs']  = h5file.create_dataset('classIDs', shape=(total_number,), dtype=np.int64)
                        dset['embedding'] = h5file.create_dataset('embedding', shape=(total_number, 50, 300), dtype=np.float) # embedding
                        for k in test_outputs.keys():
                            vis_samples[k] = [None for i in range(args.test_sample_num + 1)] # +1 to fill real image
                            img_shape = test_outputs[k].size()[2::]

                            dset[k] = h5file.create_dataset(k, shape=(total_number,)+ img_shape + (3,), dtype=np.uint8)
                            data_count[k] = 0

                    init_flag = False

                for typ, img_val in test_outputs.items():
                    cpu_data = img_val.cpu().data.numpy()
                    row, col = cpu_data.shape[2],cpu_data.shape[3]
                    if t==0:
                        this_reshape = imresize_shape(np_test_images,  (row, col))
                        this_reshape = this_reshape * (2. / 255) - 1.

                        # this_reshape = this_reshape.transpose(0, 3, 1, 2)
                        vis_samples[typ][0] = this_reshape

                    vis_samples[typ][t+1] = cpu_data
                    bs = cpu_data.shape[0]

                    start = data_count[typ]
                    this_sample = ((cpu_data + 1) * 127.5 ).astype(np.uint8)
                    this_sample = this_sample.transpose(0, 2, 3, 1)

                    dset[typ][start: start + bs] = this_sample
                    dset['saveIDs'][start: start + bs] = saveIDs
                    dset['classIDs'][start: start + bs] = classIDs
                    dset['embedding'][start: start + bs] = txt_data.detach().cpu() # embedding
                    data_count[typ] = start + bs

            print('saved files [sample {}/{}]: '.format(start_count, num_examples), data_count)

        caption_array = np.array(all_choosen_caption, dtype=object)
        string_dt = h5py.special_dtype(vlen=str)
        h5file.create_dataset("captions", data=caption_array, dtype=string_dt)
        if org_dset is not None:
            org_h5.close()


# -----------------------------------------------------
# Person Search Dataset for Both Training and Testing
#
# Author: Liangqi Li
# Creating Date: Mar 28, 2018
# Latest rectified: Apr 16, 2018
# -----------------------------------------------------

import os
import os.path as osp
import time

import pandas as pd
import cv2
import PIL
import random
import numpy as np
from sklearn.metrics import average_precision_score, precision_recall_curve
import pickle
import matplotlib.pyplot as plt

def _compute_iou(a, b):
    x1 = max(a[0], b[0])
    y1 = max(a[1], b[1])
    x2 = min(a[2], b[2])
    y2 = min(a[3], b[3])
    inter = max(0, x2 - x1) * max(0, y2 - y1)
    union = (a[2] - a[0]) * (a[3] - a[1]) + \
            (b[2] - b[0]) * (b[3] - b[1]) - inter
    return inter * 1.0 / union

class PersonSearchDataset:

    def __init__(self, root_dir, split_name='train', gallery_size=100,
                 test_mode='gallery'):
        """
        create person search dataset
        ---
        param:
            root_dir: root direction for the raw dataset
            split_name: 'train' or 'test'
        """
        self.root_dir = root_dir
        self.split = split_name
        self.images_dir = osp.join(self.root_dir, 'Image/SSM')
        self.annotation_dir = osp.join(self.root_dir, 'annotation')
        self.cache_dir = osp.join(self.root_dir, '..', 'cache')

        self.train_imnames_file = 'trainImnamesSe.csv'
        self.train_imnamesDF_file = 'trainImnamesDF.csv'
        self.test_imnames_file = 'testImnamesSe.csv'
        self.train_all_file = 'trainAllDF.csv'
        self.test_all_file = 'testAllDF.csv'
        self.query_file = 'queryDF.csv'
        # self.gallery_sizes = [50, 100, 500, 1000, 2000, 4000]
        self.gallery_sizes = [200, 500, 1000, 2000, 4000]

        if not osp.exists(self.cache_dir):
            os.mkdir(self.cache_dir)

        if self.split == 'train':
            if self.train_imnamesDF_file in os.listdir(self.cache_dir) and \
                    self.train_all_file in os.listdir(self.cache_dir):
                self.train_imnames = pd.read_csv(
                    osp.join(self.cache_dir, self.train_imnamesDF_file))
                self.train_all = pd.read_csv(
                    osp.join(self.cache_dir, self.train_all_file))
            else:
                self.train_imnames, self.train_all = self.prepare_training()
            self.train_imnames_list = list(range(self.train_imnames.shape[0]))
            random.shuffle(self.train_imnames_list)  # shuffle the list
            self.train_imnames_list_equip = self.train_imnames_list[:]
            self.num_train_images = self.train_imnames.shape[0]

        elif self.split == 'test':
            self.test_imnames = pd.read_csv(
                osp.join(self.annotation_dir, self.test_imnames_file),
                header=None, squeeze=True)
            self.test_all = pd.read_csv(
                osp.join(self.annotation_dir, self.test_all_file))
            q_to_g_file = 'q_to_g' + str(gallery_size) + 'DF.csv'
            self.queries_to_galleries = pd.read_csv(
                osp.join(self.annotation_dir, q_to_g_file))
            self.query_boxes = pd.read_csv(osp.join(self.annotation_dir,
                                                    self.query_file))
            self.delta_to_coordinates()
            self.num_test_images = self.test_imnames.shape[0]
            self.test_imnames_list = list(range(self.num_test_images))[::-1]

            # TODO: split query and gallery completely
            self.test_mode = test_mode
            if self.test_mode == 'query':
                self.query_imnames_list = list(range(
                    self.queries_to_galleries.shape[0]))[::-1]

        else:
            raise KeyError(self.split)

    def delta_to_coordinates(self):
        """change `del_x` and `del_y` to `x2` and `y2` for testing set"""
        self.test_all['del_x'] += self.test_all['x1']
        self.test_all['del_y'] += self.test_all['y1']
        self.test_all.rename(columns={'del_x': 'x2', 'del_y': 'y2'},
                             inplace=True)
        self.query_boxes['del_x'] += self.query_boxes['x1']
        self.query_boxes['del_y'] += self.query_boxes['y1']
        self.query_boxes.rename(columns={'del_x': 'x2', 'del_y': 'y2'},
                                inplace=True)

    def prepare_training(self):
        """prepare dataset for training, including flipping, resizing"""

        def flip_bbox(subset, train_al, hei_and_wid):
            """flip bboxes only"""
            subset = pd.merge(subset, hei_and_wid)
            subset['x1'] = subset['width'] - train_al['x2'] - 1
            subset['x2'] = subset['width'] - train_al['x1'] - 1
            assert (subset['x2'] >= subset['x1']).all()
            del subset['width']
            del subset['height']
            subset['flipped'] = [1 for _ in range(subset.shape[0])]

            return subset

        train_imnames = pd.read_csv(osp.join(self.annotation_dir,
                                             self.train_imnames_file),
                                    header=None, squeeze=True)
        train_all = pd.read_csv(osp.join(self.annotation_dir,
                                         self.train_all_file))

        # filter out the outliers
        # outliers = [9378, 20121, 47308, 31635, 53597]
        outliers = [528, 9875, 2709, 5517, 5418, 1178, 181, 8804, 2600, 15348]
        for outlier in outliers:
            train_all.drop(outlier, inplace=True)
        train_all.index = range(train_all.shape[0])

        # formal
        widths = [PIL.Image.open(osp.join(self.images_dir, imname)).size[0]
                  for imname in train_imnames]
        heights = [PIL.Image.open(osp.join(self.images_dir, imname)).size[1]
                   for imname in train_imnames]
        heights_and_widths = pd.DataFrame({'imname': train_imnames,
                                           'width': widths,
                                           'height': heights})

        # # ==========================debug============================
        # heights_and_widths = pd.read_csv(
        #     osp.join(self.annotation_dir, 'heights_and_widthsDF.csv'))
        # # ===========================end=============================

        # change `del_x` and `del_y` to `x2` and `y2`
        train_all['del_x'] += train_all['x1']
        train_all['del_y'] += train_all['y1']
        train_all.rename(columns={'del_x': 'x2', 'del_y': 'y2'}, inplace=True)

        # horizontally flip bounding boxes
        flipped = [0 for _ in range(train_all.shape[0])]
        train_all['flipped'] = flipped
        train_add = train_all.copy()
        train_add = flip_bbox(train_add, train_all, heights_and_widths)
        train_all = pd.concat((train_all, train_add))
        train_all.index = range(train_all.shape[0])

        flipped = [0 for _ in range(train_imnames.size)]
        train_imnames_not = pd.DataFrame({'imname': train_imnames,
                                          'flipped': flipped})
        # change the order of columns
        train_imnames_not = train_imnames_not[['imname', 'flipped']]
        flipped = [1 for _ in range(train_imnames.size)]
        train_imnames_fl = pd.DataFrame({'imname': train_imnames,
                                         'flipped': flipped})
        train_imnames_fl = train_imnames_fl[['imname', 'flipped']]
        train_imnames = pd.concat((train_imnames_fl, train_imnames_not))
        train_imnames.index = range(train_imnames.shape[0])

        train_imnames.to_csv(osp.join(self.cache_dir, 'trainImnamesDF.csv'),
                             index=False)
        train_all.to_csv(osp.join(self.cache_dir, 'trainAllDF.csv'),
                         index=False)

        return train_imnames, train_all

    def next(self):
        if self.split == 'train':
            # TODO: change with cfg parameters
            target_size = 600
            max_size = 1000
            pixel_means = np.array([[[102.9801, 115.9465, 122.7717]]])

            # Prepare for the next epoch
            if len(self.train_imnames_list) == 0:
                self.train_imnames_list = self.train_imnames_list_equip[:]
                self.train_imnames = pd.read_csv(
                    osp.join(self.cache_dir, self.train_imnamesDF_file))

            chosen = self.train_imnames_list.pop()
            im_name, flipped = self.train_imnames.loc[chosen]

            im = cv2.imread(osp.join(self.images_dir, im_name))
            if flipped == 1:
                im = im[:, ::-1, :]
            im = im.astype(np.float32, copy=False)
            im -= pixel_means
            im_shape = im.shape
            im_size_min = np.min(im_shape[0:2])
            im_size_max = np.max(im_shape[0:2])
            im_scale = float(target_size) / float(im_size_min)
            # Prevent the biggest axis from being more than MAX_SIZE
            if np.round(im_scale * im_size_max) > max_size:
                im_scale = float(max_size) / float(im_size_max)
            im = cv2.resize(im, None, None, fx=im_scale, fy=im_scale,
                            interpolation=cv2.INTER_LINEAR)
            im = im[np.newaxis, :]  # add batch dimension

            df = self.train_all.copy()
            # TODO: use pandas.query
            df = df[df['imname'] == im_name]
            df = df[df['flipped'] == flipped]
            gt_boxes = df.loc[:, 'x1': 'pid']
            gt_boxes.loc[:, 'x1': 'y2'] *= im_scale
            gt_boxes = gt_boxes.values

            im_info = np.array([im.shape[1], im.shape[2], im_scale],
                               dtype=np.float32)

            return im, gt_boxes, im_info

        elif self.split == 'test':
            if self.test_mode == 'gallery':
                # TODO: change with cfg parameters
                target_size = 600
                max_size = 1000
                pixel_means = np.array([[[102.9801, 115.9465, 122.7717]]])

                chosen = self.test_imnames_list.pop()
                im_name = self.test_imnames.loc[chosen]

                im = cv2.imread(osp.join(self.images_dir, im_name))
                orig_shape = im.shape
                im = im.astype(np.float32, copy=True)
                im -= pixel_means
                im_shape = im.shape
                im_size_min = np.min(im_shape[0:2])
                im_size_max = np.max(im_shape[0:2])
                im_scale = float(target_size) / float(im_size_min)
                # Prevent the biggest axis from being more than MAX_SIZE
                if np.round(im_scale * im_size_max) > max_size:
                    im_scale = float(max_size) / float(im_size_max)
                im = cv2.resize(im, None, None, fx=im_scale, fy=im_scale,
                                interpolation=cv2.INTER_LINEAR)
                im = im[np.newaxis, :]  # add batch dimension

                im_info = np.array([im.shape[1], im.shape[2], im_scale],
                                   dtype=np.float32)

                return im, im_info, orig_shape

            elif self.test_mode == 'query':
                # TODO: change with cfg parameters
                target_size = 600
                max_size = 1000
                pixel_means = np.array([[[102.9801, 115.9465, 122.7717]]])

                chosen = self.query_imnames_list.pop()
                im_name = self.queries_to_galleries.iloc[chosen, 0]

                im = cv2.imread(osp.join(self.images_dir, im_name))
                im = im.astype(np.float32, copy=True)
                im -= pixel_means
                im_shape = im.shape
                im_size_min = np.min(im_shape[0:2])
                im_size_max = np.max(im_shape[0:2])
                im_scale = float(target_size) / float(im_size_min)
                # Prevent the biggest axis from being more than MAX_SIZE
                if np.round(im_scale * im_size_max) > max_size:
                    im_scale = float(max_size) / float(im_size_max)
                im = cv2.resize(im, None, None, fx=im_scale, fy=im_scale,
                                interpolation=cv2.INTER_LINEAR)
                im = im[np.newaxis, :]  # add batch dimension

                df = self.query_boxes.copy()
                gt_boxes = df.ix[chosen, 'x1': 'y2'] * im_scale
                gt_boxes = gt_boxes.as_matrix().astype(np.float64)

                im_info = np.array([im.shape[1], im.shape[2], im_scale],
                                   dtype=np.float32)

                return im, gt_boxes, im_info

            else:
                raise KeyError(self.test_mode)

        else:
            raise KeyError(self.split)

    def get_query_im(self, query_name, pid):
        target_size = 600
        max_size = 1000
        pixel_means = np.array([[[102.9801, 115.9465, 122.7717]]])

        im = cv2.imread(osp.join(self.images_dir, query_name))
        im = im.astype(np.float32, copy=False)
        im -= pixel_means
        im_shape = im.shape
        im_size_min = np.min(im_shape[0:2])
        im_size_max = np.max(im_shape[0:2])
        im_scale = float(target_size) / float(im_size_min)
        # Prevent the biggest axis from being more than MAX_SIZE
        if np.round(im_scale * im_size_max) > max_size:
            im_scale = float(max_size) / float(im_size_max)
        im = cv2.resize(im, None, None, fx=im_scale, fy=im_scale,
                        interpolation=cv2.INTER_LINEAR)
        im = im[np.newaxis, :]  # add batch dimension

        df = self.train_all.copy()
        # TODO: use pandas.query
        df = df[df['imname'] == query_name]
        df = df[df['flipped'] == False]
        df = df[df['pid'] == pid]
        assert df.shape[0] == 1
        gt_boxes = df.loc[:, 'x1': 'y2'] * im_scale
        gt_boxes = gt_boxes.as_matrix()

        im_info = np.array([im.shape[1], im.shape[2], im_scale],
                           dtype=np.float32)

        return im, gt_boxes, im_info

    def get_gallery_im(self, gallery_name, flipped):
        target_size = 600
        max_size = 1000
        pixel_means = np.array([[[102.9801, 115.9465, 122.7717]]])

        im = cv2.imread(osp.join(self.images_dir, gallery_name))
        if flipped:
            im = im[:, ::-1, :]
        im = im.astype(np.float32, copy=False)
        im -= pixel_means
        im_shape = im.shape
        im_size_min = np.min(im_shape[0:2])
        im_size_max = np.max(im_shape[0:2])
        im_scale = float(target_size) / float(im_size_min)
        # Prevent the biggest axis from being more than MAX_SIZE
        if np.round(im_scale * im_size_max) > max_size:
            im_scale = float(max_size) / float(im_size_max)
        im = cv2.resize(im, None, None, fx=im_scale, fy=im_scale,
                        interpolation=cv2.INTER_LINEAR)
        im = im[np.newaxis, :]  # add batch dimension

        df = self.train_all.copy()
        # TODO: use pandas.query
        df = df[df['imname'] == gallery_name]
        df = df[df['flipped'] == flipped]
        gt_boxes = df.loc[:, 'x1': 'pid']
        gt_boxes.loc[:, 'x1': 'y2'] *= im_scale
        gt_boxes = gt_boxes.values

        im_info = np.array([im.shape[1], im.shape[2], im_scale],
                           dtype=np.float32)

        return im, gt_boxes, im_info

    def get_neg_g_im(self, g_imnames):
        """Select a negative image that does not contain person of pid"""

        choose_from = set(self.train_imnames['imname']) - g_imnames
        im_name = random.choice(list(choose_from))
        flipped = random.choice([True, False])

        target_size = 600
        max_size = 1000
        pixel_means = np.array([[[102.9801, 115.9465, 122.7717]]])

        im = cv2.imread(osp.join(self.images_dir, im_name))
        if flipped:
            im = im[:, ::-1, :]
        im = im.astype(np.float32, copy=False)
        im -= pixel_means
        im_shape = im.shape
        im_size_min = np.min(im_shape[0:2])
        im_size_max = np.max(im_shape[0:2])
        im_scale = float(target_size) / float(im_size_min)
        # Prevent the biggest axis from being more than MAX_SIZE
        if np.round(im_scale * im_size_max) > max_size:
            im_scale = float(max_size) / float(im_size_max)
        im = cv2.resize(im, None, None, fx=im_scale, fy=im_scale,
                        interpolation=cv2.INTER_LINEAR)
        im = im[np.newaxis, :]  # add batch dimension

        df = self.train_all.copy()
        # TODO: use pandas.query
        df = df[df['imname'] == im_name]
        df = df[df['flipped'] == flipped]
        gt_boxes = df.loc[:, 'x1': 'pid']
        gt_boxes.loc[:, 'x1': 'y2'] *= im_scale
        gt_boxes = gt_boxes.values

        im_info = np.array([im.shape[1], im.shape[2], im_scale],
                           dtype=np.float32)

        return im, gt_boxes, im_info

    def test_image_path(self, i):
        return osp.join(self.images_dir, self.test_imnames.iloc[i])

    def __len__(self):
        if self.split == 'train':
            return self.num_train_images
        elif self.split == 'test':
            return self.num_test_images
        else:
            raise KeyError(self.split)

    def evaluate_detections(self, gallery_det, det_thresh=0.5, iou_thresh=0.5,
                            labeled_only=False):
        """evaluate the results of the detection"""
        assert self.test_imnames.shape[0] == len(gallery_det)

        y_true, y_score = [], []
        count_gt, count_tp = 0, 0
        df = self.test_all.copy()
        for k in range(len(gallery_det)):
            im_name = self.test_imnames.iloc[k]
            gt = df[df['imname'] == im_name]
            gt_boxes = gt.loc[:, 'x1': 'y2'].values
            if labeled_only:
                pass # TODO
            det = np.asarray(gallery_det[k])
            inds = np.where(det[:, 4].ravel() >= det_thresh)[0]
            det = det[inds]
            num_gt = gt_boxes.shape[0]
            num_det = det.shape[0]
            if num_det == 0:
                count_gt += num_gt
                continue
            ious = np.zeros((num_gt, num_det), dtype=np.float32)
            for i in range(num_gt):
                for j in range(num_det):
                    ious[i, j] = _compute_iou(gt_boxes[i], det[j, :4])
            tfmat = (ious >= iou_thresh)
            # for each det, keep only the largest iou of all the gt
            for j in range(num_det):
                largest_ind = np.argmax(ious[:, j])
                for i in range(num_gt):
                    if i != largest_ind:
                        tfmat[i, j] = False
            # for each gt, keep only the largest iou of all the det
            for i in range(num_gt):
                largest_ind = np.argmax(ious[i, :])
                for j in range(num_det):
                    if j != largest_ind:
                        tfmat[i, j] = False
            for j in range(num_det):
                y_score.append(det[j, -1])
                if tfmat[:, j].any():
                    y_true.append(True)
                else:
                    y_true.append(False)
            count_tp += tfmat.sum()
            count_gt += num_gt

        det_rate = count_tp * 1.0 / count_gt
        ap = average_precision_score(y_true, y_score) * det_rate
        precision, recall, __ = precision_recall_curve(y_true, y_score)
        recall *= det_rate

        plt.plot(recall, precision)
        plt.savefig('pr.jpg')

        print(
            '{} detection:'.format('labeled only' if labeled_only else 'all'))
        print('  recall = {:.2%}'.format(det_rate))
        if not labeled_only:
            print('  ap = {:.2%}'.format(ap))

    def evaluate_search(self, gallery_det, gallery_feat, probe_feat,
                        det_thresh=0.5, gallery_size=100, dump_json=None):
        assert self.test_imnames.shape[0] == len(gallery_det)
        assert self.test_imnames.shape[0] == len(gallery_feat)
        assert self.queries_to_galleries.shape[0] == len(probe_feat)

        use_full_set = gallery_size == -1
        df = self.test_all.copy()

        # ====================formal=====================
        name_to_det_feat = {}
        for name, det, feat in zip(self.test_imnames,
                                   gallery_det, gallery_feat):
            scores = det[:, 4].ravel()
            inds = np.where(scores >= det_thresh)[0]
            if len(inds) > 0:
                name_to_det_feat[name] = (det[inds], feat[inds])

        # # =====================debug=====================
        # f = open('name_to_det_feat.pkl', 'rb+')
        # name_to_det_feat = pickle.load(f)
        # # ======================end======================

        aps = []
        accs = []
        topk = [1, 5, 10]
        # ret  # TODO: save json
        for i in range(len(probe_feat)):
            pid = self.query_boxes.ix[i, 'pid']
            y_true, y_score = [], []
            imgs, rois = [], []
            count_gt, count_tp = 0, 0
            # Get L2-normalized feature vector
            feat_p = probe_feat[i].ravel()
            # Ignore the probe image
            start = time.time()
            probe_imname = self.queries_to_galleries.iloc[i, 0]
            # probe_roi = df[df['imname'] == probe_imname]
            # probe_roi = probe_roi[probe_roi['is_query'] == 1]
            # probe_roi = probe_roi[probe_roi['pid'] == pid]
            # probe_roi = probe_roi.loc[:, 'x1': 'y2'].as_matrix()
            probe_gt = []
            tested = set([probe_imname])
            # 1. Go through the gallery samples defined by the protocol
            for g_i in range(1, gallery_size + 1):
                gallery_imname = self.queries_to_galleries.iloc[i, g_i]
                # gt = df[df['imname'] == gallery_imname]
                # gt = gt[gt['pid'] == pid]  # important
                # gt = gt.loc[:, 'x1': 'y2']
                gt = df.query('imname==@gallery_imname and pid==@pid')
                gt = gt.loc[:, 'x1': 'y2'].as_matrix().ravel()
                count_gt += (gt.size > 0)
                # compute distance between probe and gallery dets
                if gallery_imname not in name_to_det_feat: continue
                det, feat_g = name_to_det_feat[gallery_imname]
                # get L2-normalized feature matrix NxD
                assert feat_g.size == np.prod(feat_g.shape[:2])
                feat_g = feat_g.reshape(feat_g.shape[:2])

                # # compute cosine similarities
                # sim = feat_g.dot(feat_p).ravel()

                # compute euclidean distance
                dist = np.square(feat_p - feat_g).sum(1)
                sim = -dist

                # assign label for each det
                label = np.zeros(len(sim), dtype=np.int32)
                if gt.size > 0:
                    w, h = gt[2] - gt[0], gt[3] - gt[1]
                    probe_gt.append({'img': str(gallery_imname),
                                     'roi': list(gt.astype('float'))})
                    iou_thresh = min(.5, (w * h * 1.0) / ((w + 10) * (h + 10)))
                    inds = np.argsort(sim)[::-1]
                    sim = sim[inds]
                    det = det[inds]
                    # only set the first matched det as true positive
                    for j, roi in enumerate(det[:, :4]):
                        if _compute_iou(roi, gt) >= iou_thresh:
                            label[j] = 1
                            count_tp += 1
                            break
                y_true.extend(list(label))
                y_score.extend(list(sim))
                imgs.extend([gallery_imname] * len(sim))
                rois.extend(list(det))
                tested.add(gallery_imname)
            # 2. Go through the remaining gallery images if using full set
            if use_full_set:
                pass  # TODO
            # 3. Compute AP for this probe (need to scale by recall rate)
            y_score = np.asarray(y_score)
            y_true = np.asarray(y_true)
            assert count_tp <= count_gt
            if count_gt == 0:
                print(probe_imname, i)
                break
            recall_rate = count_tp * 1.0 / count_gt
            ap = 0 if count_tp == 0 else average_precision_score(
                y_true, y_score) * recall_rate
            aps.append(ap)
            inds = np.argsort(y_score)[::-1]
            y_score = y_score[inds]
            y_true = y_true[inds]
            accs.append([min(1, sum(y_true[:k])) for k in topk])
            # compute time cost
            end = time.time()
            print('{}-th loop, cost {:.4f}s'.format(i, end - start))

        print('search ranking:')
        print('  mAP = {:.2%}'.format(np.mean(aps)))
        accs = np.mean(accs, axis=0)
        for i, k in enumerate(topk):
            print('  top-{:2d} = {:.2%}'.format(k, accs[i]))

if __name__ == '__main__':

    rtdir = '/Users/liliangqi/Desktop/person_search/dataset'
    size = 100
    sysu = PersonSearchDataset(rtdir, split_name='test', gallery_size=size)

    output_dir = './data/test_result'
    gboxes_file = open(os.path.join(output_dir, 'gboxes.pkl'), 'rb')
    gboxes = pickle.load(gboxes_file, encoding='iso-8859-1')
    gfeatures_file = open(os.path.join(output_dir, 'gfeatures.pkl'), 'rb')
    gfeatures = pickle.load(gfeatures_file, encoding='iso-8859-1')
    pfeatures_file = open(os.path.join(output_dir, 'pfeatures.pkl'), 'rb')
    pfeatures = pickle.load(pfeatures_file, encoding='iso-8859-1')
    # sysu.evaluate_detections(gboxes)
    sysu.evaluate_search(gboxes, gfeatures['feat'], pfeatures['feat'],
                         det_thresh=0.5, gallery_size=size, dump_json=None)
    gboxes_file.close()
    gfeatures_file.close()
    pfeatures_file.close()


import cv2
import numpy as np
import dlib
import os
class CaffeModel(object):
    def __init__(self):
        cwd = os.path.abspath(os.path.dirname(__file__))
        trainmodel = os.path.abspath(os.path.join(cwd, "facedetector/deploy.prototxt"))
        facedetect = os.path.abspath(os.path.join(cwd, "facedetector/res10_300x300_ssd_iter_140000_fp16.caffemodel"))
        self.net = cv2.dnn.readNetFromCaffe(trainmodel,
                               facedetect)

    def _analyze(self, inputframe):
        net = self.net
        frame = inputframe.copy()
        (h, w) = frame.shape[:2]
        blob = cv2.dnn.blobFromImage(cv2.resize(frame, (300, 300)),  # the fixed-size input image
                                     1.0,  # scaleFactor
                                     (300, 300),  # spatial size of the CNN
                                     (104.0, 177.0, 123.0))  # mean subtraction values

        # Pass the blob to the network
        net.setInput(blob)
        detections = net.forward()

        # Loop over the resultant detections
        for i in range(0, detections.shape[2]):
            # Extract the confidence (i.e., probability)
            confidence = detections[0, 0, i, 2]

            # Filter out weak detections by ensuring the confidence is greater than the threshold
            if confidence < 0.5:
                continue

            # Compute the (x, y)-coordinates of the bounding box
            box = detections[0, 0, i, 3:7] * np.array([w, h, w, h])
            (startX, startY, endX, endY) = box.astype('int')

            # Draw the bounding box of the face with the associated probability
            #text = '{:.2f}%'.format(confidence * 100)
            #y = startY - 10 if startY - 10 > 10 else startY + 10
            # cv2.rectangle(frame, (startX, startY), (endX, endY), (0, 0, 255), 2)
            # cv2.putText(frame, text, (startX, y), cv2.FONT_HERSHEY_SIMPLEX, 0.45, (0, 0, 255), 2)
            return dlib.rectangle(startX,startY, endX, endY)


import numpy as np
import mxnet as mx

class SigmoidCELossMetric(mx.metric.EvalMetric):
    def __init__(self, name, output_names, label_names):
        super().__init__(name, output_names, label_names)

    def update(self, labels, preds):
        self.sum_metric += preds[0].mean().asscalar()
        self.num_inst += 1


# coding=utf-8
# Copyright 2022 The Google Research Authors.
#
# 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
#
#     http://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.

"""Fit a volume with CLIP and diffusion guidance."""

# pylint: disable=g-bare-generic
# pylint: disable=g-doc-args
# pylint: disable=invalid-name
# pylint: disable=protected-access
# pylint: disable=redefined-outer-name
# pylint: disable=missing-class-docstring
# pylint: disable=missing-function-docstring

import argparse
import random
from typing import Tuple

import clip
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
import torchvision
from tqdm.auto import tqdm
from tqdm.auto import trange
import wandb

# pylint: disable=g-bad-import-order
# pylint: disable=g-multiple-import
# Data structures and functions for rendering.
from pytorch3d.renderer import (FoVPerspectiveCameras, look_at_view_transform,
                                RayBundle, VolumeRenderer, VolumeSampler,
                                NDCGridRaysampler, EmissionAbsorptionRaymarcher)
from pytorch3d.structures import Volumes

# Import GLIDE.
from glide_text2im.download import load_checkpoint
from glide_text2im.model_creation import create_model_and_diffusion
from glide_text2im.model_creation import model_and_diffusion_defaults
# pylint: enable=g-multiple-import

from coco_val_queries import queries
from dreamfields_voxel import clamp_and_detach
from dreamfields_voxel import compute_query_rank
from dreamfields_voxel import DEFAULT_VOLUME_SIZE
from dreamfields_voxel import load_clip
from dreamfields_voxel import render_rotating_volume
from dreamfields_voxel import render_validation_view
from dreamfields_voxel import sample_backgrounds
# pylint: enable=g-bad-import-order

parser = argparse.ArgumentParser()
# Logging.
parser.add_argument("--wandb-entity", type=str, default=None)
parser.add_argument("--wandb-project", type=str, default=None)
parser.add_argument("--exp-name-prefix", type=str, default="")
# Query and seed.
parser.add_argument("--query-idx", type=int)
parser.add_argument("--seed", type=int, default=0)
# Scene parameterization.
parser.add_argument("--volume-size", type=int, default=128)
parser.add_argument(
    "--sample-mode",
    type=str,
    default="bilinear",
    choices=["bilinear", "nearest", "bicubic"],
    help="RGB and density interpolation mode for"
    "torch.nn.functional.grid_sample")
parser.add_argument("--voxel-features-dim", type=int, default=32)
parser.add_argument("--mlp-hidden-dim", type=int, default=64)
# Data augmentation
parser.add_argument("--n-aug", type=int, default=8)
parser.add_argument("--batch-size", type=int, default=20)
parser.add_argument("--crop-scale-range", default=(0.5, 1.0))
parser.add_argument("--nsq", type=int, default=8)
parser.add_argument("--bg-blur-std-range", default=(0., 10.))
# Optimization.
parser.add_argument("--loss-model", type=str, default="ViT-B/16")
parser.add_argument("--retrieve-model", type=str, default="ViT-B/32")
parser.add_argument("--lr", type=float, default=0.1)
parser.add_argument("--n-iter", type=int, default=5000)
parser.add_argument("--target-transmittance", type=float, default=0.95)
parser.add_argument("--transmittance-lam", type=float, default=1.0)
parser.add_argument("--density-tvd-lam", type=float, default=0.)
parser.add_argument("--color-tvd-lam", type=float, default=0.)
parser.add_argument("--clip-lam", type=float, default=1.0)
# GLIDE.
parser.add_argument("--guidance-scale", type=float, default=5.0)
parser.add_argument(
    "--t-schedule", type=str, choices=["anneal", "uniform"], default="anneal")
parser.add_argument("--t-respace", type=int, default=100)
parser.add_argument("--t-max", type=int, default=40)  # initial
parser.add_argument("--t-min", type=int, default=1)  # final
parser.add_argument("--denoise-n-steps", type=int, default=1)
parser.add_argument("--denoise-stop-grad", action="store_true")
parser.add_argument("--denoise-augmented", action="store_true")
parser.add_argument("--diffusion-lam", type=float, default=0.1)
# Camera.
# Our rendered scene is centered around (0,0,0)
# and is enclosed inside a bounding box
# whose side is roughly equal to 2.0 (world units).
parser.add_argument("--volume-extent-world", type=float, default=2.0)
parser.add_argument("--camera-radius", type=float, default=4.0)
# render_size describes the size of both sides of the
# rendered images in pixels.
parser.add_argument("--render-size", type=int, default=168)
parser.add_argument("--n-pts-per-ray", type=int, default=192)
args = parser.parse_args()

class FeatureVolumeSampler(VolumeSampler):
  """Sample a batch of volumes `Volumes` at points along projection rays."""

  def __init__(self,
               volumes,
               sample_mode = "bilinear",
               densities_dim = 1,
               output_dim = 3,
               features_dim = args.voxel_features_dim,
               hidden_dim = args.mlp_hidden_dim):
    """Feature volume sampler.

    Args:
      volumes: An instance of the `Volumes` class representing a batch of
        volumes that are being rendered.
      sample_mode: Defines the algorithm used to sample the volumetric voxel
        grid. Can be either "bilinear" or "nearest".
    """
    super().__init__(volumes, sample_mode)

    self.output_dim = output_dim
    self.densities_dim = densities_dim
    self.features_dim = features_dim

    self.output = nn.Linear(features_dim, output_dim)

  def forward(self, ray_bundle,
              **kwargs):
    """Samples `self._volumes` at the respective 3D ray-points.

    Args:
      ray_bundle: A RayBundle object with the following fields:
        rays_origins_world: A tensor of shape `(minibatch, ..., 3)` denoting
          the origins of the sampling rays in world coords.
        rays_directions_world: A tensor of shape `(minibatch, ..., 3)`
          containing the direction vectors of sampling rays in world coords.
        rays_lengths: A tensor of shape `(minibatch, ...,
          num_points_per_ray)` containing the lengths at which the rays are
          sampled.

    Returns:
      rays_densities: A tensor of shape
        `(minibatch, ..., num_points_per_ray, opacity_dim)` containing the
        density vectors sampled from the volume at the locations of the
        ray points.
      rays_features: A tensor of shape
        `(minibatch, ..., num_points_per_ray, feature_dim)` containing
        the feature vectors sampled from the volume at the locations of the
        ray points.
    """
    rays_densities, rays_features = super().forward(ray_bundle, **kwargs)

    # project to output feature dimension
    rays_features = self.output(rays_features)

    return rays_densities, rays_features

class FeatureVolumeRenderer(VolumeRenderer):
  """A class for rendering a batch of Volumes.

  The class should be initialized with a raysampler and a raymarcher class
  which both have to be a `Callable`.
  """

  def forward(self, cameras, volumes,
              **kwargs):
    """Render images using raymarching over rays cast through input `Volumes`.

    Args:
      cameras: A batch of cameras that render the scene. A `self.raysampler`
        takes the cameras as input and samples rays that pass through the
        domain of the volumetric function.
      volumes: An instance of the `Volumes` class representing a batch of
        volumes that are being rendered.

    Returns:
      images: A tensor of shape `(minibatch, ..., (feature_dim +
      opacity_dim)`
          containing the result of the rendering.
      ray_bundle: A `RayBundle` containing the parametrizations of the
          sampled rendering rays.
    """
    volumetric_function = FeatureVolumeSampler(
        volumes, sample_mode=self._sample_mode)
    volumetric_function = volumetric_function.to(volumes.device)
    return self.renderer(
        cameras=cameras, volumetric_function=volumetric_function, **kwargs)

class VolumeModel(torch.nn.Module):

  def __init__(self,
               renderer,
               volume_size=DEFAULT_VOLUME_SIZE,
               voxel_size=0.1,
               features_dim=args.voxel_features_dim):
    super().__init__()
    # Densities close to zero.
    self.log_densities = torch.nn.Parameter(-4.0 * torch.ones(1, *volume_size))
    self.features = torch.nn.Parameter(torch.randn(features_dim, *volume_size))
    self._voxel_size = voxel_size
    # Store the renderer module as well.
    self._renderer = renderer

  def forward(self, cameras, custom_renderer=None):
    batch_size = cameras.R.shape[0]

    # Convert the log-space values to the densities/colors.
    densities = torch.sigmoid(self.log_densities)

    # Instantiate the Volumes object, making sure the densities and colors are
    # correctly expanded batch_size-times.
    volumes = Volumes(
        densities=densities[None].expand(batch_size, *self.log_densities.shape),
        features=self.features[None].expand(batch_size, *self.features.shape),
        voxel_size=self._voxel_size,
    )

    # Given cameras and volumes, run the renderer
    # and return only the first output value (the 2nd output is a representation
    # of the sampled rays which can be omitted for our purpose).
    if custom_renderer is not None:
      return custom_renderer(cameras=cameras, volumes=volumes)[0]

    rendered = self._renderer(cameras=cameras, volumes=volumes)[0]
    rgb, silhouette = rendered.split([3, 1], dim=-1)
    return rgb, silhouette, densities

def get_renderer(resolution, n_pts_per_ray):
  ## Initialize the volumetric renderer
  # The following initializes a volumetric renderer that emits a ray from each
  # pixel of a target image and samples a set of uniformly-spaced points along
  # the ray. At each ray-point, the corresponding density and color value is
  # obtained by querying the corresponding location in the volumetric model of
  # the scene (the model is described & instantiated in a later cell).

  # The renderer is composed of a *raymarcher* and a *raysampler*.
  # - The *raysampler* is responsible for emitting rays from image pixels and
  # sampling the points along them. Here, we use the `NDCGridRaysampler` which
  # follows the standard PyTorch3D coordinate grid convention (+X from right to
  # left; +Y from bottom to top; +Z away from the user).
  # - The *raymarcher* takes the densities and colors sampled along each ray and
  # renders each ray into a color and an opacity value of the ray's source
  # pixel. Here we use the `EmissionAbsorptionRaymarcher` which implements the
  # standard Emission-Absorption raymarching algorithm.

  # Next we instantiate a volumetric model of the scene. This quantizes the 3D
  # space to cubical voxels, where each voxel is described with a 3D vector
  # representing the voxel's RGB color and a density scalar which describes the
  # opacity of the voxel (ranging between [0-1], the higher the more opaque).

  # In order to ensure the range of densities and colors is between [0-1], we
  # represent both volume colors and densities in the logarithmic space. During
  # the forward function of the model, the log-space values are passed through
  # the sigmoid function to bring the log-space values to the correct range.

  # Additionally, `VolumeModel` contains the renderer object. This object stays
  # unaltered throughout the optimization.

  # 1) Instantiate the raysampler.
  # Here, NDCGridRaysampler generates a rectangular image
  # grid of rays whose coordinates follow the PyTorch3D
  # coordinate conventions.
  # Since we use a volume of size 128^3, we sample n_pts_per_ray=150,
  # which roughly corresponds to a one ray-point per voxel.
  # We further set the min_depth=0.1 since there is no surface within
  # 0.1 units of any camera plane.

  # Changing the rendering resolution is a bit involved.
  raysampler = NDCGridRaysampler(
      image_width=resolution,
      image_height=resolution,
      n_pts_per_ray=n_pts_per_ray,
      min_depth=args.camera_radius - args.volume_extent_world * np.sqrt(3) / 2,
      max_depth=args.camera_radius + args.volume_extent_world * np.sqrt(3) / 2,
  )

  # 2) Instantiate the raymarcher.
  # Here, we use the standard EmissionAbsorptionRaymarcher
  # which marches along each ray in order to render
  # each ray into a single 3D color vector
  # and an opacity scalar.
  raymarcher = EmissionAbsorptionRaymarcher()

  # Finally, instantiate the volumetric render
  # with the raysampler and raymarcher objects.

  renderer = FeatureVolumeRenderer(
      raysampler=raysampler,
      raymarcher=raymarcher,
      sample_mode=args.sample_mode)
  return renderer

def load_diffusion(name, device):
  """Create base model."""
  options = model_and_diffusion_defaults()
  options["inpaint"] = False
  options["use_fp16"] = has_cuda
  options["timestep_respacing"] = str(
      args.t_respace)  # Use few diffusion steps for fast sampling.
  model, diffusion = create_model_and_diffusion(**options)
  model.eval()
  if has_cuda:
    model.convert_to_fp16()
  model.to(device)
  model.load_state_dict(load_checkpoint(name, device))
  print("total base parameters", sum(x.numel() for x in model.parameters()))
  return model, diffusion, options

if __name__ == "__main__":
  ## 1. Setup logging.
  query = queries[args.query_idx]
  exp_name = f"{args.exp_name_prefix}_{args.query_idx:03d}_{query}"
  wandb.init(
      entity=args.wandb_entity,
      project=args.wandb_project,
      name=exp_name,
      config=args)

  # Reproducibility.
  torch.manual_seed(args.seed)
  random.seed(args.seed)
  np.random.seed(args.seed)

  ## 2. Obtain the utilized device.
  if torch.cuda.is_available():
    device = torch.device("cuda:0")
    torch.cuda.set_device(device)
    has_cuda = True
  else:
    device = torch.device("cpu")
    has_cuda = False

  ## Initialize CLIP.
  image_mean = torch.tensor((0.48145466, 0.4578275, 0.40821073), device=device)
  image_mean = image_mean[None, :, None, None]
  image_std = torch.tensor((0.26862954, 0.26130258, 0.27577711), device=device)
  image_std = image_std[None, :, None, None]

  if args.clip_lam:
    model, preprocess, clip_size = load_clip(args.loss_model, device)

  if args.retrieve_model == args.loss_model and args.clip_lam:
    test_model, test_preprocess, test_clip_size = model, preprocess, clip_size
  else:
    test_model, test_preprocess, test_clip_size = load_clip(
        args.retrieve_model, device)

  ## Initialize GLIDE.
  # Create base model.
  base_glide_model, diffusion, base_glide_options = load_diffusion(
      "base", device)
  base_glide_model.eval()

  ## Initialize the volumetric model.
  renderer = get_renderer(args.render_size, args.n_pts_per_ray)

  # Here we carry out the volume fitting with differentiable rendering.

  # Instantiate the volumetric model.
  volume_model = VolumeModel(
      renderer,
      volume_size=[args.volume_size] * 3,
      voxel_size=args.volume_extent_world / args.volume_size,
  ).to(device)

  def get_parameters():
    return volume_model.parameters()

  # Instantiate the Adam optimizer.
  optimizer = torch.optim.Adam(get_parameters(), lr=args.lr)

  ## Embed the target caption with CLIP.
  if args.clip_lam:
    query_tok = clip.tokenize(query).to(device)
    z_clip = model.encode_text(query_tok).detach()
    z_clip = F.normalize(z_clip, dim=-1)

  ## Embed the target caption with GLIDE.
  denoise_batch_size = (
      args.n_aug *
      args.batch_size if args.denoise_augmented else args.batch_size)
  tokens = base_glide_model.tokenizer.encode(query)
  tokens, mask = base_glide_model.tokenizer.padded_tokens_and_mask(
      tokens, base_glide_options["text_ctx"])

  # Create the classifier-free guidance tokens (empty).
  full_batch_size = denoise_batch_size * 2
  uncond_tokens, uncond_mask = (
      base_glide_model.tokenizer.padded_tokens_and_mask(
          [], base_glide_options["text_ctx"]))

  # Pack the tokens together into model kwargs.
  base_model_kwargs = dict(
      tokens=torch.tensor(
          [tokens] * denoise_batch_size + [uncond_tokens] * denoise_batch_size,
          device=device),
      mask=torch.tensor(
          [mask] * denoise_batch_size + [uncond_mask] * denoise_batch_size,
          dtype=torch.bool,
          device=device,
      ),
  )

  # Create an classifier-free guidance sampling function.
  def base_model_fn(x_t, ts, **kwargs):
    half = x_t[:len(x_t) // 2]
    combined = torch.cat([half, half], dim=0)
    model_out = base_glide_model(combined, ts, **kwargs)
    eps, rest = model_out[:, :3], model_out[:, 3:]
    cond_eps, uncond_eps = torch.split(eps, len(eps) // 2, dim=0)
    half_eps = uncond_eps + args.guidance_scale * (cond_eps - uncond_eps)
    eps = torch.cat([half_eps, half_eps], dim=0)
    return torch.cat([eps, rest], dim=1)

  def preprocess_glide(x, order="NHWC"):
    if order == "NHWC":
      # x is [NHWC]. Reshape to NCHW.
      x = x.movedim(-1, 1)
    x = x * 2 - 1  # Scale from [0, 1] to [-1, 1].
    x = F.interpolate(x, (64, 64), mode="bilinear")
    return x

  def unprocess_glide(x):
    # x is [NCHW] Reshape to [NHWC].
    x = x.movedim(1, -1)
    x = (x + 1) / 2  # Scale from [-1, 1] to [0, 1].
    return x

  denoised_fn = lambda x_start: x_start

  glide_context_manager = (
      torch.no_grad if args.denoise_stop_grad else torch.enable_grad)

  ## We do Adam iterations and sample 10 random images in each minibatch.
  if args.clip_lam:
    clip_aug_fn = torchvision.transforms.RandomResizedCrop(
        clip_size, scale=args.crop_scale_range, ratio=(1.0, 1.0))
  denoise_aug_fn = torchvision.transforms.RandomResizedCrop(
      64, scale=args.crop_scale_range, ratio=(1.0, 1.0))

  pbar = trange(1, args.n_iter + 1)
  for iteration in pbar:
    metrics = {}
    visualize_images = iteration % 100 == 0 or iteration == 1

    # In case we reached the last 75% of iterations,
    # decrease the learning rate of the optimizer 10-fold.
    if iteration == round(args.n_iter * 0.75):
      tqdm.write("Decreasing LR 10-fold ...")
      for g in optimizer.param_groups:
        g["lr"] = args.lr * 0.1

    # Zero the optimizer gradient.
    optimizer.zero_grad()

    # Get a batch of viewing angles.
    elev = torch.rand(args.batch_size) * 10 + 20.
    azim = torch.rand(args.batch_size) * 360
    radius_mult = 1 + 0.1 * (torch.rand(args.batch_size) - 0.5)
    radius = radius_mult * args.camera_radius / 1.2

    R, T = look_at_view_transform(dist=radius, elev=elev, azim=azim)
    batch_cameras = FoVPerspectiveCameras(device=device, R=R, T=T)

    # Evaluate the volumetric model.
    rendered_images, rendered_silhouettes, densities = volume_model(
        batch_cameras)

    # Transmittance loss.
    tau = 1 - (1 - args.target_transmittance) / (radius_mult**2)
    tau = tau.to(device)
    avg_transmittance = 1 - rendered_silhouettes.mean(
        dim=tuple(range(1, rendered_silhouettes.ndim)))
    transmittance_loss = -torch.min(tau, avg_transmittance).mean()

    if (args.diffusion_lam > 0 and args.denoise_augmented) or args.clip_lam > 0:
      # NOTE: this background is at the render resolution, not the resize.
      # Generate random backgrounds.
      bgs = sample_backgrounds(
          args.n_aug * args.batch_size, args.render_size, device=device)

      # Composite renderings with backgrounds.
      bgs = bgs.view(args.n_aug, args.batch_size, *bgs.shape[1:])  # A,N,C,H,W
      bgs = bgs.movedim(2, -1)  # ANCHW to ANHWC
      composite_images = (
          rendered_silhouettes[None] * rendered_images[None] +
          (1 - rendered_silhouettes[None]) * bgs)
      composite_images = composite_images.reshape(  # to A*N,H,W,C
          args.n_aug * args.batch_size, args.render_size, args.render_size, 3)
      composite_images = composite_images.movedim(3, 1)  # NHWC to NCHW.

    # Compute GLIDE loss.
    # Sample from the base model.
    if args.diffusion_lam:
      base_glide_model.del_cache()
      if args.denoise_augmented:
        denoise_aug_images = denoise_aug_fn(composite_images)
        x = preprocess_glide(denoise_aug_images, order="NCHW")
      else:
        # White background.
        x = rendered_silhouettes * rendered_images + 1 - rendered_silhouettes
        x = preprocess_glide(x, order="NHWC")

      inp = x.detach() if args.denoise_stop_grad else x

      assert args.t_schedule == "anneal"
      progress = (iteration - 1) / (args.n_iter - 1)
      start_t = int(args.t_max - (args.t_max - args.t_min) * progress)
      diffuse_steps = []
      denoise_steps = []
      with glide_context_manager():
        for t in torch.linspace(
            start_t, args.t_min, args.denoise_n_steps, dtype=int):
          inp = diffusion.q_sample(
              inp, torch.tensor([t] * denoise_batch_size, device=device))
          diffuse_steps.append(inp.detach().to("cpu", non_blocking=True))
          inp = diffusion.p_sample(
              base_model_fn,
              x=torch.cat([inp, inp], dim=0),
              t=torch.tensor([t] * full_batch_size, device=device),
              clip_denoised=True,
              model_kwargs=base_model_kwargs,
              cond_fn=None,
              denoised_fn=denoised_fn,
          )["pred_xstart"][:denoise_batch_size]
          denoise_steps.append(inp.detach().to("cpu", non_blocking=True))
      base_glide_model.del_cache()
      diffusion_loss = F.mse_loss(x, inp)
      metrics["diffusion/start_t"] = start_t
      metrics["loss/diffusion_mse"] = diffusion_loss
    else:
      diffusion_loss = torch.tensor([0.], device=device)

    # Compute the CLIP loss.
    if args.clip_lam:
      clip_aug_images = clip_aug_fn(composite_images)
      x = preprocess(clip_aug_images)  # Resize and normalize.
      z_est = model.encode_image(x)
      z_est = F.normalize(z_est, dim=-1)
      clip_loss = -torch.sum(z_est * z_clip, dim=-1).mean()
    else:
      clip_loss = torch.tensor([0.], device=device)

    # Compute total loss and take an optimization step.
    loss = (
        args.clip_lam * clip_loss +
        args.transmittance_lam * transmittance_loss +
        args.diffusion_lam * diffusion_loss)

    if args.density_tvd_lam:
      loss.backward()

      with torch.no_grad():
        d = densities.squeeze(0)
        dc = torch.cat([d, d.transpose(1, 2), d.transpose(0, 2)], axis=-1)
        density_tvd_loss = torch.square(dc[Ellipsis, :-1] - dc[Ellipsis, 1:]).sum()
        loss = loss + args.density_tvd_lam * density_tvd_loss
        metrics["loss/density_tvd"] = density_tvd_loss.item()

      grad_sigmoid = lambda x: torch.exp(-x) / torch.square(1 + torch.exp(-x))
      log_d = volume_model.log_densities
      grad = log_d.grad
      grad_in = grad_sigmoid(log_d.detach())  # d log(density) / d density.
      s = args.density_tvd_lam * 2  # gradient below assumes scale is 1/2.

      d = densities.detach().squeeze(0)

      grad[0, :-1] += s * (d[:-1] - d[1:]) * grad_in[0, :-1]
      grad[0, 1:] += s * (d[1:] - d[:-1]) * grad_in[0, 1:]

      grad[0, :, :-1] += s * (d[:, :-1] - d[:, 1:]) * grad_in[0, :, :-1]
      grad[0, :, 1:] += s * (d[:, 1:] - d[:, :-1]) * grad_in[0, :, 1:]

      grad[0, :, :, :-1] += s * (d[:, :, :-1] -
                                 d[:, :, 1:]) * grad_in[0, :, :, :-1]
      grad[0, :, :,
           1:] += s * (d[:, :, 1:] - d[:, :, :-1]) * grad_in[0, :, :, 1:]
    else:
      loss.backward()

    optimizer.step()

    ## Logging.
    with torch.no_grad():
      metrics.update({
          "loss/total_loss": loss.item(),
          "loss/clip": clip_loss.item(),
          "loss/transmittance": transmittance_loss.item(),
      })

      # Print the current values of the losses.
      if iteration % 10 == 0:
        avg_transmittance_float = float(avg_transmittance.mean().item())
        pbar.set_description(
            f"Iteration {iteration:05d}:" +
            f" clip_loss = {float(clip_loss.item()):1.2f}" +
            f" diffusion_loss = {float(diffusion_loss.item()):1.5f}" +
            f" avg transmittance = {avg_transmittance_float:1.2f}" +
            f" tau = {str(tau)}")

      # Visualize the renders every 100 iterations.
      if visualize_images:
        wandb_image = lambda x: wandb.Image(clamp_and_detach(x))

        # Visualize only a single randomly selected element of the batch.
        im_show_idx = int(torch.randint(low=0, high=args.batch_size, size=(1,)))

        rendering = rendered_images[im_show_idx]
        silhouette = rendered_silhouettes[im_show_idx]
        metrics["render/rendered"] = wandb.Image(clamp_and_detach(rendering))

        if args.clip_lam > 0:
          aug_image = clip_aug_images[im_show_idx].movedim(0, 2)
          metrics["render/augmented"] = wandb_image(aug_image)

        render_white_bg = rendering * silhouette + 1 - silhouette
        metrics["render/rendered_white_bg"] = wandb_image(render_white_bg)
        metrics["render/rendered_silhouettes"] = wandb_image(
            silhouette.squeeze(-1))

        diffused = torch.stack(
            diffuse_steps, dim=1)  # [batch_size, n_steps, C, 64, 64].
        diffused = unprocess_glide(
            diffused[im_show_idx])  # [n_steps, 64, 64, C]
        diffused = diffused.movedim(0, 1).reshape(64, args.denoise_n_steps * 64,
                                                  3)
        metrics["render/diffused"] = wandb_image(diffused)

        denoised = torch.stack(
            denoise_steps, dim=1)  # [batch_size, n_steps, C, 64, 64].
        denoised = unprocess_glide(
            denoised[im_show_idx])  # [n_steps, 64, 64, C]
        denoised = denoised.movedim(0, 1).reshape(64, args.denoise_n_steps * 64,
                                                  3)
        metrics["render/denoised"] = wandb_image(denoised)

      # Validate from a held-out view.
      if iteration % 250 == 0 or iteration == 1:
        validation_view = render_validation_view(
            volume_model, test_clip_size, device=device)
        metrics["val/render"] = wandb.Image(clamp_and_detach(validation_view))

        rank, cosine_sim = compute_query_rank(
            test_model,
            test_preprocess,
            render=validation_view.movedim(-1, 0).unsqueeze(0),
            query=query,
            queries_r=queries,
            device=device)

        metrics["val/rank"] = rank
        metrics["val/acc"] = int(rank == 0)
        metrics["val/cosine_sim"] = cosine_sim

      if iteration % 500 == 0 or iteration == 1:
        # Visualize the optimized volume by rendering from multiple viewpoints
        # that rotate around the volume's y-axis.
        rotating_volume_frames = render_rotating_volume(
            volume_model, n_frames=240, device=device)
        metrics["render/video"] = wandb.Video(
            (rotating_volume_frames * 255.).astype(np.uint8),
            fps=30,
            format="mp4")

      wandb.log(metrics, iteration)


"""
This is super complex, but prints out hello world.
"""
import numpy as np
import xarray as xr
import climpred as cp
def main():
    print("CONFLICT!")
    hello = 'hello'
    world = 'world'
    print(f'{hello} {world}: this is extra stuff.')
    X = 3
    Y = 3
    print("tons of stuff")
    np.array([3])

if __name__ == '__main__':
    main()


# -*- coding: utf-8 -*-
from __future__ import print_function
import click
import os
import re
import api as face_recognition
import multiprocessing
import itertools
import sys
import PIL.Image
import numpy as np
import cv2

def scan_known_people(known_people_folder):
	"""Scans the dataset containing images of known people."""
	known_names = []
	known_face_encodings = []

	for file in image_files_in_folder(known_people_folder):
		basename = os.path.splitext(os.path.basename(file))[0]
		img = face_recognition.load_image_file(file)
		encodings = face_recognition.face_encodings(img)

		if len(encodings) > 1:
			click.echo("WARNING: More than one face found in {}. Only considering the first face.".format(file))

		if len(encodings) == 0:
			click.echo("WARNING: No faces found in {}. Ignoring file.".format(file))
		else:
			known_names.append(basename)
			known_face_encodings.append(encodings[0])

	return known_names, known_face_encodings

def test_image(known_names, known_face_encodings):
	"""Compares the captured image to the images in the dataset."""
	cam = cv2.VideoCapture(0)
	ret, image_to_check = cam.read()

	cam.release()
	cv2.destroyAllWindows()

	unknown_image = image_to_check

	# Scale down image if it's giant so things run a little faster
	if max(unknown_image.shape) > 1600:
		pil_img = PIL.Image.fromarray(unknown_image)
		pil_img.thumbnail((1600, 1600), PIL.Image.LANCZOS)
		unknown_image = np.array(pil_img)

	unknown_encodings = face_recognition.face_encodings(unknown_image)

	for unknown_encoding in unknown_encodings:
		distances = face_recognition.face_distance(known_face_encodings, unknown_encoding)
		result = list(distances <= 0.5625) #0.5625 is tolerance

		if True in result:
			continue
		else:
			# No match
			sys.exit()

	if not unknown_encodings:
		# No faces were found in image
		sys.exit()

def image_files_in_folder(folder):
	"""Return a list containing the path of the images in the folder."""
	return [os.path.join(folder, f) for f in os.listdir(folder) if re.match(r'.*\.(jpg|jpeg|png)', f, flags=re.I)]

def main():
	known_people_folder = "C:/dataset"
	known_names, known_face_encodings = scan_known_people(known_people_folder)
	test_image(known_names, known_face_encodings)

if __name__ == "__main__":
	main()