import torch
import torch.nn as nn
from cnn_finetune import make_model
from timm import create_model


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

    for param in self._classifier.parameters():
        param.requires_grad = True


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


def make_classifier(in_features, num_classes):
    return nn.Sequential(
        nn.Linear(in_features, 512),
        nn.Dropout(0.3),
        nn.Linear(512, num_classes),
    )


class MultiModals(nn.Module):
    def __init__(self, model_name, pretrained=True, num_classes=6, dropout_p=None):
        super(MultiModals, self).__init__()
        self.model = make_model(
            model_name=model_name,
            num_classes=num_classes,
            pretrained=pretrained,
            dropout_p=dropout_p,
            # classifier_factory=make_classifier
        )

        in_features = self.model._classifier.in_features

        self._classifier = nn.Sequential(
            nn.Linear(in_features + 8, 512),
            nn.Dropout(0.3),
            nn.Linear(512, num_classes),
        )

        setattr(self, 'freeze', cnnfinetune_freeze)
        setattr(self, 'unfreeze', cnnfinetune_unfreeze)

    def forward(self, images, meta):
        x = self.model._features(images)
        x = self.model.pool(x)
        x = x.view(x.size(0), -1)
        # import pdb
        # pdb.set_trace()
        # if isinstance(x, torch.HalfTensor):
        #     meta = meta.half()
        x = torch.cat([x, meta], dim=1)
        return self._classifier(x)