from collections import namedtuple
import sys
import re
import warnings
import array

import numpy as np

from scipy._lib._util import check_random_state
from scipy.optimize import minimize, differential_evolution, least_squares
import scipy.optimize as sciopt

from refnx.analysis import Objective, Interval, PDF, is_parameter
from refnx._lib import (
    unique as f_unique,
    MapWrapper,
    possibly_open_file,
    flatten,
)
from refnx._lib.util import getargspec
from refnx._lib._qmc import LatinHypercube
from refnx._lib import emcee
from refnx._lib.emcee.state import State
from refnx._lib.emcee.pbar import get_progress_bar

MCMCResult = namedtuple(
    "MCMCResult", ["name", "param", "stderr", "chain", "median"]
)

class PTSampler:
    def __init__(self, ntemps, nwalkers, ndim, logl, logp, **kwargs):
        """
        Shim class for a ptemcee.PTSampler.

        Parameters
        ----------
        ntemps: int, np.array
            Specifies the number of parallel tempering temperatures.
            If an array specifies a ladder of Beta values.
        nwalkers: int
            Number of walkers
        ndim: int
            Dimensionality of the problem space
        logl: callable
            log-likelihood function
        logp: callable
            log-prior function
        kwargs:
            Other keyword arguments supplied to construct the ptemcee.Sampler.
        """
        from refnx._lib.ptemcee import Sampler as _PTSampler

        self.ntemps = ntemps
        self.nwalkers = nwalkers
        self.ndim = ndim
        self.logl = logl
        self.logp = logp
        self.kwargs = kwargs

        sig = {
            "betas": ntemps,
            "nwalkers": nwalkers,
            "ndim": ndim,
            "logl": logl,
            "logp": logp,
        }
        sig.update(kwargs)
        self.sampler = _PTSampler(**sig)

        # chain stepper
        self._ptchain = None
        self._state = None

    def sample(
        self,
        initial_state,
        iterations=1,
        thin_by=1,
        progress=False,
        mapper=None,
        **kwds
    ):
        """
        Runs the PTSampler for a given number of iterations.

        Parameters
        ----------
        initial_state: emcee.state.State
            Holds the coordinates of the initial state
        iterations: int
            Number of steps to save into the chain
        thin_by: int
            The saved steps are separated by this many discarded steps.
        progress: bool
            Display a progress bar.
        mapper: map-like callable
            For parallelisation
        kwds: dict
            Unknown keywords

        Yields
        -------
        state: emcee.state.State
            The coordinates of the current state
        """
        if isinstance(initial_state, State):
            init_x = initial_state.coords
            rstate0 = initial_state.random_state
        else:
            init_x = initial_state
            rstate0 = np.random.RandomState().get_state()

        if self._ptchain is None:
            self._ptchain = self.sampler.chain(init_x)
        else:
            self._ptchain.ensemble.x = init_x

        # set random state of stateful chain
        self.random_state = rstate0

        self._ptchain.thin_by = thin_by

        if mapper is not None:
            self._ptchain.ensemble._mapper = mapper

        try:
            with get_progress_bar(progress, iterations * thin_by) as pbar:
                for e in self._ptchain.iterate(iterations):
                    self._state = State(e.x, log_prob=e.logl + e.logP)
                    yield self._state
                    pbar.update(thin_by)
        finally:
            self._ptchain.ensemble._mapper = map

    def thermodynamic_integration_log_evidence(self, fburnin=0.1):
        if self._ptchain is not None:
            return self._ptchain.log_evidence_estimate(fburnin)
        return None, None

    def reset(self):
        if self._state is not None:
            self._ptchain = self.sampler.chain(self._state.coords)

    def get_chain(self):
        if self._ptchain is not None:
            return self._ptchain.x
        return None

    @property
    def chain(self):
        if self._ptchain is not None:
            return self._ptchain.x
        return None

    def get_log_prob(self):
        if self._ptchain is not None:
            return self._ptchain.logP
        return None

    @property
    def random_state(self):
        if self._ptchain is not None:
            self._ptchain.ensemble._random.get_state()

    @random_state.setter
    def random_state(self, rstate0):
        if self._ptchain is not None:
            self._ptchain.ensemble._random.set_state(rstate0)

