import mikeio
from mikecore.Projections import MapProjection
import numpy as np
import pandas as pd
from matplotlib.colors import Normalize, LogNorm
from matplotlib import tri
from scipy.sparse import csr_matrix
import functools
from pyproj import CRS, Transformer
import os
import re
from dataclasses import dataclass, field
from pathlib import Path, PureWindowsPath


from ..plotting import subplots

_WGS84 = CRS.from_epsg(4326)
# A pipe-delimited clstring on a single line. Windows paths cannot contain '|'
# or a newline, so this brackets each link exactly. Empty links are ``||``.
_PIPE = re.compile(r"\|([^|\r\n]*)\|")
_UTF8_BOM = b"\xef\xbb\xbf"


def set_aspect_ratio(ax, nc, projection):
    is_geo = projection == "LONG/LAT"
    if is_geo:
        mean_lat = np.mean(nc[:, 1])
        ax.set_aspect(1.0 / np.cos(np.pi * mean_lat / 180))
    else:
        ax.set_aspect("equal")

def get_cmap_properties(values, vmin=None, vmax=None, cbar_n_levels=5, scale="linear"):
    if vmin is None:
        vmin = np.nanmin(values)
    if vmax is None:
        vmax = np.nanmax(values)
    if scale == "linear":
        norm = Normalize(vmin=vmin, vmax=vmax)
        levels = np.linspace(vmin, vmax, 100)
        cbar_levels = np.linspace(vmin, vmax, cbar_n_levels)
    elif scale == "log":
        norm = LogNorm(vmin=max(vmin, 1e-10), vmax=vmax)  # avoid log(0)
        levels = np.logspace(np.log10(max(vmin, 1e-10)), np.log10(vmax), 100)
        cbar_levels = np.logspace(np.log10(max(vmin, 1e-10)), np.log10(vmax), cbar_n_levels)

    data_min = np.nanmin(values)
    data_max = np.nanmax(values)
    extend_min = data_min < vmin
    extend_max = data_max > vmax
    if extend_min and extend_max:
        extend = "both"
    elif extend_min:
        extend = "min"
    elif extend_max:
        extend = "max"
    else:
        extend = "neither"

    return norm, levels, cbar_levels, extend

def is_tri_only(element_table):
    return max([len(el) for el in element_table]) == 3

def mesh2tri(nc, ec, et, z):
    if is_tri_only(et):
        return ec, np.stack(et), z
    
    elem_table = [list(et[i]) for i in range(len(et))]
    tmp_elmnt_nodes = elem_table.copy()
    for el, item in enumerate(tmp_elmnt_nodes):
        if len(item) == 4:
            elem_table.pop(el)  # remove quad element

            # insert two new tri elements in table
            elem_table.insert(el, item[:3])
            tri2_nodes = [item[i] for i in [2, 3, 0]]
            elem_table.append(tri2_nodes)

            # new center coordinates for new tri-elements
            ec[el] = nc[item[:3]].mean(axis=1)
            tri2_ec = nc[tri2_nodes].mean(axis=1)
            ec = np.append(ec, tri2_ec.reshape(1, -1), axis=0)

            # use same data in two new tri elements
            data = np.append(data, data[el])

    return ec, np.asarray(elem_table), z

def get_connectivity_matrix(et, n_nodes):
    row_ind = et.ravel()
    col_ind = np.repeat(np.arange(et.shape[0]), et.shape[1])
    data = np.ones(len(row_ind), dtype=int)
    connectivity_matrix = csr_matrix(
        (data, (row_ind, col_ind)), shape=(n_nodes, et.shape[0])
    )
    return connectivity_matrix

def get_node_data(nc, ec, et, z):
    ec, et, z = mesh2tri(nc, ec, et, z)
    connectivity_matrix = get_connectivity_matrix(et, nc.shape[0])

    node_centered_data = np.zeros(shape=nc.shape[0])
    for n in range(connectivity_matrix.shape[0]):
        item = connectivity_matrix.getrow(n).indices
        I = ec[item][:, :2] - nc[n][:2]
        I2 = (I**2).sum(axis=0)
        Ixy = (I[:, 0] * I[:, 1]).sum(axis=0)
        lamb = I2[0] * I2[1] - Ixy**2
        omega = np.zeros(1)
        if lamb > 1e-10 * (I2[0] * I2[1]):
            # Standard case - Pseudo
            lambda_x = (Ixy * I[:, 1] - I2[1] * I[:, 0]) / lamb
            lambda_y = (Ixy * I[:, 0] - I2[0] * I[:, 1]) / lamb
            omega = 1.0 + lambda_x * I[:, 0] + lambda_y * I[:, 1]
            omega[np.where(omega > 2)] = 2
            omega[np.where(omega < 0)] = 0
        if omega.sum() > 0:
            node_centered_data[n] = np.sum(omega * z[item]) / np.sum(omega)
        else:
            # We did not succeed using pseudo laplace procedure, use inverse distance instead
            InvDis = [
                1 / np.hypot(case[0], case[1]) for case in ec[item][:, :2] - nc[n][:2]
            ]
            node_centered_data[n] = np.sum(InvDis * z[item]) / np.sum(InvDis)

    return node_centered_data

