# Copyright (C) 2018-2022 by BDE ENS Paris-Saclay
# SPDX-License-Identifier: GPL-3.0-or-later

import time
from functools import lru_cache
from random import Random

from django import forms
from django.db import transaction
from django.db.models import Q
from django.utils.translation import gettext_lazy as _

from .base import WEISurvey, WEISurveyInformation, WEISurveyAlgorithm, WEIBusInformation
from ...models import WEIMembership

WORDS = [
	'ABBA', 'After', 'Alcoolique anonyme', 'Ambiance festive', 'Années 2000', 'Apéro', 'Art',
	'Baby foot billard biere pong', 'BBQ', 'Before', 'Bière pong', 'Bon enfant', 'Calme', 'Canapé',
	'Chanson paillarde', 'Chanson populaire', 'Chartreuse', 'Cheerleader', 'Chill', 'Choré',
	'Cinéma', 'Cocktail', 'Comédie musicle', 'Commercial', 'Copaing', 'Danse', 'Dancefloor',
	'Electro', 'Fanfare', 'Gin tonic', 'Inclusif', 'Jazz', "Jeux d'alcool", 'Jeux de carte',
	'Jeux de rôle', 'Jeux de société', 'JUL', 'Jus de fruit', 'Kfet', 'Kleptomanie assurée',
	'LGBTQ+', 'Livre', 'Morning beer', 'Musique', 'NAPS', 'Paillettes', 'Pastis', 'Paté Hénaff',
	'Peluche', 'Pena baiona', "Peu d'alcool", 'Pilier de bar', 'PMU', 'Poulpe', 'Punch', 'Rap',
	'Réveil', 'Rock', 'Rugby', 'Sandwich', 'Serge', 'Shot', 'Sociable', 'Spectacle', 'Techno',
	'Techno house', 'Thérapie Taxi', 'Tradition kchanaises', 'Troisième mi-temps', 'Turn up',
	'Vodka', 'Vodka pomme', 'Volley', 'Vomi stratégique'
]


class WEISurveyForm2022(forms.Form):
    """
    Survey form for the year 2022.
    Members choose 20 words, from which we calculate the best associated bus.
    """

    word = forms.ChoiceField(
        label=_("Choose a word:"),
        widget=forms.RadioSelect(),
    )

    def set_registration(self, registration):
        """
        Filter the bus selector with the buses of the current WEI.
        """
        information = WEISurveyInformation2022(registration)
        if not information.seed:
            information.seed = int(1000 * time.time())
            information.save(registration)
            registration._force_save = True
            registration.save()

        if self.data:
            self.fields["word"].choices = [(w, w) for w in WORDS]
            if self.is_valid():
                return

        rng = Random((information.step + 1) * information.seed)

        words = None

        buses = WEISurveyAlgorithm2022.get_buses()
        informations = {bus: WEIBusInformation2022(bus) for bus in buses}
        scores = sum((list(informations[bus].scores.values()) for bus in buses), [])
        average_score = sum(scores) / len(scores)

        preferred_words = {bus: [word for word in WORDS
                                 if informations[bus].scores[word] >= average_score]
                           for bus in buses}
        while words is None or len(set(words)) != len(words):
            # Ensure that there is no the same word 2 times
            words = [rng.choice(words) for _ignored2, words in preferred_words.items()]
        rng.shuffle(words)
        words = [(w, w) for w in words]
        self.fields["word"].choices = words


class WEIBusInformation2022(WEIBusInformation):
    """
    For each word, the bus has a score
    """
    scores: dict

    def __init__(self, bus):
        self.scores = {}
        for word in WORDS:
            self.scores[word] = 0.0
        super().__init__(bus)


class WEISurveyInformation2022(WEISurveyInformation):
    """
    We store the id of the selected bus. We store only the name, but is not used in the selection:
    that's only for humans that try to read data.
    """
    # Random seed that is stored at the first time to ensure that words are generated only once
    seed = 0
    step = 0

    def __init__(self, registration):
        for i in range(1, 21):
            setattr(self, "word" + str(i), None)
        super().__init__(registration)


class WEISurvey2022(WEISurvey):
    """
    Survey for the year 2022.
    """

    @classmethod
    def get_year(cls):
        return 2022

    @classmethod
    def get_survey_information_class(cls):
        return WEISurveyInformation2022

    def get_form_class(self):
        return WEISurveyForm2022

    def update_form(self, form):
        """
        Filter the bus selector with the buses of the WEI.
        """
        form.set_registration(self.registration)

    @transaction.atomic
    def form_valid(self, form):
        word = form.cleaned_data["word"]
        self.information.step += 1
        setattr(self.information, "word" + str(self.information.step), word)
        self.save()

    @classmethod
    def get_algorithm_class(cls):
        return WEISurveyAlgorithm2022

    def is_complete(self) -> bool:
        """
        The survey is complete once the bus is chosen.
        """
        return self.information.step == 20

    @classmethod
    @lru_cache()
    def word_mean(cls, word):
        """
        Calculate the mid-score given by all buses.
        """
        buses = cls.get_algorithm_class().get_buses()
        return sum([cls.get_algorithm_class().get_bus_information(bus).scores[word] for bus in buses]) / buses.count()

    @lru_cache()
    def score(self, bus):
        if not self.is_complete():
            raise ValueError("Survey is not ended, can't calculate score")

        bus_info = self.get_algorithm_class().get_bus_information(bus)
        # Score is the given score by the bus subtracted to the mid-score of the buses.
        s = sum(bus_info.scores[getattr(self.information, 'word' + str(i))]
                - self.word_mean(getattr(self.information, 'word' + str(i))) for i in range(1, 21)) / 20
        return s

    @lru_cache()
    def scores_per_bus(self):
        return {bus: self.score(bus) for bus in self.get_algorithm_class().get_buses()}

    @lru_cache()
    def ordered_buses(self):
        values = list(self.scores_per_bus().items())
        values.sort(key=lambda item: -item[1])
        return values

    @classmethod
    def clear_cache(cls):
        cls.word_mean.cache_clear()
        return super().clear_cache()


