import numpy as np
import chainer.links as L
import chainer.functions as F

from chainer.serializers import npz
from chainer.links.model.vision.vgg import prepare

from chainer_addons.models import PretrainedModelMixin

class VGG19Layers(PretrainedModelMixin, L.VGG16Layers):
	"""
		adds some additional layers to the original VGG16
		model definition. the last linear layer is also replaced
		if number of classes is given. this is useful in case of
		fine-tuning
	"""

	class meta:
		classifier_layers = ["fc6", "fc7", "fc8"]
		conv_map_layer = "conv5_4"
		n_conv_maps = 512
		mean = np.array([103.939, 116.779, 123.68], dtype=np.float32).reshape(3,1,1)

	@staticmethod
	def prepare(*args, **kw):
		return prepare(*args, **kw)

	@staticmethod
	def pooling(x):
		return F.max_pooling_2d(x, ksize=2, stride=2)

	def __init__(self, pretrained_model, n_classes=None):
		"""
			First we load vanilla VGG16 model, then additional layers are
			added and finally the model weights are loaded
		"""
		super(VGG19Layers, self).__init__(pretrained_model=None)

		self.add_link('conv3_4', L.Convolution2D(256, 256, 3, 1, 1))
		self.add_link('conv4_4', L.Convolution2D(512, 512, 3, 1, 1))
		self.add_link('conv5_4', L.Convolution2D(512, 512, 3, 1, 1))

		if n_classes is not None:
			self.fc8 = L.Linear(4096, n_classes, **init_args)

		self.functions = OrderedDict([
			('conv1_1',  [self.conv1_1, F.relu]),
			('conv1_2',  [self.conv1_2, F.relu]),
			('pool1',    [VGG19Layers.pooling]),
			('conv2_1',  [self.conv2_1, F.relu]),
			('conv2_2',  [self.conv2_2, F.relu]),
			('pool2',    [VGG19Layers.pooling]),
			('conv3_1',  [self.conv3_1, F.relu]),
			('conv3_2',  [self.conv3_2, F.relu]),
			('conv3_3',  [self.conv3_3, F.relu]),
			('conv3_4',  [self.conv3_4, F.relu]),
			('pool3',    [VGG19Layers.pooling]),
			('conv4_1',  [self.conv4_1, F.relu]),
			('conv4_2',  [self.conv4_2, F.relu]),
			('conv4_3',  [self.conv4_3, F.relu]),
			('conv4_4',  [self.conv4_4, F.relu]),
			('pool4',    [VGG19Layers.pooling]),
			('conv5_1',  [self.conv5_1, F.relu]),
			('conv5_2',  [self.conv5_2, F.relu]),
			('conv5_3',  [self.conv5_3, F.relu]),
			('conv5_4',  [self.conv5_4, F.relu]),
			('pool5',    [VGG19Layers.pooling]),
			('fc6',      [self.fc6, F.relu, F.dropout]),
			('fc7',      [self.fc7, F.relu, F.dropout]),
			('fc8',      [self.fc8]),
			('prob',     [F.softmax]),
		])

		npz.load_npz(pretrained_model, self)

