# Copyright (c) OpenMMLab. All rights reserved.
import argparse
from collections import OrderedDict

import torch

convert_dict_s = {
    # backbone
    'model.0': 'backbone.stem',
    'model.1': 'backbone.stage1.0',
    'model.2': 'backbone.stage1.1',
    'model.3': 'backbone.stage2.0',
    'model.4': 'backbone.stage2.1',
    'model.5': 'backbone.stage3.0',
    'model.6': 'backbone.stage3.1',
    'model.7': 'backbone.stage4.0',
    'model.8': 'backbone.stage4.1',
    'model.9': 'backbone.stage4.2',

    # neck
    'model.12': 'neck.top_down_layers.0',
    'model.15': 'neck.top_down_layers.1',
    'model.16': 'neck.downsample_layers.0',
    'model.18': 'neck.bottom_up_layers.0',
    'model.19': 'neck.downsample_layers.1',
    'model.21': 'neck.bottom_up_layers.1',

    # Detector
    'model.22': 'bbox_head.head_module',
}


def convert(src, dst):
    """Convert keys in pretrained YOLOv8 models to mmyolo style."""
    convert_dict = convert_dict_s

    try:
        yolov8_model = torch.load(src)['model']
        blobs = yolov8_model.state_dict()
    except ModuleNotFoundError:
        raise RuntimeError(
            'This script must be placed under the ultralytics repo,'
            ' because loading the official pretrained model need'
            ' `model.py` to build model.'
            'Also need to install hydra-core>=1.2.0 and thop>=0.1.1')
    state_dict = OrderedDict()

    for key, weight in blobs.items():
        num, module = key.split('.')[1:3]
        prefix = f'model.{num}'
        new_key = key.replace(prefix, convert_dict[prefix])

        if '.m.' in new_key:
            new_key = new_key.replace('.m.', '.blocks.')
            new_key = new_key.replace('.cv', '.conv')
        elif 'bbox_head.head_module.proto.cv' in new_key:
            new_key = new_key.replace(
                'bbox_head.head_module.proto.cv',
                'bbox_head.head_module.proto_preds.conv')
        elif 'bbox_head.head_module.proto' in new_key:
            new_key = new_key.replace('bbox_head.head_module.proto',
                                      'bbox_head.head_module.proto_preds')
        elif 'bbox_head.head_module.cv4.' in new_key:
            new_key = new_key.replace(
                'bbox_head.head_module.cv4',
                'bbox_head.head_module.mask_coeff_preds')
            new_key = new_key.replace('.2.weight', '.2.conv.weight')
            new_key = new_key.replace('.2.bias', '.2.conv.bias')
        elif 'bbox_head.head_module' in new_key:
            new_key = new_key.replace('.cv2', '.reg_preds')
            new_key = new_key.replace('.cv3', '.cls_preds')
        elif 'backbone.stage4.2' in new_key:
            new_key = new_key.replace('.cv', '.conv')
        else:
            new_key = new_key.replace('.cv1', '.main_conv')
            new_key = new_key.replace('.cv2', '.final_conv')

        if 'bbox_head.head_module.dfl.conv.weight' == new_key:
            print('Drop "bbox_head.head_module.dfl.conv.weight", '
                  'because it is useless')
            continue
        state_dict[new_key] = weight
        print(f'Convert {key} to {new_key}')

    # save checkpoint
    checkpoint = dict()
    checkpoint['state_dict'] = state_dict
    torch.save(checkpoint, dst)


# Note: This script must be placed under the ultralytics repo to run.
def main():
    parser = argparse.ArgumentParser(description='Convert model keys')
    parser.add_argument(
        '--src', default='yolov8s.pt', help='src YOLOv8 model path')
    parser.add_argument('--dst', default='mmyolov8s.pth', help='save path')
    args = parser.parse_args()
    convert(args.src, args.dst)


if __name__ == '__main__':
    main()