# flake8: noqa

from PIL import Image
import torchvision.transforms.functional as FT
import torchvision
import torch
import numpy as np
from scipy.spatial import distance
import random
# import torchvision.transforms.functional.pad
# torch.backends.cudnn.enabled = True
# torch.backends.cudnn.deterministic = True
import itertools
from PIL import Image
from watch.tasks.rutgers_material_seg.utils import utils
from tifffile import tifffile
# if 1:
#     torch.set_printoptions(precision=6, sci_mode=False)

IMG_EXTENSIONS = ['*.png', '*.jpeg', '*.jpg', '*.npy']


# FIXME: Hard coding mean/std is not a good idea
mean_std = ([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])

channels_dict =  {'B01':0, #(Coastal)
            'B02':1, #(Blue)
            'B03':2, #(Green)
            'B04':3, #(Red)
            'B05':4, #(RedEdge-1)
            'B06':5, #(RedEdge-2)
            'B07':6, #(RedEdge-3)
            'B08':7, #(NIR)
            'B8A':8, #(Narrow NIR)
            'B09':9, #(Water Vapor)
            'B10':10, #(Cirrus)
            'B11':11, #(SWIR-1)
            'B12':12 #(SWIR-2)
            }

class S2SelfCollectDataset(object):
    def __init__(self, root, transforms, split=False, crop_size=300, channels='B02|B03|B04'):
        self.root = root
        self.transforms = transforms
        self.split = split
        self.crop_size = crop_size
        self.randomcrop_transform = torchvision.transforms.RandomCrop(size=(crop_size, crop_size))

        self.images_root = f"{self.root}regions"
        self.images1_paths = utils.dictionary_contents(path=self.images_root, types=['*.tif'], recursive=True)
        self.channels_idx = [channels_dict[x] for x in channels.split('|')]
        self.channels_delete = [e for e in channels_dict.values() if e not in self.channels_idx]
        # print(self.channels_delete)
    
    def __getitem__(self, idx):

        img1_path = self.images1_paths[idx]
        base_path = f"{'/'.join(img1_path.split('/')[:-1])}/"
        base_region_path = f"{'/'.join(img1_path.split('/')[:-3])}/"
        
        unusable_negatives_path = utils.dictionary_contents(path=base_region_path, types=['*.tif'])
        image_name = img1_path.split('/')[-1].split('.')[0].split('_')[0]
        same_region_crops = utils.dictionary_contents(path=base_path, types=['*.tif'])
        negative_region_crops = [e for e in self.images1_paths if e not in unusable_negatives_path]
        same_region_crops.remove(img1_path)

        img2_path = random.sample(same_region_crops, 1)[0]
        negative_image_path = random.sample(negative_region_crops, 1)[0]
        
        # img1 = Image.open(img1_path)
        img1 = np.delete(tifffile.imread(img1_path),self.channels_delete, axis=-1)/2
        img2 = np.delete(tifffile.imread(img2_path),self.channels_delete, axis=-1)/2
        negative_img = np.delete(tifffile.imread(negative_image_path),self.channels_delete, axis=-1)/2

        new_image1 = self.transforms(img1)
        new_image2 = self.transforms(img2)
        new_negative_image = self.transforms(negative_img)

        crop_params = self.randomcrop_transform.get_params(new_image1, output_size=(self.crop_size, self.crop_size))
        new_image1 = FT.crop(new_image1, *crop_params)
        new_image2 = FT.crop(new_image2, *crop_params)
        new_negative_image = FT.crop(new_negative_image, *crop_params)

        # print(new_image1.dtype)
        # new_image1 = new_image1.double()
        # print(new_image1.dtype)
        # import matplotlib.pyplot as plt
        # plt.imshow(mask)
        # plt.show()

        outputs = {}
        outputs['visuals'] = {'image1': new_image1, 'image2':new_image2, 'image_name': image_name}
        outputs['inputs'] = {'image1': new_image1, 'image2':new_image2, 'negative_image': new_negative_image}

        return outputs

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