# -*- coding: utf-8 -*-

from __future__ import print_function
from pprint import pprint


"""
Routines to calculate mass-spectra of peptide sequences.
"""


atom_mass = {
  'H': 1.007825032,
  'H+': 1.007276,
  'O': 15.994914622,
}

ave_atom_mass = {
  'H': 1.00794,
  'O': 15.9994,
}

aa_monoisotopic_mass = {
  'A': 71.037114,
  'B': 110.000000,
  'C': 103.009184,
  'D': 115.026943,
  'E': 129.042593,
  'F': 147.068414,
  'G': 57.021464,
  'H': 137.058912,
  'I': 113.084064,
  'J': 110.000000,
  'K': 128.094963,
  'L': 113.084064,
  'M': 131.040485,
  'N': 114.042927,
  'O': 110.000000,
  'P': 97.052764,
  'Q': 128.058578,
  'R': 156.101111,
  'S': 87.032028,
  'T': 101.047678,
  'U': 110.000000,
  'V': 99.068414,
  'W': 186.079313,
  'X': 110.000000,
  'Y': 163.063329,
  'Z': 110.000000
}

aa_average_mass = {
  'A': 71.0788,
  'B': 110.0000,
  'C': 103.1388,
  'D': 115.0886,
  'E': 129.1155,
  'F': 147.1766,
  'G': 57.0519,
  'H': 137.1411,
  'I': 113.1594,
  'J': 110,
  'K': 128.1741,
  'L': 113.1594,
  'M': 131.1926,
  'N': 114.1038,
  'O': 110,
  'P': 97.1167,
  'Q': 128.1307,
  'R': 156.1875,
  'S': 87.0782,
  'T': 101.1051,
  'U': 110,
  'V': 99.1326,
  'W': 186.2132,
  'X': 110,
  'Y': 163.1760,
  'Z': 110
}

# the atoms here artifically suppress a H+ atom so that
# the algorithms to incoporate the charge can be calculated
# as m/z = (mass+H+*charge)/charge
ion_type_props = {
  'b': { 
    # neutral: nterm (NH)-, electron is lost to form cterm CO(+) 
    # charged: gain a H to neutralize nterm NH2, cterm triple CO(+) bond
    'atoms_to_add': [], 
    'atoms_to_sub': [],
    'is_nterm': True, 
  },
  'y': { 
    # netural: nterm NH2, neutral cterm COOH
    # charged: gain a H+ on the nterm to form NH3(+)
    'atoms_to_add': ['H', 'H', 'O'], 
    'atoms_to_sub': [],
    'is_nterm': False,
  },
}


def parse_aa_masses(sequence, aa_mass, modified_masses=[]):
  """
  Returns a list of masses for a given sequence.

  aa_mass is a dictionary of amino-acid masses.
  modified_masses is a dictionary of positions and modified masses,
  """
  masses = [aa_mass[aa] for aa in sequence]
  for modified in modified_masses:
    masses[modified['i']] = modified['mass']
  return masses


def ion_label(i, charge, ion_type):
  if charge == 1:
    return '%s%d' % (ion_type, i)
  return '%s%d(%d+)' % (ion_type, i, charge)


def calculate_neutral_mass(aa_masses, ion_type):
  props = ion_type_props[ion_type]
  mass = sum(aa_masses) 
  mass += sum([atom_mass[a] for a in props['atoms_to_add']])
  mass -= sum([atom_mass[a] for a in props['atoms_to_sub']])
  return mass


def calculate_mz(aa_masses, charge, ion_type):
  charged_mass = calculate_neutral_mass(aa_masses, ion_type)
  charged_mass += atom_mass['H+']*charge 
  if charge == 1:
    return charged_mass
  else:
    return charged_mass/charge


