import itertools
import math as M
import random as R

from otree.api import *


author = 'Corey Scheinfeld'
doc = """
Your app description
"""


class Constants(BaseConstants):
    name_in_url = 'PGG'
    players_per_group = 3
    num_rounds = 10
    instructions_reg = 'PGG/instructions_reg.html'
    instructions_pun = 'PGG/instructions_pun.html'
    a_role = 'A'
    b_role = 'B'
    c_role = 'C'


class Subsession(BaseSubsession):
    pass


class Group(BaseGroup):
    type = models.StringField()
    is_partisan = models.BooleanField()
    is_pun = models.BooleanField()

    group_pot = models.FloatField(initial=0)
    A_cont = models.FloatField()
    B_cont = models.FloatField()
    C_cont = models.FloatField()
    A_payoff = models.FloatField()
    B_payoff = models.FloatField()
    C_payoff = models.FloatField()
    A_punished = models.FloatField(initial=0)
    B_punished = models.FloatField(initial=0)
    C_punished = models.FloatField(initial=0)


class Player(BasePlayer):
    original_label = models.StringField()
    first_payoff = models.FloatField(initial=0)
    final_payoff = models.FloatField(initial=0)
    us_payoff = models.FloatField(initial=0)
    base_combo_payoff = models.FloatField(initial=0)
    punished = models.FloatField()
    round_payoff = models.FloatField(initial=0)
    is_dem = models.BooleanField(
        choices=[[False, 'Republican'], ['True', 'Democrat']], label=""
    )
    label = models.StringField()
    group_contribution = models.FloatField(
        label="Your Contribution to the Group Project:", min=0, max=20
    )
    individual_share = models.FloatField()
    kept = models.FloatField()
    punishA = models.FloatField(label="", min=0, max=5)
    punishB = models.FloatField(label="", min=0, max=5)
    punishC = models.FloatField(label="", min=0, max=5)
    reduce = models.FloatField()
    punishment_field_name = models.StringField(
        doc="The field other players use to punish this player"
    )


# FUNCTIONS
def creating_session(subsession: Subsession):
    subsession.conditions = "start"
    if subsession.round_number == 1:
        rand_val = itertools.cycle(
            [
                [True, True],
                [False, True],
                [True, False],
                [True, False],
                [False, False],
                [True, True],
                [False, True],
                [True, True],
                [False, True],
                [True, True],
                [False, True],
            ]
        )
        for g in subsession.get_groups():
            [g.is_pun, g.is_partisan] = next(rand_val)

        for p in subsession.get_players():
            p.punishment_field_name = f'punish{p.role}'


def adjust_group(group: Group):
    """what does this do??? seems redundant with just asking people their affiliation"""
    if group.subsession.round_number == 1:
        for player in group.get_players():
            player.original_label = player.participant.label
            if (
                player.original_label[
                    (len(player.original_label) - 8) : (len(player.original_label))
                ]
                == 'Democrat'
            ):
                player.affiliation = player.original_label[
                    (len(player.original_label) - 8) : (len(player.original_label))
                ]
                player.participant.label = player.original_label[
                    0 : (len(player.original_label) - 8)
                ]
            elif (
                player.original_label[
                    (len(player.original_label) - 10) : (len(player.original_label))
                ]
                == 'Republican'
            ):
                player.affiliation = player.original_label[
                    (len(player.original_label) - 10) : (len(player.original_label))
                ]
                player.participant.label = player.original_label[
                    0 : (len(player.original_label) - 10)
                ]


def set_pot(group: Group):
    for players in group.get_players():
        group.group_pot = group.group_pot + players.group_contribution
        players.kept = 20 - players.group_contribution
    for players in group.get_players():
        players.individual_share = round(group.group_pot * (2 / 3), 2)
        players.first_payoff = (
            20 - players.group_contribution + round(players.individual_share, 2)
        )


def distribute_punishments(group: Group):
    players = group.get_players()

    for player in group.get_players():
        setattr(player, player.punishment_field_name, 0)
        player.punished = (
            sum(getattr(p, player.punishment_field_name) for p in players) * 3
        )
        player.reduce = player.punishA + player.punishB + player.punishC
        player.round_payoff = max(0, player.first_payoff - player.punished)
        player.round_payoff = round(player.round_payoff - player.reduce, 2)


def set_final_payoff(group: Group):
    for player in group.get_players():

        player_in_rounds = player.in_rounds(1, group.round_number + 1)
        if group.is_pun:
            player.final_payoff = sum(p.round_payoff for p in player_in_rounds)
        else:
            player.final_payoff = sum(p.first_payoff for p in player_in_rounds)

        player.us_payoff = round(player.final_payoff / 100, 2)
        player.base_combo_payoff = player.us_payoff + 2.00


# PAGES
class MyWaitPage(WaitPage):
    group_by_arrival_time = True

    @staticmethod
    def is_displayed(player: Player):
        return player.round_number == 1


class GroupWaitPage(WaitPage):
    @staticmethod
    def is_displayed(player: Player):
        return player.round_number == 1

    after_all_players_arrive = adjust_group


class GroupDisplay(Page):
    @staticmethod
    def is_displayed(player: Player):
        group = player.group
        return player.round_number == 1 and group.is_partisan


class GroupDisplayCon(Page):
    @staticmethod
    def is_displayed(player: Player):
        group = player.group
        return player.round_number == 1 and not group.is_partisan


class Introduction(Page):
    @staticmethod
    def is_displayed(player: Player):
        return player.round_number == 1


class NextWait(WaitPage):
    after_all_players_arrive = adjust_group


class Contribution(Page):
    form_model = 'player'
    form_fields = ['group_contribution']


class ResultsWaitPage(WaitPage):
    after_all_players_arrive = set_pot


class Results1(Page):
    pass


class Punishment(Page):
    form_model = 'player'

    @staticmethod
    def get_form_fields(player: Player):
        group = player.group
        if group.is_pun:
            return [p.punishment_field_name for p in player.get_others_in_group()]
        return []


class PunishmentWait(WaitPage):
    @staticmethod
    def is_displayed(player: Player):
        group = player.group
        return group.is_pun

    after_all_players_arrive = distribute_punishments


class Results2(Page):
    @staticmethod
    def is_displayed(player: Player):
        group = player.group
        return group.is_pun


class FinalWait(WaitPage):
    after_all_players_arrive = set_final_payoff

    @staticmethod
    def is_displayed(player: Player):
        return player.round_number == Constants.num_rounds


class FinalResults(Page):
    @staticmethod
    def is_displayed(player: Player):
        return player.round_number == Constants.num_rounds


page_sequence = [
    MyWaitPage,
    GroupWaitPage,
    GroupDisplay,
    GroupDisplayCon,
    Introduction,
    NextWait,
    Contribution,
    ResultsWaitPage,
    Results1,
    Punishment,
    PunishmentWait,
    Results2,
    FinalWait,
    FinalResults,
]