from treatments import *


# Match players and identify the relevant partner identity treatment
def set_partner_based_on_identity(still_waiting_for_task, task, num_players):
    # Look for group in condition other_region or other_school
    for t in ['region', 'school']:
        s = set(p.participant.vars['data'].get(t) for p in still_waiting_for_task if p.participant.vars.get('data'))
        if len(s) == num_players:
            i = [s.pop() for _ in range(num_players)]
            group = [next(p for p in still_waiting_for_task if p.participant.vars['data'].get(t) == j) for j in i]
            if t == 'region':
                partner = 'other_region'
            else:
                partner = 'other_school'
            for p in group:
                p.participant.vars['partner'] = partner
                setattr(p, task + '_partner', partner)
            return group
    # If region and school conditions aren't met, then randomize to other three conditions: class, name, ctrl
    group = still_waiting_for_task[:num_players]
    partner = next(treatments['partner_interactive_tasks'])
    # Skip 'name' treatment if participants' data are not available (e.g. when testing)
    if partner == 'name' and not len(group) == len([p for p in group if p.participant.vars.get('data')]):
        partner = next(treatments['partner_interactive_tasks'])
    for p in group:
        p.participant.vars['partner'] = partner
        setattr(p, task + '_partner', partner)
    return group