
"""A class for the Lemke Howson algorithm with lexicographical ordering"""
from itertools import cycle

import numpy as np

from nashpy.integer_pivoting import make_tableau, pivot_tableau_lex


def shift_tableau(tableau, shape):
    """
    Shift a tableau to ensure labels of pairs of tableaux coincide
    Parameters
    ----------
        tableau: a numpy array
        shape: a tuple
    Returns
    -------
        tableau: a numpy array
    """
    return np.append(
        np.roll(tableau[:, :-1], shape[0], axis=1),
        np.ones((shape[0], 1)),
        axis=1,
    )


def tableau_to_strategy(tableau, basic_labels, strategy_labels):
    """
    Return a strategy vector from a tableau
    Parameters
    ----------
        tableau: a numpy array
        basic_labels: a set
        strategy_labels: a set
    Returns
    -------
        strategy: a numpy array
    """
    vertex = []
    for column in strategy_labels:
        if column in basic_labels:
            for i, row in enumerate(tableau[:, column]):
                if row != 0:
                    vertex.append(tableau[i, -1] / row)
        else:
            vertex.append(0)
    strategy = np.array(vertex)
    return strategy / sum(strategy)

def lemke_howson_lex(A, B, initial_dropped_label=0):
    """
     Obtain the Nash equilibria using the Lemke Howson algorithm implemented
     using lexicographical integer pivoting. (Able to solve degenerate games)
     1. Start at the artificial equilibrium (which is fully labeled)
     2. Choose an initial label to drop and move in the polytope for which
        the vertex has that label to the edge that does not share that label.
    (This is implemented using integer pivoting and the choice of label
    to drop is implemented using lexicographical ordering)
     3. A label will now be duplicated in the other polytope, drop it in a
        similar way.
     4. Repeat steps 2 and 3 until have Nash Equilibrium.
     Parameters
     ----------
             initial_dropped_label: int
     Returns
     -------
             equilibria: A tuple.
    """

    if np.min(A) <= 0:
        A = A + abs(np.min(A)) + 1
    if np.min(B) <= 0:
        B = B + abs(np.min(B)) + 1

    # build tableaux
    col_tableau = make_tableau(A)
    col_tableau = shift_tableau(col_tableau, A.shape)
    row_tableau = make_tableau(B.transpose())
    full_labels = set(range(sum(A.shape)))

    # slack variables
    row_slack_variables = range(B.shape[0], sum(B.shape))
    col_slack_variables = range(A.shape[0])

    # non-basic variables
    row_non_basic_variables = full_labels - set(row_slack_variables)
    col_non_basic_variables = full_labels - set(col_slack_variables)

    # print(initial_dropped_label)
    if initial_dropped_label in row_non_basic_variables:
        tableaux = cycle(
            (
                (row_tableau, row_slack_variables, row_non_basic_variables),
                (col_tableau, col_slack_variables, col_non_basic_variables),
            )
        )
    else:
        tableaux = cycle(
            (
                (col_tableau, col_slack_variables, col_non_basic_variables),
                (row_tableau, row_slack_variables, row_non_basic_variables),
            )
        )

    # First pivot (to drop a label)
    next_tableau, next_slack_variables, next_non_basic_variables = next(
        tableaux
    )

    entering_label = pivot_tableau_lex(
        next_tableau,
        initial_dropped_label,
        next_slack_variables,
        next_non_basic_variables,
    )

    # keeps track of each tableau's non-basic variables
    next_non_basic_variables.add(entering_label)
    next_non_basic_variables.remove(initial_dropped_label)

    while col_non_basic_variables.union(row_non_basic_variables) != full_labels:
        next_tableau, next_slack_variables, next_non_basic_variables = next(
            tableaux
        )

        # the first label is 'entering' in the sense that it will enter the next
        # tableau's set of basic variables
        just_entered_label = entering_label
        entering_label = pivot_tableau_lex(
            next_tableau,
            entering_label,
            next_slack_variables,
            next_non_basic_variables,
        )

        next_non_basic_variables.add(entering_label)
        next_non_basic_variables.remove(just_entered_label)

    row_strategy = tableau_to_strategy(
        row_tableau,
        full_labels - row_non_basic_variables,
        range(A.shape[0]),
    )

    col_strategy = tableau_to_strategy(
        col_tableau,
        full_labels - col_non_basic_variables,
        range(A.shape[0], sum(A.shape)),
    )

    return row_strategy, col_strategy