#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""Mean filtering as used for simultaneous EEG/fMRI"""

import eegpy
from eegpy.misc import FATALERROR
from eegpy.misc import debug
from eegpy.helper import upsample,downsample,find_max_overlap
#from eegpy.filter.filt_misc import filterRecursively
#if debug:
#    import pylab

#################
# Module-Import #
#################

try:
    import numpy as n
except ImportError:
    raise FATALERROR('SciPy or NumPy not found!\nPlease visit www.scipy.org or numeric.scipy.org for more information.')

# Global variables
_t_before = 50 #Time to reduce triggers by

########################
# Function definitions #
########################

class MeanFilter():
    """Implementation for the filtering by subtracting the mean.
    Should work with EEG-files and with arrays, first only with arrays.
    Works in-place!!!"""

    _data=None
    _ts = None
    _len = None
    _mean_data = None
    _filter = None#lambda x:x
    _corrts = None
    _t_before = _t_before

    def __init__(self, x=None, ts=None, len=None, filter=None, t_before=None):
        if x!=None:
            self.set_data(x)
        if t_before != None:
            self._t_before = t_before
            #print "t_before in meanfilt:"
        if ts!=None:
            self.set_timepoints(ts)
        if len!=None:
            self.set_length(len)
        if filter != None:
            self._filter = filter

    def get_data(self):
        return self._data

    def set_data(self, x):
        try:
            a = x.shape
        except Exception, e:
            raise ValueError("Only arrays / memmaps can be used as input for MeanFilter")
        self._data = x

    def set_timepoints(self,ts):
        self._ts = [int(x)-self._t_before for x in ts]
        self._ts.sort()

    def get_length(self):
        return self._len

    def set_length(self,l):
        self._len = int(l)

    def set_timepoints_auto(self,start,end=10e50,step=5500,ch_num=0,width=100):
        """If triggers are not good, one can try to automatically find the timepoints"""
        #print "Setting timepoints automatically"
        assert ch_num<self._data.shape[1] and ch_num>=0, "ch_num is not valid"
        ts = []
        t = int(start)
        offset=0
        while t<self._data.shape[0]-step and t<end:
            if t==int(start):
                #template = self._data[t-width/2:t+width/2,ch_num]
                template = self._data[t:t+step,ch_num]
                ts.append(t)
            else:
                #offset = find_max_overlap(template, self._data[t-width/2:t+width/2,ch_num], width/2)
                offset = find_max_overlap(template, self._data[t:t+step,ch_num], width/2)
                #print offset
                ts.append(t+offset)
            if debug:
                print ts[-1],
            t+=step+offset
        self.set_timepoints(ts)
        return ts

    def check_before_filtering(self):
        if self._data == None:
            raise RuntimeError("No data were set yet.")
        if self._ts == None:
            raise RuntimeError("No timepoints were set yet.")
        if self._len == None:
            #Auto calculate _len as the minimal distance between subsequent timepoints
            self._len = n.diff(n.array(self._ts)).min()
            if self._len/n.diff(n.array(self._ts)).mean()<0.1:
                print "Auto-setting of length might have failed..., self._len = ", self._len
        #Now all three values needed should be o.k.
        if debug:
            print "self._data.shape:", self._data.shape
            print "self._ts:", self._ts, n.diff(n.array(self._ts)).mean()
            print "self._len:", self._len

    def calc_mean(self):
        #self._corrts = n.zeros((len(self._ts)),"i")
        mean_shape = []
        mean_shape.append(self._len)
        for i in range(1,len(self._data.shape)):
            mean_shape.append(self._data.shape[i])
        #print "mean_shape = ", mean_shape
        #print "Sollte sein: ",self._data[self._ts[0]:self._ts[0]+self._len,...].shape
        self._mean_data = n.zeros(mean_shape,"d")
        #print self._mean_data.shape, self._data[self._ts[0]:self._ts[0]+self._len,...].shape
        for i, t in enumerate(self._ts): #First walk-through, calculating average
            #self._mean_data += self._filter(self._data[t:t+self._len,...])
            try:
                d = self._data[t:t+self._len,...]
                #d_up = upsample(d)
                self._mean_data += d
                #if i==0:
                    #self._mean_data += d#_up
                #else:
                    #self._corrts[i] = find_max_overlap(self._mean_data[:,10],d_up[:,10])
                    #if self._corrts[i]<0:
                    #    self._mean_data[:i,...] += d_up[-i:,...]
                    #elif self._corrts[i]==0:
                    #    self._mean_data += d_up
                    #else:
                    #    self._mean_data[i:,...] += d_up[:-i,...]
                #print "maxOverlap:", self._corrts[i]
                #pylab.clf()
                #pylab.ioff()
                #for i in range(d.shape[1]):
                #    pylab.subplot(8,d.shape[1]/8,i+1)
                #    pylab.plot(d[:,i])
                #pylab.show()
            except ValueError,e:
                print self._mean_data.shape, self._data[t:t+self._len,...].shape#, d_up.shape
                print "Shape missmatch occured. Seems like the last timepoint isn't completely recorded in the eeg and is therefore ignored for template-creation."
                if debug:
                     print e
        self._mean_data /= len(self._ts)
        #self._mean_data = downsample(self._mean_data)
        #if debug:
            #pylab.clf()
            #pylab.ioff()
            #for i in range(self._mean_data.shape[1]):
            #    pylab.subplot(8,self._mean_data.shape[1]/8,i+1)
            #    pylab.plot(self._mean_data[:,i])
            #pylab.show()

    def subtract_mean(self, n_ma=None):
        """Subtract the mean artifact from all artifact occurences
        n_ma: Number of nearby artifacts to include in average.
          If None (default), do global subtraction
        """
        if n_ma==None:
            assert self._mean_data != None, "Before subtracting, first calulate the mean!"
            for i,t in enumerate(self._ts): #Second walk: Substracting mean
                try:
                    #self._data[t:t+self._len,...] -= downsample(self._mean_data,start=self._corrts[i])
                    self._data[t:t+self._len,...] -= self._mean_data
                except ValueError,e:
                    print "Shape missmatch occured during writing.", t
        else:
            n_ma = int(n_ma)
            assert n_ma>0 and n_ma<len(self._ts), "Wrong value for n_ma; need 0<n_ma<=len(timepoints)"
            for i,t in enumerate(self._ts): #Second walk: Substracting mean
                try:
                    i1 = i-n_ma/2 #Startindex for mean
                    i2 = i+n_ma/2 #Endindex for mean
                    #Correct these indices
                    if i1<0:
                        i1=0
                        i2 = n_ma
                    elif i2>len(self._ts):
                        i1 = len(self._ts)-n_ma
                        i2 = len(self._ts)
                    self._data[t:t+self._len,...] -= self._mean_data
                except ValueError,e:
                    print "Shape missmatch occured during writing.", t
        return self._data

    def frequency_filter(self):
        """Applies the filter that is given to the constructor using the keyword filter.
        If no filter is given, pass.
        """
        #TODO: Change behaviour. Filter each channel as a whole and use lowpass by default.
        if self._filter == None:
            pass
        else:
            for i in range(0,self._data.shape[0],10000):
                try:
                    self._data[i:i+10000,:] = self._filter(self._data[i:i+10000,:])
                except Exception, e:
                    if debug:
                        print "Error in MeanFilter.frequency_filter:", e

    def filter(self):
        """Do the filtering"""
        self.check_before_filtering()
        self.calc_mean()
        self.subtract_mean()
        #self.frequency_filter()

    data = property(get_data,set_data)
    length = property(get_length,set_length)

#########################
# Convenience-functions #
#########################
_meanfilter = None #Global Instance of MeanFilter for use with methods below

def filter(x,ts,len=None,filter=None,t_before=None):
    """Do the PCA, with array x."""
    global _meanfilter
    _meanfilter = MeanFilter(x,ts,len,filter,t_before)
    return _meanfilter.filter()

#######################################
# If called directly, do some example #
#######################################
if __name__=='__main__':

    #TODO: Make some example use
    print "Example code not implemented yet."


import argparse
import base64
import json
import logging
import os
import re
import subprocess
import warnings
import zlib
from collections import Counter
from functools import lru_cache
from math import ceil

from ScientificTopics.download import default_model_name, _get_default_model

try:
    import matplotlib.pyplot as plt
except ImportError:
    warnings.warn("Unable to find matplotlib, plotting is disabled")
    plt = None
import numpy
import pyximport

from ScientificTopics.sentence_segmentation import Tokenizer

pyximport.install(
    setup_args={'include_dirs': numpy.get_include()})
from ScientificTopics import lda_infer

