import torch
import torch.nn.functional as F
from mmdet.datasets import build_dataloader, DATASETS
from mmdet.models import clip
from tqdm import tqdm
import matplotlib.pyplot as plt
import os
from mmcv.ops import roi_align
import argparse


CLASSES = ('person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus',
               'train', 'truck', 'boat', 'traffic light', 'fire hydrant',
               'stop sign', 'parking meter', 'bench', 'bird', 'cat', 'dog',
               'horse', 'sheep', 'cow', 'elephant', 'bear', 'zebra', 'giraffe',
               'backpack', 'umbrella', 'handbag', 'tie', 'suitcase', 'frisbee',
               'skis', 'snowboard', 'sports ball', 'kite', 'baseball bat',
               'baseball glove', 'skateboard', 'surfboard', 'tennis racket',
               'bottle', 'wine glass', 'cup', 'fork', 'knife', 'spoon', 'bowl',
               'banana', 'apple', 'sandwich', 'orange', 'broccoli', 'carrot',
               'hot dog', 'pizza', 'donut', 'cake', 'chair', 'couch',
               'potted plant', 'bed', 'dining table', 'toilet', 'tv', 'laptop',
               'mouse', 'remote', 'keyboard', 'cell phone', 'microwave',
               'oven', 'toaster', 'sink', 'refrigerator', 'book', 'clock',
               'vase', 'scissors', 'teddy bear', 'hair drier', 'toothbrush')


SEEN_CLASSES = ('truck', 'traffic light', 'fire hydrant', 'stop sign', 'parking meter', 'bench',
                'bear', 'zebra', 'backpack', 'umbrella', 'tie', 'suitcase', 'frisbee', 'skis', 'sports ball',
                'kite', 'baseball bat', 'baseball glove', 'skateboard', 'cup', 'knife',
                'spoon', 'apple', 'sandwich', 'broccoli', 'hot dog', 'pizza', 'donut', 'bed', 'toilet',
                'laptop', 'mouse', 'keyboard', 'cell phone', 'microwave', 'toaster',
                'sink', 'book', 'vase', 'toothbrush')


parser = argparse.ArgumentParser()
parser.add_argument('--dataset_type', type=str, default='CocoCaptionDataset')
parser.add_argument('--prefix', type=str, default='captions')
parser.add_argument('--subset', type=str, default='val')
parser.add_argument('--with_bbox', type=int, default=0)
parser.add_argument('--img_affine', type=bool, default=False)
parser.add_argument('--seen_classes', type=str, default=None)
parser.add_argument('--suffix', type=str, default='coco_caption')
parser.add_argument('--color', default='r')
parser.add_argument('--clip_model', type=str, default='ViT-B/32')
parser.add_argument('--embeddings', type=str, default=None)
args = parser.parse_args()


def generate_label_embeddings(labels_training, labels_testing, model, save_file):
    labels_training_prompts = [f'A photo of a {label}' for label in labels_training]
    labels_testing_prompts = [f'A photo of a {label}' for label in labels_testing]

    embeddings_training = model.encode_text(clip.tokenize(labels_training_prompts).cuda())
    embeddings_testing = model.encode_text(clip.tokenize(labels_testing_prompts).cuda())
    print(embeddings_training.shape)
    torch.save(dict(embeddings_training=embeddings_training,
                    embeddings_testing=embeddings_testing), save_file)


