import unittest
import numpy as np

from .. import query_disc, boundaries, nside2npix, nside2resol, ang2vec

try:
    from exceptions import ValueError
except:
    pass


class TestQueryDisc(unittest.TestCase):
    def setUp(self):
        self.NSIDE = 8
        self.vec = np.array([0.17101007, 0.03015369, 0.98480775])
        self.radius = np.radians(6)
        self.nside2_55_corners_precomp = np.array(
            [
                [
                    [2.44708573e-17, 5.27046277e-01, 3.60797400e-01, 4.56383842e-17],
                    [3.99652627e-01, 5.27046277e-01, 8.71041977e-01, 7.45355992e-01],
                    [9.16666667e-01, 6.66666667e-01, 3.33333333e-01, 6.66666667e-01],
                ],
                [
                    [2.44708573e-17, 5.27046277e-01, 3.60797400e-01, 4.56383842e-17],
                    [3.99652627e-01, 5.27046277e-01, 8.71041977e-01, 7.45355992e-01],
                    [9.16666667e-01, 6.66666667e-01, 3.33333333e-01, 6.66666667e-01],
                ],
            ]
        )

    def test_not_inclusive(self):
        # HIDL> query_disc, 8, [ 0.17101007,  0.03015369,  0.98480775],6,listpix,/DEG,NESTED=0
        # HIDL> print,listpix
        #           4
        np.testing.assert_array_equal(
            query_disc(self.NSIDE, self.vec, self.radius, inclusive=False),
            np.array([4]),
        )

    def test_inclusive(self):
        # HIDL> query_disc, 8, [ 0.17101007,  0.03015369,  0.98480775],6,listpix,/DEG,NESTED=0,/inclusive
        # HIDL> print,listpix
        #           0           3           4           5          11          12          13          23
        np.testing.assert_array_equal(
            query_disc(self.NSIDE, self.vec, self.radius, inclusive=True),
            np.array([0, 3, 4, 5, 11, 12, 13, 23]),
        )

    def test_boundaries(self):
        nside = 2
        corners = boundaries(nside, 5)
        corners_precomp = np.array(
            [
                [2.44708573e-17, 5.27046277e-01, 3.60797400e-01, 4.56383842e-17],
                [3.99652627e-01, 5.27046277e-01, 8.71041977e-01, 7.45355992e-01],
                [9.16666667e-01, 6.66666667e-01, 3.33333333e-01, 6.66666667e-01],
            ]
        )
        np.testing.assert_array_almost_equal(corners, corners_precomp, decimal=8)

    def test_boundaries_list(self):
        nside = 2
        corners = boundaries(nside, [5, 5])
        np.testing.assert_array_almost_equal(
            corners, self.nside2_55_corners_precomp, decimal=8
        )

    def test_boundaries_phi_theta(self):
        nside = 2
        corners = boundaries(nside, np.array([5, 5]))
        np.testing.assert_array_almost_equal(
            corners, self.nside2_55_corners_precomp, decimal=8
        )

    def test_nside_non_power_of_two(self):
        # For RING scheme, nside should not need to be a power of two.

        nside = 1
        resolution = 1.0
        theta=0.0
        phi=0.0
        radius = np.radians(1)
        while True:
            nside = nside + 1 
            res = nside2resol(nside, arcmin = True)
            print("nside={} res={} arcmin".format(nside, res))
            if res < resolution:
                break
            
        self.assertEqual(nside, 3518)
        
        x0 = ang2vec(theta, phi)
        
        pixel_indices = query_disc(nside, x0, radius, inclusive=False, nest=False)
        self.assertEqual(pixel_indices.shape[0], 11400)
        
    def test_boundaries_floatpix_array(self):
        self.assertRaises(ValueError, boundaries, 2, np.array([5.0, 5]))

    def test_boundaries_floatpix_scalar(self):
        self.assertRaises(ValueError, boundaries, 2, 1 / 2.0)

    def test_buffer_mode(self):

        # allocate something manifestly too short, should raise a value error
        buff = np.empty(0, dtype=np.int64)
        self.assertRaises(
            ValueError,
            query_disc,
            self.NSIDE,
            self.vec,
            self.radius,
            inclusive=True,
            buff=buff,
        )

        # allocate something of wrong type, should raise a value error
        buff = np.empty(nside2npix(self.NSIDE), dtype=np.float64)
        self.assertRaises(
            ValueError,
            query_disc,
            self.NSIDE,
            self.vec,
            self.radius,
            inclusive=True,
            buff=buff,
        )

        # allocate something acceptable, should succeed and return a subview
        buff = np.empty(nside2npix(self.NSIDE), dtype=np.int64)
        result = query_disc(
            self.NSIDE, self.vec, self.radius, inclusive=True, buff=buff
        )

        assert result.base is buff

        np.testing.assert_array_equal(result, np.array([0, 3, 4, 5, 11, 12, 13, 23]))