from __future__ import annotations
from typing import Dict, Mapping, List, Tuple

import demes
import numpy as np
import matplotlib
import matplotlib.patheffects

from . import utils


class Tube:
    """
    A deme represented as a tube. The tube has a length along the time
    dimension and a width equal to the deme's population size, which may
    change over time.

    :ivar list[float] time: Coordinates along the time dimension.
    :ivar list[float] size1: Coordinates along the non-time dimension,
        corresponding to the first side of the tube.
    :ivar list[float] size2: Coordinates along the non-time dimension,
        corresponding to the second side of the tube.
    """

    def __init__(
        self,
        deme: demes.Deme,
        mid: float,
        inf_start_time: float,
        log_time: bool = False,
    ):
        """
        :param demes.Deme deme: The deme for which to calculate coordinates.
        :param float mid: The mid point of the deme along the non-time dimension.
        :param float inf_start_time: The value along the time dimension which
            is used instead of infinity (for epochs with infinite start times).
        :param bool log_time: The time axis uses a log-10 scale.
        """
        self.deme = deme
        self.mid = mid
        self._coords(deme, mid, inf_start_time, log_time)

    def _coords(
        self,
        deme: demes.Deme,
        mid: float,
        inf_start_time: float,
        log_time: bool,
        num_points: int = 100,
    ) -> None:
        """Calculate tube coordinates."""
        time: List[float] = []
        size1: List[float] = []
        size2: List[float] = []
        for k, epoch in enumerate(deme.epochs):
            start_time = epoch.start_time
            if np.isinf(start_time):
                start_time = inf_start_time
            end_time = epoch.end_time

            if epoch.size_function == "constant":
                t = np.array([start_time, end_time])
                N1 = np.array([mid - epoch.start_size / 2] * 2)
                N2 = np.array([mid + epoch.end_size / 2] * 2)
            elif epoch.size_function == "exponential":
                if log_time:
                    t = np.exp(
                        np.linspace(
                            np.log(start_time), np.log(max(1, end_time)), num=num_points
                        )
                    )
                else:
                    t = np.linspace(start_time, end_time, num=num_points)
                dt = (start_time - t) / (start_time - end_time)
                r = np.log(epoch.end_size / epoch.start_size)
                N = epoch.start_size * np.exp(r * dt)
                N1 = mid - N / 2
                N2 = mid + N / 2
            elif epoch.size_function == "linear":
                if log_time:
                    t = np.exp(
                        np.linspace(
                            np.log(start_time), np.log(max(1, end_time)), num=num_points
                        )
                    )
                else:
                    t = np.linspace(start_time, end_time, num=num_points)
                dt = (start_time - t) / (start_time - end_time)
                N = epoch.start_size + (epoch.end_size - epoch.start_size) * dt
                N1 = mid - N / 2
                N2 = mid + N / 2
            else:
                raise ValueError(
                    f"Don't know how to draw epoch {k} with "
                    f'"{epoch.size_function}" size_function.'
                )

            time.extend(t)
            size1.extend(N1)
            size2.extend(N2)

        self.time = time
        self.size1 = size1
        self.size2 = size2

    def sizes_at(self, time: float) -> Tuple[float, float]:
        """Return the size coordinates of the tube at the given time."""
        N = self.deme.size_at(time)
        N1 = self.mid - N / 2
        N2 = self.mid + N / 2
        return N1, N2

    def to_path(self) -> matplotlib.path.Path:
        """Return a matplotlib.path.Path for the tube outline."""
        # Compared with separately drawing each side of the tube,
        # drawing a path that includes both sides of the tube means:
        #  * vector formats use one "stroke" rather than two; and
        #  * matplotlib path effects apply once, avoiding artifacts for
        #    very narrow tubes, where the path effect from the side
        #    drawn second is on top of the line drawn for the first side.
        N = list(self.size1) + list(self.size2)
        t = list(self.time) * 2
        vertices = list(zip(N, t))
        codes = (
            [matplotlib.path.Path.MOVETO]
            + [matplotlib.path.Path.LINETO] * (len(self.time) - 1)
        ) * 2
        return matplotlib.path.Path(vertices, codes)


