from pathlib import Path
import xarray
import numpy as np
from datetime import datetime
from matplotlib.pyplot import figure, draw, pause
from matplotlib.colors import LogNorm


def jointazel(cam: xarray.Dataset, ofn: Path = None, ttxt: str = ""):

    fg, axs = plotazel(cam, ttxt)
    # %% plot line from other camera to magnetic zenith
    overlayrowcol(axs[0], cam.rows, cam.cols, "g")
    overlayrowcol(axs[1], cam.rows, cam.cols, "g")
    # %% plot plane including this camera and line to magnetic zenith from other camera
    try:
        axs[0] = overlayrowcol(axs[0], cam.cutrow, cam.cutcol, "r", "1-D cut")
        axs[1] = overlayrowcol(axs[1], cam.cutrow, cam.cutcol, "r", "1-D cut")
    except AttributeError:
        pass  # this was not a 1-D plane case

    if ofn:
        ofn = Path(ofn).expanduser()
        print("saving", ofn)
        fg.savefig(ofn, bbox_inches="tight", dpi=100)

    if "Baz" in cam.attrs:
        axs[0].scatter(cam.Bcol, cam.Brow, 250, "red", "*", label="Mag. Zen.")
        axs[1].scatter(cam.Bcol, cam.Brow, 250, "red", "*", label="Mag. Zen.")

    for a in axs:
        a.legend()


def plotazel(data: xarray.Dataset, ttxt: str = ""):

    if "az" not in data or "el" not in data:
        return

    fg = figure(figsize=(12, 6))
    ax = fg.subplots(1, 2, sharey=True)

    ttxt = f"{data.filename} {ttxt}"
    if data.caltime is not None:
        ttxt += f"{data.caltime}"

    fg.suptitle(ttxt)

    c = ax[0].contour(data["az"].x, data["az"].y, data["az"])
    ax[0].clabel(c, fmt="%0.1f")
    ax[0].set_title("azimuth")
    ax[0].set_xlabel("x-pixels")
    ax[0].set_ylabel("y-pixels")

    c = ax[1].contour(data["el"].x, data["el"].y, data["el"])
    ax[1].clabel(c, fmt="%0.1f")
    ax[1].set_title("elevation")
    ax[1].set_xlabel("x-pixels")

    return fg, ax


def plottimeseries(data: np.ndarray, time: datetime, ttxt: str = ""):

    assert data.ndim == 2

    ax = figure().gca()

    ax.plot(time, data)
    ax.set_xlabel("time")
    ax.set_ylabel("brightness [data numbers]")
    ax.set_title(ttxt)


def overlayrowcol(ax, rows, cols, color: str = None, label: str = None):
    """
    plot FOV outline onto image via the existing axis "ax"

    inputs:
    ax: existing plot axis to overlay lines outlining FOV
    rows,cols: indices to plot
    """
    if rows is None or cols is None:
        return

    if len(rows) == 1 or isinstance(rows, (np.ndarray, xarray.DataArray)) and rows.ndim == 1:
        ax.scatter(cols, rows, color=color, alpha=0.5, marker=".", label=label)
    else:
        raise ValueError("unknonn row/col layout, was expecting 1-D")

    return ax


# %%


def plotasi(data: xarray.Dataset, ofn: Path = None):
    """
    rows,cols expect lines to be along rows Nlines x len(line)
    list of 1-D arrays or 2-D array
    """

    if data["imgs"].shape[0] == 0:
        return

    if ofn:
        ofn = Path(ofn).expanduser()
        odir = ofn.parent

    fg = figure()
    ax = fg.gca()

    hi = ax.imshow(
        data["imgs"][0], cmap="gray", origin="lower", norm=LogNorm(), interpolation="none"
    )  # priming
    ttxt = f"Themis ASI {data.site}\n"  # FOV vs. HST0,HST1: green,red '
    ht = ax.set_title(ttxt, color="g")
    ax.set_xlabel("x-pixels")
    ax.set_ylabel("y-pixels")
    ax.autoscale(True, tight=True)
    ax.grid(False)
    # %% plot narrow FOV outline
    if "imgs2" in data:
        overlayrowcol(ax, data.rows, data.cols)
    # %% play video
    try:
        for im in data["imgs"]:
            ts = str(im.time.astype("datetime64[us]").astype(datetime))
            hi.set_data(im)
            ht.set_text(ttxt + ts)
            draw()
            pause(0.01)
            if ofn:
                fn = odir / (ofn.stem + ts + ofn.suffix)
                print("saving", fn, end="\r")
                fg.savefig(fn, bbox_inches="tight", facecolor="k")
    except KeyboardInterrupt:
        return


def pcolormesh_nan(x: np.ndarray, y: np.ndarray, c: np.ndarray, cmap=None, axis=None):
    """handles NaN in x and y by smearing last valid value in column or row out,
    which doesn't affect plot because "c" will be masked too
    """

    mask = np.isfinite(x) & np.isfinite(y)
    top = None
    bottom = None

    for i, m in enumerate(mask):
        good = m.nonzero()[0]

        if good.size == 0:
            continue
        elif top is None:
            top = i
        else:
            bottom = i

        x[i, good[-1] :] = x[i, good[-1]]
        y[i, good[-1] :] = y[i, good[-1]]

        x[i, : good[0]] = x[i, good[0]]
        y[i, : good[0]] = y[i, good[0]]

    x[:top, :] = np.nanmean(x[top, :])
    y[:top, :] = np.nanmean(y[top, :])

    x[bottom:, :] = np.nanmean(x[bottom, :])
    y[bottom:, :] = np.nanmean(y[bottom, :])

    if axis is None:
        axis = figure().gca()

    return axis.pcolormesh(x, y, np.ma.masked_where(~mask, c), cmap=cmap)
