# Copyright (c) 2013, Bosco Ho. All rights reserved.

# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:

# 1. Redistributions of source code must retain the above copyright notice,
# this list of conditions and the following disclaimer.

# 2. Redistributions in binary form must reproduce the above copyright notice,
# this list of conditions and the following disclaimer in the documentation
# and/or other materials provided with the distribution.

# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
# ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
# LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
# CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
# SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
# INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
# CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
# ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
# POSSIBILITY OF SUCH DAMAGE.

import re
import sys
import os
import textwrap
import time
import json
import shutil
from copy import deepcopy
from functools import cmp_to_key

import httpx



"""
This is my uniprot python library. It provides a bunch of
routines to access the uniprot.org website's RESTful API.

Principally, you want to do two things:
1) Map different types of sequence identifiers
2) Fetch uniprot metadata for certain sequence id's.

Given a set of seqid's, this library allows a programmic
approach to transforming id's and if the uniprot_id is
obtained, it allows the extraction of the relevant metadata,
including the protein sequence.
"""


def logging(s):
  sys.stdout.write(s)


def is_html(text):
  if re.search('<html', text):
    return True
  return False


def get_uniprot_id_mapping_pairs(
    from_type, to_type, seqids, cache_fname=None, session=None):
  """
  Returns a list of matched pairs of identifiers using the new UniProt async API.

  from_type and to_type can be obtained from:
    https://www.uniprot.org/help/api_idmapping
  """
  if cache_fname and os.path.isfile(cache_fname):
    logging("Loading cached (%s->%s) mappings from %s\n" % (from_type, to_type, cache_fname))
    text = open(cache_fname).read()
    pairs = [l.split('\t') for l in text.strip().split('\n') if l]
    return pairs
  
  if session is None:
    client = httpx.Client(timeout=30.0)
    close_client = True
  else:
    client = session
    close_client = False
  
  try:
    try:
      # Submit mapping job
      job_response = client.post(
          'https://rest.uniprot.org/idmapping/run',
          data={
            'from': from_type,
            'to': to_type,
            'ids': ' '.join(seqids)
          })
    except Exception as e:
      logging("Error: Failed to submit mapping job: %s\n" % str(e))
      return []
    
    if job_response.status_code != 200:
      logging("Error: %s->%s mapping failed (%d): %s\n" % (from_type, to_type, job_response.status_code, job_response.text[:100]))
      return []
    
    job_id = job_response.json()['jobId']
    
    # Poll for job completion
    max_attempts = 60
    for attempt in range(max_attempts):
      try:
        status_response = client.get('https://rest.uniprot.org/idmapping/status/%s' % job_id)
      except Exception as e:
        logging("Error: Failed to check job status: %s\n" % str(e))
        return []
      
      if status_response.status_code not in [200, 303]:
        logging("Error: Failed to check job status (HTTP %d)\n" % status_response.status_code)
        return []
      
      status = status_response.json()
      if status.get('jobStatus') == 'FINISHED':
        break
      elif status.get('jobStatus') == 'FAILURE':
        logging("Error: ID mapping job failed: %s\n" % status.get('failureReason', 'Unknown reason'))
        return []
      
      if attempt < max_attempts - 1:
        time.sleep(1)
    else:
      logging("Error: ID mapping job timed out after %d attempts\n" % max_attempts)
      return []
    
    # Fetch results
    try:
      results_response = client.get('https://rest.uniprot.org/idmapping/stream/%s' % job_id)
    except Exception as e:
      logging("Error: Failed to fetch mapping results: %s\n" % str(e))
      return []
    
    if results_response.status_code != 200:
      logging("Error: Failed to fetch mapping results (HTTP %d): %s\n" % (results_response.status_code, results_response.text[:100]))
      return []
    
    text = results_response.text
    if cache_fname:
      with open(cache_fname, 'w') as f:
        f.write(text)
    
    # Parse results (JSON format with "results" array)
    pairs = []
    try:
      data = json.loads(text)
      if 'results' in data:
        for item in data['results']:
          if 'from' in item and 'to' in item:
            pairs.append([item['from'], item['to']])
    except (json.JSONDecodeError, KeyError):
      # Try legacy tab-delimited format as fallback
      for line in text.strip().split('\n'):
        if line and not line.startswith('From'):
          parts = line.split('\t')
          if len(parts) >= 2:
            pairs.append([parts[0], parts[1]])
    
    return pairs
  except Exception as e:
    logging("Error: Unexpected error in ID mapping: %s\n" % str(e))
    return []
  finally:
    if close_client:
      client.close()


