src/pipeline.py

import os
import random
import subprocess
import urllib.request
from pathlib import Path

CACHE_DIR = Path.home() / ".cache" / "insilicocontrols"
CACHE_DIR.mkdir(parents=True, exist_ok=True)

MANE_BB = CACHE_DIR / "MANE.bb"
MANE_BED12 = CACHE_DIR / "MANE.bed12"
HG38_FA = CACHE_DIR / "hg38.fa"
HG38_FAI = CACHE_DIR / "hg38.fa.fai"
HG19_FA = CACHE_DIR / "hg19.fa"
HG19_FAI = CACHE_DIR / "hg19.fa.fai"
BIGBEDTOBED_PATH = CACHE_DIR / "bigBedToBed"


def log(msg, log_func=None):
    if log_func:
        log_func(msg)
    else:
        print(msg)


def run_cmd(cmd, log_func=None):
    log(f"$ {cmd}", log_func)
    result = subprocess.run(cmd, shell=True, capture_output=True, text=True)
    if result.stdout:
        log(result.stdout.strip(), log_func)
    if result.stderr:
        log(result.stderr.strip(), log_func)
    return result.returncode == 0


def ensure_bigbedtobed(log_func=None):
    if BIGBEDTOBED_PATH.exists():
        return str(BIGBEDTOBED_PATH)
    log("Downloading bigBedToBed...", log_func)
    url = "https://hgdownload.cse.ucsc.edu/admin/exe/linux.x86_64/bigBedToBed"
    urllib.request.urlretrieve(url, str(BIGBEDTOBED_PATH))
    os.chmod(str(BIGBEDTOBED_PATH), 0o755)
    log("bigBedToBed ready.", log_func)
    return str(BIGBEDTOBED_PATH)


def ensure_mane(log_func=None):
    if MANE_BED12.exists():
        log("MANE annotation already cached.", log_func)
        return True
    log("Downloading MANE annotation...", log_func)
    bigbedtobed = ensure_bigbedtobed(log_func)
    url = "https://ftp.ncbi.nlm.nih.gov/refseq/MANE/trackhub/data/release_1.0/MANE.GRCh38.v1.0.refseq.bb"
    urllib.request.urlretrieve(url, str(MANE_BB))
    log("Converting bigBed to BED12...", log_func)
    ok = run_cmd(f"{bigbedtobed} {MANE_BB} {MANE_BED12}", log_func)
    if not ok:
        raise RuntimeError("Failed to convert MANE bigBed")
    log("MANE annotation ready.", log_func)
    return True


def ensure_reference(genome_version="hg38", log_func=None):
    if genome_version == "hg19":
        fa_path = HG19_FA
        fai_path = HG19_FAI
        url = "https://hgdownload.soe.ucsc.edu/goldenPath/hg19/bigZips/hg19.fa.gz"
    else:
        fa_path = HG38_FA
        fai_path = HG38_FAI
        url = "https://hgdownload.soe.ucsc.edu/goldenPath/hg38/bigZips/hg38.fa.gz"

    if fa_path.exists() and fai_path.exists():
        log(f"{genome_version} reference already cached.", log_func)
        return fa_path

    log(f"Downloading {genome_version}.fa.gz... This will take several minutes.", log_func)
    gz_path = CACHE_DIR / f"{genome_version}.fa.gz"
    urllib.request.urlretrieve(url, str(gz_path))

    log(f"Decompressing {genome_version}.fa.gz...", log_func)
    run_cmd(f"gunzip -f {gz_path}", log_func)

    log(f"Indexing {genome_version}.fa with samtools faidx...", log_func)
    run_cmd(f"samtools faidx {fa_path}", log_func)

    log(f"{genome_version} reference ready.", log_func)
    return fa_path


def parse_mane_exons(work_dir, log_func=None):
    exons_raw = work_dir / "hg38_exons_raw.bed"
    exons_sorted = work_dir / "hg38_exons.bed"

    log("Parsing MANE BED12 into individual exons...", log_func)
    with open(MANE_BED12, "r") as infile, open(exons_raw, "w") as outfile:
        for line in infile:
            cols = line.strip().split()
            if len(cols) < 12:
                continue
            chrom = cols[0]
            chromStart = int(cols[1])
            thickStart, thickEnd = int(cols[6]), int(cols[7])
            sizes = [int(x) for x in cols[10].strip(",").split(",")]
            starts = [int(x) for x in cols[11].strip(",").split(",")]
            num_exons = len(sizes)
            for i in range(num_exons):
                exon_start = chromStart + starts[i]
                exon_end = exon_start + sizes[i]
                if exon_end > thickStart and exon_start < thickEnd:
                    cds_start = max(exon_start, thickStart)
                    cds_end = min(exon_end, thickEnd)
                    has_left_intron = i > 0
                    has_right_intron = i < num_exons - 1
                    outfile.write(
                        f"{chrom}\t{exon_start}\t{exon_end}\t{cds_start}\t{cds_end}\t{has_left_intron}\t{has_right_intron}\n"
                    )

    run_cmd(f"bedtools sort -i {exons_raw} > {exons_sorted}", log_func)
    log("Exon parsing complete.", log_func)
    return exons_sorted


