# Copyright 2022 The TEMPO Collaboration
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file  in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""
Module on for the original time evolving matrix product operator (TEMPO)
algorithm. This module is based on [Strathearn2017] and [Strathearn2018].

**[Strathearn2017]**
A. Strathearn, B.W. Lovett, and P. Kirton, *Efficient real-time path integrals
for non-Markovian spin-boson models*. New Journal of Physics, 19(9),
p.093009 (2017).

**[Strathearn2018]**
A. Strathearn, P. Kirton, D. Kilda, J. Keeling and
B. W. Lovett,  *Efficient non-Markovian quantum dynamics using
time-evolving matrix product operators*, Nat. Commun. 9, 3322 (2018).
"""

import sys
from typing import Callable, Dict, Optional, Text
import warnings
from copy import copy

import numpy as np
from numpy import ndarray
from scipy.linalg import expm
from oqupy.correlations import BaseCorrelations

from oqupy.bath import Bath
from oqupy.base_api import BaseAPIClass
from oqupy.config import NpDtype, MAX_DKMAX, DEFAULT_TOLERANCE
from oqupy.config import TEMPO_BACKEND_CONFIG
from oqupy.dynamics import Dynamics
from oqupy.system import BaseSystem
from oqupy.tempo.backends.tempo_backend import TempoBackend
from oqupy.operators import commutator, acommutator
from oqupy.util import get_progress

class TempoParameters(BaseAPIClass):
    r"""
    Parameters for the TEMPO computation.

    Parameters
    ----------
    dt: float
        Length of a time step :math:`\delta t`. - It should be small enough
        such that a trotterisation between the system Hamiltonian and the
        environment it valid, and the environment auto-correlation function
        is reasonably well sampled.
    dkmax: int
        Number of time steps :math:`K\in\mathbb{N}` that should be included in
        the non-Markovian memory. - It must be large
        enough such that :math:`\delta t \times K` is larger than the
        necessary memory time :math:`\tau_\mathrm{cut}`.
    epsrel: float
        The maximal relative error in the singular value truncation (done
        in the underlying tensor network algorithm). - It must be small enough
        such that the numerical compression (using tensor network algorithms)
        does not truncate relevant correlations.
    add_correlation_time: float
        Additional correlation time to include in the last influence
        functional as explained in [Strathearn2017].
    name: str (default = None)
        An optional name for the tempo parameters object.
    description: str (default = None)
        An optional description of the tempo parameters object.
    """
    def __init__(
            self,
            dt: float,
            dkmax: int,
            epsrel: float,
            add_correlation_time: Optional[float] = None,
            name: Optional[Text] = None,
            description: Optional[Text] = None) -> None:
        """Create a TempoParameters object."""
        self.dt = dt
        self.dkmax = dkmax
        self.epsrel = epsrel
        self.add_correlation_time = add_correlation_time
        super().__init__(name, description)

    def __str__(self) -> Text:
        ret = []
        ret.append(super().__str__())
        ret.append("  dt                   = {} \n".format(self.dt))
        ret.append("  dkmax                = {} \n".format(self.dkmax))
        ret.append("  epsrel               = {} \n".format(self.epsrel))
        ret.append("  add_correlation_time = {} \n".format(
            self.add_correlation_time))
        return "".join(ret)

    @property
    def dt(self) -> float:
        """Length of a time step."""
        return self._dt

    @dt.setter
    def dt(self, new_dt: float) -> None:
        try:
            tmp_dt = float(new_dt)
        except Exception as e:
            raise AssertionError("Argument 'dt' must be float.") from e
        assert tmp_dt > 0.0, \
            "Argument 'dt' must be bigger than 0."
        self._dt = tmp_dt

    @property
    def dkmax(self) -> float:
        """Number of time steps that should be included in the non-Markovian
        memory. """
        return self._dkmax

    @dkmax.setter
    def dkmax(self, new_dkmax: float) -> None:
        try:
            if new_dkmax is None:
                tmp_dkmax = None
            else:
                tmp_dkmax = int(new_dkmax)
        except Exception as e:
            raise AssertionError("Argument 'dkmax' must be int or None.") \
                from e
        assert tmp_dkmax is None or tmp_dkmax > 0, \
            "Argument 'dkmax' must be bigger than or equal to 0 or None."
        self._dkmax = tmp_dkmax

    @dkmax.deleter
    def dkmax(self) -> None:
        self._dkmax = None

    @property
    def epsrel(self) -> float:
        """The maximal relative error in the singular value truncation."""
        return self._epsrel

    @epsrel.setter
    def epsrel(self, new_epsrel: float) -> None:
        try:
            tmp_epsrel = float(new_epsrel)
        except Exception as e:
            raise AssertionError("Argument 'epsrel' must be float.") from e
        assert tmp_epsrel > 0.0, \
            "Argument 'epsrel' must be bigger than 0."
        self._epsrel = tmp_epsrel

    @property
    def add_correlation_time(self) -> float:
        """
        Additional correlation time to include in the last influence
        functional.
        """
        return self._add_correlation_time

    @add_correlation_time.setter
    def add_correlation_time(self, new_tau: Optional[float] = None) -> None:
        if new_tau is None:
            del self.add_correlation_time
        else:
            # check input: cutoff
            try:
                tmp_new_tau = float(new_tau)
            except Exception as e:
                raise AssertionError( \
                    "Additional correlation time must be a float.") from e
            if tmp_new_tau < 0:
                raise ValueError(
                    "Additional correlation time must be non-negative.")
            self._add_correlation_time = tmp_new_tau

    @add_correlation_time.deleter
    def add_correlation_time(self) -> None:
        self._add_correlation_time = None

class Tempo(BaseAPIClass):
    """
    Class representing the entire TEMPO tensornetwork as introduced in
    [Strathearn2018].

    Parameters
    ----------
    system: BaseSystem
        The system.
    bath: Bath
        The Bath (includes the coupling operator to the system).
    parameters: TempoParameters
        The parameters for the TEMPO computation.
    initial_state: ndarray
        The initial density matrix of the system.
    start_time: float
        The start time.
    backend: str (default = None)
        The name of the backend to use for the computation. If
        `backend` is ``None`` then the default backend is used.
    backend_config: dict (default = None)
        The configuration of the backend. If `backend_config` is
        ``None`` then the default backend configuration is used.
    name: str (default = None)
        An optional name for the tempo object.
    description: str (default = None)
        An optional description of the tempo object.
    """
    def __init__(
            self,
            system: BaseSystem,
            bath: Bath,
            parameters: TempoParameters,
            initial_state: ndarray,
            start_time: float,
            backend_config: Optional[Dict] = None,
            name: Optional[Text] = None,
            description: Optional[Text] = None) -> None:
        """Create a Tempo object. """

        assert isinstance(system, BaseSystem), \
            "Argument 'system' must be an instance of BaseSystem."
        self._system = system

        assert isinstance(bath, Bath), \
            "Argument 'bath' must be an instance of Bath."
        self._bath = bath

        self._correlations = self._bath.correlations

        assert isinstance(parameters, TempoParameters), \
            "Argument 'parameters' must be an instance of TempoParameters."
        self._parameters = parameters

        try:
            tmp_initial_state = np.array(initial_state, dtype=NpDtype)
            tmp_initial_state.setflags(write=False)
        except Exception as e:
            raise AssertionError("Initial state must be numpy array.") from e
        assert len(tmp_initial_state.shape) == 2, \
            "Initial state is not a matrix."
        assert tmp_initial_state.shape[0] == \
            tmp_initial_state.shape[1], \
            "Initial state is not a square matrix."
        self._initial_state = tmp_initial_state
        self._dimension = self._initial_state.shape[0]

        try:
            tmp_start_time = float(start_time)
        except Exception as e:
            raise AssertionError("Start time must be a float.") from e
        self._start_time = tmp_start_time

        if backend_config is None:
            self._backend_config = TEMPO_BACKEND_CONFIG
        else:
            self._backend_config = backend_config

        assert self._bath.dimension == self._dimension and \
            self._system.dimension == self._dimension, \
            "Hilbertspace dimensions are unequal: " \
            + "system ({}), ".format(self._system.dimension) \
            + "initial state ({}), ".format(self._dimension) \
            + "and bath coupling ({}), ".format(self._bath.dimension)

        super().__init__(name, description)

        tmp_coupling_comm = commutator(self._bath._coupling_operator)
        tmp_coupling_acomm = acommutator(self._bath._coupling_operator)
        self._coupling_comm = tmp_coupling_comm.diagonal()
        self._coupling_acomm = tmp_coupling_acomm.diagonal()

        self._dynamics = None
        self._backend_instance = None

        self._init_tempo_backend()

    def _init_tempo_backend(self):
        """Create and initialize the tempo backend. """
        dim = self._dimension
        initial_state = self._initial_state.reshape(dim**2)
        influence = self._influence
        unitary_transform = self._bath.unitary_transform
        propagators = self._propagators
        sum_north = np.array([1.0]*(dim**2))
        sum_west = np.array([1.0]*(dim**2))
        dkmax = self._parameters.dkmax
        epsrel = self._parameters.epsrel
        self._backend_instance = TempoBackend(
                initial_state,
                influence,
                unitary_transform,
                propagators,
                sum_north,
                sum_west,
                dkmax,
                epsrel,
                config=self._backend_config)

    def _init_dynamics(self):
        """Create a Dynamics object with metadata from the Tempo object. """
        name = None
        description = "computed from '{}' tempo".format(self.name)
        self._dynamics = Dynamics(name=name,
                                  description=description)

    def _influence(self, dk: int):
        """Create the influence functional matrix for a time step distance
        of dk. """
        return influence_matrix(
            dk,
            parameters=self._parameters,
            correlations=self._correlations,
            coupling_acomm=self._coupling_acomm,
            coupling_comm=self._coupling_comm)

    def _propagators(self, step: int):
        """Create the system propagators (first and second half) for the time
        step `step`. """
        dt = self._parameters.dt
        t = self._time(step)
        first_step = expm(self._system.liouvillian(t+dt/4.0)*dt/2.0)
        second_step = expm(self._system.liouvillian(t+dt*3.0/4.0)*dt/2.0)
        return first_step, second_step

    def _time(self, step: int):
        """Return the time that corresponds to the time step `step`. """
        return self._start_time + float(step)*self._parameters.dt

    @property
    def dimension(self) -> ndarray:
        """Hilbert space dimension. """
        return copy(self._dimension)

    def compute(
            self,
            end_time: float,
            progress_type: Text = None) -> Dynamics:
        """
        Propagate (or continue to propagate) the TEMPO tensor network to
        time `end_time`.

        Parameters
        ----------
        end_time: float
            The time to which the TEMPO should be computed.
        progress_type: str (default = None)
            The progress report type during the computation. Types are:
            {``'silent'``, ``'simple'``, ``'bar'``}. If `None` then
            the default progress type is used.

        Returns
        -------
        dynamics: Dynamics
            The instance of Dynamics associated with the TEMPO object.
        """
        try:
            tmp_end_time = float(end_time)
        except Exception as e:
            raise AssertionError("End time must be a float.") from e

        dim = self._dimension
        if self._backend_instance.step is None:
            step, state = self._backend_instance.initialize()
            self._init_dynamics()
            self._dynamics.add(self._time(step), state.reshape(dim, dim))

        start_step = self._backend_instance.step
        end_step = int((end_time - self._start_time)/self._parameters.dt)
        num_step = max(0, end_step - start_step)

        progress = get_progress(progress_type)
        with progress(num_step) as prog_bar:
            while self._time(self._backend_instance.step) < tmp_end_time:
                step, state = self._backend_instance.compute_step()
                self._dynamics.add(self._time(step), state.reshape(dim, dim))
                prog_bar.update(self._backend_instance.step - start_step)
            prog_bar.update(self._backend_instance.step - start_step)

        return self._dynamics

    def get_dynamics(self) -> Dynamics:
        """Returns the instance of Dynamics associated with the Tempo object.
        """
        return self._dynamics

def influence_matrix(
        dk: int,
        parameters: TempoParameters,
        correlations: BaseCorrelations,
        coupling_acomm: ndarray,
        coupling_comm: ndarray):
    """Compute the influence functional matrix. """
    dt = parameters.dt
    dkmax = parameters.dkmax

    if dk == 0:
        time_1 = 0.0
        time_2 = None
        shape = "upper-triangle"
    elif dk < 0:
        time_1 = float(dkmax) * dt
        if parameters.add_correlation_time is not None:
            time_2 = float(dkmax) * dt \
                + np.min([float(-dk) * dt,
                            1.0*dt + parameters.add_correlation_time])
        else:
            return None
        shape = "rectangle"
    else:
        time_1 = float(dk) * dt
        time_2 = None
        shape = "square"

    eta_dk = correlations.correlation_2d_integral( \
        delta=dt,
        time_1=time_1,
        time_2=time_2,
        shape=shape,
        epsrel=parameters.epsrel)
    op_p = coupling_acomm
    op_m = coupling_comm

    if dk == 0:
        infl = np.diag(np.exp(-op_m*(eta_dk.real*op_m \
                                        + 1j*eta_dk.imag*op_p)))
    else:
        infl = np.exp(-np.outer(eta_dk.real*op_m \
                                + 1j*eta_dk.imag*op_p, op_m))

    return infl

def _analyse_correlation(
        corr_func: Callable[[np.ndarray],np.ndarray],
        times: np.ndarray,
        corr_vals: np.ndarray):
    """Check correlation function on a finer grid."""
    additional_times = (times[:-1] + times[1:])/2.0
    additional_corr_vals = corr_func(additional_times)
    new_times = list(times)
    new_corr_vals = list(corr_vals)
    for i in range(len(additional_times)):
        new_times.insert(2*i+1,additional_times[i])
        new_corr_vals.insert(2*i+1,additional_corr_vals[i])

    errors = []
    integrals = []
    integral = 0.0

    for i in range(len(times)-1):
        dt = new_times[2*i+2] - new_times[2*i]

        rough_int = 0.5 * dt * (new_corr_vals[2*i] + new_corr_vals[2*i+2])
        fine_int = 0.5 * (rough_int + dt * new_corr_vals[2*i+1])
        error = np.abs(rough_int-fine_int)
        errors.append(error)

        rough_abs_int = 0.5 * dt \
                * (np.abs(new_corr_vals[2*i]) + np.abs(new_corr_vals[2*i+2]))
        fine_abs_int = 0.5 * (rough_abs_int + dt * np.abs(new_corr_vals[2*i+1]))
        integral += fine_abs_int
        integrals.append(integral)

    full_abs_integral = integrals[-1]

    new_times = np.array(new_times)
    new_corr_val = np.array(new_corr_vals)
    errors = np.array(errors) / full_abs_integral
    integrals = np.array(integrals) / full_abs_integral

    return new_times, new_corr_val, errors, integrals

def _estimate_epsrel(
        dkmax: int,
        tolerance: float) -> float:
    """Heuristic estimation of appropriate epsrel for TEMPO."""
    power = np.log(dkmax)/np.log(4)-np.log(tolerance)/np.log(10)
    return np.power(10,-power)

GUESS_WARNING_MSG = "Estimating parameters for TEMPO computation. " \
    + "No guarantee that resulting TEMPO computation converges towards " \
    + "the correct dynamics! " \
    + "Please refer to the TEMPO documentation and check convergence by " \
    + "varying the parameters for TEMPO manually."

MAX_DKMAX_WARNING_MSG = f"Reached maximal recommended `dkmax` ({MAX_DKMAX})! " \
    + "Interrupt TEMPO parameter estimation. "\
    + "Please choose a lower tolerance, or analyse the correlation function " \
    + "to choose TEMPO parameters manually. " \
    + "Could not reach specified tolerance! "

def guess_tempo_parameters(
        bath: Bath,
        start_time: float,
        end_time: float,
        system: Optional[BaseSystem] = None,
        tolerance: Optional[float] = DEFAULT_TOLERANCE) -> TempoParameters:
    """
    Function to roughly estimate appropriate parameters for a TEMPO
    computation.

    .. warning::

        No guarantee that resulting TEMPO calculation converges towards the
        correct dynamics! Please refer to the TEMPO documentation and check
        convergence by varying the parameters for TEMPO manually.

    Parameters
    ----------
    bath: Bath
        The bath.
    start_time: float
        The start time.
    end_time: float
        The time to which the TEMPO should be computed.
    system: BaseSystem
        The system.
    tolerance: float
        Tolerance for the parameter estimation.

    Returns
    -------
    tempo_parameters : TempoParameters
        Estimate of appropriate tempo parameters.
    """
    assert isinstance(bath, Bath), \
        "Argument 'bath' must be a oqupy.Bath object."
    try:
        tmp_start_time = float(start_time)
        tmp_end_time = float(end_time)
    except Exception as e:
        raise AssertionError("Start and end time must be a float.") from e
    if tmp_end_time <= tmp_start_time:
        raise ValueError("End time must be bigger than start time.")
    assert isinstance(system, (type(None), BaseSystem)), \
        "Argument 'system' must be 'None' or a oqupy.BaseSystem object."
    try:
        tmp_tolerance = float(tolerance)
    except Exception as e:
        raise AssertionError("Argument 'tolerance' must be float.") from e
    assert tmp_tolerance > 0.0, \
        "Argument 'tolerance' must be larger then 0."
    warnings.warn(GUESS_WARNING_MSG, UserWarning)
    print("WARNING: "+GUESS_WARNING_MSG, file=sys.stderr, flush=True)

    max_tau = tmp_end_time - tmp_start_time

    corr_func = np.vectorize(bath.correlations.correlation)
    new_times = np.linspace(0, max_tau, 11, endpoint=True)
    new_corr_vals = corr_func(new_times)
    times = new_times
    corr_vals = new_corr_vals

    while True:
        if len(new_times) > MAX_DKMAX:
            warnings.warn(MAX_DKMAX_WARNING_MSG, UserWarning)
            break
        times = new_times
        corr_vals = new_corr_vals
        new_times, new_corr_vals, errors, integrals = \
                _analyse_correlation(corr_func, times, corr_vals)
        cut = np.where(integrals>(1-tolerance))[0][0]
        cut = cut+2 if cut+2<=len(times) else len(times)
        times = times[:cut]
        corr_vals = corr_vals[:cut]
        new_times = new_times[:2*cut-1]
        new_corr_vals = new_corr_vals[:2*cut-1]
        if (errors < tolerance).all():
            break

    dt = np.min(times[1:] - times[:-1])
    dkmax = len(times)
    epsrel = _estimate_epsrel(dkmax, tolerance)
    sys.stderr.flush()

    return TempoParameters(
        dt=dt,
        dkmax=dkmax,
        epsrel=epsrel,
        name="Roughly estimated parameters",
        description="Estimated with 'guess_tempo_parameters()'")

def tempo_compute(
        system: BaseSystem,
        bath: Bath,
        initial_state: ndarray,
        start_time: float,
        end_time: float,
        parameters: Optional[TempoParameters] = None,
        tolerance: Optional[float] = DEFAULT_TOLERANCE,
        backend_config: Optional[Dict] = None,
        progress_type: Optional[Text] = None,
        name: Optional[Text] = None,
        description: Optional[Text] = None) -> Dynamics:
    """
    Shortcut for creating a Tempo object and running the computation.

    Parameters
    ----------
    system: BaseSystem
        The system.
    bath: Bath
        The Bath (includes the coupling operator to the system).
    initial_state: ndarray
        The initial density matrix of the system.
    start_time: float
        The start time.
    end_time: float
        The time to which the TEMPO should be computed.
    parameters: TempoParameters
        The parameters for the TEMPO computation.
    tolerance: float
        Tolerance for the parameter estimation (only applicable if
        `parameters` is None).
    backend_config: dict (default = None)
        The configuration of the backend. If `backend_config` is
        ``None`` then the default backend configuration is used.
    progress_type: str (default = None)
        The progress report type during the computation. Types are:
        {``'silent'``, ``'simple'``, ``'bar'``}.  If `None` then
        the default progress type is used.
    name: str (default = None)
        An optional name for the tempo object.
    description: str (default = None)
        An optional description of the tempo object.
    """
    if parameters is None:
        assert tolerance is not None, \
            "If 'parameters' is 'None' then 'tolerance' must be " \
            + "a positive float."
        parameters = guess_tempo_parameters(bath=bath,
                                            start_time=start_time,
                                            end_time=end_time,
                                            system=system,
                                            tolerance=tolerance)
    tempo = Tempo(system,
                  bath,
                  parameters,
                  initial_state,
                  start_time,
                  backend_config,
                  name,
                  description)
    tempo.compute(end_time, progress_type=progress_type)
    return tempo.get_dynamics()


# It defines a new Observation class, subclassed from CompositeSourceSpectrum,
# that has some special methods and attributes and explicitly removes
# certain other methods.
"""This module handles an observation and related calculations."""
from __future__ import division

import logging
import numpy as np
import math

from . import spectrum
from . import units
from . import binning
from . import exceptions

from .obsbandpass import pixel_range, wave_range
from .spectrum import ArraySourceSpectrum

try:
    import pysynphot_utils
    utils_imported = True
except ImportError:
    utils_imported = False

def check_overlap(a, b):
    """Check for wavelength overlap between two spectra.

    .. note::

        Generalized from
        :meth:`pysynphot.spectrum.SpectralElement.check_overlap`.

    Parameters
    ----------
    a, b : `~pysynphot.spectrum.SourceSpectrum` or `~pysynphot.spectrum.SpectralElement`
       Typically a source spectrum, spectral element, observation,
       or bandpass from observation mode.

    Returns
    -------
    result : {'full', 'partial', 'none'}
        Full, partial, or no overlap.

    Raises
    ------
    AttributeError
        Given spectrum does not have flux or throughput.

    """
    if a.isAnalytic or b.isAnalytic:
        #then it's defined everywhere
        result = 'full'
    else:
        #get the wavelength arrays
        waves = list()
        for x in (a, b):
            if hasattr(x,'throughput'):
                wv = x.wave[np.where(x.throughput != 0)]
            elif hasattr(x,'flux'):
                wv = x.wave
            else:
                raise AttributeError("neither flux nor throughput in %s"%x)
            waves.append(wv)

        #get the endpoints
        a1,a2 = waves[0].min(), waves[0].max()
        b1,b2 = waves[1].min(), waves[1].max()

        #do the comparison
        if (a1>=b1 and a2<=b2):
            result = 'full'
        elif (a2<b1) or (b2<a1):
            result = 'none'
        else:
            result = 'partial'

    return result

def validate_overlap(comp1, comp2, force):
    """Validate the overlap between the wavelength sets
    of the two given components.

    Parameters
    ----------
    comp1, comp2 : `~pysynphot.spectrum.SourceSpectrum` or `~pysynphot.spectrum.SpectralElement`
        Source spectrum and bandpass of an observation.

    force : {'extrap', 'taper', `None`}
        If not `None`, the components may be adjusted by
        extrapolation or tapering.

    Returns
    -------
    comp1, comp2
        Same as inputs. However, ``comp1`` might be tapered
        if that option is selected.

    warnings : dict
        Maps warning keyword to its description.

    Raises
    ------
    KeyError
        Invalid ``force``.

    pysynphot.exceptions.DisjointError
        No overlap detected when ``force`` is `None`.

    pysynphot.exceptions.PartialOverlap
        Partial overlap detected when ``force`` is `None`.

    """
    warnings = dict()
    if force is None:
        stat = comp2.check_overlap(comp1)
        if stat=='full':
            pass
        elif stat == 'partial':
            raise(exceptions.PartialOverlap('Spectrum and bandpass do not fully overlap. You may use force=[extrap|taper] to force this Observation anyway.'))
        elif stat == 'none':
            raise(exceptions.DisjointError('Spectrum and bandpass are disjoint'))

    elif force.lower() == 'taper':
        try:
            comp1=comp1.taper()
        except AttributeError:
            comp1=comp1.tabulate().taper()
            warnings['PartialOverlap']=force

    elif force.lower().startswith('extrap'):
        #default behavior works, but check the overlap so we can set the warning
        stat=comp2.check_overlap(comp1)
        if stat == 'partial':
            warnings['PartialOverlap']=force

    else:
        raise(KeyError("Illegal value force=%s; legal values=('taper','extrap')"%force))
    return comp1, comp2, warnings

class Observation(spectrum.CompositeSourceSpectrum):
    """Class to handle an :ref:`observation <pysynphot-observation>`.
    An observation is the end point of a chain of spectral manipulation.

    Most `~pysynphot.obsbandpass.ObsBandpass` objects have a built-in
    ``binset`` that is optimized for use with the specified observing
    mode (also see :ref:`pysynphot-wavelength-table`).
    Specifying the ``binset`` here would override the built-in one.

    Parameters
    ----------
    spec : `~pysynphot.spectrum.SourceSpectrum`
        Source spectrum.

    band : `~pysynphot.spectrum.SpectralElement`
        Bandpass.

    binset : array_like or `None`
        Wavelength values to be used for binning when converting to counts.
        See :meth:`initbinset`.

    force
        See :meth:`~pysynphot.observation.Observation.validate_overlap`.

    Attributes
    ----------
    spectrum
        Same as input ``spec``.

    bandpass
        Same as input ``band``.

    binset
        Same as input ``binset``.

    component1, component2 : `~pysynphot.spectrum.SourceSpectrum` or `~pysynphot.spectrum.SpectralElement`
        Components and sub-components that belong to the observation.

    operation : str
        This is always "multiply".

    name : str
        Short description of the observation.

    warnings : dict
        To store warnings, which are inherited from all inputs. If they have the same warning keyword, the one from most recently read component is used.

    isAnalytic : bool
        Flag to indicate whether this is an analytic spectrum. This is only `True` if both inputs are analytic.

    primary_area : number or `None`
        :ref:`pysynphot-area` of the telescope. This is inherited from either of the inputs, if available (not `None`). If inputs have different values, an exception is raised.

    waveunits, fluxunits : `~pysynphot.units.Units`
        User units inherited from source spectrum.

    wave, flux : array_like
        Wavelength set and associated flux in user units. This is the native dataset.

    binwave, binflux : array_like
        Binned dataset.

    Raises
    ------
    pysynphot.exceptions.IncompatibleSources
        Input spectra have different telescope areas defined.

    """
    def __init__(self,spec,band,binset=None,force=None):
        self.spectrum = spec
        self.bandpass = band
        self.warnings={}
        self.validate_overlap(force)
        self.binset = binset

        keep=self.warnings
        spectrum.CompositeSourceSpectrum.__init__(self,
                                                  self.spectrum,
                                                  self.bandpass,
                                                  'multiply')

        self.warnings.update(keep)

        #The natural waveset of the observation is the merge of the
        #natural waveset of the spectrum with the natural waveset of the
        #bandpass. Because the Observation inherits from a
        #CompositeSourceSpectrum, this will be handled correctly.
#        self._binwave = None
        self._binflux = None

        self.initbinset(binset)
        #self.initbinflux()

    def validate_overlap(self,force):
        """Validate that spectrum and bandpass overlap.
        Warnings are stored in ``self.warnings``.

        Parameters
        ----------
        force : {'extrap', 'taper', `None`}
            If `None`, it is required that the spectrum and bandpass fully
            overlap. Partial overlap is allowed if this is set to
            ``'extrap'`` or ``'taper'``. See :func:`validate_overlap`.

        """

        #Wrap the function for convenience
        self.spectrum, self.bandpass, warn = validate_overlap(self.spectrum,
                                                              self.bandpass, force)
        self.warnings.update(warn)

    def initbinset(self,binset=None):
        """Set ``self.binwave``.

        By default, wavelength values for binning are inherited
        from bandpass. If the bandpass has no binning information,
        then source spectrum wavelengths are used. However, if
        user provides values, then those are used without question.

        Parameters
        ----------
        binset : array_like or `None`
            Wavelength values to be used for binning when converting to counts.

        """
        if binset is None:
            msg="(%s) does not have a defined binset in the wavecat table. The waveset of the spectrum will be used instead."%str(self.bandpass)

            try:
                self.binwave = self.bandpass.binset
            except (KeyError, AttributeError):
                self.binwave = self.spectrum.wave
                logging.warning(msg)

            if self.binwave is None:
                self.binwave = self.spectrum.wave
                logging.warning(msg)
        else:
            self.binwave=binset

    def initbinflux(self):
        """Calculate binned flux and edges.

        Flux is computed by integrating the spectrum
        on the specified binned wavelength set, using
        information from the natural wavelength set.

        Native wave/flux arrays should be considered samples
        of a continuous function, but not their binned counterparts.
        Thus, it makes sense to interpolate ``(wave, flux)`` but not
        ``(binwave, binflux)``.

        .. note::

            Assumes that the wavelength values in the binned
            wavelength set are the *centers* of the bins.

            Uses ``pysynphot.pysynphot_utils.calcbinflux()`` C-extension,
            if available, for binned flux calculation.

        """
        endpoints = binning.calculate_bin_edges(self.binwave)

        # merge these endpoints in with the natural waveset
        spwave = spectrum.MergeWaveSets(self.wave, endpoints)
        spwave = spectrum.MergeWaveSets(spwave,self.binwave)

        # compute indices associated to each endpoint.
        indices = np.searchsorted(spwave, endpoints)
        self._indices = indices[:-1]
        self._indices_last = indices[1:]

        # prepare integration variables.
        flux = self(spwave)
        avflux = (flux[1:] + flux[:-1]) / 2.0
        self._deltaw = spwave[1:] - spwave[:-1]

        # sum over each bin.
        if utils_imported is True:
            self._binflux, self._intwave = \
              pysynphot_utils.calcbinflux(len(self.binwave),
                                          self._indices,
                                          self._indices_last,
                                          avflux,
                                          self._deltaw)
        else:
            #Note that, like all Python striding, the range over which
            #we integrate is [first:last).
            self._binflux = np.empty(shape=self.binwave.shape,dtype=np.float64)
            self._intwave = np.empty(shape=self.binwave.shape,dtype=np.float64)
            for i in range(len(self._indices)):
                first = self._indices[i]
                last = self._indices_last[i]
                self._binflux[i]=(avflux[first:last]*self._deltaw[first:last]).sum()/self._deltaw[first:last].sum()
                self._intwave[i]=self._deltaw[first:last].sum()

        #Save the endpoints for future use
        self._bin_edges = endpoints

    def _getBinfluxProp(self):
        if self._binflux is None:
            self.initbinflux()

        if hasattr(self.bandpass, 'primary_area'):
            area = self.bandpass.primary_area
        else:
            area = None

        binflux = units.Photlam().Convert(self.binwave,
                                          self._binflux,
                                          self.fluxunits.name,
                                          area=area)
        return binflux

    def _getBinwaveProp(self):
        if self._binwave is None:
            self.initbinset(self.binset)
        return self._binwave

    binflux = property(_getBinfluxProp,doc='Flux of binned wavelength set.')
#    binwave = property(_getBinwaveProp,doc='Waveset for binned flux')

    # Multiplication is handled by performing the operation on
    # the spectral component of the Observation, and then creating a
    # new Observation as the result.
    #
    # This is because Observation is a subclass of CompositeSourceSpectrum
    # but with *a lot* of extra functionality involved in handling the
    # binned wave and flux arrays. Simply inheriting the parent class's
    # methods for multiplication does not return an Observation.
    #
    # Note that the order of operations actually implemented therefore varies
    # from what is expected, which naively would be
    #    (self.spectrum*self.bandpass) * other
    #
    def __mul__(self, other):
        # If the original object has partial overlap warnings, then
        # the forcing behavior also needs to be propagated.

        force = self.warnings.get('PartialOverlap', None)

        result = Observation(self.spectrum,
                             self.bandpass * other,
                             binset=self.binwave,
                             force=force)
        return result

    def __rmul__(self, other):
        return self.__mul__(other)

    #Disable methods that should not be supported by this class
    def __add__(self, other):
        raise NotImplementedError('Observations cannot be added')

    def __radd__(self, other):
        raise NotImplementedError('Observations cannot be added')

    def redshift(self,z):
        """Observations cannot be redshifted."""
        raise NotImplementedError('Observations cannot be redshifted')

    def writefits(self,fname,clobber=True, trimzero=True, binned=True,
                  hkeys=None):
        """Like :meth:`pysynphot.spectrum.SourceSpectrum.writefits`
        but with ``binned=True`` as default.

        """
        spectrum.CompositeSourceSpectrum.writefits(self,fname,
                                                   clobber=clobber,
                                                   trimzero=trimzero,
                                                   binned=binned,
                                                   hkeys=hkeys)

    def countrate(self,binned=True,range=None,force=False):
        """Calculate effective stimulus in count/s.
        Also see :ref:`pysynphot-formula-countrate` and
        :ref:`pysynphot-formula-effstim`.

        .. note::

            This is the calculation performed when the ETC invokes
            ``countrate``.

        Parameters
        -----------
        binned : bool
            If `True` (default), use binned data.
            Otherwise, use native data.

        range : tuple or `None`
            If not `None`, it must be a sequence with two floating-point
            elements specifying the wavelength range (*inclusive*) in the
            unit of ``self.waveunits`` in the form of ``(low, high)``;
            This is the range over which the integration will be performed.
            If the specified range does not exactly match a value in the
            wavelength set:

                * If ``binned=True``, the bin containing the range value will
                  be used. This assumes ``self.binwave`` contains bin centers.
                * If ``binned=False``, native dataset will be interpolated to
                  the specified values. (*Not Implemented.*)

        force : bool
            If `False` (default), partially overlapping ranges
            will raise an exception. If `True`, a partial overlap will
            return the calculated value instead. Disjoint ranges raise
            an exception regardless.

        Returns
        -------
        ans : float
            Count rate.

        Raises
        ------
        NotImplementedError
            Wavelength range is defined for unbinned data.

        pysynphot.exceptions.DisjointError
            Wavelength range does not overlap with observation.

        pysynphot.exceptions.PartialOverlap
            Wavelength range only partially overlaps with observation.

        """

        if self._binflux is None:
          self.initbinflux()

        myfluxunits = self.fluxunits.name
        self.convert('counts')
        warn=False
        if binned:
            #No range specified - use full range
            if range is None:
                lx,ux=(None,None)
            #Range is disjoint from binwave
            elif (range[0]>self.binwave[-1] or
                  range[1]<self.binwave[0]):
                raise exceptions.DisjointError("%s is disjoint from obs.binwave %s"%(range,
                                                                       [self.binwave[0],self.binwave[-1]]))
            #Partial overlap
            else:

                if range[0] < self._bin_edges[0]:
                    warn=True
                    lx=None
                else:
                    lx=np.searchsorted(self._bin_edges,range[0])-1

                if range[1] > self._bin_edges[-1]:
                    warn=True
                    ux=None
                else:
                    ux=np.searchsorted(self._bin_edges,range[1])

            ans = math.fsum(self.binflux[lx:ux])
            if warn and not force:
                raise exceptions.PartialOverlap("%s does not fully overlap binwave range %s. Countrate in overlap area is %f"%(range,[self.binwave[0],self.binwave[-1]],ans))

        else:
            if range is None:
                ans = math.fsum(self.flux)
            else:
                raise NotImplementedError("Sorry, range+binned=False not yet implemented")
        self.convert(myfluxunits)
        return ans

    def effstim(self,fluxunits='photlam'):
        """Compute :ref:`effective stimulus <pysynphot-formula-effstim>`.

        Calculations are done in given flux unit, and wavelengths
        in Angstrom. Native dataset is used.

        Parameters
        ----------
        fluxunits : str
            Flux unit.

        Returns
        -------
        ans : float
            Effective stimulus.

        Raises
        ------
        ValueError
            Invalid integrated flux.

        """
        oldunits=self.fluxunits
        self.convert(fluxunits)
        x=units.Units(fluxunits)
        try:
            if x.isDensity:
                rate=self.integrate()
                self._fluxcheck(rate)
                if x.isMag:
                    ans=x.unitResponse(self.bandpass) - 2.5*math.log10(rate)
                else:
                    ans=rate*x.unitResponse(self.bandpass)
            else:
                if x.isMag:
                    #its linear unit must be counts
                    self.convert('counts')
                    total=self.flux.sum()
                    self._fluxcheck(total)
                    ans=-2.5*math.log10(total)
                else:
                    ans=self.flux.sum()
                    self._fluxcheck(ans)
        finally:
            self.convert(oldunits)
            del x

        return ans

    def _fluxcheck(self,totalflux):
        if totalflux <= 0.0:
            raise ValueError('Integrated flux is <= 0')
        if np.isnan(totalflux):
            raise ValueError('Integrated flux is NaN')
        if np.isinf(totalflux):
            raise ValueError('Integrated flux is infinite')

    def pivot(self,binned=True):
        """Calculate :ref:`pivot wavelength <pysynphot-formula-pivwv>`
        of the observation.

        .. note::

            This is the calculation performed when ETC invokes ``calcphot``.

        Parameters
        ----------
        binned : bool
            Use binned dataset for calculations. Otherwise, use native dataset.

        Returns
        -------
        ans : float
            Pivot wavelength.

        """
        if binned:
            wave = self.binwave
        else:
            wave = self.wave

        countmulwave = self(wave)*wave
        countdivwave = self(wave)/wave

        num = self.trapezoidIntegration(wave,countmulwave)
        den = self.trapezoidIntegration(wave,countdivwave)

        if num == 0.0 or den == 0.0:
            return 0.0

        return math.sqrt(num/den)

    def efflam(self,binned=True):
        """Calculate :ref:`effective wavelength <pysynphot-formula-efflam>`
        of the observation.
        Calculation is done in the flux unit of ``flam``.

        .. note::

            Similar to IRAF STSDAS SYNPHOT ``efflphot`` task.

        Parameters
        ----------
        binned : bool
            Use binned dataset for calculations. Otherwise, use native dataset.

        Returns
        -------
        ans : float
            Effective wavelength.

        """
        myfluxunits=self.fluxunits.name
        self.convert('flam')
        if binned:
            wave=self.binwave
            flux=self.binflux
        else:
            wave=self.wave
            flux=self.flux

        num = self.trapezoidIntegration(wave,flux*wave*wave)
        den = self.trapezoidIntegration(wave,flux*wave)
        self.convert(myfluxunits)

        if num == 0.0 or den == 0.0:
            return 0.0

        return num/den

    def sample(self, swave, binned=True, fluxunits='counts'):
        """Sample the observation at the given wavelength.
        Also see :ref:`pysynphot-command-sample`.

        Parameters
        ----------
        swave : float
            Wavelength to sample.

        binned : bool
            Sample binned dataset (no interpolation).
            Otherwise, native (perform interpolation).

        fluxunits : {'counts'}
            Only the unit of counts is supported for now.

        Returns
        -------
        ans : float
            Sampled flux in given unit.

        Raises
        ------
        NotImplementedError
            Flux unit is not supported or non-scalar wavelength is given.

        ValueError
            Given wavelength out of range.

        """
        if self._binflux is None:
          self.initbinflux()

        if fluxunits != 'counts':
            s = "Sorry, only counts are supported at this time"
            raise NotImplementedError(s)
        else:
            #Save current fluxunits, in case they're different
            saveunits = None
            if not units.ismatch('counts', self.fluxunits):
                saveunits = self.fluxunits
                self.convert('counts')

        if binned:
            #Then we don't interpolate, just return the appropriate values
            #from binflux
            if np.isscalar(swave):
                #Find the bin in which it belongs.
                #_bin_edge[i] is the low edge of the bin centered
                #at binwave[i].

                idx = np.where(swave >= self._bin_edges)[0]
                #idx[-1] is the largest edge that is still smaller
                #than swave
                try:
                    ans = self.binflux[idx[-1]]
                except IndexError:
                    s = 'Value out of range: wavelength %g not contained in range [%g, %g]'
                    s = s % (swave, self.binwave[0], self.binwave[-1])
                    raise ValueError(s)

            else:
                #The logic for this case doesn't yet work on arrays
                s = "Sorry, only scalar samples are supported at this time"
                raise NotImplementedError(s)

        else:
            #Then we do interpolate on wave/flux
            if np.isscalar(swave):
                delta = 0.00001
                wv = np.array([swave - delta, swave, swave + delta])
                ans = np.interp(wv, self.wave, self.flux)[1]
            else:
                # This raises UnboundLocalError -- needs to be fixed!
                ans = np.interp(wv, self.wave, self.flux)

        #Change units back, if necessary, then return
        if saveunits is not None:
            self.convert(saveunits)
        return ans

    def pixel_range(self, waverange, waveunits=None, round='round'):
        """Calculate the number of wavelength bins within given
        wavelength range.

        .. note::

            This calls :func:`pysynphot.obsbandpass.pixel_range` with
            ``self.binwave`` as the first argument.

        Parameters
        ----------
        waverange, round
            See :func:`pysynphot.obsbandpass.pixel_range`.

        waveunits : str, optional
            The unit of the wavelength range.
            If `None` (default), the wavelengths are assumed to be
            in the units of ``self.waveunits``.

        Returns
        -------
        num : int or float
            Number of wavelength bins within ``waverange``.

        Raises
        ------
        pysynphot.exceptions.UndefinedBinset
            No binned dataset.

        """
        # make sure we have a binset to work with
        if self.binwave is None:
            raise exceptions.UndefinedBinset('No binset specified for this bandpass.')

        # start by converting waverange to self.waveunits, if necessary
        if waveunits is not None:
            waveunits = units.Units(waveunits)

            if not isinstance(waverange, np.ndarray):
                waverange = np.array(waverange)

            # convert to angstroms and then whatever self.waveunits is
            waverange = waveunits.ToAngstrom(waverange)

            waverange = units.Angstrom().Convert(waverange, self.waveunits.name)

        return pixel_range(self.binwave, waverange, round=round)

    def wave_range(self, cenwave, npix, waveunits=None, round='round'):
        """Calculate the wavelength range covered by the given
        number of pixels, centered on the given wavelength.

        .. note::

            This calls :func:`pysynphot.obsbandpass.wave_range` with
            ``self.binwave`` as the first argument.

        Parameters
        ----------
        cenwave, npix, round
            See :func:`pysynphot.obsbandpass.wave_range`.

        waveunits : str, optional
            Wavelength unit of the given and the returned wavelength values.
            If `None` (default), the wavelengths are assumed to be in
            the unit of ``self.waveunits``.

        Returns
        -------
        waverange : tuple of floats
            The range of wavelengths spanned by ``npix`` centered on
            ``cenwave``.

        Raises
        ------
        pysynphot.exceptions.UndefinedBinset
            No binned dataset.

        """
        # make sure we have a binset to work with
        if self.binwave is None:
            raise exceptions.UndefinedBinset('No binset specified for this bandpass.')

        # convert cenwave from waveunits to self.waveunits, if necessary
        if waveunits is not None:
            waveunits = units.Units(waveunits)

            # convert to angstroms and then whatever self.waveunits is
            cenwave = waveunits.ToAngstrom(cenwave)
            cenwave = units.Angstrom().Convert(cenwave, self.waveunits.name)

        wave1, wave2 = wave_range(self.binwave, cenwave, npix, round=round)

        # translate ends to waveunits, if necessary
        if waveunits is not None:
            # convert to angstroms
            wave1 = self.waveunits.ToAngstrom(wave1)
            wave2 = self.waveunits.ToAngstrom(wave2)

            # then to waveunits
            wave1 = units.Angstrom().Convert(wave1, waveunits.name)
            wave2 = units.Angstrom().Convert(wave2, waveunits.name)

        return wave1, wave2

    def as_spectrum(self, binned=True):
        """Reduce the observation to a simple spectrum object.

        An observation is a complex object with some restrictions on its
        capabilities. At times, it would be useful to work with the
        simulated observation as a simple object that is easier to
        manipulate and takes up less memory.

        Parameters
        ----------
        binned : bool
            If `True` (default), export binned dataset. Otherwise, native.

        Returns
        -------
        result : `~pysynphot.spectrum.ArraySourceSpectrum`
            Observation dataset as a simple spectrum object.

        """
        if binned:
            wave, flux = self.binwave, self.binflux
        else:
            wave, flux = self.wave, self.flux

        result = ArraySourceSpectrum(wave, flux,
                                     self.waveunits,
                                     self.fluxunits,
                                     name = self.name,
                                     keepneg = True)

        return result


#   Copyright 2020 The PyMC Developers
#
#   Licensed under the Apache License, Version 2.0 (the "License");
#   you may not use this file except in compliance with the License.
#   You may obtain a copy of the License at
#
#       http://www.apache.org/licenses/LICENSE-2.0
#
#   Unless required by applicable law or agreed to in writing, software
#   distributed under the License is distributed on an "AS IS" BASIS,
#   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#   See the License for the specific language governing permissions and
#   limitations under the License.

import logging

from copy import copy
from typing import Any, Dict, List, Tuple

import aesara
import numpy as np

from aesara import function as aesara_function

from pymc.aesaraf import inputvars, join_nonshared_inputs, make_shared_replacements
from pymc.bart.bart import BARTRV
from pymc.bart.tree import LeafNode, SplitNode, Tree
from pymc.blocking import RaveledVars
from pymc.model import modelcontext
from pymc.step_methods.arraystep import ArrayStepShared, Competence

_log = logging.getLogger("pymc")

class ParticleTree:
    """
    Particle tree
    """

    def __init__(self, tree, log_weight, likelihood):
        self.tree = tree.copy()  # keeps the tree that we care at the moment
        self.expansion_nodes = [0]
        self.log_weight = log_weight
        self.old_likelihood_logp = likelihood
        self.used_variates = []

    def sample_tree_sequential(
        self,
        ssv,
        available_predictors,
        prior_prob_leaf_node,
        X,
        missing_data,
        sum_trees_output,
        mean,
        linear_fit,
        m,
        normal,
        mu_std,
        response,
    ):
        tree_grew = False
        if self.expansion_nodes:
            index_leaf_node = self.expansion_nodes.pop(0)
            # Probability that this node will remain a leaf node
            prob_leaf = prior_prob_leaf_node[self.tree[index_leaf_node].depth]

            if prob_leaf < np.random.random():
                tree_grew, index_selected_predictor = grow_tree(
                    self.tree,
                    index_leaf_node,
                    ssv,
                    available_predictors,
                    X,
                    missing_data,
                    sum_trees_output,
                    mean,
                    linear_fit,
                    m,
                    normal,
                    mu_std,
                    response,
                )
                if tree_grew:
                    new_indexes = self.tree.idx_leaf_nodes[-2:]
                    self.expansion_nodes.extend(new_indexes)
                    self.used_variates.append(index_selected_predictor)

        return tree_grew

class PGBART(ArrayStepShared):
    """
    Particle Gibss BART sampling step

    Parameters
    ----------
    vars: list
        List of value variables for sampler
    num_particles : int
        Number of particles for the conditional SMC sampler. Defaults to 10
    max_stages : int
        Maximum number of iterations of the conditional SMC sampler. Defaults to 100.
    batch : int
        Number of trees fitted per step. Defaults to  "auto", which is the 10% of the `m` trees
        during tuning and 20% after tuning.
    model: PyMC Model
        Optional model for sampling step. Defaults to None (taken from context).

    Note
    ----
    This sampler is inspired by the [Lakshminarayanan2015] Particle Gibbs sampler, but introduces
    several changes. The changes will be properly documented soon.

    References
    ----------
    .. [Lakshminarayanan2015] Lakshminarayanan, B. and Roy, D.M. and Teh, Y. W., (2015),
        Particle Gibbs for Bayesian Additive Regression Trees.
        ArviX, `link <https://arxiv.org/abs/1502.04622>`__
    """

    name = "bartsampler"
    default_blocked = False
    generates_stats = True
    stats_dtypes = [{"variable_inclusion": np.ndarray, "bart_trees": np.ndarray}]

    def __init__(self, vars=None, num_particles=10, max_stages=100, batch="auto", model=None):
        _log.warning("BART is experimental. Use with caution.")
        model = modelcontext(model)
        initial_values = model.recompute_initial_point()
        value_bart = inputvars(vars)[0]
        self.bart = model.values_to_rvs[value_bart].owner.op

        self.X = self.bart.X
        self.Y = self.bart.Y
        self.missing_data = np.any(np.isnan(self.X))
        self.m = self.bart.m
        self.alpha = self.bart.alpha
        self.k = self.bart.k
        self.response = self.bart.response
        self.split_prior = self.bart.split_prior
        if self.split_prior is None:
            self.split_prior = np.ones(self.X.shape[1])

        self.init_mean = self.Y.mean()
        # if data is binary
        Y_unique = np.unique(self.Y)
        if Y_unique.size == 2 and np.all(Y_unique == [0, 1]):
            self.mu_std = 6 / (self.k * self.m ** 0.5)
        # maybe we need to check for count data
        else:
            self.mu_std = self.Y.std() / (self.k * self.m ** 0.5)

        self.num_observations = self.X.shape[0]
        self.num_variates = self.X.shape[1]
        self.available_predictors = list(range(self.num_variates))

        sum_trees_output = np.full_like(self.Y, self.init_mean).astype(aesara.config.floatX)
        self.a_tree = Tree.init_tree(
            tree_id=0,
            leaf_node_value=self.init_mean / self.m,
            idx_data_points=np.arange(self.num_observations, dtype="int32"),
            m=self.m,
        )
        self.mean = fast_mean()
        self.linear_fit = fast_linear_fit()

        self.normal = NormalSampler()
        self.prior_prob_leaf_node = compute_prior_probability(self.alpha)
        self.ssv = SampleSplittingVariable(self.split_prior)

        self.tune = True
        self.idx = 0
        self.batch = batch

        if self.batch == "auto":
            self.batch = max(1, int(self.m * 0.1))
        self.log_num_particles = np.log(num_particles)
        self.indices = list(range(1, num_particles))
        self.len_indices = len(self.indices)
        self.max_stages = max_stages

        shared = make_shared_replacements(initial_values, vars, model)
        self.likelihood_logp = logp(initial_values, [model.datalogpt], vars, shared)
        self.init_likelihood = self.likelihood_logp(sum_trees_output)
        self.init_log_weight = self.init_likelihood - self.log_num_particles
        self.all_particles = []
        for i in range(self.m):
            self.a_tree.tree_id = i
            p = ParticleTree(
                self.a_tree,
                self.init_log_weight,
                self.init_likelihood,
            )
            self.all_particles.append(p)
        self.all_trees = np.array([p.tree for p in self.all_particles])
        super().__init__(vars, shared)

    def astep(self, q: RaveledVars) -> Tuple[RaveledVars, List[Dict[str, Any]]]:
        point_map_info = q.point_map_info
        sum_trees_output = q.data
        variable_inclusion = np.zeros(self.num_variates, dtype="int")

        if self.idx == self.m:
            self.idx = 0

        for tree_id in range(self.idx, self.idx + self.batch):
            if tree_id >= self.m:
                break
            # Generate an initial set of SMC particles
            # at the end of the algorithm we return one of these particles as the new tree
            particles = self.init_particles(tree_id)
            # Compute the sum of trees without the tree we are attempting to replace
            self.sum_trees_output_noi = sum_trees_output - particles[0].tree.predict_output()

            # The old tree is not growing so we update the weights only once.
            self.update_weight(particles[0])
            for t in range(self.max_stages):
                # Sample each particle (try to grow each tree), except for the first one.
                for p in particles[1:]:
                    tree_grew = p.sample_tree_sequential(
                        self.ssv,
                        self.available_predictors,
                        self.prior_prob_leaf_node,
                        self.X,
                        self.missing_data,
                        sum_trees_output,
                        self.mean,
                        self.linear_fit,
                        self.m,
                        self.normal,
                        self.mu_std,
                        self.response,
                    )
                    if tree_grew:
                        self.update_weight(p)
                # Normalize weights
                W_t, normalized_weights = self.normalize(particles)

                # Resample all but first particle
                re_n_w = normalized_weights[1:] / normalized_weights[1:].sum()
                new_indices = np.random.choice(self.indices, size=self.len_indices, p=re_n_w)
                particles[1:] = particles[new_indices]

                # Set the new weights
                for p in particles:
                    p.log_weight = W_t

                # Check if particles can keep growing, otherwise stop iterating
                non_available_nodes_for_expansion = []
                for p in particles[1:]:
                    if p.expansion_nodes:
                        non_available_nodes_for_expansion.append(0)
                if all(non_available_nodes_for_expansion):
                    break

            # Get the new tree and update
            new_particle = np.random.choice(particles, p=normalized_weights)
            new_tree = new_particle.tree
            self.all_trees[self.idx] = new_tree
            new_particle.log_weight = new_particle.old_likelihood_logp - self.log_num_particles
            self.all_particles[tree_id] = new_particle
            sum_trees_output = self.sum_trees_output_noi + new_tree.predict_output()

            if self.tune:
                for index in new_particle.used_variates:
                    self.split_prior[index] += 1
                    self.ssv = SampleSplittingVariable(self.split_prior)
            else:
                self.batch = max(1, int(self.m * 0.2))
                for index in new_particle.used_variates:
                    variable_inclusion[index] += 1
            self.idx += 1

        stats = {"variable_inclusion": variable_inclusion, "bart_trees": copy(self.all_trees)}
        sum_trees_output = RaveledVars(sum_trees_output, point_map_info)
        return sum_trees_output, [stats]

    @staticmethod
    def competence(var, has_grad):
        """
        PGBART is only suitable for BART distributions
        """
        dist = getattr(var.owner, "op", None)
        if isinstance(dist, BARTRV):
            return Competence.IDEAL
        return Competence.INCOMPATIBLE

    def normalize(self, particles: List[ParticleTree]) -> Tuple[float, np.ndarray]:
        """
        Use logsumexp trick to get W_t and softmax to get normalized_weights
        """
        log_w = np.array([p.log_weight for p in particles])
        log_w_max = log_w.max()
        log_w_ = log_w - log_w_max
        w_ = np.exp(log_w_)
        w_sum = w_.sum()
        W_t = log_w_max + np.log(w_sum) - self.log_num_particles
        normalized_weights = w_ / w_sum
        # stabilize weights to avoid assigning exactly zero probability to a particle
        normalized_weights += 1e-12

        return W_t, normalized_weights

    def init_particles(self, tree_id: int) -> np.ndarray:
        """
        Initialize particles
        """
        p = self.all_particles[tree_id]
        p.log_weight = self.init_log_weight
        p.old_likelihood_logp = self.init_likelihood
        particles = [p]

        for _ in self.indices:
            self.a_tree.tree_id = tree_id
            particles.append(
                ParticleTree(
                    self.a_tree,
                    self.init_log_weight,
                    self.init_likelihood,
                )
            )

        return np.array(particles)

    def update_weight(self, particle: List[ParticleTree]) -> None:
        """
        Update the weight of a particle

        Since the prior is used as the proposal,the weights are updated additively as the ratio of
        the new and old log-likelihoods.
        """
        new_likelihood = self.likelihood_logp(
            self.sum_trees_output_noi + particle.tree.predict_output()
        )
        particle.log_weight += new_likelihood - particle.old_likelihood_logp
        particle.old_likelihood_logp = new_likelihood

class SampleSplittingVariable:
    def __init__(self, alpha_prior):
        """
        Sample splitting variables proportional to `alpha_prior`.

        This is equivalent as sampling weights from a Dirichlet distribution with `alpha_prior`
        parameter and then using those weights to sample from the available spliting variables.
        This enforce sparsity.
        """
        self.enu = list(enumerate(np.cumsum(alpha_prior / alpha_prior.sum())))

    def rvs(self):
        r = np.random.random()
        for i, v in self.enu:
            if r <= v:
                return i

def compute_prior_probability(alpha):
    """
    Calculate the probability of the node being a LeafNode (1 - p(being SplitNode)).
    Taken from equation 19 in [Rockova2018].

    Parameters
    ----------
    alpha : float

    Returns
    -------
    list with probabilities for leaf nodes

    References
    ----------
    .. [Rockova2018] Veronika Rockova, Enakshi Saha (2018). On the theory of BART.
    arXiv, `link <https://arxiv.org/abs/1810.00787>`__
    """
    prior_leaf_prob = [0]
    depth = 1
    while prior_leaf_prob[-1] < 1:
        prior_leaf_prob.append(1 - alpha ** depth)
        depth += 1
    return prior_leaf_prob

def grow_tree(
    tree,
    index_leaf_node,
    ssv,
    available_predictors,
    X,
    missing_data,
    sum_trees_output,
    mean,
    linear_fit,
    m,
    normal,
    mu_std,
    response,
):
    current_node = tree.get_node(index_leaf_node)
    idx_data_points = current_node.idx_data_points

    index_selected_predictor = ssv.rvs()
    selected_predictor = available_predictors[index_selected_predictor]
    available_splitting_values = X[idx_data_points, selected_predictor]
    if missing_data:
        idx_data_points = idx_data_points[~np.isnan(available_splitting_values)]
        available_splitting_values = available_splitting_values[
            ~np.isnan(available_splitting_values)
        ]

    if available_splitting_values.size == 0:
        return False, None

    idx_selected_splitting_values = discrete_uniform_sampler(len(available_splitting_values))
    split_value = available_splitting_values[idx_selected_splitting_values]
    new_split_node = SplitNode(
        index=index_leaf_node,
        idx_split_variable=selected_predictor,
        split_value=split_value,
    )

    left_node_idx_data_points, right_node_idx_data_points = get_new_idx_data_points(
        split_value, idx_data_points, selected_predictor, X
    )

    if response == "mix":
        response = "linear" if np.random.random() >= 0.5 else "constant"

    left_node_value, left_node_linear_params = draw_leaf_value(
        sum_trees_output[left_node_idx_data_points],
        X[left_node_idx_data_points, selected_predictor],
        mean,
        linear_fit,
        m,
        normal,
        mu_std,
        response,
    )
    right_node_value, right_node_linear_params = draw_leaf_value(
        sum_trees_output[right_node_idx_data_points],
        X[right_node_idx_data_points, selected_predictor],
        mean,
        linear_fit,
        m,
        normal,
        mu_std,
        response,
    )

    new_left_node = LeafNode(
        index=current_node.get_idx_left_child(),
        value=left_node_value,
        idx_data_points=left_node_idx_data_points,
        linear_params=left_node_linear_params,
    )
    new_right_node = LeafNode(
        index=current_node.get_idx_right_child(),
        value=right_node_value,
        idx_data_points=right_node_idx_data_points,
        linear_params=right_node_linear_params,
    )
    tree.grow_tree(index_leaf_node, new_split_node, new_left_node, new_right_node)

    return True, index_selected_predictor

def get_new_idx_data_points(split_value, idx_data_points, selected_predictor, X):

    left_idx = X[idx_data_points, selected_predictor] <= split_value
    left_node_idx_data_points = idx_data_points[left_idx]
    right_node_idx_data_points = idx_data_points[~left_idx]

    return left_node_idx_data_points, right_node_idx_data_points

def draw_leaf_value(Y_mu_pred, X_mu, mean, linear_fit, m, normal, mu_std, response):
    """Draw Gaussian distributed leaf values"""
    linear_params = None
    if Y_mu_pred.size == 0:
        return 0, linear_params
    else:
        norm = normal.random() * mu_std
        if Y_mu_pred.size == 1:
            mu_mean = Y_mu_pred.item() / m
        elif response == "constant":
            mu_mean = mean(Y_mu_pred) / m
        elif response == "linear":
            Y_fit, linear_params = linear_fit(X_mu, Y_mu_pred)
            mu_mean = Y_fit / m
            linear_params[2] = norm

        draw = norm + mu_mean
        return draw, linear_params

def fast_mean():
    """If available use Numba to speed up the computation of the mean."""
    try:
        from numba import jit
    except ImportError:
        return np.mean

    @jit
    def mean(a):
        count = a.shape[0]
        suma = 0
        for i in range(count):
            suma += a[i]
        return suma / count

    return mean

def fast_linear_fit():
    """If available use Numba to speed up the computation of the linear fit"""

    def linear_fit(X, Y):

        n = len(Y)
        xbar = np.sum(X) / n
        ybar = np.sum(Y) / n

        den = X @ X - n * xbar ** 2
        if den > 1e-10:
            b = (X @ Y - n * xbar * ybar) / den
        else:
            b = 0
        a = ybar - b * xbar
        Y_fit = a + b * X
        return Y_fit, [a, b, 0]

    try:
        from numba import jit

        return jit(linear_fit)
    except ImportError:
        return linear_fit

def discrete_uniform_sampler(upper_value):
    """Draw from the uniform distribution with bounds [0, upper_value).

    This is the same and np.random.randit(upper_value) but faster.
    """
    return int(np.random.random() * upper_value)

class NormalSampler:
    """
    Cache samples from a standard normal distribution
    """

    def __init__(self):
        self.size = 1000
        self.cache = []

    def random(self):
        if not self.cache:
            self.update()
        return self.cache.pop()

    def update(self):
        self.cache = np.random.normal(loc=0.0, scale=1, size=self.size).tolist()

def logp(point, out_vars, vars, shared):
    """Compile Aesara function of the model and the input and output variables.

    Parameters
    ----------
    out_vars: List
        containing :class:`pymc.Distribution` for the output variables
    vars: List
        containing :class:`pymc.Distribution` for the input variables
    shared: List
        containing :class:`aesara.tensor.Tensor` for depended shared data
    """
    out_list, inarray0 = join_nonshared_inputs(point, out_vars, vars, shared)
    f = aesara_function([inarray0], out_list[0])
    f.trust_input = True
    return f


# Copyright 2018 The Fragpy Developers. All Rights Reserved.
#
# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License.  You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied.  See the License for the
# specific language governing permissions and limitations
# under the License.
"""\
 Bond lengths;  angles; torsionals
