"""
Convert weights from pretrained models to format used in this project
"""
import os
os.environ["CUDA_VISIBLE_DEVICES"] = "-1"

import torch
import numpy as np
import h5py

from modules.config import CFG
from modules.yolov5 import YOLOv5
from modules.efficientdet import EfficientDet


def get_layer(model, torch_name):
    """Get YOLOv5 layer by name from torch weights

    Args:
        model (tf.keras.Model): native model
        torch_name (str): torch weight name

    Raises:
        ValueError: If weight name is not recognized and native layer equivalent is not found

    Returns:
        tf.keras.Layer: found layer
    """
    keys = torch_name.split('.')
    name = f"yolov5.{'.'.join(x for x in keys[1:-1])}"
    for submodel in model.layers[1:]:
        try:
            layer = submodel.get_layer(name)
            break
        except ValueError:
            continue
    else:
        raise ValueError(f'Unrecognised layer {torch_name}')
    return layer


def convert_yolov5_weights(model_cfg):
    """Convert YOLOv5 weights from torch pretrained format to native tensorflow format

    Args:
        model_cfg (str): YOLOv5 model configuration, see model implementation for implemented configurations
    """
    # load torch weights
    torch_path = os.path.join(CFG.WEIGHTS_PATH, f'pretrained_{model_cfg}.tar')
    torch_state_dict = torch.load(torch_path, map_location=torch.device('cpu'))
    # load tensorflow model
    tf_model = YOLOv5(CFG.YOLO_PRETRAINED_IMG_SHAPE, CFG.YOLO_PRETRAINED_CLASSES, CFG.FACES_PARAMS.ANCHORS, model=model_cfg, training=False)
    # iterate over torch weight items
    weights = list()
    for name, tensor in torch_state_dict.items():
        if 'anchor' in name:
            continue
        keys = name.split('.')
        # convolutional layers without bias -> one weight tensor
        if keys[-2] == 'conv':
            # get layer and set weights
            layer = get_layer(tf_model, name)
            layer.set_weights([np.transpose(tensor, (2, 3, 1, 0))])
        # batch norm layers without -> four weight tensors next to each other
        elif keys[-2] == 'bn':
            if keys[-1] == 'num_batches_tracked':
                continue
            # save weight to list
            weights.append(tensor)
            # when four weights are saved, get layer, set weights and reset weights list
            if len(weights) == 4:
                layer = get_layer(tf_model, name)
                layer.set_weights(weights)
                weights = list()
        # convolutional layers with bias -> two weight tensors
        else:
            # save weight to list
            weights.append(tensor)
            # when two weights are saved, get layer, set weights and reset weights list
            if len(weights) == 2:
                layer = get_layer(tf_model, name)
                weights[0] = np.transpose(weights[0], (2, 3, 1, 0))
                layer.set_weights(weights)
                weights = list()
    # test model inference and save converted weights
    tf_model(np.random.rand(1, *CFG.YOLO_PRETRAINED_IMG_SHAPE))
    tf_model.save_weights(CFG.PRETRAINED_TF_WEIGHTS_PATH[model_cfg])


def assign_weights(layer, weights_dict, override_name=None):
    """Assign H5 weights to a tensorflow layer

    Args:
        layer (tf.keras.Layer): target layer
        weights_dict (dict): H5 weight dictionary
        override_name (str, optional): if weight name does not clearly correspond to a specific layer, this argument can override it. Defaults to None.

    Raises:
        ValueError: if weight format is not understood
    """
    # don't assign output layer weights
    if 'predict' in layer.name:
        return
    # normalize weight names, if override is not set, find some pattern in weight name
    aliases = {
        'bn': 'batchnorm',
        'box': 'separable',
        'class': 'separable',
        'depth': 'depth'
    }
    if override_name is None:
        name = layer.name
        for key, val in aliases.items():
            if key in name:
                name = val
                break
    else:
        name = override_name
    # dict specifying weight order
    weights_order = {
        'batch': ['gamma:0', 'beta:0', 'moving_mean:0', 'moving_variance:0'],
        'separable': ['depthwise_kernel:0', 'pointwise_kernel:0', 'bias:0'],
        'depth': ['depthwise_kernel:0'],
        'conv': ['kernel:0', 'bias:0'],
    }
    # get the weights in the right order and set them to the layer
    weights = list()
    for key, val in weights_order.items():
        if key in name:
            for v in val:
                try:
                    weights.append(weights_dict[v])
                except KeyError:
                    continue
            try:
                layer.set_weights(weights)
            except ValueError:
                return
            return
    raise ValueError(f'Unknown weights - {list(weights_dict.keys())} at {layer.name}')


