"""
Class to perform standard rejection-sampling and local-linear regression
analysis for ABC.
"""

__license__ = """
    Copyright 2015 Stephane De Mita, Mathieu Siol

    This file is part of EggLib.

    EggLib 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.

    EggLib 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 EggLib.  If not, see <http://www.gnu.org/licenses/>.
"""

from .. import _eggwrapper
from .. import misc

########################################################################

class ABC(object):

    """
    Class to perform standard rejection-sampling and local-linear
    regression analysis for ABC. Methods described in Beaumont et al.
    2002 Approximate Bayesian computation in population genetics
    (Genetics 162:2025-2035).
    """

    def __init__(self):
        try: self._static_abc = _eggwrapper.ABC()
        except AttributeError: raise AttributeError, 'egglib module was not built with support for the ABC class'
        self.reset()

    ####################################################################

    def reset(self):

        """
        Delete all information contained in the instance and restore all
        default values.
        """

        self.clear_io()
        self.clear_data()

    ####################################################################

    def clear_io(self):

        """
        Reset input and output arguments.
        """

        self._outf = None
        self._tmp = []
        self._inf = []
        self._nparams = []
        self._nstats = None

    ####################################################################

    def clear_data(self):

        """
        Reset data generated by fitting methods.
        """

        self._posterior = None
        self._proba = None

    ####################################################################

    def set_outf(self, fname):

        """
        Specify the output file name. If ``None`` (which is the default)
        the results are returned as lists by the fitting method
        :meth:`fit`.
        """

        self._outf = fname

    ####################################################################

    def add_inf(self, fname):

        """
        Add an input file name to process. It is possible to process
        several input datasets (allowing different numbers of
        parameters, but only if local-linear regression is not applied).

        :param fname: The name of an input file, containing simulated
            data. The data are formatted as follow: one line per
            simulation; each line contains space- or tab-separated
            values; the values are all parameter values, followed by a
            hash (``#``) as separator, followed by all statistic values.
            The order of parameters and statistics is required to be
            consistent over lines and over files when they correspond to
            the same model.
        """

        try: f = open(fname)
        except IOError: raise ValueError, 'cannot open this input file: {0}'.format(fname)
        line = f.readline()
        f.close()
        if line.count('#') != 1: raise ValueError, 'cannot read data from this file: {0}'.format(fname)
        p, s = line.split('#')
        np = len(p.split())
        if np == 0: raise ValueError, 'cannot read data from this file: {0} (no parameters found in first line)'.format(fname)
        ns = len(s.split())
        if ns == 0: raise ValueError, 'cannot read data from this file: {0} (no statistics found in first line)'.format(fname)
        if self._nstats == None: self._nstats = ns
        elif ns != self._nstats: raise ValueError, 'number of statistics does not match previous value in this file: {0}'.format(fname)
        self._inf.append(fname)
        self._nparams.append(np)

    ####################################################################

    def add_data(self, params, stats):

        """
        Add a dataset name to process. It is possible to process several
        datasets (allowing different numbers of parameters, but only if
        local-linear regression is not applied).

        :param params: A list of parameters for each pattern. The number
            of parameters can be zero but must be consistent over
            patterns. The number of patterns must be consistent over the
            two arguments.
        :param stats: A list of statistics for each pattern (one pattern
            is a dataset entry). The number of statistics must be at
            least one and consistent over patterns.
        """

        if len(stats) != len(params): raise ValueError, 'length of `stats` and `params` lists must be equal'
        if len(stats) == 0: raise ValueError, 'cannot proceed: empty dataset'

        ns = len(stats[0])
        if ns == 0: raise ValueError, 'cannot proceed: cannot use provided dataset (number of statistics is 0)'
        if self._nstats == None: self._nstats = ns
        elif ns != self._nstats: raise ValueError, 'number of statistics does not match previous value'

        self._tmp.append(misc.TempFile())
        np = len(params[0])
        self._nparams.append(np)
        with open(self._tmp[-1].fname, 'w') as f:
            for p, s in zip(params, stats):
                if len(p) != np: 'cannot use provided dataset (number of parameters is not consistent)'
                if len(s) != ns: 'cannot use provided dataset (number of statistics is not consistent)'
                f.write(' '.join(map(str, p)) + ' # ' + ' '.join(map(str, s)) + '\n')

        self._inf.append(self._tmp[-1].fname)
        self._nparams.append(np)

    ####################################################################

    @property
    def num_ds(self):

        """
        Number of loaded datasets.
        """

        return len(self._nparams)

    ####################################################################

    @property
    def tot_samples(self):

        """
        Total number of samples used in the last fit. Requires that fit
        has been performed by :meth:`fit`.
        """

        return self._static_abc.number_of_samples()

    ####################################################################

    @property
    def num_samples(self):

        """
        List of number of samples for each datasets used in the last
        fit. The order of values in the returned list follows the order
        in which datasets have been loaded. Requires that fit has been
        performed by :meth:`fit`.
        """

        return [self._static_abc.number_of_samples_part(i) for i in range(self._nd)]

    ####################################################################

    @property
    def posterior(self):

        """
        This attribute returns the list of accepted points as a list of
        lists of floating-point values. The value is ``None`` by default
        (if no fit has been performed, or if data have not been
        requested to be imported). If local-linear regression has been
        applied, the points are given after final correction.
        """

        return self._posterior

    ####################################################################

    @property
    def proba(self):

        """
        This attribute returns model posterior probability, as a list of
        floating-point values corresponding of the frequency of each
        model among accepted points. The value is ``None`` by default
        (if no fit has been performed, if no model choice has been
        performed, or if probabilities have not been requested to be
        imported).
        """

        return self._proba

    ####################################################################

    def fit(self, obs, tol, regr=False, transf=None, choice=False, get_data=True, get_proba=True):

        """
        Perform rejection-sampling (without local-linear regression).

        :param obs: List of observed statistic values.
        :param tol: Proportion of points to keep.
        :param regr: Perform local-linear regression (does not allow
            model choice).
        :param transf: Model of parameter transformation (only if *regr*
            is ``True``). One of ``None``, ``"log"`` and ``"tan"``.
        :param choice: Boolean indicating whether model choice should be
            performed. This option is not compatible with *regr*.
        :param get_data: Get the list of posterior points (available as
            attribute :attr:`~.posterior`). If *regr* is ``False``, this
            corresponds to accepted points; if *regr* is ``True``, to
            accepted points after correction. For very large datasets,
            it is advisable to specify an output file and set this
            option to ``False``.
        :param get_proba: Get model probabilities (available as
            attribute :attr:`~.proba`). Ignored if *choice* is
            ``False``.
        :return: Actual number of accepted samples.
        """

        # clear output data
        self.clear_data()

        # some parameter check
        if choice == True and regr == True: raise ValueError, 'model choice is not supported in local-linear regression mode'
        if transf == None: transf = self._static_abc.NONE
        elif transf == 'log': transf = self._static_abc.LOG
        elif transf == 'tan': transf = self._static_abc.TAN
        else: raise ValueError, 'invalid transform mode: {0}'.format(transf)
        if regr and len(set(self._nparams)) != 1: raise ValueError, 'the number of parameters must be constant over input datasets if regression is applied'

        # preparation of input files
        if len(self._nparams) == 0: raise ValueError, 'no input file loaded'
        if choice == False and len(set(self._nparams)) > 1: raise ValueError, 'the number of parameters must be constant over input datasets'
        self._nd = len(self._inf)

        # load observed statistics
        if len(obs) != self._nstats: raise ValueError, 'the list of observed statistic values does not have the correct number of items'
        self._static_abc.number_of_statistics(self._nstats)
        for i,v in enumerate(obs): self._static_abc.obs(i, v)

        # pass file names
        for inf, np in zip(self._inf, self._nparams): self._static_abc.add_fname(inf, np)

        # load tolerance parameter
        if tol <= 0.0 or tol >= 1.0: raise ValueError, 'value for argument `tol` is out of acceptable range: {0}'.format(tol)
        self._static_abc.get_threshold(tol)

        # prepare output file for rejection step
        if self._outf == None or regr == True:
            self._tmp.append(misc.TempFile())
            outf = self._tmp[-1].fname
        else: outf = self._outf

        # perform rejection
        num_accp = self._static_abc.rejection(outf, choice, regr == False)

        # get data

        if get_data == True and regr == False:
            if choice and get_proba: self._proba = [0] * self._nd
            f = open(outf)
            self._posterior = []
            for line in f:
                line = line.split()
                if choice == True:
                    if get_proba == True: self._proba[int(line[0][1:-1])-1] += 1
                    del line[0]
                self._posterior.append(map(float, line))
            f.close()

        # get model probabilities

        if choice == True and get_proba == True:
            if not get_data:
                self._proba = [0] * self._nd
                f = open(outf)
                for line in f:
                    line = line.split()
                    self._proba[int(line[0][1:-1])-1] += 1
                f.close()
            if sum(self._proba) != num_accp: raise RuntimeError, 'inconsistency detected in rejection-sampling model choice output'
            self._proba = [float(i) / num_accp for i in self._proba]

        # perform local-linear regression

        if regr:
            tmpf = outf
            if self._outf == None:
                self._tmp.append(misc.TempFile())
                outf = self._tmp[-1].fname
            else: outf = self._outf
            self._static_abc.regression(tmpf, outf, transf, '')

            # get data

            if get_data == True:
                f = open(outf)
                self._posterior = []
                for line in f:
                    line = line.split()
                    self._posterior.append(map(float, line))
                f.close()

        # return
        return num_accp