"""

from fragpy.odict import *
from math import *
from numpy import subtract
from numpy.linalg import norm
import numpy as np

checker1=[]
checker2=[]
checker3=0

def hbon(list1,list2):

    tmpii=[]
    electNEG = ['O','N','F','Cl','S']
    for i in list1:
        closeN = closest(i,list2,checker1,checker2,checker3)
        for j in list2:
            if not closeN == j:
                if j.atm in electNEG:
                    bond1 = GetBond(i.q,j.q)
                    bondc = covrad[j.atm] + covrad[i.atm]
                    bondv = vanrad[j.atm] + vanrad[i.atm]
                    if (bond1 > bondc and bond1 < 0.9000*(bondv)):
                        angle1 = degrees(GetAngle(closeN.q,i.q,j.q))
                        if angle1 > 90:
                            tmpii.append([i.n,j.n])

    lentmpii = len(tmpii)

    for i in list2:
        hclose1=closest(i,list1,checker1,checker2,checker3)
        if not hclose1==[]:
            tmpii.append([hclose1.n,i.n])
            hclose2=closest(i,list1,hclose1,checker2,checker3)
            if not hclose2==[]:
                tmpii.append([hclose2.n,i.n])
                hclose3=closest(i,list1,hclose1,hclose2,hclose2)
                if not hclose3==[]:
                    tmpii.append([hclose3.n,i.n])

    return(tmpii,lentmpii)

def hangle(list1,list2,lenHbon):  # Hbondlist, nonHbondlist
    hangl=[]
    for i in list1[lenHbon:]:
        tmpii=[]
        for j in list2:
            if i[1]==j[0]:
                tmpii.append(j[1])
            elif i[1]==j[1]:
                tmpii.append(j[0])
                #tmpii=j[1]
        for j in tmpii:
            if not [j,i[1],i[0]] in hangl:
                hangl.append([i[0],i[1],j])
        tmpii = []
        for j in list1[lenHbon:]:
            if (i[1] == j[1] and not i[0] == j[0]):
                tmpii.append(j[0])
        for j in tmpii:
            if not [j,i[1],i[0]] in hangl:
                hangl.append([i[0],i[1],j])
    return hangl

def hdih(list1,list2):
    hdi=[]
    for i in list1:
        tmpii = []
        for j in list2:
            if (i[1] == j[0] and i[2] == j[1]):
                tmpii.append([i[0],i[1],i[2],j[2]])
            elif (i[1] == j[2] and i[2] == j[1]):
                tmpii.append([i[0],i[1],i[2],j[0]])
        for j in tmpii:
            if not [j[3],j[2],j[1],j[0]] in hdi:
                hdi.append(j)
        tmpii = []
        for j in list1:
            if (j[2] == i[1] and j[1] == i[2]):
                tmpii.append([i[0],i[1],i[2],j[0]])
        for j in tmpii:
            if not [j[3],j[2],j[1],j[0]] in hdi:
                hdi.append(j)
    return hdi

def angl(list1):
    condlist=[]
    tmpang=[]
    ange=[]
    coun=0
    for i in list1:
        tmpii = []
        for j in list1:
            if (i[0] == j[0] and not i[1] == j[1]):
                tmpii.append([i[1],i[0],j[1]])
            elif i[0] == j[1]:
                tmpii.append([i[1],i[0],j[0]])
            elif i[1] == j[0]:
                tmpii.append([i[0],i[1],j[1]])
            elif (i[1] == j[1] and not i[0] == j[0]):
                tmpii.append([i[0],i[1],j[0]])
        for j in tmpii:
            if not sorted(j) in tmpang:
                ange.append(j)
                tmpang.append(sorted(j))
    return ange

def dihed(list1):
    condlist=[]
    dihlist=[]
    for i in list1:
        tmpii = []
        for j in list1:
            if (i[1]==j[0] and i[2]==j[1]):
                tmpii.append([i[0],i[1],i[2],j[2]])
            elif (i[1]==j[2] and i[2]==j[1]):
                tmpii.append([i[0],i[1],i[2],j[0]])
            elif (i[1]==j[2] and i[0]==j[1]):
                tmpii.append([i[2],i[1],i[0],j[0]])
            elif (i[1]==j[0] and i[0]==j[1]):
                tmpii.append([i[2],i[1],i[0],j[2]])
        for j in tmpii:
            if not sorted(j) in condlist:
                dihlist.append(j)
                condlist.append(sorted(j))
    return dihlist

def closest(x,list,check1,check3,check2):
    condition=0
    for i in list:
        if not i==check1:
            if not i==check2:
                if not i == check3:
                    d1 = GetBond(x.q,i.q)
                    d2 = covrad[x.atm]+covrad[i.atm]
                    if d1 <= d2+0.1:
                        close=i
                        condition=1
                        break
                    else:
                        close=[]
                        condition=2
    if condition==0:
        close=[]
    return (close)

def GetBond(list1,list2):
    bond1 = subtract(list1,list2).tolist()
    bond2 = norm(bond1)
    return bond2
def GetAngle(s1,s2,s3):
    import numpy

    u = subtract(s1,s2)
    v = subtract(s3,s2)
    ru = norm(u)
    rv = norm(v)
    u = (u/ru).tolist()
    v = (v/rv).tolist()
    angle1 = numpy.arccos(np.dot(u,v))
    return angle1

def GetDihedral(p0,p1,p2,p3):
    b0 = -1.0 * subtract(p1,p0)
    b1 = subtract(p2,p1)
    b2 = subtract(p3,p2)
    b1 /= norm(b1)
    v = subtract(b0,np.dot(b0,b1)*b1)
    w = subtract(b2,np.dot(b2,b1)*b1)
    x = np.dot(v,w)
    y = np.dot(np.cross(b1,v),w)
    return np.arctan2(y,x)