def batch_uniprot_id_mapping_pairs(
    from_type, to_type, seqids, batch_size=400, cache_dir=None):
  """
  Returns a list of matched pairs of identifiers.
  Converts identifiers using above function 'get_uniprot_id_mapping_pairs'
  but launches url requests in safe batches of 100 seqids at a time.

  from_type and to_type can be obtained from:
    http://www.uniprot.org/faq/28#mapping-faq-table
  """
  if cache_dir:
    seqids_txt = os.path.join(cache_dir, 'seqids.json')
    if os.path.isfile(seqids_txt):
        saved_seqids = json.load(open(seqids_txt))
        if saved_seqids != seqids:
            shutil.rmtree(cache_dir)

    if not os.path.isdir(cache_dir):
      os.makedirs(cache_dir)

    if not os.path.isfile(seqids_txt):
      json.dump(seqids, open(seqids_txt, 'w'))

  pairs = []
  i_seqid = 0
  if batch_size is None:
    batch_size = len(seqids)
  with httpx.Client(timeout=30.0) as session:
    while i_seqid <= len(seqids):
      seqids_subset = seqids[i_seqid:i_seqid+batch_size]
      if cache_dir:
        subset_cache = os.path.join(cache_dir, 'mapping.%d.txt' % i_seqid)
      else:
        subset_cache = None
      subset_pairs = get_uniprot_id_mapping_pairs(
          from_type, to_type, seqids_subset, cache_fname=subset_cache, session=session)
      pairs.extend(subset_pairs)
      i_seqid += batch_size
  return pairs


def parse_isoforms(text):
  """
  Returns a dictionary of uniprot_acc and entries.

  Each entry contains:
    - var_seq: the sequence variations in the file
    - isoforms: the interpreted isoform sequences with
                the seqids
  """
  tag = None
  uniprot_data = {}
  var_seq = None
  in_isoform_section = False
  isoform_id = None
  for l in text.splitlines():
    test_tag = l[:5].strip()
    if test_tag and test_tag != tag:
      tag = test_tag
    line = l[5:].strip()
    words = line.split()
    if tag == "ID":
      uniprot_id = words[0]
      uniprot_data[uniprot_id] = {
        'var_seqs': [],
        'isoforms': {},
        'sequence': '',
      }
    if tag == "SQ":
      if words[0] != "SEQUENCE":
        uniprot_data[uniprot_id]['sequence'] += ''.join(words)
    if tag == 'FT':
      if var_seq is not None and l[5] != ' ':
        var_seq = None
      if line.startswith('VAR_SEQ'):
        range_str = words[1]
        if '..' in range_str:
          parts = range_str.split('..')
          var_seq = {
            'i': int(parts[0]),
            'j': int(parts[1]),
            'block': ''
          }
        elif len(words) >= 3 and words[2].isdigit():
          var_seq = {
            'i': int(words[1]),
            'j': int(words[2]),
            'block': ''
          }
        else:
          pos = int(words[1])
          var_seq = {
            'i': pos,
            'j': pos,
            'block': ''
          }
        uniprot_data[uniprot_id]['var_seqs'].append(var_seq)
      if var_seq is not None:
        var_seq['block'] += l[34:]
        if l.endswith('isoform'):
          # needed later to search isoform references
          var_seq['block'] += ' '
    if tag == 'CC':
      if words[0] == '-!-':
        if 'ALTERNATIVE' in words[1] and 'PRODUCTS' in words[2]:
          in_isoform_section = True
        else:
          in_isoform_section = False
      if in_isoform_section:
        if words[0].startswith('Name'):
          isoform_id = str(words[0][:-1].split('=')[1])
        for word in words:
          if word.startswith('IsoId='):
            seqid = word[:-1].split('=')[1]
            uniprot_data[uniprot_id]['isoforms'][isoform_id] = {
              'seqid': seqid
            }
            break
  for uniprot_id in uniprot_data:
    var_seqs = uniprot_data[uniprot_id]['var_seqs']
    isoforms = uniprot_data[uniprot_id]['isoforms']
    original_sequence = uniprot_data[uniprot_id]['sequence']
    for var_seq in var_seqs:
      block = var_seq['block']
      match = re.search(r'\(.*\)', block)
      isoform_ids = []
      if match:
        isoform_tokens = match.group()[1:-1].split()
        for i in range(len(isoform_tokens)):
          if 'isoform' in isoform_tokens[i]:
            isoform_ids.append(str(isoform_tokens[i+1]))
      var_seq['isoform_ids'] = isoform_ids
      if block.startswith('Missing'):
        var_seq['deletion'] = True
      else:
        var_seq['deletion'] = False
        transition = block.split('(')[0]
        if '->' in transition:
          original, mutation = transition.split('->')
          var_seq['sequence'] = original.strip()
          var_seq['mutation'] = mutation.strip()
        else:
          var_seq['sequence'] = ''
          var_seq['mutation'] = transition.strip()
    var_seqs.sort(key=lambda v:-v['i'])
    for isoform_id in isoforms:
      sequence = original_sequence
      for var_seq in uniprot_data[uniprot_id]['var_seqs']:
        if isoform_id in var_seq['isoform_ids']:
          i = var_seq['i']-1
          j = var_seq['j']
          if var_seq['deletion']:
            sequence = sequence[:i] + sequence[j:]
          else:
            sequence = sequence[:i] + var_seq['mutation'] + sequence[j:]
      isoforms[isoform_id]['sequence'] = sequence
  return uniprot_data


