"abstract classes for deep models"

import inspect, logging, time, cPickle, abc
from collections import Sequence
from operator import itemgetter, attrgetter
from functools import partial

import numpy as np
import pandas as pd
import theano

from . import config_spec
from .. import archive, deprecated

__spec_cls__ = lambda t: (t[0].endswith("Spec") and inspect.isclass(t[1]))
spec_list = map(itemgetter(1), 
                filter(__spec_cls__, inspect.getmembers(config_spec)))


log = logging.getLogger(__name__)

class Layer( object ):
    """abstract layer class.

    A layer is a list of weights. Weights are
    theano shared tensors that are needed for example to compute activation.
    You can use them also for other reasons e.g., speeds to implement
    momentum GSD.

    """

    __metaclass__ = abc.ABCMeta

    def __init__(self, wshape_lst, wnames, fanin_lst, rng, wdtype):
        """initialize by the given weight properties, all params are lists
        of the same length

        @param wshape_lst: shapes of weights
        @param wnames: names of weights
        @param fanin_lst: fan in of weights
        @param rng: a random number generator or init value (see alloc_shared_weights_)
        @param wdtype: data type of weights.
        """

        if len(wshape_lst) != len(wnames):
            raise ValueError("length names and shapes differ")

        self._weights_ = []
        for (sh, n, fi, r, dt) in zip(wshape_lst, wnames, fanin_lst, rng, wdtype):
            self._weights_.append( alloc_shared_weights_(sh, dt, n, fi, r) )

    @abc.abstractmethod
    def get_params(self):
        """weights as symbolic variables"""
        pass

    @abc.abstractmethod
    def activation(self):
        "the symbolic activation function"
        pass

    def get_weights(self):
        """weights as ndarrays"""

        return map(lambda w: w.get_value(), self.get_params())

    @abc.abstractmethod
    def set_params(self, params):
        """set weights

        @param params: tuple of (ndarray)
        @return none"""

        pass

    def set_weights(self, vlst):
        """set the weight values of this layer"""

        assert len(vlst) == len(self.get_params())

        for v, sym_w in zip(vlst, self.get_params()):
            sym_w.set_value( v )


    def weight_norm(self, degree):
        """compute the L1 or L2 norm of weights

        @param degree: keyword `l1` or `l2`
        @return  : the L1 or L2 norm of weights (float)"""

        f = {"l1" : lambda w: np.abs( w ).sum(),
                "l2" : lambda w: (np.abs( w )**2).sum()}[degree]

        return sum( map(f, self.get_weights()) )

    def get_flat_weights(self):
        """get a list of the flattened version of the weights of this layer

        @return: list[ndarray]"""

        olst = []
        for w in self.get_weights():
            olst.append( w.reshape(-1,) ) 
        return olst
    
    def load_flat_weights(self, flatp):
        """reshape the given flattened weights into the correct dimensions
        for this layer
        
        @param flatp: flattened weights (list[ndarray])
        @return: the same given weights reshaped""" 

        for i in range(len(flatp)):
            flatp[i] = flatp[i].reshape( self.get_weights()[i].shape )
        return flatp
    
