""" UMCCR post-bcbio patient analysis workflow
"""
import os
from os.path import join, abspath, dirname, isfile, basename, splitext
from ngs_utils.file_utils import splitext_plus, verify_file, verify_dir
from ngs_utils.bcbio import BcbioProject
from ngs_utils.utils import flatten
from ngs_utils.logger import critical, info, debug
from ngs_utils import logger as ngs_utils_logger
from hpc_utils.hpc import find_loc, get_ref_file
from ngs_utils.reference_data import get_key_genes_bed

shell.executable(os.environ.get('SHELL', 'bash'))
shell.prefix("")
threads_max = 32  # Use up to 32 cores at once, if available


loc = find_loc()
is_spartan = loc and loc.name == 'spartan'
is_raijin = loc and loc.name == 'raijin'
is_hpc = is_spartan or is_raijin
upload_proxy = ''
if is_spartan:
    upload_proxy = 'HTTPS_PROXY=http://wwwproxy.unimelb.edu.au:8000 '

upload_igv = is_hpc and config.get('no_s3', 'no') == 'no'

within_sub_workflow = 'within_sub_workflow' in config
config['within_sub_workflow'] = 'yes'  # to avoid redundant logging in cluster sub-executions of the Snakefile


def prep_inputs(config_):

    ###############################
    #### Parsing bcbio project ####
    ###############################

    # Parsing bcbio project and including/excluding samples
    include_names = config_.get('batch') or config_.get('sample')
    if include_names:
        include_names = str(include_names).split(',')
        include_names = [v for v in flatten([sn.split('__') for sn in include_names])]  # support "batch__sample" notation
    exclude_names = config_.get('exclude')
    if exclude_names:
        exclude_names = str(exclude_names).split(',')
        exclude_names = [v for v in flatten([sn.split('__') for sn in exclude_names])]  # support "batch__sample" notation

    ngs_utils_logger.is_silent = within_sub_workflow  # to avoid redundant logging in cluster sub-executions of the Snakefile
    run = BcbioProject(config_.get('bcbio_project', abspath(os.getcwd())),
                       exclude_samples=exclude_names,
                       include_samples=include_names)
    ngs_utils_logger.is_silent = False
    run.project_name = splitext(basename(run.bcbio_yaml_fpath))[0]

    if len(run.batch_by_name) == 0:
        if exclude_names:
            critical(f'Error: no samples left with the exclusion of batch/sample name(s): {", ".join(exclude_names)}.'
                     f'Check yaml file for available options: {run.bcbio_yaml_fpath}.')
        if include_names:
            critical(f'Error: could not find a batch or a sample with the name(s): {", ".join(include_names)}. '
                     f'Check yaml file for available options: {run.bcbio_yaml_fpath}')
        critical(f'Error: could not parse any batch or samples in the bcbio project. '
                 f'Please check the bcbio yaml file: {run.bcbio_yaml_fpath}.')

    # Batch objects index by tumor sample names
    batches = [b for b in run.batch_by_name.values() if not b.is_germline() and b.tumor and b.normal]
    assert batches

    batch_by_name = {b.name + '__' + b.tumor.name: b for b in batches}

    ############################################
    #### Reference files provided directly? ####
    ############################################
    ref_fa = config_.get('ref_fasta')
    truth_regions = config_.get('truth_regions')
    pon_dir = config_.get('panel_of_normals')
    pcgr_installation = config_.get('pcgr')

    ##################################################
    #### Bcbio reference genomes folder provided? ####
    ##################################################
    bcbio_genomes_dir = config_.get('bcbio_genomes')
    if bcbio_genomes_dir:
        # Looking for reference fasta
        for fp in [
            join(bcbio_genomes_dir, f'Hsapiens/{run.genome_build}/seq/{run.genome_build}.fa'),
            join(bcbio_genomes_dir, f'{run.genome_build}/seq/{run.genome_build}.fa'),
            join(bcbio_genomes_dir, f'seq/{run.genome_build}.fa')]:
            if isfile(fp):
                ref_fa = fp
        if not ref_fa:
            critical(f'Not found {run.genome_build}.fa in bcbio genomes directory {bcbio_genomes_dir}')

        # Looking for truth regions
        for fp in [
            join(bcbio_genomes_dir, f'Hsapiens/{run.genome_build}/validation/giab-NA12878/truth_regions.bed'),
            join(bcbio_genomes_dir, f'{run.genome_build}/validation/giab-NA12878/truth_regions.bed'),
            join(bcbio_genomes_dir, f'validation/giab-NA12878/truth_regions.bed')]:
            if isfile(fp):
                truth_regions = fp
        if not truth_regions:
            critical(f'Not found GiaB truth regions in bcbio genomes directory {bcbio_genomes_dir}')

    #################################################################
    #### Reference files not provided, but we are on known host? ####
    #################################################################
    if not ref_fa:
        ref_fa = get_ref_file(run.genome_build)
    if not truth_regions:
        truth_regions = get_ref_file(run.genome_build, ['truth_sets', 'giab', 'bed'])
    if not pon_dir:
        pon_dir = get_ref_file(run.genome_build, 'panel_of_normals_dir')
    if not pcgr_installation and loc:
        pcgr_installation = loc.pcgr_dir

    ###########################################################
    #### Done looking for refernece files, now some checks ####
    ###########################################################
    truth_regions = verify_file(truth_regions, is_critical=True, description='GiaB truth regions')

    ref_fa = verify_file(ref_fa, is_critical=True, description='Reference fasta')
    verify_file(ref_fa + '.fai', is_critical=True, description='Reference fasta fai index')

    if pon_dir:
        verify_dir(pon_dir, 'Panel of normals directory', is_critical=True)
        verify_file(join(pon_dir, 'panel_of_normals.snps.vcf.gz'), is_critical=True, description='Panel of normals SNPs file in user provided folder')
        verify_file(join(pon_dir, 'panel_of_normals.snps.vcf.gz.tbi'), is_critical=True, description='Please index panel of normal files with tabix')
        verify_file(join(pon_dir, 'panel_of_normals.indels.vcf.gz'), is_critical=True, description='Panel of normals indels file in user provided folder')
        verify_file(join(pon_dir, 'panel_of_normals.indels.vcf.gz.tbi'), is_critical=True, description='Please index panel of normal files with tabix')

    if pcgr_installation:
        verify_dir(pcgr_installation, is_critical=True, description='PCGR installation directory')
        verify_dir(join(pcgr_installation, 'data'), is_critical=True, description='PCGR data directory')
        if not within_sub_workflow:
            info('PCGR is available')

    ####################################
    #### Also loading key genes bed ####
    ####################################
    key_genes_bed = get_key_genes_bed(run.genome_build)

    return run, batch_by_name, ref_fa, truth_regions, pon_dir, key_genes_bed, pcgr_installation


