import numpy as np
from scipy.optimize import linprog
import matplotlib.pyplot as plt

# -------- 1. Discretize μ and ν --------

# μ: uniform on segment x = 0, y in [-1, 1], total mass 1
m = 101
y_mu = np.linspace(-1.0, 1.0, m)
x_mu = np.zeros_like(y_mu)
mu_pts = np.stack([x_mu, y_mu], axis=1)
mu_w = np.ones(m) / m  # weights sum to 1

# ν: 1/2 at (1, 0) + 1/2 uniform on segment x = -1, y in [-2, 2]
n_line = 201
y_nu_line = np.linspace(-2.0, 2.0, n_line)
x_nu_line = -np.ones_like(y_nu_line)
nu_line_pts = np.stack([x_nu_line, y_nu_line], axis=1)
nu_line_w = (0.5 / n_line) * np.ones(n_line)  # weights sum to 1/2

nu_pts = np.vstack([np.array([[1.0, 0.0]]), nu_line_pts])
nu_w = np.concatenate([[0.5], nu_line_w])  # total mass 1
n = nu_pts.shape[0]

# -------- 2. Dual LP for W1 with cost ||x - y|| --------

# Variables: (u_0, ..., u_{m-1}, v_0, ..., v_{n-1})
# Maximize sum_i μ_i u_i + sum_j ν_j v_j
# subject to u_i + v_j <= ||x_i - y_j|| for all i, j
# linprog minimizes, so minimize -objective

c = np.concatenate([-mu_w, -nu_w])  # coefficients for [u, v]

rows = []
rhs = []
for i in range(m):
    for j in range(n):
        row = np.zeros(m + n)
        row[i] = 1.0            # u_i
        row[m + j] = 1.0        # v_j
        cost_ij = np.linalg.norm(mu_pts[i] - nu_pts[j])
        rows.append(row)
        rhs.append(cost_ij)

A_ub = np.array(rows)
b_ub = np.array(rhs)

bounds = [(-np.inf, np.inf)] * (m + n)  # u, v free

res = linprog(c, A_ub=A_ub, b_ub=b_ub, bounds=bounds, method="highs")
if not res.success:
    raise RuntimeError(f"LP failed: {res.message}")

x_opt = res.x
u = x_opt[:m]
v = x_opt[m:]

# -------- 3. Build 1-Lipschitz Kantorovich potential via McShane --------

# Support S = supp(μ) ∪ supp(ν), with f defined by dual variables:
#   f(x_i) = u_i on μ-points
#   f(y_j) = -v_j on ν-points
S_pts = np.vstack([mu_pts, nu_pts])
f_S = np.concatenate([u, -v])

def phi_on_grid(xs, ys):
    X, Y = np.meshgrid(xs, ys)
    grid = np.stack([X, Y], axis=-1)  # shape (ny, nx, 2)
    phi = np.full(X.shape, np.inf)
    for k, z in enumerate(S_pts):
        dx = grid[..., 0] - z[0]
        dy = grid[..., 1] - z[1]
        dist = np.sqrt(dx * dx + dy * dy)
        phi = np.minimum(phi, f_S[k] + dist)
    return X, Y, phi

# -------- 4. Evaluate and plot on y >= 0 half-plane, higher resolution --------

xs = np.linspace(-2.0, 2.0, 241)
ys = np.linspace(0.0, 3.0, 181)
X, Y, phi = phi_on_grid(xs, ys)

plt.figure(figsize=(12, 10))
cs = plt.contour(X, Y, phi, levels=40)
plt.clabel(cs, inline=True, fontsize=7)

# overlay μ and ν (only y >= 0 shown)
mask_mu = mu_pts[:, 1] >= 0.0
mask_nu = nu_pts[:, 1] >= 0.0
plt.plot(mu_pts[mask_mu, 0], mu_pts[mask_mu, 1], 'k.', markersize=3, label='mu')
plt.plot(nu_pts[mask_nu, 0], nu_pts[mask_nu, 1], 'r.', markersize=3, label='nu (y>=0)')

plt.axhline(0.0, linewidth=0.5)
plt.axvline(0.0, linewidth=0.5)
plt.gca().set_aspect('equal', 'box')
plt.title("Kantorovich potential φ(x), y ≥ 0")
plt.xlabel("x")
plt.ylabel("y")
plt.legend(loc='upper right', fontsize=7)
plt.tight_layout()
plt.savefig("Kantorovich potential.svg")
plt.show()