from otree.api import *
from .receivers import *
from .matching import *
from datetime import datetime
from shared import helpers
from itertools import count
from shared.helpers import NOCPARTNER_FEE, NOCPARTNER_URL

doc = """
Your app description
"""

# THIS APP DEPENDS ON APP `intro2`
# ALWAYS SET player.participant.wait_page_arrival = time.time() ON THE LAST PAGE OF THE ROUND FOR EACH MODE (ALONE, NORMAL)

# TEST THE FOLLOWING PATHS
# Consent was not given (x),
# Comprehension1 failed (x),
# Comprehension2 failed (x),
# Attention1 and Attention2 failed (x),
# Common Account is 0.00$ (x),
# Player 1 of the group is still in Round 1, but Player 2 is already in a new group in Round 2. See that everything works as required for player 1 (x)
# Round 1:
# - no Account Partner (x),
# - Account Partner and no Communication Partner (x),
# - Account Partner and Communication Partner (x),
# - No Allocation to Account Partner made (x),
# - No Allocation to Receiver Couple made (x),
# Round 2:
# - no Account Partner (x),
# - Account Partner and no Communication Partner (x),
# - Account Partner and Communication Partner (x),
# - No Allocation to Account Partner made (x),
# - No Allocation to Receiver Couple made (x),
# Round 3:
# - no Account Partner (x),
# - Account Partner and no Communication Partner (x),
# - Account Partner and Communication Partner (x),
# - One people in group did not pick preferred Communication Partner (x),
# - Two people in group did not pick preferred Communication Partner (),
# - Three people in group did not pick preferred Communication Partner (x),
# - Noone picked preferred Communication Partner (x),
# - No Allocation to Account Partner made (x),
# - No Allocation to Receiver Couple made (x),

# need to: retest all above cases

class C(BaseConstants):
    NAME_IN_URL = "allocation"
    PLAYERS_PER_GROUP = None
    NUM_ROUNDS = helpers.NUM_ROUNDS
    NOCPARTNER_URL = NOCPARTNER_URL

    MIN_ON_CHAT = helpers.MIN_ON_CHAT
    TIME_ON_CHAT = MIN_ON_CHAT * 60
    SOFT_TIMEOUT = 3 * 60

    DROPOUT_ERR_MSG = "you did not specify an allocation with the slider on time"
    MORAL_CHOICES = [
        "Very morally inappropriate",
        "Somewhat morally inappropriate",
        "Somewhat morally appropriate",
        "Very morally appropriate",
    ]
    SOCIAL_CHOICES = [
        "Very socially inappropriate",
        "Somewhat socially inappropriate",
        "Somewhat socially appropriate",
        "Very socially appropriate",
    ]
    ENDOGENEOUS_PICKS = dict(A1=["A2", "D1"], A2=["A1", "D2"], D1=["A1", "D2"], D2=["A2", "D1"])

    ACCOUNT_PARTNER = ReceiverType.ACCOUNT_PARTNER.value
    THIRD_PARTY = ReceiverType.THIRD_PARTY.value

    # For incentivized Social Appropriateness
    # Tuples are (field_name, most_common_answer)
    MOST_COMMON_SOCIAL_APP_ANSWERS = [
        ("social_split_back", SOCIAL_CHOICES[0]),
        ("social_split_equal", SOCIAL_CHOICES[3]),
        ("social_split_correct", SOCIAL_CHOICES[1]),
        ("social_split_one", SOCIAL_CHOICES[0]),
    ]
    SOCIAL_APP_BONUS = cu(1)


def vars_for_admin_report(subsession):
    unassigned_receivers = Receiver.filter(subsession=subsession, type=C.ACCOUNT_PARTNER, dictator=None)
    return dict(unassigned_receivers=unassigned_receivers, nr_unassigned_receivers=len(unassigned_receivers))


class Subsession(BaseSubsession):
    last_matching_run = models.FloatField()


