from google.appengine.ext import ndb

from models.grading import SITTING_ROOM_TYPES
from models.team import (
    Team,
    Contestant,
)
from models.sitting import Sitting
import random

SINGLE_TEST_SITTING = {
    'alg' : 'alg,at',
    'at' : 'at,calc',
    'calc' : 'calc,geo',
    'geo' : 'at,geo',
    'team' : 'power,team',
    'power' : 'power,team'
}

class RoomAssignment(ndb.Model):
    testing_id = ndb.StringProperty(required=True)
    sitting_nid = ndb.IntegerProperty(required=True)

    @classmethod
    def assign_team(cls, team_id):
        team = Team.get_by_team_id(team_id)
        contestants = Contestant.fetch_by_team(team.key)

        cls.delete_team_assignments(team.key)

        def get_sitting_type(tests):
            if len(tests) == 1 and tests[0] != 'gen':
                return SINGLE_TEST_SITTING[tests[0]]
            return ','.join(sorted(list(tests)))

        def select_sitting(sitting_type):
            sittings = Sitting.fetch_by_exam(sitting_type)
            weights = [s.capacity for s in sittings]
            total = sum(weights)
            index = random.randint(1, total)

            counter = 0
            i = 0
            while (counter < index):
                counter += weights[i]
                i += 1

            return sittings[i-1]

        def assign_to_sitting(testing_id, tests, size):
            sitting_type = get_sitting_type(tests)
            sitting = select_sitting(sitting_type)

            assignment = RoomAssignment(
                testing_id=testing_id,
                sitting_nid=sitting.nid,
                parent=team.key,
            )
            assignment.put()

        if len(team.team_tests) > 0:
            assign_to_sitting(str(team_id), team.team_tests, len(contestants))

        for c in contestants:
            if len(c.tests) == 0: continue
            assign_to_sitting(c.contestant_id, c.tests, 1)

    @staticmethod
    def get_assigned_team_ids():
        team_ids = set()
        all_room_assignments = RoomAssignment.query().fetch()
        for ra in all_room_assignments:
            try:
                team_id = int(ra.testing_id)
                team_ids.add(team_id)
            except ValueError as e:
                continue
        return list(team_ids)

    @staticmethod
    def delete_all():
        ndb.delete_multi(RoomAssignment.query().iter(keys_only=True))
    
    @staticmethod
    def delete_team_assignments(team_key):
        for a in RoomAssignment.query(ancestor=team_key).fetch():   
            a.key.delete()

    @staticmethod
    def fetch_by_team(team_key):
        return RoomAssignment.query(ancestor=team_key).fetch()

    @staticmethod
    def fetch_all():
        return RoomAssignment.query().fetch()