class CurveFitter:
    """
    Analyse a curvefitting system (with MCMC sampling)

    Parameters
    ----------
    objective : refnx.analysis.Objective
        The :class:`refnx.analysis.Objective` to be analysed.
    nwalkers : int, optional
        How many walkers you would like the sampler to have. Must be an
        even number. The more walkers the better.
    ntemps : int or None, optional
        If `ntemps == -1`, then an :class:`emcee.EnsembleSampler` is used
        during the `sample` method.
        Otherwise, or if `ntemps is None` then parallel tempering is
        used with a :class:`ptemcee.sampler.Sampler` object during the `sample`
        method, with `ntemps` specifing the number of temperatures. Can be
        `None`, in which case the `Tmax` keyword argument sets the maximum
        temperature. Parallel Tempering is useful if you expect your
        posterior distribution to be multi-modal.

    mcmc_kws : dict
        Keywords used to create the :class:`emcee.EnsembleSampler` or
        :class:`ptemcee.sampler.Sampler` objects.

    Notes
    -----
    See the documentation at http://dan.iel.fm/emcee/current/api/ for
    further details on what keywords are permitted, and for further
    information on Parallel Tempering. The `pool` and `threads` keywords
    are ignored here. Specification of parallel threading is done with the
    `pool` argument in the `sample` method.
    """

    def __init__(self, objective, nwalkers=200, ntemps=-1, **mcmc_kws):
        """
        Parameters
        ----------
        objective : refnx.analysis.Objective
            The :class:`refnx.analysis.Objective` to be analysed.
        nwalkers : int, optional
            How many walkers you would like the sampler to have. Must be an
            even number. The more walkers the better.
        ntemps : int or None, optional
            If `ntemps == -1`, then an :class:`emcee.EnsembleSampler` is used
            during the `sample` method.
            Otherwise, or if `ntemps is None` then parallel tempering is
            used with a :class:`ptemcee.sampler.Sampler` object during the
            `sample` method, with `ntemps` specifing the number of
            temperatures. Can be `None`, in which case the `Tmax` keyword
            argument sets the maximum temperature. Parallel Tempering is
            useful if you expect your posterior distribution to be multi-modal.
        mcmc_kws : dict
            Keywords used to create the :class:`emcee.EnsembleSampler` or
            :class:`ptemcee.sampler.PTSampler` objects.

        Notes
        -----
        See the documentation at http://dan.iel.fm/emcee/current/api/ for
        further details on what keywords are permitted. The `pool` and
        keyword is ignored here. Specification of parallel threading is done
        with the `pool` argument in the `sample` method.
        To use parallel tempering you will need to install the
        :package:`ptemcee` package.
        """
        self.objective = objective
        self._varying_parameters = []
        self.__var_id = []

        self.mcmc_kws = {}
        if mcmc_kws is not None:
            self.mcmc_kws.update(mcmc_kws)

        if "pool" in self.mcmc_kws:
            self.mcmc_kws.pop("pool")
        if "threads" in self.mcmc_kws:
            self.mcmc_kws.pop("threads")

        self._nwalkers = nwalkers
        self._ntemps = ntemps
        self.make_sampler()
        self._state = None

    def __setstate__(self, state):
        self.__dict__.update(state)
        self.__var_id = [
            id(obj) for obj in self.objective.varying_parameters()
        ]

    @property
    def nvary(self):
        return len(self._varying_parameters)

    def __repr__(self):
        # attempt to get a minimum repr for a CurveFitter. However,
        # it has so much state when the sampling has been done, that
        # will be ignored.
        d = {
            "objective": self.objective,
            "_nwalkers": self._nwalkers,
            "_ntemps": self._ntemps,
            "mcmc_kws": self.mcmc_kws,
        }
        return (
            "CurveFitter({objective!r},"
            " nwalkers={_nwalkers},"
            " ntemps={_ntemps},"
            " {mcmc_kws!r})".format(**d)
        )

    def make_sampler(self):
        """
        Make the samplers for the Objective.

        Use this method if the number of varying parameters changes.
        """
        self._varying_parameters = self.objective.varying_parameters()
        self.__var_id = [id(obj) for obj in self._varying_parameters]

        if not self.nvary:
            raise ValueError("No parameters are being fitted")

        if self._ntemps == -1:
            self.sampler = emcee.EnsembleSampler(
                self._nwalkers,
                self.nvary,
                self.objective.logpost,
                **self.mcmc_kws
            )
        # Parallel Tempering was requested.
        else:
            sig = {
                "ntemps": self._ntemps,
                "nwalkers": self._nwalkers,
                "ndim": self.nvary,
                "logl": self.objective.logl,
                "logp": self.objective.logp,
            }
            sig.update(self.mcmc_kws)
            self.sampler = PTSampler(**sig)

        self._state = None

    def _check_vars_unchanged(self):
        """
        Keep track of whether the varying parameters have changed after
        construction of CurveFitter object.

        """
        var_ids = [id(obj) for obj in self.objective.varying_parameters()]
        if not (np.array_equal(var_ids, self.__var_id)):
            raise RuntimeError(
                "The Objective.varying_parameters() have"
                " changed since the CurveFitter was created."
                " To keep on using the CurveFitter call"
                " the CurveFitter.make_samplers() method."
            )

    def initialise(self, pos="covar", random_state=None):
        """
        Initialise the emcee walkers.

        Parameters
        ----------
        pos : str or np.ndarray
            Method for initialising the emcee walkers. One of:

            - 'covar', use the estimated covariance of the system.
            - 'jitter', add a small amount of gaussian noise to each parameter
            - 'prior', sample random locations from the prior using Latin
                Hyper Cube.
            - pos, an array that specifies a snapshot of the walkers. Has shape
                `(nwalkers, ndim)`, or `(ntemps, nwalkers, ndim)` if parallel
                 tempering is employed. You can also provide a previously
                 created chain.
        random_state : {int, `np.random.RandomState`, `np.random.Generator`}
            If `random_state` is not specified the `~np.random.RandomState`
            singleton is used.
            If `random_state` is an int, a new ``RandomState`` instance is
            used, seeded with random_state.
            If `random_state` is already a ``RandomState`` or a ``Generator``
            instance, then that object is used.
            Specify `random_state` for repeatable initialisations.
        """
        nwalkers = self._nwalkers
        nvary = self.nvary

        # acquire a random number generator
        rng = check_random_state(random_state)

        # account for parallel tempering
        _ntemps = self._ntemps

        # If you're not doing parallel tempering, temporarily set the number of
        # temperatures to be created to 1, thereby producing initial positions
        # of (1, nwalkers, nvary), this first dimension should be removed at
        # the end of the method
        if self._ntemps == -1:
            _ntemps = 1

        # position is specified with array (no parallel tempering)
        if (
            isinstance(pos, np.ndarray)
            and self._ntemps == -1
            and pos.shape == (nwalkers, nvary)
        ):
            init_walkers = np.copy(pos)[np.newaxis]

        # position is specified with array (with parallel tempering)
        elif (
            isinstance(pos, np.ndarray)
            and self._ntemps > -1
            and pos.shape == (_ntemps, nwalkers, nvary)
        ):
            init_walkers = np.copy(pos)

        # position is specified with existing chain
        elif isinstance(pos, np.ndarray):
            self.initialise_with_chain(pos)
            return

        # position is to be created from covariance matrix
        elif pos == "covar":
            p0 = np.array(self._varying_parameters)
            cov = self.objective.covar()
            init_walkers = rng.multivariate_normal(
                np.atleast_1d(p0), np.atleast_2d(cov), size=(_ntemps, nwalkers)
            )

        # position is specified by jittering the parameters with gaussian noise
        elif pos == "jitter":
            var_arr = np.array(self._varying_parameters)
            pos = 1 + rng.standard_normal((_ntemps, nwalkers, nvary)) * 1.0e-4
            pos *= var_arr
            init_walkers = pos

        # use the prior to initialise position
        elif pos == "prior":
            arr = np.zeros((_ntemps, nwalkers, nvary))
            LHC = LatinHypercube(nvary, seed=random_state)
            samples = LHC.random(n=_ntemps * nwalkers).reshape(
                _ntemps, nwalkers, nvary
            )

            for i, param in enumerate(self._varying_parameters):
                # bounds are not a closed interval, just jitter it.
                if (
                    isinstance(param.bounds, Interval)
                    and not param.bounds._closed_bounds
                ):
                    vals = (
                        1 + rng.standard_normal((_ntemps, nwalkers)) * 1.0e-1
                    )
                    vals *= param.value
                    arr[..., i] = vals
                else:
                    sample_arr = samples[..., i]
                    transformed = param.bounds.invcdf(sample_arr)
                    arr[..., i] = transformed

            init_walkers = arr

        else:
            raise RuntimeError(
                "Didn't use any known method for " "CurveFitter.initialise"
            )

        # if you're not doing parallel tempering then remove the first
        # dimension
        if self._ntemps == -1:
            init_walkers = init_walkers[0]

        # now validate initialisation, ensuring all init pos have finite
        # logpost
        for i, param in enumerate(self._varying_parameters):
            init_walkers[..., i] = param.valid(init_walkers[..., i])

        rstate0 = None
        if isinstance(rng, np.random.RandomState):
            rstate0 = rng.get_state()

        self._state = State(init_walkers, random_state=rstate0)

        # finally reset the sampler to reset the chain
        # you have to do this at the end, not at the start because resetting
        # makes self.sampler.chain == None and the PTsampler creation doesn't
        # work
        self.sampler.reset()

    def initialise_with_chain(self, chain):
        """
        Initialise sampler with a pre-existing chain

        Parameters
        ----------
        chain : array
            Array of size `(steps, ntemps, nwalkers, ndim)` or
            `(steps, nwalkers, ndim)`, containing a chain from a previous
            sampling run.
        """
        # we should be left with (nwalkers, ndim) or (ntemp, nwalkers, ndim)

        if self._ntemps == -1:
            required_shape = (self._nwalkers, self.nvary)
        else:
            required_shape = (self._ntemps, self._nwalkers, self.nvary)

        chain_shape = chain.shape[1:]

        # if the shapes are the same, then we can initialise
        if required_shape == chain_shape:
            self.initialise(pos=chain[-1])
        else:
            raise ValueError(
                "You tried to initialise with a chain, but it was"
                " the wrong shape"
            )

    @property
    def chain(self):
        """
        MCMC chain belonging to CurveFitter.sampler

        Returns
        -------
        chain : array
            The MCMC chain with shape `(steps, nwalkers, ndim)` or
            `(steps, ntemps, nwalkers, ndim)`.
        """
        return self.sampler.get_chain()

    @property
    def logpost(self):
        """
        Log-probability for each of the entries in `self.chain`
        """
        return self.sampler.get_log_prob()

    @property
    def index_max_prob(self):
        """
        The index of the highest log-probability for the samples
        """
        log_probs = self.sampler.get_log_prob()
        if isinstance(self.sampler, PTSampler):
            log_probs = log_probs[:, 0]

        loc = np.argmax(log_probs)
        idx = np.unravel_index(loc, log_probs.shape)

        if isinstance(self.sampler, PTSampler):
            idx = list(idx)
            idx.insert(1, 0)
            return tuple(idx)

        return idx

    def reset(self):
        """
        Reset the sampled chain.

        Typically used on a sampler after a burn-in period.
        """
        self.sampler.reset()

    def acf(self, nburn=0, nthin=1):
        """
        Calculate the autocorrelation function

        Returns
        -------
        acfs : np.ndarray
            The autocorrelation function, acfs.shape=(lags, nvary)
        """
        return autocorrelation_chain(self.chain, nburn=nburn, nthin=nthin)

    def sample(
        self,
        steps,
        nthin=1,
        random_state=None,
        f=None,
        callback=None,
        verbose=True,
        pool=-1,
    ):
        """
        Performs sampling from the objective.

        Parameters
        ----------
        steps : int
            Collect `steps` samples into the chain. The sampler will run a
            total of `steps * nthin` moves.
        nthin : int, optional
            Each chain sample is separated by `nthin` iterations.
        random_state : {int, `np.random.RandomState`, `np.random.Generator`}
            If `random_state` is not specified the `~np.random.RandomState`
            singleton is used.
            If `random_state` is an int, a new ``RandomState`` instance is
            used, seeded with random_state.
            If `random_state` is already a ``RandomState`` or a ``Generator``
            instance, then that object is used.
            Specify `random_state` for repeatable minimizations.
        f : file-like or str
            File to incrementally save chain progress to. Each row in the file
            is a flattened array of size `(nwalkers, ndim)` or
            `(ntemps, nwalkers, ndim)`. There are `steps` rows in the
            file.
        callback : callable
            callback function to be called at each iteration step. Has the
            signature `callback(coords, logprob)`.
        verbose : bool, optional
            Gives updates on the sampling progress
        pool : int or map-like object, optional
            If `pool` is an `int` then it specifies the number of threads to
            use for parallelization. If `pool == -1`, then all CPU's are used.
            If pool is a map-like callable that follows the same calling
            sequence as the built-in map function, then this pool is used for
            parallelisation.

        Notes
        -----
        Please see :class:`emcee.EnsembleSampler` for its detailed behaviour.

        >>> # we'll burn the first 500 steps
        >>> fitter.sample(500)
        >>> # after you've run those, then discard them by resetting the
        >>> # sampler.
        >>> fitter.sampler.reset()
        >>> # Now collect 40 steps, each step separated by 50 sampler
        >>> # generations.
        >>> fitter.sample(40, nthin=50)

        One can also burn and thin in `Curvefitter.process_chain`.
        """
        self._check_vars_unchanged()

        # setup a random number generator
        rng = check_random_state(random_state)

        if self._state is None:
            self.initialise(random_state=rng)

        # for saving progress to file
        def _callback_wrapper(state, h=None):
            if callback is not None:
                callback(state.coords, state.log_prob)

            if h is not None:
                h.write(" ".join(map(str, state.coords.ravel())))
                h.write("\n")

        # remove chains from each of the parameters because they slow down
        # pickling but only if they are parameter objects.
        flat_params = f_unique(flatten(self.objective.parameters))
        flat_params = [param for param in flat_params if is_parameter(param)]
        # zero out all the old parameter stderrs
        for param in flat_params:
            param.stderr = None
            param.chain = None

        # make sure the checkpoint file exists
        if f is not None:
            with possibly_open_file(f, "w") as h:
                # write the shape of each step of the chain
                h.write("# ")
                shape = self._state.coords.shape
                h.write(", ".join(map(str, shape)))
                h.write("\n")

        # set the random state of the sampler
        # normally one could give this as an argument to the sample method
        # but PTSampler didn't historically accept that...
        if isinstance(rng, np.random.RandomState):
            rstate0 = rng.get_state()
            self._state.random_state = rstate0
            self.sampler.random_state = rstate0

        # using context manager means we kill off zombie pool objects
        # but does mean that the pool has to be specified each time.
        with MapWrapper(pool) as g, possibly_open_file(f, "a") as h:
            # these kwargs are provided to the sampler.sample method
            kwargs = {"iterations": steps, "thin": nthin}

            # if you're not creating more than 1 thread, then don't bother with
            # a pool.
            if isinstance(self.sampler, emcee.EnsembleSampler):
                if pool == 1:
                    self.sampler.pool = None
                else:
                    self.sampler.pool = g
            else:
                kwargs["mapper"] = g

            # new emcee arguments
            sampler_args = getargspec(self.sampler.sample).args
            if "progress" in sampler_args and verbose:
                kwargs["progress"] = True
                verbose = False

            if "thin_by" in sampler_args:
                kwargs["thin_by"] = nthin
                kwargs.pop("thin", 0)

            # perform the sampling
            for state in self.sampler.sample(self._state, **kwargs):
                self._state = state
                _callback_wrapper(state, h=h)

        if isinstance(self.sampler, emcee.EnsembleSampler):
            self.sampler.pool = None

        # sets parameter value and stderr
        return process_chain(self.objective, self.chain)

    def fit(self, method="L-BFGS-B", target="nll", verbose=True, **kws):
        """
        Obtain the maximum log-likelihood, or log-posterior, estimate (mode)
        of the objective. Maximising the log-likelihood is equivalent to
        minimising chi2 in a least squares fit.

        Parameters
        ----------
        method : str
            which method to use for the optimisation. One of:

            - `'least_squares'`: :func:`scipy.optimize.least_squares`.
            - `'L-BFGS-B'`: L-BFGS-B.
            - `'differential_evolution'`:
              :func:`scipy.optimize.differential_evolution`
            - `'dual_annealing'`:
              :func:`scipy.optimize.dual_annealing` (SciPy >= 1.2.0)
            - `'shgo'`: :func:`scipy.optimize.shgo` (SciPy >= 1.2.0)

            You can also choose many of the minimizers from
            :func:`scipy.optimize.minimize`.

        target : {'nll', 'nlpost'}, optional
            Minimize the negative log-likelihood (`'nll'`) or the negative
            log-posterior (`'nlpost'`). This is equivalent to maximising the
            likelihood or posterior probabilities respectively.
            Maximising the likelihood is equivalent to minimising chi^2 in a
            least-squares fit.
            This option only applies to the `differential_evolution`, `shgo`,
            `dual_annealing` or `L-BFGS-B` methods.
            These optimisers require lower and upper (box) bounds for each
            parameter. If the `Bounds` on a parameter are not an `Interval`,
            but a `PDF` specifying a statistical distribution, then the lower
            and upper bounds are approximated as
            ``PDF.rv.ppf([0.005, 0.995])``, covering 99 % of the statistical
            distribution.
        verbose : bool, optional
            Gives fitting progress. To see a progress bar tqdm has to be
            installed.
        kws : dict
            Additional arguments are passed to the underlying minimization
            method.

        Returns
        -------
        result, covar : :class:`scipy.optimize.OptimizeResult`, np.ndarray
            `result.x` contains the best fit parameters
            `result.covar` is the covariance matrix for the fit.
            `result.stderr` is the uncertainties on each of the fit parameters.

        Notes
        -----
        If the `objective` supplies a `residuals` method then `least_squares`
        can be used. Otherwise the `nll` method of the `objective` is
        minimised. Use this method just before a sampling run.
        If `self.objective.parameters` is a `Parameters` instance, then each
        of the varying parameters has its value updated by the fit, and each
        `Parameter` has a `stderr` attribute which represents the uncertainty
        on the fit parameter.

        The use of `dual annealing` and `shgo` requires that `scipy >= 1.2.0`
        be installed.

        """
        _varying_parameters = self.objective.varying_parameters()
        init_pars = np.array(_varying_parameters)

        _min_kws = {}
        _min_kws.update(kws)
        _bounds = bounds_list(self.objective.varying_parameters())
        _min_kws["bounds"] = _bounds

        # setup callback default
        _min_kws.setdefault("callback", None)

        cost = self.objective.nll
        if target == "nlpost":
            cost = self.objective.nlpost

        # a decorator for the progress bar updater
        def _callback_wrapper(callback_func, pbar):
            def callback(*args, **kwds):
                pbar.update(1)
                if callback_func is None:
                    return None
                else:
                    return callback_func(*args, **kwds)

            return callback

        # least_squares Trust Region Reflective by default
        if method == "least_squares":
            b = np.array(_bounds)
            _min_kws["bounds"] = (b[..., 0], b[..., 1])

            # least_squares doesn't have a callback
            _min_kws.pop("callback", None)

            res = least_squares(
                self.objective.residuals, init_pars, **_min_kws
            )
        # differential_evolution, dual_annealing, shgo require lower and upper
        # bounds
        elif method in ["differential_evolution", "dual_annealing", "shgo"]:
            mini = getattr(sciopt, method)

            if method == "shgo":
                if "n" not in _min_kws:
                    _min_kws["n"] = 100
                if "iters" not in kws:
                    _min_kws["iters"] = 5

            with get_progress_bar(verbose, None) as pbar:
                _min_kws["callback"] = _callback_wrapper(
                    _min_kws["callback"], pbar
                )

                res = mini(cost, **_min_kws)
        else:
            # otherwise stick it to minimizer. Default being L-BFGS-B
            _min_kws["method"] = method
            _min_kws["bounds"] = _bounds

            with get_progress_bar(verbose, None) as pbar:
                _min_kws["callback"] = _callback_wrapper(
                    _min_kws["callback"], pbar
                )

                res = minimize(cost, init_pars, **_min_kws)

        # OptimizeResult.success may not be present (dual annealing)
        if hasattr(res, "success") and res.success:
            self.objective.setp(res.x)

            # Covariance matrix estimation
            covar = self.objective.covar()
            errors = np.sqrt(np.diag(covar))
            res["covar"] = covar
            res["stderr"] = errors

            # check if the parameters are all Parameter instances.
            flat_params = list(f_unique(flatten(self.objective.parameters)))
            if np.all([is_parameter(param) for param in flat_params]):
                # zero out all the old parameter stderrs
                for param in flat_params:
                    param.stderr = None
                    param.chain = None

                for i, param in enumerate(_varying_parameters):
                    param.stderr = errors[i]

            # need to touch up the output to check we leave
            # parameters as we found them
            self.objective.setp(res.x)

        return res

