configfile: "config/config.yaml"

rule all:
    input:
        "logs/pseudobulk_pilot_summary.json",
        "logs/expression_pilot_summary.json",
        "config/gene_universe_summary.json",
        "data/processed/response/response_consolidation_summary.json",
        "data/processed/response/m1_response_audit.json",
        "data/processed/basal/basal_normalization_summary.json",

rule freeze_version:
    output:
        manifest="bootstrap/hf_file_manifest.parquet",
        version="bootstrap/dataset_version.json",
    shell:
        "uv run tahoe-benchmark freeze-version --config config/config.yaml --output-dir bootstrap"

rule pilot_pseudobulk:
    input:
        manifest="bootstrap/hf_file_manifest.parquet",
        config="config/config.yaml",
    output:
        summary="logs/pseudobulk_pilot_summary.json",
        detail="logs/pseudobulk_pilot_shards.csv",
    shell:
        "uv run tahoe-benchmark pilot-pseudobulk --config {input.config} --manifest {input.manifest} --raw-dir data/raw --output-dir logs"

rule pilot_expression:
    input:
        manifest="bootstrap/hf_file_manifest.parquet",
        config="config/config.yaml",
    output:
        summary="logs/expression_pilot_summary.json",
        detail="logs/expression_pilot_shards.csv",
    shell:
        "uv run tahoe-benchmark pilot-expression --config {input.config} --manifest {input.manifest} --raw-dir data/raw --output-dir logs"

rule gene_universe:
    input:
        manifest="bootstrap/hf_file_manifest.parquet",
        config="config/config.yaml",
    output:
        summary="config/gene_universe_summary.json",
        universe="config/gene_universe.csv",
        genes="config/gene_universe.txt",
    shell:
        "uv run python -m tahoe_axis_benchmark.gene_universe --config {input.config} --manifest {input.manifest} --raw-dir data/raw --checkpoint-dir data/checkpoints/gene_stats --output-dir config"

rule response_parts:
    input:
        manifest="bootstrap/hf_file_manifest.parquet",
        config="config/config.yaml",
        universe="config/gene_universe.csv",
        plate_map="config/development_plate_map.csv",
        scope="config/cell_line_scope.csv",
    output:
        summary="data/interim/response_parts/response_parts_summary.json",
    shell:
        "uv run python -m tahoe_axis_benchmark.response --config {input.config} --manifest {input.manifest} --gene-universe {input.universe} --plate-map {input.plate_map} --cell-scope {input.scope} --raw-dir data/raw --parts-dir data/interim/response_parts"

rule consolidate_response:
    input:
        summary="data/interim/response_parts/response_parts_summary.json",
    output:
        condition="data/processed/response/condition_response.parquet",
        repeat="data/processed/response/repeat_response.parquet",
        qc="data/processed/response/condition_qc.parquet",
        summary="data/processed/response/response_consolidation_summary.json",
    shell:
        "uv run python -m tahoe_axis_benchmark.consolidate_response --parts-dir data/interim/response_parts --output-dir data/processed/response"

rule audit_response:
    input:
        condition="data/processed/response/condition_response.parquet",
        repeat="data/processed/response/repeat_response.parquet",
        qc="data/processed/response/condition_qc.parquet",
        plate_map="config/development_plate_map.csv",
        scope="config/cell_line_scope.csv",
    output:
        summary="data/processed/response/m1_response_audit.json",
        condition_audit="data/processed/response/condition_key_audit.csv",
        repeat_audit="data/processed/response/repeat_key_audit.csv",
        coverage="data/processed/response/drug_context_coverage.csv",
    shell:
        "uv run python -m tahoe_axis_benchmark.audit_response --condition {input.condition} --repeat {input.repeat} --qc {input.qc} --plate-map {input.plate_map} --scope {input.scope} --output-dir data/processed/response"

rule aggregate_full_dmso:
    input:
        manifest="bootstrap/hf_file_manifest.parquet",
        config="config/config.yaml",
        universe="config/gene_universe.csv",
    output:
        summary="data/interim/basal/full_dmso_summary.json",
        tokens="data/interim/basal/frozen_gene_token_map.parquet",
        cells="data/interim/basal/dmso_cell_counts.parquet",
    shell:
        "uv run python -m tahoe_axis_benchmark.full_basal --config {input.config} --manifest {input.manifest} --gene-universe {input.universe} --raw-dir data/raw --checkpoint-dir data/checkpoints/dmso_parts --output-dir data/interim/basal"

rule normalize_basal:
    input:
        summary="data/interim/basal/full_dmso_summary.json",
        tokens="data/interim/basal/frozen_gene_token_map.parquet",
        cells="data/interim/basal/dmso_cell_counts.parquet",
    output:
        summary="data/processed/basal/basal_normalization_summary.json",
        basal="data/processed/basal/basal_context.parquet",
        plate="data/processed/basal/dmso_plate_logcpm.parquet",
        tmm="data/processed/basal/dmso_tmm_factors.parquet",
    shell:
        "uv run python -m tahoe_axis_benchmark.normalize_basal --counts-dir data/interim/basal/dmso_counts_by_plate --token-map {input.tokens} --cell-counts {input.cells} --output-dir data/processed/basal"
