import pytest
import numpy as np
import pandas as pd
import numpy.testing as npt
import pandas.testing as pdt
from scipy.stats import logistic

import zepid as ze
from zepid import (RiskRatio, RiskDifference, OddsRatio, NNT, IncidenceRateRatio, IncidenceRateDifference,
                   Sensitivity, Specificity, Diagnostics, interaction_contrast, interaction_contrast_ratio, spline,
                   table1_generator)
from zepid.calc import sensitivity, specificity

@pytest.fixture
def data_set():
    df = pd.DataFrame()
    df['exp'] = [1]*50 + [0]*50
    df['dis'] = [1]*25 + [0]*25 + [1]*25 + [0]*25
    return df

@pytest.fixture
def multi_exposures():
    df = pd.DataFrame()
    df['exp'] = [1]*50 + [0]*50 + [2]*50
    df['dis'] = [1]*25 + [0]*25 + [1]*25 + [0]*25 + [1]*25 + [0]*25
    return df

@pytest.fixture
def time_data():
    df = pd.DataFrame()
    df['exp'] = [1]*50 + [0]*50
    df['dis'] = [1]*6 + [0]*44 + [1]*14 + [0]*36
    df['t'] = [2]*50 + [8]*50
    return df

class TestRiskRatio:

    def test_risk_ratio_reference_equal_to_1(self, data_set):
        rr = RiskRatio()
        rr.fit(data_set, exposure='exp', outcome='dis')
        assert rr.risk_ratio[0] == 1

    def test_risk_ratio_equal_to_1(self, data_set):
        rr = RiskRatio()
        rr.fit(data_set, exposure='exp', outcome='dis')
        assert rr.risk_ratio[1] == 1

    def test_multiple_exposures(self, multi_exposures):
        rr = RiskRatio()
        rr.fit(multi_exposures, exposure='exp', outcome='dis')
        assert rr.results.shape[0] == 3
        assert list(rr.results.index) == ['Ref:0', '1', '2']

    def test_match_sas_ci(self, data_set):
        sas_ci = 0.6757, 1.4799
        rr = RiskRatio()
        rr.fit(data_set, exposure='exp', outcome='dis')
        df = rr.results
        npt.assert_allclose(np.round(df.loc[df.index == '1'][['RR_LCL', 'RR_UCL']], 4), [sas_ci])

    def test_match_sas_sampledata(self):
        sas_rd = 0.742118331
        sas_se = 0.312612740
        sas_ci = 0.402139480, 1.369523870
        df = ze.load_sample_data(False)
        rr = RiskRatio()
        rr.fit(df, exposure='art', outcome='dead')
        npt.assert_allclose(rr.risk_ratio[1], sas_rd, rtol=1e-5)
        rf = rr.results
        npt.assert_allclose(rf.loc[rf.index == '1'][['RR_LCL', 'RR_UCL']], [sas_ci], rtol=1e-5)
        npt.assert_allclose(rf.loc[rf.index == '1'][['SD(RR)']], sas_se, rtol=1e-5)

class TestRiskDifference:

    def test_risk_difference_reference_equal_to_0(self, data_set):
        rd = RiskDifference()
        rd.fit(data_set, exposure='exp', outcome='dis')
        assert rd.risk_difference[0] == 0

    def test_risk_difference_equal_to_0(self, data_set):
        rd = RiskDifference()
        rd.fit(data_set, exposure='exp', outcome='dis')
        assert rd.risk_difference[1] == 0

    def test_multiple_exposures(self, multi_exposures):
        rd = RiskDifference()
        rd.fit(multi_exposures, exposure='exp', outcome='dis')
        assert rd.results.shape[0] == 3
        assert list(rd.results.index) == ['Ref:0', '1', '2']

    def test_match_sas_ci(self, data_set):
        sas_ci = -0.195996398, 0.195996398
        rd = RiskDifference()
        rd.fit(data_set, exposure='exp', outcome='dis')
        df = rd.results
        npt.assert_allclose(df.loc[df.index == '1'][['RD_LCL', 'RD_UCL']], [sas_ci])

    def test_match_sas_se(self, data_set):
        sas_se = 0.1
        rd = RiskDifference()
        rd.fit(data_set, exposure='exp', outcome='dis')
        df = rd.results
        npt.assert_allclose(df.loc[df.index == '1'][['SD(RD)']], sas_se)

    def test_match_sas_sampledata(self):
        sas_rr = -0.045129870
        sas_se = 0.042375793
        sas_ci = -0.128184899, 0.037925158
        df = ze.load_sample_data(False)
        rd = RiskDifference()
        rd.fit(df, exposure='art', outcome='dead')
        npt.assert_allclose(rd.risk_difference[1], sas_rr)
        rf = rd.results
        npt.assert_allclose(rf.loc[rf.index == '1'][['RD_LCL', 'RD_UCL']], [sas_ci])
        npt.assert_allclose(rf.loc[rf.index == '1'][['SD(RD)']], sas_se)

    def test_frechet_bounds(self):
        df = ze.load_sample_data(False)
        rd = RiskDifference()
        rd.fit(df, exposure='art', outcome='dead')
        npt.assert_allclose(rd.results['UpperBound'][1] - rd.results['LowerBound'][1], 1.0000)

    def test_frechet_bounds2(self, multi_exposures):
        rd = RiskDifference()
        rd.fit(multi_exposures, exposure='exp', outcome='dis')
        npt.assert_allclose(rd.results['UpperBound'][1:] - rd.results['LowerBound'][1:], [1.0000, 1.0000])

class TestOddsRatio:

    def test_odds_ratio_reference_equal_to_1(self, data_set):
        ord = OddsRatio()
        ord.fit(data_set, exposure='exp', outcome='dis')
        assert ord.odds_ratio[0] == 1

    def test_odds_ratio_equal_to_1(self, data_set):
        ord = OddsRatio()
        ord.fit(data_set, exposure='exp', outcome='dis')
        assert ord.odds_ratio[1] == 1

    def test_multiple_exposures(self, multi_exposures):
        ord = OddsRatio()
        ord.fit(multi_exposures, exposure='exp', outcome='dis')
        assert ord.results.shape[0] == 3
        assert list(ord.results.index) == ['Ref:0', '1', '2']

    def test_match_sas_ci(self, data_set):
        sas_ci = 0.4566, 2.1902
        ord = OddsRatio()
        ord.fit(data_set, exposure='exp', outcome='dis')
        df = ord.results
        npt.assert_allclose(df.loc[df.index == '1'][['OR_LCL', 'OR_UCL']], [sas_ci], rtol=1e-4)

    def test_match_sas_sampledata(self):
        sas_or = 0.7036
        sas_se = 0.361479191
        sas_ci = 0.3465, 1.4290
        df = ze.load_sample_data(False)
        ord = OddsRatio()
        ord.fit(df, exposure='art', outcome='dead')
        npt.assert_allclose(ord.odds_ratio[1], sas_or, rtol=1e-4)
        rf = ord.results
        npt.assert_allclose(rf.loc[rf.index == '1'][['OR_LCL', 'OR_UCL']], [sas_ci], rtol=1e-3)
        npt.assert_allclose(rf.loc[rf.index == '1'][['SD(OR)']], sas_se, rtol=1e-4)

class TestNNT:

    def test_return_infinity(self, data_set):
        nnt = NNT()
        nnt.fit(data_set, exposure='exp', outcome='dis')
        assert np.isinf(nnt.number_needed_to_treat[1])

    def test_match_inverse_of_risk_difference(self):
        df = ze.load_sample_data(False)

        rd = RiskDifference()
        rd.fit(df, exposure='art', outcome='dead')

        nnt = NNT()
        nnt.fit(df, exposure='art', outcome='dead')

        npt.assert_allclose(nnt.number_needed_to_treat[1], 1/rd.risk_difference[1])
        rf = rd.results
        nf = nnt.results
        npt.assert_allclose(nf.loc[nf.index == '1'][['NNT_LCL', 'NNT_UCL']],
                            1 / rf.loc[rf.index == '1'][['RD_LCL', 'RD_UCL']])
        npt.assert_allclose(nf.loc[nf.index == '1'][['SD(RD)']], rf.loc[rf.index == '1'][['SD(RD)']])

    def test_multiple_exposures(self, multi_exposures):
        nnt = NNT()
        nnt.fit(multi_exposures, exposure='exp', outcome='dis')
        assert nnt.results.shape[0] == 3
        assert list(nnt.results.index) == ['Ref:0', '1', '2']

class TestIncidenceRateRatio:

    def test_incidence_rate_ratio_reference_equal_to_1(self, time_data):
        irr = IncidenceRateRatio()
        irr.fit(time_data, exposure='exp', outcome='dis', time='t')
        assert irr.incidence_rate_ratio[0] == 1

    def test_incidence_rate_ratio_equal_to_expected(self, time_data):
        sas_irr = 1.714285714
        sas_se = 0.487950036
        sas_ci = 0.658778447, 4.460946657
        irr = IncidenceRateRatio()
        irr.fit(time_data, exposure='exp', outcome='dis', time='t')
        npt.assert_allclose(irr.incidence_rate_ratio[1], sas_irr, rtol=1e-4)
        rf = irr.results
        npt.assert_allclose(rf.loc[rf.index == '1'][['IRR_LCL', 'IRR_UCL']], [sas_ci], rtol=1e-4)
        npt.assert_allclose(rf.loc[rf.index == '1'][['SD(IRR)']], sas_se, rtol=1e-4)

    def test_multiple_exposures(self):
        df = pd.DataFrame()
        df['exp'] = [1]*50 + [0]*50 + [2]*50
        df['dis'] = [1]*25 + [0]*25 + [1]*25 + [0]*25 + [1]*25 + [0]*25
        df['t'] = 2
        irr = IncidenceRateRatio()
        irr.fit(df, exposure='exp', outcome='dis', time='t')
        assert irr.results.shape[0] == 3
        assert list(irr.results.index) == ['Ref:0', '1', '2']

    def test_match_sas_sampledata(self):
        sas_irr = 0.753956
        sas_se = 0.336135409
        sas_ci = 0.390146, 1.457017
        df = ze.load_sample_data(False)
        irr = IncidenceRateRatio()
        irr.fit(df, exposure='art', outcome='dead', time='t')
        npt.assert_allclose(irr.incidence_rate_ratio[1], sas_irr, rtol=1e-5)
        rf = irr.results
        npt.assert_allclose(rf.loc[rf.index == '1'][['IRR_LCL', 'IRR_UCL']], [sas_ci], rtol=1e-5)
        npt.assert_allclose(rf.loc[rf.index == '1'][['SD(IRR)']], sas_se, rtol=1e-5)

class TestIncidenceRateDifference:

    def test_incidence_rate_difference_reference_equal_to_0(self, time_data):
        ird = IncidenceRateDifference()
        ird.fit(time_data, exposure='exp', outcome='dis', time='t')
        assert ird.incidence_rate_difference[0] == 0

    def test_multiple_exposures(self):
        df = pd.DataFrame()
        df['exp'] = [1]*50 + [0]*50 + [2]*50
        df['dis'] = [1]*25 + [0]*25 + [1]*25 + [0]*25 + [1]*25 + [0]*25
        df['t'] = 2
        ird = IncidenceRateDifference()
        ird.fit(df, exposure='exp', outcome='dis', time='t')
        assert ird.results.shape[0] == 3
        assert list(ird.results.index) == ['Ref:0', '1', '2']

    def test_match_openepi_sampledata(self):
        oe_irr = -0.001055
        oe_ci = -0.003275, 0.001166
        df = ze.load_sample_data(False)
        ird = IncidenceRateDifference()
        ird.fit(df, exposure='art', outcome='dead', time='t')
        npt.assert_allclose(ird.incidence_rate_difference[1], oe_irr, atol=1e-5)
        rf = ird.results
        npt.assert_allclose(rf.loc[rf.index == '1'][['IRD_LCL', 'IRD_UCL']], [oe_ci], atol=1e-5)

