import pandas as pd
import os
import cooler
import numpy as np
import tqdm
from scipy.stats import norm
import cooltools
import warnings

warnings.filterwarnings("ignore")
from sc3dg.analysis.embedding import scale_matrix, gaussian_distance, cosine_similarity, PCAEmbedding, MDSEmbedding, tSNEEmbedding, \
    UMAPEmbedding, LLEEmbedding, IsomapEmbedding
import bioframe
import sys
import sc3dg.analysis.genome as genome
from scipy.special import comb, perm
import random
from cooltools import insulation
import random
import tqdm
import shutil
import time
from joblib import Parallel, delayed
hg38_cens = pd.DataFrame({'chrom': ['chr1', 'chr10', 'chr11', 'chr12', 'chr13', 'chr14', 'chr15', 'chr16', 'chr17',
                                    'chr18', 'chr19', 'chr2', 'chr20', 'chr21', 'chr22', 'chr3', 'chr4', 'chr5', 'chr6',
                                    'chr7', 'chr8', 'chr9', 'chrX', 'chrY'],
                          'start': [122026459, 39686682, 51078348, 34769407, 16000000, 16000000, 17083673, 36311158,
                                    22813679, 15460899, 24498980, 92188145, 26436232, 10864560, 12954788, 90772458,
                                    49712061, 46485900, 58553888, 58169653, 44033744, 43389635, 58605579, 10316944],
                          'end': [124932724, 41593521, 54425074, 37185252, 18051248, 18173523, 19725254, 38265669,
                                  26616164, 20861206, 27190874, 94090557, 30038348, 12915808, 15054318, 93655574,
                                  51743951, 50059807, 59829934, 61528020, 45877265, 45518558, 62412542, 10544039],
                          'mid': [123479591, 40640101, 52751711, 35977329, 17025624, 17086761, 18404463, 37288413,
                                  24714921, 18161052, 25844927, 93139351, 28237290, 11890184, 14004553, 92214016,
                                  50728006, 48272853, 59191911, 59848836, 44955504, 44454096, 60509060, 10430491]})

mm10_cens = pd.DataFrame(columns=['chrom', 'start', 'end', 'mid'])
mm10_cens['start'] = mm10_cens['start'].astype('int64')
mm10_cens['end'] = mm10_cens['end'].astype('int64')
mm10_cens['mid'] = mm10_cens['mid'].astype('int64')

mm10_chromsize = pd.DataFrame({'chrom': ['chr1', 'chr2', 'chr3', 'chr4', 'chr5', 'chr6', 'chr7', 'chr8', 'chr9',
                                         'chr10', 'chr11', 'chr12', 'chr13', 'chr14', 'chr15', 'chr16', 'chr17',
                                         'chr18', 'chr19', 'chrX', 'chrY', 'chrM'],
                               'length': [195471971, 182113224, 160039680, 156508116, 151834684, 149736546, 145441459,
                                          129401213, 124595110, 130694993, 122082543, 120129022, 120421639, 124902244,
                                          104043685, 98207768, 94987271, 90702639, 61431566, 171031299, 91744698,
                                          16299]})
hg38_chromsize = pd.DataFrame({'chrom': ['chr1', 'chr2', 'chr3', 'chr4', 'chr5', 'chr6', 'chr7', 'chr8', 'chr9',
                                         'chr10', 'chr11', 'chr12', 'chr13', 'chr14', 'chr15', 'chr16', 'chr17',
                                         'chr18', 'chr19', 'chr20', 'chr21', 'chr22', 'chrX', 'chrY', 'chrM'],
                               'length': [248956422, 242193529, 198295559, 190214555, 181538259, 170805979, 159345973,
                                          145138636, 138394717, 133797422, 135086622, 133275309, 114364328, 107043718,
                                          101991189, 90338345, 83257441, 80373285, 58617616, 64444167, 46709983,
                                          50818468, 156040895, 57227415, 16569]})
data_info = {'hg38': [hg38_chromsize, hg38_cens], 'mm10': [mm10_chromsize, mm10_cens]}


