from transformers import PreTrainedModel
import torch
import os

url_map = {
	"inception_v3": "https://download.pytorch.org/models/inception_v3_google-1a9a5a14.pth"
}

class InceptionV3ModelForImageClassification(PreTrainedModel):
	def __init__(self, config):
		super().__init__(config)

		model_path = f"{self.config.model_name}.bin".replace("/","_")

		if self.config.model_name == "google-safesearch-mini":
			self.model = torch.jit.load(model_path)
		elif self.config.model_name == "inception_v3":
			self.model = torch.hub.load('pytorch/vision:v0.6.0', 'inception_v3', pretrained=True)
		else:
			if not os.path.exists(model_path):
				from urllib.request import urlretrieve
				urlretrieve(f"https://huggingface.co/{self.config.model_name}/resolve/main/pytorch_model.bin", model_path)
			self.model = torch.jit.load(model_path) if self.config.use_jit else torch.load(model_path)

	def forward(self, input_ids):
		out, aux = self.model(input_ids)
		return out, aux

	def freeze(self):
		for param in self.model.parameters():
			param.requires_grad = False

	def unfreeze(self):
		for param in self.model.parameters():
			param.requires_grad = True

	def train(self, mode=True):
		super().train(mode)
		self.model.train(mode)

	def eval(self):
		return self.train(False)

	def to(self, device):
		self.model.to(device)
		return self

	def cuda(self, device=None):
		return self.to("cuda")

	def cpu(self):
		return self.to("cpu")

	def state_dict(self, destination=None, prefix='', keep_vars=False):
		return self.model.state_dict(destination, prefix, keep_vars)

	def load_state_dict(self, state_dict, strict=True):
		return self.model.load_state_dict(state_dict, strict)

	def parameters(self, recurse=True):
		return self.model.parameters(recurse)

	def named_parameters(self, prefix='', recurse=True):
		return self.model.named_parameters(prefix, recurse)

	def children(self):
		return self.model.children()

	def named_children(self):
		return self.model.named_children()

	def modules(self):
		return self.model.modules()

	def named_modules(self, memo=None, prefix=''):
		return self.model.named_modules(memo, prefix)

	def zero_grad(self, set_to_none=False):
		return self.model.zero_grad(set_to_none)

	def share_memory(self):
		return self.model.share_memory()

	def transform(self, image):
		from torchvision import transforms
		transform = transforms.Compose([
			transforms.Resize(299),
			transforms.ToTensor(),
			transforms.Normalize(mean=self.config.mean, std=self.config.std)
		])
		image = transform(image)
		return image

	def open_image(self, path):
		from PIL import Image
		path = 'https://images.unsplash.com/photo-1594568284297-7c64464062b1'
		if path.startswith('http://') or path.startswith('https://'):
			import requests
			from io import BytesIO
			response = requests.get(path)
			image = Image.open(BytesIO(response.content)).convert('RGB')
		else:
			image = Image.open(path).convert('RGB')
		return image

	def predict(self, path, device="cuda", print_tensor=True):
		image = self.open_image(path)
		image = self.transform(image)
		image = image.unsqueeze(0)
		self.eval()
		if device == "cuda":
			image = image.cuda()
			self.cuda()
		else:
			image = image.cpu()
			self.cpu()
		with torch.no_grad():
			out, aux = self(image)
			if print_tensor:
				print(out)
			_, predicted = torch.max(out.data, 1)
		return self.config.classes[str(predicted.item())]