class TestDiagnostics:

    @pytest.fixture
    def test_data(self):
        df = pd.DataFrame()
        df['test'] = [1]*50 + [0]*50
        df['case'] = [1]*40 + [0]*10 + [1]*15 + [0]*35
        return df

    def test_sensitivity_same_as_calc(self, test_data):
        se = Sensitivity()
        se.fit(test_data, test='test', disease='case')
        sens = sensitivity(40, 50)
        npt.assert_allclose(se.sensitivity, sens[0])

    def test_specificity_same_as_calc(self, test_data):
        sp = Specificity()
        sp.fit(test_data, test='test', disease='case')
        spec = specificity(15, 50)
        npt.assert_allclose(sp.specificity, spec[0])

    def test_diagnostic_same_as_compositions(self, test_data):
        se = Sensitivity()
        se.fit(test_data, test='test', disease='case')

        sp = Specificity()
        sp.fit(test_data, test='test', disease='case')

        diag = Diagnostics()
        diag.fit(test_data, test='test', disease='case')

        npt.assert_allclose(diag.sensitivity.sensitivity, se.sensitivity)
        npt.assert_allclose(diag.specificity.specificity, sp.specificity)

    def test_match_sas_sensitivity_ci(self, test_data):
        sas_ci = [0.689127694, 0.910872306]
        diag = Diagnostics()
        diag.fit(test_data, test='test', disease='case')
        npt.assert_allclose(diag.sensitivity.results[['Se_LCL', 'Se_UCL']], [sas_ci])

    def test_match_sas_specificity_ci(self, test_data):
        sas_ci = [0.572979816, 0.827020184]
        diag = Diagnostics()
        diag.fit(test_data, test='test', disease='case')
        npt.assert_allclose(diag.specificity.results[['Sp_LCL', 'Sp_UCL']], [sas_ci])

class TestInteractionContrasts:

    @pytest.fixture
    def data_ic(self, n=10000):
        df = pd.DataFrame()
        np.random.seed(111)
        df['exp'] = np.random.binomial(1, 0.5, size=n)
        df['mod'] = np.random.binomial(1, 0.5, size=n)
        df['y'] = np.random.binomial(1, size=n, p=logistic.cdf(0.1 + 0.2*df['exp'] + 0.3*df['mod'] -
                                                               0.4*df['mod']*df['exp']))
        # Note: IC will not be equal to ICR
        return df

    def test_interaction_contrast(self, data_ic):
        ic = interaction_contrast(data_ic, exposure='exp', outcome='y', modifier='mod', print_results=False)
        npt.assert_allclose(np.round(ic[0], 4), -0.1009)

    def test_interaction_contrast_ci(self, data_ic):
        ic = interaction_contrast(data_ic, exposure='exp', outcome='y', modifier='mod', print_results=False)
        assert ic[1] < -0.1009 < ic[2]

    def test_interaction_contrast_ratio(self, data_ic):
        icr = interaction_contrast_ratio(data_ic, exposure='exp', outcome='y', modifier='mod', print_results=False)
        npt.assert_allclose(np.round(icr[0], 4), -0.4908)

    def test_interaction_contrast_ratio_delta_ci(self, data_ic):
        icr = interaction_contrast_ratio(data_ic, exposure='exp', outcome='y', modifier='mod', print_results=False)
        assert icr[1] < -0.4908 < icr[2]

    def test_interaction_contrast_ratio_bootstrap_ci(self, data_ic):
        icr = interaction_contrast_ratio(data_ic, exposure='exp', outcome='y', modifier='mod',
                                         ci='bootstrap', print_results=False)
        assert icr[1] < -0.4908 < icr[2]

class TestSplines:

    @pytest.fixture
    def spline_data(self):
        df = pd.DataFrame()
        df['v'] = [1, 5, 10, 15, 20]
        return df

    def test_error_for_bad_nknots(self, spline_data):
        with pytest.raises(ValueError):
            spline_data['sp'] = spline(spline_data, 'v', n_knots=1.5)
        with pytest.raises(ValueError):
            spline_data['sp'] = spline(spline_data, 'v', n_knots=0)
        with pytest.raises(ValueError):
            spline_data['sp'] = spline(spline_data, 'v', n_knots=-1)
        with pytest.raises(ValueError):
            spline_data['sp'] = spline(spline_data, 'v', n_knots=8)

    def test_error_for_unequal_numbers(self, spline_data):
        with pytest.raises(ValueError):
            spline_data['sp'] = spline(spline_data, 'v', n_knots=1, knots=[1, 3])
        with pytest.raises(ValueError):
            spline_data['sp'] = spline(spline_data, 'v', n_knots=3, knots=[1, 3])

    def test_error_for_bad_order(self, spline_data):
        with pytest.raises(ValueError):
            spline_data['sp'] = spline(spline_data, 'v', n_knots=3, knots=[3, 1, 2])

    def test_auto_knots1(self, spline_data):
        spline_data['sp'] = spline(spline_data, 'v', n_knots=1, restricted=False)
        expected_splines = pd.DataFrame.from_records([{'sp': 0.0},
                                                      {'sp': 0.0},
                                                      {'sp': 0.0},
                                                      {'sp': 5.0},
                                                      {'sp': 10.0}])
        pdt.assert_series_equal(spline_data['sp'], expected_splines['sp'])

    def test_auto_knots2(self, spline_data):
        spline_data[['sp1', 'sp2']] = spline(spline_data, 'v', n_knots=2, restricted=False)
        expected_splines = pd.DataFrame.from_records([{'sp1': 0.0, 'sp2': 0.0},
                                                      {'sp1': 0.0, 'sp2': 0.0},
                                                      {'sp1': 10 - 20/3, 'sp2': 0.0},
                                                      {'sp1': 15 - 20/3, 'sp2': 15 - 40/3},
                                                      {'sp1': 20 - 20/3, 'sp2': 20 - 40/3}])
        pdt.assert_frame_equal(spline_data[['sp1', 'sp2']], expected_splines[['sp1', 'sp2']])

    def test_user_knots1(self, spline_data):
        spline_data['sp'] = spline(spline_data, 'v', n_knots=1, knots=[16], restricted=False)
        expected_splines = pd.DataFrame.from_records([{'sp': 0.0},
                                                      {'sp': 0.0},
                                                      {'sp': 0.0},
                                                      {'sp': 0.0},
                                                      {'sp': 4.0}])
        pdt.assert_series_equal(spline_data['sp'], expected_splines['sp'])

    def test_user_knots2(self, spline_data):
        spline_data[['sp1', 'sp2']] = spline(spline_data, 'v', n_knots=2, knots=[10, 16], restricted=False)
        expected_splines = pd.DataFrame.from_records([{'sp1': 0.0, 'sp2': 0.0},
                                                      {'sp1': 0.0, 'sp2': 0.0},
                                                      {'sp1': 0.0, 'sp2': 0.0},
                                                      {'sp1': 5.0, 'sp2': 0.0},
                                                      {'sp1': 10.0, 'sp2': 4.0}])
        pdt.assert_frame_equal(spline_data[['sp1', 'sp2']], expected_splines[['sp1', 'sp2']])

    def test_quadratic_spline1(self, spline_data):
        spline_data['sp'] = spline(spline_data, 'v', n_knots=1, knots=[16], term=2, restricted=False)
        expected_splines = pd.DataFrame.from_records([{'sp': 0.0},
                                                      {'sp': 0.0},
                                                      {'sp': 0.0},
                                                      {'sp': 0.0},
                                                      {'sp': 4.0**2}])
        pdt.assert_series_equal(spline_data['sp'], expected_splines['sp'])

    def test_quadratic_spline2(self, spline_data):
        spline_data[['sp1', 'sp2']] = spline(spline_data, 'v', n_knots=2, knots=[10, 16], term=2, restricted=False)
        expected_splines = pd.DataFrame.from_records([{'sp1': 0.0, 'sp2': 0.0},
                                                      {'sp1': 0.0, 'sp2': 0.0},
                                                      {'sp1': 0.0, 'sp2': 0.0},
                                                      {'sp1': 5.0**2, 'sp2': 0.0},
                                                      {'sp1': 10.0**2, 'sp2': 4.0**2}])
        pdt.assert_frame_equal(spline_data[['sp1', 'sp2']], expected_splines[['sp1', 'sp2']])

    def test_cubic_spline1(self, spline_data):
        spline_data['sp'] = spline(spline_data, 'v', n_knots=1, knots=[16], term=3, restricted=False)
        expected_splines = pd.DataFrame.from_records([{'sp': 0.0},
                                                      {'sp': 0.0},
                                                      {'sp': 0.0},
                                                      {'sp': 0.0},
                                                      {'sp': 4.0**3}])
        pdt.assert_series_equal(spline_data['sp'], expected_splines['sp'])

    def test_cubic_spline2(self, spline_data):
        spline_data[['sp1', 'sp2']] = spline(spline_data, 'v', n_knots=2, knots=[10, 16], term=3, restricted=False)
        expected_splines = pd.DataFrame.from_records([{'sp1': 0.0, 'sp2': 0.0},
                                                      {'sp1': 0.0, 'sp2': 0.0},
                                                      {'sp1': 0.0, 'sp2': 0.0},
                                                      {'sp1': 5.0**3, 'sp2': 0.0},
                                                      {'sp1': 10.0**3, 'sp2': 4.0**3}])
        pdt.assert_frame_equal(spline_data[['sp1', 'sp2']], expected_splines[['sp1', 'sp2']])

    def test_higher_order_spline(self, spline_data):
        spline_data[['sp1', 'sp2']] = spline(spline_data, 'v', n_knots=2, knots=[10, 16], term=3.7, restricted=False)
        expected_splines = pd.DataFrame.from_records([{'sp1': 0.0, 'sp2': 0.0},
                                                      {'sp1': 0.0, 'sp2': 0.0},
                                                      {'sp1': 0.0, 'sp2': 0.0},
                                                      {'sp1': 5.0**3.7, 'sp2': 0.0},
                                                      {'sp1': 10.0**3.7, 'sp2': 4.0**3.7}])
        pdt.assert_frame_equal(spline_data[['sp1', 'sp2']], expected_splines[['sp1', 'sp2']])

    def test_restricted_spline1(self, spline_data):
        spline_data['rsp'] = spline(spline_data, 'v', n_knots=2, knots=[10, 16], restricted=True)
        expected_splines = pd.DataFrame.from_records([{'rsp': 0.0},
                                                      {'rsp': 0.0},
                                                      {'rsp': 0.0},
                                                      {'rsp': 5.0},
                                                      {'rsp': 6.0}])
        pdt.assert_series_equal(spline_data['rsp'], expected_splines['rsp'])

    def test_restricted_spline2(self, spline_data):
        spline_data['rsp'] = spline(spline_data, 'v', n_knots=2, knots=[5, 16], restricted=True)
        expected_splines = pd.DataFrame.from_records([{'rsp': 0.0},
                                                      {'rsp': 0.0},
                                                      {'rsp': 10.0 - 5.0},
                                                      {'rsp': 15.0 - 5.0},
                                                      {'rsp': (20.0 - 5.0) - (20.0 - 16.0)}])
        pdt.assert_series_equal(spline_data['rsp'], expected_splines['rsp'])

    def test_restricted_spline3(self, spline_data):
        spline_data['rsp'] = spline(spline_data, 'v', n_knots=2, knots=[5, 16], term=2, restricted=True)
        expected_splines = pd.DataFrame.from_records([{'rsp': 0.0},
                                                      {'rsp': 0.0},
                                                      {'rsp': (10.0 - 5.0)**2 - 0},
                                                      {'rsp': (15.0 - 5.0)**2 - 0},
                                                      {'rsp': (20.0 - 5.0)**2 - (20.0 - 16.0)**2}])
        pdt.assert_series_equal(spline_data['rsp'], expected_splines['rsp'])

