Protein Sequence Similarity Search

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

"""Quick protein homologue search via ColabFold MMseqs2 API.

Submits a protein sequence to the ColabFold MMseqs2 server, downloads the
resulting MSA archive, parses the A3M alignment headers, and writes a
Markdown-formatted table of sequence homologues sorted by E-value to an
output file specified via the required --output flag.
"""

# /// script
# requires-python = ">=3.10"
# dependencies = [
#   "scienceskillscommon",
# ]
# [tool.uv.sources]
# scienceskillscommon = { path = "../../scienceskillscommon" }
# ///

import argparse
import json
import os
import shutil
import sys
import tarfile
import tempfile
import time
import urllib.error
import urllib.parse

from science_skills.skills.scienceskillscommon import http_client

MAX_ALIGNMENT_HITS = 300
POLLING_TIMEOUT = 15 * 60  # 15 minutes.
COLABFOLD_HOST = "https://api.colabfold.com"
FASTA_COLUMNS = [
    "target",
    "bit_score",
    "identity",
    "e_value",
    "q_start",
    "q_end",
    "q_len",
    "t_start",
    "t_end",
    "t_len",
]
CLIENT = http_client.HttpClient(COLABFOLD_HOST, qps=2)


def read_sequence(query_input, tty=None):
  """Read sequence from a file path or raw string."""
  if os.path.isfile(query_input):
    print(f"[*] Reading sequence from file: {query_input}")
    if tty:
      print(f"[*] Reading sequence from file: {query_input}", file=tty)
    with open(query_input, "r") as f:
      sequence = f.read().strip()
      if sequence.startswith(">"):
        print("[*] Sequence is in FASTA format")
        if tty:
          print("[*] Sequence is in FASTA format", file=tty)
        sequence = "".join(sequence.split("\n")[1:])  # Remove FASTA header
      else:
        print("[*] Sequence is in raw format")  # No further processing needed
        if tty:
          print("[*] Sequence is in raw format", file=tty)
  else:
    print("[*] Using raw sequence string provided via command line.")
    sequence = query_input.strip()

  if not sequence:
    print("[!] Error: Empty sequence provided.")
    if tty:
      print("[!] Error: Empty sequence provided.", file=tty)
    sys.exit(1)

  return sequence


def parse_a3m(file_path, q_len):
  """Parse ColabFold-annotated A3M headers into hit dictionaries."""
  homologues = []

  with open(file_path, "r") as f:
    for line in f:
      if not line.startswith(">"):
        continue

      parts = line.strip().split()

      # Skip query header (no stat columns) or malformed lines
      if len(parts) < 10:
        continue

      try:
        hit = dict(zip(FASTA_COLUMNS, parts, strict=True))
        for col in ["q_start", "q_end", "q_len", "t_start", "t_end", "t_len"]:
          hit[col] = int(hit[col])
        for col in ["bit_score", "identity", "e_value"]:
          hit[col] = float(hit[col])

        # Query Coverage
        if q_len > 0 and hit["q_end"] > hit["q_start"]:
          aligned_residues = hit["q_end"] - hit["q_start"] + 1
          cov = min((aligned_residues / q_len) * 100, 100.0)
        else:
          cov = 0.0

        # Alignment length (target span)
        if hit["t_end"] > hit["t_start"]:
          aln_len = hit["t_end"] - hit["t_start"] + 1
        else:
          aln_len = 0

        hit |= {
            "target_id": hit["target"][1:],  # Strip leading '>'
            "q_cov": cov,
            "aln_len": aln_len,
        }
        homologues.append(hit)
      except (ValueError, IndexError):
        print(
            f"[!] Warning: Skipping malformed hit: {line.strip()}",
            file=sys.stderr,
        )
        continue

  homologues.sort(key=lambda x: x["e_value"])
  return homologues


def search_mmseqs2(query_input, output_md, json_file, include_mgnify=False):
  """Submit sequence to ColabFold MMseqs2, parse results, write output."""
  # Redirect all print output to the .md file
  original_stdout = sys.stdout
  md_file = open(output_md, "w")
  sys.stdout = md_file

  try:
    _run_mmseqs2(query_input, json_file, include_mgnify, tty=original_stdout)
  finally:
    sys.stdout = original_stdout
    md_file.close()

  print(f"[*] Results saved to: {output_md}")
  if json_file:
    print(f"[*] Results saved to: {json_file}")
  print("...done!")


