# From EsVIT repo

from esvit.models import build_model

from esvit.config import config, update_config, save_config

import torch

from torch.utils.data import DataLoader, Dataset

from torchvision import transforms

# From EsVIT repo. Need to download full repo. https://github.com/microsoft/esvit

def load_encoder_esVIT(args, device):

    # ============ building network ... ============

    num_features = []

    # if the network is a 4-stage vision transformer (i.e. swin)

    if 'swin' in args.arch :

        update_config(config, args)

        model = build_model(config, is_teacher=True)

        swin_spec = config.MODEL.SPEC

        embed_dim=swin_spec['DIM_EMBED']

        depths=swin_spec['DEPTHS']

        num_heads=swin_spec['NUM_HEADS'] 

        # For each stage, we have n stacked models (d)

        # Each model takes embeddings of dimension embed_dim (the first param), 

        # And then the stage i, input dim is input dim(i-1)*2

        for i, d in enumerate(depths):

            num_features += [int(embed_dim * 2 ** i)] * d 

    # if the network is a 4-stage vision transformer (i.e. longformer)

    elif 'vil' in args.arch :

        update_config(config, args)

        model = build_model(config, is_teacher=True)

        msvit_spec = config.MODEL.SPEC

        arch = msvit_spec.MSVIT.ARCH

        layer_cfgs = model.layer_cfgs

        num_stages = len(model.layer_cfgs)

        depths = [cfg['n'] for cfg in model.layer_cfgs]

        dims = [cfg['d'] for cfg in model.layer_cfgs]

        out_planes = model.layer_cfgs[-1]['d']

        Nglos = [cfg['g'] for cfg in model.layer_cfgs]

        print(dims)

        for i, d in enumerate(depths):

            num_features += [ dims[i] ] * d

    # if the network is a 4-stage vision transformer (i.e. CvT)

    elif 'cvt' in args.arch :

        update_config(config, args)

        model = build_model(config, is_teacher=True)

        cvt_spec = config.MODEL.SPEC

        embed_dim=cvt_spec['DIM_EMBED']

        depths=cvt_spec['DEPTH']

        num_heads=cvt_spec['NUM_HEADS'] 

        print(f'embed_dim {embed_dim} depths {depths}')

        for i, d in enumerate(depths):

            num_features += [int(embed_dim[i])] * int(d) 

    # if the network is a vanilla vision transformer (i.e. deit_tiny, deit_small, vit_base)

    else:

        raise ValueError(f'{args.arch} not supported yet.')

    model.to(device)

    # load weights to evaluate

    state_dict = torch.load(args.checkpoint, map_location=device)

    # Technically we can also load the weights of the student but in knowledge distillation, I think it's more common to take the teacher 

    # and in DINO paper, they show that the teacher learns better. 

    state_dict = state_dict['teacher']

    #Line below was initally in the code but I think it's usefless in our case (swin-t)

    #state_dict = {k.replace("module.", ""): v for k, v in state_dict.items()}

    #in trained model, you probably have the dense DINO head and in the loaded one a regular head. Those keys won't be matching.

    #IncompatibleKeys(missing_keys=['head.weight', 'head.bias'], unexpected_keys=['head_dense.mlp.0.weight', 'head_dense.mlp.0.bias', 'head_dense.mlp.2.weight', 'head_dense.mlp.2.bias', 'head_dense.mlp.4.weight', 'head_dense.mlp.4.bias', 'head_dense.last_layer.weight_g', 'head_dense.last_layer.weight_v', 'head.mlp.0.weight', 'head.mlp.0.bias', 'head.mlp.2.weight', 'head.mlp.2.bias', 'head.mlp.4.weight', 'head.mlp.4.bias', 'head.last_layer.weight_g', 'head.last_layer.weight_v'])

    #in any case, we do not use the heads but the out features of each stage.

    msg = model.load_state_dict(state_dict, strict=False)

    print(msg)

    model.eval()

    print(f"Model {args.arch} {args.patch_size}x{args.patch_size} built with pretrained weigths {args.checkpoint}.")

    ##a choice, 4 will take the last 2 stages for instance

    #if n>1, they are just stacked features.

    #paper says : For all transformers architecture, we use the concatenation of view-level features

    # in the last layers (results are similar to the use of 3 or 5 layers in our initial experiments)

    num_features_linear = sum(num_features[-args.n_last_blocks:])

    print(f'num_features_linear {num_features_linear}')

    return model, num_features_linear, depths

# Regular resnet encoder. 

def load_encoder_resnet(backbone, checkpoint_file, use_imagenet_weights, device):

    import torch.nn as nn

    import torchvision.models as models

    class DecapitatedResnet(nn.Module):

        def __init__(self, base_encoder, pretrained):

            super(DecapitatedResnet, self).__init__()

            self.encoder = base_encoder(pretrained=pretrained)

        def forward(self, x):

            # Same forward pass function as used in the torchvision 'stock' ResNet code

            # but with the final FC layer removed.

            x = self.encoder.conv1(x)

            x = self.encoder.bn1(x)

            x = self.encoder.relu(x)

            x = self.encoder.maxpool(x)

            x = self.encoder.layer1(x)

            x = self.encoder.layer2(x)

            x = self.encoder.layer3(x)

            x = self.encoder.layer4(x)

            x = self.encoder.avgpool(x)

            x = torch.flatten(x, 1)

            return x

    model = DecapitatedResnet(models.__dict__[backbone], use_imagenet_weights)

    if use_imagenet_weights:

        if checkpoint_file is not None:

            raise Exception(

                "Either provide a weights checkpoint or the --imagenet flag, not both."

            )

        print(f"Created encoder with Imagenet weights")

    else:

        checkpoint = torch.load(checkpoint_file, map_location="cpu")

        state_dict = checkpoint["state_dict"]

        for k in list(state_dict.keys()):

            # retain only encoder_q up to before the embedding layer

            if k.startswith("module.encoder_q") and not k.startswith(

                "module.encoder_q.fc"

            ):

                # remove prefix from key names

                state_dict[k[len("module.encoder_q.") :]] = state_dict[k]

            # delete renamed or unused k

            del state_dict[k]

        # Verify that the checkpoint did not contain data for the final FC layer

        msg = model.encoder.load_state_dict(state_dict, strict=False)

        assert set(msg.missing_keys) == {"fc.weight", "fc.bias"}

        print(f"Loaded checkpoint {checkpoint_file}")

    model = model.to(device)

    if torch.cuda.device_count() > 1:

        model = torch.nn.DataParallel(model)

    model.eval()

    return model