from collections import defaultdict as dd
import json

roles = [u'oberon', u'merlin', u'mordred', u'morgana', u'percival', u'assassin']

ANY_GOOD = [u'merlin', u'percival', u'good']
ANY_EVIL = [u'oberon', u'mordred', u'morgana', u'assassin', u'evil']


def isSpy(x, game):
    return game["players"][x]["spy"]

def approval(x):
    if x["Approve"] + x["Reject"] == 0:
        return 0.0
    return (x["Approve"] * 1.0) / (x["Approve"] + x["Reject"])

def getRole(x, game):
    if str(x) not in game["role_by_seat"]:
        if isSpy(x, game):
            return "evil"
        return "good"
    return game["role_by_seat"][str(x)]

def sumKeys(a, b):
    out = dd(int)
    for k, v in a.items():
        out[k] += v
    for k, v in b.items():
        out[k] += v
    return out

def sum_helper(table, constraint_list):
    if len(constraint_list) == 0:
        return table
    c = constraint_list[0]
    out = dd(int)
    if isinstance(c, type([])):
        for v in c:
            out = sumKeys(out, sum_helper(table[v], constraint_list[1:]))
    elif c is None:
        for k in table:
            out = sumKeys(out, sum_helper(table[k], constraint_list[1:]))
    else:
        return sumKeys(out, sum_helper(table[c], constraint_list[1:]))
    return out

