source/gconvex_authors/tools/visualize.py

from __future__ import annotations
"""Evaluation and plotting helpers for mechanism visualization tasks."""

import os
from typing import Optional, Tuple

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import torch
from matplotlib.colors import BoundaryNorm, ListedColormap
from matplotlib.patches import Patch

from config import WRITING_ROOT
from tools.feedback import logger
from tools.utils import loader

# ---------- core save/show helper ----------

def _save_or_show(fig: plt.Figure, filename: str, save_dir: Optional[str] = WRITING_ROOT):
    """
    Save the figure under `save_dir` (defaults to config.WRITING_ROOT), or
    display it interactively if save_dir is None.
    """
    if save_dir is not None:
        os.makedirs(save_dir, exist_ok=True)
        path = os.path.join(save_dir, filename)
        fig.savefig(path, dpi=150, bbox_inches="tight")
        plt.close(fig)
        logger.info(f"[viz] saved {path}")
    else:
        fig.show()

# ---------- core evaluation helpers ----------

@torch.no_grad()
def eval_mech_1d(mechanism: torch.nn.Module, N: int = 801, device: Optional[torch.device] = None):
    """
    Evaluate mechanism on a 1D grid x in [0,1].
    Returns xs, q (N,), t (N,), v (N or None), k (N or None).
    """
    xs = np.linspace(0.0, 1.0, N, dtype=np.float32)
    X = torch.from_numpy(xs[:, None])  # (N,1)

    if device is None:
        try:
            device = next(mechanism.parameters()).device
        except StopIteration:
            device = torch.device("cpu")
    X = X.to(device)

    out = mechanism.compute_mechanism(X, mode="hard")
    q = out["choice"].detach().cpu().numpy().reshape(-1)          # (N,)
    t = out["revenue"].detach().cpu().numpy().reshape(-1)         # (N,)
    v = out.get("v", None)
    k = out.get("kernel", None)
    if v is not None: v = v.detach().cpu().numpy().reshape(-1)
    if k is not None: k = k.detach().cpu().numpy().reshape(-1)
    return xs, q, t, v, k


@torch.no_grad()
def eval_mech_2d(mechanism: torch.nn.Module, N: int = 1000, device: Optional[torch.device] = None):
    """
    Evaluate mechanism on a 2D grid (x1,x2) in [0,1]^2.
    Returns (xs, ys, X, Y, q1, q2, t, v, k)
    """
    xs = np.linspace(0.0, 1.0, N, dtype=np.float32)
    ys = np.linspace(0.0, 1.0, N, dtype=np.float32)
    X, Y = np.meshgrid(xs, ys)
    grid = np.stack([X, Y], axis=-1).reshape(-1, 2)
    T = torch.from_numpy(grid)

    if device is None:
        try:
            device = next(mechanism.parameters()).device
        except StopIteration:
            device = torch.device("cpu")
    T = T.to(device)

    out = mechanism.compute_mechanism(T, mode="hard")
    Q = out["choice"].detach().cpu().numpy().reshape(N, N, 2)     # (N,N,2)
    t = out["revenue"].detach().cpu().numpy().reshape(N, N)       # (N,N)
    v = out.get("v", None)
    k = out.get("kernel", None)
    if v is not None: v = v.detach().cpu().numpy().reshape(N, N)
    if k is not None: k = k.detach().cpu().numpy().reshape(N, N)
    q1, q2 = Q[..., 0], Q[..., 1]
    return xs, ys, X, Y, q1, q2, t, v, k

# ---------- math helpers ----------

def bundle_probs_from_marginals(q1: np.ndarray, q2: np.ndarray) -> Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
    """
    Independent implementation map: (q1,q2) -> probs over {00,10,01,11}.
    """
    p11 = q1 * q2
    p10 = q1 * (1.0 - q2)
    p01 = (1.0 - q1) * q2
    p00 = (1.0 - q1) * (1.0 - q2)
    return p00, p10, p01, p11

