# %%
import geowombat as gw
import rasterio
from geowombat.data import l8_224077_20200518_B2, l8_224078_20200518_B2
import matplotlib.pyplot as plt

fig, ax = plt.subplots(dpi=200)

with gw.open(
    [l8_224077_20200518_B2, l8_224078_20200518_B2],
    mosaic=True,
    bounds_by="union",
) as src:
    print(src)
src.where(src != 0).gw.imshow(robust=True, ax=ax)


# %%
fig, ax = plt.subplots(dpi=200)

with gw.open([l8_224077_20200518_B2], mosaic=True) as src:
    attrs = src.attrs
    src = src.where(src > 0)
    print(src)
    src.attrs = attrs
    src.gw.save(
        "l8_224077_20200518_B2.tif",
    )

src.gw.imshow(robust=True, ax=ax)
# %%
fig, ax = plt.subplots(dpi=200)

with gw.open(
    [l8_224078_20200518_B2],
    mosaic=True,
) as src:
    attrs = src.attrs
    src = src.where(src > 0)
    print(src)
    src.attrs = attrs
    src.gw.save(
        "l8_224078_20200518_B2.tif",
    )
src.gw.imshow(robust=True, ax=ax)


# %%
fig, ax = plt.subplots(dpi=200)
import numpy as np
with gw.open(
    ["l8_224077_20200518_B2.tif", "l8_224078_20200518_B2.tif"],
    mosaic=True,
    bounds_by="union",
    nodata=np.nan,
) as src:
    print(src)
src.gw.imshow(robust=True, ax=ax)

# %%