# Obtain run_ids from centroided DIA mzXML files in dia_data folder
run_ids, = glob_wildcards("../data_dia/{run}.mzXML")

rule all:
    input:
        "tric_feature_alignment.done"

rule irt_spectrast_consensus:
    input:
        "../results/library/rtn-{run}_Q1.splib"
    output:
        name = touch("../results/openswath/irt_{run}/consensus"),
        splib = "../results/openswath/irt_{run}/consensus.splib"
    shell:
        "spectrast -cQ20 -cN{output.name} -cAC {input}"

rule irt_spectrast_filter:
    input:
        fasta = "../data_library/library.fasta",
        splib = rules.irt_spectrast_consensus.output.splib
    output:
        name = touch("../results/openswath/irt_{run}/proteotypic"),
        mrm = "../results/openswath/irt_{run}/proteotypic.mrm"
    shell:
        "spectrast -cN{output.name} -cM -cu -cd -cD{input.fasta} {input.splib}"

rule irt_generate_swath_assays:
    input:
        rules.irt_spectrast_filter.output.mrm
    output:
        "../results/openswath/irt_{run}/assays.pqp"
    shell:
        "OpenSwathAssayGenerator -in {input} -out {output}"

rule irt_filter_swath_assays:
    input:
        rules.irt_generate_swath_assays.output
    params:
        bins = 10,
        peptides = 5
    output:
        "../results/openswath/irt_{run}.pqp"
    script:
        "scripts/hrirt.R"

rule irt_nonlinear_filter_swath_assays:
    input:
        rules.irt_generate_swath_assays.output
    params:
        bins = 1000,
        peptides = 5
    output:
        "../results/openswath/irt_nonlinear_{run}.pqp"
    script:
        "scripts/hrirt.R"

rule openswath:
    input:
        irt_linear=rules.irt_filter_swath_assays.output,
        irt_nonlinear=rules.irt_nonlinear_filter_swath_assays.output,
        mzxml="../data_dia/{run}.mzXML",
    output:
        osw="../results/openswath/{run}.osw",
        cache=temp(directory("/tmp/{run}/"))
    threads: 4
    resources:
        mem_mb=lambda wildcards, attempt: attempt * 8192
    shell:
        "OpenSwathWorkflow -in {input.mzxml} -tr ../data_library/library.pqp -tr_irt {input.irt_linear} -tr_irt_nonlinear {input.irt_nonlinear} -out_osw {output.osw} -threads {threads} -min_upper_edge_dist 1 -mz_extraction_window 30 -mz_extraction_window_unit ppm -mz_extraction_window_ms1 20 -mz_extraction_window_ms1_unit ppm -mz_correction_function regression_delta_ppm -use_ms1_traces -irt_mz_extraction_window 50 -irt_mz_extraction_window_unit ppm -rt_extraction_window 600 -RTNormalization:estimateBestPeptides -RTNormalization:alignmentMethod lowess -RTNormalization:outlierMethod none -Scoring:stop_report_after_feature 5 -Scoring:TransitionGroupPicker:compute_peak_quality false -Scoring:Scores:use_ms1_mi -Scoring:Scores:use_mi_score -threads {threads} -readOptions cache -tempDirectory {output.cache} -batchSize 1000 -ms1_isotopes 3 -Scoring:max_transitions 6 -Scoring:min_transitions 3 -Scoring:max_ms1_isotopes 4 -Scoring:min_ms1_isotopes 2" # -enable_uis_scoring

rule pyprophet_model:
    input:
        expand("../results/openswath/{run}.osw", run=run_ids)
    params:
        subsample_ratio = lambda wildcards, input: round(1.0/len(input),3)
    output:
        "../results/pyprophet/runspecific/model.osw",
    shell:
        "pyprophet merge --subsample_ratio={params.subsample_ratio} --out {output} {input} && "
        "pyprophet score --in {output} --level=ms1ms2"
        # "pyprophet score --in {output} --level=ms1 && "
        # "pyprophet score --in {output} --level=transition"

rule pyprophet_apply:
    input:
        osw = rules.openswath.output.osw,
        model = rules.pyprophet_model.output
    output:
        osw = "../results/pyprophet/runspecific/{run}.osw",
        oswr = "../results/pyprophet/runspecific/{run}.oswr"
    shell:
        "pyprophet score --in {input.osw} --out {output.osw} --apply_weights {input.model} --level=ms1ms2 && "
        # "pyprophet score --in {output.osw} --apply_weights {input.model} --level=ms1 && "
        # "pyprophet score --in {output.osw} --apply_weights {input.model} --level=transition && "
        # "pyprophet ipf --no-ipf_ms1_scoring --in {output.osw} && "
        "pyprophet reduce --in {output.osw} --out {output.oswr}"

rule pyprophet_global:
    input:
        template = rules.pyprophet_model.output,
        oswr = expand("../results/pyprophet/runspecific/{run}.oswr", run=run_ids)
    output:
        "../results/pyprophet/global/model.oswr"
    shell:
        "pyprophet merge --template {input.template} --out {output} {input.oswr} && "
        "pyprophet peptide --context=global --in {output} && "
        "pyprophet protein --context=global --in {output}"

rule pyprophet_backpropagate:
    input:
        osw = rules.pyprophet_apply.output.osw,
        model = rules.pyprophet_global.output
    output:
        "../results/pyprophet/integrated/{run}.osw"
    shell:
        "pyprophet backpropagate --apply_scores {input.model} --in {input.osw}  --out {output} && "
        "pyprophet export --in {output}"

rule tric_prepare:
    input:
        rules.pyprophet_backpropagate.output
    output:
        "../results/tric/{run}.mzXML.tsv"
    shell:
        "pyprophet export --in={input} && mv $(basename {output}) {output}"

rule tric_feature_alignment:
    input:
        expand("../results/tric/{run}.mzXML.tsv", run=run_ids)
    output:
        long="../results/tric/feature_alignment.tsv",
        matrix="../results/tric/feature_alignment_matrix.tsv",
        check=temp(touch("tric_feature_alignment.done"))
    shell:
        "feature_alignment.py --in {input} --out {output.long} --out_matrix {output.matrix} --method LocalMST --realign_method lowess_cython --max_rt_diff 60 --mst:useRTCorrection True --mst:Stdev_multiplier 3.0 --target_fdr 0.01 --max_fdr_quality 0.05"