def get_tris(nc, ec, et, z):
    ec, et, z = mesh2tri(nc, ec, et, z)
    triang = tri.Triangulation(nc[:, 0], nc[:, 1], et)
    zn = get_node_data(nc, ec, et, z)
    return triang, zn


def _proj2geo(mp, rx, ry):
    """Vectorized wrapper around MapProjection.Proj2Geo (scalar-only in mikecore)."""
    rx = np.asarray(rx, dtype=float)
    ry = np.asarray(ry, dtype=float)
    lon = np.empty(rx.shape, dtype=float)
    lat = np.empty(ry.shape, dtype=float)
    fx, fy, olon, olat = rx.ravel(), ry.ravel(), lon.ravel(), lat.ravel()
    for k in range(fx.size):
        olon[k], olat[k] = mp.Proj2Geo(float(fx[k]), float(fy[k]))
    return lon, lat


def _fmt_timedelta(td):
    """Human-friendly timestep label, e.g. 'hourly (1 hour)' or '30-minute (30 minutes)'."""
    s = int(td.total_seconds())
    days, rem = divmod(s, 86400)
    hours, rem = divmod(rem, 3600)
    mins, secs = divmod(rem, 60)
    parts = []
    for val, name in ((days, "day"), (hours, "hour"), (mins, "minute"), (secs, "second")):
        if val:
            parts.append(f"{val} {name}{'s' if val != 1 else ''}")
    detail = ", ".join(parts) if parts else "0 seconds"
    nice = {86400: "daily", 3600: "hourly", 1800: "30-minute", 900: "15-minute",
            600: "10-minute", 300: "5-minute", 60: "1-minute"}.get(s)
    return f"{nice} ({detail})" if nice else detail


def _describe_step(times):
    """Describe the timestep interval of a time index."""
    if len(times) < 2:
        return "single timestep (no interval)"
    diffs = np.diff(times.values).astype("timedelta64[s]").astype(np.int64)
    if np.all(diffs == diffs[0]):
        return _fmt_timedelta(pd.Timedelta(seconds=int(diffs[0])))
    return "irregular (variable timestep)"


def _resolve_time_index(times, t):
    """Resolve `t` (int index OR timestamp) to a valid time index, or raise.

    Raises ValueError with the available range + timestep interval if `t` is
    out of range, and with the surrounding timesteps if a timestamp falls
    between available steps.
    """
    interval = _describe_step(times)
    n = len(times)

    if isinstance(t, (int, np.integer)) and not isinstance(t, bool):
        idx = t + n if t < 0 else t
        if not (0 <= idx < n):
            raise ValueError(
                f"timestep index {t} is out of range (valid 0..{n - 1}). "
                f"File covers {times[0]} to {times[-1]}; timestep: {interval}."
            )
        return int(idx)

    try:
        ts = pd.Timestamp(t)
    except (ValueError, TypeError) as e:
        raise ValueError(f"Could not read '{t}' as a timestep index or timestamp.") from e

    if ts < times[0] or ts > times[-1]:
        raise ValueError(
            f"Time {ts} is outside the file's range "
            f"{times[0]} to {times[-1]}; timestep: {interval}."
        )
    exact = np.nonzero(times == ts)[0]
    if exact.size:
        return int(exact[0])

    pos = int(times.searchsorted(ts))  # ts lies strictly between pos-1 and pos
    raise ValueError(
        f"Time {ts} is within the file's range but is not an available timestep. "
        f"Nearest available: previous {times[pos - 1]}, next {times[pos]}; "
        f"timestep: {interval}."
    )


