SGJP / Kuznia

# -*- coding: utf-8 -*-
from hashlib import md5

from django.db.models import Count, Q
from django.core.cache import cache
from django.utils.encoding import force_unicode

from dictionary.models import Lexeme, LexemeAttribute, PartOfSpeech, LexemeList, \
    LexemeIndex
from dictionary.ajax_slickgrid import SlickGridQuery
from common.decorators import ajax
from common.util import bisect_left, reverse, json_encode

import locale
locale.setlocale(locale.LC_ALL, 'pl_PL.UTF-8')

ASPECT = LexemeAttribute.objects.get(name=u'aspekt')

COLOR_SCHEMES = dict(PartOfSpeech.objects.values_list('symbol', 'color_scheme'))


class LexemeQuery(SlickGridQuery):
    model = Lexeme
    sort_field = 'entry'
    default_columns = ('entry', 'abbr_pos', 'genders')
    column_data = {
        'id': lambda lexeme: lexeme.id,
        'scheme': lambda lexeme: COLOR_SCHEMES[lexeme.part_of_speech_id],
        'entry': lambda lexeme: lexeme.entry,
        'abbr_pos': lambda lexeme: lexeme.part_of_speech.abbr_name,
        'pos': lambda lexeme: lexeme.part_of_speech_id,
        'patterns': lambda lexeme: lexeme.lip_data()['patterns'],
        'genders': lambda lexeme:
            getattr(lexeme.attribute_value(ASPECT), 'value', '')
            or lexeme.lip_data()['genders'],
        'vocabs': lambda lexeme:
            '/'.join(lexeme.vocabularies.values_list('id', flat=True)),
        'owner': lambda lexeme: lexeme.owner_vocabulary_id,
        'status': lambda lexeme:
            force_unicode(dict(Lexeme.STATUS_CHOICES).get(lexeme.status)),
        'qualifiers': lambda lexeme:
            ', '.join(lexeme.qualifiers.values_list('label', flat=True)),
    }

    filter_field_translation = {
        'form': 'lexemeform__form',
        'lexeme_qualifier': 'qualifiers__id',
        'lip_qualifier': 'lexemeinflectionpattern__qualifiers__id',
        'qualifier': 'qualifiers_cache__id',
        'classification_value': 'classificationvalue__id',
        'pattern_name': 'lexemeinflectionpattern__pattern__name',
        'pattern_type': 'lexemeinflectionpattern__pattern__type_id',
        'gender': 'lexemeinflectionpattern__gender_id',
        'containing_vocabulary': 'vocabularies__id',
        'owner_vocabulary': 'owner_vocabulary_id',
        'pattern_count': 'pc',
        'gender_count': 'gc',
        'cr_type': 'refs_to__type_id',
        'borrowing_source': 'borrowing_source_id',
        'responsible': 'responsible_id',
        'abbr_pos': 'part_of_speech__abbr_name',
    }

    def __init__(self, visible_vocabs, *args, **kwargs):
        self.visible_vocabs = list(visible_vocabs)
        super(LexemeQuery, self).__init__(*args, **kwargs)
        self.lexeme_list = self.get_lexeme_list()

    def a_fronte(self):
        return self.sort_rules[0] == 'a_fronte'

    def sort_queryset(self, queryset):
        order_list = [self.sort_field, 'id']
        if not self.a_fronte():
            queryset = queryset.extra(select={'rev': "reverse(haslo)"})
            order_list[0] = 'rev'
        return queryset.extra(order_by=order_list)

    def apply_complex_filter_rule(self, queryset, rule):
        lookup = self.lookup_translation[rule['op']]
        negated = (lookup[0] == '-')
        lookup = lookup.lstrip('-')
        fields = [
            self.translate_filter_field(field) for field in rule['fields']]
        data = rule['data']
        q = reduce(
            lambda q1, q2: q1 | q2,
            (Q(**{(field + '__' + lookup): data}) for field in fields))
        if negated:
            queryset = queryset.exclude(q)
        else:
            queryset = queryset.filter(q).distinct()
        return queryset

    def apply_filter_rule(self, queryset, rule):
        lookup = self.lookup_translation[rule['op']]
        negated = (lookup[0] == '-')
        field, data = rule['field'], rule['data']
        new_rule = dict(rule)
        if field == 'pattern_count':
            queryset = queryset.annotate(
                pc=Count('lexemeinflectionpattern__pattern', distinct=True))
        elif field == 'gender_count':
            queryset = queryset.annotate(
                gc=Count('lexemeinflectionpattern__gender', distinct=True))
        elif not field.startswith('extra') and lookup in ['in', '-in']:
            if not data:
                new_rule['op'] = 'isnull'
                lookup = new_rule['op']
                new_rule['data'] = not negated
        elif field == 'pronunciation':
            new_rule['fields'] = [
                'pronunciation',
                'lexemeinflectionpattern__pronunciation'
            ]
            return self.apply_complex_filter_rule(queryset, new_rule)
        elif field.startswith('extra'):
            attr = LexemeAttribute.objects.get(id=int(field.split('-')[1]))
            new_rule['field'] = 'lexemeav__attribute_value_id'
            if attr.closed and not data:
                new_rule['op'] = 'in' if negated else '-in'
                lookup = new_rule['op']
                values = attr.values.all().values_list('id', flat=True)
                new_rule['data'] = tuple(values)
            elif not attr.closed:
                new_rule['field'] = 'value'
                # lekki abuse
                matching_values = super(LexemeQuery, self).apply_filter_rule(
                    attr.values, new_rule).values_list('id', flat=True)
                new_rule = {
                    'field': 'lexemeav__attribute_value_id',
                    'op': 'in',
                    'data': tuple(matching_values),
                }
        # optymalizacja warunków "różne od"
        ident = lambda x: x
        int_list = lambda l: [int(i) for i in l] if l else []
        indirect_fields = {
            'containing_vocabulary': (
                'slownik', 'leksemy_w_slownikach', 'l_id', ident),
            'lexeme_qualifier': (
                'qualifier_id', 'kwalifikatory_leksemow', 'lexeme_id',
                int_list),
            'qualifier': (
                'qualifier_id', 'dictionary_lexemeformqualifier', 'lexeme_id',
                int_list),
            'classification_value': (
                'classification_value_id', 'dictionary_lexemecv', 'lexeme_id',
                int_list),
            'gender': (
                'gender_id', 'odmieniasie', 'l_id', int_list),
            'cr_type': (
                'typods_id', 'odsylacze', 'l_id_od', int_list),
            'extra': (
                'attribute_value_id', 'dictionary_lexemeav', 'lexeme_id',
                int_list),
        }
        if lookup == '-in' and (
                field in indirect_fields
                or field.startswith('extra') and attr.closed):
            key = 'extra' if field.startswith('extra') else field
            field1, table, field2, f = indirect_fields[key]
            for value in f(new_rule['data']):
                queryset = queryset.extra(where=[
                    '''%%s NOT IN (SELECT %s FROM %s WHERE %s = leksemy.id)'''
                    % (field1, table, field2)], params=[value])
            return queryset
        indirect_fields2 = {
            'pattern_name': (
                'w_id', 'odmieniasie', 'wzory', 'w_id', 'id', 'l_id', ident),
            'pattern_type': (
                'typ', 'odmieniasie', 'wzory', 'w_id', 'id', 'l_id', int_list),
            'lip_qualifier': (
                'qualifier_id', 'odmieniasie', 'kwalifikatory_odmieniasiow',
                'id', 'lexemeinflectionpattern_id', 'l_id', int_list),
        }
        if field in indirect_fields2 and lookup == '-in':
            field1, table1, table2, join1, join2, field2, f = \
                indirect_fields2[field]
            for value in f(new_rule['data']):
                queryset = queryset.extra(where=[
                    '''%%s NOT IN (SELECT b.%s FROM %s a
                    INNER JOIN %s b ON (a.%s = b.%s) WHERE a.%s = leksemy.id)'''
                    % (field1, table1, table2, join1, join2, field2)
                ], params=[value])
            return queryset
        return super(LexemeQuery, self).apply_filter_rule(queryset, new_rule)

    def get_queryset(self):
        lexemes = super(LexemeQuery, self).get_queryset()
        return lexemes.distinct().filter(
            vocabularies__id__in=self.visible_vocabs)

    # nieużywane
    def filter_from(self, queryset, from_value, upward):
        if self.a_fronte():
            return super(LexemeQuery, self).filter_from(
                queryset, from_value, upward)
        else:
            if upward:
                comp = '>='
            else:
                comp = '<='
            return queryset.extra(
                where=["reverse(haslo) " + comp + " %s"],
                params=[reverse(from_value)])

    # indeks wiersza w danym sortowaniu, w którym
    # znajdzie się rekord o danym id
    def row_index(self, lexeme_id):
        if self.lexeme_list:
            try:
                li = self.lexeme_list.lexemeindex_set.get(lexeme_id=lexeme_id)
                return li.index
            except LexemeIndex.DoesNotExist:
                return None
        else:
            id_list = self.get_id_list()
            try:
                return id_list.index(lexeme_id)
            except ValueError:
                return None

    def search_index(self, mask):
        # TODO zrobić sprytniej (pytać bazę o poszczególne wiersze w bisect)
        if self.lexeme_list:
            id_list = self.lexeme_list.lexemeindex_set.values_list(
                'lexeme_id', flat=True)
        else:
            id_list = self.get_id_list()
        count = len(id_list)
        if count == 0:
            return 0

        index = bisect_left(id_list, mask, cmp=self.lexeme_cmp())
        if index == count:
            index -= 1
        return index

    def cache_key(self):
        key = json_encode(
            (self.sort_rules, self.filter, self.visible_vocabs),
            ensure_ascii=True)
        return md5(key.encode()).hexdigest()

    def get_cached_lexemes(self, refresh=True):
        key = self.cache_key()
        cached = cache.get(key)
        if refresh:
            cache.set(key, cached)
        return cached

    def cache_lexemes(self, id_list):
        key = self.cache_key()
        cache.set(key, id_list)
        key_list = cache.get('key_list', [])
        if key not in key_list:
            key_list.append(key)
        cache.set('key_list', key_list)

    def get_id_list(self, force_reload=False, refresh=True):
        if not force_reload:
            id_list = self.get_cached_lexemes(refresh=refresh)
        else:
            id_list = None
        if id_list is None:
            lexemes = self.get_sorted_queryset()
            if 'rev' in lexemes.query.extra_select:
                id_list = list(
                    row[0] for row in lexemes.values_list('id', 'rev'))
            else:
                id_list = list(lexemes.values_list('id', flat=True))
            self.cache_lexemes(id_list)
        return id_list

    def get_lexeme_list(self):
        try:
            return LexemeList.objects.get(
                filter=json_encode(self.filter), a_fronte=self.a_fronte(),
                vocabularies=json_encode(self.visible_vocabs))
        except LexemeList.DoesNotExist:
            return None

    def reload_lexeme_list(self):
        LexemeIndex.objects.filter(lexeme_list=self.lexeme_list).delete()
        lexeme_ids = self.get_sorted_queryset().values_list('id', flat=True)
        lexeme_indices = (
            LexemeIndex(
                lexeme_list=self.lexeme_list, lexeme_id=lexeme_id, index=i)
            for i, lexeme_id in enumerate(lexeme_ids))
        LexemeIndex.objects.bulk_create(lexeme_indices, batch_size=4000)

    def lexeme_cmp(self):
        def fun(lexeme_id, mask):
            e1 = Lexeme.objects.get(id=lexeme_id).entry
            e2 = mask
            if not self.a_fronte():
                e1 = reverse(e1)
                e2 = reverse(e2)
            result = locale.strcoll(e1, e2)
            return result

        return fun

    def export_list(self, columns, output_file):
        column_translation = {
            'pos': 'part_of_speech_id',
            'abbr_pos': 'part_of_speech__abbr_name',
            'patterns': 'pattern_list',
            'genders': 'gender_list',
            'vocabs': 'vocab_list',
            'owner': 'owner_vocabulary_id',
            'qualifiers': 'qualifier_list',
        }
        lexemes = Lexeme.objects.filter(id__in=self.get_id_list())
        if 'patterns' in columns:
            lexemes = lexemes.extra(select={
                'pattern_list':
                    "select coalesce(string_agg(distinct w.w_id, '/'), '') "
                    "from wzory w join odmieniasie o on w.id = o.w_id "
                    "where o.l_id = leksemy.id",
            })
        if 'genders' in columns:
            lexemes = lexemes.extra(select={
                'gender_list':
                    """select coalesce(
                        (select val.value
                          from dictionary_lexemeav lav
                          join dictionary_lexemeattributevalue val
                            on lav.attribute_value_id = val.id
                          join dictionary_lexemeattribute attr
                            on (val.attribute_id = attr.id
                                and attr.name = 'aspekt')
                          where (leksemy.id = lav.lexeme_id)),
                        (select string_agg(distinct g.symbol, '/')
                          from dictionary_gender g
                          join odmieniasie o on g.id = o.gender_id
                          where o.l_id = leksemy.id),
                        '')"""
            })
        if 'vocabs' in columns:
            lexemes = lexemes.extra(select={
                'vocab_list':
                    "select coalesce(string_agg(ls.slownik, '/'), '') "
                    "from leksemy_w_slownikach ls "
                    "where ls.l_id = leksemy.id",
            })
        if 'qualifiers' in columns:
            lexemes = lexemes.extra(select={
                'qualifier_list':
                    "select coalesce(string_agg(kwal.kwal, ', '), '') "
                    "from kwalifikatory_leksemow kl "
                    "join kwalifikatory kwal on kl.qualifier_id = kwal.id "
                    "where kl.lexeme_id = leksemy.id",
            })
        fields = [
            column_translation.get(column, column) for column in columns]
        for row in lexemes.values(*fields):
            print >>output_file,  '\t'.join(
                unicode(row[field]) for field in fields)


