seq_dir = './example/reads'
!ls -l {seq_dir} | head

"""
Explanation: Steps: 

Get samples
read in sample sheet
guess samples from sequence folder


Get params
filter db per sample
samples for binning
samples for abundance profiling
tmp dir
envs
software
assembler
trimmer
params
atropos
maxbin 
humann2
metaphlan
kraken
shogun?




format yaml

Step 1: get samples
First, we need to read in a set of samples for analysis. 
There are two options for this: guess the sample names, or read them in from a sample sheet or manifest. 
Option 1: guess sample names from sequence output folder.
Enter the correct sequencing directory below, and a list of some example filenames should pop up. Make sure these look correct.
End of explanation
"""


sample_paths = get_sample_paths(seq_dir)

print(yaml.dump(sample_paths, default_flow_style = False))

"""
Explanation: If that looks correct, execute the following cell. It should result in a dictionary associating samples to filepaths.
To validate, a portion of the dictionary is converted to yaml and printed below:
End of explanation
"""


# read in sample sheet

def read_sample_sheet(f, sep='\t'):
    data = False
    data_lines = ''
    for line in f:
        if data:
            if line.startswith('[') or line.strip() == '':
                data = False
                continue
            data_lines += line
        elif line.startswith('[Data]'):
            data = True
            continue
        else:
            continue
            
    data_df = pd.read_csv(StringIO(data_lines), sep = '\t', comment='#')
    return(data_df)

with open('./example/example_sample_sheet.txt', 'r') as f:
    ss_df = read_sample_sheet(f)

# column for read prefixes
# column for sample names
# aggregate into dict

ss_df.head()


seq_dir = './example/reads/'

fps = os.listdir(seq_dir)

files_df = parse_ilm_fps_to_df(fps)

files_df

def get_sample_reads_df_from_sample_sheet(ss_df, seq_dir, sample_col='Description', prefix_col='Sample_ID'):
    sample_reads_dict = {}
    
    samples = list(set(ss_df[sample_col]))
    
    fps = os.listdir(seq_dir)

    files_df = parse_ilm_fps_to_df(fps)
    
    for s in samples:
        # get the subset of sample sheet rows for that sample
        sample_rows = ss_df.loc[ss_df[sample_col] == s,]
        
        f_fps = []
        r_fps = []
        
        # iter across subset table and find file corresponding to each row
        for idx, row in sample_rows.iterrows():
            lane = 'L{0:03d}'.format(row['Lane']) 
            
            f_fps.extend(files_df.loc[(files_df['Sample'] == row['Sample_ID']) & 
                                      (files_df['Lane'] == lane) &
                                      (files_df['Read'] == 'R1'), 'File'].values)
            
            r_fps.extend(files_df.loc[(files_df['Sample'] == row['Sample_ID']) & 
                                      (files_df['Lane'] == lane) &
                                      (files_df['Read'] == 'R2'), 'File'].values)
        
        f_fps = sorted(f_fps)
        r_fps = sorted(r_fps)
        
        sample_reads_dict[s] = {'forward': [os.path.join(seq_dir, x) for x in f_fps],
                                'reverse': [os.path.join(seq_dir, x) for x in r_fps]}
    
    return(sample_reads_dict)

sample_reads_dict = get_sample_reads_df_from_sample_sheet(ss_df, seq_dir,  sample_col='Description', prefix_col='Sample_ID')
print(yaml.dump(sample_reads_dict, default_flow_style = False))

"""
Explanation: Option 2: read sample names and prefixes from sample sheet
This option is in development
End of explanation
"""


config_dict = {}

"""
Explanation: Step 2: define parameters
- filter db per sample
- samples for binning
- samples for abundance profiling
- tmp dir
- envs
- software
- assembler
- trimmer
- params
    * atropos
    * maxbin 
    * humann2
    * metaphlan
    * kraken
    * shogun?
End of explanation
"""


filter_db = '/home/jgsanders/ref_data/genomes/Homo_sapiens_Bowtie2_v0.1/Homo_sapiens'

samples_dict = add_filter_db(sample_paths, filter_db)

print(yaml.dump(samples_dict, default_flow_style = False))

"""
Explanation: Add filter databases per sample
These are the filepaths to the databases used for host filtering. 
The method 'add host filter db' adds the provided filepath to the sample dictionary you created above.
By default, all samples get same database.
You can also provide a list of specific sample names to update if you want to set different filter dbs for some samples. To do this, just re-execute add_filter_db providing a list of samples to update; the existing filter_db paths will remain unchanged if they are not in this list. 
Example:
```
filter_db1 = '/home/jgsanders/ref_data/genomes/Homo_sapiens_Bowtie2_v0.1/Homo_sapiens'
filter_db2 = '/home/jgsanders/ref_data/genomes/mouse/mouse'
samples_dict = add_filter_db(sample_paths, filter_db1)
samples_dict = add_filter_db(samples_dict, filter_db2, samples = ['Vf'])
```
End of explanation
"""


# binning_samples = ['Vf','Bs']
binning_samples = list(samples_dict.keys())

config_dict['binning_samples'] = binning_samples

"""
Explanation: Add samples for binning
These are the samples from which assembled contigs will be binned into putative draft genomes.
By default, bin all samples.
End of explanation
"""


#abundance_samples = ['Vf','Bs']

# This line filters out BLANK* samples from abundance mapping
abundance_samples = [x for x in samples_dict if not x.startswith('BLANK')]

# This line uses all samples for abundance mapping
abundance_samples = list(samples_dict.keys())

config_dict['abundance_samples'] = abundance_samples

