def get_model_name(wildcards):
    config_param=f'{wildcards.model}_model_name'
    return config[config_param]

def get_model_requirement(wildcards):
    config_param=f'{wildcards.model}_model_name'
    return config['dorado_model_dir']+'/'+config[config_param]

def get_benchmarking_file(wildcards):
    config_param = f'{wildcards.model}_benchmarks_file'
    if config[config_param]:
        return f'--batchsize-benchmarks-file {config[config_param]}'
    return ''

def get_qscore(wildcards):
    config_param = f'{wildcards.model}_min_q'
    if config[config_param]:
        return f'--min-qscore {config[config_param]}'
    return ''

def get_batch_names(input_file):
    with open(input_file,'rt') as ih:
        return [line.strip().split('/')[-1].split('.')[0] for line in ih]

rule download_pod5:
    output:
        '../data/raw_pod5/PBK98658_853a956f_57f83f46_{batch}.pod5'
    threads: 1
    wildcard_constraints:
        batch="\d+"
    shell:
        """
        aws s3 cp --no-sign-request s3://ont-open-data/nomiss_96BC_P2I_SUP_2026/raw/pod5/PBK98658_853a956f_57f83f46_{wildcards.batch}.pod5 {output}
        """

rule basecall_pod5:
    input:
        pod5='../data/raw_pod5/PBK98658_853a956f_57f83f46_{batch}.pod5',
        model=get_model_requirement
    output:
        '../data/basecalled_reads/{model}/PBK98658_853a956f_57f83f46_{batch}.fastq.gz'
    threads:
        32
    resources:
        gpu=1
    params:
        benchmarking=get_benchmarking_file,
        min_qscore=get_qscore,
        dorado_model=get_model_name
    wildcard_constraints:
        batch="\d+",
        model="hac|fast"
    shell:
        """
        dorado basecaller --models-directory {config[dorado_model_dir]} --emit-fastq {params.benchmarking} {params.min_qscore} {params.dorado_model} {input.pod5} > {output}
        """

rule concatenate_basecalled_fastq:
    input:
        expand('../data/basecalled_reads/{{model}}/PBK98658_853a956f_57f83f46_{batch}.fastq.gz',batch=range(1,config["pod5_dataset_size"]+1,config["pod5_stride"]))
    output:
        '../data/basecalled_reads/{model}.fastq.gz'
    threads: 1
    wildcard_constraints:
        model="hac|fast"
    shell:
        """
        zcat {input} > {output}
        """