def estimate_cutoff(xs: np.ndarray, q: np.ndarray, level: float = 0.5) -> float:
    """
    First x where q(x) crosses 'level' (linear interpolation). Returns nan if no crossing.
    """
    idx = np.where(q >= level)[0]
    if len(idx) == 0: return float("nan")
    k = idx[0]
    if k == 0: return float(xs[0])
    x0, x1 = xs[k-1], xs[k]
    q0, q1 = q[k-1], q[k]
    if q1 == q0: return float(x1)
    a = (level - q0) / (q1 - q0)
    return float(x0 + a * (x1 - x0))

def theory_step_allocation(xs: np.ndarray, reserve: float = 0.5):
    """
    Posted price theory for x ~ U[0,1]: q*(x)=1{x >= reserve}, t*(x)=reserve*q*(x).
    """
    q_star = (xs >= reserve).astype(float)
    t_star = reserve * q_star
    return q_star, t_star

# ---------- small plotting primitives ----------

def _heatmap(Z, xs, ys, title, clabel=None, filename=None, save_dir: Optional[str] = WRITING_ROOT):
    fig, ax = plt.subplots(figsize=(6, 6))
    
    im = ax.imshow(Z, origin="lower", extent=[xs[0], xs[-1], ys[0], ys[-1]], aspect="equal")
    ax.set_title(title); ax.set_xlabel("x1"); ax.set_ylabel("x2")
    if clabel: fig.colorbar(im, ax=ax, label=clabel, shrink=0.73)
    fig.tight_layout()
    if filename is None:
        # derive a safe filename from the title
        filename = title.lower().replace(" ", "_").replace("(", "").replace(")", "") + ".png"
    _save_or_show(fig, filename, save_dir)

def _quiver(X, Y, U, V, title, step=8, filename=None, save_dir: Optional[str] = WRITING_ROOT):
    fig, ax = plt.subplots(figsize=(6, 6))
    ax.quiver(X[::step, ::step], Y[::step, ::step], U[::step, ::step], V[::step, ::step],
              angles="xy", scale_units="xy", scale=1.0)
    ax.set_title(title); ax.set_xlabel("x1"); ax.set_ylabel("x2")
    ax.set_xlim(0, 1); ax.set_ylim(0, 1); ax.set_aspect("equal", adjustable="box")
    fig.tight_layout()
    if filename is None:
        filename = title.lower().replace(" ", "_") + ".png"
    _save_or_show(fig, filename, save_dir)

# ---------- 1D high-level visualization ----------

def plot_mechanism_1d(mechanism: torch.nn.Module,
                      reserve: float = 0.5,
                      N: int = 1201,
                      show_revenue_curve: bool = True,
                      save_dir: Optional[str] = WRITING_ROOT):
    """
    Plots learned allocation/payment vs. posted-price theory, reports reserve & revenue.
    """
    xs, q, t, v, k = eval_mech_1d(mechanism, N=N)
    q_star, t_star = theory_step_allocation(xs, reserve)
    r_hat = estimate_cutoff(xs, q, level=0.5)

    # allocation
    fig, ax = plt.subplots(figsize=(6, 4.5))
    ax.plot(xs, q, label="learned a(x)", lw=2)
    ax.plot(xs, q_star, "--", label=f"optimal $a^*(x) = 1${{x ≥ {reserve}}}", lw=2)
    if not np.isnan(r_hat):
        ax.axvline(r_hat, color="k", ls=":", label=f"learned price $\hat{{p}}$ ≈ {r_hat:.3f}")
    ax.axvline(reserve, color="gray", ls="--", alpha=0.8, label=f"optimal price $\hat{{p}}$={reserve}")
    ax.set_ylim(-0.05, 1.05); ax.set_xlabel("type x"); ax.set_ylabel("allocation a(x)")
    ax.set_title("Single Good Allocation"); ax.legend(); fig.tight_layout()
    _save_or_show(fig, "1d_allocation.png", save_dir)

    # payment

    fig, ax = plt.subplots(figsize=(6, 6))
    ax.plot(xs, t, label="learned t(x)", lw=2)
    ax.plot(xs, t_star, "--", label="theory t*(x)", lw=2)
    ax.axvline(reserve, color="gray", ls="--", alpha=0.8)
    ax.set_xlabel("type x"); ax.set_ylabel("payment t(x)")
    ax.set_title("1D payment: learned vs theory"); ax.legend(); fig.tight_layout()
    _save_or_show(fig, "1d_payment.png", save_dir)

    # optional revenue curve context
    if show_revenue_curve:
        ps = np.linspace(0, 1, 400)
        R = ps * (1 - ps)
        fig, ax = plt.subplots(figsize=(6, 4))
        ax.plot(ps, R, lw=2, label="Theoretical Profit")
        ax.axvline(reserve, ls="--")
        ax.scatter([reserve], [reserve * (1 - reserve)], zorder=3, color="red", label=f"Posted Price Found")
        ax.set_xlabel("posted price p"); ax.set_ylabel("expected revenue R(p)")
        ax.set_title("Monopoly revenue (Uniform[0,1])"); fig.tight_layout()
        ax.legend()
        _save_or_show(fig, "1d_revenue_curve.png", save_dir)

    # metrics
    rev_learned = float(t.mean())
    rev_theory = float(reserve * (1 - reserve))
    l2_q = float(np.sqrt(np.mean((q - q_star)**2)))
    l2_t = float(np.sqrt(np.mean((t - t_star)**2)))

    logger.info(f"[1D] learned cutoff r̂: {r_hat:.3f}")
    logger.info(f"[1D] expected revenue — learned: {rev_learned:.4f} | theory(opt): {rev_theory:.4f} | gap: {rev_learned - rev_theory:+.4f}")
    logger.info(f"[1D] L2 error: q vs theory = {l2_q:.4f},  t vs theory = {l2_t:.4f}")