def parse_uniprot_txt_file(cache_txt):
  """
  Parses the text of metadata retrieved from uniprot.org.

  Only a few fields have been parsed, but this provides a
  template for the other fields.

  A single description is generated from joining alternative
  descriptions.

  Returns a dictionary with the main UNIPROT ACC as keys.
  """

  tag = None
  uniprot_id = None
  metadata_by_seqid = {}
  for l in cache_txt.splitlines():
    test_tag = l[:5].strip()
    if test_tag and test_tag != tag:
      tag = test_tag
    line = l[5:].strip()
    words = line.split()
    if tag == "ID":
      uniprot_id = words[0]
      is_reviewed = words[1].startswith('Reviewed')
      length = int(words[2])
      metadata_by_seqid[uniprot_id] = {
        'id': uniprot_id,
        'is_reviewed': is_reviewed,
        'length': length,
        'sequence': '',
        'accs': [],
      }
      entry = metadata_by_seqid[uniprot_id]
    if tag == "SQ":
      if words[0] != "SEQUENCE":
        entry['sequence'] += ''.join(words)
    if tag == "AC":
      accs = [w.replace(";", "") for w in words]
      entry['accs'].extend(accs)
    if tag == "DR":
      if 'PDB' in words[0]:
        if 'pdb' not in entry:
          entry['pdb'] = words[1][:-1]
        if 'pdbs' not in entry:
          entry['pdbs'] = []
        entry['pdbs'].append(words[1][:-1])
      if 'RefSeq' in words[0]:
        if 'refseq' not in entry:
          entry['refseq'] = []
        ids = [w[:-1] for w in words[1:]]
        entry['refseq'].extend(ids)
      if 'KEGG' in words[0]:
        if 'kegg' not in entry:
          entry['kegg'] = []
        ids = [w[:-1] for w in words[1:]]
        ids = list(filter(lambda w: len(w) > 1, ids))
        entry['kegg'].extend(ids)
      if 'GO' in words[0]:
        if 'go' not in entry:
          entry['go'] = []
        entry['go'].append(' '.join(words[1:]))
      if 'Pfam' in words[0]:
        if 'pfam' not in entry:
          entry['pfam'] = []
        entry['pfam'].append(words[1][:-1])
    if tag == "GN":
      if 'gene' not in entry and len(words) > 0:
        pieces = words[0].split("=")
        if len(pieces) > 1 and 'name' in pieces[0].lower():
          entry['gene'] = pieces[1].replace(';', '').replace(',', '')
    if tag == "OS":
      if 'organism' not in entry:
        entry['organism'] = ""
      entry['organism'] += line
    if tag == "DE":
      if 'descriptions' not in entry:
        entry['descriptions'] = []
      entry['descriptions'].append(line)
    if tag == "CC":
      if 'comment' not in entry:
        entry['comment'] = ''
      entry['comment'] += line + '\n'

  for entry in metadata_by_seqid.values():
    descriptions = entry['descriptions']
    for i in reversed(range(len(descriptions))):
      description = descriptions[i]
      if 'Short' in description or 'Full' in description:
        j = description.find('=')
        descriptions[i] = description[j+1:].replace(';', '')
      else:
        del descriptions[i]
    entry['description'] = '; '.join(descriptions)

  return metadata_by_seqid


