
import os
import sys
sys.path.append(os.path.abspath("/users/amtseng/tfmodisco/src/"))
from tfmodisco.run_tfmodisco import import_shap_scores, import_tfmodisco_results
from motif.read_motifs import pfm_info_content, pfm_to_pwm, trim_motif_by_ic
from motif.match_motifs import match_motifs_to_database
from util import figure_to_vdom_image
import plot.viz_sequence as viz_sequence
import numpy as np
import h5py
import sklearn.cluster
import scipy.cluster.hierarchy
import matplotlib.pyplot as plt
import vdom.helpers as vdomh
from IPython.display import display


# Define parameters/fetch arguments
tf_name = os.environ["TFM_TF_NAME"]
tfm_results_path = os.environ["TFM_TFM_PATH"]
if "TFM_MOTIF_CACHE" in os.environ:
    tfm_motifs_cache_dir = os.environ["TFM_MOTIF_CACHE"]
else:
    tfm_motifs_cache_dir = None

print("TF name: %s" % tf_name)
print("TF-MoDISco results path: %s" % tfm_results_path)
print("Saved TF-MoDISco-derived motifs cache: %s" % tfm_motifs_cache_dir)

TF name: Sox2
TF-MoDISco results path: /users/amtseng/tfmodisco/results/tfmodisco/BPNet/BPNet_Sox2_ChIPseq/BPNet_Sox2_ChIPseq_profile_tfm.h5
Saved TF-MoDISco-derived motifs cache: /users/amtseng/tfmodisco/results/reports/tfmodisco_results//cache/BPNet/BPNet_Sox2_ChIPseq/BPNet_Sox2_ChIPseq_profile


# Define paths and constants
min_ic = 0.3


if tfm_motifs_cache_dir:
    os.makedirs(tfm_motifs_cache_dir, exist_ok=True)


motif_pfms, motif_hcwms, motif_cwms = [], [], []  # Save the trimmed PFMs, hCWMs, and CWMs
motif_pfms_short = []  # PFMs that are even more trimmed (for TOMTOM)
num_seqlets = []  # Number of seqlets for each motif
motif_seqlets = []  # Save seqlets of each motif
if tfm_motifs_cache_dir:
    motif_hdf5 = h5py.File(os.path.join(tfm_motifs_cache_dir, "all_motifs.h5"), "w")

