"""Module holding redis models for the ``surveys`` app."""

from datetime import datetime
from itertools import chain
import logging

from redis import WatchError

from limpyd.fields import PKField, SortedSetField
from limpyd_extensions.dynamic.fields import DynamicSortedSetField
from limpyd_extensions.dynamic.model import ModelWithDynamicFieldMixin

from django.utils.functional import cached_property

from mtp.apps.common.limpyd_base import BaseModel, DjangoBackedRedisModel
from mtp.apps.stats.helpers import day_range

logger = logging.getLogger('mtp.surveys')


class BaseStatModel(ModelWithDynamicFieldMixin, BaseModel):
    """Base model for surveys stats.

    Used to pre-compute stats for question (see ``QuestionData``) or global satisfaction (see
    ``GlobalSatisfaction``) on a daily basis for different kind of data:
    - for a session
    - for a training
    - for a trainer
    - for a couple (training, trainer)

    Attributes
    ----------
    answers_for_session : DynamicSortedSetField
        A SortedSet for each session (session pk).
        Keys are a token representing session+user+question (see ``get_answer_key``)
        Score are the answered value.
        It allows to easily get the number of answer, and the average (sum of all values / number
        of answer)
        There is no notion of day here because a session is a one shot. So it's stored differently
        than all other precomputed data.
    nb_by_day : SortedSetField
        A SortedSet storing the number of answer (score) for each day (key)
    avg_by_day : SortedsetField
        A SortedSet storing the average answer value (score) for each day (key)
    nb_by_day_for_trainer : DynamicSortedSetField
        A SortedSet for each trainer (trainer pk), storing the number of answer (score) for each
        day (key)
    avg_by_day_for_trainer : DynamicSortedSetField
        A SortedSet for each trainer (trainer pk), storing the average answer value (score) for each
        day (key)
    nb_by_day_for_training : DynamicSortedSetField
        A SortedSet for each training (training pk), storing the number of answer (score) for each
        day (key)
    avg_by_day_for_training : DynamicSortedSetField
        A SortedSet for each training (training pk), storing the average answer value (score) for
        each day (key)
    nb_by_day_for_training_and_trainer : DynamicSortedSetField
        A SortedSet for each couple (training, trainer) (see ``get_training_and_trainer_key``),
        storing the number of answer (score) for each day (key)
    avg_by_day_for_training_and_trainer : DynamicSortedSetField
        A SortedSet for each couple (training, trainer) (see ``get_training_and_trainer_key``),
        storing the average answer value (score) for each day (key)

    Example
    ------

    >>> session_zset = self.answers_for_session(session.pk)
    >>> key = self.get_answer_key(session.pk, user.pk, question.pk)
    >>> session_zset.zadd(answer_value, key)

    >>> self.nb_by_day.zadd(nb_answers_for_this_day, day)
    >>> self.avg_by_day.zadd(avg_of_answers_for_this_day, day)

    >>> trainer_zset_nb = self.nb_by_day_for_trainer(trainer.pk)
    >>> trainer_zset_nb.zadd(nb_answers_for_this_trainer_for_the_day, day)
    >>> trainer_zset_avg = self.avg_by_day_for_trainer(trainer.pk)
    >>> trainer_zset_avg.zadd(avg_of_answers_for_this_trainer_for_the_day, day)

    >>> training_zset_nb = self.nb_by_day_for_training(training.pk)
    >>> training_zset_nb.zadd(nb_answers_for_this_training_for_the_day, day)
    >>> training_zset_avg = self.avg_by_day_for_training(training.pk)
    >>> training_zset_avg.zadd(avg_of_answers_for_this_training_for_the_day, day)

    >>> train_zset_nb = self.nb_by_day_for_training_and_trainer(train.pk)
    >>> train_zset_nb.zadd(nb_answers_for_this_train_for_the_day, day)
    >>> train_zset_avg = self.avg_by_day_for_training_and_trainer(train.pk)
    >>> train_zset_avg.zadd(avg_of_answers_for_this_train_for_the_day, day)

    Notes
    -----
    When we talk about ``day`` to use in sorted sets, it's the string representation following this
    format: 'YYYY-MM-DD'.

    """

    abstract = True

    id = PKField()

    answers_for_session = DynamicSortedSetField()

    avg_by_day = SortedSetField()
    nb_by_day = SortedSetField()

    avg_by_day_for_trainer = DynamicSortedSetField()
    nb_by_day_for_trainer = DynamicSortedSetField()

    avg_by_day_for_training = DynamicSortedSetField()
    nb_by_day_for_training = DynamicSortedSetField()

    avg_by_day_for_training_and_trainer = DynamicSortedSetField()
    nb_by_day_for_training_and_trainer = DynamicSortedSetField()

    @staticmethod
    def get_answer_key(session_id, user_id, question_id):
        """Return the key to uses as key in``answers_for_session``.

        It uniquely identify an answer for a question by a user during a training session.

        Example
        -------

        >>> sorted_set = self.answers_for_session(session.pk)
        >>> key = self.get_answer_key(session.pk, user.pk, question.pk)
        >>> sorted_set.zadd(answer_value, key)

        Notes
        -----
        We could have used a hash of the string to reduce the memory footprint, but it's easier
        this way to debug data when checking directly in redis.

        """

        return '%s:%s:%s' % (session_id, user_id, question_id)

    @staticmethod
    def get_training_and_trainer_key(training_id, trainer_id):
        """Return the key to uses as key in``*_by_day_for_training_and_trainer``.

        Example
        -------

        >>> key = self.get_training_and_trainer_key(training.pk, trainer.pk)
        >>> sorted_set = self.nb_by_day_for_training_and_trainer(key)
        >>> sorted_set.zadd(nb_answers_for_this_training_and_trainer_for_this_day, day)

        Notes
        -----
        We could have used a hash of the string to reduce the memory footprint, but it's easier
        this way to debug data when checking directly in redis.

        """

        return '%s:%s' % (training_id, trainer_id)

    def get_stats(self, axis, start_date=None, end_date=None, *axis_values):
        """Get the statistics on the given axis.

        Parameters
        ----------
        axis : str
            Name of the axis for which we want the stats: one of:
            - global
            - session
            - trainer
            - training
            - training_and_trainer
        start_date : date, optional
            For axis other than "session", the start of the period to get the statistics.
            If ommited, will compute stats from the beginning.
            Should not be set if axis is "session".
        end_date : date, optional
            For axis other than "session", the end of the period to get the statistics.
            If ommited, will compute stats until the end.
            Should not be set if axis is "session".
        axis_values : list
            The values to precise the axis:
            - ``[session_id]`` if "session"
            - ``[trainer_id]`` if "trainer"
            - ``[training_id]`` if "training"
            - ``[training_id, trainer_id]`` if "training_and_trainer"

        Returns
        -------
        dict
            A dict with these entries
            - nb: the number of answers on this period
            - avg: the average answer value on this period
            - start_date: the start date in format 'YYYY-MM-DD'
            - end_date: the start date in format 'YYYY-MM-DD'
            - days: list with one entry by day for this period, each entry being a tuple with
                - day: the day in format 'YYY-MM-DD'
                - nb: the nb of answers on this day
                - avg: the avg answer on this day
            For the "session" axis, only "nb" and "avg" are present.

        Raises
        ------
        TypeError
            When dates are not expected or not correct.

        """

        # Specific case for session: we don't need dates, and use data from the simple sorted
        # set for the given session.
        if axis == 'session':
            if start_date or end_date:
                raise TypeError('`start_date` and `end_date` are unexpected for the '
                                '"session" axis')

            # Get the data.
            data = dict(self.answers_for_session(axis_values[0]).zrange(0, -1, withscores=True))

            # Handle the case with no data for this session.
            if not data:
                return {
                    'nb': 0,
                    'avg': 0,
                }

            # We have one entry by answer, so we know the number of answers.
            nb = len(data)
            # Each score (value of the dict) is the answer, so easy to compute the average.
            avg = sum(data.values()) / nb

            # We can now return the result, composed of only ``nb`` and ``avg`` (days related
            # entries don't have any sense here).
            return {
                'nb': nb,
                'avg': avg
            }

        # Now work on axis other than "session".

        # Ensure we have both date or none
        if start_date and not end_date or end_date and not start_date:
            raise TypeError('Both `start_date` and `end_date` are expected, or none of them.')
        if start_date and start_date > end_date:
            raise TypeError('`start_date` must be before `end_date`.')

        # Get the keys to fetch data from.
        if axis == 'global':
            # In this case, we directly have the sorted set
            fields = (self.nb_by_day, self.avg_by_day)
        else:
            # All other axises need a key to get the real sorted sets.

            if axis == 'training_and_trainer':
                # For this one, we compute the key from the ``axis_values``.
                key = self.get_training_and_trainer_key(*axis_values)
            else:
                # For all others, the key is the first entry in ``axis_values``.
                key = axis_values[0]

            # Get the sorted sets depending on the axis and the key
            fields = (
                getattr(self, 'nb_by_day_for_%s' % axis)(key),
                getattr(self, 'avg_by_day_for_%s' % axis)(key)
            )

        if start_date and end_date:
            # If we have a period, we have to retrieve all dates

            # Get all the days we need.
            days = [str(day) for day in day_range(start_date, end_date)]

            # We fetch data in one pass.
            with self.database.pipeline() as pipeline:
                for day in days:
                    # First, the number.
                    fields[0].zscore(day)
                    # Then, the average.
                    fields[1].zscore(day)

                # Data will have numbers in odd entries (0, 2, 4...),
                # and averages in even ones (1, 3, 5...)
                data = pipeline.execute()

        else:
            # We don't have a period, we can retrieve everything
            data_nb = dict(fields[0].zrange(0, -1, withscores=True))
            data_avg = dict(fields[1].zrange(0, -1, withscores=True))

            # We need the whole day set to fill the gaps (for days without data)
            days = set(data_nb.keys()) | set(data_avg.keys())

            if not days:
                data = []
                start_date = None
                end_date = None
            else:
                # Start date is the minimum date (key) in both dicts, end date is the max
                start_date = datetime.strptime(min(days), '%Y-%m-%d').date()
                end_date = datetime.strptime(max(days), '%Y-%m-%d').date()
                # We can get the list of all days
                days = [str(day) for day in day_range(start_date, end_date)]

                # Fill ``data`` to have the same format as when fetching from a period
                # using 0 for data where we don't have one.
                data = []
                for day in days:
                    data.extend([data_nb.get(day, 0), data_avg.get(day, 0)])

        # Handle the case with no data.
        if not data:
            return {
                'start_date': str(start_date) if start_date else None,
                'end_date': str(end_date) if end_date else None,
                'nb': 0,
                'avg': 0,
                'days': []
            }

        # A list with the stats for each days.
        stats_by_day = []
        # Total number of answers.
        total_nb_answers = 0
        # To compute the global average.
        total = 0

        # We can now compute the data.
        for index, day in enumerate(days):
            nb = data[index*2]  # 0, 2, 4...
            avg = data[index*2+1]  # 1, 3, 5...

            if nb is not None and avg is not None:
                total_nb_answers += nb
                total += nb * avg

            stats_by_day.append((day, nb or 0, avg or 0))

        # The global average
        total_avg = total / total_nb_answers if total_nb_answers else 0

        # We can return the data
        return {
            'start_date': str(start_date),
            'end_date': str(end_date),
            'nb': total_nb_answers,
            'avg': total_avg,
            'days': stats_by_day
        }


