# selexprep benchmark scaffolding
#
# Per-accession pipeline: fetch -> detect, then aggregate -> figure_a.
# Reads ``ground_truth.tsv``; processes only ``verified=true`` rows.
#
# Implementation notes:
#  1. ``verified=false`` rows are excluded from the DAG entirely
#     (they also get a stderr warning from the metrics aggregator).
#  2. Only ENA/SRA/DDBJ-fetchable accessions appear here. Non-accession
#     sources (figshare/zenodo/local-FASTQ) are out of scope.
#  3. ``round_map_source == "curated"`` rows use the curator's
#     ``round_map_path`` (resolved relative to ``ground_truth.tsv``)
#     and pass ``--allow-manual-review`` to ``selexprep fetch`` so the
#     FASTQs download even when ENA metadata is too sparse for the
#     auto round-parser.
#
# Reproducibility:
#
#   snakemake -s benchmarks/Snakefile --cores 1
#
# Or dry-run:
#
#   snakemake -s benchmarks/Snakefile --cores 1 --dry-run

from pathlib import Path

import pandas as pd

# Anchor paths to this Snakefile's directory so both invocation styles work:
#
#   snakemake -s benchmarks/Snakefile ...
#   cd benchmarks && snakemake -s Snakefile ...
#
# This mirrors the audit.smk path discipline.
_BENCHMARKS_DIR = Path(workflow.basedir).resolve()
GROUND_TRUTH = str(_BENCHMARKS_DIR / "ground_truth.tsv")
OUTROOT = str(_BENCHMARKS_DIR / "results")

# filter to verified-only rows for the Snakemake DAG.
_gt = pd.read_csv(GROUND_TRUTH, sep="\t", dtype=str).fillna("")
_gt = _gt[_gt["verified"].str.strip().str.lower() == "true"]
ACCESSIONS = sorted(_gt["accession"].tolist())
ROUND_MAP_SOURCE = dict(zip(_gt["accession"], _gt["round_map_source"].replace("", "auto")))
ROUND_MAP_PATH = dict(zip(_gt["accession"], _gt["round_map_path"]))

# resolve curated round-map paths relative to
# ``ground_truth.tsv``'s directory so Snakemake invocations from any
# cwd still find the files.
_GROUND_TRUTH_DIR = _BENCHMARKS_DIR


rule all:
    input:
        OUTROOT + "/table_1.md",
        OUTROOT + "/metrics.json",


rule fetch:
    output:
        # snakemake 7's Snakefile preprocessor errors on wildcards buried
        # inside an f-string (``f"{OUTROOT}/{{accession}}/..."``); plain
        # string concatenation works in any snakemake version. Snakemake 9
        # tolerated the f-string form but the bench extra now pins
        # ``snakemake<8`` (Python 3.10 compat with pulp<2.8).
        marker=OUTROOT + "/{accession}/fetch_metadata.json",
        # explicit FASTQ manifest, written AFTER fetch completes.
        # Replaces the old ``Path.glob`` input function in rule detect, which
        # Snakemake evaluated before fetch ran → empty file list on a fresh
        # run. The manifest is the single source of truth for which FASTQs
        # detect consumes as the primary stream. R2 mates are recorded in a
        # separate manifest and passed via ``selexprep detect --paired-r2`` so
        # paired split-primer datasets are evaluated as paired-end instead of
        # being forced into R1-only partial recovery.
        manifest=OUTROOT + "/{accession}/fastqs.manifest",
        r2_manifest=OUTROOT + "/{accession}/fastqs.r2.manifest",
    params:
        # Pass --allow-manual-review on curated rows so fetch downloads
        # FASTQs even when ENA metadata is too sparse to auto-assign
        # rounds. NONE-confidence runs land in round_unknown/ and the
        # curator's round_map_path routes them by basename. The runner
        # gates the all-unassigned refusal on this flag, so curated rows
        # whose runs are ALL unassigned still download.
        extra_flags=lambda w: (
            "--allow-manual-review" if ROUND_MAP_SOURCE[w.accession] == "curated" else ""
        ),
    shell:
        # partial-fetch fix. Snakemake runs
        # shell: under ``set -euo pipefail``, so a bare ``fetch; rc=$?``
        # would abort AT the fetch line on any non-zero exit (rc never
        # captured). Scope errexit OFF only around the fetch with
        # ``set +e`` / ``set -e`` (pipefail + nounset stay on), capture
        # the code, then gate on evidence: require BOTH a non-empty
        # manifest AND a non-empty fetch_metadata.json. A genuine partial
        # download (some runs missing, e.g. PRJEB70964 17/27) proceeds
        # with the rc surfaced to stderr; a total fetch crash (empty
        # manifest or missing metadata) still fails the rule.
        # if/fi only — bash brace groups collide with Snakemake's {}.
        """
        set +e
        selexprep fetch {wildcards.accession} --outdir {OUTROOT}/{wildcards.accession} {params.extra_flags}
        rc=$?
        set -e
        find {OUTROOT}/{wildcards.accession} -name '*.fastq.gz' ! -name '*_2.fastq.gz' | sort > {output.manifest}
        find {OUTROOT}/{wildcards.accession} -name '*_2.fastq.gz' | sort > {output.r2_manifest}
        if [ ! -s {output.manifest} ] || [ ! -s {output.marker} ]; then echo "fetch failed for {wildcards.accession} (rc=$rc): empty manifest or missing fetch_metadata.json" >&2; exit 1; fi
        if [ "$rc" -ne 0 ]; then echo "fetch rc=$rc for {wildcards.accession} (partial download; proceeding)" >&2; fi
        """


