import json


def load_bigbird_config(config_filepath):
    """Load bigbird config from original tf checkpoint. (pretrain_config.json)"""
    config_key_lst = [
        "attention_probs_dropout_prob",
        "attention_type",
        "block_size",
        "pad_token_id",
        "bos_token_id",
        "eos_token_id",
        "sep_token_id",
        "gradient_checkpointing",  # Originally `use_gradient_checkpointing`
        "hidden_dropout_prob",
        "hidden_size",
        "initializer_range",
        "intermediate_size",
        "max_position_embeddings",
        "num_attention_heads",
        "num_hidden_layers",
        "num_random_blocks",
        "pad_token_id",
        "rescale_embeddings",
        "type_vocab_size",
        "use_bias",
        "vocab_size",
    ]

    config = {
        "position_embedding_type": "absolute",
        "tokenizer_class": "BertTokenizer",  # NOTE Remove this one if you use other tokenizer.
    }

    with open(config_filepath, "r", encoding="utf-8") as f:
        tf_config = json.load(f)

    for config_key in config_key_lst:
        if config_key in tf_config:
            config[config_key] = tf_config[config_key]
        else:
            if config_key == "gradient_checkpointing":
                config[config_key] = tf_config["use_gradient_checkpointing"]
            elif config_key == "num_random_blocks":
                config[config_key] = tf_config["num_rand_blocks"]
            elif config_key == "rescale_embeddings":
                config[config_key] = tf_config["rescale_embedding"]
            else:
                raise KeyError(f"{config_key} not in tensorflow config!!")

    return dict(sorted(config.items()))