#!/usr/bin/env python3
"""Plot the KKT trajectory mu -> x(mu) of the projection of x0 onto a 2D quadric.

Panels: the trajectory in the plane, the secular function f(mu), the distance ||x0 - x(mu)||
and f'(mu). Colors encode mu, from mu -> -inf (blue) to mu -> +inf (red). The quadric, x(mu),
f, f' and the degenerate points come from quadproj.
"""

import argparse
import shutil
from pathlib import Path

import matplotlib.pyplot as plt
import numpy as np
from matplotlib import colormaps
from matplotlib.cm import ScalarMappable
from matplotlib.collections import LineCollection
from matplotlib.colors import Normalize
from matplotlib.lines import Line2D
from quadproj.fun import Fun
from quadproj.project import get_squared_radius_k, get_xd_k
from quadproj.quadrics import Quadric
from scipy.optimize import brentq

# One figure per case, for the quadric x'Ax + b'x + c = 0 (b = 0 and c = -1 by default).
CASES = {
    'x_trajectories_mu_1': dict(A=np.diag([2.0, 1.0]), x0=[0.0, -0.1]),
    'x_trajectories_mu_2': dict(A=np.diag([2.0, 0.75]), x0=[0.0, -0.6]),
    'x_trajectories_mu_parabola': dict(A=np.diag([1.0, 0.0]), b=[0.0, -1.0], c=0.0, x0=[0.3, 0.2]),
}

TOL = 1e-6            # eigenvalues and normalized components of x0 below TOL are considered zero

N_PER_BRANCH = 266    # mu samples between consecutive poles
N_TAIL = 100          # extra samples extending the outer branches geometrically up to |mu| ~ 1e4 / min |lambda|
EPS = 0.01            # gap left on each side of a pole
MU_MARGIN = 2.5       # the mu axis extends MU_MARGIN beyond the poles and 0
N_DOTTED_OUTER = 2    # points drawn dotted at the ends of the outer branches
Y_PADDING = 0.1       # fraction of the y-range of the features added below and above them
WINDOW_SCALE = 2      # the trajectory panel is WINDOW_SCALE times as wide as x0 and the KKT points

COLORMAP = 'RdYlBu_r'
QUADRIC_COLOR = 'royalblue'
STAR_STYLE = dict(color='purple', marker='v', s=12, zorder=2)
ROOT_STYLE = dict(color='purple', marker='s', s=12, zorder=2)
POLE_STYLE = dict(linestyle='dashed', color='grey')
TEXT_STYLE = dict(ha='center', va='center', fontsize=12, clip_on=True)
DIST_D_LABEL = r'$||\mathbf{x}^0 - \mathbf{x}^\mathrm{d}||_2$'


def f_poles(Q, x0):
    """Return the poles of f as (mu, eigenvalue indices, active) tuples, sorted by mu.

    Equal eigenvalues share a pole, which is active when x0 has a normalized component along them.
    """
    x0_norm = Q.to_normalized(x0)
    groups = []
    for k in np.flatnonzero(np.abs(Q.eig) > TOL):
        if groups and abs(Q.eig[k] - Q.eig[groups[-1][0]]) < TOL:
            groups[-1].append(k)
        else:
            groups.append([k])
    poles = [(-1 / Q.eig[ks[0]], ks, bool(np.any(np.abs(x0_norm[ks]) > TOL))) for ks in groups]
    return sorted(poles, key=lambda pole: pole[0])