def creating_session(subsession: Subsession):
    session = subsession.session
    if subsession.round_number == 1:
        # Create a receiver dict to share names and profile_pictures of the same receiver among subsessions
        session.receivers = dict()
        session.id_gen = count()
        for p in subsession.get_players():
            participant = p.participant
            # Set default values
            participant.is_dropout = False
            participant.display_info = None
            # Set round to implement for payment
            participant.paid_round = random.choice(range(1, C.NUM_ROUNDS + 1))
            # List of participant ids of matched partners to enforce stranger matching
            participant.previous_comm_partners = []
            participant.previous_acc_partners = []
            participant.previous_third_party_partners = []
            # Dict to keep track of allocations in each round
            participant.spent = []
            participant.kept = []

    set_receivers(subsession, session)


class Group(BaseGroup):
    pass


class Player(BasePlayer):
    # Info models
    name = models.StringField()
    profile_picture = models.StringField()
    reward_group = models.StringField()
    reward = models.CurrencyField()

    # Account models
    account_partner_id = models.IntegerField(initial=-1)
    common_account = models.CurrencyField(initial=0)
    intention = models.CurrencyField(
        label="Please enter below the amount you want to keep for yourself. Once you are done, click on the button to continue to the chat."
    )

    # Communication models
    is_grouped = models.BooleanField()
    comm_partner_id = models.IntegerField(initial=-1)
    comm_partner_reward_group = models.StringField(initial="")
    allocation = models.CurrencyField(blank=True)
    endogeneous_role = models.StringField()
    preferred_role = models.StringField()
    selected_role = models.StringField()
    channel_code = models.StringField()
    last_move = models.CurrencyField()
    partners_last_move = models.CurrencyField()

    # Third-party Allocation models
    receiver1_id = models.IntegerField(initial=-1)
    receiver2_id = models.IntegerField(initial=-1)
    receivers_comm_account = models.CurrencyField(initial=0)
    receivers_allocation = models.CurrencyField()

    # Moral appropriateness models
    allocation_reasoning = models.LongStringField()
    moral_split_back = models.StringField(choices=C.MORAL_CHOICES)
    moral_split_equal = models.StringField(choices=C.MORAL_CHOICES)
    moral_split_correct = models.StringField(choices=C.MORAL_CHOICES)
    moral_split_one = models.StringField(choices=C.MORAL_CHOICES, label="Keeping everything for oneself")

    # Social appropriateness models
    partner_selection = models.LongStringField()
    social_split_back = models.StringField(choices=C.SOCIAL_CHOICES)
    social_split_equal = models.StringField(choices=C.SOCIAL_CHOICES)
    social_split_correct = models.StringField(choices=C.SOCIAL_CHOICES)
    social_split_one = models.StringField(choices=C.SOCIAL_CHOICES, label="Keeping everything for oneself")
    social_bonus = models.CurrencyField(initial=0)

    # Endogeneous models
    did_not_pick_preferred_partner = models.BooleanField()

    def has_account_partner(self):
        return self.account_partner_id >= 0

    def get_account_partner(self):
        # , receiver_id=player.account_partner_id)
        [partner] = Receiver.filter(subsession=self.subsession, dictator=self)
        return partner

    def get_third_party_receivers(self):
        [receiver1] = Receiver.filter(subsession=self.subsession, receiver_id=self.receiver1_id)
        [receiver2] = Receiver.filter(subsession=self.subsession, receiver_id=self.receiver2_id)
        return receiver1, receiver2

    def has_comm_partner(self):
        return self.comm_partner_id >= 0

    def get_comm_partner(self):
        players = self.group.get_players()
        comm_partner_id = self.comm_partner_id
        [comm_partner] = [p for p in players if p.id_in_subsession == comm_partner_id]
        return comm_partner

    def has_group(self):
        has_group = self.field_maybe_none("is_grouped")
        if has_group is None:
            has_group = bool(self.get_others_in_group())
            self.is_grouped = has_group
        return has_group

    def is_endogeneous_round(self) -> bool:
        return self.round_number == 3


def intention_max(player):
    return player.common_account


def allocation_max(player):
    return player.common_account


def receivers_allocation_max(player):
    return player.receivers_comm_account


