# Copyright (C) 2018-2023 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 = {
    "ambiance":["Question", {1:"réponse 1", 2:"réponse 2", 3:"réponse 3", 4:"réponse 4", 5:"réponse 5"}],
    "musique":["Question", {1:"réponse 1", 2:"réponse 2", 3:"réponse 3", 4:"réponse 4", 5:"réponse 5"}],
    "boisson":["Question", {1:"réponse 1", 2:"réponse 2", 3:"réponse 3", 4:"réponse 4", 5:"réponse 5"}],
    "beauferie":["Question", {1:"réponse 1", 2:"réponse 2", 3:"réponse 3", 4:"réponse 4", 5:"réponse 5"}],
    "sommeil":["Question", {1:"réponse 1", 2:"réponse 2", 3:"réponse 3", 4:"réponse 4", 5:"réponse 5"}],
    "vacances":["Question", {1:"réponse 1", 2:"réponse 2", 3:"réponse 3", 4:"réponse 4", 5:"réponse 5"}],
    "activite":["Question", {1:"réponse 1", 2:"réponse 2", 3:"réponse 3", 4:"réponse 4", 5:"réponse 5"}],
    "hygiene":["Question", {1:"réponse 1", 2:"réponse 2", 3:"réponse 3", 4:"réponse 4", 5:"réponse 5"}],
    "animal":["Question", {1:"réponse 1", 2:"réponse 2", 3:"réponse 3", 4:"réponse 4", 5:"réponse 5"}],
    "fensfoire":["Question", {1:"réponse 1", 2:"réponse 2", 3:"réponse 3", 4:"réponse 4", 5:"réponse 5"}],
    "kokarde":["Question", {1:"réponse 1", 2:"réponse 2", 3:"réponse 3", 4:"réponse 4", 5:"réponse 5"}],
    "copain":["Question", {1:"réponse 1", 2:"réponse 2", 3:"réponse 3", 4:"réponse 4", 5:"réponse 5"}],
    "vie":["Question", {1:"réponse 1", 2:"réponse 2", 3:"réponse 3", 4:"réponse 4", 5:"réponse 5"}],
    "jeux":["Question", {1:"réponse 1", 2:"réponse 2", 3:"réponse 3", 4:"réponse 4", 5:"réponse 5"}],
    "calin":["Question", {1:"réponse 1", 2:"réponse 2", 3:"réponse 3", 4:"réponse 4", 5:"réponse 5"}],
    "vommi":["Question", {1:"réponse 1", 2:"réponse 2", 3:"réponse 3", 4:"réponse 4", 5:"réponse 5"}],
    "kfet":["Question", {1:"réponse 1", 2:"réponse 2", 3:"réponse 3", 4:"réponse 4", 5:"réponse 5"}],
    "fatigue":["Question", {1:"réponse 1", 2:"réponse 2", 3:"réponse 3", 4:"réponse 4", 5:"réponse 5"}],
    "duree trajet":["Question", {1:"réponse 1", 2:"réponse 2", 3:"réponse 3", 4:"réponse 4", 5:"réponse 5"}],
    "scolarite":["Question", {1:"réponse 1", 2:"réponse 2", 3:"réponse 3", 4:"réponse 4", 5:"réponse 5"}]
}


class WEISurveyForm2023(forms.Form):
    """
    Survey form for the year 2023.
    Members answer 20 question, from which we calculate the best associated bus.
    """

    def __init__(self,**kwargs):
        super().__init__(**kwargs)
        for question in WORDS:
            self.fields[question] = forms.ChoiceField(
                label=WORDS[question][0]+question,
                widget=forms.RadioSelect(),
            )

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

#        if self.data:
#            for question in WORDS:
#                self.fields[question].choices = [answer for answer in WORDS[question][1]]
#            if self.is_valid():
#                return

#        rng = Random(information.seed) # add someting if it's alwais the same
#        questions = list(WORDS.keys())
#        rng.shuffle(questions)

#        for question in questions:
        for question in WORDS:
            answers = [(answer, WORDS[question][1][answer]) for answer in WORDS[question][1]]
            self.fields[question].choices = answers


class WEIBusInformation2023(WEIBusInformation):
    """
    For each question, the bus has ordered answers
    """
    scores: dict

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


class WEISurveyInformation2023(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

    def __init__(self, registration):
        for question in WORDS:
            setattr(self, str(question), None)
        super().__init__(registration)


class WEISurvey2023(WEISurvey):
    """
    Survey for the year 2023.
    """

    @classmethod
    def get_year(cls):
        return 2023

    @classmethod
    def get_survey_information_class(cls):
        return WEISurveyInformation2023

    def get_form_class(self):
        return WEISurveyForm2023

    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):
        for question in WORDS:
            answer = form.cleaned_data[question]
            setattr(self.information, question, answer)
        self.save()

    @classmethod
    def get_algorithm_class(cls):
        return WEISurveyAlgorithm2023

    def is_complete(self) -> bool:
        """
        The survey is complete once the bus is chosen.
        """
        for question in WORDS:
            if not getattr(self.information, question):
                return False
        return True

    @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 WEISurveyAlgorithm2023(WEISurveyAlgorithm):
    """
    The algorithm class for the year 2023.
    We use Gale-Shapley algorithm to attribute 1y students into buses.
    """

    @classmethod
    def get_survey_class(cls):
        return WEISurvey2023

    @classmethod
    def get_bus_information_class(cls):
        return WEIBusInformation2023

    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
        WEISurvey2023.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()