import os

import torch
import numpy as np

from modules.config import CFG
from modules.interhand import InterhandEstimator


def main():
    """Convert Interhand weights from native torch format to tensorflow format

    Raises:
        KeyError: when native layer name is not recognized
    """
    # load torch weights
    torch_path = os.path.join(CFG.WEIGHTS_PATH, 'interhand.tar')
    torch_model_dict = torch.load(torch_path, map_location=torch.device('cpu'))
    # load tensorflow model
    tf_model = InterhandEstimator(CFG.INTERHAND_PARAMS.IMG_SHAPE, CFG.INTERHAND_PARAMS.JOINT_HM_DIMS,
                                  CFG.INTERHAND_PARAMS.DEPTH_DIM, CFG.INTERHAND_PARAMS.ROOT_HM_DEPTH)
    tf_model.summary()
    
    weights = []
    # for every weight entry in weight dict
    for key, val in torch_model_dict['network'].items():
        keys = key.split('.')
        # batch norm stats
        if keys[-1] == 'num_batches_tracked':
            continue
        # backbone submodel
        if keys[1] == 'backbone_net':
            # get submodel
            submodel = tf_model.get_layer('ResnetBackbone')
            # convolutional layers - no bias, set weights right away
            if keys[-2].startswith('conv'):
                name = '.'.join(part for part in keys[2:-1])
                layer = submodel.get_layer(name)
                layer.set_weights([np.transpose(val, (2, 3, 1, 0))])
            # batch norm layer - 4 weight objects, save current weights to list
            elif keys[-2].startswith('bn'):
                weights.append(val)
                # if 4 weights in list -> set weights
                if len(weights) == 4:
                    name = '.'.join(part for part in keys[2:-1])
                    layer = submodel.get_layer(name)
                    layer.set_weights(weights)
                    weights = []
            # downsamples layers
            elif keys[-3] == 'downsample':
                # convolutional layer, again no bias -> set right away
                if int(keys[-2]) == 0:
                    name = '.'.join(part for part in keys[2:-1])
                    layer = submodel.get_layer(name)
                    layer.set_weights([np.transpose(val, (2, 3, 1, 0))])
                # batch norm, wait for 4 weights
                else:
                    weights.append(val)
                    # if 4 weights in list -> set weights
                    if len(weights) == 4:
                        name = '.'.join(part for part in keys[2:-1])
                        layer = submodel.get_layer(name)
                        layer.set_weights(weights)
                        weights = []
            else:
                raise KeyError(f'Unrecognised resnet layer name: {key}')
        # decoder branches
        elif keys[1] == 'pose_net':
            # deconvolutional part of heatmap regression
            if keys[2].startswith('joint_deconv'):
                submodel = tf_model.get_layer('PoseNet')
                # deconv layer without bias -> set right away
                if int(keys[-2]) % 3 == 0:
                    name = '.'.join(part for part in keys[1:-1])
                    layer = submodel.get_layer(name)
                    layer.set_weights([np.transpose(val, (2, 3, 1, 0))])
                # batch norm, wait for 4 weights
                else:
                    weights.append(val)
                    # if 4 weights in list -> set weights
                    if len(weights) == 4:
                        name = '.'.join(part for part in keys[1:-1])
                        layer = submodel.get_layer(name)
                        layer.set_weights(weights)
                        weights = []
            # convolutional part of heatmap regression
            elif keys[2].startswith('joint_conv'):
                # convolutional layer with bias, wait for 2 weight objects, then set weights
                weights.append(val)
                if len(weights) == 2:
                    submodel = tf_model.get_layer('PoseNet')
                    name = '.'.join(part for part in keys[1:-1])
                    layer = submodel.get_layer(name)
                    layer.set_weights([np.transpose(weights[0], (2, 3, 1, 0)), weights[1]])
                    weights = []
            # fully connected layers with bias
            elif keys[2].endswith('fc'):
                weights.append(val)
                # if 2 weights in list -> set weights
                if len(weights) == 2:
                    # either handedness classification or relative root depth regression
                    submodel = tf_model.get_layer('HandNet' if keys[2] == 'hand_fc' else 'DepthNet')
                    name = '.'.join(part for part in keys[1:-1])
                    layer = submodel.get_layer(name)
                    layer.set_weights([np.transpose(weights[0]), weights[1]])
                    weights = []
            else:
                raise KeyError(f'Unrecognised layer name: {key}')
        else:
            raise KeyError(f'Unrecognised layer name: {key}')
    # test model and save weights
    tf_model(np.random.rand(1, *CFG.INTERHAND_PARAMS.IMG_SHAPE, 3))
    tf_model.save_weights(CFG.INTERHAND_TF_PATH)


if __name__ == '__main__':
    main()
