import pytest
import xarray
from distributed import Client

from odc.stac.bench import (
    BenchLoadParams,
    collect_context_info,
    load_from_json,
    run_bench,
)

CFG = {"*": {"warnings": "ignore"}}


@pytest.fixture(scope="module")
def dask_client():
    client = Client(
        n_workers=1,
        threads_per_worker=2,
        memory_limit="500MiB",
        local_directory="/tmp/",
        memory_target_fraction=False,
        memory_spill_fraction=False,
        memory_pause_fraction=False,
        dashboard_address=None,
        processes=False,
    )
    yield client
    client.shutdown()
    del client


def test_load_from_json_stackstac(fake_dask_client, bench_site1, bench_site2):
    dask_client = fake_dask_client
    params = BenchLoadParams(
        scenario="test1",
        method="stackstac",
        bands=("B04", "B02", "B03"),
        chunks=(2048, 2048),
        resampling="nearest",
        extra={
            "odc-stac": {"groupby": "solar_day", "stac_cfg": CFG},
            "stackstac": {"dtype": "uint16", "fill_value": 0},
        },
    )
    xx = load_from_json(bench_site1, params)
    assert "band" in xx.dims
    assert xx.shape == (1, 3, 90978, 10980)
    assert xx.dtype == "uint16"
    assert xx.spec.epsg == 32735

    yy = load_from_json(bench_site1, params.with_method("odc-stac"))

    rrx = collect_context_info(dask_client, xx)
    rry = collect_context_info(dask_client, yy)
    assert rrx.shape == rry.shape
    assert rrx == rry

    xx = load_from_json(bench_site2, params)
    assert "band" in xx.dims
    assert xx.dtype == "uint16"
    assert xx.spec.epsg == 32735

    params.crs = "epsg:32736"
    xx = load_from_json(bench_site2, params)
    assert "band" in xx.dims
    assert xx.dtype == "uint16"
    assert xx.spec.epsg == 32736

    with pytest.raises(ValueError):
        load_from_json(bench_site1, params.with_method("wroNg"))


def test_bench_context(fake_dask_client, bench_site1, bench_site2):
    params = BenchLoadParams(
        scenario="test1",
        method="odc-stac",
        bands=("red", "green", "blue"),
        chunks=(2048, 2048),
        extra={"odc-stac": {"groupby": "solar_day", "stac_cfg": CFG}},
    )
    xx = load_from_json(bench_site1, params)
    nt, ny, nx = xx.red.shape
    nb = len(xx.data_vars)

    # Check normal case Dataset, with time coords
    rr = collect_context_info(
        fake_dask_client, xx, method=params.method, scenario="site1"
    )
    assert rr.shape == (nt, nb, ny, nx)
    assert rr.chunks == (1, 1, 2048, 2048)
    assert rr.crs == f"epsg:{xx.geobox.crs.epsg}"
    assert rr.crs == xx.geobox.crs
    assert rr.nthreads == 2
    assert rr.total_ram == 500 * (1 << 20)

    header_txt = rr.render_txt()
    assert "T.slice   : 2020-06-06" in header_txt
    assert "Data      : 1.3.90978.10980.uint16,  5.58 GiB" in header_txt

    run_txt = rr.render_timing_info((0, 0.1, 30))

    pd_dict = rr.to_pandas_dict()
    assert pd_dict["resolution"] == rr.resolution
    assert pd_dict["data"] == "1.3.90978.10980.uint16"
    assert pd_dict["chunks_x"] == 2048
    assert pd_dict["chunks_y"] == 2048

    # Check DataArray case
    rr = collect_context_info(
        fake_dask_client, xx.red, method="odc-stac", scenario="site1"
    )
    assert rr.shape == (nt, 1, ny, nx)
    assert rr.crs == xx.geobox.crs

    # Check Dataset with 0 dimension time axis and extras field
    rr = collect_context_info(
        fake_dask_client,
        xx.isel(time=0),
        method=params.method,
        scenario=params.scenario,
        extras={"custom": 2},
    )
    assert rr.extras == {"custom": 2}
    assert rr.shape == (1, nb, ny, nx)

    header_txt = rr.render_txt()
    assert "GEO       : epsg:32735" in header_txt
    assert "T.slice   : 2020-06-06" in header_txt

    # Check no time info at all
    rr = collect_context_info(
        fake_dask_client,
        xx.isel(time=0, drop=True),
        method=params.method,
        scenario=params.scenario,
    )
    assert rr.shape == (nt, nb, ny, nx)
    assert rr.dtype == xx.red.dtype
    assert rr.temporal_id == "-"

    # Check wrong type
    with pytest.raises(ValueError):
        collect_context_info(fake_dask_client, "wrong input type")  # type: ignore

    # Check multi-time axis
    xx = load_from_json(bench_site2, params)
    nt, ny, nx = xx.red.shape
    nb = len(xx.data_vars)

    assert nt > 1

    rr = collect_context_info(
        fake_dask_client,
        xx,
        method=params.method,
        scenario=params.scenario,
    )
    assert rr.shape == (nt, nb, ny, nx)
    assert rr.temporal_id == "2020-06-01__2020-07-31"

    # Check missing GEO info
    no_geo = _strip_geo(xx.red)
    assert no_geo.geobox is None
    with pytest.raises(ValueError):
        # no geobox
        collect_context_info(fake_dask_client, no_geo)


