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



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 _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)}")