def analyze_coverage(work_dir, probes_bed, exons_bed, log_func=None):
    merged = work_dir / "merged_probes.bed"
    fully = work_dir / "fully_covered_exons.bed"
    any_cov = work_dir / "any_covered_exons.bed"
    partial = work_dir / "partially_covered_exons.bed"
    probes_no_exons = work_dir / "probes_without_exons.bed"
    unused = work_dir / "unused_probes.bed"

    log("Merging contiguous probes...", log_func)
    run_cmd(f"bedtools sort -i {probes_bed} | bedtools merge -i - > {merged}", log_func)

    log("Intersecting probes with exons...", log_func)
    run_cmd(f"bedtools intersect -a {exons_bed} -b {merged} -f 0.95 -wa -u > {fully}", log_func)
    run_cmd(f"bedtools intersect -a {exons_bed} -b {merged} -wa -u > {any_cov}", log_func)
    run_cmd(f"bedtools intersect -a {any_cov} -b {fully} -v > {partial}", log_func)
    run_cmd(f"bedtools intersect -a {probes_bed} -b {exons_bed} -v > {probes_no_exons}", log_func)
    run_cmd(f"bedtools intersect -a {merged} -b {any_cov} -v > {unused}", log_func)

    stats = {
        "fully_covered": int(subprocess.check_output(f"wc -l < {fully}", shell=True).strip() or 0),
        "partially_covered": int(subprocess.check_output(f"wc -l < {partial}", shell=True).strip() or 0),
        "probes_no_exons": int(subprocess.check_output(f"wc -l < {probes_no_exons}", shell=True).strip() or 0),
        "unused_probes": int(subprocess.check_output(f"wc -l < {unused}", shell=True).strip() or 0),
    }
    return stats, fully, partial, unused


def generate_target_snvs(work_dir, fully_bed, partial_bed, unused_bed,
                         include_cds=True, include_intron=True, include_offtarget=True,
                         mode="mane", probes_bed=None, direct_window_size=0,
                         log_func=None):
    output = work_dir / "igv_variant_navigator.bed"
    log("Generating target SNVs...", log_func)

    if mode == "direct_bed":
        if not probes_bed:
            raise ValueError("probes_bed must be provided when mode='direct_bed'")
        with open(probes_bed, "r") as infile, open(output, "w") as outfile:
            for idx, line in enumerate(infile):
                if not line.strip() or line.startswith("#") or line.startswith("track"):
                    continue
                parts = line.strip().split("\t")
                if len(parts) < 3:
                    continue
                chrom = parts[0]
                start, end = int(parts[1]), int(parts[2])
                
                if direct_window_size > 0:
                    w_start = start
                    w_idx = 1
                    while w_start < end:
                        w_end = min(w_start + direct_window_size, end)
                        pos = random.randint(w_start, w_end - 1) if w_end > w_start else w_start
                        probe_label = f"ProbeLocus_{idx+1}_{chrom}_{start}_{end}_w{w_idx}"
                        outfile.write(f"{chrom}\t{pos}\t{pos+1}\t{probe_label}_direct\n")
                        w_start = w_end
                        w_idx += 1
                else:
                    pos = random.randint(start, end - 1) if end > start else start
                    probe_label = f"ProbeLocus_{idx+1}_{chrom}_{start}_{end}"
                    outfile.write(f"{chrom}\t{pos}\t{pos+1}\t{probe_label}_direct\n")
    else:
        def process_exons(bed_file, outfile, label_prefix):
            with open(bed_file, "r") as infile:
                for idx, line in enumerate(infile):
                    if not line.strip():
                        continue
                    parts = line.strip().split("\t")
                    if len(parts) < 7:
                        continue
                    chrom = parts[0]
                    start, end = int(parts[1]), int(parts[2])
                    cds_start, cds_end = int(parts[3]), int(parts[4])
                    has_left_intron = parts[5] == "True"
                    has_right_intron = parts[6] == "True"
                    exon_label = f"{label_prefix}_{idx+1}_{chrom}_{start}_{end}"
                    if include_intron and has_left_intron:
                        left_snv = max(0, start - random.randint(1, 10))
                        outfile.write(f"{chrom}\t{left_snv}\t{left_snv+1}\t{exon_label}_left_intron\n")
                    if include_cds and cds_end > cds_start:
                        exon_snv = random.randint(cds_start, cds_end - 1)
                        outfile.write(f"{chrom}\t{exon_snv}\t{exon_snv+1}\t{exon_label}_cds\n")
                    if include_intron and has_right_intron:
                        right_snv = end + random.randint(1, 10)
                        outfile.write(f"{chrom}\t{right_snv}\t{right_snv+1}\t{exon_label}_right_intron\n")

        with open(output, "w") as outfile:
            process_exons(fully_bed, outfile, "FullyCovered")
            process_exons(partial_bed, outfile, "PartiallyCovered")
            if include_offtarget:
                with open(unused_bed, "r") as infile:
                    for idx, line in enumerate(infile):
                        if not line.strip():
                            continue
                        parts = line.strip().split("\t")
                        if len(parts) < 3:
                            continue
                        chrom = parts[0]
                        start, end = int(parts[1]), int(parts[2])
                        midpoint = (start + end) // 2
                        probe_label = f"UnusedProbe_{idx+1}_{chrom}_{start}_{end}"
                        outfile.write(f"{chrom}\t{midpoint}\t{midpoint+1}\t{probe_label}_midpoint\n")

    total = int(subprocess.check_output(f"wc -l < {output}", shell=True).strip() or 0)
    log(f"Generated {total} total SNVs.", log_func)
    return output, total


