import json, os, random, math
from collections import defaultdict

import torch
from torch.utils.data import Dataset
import torchvision.transforms as T

import numpy as np
import PIL
from skimage.transform import resize as imresize
import pycocotools.mask as mask_utils
from random import shuffle
from data.preprocess import *
from PIL import Image
class CustomDataset(Dataset):
    def __init__(self, image_dir, instances_json, stuff_json=None,
                 stuff_only=True, image_size=(64, 64), mask_size=16,
                 normalize_images=False, max_samples=None,
                 include_relationships=True, min_object_size=0.02,
                 min_objects_per_image=3, max_objects_per_image=25, left_right_flip=False,
                 include_other=False, instance_whitelist=None, stuff_whitelist=None, data_mode='train', classes=None):
        """
        A PyTorch Dataset for loading Coco and Coco-Stuff annotations and converting
        them to scene graphs on the fly.
    
        Inputs:
        - image_dir: Path to a directory where images are held
        - instances_json: Path to a JSON file giving COCO annotations
        - stuff_json: (optional) Path to a JSON file giving COCO-Stuff annotations
        - stuff_only: (optional, default True) If True then only iterate over
          images which appear in stuff_json; if False then iterate over all images
          in instances_json.
        - image_size: Size (H, W) at which to load images. Default (64, 64).
        - mask_size: Size M for object segmentation masks; default 16.
        - normalize_image: If True then normalize images by subtracting ImageNet
          mean pixel and dividing by ImageNet std pixel.
        - max_samples: If None use all images. Other wise only use images in the
          range [0, max_samples). Default None.
        - include_relationships: If True then include spatial relationships; if
          False then only include the trivial __in_image__ relationship.
        - min_object_size: Ignore objects whose bounding box takes up less than
          this fraction of the image.
        - min_objects_per_image: Ignore images which have fewer than this many
          object annotations.
        - max_objects_per_image: Ignore images which have more than this many
          object annotations.
        - include_other: If True, include COCO-Stuff annotations which have category
          "other". Default is False, because I found that these were really noisy
          and pretty much impossible for the system to model.
        - instance_whitelist: None means use all instance categories. Otherwise a
          list giving a whitelist of instance category names to use.
        - stuff_whitelist: None means use all stuff categories. Otherwise a list
          giving a whitelist of stuff category names to use.
        """
        super(Dataset, self).__init__()
        class_names = classes if classes is not ['all'] else ['cow', 'person', 'sheep', 'bird', 'cat', 'dog', 'horse', 'aeroplane', 'motorbike', 'bicycle']
        data_dict =  preprocess_make_data(class_names, data_mode)
        self.data_dict = data_dict

        if stuff_only and stuff_json is None:
            print('WARNING: Got stuff_only=True but stuff_json=None.')
            print('Falling back to stuff_only=False.')

        self.image_dir = image_dir
        self.mask_size = mask_size
        self.max_samples = max_samples
        self.max_objects_per_image = max_objects_per_image
        self.normalize_images = normalize_images
        self.include_relationships = include_relationships
        self.left_right_flip = left_right_flip
        self.set_image_size(image_size)

#        with open(instances_json, 'r') as f:
#            instances_data = json.load(f)

        stuff_data = None
#        if stuff_json is not None and stuff_json != '':
#            with open(stuff_json, 'r') as f:
#                stuff_data = json.load(f)

#        self.image_ids = []
#        self.image_id_to_filename = {}
#        self.image_id_to_size = {}
#        for image_data in instances_data['images']:
#            image_id = image_data['id']
#            filename = image_data['file_name']
#            width = image_data['width']
#            height = image_data['height']
#            self.image_ids.append(image_id)
#            self.image_id_to_filename[image_id] = filename
#            self.image_id_to_size[image_id] = (width, height)

        self.vocab = {
            'object_name_to_idx': {},
            'pred_name_to_idx': {},
        }


        # COCO category labels start at 1, so use 0 for __image__
        self.vocab['object_name_to_idx']['__image__'] = 0

        # Build object_idx_to_name
#        name_to_idx = self.vocab['object_name_to_idx']
#        assert len(name_to_idx) == len(set(name_to_idx.values()))
#        max_object_idx = max(name_to_idx.values())
#        idx_to_name = ['NONE'] * (1 + max_object_idx)
#        for name, idx in self.vocab['object_name_to_idx'].items():
#            idx_to_name[idx] = name
#        self.vocab['object_idx_to_name'] = idx_to_name

        # Prune images that have too few or too many objects
