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 = 0
sample_name = 'O_Astrocyte'

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: 3761 SequenceContrib ScoresHyp_Contrib Scores	Ctcf, Ctcfl
metacluster_1/pattern_1 # seqlets: 2098 SequenceContrib ScoresHyp_Contrib Scores	Sp2, Sp1, Sp3, Klf3, Sp4, Sp5, Klf6, Wt1, Maz, Zbtb17, Klf1, Klf15, Egr1, Zfp281, Klf5, E2f4, Klf4, Klf7, Egr2, Klf8, Klf12, Zfx, Rela, E2f1, Fli1, Mxi1, E2f7, E2f3, Ptf1a
metacluster_1/pattern_2 # seqlets: 1954 SequenceContrib ScoresHyp_Contrib Scores	Nfib, Nfic, Nfia
metacluster_1/pattern_3 # seqlets: 1365 SequenceContrib ScoresHyp_Contrib Scores	Nrf1
metacluster_1/pattern_4 # seqlets: 1047 SequenceContrib ScoresHyp_Contrib Scores	Nfyb, Nfyc, Foxi1, Nfya, Pbx3
metacluster_1/pattern_5 # seqlets: 935 SequenceContrib ScoresHyp_Contrib Scores	Gabpa, Ets1, Elk4, Erg, Elk3, Etv4, Elf1, Elk1, Etv1, Etv5, Etv2, Fev, Etv3, Etv6, XP_911724.4, Elf4, Ehf, Elf2, Spi1, Elf5, Spdef, Elf3, Spib
metacluster_1/pattern_6 # seqlets: 753 SequenceContrib ScoresHyp_Contrib Scores	Creb1, Atf7, Jdp2, Creb5, Atf2, Crem, Mafb, Atf1, Atf3, Fosl2, Fosl1, Nfil3, Tef, Tfe3, Jun, Hlf, Dbp
metacluster_1/pattern_7 # seqlets: 700 SequenceContrib ScoresHyp_Contrib Scores	Vsx2, Lhx8, Prrx2, Lhx2, En2, Shox2, Rhox6, Alx1, Rax, Arx, Emx2, Lhx9, Pou1f1, Uncx, Pax6, Lhx3, Lbx2, Vsx1, Lhx4, Noto, Alx4, Hoxa2, Gbx1, Alx3, Sebox, Nkx1-2, Gbx2, Gsx2, Prop1, Msx3, Pax4, Phox2b, Hoxd3, Evx1, Phox2a, Nkx1-1, Vax2, Esx1, Prrx1, Lhx6, Otp, En1, Hoxa1, Vax1, Nobox, Lmx1b, Msx1, Pax7, Pou2f3, Evx2, Pou2f2, Isx, Mnx1, Pdx1, Hoxc5, Hoxd1, Hoxb3, Hoxa3, Pou3f1, Lhx1, Pou3f2, Lhx5, Pou2f1, Msx2, Dlx3, Lmx1a, Barhl1, Hoxa5, Isl1, Barhl2, Hoxc6, Dlx5, Hlx, Dbx1, Dbx2, Hoxb5, Tlx2, Pou3f4, Isl2, Meox2, Pou4f3, Meox1, Pou6f1, Nkx6-3, Bsx, Hnf1b, Nkx6-1, Hoxb4
metacluster_1/pattern_8 # seqlets: 682 SequenceContrib ScoresHyp_Contrib Scores	Rfx2, Rfx3, Rfx1, Rfx4, Rfx6, Rfx7
metacluster_1/pattern_9 # seqlets: 575 SequenceContrib ScoresHyp_Contrib Scores	Irf2, Irf9, Irf1, Prdm1, Stat2, Irf7, Irf8, Irf4
metacluster_1/pattern_10 # seqlets: 541 SequenceContrib ScoresHyp_Contrib Scores	Sox9, Sox3, Sox10, Sox4, Sox2, Sox6, Sox17
metacluster_1/pattern_11 # seqlets: 537 SequenceContrib ScoresHyp_Contrib Scores	Usf2, Usf1, Tfeb, Arntl, Mitf, Tfec, Arnt, Clock, Srebf1, Bhlhe40, Mlx, Npas2, Mycn, Myc, Nr1h4, Max, Srebf2, Rxrg, Tcfl5, Creb3l2, Rxrb
metacluster_1/pattern_12 # seqlets: 498 SequenceContrib ScoresHyp_Contrib Scores	Thap11, Zfp143, Tbx2
metacluster_1/pattern_13 # seqlets: 316 SequenceContrib ScoresHyp_Contrib Scores	Cebpb, Cebpa, Cebpd, Cebpg, Atf4, Ddit3, Cebpe, Pou5f1, Pou3f3
metacluster_1/pattern_14 # seqlets: 201 SequenceContrib ScoresHyp_Contrib Scores	Zbtb33
metacluster_1/pattern_15 # seqlets: 147 SequenceContrib ScoresHyp_Contrib Scores	Yy1, Taf1, Zfp42
metacluster_1/pattern_16 # seqlets: 78 SequenceContrib ScoresHyp_Contrib Scores	Rest, Mafk
metacluster_1/pattern_17 # seqlets: 66 SequenceContrib ScoresHyp_Contrib Scores	Zfp335, Mafa, Bach2
metacluster_1/pattern_18 # seqlets: 63 SequenceContrib ScoresHyp_Contrib Scores	Fosb, Junb, Fos, Bach1, Batf, Jund, Tal1, Pknox1, Nfe2l2
metacluster_1/pattern_19 # seqlets: 38 SequenceContrib ScoresHyp_Contrib Scores
metacluster_1/pattern_20 # seqlets: 33 SequenceContrib ScoresHyp_Contrib Scores