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 = 4
sample_name = 'O_qNSC'

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: 3325 SequenceContrib ScoresHyp_Contrib Scores	Ctcf, Ctcfl
metacluster_1/pattern_1 # seqlets: 2471 SequenceContrib ScoresHyp_Contrib Scores	Sp2, Sp3, Klf3, Sp4, Sp1, Klf6, Klf15, Zfx, Zbtb17, Wt1, Klf1, Maz, Sp5, Egr1, Zfp281, E2f4, Klf5, Klf4, Klf7, Egr2, Klf8, Klf12, Usf2, Tfap2c, Rela, E2f1, Fli1, Nr1h4, E2f7, Ptf1a, Mxi1
metacluster_1/pattern_2 # seqlets: 1786 SequenceContrib ScoresHyp_Contrib Scores	Nrf1
metacluster_1/pattern_3 # seqlets: 1503 SequenceContrib ScoresHyp_Contrib Scores	Nfib, Nfic
metacluster_1/pattern_4 # seqlets: 1260 SequenceContrib ScoresHyp_Contrib Scores	Nfyb, Nfyc, Nfya, Foxi1, Pbx3
metacluster_1/pattern_5 # seqlets: 1256 SequenceContrib ScoresHyp_Contrib Scores	Jdp2, Creb1, Atf7, Atf2, Creb5, Mafb, Crem, Atf3, Fosl1, Atf1, Fosl2, Tef, Jun, Nfil3, Hlf, Gmeb2, Dbp, Cebpg, Cebpb, Rorb, Gmeb1, Atf4, Junb, Tfe3, Jund, Cebpa
metacluster_1/pattern_6 # seqlets: 1092 SequenceContrib ScoresHyp_Contrib Scores	Gabpa, Ets1, Elk4, Erg, Elk1, Elk3, Etv1, Etv4, Elf1, Etv5, Fev, Etv2, Etv3, XP_911724.4, Etv6, Elf4, Ehf, Elf2, Elf5, Spi1, Spdef, Elf3
metacluster_1/pattern_7 # seqlets: 848 SequenceContrib ScoresHyp_Contrib Scores	Vsx2, Lhx8, En2, Prrx2, Rhox6, Lhx2, Shox2, Alx1, Lhx9, Emx2, Pou1f1, Lhx3, Rax, Arx, Pax6, Uncx, Hoxa2, Lbx2, Vsx1, Lhx4, Noto, Gbx1, Alx3, Alx4, Sebox, Gsx2, Pax4, Gbx2, Nkx1-2, Hoxd3, Msx3, Prop1, Evx1, Vax2, Lhx6, Phox2b, Phox2a, Nkx1-1, Esx1, Prrx1, Hoxa1, En1, Vax1, Otp, Lmx1b, Evx2, Nobox, Pou2f3, Msx1, Mnx1, Pax7, Pou2f2, Isx, Hoxc5, Pdx1, Hoxa3, Hoxd1, Hoxb3, Pou3f1, Pou3f2, Lhx5, Pou2f1, Lhx1, Msx2, Lmx1a, Hlx, Dlx3, Hoxa5, Hoxc6, Barhl1, Dlx5, Barhl2, Isl1, Isl2, Hoxb5, Meox2, Dbx1, Pou3f4, Tlx2, Pou4f3, Dbx2, Meox1, Nkx6-3, Pou6f1, Nkx6-1, Hnf1b, Bsx, Hoxb4, Dlx2, Arid3b
metacluster_1/pattern_8 # seqlets: 699 SequenceContrib ScoresHyp_Contrib Scores	Thap11, Zfp143, Tbx2
metacluster_1/pattern_9 # seqlets: 685 SequenceContrib ScoresHyp_Contrib Scores	Usf1, Tfeb, Mitf, Arntl, Tfec, Clock, Srebf1, Bhlhe40, Arnt, Mlx, Npas2, Mycn, Max, Myc, Srebf2, Rfx1, Tcfl5, Creb3l2, Xbp1, Mlxip, Bhlhe41, Tfap2b
metacluster_1/pattern_10 # seqlets: 634 SequenceContrib ScoresHyp_Contrib Scores	Rfx3, Rfx2, Rfx4, Rfx7, Rfx6
metacluster_1/pattern_11 # seqlets: 398 SequenceContrib ScoresHyp_Contrib Scores	Sox9, Sox3, Sox4, Sox2, Sox6, Sox17, Tcf7
metacluster_1/pattern_12 # seqlets: 349 SequenceContrib ScoresHyp_Contrib Scores	Irf2, Irf1, Irf7, Irf9, Stat1, Prdm1, Stat2, Irf8, Irf4, Irf3, Spib
metacluster_1/pattern_13 # seqlets: 307 SequenceContrib ScoresHyp_Contrib Scores	Taf1, Yy1, Zfp42
metacluster_1/pattern_14 # seqlets: 184 SequenceContrib ScoresHyp_Contrib Scores	Zbtb33
metacluster_1/pattern_15 # seqlets: 50 SequenceContrib ScoresHyp_Contrib Scores	Rest