# ---------- 2D high-level visualization ----------

def plot_mechanism_2d(mechanism: torch.nn.Module,
                      N: int = 201,
                      tau: float = 1e-3,
                      show_q_heatmaps: bool = True,
                      show_bundle_map: bool = True,
                      show_payment_heatmap: bool = True,
                      show_value_heatmap: bool = True,
                      show_quiver_field: bool = True,
                      quiver_step: int = 10,
                      kernel_label: str = "gross value (kernel)",
                      save_dir: Optional[str] = WRITING_ROOT):
    """
    Visual diagnostics on [0,1]^2:
      - q1, q2 heatmaps
      - implied bundle regions & tie (Maxwell) set from independent implementation
      - payment heatmap
      - optional v(x) and kernel(x,q(x)) heatmaps
      - optional quiver of q(x)
    """
    xs, ys, X, Y, q1, q2, t, v, k = eval_mech_2d(mechanism, N=N)

    if show_q_heatmaps:
        _heatmap(q1, xs, ys, "$a_1(x)$: probability of getting good 1", "$a_1$", filename="a_1_heatmap.png", save_dir=save_dir)
        _heatmap(q2, xs, ys, "$a_2(x)$: probability of getting good 2", "$a_2$", filename="a_2_heatmap.png", save_dir=save_dir)

    if show_bundle_map:
        p00, p10, p01, p11 = bundle_probs_from_marginals(q1, q2)
        stacks = np.stack([p00, p10, p01, p11], axis=-1)  # (N,N,4)
        labels = stacks.argmax(axis=-1)

        # tie curve
        sv = np.sort(stacks, axis=-1)
        gap = sv[..., -1] - sv[..., -2]
        tie = (gap < tau).astype(float)

        names  = {0: "∅ (00)", 1: "{1} (10)", 2: "{2} (01)", 3: "{1,2} (11)"}
        colors = ["#d9d9d9", "#6baed6", "#74c476", "#fd8d3c"]
        cmap   = ListedColormap(colors)
        norm   = BoundaryNorm([-0.5, 0.5, 1.5, 2.5, 3.5], cmap.N)

        fig, ax = plt.subplots(figsize=(7, 7))
        im = ax.imshow(labels, origin="lower", extent=[xs[0], xs[-1], ys[0], ys[-1]],
                       interpolation="nearest", aspect="equal", cmap=cmap, norm=norm)
        # draw boundaries + tie
        ax.contour(X, Y, labels, levels=[0.5, 1.5, 2.5], colors="k", linewidths=1.0)
        if tie.any():
            ax.contour(X, Y, tie, levels=[0.5], colors="k", linewidths=0.8, linestyles="--")
        ax.set_title("Implied bundle regions (indep. implementation)")
        ax.set_xlabel("x1"); ax.set_ylabel("x2")
        patches = [Patch(facecolor=colors[k], edgecolor='none', label=names[k]) for k in range(4)]
        ax.legend(handles=patches, title="Chosen bundle", loc="upper left", frameon=True)
        fig.tight_layout()
        _save_or_show(fig, "bundle_map.png", save_dir)

    if show_payment_heatmap:
        _heatmap(t, xs, ys, "Payment t(x)", "t(x)", filename="payment_heatmap.png", save_dir=save_dir)

    if show_value_heatmap and (v is not None):
        _heatmap(v, xs, ys, "Indirect utility v(x)", "v(x)", filename="indirect_utility_heatmap.png", save_dir=save_dir)

    if (k is not None):
        _heatmap(k, xs, ys, kernel_label, kernel_label, filename="kernel_heatmap.png", save_dir=save_dir)

    if show_quiver_field:
        _quiver(X, Y, q1, q2, "Allocation vector field q(x)", step=quiver_step, filename="allocation_quiver.png", save_dir=save_dir)


