#   Copyright 2020 The PyMC Developers
#
#   Licensed under the Apache License, Version 2.0 (the "License");
#   you may not use this file except in compliance with the License.
#   You may obtain a copy of the License at
#
#       http://www.apache.org/licenses/LICENSE-2.0
#
#   Unless required by applicable law or agreed to in writing, software
#   distributed under the License is distributed on an "AS IS" BASIS,
#   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#   See the License for the specific language governing permissions and
#   limitations under the License.

import numpy as np
import pandas as pd
import theano.tensor as tt


def any_to_tensor_and_labels(x, labels=None):
    """Util for converting input x to tensor trying to
    create labels for columns if they are not provided.

    Default names for columns are ['x0', 'x1', ...], for mappable
    arrays (e.g. pd.DataFrame) their names are treated as labels.
    You can override them with `labels` argument.

    If you have tensor input you should provide labels as we
    cannot get their shape directly

    If you pass dict input we cannot rely on labels order thus dict
    keys are treated as labels anyway

    Parameters
    ----------
    x: np.ndarray | pd.DataFrame | tt.Variable | dict | list
    labels: list - names for columns of output tensor

    Returns
    -------
    (x, labels) - tensor and labels for its columns
    """
    if isinstance(labels, str):
        labels = [labels]
    # pandas.DataFrame
    # labels can come from here
    # we can override them
    if isinstance(x, pd.DataFrame):
        if not labels:
            labels = x.columns
        x = x.to_numpy()

    # pandas.Series
    # there can still be a label
    # we can override labels
    elif isinstance(x, pd.Series):
        if not labels:
            labels = [x.name]
        x = x.to_numpy()[:, None]

    # dict
    # labels are keys,
    # cannot override them
    elif isinstance(x, dict):
        # try to do it via pandas
        try:
            x = pd.DataFrame.from_dict(x)
            labels = x.columns
            x = x.to_numpy()
        # some types fail there
        # another approach is to construct
        # variable by hand
        except (ValueError, TypeError):
            res = []
            labels = []
            for k, v in x.items():
                res.append(v)
                labels.append(k)
            x = tt.stack(res, axis=1)
            if x.ndim == 1:
                x = x[:, None]
    # case when it can appear to be some
    # array like value like lists of lists
    # numpy deals with it
    elif not isinstance(x, tt.Variable):
        x = np.asarray(x)
        if x.ndim == 0:
            raise ValueError("Cannot use scalars")
        elif x.ndim == 1:
            x = x[:, None]
    # something really strange goes here,
    # but user passes labels trusting seems
    # to be a good option
    elif labels is not None:
        x = tt.as_tensor_variable(x)
        if x.ndim == 0:
            raise ValueError("Cannot use scalars")
        elif x.ndim == 1:
            x = x[:, None]
    else:  # trust input
        pass
    # we should check that we can extract labels
    if labels is None and not isinstance(x, tt.Variable):
        labels = ["x%d" % i for i in range(x.shape[1])]
    # for theano variables we should have labels from user
    elif labels is None:
        raise ValueError("Please provide labels as " "we cannot infer shape of input")
    else:  # trust labels, user knows what he is doing
        pass
    # it's time to check shapes if we can
    if not isinstance(x, tt.Variable):
        if not len(labels) == x.shape[1]:
            raise ValueError(
                "Please provide full list "
                "of labels for coefficients, "
                "got len(labels)=%d instead of %d" % (len(labels), x.shape[1])
            )
    else:
        # trust labels, as we raised an
        # error in bad case, we have labels
        pass
    # convert labels to list
    if isinstance(labels, pd.RangeIndex):
        labels = ["x%d" % i for i in labels]
    # maybe it was a tuple ot whatever
    elif not isinstance(labels, list):
        labels = list(labels)
    # as output we need tensor
    if not isinstance(x, tt.Variable):
        x = tt.as_tensor_variable(x)
        # finally check dimensions
        if x.ndim == 0:
            raise ValueError("Cannot use scalars")
        elif x.ndim == 1:
            x = x[:, None]
    return x, labels
