import mikeio
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.colors import Normalize, LogNorm
from matplotlib import tri
from mpl_toolkits.axes_grid1 import make_axes_locatable
from scipy.sparse import csr_matrix

from ..plotting import subplots
from ._utils import *
from ._utils import _resolve_time_index, _proj2geo, _add_basemap, _add_matched_colorbar

def _contourf_geometryfm2d(geometry, values, vmin=None, vmax=None, cbar_n_levels=5, scale="linear", cmap="jet", ax=None, add_colorbar=True):
    if ax is None:
        fig, ax = subplots()
    nc = geometry.node_coordinates
    ec = geometry.element_coordinates
    et = geometry.element_table
    projection = geometry.projection

    set_aspect_ratio(ax, nc, projection)
    assert len(values) == ec.shape[0], "Values length must match number of elements"
    norm, levels, cbar_levels, extend = get_cmap_properties(values, vmin=vmin, vmax=vmax, cbar_n_levels=cbar_n_levels, scale=scale)
    triang, zn = get_tris(nc, ec, et, values)
    fig_obj = ax.tricontourf(triang, zn, cmap=cmap, norm=norm, levels=levels, extend=extend)
    if add_colorbar:
        cax = make_axes_locatable(ax).append_axes("right", size="5%", pad=0.05)
        colorbar = plt.colorbar(
            fig_obj,
            cax=cax,
            ticks=cbar_levels,
            boundaries=cbar_levels,
            extend=extend,
        )
        colorbar.set_ticks(cbar_levels)

        return ax, colorbar
    return ax
        

def contourf(
        geometry,
        values,
        vmin=None,
        vmax=None,
        cbar_n_levels=5,
        cmap="jet",
        ax=None,
        scale="linear",
        add_colorbar=True,
):
    if isinstance(geometry, mikeio.spatial.GeometryFM2D):
        return _contourf_geometryfm2d(geometry, values, vmin=vmin, vmax=vmax, cbar_n_levels=cbar_n_levels, cmap=cmap, ax=ax, scale=scale, add_colorbar=add_colorbar)
    else:
        raise NotImplementedError(f"contourf not implemented for geometry type {type(geometry)}")
    