class SpeedLayer( Layer ):
    """This layer provides an extra set of weights as a support the momentum algorithm for SGD.
    At time point t, we need weights at t-1 and the gradient at t to update weights. Namely

    w(t+1) - w(t) = - rho * dE(w)/dw + p s(t), for t = 0, 1, ... and s(0) = 0

    I call w(t) - w(t-1) = s(t) (speed)
    """


    __metaclass__ = abc.ABCMeta

    def __init__(self, ws, wn, rng, wdtype):
        super(SpeedLayer, self).__init__(ws+ws,
                wn+map(lambda n: "%s_speed" % n, wn),
                map(sum, ws) + map(sum, ws),
                rng + [0.0]*len(rng),
                wdtype + wdtype)
        self.__wl = len(ws)

    def get_speeds(self):
        """get speeds
        
        @return: a tuple of (ndarray) speeds """
        
        return map(lambda i: self._weights_[self.__wl+i].get_value(),
                range(self.__wl))

    def set_speeds(self, v):
        """set speeds
        
        @param v: tuple of (ndarray)
        @return: none"""

        for i in range(self.__wl):
            self._weights_[self.__wl + i].set_value( v[i] )

    def get_params(self):
        """get weights
        
        @return: a tuple of (theano.tensor) weights"""

        return self._weights_[0:self.__wl]
    
    def set_params(self, params):
        """set weights
        
        @param params: tuple of (ndarray)
        @return none"""

        for i in range(len(self.__wl)):
            self._weights_[i] = params[i]
            
    @classmethod
    def _speed_update_f(cls, s, g, mu, rho):
        """speed update formula
        
        new_speed = -rho * g + mu*s

        @param s: cur speed
        @param g: gradient
        @param mu: momentum
        @param rho: learning rate
        @return: new speed"""
        
        return (s*mu) - (rho * g)

    @classmethod
    def _weight_update_f(cls, cur_weight, speed):
        """speed update for the given gradient
        
        new_weight = cur_weight + speed
        
        @param cur_weight: current weight (ndarray)
        @param speed: speed (ndarray)
        @return: new_weight (ndarray)
        """
        
        return cur_weight + speed 

    def speed_update(self, gradient, mu, rho):
        """update speeds for the given gradients
        
        new_speed = -rho * gradient + mu*cur_speed
        
        @param gradient: gradient (list of ndarray)
        @param mu: momentum (float)
        @param rho: learning rate (float)
        @return: new_speed (ndarray)
        """
        
        upd_f = self._speed_update_f
        new_spd = []
        for i in range(self.__wl):
            new_spd.append( upd_f(self.get_speeds()[i], gradient[i], mu, rho) )
        self.set_speeds(new_spd)

    def weight_update(self):
        """update speeds for the current speed
        
        new_weight = cur_weight + cur_speed
        
        """
        upd_f = self._weight_update_f
        new_w = []
        for i in range(self.__wl):
            new_w.append( upd_f(self.get_weights()[i], self.get_speeds()[i]) )
        self.set_weights(new_w)


class Model( object ):
    """generic model class with basic functionality"""

    __metaclass__ = abc.ABCMeta

    def __init__(self, layers):
        self.__layers = layers
        log.info("declared model\n%s", str(self))

    def __getitem__(self, i):
        return self.__layers[i]

    def __len__(self):
        return len(self.__layers)

    def __setitem__(self, i, it):
        self.__layers[i] = it

    def __delitem__(self, i):
        del self.__layers[i]

    def load(self, path):
        archpath = archive.archname(path)

        def load_by_name(layer_key, wn, what):
            kp = {"weights" : "%s/%s", "shape" : "%s/%s_shape"}[what]
            return archive.load_object(archpath, kp % (layer_key, wn) ).values

        for i, layer in enumerate(self):
            layer_key = archive.basename( path ) + ("/layer_%d" % i)
            layer_weights = []
            for wn in map(attrgetter("name"), layer.get_params()):
                layer_weights.append( load_by_name(layer_key, wn, "weights") )
            layer.set_weights( layer.load_flat_weights(layer_weights) )

    def save(self, path):
        def weights_names_shapes(layer):
            "[(wname, w_flat, w_shape), ...]"
            layer_weights = layer.get_flat_weights()
            layer_weight_names = map(attrgetter("name"), layer.get_params())
            layer_shapes = map(lambda w: w.shape, layer.get_weights())
            return zip(layer_weight_names, layer_weights, layer_shapes)

        log.info("saving model to %s", path)
        for i, layer in enumerate(self):
            layer_key = archive.basename( path ) + ("/layer_%d" % i)
            for wn, w, wsh in weights_names_shapes(layer):
                key = layer_key + "/" + wn
                archive.save_object( archive.archname(path), key, pd.Series(w) )
                key = layer_key + "/" + wn + "_shape"
                archive.save_object( archive.archname(path), key,
                                     pd.Series(wsh, index=range(len(wsh))) )

    def weight_norm(self, degree):
        """compute the L1 or L2 norm of weights

        @param degree: keyword `l1` or `l2`
        @return: the L1 or L2 norm of weights (float)"""

        return sum( map(lambda l: l.weight_norm(degree), self) )


    def get_weights(self):
        """numpy arrays for the weights

        @return: the list of parameters of layers
        """

        return map(lambda l: l.get_weights(), self.__layers)

    def set_weights(self, wlst):
        for w, i in enumerate(wlst):
            self.__layers[i].set_weights( w )

    def get_params(self):
        """(theano) symbolic variables for the weights

        @return: the list of parameters of layers
        """

        return map(lambda l: l.get_params(), self.__layers)

    def __str__(self):
        return "\n".join( map(str, self) )



