"""Multiple EOF analysis for :py:mod:`numpy` array data."""
# (c) Copyright 2010-2012 Andrew Dawson. All Rights Reserved.
#     
# This file is part of eof2.
# 
# eof2 is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
# 
# eof2 is distributed in the hope that it will be useful, but WITHOUT
# ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
# FITNESS FOR A PARTICULAR PURPOSE.  See the GNU General Public License
# for more details.
# 
# You should have received a copy of the GNU General Public License
# along with eof2.  If not, see <http://www.gnu.org/licenses/>.
import numpy as np
import numpy.ma as ma

from eofsolve import EofSolver
from errors import EofError


class MultipleEofSolver(object):
    """Multiple EOF analysis (:py:mod:`numpy` interface)."""

    def __init__(self, *datasets, **kwargs): 
        """Create a MultipleEofSolver object.

        The EOF solution is computed at initialization time. Method
        calls are used to retrieve computed quantities.
        
        **Arguments:**

        *\*datasets*
            One or more :py:class:`numpy.ndarray`s or
            :py:class:`numpy.ma.core.MasekdArray`s with two or more
            dimensions containing the data to be analysed. The first
            dimension of each array is assumed to represent time.
            Missing values are permitted, either in the form of masked
            arrays, or the value :py:attr:`numpy.nan`. Missing values
            must be constant with time (e.g., values of an oceanographic
            field over land).

        **Optional arguments:**

        *weights*
            A sequence of arrays of weights whose shapes are compatible
            with those of the input data sets. The weights can have the
            same shape as the input data set or a shape compatible with
            a an array broadcast operation (ie. the shape of the weights
            can match the rightmost parts of the shape of the input data
            set). If none of the input data sets require weighting then
            the single value *None* may be used. Defaults to *None* (no
            weighting for any data set).

        *center*
            If *True*, the mean along the first axis of the input data
            set (the time-mean) will be removed prior to analysis. If
            *False*, the mean along the first axis will not be removed.
            Defaults to *True* (mean is removed). Generally this option
            should be set to *True* as the covariance interpretation
            relies on input data being anomalies with a time-mean of 0.
            A valid reson for turning this off would be if you have
            already generated an anomaly data set. Setting to *True* has
            the useful side-effect of propagating missing values along
            the time-dimension, ensuring the solver will work even if
            missing values occur at different locations at different
            times.

        *ddof*
            'Delta degrees of freedom'. The divisor used to normalize
            the covariance matrix is *N - ddof* where *N* is the
            number of samples. Defaults to *1*.

        """

        # Define valid keyword arguments and their default values. This method
        # is required since Python 2.7 cannot accept a variable argument list
        # followed by a set of keyword arguments. For some reason both must be
        # variable.
        keywords = {"weights": None, "center": True, "ddof": 1}
        for kwarg in kwargs.keys():
            if kwarg not in keywords.keys():
                
raise EofError("invalid argument: %s" % kwarg)
        weights = kwargs.get("weights", keywords["weights"])
        center = kwargs.get("center", keywords["center"])
        ddof = kwargs.get("ddof", keywords["ddof"])
        # Record the number of datasets provided.
        self._ndatasets = len(datasets)
        # Initialise instance variables dealing with dataset shapes.
        self._
multirecords = list()
        self._multishapes = list()
        self._multislicers = list()
        self._multichannels = list()
        self._multidtypes = list()
        slicebegin = 0
        for dataset in datasets:
            records = dataset.shape[0]
            shape = dataset.shape[1:]
            channels = np.product(shape)
            slicer = slice(slicebegin, slicebegin+channels)
            slicebegin += channels
            self._multirecords.append(records)
            self._multishapes.append(shape)
            self._multislicers.append(slicer)
            self._multichannels.append(channels)
            self._multidtypes.append(dataset.dtype)
        # Check that all fields have the same time dimension.
        if not (np.array(self._multirecords) == self._multirecords[0]).all():
            raise EofError("all datasets must have the same first dimension")
        # Get the dtype that will be used for the data and weights. This will
        # be the 'highest' dtype of those passed.
        dtype = sorted(self._multidtypes, reverse=True)[0]
        # Form a full array to pass to the EOF solver consisting of all the
        # flat inputs.
        nt = self._multirecords[0]
        ns = self._multichannels.sum()
        dataset = ma.empty([nt, ns], dtype=dtype)
        for iset in xrange(self._ndatasets):
            slicer = self._multislicers[iset]
            channels = self._multichannels[iset]
            dataset[:, slicer] = datasets[iset].reshape([nt, channels])
        # Construct an array of weights the same shape as the data array.
        if weights is not None:
            if len(weights) != self._ndatasets:
                raise EofError("number of weights and datasets differs")
            if not filter(lambda i: False if i is None else True, weights):
                # If every entry in the weights list is None then just pass
                # None to the EofSolver __init__ method.
                warr = None
            else:
                # Construct a spatial weights array.
                warr = np.empty([1, ns], dtype=dtype)
                for iset in xrange(self._ndatasets):
                    slicer = self._multislicers[iset]
                    if weights[iset] is None:
                        # If this dataset has no weights use 1 for the weight
                        # of all elements.
                        warr[:, slicer] = 1.
                    else:
                        # Otherwise use the weights. These need to be
                        # conformed to the correct dimensions.
                        channels = self._multichannels[iset]
                        try:
                            warr[:, slicer] = np.broadcast_arrays(
                                    datasets[iset][0],
                                    weights[iset])[1].reshape([channels])
                        except ValueError:
                            raise EofError("weights are invalid")
        else:
            # Just pass None if none of the input datasets have associated
            # weights.
            warr = None
        # Create an EofSolver object to handle the computations.
        self._solver = EofSolver(dataset, weights=warr, center=center, ddof=1)

    def _unwrap(self, modes):
        """Split a returned mode field into component parts."""
        nmodes = modes.shape[0]
        modeset = [modes[:, slicer].reshape((nmodes,)+shape) \
                for slicer, shape in zip(self._multislicers, self._multishapes)]
        return modeset
        
    def pcs(self, pcscaling=0, npcs=None):
        """Principal component time series (PCs).

        Returns an array where the columns are the ordered PCs.
        
        **Optional arguments:**

        *pcscaling*
            Set the scaling of the retrieved PCs. The following
            values are accepted:

            * *0* : Un-scaled PCs (default).
            * *1* : PCs are scaled to unit variance (divided by the
              s