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 = 5
sample_name = 'Y_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: 4270 SequenceContrib ScoresHyp_Contrib Scores	Ctcf, Ctcfl
metacluster_1/pattern_1 # seqlets: 2038 SequenceContrib ScoresHyp_Contrib Scores	Sp2, Sp3, Sp4, Klf3, Sp1, Klf6, Klf15, Zfx, Wt1, Maz, Klf1, Sp5, Zbtb17, Egr1, Zfp281, E2f4, Klf5, Klf4, Klf7, Egr2, Klf8, Klf12, Tfap2c, E2f1, Usf2, Rela, Nr1h4, E2f7, Ptf1a, Egr3
metacluster_1/pattern_2 # seqlets: 1793 SequenceContrib ScoresHyp_Contrib Scores	Nfib, Nfic
metacluster_1/pattern_3 # seqlets: 1490 SequenceContrib ScoresHyp_Contrib Scores	Nrf1
metacluster_1/pattern_4 # seqlets: 1145 SequenceContrib ScoresHyp_Contrib Scores	Jdp2, Creb1, Atf7, Atf2, Creb5, Mafb, Crem, Atf1, Atf3, Tef, Nfil3, Fosl1, Fosl2, Jun, Hlf, Dbp, Cebpg, Gmeb2, Tfe3, Cebpb, Rorb, Gmeb1, Atf4, Junb, Cebpa, Jund, Cebpd
metacluster_1/pattern_5 # seqlets: 1078 SequenceContrib ScoresHyp_Contrib Scores	Nfyb, Nfyc, Nfya, Foxi1, Pbx3
metacluster_1/pattern_6 # seqlets: 900 SequenceContrib ScoresHyp_Contrib Scores	Fli1, Gabpa, Elk4, Ets1, Elk1, Erg, Elf1, Etv4, Elk3, Etv1, Etv2, Etv5, Fev, Etv3, Etv6, XP_911724.4, Ehf, Elf4, Elf3, Elf5, Elf2, Spi1, Spdef
metacluster_1/pattern_7 # seqlets: 673 SequenceContrib ScoresHyp_Contrib Scores	En2, Lhx2, Pax6, Vsx2, Emx2, Sebox, Uncx, Phox2b, Prrx2, Alx1, Evx1, Evx2, Otp, Alx3, Esx1, Prrx1, Phox2a, Lhx8, Rhox6, Lhx3, Gbx1, Lhx9, Hoxa2, Gsx2, Pou1f1, Shox2, Prop1, Pou2f2, Lhx4, Gbx2, Vax2, Arx, En1, Vax1, Vsx1, Pdx1, Isx, Lbx2, Hoxd3, Msx1, Rax, Pou2f1, Alx4, Hoxc5, Noto, Hoxa3, Pou3f1, Nkx1-2, Hoxb3, Hoxa4, Hoxa1, Nkx1-1, Msx3, Nobox, Mnx1, Lmx1b, Msx2, Pax4, Pou3f2, Hoxb4, Hoxc6, Lhx6, Hoxd1, Pou6f1, Hoxa5, Meox2, Lhx5, Dlx3, Hoxb5, Hoxc4, Lhx1, Pou2f3, Pax7, Dlx1, Pou4f3, Hlx, Meox1, Nkx6-3, Nkx6-1, Isl1, Pou3f4, Isl2, Dbx1, Dlx4, Dlx5, Dlx2, Hoxc9, Hoxa7, Barhl1, Barx2, Barx1, Hoxc10, Lmx1a, Arid3b, Tlx2, Hoxa6, Bsx, Dbx2
metacluster_1/pattern_8 # seqlets: 661 SequenceContrib ScoresHyp_Contrib Scores	Rfx3, Rfx2, Rfx1, Rfx4, Rfx6, Rfx7
metacluster_1/pattern_9 # seqlets: 655 SequenceContrib ScoresHyp_Contrib Scores	Usf1, Tfeb, Mitf, Arntl, Tfec, Clock, Mlx, Arnt, Srebf1, Bhlhe40, Mycn, Max, Myc, Npas2, Creb3l2, Tcfl5, Srebf2, Nr2f6, Mlxip, Xbp1
metacluster_1/pattern_10 # seqlets: 559 SequenceContrib ScoresHyp_Contrib Scores	Thap11, Zfp143, Tbx2
metacluster_1/pattern_11 # seqlets: 433 SequenceContrib ScoresHyp_Contrib Scores	Sox9, Sox10, Sox3, Sox2, Sox4, Sox6, Sox17
metacluster_1/pattern_12 # seqlets: 219 SequenceContrib ScoresHyp_Contrib Scores	Taf1, Yy1, Zfp42
metacluster_1/pattern_13 # seqlets: 144 SequenceContrib ScoresHyp_Contrib Scores	Zbtb33
metacluster_1/pattern_14 # seqlets: 106 SequenceContrib ScoresHyp_Contrib Scores	Mafk, Fos, Bach2, Fosb, Mafa, Nhlh1, Bach1, Nfe2l2, Batf, Zfp335, Mafg, Tal1, Mxi1
metacluster_1/pattern_15 # seqlets: 85 SequenceContrib ScoresHyp_Contrib Scores
metacluster_1/pattern_16 # seqlets: 77 SequenceContrib ScoresHyp_Contrib Scores
metacluster_1/pattern_17 # seqlets: 70 SequenceContrib ScoresHyp_Contrib Scores	Rest
metacluster_1/pattern_18 # seqlets: 45 SequenceContrib ScoresHyp_Contrib Scores