class TestTable1:

    @pytest.fixture
    def data(self, n=1000):
        df = pd.DataFrame()
        np.random.seed(111)
        df['exp'] = np.random.binomial(1, 0.5, size=n)
        df['mod'] = np.random.binomial(1, 0.5, size=n)
        df['y'] = np.random.binomial(1, size=n, p=logistic.cdf(0.1 + 0.2*df['exp'] + 0.3*df['mod'] -
                                                               0.4*df['mod']*df['exp']))
        df['continuous'] = np.random.normal(size=n)
        return df

    def test_unstratified_median(self, data):
        t = table1_generator(data, cols=['exp', 'mod', 'y', 'continuous'],
                             variable_type=['category', 'category', 'category', 'continuous'])
        assert isinstance(t, type(pd.DataFrame()))

    def test_unstratified_mean(self, data):
        t = table1_generator(data, cols=['exp', 'mod', 'y', 'continuous'],
                             variable_type=['category', 'category', 'category', 'continuous'],
                             continuous_measure='mean')
        assert isinstance(t, type(pd.DataFrame()))

    def test_stratified_median(self, data):
        t = table1_generator(data, cols=['mod', 'y', 'continuous'],
                             variable_type=['category', 'category', 'continuous'], strat_by='exp')
        assert isinstance(t, type(pd.DataFrame()))

    def test_stratified_mean(self, data):
        t = table1_generator(data, cols=['mod', 'y', 'continuous'],
                             variable_type=['category', 'category', 'continuous'],
                             continuous_measure='mean', strat_by='exp')
        assert isinstance(t, type(pd.DataFrame()))

    def test_catch_different_lengths(self, data):
        with pytest.raises(ValueError):
            table1_generator(data, cols=['exp', 'mod', 'y', 'continuous'],
                             variable_type=['category', 'category', 'continuous'],
                             continuous_measure='A')

    def test_wrong_continuous_measure_error(self, data):
        with pytest.raises(ValueError):
            table1_generator(data, cols=['exp', 'mod', 'y', 'continuous'],
                             variable_type=['category', 'category', 'category', 'continuous'],
                             continuous_measure='A')
        with pytest.raises(ValueError):
            table1_generator(data, cols=['mod', 'y', 'continuous'],
                             variable_type=['category', 'category', 'continuous'],
                             continuous_measure='A', strat_by='exp')


"""
trkd_roi_analysis.py

This module run various anlyses on the tracked ROI USIs (Unexpected event
Selectivity Indices).

Authors: <NAME>

Last modified: 29 May 2021
"""

import numpy as np
import pandas as pd
import scipy.stats as scist
import sys, copy, time
import itertools as it

sys.path.extend(['../..'])
from analysis import session
from util import gen_util, math_util

#############################################

def set_up_mouse_info(mouse_df_fnm):
    """
    Returns mouse dataframe and mouse numbers.

    Parameters
    ----------
    mouse_df_fnm : string
        File name for the dataframe that contains experiment information

    Returns
    -------
    mouse_df : Pandas DataFrame
        Dataframe with mouse/experiment information
    mouse_ns_full : 1-D array of numbers
        All mouse numbers
    mouse_ns_sess_123 : 1-D array of numbers
        Mouse numbers for mice with data from all of sessions 1-3
    """

    # declarations/initializations
    mouse_df = pd.read_csv(mouse_df_fnm)
    # Restrict to production mice that passed and have all files
    mouse_df = mouse_df[(mouse_df.runtype=='prod') & (mouse_df.pass_fail=='P')
                        & (mouse_df.all_files)]
    # Obtain all mouse numbers
    mouse_ns_full = mouse_df['mouse_n'].unique()
    mouse_ns_sess_123 = []

    print('All mice: ', mouse_ns_full)

    for mouse_n in mouse_ns_full:
        # Obtain session numbers for each mouse
        sess_ns = mouse_df[(mouse_df.mouse_n==mouse_n)]['sess_n'].values
        # Add mouse to mouse numbers with data from sess. 1-3
        if np.sum(np.isin(sess_ns, np.asarray([1,2,3])))>=3:
            mouse_ns_sess_123.append(mouse_n)
        # Want mouse 8 as well for running.  will update pupil diameter later
        elif mouse_n==8:
            mouse_ns_sess_123.append(mouse_n)

    mouse_ns_sess_123 = np.asarray(mouse_ns_sess_123)
    print('Mice with all of sessions 1-3 :', mouse_ns_sess_123)

    return mouse_df, mouse_ns_full, mouse_ns_sess_123

#############################################

def make_usi_df(stimtype, mouse_df, mouse_df_fnm, mouse_ns, datadir,
                brk_dir = 'any',
                op = 'd-prime',
                only_matched_rois = True,
                sess_ns = [1,2,3],
                scale = True,
                remnans = False):

    """
    Returns a dataframe that contains, for each tracked ROI,
    USI, mean and std of the df/f for (un)expected events over each session.

    Parameters
    ----------
    stimtype : string
        Either 'gabors' or 'bricks'.
    mouse_df : Pandas DataFrame
        Dataframe with mouse information
    mouse_df_fnm : string
        File name for 'mouse_df'
    mouse_ns : 1-D arraylike of numbers
        Contains all mouse numbers for which to obtain data
    datadir : string
        Data directory
    brk_dir : string;  optional, default = 'any'
        Options: 'right', 'left', 'temp', 'nasal' or 'any' (default).  Specify
        brick direction if desired.
    op : string; optional, default = 'd-prime'
        Options: 'diff', 'discr', 'd-prime' (default).  Operation to use when
        computing USIs.
    only_matched_rois : boolean; optional, default = True
        If 'True' (default), only USIs for ROIs that have been matched across
        sessions (and contain no NaNs in any recorded session) are included.
        Else, USIs for all ROIs are included (excluding ROIs with NaN df/f
        values for any particular session if 'remnans' argument is set to True).
    sess_ns : 1-D arraylike; optional, default = [1,2,3]
        Which sessions to include for selected mice.
    scale : boolean; optional, default = True
        Option to scale each ROI by its statistics for each session.  Rescales
        through an analog to a z-score, where the mean is replaced by the
        median, and the standard deviation is replaced by the 5% to 95%
        interval.  This allows ROIs to be treated more equally to others in a
        given session, compensating for calcium signal issues, e.g.
    remnans : boolean; optional, default = False
        Option to discount ROIs that have NaN df/f values.  Note, this is
        redundant if 'only_matched_rois' is set to True.

    Returns
    -------
    usi_df (pd DataFrame) : Pandas dataframe
        Contains, for each tracked ROI, USI, mean and std of the df/f for
        (un)expected events over each session

        Columns:
            'mouse_n', 'layer', 'compartment' :
                mouse number, layer, and compartment, respectively
            'sess_<session>_usi' for <session> in [1,2,3] :
                USI values for given sessions
            'sess_<session>_expec_evnt_<statistic>' for <session> in [1,2,3],
            <statistic> in ['mn', 'std'] :
                mean and std of integrated df/f for each ROI over all
                expected events in sess. (D/G frames or expected flow)
            'sess_<session>_unexp_evnt_<statistic>' for <session> in [1,2,3],
            <statistic> in ['mn', 'std'] :
                mean and std of integrated df/f for each ROI over all expected
                events in sess. (U/G frames or unexpected flow)
    """

    # declarations/initializations
    t = time.time()
    layers_dict = {'L23-Cux2': 'L2/3', 'L5-Rbp4': 'L5'}

    tracked_roi_usi_df = \
        pd.DataFrame(columns=['mouse_n', 'layer', 'compartment',
            'sess_1_usi', 'sess_2_usi', 'sess_3_usi',
            'sess_1_expec_evnt_mn', 'sess_2_expec_evnt_mn',
            'sess_3_expec_evnt_mn',
            'sess_1_expec_evnt_std', 'sess_2_expec_evnt_std',
            'sess_3_expec_evnt_std',
            'sess_1_unexp_evnt_mn', 'sess_2_unexp_evnt_mn',
            'sess_3_unexp_evnt_mn',
            'sess_1_unexp_evnt_std', 'sess_2_unexp_evnt_std',
            'sess_3_unexp_evnt_std'])

    # Loop over all given mouse numbers
    for idx, mouse_n in enumerate(mouse_ns):
        # declarations/initializations
        unexp_mn = []
        unexp_std = []
        expec_mn = []
        expec_std = []
        usi = []
        # Determine layer and compartment
        layer = layers_dict[mouse_df[mouse_df['mouse_n']==mouse_n]['line'].
            values[0]]
        compartment = mouse_df[mouse_df['mouse_n']==mouse_n]['plane'].values[0]
        for sess_n in sess_ns:
            print('Mouse ', mouse_n, ', sess ', sess_n)
            # Obtain session object
            sessid = mouse_df[(mouse_df['mouse_n']==mouse_n) &
                              (mouse_df['sess_n']==sess_n)]['sessid'].values[0]
            sess = session.Session(datadir, sessid,
                only_matched_rois=only_matched_rois)
            sess.extract_sess_attribs(mouse_df=mouse_df_fnm)
            sess.extract_info(fulldict=False, roi=True, run=False, pupil=False)
            roi_data = []
            # Get expected event data
            expec_data = get_roi_data(sess, stimtype, surp=0, brk_dir=brk_dir,
                scale=scale, remnans=remnans)
            expec_data_mn = np.nanmean(expec_data, axis=1)
            expec_data_std = np.nanstd(expec_data, axis=1)
            expec_mn.append(expec_data_mn)
            expec_std.append(expec_data_std)
            roi_data.append(expec_data)
            # Get unexpected event data
            unexp_data = get_roi_data(sess, stimtype, surp=1, brk_dir=brk_dir,
                scale=scale, remnans=remnans)
            unexp_data_mn = np.nanmean(unexp_data, axis=1)
            unexp_data_std = np.nanstd(unexp_data, axis=1)
            unexp_mn.append(unexp_data_mn)
            unexp_std.append(unexp_data_std)
            roi_data.append(unexp_data)
            print('unexp_data.shape = ', unexp_data.shape)
            # Compute USI
            # Direct computation
            usi.append((unexp_data_mn - expec_data_mn) /
                        np.sqrt(0.5*(unexp_data_std**2 + expec_data_std**2)))

            # Below is equivalent to above directly computed line (no
            # permutations needed here). Note: no need for distribution
#             unexp_idxs, _, _ = surp_idx_by_sess(
#                 roi_data, n_perms=0, datatype='roi', op=op, nanpol=None)
#             usi.append(unexp_idxs)

            print('{:.2f} sec'.format(time.time()-t))

        tracked_roi_usi_df.loc[idx] = [mouse_n, layer, compartment,
                                       usi[0], usi[1], usi[2],
                                       expec_mn[0], expec_mn[1], expec_mn[2],
                                       expec_std[0], expec_std[1], expec_std[2],
                                       unexp_mn[0], unexp_mn[1], unexp_mn[2],
                                       unexp_std[0], unexp_std[1], unexp_std[2]]

    return tracked_roi_usi_df

#############################################