def _run_mmseqs2(query_input, json_file, include_mgnify, tty):
  """Internal: runs the actual MMseqs2 search with print going to file."""
  sequence = read_sequence(query_input, tty=tty)
  q_len = len(sequence)

  # --- 1. Submit ---
  print("[*] Search method: MMseqs2 (ColabFold API)")
  print(
      f"[*] Submitting sequence ({q_len} residues) to ColabFold MMseqs2 API..."
  )
  print("[*] Search method: MMseqs2 (ColabFold API)", file=tty)

  query_fasta = f">Query_1\n{sequence}\n"
  data = urllib.parse.urlencode({
      "q": query_fasta,
      "mode": "all",
  }).encode("ascii")

  try:
    ticket = CLIENT.fetch_json(
        f"{COLABFOLD_HOST}/ticket/msa", method="POST", data=data, timeout=60
    )
  except (
      urllib.error.URLError,
      urllib.error.HTTPError,
      TimeoutError,
      RuntimeError,
  ) as e:
    print(f"[!] MMseqs2 API Submission Failed: {e}")
    print(f"[!] MMseqs2 API Submission Failed: {e}", file=tty)
    sys.exit(2)

  ticket_id = ticket.get("id")
  if not ticket_id:
    status = ticket.get("status", "UNKNOWN")
    if status == "RATELIMIT":
      print("[!] MMseqs2 rate limit hit on submission.")
      print("[!] MMseqs2 rate limit hit on submission.", file=tty)
    else:
      print(f"[!] MMseqs2 submission failed: {ticket}")
      print(f"[!] MMseqs2 submission failed: {ticket}", file=tty)
    sys.exit(2)

  print(f"[*] Ticket ID generated: {ticket_id}")
  print(f"[*] Ticket ID: {ticket_id}", file=tty)

  # --- 2. Poll ---
  print("[*] Polling server for completion...")
  print("[*] Polling server for completion...", file=tty)

  start_time = time.monotonic()
  while time.monotonic() - start_time < POLLING_TIMEOUT:
    state = CLIENT.fetch_json(
        f"{COLABFOLD_HOST}/ticket/{ticket_id}", timeout=20
    ).get("status", "UNKNOWN")

    # States below are taken from the original ColabFold code:
    # https://github.com/sokrypton/ColabFold/blob/main/colabfold/colabfold.py:201
    if state == "COMPLETE":
      print("\n[*] Job finished successfully!")
      print("\n[*] Job finished successfully!", file=tty)
      break

    elif state in ("ERROR", "MAINTENANCE"):
      print(f"[!] MMseqs2 job failed with status: {state}")
      print(f"[!] MMseqs2 job failed with status: {state}", file=tty)
      sys.exit(2)

    elif state == "RATELIMIT":
      print("[!] Rate limit hit. Waiting...")
      print("[!] Rate limit hit. Waiting...", file=tty)

    elif state in ("PENDING", "RUNNING", "UNKNOWN"):
      sys.stdout.write(".")
      sys.stdout.flush()
      tty.write(".")
      tty.flush()
    time.sleep(10)

  else:
    print("[!] Polling timed out.")
    print("[!] Polling timed out.", file=tty)
    sys.exit(2)

  # --- 3. Download & Extract ---
  print("[*] Downloading and extracting MSA files...")
  tmp_dir = tempfile.mkdtemp(prefix="mmseqs2_")
  tar_path = os.path.join(tmp_dir, f"{ticket_id}.tar.gz")

  try:
    raw_data = CLIENT.fetch_bytes(
        f"{COLABFOLD_HOST}/result/download/{ticket_id}", timeout=120
    )
    with open(tar_path, "wb") as f:
      f.write(raw_data)

    with tarfile.open(tar_path, "r:gz") as tar:
      for member in tar.getmembers():
        if member.name.startswith("/") or ".." in member.name:
          continue
        tar.extract(member, path=tmp_dir)

    os.remove(tar_path)
  except (
      urllib.error.URLError,
      urllib.error.HTTPError,
      TimeoutError,
      RuntimeError,
      OSError,
      tarfile.ReadError,
  ) as e:
    print(f"[!] Failed to download/extract results: {e}")
    print(f"[!] Failed to download/extract results: {e}", file=tty)
    shutil.rmtree(tmp_dir, ignore_errors=True)
    sys.exit(2)

  # --- 4. Parse A3M files ---
  all_homologues = []

  uniref_path = os.path.join(tmp_dir, "uniref.a3m")
  if os.path.exists(uniref_path):
    uniref_hits = parse_a3m(uniref_path, q_len)
    print(f"[*] Parsed {len(uniref_hits)} hits from uniref.a3m")
    all_homologues.extend(uniref_hits)
  else:
    print("[!] Warning: uniref.a3m not found in results archive.")

  if include_mgnify:
    mgnify_path = os.path.join(tmp_dir, "bfd.mgnify30.metaeuk30.smag30.a3m")
    if os.path.exists(mgnify_path):
      mgnify_hits = parse_a3m(mgnify_path, q_len)
      print(f"[*] Parsed {len(mgnify_hits)} hits from mgnify a3m")
      all_homologues.extend(mgnify_hits)
    else:
      print("[!] Warning: mgnify a3m file not found in results archive.")
      print(
          "[!] Warning: mgnify a3m file not found in results archive.", file=tty
      )

  # Cleanup temp directory
  shutil.rmtree(tmp_dir, ignore_errors=True)

  if not all_homologues:
    print("[!] No homologues found.")
    print("[!] No homologues found.", file=tty)
    sys.exit(0)

  # Sort combined results by E-value, take top hits
  all_homologues.sort(key=lambda x: x["e_value"])
  all_homologues = all_homologues[:MAX_ALIGNMENT_HITS]

  # --- 5. Save JSON ---
  if json_file:
    with open(json_file, "w") as f:
      json.dump(all_homologues, f, indent=4)
    print(f"[*] JSON results successfully saved to: {json_file}")

  # --- 6. Output Markdown Table ---
  print(f"\n### Top {len(all_homologues)} Sequence Homologues (MMseqs2)")
  print("| Target ID | Q-Cov | E-value | Seq Identity(%) | Aln Length |")
  print("|---|---|---|---|---|")

  for hit in all_homologues:
    target_id = hit["target_id"]
    q_cov = f"{hit['q_cov']:.1f}%"
    e_value = f"{hit['e_value']:.2e}"
    identity = f"{hit['identity'] * 100:.1f}%"
    aln_len = str(hit["aln_len"])

    print(f"| {target_id} | {q_cov} | {e_value} | {identity} | {aln_len} |")


