Ncbi Sequence Fetch

scripts/ncbi_fetch.py

# Copyright 2026 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""NCBI E-utilities wrapper for protein and nucleotide sequence retrieval.

Provides subcommands for the core NCBI E-utilities operations needed for
protein sequence retrieval: efetch, esearch, elink, CDS translation,
patent sequence search, and gene-to-protein lookup.

Rate-limited to 3 requests/second (10/s with NCBI_API_KEY).
"""

# /// script
# requires-python = ">=3.10"
# dependencies = [
#   "polite-http",
#   "python-dotenv",
# ]
# ///

from __future__ import annotations

import argparse
import gzip
import json
import os
import re
import sys
from typing import Any
import urllib.parse
import urllib.request
import xml.etree.ElementTree as ET

import dotenv
from polite_http import http_client

_CODON_TABLE = {
    'ATA': 'I',
    'ATC': 'I',
    'ATT': 'I',
    'ATG': 'M',
    'ACA': 'T',
    'ACC': 'T',
    'ACG': 'T',
    'ACT': 'T',
    'AAC': 'N',
    'AAT': 'N',
    'AAA': 'K',
    'AAG': 'K',
    'AGC': 'S',
    'AGT': 'S',
    'AGA': 'R',
    'AGG': 'R',
    'CTA': 'L',
    'CTC': 'L',
    'CTG': 'L',
    'CTT': 'L',
    'CCA': 'P',
    'CCC': 'P',
    'CCG': 'P',
    'CCT': 'P',
    'CAC': 'H',
    'CAT': 'H',
    'CAA': 'Q',
    'CAG': 'Q',
    'CGA': 'R',
    'CGC': 'R',
    'CGG': 'R',
    'CGT': 'R',
    'GTA': 'V',
    'GTC': 'V',
    'GTG': 'V',
    'GTT': 'V',
    'GCA': 'A',
    'GCC': 'A',
    'GCG': 'A',
    'GCT': 'A',
    'GAC': 'D',
    'GAT': 'D',
    'GAA': 'E',
    'GAG': 'E',
    'GGA': 'G',
    'GGC': 'G',
    'GGG': 'G',
    'GGT': 'G',
    'TCA': 'S',
    'TCC': 'S',
    'TCG': 'S',
    'TCT': 'S',
    'TTC': 'F',
    'TTT': 'F',
    'TTA': 'L',
    'TTG': 'L',
    'TAC': 'Y',
    'TAT': 'Y',
    'TAA': '*',
    'TAG': '*',
    'TGC': 'C',
    'TGT': 'C',
    'TGA': '*',
    'TGG': 'W',
}


EUTILS_BASE = 'https://eutils.ncbi.nlm.nih.gov/entrez/eutils'

_CLIENT = None


def get_api_client():
  """Returns the lazily initialized HttpClient."""
  global _CLIENT
  if _CLIENT is None:
    api_key = os.environ.get('NCBI_API_KEY', '')
    qps = 10.0 if api_key else 3.0
    _CLIENT = http_client.HttpClient(EUTILS_BASE, qps=qps)
  return _CLIENT


def _eutils_get(endpoint: str, params: dict[str, str | int]) -> str | None:
  """Sends a GET request to an NCBI E-utilities endpoint.

  Handles rate limiting, retries with backoff, and API key
  injection.

  Args:
    endpoint: E-utilities endpoint name (e.g. 'efetch.fcgi').
    params: Query parameters as a dict.

  Returns:
    Response text or None on failure.
  """
  api_key = os.environ.get('NCBI_API_KEY', '')
  if api_key:
    params['api_key'] = api_key

  query_string = urllib.parse.urlencode(params)
  full_url = f'{EUTILS_BASE}/{endpoint}'
  if query_string:
    full_url += f'?{query_string}'

  try:
    res = get_api_client().fetch(full_url)
    data = res.data
    if data.startswith(b'\x1f\x8b'):
      try:
        data = gzip.decompress(data)
      except gzip.BadGzipFile:
        pass
    encoding = getattr(res, 'encoding', None) or 'utf-8'
    return data.decode(encoding, errors='replace')
  except http_client.HttpError as e:
    print(f'{endpoint} error after all retries: {e}', file=sys.stderr)
    return None


def efetch(
    db: str,
    db_id: str | int,
    retmode: str = 'text',
    rettype: str = 'fasta',
) -> str | None:
  """Fetches data from NCBI efetch endpoint.

  Args:
    db: Database name (protein, nuccore, gene, pubmed, etc.)
    db_id: One or more comma-separated IDs.
    retmode: Return mode (text, xml, json).
    rettype: Return type (fasta, gb, gp, fasta_cds_aa, etc.)

  Returns:
    Response text or None on failure.
  """
  return _eutils_get(
      'efetch.fcgi',
      {
          'db': db,
          'id': str(db_id),
          'rettype': rettype,
          'retmode': retmode,
      },
  )


def esearch(
    db: str,
    term: str,
    retmax: int = 20,
) -> tuple[list[str], int]:
  """Searches NCBI esearch endpoint.

  Args:
    db: Database name (protein, nuccore, gene, pubmed, etc.)
    term: Search query using Entrez syntax.
    retmax: Maximum results to return.

  Returns:
    Tuple of (list of ID strings, total count).
  """
  text = _eutils_get(
      'esearch.fcgi',
      {
          'db': db,
          'term': term,
          'retmode': 'json',
          'retmax': retmax,
      },
  )
  if text:
    try:
      data = json.loads(text)
      result = data.get('esearchresult', {})
      return result.get('idlist', []), int(result.get('count', 0))
    except (json.JSONDecodeError, ValueError) as e:
      print(f'esearch parse error: {e}', file=sys.stderr)
  return [], 0


def elink(dbfrom: str, db: str, link_id: str) -> list[str]:
  """Follows cross-database links via NCBI elink.

  Args:
    dbfrom: Source database (e.g., pubmed).
    db: Target database (e.g., protein).
    link_id: ID in the source database.

  Returns:
    List of linked IDs in the target database.
  """
  text = _eutils_get(
      'elink.fcgi',
      {
          'dbfrom': dbfrom,
          'db': db,
          'id': link_id,
          'retmode': 'json',
      },
  )
  if text:
    try:
      data = json.loads(text)
      linksets = data.get('linksets', [])
      if linksets:
        linksetdbs = linksets[0].get('linksetdbs', [])
        if linksetdbs:
          return linksetdbs[0].get('links', [])
    except (json.JSONDecodeError, ValueError) as e:
      print(f'elink parse error: {e}', file=sys.stderr)
  return []


def esummary(db: str, ids: str | list[str]) -> dict[str, Any]:
  """Fetches summary info for one or more IDs.

  Args:
    db: Database name.
    ids: Comma-separated ID string or list.

  Returns:
    Dict of id -> summary info.
  """
  if isinstance(ids, list):
    ids = ','.join(str(i) for i in ids)
  text = _eutils_get(
      'esummary.fcgi',
      {
          'db': db,
          'id': ids,
          'retmode': 'json',
      },
  )
  if text:
    try:
      data = json.loads(text)
      return data.get('result', {})
    except (json.JSONDecodeError, ValueError) as e:
      print(f'esummary parse error: {e}', file=sys.stderr)
  return {}


def parse_fasta(text: str | None) -> list[tuple[str, str]]:
  """Parses FASTA text into list of (header, sequence) tuples."""
  if not text:
    return []
  entries = []
  current_header = None
  current_seq = []
  for line in text.strip().splitlines():
    line = line.strip()
    if line.startswith('>'):
      if current_header is not None:
        entries.append((current_header, ''.join(current_seq)))
      current_header = line
      current_seq = []
    else:
      current_seq.append(line)
  if current_header is not None:
    entries.append((current_header, ''.join(current_seq)))
  return entries


def translate_dna(seq: str) -> str:
  """Translates a DNA sequence to protein using the standard codon table."""
  protein = []
  seq = seq.upper().replace('\n', '').replace(' ', '')
  for i in range(0, len(seq) - len(seq) % 3, 3):
    codon = seq[i : i + 3]
    if 'N' in codon:
      protein.append('X')
    else:
      aa = _CODON_TABLE.get(codon, 'X')
      protein.append(aa)
  return ''.join(protein).rstrip('*')


def get_longest_orf(seq: str, min_len: int = 75) -> str | None:
  """Finds the longest open reading frame in a DNA sequence.

  Searches all 6 reading frames (3 forward + 3 reverse complement).

  Args:
    seq: DNA sequence string.
    min_len: Minimum ORF length in amino acids.

  Returns:
    Longest ORF protein sequence, or None.
  """
  rev_comp = {'A': 'T', 'T': 'A', 'C': 'G', 'G': 'C', 'N': 'N'}
  rev_seq = ''.join(rev_comp.get(b, 'N') for b in reversed(seq.upper()))

  longest = ''
  for s in [seq.upper(), rev_seq]:
    for frame in range(3):
      translated = translate_dna(s[frame:])
      # Split on stop codons
      peptides = translated.split('*')
      for pep in peptides:
        idx = pep.find('M')
        if idx != -1:
          orf = pep[idx:]
          if len(orf) >= min_len and len(orf) > len(longest):
            longest = orf
  return longest if longest else None


def extract_cds_translations(
    xml_text: str | None,
    identifier: str | None = None,
    target_len: int = 0,
) -> list[dict[str, str]]:
  """Extracts CDS translations from GenBank XML.

  Args:
    xml_text: GenBank XML string from efetch.
    identifier: Optional gene name to filter by.
    target_len: Optional target protein length for best-match selection.

  Returns:
    List of dicts with keys: translation, product,
    gene, locus_tag, protein_id
  """
  del identifier, target_len  # Reserved for future filtering.
  if not xml_text:
    return []
  results = []
  try:
    root = ET.fromstring(xml_text)
    for seq in root.findall('.//GBSeq'):
      for feat in seq.findall('.//GBFeature'):
        key = feat.find('GBFeature_key')
        if key is None or key.text != 'CDS':
          continue
        entry = {
            'translation': '',
            'product': '',
            'gene': '',
            'locus_tag': '',
            'protein_id': '',
            'note': '',
        }
        for qual in feat.findall('.//GBQualifier'):
          name = qual.find('GBQualifier_name')
          val = qual.find('GBQualifier_value')
          if name is not None and val is not None and val.text:
            if name.text == 'translation':
              entry['translation'] = val.text.upper()
            elif name.text in [
                'product',
                'gene',
                'locus_tag',
                'protein_id',
                'note',
            ]:
              entry[name.text] = val.text
        if entry['translation']:
          results.append(entry)
  except ET.ParseError as e:
    print(f'XML parse error: {e}', file=sys.stderr)
  return results


def is_genomic_record(xml_text: str) -> bool:
  """Checks if a GenBank XML record is genomic DNA (not mRNA/CDS)."""
  try:
    root = ET.fromstring(xml_text)
    moltype = root.find('.//GBSeq_moltype')
    if moltype is not None and moltype.text:
      if 'genomic' in moltype.text.lower():
        return True
    for feat in root.iter('GBFeature'):
      k = feat.find('GBFeature_key')
      if k is not None and k.text == 'source':
        for qual in feat.iter('GBQualifier'):
          n = qual.find('GBQualifier_name')
          v = qual.find('GBQualifier_value')
          if (
              n is not None
              and n.text == 'mol_type'
              and v is not None
              and v.text
              and 'genomic' in v.text.lower()
          ):
            return True
        break
  except ET.ParseError:
    pass
  return False


def cmd_fetch_protein(args: argparse.Namespace) -> None:
  """Fetches protein sequence by accession ID."""
  results = []
  for acc in args.accession:
    fasta = efetch('protein', acc)
    entries = parse_fasta(fasta) if fasta else []
    if entries:
      for header, seq in entries:
        results.append({
            'accession': acc,
            'header': header,
            'sequence': seq,
            'length': len(seq),
        })
    else:
      # Try with dot variant
      if '_' in acc:
        fasta = efetch('protein', acc.replace('_', '.'))
        entries = parse_fasta(fasta) if fasta else []
        for header, seq in entries:
          results.append({
              'accession': acc,
              'header': header,
              'sequence': seq,
              'length': len(seq),
          })
      if not any(r['accession'] == acc for r in results):
        results.append({
            'accession': acc,
            'error': 'Not found',
        })
  _write_output(results, args.output)


def cmd_fetch_nucleotide(args: argparse.Namespace) -> None:
  """Fetches nucleotide sequence by accession ID."""
  results = []
  for acc in args.accession:
    fasta = efetch('nuccore', acc)
    entries = parse_fasta(fasta) if fasta else []
    if entries:
      for header, seq in entries:
        results.append({
            'accession': acc,
            'header': header,
            'sequence': seq,
            'length': len(seq),
        })
    else:
      results.append({'accession': acc, 'error': 'Not found'})
  _write_output(results, args.output)


def cmd_cds_translate(args: argparse.Namespace) -> None:
  """Fetches CDS and translates to protein sequence.

  Tries multiple approaches:
  1. NCBI fasta_cds_aa (pre-translated CDS protein)
  2. GenBank XML CDS translation annotations
  3. Raw nucleotide to 6-frame ORF finding (if not genomic)

  Args:
    args: Parsed CLI args with accession, target_length, and output attributes.
  """
  results = []
  for acc in args.accession:
    result = {'accession': acc}

    # Approach 1: Pre-translated CDS protein
    fasta = efetch('nuccore', acc, rettype='fasta_cds_aa')
    if fasta:
      entries = parse_fasta(fasta)
      if entries:
        best = _pick_best_by_length(entries, args.target_length)
        if best:
          header, seq = best
          seq = seq.replace('*', '')
          # Check for dominant isoform if multiple entries
          if len(entries) > 1 and not args.target_length:
            dom = _get_dominant_cds(acc)
            if dom:
              seq = dom
              header = f'>dominant_isoform_{acc}'
          result['header'] = header
          result['sequence'] = seq
          result['length'] = len(seq)
          result['method'] = 'fasta_cds_aa'
          results.append(result)
          continue

    # Approach 2: GenBank XML CDS translations
    xml_data = efetch('nuccore', acc, retmode='xml', rettype='gb')
    if xml_data:
      genomic = is_genomic_record(xml_data)
      cds_list = extract_cds_translations(xml_data)
      if cds_list:
        if args.target_length:
          best_cds = min(
              cds_list,
              key=lambda c: abs(len(c['translation']) - args.target_length),
          )
        else:
          best_cds = max(
              cds_list,
              key=lambda c: len(c['translation']),
          )
        result['header'] = (
            f'>{best_cds["gene"]}'
            f'|{best_cds["product"]}'
            f'|{best_cds["protein_id"]}'
        )
        result['sequence'] = best_cds['translation']
        result['length'] = len(best_cds['translation'])
        result['method'] = 'genbank_xml_cds'
        result['cds_info'] = best_cds
        if genomic:
          result['is_genomic'] = True
        results.append(result)
        continue
      if genomic:
        result['is_genomic'] = True

    # Approach 3: Raw nucleotide ORF finding
    fasta = efetch('nuccore', acc)
    if fasta:
      entries = parse_fasta(fasta)
      if entries:
        _, dna_seq = entries[0]
        orf = get_longest_orf(dna_seq)
        if orf:
          tl = args.target_length
          if tl == 0 or abs(len(orf) - tl) < 50:
            result['header'] = f'>ORF_{acc}'
            result['sequence'] = orf
            result['length'] = len(orf)
            result['method'] = 'orf_translation'
            results.append(result)
            continue

    result['error'] = 'No CDS translation found'
    results.append(result)

  _write_output(results, args.output)


def cmd_search(args: argparse.Namespace) -> None:
  """Searches an NCBI database by query terms."""
  ids, count = esearch(args.database, args.query, retmax=args.retmax)
  result = {
      'database': args.database,
      'query': args.query,
      'total_count': count,
      'returned_count': len(ids),
      'ids': ids,
  }

  # Optionally fetch FASTA for returned IDs
  if args.fetch_sequences and ids:
    sequences = []
    batch_size = 20
    for i in range(0, len(ids), batch_size):
      batch = ids[i : i + batch_size]
      fasta = efetch(args.database, ','.join(batch))
      if fasta:
        entries = parse_fasta(fasta)
        for header, seq in entries:
          sequences.append({
              'header': header,
              'sequence': seq,
              'length': len(seq),
          })
    result['sequences'] = sequences
  _write_output(result, args.output)


def cmd_elink(args: argparse.Namespace) -> None:
  """Follows cross-database links from one NCBI database to another."""
  linked_ids = elink(args.dbfrom, args.db, args.id)
  result = {
      'source_db': args.dbfrom,
      'target_db': args.db,
      'source_id': args.id,
      'linked_ids': linked_ids,
      'count': len(linked_ids),
  }

  if args.fetch_sequences and linked_ids:
    sequences = []
    batch_size = 20
    for i in range(0, len(linked_ids), batch_size):
      batch = linked_ids[i : i + batch_size]
      fasta = efetch(args.db, ','.join(batch))
      if fasta:
        entries = parse_fasta(fasta)
        for header, seq in entries:
          sequences.append({
              'header': header,
              'sequence': seq,
              'length': len(seq),
          })
    result['sequences'] = sequences

  _write_output(result, args.output)


def cmd_gene_protein(args: argparse.Namespace) -> None:
  """Searches for protein sequence by gene name and organism.

  Tries NCBI Protein database with gene name + organism
  filter. Optionally filters by sequence length.

  Args:
    args: Parsed CLI args with gene, organism, target_length, retmax, and output
      attributes.
  """
  gene = args.gene
  organism = args.organism
  results = []

  # Search NCBI Protein
  query = f'{gene}[Gene Name] AND {organism}[Organism]'
  if args.target_length:
    lo = args.target_length - 25
    hi = args.target_length + 25
    query += f' AND {lo}:{hi}[Sequence Length]'

  ids, count = esearch('protein', query, retmax=args.retmax)
  if ids:
    for uid in ids[: args.retmax]:
      fasta = efetch('protein', uid)
      if fasta:
        entries = parse_fasta(fasta)
        for header, seq in entries:
          results.append({
              'uid': uid,
              'header': header,
              'sequence': seq,
              'length': len(seq),
          })

  output = {
      'gene': gene,
      'organism': organism,
      'query': query,
      'total_count': count,
      'results': results,
  }
  _write_output(output, args.output)


def cmd_locus_protein(args: argparse.Namespace) -> None:
  """Searches for protein sequence by locus tag.

  Tries:
  1. NCBI Protein search by gene name + organism
  2. NCBI Nuccore search for GenBank records with CDS
     translations

  Args:
    args: Parsed CLI args with locus, organism, retmax, and output attributes.
  """
  locus = args.locus
  organism = args.organism
  results = []

  # Try NCBI Protein
  query = f'{locus}[Gene Name]'
  if organism:
    query += f' AND {organism}[Organism]'
  ids, _ = esearch('protein', query, retmax=args.retmax)
  if ids:
    for uid in ids[: args.retmax]:
      fasta = efetch('protein', uid)
      if fasta:
        entries = parse_fasta(fasta)
        for header, seq in entries:
          if len(seq) >= 50:
            results.append({
                'uid': uid,
                'header': header,
                'sequence': seq,
                'length': len(seq),
                'source': 'protein_db',
            })

  # Try NCBI Nuccore for CDS translations
  nuc_query = f'{locus}[Accession] OR {locus}[All Fields]'
  nuc_ids, _ = esearch('nuccore', nuc_query, retmax=10)
  if nuc_ids:
    xml_data = efetch(
        'nuccore',
        ','.join(nuc_ids),
        retmode='xml',
        rettype='gb',
    )
    if xml_data:
      cds_list = extract_cds_translations(xml_data, identifier=locus)
      for cds in cds_list:
        results.append({
            'header': f'>{cds["gene"]}|{cds["product"]}|{cds["protein_id"]}',
            'sequence': cds['translation'],
            'length': len(cds['translation']),
            'source': 'nuccore_cds',
            'gene': cds['gene'],
            'locus_tag': cds['locus_tag'],
        })

  output = {
      'locus': locus,
      'organism': organism or '',
      'total_results': len(results),
      'results': results,
  }
  _write_output(output, args.output)


def cmd_pubmed_proteins(args: argparse.Namespace) -> None:
  """Finds protein sequences linked to a PubMed article.

  Searches:
  1. NCBI Protein records with PMID filter
  2. NCBI elink PubMed to Protein
  3. NCBI Nuccore records with PMID for CDS
     translations

  Args:
    args: Parsed CLI args with pmid, identifier, and output attributes.
  """
  pmid = args.pmid
  identifier = args.identifier
  results = []

  # Step 1: Protein db search by PMID
  prot_ids, _ = esearch('protein', f'{pmid}[PMID]', retmax=100)

  # Step 2: elink PubMed→Protein if no direct hits
  if not prot_ids:
    prot_ids = elink('pubmed', 'protein', pmid)

  if prot_ids:
    batch_size = 20
    for i in range(0, len(prot_ids), batch_size):
      batch = prot_ids[i : i + batch_size]
      xml_data = efetch('protein', ','.join(batch), retmode='xml', rettype='gp')
      if xml_data:
        try:
          root = ET.fromstring(xml_data)
          # NCBI returns GBSeq or INSDSeq depending on format
          seq_els = root.findall('.//GBSeq')
          pfx = 'GBSeq'
          if not seq_els:
            seq_els = root.findall('.//INSDSeq')
            pfx = 'INSDSeq'
          for gb_seq in seq_els:
            defn = gb_seq.find(f'{pfx}_definition')
            acc = gb_seq.find(f'{pfx}_primary-accession')
            seq_el = gb_seq.find(f'{pfx}_sequence')
            if seq_el is not None and seq_el.text:
              defn_text = (defn.text or '') if defn is not None else ''
              entry = {
                  'accession': acc.text if acc is not None else '',
                  'definition': defn_text,
                  'sequence': seq_el.text.upper(),
                  'length': len(seq_el.text),
                  'source': 'protein_db',
              }
              if identifier:
                entry['matches_identifier'] = (
                    identifier.lower() in defn_text.lower()
                )
              results.append(entry)
        except (ET.ParseError, AttributeError):
          pass

  # Step 3: Nuccore records for CDS translations
  nuc_ids, _ = esearch('nuccore', f'{pmid}[PMID]', retmax=100)
  if nuc_ids:
    for i in range(0, len(nuc_ids), 10):
      batch = nuc_ids[i : i + 10]
      xml_data = efetch('nuccore', ','.join(batch), retmode='xml', rettype='gb')
      if xml_data:
        cds_list = extract_cds_translations(xml_data, identifier=identifier)
        for cds in cds_list:
          entry = {
              'gene': cds['gene'],
              'product': cds['product'],
              'protein_id': cds['protein_id'],
              'sequence': cds['translation'],
              'length': len(cds['translation']),
              'source': 'nuccore_cds',
          }
          if identifier:
            entry['matches_identifier'] = (
                identifier.lower() in cds['gene'].lower()
                or identifier.lower() in cds['product'].lower()
                or identifier.lower() in cds['protein_id'].lower()
            )
          results.append(entry)

  output = {
      'pmid': pmid,
      'identifier': identifier or '',
      'total_results': len(results),
      'results': results,
  }
  _write_output(output, args.output)


def cmd_patent_search(args: argparse.Namespace) -> None:
  """Searches for protein sequences from patents.

  Two modes:
  1. By patent number: fetches all protein sequences
     from a specific patent
  2. By keywords: searches NCBI Protein with
     patent[Properties] filter

  Args:
    args: Parsed CLI args with patent_number, keywords, organism, retmax, and
      output attributes.
  """
  results = []

  if args.patent_number:
    # Mode 1: Search by patent number
    patent = args.patent_number.strip()

    # Extract the core 7-10 digit number, ignoring prefixes like US/WO/EP
    # and suffixes like B2/A1
    digit_match = re.search(r'(\d{7,10})', patent)
    digits = digit_match.group(1) if digit_match else ''

    # Build queries in priority order: most specific first.
    # Avoid broad text searches like 'patent {patent}' which match millions
    # of unrelated records and pollute the ID list.
    queries = []
    if digits:
      queries.append(f'{digits}[Patent Publication Number]')

    # Extract country prefix (e.g. US, WO, EP) for exact phrase search
    prefix_match = re.match(r'^([A-Za-z]{2})\s*', patent)
    if prefix_match and digits:
      country = prefix_match.group(1).upper()
      queries.append(f'"{country} {digits}"')

    # Handle international patents
    if re.search(r'^(WO|EP)', patent, re.IGNORECASE):
      stripped = re.sub(r'^(WO|EP)', '', patent)
      queries.append(f'patent {stripped}')

    # Fallback: broad text search only if no specific queries matched
    if not queries:
      queries.append(f'patent {patent}')

    all_ids = []
    for q in queries:
      ids, _ = esearch('protein', q, retmax=args.retmax)
      all_ids.extend(ids)
      # If we already have enough IDs from a specific query, stop
      if len(set(all_ids)) >= args.retmax:
        break
    all_ids = list(dict.fromkeys(all_ids))
    all_ids = all_ids[: args.retmax]  # dedupe preserving order

    if all_ids:
      batch_size = 20
      for i in range(0, len(all_ids), batch_size):
        batch = all_ids[i : i + batch_size]
        fasta = efetch('protein', ','.join(batch))
        if fasta:
          entries = parse_fasta(fasta)
          for header, seq in entries:
            # Filter results that don't mention our patent's digits
            if digits and digits not in header:
              continue

            seq_id_match = re.search(r'Sequence (\d+)', header)
            pat_match = re.search(
                r'patent\s+(?:[a-zA-Z]{2}\s*)?(\d{7,10})',
                header,
                re.IGNORECASE,
            )
            results.append({
                'header': header,
                'sequence': seq,
                'length': len(seq),
                'seq_id_no': (
                    int(seq_id_match.group(1)) if seq_id_match else None
                ),
                'patent_number': pat_match.group(1) if pat_match else patent,
            })

  elif args.keywords:
    # Mode 2: Search by keywords with patent filter
    keywords = args.keywords
    query_parts = []
    for kw in keywords:
      query_parts.append(kw)
    query = ' AND '.join(query_parts)

    if args.organism:
      query += f' AND {args.organism}[Organism]'
    query += ' AND patent[Properties]'

    ids, _ = esearch('protein', query, retmax=args.retmax)
    if ids:
      batch_size = 20
      for i in range(0, len(ids), batch_size):
        batch = ids[i : i + batch_size]
        fasta = efetch('protein', ','.join(batch))
        if fasta:
          entries = parse_fasta(fasta)
          for header, seq in entries:
            pat_match = re.search(
                r'patent\s+US\s*(\d{7,8})',
                header,
                re.IGNORECASE,
            )
            if pat_match:
              results.append({
                  'header': header,
                  'sequence': seq,
                  'length': len(seq),
                  'patent_number': pat_match.group(1),
              })

  output = {
      'patent_number': args.patent_number or '',
      'keywords': args.keywords or [],
      'organism': args.organism or '',
      'total_results': len(results),
      'results': results,
  }
  _write_output(output, args.output)


def cmd_organism_length(args: argparse.Namespace) -> None:
  """Searches for protein by organism and exact length.

  Uses NCBI SLEN (Sequence Length) filter to find
  proteins of exact length from a given organism.
  Useful as a last-resort lookup when only organism
  and expected AA length are known.

  Args:
    args: Parsed CLI args with organism, length, retmax, and output attributes.
  """
  organism = args.organism
  length = args.length
  results = []

  query = f'"{organism}"[Organism] AND {length}[SLEN]'
  ids, count = esearch('protein', query, retmax=args.retmax)
  if ids:
    batch_size = 20
    for i in range(0, len(ids), batch_size):
      batch = ids[i : i + batch_size]
      fasta = efetch('protein', ','.join(batch))
      if fasta:
        entries = parse_fasta(fasta)
        for header, seq in entries:
          results.append({
              'header': header,
              'sequence': seq,
              'length': len(seq),
          })

  output = {
      'organism': organism,
      'target_length': length,
      'query': query,
      'total_count': count,
      'results': results,
  }
  _write_output(output, args.output)


def _pick_best_by_length(
    entries: list[tuple[str, str]],
    target_len: int,
) -> tuple[str, str] | None:
  """Picks the FASTA entry closest to target length, or longest if no target."""
  if not entries:
    return None
  valid = [(h, s) for h, s in entries if len(s) >= 50]
  if not valid:
    return None
  if target_len and target_len > 0:
    return min(valid, key=lambda x: abs(len(x[1]) - target_len))
  return max(valid, key=lambda x: len(x[1]))


def _get_dominant_cds(accession: str) -> str | None:
  """Returns the 'dominant isoform' CDS translation if annotated in GenBank."""
  xml_data = efetch('nuccore', accession, retmode='xml', rettype='gb')
  if not xml_data:
    return None
  cds_list = extract_cds_translations(xml_data)
  for cds in cds_list:
    if 'dominant' in cds.get('note', '').lower():
      return cds['translation']
  return None


def _write_output(
    data: dict[str, Any] | list[Any],
    output_path: str | None,
) -> None:
  """Writes results to JSON file and prints summary to stdout."""
  if output_path:
    with open(output_path, 'w') as f:
      json.dump(data, f, indent=2)
    print(f'Results written to: {output_path}')
  else:
    print(json.dumps(data, indent=2))

  # Print summary to stdout
  if isinstance(data, list):
    for item in data:
      if 'error' in item:
        print(f"  {item.get('accession', '?')}: ERROR — {item['error']}")
      elif 'sequence' in item:
        print(f"  {item.get('accession', '?')}: {item['length']} AA")
  elif isinstance(data, dict):
    if 'results' in data:
      n = len(data['results'])
      print(f'  Found {n} result(s)')
      for r in data['results'][:5]:
        if 'sequence' in r:
          desc = r.get('header', r.get('definition', ''))[:80]
          print(f"    {r['length']} AA — {desc}")
    elif 'sequences' in data:
      n = len(data['sequences'])
      print(f'  Found {n} sequence(s)')
    elif 'ids' in data:
      print(
          f'  {data.get("total_count", 0)} total,'
          f' {data.get("returned_count", 0)} returned'
      )


def main() -> None:
  dotenv.load_dotenv(os.path.expanduser('~/.env'))
  parser = argparse.ArgumentParser(
      description='NCBI E-utilities wrapper for sequence retrieval.',
  )
  sub = parser.add_subparsers(dest='command', help='Subcommand')

  # fetch-protein
  p = sub.add_parser('fetch-protein', help='Fetch protein by accession')
  p.add_argument('accession', nargs='+', help='Protein accession(s)')
  p.add_argument('--output', '-o', help='Output JSON file')

  # fetch-nucleotide
  p = sub.add_parser('fetch-nucleotide', help='Fetch nucleotide by accession')
  p.add_argument('accession', nargs='+', help='Nucleotide accession(s)')
  p.add_argument('--output', '-o', help='Output JSON file')

  # cds-translate
  p = sub.add_parser(
      'cds-translate',
      help='Fetch CDS accession and translate to protein',
  )
  p.add_argument('accession', nargs='+', help='CDS/nucleotide accession(s)')
  p.add_argument(
      '--target-length',
      type=int,
      default=0,
      help='Expected protein length (AA) for best-match selection',
  )
  p.add_argument('--output', '-o', help='Output JSON file')

  # search
  p = sub.add_parser('search', help='Search an NCBI database')
  p.add_argument('query', help='Entrez search query')
  p.add_argument(
      '--database',
      '-d',
      default='protein',
      help='Database (protein, nuccore, gene, pubmed, etc.)',
  )
  p.add_argument(
      '--retmax', type=int, default=20, help='Maximum results to return'
  )
  p.add_argument(
      '--fetch-sequences',
      action='store_true',
      help='Also fetch FASTA sequences for returned IDs',
  )
  p.add_argument('--output', '-o', help='Output JSON file')

  # elink
  p = sub.add_parser('elink', help='Follow cross-database links')
  p.add_argument('id', help='Source ID')
  p.add_argument(
      '--dbfrom', required=True, help='Source database (e.g., pubmed)'
  )
  p.add_argument('--db', required=True, help='Target database (e.g., protein)')
  p.add_argument(
      '--fetch-sequences',
      action='store_true',
      help='Also fetch FASTA sequences for linked IDs',
  )
  p.add_argument('--output', '-o', help='Output JSON file')

  # gene-protein
  p = sub.add_parser(
      'gene-protein',
      help='Search protein by gene name + organism',
  )
  p.add_argument('gene', help='Gene name')
  p.add_argument('--organism', required=True, help='Organism name')
  p.add_argument(
      '--target-length',
      type=int,
      default=0,
      help='Expected protein length (AA)',
  )
  p.add_argument('--retmax', type=int, default=10, help='Max results')
  p.add_argument('--output', '-o', help='Output JSON file')

  # locus-protein
  p = sub.add_parser(
      'locus-protein',
      help='Search protein by locus tag',
  )
  p.add_argument('locus', help='Locus tag or accession')
  p.add_argument('--organism', help='Organism name')
  p.add_argument('--retmax', type=int, default=10, help='Max results')
  p.add_argument('--output', '-o', help='Output JSON file')

  # pubmed-proteins
  p = sub.add_parser(
      'pubmed-proteins',
      help='Find protein sequences linked to a PubMed article',
  )
  p.add_argument('pmid', help='PubMed ID')
  p.add_argument('--identifier', help='Gene/protein name to match')
  p.add_argument('--output', '-o', help='Output JSON file')

  # patent-search
  p = sub.add_parser(
      'patent-search',
      help='Search for protein sequences from patents',
  )
  p.add_argument('--patent-number', help='Patent number (e.g., US12345678)')
  p.add_argument(
      '--keywords', nargs='+', help='Search keywords for patent protein search'
  )
  p.add_argument('--organism', help='Organism filter')
  p.add_argument('--retmax', type=int, default=50, help='Max results')
  p.add_argument('--output', '-o', help='Output JSON file')

  # organism-length
  p = sub.add_parser(
      'organism-length',
      help='Search protein by organism and exact AA length',
  )
  p.add_argument('--organism', required=True, help='Organism name')
  p.add_argument(
      '--length', type=int, required=True, help='Exact protein length in AA'
  )
  p.add_argument('--retmax', type=int, default=50, help='Max results')
  p.add_argument('--output', '-o', help='Output JSON file')

  args = parser.parse_args()
  if not args.command:
    parser.print_help()
    sys.exit(1)

  cmd_map = {
      'fetch-protein': cmd_fetch_protein,
      'fetch-nucleotide': cmd_fetch_nucleotide,
      'cds-translate': cmd_cds_translate,
      'search': cmd_search,
      'elink': cmd_elink,
      'gene-protein': cmd_gene_protein,
      'locus-protein': cmd_locus_protein,
      'pubmed-proteins': cmd_pubmed_proteins,
      'patent-search': cmd_patent_search,
      'organism-length': cmd_organism_length,
  }
  cmd_map[args.command](args)


if __name__ == '__main__':
  main()