#!/usr/bin/env python3

from snakemake.remote.HTTP import RemoteProvider as HTTPRemoteProvider
from pathlib import Path

bbmap = 'docker://quay.io/biocontainers/bbmap:39.01--h5c4e2a8_0'
ham = 'docker://quay.io/biocontainers/hifiasm_meta:hamtv0.3.1--h5b5514e_0'

test_data_url = (
    'https://zenodo.org/record/5908204/files/'
    'zymoD6331std-ecoli-ten-percent.fq.gz')

HTTP = HTTPRemoteProvider()
test_input = HTTP.remote(
    test_data_url,
    keep_local=True)

# try different subsets of different sizes to find one that doesn't segfault
rule target:
    input:
        expand('test_pipeline/{subset}.{redundancy}.{bits}/asm.a_ctg.gfa',
               redundancy=['high', 'normal'],
               # add 'full' to run with the whole dataset
               subset=['14.1', '42.1', '42.10', '42.2'],
               bits=['0', '37'])

rule hifiasm_meta:
    input:
        'data/zymoD6331std-ecoli-ten-percent.{subset}.fq.gz'
    output:
        'test_pipeline/{subset}.{redundancy}.{bits}/asm.a_ctg.gfa',
        'test_pipeline/{subset}.{redundancy}.{bits}/asm.p_ctg.gfa'
    params:
        prefix = directory('test_pipeline/{subset}.{redundancy}.{bits}/asm'),
        force = lambda wildcards:
            '--force-rs' if wildcards.redundancy == 'high' else ''
    threads:
        workflow.cores
    log:
        'test_pipeline/logs/hifiasm_meta.{subset}.{redundancy}.{bits}.log'
    benchmark:
        'test_pipeline/benchmarks/hifiasm_meta.{subset}.{redundancy}.{bits}.log'
    container:
        ham
    shell:
        'hifiasm_meta '
        '-f {wildcards.bits} ' # f 0 disables bloom filter to same RAM
        '-t {threads} '
        '{params.force} '
        '-o {params.prefix} '
        '{input} '
        '&> {log}'

rule subset_reads:
    input:
        'data/zymoD6331std-ecoli-ten-percent.full.fq.gz'
    output:
        'data/zymoD6331std-ecoli-ten-percent.{seed}.{multiplier}.fq.gz'
    params:
        bases = lambda wildcards:
            str(2500000 * int(wildcards.multiplier))
    log:
        'test_pipeline/logs/subset_reads.{seed}.{multiplier}.log'
    threads:
        4
    container:
        bbmap
    shell:
        'reformat.sh '
        'in={input} '
        'out={output} '
        'sampleseed={wildcards.seed} '
        'samplebasestarget={params.bases} '
        'zl=9 '
        '2> {log}'

rule get_input:
    input:
        test_input
    output:
        'data/zymoD6331std-ecoli-ten-percent.full.fq.gz'
    shell:
        'cp {input} {output}'
