import math
import os
import pickle

import crick
import pandas as pd
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_info_table_path =  '/scDifformer/data/240412_test/material/gene_info_table.csv'

gene_info = pd.read_csv(gene_info_table_path)
func_gene_list = [i for i in gene_info[(gene_info["gene_type"] == "protein_coding") | (gene_info["gene_type"] == "miRNA")]["ensembl_id"]]
home_path = r'/scDifformer/data/240412_test/tdigest/'
gene_median_dict_path = "/scDifformer/data/240412_test/material/gene_median_dictionary.pkl"


def combine_digests(rootdir: str) -> None:
    # merge new tdigests into total tdigest dict
    def merge_digest(dict_key_ensembl_id, dict_value_tdigest, new_tdigest_dict):
        new_gene_tdigest = new_tdigest_dict.get(dict_key_ensembl_id)
        if new_gene_tdigest is not None:
            dict_value_tdigest.merge(new_gene_tdigest)
            return dict_value_tdigest
        elif new_gene_tdigest is None:
            return dict_value_tdigest


    # initiate tdigests
    median_digests = [crick.tdigest.TDigest() for _ in range(len(func_gene_list))]
    total_tdigest_dict = dict(zip(func_gene_list, median_digests))
    for subdir, dirs, files in os.walk(rootdir):
        for file in files:
            if file.endswith(".gene_median_digest_dict.pickle"):
                logger.info(f"{rootdir}{file} is walking")
                with open(f"{rootdir}{file}", "rb") as fp:
                    tdigest_dict = pickle.load(fp)
                total_tdigest_dict = {k: merge_digest(k, v, tdigest_dict) for k, v in total_tdigest_dict.items()}
    # save dict of merged tdigests
    # with open(f"/mnt/nfs/geneformer/data/pkl/total_gene_tdigest_dict.pkl", "wb") as fp:
    #     pickle.dump(total_tdigest_dict, fp)

    # extract medians and save dict
    total_median_dict = {k: v.quantile(0.5) for k, v in total_tdigest_dict.items()}
    # with open(f"/mnt/nfs/geneformer/data/pkl/total_gene_median_dict.pkl", "wb") as fp:
    #     pickle.dump(total_median_dict, fp)

    # save dict of only detected genes' medians
    detected_median_dict = {k: v for k, v in total_median_dict.items() if not math.isnan(v)}
    with open(gene_median_dict_path, "wb") as fp:
        pickle.dump(detected_median_dict, fp)
    del detected_median_dict, total_median_dict, total_tdigest_dict


if __name__ == '__main__':
    combine_digests(rootdir=home_path)