import mikeio
import numpy as np
from matplotlib.collections import PolyCollection
import matplotlib.pyplot as plt
from matplotlib import tri
from matplotlib.animation import FuncAnimation
from matplotlib.axis import Axis
from tqdm import tqdm
from matplotlib.colors import Normalize, LogNorm
from scipy.sparse import csr_matrix
from mpl_toolkits.axes_grid1 import make_axes_locatable

from .subplots import subplots, Colors, show

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 dist_in_meters(coords, pt):
    coords = np.atleast_2d(coords)
    xe = coords[:, 0]
    ye = coords[:, 1]
    xp = pt[0]
    yp = pt[1]
    d = np.sqrt(np.square(xe - xp) + np.square(ye - yp))
    return d

def relative_cumulative_distance(coords):
    coords = np.atleast_2d(coords)
    d = np.zeros_like(coords[:, 0])
    for j in range(1, len(d)):
        d[j] = d[j - 1] + dist_in_meters(coords[j, 0:2], coords[j - 1, 0:2])[0]
    return d

def plot_vertical_profile(da, timestep, s_coordinate, elements, cmin = None, cmax = None, label = "", add_colorbar = True, title = None, cmap = None, ax = None):
    values = da.values[timestep]
    z_coordinate = da._zn[timestep]
    sz = np.c_[s_coordinate, z_coordinate]
    verts = sz[elements]

    if cmap is None:
        cmap = "jet"
    norm = Normalize(vmin=cmin, vmax=cmax)
    pc = PolyCollection(verts, cmap=cmap, norm=norm, edgecolor="none", linewidth=0, antialiased=False)

    if cmin is None:
        cmin = np.nanmin(values)
    if cmax is None:
        cmax = np.nanmax(values)
    pc.set_clim(cmin, cmax)

    if add_colorbar:
        plt.colorbar(pc, ax=ax, label=label, orientation="vertical")
    pc.set_array(values)

    ax.add_collection(pc)
    ax.autoscale()
    ax.set_xlabel("Relative distance [m]")
    ax.set_ylabel("Z [m]")

    if title is not None:
        ax.set_title(title)

    return ax

def init_vertical_profile_plot(ax, da, s_coordinate, elements, cmin=None, cmax=None, label="", title_prefix="", cmap="jet", add_colorbar=True):
    # frame 0 init
    t0 = 0
    values0 = da.values[t0]
    z0 = da._zn[t0]

    verts0 = np.c_[s_coordinate, z0][elements]  # (nElem, nVert, 2)

    norm = Normalize(vmin=cmin, vmax=cmax)
    pc = PolyCollection(verts0, cmap=cmap, norm=norm, edgecolor="none", linewidth=0, antialiased=False)

    if cmin is None:
        cmin = np.nanmin(da.values)  # global limits for stable color scaling
    if cmax is None:
        cmax = np.nanmax(da.values)
    pc.set_clim(cmin, cmax)

    pc.set_array(values0)
    ax.add_collection(pc)
    ax.autoscale()
    ax.set_xlabel("Relative distance [m]")
    ax.set_ylabel("Z [m]")

    if title_prefix:
        ax.set_title(f"{title_prefix} (t={t0})")

    cbar = None
    if add_colorbar:
        cbar = plt.colorbar(pc, ax=ax, label=label, orientation="vertical")

    return pc, cbar

def animate_vertical_profile(da, s_coordinate, elements, ax=None, frames=None, interval=100, cmin=None, cmax=None, label="", title_prefix="", cmap="jet", repeat=True, blit=False, save_path=None, fps=10, xlim=None, ylim=None):
    if frames is None:
        frames = range(da.shape[0])  # all timesteps

    if ax is None:
        fig, ax = subplots(figwidth=20, figheight=10, nrows=1, ncols=1)
    else:
        fig = ax.get_figure()
    pc, _ = init_vertical_profile_plot(ax=ax, da=da, s_coordinate=s_coordinate, elements=elements, cmin=cmin, cmax=cmax, label=label, title_prefix=title_prefix, cmap=cmap, add_colorbar=True)
    if xlim is not None:
        ax.set_xlim(xlim)
    if ylim is not None:
        ax.set_ylim(ylim)
    pbar = tqdm(total=len(frames), desc="Animating timesteps", unit="frame")
    def update(t):
        pbar.update(1)
        values = da.values[t]
        z = da._zn[t]
        verts = np.c_[s_coordinate, z][elements]

        # update geometry + values
        pc.set_verts(verts)
        pc.set_array(values)
        ax.set_title(da.time[t].strftime("%Y-%m-%d %H:%M:%S"))

        return (pc,)

    anim = FuncAnimation(
        fig, update, frames=frames, interval=interval,
        repeat=repeat, blit=blit
    )
    if save_path is not None:
        # Choose writer based on extension
        if save_path.lower().endswith(".gif"):
            anim.save(save_path, writer="pillow", fps=fps)
        else:
            anim.save(save_path, writer="ffmpeg", fps=fps)
    else:
        plt.show()
    return anim

def animate(
        dfs: mikeio.Dataset|mikeio.DataArray,
        item: int = 0,
        frames: list = None,
        ax: Axis = None,
        xlim: list = None,
        ylim: list = None,
        title_prefix: str = "",
        label: str = None,
        fps: int = 10,
        save_path: str = None,
        cmin: float = None,
        cmax: float = None,
        cmap: str = None,
        add_bathy: bool = True,
        water_color: str = "#5ABCD8",
        bathy_color: str = "#C2B280",
):
    if not isinstance(dfs.geometry, mikeio.spatial.GeometryFMVerticalProfile):
        raise NotImplementedError(f"animate function currently only works with vertical profile type dfsu files")
    if isinstance(dfs, mikeio.DataArray):
        dfs = mikeio.Dataset(dfs)
        item = 0
    node_coordinates = dfs.geometry.node_coordinates
    element_table = dfs.geometry.element_table
    s_coordinate = relative_cumulative_distance(node_coordinates)
    elements = np.asarray([list(element_table[i]) for i in range(len(list(element_table)))])
    label = label if label is not None else f"{dfs[item].item.name} [{dfs[item].item.unit.short_name}]"
    if ax is None:
        fig, ax = subplots(figwidth=20, figheight=10, nrows=1, ncols=1)
    if add_bathy:
        ax.plot(dfs.geometry.element_coordinates[:, 0], dfs.geometry.element_coordinates[:, 2], color="black", linewidth=1.5)
    animate_vertical_profile(dfs[item], s_coordinate, elements, ax=ax, frames=frames, xlim=xlim, ylim=ylim, title_prefix=title_prefix, label=label, fps=fps, save_path=save_path, cmap=cmap, cmin=cmin, cmax=cmax)

def _contourf_geometryfm2d(geometry, values, vmin=None, vmax=None, cbar_n_levels=5, scale="linear", cmap="jet", ax=None, ):
    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)
    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
    
    

def contourf(
        geometry,
        values,
        vmin=None,
        vmax=None,
        cbar_n_levels=5,
        cmap=None,
        ax=None,
        scale="linear"
):
    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)