
from functools import partial

import numpy as np
from scipy import stats

import statsmodels.api as sm
from statsmodels.base.model import GenericLikelihoodModel
from statsmodels.tools.numdiff import approx_fprime, approx_hess

data = sm.datasets.spector.load()
data.exog = sm.add_constant(data.exog, prepend=False)
# in this dir

probit_mod = sm.Probit(data.endog, data.exog)
probit_res = probit_mod.fit()
loglike = probit_mod.loglike
score = probit_mod.score
mod = GenericLikelihoodModel(data.endog, data.exog*2, loglike, score)
res = mod.fit(method="nm", maxiter = 500)

def probitloglike(params, endog, exog):
    """
    Log likelihood for the probit
    """
    q = 2*endog - 1
    X = exog
    return np.add.reduce(stats.norm.logcdf(q*np.dot(X,params)))


model_loglike = partial(probitloglike, endog=data.endog, exog=data.exog)
mod = GenericLikelihoodModel(data.endog, data.exog, loglike=model_loglike)
res = mod.fit(method="nm", maxiter=500)
print(res)


np.allclose(res.params, probit_res.params, rtol=1e-4)
print(res.params, probit_res.params)

#datal = sm.datasets.longley.load()
datal = sm.datasets.ccard.load()
datal.exog = sm.add_constant(datal.exog, prepend=False)
# Instance of GenericLikelihood model does not work directly, because loglike
# cannot get access to data in self.endog, self.exog

nobs = 5000
rvs = np.random.randn(nobs,6)
datal.exog = rvs[:,:-1]
datal.exog = sm.add_constant(datal.exog, prepend=False)
datal.endog = 1 + rvs.sum(1)

show_error = False
show_error2 = 1#False
if show_error:
    def loglike_norm_xb(self, params):
        beta = params[:-1]
        sigma = params[-1]
        xb = np.dot(self.exog, beta)
        return stats.norm.logpdf(self.endog, loc=xb, scale=sigma)

    mod_norm = GenericLikelihoodModel(datal.endog, datal.exog, loglike_norm_xb)
    res_norm = mod_norm.fit(method="nm", maxiter = 500)

    print(res_norm.params)

if show_error2:
    def loglike_norm_xb(params, endog, exog):
        beta = params[:-1]
        sigma = params[-1]
        #print exog.shape, beta.shape
        xb = np.dot(exog, beta)
        #print xb.shape, stats.norm.logpdf(endog, loc=xb, scale=sigma).shape
        return stats.norm.logpdf(endog, loc=xb, scale=sigma).sum()

    model_loglike3 = partial(loglike_norm_xb,
                             endog=datal.endog, exog=datal.exog)
    mod_norm = GenericLikelihoodModel(datal.endog, datal.exog, model_loglike3)
    res_norm = mod_norm.fit(start_params=np.ones(datal.exog.shape[1]+1),
                            method="nm", maxiter = 5000)

    print(res_norm.params)

class MygMLE(GenericLikelihoodModel):
    # just for testing
    def loglike(self, params):
        beta = params[:-1]
        sigma = params[-1]
        xb = np.dot(self.exog, beta)
        return stats.norm.logpdf(self.endog, loc=xb, scale=sigma).sum()

    def loglikeobs(self, params):
        beta = params[:-1]
        sigma = params[-1]
        xb = np.dot(self.exog, beta)
        return stats.norm.logpdf(self.endog, loc=xb, scale=sigma)

mod_norm2 = MygMLE(datal.endog, datal.exog)
#res_norm = mod_norm.fit(start_params=np.ones(datal.exog.shape[1]+1), method="nm", maxiter = 500)
res_norm2 = mod_norm2.fit(start_params=[1.]*datal.exog.shape[1]+[1], method="nm", maxiter = 500)
np.allclose(res_norm.params, res_norm2.params)
print(res_norm2.params)

res2 = sm.OLS(datal.endog, datal.exog).fit()
start_params = np.hstack((res2.params, np.sqrt(res2.mse_resid)))
res_norm3 = mod_norm2.fit(start_params=start_params, method="nm", maxiter = 500,
                          retall=0)
print(start_params)
print(res_norm3.params)
print(res2.bse)
print(res_norm3.bse)
print('llf', res2.llf, res_norm3.llf)

bse = np.sqrt(np.diag(np.linalg.inv(res_norm3.model.hessian(res_norm3.params))))
res_norm3.model.score(res_norm3.params)