def sample_mu(Q, poles):
    """Return the mu samples of each segment between consecutive poles, and the range of the mu axis.

    x(mu) only reaches its limit as |mu| -> inf, so the outer segments continue geometrically
    beyond the mu axis; these samples show up in the trajectory but not on the mu axes.
    """
    pole_mus = [mu for mu, _, _ in poles]
    mu_min, mu_max = min(pole_mus + [0]) - MU_MARGIN, max(pole_mus + [0]) + MU_MARGIN
    tail = np.geomspace(MU_MARGIN / N_PER_BRANCH, 1e4 / np.abs(Q.eig[np.abs(Q.eig) > TOL]).min(), N_TAIL)
    starts = [mu_min] + [mu + EPS for mu in pole_mus]
    stops = [mu - EPS for mu in pole_mus] + [mu_max]
    segments = [np.linspace(start, stop, N_PER_BRANCH) for start, stop in zip(starts, stops)]
    segments[0] = np.concatenate((mu_min - tail[::-1], segments[0]))
    segments[-1] = np.concatenate((segments[-1], mu_max + tail))
    return segments, (mu_min, mu_max)


def off_pole(fun):
    """Wrap a function of mu from quadproj, which is 0/0 exactly at an inactive pole, where it is continuous."""
    def wrapped(mu):
        with np.errstate(divide='ignore', invalid='ignore'):
            value = fun(mu)
        return fun(np.nextafter(mu, np.inf)) if np.any(np.isnan(value)) else value
    return wrapped


def kkt_roots(Q, x0):
    """Return (mu*, the other roots of f), found from the sign changes of f between samples.

    mu* is the root between the active poles surrounding 0, as in quadproj, or None.
    """
    poles = f_poles(Q, x0)
    segments, _ = sample_mu(Q, poles)
    f = off_pole(Fun(Q, Q.to_normalized(x0)).f)
    # f is continuous across inactive poles, so roots may lie in the gaps left around them.
    branches = [segments[0]]
    for samples, (_, _, is_active) in zip(segments[1:], poles):
        if is_active:
            branches.append(samples)
        else:
            branches[-1] = np.concatenate((branches[-1], samples))
    roots = []
    for samples in branches:
        signs = np.sign([f(mu) for mu in samples])
        roots += [brentq(f, samples[i], samples[i + 1]) for i in np.flatnonzero(signs[:-1] * signs[1:] < 0)]

    active = [mu for mu, _, is_active in poles if is_active]
    e1 = max((mu for mu in active if mu < 0), default=-np.inf)
    e2 = min((mu for mu in active if mu > 0), default=np.inf)
    mu_star = min((mu for mu in roots if e1 < mu < e2), key=abs, default=None)
    return mu_star, [mu for mu in roots if mu != mu_star]


def degenerate_point(Q, x0, k):
    """Return quadproj's degenerate KKT point at mu = -1/eig_k, or None if it does not exist."""
    if abs(Q.eig[k]) < TOL:
        return None
    x0_norm = Q.to_normalized(x0)
    L_k = [i for i in range(Q.dim) if abs(Q.eig[i] - Q.eig[k]) < TOL]
    if np.any(np.abs(x0_norm[L_k]) > TOL) or get_squared_radius_k(Q, x0_norm, L_k, Q.eig[k]) <= 0:
        return None
    return Q.to_non_normalized(get_xd_k(Q, x0_norm, L_k, Q.eig[k]))


def pole_line(Q, x0, k):
    """Return (point, direction) of the line through lim x(mu) at mu = -1/eig_k.

    The limit is taken with the components of the eig_k eigenspace set to 0. When x0 has a
    component in that eigenspace, x(mu) diverges along it and the line is the asymptote;
    otherwise the line is along eigenvector k and contains the degenerate points.
    """
    x0_norm = Q.to_normalized(x0)
    on_pole = np.abs(Q.eig - Q.eig[k]) < TOL
    with np.errstate(divide='ignore', invalid='ignore'):
        u = Fun(Q, x0_norm).x_std(-1 / Q.eig[k])
    u[on_pole] = 0
    direction = np.where(on_pole, x0_norm, 0) if np.any(np.abs(x0_norm[on_pole]) > TOL) else np.eye(Q.dim)[k]
    origin = Q.to_non_normalized(np.zeros(Q.dim))
    return Q.to_non_normalized(u), Q.to_non_normalized(direction) - origin