def plot_revenue_fit_1d(mechanism: torch.nn.Module,
                        reserve: float | None = None,
                        N_price_grid: int = 400,
                        price_from: str = "cutoff",
                        N_eval: int = 1201,
                        save_dir: Optional[str] = WRITING_ROOT):
    """
    Show how your posted price fits the revenue curve R(p) for x ~ Uniform[0,1].

    Parameters
    ----------
    mechanism : trained mechanism with compute_mechanism(X)
    reserve   : if provided, use this as the posted price p (overrides price_from)
    N_price_grid : resolution for plotting the revenue curve
    price_from : "cutoff" (estimate from q crossing 0.5) or "argmax" (maximizes empirical revenue curve)
    N_eval    : resolution for evaluating the learned mechanism in 1D

    Notes
    -----
    - Theory for Uniform[0,1]: demand at price p is 1 - F(p) = 1 - p, so R(p) = p*(1-p).
    - If reserve is None and price_from == "cutoff", we use the learned cutoff r̂ from q(x).
    - If price_from == "argmax", we pick p that maximizes p*(1-p) on the same grid (≈ 0.5).
    """
    # 1) Evaluate learned mechanism in 1D (for reserve-from-cutoff & reporting)
    xs, q, t, _, _ = eval_mech_1d(mechanism, N=N_eval)
    r_hat = estimate_cutoff(xs, q, level=0.5)

    # 2) Theoretical revenue curve
    ps = np.linspace(0.0, 1.0, N_price_grid)
    R = ps * (1.0 - ps)  # Uniform[0,1]

    # 3) Choose the posted price to mark
    if reserve is not None:
        p_mark = float(reserve)
        mark_label = f"your posted price p = {p_mark:.3f}"
    else:
        if price_from == "argmax":
            j = int(np.argmax(R))
            p_mark = float(ps[j])
            mark_label = f"argmax R(p) ≈ {p_mark:.3f}"
        else:  # "cutoff"
            p_mark = float(r_hat) if not np.isnan(r_hat) else 0.5
            mark_label = "learned price $\hat{p}$ ≈ "+ f"{p_mark:.3f}"

    R_mark = float(p_mark * (1.0 - p_mark))

    # 4) Plot & save
    fig, ax = plt.subplots(figsize=(6, 4.5))
    ax.plot(ps, R, lw=2, label="$R(p)=p(1-p)$")
    ax.axvline(0.5, ls="--", color="gray", alpha=0.8, label="optimal price $\hat{p}=0.5$")
    ax.scatter([p_mark], [R_mark], s=50, zorder=3, label=f"{mark_label}", color="red")#\nR(p)={R_mark:.3f}
    ax.set_xlabel("posted price p"); ax.set_ylabel("expected revenue R(p)")
    ax.set_title("Posted Price and Single Good")
    ax.legend()
    fig.tight_layout()
    _save_or_show(fig, "revenue_fit.png", save_dir)

    # 5) Print a quick summary vs learned mechanism revenue
    rev_learned = float(t.mean())  # uniform types on [0,1]
    rev_theory_opt = 0.25
    logger.info(f"[Revenue fit] marked price p = {p_mark:.3f}, R(p) = {R_mark:.4f}")
    logger.info(f"[Revenue fit] learned mechanism revenue (mean t): {rev_learned:.4f}")
    logger.info(f"[Revenue fit] theory optimum: {rev_theory_opt:.4f}, gap to marked: {R_mark - rev_theory_opt:+.4f}")
    if not np.isnan(r_hat):
        logger.info(f"[Revenue fit] learned cutoff r̂ from q(x): {r_hat:.3f}")


