# This file is part of the Open Data Cube, see https://opendatacube.org for more information
#
# Copyright (c) 2015-2023 ODC Contributors
# SPDX-License-Identifier: Apache-2.0
import pytest
import xarray as xr
import pandas as pd
import numpy as np
from datacube.testutils.geom import epsg4326, epsg3857
from datacube.testutils import mk_sample_xr_dataset, remove_crs
from datacube.utils.geometry import assign_crs
from datacube.utils.xarray_geoextensions import (
    _norm_crs,
    _xarray_affine,
    _xarray_geobox,
    _xarray_extent,
    _get_crs_from_attrs,
    _get_crs_from_coord,
)

multi_coords = xr.DataArray(
    np.zeros(1),
    [
        (
            "spec",
            pd.MultiIndex.from_arrays(
                np.array([["2001-01-01"], ["2001-01-01"]]), names=("time", "solar_day")
            ),
        )
    ],
).coords
single_coord = dict(time=np.array(["2001-01-01"]))


@pytest.mark.parametrize(
    "odc_style_xr_dataset", [single_coord, multi_coords], indirect=True
)
def test_xr_extension(odc_style_xr_dataset):
    xx = odc_style_xr_dataset
    assert _norm_crs(None) is None
    assert _norm_crs(epsg4326) is epsg4326
    assert _norm_crs(str(epsg4326)) == epsg4326

    with pytest.raises(ValueError):
        _norm_crs([])

    assert (1,) + xx.geobox.shape == xx.B10.shape

    (sx, zz0, tx, zz1, sy, ty) = xx.affine[:6]
    assert (zz0, zz1) == (0, 0)

    xx = remove_crs(xx)

    assert _xarray_geobox(xx) is None
    assert _xarray_extent(xx) is None

    # affine should still be valid
    A = _xarray_affine(xx)
    assert A is not None
    assert A * (0.5, 0.5) == (xx.longitude[0], xx.latitude[0])
    assert A * (0.5, 1.5) == (xx.longitude[0], xx.latitude[1])


def test_xr_geobox():
    xy = (10, 111)
    rxy = (10, -100)
    resolution = rxy[::-1]

    ds = mk_sample_xr_dataset(crs=epsg3857, xy=xy, resolution=resolution)

    assert ds.geobox.crs == epsg3857
    assert ds.band.geobox.crs == epsg3857
    assert ds.band.affine * (0, 0) == xy
    assert ds.band.affine * (1, 1) == tuple(a + b for a, b in zip(xy, rxy))

    assert ds.band[:, :2, :2].affine * (0, 0) == xy
    assert ds.band[:, :2, :2].affine * (1, 1) == tuple(a + b for a, b in zip(xy, rxy))
    assert ds.band.isel(time=0, x=0).affine is None

    xx = ds.band + 1000
    assert xx.geobox is not None
    assert xx.geobox == ds.band.geobox

    assert mk_sample_xr_dataset(crs=epsg4326).geobox.crs == epsg4326
    assert mk_sample_xr_dataset(crs=epsg4326).band.geobox.crs == epsg4326

    assert mk_sample_xr_dataset(crs=None).geobox is None
    assert mk_sample_xr_dataset(crs=None).affine is not None


def test_xr_geobox_unhappy():
    xx = mk_sample_xr_dataset(crs=None)

    # test duplicates behaviour in get_crs_from_{coord,attrs} are caught
    xx.band.attrs.update(grid_mapping="x")  # force exception in coord
    xx.x.attrs.update(crs="EPSG:4326")  # force exception in attr
    xx.y.attrs.update(crs="EPSG:3857")
    with pytest.warns(UserWarning):
        assert xx.band.geobox is not None

    # test crs not being a string
    xx = mk_sample_xr_dataset(crs=None)
    xx.attrs["crs"] = ["this will not parse"]
    with pytest.warns(UserWarning):
        assert xx.geobox is None

    xx.attrs["crs"] = "this will fail CRS() constructor"
    with pytest.warns(UserWarning):
        assert xx.geobox is None

    xx = mk_sample_xr_dataset(crs=epsg3857)
    xx.attrs.pop("crs", None)
    xx.band.attrs.pop("crs", None)
    xx.spatial_ref.attrs["spatial_ref"] = "this will fail CRS() constructor"
    with pytest.warns(UserWarning):
        assert xx.geobox is None