def auto_ylim(mu, values, branches, xlim, extra=(), floor=None):
    """Return y-limits framing the finite features of a curve plotted over xlim.

    The features are the values at both ends of xlim, the local extrema of each branch and
    `extra` (markers, reference lines). Divergences at the poles are not features: the curve
    runs off the axes there.
    """
    visible = (mu >= xlim[0]) & (mu <= xlim[1])
    features = [values[visible][0], values[visible][-1], *extra]
    for branch in branches:
        v = values[branch[visible[branch]]]
        slopes = np.diff(v)
        features.extend(v[1:-1][slopes[:-1] * slopes[1:] < 0])
    low, high = min(features), max(features)
    padding = Y_PADDING * (high - low) if high > low else 1
    low -= padding
    if floor is not None:
        low = max(low, floor)
    return low, high + padding


def square_window(points):
    """Return (xlim, ylim) of a square centred on the points and WINDOW_SCALE times as wide (at least 2)."""
    points = np.asarray(points)
    low, high = points.min(axis=0), points.max(axis=0)
    half = max(WINDOW_SCALE * (high - low).max() / 2, 1)
    centre = (low + high) / 2
    return (centre[0] - half, centre[0] + half), (centre[1] - half, centre[1] + half)


def draw_quadric(ax, Q, xlim, ylim):
    if Q.is_single_plane:  # psi >= 0 vanishes on a line, without sign change for the contour
        k = np.flatnonzero(np.abs(Q.eig) <= TOL)[0]
        ax.axline(Q.d, Q.d + Q.V[:, k], color=QUADRIC_COLOR, zorder=1)
    elif Q.is_conical and Q.n == 0:  # psi >= 0 vanishes at a single point
        ax.scatter(*Q.d, color=QUADRIC_COLOR, zorder=1)
    else:
        X, Y = np.meshgrid(np.linspace(*xlim, 801), np.linspace(*ylim, 801))
        P = np.stack((X, Y), axis=-1)
        psi = np.einsum('...i,ij,...j->...', P, Q.A, P) + P @ Q.b + Q.c
        ax.contour(X, Y, psi, levels=[0], colors=QUADRIC_COLOR, zorder=1)


def colored_line(ax, x, y, idx, cmap, norm, n_dotted=0):
    """Draw y(x) colored by idx, with the first and last n_dotted points joined by dotted segments."""
    points = np.column_stack((x, y))
    n = len(points)
    if n_dotted == 0:
        parts = [(0, n, 'solid')]
    else:
        parts = [(0, n_dotted, 'dotted'),
                 (n_dotted - 1, n - n_dotted + 1, 'solid'),
                 (n - n_dotted, n, 'dotted')]
    for start, stop, linestyle in parts:
        p = points[start:stop]
        lc = LineCollection(np.stack((p[:-1], p[1:]), axis=1), cmap=cmap, norm=norm, linestyle=linestyle)
        lc.set_array(idx[start:stop - 1])
        ax.add_collection(lc)


