# %%
# GBR tools
# ---------
# This contans additional functions used in the GBR testcase

from typing import List
import numpy as np
from slim4 import slim2d, tools
from osgeo import ogr


def recursive_rule(depth=0):
    if depth == 0:
        points = np.array(
            [
                [2 / 3, 1 / 6],
                [1 / 6, 2 / 3],
                [1 / 6, 1 / 6],
            ]
        )
        weights = np.array([1 / 6, 1 / 6, 1 / 6])
        # points = np.array([[1 / 3, 1 / 3]])
        # weights = np.array([1 / 2])
        return points, weights
    else:
        points, weights = recursive_rule(depth - 1)
        points = np.vstack(
            [
                points,
                points + np.array([1, 0]),
                points + np.array([0, 1]),
                np.array([1, 1]) - points,
            ]
        )
        weights = np.hstack([weights, weights, weights, weights])
        return points / 2, weights / 4


# bathymetry
def bathy_on_mesh(
    mesh: slim2d.Mesh,
    filename: str | List[str],
    nodal=True,
    cg=True,
    rule_depth=1,
):
    if nodal:
        xn, yn = mesh.xnodes.astype(np.float64)
        bath_cg = -tools.interpolate_tif(filename, xn, yn, mesh.projection)
        bath_dg = bath_cg[mesh.element_nodes][:, :, None].T
        return bath_dg
    else:
        x, y = mesh.x.astype(np.float64)
        points, weights = recursive_rule(rule_depth)
        xsi, eta = points[:, 0], points[:, 1]
        phi = np.array([1 - xsi - eta, xsi, eta])

        xp = np.einsum("jn,jp->np", x, phi, optimize="optimal")
        yp = np.einsum("jn,jp->np", y, phi, optimize="optimal")
        bath_data = -tools.interpolate_tif(filename, xp, yp, mesh.projection)
        bath_mass = np.einsum(
            "ep,np,p->en", bath_data, phi, weights, optimize="optimal"
        )
        imass = 24 * np.eye(3) - 6
        bath_dg = np.einsum("ej,ij->ie", bath_mass, imass, optimize="optimal")[
            None, :, :
        ]

        if cg:
            bath_cg = mesh.gather_dg_to_cg("mean_area", bath_dg)
            bath_dg = bath_cg[0][mesh.element_nodes][:, :, None].T
        return bath_dg


def shapefile_to_flat_gpkg(input_shp, output_gpkg, output_layer_name, bounds):
    """
    This is mostly ChatGPT generated code
    """
    minlon, maxlon, minlat, maxlat = bounds

    src_ds = ogr.Open(input_shp)
    if src_ds is None:
        raise RuntimeError(f"Could not open {input_shp}")

    src_lyr = src_ds.GetLayer()
    src_srs = src_lyr.GetSpatialRef()

    # Build bbox polygon
    ring = ogr.Geometry(ogr.wkbLinearRing)
    ring.AddPoint(minlon, minlat)
    ring.AddPoint(maxlon, minlat)
    ring.AddPoint(maxlon, maxlat)
    ring.AddPoint(minlon, maxlat)
    ring.AddPoint(minlon, minlat)

    bbox_geom = ogr.Geometry(ogr.wkbPolygon)
    bbox_geom.AddGeometry(ring)

    # Coarse filter
    src_lyr.SetSpatialFilterRect(minlon, minlat, maxlon, maxlat)
    drv = ogr.GetDriverByName("GPKG")
    try:
        drv.DeleteDataSource(output_gpkg)
    except Exception:
        pass

    def polygon_parts(geom):
        if geom is None:
            return
        name = geom.GetGeometryName().upper()
        if name == "POLYGON":
            yield geom.Clone()
        elif name == "MULTIPOLYGON":
            for i in range(geom.GetGeometryCount()):
                part = geom.GetGeometryRef(i)
                if part is not None:
                    yield part.Clone()

    all_parts = ogr.Geometry(ogr.wkbMultiPolygon)
    count = 0
    for src_feat in src_lyr:
        geom = src_feat.GetGeometryRef()
        if geom is None:
            continue
        for part in polygon_parts(geom):
            if part.Within(bbox_geom):
                all_parts.AddGeometry(part)
                count += 1
    all_parts = all_parts.UnionCascaded()

    dst_ds = drv.CreateDataSource(output_gpkg)
    dst_lyr = dst_ds.CreateLayer(
        output_layer_name, srs=src_srs, geom_type=ogr.wkbMultiPolygon
    )
    if count > 0:
        dst_defn = dst_lyr.GetLayerDefn()
        out_feat = ogr.Feature(dst_defn)
        out_feat.SetGeometry(all_parts)

        if dst_lyr.CreateFeature(out_feat) != 0:
            raise RuntimeError("Failed to create output feature")
        out_feat = None
    src_ds = None
    dst_ds = None
    print(f"Wrote {output_gpkg} with {count} polygons in one feature")


def poly_to_line(gpkg_in, gpkg_out, physical_name):
    src = ogr.Open(gpkg_in)
    lyr = src.GetLayer()
    drv = ogr.GetDriverByName("GPKG")
    dst = drv.CreateDataSource(gpkg_out)
    out = dst.CreateLayer("lines", lyr.GetSpatialRef(), ogr.wkbMultiLineString)
    out.CreateField(ogr.FieldDefn("physical", ogr.OFTString))
    defn = out.GetLayerDefn()
    for f in lyr:
        g = f.GetGeometryRef()
        b = g.Boundary()
        if b.GetGeometryType() == ogr.wkbLineString:
            m = ogr.Geometry(ogr.wkbMultiLineString)
            m.AddGeometry(b)
            b = m
        nf = ogr.Feature(defn)
        nf.SetGeometry(b)
        nf.SetField("physical", physical_name)
        out.CreateFeature(nf)