def plot_wind_vectors(
    dfs2_path,
    u_item=0,
    v_item=1,
    timestep=0,
    skip=5,
    scale=None,
    outfile=None,
    bbox=None,
    add_basemap=False,
    basemap_source="Esri.WorldImagery",
):
    """
    Parameters
    ----------
    dfs2_path : str
        Path to the dfs2 file.
    u_item, v_item : int or str
        Item index or name for the U and V wind components (true east/north).
    timestep : int or str/datetime
        Either an integer time index, or an actual timestamp (e.g.
        "2018-01-01 01:00"). A timestamp outside the file range, or one that
        falls between available steps, raises ValueError describing the range,
        the timestep interval, and the nearest available steps.
    skip : int
        Plot every `skip`-th vector in each direction (thins a dense grid).
    scale : float or None
        matplotlib quiver scale (None = autoscale; larger = shorter arrows).
    outfile : str or None
        If given, save the figure here instead of showing it.
    bbox : (lon_min, lat_min, lon_max, lat_max) or None
        Restrict the plot to this lon/lat box. None plots the whole domain.
        Raises ValueError if the box does not overlap the domain.
    add_basemap : bool
        If True, overlay a web-tile basemap under the vectors (needs contextily
        and an internet connection).
    basemap_source : str
        Dotted contextily provider name, e.g. "Esri.WorldImagery" (satellite,
        default), "OpenStreetMap.Mapnik", or "CartoDB.Positron".
    """
    # Resolve `timestep` (integer index OR timestamp) against the file's time axis.
    tidx = _resolve_time_index(mikeio.open(dfs2_path).time, timestep)
    ds = mikeio.read(dfs2_path, items=[u_item, v_item], time=tidx)
    u = np.squeeze(ds[0].to_numpy())  # (ny, nx)
    v = np.squeeze(ds[1].to_numpy())
 
    g = ds.geometry
    x0, y0 = getattr(g, "origin", (0.0, 0.0))
 
    # g.x/g.y may be stored absolute (geographic grids, where g.x[0] == origin)
    # or local/0-based (projected grids). Normalize to 0-based so that
    # "origin + local" gives the right coordinate in both cases.
    gx = np.asarray(g.x, dtype=float) - float(g.x[0])
    gy = np.asarray(g.y, dtype=float) - float(g.y[0])
 
    # Thin first, so we only project the nodes we actually draw.
    sl = (slice(None, None, skip), slice(None, None, skip))
    Xl, Yl = np.meshgrid(gx, gy)
    Xl, Yl, u, v = Xl[sl], Yl[sl], u[sl], v[sl]
 
    # Native projection coords are axis-aligned with the grid (origin + local);
    # convert them to true WGS84 lon/lat via the file's projection.
    mp = MapProjection(g.projection_string)
    lon, lat = _proj2geo(mp, x0 + Xl, y0 + Yl)
    lon = (lon + 180.0) % 360.0 - 180.0   # wrap 0..360 longitudes to [-180, 180]
    mag = np.hypot(u, v)
 
    # Optional lon/lat bounding box.
    view_bbox = None
    if bbox is not None:
        xmin, ymin, xmax, ymax = bbox
        if xmin > xmax or ymin > ymax:
            raise ValueError("bbox must be (lon_min, lat_min, lon_max, lat_max) with min <= max")
        dxmin, dxmax = float(lon.min()), float(lon.max())
        dymin, dymax = float(lat.min()), float(lat.max())
        if not (xmin <= dxmax and xmax >= dxmin and ymin <= dymax and ymax >= dymin):
            raise ValueError(
                f"bbox {bbox} does not overlap the domain "
                f"(lon: {dxmin:g}..{dxmax:g}, lat: {dymin:g}..{dymax:g})."
            )
        keep = (lon >= xmin) & (lon <= xmax) & (lat >= ymin) & (lat <= ymax)
        lon, lat, u, v, mag = lon[keep], lat[keep], u[keep], v[keep], mag[keep]
        view_bbox = (max(xmin, dxmin), min(xmax, dxmax),
                     max(ymin, dymin), min(ymax, dymax))
 
    # Plot.
    fig, ax = plt.subplots(figsize=(10, 8))
    # A thin outline keeps arrows legible over dark satellite imagery.
    edge = dict(linewidth=0.4, edgecolor="white") if add_basemap else {}
    q = ax.quiver(lon, lat, u, v, mag, cmap="viridis", scale=scale,
                  scale_units="xy" if scale else None, width=0.0025,
                  pivot="middle", **edge)
 
    ref = np.nanpercentile(mag, 90)
    if np.isfinite(ref) and ref > 0:
        ax.quiverkey(q, 0.9, 1.02, ref, f"{ref:.1f} m/s",
                     labelpos="E", coordinates="axes")
 
    t = ds.time[0]
    ax.set_title(f"Wind vectors — {t:%Y-%m-%d %H:%M}")
    ax.set_xlabel("Longitude")
    ax.set_ylabel("Latitude")
 
    # Correct longitude compression for the map's mean latitude.
    mlat = float(np.mean(lat)) if lat.size else 0.0
    ax.set_aspect(1.0 / max(np.cos(np.deg2rad(mlat)), 1e-6))
 
    if view_bbox is not None:
        ax.set_xlim(view_bbox[0], view_bbox[1])
        ax.set_ylim(view_bbox[2], view_bbox[3])
    else:
        # Explicit limits from the data so the basemap fetches the right extent.
        ax.set_xlim(float(lon.min()), float(lon.max()))
        ax.set_ylim(float(lat.min()), float(lat.max()))
 
    if add_basemap:
        _add_basemap(ax, basemap_source)
 
    # Reserve margins for labels + colorbar (tight_layout fights the fixed
    # aspect and the manually placed colorbar, so we set margins directly).
    fig.subplots_adjust(left=0.09, right=0.86, top=0.93, bottom=0.10)
    cbar, sync_cbar = _add_matched_colorbar(fig, ax, q, "Wind speed [m/s]")
 
    if outfile:
        fig.canvas.draw()   # apply aspect so the axis box is final
        sync_cbar()         # match the colorbar to it before saving
        fig.savefig(outfile, dpi=150, bbox_inches="tight")
        print(f"Saved {outfile}")
    else:
        plt.show()
    return fig, ax


if __name__ == "__main__":
    fname = r"C:\Users\abja\Downloads\Test.dfs2"
    u_item = 0
    v_item = 1
    plot_wind_vectors(fname, u_item, v_item, timestep=0, skip=1, scale=None, outfile=None, bbox=None, add_basemap=False)
