#!/usr/bin/env python3

import numpy as np

def main():
    u1 = np.array([[1,2,5],[4,3,4],[5,2,1]]) # Player 1's utility function.
    u2 = np.array([[5,4,1],[2,3,2],[1,4,5]]) # Player 2's utility function.

    u, x1, s2 = stackelberg(u1,u2)

    # Expected values:
    u_ = 3.
    x1_ = np.array([0., 1., 0.])
    s2_ = 1

    np.set_printoptions(precision=2)
    print('\n=== u1 ===')
    print(u1)
    print('\n=== u2 ===')
    print(u2)

    print('\n===============')
    print("\nExpected utility: {:.4f}".format(u_))
    print("    Your utility: {:.4f}".format(u))
    print('\n===============')
    print("\nExpected x1:")
    print(x1_)
    print("\n\nYour x1:")
    print(x1)
    print('\n===============\n')
    print("Expected s2: {}".format(s2_))
    print("    Your s2: {}".format(s2))
    print('\n===============\n')

def stackelberg(u1, u2):
    '''
    Compute the optimal Stackelberg strategy for the given 2-player normal form game.

    Arguments:
        u1: a 2D numpy array representing the utility function for player 1
        u2: a 2D numpy array representing the utility function for player 2

    Return:
        (u,x1,s2) where
            u is the expected utility for player 1 of mixed strategy x1 and s2
            x1 is a 1D numpy array of probabilities representing the mixed strategy of player 1
            s2 is the index (0-based) of the pure strategy for player 2
    '''

    n1, n2 = u1.shape

    # TODO: These are placeholder values that have the correct type and shape.
    #       Replace these with your implementation.
    best_util = -1
    best_x1 = np.zeros(n1)
    best_s2 = -1

    return (best_util, np.array(best_x1), best_s2)

if __name__ == '__main__':
    main()