import os.path
import random
import torchvision.transforms as transforms
import torch
from data.base_dataset import BaseDataset
from data.image_folder import make_dataset
from PIL import Image
import json


class AlignedDataset(BaseDataset):
    def initialize(self, opt):
        self.opt = opt
        self.root = opt.dataroot
        self.dir_AB = os.path.join(opt.dataroot, 'images', opt.phase)
        self.dir_bbox = os.path.join(opt.dataroot, 'bbox', opt.phase)

        #self.AB_paths, self.bbox_paths = sorted(make_dataset(self.dir_AB, self.dir_bbox))
        self.AB_paths, self.bbox_paths = make_dataset(self.dir_AB, self.dir_bbox)
        self.AB_paths = sorted(self.AB_paths)
        self.bbox_paths = sorted(self.bbox_paths)

        assert(opt.resize_or_crop == 'resize_and_crop')

        transform_list = [transforms.ToTensor(),
                          transforms.Normalize((0.5, 0.5, 0.5),
                                               (0.5, 0.5, 0.5))]

        self.transform = transforms.Compose(transform_list)

    def __getitem__(self, index):
        AB_path = self.AB_paths[index]
        bbox_path = self.bbox_paths[index]

        w = self.opt.loadSize
        h = self.opt.loadSize
        fine_size = self.opt.fineSize
        w_offset = random.randint(0, max(0, w - fine_size - 1))
        h_offset = random.randint(0, max(0, h - fine_size - 1))

        bbox = json.load(open(bbox_path))
        bbox_x = max(int((bbox['x'] / fine_size) * self.opt.loadSize), 0)
        bbox_y = max(int((bbox['y'] / fine_size) * self.opt.loadSize), 0)
        bbox_w = max(int((bbox['w'] / fine_size) * self.opt.loadSize), 0)
        bbox_h = max(int((bbox['h'] / fine_size) * self.opt.loadSize), 0)

        if bbox_y <= h_offset or bbox_x <= w_offset:
            # shift transform not applicable
            AB = Image.open(AB_path).convert('RGB')
            AB = self.transform(AB)
            A = AB[:, :fine_size, :fine_size]
            B = AB[:, :fine_size, fine_size:2 * fine_size]
            bbox = [bbox['y'], bbox['x'], bbox['w'], bbox['h']]
        else:
            AB = Image.open(AB_path).convert('RGB')
            AB = AB.resize((self.opt.loadSize * 2, self.opt.loadSize), Image.BICUBIC)
            AB = self.transform(AB)
            A = AB[:, h_offset:h_offset + fine_size, w_offset:w_offset + fine_size]
            B = AB[:, h_offset:h_offset + fine_size, w + w_offset:w + w_offset + fine_size]
            bbox = [bbox_y - h_offset, bbox_x - w_offset, bbox_w - w_offset, bbox_h - h_offset]

        if (not self.opt.no_flip) and random.random() < 0.5:
            A = A.flip(2)
            B = B.flip(2)
            bbox = [bbox[0], fine_size - bbox[2], fine_size - bbox[1], bbox[3]]

        return {'A': A, 'B': B, 'bbox': bbox,
                'A_paths': AB_path, 'B_paths': AB_path}

    def __len__(self):
        return len(self.AB_paths)

    def name(self):
        return 'AlignedDataset'
