
import sys, os
from pathlib import Path

import random

from tqdm import tqdm
try:
	import wget
except ImportError:
	print('WARNING: unable to import wget')
import shutil
import h5py as hf

import cv2
import numpy as np
import torch

from omnibelt import create_dir
import omnifig as fig

from ... import util

from ...data import Dataset, Deviced, Batchable, Image_Dataset

DATASET_URL = 'ftp://guest:GU.205dldo@ftp.softronics.ch/mvtec_anomaly_detection/mvtec_anomaly_detection.tar.xz'

def get_images(root, size=None, is_mask=False):
	raw = [util.read_raw_bytes(name, root=root) for name in os.listdir(root)]
	if size is None:
		return raw
	
	images = []
	for data in raw:
		image = util.str_to_byte_img(data) if is_mask else util.str_to_rgb(data)[..., ::-1]
		image = cv2.resize(image, (size,size))
		images.append(image)
	return images
	
@fig.Script('download-mvtec', description='Download and Format MVTec Anomaly Detection dataset')
def download_mvtec(A):
	raise NotImplementedError
	root = A.pull('root', os.environ.get('OMNILEARN_DATA_DIR', None))
	
	if root is None:
		raise Exception('Must specify a OMNILEARN_DATA_DIR as env variable (or pass in a `root`)')
	
	root = os.path.join(root, 'mvtec')
	
	if not os.path.isdir(root):
		create_dir(root)

	raw_name = A.pull('raw-name', 'mvtec_anomaly_detection.tar.xz', silent=True)
	
	size = A.pull('size', None)
	name = 'full' if size is None else f'size{size}'
	
	path = os.path.join(root, name)
	
	cleanup = A.pull('cleanup', False)
	cleanup_compressed = A.pull('cleanup-compressed', False)
	
	base = os.path.join(root, 'raw')
	if not os.path.isdir(base):
		rawpath = os.path.join(root, raw_name)
		if not os.path.isfile(rawpath):
			url = A.pull('url', DATASET_URL)
			print(f'Downloading MVTec AD dataset from {url}')
			
			rawpath = wget.download(url, out=root)
		
		create_dir(base)
		
		print(f'Extracting MVTec AD dataset from {rawpath} file')
		
		shutil.unpack_archive(rawpath, extract_dir=base)
		
		if cleanup_compressed:
			os.remove(rawpath)
	
	create_dir(path)
	
	contents = os.listdir(path)
	
	todo = [name for name in os.listdir(base)
	        if os.path.isdir(os.path.join(base, name))
	        and f'{name}.h5' not in contents]
	
	print(f'Formatting {len(todo)} object directories')
	
	pbar = A.pull('pbar', True, silent=True)
	if pbar:
		todo = tqdm(todo)
	
	for cat in todo:
		
		croot = os.path.join(base, cat)
		
		if pbar:
			todo.set_description(cat)
		
		with hf.File(os.path.join(path, f'{cat}.h5'), 'w') as f:
		
			mroot = os.path.join(croot, 'ground_truth')
		
			defects = os.listdir(mroot)
			f.attrs['defects'] = defects
		
			for defect in defects:
				f.create_dataset(f'mask_{defect}', data=get_images(os.path.join(mroot, defect),
				                                                   size=size, is_mask=True))
		
			f.create_dataset('train_good', data=get_images(os.path.join(croot, 'train', 'good'), size=size))
			
			mroot = os.path.join(croot, 'test')
			for test in os.listdir(mroot):
				f.create_dataset(f'test_{test}', data=get_images(os.path.join(mroot, test), size=size))
		
		if cleanup:
			print(f'removing: {cat}')
			# shutil.rmtree(croot)
	
	print('Formatted all categories')
	
	return path