def plot_profit(sample, dim_to_mechs):
    """
    Plot profit curves for the given mechanism.
    """
    n_items = []
    P = []
    for dim, (mech, mech_date) in sorted(dim_to_mechs.items()):
        n_items.append(dim)
        p = mech_date["profits"].mean().item()
        logger.debug(f"{dim} {p}")
        P.append(p/dim)

    fig, ax = plt.subplots(figsize=(6, 6))
    ax.plot(n_items, P, lw=2, label="profits per item")
    ax.set_xlabel("number of items"); ax.set_ylabel("profit")
    ax.set_title("Profit curve")
    ax.legend()
    fig.tight_layout()
    _save_or_show(fig, "profit_per_item_curve.png", save_dir)


def make_table(dim_to_mechs, sample, save_dir):
    """
    Create a profit table for the given mechanisms and sample.
    """
    dims = []
    mp = []
    ms = []
    mv = []
    for dim, (mech, mech_data) in sorted(dim_to_mechs.items()):
        profit_per_item = mech_data["profits"].mean().item() / dim
        surplus_per_item = mech_data["kernel"].mean().item() / dim
        v_per_item = mech_data["v"].mean().item() / dim
        dims.append(int(dim))
        mp.append(profit_per_item)
        ms.append(surplus_per_item)
        mv.append(v_per_item)
        
    df = pd.DataFrame({
        "Number of Items": dims,
        "Mean Profit per Item": mp,
        "Mean Surplus per Item": ms,
        "Mean Utility per Item": mv
    })
    
    fig, ax = plt.subplots(figsize=(7, 2))
    ax.axis('off')
    table = ax.table(cellText=df.values, colLabels=df.columns, loc='center', cellLoc='center')
    table.auto_set_font_size(False)
    table.set_fontsize(10)
    table.auto_set_column_width(col=list(range(len(df.columns))))
    _save_or_show(fig, "mech_table.png", save_dir)

    latex_table = df.to_latex(index=False,
                              caption="Optimized Mechanisms", label="tab:results",
                              float_format="%.3f",
                              column_format="|l|r|r|")
    with open(f"{save_dir}/table.tex", "w") as f:
        f.write(latex_table)


if __name__=="__main__":    
    from config import WRITING_ROOT
    data = loader()
    for tag, (sample, models) in data.items():
        logger.info(f"working on {tag}...")
        save_dir = f"{WRITING_ROOT}/plots/{tag}"
        logger.info(f"saving to {save_dir}")
        dim_to_mechs = {}
        for m in models:
            dim = m.full_Y().shape[2]
            dim_to_mechs[dim] = (m, m.compute_mechanism(sample[:, :dim], mode="hard"))
        if 1 in dim_to_mechs:
            plot_revenue_fit_1d(dim_to_mechs[1][0], save_dir=save_dir)
            plot_mechanism_1d(dim_to_mechs[1][0], reserve=0.5, N=1000, show_revenue_curve=True, save_dir=save_dir)
        if 2 in dim_to_mechs:
            plot_mechanism_2d(dim_to_mechs[2][0], N=1000, tau=1e-3,
                      show_q_heatmaps=True,
                      show_bundle_map=True,
                      show_payment_heatmap=True,
                      show_value_heatmap=False,
                      show_quiver_field=False,
                      quiver_step=10,
                      save_dir=save_dir)            
        plot_profit(sample, dim_to_mechs)
        make_table(dim_to_mechs, sample, save_dir)
        # Collect profit per item for each mechanism


