# This file was *autogenerated* from the file correlated_polytope_cone.sage
from sage.all_cmdline import *   # import sage library

_sage_const_1 = Integer(1); _sage_const_0 = Integer(0)


def linear_program(N_i, X, print_constraints = False, bounded = True):
    def get_variables(N_i):
        # generate a very long list of variables
        variables = []
        # c is a list of stategies, one entry per player fixiing a stategy for each player
        for c in cartesian_product([range(_sage_const_1 ,n+_sage_const_1 ) for n in N_i]): # c = [strategy of player 1, ... , strategy of player n]
            index = [str(i+_sage_const_1 )+str(c[i]) for i in range(len(c))] # [1StategyOfPlayer1, ..., nStategyOfPlayern]
            variables.append("_".join(["s"+tupl for tupl in index]))# s1StategyOfPlayer1_..._snStategyOfPlayern
        return variables

    def get_S_i(N_i):
        N = len(N_i)
        res = []
        for i in range(N):
            S_i = ["s{i}{j}".format(i = i+_sage_const_1 , j = j) for j in range(_sage_const_1 ,N_i[i]+_sage_const_1 )] #[si1, ..., sid_i]
            res.append(S_i)
        return res

    def remove_duplicates(liste):
        res = []
        for l in liste:
            if not l in res:
                res.append(l)
        return res

    def get_S_mi(N, variables):
        # delete ith strategy
        res = []
        split = [v.split("_") for v in variables]
        for i in range(N):
            if i == _sage_const_0 :
                 delete_i = [s[i+_sage_const_1 :] for s in split]
            else:
                delete_i = [s[:i]+s[i+_sage_const_1 :] for s in split]
            join = ["_".join(s) for s in remove_duplicates(delete_i)]
            res.append(join)
        return res

    def get_variable(i,s_i,s_mi):
        split = s_mi.split("_")
        insert = split[:i-_sage_const_1 ]+[s_i]+split[i-_sage_const_1 :]
        return "_".join(insert)

    N = len(N_i)
    variables = get_variables(N_i)
    S_i = get_S_i(N_i)
    S_mi = get_S_mi(N, variables)
    lp = MixedIntegerLinearProgram()
    p = lp.new_variable(real=True, nonnegative=True)
    if bounded:
        lp.add_constraint(sum([p[s] for s in variables]) == _sage_const_1 )
    if print_constraints:
        print(sum([p[s] for s in variables]) == _sage_const_1 )
    for i in range(_sage_const_1 ,N+_sage_const_1 ):
        for s_i in S_i[i-_sage_const_1 ]:
            for d_i in S_i[i-_sage_const_1 ]:
                if not s_i == d_i:
                    res_sum = _sage_const_0 
                    for s_mi in S_mi[i-_sage_const_1 ]:
                        s = get_variable(i,s_i,s_mi)
                        d = get_variable(i,d_i,s_mi)
                        ui_s = X[variables.index(s)][i-_sage_const_1 ]
                        ui_d = X[variables.index(d)][i-_sage_const_1 ]
                        res_sum += p[s]*(ui_s - ui_d)
                    if print_constraints:
                        print(res_sum >= _sage_const_0 )
                    lp.add_constraint(res_sum >= _sage_const_0 )
    return lp

def correlated_equilibrium_polytope(N_i, X, base_ring = QQ, print_constraints = False):
    lp = linear_program(N_i, X, print_constraints = print_constraints, bounded = True)
    return lp.polyhedron(base_ring = QQ)


def correlated_equilibrium_cone(N_i, X, base_ring = QQ, print_constraints = False):
    lp = linear_program(N_i, X, print_constraints = print_constraints, bounded = False)
    return lp.polyhedron(base_ring = base_ring)