
from torch import nn
import torch.nn.functional as F
from utils.fmodule import FModule


class Model(FModule):
    def __init__(self):
        super(Model, self).__init__()
        self.input_require_shape = [1, -1, -1]
        self.task = 'fashion'
        # feature extractor
        self.fc1 = nn.Linear(784, 200)
        self.fc2 = nn.Linear(200, 200)
        self.fc3 = nn.Linear(200, 10)
        self.activation = nn.ReLU()

        self.shared_model = nn.Sequential(
            nn.Linear(784, 200),
            nn.ReLU(),
            nn.Linear(200, 200),
            nn.ReLU(),
        )
        self.local_model = nn.Sequential(
            nn.Linear(200, 10)
        )
        # self.fc1 = nn.Linear(784, 200)
        # self.fc2 = nn.Linear(200, 200)
        # self.fc3 = nn.Linear(200, 10)

    def features_extras(self, x):
        x = x.view(-1, x.shape[1] * x.shape[-2] * x.shape[-1])
        x = self.activation(self.fc1(x))
        x = self.activation(self.fc2(x))
        return x

    def classifier_head(self, x):
        return self.fc3(x)

    def forward(self, x):
        x = self.features_extras(x)
        x = self.classifier_head(x)
        return x
    #
    # def get_embedding(self, x):
    #     x = x.view(-1, x.shape[1] * x.shape[-2] * x.shape[-1])
    #     x = self.fc1(x)
    #     x = F.relu(x)
    #     x = self.fc2(x)
    #     x = F.relu(x)
    #     return x