#        new_image_ids = []
#        total_objs = 0
#        for image_id in self.image_ids:
#            num_objs = len(self.image_id_to_objects[image_id])
#            total_objs += num_objs
#            if min_objects_per_image <= num_objs <= max_objects_per_image:
#                new_image_ids.append(image_id)
#        self.image_ids = new_image_ids

#        self.vocab['pred_idx_to_name'] = [
#            '__in_image__',
#            'left of',
#            'right of',
#            'above',
#            'below',
#            'inside',
#            'surrounding',
#        ]
#        self.vocab['pred_name_to_idx'] = {}
#        for idx, name in enumerate(self.vocab['pred_idx_to_name']):
#            self.vocab['pred_name_to_idx'][name] = idx
        #self.max_objects_per_image = 25

    def set_image_size(self, image_size):
        print('called set_image_size', image_size)
        transform = [Resize(image_size), T.ToTensor()]
        if self.normalize_images:
            transform.append(imagenet_preprocess())
        self.transform = T.Compose(transform)
        self.image_size = image_size

    def total_objects(self):
        total_objs = 0
        for id in self.data_dict.keys():
            total_objs = total_objs + self.data_dict[id]['boxes'].shape[0]
        return total_objs


    def __len__(self):
        return len(self.data_dict.keys())

    def __getitem__(self, index):
        """
        Get the pixels of an image, and a random synthetic scene graph for that
        image constructed on-the-fly from its COCO object annotations. We assume
        that the image will have height H, width W, C channels; there will be O
        object annotations, each of which will have both a bounding box and a
        segmentation mask of shape (M, M). There will be T triples in the scene
        graph.
    
        Returns a tuple of:
        - image: FloatTensor of shape (C, H, W)
        - objs: LongTensor of shape (O,)
        - boxes: FloatTensor of shape (O, 4) giving boxes for objects in
          (x0, y0, x1, y1) format, in a [0, 1] coordinate system
        - masks: LongTensor of shape (O, M, M) giving segmentation masks for
          objects, where 0 is background and 1 is object.
        - triples: LongTensor of shape (T, 3) where triples[t] = [i, p, j]
          means that (objs[i], p, objs[j]) is a triple.
        """
        image_id = str(index)
        img = self.data_dict[image_id]['image']

        #flip = False
        #if index >= len(self.image_ids):
        #    index = index - len(self.image_ids)
        #    flip = True
        #image_id = self.image_ids[index]

        #filename = self.image_id_to_filename[image_id]
        #image_path = os.path.join(self.image_dir, filename)
        #with open(image_path, 'rb') as f:
        #    with PIL.Image.open(f) as image:
        #        if flip:
        #            image = PIL.ImageOps.mirror(image)
        #        WW, HH = image.size
        #        image = self.transform(image.convert('RGB'))
        img = Image.fromarray(np.uint8(img))
        WW, HH = img.size
        image = self.transform(img) 

        bbxs = self.data_dict[image_id]['boxes'].tolist()
        
        class_name = self.data_dict[image_id]['class']
        H, W = self.image_size


        objs, boxes, masks = [], [], []
        for bbx in bbxs:
            x0, y0, x1, y1 = bbx
            #x0 = x0 * 550.0/WW 
            #y0 = y0 * 550.0/HH
            #x1 = x1 * 550.0/WW
            #y1 = y1 * 550.0/HH
            x1 = (x1*550.0 - x0*550.0 + 1.0)/550.0
            y1 = (y1*550.0 - y0*550.0 + 1.0)/550.0
            boxes.append(torch.FloatTensor([x0, y0, x1, y1]))
            #boxes.append(torch.FloatTensor([0.0, 0.0, 1.0, 1.0]))

            # This will give a numpy array of shape (HH, WW)
            #mask = torch.zeros(1, H, W)
            # mask = seg_to_mask(object_data['segmentation'], WW, HH)
            #mask[:, round(y0 * H):max(round(y0 * H)+1, round(y1 * H)), round(x0 * W):max(round(x0 * W)+1, round(x1 * W))] = 1

            # Crop the mask according to the bounding box, being careful to
            # ensure that we don't crop a zero-area region
            # mx0, mx1 = int(round(x)), int(round(x + w))
            # my0, my1 = int(round(y)), int(round(y + h))
            # mx1 = max(mx0 + 1, mx1)
            # my1 = max(my0 + 1, my1)
            # mask = mask[my0:my1, mx0:mx1]
            # mask = imresize(255 * mask, (self.mask_size, self.mask_size),
            #                 mode='constant')
            #mask = torch.from_numpy((mask > 128).astype(np.int64))
            #masks.append(mask)
            #obj_to_cls.append(torch.from_numpy(class_name))
        for obj in self.data_dict[image_id]['labels']:
            objs.append(obj)

        # obj_masks = []
        #for object_data in self.image_id_to_objects[image_id]:
        #    objs.append(object_data['category_id'])
        #    x, y, w, h = object_data['bbox']
        #    x0 = x / WW
        #    y0 = y / HH
        #    x1 = (w) / WW
        #    y1 = (h) / HH
        #    if flip:
        #        x0 = 1 - (x0 + x1)
        #    boxes.append(np.array([x0, y0, x1, y1]))

            # # This will give a numpy array of shape (HH, WW)
            # mask = seg_to_mask(object_data['segmentation'], WW, HH)
            # if flip:
            #     mask = mask[:, ::-1]
            #     x = WW - x - w
            # # Crop the mask according to the bounding box, being careful to
            # # ensure that we don't crop a zero-area region
            # # mx0, mx1 = int(round(x)), int(round(x + w))
            # # my0, my1 = int(round(y)), int(round(y + h))
            # # mx1 = max(mx0 + 1, mx1)
            # # my1 = max(my0 + 1, my1)
            # # obj_mask = mask[my0:my1, mx0:mx1]
            # # obj_mask = imresize(255.0 * obj_mask, (self.mask_size, self.mask_size),
            # #                  mode='constant')
            # # obj_mask = torch.from_numpy((obj_mask > 128).astype(np.int64))
            # # obj_masks.append(obj_mask)
            #
            # mask = imresize(255.0 * mask, (self.image_size[0], self.image_size[1]),
            #                 mode='constant')
            # mask = torch.from_numpy((mask > 128).astype(np.int64))
            # masks.append(mask)

        # shuffle(objs)
        # shuffle(boxes)
        # shuffle(masks)
        # Add dummy __image__ object
        # objs.append(184)
        # boxes.append(torch.FloatTensor([0, 0, 1, 1]))
        # masks.append(torch.ones(self.mask_size, self.mask_size).long())

        # add 0 for number of objects
        for _ in range(len(objs), self.max_objects_per_image):
            objs.append(self.vocab['object_name_to_idx']['__image__'])
            boxes.append(torch.FloatTensor([0.0, 0.0, 1.0, 1.0]))
            # masks.append(torch.zeros((self.image_size[0], self.image_size[1])).long())
            # obj_masks.append(torch.zeros((self.mask_size, self.mask_size)).long())

        #objs = torch.LongTensor(objs)
        #boxes = np.vstack(boxes)
        #O = len(objs)
        #rand_idx = list(range(O))
        #random.shuffle(rand_idx)

        #objs = [objs[i] for i in rand_idx]
        #boxes = [boxes[i] for i in rand_idx]

        objs = torch.LongTensor(objs)
        boxes = torch.stack(boxes, dim=0)
        

        # masks = torch.stack(masks, dim=0)
        # obj_masks = torch.stack(obj_masks, dim=0)
        # b_map = self.get_bbox_map_p(boxes)

        # box_areas = (boxes[:, 2] - boxes[:, 0]) * (boxes[:, 3] - boxes[:, 1])

        # Compute centers of all objects
        """
        obj_centers = []
        _, MH, MW = masks.size()
        for i, obj_idx in enumerate(objs):
            x0, y0, x1, y1 = boxes[i]
            mask = (masks[i] == 1)
            xs = torch.linspace(x0, x1, MW).view(1, MW).expand(MH, MW)
            ys = torch.linspace(y0, y1, MH).view(MH, 1).expand(MH, MW)
            if mask.sum() == 0:
                mean_x = 0.5 * (x0 + x1)
                mean_y = 0.5 * (y0 + y1)
            else:
                mean_x = xs[mask].mean()
                mean_y = ys[mask].mean()
            obj_centers.append([mean_x, mean_y])
        obj_centers = torch.FloatTensor(obj_centers)

        # Add triples
        
        triples = []
        num_objs = objs.size(0)
        __image__ = self.vocab['object_name_to_idx']['__image__']
        real_objs = []
        if num_objs > 1:
            real_objs = (objs != __image__).nonzero().squeeze(1)
        for cur in real_objs:
            choices = [obj for obj in real_objs if obj != cur]
            if len(choices) == 0 or not self.include_relationships:
                break
            other = random.choice(choices)
            if random.random() > 0.5:
                s, o = cur, other
            else:
                s, o = other, cur

            # Check for inside / surrounding
            sx0, sy0, sx1, sy1 = boxes[s]
            ox0, oy0, ox1, oy1 = boxes[o]
            d = obj_centers[s] - obj_centers[o]
            theta = math.atan2(d[1], d[0])

            if sx0 < ox0 and sx1 > ox1 and sy0 < oy0 and sy1 > oy1:
                p = 'surrounding'
            elif sx0 > ox0 and sx1 < ox1 and sy0 > oy0 and sy1 < oy1:
                p = 'inside'
            elif theta >= 3 * math.pi / 4 or theta <= -3 * math.pi / 4:
                p = 'left of'
            elif -3 * math.pi / 4 <= theta < -math.pi / 4:
                p = 'above'
            elif -math.pi / 4 <= theta < math.pi / 4:
                p = 'right of'
            elif math.pi / 4 <= theta < 3 * math.pi / 4:
                p = 'below'
            p = self.vocab['pred_name_to_idx'][p]
            triples.append([s, p, o])

        # Add __in_image__ triples
        O = objs.size(0)
        in_image = self.vocab['pred_name_to_idx']['__in_image__']
        for i in range(O - 1):
            triples.append([i, in_image, O - 1])
        """
        cls = torch.from_numpy(class_name)
        print("Image {}".format(image.shape))
        print("Labels {}".format(objs.shape))
        print("Bbox {}".format(boxes.shape))
        print("Class {}".format(cls.shape))
        # triples = torch.LongTensor(triples)
        return image, objs, boxes, cls## , b_map #, None # obj_masks #, obj_masks # , b_map # masks # , triples

    def get_bbox_map_p(self, bbox):
        mapping = np.zeros((len(bbox), self.image_size[0], self.image_size[0]))
        for idx in range(self.max_objects_per_image):
            if min(bbox[idx]) < 0:
                continue
            line_space = np.linspace(0, self.image_size[0]-1, num=self.image_size[0])
            xv, yv = np.meshgrid(line_space, line_space)
            mapping[idx][(xv < int((bbox[idx][0] + bbox[idx][2]) * self.image_size[0])) * (xv > int(bbox[idx][0] * self.image_size[0])) *
                    (yv < int((bbox[idx][1] + bbox[idx][3]) * self.image_size[0])) * (yv > int(bbox[idx][1] * self.image_size[0]))] = 1
        return mapping


