#module for Matrix_Single


from misc import *


#following class simulates a single-stage matrix game
class Matrix_Single:
    
    def __init__(self, player_list, strat_mat, payoff_mat):
        #argument players expected to be a list of string
        #with 1 to 1 correspondence with strat_mat such that
        #player[0]'s strats are strat_mat[0]
        #argument strat_mat expected to be a mat of Strat
        #argument payoff_mat expected to be a mat of float
        #with a len equal to that of the product of the len of
        #every list of Strategy within strat_mat
        
        self.check_validity(player_list, strat_mat, payoff_mat)
        self.player_list = player_list
        self.strat_mat = strat_mat
        self.strats_combined_list = combine(strat_mat)    
        self.payoff_mat = payoff_mat
        self.matrix_dict = self.make_matrix_dict()
        self.br_dict = self.make_br_dict()
        self.wd_dict = self.make_wd_dict()
        self.sd_dict = self.make_sd_dict()
        self.player_symbol_dict = self.make_player_symbol_dict()
        self.symbol_dict = self.make_strat_symbol_dict()
        self.pure_ne_list = self.make_pure_ne_list()
        self.mixed_ne_list = self.make_mixed_ne_list()
        self.logical_payoff_mat = self.make_logical_payoff_mat()

    #following function checks if the constructor arguments are comprehensible
    def check_validity(self, player_list, strat_mat, payoff_mat):  #FUNCTION DOES NOT THROW EXCEPTIONS YET
        #arguments here expected to be the same as __init__ arguments of the same name
        
        if len(player_list) < 1 or len(strat_mat) < 1 or len(payoff_mat) < 1 or \
           not len(player_list) == len(strat_mat):
            print("ERROR")  #PLACEHOLDER CODE. THROW AN EXCEPTION HERE

    #following code maps combined strats to payoff
    def make_matrix_dict(self):
        
        matrix_dict = {strat_combo[1] : Payoff_Dict(self.player_list, self.payoff_mat[strat_combo[0]]) \
                       for strat_combo in enumerate(self.strats_combined_list)}
        
        #return value is a dict of the form list of Strategy: list of float
        return matrix_dict

    #following function finds BR sub i(s sub -i), or best s sub i for the player to
    #choose, given player i and s sub -i
    def make_br_list(self, player, other_strats):
        #argument player expected to be a string in self.player_list
        #argument other_strats expected to be list of Stratgey
        #with one strategy for every player but i: s sub -i
        
        br_list = []
        br_payoff = None
        player_offset = self.player_list.index(player)
    
        #following code creates a list of all possible s with every s sub i
        #(i is player) and s sub -i being other_strats.
        #naturally, the len of context_list is the number of strats player has
        context_list = []
        for strats_combined in self.strats_combined_list:
            if other_strats == strats_combined[:player_offset] + \
                               strats_combined[player_offset + 1:]:
                context_list.append(strats_combined)
        #end
                
        #following code finds the highest U(s), where s is generated by previous code block,
        #then goes through context_list and finds the U(s) that matches the highest U(s)
        #found U(s) are s with BR sub i(s sub -i) and have their s sub i appended to br_list
        #because there may be multiple highest U(s), this step is necessary
        for strats_combined in context_list:
            if br_payoff == None or \
               br_payoff < self.matrix_dict[strats_combined].get_payoff(player):
                br_payoff = self.matrix_dict[strats_combined].get_payoff(player)
        for strats_combined in context_list:
            if br_payoff == self.matrix_dict[strats_combined].get_payoff(player):
                br_list.append(strats_combined[player_offset])
        #end

        #return value is a list of Strat where each strategy is player's strategy
        return br_list

    #following function links players to their best responses in a dict
    def make_br_dict(self):    
        br_dict = {}

        #following code goes through all s sub -i for each player i
        #dictionary values are sets to avoid duplicate entries
        for player in enumerate(self.player_list):
            br_dict[player[1]] = set()
            other_players_combos = combine(self.strat_mat[:player[0]] + \
                                           self.strat_mat[player[0] + 1:])
            for combo in other_players_combos:
                for br in self.make_br_list(player[1], combo):
                    br_dict[player[1]].add(br)
        #end

        #return value is a dict of the form player: list of Strat
        return br_dict
    
    #following function links players to their weakly dominated (wd) strategies
    def make_wd_dict(self):  

        #following function checks if, for a certain player
        #first_strat weakly dominates second_strat.
        #following function is also nested within make_wd_dict
        def check_weak_dominance(player, first_strat, second_strat):
            #argument first_strat expected to be Strat
            #argument second_strat expected to be Strat that is
            #not the same as first_strat

            #following code generates all possible s sub -i
            #then compares first_strat: s sub ij, and second strat: s sub ik
            #s sub ij >= s sub ik for all s sub -i, and s sub ij > s sub ik
            #for at least one s sub -i for s sub ij to weakly dominate s sub ik
            other_players_combos = combine(self.strat_mat[:player[0]] + \
                                           self.strat_mat[player[0] + 1:])
            flag = False
            for combo in other_players_combos:
                first_combo = combo[:player[0]] + (first_strat,) + combo[player[0]:]
                second_combo = combo[:player[0]] + (second_strat,) + combo[player[0]:]
                if self.matrix_dict[first_combo].get_payoff(player[1]) < \
                   self.matrix_dict[second_combo].get_payoff(player[1]):
                    return False
                if self.matrix_dict[first_combo].get_payoff(player[1]) > \
                   self.matrix_dict[second_combo].get_payoff(player[1]):
                    flag = True
            #end

            #return value is a boolean
            return flag
        #end nested function

        wd_dict = {}

        #following code goes through all players, for each player i
        #it creates all combinations of two strategies picked from s sub i,
        #s sub ij and s sub ik, and checks if s sub ij weakly dominates s sub ik,
        #or vice versa. the dominant strategy becomes the player's dict value
        for player in enumerate(self.player_list):
            wd_list = []
            strats_compared = itertools.combinations(self.strat_mat[player[0]], 2)
            for comparison in strats_compared:
                if check_weak_dominance(player, comparison[0], comparison[1]):
                    wd_list.append(comparison[1])
                if check_weak_dominance(player, comparison[1], comparison[0]):
                    wd_list.append(comparison[0])
            wd_dict[player[1]] = wd_list
        #end

        #return value is a dict of the form player : list of Strat
        return wd_dict

    #following function links players to their strictly dominated (sd) strategies
    def make_sd_dict(self): 

        #following function checks if, for a certain player
        #first_strat strictly dominates second_strat.
        #following function is also nested within make_sd_dict
        def check_strict_dominance(player, first_strat, second_strat):
            #argument first_strat expected to be Strat
            #argument second_strat expected to be Strat that is
            #not the same as first_strat

            #following code generates all possible s sub -i
            #then compares first_strat: s sub ij, and second strat: s sub ik
            #s sub ij > s sub ik for all s sub -i for s sub ij to strictly dominate s sub ik
            other_players_combos = combine(self.strat_mat[:player[0]] + \
                                           self.strat_mat[player[0] + 1:])
            for combo in other_players_combos:
                first_combo = combo[:player[0]] + (first_strat,) + combo[player[0]:]
                second_combo = combo[:player[0]] + (second_strat,) + combo[player[0]:]
                if self.matrix_dict[first_combo].get_payoff(player[1]) <= \
                   self.matrix_dict[second_combo].get_payoff(player[1]):
                    return False
            return True
            #end
        #end nested function
        
        sd_dict = {}

        #following code goes through all players, for each player i
        #it creates all combinations of two strategies picked from s sub i,
        #s sub ij and s sub ik, and checks if s sub ij strictly dominates s sub ik,
        #or vice versa. the dominant strategy becomes the player's dict value
        for player in enumerate(self.player_list):
            sd_list = []
            strats_compared = itertools.combinations(self.strat_mat[player[0]], 2)
            for comparison in strats_compared:
                if check_strict_dominance(player, comparison[0], comparison[1]):
                    sd_list.append(comparison[1])
                if check_strict_dominance(player, comparison[1], comparison[0]):
                    sd_list.append(comparison[0])
            sd_dict[player[1]] = sd_list
        #end

        #return value is a dict of the form player: list of Strat
        return sd_dict

    #following function links a player to their Sympy symbols
    #for each strat a player has, that strat has a distinct Sympy symbol
    #form is a0, a1, a2...for player 1, b0, b1, b2... for player 2 and so on
    #a player's final strat (what comes last) does not get its own symbol
    #but it is represented by 1 - n0 - n1 - n2... (where n could be a, b, etc.)
    def make_player_symbol_dict(self):
        
        player_symbol_dict = {}

        #following code goes through all players i and for each i scraps 
        #s sub ij if s sub ij is not weakly dominated (implicitly meaning if
        #it is not strictly dominated either) to make s' sub i
        for player in enumerate(self.player_list):
            player_symbol_list = []
            end_symbol = 1
            considered_strats = [strat for strat in self.strat_mat[player[0]] \
                                 if strat not in self.wd_dict[player[1]]]

            #following code goes through each s' sub i and creates a corresponding symbol
            #if s' sub i is not at the end, else it goes through all the symbols in
            #player_symbol_list and creates 1 - s' sub ia - s' sub ib...
            for strat in enumerate(considered_strats):
                if 0 < strat[0] == len(considered_strats) - 1:
                    for symbol in player_symbol_list:
                        end_symbol -= symbol
                    player_symbol_list.append(end_symbol)
                else:
                    symbol = Symbol(ALPHABET[player[0]] + NUMBERS[strat[0]])
                    player_symbol_list.append(symbol)
            #end

            player_symbol_dict[player[1]] = player_symbol_list
        #end

        #return value is a dict of the form player: list of Symbol
        return player_symbol_dict

    #following function links sympy Symbols to all strats of all players
    #form is the same as self.make_player_symbol_dict
    #the difference between this function and self.make_player_symbol_dict is that
    #this maps Symbols directly to strats, 1 to 1
    #whereas self.make_player_symbol_dict maps 1 player to all their Symbols
    def make_strat_symbol_dict(self):
        
        symbol_dict = {}
        
        #following code goes through all players i and for each i scraps
        #s sub ij if s sub ij is weakly dominated (implicitly meaning if
        #it is not strictly dominated either) to make s' sub i
        for player in enumerate(self.player_list):
            player_symbol_list = []
            end_symbol = 1
            considered_strats = [strat for strat in self.strat_mat[player[0]] \
                                     if strat not in self.wd_dict[player[1]]]

            #following code goes through each s' sub i and creates a corresponding symbol
            #if s' sub i is not at the end, else it goes through all the symbols in
            #player_symbol_list and creates 1 - s' sub ia - s' sub ib...
            #unlike in make_player_symbol_dict, this put the symbols 
            #into a dict of the form Symbol: Strat
            for strat in enumerate(considered_strats):
                if 0 < strat[0] == len(considered_strats) - 1:
                    for symbol in player_symbol_list:
                        end_symbol -= symbol
                    symbol_dict[end_symbol] = strat[1]               
                else:
                    symbol = Symbol(ALPHABET[player[0]] + NUMBERS[strat[0]])
                    symbol_dict[symbol] = strat[1]
                    player_symbol_list.append(symbol)
            #end
                    
        #end

        #return value is a dict of the form Symbol: Strat
        return symbol_dict

    #following function finds the pure-strat NE of this matrix subgame
    def make_pure_ne_list(self):
        
        pure_equilibrium_list = []
        candidates_counter = collections.Counter()

        #following code goes through all the possible s sub -i for each player i
        #then inserts BR sub i(s sub -i) into the proper position to make s, a complete
        #strategy profile, before incrementing that s in a Counter. Since NE are 
        #exclusively made up of BRs, they will have the highest "score".
        #casting to list is done for mutability, to allow an insertion
        #casting back to a tuple is to allow the counter to accept it
        for player in enumerate(self.player_list):
            other_players_combos = combine(self.strat_mat[:player[0]] + \
                                           self.strat_mat[player[0] + 1:])
            for combo in other_players_combos:
                for br in self.make_br_list(player[1], combo):
                    combo_copy = list(combo)
                    combo_copy.insert(player[0], br)
                    candidates_counter[tuple(combo_copy)] += 1
        #end

        #following code goes through all entries in the counter, from
        #highest to lowest score. Since candidates are scored by number of
        #BRs, all NE have a score that is the number of players. thus, the code
        #can stop searching when it sees a score lower than the player count
        for pure_equilibrium in candidates_counter.most_common():
            if pure_equilibrium[1] != len(self.player_list):
                break
            pure_equilibrium_list.append(tuple(pure_equilibrium[0]))

        #return value is a matrix of Strat
        return pure_equilibrium_list
    
    #following function finds the mixed-strat NE of this matrix subgame
    def make_mixed_ne_list(self):
        
        #following function creates a system of equations for sympy.solve()
        #to find the solution values for.
        #following function is also nested within make_mixed_ne_list
        def make_equation_list():   

            #following function creates terms, or sums of primaries, which are mononomials.
            #following function is also nested within make_equation_list
            def make_term_list(player, other_symbol_lists_combined):
                #argument player expected to be an enumeration: (int, string)
                #argument other_symbol_lists_combined expected to be mat of Strategy
                
                term_list = []
                
                #following code goes through all s sub j,k,l..., for all s sub -i,
                #for all s sub i, all in symbol form. for each iteration of s sub -i, this function
                #constructs a strategy profile s sub i + s sub -i and creates a mononomial
                #from multiplying the U sub i(s) by the symbols in s sub -i. It adds this mononomial
                #to the one created in the next iteration of s sub -i until a term is made.
                #a term is made for each iteration of s sub i
                for player_symbol in self.player_symbol_dict[player[1]]:
                    term = 0
                    for symbol_combo in other_symbol_lists_combined:
                        strat_profile = []
                        for symbol in enumerate(symbol_combo):
                            strat_profile.append(self.symbol_dict[symbol[1]])
                        strat_profile.insert(player[0], self.symbol_dict[player_symbol])
                        coefficient = self.matrix_dict[tuple(strat_profile)].get_payoff(player[1])
                        primary = coefficient
                        for equation_symbol in symbol_combo:
                            primary *= equation_symbol
                        term += primary
                    term_list.append(term)
                #end

                #return value is a list of Add, which can be thought of as a summation of a
                #list of Mul. the len of term_list is equal to the number of i's strats
                return term_list
            #end nested function
                
            equation_list = []
            
            #following code creates all s sub -i for each player i so make_term_list can
            #generate a term list for each i
            for player in enumerate(self.player_list):
                other_symbol_lists_combined = combine(symbol_mat[:player[0]] + symbol_mat[player[0] + 1:])
                term_list = make_term_list(player, other_symbol_lists_combined)

                #following code goes through each term in term_list, creating the expressions
                #term[0] - term[1], term[1] - term[2]...term[n-1] - term[n]. For sympy.solve,
                #x - y means x = y, so this checks if term[0] == term[1]...term[n-1] == term[n]
                #these expressions are put into equation_list
                for term in enumerate(term_list):
                    if term[0] == len(term_list) - 1:
                        break
                    equation = term[1]
                    equation -= term_list[term[0] + 1]
                    equation_list.append(equation)
                #end

            #end

            #return value is Add which can be thought of as a summation of a list of make_term_list
            return equation_list
        #end nested function

        #following function properly formats a list of solve() solutions
        #to be congruent with a list of Symbol. in short, it organizes solutions
        #by player and adds the extra symbols.
        #following function is also nested within make_mixed_ne_list
        def properly_format(solution_list):
            #argument solution_list expected to be a list of floats x, such that 0 <= x <= 1
            
            res_mat = [[]]
            initial = True

            #following code goes through each that was used for the sympy equation
            #and splits solution_list into multiple sub-lists, creating a res_mat such that
            #there is a 1 to 1 correspondence between res_mat and symbol_mat, and that each
            #solution is paired with the right symbol
            for symbol in enumerate(unique_symbol_list):
                if initial:
                    res_mat[-1].append(solution_list[symbol[0]])
                    initial = False
                    
                #following code detects when all the symbols for a certain player i are
                #accounted for and another player's i* symbols are coming up
                #before packaging all the old player's solutions into a list, it
                #creates a 1 - i's symbols symbol to pad the list, to accound for i's
                #last strategy that is only played when all other strategies aren't
                elif str(symbol[1])[0] != str(unique_symbol_list[symbol[0] - 1])[0]:
                    res_mat[-1].append(1 - sum(res_mat[-1]))
                    res_mat.append([])
                    res_mat[-1].append(solution_list[symbol[0]])
                #end
                    
                else:
                    res_mat[-1].append(solution_list[symbol[0]])
            res_mat[-1].append(1 - sum(res_mat[-1]))
            #end

            #return value is a matrix of floats, solutions organized by player 
            return res_mat
        #end nested function

        #following funtions determine if a would-be mixed strategy is really mixed or not.
        #following function is also nested within make_mixed_ne_list
        def is_mixed(solution):
            #argument solution expected to be a list of floats x, such that 0 <= x <= 1
            
            for prob in solution:
                if 0 != prob != 1:
                    return True
            return False
        #end nested function

        mixed_ne_list = []

        #following code creates lists that are ordered with respect to each other
        #and have a 1 to 1 correspondence, ex. symbol_mat[0] holds i's symbols, and
        #strat_mat[0] holds i's strats.
        #difference betwee unique_symbol_list and symbol_mat is that unique_symbol_list
        #does not include the 1 - s sub i symbol, because solve would not comprehend it
        symbol_mat = []
        strat_mat = []
        unique_symbol_list = []
        for player in self.player_list:
            symbol_list = []
            for symbol in self.player_symbol_dict[player]:
                symbol_list.append(symbol)
            symbol_mat.append(symbol_list)
        for player in self.player_list:
            strat_list = []
            for symbol in self.player_symbol_dict[player]:
                strat_list.append(self.symbol_dict[symbol])
            strat_mat.append(strat_list)
        for symbol_list in symbol_mat:
            for symbol in symbol_list:
                if len(str(symbol)) == 2:
                    unique_symbol_list.append(symbol)
        #end
                    
        equation_list = make_equation_list()

        #following code runs if there is no mixed equilibrium
        if len(equation_list) == 0:
            return []
        #end

        #following code solves the system of equations created by make_equation_list,
        #makes everything mutable, and formats the strat_mat to link with solution_mat
        #from there it creates mixed NE for mixed_ne_list
        solution_mat = list(solve(equation_list, unique_symbol_list, set = True))
        solution_mat[1] = [list(solution_tuple) for solution_tuple in solution_mat[1]]
        unique_symbol_list = [symbol for symbol in unique_symbol_list if symbol in solution_mat[0]]
        solution_mat = [properly_format(solution_list) for solution_list in solution_mat[1]]
        for solution_list in solution_mat:
            mixed_ne = []
            for strat_list, solution in zip(strat_mat, solution_list):
                if is_mixed(solution):
                    mixed_ne.append(Mixed_Strat(strat_list, solution))
            if len(mixed_ne):
                mixed_ne_list.append(tuple(mixed_ne))
        #end

        #return value is a list of Mixed_Strat
        return mixed_ne_list
    
    #following function returns the payoff for each player in a mixed ne
    def make_mixed_payoff_list(self, mixed_ne):
        #argument mixed_ne expected to be an element of self.mixed_ne_list

        mixed_payoff_list = [0 for player in self.player_list]
        player_strat_list = []
        player_prob_list = []

        #following code put all combinations of strat prob pairs for the mixed ne
        #into player_strat_list and player_prob_list
        for mixed_strat in mixed_ne:
            strat_list = [strat_prob[0] for strat_prob in mixed_strat.strat_prob_list]
            prob_list = [strat_prob[1] for strat_prob in mixed_strat.strat_prob_list]
            player_strat_list.append(strat_list)
            player_prob_list.append(prob_list)
        #end
            
        strat_combo_list = combine(player_strat_list)
        prob_combo_list = []

        #following code takes the product of the floats in each prob_combo and
        #puts it into prob_combo_list (1 to 1 correspondence with strat_combo_list)
        for prob_combo in combine(player_prob_list):
            prob_combo_list.append(functools.reduce(lambda x, y: x*y, prob_combo))
        #end

        #following code multiplies the probability that a strat combo will be played by
        #the payoff of the strat combo for player i, then adds that to player i's position
        #in mixed_payoff_list; for each strat_combo, for each player
        for strat_combo, prob_combo in zip(strat_combo_list, prob_combo_list):
            for player in enumerate(self.player_list):
                player_payoff = prob_combo * self.matrix_dict[strat_combo].get_payoff(player[1])
                mixed_payoff_list[player[0]] += player_payoff
        #end
                
        #return value is a list of payoff, with a 1 to 1 correspondence with self.player_list
        return mixed_payoff_list
    
    #following function returns the payoff that would be given to each player
    #assuming all players play this game logically
    def make_logical_payoff_mat(self):  
        
        logical_payoff_mat = []
        for pure_ne in self.pure_ne_list:
            logical_payoff_mat.append(self.matrix_dict[pure_ne].payoff_list)
        for mixed_ne in self.mixed_ne_list:
            logical_payoff_mat.append(self.make_mixed_payoff_list(mixed_ne))
        #return value is a matrix of payoff
        return logical_payoff_mat

    #following function returns the payoffs of the tree single, played logically
    #that is, the payoffs if the acting player plays a BR
    def get_logical_payoff_mat(self):

        #return value is a payoff
        return self.logical_payoff_mat                
    
    #following function prints relevant information 
    def print_state(self):

        #following code prints all strategy combinations and their yields to each player
        #in the form "s yields U sub i(s) for i, U sub j(s) for j..." 
        print("strats:")
        for strat_list in self.strats_combined_list:
            print(show_strats(strat_list, sep = " and "), "yield", end = " ")
            for player in self.matrix_dict[strat_list].player_list:
                print(self.matrix_dict[strat_list].get_payoff(player),
                      "for", player, end = ", ")
            print()
        #end

        #following code goes through all s sub -i for each player i
        #and prints the BR sub i(s sub -i) for each s sub -i
        for player in enumerate(self.player_list):
            other_players_combos = combine(self.strat_mat[:player[0]] + \
                                                 self.strat_mat[player[0] + 1:])
            for combo in other_players_combos:
                print("best response(s) for player", player[1], "to strat(ies)", show_strats(combo), "are: ", end ="")
                for br in self.make_br_list(player[1], combo):
                    print(br, end = ", ")
                print()
        #end

        #following code prints out all NE, pure and mixed    
        print("pure-strat nash equilibrium(a): ", end = "")
        for ne in self.pure_ne_list:
            print(show_strats(ne, " and "), end = ", ")
        print()
        print("mixed-strat nash equilibrium(a): ", end = "")
        for ne in self.mixed_ne_list:
            print(ne, end = ", ")
        print("\n")
        #end