"""
Explanation: Add samples for abundance profiling
These are the samples from which abundance information will be used for binning of binning_samples
By default, use all samples.
End of explanation
"""


config_dict['trimmer'] = 'atropos'

"""
Explanation: Add trimmer to use
These are the read trimmers to use for qc.
Currently, the available options are:
- atropos
- skewer
End of explanation
"""


assemblers = ['megahit','metaspades']

config_dict['assemblers'] = assemblers

"""
Explanation: Add assemblers to use
These are the assemblers to use for assembly.
Currently, the available options are:
- spades
- metaspades
- megahit
End of explanation
"""


config_dict['mapping_assembler'] = 'metaspades'

"""
Explanation: Add assembly to use for binning
This is the assembler to use for binning.
Currently, the available options are:
- spades
- metaspades
- megahit
End of explanation
"""


tmp_dir_root = '/localscratch'

config_dict['tmp_dir_root'] = tmp_dir_root

"""
Explanation: Add temporary directory
This is the directory used during execution of rules. Should ideally be local scratch storage.
End of explanation
"""


config_dict['env'] = 'source activate snakemake_assemble'
config_dict['anvi_env'] = 'source activate sn_anvio'
config_dict['metaphlan_env'] = 'source activate humann2'
config_dict['humann2_env'] = 'source activate humann2'
config_dict['kraken_env'] = 'source activate kraken'
config_dict['shogun_env'] = 'source activate shogun'

"""
Explanation: Add environments
These are the conda environments to use for various rules. 
env is the primary snakemake_assemble environment. Change others as necessary for other rules.
End of explanation
"""


software = {'snakemake_folder': '/home/jgsanders/git_sw/snakemake_assemble',
            'seqtk': 'seqtk'}

config_dict['software'] = software

"""
Explanation: Add software paths
Some software binaries are specified directly. 
Also, the installation folder for the snakemake_assemble repository is provided as snakemake_folder
End of explanation
"""


resources =  {'centrifuge_base': '/home/jgsanders/miniconda/envs/anvio2/centrifuge',
              'centrifuge_models': '/home/jgsanders/miniconda/envs/anvio2/centrifuge/b+h+v/b+h+v'}
config_dict['resources'] = resources

"""
Explanation: anvio resources
Sets paths to resources for anvio
End of explanation
"""


params = {}

"""
Explanation: Add parameters definitions
These are parameters passed to the various tools used.
End of explanation
"""


params['atropos'] = ' -a GATCGGAAGAGCACACGTCTGAACTCCAGTCAC -A GATCGGAAGAGCGTCGTGTAGGGAAAGGAGTGT -q 15 --minimum-length 100 --pair-filter any'

"""
Explanation: Atropos params
Sets quality and adapter trimming parameters. 
Some suggested defaults: 
Nextera: -a CTGTCTCTTATACACATCT -A CTGTCTCTTATACACATCT -q 15 --minimum-length 100 --pair-filter any
KapaHyperPlus: -a GATCGGAAGAGCACACGTCTGAACTCCAGTCAC -A GATCGGAAGAGCGTCGTGTAGGGAAAGGAGTGT -q 15 --minimum-length 100 --pair-filter any
End of explanation
"""


params['maxbin'] = '-plotmarker'

"""
Explanation: maxbin params
Sets binning parameters.
End of explanation
"""


params['kraken'] = {'kraken_db': '/home/qiz173/Databases/Kraken/stdb'}

"""
Explanation: Kraken params
Sets parameters for Kraken taxonomic profiling.
End of explanation
"""


params['shogun'] = "--utree_indx /home/qiz173/Databases/SHOGUN/annotated/utree/stdb.ctr"

"""
Explanation: SHOGUN params
Sets parameters for Shogun taxonomic profiling.
End of explanation
"""


params['metaphlan'] = {'metaphlan_dir': '/home/jgsanders/git_sw/metaphlan2'}

"""
Explanation: MetaPhlAn2 params
Sets parameters for MetaPhlAn2 taxonomic profilling.
End of explanation
"""


humann2 = {'humann2_aa_db': 'humann2_aa_db: /home/jgsanders/ref_data/humann2/uniref',
           'humann2_nt_db': '/home/jgsanders/ref_data/humann2/chocophlan',
           'metaphlan_dir': '/home/jgsanders/share/metaphlan2',
           'norms': ['cpm','relab'],
           'other': ''}

params['humann2'] = humann2

"""
Explanation: humann2
Sets various parameters for HUMAnN2 funcitonal profiling. 
humann2_aa_db: Amino acid database for translated amino acid search
humann2_nt_db: ChocoPhlAn database for nucleotide search
metaphan_dir: path to metaphlan2 directory install, should also have the default metaphlan2 db
norms: normalizations for which to output tables
other: any other HUMAnN2 parameters to use
End of explanation
"""


params['cram'] = ''
params['skewer'] = '-x ./adapters/HyperPlus.fa -n -l 100 -m any -q 15'

"""
Explanation: other
Various other legacy parameters that I need to fix
End of explanation
"""


config_dict['params'] = params
config_dict['samples'] = samples_dict

"""
Explanation: Step 3: format yaml
Now the entire configuration dictionary can be formatted as a yaml file and saved for use in the pipeline.
End of explanation
"""


config_yaml = yaml.dump(config_dict, default_flow_style = False)

with open('config.yaml', 'w') as f:
    f.write(config_yaml)

"""
Explanation: Format as yaml and write to filepath:
End of explanation
"""


print(config_yaml)

"""
Explanation: This is the complete config yaml:
End of explanation
"""