def get_roi_data(sess, stimtype, surp, brk_dir, scale, remnans):
    """
    Returns 2-D array of integrated df/f for each ROI and segment.  Segments are
    frames D-G or U-G (both comprise 0.6 sec each) for Gabors and the 2 sec
    immediately pre- or proceeding the onset of unexpected flow.

    Parameters
    ----------
    sess : session object
        Session object for particular mouse/session
    stimtype : string
        Options: 'gabors', 'bricks'
    brk_dir : string
        Options: 'right', 'left', 'temp', 'nasal' or 'any' (default).  Specify
        brick direction if desired
    scale : boolean; optional, default = True
        Option to scale each ROI by its statistics for each session.  Rescales
        through an analog to a z-score, where the mean is replaced by the
        median, and the standard deviation is replaced by the 5% to 95%
        interval.  This allows ROIs to be treated more equally to others in a
        given session, compensating for calcium signal issues, e.g.
    remnans : boolean; optional, default = False
        Option to discount ROIs that have NaN df/f values.  Note, this is
        redundant if 'only_matched_rois' is set to True.

    Returns
    -------
    roi_data : 2-D array of numbers
        Integated df/f for each ROI for each segment (frames D/U-G or 2s
        immediately before or after onset of unexpected flow).
        Rows: ROIs.  Columns: Segments (/sequences)
    """

    # Warameters:
    # Width for brick stim-locked traces = 2*half_width
    half_width = 2
    # Pre / post times, organized by stimulus and surprise
    pre_dict = {'gabors':[0,0], 'bricks':[half_width,0]}
    post_dict = {'gabors':[0.6,0.6], 'bricks':[0,half_width]}
    pre = pre_dict[stimtype][surp]
    post = post_dict[stimtype][surp]
    # Only get segements around counterflow onset
    surp = 1 if stimtype=='bricks' else surp

    # Obtain ROI data
    stim = sess.get_stim(stimtype)
    seg_ns = \
        stim.get_segs_by_criteria(surp=surp, gabk=16, gabfr=3,
                                  bri_size=128, bri_dir=brk_dir, by="seg",
                                  remconsec=(stimtype == "bricks"))
    twop_fr_ns = \
        stim.get_twop_fr_by_seg(seg_ns, first=True,
                                ch_fl=[pre, post])['first_twop_fr']
    roi_data_df = \
        stim.get_roi_data(twop_fr_ns, pre=pre, post=post, integ=True,
                          scale=scale, remnans=remnans)
    # Roi x seq (time is already integrated out):
    roi_data = gen_util.reshape_df_data(roi_data_df, squeeze_rows=False,
        squeeze_cols=True)

    return roi_data

#############################################

def surp_idx_by_sess(data, n_perms=1000, datatype='roi', op='diff',
                     stats='mean', nanpol=None):
    """
    surp_idx_by_sess(data)

    Returns session item (ROIs or 1 for running) indices for difference between
    surprise and regular sequences, as well as their percentiles based on
    random permutations for each item.

    Required args:
        - data (3D array): data array, structured as
                               reg, surp [x ROIs] x sequences

    Optional args:
        - n_perms (int)         : number of permutations for CI estimation
                                  default: 1000
        - datatype (str)        : type of data (e.g., 'roi', 'run')
                                  default: 'roi'
        - op (str)              : operation to use in measuring surprise
                                  indices ('diff', 'rel_diff', 'discr')
                                  default: 'diff'
        - stats (str)           : statistic used across sequences
                                  default: 'mean'
        - nanpol (str)          : NaN policy ('omit' or None)
                                  default: None

    Returns:
        - item_idxs (1-D array) : item (ROIs or 1 for running) surprise indices
                                 for the session
        - item_percs (1-D array): item (ROIs or 1 for running) surprise index
                                 percentiles for the session, based on
                                 each item's random permutations
        - all_rand (2-D array)  : item (ROIs or 1 for running) indices
                                 calculated through randomized permutation,
                                    structured as item x n_perms
    """

    # take statistic across sequences, unless the index id discr (D')
    if op != 'discr' and op != 'd-prime':
        seq_mes = np.stack([math_util.mean_med(
            subdata, stats=stats, axis=-1, nanpol=nanpol)
            for subdata in data])
        axis = None
    else:
        seq_mes = data
        axis = -1

    # calculate index
    item_idxs = math_util.calc_op(seq_mes, op=op, nanpol=nanpol)

    # reshape to add an item/channel dimension if datatype isn't ROIs
    last_dim = np.sum([sub.shape[-1] for sub in data])
    if datatype != 'roi':
        item_idxs = np.asarray(item_idxs).reshape(-1)
        targ = (1, last_dim)
    else:
        targ = (-1, last_dim)

    # get CI
    div = data[0].shape[-1] # length of reg
    # perms (items x perms)
    all_rand = math_util.permute_diff_ratio(
        np.concatenate(data, axis=-1).reshape(targ), div=div,
        n_perms=n_perms, stats=stats, nanpol=nanpol, op=op)

    item_percs = np.empty(len(item_idxs))
    for r, (item_idx, item_rand) in enumerate(zip(item_idxs, all_rand)):
        item_percs[r] = scist.percentileofscore(
            item_rand, item_idx, kind='mean')

    return item_idxs, item_percs, all_rand

#############################################

def make_usi_abs_mean_df(gab_df, brk_df, stimtype_list, n_perm):
    """
    Returns dataframe of the mean |USI| values over ROIs for each session, the
    p-values for each session pair comparison, and the Bonferroni-corrected
    significance level of each.

    Parameters
    ----------
    gab_df : Pandas DataFrame
        Dataframe of USIs for Gabor stimulus
    brk_df : Pandas DataFrame
        Dataframe of USIs for visual flow (bricks) stimulus
    stimtype_list : list of strings
        Contains stimuli over which to compute values
    n_perm : number
        Number of permutations to use to compute p-values for intersession
        comparisons

    Returns
    -------
    df : Pandas DataFrame
        Dataframe with mean |USI| session values and comparisons

        Columns:
            'layer', 'compartment', 'stimtype' :
                layer, compartment, and stimulus type, respectively
            'usi_abs_mn_<session n>' :
                <|USI|>, where the average is taken over ROIs
            'usi_abs_mn_raw_p__<sess m>_<sess n>' :
                raw p-value of <|USI|> between sessions m and n
            'usi_sig_for_abs_mn__<sess m>_<sess n>' :
                (Bonferroni-) corrected significance level for <|USI|> between
                sessions m and n
    """

    # declarations/initializations
    sess_ns = [1,2,3]
    bonf_n = 12
    alpha_001 = 0.001/bonf_n
    alpha_01  = 0.01/bonf_n
    alpha_05  = 0.05/bonf_n
    df = pd.DataFrame(
        columns = ['layer', 'compartment', 'stimtype',
                   'usi_abs_mn_1', 'usi_abs_mn_2', 'usi_abs_mn_3',
                   'usi_abs_mn_raw_p__1_2', 'usi_abs_mn_raw_p__2_3',
                   'usi_abs_mn_raw_p__1_3',
                   'usi_sig_for_abs_mn__1_2', 'usi_sig_for_abs_mn__2_3',
                   'usi_sig_for_abs_mn__1_3'])
    df_dict = {'gabors':gab_df, 'bricks':brk_df}
    idx = -1

    # Loop through stimulus types, compartments/layers and sessions
    for stimtype in stimtype_list:
        for compartment in ['dend', 'soma']:
            for layer in ['L2/3', 'L5']:
                # declarations/initializations
                usi_df = df_dict[stimtype]
                idx+=1
                print(layer, compartment)
                # Set masks to desired layer/compartment
                mask0 = usi_df['layer']==layer
                mask1 = usi_df['compartment']==compartment
                usis = []
                usi_abs_mn = []
                usi_abs_mn_pval = []
                usi_sig_abs_mn_pval = []
                # Obtain USIs and calculate the mean over ROIs of their absolute
                #  values for each session
                for sess_n in sess_ns:
                    usis.append(np.hstack(usi_df[(mask0 & mask1)]
                                          ['sess_{}_usi'.format(sess_n)].
                                          values))
                    usi_abs_mn.append(np.nanmean(np.abs(usis[-1])))
                # Compute the p-value for the mean of the absolute values by
                #  shuffling all session pairs
                usi_abs_mn_pval = compute_usi_abs_mn_pval(usis, usi_abs_mn,
                    n_perm, metric='abs_mean')
                # Determine the significance levels
                for pval in usi_abs_mn_pval:
                    if pval <= alpha_001:
                        pval_str = '<= 0.001'
                    elif pval <= alpha_01:
                        pval_str = '<= 0.01'
                    elif pval <= alpha_05:
                        pval_str = '<= 0.05'
                    else:
                        pval_str = 'False'
                    usi_sig_abs_mn_pval.append(pval_str)
                df.loc[idx] = [layer, compartment, stimtype,
                               usi_abs_mn[0], usi_abs_mn[1], usi_abs_mn[2],
                               usi_abs_mn_pval[0], usi_abs_mn_pval[1],
                               usi_abs_mn_pval[2],
                               usi_sig_abs_mn_pval[0], usi_sig_abs_mn_pval[1],
                               usi_sig_abs_mn_pval[2]]

    return df

#############################################

def compute_usi_abs_mn_pval(usis, usi_scalar, n_perm, metric='abs_mean'):
    """
    Compute p-values of <|USI|> across sessions by shuffling pairs

    Parameters
    ----------
    usis : list of 1-D arrays of numbers
        usis[i] = 1-D array of USIs for session i for all mice in a particular
        layer/compartment/stimulus type
    usi_scalar : list of numbers
        Specifies the scalar metric for a given session.  Can be, e.g.,
        variance.  See options below in 'metric'.  E.g., usi_scalar[i] can be
        mean of USIs for session i
    n_perm : number
        Number of permutations
    metric : string; optional, default = 'abs_mean'
        Options: 'var', 'mean', 'abs_mean'. Scalar metric of USIs used in
        'usi_scalar', above. Used to determine what scalar metric to use with
        shuffled USIs.

    Returns
    -------
    usi_scalar_pval : list of numbers
        P-values for chosen scalar (var, mean, |mean|).  In order, compares
        sessions 1&2, 2&3, and 1&3
    """

    # declarations/initializations
    diff_distro = [];
    usi_scalar_pval = [];
    # Compute differences:
    usi_diff = [usi_scalar[0]-usi_scalar[1],
                usi_scalar[1]-usi_scalar[2],
                usi_scalar[0]-usi_scalar[2]]
    for _ in range(int(n_perm)):
        usi_perm = []
        # Make into array, with each column corresponding to a different session
        usi_arr = np.asarray(usis).transpose()
        # Vector of which session (column) to take from usi_arr:
        rand_sess = \
            np.random.choice(range(usi_arr.shape[1]),
                             size=usi_arr.shape[0], replace=True)
        roi_idx = range(usi_arr.shape[0])
        # Use above to select 2 of the 3 random sessions from usi_arr
        usi_perm.append(usi_arr[roi_idx, rand_sess])
        direc = np.random.choice([-1,1], size=usi_arr.shape[0], replace=True)
        usi_perm.append(usi_arr[roi_idx, np.mod(rand_sess + direc, 3)])
        # Compute a differences of metrics from the distributions
        if metric=='var':
            diff_distro.append(np.nanvar(usi_perm[0]) - np.nanvar(usi_perm[1]))
        elif metric=='mean':
            diff_distro.append(np.nanmean(usi_perm[0]) -
                               np.nanmean(usi_perm[1]))
        elif metric=='abs_mean':
            diff_distro.append(np.nanmean(np.abs(usi_perm[0])) -
                               np.nanmean(np.abs(usi_perm[1])))
    diff_distro = np.asarray(diff_distro)

    for i in range(len(usi_scalar)):
        usi_scalar_pval.append(compute_2_tailed_pval(usi_diff[i], diff_distro))

    return usi_scalar_pval

#############################################