class Receiver(ExtraModel):
    """A Receiver for a subsession"""

    receiver_id = models.IntegerField()  # An ID for single Receiver across subsessions
    subsession = models.Link(Subsession)  # round_nr or subsession maps to round_nr of Part 1
    dictator = models.Link(Player)  # if None: Receiver is not assigned to Dictator yet
    type = models.IntegerField()  # 1: Account Partner, 2: Third Party Partner
    receiver_session = models.StringField()  # Session code of Receiver session
    receiver = models.StringField()  # ID of Receiver from Receiver session
    name = models.StringField()  # Display name
    reward_group = models.StringField()  # Reward group in Part 1
    reward = models.CurrencyField()  # Reward in Part 1 of the corresponding round
    profile_picture = models.StringField()  # Display profile picture path
    payoff = models.CurrencyField()  # The payoff that was implemented for them
    is_round_paid = models.BooleanField()  # Whether the round of this subsession is paid


class SliderMove(ExtraModel):
    dictator = models.Link(Player)
    timestamp = models.IntegerField()
    amount = models.CurrencyField()


def custom_export(players):
    yield [
        "receiver_id",
        "receiver_session",
        "receiver",
        "round_nr",
        "is_round_paid",
        "payoff",
        "type",
        "name",
        "reward_group",
        "reward",
        "profile_picture",
        "dictator_session",
        "dictator",
    ]
    for r in Receiver.filter():
        round_nr = r.subsession.round_number
        dictator_session = r.subsession.session.code
        dictator = r.dictator.id_in_subsession if r.dictator else None
        yield [
            r.receiver_id,
            r.receiver_session,
            r.receiver,
            round_nr,
            r.is_round_paid,
            r.payoff,
            r.type,
            r.name,
            r.reward_group,
            r.reward,
            r.profile_picture,
            dictator_session,
            dictator,
        ]

    yield ["session", "participant", "timestamp", "amount"]
    for m in SliderMove.filter():
        session = m.dictator.session
        yield [session.code, m.dictator.id_in_subsession, datetime.fromtimestamp(int(m.timestamp)), m.amount]


# PAGES
class Matching(WaitPage):
    body_text = f"We are matching you with participants. <b>Please do not leave this page.</b> The maximum waiting time is {MAX_WAIT_TIME_MIN} minutes."
    # Option 1: Acc and Comm -> Play normally (x)
    # Option 2: Acc and no Comm -> Play alone (x)
    # Option 3: no Acc and Comm -> Not matched, fall through to Option 4 (x)
    # Option 4: no Acc and no Comm -> Continue to next round (x)
    group_by_arrival_time = True


def group_by_arrival_time_method(subsession: Subsession, waiting_players):
    # Match players with an Account partner, if they don't have one yet
    unassigned_receivers = Receiver.filter(subsession=subsession, type=C.ACCOUNT_PARTNER, dictator=None)
    matching.find_account_partner(waiting_players, unassigned_receivers)

    # Drop players waiting too long
    if g := matching.waiting_too_long(waiting_players):
        return g

    # Run Matching only once per minute
    last_run = subsession.field_maybe_none("last_matching_run")
    if last_run and (timespan := time.time() - last_run) < 10:
        print("too close to last matching:", timespan, "seconds")
        return
    print("try matching")
    subsession.last_matching_run = time.time()
    # Try and match each player with a Communication partner
    if match := matching.try_match(subsession, waiting_players):
        return match


class AccountPage(Page):
    """This Page is only shown to Dictators with an Account Partner"""

    @staticmethod
    def is_displayed(player: Player):
        return player.has_account_partner()


class CommunicationPage(AccountPage):
    """This Page is only shown to Dictators with an Account Partner and a group"""

    @staticmethod
    def is_displayed(player: Player):
        return AccountPage.is_displayed(player) and player.has_group()


class EndogeneousPageWithPartners(CommunicationPage):
    """This Page is only shown in the endogeneous round to Dictators with an Account Partner and a group"""

    @staticmethod
    def is_displayed(player: Player):
        return player.is_endogeneous_round() and CommunicationPage.is_displayed(player)


class EndogeneousPage(AccountPage):
    """This Page is only shown in the endogeneous round to Dictators with an Account Partner"""

    @staticmethod
    def is_displayed(player: Player):
        return player.is_endogeneous_round() and AccountPage.is_displayed(player)