def load_chain(f):
    """
    Loads a chain from disk. Does not change the state of a CurveFitter
    object.

    Parameters
    ----------
    f : str or file-like
        File containing the chain.

    Returns
    -------
    chain : array
        The loaded chain - `(nsteps, nwalkers, ndim)` or
        `(nsteps, ntemps, nwalkers, ndim)`
    """
    with possibly_open_file(f, "r") as g:
        # read header
        header = g.readline()
        expr = re.compile(r"(\d+)")
        matches = expr.findall(header)
        if matches:
            if len(matches) == 3:
                ntemps, nwalkers, ndim = map(int, matches)
            elif len(matches) == 2:
                ntemps = None
                nwalkers, ndim = map(int, matches)
        else:
            raise ValueError("Couldn't read header line of chain file")

    chain = np.loadtxt(f)

    if ntemps is not None:
        chain = np.reshape(chain, (-1, ntemps, nwalkers, ndim))
    else:
        chain = np.reshape(chain, (-1, nwalkers, ndim))

    return chain

def process_chain(objective, chain, nburn=0, nthin=1, flatchain=False):
    """
    Process the chain produced by a sampler for a given Objective

    Parameters
    ----------
    objective : refnx.analysis.Objective
        The Objective function that the Posterior was sampled for
    chain : array
        The MCMC chain
    nburn : int, optional
        discard this many steps from the start of the chain
    nthin : int, optional
        only accept every `nthin` samples from the chain
    flatchain : bool, optional
        collapse the walkers down into a single dimension.

    Returns
    -------
    [(param, stderr, chain)] : list
        List of (param, stderr, chain) tuples.
        If `isinstance(objective.parameters, Parameters)` then `param` is a
        `Parameter` instance. `param.value`, `param.stderr` and
        `param.chain` will contain the median, stderr and chain samples,
        respectively. Otherwise `param` will be a float representing the
        median of the chain samples.
        `stderr` is the half width of the [15.87, 84.13] spread (similar to
        standard deviation) and `chain` is an array containing the MCMC
        samples for that parameter.

    Notes
    -----
    The chain should have the shape `(iterations, nwalkers, nvary)` or
    `(iterations, ntemps, nwalkers, nvary)` if parallel tempering was
    employed.
    The burned and thinned chain is created via:
    `chain[nburn::nthin]`.
    Note, if parallel tempering is employed, then only the lowest temperature
    of the parallel tempering chain is processed and returned as it
    corresponds to the (lowest energy) target distribution.
    If `flatten is True` then the burned/thinned chain is reshaped and
    `arr.reshape(-1, nvary)` is returned.
    This function has the effect of setting the parameter stderr's.
    """
    chain = chain[nburn::nthin]
    shape = chain.shape
    nvary = shape[-1]

    # nwalkers = shape[1]
    if len(shape) == 4:
        ntemps = shape[1]
    elif len(shape) == 3:
        ntemps = -1

    if ntemps != -1:
        # PTSampler, we require the target distribution in the first row.
        chain = chain[:, 0]

    _flatchain = chain.reshape((-1, nvary))
    if flatchain:
        chain = _flatchain

    flat_params = list(f_unique(flatten(objective.parameters)))
    varying_parameters = objective.varying_parameters()

    # set the stderr of each of the Parameters
    result_list = []
    if np.all([is_parameter(param) for param in flat_params]):
        # zero out all the old parameter stderrs
        for param in flat_params:
            param.stderr = None
            param.chain = None

        # do the error calcn for the varying parameters and set the chain
        quantiles = np.percentile(_flatchain, [15.87, 50, 84.13], axis=0)
        for i, param in enumerate(varying_parameters):
            std_l, median, std_u = quantiles[:, i]
            param.value = median
            param.stderr = 0.5 * (std_u - std_l)

            # copy in the chain
            param.chain = np.copy(chain[..., i])
            res = MCMCResult(
                name=param.name,
                param=param,
                median=param.value,
                stderr=param.stderr,
                chain=param.chain,
            )
            result_list.append(res)

        fitted_values = np.array(varying_parameters)

        # give each constrained param a chain (to be reshaped later)
        constrained_params = [
            param for param in flat_params if param.constraint is not None
        ]

        for constrain_param in constrained_params:
            constrain_param.chain = np.empty(chain.shape[:-1], float)

        # now iterate through the varying parameters, set the values, thereby
        # setting the constraint value
        if len(constrained_params):
            for index in np.ndindex(chain.shape[:-1]):
                # iterate over parameter vectors
                pvals = chain[index]
                objective.setp(pvals)

                for constrain_param in constrained_params:
                    constrain_param.chain[index] = constrain_param.value

            for constrain_param in constrained_params:
                quantiles = np.percentile(
                    constrain_param.chain, [15.87, 50, 84.13]
                )

                std_l, median, std_u = quantiles
                constrain_param.value = median
                constrain_param.stderr = 0.5 * (std_u - std_l)

        # now reset fitted parameter values (they would've been changed by
        # constraints calculations
        objective.setp(fitted_values)

    # the parameter set are not Parameter objects, an array was probably
    # being used with BaseObjective.
    else:
        for i in range(nvary):
            c = np.copy(chain[..., i])
            median, stderr = uncertainty_from_chain(c)
            res = MCMCResult(
                name="", param=median, median=median, stderr=stderr, chain=c
            )
            result_list.append(res)

    return result_list

def uncertainty_from_chain(chain):
    """
    Calculates the median and uncertainty of MC samples.

    Parameters
    ----------
    chain : array-like

    Returns
    -------
    median, stderr : float, float
        `median` of the chain samples. `stderr` is half the width of the
        [15.87, 84.13] spread.
    """
    flatchain = chain.flatten()
    std_l, median, std_u = np.percentile(flatchain, [15.87, 50, 84.13])
    return median, 0.5 * (std_u - std_l)

def autocorrelation_chain(chain, nburn=0, nthin=1):
    """
    Calculate the autocorrelation function

    Parameters
    ----------
    chain : np.ndarray
        The MCMC chain - `(nsteps, nwalkers, ndim)` or
        `(nsteps, ntemps, nwalkers, ndim)`

    Returns
    -------
    acfs : np.ndarray
        The autocorrelation function, acfs.shape=(lags, nvary)
    """
    lchain = chain
    # parallel tempered chain
    if len(chain.shape) == 4:
        lchain = lchain[:, 0]

    lchain = lchain[nburn::nthin]
    # (iterations, walkers, vary) -> (vary, walkers, iterations)
    lchain = np.swapaxes(lchain, 0, 2)
    shape = lchain.shape[:-1]

    acfs = np.zeros_like(lchain)

    # iterate over each parameter/walker
    for index in np.ndindex(*shape):
        s = _function_1d(lchain[index])
        acfs[index] = s

    # now average over walkers
    acfs = np.mean(acfs, axis=1)
    return np.transpose(acfs)

def bounds_list(parameters):
    """
    Approximates interval bounds for a parameter set.

    Parameters
    ----------
    parameters : sequence
        A sequence containing individual parameters

    Returns
    -------
    bounds: tuple
        ``(min, max)`` pairs that define the finite lower and upper bounds
        every element in ``parameters``.

    If the `Bounds` applied by a parameter are a `PDF` instance then the upper
    and lower bound are approximated by ``PDF.rv.ppf([0.005, 0.995])``, which
    covers 99% of the statistical distribution.
    """
    bounds = []
    for param in parameters:
        if hasattr(param, "bounds") and isinstance(param.bounds, Interval):
            bnd = param.bounds
            bounds.append((bnd.lb, bnd.ub))
        elif (
            hasattr(param, "bounds")
            and isinstance(param.bounds, PDF)
            and hasattr(param.bounds.rv, "ppf")
        ):
            bounds.append(param.bounds.rv.ppf([0.005, 0.995]))
        else:
            # We can't handle this bound
            bounds.append((-np.inf, np.inf))

    return bounds

# Following code is for autocorrelation analysis of chains and is taken from
# emcee.autocorr
def _next_pow_two(n):
    """Returns the next power of two greater than or equal to `n`"""
    i = 1
    while i < n:
        i = i << 1
    return i

def _function_1d(x):
    """Estimate the normalized autocorrelation function of a 1-D series

    Args:
        x: The series as a 1-D numpy array.

    Returns:
        array: The autocorrelation function of the time series.

    """
    x = np.atleast_1d(x)
    if len(x.shape) != 1:
        raise ValueError("invalid dimensions for 1D autocorrelation function")
    n = _next_pow_two(len(x))

    # Compute the FFT and then (from that) the auto-correlation function
    f = np.fft.fft(x - np.mean(x), n=2 * n)
    acf = np.fft.ifft(f * np.conjugate(f))[: len(x)].real
    acf /= acf[0]
    return acf