def test_crs_from_coord():
    xx_none = mk_sample_xr_dataset(crs=None)
    xx_3857 = mk_sample_xr_dataset(crs=epsg3857)
    xx_4326 = mk_sample_xr_dataset(crs=epsg4326)
    cc_4326 = xx_4326.geobox.xr_coords(with_crs="epsg_4326")["epsg_4326"]
    cc_3857 = xx_3857.geobox.xr_coords(with_crs="epsg_3857")["epsg_3857"]

    assert _get_crs_from_coord(xx_none.band) is None
    assert _get_crs_from_coord(xx_none) is None
    assert _get_crs_from_coord(xx_3857.band) == epsg3857
    assert _get_crs_from_coord(xx_3857) == epsg3857
    assert _get_crs_from_coord(xx_4326.band) == epsg4326
    assert _get_crs_from_coord(xx_4326) == epsg4326

    xx = xx_none.band.assign_attrs(grid_mapping="x")
    with pytest.raises(ValueError):
        _get_crs_from_coord(xx)
    xx = xx_none.band.assign_attrs(grid_mapping="no_such_coord")
    assert _get_crs_from_coord(xx) is None

    xx_2crs = xx_none.assign_coords(cc_4326=cc_4326, cc_3857=cc_3857)
    assert xx_2crs.geobox is None

    # two coords, no grid mapping, strict mode
    with pytest.raises(ValueError):
        _get_crs_from_coord(xx_2crs)
    with pytest.raises(ValueError):
        _get_crs_from_coord(xx_2crs.band)

    # any should just return "first" guess, we not sure which one
    crs = _get_crs_from_coord(xx_2crs, "any")
    assert epsg4326 == crs or epsg3857 == crs

    # all should return a list of length 2
    crss = _get_crs_from_coord(xx_2crs, "all")
    assert len(crss) == 2
    assert any(crs == epsg3857 for crs in crss)
    assert any(crs == epsg4326 for crs in crss)

    with pytest.raises(ValueError):
        _get_crs_from_coord(xx_2crs, "no-such-mode")


def test_crs_from_attrs():
    xx_none = mk_sample_xr_dataset(crs=None)
    xx_3857 = mk_sample_xr_dataset(crs=epsg3857)
    xx_4326 = mk_sample_xr_dataset(crs=epsg4326)

    assert _get_crs_from_attrs(xx_none, ("y", "x")) is None
    assert _get_crs_from_attrs(xx_none.band, ("y", "x")) is None
    assert _get_crs_from_attrs(xx_3857.band, ("y", "x")) == epsg3857
    assert _get_crs_from_attrs(xx_3857, ("y", "x")) == epsg3857
    assert _get_crs_from_attrs(xx_4326.band, ("latitude", "longitude")) == epsg4326
    assert _get_crs_from_attrs(xx_4326, ("latitude", "longitude")) == epsg4326

    assert _get_crs_from_attrs(xr.Dataset(), None) is None

    # check inconsistent CRSs
    xx = xx_3857.copy()
    xx.x.attrs["crs"] = xx_4326.attrs["crs"]
    with pytest.warns(UserWarning):
        _get_crs_from_attrs(xx, ("y", "x"))

    xx = xx_none.copy()
    xx.attrs["crs"] = epsg3857
    assert xx.geobox is not None
    assert xx.geobox.crs is epsg3857


@pytest.mark.parametrize("odc_style_xr_dataset", [{}, multi_coords], indirect=True)
def test_assign_crs(odc_style_xr_dataset):
    xx = odc_style_xr_dataset
    assert xx.geobox is not None

    xx_nocrs = remove_crs(xx)
    assert xx_nocrs.geobox is None
    yy = assign_crs(xx_nocrs, "epsg:4326")
    assert xx_nocrs.geobox is None  # verify source is not modified in place
    assert yy.geobox.crs == "epsg:4326"

    yy = assign_crs(xx_nocrs.B10, "epsg:4326")
    assert yy.geobox.crs == "epsg:4326"

    xx_xr_style_crs = xx_nocrs.copy()
    xx_xr_style_crs.attrs.update(crs="epsg:3857")
    yy = assign_crs(xx_xr_style_crs)
    assert yy.geobox.crs == "epsg:3857"

    with pytest.raises(ValueError):
        assign_crs(xx_nocrs)


def test_xr_opendataset(example_netcdf_path):
    xx = xr.open_dataset(example_netcdf_path)
    assert xx.geobox is None
