import unittest

import numpy as np
import pytest
from numpy.testing import assert_array_almost_equal, assert_array_equal

import ldsc.regressions as reg

np.set_printoptions(precision=4)


def test_update_separators():
    ii1 = [True, True, False, True, True, False, True]
    ii2 = [True, True, False, True, True, False, False]
    ii3 = [False, True, False, True, True, False, False]
    ii4 = [False, True, False, True, True, False, True]
    ii5 = [True, True, True, True, True, True, True]
    iis = map(np.array, [ii1, ii2, ii3, ii4, ii5])
    ids = np.arange(len(ii1))
    for ii in iis:
        s = np.arange(np.sum(ii) + 1)
        t = reg.update_separators(s, ii)
        assert t[0] == 0
        assert t[-1] == len(ii)
        assert_array_equal(ids[ii][(s[1:-2])], ids[(t[1:-2])])


def test_p_z_norm():
    est = 10
    se = 1
    p, z = reg.p_z_norm(est, se)
    assert z == 10
    assert_array_almost_equal(p * 1e23, 1.523971)
    se = 0
    p, z = reg.p_z_norm(est, se)
    assert p == 0
    assert np.isinf(z)


def test_append_intercept():
    x = np.ones((5, 2))
    correct_x = np.ones((5, 3))
    assert_array_equal(reg.append_intercept(x), correct_x)


def test_remove_brackets():
    x = " [] [] asdf [] "
    assert reg.remove_brackets(x) == "asdf"


class Test_h2_obs_to_liab(unittest.TestCase):
    def test_bad_data(self):
        with pytest.raises(ValueError):
            reg.h2_obs_to_liab(1, 1, 0.5)
        with pytest.raises(ValueError):
            reg.h2_obs_to_liab(1, 0.5, 1)
        with pytest.raises(ValueError):
            reg.h2_obs_to_liab(1, 0, 0.5)
        with pytest.raises(ValueError):
            reg.h2_obs_to_liab(1, 0.5, 0)

    def test_approx_scz(self):
        # conversion for a balanced study of a 1% phenotype is about 1/2
        x = reg.h2_obs_to_liab(1, 0.5, 0.01)
        assert_array_almost_equal(x, 0.551907298063)


class Test_gencov_obs_to_liab(unittest.TestCase):
    def test_qt(self):
        self.assertEqual(reg.gencov_obs_to_liab(1, None, None, None, None), 1)

    def test_approx_scz(self):
        x = reg.gencov_obs_to_liab(1, 0.5, None, 0.01, None)
        assert_array_almost_equal(x, np.sqrt(0.551907298063))
        x = reg.gencov_obs_to_liab(1, None, 0.5, None, 0.01)
        assert_array_almost_equal(x, np.sqrt(0.551907298063))
        x = reg.gencov_obs_to_liab(1, 0.5, 0.5, 0.01, 0.01)
        assert_array_almost_equal(x, 0.551907298063)