def parse_uniprot_metadata_with_seqids(seqids, cache_txt):
  """
  Returns a dictionary of metadata of given seqids, doing the requisite
  lookup of seqids in the ACCs and IDs, and more importantly handles
  isoform seqid's
  """
  metadata = parse_uniprot_txt_file(cache_txt)
  tmp = metadata.copy()
  for uniprot_id in metadata.keys():
    for seqid in metadata[uniprot_id]['accs'] + [metadata[uniprot_id]['id']]:
      tmp[seqid] = metadata[uniprot_id]
  metadata = tmp
  results = {}
  isoform_dict = parse_isoforms(cache_txt)
  for seqid in seqids:
    if seqid in metadata:
      results[seqid] = metadata[seqid]
    else:
      primary_seqid = seqid[:6]
      if primary_seqid in metadata:
        protein_metadata = metadata[primary_seqid]
        uniprot_id = protein_metadata['id']
        if uniprot_id in isoform_dict:
          isoforms = isoform_dict[uniprot_id]['isoforms'].values()
          for isoform in isoforms:
            if isoform['seqid'] == seqid:
              results[seqid] = deepcopy(protein_metadata)
              results[seqid]['accs'] = [seqid]
              results[seqid]['sequence'] = isoform['sequence']
  return results


def fetch_uniprot_metadata(seqids, cache_fname=None):
  """
  Returns a dictonary of the uniprot metadata (as parsed
  by parse_uniprot_txt_file) of the given seqids. The seqids
  must be valid uniprot identifiers.

  Now handles isoform versions of accession id's!
  """

  primary_seqids = [s[:6] for s in seqids]
  if cache_fname and os.path.isfile(cache_fname):
    logging("Loading cached metadata from " + cache_fname + "\n")
    cache_txt = open(cache_fname).read()
  else:
    logging("Fetching metadata for %d Uniprot IDs from UniProt ...\n" % len(primary_seqids))
    client = httpx.Client(timeout=30.0)
    try:
      r = client.post(
          'https://rest.uniprot.org/uniprotkb/search',
          params={
            'query': ' OR '.join(['accession:%s' % s for s in primary_seqids]),
            'format': 'txt'
          })
      
      if r.status_code == 200:
        cache_txt = r.text
      else:
        logging("Error fetching metadata: HTTP %d\n" % r.status_code)
        return {}
      
      if cache_fname:
        with open(cache_fname, 'w') as f:
          f.write(cache_txt)
      
      if is_html(cache_txt):
        # Got HTML response -> error
        logging("Error in fetching metadata\n")
        return {}
    finally:
      client.close()

  return parse_uniprot_metadata_with_seqids(seqids, cache_txt)


def batch_uniprot_metadata(seqids, cache_dir=None, batch_size=400):
  """
  Returns a dictonary of the uniprot metadata (as parsed
  by parse_uniprot_txt_file) of the given seqids. The seqids
  must be valid uniprot identifiers.

  Now handles isoform versions of accession id's!
  """

  unique_seqids = list(set(seqids))

  if cache_dir:
    seqids_txt = os.path.join(cache_dir, 'seqids.json')
    if os.path.isfile(seqids_txt):
        saved_seqids = json.load(open(seqids_txt))
        if saved_seqids != unique_seqids:
            shutil.rmtree(cache_dir)

    if not os.path.isdir(cache_dir):
      os.makedirs(cache_dir)

    if not os.path.isfile(seqids_txt):
      json.dump(unique_seqids, open(seqids_txt, 'w'))

  metadata = {}
  i_seqid = 0
  if batch_size is None:
    batch_size = len(unique_seqids)
  while i_seqid <= len(unique_seqids):
    seqids_subset = unique_seqids[i_seqid:i_seqid+batch_size]
    if cache_dir:
      subset_cache = os.path.join(cache_dir, 'metadata.%d.txt' % i_seqid)
    else:
      subset_cache = None
    metadata_subset = fetch_uniprot_metadata(
        seqids_subset, cache_fname=subset_cache)
    metadata.update(metadata_subset)
    i_seqid += batch_size
  return metadata


def is_text(seqid):
  if re.match('[A-Z,a-z,_]+$', seqid):
    return True
  return False


def is_uniprot(seqid):
  if re.match('[A-N,R-Z][0-9][A-Z][A-Z,0-9][A-Z,0-9][0-9]$', seqid):
    return True
  if re.match('[O,P,Q][0-9][A-Z,0-9][A-Z,0-9][A-Z,0-9][0-9]$', seqid):
    return True
  return False


def is_uniprot_variant(seqid):
  if is_uniprot(seqid[:6]):
    if len(seqid) == 6:
      return True
    variant = seqid[6:]
    if re.match(r'[-]\d+', variant):
      return True
  return False


