"""Lesson 10: balanced transport between independent population samples."""
import argparse
import json
from pathlib import Path
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import numpy as np
import ot
from common import make_dataset, prepare_embedding


def population_snapshots(data, Z):
    sessions = np.unique(data["session"][data["train"]])
    def snapshot(sids, target_time):
        result = []
        for sid in sids:
            ids = np.flatnonzero(data["session"] == sid)
            result.append(ids[np.argmin(abs(data["time"][ids] - target_time))])
        return np.asarray(result)
    # Disjoint session sets intentionally remove individual correspondence.
    source_ids = snapshot(sessions[::2], .33)
    target_ids = snapshot(sessions[1::2], .80)
    A, B = Z[source_ids], Z[target_ids]
    return A, B


def run(output_dir: Path) -> dict:
    output_dir.mkdir(parents=True, exist_ok=True)
    data = make_dataset()
    _, Z, _, _ = prepare_embedding(data)
    A, B = population_snapshots(data, Z)
    A, B = A[np.argsort(A[:, 1])], B[np.argsort(B[:, 1])]
    a, b = np.full(len(A), 1 / len(A)), np.full(len(B), 1 / len(B))
    cost = ot.dist(A, B, metric="sqeuclidean")
    scale = float(np.median(cost))
    normalized_cost = cost / scale
    exact_plan = ot.emd(a, b, cost)
    exact_w2 = float(np.sqrt(np.sum(exact_plan * cost)))
    epsilons = [.05, .2, 1.0]
    fig, axes = plt.subplots(2, 3, figsize=(12, 7), constrained_layout=True)
    metrics = {"exact_w2": exact_w2, "cost_scale": scale,
               "source_n": int(len(A)), "target_n": int(len(B))}
    for col, eps in enumerate(epsilons):
        ax = axes[0, col]
        plan = ot.sinkhorn(a, b, normalized_cost, reg=eps,
                           method="sinkhorn_log", numItermax=5000, stopThr=1e-10)
        residual = max(np.max(abs(plan.sum(axis=1) - a)),
                       np.max(abs(plan.sum(axis=0) - b)))
        assert residual < 1e-6, "Sinkhorn did not satisfy marginals"
        transported_cost = float(np.sum(plan * cost))
        metrics[f"epsilon_{eps}_marginal_error"] = float(residual)
        metrics[f"epsilon_{eps}_transport_cost"] = transported_cost
        entropy = float(-np.sum(plan * np.log(np.maximum(plan, 1e-300))))
        metrics[f"epsilon_{eps}_entropy"] = entropy
        # Show two largest links per source; the heatmap includes all entries.
        for i in range(len(A)):
            for j in np.argsort(plan[i])[-2:]:
                ax.plot([A[i, 0], B[j, 0]], [A[i, 1], B[j, 1]], color="gray",
                        alpha=.15 + .65 * float(plan[i, j] / plan.max()), linewidth=.8)
        ax.scatter(*A.T, c="#167b8c", s=15, label="Earlier population")
        ax.scatter(*B.T, c="#dc5a32", s=15, label="Later population")
        ax.set(title=f"epsilon = {eps} (normalized cost)", xlabel="PC 1", ylabel="PC 2")
        heat = axes[1, col].imshow(plan * len(A), aspect="auto",
                                    cmap="viridis", vmin=0, vmax=.12)
        axes[1, col].set(xlabel="Target (ordered by PC 2)",
                         ylabel="Source (ordered by PC 2)",
                         title=f"All conditional masses; entropy {entropy:.2f}")
    axes[0, 0].legend(fontsize=7)
    fig.colorbar(heat, ax=axes[1, :], label="Conditional mass P(target | source)", shrink=.8)
    fig.savefig(output_dir / "lesson10.png", dpi=160)
    plt.close(fig)
    return metrics


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--output", type=Path, default=Path("outputs"))
    args = parser.parse_args()
    print(json.dumps(run(args.output), indent=2, allow_nan=False))