def main():
  parser = argparse.ArgumentParser(
      description=(
          "Quick protein homologue search via ColabFold MMseqs2 API. "
          "Exits with code 2 on API failures to signal fallback to BLAST."
      )
  )
  parser.add_argument(
      "query_input", help="Path to a FASTA file or raw sequence string"
  )
  parser.add_argument(
      "-o",
      "--output",
      required=True,
      help="Path to save the output Markdown (.md) file (required)",
  )
  parser.add_argument(
      "-j",
      "--json",
      help="Path to save the output JSON file (optional)",
      default=None,
  )
  parser.add_argument(
      "--include-mgnify",
      action="store_true",
      help="Also parse and include hits from the mgnify/environmental database",
  )
  args = parser.parse_args()

  print(f"[*] Output: {args.output}")
  if args.json:
    print(f"[*] Output: {args.json}")

  search_mmseqs2(args.query_input, args.output, args.json, args.include_mgnify)


if __name__ == "__main__":
  main()

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

"""Runs EBI BLAST for a FASTA file or raw sequence string against UniProt.

Writes a Markdown-formatted results table to the file specified by the
required --output flag. Optionally saves raw JSON via -j/--json.
"""

# /// script
# requires-python = ">=3.10"
# dependencies = [
#   "scienceskillscommon",
#   "python-dotenv",
# ]
# [tool.uv.sources]
# scienceskillscommon = { path = "../../scienceskillscommon" }
# ///

import argparse
import json
import os
import sys
import time
import urllib.error
import urllib.parse

import dotenv
from science_skills.skills.scienceskillscommon import http_client

ALLOWED_DATABASES = [
    "uniprotkb",
    "uniprotkb_swissprot",
    "uniprotkb_swissprotsv",
    "uniprotkb_reference_proteomes",
    "uniprotkb_trembl",
    "uniprotkb_refprotswissprot",
    "uniprotkb_archaea",
    "uniprotkb_arthropoda",
    "uniprotkb_bacteria",
    "uniprotkb_complete_microbial_proteomes",
    "uniprotkb_eukaryota",
    "uniprotkb_fungi",
    "uniprotkb_human",
    "uniprotkb_mammals",
    "uniprotkb_nematoda",
    "uniprotkb_rodents",
    "uniprotkb_vertebrates",
    "uniprotkb_viridiplantae",
    "uniprotkb_viruses",
    "uniprotkb_enzyme",
    "uniprotkb_covid19",
    "uniref100",
    "uniref90",
    "uniref50",
    "pdb",
]