def apply_mutations(read_start, read_length, ref_fasta, chrom, variants_to_apply):
    import pysam
    try:
        ref_seq = ref_fasta.fetch(chrom, read_start, read_start + read_length + 50).upper()
    except Exception:
        return None, None, 0

    cigar = []
    out_seq = []
    nm = 0
    ref_idx = 0
    read_idx = 0

    variants_to_apply = sorted(variants_to_apply, key=lambda x: x["pos"])

    for v in variants_to_apply:
        v_ref_idx = v["pos"] - read_start
        if v_ref_idx < ref_idx:
            continue
        if v_ref_idx > ref_idx:
            dist = min(v_ref_idx - ref_idx, read_length - read_idx)
            if dist > 0:
                out_seq.append(ref_seq[ref_idx:ref_idx + dist])
                cigar.append([0, dist])
                ref_idx += dist
                read_idx += dist
        if read_idx >= read_length:
            break

        ref_len = len(v["ref"])
        alt_len = len(v["alt"])
        if ref_len == 1 and alt_len == 1:
            out_seq.append(v["alt"])
            cigar.append([0, 1])
            ref_idx += 1
            read_idx += 1
            nm += 1
        elif alt_len > ref_len:
            ins_seq = v["alt"]
            avail = read_length - read_idx
            if avail > 0:
                added_seq = ins_seq[:avail]
                out_seq.append(added_seq)
                cigar.append([0, 1])
                if len(added_seq) > 1:
                    cigar.append([1, len(added_seq) - 1])
                ref_idx += 1
                read_idx += len(added_seq)
                nm += len(added_seq) - 1
        elif ref_len > alt_len:
            out_seq.append(v["alt"])
            cigar.append([0, 1])
            cigar.append([2, ref_len - 1])
            ref_idx += ref_len
            read_idx += 1
            nm += ref_len - 1

    if read_idx < read_length:
        rem = read_length - read_idx
        out_seq.append(ref_seq[ref_idx:ref_idx + rem])
        cigar.append([0, rem])

    final_cigar = []
    for op, length in cigar:
        if length > 0:
            if final_cigar and final_cigar[-1][0] == op:
                final_cigar[-1][1] += length
            else:
                final_cigar.append([op, length])

    return "".join(out_seq), [tuple(c) for c in final_cigar], nm


