import os
import pydicom
import numpy as np
import json
import PIL
import skimage
import logging

import torchxrayvision as xrv
from PIL import Image, ImageFile
from torch.utils.data import Dataset
from torchvision.transforms import *
from .utils import load_file
from .papers.open_image import *

ImageFile.LOAD_TRUNCATED_IMAGES = True  # Are we sure ?
PIL.Image.MAX_IMAGE_PIXELS = None  # Are we sure ?
logging.getLogger('PIL').setLevel(logging.WARNING)


def read_images(root, image_path, split, file):
    lines = load_file(os.path.join(root, split + '.' + file))
    return [[os.path.join(image_path, image) for image in line.split(',')] for line in lines]


def get_transforms(split, resize, crop, custom_transform_train, custom_transform_val, ext):
    if ext in [".npy", ".npz"]:
        return lambda x: x

    assert not (custom_transform_train is not None and (ext in ['.xrv'])), \
        'Cant specify custom transform when using xrv'

    if ext == '.xrv':
        return transforms.Compose([xrv.datasets.XRayCenterCrop(),
                                   xrv.datasets.XRayResizer(crop),
                                   lambda x: torch.from_numpy(x)])

    if split == 'train':
        return transforms.Compose([
            transforms.Resize(resize),
            transforms.RandomCrop(crop),
            transforms.RandomHorizontalFlip(),
            transforms.ToTensor(),
            transforms.Normalize((0.485, 0.456, 0.406),
                                 (0.229, 0.224, 0.225))]) if custom_transform_train is None else eval(
            custom_transform_train)
    else:
        return transforms.Compose([
            transforms.Resize((crop, crop)),
            transforms.ToTensor(),
            transforms.Normalize((0.485, 0.456, 0.406),
                                 (0.229, 0.224, 0.225))]) if custom_transform_val is None else eval(
            custom_transform_val)


def open_image(image, ext):
    if ext == '.jpg':
        return Image.open(image).convert('RGB')

    if ext == '.xrv':
        img = skimage.io.imread(image)
        img = xrv.datasets.normalize(img, 255)
        if len(img.shape) > 2:
            img = img[:, :, 0]
        if len(img.shape) < 2:
            print("error, dimension lower than 2 for image")
        return img[None, :, :]

    if ext == '.png':
        return Image.open(image).convert('RGB')

    if ext == '.dcm':
        ds = pydicom.dcmread(image)
        img = ds.pixel_array
        return Image.fromarray(np.uint8(img)).convert('RGB')

    if ext in ['.npy', '.npz']:
        return torch.from_numpy(np.load(image))

    # Special cases
    if ext in PAPER_EXT.keys():
        return eval(PAPER_EXT[ext])(image)

    raise NotImplementedError("Image extension {} not implemented".format(ext))


def do_image(image, transform, ext):
    opened_images = [open_image(im, ext) for im in image]
    transformed_images = [transform(im) for im in opened_images]
    return transformed_images


class ImageDataset(Dataset):
    def __init__(self,
                 root=None,
                 file=None,
                 split=None,
                 image_path=None,
                 load_memory=False,
                 custom_transform_train=None,
                 custom_transform_val=None,
                 resize=256,
                 crop=224,
                 ext='.jpg',
                 multi_image=None,
                 **kwargs):

        assert split is not None, "Argument split cant be None"

        self.root = root
        self.file = file
        self.split = split
        self.image_path = image_path
        self.load_memory = load_memory
        self.resize = eval(str(resize))
        self.crop = eval(str(crop))
        self.ext = ext
        self.multi_image = multi_image or 0
        self.images = None

        if file is not None:
            self.images = read_images(root, image_path, split, file)

        self.transform = get_transforms(split,
                                        self.resize,
                                        self.crop,
                                        custom_transform_train,
                                        custom_transform_val,
                                        self.ext)

    def __len__(self):
        return len(self.images or [])

    def __getitem__(self, index):
        return {'image': do_image(self.images[index], self.transform, self.ext)}

    def inference(self, image):
        if not isinstance(image, list):
            image = [image]
        batch = [{'image': do_image(i, self.transform, self.ext)} for i in image]
        return self.get_collate_fn()(batch)

    def get_collate_fn(self):
        def collate_fn(batch):
            # Return one image
            if not self.multi_image or self.multi_image == 1:
                return {'images': torch.stack([s['image'][0] for s in batch])}

            # Return multiple image
            new_batch = []
            new_masks = []
            for sample in batch:
                sample_images = sample['image']
                # Remove image to get to self.multi_image
                if len(sample_images) > self.multi_image:
                    sample_images = sample_images[:self.multi_image]
                # Pad with zeros to get to self.multi_image
                if len(sample_images) < self.multi_image:
                    first_image = sample_images[0]
                    for _ in range(self.multi_image - len(sample_images)):
                        sample_images.append(first_image.new_zeros(first_image.size()))
                # Stack
                sample_images = torch.cat([s.unsqueeze(dim=0) for s in sample_images], dim=0)
                sample_mask = (sample_images.sum(dim=(1, 2, 3)) != 0)
                new_batch.append(sample_images)
                new_masks.append(sample_mask)

            collated = {'images': torch.stack(new_batch),
                        'images_mask': torch.stack(new_masks)}
            return collated

        return collate_fn

    def __repr__(self):
        transform = self.transform
        if hasattr(self.transform, "transforms"):
            transform = self.transform.transforms

        return "ImageDataset\n" + \
               json.dumps({
                   "image_path": self.image_path,
                   "root": self.root,
                   "file": self.file,
                   "transforms": transform,
                   "ext": self.ext,
               }, indent=4, sort_keys=False, default=str)