class Test_Hsq_1D(unittest.TestCase):
    def setUp(self):
        self.chisq = np.ones((4, 1)) * 4
        self.ld = np.ones((4, 1))
        self.w_ld = np.ones((4, 1))
        self.N = 9 * np.ones((4, 1))
        self.M = np.array([[7]])
        self.hsq = reg.Hsq(
            self.chisq, self.ld, self.w_ld, self.N, self.M, n_blocks=3, intercept=1
        )

    def test_weights(self):
        hsq = 0.5
        w = reg.Hsq.weights(self.ld, self.w_ld, self.N, self.M, hsq)
        assert_array_equal(w.shape, self.ld.shape)
        assert_array_almost_equal(w[0, 0], 0.5 / np.square(1 + hsq * self.N / self.M))
        # test that it deals correctly with out-of-bounds h2
        assert_array_almost_equal(
            reg.Hsq.weights(self.ld, self.w_ld, self.N, self.M, 1),
            reg.Hsq.weights(self.ld, self.w_ld, self.N, self.M, 2),
        )
        assert_array_almost_equal(
            reg.Hsq.weights(self.ld, self.w_ld, self.N, self.M, 0),
            reg.Hsq.weights(self.ld, self.w_ld, self.N, self.M, -1),
        )

    def test_summarize_chisq(self):
        chisq = np.arange(100).reshape((100, 1))
        mean_chisq, lambda_gc = self.hsq._summarize_chisq(chisq)
        assert_array_almost_equal(mean_chisq, 49.5)
        assert_array_almost_equal(lambda_gc, 108.81512420312156)

    def test_summary(self):
        # not much to test; we can at least make sure no errors at runtime
        self.hsq.summary(["asdf"])
        self.ld += np.arange(4).reshape((4, 1))
        self.chisq += np.arange(4).reshape((4, 1))
        hsq = reg.Hsq(self.chisq, self.ld, self.w_ld, self.N, self.M, n_blocks=3)
        hsq.summary(["asdf"])
        # test ratio printout with mean chi^2 < 1
        hsq.mean_chisq = 0.5
        hsq.summary(["asdf"])

    def test_update(self):
        pass

    def test_aggregate(self):
        chisq = np.ones((10, 1)) * 3 / 2
        ld = np.ones((10, 1)) * 100
        N = np.ones((10, 1)) * 100000
        M = 1e7
        agg = reg.Hsq.aggregate(chisq, ld, N, M)
        assert_array_almost_equal(agg, 0.5)
        agg = reg.Hsq.aggregate(chisq, ld, N, M, intercept=1.5)
        assert_array_almost_equal(agg, 0)


class Test_Coef(unittest.TestCase):
    def setUp(self):
        self.hsq1 = 0.2
        self.hsq2 = 0.7
        ld = (np.abs(np.random.normal(size=800)) + 1).reshape((400, 2))
        N = np.ones((400, 1)) * 1e5
        self.M = np.ones((1, 2)) * 1e7 / 2.0
        chisq = 1 + 1e5 * (
            ld[:, 0] * self.hsq1 / self.M[0, 0] + ld[:, 1] * self.hsq2 / self.M[0, 1]
        ).reshape((400, 1))
        w_ld = np.ones_like(chisq)
        self.hsq_noint = reg.Hsq(chisq, ld, w_ld, N, self.M, n_blocks=3, intercept=1)
        self.hsq_int = reg.Hsq(chisq, ld, w_ld, N, self.M, n_blocks=3)
        print(self.hsq_noint.summary())
        print(self.hsq_int.summary())

    def test_coef(self):
        a = [self.hsq1 / self.M[0, 0], self.hsq2 / self.M[0, 1]]
        assert_array_almost_equal(self.hsq_noint.coef, a)
        assert_array_almost_equal(self.hsq_int.coef, a)

    def test_cat_hsq(self):
        a = [[self.hsq1, self.hsq2]]
        assert_array_almost_equal(self.hsq_noint.cat, a)
        assert_array_almost_equal(self.hsq_int.cat, a)

    def test_tot(self):
        a = self.hsq1 + self.hsq2
        assert_array_almost_equal(self.hsq_noint.tot, a)
        assert_array_almost_equal(self.hsq_int.tot, a)

    def test_prop_hsq(self):
        d = self.hsq1 + self.hsq2
        a = [[self.hsq1 / d, self.hsq2 / d]]
        assert_array_almost_equal(self.hsq_noint.prop, a)
        assert_array_almost_equal(self.hsq_int.prop, a)

    def test_enrichment(self):
        d = self.hsq1 + self.hsq2
        a = [[self.hsq1 / (0.5 * d), self.hsq2 / (0.5 * d)]]
        assert_array_almost_equal(self.hsq_noint.enrichment, a)
        assert_array_almost_equal(self.hsq_int.enrichment, a)

    def test_intercept(self):
        assert self.hsq_int.intercept is not None
        assert self.hsq_int.ratio is not None
        assert isinstance(self.hsq_int.ratio, float)
        assert_array_almost_equal(self.hsq_int.intercept, 1)
        assert_array_almost_equal(self.hsq_int.ratio, 0)


