#  Copyright (c) 2021 zfit
import contextlib
import itertools
import warnings
from collections import OrderedDict
from typing import Dict, Union, Callable, Optional, Tuple, Iterable

import colored
import iminuit
import numpy as np
import scipy.stats
from colorama import Style, init
from ordered_set import OrderedSet
from tabulate import tabulate

from .errors import compute_errors, covariance_with_weights, dict_to_matrix, matrix_to_dict
from .interface import ZfitMinimizer, ZfitResult
from ..core.interfaces import ZfitLoss, ZfitParameter
from ..core.parameter import set_values
from ..settings import run
from ..util.container import convert_to_container
from ..util.deprecation import deprecated_args
from ..util.warnings import ExperimentalFeatureWarning
from ..util.ztyping import ParamsTypeOpt

init(autoreset=True)

def _minos_minuit(result, params, cl=None, *, minuit_minimizer=None):
    from zfit.minimize import Minuit

    fitresult = result
    minimizer = fitresult.minimizer
    if not isinstance(minimizer, Minuit):
        if minuit_minimizer is None:
            raise TypeError("To use `minos_error` method with a minimizer other than `Minuit`, a Minuit minimizer"
                            " needs to be passed to the method (it doesn't has to be minimized though).")
        minimizer = minuit_minimizer

    merror_result = minimizer._minuit_minimizer.minos(*(p.name for p in params), cl=cl).merrors  # returns every var
    attrs = ['lower', 'upper', 'is_valid', 'upper_valid', 'lower_valid', 'at_lower_limit', 'at_upper_limit', 'nfcn']
    result = {}
    for p, merror in zip(params, merror_result.values()):
        result[p] = {attr: getattr(merror, attr) for attr in attrs}
        result[p]['original'] = result
    # result = OrderedDict((p, result[p.name]) for p in params)
    new_result = None
    return result, new_result

def _covariance_minuit(result, params):
    if any(data.weights is not None for data in result.loss.data):
        warnings.warn("The computation of the covariance matrix with weights is still experimental.",
                      ExperimentalFeatureWarning)

    fitresult = result
    minimizer = fitresult.minimizer

    from zfit.minimizers.minimizer_minuit import Minuit

    if not isinstance(minimizer, Minuit):
        raise TypeError("Cannot compute the covariance matrix with 'covariance_minuit' with a different"
                        " minimizer than `Minuit`.")

    covariance = result.minimizer._minuit_minimizer.covariance

    covariance_dict = {}
    for p1 in params:
        for p2 in params:
            key = (p1, p2)
            covariance_dict[key] = covariance[tuple(k.name for k in key)]

    return covariance_dict

def _covariance_np(result, params):
    if any(data.weights is not None for data in result.loss.data):
        warnings.warn("The computation of the covariance matrix with weights is still experimental.",
                      ExperimentalFeatureWarning)

    # TODO: maybe activate again? currently fails due to numerical problems
    # numgrad_was_none = settings.options.numerical_grad is None
    # if numgrad_was_none:
    #     settings.options.numerical_grad = True

    _, gradient, hessian = result.loss.value_gradients_hessian(params)
    covariance = np.linalg.inv(hessian)

    # if numgrad_was_none:
    #     settings.options.numerical_grad = None

    return matrix_to_dict(params, covariance)

