"""
Visualization module based on holoviews

"""
# Copyright 2018 European Union
# This file is part of pyposeidon.
# Licensed under the EUPL, Version 1.2 or – as soon they will be approved by the European Commission - subsequent versions of the EUPL (the "Licence").
# Unless required by applicable law or agreed to in writing, software distributed under the Licence is distributed on an "AS IS" basis, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the Licence for the specific language governing permissions and limitations under the Licence.
import cartopy.crs as ccrs
import geopandas as gp
import geoviews as gv
import holoviews as hv
import hvplot.xarray
import numpy as np
import pandas as pd
import shapely
import spatialpandas
import xarray as xr
from holoviews import opts
from holoviews.operation.datashader import datashade
from holoviews.operation.datashader import dynspread
from holoviews.operation.datashader import rasterize

from .utils import quads_to_tris

gv.extension("bokeh")
hv.output(widget_location="bottom")


@xr.register_dataset_accessor("hplot")
# @xr.register_dataarray_accessor('pplot')


class hplot(object):
    def __init__(self, xarray_obj):
        self._obj = xarray_obj

    def contourf(self, var="depth", it=None, tiles=False, **kwargs):
        x_var = kwargs.get("x", "SCHISM_hgrid_node_x")
        y_var = kwargs.get("y", "SCHISM_hgrid_node_y")
        tes_var = kwargs.get("e", "SCHISM_hgrid_face_nodes")
        t_var = kwargs.get("t", "time")

        x = self._obj[x_var][:].values
        y = self._obj[y_var][:].values
        try:
            t = self._obj[t_var][:].values
        except:
            pass

        tes = self._obj[tes_var].values[:, :4]

        # sort out quads
        try:
            mask = np.isnan(tes)[:, 3]
            tr3 = tes[mask][:, :3]
            tr3_ = quads_to_tris(tes[~mask])
            if tr3_:
                tri3 = np.append(tr3, tr3_, axis=0).astype(int)
            else:
                tri3 = tr3.astype(int)
        except:
            tri3 = tes.astype(int)

        if tri3.min() > 0:
            tri3 = tri3 - 1

        z = kwargs.get("z", self._obj[var].values[it, :].flatten())

        nodes = pd.DataFrame({"longitude": x, "latitude": y, "{}".format(var): z})
        elems = pd.DataFrame(tri3, columns=["a", "b", "c"])

        width = kwargs.get("width", 800)
        height = kwargs.get("height", 600)
        opts.defaults(opts.WMTS(width=width, height=height))

        tile = gv.WMTS("https://b.tile.openstreetmap.org/{Z}/{X}/{Y}.png")

        points = gv.operation.project_points(gv.Points(nodes, vdims=["{}".format(var)]))

        trimesh = gv.TriMesh((elems, points))

        if tiles:
            return tile * rasterize(trimesh, aggregator="mean").opts(
                colorbar=True, cmap="Viridis", padding=0.1, tools=["hover"]
            )
        else:
            return rasterize(trimesh, aggregator="mean").opts(
                colorbar=True,
                cmap="Viridis",
                padding=0.1,
                tools=["hover"],
                width=width,
                height=height,
            )

    def meteo(self, **kwargs):
        m = self._obj.set_coords(["lon", "lat"])
        m["lon"] = m["lon"].isel(time=1, drop=True)
        m["lat"] = m["lat"].isel(time=1, drop=True)

        width = kwargs.get("width", 800)
        height = kwargs.get("height", 600)

        return m.prmsl.hvplot.quadmesh(
            x="lon",
            y="lat",
            geo=True,
            widget_location="bottom",
            datashade=True,
            tiles=True,
            width=width,
            height=height,
        )

    def quiver(self, **kwargs):
        return

    def mesh(self, tiles=False, **kwargs):
        x_var = kwargs.get("x", "SCHISM_hgrid_node_x")
        y_var = kwargs.get("y", "SCHISM_hgrid_node_y")
        tes_var = kwargs.get("e", "SCHISM_hgrid_face_nodes")

        x = self._obj[x_var][:].values
        y = self._obj[y_var][:].values
        tes = self._obj[tes_var].values[:, :4]

        # check tesselation

        try:
            tes = tes.fillna(-1.0)  # adapt to schism >=5.10
            tes = tes.values
        except:
            pass  # given as numpy array

        try:
            if np.isnan(tes[:, 3]).all():
                tes = tes[:, :3]
        except:
            pass

        try:
            if (tes[:, 3] == -1).all():
                tes = tes[:, :3]
        except:
            pass

        width = kwargs.get("width", 800)
        height = kwargs.get("height", 600)
        opts.defaults(opts.WMTS(width=width, height=height))

        tile = gv.WMTS("https://b.tile.openstreetmap.org/{Z}/{X}/{Y}.png")

        nodes = pd.DataFrame({"longitude": x, "latitude": y})

        points = gv.operation.project_points(gv.Points(nodes), projection=ccrs.PlateCarree())

        if tes.shape[1] == 3:
            elems = pd.DataFrame(tes, columns=["a", "b", "c"])

            if elems.min().min() > 0:
                elems = elems - 1

            trimesh = gv.TriMesh((elems, points)).edgepaths

            if tiles:
                return tile * datashade(trimesh, precompute=True, cmap=["black"])
            else:
                return datashade(trimesh, precompute=True, cmap=["green"]).opts(width=width, height=height)

        else:  # there are quads
            elems = pd.DataFrame(tes, columns=["a", "b", "c", "d"])

            if elems.min().min() > 0:
                elems = elems - 1

            quads = elems.loc[~elems.d.isna()].copy()
            quads = quads.reset_index(drop=True)
            ap = nodes.loc[quads.a, ["longitude", "latitude"]]
            bp = nodes.loc[quads.b, ["longitude", "latitude"]]
            cp = nodes.loc[quads.c, ["longitude", "latitude"]]
            dp = nodes.loc[quads.d, ["longitude", "latitude"]]
            quads["ap"] = ap.values.tolist()
            quads["bp"] = bp.values.tolist()
            quads["cp"] = cp.values.tolist()
            quads["dp"] = dp.values.tolist()

            n = 2
            al = quads.ap + quads.bp + quads.cp + quads.dp + quads.ap
            coords = [[l[i : i + n] for i in range(0, len(l), n)] for l in al]

            quads["coordinates"] = coords

            qpolys = shapely.polygons(quads.coordinates.to_list())

            df = gp.GeoDataFrame(geometry=qpolys)

            df_ = spatialpandas.GeoDataFrame(df)

            q = gv.Path(df_)

            qdf = dynspread(rasterize(q, precompute=True).options(cmap=["black"]))

            triangles = elems.loc[elems.d.isna()]

            trimesh = gv.TriMesh((triangles, points)).edgepaths

            wireframe = dynspread(rasterize(trimesh, precompute=True).opts(cmap=["black"]))

            if tiles:
                return tile * wireframe * qdf
            else:
                return wireframe.opts(cmap=["green"]) * qdf.opts(cmap=["green"])

    def qframes(self, **kwargs):
        return

    def frames(self, var="depth", tiles=False, **kwargs):
        x_var = kwargs.get("x", "SCHISM_hgrid_node_x")
        y_var = kwargs.get("y", "SCHISM_hgrid_node_y")
        tes_var = kwargs.get("e", "SCHISM_hgrid_face_nodes")
        t_var = kwargs.get("t", "time")

        x = self._obj[x_var][:].values
        y = self._obj[y_var][:].values
        t = self._obj[t_var][:].values
        tes = self._obj[tes_var].values[:, :4]

        # check tesselation

        # sort out quads
        try:
            mask = tes[:, 3] == -1
            tr3 = tes[mask][:, :3]
            tr3_ = quads_to_tris(tes[~mask])
            if tr3_:
                tri3 = np.append(tr3, tr3_, axis=0).astype(int)
            else:
                tri3 = tr3.astype(int)
        except:
            tri3 = tes[:, :3].astype(int)

        if tri3.min() > 0:
            tri3 = tri3 - 1

        var = kwargs.get("var", "depth")
        z = kwargs.get("z", self._obj[var].values[0, :].flatten())

        zmin = self._obj[var].values.min()
        zmax = self._obj[var].values.max()

        nodes = pd.DataFrame({"longitude": x, "latitude": y, "{}".format(var): z})
        elems = pd.DataFrame(tri3, columns=["a", "b", "c"])

        width = kwargs.get("width", 800)
        height = kwargs.get("height", 600)
        opts.defaults(opts.WMTS(width=width, height=height))

        tile = gv.WMTS("https://b.tile.openstreetmap.org/{Z}/{X}/{Y}.png")

        points = gv.operation.project_points(gv.Points(nodes, vdims=["{}".format(var)]))

        trimesh = gv.TriMesh((elems, points))

        def time_mesh(time):
            points.data[var] = self._obj[var].sel(time=time).values
            return gv.TriMesh((elems, points))  # , crs=ccrs.GOOGLE_MERCATOR)

        meshes = hv.DynamicMap(time_mesh, kdims=["Time"]).redim.values(Time=t)

        imesh = rasterize(meshes, aggregator="mean").opts(
            cmap="viridis",
            colorbar=True,
            padding=0.1,
            tools=["hover"],
            clim=(zmin, zmax),
        )

        if tiles:
            return hv.output(tile * imesh, holomap="scrubber", fps=1)
        else:
            return hv.output(imesh.opts(width=width, height=height), holomap="scrubber", fps=1)


