imaginaryfriend/src/domain/pair.py

from datetime import datetime, timedelta

from orator.orm import Model
from orator.orm import belongs_to
from orator.orm import has_many

from src.utils import *

import src.domain.chat
import src.domain.reply
import src.domain.word


class Pair(Model):
    __guarded__ = ['id']

    # TODO Move to config
    end_sentence = '.....!!?'
    all = '.!?,;:()"\'

    @has_many
    def replies(self):
        return src.domain.reply.Reply

    @belongs_to
    def chat(self):
        return src.domain.chat.Chat

    @belongs_to
    def first(self):
        return src.domain.word.Word

    @belongs_to
    def second(self):
        return src.domain.word.Word

    @staticmethod
    def generate(message):
        return Pair.generate_story(message, message.words, random.randint(0, 2) + 1)

    @staticmethod
    def generate_story(message, words, sentences):
        words_ids = src.domain.word.Word.where_in('word', words).get().pluck('id').all()

        result = []
        for _ in range(0, sentences):
            result.append(Pair.__generate_sentence(message, words_ids))

        return ' '.join(result)

    @staticmethod
    def learn(message):
        src.domain.word.Word.learn(message.words)

        words = [None]
        for word in message.words:
            words.append(word)
            if word[-1] in Pair.end_sentence:
                words.append(None)
        if words[-1] is not None:
            words.append(None)

        while any(word for word in words):
            trigram = words[:3]
            first_word_id, second_word_id, *third_word_id = list(map(
                lambda x: None if x is None else src.domain.word.Word.where('word', x).first().id,
                trigram
            ))
            third_word_id = None if len(third_word_id) == 0 else third_word_id[0]

            words.pop(0)

            pair = Pair.where('chat_id', message.chat.id)\
                .where('first_id', first_word_id)\
                .where('second_id', second_word_id)\
                .first()
            if pair is None:
                pair = Pair.create(chat_id=message.chat.id,
                                   first_id=first_word_id,
                                   second_id=second_word_id)

            reply = pair.replies().where('word_id', third_word_id).first()

            if reply is not None:
                reply.count += 1
                reply.save()
            else:
                pair.replies().create(pair_id=pair.id, word_id=third_word_id)

    @staticmethod
    def __generate_sentence(message, word_ids):
        sentences = []
        safety_counter = 50
        first_word = None
        second_words = list(word_ids)

        while safety_counter > 0:
            pair = Pair.__get_pair(chat_id=message.chat.id, first_id=first_word, second_ids=second_words)
            replies = getattr(pair, 'replies', [])
            safety_counter -= 1

            if pair is None or len(replies) == 0:
                continue

            reply = random.choice(replies.all())
            first_word = pair.second.id

            # TODO. WARNING! Do not try to fix, it's magic, i have no clue why
            try:
                second_words = [reply.word.id]
            except AttributeError:
                second_words = None

            if len(sentences) == 0:
                sentences.append(capitalize(pair.second.word))
                word_ids.remove(pair.second.id)

            # TODO. WARNING! Do not try to fix, it's magic, i have no clue why
            try:
                reply_word = reply.word.word
            except AttributeError:
                reply_word = None

            if reply_word is not None:
                sentences.append(reply_word)
            else:
                break

        sentence = ' '.join(sentences).strip()
        if sentence[-1:] not in Pair.end_sentence:
            sentence += random.choice(list(Pair.end_sentence))

        return sentence

    @staticmethod
    def __get_pair(chat_id, first_id, second_ids):
        ten_minutes_ago = datetime.now() - timedelta(seconds=10 * 60)
        pairs = Pair.with_('replies')\
            .where('chat_id', chat_id)\
            .where('first_id', first_id)\
            .where_in('second_id', second_ids)\
            .where('created_at', '<', ten_minutes_ago)\
            .limit(3)\
            .get()\
            .all()

        return random_element(pairs)