class LDAInfer(object):
    """
    LDAInfer is a class that reads LDA models and infer topics for new documents.
    To use LDAInfer, four model files are required:

    1. Punkt parameter file, for segmenting paragraphs into sentences.
    2. SentencePieceModel file, for tokenizing sentences.
    3. Stopwords file.
    4. LDA result directory, for collecting LDA doc-topic/topic-word tables.

    Two methods, tokenize_sentences() and tokenize_paragraph() are used to
    tokenize text into tokens and convert into vocabulary ids.

    Two methods, infer_topic_fast() and infer_topic() are used to infer topics
    for list of token ids representing a document.
    """

    def __init__(self, lda_result_dir,
                 punkt_model, stopwords, spm_model,
                 beta=0.01, alpha=0.1, num_vocab=None):
        """
        Initialize a LDAInfer class to infer topics for new documents.

        :param lda_result_dir: Directory containing tables of a LDA model.
        :param punkt_model: Filename of Punkt model.
        :param stopwords: Filename of stopwords list.
        :param spm_model: Filename of SentencePiece model.
        :param beta: Beta of LDA model.
        :param alpha: Alpha of LDA model.
        :param num_vocab: Number of vocabularies.
        """
        self.lda_result_dir = lda_result_dir
        self.beta = beta
        self.alpha = alpha
        self.num_vocab = num_vocab
        self._ntopics = None

        if self.num_vocab is not None:
            self._betasum = self.beta * self.num_vocab

        logging.info("Loading tokenizer...")
        self.tokenizer = Tokenizer(punkt_model, spm_model)
        self.spm_tokenizer = self.tokenizer.spm_tokenizer
        logging.info("Loading stopwords...")
        with open(stopwords, encoding='utf8') as f:
            self.stopwords = [x.strip().split()[0] for x in f if x.strip()]
        self.stopwords = set(self.spm_tokenizer.PieceToId(x) for x in self.stopwords)

        self.n_t, self.n_tw = self._read_model_parameters()

    @property
    def ntopics(self):
        """
        Return number of topics.
        """
        return self._ntopics

    @lru_cache(maxsize=1000)
    def _get_p(self, word_id):
        p_array = []
        for i in range(self.ntopics):
            p = self.n_tw[i, word_id] * 1. * self.beta / (self.n_t[i] + self.beta * self.num_vocab)
            p_array.append(p)

        p_array = numpy.asarray(p_array)

        summation = numpy.sum(p_array)
        if summation == 0:
            return numpy.ones(shape=(self.ntopics,)) / self.ntopics
        else:
            return p_array / numpy.sum(p_array)

    def _propose_word(self, word_id, size=None):
        return numpy.random.choice(self.ntopics, p=self._get_p(word_id), size=size)

    def _propose_doc(self, total_tokens, size=None):
        """-Int for doc_topics[-i]"""
        rand_numbers = numpy.random.rand(*size) * (total_tokens + self.alpha * self.ntopics)
        proposal = numpy.empty(shape=size, dtype=numpy.int)

        cond = rand_numbers < total_tokens

        take_doc_topics = numpy.where(cond)
        proposal[take_doc_topics] = -(rand_numbers[take_doc_topics]).astype(numpy.int)

        take_random_topics = numpy.where(numpy.logical_not(cond))
        proposal[take_random_topics] = numpy.random.choice(self.ntopics, size=len(take_random_topics[0]))
        return proposal

    def tokenize_sentences(self, text, filter_stopwords=True):
        """
        Tokenize words in a paragraph, and return a list of lists, each list representing
        a sentence. All words are converted to id in vocabulary.

        :param text: Paragraph to tokenize
        :param filter_stopwords: Switch to filter out all stopwords
        :return: a list of lists of token ids.
        """
        token_ids_all = self.tokenizer.tokenize_ids(text)
        valid_tokens = []
        for sentence in token_ids_all:
            sentence = [x for x in sentence if x not in self.stopwords or not filter_stopwords]
            valid_tokens.append(sentence)
        return valid_tokens

    def tokenize_paragraph(self, text, filter_stopwords=True):
        """
        Tokenize words in a paragraph, and return a list of token ids.
        See also @tokenize_sentences()

        :param text: Paragraph to tokenize
        :param filter_stopwords: Switch to filter out all stopwords
        :return: a list of token ids.
        """
        token_ids_all = self.tokenize_sentences(text, filter_stopwords)
        # Collect
        token_ids_all = sum(token_ids_all, [])

        return token_ids_all

    def generate_html(self, original_tokens, doc_topics, monte_carlo_states, doc_topics_states):
        html_template_filename = os.path.join(
            os.path.dirname(os.path.realpath(__file__)),
            'plot.template.html'
        )
        with open(html_template_filename) as f:
            html_template = f.read()

        top_words = [x[1] for x in sorted(self.get_top_words().items(), key=lambda x: x[0])]

        data_base64 = base64.b64encode(zlib.compress(json.dumps({
            'top_words': top_words,
            'states': monte_carlo_states.tolist(),
            'n_tw': doc_topics_states.tolist(),
            'tokens': [(x not in self.stopwords, self.spm_tokenizer.IdToPiece(x)) for x in original_tokens]
        }).encode())).decode()
        return html_template.replace('{{data_base64}}', data_base64)

    def infer_topic_fast(self, doc, iterations=1000, mh_steps=2, plot_mc_states=False):
        """
        Infer topics for a new document. This method uses Cython implementation so it is
        relatively faster than infer_topic().

        Upon finishing, three objects will be returned:
        1. Topic assignments for each word.
        2. Topic assignments of each iteration of the entire Monte-Carlo process.
        3. Number of topic assignments of each iteration.

        :param doc: List of tokens representing a document.
        :param iterations: Number of Monte-Carlo iterations.
        :param mh_steps: Number of Metropolis-Hastings steps for each iteration.
        :param plot_mc_states: Whether to plot the MC states of all iterations.
        :return:
        """
        doc_topics, monte_carlo_states, doc_topics_states, doc_probabilities = lda_infer.infer_topic(
            numpy.asarray(doc, dtype=numpy.uint32),
            self.n_t, self.n_tw,
            self.alpha, self.beta, iterations, mh_steps
        )

        if plot_mc_states:
            for i in range(self.ntopics):
                t_arr = doc_topics_states[:, i]
                if numpy.sum(t_arr) > iterations:
                    plt.plot(list(range(iterations)), t_arr, linewidth=1, label='Topic %d' % i)
            plt.legend()
            plt.show()

        return doc_topics, monte_carlo_states, doc_topics_states, doc_probabilities

    def infer_topic(self, doc, iterations=1000, mh_steps=2, plot_mc_states=False):
        """
        Infer topics for a new document. This method uses Cython implementation so it is
        relatively faster than infer_topic().

        Upon finishing, three objects will be returned:
        1. Topic assignments for each word.
        2. Topic assignments of each iteration of the entire Monte-Carlo process.
        3. Number of topic assignments of each iteration.

        :param doc: List of tokens representing a document.
        :param iterations: Number of Monte-Carlo iterations.
        :param mh_steps: Number of Metropolis-Hastings steps for each iteration.
        :param plot_mc_states: Whether to plot the MC states of all iterations.
        :return:
        """
        doc_topics = numpy.random.choice(self.ntopics, size=(len(doc),))
        topic_counter = numpy.zeros((self.ntopics,), dtype=numpy.int)
        for t, c in Counter(doc_topics).items():
            topic_counter[t] += c

        monte_carlo_states = numpy.empty((iterations, len(doc)), dtype=numpy.int)
        doc_topics_states = numpy.empty((iterations, self.ntopics), dtype=numpy.int)

        word_proposals = [self._propose_word(x, size=(iterations, mh_steps)) for x in doc]
        doc_proposals = [self._propose_doc(len(doc), size=(iterations, mh_steps)) for _ in doc]
        random_numbers = numpy.random.rand(iterations, len(doc), mh_steps, 2)

        for i in range(iterations):
            monte_carlo_states[i, :] = doc_topics
            doc_topics_states[i, :] = topic_counter

            for j, word_id in enumerate(doc):
                original_topic = doc_topics[j]

                topic = original_topic
                for k in range(mh_steps):
                    new_topic = word_proposals[j][i, k]

                    if new_topic != topic:
                        # Simplified version of cxx::LightDocSampler::Sample()
                        n_td_alpha = topic_counter[new_topic] + self.alpha
                        n_sd_alpha = topic_counter[topic] + self.alpha

                        if topic == original_topic:
                            n_sd_alpha -= 1
                        if new_topic == original_topic:
                            n_td_alpha -= 1

                        pi = n_td_alpha / n_sd_alpha

                        if random_numbers[i, j, k, 0] < pi:
                            topic = new_topic

                    doc_proposal = doc_proposals[j][i, k]
                    new_topic = doc_topics[-doc_proposal] if doc_proposal < 0 else doc_proposal

                    if new_topic != topic:
                        if new_topic != original_topic:
                            n_tw_beta = self.n_tw[new_topic, word_id] + self.beta
                            n_t_beta_sum = self.n_t[new_topic] + self._betasum
                            pi_t = n_tw_beta / n_t_beta_sum
                        else:
                            n_td_alpha = topic_counter[new_topic] + self.alpha - 1
                            n_tw_beta = self.n_tw[new_topic, word_id] + self.beta
                            n_t_beta_sum = self.n_t[new_topic] + self._betasum
                            proposal_t = topic_counter[new_topic] + self.alpha
                            pi_t = (n_td_alpha * n_tw_beta) / (n_t_beta_sum * proposal_t)

                        if topic != original_topic:
                            n_sw_beta = self.n_tw[topic, word_id] + self.beta
                            n_s_beta_sum = self.n_t[topic] + self._betasum
                            pi_s = n_sw_beta / n_s_beta_sum
                        else:
                            n_sd_alpha = topic_counter[topic] + self.alpha
                            n_sw_beta = self.n_tw[topic, word_id] + self.beta
                            n_s_beta_sum = self.n_t[topic] + self._betasum
                            proposal_s = topic_counter[topic] + self.alpha
                            pi_s = (n_sd_alpha * n_sw_beta) / (n_s_beta_sum * proposal_s)

                        pi = pi_t / pi_s

                        if random_numbers[i, j, k, 1] < pi:
                            topic = new_topic

                if topic != original_topic:
                    topic_counter[original_topic] -= 1
                    topic_counter[topic] += 1

                doc_topics[j] = topic

        if plot_mc_states:
            for i in range(self.ntopics):
                t_arr = doc_topics_states[:, i]
                if numpy.sum(t_arr) > iterations:
                    plt.plot(list(range(iterations)), t_arr, linewidth=1, label='Topic %d' % i)
            plt.legend()
            plt.show()

        return doc_topics, monte_carlo_states, doc_topics_states

    def get_top_words(self):
        top_words = {}
        for topic in range(len(self.n_tw)):
            sort_index = self.n_tw[topic].argsort()[:-21:-1]
            top_count = self.n_tw[topic][sort_index]
            string_builder = []
            for word_id, count in zip(sort_index, top_count):
                piece_id = self.spm_tokenizer.IdToPiece(int(word_id)) if self.spm_tokenizer is not None else str(
                    word_id)
                string_builder.append(piece_id)
            top_words[topic] = string_builder
        return top_words

    def find_possible_stopwords(self):
        stopwords = Counter()
        for topic in range(len(self.n_tw)):
            total_counts = sum(self.n_tw[topic])
            sort_index = self.n_tw[topic].argsort()[:-21:-1]
            top_count = self.n_tw[topic][sort_index]
            string_builder = []
            for word_id, count in zip(sort_index, top_count):
                piece_id = self.spm_tokenizer.IdToPiece(int(word_id)) if self.spm_tokenizer is not None else str(
                    word_id)
                stopwords[piece_id] += 1
                string_builder.append('%s (%.4f)' % (piece_id, count / total_counts))
            logging.info('Topic %d: %s', topic, ', '.join(string_builder))

        for word, count in sorted(stopwords.items(), key=lambda x: x[1], reverse=True):
            if count > self.ntopics / 10:
                logging.info('Possible stopword: %s in %d topics', word, count)

    def _read_model_parameters_try_cache(self):
        """
        Read cached LDA model tables n_tw and n_t.
        """
        cache_filename = os.path.join(self.lda_result_dir, 'cached_parameters.npz')

        if not os.path.exists(cache_filename):
            raise FileNotFoundError()

        cache_timestamp = os.path.getmtime(cache_filename)
        model_parameters = [x for x in os.listdir(self.lda_result_dir) if re.match(r'server_\d_table_\d.model', x)]
        total_nodes = max(int(x.split('_')[1]) for x in model_parameters)
        for i in range(total_nodes):
            f_timestamp = os.path.getmtime(os.path.join(self.lda_result_dir, 'server_%d_table_1.model' % i))
            if f_timestamp > cache_timestamp:
                raise FileNotFoundError()
            f_timestamp = os.path.getmtime(os.path.join(self.lda_result_dir, 'server_%d_table_0.model' % i))
            if f_timestamp > cache_timestamp:
                raise FileNotFoundError()

        logging.info('Found cached parameters: %s, loading cache...', cache_filename)
        data = numpy.load(cache_filename)
        n_t = data['n_t']
        n_tw = data['n_tw']
        self._ntopics = n_t.shape[0]

        return n_t, n_tw

    def _read_model_parameters(self):
        """
        Read LDA model tables n_tw and n_t.
        """
        try:
            return self._read_model_parameters_try_cache()
        except FileNotFoundError:
            pass

        model_parameters = [x for x in os.listdir(self.lda_result_dir) if re.match(r'server_\d_table_\d.model', x)]
        total_nodes = max(int(x.split('_')[1]) for x in model_parameters)

        # Alphas
        logging.info('Reading word count by topic...')
        n_t = {}
        for i in range(total_nodes):
            with open(os.path.join(self.lda_result_dir, 'server_%d_table_1.model' % i)) as f:
                line = f.readline()
                if not line:
                    continue
                for topic, topic_count in re.findall(r'(\d+):(\d+)', line):
                    n_t[int(topic)] = int(topic_count)
                logging.info('Loaded topic summary table %s', 'server_%d_table_1.model' % i)

        self._ntopics = max(n_t.keys()) + 1
        logging.info('Found %d topics', self.ntopics)
        n_t = numpy.asarray([x[1] for x in sorted(n_t.items())], dtype=numpy.uint32)
        assert len(n_t) == self.ntopics
        logging.debug('Word count by topics: %r', n_t)

        # Word beta
        logging.info('Reading word count by topic and word id...')
        n_tw = numpy.zeros((self.ntopics, self.num_vocab), dtype=numpy.uint32)
        for i in range(total_nodes):
            with open(os.path.join(self.lda_result_dir, 'server_%d_table_0.model' % i)) as f:
                for line in f:
                    if not line.strip():
                        break

                    matches = re.findall(r'(\d+)(?=\s)|(\d+):(\d+)', line)
                    word_id = int(matches[0][0])
                    for _, topic, topic_count in matches[1:]:
                        n_tw[int(topic), word_id] = int(topic_count)

                logging.info('Loaded word topic table %s', 'server_%d_table_0.model' % i)

        numpy.savez(
            os.path.join(self.lda_result_dir, 'cached_parameters.npz'),
            n_t=n_t,
            n_tw=n_tw
        )

        return n_t, n_tw

