import torch
import torch.nn as nn

from nets.siamese import Siamese


class stacking(nn.Module):
    def __init__(self):
        super(stacking, self).__init__()
        model_path = "weight/{}_best_epoch_weights.pth"
        device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

        resnet18 = Siamese(pretrained=True, model_name="resnet18")
        resnet18.load_state_dict(torch.load(model_path.format("resnet18"), map_location=device))
        self.resnet18 = resnet18.eval()

        resnet34 = Siamese(model_name="resnet34")
        resnet34.load_state_dict(torch.load(model_path.format("resnet34"), map_location=device))
        self.resnet34 = resnet34.eval()

        resnet50 = Siamese(pretrained=True, model_name="resnet50")
        resnet50.load_state_dict(torch.load(model_path.format("resnet50"), map_location=device))
        self.resnet50 = resnet50.eval()

    def forward(self, x):

        with torch.no_grad():
            output1 = self.resnet18(x).detach()
            output2 = self.resnet34(x).detach()
            output3 = self.resnet50(x).detach()

            pass
        return x
