from matlas.matches import DenovoModisco, DenovoHomer
from vdom.helpers import (b, summary)
from IPython.display import display
import numpy as np

def adjust_motifs(ob):
    for motif_name in ob.denovo_pwms.keys():
        a = ob.get_pwm(motif_name)
        ones = np.sum(a, axis=1)
        for i, one in enumerate(ones):
            if (one-1.0)!=1:
            #print('decrease', motif_name, i)
                part = (one-1.0)/4
                ob.denovo_pwms[motif_name]['PPM'][i] -= part
    
    return ob

def display_paiwise_pattern_comparison(task_idx, sample_name, modiscodir, homerdir):
    display(summary(b(sample_name)))
    
    #homerdir = "{0}/task_{1}".format(homer_root, task_idx)
    homer_ob = DenovoHomer(homerdir)
    homer_ob = adjust_motifs(homer_ob)
    homer_ob.load_matched_motifs()
    homer_ob.get_motif_per_celltype()

    #modiscodir = "{0}/task_{1}".format(modisco_root, task_idx)
    modisco_ob = DenovoModisco(modiscodir)
    modisco_ob.load_matched_motifs()
    modisco_ob.get_motif_per_celltype()

    modisco_ob.display_pairwise_tf_table(homer_ob)
    return None

def display_denovo_patterns(sample_name, modiscodir, match_threshold=0.05):
    display(summary(b(sample_name)))
    
    ob = DenovoModisco(modiscodir)
    ob.load_matched_motifs()
    ob.get_motif_per_celltype(match_threshold=match_threshold)
    ob.display_individual_table()
    
    return None

root = "/mnt/lab_data/kundaje/users/msharmin/NSC_ATAC_PEAKS_to_share"

task_idx = 6
sample_name = 'Y_Endo'

display_denovo_patterns(
    sample_name,
    modiscodir="{0}/gw_modisco/multi_tasks/task_0.{1}".format(root, task_idx)
)

display_paiwise_pattern_comparison(
    task_idx, sample_name, 
    modiscodir="{0}/gw_modisco/multi_tasks/task_0.{1}".format(root, task_idx),
    homerdir="{0}/gw_homer/task_{1}".format(root, task_idx),
)

Pattern Name	TF Name(s)	Modisco
metacluster_0/pattern_0 # seqlets: 2160 SequenceContrib ScoresHyp_Contrib Scores	Sp2, Sp3, Sp1, Klf3, Klf6, Sp4, Sp5, Zfx, Klf1, Wt1, Klf15, Zbtb17, Maz, Egr1, Klf5, Zfp281, Klf4, Klf7, Egr2, Usf2, E2f4, Klf8, Klf12, Ctcfl, E2f1, Mxi1, Nr1h4, Tfap2c, Zfp740, Sall4, Mbd2, Srebf2, Fli1
metacluster_0/pattern_1 # seqlets: 1425 SequenceContrib ScoresHyp_Contrib Scores	Nfyb, Nfyc, Nfya, Foxi1, Pbx3
metacluster_0/pattern_2 # seqlets: 1337 SequenceContrib ScoresHyp_Contrib Scores	Rela, E2f3
metacluster_0/pattern_3 # seqlets: 1162 SequenceContrib ScoresHyp_Contrib Scores	Nrf1
metacluster_0/pattern_4 # seqlets: 979 SequenceContrib ScoresHyp_Contrib Scores	Etv2, Etv4, Ets1, Erg, Etv6, Elf2, Etv5, Etv3, Gabpa, Elk3, XP_911724.4, Etv1, Elk4, Elk1, Fev, Spi1, Ehf, Elf4, Elf3, Spdef, Elf5, Spic, Spib
metacluster_0/pattern_5 # seqlets: 929 SequenceContrib ScoresHyp_Contrib Scores	Creb1, Crem, Jdp2, Atf7, Atf2, Atf1, Creb5, Atf3, Jun
metacluster_0/pattern_6 # seqlets: 782 SequenceContrib ScoresHyp_Contrib Scores	Elf1, Irf8
metacluster_0/pattern_7 # seqlets: 680 SequenceContrib ScoresHyp_Contrib Scores	Thap11, Zfp143, Tbx2
metacluster_0/pattern_8 # seqlets: 501 SequenceContrib ScoresHyp_Contrib Scores	Ctcf, Rfx1
metacluster_0/pattern_9 # seqlets: 408 SequenceContrib ScoresHyp_Contrib Scores	Foxd3, Foxj3, Foxj2, Gm5294, Foxc2, Foxc1, Foxj1, Foxo4, Foxd1, Foxo3, Foxq1, Foxl1, Foxp2, Foxk1, Foxa1, Foxl2, Foxa3, Foxg1, Foxa2, Foxb1, Foxm1, Foxp1, Zfp105, Foxo1, Onecut1
metacluster_0/pattern_10 # seqlets: 279 SequenceContrib ScoresHyp_Contrib Scores	Usf1, Tfeb, Arntl, Tfe3, Mitf, Clock, Tfec, Arnt, Srebf1, Bhlhe40, Mlx, Mycn, Npas2, Max, Myc, Tcfl5, Creb3l2, Rxrg
metacluster_0/pattern_11 # seqlets: 173 SequenceContrib ScoresHyp_Contrib Scores	A1JVI6_MOUSE
metacluster_0/pattern_12 # seqlets: 109 SequenceContrib ScoresHyp_Contrib Scores	Yy1, Taf1, Zfp42
metacluster_0/pattern_13 # seqlets: 105 SequenceContrib ScoresHyp_Contrib Scores	Bcl6
metacluster_0/pattern_14 # seqlets: 99 SequenceContrib ScoresHyp_Contrib Scores	Jund, Fosl2, Junb, Fos, Fosb, Fosl1, Batf, Bach2, Pknox1, Mafk, Bach1
metacluster_0/pattern_15 # seqlets: 84 SequenceContrib ScoresHyp_Contrib Scores	Irf1, Irf2, Irf7, Irf9, Irf4, Stat2, Prdm1, Stat1, Irf5, Irf3, Irf6
metacluster_0/pattern_16 # seqlets: 67 SequenceContrib ScoresHyp_Contrib Scores	Zbtb33
metacluster_0/pattern_17 # seqlets: 51 SequenceContrib ScoresHyp_Contrib Scores	Nfib
metacluster_0/pattern_18 # seqlets: 39 SequenceContrib ScoresHyp_Contrib Scores	Tead4, Pou3f1