__init__.py

#
# Copyright (C) 2012 INRA
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.
#

import os
import glob
import sys

from ng6.ng6workflow import NG6Workflow
from ng6.utils import Utils

class RADseq (NG6Workflow):

    ENZYMES = {
              'sbfI' : {
                        'rad'    : 'CCTGCAGG',
                        'radtag' : 'TGCAGG'
            }
              ,'ecoRI' : {
                        'rad'    : 'GAATTC',
                        'radtag' : 'AATTC'
	        }
	          ,'taqI' : {
                        'rad'    : 'TCGA',
                        'radtag' : 'GCA'
	        }
    }


    def get_enzyme (self, name):
        if name not in self.ENZYMES :
            raise ValueError("The enzyme name " + str(name) + " does not exists. Accepted names are " + str (list(self.ENZYMES.keys())))

        return (self.ENZYMES[name]['rad'], self.ENZYMES[name]['radtag'])

    def get_name(self):
        return 'radseq'

    def get_description(self):
        return "RADseq data quality analysis workflow"

    def define_parameters(self, function="process"):

        # if metadata.barcode faire splitbc. Sinon liste read1 = liste des échantillons.

        # split_bc
        self.add_multiple_parameter_list("pool", "Sequence pool", required = False)
        self.add_parameter("id", "Pool uniq identifier", add_to = "pool")
        # required = True que si pool précisé
        #TODO check rules
        self.add_input_file("read1", "Read 1 fastq file path", rules="RequiredIf?ALL[poll!=None]", add_to = "pool")
        self.add_input_file("read2", "Read 2 fastq file path", add_to = "pool")
        # options bougées dans process_radtag
        #self.add_parameter("enzyme", "Restriction enzyme name", required = True , choices = list(self.ENZYMES.keys()))
        #self.add_parameter("enzyme2", "Restriction enzyme 2 name", required = False, choices = list(self.ENZYMES.keys()))
        # option inutile maintenant, puisque potentiel rescue et check RAD dans process_radtag
        #self.add_parameter("tag_mismatch", "Max. number of mismatches allowed in the radTAG sequence",default = 0, type=int)
        self.add_parameter("mismatches", "Max. number of mismatches allowed.",default = 0, type=int)
        self.add_parameter("trim_barcode", "Should the barecode be trimmed",default = True, type=bool, rules="Exclude=trim_reads2")
        self.add_parameter("trim_reads2", " Shoud the read 2 be trimmed to have the same length as the read1",default = False, type=bool)

        #clone_filter
        self.add_parameter("dereplicate", "Should we remove PCR duplcat. Only for paired end - single RADSeq protocol",default = False, type=bool)

        # process_radtag
        self.add_parameter("enzyme", "Restriction enzyme name", required = True , choices = list(self.ENZYMES.keys()))
        self.add_parameter("enzyme2", "Restriction enzyme 2 name", required = False, default="", choices = list(self.ENZYMES.keys()))
        self.add_parameter("uncall_remove", "clean data, remove any read with an uncalled base.", default=True, type=bool)
        self.add_parameter("discard_low_qual", "discard reads with low quality scores.", default=True, type=bool)
        self.add_parameter("rescue_radtag", "rescue barcodes and RAD-Tags.", default=False, type=bool)
        self.add_parameter("max_length", "truncate final read length to this value. (default none)", required = False, type=int)
        self.add_parameter("quality_encode", "specify how quality scores are encoded, 'phred33' (Illumina 1.8+, Sanger, default) or 'phred64' (Illumina 1.3 - 1.5).", choices = ['phred33', 'phred64'], default='phred33')
        self.add_parameter("keep_discard_read", "capture discarded reads to a file.", default = True, type=bool)
        self.add_parameter("window_size", "set the size of the sliding window as a fraction of the read length, between 0 and 1 (default 0.15).", default=0.15, type=float)
        self.add_parameter("limit_score", "set the score limit. If the average score within the sliding window drops below this value, the read is discarded (default 10).", default=10, type=int)

        # ustacks
        self.add_parameter("min_cov", "stacks minimum coverage", type=int)
        self.add_parameter("primary_mismatch", "stacks max mismatch", type=int)
        self.add_parameter("secondary_mismatch", "2nd read max mismatch", type=int)
        self.add_parameter("uncall_hap", "uncall hap with 2nd read", type=bool)
        self.add_parameter("disable_removal_algo", "removal cluster correction algo", type=bool)
        self.add_parameter("disable_deleveraging_algo", "deleveraging cluster correction algo", type=bool)
        self.add_parameter("max_locus", "max stacks per cluster", type=int)

        # cstacks
        self.add_parameter("batch_id", " batch_id of this radseq analysis",default = 1, type=int)
        self.add_parameter("catalog_mismatches", "How many mismatches allowed between sample to generate the cluster catalog",default = 1, type=int)

        # ng6
        self.add_parameter("compression", "How should data be compressed once archived (none|gz|bz2)",default = 'gz', choices=['none', 'gz', 'bz2'])
        self.add_parameter("databank", "Which databank should be used to seek contamination (as to be nr databank)")


    def process(self):

        read1_list=[]
        read2_list= []

        # plus besoins on fait le check du rad dans process_radtag, sinon ça ne fonctionne pas dans le cas des double digest.
        # rad, rad_tag = self.get_enzyme(self.enzyme)
        # si aucun pool, alors données démultipléxées et donc les sample read1 sont les vraies sample read1
        if not self.pool:
            for s in self.samples:
                for r in s.reads1 :
                    read1_list.append(r)
                for r in s.reads2:
                    read2_list.append(r)
        # si démultiplexage alors définition de pool, et les sample read1 sont égaux au pool read1
        else:
            # group all individual by pool
            pools = {}
            for p in self.pool :
                if p["id"] in pools :
                    raise ValueError("Duplicated pool id." + p['id'])
                pools[p["id"]] = (p, [])

            for sample in self.samples :
                pool_id = sample.get_metadata('pool_id')
                if pool_id not in pools:
                    raise ValueError("The pool id " + pool_id + " does not exists in (individual " + sample.name + ")")
                pools[pool_id][1].append(sample)

            # prepare fastq files
            fastq_files_1 = []
            fastq_files_2 = []
            barcode_files = []

            for pool_id, data in pools.items() :
                p = data[0]
                samples = data[1]
                tmp_barcode = self.get_temporary_file()
                barcode_files.append(tmp_barcode)
                fastq_files_1.append(p['read1']);
                if p['read2'] :
                    fastq_files_2.append(p['read2'])

                # write barcode file
                with open(tmp_barcode, "w") as ff:
                    for sample in samples :
                        ff.write(sample.name + "\t" + sample.get_metadata('barcode') +"\n")

            #splitbc = self.add_component("Splitbc", [ fastq_files_1, barcode_files, fastq_files_2, rad, rad_tag, self.mismatches, self.tag_mismatch, self.trim_barcode, self.trim_reads2])
            splitbc = self.add_component("Splitbc", [ fastq_files_1, barcode_files, fastq_files_2, self.mismatches, self.trim_barcode, self.trim_reads2])
            read1_list=splitbc.output_read1
            read2_list=splitbc.output_read2

        # PROCESS_RADTAG
        print("before process",read1_list,read2_list)
        process_radtag = self.add_component("ProcessRadtag", [read1_list, read2_list, self.dereplicate, self.enzyme, self.enzyme2, self.window_size, self.limit_score, self.quality_encode, self.max_length, self.uncall_remove, self.discard_low_qual, self.rescue_radtag, self.keep_discard_read ])
