import mikeio
import os
import geopandas as gpd
from shapely import polygons
import numpy as np

from ._utils import _element_corners

def mesh2poly(input_fname, output_fname, verbose=True):
    """
    Convert a MIKE mesh file to a GeoDataFrame of polygons.
    Parameters:
    - input_fname: Path to the MIKE mesh file (.mesh).
    - output_fname: Path to save the output GeoDataFrame (.shp).
    - verbose: If True, print progress messages.
    Returns:
    - A GeoDataFrame containing the polygons of the mesh elements.
    """
    # Check if file exists
    if not os.path.exists(input_fname):
        raise FileNotFoundError(f"Input file not found: {input_fname}")
    # Check file extensions
    if not input_fname.endswith(".mesh"):
        raise ValueError("Input file must have a .mesh extension")
    if not output_fname.endswith(".shp"):
        raise ValueError("Output file must have a .shp extension")
    # Make sure output directory exists
    output_dir = os.path.dirname(output_fname)
    if output_dir and not os.path.exists(output_dir):
        if verbose:
            print(f"WARNING: Output directory does not exist. Creating: {output_dir}")
        os.makedirs(output_dir)
    # Load mesh
    mesh = mikeio.Mesh(input_fname)
    geom = mesh.geometry
    projection = geom.projection if geom.projection != "LONG/LAT" else "EPSG:4326"
    z = mesh.element_coordinates[:, 2]
    mpoly = geom.to_shapely()
    polygons = list(mpoly.geoms)
    values = z.flatten()

    gdf = gpd.GeoDataFrame(data={"ID": range(len(polygons)), "Bathymetry": values}, geometry=polygons, crs=projection)
    gdf.to_file(output_fname)
    if verbose:
        print(f"Mesh converted to polygons and saved to: {output_fname}")


 
 
def dfs2_to_geodataframe(dfs2_path):
    """Return a GeoDataFrame of every element with x/y indices, in the most
    appropriate CRS for the file (see module docstring)."""
    g = mikeio.open(dfs2_path).geometry
    nx, ny = g.nx, g.ny
    X, Y, crs = _element_corners(g)
 
    # Corner arrays for every cell, each shape (ny, nx). Clockwise ring order:
    # bottom-left, top-left, top-right, bottom-right (+ closing point).
    def corner(js, is_):
        return np.stack([X[js, is_], Y[js, is_]], axis=-1)  # (ny, nx, 2)
 
    bl = corner(slice(0, ny),     slice(0, nx))
    tl = corner(slice(1, ny + 1), slice(0, nx))
    tr = corner(slice(1, ny + 1), slice(1, nx + 1))
    br = corner(slice(0, ny),     slice(1, nx + 1))
 
    rings = np.stack([bl, tl, tr, br, bl], axis=2).reshape(nx * ny, 5, 2)
    geom = polygons(rings)
 
    jj, ii = np.meshgrid(np.arange(ny), np.arange(nx), indexing="ij")
    return gpd.GeoDataFrame(
        {"x_index": ii.ravel(), "y_index": jj.ravel(), "geometry": geom},
        crs=crs,
    )
 
 
def dfs2_to_shapefile(input_fname, output_fname, verbose=True):
    """Write every element of `input_fname` as a polygon to `output_fname`."""
    gdf = dfs2_to_geodataframe(input_fname)
    gdf.to_file(output_fname)  # geopandas writes the .prj from the GeoDataFrame CRS
    if verbose:
        epsg = gdf.crs.to_epsg()
        crs_desc = f"EPSG:{epsg}" if epsg else gdf.crs.name
        print(f"Wrote {len(gdf)} elements to {output_fname}  (CRS: {crs_desc})")