def _strip_geo(xx: xarray.DataArray) -> xarray.DataArray:
    no_geo = xx.drop_vars("spatial_ref")
    no_geo.attrs.pop("crs", None)
    no_geo.attrs.pop("grid_mapping", None)
    no_geo.x.attrs.pop("crs", None)
    no_geo.y.attrs.pop("crs", None)
    assert no_geo.geobox is None
    return no_geo


def test_run_bench(fake_dask_client, bench_site1, capsys):
    dask_client = fake_dask_client
    params = BenchLoadParams(
        scenario="test1",
        method="odc-stac",
        bands=("red", "green", "blue"),
        chunks=(2048, 2048),
        extra={"odc-stac": {"groupby": "solar_day", "stac_cfg": CFG}},
    )
    xx = load_from_json(bench_site1, params)

    rr, timing = run_bench(xx, dask_client, 10)

    assert rr.scenario == params.scenario
    assert rr.method == params.method
    assert len(timing) == 10
    _io = capsys.readouterr()
    assert len(_io.out) > 0


def test_bench_params_json():
    params = BenchLoadParams(
        scenario="test1",
        method="odc-stac",
        bands=("red", "green", "blue"),
        chunks=(100, 200),
        extra={"odc-stac": {"groupby": "solar_day", "stac_cfg": CFG}},
    )

    assert params == BenchLoadParams.from_json(params.to_json())
    assert params.to_json() == BenchLoadParams.from_json(params.to_json()).to_json()

    # function should round-trip too
    params.patch_url = load_from_json
    assert params == BenchLoadParams.from_json(params.to_json())

import pytest
import xarray
from distributed import Client

from odc.stac.bench import (
    BenchLoadParams,
    collect_context_info,
    load_from_json,
    run_bench,
)

CFG = {"*": {"warnings": "ignore"}}


@pytest.fixture(scope="module")
def dask_client():
    client = Client(
        n_workers=1,
        threads_per_worker=2,
        memory_limit="500MiB",
        local_directory="/tmp/",
        memory_target_fraction=False,
        memory_spill_fraction=False,
        memory_pause_fraction=False,
        dashboard_address=None,
        processes=False,
    )
    yield client
    client.shutdown()
    del client


def test_load_from_json_stackstac(fake_dask_client, bench_site1, bench_site2):
    dask_client = fake_dask_client
    params = BenchLoadParams(
        scenario="test1",
        method="stackstac",
        bands=("B04", "B02", "B03"),
        chunks=(2048, 2048),
        resampling="nearest",
        extra={
            "odc-stac": {"groupby": "solar_day", "stac_cfg": CFG},
            "stackstac": {"dtype": "uint16", "fill_value": 0},
        },
    )
    xx = load_from_json(bench_site1, params)
    assert "band" in xx.dims
    assert xx.shape == (1, 3, 90978, 10980)
    assert xx.dtype == "uint16"
    assert xx.spec.epsg == 32735

    yy = load_from_json(bench_site1, params.with_method("odc-stac"))

    rrx = collect_context_info(dask_client, xx)
    rry = collect_context_info(dask_client, yy)
    assert rrx.shape == rry.shape
    assert rrx == rry

    xx = load_from_json(bench_site2, params)
    assert "band" in xx.dims
    assert xx.dtype == "uint16"
    assert xx.spec.epsg == 32735

    params.crs = "epsg:32736"
    xx = load_from_json(bench_site2, params)
    assert "band" in xx.dims
    assert xx.dtype == "uint16"
    assert xx.spec.epsg == 32736

    with pytest.raises(ValueError):
        load_from_json(bench_site1, params.with_method("wroNg"))


