import os
import numpy as np
# crystalball is single sided, local reimplementation of double-sided here until
# the PR can be merged
# from scipy.stats import crystalball
from .doublecrystalball import doublecrystalball
from awkward import JaggedArray
from coffea.lookup_tools.dense_lookup import dense_lookup


class rochester_lookup():

    def __init__(self, wrapped_values):
        self._build(wrapped_values)

    def _build(self, wrapped_values):

        self._nsets = wrapped_values['nsets']
        self._members = wrapped_values['members']

        # now build the lookup tables
        # for data scale, simple, just M A in bins of eta,phi
        edges = wrapped_values['edges']['scales']
        M = wrapped_values['values']['M']
        A = wrapped_values['values']['A']
        self._M = {s: {m: {t: dense_lookup(M[s][m][t], edges) for t in M[s][m]} for m in M[s]} for s in M}
        self._A = {s: {m: {t: dense_lookup(A[s][m][t], edges) for t in A[s][m]} for m in A[s]} for s in A}

        # for mc scale, more complicated
        # version 1 if gen pt available
        # only requires the kRes lookup
        edges = wrapped_values['edges']['res']
        kRes = wrapped_values['values']['kRes']
        self._kRes = {s: {m: {t: dense_lookup(kRes[s][m][t], edges) for t in kRes[s][m]} for m in kRes[s]} for s in kRes}

        # version 2 if gen pt not available
        edges = wrapped_values['edges']['cb']
        rsPars = wrapped_values['values']['rsPars']
        cbS = wrapped_values['values']['cbS']
        cbA = wrapped_values['values']['cbA']
        cbN = wrapped_values['values']['cbN']
        self._rsPars = {s: {m: {t: dense_lookup(rsPars[s][m][t], edges) for t in rsPars[s][m]} for m in rsPars[s]} for s in rsPars}
        self._cbS = {s: {m: dense_lookup(cbS[s][m], edges) for m in cbS[s]} for s in cbS}
        self._cbA = {s: {m: dense_lookup(cbA[s][m], edges) for m in cbA[s]} for s in cbA}
        self._cbN = {s: {m: dense_lookup(cbN[s][m], edges) for m in cbN[s]} for s in cbN}

        self._loaduncs = len(self._M.keys()) > 1

    def _error(self, func, *args):
        if not self._loaduncs:
            return None

        newargs = args + (0, 0)
        default = func(*newargs)
        result = np.zeros_like(default)
        for s in range(self._nsets):
            oneOver = 1.0 / self._members[s]
            for m in range(self._members[s]):
                newargs = args + (s, m)
                d = func(*newargs) - default
                result = result + d * d * oneOver
        return result**0.5

    def kScaleDT(self, charge, pt, eta, phi, s=0, m=0):
        '''Momentum scale correction for data
            required:
                charge
                pt
                eta
                phi
            optional:
                s: Rochester correction set to use (default 0)
                m: Rochester correction member to use (default 0)
        '''
        # type = 1 corresponds to data
        M = self._M[s][m][1](eta, phi)
        A = self._A[s][m][1](eta, phi)
        return 1.0 / (M + charge * A * pt)

    def kScaleDTerror(self, charge, pt, eta, phi):
        '''Momentum scale correction uncertainty for data
            required:
                charge
                pt
                eta
                phi
        '''
        return self._error(self.kScaleDT, charge, pt, eta, phi)

    def kScaleMC(self, charge, pt, eta, phi, s=0, m=0):
        '''Momentum scale correction for mc (not recommended, use kSpreadMC instead)
            required:
                charge
                pt
                eta
                phi
            optional:
                s: Rochester correction set to use (default 0)
                m: Rochester correction member to use (default 0)
        '''
        # type = 0 corresponds to mc
        M = self._M[s][m][0](eta, phi)
        A = self._A[s][m][0](eta, phi)
        return 1.0 / (M + charge * A * pt)

    def kScaleMCerror(self, charge, pt, eta, phi):
        '''Momentum scale correction uncertainty for mc (not recommended, use kSpreadMC instead)
            required:
                charge
                pt
                eta
                phi
        '''
        return self._error(self.kScaleMC, charge, pt, eta, phi)

    def kSpreadMC(self, charge, pt, eta, phi, genpt, s=0, m=0):
        '''Momentum scale correction for mc (if genpt not availble, use kSmearMC)
            required:
                charge
                pt
                eta
                phi
                genpt: the match gen particle pt
            optional:
                s: Rochester correction set to use (default 0)
                m: Rochester correction member to use (default 0)
        '''
        k = self.kScaleMC(charge, pt, eta, phi, s, m)
        return k * self._kSpread(genpt, k * pt, eta, s, m)

    def kSpreadMCerror(self, charge, pt, eta, phi, genpt):
        '''Momentum scale correction uncertainty for mc (if genpt not availble, use kSmearMC)
            required:
                charge
                pt
                eta
                phi
                genpt: the match gen particle pt
        '''
        return self._error(self.kSpreadMC, charge, pt, eta, phi, genpt)

    def _kSpread(self, genpt, kpt, eta, s=0, m=0):
        x = genpt / kpt
        abseta = abs(eta)
        kData = self._kRes[s][m][1](abseta)  # type 1 is data
        kMC = self._kRes[s][m][0](abseta)    # type 0 is MC
        return x / (1.0 + (x - 1.0) * kData / kMC)

    def kSmearMC(self, charge, pt, eta, phi, nl, u, s=0, m=0):
        '''Momentum scale correction for mc (should prefer to use kSpreadMC)
            required:
                charge
                pt
                eta
                phi
                nl: number of tracker layers with measurements
                u: random float between 0 and 1
            optional:
                s: Rochester correction set to use (default 0)
                m: Rochester correction member to use (default 0)
        '''
        k = self.kScaleMC(charge, pt, eta, phi, s, m)
        return k * self._kExtra(k * pt, eta, nl, u, s, m)

    def kSmearMCerror(self, charge, pt, eta, phi, nl, u):
        '''Momentum scale correction uncertainty for mc (should prefer to use kSpreadMC)
            required:
                charge
                pt
                eta
                phi
                nl: number of tracker layers with measurements
                u: random float between 0 and 1
        '''
        return self._error(self.kSmearMC, charge, pt, eta, phi, nl, u)

    def _sigma(self, pt, eta, nl, s=0, m=0):
        dpt = pt - 45
        abseta = abs(eta)
        return self._rsPars[s][m][0](abseta, nl) + self._rsPars[s][m][1](abseta, nl) * dpt + self._rsPars[s][m][2](abseta, nl) * dpt * dpt

    def _kExtra(self, kpt, eta, nl, u, s=0, m=0):
        # if it is a jagged array, save the offsets then flatten everything
        # needed for the ternary conditions later
        offsets = None
        if isinstance(kpt, JaggedArray):
            offsets = kpt.offsets
            kpt = kpt.flatten()
            eta = eta.flatten()
            nl = nl.flatten()
            u = u.flatten()
        abseta = abs(eta)
        kData = self._kRes[s][m][1](abseta)  # type 1 is data
        kMC = self._kRes[s][m][0](abseta)    # type 0 is MC
        mask = (kData > kMC)
        x = np.zeros_like(kpt)
        sigma = self._sigma(kpt, eta, nl, s, m)
        # Rochester cbA = beta, cbN = m, as well as cbM (always 0?) = loc and cbS = scale to transform y = (x-loc)/scale in the pdf method
        cbA = self._cbA[s][m](abseta, nl)
        cbN = self._cbN[s][m](abseta, nl)
        loc = np.zeros_like(u)
        cbS = self._cbS[s][m](abseta, nl)
        invcdf = doublecrystalball.ppf(u, cbA, cbA, cbN, cbN, loc, cbS)
        x[mask] = np.sqrt(kData[mask] * kData[mask] - kMC[mask] * kMC[mask]) * sigma[mask] * invcdf[mask]
        result = np.ones_like(kpt)
        result[(x > -1)] = 1.0 / (1.0 + x[x > -1])
        if offsets is not None:
            result = JaggedArray.fromoffsets(offsets, result)
        return result