def seg_to_mask(seg, width=1.0, height=1.0):
    """
    Tiny utility for decoding segmentation masks using the pycocotools API.
    """
    if type(seg) == list:
        rles = mask_utils.frPyObjects(seg, height, width)
        rle = mask_utils.merge(rles)
    elif type(seg['counts']) == list:
        rle = mask_utils.frPyObjects(seg, height, width)
    else:
        rle = seg
    return mask_utils.decode(rle)


def coco_collate_fn(batch):
    """
    Collate function to be used when wrapping CocoSceneGraphDataset in a
    DataLoader. Returns a tuple of the following:
  
    - imgs: FloatTensor of shape (N, C, H, W)
    - objs: LongTensor of shape (O,) giving object categories
    - boxes: FloatTensor of shape (O, 4)
    - masks: FloatTensor of shape (O, M, M)
    - triples: LongTensor of shape (T, 3) giving triples
    - obj_to_img: LongTensor of shape (O,) mapping objects to images
    - triple_to_img: LongTensor of shape (T,) mapping triples to images
    """
    all_imgs, all_objs, all_boxes, all_masks, all_triples = [], [], [], [], []
    all_obj_to_img, all_triple_to_img = [], []
    obj_offset = 0
    for i, (img, objs, boxes, masks, triples) in enumerate(batch):
        all_imgs.append(img[None])
        if objs.dim() == 0 or triples.dim() == 0:
            continue
        O, T = objs.size(0), triples.size(0)
        all_objs.append(objs)
        all_boxes.append(boxes)
        all_masks.append(masks)
        triples = triples.clone()
        triples[:, 0] += obj_offset
        triples[:, 2] += obj_offset
        all_triples.append(triples)

        all_obj_to_img.append(torch.LongTensor(O).fill_(i))
        all_triple_to_img.append(torch.LongTensor(T).fill_(i))
        obj_offset += O

    all_imgs = torch.cat(all_imgs)
    all_objs = torch.cat(all_objs)
    all_boxes = torch.cat(all_boxes)
    all_masks = torch.cat(all_masks)
    all_triples = torch.cat(all_triples)
    all_obj_to_img = torch.cat(all_obj_to_img)
    all_triple_to_img = torch.cat(all_triple_to_img)

    out = (all_imgs, all_objs, all_boxes, all_masks, all_triples,
           all_obj_to_img, all_triple_to_img)
    return out


