import os
import numpy as np
import rasterio
from rasterio.transform import from_origin
from rasterio.warp import reproject, Resampling

_RESAMP = {
    "nearest": Resampling.nearest,
    "bilinear": Resampling.bilinear,
    "cubic": Resampling.cubic,
    "cubic_spline": Resampling.cubic_spline,
    "average": Resampling.average,
}


def read_funwave_depth(path, mglob, nglob):
    """Read a FUNWAVE depth file into a (nglob, mglob) array.

    Returned array is in FILE order: row 0 == J=1 == southern row,
    column 0 == I=1 == western column.
    """
    with open(path) as f:
        tokens = f.read().split()          # splits on any whitespace incl. newlines

    expected = mglob * nglob
    if len(tokens) != expected:
        raise ValueError(
            f"Token count mismatch: file has {len(tokens)} values but "
            f"Mglob*Nglob = {mglob}*{nglob} = {expected}. "
            "Check that Mglob/Nglob are correct and the file is not truncated."
        )

    flat = np.asarray(tokens, dtype=np.float64)
    # Outer loop is J (slow), inner loop is I (fast) -> C-order reshape gives
    # arr[j, i] = Depth(I=i+1, J=j+1).
    return flat.reshape(nglob, mglob)

def generate_grid(
    xmin, xmax, ymin, ymax,
    dphi, dtheta,
    raster_path,
    out_txt_path,
    parent=None,                 # None, or dict(lon_west=, lat_south=, dphi=, dtheta=)
    flip_sign=False,             # True: negate raster (elevation +up -> FUNWAVE depth +down)
    nodata_fill=None,            # value used where the raster does not cover the grid
    resampling="bilinear",       # nearest | bilinear | cubic | cubic_spline | average
    fmt="%.6f",                  # numeric format for the .txt
    export_shapefile=False,      # also write <out>.shp of cell polygons (ix, jy)
    export_clipped_raster=False, # also write <out>.tif of the interpolated grid
    grid_crs="EPSG:4326",
    verbose=True,
):
    
    """
    Generate a FUNWAVE-TVD computational grid (depth file) from a raster.

    Parameters
    ----------
    xmin, xmax, ymin, ymax : float
        Bounding box in grid CRS (lon/lat for EPSG:4326)
    dphi, dtheta : float
        Grid spacing in degrees (or CRS units)
    raster_path : str
        Path to input raster with depth/elevation data
    out_txt_path : str
        Output path for FUNWAVE depth .txt file
    parent : dict, optional
        Parent grid info for nesting: dict(lon_west=, lat_south=, dphi=, dtheta=)
    flip_sign : bool, optional
        If True, negate raster values (elevation +up -> FUNWAVE depth +down)
    nodata_fill : float, optional
        Fill value where raster does not cover grid
    resampling : str, optional
        Resampling method: "nearest" | "bilinear" | "cubic" | "cubic_spline" | "average"
    fmt : str, optional
        Numeric format for .txt output (default: "%.6f")
    export_shapefile : bool, optional
        Write cell polygon shapefile with (ix, jy) indices
    export_clipped_raster : bool, optional
        Write interpolated grid as GeoTIFF
    grid_crs : str, optional
        Coordinate reference system (default: "EPSG:4326")
    verbose : bool, optional
        Print progress messages

    Returns
    -------
    dict
        Grid metadata: Mglob, Nglob, dphi, dtheta, xmin, xmax, ymin, ymax,
        txt (output path), shapefile, raster, n_nodata, flip_sign

    Example
    -------
    >>> parent = dict(lon_west=-129, lat_south=46, dphi=30/3600.0, dtheta=30/3600.0)
    >>> result = generate_grid(
    ...     xmin=-127.001, xmax=-126, ymin=48.1, ymax=48.5,
    ...     dphi=30/3600.0, dtheta=30/3600.0,
    ...     raster_path="dem.tif",
    ...     out_txt_path="depth.txt",
    ...     parent=parent,
    ...     flip_sign=False,
    ...     export_clipped_raster=True,
    ...     export_shapefile=True,
    ... )
    """
    def log(msg):
        if verbose:
            print(msg)

    xmin, xmax, ymin, ymax = map(float, (xmin, xmax, ymin, ymax))
    dphi, dtheta = float(dphi), float(dtheta)
    if resampling not in _RESAMP:
        raise ValueError(f"resampling must be one of {list(_RESAMP)}")

    # ---------------------------------------------------------------
    # 1. Nested-grid alignment
    # ---------------------------------------------------------------
    if parent is not None:
        px0 = float(parent["lon_west"]); py0 = float(parent["lat_south"])
        pdx = float(parent["dphi"]);     pdy = float(parent["dtheta"])

        sx, sy = pdx / dphi, pdy / dtheta
        if abs(round(sx) - sx) > 1e-6 or abs(round(sy) - sy) > 1e-6:
            log(f"WARNING: refinement ratio is not an integer "
                f"(parent/child = {sx:.5f} in x, {sy:.5f} in y). Child grid points "
                f"will NOT coincide with parent grid points -- nested-grid "
                f"alignment cannot be guaranteed. Choose dphi/dtheta that evenly "
                f"divide the parent cell size.")
        else:
            log(f"Refinement ratio: {int(round(sx))} (x) x {int(round(sy))} (y) "
                f"relative to parent  [OK, integer].")

        def snap(v, origin, step):
            return origin + round((v - origin) / step) * step

        nx0, nx1 = snap(xmin, px0, pdx), snap(xmax, px0, pdx)
        ny0, ny1 = snap(ymin, py0, pdy), snap(ymax, py0, pdy)

        tol = 1e-7
        adj = []
        if abs(nx0 - xmin) > tol: adj.append(f"xmin {xmin:.8f} -> {nx0:.8f}")
        if abs(nx1 - xmax) > tol: adj.append(f"xmax {xmax:.8f} -> {nx1:.8f}")
        if abs(ny0 - ymin) > tol: adj.append(f"ymin {ymin:.8f} -> {ny0:.8f}")
        if abs(ny1 - ymax) > tol: adj.append(f"ymax {ymax:.8f} -> {ny1:.8f}")
        if adj:
            log("WARNING: bounding box adjusted to align with the parent grid:\n  "
                + "\n  ".join(adj))
        else:
            log("Bounding box already aligned with the parent grid (no adjustment).")
        xmin, xmax, ymin, ymax = nx0, nx1, ny0, ny1
    else:
        log("No parent grid supplied: standalone grid, no alignment snapping applied.")

    # ---------------------------------------------------------------
    # 2. Grid dimensions (node counts)
    # ---------------------------------------------------------------
    Mglob = int(round((xmax - xmin) / dphi)) + 1
    Nglob = int(round((ymax - ymin) / dtheta)) + 1
    # snap extent to the exact node count
    xmax = xmin + (Mglob - 1) * dphi
    ymax = ymin + (Nglob - 1) * dtheta
    log(f"Grid: Mglob={Mglob} (x) x Nglob={Nglob} (y) = {Mglob * Nglob:,} cells")
    log(f"Extent: lon [{xmin:.8f}, {xmax:.8f}]  lat [{ymin:.8f}, {ymax:.8f}]")

    # ---------------------------------------------------------------
    # 3. Interpolate the raster onto the grid nodes
    # ---------------------------------------------------------------
    dst = np.full((Nglob, Mglob), np.nan, dtype="float64")
    # north-up transform with node positions at pixel centres
    dst_transform = from_origin(xmin - dphi / 2.0, ymax + dtheta / 2.0, dphi, dtheta)
    with rasterio.open(raster_path) as src:
        src_crs = src.crs or grid_crs
        reproject(
            source=rasterio.band(src, 1),
            destination=dst,
            src_transform=src.transform, src_crs=src_crs,
            dst_transform=dst_transform, dst_crs=grid_crs,
            src_nodata=src.nodata, dst_nodata=np.nan,
            resampling=_RESAMP[resampling],
        )

    nan_mask = ~np.isfinite(dst)
    n_nan = int(nan_mask.sum())
    if n_nan:
        if nodata_fill is None:
            raise ValueError(
                f"{n_nan:,} of {dst.size:,} grid nodes fall outside the raster "
                f"coverage. Supply a raster covering the whole domain, or pass "
                f"nodata_fill=<value>.")
        log(f"WARNING: {n_nan:,} nodes are outside raster coverage; "
            f"filled with nodata_fill={nodata_fill}.")
        dst[nan_mask] = nodata_fill

    if flip_sign:
        dst = -dst  # elevation (+up) -> FUNWAVE depth (+down)

    # ---------------------------------------------------------------
    # 4. Write the FUNWAVE depth .txt  (south row first, west->east)
    # ---------------------------------------------------------------
    out_dir = os.path.dirname(os.path.abspath(out_txt_path))
    os.makedirs(out_dir, exist_ok=True)
    np.savetxt(out_txt_path, np.flipud(dst), fmt=fmt)
    log(f"Wrote depth file: {out_txt_path}")

    base = os.path.splitext(out_txt_path)[0]
    out = dict(Mglob=Mglob, Nglob=Nglob, dphi=dphi, dtheta=dtheta,
               xmin=xmin, xmax=xmax, ymin=ymin, ymax=ymax,
               txt=out_txt_path, shapefile=None, raster=None,
               n_nodata=n_nan, flip_sign=flip_sign)

    # ---------------------------------------------------------------
    # 5. Optional GeoTIFF of the interpolated grid
    # ---------------------------------------------------------------
    if export_clipped_raster:
        tif = base + ".tif"
        prof = dict(driver="GTiff", height=Nglob, width=Mglob, count=1,
                    dtype="float32", crs=grid_crs, transform=dst_transform,
                    compress="deflate")
        if nodata_fill is not None:
            prof["nodata"] = float(nodata_fill)
        with rasterio.open(tif, "w", **prof) as d:
            d.write(dst.astype("float32"), 1)
            d.update_tags(AREA_OR_POINT="Area")
        out["raster"] = tif
        log(f"Wrote raster:    {tif}")

    # ---------------------------------------------------------------
    # 6. Optional cell-polygon shapefile (ix, jy)
    # ---------------------------------------------------------------
    if export_shapefile:
        import geopandas as gpd
        import shapely
        shp = base + ".shp"
        ncell = Mglob * Nglob
        if ncell > 2_000_000:
            log(f"WARNING: exporting {ncell:,} cell polygons to a shapefile; this "
                f"is slow and may approach/exceed the 2 GB shapefile limit. "
                f"Consider GeoPackage for very large grids.")
        for ext in (".shp", ".shx", ".dbf", ".prj", ".cpg"):
            p = base + ext
            if os.path.exists(p):
                os.remove(p)
        ix = np.arange(1, Mglob + 1)
        lon_c = xmin + (ix - 1) * dphi
        X0, X1 = lon_c - dphi / 2.0, lon_c + dphi / 2.0
        rows_per = max(1, 500_000 // Mglob)
        first = True
        for j0 in range(1, Nglob + 1, rows_per):
            j1 = min(Nglob, j0 + rows_per - 1)
            jr = np.arange(j0, j1 + 1)
            lat_c = ymin + (jr - 1) * dtheta
            Y0, Y1 = lat_c - dtheta / 2.0, lat_c + dtheta / 2.0
            nb = len(jr)
            geoms = shapely.box(np.tile(X0, nb), np.repeat(Y0, Mglob),
                                np.tile(X1, nb), np.repeat(Y1, Mglob))
            gdf = gpd.GeoDataFrame(
                {"ix": np.tile(ix, nb), "jy": np.repeat(jr, Mglob)},
                geometry=geoms, crs=grid_crs)
            gdf.to_file(shp, append=not first)
            first = False
        out["shapefile"] = shp
        log(f"Wrote shapefile: {shp}")

    return out


def txt2tif(depth_path, out_path, mglob, nglob, dphi, dtheta,
                     lon_west=0.0, lat_south=0.0, crs="EPSG:4326",
                     dtype="float32", nodata=None, flip_sign=False,
                     area_or_point="Area"):
    """Convert a FUNWAVE depth file to a GeoTIFF."""
    arr = read_funwave_depth(depth_path, mglob, nglob)   # (nglob, mglob), south-up

    if flip_sign:
        arr = -arr   # FUNWAVE depth (positive down) -> elevation (positive up)

    # FUNWAVE row 0 is the SOUTH edge; GeoTIFF row 0 is the NORTH edge -> flip.
    arr_north_up = np.flipud(arr).astype(dtype)

    # Node (I=1, J=1) at (lon_west, lat_south) is the centre of the SW pixel.
    # Geotransform origin = top-left CORNER of the top-left (NW) pixel.
    west_corner = lon_west - dphi / 2.0
    north_corner = lat_south + (nglob - 0.5) * dtheta
    transform = from_origin(west_corner, north_corner, dphi, dtheta)

    profile = {
        "driver": "GTiff",
        "height": nglob,
        "width": mglob,
        "count": 1,
        "dtype": dtype,
        "crs": crs,
        "transform": transform,
        "compress": "deflate",   # lossless; harmless to drop if undesired
    }
    if nodata is not None:
        profile["nodata"] = nodata

    with rasterio.open(out_path, "w", **profile) as dst:
        dst.write(arr_north_up, 1)
        dst.update_tags(AREA_OR_POINT=area_or_point)

    return arr_north_up, transform