#  Copyright (c) 2025 by EOPF Sample Service team and contributors
#  Permissions are hereby granted under the terms of the Apache 2.0 License:
#  https://opensource.org/license/apache-2-0.

from collections.abc import Sequence
from unittest import TestCase

import fsspec
import numpy as np
import pyproj
import pytest
import xarray as xr
import zarr.storage

from tests.helpers import make_s2_msi_l1c, make_s2_msi_l2a
from xarray_eopf.amode import AnalysisModeRegistry
from xarray_eopf.amodes.sentinel2 import MsiL1c, MsiL2a, register


class Sentinel2AnalysisModeTest(TestCase):
    def test_register(self):
        registry = AnalysisModeRegistry()
        register(registry)
        self.assertEqual(2, len(list(registry.keys())))


# noinspection PyUnresolvedReferences
class MsiTestMixin:
    def test_is_valid_source(self: TestCase):
        pass

    def test_get_applicable_params(self: TestCase):
        self.assertEqual(
            {},
            self.mode.get_applicable_params(),
        )
        self.assertEqual(
            {
                "resolution": 10,
                "interp_methods": 1,
                "bbox": [1, 3, 4, 5],
                "crs": pyproj.CRS.from_string("EPSG:4326"),
                "agg_methods": {"scl": "mode"},
            },
            self.mode.get_applicable_params(
                resolution=10,
                interp_methods=1,
                bbox=[1, 3, 4, 5],
                crs="EPSG:4326",
                agg_methods={"scl": "mode"},
            ),
        )

    def test_process_metadata(self: TestCase):
        self.assertEqual({}, self.mode.process_metadata(xr.DataTree()))

    def test_assign_grid_mapping(self: TestCase):
        def make_band():
            return xr.DataArray(np.zeros((10, 10)), dims=("y", "x"))

        # in dataset attrs with key "horizontal_CRS_code"
        dataset = self.mode.assign_grid_mapping(
            xr.Dataset(
                dict(
                    b01=make_band(),
                    b02=make_band(),
                    b03=make_band(),
                ),
                attrs={"horizontal_CRS_code": "ESPG:32632"},
            ),
            {"test": "dict"},
        )
        self.assertIn("spatial_ref", dataset)
        self.assertEqual(
            "transverse_mercator", dataset.spatial_ref.attrs.get("grid_mapping_name")
        )
        self.assertEqual("spatial_ref", dataset.b01.attrs.get("grid_mapping"))
        self.assertEqual("spatial_ref", dataset.b02.attrs.get("grid_mapping"))
        self.assertEqual("spatial_ref", dataset.b03.attrs.get("grid_mapping"))

        # in band attrs with key "proj:epsg"
        ds = xr.Dataset(
            dict(
                b01=make_band(),
                b02=make_band(),
                b03=make_band(),
            ),
        )
        ds["b01"].attrs = {"proj:epsg": 32632}
        dataset = self.mode.assign_grid_mapping(ds, {"test": "dict"})
        self.assertIn("spatial_ref", dataset)
        self.assertEqual(
            "transverse_mercator", dataset.spatial_ref.attrs.get("grid_mapping_name")
        )
        self.assertEqual("spatial_ref", dataset.b01.attrs.get("grid_mapping"))
        self.assertEqual("spatial_ref", dataset.b02.attrs.get("grid_mapping"))
        self.assertEqual("spatial_ref", dataset.b03.attrs.get("grid_mapping"))

        # in band attrs with key "proj:epsg"
        dataset = self.mode.assign_grid_mapping(
            xr.Dataset(
                dict(
                    b01=make_band(),
                    b02=make_band(),
                    b03=make_band(),
                ),
            ),
            {"properties": {"proj:code": "EPSG:32632"}},
        )
        self.assertIn("spatial_ref", dataset)
        self.assertEqual(
            "transverse_mercator", dataset.spatial_ref.attrs.get("grid_mapping_name")
        )
        self.assertEqual("spatial_ref", dataset.b01.attrs.get("grid_mapping"))
        self.assertEqual("spatial_ref", dataset.b02.attrs.get("grid_mapping"))
        self.assertEqual("spatial_ref", dataset.b03.attrs.get("grid_mapping"))

    def test_assign_grid_mapping_fail(self: TestCase):
        def make_band():
            return xr.DataArray(np.zeros((10, 10)), dims=("y", "x"))

        dataset = self.mode.assign_grid_mapping(
            xr.Dataset(
                dict(
                    b01=make_band(),
                    b02=make_band(),
                    b03=make_band(),
                ),
                attrs={"horizontal_CRS_code": "ESPG:-1"},
            ),
            {},
        )
        self.assertNotIn("spatial_ref", dataset)
        self.assertEqual(None, dataset.b01.attrs.get("grid_mapping"))
        self.assertEqual(None, dataset.b02.attrs.get("grid_mapping"))
        self.assertEqual(None, dataset.b03.attrs.get("grid_mapping"))

        # test with wrong parameters in data variable attributes
        def make_band():
            return xr.DataArray(
                np.zeros((10, 10)),
                dims=("y", "x"),
                attrs={"proj:epsg": -1},
            )

        dataset = self.mode.assign_grid_mapping(
            xr.Dataset(
                dict(
                    b01=make_band(),
                    b02=make_band(),
                    b03=make_band(),
                )
            ),
            {},
        )
        self.assertNotIn("spatial_ref", dataset)
        self.assertEqual(None, dataset.b01.attrs.get("grid_mapping"))
        self.assertEqual(None, dataset.b02.attrs.get("grid_mapping"))
        self.assertEqual(None, dataset.b03.attrs.get("grid_mapping"))

    def assert_transform_datatree_ok(self, original_dt: xr.DataTree):
        dt = self.mode.transform_datatree(original_dt, resolution=10)
        self.assertIs(dt, original_dt)

    def assert_convert_datatree_ok(
        self,
        original_dt: xr.DataTree,
        expected_var_names: str | list[str],
        expected_size: int | tuple[int, int],
        bbox: Sequence[float | int] | None = None,
        crs: pyproj.CRS | str | None = None,
        resolution: int | float | None = None,
    ):
        ds = self.mode.convert_datatree(
            original_dt,
            includes=expected_var_names,
            resolution=resolution,
            bbox=bbox,
            crs=crs,
        )
        self.assertIsInstance(ds, xr.Dataset)
        if isinstance(expected_var_names, str):
            expected_var_names = [expected_var_names]
        self.assertCountEqual(expected_var_names, ds.data_vars.keys())
        if isinstance(expected_size, int):
            expected_size = (expected_size, expected_size)
        for var_name, var in ds.data_vars.items():
            self.assertEqual(
                (expected_size[0], expected_size[1]), var.shape, msg=var_name
            )

        # noinspection PyTypeChecker
        if crs is not None and crs.is_geographic:
            self.assertCountEqual(["spatial_ref", "lat", "lon"], ds.coords.keys())
            self.assertEqual((expected_size[1],), ds.lon.shape, msg="x")
            self.assertEqual((expected_size[0],), ds.lat.shape, msg="y")
        else:
            self.assertCountEqual(["spatial_ref", "x", "y"], ds.coords.keys())
            self.assertEqual((expected_size[1],), ds.x.shape, msg="x")
            self.assertEqual((expected_size[0],), ds.y.shape, msg="y")

    def assert_convert_datatree_fail(self, original_dt: xr.DataTree):
        with pytest.raises(ValueError, match="No variables selected"):
            self.mode.convert_datatree(original_dt, includes="bibo")


