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 = 3
sample_name = 'O_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: 5985 SequenceContrib ScoresHyp_Contrib Scores	Ctcf, Ctcfl
metacluster_1/pattern_1 # seqlets: 2723 SequenceContrib ScoresHyp_Contrib Scores	Sp2, Sp3, Sp4, Sp1, Klf3, Klf6, Wt1, Klf1, Sp5, Zbtb17, Egr1, Zfp281, E2f4, Maz, Klf4, Klf5, Klf7, Klf8, Klf12, Klf15, Zfx, Egr2, E2f1, Ptf1a, E2f7
metacluster_1/pattern_2 # seqlets: 2087 SequenceContrib ScoresHyp_Contrib Scores	Nfib, Nfic, Nfia
metacluster_1/pattern_3 # seqlets: 1929 SequenceContrib ScoresHyp_Contrib Scores	Nrf1
metacluster_1/pattern_4 # seqlets: 1848 SequenceContrib ScoresHyp_Contrib Scores	Jdp2, Creb1, Atf7, Atf2, Creb5, Crem, Mafb, Atf3, Atf1, Fosl1, Nfil3, Fosl2, Tef, Jun, Hlf, Dbp, Cebpg, Gmeb2, Cebpb, Atf4, Tfe3, Jund, Gmeb1, Junb, Cebpa, Cebpd
metacluster_1/pattern_5 # seqlets: 1302 SequenceContrib ScoresHyp_Contrib Scores	Nfyb, Nfyc, Nfya, Foxi1, Pbx3
metacluster_1/pattern_6 # seqlets: 1216 SequenceContrib ScoresHyp_Contrib Scores	Rfx3, Rfx2, Rfx1, Rfx4, Rfx6, Rfx7, Sox2
metacluster_1/pattern_7 # seqlets: 970 SequenceContrib ScoresHyp_Contrib Scores	Sox9, Sox10, Sox3, Sox4, Sox6, Sox17
metacluster_1/pattern_8 # seqlets: 939 SequenceContrib ScoresHyp_Contrib Scores	Ascl2, Ascl1, Nhlh1, Myf6, Tcf3, Myog, Bhlha15, Tcf12, Myod1, Nhlh2, Tal1, Tcf7l1, Atoh1, Tcf4, Snai2
metacluster_1/pattern_9 # seqlets: 812 SequenceContrib ScoresHyp_Contrib Scores	Gabpa, Fli1, Elk1, Ets1, Elk4, Erg, Elf1, Etv4, Elk3, Etv1, Etv5, Etv2, Etv3, Fev, Etv6, XP_911724.4, Elf4, Elf2, Ehf, Elf5, Spi1, Elf3, Spdef
metacluster_1/pattern_10 # seqlets: 806 SequenceContrib ScoresHyp_Contrib Scores	Usf2, Usf1, Tfeb, Tfec, Arntl, Mitf, Srebf1, Bhlhe40, Clock, Npas2, Arnt, Mlx, Creb3l2, Tcfl5, Mycn, Max, Srebf2, Myc, Xbp1, Bhlhe41, Hes1, Mlxip
metacluster_1/pattern_11 # seqlets: 787 SequenceContrib ScoresHyp_Contrib Scores	Pou2f1, Lhx2, Phox2a, Phox2b, Gbx1, Prrx1, Otp, Vsx2, Evx2, Isx, Esx1, En2
metacluster_1/pattern_12 # seqlets: 717 SequenceContrib ScoresHyp_Contrib Scores	Thap11, Zfp143, Tbx2
metacluster_1/pattern_13 # seqlets: 227 SequenceContrib ScoresHyp_Contrib Scores	Irf1, Irf2, Irf7, Irf8, Irf9, Prdm1, Stat1, Stat2, Irf4, Irf3, Irf5
metacluster_1/pattern_14 # seqlets: 172 SequenceContrib ScoresHyp_Contrib Scores	Zbtb33
metacluster_1/pattern_15 # seqlets: 129 SequenceContrib ScoresHyp_Contrib Scores	Taf1, Yy1, Zfp42