import argparse
import os
import subprocess


def main():

    parser = argparse.ArgumentParser(
                description='Run Jupyter notebook, create Markdown output')
    parser.add_argument('nb', help='Name of Jupyter notebook.')
    parser.add_argument('md_output', help='Name of created Markdown output.')
    args = parser.parse_args()

    nb = args.nb
    if not (os.path.isfile(nb) and os.path.splitext(nb)[-1] == '.ipynb'):
        raise IOError(f"not a valid Jupyter notebook: {nb}")

    md_output = args.md_output
    if os.path.splitext(md_output)[-1] not in {'.md', '.markdown'}:
        raise IOError(f"not a valid Markdown file: {md_output}")

    if os.path.isfile(md_output):
        os.remove(md_output)

    subprocess.check_call(['jupyter', 'nbconvert',
                           '--to', 'notebook',
                           '--execute',
                           '--inplace',
                           '--ExecutePreprocessor.timeout=-1',
                           nb,
                           ])

    subprocess.check_call(['jupyter', 'nbconvert',
                           '--output-dir', os.path.dirname(md_output),
                           '--output', os.path.basename(md_output),
                           '--to', 'markdown',
                           nb,
                           ])


if __name__ == '__main__':
    main()

run:
    def path(f):
        """Get path relative to `summary_dir`."""
        return os.path.relpath(f, config['summary_dir'])
    with open(output.summary, 'w') as f:
        f.write(textwrap.dedent(f"""
        # Summary

        Analysis run by [Snakefile]({path(workflow.snakefile)})
        using [this config file]({path(workflow.configfiles[0])}).
        See the [README in the top directory]({path('README.md')})
        for details.

        Here is the rule graph of the computational workflow:
        ![{path(input.rulegraph)}]({path(input.rulegraph)})

        Here is the Markdown output of each notebook in the workflow:

        1. Get codon-variant-table from [here]({config['codon_variant_table_url']}).

        2. Count variants and then
           [aggregate counts]({path(input.aggregate_variant_counts)}) to create
           to create [variant counts file]({path(input.variant_counts)}).

        3. [Analyze sequencing counts to cells ratio]({path(input.counts_to_cells_ratio)});
           this prints a list of any samples where this ratio too low. Also
           creates [a CSV]({path(input.counts_to_cells_csv)}) with the
           sequencing counts, number of sorted cells, and ratios for
           all samples.

        4. [Escape scores from variant counts]({path(input.counts_to_scores)}).

        5. [Escape fractions for mutations and homologs]({path(input.scores_to_frac_escape)});
           creating [mutation escape fraction file]({path(input.escape_fracs)})
           and [homolog escape fraction file]({path(input.escape_fracs_homologs)}).

        6. [Call sites of strong escape]({path(input.call_strong_escape_sites)}),
           and write to [a CSV file]({path(input.strong_escape_sites)}).

        7. Plot [escape profiles]({path(input.escape_profiles)}).

        9. [Multidimensional scaling]({path(input.mds_escape_profiles)}) on escape profiles.

        10. Map escape profiles to ``*.pdb`` files using [this notebook]({path(input.output_pdbs)})

        11. Output a list of RBD structural contacts for each antibody with a high-resolution structure using [this notebook]({path(input.structural_contacts)})

        12. [Count mutations in GISAID RBD sequences]({path(input.gisaid_rbd_mutations)}).
           to create [this counts file]({path(input.gisaid_mutation_counts)}).

        13. [Analyze GISAID mutations at sites of escape]({path(input.natural_mutations)}).

        14. [Escape from binding by homologs]({path(input.homolog_escape)}).

        15. [Analyze viral escape selections]({path(input.escape_selections)}).

        16. [Make supplementary data files]({path(input.make_supp_data)}),
            which are [here]({path(config['supp_data_dir'])}). These include
            `dms-view` input files.

        17. Make custom plots for this project that fall out of the core pipeline, summarzied in [this notebook]({path(input.custom_plots)})

        """
        ).strip())

shell:
    """
    R -e \"rmarkdown::render(input=\'{params.nb}\')\";
    mv {params.md} {output.md};
    mv {params.md_files} {output.md_files}
    """

shell:
    """
    R -e \"rmarkdown::render(input=\'{params.nb}\')\";
    mv {params.md} {output.md}
    """

shell:
    "python scripts/run_nb.py {params.nb} {output.nb_markdown}"

shell:
    "python scripts/run_nb.py {params.nb} {output.nb_markdown}"

shell:
    "python scripts/run_nb.py {params.nb} {output.nb_markdown}"

shell:
    "python scripts/run_nb.py {params.nb} {output.nb_markdown}"

shell:
    "python scripts/run_nb.py {params.nb} {output.nb_markdown}"

shell:
    "python scripts/run_nb.py {params.nb} {output.nb_markdown}"

shell:
    "python scripts/run_nb.py {params.nb} {output.nb_markdown}"

shell:
    "python scripts/run_nb.py {params.nb} {output.nb_markdown}"

shell:
    "python scripts/run_nb.py {params.nb} {output.nb_markdown}"

shell:
    "python scripts/run_nb.py {params.nb} {output.nb_markdown}"

shell:
    "python scripts/run_nb.py {params.nb} {output.nb_markdown}"

shell:
    "python scripts/run_nb.py {params.nb} {output.nb_markdown}"

shell:
    "python scripts/run_nb.py {params.nb} {output.nb_markdown}"

run:
    # parse sample and library from `sample_lib` wildcard
    lib = params.sample_lib.split('_')[-1]
    sample = params.sample_lib[: -len(lib) - 1]
    assert sample == (barcode_runs_expandR1
                      .set_index('sample_lib')
                      .at[params.sample_lib, 'sample']
                      )
    assert lib == (barcode_runs_expandR1
                   .set_index('sample_lib')
                   .at[params.sample_lib, 'library']
                   )
    # initialize `CodonVariantTable` (used to get valid barcodes)
    wt_seqrecord = Bio.SeqIO.read(input.wt_seq, 'fasta')
    geneseq = str(wt_seqrecord.seq)
    primary_target = wt_seqrecord.name
    variants=dms_variants.codonvarianttable.CodonVariantTable(
                geneseq=geneseq,
                barcode_variant_file=input.variant_table,
                substitutions_are_codon=True,
                substitutions_col='codon_substitutions',
                primary_target=primary_target)
    # initialize `IlluminaBarcodeParser`
    parser = dms_variants.illuminabarcodeparser.IlluminaBarcodeParser(
                valid_barcodes=variants.valid_barcodes(lib),
                **config['illumina_barcode_parser_params'])
    # parse barcodes
    counts, fates = parser.parse(input.r1s,
                                 add_cols={'library': lib,
                                           'sample': sample})
    # write files
    counts.to_csv(output.counts, index=False)
    fates.to_csv(output.fates, index=False)

run:
    urllib.request.urlretrieve(config['codon_variant_table_url'],
                               output.codon_variant_table)

run:
    urllib.request.urlretrieve(config['mut_bind_expr_url'], output.file)

run:
    urllib.request.urlretrieve(config['variant_expr_url'], output.file)

run:
    urllib.request.urlretrieve(config['variant_bind_url'], output.file)