from os.path import join, expandvars
from joblib import Parallel, delayed
from glob import glob
from os import system
from tax_credit.framework_functions import (parameter_sweep,
                                            generate_per_method_biom_tables,
                                            move_results_to_repository)


project_dir = expandvars("$HOME/Desktop/projects/tax-credit")
analysis_name= "mock-community"
data_dir = join(project_dir, "data", analysis_name)

reference_database_dir = expandvars("$HOME/Desktop/ref_dbs/")
results_dir = expandvars("$HOME/Desktop/projects/mock-community/")

"""
Explanation: Data generation: using python to sweep over methods and parameters
This notebook demonstrates taxonomy classification using vsearch followed by consensus assignment in QIIME2's q2-feature-classifier.
Environment preparation
End of explanation
"""


dataset_reference_combinations = [
 ('mock-1', 'gg_13_8_otus'), # formerly S16S-1
 ('mock-2', 'gg_13_8_otus'), # formerly S16S-2
 ('mock-3', 'gg_13_8_otus'), # formerly Broad-1
 ('mock-4', 'gg_13_8_otus'), # formerly Broad-2
 ('mock-5', 'gg_13_8_otus'), # formerly Broad-3
 ('mock-6', 'gg_13_8_otus'), # formerly Turnbaugh-1
 ('mock-7', 'gg_13_8_otus'), # formerly Turnbaugh-2
 ('mock-8', 'gg_13_8_otus'), # formerly Turnbaugh-3
 ('mock-9', 'unite_20.11.2016_clean_fullITS'), # formerly ITS1
 ('mock-10', 'unite_20.11.2016_clean_fullITS'), # formerly ITS2-SAG
 ('mock-12', 'gg_13_8_otus'), # Extreme
 ('mock-13', 'gg_13_8_otus_full16S_clean'), # kozich-1
 ('mock-14', 'gg_13_8_otus_full16S_clean'), # kozich-2
 ('mock-15', 'gg_13_8_otus_full16S_clean'), # kozich-3
 ('mock-16', 'gg_13_8_otus'), # schirmer-1
 ('mock-18', 'gg_13_8_otus'),
 ('mock-19', 'gg_13_8_otus'),
 ('mock-20', 'gg_13_8_otus'),
 ('mock-21', 'gg_13_8_otus'),
 ('mock-22', 'gg_13_8_otus'),
 ('mock-23', 'gg_13_8_otus'),
 ('mock-24', 'unite_20.11.2016_clean_fullITS'),
 ('mock-25', 'unite_20.11.2016_clean_fullITS'),
 ('mock-26-ITS1', 'unite_20.11.2016_clean_fullITS'),
 ('mock-26-ITS9', 'unite_20.11.2016_clean_fullITS'),
]

reference_dbs = {'gg_13_8_otus_clean' : (join(reference_database_dir, 'gg_13_8_otus/99_otus_clean_515f-806r.qza'),
                                         join(reference_database_dir, 'gg_13_8_otus/taxonomy/99_otu_taxonomy.qza')),
                 'gg_13_8_otus' : (join(reference_database_dir, 'gg_13_8_otus/rep_set/99_otus_515f-806r_trim250.qza'), 
                                   join(reference_database_dir, 'gg_13_8_otus/taxonomy/99_otu_taxonomy.qza')),
                 'gg_13_8_otus_full16S_clean' : (join(reference_database_dir, 'gg_13_8_otus/99_otus_clean.qza'), 
                                                 join(reference_database_dir, 'gg_13_8_otus/taxonomy/99_otu_taxonomy.qza')),
                 'gg_13_8_otus_full16S' : (join(reference_database_dir, 'gg_13_8_otus/rep_set/99_otus.qza'), 
                                           join(reference_database_dir, 'gg_13_8_otus/taxonomy/99_otu_taxonomy.qza')),
                 'unite_20.11.2016_clean_fullITS' : (join(reference_database_dir, 'sh_qiime_release_20.11.2016/developer/sh_refs_qiime_ver7_99_20.11.2016_dev_clean.qza'), 
                                   join(reference_database_dir, 'sh_qiime_release_20.11.2016/developer/sh_taxonomy_qiime_ver7_99_20.11.2016_dev_clean.qza')),
                 'unite_20.11.2016_clean' : (join(reference_database_dir, 'sh_qiime_release_20.11.2016/developer/sh_refs_qiime_ver7_99_20.11.2016_dev_clean_ITS1Ff-ITS2r.qza'), 
                                             join(reference_database_dir, 'sh_qiime_release_20.11.2016/developer/sh_taxonomy_qiime_ver7_99_20.11.2016_dev.qza')),
                 'unite_20.11.2016' : (join(reference_database_dir, 'sh_qiime_release_20.11.2016/developer/sh_refs_qiime_ver7_99_20.11.2016_dev_ITS1Ff-ITS2r_trim250.qza'), 
                                       join(reference_database_dir, 'sh_qiime_release_20.11.2016/developer/sh_taxonomy_qiime_ver7_99_20.11.2016_dev.qza'))}