class TrainTable(object):
    """Docstring for TrainTable. """

    # good: round, vote, proposer_role, evils on team
    def __init__(self, path):
        self.good_vote = dd(lambda: dd(lambda: dd(lambda: dd(lambda: dd(lambda: dd(int))))))
        self.percy_vote = dd(lambda: dd(lambda: dd(lambda: dd(lambda: dd(lambda: dd(lambda: dd(int)))))))
        self.merlin_vote = dd(lambda: dd(lambda: dd(lambda: dd(lambda: dd(lambda: dd(lambda: dd(int)))))))
        self.evil_vote = dd(lambda: dd(lambda: dd(lambda: dd(lambda: dd(lambda: dd(int))))))
        self.morgana_vote = dd(lambda: dd(lambda: dd(lambda: dd(lambda: dd(lambda: dd(int))))))
        self.assassin_vote = dd(lambda: dd(lambda: dd(lambda: dd(lambda: dd(lambda: dd(int))))))
        self.oberon_vote = dd(lambda: dd(lambda: dd(lambda: dd(lambda: dd(lambda: dd(int))))))
        self.mordred_vote = dd(lambda: dd(lambda: dd(lambda: dd(lambda: dd(lambda: dd(int))))))
        self._path = path
        self._build_prob_table()

    def GoodVote(self, round=None, vote=None, proposer_role=None, evil_count=None, size=None):
        return sum_helper(self.good_vote, [size, round, vote, proposer_role, evil_count])

    def PercivalVote(self, round=None, vote=None, proposer_role=None, evil_count=None, team_contains=None, size=None):
        return sum_helper(self.percy_vote, [size, round, vote, proposer_role, evil_count, team_contains])

    def EvilVote(self, round=None, vote=None, proposer_role=None, evil_count=None, size=None):
        return sum_helper(self.evil_vote, [size, round, vote, proposer_role, evil_count])

    def MordredVote(self, round=None, vote=None, proposer_role=None, evil_count=None, size=None):
        return sum_helper(self.mordred_vote, [size, round, vote, proposer_role, evil_count])

    def MorganaVote(self, round=None, vote=None, proposer_role=None, evil_count=None, size=None):
        return sum_helper(self.morgana_vote, [size, round, vote, proposer_role, evil_count])

    def AssassinVote(self, round=None, vote=None, proposer_role=None, evil_count=None, size=None):
        return sum_helper(self.asassin_vote, [size, round, vote, proposer_role, evil_count])

    def OberonVote(self, round=None, vote=None, proposer_role=None, evil_count=None, size=None):
        return sum_helper(self.oberon_vote, [size, round, vote, proposer_role, evil_count])

    def MerlinVote(self, round=None, vote=None, proposer_role=None, evil_count=None, known_evil=None, size=None):
        return sum_helper(self.merlin_vote, [size, round, vote, proposer_role, evil_count, known_evil])


    def _build_prob_table(self):
        f = open(self._path)
        self._data = json.load(f)
        f.close()
        self._build_good_vote(self.good_vote, "good")
        self._build_percy_vote(self.percy_vote)
        self._build_merlin_vote(self.merlin_vote)
        self._build_evil_vote(self.evil_vote, "evil")
        self._build_evil_vote(self.mordred_vote, "mordred")
        self._build_evil_vote(self.morgana_vote, "morgana")
        self._build_evil_vote(self.assassin_vote, "assassin")
        self._build_evil_vote(self.oberon_vote, "oberon")

    def _build_good_vote(self, table, exprole):
        for game in self._data:
            size = len(game["players"])
            for player in game["players"]:
                idx = player["seat"]
                if player["spy"]:
                    continue
                if getRole(idx, game) != exprole:
                    continue
                for round_n, round in enumerate(game["log"]):
                    for vote_n in range(1, 6):
                        if str(vote_n) not in round:
                            break
                        vote = round[str(vote_n)]
                        prole = getRole(vote["proposer"], game)
                        if vote["proposer"] == idx:
                            prole = "self"
                        evils = sum([int(isSpy(x, game)) for x in vote["team"]])
                        action = vote["votes"][idx]
                        table[size][round_n + 1][vote_n][prole][evils][action] += 1

    def _build_evil_vote(self, table, exprole):
        for game in self._data:
            size = len(game["players"])
            for player in game["players"]:
                idx = player["seat"]
                if not player["spy"]:
                    continue
                if getRole(idx, game) != exprole:
                    continue
                for round_n, round in enumerate(game["log"]):
                    for vote_n in range(1, 6):
                        if str(vote_n) not in round:
                            break
                        vote = round[str(vote_n)]
                        prole = getRole(vote["proposer"], game)
                        if vote["proposer"] == idx:
                            prole = "self"
                        evils = sum([int(isSpy(x, game)) for x in vote["team"]])
                        action = vote["votes"][idx]
                        table[size][round_n + 1][vote_n][prole][evils][action] += 1

    def _build_percy_vote(self, table):
        for game in self._data:
            size = len(game["players"])
            for player in game["players"]:
                idx = player["seat"]
                if player["spy"]:
                    continue
                if getRole(idx, game) != "percival":
                    continue
                for round_n, round in enumerate(game["log"]):
                    for vote_n in range(1, 6):
                        if str(vote_n) not in round:
                            break
                        vote = round[str(vote_n)]
                        prole = getRole(vote["proposer"], game)
                        evils = sum([int(isSpy(x, game)) for x in vote["team"]])
                        action = vote["votes"][idx]
                        tc = ""
                        rolesOnTeam = [getRole(x, game) for x in vote["team"]]
                        if "merlin" in rolesOnTeam:
                            tc += "merlin"
                        if "morgana" in rolesOnTeam:
                            tc += "morgana"
                        if tc == "":
                            tc = "neither"
                        table[size][round_n + 1][vote_n][prole][evils][tc][action] += 1

    def _build_merlin_vote(self, table):
        for game in self._data:
            size = len(game["players"])
            for player in game["players"]:
                idx = player["seat"]
                if player["spy"]:
                    continue
                if getRole(idx, game) != "merlin":
                    continue
                for round_n, round in enumerate(game["log"]):
                    for vote_n in range(1, 6):
                        if str(vote_n) not in round:
                            break
                        vote = round[str(vote_n)]
                        prole = getRole(vote["proposer"], game)
                        evils = sum([int(isSpy(x, game)) for x in vote["team"]])
                        action = vote["votes"][idx]
                        known_evils = sum([int(isSpy(x, game) and (getRole(x, game) != "mordred")) for x in vote["team"]])
                        table[size][round_n + 1][vote_n][prole][evils][known_evils][action] += 1









def main():
    table = TrainTable("dataset.json")


if __name__ == '__main__':
    main()