class Test_Hsq_2D(unittest.TestCase):
    def setUp(self):
        self.chisq = np.ones((17, 1)) * 4
        self.ld = np.hstack([np.ones((17, 1)), np.arange(17).reshape((17, 1))]).reshape(
            (17, 2)
        )
        self.w_ld = np.ones((17, 1))
        self.N = 9 * np.ones((17, 1))
        self.M = np.array([[7, 2]])
        self.hsq = reg.Hsq(
            self.chisq, self.ld, self.w_ld, self.N, self.M, n_blocks=3, intercept=1
        )

    def test_summary(self):
        # not much to test; we can at least make sure no errors at runtime
        self.hsq.summary(["asdf", "qwer"])
        # change to random 7/30/2019 to avoid inconsistent singular matrix errors
        self.ld += np.random.normal(scale=0.1, size=(17, 2))
        self.chisq += np.arange(17).reshape((17, 1))
        hsq = reg.Hsq(self.chisq, self.ld, self.w_ld, self.N, self.M, n_blocks=3)
        hsq.summary(["asdf", "qwer"])
        # test ratio printout with mean chi^2 < 1
        hsq.mean_chisq = 0.5
        hsq.summary(["asdf", "qwer"])


class Test_Gencov_1D(unittest.TestCase):
    def setUp(self):
        self.z1 = np.ones((4, 1)) * 4
        self.z2 = np.ones((4, 1))
        self.ld = np.ones((4, 1))
        self.w_ld = np.ones((4, 1))
        self.N1 = 9 * np.ones((4, 1))
        self.N2 = 7 * np.ones((4, 1))
        self.M = np.array([[7]])
        self.hsq1 = 0.5
        self.hsq2 = 0.6
        self.gencov = reg.Gencov(
            self.z1,
            self.z2,
            self.ld,
            self.w_ld,
            self.N1,
            self.N2,
            self.M,
            self.hsq1,
            self.hsq2,
            1.0,
            1.0,
            n_blocks=3,
            intercept_gencov=1,
        )

    def test_weights(self):
        # check that hsq weights = gencov weights when z1 = z2
        ld = np.abs(np.random.normal(size=100)).reshape((100, 1))
        w_ld = np.abs(np.random.normal(size=100)).reshape((100, 1))
        N1 = np.abs(np.random.normal(size=100)).reshape((100, 1))
        N2 = N1
        M = 10
        h1, h2, rho_g = 0.5, 0.5, 0.5
        wg = reg.Gencov.weights(
            ld, w_ld, N1, N2, M, h1, h2, rho_g, intercept_gencov=1.0
        )
        wh = reg.Hsq.weights(ld, w_ld, N1, M, h1, intercept=1.0)
        assert_array_almost_equal(wg, wh)

    def test_update(self):
        pass

    def test_summary(self):
        # not much to test; we can at least make sure no errors at runtime
        self.gencov.summary(["asdf"])
        self.ld += np.arange(4).reshape((4, 1))
        self.z1 += np.arange(4).reshape((4, 1))
        gencov = reg.Gencov(
            self.z1,
            self.z2,
            self.ld,
            self.w_ld,
            self.N1,
            self.N2,
            self.M,
            self.hsq1,
            self.hsq2,
            1.0,
            1.0,
            n_blocks=3,
        )
        gencov.summary(["asdf"])

    def test_aggregate(self):
        z1z2 = np.ones((10, 1)) / 2
        ld = np.ones((10, 1)) * 100
        N = np.ones((10, 1)) * 100000
        M = 1e7
        agg = reg.Gencov.aggregate(z1z2, ld, N, M)
        assert_array_almost_equal(agg, 0.5)
        agg = reg.Gencov.aggregate(z1z2, ld, N, M, intercept=0.5)
        assert_array_almost_equal(agg, 0)


