from scipy.integrate import nquad, qmc_quad
import numpy as np
import logging
from scipy.stats import rel_breitwigner, norm

logging.basicConfig(format='%(levelname)s:%(message)s')
logger = logging.getLogger(__name__)
logger.addHandler(logging.NullHandler())

try:
    import vegas

    def vegas_int(f,ranges,**kwargs):

        # @vegas.batchintegrand
        def fwrap(x):
            # print("test", x.shape)
            return f(*x)

        @vegas.batchintegrand
        def fbatch(x):
            return np.apply_along_axis(fwrap, axis=1, arr=x)

        if 'peak_sampling' in kwargs:
            del kwargs['peak_sampling']

        map = vegas.AdaptiveMap(ranges)
        peak_pos = kwargs.get('peak_pos', None)
        if 'peak_pos' in kwargs:
            del kwargs['peak_pos']
        if peak_pos:
            neval = kwargs.get('neval', 100)
            shat_ind_neval = int(neval/(len(peak_pos) + 1))
            shat_data = np.concatenate(
                            [rel_breitwigner.rvs(
                                rho=p[0]/p[1],
                                scale=p[1],
                                size=shat_ind_neval
                                )**2/p[2] for p in peak_pos] # normalise by S to get distribution in tau
                            + [np.linspace(ranges[1][0], ranges[1][1], 10)]
                        )
            np.random.shuffle(shat_data)
            shat_data = np.clip(shat_data, ranges[1][0], ranges[1][1])
            that_data = norm.rvs(0, np.abs(0.0005*ranges[0][0]), size=len(shat_data) - 10)
            that_data = -np.abs(that_data)
            that_data = np.concatenate((that_data, np.linspace(ranges[0][0], ranges[0][1], 10)))
            np.random.shuffle(that_data)
            x = np.column_stack((that_data, shat_data))
            map.adapt_to_samples(x, fbatch(x), nitn=kwargs.get('nitn', 5))
        integ = vegas.Integrator(map)

        summary = kwargs.get('summary', False)
        if 'summary' in kwargs:
            del kwargs['summary']
        discard = kwargs.get('discard', True)
        if 'discard' in kwargs:
            del kwargs['discard']
        warn = kwargs.get('warn', True)
        if 'warn' in kwargs:
            del kwargs['warn']

        if discard:
            integ(fwrap, **kwargs)
        result = integ(fwrap, **kwargs)
        if summary:
            print(result.summary())
            print(f'result = {result:.3f}; Q={result.Q}; chi2={result.chi2}; dof={result.dof}; sum_neval={result.sum_neval}; avg_neval={result.avg_neval}')
        if warn and result.Q < 0.01:
            logger.warning(f'Error estimate of integral may be unreliable (Q={result.Q}). Set summary=True for more information. Consider increasing "neval" by a factor 5-10 and check if results agree within errors!')
        return result.val, result.sdev
except ImportError:
    def vegas_int(f,ranges,**kwargs):
        return nquad(f,ranges)
    logger.error('Vegas integrator not available. Falling back to scipy.nquad. If you want to use vega, install it using "pip install --user vegas"')

def qmc_quad_int(f,ranges,**kwargs):

    def fwrap(x):
        return f(*x)
    r = np.transpose(ranges)
    return qmc_quad(fwrap, r[0], r[1], **kwargs)


class anyIntegrator(object):
    defaults = {
            'nquad': {'opts':{'epsabs': 1e-5, 'epsrel': 1e-5, 'limit': 200}},
            'qmc_quad': {'n_estimates': 8, 'n_points': 1024},
            'vegas': {'nitn':10, 'neval': 1000,'summary': False, 'discard': True, 'warn': True}
            }
    integrators = {'nquad': nquad, 'qmc_quad': qmc_quad_int, 'vegas': vegas_int}

    def __init__(self, integrator='vegas', **options):
        self.options = anyIntegrator.defaults.get(integrator, {})
        self.options.update(options)
        self.integrator = anyIntegrator.integrators.get(integrator,integrator)

    def __call__(self,*args,integrator=None,**kwargs):
        func = self.integrator if not integrator else anyIntegrator.integrators.get(integrator,integrator)
        options = self.options if not integrator else anyIntegrator.defaults.get(integrator,{})
        options = dict(options)
        options.update(**kwargs)
        result = func(*args,**options)
        if hasattr(result, 'integral') and hasattr(result, 'standard_error'):
            return result.integral, result.standard_error
        return result
