import numpy as np
from statsmodels.discrete.conditional_models import (
      ConditionalLogit, ConditionalPoisson)
from statsmodels.tools.numdiff import approx_fprime
from numpy.testing import assert_allclose
import pandas as pd


def test_logit_1d():

    y = np.r_[0, 1, 0, 1, 0, 1, 0, 1, 1, 1]
    g = np.r_[0, 0, 0, 1, 1, 1, 2, 2, 2, 2]

    x = np.r_[0, 1, 0, 0, 1, 1, 0, 0, 1, 0]
    x = x[:, None]

    model = ConditionalLogit(y, x, groups=g)

    # Check the gradient for the denominator of the partial likelihood
    for x in -1, 0, 1, 2:
        params = np.r_[x, ]
        _, grad = model._denom_grad(0, params)
        ngrad = approx_fprime(params, lambda x: model._denom(0, x))
        assert_allclose(grad, ngrad)

    # Check the gradient for the loglikelihood
    for x in -1, 0, 1, 2:
        grad = approx_fprime(np.r_[x, ], model.loglike)
        score = model.score(np.r_[x, ])
        assert_allclose(grad, score, rtol=1e-4)

    result = model.fit()

    # From Stata
    assert_allclose(result.params, np.r_[0.9272407], rtol=1e-5)
    assert_allclose(result.bse, np.r_[1.295155], rtol=1e-5)


def test_logit_2d():

    y = np.r_[0, 1, 0, 1, 0, 1, 0, 1, 1, 1]
    g = np.r_[0, 0, 0, 1, 1, 1, 2, 2, 2, 2]

    x1 = np.r_[0, 1, 0, 0, 1, 1, 0, 0, 1, 0]
    x2 = np.r_[0, 0, 1, 0, 0, 1, 0, 1, 1, 1]
    x = np.empty((10, 2))
    x[:, 0] = x1
    x[:, 1] = x2

    model = ConditionalLogit(y, x, groups=g)

    # Check the gradient for the denominator of the partial likelihood
    for x in -1, 0, 1, 2:
        params = np.r_[x, -1.5*x]
        _, grad = model._denom_grad(0, params)
        ngrad = approx_fprime(params, lambda x: model._denom(0, x))
        assert_allclose(grad, ngrad, rtol=1e-5)

    # Check the gradient for the loglikelihood
    for x in -1, 0, 1, 2:
        params = np.r_[-0.5*x, 0.5*x]
        grad = approx_fprime(params, model.loglike)
        score = model.score(params)
        assert_allclose(grad, score, rtol=1e-4)

    result = model.fit()

    # From Stata
    assert_allclose(result.params, np.r_[1.011074, 1.236758], rtol=1e-3)
    assert_allclose(result.bse, np.r_[1.420784, 1.361738], rtol=1e-5)

    result.summary()


def test_formula():

    for j in 0, 1:

        np.random.seed(34234)
        n = 200
        y = np.random.randint(0, 2, size=n)
        x1 = np.random.normal(size=n)
        x2 = np.random.normal(size=n)
        g = np.random.randint(0, 25, size=n)

        x = np.hstack((x1[:, None], x2[:, None]))
        if j == 0:
            model1 = ConditionalLogit(y, x, groups=g)
        else:
            model1 = ConditionalPoisson(y, x, groups=g)
        result1 = model1.fit()

        df = pd.DataFrame({"y": y, "x1": x1, "x2": x2, "g": g})
        if j == 0:
            model2 = ConditionalLogit.from_formula(
                        "y ~ 0 + x1 + x2", groups="g", data=df)
        else:
            model2 = ConditionalPoisson.from_formula(
                        "y ~ 0 + x1 + x2", groups="g", data=df)
        result2 = model2.fit()

        assert_allclose(result1.params, result2.params, rtol=1e-5)
        assert_allclose(result1.bse, result2.bse, rtol=1e-5)
        assert_allclose(result1.cov_params(), result2.cov_params(), rtol=1e-5)
        assert_allclose(result1.tvalues, result2.tvalues, rtol=1e-5)


