#!/usr/bin/python3

"""
This module defines the evolutionary Model that can be linked
to phylogeny, and computed by one of codeml, gerp, slr.
"""

from re       import sub
from warnings import warn

from ..evol.control import PARAMS, AVAIL
from ..evol.parser  import parse_paml, parse_rst, get_ancestor, parse_slr
try:
    from ..treeview.faces import SequencePlotFace
except ImportError:
    TREEVIEW = False
else:
    TREEVIEW = True

class Model:
    '''Evolutionary model.
    "omega" stands for starting value of omega, in the computation. As
    Zihen Yang says, it is good to try with different starting values...
    model linked to tree by _tree variable
    results of calculation are stored in dictionaries:
     * branches: w dN dS bL by mean of their node_id
     * sites   : values at each site.
     * classes : classes of sites and proportions
     * stats   : lnL number of parameters kappa value and codon frequencies stored here.

    available models are:
        =========== ============================= ==================
        Model name  Description                   Model kind
        =========== ============================= ==================\n%s
        =========== ============================= ==================\n

    :argument model_name: string with model name. Add a dot followed by anything at the end of the string in order to extend the name of the model and avoid overwriting.
    :argument None tree: a Tree object
    :argument None path: path to outfile, were model computation output can be found.

    '''
    def __init__(self, model_name, tree=None, path=None, **kwargs):
        self._tree      = tree
        self.name, args = check_name(model_name)
        self.sites      = None
        self.classes    = None
        self.n_classes  = None
        self.branches   = {}
        self.stats      = {}
        self.properties = {}
        for a, b in list(args.items()):
            self.properties [a] = b
        params = dict(list(PARAMS.items()))
        self._change_params(params)
        for key, arg in list(kwargs.items()):
            if key not in params:
                warn('WARNING: unknown param %s, can cause problems...'% (key))
            if key == 'gappy':
                arg = not arg
            params[key] = arg
        self.__check_marks()
        if path:
            self._load(path)

    def __str__(self):
        '''
        to print nice info
        '''
        str_mark = ''
        str_line = '\n        mark:%-5s, omega: %-10s, node_ids: %-4s, name: %s'
        for i, node in enumerate(self._tree.traverse()):
            if node.is_root:
                str_mark += str_line % (self.branches[node.props.get('node_id')]['mark'],
                                        'None',
                                        node.props.get('node_id'), node.name or 'ROOT')
            else:
                str_mark += str_line % (self.branches[node.props.get('node_id')]['mark'],
                                        self.branches[node.props.get('node_id')].get('w',
                                                                        'None'),
                                        node.props.get('node_id'), node.name or 'EDGE')
        str_site = ''
        str_line = '\n        %-12s: %s '
        if self.classes:
            for t in [t for t in sorted(self.classes)]:
                str_site += str_line % (t, ' '.join(['%s%s=%-9s' % (t[0], j, i)\
                                                     for j, i in \
                                                     enumerate(self.classes[t])]
                                                ))
        return ''' Evolutionary Model %s:
        log likelihood       : %s
        number of parameters : %s
        sites inference      : %s
        sites classes        : %s
        branches             : %s
        ''' % (self.name,
               self.lnL if 'lnL' in self.stats else 'None',
               self.np  if 'np'  in self.stats else 'None',
               ', '.join(sorted(list(self.sites.keys())))  if self.sites else 'None',
               str_site if self.classes else 'None',
               str_mark if self.branches else 'None'
           )


    def __check_marks(self):
        """
        checks if tree is marked and if model allows marks.
        fill up branches dict with marks
        """
        has_mark = any(n.props.get('mark') for n in self._tree.descendants())
        for i, node in enumerate(self._tree.traverse()):
            if has_mark and self.properties['allow_mark']:
                self.branches[node.props.get('node_id')] = {'mark': node.props.get('mark') or ' #0'}
            elif 'branch' in self.properties['typ']:
                self.branches[node.props.get('node_id')] = {'mark': ' #'+str(i)}
            else:
                self.branches[node.props.get('node_id')] = {'mark': ''}

    def _load(self, path):
        '''
        parse outfiles and load in model object
        '''
        if self.properties['exec'] == 'codeml':
            parse_paml(path, self)
            # parse rst file if site or branch-site model
            if 'site' in self.properties['typ']:
                # sites and classes attr
                for key, val in parse_rst(path).items():
                    setattr(self, key, val)
            if 'ancestor' in self.properties['typ']:
                get_ancestor(path, self)
            vars(self) ['lnL'] = self.stats ['lnL']
            vars(self) ['np']  = self.stats ['np']
        elif self.properties['exec'] == 'Slr':
            for key, val in parse_slr(path).items():
                setattr (self, key, val)
            vars(self) ['lnL'] = 0
            vars(self) ['np']  = 0

    def _change_params(self, params):
        '''
        change model specific values
        '''
        for key, change in self.properties ['changes']:
            params[key] = change
        self.properties ['params'] = params

    def set_histface(self, up=True, hlines=(1.0, 0.3), kind='bar',
                      errors=False, colors=None, **kwargs):
        """
        To add histogram face for a given site mdl (M1, M2, M7, M8)
        can choose to put it up or down the tree.
        2 types are available:
           * stick: to draw histogram.
           * curve: to draw plot.
        You can define color scheme by passing a diccionary, default is:
            col = {'NS' : 'grey'  ,
                   'RX' : 'green' ,
                   'RX+': 'green' ,
                   'CN' : 'cyan'  ,
                   'CN+': 'blue'  ,
                   'PS' : 'orange',
                   'PS+': 'red'    }
        """
        if self.sites is None:
            warn("WARNING: model %s not computed." % (self.name))
            return None
        if not 'header' in kwargs:
            kwargs['header'] = 'Omega value for sites under %s model' % \
                               (self.name)
        if 'BEB' in self.sites:
            val = 'BEB'
        elif 'NEB' in self.sites:
            val = 'NEB'
        else:
            val = 'SLR'
        colors = self.colorize_rst(val, col=colors)
        if not 'ylim' in kwargs:
            kwargs['ylim'] = (0, 2)
        if errors:
            errors = self.sites[val].get('se', None)
        if TREEVIEW:
            try:
                hist = SequencePlotFace(self.sites[val]['w'], hlines=hlines,
                                        colors=colors, errors=errors,
                                        ylabel=u'Omega (\u03c9)', kind=kind,
                                        **kwargs)
            except KeyError:
                raise Exception('ERROR: no sites to display, only available ' +
                                'histfaces for site models\n')
            if up:
                setattr(hist, 'up', True)
            else:
                setattr(hist, 'up', False)
        else:
            hist = None
        self.properties['histface'] = hist


    def get_ctrl_string(self, outfile=None):
        '''
        generate ctrl string to write to a file, if file is given,
        write it, otherwise returns the string

        :argument None outfile: if a path is given here, write control string into it.

        :returns: the control string

        '''
        string = ''
        if 'sep' in self.properties:
            sep = self.properties ['sep']
        else:
            sep = ' = '
        for prm in ['seqfile', 'treefile', 'outfile']:
            string += '%15s%s%s\n' % (prm, sep,
                                      str(self.properties['params'][prm]))
        string += '\n'
        for prm in sorted(list(self.properties ['params'].keys()), key=lambda x:
                          sub('fix_', '', x.lower())):
            if prm in ['seqfile', 'treefile', 'outfile']:
                continue
            if str(self.properties ['params'][prm]).startswith('*'):
                continue
                #string += ' *'+'%13s = %s\n' \
                #          % (p, str(self.properties ['params'][p])[1:])
            else:
                string += '%15s%s%s\n' % (prm, sep,
                                          str(self.properties ['params'][prm]))
        if outfile is None:
            return string
        else:
            open(outfile, 'w').write(string)

    def colorize_rst(self, val, col=None):
        '''
        Colorize function, that take in argument a list of values
        corresponding to a list of classes and returns a list of
        colors to paint histogram.

        :param val: type of estimation, can be BEB or NEB (only
           positive-selection models have BEB)
        :param None col: a dictionary of colors that by default is:
           {"NS" : "grey",
            "RX" : "green",
            "RX+": "green",
            "CN" : "cyan",
            "CN+": "blue",
            "PS" : "orange",
            "PS+": "red"}

        :returns: a list of colors dependending categories of sites that are among:
          - CN+ > 0.99 probabylity of beloging to conserved class of site
          - CN  > 0.95 probabylity of beloging to conserved class of site
          - NS  not significant
          - RX+ > 0.99 probabylity of beloging to relaxed class of site
          - RX  > 0.95 probabylity of beloging to relaxed class of site
          - PS+ > 0.99 probabylity of beloging to positively-selected class of site
          - PS  > 0.95 probabylity of beloging to positively-selected class of site
        '''
        col = col or {'NS' : 'grey',
                      'RX' : 'green',
                      'RX+': 'green',
                      'CN' : 'cyan',
                      'CN+': 'blue',
                      'PS' : 'orange',
                      'PS+': 'red'}
        if not 'site' in self.properties['typ']:
            raise Exception('ERROR: histogram are only for site and '
                            'branch-site models.')
        categories = self.significance_by_site(val)
        return [col[cat] for cat in categories]

    def significance_by_site(self, val):
        '''
        Summarize significance of site models.

        :param val: type of estimation, can be BEB or NEB (only
           positive-selection models have BEB)

        :returns: a list of categories among:
          - CN+ > 0.99 probabylity of beloging to conserved class of site
          - CN  > 0.95 probabylity of beloging to conserved class of site
          - NS  not significant
          - RX+ > 0.99 probabylity of beloging to relaxed class of site
          - RX  > 0.95 probabylity of beloging to relaxed class of site
          - PS+ > 0.99 probabylity of beloging to positively-selected class of site
          - PS  > 0.95 probabylity of beloging to positively-selected class of site
        '''
        if not 'site' in self.properties['typ']:
            raise Exception('ERROR: only for site and '
                            'branch-site models.')
        ps_model = 'positive' in self.properties['evol']
        categories = []
        for pval, curr_class in zip(self.sites[val]['pv'],
                                    self.sites[val]['class']):
            if pval < 0.95:
                categories.append('NS')
            elif curr_class == self.n_classes[val] and not ps_model:
                if pval < 0.99:
                    categories.append('RX')
                else:
                    categories.append('RX+')
            elif curr_class == 1:
                if pval < 0.99:
                    categories.append('CN')
                else:
                    categories.append('CN+')
            elif curr_class >= self.n_classes[val] and ps_model:
                if pval < 0.99:
                    categories.append('PS')
                else:
                    categories.append('PS+')
            elif curr_class == self.n_classes[val]:
                if pval < 0.99:
                    categories.append('RX')
                else:
                    categories.append('RX+')
            else:
                categories.append('NS')
        return categories


def check_name(model):
    '''
    check that model name corresponds to one of the available
    '''
    if sub(r'\..*', '', model) in AVAIL:
        return model, AVAIL [sub(r'\..*', '', model)]



Model.__doc__ = Model.__doc__ % \
                ('\n'.join([ '          %-8s   %-27s   %-15s  ' % \
                             ('%s' % (x), AVAIL[x]['evol'], AVAIL[x]['typ']) \
                             for x in sorted(sorted(AVAIL.keys()),key=lambda x: \
                                AVAIL[x]['typ'],
                                reverse=True)]))
