from numpy.testing import assert_almost_equal
import pandas as pd

from statsmodels.multivariate.cancorr import CanCorr

data_fit = pd.DataFrame([[191, 36, 50,  5, 162,  60],
                         [189, 37, 52,  2, 110,  60],
                         [193, 38, 58, 12, 101, 101],
                         [162, 35, 62, 12, 105,  37],
                         [189, 35, 46, 13, 155,  58],
                         [182, 36, 56,  4, 101,  42],
                         [211, 38, 56,  8, 101,  38],
                         [167, 34, 60,  6, 125,  40],
                         [176, 31, 74, 15, 200,  40],
                         [154, 33, 56, 17, 251, 250],
                         [169, 34, 50, 17, 120,  38],
                         [166, 33, 52, 13, 210, 115],
                         [154, 34, 64, 14, 215, 105],
                         [247, 46, 50,  1,  50,  50],
                         [193, 36, 46,  6,  70,  31],
                         [202, 37, 62, 12, 210, 120],
                         [176, 37, 54,  4,  60,  25],
                         [157, 32, 52, 11, 230,  80],
                         [156, 33, 54, 15, 225,  73],
                         [138, 33, 68,  2, 110,  43]])


def test_cancorr():
    # Compare results to SAS example:
    # https://support.sas.com/documentation/cdl/en/statug/63347/HTML/default/
    # viewer.htm#statug_cancorr_sect020.htm

    X1 = data_fit.iloc[:, :3]
    Y1 = data_fit.iloc[:, 3:]
    mod = CanCorr(Y1, X1)
    r = mod.corr_test()
    assert_almost_equal(r.stats_mv.loc["Wilks' lambda", "Value"], 0.35039053, decimal=8)
    assert_almost_equal(
        r.stats_mv.loc["Pillai's trace", "Value"], 0.67848151, decimal=8
    )
    assert_almost_equal(
        r.stats_mv.loc["Hotelling-Lawley trace", "Value"], 1.77194146, decimal=8
    )
    assert_almost_equal(
        r.stats_mv.loc["Roy's greatest root", "Value"], 1.72473874, decimal=8
    )
    assert_almost_equal(r.stats_mv.loc["Wilks' lambda", "F Value"], 2.05, decimal=2)
    assert_almost_equal(r.stats_mv.loc["Pillai's trace", "F Value"], 1.56, decimal=2)
    assert_almost_equal(
        r.stats_mv.loc["Hotelling-Lawley trace", "F Value"], 2.64, decimal=2
    )
    assert_almost_equal(
        r.stats_mv.loc["Roy's greatest root", "F Value"], 9.20, decimal=2
    )
    assert_almost_equal(r.stats_mv.loc["Wilks' lambda", "Num DF"], 9, decimal=3)
    assert_almost_equal(r.stats_mv.loc["Pillai's trace", "Num DF"], 9, decimal=3)
    assert_almost_equal(
        r.stats_mv.loc["Hotelling-Lawley trace", "Num DF"], 9, decimal=3
    )
    assert_almost_equal(r.stats_mv.loc["Roy's greatest root", "Num DF"], 3, decimal=3)
    assert_almost_equal(r.stats_mv.loc["Wilks' lambda", "Den DF"], 34.223, decimal=3)
    assert_almost_equal(r.stats_mv.loc["Pillai's trace", "Den DF"], 48, decimal=3)
    assert_almost_equal(
        r.stats_mv.loc["Hotelling-Lawley trace", "Den DF"], 19.053, decimal=3
    )
    assert_almost_equal(r.stats_mv.loc["Roy's greatest root", "Den DF"], 16, decimal=3)
    assert_almost_equal(r.stats_mv.loc["Wilks' lambda", "Pr > F"], 0.0635, decimal=4)
    assert_almost_equal(r.stats_mv.loc["Pillai's trace", "Pr > F"], 0.1551, decimal=4)
    assert_almost_equal(
        r.stats_mv.loc["Hotelling-Lawley trace", "Pr > F"], 0.0357, decimal=4
    )
    assert_almost_equal(
        r.stats_mv.loc["Roy's greatest root", "Pr > F"], 0.0009, decimal=4
    )
    assert_almost_equal(r.stats.loc[0, "Wilks' lambda"], 0.35039053, decimal=8)
    assert_almost_equal(r.stats.loc[1, "Wilks' lambda"], 0.95472266, decimal=8)
    assert_almost_equal(r.stats.loc[2, "Wilks' lambda"], 0.99473355, decimal=8)
    assert_almost_equal(r.stats.loc[0, "F Value"], 2.05, decimal=2)
    assert_almost_equal(r.stats.loc[1, "F Value"], 0.18, decimal=2)
    assert_almost_equal(r.stats.loc[2, "F Value"], 0.08, decimal=2)
    assert_almost_equal(r.stats.loc[0, "Num DF"], 9, decimal=2)
    assert_almost_equal(r.stats.loc[1, "Num DF"], 4, decimal=2)
    assert_almost_equal(r.stats.loc[2, "Num DF"], 1, decimal=2)
    assert_almost_equal(r.stats.loc[0, "Den DF"], 34.223, decimal=3)
    assert_almost_equal(r.stats.loc[1, "Den DF"], 30, decimal=2)
    assert_almost_equal(r.stats.loc[2, "Den DF"], 16, decimal=2)
    assert_almost_equal(r.stats.loc[0, "Pr > F"], 0.0635, decimal=4)
    assert_almost_equal(r.stats.loc[1, "Pr > F"], 0.9491, decimal=4)
    assert_almost_equal(r.stats.loc[2, "Pr > F"], 0.7748, decimal=4)