def test_poisson_1d():

    y = np.r_[3, 1, 1, 4, 5, 2, 0, 1, 6, 2]
    g = np.r_[0, 0, 0, 0, 1, 1, 1, 1, 1, 1]

    x = np.r_[0, 1, 0, 0, 1, 1, 0, 0, 1, 0]
    x = x[:, None]

    model = ConditionalPoisson(y, x, groups=g)

    # Check the gradient for the loglikelihood
    for x in -1, 0, 1, 2:
        grad = approx_fprime(np.r_[x, ], model.loglike)
        score = model.score(np.r_[x, ])
        assert_allclose(grad, score, rtol=1e-4)

    result = model.fit()

    # From Stata
    assert_allclose(result.params, np.r_[0.6466272], rtol=1e-4)
    assert_allclose(result.bse, np.r_[0.4170918], rtol=1e-5)


def test_poisson_2d():

    y = np.r_[3, 1, 4, 8, 2, 5, 4, 7, 2, 6]
    g = np.r_[0, 0, 0, 1, 1, 1, 2, 2, 2, 2]

    x1 = np.r_[0, 1, 0, 0, 1, 1, 0, 0, 1, 0]
    x2 = np.r_[2, 1, 0, 0, 1, 2, 3, 2, 0, 1]
    x = np.empty((10, 2))
    x[:, 0] = x1
    x[:, 1] = x2

    model = ConditionalPoisson(y, x, groups=g)

    # Check the gradient for the loglikelihood
    for x in -1, 0, 1, 2:
        params = np.r_[-0.5*x, 0.5*x]
        grad = approx_fprime(params, model.loglike)
        score = model.score(params)
        assert_allclose(grad, score, rtol=1e-4)

    result = model.fit()

    # From Stata
    assert_allclose(result.params, np.r_[-.9478957, -.0134279], rtol=1e-3)
    assert_allclose(result.bse, np.r_[.3874942, .1686712], rtol=1e-5)

    result.summary()