# IMAGENET_MEAN = [0.485, 0.456, 0.406]
# IMAGENET_STD = [0.229, 0.224, 0.225]
IMAGENET_MEAN = [0.5, 0.5, 0.5]
IMAGENET_STD = [0.5, 0.5, 0.5]

INV_IMAGENET_MEAN = [-m for m in IMAGENET_MEAN]
INV_IMAGENET_STD = [1.0 / s for s in IMAGENET_STD]


def imagenet_preprocess():
    return T.Normalize(mean=IMAGENET_MEAN, std=IMAGENET_STD)


def rescale(x):
    lo, hi = x.min(), x.max()
    return x.sub(lo).div(hi - lo)


def imagenet_deprocess(rescale_image=True):
    transforms = [
        T.Normalize(mean=[0, 0, 0], std=INV_IMAGENET_STD),
        T.Normalize(mean=INV_IMAGENET_MEAN, std=[1.0, 1.0, 1.0]),
    ]
    if rescale_image:
        transforms.append(rescale)
    return T.Compose(transforms)


def imagenet_deprocess_batch(imgs, rescale=True):
    """
    Input:
    - imgs: FloatTensor of shape (N, C, H, W) giving preprocessed images
    
    Output:
    - imgs_de: ByteTensor of shape (N, C, H, W) giving deprocessed images
    in the range [0, 255]
    """
    if isinstance(imgs, torch.autograd.Variable):
        imgs = imgs.data
    imgs = imgs.cpu().clone()
    deprocess_fn = imagenet_deprocess(rescale_image=rescale)
    imgs_de = []
    for i in range(imgs.size(0)):
        img_de = deprocess_fn(imgs[i])[None]
        img_de = img_de.mul(255).clamp(0, 255).byte()
        imgs_de.append(img_de)
    imgs_de = torch.cat(imgs_de, dim=0)
    return imgs_de


