""" Tests for yatsm.regression.structural_break
"""
import numpy as np
import pandas as pd
import patsy
import pytest
import xarray as xr

try:
    from rpy2 import robjects
    from rpy2.robjects import pandas2ri
    from rpy2.robjects.packages import importr
    pandas2ri.activate()
except ImportError:
    has_rpy2 = False
else:
    has_rpy2 = True

from yatsm.regression import structural_break as sb


@pytest.fixture('module')
def test_data():
    """ A very silly regression with an obvious change point
    """
    np.random.seed(0)
    y = np.concatenate((np.random.rand(250), np.random.rand(250) + 1))
    df = pd.DataFrame(np.column_stack((np.ones(500), np.arange(500), y)),
                      columns=['int', 'slope', 'y'])

    return df


@pytest.fixture('module')
def rpy2_strucchange():
    if has_rpy2:
        try:
            utils = importr('utils')
            utils.install_packages(
                'strucchange',
                repos='http://cran.revolutionanalytics.com/'
            )
        except Exception as exc:
            pytest.skip('Unable to install "strucchange"')
        else:
            return True
    return False


@pytest.fixture('module')
def strucchange_cusum_OLS(rpy2_strucchange, test_data):
    if has_rpy2 and rpy2_strucchange:
        try:
            rstr = """
            function(dat) {
                library(strucchange)
                model <- efp(y ~ ., data=dat, type="OLS-CUSUM")
                sctest(model)
            }
            """
            rfunc = robjects.r(rstr)
            rdat = robjects.r.ts(test_data)
            result = rfunc(rdat)
            result = dict(zip(result.names, list(result)))
        except Exception as exc:
            pass
        else:
            results = (result['statistic'][0], result['p.value'][0])
            return results

    # Validated against strucchange==1.5.1
    return (3.706774647551639, 2.325251102774928e-12)


def test_cusum_OLS(test_data, strucchange_cusum_OLS):
    """ Tested against strucchange 1.5.1
    """
    y = test_data.pop('y')
    X = test_data
    # Test sending pandas
    result = sb.cusum_OLS(X, y)
    assert np.allclose(result.score, strucchange_cusum_OLS[0])
    assert np.allclose(result.pvalue, strucchange_cusum_OLS[1])

    # And ndarray and xarray
    result = sb.cusum_OLS(X.values, xr.DataArray(y, dims=['time']))
    assert np.allclose(result.score, strucchange_cusum_OLS[0])
    assert np.allclose(result.pvalue, strucchange_cusum_OLS[1])


def test_cusum_rec_efp_sctest_pvalue(airquality):
    y = airquality['Ozone'].values
    X = patsy.dmatrix('1 + SolarR + Wind + Temp', data=airquality)

    process = sb._cusum_rec_efp(X, y)
    stat = sb._cusum_rec_sctest(process)
    pvalue = sb._brownian_motion_pvalue(0.2335143, 1)

    np.testing.assert_allclose(process.sum(), 1.9446869, rtol=1e-5)
    np.testing.assert_allclose(stat, 0.2335143, rtol=1e-5)
    np.testing.assert_allclose(pvalue, 0.9657902, rtol=1e-4)


@pytest.mark.parametrize(('alpha', 'truth'), [
    # pvalues are obtained from statsmodels and are estimates
    (0.10, 0.850),
    (0.05, 0.948),
    (0.01, 1.143)
])
def test_cusum_rec_test_crit(alpha, truth):
    test = sb._cusum_rec_test_crit(alpha)
    np.testing.assert_allclose(truth, test, rtol=1e-3)

@@ -362,48 +362,43 @@
 tsm.
-regression import structural_break as sb
+structural_break import cusum as cu
 %0A%0A%0A@
@@ -817,32 +817,67 @@
 2:%0A        try:%0A
+            base = importr('base')%0A
             util
@@ -897,16 +897,101 @@
 utils')%0A
+            has_pkg = base.require('strucchange')%5B0%5D%0A            if not has_pkg:%0A    
         
@@ -1034,16 +1034,20 @@
         
+    
 'strucch
@@ -1053,16 +1053,20 @@
 hange',%0A
+    
         
@@ -1118,16 +1118,20 @@
 s.com/'%0A
+    
         
@@ -2242,34 +2242,34 @@
 as%0A    result = 
-sb
+cu
 .cusum_OLS(X, y)
@@ -2439,18 +2439,18 @@
 esult = 
-sb
+cu
 .cusum_O
@@ -2790,18 +2790,18 @@
 ocess = 
-sb
+cu
 ._cusum_
@@ -2821,26 +2821,26 @@
 %0A    stat = 
-sb
+cu
 ._cusum_rec_
@@ -2868,18 +2868,18 @@
 value = 
-sb
+cu
 ._browni
@@ -3327,10 +3327,10 @@
 t = 
-sb
+cu
 ._cu

Fix test and improve test speed