def generate_synthetic_bam(
    work_dir, snvs_bed, fasta_path, depth, vaf, rg_id, rg_sm,
    insert_size, insert_std, indel_interval, read_length=150,
    sequencing_mode="hybrid_capture",
    log_func=None, progress_func=None
):
    import pysam
    from collections import defaultdict

    output_bam = work_dir / "synthetic.bam"
    output_vcf = work_dir / "synthetic.vcf"

    log(f"Loading reference {fasta_path}...", log_func)
    fasta = pysam.FastaFile(str(fasta_path))

    header = {
        "HD": {"VN": "1.0", "SO": "unsorted"},
        "SQ": [{"SN": chrom, "LN": length} for chrom, length in zip(fasta.references, fasta.lengths)],
        "RG": [{"ID": rg_id, "SM": rg_sm, "PL": "ILLUMINA"}],
    }
    chrom_to_tid = {chrom: i for i, chrom in enumerate(fasta.references)}

    groups = defaultdict(list)
    variant_count = 0

    log("Reading variant positions...", log_func)
    with open(snvs_bed, "r") as f:
        for line in f:
            if line.startswith("#") or not line.strip():
                continue
            parts = line.strip().split("\t")
            if len(parts) < 4:
                continue
            chrom = parts[0]
            pos = int(parts[1])
            label = parts[3]
            group_key = "_".join(label.split("_")[:5])
            try:
                ref_base = fasta.fetch(chrom, pos, pos + 1).upper()
            except Exception:
                continue
            if indel_interval > 0 and (variant_count + 1) % indel_interval == 0:
                indel_len = random.randint(1, 4)
                if random.choice(["INS", "DEL"]) == "INS":
                    ins_bases = "".join(random.choices(["A", "C", "G", "T"], k=indel_len))
                    ref_seq, alt_seq = ref_base, ref_base + ins_bases
                else:
                    try:
                        ref_seq = fasta.fetch(chrom, pos, pos + indel_len + 1).upper()
                    except Exception:
                        ref_seq = ref_base
                    alt_seq = ref_base
            else:
                alt_bases = [b for b in ["A", "C", "G", "T"] if b != ref_base]
                alt_seq = random.choice(alt_bases) if alt_bases else "A"
                ref_seq = ref_base

            groups[group_key].append({
                "chrom": chrom,
                "pos": pos,
                "vcf_pos": pos + 1,
                "ref": ref_seq,
                "alt": alt_seq,
            })
            variant_count += 1

    log(f"Generating read clusters across {len(groups)} regions...", log_func)
    total_groups = len(groups)

    with pysam.AlignmentFile(str(output_bam), "wb", header=header) as out_bam, \
         open(output_vcf, "w") as vcf:
        vcf.write("##fileformat=VCFv4.2\n")
        vcf.write('##INFO=<ID=DP,Number=1,Type=Integer,Description="Target Depth">\n')
        vcf.write('##INFO=<ID=AF,Number=A,Type=Float,Description="Target Allele Frequency">\n')
        vcf.write("#CHROM\tPOS\tID\tREF\tALT\tQUAL\tFILTER\tINFO\n")

        for group_idx, (group_key, variants) in enumerate(groups.items()):
            if progress_func:
                progress_func(group_idx / total_groups, f"Processing region {group_idx+1}/{total_groups}")

            chrom = variants[0]["chrom"]
            if chrom not in fasta.references and f"chr{chrom}" in fasta.references:
                chrom = f"chr{chrom}"
                for v in variants:
                    v["chrom"] = chrom
            if chrom not in fasta.references:
                continue

            tid = chrom_to_tid[chrom]
            min_pos = min(v["pos"] for v in variants)
            max_pos = max(v["pos"] for v in variants)
            center = (min_pos + max_pos) // 2
            sigma = max(30.0, (max_pos - min_pos) / 3.0)
            variant_coverage = {v["pos"]: 0 for v in variants}
            generated_pairs = []

            if sequencing_mode == "pcr_amplicon":
                # For PCR amplicon mode, fragments start and end exactly at target coordinates
                parts = group_key.split("_")
                try:
                    target_start = int(parts[3])
                    target_end = int(parts[4])
                except (IndexError, ValueError):
                    target_start = min_pos - 50
                    target_end = max_pos + 50
                
                tlen = target_end - target_start
                fwd_start = target_start
                current_read_length = tlen
                rev_start = fwd_start
                
                # All pairs in this region are identical
                for _ in range(depth):
                    generated_pairs.append({
                        "fwd_start": fwd_start,
                        "rev_start": rev_start,
                        "tlen": tlen,
                        "alt_positions": set(),
                        "read_length": current_read_length,
                    })
            else: # hybrid_capture mode
                max_pairs = depth * 500
                while len(generated_pairs) < max_pairs:
                    if all(cov >= depth for cov in variant_coverage.values()):
                        break
                    frag_center = int(random.gauss(center, sigma))
                    tlen = max(read_length + 10, int(random.gauss(insert_size, insert_std)))
                    frag_start = int(frag_center - tlen / 2)
                    frag_end = frag_start + tlen
                    fwd_start = frag_start
                    rev_start = frag_end - read_length

                    covers_any = False
                    for v in variants:
                        pos = v["pos"]
                        if (fwd_start <= pos < fwd_start + read_length) or (rev_start <= pos < rev_start + read_length):
                            variant_coverage[pos] += 1
                            covers_any = True

                    if covers_any or (min_pos - tlen <= frag_center <= max_pos + tlen):
                        generated_pairs.append({
                            "fwd_start": fwd_start,
                            "rev_start": rev_start,
                            "tlen": tlen,
                            "alt_positions": set(),
                            "read_length": read_length,
                        })

            for v in variants:
                pos = v["pos"]
                covering_pairs = [
                    p for p in generated_pairs
                    if (p["fwd_start"] <= pos < p["fwd_start"] + p.get("read_length", read_length))
                    or (p["rev_start"] <= pos < p["rev_start"] + p.get("read_length", read_length))
                ]
                num_alts = int(len(covering_pairs) * vaf)
                for pair in random.sample(covering_pairs, num_alts):
                    pair["alt_positions"].add(pos)
                if covering_pairs:
                    vcf.write(f"{chrom}\t{v['vcf_pos']}\t.\t{v['ref']}\t{v['alt']}\t.\tPASS\tDP={len(covering_pairs)};AF={vaf}\n")

            for i, pair in enumerate(generated_pairs):
                fwd_start = pair["fwd_start"]
                rev_start = pair["rev_start"]
                tlen = pair["tlen"]
                current_read_length = pair.get("read_length", read_length)
                active = [v for v in variants if v["pos"] in pair["alt_positions"]]

                fwd_seq, fwd_cigar, fwd_nm = apply_mutations(fwd_start, current_read_length, fasta, chrom, active)
                rev_seq, rev_cigar, rev_nm = apply_mutations(rev_start, current_read_length, fasta, chrom, active)
                if not fwd_seq or not rev_seq:
                    continue

                read_name = f"A00979:882:H7JWLDRX7:1:synth_{group_key}:{i}"
                a_fwd = pysam.AlignedSegment()
                a_fwd.query_name = read_name
                a_fwd.query_sequence = fwd_seq
                a_fwd.query_qualities = pysam.qualitystring_to_array("I" * current_read_length)
                a_fwd.reference_id = tid
                a_fwd.reference_start = fwd_start
                a_fwd.mapping_quality = 60
                a_fwd.cigartuples = fwd_cigar
                a_fwd.next_reference_id = tid
                a_fwd.next_reference_start = rev_start
                a_fwd.template_length = tlen
                a_fwd.set_tags([("RG", rg_id), ("MC", f"{current_read_length}M"), ("AS", current_read_length), ("NM", fwd_nm)])

                a_rev = pysam.AlignedSegment()
                a_rev.query_name = read_name
                a_rev.query_sequence = rev_seq
                a_rev.query_qualities = pysam.qualitystring_to_array("I" * current_read_length)
                a_rev.reference_id = tid
                a_rev.reference_start = rev_start
                a_rev.mapping_quality = 60
                a_rev.cigartuples = rev_cigar
                a_rev.next_reference_id = tid
                a_rev.next_reference_start = fwd_start
                a_rev.template_length = -tlen
                a_rev.set_tags([("RG", rg_id), ("MC", f"{current_read_length}M"), ("AS", current_read_length), ("NM", rev_nm)])

                if random.choice([True, False]):
                    a_fwd.flag = 99
                    a_rev.flag = 147
                else:
                    a_fwd.flag = 163
                    a_rev.flag = 83

                out_bam.write(a_fwd)
                out_bam.write(a_rev)

    if progress_func:
        progress_func(1.0, "Sorting and indexing BAM...")

    sorted_bam = work_dir / "synthetic.sorted.bam"
    run_cmd(f"samtools sort {output_bam} -o {sorted_bam}", log_func)
    run_cmd(f"samtools index {sorted_bam}", log_func)
    output_bam.unlink(missing_ok=True)

    log("BAM and VCF generation complete.", log_func)
    return sorted_bam, output_vcf