class FitResult(ZfitResult):
    _default_hesse = "hesse_np"
    _hesse_methods = {"minuit_hesse": _covariance_minuit, "hesse_np": _covariance_np}
    _default_error = "zfit_error"
    _error_methods = {"minuit_minos": _minos_minuit, "zfit_error": compute_errors}

    def __init__(self, params: Dict[ZfitParameter, float], edm: float, fmin: float, status: int, converged: bool,
                 info: dict, loss: ZfitLoss, minimizer: "ZfitMinimizer"):
        """Create a `FitResult` from a minimization. Store parameter values, minimization infos and calculate errors.

        Any errors calculated are saved under `self.params` dictionary with::

            {parameter: {error_name1: {'low': value, 'high': value or similar}}

        Args:
            params: Result of the fit where each
               :py:class:`~zfit.Parameter` key has the value from the minimum found by the minimizer.
            edm: The estimated distance to minimum, estimated by the minimizer (if available)
            fmin: The minimum of the function found by the minimizer
            status: A status code (if available)
            converged: Whether the fit has successfully converged or not.
            info: Additional information (if available) like *number of function calls* and the
                original minimizer return message.
            loss: The loss function that was minimized. Contains also the pdf, data etc.
            minimizer: Minimizer that was used to obtain this `FitResult` and will be used to
                calculate certain errors. If the minimizer is state-based (like "iminuit"), then this is a copy
                and the state of other `FitResults` or of the *actual* minimizer that performed the minimization
                won't be altered.
        """
        super().__init__()

        self._status = status
        self._converged = converged
        self._params = self._input_convert_params(params)
        self._params_at_limit = any(param.at_limit for param in self.params)
        self._edm = edm
        self._fmin = fmin
        self._info = info
        self._loss = loss
        self._minimizer = minimizer
        self._valid = True
        self._covariance_dict = {}

    def _input_convert_params(self, params):
        return ParamHolder((p, {"value": v}) for p, v in params.items())

    def _get_uncached_params(self, params, method_name):
        return [p for p in params if self.params[p].get(method_name) is None]

    @classmethod
    def from_minuit(cls, loss: ZfitLoss, params: Iterable[ZfitParameter], result: iminuit.util.FMin,
                    minimizer: Union[ZfitMinimizer, iminuit.Minuit]) -> 'FitResult':
        """Create a `FitResult` from a :py:class:~`iminuit.util.MigradResult` returned by
        :py:meth:`iminuit.Minuit.migrad` and a iminuit :py:class:~`iminuit.Minuit` instance with the corresponding
        zfit objects.

        Args:
            loss: zfit Loss that was minimized.
            params: Iterable of the zfit parameters that were floating during the minimization.
            result: Return value of the iminuit migrad command.
            minimizer: Instance of the iminuit Minuit that was used to minimize the loss.

        Returns:
            A `FitResult` as if zfit Minuit was used.
        """

        from .minimizer_minuit import Minuit
        if not isinstance(minimizer, Minuit):
            if isinstance(minimizer, iminuit.Minuit):
                minimizer_new = Minuit()
                minimizer_new._minuit_minimizer = minimizer
                minimizer = minimizer_new
            else:
                raise ValueError(f"Minimizer {minimizer} not supported. Use `Minuit` from zfit or from iminuit.")
        params_result = [p_dict for p_dict in result.params]
        result_vals = [res.value for res in params_result]
        set_values(params, values=result_vals)
        fmin_object = result.fmin
        info = {'n_eval': fmin_object.nfcn,
                'n_iter': "REMOVED FROM MINUIT",  # TODO 6.x: remove completely
                # 'grad': result['jac'],
                # 'message': result['message'],
                'original': fmin_object}
        edm = fmin_object.edm
        fmin = fmin_object.fval
        status = -999
        converged = fmin_object.is_valid
        params = OrderedDict((p, res.value) for p, res in zip(params, params_result))
        return cls(params=params, edm=edm, fmin=fmin, info=info, loss=loss,
                   status=status, converged=converged,
                   minimizer=minimizer)

    @property
    def params(self):
        return self._params

    @property
    def edm(self):
        """The estimated distance to the minimum.

        Returns:
            Numeric
        """
        return self._edm

    @property
    def minimizer(self):
        return self._minimizer

    @property
    def loss(self) -> ZfitLoss:
        # TODO(Mayou36): this is currently a reference, should be a copy of the loss?
        return self._loss

    @property
    def fmin(self):
        """Function value at the minimum.

        Returns:
            Numeric
        """
        return self._fmin

    @property
    def status(self):
        return self._status

    @property
    def info(self):
        return self._info

    @property
    def converged(self):
        return self._converged

    @property
    def valid(self):
        return self._valid and not self.params_at_limit and self.converged

    @property
    def params_at_limit(self):
        return self._params_at_limit

    @contextlib.contextmanager
    def _input_check_reset_params(self, params):
        params = self._input_check_params(params=params)
        old_values = run(params)
        try:
            yield params
        except Exception:
            warnings.warn("Exception occurred, parameter values are not reset and in an arbitrary, last"
                          " used state. If this happens during normal operation, make sure you reset the values.",
                          RuntimeWarning)
            raise
        set_values(params=params, values=old_values)

    def _input_check_params(self, params):
        if params is not None:
            params = convert_to_container(params)
        else:
            params = list(self.params.keys())
        return params

    def hesse(self, params: ParamsTypeOpt = None, method: Union[str, Callable] = None,
              error_name: Optional[str] = None) -> OrderedDict:
        """Calculate for `params` the symmetric error using the Hessian/covariance matrix.

        Args:
            params: The parameters to calculate the
                Hessian symmetric error. If None, use all parameters.
            method: the method to calculate the covariance matrix. Can be {'minuit_hesse', 'hesse_np'} or a callable.
            error_name: The name for the error in the dictionary.

        Returns:
            Result of the hessian (symmetric) error as dict with each parameter holding
                the error dict {'error': sym_error}.

                So given param_a (from zfit.Parameter(.))
                `error_a = result.hesse(params=param_a)[param_a]['error']`
                error_a is the hessian error.
        """
        if method is None:
            # LEGACY START
            method = self._default_hesse
            from zfit.minimizers.minimizer_minuit import Minuit

            if isinstance(self.minimizer, Minuit):
                method = "minuit_hesse"
            # LEGACY END
        if error_name is None:
            if not isinstance(method, str):
                raise ValueError("Need to specify `error_name` or use a string as `method`")
            error_name = method

        with self._input_check_reset_params(params) as params:
            uncached_params = self._get_uncached_params(params=params, method_name=error_name)
            if uncached_params:
                error_dict = self._hesse(params=uncached_params, method=method)
                self._cache_errors(error_name=error_name, errors=error_dict)

        errors = OrderedDict((p, self.params[p][error_name]) for p in params)
        return errors

    def _cache_errors(self, error_name, errors):
        for param, error in errors.items():
            self.params[param][error_name] = error

    def _hesse(self, params, method):
        covariance_dict = self.covariance(params, method, as_dict=True)
        return OrderedDict((p, {"error": covariance_dict[(p, p)] ** 0.5}) for p in params)

    def error(self, params: ParamsTypeOpt = None, method: Union[str, Callable] = None, error_name: str = None,
              sigma: float = 1.0) -> OrderedDict:
        r"""

        .. deprecated:: unknown
            Use :func:`errors` instead.

        Args:
            params: The parameters or their names to calculate the
                 errors. If `params` is `None`, use all *floating* parameters.
            method: The method to use to calculate the errors. Valid choices are
                {'minuit_minos'} or a Callable.
            sigma: Errors are calculated with respect to `sigma` std deviations. The definition
                of 1 sigma depends on the loss function and is defined there.

                For example, the negative log-likelihood (without the factor of 2) has a correspondents
                of :math:`\Delta` NLL of 1 corresponds to 1 std deviation.
            error_name: The name for the error in the dictionary.

        Returns:
            A `OrderedDict` containing as keys the parameter names and as value a `dict` which
                contains (next to probably more things) two keys 'lower' and 'upper',
                holding the calculated errors.
                Example: result['par1']['upper'] -> the asymmetric upper error of 'par1'
        """
        warnings.warn("`error` is deprecated, use `errors` instead. This will return not only the errors but also "
                      "(a possible) new FitResult if a minimum was found. So change"
                      "errors = result.error()"
                      "to"
                      "errors, new_res = result.errors()", DeprecationWarning)
        return self.errors(params=params, method=method, error_name=error_name, sigma=sigma)[0]

    @deprecated_args(None, "Use cl for confidence level instead.", 'sigma')
    def errors(self, params: ParamsTypeOpt = None, method: Union[str, Callable] = None, error_name: str = None,
               cl: Optional[float] = None, sigma=None) -> Tuple[OrderedDict, Union[None, 'FitResult']]:
        r"""Calculate and set for `params` the asymmetric error using the set error method.

            Args:
                params: The parameters or their names to calculate the
                     errors. If `params` is `None`, use all *floating* parameters.
                method: The method to use to calculate the errors. Valid choices are
                    {'minuit_minos'} or a Callable.
                cl: Uncertainties are calculated with respect to the confidence level cl. The default is 68.3%.
                    For example, the negative log-likelihood (without the factor of 2) has a correspondents
                    of :math:`\Delta` NLL of 1 corresponds to 1 std deviation.
                sigma: Errors are calculated with respect to `sigma` std deviations. The definition
                    of 1 sigma depends on the loss function and is defined there.
                error_name: The name for the error in the dictionary.

            Returns:
                A `OrderedDict` containing as keys the parameter and as value a `dict` which
                    contains (next to probably more things) two keys 'lower' and 'upper',
                    holding the calculated errors.
                    Example: result[par1]['upper'] -> the asymmetric upper error of 'par1'
        """
        if sigma is not None:
            if cl is not None:
                raise ValueError("Cannot define sigma and cl, use cl only.")
            else:
                cl = scipy.stats.chi2(1).ppf(cl)

        if cl is None:
            cl = 0.68268949  # scipy.stats.chi2(1).cdf(1)

        if method is None:
            # TODO: legacy, remove 0.6
            from zfit.minimize import Minuit
            if isinstance(self.minimizer, Minuit):
                method = 'minuit_minos'
                warnings.warn("'minuit_minos' will be changed as the default errors method to a custom implementation"
                              "with the same functionality. If you want to make sure that 'minuit_minos' will be used "
                              "in the future, add it explicitly as in `errors(method='minuit_minos')`", FutureWarning)
            else:
                method = self._default_error
        if error_name is None:
            if not isinstance(method, str):
                raise ValueError("Need to specify `error_name` or use a string as `method`")
            error_name = method

        if method == 'zfit_error':
            warnings.warn("'zfit_error' is still experimental and may fails.", ExperimentalFeatureWarning)

        params = self._input_check_params(params)

        with self._input_check_reset_params(self.params.keys()):
            # TODO: cache with cl!
            uncached_params = self._get_uncached_params(params=params, method_name=error_name)

            new_result = None

            if uncached_params:
                error_dict, new_result = self._error(params=uncached_params, method=method, cl=cl)
                if new_result is None:
                    self._cache_errors(error_name=error_name, errors=error_dict)
                else:
                    msg = "Invalid, a new minimum was found."
                    self._cache_errors(error_name=error_name, errors={p: msg for p in params})
                    self._valid = False
                    new_result._cache_errors(error_name=error_name, errors=error_dict)
        all_errors = OrderedDict((p, self.params[p][error_name]) for p in params)

        return all_errors, new_result

    def _error(self, params, method, cl):
        if not callable(method):
            try:
                method = self._error_methods[method]
            except KeyError:
                raise KeyError("The following method is not a valid, implemented method: {}".format(method))
        return method(result=self, params=params, cl=cl)

    def covariance(self, params: ParamsTypeOpt = None, method: Union[str, Callable] = None, as_dict: bool = False):
        """Calculate the covariance matrix for `params`.

            Args:
                params: The parameters to calculate
                    the covariance matrix. If `params` is `None`, use all *floating* parameters.
                method: The method to use to calculate the covariance matrix. Valid choices are
                    {'minuit_hesse', 'hesse_np'} or a Callable.
                as_dict: Default `False`. If `True` then returns a dictionnary.

            Returns:
                2D `numpy.array` of shape (N, N);
                `dict`(param1, param2) -> covariance if `as_dict == True`.
        """
        if method is None:
            # LEGACY START
            method = self._default_hesse
            from zfit.minimizers.minimizer_minuit import Minuit

            if isinstance(self.minimizer, Minuit):
                method = "minuit_hesse"
            # LEGACY END

        if method not in self._covariance_dict:
            with self._input_check_reset_params(params) as params:
                self._covariance_dict[method] = self._covariance(method=method)

        params = self._input_check_params(params)
        covariance = {k: self._covariance_dict[method][k] for k in itertools.product(params, params)}

        if as_dict:
            return covariance
        else:
            return dict_to_matrix(params, covariance)

    def _covariance(self, method):
        if not callable(method):
            try:
                method = self._hesse_methods[method]
            except KeyError:
                raise KeyError("The following method is not a valid, implemented method: {}".format(method))

        params = list(self.params.keys())

        if any(data.weights is not None for data in self.loss.data):
            return covariance_with_weights(method=method, result=self, params=params)
        else:
            return method(result=self, params=params)

    def correlation(self, params: ParamsTypeOpt = None, method: Union[str, Callable] = None, as_dict: bool = False):
        """Calculate the correlation matrix for `params`.

            Args:
                params: The parameters to calculate
                    the correlation matrix. If `params` is `None`, use all *floating* parameters.
                method: The method to use to calculate the correlation matrix. Valid choices are
                    {'minuit_hesse', 'hesse_np'} or a Callable.
                as_dict: Default `False`. If `True` then returns a dictionnary.

            Returns:
                2D `numpy.array` of shape (N, N);
                `dict`(param1, param2) -> correlation if `as_dict == True`.
        """

        covariance = self.covariance(params=params, method=method, as_dict=False)
        correlation = covariance_to_correlation(covariance)

        if as_dict:
            params = self._input_check_params(params)
            return matrix_to_dict(params, correlation)
        else:
            return correlation

    def __str__(self):
        string = Style.BRIGHT + f'FitResult' + Style.NORMAL + f' of\n{self.loss} \nwith\n{self.minimizer}\n\n'
        string += tabulate(
            [[color_on_bool(self.valid), color_on_bool(self.converged, on_true=False),
              color_on_bool(self.params_at_limit, on_true=colored.bg(9), on_false=False),
              format_value(self.edm, highprec=False),
              format_value(self.fmin)]],
            ['valid', 'converged', 'param at limit', 'edm', 'min value'],
            tablefmt='fancy_grid',
            disable_numparse=True)
        string += '\n\n' + Style.BRIGHT + "Parameters\n" + Style.NORMAL
        string += str(self.params)
        return string

    def _repr_pretty_(self, p, cycle):
        if cycle:
            p.text(self.__repr__())
            return
        p.text(self.__str__())

