from pytadbit.parsers.hic_parser import load_hic_data_from_reads

r_enz_1 = 'HindIII'
r_enz_2 = 'MboI'
reso = 1000000

hic_data_1 = load_hic_data_from_reads(
    'results/{0}/03_filtering/valid_reads12_{0}.tsv'.format(r_enz_1),
    reso)
hic_data_2 = load_hic_data_from_reads(
    'results/{0}/03_filtering/valid_reads12_{0}.tsv'.format(r_enz_2),
    reso)

"""
Explanation: Table of Contents

The HiC_data object
Filter columns with too few interaction count
Normalization
Other normalizations


Save raw and normalized matrices
Save biases and bin filtering
Save normalized chromosome matrices



The HiC_data object
End of explanation
"""


hic_data_1.filter_columns(draw_hist=True, min_count=10, by_mean=True)
hic_data_2.filter_columns(draw_hist=True, min_count=10, by_mean=True)

print len(hic_data_1)
print len(hic_data_1.bads)
print len(hic_data_2)
print len(hic_data_2.bads)

"""
Explanation: Filter columns with too few interaction count
For this, very sparse, example dataset we are going to ask for very few interactions per bin.
This can be done either by setting min_perc asking for each bin to contain a minimum percentage of cells with interaction data. Or but setting min_count asking that the number of cells with interaction data of each bin is above the defined cutoff.
End of explanation
"""


hic_data_1.normalize_hic(iterations=10, max_dev=0.1)
hic_data_2.normalize_hic(iterations=10, max_dev=0.1)

"""
Explanation: Normalization
As normalization we use either the ICE normalization <a name="ref-1"/>(Imakaev et al., 2012) with a "full" normalization until the sum of all columns of the matrix are equal, or something more similar to the vanilla normalization used in <a name="ref-2"/>(Rao et al., 2014) which is exactly like running the ICE normalization without iterations.
Note: if columns with a lot of zeroes are present the ICE normalization will last very long to converge, and these low-coverage columns will present, at the end of the normalization, few cells with very high values of interaction
End of explanation
"""


from pytadbit.mapping.analyze import hic_map

hic_map(hic_data_1, normalized=True, focus='chr18', show=True)
hic_map(hic_data_2, normalized=True, focus='chr18', show=True)

"""
Explanation: At 100 kb itmakes no sense to view the full matrix (your matrix will have more cells than your screen have pixels).
Try plotting a region of the genome only, with the focus parameter.
End of explanation
"""


from cPickle import dump

! mkdir -p results/$r_enz_1/04_normalizing
! mkdir -p results/$r_enz_2/04_normalizing

"""
Explanation: Other normalizations
ICE normalization is widely used however other, more convaluted, normalizations <a name="ref-1"/>(Hu et al., 2012) <a name="ref-2"/>(Yaffe and Tanay, 2011) can be used outside TADbit and then added as normalized matrices.
Save raw and normalized matrices
Save biases and bin filtering
End of explanation
"""


out = open('results/{1}/04_normalizing/biases_{0}_{1}.pick'.format(reso, r_enz_1), 'w')
dump(hic_data_1.bias, out)
out.close()
out = open('results/{1}/04_normalizing/biases_{0}_{1}.pick'.format(reso, r_enz_2), 'w')
dump(hic_data_2.bias, out)
out.close()

"""
Explanation: Save biases to separate file.
End of explanation
"""


out = open('results/{1}/04_normalizing/bad_columns_{0}_{1}.pick'.format(reso, r_enz_1), 'w')
dump(hic_data_1.bads, out)
out.close()
out = open('results/{1}/04_normalizing/bad_columns_{0}_{1}.pick'.format(reso, r_enz_2), 'w')
dump(hic_data_2.bads, out)
out.close()

"""
Explanation: Save "bad" columns to an other file
End of explanation
"""


! dryhic3.r results/HindIII/03_filtering/valid_reads12_HindIII.tsv \
       results/HindIII/04_normalizing/bad_columns_1000000_HindIII.pick \
       HindIII hg38 1000000 results/HindIII/04_normalizing/biases_dryhic_1000000_HindIII.pick

! dryhic3.r results/MboI/03_filtering/valid_reads12_MboI.tsv \
       results/MboI/04_normalizing/bad_columns_1000000_MboI.pick \
       MboI hg38 1000000 results/MboI/04_normalizing/biases_dryhic_1000000_MboI.pick

bias_dry_path = 'results/{1}/04_normalizing/biases_dryhic_{0}_{1}.tsv'

hic_data_1.bias = dict([(int(l.split()[0]), float(l.split()[1])) for l in open(bias_dry_path.format(reso, 'HindIII'))])
hic_data_2.bias = dict([(int(l.split()[0]), float(l.split()[1])) for l in open(bias_dry_path.format(reso, 'MboI'))])

hic_map(hic_data_1, normalized=True, focus='chr18', show=True)
hic_map(hic_data_2, normalized=True, focus='chr18', show=True)

"""
Explanation: Dryhic normalization
We can launch the external script to use dryhic normalization
End of explanation
"""


hic_map(hic_data_1, by_chrom='intra', normalized=False,
       savedata='results/{1}/04_normalizing/{0}_raw'.format(reso, r_enz_1))
hic_map(hic_data_2, by_chrom='intra', normalized=False,
       savedata='results/{1}/04_normalizing/{0}_raw'.format(reso, r_enz_2))

hic_map(hic_data_1, by_chrom='intra', normalized=True,
       savedata='results/{1}/04_normalizing/{0}_norm'.format(reso, r_enz_1))
hic_map(hic_data_2, by_chrom='intra', normalized=True,
       savedata='results/{1}/04_normalizing/{0}_norm'.format(reso, r_enz_2))

"""
Explanation: Save normalized chromosome/genome matrices
This time we do not need to save appart the normalization biases and the list of columns with poor signal.
End of explanation
"""


if reso >= 300000:
    hic_map(hic_data_1, by_chrom=False, normalized=False,
            savedata='results/{1}/04_normalizing/{0}_raw.mat'.format(reso, r_enz_1))

    hic_map(hic_data_1, by_chrom=False, normalized=True,
            savedata='results/{1}/04_normalizing/{0}_norm.mat'.format(reso, r_enz_1))
    hic_map(hic_data_2, by_chrom=False, normalized=False,
            savedata='results/{1}/04_normalizing/{0}_raw.mat'.format(reso, r_enz_2))

    hic_map(hic_data_2, by_chrom=False, normalized=True,
            savedata='results/{1}/04_normalizing/{0}_norm.mat'.format(reso, r_enz_2))    

"""
Explanation: If the resolution is not to low, we could also save genomic matrices:
End of explanation
"""


