'''
support for dataset IO and manipulation

@author: odenas

'''
import logging

import numpy as np
import pandas as pd
import theano

import archive, ops

log = logging.getLogger(__name__)


class aaHDFArchive(object):
    def __init__(self, path):
        self.archive, self.datapath = archive.split(path)
        if not hasattr(self, "pX") or not hasattr(self, "X"):
            raise ValueError("this mixin requires ")

    def dump(self, path):
        arch = self.archive
        key = self.datapath
        X, Y, T = self.pX, self.sY, self.dfT
        (nsamp, ntrack, width) = X.values.shape

        if not (X is None):
            archive.save_object( arch, "%s/rawX" % key, X )
            normX, meanX, sdX = self.normalize_features(X.values.reshape(nsamp, -1))

            archive.save_object( arch, "%s/X" % key,
                                 pd.Panel(normX.reshape(nsamp, ntrack, width),
                                          items = X.items,
                                          major_axis = X.major_axis,
                                          minor_axis = X.minor_axis) )
            archive.save_object( arch, "%s/meanX" % key,
                                 pd.DataFrame(meanX, index = X.major_axis,
                                              columns=X.minor_axis ) )
            archive.save_object( arch, "%s/sdX" % key,
                                 pd.DataFrame(sdX, index = X.major_axis,
                                              columns=X.minor_axis ) )

        if not (Y is None):
            archive.save_object( arch, "%s/Y" % key, Y )

        if not (T is None):
            archive.save_object( arch, "%s/T" % key, T )

    @classmethod
    def _from_archive(cls, path, raw, *args):
        ap, did = archive.split(path)

        key = "%s/%s" % (did, (raw and "rawX" or "X"))
        X = archive.load_object(ap, key)

        def load_none(k, p=ap, did=did):
            try:
                return archive.load_object(p, "%s/%s" % (did, k))
            except Exception:
                log.info("%s not found for this dataset" % k)

        return cls(X, load_none("Y"), load_none("T"), *args)