def compute_2_tailed_pval(value, distro):
    """
    Compute 2-tailed p-value

    Parameters
    ----------
    value : number
        value for which to ascertain the p-value
    distro : 1-D array of numbers
        computed distribution against which to compare 'value' to ascertain the
        p-value

    Returns
    -------
    pval : number
        computed 2-tailed p-value
    """
    distro = np.asarray(distro)
    n_perm_idcs = distro.size
    # Form array of indices where 'value' is no greater than the distribution
    #  values
    perm_idcs_larger = np.where(distro >= value)[0]
    # The probability is then the raio of the length of this array to the
    #  distribution size
    pval = len(perm_idcs_larger) / n_perm_idcs
    # 2-tailed correction
    if pval > 0.5:
        pval = 1-pval
    pval *= 2
    if np.isnan(value):
        pval = np.nan

    return pval

#############################################

def make_usi_abs_frac_chng_df(gab__tracked_roi_usi_df, brk__tracked_roi_usi_df,
                              n_perm, n_bstrap):
    """
    Returns dataframe containing fractional chng from session 1 to 3 of the
    <|USIs|> over ROIs for Gabors and visual flow, along with their
    uncertainties (bootstrapped stdev) and p-value of the difference between the
    Gabor and visual flow chng for each layer/compartment

    Parameters
    ----------
    gab__tracked_roi_usi_df : Pandas DataFrame
        Dataframe containing tracked usi for each
        mouse/layer/compartment/session for the Gabor stimulus
    brk__tracked_roi_usi_df : Pandas DataFrame
        Dataframe containing tracked usi for each
        mouse/layer/compartment/session for the bricks stimulus
    n_perm : number
        Number of permutations to perform to compare Gabors against visual flow
    n_bstrap : number
        Number of resamplings with replacement to do

    Returns
    -------
    usi_chng_df : Pandas DataFrame
        Dataframe with <|USIs|> (over ROIs) fractional chng from sess 1 to 3,
        and comparisons across stimuli

        Columns:
            'layer', compartment, 'sess_compare' :
                layer, compartment, and sessions compared, respectively
            'gab_mn_abs_frac_chng' :
                Gabor fractional chng in <|USIs|> from sess 1 to 3
            'brk_mn_abs_frac_chng' :
                Bricks fractional chng in <|USIs|> from sess 1 to 3
            'pval_raw' :
                Raw p-value comparing Gabor and brick fractional changes
            'gab_bstrap_std' :
                Bootstrapped stdev for 'gab_mn_abs_frac_chng'
            'brk_bstrap_std' :
                Bootstrapped stdev for 'brk_mn_abs_frac_chng'
    """

    # declarations/initializations
    sess_compare = [1,3]
    compartment_list = ['dend', 'soma', 'all']
    layer_list = ['L2/3', 'L5', 'all']
    usi_chng_df = pd.DataFrame(columns=['layer', 'compartment', 'sess_compare',
                                        'gab_mn_abs_frac_chng',
                                        'brk_mn_abs_frac_chng',
                                        'pval_raw', 'gab_bstrap_std',
                                        'brk_bstrap_std'])
    idx = -1
    # Loop through compartments and layers to get data for each
    for layer, compartment in it.product(layer_list, compartment_list):
        if layer == 'all' and compartment != 'all':
            continue
        if layer != 'all' and compartment == 'all':
            continue
        print(layer, compartment)
        idx+=1
        # Obtain fractional changes and ROI pairs across sessions and
        #  layers/compartments if layer is 'all' for statistical testing
        gab_frac_chng, brk_frac_chng, gab_usi_lay_comp, brk_usi_lay_comp = \
            get_usi_abs_frac_chng(gab__tracked_roi_usi_df,
                                  brk__tracked_roi_usi_df,
                                  layer, compartment, sess_compare)
        # Compute p-value
        pval = \
            compute_usi_abs_frac_chng_pval(gab_usi_lay_comp, brk_usi_lay_comp,
                                           gab_frac_chng, brk_frac_chng,
                                           n_perm=n_perm)
        # Compute uncertainty
        gab_std, brk_std = \
            compute_usi_abs_frac_chng_err(gab_usi_lay_comp, brk_usi_lay_comp,
                                          n_bstrap=n_bstrap)
        # Add data to dataframe
        usi_chng_df.loc[idx] = [layer, compartment, sess_compare,
                                gab_frac_chng, brk_frac_chng,
                                pval, gab_std, brk_std]

    return usi_chng_df

#############################################

def get_usi_abs_frac_chng(gab__tracked_roi_usi_df, brk__tracked_roi_usi_df,
                          layer, compartment, sess_compare):
    '''
    Get absolute fractional changes of <|USIs|> (over ROIs) from sess
    sess_compare[0] (usually sess 1) to sess_compare[1] (usually sess 3) for
    Gabors and visual flow, along with the associated USI pairs for
    Gabors/bricks for each session in order to shuffle stimulus labels to later
    obtain p-values when comparing Gabors and visual flow.

    Parameters
    ----------
    gab__tracked_roi_usi_df : Pandas DataFrame
        Dataframe containing tracked usi for each
        mouse/layer/compartment/session for the Gabor stimulus
    brk__tracked_roi_usi_df : Pandas DataFrame
        Dataframe containing tracked usi for each
        mouse/layer/compartment/session for the bricks stimulus
    layer : string
        Layer for which to obtain data
    compartment : string
        Compartment for which to obtain data
    sess_compare : 1-D arraylike of numbers
        Sessions for which to obtain data

    Returns
    -------
    gab_abs_usi_frac_chng : 1-D list of numbers
        Absolute fractional change from sess_compare[0] to sess_compare[1] in
        <|USIs|> over ROIs for Gabors. Note, only 1 number in list unless
        looking at all layers/compartments
    brk_usi_abs_frac_chng : 1-D list of numbers
        Absolute fractional change from sess_compare[0] to sess_compare[1] in
        <|USIs|> over ROIs for visual flow. Note, only 1 number in list unless
        looking at all layers/compartments
    gab_usi_lay_comp : 3-D list of numbers
        Outer: Layers/compartments (just 1 entry if layer not 'all').
        Next: Sessions. Inner: USIs for tracked ROIs
    brk_usi_lay_comp : 3-D list of numbers
        Outer: Layers/compartments (just 1 entry if layer not 'all').
        Next: Sessions. Inner: USIs for tracked ROIs
    '''

    # declarations/initializations
    layer_list = ['L2/3', 'L5']
    compartment_list = ['dend', 'soma']
    gab_abs_usi_frac_chng = []
    brk_usi_abs_frac_chng = []
    gab_usi_lay_comp = []
    brk_usi_lay_comp = []

    # Unless we're processing all of the layers/compartments together, no reason
    #  to loop
    if layer != 'all':
        layer_list = [layer]
        compartment_list = [compartment]
    # Loop over chosen layers/compartments (trivial loops unless looking at all
    #  layers/compartments)
    for layer, compartment in it.product(layer_list, compartment_list):
        # declarations/initializations
        gab_usi_pairs = []
        brk_usi_pairs = []
        # Obtain data for each session that we want to compare
        for sess_n in sess_compare:
            # Obtain sub-dataframes limited to chosen layer/compartment
            mask0 = gab__tracked_roi_usi_df['layer']==layer
            mask1 = gab__tracked_roi_usi_df['compartment']==compartment
            gab_df = gab__tracked_roi_usi_df[mask0 & mask1]
            mask0 = brk__tracked_roi_usi_df['layer']==layer
            mask1 = brk__tracked_roi_usi_df['compartment']==compartment
            brk_df = brk__tracked_roi_usi_df[mask0 & mask1]
            # Obtain the USI values for Gabor / bricks for each session
            gab_usi_pairs.append(np.hstack(gab_df['sess_{}_usi'.format(sess_n)].
                                           values))
            brk_usi_pairs.append(np.hstack(brk_df['sess_{}_usi'.format(sess_n)].
                                           values))
        # Amalgamate USIs across layers/compartments (for 'all')
        gab_usi_lay_comp.append(gab_usi_pairs)
        brk_usi_lay_comp.append(brk_usi_pairs)
        # compute absolute fractional change across sessions for <|USI|> over
        #  ROIs for Gabors
        mean0 = np.nanmean(np.abs(gab_usi_pairs[0]))
        mean1 = np.nanmean(np.abs(gab_usi_pairs[1]))
        gab_abs_usi_frac_chng.append(np.abs((mean1-mean0)/mean0))
        # compute absolute fractional change across sessions for <|USI|> over
        #  ROIs for visual flow
        mean0 = np.nanmean(np.abs(brk_usi_pairs[0]))
        mean1 = np.nanmean(np.abs(brk_usi_pairs[1]))
        brk_usi_abs_frac_chng.append(np.abs((mean1-mean0)/mean0))

    return gab_abs_usi_frac_chng, brk_usi_abs_frac_chng, gab_usi_lay_comp, \
           brk_usi_lay_comp

#############################################

def compute_usi_abs_frac_chng_pval(gab_usi_lay_comp, brk_usi_lay_comp,
                                   gab_frac_chng, brk_frac_chng, n_perm):
    """
    Compute p-values for <|USI|> (over ROIs) |fractional changes| across
    sessions by shuffling stimulus labels

    Parameters
    ----------
    gab_usi_lay_comp : 3-D list of numbers
        Outer: Layers/compartments (just 1 entry if layer not 'all').
        Next: Sessions. Inner: USIs for tracked ROIs
    brk_usi_lay_comp : 3-D list of numbers
        Outer: Layers/compartments (just 1 entry if layer not 'all').
        Next: Sessions. Inner: USIs for tracked ROIs
    gab_abs_usi_frac_chng : 1-D list of numbers
        Absolute fractional change from sess_compare[0] to sess_compare[1] in
        <|USIs|> over ROIs for Gabors. Note, only 1 number in list unless
        looking at all layers/compartments
    brk_usi_abs_frac_chng : 1-D list of numbers
        Absolute fractional change from sess_compare[0] to sess_compare[1] in
        <|USIs|> over ROIs for visual flow. Note, only 1 number in list unless
        looking at all layers/compartments
    n_perm : number
        Number of shuffles to perform

    Returns
    -------
    pval : number
        P-value of difference between |fractional changes| across sessions of
        <|USI|> over ROIs for Gabors vs. bricks
    """

    # declarations/initializations
    n_perm = int(n_perm)
    gab_brk_distro = []
    # Permutations loop
    for _ in range(n_perm):
        # declarations/initializations
        gab_perm_pair_mn = []
        gab_frac_chng_perm = []
        brk_perm_pair_mn = []
        brk_frac_chng_perm = []
        # Get USIs for each layer/compartment (just one iteration if layer not
        #  'all')
        for i in range(len(gab_usi_lay_comp)):
            gab_usi_pair = gab_usi_lay_comp[i]
            brk_usi_pair = brk_usi_lay_comp[i]
            # Shuffle stimulus labels for each session
            for j in range(len(gab_usi_pair)):
                # Arrange USIs as columns, one column per stimulus type
                arr = np.vstack((gab_usi_pair[j],brk_usi_pair[j])).transpose()
                # Index ROIs
                roi_idx = range(arr.shape[0])
                # Randomly pick 1st or 2nd column of array
                col = np.random.choice(range(arr.shape[1]), size=len(roi_idx),
                                       replace=True)
                # Obtain <|USI|> (over ROIs) with shuffled stimulus labels
                gab_perm_pair_mn.append(np.nanmean(np.abs(arr[roi_idx, col])))
                brk_perm_pair_mn.append(np.nanmean(np.abs(arr[roi_idx,
                                                          np.mod(col+1,2)])))
            # Compute |fractional change| across session for shuffled <|USI|>
            gab_frac_chng_perm.append(
                np.abs((gab_perm_pair_mn[1]-gab_perm_pair_mn[0]) /
                       gab_perm_pair_mn[0]))
            brk_frac_chng_perm.append(
                np.abs((brk_perm_pair_mn[1]-brk_perm_pair_mn[0]) /
                       brk_perm_pair_mn[0]))
        # Add difference of distributions (averaged over layers/compartments if
        #  layer is 'all') to shuffled distribution
        gab_brk_distro.append(
            np.mean(gab_frac_chng_perm) - np.mean(brk_frac_chng_perm))

    gab_brk_distro = np.asarray(gab_brk_distro)
    # Compute difference and p-value
    diff = np.mean(gab_frac_chng) - np.mean(brk_frac_chng)
    pval = compute_2_tailed_pval(diff, gab_brk_distro)

    return pval