"""
Created on 01 Apr 2021 21:40:47
@author: jiahuei
"""
import logging
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch import Tensor
from argparse import ArgumentParser, _ArgumentGroup
from typing import Optional, Union, Dict
from copy import deepcopy
from models import register_model
from models.transformer import (
    CachedMultiHeadedAttention, MultiHeadedAttention, PositionwiseFeedForward, PositionalEncoding,
    InputEmbedding as Embeddings, OutputEmbedding as Generator,
    LayerNorm, SublayerConnection,
    Decoder, DecoderLayer, CachedTransformerBase
)
from data.collate import ObjectRelationCollate
from utils.model_utils import repeat_tensors, pack_wrapper, clones
from utils.misc import str_to_bool

logger = logging.getLogger(__name__)

# noinspection PyAbstractClass
class EncoderComputerDecoder(nn.Module):
    """
    A standard Encoder-Decoder architecture. Base for this and many
    other models.
    """

    def __init__(self, encoder, computer, decoder, src_embed, tgt_embed, generator):
        super().__init__()
        self.encoder = encoder
        self.computer = computer
        self.decoder = decoder
        self.src_embed = src_embed
        self.tgt_embed = tgt_embed
        self.generator = generator

    def forward(self, src, boxes, tgt, src_mask, tgt_mask):
        """Take in and process masked src and target sequences."""
        enc_out = self.encode(src, boxes, src_mask)
        assert enc_out.size(0) == src_mask.size(0)
        if logger.isEnabledFor(logging.DEBUG):
            logger.debug(
                f"{self.__class__.__name__}: "
                f"Encoder output shape = `{enc_out.shape}`    "
                f"Target shape = `{tgt.shape}`"
            )
        com_out = self.compute(enc_out, src_mask)
        if com_out.size(0) != tgt.size(0):
            assert tgt.size(0) % com_out.size(0) == 0
            seq_per_img = int(tgt.size(0) / com_out.size(0))
            com_out = repeat_tensors(seq_per_img, com_out)
        return self.decode(com_out, None, tgt, tgt_mask), com_out

    def encode(self, src, boxes, src_mask):
        return self.encoder(self.src_embed(src), boxes, src_mask)

    def compute(self, memory, src_mask):
        return self.computer(memory, src_mask)

    def decode(self, memory, src_mask, tgt, tgt_mask):
        return self.decoder(self.tgt_embed(tgt), memory, src_mask, tgt_mask)

# noinspection PyAbstractClass
class Computer(nn.Module):
    """Generic N layer decoder with masking."""

    def __init__(self, layer, N):
        super().__init__()
        self.inputs = nn.Parameter(torch.Tensor(1, 5, layer.size))
        self.layers = clones(layer, N)
        self.norm = LayerNorm(layer.size)

    def forward(self, memory, src_mask):
        x = self.inputs
        x = repeat_tensors(memory.size(0), x)
        for layer in self.layers:
            x = layer(x, memory, src_mask)
        return self.norm(x)

# noinspection PyAbstractClass
class ComputerLayer(nn.Module):
    """Decoder is made of self-attn, src-attn, and feed forward (defined below)"""

    def __init__(self, size, self_attn, src_attn, feed_forward, dropout):
        super().__init__()
        self.size = size
        self.self_attn = self_attn
        self.src_attn = src_attn
        self.feed_forward = feed_forward
        self.sublayer = clones(SublayerConnection(size, dropout), 3)

    def forward(self, x, memory, src_mask):
        """Follow Figure 1 (right) for connections."""
        m = memory
        x = self.sublayer[0](x, lambda x: self.self_attn(x, x, x))
        x = self.sublayer[1](x, lambda x: self.src_attn(x, m, m, src_mask))
        return self.sublayer[2](x, self.feed_forward)

# noinspection PyAbstractClass
class Encoder(nn.Module):
    """Core encoder is a stack of N layers"""

    def __init__(self, layer, N):
        super().__init__()
        self.layers = clones(layer, N)
        self.norm = LayerNorm(layer.size)

    def forward(self, x, box, mask):
        """Pass the input (and mask) through each layer in turn."""
        for layer in self.layers:
            x = layer(x, box, mask)
        return self.norm(x)

# noinspection PyAbstractClass
class EncoderLayer(nn.Module):
    """Encoder is made up of self-attn and feed forward (defined below)"""

    def __init__(self, size, self_attn, feed_forward, dropout):
        super().__init__()
        self.self_attn = self_attn
        self.feed_forward = feed_forward
        self.sublayer = clones(SublayerConnection(size, dropout), 2)
        self.size = size

    def forward(self, x, box, mask):
        """Follow Figure 1 (left) for connections."""
        x = self.sublayer[0](x, lambda x: self.self_attn(x, x, x, box, mask))
        return self.sublayer[1](x, self.feed_forward)

# noinspection PyAbstractClass
class BoxMultiHeadedAttention(nn.Module):
    """
    Self-attention layer with relative position weights.
    Following the paper "Relation Networks for Object Detection" in https://arxiv.org/pdf/1711.11575.pdf
    """

    def __init__(self, h, d_model, trigonometric_embedding=True, dropout=0.1):
        """Take in model size and number of heads."""
        super().__init__()

        assert d_model % h == 0
        self.trigonometric_embedding = trigonometric_embedding

        # We assume d_v always equals d_k
        self.h = h
        self.d_k = d_model // h
        if self.trigonometric_embedding:
            self.dim_g = 64
        else:
            self.dim_g = 4
        geo_feature_dim = self.dim_g

        # matrices W_q, W_k, W_v, and one last projection layer
        self.linears = clones(nn.Linear(d_model, d_model), 4)
        self.WGs = clones(nn.Linear(geo_feature_dim, 1, bias=True), 8)

        # self.attn = None
        self.dropout = nn.Dropout(p=dropout)

    def forward(self, input_query, input_key, input_value, input_box, mask=None):
        """Implements Figure 2 of Relation Network for Object Detection"""
        if mask is not None:
            # Same mask applied to all h heads.
            mask = mask.unsqueeze(1)
        nbatches = input_query.size(0)

        # tensor with entries R_mn given by a hardcoded embedding of the relative position between bbox_m and bbox_n
        relative_geometry_embeddings = self.BoxRelationalEmbedding(
            input_box,
            trigonometric_embedding=self.trigonometric_embedding
        )
        flatten_relative_geometry_embeddings = relative_geometry_embeddings.view(-1, self.dim_g)

        # 1) Do all the linear projections in batch from d_model => h x d_k
        query, key, value = [
            l(x).view(nbatches, -1, self.h, self.d_k).transpose(1, 2)
            for l, x in zip(self.linears, (input_query, input_key, input_value))
        ]
        box_size_per_head = list(relative_geometry_embeddings.shape[:3])
        box_size_per_head.insert(1, 1)
        relative_geometry_weights_per_head = [
            ly(flatten_relative_geometry_embeddings).view(box_size_per_head) for ly in self.WGs
        ]
        relative_geometry_weights = torch.cat(relative_geometry_weights_per_head, 1)
        relative_geometry_weights = F.relu(relative_geometry_weights)

        # 2) Apply attention on all the projected vectors in batch.
        x, box_attn = self.box_attention(
            query, key, value, relative_geometry_weights, mask=mask, dropout=self.dropout
        )

        # 3) "Concat" using a view and apply a final linear.
        x = x.transpose(1, 2).contiguous().view(nbatches, -1, self.h * self.d_k)

        # # Legacy
        # x = input_value + x

        return self.linears[-1](x)

    @staticmethod
    def BoxRelationalEmbedding(f_g, dim_g=64, wave_len=1000, trigonometric_embedding=True):
        """
        Given a tensor with bbox coordinates for detected objects on each batch image,
        this function computes a matrix for each image

        with entry (i,j) given by a vector representation of the
        displacement between the coordinates of bbox_i, and bbox_j

        input: np.array of shape=(batch_size, max_nr_bounding_boxes, 4)
        output: np.array of shape=(batch_size, max_nr_bounding_boxes, max_nr_bounding_boxes, 64)
        """
        # returns a relational embedding for each pair of bboxes, with dimension = dim_g
        # follow implementation of https://github.com/heefe92/Relation_Networks-pytorch/blob/master/model.py#L1014-L1055

        batch_size = f_g.size(0)

        x_min, y_min, x_max, y_max = torch.chunk(f_g, 4, dim=-1)

        cx = (x_min + x_max) * 0.5
        cy = (y_min + y_max) * 0.5
        w = (x_max - x_min) + 1.
        h = (y_max - y_min) + 1.

        # cx.view(1,-1) transposes the vector cx, and so dim(delta_x) = (dim(cx), dim(cx))
        delta_x = cx - cx.view(batch_size, 1, -1)
        delta_x = torch.clamp(torch.abs(delta_x / w), min=1e-3)
        delta_x = torch.log(delta_x)

        delta_y = cy - cy.view(batch_size, 1, -1)
        delta_y = torch.clamp(torch.abs(delta_y / h), min=1e-3)
        delta_y = torch.log(delta_y)

        delta_w = torch.log(w / w.view(batch_size, 1, -1))
        delta_h = torch.log(h / h.view(batch_size, 1, -1))

        matrix_size = delta_h.size()
        delta_x = delta_x.view(batch_size, matrix_size[1], matrix_size[2], 1)
        delta_y = delta_y.view(batch_size, matrix_size[1], matrix_size[2], 1)
        delta_w = delta_w.view(batch_size, matrix_size[1], matrix_size[2], 1)
        delta_h = delta_h.view(batch_size, matrix_size[1], matrix_size[2], 1)

        position_mat = torch.cat((delta_x, delta_y, delta_w, delta_h), -1)

        if trigonometric_embedding:
            feat_range = torch.arange(dim_g / 8, device=f_g.device)
            dim_mat = feat_range / (dim_g / 8)
            dim_mat = 1. / (torch.pow(wave_len, dim_mat))

            dim_mat = dim_mat.view(1, 1, 1, -1)
            position_mat = position_mat.view(batch_size, matrix_size[1], matrix_size[2], 4, -1)
            position_mat = 100. * position_mat

            mul_mat = position_mat * dim_mat
            mul_mat = mul_mat.view(batch_size, matrix_size[1], matrix_size[2], -1)
            sin_mat = torch.sin(mul_mat)
            cos_mat = torch.cos(mul_mat)
            embedding = torch.cat((sin_mat, cos_mat), -1)
        else:
            embedding = position_mat
        return embedding

    @staticmethod
    def box_attention(query, key, value, box_relation_embds_matrix, mask=None, dropout=None):
        """
        Compute 'Scaled Dot Product Attention as in paper Relation Networks for Object Detection'.
        Follow the implementation in
        https://github.com/heefe92/Relation_Networks-pytorch/blob/master/model.py#L1026-L1055
        """

        N = value.size()[:2]
        dim_k = key.size(-1)
        dim_g = box_relation_embds_matrix.size()[-1]

        w_q = query
        w_k = key.transpose(-2, -1)
        w_v = value

        # attention weights
        scaled_dot = torch.matmul(w_q, w_k)
        scaled_dot = scaled_dot / np.sqrt(dim_k)
        if mask is not None:
            scaled_dot = scaled_dot.masked_fill(mask == 0, -1e9)

        # w_g = box_relation_embds_matrix.view(N,N)
        w_g = box_relation_embds_matrix
        w_a = scaled_dot
        # w_a = scaled_dot.view(N,N)

        # multiplying log of geometric weights by feature weights
        w_mn = torch.log(torch.clamp(w_g, min=1e-6)) + w_a
        w_mn = torch.nn.Softmax(dim=-1)(w_mn)
        if dropout is not None:
            w_mn = dropout(w_mn)

        output = torch.matmul(w_mn, w_v)

        return output, w_mn