def load_model(model_name=default_model_name, topic_model=0):
    params = _get_default_model(model_name, topic_model)
    return LDAInfer(**params)

def gen_blocks(input_file, nodes, block_size, dump_binary):
    logging.info('Scanning input file...')

    total_docs = 0
    total_tokens = 0
    with open(input_file + '.libsvm') as f:
        for line in f:
            tokens = sum([int(x.split(':')[1]) for x in line.split('\t')[1].split(' ')])
            total_tokens += tokens
            total_docs += 1

    total_block_size = (total_tokens * 2 + total_docs + 2) * 4 / 1024 / 1024
    logging.info('In total %d documents, %d tokens, estimated total block size: %.2f MB',
                 total_docs, total_tokens, total_block_size)

    num_blocks = int(ceil(total_block_size / nodes / block_size))
    tokens_per_block = int(ceil(total_tokens / nodes / num_blocks))

    logging.info('Scheduling %d nodes with %d blocks, each block has %d tokens',
                 nodes, num_blocks, tokens_per_block)

    dump_binary_results = []

    with open(input_file + '.libsvm') as f:
        for node in range(nodes):
            for block in range(num_blocks):
                tokens_this_block = 0
                docs_this_block = 0

                block_input = input_file + '.libsvm.%d.%d' % (node, block)
                with open(block_input, 'w') as f_output:
                    while tokens_this_block < tokens_per_block:
                        try:
                            line = f.readline()
                            if not line.strip():
                                raise EOFError('empty of file reached')
                        except (EOFError, IOError):
                            break

                        tokens = sum([int(x.split(':')[1]) for x in line.split('\t')[1].split(' ')])
                        tokens_this_block += tokens
                        f_output.write(line)
                        docs_this_block += 1

                    logging.info('Node %d, block %d has %d docs and %d tokens',
                                 node, block, docs_this_block, tokens_this_block)

                    if dump_binary is not None:
                        cmd = [
                            dump_binary,
                            block_input, input_file + '.vocab',
                            os.path.realpath(os.path.dirname(block_input)), str(node), str(block)
                        ]
                        logging.info('Executing dump_binary: %r', cmd)
                        dump_binary_results.append(
                            subprocess.Popen(cmd))

    logging.info('Waiting for dump_binary to finish.')
    for i in dump_binary_results:
        i.wait()

def analyze_results(input_dir, spm_model):
    inferer = LDAInfer(input_dir, spm_model=spm_model)
    inferer.find_possible_stopwords()

def infer(input_dir,
          infer_input, infer_output,
          params, spm_model, stopwords,
          alpha, beta, num_vocab, generate_html):
    logging.info('Loading model...')
    inferer = LDAInfer(input_dir, punkt_model=params, spm_model=spm_model, stopwords=stopwords,
                       beta=beta, alpha=alpha, num_vocab=num_vocab)

    # top_words = [x[1] for x in sorted(inferer.get_top_words().items(), key=lambda x: x[0])]
    # inferer.find_possible_stopwords()

    with open(infer_input, encoding='utf8') as input_file, open(infer_output, 'w') as output_file:
        logging.info('Starting to infer topics on new documents.')
        for i, line in enumerate(input_file):
            if not line.strip():
                logging.warning('Found an empty line at line %d', i)
                continue

            logging.info('Tokenizing paragraph.')
            token_ids_all = inferer.tokenize_paragraph(line, filter_stopwords=False)
            token_ids_stop = inferer.tokenize_paragraph(line, filter_stopwords=True)

            logging.info('Inferring paragraph.')
            topics, mc_states, n_tw_states, _ = inferer.infer_topic_fast(token_ids_stop)

            if generate_html:
                with open('plot.%d.html' % i, 'w', encoding='utf8') as f:
                    f.write(inferer.generate_html(token_ids_all, topics, mc_states, n_tw_states))

            topic_count = Counter(topics)
            builder = []
            for topic in sorted(topic_count):
                builder.append('%d:%d' % (topic, topic_count[topic]))
            output_file.write(' '.join(builder))
            output_file.write('\n')
            logging.info('Inferred %d tokens: %s',
                         len(token_ids_stop), ' '.join('%d:%d' % x for x in Counter(token_ids_stop).items()))

