# Copyright (c) 2023, NVIDIA Corporation & Affiliates. All rights reserved.
#
# This work is made available under the Nvidia Source Code License-NC.
# To view a copy of this license, visit
# https://github.com/NVlabs/prismer/blob/main/LICENSE

import torch
import os
import PIL.Image as Image
import numpy as np
import cv2
import clip
import pickle as pk
try:
    import ruamel_yaml as yaml
except ModuleNotFoundError:
    import ruamel.yaml as yaml

from experts.model_bank import load_expert_model
from experts.ocr_detection.generate_dataset import Dataset
from accelerate import Accelerator
from tqdm import tqdm


model, transform = load_expert_model(task='ocr_detection')
accelerator = Accelerator(mixed_precision='fp16')
pca_clip = pk.load(open('dataset/clip_pca.pkl', 'rb'))

config = yaml.load(open('configs/experts.yaml', 'r'), Loader=yaml.Loader)
data_path = config['data_path']
save_path = os.path.join(config['save_path'], 'ocr_detection')

batch_size = 32
dataset = Dataset(data_path, transform)
data_loader = torch.utils.data.DataLoader(
    dataset=dataset,
    batch_size=batch_size,
    shuffle=False,
    num_workers=4,
    pin_memory=True,
)

clip_model, _ = clip.load("ViT-L/14", device=accelerator.device)
model, data_loader = accelerator.prepare(model, data_loader)


def get_label(w, h, word_instances):
    word_lists = []
    final_mask = np.ones([h, w],  dtype=np.uint8) * 255
    counter = 0
    for word_instance in word_instances[::-1]:
        mask = np.zeros([h ,w])
        mask = cv2.fillPoly(mask, [np.int32(word_instance.word_bbox.reshape(-1, 2))], 1)
        text = word_instance.text.lower()
        final_mask[mask > 0] = counter
        word_lists.append(text)
        counter += 1
    return final_mask, word_lists


with torch.no_grad():
    for i, (test_data, image_path, scale_w, scale_h, original_w, original_h) in enumerate(tqdm(data_loader)):
        word_instance_lists = model(test_data, scale_w, scale_h, original_w, original_h)
        for k in range(len(word_instance_lists)):
            word_instance = word_instance_lists[k]
            if len(word_instance) == 0:
                continue
            else:
                final_mask, word_lists = get_label(original_w[k], original_h[k], word_instance)

            final_mask = Image.fromarray(final_mask)
            img_path_split = image_path[k].split('/')
            ps = image_path[k].split('.')[-1]
            im_save_path = os.path.join(save_path, img_path_split[-3], img_path_split[-2])
            os.makedirs(im_save_path, exist_ok=True)

            final_mask.save(os.path.join(im_save_path, img_path_split[-1].replace(f'.{ps}', '.png')))

            if len(word_lists) > 0:
                word_embed = clip.tokenize(word_lists).to(accelerator.device)
                word_features = pca_clip.transform(clip_model.encode_text(word_embed).float().cpu())
                word_lists = {j: {'features': torch.from_numpy(word_features[j]).float(),
                                  'text': word_lists[j]} for j in range(len(word_lists))}
                torch.save(word_lists, os.path.join(im_save_path, img_path_split[-1].replace(f'.{ps}', '.pt')))