#############################################

def compute_usi_abs_frac_chng_err(gab_usi_lay_comp, brk_usi_lay_comp,
                              n_bstrap=1e3):

    """
    Compute uncertainty for <|USI|> (over ROIs) |fractional change|

    Parameters
    ----------
    gab_usi_lay_comp : 3-D list of numbers
        Outer: Layers/compartments (just 1 entry if layer not 'all').
        Next: Sessions. Inner: USIs for tracked ROIs
    brk_usi_lay_comp : 3-D list of numbers
        Outer: Layers/compartments (just 1 entry if layer not 'all').
        Next: Sessions. Inner: USIs for tracked ROIs
    n_bstrap : number; optional, default = 1e3
        Number of resamplings with replacement to do

    Returns
    -------
    gab_bstrap_std : number
        Standard deviation of distribution of <|USI|> (over ROIs)
        |fractional change| obtained via resampling with replacement from each
        session's USIs for Gabors.
    brk_bstrap_std :
        Standard deviation of distribution of <|USI|> (over ROIs)
        |fractional change| obtained via resampling with replacement from each
        session's USIs for Gabors.
    """

    # declarations/initializations
    n_bstrap = int(n_bstrap)
    gab_brk_distro = []
    gab_bstrap_pair_mn = []
    brk_bstrap_pair_mn = []
    # Loop over layers/compartments (1 iteration if layer not 'all')
    for i in range(len(gab_usi_lay_comp)):
        # Obtain USI pairs (across sessions) for each stimulus for
        #  layer/compartment
        gab_usi_pairs = gab_usi_lay_comp[i]
        brk_usi_pairs = brk_usi_lay_comp[i]
        # Loop over sessions
        for j in range(len(gab_usi_pairs)):
            # Sample with replacement
            gab_bstrap = \
                np.random.choice(gab_usi_pairs[j], size=(len(gab_usi_pairs[j]),
                                 n_bstrap), replace=True)
            brk_bstrap = \
                np.random.choice(brk_usi_pairs[j], size=(len(brk_usi_pairs[j]),
                                 n_bstrap), replace=True)
            # Compute average over ROIs of |USIs|
            gab_bstrap_pair_mn.append(np.mean(np.abs(gab_bstrap), axis=0))
            brk_bstrap_pair_mn.append(np.mean(np.abs(brk_bstrap), axis=0))
        # Compute |fractional change| of <|USI|>
        gab_frac_chng = \
            np.abs((gab_bstrap_pair_mn[1]-gab_bstrap_pair_mn[0]) /
                   gab_bstrap_pair_mn[0])
        brk_frac_chng = \
            np.abs((brk_bstrap_pair_mn[1]-brk_bstrap_pair_mn[0]) /
                   brk_bstrap_pair_mn[0])
    # Uncertainty
    gab_bstrap_std = np.std(gab_frac_chng)
    brk_bstrap_std = np.std(brk_frac_chng)

    return gab_bstrap_std, brk_bstrap_std

#############################################

def make_usi_corr_df(tracked_roi_usi_df, usi_corr_df, stimtype, usi_base_sess=1,
                     n_perm=1e4, n_bstrap=1e4):
    """
    Returns a dataframe that contains, for each compartment and layer, the raw
    and residual USI correlations (i.e., correlations between usi_{day m} vs.
    Delta(USI) = usi_{day m+1} - usi_{day m} for m = 1 or 2), the bounds of the
    standard deviation for the residual correlations and 95% confidence interval
    for the null distribution of the residual correlations, the raw p-value
    (multiply by 8 = Bonferroni N for corrected p-value), and the corrected
    significance level.

    For inspection convenience, the variables used to determine the
    correlations, the raw shuffled correlation distribution, and the raw
    bootstrapped distribution are included as well.

    Parameters
    ----------
    tracked_roi_usi_df : Pandas DataFrame
        Dataframe with tracked ROI USIs and statistics
    usi_corr_df : Pandas DataFrame
        Dataframe with correlation information of USIs across sessions.  See
        below under 'Returns'
    stimtype : string
        Stimulus type ('gabors' or 'bricks')
    usi_base_sess : number; optional, default = 1
        Session to obtain USIs and against which to compare delta(USI) with
        usi_base_sess+1
    n_perm : number; optional, default = 1e4
        Number of permutations to perform to compare Gabors against visual flow
    n_bstrap : number; optional, default =1e4
        Number of resamplings with replacement to do

    Returns
    -------
    usi_corr_df (pd DataFrame) : dataframe that contains, for each
        layer/compartment, data on the correlations between the USI in 1 session
        and change in USI by the next session

        Columns:
            stimtype, layer, compartment :
                stimulus type (gabor / bricks), layer, and compartment,
                respectively
            usi_base_sess :
                session USI obtained from difference in USIs taken from this and
                subsequent sessions
            corr_raw :
                raw correlation between USI and Delta(USI)
            corr_resid :
                normalized residual correlation
            corr_resid_low_sd, corr_resid_high_sd :
                interval of bootstrapped correlation std dev (/err)
            corr_resid_null_low_ci, corr_resid_null_high_ci :
                interval of 95% CI of null resid corr
            pval_raw :
                raw p-value of correlation.  Corrected p-value = 8*pval_raw,
                where 8 = the Bonferroni N
            sig_correc :
                (bonferroni-) corrected significance level
            usi, delta_usi :
                USI and Delta(USI).  Correlation of these variables =
                'corr_raw'.  For inspection if desired
            corr_raw_distro :
                raw correlation distribution obtained by shuffling session
                labels.  Used to obtain p-value. For inspection if desired
            bstrap_raw_distro :
                raw correlation bootstrapped distribution, used to obtain
                standard error. For inspection if desired
    """

    # declarations/initializations
    bonf_n = 8
    alpha_001 = 0.001/bonf_n
    alpha_01  = 0.01/bonf_n
    alpha_05  = 0.05/bonf_n

    idx = usi_corr_df.shape[0]
    layers = {'L23-Cux2':'L2/3', 'L5-Rbp4':'L5'}
    lc = list(it.product(layers.values(), ['dend', 'soma']))

    # Loop through layers and compartments
    for i, (l,c) in enumerate(lc):
        print(c, l)

        # declarations/initializations
        usi                   = []
        corr_raw_distro       = []
        error_corr_raw_distro = []
        mask0 = tracked_roi_usi_df['compartment']==c
        mask1 = tracked_roi_usi_df['layer']==l
        usi.append( np.hstack(tracked_roi_usi_df[mask0 & mask1]['sess_1_usi']) )
        usi.append( np.hstack(tracked_roi_usi_df[mask0 & mask1]['sess_2_usi']) )
        usi.append( np.hstack(tracked_roi_usi_df[mask0 & mask1]['sess_3_usi']) )

        # Find correlation between USI and delta(USI)
        x = usi[usi_base_sess-1]
        x2 = usi[usi_base_sess]
        y = x2 - x
        corr_raw = scist.pearsonr(x, y)[0]

        # Permuation and bootstrap loop
        roi_idx = range(x.size)
        loop_length = int(max(n_perm, n_bstrap))
        for i_perm in range(loop_length):
            # Shuffle session labels for USIs
            if i_perm < n_perm:
                # ROIs x sessions
                arr = np.asarray(usi).transpose()
                # Row 1: roi indexes; row 2: randomly chosen session
                sess_col = np.random.choice([0,1,2], x.size)
                # USIs from randomly chosen session
                usi1_perm = arr[roi_idx, sess_col]
                # USIs from randomly chosen remaining session (periodic boundary
                #  conditions)
                direction  = np.random.choice([-1,1], x.size)
                usi2_perm = arr[roi_idx,np.mod(sess_col+direction,3)]
                # compute pearson from shuffled USIs and add to shuffled distro
                x_perm = usi1_perm
                y_perm = usi2_perm - usi1_perm
                corr_raw_distro.append(scist.pearsonr(x_perm, y_perm)[0])
            # Bootstrapped standard deviation (~ standard error of correlation)
            if i_perm < n_bstrap:
                # sample from ROI indices with replacement
                samp_idxs = np.random.choice(roi_idx, x.size, replace=True)
                x_bstrap = x[samp_idxs]
                y_bstrap = y[samp_idxs]
                error_corr_raw_distro.append(scist.pearsonr(x_bstrap, y_bstrap)
                                             [0])

        # Compute correlation p-value, normalized value, and standard deviation
        corr_raw_distro = np.asarray(corr_raw_distro)
        error_corr_raw_distro = np.asarray(error_corr_raw_distro)
        pval_raw = compute_2_tailed_pval(corr_raw, corr_raw_distro)
        if pval_raw <= alpha_001:
            pval_correc_str = '<= 0.001'
        elif pval_raw <= alpha_01:
            pval_correc_str = '<= 0.01'
        elif pval_raw <= alpha_05:
            pval_correc_str = '<= 0.05'
        else:
            pval_correc_str = 'False'
        corr_med = np.median(corr_raw_distro)
        corr_resid = corr_raw - corr_med
        sigma = np.nanstd(error_corr_raw_distro)
        # Normalize to be in [-1, 1] by dividing difference by distance between
        #  median and left (if corr_raw < corr_med) or right
        #  (if corr_raw > corr_med) bound
        if corr_resid < 0:
            corr_resid /= (corr_med+1)
            sigma     /= (corr_med+1)
        elif corr_resid > 0:
            corr_resid /= (1-corr_med)
            sigma     /= (1-corr_med)
        corr_resid_low_sd  = corr_resid - sigma
        corr_resid_high_sd = corr_resid + sigma

        # Compute normalized correlation CI
        corr_resid_low_arg = int(np.round(alpha_05*corr_raw_distro.size))
        corr_resid_low_nom = np.sort(corr_raw_distro)[corr_resid_low_arg]
        corr_resid_high_arg = int(np.round((1-alpha_05)*corr_raw_distro.size))
        corr_resid_high_nom = np.sort(corr_raw_distro)[corr_resid_high_arg]
        corr_resid_null_low_ci = (corr_resid_low_nom - corr_med) / (corr_med+1)
        corr_resid_null_high_ci = (corr_resid_high_nom - corr_med) / \
                                  (1-corr_med)

        # Add row to dataframe
        usi_corr_df.loc[idx] = \
            [stimtype, l, c, usi_base_sess,
            corr_raw, corr_resid, corr_resid_low_sd, corr_resid_high_sd,
            corr_resid_null_low_ci, corr_resid_null_high_ci,
            pval_raw, pval_correc_str,
            x, y, corr_raw_distro, error_corr_raw_distro]
        idx+=1

    return usi_corr_df

#############################################

