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

import time
import json
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 django.utils.safestring import mark_safe

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

WORDS = {
    'list': [
        'Fiesta', 'Graillance', 'Move it move it', 'Calme', 'Nerd et geek', 'Jeux de rôles et danse rock',
        'Strass et paillettes', 'Spectaculaire', 'Splendide', 'Flow inégalable', 'Rap', 'Battles légendaires',
        'Techno', 'Alcool', 'Kiffeur·euse', 'Rugby', 'Médiéval', 'Festif',
        'Stylé', 'Chipie', 'Rétro', 'Vache', 'Farfadet', 'Fanfare',
    ],
    'questions': {
        "alcool": [
            """Sur une échelle allant de 0 (= 0 alcool ou très peu) à 5 (= la fontaine de jouvence alcoolique),
            quel niveau de consommation d’alcool souhaiterais-tu ?""",
            {
                42: 4,
                47: 1,
                48: 3,
                45: 3.5,
                44: 4,
                46: 5,
                43: 3,
                49: 3
            }
        ],
        "voie_post_bac": [
            """Si la DA du bus de ton choix correspondait à une voie post-bac, laquelle serait-elle ?""",
            {
                42: "Double licence cuisine/arts du cirque option burger",
                47: "BTS Exploration de donjon",
                48: "Ecole des stars en herbe",
                45: "Déscolarisation précoce",
                44: "Rattrapage pour excès de kiff",
                46: "Double cursus STAPS / Licence d’histoire",
                43: "Recherche active d’un sugar daddy/d’un sugar mommy",
                49: "Licence de musicologie"
            }
        ],
        "boite": [
            """Tu es seul·e sur une île déserte et devant toi il y a une sombre boîte de taille raisonnable.
            Qu’y a-t-il à l’intérieur ?""",
            {
                42: "Un burgouzz de valouzz",
                47: "Un ocarina (pour me téléporter hors de ce bourbier)",
                48: "Des paillettes, un micro de karaoké et une enceinte bluetooth",
                45: "Un kebab",
                44: "Une 86 et un caisson pour taper du pied",
                46: "Une épée, un ballon et une tireuse",
                43: "Des lunettes de soleil",
                49: "Mon instrument de musique"
            }
        ],
        "tardif": [
            """Il est 00h, tu as passé la journée à la plage avec tes copains et iels te proposent de prolonger parce
            qu’après tout, il n’y a plus personne sur la plage à cette heure-ci. Tu n’habites pas loin mais t’enchaînes
            demain avec une journée similaire avec un autre groupe d’amis parce que t’es trop #busy. Que fais-tu ?""",
            {
                42: "On veut se déchaîner toute la nuit !!",
                47: "Je prends une glace et chill un moment avant d’aller dormir",
                48: "J’enfile mes boogie shoes pour enflammer le dancefloor avec elleux et lancer un concours de slay, le perdant finit la bouteille de rhum",
                45: "La fête continuuuuue",
                44: "Soirée sangria plage → boîte → lever de soleil sur la plage",
                46: "Minuit ? C’est l’heure du genepi. On commence les alcools forts !!",
                43: "T’enchaînes direct (faut pas les priver de ta présence)",
                49: "On continue en mode chill (soirée potins)"
            }
        ],
        "cohesion": [
            """C’est la rentrée de Seconde et tu découvres ta classe, tes camarades et ta prof principale!!!
            qui vous propose une activité de cohésion. Laquelle est-elle ?""",
            {
                42: "Un relais cubi en ventriglisse",
                47: "Un jeu de rôle",
                48: "Organiser la soirée de l’année dans le lycée. Le thème : SLAY (Spotlight, Love, Amaze/All-night, Yeah), paillettes, disco",
                45: "La prof de français propose un slam parce qu'elle pense que c'est du rap littéraire qui fera plaisir aux élèves",
                44: "P’tit escape game + apéro",
                46: "Joute avec des boucliers Gilbert",
                43: "Tournage d’un clip de confessions nocturnes de Diam’s",
                49: "Je sais pas j’ai raté mon BAFA"
            }
        ],
        "artiste": [
            """C’est l’été et la saison des festivals a commencé. Tu regardes la programmation du festival
            pas loin de chez toi et tu découvres avec joie la présence d’un·e artiste. De qui s’agit-il ?""",
            {
                42: "Moto-Moto (il chantera son fameux tube “je les aime grosses, je les aime bombées”)",
                47: "Hatsune Miku",
                48: "Rihanna",
                45: "Vald",
                44: "Qui connaît vraiment les noms des artistes de tech ?",
                46: "Perceval",
                43: "Fatal bazooka",
                49: "Måneskin"
            }
        ],
        "annonce_noel": [
            """C’est Noël et tu revois toute ta famille, oncles, tantes, cousin·e·s, grands-parents, la totale.
            D’un coup, tu te lèves, tapotes de manière pompeuse sur ton verre avec un de tes couverts.
            Qu’annonces-tu ?""",
            {
                42: """« Chère famille. Je sais bien que nous avions dit : pas de politique à table.
                    Je ne peux toutefois me retenir de vous annoncer une grande nouvelle…
                    j’ai décidé de quitter la ville pour consacrer ma vie au culte du Roi Julian.
                    A moi la jungle luxuriante, là où le soleil chaud caresse les palmiers,
                    où les lémuriens dansent avec frénésie et où chaque repas est une ode au burger sauvage.
                    Longue vie à Sa Majesté le Roi Julian ! »""",
                47: "« J’ai perdu »",
                48: "« Mes chers parents je pars, j’arrête l’ENS pour devenir DJ slay à Ibiza »",
                45: "J’interromps le repas pour rapper les 6min de bande organisée",
                44: "« Digestif ? Pétanque ? Les deux ? »",
                46: "« Montjoie St Denis à bas la Macronie »",
                43: "« Je suis enceinte » (c’est faux j’ai juste besoin d’attention)",
                49: """Discours de remerciement :
                    je lance un powerpoint de 65 slides et sors une feuille A4 blanche (je fais semblant de lire mon discours dessus)"""
            }
        ],
        "vacances": [
            """Les vacances sont là et t’aimerais bien partir quelque part, mais où ?""",
            {
                42: "A Madagascar, à bord d’un bus conduit par des pingouins",
                47: "Dans ma chambre",
                48: "Rio de Janeiro",
                45: "N'importe où tant qu'on peut sortir tous les soirs",
                44: "Tu suis les plans du club ski ou de piratens",
                46: "Carcassonne",
                43: "Coachella",
                49: "Dans les montagnes de la république populaire d’Auvergne-Rhônes-Alpes pour profiter de la fraîcheur, de la nature et de mes ami·e·s"
            }
        ],
        "loisir": [
            """T’as fini ta journée de cours et tu t’apprêtes à profiter d’une activité/hobby/loisir de ton choix.
            Laquelle est-ce ?""",
            {
                42: "Cueillir des noix de coco",
                47: "Essayer de travailler puis chill avec des potes autour d’un jeu en buvant du thé",
                48: "Repet du nouveau spectacle de mon club, before (potins) puis sortie avec les potes jusqu’au bout de la night",
                45: "Zoner avec les copaings jusqu’à pas d’heure",
                44: "Go Kfet pour se faire traquenard jusqu’à 3h du mat",
                46: "Déterminer ce qui est le plus solide entre mon crâne et une ecocup",
                43: "Revoir pour la 6e fois gossip girl au fond de ton lit",
                49: "Jouer de mon instrument préféré avec les copains/copines pour préparer le prochain concert #solidays"
            }
        ],
        "plan": [
            """Tu reçois un message sur la conversation de groupe que tu partages avec tes potes :
            vous êtes chaud·e·s pour vous retrouver. Quel plan t’attire le plus ?""",
            {
                42: """Après-midi piscine, puis before arrosé de mojito,
                    avant d’aller s’éclater en pot avec toute la savane et de finir sur un after spécial pina colada""",
                47: """(matin) : Ptit jeu de rôle
                    (repas) : le traditionnel poké-tacos
                    (juste après le repas) : combat avec des épées en mousse avec les copains!
                    (16h00) : pause thé
                    (fin d’après midi) : initiation à la danse rock
                    (soirée) : découverte d’un jeu de société avec des règles obscures
                    """,
                48: "Soirée champagne and chic : spectacle et dîner au moulin rouge puis soirée sur les champs",
                45: "Se regrouper pour une soirée, même si il n’est encore que 10h",
                44: "P’tit poké qui termine en koin koin avec after poker",
                46: "Une dégustation de bière, un rugby et toute autre activité joviale",
                43: "Un brunch de pour papoter puis friperies",
                49: "Soirée raclette !"
            }
        ]
    },
    'stats': [
        {
            "question": """Le WEI est structuré par bus, et au sein de chaque bus, par équipes.
                         Pour toi, être dans une équipe où tout le monde reste sobre (primo-entrants comme encadrants) c'est :""",
            "answers": [
                (1, "Inenvisageable"),
                (2, "À contre cœur"),
                (3, "Pourquoi pas"),
                (4, "Souhaitable"),
                (5, "Nécessaire"),
            ],
            "help_text": "(De toute façon aucun alcool n'est consommé pendant les trajets du bus, ni aller, ni retour.)",
        },
        {
            "question": "Faire partie d'un bus qui n'apporte pas de boisson alcoolisée pour ses membres, pour toi c'est :",
            "answers": [
                (1, "Inenvisageable"),
                (2, "À contre cœur"),
                (3, "Pourquoi pas"),
                (4, "Souhaitable"),
                (5, "Nécessaire"),
            ],
            "help_text": """(Tout les bus apportent de l'alcool cette année, cette question sert à l'organisation pour l'année prochaine.
                         De plus il y aura de toute façon de l'alcool commun au WEI et aucun alcool n'est consommé pendant les trajets en bus.)""",
        },
    ]
}