class NoPartner(Page):
    timeout_seconds = C.SOFT_TIMEOUT

    @staticmethod
    def before_next_page(player: Player, timeout_happened):
        participant = player.participant
        participant.wait_page_arrival = time.time()

        # If this round was supposed to be paid for a Dictator who has not found an Account partner
        if not player.has_account_partner():
            # Set allocations to zero
            participant.spent.append(0)
            participant.kept.append(0)

    @staticmethod
    def is_displayed(player: Player):
        return not player.has_account_partner() or not player.has_group()
        
    @staticmethod
    def vars_for_template(player: Player):
        return dict(nocpartner_fee=cu(helpers.NOCPARTNER_FEE))

class AccountPartner(AccountPage):
    form_model = "player"
    form_fields = ["intention"]

    timeout_seconds = C.SOFT_TIMEOUT

    @staticmethod
    def vars_for_template(player: Player):
        if player.field_maybe_none("reward") is None:
            helpers.set_info([player])
        return dict(
            acc_partner=player.get_account_partner(),
            task_path=f"global/games/game{player.round_number}.png",
            task_name=helpers.get_task_name(player),
            comm_account=float(player.common_account),
        )

    @staticmethod
    def error_message(player: Player, values):
        if values["intention"] < 0:
            return "Please state your intention with the slider."

    @staticmethod
    def before_next_page(player: Player, timeout_happened):
        # If timeout_happened, we set the intention to None, bc oTree would would set it
        # to 0 per default after timeout, but we don't want that
        if timeout_happened:
            player.intention = None


class EndogeneousInstructions(EndogeneousPageWithPartners):
    timeout_seconds = C.SOFT_TIMEOUT


class PickCommunicationPartner(EndogeneousPageWithPartners):
    form_model = "player"
    form_fields = ["preferred_role"]

    timeout_seconds = C.SOFT_TIMEOUT

    @staticmethod
    def vars_for_template(player: Player):
        partner_roles = C.ENDOGENEOUS_PICKS[player.endogeneous_role]
        [p1, p2] = [p for p in player.get_others_in_group() if p.endogeneous_role in partner_roles]
        return dict(
            p1=p1,
            p2=p2,
            p1_r1_account=(p1.participant.kept[0] + p1.participant.spent[0]),
            p1_r2_account=(p1.participant.kept[1] + p1.participant.spent[1]),
            p2_r1_account=(p2.participant.kept[0] + p2.participant.spent[0]),
            p2_r2_account=(p2.participant.kept[1] + p2.participant.spent[1]),
        )

    @staticmethod
    def before_next_page(player: Player, timeout_happened):
        # If someone has not picked a preferred_role in time, mark that in the data
        if timeout_happened and not player.field_maybe_none("preferred_role"):
            print(player, "did not pick preference")
            player.did_not_pick_preferred_partner = True


class SelectCommunicationPartner(WaitPage):
    @staticmethod
    def after_all_players_arrive(group: Group):
        players = group.get_players()
        [print(p.preferred_role) for p in players]
        # Select the players who picked their preferred partner
        players_who_picked = [p for p in players if p.field_maybe_none("preferred_role")]
        print("players who picked:", players_who_picked)
        # If no player picked, set random preferred roles for everyone
        if not len(players_who_picked):
            print("Noone picked, choosing random preferred roles")
            for p in players:
                p.did_not_pick_preferred_partner = True
                p.preferred_role = random.choice(C.ENDOGENEOUS_PICKS[p.endogeneous_role])
            players_who_picked = players
        # Select random player who's choice is implemented
        p1, p2, p3, p4 = random.choice(players_who_picked), None, None, None
        for p in players:
            if p == p1:
                continue
            # Found player's choice partner
            elif p.endogeneous_role == p1.preferred_role:
                p2 = p
            elif not p3:
                p3 = p
            else:
                p4 = p

        set_communication_partner(p1, p2)
        set_communication_partner(p3, p4)

    @staticmethod
    def is_displayed(player: Player):
        return EndogeneousPageWithPartners.is_displayed(player)


class CommunicationPartner(CommunicationPage):
    timeout_seconds = C.SOFT_TIMEOUT

    @staticmethod
    def vars_for_template(player: Player):
        return dict(
            partner=player.get_comm_partner(),
            acc_partner=player.get_account_partner(),
            is_endogeneous_round=player.is_endogeneous_round(),
        )


class ChatWait(WaitPage):
    @staticmethod
    def is_displayed(player: Player):
        return player.has_group()