class aAnchorDataset(object):
    """this dataset contains various tracks of epigenetic signal in the
    for a set of genome sites (e.g., TSS-cenetered regions) all of the
    same width. an instance maintains references to X (as a panel and ndarray),
    Y, and T
    """

    def __init__(self, X, Y, T, batch_size, valid_s=None, valid_idx=None, rng=None):
        self.pX = X
        self.X = X.values

        self.sY = Y
        if not (Y is None):
            self.Y = Y.values

        self.dfT = T
        if not (T is None):
            self.T = T["label_code"].values

        ## theano shared vars
        self.__shX = None
        self.__shY = None
        self.__shT = None

        self.batch_size = (batch_size or self.X.shape[0])

        nb = self.X.shape[0] / self.batch_size ## nr. of batches
        self.train_batches, self.valid_batches = self.__batches(self.X.shape[0],
                                                            self.batch_size,
                                                            (valid_s or 0.25),
                                                            (valid_idx or nb - nb/4),
                                                             rng)
        self.label_names = None
        if not (self.dfT is None):
            self.label_names = np.unique( self.dfT["label_name"].values ).tolist()

        self.track_names = self.pX.major_axis.values.tolist()

    def __sh_anon(self, what, shape=None, borrow=True):
        if getattr(self, what) is None:
            raise ValueError("cannot share non-existent member %s" % what)

        if getattr(self, "_AnchorDataset__sh%s" % what) is None:
            init_val = getattr(self, what)
            if shape:
                init_val = init_val.reshape( shape )

            setattr(self, "_AnchorDataset__sh%s" % what,
                    theano.shared(init_val, borrow=borrow))
        return getattr(self, "_AnchorDataset__sh%s" % what)

    @property
    def shX(self):
        return self.__sh_anon("X")

    @property
    def shY(self):
        return self.__sh_anon("Y")

    @property
    def shT(self):
        return self.__sh_anon("T")

    def share(self, which, shape=None, borrow=True):
        """wrap the data on a thean.shared variable

        @param which: what component to wrap (str, typically 'X', 'T', 'Y')
        @param shape: reshape the array to this shape
        @param borrow: passer to theano.share
        @return: theano.shared instance initialized to the required data"""

        val = getattr(self, which)
        if not (shape is None):
            val = val.reshape( shape )
        return theano.shared(val, borrow=borrow)



    def __iter_batches(self, which, nepochs):
        """infinite loop over train/valid batches

        @param nepochs: loop this many times over train batches (0 will loop forever)
        @return: iterator """

        assert which in ("train_batches", "valid_batches")

        batches = getattr(self, which)
        epoch = 0
        while True:
            for i in batches:
                yield i
            epoch = epoch + 1
            if epoch == nepochs:
                break

    def iter_train(self, nepochs):
        return self.__iter_batches("train_batches", nepochs)

    def iter_valid(self, nepochs):
        return self.__iter_batches("valid_batches", nepochs)

    @staticmethod
    def __batches(tot_size, batch_s, valid_s, valid_idx, rng):
        """create train and validation batches from the given params.

        the idea is to split the data into batches and allocate a 'valid_s'
        portion of them for validation. the position of the (continuous) validation
        block is w.r.t batches. E.g., for tot_size = 10, batch_s = 2,
        valid_idx=3, valid_s = 0.3 you get 4 + 1 train + valid batches: T T T V T

        @param tot_size: nr. of examples
        @param batch_s : batch size
        @param valid_s : fraction of data to allocate for validation
        @param valid_idx: batch index at which allocate validation data
        @param rng : numpy.RandomState used to shuffle batches or None (no shuffle)
        @return : (train_batches, valid_batches)"""

        if valid_s <= 0 or valid_s >= 1:
            raise ValueError("valid_s (%f) should be between (0, 1) ", valid_s)

        if batch_s > tot_size * min(valid_s, 1-valid_s):
            raise ValueError("batch size (%d) too big > min(valid_s=%d, train_s=%d)"% (batch_s,
                             tot_size * valid_s, tot_size *(1-valid_s)))

        all_batches = range( tot_size / batch_s )
        try:
            valid_batches = all_batches[valid_idx:valid_idx+int(len(all_batches)*valid_s)]
        except IndexError:
            raise ValueError("valid_idx (%d) should be between 0 and %d",
                             valid_idx, len(all_batches)-1)
        train_batches = list( set(all_batches) - set(valid_batches) )
        assert set(train_batches + valid_batches) == set(all_batches)
        assert len( set(train_batches) & set(valid_batches) ) == 0
        if not (rng is None):
            rng.shuffle(train_batches)
            rng.shuffle(valid_batches)
        logging.info("train batches: %s", str(train_batches))
        logging.info("valid batches: %s", str(valid_batches))
        return (train_batches, valid_batches)

    @property
    def labels(self):
        return len(self.label_names)

    @property
    def tracks(self):
        return len(self.track_names)

    @property
    def width(self):
        return self.X.shape[2]

    @staticmethod
    def normalize_features(x):
        """transform each component of flattened X examples to 0 mean and 1 std
        So the values of track t at position i are 0 mean and 1 std

        x: a pandas data panel of the form <anchors> X <tracks> X <genome position>
        return: (the shifted input,
                the mean for each input component, the sd of each input component)
                the latter 2 are arrays of shape(<tracks>, <genome position>)
        """

        normX, m,v = ops.standardize(x.values, axis=0)
        pX = pd.Panel( normX, items = x.items, major_axis = x.major_axis,
                          minor_axis = x.minor_axis)
        return pX, m.reshape(x.shape[1], -1), v.reshape(x.shape[1], -1)

    @staticmethod
    def fit_features(x):
        """transform each **component** of X so that it fits on an interval [-1, 1].
        So the values of track t at position i are all in [-1,1]

        @param x: a pandas data panel of the form <anchors> X <tracks> X <genome position>
        @return: the scaled input
        """

        fitX = ops.fit(x.values.reshape(x.shape[0], -1), axis=0)

        return pd.Panel( fitX.reshape(x.shape),
                         items = x.items, major_axis = x.major_axis,
                          minor_axis = x.minor_axis)




#### TODO ####
class Dataset(object):
    def __init__(self, X, Y, T):
        self.X, self.Y, self.T = X, Y, T

        if (not (Y is None)) and X.shape[0] != Y.shape[0]:
            raise ValueError("|X| (%d) != |Y| (%d)" % (X.shape[0],
                                                       Y.shape[0]))
        if (not (T is None)) and X.shape[0] != T.shape[0]:
            raise ValueError("|X| (%d) != |T| (%d)" % (X.shape[0],
                                                       T.shape[0]))
        log.info("allocated dataset. X of shape %s, Y %s, T %s",
                 str(self.X.shape),
                 (self.Y is None and "missing" or str(self.Y.shape)),
                 (self.T is None and "missing" or str(self.T.shape)))

    @property
    def is_labeled(self):
        return not (self.T is None)

    @property
    def labels(self):
        if self.is_labeled:
            return np.unique( self.T ).shape[0]
        else:
            raise AttributeError("unlabeled dataset")

    @staticmethod
    def normalize_features(x):
        """transform each component of flattened X examples to 0 mean and 1 std
        So the values of feature f (from all examples) are 0 mean and 1 std

        x: a ndarray of shape (nr. examples, nr. of features)
        return: (the shifted input,
                the mean for each input component, the sd of each input component)
                the latter 2 are arrays of shape(<tracks>, <genome position>)
        """

        return ops.standardize(x, axis=0)

    @staticmethod
    def fit_features(x):
        """transform each **component** of X so that it fits on an interval [-1, 1].
        So the values of track t at position i are all in [-1,1]

        @param x: a ndarray of shape (nr. examples, nr. features)
        @return: the fitted input
        """

        return ops.fit(x, axis=0)