IMAGES = {
    "vacances": {
        49: "/static/wei/img/logo_auvergne_rhone_alpes.jpg",
    }
}

NB_WORDS = 5


class OptionalImageRadioSelect(forms.RadioSelect):
    def __init__(self, images=None, *args, **kwargs):
        self.images = images or {}
        super().__init__(*args, **kwargs)

    def create_option(self, name, value, label, selected, index, subindex=None, attrs=None):
        option = super().create_option(name, value, label, selected, index, subindex=subindex, attrs=attrs)
        img_url = self.images.get(value)
        if img_url:
            option['label'] = mark_safe(f'{label} <img src="{img_url}" style="height:32px;vertical-align:middle;">')
        else:
            option['label'] = label
        return option


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

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

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

        if information.step == 0:
            self.fields["words"] = forms.MultipleChoiceField(
                label=_(f"Select {NB_WORDS} words that describe the WEI experience you want to have."),
                choices=[(w, w) for w in WORDS['list']],
                widget=forms.CheckboxSelectMultiple(),
                required=True,
            )
            if self.is_valid():
                return

            all_preferred_words = WORDS['list']
            rng.shuffle(all_preferred_words)
            self.fields["words"].choices = [(w, w) for w in all_preferred_words]
        elif information.step <= len(WORDS['questions']):
            questions = list(WORDS['questions'].items())
            idx = information.step - 1
            if idx < len(questions):
                q, (desc, answers) = questions[idx]
                if q == 'alcool':
                    choices = [(i / 2, str(i / 2)) for i in range(11)]
                else:
                    choices = [(k, v) for k, v in answers.items()]
                    rng.shuffle(choices)
                self.fields[q] = forms.ChoiceField(
                    label=desc,
                    choices=choices,
                    widget=OptionalImageRadioSelect(images=IMAGES.get(q, {})),
                    required=True,
                )
        elif information.step == len(WORDS['questions']) + 1:
            for i, v in enumerate(WORDS['stats']):
                self.fields[f'stat_{i}'] = forms.ChoiceField(
                    label=v['question'],
                    choices=v['answers'],
                    widget=forms.RadioSelect(),
                    required=False,
                    help_text=_(v.get('help_text', ''))
                )

    def clean_words(self):
        data = self.cleaned_data['words']
        if len(data) != NB_WORDS:
            raise forms.ValidationError(_(f"Please choose exactly {NB_WORDS} words"))
        return data


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

    def __init__(self, bus):
        self.scores = {}
        super().__init__(bus)


