from scale.dataset import read_mtx

from argparse import ArgumentParser
import pandas as pd
import numpy as np

import os


def parse_args():
    parser = ArgumentParser('Preparing raw peaks from cicero pipeline')
    parser.add_argument('--dataset_path', help='Path to Scale dataset: count, feature, barcode folder')
    parser.add_argument('--label_path', help='Path to cell labels')
    parser.add_argument('--num_peaks_threshold', type=int, help='Num peaks to filter')
    parser.add_argument('--output_path', help='Path to save peaks in bed folder')
    parser.add_argument('--suffix', help='Suffix to path')
    return parser.parse_args()


def main():
    args = parse_args()
    labels = pd.read_csv(args.label_path, sep='\t', header=None)
    count, feature, barcode = read_mtx(args.dataset_path)

    os.makedirs(args.output_path, exist_ok=True)
    cell_types = labels[1].unique()
    cell_barcodes = {}
    for cell_type in cell_types:
        cell_barcodes[cell_type] = list(labels[labels[1] == cell_type].index)

    for cell_type, barcode in cell_barcodes.items():
        cell_by_feature = np.asarray(count[barcode].sum(axis=0)).flatten()
        feature_threshold = cell_by_feature[np.argsort(cell_by_feature)[-args.num_peaks_threshold]]
        print(f'{cell_type}: {feature_threshold}')
        filtered_features = (cell_by_feature > 0) & (cell_by_feature >= feature_threshold)
        print(f'{cell_type}: filtered {np.sum(filtered_features)}')
        output = pd.DataFrame(feature[filtered_features])
        # print(cell_type, cell_by_feature[np.argsort(cell_by_feature)[-args.num_peaks_threshold:]][:10])
        output['chr'] = output[0].apply(lambda x: x.split('_')[0])
        output['start'] = output[0].apply(lambda x: x.split('_')[1])
        output['end'] = output[0].apply(lambda x: x.split('_')[2])
        output.drop(0, axis=1).to_csv(
            os.path.join(args.output_path, f'{cell_type.replace(" ", "_").replace("/", "_")}_{args.suffix}.bed'),
            header=None,
            index=None,
            sep='\t'
        )

if __name__ == '__main__':
    main()


[((11, 13, 11, 71), 'argparse.ArgumentParser', 'ArgumentParser', ({(11, 28, 11, 70): '"""Preparing raw peaks from cicero pipeline"""'}, {}), "('Preparing raw peaks from cicero pipeline')", False, 'from argparse import ArgumentParser\n'), ((22, 13, 22, 64), 'pandas.read_csv', 'pd.read_csv', (), '', True, 'import pandas as pd\n'), ((23, 30, 23, 57), 'scale.dataset.read_mtx', 'read_mtx', ({(23, 39, 23, 56): 'args.dataset_path'}, {}), '(args.dataset_path)', False, 'from scale.dataset import read_mtx\n'), ((25, 4, 25, 48), 'os.makedirs', 'os.makedirs', (), '', False, 'import os\n'), ((37, 17, 37, 57), 'pandas.DataFrame', 'pd.DataFrame', ({(37, 30, 37, 56): 'feature[filtered_features]'}, {}), '(feature[filtered_features])', True, 'import pandas as pd\n'), ((33, 44, 33, 71), 'numpy.argsort', 'np.argsort', ({(33, 55, 33, 70): 'cell_by_feature'}, {}), '(cell_by_feature)', True, 'import numpy as np\n'), ((36, 39, 36, 64), 'numpy.sum', 'np.sum', ({(36, 46, 36, 63): 'filtered_features'}, {}), '(filtered_features)', True, 'import numpy as np\n')]