def main():
    parser = argparse.ArgumentParser(description='LDA tools for sentence segmenter and tokenizer.')
    parser.add_argument('--verbose', action='store_true', default=False, help='More debugging logging.')

    subparsers = parser.add_subparsers(help='Action to execute.', dest='action')

    download_paragraph = subparsers.add_parser("gen-blocks", help='Generate blocks for LDA use.')
    download_paragraph.add_argument('--input', action='store', type=str, required=True,
                                    help='Input corpus, for example corpus.lightlda, '
                                         'must has *.lightlda.libsvm and *.lightlda.vocab.')
    download_paragraph.add_argument('--nodes', action='store', type=int, required=True,
                                    help='How many nodes to use.')
    download_paragraph.add_argument('--block-size', action='store', type=int, default=500,
                                    help='The average block size in MB. See lightlda/dump_binary.cpp for block format.')
    download_paragraph.add_argument('--dump-binary', action='store', type=str, default=None,
                                    help='If provided the executive path for dump-binary, I will help you execute it.')

    analyze_arguments = subparsers.add_parser("analyze", help='Analyze LDA result.')
    analyze_arguments.add_argument('--lda-result-dir', action='store', type=str, required=True,
                                   help='LDA result dir')
    analyze_arguments.add_argument('--spm', action='store', type=str, required=True,
                                   help='SPM model file')

    infer_arguments = subparsers.add_parser("infer", help='Infer LDA topics.')
    default_model = {}
    try:
        default_model = _get_default_model()
    except FileNotFoundError:
        pass
    infer_arguments.add_argument('--lda-result-dir', action='store', type=str,
                                 default=default_model['lda_result_dir'],
                                 help='LDA result dir')
    infer_arguments.add_argument('--input', action='store', type=str, required=True,
                                 help='Paragraphs to infer')
    infer_arguments.add_argument('--output', action='store', type=str, required=True,
                                 help='Inference output')
    infer_arguments.add_argument('--model-dir', action='store', type=str, default=None,
                                 help='Directory of the model that contains a config.json')
    infer_arguments.add_argument('--params', action='store', type=str,
                                 default=default_model['punkt_model'],
                                 help='Punkt parameters')
    infer_arguments.add_argument('--spm', action='store', type=str,
                                 default=default_model['spm_model'],
                                 help='SPM model file')
    infer_arguments.add_argument('--stopwords', action='store', type=str,
                                 default=default_model['stopwords'],
                                 help='List of stopwords')
    infer_arguments.add_argument('--num-vocab', action='store', type=int,
                                 default=default_model['num_vocab'],
                                 help='Number of vocabulary')
    infer_arguments.add_argument('--alpha', action='store', type=float,
                                 default=default_model['alpha'],
                                 help='Topics prior')
    infer_arguments.add_argument('--beta', action='store', type=float,
                                 default=default_model['beta'],
                                 help='Words prior')
    infer_arguments.add_argument('--generate-html', action='store_true', default=False,
                                 help='Generate HTML visualizations')

    args = parser.parse_args()

    if args.verbose:
        logging.basicConfig(format='%(asctime)s - %(levelname)s - %(message)s',
                            level=logging.DEBUG)
    else:
        logging.basicConfig(format='%(asctime)s - %(levelname)s - %(message)s',
                            level=logging.INFO)

    if args.action == 'gen-blocks':
        gen_blocks(args.input, args.nodes, args.block_size, args.dump_binary)
    elif args.action == 'analyze':
        analyze_results(args.lda_result_dir, args.spm)
    elif args.action == 'infer':
        if args.model_dir is not None:
            d = os.path.realpath(args.model_dir)
            with open(os.path.join(d, 'config.json')) as f:
                config = json.load(f)
                kwargs = {
                    'input_dir': d,
                    'params': os.path.join(d, config['params']),
                    'spm_model': os.path.join(d, config['spm']),
                    'stopwords': os.path.join(d, config['stopwords']),
                    'alpha': config['alpha'],
                    'beta': config['beta'],
                    'num_vocab': config['num_vocab'],
                }
        else:
            kwargs = {
                'input_dir': args.lda_result_dir,
                'params': args.params,
                'spm_model': args.spm,
                'stopwords': args.stopwords,
                'alpha': args.alpha,
                'beta': args.beta,
                'num_vocab': args.num_vocab,
            }
        input_dir = kwargs.pop('input_dir')
        infer(input_dir,
              args.input, args.output,
              generate_html=args.generate_html,
              **kwargs)
    else:
        parser.print_usage()

if __name__ == '__main__':
    main()


# -*- coding: utf-8 -*-
"""
.. wikisection:: overview
    :title: (6) Alignment-free local similarity detection with Word-Blot

    The :mod:`biseqt.blot` module implements the WordBlot algorithm.

    >>> from biseqt.blot import WordBlot
    >>> from biseqt.sequence import Sequence, Alphabet
    >>> from biseqt.stochastics import rand_seq, MutationProcess
    >>> A = Alphabet('ACGT')
    >>> S = rand_seq(A, 100)
    >>> M = MutationProcess(A, go_prob=.1, ge_prob=.1, subst_probs=.2)
    >>> T, _ = M.mutate(S)
    >>> WB = WordBlot(S, T, wordlen=5, alphabet=A, g_max=.5, \
                            sensitivity=.99)
    >>> list(WB.similar_segments(50, .7))
    [(((-28, 7), (0, 99)), 1.1988215486845972, 0.7630903108462143)]
    >>>
"""
import sys
import warnings
import numpy as np
import logging
from itertools import groupby, product
from scipy.special import erfcinv
from scipy.spatial import cKDTree
from .seeds import SeedIndex, SeedIndexMultiple
from .kmers import as_kmer_seq
from .util import Logger

# so we can catch numpy warnings
warnings.filterwarnings('error')

# A band (i-r, i+r) is considered of interest if xs[i] >threshold. This
# function returns a maximal (disjoint) set of bands of interest (i.e if two
# bands overlap they are reported as one bigger band).
def find_peaks(xs, rs, threshold):
    """Finds maximal (disjoint) peak regions in a sequence of real numbers.
    Each value that is at least as large as the threshold constitutes the
    center of a peak with radius according to its position. In the output all
    overlapping peaks are merged into maximal peaks.

    Args:
        xs: the 1D data sequence of interest
        rs: the radii for peaks defined at every point exceeding threshold,
            could be a 1D sequence or a fixed number,
        threshold: cutoff value to compare with ``xs`` values.

    Returns:
        list (tuple): A least of "peaks", each a tuple of ``(left, right)``
        coordinates in ``xs``. Returned peaks are guaranteed to be disjoint.
    """
    peaks = []
    cur_peak = None
    for idx, x in enumerate(xs):
        radius = rs[idx] if isinstance(rs, list) else rs
        assert isinstance(radius, int)
        if x < threshold:
            continue
        peak_l, peak_r = max(0, idx - 1), min(len(xs) - 1, idx + 1)
        if cur_peak is None:
            cur_peak = (peak_l, peak_r)
            continue
        if peak_l < cur_peak[1] + radius:  # overlaps with cur_peak
            assert peak_r >= cur_peak[1]
            cur_peak = (cur_peak[0], peak_r)
        else:
            peaks.append(cur_peak)
            cur_peak = (peak_l, peak_r)
    if cur_peak is not None:
        peaks.append(cur_peak)
    return [(int(l), int(r)) for (l, r) in peaks]

def wall_to_wall_distance(len0, len1, diag):
    """Wall to wall distance for a diagonal position

    .. math::
        L = \\min(l_0 - d, l_1) + \\min(d, 0)

    with :math:`l_0,l_1` being the length of the sequences (i.e ``len0`` and
    ``len1`` arguments) and :math:`d` the starting diagonal (i.e ``diag``
    argument).
    """
    return min(len0 - diag, len1) + min(diag, 0)

def expected_overlap_len(len0, len1, diag, gap_prob):
    """Calculates the expected length of an overlap alignment given its starting
    coordinates:

    .. math::
        K = \\left(\\frac{2}{2 - g}\\right) L

    where :math:`L` is the :func:`wall_to_wall_distance`.

    Args:
        len0 (int): Length of the 1st sequence.
        len1 (int): Length of the 2nd sequence.
        diag (int): Starting diagonal of alignments to consider.
        gap_prob (float): Probability of indels occuring at any position.
    Returns:
        int: Expected length of an overlap alignment.
    """
    L = wall_to_wall_distance(len0, len1, diag)
    expected_len = (2. / (2 - gap_prob)) * L
    assert expected_len >= 0
    return int(np.ceil(expected_len))

