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 = 2
sample_name = 'O_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: 4644 SequenceContrib ScoresHyp_Contrib Scores	Ctcf, Ctcfl
metacluster_1/pattern_1 # seqlets: 2979 SequenceContrib ScoresHyp_Contrib Scores	Sp2, Sp3, Sp1, Sp4, Klf3, E2f4, Klf6, Klf4, Zbtb17, Sp5, Egr1, Klf5, Klf7, Wt1, Klf8, Klf12, Egr2, Zfp281, Klf1, Maz, E2f1, Klf15, Egr3, Srebf2
metacluster_1/pattern_2 # seqlets: 1926 SequenceContrib ScoresHyp_Contrib Scores	Nrf1
metacluster_1/pattern_3 # seqlets: 1785 SequenceContrib ScoresHyp_Contrib Scores	Atf2, Atf7, Creb1, Jdp2, Creb5, Crem, Atf1, Atf3, Mafb, Fosl2, Fosl1, Nfil3, Jun, Tef, Dbp, Hlf, Cebpg, Tfe3
metacluster_1/pattern_4 # seqlets: 1637 SequenceContrib ScoresHyp_Contrib Scores	Nfib, Nfia, Nfic
metacluster_1/pattern_5 # seqlets: 1260 SequenceContrib ScoresHyp_Contrib Scores	Nfyb, Nfyc, Nfya, Foxi1, Pbx3
metacluster_1/pattern_6 # seqlets: 1182 SequenceContrib ScoresHyp_Contrib Scores	Lhx2, Phox2b, Gbx1, Pou2f1, Phox2a, Evx2, Prrx1, Otp, Isx, En2, Esx1, Lhx3, Rax, Hoxa2, Nkx1-2, Dlx3, Lhx8, Emx2, Hoxd1, Msx2, Gsx2, Gbx2, Hoxd3, Nobox, Pax6, Vsx2, Prrx2, Pax7, Msx3, Alx3, Noto, Hoxa1, Nkx1-1, En1, Lhx4, Vsx1, Vax2, Prop1, Pou1f1, Lhx9, Uncx, Msx1, Vax1, Evx1, Dlx1, Pou2f3, Dlx2, Rhox6, Lbx2, Shox2, Barhl1, Lhx6, Barx1, Meox1, Arx, Pou2f2, Alx1, Pou6f1, Hoxb3, Hoxc5, Hnf1a, Alx4, Pdx1, Hoxa5, Bsx, Sebox
metacluster_1/pattern_7 # seqlets: 1041 SequenceContrib ScoresHyp_Contrib Scores	Ascl2, Ascl1, Nhlh1, Myf6, Tcf3, Myog, Ptf1a, Nhlh2, Myod1, Tcf12, Bhlha15, Tal1, Tcf7l1, Atoh1, Tcf4, Snai2
metacluster_1/pattern_8 # seqlets: 863 SequenceContrib ScoresHyp_Contrib Scores	Rfx3, Rfx2, Rfx1, Rfx4, Rfx7, Rfx6, Sox2, Pou5f1
metacluster_1/pattern_9 # seqlets: 812 SequenceContrib ScoresHyp_Contrib Scores	Gabpa, Fli1, Elk1, Ets1, Elk4, Erg, Elf1, Elk3, Etv4, Etv1, Etv2, Etv5, Etv3, XP_911724.4, Etv6, Fev, Elf4, Ehf, Spi1, Elf2, Elf3, Elf5, Spdef, Spib, Irf8
metacluster_1/pattern_10 # seqlets: 788 SequenceContrib ScoresHyp_Contrib Scores	Sox9, Sox3, Sox4, Sox6, Sox17, Tcf7
metacluster_1/pattern_11 # seqlets: 780 SequenceContrib ScoresHyp_Contrib Scores	Usf2, Usf1, Tfeb, Tfec, Arntl, Srebf1, Mitf, Bhlhe40, Mlx, Clock, Arnt, Npas2, Creb3l2, Tcfl5, Mycn, Max, Myc, Mlxip, Bhlhe41, Xbp1
metacluster_1/pattern_12 # seqlets: 705 SequenceContrib ScoresHyp_Contrib Scores	Thap11, Zfp143, Tbx2
metacluster_1/pattern_13 # seqlets: 326 SequenceContrib ScoresHyp_Contrib Scores	Zbtb33
metacluster_1/pattern_14 # seqlets: 244 SequenceContrib ScoresHyp_Contrib Scores	Irf2, Irf1, Irf9, Irf7, Prdm1, Irf4, Stat2, Stat1, Irf5, Spic, Irf3
metacluster_1/pattern_15 # seqlets: 156 SequenceContrib ScoresHyp_Contrib Scores	Zfx, Fos, Junb, Bach2, Fosb, Jund, Mafk, Mafa, Bach1, Zfp335, Nfe2l2, Batf, Nr1h4
metacluster_1/pattern_16 # seqlets: 141 SequenceContrib ScoresHyp_Contrib Scores	Taf1, Yy1, Zfp42