SGJP / Kuznia

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

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, reader_lexemes
from dictionary.ajax_slickgrid import SlickGridQuery
from common.decorators import ajax
from common.util import bisect_left, reverse, json_encode
from dictionary.util import check_query_params

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 DBList(object):
    def __init__(self, lexeme_grid):
        self.lexeme_indices = lexeme_grid.lexeme_list.lexemeindex_set
        self.count = self.lexeme_indices.count()

    def __getitem__(self, item):
        return self.lexeme_indices.get(index=item).lexeme

    def __len__(self):
        return self.count


def ident(x):
    return x


def int_list(l):
    return [int(i) for i in l] if l else []


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.pos_name.abbr,
        '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.vocab_list()),
        '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)),
    }

    literal_filter_fields = [
        'entry', 'part_of_speech', 'status', 'comment', 'qualifiers_dor',
        'gloss', 'note', 'extended_note', 'pronunciation', 'valence',
        'lexemeav__attribute_value_id',
        'lexemeinflectionpattern__pronunciation',
        'value',
    ]

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

    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),
    }

    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),
    }

    def __init__(self, *args, **kwargs):
        super(LexemeQuery, self).__init__(*args, **kwargs)
        self.visible_vocabs = self.query_params['visible_vocabs']
        self.reader = self.query_params.get('reader', True)
        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 optimize_not_in(self, queryset, field, new_rule, attr=None):
        if field in self.indirect_fields or attr and attr.closed:
            key = 'extra' if field.startswith('extra') else field
            field1, table, field2, f = self.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
        if field in self.indirect_fields2:
            field1, table1, table2, join1, join2, field2, f = \
                self.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

    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)
        attr = None
        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),
                }
        if lookup == '-in':
            queryset_not_in = self.optimize_not_in(
                queryset, field, new_rule, attr)
            if queryset_not_in:
                return queryset_not_in

        return super(LexemeQuery, self).apply_filter_rule(queryset, new_rule)

    def get_queryset(self):
        lexemes = super(LexemeQuery, self).get_queryset()
        if self.reader:
            return reader_lexemes(lexemes)
        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):
        if self.lexeme_list:
            try:
                if not self.a_fronte():
                    mask = reverse(mask)
                li = self.lexeme_list.lexemeindex_set.filter(entry__gte=mask)[0]
                return li.index
            except IndexError:
                count = self.lexeme_list.lexemeindex_set.count()
                return count - 1 if count else 0
        else:
            id_list = self.get_id_list()
            count = len(id_list)
            index = bisect_left(id_list, mask, cmp=self.lexeme_cmp())
        if index == count > 0:
            index -= 1
        return index

    def cache_key(self):
        key = json_encode(self.query_params, 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 self.lexeme_list:
            return self.lexeme_list.lexemeindex_set.values_list(
                'lexeme_id', flat=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', 'entry')
        lexeme_indices = (
            LexemeIndex(
                lexeme_list=self.lexeme_list, lexeme_id=lexeme_id, index=i,
                entry=entry if self.a_fronte() else reverse(entry))
            for i, (lexeme_id, entry) in enumerate(lexeme_ids))
        LexemeIndex.objects.bulk_create(lexeme_indices, batch_size=4000)

    def lexeme_cmp(self):
        def fun(lexeme_id, mask):
            entry = Lexeme.all_objects.get(id=lexeme_id).entry
            return self.strcoll(entry, mask)
        return fun

    def strcoll(self, s1, s2):
        if not self.a_fronte():
            s1 = reverse(s1)
            s2 = reverse(s2)
        return locale.strcoll(s1, s2)

    def export_list(self, columns, output_file):
        column_translation = {
            'pos': 'part_of_speech_id',
            'abbr_pos': 'part_of_speech__pos_name__abbr',
            '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, query_params):
    check_query_params(request, query_params)
    query = LexemeQuery(query_params)
    return {'index': query.row_index(id)}


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


@ajax(login_required=False, method='get')
def get_lexemes(request, from_page, to_page, rows, query_params, columns,
                force_reload=False):
    check_query_params(request, query_params)
    session_prefix = 'reader_' if query_params.get('reader', True) else ''
    request.session[session_prefix + 'sort_rules'] = query_params['sort_rules']
    request.session[session_prefix + 'filter'] = query_params['filter']
    request.session[session_prefix + 'columns'] = columns
    query = LexemeQuery(query_params, 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]
    # brzydka łata
    lexemes_qs = prefetch(Lexeme.all_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, query_params, columns, exponent):
    check_query_params(request, query_params)
    query = LexemeQuery(query_params, columns)
    lexemes = query.get_sorted_queryset().filter(lexemeform__form=exponent)
    rows = [row for row in query.prepare_rows(lexemes)]
    for row in rows:
        row_idx = query.row_index(row['id'])
        if row_idx is not None:
            row['row'] = row_idx
    return {
        'rows': rows,
    }


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