class WEISurveyAlgorithm2022(WEISurveyAlgorithm):
    """
    The algorithm class for the year 2022.
    We use Gale-Shapley algorithm to attribute 1y students into buses.
    """

    @classmethod
    def get_survey_class(cls):
        return WEISurvey2022

    @classmethod
    def get_bus_information_class(cls):
        return WEIBusInformation2022

    def run_algorithm(self, display_tqdm=False):
        """
        Gale-Shapley algorithm implementation.
        We modify it to allow buses to have multiple "weddings".
        """
        surveys = list(self.get_survey_class()(r) for r in self.get_registrations())  # All surveys
        surveys = [s for s in surveys if s.is_complete()]  # Don't consider invalid surveys
        # Don't manage hardcoded people
        surveys = [s for s in surveys if not hasattr(s.information, 'hardcoded') or not s.information.hardcoded]

        # Reset previous algorithm run
        for survey in surveys:
            survey.free()
            survey.save()

        non_men = [s for s in surveys if s.registration.gender != 'male']
        men = [s for s in surveys if s.registration.gender == 'male']

        quotas = {}
        registrations = self.get_registrations()
        non_men_total = registrations.filter(~Q(gender='male')).count()
        for bus in self.get_buses():
            free_seats = bus.size - WEIMembership.objects.filter(bus=bus, registration__first_year=False).count()
            # Remove hardcoded people
            free_seats -= WEIMembership.objects.filter(bus=bus, registration__first_year=True,
                                                       registration__information_json__icontains="hardcoded").count()
            quotas[bus] = 4 + int(non_men_total / registrations.count() * free_seats)

        tqdm_obj = None
        if display_tqdm:
            from tqdm import tqdm
            tqdm_obj = tqdm(total=len(non_men), desc="Non-hommes")

        # Repartition for non men people first
        self.make_repartition(non_men, quotas, tqdm_obj=tqdm_obj)

        quotas = {}
        for bus in self.get_buses():
            free_seats = bus.size - WEIMembership.objects.filter(bus=bus, registration__first_year=False).count()
            free_seats -= sum(1 for s in non_men if s.information.selected_bus_pk == bus.pk)
            # Remove hardcoded people
            free_seats -= WEIMembership.objects.filter(bus=bus, registration__first_year=True,
                                                       registration__information_json__icontains="hardcoded").count()
            quotas[bus] = free_seats

        if display_tqdm:
            tqdm_obj.close()

            from tqdm import tqdm
            tqdm_obj = tqdm(total=len(men), desc="Hommes")

        self.make_repartition(men, quotas, tqdm_obj=tqdm_obj)

        if display_tqdm:
            tqdm_obj.close()

        # Clear cache information after running algorithm
        WEISurvey2022.clear_cache()

    def make_repartition(self, surveys, quotas=None, tqdm_obj=None):
        free_surveys = surveys.copy()  # Remaining surveys
        while free_surveys:  # Some students are not affected
            survey = free_surveys[0]
            buses = survey.ordered_buses()  # Preferences of the student
            for bus, current_score in buses:
                if self.get_bus_information(bus).has_free_seats(surveys, quotas):
                    # Selected bus has free places. Put student in the bus
                    survey.select_bus(bus)
                    survey.save()
                    free_surveys.remove(survey)
                    break
                else:
                    # Current bus has not enough places. Remove the least preferred student from the bus if existing
                    least_preferred_survey = None
                    least_score = -1
                    # Find the least student in the bus that has a lower score than the current student
                    for survey2 in surveys:
                        if not survey2.information.valid or survey2.information.get_selected_bus() != bus:
                            continue
                        score2 = survey2.score(bus)
                        if current_score <= score2:  # Ignore better students
                            continue
                        if least_preferred_survey is None or score2 < least_score:
                            least_preferred_survey = survey2
                            least_score = score2

                    if least_preferred_survey is not None:
                        # Remove the least student from the bus and put the current student in.
                        # If it does not exist, choose the next bus.
                        least_preferred_survey.free()
                        least_preferred_survey.save()
                        free_surveys.append(least_preferred_survey)
                        survey.select_bus(bus)
                        survey.save()
                        free_surveys.remove(survey)
                        break
            else:
                raise ValueError(f"User {survey.registration.user} has no free seat")

            if tqdm_obj is not None:
                tqdm_obj.n = len(surveys) - len(free_surveys)
                tqdm_obj.refresh()