def plot_trajectories(Q, x0):
    if Q.dim != 2 or Q.is_empty:
        raise ValueError('only non-empty 2D quadrics can be plotted')
    F = Fun(Q, Q.to_normalized(x0))
    x, f, df = off_pole(F.x_not_std), off_pole(F.f), off_pole(F.d_f)
    poles = f_poles(Q, x0)
    pole_mus = np.array([mu for mu, _, _ in poles])
    pole_active = np.array([is_active for _, _, is_active in poles], dtype=bool)
    degenerate = not pole_active.all()

    segments, xlim = sample_mu(Q, poles)
    mu = np.concatenate(segments)
    idx = np.arange(len(mu))
    pole_indices = np.cumsum([len(samples) for samples in segments[:-1]], dtype=int)
    # A pole where x0 has a zero normalized component vanishes and joins its neighbouring branches.
    branches = np.split(idx, pole_indices[pole_active])

    cmap = colormaps[COLORMAP]
    norm = Normalize(0, len(mu) - 1)
    pole_labels = [rf'$\frac{{-1}}{{\lambda_{ks[0] + 1}}}$' for _, ks, _ in poles]
    ticks = [xlim[0], *pole_mus, 0, xlim[1]]
    tick_labels = [r'$\mu \to -\infty$', *pole_labels, r'$0$', r'$\mu \to \infty$']

    mu_star, other_roots = kkt_roots(Q, x0)
    markers = ([] if mu_star is None else [(mu_star, STAR_STYLE)]) + [(root, ROOT_STYLE) for root in other_roots]

    # (point, direction, degenerate point or None) for each pole cancelled by a zero component of x0
    inactive = [(*pole_line(Q, x0, ks[0]), degenerate_point(Q, x0, ks[0]))
                for _, ks, is_active in poles if not is_active]
    xds = [xd for _, _, xd in inactive if xd is not None]
    if Q.is_conical and not Q.is_single_plane:  # the apex is a KKT point
        xds.append(Q.d)
    xd_dists = [np.linalg.norm(xd - x0) for xd in xds]

    fig, axs = plt.subplots(2, 2, figsize=(10, 7))
    ax_x, ax_f, ax_dist, ax_df = axs.flat

    # (axis, function, label, dotted points of the middle branches, extra values to show, lower bound)
    panels = [
        (ax_f, f, r'$f(\mu)$', 5, [0], None),
        (ax_dist, lambda m: F.dist(x(m)), r'$||\mathbf{x}^0 - \mathbf{x}(\mu)||_2$', 6,
         [dist + 0.3 for dist in xd_dists], 0),  # green line and its label above it
        (ax_df, df, r'$f^\prime(\mu)$', 50, [], None),
    ]
    for ax, fun, ylabel, n_dotted_middle, extra, floor in panels:
        values = np.array([fun(m) for m in mu])
        for i, branch in enumerate(branches):
            if degenerate:
                n_dotted = 0
            else:
                n_dotted = n_dotted_middle if 0 < i < len(branches) - 1 else N_DOTTED_OUTER
            colored_line(ax, mu[branch], values[branch], branch, cmap, norm, n_dotted)
        for pole in pole_mus[pole_active]:
            ax.axvline(pole, **POLE_STYLE)
        for root, style in markers:
            ax.scatter(root, fun(root), **style)
        marker_values = [fun(root) for root, _ in markers if xlim[0] <= root <= xlim[1]]
        ylim = auto_ylim(mu, values, branches, xlim, extra + marker_values + [fun(0)], floor)  # x(0) = x0
        ax.set(xlabel=r'$\mu$', ylabel=ylabel, xlim=xlim, ylim=ylim)
        ax.set_xticks(ticks, tick_labels)
    ax_f.axhline(0, color='red', alpha=0.5, zorder=1)
    for ax in (ax_f, ax_df):
        ax.yaxis.set_ticks_position('right')

    xs = np.array([x(m) for m in mu])
    kkt_points = [x(root) for root, _ in markers] + xds
    has_centre = not (Q.is_parabolic or Q.is_cylindrical)  # otherwise the centre is not unique
    window_points = [x0, *kkt_points]
    if has_centre:
        window_points.append(Q.center)
    if Q.is_central:  # the vertices, so that the whole shape (e.g. both arcs of a hyperbola) is shown
        window_points += [Q.to_non_normalized(sign * np.eye(Q.dim)[k] / np.sqrt(Q.eig[k]))
                          for k in np.flatnonzero(Q.eig > TOL) for sign in (-1, 1)]
    window_xlim, window_ylim = square_window(window_points)

    draw_quadric(ax_x, Q, window_xlim, window_ylim)
    if has_centre:
        ax_x.scatter(*Q.center, color=QUADRIC_COLOR)
    ax_x.scatter(xs[:, 0], xs[:, 1], c=idx, cmap=cmap, norm=norm, s=4)
    ax_x.scatter(*x0, color='red', marker='x')
    for (_, style), point in zip(markers, kkt_points):
        ax_x.scatter(*point, **style)

    if not degenerate:
        for pole_index, label in zip(pole_indices, pole_labels):
            for i in (pole_index - 1, pole_index):
                ax_x.text(*xs[i], label, **TEXT_STYLE)
        for _, ks, _ in poles:
            point, direction = pole_line(Q, x0, ks[0])
            ax_x.axline(point, point + direction, linestyle='dotted', color='grey')

    for point, direction, xd in inactive:
        if xd is None:
            ax_x.plot(*np.column_stack((point - direction, point + direction)), color='grey', zorder=0)
        else:
            ax_x.plot(*np.column_stack((point, xd)), color='grey', zorder=0)
    for xd, dist in zip(xds, xd_dists):
        ax_x.scatter(*xd, color='green', zorder=2)
        ax_dist.axhline(dist + 0.1, color='green')
        ax_dist.text(xlim[1] + (pole_mus[-1] - xlim[1]) * 2 / 3, dist + 0.2, DIST_D_LABEL, color='green')

    legend = [Line2D([], [], color=QUADRIC_COLOR, lw=2, label=r'$\mathcal{Q}(\Psi)$')]
    if mu_star is not None:
        legend.append(Line2D([], [], linestyle='none', marker='v', color='purple', markersize=6,
                             label=r'$\mathbf{x}(\mu^*)$'))
    if other_roots:
        legend.append(Line2D([], [], linestyle='none', marker='s', color='purple', markersize=6,
                             label=r'$\mathbf{x}(\mu)$, $f(\mu) = 0$'))
    legend.append(Line2D([], [], linestyle='none', marker='x', color='red', markersize=8,
                         label=r'$\mathbf{x}^0 = \mathbf{x}(0)$'))
    if xds:
        legend.append(Line2D([], [], linestyle='none', marker='o', color='green', markersize=8,
                             label=r'$\mathbf{x}^\mathrm{d}$'))
    ax_x.legend(handles=legend, loc='best', fontsize=9)
    ax_x.set(xlabel=r'$x_1$', ylabel=r'$x_2$', xlim=window_xlim, ylim=window_ylim, aspect='equal')

    fig.tight_layout(pad=2.1)
    colorbar = fig.colorbar(ScalarMappable(norm, cmap), ax=ax_x, ticks=[0, len(mu) - 1])
    colorbar.ax.set_yticklabels([r'$\mu \to -\infty$', r'$\mu \to \infty$'])
    return fig


