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 = 8
sample_name = 'Y_aNSC'

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: 5737 SequenceContrib ScoresHyp_Contrib Scores	Ctcf, Ctcfl
metacluster_1/pattern_1 # seqlets: 3010 SequenceContrib ScoresHyp_Contrib Scores	Sp2, Sp3, Sp1, Klf3, Sp4, Klf6, Klf1, Wt1, E2f4, Sp5, Egr1, Klf4, Klf7, Maz, Klf5, Zfx, Klf8, Zbtb17, Klf12, Klf15, Egr2, Zfp281, E2f1, E2f7, Egr3
metacluster_1/pattern_2 # seqlets: 2079 SequenceContrib ScoresHyp_Contrib Scores	Nfib, Nfic, Nfia
metacluster_1/pattern_3 # seqlets: 1909 SequenceContrib ScoresHyp_Contrib Scores	Nrf1
metacluster_1/pattern_4 # seqlets: 1552 SequenceContrib ScoresHyp_Contrib Scores	Atf7, Creb1, Atf2, Jdp2, Creb5, Crem, Atf1, Atf3, Mafb, Fosl2, Fosl1, Nfil3, Jun, Tef, Hlf, Dbp, Cebpg, Tfe3, Gmeb2
metacluster_1/pattern_5 # seqlets: 1356 SequenceContrib ScoresHyp_Contrib Scores	Nfyb, Nfyc, Nfya, Foxi1, Pbx3
metacluster_1/pattern_6 # seqlets: 1192 SequenceContrib ScoresHyp_Contrib Scores	Rfx3, Rfx2, Rfx1, Rfx4, Rfx6, Rfx7, Sox2, Pou5f1
metacluster_1/pattern_7 # seqlets: 1168 SequenceContrib ScoresHyp_Contrib Scores	Ptf1a, Ascl2, Ascl1, Bhlha15, Nhlh1, Myf6, Myog, Tcf3, Tcf12, Myod1, Tal1, Atoh1, Tcf7l1, Nhlh2, Tcf4, Lyl1, Neurod2, Olig2
metacluster_1/pattern_8 # seqlets: 1043 SequenceContrib ScoresHyp_Contrib Scores	Pou2f1, Lhx2, Phox2b, Phox2a, Gbx1, Prrx1, Otp, Vsx2, Evx2, Isx, Esx1, En2
metacluster_1/pattern_9 # seqlets: 1036 SequenceContrib ScoresHyp_Contrib Scores	Sox9, Sox10, Sox3, Sox4, Sox6
metacluster_1/pattern_10 # seqlets: 874 SequenceContrib ScoresHyp_Contrib Scores	Usf2, Usf1, Tfeb, Tfec, Arntl, Mitf, Srebf1, Bhlhe40, Clock, Mlx, Arnt, Npas2, Creb3l2, Tcfl5, Mycn, Srebf2, Max, Myc, Xbp1, Bhlhe41, Mlxip, Hes1
metacluster_1/pattern_11 # seqlets: 787 SequenceContrib ScoresHyp_Contrib Scores	Gabpa, Fli1, Elk1, Ets1, Elk4, Erg, Elf1, Elk3, Etv4, Etv1, Etv5, Fev, Etv2, Etv3, Etv6, XP_911724.4, Elf4, Ehf, Elf5, Elf2, Spi1, Spdef, Elf3
metacluster_1/pattern_12 # seqlets: 738 SequenceContrib ScoresHyp_Contrib Scores	Zfp143, Thap11, Tbx2
metacluster_1/pattern_13 # seqlets: 160 SequenceContrib ScoresHyp_Contrib Scores	Irf2, Irf7, Irf1, Irf9, Prdm1, Stat2, Stat1, Irf8, Irf3, Spib, Irf4, Irf5, Spic
metacluster_1/pattern_14 # seqlets: 152 SequenceContrib ScoresHyp_Contrib Scores	Zbtb33
metacluster_1/pattern_15 # seqlets: 118 SequenceContrib ScoresHyp_Contrib Scores	Yy1, Taf1, Zfp42
metacluster_1/pattern_16 # seqlets: 85 SequenceContrib ScoresHyp_Contrib Scores	Bach2, Mafk, Junb, Fosb, Bach1, Mafa, Fos, Nfe2l2, Jund, Mxi1, Batf