# noinspection PyAbstractClass,PyAttributeOutsideInit
@register_model("ort_computer")
class ORTComputerModel(CachedTransformerBase):
    COLLATE_FN = ObjectRelationCollate

    def __init__(self, config):
        super().__init__(config)
        self.box_trigonometric_embedding = True
        self.make_model()

    def make_model(self, h=8, dropout=0.1):
        """Helper: Construct a model from hyperparameters."""
        bbox_attn = BoxMultiHeadedAttention(h, self.d_model, self.box_trigonometric_embedding)
        attn = CachedMultiHeadedAttention(h, self.d_model)
        self_attn = deepcopy(attn)
        self_attn.self_attention = True
        ff = PositionwiseFeedForward(self.d_model, self.dim_feedforward, dropout)
        position = PositionalEncoding(self.d_model, dropout)
        model = EncoderComputerDecoder(
            Encoder(EncoderLayer(
                self.d_model, deepcopy(bbox_attn), deepcopy(ff), dropout), self.num_layers
            ),
            Computer(ComputerLayer(
                self.d_model, MultiHeadedAttention(h, self.d_model), MultiHeadedAttention(h, self.d_model),
                deepcopy(ff), dropout), self.num_layers
            ),
            Decoder(DecoderLayer(
                self.d_model, self_attn, attn, deepcopy(ff), dropout), self.num_layers
            ),
            lambda x: x,  # nn.Sequential(Embeddings(self.d_model, src_vocab), deepcopy(position)),
            nn.Sequential(Embeddings(self.d_model, self.vocab_size), deepcopy(position)),
            Generator(self.d_model, self.vocab_size)
        )
        self.att_embed = nn.Sequential(
            nn.Linear(self.att_feat_size, self.d_model),
            nn.ReLU(),
            nn.Dropout(self.drop_prob_src)
        )
        # This was important from their code.
        # Initialize parameters with Glorot / fan_avg.
        for p in model.parameters():
            if p.dim() > 1:
                nn.init.xavier_uniform_(p)
        self.model = model

    def _prepare_feature(
            self,
            att_feats: Tensor,
            att_masks: Optional[Tensor] = None,
            boxes: Optional[Tensor] = None,
            seq: Optional[Tensor] = None
    ):

        att_feats, att_masks = self.clip_att(att_feats, att_masks)
        att_feats = pack_wrapper(self.att_embed, att_feats, att_masks)

        if att_masks is None:
            att_masks = att_feats.new_ones(att_feats.shape[:2], dtype=torch.long)
        att_masks = att_masks.unsqueeze(-2)

        if seq is not None:
            # crop the last one
            seq = seq[:, :-1]
            seq_mask = seq.data.ne(self.pad_idx)  # seq_mask: torch.Tensor
            seq_mask = seq_mask.unsqueeze(-2)
            seq_mask = seq_mask & self.subsequent_mask(seq.size(-1)).to(seq_mask)
        else:
            seq_mask = None

        return att_feats, boxes, seq, att_masks, seq_mask

    # noinspection PyMethodOverriding
    def _forward(self, att_feats, boxes, seqs, att_masks=None, **kwargs):
        att_feats, boxes, seq, att_masks, seq_mask = self._prepare_feature(att_feats, att_masks, boxes, seqs)
        out = self.model(att_feats, boxes, seq, att_masks, seq_mask)
        outputs = self.model.generator(out[0])
        return outputs #, self.model.generator(out[1])

    def get_logprobs_state(self, it, memory, mask, state):
        """
        state = [ys.unsqueeze(0)]
        """
        ys = it.unsqueeze(1)
        if state is None:
            pass
        else:
            # Retrieve reordered cache from state, and update them
            self._update_caches(state[1:])
        out = self.model.decode(
            memory, mask, ys, self.subsequent_mask(ys.size(1)).to(memory.device)
        )
        logprobs = self.model.generator(out[:, -1])
        # Add layer cache into state list, transposed so that beam_step can reorder them
        return logprobs, [ys.unsqueeze(0)] + self._retrieve_caches()

    # noinspection PyMethodOverriding
    def _sample(self, att_feats, boxes, att_masks=None, opt=None, **kwargs):
        if opt is None:
            opt = {}
        att_feats, boxes, seq, att_masks, seq_mask = self._prepare_feature(att_feats, att_masks, boxes)
        memory = self.model.encode(att_feats, boxes, att_masks)
        memory = self.model.compute(memory, att_masks)
        state = None
        mask = torch.ones(size=(memory.size(0), 1, memory.size(1)), dtype=memory.dtype, device=memory.device)
        return self._generate_captions(att_feats, mask, memory, state, opt)

    @staticmethod
    def clip_att(att_feats, att_masks):
        # Clip the length of att_masks and att_feats to the maximum length
        if att_masks is not None:
            max_len = att_masks.data.long().sum(1).max()
            att_feats = att_feats[:, :max_len].contiguous()
            att_masks = att_masks[:, :max_len].contiguous()
        return att_feats, att_masks

    @staticmethod
    def subsequent_mask(size):
        """Mask out subsequent positions."""
        attn_shape = (1, size, size)
        mask = torch.triu(torch.ones(attn_shape), diagonal=1).eq(0)
        return mask

    @staticmethod
    def add_argparse_args(parser: Union[_ArgumentGroup, ArgumentParser]):
        # fmt: off
        ORTComputerModel.COLLATE_FN.add_argparse_args(parser)
        CachedTransformerBase.add_argparse_args(parser)
        # Relation args
        parser.add_argument(
            "--box_trigonometric_embedding", type=str_to_bool,
            default=True
        )
        # fmt: on
        # return parser


import sys
import os
import ctypes
import glob
import warnings
import numpy as np
from math import fsum, sqrt

import MulensModel
from MulensModel.utils import Utils
try:
    import MulensModel.VBBL as mm_vbbl
except Exception:
    _vbbl_wrapped = False
else:
    _vbbl_wrapped = True
try:
    import MulensModel.AdaptiveContouring as mm_ac
except Exception:
    _adaptive_contouring_wrapped = False
else:
    _adaptive_contouring_wrapped = True

def _try_load(path, name):
    """
    Try loading compiled C library.
    Input is *str* or *list* of *str*.
    """
    if isinstance(path, str):
        path = [path]
    for path_ in path:
        try:
            out = ctypes.cdll.LoadLibrary(path_)
        except OSError:
            print("WARNING - File not loaded:", path_)
            print("Everything should work except:", name)
            pass
        else:
            return out
    return None

def _get_path_2(name_1, name_2):
    """convenience function"""
    module_path = os.path.abspath(__file__)
    for i in range(3):
        module_path = os.path.dirname(module_path)
    return os.path.join(module_path, 'source', name_1, name_2)

def _import_compiled_VBBL():
    """try importing manually compiled VBBL package"""
    vbbl = _try_load(
        _get_path_2('VBBL', "VBBinaryLensingLibrary_wrapper.so"), "VBBL")
    _vbbl_wrapped = (vbbl is not None)
    if not _vbbl_wrapped:
        return (_vbbl_wrapped, None, None)

    vbbl.VBBinaryLensing_BinaryMagDark.argtypes = 7 * [ctypes.c_double]
    vbbl.VBBinaryLensing_BinaryMagDark.restype = ctypes.c_double

    vbbl.VBBL_SG12_5.argtypes = 12 * [ctypes.c_double]
    vbbl.VBBL_SG12_5.restype = np.ctypeslib.ndpointer(
        dtype=ctypes.c_double, shape=(10,))

    return (_vbbl_wrapped,
            vbbl.VBBinaryLensing_BinaryMagDark, vbbl.VBBL_SG12_5)

def _import_compiled_AdaptiveContouring():
    """try importing manually compiled AdaptiveContouring package"""
    ac = "AdaptiveContouring"
    adaptive_contour = _try_load(_get_path_2(ac, ac + "_wrapper.so"), ac)
    _adaptive_contouring_wrapped = (adaptive_contour is not None)
    if not _adaptive_contouring_wrapped:
        return (_adaptive_contouring_wrapped, None)
    adaptive_contour.Adaptive_Contouring_Linear.argtypes = (
        8 * [ctypes.c_double])
    adaptive_contour.Adaptive_Contouring_Linear.restype = ctypes.c_double
    return (_adaptive_contouring_wrapped,
            adaptive_contour.Adaptive_Contouring_Linear)

# Check import and try manually compiled versions.
if _vbbl_wrapped:
    _vbbl_binary_mag_dark = mm_vbbl.VBBinaryLensing_BinaryMagDark
    _vbbl_SG12_5 = mm_vbbl.VBBL_SG12_5
else:
    out = _import_compiled_VBBL()
    _vbbl_wrapped = out[0]
    _vbbl_binary_mag_dark = out[1]
    _vbbl_SG12_5 = out[2]
if not _vbbl_wrapped:
    _solver = 'numpy'
else:
    _solver = 'Skowron_and_Gould_12'
if _adaptive_contouring_wrapped:
    _adaptive_contouring_linear = mm_ac.Adaptive_Contouring_Linear
else:
    out = _import_compiled_AdaptiveContouring()
    _adaptive_contouring_wrapped = out[0]
    _adaptive_contouring_linear = out[1]

