from typing import OrderedDict


def backwards_state_dict(state_dict):
    """
    Modify the state dict of older models for backwards compatibility

    Parameters
    ----------
    state_dict : dict
        Model state dict with pretrained weights

    Returns
    -------
    state_dict : dict
        Updated model state dict with modified layer names
    """
    # List of layer names to change
    changes = (('model.model', 'model'),
               ('pose_network', 'pose_net'),
               ('disp_network', 'depth_net'))
    # Iterate over all keys and values
    updated_state_dict = OrderedDict()
    for key, val in state_dict.items():
        # Ad hoc changes due to version changes
        key = '{}.{}'.format('model', key)
        if 'disp_network' in key:
            key = key.replace('conv3.0.weight', 'conv3.weight')
            key = key.replace('conv3.0.bias', 'conv3.bias')
        # Change layer names
        for change in changes:
            key = key.replace('{}.'.format(change[0]),
                              '{}.'.format(change[1]))
        updated_state_dict[key] = val
    # Return updated state dict
    return updated_state_dict