def usi_corr_permutation(tracked_roi_usi_df, stimtype, l, c, usi_base_sess=1):
    """
    Returns the median of the shuffled distribution of correlations between
    USI and Delta(USI)

    Parameters
    ----------
    tracked_roi_usi_df : Pandas DataFrame
        Dataframe with tracked ROI USIs and statistics
    stimtype : string
        Stimulus type ('gabors' or 'bricks')
    l : string
        Layer from which to get data
    c : string
        Compartment from which to get data
    usi_base_sess : number; optional, default = 1
        Session to obtain USIs and against which to compare delta(USI) with
        usi_base_sess+1

    Returns
    -------
    median_array : 2-D array of numbers
        Array with data with a correlation closest to median of a shuffled
        distro.  Size: ROIs x 2.  1st column: USIs; 2nd column: Delta(USI)
    """

    print(l, c)

    # declarations/initializations
    usi                   = []
    corr_raw_distro       = []
    mask0 = tracked_roi_usi_df['compartment']==c
    mask1 = tracked_roi_usi_df['layer']==l
    usi.append( np.hstack(tracked_roi_usi_df[mask0 & mask1]['sess_1_usi']) )
    usi.append( np.hstack(tracked_roi_usi_df[mask0 & mask1]['sess_2_usi']) )
    usi.append( np.hstack(tracked_roi_usi_df[mask0 & mask1]['sess_3_usi']) )

    # Find correlation between USI and delta(USI)
    x = usi[usi_base_sess-1]
    x2 = usi[usi_base_sess]
    y = x2 - x
    corr_raw = scist.pearsonr(x, y)[0]

    # Permuation loop
    roi_idx = range(x.size)
    for i_perm in range(int(1e4)):
        # Shuffle session labels for USIs
        # ROIs x sessions
        arr = np.asarray(usi).transpose()
        # Row 1: roi indexes; row 2: randomly chosen session
        sess_col = np.random.choice([0,1,2], x.size)
        # USIs from randomly chosen session
        usi1_perm = arr[roi_idx, sess_col]
        # USIs from randomly chosen remaining session (periodic boundary
        #  conditions)
        direction  = np.random.choice([-1,1], x.size)
        usi2_perm = arr[roi_idx,np.mod(sess_col+direction,3)]
        # compute pearson from shuffled USIs and add to shuffled distro
        x_perm = usi1_perm
        y_perm = usi2_perm - usi1_perm
        corr_raw_distro.append(scist.pearsonr(x_perm, y_perm)[0])

    corr_med = np.median(corr_raw_distro)
    corr_resid = corr_raw - corr_med

    print('Corr raw, median = ', corr_raw, corr_med)

    ###############

    # Second run to grab the permutation data with correlation
    #  closest to median from above

    # declarations/initializations
    corr_raw_distro       = []
    # Just make sure it's farther away from median corr than anything possible
    prev_pearson = -2

    for i_perm in range(int(1e4)):
        # Shuffle session labels for USIs
        # ROIs x sessions
        arr = np.asarray(usi).transpose()
        # Row 1: roi indexes; row 2: randomly chosen session
        sess_col = np.random.choice([0,1,2], x.size)
        # USIs from randomly chosen session
        usi1_perm = arr[roi_idx, sess_col]
        # USIs from randomly chosen remaining session (periodic boundary
        #  conditions)
        direction  = np.random.choice([-1,1], x.size)
        usi2_perm = arr[roi_idx,np.mod(sess_col+direction,3)]
        # Compute pearson from shuffled USIs and add to shuffled distro
        x_perm = usi1_perm
        y_perm = usi2_perm - usi1_perm
        corr_raw_distro.append(scist.pearsonr(x_perm, y_perm)[0])
        # Find data for median of distro
        if np.abs(corr_med - corr_raw_distro[-1]) < np.abs(corr_med - prev_pearson):
            x_keep_perm = x_perm
            y_keep_perm = y_perm
            prev_pearson = corr_raw_distro[-1]

    # Median data
    median_array = np.stack((x_keep_perm, y_keep_perm), axis=1)

    return median_array


import numpy as np
import scipy.sparse as sp
import properties

from .. import utils
from .base import BaseSimilarityMeasure

class LinearCorrespondence(BaseSimilarityMeasure):
    """
    The petrophysical linear constraint for joint inversions.

    ..math::
        \\phi_c({\\mathbf m}_{\\mathbf1},{\\mathbf m}_{\\mathbf2})=\\lambda\\sum_{i=1}^M
        (k_1*m_1 + k_2*m_2 + k_3)

    Assuming that we are working with two models only.

    """

    coefficients = properties.Array(
        "coefficients for the linear relationship between parameters",
        shape=(3,),
        default=np.array([1.0, -1.0, 0.0]),
    )

    def relation(self, model):
        """
        Computes the values of petrophysical linear relationship between two different
        geophysical models.

        The linear relationship is defined as:

        f(m1, m2)  = k1*m1 + k2*m2 + k3

        :param numpy.ndarray model: stacked array of individual models
                                    np.c_[model1, model2,...]

        :rtype: float
        :return: linearly related petrophysical values of two different models,
                  dimension: M by 1, :M number of model parameters.

        """
        m1, m2 = self.wire_map * model
        k1, k2, k3 = self.coefficients

        return k1 * m1 + k2 * m2 + k3

    def __call__(self, model):
        """
        Computes the sum of values of petrophysical linear relationship
        between two different geophysical models.

        :param numpy.ndarray model: stacked array of individual models
                                    np.c_[model1, model2,...]

        :rtype: float
        :return: a scalar value.
        """

        result = self.relation(model)
        return 0.5 * result.T @ result

    def deriv(self, model):
        """Computes the Jacobian of the coupling term.

        :param list of numpy.ndarray ind_models: [model1, model2,...]

        :rtype: numpy.ndarray
        :return: result: gradient of the coupling term with respect to model1, model2,
                 :dimension 2M by 1, :M number of model parameters.
        """
        k1, k2, k3 = self.coefficients
        r = self.relation(model)
        dc_dm1 = k1 * r
        dc_dm2 = k2 * r

        result = np.r_[dc_dm1, dc_dm2]

        return result

    def deriv2(self, model, v=None):
        """Computes the Hessian of the linear coupling term.

        :param list of numpy.ndarray ind_models: [model1, model2, ...]
        :param numpy.ndarray v: vector to be multiplied by Hessian
        :rtype: scipy.sparse.csr_matrix if v is None
                numpy.ndarray if v is not None
        :return Hessian matrix: | h11, h21 | :dimension 2M*2M.
                                |          |
                                | h12, h22 |
        """

        k1, k2, k3 = self.coefficients
        if v is not None:
            v1, v2 = self.wire_map * v
            p1 = k1**2 * v1 + k2 * k1 * v2
            p2 = k2 * k1 * v1 + k2**2 * v2
            return np.r_[p1, p2]
        else:
            n = self.regmesh.nC
            A = utils.sdiag(np.ones(n) * (k1**2))
            B = utils.sdiag(np.ones(n) * (k2**2))
            C = utils.sdiag(np.ones(n) * (k1 * k2))
            return sp.bmat([[A, C], [C, B]], format="csr")


from SimpleCV.base import *
from SimpleCV.Features.Features import Feature, FeatureSet
from SimpleCV.Color import Color
from SimpleCV.ImageClass import Image
from SimpleCV.Features.Detection import ShapeContextDescriptor
import math
import scipy.stats as sps
try:
    from sklearn import neighbors
except:
    print "Need scikits learn installed"
    #exit()

"""
Classify an object based on shape context
"""
class ShapeContextClassifier():

    def  __init__(self,images,labels):
        """
        Create a shape context classifier.

        * *images* - a list of input binary images where the things
          to be detected are white.

        * *labels* - the names of each class of objects.
        """
        self.imgMap = {}
        self.ptMap = {}
        self.descMap = {}
        self.knnMap = {}
        self.blobCount = {}
        self.labels = labels
        self.images = images
        import warnings
        warnings.simplefilter("ignore")
        for i in range(0,len(images)):
            print "precomputing " + images[i].filename
            self.imgMap[labels[i]] = images[i]

            pts,desc,count  = self._image2FeatureVector(images[i])
            self.blobCount[labels[i]] = count
            self.ptMap[labels[i]] = pts
            self.descMap[labels[i]] = desc
            knn = neighbors.KNeighborsClassifier()
            knn.fit(desc,range(0,len(pts)))
            self.knnMap[labels[i]] = knn

    def _image2FeatureVector(self,img):
        """
        generate a list of points, SC descriptors, and the count of points
        """
        #IMAGES MUST BE WHITE ON BLACK!
        fulllist = []
        raw_descriptors = []
        blobs = img.findBlobs(minsize=50)
        count = 0
        if( blobs is not None ):
            count = len(blobs)
            for b in blobs:
                fulllist += b._filterSCPoints()
                raw_descriptors = blobs[0]._generateSC(fulllist)
        return fulllist,raw_descriptors,count

    def _getMatch(self,model_scd,test_scd):
        correspondence,distance = self._doMatching(model_scd,test_scd)
        return self._matchQuality(distances)

    def _doMatching(self,model_name,test_scd):
        myPts = len(test_scd)
        otPts = len(self.ptMap[model_name])
        # some magic metric that keeps features
        # with a lot of points from dominating
        #metric = 1.0 + np.log10( np.max([myPts,otPts])/np.min([myPts,otPts])) # <-- this could be moved to after the sum
        otherIdx = []
        distance = []
        import warnings
        warnings.simplefilter("ignore")
        results = []
        for sample in test_scd:
            best = self.knnMap[model_name].predict(sample)
            idx = best[0] # this is where we can play with k
            scd = self.descMap[model_name][idx]
            temp = np.sqrt(np.sum(((sample-scd)**2)))
            #temp = 0.5*np.sum((sample-scd)**2)/np.sum((sample+scd))
            if( math.isnan(temp) ):
                temp = sys.maxint
            distance.append(temp)
        return [otherIdx,distance]

    def _matchQuality(self,distances):
        #distances = np.array(distances)
        #sd = np.std(distances)
        #x = np.mean(distances)
        #min = np.min(distances)
        # not sure trimmed mean is perfect
        # realistically we should have some bimodal dist
        # and we want to throw away stuff with awful matches
        # so long as the number of points is not a huge
        # chunk of our points.
        #tmean = sps.tmean(distances,(min,x+sd))
        tmean = np.mean(distances)
        std = np.std(distances)
        return tmean,std

    def _buildMatchDict(self,image, countBlobs):
        # we may want to base the count on the num best_matchesber of large blobs
        points,descriptors,count = self._image2FeatureVector(image)
        matchDict = {}
        matchStd = {}
        for key,value in self.descMap.items():
            if( countBlobs and self.blobCount[key] == count ): # only do matching for similar number of blobs
                #need to hold on to correspondences
                correspondence, distances = self._doMatching(key,descriptors)
                result,std = self._matchQuality(distances)
                matchDict[key] = result
                matchStd[key] = std
            elif( not countBlobs ):
                correspondence, distances = self._doMatching(key,descriptors)
                result,std = self._matchQuality(distances)
                matchDict[key] = result
                matchStd[key] = std

        return points,descriptors,count,matchDict, matchStd

    def classify(self,image, blobFilter=True):
        """
        Classify an input image.

        * *image* - the input binary image.
        * *blobFilter* - Do a first pass where you only match objects
          that have the same number of blobs - speeds up computation
          and match quality.
        """
        points,descriptors,count,matchDict,matchStd = self._buildMatchDict(image, blobFilter)
        best = sys.maxint
        best_name = "No Match"
        for k,v in matchDict.items():
            if ( v < best ):
                best = v
                best_name = k

        return best_name, best, matchDict, matchStd

    def getTopNMatches(self,image,n=3, blobFilter = True):
        """
        Classify an input image and return the top n results.

        * *image* - the input binary image.
        * *n* - the number of results to return.
        * *blobFilter* - Do a first pass where you only match objects
          that have the same number of blobs - speeds up computation
          and match quality.
        """
        n = np.clip(n,1,len(self.labels))
        points,descriptors,count,matchDict,matchStd = self._buildMatchDict(image,blobFilter)
        best_matches = list(sorted(matchDict, key=matchDict.__getitem__))
        retList = []
        for k in best_matches:
            retList.append((k,matchDict[k]))
        return retList[0:n], matchDict, matchStd


<filename>quadpy/sphere/stroud.py
# -*- coding: utf-8 -*-
#
import numpy
import sympy

from .helpers import cartesian_to_spherical

from .albrecht_collatz import AlbrechtCollatz
from .mclaren import McLaren
from ..nsphere.stroud1969 import Stroud1969