# band radius for edit path of length K
def band_radius(expected_len, gap_prob, sensitivity):
    """Calculates the smallest band radius in the dynamic programming table
    such that an alignment of given expected length, with the given gap
    probability, stays entirely within the diagonal band with probability given
    as sensitivity. This is given by:

    .. math::
        r = \\mathrm{erf}^{-1}\\left(1-\\epsilon\\right)
            \\sqrt{2gK}

    where :math:`g` is the gap probability, :math:`1-\\epsilon` is the desired
    sensitivity, and :math:`K` is the given expected length.

    Args:
        expected_len (int):
            minimum expected length of similar region.
        gap_prob (float): Probability of indels occuring at any position.
        sensitivity (float): The probability that an alignment with given gap
            probability remains entirely within the band.
    Returns:
        int: The smallest band radius guaranteeing the required sensitivity.
    """
    assert 0 < gap_prob < 1 and 0 < sensitivity < 1
    epsilon = 1. - sensitivity
    C = erfcinv(epsilon) * np.sqrt(2 * gap_prob)
    radius = C * np.sqrt(expected_len)
    return max(1, int(np.ceil(radius)))

def band_radii(expected_lens, gap_prob, sensitivity):
    """Same as :func:`band_radius` but for bulk calculations (e.g. finding
    overlaps).

    Args:
        expected_lens (list):
            List of integers for which to calculate band radii.
        gap_prob (float):
            As in :func:`band_radius`.
        sensitivity (float):
            As in :func:`band_radius`.
    """
    assert 0 < gap_prob < 1 and 0 < sensitivity < 1
    epsilon = 1. - sensitivity
    C = erfcinv(epsilon) * np.sqrt(2 * gap_prob)
    return np.array([max(1, int(np.ceil(C * np.sqrt(K))))
                     for K in expected_lens])

def H0_moments(alphabet_len, wordlen, area):
    """The mean and standrad deviation of the limiting normal distribution
    under the :math:`H_0` (unrelated) model given by:

    .. math::
        \\begin{aligned}
            \mu_0 & = Ap^w \\\\
            \sigma_0^2 & =
                A\\left[(1 - p^w)\left(p^w + \\frac{2p^{w+1}}{1 - p}\\right)
                - 2wp^{2w}\\right]
        \\end{aligned}

    where :math:`w` is the word length, :math:`A` is the area of the ROI, and
    :math:`p = \\frac{1}{|\Sigma|}` with :math:`|\Sigma|` being the alphabet
    length.
    """
    p_H0 = 1. / alphabet_len
    pw_H0 = p_H0 ** wordlen

    mu_H0 = area * pw_H0
    sd_H0 = np.sqrt(area * (
        (1 - pw_H0) * (pw_H0 + 2 * p_H0 * pw_H0 / (1 - p_H0)) -
        2 * wordlen * pw_H0 ** 2
    ))
    return mu_H0, sd_H0

def H1_moments(alphabet_len, wordlen, area, seglen, p_match):
    """The mean and standrad deviation of the limiting normal distribution under
    the :math:`H_1` (related) model given by:

    .. math::
        \\begin{aligned}
            \mu_1 & = \mu_0 + Kp^w \\\\
            \sigma_1^2 & = \sigma_0^2
                + 2K\\left[(1 - p^w)\left(p^w + \\frac{2p^{w+1}}{1 - p}\\right)
                - 2wp^{2w}\\right]
        \\end{aligned}

    where :math:`w` is the word length, :math:`K` is the similarity length, and
    :math:`p` is the match probability.
    """
    mu_H0, sd_H0 = H0_moments(alphabet_len, wordlen, area)

    p_H1 = p_match
    if p_H1 == 1.:
        # we can't let p_H1 == 1 because we get division by zero below
        p_H1 = 1 - np.finfo(float).eps
    pw_H1 = p_H1 ** wordlen

    mu_H1 = mu_H0 + seglen * pw_H1
    sd_H1 = np.sqrt(sd_H0 ** 2 + seglen * (
        (1 - pw_H1) * (pw_H1 + 2 * p_H1 * pw_H1 / (1 - p_H1)) -
        2 * wordlen * pw_H1 ** 2
    ))
    return mu_H1, sd_H1

# FIXME the fact that we have organized our data as self.S and self.T is the
# main blocker for merging pairwise and multiple sequence implementations.
# Also involved: SeedIndexMultiple
class WordBlot(SeedIndex):
    """A similarity finder based on m-dependent CLT statistics.

    Attributes:
        g_max (float):
            Upper bound for indel probabilities in mutation model.
        sensitivity (float):
            Desired sensitivity of bands.
    """
    def __init__(self, S, T, g_max=None, sensitivity=None, **kw):
        assert 0 < g_max < 1 and 0 < sensitivity < 1
        self.g_max = g_max
        self.sensitivity = sensitivity
        super(WordBlot, self).__init__(S, T, **kw)

    def score_num_seeds(self, **kw):
        """Calculates our key central statistics based on m-dependent CLT. For
        a given observation of number of seeds in a region of interest (ROI),
        calculates the z-core against the :math:`H_0` (unrelated) and
        :math:`H_1` (related) models. In either case, the distribution of
        number of seeds, being a sum of m-dependent identically distributed
        random variables, is asymptotically normal with parameters given by
        :func:`H0_moments`, :func:`H1_moments`.

        Keyword Args:
            num_seeds (int):
                Number of observed seed in the ROI.
            area (int|float):
                Area of the ROI.
            seglen (int):
                Similar segment length in H1 model.
            p_match (float):
                Expected match probability at any given position.

        Returns:
            tuple (float): z-scores in H0 and H1 models
        """
        num_seeds = kw['num_seeds']
        area = kw['area']

        if area == 0:
            return float('-inf'), float('-inf')

        mu_H0, sd_H0 = H0_moments(len(self.alphabet), self.wordlen, area)
        mu_H1, sd_H1 = H1_moments(len(self.alphabet), self.wordlen, area,
                                  kw['seglen'], kw['p_match'])

        z_H0 = (num_seeds - mu_H0) / sd_H0  # score under H0
        z_H1 = (num_seeds - mu_H1) / sd_H1  # score under H1
        return z_H0, z_H1

    def band_radius(self, K):
        """Wraps :func:`band_radius` with our mutation parameters and sequence
        lengths.

        Args:
            K (int): expected alignment length of interest.

        Returns:
            int: radius of band for desired :attr:`sensitivity`.
        """
        return band_radius(K, self.g_max, self.sensitivity)

    def segment_dims(self, d_band=None, a_band=None):
        """Calculate the edit path length :math:`K` and the area of the given
        diagonal/antiodiagonal segment.

        Keyword Args:
            d_band (tuple): lower and upper diagonals limiting the segment.
            a_band (tuple): lower and upper antidiagonal positions limiting the
                segment.

        Returns:
            tuple: edit path length and segment area.
        """
        a_min, a_max = a_band
        d_min, d_max = d_band
        K = (a_max - a_min) / 2
        A = (d_max - d_min) * K
        return K, A

    def estimate_match_probability(self, num_seeds, d_band=None, a_band=None):
        """Estimate the edit path match probability given the provided observed
        number of seeds :math:`n` in given sigment:

        .. math::
            \hat{p} = \\left(\\frac{n - Ap_\circ^w}{K}\\right)^{\\frac{1}{w}}

        where :math:`n, A, p_\circ, K, w` are the number of seeds, segment
        area, null probability of matching nucleotides, segment length, and the
        word length, respectively.

        Args:
            num_seeds (int): number of seeds observed in segment.

        Keywords Args:
            d_band (tuple):
                lower and upper diagonals limiting the segment.
            a_band (tuple):
                lower and upper antidiagonal positions limiting the segment.

        Returns:
            float: estimated match probability
        """
        K, area = self.segment_dims(d_band=d_band, a_band=a_band)
        word_p_null = (1./len(self.alphabet)) ** self.wordlen
        # NOTE K effectively has become the projected alignment length not the
        # full length. This is REALLY important in interpretation but I think
        # must things are currently consistent (TODO full review needed).
        word_p = (num_seeds - area * word_p_null) / K
        try:
            match_p = np.exp(np.log(word_p) / self.wordlen)
        except Warning:
            # presumably this happened because word_p was too small for log
            match_p = 0
        return min(match_p, 1)

    @classmethod
    def find_all_neighbors(cls, seeds, d_radius, a_radius):
        """For each seed finds all seeds in its neighborhood defined by:

        .. math::

            U_{(d, a)} = \\{(d', a'): |d - d'| < r_d,  |a - a'| < r_a \\}

        This is done using a Quad-Tree in :math:`O(m lg m)` time where m is the
        number of seeds.

        Returns:
            list: tuples ``((d, a), neighs)`` where ``neighs`` is a list of
                  neighbor indices.
        """
        # normalize the two diameters so we can use a standard L∞ neighborhood.
        # typically a_diam is larger, so scale up d values proportionally
        d_coeff = 1. * a_radius / d_radius
        radius = a_radius

        all_seeds = list(cls.to_diagonal_coordinates(i, j) for i, j in seeds)
        if not all_seeds:
            return []
        all_seeds_scaled = np.array([(d * d_coeff, a) for d, a in all_seeds])
        quad_tree = cKDTree(all_seeds_scaled)
        all_neighs = quad_tree.query_ball_tree(quad_tree, radius,
                                               p=float('inf'))
        # all_neighs[i] is the indices of the neighbors of all_seeds[i]; this
        # always contains the seed itself (i.e always: i in neighs[i])
        for idx, _ in enumerate(all_neighs):
            all_neighs[idx].remove(idx)
        return zip(all_seeds, all_neighs)

    def score_seeds(self, K):
        """Find the neighbors of each seed in the sense of
        :func:`find_all_neighbors` and estimates the match probability of the
        segment centered at the coordinates of each seed.

        Args:
            K (int): the similarity legnth of interest that dictates
                neighborhood shapes and estimated match probabilities.

        Returns:
            list(dict): List of dictionaries with keys: ``seed`` (coordinates
            of exactly matching kmer in diagonal coordinates), ``neighs`` (list
            of indices of neighbors of this seed in the appropriate diagonal
            strip), ``p`` the estimated match probability
            of a segment centered at the seed.
        """
        d_radius = int(np.ceil(self.band_radius(K)))
        a_radius = K

        seeds_with_neighs = self.find_all_neighbors(
            self.seeds(exclude_trivial=True), d_radius, a_radius
        )

        # n is the number of seeds in neighborhood excluding the center seed,
        def _p(d, a, n):
            d_band = (d - d_radius, d + d_radius)
            a_band = (a - a_radius, a + a_radius)
            return self.estimate_match_probability(n + 1, d_band=d_band,
                                                   a_band=a_band)

        return [{'seed': (d, a), 'neighs': neighs, 'p': _p(d, a, len(neighs))}
                for (d, a), neighs in seeds_with_neighs]

    def similar_segments(self, K_min, p_min, at_least_one=False):
        """Find all maximal local similarities of given minium length and match
        probability. Additionally for each segment, the match probability is
        estimated and H0/H1 scores are calculated.

        Args:
            K_min (int):
                minimum required length of similarity.
            p_min (float):
                Minimum required match probability at each position.
            score (bool):
                Whether to score the segment against H1 by counting seeds
                inside.

        Yields:
            dict: dictionary with keys: ``segment`` (coordinates of similar
            region in diagonal coordinates ``((d_min, d_max), (a_min,
            a_max))``)), ``p`` the estimated match probability, and ``score``
            the H1 z-score if keyword argument ``score`` is true.
        """
        self.log('finding local similarities between %s and %s' %
                 (self.S.content_id[:8], self.T.content_id[:8]))
        d_radius = int(np.ceil(self.band_radius(K_min)))
        a_radius = K_min
        scored_seeds = self.score_seeds(K_min)

        def _update_seg(seg, seed):
            d, a = seed
            if seg is None:
                d_min, d_max = d - d_radius, d + d_radius
                a_min, a_max = a - a_radius, a + a_radius
            else:
                (d_min, d_max), (a_min, a_max) = seg
                d_min = min(d - d_radius, d_min)
                d_max = max(d + d_radius, d_max)
                a_min = min(a - a_radius, a_min)
                a_max = max(a + a_radius, a_max)
            return (d_min, d_max), (a_min, a_max)

        avail = [rec['p'] >= p_min for rec in scored_seeds]
        if not any(avail) and at_least_one:
            # we're obliged to return something, let the highest probability
            # seed go through.
            assert len(scored_seeds), 'no seeds found while at_least_one=True'
            avail[np.argmax([rec['p'] for rec in scored_seeds])] = True
        while True:
            try:
                seed_idx = avail.index(True)
            except ValueError:
                break
            stack = [seed_idx]
            avail[seed_idx] = False
            ps_in_seg = [scored_seeds[seed_idx]['p']]
            seg = None
            while stack:
                idx = stack.pop()
                ps_in_seg.append(scored_seeds[idx]['p'])
                seg = _update_seg(seg, scored_seeds[idx]['seed'])
                for neigh in scored_seeds[idx]['neighs']:
                    if avail[neigh]:
                        stack.append(neigh)
                        avail[neigh] = False
            if seg is None:
                break
            else:
                (d_min, d_max), (a_min, a_max) = seg
                d_min = min(len(self.S), max(d_min, -len(self.T)))
                d_max = min(len(self.S), max(d_max, -len(self.T)))
                a_min = max(a_min, 0)
                a_max = min(a_max, len(self.S) + len(self.T))
                seg = (d_min, d_max), (a_min, a_max)
            # NOTE the following is more justifiable but it matches the
            # average. TODO turn this in into an experiment to justify
            # p_hat = self.estimate_match_probability(
            #   n, d_band=seg[0], a_band=seg[1])
            p_hat = sum(ps_in_seg) / len(ps_in_seg)
            res = {'segment': seg, 'p': p_hat}
            n = self.seed_count(d_band=seg[0], a_band=seg[1])
            K_hat, area_hat = self.segment_dims(d_band=seg[0],
                                                a_band=seg[1])
            scores = self.score_num_seeds(num_seeds=n, area=area_hat,
                                          seglen=K_hat, p_match=p_hat)
            res['scores'] = scores
            yield res