def covariance_to_correlation(covariance):
    diag = np.diag(1 / np.diag(covariance) ** 0.5)
    return np.matmul(diag, np.matmul(covariance, diag))

def format_value(value, highprec=True):
    try:
        import iminuit
        m_error_class = iminuit.util.MError
    except ImportError:
        m_error_class = dict

    if isinstance(value, dict) and 'error' in value:
        value = value['error']
        value = f"{value:> 6.2g}"
        value = f'+/-{" " * (8 - len(value))}' + value
    if isinstance(value, m_error_class) or (isinstance(value, dict) and 'lower' in value and 'upper' in value):
        if isinstance(value, m_error_class):
            lower = value.lower
            upper = value.upper
        else:
            lower = value['lower']
            upper = value['upper']
        lower_sign = f"{np.sign(lower): >+}"[0]
        upper_sign = f"{np.sign(upper): >+}"[0]
        lower, upper = f"{np.abs(lower): >6.2g}", f"{upper: >6.2g}"
        lower = lower_sign + " " * (7 - len(lower)) + lower
        upper = upper_sign + " " * (7 - len(upper)) + upper
        # lower += " t" * (11 - len(lower))
        value = lower + " " * 3 + upper

    if isinstance(value, float):
        if highprec:
            value = f"{value:> 6.4g}"
        else:
            value = f"{value:> 6.2g}"
    return value