#fprime in fit option cannot be overwritten, set to None, when score is defined
# exception is fixed, but I do not think score was supposed to be called

res_bfgs = mod_norm2.fit(start_params=start_params, method="bfgs", fprime=None,
                         maxiter=500, retall=0)

hb=-approx_hess(res_norm3.params, mod_norm2.loglike, epsilon=-1e-4)
hf=-approx_hess(res_norm3.params, mod_norm2.loglike, epsilon=1e-4)
hh = (hf+hb)/2.
print(np.linalg.eigh(hh))

grad = -approx_fprime(res_norm3.params, mod_norm2.loglike, epsilon=-1e-4)
print(grad)
gradb = -approx_fprime(res_norm3.params, mod_norm2.loglike, epsilon=-1e-4)
gradf = -approx_fprime(res_norm3.params, mod_norm2.loglike, epsilon=1e-4)
print((gradb+gradf)/2.)

print(res_norm3.model.score(res_norm3.params))
print(res_norm3.model.score(start_params))
mod_norm2.loglike(start_params/2.)
print(np.linalg.inv(-1*mod_norm2.hessian(res_norm3.params)))
print(np.sqrt(np.diag(res_bfgs.cov_params())))
print(res_norm3.bse)

print("MLE - OLS parameter estimates")
print(res_norm3.params[:-1] - res2.params)
print("bse diff in percent")
print((res_norm3.bse[:-1] / res2.bse)*100. - 100)

'''
Optimization terminated successfully.
         Current function value: 12.818804
         Iterations 6
Optimization terminated successfully.
         Current function value: 12.818804
         Iterations: 439
         Function evaluations: 735
Optimization terminated successfully.
         Current function value: 12.818804
         Iterations: 439
         Function evaluations: 735
<statsmodels.model.LikelihoodModelResults object at 0x02131290>
[ 1.6258006   0.05172931  1.42632252 -7.45229732] [ 1.62581004  0.05172895  1.42633234 -7.45231965]
Warning: Maximum number of function evaluations has been exceeded.
[  -1.18109149  246.94438535  -16.21235536   24.05282629 -324.80867176
  274.07378453]
Warning: Maximum number of iterations has been exceeded
[  17.57107    -149.87528787   19.89079376  -72.49810777  -50.06067953
  306.14170418]
Optimization terminated successfully.
         Current function value: 506.488765
         Iterations: 339
         Function evaluations: 550
[  -3.08181404  234.34702702  -14.99684418   27.94090839 -237.1465136
  284.75079529]
[  -3.08181304  234.34701361  -14.99684381   27.94088692 -237.14649571
  274.6857294 ]
[   5.51471653   80.36595035    7.46933695   82.92232357  199.35166485]
llf -506.488764864 -506.488764864
Optimization terminated successfully.
         Current function value: 506.488765
         Iterations: 9
         Function evaluations: 13
         Gradient evaluations: 13
(array([  2.41772580e-05,   1.62492628e-04,   2.79438138e-04,
         1.90996240e-03,   2.07117946e-01,   1.28747174e+00]), array([[  1.52225754e-02,   2.01838216e-02,   6.90127235e-02,
         -2.57002471e-04,  -5.25941060e-01,  -8.47339404e-01],
       [  2.39797491e-01,  -2.32325602e-01,  -9.36235262e-01,
          3.02434938e-03,   3.95614029e-02,  -1.02035585e-01],
       [ -2.11381471e-02,   3.01074776e-02,   7.97208277e-02,
         -2.94955832e-04,   8.49402362e-01,  -5.20391053e-01],
       [ -1.55821981e-01,  -9.66926643e-01,   2.01517298e-01,
          1.52397702e-03,   4.13805882e-03,  -1.19878714e-02],
       [ -9.57881586e-01,   9.87911166e-02,  -2.67819451e-01,
          1.55192932e-03,  -1.78717579e-02,  -2.55757014e-02],
       [ -9.96486655e-04,  -2.03697290e-03,  -2.98130314e-03,
         -9.99992985e-01,  -1.71500426e-05,   4.70854949e-06]]))
[[ -4.91007768e-05  -7.28732630e-07  -2.51941401e-05  -2.50111043e-08
   -4.77484718e-08  -9.72022463e-08]]
[[ -1.64845915e-08  -2.87059265e-08  -2.88764568e-07  -6.82121026e-09