import os
import pickle
from collections import OrderedDict

from loguru import logger


def get_filename() -> str:
    current_file_name = os.path.basename(__file__)
    log_name = "{}.log"
    return log_name.format(current_file_name.split('.')[0])


def set_log() -> None:
    filename = get_filename()
    logger.add(f'../log/{filename}')


set_log()

gene_median_path = "/scDifformer/data/240412_test/material/gene_median_dictionary.pkl"
token_dictionary_path = "/scDifformer/data/240412_test/material/token_dictionary.pkl"

if __name__ == '__main__':

    with open(gene_median_path, 'rb') as fb:
        median_data = pickle.load(fb)

    pkl_dict = OrderedDict()
    pkl_dict['<pad>'] = 0
    pkl_dict['<mask>'] = 1

    i = 2
    for gene_id, val in median_data.items():
        pkl_dict[gene_id] = i
        i += 1
    with open(token_dictionary_path, "wb") as fp:
        pickle.dump(dict(pkl_dict), fp)

    with open(token_dictionary_path, 'rb') as fb:
        data = pickle.load(fb)

    print(i)