# -*- Mode: python; tab-width: 4; indent-tabs-mode:nil; coding:utf-8 -*-
# vim: tabstop=4 expandtab shiftwidth=4 softtabstop=4
#
# RecSQL -- a simple mash-up of sqlite and numpy.recsql
# Copyright (C) 2007-2016 Oliver Beckstein <orbeckst@gmail.com>
# Released under the GNU Public License, version 3 or higher (your choice)

"""
:mod:`recsql.sqlfunctions` --- Functions that enhance a SQLite db
=================================================================

This module contains new SQL functions to be added to a SQLite database that
can be used in the same way as the builtin functions.

Example:

  Add the functions to an existing connection in the following way (assuming
  that the db connection is available in ``self.connection``)::

     from sqlfunctions import *
     self.connection.create_function("sqrt", 1, _sqrt)
     self.connection.create_function("sqr", 1, _sqr)
     self.connection.create_function("periodic", 1, _periodic)
     self.connection.create_function("pow", 2, _pow)
     self.connection.create_function("match", 2, _match)   # implements MATCH
     self.connection.create_function("regexp", 2, _regexp) # implements REGEXP
     self.connection.create_function("fformat",2,_fformat)
     self.connection.create_aggregate("std",1,_Stdev)
     self.connection.create_aggregate("stdN",1,_StdevN)
     self.connection.create_aggregate("median",1,_Median)
     self.connection.create_aggregate("array",1,_NumpyArray)
     self.connection.create_aggregate("histogram",4,_NumpyHistogram)
     self.connection.create_aggregate("distribution",4,_NormedNumpyHistogram)
     self.connection.create_aggregate("meanhistogram",5,_MeanHistogram)
     self.connection.create_aggregate("stdhistogram",5,_StdHistogram)
     self.connection.create_aggregate("minhistogram",5,_MinHistogram)
     self.connection.create_aggregate("maxhistogram",5,_MaxHistogram)
     self.connection.create_aggregate("medianhistogram",5,_MedianHistogram)
     self.connection.create_aggregate("zscorehistogram",5,_ZscoreHistogram)

Module content
--------------
.. See the autogenerated content in the online docs or the source code.

"""
from itertools import izip
import re
import numpy


from sqlutil import adapt_numpyarray, convert_numpyarray,\
    adapt_object, convert_object


def _sqrt(x):
    try:
        x = float(x)
    except TypeError:
        return None
    return numpy.sqrt(x)

def _sqr(x):
    return x*x

def _periodic(x):
    """Wrap angle in degree between -180 and +180"""
    x = x % 360
    if x <= -180:
        return 360 + x
    elif x > 180:
        return -360 + x
    return x

def _pow(x,y):
    try:
        x = float(x)
        y = float(y)
    except TypeError:
        return None
    return numpy.power(x,y)

def _regexp(pattern, string):
    """string REGEXP pattern == re.search(pattern, string)"""
    return re.search(unicode(pattern), unicode(string)) != None

def _match(pattern, string):
    """string MATCH pattern == re.match(pattern, string)"""
    return re.match(unicode(pattern), unicode(string)) != None

def _fformat(format,x):
    return unicode(format) % x

class _Stdev(object):
    """Implement standard deviation of the sample as SQL aggregate function.
    (Uses N-1 variance.)
    Do it in one pass (see eg
    http://smallcode.weblogs.us/2006/11/27/calculate-standard-deviation-in-one-pass/
    though we may run in an underflow by calculating N/N-1<X^2>-<X>^2>.).

    Also, we don't check if our arguments are valid as numbers.
    """
    def __init__(self):
        self.x2 = 0
        self.x = 0
        self.n = 0
    def step(self,x):
        try:
            x = float(x)
            self.x2 += x*x
            self.x  += x
            self.n  += 1
        except TypeError:
            pass        # don't contribute to average
    def finalize(self):
        if self.n<2: return 0.0
        return numpy.sqrt((self.n*self.x2 - self.x*self.x)/(self.n*(self.n-1)))

class _StdevN(_Stdev):
    """Implement standard deviation as SQL aggregate function.
    (Uses N variance.)
    Do it in one pass (see eg
    http://smallcode.weblogs.us/2006/11/27/calculate-standard-deviation-in-one-pass/
    though we may run in an underflow by calculating <X^2>-<X>^2>.).

    Also, we don't check if our arguments are valid as numbers.
    """
    def finalize(self):
        X2 = self.x2/self.n
        X = self.x/self.n
        return numpy.sqrt(X2 - X**2)

class _Median(object):
    def __init__(self):
        self.data = []
    def step(self,x):
        try:
            x = float(x)
            self.data.append(x)
        except TypeError:
            pass        # don't contribute
    def finalize(self):
        return numpy.median(self.data)

class _NumpyArray(object):
    def __init__(self):
        self.data = []
    def step(self,x):
        self.data.append(x)
    def finalize(self):
        return adapt_numpyarray(numpy.array(self.data))

class _NumpyHistogram(object):
    def __init__(self):
        self.is_initialized = False
        self.data = []
    def step(self,x,bins,xmin,xmax):
        if not self.is_initialized:
            self.bins = bins
            self.range = (xmin,xmax)
            self.is_initialized = True
        self.data.append(x)
    def finalize(self):
        hist,edges = numpy.histogram1d(self.data,bins=self.bins,range=self.range,
                                       normed=False)
        return adapt_object((hist,edges))