def is_sgd(seqid):
  if re.match(r'Y[A-Z][L,R]\d\d\d[W|C]$', seqid):
    return True
  return False


def is_refseq(seqid):
  if re.match(r'[N,X,Y,Z][P,M]_\d+([.]\d+)?$', seqid):
    return True
  return False


def is_ensembl(seqid):
  if seqid.startswith("ENS"):
    return True
  return False


def is_maybe_uniprot_id(seqid):
  return '_' in seqid


def get_naked_seqid(seqid):
  if '|' not in seqid:
    return seqid
  pieces = seqid.split('|')
  if is_text(pieces[0]):
    return pieces[1]
  return seqid





def probe_id_type(entries, is_id_fn, name, uniprot_mapping_type, cache_fname):
  alternative_ids = []
  for entry in entries:
    if entry['id_type'] != '':
      continue
    if is_id_fn(entry['seqid']):
      entry['id_type'] = uniprot_mapping_type
      alternative_ids.append(entry['seqid'])
  if len(alternative_ids) == 0:
    return
  n_id = len(alternative_ids)
  pairs = batch_uniprot_id_mapping_pairs(
      uniprot_mapping_type, 'UniProtKB', alternative_ids, cache_dir=cache_fname)
  alternative_to_uniprot = { p[0]:p[1] for p in pairs }
  for entry in entries:
    if entry['seqid'] in alternative_to_uniprot:
      entry['uniprot_acc'] = alternative_to_uniprot[entry['seqid']]


def get_metadata_with_some_seqid_conversions(seqids, cache_dir=None):
  logging("Looking up uniprot metadata for %d seqids\n" % len(seqids))
  entries = []
  for seqid in seqids:
    entries.append({
      'raw_seqid':seqid,
      'seqid':'',
      'id_type':'',
      'uniprot_acc':'',
      'metadata':''})

  # break up pieces when in form xx|xxxxx|xxx
  for entry in entries:
    entry['seqid'] = get_naked_seqid(entry['raw_seqid'])

  # convert a few types into uniprot_ids
  id_types = [
    (is_sgd, 'locustag', 'SGD'),
    (is_refseq, 'refseqp', 'RefSeq_Protein'),
    (is_refseq, 'refseqnt', 'RefSeq_Nucleotide'),
    (is_ensembl, 'ensembl', 'Ensembl'),
    (is_maybe_uniprot_id, 'uniprotid', 'UniProtKB_AC-ID')]
  for is_id_fn, name, uniprot_mapping_type in id_types:
    if cache_dir:
      seqid_cache_fname = os.path.join(cache_dir, name)
    else:
      seqid_cache_fname = None
    probe_id_type(entries, is_id_fn, name, uniprot_mapping_type, seqid_cache_fname)

  # delete the variant suffixes in some uniprot id's as probe_id_type
  # can't cope with uniprot variant for id mapping lookup (bad!)
  for entry in entries:
    if entry['id_type'] == '' and is_uniprot_variant(entry['seqid']):
      entry['seqid'] = entry['seqid'][:6]

  # map UNIPROT ID's to their current best entry
  if cache_dir:
    seqid_cache_fname = os.path.join(cache_dir, 'uniprotuniprot')
  else:
    seqid_cache_fname = None
  probe_id_type(entries, is_uniprot, 'UNIPROT-ACC', 'UniProtKB_AC-ID', seqid_cache_fname)

  uniprot_seqids = []
  for entry in entries:
    if entry.get('uniprot_acc'):
      uniprot_seqids.append(entry['uniprot_acc'])
    if is_uniprot_variant(entry['raw_seqid']):
      # put the isoform variants back up id name as
      # batch_uniprot_metadata can handle isoforms!
      uniprot_seqids.append(entry['raw_seqid'])
      entry['uniprot_acc'] = entry['raw_seqid']
  uniprot_dict = batch_uniprot_metadata(uniprot_seqids, cache_dir)

  result = {}
  for entry in entries:
    if entry['uniprot_acc'] == "":
      continue
    uniprot_acc = entry['uniprot_acc']
    if uniprot_acc not in uniprot_dict:
      continue
    result[entry['raw_seqid']] = uniprot_dict[uniprot_acc]

  return result


