# ----------------------------------------------------------------------------
# Copyright (c) 2016--, gneiss development team.
#
# Distributed under the terms of the GPLv3 License.
#
# The full license is in the file COPYING.txt, distributed with this software.
# ----------------------------------------------------------------------------
import os
import numpy as np
import numpy.testing as npt
import pandas as pd
import pandas.util.testing as pdt
import statsmodels.formula.api as smf
import unittest
import skbio
from gneiss.regression import RegressionResults
from skbio.stats.composition import _gram_schmidt_basis, ilr_inv


class TestRegressionResults(unittest.TestCase):

    def setUp(self):
        self.data = pd.DataFrame([[1, 1, 1],
                                  [3, 2, 3],
                                  [4, 3, 2],
                                  [5, 4, 4],
                                  [2, 5, 3],
                                  [3, 6, 5],
                                  [4, 7, 4]],
                                 index=['s1', 's2', 's3', 's4',
                                        's5', 's6', 's7'],
                                 columns=['Y1', 'Y2', 'X'])
        model1 = smf.ols(formula="Y1 ~ X", data=self.data)
        model2 = smf.ols(formula="Y2 ~ X", data=self.data)
        self.results = [model1.fit(), model2.fit()]
        self.pickle_fname = "test.pickle"
        self.tree = '(a,b)r;'

    def tearDown(self):
        if os.path.exists(self.pickle_fname):
            os.remove(self.pickle_fname)

    def test_balances(self):
        res = RegressionResults(self.results, balances=np.arange(7))
        npt.assert_almost_equal(res.balances, np.arange(7))

    def test_tree(self):
        t = skbio.TreeNode.read([u"a;"])
        res = RegressionResults(self.results, tree=t)
        self.assertEqual(str(res.tree), "a;\n")

    def test_r2(self):
        fittedvalues = pd.DataFrame({'s1': [1.986842, 1.236842],
                                     's2': [3.065789, 3.815789],
                                     's3': [2.526316, 2.526316],
                                     's4': [3.605263, 5.105263],
                                     's5': [3.065789, 3.815789],
                                     's6': [4.144737, 6.394737],
                                     's7': [3.605263, 5.105263]},
                                    index=['Y1', 'Y2']).T
        m = self.data.mean(axis=0)
        sse = ((fittedvalues - self.data.iloc[:, :2])**2).sum().sum()
        # ssr = ((fittedvalues - m)**2).sum().sum()
        sst = ((m - self.data.iloc[:, :2])**2).sum().sum()
        exp_r2 = 1 - (sse / sst)

        res = RegressionResults(self.results)
        self.assertAlmostEqual(exp_r2, res.r2)

    def test_regression_results_pvalues(self):
        # checks to see if pvalues are calculated correctly.
        res = RegressionResults(self.results)
        exp = pd.DataFrame({'Intercept': [0.307081, 0.972395],
                            'X': [0.211391, 0.029677]},
                           index=['Y1', 'Y2'])
        pdt.assert_frame_equal(res.pvalues, exp,
                               check_exact=False,
                               check_less_precise=True)

    def test_check_projection(self):
        feature_names = ['Z1', 'Z2', 'Z3']
        basis = _gram_schmidt_basis(3)
        res = RegressionResults(self.results, basis=basis,
                                feature_names=feature_names)

        feature_names = ['Z1', 'Z2', 'Z3']
        basis = _gram_schmidt_basis(3)

        # Test if feature_names is checked for
        res = RegressionResults(self.results, basis=basis)
        with self.assertRaises(ValueError):
            res._check_projection(True)

        # Test if basis is checked for
        res = RegressionResults(self.results, feature_names=feature_names)
        with self.assertRaises(ValueError):
            res._check_projection(True)

    def test_regression_results_coefficient(self):
        exp_coef = pd.DataFrame({'Intercept': [1.447368, -0.052632],
                                 'X': [0.539474, 1.289474]},
                                index=['Y1', 'Y2'])
        res = RegressionResults(self.results)
        pdt.assert_frame_equal(res.coefficients(), exp_coef,
                               check_exact=False,
                               check_less_precise=True)

    def test_regression_results_coefficient_projection(self):
        exp_coef = pd.DataFrame(
            {'Intercept': ilr_inv(np.array([[1.447368, -0.052632]])),
             'X': ilr_inv(np.array([[0.539474, 1.289474]]))},
            index=['Z1', 'Z2', 'Z3'])
        feature_names = ['Z1', 'Z2', 'Z3']
        basis = _gram_schmidt_basis(3)
        res = RegressionResults(self.results, basis=basis,
                                feature_names=feature_names)

        pdt.assert_frame_equal(res.coefficients(project=True), exp_coef,
                               check_exact=False,
                               check_less_precise=True)

    def test_regression_results_coefficient_project_error(self):
        res = RegressionResults(self.results)
        with self.assertRaises(ValueError):
            res.coefficients(project=True)

    def test_regression_results_residuals_projection(self):
        A = np.array  # aliasing np.array for the sake of pep8
        exp_resid = pd.DataFrame({'s1': ilr_inv(A([-0.986842, -0.236842])),
                                  's2': ilr_inv(A([-0.065789, -1.815789])),
                                  's3': ilr_inv(A([1.473684, 0.473684])),
                                  's4': ilr_inv(A([1.394737, -1.105263])),
                                  's5': ilr_inv(A([-1.065789, 1.184211])),
                                  's6': ilr_inv(A([-1.144737, -0.394737])),
                                  's7': ilr_inv(A([0.394737, 1.894737]))},
                                 index=['Z1', 'Z2', 'Z3']).T
        feature_names = ['Z1', 'Z2', 'Z3']
        basis = _gram_schmidt_basis(3)
        res = RegressionResults(self.results, basis=basis,
                                feature_names=feature_names)
        pdt.assert_frame_equal(res.residuals(project=True), exp_resid,
                               check_exact=False,
                               check_less_precise=True)

    def test_regression_results_residuals(self):
        exp_resid = pd.DataFrame({'s1': [-0.986842, -0.236842],
                                  's2': [-0.065789, -1.815789],
                                  's3': [1.473684, 0.473684],
                                  's4': [1.394737, -1.105263],
                                  's5': [-1.065789, 1.184211],
                                  's6': [-1.144737, -0.394737],
                                  's7': [0.394737, 1.894737]},
                                 index=['Y1', 'Y2']).T
        res = RegressionResults(self.results)
        pdt.assert_frame_equal(res.residuals(), exp_resid,
                               check_exact=False,
                               check_less_precise=True)

    def test_regression_results_predict(self):
        model = RegressionResults(self.results)
        res_predict = model.predict(self.data[['X']])

        exp_predict = pd.DataFrame({'s1': [1.986842, 1.236842],
                                    's2': [3.065789, 3.815789],
                                    's3': [2.526316, 2.526316],
                                    's4': [3.605263, 5.105263],
                                    's5': [3.065789, 3.815789],
                                    's6': [4.144737, 6.394737],
                                    's7': [3.605263, 5.105263]},
                                   index=['Y1', 'Y2']).T

        pdt.assert_frame_equal(res_predict, exp_predict)

    def test_regression_results_predict_extrapolate(self):
        model = RegressionResults(self.results)
        extrapolate = pd.DataFrame({'X': [8, 9, 10]},
                                   index=['k1', 'k2', 'k3'])
        res_predict = model.predict(extrapolate)

        exp_predict = pd.DataFrame({'k1': [5.76315789, 10.26315789],
                                    'k2': [6.30263158, 11.55263158],
                                    'k3': [6.84210526, 12.84210526]},
                                   index=['Y1', 'Y2']).T

        pdt.assert_frame_equal(res_predict, exp_predict)

    def test_regression_results_predict_projection(self):
        feature_names = ['Z1', 'Z2', 'Z3']
        basis = _gram_schmidt_basis(3)
        model = RegressionResults(self.results, basis=basis,
                                  feature_names=feature_names)

        res_predict = model.predict(self.data[['X']], project=True)
        A = np.array  # aliasing np.array for the sake of pep8
        exp_predict = pd.DataFrame({'s1': ilr_inv(A([1.986842, 1.236842])),
                                    's2': ilr_inv(A([3.065789, 3.815789])),
                                    's3': ilr_inv(A([2.526316, 2.526316])),
                                    's4': ilr_inv(A([3.605263, 5.105263])),
                                    's5': ilr_inv(A([3.065789, 3.815789])),
                                    's6': ilr_inv(A([4.144737, 6.394737])),
                                    's7': ilr_inv(A([3.605263, 5.105263]))},
                                   index=feature_names).T

        pdt.assert_frame_equal(res_predict, exp_predict)

    def test_regression_results_predict_none(self):
        model = RegressionResults(self.results)
        res_predict = model.predict()

        exp_predict = pd.DataFrame({'s1': [1.986842, 1.236842],
                                    's2': [3.065789, 3.815789],
                                    's3': [2.526316, 2.526316],
                                    's4': [3.605263, 5.105263],
                                    's5': [3.065789, 3.815789],
                                    's6': [4.144737, 6.394737],
                                    's7': [3.605263, 5.105263]},
                                   index=['Y1', 'Y2']).T

        pdt.assert_frame_equal(res_predict, exp_predict)

    def test_read_write(self):
        exp = RegressionResults(self.results, tree=self.tree)
        exp.write_pickle(self.pickle_fname)
        res = RegressionResults.read_pickle(self.pickle_fname)
        self.assertEqual(str(res.tree), str(exp.tree))
        pdt.assert_frame_equal(res.pvalues, exp.pvalues)


if __name__ == "__main__":
    unittest.main()