def test_bench_context(fake_dask_client, bench_site1, bench_site2):
    params = BenchLoadParams(
        scenario="test1",
        method="odc-stac",
        bands=("red", "green", "blue"),
        chunks=(2048, 2048),
        extra={"odc-stac": {"groupby": "solar_day", "stac_cfg": CFG}},
    )
    xx = load_from_json(bench_site1, params)
    nt, ny, nx = xx.red.shape
    nb = len(xx.data_vars)

    # Check normal case Dataset, with time coords
    rr = collect_context_info(
        fake_dask_client, xx, method=params.method, scenario="site1"
    )
    assert rr.shape == (nt, nb, ny, nx)
    assert rr.chunks == (1, 1, 2048, 2048)
    assert rr.crs == f"epsg:{xx.geobox.crs.epsg}"
    assert rr.crs == xx.geobox.crs
    assert rr.nthreads == 2
    assert rr.total_ram == 500 * (1 << 20)

    header_txt = rr.render_txt()
    assert "T.slice   : 2020-06-06" in header_txt
    assert "Data      : 1.3.90978.10980.uint16,  5.58 GiB" in header_txt

    run_txt = rr.render_timing_info((0, 0.1, 30))

    pd_dict = rr.to_pandas_dict()
    assert pd_dict["resolution"] == rr.resolution
    assert pd_dict["data"] == "1.3.90978.10980.uint16"
    assert pd_dict["chunks_x"] == 2048
    assert pd_dict["chunks_y"] == 2048

    # Check DataArray case
    rr = collect_context_info(
        fake_dask_client, xx.red, method="odc-stac", scenario="site1"
    )
    assert rr.shape == (nt, 1, ny, nx)
    assert rr.crs == xx.geobox.crs

    # Check Dataset with 0 dimension time axis and extras field
    rr = collect_context_info(
        fake_dask_client,
        xx.isel(time=0),
        method=params.method,
        scenario=params.scenario,
        extras={"custom": 2},
    )
    assert rr.extras == {"custom": 2}
    assert rr.shape == (1, nb, ny, nx)

    header_txt = rr.render_txt()
    assert "GEO       : epsg:32735" in header_txt
    assert "T.slice   : 2020-06-06" in header_txt

    # Check no time info at all
    rr = collect_context_info(
        fake_dask_client,
        xx.isel(time=0, drop=True),
        method=params.method,
        scenario=params.scenario,
    )
    assert rr.shape == (nt, nb, ny, nx)
    assert rr.dtype == xx.red.dtype
    assert rr.temporal_id == "-"

    # Check wrong type
    with pytest.raises(ValueError):
        collect_context_info(fake_dask_client, "wrong input type")  # type: ignore

    # Check multi-time axis
    xx = load_from_json(bench_site2, params)
    nt, ny, nx = xx.red.shape
    nb = len(xx.data_vars)

    assert nt > 1

    rr = collect_context_info(
        fake_dask_client,
        xx,
        method=params.method,
        scenario=params.scenario,
    )
    assert rr.shape == (nt, nb, ny, nx)
    assert rr.temporal_id == "2020-06-01__2020-07-31"

    # Check missing GEO info
    no_geo = _strip_geo(xx.red)
    assert no_geo.geobox is None
    with pytest.raises(ValueError):
        # no geobox
        collect_context_info(fake_dask_client, no_geo)


def _strip_geo(xx: xarray.DataArray) -> xarray.DataArray:
    no_geo = xx.drop_vars("spatial_ref")
    no_geo.attrs.pop("crs", None)
    no_geo.attrs.pop("grid_mapping", None)
    no_geo.x.attrs.pop("crs", None)
    no_geo.y.attrs.pop("crs", None)
    assert no_geo.geobox is None
    return no_geo


def test_run_bench(fake_dask_client, bench_site1, capsys):
    dask_client = fake_dask_client
    params = BenchLoadParams(
        scenario="test1",
        method="odc-stac",
        bands=("red", "green", "blue"),
        chunks=(2048, 2048),
        extra={"odc-stac": {"groupby": "solar_day", "stac_cfg": CFG}},
    )
    xx = load_from_json(bench_site1, params)

    rr, timing = run_bench(xx, dask_client, 10)

    assert rr.scenario == params.scenario
    assert rr.method == params.method
    assert len(timing) == 10
    _io = capsys.readouterr()
    assert len(_io.out) > 0


def test_bench_params_json():
    params = BenchLoadParams(
        scenario="test1",
        method="odc-stac",
        bands=("red", "green", "blue"),
        chunks=(100, 200),
        extra={"odc-stac": {"groupby": "solar_day", "stac_cfg": CFG}},
    )

    assert params == BenchLoadParams.from_json(params.to_json())
    assert params.to_json() == BenchLoadParams.from_json(params.to_json()).to_json()

    # function should round-trip too
    params.patch_url = load_from_json
    assert params == BenchLoadParams.from_json(params.to_json())

# Check normal case Dataset, with time coords # Check DataArray case # Check Dataset with 0 dimension time axis and extras field # Check no time info at all # Check wrong type # type: ignore # Check multi-time axis # Check missing GEO info # no geobox # function should round-trip too