"""
****************** COPYRIGHT AND CONFIDENTIALITY INFORMATION ******************
Copyright (c) 2018 [Thomson Licensing]
All Rights Reserved
This program contains proprietary information which is a trade secret/business \
secret of [Thomson Licensing] and is protected, even if unpublished, under \
applicable Copyright laws (including French droit d'auteur) and/or may be \
subject to one or more patent(s).
Recipient is to retain this program in confidence and is not permitted to use \
or make copies thereof other than as permitted in a written agreement with \
[Thomson Licensing] unless otherwise expressly allowed by applicable laws or \
by [Thomson Licensing] under express agreement.
Thomson Licensing is a company of the group TECHNICOLOR
*******************************************************************************
This scripts permits one to reproduce training and experiments of:
    Engilberge, M., Chevallier, L., Pérez, P., & Cord, M. (2018, April).
    Finding beans in burgers: Deep semantic-visual embedding with localization.
    In Proceedings of CVPR (pp. 3984-3993)

Author: Martin Engilberge
"""

import argparse
import re
import time

import numpy as np
from numpy.__config__ import show
import torch


from misc.model import img_embedding, joint_embedding
from torch.utils.data import DataLoader, dataset

from misc.dataset import TextDataset
from misc.utils import collate_fn_cap_padded
from torch.utils.data import DataLoader
from misc.utils import load_obj 
from misc.evaluation import recallTopK 

from misc.utils import show_imgs
import sys 
from misc.dataset import TextEncoder 

device = torch.device("cuda")
# device = torch.device("cpu") # uncomment to run with cpu

if __name__ == '__main__':

    parser = argparse.ArgumentParser(description='Extract embedding representation for images')
    parser.add_argument("-p", '--path', dest="model_path", help='Path to the weights of the model to evaluate')
    parser.add_argument("-d", '--data', dest="data_path", help='path to the file containing the sentence to embed')
    parser.add_argument("-bs", "--batch_size", help="The size of the batches", type=int, default=1)

    args = parser.parse_args()

    print("Loading model from:", args.model_path)
    checkpoint = torch.load(args.model_path, map_location=lambda storage, loc: storage)

    join_emb = joint_embedding(checkpoint['args_dict'])
    join_emb.load_state_dict(checkpoint["state_dict"])

    for param in join_emb.parameters():
        param.requires_grad = False

    join_emb.to(device)
    join_emb.eval()

    encoder = TextEncoder()
    print("Loading model done")
    # (4) design intersection mode.
    print("Please input your description of the image that you wanna search >>>")
    for line in sys.stdin:
        
        t0 = time.time() 
        cap_str = line.strip() 
        # with open(args.data_path, 'w') as cap_file:
        #     cap_file.writelines(cap_str)
        t1 = time.time()
        print("text is embedding ...")
        dataset = torch.Tensor(encoder.encode(cap_str)).unsqueeze(dim=0)
        t111 = time.time()
        dataset_loader = DataLoader(dataset, batch_size=args.batch_size, num_workers=1, pin_memory=True, collate_fn=collate_fn_cap_padded)
        t11 = time.time() 
        caps_enc = list()
        for i, (caps, length) in enumerate(dataset_loader, 0):
            input_caps = caps.to(device)
            with torch.no_grad():
                _, output_emb = join_emb(None, input_caps, length)
            caps_enc.append(output_emb.cpu().data.numpy())

        t12 = time.time() 
        caps_stack = np.vstack(caps_enc)
        # print(t11 - t1, t12 - t11, t111 - t1)
        
        t2 = time.time()
        print("recall from resources ...")
        # (1) load candidate imgs from saved embeding pkl file. 
        imgs_emb_file_path = "/home/atticus/proj/matching/DSVE/imgs_embed/v20210915_01_9408/allImg"
        # imgs_emb(40775, 2400)
        imgs_emb, imgs_path = load_obj(imgs_emb_file_path)
        # (2) calculate the sim between cap and imgs. 
        # (3) rank imgs and display the searching result. 
        recall_imgs = recallTopK(caps_stack, imgs_emb, imgs_path, ks=5)

        t3 = time.time() 
        show_imgs(imgs_path=recall_imgs)
        
        # print("input stage time: {} \n text embedding stage time: {} \n recall stage time: {}".format(t1 - t0, t2 - t1, t3 - t2))

        print("======== current epoch done ========")
        print("Please input your description of the image that you wanna search >>>")