def convert_efficientdet_weights(model_cfg):
    """Convert EfficientDet weights from H5 pretrained format to native tensorflow format

    Args:
        model_cfg (str): EfficientDet model configuration, see model implementation for implemented configurations
    """
    # h5 weights path
    h5_path = os.path.join(CFG.WEIGHTS_PATH, f'pretrained_{model_cfg}.h5')
    # load native tensorflow model
    model = EfficientDet(
        (*CFG.COMBINED_PARAMS.IMG_SHAPE, 3),
        2, CFG.COMBINED_PARAMS.EFFDET_ANCHORS, model=model_cfg
    )
    model.build((1, *CFG.COMBINED_PARAMS.IMG_SHAPE, 3))
    # read h5 file
    with h5py.File(h5_path, 'r') as file:
        for submodel_name in file.keys():
            try:
                submodel_dict = file[submodel_name][submodel_name]
            except KeyError:
                continue
            # organized into submodels
            submodel = model.get_layer(submodel_name)
            # efficientnet backbone submodel
            if 'efficientnet' in submodel_name:
                # for every layer block in submodel
                for block_name, block_dict in submodel_dict.items():
                    # for every layer in block
                    for layer_name, layer_dict in block_dict.items():
                        # get layer and assign its weights
                        if layer_name == 'se':
                            for sublayer_name, sublayer_dict in layer_dict.items():
                                sublayer = submodel.get_layer(f'{block_name}.{layer_name}.{sublayer_name}')
                                assign_weights(sublayer, sublayer_dict)
                        else:
                            layer = submodel.get_layer(f'{block_name}.{layer_name}')
                            assign_weights(layer, layer_dict)
            # FPN submodel
            elif submodel_name == 'fpn_cells':
                # for every cell in FPN
                for cell_name, cell_dict in submodel_dict.items():
                    cell = submodel.get_layer(cell_name)
                    # for every node in cell
                    for node_name, node_dict in cell_dict.items():
                        node = cell.get_layer(node_name)
                        # for every layer in node
                        for layer_name, layer_dict in node_dict.items():
                            if 'resample' in layer_name:
                                layer = node.get_layer(layer_name)
                                # for every sublayer in layer
                                for sublayer_name, sublayer_dict in layer_dict.items():
                                    # get layer and assign its weights
                                    sublayer = layer.get_layer(sublayer_name)
                                    assign_weights(sublayer, sublayer_dict)
                            if 'op_after_combine' in layer_name:
                                # for every sublayer in layer
                                for sublayer_name, sublayer_dict in layer_dict.items():
                                    # get layer and assign its weights
                                    sublayer = node.get_layer(sublayer_name)
                                    assign_weights(sublayer, sublayer_dict, 'separable' if 'bn' not in sublayer.name else None)
                            else:
                                pass  # dont transfer attention weights
            # other submodels
            else:
                # for every layer in submodel
                for layer_name, layer_dict in submodel_dict.items():
                    # get layer and assign its weights
                    layer = submodel.get_layer(layer_name)
                    assign_weights(layer, layer_dict)
    # test model inference and save converted weights
    model(np.random.rand(1, *CFG.COMBINED_PARAMS.IMG_SHAPE, 3))
    model.save_weights(CFG.PRETRAINED_TF_WEIGHTS_PATH[model_cfg])


if __name__ == '__main__':
    convert_yolov5_weights('yolov5s')
    convert_yolov5_weights('yolov5m')
    convert_yolov5_weights('yolov5l')
    convert_efficientdet_weights('efficientdet-d0')
    convert_efficientdet_weights('efficientdet-d1')
    convert_efficientdet_weights('efficientdet-d2')
    convert_efficientdet_weights('efficientdet-d3')
    convert_efficientdet_weights('efficientdet-d4')
    convert_efficientdet_weights('efficientdet-d5')
    convert_efficientdet_weights('efficientdet-d6')