def verbose_compile(func):
    """decorator that explains what is being theano-compiled, by logging the functions __doc__"""
    def newfunc(*args, **kwargs):
        log.info( func.__doc__ )
        return func(*args, **kwargs)
    return newfunc


def alloc_shared_weights_(shape, dtype, name, fan_in, rng):
    """alloc a matrix of weights and return a theano shared variable
    Note : optimal initialization of weights is dependent on the
    activation function used (among other things). Here  I replicate
    http://www.deeplearning.net/tutorial/mlp.html#mlp
    """

    if type(rng) == float:
        val = str(rng)
        var = theano.shared(value=np.zeros(shape, dtype=dtype) + rng, name=name)
    elif type(rng) == np.random.RandomState:
        val = "rnd"
        thr = np.sqrt(3./fan_in)
        var = theano.shared(value=np.asarray(rng.uniform(low=-thr, high=thr,
            size=shape), dtype=dtype), name=name)
    else:
        raise ValueError("cannot understant `rng`")

    log.debug("%s (%s) = %s (%s)", str(var), str(var.get_value().shape),
        val, str(var.get_value().dtype))

    return var

def adjust_lr(err, lrmax):
    "adjust error rate for the next batch"

    ## another strategy is to keep the learning rate fixed
    ## if it does not increase the error
    ## and to sqrt(half) it otherwise

    if len(err) < 2:
        return lrmax
    E = max( err[-2], err[-1] )
    if E == 0:
        return lrmax
    if (err[-2] - err[-1]) / E < -0.001: ## went up by more than 1%
        return float( np.array([lrmax * np.sqrt((len(err)-1) / float(len(err)))], dtype=np.float32)[0] )
    return lrmax

@deprecated
def get_data(low_idx, high_idx, path, borrow=True, as_dtypeX=None, as_dtypeY=None):
    """load rasterized data into shared theano variables"""

    x, y = archive.get_target_dataset(path, low_idx, high_idx)

    SH = partial(theano.shared, borrow=borrow)

    if as_dtypeX:
        x = np.asarray(x, dtype=as_dtypeX)
    if as_dtypeY:
        y = np.asarray(y, dtype=as_dtypeY)

    return SH(x.reshape(x.shape[0], np.prod(x.shape[1:]))), SH(y)


def shift_data(x):
    """transform each component of X so that it has zero mean

    x: a pandas data panel of the form <anchors> X <tracks> X <genome position>
    return: (the shifted input, the mean for each input component. this is an array of shape(<tracks>, <genome position>))
    """

    m = np.empty( x.shape[1:] )
    for i in xrange(m.shape[0]):
        for j in xrange(m.shape[1]):
            m[i, j] = x.values[:, i, j].mean()
    for a in x:
        x[a] -= m
    return x, m

def scale_data(x):
    """transform each component of X so that it has unit variance

    @param x: a pandas data panel of the form <anchors> X <tracks> X <genome position>
    @return: (the scaled input, the std for each input component. this is an array of shape(<tracks>, <genome position>))
    """

    m = np.empty( x.shape[1:] )
    for i in xrange(m.shape[0]):
        for j in xrange(m.shape[1]):
            m[i, j] = x.values[:, i, j].std()
    if np.any( m == 0 ):
        log.warning("zero variance components")
        m[ m == 0 ] = 1
    for a in x:
        x[a] /= m
    return x, m

def fit_data(x):
    """transform each component of X so that it fits on an interval [-1, 1]

    @param x: a pandas data panel of the form <anchors> X <tracks> X <genome position>
    @return: (the scaled input, the std for each input component. this is an array of shape(<tracks>, <genome position>))
    """

    md = np.empty( x.shape[1:] )
    rg = np.empty( x.shape[1:] )
    for i in xrange(md.shape[0]):
        for j in xrange(md.shape[1]):
            md[i, j] = (np.max(x.values[:, i, j]) + np.min(x.values[:, i, j])) / 2
            rg = (np.max(x.values[:, i, j]) - np.min(x.values[:, i, j]))
    if np.any( rg == 0 ):
        log.warning("zero variability components")
        rg[ rg == 0 ] = 1

    ## x has mid=0 and rg = 1
    for a in x:
        x[a] = (x[a] - md) / rg
    return x