def use_latex_fonts():
    """Typeset text with LaTeX when available, otherwise with matplotlib's Computer Modern."""
    plt.rcParams.update({'font.family': 'serif', 'font.size': 12})
    if shutil.which('latex'):
        plt.rcParams.update({'text.usetex': True, 'text.latex.preamble': r'\usepackage{amsmath}'})
    else:
        plt.rcParams.update({'mathtext.fontset': 'cm', 'font.serif': ['cmr10'],
                             'axes.formatter.use_mathtext': True})


def main():
    use_latex_fonts()
    parser = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    parser.add_argument('cases', nargs='*', metavar='CASE',
                        help=f'cases to plot, among {", ".join(CASES)} (default: all)')
    parser.add_argument('--output-dir', type=Path, default=Path('figures'))
    parser.add_argument('--show', action='store_true')
    args = parser.parse_args()
    unknown = set(args.cases) - set(CASES)
    if unknown:
        parser.error(f'unknown cases: {", ".join(sorted(unknown))}')

    args.output_dir.mkdir(parents=True, exist_ok=True)
    for name in args.cases or CASES:
        case = CASES[name]
        Q = Quadric(np.asarray(case['A'], dtype=float), np.asarray(case.get('b', np.zeros(2)), dtype=float),
                    case.get('c', -1.0))
        fig = plot_trajectories(Q, np.array(case['x0'], dtype=float))
        for extension in ('.png', '.pdf'):
            fig.savefig(args.output_dir / (name + extension))
        print('wrote', args.output_dir / name)
    if args.show:
        plt.show()


if __name__ == '__main__':
    main()