# import hydromt and setup logging
import geopandas as gpd
import pandas as pd
import hydromt
from hydromt.log import setuplog

logger = setuplog("read vector data", log_level=10)

# Download artifacts for the Piave basin to `~/.hydromt_data/`.
data_catalog = hydromt.DataCatalog(logger=logger, data_libs=["artifact_data"])

# inspect data source entry in data catalog yaml file
data_catalog["gadm_level3"]

gdf = data_catalog.get_geodataframe("gadm_level3")
print(f"number of rows: {gdf.index.size}")
gdf.head()

gdf_subset = data_catalog.get_geodataframe(
    "gadm_level3", bbox=gdf[:5].total_bounds, variables=["GID_0", "NAME_3"]
)
print(f"number of rows: {gdf_subset.index.size}")
gdf_subset.head()

# create example point CSV data with funny `x` coordinate name and additional column
import numpy as np
import pandas as pd

fn = "tmpdir/xy.csv"
df = pd.DataFrame(
    columns=["x_centroid", "y"],
    data=np.vstack([gdf_subset.centroid.x, gdf_subset.centroid.y]).T,
)
df["name"] = gdf_subset["NAME_3"]
df.to_csv(fn)  # write to file
df.head()

# Create data source entry for the data catalog for the new csv data
# NOTE that we add specify the name of the x coordinate with the `x_dim` argument, while the y coordinate is understood by HydroMT.
data_source = {
    "GADM_level3_centroids": {
        "path": fn,
        "data_type": "GeoDataFrame",
        "driver": "vector_table",
        "crs": 4326,
        "driver_kwargs": {"x_dim": "x_centroid"},
    }
}
data_catalog.from_dict(data_source)
data_catalog["GADM_level3_centroids"]

# we can then read the data back as a GeoDataFrame
gdf_centroid = data_catalog.get_geodataframe("GADM_level3_centroids")
print(f"CRS: {gdf_centroid.crs}")
gdf_centroid.head()

# m = gdf.explore(width='20%', height='50%')
# gdf_subset.explore(m=m, color='red')  # subset in red
# m

ax = gdf.plot()
gdf_subset.plot(ax=ax, color="red")
gdf_centroid.plot(ax=ax, color="k")