import argparse

import jax
import optax
from jax import numpy as jnp, random
from matplotlib import collections, pyplot as plt, rcParams


def parse_args():
    p = argparse.ArgumentParser()
    p.add_argument("--seed", type=int, default=0)
    p.add_argument("--points", type=int, default=200)
    p.add_argument("--power", type=float, default=2.0)
    p.add_argument("--markersize", type=float, default=4.0)
    p.add_argument("--linewidth", type=float, default=1.0)
    return p.parse_args()


def get_wasserstein_distance(x, y, power=2):
    displacements = x[:, None] - y[None, :]
    pow_distances = (jnp.abs(displacements) ** power).sum(-1)
    i, j = optax.assignment.hungarian_algorithm(pow_distances)
    distance = pow_distances[i, j].mean() ** (1 / power)
    return distance, (i, j)


def estimate_wasserstein_distance(key, sample_fn_1, sample_fn_2, n_samples):
    keys = random.split(key)
    x = jax.vmap(sample_fn_1)(random.split(keys[0], n_samples))
    y = jax.vmap(sample_fn_2)(random.split(keys[1], n_samples))
    distance, (i, j) = get_wasserstein_distance(x, y)
    return distance


def main():
    args = parse_args()

    key = random.key(args.seed)

    keys = random.split(key)
    x = random.normal(keys[0], (args.points, 2))
    y = random.normal(keys[1], (args.points, 2)) + jnp.array([0.2, 0.0])

    distance, (i, j) = get_wasserstein_distance(x, y, args.power)

    fig, ax = plt.subplots(constrained_layout=True)

    ax.scatter(*x.T, s=args.markersize**2, label="facility", edgecolor="none")
    ax.scatter(*y.T, s=args.markersize**2, label="client", edgecolor="none")

    data = jnp.stack((x[i], y[j]), 1)
    lc = collections.LineCollection(
        list(data),
        linewidth=args.linewidth,
        color="lightgrey",
        zorder=-10,
        label="assignment",
    )
    ax.add_collection(lc)

    ax.set(title=f"Wasserstein {args.power}-distance: {distance:g}")

    ax.legend()

    rcParams["savefig.dpi"] = 300
    plt.show()


if __name__ == "__main__":
    main()