# Path: econtools/metrics/api.py
# def reg(df: pd.DataFrame,
#         y_name: str, x_name: Union[str, List[str]],
#         fe_name: Optional[str]=None, a_name: Optional[str]=None,
#         nosingles: bool=True,
#         vce_type: Optional[str]=None, cluster: Optional[str]=None,
#         shac: Optional[dict]=None,
#         addcons: Optional[bool]=None, nocons: bool=False,
#         awt_name: Optional[str]=None,
#         save_mem: bool=False,
#         check_colinear: bool=False,
#         ) -> Results:
#     """OLS Regression.
# 
#     Args:
#         df (DataFrame): Data with any relevant variables.
#         y_name (str): Column name in ``df`` of the dependent variable.
#         x_name (str or list): Column name(s) in ``df`` of the independent
#                 variables/regressors
# 
#     Keyword Args:
#         vce_type (str): Type of estimator to use for variance-covariance matrix
#             of estimated coefficients. Default is standard OLS. Possible
#             choices are:
#                 - 'robust' or 'hc1'
#                 - 'hc2'
#                 - 'hc3'
#                 - 'cluster' (requires kwarg ``cluster``)
#                 - 'shac' (requires kwarg ``shac``)
#         cluster (str): Column name in ``df`` used to cluster standard errors.
#         shac (dict): Arguments to pass to spatial HAC estimator.
#             Requires:
#                 - **x** (*str*): Column name in ``df`` to serve as longitude.
#                 - **y** (*str*): Column name in ``df`` to serve as latitude.
#                 - **kern** (*str*): Kernel to use in estimation. May be
#                   triangle (``tria``) or uniform (``unif``).
#                 - **band** (float): Bandwidth for kernel.
#         fe_name (str) - Column name in ``df`` that defines groups for within
#             transformation (demeaning).
#         a_name (str) - Deprecated. See ``fe_name``.
#         awt_name (str): Column name in ``df`` to use for analytic weights in
#             regression.
#         addcons (bool): Defaults to False. Add a constant to independent
#             variables. Has no effect if ``a_name`` is passed.
#         nocons (bool): Defaults to False. Flag so estimators know that
#             independent variables ``df`` do not include a constant. Only
#             affects degrees of freedom.
#         nosingles (bool): Defaults to True. Drop observations that are obsorbed
#             by the within transformation. Has no effect if ``a_name=None``.
#         save_mem (bool): Defaults to False. If True, the returned
#             :py:class:`~econtools.metrics.core.Results` object will not save
#             large objects, specifically ``yhat``, ``sample``, and ``resid``.
#         check_colinear (bool): Default False. Checks rank of regressor matrix,
#             X. If X is rank deficient, an error is raised that prints the
#             colinear columns.
# 
#     Returns:
#         A :py:class:`~econtools.metrics.core.Results` object
#     """
# 
#     fe_name = _a_name_deprecation_handling(a_name, fe_name)
# 
#     RegWorker = Regression(
#         df, y_name, x_name,
#         fe_name=fe_name, nosingles=nosingles, addcons=addcons, nocons=nocons,
#         vce_type=vce_type, cluster=cluster, shac=shac,
#         awt_name=awt_name,
#         save_mem=save_mem,
#         check_colinear=check_colinear,
#     )
# 
#     results = RegWorker.main()
#     return results
# 
# Path: econtools/metrics/core.py
# def _get_colinear_cols(matrix: pd.DataFrame) -> List[str]:
#     K = matrix.shape[1]
#     rank = la.matrix_rank(matrix)
#     if rank < K:
#         colinear_idx = find_colinear_columns(matrix.values,
#                                              arr_rank=rank)
#         colinear_cols = matrix.columns[colinear_idx]
#         return colinear_cols.tolist()
#     elif rank > K:
#         raise ValueError("Something has gone very, very wrong.")
#     else:
#         return []


from os import path
from econtools.metrics.api import reg
from econtools.metrics.core import _get_colinear_cols
import pandas as pd
import numpy as np
import pytest


class TestCheckColinear(object):

    def setup(self):
        """Stata reg output from `sysuse auto; reg price mpg`"""
        test_path = path.split(path.relpath(__file__))[0]
        auto_path = path.join(test_path, 'data', 'auto.dta')
        autodata = pd.read_stata(auto_path)
        self.df = autodata

    def test_one_colin_column(self):
        df = self.df.copy()
        df['neg_mpg'] = -1 * df['mpg']
        result = _get_colinear_cols(df[['price', 'mpg', 'neg_mpg']])
        expected = ['neg_mpg']
        assert result == expected

    def test_two_colin_column(self):
        df = self.df.copy()
        df['neg_mpg'] = -1 * df['mpg']
        df['mpg2'] = 2 * df['mpg']
        result = _get_colinear_cols(df[['price', 'mpg', 'neg_mpg', 'mpg2']])
        expected = ['neg_mpg', 'mpg2']
        assert result == expected

    def test_two_colin_column_reorder(self):
        df = self.df.copy()
        df['neg_mpg'] = -1 * df['mpg']
        df['mpg2'] = 2 * df['mpg']