class BinaryLens(object):
    """
    The binary lens equation - its solutions, images, parities,
    magnifications, etc.

    The binary lens equation is a 5th order complex polynomial.

    Attributes :
        mass_1: *float*
            mass of the primary (left-hand object) as a fraction of
            the total mass.

        mass_2: *float*
            mass of the secondary (right-hand object) as a fraction of the
            total mass.

        separation: *float*
            separation between the two bodies as a fraction of the Einstein
            ring.

    Note: mass_1 and mass_2 may be defined as a fraction of some other
    mass than the total mass. This is possible but not recommended -
    make sure you know what you're doing before you start using this
    possibility.

    """
    def __init__(self, mass_1=None, mass_2=None, separation=None):
        self.mass_1 = float(mass_1)  # This speeds-up code for np.float input.
        self.mass_2 = float(mass_2)
        self.separation = float(separation)
        self._total_mass = None
        self._mass_difference = None
        self._position_z1 = None
        self._position_z2 = None
        self._last_polynomial_input = None
        self._solver = _solver
        self._use_planet_frame = True

    def _calculate_variables(self, source_x, source_y):
        """calculates values of constants needed for polynomial coefficients"""
        self._total_mass = 0.5 * (self.mass_1 + self.mass_2)
        # This is total_mass in WM95 paper.

        self._mass_difference = 0.5 * (self.mass_2 - self.mass_1)
        self._zeta = source_x + source_y * 1.j
        if self._use_planet_frame:
            self._position_z1 = -self.separation + 0.j
            self._position_z2 = 0. + 0.j
        else:
            self._position_z1 = -0.5 * self.separation + 0.j
            self._position_z2 = 0.5 * self.separation + 0.j

    def _get_polynomial(self, source_x, source_y):
        """get polynomial coefficients"""
        if self._use_planet_frame:
            return self._get_polynomial_planet_frame(source_x, source_y)
        else:
            return self._get_polynomial_WM95(source_x, source_y)

    def _get_polynomial_WM95(self, source_x, source_y):
        """
        calculate coefficients of the polynomial in geometric center frame
        """
        # Calculate constants
        self._calculate_variables(source_x=source_x, source_y=source_y)
        total_m = self._total_mass
        total_m_pow2 = total_m * total_m

        m_diff = self._mass_difference
        m_diff_pow2 = m_diff * m_diff

        pos_z1 = self._position_z1

        z1_pow2 = pos_z1 * pos_z1
        z1_pow3 = z1_pow2 * pos_z1
        z1_pow4 = z1_pow2 * z1_pow2

        zeta = self._zeta
        zeta_conj = zeta.conjugate()
        zeta_conj_pow2 = zeta_conj * zeta_conj

        # Calculate the coefficients of the 5th order complex polynomial
        coeff_5 = Utils.complex_fsum([z1_pow2, -zeta_conj_pow2])
        coeff_4 = Utils.complex_fsum(
            [-2. * total_m * zeta_conj,
             zeta * zeta_conj_pow2, -2. * m_diff * pos_z1,
             -zeta * z1_pow2])
        coeff_3 = Utils.complex_fsum(
            [4. * total_m * zeta * zeta_conj,
             4. * m_diff * zeta_conj * pos_z1,
             2. * zeta_conj_pow2 * z1_pow2, -2. * z1_pow4])
        coeff_2 = Utils.complex_fsum(
            [4. * total_m_pow2 * zeta,
             4. * total_m * m_diff * pos_z1,
             -4. * m_diff * zeta * zeta_conj * pos_z1,
             -2. * zeta * zeta_conj_pow2 * z1_pow2,
             4. * m_diff * z1_pow3, 2. * zeta * z1_pow4])
        coeff_1 = Utils.complex_fsum(
            [-8. * total_m * m_diff * zeta * pos_z1,
             -4. * m_diff_pow2 * z1_pow2,
             -4. * total_m_pow2 * z1_pow2,
             -4. * total_m * zeta * zeta_conj * z1_pow2,
             -4. * m_diff * zeta_conj * z1_pow3,
             -zeta_conj_pow2 * z1_pow4, z1_pow3 * z1_pow3])
        coeff_0 = Utils.complex_fsum(
            [4. * m_diff_pow2 * zeta,
             4. * total_m * m_diff * pos_z1,
             4. * m_diff * zeta * zeta_conj * pos_z1,
             2. * total_m * zeta_conj * z1_pow2,
             zeta * zeta_conj_pow2 * z1_pow2,
             -2. * m_diff * z1_pow3 - zeta * z1_pow4])
        coeff_0 *= z1_pow2

        # Return the coefficients of the polynomial
        coeffs_list = [coeff_0, coeff_1, coeff_2, coeff_3, coeff_4, coeff_5]
        return np.array(coeffs_list).reshape(6)

    def _get_polynomial_planet_frame(self, source_x, source_y):
        """calculate coefficients of the polynomial in planet frame"""
        # Calculate constants
        self._calculate_variables(source_x=source_x, source_y=source_y)
        total_m = self._total_mass

        m_diff = self._mass_difference

        zeta = self._zeta
        zeta_conj = zeta.conjugate()

        c_sum = Utils.complex_fsum

        z1 = self._position_z1

        coeff_5 = c_sum([z1, -zeta_conj]) * zeta_conj
        coeff_4 = c_sum([
            (-m_diff + total_m) * z1,
            -c_sum([2. * total_m, z1 * c_sum([2. * z1, zeta])]) * zeta_conj,
            c_sum([2. * z1 + zeta]) * zeta_conj**2
            ])
        coeff_3 = c_sum([
            z1 * c_sum([m_diff * z1, -total_m * c_sum([z1, 2. * zeta])]),
            zeta_conj * c_sum([
                2. * m_diff * z1,
                c_sum([2. * total_m, z1**2]) * c_sum([z1, 2. * zeta])
                ]),
            -z1 * c_sum([z1, 2. * zeta]) * zeta_conj**2
            ])
        coeff_2 = c_sum([
            m_diff * z1 * c_sum([2. * total_m, z1 * zeta]),
            total_m * c_sum([
                -2. * total_m * z1, 4. * total_m * zeta, 3. * z1**2 * zeta]),
            -z1 * zeta_conj * c_sum([
                zeta * c_sum([6. * total_m, z1**2]),
                2. * m_diff * c_sum([z1, zeta])
                ]),
            z1**2 * zeta * zeta_conj**2
            ])
        coeff_1 = -z1 * (m_diff + total_m) * c_sum([
            m_diff * z1, -total_m * z1, 4. * total_m * zeta, z1**2 * zeta,
            -2. * z1 * zeta * zeta_conj
            ])
        coeff_0 = (m_diff + total_m)**2 * z1**2 * zeta

        coeffs_list = [coeff_0, coeff_1, coeff_2, coeff_3, coeff_4, coeff_5]
        return np.array(coeffs_list).reshape(6)

    def _get_polynomial_roots(self, source_x, source_y):
        """roots of the polynomial"""
        polynomial_input = [self.mass_1, self.mass_2, self.separation,
                            source_x, source_y]

        if polynomial_input == self._last_polynomial_input:
            return self._polynomial_roots

        polynomial = self._get_polynomial(
            source_x=source_x, source_y=source_y)

        np_polyroots = np.polynomial.polynomial.polyroots
        if self._solver == 'numpy':
            self._polynomial_roots = np_polyroots(polynomial)
        elif self._solver == 'Skowron_and_Gould_12':
            args = polynomial.real.tolist() + polynomial.imag.tolist()
            try:
                out = _vbbl_SG12_5(*args)
            except ValueError as err:
                err2 = "\n\nSwitching from Skowron & Gould 2012 to numpy"
                warnings.warn(str(err) + err2, UserWarning)
                self._solver = 'numpy'
                self._polynomial_roots = np_polyroots(polynomial)
            else:
                self._polynomial_roots = np.array([
                    out[0]+out[5]*1.j, out[1]+out[6]*1.j, out[2]+out[7]*1.j,
                    out[3]+out[8]*1.j, out[4]+out[9]*1.j])
        else:
            raise ValueError('Unknown solver: {:}'.format(self._solver))
        self._last_polynomial_input = polynomial_input

        return self._polynomial_roots

    def _polynomial_roots_ok(
            self, source_x, source_y, return_distances=False):
        """verified roots of polynomial i.e. roots of lens equation"""
        roots = self._get_polynomial_roots(
            source_x=source_x, source_y=source_y)

        # Two lines below are simplified assuming
        # self._position_z1.imag = 0 and same for z2.
        roots_conj = np.conjugate(roots)
        solutions = (self._zeta +
                     self.mass_1 / (roots_conj - self._position_z1) +
                     self.mass_2 / (roots_conj - self._position_z2))
        # This backs-up the lens equation.

        out = []
        distances = []
        for (i, root) in enumerate(roots):
            distances_from_root = abs((solutions-root)**2)
            min_distance_arg = np.argmin(distances_from_root)

            if i == min_distance_arg:
                out.append(root)
                distances.append(distances_from_root[min_distance_arg])
            # The values in distances[] are a diagnostic on how good the
            # numerical accuracy is.

        # If the lens equation is solved correctly, there should be
        # either 3 or 5 solutions (corresponding to 3 or 5 images)
        if len(out) not in [3, 5]:
            msg = ("Wrong number of solutions to the lens equation of binary" +
                   " lens.\nGot {:} and expected 3 or 5.\nThe parameters " +
                   "(m1, m2, s, source_x, source_y, solver) are:\n" +
                   "{:} {:} {:} {:} {:}  {:}\n\n" +
                   "Consider using 'point_source_point_lens' method for " +
                   "epochs when the source is very far from the lens. Note " +
                   "that it's different from 'point_source' method.")
            txt = msg.format(
                len(out), repr(self.mass_1), repr(self.mass_2),
                repr(self.separation), repr(source_x), repr(source_y),
                self._solver)

            if self._solver != "Skowron_and_Gould_12":
                txt += (
                    "\n\nYou should switch to using Skowron_and_Gould_12" +
                    " polynomial root solver. It is much more accurate than " +
                    "numpy.polynomial.polynomial.polyroots(). " +
                    "Skowron_and_Gould_12 method is selected in automated " +
                    "way if VBBL is imported properly.")
            else:
                distance = sqrt(source_x**2 + source_y**2)
                if (self.mass_2 > 1.e-6 * self.mass_1 and
                        (distance < 15. or distance < 2. * self.separation)):
                    txt += ("\n\nThis is surprising error - please contact " +
                            "code authors and provide the above error " +
                            "message.")
            txt += "\nMulensModel version: {:}".format(MulensModel.__version__)

            raise ValueError(txt)

        if return_distances:
            return (np.array(out), np.array(distances))
        else:
            return np.array(out)

    def _jacobian_determinant_ok(self, source_x, source_y):
        """determinants of lens equation Jacobian for verified roots"""
        roots_ok_bar = np.conjugate(self._polynomial_roots_ok(
                                   source_x=source_x, source_y=source_y))
        # Variable X_bar is conjugate of variable X.
        add_1 = self.mass_1 / (self._position_z1 - roots_ok_bar)**2
        add_2 = self.mass_2 / (self._position_z2 - roots_ok_bar)**2
        derivative = add_1 + add_2

        return 1. - derivative * np.conjugate(derivative)

    def _signed_magnification(self, source_x, source_y):
        """signed magnification for each image separately"""
        return 1. / self._jacobian_determinant_ok(
                source_x=source_x, source_y=source_y)

    def _point_source(self, source_x, source_y):
        """calculate point source magnification"""
        signed_magnification = self._signed_magnification(
            source_x=source_x, source_y=source_y)
        return fsum(abs(signed_magnification))

    def _point_source_Witt_Mao_95(self, source_x, source_y):
        """calculate point source magnification"""
        return self._point_source(source_x=source_x, source_y=source_y)

    def point_source_magnification(self, source_x, source_y):
        """
        Calculate point source magnification for given position. The
        origin of the coordinate system is at the center of mass and
        both masses are on X axis with higher mass at negative X; this
        means that the higher mass is at (X, Y)=(-s*q/(1+q), 0) and
        the lower mass is at (s/(1+q), 0).

        Parameters :
            source_x: *float*
                X-axis coordinate of the source.

            source_y: *float*
                Y-axis coordinate of the source.

        Returns :
            magnification: *float*
                Point source magnification.
        """
        if self._use_planet_frame:
            x_shift = -self.mass_1 / (self.mass_1 + self.mass_2)
        else:
            x_shift = self.mass_2 / (self.mass_1 + self.mass_2) - 0.5
        x_shift *= self.separation
        # We need to add this because in order to shift to correct frame.
        return self._point_source_Witt_Mao_95(
                source_x=float(source_x)+x_shift, source_y=float(source_y))
        # Casting to float speeds-up code for np.float input.

    def _get_magnification_w_plus(self, source_x, source_y, radius,
                                  magnification_center=None):
        """Evaluates Gould (2008) eq. 7"""
        dx = [1., 0., -1., 0.]
        dy = [0., 1., 0., -1.]
        out = []
        for (i, dxval) in enumerate(dx):
            x = source_x + dxval * radius
            y = source_y + dy[i] * radius
            out.append(self.point_source_magnification(
                                              source_x=x, source_y=y))
        if magnification_center is None:
            magnification_center = self.point_source_magnification(
                                    source_x=source_x, source_y=source_y)
        return 0.25 * fsum(out) - magnification_center

    def _get_magnification_w_times(self, source_x, source_y, radius,
                                   magnification_center=None):
        """Evaluates Gould (2008) eq. 8"""
        shift = radius / sqrt(2.)
        dx = [1., -1., -1., 1.]
        dy = [1., 1., -1., -1.]
        out = []
        for (i, dxval) in enumerate(dx):
            x = source_x + dxval * shift
            y = source_y + dy[i] * shift
            out.append(self.point_source_magnification(
                                              source_x=x, source_y=y))
        if magnification_center is None:
            magnification_center = self.point_source_magnification(
                                    source_x=source_x, source_y=source_y)
        return 0.25 * fsum(out) - magnification_center

    def _rho_check(self, rho):
        """
        Check if rho is float and positive.
        """
        if rho is None:
            raise TypeError(
                'rho must be positive float, but None was provided')
        if not isinstance(rho, float):
            raise TypeError(
                'rho must be positive float, but ' + str(rho) +
                str(type(rho)) + ' was provided')
        if rho < 0:
            raise ValueError(
                'rho must be positive, got: {:}'.format(rho))

    def hexadecapole_magnification(self, source_x, source_y, rho, gamma,
                                   quadrupole=False, all_approximations=False):
        """
        Magnification in hexadecapole approximation of the
        binary-lens/finite-source event - based on `Gould 2008 ApJ
        681, 1593
        <https://ui.adsabs.harvard.edu/abs/2008ApJ...681.1593G/abstract>`_.

        For coordinate system convention see
        :py:func:`point_source_magnification()`

        Parameters :
            source_x: *float*
                X-axis coordinate of the source.

            source_y: *float*
                Y-axis coordinate of the source.

            rho: *float*
                Source size relative to Einstein ring radius.

            gamma: *float*
                Linear limb-darkening coefficient in gamma convention.

            quadrupole: *boolean*, optional
                Return quadrupole approximation instead of hexadecapole?
                Default is *False*.

            all_approximations: *boolean*, optional
                Return hexadecapole, quadrupole, and point source
                approximations? Default is *False*.

        Returns :
            magnification: *float* or *sequence* of three *floats*
                Hexadecapole approximation (*float*) by default.
                Quadrupole approximation (*float*) if *quadrupole*
                parameter is *True*. Hexadecapole, quadrupole, and
                point source approximations (*sequence* of three
                *floats*) if *all_approximations* parameter is *True*.
        """
        # In this function, variables named a_* depict magnification.
        if quadrupole and all_approximations:
            raise ValueError('Inconsistent parameters of ' +
                             'BinaryLens.hexadecapole_magnification()')
        self._rho_check(rho)

        a_center = self.point_source_magnification(
            source_x=source_x, source_y=source_y)
        a_rho_half_plus = self._get_magnification_w_plus(
            source_x=source_x, source_y=source_y, radius=0.5*rho,
            magnification_center=a_center)
        a_rho_plus = self._get_magnification_w_plus(
            source_x=source_x, source_y=source_y, radius=rho,
            magnification_center=a_center)

        # This is Gould 2008 eq. 9:
        a_2_rho_square = (16. * a_rho_half_plus - a_rho_plus) / 3.

        # Gould 2008 eq. 6 (part 1/2):
        a_quadrupole = a_center + a_2_rho_square * (1. - 0.2 * gamma)

        # At this point is quadrupole approximation is finished
        if quadrupole:
            return a_quadrupole

        a_rho_times = self._get_magnification_w_times(
            source_x=source_x, source_y=source_y, radius=rho,
            magnification_center=a_center)

        # This is Gould (2008) eq. 9:
        a_4_rho_power4 = 0.5 * (a_rho_plus + a_rho_times) - a_2_rho_square
        # This is Gould (2008) eq. 6 (part 2/2):
        a_add = a_4_rho_power4 * (1. - 11. * gamma / 35.)
        a_hexadecapole = a_quadrupole + a_add

        if all_approximations:
            return (a_hexadecapole, a_quadrupole, a_center)
        else:
            return a_hexadecapole

    def adaptive_contouring_magnification(
            self, source_x, source_y, rho, gamma=None, u_limb_darkening=None,
            accuracy=0.1, ld_accuracy=0.001):
        """
        Binary lens finite source magnification calculated using
        Adaptive Contouring method by `Dominik 2007 MNRAS, 377, 1679
        <https://ui.adsabs.harvard.edu/abs/2007MNRAS.377.1679D/abstract>`_

        See also
        `AdaptiveContouring website by Martin Dominik
        <http://star-www.st-and.ac.uk/~md35/Software.html>`_

        For coordinate system convention see
        :py:func:`point_source_magnification()`

        Parameters :
            source_x: *float*
                X-axis coordinate of the source.

            source_y: *float*
                Y-axis coordinate of the source.

            rho: *float*
                Source size relative to Einstein ring radius.

            gamma: *float*, optional
                Linear limb-darkening coefficient in gamma convention.

            u_limb_darkening: *float*
                Linear limb-darkening coefficient in u convention.
                Note that either *gamma* or *u_limb_darkening* can be
                set.  If neither of them is provided then limb
                darkening is ignored.

            accuracy: *float*, optional
                Requested accuracy of the result defined as the sum of
                the area of the squares that determine the contour
                line and the estimated total enclosed area (see sec. 4
                of the paper).  As M. Dominik states: *"this vastly
                overestimates the fractional error, and a suitable
                value should be chosen by testing how its variation
                affects the final results - I recommend starting at
                acc = 0.1."* It significantly affects execution time.

            ld_accuracy: *float*, optional
                Requested limb-darkening accuracy. As M. Dominik
                states: *" Fractional uncertainty for the adaptive
                Simpson integration of the limb-darkening
                profile-related function during application of Green's
                theorem."* It does not add execution time so can be
                set to very small value.

        Returns :
            magnification: *float*
                Magnification.

        """
        if accuracy <= 0.:
            raise ValueError('adaptive_contouring requires accuracy > 0')
        if ld_accuracy <= 0.:
            raise ValueError('adaptive_contouring requires ld_accuracy > 0')
        # Note that this accuracy is not guaranteed.
        self._rho_check(rho)

        if not _adaptive_contouring_wrapped:
            raise ValueError('Adaptive Contouring was not imported properly')

        if gamma is not None and u_limb_darkening is not None:
            raise ValueError(
                'Only one limb darkening parameters can be set' +
                ' in BinaryLens.adaptive_contouring_magnification()')
        elif gamma is not None:
            gamma = float(gamma)
        elif u_limb_darkening is not None:
            gamma = float(Utils.u_to_gamma(u_limb_darkening))
        else:
            gamma = float(0.0)

        s = float(self.separation)
        q = float(self.mass_2 / self.mass_1)
        # AdaptiveContouring uses different coordinates conventions,
        # so we have to transform the coordinates below.
        x = float(-source_x)
        y = float(-source_y)
        rho = float(rho)
        accuracy = float(accuracy)
        ld_accuracy = float(ld_accuracy)

        magnification = _adaptive_contouring_linear(
            s, q, x, y, rho, gamma, accuracy, ld_accuracy)

        return magnification

    def vbbl_magnification(self, source_x, source_y, rho,
                           gamma=None, u_limb_darkening=None,
                           accuracy=0.001):
        """
        Binary lens finite source magnification calculated using VBBL
        library that implements advanced contour integration algorithm
        presented by `Bozza 2010 MNRAS, 408, 2188
        <https://ui.adsabs.harvard.edu/abs/2010MNRAS.408.2188B/abstract>`_.
        See also `VBBL website by Valerio Bozza
        <http://www.fisica.unisa.it/GravitationAstrophysics/VBBinaryLensing.htm>`_.

        For coordinate system convention see
        :py:func:`point_source_magnification()`

        Parameters :
            source_x: *float*
                X-axis coordinate of the source.

            source_y: *float*
                Y-axis coordinate of the source.

            rho: *float*
                Source size relative to Einstein ring radius.

            gamma: *float*, optional
                Linear limb-darkening coefficient in gamma convention.

            u_limb_darkening: *float*
                Linear limb-darkening coefficient in u convention.
                Note that either *gamma* or *u_limb_darkening* can be
                set.  If neither of them is provided then limb
                darkening is ignored.

            accuracy: *float*, optional
                Requested accuracy of the result.

        Returns :
            magnification: *float*
                Magnification.

        """
        self._rho_check(rho)
        if accuracy <= 0.:
            raise ValueError(
                "VBBL requires accuracy > 0 e.g. 0.01 or 0.001;" +
                "\n{:} was  provided".format(accuracy))

        if not _vbbl_wrapped:
            raise ValueError('VBBL was not imported properly')

        if gamma is not None and u_limb_darkening is not None:
            raise ValueError('Only one limb darkening parameters can be set' +
                             ' in BinaryLens.vbbl_magnification()')
        elif gamma is not None:
            u_limb_darkening = float(Utils.gamma_to_u(gamma))
        elif u_limb_darkening is not None:
            u_limb_darkening = float(u_limb_darkening)
        else:
            u_limb_darkening = float(0.0)

        s = float(self.separation)
        q = float(self.mass_2 / self.mass_1)
        x = float(source_x)
        y = float(source_y)
        rho = float(rho)
        accuracy = float(accuracy)

        magnification = _vbbl_binary_mag_dark(
            s, q, x, y, rho, u_limb_darkening, accuracy)

        return magnification


