# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
# SPDX-License-Identifier: CC-BY-NC-4.0


import io
import torch
import jsonlines
import numpy as np
from tqdm import tqdm
from pathlib import Path


def xywh2xyxy(bbox):
    bbox = bbox.copy()
    bbox[2] += bbox[0]
    bbox[3] += bbox[1]
    return bbox


def get_refcoco_refId_to_boxIdx(path):
    xs = np.load(path, allow_pickle=True)
    id2box = dict()
    for x in xs:
        img_id = int(Path(x['file_name']).stem.split('_')[-1])
        id2box[img_id] = dict()
        for idx, rid in enumerate(x['ref_ids']):
            id2box[img_id][rid] = idx
    return id2box


def add_global_vilbert_feats(feats, bboxs, input_torch=False):
    if input_torch:
        g_feat = feats.sum(dim=0) / feats.size(0)
        g_bbox = torch.tensor([0, 0, 1, 1, 1]).float()
        feats = torch.cat([g_feat.unsqueeze(0), feats], dim=0)
        bboxs = torch.cat([g_bbox.unsqueeze(0), bboxs], dim=0)
    else:
        g_feat = np.sum(feats, axis=0) / len(feats)
        g_bbox = np.array([0, 0, 1, 1, 1])
        feats = np.concatenate([np.expand_dims(g_feat, axis=0), feats], axis=0)
        bboxs = np.concatenate([np.expand_dims(g_bbox, axis=0), bboxs], axis=0)
    return feats, bboxs


def bbox2spatial_vilbert(bbox, image_width, image_height, mode='xywh'):
    if mode == 'xywh':
        bbox = xywh2xyxy(bbox)
    image_location = list(bbox).copy()
    image_location += [
        (image_location[3] - image_location[1])
        * (image_location[2] - image_location[0])
        / (float(image_width) * float(image_height))
    ]
    # image_location_ori = image_location.copy()
    image_location[0] /= float(image_width)
    image_location[1] /= float(image_height)
    image_location[2] /= float(image_width)
    image_location[3] /= float(image_height)
    return image_location


def bbox2spatial_gw(bbox, image_width, image_height, mode='xywh'):
    if mode == 'xywh':
        x_width = bbox[2]
        y_height = bbox[3]
    elif mode == 'xyxy':
        x_width = bbox[2] - bbox[0]
        y_height = bbox[3] - bbox[1]

    x_left = bbox[0]
    y_upper = bbox[1]
    x_right = x_left + x_width
    y_lower = y_upper + y_height

    x_center = x_left + 0.5 * x_width
    y_center = y_upper + 0.5 * y_height

    # Rescale features fom -1 to 1

    x_left = (1.*x_left / image_width) * 2 - 1
    x_right = (1.*x_right / image_width) * 2 - 1
    x_center = (1.*x_center / image_width) * 2 - 1

    y_lower = (1.*y_lower / image_height) * 2 - 1
    y_upper = (1.*y_upper / image_height) * 2 - 1
    y_center = (1.*y_center / image_height) * 2 - 1

    x_width = (1.*x_width / image_width) * 2
    y_height = (1.*y_height / image_height) * 2

    # Concatenate features
    spatial = [
        x_left, y_upper, x_right, y_lower,
        x_center, y_center, x_width, y_height
    ]
    return spatial


def load_fasttext(fname):
    fin = io.open(fname, 'r', encoding='utf-8', newline='\n', errors='ignore')
    num_words, dim = map(int, fin.readline().split())
    data = {}
    # cnt = 0
    for line in tqdm(fin):
        tokens = line.rstrip().split(' ')
        data[tokens[0]] = np.array(list(map(float, tokens[1:])))
        # cnt += 1
        # if cnt >= 40000: break
    return num_words, dim, data


def load_glove(fname, only_vocab=False, norm=True):
    with open(fname, 'r') as f:
        ivocab = {}
        vocab = {}
        vectors = []
        print("[INFO] Loading glove vectors %s..." % ("(only vocab)" if only_vocab else ""))
        for line in tqdm(f):
            vals = line.rstrip().split(' ')
            wrd, vec = vals[0], vals[1:]
            vocab[wrd] = len(vocab)
            ivocab[len(ivocab)] = wrd
            if not only_vocab:
                vec = np.array(list(map(float, vec)))
                if norm:
                    vec = vec / np.linalg.norm(vec)
                vectors.append(vec)
            # if len(vocab) > 100: break
    if only_vocab:
        return vocab, ivocab
    else:
        return vocab, ivocab, np.array(vectors)


def load_mscoco_obj_vocab(path):
    mscoco_obj_vocab = dict()
    mscoco_obj_vocab[0] = ''
    with open(path) as f:
        for i, line in enumerate(f, start=1):
            obj_name = line.strip() 
            mscoco_obj_vocab[i] = obj_name
    return mscoco_obj_vocab


def load_rcnn_obj_vocab(path):
    rcnn_obj_vocab = dict()
    rcnn_obj_vocab[0] = 'other'
    with open(path) as f:
        for i, line in enumerate(f, start=1):
            obj_name = line.strip()
            if ',' in obj_name:
                obj_name = obj_name.split(',')[0]
            rcnn_obj_vocab[i] = obj_name
    return rcnn_obj_vocab


def get_mscoco_id(path):
    return int(Path(path).stem.split('_')[-1])


def get_mscoco_idmap(mscoco_dir):
    imgs = Path(mscoco_dir).rglob("*.jpg")
    idmap = dict()
    for path in imgs:
        idmap[int(path.stem.split('_')[-1])] = path
    return idmap


def load_all_gt_games(annotation_path, verbose=False, debug=False):
    if verbose:
        print("[INFO] Loading %s..." % annotation_path)
    games = []
    cnt = 0
    with jsonlines.open(annotation_path) as reader:
        iterator = tqdm(reader) if verbose else reader
        for annotation in iterator:
            if debug and cnt >= 403: break
            g = Game.from_annotation(annotation)
            games.append(g)
            cnt += 1
    return games


class Game(object):
    def __init__(self, game_id, object_id, image_info, objects, qas, status):
        self.id = game_id
        self.object_id = object_id
        self.image_info = image_info
        self.image_id = image_info['id']
        self.image_height = image_info['height']
        self.image_width = image_info['width']
        self.objects = objects
        self.qas = qas
        self.status = status
        tgt_index, bboxs, cats = -1, [], []
        for index, obj in enumerate(self.objects):
            bboxs.append(obj['bbox'])
            cats.append(obj['category_id'])
            if obj['id'] == self.object_id:
                assert tgt_index == -1, "Too many target object in game %d." % self.id
                tgt_index = index
        assert tgt_index != -1, "There is no target object in game %d" % self.id
        self.target_index = tgt_index
        self.bboxs = bboxs
        self.categories = cats
        self.target_category = cats[tgt_index]

    @classmethod
    def from_annotation(cls, annotation):
        return cls(
            game_id=annotation['id'],
            object_id=annotation['object_id'],
            image_info=annotation['image'],
            objects=annotation['objects'],
            qas=annotation['qas'],
            status=annotation['status'])