def _round_map_for(accession: str) -> str:
    """Resolve which rounds.tsv ``selexprep detect`` should consume.

    - ``auto`` rows: fetch's auto-emitted ``rounds.tsv``.
    - ``curated`` rows: the curator-supplied TSV path from
      ``ground_truth.tsv``, resolved relative to that file's directory.
    """
    if ROUND_MAP_SOURCE[accession] == "curated":
        return str((_GROUND_TRUTH_DIR / ROUND_MAP_PATH[accession]).resolve())
    return OUTROOT + "/" + accession + "/rounds.tsv"


rule detect:
    input:
        # depend on the manifest (an output of rule fetch),
        # not a Path.glob input function. This guarantees fetch ran first
        # AND that the FASTQ list reflects what's actually on disk. The old
        # glob was evaluated by Snakemake before fetch executed → empty
        # fastq_args on a fresh run (blocker).
        manifest=OUTROOT + "/{accession}/fastqs.manifest",
        r2_manifest=OUTROOT + "/{accession}/fastqs.r2.manifest",
    output:
        OUTROOT + "/{accession}/library_report.json",
    params:
        round_map=lambda w: _round_map_for(w.accession),
    shell:
        # if/fi guard (NOT ``|| { ... }``): bash brace groups collide with
        # Snakemake's {} placeholder syntax. Fail fast + readably if fetch
        # produced no FASTQs rather than calling detect with no arguments.
        "if [ ! -s {input.manifest} ]; then "
        "echo 'empty manifest: fetch produced no FASTQs for {wildcards.accession}' >&2; "
        "exit 1; fi; "
        "r2_args=''; "
        "if [ -s {input.r2_manifest} ]; then "
        "while IFS= read -r fq; do r2_args=\"$r2_args --paired-r2 $fq\"; done < {input.r2_manifest}; "
        "fi; "
        "selexprep detect $(cat {input.manifest}) $r2_args "
        "--round-map {params.round_map} "
        "--outdir {OUTROOT}/{wildcards.accession}"


rule compute_metrics:
    input:
        ground_truth=GROUND_TRUTH,
        reports=expand(OUTROOT + "/{acc}/library_report.json", acc=ACCESSIONS),
    output:
        OUTROOT + "/metrics.json",
    shell:
        "python -m selexprep.benchmark.metrics "
        "--ground-truth {input.ground_truth} "
        "--reports-dir {OUTROOT} "
        "--out {output}"


rule figure_a:
    input:
        metrics=OUTROOT + "/metrics.json",
        ground_truth=GROUND_TRUTH,
    output:
        OUTROOT + "/table_1.md",
    shell:
        "python -m selexprep.benchmark.figure_a "
        "--metrics {input.metrics} --ground-truth {input.ground_truth} --outdir {OUTROOT}"