class Chat(AccountPage):
    form_model = "player"
    form_fields = ["allocation"]

    timeout_seconds = C.TIME_ON_CHAT

    @staticmethod
    def vars_for_template(player: Player):
        if player.has_comm_partner():
            partner = player.get_comm_partner()
            player_account = float(player.common_account)
            partner_account = float(partner.common_account)

            return dict(
                plays_alone=False,
                player_account=player_account,
                partner_account=partner_account,
                partner_intention=float(partner.intention)
                if partner.field_maybe_none("intention") is not None
                else -1,
                player_intention=float(player.intention) if player.field_maybe_none("intention") is not None else -1,
                acc_partner=player.get_account_partner(),
                partner=partner,
            )
        else:
            return dict(
                plays_alone=True,
                player_account=float(player.common_account),
                player_intention=float(player.intention) if player.field_maybe_none("intention") is not None else -1,
            )

    @staticmethod
    def js_vars(player):
        last_move = player.field_maybe_none("last_move")
        last_move = float(last_move) if last_move is not None else last_move
        partners_last_move = player.field_maybe_none("partners_last_move")
        partners_last_move = float(partners_last_move) if partners_last_move is not None else partners_last_move
        return dict(last_move=last_move, partners_last_move=partners_last_move)

    @staticmethod
    def live_method(player: Player, amount):
        now = time.time()
        try:
            amount = float(amount)
        except:
            return
        SliderMove.create(
            dictator=player,
            timestamp=now,
            amount=amount,
        )
        player.last_move = amount

        if player.has_comm_partner():
            partner = player.get_comm_partner()
            partner.partners_last_move = amount
            return {partner.id_in_group: amount}

    @staticmethod
    def app_after_this_page(player: Player, upcoming_apps):
        if helpers.is_dropout(player):
            return "dropout"

    @staticmethod
    def error_message(player: Player, values):
        if not player.has_comm_partner():
            if values["allocation"] is None:
                return "You must make an allocation before proceeding."

    @staticmethod
    def before_next_page(player: Player, timeout_happened):
        if player.field_maybe_none("allocation") is None:
            if player.field_maybe_none("last_move") is None:
                # Drop player and re-enqueue allocation partner
                helpers.set_dropout(player, C.DROPOUT_ERR_MSG)
                receiver = player.get_account_partner()
                receiver.dictator = None
                return
            # Recover allocation from last_move of player
            player.allocation = player.last_move
        player.last_move = None
        if helpers.is_dropout(player):
            return

        # Assign two Receivers to the Dictator to allocate between in next Page
        receivers = Receiver.filter(
            subsession=player.subsession, type=C.THIRD_PARTY
        )  # Fetch third-party Receiver pool
        random.shuffle(receivers)
        matching.find_receiver_couple(player, receivers)

        participant = player.participant
        participant.wait_page_arrival = time.time()
        alloc = player.allocation
        participant.kept.append(alloc)
        alloc_to_rec = player.common_account - alloc
        participant.spent.append(alloc_to_rec)

        # Set payoff for Dictator and Receiver
        round_nr = player.round_number
        if participant.paid_round == round_nr:
            participant.payoff = alloc
        receiver = player.get_account_partner()
        if receiver.is_round_paid:
            print(f"Set payoff for R{receiver.receiver_id}: {alloc_to_rec}")
            receiver.payoff = alloc_to_rec


