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