class WordBlotOverlap(WordBlot):
    """A specialized version of WordBlot for detecting overlap
    (suffix-prefix) similarities between sequences (e.g. in a sequencing
    context)."""
    def score_seeds(self):
        """For each seed finds all seeds in its neighborhood defined by:

        .. math::

            U_{(d, a)} = \\{(d', a'): |d - d'| < r_d(d),  |a - a'| < r_a(d) \\}

        in such a way that each seed's neighborhood is the entire diagonal band
        containing it. Each seed recieves an estimated match probability for a
        similarity on its diagonal band and a z-score with respect to the H1
        model.

        Returns:
            list: dicts with keys ``seed`` (diagonal coordinates of seed),
                  ``p`` (estimated match probability of overlap alignment),
                  ``L`` (the length of an alignment through the seed's band
                  according to sequence lengths), ``score`` (the z-score with
                  respect to H1), and ``r`` (band radius at the seed's
                  coordinates).
        """
        def _len(d):
            return expected_overlap_len(
                len(self.S), len(self.T), d, self.g_max
            )

        def _rad(d):
            return np.ceil(self.band_radius(_len(d)))

        all_seeds = list(self.to_diagonal_coordinates(i, j)
                         for i, j in self.seeds(exclude_trivial=True))
        if not all_seeds:
            return []
        all_seeds_scaled = np.array([(d / _rad(d), ) for d, a in all_seeds])
        quad_tree = cKDTree(all_seeds_scaled)
        all_neighs = quad_tree.query_ball_tree(quad_tree, 1, p=float('inf'))
        # all_neighs[i] is the indices of the neighbors of all_seeds[i]; this
        # always contains the seed itself (i.e always: i in neighs[i])
        for idx, _ in enumerate(all_neighs):
            all_neighs[idx].remove(idx)
        seeds_with_neighs = zip(all_seeds, all_neighs)

        # n excludes the center seed itself
        def _p(d, n):
            L = _len(d)
            d_radius = int(np.ceil(self.band_radius(L)))
            area = 2 * d_radius * L
            word_p_null = (1./len(self.alphabet)) ** self.wordlen
            word_p = (n + 1 - area * word_p_null) / L
            try:
                match_p = np.exp(np.log(word_p) / self.wordlen)
            except Warning:
                # presumably this happened because word_p was too small for log
                match_p = 0
            return min(match_p, 1)

        return [{'seed': (d, a), 'r': _rad(d), 'L': _len(d),
                 'p': _p(d, len(neighs))}
                for (d, a), neighs in seeds_with_neighs]

    def highest_scoring_overlap_band(self):
        """Finds the highest scoring diagonal band according to probabiliy
        estimations of :func:`score_seeds`.

        Returns:
            dict: with same keys ``p, score, d_band`` consistent with the
                  output of :func:`score_seeds` for the highest scoring
                  diagonal band.
        """
        scored_seeds = self.score_seeds()
        if not scored_seeds:
            return None
        idx = max(range(len(scored_seeds)), key=lambda i: scored_seeds[i]['p'])
        seed, rad = scored_seeds[idx]['seed'], scored_seeds[idx]['r']
        p_hat, overlap_len = scored_seeds[idx]['p'], scored_seeds[idx]['L']
        d_band = seed[0] - rad, seed[0] + rad
        res = {'d_band': d_band, 'p': p_hat, 'len': overlap_len}
        area = 2 * rad * overlap_len
        mu_H1, sd_H1 = H1_moments(len(self.alphabet), self.wordlen, area,
                                  overlap_len, p_hat)
        num_seeds = self.seed_count(d_band=d_band)
        z_H1 = (num_seeds - mu_H1) / sd_H1
        res['score'] = z_H1
        return res