# Path: econtools/metrics/api.py
# def reg(df: pd.DataFrame,
#         y_name: str, x_name: Union[str, List[str]],
#         fe_name: Optional[str]=None, a_name: Optional[str]=None,
#         nosingles: bool=True,
#         vce_type: Optional[str]=None, cluster: Optional[str]=None,
#         shac: Optional[dict]=None,
#         addcons: Optional[bool]=None, nocons: bool=False,
#         awt_name: Optional[str]=None,
#         save_mem: bool=False,
#         check_colinear: bool=False,
#         ) -> Results:
#     """OLS Regression.
# 
#     Args:
#         df (DataFrame): Data with any relevant variables.
#         y_name (str): Column name in ``df`` of the dependent variable.
#         x_name (str or list): Column name(s) in ``df`` of the independent
#                 variables/regressors
# 
#     Keyword Args:
#         vce_type (str): Type of estimator to use for variance-covariance matrix
#             of estimated coefficients. Default is standard OLS. Possible
#             choices are:
#                 - 'robust' or 'hc1'
#                 - 'hc2'
#                 - 'hc3'
#                 - 'cluster' (requires kwarg ``cluster``)
#                 - 'shac' (requires kwarg ``shac``)
#         cluster (str): Column name in ``df`` used to cluster standard errors.
#         shac (dict): Arguments to pass to spatial HAC estimator.
#             Requires:
#                 - **x** (*str*): Column name in ``df`` to serve as longitude.
#                 - **y** (*str*): Column name in ``df`` to serve as latitude.
#                 - **kern** (*str*): Kernel to use in estimation. May be
#                   triangle (``tria``) or uniform (``unif``).
#                 - **band** (float): Bandwidth for kernel.
#         fe_name (str) - Column name in ``df`` that defines groups for within
#             transformation (demeaning).
#         a_name (str) - Deprecated. See ``fe_name``.
#         awt_name (str): Column name in ``df`` to use for analytic weights in
#             regression.
#         addcons (bool): Defaults to False. Add a constant to independent
#             variables. Has no effect if ``a_name`` is passed.
#         nocons (bool): Defaults to False. Flag so estimators know that
#             independent variables ``df`` do not include a constant. Only
#             affects degrees of freedom.
#         nosingles (bool): Defaults to True. Drop observations that are obsorbed
#             by the within transformation. Has no effect if ``a_name=None``.
#         save_mem (bool): Defaults to False. If True, the returned
#             :py:class:`~econtools.metrics.core.Results` object will not save
#             large objects, specifically ``yhat``, ``sample``, and ``resid``.
#         check_colinear (bool): Default False. Checks rank of regressor matrix,
#             X. If X is rank deficient, an error is raised that prints the
#             colinear columns.
# 
#     Returns:
#         A :py:class:`~econtools.metrics.core.Results` object
#     """
# 
#     fe_name = _a_name_deprecation_handling(a_name, fe_name)
# 
#     RegWorker = Regression(
#         df, y_name, x_name,
#         fe_name=fe_name, nosingles=nosingles, addcons=addcons, nocons=nocons,
#         vce_type=vce_type, cluster=cluster, shac=shac,
#         awt_name=awt_name,
#         save_mem=save_mem,
#         check_colinear=check_colinear,
#     )
# 
#     results = RegWorker.main()
#     return results
# 
# Path: econtools/metrics/core.py
# def _get_colinear_cols(matrix: pd.DataFrame) -> List[str]:
#     K = matrix.shape[1]
#     rank = la.matrix_rank(matrix)
#     if rank < K:
#         colinear_idx = find_colinear_columns(matrix.values,
#                                              arr_rank=rank)
#         colinear_cols = matrix.columns[colinear_idx]
#         return colinear_cols.tolist()
#     elif rank > K:
#         raise ValueError("Something has gone very, very wrong.")
#     else:
#         return []


# Path: econtools/metrics/tests/test_colinear.py
from os import path
from econtools.metrics.api import reg
from econtools.metrics.core import _get_colinear_cols
import pandas as pd
import numpy as np
import pytest



class TestCheckColinear(object):

    def setup(self):
        """Stata reg output from `sysuse auto; reg price mpg`"""
        test_path = path.split(path.relpath(__file__))[0]
        auto_path = path.join(test_path, 'data', 'auto.dta')
        autodata = pd.read_stata(auto_path)
        self.df = autodata

    def test_one_colin_column(self):
        df = self.df.copy()
        df['neg_mpg'] = -1 * df['mpg']
        result = _get_colinear_cols(df[['price', 'mpg', 'neg_mpg']])
        expected = ['neg_mpg']
        assert result == expected

    def test_two_colin_column(self):
        df = self.df.copy()
        df['neg_mpg'] = -1 * df['mpg']
        df['mpg2'] = 2 * df['mpg']
        result = _get_colinear_cols(df[['price', 'mpg', 'neg_mpg', 'mpg2']])
        expected = ['neg_mpg', 'mpg2']
        assert result == expected

    def test_two_colin_column_reorder(self):
        df = self.df.copy()
        df['neg_mpg'] = -1 * df['mpg']
        df['mpg2'] = 2 * df['mpg']

        result = _get_colinear_cols(df[['mpg2', 'mpg', 'price', 'neg_mpg']])