class MMCooler:
    def __init__(self,
                 directory,  # 目录
                 resolution,  # 分辨率
                 describe=None,  # 技术名称
                 merge=True,  # 是否将读取的cell全部merge
                 genome: str = 'mm10',
                 balance: bool = True,
                 read_cache=True,
                 n_jobs=6,
                 ):
        '''
            directory:
            resolution:如果int,那没什么，就是直接读取；如果是list，list第一个必须要有，切是最小的，后面的分辨率会按照这个去合并或者读取
            describe:
            merge:
            genome:
            balance
        '''
        self.directory = directory
        self.n_jobs = n_jobs
        self.mcool_files = []
        self.mcool_names = []
        self.cache_file = []
        self.clr_obj = None

        self.__merge = merge
        self.genome = genome
        self.__balance = balance

        self.resolution = resolution if isinstance(resolution, list) else [resolution]
        self.describe = describe
        self.chromosome = ['chr1', 'chr2', 'chr3', 'chr4', 'chr5', 'chr6',
                           'chr7', 'chr8', 'chr9', 'chr10', 'chr11', 'chr12',
                           'chr13', 'chr14', 'chr15', 'chr16', 'chr17', 'chr18',
                           'chr19', 'chr20', 'chr21', 'chr22', 'chrX', 'chrY', 'chrM']
        self.__chrom_terr = ['chr1', 'chr2', 'chr3', 'chr4', 'chr5', 'chr6',
                             'chr7', 'chr8', 'chr9', 'chr10', 'chr11', 'chr12',
                             'chr13', 'chr14', 'chr15', 'chr16', 'chr17', 'chr18',
                             'chr19']
        self.read_cache = read_cache
        self.__pool = False

        if not os.path.exists('./cache'):
            os.makedirs('./cache')
        self.extract_mcool_files()
        self.read_mcool_files()
        if not os.path.exists('./file'):
            os.makedirs('./file')

    def stat_merge(self):
        return self.__merge

    def stat_balance(self):
        return self.__balance

    def stat_resolution(self):
        return isinstance(self.resolution, list)

    def shape(self):
        shape = np.array(self.clr_obj)
        return shape

    def stat_file_name(self):
        return self.mcool_names

    def __read_zoomify(self, mcool_file, mcool_name,cache_name):
        raw_res = self.resolution[0] 
        if not os.path.exists(cache_name):
            cooler.zoomify_cooler(base_uris=mcool_file + "::resolutions/" + str(raw_res), outfile=cache_name,
                                    resolutions=self.resolution,
                                    chunksize=10000, nproc=16)

    # 读取目录下所有的mcool文件
    def extract_mcool_files(self):
        if self.directory.endswith('.mcool'):
            self.mcool_files.append(self.directory)
            self.mcool_names.append(self.directory.split('/')[-1].split('.')[0])

        for root, dir_, files in os.walk(self.directory):
            for f in files:
                if f.endswith('.mcool'):
                    self.mcool_files.append(root + '/' + f)
                    self.mcool_names.append(f.split('.')[0])

    # 加载所有mcool文件到cooler.Cooler对象,如果merge不是none，那就merge到这个
    def read_mcool_files(self):
        '''
            self.coolers: list of cooler objects
        '''
        # 如果合并
        if self.__merge:
            print('run merge mode')
            # 多个分辨率
            self.clr_obj = []
            raw_res = self.resolution[0]  # 读取有这个分辨率,并以此分辨率座位基础

            # 先把每个样本bin化，然后再读取
            print('binning the raw mcool files')
            for mcool_file, mcool_name in zip(self.mcool_files, self.mcool_names):
                # 如果美有cache，那就创建，读取
                # 创建对应样本cache的名字
                cache_name = './cache/%s_multi_res_%s' % (self.describe, mcool_name)
                for res in self.resolution:
                    cache_name += '-'
                    cache_name += str(res)
                cache_name += '.mcool'

                # 如果没有这个文件，那就生成
                if not os.path.exists(cache_name):
                    cooler.zoomify_cooler(base_uris=mcool_file + "::resolutions/" + str(raw_res), outfile=cache_name,
                                          resolutions=self.resolution,
                                          chunksize=10000, nproc=32)

            # 都存到cache了，然后按照分辨率读取,并 merge
            print('reading the binned mcool files')
            for res in self.resolution:
                tmp = []
                for mcool_file, mcool_name in zip(self.mcool_files, self.mcool_names):
                    cache_name = './cache/%s_multi_res_%s' % (self.describe, mcool_name)
                    for rr in self.resolution:
                        cache_name += '-'
                        cache_name += str(rr)
                    cache_name += '.mcool'
                    tmp.append(cache_name + '::resolutions/' + str(res))

                print('merging the binned mcool files/resolution=%s' % res)
                if not os.path.exists('./cache/%s_merge_all_res-%s' % (self.describe, res)):
                    cooler.merge_coolers(input_uris=tmp,
                                         output_uri='./cache/%s_merge_all_res-%s' % (self.describe, res),
                                         mergebuf=1000000)
                self.clr_obj.append(cooler.Cooler('./cache/%s_merge_all_res-%s' % (self.describe, res)))

        # 如果不合并
        else:
            print('run nomerge mode')
            self.clr_obj = []
            # 多个分辨率

            # 先把每个样本bin化，然后再读取
            raw_res = self.resolution[0]  # 读取有这个分辨率,并以此分辨率座位基础
            print('binning the raw mcool files')
            task = []
            for mcool_file, mcool_name in zip(self.mcool_files, self.mcool_names):
                cache_name = './cache/%s_multi_nomerge_%s_res' % (self.describe, mcool_name)
                for res in self.resolution:
                    cache_name += '-'
                    cache_name += str(res)
                cache_name += '.mcool'
                task.append([ mcool_file, mcool_name, cache_name])
            Parallel(n_jobs=self.n_jobs)(delayed(self.__read_zoomify)(*t) for t in task)
            
            # 都存到cache了，然后按照分辨率读取
            # 外层是分辨率，里层是样本
            print('reading the binned mcool files')
            for res in self.resolution:
                print(res)
                tmp = []
                for mcool_file, mcool_name in zip(self.mcool_files, self.mcool_names):
                    cache_name = './cache/%s_multi_nomerge_%s_res' % (self.describe, mcool_name)
                    for r in self.resolution:
                        cache_name += '-'
                        cache_name += str(r)
                    cache_name += '.mcool'
                    tmp.append(cooler.Cooler(cache_name + '::resolutions/' + str(res)))
                    if cache_name not in self.cache_file:
                        self.cache_file.append(cache_name)
                self.clr_obj.append(tmp)

    # 获取单个细胞所有contact数量
    def get_total_contact(self, res=None):
        '''
            res: resolution
        '''
        if res is None:
            clr = self.clr_obj
        else:
            clr = [self.clr_obj[self.resolution.index(res)]]

        res = []
        for cr in clr[0]:
            res.append(cr.info['sum'])
        df = pd.DataFrame(res, index=self.mcool_names, columns=['total_contact'])
        df.to_csv('./file/{}_total_contact.csv'.format(self.describe), sep='\t')
        return res

    # 输入一个cooler，返回一组结果
    def __get_contact_vs_distance(self, clr, res, arms):
        resolution = res
        arms = arms[arms.chrom.isin(clr.chromnames)].reset_index(drop=True)
        if not cooltools.lib.is_cooler_balanced(clr):
            cooler.balance_cooler(clr, store=True)
        cvd_smooth_agg = cooltools.expected_cis(
            clr=clr,
            view_df=arms,
            smooth=True,
            aggregate_smoothed=True,
            nproc=16
        )

        cvd_smooth_agg['s_bp'] = cvd_smooth_agg['dist'] * resolution
        cvd_smooth_agg['balanced.avg.smoothed.agg'].loc[cvd_smooth_agg['dist'] < 2] = np.nan
        cvd_smooth_agg = cvd_smooth_agg.groupby('s_bp').sum().reset_index()
        cvd_smooth_agg = cvd_smooth_agg.set_index('s_bp')

        return cvd_smooth_agg.loc[:, ['balanced.avg.smoothed.agg']]

    # 获得绘制P（s）图像的数据
    def get_contact_vs_distance(self, res=None, version='hg38'):
        '''
            dismiss:舍弃对角线开始的第几个细胞
            mode:
                cis:   interaction decay with distance
                 or trans:teraction frequencies for inter-chromosomal blocks
            intra_only:  Return expected only for symmetric intra-regions defined by view_df,
            i.e. chromosomes, chromosomal-arms, intra-domains, etc.
            When False returns expected both for symmetric intra-regions and
            assymetric inter-regions.
        '''

        if res is None:
            sys.exit('res is None')

        chromsizes, cens = data_info[version]
        arms = bioframe.make_chromarms(chromsizes, cens)

        clr = self.clr_obj[self.resolution.index(res)]
        result = []
        for name, val in zip(self.mcool_names, clr):
            tmp = self.__get_contact_vs_distance(val, res, arms)
            tmp.columns = [name]
            result.append(tmp)
        result = pd.concat(result, axis=1)

        return result

    # 获得intra-chrom contact个数
    def get_intra_contact_count(self, split_chorm: bool = True):
        '''
            merge + split: 返回染色体dict
            merge + nosolt:返回一个值
            nomerge + nosplit:返回list of count
            nomerge + split:返回 list of dict
        '''

        counts = []

        for val in self.clr_obj[-1]:
            counts.append(self.__get_intra_contact_count(val, split_chorm))

        df = pd.DataFrame(counts)
        df.index = self.mcool_names

        df.to_csv('./file/{}_splitChrom:{}_intra_contact_count.csv'.format(self.describe, split_chorm), sep='\t')

        return counts

    # 一个cool，一个dict
    def __get_intra_contact_count(self, clr, split_chorm: bool = True):

        bins = clr.bins()[:]
        pixels = clr.pixels()[:]
        pixels['chr1'] = list(bins.iloc[list(pixels['bin1_id']), 0])
        pixels['chr2'] = list(bins.iloc[list(pixels['bin2_id']), 0])
        pixels = pixels[pixels['chr1'] == pixels['chr2']]
        if split_chorm:
            # 如果merge，那就返回一个dict
            counts = {}
            for chrom in clr.chromnames:
                counts[chrom] = pixels[pixels['chr1'] == chrom].shape[0]
        # 不分染色体，那就只返回一个数值
        else:
            counts = {'total': 0}
            for chrom in clr.chromnames:
                counts['total'] += pixels[pixels['chr1'] == chrom].shape[0]

        return counts

    # 获得contact agg,只算上三角+对角线，求和
    def get_intra_contact_agg(self, split_chorm: bool = True, res=None):
        '''
            merge + split: 返回染色体dict
            merge + nosolt:返回一个值
            nomerge + nosplit:返回list of count
            nomerge + split:返回 list of dict
        '''
        counts = []

        for val in self.clr_obj[-1]:
            counts.append(self.__get_intra_contact_agg(val, split_chorm))

        df = pd.DataFrame(counts)
        df.index = self.mcool_names

        df.to_csv('./file/{}_splitChrom:{}_intra_contact_agg.csv'.format(self.describe, split_chorm), sep='\t')
        return counts

    # 一个cooelr，一个result
    def __get_intra_contact_agg(self, clr, split_chorm: bool = True):
        # 分染色体计数
        bins = clr.bins()[:]
        pixels = clr.pixels()[:]
        pixels['chr1'] = list(bins.iloc[list(pixels['bin1_id']), 0])
        pixels['chr2'] = list(bins.iloc[list(pixels['bin2_id']), 0])
        pixels['start'] = list(bins.iloc[list(pixels['bin1_id']), 1])
        pixels['end'] = list(bins.iloc[list(pixels['bin1_id']), 2])
        pixels = pixels[pixels['chr1'] == pixels['chr2']]
        if split_chorm:
            # 如果merge，那就返回一个dict

            counts = {}
            for chrom in self.chromosome:
                try:
                    tmp = pixels[(pixels['chr1'] == chrom)]
                except:
                    continue

                counts[chrom] = sum(tmp['count'])
        # 不分染色体，那就只返回一个数值
        else:

            counts = {'total': sum(pixels[pixels['chr1'] == pixels['chr2']]['count'])}

        return counts

    # 获得inter-chrom contact个数
    def get_inter_contact_count(self):
        counts = []

        for val in self.clr_obj[-1]:
            counts.append(self.__get_inter_contact_count(val))
        df = pd.DataFrame(counts)
        df.index = self.mcool_names

        df.to_csv('./file/{}_inter_contact_count.csv'.format(self.describe), sep='\t')
        return counts

    def __get_inter_contact_count(self, clr):
        # 分染色体计数
        bins = clr.bins()[:]
        pixels = clr.pixels()[:]
        pixels['chr1'] = list(bins.iloc[list(pixels['bin1_id']), 0])
        pixels['chr2'] = list(bins.iloc[list(pixels['bin2_id']), 0])
        pixels['start'] = list(bins.iloc[list(pixels['bin1_id']), 1])
        pixels['end'] = list(bins.iloc[list(pixels['bin1_id']), 2])
        pixels = pixels[pixels['chr1'] != pixels['chr2']]
        counts = {'total': pixels.shape[0]}

        return counts

    # 获得inter-chrom contact总数
    def get_inter_contact_agg(self, split_chrom: bool = True, res=None):
        counts = []

        for val in self.clr_obj[-1]:
            counts.append(self.__get_inter_contact_agg(val))
        df = pd.DataFrame(counts)
        df.index = self.mcool_names

        df.to_csv('./file/{}_inter_contact_agg.csv'.format(self.describe), sep='\t')
        return counts

    def __get_inter_contact_agg(self, clr, split_chorm: bool = True):
        # 分染色体计数
        bins = clr.bins()[:]
        pixels = clr.pixels()[:]
        pixels['chr1'] = list(bins.iloc[list(pixels['bin1_id']), 0])
        pixels['chr2'] = list(bins.iloc[list(pixels['bin2_id']), 0])
        pixels['start'] = list(bins.iloc[list(pixels['bin1_id']), 1])
        pixels['end'] = list(bins.iloc[list(pixels['bin1_id']), 2])
        pixels = pixels[pixels['chr1'] != pixels['chr2']]
        counts = {'total': sum(pixels['count'])}

        return counts

    # 获得clr_obj对象中的embeeding,仅在merge=False时候使用
    def get_clustering_embedding(self, res=None, cache=True,
                                 similarityMethod='gaussian_distance',
                                 embeddingMethod='UMAP',
                                 dim=2):
        if self.__merge:
            print('merged,no embedding')
            return

        if res is None:
            print('please input resolution')
            return

        CellSimilarity = []
        CellFeature = []

        clr = self.clr_obj[self.resolution.index(res)]

        fileCount = len(clr)

        if not cache or not os.path.exists('./cache/cache_%s_%s.npz' % (self.describe, res)):
            # 不存在cache，那还是要做一遍，但是存下来

            for cr in clr:
                ####获取该细胞的HIC矩阵
                TargetMatrix_1 = cr.matrix(balance=False)[:]

                ###获取距离对角线0-10的数据并降成一维
                diagonal_data = [np.diag(TargetMatrix_1, k) for k in range(0, 10)]

                ###获取降维后的矩阵
                SampledTargetMatrix_1 = scale_matrix(TargetMatrix_1, 200)

                ###整合整体特征与染色体内部特征
                CellFeature.append(np.concatenate(
                    [np.concatenate([arr.flatten() for arr in diagonal_data]), SampledTargetMatrix_1.flatten()],
                    axis=0))

                ####计算细胞间的similarity
                CurCellSimilarity = []
                for j in range(fileCount):
                    TargetMatrix_2 = clr[j].matrix(balance=False)[:]
                    SampledTargetMatrix_2 = scale_matrix(TargetMatrix_2, 200)
                    if similarityMethod == 'gaussian_distance':
                        CurCellSimilarity.append(gaussian_distance(SampledTargetMatrix_1, SampledTargetMatrix_2, 1))
                    if similarityMethod == 'EuclideanDistance':
                        CurCellSimilarity.append(np.linalg.norm(SampledTargetMatrix_1 - SampledTargetMatrix_2))
                    if similarityMethod == 'CosineSimilarity':
                        CurCellSimilarity.append(1 - cosine_similarity(SampledTargetMatrix_1, SampledTargetMatrix_2))
                ###记录每个细胞与所有细胞间的similarity
                CellSimilarity.append(CurCellSimilarity)
        else:
            ###存在cache，那就直接读取
            print('load cache')
            cache = np.load('./cache/cache_%s_%s.npz' % (self.describe, res))
            CellFeature = cache['CellFeature']
            CellSimilarity = cache['CellSimilarity']

        # print(CellFeature)
        ##对矩阵正态分布处理
        CellSimilarity = np.array(CellSimilarity)
        smoothed_CellSimilarity = 1 - norm.cdf(CellSimilarity, np.mean(CellSimilarity), np.std(CellSimilarity))
        smoothed_CellSimilarity = (smoothed_CellSimilarity - np.min(smoothed_CellSimilarity)) / (
                    np.max(smoothed_CellSimilarity) - np.min(smoothed_CellSimilarity))

        ###计算每个细胞的embadding
        CellFeature = np.array(CellFeature)
        self.CellFeature = CellFeature
        self.cellSimilarity = smoothed_CellSimilarity

        if cache and not os.path.exists('./cache/cache_%s_%s.npz' % (self.describe, res)):
            if not os.path.exists('cache'):
                os.makedirs('cache')
            print('save cache')
            cache = {}
            cache['CellFeature'] = CellFeature
            cache['CellSimilarity'] = smoothed_CellSimilarity
            cache['mcool_names'] = self.mcool_names

            np.savez('./cache/cache_%s_%s.npz' % (self.describe, res), **cache)

        print(CellFeature)
        if embeddingMethod == 'PCA':
            CellsEmbedding = PCAEmbedding(CellFeature, dim=dim)

        if embeddingMethod == 'MDS':
            CellsEmbedding = MDSEmbedding(CellFeature, dim=dim)

        if embeddingMethod == 'tSNE':
            CellsEmbedding = tSNEEmbedding(CellFeature, dim=dim)

        if embeddingMethod == 'UMAP':
            CellsEmbedding = UMAPEmbedding(CellFeature, dim=dim)

        if embeddingMethod == 'LLE':
            CellsEmbedding = LLEEmbedding(CellFeature, dim=dim)

        if embeddingMethod == 'Isomap':
            CellsEmbedding = IsomapEmbedding(CellFeature, dim=dim)
        else:
            print('please input right embedding method')
            return smoothed_CellSimilarity

        return smoothed_CellSimilarity, CellsEmbedding

    # 获得每个cell的对角线元素flateen的feature
    def get_clustering_diagnal(self, dim=10, res=None):

        shape = np.array(self.clr_obj).shape

        if len(shape) == 1:
            CellFeature = []
            for clr in self.clr_obj:
                ####获取该细胞的HIC矩阵
                TargetMatrix_1 = clr.matrix(balance=False)[:]

                ###获取距离对角线0-10的数据并降成一维
                diagonal_data = [np.diag(TargetMatrix_1, k) for k in range(0, dim)]

                ###获取降维后的矩阵
                SampledTargetMatrix_1 = scale_matrix(TargetMatrix_1, 200)

                ###整合整体特征与染色体内部特征
                result = np.concatenate(
                    [np.concatenate([arr.flatten() for arr in diagonal_data]), SampledTargetMatrix_1.flatten()], axis=0)
                CellFeature.append(result)
        else:
            CellFeature = []
            for lt in self.clr_obj:
                tmp_feature = []
                for clr in lt:
                    ####获取该细胞的HIC矩阵
                    TargetMatrix_1 = clr.matrix(balance=True)[:]

                    ###获取距离对角线0-10的数据并降成一维
                    diagonal_data = [np.diag(TargetMatrix_1, k) for k in range(0, dim)]

                    ###获取降维后的矩阵
                    SampledTargetMatrix_1 = scale_matrix(TargetMatrix_1, 200)

                    ###整合整体特征与染色体内部特征
                    result = np.concatenate(
                        [np.concatenate([arr.flatten() for arr in diagonal_data]), SampledTargetMatrix_1.flatten()],
                        axis=0)
                    tmp_feature.append(result)
                CellFeature.append(tmp_feature)
        return CellFeature

    # 获得每个mcool文件的contact per bin
    def get_avg_contact_per_bin(self, ignore_diag=False, res=None):
        if res is None:
            clr = self.clr_obj
        else:
            clr = self.clr_obj[self.resolution.index(res)]

        shape = np.array(clr).shape

        if len(shape) == 1:
            avgs = []
            for cr in clr:
                avgs.append(self.__get_avg_contact_per_bin(cr, ignore_diag))
        else:
            avgs = []
            for val in clr:
                tmp_avg = []
                for cr in val:
                    tmp_avg.append(self.__get_avg_contact_per_bin(cr, ignore_diag))
                avgs.append(tmp_avg)

        return avgs

    def __get_avg_contact_per_bin(self, clr, ignore_diag=False):
        top = 0
        down = 0
        for chrom in self.chromosome:

            try:
                mat = clr.matrix(balance=False).fetch(chrom)
            except:
                continue
            if ignore_diag:
                top += np.triu(mat, 1).sum()
            else:
                top += np.triu(mat).sum()

            down += mat.shape[0]

        return top / down

    # 获得不同分辨的覆盖率
    def get_coverage(self, res=None):

        if res is None:
            clr = self.clr_obj
        else:
            clr = [self.clr_obj[self.resolution.index(res)]]

        coverages = []
        for val in clr:
            tmp_coverage = []
            for cr in val:
                tmp_coverage.append(self.__get_coverage(cr))
            coverages.append(tmp_coverage)

        if res is None:
            for res in self.resolution:
                df = pd.DataFrame(coverages[self.resolution.index(res)])
                df.index = self.mcool_names
                df.to_csv('./file/{}_{}_coverages.csv'.format(self.describe, res), sep='\t')

        else:
            df = pd.DataFrame(coverages[0])
            df.index = self.mcool_names
            df.to_csv('./file/{}_{}_coverages.csv'.format(self.describe, res), sep='\t')

        return coverages

    def __get_coverage(self, clr):

        bins = clr.bins()[:]
        pixels = clr.pixels()[:]
        cover = set(set(np.unique(pixels['bin1_id'])).union(set(np.unique(pixels['bin2_id']))))
        freq = len(cover) / bins.shape[0]
        return freq

    # 获得累计覆盖率
    def get_acum_coverage(self, res=None):
        if self.__merge:
            sys.exit('merge is True, can not get acum_coverage')

        if res is None:
            clr = self.clr_obj
        else:
            clr = [self.clr_obj[self.resolution.index(res)]]

        coverages = []
        for val in clr:
            tmp_coverage = []
            for cr in val:
                tmp_coverage.append(self.__get_acum_coverage(cr))
            coverages.append(tmp_coverage)

        return coverages

    def __get_acum_coverage(self, clr):
        freq = []
        bins = clr[0].bins()[:]
        pixels = clr[0].pixels()[:]
        cover = set(set(np.unique(pixels['bin1_id'])).union(set(np.unique(pixels['bin2_id']))))
        freq.append(len(cover) / bins.shape[0])
        for cr in clr[1:]:
            bins = cr.bins()[:]
            pixels = cr.pixels()[:]
            cover = cover.union(set(np.unique(pixels['bin1_id'])).union(set(np.unique(pixels['bin2_id']))))
            freq.append(len(cover) / bins.shape[0])
        freq = sorted(freq)
        return freq

    # 计算染色质疆域的nDS
    def get_nDS_chrom_terr(self, res=None, epoch=100):
        if self.__merge:
            sys.exit('merge is True, can not get nDS_chrom_terr_per_cell')

        if res is None:
            clr = self.clr_obj
        else:
            clr = [self.clr_obj[self.resolution.index(res)]]

        self.chrom_pair = []
        for cr in self.__chrom_terr:
            for cr2 in self.__chrom_terr[self.__chrom_terr.index(cr) + 1:]:
                self.chrom_pair.append([cr, cr2])

        result = []
        for i, val in enumerate(clr):
            # print(i)
            result.append(self.__get_nDS_chrom_terr(val, i, epoch))

        return result

    # 获得单个细胞的每个pair的nDS和result
    def __get_nDS_chrom_terr(self, clr, idx, epoch):
        print('cal DS')
        DS = []
        for cell in clr:
            bins = cell.bins()[:]
            pair_DS = []
            pixels = cell.pixels()[:]
            # print(bins)
            pixels['chr1'] = list(bins.iloc[list(pixels['bin1_id']), 0])
            pixels['chr2'] = list(bins.iloc[list(pixels['bin2_id']), 0])
            for pair in self.chrom_pair:
                chr1 = list(bins[bins['chrom'] == pair[0]].index)
                chr2 = list(bins[bins['chrom'] == pair[1]].index)
                pair_DS.append(self.__terri_detection_score(pixels, chr1, chr2))
            DS.append(pair_DS)
        df = pd.DataFrame(DS)
        df.columns = [x[0] + '_' + x[1] for x in self.chrom_pair]

        df.to_csv('./file/%s_DS_%s.csv' % (self.describe, str(self.resolution[idx])), sep='\t')
        print('cal random DS')
        # random ds
        random_DS = []
        for cell in clr:

            pair_DS = [0] * 171
            for i in tqdm.tqdm(range(epoch)):
                pixels = cell.pixels()[:]
                for pair in self.chrom_pair:
                    chr1 = list(bins[bins['chrom'] == pair[0]].index)
                    chr2 = list(bins[bins['chrom'] == pair[1]].index)

                    tmp = pixels.sample(len(chr1) + len(chr2), replace=True)
                    pair_DS[self.chrom_pair.index(pair)] += self.__terri_detection_score(tmp, chr1, chr2)
            random_DS.append([x / epoch for x in pair_DS])
        df2 = pd.DataFrame(random_DS)
        df2.columns = [x[0] + '_' + x[1] for x in self.chrom_pair]

        df2.to_csv('./file/%s_terr_rDS_%s.csv' % (self.describe, str(self.resolution[idx])), sep='\t', index=None)

        nDS = df - df2

        nDS.to_csv('./file/%s_terr_nDS_%s.csv' % (self.describe, str(self.resolution[idx])), sep='\t', index=None)
        return nDS

    def __terri_detection_score(self, pixels: pd.DataFrame, chr1: list, chr2: list):
        AA = pixels[pixels['bin1_id'].isin(chr1) & pixels['bin2_id'].isin(chr1)].shape[0]
        BB = pixels[pixels['bin1_id'].isin(chr2) & pixels['bin2_id'].isin(chr2)].shape[0]
        AB = pixels[pixels['bin1_id'].isin(chr1) & pixels['bin2_id'].isin(chr2)].shape[0]
        BA = pixels[pixels['bin1_id'].isin(chr2) & pixels['bin2_id'].isin(chr1)].shape[0]
        Eintra = len(chr1) * len(chr1) + len(chr2) * len(chr2)
        Einter = len(chr1) * len(chr2) * 2
        # print(AA,BB,AB,BA,Eintra,Einter)
        DS = (AA + BB) / Eintra - (AB + BA) / Einter
        return DS

    # 计算每个bin的A/B compartment
    def get_AB_compartment(self, genome='', res=None, epoch=100):
        if self.__merge:
            sys.exit('merge is True, can not get AB_compartmen')
        if res is None:
            sys.exit('res is None, can not get AB_compartmen')
        if res not in self.resolution:
            sys.exit('res is not in resolution, can not get AB_compartmen')

        clr = self.clr_obj[self.resolution.index(res)]

        print('merging cooler')
        merge = []
        for val in self.cache_file:
            name = val + '::resolutions/' + str(res)
            merge.append(name)
        merge_name = './cache/merge_' + str(res) + '_' + self.describe + '.mcool'
        if not os.path.exists(merge_name):
            cooler.merge_coolers(input_uris=merge, output_uri=merge_name, mergebuf=1000000)
        tmp_merge = cooler.Cooler(merge_name)
        cooler.balance_cooler(tmp_merge, store=True)
        print('done')
        genome = bioframe.load_fasta(genome)
        bins = tmp_merge.bins()[:]
        gc_cov = bioframe.frac_gc(bins[['chrom', 'start', 'end']], genome)
        view_df = pd.DataFrame({'chrom': tmp_merge.chromnames,
                                'start': 0,
                                'end': tmp_merge.chromsizes.values,
                                'name': tmp_merge.chromnames}
                               )
        cis_eigs = cooltools.eigs_cis(
            tmp_merge,
            gc_cov,
            view_df=view_df,
            n_eigs=3,
        )

        # cis_eigs[0] returns eigenvalues, here we focus on eigenvectors
        eigenvector_track = cis_eigs[1][['chrom', 'start', 'end', 'E1']]
        store_name = 'E1'
        with tmp_merge.open("r+") as grp:
            if store_name in grp["bins"]:
                del grp["bins"][store_name]
            h5opts = {"compression": "gzip", "compression_opts": 6}
            grp["bins"].create_dataset(store_name, data=eigenvector_track['E1'], **h5opts)

        data = tmp_merge.bins()[:]

        data['E1'] = np.nan_to_num(data['E1'])
        data['sign'] = '$'
        data['sign'][data['E1'] > 0] = 'A'
        data['sign'][data['E1'] < 0] = 'B'
        a_list = []
        b_list = []
        for c in set(data['chrom']):
            seq = list(data[data['chrom'] == c]['sign'])
            base = data[data['chrom'] == c].index[0]

            seq = ''.join(seq)
            #     print(seq)
            new_seq = ''
            idx = []
            tmp_idx = []
            for i in range(len(seq)):
                sym = seq[i]
                if new_seq == '':
                    new_seq += sym
                    tmp_idx.append(i + base)
                else:
                    if sym == seq[i - 1]:
                        tmp_idx.append(i + base)
                    else:
                        idx.append(tmp_idx)
                        tmp_idx = [i + base]
                        new_seq += sym
            for i in range(len(new_seq) - 3):
                if new_seq[i:i + 3] == 'ABA':
                    tmp_A = idx[i] + idx[i + 2]
                    tmp_B = idx[i + 1]
                    a_list.append(tmp_A)
                    b_list.append(tmp_B)

                elif new_seq[i:i + 3] == 'BAB':
                    tmp_A = idx[i + 1]
                    tmp_B = idx[i] + idx[i + 2]
                    a_list.append(tmp_A)
                    b_list.append(tmp_B)
        result = []
        for i in tqdm.tqdm(range(len(clr))):
            cr = clr[i]
            bins = cr.bins()[:]
            pixels = cr.pixels()[:]
            pixels['chr1'] = list(bins.iloc[list(pixels['bin1_id']), 0])
            pixels['chr2'] = list(bins.iloc[list(pixels['bin2_id']), 0])
            pixels['start'] = list(bins.iloc[list(pixels['bin1_id']), 1])
            pixels['end'] = list(bins.iloc[list(pixels['bin1_id']), 2])
            pixels = pixels[pixels['chr1'] == pixels['chr2']]
            DS = []
            for i, j in zip(a_list, b_list):
                DS.append(self.__terri_detection_score(pixels, i, j))
            rDS = self.__return_AB_rds(a_list, b_list, pixels, bins, epoch)
            result.append((DS - rDS).mean())

        return result

    def __return_AB_rds(self, a_list, b_list, pixels, bins, epoch):

        pixels['chr1'] = list(bins.iloc[list(pixels['bin1_id']), 0])
        pixels['chr2'] = list(bins.iloc[list(pixels['bin2_id']), 0])
        pixels['start'] = list(bins.iloc[list(pixels['bin1_id']), 1])
        pixels['end'] = list(bins.iloc[list(pixels['bin1_id']), 2])
        pixels = pixels[pixels['chr1'] == pixels['chr2']]
        bins['index'] = list(bins.index)
        bins = bins.groupby('chrom').describe()['index'].loc[:, ['min', 'max']].astype(int)
        res = []
        for _ in range(epoch):
            for i, j in zip(a_list, b_list):
                idx = np.random.randint(bins['min'], bins['max'])
                idx = random.sample(list(idx), 1)[0]
                tmp = list(sorted(i + j))

                if tmp.index(j[0]) != 0:
                    ch1 = list(range(idx, idx + tmp.index(j[0]))) + list(
                        range(idx + tmp.index(j[-1]) + 1, idx + len(tmp)))
                    ch2 = list(range(idx + tmp.index(j[0]), idx + tmp.index(j[-1]) + 1))
                else:
                    ch1 = list(range(idx, idx + tmp.index(j[-1]) + 1))
                    ch2 = list(range(idx + tmp.index(j[-1]) + 1, idx + len(tmp)))
                try:
                    rds = self.__terri_detection_score(pixels, ch1, ch2)
                except:
                    continue
                res.append(rds)
        return np.array(res).mean()

    def get_TAD_nDS(self, res=500000, epoch=500, nbins=10, ):
        chrom = ['chr' + str(i) for i in range(20)]
        if res not in self.resolution:
            sys.exit('resolution not in %s' % str(self.resolution))

        merge = []
        for val in self.cache_file:
            name = val + '::resolutions/' + str(res)
            merge.append(name)
        merge_name = './cache/merge_' + str(res) + '_' + self.describe + '.mcool'
        if not os.path.exists(merge_name):
            cooler.merge_coolers(input_uris=merge, output_uri=merge_name, mergebuf=1000000)
        tmp_merge = cooler.Cooler(merge_name)
        cooler.balance_cooler(tmp_merge, store=True)
        resolution = res
        win = [10 * resolution]
        insulation_table = insulation(tmp_merge, win, verbose=False)
        boundary = insulation_table[insulation_table['is_boundary_' + str(win[0])] == True]
        
        boundary = boundary[boundary['chrom'].isin(chrom)]

        tad_index = list(boundary.index)

        result = []
        clr = self.clr_obj[self.resolution.index(res)]
        for i in tqdm.tqdm(range(len(clr)), desc='cell'):
            DS = self.__get_TAD_DS(clr[i], tad_index, epoch, nbins)
            result.append(DS)

        return result

    def __get_TAD_DS(self, clr, data, epoch, nbins):
        chrom = ['chr' + str(i) for i in range(20)]
        bins = clr.bins()[:]
        
        bins['index'] = list(bins.index)
        

        thre = bins.groupby('chrom').describe()['index'].loc[:, ['min', 'max']]#.astype(int)
        thre = thre[thre.index.isin(chrom)]
        thre['min'] = thre['min'].astype(int)
        thre['max'] = thre['max'].astype(int)
   
        pixels = clr.pixels()[:]
        pixels['chr1'] = list(bins.iloc[list(pixels['bin1_id']), 0])
        pixels['chr2'] = list(bins.iloc[list(pixels['bin2_id']), 0])
        pixels = pixels[pixels['chr1'] == pixels['chr2']]
        pixels = pixels[pixels['chr1'].isin(chrom)]
        pixels = pixels[pixels['chr2'].isin(chrom)]

        result = pd.DataFrame(columns=['tad_boundary', 'tad_left', 'tad_right', 'chrom'])
        result['tad_boundary'] = list(data)
        result['chrom'] = list(bins.loc[data, :]['chrom'])
        result = result.merge(thre, on='chrom', how='left')
        result['tad_left'] = [max(a, b - nbins) for a, b in zip(list(result['min']), data)]
        result['tad_right'] = [min(a, b + nbins) for a, b in zip(list(result['max']), data)]
        intra_pixels = pixels[pixels['bin1_id'] != pixels['bin2_id']]
        inter_pixels = pixels[pixels['bin1_id'] != pixels['bin2_id']]

        result['DS'] = result.apply(lambda x: self.__get_tad_ds(x, intra_pixels, inter_pixels, nbins), axis=1)

        thre['r_min'] = list(thre['min'] + nbins)
        thre['r_max'] = list(thre['max'] - nbins)

        rds = self.__get_tad_rds(thre, intra_pixels, inter_pixels, nbins, epoch)

        result = np.mean(result['DS']) - rds

        return result

    def __get_tad_rds(self, df, intra_pixels, inter_pixels, nbins, epoch):
        
        res = []
        for i in range(epoch):
            # print(nbins)
            idx = np.random.randint(df['r_min'], df['r_max'])
            idx = random.sample(list(idx), 1)[0]
            idx1 = range(idx - nbins, idx)
            idx2 = range(idx + 1, idx + nbins + 1)
            intra = intra_pixels[((intra_pixels['bin1_id'].isin(idx1)) & (intra_pixels['bin2_id'].isin(idx1))) | (
                        (intra_pixels['bin1_id'].isin(idx2)) & (intra_pixels['bin2_id'].isin(idx2)))]['count'].sum()
            inter = inter_pixels[(inter_pixels['bin1_id'].isin(idx1)) & (inter_pixels['bin2_id'].isin(idx2)) | (
                inter_pixels['bin1_id'].isin(idx2)) & (inter_pixels['bin2_id'].isin(idx1))]['count'].sum()
            # print(idx1,idx2)
            res.append((intra - inter) / (nbins * nbins * 2))

        return np.mean(res)

    def __get_tad_ds(self, x, intra_pixels, inter_pixels, nbins):
        x = list(x)

        idx1 = range(x[1], x[0])
        idx2 = range(x[0] + 1, x[2] + 1)

        intra = intra_pixels[((intra_pixels['bin1_id'].isin(idx1)) & (intra_pixels['bin2_id'].isin(idx1))) | (
                    (intra_pixels['bin1_id'].isin(idx2)) & (intra_pixels['bin2_id'].isin(idx2)))]['count'].sum()
        inter = inter_pixels[(inter_pixels['bin1_id'].isin(idx1)) & (inter_pixels['bin2_id'].isin(idx2)) | (
            inter_pixels['bin1_id'].isin(idx2)) & (inter_pixels['bin2_id'].isin(idx1))]['count'].sum()
        # print(intra, inter, 1)
        return (intra - inter) / (nbins * nbins * 2)

    def get_Centromere_nDs(self, res=None, epoch=100):
        if res is None:
            sys.exit('resoluiton is required')

        clr = self.clr_obj[self.resolution.index(res)]

        result = []
        for i, cr in enumerate(clr):
            log = self.__get_Centromere_nDs(cr, res, epoch=epoch)
            result.append([log[0], log[1], log[2]])
        result = pd.DataFrame(result)
        result.columns = ['DS', 'rDS', 'nDS']
        result.to_csv('./file/%s_centromere_DS_%s_%s.csv' % (self.describe, str(res), self.mcool_names[i]), sep='\t')

        return list(result['nDS'])

    def __get_Centromere_nDs(self, clr, res, epoch):
        centro_pos1 = 3000000  # 3mb
        centro_pos2 = 13000000  # 13mb

        # 获得binid的最大值和最小值，每个染色体不一定有这么长的bp区间
        bin_idx1 = int(centro_pos1 / res)
        bins_idx2 = int(centro_pos2 / res)

        bins = clr.bins()[:]
        pixels = clr.pixels()[:]

        chrom, n_bins = list(bins.groupby('chrom').count()['start'].index), list(bins.groupby('chrom').count()['start'])
        interval_to_cal = {}
        for c, n in zip(chrom, n_bins):
            if bin_idx1 > int(n):
                continue
            interval_to_cal[c] = [bin_idx1, min(n, bins_idx2)]

        ac = 0
        for i, k in enumerate(interval_to_cal.keys()):
            if i == 0:
                ac += n_bins[i]
            else:
                id1, id2 = interval_to_cal[k]
                id1 += ac
                id2 += ac
                ac += n_bins[i]
                interval_to_cal[k] = [id1, id2]
        interval_to_cal_ = []
        for k in interval_to_cal.keys():
            interval_to_cal_ += list(range(interval_to_cal[k][0], interval_to_cal[k][1] + 1))

        chr1 = interval_to_cal_
        chr2 = list(set(list(bins.index)).difference(set(chr1)))

        DS = self.__terri_detection_score(pixels, chr1, chr2)

        r_DS = []
        for i in range(epoch):
            random_interval = list(bins.sample(len(chr1)).index, replace=True)
            random_interval2 = list(set(list(bins.index)).difference(set(random_interval)))
            r_DS.append(self.__terri_detection_score(pixels, random_interval, random_interval2))
        rDS = np.mean(r_DS)
        print(DS - rDS)

        return DS, rDS, DS - rDS

    def get_insulation_score(self, res=None, win=None):
        if res is None:
            sys.exit('resoluiton is required')

        clr = self.clr_obj[self.resolution.index(res)]
        result = []
        for i, cr in enumerate(clr):
            score = self.__get_insulation_score(cr, win)

        return result

    def __get_insulation_score(self, clr, win):
        win = [win]
        if not cooltools.lib.is_cooler_balanced(clr):
            cooler.balance_cooler(clr, store=True)
        # 获得 insulation score，以及是不是边界
        insulation_table = insulation(clr, win, verbose=True)
        s = insulation_table['log2_insulation_score_{}'.format(win[0])]
        s = np.nan_to_num(s)
        return s

    def pooled_cell(self, pool_size=None, num=None, sort=True):
        if self.__merge:
            sys.exit('merged can not pool the cell')
        if self.__pool:
            print('already pooled')
            return

        if pool_size is None or num is None:
            sys.exit('pool_size and num is required')

        # 一共多少个细胞
        element = len(self.clr_obj[0])
        #
        # if sort and num > element:
        #     print('sorted mode: num is bigger than element, set num = element')
        #     num = element
        new_clr_obj = []
        # 循环添加不同分辨率的空list
        for i in range(len(self.resolution)):
            new_clr_obj.append([])

        if not sort:
            for i in tqdm.tqdm(range(num), desc='pooling'):
                pool = random.choices(list(range(element)), k=pool_size)
                mcool_pool = [self.cache_file[idx] for idx in pool]
                for j in range(len(self.resolution)):
                    clr_list = [val + '::/resolutions/' + str(self.resolution[j]) for val in mcool_pool]

                    new_clr_obj[j].append(self.__merge_pool(clr_list, str(i)))
        else:
            # 从大到小排序total contact
            idx = [i[0] for i in sorted(enumerate(self.get_total_contact()), key=lambda x: x[1], reverse=True)]
            # 如果num大于element，那么就随机插入
            if num > element:
                # 假设这是原始列表
                original_list = self.get_total_contact()
                selected_elements = original_list.copy()
                selected_elements.extend(random.sample(original_list, 5))
                sorted_elements = sorted(selected_elements, reverse=True)
                idx = [original_list.index(x) for x in sorted_elements]

            for i in tqdm.tqdm(range(num), desc='pooling::sorted'):
                pool = idx[:i + 1]

                mcool_pool = [self.cache_file[idx] for idx in pool]
                for j in range(len(self.resolution)):
                    clr_list = [val + '::/resolutions/' + str(self.resolution[j]) for val in mcool_pool]
                    new_clr_obj[j].append(self.__merge_pool(clr_list, str(i)))

        self.clr_obj = new_clr_obj
        self.mcool_names = ['pooled_cell_%s' % str(i) for i in range(num)]

    def __merge_pool(self, clr_list, cell_ids):

        cooler.merge_coolers(input_uris=clr_list,
                             output_uri='./cache/%s_pool_all_res-%s.mcool' % (self.describe, cell_ids),
                             mergebuf=1000000)

        res = cooler.Cooler('./cache/%s_pool_all_res-%s.mcool' % (self.describe, cell_ids))

        return res

    def filter_by_total_contact(self, max_=np.inf, min_=0):
        if max_ is None and min_ is None:
            sys.exit('max and min is required')

        tmp_mcool_files = []
        tmp_mcool_names = []
        tmp_cache_file = []
        tmp_clr_obj = []
        for i in range(len(self.resolution)):
            tmp_clr_obj.append([])
        total_contact_list = self.get_total_contact()
        for a, b, c, in zip(self.mcool_files, self.mcool_names, self.cache_file):
            # print(c)
            total_contact = total_contact_list[self.mcool_files.index(a)]
            if total_contact > max_ or total_contact < min_:
                continue
            tmp_mcool_files.append(a)
            tmp_mcool_names.append(b)
            tmp_cache_file.append(c)
            for i in range(len(self.resolution)):
                tmp_clr_obj[i].append(self.clr_obj[i][self.mcool_files.index(a)])

        self.mcool_files = tmp_mcool_files
        self.mcool_names = tmp_mcool_names
        self.cache_file = tmp_cache_file
        self.clr_obj = tmp_clr_obj

    def get_cis_trans_ratio(self, distance=20000):
        if self.resolution[0] > distance:
            sys.exit('distance must be smaller than the samllest resolution')

        result = []
        gap = distance / self.resolution[0]

        for clr in self.clr_obj[0]:
            cis, trans = self.__get_cis_trans_ratio(clr, gap)
            result.append([cis, trans])

        result = pd.DataFrame(result)
        result.columns = ['cis', 'trans']
        result.index = self.mcool_names
        return result

    def __get_cis_trans_ratio(self, clr, distance):
        bins = clr.bins()[:]
        pixels = clr.pixels()[:]
        pixels['chr1'] = list(bins.iloc[list(pixels['bin1_id']), 0])
        pixels['chr2'] = list(bins.iloc[list(pixels['bin2_id']), 0])
        pixels = pixels[pixels['chr1'] == pixels['chr2']]

        cis = sum(pixels[pixels['bin2_id'] - pixels['bin1_id'] <= distance]['count'])
        trans = sum(pixels[pixels['bin2_id'] - pixels['bin1_id'] > distance]['count'])
        return cis, trans

    def sorted_cell(self, reverse=True, num=None):
        idx = [i[0] for i in sorted(enumerate(self.get_total_contact()), key=lambda x: x[1], reverse=reverse)]

        self.cache_file = [self.cache_file[i] for i in idx]
        self.mcool_files = [self.mcool_files[i] for i in idx]
        self.mcool_names = [self.mcool_names[i] for i in idx]
        for i in range(len(self.resolution)):
            self.clr_obj[i] = [self.clr_obj[i][j] for j in idx]

        if num is not None and num < len(self.cache_file):
            self.cache_file = self.cache_file[:num]
            self.mcool_files = self.mcool_files[:num]
            self.mcool_names = self.mcool_names[:num]
            for i in range(len(self.resolution)):
                self.clr_obj[i] = self.clr_obj[i][:num]



   

































