#!/usr/bin/env python
#
# Copyright (c) 2017 10X Genomics, Inc. All rights reserved.
#

import cellranger.analysis.diffexp as cr_diffexp
import cellranger.analysis.io as analysis_io
from cellranger.analysis.singlegenome import SingleGenomeAnalysis
import cellranger.h5_constants as h5_constants
import cellranger.analysis.constants as analysis_constants
import cellranger.matrix as cr_matrix
import cellranger.io as cr_io
import cellranger.library_constants as lib_constants

NUM_THREADS_MIN = 4
#TODO Not clear why this stage takes > 1 thread. Martian thinks it does and kills it on long jobs

__MRO__ = """
stage RUN_DIFFERENTIAL_EXPRESSION(
    in  h5     matrix_h5,
    in  h5     clustering_h5,
    in  bool   skip,
    in  int    random_seed,
    in  int    max_clusters,
    out h5     diffexp_h5,
    out path   diffexp_csv,
    src py     "stages/analyzer/run_differential_expression",
) split using (
    in  string clustering_key,
)
"""

def split(args):
    if args.skip:
        return {'chunks': [{'__mem_gb': h5_constants.MIN_MEM_GB}]}

    chunks = []
    # FIXME: Add one for reasons unknown
    matrix_mem_gb = 1.8 * cr_matrix.CountMatrix.get_mem_gb_from_matrix_h5(args.matrix_h5)
    chunk_mem_gb = int(max(matrix_mem_gb, h5_constants.MIN_MEM_GB))

    # HACK - give big jobs more threads in order to avoid overloading a node
    threads = min(cr_io.get_thread_request_from_mem_gb(chunk_mem_gb), NUM_THREADS_MIN)
    threads = 4

    for key in SingleGenomeAnalysis.load_clustering_keys_from_h5(args.clustering_h5):
        chunks.append({
            'clustering_key': key,
            '__mem_gb': chunk_mem_gb,
            '__threads': threads,
        })

    return {'chunks': chunks, 'join': {'__mem_gb' : 1}}

def main(args, outs):
    if args.skip:
        return

    matrix = cr_matrix.CountMatrix.load_h5_file(args.matrix_h5)

    # For now, only compute for gene expression features
    matrix = matrix.select_features_by_type(lib_constants.GENE_EXPRESSION_LIBRARY_TYPE)

    clustering = SingleGenomeAnalysis.load_clustering_from_h5(args.clustering_h5, args.clustering_key)

    diffexp = cr_diffexp.run_differential_expression(matrix, clustering.clusters)

    with analysis_io.open_h5_for_writing(outs.diffexp_h5) as f:
        cr_diffexp.save_differential_expression_h5(f, args.clustering_key, diffexp)

    cr_diffexp.save_differential_expression_csv(args.clustering_key, diffexp, matrix, outs.diffexp_csv)

def join(args, outs, chunk_defs, chunk_outs):
    if args.skip:
        return

    chunk_h5s = [chunk_out.diffexp_h5 for chunk_out in chunk_outs]
    chunk_csv_dirs = [chunk_out.diffexp_csv for chunk_out in chunk_outs]

    analysis_io.combine_h5_files(chunk_h5s, outs.diffexp_h5, [analysis_constants.ANALYSIS_H5_DIFFERENTIAL_EXPRESSION_GROUP,
                                                        analysis_constants.ANALYSIS_H5_KMEANS_DIFFERENTIAL_EXPRESSION_GROUP])

    for csv_dir in chunk_csv_dirs:
        cr_io.copytree(csv_dir, outs.diffexp_csv, allow_existing=True)