class MsiL1CTest(MsiTestMixin, TestCase):
    mode = MsiL1c()

    def test_is_valid_source_ok(self):
        self.assertTrue(self.mode.is_valid_source("data/S2A_MSIL1C_20240201.zarr"))
        self.assertTrue(
            self.mode.is_valid_source(
                zarr.storage.DirectoryStore("data/S2A_MSIL1C_20240201.zarr")
            )
        )
        fs: fsspec.AbstractFileSystem = fsspec.filesystem("local")
        self.assertTrue(
            self.mode.is_valid_source(
                fs.get_mapper(root="data/S2A_MSIL1C_20240201.zarr")
            )
        )

    def test_is_no_valid_source(self):
        self.assertFalse(self.mode.is_valid_source("data/S2A_MSIL2A_20240201.zarr"))
        self.assertFalse(self.mode.is_valid_source(dict()))

    def test_transform_datatree(self):
        self.assert_transform_datatree_ok(make_s2_msi_l1c())

    def test_convert_datatree(self):
        self.assert_convert_datatree_ok(
            make_s2_msi_l1c(r10m_size=48),
            expected_var_names=[
                "b01",
                "b02",
                "b03",
                "b04",
                "b05",
                "b06",
                "b07",
                "b08",
                "b09",
                "b10",
                "b11",
                "b12",
                "b8a",
            ],
            expected_size=48,
        )

    def test_convert_datatree_common_bandname(self):
        self.assert_convert_datatree_ok(
            make_s2_msi_l1c(r10m_size=48),
            expected_var_names="blue",
            expected_size=48,
        )

    def test_convert_datatree_common_bandnames(self):
        self.assert_convert_datatree_ok(
            make_s2_msi_l1c(r10m_size=48),
            expected_var_names=[
                "coastal",
                "blue",
                "green",
                "red",
                "rededge071",
                "rededge075",
                "rededge078",
                "nir",
                "nir08",
                "nir09",
                "cirrus",
                "swir16",
                "swir22",
            ],
            expected_size=48,
        )

    def test_convert_datatree_no_refds(self):
        self.assert_convert_datatree_ok(
            make_s2_msi_l1c(r10m_size=10000),
            expected_var_names=[
                "b05",
                "b06",
                "b07",
                "b8a",
            ],
            expected_size=10000,
        )

    def test_convert_datatree_bbox(self):
        self.assert_convert_datatree_ok(
            make_s2_msi_l1c(r10m_size=10000),
            expected_var_names=[
                "b02",
                "b03",
                "b04",
            ],
            expected_size=2000,
            bbox=[20000, 40000, 40000, 60000],
        )

    def test_convert_datatree_crs(self):
        self.assert_convert_datatree_ok(
            make_s2_msi_l1c(r10m_size=10000),
            expected_var_names=[
                "b02",
                "b03",
                "b04",
            ],
            expected_size=(10052, 9986),
            crs=pyproj.CRS.from_string("EPSG:4326"),
        )

    def test_convert_datatree_fail(self):
        self.assert_convert_datatree_fail(make_s2_msi_l1c(r10m_size=48))