# from https://pandas.pydata.org/pandas-docs/stable/development/extending.html
@pd.api.extensions.register_dataframe_accessor("hplot")
class BoundaryAccessor:
    def __init__(self, pandas_obj):
        self._validate(pandas_obj)
        self._obj = pandas_obj

    @staticmethod
    def _validate(obj):
        # verify there is a column latitude and a column longitude
        if "lat" not in obj.columns or "lon" not in obj.columns:
            raise AttributeError("Must have 'lat' and 'lon'.")

    @property
    def center(self):
        # return the geographic center point of this DataFrame
        lat = self._obj.lat
        lon = self._obj.lon
        return (float(lon.mean()), float(lat.mean()))

    def show(self, **kwargs):
        width = kwargs.get("width", 800)
        height = kwargs.get("height", 600)
        opts.defaults(opts.WMTS(width=width, height=height))

        tile = gv.WMTS("https://b.tile.openstreetmap.org/{Z}/{X}/{Y}.png")

        ps = [tile]

        # external
        ds = self._obj.loc["line0"]
        ds = ds.append(ds.iloc[0]).reset_index(drop=True)

        dland = ds.loc[ds.tag == "land"]
        dland_ = np.split(dland, np.flatnonzero(np.diff(dland.index) != 1) + 1)
        for seg in dland_:
            ps.append(seg.hvplot.paths(x="lon", y="lat", geo=True, color="brown"))

        dwater = ds.loc[ds.tag == "open"]
        dwater_ = np.split(dwater, np.flatnonzero(np.diff(dwater.index) != 1) + 1)
        for seg in dwater_:
            ps.append(seg.hvplot.paths(x="lon", y="lat", geo=True, color="blue"))

        # islands
        for line in self._obj.index.levels[0][1:]:
            ds = self._obj.loc[line]
            ds = ds.append(ds.iloc[0]).reset_index(drop=True)
            ps.append(ds.hvplot.paths(x="lon", y="lat", geo=True, color="green"))

        return hv.Overlay(ps)
