import numpy as np
from scipy import stats
import statsmodels.api as sm
from statsmodels.base.model import GenericLikelihoodModel

"""
Explanation: Maximum Likelihood Estimation (Generic models)
This tutorial explains how to quickly implement new maximum likelihood models in statsmodels. We give two examples: 

Probit model for binary dependent variables
Negative binomial model for count data

The GenericLikelihoodModel class eases the process by providing tools such as automatic numeric differentiation and a unified interface to scipy optimization functions. Using statsmodels, users can fit new MLE models simply by "plugging-in" a log-likelihood function. 
Example 1: Probit model
End of explanation
"""


data = sm.datasets.spector.load_pandas()
exog = data.exog
endog = data.endog
print(sm.datasets.spector.NOTE)
print(data.exog.head())

"""
Explanation: The Spector dataset is distributed with statsmodels. You can access a vector of values for the dependent variable (endog) and a matrix of regressors (exog) like this:
End of explanation
"""


exog = sm.add_constant(exog, prepend=True)

"""
Explanation: Them, we add a constant to the matrix of regressors:
End of explanation
"""


class MyProbit(GenericLikelihoodModel):
    def loglike(self, params):
        exog = self.exog
        endog = self.endog
        q = 2 * endog - 1
        return stats.norm.logcdf(q*np.dot(exog, params)).sum()

"""
Explanation: To create your own Likelihood Model, you simply need to overwrite the loglike method.
End of explanation
"""


sm_probit_manual = MyProbit(endog, exog).fit()
print(sm_probit_manual.summary())

"""
Explanation: Estimate the model and print a summary:
End of explanation
"""


sm_probit_canned = sm.Probit(endog, exog).fit()

print(sm_probit_canned.params)
print(sm_probit_manual.params)

print(sm_probit_canned.cov_params())
print(sm_probit_manual.cov_params())

"""
Explanation: Compare your Probit implementation to statsmodels' "canned" implementation:
End of explanation
"""


import numpy as np
from scipy.stats import nbinom

def _ll_nb2(y, X, beta, alph):
    mu = np.exp(np.dot(X, beta))
    size = 1/alph
    prob = size/(size+mu)
    ll = nbinom.logpmf(y, size, prob)
    return ll

"""
Explanation: Notice that the GenericMaximumLikelihood class provides automatic differentiation, so we did not have to provide Hessian or Score functions in order to calculate the covariance estimates.
Example 2: Negative Binomial Regression for Count Data
Consider a negative binomial regression model for count data with
log-likelihood (type NB-2) function expressed as:
$$
    \mathcal{L}(\beta_j; y, \alpha) = \sum_{i=1}^n y_i ln 
    \left ( \frac{\alpha exp(X_i'\beta)}{1+\alpha exp(X_i'\beta)} \right ) -
    \frac{1}{\alpha} ln(1+\alpha exp(X_i'\beta)) + ln \Gamma (y_i + 1/\alpha) - ln \Gamma (y_i+1) - ln \Gamma (1/\alpha)
$$
with a matrix of regressors $X$, a vector of coefficients $\beta$,
and the negative binomial heterogeneity parameter $\alpha$. 
Using the nbinom distribution from scipy, we can write this likelihood
simply as:
End of explanation
"""


from statsmodels.base.model import GenericLikelihoodModel

class NBin(GenericLikelihoodModel):
    def __init__(self, endog, exog, **kwds):
        super(NBin, self).__init__(endog, exog, **kwds)
        
    def nloglikeobs(self, params):
        alph = params[-1]
        beta = params[:-1]
        ll = _ll_nb2(self.endog, self.exog, beta, alph)
        return -ll 
    
    def fit(self, start_params=None, maxiter=10000, maxfun=5000, **kwds):
        # we have one additional parameter and we need to add it for summary
        self.exog_names.append('alpha')
        if start_params == None:
            # Reasonable starting values
            start_params = np.append(np.zeros(self.exog.shape[1]), .5)
            # intercept
            start_params[-2] = np.log(self.endog.mean())
        return super(NBin, self).fit(start_params=start_params, 
                                     maxiter=maxiter, maxfun=maxfun, 
                                     **kwds) 

"""
Explanation: New Model Class
We create a new model class which inherits from GenericLikelihoodModel:
End of explanation
"""


import statsmodels.api as sm

medpar = sm.datasets.get_rdataset("medpar", "COUNT", cache=True).data

medpar.head()

"""
Explanation: Two important things to notice: 

nloglikeobs: This function should return one evaluation of the negative log-likelihood function per observation in your dataset (i.e. rows of the endog/X matrix). 
start_params: A one-dimensional array of starting values needs to be provided. The size of this array determines the number of parameters that will be used in optimization.

That's it! You're done!
Usage Example
The Medpar
dataset is hosted in CSV format at the Rdatasets repository. We use the read_csv
function from the Pandas library to load the data
in memory. We then print the first few columns:
End of explanation
"""


y = medpar.los
X = medpar[["type2", "type3", "hmo", "white"]].copy()
X["constant"] = 1

"""
Explanation: The model we are interested in has a vector of non-negative integers as
dependent variable (los), and 5 regressors: Intercept, type2,
type3, hmo, white.
For estimation, we need to create two variables to hold our regressors and the outcome variable. These can be ndarrays or pandas objects.
End of explanation
"""


mod = NBin(y, X)
res = mod.fit()

"""
Explanation: Then, we fit the model and extract some information:
End of explanation
"""


print('Parameters: ', res.params)
print('Standard errors: ', res.bse)
print('P-values: ', res.pvalues)
print('AIC: ', res.aic)

"""
Explanation: Extract parameter estimates, standard errors, p-values, AIC, etc.:
End of explanation
"""


print(res.summary())

"""
Explanation: As usual, you can obtain a full list of available information by typing
dir(res).
We can also look at the summary of the estimation results.
End of explanation
"""


res_nbin = sm.NegativeBinomial(y, X).fit(disp=0)
print(res_nbin.summary())

print(res_nbin.params)

print(res_nbin.bse)

"""
Explanation: Testing
We can check the results by using the statsmodels implementation of the Negative Binomial model, which uses the analytic score function and Hessian.
End of explanation
"""