def find_positions(graph: demes.Graph) -> Dict[str, float]:
    """
    Find optimal deme positions along a single dimension by:

      - ordering the demes to minimise the number of lines crossing demes,
      - minimising the distance from each parent deme to the mean position of
        its children,
      - minimising the distance between interacting demes (where interactions
        are either migrations or pulses).

    In addition, the solution is constrained so that contemporary demes
    have a minimum separation distance.

    :param graph:
        The graph for which positions should be obtained.
    :return:
        A dictionary mapping deme names to positions.
    """
    sep = utils.separation_heuristic(graph)
    positions = utils.minimal_crossing_positions(
        graph, sep=sep, unique_interactions=False
    )
    return utils.optimise_positions(
        graph, positions=positions, sep=sep, unique_interactions=False
    )


def tubes(
    graph: demes.Graph,
    ax: matplotlib.axes.Axes | None = None,
    colours: utils.ColourOrColourMapping | None = None,
    log_time: bool = False,
    title: str | None = None,
    inf_ratio: float = 0.2,
    positions: Mapping[str, float] | None = None,
    num_lines_per_migration: int = 10,
    seed: int | None = None,
    max_time: float | None = None,
    labels: str = "xticks-mid",
    fill: bool = True,
    scale_bar: bool = False,
) -> matplotlib.axes.Axes:
    """
    Plot a demes-as-tubes schematic of the graph and the demes' relationships.

    Each deme is depicted as a tube, where the tube’s width is
    proportional to the deme’s size at any given time.
    Horizontal lines with arrows indicate either:

     * an ancestor/descendant relation (thick solid lines, open arrow heads),
     * an admixture pulse (dashed lines, closed arrow heads),
     * or a period of continuous migration (thin solid lines, closed arrow heads).

    Lines are drawn in the colour of the ancestor or source deme, and arrows
    point from ancestor to descendant or from source to dest.
    For each period of continuous migration, multiple thin lines are drawn.
    The time for each migration line is sampled uniformly at
    random from the migration's time interval (or log-uniformly for a
    log-scaled time axis). Symmetric migrations have lines in both directions.

    If ``positions`` are not specified, the positions will be chosen
    automatically such that line crossings are minimised and related demes
    are close together. If the automatically chosen positions are unexpected,
    please open an issue at https://github.com/grahamgower/demesdraw/issues/.

    :param demes.Graph graph:
        The demes graph to plot.
    :param ax:
        The matplotlib axes onto which the figure
        will be drawn. If None, an empty axes will be created for the figure.
    :type ax: Optional[matplotlib.axes.Axes]
    :param colours:
        A mapping from deme name to matplotlib colour. Alternately,
        ``colours`` may be a named colour that will be used for all demes.
    :type colours: Optional[dict or str]
    :param log_time:
        If True, use a log-10 scale for the time axis.
        If False (*default*), a linear scale will be used.
    :param title:
        The title of the figure.
    :param inf_ratio:
        The proportion of the time axis that will be
        used for the time interval which stretches towards infinity.
    :param positions:
        A dictionary mapping deme names to horizontal coordinates
        (the mid point of the deme's tube).
        Note that the width of a deme is the deme's (max) size,
        so the positions should allow sufficient space to avoid overlapping.
    :param num_lines_per_migration:
        The number of lines to draw per migration. For symmetric migrations,
        this number of lines will be drawn in each direction.
    :param seed:
        Seed for the random number generator. The generator is
        used to sample times for migration lines.
    :param max_time:
        The maximum time value shown in the figure.
        If demographic events (e.g. size changes, migrations, common ancestor
        events) occur before this time, those events will not be visible.
        If no demographic events occur before this time, the root demes will be
        drawn so they extend to the given time.
    :param labels:
        A string indicating where the deme names should be drawn, or ``None``.
        The options are:

         * "xticks": for extant demes, labels are written under the x-axis
           as matplotlib xticklabels. Labels are not written for extinct demes.
         * "legend": labels for colour patches are written in a legend.
         * "mid": labels are written in the middle of each deme's tube.
         * "xticks-legend": for extant demes, labels are written under the
           x-axis as matplotlib xticklabels. For extinct demes, labels for
           colour patches are written in a legend.
         * "xticks-mid" (*default*): for extant demes, labels are written
           under the x-axis as matplotlib xticklabels. For extinct demes,
           labels are written in the middle of each deme's tube.
         * ``None``: no labels are written.

    :type labels: str or None
    :param fill:
        If True, the inside of the tubes will be painted.
        If False, only the outline of the tubes will be drawn.
    :param scale_bar:
        If True, draw a scale bar that indicates population size.
        If False (*default*), no scale bar will be shown.
    :return:
        The matplotlib axes onto which the figure was drawn.
    """
    if labels not in ("xticks", "legend", "mid", "xticks-legend", "xticks-mid", None):
        raise ValueError(f"Unexpected value for labels: '{labels}'")

    if ax is None:
        _, ax = utils.get_fig_axes()

    if log_time:
        ax.set_yscale("log", base=10)

    colours = utils._get_colours(graph, colours)

    rng = np.random.default_rng(seed)
    inf_start_time = max_time
    if inf_start_time is None:
        inf_start_time = utils._inf_start_time(graph, inf_ratio, log_time)

    if positions is None:
        positions = find_positions(graph)

    tubes = {}
    ancestry_arrows = []

    for j, deme in enumerate(graph.demes):
        colour = colours[deme.name]
        plot_kwargs = dict(
            color=colour,
            # solid_capstyle="butt",
            # capstyle="butt",
            zorder=1,
            path_effects=[
                matplotlib.patheffects.withStroke(linewidth=3, foreground="white")
            ],
        )

        mid = positions[deme.name]

        tube = Tube(deme, mid, inf_start_time, log_time=log_time)
        tubes[deme.name] = tube

        path_patch = matplotlib.patches.PathPatch(
            tube.to_path(), capstyle="butt", fill=False, **plot_kwargs
        )
        ax.add_patch(path_patch)
        if fill:
            poly = ax.fill_betweenx(
                tube.time,
                tube.size1,
                tube.size2,
                facecolor=colour,
                edgecolor="none",
                alpha=0.5,
                zorder=1,
            )
            poly._demesdraw_data = dict(deme=deme)  # type: ignore[attr-defined]

        # Indicate ancestry from ancestor demes.
        tube_frac = np.linspace(
            tube.size1[0], tube.size2[0], max(2, len(deme.ancestors))
        )
        left = 0
        right = 0
        arrow_kwargs = dict(
            arrowstyle="-|>",
            mutation_scale=10,
            alpha=1,
            facecolor=(1, 1, 1, 0),
            path_effects=[
                matplotlib.patheffects.withStroke(linewidth=3, foreground="white")
            ],
            zorder=2,
        )
        for ancestor_id, proportion in zip(deme.ancestors, deme.proportions):
            anc_size1, anc_size2 = tubes[ancestor_id].sizes_at(deme.start_time)
            if anc_size2 < tube.size1[0]:
                # Ancestor is to the left.
                x_pos = [anc_size2, tube_frac[left]]
                ancestry_arrows.append(
                    (
                        tube.time[0],
                        min(x_pos),
                        max(x_pos),
                        colours[ancestor_id],
                    )
                )
                left += 1
            elif tube.size2[0] < anc_size1:
                # Ancestor is to the right.
                x_pos = [anc_size1, tube_frac[len(tube_frac) - right - 1]]
                ancestry_arrows.append(
                    (
                        tube.time[0],
                        max(x_pos),
                        min(x_pos),
                        colours[ancestor_id],
                    )
                )
                right += 1

    slots: Dict[int, int] = {}
    for j, (time, x1, x2, colour) in enumerate(ancestry_arrows):
        taken = set()
        for k, slot in slots.items():
            k_time, k_x1, k_x2, _ = ancestry_arrows[k]
            if np.isclose(time, k_time) and not (
                min(x1, x2) > max(k_x1, k_x2) or min(k_x1, k_x2) > max(x1, x2)
            ):
                taken.add(slot)
        slot = 0
        i = 0
        while slot in taken:
            i += 1
            odd = i % 2
            delta = i // 2
            slot = -odd * delta + (1 - odd) * delta
        slots[j] = slot
        radius = slot * 0.15
        arr = matplotlib.patches.FancyArrowPatch(
            (x1, time),
            (x2, time),
            connectionstyle=f"arc3,rad={radius}",
            edgecolor=colour,
            **arrow_kwargs,  # type: ignore[arg-type]
        )
        arr.set_sketch_params(1, 100, 2)
        # Ignore potential overflow when using log scale. Infinity is fine.
        with np.errstate(over="ignore"):
            ax.add_patch(arr)

    # Update the axes view. ax.add_patch() doesn't do this itself.
    ax.autoscale_view()

    # Calculate an offset for the space between arrowhead and tube.
    xlim = ax.get_xlim()
    offset = 0.01 * (xlim[1] - xlim[0])

    def migration_line(source, dest, time, metadata, **mig_kwargs):
        """Draw a migration line from source to dest at the given time."""
        source_size1, source_size2 = tubes[source].sizes_at(time)
        dest_size1, dest_size2 = tubes[dest].sizes_at(time)

        if source_size2 < dest_size1:
            x = [source_size2 + offset, dest_size1 - offset]
            arrow = ">"
            lr = "r"
        else:
            x = [source_size1 - offset, dest_size2 + offset]
            arrow = "<"
            lr = "l"

        metadata = dict(lr=lr, **metadata)

        colour = colours[source]
        lines = ax.plot(
            x,
            [time, time],
            color=colour,
            path_effects=[matplotlib.patheffects.Normal()],
            zorder=-1,
            **mig_kwargs,
        )
        for line in lines:
            line.set_sketch_params(1, 100, 2)
            line._demesdraw_data = metadata
        lines = ax.plot(
            x[1],
            time,
            arrow,
            markersize=2,
            color=colour,
            path_effects=[matplotlib.patheffects.Normal()],
            **mig_kwargs,
        )
        for line in lines:
            line._demesdraw_data = metadata

    def random_migration_time(migration, log_scale: bool) -> float:
        start_time = migration.start_time
        if np.isinf(start_time):
            start_time = inf_start_time
        if log_scale:
            t = np.exp(
                rng.uniform(np.log(max(1, migration.end_time)), np.log(start_time))
            )
        else:
            t = rng.uniform(migration.end_time, start_time)
        return t

    # Plot migration lines.
    migration_kwargs = dict(linewidth=0.2, alpha=0.5)
    for migration in graph.migrations:
        for _ in range(num_lines_per_migration):
            t = random_migration_time(migration, log_time)
            migration_line(
                migration.source,
                migration.dest,
                t,
                dict(migration=migration),
                **migration_kwargs,
            )

    # Plot pulse lines.
    for pulse in graph.pulses:
        for j, source in enumerate(pulse.sources):
            migration_line(
                source,
                pulse.dest,
                pulse.time,
                dict(pulse=pulse, source=j),
                linestyle="--",
                linewidth=1,
            )

    xticks = []
    xticklabels = []
    deme_labels = [deme.name for deme in graph.demes]

    if labels in ("xticks", "xticks-legend", "xticks-mid"):
        # Put labels for the leaf nodes underneath the figure.
        leaves = [deme.name for deme in graph.demes if deme.end_time == 0]
        deme_labels = [deme.name for deme in graph.demes if deme.end_time != 0]
        xticks = [positions[leaf] for leaf in leaves]
        xticklabels = leaves

    if len(deme_labels) > 0 and labels in ("legend", "xticks-legend"):
        # Add legend.
        ax.legend(
            handles=[
                matplotlib.patches.Patch(
                    edgecolor=matplotlib.colors.to_rgba(colours[deme_name], 1.0),
                    facecolor=(
                        matplotlib.colors.to_rgba(colours[deme_name], 0.3)
                        if fill
                        else None
                    ),
                    label=deme_name,
                )
                for deme_name in deme_labels
            ],
            # Use a horizontal layout, rather than vertical.
            ncol=max(1, len(deme_labels) // 2),
        )

    if labels in ("mid", "xticks-mid"):
        for deme_name in deme_labels:
            if log_time:
                tmid = np.exp(
                    (
                        np.log(tubes[deme_name].time[0])
                        + np.log(max(1, tubes[deme_name].time[-1]))
                    )
                    / 2
                )
            else:
                tmid = (tubes[deme_name].time[0] + tubes[deme_name].time[-1]) / 2
            ax.text(
                positions[deme_name],
                tmid,
                deme_name,
                ha="center",
                va="center",
                # Give the text some contrast with its background.
                bbox=dict(boxstyle="round", fc="white", ec="none", alpha=0.6, pad=0.2),
            )

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

    ax.set_xticks(xticks)
    ax.set_xticklabels(xticklabels)
    ax.tick_params(axis="x", length=0)
    ax.tick_params(axis="x", which="minor", bottom=False)
    ax.spines[["top", "right", "bottom"]].set_visible(False)

    ax.set_ylabel(f"time ago ({graph.time_units})")

    ax.set_ylim(1 if log_time else 0, inf_start_time)

    if scale_bar:
        # Plot a scale bar that indicates the population size.
        width = utils.size_max(graph)
        transform = matplotlib.transforms.blended_transform_factory(
            ax.transData, ax.transAxes
        )
        x_offset = min(min(t.size1) for t in tubes.values())
        ax_sb = ax.inset_axes((x_offset, -0.16, width, 1e-6), transform=transform)

        ax_sb.yaxis.set_major_locator(matplotlib.ticker.NullLocator())
        ax_sb.spines[["left", "right", "top"]].set_visible(False)
        ax_sb.xaxis.set_ticks_position("top")
        locator = matplotlib.ticker.AutoLocator()
        locator.set_params(integer=True)
        ax_sb.xaxis.set_major_locator(locator)
        ax_sb.xaxis.set_minor_locator(matplotlib.ticker.AutoMinorLocator())
        ax_sb.set_xlim(0, width)
        ax_sb.set_xlabel("deme size (individuals)")
        ax_sb.xaxis.set_label_position("bottom")

    # Status bar text when in interactive mode.
    def format_coord(x, y):
        return f"time={y:.1f}"

    ax.format_coord = format_coord  # type: ignore[method-assign]

    # Note: we must hold a reference to the motion_notify_event callback,
    # or the callback will get garbage collected and thus disabled.
    # So we just add an attribute to the returned Axes object.
    ax._demesdraw_mouseoverpopup = (  # type: ignore[attr-defined]
        _MouseoverPopup(ax.figure, ax)
    )

    return ax


# Copied and modified from:
# https://matplotlib.org/stable/tutorials/advanced/blitting.html
class _BlitManager:
    def __init__(self, canvas, animated_artists=()):
        """
        Parameters
        ----------
        canvas : FigureCanvasAgg
            The canvas to work with, this only works for sub-classes of the Agg
            canvas which have the `~FigureCanvasAgg.copy_from_bbox` and
            `~FigureCanvasAgg.restore_region` methods.

        animated_artists : Iterable[Artist]
            List of the artists to manage
        """
        self.canvas = canvas
        self._bg = None
        self._artists = []

        if hasattr(canvas, "copy_from_bbox") and hasattr(canvas, "restore_region"):
            # Webagg and related backends support blitting as of matplotlib 3.4,
            # but supports_blit is False because of some rendering artifacts.
            # However, not using blitting is very laggy, so we use it anyway.
            self.canvas.supports_blit = True
        if not self.canvas.supports_blit:
            return

        for a in animated_artists:
            self.add_artist(a)
        # grab the background on every draw
        self.cid = canvas.mpl_connect("draw_event", self.on_draw)

    def on_draw(self, event):
        """Callback to register with 'draw_event'."""
        cv = self.canvas
        if event is not None:
            if event.canvas != cv:
                return
        self._bg = cv.copy_from_bbox(cv.figure.bbox)
        self._draw_animated()

    def add_artist(self, art):
        """
        Add an artist to be managed.

        Parameters
        ----------
        art : Artist

            The artist to be added.  Will be set to 'animated' (just
            to be safe).  *art* must be in the figure associated with
            the canvas this class is managing.

        """
        assert art.figure == self.canvas.figure
        art.set_animated(True)
        self._artists.append(art)

    def _draw_animated(self):
        """Draw all of the animated artists."""
        fig = self.canvas.figure
        for a in self._artists:
            fig.draw_artist(a)

    def update(self):
        """Update the screen with animated artists."""
        if not self.canvas.supports_blit:
            self.canvas.draw_idle()
            return
        cv = self.canvas
        fig = cv.figure
        # paranoia in case we missed the draw event,
        if self._bg is None:
            self.on_draw(None)
        else:
            # restore the background
            cv.restore_region(self._bg)
            # draw all of the animated artists
            self._draw_animated()
            # update the GUI state
            cv.blit(fig.bbox)


class _MouseoverPopup:
    def __init__(self, fig, ax):
        """Setup interactive mouseover annotations."""

        self.ax = ax
        self.popup = ax.annotate(
            "",
            xy=(0, 0),
            xytext=(20, 20),
            textcoords="offset points",
            multialignment="left",
            horizontalalignment="center",
            verticalalignment="center",
            bbox=dict(boxstyle="round", fc="yellow"),
            arrowprops=dict(arrowstyle="->"),
        )
        self.popup.set_visible(False)
        self.bm = _BlitManager(fig.canvas, [self.popup])

        # Detect mouseover for artists with a _demesdraw_data attribute,
        # which were added in demesdraw.tubes().
        self.artists = []
        for artist in reversed(ax.get_children()):
            if hasattr(artist, "_demesdraw_data"):
                self.artists.append(artist)

        # Get the midpoint of the figure in data coords.
        axis_to_data = self.ax.transAxes + self.ax.transData.inverted()
        self.xmid, self.ymid = axis_to_data.transform((0.5, 0.5))

        # Save a reference to the callback so that it doesn't get garbage collected.
        self.cid = fig.canvas.mpl_connect("motion_notify_event", self.motion_callback)

    def motion_callback(self, event):
        visible = self.popup.get_visible()
        mouse_inside_artist = False
        if event.inaxes == self.ax:
            for artist in self.artists:
                contained, _ = artist.contains(event)
                if contained:
                    data = artist._demesdraw_data
                    label = self.get_mouseover_label(data, event)
                    x = event.xdata
                    y = event.ydata
                    xt = 50
                    yt = 50
                    if x > self.xmid:
                        xt *= -1
                    if y > self.ymid:
                        yt *= -1
                    self.popup.xy = (x, y)
                    self.popup.xyann = (xt, yt)  # xytext
                    self.popup.set_text(label)
                    self.popup.set_visible(True)
                    mouse_inside_artist = True
                    self.bm.update()
                    break

        if not mouse_inside_artist and visible:
            self.popup.set_visible(False)
            self.bm.update()

    def get_arrow(self, lr, source, dest):
        assert lr in ("l", "r")
        arrow = r"$\rightarrow$" if lr == "r" else r"$\leftarrow$"
        if lr == "r":
            return f"{source} {arrow} {dest}"
        else:
            return f"{dest} {arrow} {source}"

    def get_mouseover_label(self, data, event):
        if "deme" in data:
            deme = data["deme"]
            time = max(min(event.ydata, deme.start_time - 1e-6), deme.end_time)
            size = deme.size_at(time)
            ancestors = ", ".join(deme.ancestors)
            proportions = ", ".join(map(str, deme.proportions))
            label = (
                f"deme: {deme.name}\n"
                f"time: ({deme.start_time}, {deme.end_time}]\n"
                f"ancestors: {ancestors}\n"
                f"proportions: {proportions}\n"
                f"size: {int(size)}"
            )
        elif "migration" in data:
            migration = data["migration"]
            arrow = self.get_arrow(data["lr"], migration.source, migration.dest)
            label = (
                f"migration: {arrow}\n"
                f"time: ({migration.start_time}, {migration.end_time}]\n"
                f"rate: {migration.rate}"
            )
        elif "pulse" in data:
            pulse = data["pulse"]
            j = data["source"]
            arrow = self.get_arrow(data["lr"], pulse.sources[j], pulse.dest)
            label = (
                f"pulse: {arrow}\n"
                f"time: {pulse.time}\n"
                f"proportion: {pulse.proportions[j]}"
            )
        return label