"""

#         # USTACKS
#         # creation des listes pour conserver l'ordre entre read1 read2 et sample id
#         samples_id = []
#         read1_files = []
#         read2_files = []
#         # check non empty and MAJ sample_id
#         for idx, read1 in enumerate(process_radtag.output_read_1):
#             if os.path.getsize(read1) > 0:
#                 read1_files.append(read1)
#                 read2_files.append(process_radtag.output_read_2[idx])
#                 samples_id.append(idx)
#
        ustacks_opt={"samples_id":samples_id,"read1_files":read1_list}

        if self.min_cov :
            ustacks_opt["min_cov"] = self.min_cov
        if self.primary_mismatch :
            ustacks_opt["primary_mismatch"] = self.primary_mismatch
        if self.secondary_mismatch :
            ustacks_opt["secondary_mismatch"] =  self.secondary_mismatch
        if self.uncall_hap :
            ustacks_opt["uncall_hap"] = self.uncall_hap
        if self.disable_removal_algo :
            ustacks_opt["removal_algo"] = False
        if self.disable_deleveraging_algo :
            ustacks_opt["deleveraging_algo"] = False
        if self.max_locus :
            ustacks_opt["max_locus"] = self.max_locus

        ustacks = self.add_component("Ustacks", [],ustacks_opt)

        # CSTACKS
#         cstacks = self.add_component("Cstacks", [ustacks.alleles, ustacks.snps, ustacks.tags, self.batch_id, self.catalog_mismatches])

"""