import mikeio
import numpy as np
from matplotlib.collections import PolyCollection
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation
from matplotlib.axis import Axis
from tqdm import tqdm
from matplotlib.colors import Normalize

from .subplots import subplots, Colors, show


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)
    