# %%
# Generate the horizontal mesh
# ----------------------------
# This is the second step of the Great Barrier Reef workflow. It uses the static
# datasets downloaded by download.py to build an unstructured horizontal mesh.
#
# The mesh-size field is refined near reefs, coastlines, and narrow openings,
# while also respecting a bathymetry-based stability constraint. The resulting
# mesh is written to data/gbr.msh and is used by the preprocessing step.

import numpy as np
import pyogrio
import pyproj
import seamsh
import shapely
from osgeo import osr

CurveType = seamsh.geometry.CurveType
gmsh = seamsh.gmsh.gmsh


# %%
# The computational mesh is generated in a projected coordinate system, while
# some of the the input boundary curves are specified in longitude/latitude
# coordinates.

proj = osr.SpatialReference()
proj.ImportFromProj4("+ellps=WGS84 +proj=utm +zone=55 +south")

projlonlat = osr.SpatialReference()
projlonlat.ImportFromEPSG(4326)
projlonlat.SetAxisMappingStrategy(osr.OAMS_TRADITIONAL_GIS_ORDER)


# %%
# The bathymetry raster and reef polygons are used by the mesh-size field. The
# reef distance query is cached because it is called many times during meshing.

bath_interp = seamsh.field.Raster("data/bathymetry/merged/merge_d100.tif")
reefs = "data/reefs.gpkg"


def reefs_dist(xy, projection):
    if not hasattr(reefs_dist, "tree"):
        gdf = pyogrio.read_dataframe(reefs)
        gdf = gdf.to_crs(pyproj.CRS.from_wkt(proj.ExportToWkt()))
        polys = gdf.explode(index_parts=False).geometry.to_numpy()
        reefs_dist.tree = shapely.STRtree(polys)
    x, y = xy.T
    out = np.full(x.shape[0], np.nan, dtype="float64")
    pts = shapely.points(x, y)
    [pt_idx, poly_idx], dist = reefs_dist.tree.query_nearest(pts, return_distance=True)
    out[pt_idx] = dist
    return out


# %%
# We define two open-boundary curves that close the offshore side of the domain:
# one through the Coral Sea and one through Torres Strait. Coastline curves are
# read from the GeoPackage created in download.py.

coral_sea = np.array(
    [
        (152.909557, -26.515829),
        (154.342362, -25.847186),
        (155.297566, -25.297944),
        (155.934369, -24.255180),
        (156.181130, -23.077095),
        (156.006009, -20.784606),
        (155.361246, -19.312000),
        (154.676684, -18.420477),
        (153.355319, -16.482209),
        (152.081714, -14.094200),
        (150.999149, -11.761911),
        (150.541448, -10.551986),
    ]
)

torres_strait = np.array(
    [
        (142.306795, -9.097290),
        (142.408286, -9.383851),
        (142.489876, -9.674392),
        (142.547586, -9.978864),
        (142.573456, -10.14403),
        (142.557536, -10.37686),
        (142.529676, -10.60372),
        (142.444106, -10.86441),
        (142.396345, -10.98381),
    ]
)

domain = seamsh.geometry.Domain(proj)
domain.add_boundary_curve(coral_sea, "coral_sea", projlonlat, CurveType.BSPLINE)
domain.add_boundary_curve(torres_strait, "torres_strait", projlonlat, CurveType.BSPLINE)
domain.add_boundary_curves_shp("data/contour/coastlines.gpkg", "physical")

coast_dist = seamsh.field.Distance(domain, 200, ["coast"])
coral_dist = seamsh.field.Distance(domain, 200, ["coral_sea"])
torres_strait_dist = seamsh.field.Distance(domain, 200, ["torres_strait"])


# %%
# The following functions define the spatially varying target mesh size. The
# final size is the result of geometric refinement constraints and a minimum size
# imposed by the local water depth and target time step.


def expogrowth(d, dnear: float, dgrowth: float, growthfactor: float):
    return np.maximum(np.exp((np.log(growthfactor) / dgrowth) * (d - dnear)), 1)


def mesh_size(x, projection):
    dreef = reefs_dist(x, projection)
    dreef = np.clip(dreef, 0, 1e5)

    lc = 10000  # largest mesh size
    f = np.ones(x.shape[0])

    eg = expogrowth(dreef, 600, 2000, 1.5)
    freef = 200 / lc * eg
    f = np.minimum(f, freef)

    c_dist = coral_dist(x, projection)
    fopen = 1 / expogrowth(c_dist, lc, lc, 2)
    f = np.maximum(f, fopen)

    ts_dist = torres_strait_dist(x, projection)
    fopen = 1 / expogrowth(ts_dist, 4000, 4000, 2) * (4000 / lc)
    f = np.maximum(f, fopen)

    size = f * lc
    size = np.minimum(size, lc)

    bath = -bath_interp(x, projection)
    bath = np.clip(bath, 2, np.inf)
    g = 9.81
    dtmax = 7

    sizemin = np.sqrt(12 * g * bath) * dtmax
    sizemin = np.clip(np.clip(sizemin, 200, np.inf), 0, lc)
    sizemax = np.sqrt(12 * g * bath) * dtmax * 5

    size = np.minimum(size, sizemax)
    size = np.maximum(size, sizemin)
    return size * 10


# %%
# Finally, we coarsen the input boundaries according to the same size field and
# call gmsh through seamsh. The resulting mesh is the horizontal mesh used by
# prepro.py.

coarse_initial = seamsh.geometry.coarsen_boundaries(
    domain, (145.440, -14.746), projlonlat, mesh_size
)
seamsh.gmsh.mesh(coarse_initial, "data/gbr.msh", mesh_size, version=4.1, smoothness=0.3)
# seamsh.gmsh.convert_to_gis("data/gbr.msh", proj, "data/gbr.gpkg")