def visualize_bbox_scores(model, data_loader, save_dir, suffix='all', color='r', embeddings=''):
    embeddings = torch.load(embeddings)['embeddings_testing'].cuda()
    embeddings = embeddings / embeddings.norm(dim=-1, keepdim=True)
    scale = model.visual.input_resolution
    with torch.no_grad():
        scores = []
        logits = []
        correct = []
        correct_5 = []
        for data in tqdm(data_loader):
            img = data['img'].data[0].cuda()
            gt_label = data['gt_labels'].data[0]
            gt_bboxes = data['gt_bboxes'].data[0]
            gt_label = torch.cat(gt_label, dim=0).long().cuda()
            # text_feature = embeddings[torch.cat(gt_label, dim=0).long()].cuda()
            rois = []
            for batch_id, bbox in enumerate(gt_bboxes):
                roi = bbox.new_ones(bbox.shape[0], 5) * batch_id
                roi[:, 1:] = bbox
                rois.append(roi)
            rois = torch.cat(rois, dim=0).cuda()
            img_crops = roi_align(img, rois, (scale, scale), 1.0, 1)
            img_feature = model.encode_image(img_crops, True)
            score = (img_feature[:, None] * embeddings[None]).sum(-1)
            logits.append(score[range(score.shape[0]), gt_label])

            score = F.softmax(score, dim=-1)
            _, preds = torch.topk(score, k=1, dim=-1)
            _, preds_5 = torch.topk(score, k=5, dim=-1)
            correct.append((preds[:, 0] == gt_label).float())
            correct_5.append((preds_5 == gt_label[:, None]).sum(-1).float())
            # print(score)
            scores.append(score[range(score.shape[0]), gt_label])

        scores = torch.cat(scores)
        correct = torch.cat(correct)
        correct_5 = torch.cat(correct_5)
        logits = torch.cat(logits)

        print(scores.mean(), flush=True)
        print(correct.mean(), flush=True)
        print(logits.mean(), flush=True)
    scores = scores.view(-1).cpu().numpy()
    correct_5 = correct_5.view(-1).cpu().numpy()
    correct = correct.view(-1).cpu().numpy()
    logits = logits.view(-1).cpu().numpy()
    acc = correct.sum() / len(logits)
    acc_5 = correct_5.sum() / len(logits)
    print(acc, acc_5)
    # plt.subplot(121)
    plt.hist(logits, bins=1000, color=color)
    plt.title(f'scores_{suffix}')
    # plt.subplot(122)
    # plt.hist(logits, bins=1000, color=color)
    # plt.title(f'acc_{acc}_scores_{logits.mean().item()}_{suffix}')
    plt.savefig(os.path.join(save_dir, f'bbox_scores_{suffix}.png'))
    plt.close()


def get_dataloader_bbox(data_root, dataset_type, prefix='captions', subset='val',
                   with_bbox=True, seen_classes=None):
    dataset_class = DATASETS.get(dataset_type)
    # dataset settings

    img_norm_cfg = dict(
        mean=[122.7709383, 116.7460125, 104.09373615],
        std=[68.5005327, 66.6321579, 70.32316305], to_rgb=True)
    load_anns = \
        dict(type='LoadOpenAnnotations',
             with_bbox=False,
             with_img_id=False,
             with_label=False,
             with_unseen=False,
             with_ann_id=False)
    meta_keys = ['filename', 'caption']
    if with_bbox:
        load_anns.update(dict(
            with_bbox=True,
            with_img_id=True,
            with_label=True,
            with_unseen=True,
            with_ann_id=True))
        meta_keys.extend(['gt_bboxes', 'gt_labels', 'ann_ids', 'img_id'])
    resize = dict(type='Resize', img_scale=(1333, 800), keep_ratio=True)

    pipeline = [
        dict(type='LoadImageFromFile'),
        dict(type='LoadAnnotations', with_bbox=True),
        dict(type='Resize', img_scale=(1333, 800), keep_ratio=True),
        dict(type='Normalize', **img_norm_cfg),
        dict(type='Pad', size_divisor=32),
        dict(type='DefaultFormatBundle'),
        dict(type='Collect', keys=['img', 'gt_bboxes', 'gt_labels'],
             meta_keys=['filename'])]

    dataset_cfg = dict(
        ann_file=data_root + f'annotations/{prefix}_{subset}2017.json',
        img_prefix=data_root + f'{subset}2017/',
        pipeline=pipeline,
        test_mode=False)
    if seen_classes is not None:
        dataset_cfg.update(dict(seen_classes=seen_classes))

    dataset = dataset_class(**dataset_cfg)

    return build_dataloader(dataset,
                            seed=None, dist=False, shuffle=False,
                            drop_last=False, workers_per_gpu=4,
                            samples_per_gpu=32)


if __name__ == '__main__':

    device = "cuda" if torch.cuda.is_available() else "cpu"
    model, _ = clip.load(args.clip_model, device=device, download_root='models', use_text_encoder=False)
    model.init_weights()
    model.eval()
    input_scale = model.visual.input_resolution
    print(input_scale)

    data_root = r'data/coco/'
    save_dir = r'data/coco'
    os.makedirs(save_dir, exist_ok=True)

    # dataset_type = 'CocoOpenDataset'
    # prefix = 'instances'
    dataset_type = args.dataset_type
    prefix = args.prefix
    subset = args.subset
    with_bbox = args.with_bbox
    seen_classes = args.seen_classes
    img_affine = args.img_affine
    suffix = args.suffix
    color = args.color

    data_loader = get_dataloader_bbox(data_root, dataset_type, prefix=prefix,
                                 subset=subset, with_bbox=with_bbox,
                                 seen_classes=seen_classes)
    visualize_bbox_scores(model, data_loader, save_dir, suffix=suffix, color=color,  embeddings=args.embeddings)

    # generate_label_embeddings(SEEN_CLASSES, CLASSES, model,
    #                           os.path.join(save_dir, 'class_embeddings_vitB16.pt'))
