import numpy as np
import subprocess
import argparse
import h5py
import tempfile
import os

def fetch_tomtom_args():
    parser = argparse.ArgumentParser()
    parser.add_argument("-m", "--modisco_h5py", required=True, type=str, help="path to the output .h5py file generated by the run_modisco.py script")
    parser.add_argument("-o", "--output_prefix", required=True, type=str, help="Path and name of the TSV file to store the tomtom output")
    parser.add_argument("-d", "--meme_motif_db", required=True, type=str, help="path to motif database")
    parser.add_argument("-n", "--top_n_matches", type=int, default=3, help="Max number of matches to return from TomTom")
    parser.add_argument("-tt", "--tomtom_exec", type=str, default='tomtom', help="Command to use to execute tomtom")
    parser.add_argument("-th", "--trim_threshold", type=float, default=0.3, help="Trim threshold for trimming long motif, trim to those with at least prob trim_threshold on both ends")
    parser.add_argument("-tm", "--trim_min_length", type=int, default=3, help="Minimum acceptable length of motif after trimming")
    args = parser.parse_args()
    return args


def write_meme_file(ppm, bg, fname):
    f = open(fname, 'w')
    f.write('MEME version 4\n\n')
    f.write('ALPHABET= ACGT\n\n')
    f.write('strands: + -\n\n')
    f.write('Background letter frequencies (from unknown source):\n')
    f.write('A %.3f C %.3f G %.3f T %.3f\n\n' % tuple(list(bg)))
    f.write('MOTIF 1 TEMP\n\n')
    f.write('letter-probability matrix: alength= 4 w= %d nsites= 1 E= 0e+0\n' % ppm.shape[0])
    for s in ppm:
        f.write('%.5f %.5f %.5f %.5f\n' % tuple(s))
    f.close()


def fetch_tomtom_matches(ppm, cwm, background=[0.25, 0.25, 0.25, 0.25], tomtom_exec_path='tomtom', motifs_db='HOCOMOCOv11_core_HUMAN_mono_meme_format.meme', n=5, trim_threshold=0.3, trim_min_length=3):

    """Fetches top matches from a motifs database using TomTom.
    Args:
        ppm: position probability matrix- numpy matrix of dimension (N,4)
        background: list with ACGT background probabilities
        tomtom_exec_path: path to TomTom executable
        motifs_db: path to motifs database in meme format
        n: number of top matches to return, ordered by p-value
        temp_dir: directory for storing temp files
        trim_threshold: the ppm is trimmed from left till first position for which
            probability for any base pair >= trim_threshold. Similarly from right.
    Returns:
        list: a list of up to n results returned by tomtom, each entry is a
            dictionary with keys 'Target ID', 'p-value', 'E-value', 'q-value'
    """

    _, fname = tempfile.mkstemp()

    score = np.sum(np.abs(cwm), axis=1)
    trim_thresh = np.max(score) * trim_threshold  # Cut off anything less than 30% of max score
    pass_inds = np.where(score >= trim_thresh)[0]
    trimmed = ppm[np.min(pass_inds): np.max(pass_inds) + 1]

    # can be None of no base has prob>t
    if trimmed is None:
        return []

    # trim and prepare meme file
    write_meme_file(trimmed, background, fname)

    # run tomtom
    cmd = '%s -no-ssc -oc . -verbosity 1 -text -min-overlap 5 -mi 1 -dist pearson -evalue -thresh 10.0 %s %s' % (tomtom_exec_path, fname, motifs_db)
    #print(cmd)
    out = subprocess.check_output(cmd, shell=True)

    # prepare output
    dat = [x.split('\\t') for x in str(out).split('\\n')]
    schema = dat[0]

    # meme v4 vs v5:
    if 'Target ID' in schema:
        tget_idx = schema.index('Target ID')
    else:
        tget_idx = schema.index('Target_ID')

    pval_idx, eval_idx, qval_idx =schema.index('p-value'), schema.index('E-value'), schema.index('q-value')

    r = []
    for t in dat[1:min(1+n, len(dat)-1)]:
        if t[0]=='':
            break

        mtf = {}
        mtf['Target_ID'] = t[tget_idx]
        mtf['p-value'] = float(t[pval_idx])
        mtf['E-value'] = float(t[eval_idx])
        mtf['q-value'] = float(t[qval_idx])
        r.append(mtf)

    os.system('rm ' + fname)
    return r


def main(): 
    args = fetch_tomtom_args()

    modisco_results = h5py.File(args.modisco_h5py, 'r')

    # get pfms
    ppms = []
    cwms = []
    seqlet_tally = []
    names = []

    for metacluster_name in modisco_results['metacluster_idx_to_submetacluster_results']:
        metacluster = modisco_results['metacluster_idx_to_submetacluster_results'][metacluster_name]
        all_pattern_names = [x.decode("utf-8") for x in list(metacluster["seqlets_to_patterns_result"]["patterns"]["all_pattern_names"][:])]

        for pattern_name in all_pattern_names:

            ppm = np.array(metacluster['seqlets_to_patterns_result']['patterns'][pattern_name]['sequence']['fwd'])
            num_seqlets = len(metacluster['seqlets_to_patterns_result']['patterns'][pattern_name]['seqlets_and_alnmts']['seqlets'])
            cwm = np.array(metacluster['seqlets_to_patterns_result']['patterns'][pattern_name]["task0_contrib_scores"]['fwd'])

            ppms.append(ppm)
            seqlet_tally.append(num_seqlets)
            cwms.append(cwm)
            names.append(metacluster_name + '.' + pattern_name)

    modisco_results.close()

    res = []

    for i,x in enumerate(ppms):
        res.append(fetch_tomtom_matches(x, cwms[i], tomtom_exec_path=args.tomtom_exec, motifs_db=args.meme_motif_db,
                   n=args.top_n_matches, trim_threshold=args.trim_threshold, trim_min_length=args.trim_min_length))

    # write output. Skipping those patterns which disappear after trimming or have no matches
    with open(args.output_prefix, 'w') as f:
        # write header
        f.write("Pattern")
        f.write("\tNum_Seqlets")

        for i in range(args.top_n_matches):
            f.write("\tMatch_{}\tq-value".format(i+1))
        f.write("\n")

        assert len(res) == len(names)

        for i,r in enumerate(res):
            f.write(names[i])
            f.write("\t{}".format(seqlet_tally[i]))
            for match in r:
                f.write("\t{}\t{}".format(match['Target_ID'], match['q-value']))

            # when fewer than n matches are found
            if len(r) != args.top_n_matches:
                f.write("\t\t"*(args.top_n_matches-len(r)))
            f.write("\n")
if __name__=="__main__":
    main()