class TheanoShare( object ):
    """a dataset that can return its data as theano shared variables"""

    def __init__(self):
        self.__shX, self.__shY, self.__shT = None, None, None

    def __sh_anon(self, what, shape=None, borrow=True):
        if getattr(self, what) is None:
            raise ValueError("cannot share non-existent member %s" % what)

        if getattr(self, "_TheanoShare__sh%s" % what) is None:
            init_val = getattr(self, what)
            if shape:
                init_val = init_val.reshape( shape )

            setattr(self, "_TheanoShare__sh%s" % what,
                    theano.shared(init_val, borrow=borrow))
        return getattr(self, "_TheanoShare__sh%s" % what)

    @property
    def shX(self):
        return self.__sh_anon("X")

    @property
    def shY(self):
        return self.__sh_anon("Y")

    @property
    def shT(self):
        return self.__sh_anon("T")

    def share(self, which, shape=None, borrow=True):
        """wrap the data on a thean.shared variable

        @param which: what component to wrap (str, typically 'X', 'T', 'Y')
        @param shape: reshape the array to this shape
        @param borrow: passer to theano.share
        @return: theano.shared instance initialized to the required data"""

        val = getattr(self, which)
        if not (shape is None):
            val = val.reshape( shape )
        return theano.shared(val, borrow=borrow)

class TrainDataset( object ):
    """a mixin for batch functionality, valid and train sub-dataset"""

    def __init__(self, batch_s, tot_s=None, valid_s=None, valid_idx=None, rng=None):
        """Dataset that will
        create train and validation batches from the given params.

        the idea is to split the data into batches and allocate a 'valid_s'
        portion of them for validation. the position of the (continuous) validation
        block is w.r.t batches. E.g., for tot_size = 10, batch_s = 2,
        valid_idx=3, valid_s = 0.3 you get 4 + 1 train + valid batches: T T T V T

        @param tot_s: nr. of examples
        @param batch_s : batch size
        @param valid_s : fraction of data to allocate for validation
        @param valid_idx: batch index at which allocate validation data
        @param rng : numpy.RandomState used to shuffle batches or None (no shuffle)
        @return : (train_batches, valid_batches)"""

        if tot_s is None:
            tot_s = self.X.shape[0]
        if self.X.shape[0] < tot_s:
            log.warning("total size (%d) > dataset size (%d). adjusting ...",
                        tot_s,self.X.shape[0])
        self.total_size = min( tot_s, self.X.shape[0] )

        self.batch_size = batch_s

        if valid_s is None:
            valid_s = 0.25
        self.valid_size = valid_s

        self.n_batches  = self.X.shape[0] / self.batch_size ## nr. of batches
        if valid_idx is None:
            valid_idx = self.n_batches - int(self.n_batches * self.valid_size)
        self.valid_idx = valid_idx

        self.rng = rng

        self.train_batches, self.valid_batches = self.__batches()

    def __batches(self):
        tot_size = self.total_size
        batch_s = self.batch_size
        valid_s = self.valid_size
        valid_idx = self.valid_idx
        rng = self.rng

        if valid_s <= 0 or valid_s >= 1:
            raise ValueError("valid_s (%f) should be between (0, 1) ", valid_s)

        if batch_s > tot_size * min(valid_s, 1-valid_s):
            raise ValueError("batch_s (%d) > min(valid_s=%d, train_s=%d)" % (batch_s,
                             tot_size * valid_s, tot_size * (1-valid_s)) )

        all_batches = range( tot_size / batch_s )
        try:
            valid_batches = all_batches[valid_idx:valid_idx+int(len(all_batches)*valid_s)]
        except IndexError:
            raise ValueError("valid_idx (%d) should be between 0 and %d",
                             valid_idx, len(all_batches)-1)
        train_batches = list( set(all_batches) - set(valid_batches) )
        assert set(train_batches + valid_batches) == set(all_batches)
        assert len( set(train_batches) & set(valid_batches) ) == 0

        if not (rng is None):
            rng.shuffle(train_batches)
            rng.shuffle(valid_batches)

        logging.info("train batches: %s", str(train_batches))
        logging.info("valid batches: %s", str(valid_batches))

        return (train_batches, valid_batches)


    def __iter_batches(self, which, nepochs):
        """infinite loop over train/valid batches

        @param nepochs: loop this many times over train batches (0 will loop forever)
        @return: iterator """

        assert which in ("train_batches", "valid_batches")

        batches = getattr(self, which)
        epoch = 0
        while True:
            for i in batches:
                yield i
            epoch = epoch + 1
            if epoch == nepochs:
                break

    def iter_train(self, nepochs):
        return self.__iter_batches("train_batches", nepochs)

    def iter_valid(self, nepochs):
        return self.__iter_batches("valid_batches", nepochs)