def _add_basemap(ax, source="Esri.WorldImagery"):
    """Overlay a web-tile basemap, keeping the axes in lon/lat (EPSG:4326).

    `source` is a dotted contextily provider name, e.g. "Esri.WorldImagery"
    (satellite), "OpenStreetMap.Mapnik", or "CartoDB.Positron".
    """
    try:
        import contextily as ctx
    except ImportError as e:
        raise ImportError(
            "add_basemap=True needs contextily: pip install contextily"
        ) from e
    provider = functools.reduce(getattr, source.split("."), ctx.providers)
    # crs="EPSG:4326" tells contextily to warp the tiles to our lon/lat axes.
    ctx.add_basemap(ax, crs="EPSG:4326", source=provider, attribution_size=6)


def _add_matched_colorbar(fig, ax, mappable, label, width=0.02, pad=0.02):
    """Add a colorbar whose height always tracks the axis's actual drawn box.
 
    With a fixed aspect ratio the axis is letterboxed inside its slot, so the
    default colorbar is too tall. This reads the axis's post-aspect position on
    every draw/resize and matches the colorbar to it. Returns (cbar, sync_fn);
    call sync_fn() after a manual fig.canvas.draw() before saving.
    """
    cax = fig.add_axes([0, 0, 1e-6, 1e-6])
    cbar = fig.colorbar(mappable, cax=cax)
    cbar.set_label(label)
 
    def _sync(_event=None):
        p = ax.get_position()  # actual drawn box (aspect already applied)
        target = (p.x1 + pad, p.y0, width, p.height)
        c = cax.get_position()
        if (abs(c.x0 - target[0]) > 1e-4 or abs(c.y0 - target[1]) > 1e-4 or
                abs(c.width - target[2]) > 1e-4 or abs(c.height - target[3]) > 1e-4):
            cax.set_position(target)
            fig.canvas.draw_idle()
 
    fig.canvas.mpl_connect("draw_event", _sync)
    fig.canvas.mpl_connect("resize_event", _sync)
    return cbar, _sync


def _usable_native_crs(projection_string):
    """Return a pyproj CRS if the projection is a GIS-usable *projected* CRS,
    otherwise None (short names, rotated-pole, or anything PROJ can't transform)."""
    if projection_string == "LONG/LAT":
        return None
    try:
        crs = CRS.from_wkt(projection_string)
        # A parseable WKT isn't enough; PROJ must be able to transform it.
        Transformer.from_crs(crs, _WGS84, always_xy=True)
        return crs
    except Exception:
        return None
 
 
def _element_corners(g):
    """Corner lattices (ny+1, nx+1) plus the output CRS.
 
    Returns (X, Y, crs): native projected coords when a usable CRS exists,
    otherwise true WGS84 lon/lat.
    """
    x0, y0 = g.origin
    dx, dy = g.dx, g.dy
    nx, ny = g.nx, g.ny
    # Cell centres are at origin + i*dx (origin == first node in every MIKE
    # convention we've seen); corners lie half a cell out on each side.
    xe = x0 + (np.arange(nx + 1) - 0.5) * dx
    ye = y0 + (np.arange(ny + 1) - 0.5) * dy
    XE, YE = np.meshgrid(xe, ye)  # (ny+1, nx+1)
 
    if g.is_geo:
        # Native coordinates are already true lon/lat (may be in 0..360).
        lon = (XE + 180.0) % 360.0 - 180.0
        return lon, YE, _WGS84
 
    native = _usable_native_crs(g.projection_string)
    if native is not None:
        return XE, YE, native  # keep native projected coordinates
 
    # Projected but not GIS-usable: convert to lon/lat with MIKE's engine.
    mp = MapProjection(g.projection_string)
    lon = np.empty(XE.shape, dtype=float)
    lat = np.empty(XE.shape, dtype=float)
    fx, fy, olon, olat = XE.ravel(), YE.ravel(), lon.ravel(), lat.ravel()
    for k in range(fx.size):
        olon[k], olat[k] = mp.Proj2Geo(float(fx[k]), float(fy[k]))
    lon = (lon + 180.0) % 360.0 - 180.0
    return lon, lat, _WGS84


