# ---
# jupyter:
#   jupytext:
#     formats: ipynb,py:percent
#     text_representation:
#       extension: .py
#       format_name: percent
#       format_version: '1.3'
#       jupytext_version: 1.13.0
#   kernelspec:
#     display_name: ODC
#     language: python
#     name: odc
# ---

# %%
import json
import time
from timeit import default_timer as t_now

import datacube
import numpy as np
import odc.stac
import planetary_computer as pc
import pystac.item
import pystac_client
import rioxarray
import stackstac
import xarray as xr
from distributed import Client
from distributed import wait as dask_wait
from odc.stac.bench import (
    BenchLoadParams,
    collect_context_info,
    load_from_json,
    run_bench,
)

# generated by bench-prep-query notebook
fname = [
    "site1-20200606-tall-strip-africa.geojson",
    "site2-2020_jun_jul-35MNM.geojson",
][0]

print(f"Load from file: {fname}")
all_features = json.load(open(fname, "rt"))
print(json.dumps(all_features.get("properties", {}), indent=2))

# %%
params = BenchLoadParams(
    scenario=fname,
    method="stackstac",
    bands=("B02", "B03", "B04"),
    chunks=(2048, 2048),
    patch_url=pc.sign,
    extra={
        "stackstac": {"dtype": "uint16", "fill_value": 0},
        "odc-stac": {"groupby": "solar_day", "stac_cfg": {"*": {"warnings": "ignore"}}},
    },
    # resolution = (2**5)*10,
    # crs="epsg:3857",
)

xx = load_from_json(
    all_features,
    params.with_method("stackstac"),
)

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

assert xx.spec.transform == xx.geobox.transform
assert xx.geobox == yy.geobox
xx

# %%
client = Client("tcp://127.0.0.1:8786")
if len(client.futures) > 0:
    print("Restarting Client")
    client.restart()
display(client)

# %%
rr_odc, results_odc = run_bench(yy, client, ntimes=3)

# %%
rr_stc, results_stc = run_bench(xx, client, ntimes=3)