class WordBlotOverlapRef(WordBlotOverlap):
    """An in-memory, SQL-free version of :class:`WordBlotOverlap` for faster
    comparisons.  Due to implementation details the word length is constrained
    above by the available memory.

    Attributes:
        allowed_memory (int|float): allocatable memory in GB for kmers index.
    """
    def __init__(self, ref, allowed_memory=1, **kw):
        self.wordlen = kw['wordlen']
        self.alphabet = kw['alphabet']
        self.g_max = kw['g_max']
        self.sensitivity = kw['sensitivity']
        self.log_level = kw.get('log_level', logging.INFO)
        self.S = ref
        num_kmers = len(self.alphabet) ** self.wordlen
        mem_needed = sys.getsizeof(num_kmers) * num_kmers
        mem_needed_gb = np.power(2, np.log2(mem_needed) - 30)
        assert allowed_memory > 0, 'allowed memory must be positive'
        self.allowed_memory = allowed_memory
        if mem_needed_gb > self.allowed_memory:
            msg = 'not enough memory (max = %.2f GB) ' % allowed_memory
            msg += 'to store %d-mers ' % self.wordlen
            msg += '(%.2f GB needed)' % mem_needed_gb
            raise MemoryError(msg)
        self.kmer_hits = [[] for _ in range(num_kmers)]
        for pos, kmer in enumerate(as_kmer_seq(ref, self.wordlen)):
            self.kmer_hits[kmer].append(pos)
        self.T = None
        relpath = 'python-object'
        log_header = '%d-mer cache (%s)' % (self.wordlen, relpath)
        self._logger = Logger(log_level=self.log_level, header=log_header)
        self._seeds = {}

    def seeds(self, exclude_trivial=True):
        assert self.T is not None
        if self.T.content_id not in self._seeds:
            self._seeds = {self.T.content_id: []}
            for pos, kmer in enumerate(as_kmer_seq(self.T, self.wordlen)):
                for pos_ref in self.kmer_hits[kmer]:
                    if self.S == self.T and exclude_trivial and pos == pos_ref:
                        continue
                    self._seeds[self.T.content_id].append((pos_ref, pos))
        return self._seeds[self.T.content_id]

    def seed_count(self, d_band=None, a_band=None):
        assert self.T is not None
        seeds = self.seeds()
        cnt = 0
        if d_band:
            d_min, d_max = d_band
        if a_band:
            a_min, a_max = a_band
        for i, j in seeds:
            d, a = self.to_diagonal_coordinates(i, j)
            if d_band and not d_min <= d <= d_max:
                continue
            if a_band and not a_min <= a <= a_max:
                continue
            cnt += 1
        return cnt

    def score_seeds_(self, seq):
        self.T = seq
        return super(WordBlotOverlapRef, self).score_seeds()

    def highest_scoring_overlap_band(self, seq):
        self.T = seq
        return super(WordBlotOverlapRef, self).highest_scoring_overlap_band()

class WordBlotLocalRef(WordBlot):
    """An in-memory, SQL-free version of :class:`WordBlot` for faster
    comparisons. Due to implementation details the word length is constrained
    above by the available memory.

    Attributes:
        allowed_memory (int|float): allocatable memory in GB for kmers index.
    """
    def __init__(self, ref, allowed_memory=1, **kw):
        self.wordlen = kw['wordlen']
        self.alphabet = kw['alphabet']
        self.g_max = kw['g_max']
        self.sensitivity = kw['sensitivity']
        self.log_level = kw.get('log_level', logging.INFO)
        self.S = ref
        num_kmers = len(self.alphabet) ** self.wordlen
        mem_needed = sys.getsizeof(num_kmers) * num_kmers
        mem_needed_gb = np.power(2, np.log2(mem_needed) - 30)
        assert allowed_memory > 0, 'allowed memory must be positive'
        self.allowed_memory = allowed_memory
        if mem_needed_gb > self.allowed_memory:
            msg = 'not enough memory (max = %.2f GB) ' % allowed_memory
            msg += 'to store %d-mers ' % self.wordlen
            msg += '(%.2f GB needed)' % mem_needed_gb
            raise MemoryError(msg)
        self.kmer_hits = [[] for _ in range(num_kmers)]
        for pos, kmer in enumerate(as_kmer_seq(ref, self.wordlen)):
            self.kmer_hits[kmer].append(pos)
        self.T = None
        relpath = 'python-object'
        log_header = '%d-mer cache (%s)' % (self.wordlen, relpath)
        self._logger = Logger(log_level=self.log_level, header=log_header)
        self._seeds = {}

    def seeds(self, exclude_trivial=True):
        assert self.T is not None
        if self.T.content_id not in self._seeds:
            self._seeds = {self.T.content_id: []}
            for pos, kmer in enumerate(as_kmer_seq(self.T, self.wordlen)):
                for pos_ref in self.kmer_hits[kmer]:
                    if self.S == self.T and exclude_trivial and pos == pos_ref:
                        continue
                    self._seeds[self.T.content_id].append((pos_ref, pos))
        return self._seeds[self.T.content_id]

    def seed_count(self, d_band=None, a_band=None):
        assert self.T is not None
        seeds = self.seeds()
        cnt = 0
        if d_band:
            d_min, d_max = d_band
        if a_band:
            a_min, a_max = a_band
        for i, j in seeds:
            d, a = self.to_diagonal_coordinates(i, j)
            if d_band and not d_min <= d <= d_max:
                continue
            if a_band and not a_min <= a <= a_max:
                continue
            cnt += 1
        return cnt

    def score_seeds_(self, seq, K):
        self.T = seq
        return super(WordBlotLocalRef, self).score_seeds(K)

    def similar_segments(self, seq, K_min, p_min, at_least_one=False):
        self.T = seq
        kw = {'at_least_one': at_least_one}
        for res in super(WordBlotLocalRef, self).similar_segments(K_min, p_min,
                                                                  **kw):
            yield res