# --------------------------------------------------------------------------- #
# Reporting
# --------------------------------------------------------------------------- #
@dataclass
class PackageReport:
    """Summary of a packaging run."""
 
    zip_path: str | None = None
    setup_name: str = ""
    folder_name: str = ""
    # source (resolved) path -> destination filename inside the folder
    copied: dict[str, str] = field(default_factory=dict)
    # links found but not on disk: stored_path -> resolved_path
    missing: dict[str, str] = field(default_factory=dict)
    # destination filename -> [source paths] where a rename was needed
    renamed_for_collision: dict[str, str] = field(default_factory=dict)
    # count of link tokens rewritten in the setup text
    links_rewritten: int = 0
    # how many distinct source files were referenced more than once
    duplicate_links_collapsed: int = 0
 
    def summary(self) -> str:
        lines = [
            f"Setup      : {self.setup_name}",
            f"Zip        : {self.zip_path}",
            f"Folder     : {self.folder_name}/",
            f"Files packaged   : {len(self.copied)}",
            f"Links rewritten  : {self.links_rewritten}",
        ]
        if self.duplicate_links_collapsed:
            lines.append(
                f"Duplicate links collapsed : {self.duplicate_links_collapsed}"
            )
        if self.renamed_for_collision:
            lines.append("Renamed to avoid basename clashes:")
            for dest, src in self.renamed_for_collision.items():
                lines.append(f"    {os.path.basename(src)}  ->  {dest}   (from {src})")
        if self.missing:
            lines.append("Missing (not on disk, left untouched):")
            for stored, resolved in self.missing.items():
                lines.append(f"    {stored}   (looked in: {resolved})")
        return "\n".join(lines)
 

# --------------------------------------------------------------------------- #
# Reading helpers
# --------------------------------------------------------------------------- #
 
 
def _read_text_preserving(path: str) -> tuple[str, str, bool]:
    """Read a PFS file without altering its bytes outside of what we edit.
 
    Returns ``(text, encoding, had_bom)``. We read raw bytes so that line
    endings (CR/LF) and any leading UTF-8 BOM are preserved exactly: the BOM is
    tracked separately rather than folded into the codec, because ``utf-8-sig``
    would strip a BOM on read yet *add* one on write, corrupting a BOM-less file.
    """
    raw = open(path, "rb").read()
    if raw.startswith(_UTF8_BOM):
        return raw[len(_UTF8_BOM):].decode("utf-8"), "utf-8", True
    for enc in ("utf-8", "cp1252", "latin-1"):
        try:
            return raw.decode(enc), enc, False
        except UnicodeDecodeError:
            continue
    return raw.decode("latin-1"), "latin-1", False  # latin-1 never fails
 
 
def _write_text_preserving(path: str, text: str, encoding: str, had_bom: bool) -> None:
    """Write text back byte-faithfully (no newline translation, BOM restored)."""
    data = text.encode(encoding)
    if had_bom:
        data = _UTF8_BOM + data
    with open(path, "wb") as fh:
        fh.write(data)
 
 
def _resolve(stored_path: str, setup_dir: str) -> str:
    """Resolve a stored (Windows-style) link against the setup's directory.
 
    Relative links (the normal case, ``..\\..\\...``) resolve against
    ``setup_dir``. Absolute links are taken as-is. Separators are normalised to
    the host OS so this works whether the code runs on Windows or POSIX.
    """
    pw = PureWindowsPath(stored_path)
    native = stored_path.replace("\\", os.sep).replace("/", os.sep)
    if pw.drive or pw.root:  # absolute (e.g. C:\... or \\server\share)
        candidate = native
    else:
        candidate = os.path.join(setup_dir, native)
    return os.path.normpath(candidate)
 
 
def _identity(resolved_path: str) -> str:
    """Filesystem-identity key so the same file via different links dedupes."""
    return os.path.normcase(os.path.realpath(resolved_path))


def _crosscheck_with_mikeio(setup_file: str, report: PackageReport) -> None:
    """Best-effort: warn if mikeio sees links we didn't (or vice versa)."""
    try:
        import mikeio
    except Exception:
        print("[verify] mikeio not installed; skipping cross-check.")
        return
 
    setup_dir = os.path.dirname(setup_file)
    ours = {
        _identity(_resolve(s, setup_dir))
        for s, _, exists in find_linked_files(setup_file)
        if exists
    }
 
    theirs: set[str] = set()
 
    def walk(sec):
        for _, v in sec.items():
            if hasattr(v, "items"):
                walk(v)
            elif isinstance(v, str) and v.startswith("|") and v.endswith("|") and len(v) > 2:
                r = _resolve(v[1:-1], setup_dir)
                if os.path.isfile(r):
                    theirs.add(_identity(r))
 
    try:
        walk(mikeio.read_pfs(setup_file))
    except Exception as exc:  # noqa: BLE001
        print(f"[verify] mikeio could not parse the file: {exc}")
        return
 
    only_ours = ours - theirs
    only_theirs = theirs - ours
    if not only_ours and not only_theirs:
        print(f"[verify] OK: mikeio agrees on all {len(ours)} linked files.")
    else:
        if only_theirs:
            print(f"[verify] mikeio found links we missed: {sorted(only_theirs)}")
        if only_ours:
            print(f"[verify] we found links mikeio missed: {sorted(only_ours)}")
 
