# Copyright (c) OpenMMLab. All rights reserved.

import argparse

import os.path as osp

from collections import OrderedDict

import mmcv

import torch

from mmcv.runner import CheckpointLoader

def convert_mit(ckpt):

    new_ckpt = OrderedDict()

    # Process the concat between q linear weights and kv linear weights

    for k, v in ckpt.items():

        if k.startswith('head'):

            continue

        # patch embedding conversion

        elif k.startswith('patch_embed'):

            stage_i = int(k.split('.')[0].replace('patch_embed', ''))

            new_k = k.replace(f'patch_embed{stage_i}', f'layers.{stage_i-1}.0')

            new_v = v

            if 'proj.' in new_k:

                new_k = new_k.replace('proj.', 'projection.')

        # transformer encoder layer conversion

        elif k.startswith('block'):

            stage_i = int(k.split('.')[0].replace('block', ''))

            new_k = k.replace(f'block{stage_i}', f'layers.{stage_i-1}.1')

            new_v = v

            if 'attn.q.' in new_k:

                sub_item_k = k.replace('q.', 'kv.')

                new_k = new_k.replace('q.', 'attn.in_proj_')

                new_v = torch.cat([v, ckpt[sub_item_k]], dim=0)

            elif 'attn.kv.' in new_k:

                continue

            elif 'attn.proj.' in new_k:

                new_k = new_k.replace('proj.', 'attn.out_proj.')

            elif 'attn.sr.' in new_k:

                new_k = new_k.replace('sr.', 'sr.')

            elif 'mlp.' in new_k:

                string = f'{new_k}-'

                new_k = new_k.replace('mlp.', 'ffn.layers.')

                if 'fc1.weight' in new_k or 'fc2.weight' in new_k:

                    new_v = v.reshape((*v.shape, 1, 1))

                new_k = new_k.replace('fc1.', '0.')

                new_k = new_k.replace('dwconv.dwconv.', '1.')

                new_k = new_k.replace('fc2.', '4.')

                string += f'{new_k} {v.shape}-{new_v.shape}'

        # norm layer conversion

        elif k.startswith('norm'):

            stage_i = int(k.split('.')[0].replace('norm', ''))

            new_k = k.replace(f'norm{stage_i}', f'layers.{stage_i-1}.2')

            new_v = v

        else:

            new_k = k

            new_v = v

        new_ckpt[new_k] = new_v

    return new_ckpt

def main():

    parser = argparse.ArgumentParser(

        description='Convert keys in official pretrained segformer to '

        'MMSegmentation style.')

    parser.add_argument('src', help='src model path or url')

    # The dst path must be a full path of the new checkpoint.

    parser.add_argument('dst', help='save path')

    args = parser.parse_args()

    checkpoint = CheckpointLoader.load_checkpoint(args.src, map_location='cpu')

    if 'state_dict' in checkpoint:

        state_dict = checkpoint['state_dict']

    elif 'model' in checkpoint:

        state_dict = checkpoint['model']

    else:

        state_dict = checkpoint

    weight = convert_mit(state_dict)

    mmcv.mkdir_or_exist(osp.dirname(args.dst))

    torch.save(weight, args.dst)

if __name__ == '__main__':

    main()