# Python standard library
from os.path import join
from os import listdir
import os, sys, re, json, glob

# 3rd party imports from pypi
from snakemake.workflow import workflow as wf_api
from snakemake.utils import R

# Local imports
from scripts.common import (
    allocated,
    provided,
    references,
    str_bool
)
from scripts.ambient import validate_ambient_options, validate_soupx_import

#Global workflow variables
configfile: 'config.json'
ambient_rna_method = validate_ambient_options(config['options'])
if ambient_rna_method == 'soupx':
    for sample_dir in config['options']['input']:
        if os.path.isdir(sample_dir):
            validate_soupx_import(sample_dir)
tmpdir = config['options']['tmp_dir']
pipeline = config['options']['pipeline']
CELLRANGER = config['options']['cellranger']
inputs = config['options']['input']
samples  = config['samples']
workpath = config['project']['workpath']
genome = config['options']['genome']
features = config['options']['features']       # Features files for cellranger (not used in all pipelines)
libraries = config['options']['libraries']     # Libraries files for cellranger (not used in all pipelines)
CMO_REF = config['options']['cmo_reference']     # CMO Reference files for cellranger (only used in multi pipeline)
CMO_SAMPLE = config['options']['cmo_sample']     # CMO Sample files for cellranger (only used in multi pipelines)
HTO_SAMPLE = config['options']['hto_sample']     # HTO Sample files for cellranger (only used in multi pipelines)
OCM_SAMPLE = config['options']['ocm_sample']     # OCM Sample files for cellranger (only used in multi pipelines)
PROBE_SAMPLE = config['options']['probe_sample']     # Probe Sample files for cellranger (only used in multi pipelines)
PROBE_SET = config['options']['probe_set']     # Probe Set file for cellranger (only used in multi pipelines)
RENAME = config['options']['rename']            # File containing how to rename samples when running Cell Ranger analysis
FORCECELLS = config['options']['forcecells']                # File containing the cell counts to use for running Cell Ranger analysis
CHAIN = config['options']['chain']                # Chain type for cellranger (only used in vdj pipeline)
exclude_introns = str_bool(                           # Use introns for pre mRNA,
    config['options']['exclude_introns']              # default: False
)
create_bam = str_bool(                           # Create BAM files during Cell Ranger analysis,
    config['options']['create_bam']              # default: False
)
aggr = config['options']['aggregate']
input_fastq = config['options']['input']
filter_file = config['options']['filter']	# Filter threshold file for QC analysis (not used in all pipelines)
METADATA_FILE = config['options']['metadata']	# Metadata file for QC analysis (not used in all pipelines)
if 'libraries' in config:
    lib_samples = list(config['libraries'].keys()) # Libraries file samples
else:
    lib_samples = samples # Handling the situation where cellranger outputs is used as input and no libraries file is provided
pipeline_output = []


RENAME_DICT = dict() #Dictionary containing information on renamed samples where the keys are the FASTQ file names and the values are the Cell Ranger sample names
if RENAME != 'None':
    with open(RENAME) as f:
        tabs = [i.lower() for i in next(f).strip().split(',')]
        index_fastq = [i for i in range(len(tabs)) if 'fastq' in tabs[i]][0]
        index_sample = [i for i in range(len(tabs)) if 'name' in tabs[i]][0]
        for line in f:
            line = line.strip().split(',')
            RENAME_DICT[line[index_fastq]] = line[index_sample]
    samples = list(set(RENAME_DICT.values()))
#    samples = list((set(samples) - set(RENAME_DICT.keys())).union(RENAME_DICT.values()))
    samples.sort()

CELLCOUNT_DICT = dict() #Dictionary containing information on renamed samples where the keys are the FASTQ file names and the values are the Cell Ranger sample names
CELLCOUNT_LIBRARIES = False
if FORCECELLS != 'None':
    with open(FORCECELLS) as f:
        tabs = [i.lower() for i in next(f).strip().split(',')]
        index_cell = tabs.index('cells')
        index_sample = tabs.index('sample')
        if "name" in tabs:
            index_library = tabs.index('name')
            CELLCOUNT_LIBRARIES = True
            for line in f:
                line = line.strip().split(',')
                values = CELLCOUNT_DICT.get(line[index_library], list())
                values.append((line[index_sample], line[index_cell]))
                CELLCOUNT_DICT[line[index_library]] = values
        else:
            for line in f:
                line = line.strip().split(',')
                #if line[index_sample] not in CELLCOUNT_DICT:
                CELLCOUNT_DICT[line[index_sample]] = line[index_cell]

# Import rules
include: join("rules", "common.smk")

if pipeline == 'cite':
    include: join("rules", "cite.smk")
if pipeline == "gex":
    include: join("rules", "gex.smk")
if pipeline == 'multi':
    include: join("rules", "multi.smk")
if pipeline == 'atac':
    include: join("rules", "atac.smk")
if pipeline == 'vdj':
    include: join("rules", "vdj.smk")
if pipeline == 'multiome':
    include: join("rules", "multiome.smk")

# Final output files of the pipeline (fully definited within the pipeline smk files)
rule all:
    input:
        pipeline_output