"""
General utility functions for the pymodulon package
"""
import json
import logging
import re
from itertools import combinations

import numpy as np
import pandas as pd
from scipy import stats
from scipy.special import digamma
from sklearn.neighbors import BallTree, KDTree

from pymodulon.enrichment import FDR

################
# Type Aliases #
################

def _check_table(table, name, index=None, index_col=0):
    # Set as empty dataframe if not input given
    if table is None:
        return pd.DataFrame(index=index)

    # Load table if necessary
    elif isinstance(table, str):
        try:
            table = pd.read_json(table)
        except ValueError:
            sep = "\t" if table.endswith(".tsv") else ","
            table = pd.read_csv(table, index_col=index_col, sep=sep)

    # Coerce indices and columns to ints if necessary
    newcols = []
    for col in table.columns:
        try:
            newcols.append(int(col))
        except ValueError:
            newcols.append(col)
    table.columns = newcols

    newrows = []
    for row in table.index:
        try:
            newrows.append(int(row))
        except ValueError:
            newrows.append(row)
    table.index = newrows

    # Replace empty strings with None
    table = table.replace("", np.nan)

    if isinstance(table, pd.DataFrame):
        # dont run _check_table_helper if no index is passed
        return table if index is None else _check_table_helper(table, index, name)
    else:
        raise TypeError(
            "{}_table must be a pandas DataFrame "
            "filename or a valid JSON string".format(name)
        )

def _check_table_helper(table, index, name):
    if table.shape == (0, 0):
        return pd.DataFrame(index=index)

    # Check if all indices are in table
    missing_index = list(set(index) - set(table.index))
    if len(missing_index) > 0:
        logging.warning(
            "Some {} are missing from the {} table: {}".format(
                name, name, missing_index
            )
        )

    # Remove extra indices from table
    table = table.reindex(index)
    return table

def _check_dict(table, index_col=0):
    try:
        table = json.loads(table.replace("'", '"'))
    except ValueError:
        sep = "\t" if table.endswith(".tsv") else ","
        table = pd.read_csv(table, index_col=index_col, header=None, sep=sep)
        table = table.to_dict()[1]
    return table

