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 + " (differential regions)")))
    
    #homerdir = "{0}/task_{1}".format(homer_root, task_idx)
    homer_ob = DenovoHomer(homerdir)
    homer_ob.fetch_tomtom_matches(save_report=True, 
                                  tomtom_dir= "{0}/{1}_tomtomout".format(homerdir, "CISBP_2.00"))

    
    homer_ob = adjust_motifs(homer_ob)
    homer_ob.load_matched_motifs()
    homer_ob.get_motif_per_celltype()

    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 = 'aNSC_Old'

display_denovo_patterns(
    sample_name,
    modiscodir="{0}/gw_modisco/multi_tasks/{1}".format(root, sample_name)
)

root = "/mnt/lab_data/kundaje/users/msharmin/NSC_ATAC_PEAKS_to_share"
task_idx = 3
sample_name = 'aNSC_Old'
display_paiwise_pattern_comparison(
    task_idx, sample_name, 
    modiscodir="{0}/gw_modisco/multi_tasks/{1}".format(root, sample_name),
    homerdir="{0}/gw_homer/{1}".format(root, sample_name),
)

Pattern Name	TF Name(s)	Modisco
metacluster_1/pattern_0 # seqlets: 412 SequenceContrib ScoresHyp_Contrib Scores	Nfib, Nfic
metacluster_1/pattern_1 # seqlets: 290 SequenceContrib ScoresHyp_Contrib Scores	Ctcf, Ctcfl
metacluster_1/pattern_2 # seqlets: 258 SequenceContrib ScoresHyp_Contrib Scores	Sp2, Sp1, Sp3, Sp4, Klf3, Sp5, Klf6, Klf15, Wt1, Zbtb17, Maz, Klf1, Zfx, Egr1, Zfp281, Egr2, E2f4, Klf7, Klf4, Klf5, Usf2, Klf8, Klf12, Rela, E2f7, E2f1, Mxi1, Ptf1a, Mbd2, Fli1, E2f6
metacluster_1/pattern_3 # seqlets: 130 SequenceContrib ScoresHyp_Contrib Scores	Sox9, Sox3, Sox2, Sox4, Sox10, Sox6, Sox17
metacluster_1/pattern_4 # seqlets: 125 SequenceContrib ScoresHyp_Contrib Scores	Ascl2, Ascl1, Tcf3, Myf6, Myog, Nhlh1, Bhlha15, Myod1, Tcf12, Tcf7l1, Nhlh2, Tal1, Atoh1, Tcf4, Snai2, Myf5, Lyl1, Olig2
metacluster_1/pattern_5 # seqlets: 117 SequenceContrib ScoresHyp_Contrib Scores	Prrx1, Phox2b, Lhx2, Gbx1, Nobox, Phox2a, Pax7, En2, Vsx2, Evx2, Msx2, Otp, Hoxd1, Isx, Emx2, Msx3, En1, Rax, Prrx2, Dlx3, Esx1, Gsx2, Dlx1, Lhx3, Dlx2, Hoxa1, Nkx1-2, Lhx8, Hoxa2, Nkx1-1, Vax2, Gbx2, Hoxd3, Alx3, Lhx4, Vax1, Lhx9, Evx1, Bsx, Shox2, Vsx1, Prop1, Noto, Alx1, Lbx2, Arx, Meox1, Hoxb3, Alx4, Pax6, Hoxa5, Mnx1, Pou2f1, Barx1, Barhl1, Pou6f1, Dlx5, Hoxb5, Msx1, Pou1f1, Lhx6, Uncx, Hoxa3, Barhl2, Hmx2, Dlx4, Nfyb, Pou2f3, Rhox6, Hoxa6, Sebox, Nfyc, Lhx1, Lhx5, Nkx6-3, Nkx2-5, Hoxa7, Nkx2-3, Nkx3-1, Pdx1, Lmx1b, Nkx3-2, Isl1, Pou2f2
metacluster_1/pattern_6 # seqlets: 60 SequenceContrib ScoresHyp_Contrib Scores	Usf1, Tfeb, Tfec, Tfe3, Arntl, Srebf1, Mitf, Bhlhe40, Mlx, Clock, Arnt, Rfx1, Tcfl5, Srebf2, Mycn, Npas2, Jdp2, Myc, Max, Hlf, Atf2, Creb1, Atf7, Creb5, Creb3l2, Bhlhe41, Mlxip, Gabpa, Rxrg, Atf1, Rxrb
metacluster_1/pattern_7 # seqlets: 50 SequenceContrib ScoresHyp_Contrib Scores	Etv4, Erg, Etv1, Elk4, Etv5, Ets1, Elk3, Elk1, Etv3, XP_911724.4, Etv6, Fev, Etv2, Elf1, Elf4, Elf2, Spi1, Ehf, Spdef, Elf5, Spic
metacluster_1/pattern_8 # seqlets: 53 SequenceContrib ScoresHyp_Contrib Scores	Nrf1
metacluster_1/pattern_9 # seqlets: 48 SequenceContrib ScoresHyp_Contrib Scores
metacluster_1/pattern_10 # seqlets: 64 SequenceContrib ScoresHyp_Contrib Scores	Crem, Atf3, Fosl2, Nfil3, Jun, Tef, Fosl1, Gmeb2, Dbp, Cebpg
metacluster_1/pattern_11 # seqlets: 33 SequenceContrib ScoresHyp_Contrib Scores	Tfap2b