"""
Vertical extrapolation
===========================

Example of using PyWAsP to perform vertical extrapolation.

"""

# %%
# Prepare TopographyMap
# ------------------------
# First, we need to prepare the topography map. This is done by reading the
# elevation and roughness maps and creating a TopographyMap object.

import numpy as np
import windkit as wk
import pywasp as pw

ssl = wk.load_tutorial_data("serra_santa_luzia")
# The wind directions are relative to the grid north of the UTM 29N maps
bwc = ssl.bwc.assign_attrs(wind_dir_crs="EPSG:32629")

topo_map = pw.wasp.TopographyMap(ssl.elev, ssl.rgh)


# %%
# Define output locations
# ------------------------
# Second, we need to define the output locations. This is done by creating a
# dataset with the coordinates of the output locations.

output_locs = wk.spatial.create_dataset(
    west_east=bwc.west_east.values,
    south_north=bwc.south_north.values,
    height=np.linspace(10.0, 100.0, 10),
    crs="EPSG:32629",
    struct="stacked_point",
)
print(output_locs)

# %%
# Calculate resource grid
# ------------------------
# Finally, we can calculate the resource grid. This is done by calling the
# predict_wwc function. This function takes the output locations,
# the boundary conditions, and the topography map as input.
# The output is a weibull wind climate dataset at the output locations.

wwc = pw.wasp.predict_wwc(bwc, topo_map, output_locs)


# %%
# Plot the mean wind speed
# ------------------------
# We can plot the mean wind speed to see the result.

wwc["wspd"].plot.line(y="height")