class _NormedNumpyHistogram(_NumpyHistogram):
    def finalize(self):
        hist,edges = numpy.histogram1d(self.data,bins=self.bins,range=self.range,
                                       normed=True)
        return adapt_object((hist,edges))

class _FunctionHistogram(_NumpyHistogram):
    """Baseclass for histogrammed functions.

    A histogrammed function is created by applying a function
    to all values y that have been accumulated in a bin x.
    """
    def __init__(self):
        _NumpyHistogram.__init__(self)
        self.y = []
    def step(self,x,y,bins,xmin,xmax):
        _NumpyHistogram.step(self,x,bins,xmin,xmax)
        self.y.append(y)
    def finalize(self):
        raise NotImplementedError("_FunctionHistogram must be inherited from.")
        # return adapt_object( (...,...,...) )

class _MeanHistogram(_FunctionHistogram):
    """Mean of the weights in each bin.
    Takes TWO column arguments: value and weight"""
    def finalize(self):
        return adapt_object(regularized_function(\
                self.data,self.y,numpy.mean,bins=self.bins,range=self.range))

class _StdHistogram(_FunctionHistogram):
    """Standard deviation of the weights in each bin.
    Takes TWO column arguments: value and weight"""
    def finalize(self):
        return adapt_object(regularized_function(\
                self.data,self.y,numpy.std,bins=self.bins,range=self.range))

class _MinHistogram(_FunctionHistogram):
    """Min value of the weights in each bin.
    Takes TWO column arguments: value and weight"""
    def _min(self,v):
        try:
            return numpy.min(v)
        except ValueError:  # empty array
            return numpy.nan

    def finalize(self):
        return adapt_object(regularized_function(\
                self.data,self.y,self._min,bins=self.bins,range=self.range))

class _MaxHistogram(_FunctionHistogram):
    """Max value of the weights in each bin.
    Takes TWO column arguments: value and weight"""
    def _max(self,v):
        try:
            return numpy.max(v)
        except ValueError:  # empty array
            return numpy.nan

    def finalize(self):
        return adapt_object(regularized_function(\
                self.data,self.y,self._max,bins=self.bins,range=self.range))

class _MedianHistogram(_FunctionHistogram):
    """Median value of the weights in each bin.
    Takes TWO column arguments: value and weight"""
    def finalize(self):
        return adapt_object(regularized_function(\
                self.data,self.y,numpy.median,bins=self.bins,range=self.range))

class _ZscoreHistogram(_FunctionHistogram):
    """Z-score of the weights in each bin abs(Y - <Y>)/std(Y).
    Takes TWO column arguments: value and weight"""
    def Zscore(self,v):
        m = v.mean()
        s = v.std()
        return numpy.nan_to_num( numpy.mean(numpy.abs(v - m))/s )

    def finalize(self):
        return adapt_object(\
            regularized_function(self.data,self.y,self.Zscore,bins=self.bins,range=self.range))


# Helper functions

def regularized_function(x,y,func,bins=None,range=None):
    """Compute func() over data aggregated in bins.

    (x,y) --> (x', func(Y'))  with Y' = {y: y(x) where x in x' bin}

    First the data is collected in bins x' along x and then func is applied to
    all data points Y' that have been collected in the bin.

    :Arguments:
       x
          abscissa values (for binning)
       y
          ordinate values (func is applied)
       func
          a numpy ufunc that takes one argument, func(Y')
       bins
          number or array
       range
          limits (used with number of bins)

    :Returns:
       F,edges
          function and edges (midpoints = 0.5*(edges[:-1]+edges[1:]))
    """
    _x = numpy.asarray(x)
    _y = numpy.asarray(y)

    # setup of bins from numpy.histogram
    if (range is not None):
        mn, mx = range
        if (mn > mx):
            raise ValueError('max must be larger than min in range parameter.')

    if not numpy.iterable(bins):
        if range is None:
            range = (_x.min(), _x.max())
        mn, mx = [float(mi) for mi in range]
        if mn == mx:
            mn -= 0.5
            mx += 0.5
        bins = numpy.linspace(mn, mx, bins+1, endpoint=True)
    else:
        bins = numpy.asarray(bins)
        if (numpy.diff(bins) < 0).any():
            raise ValueError('bins must increase monotonically.')

    sorting_index = numpy.argsort(_x)
    sx = _x[sorting_index]
    sy = _y[sorting_index]

    # boundaries in SORTED data that demarcate bins; position in bin_index is the bin number
    bin_index = numpy.r_[sx.searchsorted(bins[:-1], 'left'),
                         sx.searchsorted(bins[-1], 'right')]

    # naive implementation: apply operator to each chunk = sy[start:stop] separately
    #
    # It's not clear to me how one could effectively block this procedure (cf
    # block = 65536 in numpy.histogram) because there does not seem to be a
    # general way to combine the chunks for different blocks, just think of
    # func=median
    F = numpy.zeros(len(bins)-1)  # final function
    F[:] = [func(sy[start:stop]) for start,stop in izip(bin_index[:-1],bin_index[1:])]
    return F,bins
