# IMPORT LIBRARIES
import sys
import pygmt
import rasterio
import xarray as xr
import numpy as np

# Show versions
print("Python version:", sys.version)
print("PyGMT version:", pygmt.__version__)
print("Rasterio version:", rasterio.__version__)
print("Xarray version:", xr.__version__)
print("NumPy version:", np.__version__)

Python version: 3.11.13 | packaged by conda-forge | (main, Jun  4 2025, 14:39:58) [MSC v.1943 64 bit (AMD64)]
PyGMT version: v0.17.0
Rasterio version: 1.4.3
Xarray version: 2025.10.1
NumPy version: 2.3.3

#Sri Lanka bounding box
region = [79.0, 82.5, 5.0, 10.5]

# Download elevation grid
grid = pygmt.datasets.load_earth_relief(resolution="30s", region=region)

# FIRST FIGURE
fig = pygmt.Figure()

# Create 3-D perspective view with grdview
fig.grdview(
    grid=grid,
    shading=True, # Default shading
    perspective=[130, 30], # Sets the view azimuth as 130 degrees, and the view angle as 30 degrees 
    frame=["xa", "ya", "SnE"], # Sets the x- and y-axis labels, and annotates the west, south, and east axes
    projection="M12c", # Sets a Mercator projection on a 12-centimeter figure
    zsize="2c", # Sets the height of the three-dimensional relief at 2 centimeters
    surftype="s",
    cmap="topo1",
    contourpen="0.5p",
)

# Add a vertical colorbar to the right
fig.colorbar(
    perspective=True,
    position="JMR+v+w12c/0.5c+o1c/0c",  # right side, vertical, width/height, offset
    frame=["a1000", "x+lElevation (m)"],
)

#Add map scale
fig.basemap(perspective=True, map_scale="jBM+w300k+o0c/-11c+f+l")

# Shift origin for second plot
fig.shift_origin(xshift="25c")

# SECOND FIGURE

# Create Xarray DataArray from image of Sri Lanka so it is usable by PyGMT
flag_to_image = "Flag_of_Sri_Lanka_vertical.png"
with rasterio.open(flag_to_image) as dataset:
    data = dataset.read()
    drapegrid = xr.DataArray(data, dims=("band", "y", "x"))

# Flip the image on both axes so it is facing the viewer
drapegrid_rot180 = drapegrid.isel(y=slice(None, None, -1), x=slice(None, None, -1))

# Plot 3D flag draped over relief
fig.grdview(
    projection="M12c",
    region=region,
    grid=grid,
    drapegrid=drapegrid_rot180,
    surftype="i",  # "image" mode needed
    perspective=[130, 30],
    zsize="1.5c",
    plane="+gdarkgray",
    frame=True,
)

# Create a masked version of the elecation grid where values above 0 become NaN (transparent)
grid_masked = grid.where(grid <= 0, np.nan)

#Plot elevation grid on top of image, omitting Sri Lanka land area
fig.grdview(
    projection="M12c",
    shading=True,
    region=region,
    grid=grid_masked,
    cmap=True,
    surftype="s", 
    perspective=[130, 30],
    zsize="1.5c",
    plane="+gdarkgray",
    frame=True,
)

#Show and save figure
#fig.show()
#fig.savefig("fountains_of_paradise.jpg")

C:\Users\timwu\anaconda3\envs\pygmt\Lib\site-packages\rasterio\__init__.py:356: NotGeoreferencedWarning: Dataset has no geotransform, gcps, or rpcs. The identity matrix will be returned.
  dataset = DatasetReader(path, driver=driver, sharing=sharing, **kwargs)
C:\Users\timwu\anaconda3\envs\pygmt\Lib\site-packages\rasterio\__init__.py:366: NotGeoreferencedWarning: The given matrix is equal to Affine.identity or its flipped counterpart. GDAL may ignore this matrix and save no geotransform without raising an error. This behavior is somewhat driver-specific.
  dataset = writer(