import torch.nn as nn
from efficientnet_pytorch import EfficientNet

from config import *


class Efficientnet(nn.Module):
    def __init__(self, args):
        super().__init__()

        # EfficientNetの学習済みモデルをロード
        self.network = EfficientNet.from_pretrained('efficientnet-b0', num_classes=NUM_CLASS)

        self.network._fc = nn.Identity()

        '''
        # extract_features()でmodel.network.eval()を行うから分ける必要がない気もする(現状は一応やってる)
        if args.metric == 'softmax':
            self.network._fc = nn.Sequential(nn.Linear(self.network._fc.in_features, args.num_features),
                                             nn.ReLU(),
                                             nn.Dropout(0.50))
        else:
            self.network._fc = nn.Sequential(nn.Linear(self.network._fc.in_features, args.num_features),
                                             nn.ReLU())
        '''

    def forward(self, x):
        outputs = self.network(x)
        # x = self.network._avg_pooling(x)
        # x = x.flatten(start_dim=1)
        # outputs = self.network._dropout(x)
        return outputs

    # 学習した特徴量を抽出する
    def extract_features(self, data_loader):

        # EfficientNetの最終層を切り落とす(上のコメントの関係でいらなくなるかも)
        # model_lt = list(self.network._fc)
        # self.network._fc = nn.Sequential(model_lt[0], model_lt[1])

        self.network.eval()
        with torch.no_grad():
            features_list = []
            for data in data_loader:
                images, labels = data

                images = images.to(DEVICE)
                labels = labels.to(DEVICE)

                feature = self.network.extract_features(images)
                feature = feature.flatten(start_dim=1)

                feature = feature.cpu()
                features_list.append(feature)

            features = torch.cat(features_list, dim=0)

        return features
