import numpy as np
import unittest
import pmsco.reports.rfactor as rp_rfactor


class TestGridMethods(unittest.TestCase):
    def test_triplet_to_grid__basic(self):
        x = np.array([-1, 0, 1, 1, 0, -1])
        y = np.array([2, 2, 2, 3, 3, 3])
        z = np.array([0.1, 0.2, 0.3, 0.6, 0.5, 0.4])
        gx, gy, gz = rp_rfactor.triplet_to_grid(x, y, z)
        expected_gx = np.array([[-1, 0, 1], [-1, 0, 1]]).T
        expected_gy = np.array([[2, 2, 2], [3, 3, 3]]).T
        expected_gz = np.array([[0.1, 0.2, 0.3], [0.4, 0.5, 0.6]]).T
        np.testing.assert_array_almost_equal(gx, expected_gx, 1, "grid_x")
        np.testing.assert_array_almost_equal(gy, expected_gy, 1, "grid_y")
        np.testing.assert_array_almost_equal(gz, expected_gz, 2, "grid_z")

    def test_triplet_to_grid__imprecise(self):
        x = np.array([-0.99, 0, 1, 1.001, 0, -1])
        y = np.array([1.999, 2.00001, 2, 3.01, 2.98, 3])
        z = np.array([0.1, 0.2, 0.3, 0.6, 0.5, 0.4])
        gx, gy, gz = rp_rfactor.triplet_to_grid(x, y, z)
        expected_gx = np.array([[-1, 0, 1], [-1, 0, 1]]).T
        expected_gy = np.array([[2, 2, 2], [3, 3, 3]]).T
        expected_gz = np.array([[0.1, 0.2, 0.3], [0.4, 0.5, 0.6]]).T
        np.testing.assert_array_almost_equal(gx, expected_gx, 1, "grid_x")
        np.testing.assert_array_almost_equal(gy, expected_gy, 1, "grid_y")
        np.testing.assert_array_almost_equal(gz, expected_gz, 2, "grid_z")

    def test_triplet_to_grid__missing(self):
        x = np.array([-1, 0, 1, 0, -1])
        y = np.array([2, 2, 3, 3, 3])
        z = np.array([0.1, 0.2, 0.6, 0.5, 0.4])
        gx, gy, gz = rp_rfactor.triplet_to_grid(x, y, z)
        expected_gx = np.array([[-1, 0, 1], [-1, 0, 1]]).T
        expected_gy = np.array([[2, 2, 2], [3, 3, 3]]).T
        expected_gz = np.array([[0.1, 0.2, 0.2], [0.4, 0.5, 0.6]]).T
        np.testing.assert_array_almost_equal(gx, expected_gx, 1, "grid_x")
        np.testing.assert_array_almost_equal(gy, expected_gy, 1, "grid_y")
        np.testing.assert_array_almost_equal(gz, expected_gz, 2, "grid_z")

    def test_triplet_to_grid__extra(self):
        x = np.array([-1, 0, 1, 1, 1, 0, -1])
        y = np.array([2, 2, 2, 2.01, 3, 3, 3])
        z = np.array([0.1, 0.2, 0.3, 0.35, 0.6, 0.5, 0.4])
        gx, gy, gz = rp_rfactor.triplet_to_grid(x, y, z)
        expected_gx = np.array([[-1, 0, 1], [-1, 0, 1]]).T
        expected_gy = np.array([[2, 2, 2], [3, 3, 3]]).T
        expected_gz = np.array([[0.1, 0.2, 0.3], [0.4, 0.5, 0.6]]).T
        np.testing.assert_array_almost_equal(gx, expected_gx, 1, "grid_x")
        np.testing.assert_array_almost_equal(gy, expected_gy, 1, "grid_y")
        np.testing.assert_array_almost_equal(gz, expected_gz, 2, "grid_z")

    def test_triplet_to_grid__split_column(self):
        x = np.array([-1, 0, 0.5, 1, 1, 0.5, 0, -1])
        y = np.array([2, 2, 2, 2, 3, 3, 3, 3])
        z = np.array([0.1, 0.2, 0.24, 0.3, 0.6, 0.45, 0.5, 0.4])
        gx, gy, gz = rp_rfactor.triplet_to_grid(x, y, z)
        expected_gx = np.array([[-1, -0.5, 0, 0.5, 1], [-1, -0.5, 0, 0.5, 1]]).T
        expected_gy = np.array([[2, 2, 2, 2, 2], [3, 3, 3, 3, 3]]).T
        expected_gz = np.array([[0.1, 0.1, 0.2, 0.24, 0.3], [0.4, 0.5, 0.5, 0.45, 0.6]]).T
        np.testing.assert_array_almost_equal(gx, expected_gx, 1, "grid_x")
        np.testing.assert_array_almost_equal(gy, expected_gy, 1, "grid_y")
        np.testing.assert_array_almost_equal(gz, expected_gz, 2, "grid_z")


if __name__ == '__main__':
    unittest.main()