from pyod.utils.data import generate_data
import numpy as np
import pandas as pd
import os

def generate_3D_data(n_sys=1, n_train=1000, n_test=500, n_features=2, contamination=0.1,
                  train_only=False, offset=10, behaviour='old',
                  random_state=None):

    """Utility function to generate synthesized data.
    Normal data is generated by a multivariate Gaussian distribution and
    outliers are generated by a uniform distribution.

    Parameters
    ----------
    n_sys : int, (default=1)
        The number of systems to generate.

    n_train : int, (default=1000)
        The number of training points to generate.

    n_test : int, (default=500)
        The number of test points to generate.

    n_features : int, optional (default=2)
        The number of features (dimensions).

    contamination : float in (0., 0.5), optional (default=0.1)
        The amount of contamination of the data set, i.e.
        the proportion of outliers in the data set. Used when fitting to
        define the threshold on the decision function.

    train_only : bool, optional (default=False)
        If true, generate train data only.

    offset : int, optional (default=10)
        Adjust the value range of Gaussian and Uniform.

    behaviour : str, default='old'
        Behaviour of the returned datasets which can be either 'old' or
        'new'. Passing ``behaviour='new'`` returns
        "X_train, y_train, X_test, y_test", while passing ``behaviour='old'``
        returns "X_train, X_test, y_train, y_test".

    random_state : int, RandomState instance or None, optional (default=None)
        If int, random_state is the seed used by the random number generator;
        If RandomState instance, random_state is the random number generator;
        If None, the random number generator is the RandomState instance used
        by `np.random`.

    Returns
    -------
    X_train : numpy array of shape (n_sys, n_train, n_features) if n_sys > 1 else shape (n_train, n_features)
        Training data.

    y_train : numpy array of shape (n_sys, n_train) if n_sys > 1 else shape (n_train, )
        Training ground truth.

    X_test : numpy array of shape (n_sys, n_test, n_features) if n_sys > 1 else shape (n_test, n_features)
        Test data.

    y_test : numpy array of shape (n_sys, n_test,) if n_sys > 1 else shape (n_test, )
        Test ground truth.

    """

    if n_sys == 1:

        return generate_data(n_train=n_train,
                              n_test=n_test,
                              n_features=n_features,
                              contamination=contamination,
                              train_only=train_only,
                              offset=offset,
                              behaviour=behaviour,
                              random_state=random_state)

    else:
        X_train_sys = []
        y_train_sys = []
        X_test_sys = []
        y_test_sys = []

        for sys_idx in range(n_sys):

            X_train, y_train, X_test, y_test = generate_data(n_train=n_train,
                                                             n_test=n_test,
                                                             n_features=n_features,
                                                             contamination=contamination,
                                                             train_only=train_only,
                                                             offset=offset,
                                                             behaviour=behaviour,
                                                             random_state=random_state)
            X_train_sys.append(X_train)
            y_train_sys.append(y_train)
            X_test_sys.append(X_test)
            y_test_sys.append(y_test)

        # X_train_sys = np.array(X_train_sys)
        # y_train_sys = np.array(y_train_sys)
        # X_test_sys = np.array(X_test_sys)
        # y_test_sys = np.array(y_test_sys)

        return X_train_sys, y_train_sys, X_test_sys, y_test_sys


def load_sys_data(sys_dir, data_dir):

    """Utility function to load system feature.

    Parameters
    ----------
    data_dir: string
        directory of system description

    data_dir: string
        directory of system feature

    Returns
    -------
    X : list
        system features

    Y: numpy array of shape (n_sys, 1),
        groundtruth, 1 marks abnormal system

    sys_info: dictionary
        system number, time-series length and dimention number

    """

    sys_buf = pd.read_csv(sys_dir)

    X_csv_name = sys_buf['system'].values
    Y = sys_buf['label'].values
    X = []
    ts_len = []
    ts_dim = []

    for x_name in X_csv_name:
        ts_data_fname = os.path.join(data_dir, x_name)
        ts_data = pd.read_csv(ts_data_fname).values
        X.append(ts_data)
        ts_len.append(ts_data.shape[0])
        ts_dim.append(ts_data.shape[1])

    sys_info = {'sys_num': Y.shape[0],
                'ts_len': np.array(ts_len),
                'ts_ndim': np.array(ts_dim),
                }

    # print(type(X))
    # print(type(Y))

    return X, Y, sys_info


def generate_sys_feature(sys_feature: list):

    """Utility function to generate system feature.

    Parameters
    ----------
    sys_feature : list
        each element of the sys_feature is feature of each system

    Returns
    -------
    sys_feature_array_with_id : numpy array of shape (n, 2),
        the first column is system id, second column is its feature, n = feature number * system number.

    """

    # ts_len = np.array([ts.shape[0] for ts in sys_feature]).min()
    # ndim = np.array([ts.shape[1] for ts in sys_feature]).min()
    # sys_feature = [ts[:ts_len, :ndim].T.reshape((ts_len*ndim, 1)) for ts in sys_feature]

    sys_id = np.concatenate([idx*np.ones((len(sys_feature[idx]),1)).astype(np.int)
                             for idx in range(len(sys_feature))], axis=0)
    sys_feature_array = np.concatenate(sys_feature, axis=0)
    sys_feature_array_with_id = np.concatenate((sys_id, sys_feature_array), axis=1)

    return sys_feature_array_with_id


# load_sys_data('/Users/didi/Documents/hegsns/hegsns/TODS/tods_Guanchu/tods/datasets/anomaly/system_wise/sample/train.csv',
#               '/Users/didi/Documents/hegsns/hegsns/TODS/tods_Guanchu/tods/datasets/anomaly/system_wise/sample/systems')
# generate_3D_data(n_sys=5)