def get_filtered_uniprot_metadata(seqids, cache_txt):
  """
  Returns a dictionary of uniprot data, but filters
  seqids for uniprot identifiers by using a mapping call
  to uniprot first.
  """

  stripped_seqids = [s[:6] for s in seqids]
  pairs = batch_uniprot_id_mapping_pairs(
      'UniProtKB_AC-ID', 'UniProtKB', stripped_seqids)
  uniprot_seqids = []
  for seqid1, seqid2 in pairs:
    if seqid1 in stripped_seqids and seqid1 not in uniprot_seqids:
      uniprot_seqids.append(seqid1)
  uniprot_dict = batch_uniprot_metadata(uniprot_seqids, cache_txt)
  for seqid in seqids:
    if seqid not in uniprot_seqids and seqid[:6] in uniprot_seqids:
      uniprot_dict[seqid] = uniprot_dict[seqid[:6]]
  return uniprot_dict


def sort_seqids_by_uniprot(seqids, uniprot_data):
  """
  Returns a sorted list of seqids such that longest proteins that are
  reviewed appear first in the list.
  """

  def cmp_longer_protein_is_first(seqid1, seqid2):
    return uniprot_data[seqid2]['length'] \
        - uniprot_data[seqid1]['length']

  def diff_list(orig_list, other_list):
    return [v for v in orig_list if v not in other_list]

  uniprot_seqids = list(filter(lambda s: s in uniprot_data, seqids))
  uniprot_seqids.sort(key=cmp_to_key(cmp_longer_protein_is_first))

  remainder_seqids = diff_list(seqids, uniprot_seqids)

  is_reviewed = lambda s: uniprot_data[s]['is_reviewed']
  reviewed_seqids = list(filter(is_reviewed, uniprot_seqids))
  reviewed_seqids.sort(key=cmp_to_key(cmp_longer_protein_is_first))

  unreviewed_seqids = diff_list(uniprot_seqids, reviewed_seqids)

  return reviewed_seqids + unreviewed_seqids + remainder_seqids


def parse_fasta_header(header, seqid_fn=None):
  """
  Parses a FASTA format header (with our without the initial '>').

  If NCBI SeqID format (gi|gi-number|gb|accession etc, is detected
  the first id in the list is used as the canonical id (see see
  http://www.ncbi.nlm.nih.gov/books/NBK21097/#A631 ).

  Extra processing to parse the seqid can be provided
  by giving a seqid_fn function.

  Returns a tuple of sequence id and sequence name/description.
  """
  # check to see if we have an NCBI-style header
  if header[0] == '>':
    header = header[1:]
  if header.find("|") != -1 and '|' in header.split()[0]:
    tokens = header.split('|')
    # "gi|ginumber|gb|accession bla bla" becomes "gi|ginumber"
    seqid = "%s|%s" % (tokens[0], tokens[1].split()[0])
    name = seqid + ' ' + tokens[-1].strip()
  else:
    # otherwise just split on spaces & hope for the best
    tokens = header.split()
    seqid = tokens[0]
    name = header[0:-1].strip()

  if seqid_fn is not None:
    seqid = seqid_fn(seqid)

  return seqid, name



def read_selected_fasta(seqids, fasta_db, seqid_fn=None):
  """
  Extracts protein sequences from a fasta database, given
  a list of desired seqids. A seqid_fn can be given that
  parses both the input seqids and the seqid in the fasta
  file to faciliate matching, but the keys in the returned
  structure corresponds to entries in 'seqids'.
  """
  live_seqid = None
  proteins = {}
  if seqid_fn is not None:
     original_seqid_map = { seqid_fn(s):s for s in seqids }
     seqids = list(original_seqid_map.keys())
  for i, line in enumerate(open(fasta_db)):
    if line.startswith(">"):
      fasta_seqid, description = \
          parse_fasta_header(line, seqid_fn)
      live_seqid = None
      for seqid in seqids:
        if fasta_seqid == seqid:
          live_seqid = fasta_seqid
          if seqid_fn:
            live_seqid = original_seqid_map[fasta_seqid]
          proteins[live_seqid] = {
            'sequence': "",
            'description': description,
          }
          break
    elif live_seqid:
      proteins[live_seqid]['sequence'] += line.strip()
  for seqid in proteins:
    sequence = proteins[seqid]['sequence']
    if sequence:
      proteins[seqid]['length'] = len(sequence)
  return proteins