class Allocation(AccountPage):
    form_model = "player"
    form_fields = ["receivers_allocation"]

    @staticmethod
    def live_method(player: Player, amount):
        now = time.time()
        try:
            amount = float(amount)
        except:
            return

        SliderMove.create(
            dictator=player,
            timestamp=now,
            amount=amount,
        )
        player.last_move = amount

    @staticmethod
    def vars_for_template(player: Player):
        receiver1, receiver2 = player.get_third_party_receivers()
        account = float(player.receivers_comm_account)
        return dict(receiver1=receiver1, receiver2=receiver2, account=account, task_name=helpers.get_task_name(player))

    @staticmethod
    def js_vars(player):
        last_move = player.field_maybe_none("last_move")
        last_move = float(last_move) if last_move is not None else last_move
        return dict(last_move=last_move)

    @staticmethod
    def app_after_this_page(player: Player, upcoming_apps):
        if helpers.is_dropout(player):
            return "dropout"

    @staticmethod
    def before_next_page(player: Player, timeout_happened):
        def set_payoff(r, payoff):
            print(f"New allocation for R{r.receiver_id}: {payoff}")
            receiver_key = helpers.get_receiver_key(r.receiver_session, r.receiver)
            receiver_data = player.session.receivers.get(receiver_key)
            receiver_data["allocations"].append(payoff)
            allocs = receiver_data["allocations"]
            print(f"All Allocations of R{r.receiver_id}: {allocs}")
            r.payoff = random.choice(allocs)
            print(f"Drew payoff for R{r.receiver_id}: {r.payoff}")

        if player.field_maybe_none("receivers_allocation") is None:
            helpers.set_dropout(player, C.DROPOUT_ERR_MSG)
            return

        receiver1, receiver2 = player.get_third_party_receivers()
        if receiver1.is_round_paid:
            set_payoff(receiver1, player.receivers_allocation)
        if receiver2.is_round_paid:
            set_payoff(receiver2, player.receivers_comm_account - player.receivers_allocation)


class MoralAppropiateness(AccountPage):
    form_model = "player"
    form_fields = [
        "allocation_reasoning",
        "moral_split_back",
        "moral_split_equal",
        "moral_split_correct",
        "moral_split_one",
    ]

    @staticmethod
    def vars_for_template(player: Player):
        partner = player.get_account_partner()
        receiver1, receiver2 = player.get_third_party_receivers()
        return dict(
            allocation_reasoning_label=f"How did you make your decisions of how to divide the common account between you and your matched worker {partner.name}, and between {receiver1.name} and {receiver2.name}? Please write which considerations came to your mind when you made both choices.",
            moral_split_back_label=f"Giving to {partner.name} the monetary contribution he/she produced.",
            moral_split_equal_label=f"Giving an equal amount to you and {partner.name}.",
            moral_split_correct_label=f"Splitting the amount according to the best guess of the number of correct answers of both participants.",
            task_name=helpers.get_task_name(player),
        )


class BetweenRounds(Page):
    @staticmethod
    def before_next_page(player: Player, timeout_happened):
        # TODO: ALWAYS PUT THIS ON THE LAST PAGE OF THE ROUND BEFORE ROUNDS 1/2/3 TO ENSURE THE CORRECT WAIT TIME IN THE MATCHING
        player.participant.wait_page_arrival = time.time()

    @staticmethod
    def is_displayed(player: Player):
        return player.round_number < C.NUM_ROUNDS


class SocialAppropiateness(EndogeneousPage):
    form_model = "player"
    form_fields = [
        "partner_selection",
        "social_split_back",
        "social_split_equal",
        "social_split_correct",
        "social_split_one",
    ]

    @staticmethod
    def vars_for_template(player: Player):
        return dict(
            partner_selection_label=f"Please state your reasons behind choosing your conversation partner in the third round.",
            social_split_back_label=f"Giving back to the other participant the monetary contribution he/she produced",
            social_split_equal_label=f"Giving an equal amount to both participants",
            social_split_correct_label=f"Splitting the amount according to the best guess of the number of correct answers of both participants",
        )

    # UNCOMMENT THIS TO INCENTIVIZE SOCIAL APPROPIATENESS QUESTIONS
    # @staticmethod
    # def before_next_page(player: Player, timeout_happened):
    #     selected_field, most_common_answer = random.choice(C.MOST_COMMON_SOCIAL_APP_ANSWERS)
    #     if player.field_maybe_none(selected_field) == most_common_answer:
    #         player.social_bonus = C.SOCIAL_APP_BONUS
    #         player.payoff += C.SOCIAL_APP_BONUS


class EndRounds(Page):
    @staticmethod
    def is_displayed(player: Player):
        return player.round_number == C.NUM_ROUNDS


class End(Page):
    @staticmethod
    def is_displayed(player: Player):
        return player.round_number == C.NUM_ROUNDS


page_sequence = [
    Matching,
    NoPartner,
    AccountPartner,
    EndogeneousInstructions,
    PickCommunicationPartner,
    SelectCommunicationPartner,
    CommunicationPartner,
    ChatWait,
    Chat,
    Allocation,
    MoralAppropiateness,
    BetweenRounds,
    SocialAppropiateness,
    EndRounds,
    End,
]