class MsiL2aTest(MsiTestMixin, TestCase):
    mode = MsiL2a()

    def test_is_valid_source(self):
        self.assertTrue(self.mode.is_valid_source("S2A_MSIL2A_20240201.zarr"))
        self.assertFalse(self.mode.is_valid_source("S2A_MSIL1C_20240201.zarr"))
        self.assertFalse(self.mode.is_valid_source(dict()))

    def test_transform_datatree(self):
        self.assert_transform_datatree_ok(make_s2_msi_l2a())

    def test_convert_datatree(self):
        self.assert_convert_datatree_ok(
            make_s2_msi_l2a(r10m_size=48),
            expected_var_names=[
                "b01",
                "b02",
                "b03",
                "b04",
                "b05",
                "b06",
                "b07",
                "b08",
                "b11",
                "b12",
                "b8a",
                "cld",
                "scl",
                "snw",
            ],
            expected_size=48,
        )

    def test_convert_datatree_common_bandnames(self):
        self.assert_convert_datatree_ok(
            make_s2_msi_l2a(r10m_size=48),
            expected_var_names=[
                "coastal",
                "blue",
                "green",
                "red",
                "rededge071",
                "rededge075",
                "rededge078",
                "nir",
                "nir08",
                "swir16",
                "swir22",
                "cld",
                "scl",
                "snw",
            ],
            expected_size=48,
        )

    def test_convert_datatree_no_refds(self):
        self.assert_convert_datatree_ok(
            make_s2_msi_l2a(r10m_size=10000),
            expected_var_names=[
                "b05",
                "b06",
                "b07",
                "b8a",
                "cld",
                "scl",
                "snw",
            ],
            expected_size=10000,
        )

    def test_convert_datatree_bbox(self):
        self.assert_convert_datatree_ok(
            make_s2_msi_l2a(r10m_size=10000),
            expected_var_names=[
                "b02",
                "b03",
                "b04",
            ],
            expected_size=2000,
            bbox=[20000, 40000, 40000, 60000],
        )

    def test_convert_datatree_fail(self):
        self.assert_convert_datatree_fail(make_s2_msi_l2a(r10m_size=48))
