#!/usr/bin/env python3

import argparse
import pandas as pd
import numpy as np
import json
import sys
import os
import time

import matplotlib
matplotlib.use('Agg')
import logging
from pkg_resources import Requirement, resource_filename

from miner import miner, util, GIT_SHA
from miner import __version__ as MINER_VERSION


DESCRIPTION = """miner3-mechinf - MINER compute mechanistic inference
MINER Version %s (Git SHA %s)""" % (str(MINER_VERSION).replace('miner3 ', ''),
                                    GIT_SHA.replace('$Id: ', '').replace(' $', ''))

NUM_CORES = 5
MIN_REGULON_GENES = 5

if __name__ == '__main__':
    logging.getLogger('matplotlib.font_manager').disabled = True
    LOG_FORMAT = '%(asctime)s %(message)s'
    logging.basicConfig(format=LOG_FORMAT, level=logging.DEBUG,
                        datefmt='%Y-%m-%d %H:%M:%S \t')

    parser = argparse.ArgumentParser(formatter_class=argparse.RawDescriptionHelpFormatter,
                                     description=DESCRIPTION)
    parser.add_argument('expfile', help="input matrix")
    parser.add_argument('mapfile', help="identifier mapping file")
    parser.add_argument('coexprdict', help="coexpressionDictionary.json file from miner-coexpr")
    parser.add_argument('outdir', help="output directory")
    parser.add_argument('-mc', '--mincorr', type=float, default=0.2,
                        help="minimum correlation")
    parser.add_argument('--skip_tpm', action="store_true",
                        help="overexpression threshold")
    parser.add_argument('--firmout', default='miner_exported_regulons.sgn',
                        help='file name for FIRM input file, will be stored in outdir')
    parser.add_argument('--genelist', default='all_genes.txt',
                        help='file name for the gene file, will be stored in outdir')
    parser.add_argument('--tfs2genes', default=None,
                        help='Transcription factor to gene mapping file, either JSON or pkl')

    args = parser.parse_args()

    if not os.path.exists(args.expfile):
        sys.exit("expression file not found")
    if not os.path.exists(args.mapfile):
        sys.exit("identifier mapping file not found")
    if not os.path.exists(args.coexprdict):
        sys.exit("revised clusters file not found")

    if not os.path.exists(args.outdir):
        os.makedirs(args.outdir)

    with open(os.path.join(args.outdir, 'run_info.txt'), 'w') as outfile:
        util.write_dependency_infos(outfile)

    exp_data, conv_table = miner.preprocess(args.expfile, args.mapfile, do_preprocess_tpm=(not args.skip_tpm))

    with open(args.coexprdict) as infile:
        revised_clusters = json.load(infile)

    # get first principal component axes of clusters
    t1 = time.time()
    if args.tfs2genes is not None:
        database_path = args.tfs2genes
    else:
        try:
            database_path = resource_filename(Requirement.parse("isb_miner3"),
                                              'miner/data/network_dictionaries/{}'.format("tfbsdb_tf_to_genes.pkl"))
        except:
            # running from source
            database_path = os.path.join('miner', 'data', 'network_dictionaries', 'tfbsdb_tf_to_genes.pkl')

    axes = miner.principalDf(revised_clusters, exp_data,
                             subkey=None, minNumberGenes=1)

    # analyze revised clusters for enrichment in relational database
    # (default: transcription factor binding site database)
    mechanistic_output = miner.mechanisticInference(axes, revised_clusters, exp_data,
                                                    correlationThreshold=args.mincorr,
                                                    numCores=NUM_CORES,
                                                    database_path=database_path)

    # write mechanistic output to .json file
    with open(os.path.join(args.outdir, "mechanisticOutput.json"), 'w') as outfile:
        json.dump(mechanistic_output, outfile)

    # order mechanisticOutput as {tf:{coexpressionModule:genes}}
    coregulation_modules = miner.getCoregulationModules(mechanistic_output)

    # write coregulation modules to .json file
    with open(os.path.join(args.outdir, "coregulationModules.json"), 'w') as outfile:
        json.dump(coregulation_modules, outfile)

    # get final regulons by keeping genes that requently appear coexpressed and associated
    # to a common regulator
    regulons = miner.getRegulons(coregulation_modules,
                                 minNumberGenes=MIN_REGULON_GENES,
                                 freqThreshold=0.333)

    # reformat regulon dictionary for consistency with revisedClusters and coexpressionModules
    regulon_modules, regulon_df = miner.regulonDictionary(regulons)

    # FIRM export: note that we do not check whether we have RefSeq or Entrez, maybe this should be
    # checked in the glue
    entrez_map = miner.make_entrez_map(args.mapfile)
    with open(os.path.join(args.outdir, args.firmout), 'w') as outfile:
        outfile.write('Gene\tGroup\n')
        for regulon, genes in regulon_modules.items():
            for gene in genes:
                if gene in entrez_map:
                    outfile.write('%s\t%s\n' % (entrez_map[gene], regulon))

    # OpenTargets export
    with open(os.path.join(args.outdir, args.genelist), 'w') as outfile:
        all_genes = set()
        for regulon, genes in regulon_modules.items():
            all_genes.update(genes)
        for gene in sorted(all_genes):
            outfile.write('%s\n' % gene)

    # write regulons to json file
    with open(os.path.join(args.outdir, "regulons.json"), 'w') as outfile:
        json.dump(regulon_modules, outfile)
    regulon_df.to_csv(os.path.join(args.outdir, "regulonDf.csv"))

    # define coexpression modules as composite of coexpressed regulons
    coexpression_modules = miner.getCoexpressionModules(mechanistic_output)

    # write coexpression modules to .json file
    with open(os.path.join(args.outdir, "coexpressionModules.json"), 'w') as outfile:
        json.dump(coexpression_modules, outfile)

    # write annotated coexpression clusters to .json file
    with open(os.path.join(args.outdir, "coexpressionDictionary_annotated.json"), 'w') as outfile:
        json.dump(revised_clusters, outfile)

    """
    # write annotated regulons to .json file
    with open(os.path.join(args.outdir, "regulons_annotated.json"), 'w') as outfile:
        json.dump(regulons, outfile)

    # reconvert coexpression modules
    annotated_coexpression_modules = mechanistic_inference.convert_dictionary(coexpression_modules, conv_table)

    # write annotated coexpression modules to .json file
    with open(os.path.join(args.outdir, "coexpressionModules_annotated.json"), 'w') as outfile:
        json.dump(annotated_coexpression_modules, outfile)"""


    # Get eigengenes for all modules
    eigengenes = miner.getEigengenes(regulon_modules, exp_data, regulon_dict=None,
                                     saveFolder=None)
    eigen_scale = np.percentile(exp_data, 95) / np.percentile(eigengenes, 95)
    eigengenes = eigen_scale * eigengenes
    eigengenes.index = np.array(eigengenes.index).astype(str)

    # write eigengenes to .csv
    eigengenes.to_csv(os.path.join(args.outdir, "eigengenes.csv"))

    t2 = time.time()
    logging.info("Completed mechanistic inference in {:.2f} minutes".format((t2 - t1) / 60.))
    logging.info("Inferred network with {:d} regulons, {:d} regulators, and {:d} co-regulated genes".format(len(regulon_df.Regulon_ID.unique()), len(regulon_df.Regulator.unique()),len(regulon_df.Gene.unique())))