class BusInformationForm2025(forms.ModelForm):
    class Meta:
        model = Bus
        fields = ['information_json']
        widgets = {
            'information_json': forms.HiddenInput(),
        }

    def __init__(self, *args, words=None, **kwargs):
        super().__init__(*args, **kwargs)

        initial_scores = {}
        if self.instance and self.instance.information_json:
            try:
                info = json.loads(self.instance.information_json)
                initial_scores = info.get("scores", {})
            except (json.JSONDecodeError, TypeError, AttributeError):
                initial_scores = {}
        if words is None:
            words = WORDS['list']
        self.words = words

        choices = [(i, str(i)) for i in range(6)]  # [(0, '0'), (1, '1'), ..., (5, '5')]
        for word in words:
            self.fields[word] = forms.TypedChoiceField(
                label=word,
                choices=choices,
                coerce=int,
                initial=initial_scores.get(word, 0) if word in initial_scores else None,
                required=True,
                widget=forms.RadioSelect,
                help_text=_("Rate between 0 and 5."),
            )

    def clean(self):
        cleaned_data = super().clean()
        scores = {}
        for word in self.words:
            value = cleaned_data.get(word)
            if value is not None:
                scores[word] = value
        # On encode en JSON
        cleaned_data['information_json'] = json.dumps({"scores": scores})
        return cleaned_data