def color_on_bool(value, on_true=colored.bg(10), on_false=colored.bg(9)):
    if not value and on_false:
        value_add = on_false
    elif value and on_true:
        value_add = on_true
    else:
        value_add = ''
    value = value_add + str(value) + Style.RESET_ALL
    return value

class ParamHolder(dict):  # no UserDict, we only want to change the __str__

    def __str__(self) -> str:
        order_keys = ['value', 'hesse']
        keys = OrderedSet()
        for pdict in self.values():
            keys.update(OrderedSet(pdict))
        order_keys = OrderedSet([key for key in order_keys if key in keys])
        order_keys.update(keys)

        rows = []
        for param, pdict in self.items():
            row = [param.name]
            row.extend(format_value(pdict.get(key, ' ')) for key in order_keys)
            row.append(color_on_bool(run(param.at_limit), on_true=colored.bg('light_red'), on_false=False))
            rows.append(row)

        order_keys = ['name'] + list(order_keys) + ['at limit']
        table = tabulate(rows, order_keys, numalign="right", stralign='right', colalign=('left',))
        return table


import cv2
import numpy as np
tmp_file_idx = 0
# 这个类每次进行预处理的时候都会重新声明,所以init每次预处理都会执行,随机数并不会被固定
class VideoTransform():
    def __init__(self, size,
                 flip_p=0.5,  # 反转概率
                 mean=127,  # 通道均值
                 std=1.,
                 horizontal_flip= False,  # 随机水平翻转
                 vertical_flip = False,  # 随机垂直翻转
                 random_sample = False,  # 随机重采样
                 use_bright_contrast = False,  # 随机调整亮度对比度

                 ):
        if isinstance(size, tuple):
            self.size = size
        elif isinstance(size, list):
            self.size = (size[0], size[1])
        else:
            self.size = (size, size)
        self.flip_p = flip_p
        self.mean = mean
        self.std = std
        self.use_horizontal_flip = horizontal_flip
        self.use_vertical_flip =vertical_flip
        self.use_random_sample =random_sample
        self.use_bright_contrast = use_bright_contrast

        self.constrast_p = np.random.rand()
        self.brightness_p = np.random.rand()
        self.horizontal_p = np.random.rand()  # 翻转的随机数
        self.vertical_p = np.random.rand()
        self.sample_p = np.random.rand()

        if self.brightness_p>0.5: # 亮度
            self.beta = 0
        else:
            self.beta = np.random.uniform(-50,50)
        if self.constrast_p >0.5: # 对比度
            self.alpha = 1
        else:
            self.alpha = np.random.uniform(0.7, 1.3)

        # 随机裁剪大小
        self.sample_new_h = np.random.randint(int(self.size[1] / 2), self.size[1])
        self.sample_new_w = np.random.randint(int(self.size[0] / 2), self.size[0])
        # 随机裁剪起点坐标
        self.sample_new_y = np.random.randint(0, self.size[1] - self.sample_new_h)
        self.sample_new_x = np.random.randint(0, self.size[0] - self.sample_new_w)
        # 随即裁剪填充起点
        self.sample_padding_y = np.random.randint(0, self.size[1] - self.sample_new_h)
        self.sample_padding_x = np.random.randint(0, self.size[0] - self.sample_new_w)

    # 随机亮度&对比度变换
    def bright_constrast_adjust(self,data):

        blank = np.zeros_like(data)
        data = cv2.addWeighted(data,self.alpha,blank,1-self.alpha,self.beta)
        return data

    # 水平翻转
    def horizontal_flip(self,
                        type,  # 需要反转的类型,接受图像或者是lebel
                        data):
        if self.horizontal_p > self.flip_p:
            return data
        if type == 'img':  # 针对图像做水平反转
            img = data
            img = img[:, ::-1].copy()
            return img
        elif type == 'label':  # 针对leibel做水平反转
            label = data
            for label_idx in range(len(label)):
                label_t = label[label_idx]
                label_t[0] = 1. - label_t[0]
                label_t[2] = 1. - label_t[2]
                label_t = label_t[[2, 1, 0, 3, 4]]  # 左右翻转后会从左上右下点变成右上左下点，重新修正回左上右下的格式
                label[label_idx] = label_t
            return label.copy()
        else:
            raise Exception('wrong')

    # 垂直翻转
    def vertical_flip(self, type, data):
        if self.vertical_p > self.flip_p:
            return data
        if type == 'img':
            img = data
            img = img[::-1, :].copy()
            return img
        elif type == 'label':
            label = data
            for label_idx in range(len(label)):
                label_t = label[label_idx]
                label_t[1] = 1. - label_t[1]
                label_t[3] = 1. - label_t[3]
                label_t = label_t[[0, 3, 2, 1, 4]]  # 上下翻转后会从左上右下点变成右上左下点，重新修正回左上右下的格式
                label[label_idx] = label_t
            return label.copy()
        else:
            raise Exception('wrong')

    def __call__(self, data):
        video = data
        # todo
        res_video = np.empty((self.size[0], self.size[1],video.shape[-1]), dtype=np.float32)
        # res_video = np.empty((self.size[0], self.size[1],video.shape[-2], video.shape[-1]), dtype=np.float32)

        # img
        for idx in range(video.shape[-1]):
            # resize
            img = video[...,idx]
            if self.size != img.shape[:2]:
                img = cv2.resize(img, self.size)
            # 亮度对比度
            if self.use_bright_contrast:
                img = self.bright_constrast_adjust(img)
            # # 裁切
            # if self.use_random_sample:
            #     img = self.rand_sample("img", img)
            # 翻转
            if self.use_horizontal_flip:
                img = self.horizontal_flip('img',img)
            if self.use_vertical_flip:
                img = self.vertical_flip('img',img)

            img = img.astype(np.float32)
            img = img - self.mean
            img = img * (self.std / 255.)
            res_video[...,idx] = img.copy()
        return res_video

