Alignment.py

import os
import sys
import shlex
import subprocess
from pathlib import Path
import shutil

########################## REQUIREMENTS ##################
# REQUIRES APPLICATIONS (command line installation):  minimap2, samtools, gzip
# REQUIRED FILES in Input Directory:
# - Compressed fastq files: fastq.gz generated by nanopore or can be generat
# - library.fasta/ library.fa of reference sequences
###########################################################


def check_dependency(tool: str):
    """Check if a tool is installed and available in PATH."""
    if shutil.which(tool) is None:
        raise EnvironmentError(f"Required tool '{tool}' is not installed or not in PATH.")


def run_command(cmd, log_file="log.txt", output_file=None, cwd=None):
    """
    Run a command, log output, optionally write stdout to a file, stream text to console.
    Works with binary outputs like BAM safely and captures stderr for debugging.
    CHANGED: runs in shell to mimic .zsh/bash environment.
    """
    with open(log_file, "a") as lf:
        lf.write(f"\n[CMD] {' '.join(cmd)}\n")

        out_f = open(output_file, "wb") if output_file else None

        try:
            process = subprocess.Popen(
                " ".join([shlex.quote(c) for c in cmd]),
                cwd=cwd,
                stdout=subprocess.PIPE,
                stderr=subprocess.PIPE,  # capture stderr
                shell=True,              # CHANGED: run in shell
                env=os.environ.copy(),
            )

            stdout_data, stderr_data = process.communicate()

            # Write stdout
            if out_f:
                out_f.write(stdout_data)
            else:
                try:
                    lf.write(stdout_data.decode())
                    print(stdout_data.decode(), end="")
                except UnicodeDecodeError:
                    lf.write("[binary stdout skipped]\n")

            # Always log and print stderr
            try:
                err_text = stderr_data.decode()
                if err_text.strip():
                    print(err_text, end="")
                    lf.write(err_text)
            except UnicodeDecodeError:
                lf.write("[binary stderr skipped]\n")

            if process.returncode != 0:
                raise RuntimeError(
                    f"Command failed ({process.returncode}): {' '.join(cmd)}\nError message:\n{stderr_data.decode(errors='replace')}"
                )
        finally:
            if out_f:
                out_f.close()


def process_single_directory(directory: Path):
    """Process a single directory containing compressed fastq files."""
    if not directory.exists() or not directory.is_dir():
        raise FileNotFoundError(f"Directory {directory} does not exist")

    alignment_dir = directory / "Alignment"
    alignment_dir.mkdir(parents=True, exist_ok=True)

    # Gather files
    gz_files = list(directory.glob("*.gz"))
    fastq_files = list(directory.glob("*.fastq"))
    library_file = next((directory / lib for lib in ["library.fasta", "library.fa"] if (directory / lib).exists()), None)

    if not gz_files and not fastq_files:
        raise FileNotFoundError(f"No .gz or .fastq files found in {directory}")
    if library_file is None:
        raise FileNotFoundError(f"No library.fasta or library.fa found in {directory}")

    # Clean old fastq files
    for fq in fastq_files:
        fq.unlink()

    # Decompress gz files
    for gz in gz_files:
        run_command(["gzip", "-dk", str(gz)], cwd=directory)

    # Concatenate into single.fastq
    single_fastq_path = alignment_dir / "single.fastq"
    with open(single_fastq_path, "wb") as outfile:
        for fq in directory.glob("*.fastq"):
            with open(fq, "rb") as infile:
                shutil.copyfileobj(infile, outfile)

    # Move library to alignment directory
    aligned_library = alignment_dir / "library.fasta"
    shutil.copy(str(library_file), aligned_library)

    # Ensure FASTA index exists
    faidx_path = aligned_library.with_suffix(".fasta.fai")
    if not faidx_path.exists():
        run_command(["samtools", "faidx", str(aligned_library)], cwd=alignment_dir)

    # Alignment steps
    run_command(
        ["minimap2", "-ax", "map-ont", str(aligned_library), str(single_fastq_path)],
        output_file=str(alignment_dir / "aln.sam"),
        cwd=alignment_dir,
    )
    run_command(
        ["samtools", "view", "-bT", str(aligned_library), "-F", "16", "aln.sam"],
        output_file=str(alignment_dir / "aln.bam"),
        cwd=alignment_dir,
    )
    run_command(
        ["samtools", "sort", "aln.bam"],
        output_file=str(alignment_dir / "aln.s.bam"),
        cwd=alignment_dir,
    )
    run_command(
        ["samtools", "index", "aln.s.bam"],
        cwd=alignment_dir,
    )
    run_command(
        ["samtools", "calmd", "-bAr", "aln.s.bam", str(aligned_library)],
        output_file=str(alignment_dir / "aln.baq.bam"),
        cwd=alignment_dir,
    )
    run_command(
        ["samtools", "index", "aln.baq.bam"],
        cwd=alignment_dir,
    )

    # CHANGED/ADDED: print and return alignment directory
    print(f"Alignment directory created: {alignment_dir}")
    return alignment_dir


def process_multiple_directories(base_dir: Path):
    alignment_dirs = []  # CHANGED/ADDED: store all output dirs
    for f in base_dir.iterdir():
        if f.is_dir() and list(f.glob("*compressed")):
            print(f"Processing {f}")
            alignment_dirs.append(process_single_directory(f))
    return alignment_dirs  # CHANGED/ADDED


def process_single_file(file_path: Path):
    """Process a single .gz file, creating a working directory for it."""
    if file_path.suffix != ".gz":
        raise ValueError(f"File {file_path} is not a .gz file")

    base_dir = file_path.parent
    work_dir = base_dir / file_path.stem
    work_dir.mkdir(exist_ok=True)

    gz_copy = work_dir / file_path.name
    shutil.copy(file_path, gz_copy)

    # Copy library if exists
    for libname in ["library.fasta", "library.fa"]:
        if (base_dir / libname).exists():
            shutil.copy(base_dir / libname, work_dir / libname)

    # CHANGED/ADDED: process and return alignment directory
    alignment_dir = process_single_directory(work_dir)
    return alignment_dir


def run_alignment(input_path: str):
    # Dependency checks
    for tool in ["minimap2", "samtools", "gzip"]:
        check_dependency(tool)

    base = Path(input_path)
    output_dirs = []  # CHANGED/ADDED: store all alignment dirs

    if base.is_file():
        print(f"Detected single file mode: {base}")
        output_dirs.append(process_single_file(base))
    elif base.is_dir():
        fastq_gz = list(base.glob("*.gz"))
        subdirs_with_compressed = [d for d in base.iterdir() if d.is_dir() and list(d.glob("*.gz"))]

        if fastq_gz:
            print(f"Detected single directory mode: {base}")
            output_dirs.append(process_single_directory(base))
        elif subdirs_with_compressed:
            print(f"Detected multiple directory mode under: {base}")
            output_dirs.extend(process_multiple_directories(base))
        else:
            raise FileNotFoundError(f"No compressed .fastq files found in {base}")
    else:
        raise FileNotFoundError(f"Input path {base} does not exist")

    # CHANGED/ADDED: print all output directories
    print("\nAll alignment output directories:")
    for d in output_dirs:
        print(d)

    return output_dirs  # CHANGED/ADDED: return list of alignment directories


if __name__ == "__main__":
    # input directory with required files
    alignment_dirs = run_alignment("/Users/timshel/transcripts/fast5_422_compressed_2")