def test_lasso_logistic():

    np.random.seed(3423948)

    n = 200
    groups = np.arange(10)
    groups = np.kron(groups, np.ones(n // 10))
    group_effects = np.random.normal(size=10)
    group_effects = np.kron(group_effects, np.ones(n // 10))

    x = np.random.normal(size=(n, 4))
    params = np.r_[0, 0, 1, 0]
    lin_pred = np.dot(x, params) + group_effects

    mean = 1 / (1 + np.exp(-lin_pred))
    y = (np.random.uniform(size=n) < mean).astype(np.int)

    model0 = ConditionalLogit(y, x, groups=groups)
    result0 = model0.fit()

    # Should be the same as model0
    model1 = ConditionalLogit(y, x, groups=groups)
    result1 = model1.fit_regularized(L1_wt=0, alpha=0)

    assert_allclose(result0.params, result1.params, rtol=1e-3)

    model2 = ConditionalLogit(y, x, groups=groups)
    result2 = model2.fit_regularized(L1_wt=1, alpha=0.05)

    # Rxegression test
    assert_allclose(result2.params, np.r_[0, 0, 0.55235152, 0], rtol=1e-4)

    # Test with formula
    df = pd.DataFrame({"y": y, "x1": x[:, 0], "x2": x[:, 1], "x3": x[:, 2],
                       "x4": x[:, 3], "groups": groups})
    fml = "y ~ 0 + x1 + x2 + x3 + x4"
    model3 = ConditionalLogit.from_formula(fml, groups="groups", data=df)
    result3 = model3.fit_regularized(L1_wt=1, alpha=0.05)
    assert_allclose(result2.params, result3.params)


def test_lasso_poisson():

    np.random.seed(342394)

    n = 200
    groups = np.arange(10)
    groups = np.kron(groups, np.ones(n // 10))
    group_effects = np.random.normal(size=10)
    group_effects = np.kron(group_effects, np.ones(n // 10))

    x = np.random.normal(size=(n, 4))
    params = np.r_[0, 0, 1, 0]
    lin_pred = np.dot(x, params) + group_effects

    mean = np.exp(lin_pred)
    y = np.random.poisson(mean)

    model0 = ConditionalPoisson(y, x, groups=groups)
    result0 = model0.fit()

    # Should be the same as model0
    model1 = ConditionalPoisson(y, x, groups=groups)
    result1 = model1.fit_regularized(L1_wt=0, alpha=0)

    assert_allclose(result0.params, result1.params, rtol=1e-3)

    model2 = ConditionalPoisson(y, x, groups=groups)
    result2 = model2.fit_regularized(L1_wt=1, alpha=0.2)

    # Regression test
    assert_allclose(result2.params, np.r_[0, 0, 0.91697508, 0], rtol=1e-4)

    # Test with formula
    df = pd.DataFrame({"y": y, "x1": x[:, 0], "x2": x[:, 1], "x3": x[:, 2],
                       "x4": x[:, 3], "groups": groups})
    fml = "y ~ 0 + x1 + x2 + x3 + x4"
    model3 = ConditionalPoisson.from_formula(fml, groups="groups", data=df)
    result3 = model3.fit_regularized(L1_wt=1, alpha=0.2)
    assert_allclose(result2.params, result3.params)

[((17, 12, 17, 44), 'statsmodels.discrete.conditional_models.ConditionalLogit', 'ConditionalLogit', (), '', False, 'from statsmodels.discrete.conditional_models import ConditionalLogit, ConditionalPoisson\n'), ((35, 4, 35, 63), 'numpy.testing.assert_allclose', 'assert_allclose', (), '', False, 'from numpy.testing import assert_allclose\n'), ((36, 4, 36, 59), 'numpy.testing.assert_allclose', 'assert_allclose', (), '', False, 'from numpy.testing import assert_allclose\n'), ((46, 8, 46, 25), 'numpy.empty', 'np.empty', ({(46, 17, 46, 24): '(10, 2)'}, {}), '((10, 2))', True, 'import numpy as np\n'), ((50, 12, 50, 44), 'statsmodels.discrete.conditional_models.ConditionalLogit', 'ConditionalLogit', (), '', False, 'from statsmodels.discrete.conditional_models import ConditionalLogit, ConditionalPoisson\n'), ((69, 4, 69, 72), 'numpy.testing.assert_allclose', 'assert_allclose', (), '', False, 'from numpy.testing import assert_allclose\n'), ((70, 4, 70, 69), 'numpy.testing.assert_allclose', 'assert_allclose', (), '', False, 'from numpy.testing import assert_allclose\n'), ((116, 12, 116, 46), 'statsmodels.discrete.conditional_models.ConditionalPoisson', 'ConditionalPoisson', (), '', False, 'from statsmodels.discrete.conditional_models import ConditionalLogit, ConditionalPoisson\n'), ((127, 4, 127, 63), 'numpy.testing.assert_allclose', 'assert_allclose', (), '', False, 'from numpy.testing import assert_allclose\n'), ((128, 4, 128, 60), 'numpy.testing.assert_allclose', 'assert_allclose', (), '', False, 'from numpy.testing import assert_allclose\n'), ((138, 8, 138, 25), 'numpy.empty', 'np.empty', ({(138, 17, 138, 24): '(10, 2)'}, {}), '((10, 2))', True, 'import numpy as np\n'), ((142, 12, 142, 46), 'statsmodels.discrete.conditional_models.ConditionalPoisson', 'ConditionalPoisson', (), '', False, 'from statsmodels.discrete.conditional_models import ConditionalLogit, ConditionalPoisson\n'), ((154, 4, 154, 74), 'numpy.testing.assert_allclose', 'assert_allclose', (), '', False, 'from numpy.testing import assert_allclose\n'), ((155, 4, 155, 69), 'numpy.testing.assert_allclose', 'assert_allclose', (), '', False, 'from numpy.testing import assert_allclose\n'), ((162, 4, 162, 27), 'numpy.random.seed', 'np.random.seed', ({(162, 19, 162, 26): '(3423948)'}, {}), '(3423948)', True, 'import numpy as np\n'), ((165, 13, 165, 26), 'numpy.arange', 'np.arange', ({(165, 23, 165, 25): '10'}, {}), '(10)', True, 'import numpy as np\n'), ((167, 20, 167, 45), 'numpy.random.normal', 'np.random.normal', (), '', True, 'import numpy as np\n'), ((170, 8, 170, 37), 'numpy.random.normal', 'np.random.normal', (), '', True, 'import numpy as np\n'), ((177, 13, 177, 50), 'statsmodels.discrete.conditional_models.ConditionalLogit', 'ConditionalLogit', (), '', False, 'from statsmodels.discrete.conditional_models import ConditionalLogit, ConditionalPoisson\n'), ((181, 13, 181, 50), 'statsmodels.discrete.conditional_models.ConditionalLogit', 'ConditionalLogit', (), '', False, 'from statsmodels.discrete.conditional_models import ConditionalLogit, ConditionalPoisson\n'), ((184, 4, 184, 62), 'numpy.testing.assert_allclose', 'assert_allclose', (), '', False, 'from numpy.testing import assert_allclose\n'), ((186, 13, 186, 50), 'statsmodels.discrete.conditional_models.ConditionalLogit', 'ConditionalLogit', (), '', False, 'from statsmodels.discrete.conditional_models import ConditionalLogit, ConditionalPoisson\n'), ((190, 4, 190, 74), 'numpy.testing.assert_allclose', 'assert_allclose', (), '', False, 'from numpy.testing import assert_allclose\n'), ((193, 9, 194, 56), 'pandas.DataFrame', 'pd.DataFrame', ({(193, 22, 194, 55): "{'y': y, 'x1': x[:, (0)], 'x2': x[:, (1)], 'x3': x[:, (2)], 'x4': x[:, (3)],\n    'groups': groups}"}, {}), "({'y': y, 'x1': x[:, (0)], 'x2': x[:, (1)], 'x3': x[:, (2)],\n    'x4': x[:, (3)], 'groups': groups})", True, 'import pandas as pd\n'), ((196, 13, 196, 73), 'statsmodels.discrete.conditional_models.ConditionalLogit.from_formula', 'ConditionalLogit.from_formula', (), '', False, 'from statsmodels.discrete.conditional_models import ConditionalLogit, ConditionalPoisson\n'), ((198, 4, 198, 51), 'numpy.testing.assert_allclose', 'assert_allclose', ({(198, 20, 198, 34): 'result2.params', (198, 36, 198, 50): 'result3.params'}, {}), '(result2.params, result3.params)', False, 'from numpy.testing import assert_allclose\n'), ((203, 4, 203, 26), 'numpy.random.seed', 'np.random.seed', ({(203, 19, 203, 25): '(342394)'}, {}), '(342394)', True, 'import numpy as np\n'), ((206, 13, 206, 26), 'numpy.arange', 'np.arange', ({(206, 23, 206, 25): '10'}, {}), '(10)', True, 'import numpy as np\n'), ((208, 20, 208, 45), 'numpy.random.normal', 'np.random.normal', (), '', True, 'import numpy as np\n'), ((211, 8, 211, 37), 'numpy.random.normal', 'np.random.normal', (), '', True, 'import numpy as np\n'), ((215, 11, 215, 27), 'numpy.exp', 'np.exp', ({(215, 18, 215, 26): 'lin_pred'}, {}), '(lin_pred)', True, 'import numpy as np\n'), ((216, 8, 216, 31), 'numpy.random.poisson', 'np.random.poisson', ({(216, 26, 216, 30): 'mean'}, {}), '(mean)', True, 'import numpy as np\n'), ((218, 13, 218, 52), 'statsmodels.discrete.conditional_models.ConditionalPoisson', 'ConditionalPoisson', (), '', False, 'from statsmodels.discrete.conditional_models import ConditionalLogit, ConditionalPoisson\n'), ((222, 13, 222, 52), 'statsmodels.discrete.conditional_models.ConditionalPoisson', 'ConditionalPoisson', (), '', False, 'from statsmodels.discrete.conditional_models import ConditionalLogit, ConditionalPoisson\n'), ((225, 4, 225, 62), 'numpy.testing.assert_allclose', 'assert_allclose', (), '', False, 'from numpy.testing import assert_allclose\n'), ((227, 13, 227, 52), 'statsmodels.discrete.conditional_models.ConditionalPoisson', 'ConditionalPoisson', (), '', False, 'from statsmodels.discrete.conditional_models import ConditionalLogit, ConditionalPoisson\n'), ((231, 4, 231, 74), 'numpy.testing.assert_allclose', 'assert_allclose', (), '', False, 'from numpy.testing import assert_allclose\n'), ((234, 9, 235, 56), 'pandas.DataFrame', 'pd.DataFrame', ({(234, 22, 235, 55): "{'y': y, 'x1': x[:, (0)], 'x2': x[:, (1)], 'x3': x[:, (2)], 'x4': x[:, (3)],\n    'groups': groups}"}, {}), "({'y': y, 'x1': x[:, (0)], 'x2': x[:, (1)], 'x3': x[:, (2)],\n    'x4': x[:, (3)], 'groups': groups})", True, 'import pandas as pd\n'), ((237, 13, 237, 75), 'statsmodels.discrete.conditional_models.ConditionalPoisson.from_formula', 'ConditionalPoisson.from_formula', (), '', False, 'from statsmodels.discrete.conditional_models import ConditionalLogit, ConditionalPoisson\n'), ((239, 4, 239, 51), 'numpy.testing.assert_allclose', 'assert_allclose', ({(239, 20, 239, 34): 'result2.params', (239, 36, 239, 50): 'result3.params'}, {}), '(result2.params, result3.params)', False, 'from numpy.testing import assert_allclose\n'), ((24, 8, 24, 36), 'numpy.testing.assert_allclose', 'assert_allclose', ({(24, 24, 24, 28): 'grad', (24, 30, 24, 35): 'ngrad'}, {}), '(grad, ngrad)', False, 'from numpy.testing import assert_allclose\n'), ((28, 15, 28, 55), 'statsmodels.tools.numdiff.approx_fprime', 'approx_fprime', ({(28, 29, 28, 39): 'np.r_[x,]', (28, 41, 28, 54): 'model.loglike'}, {}), '(np.r_[x,], model.loglike)', False, 'from statsmodels.tools.numdiff import approx_fprime\n'), ((30, 8, 30, 47), 'numpy.testing.assert_allclose', 'assert_allclose', (), '', False, 'from numpy.testing import assert_allclose\n'), ((57, 8, 57, 47), 'numpy.testing.assert_allclose', 'assert_allclose', (), '', False, 'from numpy.testing import assert_allclose\n'), ((62, 15, 62, 51), 'statsmodels.tools.numdiff.approx_fprime', 'approx_fprime', ({(62, 29, 62, 35): 'params', (62, 37, 62, 50): 'model.loglike'}, {}), '(params, model.loglike)', False, 'from statsmodels.tools.numdiff import approx_fprime\n'), ((64, 8, 64, 47), 'numpy.testing.assert_allclose', 'assert_allclose', (), '', False, 'from numpy.testing import assert_allclose\n'), ((79, 8, 79, 29), 'numpy.random.seed', 'np.random.seed', ({(79, 23, 79, 28): '(34234)'}, {}), '(34234)', True, 'import numpy as np\n'), ((81, 12, 81, 43), 'numpy.random.randint', 'np.random.randint', (), '', True, 'import numpy as np\n'), ((82, 13, 82, 37), 'numpy.random.normal', 'np.random.normal', (), '', True, 'import numpy as np\n'), ((83, 13, 83, 37), 'numpy.random.normal', 'np.random.normal', (), '', True, 'import numpy as np\n'), ((84, 12, 84, 44), 'numpy.random.randint', 'np.random.randint', (), '', True, 'import numpy as np\n'), ((86, 12, 86, 49), 'numpy.hstack', 'np.hstack', ({(86, 22, 86, 48): '(x1[:, (None)], x2[:, (None)])'}, {}), '((x1[:, (None)], x2[:, (None)]))', True, 'import numpy as np\n'), ((93, 13, 93, 63), 'pandas.DataFrame', 'pd.DataFrame', ({(93, 26, 93, 62): "{'y': y, 'x1': x1, 'x2': x2, 'g': g}"}, {}), "({'y': y, 'x1': x1, 'x2': x2, 'g': g})", True, 'import pandas as pd\n'), ((102, 8, 102, 66), 'numpy.testing.assert_allclose', 'assert_allclose', (), '', False, 'from numpy.testing import assert_allclose\n'), ((103, 8, 103, 60), 'numpy.testing.assert_allclose', 'assert_allclose', (), '', False, 'from numpy.testing import assert_allclose\n'), ((105, 8, 105, 68), 'numpy.testing.assert_allclose', 'assert_allclose', (), '', False, 'from numpy.testing import assert_allclose\n'), ((120, 15, 120, 55), 'statsmodels.tools.numdiff.approx_fprime', 'approx_fprime', ({(120, 29, 120, 39): 'np.r_[x,]', (120, 41, 120, 54): 'model.loglike'}, {}), '(np.r_[x,], model.loglike)', False, 'from statsmodels.tools.numdiff import approx_fprime\n'), ((122, 8, 122, 47), 'numpy.testing.assert_allclose', 'assert_allclose', (), '', False, 'from numpy.testing import assert_allclose\n'), ((147, 15, 147, 51), 'statsmodels.tools.numdiff.approx_fprime', 'approx_fprime', ({(147, 29, 147, 35): 'params', (147, 37, 147, 50): 'model.loglike'}, {}), '(params, model.loglike)', False, 'from statsmodels.tools.numdiff import approx_fprime\n'), ((149, 8, 149, 47), 'numpy.testing.assert_allclose', 'assert_allclose', (), '', False, 'from numpy.testing import assert_allclose\n'), ((166, 29, 166, 45), 'numpy.ones', 'np.ones', ({(166, 37, 166, 44): 'n // 10'}, {}), '(n // 10)', True, 'import numpy as np\n'), ((168, 43, 168, 59), 'numpy.ones', 'np.ones', ({(168, 51, 168, 58): 'n // 10'}, {}), '(n // 10)', True, 'import numpy as np\n'), ((172, 15, 172, 32), 'numpy.dot', 'np.dot', ({(172, 22, 172, 23): 'x', (172, 25, 172, 31): 'params'}, {}), '(x, params)', True, 'import numpy as np\n'), ((207, 29, 207, 45), 'numpy.ones', 'np.ones', ({(207, 37, 207, 44): 'n // 10'}, {}), '(n // 10)', True, 'import numpy as np\n'), ((209, 43, 209, 59), 'numpy.ones', 'np.ones', ({(209, 51, 209, 58): 'n // 10'}, {}), '(n // 10)', True, 'import numpy as np\n'), ((213, 15, 213, 32), 'numpy.dot', 'np.dot', ({(213, 22, 213, 23): 'x', (213, 25, 213, 31): 'params'}, {}), '(x, params)', True, 'import numpy as np\n'), ((88, 21, 88, 53), 'statsmodels.discrete.conditional_models.ConditionalLogit', 'ConditionalLogit', (), '', False, 'from statsmodels.discrete.conditional_models import ConditionalLogit, ConditionalPoisson\n'), ((90, 21, 90, 55), 'statsmodels.discrete.conditional_models.ConditionalPoisson', 'ConditionalPoisson', (), '', False, 'from statsmodels.discrete.conditional_models import ConditionalLogit, ConditionalPoisson\n'), ((95, 21, 96, 63), 'statsmodels.discrete.conditional_models.ConditionalLogit.from_formula', 'ConditionalLogit.from_formula', (), '', False, 'from statsmodels.discrete.conditional_models import ConditionalLogit, ConditionalPoisson\n'), ((98, 21, 99, 63), 'statsmodels.discrete.conditional_models.ConditionalPoisson.from_formula', 'ConditionalPoisson.from_formula', (), '', False, 'from statsmodels.discrete.conditional_models import ConditionalLogit, ConditionalPoisson\n'), ((174, 20, 174, 37), 'numpy.exp', 'np.exp', ({(174, 27, 174, 36): '(-lin_pred)'}, {}), '(-lin_pred)', True, 'import numpy as np\n'), ((175, 9, 175, 34), 'numpy.random.uniform', 'np.random.uniform', (), '', True, 'import numpy as np\n')]