"""
Implements all the logic for mixed logit models
"""
# pylint: disable=invalid-name
import scipy.stats
from scipy.optimize import minimize
from ._choice_model import ChoiceModel
from ._device import device as dev
import numpy as np


class MixedLogit(ChoiceModel):
    """Class for estimation of Mixed Logit Models"""

    def __init__(self):
        """Init Function"""
        super(MixedLogit, self).__init__()
        self.rvidx = None  # Boolean index of random vars in X. True = rand var
        self.rvdist = None

    # X: (N, J, K)
    def fit(self, X, y, varnames=None, alt=None, isvars=None,
            randvars=None, mixby=None, base_alt=None, fit_intercept=False,
            init_coeff=None, maxiter=2000, random_state=None,
            n_draws=200, halton=True, verbose=1):
        self._validate_inputs(X, y, alt, varnames, isvars,
                              base_alt, fit_intercept, maxiter)
        self._pre_fit(alt, varnames, isvars, base_alt,
                      fit_intercept, maxiter)

        if random_state is not None:
            np.random.seed(random_state)

        X, Xnames = self._setup_design_matrix(X)
        J, K, R = X.shape[1], X.shape[2], n_draws

        if mixby is not None:  # If panel
            X, y, panel_info = self._balance_panels(X, y, mixby)
            N, P = panel_info.shape
        else:
            N, P = X.shape[0], 1
            panel_info = np.ones((N, 1))

        X = X.reshape(N, P, J, K)
        y = y.reshape(N, P, J, 1)

        # Variable that contains a boolean index of random variables.
        self.n_draws = n_draws
        self.randvars = randvars
        rvpos = [np.where(Xnames == rv)[0][0] for rv in self.randvars.keys()]
        self.rvidx = np.zeros(K, dtype=bool)
        self.rvidx[rvpos] = True  # True: Random var, False: Fixed var
        self.rvdist = list(self.randvars.values())

        draws = self._generate_draws(N, R, halton)  # (N,Kr,R)
        n_coeff = K + len(rvpos)
        if init_coeff is None:
            betas = np.repeat(.0, K + len(rvpos))
        else:
            betas = init_coeff
            if len(init_coeff) != n_coeff:
                raise ValueError("The size of init_coeff must be: " + n_coeff)

        if dev.using_gpu:
            X, y = dev.to_gpu(X), dev.to_gpu(y)
            panel_info = dev.to_gpu(panel_info)
            draws = dev.to_gpu(draws)
            if verbose > 0:
                print("**** GPU Processing Enabled ****")

        optimizat_res = minimize(self._loglik_gradient, betas,
                                 jac=True, args=(X, y, panel_info, draws),
                                 method='BFGS', tol=1e-5,
                                 options={'gtol': 1e-4, 'maxiter': maxiter})

        fvpos = list(set(range(len(Xnames))) - set(rvpos))
        coeff_names = np.concatenate((Xnames[fvpos], Xnames[rvpos],
                                      np.char.add("sd.", Xnames[rvpos])))

        self._post_fit(optimizat_res, coeff_names, N, verbose)

    def _compute_probabilities(self, betas, X, panel_info, draws):
        Bf, Br = self._transform_betas(betas, draws)  # Get fixed and rand coef
        Xf = X[:, :, :, ~self.rvidx]  # Data for fixed coefficients
        Xr = X[:, :, :, self.rvidx]   # Data for random coefficients

        XBf = dev.np.einsum('npjk,k -> npj', Xf, Bf)  # (N,P,J)
        XBr = dev.np.einsum('npjk,nkr -> npjr', Xr, Br)  # (N,P,J,R)
        V = XBf[:, :, :, None] + XBr  # (N,P,J,R)
        V[V > 700] = 700
        eV = dev.np.exp(V)
        sumeV = dev.np.sum(eV, axis=2, keepdims=True)
        sumeV[sumeV == 0] = 1e-30
        p = eV/sumeV  # (N,P,J,R)
        p = p*panel_info[:, :, None, None]  # Zero for unbalanced panels
        return p

    def _loglik_gradient(self, betas, X, y, panel_info, draws):
        if dev.using_gpu:
            betas = dev.to_gpu(betas)
        p = self._compute_probabilities(betas, X, panel_info, draws)
        # Probability of chosen alt
        pch = (y*p).sum(axis=2)  # (N,P,R)
        pch = self._prob_product_across_panels(pch, panel_info)  # (N,R)

        # Log-likelihood
        lik = pch.mean(axis=1)  # (N,)
        loglik = dev.np.log(lik).sum()

        # Gradient
        Xf = X[:, :, :, ~self.rvidx]
        Xr = X[:, :, :, self.rvidx]

        ymp = y - p  # (N,P,J,R)
        # Gradient for fixed and random params
        gr_f = dev.np.einsum('npjr,npjk -> nkr', ymp, Xf)
        der = self._compute_derivatives(betas, draws)
        gr_b = dev.np.einsum('npjr,npjk -> nkr', ymp, Xr)*der
        gr_w = dev.np.einsum('npjr,npjk -> nkr', ymp, Xr)*der*draws
        # Multiply gradient by the chose prob. and dived by mean chose prob.
        gr_f = (gr_f*pch[:, None, :]).mean(axis=2)/lik[:, None]  # (N,Kf)
        gr_b = (gr_b*pch[:, None, :]).mean(axis=2)/lik[:, None]  # (N,Kr)
        gr_w = (gr_w*pch[:, None, :]).mean(axis=2)/lik[:, None]  # (N,Kr)
        # Put all gradients in a single array and aggregate them
        grad = dev.np.concatenate((gr_f, gr_b, gr_w), axis=1)  # (N,K,R)
        grad = grad.sum(axis=0)  # (K,)

        if dev.using_gpu:
            grad, loglik = dev.to_cpu(grad), dev.to_cpu(loglik)
        return -loglik, -grad

    def _prob_product_across_panels(self, pch, panel_info):
        if not np.all(panel_info):  # If panel unbalanced. Not all ones
            idx = panel_info == 0
            for i in range(pch.shape[2]):
                pch[:, :, i][idx] = 1  # Multiply by one when unbalanced
        pch = pch.prod(axis=1)  # (N,R)
        pch[pch == 0] = 1e-30
        return pch  # (N,R)

    def _apply_distribution(self, betas_random, draws):
        for k, dist in enumerate(self.rvdist):
            if dist == 'ln':
                betas_random[:, k, :] = dev.np.exp(betas_random[:, k, :])
        return betas_random

    def _balance_panels(self, X, y, mixby):
        _, J, K = X.shape
        _, p_obs = np.unique(mixby, return_counts=True)
        p_obs = (p_obs/J).astype(int)
        N = len(p_obs)  # This is the new N after accounting for panels
        P = np.max(p_obs)  # Panel length for all records

        if not np.all(p_obs[0] == p_obs):  # Balancing needed
            y = y.reshape(X.shape[0], J, 1)
            Xbal, ybal = np.zeros((N*P, J, K)), np.zeros((N*P, J, 1))
            panel_info = np.zeros((N, P))
            cum_p = 0  # Cumulative sum of n_obs at every iteration
            for n, p in enumerate(p_obs):
                # Copy data from original to balanced version
                Xbal[n*P:n*P + p, :, :] = X[cum_p:cum_p + p, :, :]
                ybal[n*P:n*P + p, :, :] = y[cum_p:cum_p + p, :, :]
                panel_info[n, :p] = np.ones(p)
                cum_p += p

        else:  # No balancing needed
            Xbal, ybal = X, y
            panel_info = np.ones((N, P))

        return Xbal, ybal, panel_info

    def _compute_derivatives(self, betas, draws):
        _, betas_random = self._transform_betas(betas, draws)
        Kr = self.rvidx.sum()  # Number of random coeff
        N, R = draws.shape[0], draws.shape[2]
        der = dev.np.ones((N, Kr, R))
        for k, dist in enumerate(self.rvdist):
            if dist == 'ln':
                der[:, k, :] = betas_random[:, k, :]
        return der

    def _transform_betas(self, betas, draws):
        # Extract coeffiecients from betas array
        Kr = self.rvidx.sum()   # Number of random coeff
        Kf = len(betas) - 2*Kr  # Number of fixed coeff
        betas_fixed = betas[0:Kf]  # First Kf positions
        br_mean, br_sd = betas[Kf:Kf+Kr], betas[Kf+Kr:]  # Remaining positions
        # Compute: betas = mean + sd*draws
        betas_random = br_mean[None, :, None] + draws*br_sd[None, :, None]
        betas_random = self._apply_distribution(betas_random, draws)
        return betas_fixed, betas_random

    def _generate_draws(self, sample_size, n_draws, halton=True):
        if halton:
            draws = self._get_halton_draws(sample_size, n_draws,
                                           len(self.rvdist))
        else:
            draws = self._get_random_normal_draws(sample_size, n_draws,
                                                  len(self.rvdist))
        return draws

    def _get_random_normal_draws(self, sample_size, n_draws, n_vars):
        draws = [np.random.normal(0, 1, size=(sample_size, n_draws))
                 for _ in range(n_vars)]
        draws = np.stack(draws, axis=1)
        return draws

    def _get_halton_draws(self, sample_size, n_draws, n_vars, shuffled=False):
        primes = [2, 3, 5, 7, 11, 13, 17, 19, 23, 29, 31, 37, 41, 43, 47,
                  53, 59, 61, 71, 73, 79, 83, 89, 97, 101, 103, 107,
                  109, 113, 127, 131, 137, 139, 149, 151, 157, 163, 167,
                  173, 179, 181, 191, 193, 197, 199]

        def halton_seq(length, prime=3, shuffled=False, drop=100):
            h = np.array([.0])
            t = 0
            while len(h) < length + drop:
                t += 1
                h = np.append(h, np.tile(h, prime-1) +
                              np.repeat(np.arange(1, prime)/prime**t, len(h)))
            seq = h[drop:length+drop]
            if shuffled:
                np.random.shuffle(seq)
            return seq

        draws = [halton_seq(sample_size*n_draws, prime=primes[i % len(primes)],
                            shuffled=shuffled).reshape(sample_size, n_draws)
                 for i in range(n_vars)]
        draws = np.stack(draws, axis=1)
        draws = scipy.stats.norm.ppf(draws)
        return draws
