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 = 1
sample_name = 'O_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_1/pattern_0 # seqlets: 4088 SequenceContrib ScoresHyp_Contrib Scores	Sp4, Sp2, Sp3, Klf3, Sp1, Klf6, Klf4, E2f4, Klf7, Klf1, Klf8, Klf12, Egr1, Zbtb17, Egr2, E2f1, Klf5, Sp5, Zfp281, Zfx
metacluster_1/pattern_1 # seqlets: 3238 SequenceContrib ScoresHyp_Contrib Scores	Ets1, Etv4, Erg, Ehf, Elf2, Etv2, Elf5, Elf4, Etv6, Fli1, Etv5, Elf3, Gabpa, XP_911724.4, Etv3, Etv1, Elk3, Elk4, Elk1, Elf1, Fev, Spi1, Spdef, Stat3, Spic, Stat4, Stat6, Spib
metacluster_1/pattern_2 # seqlets: 1807 SequenceContrib ScoresHyp_Contrib Scores	Nfyb, Foxi1, Nfyc, Nfya, Pbx3
metacluster_1/pattern_3 # seqlets: 1341 SequenceContrib ScoresHyp_Contrib Scores	Nrf1
metacluster_1/pattern_4 # seqlets: 1217 SequenceContrib ScoresHyp_Contrib Scores	Creb1, Atf1, Atf7, Creb5, Atf2, Crem, Jdp2, Mafb, Nfil3, Jun, Tef, Rxrb, Dbp, Gmeb2, Hlf, Cebpb, Atf3, Fosl1
metacluster_1/pattern_5 # seqlets: 718 SequenceContrib ScoresHyp_Contrib Scores	Foxj2, Foxc2, Foxm1, Foxc1, Foxd3, Foxa3, Foxa1, Foxl1, Foxa2, Foxl2, Foxb1, Foxd1, Foxk1, Gm5294, Foxg1, Foxj3, Foxq1, Foxj1, Foxp2, Foxo3, Foxo4, Foxp1
metacluster_1/pattern_6 # seqlets: 710 SequenceContrib ScoresHyp_Contrib Scores	Zfp143, Thap11, Tbx2
metacluster_1/pattern_7 # seqlets: 568 SequenceContrib ScoresHyp_Contrib Scores	Sox4, Sox3, Sry, Sox6, Sox12, Sox18, Sox14, Sox2, Sox7, Sox17, Sox9, Sox8, Sox11, Sox15, Sox21
metacluster_1/pattern_8 # seqlets: 443 SequenceContrib ScoresHyp_Contrib Scores	Junb, Fosb, Fosl2, Bach2, Jund, Fos, Nfe2l2, Bach1, Nfe2, Batf, Mafk, Mafg, Mafa, Maf, Maff, Nr1d2
metacluster_1/pattern_9 # seqlets: 335 SequenceContrib ScoresHyp_Contrib Scores	Usf2, Tfeb, Arntl, Tfe3, Mxi1, Arnt, Clock, Usf1, Mitf, Mlx, Tfec, Bhlhe40, Npas2, Srebf1, Bhlhe41, Mycn, Myc, Rxrg
metacluster_1/pattern_10 # seqlets: 245 SequenceContrib ScoresHyp_Contrib Scores	Ctcf, Ctcfl
metacluster_1/pattern_11 # seqlets: 173 SequenceContrib ScoresHyp_Contrib Scores	Nfib, Nfic, Nfia
metacluster_1/pattern_12 # seqlets: 173 SequenceContrib ScoresHyp_Contrib Scores	Irf1, Irf2, Irf8, Irf7, Prdm1, Irf9, Stat1, Stat2, Irf4, Irf3
metacluster_1/pattern_13 # seqlets: 169 SequenceContrib ScoresHyp_Contrib Scores	Mef2d, Mef2a, Mef2c
metacluster_1/pattern_14 # seqlets: 166 SequenceContrib ScoresHyp_Contrib Scores	Wt1, Maz, Klf15
metacluster_1/pattern_15 # seqlets: 94 SequenceContrib ScoresHyp_Contrib Scores	Taf1, Yy1, Zfp42, Rfx1
metacluster_1/pattern_16 # seqlets: 78 SequenceContrib ScoresHyp_Contrib Scores	Zbtb33
metacluster_1/pattern_17 # seqlets: 70 SequenceContrib ScoresHyp_Contrib Scores	Gata4, Gata2, Tal1, Gata1, Gata3, Gata6, Mecom, Gata5