# --------------------------------------------------------
# Tensorflow Faster R-CNN
# Licensed under The MIT License [see LICENSE for details]
# Written by Xinlei Chen
# --------------------------------------------------------
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

from nets.network import Network
from model.config import cfg

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.autograd import Variable
import math
import torchvision.models as models


class alexnet(Network):
    def __init__(self):
        Network.__init__(self)
        self._feat_stride = [16, ]
        self._feat_compress = [1. / float(self._feat_stride[0]), ]
        self._net_conv_channels = 256
        self._fc7_channels = 4096

    def _init_head_tail(self):
        self.alexnet = models.alexnet()
        # Remove fc8
        self.alexnet.classifier = nn.Sequential(
            *list(self.alexnet.classifier._modules.values())[:-1])

        # Fix the layers before conv3:
        for layer in range(6):
            for p in self.alexnet.features[layer].parameters():
                p.requires_grad = False

        # not using the last maxpool layer
        self._layers['head'] = nn.Sequential(
            *list(self.alexnet.features._modules.values())[:-1])

    def _image_to_head(self):
        net_conv = self._layers['head'](self._image)
        self._act_summaries['conv'] = net_conv

        return net_conv

    def _head_to_tail(self, pool5):
        pool5_flat = pool5.view(pool5.size(0), -1)
        fc7 = self.alexnet.classifier(pool5_flat)

        return fc7

    def load_pretrained_cnn(self, state_dict):
        if 'state_dict' in state_dict:
            state_dict = state_dict['state_dict']

        new_state = {}
        for k, v in state_dict.items():
            if 'features' in k:
                k = k.replace('module.', '')
            new_state[k] = v

        self.alexnet.load_state_dict(
            {k: v for k, v in new_state.items()
             if k in self.alexnet.state_dict()})