class Resize(object):
    def __init__(self, size, interp=PIL.Image.BILINEAR):
        if isinstance(size, tuple):
              H, W = size
              self.size = (W, H)
        else:
            self.size = (size, size)
        self.interp = interp

    def __call__(self, img):
        return img.resize(self.size, self.interp)


def unpack_var(v):
    if isinstance(v, torch.autograd.Variable):
        return v.data
    return v


def split_graph_batch(triples, obj_data, obj_to_img, triple_to_img):
    triples = unpack_var(triples)
    obj_data = [unpack_var(o) for o in obj_data]
    obj_to_img = unpack_var(obj_to_img)
    triple_to_img = unpack_var(triple_to_img)

    triples_out = []
    obj_data_out = [[] for _ in obj_data]
    obj_offset = 0
    N = obj_to_img.max() + 1
    for i in range(N):
        o_idxs = (obj_to_img == i).nonzero().view(-1)
        t_idxs = (triple_to_img == i).nonzero().view(-1)

        cur_triples = triples[t_idxs].clone()
        cur_triples[:, 0] -= obj_offset
        cur_triples[:, 2] -= obj_offset
        triples_out.append(cur_triples)

        for j, o_data in enumerate(obj_data):
            cur_o_data = None
            if o_data is not None:
                cur_o_data = o_data[o_idxs]
            obj_data_out[j].append(cur_o_data)

        obj_offset += o_idxs.size(0)

        return triples_out, obj_data_out



torch.cat torch.stack torch.FloatTensor torch.from_numpy torch.LongTensor