MAX_ALIGNMENT_HITS = 300
DEFAULT_EVALUE = 1000
POLLING_TIMEOUT = 15 * 60  # 15 minutes.
BASE_URL = "https://www.ebi.ac.uk/Tools/services/rest/ncbiblast"
CLIENT = http_client.HttpClient(BASE_URL, qps=2)


def search_uniprot(query_input, output_md, json_file, databases):
  """Runs EBI BLAST search, writes results to output files."""
  # Redirect all print output to the .md file
  original_stdout = sys.stdout
  with open(output_md, "w") as md_file:
    sys.stdout = md_file
    try:
      _run_blast(query_input, json_file, databases, tty=original_stdout)
    finally:
      sys.stdout = original_stdout

  print(f"[*] Results saved to: {output_md}")
  if json_file:
    print(f"[*] Results saved to: {json_file}")
  print("...done!")


def _run_blast(query_input, json_file, databases, tty):
  """Internal: runs the actual BLAST search with print going to file."""

  # Process and validate the databases
  selected_dbs = [db.strip().lower() for db in databases.split(",")]
  invalid_dbs = [db for db in selected_dbs if db not in ALLOWED_DATABASES]

  if invalid_dbs:
    print(f"[!] Error: Invalid database(s) provided: {', '.join(invalid_dbs)}")
    print(f"[*] Allowed databases are: {', '.join(ALLOWED_DATABASES)}")
    print(
        f"[!] Error: Invalid database(s) provided: {', '.join(invalid_dbs)}",
        file=tty,
    )
    sys.exit(1)

  # Determine if input is a file path or a raw sequence string
  if os.path.isfile(query_input):
    print(f"[*] Reading sequence from file: {query_input}")
    with open(query_input, "r") as f:
      sequence = f.read().strip()
      # If it's a fasta file with a header, skip the header line
      if sequence.startswith(">"):
        sequence = "".join(sequence.split("\n")[1:])
  else:
    print("[*] Using raw sequence string provided via command line.")
    sequence = query_input.strip()

  if not sequence:
    print("[!] Error: Empty sequence provided.")
    print("[!] Error: Empty sequence provided.", file=tty)
    sys.exit(1)

  # Calculate query length for later Q-Cov math
  q_len = len(sequence)

  print("[*] Search method: EBI BLAST")
  print(f"[*] Submitting sequence ({q_len} residues) to EBI NCBI BLAST API...")
  print(f"[*] Searching databases: {', '.join(selected_dbs)}")
  print(
      f"[*] Search method: EBI BLAST — databases: {', '.join(selected_dbs)}",
      file=tty,
  )

  # 1. Submit Job
  url_run = f"{BASE_URL}/run"
  params = {
      "program": "blastp",
      "stype": "protein",
      "sequence": sequence,
      "database": ",".join(selected_dbs),
  }
  user_email = os.environ.get("USER_EMAIL")
  if user_email:
    params["email"] = user_email
  else:
    print(
        "[!] Warning: USER_EMAIL environment variable not set. Request may"
        " fail."
    )
  data = urllib.parse.urlencode(params).encode("utf-8")
  try:
    job_id = CLIENT.fetch_text(url_run, method="POST", data=data).strip()
    print(f"[*] Job ID generated: {job_id}")
    print(f"[*] Job ID: {job_id}", file=tty)
  except (
      http_client.HttpError,
      TimeoutError,
      RuntimeError,
      IOError,
  ) as e:
    print(f"[!] API Submission Failed: {e}")
    print(f"[!] API Submission Failed: {e}", file=tty)
    sys.exit(1)
  # 2. Poll the server until the search finishes
  print("[*] Polling server for completion (this may take a few minutes)...")
  print("[*] Polling server for completion...", file=tty)
  url_status = f"{BASE_URL}/status/{job_id}"

  start_time = time.monotonic()
  while time.monotonic() - start_time < POLLING_TIMEOUT:

    try:
      status = CLIENT.fetch_text(url_status, timeout=20).strip()

      if status == "FINISHED":
        print("\n[*] Job marked as FINISHED.")
        print("\n[*] Job marked as FINISHED.", file=tty)
        break
      elif status in ["ERROR", "FAILURE", "NOT_FOUND"]:
        print(f"\n[!] BLAST job failed with status: {status}")
        print(f"\n[!] BLAST job failed with status: {status}", file=tty)
        sys.exit(1)

      sys.stdout.write(".")
      sys.stdout.flush()
      tty.write(".")
      tty.flush()
      time.sleep(30)

    except (
        urllib.error.URLError,
        http_client.HttpError,
        TimeoutError,
        RuntimeError,
    ) as e:
      print(f"\n[!] Polling error: {e}")
      print(f"\n[!] Polling error: {e}", file=tty)
      sys.exit(1)

  else:
    print("\n[!] Polling timed out.")
    print("\n[!] Polling timed out.", file=tty)
    sys.exit(1)

  time.sleep(3)

  # 3. Fetch and format the results
  print("\n[*] Job complete. Fetching results...\n")
  try:
    res = CLIENT.fetch_json(f"{BASE_URL}/result/{job_id}/json", timeout=60)
  except (
      urllib.error.URLError,
      http_client.HttpError,
      TimeoutError,
      RuntimeError,
      json.JSONDecodeError,
      IOError,
  ) as e:
    print(f"[!] Failed to fetch or parse results: {e}")
    sys.exit(1)

  # Save JSON to file (optional)
  if json_file:
    with open(json_file, "w") as f:
      json.dump(res, f, indent=4)
    print(f"[*] Raw JSON results successfully saved to: {json_file}\n")

  # Output as a Markdown Table for LLM Agent parsing
  print(f"### Top {MAX_ALIGNMENT_HITS} Sequence Homologues (EBI BLAST)")
  print("| Target ID | Q-Cov | E-value | Seq Identity(%) | Aln Length |")
  print("|---|---|---|---|---|")

  # The EBI JSON wraps hits in a 'hits' array
  hits = res.get("hits", [])

  if not hits:
    print("[!] No homologues found.")
    return

  for hit in hits[:MAX_ALIGNMENT_HITS]:
    target_acc = hit.get("hit_acc", "N/A")
    target_desc = hit.get("hit_desc", hit.get("hit_def", ""))

    # Combine Accession and Description for the LLM
    target = f"{target_acc} {target_desc}".strip()

    hsps = hit.get("hit_hsps", [{}])
    best_hsp = hsps[0] if hsps else {}

    # Calculate Query Coverage
    hsp_query_from = int(best_hsp.get("hsp_query_from", 0))
    hsp_query_to = int(best_hsp.get("hsp_query_to", 0))

    if q_len > 0 and hsp_query_to > hsp_query_from:
      aligned_q_residues = (hsp_query_to - hsp_query_from) + 1
      cov_percentage = min((aligned_q_residues / q_len) * 100, 100.0)
      q_cov = f"{cov_percentage:.1f}%"
    else:
      q_cov = "N/A"

    evalue = str(best_hsp.get("hsp_expect", "N/A"))
    seq_id = str(best_hsp.get("hsp_identity", "N/A")) + "%"
    aln_len = str(best_hsp.get("hsp_align_len", "N/A"))

    # Print Markdown row without truncation
    print(f"| {target} | {q_cov} | {evalue} | {seq_id} | {aln_len} |")


def main():
  dotenv.load_dotenv(os.path.expanduser("~/.env"))
  parser = argparse.ArgumentParser(
      description=(
          "Query EBI NCBI BLAST with a FASTA file or raw sequence string."
      )
  )
  parser.add_argument(
      "query_input", help="Path to the FASTA file or raw sequence string"
  )
  parser.add_argument(
      "-o",
      "--output",
      required=True,
      help="Path to save the output Markdown (.md) file (required)",
  )
  parser.add_argument(
      "-j",
      "--json",
      help="Path to save the output JSON file (optional)",
      default=None,
  )
  parser.add_argument(
      "--databases",
      help="Comma-separated list of databases to search",
      default="uniprotkb",
  )
  args = parser.parse_args()

  print(f"[*] Output: {args.output}")
  if args.json:
    print(f"[*] Output: {args.json}")

  search_uniprot(args.query_input, args.output, args.json, args.databases)


if __name__ == "__main__":
  main()