class ImgBrightAndContrastTransform():
    def __init__(self):
        self.rand_bright = np.random

def show(video, label):
    img = video[-1]
    h, w, c = img.shape
    for label_t in label:
        label_t = label_t * (w, h, w, h, 1)
        label_t = label_t.astype(np.int32)
        cv2.rectangle(img, (label_t[0], label_t[1]), (label_t[2], label_t[3]), (0, 0, 255))
    cv2.imwrite("tmp.jpg", img)


# Merge stimuli information from spike2 mat file into Kwik file

import h5py as h5
import tables
import os
import numpy as np
import argparse
import glob

try: import simplejson as json
except ImportError: import json

from klusta_pipeline.dataio import load_recordings, save_info, load_digmark, load_stim_info
from klusta_pipeline.utils import get_import_list, validate_merge, realign

def get_args():

    parser = argparse.ArgumentParser(description='Compile Spike2 epoch .mat files into KlustaKwik KWD file.')
    parser.add_argument('path', default = './', nargs='?',
                       help='directory containing all of the mat files to compile')
    parser.add_argument('dest', default = './', nargs='?',
                       help='destination directory for kwd and other files')
    return parser.parse_args()

def get_rec_samples(kwd_file,index):
    with h5.File(kwd_file, 'r') as kwd:
        return kwd['/recordings/{}/data'.format(index)].shape[0]

