### import two datasets
def reindex_xcms_by_mzrt(df):
    df.index = (df.loc[:,'mz'].astype('str') + 
               ':' + df.loc[:, 'rt'].astype('str'))
    return df


# alzheimers
local_path = '/home/irockafe/Dropbox (MIT)/Alm_Lab/'\
'projects'
alzheimers_path = local_path + '/revo_healthcare/data/processed/MTBLS72/positive_mode/'\
'mtbls_no_retcor_bw2.csv'
## Import the data and remove extraneous columns
df_alzheimers = pd.read_csv(alzheimers_path, index_col=0)
df_alzheimers = reindex_xcms_by_mzrt(df_alzheimers)
# malaria
malaria_path = local_path + ('/revo_healthcare/data/processed/MTBLS315/'+ 
            'uhplc_pos/xcms_result_4.csv')
df_malaria = pd.read_csv(malaria_path, index_col=0)
df_malaria = reindex_xcms_by_mzrt(df_malaria)
    

ppm_alz_v_malaria = ppm_matrix(df_malaria['mz'], 
                                   df_alzheimers['mz'])

rt_alz_v_malaria = pairwise_difference(df_malaria['rt'], 
                                          df_alzheimers['rt'])


"""
Explanation: <h2> Let's import a couple datasets and take them for a spin</h2>
End of explanation
"""


sns.heatmap(np.log10(ppm_alz_v_malaria))
plt.title('Log10 ppm difference')
plt.show()

# How many for differences at 30ppm?
ppm_window = 30
within_ppm = (ppm_alz_v_malaria[ppm_alz_v_malaria < 30]
                  .dropna(axis=0, how='all')
                  .dropna(axis=1, how='all')
             )
print 'shape', ppm_alz_v_malaria.shape
print ('ppm within {ppm} ppm: '.format(ppm=ppm_window) +
       '{num}'.format(num=(ppm_alz_v_malaria < 30).sum().sum()))

# Get indexes
print 'shape of htose within 30ppm:, ', within_ppm.shape


# How many m/z from one dataset could be m/z isomers from
# other dataset?
print ('\n\nMass matches between datasets (isomers and 1:1 matches)',
       (within_ppm < 30).sum().sum())

print '\nAlzheimers "isomers" in other dataset that are match >1 feature in other set', ((within_ppm < 30).sum(axis=0)>1).sum()
print 'Alzheimers total', df_alzheimers['rt'].shape


print '\n\nMalaria "isomers in other dataset that match >1 feature in other set', ((within_ppm < 30).sum(axis=1) > 1).sum()
print 'Malaria total', df_malaria['rt'].shape



# Show distribution of # of isomers per feature in both malaria and fever datasets
print (within_ppm < 30).sum(axis=0).hist(bins=30)
plt.title('Alzheimers isomers in malaria dataset')
plt.show()

(within_ppm < 30).sum(axis=1).hist(bins=30)
plt.title('Malaria isomers in alzheimers dataset')
plt.show()

"""
Explanation: <h2> Looks like There aren't too many ppm m/z overlaps </h2>
End of explanation
"""


local = '/home/irockafe/Dropbox (MIT)/Alm_Lab/projects/'
xml_file = local + 'revo_healthcare/data/external/toy_database.xml'
xml_file = local + 'revo_healthcare/data/external/serum_metabolites.xml'
#xml_tree = etree.iterparse(xml_file, tag='metabolite')
#



# namespace - at the top of file. fucks with every tag.
# very annoying, so name all tags ns + tag
ns = '{http://www.hmdb.ca}'
nsmap = {None : ns}
# If you're within a metabolite tag
count = 0
seen_mass = 0
d = {}
for event, element in etree.iterparse(xml_file, tag=ns+'metabolite'):
    tree = etree.ElementTree(element)
    # Aggregate info into a dictionary of 
    # {HMDB_ID: iso_mass}
    accession = []
    # Get accession number and masses for each metabolite
    # Could be multiple accessions. Grab all of them, 
    # sort to make unique identifier
    for elem in tree.iter():
        if elem.tag == ns+'accession':
            accession.append(elem.text)
        # If you just saw a 'mono_mass' entry,
        # get the mass value and reset, saying you
        # havent seen 'mono_mass' in the text of next metabolite
        if (elem.tag == ns+'value') & (seen_mass == 1): 
                mass = float(elem.text)
                seen_mass = 0
        if elem.text == 'mono_mass':
                seen_mass = 1
    elem.clear()
    
    # sort accession numbers and join with '_'    
    accession_key = '_'.join(sorted(accession))
    # add to dictionary
    if mass:
        d[accession_key] = mass
        
    # reset mass - only add feature if mass listed
    mass = None
    # reset accession numbers
    accession = []

    element.clear()
    count += 1
    if count % 1000 == 0:
        print('Made it through ' + str(count) + ' metabolites')

#pickle.dump(d, open('serumdb_dict.p', 'wb'))
print 'Number of metabolites: %s' % len(d.keys())

serumdb_masses = pd.Series(d, dtype='float32')

serumdb_ppm_matrix = ppm_matrix(serumdb_masses, serumdb_masses)*10**6

#df = pd.DataFrame(serumdb_ppm_matrix, index=serumdb_masses.index,
#            columns=serumdb_masses.index)*10**6
# Forget about using a dataframe - uses too much memory

"""
Explanation: <h2> So, about 1/4 of the mass-matches have potential isomers in the other dataset...? </h2>
Notice how there are more matches to the malaria set, which has more peaks. Makes sense - more peaks either means more actual molecules, 
or more adducts that could be mistakenly matched as molecules
<h2> Get masses of all hmdb serum metabolites </h2>
parse the xml file
End of explanation
"""


top_ppm = 30
pairs = np.full((top_ppm), np.nan)
print(pairs)
for i in range(1,top_ppm):
    # div by two, b/c half matrix is redundant
    # subtract length of diagonal of matrix, too
    num = ((serumdb_ppm_matrix < i).sum() / 2) - serumdb_ppm_matrix.shape[0]
    pairs[i] = num

plt.scatter(x=range(1,30), y=pairs[1:])
plt.title('Number of pairs of molecules that could overlap in human serum database\n')
plt.show()

"""
Explanation: <h2> So, we've got 6,315,000 pairs of molecules that could be isomers at 1 ppm </h2>
That's about 10% of possible pairs from 25,000 molecules
End of explanation
"""


# how to plot the number of overlaps per molecule?
num_below_1ppm = (serumdb_ppm_matrix < 1).sum(axis=1) - 1
plt.hist((serumdb_ppm_matrix < 1).sum(axis=1) - 1 )
plt.title('Pairs of overlapping mz at ppm 1')
plt.show()

num_below_1ppm

"""
Explanation: <h2> Looks like there are more isomers than 1:1 pairings, by a lot </h2>
Less than 6000 of
End of explanation
"""