class WEISurveyInformation2025(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, NB_WORDS + 1):
            setattr(self, "word" + str(i), None)
        for q in WORDS['questions']:
            setattr(self, q, None)
        super().__init__(registration)

    def reset(self, registration):
        """
        Réinitialise complètement le questionnaire : step, seed, mots choisis et réponses aux questions.
        """
        self.step = 0
        self.seed = 0
        for i in range(1, NB_WORDS + 1):
            setattr(self, f"word{i}", None)
        for q in WORDS['questions']:
            setattr(self, q, None)
        self.save(registration)
        registration._force_save = True
        registration.save()


class WEISurvey2025(WEISurvey):
    """
    Survey for the year 2025.
    """

    @classmethod
    def get_year(cls):
        return 2025

    @classmethod
    def get_survey_information_class(cls):
        return WEISurveyInformation2025

    def get_form_class(self):
        return WEISurveyForm2025

    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):
        if self.information.step == 0:
            words = form.cleaned_data['words']
            for i, word in enumerate(words, 1):
                setattr(self.information, "word" + str(i), word)
            self.information.step += 1
            self.save()
        elif 1 <= self.information.step <= len(WORDS['questions']):
            questions = list(WORDS['questions'].keys())
            idx = self.information.step - 1
            if idx < len(questions):
                q = questions[idx]
                setattr(self.information, q, form.cleaned_data[q])
                self.information.step += 1
                self.save()
        else:
            for i, __ in enumerate(WORDS['stats']):
                ans = form.cleaned_data.get(f'stat_{i}')
                if ans is not None:
                    setattr(self.information, f'stat_{i}', ans)
            self.information.step += 1
            self.save()

    @classmethod
    def get_algorithm_class(cls):
        return WEISurveyAlgorithm2025

    def is_complete(self) -> bool:
        """
        The survey is complete once the bus is chosen.
        """
        return self.information.step > len(WORDS['questions']) + 1

    @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_questions(self, bus):
        """
        The score given by the answers to the questions
        """
        if not self.is_complete():
            raise ValueError("Survey is not ended, can't calculate score")
        s = sum(1 for q in WORDS['questions'] if q != 'alcool' and getattr(self.information, q) == bus.pk)
        if 'alcool' in WORDS['questions'] and bus.pk in WORDS['questions']['alcool'][1] and hasattr(self.information, 'alcool'):
            s -= abs(float(self.information.alcool) - float(WORDS['questions']['alcool'][1][bus.pk]))
        return s

    @lru_cache()
    def score_words(self, bus):
        """
        The score given by the choice of words
        """
        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, 1 + NB_WORDS)) / self.get_algorithm_class().get_buses().count()
        return s

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

    @lru_cache()
    def ordered_buses(self):
        """
        Order the buses by the score_questions of the survey.
        """
        values = list(self.scores_per_bus().items())
        values.sort(key=lambda item: -item[1][0])
        return values

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


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

    @classmethod
    def get_survey_class(cls):
        return WEISurvey2025

    @classmethod
    def get_bus_information_class(cls):
        return WEIBusInformation2025

    @classmethod
    def get_bus_information_form(cls):
        return BusInformationForm2025

    @classmethod
    def get_buses(cls):

        if not hasattr(cls, '_buses'):
            cls._buses = Bus.objects.filter(wei__year=cls.get_survey_class().get_year(), size__gt=0).all().exclude(name='Staff')
        return cls._buses

    def run_algorithm(self, display_tqdm=False):
        """
        Gale-Shapley algorithm implementation.
        We modify it to allow buses to have multiple "weddings".
        We use lexigographical order on both scores
        """
        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
        WEISurvey2025.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_scores 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_words(bus)
                        if current_scores[1] <= 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()