class Stroud(object):
    """
    <NAME>,
    Approximate Calculation of Multiple Integrals,
    Prentice Hall, 1971.
    """

    def __init__(self, index, symbolic=False):
        self.name = "Stroud_U3({})".format(index)

        if index == "U3 3-1":
            self.set_data(McLaren(1, symbolic=symbolic))
        elif index == "U3 5-1":
            self.set_data(AlbrechtCollatz(1, symbolic=symbolic))
        elif index == "U3 5-2":
            self.set_data(AlbrechtCollatz(2, symbolic=symbolic))
        elif index == "U3 5-3":
            self.set_data(AlbrechtCollatz(3, symbolic=symbolic))
        elif index == "U3 5-4":
            self.set_data(AlbrechtCollatz(4, symbolic=symbolic))
        elif index == "U3 5-5":
            self.set_data(McLaren(2, symbolic=symbolic))
        elif index == "U3 7-1":
            self.set_data(McLaren(3, symbolic=symbolic))
        elif index == "U3 7-2":
            self.set_data(AlbrechtCollatz(5, symbolic=symbolic))
        elif index == "U3 8-1":
            self.set_data(McLaren(4, symbolic=symbolic))
        elif index == "U3 9-1":
            self.set_data(McLaren(5, symbolic=symbolic))
        elif index == "U3 9-2":
            self.set_data(McLaren(6, symbolic=symbolic))
        elif index == "U3 9-3":
            self.set_data(McLaren(7, symbolic=symbolic))
        elif index == "U3 11-1":
            self.set_data(McLaren(8, symbolic=symbolic))
        elif index == "U3 11-2":
            scheme = Stroud1969(3, symbolic=symbolic)
            self.degree = scheme.degree
            self.weights = scheme.weights
            pi = sympy.pi if symbolic else numpy.pi
            self.weights /= 4 * pi
            self.points = scheme.points
            self.azimuthal_polar = cartesian_to_spherical(self.points)
        elif index == "U3 11-3":
            self.set_data(McLaren(9, symbolic=symbolic))
        else:
            assert index == "U3 14-1", "Illegal index {}.".format(index)
            self.set_data(McLaren(10, symbolic=symbolic))
        return

    def set_data(self, scheme):
        self.degree = scheme.degree
        self.weights = scheme.weights
        self.points = scheme.points
        self.azimuthal_polar = scheme.azimuthal_polar
        return


<filename>emotionclassification.py<gh_stars>1-10
#!/usr/bin/env python3
# -*- coding = utf-8 -*-
import os
import sys
import json
import argparse
from statistics import mode

import cv2
import numpy as np
from keras.models import load_model

from data.load_data import get_fer2013_data
from models.model_factory import load_keras_model
from util.constant import fer2013_classes
from util.classifyimgops import apply_offsets
from util.classifyimgops import preprocess_input
from util.info import load_info

ap = argparse.ArgumentParser()
ap.add_argument("-d", "--detector", default = "DNN",
                help = "The detector used to detect faces: [mtcnn, dnn, cascade]")
ap.add_argument("-i", "--image", default = None,
                help = "The images that you want to detect faces from.")
args = vars(ap.parse_args())

# Get Cascade Classifier and DNN Model.
face_detector, net = load_info()

# Choose Cascade vs DNN.
dnn = True
cascade = False

# Bounding.
image_window = 10
if dnn:
   emotion_offsets = (45, 40)
elif cascade:
   emotion_offsets = (30, 40)
else:
   raise ValueError("You must choose either dnn or cascade.")

# Load Emotion Detection Model.
emotion_labels = fer2013_classes
emotion_classifier = load_keras_model('Model-27-0.6631', compile = False)
emotion_target_size = emotion_classifier.input_shape[1:3]
emotion_window = []

# If running from an IDE (not from command line), then enter images here.
savedir = "modded" # Directory to save changed images.
userimages = ["./test_imgs/amoghneutralnew.jpg"]

for image in userimages:
   file, extension = os.path.splitext(image)
   file = os.path.basename(file)
   image = np.array(cv2.imread(image), dtype = 'uint8')
   gray_image = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
   global faces
   if dnn:
      blob = cv2.dnn.blobFromImage(cv2.resize(image, (300, 300)), 1.0, (300, 300), swapRB=False, crop=False)
      net.setInput(blob)
      faces = net.forward()
   if cascade:
      faces = face_detector.detectMultiScale(gray_image, 1.3, 5)

   if cascade:
      for face_coordinates in faces:
         print(face_coordinates)
         x1, x2, y1, y2 = apply_offsets(face_coordinates, emotion_offsets)
         gray_face = gray_image[y1: y2, x1: x2]
         try:
            gray_face = cv2.resize(gray_face, (emotion_target_size))
         except:
            continue

         gray_face = preprocess_input(gray_face, False)
         gray_face = np.expand_dims(gray_face, 0)
         gray_face = np.expand_dims(gray_face, -1)
         emotion_prediction = emotion_classifier.predict(gray_face)
         emotion_probability = np.max(emotion_prediction)
         emotion_label_arg = np.argmax(emotion_prediction)
         emotion_text = emotion_labels[emotion_label_arg]
         emotion_window.append(emotion_text)

         if len(emotion_window) > image_window: emotion_window.pop(0)
         try:
            emotion_mode = mode(emotion_window)
         except:
            continue

         if emotion_text == 'happy':
            color = emotion_probability * np.asarray((255, 255, 0))
         elif emotion_text == 'angry':
            color = emotion_probability * np.asarray((255, 0, 0))
         elif emotion_text == 'sad':
            color = emotion_probability * np.asarray((0, 0, 255))
         elif emotion_text == 'surprise':
            color = emotion_probability * np.asarray((0, 255, 255))
         else:
            color = emotion_probability * np.asarray((0, 255, 0))

         color = color.astype(int)
         color = color.tolist()

         x, y, w, h = face_coordinates
         # cv2.rectangle(image, (x, y), (xe, ye), color, 3)
         cv2.rectangle(image, (x, y), (x + w, y + h), color, 3)
         cv2.putText(image, emotion_mode, (x + 0, y - 45), cv2.FONT_HERSHEY_SIMPLEX,
                     1, color, 3, cv2.LINE_AA)
         # Print a background behind the text.
         bg = np.full((image.shape), (0, 0, 0), dtype = np.uint8)
         cv2.putText(bg, emotion_mode, (x + 0, y - 45), cv2.FONT_HERSHEY_SIMPLEX,
                     1, color, 3, cv2.LINE_AA)
         x1, y1, w1, h1 = cv2.boundingRect(bg[:,:,2])
         image[y1: y1 + h1, x1: x1 + w1] = bg[y1: y1 + h1, x1: x1 + w1]

   if dnn:
      (h, w) = image.shape[:2]
      for k in range(0, faces.shape[2]):
         c = faces[0, 0, k, 2]
         if c < 0.5: continue
         box = faces[0, 0, k, 3:7] * np.array([w, h, w, h])
         (x, y, xe, ye) = box.astype("int")
         face_coordinates = (x, y, xe, ye)

         x1, x2, y1, y2 = apply_offsets(face_coordinates, emotion_offsets)
         gray_face = gray_image[y1: y2, x1: x2]
         try:
            gray_face = cv2.resize(gray_face, (emotion_target_size))
         except:
            continue

         gray_face = cv2.flip(gray_face, 1)
         # gray_face = preprocess_input(gray_face, False)
         gray_face = np.expand_dims(gray_face, 0)
         gray_face = np.expand_dims(gray_face, -1)
         emotion_prediction = emotion_classifier.predict(gray_face)
         emotion_probability = np.max(emotion_prediction)
         emotion_label_arg = np.argmax(emotion_prediction)
         emotion_text = emotion_labels[emotion_label_arg]
         emotion_window.append(emotion_text)

         if len(emotion_window) > image_window: emotion_window.pop(0)
         try:
            emotion_mode = mode(emotion_window)
         except:
            continue

         if emotion_text == 'happy':
            color = emotion_probability * np.asarray((255, 255, 0))
         elif emotion_text == 'angry':
            color = emotion_probability * np.asarray((255, 0, 0))
         elif emotion_text == 'sad':
            color = emotion_probability * np.asarray((0, 0, 255))
         elif emotion_text == 'surprise':
            color = emotion_probability * np.asarray((0, 255, 255))
         else:
            color = emotion_probability * np.asarray((0, 255, 0))

         color = color.astype(int)
         color = color.tolist()

         x, y, w, h = face_coordinates
         cv2.rectangle(image, (x, y), (xe, ye), color, 3)
         # cv2.rectangle(image, (x, y), (x + w, y + h), color, 3)
         cv2.putText(image, emotion_mode, (x + 0, y - 80), cv2.FONT_HERSHEY_SIMPLEX,
                     3, color, 3, cv2.LINE_AA)
         # Print a background behind the text.
         # bg = np.full((image.shape), (0, 0, 0), dtype = np.uint8)
         # cv2.putText(bg, emotion_mode, (x + 0, y - 45), cv2.FONT_HERSHEY_SIMPLEX,
         #             1, color, 3, cv2.LINE_AA)
         # x1, y1, w1, h1 = cv2.boundingRect(bg[:,:,2])
         # print(x1, y1, w1, h1)
         # image[y1: y1 + h1, x1: x1 + w1] = bg[y1: y1 + h1, x1: x1 + w1]

   cv2.imwrite(os.path.join(savedir or "", f"{file}-detect{extension}"), image)


<reponame>sina-moammar/2020-Fall-Computational-Physics
import numpy as np
from scipy.optimize import curve_fit
from typing import Callable, List, Tuple, TypeVar

def auto_correlation(data: List[float], lag: int) -> float:
    """Calculates auto-correlation of `data` with `lag` time delay

    Args:
        data (List[float]): data
        lag (int): lag of comparision

    Returns:
        float: auto-correlation of `data[lag:]` with `data[:-lag]
    """

    if lag == 0:
        return 1
    else:
        correlation = np.corrcoef(data[:len(data) - lag], data[lag:])[0, 1]
        return 0 if np.isnan(correlation) else correlation

def _exp_model(x: float, char_length: float) -> float:
    """Model of `f(x) = exp(- x / l)`

    Args:
        x (float): input
        char_length (float): characteristic length (l)

    Returns:
        float: output of function
    """

    return np.exp(-x / char_length)

def exp_characteristic_length(x_s: List[float], y_s: List[float]) -> Tuple[float, float]:
    """Finds characteristic length of `exp(-x / l)` function.

    Args:
        x_s (List[float]): inputs
        y_s (List[float]): outputs

    Returns:
        Tuple[float, float]: characteristic length and its error
    """

    try:
        x_length = x_s[-1] - x_s[0]
        x_length = 1 if x_length == 0 else x_length
        # should normalized x
        fit_para, fit_error = curve_fit(_exp_model, x_s / x_length, y_s, p0=(.5,))
        # because of normalization
        fit_para = fit_para[0] * x_length
        fit_error = np.sqrt(fit_error[0]) * x_length
    except:
        fit_para, fit_error = 0, 0

    return fit_para, fit_error

X = TypeVar('X')

def bootstrap_error(data: List[X], function: Callable[[List[X]], float], size: int = 100) -> float:
    """Calculates bootstrap error of `data`.

    Args:
        data (List[X]): list of any data
        function (Callable[[List[X]], float]): function gives the value that we want to calculate its error
        size (int, optional): number of batches. Defaults to 100.

    Returns:
        float: bootstrap error
    """

    ensemble_values = np.zeros(size)

    for i in range(size):
        # select `len(data)` values from data randomly
        random_numbers = np.random.randint(0, len(data), len(data))
        # get function output of this batch
        ensemble_values[i] = function(data[random_numbers])

    # return standard deviation of function output of batches
    return np.std(ensemble_values)