[((46, 15, 46, 84), 'cellranger.analysis.singlegenome.SingleGenomeAnalysis.load_clustering_keys_from_h5', 'SingleGenomeAnalysis.load_clustering_keys_from_h5', ({(46, 65, 46, 83): 'args.clustering_h5'}, {}), '(args.clustering_h5)', False, 'from cellranger.analysis.singlegenome import SingleGenomeAnalysis\n'), ((59, 13, 59, 63), 'cellranger.matrix.CountMatrix.load_h5_file', 'cr_matrix.CountMatrix.load_h5_file', ({(59, 48, 59, 62): 'args.matrix_h5'}, {}), '(args.matrix_h5)', True, 'import cellranger.matrix as cr_matrix\n'), ((64, 17, 64, 102), 'cellranger.analysis.singlegenome.SingleGenomeAnalysis.load_clustering_from_h5', 'SingleGenomeAnalysis.load_clustering_from_h5', ({(64, 62, 64, 80): 'args.clustering_h5', (64, 82, 64, 101): 'args.clustering_key'}, {}), '(args.clustering_h5, args.\n    clustering_key)', False, 'from cellranger.analysis.singlegenome import SingleGenomeAnalysis\n'), ((66, 14, 66, 81), 'cellranger.analysis.diffexp.run_differential_expression', 'cr_diffexp.run_differential_expression', ({(66, 53, 66, 59): 'matrix', (66, 61, 66, 80): 'clustering.clusters'}, {}), '(matrix, clustering.clusters)', True, 'import cellranger.analysis.diffexp as cr_diffexp\n'), ((71, 4, 71, 103), 'cellranger.analysis.diffexp.save_differential_expression_csv', 'cr_diffexp.save_differential_expression_csv', ({(71, 48, 71, 67): 'args.clustering_key', (71, 69, 71, 76): 'diffexp', (71, 78, 71, 84): 'matrix', (71, 86, 71, 102): 'outs.diffexp_csv'}, {}), '(args.clustering_key, diffexp,\n    matrix, outs.diffexp_csv)', True, 'import cellranger.analysis.diffexp as cr_diffexp\n'), ((80, 4, 81, 125), 'cellranger.analysis.io.combine_h5_files', 'analysis_io.combine_h5_files', ({(80, 33, 80, 42): 'chunk_h5s', (80, 44, 80, 59): 'outs.diffexp_h5', (80, 61, 81, 124): '[analysis_constants.ANALYSIS_H5_DIFFERENTIAL_EXPRESSION_GROUP,\n    analysis_constants.ANALYSIS_H5_KMEANS_DIFFERENTIAL_EXPRESSION_GROUP]'}, {}), '(chunk_h5s, outs.diffexp_h5, [\n    analysis_constants.ANALYSIS_H5_DIFFERENTIAL_EXPRESSION_GROUP,\n    analysis_constants.ANALYSIS_H5_KMEANS_DIFFERENTIAL_EXPRESSION_GROUP])', True, 'import cellranger.analysis.io as analysis_io\n'), ((39, 26, 39, 89), 'cellranger.matrix.CountMatrix.get_mem_gb_from_matrix_h5', 'cr_matrix.CountMatrix.get_mem_gb_from_matrix_h5', ({(39, 74, 39, 88): 'args.matrix_h5'}, {}), '(args.matrix_h5)', True, 'import cellranger.matrix as cr_matrix\n'), ((43, 18, 43, 68), 'cellranger.io.get_thread_request_from_mem_gb', 'cr_io.get_thread_request_from_mem_gb', ({(43, 55, 43, 67): 'chunk_mem_gb'}, {}), '(chunk_mem_gb)', True, 'import cellranger.io as cr_io\n'), ((68, 9, 68, 57), 'cellranger.analysis.io.open_h5_for_writing', 'analysis_io.open_h5_for_writing', ({(68, 41, 68, 56): 'outs.diffexp_h5'}, {}), '(outs.diffexp_h5)', True, 'import cellranger.analysis.io as analysis_io\n'), ((69, 8, 69, 83), 'cellranger.analysis.diffexp.save_differential_expression_h5', 'cr_diffexp.save_differential_expression_h5', ({(69, 51, 69, 52): 'f', (69, 54, 69, 73): 'args.clustering_key', (69, 75, 69, 82): 'diffexp'}, {}), '(f, args.clustering_key, diffexp)', True, 'import cellranger.analysis.diffexp as cr_diffexp\n'), ((84, 8, 84, 70), 'cellranger.io.copytree', 'cr_io.copytree', (), '', True, 'import cellranger.io as cr_io\n')]