def merge_recording_info(klu_path,mat_path):
    batch = klu_path.split('__')[-1]
    with open(os.path.join(klu_path,batch+'_info.json')) as f:
        info = json.load(f)

    assert 'recordings' not in info

    import_list = get_import_list(mat_path,info['exports'])
    for item in import_list:
        assert os.path.exists(item), item

    mat_data = validate_merge(import_list,info['omit'])
    fs = info['params']['fs']

    chans = set(mat_data[0]['chans'])
    for d2 in mat_data[1:]:
        chans = chans.intersection(d2['chans'])
    chans = list(chans)

    for i,m in zip(info['exports'],mat_data):
        i['chans'] = chans

    rec_list = []
    for import_file in import_list:
        recordings = load_recordings(import_file,chans)
        for r in recordings:
            rec = realign(r,chans,fs,'spline')
            del rec['data']
            rec_list.append(rec)

    info['recordings'] = [{k:v for k,v in rec.items() if k is not 'data'} for rec in rec_list]
    save_info(klu_path,info)
    return info

def merge(spike2mat_folder, kwik_folder):

    info_json = glob.glob(os.path.join(kwik_folder,'*_info.json'))[0]
    with open(info_json, 'r') as f:
        info = json.load(f)

    kwik_data_file = os.path.join(kwik_folder,info['name']+'.kwik')
    kwd_raw_file = os.path.join(kwik_folder,info['name']+'.raw.kwd')\

    with tables.open_file(kwik_data_file, 'r+') as kkfile:

        digmark_timesamples = []
        digmark_recording = []
        digmark_codes = []
        stimulus_timesamples = []
        stimulus_recording = []
        stimulus_codes = []
        stimulus_names = []

        spike_recording_obj = kkfile.get_node('/channel_groups/0/spikes','recording')
        spike_time_samples_obj = kkfile.get_node('/channel_groups/0/spikes','time_samples')

        spike_recording = spike_recording_obj.read()
        spike_time_samples = spike_time_samples_obj.read()

        try:
            assert 'recordings' in info
        except AssertionError:
            info = merge_recording_info(kwik_folder,spike2mat_folder)

        order = np.sort([str(ii) for ii in range(len(info['recordings']))])
        print order
	print len(spike_recording)
        is_done = np.zeros(spike_recording.shape,np.bool_)
        for rr,rid_str in enumerate(order):
            # rr: index of for-loop
            # rid: recording id
            # rid_str: string form of recording id
            rid = int(rid_str)
            rec = info['recordings'][rid]

            n_samps = get_rec_samples(kwd_raw_file,rid)

            #is_done = np.vectorize(lambda x: x not in done)

            todo = ~is_done & (spike_time_samples >= n_samps)
            print "rec {}: {} spikes done".format(rid,is_done.sum())
            print "setting {} spikes to next cluster".format(todo.sum())
            if todo.sum()>0:
                spike_recording[todo] = int(order[rr+1])
                spike_time_samples[todo] -= n_samps
            is_done = is_done | ~todo
            print is_done.sum()

            t0 = rec['start_time']
            fs = rec['fs']
            dur = float(n_samps) / fs

            s2mat = os.path.split(rec['file_origin'])[-1]
            s2mat = os.path.join(spike2mat_folder, s2mat)

            codes, times = load_digmark(s2mat)
            rec_mask = (times >= t0) * (times < (t0+dur))

            codes = codes[rec_mask]
            times = times[rec_mask] - t0
            time_samples = (times * fs).round().astype(np.uint64)
            recording = rid * np.ones(codes.shape,np.uint16)

            digmark_timesamples.append(time_samples)
            digmark_recording.append(recording)
            digmark_codes.append(codes)

            codes, times, names = load_stim_info(s2mat)
            rec_mask = (times >= t0) * (times < (t0+dur))

            codes = codes[rec_mask]
            names = names[rec_mask]
            times = times[rec_mask] - t0
            time_samples = (times * fs).round().astype(np.uint64)
            recording = rid * np.ones(codes.shape,np.uint16)

            stimulus_timesamples.append(time_samples)
            stimulus_recording.append(recording)
            stimulus_codes.append(codes)
            stimulus_names.append(names)

        digmark_timesamples = np.concatenate(digmark_timesamples)
        digmark_recording = np.concatenate(digmark_recording)
        digmark_codes = np.concatenate(digmark_codes)
        stimulus_timesamples = np.concatenate(stimulus_timesamples)
        stimulus_recording = np.concatenate(stimulus_recording)
        stimulus_codes = np.concatenate(stimulus_codes)
        stimulus_names = np.concatenate(stimulus_names)

        print digmark_timesamples.dtype
        print digmark_recording.dtype
        print digmark_codes.dtype
        print stimulus_timesamples.dtype
        print stimulus_recording.dtype
        print stimulus_codes.dtype
        print stimulus_names.dtype

        kkfile.create_group("/", "event_types", "event_types")

        kkfile.create_group("/event_types", "DigMark")
        kkfile.create_earray("/event_types/DigMark", 'time_samples', obj=digmark_timesamples)
        kkfile.create_earray("/event_types/DigMark", 'recording', obj=digmark_recording)
        kkfile.create_earray("/event_types/DigMark", 'codes', obj=digmark_codes)

        kkfile.create_group("/event_types", "Stimulus")
        kkfile.create_earray("/event_types/Stimulus", 'time_samples', obj=stimulus_timesamples)
        kkfile.create_earray("/event_types/Stimulus", 'recording', obj=stimulus_recording)
        kkfile.create_earray("/event_types/Stimulus", 'codes', obj=stimulus_codes)
        kkfile.create_earray("/event_types/Stimulus", 'text', obj=stimulus_names)

        spike_recording_obj[:] = spike_recording
        spike_time_samples_obj[:] = spike_time_samples

def main():
    args = get_args()
    spike2mat_folder = os.path.abspath(args.path)
    kwik_folder = os.path.abspath(args.dest)
    merge(spike2mat_folder, kwik_folder)

if __name__ == '__main__':
    main()


# 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