def read_fasta(fasta_db, seqid_fn=None):
  """
  Parses a fasta database file. Warning: will be slow for very large
  databases. In that case, it would be better to use
  read_selected_fasta() instead.

  Returns a lsit of seqids encountered, and a dictionary
  of sequences for each seqid.
  """
  seqids = []
  seqid = None
  proteins = {}
  for line in open(fasta_db):
    if line.startswith(">"):
      seqid, description = parse_fasta_header(line, seqid_fn)
      seqids.append(seqid)
      proteins[seqid] = {
        'sequence':"",
        'description':description,
      }
      continue
    if seqid is not None:
      words = line.split()
      if words:
        proteins[seqid]['sequence'] += words[0]
  return seqids, proteins



def write_fasta(
    fasta_fname, proteins, seqids, width=50):
  """
  Creates a fasta file of the sequences of a subset of the proteins.
  """
  f = open(fasta_fname, "w")
  for seqid in seqids:
    f.write(">" + seqid)
    if 'description' in proteins[seqid]:
      f.write(" " + proteins[seqid]['description'])
    f.write("\n")
    sequence = proteins[seqid]['sequence']
    for i in range(0, len(sequence), width):
      f.write(sequence[i:i+width] + "\n")
  f.close()


# SeqIDType tool
# Updated to use current UniProt API field names (as of 2025)
_SEQIDTYPE_SCRAPE = """
UniProt
UniProtKB                                UniProtKB                      both
UniProtKB AC/ID                          UniProtKB_AC-ID                both
UniProtKB/Swiss-Prot                     UniProtKB-Swiss-Prot           both
UniParc                                  UniParc                        both
UniRef50                                 UniRef50                       both
UniRef90                                 UniRef90                       both
UniRef100                                UniRef100                      both
Gene Name                                Gene_Name                      both
CRC64                                    CRC64                          both
Sequence databases
CCDS                                     CCDS                           both
EMBL/GenBank/DDBJ                        EMBL-GenBank-DDBJ              both
EMBL/GenBank/DDBJ CDS                    EMBL-GenBank-DDBJ_CDS          both
GI number                                GI_number                      both
PIR                                      PIR                            both
RefSeq Nucleotide                        RefSeq_Nucleotide              both
RefSeq Protein                           RefSeq_Protein                 both
3D structure databases
PDB                                      PDB                            both
Protein-protein interaction databases
BioGRID                                  BioGRID                        both
ComplexPortal                            ComplexPortal                  both
DIP                                      DIP                            both
STRING                                   STRING                         both
Chemistry
ChEMBL                                   ChEMBL                         both
DrugBank                                 DrugBank                       both
GuidetoPHARMACOLOGY                      GuidetoPHARMACOLOGY            both
SwissLipids                              SwissLipids                    both
Protein family/group databases
Allergome                                Allergome                      both
ESTHER                                   ESTHER                         both
MEROPS                                   MEROPS                         both
PeroxiBase                               PeroxiBase                     both
REBASE                                   REBASE                         both
TCDB                                     TCDB                           both
PTM databases
GlyConnect                               GlyConnect                     both
Genetic variation databases
BioMuta                                  BioMuta                        both
DMDM                                     DMDM                           both
Proteomic databases
CPTAC                                    CPTAC                          both
ProteomicsDB                             ProteomicsDB                   both
Protocols and materials databases
DNASU                                    DNASU                          both
Genome annotation databases
Ensembl                                  Ensembl                        both
Ensembl Genomes                          Ensembl_Genomes                both
Ensembl Genomes Protein                  Ensembl_Genomes_Protein        both
Ensembl Genomes Transcript               Ensembl_Genomes_Transcript     both
Ensembl Protein                          Ensembl_Protein                both
Ensembl Transcript                       Ensembl_Transcript             both
GeneID                                   GeneID                         both
KEGG                                     KEGG                           both
PATRIC                                   PATRIC                         both
UCSC                                     UCSC                           both
WBParaSite                               WBParaSite                     both
WBParaSite Transcript/Protein            WBParaSite_Transcript-Protein  both
Organism-specific databases
ArachnoServer                            ArachnoServer                  both
Araport                                  Araport                        both
CGD                                      CGD                            both
ConoServer                               ConoServer                     both
dictyBase                                dictyBase                      both
EchoBASE                                 EchoBASE                       both
euHCVdb                                  euHCVdb                        both
FlyBase                                  FlyBase                        both
GeneCards                                GeneCards                      both
GeneReviews                              GeneReviews                    both
HGNC                                     HGNC                           both
LegioList                                LegioList                      both
Leproma                                  Leproma                        both
MaizeGDB                                 MaizeGDB                       both
MGI                                      MGI                            both
MIM                                      MIM                            both
neXtProt                                 neXtProt                       both
OpenTargets                              OpenTargets                    both
Orphanet                                 Orphanet                       both
PharmGKB                                 PharmGKB                       both
PomBase                                  PomBase                        both
PseudoCAP                                PseudoCAP                      both
RGD                                      RGD                            both
SGD                                      SGD                            both
TubercuList                              TubercuList                    both
VEuPathDB                                VEuPathDB                      both
VGNC                                     VGNC                           both
WormBase                                 WormBase                       both
WormBase Protein                         WormBase_Protein               both
WormBase Transcript                      WormBase_Transcript            both
Xenbase                                  Xenbase                        both
ZFIN                                     ZFIN                           both
Phylogenomic databases
eggNOG                                   eggNOG                         both
GeneTree                                 GeneTree                       both
HOGENOM                                  HOGENOM                        both
OMA                                      OMA                            both
OrthoDB                                  OrthoDB                        both
TreeFam                                  TreeFam                        both
Enzyme and pathway databases
BioCyc                                   BioCyc                         both
PlantReactome                            PlantReactome                  both
Reactome                                 Reactome                       both
UniPathway                               UniPathway                     both
Miscellaneous
ChiTaRS                                  ChiTaRS                        both
GeneWiki                                 GeneWiki                       both
GenomeRNAi                               GenomeRNAi                     both
PHI-base                                 PHI-base                       both
Gene expression databases
CollecTF                                 CollecTF                       both
Family and domain databases
DisProt                                  DisProt                        both
IDEAL                                    IDEAL                          both
"""


