import re

from six.moves import xrange

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


# 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]
    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

# Check that a provided sequence is valid
def verify_sequence(sequence):
  assert (type(sequence) is not str), 'Sequence must be an array of amino acids, which can be produced from residues_from_sequence'

  for r in sequence:
    assert (r in [key.mod for key in AMINO_ACIDS]), 'Sequence must contain valid amino acids as defined in AMINO_ACIDS'

  return True

# 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 get_precursors(seq, 
                   min_mz=400, 
                   max_mz=2000, 
                   min_charge=1, 
                   max_charge=5):
  m = mass_from_sequence(seq)
  prec = [{
           'charge_state': c, 
           'mz': (m + c * Formula('H').mass) / c
          } for c in xrange(min_charge, max_charge + 1)]
  return filter(lambda x: min_mz <= x['mz'] <= max_mz, prec)

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