"""
Pure-strategy saddle point visualizer for a 5 x 5 two-player zero-sum
(matrix) game, using a hardcoded, monotone payoff matrix so the surface
is a genuine smooth saddle shape (like a Pringle chip / hyperbolic
paraboloid) rather than a jagged random surface: moving away from the
saddle point in any direction, the payoff changes smoothly and
monotonically, never reversing direction.

Definition used (matches the course definition of a saddle point):
    A cell (i*, j*) of the payoff matrix A (player 1's payoffs; player 2's
    payoffs are -A) is a saddle point if

        A[i*, j*] >= A[i, j*]   for every row i   -- i.e. it is the
                                    COLUMN MAXIMUM (best player 1 can do
                                    against player 2's strategy j*)

        A[i*, j*] <= A[i*, j]   for every column j -- i.e. it is the
                                    ROW MINIMUM (best player 2 can hold
                                    player 1 to, given player 1 plays i*)

Equivalently, maxmin = min_j max_i A[i,j] = max_i min_j A[i,j] = minmax,
and the saddle point is exactly the (row, column) pair achieving both.

The hardcoded matrix below is A[i,j] = (j - j*)^2 - (i - i*)^2, centered
at the middle strategy (i*, j*) = (3, 3) for both players (1-indexed).
Along the saddle row, payoff is (j-j*)^2: it strictly increases as j
moves away from j* in either direction -- a smooth row MINIMUM at j*.
Along the saddle column, payoff is -(i-i*)^2: it strictly decreases as
i moves away from i* in either direction -- a smooth column MAXIMUM at
i*. That is exactly the saddle shape, hardcoded on purpose.

Run this as a script with an interactive matplotlib backend (TkAgg,
QtAgg, or a Jupyter widget backend) to rotate the 3D plot.
"""

import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D  # noqa: F401 (registers 3d projection)

# Hardcoded 5x5 payoff matrix (player 1's payoffs; player 2's are -A).
# Rows = player 1's strategies (1..5), columns = player 2's strategies
# (1..5). Built from A[i,j] = (j-3)^2 - (i-3)^2, so it is monotone and
# smooth away from the saddle point at (row 3, col 3) in every direction.
A = np.array([
    [ 0, -3, -4, -3,  0],
    [ 3,  0, -1,  0,  3],
    [ 4,  1,  0,  1,  4],
    [ 3,  0, -1,  0,  3],
    [ 0, -3, -4, -3,  0],
], dtype=float)


def find_pure_saddle_points(A):
    """
    Return a list of (i, j) cells that are simultaneously a row minimum
    and a column maximum, i.e. all pure-strategy saddle points of A.
    """
    row_min = A.min(axis=1, keepdims=True)   # player 2's best reply value, per row
    col_max = A.max(axis=0, keepdims=True)   # player 1's best reply value, per column

    is_row_min = A == row_min   # cell equals the minimum of its row
    is_col_max = A == col_max   # cell equals the maximum of its column

    saddle_mask = is_row_min & is_col_max
    return list(zip(*np.where(saddle_mask))), row_min.flatten(), col_max.flatten()


def plot_pure_saddle(A, title="Pure-strategy saddle point of a monotone zero-sum game"):
    n, m = A.shape
    saddles, row_min, col_max = find_pure_saddle_points(A)

    maxmin = row_min.max()   # player 1's guaranteed value (best worst-case row)
    minmax = col_max.min()   # player 2's guaranteed cap (best worst-case column)

    rows = np.arange(1, n + 1)   # player 1's strategies
    cols = np.arange(1, m + 1)   # player 2's strategies
    J, I = np.meshgrid(cols, rows)   # J = player 2 axis, I = player 1 axis

    fig = plt.figure(figsize=(9, 7))
    ax = fig.add_subplot(111, projection="3d")

    surf = ax.plot_surface(J, I, A, cmap="coolwarm", edgecolor="k",
                            linewidth=0.4, alpha=0.9)

    # Highlight every saddle point (there may be more than one, all tied
    # at the same value maxmin = minmax).
    for (i, j) in saddles:
        ax.scatter([j + 1], [i + 1], [A[i, j]], color="black", s=90,
                   marker="o", depthshade=False)

    if saddles:
        i0, j0 = saddles[0]
        # Trace the saddle row (fixed player-1 strategy i0): this is
        # where the saddle value is the ROW MINIMUM.
        ax.plot(cols, np.full(m, i0 + 1), A[i0, :], color="blue", linewidth=3,
                label=f"Row {i0+1}: saddle is the row minimum (player 2's best reply)")
        # Trace the saddle column (fixed player-2 strategy j0): this is
        # where the saddle value is the COLUMN MAXIMUM.
        ax.plot(np.full(n, j0 + 1), rows, A[:, j0], color="green", linewidth=3,
                label=f"Column {j0+1}: saddle is the column maximum (player 1's best reply)")
        ax.legend(loc="upper left", fontsize=8)

    ax.set_xlabel("Player 2's strategy (column j)")
    ax.set_ylabel("Player 1's strategy (row i)")
    ax.set_zlabel("Player 1's payoff  A[i, j]")
    ax.set_xticks(cols)
    ax.set_yticks(rows)
    ax.set_title(f"{title}\nmaxmin = {maxmin:g} = minmax = {minmax:g}  ->  saddle point(s): "
                 + ", ".join(f"(row {i+1}, col {j+1})" for i, j in saddles))

    fig.colorbar(surf, shrink=0.6, aspect=12, label="Player 1's payoff")
    plt.tight_layout()
    plt.show()


if __name__ == "__main__":
    print("Payoff matrix A (player 1's payoffs, 5x5, hardcoded and monotone):")
    print(A.astype(int))

    saddles, row_min, col_max = find_pure_saddle_points(A)
    print("\nRow minima (player 2's best reply value per row):", row_min.astype(int))
    print("Column maxima (player 1's best reply value per column):", col_max.astype(int))
    print("maxmin =", int(row_min.max()), " minmax =", int(col_max.min()))
    print("Saddle point(s) at (row, col), 0-indexed:", saddles)

    plot_pure_saddle(A, title="A smooth, monotone saddle: 5 strategies each")