def _get_uniprot_mapping_rules():
  """Fetch which 'to' types are valid for each 'from' type from UniProt API."""
  try:
    response = httpx.get('https://rest.uniprot.org/configure/idmapping/fields', timeout=10.0)
    if response.status_code != 200:
      return {}
    data = response.json()
    field_to_rule = {}
    for group in data.get('groups', []):
      for item in group.get('items', []):
        if item.get('from', False):
          field_to_rule[item['name']] = item.get('ruleId')
    rule_to_tos = {rule['ruleId']: rule.get('tos', []) for rule in data.get('rules', [])}
    return {ft: rule_to_tos[rid] for ft, rid in field_to_rule.items() if rid in rule_to_tos}
  except Exception:
    return {}


def _get_seqidtype_id_types():
  """Parse seqid types from scrape data."""
  id_types = []
  for line in _SEQIDTYPE_SCRAPE.splitlines():
    words = line.split()
    if words and words[-1] in ['both', 'to']:
      id_types.append(words[-2])
  return id_types


def seqidtype_analyze(seqid, cache_fname=None):
    """
    Analyzes the type of seqid at http://uniprot.org by
    brute-force matching seqids against all seqid types.
    Maps to UniProtKB (the most common valid destination type).
    """
    id_types = _get_seqidtype_id_types()
    mapping_rules = _get_uniprot_mapping_rules()
    cache = json.load(open(cache_fname)) if (cache_fname and os.path.isfile(cache_fname)) else {seqid: {}}
    
    logging("Analyzing %s\n" % seqid)
    good_types = []
    
    for from_type in id_types:
      if seqid not in cache:
        cache[seqid] = {}
      if from_type not in cache[seqid]:
        # Only attempt mappings that are valid according to API rules
        if from_type in mapping_rules and "UniProtKB" in mapping_rules[from_type]:
          pairs = batch_uniprot_id_mapping_pairs(from_type, "UniProtKB", [seqid])
          cache[seqid][from_type] = pairs[0][1] if pairs else None
        else:
          cache[seqid][from_type] = None
        if cache_fname:
          json.dump(cache, open(cache_fname, 'w'))
      
      result = cache[seqid][from_type]
      if result is not None:
        good_types.append(from_type)
        logging('%s:%s -> %s\n' % (seqid, from_type, result))
      else:
        logging('%s:%s -> None\n' % (seqid, from_type))

    logging('\n%s is compatible with: %s\n' % (seqid, ', '.join(good_types) if good_types else 'None'))


def seqidtype_cli():
  """Entry point for seqidtype CLI tool."""
  import sys
  usage = textwrap.dedent("""\
    `seqidtype` works out the type of seqid at http://uniprot.org by
    brute-force matching seqids against all seqid types

    (c) 2013, Bosco Ho. BSD.

    seqidtype seqid1 seqid2 seqid3 ...

    (Example: seqidtype YOR261C)
    """)
  if len(sys.argv) == 1:
    print(usage)
  else:
    for seqid in sys.argv[1:]:
      seqidtype_analyze(seqid, 'seqidtype.json')