def visualize_transport(
    x, y, model, save_dir=WRITING_ROOT,
    n_arrows=200, figsize=(7, 7)
):
    """
    Visualize learned Monge map T(x)=∇f(x) with straight-line transport arrows.
    Each arrow is a single line colored from red (start) → blue (end).
    """
    assert x.shape[1] == 2, "Visualization only implemented for d=2."

    import matplotlib.pyplot as plt
    import matplotlib.cm as cm
    from matplotlib.collections import LineCollection
    import numpy as np
    import os

    # -----------------------------------------------------
    # Prepare directory + filename
    # -----------------------------------------------------
    os.makedirs(save_dir, exist_ok=True)
    model_name = model.__class__.__name__
    fig_path = os.path.join(save_dir, f"OTplot_{model_name}.png")

    # -----------------------------------------------------
    # Prepare data
    # -----------------------------------------------------
    x_cpu = x.detach().cpu()
    y_cpu = y.detach().cpu()
    Tx = model.transport_X_to_Y(x_cpu).detach().cpu()

    x_np = x_cpu.detach().numpy()
    y_np = y_cpu.detach().numpy()
    Tx_np = Tx.detach().numpy()

    # -----------------------------------------------------
    # Subsample for arrows
    # -----------------------------------------------------
    n = min(n_arrows, x_np.shape[0])
    idx = np.random.choice(x_np.shape[0], n, replace=False)
    xs = x_np[idx]
    Txs = Tx_np[idx]

    # -----------------------------------------------------
    # Build the transport lines
    # -----------------------------------------------------
    # Each line is shape (2,2): [[x1,y1],[x2,y2]]
    lines = np.stack([xs, Txs], axis=1)  # shape (n,2,2)

    # Create gradient colors: 0=red (start), 1=blue (end)
    cmap = cm.get_cmap("coolwarm")  # red→white→blue
    start_color = cmap(0.0)         # red side of coolwarm
    end_color   = cmap(1.0)         # blue side

    # Build a Nx2 array of colors: start red, end blue
    # LineCollection expects one color *per line*, but we can fade using alpha
    # So instead: fade based on distance
    dists = np.sqrt(np.sum((Txs - xs)**2, axis=1))
    if dists.max() > 0:
        d_norm = (dists - dists.min()) / (dists.max() - dists.min())
    else:
        d_norm = np.zeros_like(dists)

    # Blend: color = (1-t)*red + t*blue
    colors = (1 - d_norm)[:, None] * start_color + d_norm[:, None] * end_color
    colors[:, 3] = 0.8  # set alpha

    # -----------------------------------------------------
    # Plot
    # -----------------------------------------------------
    plt.figure(figsize=figsize)

    # Softer scatter for X, Y, and T(x)
    plt.scatter(x_np[:, 0], x_np[:, 1], s=10, color="#b33939", alpha=0.45, label="source $x$")
    plt.scatter(y_np[:, 0], y_np[:, 1], s=10, color="#3b6ea8", alpha=0.45, label="target $y$")
    plt.scatter(Tx_np[:, 0], Tx_np[:, 1], s=12, color="#f5b041", alpha=0.55, label="$T(x)$")

    ax = plt.gca()
    # lc = LineCollection(lines, colors=colors, linewidth=1.6)
    # ax.add_collection(lc)

    plt.legend()
    plt.title("Dynamic Transport: red → blue flow")
    plt.xlabel("$x_1$")
    plt.ylabel("$x_2$")
    plt.axis("equal")
    plt.tight_layout()
    plt.savefig(fig_path, dpi=240)
    plt.close()

    logger.info(f"[✓] Saved dynamic transport figure to: {fig_path}")
    return fig_path