def compute_threshold(ic, dagostino_cutoff):
    """
    Computes D'agostino-test-based threshold for a component of an M matrix

    Parameters
    ----------
    ic: ~pandas.Series
        Pandas Series containing an independent component
    dagostino_cutoff: int
        Minimum D'agostino test statistic value to determine threshold

    Returns
    -------
    iModulon threshold: list
        List of thresholds for each iModulon
    """

    i = 0

    # Sort genes based on absolute value
    ordered_genes = abs(ic).sort_values()

    # Compute k2-statistic
    k_square, p = stats.normaltest(ic)

    # Iteratively remove gene w/ largest weight until k2-statistic < cutoff
    while k_square > dagostino_cutoff:
        i -= 1
        k_square, p = stats.normaltest(ic.loc[ordered_genes.index[:i]])

    # Select genes in iModulon
    comp_genes = ordered_genes.iloc[i:]

    # Slightly modify threshold to improve plotting visibility
    if len(comp_genes) == len(ic.index):
        return max(comp_genes) + 0.05
    else:
        return np.mean([ordered_genes.iloc[i], ordered_genes.iloc[i - 1]])

def dima(ica_data, sample1, sample2, threshold=5, fdr=0.1, alternate_A=None):
    """
    Creates DIMA table of differentially expressed iModulons

    Parameters
    ----------
    ica_data: ~pymodulon.core.IcaData
        :class:`~pymodulon.core.IcaData` data object
    sample1: str or list
        List of sample IDs or name of "project:condition"
    sample2: str or list
        List of sample IDs or name of "project:condition"
    threshold: float
        Minimum activity difference to determine DiMAs (default = 5)
    fdr: float
        False Detection Rate (default = .1)
    alternate_A: ~pandas.DataFrame
        Alternate A to use (default = None)

    Returns
    -------
    results: DataFrame
        Table of differentially expressed iModulons

    """

    # use the undocumented alternate_A option to allow custom-built DIMCA
    # activity matrix to be used in lieu of standard activty matrix
    if alternate_A is not None:
        A_to_use = alternate_A
    else:
        A_to_use = ica_data.A

    _diff = pd.DataFrame()

    sample1_list = _parse_sample(ica_data, sample1)
    sample2_list = _parse_sample(ica_data, sample2)

    for name, group in ica_data.sample_table.groupby(["project", "condition"]):
        for i1, i2 in combinations(group.index, 2):
            _diff[":".join(name)] = abs(A_to_use[i1] - A_to_use[i2])
    dist = {}

    for k in A_to_use.index:
        dist[k] = stats.lognorm(*stats.lognorm.fit(_diff.loc[k].values)).cdf

    res = pd.DataFrame(index=A_to_use.index)
    for k in res.index:
        a1 = A_to_use.loc[k, sample1_list].mean()
        a2 = A_to_use.loc[k, sample2_list].mean()
        res.loc[k, "difference"] = a2 - a1
        res.loc[k, "pvalue"] = 1 - dist[k](abs(a1 - a2))
    result = FDR(res, fdr)
    return result[(abs(result.difference) > threshold)].sort_values(
        "difference", ascending=False
    )

def _parse_sample(ica_data, sample):
    """
    Parses sample inputs into a list of sample IDs

    Parameters
    ----------
    ica_data: ~pymodulon.core.IcaData
        :class:`~pymodulon.core.IcaData` data object
    sample: list
        Sequence of sample IDs or "project:condition"

    Returns
    -------
    samples: list
        A list of `samples`
    """

    sample_table = ica_data.sample_table
    if isinstance(sample, str):
        proj, cond = re.search("(.*):(.*)", sample).groups()
        samples = sample_table[
            (sample_table.project == proj) & (sample_table.condition == cond)
        ].index
        if len(samples) == 0:
            raise ValueError(
                f"No samples exist for project={proj} condition=" f"{cond}"
            )
        else:
            return samples
    else:
        return sample

def explained_variance(
    ica_data, genes=None, samples=None, imodulons=None, reference=None
):
    """
    Computes the fraction of variance explained by iModulons (from 0 to 1)

    Parameters
    ----------
    ica_data: ~pymodulon.core.IcaData
        :class:`~pymodulon.core.IcaData` data object
    genes: str or list, optional
        List of genes to use (default: all genes)
    samples: str or list, optional
        List of samples to use (default: all samples)
    imodulons: int or str or list, optional
        List of iModulons to use (default: all iModulons)
    reference: list, optional
        List of samples that represent the reference condition for the
        set. If none are provided, uses the dataset-specific reference
        condition.

    Returns
    -------
    float
        Fraction of variance explained by selected iModulons for selected
        genes/samples
    """

    # Check inputs
    if genes is None:
        genes = ica_data.X.index
    elif isinstance(genes, str):
        genes = [genes]

    gene_loci = set(genes) & set(ica_data.X.index)
    gene_names = set(genes) - set(ica_data.X.index)
    name_loci = [ica_data.name2num(gene) for gene in gene_names]
    genes = list(set(gene_loci) | set(name_loci))

    if samples is None:
        samples = ica_data.X.columns
    elif isinstance(samples, str):
        samples = [samples]

    if imodulons is None:
        imodulons = ica_data.M.columns
    elif isinstance(imodulons, str) or isinstance(imodulons, int):
        imodulons = [imodulons]

    if reference is None:
        centered = ica_data.X
    else:
        centered = ica_data.X.subtract(ica_data.X[reference].mean(axis=1), axis=0)

    # Account for normalization procedures before ICA (X=SA-x_mean)
    baseline = centered.subtract(centered.mean(axis=0), axis=1)
    baseline = baseline.loc[genes, samples]

    # Initialize variables
    base_err = np.linalg.norm(baseline) ** 2
    MA = np.zeros(baseline.shape)
    rec_var = [0]
    ma_arrs = {}
    ma_weights = {}

    # Get individual modulon contributions
    for k in imodulons:
        ma_arr = np.dot(
            ica_data.M.loc[genes, k].values.reshape(len(genes), 1),
            ica_data.A.loc[k, samples].values.reshape(1, len(samples)),
        )
        ma_arrs[k] = ma_arr
        ma_weights[k] = np.sum(ma_arr ** 2)

    # Sum components in order of most important component first
    sorted_mods = sorted(ma_weights, key=ma_weights.get, reverse=True)
    # Compute reconstructed variance
    for k in sorted_mods:
        MA = MA + ma_arrs[k]
        sa_err = np.linalg.norm(MA - baseline) ** 2
        rec_var.append((1 - sa_err / base_err))

    return np.clip(rec_var[-1], 0, 1)

def infer_activities(ica_data, data):
    """
    Infer iModulon activities for external data

    Parameters
    ----------
    ica_data: ~pymodulon.core.IcaData
        :class:`~pymodulon.core.IcaData` data object
    data: ~pandas.DataFrame
        External expression profiles (must be centered to a reference)

    Returns
    -------
    new_activities: ~pandas.DataFrame
        Inferred activities for the expression profiles
    """

    shared_genes = ica_data.M.index & data.index
    x = data.loc[shared_genes].values
    m = ica_data.M.loc[shared_genes].values
    m_inv = np.linalg.pinv(m)
    a = np.dot(m_inv, x)
    return pd.DataFrame(a, index=ica_data.imodulon_names, columns=data.columns)

def mutual_info_distance(x, y):
    x = np.asarray(x).reshape(x.shape[0], 1)
    y = np.asarray(y).reshape(x.shape[0], 1)
    h = entropy(np.hstack([x, y]))
    if h == 0:
        return 1
    else:
        return 1 - mi(x, y) / h

# the following code is taken from the NPEET package; it cannot be installed
# via pip, so the necessary functions are copied here; the package appears to
# be un-maintained, so updates are not very likely; this is the GitHub page:
# https://github.com/gregversteeg/NPEET

def mi(x, y, z=None, k=3, base=2, alpha=0):
    """Mutual information of x and y (conditioned on z if z is not None)
    x, y should be a list of vectors, e.g. x = [[1.3], [3.7], [5.1], [2.4]]
    if x is a one-dimensional scalar and we have four samples
    """
    assert len(x) == len(y), "Arrays should have same length"
    assert k <= len(x) - 1, "Set k smaller than num. samples - 1"
    x, y = np.asarray(x), np.asarray(y)
    x, y = x.reshape(x.shape[0], -1), y.reshape(y.shape[0], -1)
    x = add_noise(x)
    y = add_noise(y)
    points = [x, y]
    if z is not None:
        z = np.asarray(z)
        z = z.reshape(z.shape[0], -1)
        points.append(z)
    points = np.hstack(points)
    # Find nearest neighbors in joint space, p=inf means max-norm
    tree = build_tree(points)
    dvec = query_neighbors(tree, points, k)
    if z is None:
        a, b, c, d = (
            avgdigamma(x, dvec),
            avgdigamma(y, dvec),
            digamma(k),
            digamma(len(x)),
        )
        if alpha > 0:
            d += lnc_correction(tree, points, k, alpha)
    else:
        xz = np.c_[x, z]
        yz = np.c_[y, z]
        a, b, c, d = (
            avgdigamma(xz, dvec),
            avgdigamma(yz, dvec),
            avgdigamma(z, dvec),
            digamma(k),
        )
    return max(0, (-a - b + c + d) / np.log(base))

def entropy(x, k=3, base=2):
    """The classic K-L k-nearest neighbor continuous entropy estimator
    x should be a list of vectors, e.g. x = [[1.3], [3.7], [5.1], [2.4]]
    if x is a one-dimensional scalar and we have four samples
    """
    assert k <= len(x) - 1, "Set k smaller than num. samples - 1"
    x = np.asarray(x)
    n_elements, n_features = x.shape
    x = add_noise(x)
    tree = build_tree(x)
    nn = query_neighbors(tree, x, k)
    const = digamma(n_elements) - digamma(k) + n_features * np.log(2)
    return max(0, (const + n_features * np.log(nn).mean()) / np.log(base))

def add_noise(x, intens=1e-10):
    """Small noise to break degeneracy, see doc."""
    return x + intens * np.random.random_sample(x.shape)

def build_tree(points):
    if points.shape[1] >= 20:
        return BallTree(points, metric="chebyshev")
    return KDTree(points, metric="chebyshev")

def query_neighbors(tree, x, k):
    return tree.query(x, k=k + 1)[0][:, k]

def avgdigamma(points, dvec):
    """This part finds number of neighbors in some radius in the marginal space
    returns expectation value of <psi(nx)>"""
    tree = build_tree(points)
    dvec = dvec - 1e-15
    num_points = count_neighbors(tree, points, dvec)
    return np.mean(digamma(num_points))

def lnc_correction(tree, points, k, alpha):
    e = 0
    n_sample = points.shape[0]
    for point in points:
        # Find k-nearest neighbors in joint space, p=inf means max norm
        knn = tree.query(point[None, :], k=k + 1, return_distance=False)[0]
        knn_points = points[knn]
        # Substract mean of k-nearest neighbor points
        knn_points = knn_points - knn_points[0]
        # Calculate covariance matrix of k-nearest neighbor points, obtain
        # eigen vectors
        covr = knn_points.T @ knn_points / k
        _, v = np.linalg.eig(covr)
        # Calculate PCA-bounding box using eigen vectors
        V_rect = np.log(np.abs(knn_points @ v).max(axis=0)).sum()
        # Calculate the volume of original box
        log_knn_dist = np.log(np.abs(knn_points).max(axis=0)).sum()

        # Perform local non-uniformity checking and update correction term
        if V_rect < log_knn_dist + np.log(alpha):
            e += (log_knn_dist - V_rect) / n_sample
    return e

def count_neighbors(tree, x, r):
    return tree.query_radius(x, r, count_only=True)