# Copyright (c) 2018-2025 by xcube team and contributors
# Permissions are hereby granted under the terms of the MIT License:
# https://opensource.org/licenses/MIT.

import unittest

import numpy as np

from xcube_resampling.gridmapping import CRS_WGS84, GridMapping
from xcube_resampling.reproject import reproject_dataset

from .sampledata import (
    create_2x5x5_dataset_regular_utm,
    create_5x5_dataset_regular_utm,
    create_5x5_dataset_regular_utm_antimeridian,
    create_large_dataset_for_reproject,
)


# noinspection PyMethodMayBeStatic
class ReprojectDatasetTest(unittest.TestCase):
    def test_reproject_target_gm(self):
        source_ds = create_5x5_dataset_regular_utm()

        # test projected CRS, similar resolution
        target_gm = GridMapping.regular(
            size=(5, 5), xy_min=(4320120, 3382520), xy_res=80, crs="epsg:3035"
        )
        target_ds = reproject_dataset(source_ds, target_gm)
        np.testing.assert_almost_equal(
            target_ds.band_1.values,
            np.array(
                [
                    [1, 1, 2, 3, 4],
                    [6, 6, 7, 8, 9],
                    [11, 12, 12, 13, 14],
                    [16, 17, 17, 18, 19],
                    [21, 17, 17, 18, 19],
                ],
                dtype=target_ds.band_1.dtype,
            ),
        )

    def test_reproject_target_gm_3d(self):
        source_ds = create_2x5x5_dataset_regular_utm()

        # test projected CRS, similar resolution
        target_gm = GridMapping.regular(
            size=(5, 5), xy_min=(4320120, 3382520), xy_res=80, crs="epsg:3035"
        )
        target_ds = reproject_dataset(source_ds, target_gm)
        self.assertEqual(
            set(source_ds.variables),
            set(target_ds.variables),
        )
        np.testing.assert_almost_equal(
            target_ds.band_1.values,
            np.array(
                [
                    [
                        [1, 1, 2, 3, 4],
                        [6, 6, 7, 8, 9],
                        [11, 12, 12, 13, 14],
                        [16, 17, 17, 18, 19],
                        [21, 17, 17, 18, 19],
                    ],
                    [
                        [1, 1, 2, 3, 4],
                        [6, 6, 7, 8, 9],
                        [11, 12, 12, 13, 14],
                        [16, 17, 17, 18, 19],
                        [21, 17, 17, 18, 19],
                    ],
                ],
                dtype=target_ds.band_1.dtype,
            ),
        )

    def test_reproject_target_gm_j_axis_up(self):
        source_ds = create_5x5_dataset_regular_utm()
        target_gm = GridMapping.regular(
            size=(5, 5),
            xy_min=(4320120, 3382520),
            xy_res=80,
            crs="epsg:3035",
            is_j_axis_up=True,
        )
        target_ds = reproject_dataset(source_ds, target_gm)
        np.testing.assert_almost_equal(
            target_ds.band_1.values,
            np.array(
                [
                    [21, 17, 17, 18, 19],
                    [16, 17, 17, 18, 19],
                    [11, 12, 12, 13, 14],
                    [6, 6, 7, 8, 9],
                    [1, 1, 2, 3, 4],
                ],
                dtype=target_ds.band_1.dtype,
            ),
        )

    def test_reproject_source_gm_j_axis_up(self):
        source_ds = create_5x5_dataset_regular_utm()
        source_ds = source_ds.isel(y=slice(None, None, -1))
        target_gm = GridMapping.regular(
            size=(5, 5), xy_min=(4320120, 3382520), xy_res=80, crs="epsg:3035"
        )
        target_ds = reproject_dataset(source_ds, target_gm)
        np.testing.assert_almost_equal(
            target_ds.band_1.values,
            np.array(
                [
                    [1, 1, 2, 3, 4],
                    [6, 6, 7, 8, 9],
                    [11, 12, 12, 13, 14],
                    [16, 17, 17, 18, 19],
                    [21, 17, 17, 18, 19],
                ],
                dtype=target_ds.band_1.dtype,
            ),
        )

    def test_reproject_target_gm_finer_res(self):
        source_ds = create_5x5_dataset_regular_utm()
        target_gm = GridMapping.regular(
            size=(5, 5), xy_min=(4320090, 3382490), xy_res=20, crs="epsg:3035"
        )
        target_ds = reproject_dataset(source_ds, target_gm)
        np.testing.assert_almost_equal(
            target_ds.band_1.values,
            np.array(
                [
                    [15, 16, 16, 16, 16],
                    [15, 16, 16, 16, 16],
                    [15, 16, 16, 16, 16],
                    [20, 21, 21, 21, 21],
                    [20, 21, 21, 21, 21],
                ],
                dtype=target_ds.band_1.dtype,
            ),
        )

    def test_reproject_target_gm_coarser_res(self):
        source_ds = create_5x5_dataset_regular_utm()
        target_gm = GridMapping.regular(
            size=(3, 3), xy_min=(4320070, 3382620), xy_res=120, crs="epsg:3035"
        )
        target_ds = reproject_dataset(source_ds, target_gm)
        np.testing.assert_almost_equal(
            target_ds.band_1.values,
            np.array(
                [
                    [0, 1, 3],
                    [5, 6, 8],
                    [15, 11, 13],
                ],
                dtype=target_ds.band_1.dtype,
            ),
        )

    def test_reproject_target_gm_geographic_crs(self):
        source_ds = create_5x5_dataset_regular_utm()
        target_gm = GridMapping.regular(
            size=(5, 5), xy_min=(9.9889, 53.5502), xy_res=0.0006, crs=CRS_WGS84
        )
        target_ds = reproject_dataset(source_ds, target_gm)
        np.testing.assert_almost_equal(
            target_ds.band_1.values,
            np.array(
                [
                    [7, 8, 8, 8, 9],
                    [12, 13, 13, 13, 14],
                    [12, 13, 13, 13, 14],
                    [17, 18, 18, 18, 19],
                    [22, 23, 23, 23, 24],
                ],
                dtype=target_ds.band_1.dtype,
            ),
        )

    def test_reproject_target_gm_geographic_crs_antimeridian(self):
        source_ds = create_5x5_dataset_regular_utm_antimeridian()

        # Geographic target grid placed near 179.99°E
        target_gm = GridMapping.regular(
            size=(5, 5),
            xy_min=(179.998, 9.95),
            xy_res=0.001,
            crs=CRS_WGS84,
        )

        fill_value = -9999
        target_ds = reproject_dataset(
            source_ds,
            target_gm,
            fill_values=fill_value,
        )

        np.testing.assert_almost_equal(
            target_ds.band_1.values,
            np.array(
                [
                    [-9999, -9999, -9999, -9999, -9999],
                    [0, 1, 2, 3, 4],
                    [5, 6, 7, 8, 9],
                    [10, 11, 12, 13, 14],
                    [15, 16, 17, 18, 19],
                ],
                dtype=target_ds.band_1.dtype,
            ),
        )
        np.testing.assert_almost_equal(
            target_ds.lon.values,
            np.array(
                [179.998, 179.999, 180.0, -179.999, -179.998],
                dtype=target_ds.lon.dtype,
            ),
        )

    def test_reproject_target_gm_geographic_crs_fine_res(self):
        source_ds = create_5x5_dataset_regular_utm()

        # test geographic CRS with 1/2 resolution
        target_gm = GridMapping.regular(
            size=(5, 5), xy_min=(9.98875, 53.55005), xy_res=0.0003, crs=CRS_WGS84
        )
        target_ds = reproject_dataset(source_ds, target_gm)
        np.testing.assert_almost_equal(
            target_ds.band_1.values,
            np.array(
                [
                    [12, 12, 12, 13, 13],
                    [17, 17, 17, 18, 18],
                    [17, 17, 17, 18, 18],
                    [22, 17, 17, 18, 18],
                    [22, 22, 22, 23, 23],
                ],
                dtype=target_ds.band_1.dtype,
            ),
        )

    def test_reproject_complex_dask_array(self):
        source_ds = create_large_dataset_for_reproject()
        target_gm = GridMapping.regular(
            size=(10, 10),
            xy_min=(6.1, 48.1),
            xy_res=0.2,
            crs=CRS_WGS84,
            tile_size=(5, 5),
        )

        target_ds = reproject_dataset(source_ds, target_gm, interp_methods="triangular")
        self.assertCountEqual(["temperature", "onedim_data"], list(target_ds.data_vars))
        self.assertAlmostEqual(target_ds.temperature.values[0, 0, 0], 6216.14, places=2)
        self.assertAlmostEqual(
            target_ds.temperature.values[0, -1, -1], 2848.50, places=2
        )
        self.assertEqual(
            [2, 5, 5],
            [
                target_ds.temperature.chunksizes["time"][0],
                target_ds.temperature.chunksizes["lat"][0],
                target_ds.temperature.chunksizes["lon"][0],
            ],
        )

        target_ds = reproject_dataset(source_ds, target_gm, interp_methods=1)
        self.assertCountEqual(["temperature", "onedim_data"], list(target_ds.data_vars))
        self.assertAlmostEqual(target_ds.temperature.values[0, 0, 0], 6216.14, places=2)
        self.assertAlmostEqual(
            target_ds.temperature.values[0, -1, -1], 2848.50, places=2
        )
        self.assertEqual(
            [2, 5, 5],
            [
                target_ds.temperature.chunksizes["time"][0],
                target_ds.temperature.chunksizes["lat"][0],
                target_ds.temperature.chunksizes["lon"][0],
            ],
        )

    def test_reproject_raise_not_implemented(self):
        source_ds = create_5x5_dataset_regular_utm()
        target_gm = GridMapping.regular(
            size=(5, 5), xy_min=(4320080, 3382480), xy_res=20, crs="epsg:3035"
        )
        with self.assertRaises(NotImplementedError) as cm:
            _ = reproject_dataset(source_ds, target_gm, interp_methods="cubic")
        self.assertIn(
            "interp_methods must be one of 0, 1, 'nearest', 'bilinear', 'triangular'",
            str(cm.exception),
        )

    def test_reproject_outside_return_empty_ds(self):
        source_ds = create_5x5_dataset_regular_utm()
        target_gm = GridMapping.regular(
            size=(5, 5), xy_min=(4320550, 3382920), xy_res=80, crs="epsg:3035"
        )
        target_ds = reproject_dataset(source_ds, target_gm)
        np.testing.assert_almost_equal(
            target_ds.band_1.values,
            np.full((5, 5), -1, dtype=target_ds.band_1.dtype),
        )