class GlobalSatisfaction(BaseStatModel):
    """A redis model to store global satisfaction.

    Based on ``BaseStatModel`` it allows to save the same kind of data than for the questions,
    but in a global way.
    It's easy to have the global satisfaction for a specific trainer for example.

    It must only include data about answers of question having the ``is_for_global_satisfaction``
    flag set to ``True`` (ie single choice question with percentage)

    Notes
    -----

    This redis model must be used as a Singleton: you must not create one manually, but use
    the ``get_instance`` class method that will create it if needed.

    Example
    -------

    >>> global_satisfaction = GlobalSatisfaction.get_instance()
    >>> global_satisfaction.nb_by_day.zadd(nb_answers_for_this_day, day)

    """

    @classmethod
    def get_instance(cls):
        return cls.get_or_connect(pk=1)[0]  # [0] is the instance, [1] is the ``created`` flag


class QuestionData(BaseStatModel, DjangoBackedRedisModel):
    """Store precomputed stats day by day for a question.

    See ``BaseStatModel`` for description of attributes not described below.

    """

    @cached_property
    def django_model(self):
        """Define the django model attached to this redis model."""
        from .models import Question
        return Question

    def save_single_choice_number_answer(self, training_id, trainer_id, session_id, user_id, value,
                                         answer_time):
        """Save answer data and precompute daily stats.

        It will save the answer value and compute the daily stats for different axises, for the
        current question itself but also, if the question if marked as couting for global
        satisfaction, in the ``GlobalSatisfaction`` singleton.

        Parameters
        ----------
        training_id : UUID
            Pk of the ``Training`` instance in which the answer was created.
        trainer_id : UUID
            Pk of the trainer (a ``User`` instance) who hosted the training.
        session_id : UUID
            Pk of the ``TrainingSession`` instance) in which the answer was created.
        user_id : UUID
            Pk of the ``User`` instance) who answered the qestion.
        answer_time : datetime
            When the answer was created. Used to compute the day as stats are computed on
            a daily basis.

        """

        global_satisfaction = None
        if self.django_instance.is_for_global_satisfaction:
            global_satisfaction = GlobalSatisfaction.get_instance()

        # Stringified date, format 'YYYY-MM-DD'
        day = str(answer_time.date())

        answer_key = self.get_answer_key(session_id, user_id, self.pk.get())
        train_key = self.get_training_and_trainer_key(training_id, trainer_id)

        # All the zset we'll have to update

        # First, all for the question
        zset_tuples = [
            ('question_total',
                self.avg_by_day,
                self.nb_by_day),
            ('question_trainer',
                self.avg_by_day_for_trainer(trainer_id),
                self.nb_by_day_for_trainer(trainer_id)),
            ('question_training',
                self.avg_by_day_for_training(training_id),
                self.nb_by_day_for_training(training_id)),
            ('question_training_and_trainer',
                self.avg_by_day_for_training_and_trainer(train_key),
                self.nb_by_day_for_training_and_trainer(train_key)),
        ]

        # And if question answers count for global satisfaction, we want to update
        # the global zsets
        if global_satisfaction:
            zset_tuples += [
                ('global_total',
                    global_satisfaction.avg_by_day,
                    global_satisfaction.nb_by_day),
                ('global_trainer',
                    global_satisfaction.avg_by_day_for_trainer(trainer_id),
                    global_satisfaction.nb_by_day_for_trainer(trainer_id)),
                ('global_training',
                    global_satisfaction.avg_by_day_for_training(training_id),
                    global_satisfaction.nb_by_day_for_training(training_id)),
                ('global_training_and_trainer',
                    global_satisfaction.avg_by_day_for_training_and_trainer(train_key),
                    global_satisfaction.nb_by_day_for_training_and_trainer(train_key)),
            ]

        # Check if the user already answered this question in the session.
        existing_value = self.answers_for_session(session_id).zscore(answer_key)

        def do_save(pipeline):

            # Get avg and nb for each tuple before starting the redis transaction.
            avg, nb = {}, {}
            for name, avg_zset, nb_zset in zset_tuples:
                avg[name] = avg_zset.zscore(day) or 0
                nb[name] = nb_zset.zscore(day) or 0

            # Start the transaction to update values.
            pipeline.multi()

            # Now we can save data in redis.
            for name, avg_zset, nb_zset in zset_tuples:
                total = avg[name] * nb[name]
                total += value

                if existing_value is not None:
                    # If the user already had a value saved, remove it from the total
                    # Also, we don't change the nb of answers.
                    total -= existing_value
                else:
                    # It's a new answer, we can increment the nb of answers.
                    nb_zset.zincrby(day, 1)
                    nb[name] += 1

                # Set the new average for this day.
                avg_zset.zadd(total / nb[name], day)

            # Save data for session (we don't save nb and avg by day here as a session
            # is a one time shot, so the numbers is the cardinality of entries in the zset, and
            # the average, the average of the answers in the zset (each session has its own zset)
            self.answers_for_session(session_id).zadd(value, answer_key)
            if global_satisfaction:
                global_satisfaction.answers_for_session(session_id).zadd(value, answer_key)

        # Get the redis limpyd fieldto watch for update while we're updating them ourselves.
        # We simply iterate over ``zset_tuple`` to get every fields (all entries except the
        # first one which is the name) of each tuple.
        watched = list(chain.from_iterable((t[1:] for t in zset_tuples)))

        # We could have used the simplest `self.database.transaction` but here we want to check
        # if the answer for this question for this user was answered during the execution.
        # If yes, we stop because we don't want to save an older value for the same answer/user.
        # If it's any other change, we still want to save the answer so we continue.
        with self.database.pipeline(True) as pipeline:
            while True:
                try:
                    pipeline.watch(*watched)
                    do_save(pipeline)
                    return pipeline.execute()
                except WatchError:
                    # Get the existing value currently in redis
                    current_existing_value = \
                        self.answers_for_session(session_id).zscore(answer_key)
                    # If it has changed, stop here, a new value for this answer/user was saved
                    if current_existing_value != existing_value:
                        break

    def update_single_choice_number_answers(self, answer_dicts, force_all=False):
        """Given a list of answer for the question, update its statistics.

        Parameters
        ----------
        answer_dicts : list of dicts
            A list containing a dict for each answer. Each dict must have the following entries:
            - ``id``
            - ``training_session_id``
            - ``trainee_id``
        force_all : boolean
            If ``False``, will only use new answers to update the statistics. If ``True``, all
            answers will be used.

        Returns
        -------
        tuple
            A tuple with two numbers: the number of answers used successfuly to update the
            statistics and the number that failed.

        Notes
        -----
        If ``force_all`` is ``False`` (the default), we 'll get the missing answers by checking
        the difference between the ones we've got from the ``answer_dicts`` argument, and the
        one in redis.
        And as answers are saved by session, we manage to do the difference by session too.

        """
        if force_all:
            # If we're asked to use all answers, do not compute missing ones, use all.
            pks = [answer['id'] for answer in answer_dicts]

        else:
            # Get all answer keys, by session.
            answer_keys = {}
            answers_by_key = {}
            for answer in answer_dicts:
                key = self.get_answer_key(
                    answer['training_session_id'],
                    answer['trainee_id'],
                    self.pk.get(),
                )
                answers_by_key[key] = answer
                answer_keys.setdefault(answer['training_session_id'], set()).add(key)

            # Get all missing keys
            missing_keys = set()
            for session_id, all_keys in answer_keys.items():
                existing_keys = self.answers_for_session(session_id).zrange(0, -1)
                missing_keys.update(all_keys.difference(existing_keys))

            # Get the pks for all missing keys
            pks = [answers_by_key[key]['id'] for key in missing_keys]

        # Get answers from the pks
        from .models import Answer  # Avoid recursive import
        answers = Answer.objects.filter_for_stats(pks)

        # Now we can update stats for all these answers.
        updated, failed = 0, 0
        for answer in answers:
            try:
                self.django_instance.save_answer_data(answer)
            except Exception as e:
                failed += 1
                logger.exception(
                    'Error saving survey answer redis data. Answer=[%s], message=%s',
                    answer.pk, e
                )
            else:
                updated += 1

        return updated, failed
