# Copyright (C) 2018-2021 by BDE ENS Paris-Saclay
# SPDX-License-Identifier: GPL-3.0-or-later
from argparse import ArgumentParser, FileType

from django.core.management import BaseCommand
from django.db import transaction
from django.db.models import Q

from ...forms import CurrentSurvey
from ...models import Bus, WEIMembership


class Command(BaseCommand):
    help = "Attribute to each first year member a team for the WEI"

    def add_arguments(self, parser: ArgumentParser):
        parser.add_argument('--doit', '-d', action='store_true', help='Finally run the algorithm in non-dry mode.')
        parser.add_argument('--output', '-o', nargs='?', type=FileType('w'), default=self.stdout,
                            help='Output file for the algorithm result. Default is standard output.')

    def attribute_teams(self, bus):
        if not bus.size:
            return

        teams = bus.teams.all()

        old_members = WEIMembership.objects.filter(registration__first_year=False, bus=bus, team__in=teams)
        new_members = WEIMembership.objects.filter(registration__first_year=True, bus=bus)

        n_chef_tot = old_members.count()
        n_equipe = teams.count()
        size_goal = {}
        for team in teams:
            size_goal[team] = int(old_members.filter(team=team).count() * new_members.count() / n_chef_tot)
        i = 0
        non_men = list(new_members.filter(~Q(registration__gender='male')))
        men = list(new_members.filter(registration__gender='male'))

        print(bus.name, size_goal.values())
        return

        while non_men:
            if new_members.filter(team=teams[i]).count() < size_goal[teams[i]]:
                non_man = non_men.pop()
                non_man.team = teams[i]
                non_man.save()
            i += 1
            i %= n_equipe

        for i, team in enumerate(teams):
            print(i, new_members.filter(team=team).count())
            if new_members.filter(team=team).count() == 1:
                # Si une fille est seule on l'enlève.
                g_alone = new_members.get(team=team)
                g_alone.team = teams[i + 1]
                g_alone.save()
        remain = True
        i = 0
        while men and remain:
            if new_members.filter(team=teams[i]).count() < size_goal[teams[i]]:
                m = men.pop()
                m.team = teams[i]
                m.save()
            i += 1
            i %= n_equipe

    @transaction.atomic
    def handle(self, *args, **options):
        """
        Run the WEI algorithm to attribute a bus to each first year member.
        """
        sid = transaction.savepoint()

        algorithm = CurrentSurvey.get_algorithm_class()()

        try:
            from tqdm import tqdm
            display_tqdm = True
        except ImportError:
            display_tqdm = False

        for bus in Bus.objects.all():
            self.attribute_teams(bus)

        raise Exception()

        output = options['output']
        registrations = algorithm.get_registrations()
        per_bus = {bus: [r for r in registrations if 'selected_bus_pk' in r.information
                         and r.information['selected_bus_pk'] == bus.pk]
                   for bus in algorithm.get_buses()}
        for bus, members in per_bus.items():
            output.write(bus.name + "\n")
            output.write("=" * len(bus.name) + "\n")
            order = -1
            for r in members:
                survey = CurrentSurvey(r)
                for order, (b, _score) in enumerate(survey.ordered_buses()):
                    if b == bus:
                        break
                output.write(f"{r.user.username} ({order + 1})\n")
            output.write("\n")

        if not options['doit']:
            self.stderr.write(self.style.WARNING("Running in dry mode. "
                                                 "Use --doit option to really execute the algorithm."))
            transaction.savepoint_rollback(sid)
            return