with h5py.File(tfm_results_path, "r") as f:
    metaclusters = f["metacluster_idx_to_submetacluster_results"]
    num_metaclusters = len(metaclusters.keys())
    for metacluster_i, metacluster_key in enumerate(metaclusters.keys()):
        metacluster = metaclusters[metacluster_key]
        display(vdomh.h3("Metacluster %d/%d" % (metacluster_i + 1, num_metaclusters)))
        
        if "patterns" not in metacluster["seqlets_to_patterns_result"].keys():
            continue
        patterns = metacluster["seqlets_to_patterns_result"]["patterns"]
        num_patterns = len(patterns["all_pattern_names"][:])
        
        motif_pfms.append([])
        motif_hcwms.append([])
        motif_cwms.append([])
        motif_pfms_short.append([])
        num_seqlets.append([])
        motif_seqlets.append([])
        num_patterns = len(patterns)
        
        for pattern_i, pattern_name in enumerate(patterns["all_pattern_names"][:]):
            pattern_name = pattern_name.decode()
            pattern = patterns[pattern_name]
            seqlets = pattern["seqlets_and_alnmts"]["seqlets"]
            
            display(vdomh.h4("Pattern %d/%d" % (pattern_i + 1, num_patterns)))
            display(vdomh.p("%d seqlets" % len(seqlets)))
            
            pfm = pattern["sequence"]["fwd"][:]
            hcwm = pattern["task0_hypothetical_contribs"]["fwd"][:]
            cwm = pattern["task0_contrib_scores"]["fwd"][:]
        
            pfm_fig = viz_sequence.plot_weights(pfm, subticks_frequency=10, return_fig=True)
            hcwm_fig = viz_sequence.plot_weights(hcwm, subticks_frequency=10, return_fig=True)
            cwm_fig = viz_sequence.plot_weights(cwm, subticks_frequency=10, return_fig=True)
            pfm_fig.tight_layout()
            hcwm_fig.tight_layout()
            cwm_fig.tight_layout()

            motif_table = vdomh.table(
                vdomh.tr(
                    vdomh.td("Sequence (PFM)"),
                    vdomh.td(figure_to_vdom_image(pfm_fig))
                ),
                vdomh.tr(
                    vdomh.td("Hypothetical contributions (hCWM)"),
                    vdomh.td(figure_to_vdom_image(hcwm_fig))
                ),
                vdomh.tr(
                    vdomh.td("Actual contributions (CWM)"),
                    vdomh.td(figure_to_vdom_image(cwm_fig))
                )
            )
            display(motif_table)
            plt.close("all")  # Remove all standing figures

            # Trim motif based on information content
            short_trimmed_pfm = trim_motif_by_ic(pfm, pfm, min_ic=min_ic)
            motif_pfms_short[-1].append(short_trimmed_pfm)

            # Expand trimming to +/- 4bp on either side
            trimmed_pfm = trim_motif_by_ic(pfm, pfm, min_ic=min_ic, pad=4)
            trimmed_hcwm = trim_motif_by_ic(pfm, hcwm, min_ic=min_ic, pad=4)
            trimmed_cwm = trim_motif_by_ic(pfm, cwm, min_ic=min_ic, pad=4)

            motif_pfms[-1].append(trimmed_pfm)
            motif_hcwms[-1].append(trimmed_hcwm)
            motif_cwms[-1].append(trimmed_cwm)

            num_seqlets[-1].append(len(seqlets))

            if tfm_motifs_cache_dir:
                # Save results and figures
                motif_id = "%d_%d" % (metacluster_i, pattern_i)
                pfm_fig.savefig(os.path.join(tfm_motifs_cache_dir, motif_id + "_pfm_full.png"))
                hcwm_fig.savefig(os.path.join(tfm_motifs_cache_dir, motif_id + "_hcwm_full.png"))
                cwm_fig.savefig(os.path.join(tfm_motifs_cache_dir, motif_id + "_cwm_full.png"))
                motif_dset = motif_hdf5.create_group(motif_id)
                motif_dset.create_dataset("pfm_full", data=pfm, compression="gzip")
                motif_dset.create_dataset("hcwm_full", data=hcwm, compression="gzip")
                motif_dset.create_dataset("cwm_full", data=cwm, compression="gzip")
                motif_dset.create_dataset("pfm_trimmed", data=trimmed_pfm, compression="gzip")
                motif_dset.create_dataset("hcwm_trimmed", data=trimmed_hcwm, compression="gzip")
                motif_dset.create_dataset("cwm_trimmed", data=trimmed_cwm, compression="gzip")
                motif_dset.create_dataset("pfm_short_trimmed", data=short_trimmed_pfm, compression="gzip")
    if tfm_motifs_cache_dir:
        motif_hdf5.close()


colgroup = vdomh.colgroup(
    vdomh.col(style={"width": "5%"}),
    vdomh.col(style={"width": "5%"}),
    vdomh.col(style={"width": "45%"}),
    vdomh.col(style={"width": "45%"})
)
header = vdomh.thead(
    vdomh.tr(
        vdomh.th("#", style={"text-align": "center"}),
        vdomh.th("Seqlets", style={"text-align": "center"}),
        vdomh.th("Forward", style={"text-align": "center"}),
        vdomh.th("Reverse", style={"text-align": "center"})
    )
)

for i in range(len(motif_hcwms)):
    display(vdomh.h3("Metacluster %d/%d" % (i + 1, num_metaclusters)))
    body = []
    for j in range(len(motif_hcwms[i])):
        motif = motif_hcwms[i][j]
        if np.sum(motif[:, [0, 2]]) > 0.5 * np.sum(motif):
            # Forward is purine-rich, reverse-complement is pyrimidine-rich
            f, rc = motif, np.flip(motif, axis=(0, 1))
        else:
            f, rc = np.flip(motif, axis=(0, 1)), motif
            
        f_fig = viz_sequence.plot_weights(f, figsize=(20, 4), return_fig=True)
        f_fig.tight_layout()
        rc_fig = viz_sequence.plot_weights(rc, figsize=(20, 4), return_fig=True)
        rc_fig.tight_layout()
        
        if tfm_motifs_cache_dir:
            # Save results and figures
            motif_id = "%d_%d" % (i, j)
            f_fig.savefig(os.path.join(tfm_motifs_cache_dir, motif_id + "_hcwm_trimmed_fwd.png"))
            rc_fig.savefig(os.path.join(tfm_motifs_cache_dir, motif_id + "_hcwm_trimmed_rev.png"))

        body.append(
            vdomh.tr(
                vdomh.td(str(j + 1)),
                vdomh.td(str(num_seqlets[i][j])),
                vdomh.td(figure_to_vdom_image(f_fig)),
                vdomh.td(figure_to_vdom_image(rc_fig))
            )
        )
    display(vdomh.table(colgroup, header, vdomh.tbody(*body)))
    plt.close("all")


