from otree.api import *
from itertools import cycle
from collections import defaultdict
from time import time

c = Currency

doc = """
These are the study instructions and this app also contains
code to set the participant role and between subjects condition.
"""


class Constants(BaseConstants):
    name_in_url = 'csr_grouping'
    players_per_group = None
    num_rounds = 1

    attitudes = ['pro', 'neutral', 'anti']
    macro_group_size = 12
    max_seconds_waiting = 60
    treatments = ['A', 'B', 'C', 'D']
    treatments_requirements = dict(
        A=dict(DM=[attitudes[0]], COO=[attitudes[0]]),
        B=dict(DM=[attitudes[0]], COO=attitudes),
        C=dict(DM=[attitudes[2]], COO=attitudes),
        D=dict(DM=[attitudes[2]], COO=[attitudes[0]]),
    )
    roles = ['DM', 'COO']


class Subsession(BaseSubsession):
    pass


class Group(BaseGroup):
    pass


class Player(BasePlayer):
    treatment = models.StringField()
    task_role = models.StringField()


# FUNCTIONS
def creating_session(subsession: Subsession):
    if subsession.round_number == 1:
        const = Constants
        session = subsession.session
        config = session.config
        treatments = config.get('treatments', const.treatments.copy())
        session.macro_group_size = config.get('macro_group_size', const.macro_group_size)
        session.max_seconds_waiting = config.get('max_seconds_waiting', const.max_seconds_waiting)
        session.treatments = treatments
        # The initial number of assignments per treatment can be altered to prioritize some treatments
        session.treatment_assignments = {t: 0 for t in treatments}

        session.previous_pairs = []

        for p in subsession.get_players():
            participant = p.participant
            participant.excluded = False
            participant.unmatched = False
            participant.dropout = False


"""
Matching phase:
1.	Grouping:
a.	Macro groups of P size (e.g. 20)
b.	Max waiting time is T
c.	Create group based on cell with min groups
d.	Try to form macro group for that cell
e.	If fail and time passed < T then wait, else take next cell and repeat
f.	Else assign macro group to cell
g.	If time passed >= T, Try to form macro group for any cell (in order of preference)
h.	Send macro group forward
"""


def players_by_arrival_time(waiting_players: list):
    return sorted(waiting_players, key=lambda p: p.participant.grouping_page_arrival)


def waiting_too_long(player: Player, max_seconds_waiting):
    participant = player.participant
    return time() - participant.grouping_page_arrival > max_seconds_waiting


def get_players_per_attitude(players: list) -> dict:
    ps_per_attitude = defaultdict(list)
    for p in players:
        ps_per_attitude[p.participant.attitude].append(p)
    return ps_per_attitude


def find_macro_group(treatment, players: list, const: Constants, macro_group_size: int) -> list:
    macro_group = []
    requirement = const.treatments_requirements[treatment]
    players_per_role = macro_group_size // len(const.roles)
    players_by_role = {const.roles[0]: [], const.roles[1]: []}
    for p in players:
        attitude = p.participant.attitude
        for role in players_by_role.keys():
            if len(players_by_role[role]) >= players_per_role:
                continue
            if attitude in requirement[role] and all(p not in r for r in players_by_role.values()):
                players_by_role[role].append(p)

    if all(len(ps) >= players_per_role for ps in players_by_role.values()):
        macro_group = list(zip(*players_by_role.values()))
    return macro_group


def assign_to_treatment(macro_group: list, const: Constants, session: Subsession.session, treatment):
    roles = cycle(const.roles)
    for pair in macro_group:
        for p in pair:
            role = next(roles)
            p.participant.role = role
            p.task_role = role
            p.participant.treatment = treatment
            p.treatment = treatment
    session.treatment_assignments[treatment] += 1


def set_macro_group(players: list, const: Constants, session: Subsession.session, waited_too_long=False) -> list:
    macro_group = []
    treatments = sorted(session.treatments, key=lambda treatment: session.treatment_assignments[treatment])
    for n, t in enumerate(treatments):
        macro_group.extend(find_macro_group(t, players, const, session.macro_group_size))
        if macro_group:
            assign_to_treatment(macro_group, const, session, t)
        if macro_group or not waited_too_long:
            break
    return macro_group


def group_by_arrival_time_method(subsession: Subsession, waiting_players):
    const = Constants
    session = subsession.session
    max_seconds_waiting = session.max_seconds_waiting
    ps = players_by_arrival_time(waiting_players)
    macro_group = set_macro_group(ps, const, session)
    if macro_group:
        macro_group = [p for pair in macro_group for p in pair]
        return macro_group

    for p in ps:
        if waiting_too_long(p, max_seconds_waiting):
            macro_group = set_macro_group(ps, const, session, True)
            if macro_group:
                macro_group = [p for pair in macro_group for p in pair]
                return macro_group
            p.participant.excluded = True
            return [p]


# PAGES
class GroupingPage(WaitPage):
    group_by_arrival_time = True

    title_text = "Please Wait Here"
    body_text = "Kindly stay on this tab and DO NOT navigate to another tab/screen. Thank you!"
    @staticmethod
    def app_after_this_page(player: Player, upcoming_apps):
        if player.participant.excluded:
            return upcoming_apps[-1]


page_sequence = [GroupingPage]