def calculate_peaks(aa_masses, charge, ion_type):
  """
  Returns a list of theoretical mz peaks [mz, label]
  """
  peaks = []
  for i in range(len(aa_masses)):
    if ion_type_props[ion_type]['is_nterm']:
      fragment = aa_masses[:i+1]
    else:
      fragment = aa_masses[i:]
    mz = calculate_mz(fragment, charge, ion_type)
    # round to 4 decimal places
    mz = int(mz*1E4)/1E4
    label = ion_label(len(fragment), charge, ion_type)
    peaks.append((mz, label))
  return peaks


def map_matched_ions(
    ion_type, sequence, peaks, mz_delta_max=0.8, modified_aa_masses=[],
    aa_mass=aa_monoisotopic_mass):
  """
  Returns a list of matched peaks between sequence and mz peaks.
  A peak is [mz, intensity, label, mz_delta, mz_error].

  mz_error = mz_delta/mz_theory*1E6
  peaks is a list of [mz, intensity].
  aa_mass is a dictionary of masses for amino acids
  """
  aa_masses = parse_aa_masses(sequence, aa_mass, modified_aa_masses)
  ion = ion_type[0]
  pieces = ion_type.split('(')
  charge = int(pieces[1][0]) if len(pieces) > 1 else 1
  theory_peaks = calculate_peaks(aa_masses, charge, ion)
  matched = []
  for mz_theory, label in theory_peaks:
    for mz, intensity in peaks:
      mz_delta = mz_theory - mz
      mz_error = int(mz_delta/mz_theory*1E6)
      if abs(mz_delta) <= mz_delta_max:
        matched.append([mz, intensity, label, mz_delta, mz_error])
        break
  return matched


if __name__ == "__main__":
  seq = 'MSAFLLTKR'
  spectrum = [[622.494, 2192.0],
   [615.666, 1811.0],
   [619.456, 1810.0],
   [628.245, 1772.0],
   [616.511, 1016.0],
   [614.583, 780.0],
   [629.003, 673.0],
   [490.215, 563.0],
   [606.378, 533.0],
   [563.797, 430.0],
   [610.871, 421.0],
   [472.861, 376.0],
   [607.103, 346.0],
   [584.56, 341.0],
   [605.67, 334.0],
   [620.333, 310.0],
   [375.222, 299.0],
   [625.311, 263.0],
   [488.298, 254.0],
   [786.349, 223.0],
   [899.377, 221.0],
   [506.51, 215.0],
   [881.571, 215.0],
   [601.8, 207.0],
   [514.999, 202.0],
   [598.419, 201.0],
   [608.138, 199.0],
   [593.698, 196.0],
   [759.359, 196.0],
   [530.555, 186.0],
   [623.222, 186.0],
   [602.825, 184.0],
   [542.091, 180.0],
   [261.063, 169.0],
   [627.002, 160.0],
   [913.689, 160.0],
   [592.707, 157.0],
   [825.581, 157.0],
   [471.826, 156.0],
   [505.835, 155.0],
   [609.168, 154.0],
   [768.284, 151.0],
   [563.027, 146.0],
   [570.712, 146.0],
   [580.803, 145.0],
   [979.473, 137.0],
   [596.779, 130.0],
   [541.113, 129.0],
   [777.839, 186.0],
   [550.083, 159.0]];
  delta_mass = 0.5
  modified_aa_masses = [{ 'i':3, 'mass':100.0 }]
  modified_aa_masses = []
  aa_masses = parse_aa_masses(seq, aa_monoisotopic_mass, modified_aa_masses)
  pprint(aa_masses)
  bion_peaks = calculate_peaks(aa_masses, 1, 'b')
  yion_peaks = calculate_peaks(aa_masses, 1, 'y')
  pprint(yion_peaks)
  peaks = []
  matched = map_matched_ions('y', seq, spectrum, delta_mass, modified_aa_masses, aa_monoisotopic_mass)
  matched = map_matched_ions('b', seq, spectrum, delta_mass, modified_aa_masses, aa_monoisotopic_mass)
  pprint(matched);