from pychemy.amino_acids import AMINO_ACIDS, LOSSES
from pychemy.molmass import Formula
import re


# Retreives first amino acid from peptide sequence based on those available in AMINO_ACIDS
def get_next_aa(sequence):
  keep = []
  for key in AMINO_ACIDS:
    if sequence.startswith(key.mod):
      keep.append(key)
  
  if len(keep) == 1:
    return keep[0]
  elif len(keep) == 0:
    raise Exception('No matching residue in amino acid set')
  else:
    mod_set = [k for k in keep if k.mod != k.residue]
    print len(mod_set)
    if len(mod_set) == 1:
      return mod_set[0]
    else:
      raise Exception('Conflicting modifications in amino acid set')

# Parse string containing sequence into residues
def residues_from_sequence(sequence):
  count = 0
  max_iter = len(sequence)
  residues = []
  while len(sequence) > 0 and count < max_iter + 2:
    count += 1
    aa = get_next_aa(sequence)
    residues.append(aa.mod)
    sequence = sequence.replace(aa.mod, '', 1)
  return residues

# Calculate mass of peptide based on chemical formula entered as 'ACDE'
def mass_from_sequence(sequence = '', nterm = 'N-term', cterm = 'C-term'):
  seq_residues = residues_from_sequence(sequence)
  mass = 0

  # Include N-terminal group mass
  mass += AMINO_ACIDS[nterm].mass

  # Include each residue mass
  for r in seq_residues:
    mass += AMINO_ACIDS[r].mass
      
  # Include C-terminal group mass
  mass += AMINO_ACIDS[cterm].mass

  return mass


def update_loss(loss, new_loss):
  for idx,l in enumerate(loss):
    if l[1].formula == LOSSES[new_loss].formula:
      loss[idx][0] += 1
      return loss
  return loss + [[1, LOSSES[new_loss]]]


def loss_name(loss = []):
  out = ''
  for l in loss:
    if l[0] == 1:
      out += ' - ' + l[1].display_name
    else:
      out += ' - ' + str(l[0]) + ' ' + l[1].display_name
  return out


def predict_positive_fragments_from_sequence(sequence = [], nterm = 'N-term', cterm = 'C-term',
                                             max_charge_state = 1):
  out = []

  fragments = predict_fragments_from_sequence(sequence = sequence, cterm = cterm)

  for item in fragments:
    out.append(( round(item[0] + Formula('H').isotope.mass, 5) , item[1] + ' (+1)'))
    
    for cs in range(2,max_charge_state + 1):
      if item[0] > 200 * cs:
        out.append(( round((item[0] + cs * Formula('H').isotope.mass) / cs, 5), item[1] + ' (+' +  str(cs) + ')'))
  
  out = sorted(list(set(out)), key = lambda x:x[0])
  return out


def predict_negative_fragments_from_sequence(sequence = [], nterm = 'N-term', cterm = 'C-term',
                                             max_charge_state = 1):
  out = []

  fragments = predict_fragments_from_sequence(sequence = sequence, cterm = cterm)

  for item in fragments:
    out.append(( round(item[0] - Formula('H').isotope.mass, 5), item[1] + ' (-1)'))
    
    for cs in range(2,max_charge_state + 1):
      if item[0] > 200 * cs:
        out.append(( round((item[0] - cs * Formula('H').isotope.mass) / cs, 5), item[1] + ' (-' +  str(cs) + ')'))
  
  out = sorted(list(set(out)), key = lambda x:x[0])
  return out


# Predict fragments from peptide sequence    
def predict_fragments_from_sequence(sequence = [], nterm = 'N-term', cterm = 'C-term'):
  out = []

  # b-ions
  nterm_mass = AMINO_ACIDS[nterm].mass
  out += append_residue(mass = nterm_mass, 
                        pos = 1, 
                        ion_type = 'b_', 
                        seq = sequence, 
                        term = cterm, 
                        loss = [])
  for l in AMINO_ACIDS[nterm].losses:
    out += append_residue(mass = nterm_mass - l.mass, 
                          pos = 1, 
                          ion_type = 'b_', 
                          seq = sequence, 
                          term = cterm, 
                          loss = update_loss([],l.formula))

  # a-ions
  nterm_mass = AMINO_ACIDS[nterm].mass - Formula('CO').isotope.mass 
  out += append_residue(mass = nterm_mass, 
                        pos = 1, 
                        ion_type = 'a_', 
                        seq = sequence, 
                        term = cterm, 
                        loss = [])
  for l in AMINO_ACIDS[nterm].losses:
    out += append_residue(mass = nterm_mass - l.mass, 
                          pos = 1, 
                          ion_type = 'a_', 
                          seq = sequence, 
                          term = cterm, 
                          loss = update_loss([],l.formula))

  # y-ions
  cterm_mass = AMINO_ACIDS[cterm].mass
  out += append_residue(mass = cterm_mass, 
                        pos = 1, 
                        ion_type = 'y_', 
                        seq = sequence[::-1], 
                        term = nterm, 
                        loss = [])
  for l in AMINO_ACIDS[cterm].losses:
    out += append_residue(mass = cterm_mass - l.mass, 
                          pos = 1, 
                          ion_type = 'y_', 
                          seq = sequence[::-1], 
                          term = nterm, 
                          loss = update_loss([],l.formula))
      
  out = sorted(list(set(out)), key = lambda x:x[0])
  return out


def append_residue(mass, pos, ion_type = '', seq = [], term = '', loss = []):
  out = []

  if pos == 1 and not seq:
    raise Exception("append_residue encountered empty sequence")

  if len(seq) == 0:
    # No residues left to place, place terminus
    r = AMINO_ACIDS[term]

    if ion_type ==  'a_':
      # Special case: M - CO
      loss = update_loss(loss, 'CO')  
    
    out += [(round(mass + r.mass, 5), 'M' + loss_name(loss))]
    for l in r.losses: 
      out += [(round(mass + r.mass - l.mass, 5), 'M' + loss_name(update_loss(loss, l.formula)))]

  else:
    # Place each loss form of next residue
    r = AMINO_ACIDS[seq[0]]

    out += [(round(mass + r.mass, 5), ion_type + str(pos) + loss_name(loss))]          
    out += append_residue(mass = mass + r.mass, 
                          pos = pos + 1, 
                          ion_type = ion_type, 
                          seq = seq[1:], 
                          term = term, 
                          loss = loss)
    for l in r.losses:
      curr_loss = update_loss(loss, l.formula)
      out += [(round(mass + r.mass - l.mass, 5), ion_type + str(pos) + loss_name(curr_loss))]
      out += append_residue(mass = mass + r.mass - l.mass, 
                            pos = pos + 1, 
                            ion_type = ion_type, 
                            seq = seq[1:], 
                            term = term, 
                            loss = curr_loss)

  return out

@@ -502,31 +502,8 @@
 ue%5D%0A
-    print len(mod_set)%0A
     

Remove print statement