'''
@file: CDhitExe.py

@author: Jaume Bonet
@mail:   jaume.bonet@gmail.com
@date:   2013

@ [oliva's lab](http://sbi.imim.es)

@class: CDhitExe
'''
import os

from SBI.external       import ExternalExe
from SBI.external.CDhit import CDhitList  as CHL
from SBI.sequence       import Fasta
from SBI.beans          import Executable
from SBI                import SBIglobals as SBIg


class CDhitExe(ExternalExe):
    '''
    Executes CD-hit program over a FASTA file.

    '''
    def __init__(self, fasta, threshold, step = None, memory = None,
                 output_dir = None, execute = True):
        '''
        @param:    fasta
        @pdef:     input FASTA file
        @ptype:    {String}

        @param:    threshold
        @pdef:     maximum allowed homology. Between 0 and 1
        @ptype:    {Float}

        @param:    step
        @pdef:     execute a partial clustering to speed the process.
                   Between 0 and 1
        @ptype:    {Float}

        @param:    memory
        @pdef:     maximum memory allowed to use. default is 400 Mb
        @ptype:    {Integer}

        @param:    output_dir
        @pdef:     directory to print the output to
        @ptype:    {String}

        @param:    execute
        @pdef:     flag. automatically execute when creating the object.
        @pdefault: _True_
        @ptype:    {Boolean}
        '''

        fasta = os.path.abspath(fasta)

        if CDhitExe._EXE is None:
            self._set_default_executable('cd-hit')

        self._input     = fasta
        self._outdir    = os.path.split(fasta)[0] if output_dir is None else output_dir
        self._output    = self._build_output_file(fasta, threshold)
        self._threshold = str(threshold)
        self._word      = self._set_word_size(threshold)
        self._memory    = memory
        self._step      = step
        if step is not None:
            self._step_file = self._build_output_file(fasta, step)

        if execute:
            self._execute()

        self._cdhitlist = None

    ##############
    # ATTRIBUTES #
    ##############
    @property
    def output_file(self):
        '''
        Name of the cd-hit output file.

        @return: {String}
        '''
        return self._output + '.clstr'

    @property
    def master_fasta_file(self):
        '''
        Name of the fasta output file.

        @return: {String}
        '''
        return self._output + '.fa'

    ###########
    # METHODS #
    ###########
    def compare(self, fasta_file, output_file):
        '''
        Using the same threshold as before, creates a list of those proteins in
        fasta_file similar to the first input and a fasta file with the
        remaining proteins.

        @param:    fasta_file
        @pdef:     name of the Fasta file
        @ptype:    {String}

        @param:    output_file
        @pdef:     prefix of the output file
        @ptype:    {String}

        @return: {List} of {String}: list file name, fasta file name
        '''
        e = Executable(executable = self._EXE.executable + '-2d',
                       path       = self._EXE.path)
        e.add_attribute(self._input,      '-i')
        e.add_attribute(fasta_file,       '-i2')
        e.add_attribute(output_file,      '-o')
        e.add_attribute(self._threshold,  '-c')
        e.add_attribute('1',              '-g')
        e.add_attribute(self._word,       '-n')
        if self._memory is not None:
            e.add_attribute(self._memory, '-M')
        try:
            e.execute()
            e.clean_command()
        except SystemError, er:
            SBIg.throw(self, 'Some error occurred while executing cd-hit\n{0}\n'.format(er),
                       SystemError)

    def make_master_fasta(self, output_fasta = None, force = None):
        '''
        Directly create a fasta file from the master sequences of each
        cluster of a list.

        @param:    output_fasta
        @pdef:     name of the Fasta file
        @pdefault: input fasta + threshold + '.fa'
        @ptype:    {String}

        @param:    force
        @pdef:     overwrite previous files with the same name
        @pdefault: _SBIglobals.overwrite_
        @ptype:    {Boolean}

        @return: {Fasta}
        '''
        if output_fasta is None:
            output_fasta = self.master_fasta_file
        clist = self.parse()
        prots = set()
        fs    = Fasta(self._input)
        for c in clist.clusters:
            prots.add(c.master.name)
        seq = fs.retrieve(prots)
        return Fasta.build_multifasta(output_fasta, seq, force = force)

    def parse(self):
        '''
        Read the cd-hit output automatically generated by the execution.

        @return: {CDhitList}
        '''
        if self._cdhitlist is not None:
            return self._cdhitlist

        if self._step is None:
            self._cdhitlist = self.parse_cdhit_output(self.output_file)
        else:
            clist = self.parse_cdhit_output(self._step_file + '.clstr')
            clist.merge_clusters(self.output_file)
            self._cdhitlist = clist
        return self._cdhitlist

    @staticmethod
    def parse_cdhit_output(cdhit_file):
        '''
        Parse a given cd-hit output.

        @param:    cdhit_file
        @pdef:     name of the cd-hit output file
        @ptype:    {String}
        '''
        return CHL.CDhitList(cdhit_file)

    @staticmethod
    def dynamic(executable, path):
        '''
        Manually set the values for cd-hit path and executable .

        @param:    executable
        @pdef:     name of the executable file
        @ptype:    {String}

        @param:    path
        @pdef:     path to the executable file
        @ptype:    {String}
        '''
        CDhitExe._set_dynamic_executable(executable, path)

    ###################
    # PRIVATE METHODS #
    ###################
    def _build_output_file(self, infile, threshold):
        output_file = os.path.split(infile)[1] + '.' + str(threshold).replace('.', '_')
        return os.path.join(self._outdir, output_file)

    def _set_word_size(self, threshold):
        threshold = float(threshold)
        if threshold >= 0.7:
            word = 5
        elif threshold >= 0.6:
            word = 4
        elif threshold >= 0.5:
            word = 3
        elif threshold >= 0.4:
            word = 2
        else:
            word = 1
        return word

    def _execute(self):
        if self._step is None:
            self._simple_execute(self._threshold)
        else:
            self._word   = self._set_word_size(self._step)
            self._output = self._step_file
            self._simple_execute(self._step)
            infile       = self._input
            self._input  = self.output_file
            self._word   = self._set_word_size(self._threshold)
            self._output = self._build_output_file(infile, self._threshold)
            self._simple_execute(self._threshold)
            self._input  = infile

    def _simple_execute(self, threshold):
        self._EXE.add_attribute(self._input,     '-i')
        self._EXE.add_attribute(self._output,    '-o')
        SBIg.alert('debug', self,
                   'Setting threshold to {0}'.format(threshold))
        self._EXE.add_attribute(threshold,       '-c')
        self._EXE.add_attribute('1',             '-g')
        SBIg.alert('debug', self, 'Setting word to {0}'.format(self._word))
        self._EXE.add_attribute(self._word,      '-n')
        if self._memory is not None:
            self._EXE.add_attribute(self._memory, '-M')

        try:
            self._EXE.execute()
            self._EXE.clean_command()
        except SystemError, e:
            SBIg.throw(self, 'Some error occurred while executing cd-hit\n{0}\n'.format(e),
                       SystemError)