"""
Explanation: Preparing data set sweep
First, we're going to define the data sets that we'll sweep over. The following cell does not need to be modified unless if you wish to change the datasets or reference databases used in the sweep.
End of explanation
"""


method_parameters_combinations = {
              'vsearch' : {'p-maxaccepts': [1, 10, 100],
                           'p-perc-identity': [0.80, 0.90, 0.97, 0.99],
                           'p-min-consensus': [0.51, 0.75, 0.99]}
             }

"""
Explanation: Preparing the method/parameter combinations and generating commands
Now we set the methods and method-specific parameters that we want to sweep. Modify to sweep other methods. Note how method_parameters_combinations feeds method/parameter combinations to parameter_sweep() in the cell below.
End of explanation
"""


command_template = "mkdir -p {0}; qiime feature-classifier vsearch --i-query {1} --o-classification {0}/rep_seqs_tax_assignments.qza --i-reference-reads {2} --i-reference-taxonomy {3} {5}; qiime tools export {0}/rep_seqs_tax_assignments.qza --output-dir {0}"
        
commands = parameter_sweep(data_dir, results_dir, reference_dbs,
                           dataset_reference_combinations,
                           method_parameters_combinations, command_template,
                           infile='rep_seqs.qza', output_name='rep_seqs_tax_assignments.qza')


"""
Explanation: Now enter the template of the command to sweep, and generate a list of commands with parameter_sweep().
Fields must adhere to following format:
                  {0} = output directory
                  {1} = input data
                  {2} = reference sequences
                  {3} = reference taxonomy
                  {4} = method name
                  {5} = other parameters
End of explanation
"""


print(len(commands))
commands[0]

"""
Explanation: As a sanity check, we can look at the first command that was generated and the number of commands generated.
End of explanation
"""


Parallel(n_jobs=4)(delayed(system)(command) for command in commands)

"""
Explanation: Finally, we run our commands.
End of explanation
"""


taxonomy_glob = join(results_dir, '*', '*', '*', '*', 'taxonomy.tsv')
generate_per_method_biom_tables(taxonomy_glob, data_dir)

"""
Explanation: Generate per-method biom tables
Modify the taxonomy_glob below to point to the taxonomy assignments that were generated above. This may be necessary if filepaths were altered in the preceding cells.
End of explanation
"""


precomputed_results_dir = join(project_dir, "data", "precomputed-results", analysis_name)
method_dirs = glob(join(results_dir, '*', '*', '*', '*'))
move_results_to_repository(method_dirs, precomputed_results_dir)

"""
Explanation: Move result files to repository
Add results to the tax-credit directory (e.g., to push these results to the repository or compare with other precomputed results in downstream analysis steps). The precomputed_results_dir path and methods_dirs glob below should not need to be changed unless if substantial changes were made to filepaths in the preceding cells.
End of explanation
"""