@Dataset('mvtec')
class MVTec_Anomaly_Detection(Batchable):

	CATEGORIES = {'bottle', 'carpet', 'leather', 'screw', 'transistor', 'cable', 'grid',
        'metal_nut', 'tile', 'wood', 'capsule', 'hazelnut', 'pill','toothbrush', 'zipper'}
	GREYSCALES = {'screw', 'grid', 'zipper'}

	def __init__(self, A, **kwargs):
	
		dataroot = util.get_data_dir(A) / 'mvtec'

		mode = A.pull('mode', 'train')

		size = A.pull('size', None)
		
		dirname = 'full' if size is None else f'size{size}'
		droot = dataroot / dirname
		
		if not dataroot.is_dir() or not droot.is_dir():
			download = A.pull('download', False)
			if not download:
				raise Exception(f'This dataset ({dirname}) hasnt been downloaded and setup yet '
				                '(set the "download" argument to do so automatically)')
		
			print(f'Downloading/Formatting MVTec dataset: {dirname}')
			droot = fig.quick_run('download-mvtec', root=str(dataroot), size=size)
			droot = Path(droot)
		
		cat = A.pull('cat', 'random')
		if cat == 'random':
			cat = random.choice(list(self.CATEGORIES))

		ratio = A.pull('ratio', None)
		if ratio is not None:
			assert 0 < ratio <= 1, f'bad ratio: {ratio}'

		path = [c for c in droot.glob('*.h5') if cat == c.stem][0]
		
		include_class = A.pull('include-class', '<>include_class', True)
		include_mask = A.pull('include-masks', '<>include_masks', False)
		
		cut = A.pull('cut', 'train')
		assert cut in {None, 'normal', 'anomalies', 'all', 'train', 'test'}, f'unknown: {cut}'

		tfms = A.pull('transforms', None, silent=True)
		if tfms is not None:
			tfms = tfms.get(cat, None)

		augmenter = None
		augment_factor = 1
		if tfms is not None:
			augment_here = A.pull('augment_factor', None)

			if augment_here is not None:
				A.push('augmenter', tfms)
				A.push('augmenter._type', 'image-transform', overwrite=False)
				augmenter = A.pull('augmenter', None)
				augment_factor = augment_here
			else:
				A.push('augmentations', tfms, force_root=True)



		C = 1 if cat in self.GREYSCALES else 3
		din = (C, 1024, 1024) if size is None else (C, size, size)
		dout = din
		super().__init__(A, din=din, dout=dout, **kwargs)
		
		raw = hf.File(path, 'r')
		
		# print(raw.keys())

		uses = [key for key in raw.keys() if 'mask_' not in key]
		if cut == 'train':
			uses = [key for key in uses if 'train_' in key]
		elif cut == 'test':
			uses = [key for key in uses if 'test_' in key]
		elif cut == 'normal':
			uses = [key for key in uses if 'train_' in key or key == 'test_good']
		elif cut == 'anomalies':
			uses = [key for key in uses if 'test_' in key and key != 'test_good']

		C, H, W = din
		
		images = []
		masks = [] if include_mask else None
		labels = [] if include_class else None
		idents = []
		for key in uses:
			N = len(raw[key])
			if ratio is None:
				sel = ()
			elif mode == 'test':
				sel = slice(int(-N*(1-ratio)),None)
			else:
				sel = slice(int(N*ratio))

			imgs = raw[key][sel]
			# print(key, imgs.shape)
			images.extend(imgs)
			if labels is not None:
				labels.extend([1 if 'test_' in key else 0]*len(imgs))
			ident = '_'.join(key.split('_')[1:])
			idents.append(ident)
			if masks is not None:
				if 'test_' in key and ident != 'good':
					ident = f'mask_{ident}'
					masks.extend(raw[ident][sel])
				else:
					masks.extend([np.zeros((H,W))]*len(imgs))
		
		if masks is not None and not len(masks):
			masks = None
		if labels is not None:
			labels = torch.tensor(labels).long()
		
		if size is not None:
			# print([m.shape for m in images])
			images = torch.from_numpy(np.stack(images))
			images = images.unsqueeze(1) if C == 1 else images.permute(0,3,1,2)
			images = images.float().div(255)
			if masks is not None and len(masks):
				# print([m.shape for m in masks])
				masks = torch.from_numpy(np.stack(masks)).unsqueeze(1).bool()

		raw.close()
		# self.raw = raw

		self.transforms = tfms
		
		self.size = size
		
		self.images = images
		self.masks = masks
		self.labels = labels
		self.uses = uses
		self.idents = idents

		self.augmenter = augmenter
		self.augment_factor = augment_factor

	def get_transforms(self):
		return self.transforms
	
	def __len__(self):
		return len(self.images)#*self.augment_factor
	
	def __getitem__(self, item):

		# item %= len(self.images)
		
		if self.size is not None:
			out = [self.images[item]]

			if self.masks is not None:
				out.append(self.masks[item])

				if self.augmenter is not None:
					full = self.augmenter(torch.cat(out, 1))
					out = [full[:, :-1].contiguous(), full[:, -1:].contiguous()]

				if self.labels is not None:
					out.append(self.labels[item])
			elif self.labels is not None:
				out.append(self.labels[item])
			if self.augmenter is not None:
				out[0] = self.augmenter(out[0])

			return out
		
		raise NotImplementedError