class AnchorAnnotation( object ):
    def __init__(self, anchors, tracks, width, labels):
        self.pX = pd.Panel( self.X, items = anchors, major_axis = tracks,
                           minor_axis = width )
        if not (self.Y is None):
            self.sY = pd.Series( self.Y, index=anchors )

        self.label_names = None
        if not (self.T is None):
            coden = map(lambda v: labels[v], self.T)
            self.dfT = pd.DataFrame({"label_code" : self.T,
                                     "label_name" : coden})
            self.label_names = labels


class AnchorDataset(Dataset, AnchorAnnotation, TheanoShare):
    def __init__(self, X, Y, T):
        """X, Y, T are a Panel, Series, and DataFrame resp."""

        valY = None
        if not (Y is None):
            valY = Y.values
        valT = None
        if not (T is None):
            valT = T["label_code"].values

        Dataset.__init__(self, X.values, valY, valT)
        TheanoShare.__init__(self)
        self.pX, self.sY, self.dfT = X, Y, T

        self.label_names = None
        if not (self.T is None):
            self.label_names = np.unique( self.dfT["label_name"].values ).tolist()

    @property
    def track_names(self):
        return self.pX.major_axis.tolist()

    @property
    def tracks(self):
        return self.X.shape[1]

    @property
    def width(self):
        return self.X.shape[2]

    @property
    def labels(self):
        return len(self.label_names)

    def dump(self, path):
        arch, key = archive.split(path)

        X, Y, T = self.pX, self.sY, self.dfT
        (nsamp, ntrack, width) = X.values.shape

        if not (X is None):
            archive.save_object( arch, "%s/rawX" % key, X )
            normX, meanX, sdX = self.normalize_features(X.values.reshape(nsamp, -1))

            archive.save_object( arch, "%s/X" % key,
                                 pd.Panel(normX.reshape(nsamp, ntrack, width),
                                          items = X.items,
                                          major_axis = X.major_axis,
                                          minor_axis = X.minor_axis)
                                )
            archive.save_object( arch, "%s/meanX" % key,
                                 pd.DataFrame(meanX.reshape(ntrack, width),
                                              index = X.major_axis,
                                              columns=X.minor_axis ) )
            archive.save_object( arch, "%s/sdX" % key,
                                 pd.DataFrame(sdX.reshape(ntrack, width),
                                              index = X.major_axis,
                                              columns=X.minor_axis ) )

        if not (Y is None):
            archive.save_object( arch, "%s/Y" % key, Y )

        if not (T is None):
            archive.save_object( arch, "%s/T" % key, T )

    @classmethod
    def mean_sdX(cls, path):
        ap, did = archive.split(path)

        meanX = archive.load_object(ap, "/".join((did, "meanX")))
        sdX = archive.load_object(ap, "/".join((did, "sdX")))
        return meanX, sdX

    @classmethod
    def _from_archive(cls, path, raw, *args):
        ap, did = archive.split(path)

        key = "%s/%s" % (did, (raw and "rawX" or "X"))
        X = archive.load_object(ap, key)

        def load_none(k, p=ap, did=did):
            try:
                return archive.load_object(p, "%s/%s" % (did, k))
            except Exception:
                pass
        Y, T = load_none("Y"), load_none("T")
        return cls(X, Y, T, *args)

class TrainAnchorDataset(AnchorDataset, TrainDataset):
    def __init__(self, X, Y, T, batch_s,
                 tot_s=None, valid_s=None, valid_idx=None, rng=None):

        AnchorDataset.__init__(self, X, Y, T)
        TrainDataset.__init__(self, batch_s, tot_s, valid_s, valid_idx, rng)

    @classmethod
    def _from_archive(cls, path, raw, batch_s, **kwargs):
        ap, did = archive.split(path)

        key = "%s/%s" % (did, (raw and "rawX" or "X"))
        X = archive.load_object(ap, key)

        def load_none(k, p=ap, did=did):
            try:
                return archive.load_object(p, "%s/%s" % (did, k))
            except Exception:
                pass
        Y, T = load_none("Y"), load_none("T")
        return cls(X, Y, T, batch_s, **kwargs)

