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 = 7
sample_name = 'Y_NPC'

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: 4880 SequenceContrib ScoresHyp_Contrib Scores	Ctcf, Ctcfl
metacluster_1/pattern_1 # seqlets: 3104 SequenceContrib ScoresHyp_Contrib Scores	Sp2, Sp3, Sp1, Sp4, Klf3, E2f4, Klf4, Klf6, Zbtb17, Sp5, Egr1, Klf7, Klf5, Klf8, Klf12, Wt1, Egr2, Zfp281, Klf1, Maz, E2f1, Srebf2, Klf15, Egr3
metacluster_1/pattern_2 # seqlets: 1950 SequenceContrib ScoresHyp_Contrib Scores	Nrf1, Zfx
metacluster_1/pattern_3 # seqlets: 1669 SequenceContrib ScoresHyp_Contrib Scores	Creb1, Atf7, Atf2, Jdp2, Creb5, Crem, Atf1, Atf3, Mafb, Fosl2, Nfil3, Fosl1, Tef, Jun, Hlf, Dbp, Tfe3, Cebpg
metacluster_1/pattern_4 # seqlets: 1401 SequenceContrib ScoresHyp_Contrib Scores	Nfib, Nfic, Nfia
metacluster_1/pattern_5 # seqlets: 1257 SequenceContrib ScoresHyp_Contrib Scores	Lhx2, Pou2f1, Phox2b, Gbx1, Phox2a, Evx2, Prrx1, Otp, Isx, Esx1, En2, Lhx3, Rax, Hoxa2, Lhx8, Nkx1-2, Emx2, Hoxd1, Vsx2, Gbx2, Dlx3, Msx2, Hoxd3, Alx3, Pax6, Nobox, Gsx2, Pax7, Prrx2, Msx3, Noto, En1, Hoxa1, Lhx4, Nkx1-1, Msx1, Prop1, Vsx1, Vax2, Pou1f1, Uncx, Lhx9, Vax1, Evx1, Rhox6, Pou6f1, Lbx2, Shox2, Pou2f3, Dlx1
metacluster_1/pattern_6 # seqlets: 1187 SequenceContrib ScoresHyp_Contrib Scores	Ascl2, Ptf1a, Ascl1, Nhlh1, Myf6, Myog, Bhlha15, Tcf3, Tcf12, Myod1, Tal1, Nhlh2, Atoh1, Tcf7l1, Tcf4, Lyl1, Neurod2, Olig2, Snai2
metacluster_1/pattern_7 # seqlets: 1158 SequenceContrib ScoresHyp_Contrib Scores	Nfyb, Nfyc, Nfya, Foxi1, Pbx3
metacluster_1/pattern_8 # seqlets: 1064 SequenceContrib ScoresHyp_Contrib Scores	Rfx3, Rfx2, Rfx1, Rfx4, Rfx6, Rfx7, Sox2, Pou5f1
metacluster_1/pattern_9 # seqlets: 871 SequenceContrib ScoresHyp_Contrib Scores	Sox9, Sox3, Sox4, Sox6, Sox17
metacluster_1/pattern_10 # seqlets: 763 SequenceContrib ScoresHyp_Contrib Scores	Gabpa, Fli1, Elk1, Ets1, Elk4, Erg, Elf1, Elk3, Etv1, Etv4, Etv2, Etv5, Fev, Etv3, Etv6, XP_911724.4, Ehf, Elf4, Spi1, Elf2, Elf3, Elf5, Spdef, Spib, Irf8
metacluster_1/pattern_11 # seqlets: 745 SequenceContrib ScoresHyp_Contrib Scores	Usf1, Usf2, Tfeb, Mitf, Arntl, Tfec, Clock, Arnt, Srebf1, Bhlhe40, Mlx, Mycn, Npas2, Myc, Max, Tcfl5, Creb3l2, Rxrg, Rxrb, Mxi1
metacluster_1/pattern_12 # seqlets: 667 SequenceContrib ScoresHyp_Contrib Scores	Thap11, Zfp143, Tbx2
metacluster_1/pattern_13 # seqlets: 275 SequenceContrib ScoresHyp_Contrib Scores	Zbtb33
metacluster_1/pattern_14 # seqlets: 227 SequenceContrib ScoresHyp_Contrib Scores	Irf9, Irf2, Irf1, Irf7, Prdm1, Irf4, Stat2, Irf5
metacluster_1/pattern_15 # seqlets: 176 SequenceContrib ScoresHyp_Contrib Scores	Fos, Jund, Mafk, Bach2, Junb, Fosb, Nfe2l2, Bach1, Mafa
metacluster_1/pattern_16 # seqlets: 120 SequenceContrib ScoresHyp_Contrib Scores	Taf1, Yy1, Zfp42