num_matches_to_keep = 10
num_matches_to_show = 5

header = vdomh.thead(
    vdomh.tr(
        vdomh.th("Motif ID", style={"text-align": "center"}),
        vdomh.th("q-val", style={"text-align": "center"}),
        vdomh.th("PWM", style={"text-align": "center"})
    )
)

for i in range(len(motif_pfms)):
    display(vdomh.h3("Metacluster %d/%d" % (i + 1, num_metaclusters)))
    
    # Compute TOMTOM matches for all motifs in the metacluster at once
    out_dir = os.path.join(tfm_motifs_cache_dir, "tomtom", "metacluster_%d" % i) if tfm_motifs_cache_dir else None
    tomtom_matches = match_motifs_to_database(motif_pfms_short[i], top_k=num_matches_to_keep, temp_dir=out_dir)
    
    for j in range(len(motif_pfms[i])):
        display(vdomh.h4("Motif %d/%d" % (j + 1, len(motif_pfms[i]))))
        viz_sequence.plot_weights(motif_hcwms[i][j])
    
        body = []
        for k, (match_name, match_pfm, match_qval) in enumerate(tomtom_matches[j]):
            fig = viz_sequence.plot_weights(pfm_to_pwm(match_pfm), return_fig=True)
            fig.tight_layout()
            if k < num_matches_to_show:
                body.append(
                    vdomh.tr(
                        vdomh.td(match_name),
                        vdomh.td(str(match_qval)),
                        vdomh.td(figure_to_vdom_image(fig))
                    )
                )
                if tfm_motifs_cache_dir:
                    # Save results and figures
                    motif_id = "%d_%d" % (i, j)
                    fig.savefig(os.path.join(out_dir, motif_id + ("_hit-%d.png" % (k + 1))))
            else:
                body.append(
                    vdomh.tr(
                        vdomh.td(match_name),
                        vdomh.td(str(match_qval)),
                        vdomh.td("Not shown")
                    )
                )
        if not body:
            display(vdomh.p("No TOMTOM matches passing threshold"))
        else:
            display(vdomh.table(header, vdomh.tbody(*body)))
        plt.close("all")

#	Seqlets	Forward	Reverse
1	3890
2	1706
3	275
4	79
5	75
6	61
7	60
8	34

Motif ID	q-val	PWM
MA0142.1_Pou5f1::Sox2	6.39917e-12
PO5F1_HUMAN.H11MO.0.A	6.39917e-12
NANOG_HUMAN.H11MO.0.A	1.75909e-11
SOX2_HUMAN.H11MO.1.A	4.98961e-06
MA0792.1_POU5F1B	0.00020548799999999998
MA0789.1_POU3F4	0.000452989	Not shown
PO3F1_HUMAN.H11MO.0.C	0.0007064969999999999	Not shown
MA1115.1_POU5F1	0.00348698	Not shown
PO5F1_HUMAN.H11MO.1.A	0.00505667	Not shown
MA0786.1_POU3F1	0.0123131	Not shown

Motif ID	q-val	PWM
MA0868.2_SOX8	4.18004e-05
MA0077.1_SOX9	0.000150458
SOX13_HUMAN.H11MO.0.D	0.0007648139999999999
MA0867.2_SOX4	0.00469998
MA1120.1_SOX13	0.00469998
MA0515.1_Sox6	0.00469998	Not shown
MA1152.1_SOX15	0.00469998	Not shown
SOX9_HUMAN.H11MO.0.B	0.0067804	Not shown
MA0514.1_Sox3	0.00905843	Not shown
SOX17_HUMAN.H11MO.0.C	0.00905843	Not shown