# FIXME lots of code duplication here, can it be cleaned up? (cf. note above
# WordBlot; the main issue is self.S and self.T being used everywhere there,
# and presumably lots of hidden pairwise assumptions). The good news is:
# every function implemented below should in principle work exactly as is for
# pairwise.
class WordBlotMultiple(SeedIndexMultiple):
    """A multiple sequence similarity finder based on m-dependent CLT
    statistics.

    Attributes:
        g_max (float):
            Upper bound for indel probabilities in mutation model.
        sensitivity (float):
            Desired sensitivity of bands.
    """
    def __init__(self, *seqs, **kw):
        g_max, sensitivity = kw.pop('g_max'), kw.pop('sensitivity')
        assert 0 < g_max < 1 and 0 < sensitivity < 1
        self.g_max = g_max
        self.sensitivity = sensitivity
        super(WordBlotMultiple, self).__init__(*seqs, **kw)

    def band_radius(self, K):
        """Wraps :func:`band_radius` with our mutation parameters and sequence
        lengths.

        Args:
            K (int): expected alignment length of interest.

        Returns:
            int: radius of band for desired :attr:`sensitivity`.
        """
        return band_radius(K, self.g_max, self.sensitivity)

    def wall_to_wall_distance(self, ds):
        """Wall to wall distance :math:`L` for a diagonal position :math:`d_1,
        \ldots, d_{n-1}`.

        The distance :math:`L` is the largest possible value of the
        antidiagonal coordinate once all diagonal coordinates are fixed.
        """
        raise NotImplementedError

    def estimate_match_probability(self, num_seeds, K, volume):
        """Estimate the edit path match probability given the provided observed
        number of seeds in given sigment.

        .. math::
            \hat{p} = \\left(
                            \\frac{n - Vp_\circ^{w(N-1)}}{K}
                      \\right)^{\\frac{1}{w(N-1)}}

        where :math:`n, N, V, p_\circ, w` are the number of seeds, number of
        sequences, segment n-d volume, null probability of matching
        nucleotides, and word length, respectively.

        Args:
            num_seeds (int|float): number of seeds observed in segment.
            K (int|float): the expected length of the alignment.
            volume (int|float): the n-d volume of the region of interest.

        Returns:
            float: estimated match probability
        """
        power = self.wordlen * (len(self.seqs) - 1)
        word_p_null = (1. / len(self.alphabet)) ** power

        if num_seeds > 0:
            word_p = (num_seeds - volume * word_p_null) / K
            try:
                match_p = np.exp(np.log(word_p) / self.wordlen)
            except Warning:
                # presumably this happened because word_p was too small for log
                match_p = 0
        else:
            match_p = 0
        return min(match_p, 1)

    def score_seeds(self, K):
        """Find the neighbors of each seed in the sense of
        :func:`find_all_neighbors` and estimates the match probability of the
        segment centered at the coordinates of each seed.

        Args:
            K (int): the similarity legnth of interest that dictates
                neighborhood shapes and estimated match probabilities.

        Returns:
            list: list of tuples ``((d, a), neighs, p)`` where ``(d, a)``
                  is the diagonal coordinates of the seed, ``neighs`` is a list
                  of integer indices of seeds in its diagonal/antidiagonal
                  neighborhood, and ``p`` is the estimated match probability of
                  a segment of length ``K`` centered at the seed.
        """
        d_radius = int(np.ceil(self.band_radius(K)))
        a_radius = int(np.ceil(len(self.seqs) * K / 2.))
        seeds_with_neighs = self.find_all_neighbors(d_radius, a_radius)
        volume = (2 * d_radius) ** (len(self.seqs) - 1) * K

        # n excludes the center seed itself
        def _p(n):
            return self.estimate_match_probability(n + 1, K, volume)

        return [{'seed': (ds, a), 'neighs': neighs, 'p': _p(len(neighs))}
                for (ds, a), neighs in seeds_with_neighs]

    def find_all_neighbors(self, d_radius, a_radius):
        """For each seed finds all seeds in its neighborhood defined by:

        .. math::

            U_{(d_1,\ldots,d_{n-1}, a)} =
                \\{(d_1', \ldots, d_{n-1}', a'):
                    |d_k - d_k'| < r_d,  |a - a'| < r_a \\}

        This is done using a kD-Tree in :math:`O(nm lg m)` time where m is the
        number of seeds and n is the number of sequences.

        Returns:
            list: tuples ``((ds, a), neighs)`` where ``neighs`` is a list of
                  neighbor indices.
        """
        self.log('finding all neighbors using a kD-tree')
        # normalize the two diameters so we can use a standard L∞ neighborhood.
        # typically a_diam is larger, so scale up d values proportionally
        d_coeff = 1. * a_radius / d_radius
        radius = a_radius

        all_seeds = list(self.seeds())
        if not all_seeds:
            return []
        all_seeds_scaled = np.array([[d * d_coeff for d in ds] + [a]
                                     for ds, a in all_seeds])
        quad_tree = cKDTree(all_seeds_scaled)
        all_neighs = quad_tree.query_ball_tree(quad_tree, radius,
                                               p=float('inf'))
        # all_neighs[i] is the indices of the neighbors of all_seeds[i]; this
        # always contains the seed itself (i.e always: i in neighs[i])
        for idx, _ in enumerate(all_neighs):
            all_neighs[idx].remove(idx)
        self.log('found all neighbors')
        return zip(all_seeds, all_neighs)

    def score_num_seeds(self, num_seeds, **kw):
        """The exrtension of :func:`WordBlot.score_num_seeds` to multiple
        sequence comparisons.

        Keyword Args:
            num_seeds (int):
                Number of observed seed in the ROI.
            volume (int|float):
                Area of the ROI.
            seglen (int):
                Similar segment length in H1 model.
            p_match (float):
                Expected match probability at any given position.

        Returns:
            tuple (float): z-scores in H0 and H1 models
        """
        p_match = kw['p_match']
        seglen = kw['seglen']
        volume = kw['volume']
        if volume == 0:
            return float('-inf'), float('-inf')
        if p_match == 1.:
            # we can't let p_H1 == 1 because we get division by zero below
            p_match = 1 - np.finfo(float).eps

        p_H0 = (1. / len(self.alphabet)) ** (len(self.seqs) - 1)
        pw_H0 = p_H0 ** self.wordlen

        mu_H0 = volume * pw_H0
        sd_H0 = np.sqrt(volume * (
            (1 - pw_H0) * (pw_H0 + 2 * p_H0 * pw_H0 / (1 - p_H0)) -
            2 * self.wordlen * pw_H0 ** 2
        ))

        p_H1 = p_match
        pw_H1 = p_H1 ** self.wordlen

        mu_H1 = mu_H0 + seglen * pw_H1
        sd_H1 = np.sqrt(sd_H0 ** 2 + seglen * (
            (1 - pw_H1) * (pw_H1 + 2 * p_H1 * pw_H1 / (1 - p_H1)) -
            2 * self.wordlen * pw_H1 ** 2
        ))

        z_H0 = (num_seeds - mu_H0) / sd_H0  # score under H0
        z_H1 = (num_seeds - mu_H1) / sd_H1  # score under H1
        return z_H0, z_H1

    def similar_segments(self, K_min, p_min, at_least_one=False):
        """Find all maximal local similarities of given minium length and match
        probability. Additionally for each segment, the match probability is
        estimated and H0/H1 scores are calculated.

        Args:
            K_min (int):
                minimum required length of similarity.
            p_min (float):
                Minimum required match probability at each position.
            score (bool):
                Whether to score the segment against H1 by counting seeds
                inside.

        Yields:
            dict: dictionary with keys: ``segment`` (coordinates of similar
            region in diagonal coordinates, ``p`` the estimated match
            probability, and ``score`` the H1 z-score if keyword argument
            ``score`` is true.
        """
        d_radius = int(np.ceil(self.band_radius(K_min)))
        a_radius = int(np.ceil(len(self.seqs) * K_min / 2.))
        scored_seeds = self.score_seeds(K_min)
        self.log('finding local similarities between %d sequences' %
                 len(self.seqs))

        def _update_seg(seg, seed):
            ds, a = seed
            if seg is None:
                d_ranges = [None] * (len(self.seqs) - 1)
                for i in range(len(self.seqs) - 1):
                    d_ranges[i] = ds[i] - d_radius, ds[i] + d_radius
                a_range = a - a_radius, a + a_radius
            else:
                d_ranges, a_range = seg
                assert all(len(r) == 2 for r in d_ranges)  # pairs of min, max
                for i in range(len(self.seqs) - 1):
                    d_min, d_max = d_ranges[i]
                    d_ranges[i] = (min(ds[i], d_min),
                                   max(ds[i], d_max))
                a_min, a_max = seg[-1]
                a_range = (min(a - a_radius, a_min), max(a + a_radius, a_max))
            return d_ranges, a_range

        avail = [rec['p'] >= p_min for rec in scored_seeds]
        if not any(avail) and at_least_one:
            # we're obliged to return something, let the highest probability
            # seed go through.
            assert len(scored_seeds), 'no seeds found while at_least_one=True'
            avail[np.argmax([rec['p'] for rec in scored_seeds])] = True
        while True:
            try:
                # TODO grab the highest probability so we yield in decreasing
                # order of quality
                seed_idx = avail.index(True)
            except ValueError:
                break
            stack = [seed_idx]
            avail[seed_idx] = False
            ps_in_seg = [scored_seeds[seed_idx]['p']]
            seg = None
            while stack:
                idx = stack.pop()
                ps_in_seg.append(scored_seeds[idx]['p'])
                seg = _update_seg(seg, scored_seeds[idx]['seed'])
                for neigh in scored_seeds[idx]['neighs']:
                    if avail[neigh]:
                        stack.append(neigh)
                        avail[neigh] = False
            if seg is None:
                break
            else:
                # clip overflowing values so translating back to standard
                # coordinates gives meaningful numbers.
                ds_range, a_range = seg
                for idx in range(len(self.seqs) - 1):
                    d_min, d_max = ds_range[idx]
                    d_min = min(
                        len(self.seqs[0]),
                        max(d_min, -len(self.seqs[idx + 1]))
                    )
                    d_max = min(
                        len(self.seqs[0]),
                        max(d_max, -len(self.seqs[idx + 1]))
                    )
                    ds_range[idx] = d_min, d_max
                a_min, a_max = a_range
                a_min = max(a_min, 0)
                a_max = min(a_max, sum(len(seq) for seq in self.seqs))
                seg = ds_range, (a_min, a_max)
            # NOTE the following is more justifiable but it matches the
            # average. TODO turn this in into an experiment to justify
            # p_hat = self.estimate_match_probability(
            #   n, d_band=seg[0], a_band=seg[1])
            p_hat = sum(ps_in_seg) / len(ps_in_seg)
            res = {'segment': seg, 'p': p_hat}
            ds_band, a_band = seg
            n = self.seed_count(ds_band=ds_band, a_band=a_band)
            K_hat = np.ceil((a_band[1] - a_band[0]) / len(self.seqs))
            volume = a_band[1] - a_band[0]
            for d_band in ds_band:
                volume *= d_band[1] - d_band[0]
            scores = self.score_num_seeds(num_seeds=n, volume=volume,
                                          seglen=K_hat, p_match=p_hat)
            res['scores'] = scores
            yield res

class WordBlotMultipleFast(WordBlotMultiple):
    """An in-memory, SQL-free version of :class:`WordBlotOverlapMultiple` for
    faster comparisons. Due to implementation details the word length is
    constrained above by the available memory.

    Attributes:
        allowed_memory (int|float): allocatable memory in GB for kmers index,
            default is 1.
    """
    def __init__(self, *seqs, **kw):
        name = '_'.join(S.content_id[:8] for S in seqs)
        self.name = name
        self.wordlen = kw['wordlen']
        self.alphabet = kw['alphabet']
        self.g_max = kw['g_max']
        self.sensitivity = kw['sensitivity']
        self.log_level = kw.get('log_level', logging.INFO)
        self.seqs = seqs
        self.allowed_memory = kw.get('allowed_memory', 1)
        assert self.allowed_memory > 0, 'allowed memory must be positive'
        num_kmers = len(self.alphabet) ** self.wordlen
        mem_needed = sys.getsizeof(num_kmers) * num_kmers
        mem_needed_gb = np.power(2, np.log2(mem_needed) - 30)
        if mem_needed_gb > self.allowed_memory:
            msg = 'not enough memory (max = %.2f GB) ' % self.allowed_memory
            msg += 'to store %d-mers ' % self.wordlen
            msg += '(%.2f GB needed)' % mem_needed_gb
            raise MemoryError(msg)
        self.kmer_hits = [[] for _ in range(num_kmers)]
        for idx, seq in enumerate(self.seqs):
            for pos, kmer in enumerate(as_kmer_seq(seq, self.wordlen)):
                self.kmer_hits[kmer].append((idx, pos))
        log_header = '%d-mer word-blot (python-object)' % self.wordlen
        self._logger = Logger(log_level=self.log_level, header=log_header)

    def seeds(self):
        for hits in self.kmer_hits:
            hits = {seqid: [c[1] for c in seq_hits]
                    for seqid, seq_hits in groupby(hits,
                                                   key=lambda c: c[0])}
            # only consider kmers present in all sequences
            if len(hits) < len(self.seqs):
                continue
            for idxs in product(*hits.values()):
                ds, a = self.to_diagonal_coordinates(*idxs)
                yield list(ds), a

    def seed_count(self, ds_band=None, a_band=None):
        cnt = 0
        for ds, a in self.seeds():
            if ds_band:
                if not all(ds_band[i][0] <= d <= ds_band[i][1]
                           for i, d in enumerate(ds)):
                    continue
            if a_band and not a_band[0] <= a <= a_band[1]:
                continue
            cnt += 1
        return cnt