class Test_Gencov_2D(unittest.TestCase):
    def setUp(self):
        self.ld = np.abs(np.random.normal(size=100).reshape((50, 2))) + 2
        self.z1 = np.random.normal(size=50).reshape((50, 1))
        self.z2 = np.random.normal(size=50).reshape((50, 1))
        self.w_ld = np.random.normal(size=50).reshape((50, 1))
        self.N1 = 9 * np.ones((50, 1))
        self.N2 = 7 * np.ones((50, 1))
        self.M = np.array([[700, 222]])
        self.hsq1 = 0.5
        self.hsq2 = 0.6
        self.gencov = reg.Gencov(
            self.z1,
            self.z2,
            self.ld,
            self.w_ld,
            self.N1,
            self.N2,
            self.M,
            self.hsq1,
            self.hsq2,
            1.0,
            1.0,
            n_blocks=3,
            intercept_gencov=1,
        )

    def test_summary(self):
        # not much to test; we can at least make sure no errors at runtime
        self.gencov.summary(["asdf", "qwer"])

    def test_eq_hsq(self):
        """
        Gencov should be the same as hsq if z1 = z2, hsq + intercept_hsq are 0 and
        all intermediate rg's are > 0 (because Hsq.weights lower-bounds the hsq guess at 0
        but Gencov.weights lower-bounds the rho_g guess at -1). The setup below guarantees
        that all intermediate rho_g guesses will be 1

        """
        self.ld = np.abs(np.random.normal(size=100).reshape((50, 2))) + 2
        self.z1 = (np.sum(self.ld, axis=1) + 10).reshape((50, 1))
        gencov = reg.Gencov(
            self.z1,
            self.z1,
            self.ld,
            self.w_ld,
            self.N1,
            self.N1,
            self.M,
            0,
            0,
            0,
            0,
            n_blocks=3,
            intercept_gencov=1,
        )
        hsq = reg.Hsq(
            np.square(self.z1),
            self.ld,
            self.w_ld,
            self.N1,
            self.M,
            n_blocks=3,
            intercept=1,
        )
        print(gencov.summary(["asdf", "asdf"]))
        print()
        print(hsq.summary(["asdf", "asdf"]))
        assert_array_almost_equal(gencov.tot, hsq.tot)
        assert_array_almost_equal(gencov.tot_se, hsq.tot_se)
        assert_array_almost_equal(gencov.tot_cov, hsq.tot_cov)


class Test_RG_2D(unittest.TestCase):
    def setUp(self):
        self.ld = np.abs(np.random.normal(size=100).reshape((50, 2))) + 2
        self.z1 = (np.sum(self.ld, axis=1) * 10).reshape((50, 1))
        self.w_ld = np.random.normal(size=50).reshape((50, 1))
        self.N1 = 9 * np.ones((50, 1))
        self.N2 = 7 * np.ones((50, 1))
        self.M = np.array([[700, 222]])
        self.hsq1 = 0.5
        self.hsq2 = 0.6
        self.rg = reg.RG(
            self.z1,
            -self.z1,
            self.ld,
            self.w_ld,
            self.N1,
            self.N1,
            self.M,
            1.0,
            1.0,
            0,
            n_blocks=20,
        )

    def test_summary(self):
        # just make sure it doesn't encounter any errors at runtime
        print(self.rg.summary())
        print(self.rg.summary(silly=True))

    def test_rg(self):
        # won't be exactly 1 because the h2 values passed to Gencov aren't 0
        assert isinstance(self.rg.rg_ratio, float)
        assert np.abs(self.rg.rg_ratio + 1) < 0.01


class Test_RG_Bad(unittest.TestCase):
    def test_negative_h2(self):
        ld = np.arange(50).reshape((50, 1)) + 0.1
        z1 = (1 / np.sum(ld, axis=1) * 10).reshape((50, 1))
        w_ld = np.ones((50, 1))
        N1 = 9 * np.ones((50, 1))
        M = np.array([[-700]])
        rg = reg.RG(z1, -z1, ld, w_ld, N1, N1, M, 1.0, 1.0, 0, n_blocks=20)
        assert rg._negative_hsq
        # check no runtime errors when _negative_hsq is True
        print(rg.summary())
        print(rg.summary(silly=True))
        assert rg.rg_ratio == "NA"
        assert rg.rg_se == "NA"
        assert rg.rg == "NA"
        assert rg.p == "NA"
        assert rg.z == "NA"