Motif ID	q-val	PWM
ZIC3_HUMAN.H11MO.0.B	6.485039999999999e-09
ZIC2_HUMAN.H11MO.0.D	1.2765100000000001e-08
MA1629.1_Zic2	8.96541e-05
MA0696.1_ZIC1	0.00403624
MA0751.1_ZIC4	0.00420461
MA0697.1_ZIC3	0.027754700000000004	Not shown
MA1628.1_Zic1::Zic2	0.028470799999999997	Not shown
ZIC4_HUMAN.H11MO.0.D	0.028749	Not shown
MA1584.1_ZIC5	0.0326596	Not shown
MA1648.1_TCF12(var.2)	0.0366161	Not shown

Motif ID	q-val	PWM
SOX15_HUMAN.H11MO.0.D	0.013068200000000002
MA0077.1_SOX9	0.013068200000000002
MA0867.2_SOX4	0.015465999999999999
SOX4_HUMAN.H11MO.0.B	0.015465999999999999
MA0515.1_Sox6	0.016246
SOX13_HUMAN.H11MO.0.D	0.01999	Not shown
MA1152.1_SOX15	0.021987200000000002	Not shown
MA0442.2_SOX10	0.021987200000000002	Not shown
SOX9_HUMAN.H11MO.1.B	0.021987200000000002	Not shown
SOX2_HUMAN.H11MO.0.A	0.028968900000000002	Not shown

Direct links to results¶

Define constants and paths¶

Plot TF-MoDISco results¶

Metacluster 1/1

Pattern 1/9

Pattern 2/9

Pattern 3/9

Pattern 4/9

Pattern 5/9

Pattern 6/9

Pattern 7/9

Pattern 8/9

Summary of motifs¶

Metacluster 1/1

Top TOMTOM matches for each motif¶

Metacluster 1/1

Motif 1/8

Motif 2/8

Motif 3/8

Motif 4/8

Motif 5/8

Motif 6/8

Motif 7/8

Motif 8/8

Sequence (PFM)
Hypothetical contributions (hCWM)
Actual contributions (CWM)

Motif ID	q-val	PWM
SOX9_HUMAN.H11MO.0.B	0.14871099999999998
MA1530.1_NKX6-3	0.321355
MA0089.2_NFE2L1	0.321355
MA1142.1_FOSL1::JUND	0.321355
HXB3_HUMAN.H11MO.0.D	0.321355
HMGA2_HUMAN.H11MO.0.D	0.321355	Not shown
DLX6_HUMAN.H11MO.0.D	0.321355	Not shown
MA0489.1_JUN(var.2)	0.321355	Not shown
FEZF1_HUMAN.H11MO.0.C	0.321355	Not shown
MA0841.1_NFE2	0.321355	Not shown

Motif ID	q-val	PWM
MA0505.1_Nr5a2	0.0054961
NR6A1_HUMAN.H11MO.0.B	0.00658241
MA1541.1_NR6A1	0.10760499999999999
ERR3_HUMAN.H11MO.0.B	0.10760499999999999
ERR2_HUMAN.H11MO.0.A	0.10760499999999999
MA0071.1_RORA	0.10760499999999999	Not shown
NR5A2_HUMAN.H11MO.0.B	0.14736	Not shown
RORA_HUMAN.H11MO.0.C	0.14990499999999998	Not shown
RARA_HUMAN.H11MO.1.A	0.157195	Not shown
MA0160.1_NR4A2	0.183158	Not shown

Motif ID	q-val	PWM
ZIC3_HUMAN.H11MO.0.B	8.49922e-06
ZIC2_HUMAN.H11MO.0.D	8.49922e-06
MA1629.1_Zic2	0.00131513
MA0696.1_ZIC1	0.011888
MA0751.1_ZIC4	0.011888
MA0697.1_ZIC3	0.0264498	Not shown
ZIC4_HUMAN.H11MO.0.D	0.0282423	Not shown
MA1584.1_ZIC5	0.067971	Not shown
MA1628.1_Zic1::Zic2	0.067971	Not shown
SOX9_HUMAN.H11MO.0.B	0.0912138	Not shown