"""
Obstacle speedups
=================

This example builds flat elevation and roughness maps from scratch, adds a
building with an interior courtyard and a porous row of trees, and plots the
obstacle speedup over a resource grid.
"""

# %%
# Create simple maps
# ------------------
# Flat terrain with a single roughness length, so the only site effect in this
# example is the shelter from the obstacles.

import geopandas as gpd
import matplotlib.pyplot as plt
import numpy as np
import windkit as wk
from shapely.geometry import Polygon, box

import pywasp as pw


crs = "EPSG:32632"
center_x, center_y = 500_000.0, 6_200_000.0
half_width = 1_000.0
bbox = wk.spatial.BBox.from_cornerpts(
    minx=center_x - half_width,
    miny=center_y - half_width,
    maxx=center_x + half_width,
    maxy=center_y + half_width,
    crs=crs,
)

elev_map = wk.create_vector_map(bbox, map_type="elevation", elevation=0.0)
roughness_map = wk.create_vector_map(
    bbox, map_type="roughness", roughness_change=(0.05, 0.05)
)


# %%
# Add a building and a porous tree row
# ------------------------------------
# The 60 m x 40 m building is solid and has an interior courtyard. The row of
# trees south of it has the same 12 m height but 50% porosity, so the two
# wakes can be compared directly.

building = Polygon(
    box(center_x - 90, center_y + 20, center_x - 30, center_y + 60).exterior.coords,
    holes=[
        box(center_x - 70, center_y + 30, center_x - 50, center_y + 50).exterior.coords
    ],
)
tree_row = box(center_x - 65, center_y - 110, center_x - 55, center_y - 10)
obstacle_map = gpd.GeoDataFrame(
    {
        "height": [12.0, 12.0],
        "porosity": [0.0, 0.5],
    },
    geometry=[building, tree_row],
    crs=crs,
)

topography = pw.wasp.TopographyMap(
    elev_map,
    roughness_map,
    obstacle_map=obstacle_map,
)


# %%
# Calculate and plot the obstacle speedup
# ---------------------------------------
# Site effects are directional; this plot shows the sector with wind from the
# west at 6 m above ground, half the obstacle height. The solid building gives
# a stronger and longer wake than the porous tree row, and the courtyard is
# sheltered by the surrounding building. Grid points inside the building and
# below its height have no obstacle speedup (NaN, with a warning), so they
# are blank. This ray- and sector-based view is consistent with the
# WAsP-shelter description in Peña et al. (2015), which discusses directional
# wake spreading and the reduction of shelter with increasing porosity. See
# `Shelter models and observations
# <https://backend.orbit.dtu.dk/ws/portalfiles/portal/122110497/Report_WP3_DTUversion.pdf>`_.
#
# Increase ``NR_POINTS_RESOURCE_GRID`` for a smoother picture; this also
# increases the number of site-effect evaluations and therefore the run time.
NR_POINTS_RESOURCE_GRID = 30
HEIGHT = 6.0
grid = wk.spatial.create_cuboid(
    west_east=np.linspace(center_x - 150, center_x + 250, NR_POINTS_RESOURCE_GRID),
    south_north=np.linspace(center_y - 200, center_y + 200, NR_POINTS_RESOURCE_GRID),
    height=[HEIGHT],
    crs=crs,
)
site_effects = topography.get_site_effects(grid, n_sectors=12)
obstacle_speedup = site_effects["obstacle_speedups"].sel(height=HEIGHT, sector=270)

fig, ax = plt.subplots(figsize=(8, 6))
obstacle_speedup.plot(ax=ax, x="west_east", y="south_north", cmap="RdYlBu_r")
obstacle_map.boundary.plot(ax=ax, color="black", linewidth=1.5)
ax.set_title(f"Obstacle speedup in sector 270° at {HEIGHT:g} m a.g.l.")
ax.set_aspect("equal")
plt.tight_layout()
plt.show()