# Zapytanie o indeks wiersza o pewnym id przy danym sortowaniu
@ajax(login_required=False, method='get')
def row_index(request, id, sort_rules, filter, visible_vocabs):
    query = LexemeQuery(
        filter=filter, sort_rules=sort_rules, visible_vocabs=visible_vocabs)
    return {'index': query.row_index(id)}


@ajax(login_required=False, method='get')
def search_index(request, sort_rules, filter, visible_vocabs, search=''):
    query = LexemeQuery(
        filter=filter, sort_rules=sort_rules, visible_vocabs=visible_vocabs)
    return {'index': query.search_index(search)}


@ajax(login_required=False, method='get')
def get_lexemes(request, from_page, to_page, rows, sort_rules, filter,
                visible_vocabs, force_reload=False):
    request.session['sort_rules'] = sort_rules
    request.session['filter'] = filter
    columns = request.session.get('columns')
    query = LexemeQuery(
        filter=filter, sort_rules=sort_rules, visible_vocabs=visible_vocabs,
        columns=columns)
    start, response_rowcount = query.count_pages(from_page, to_page, rows)
    if query.lexeme_list:
        count = query.lexeme_list.lexemeindex_set.count()
        sublist = list(query.lexeme_list.lexemeindex_set.filter(
            index__gte=start, index__lt=start + response_rowcount)
           .values_list('lexeme_id', flat=True))
    else:
        id_list = query.get_id_list(force_reload)
        count = len(id_list)
        sublist = id_list[start:start + response_rowcount]
    lexemes_qs = prefetch(Lexeme.objects.filter(id__in=sublist))
    lexemes_dict = dict((l.id, l) for l in lexemes_qs)
    lexemes = [lexemes_dict[lexeme_id] for lexeme_id in sublist]
    return {
        'rows': query.prepare_rows(lexemes),
        'count': count,
        'page': from_page,
    }


@ajax(login_required=False, method='get')
def search_by_form(request, sort_rules, filter, exponent, visible_vocabs):
    query = LexemeQuery(
        filter=filter, sort_rules=sort_rules, visible_vocabs=visible_vocabs)
    id_list = query.get_id_list()
    lexemes = query.get_sorted_queryset().filter(lexemeform__form=exponent)
    rows = [row for row in query.prepare_rows(lexemes) if row['id'] in id_list]
    for row in rows:
        row['row'] = id_list.index(row['id'])
    return {
        'rows': rows,
    }


def prefetch(queryset):
    return queryset.select_related(
        'owner_vocabulary', 'part_of_speech').prefetch_related(
        'lexemeinflectionpattern_set__pattern',
        'lexemeinflectionpattern_set__gender', 'vocabularies')