import surpyval
from surpyval import np
from scipy.linalg import inv
from scipy.optimize import minimize
from scipy.special import ndtri as z

from autograd import jacobian, hessian
from surpyval import parametric as para


class MixtureModel():
    """
    Generalised from algorithm found here
    https://www.sciencedirect.com/science/article/pii/S0307904X12002545
    """
    def __init__(self, x, dist=para.Weibull, **kwargs):
        self.m = kwargs.pop('m', 2)
        raw_data = {}
        c = kwargs.pop('c', None)
        n = kwargs.pop('n', None)
        raw_data['x'] = x
        raw_data['c'] = c
        raw_data['n'] = n
        self.raw_data = raw_data
        self.dist = dist

        x, c, n = surpyval.xcn_handler(x, c, n)
        assert len(x) > self.m * (self.dist.k + 1)

        self.x = x
        self.c = c
        self.n = n

        self.N = n.sum()
        splits_x = np.array_split(x, self.m)
        splits_c = np.array_split(c, self.m)
        splits_n = np.array_split(n, self.m)
        params = np.zeros(shape=(self.m, self.dist.k))

        for i in range(self.m):
            params[i, :] = self.dist.fit(x=splits_x[i], c=splits_c[i],
                                         n=splits_n[i]).params
        self.params = params
        self.w = np.ones(shape=(self.m)) / self.m
        self.p = np.ones(shape=(self.m, len(self.x))) / self.m

    def Q(self, params):
        params = params.reshape(self.m, self.dist.k)
        f = np.zeros_like(self.p)
        for i in range(self.m):
            like = self.dist.like(self.x, self.c, self.n, *params[i])
            like += surpyval.TINIEST
            like = np.where(like < 1, like, 1)
            like = np.log(like)
            like = np.multiply(self.n, like)

            f[i] = np.multiply(self.p[i], like)
        f = -np.sum(f)
        self.loglike = f
        return f

    def expectation(self):
        for i in range(self.m):
            like = self.dist.like(self.x, self.c, self.n, *self.params[i])
            like = np.multiply(self.w[i], like)
            self.p[i] = like
        self.p = np.divide(self.p, np.sum(self.p, axis=0))
        self.w = np.sum(self.p, axis=1) / self.N

    def maximisation(self):
        bounds = self.dist.bounds * self.m
        res = minimize(self.Q, self.params, bounds=bounds)
        # if not res.success:
        #   print(res)
        #   raise Exception('Max failed')
        self.params = res.x.reshape(self.m, self.dist.k)

    def EM(self):
        self.expectation()
        self.maximisation()

    def _em(self, tol=1e-6, max_iter=1000):
        i = 0
        self.EM()
        f0 = self.loglike
        self.EM()
        f1 = self.loglike
        while (np.abs(f0 - f1) > tol) and (i < max_iter):
            f0 = f1
            self.EM()
            f1 = self.loglike
            i += 1
        if i >= 1000:
            print('Max iterations reached')

    def neg_ll(self, x, c, n, *params):
        f = np.zeros_like(self.p)
        params = np.reshape(params, (self.m, self.dist.k + 1))
        f = np.zeros_like(x)
        for i in range(self.m):
            like = self.dist.like(x, c, n, *params[i, 1::])
            like = np.multiply(params[i, 0], like)
            f = f + like
        f = np.where(f <= 0, surpyval.TINIEST, f)
        f = np.where(f < 1, f, 1)
        f = np.log(f)
        f = np.multiply(n, f)
        f = -np.sum(f)
        return f

    def _mle(self):

        def fun(x):
            return self.neg_ll(self.x, self.c, self.n, *x)

        bounds = tuple(((0, 1), *self.dist.bounds)) * self.m
        init = np.hstack([np.atleast_2d(self.w).T, self.params])

        def w_sums_to_1(x):
            xx = np.reshape(x, (self.m, self.dist.k + 1))
            return xx[:, 0].sum() - 1
        cons = {'type': 'eq', 'fun': w_sums_to_1}
        jac = jacobian(fun)
        hess = hessian(fun)
        res = minimize(fun, init.flatten(),
                       jac=jac, bounds=bounds,
                       constraints=[cons])
        self.hess_inv = inv(hess(res.x))
        self.res = res
        self.params = np.reshape(res.x, (self.m, self.dist.k + 1))[:, 1::]
        self.w = np.reshape(res.x, (self.m, self.dist.k + 1))[:, 0]

    def fit(self, how='EM'):
        if how == 'EM':
            self._em()
            self.method = 'EM'
        elif how == 'MLE':
            self._mle()
            self.method = 'MLE'

    def ll(self):
        return self.loglike

    def R_cb(self, t, cb=0.05):

        def ssf(params):
            params = np.reshape(params, (self.m, self.dist.k + 1))
            F = np.zeros_like(t)
            for i in range(self.m):
                F = F + params[i, 0] * self.dist.ff(t, *params[i, 1::])
            return 1 - F

        pvars = self.hess_inv[np.triu_indices(self.hess_inv.shape[0])]
        with np.errstate(all='ignore'):
            jac = jacobian(ssf)(self.res.x)

        var_u = []
        for i, j in enumerate(jac):
            j = np.atleast_2d(j).T * j
            j = j[np.triu_indices(j.shape[0])]
            var_u.append(np.sum(j * pvars))
        diff = (z(cb/2) * np.sqrt(np.array(var_u))
                * np.array([1., -1.]).reshape(2, 1))
        R_hat = self.sf(t)
        exponent = diff/(R_hat*(1 - R_hat))
        R_cb = R_hat / (R_hat + (1 - R_hat) * np.exp(exponent))
        return R_cb.T

    def mean(self):
        mean = 0
        for i in range(self.m):
            mean += self.w[i] * self.dist.mean(*self.params[i])
        return mean

    def random(self, size):
        sizes = np.random.multinomial(size, self.w)
        rvs = np.zeros(size)
        s_last = 0
        for i, s in enumerate(sizes):
            rvs[s_last:s+s_last] = self.dist.random(s, *self.params[i, :])
            s_last = s
        np.random.shuffle(rvs)
        return rvs

    def df(self, t):
        f = np.zeros_like(t)
        for i in range(self.m):
            f += self.w[i] * self.dist.df(t, *self.params[i])
        return f

    def ff(self, t):
        F = np.zeros_like(t)
        for i in range(self.m):
            F = F + self.w[i] * self.dist.ff(t, *self.params[i])
        return F

    def sf(self, t):
        return 1 - self.ff(t)

    def cs(self, t, T):
        return self.sf(t + T) / self.sf(T)