run, batch_by_name, ref_fa, truth_regions, pon_dir, key_genes_bed, pcgr_installation = prep_inputs(config)
GERMLINE_SUFFIX = '-germline'
if any(isfile(join(run.date_dir, batch_by_name[b].normal.name + '-ensemble-annotated.vcf.gz'))
       for b in batch_by_name.keys()):
    GERMLINE_SUFFIX = ''


rule all:
    input: 'log/umccrised.done'
    # A trick to avoid duplicating all input paths in the top "all" rule which has to be defined on top.

# TODO: try subworkflows here? http://snakemake.readthedocs.io/en/stable/snakefiles/modularization.html#sub-workflows
"""
subworkflow small_variants:
    workdir: 'small_variants'
    snakefile: 'Snakefile.small_variants'

rule all:
    input:  small_variants('log/small_variants.done')
    output: ...
    shell:  ...
"""
# Or maybe it's not reasonable and not doable here since the input file is a phony .done, and also we depend on config in subworkflows


include: "small_variants.smk"
include: "coverage.smk"
include: "structural.smk"
include: "igv.smk"
include: "pcgr.smk"
include: "rmd.smk"


localrules: multiqc, copy_logs, umccrise


rule multiqc:  # {}
    input:
        join(run.date_dir, 'multiqc/multiqc_report.html')
    output:
        run.project_name + '-multiqc_report.html'
    shell:
        'cp {input} {output}'
        # generate proper mutliqc_umccrise



## Additional information
# TODO: link it to MultiQC
rule copy_logs:  # {}
    input:
        versions = join(run.date_dir, 'data_versions.csv'),
        programs = join(run.date_dir, 'programs.txt'),
        conf_dir = run.config_dir
    output:
        versions = 'log/' + run.project_name + '-data_versions.csv',
        programs = 'log/' + run.project_name + '-programs.txt',
        conf_dir = directory(join('log/' + run.project_name + '-config'))
    shell:
        'cp -r {input.versions} {output.versions} && ' \
        'cp -r {input.programs} {output.programs} && ' \
        'cp -r {input.conf_dir} {output.conf_dir}'


rule umccrise:
    input:  # Copy here inputs of the "milestone" rules (rules without output defined in the end of each Snakemake.* file)
        rules.multiqc.output,
        rules.copy_logs.output,
        rules.coverage.output,
        rules.structural.output,
        rules.small_variants.output,
        rules.rmd.output,
        (rules.pcgr.output if pcgr_installation else rules.pcgr_prep.output),
        rules.igv.output
    output:
        temp(touch('log/umccrised.done'))
