#
# Copyright 2018 Analytics Zoo Authors.
#
# 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 pandas as pd
import numpy as np


def split_timeseries_dataframe(df,
                               id_col,
                               val_ratio=0,
                               test_ratio=0.1,
                               look_back=0,
                               horizon=1):
    """
    split input dataframe into train_df, val_df and test_df according to split ratio.
    The dataframe is splitted in its originally order in timeline.
    e.g. |......... train_df(80%) ........ | ... val_df(10%) ...| ...test_df(10%)...|
    :param df: dataframe to be splitted
    :param id_col: id column name
    :param val_ratio: validation ratio
    :param test_ratio: test ratio
    :param look_back: the length to look back
    :param horizon: num of steps to look forward

    :return: splited dataframes
    """
    id_list = list(np.unique(df[id_col]))
    train_df_list = []
    valid_df_list = []
    test_df_list = []
    for id_name in id_list:
        train_df_id, valid_df_id, test_df_id = \
            split_single_timeseries_dataframe(df=df[df[id_col] == id_name],
                                              val_ratio=val_ratio,
                                              test_ratio=test_ratio,
                                              look_back=look_back,
                                              horizon=horizon)
        train_df_list.append(train_df_id)
        valid_df_list.append(valid_df_id)
        test_df_list.append(test_df_id)
    train_df = pd.concat(train_df_list)
    valid_df = pd.concat(valid_df_list)
    test_df = pd.concat(test_df_list)
    return train_df, valid_df, test_df


def split_single_timeseries_dataframe(df,
                                      val_ratio=0,
                                      test_ratio=0.1,
                                      look_back=0,
                                      horizon=1):
    total_num = df.index.size
    test_num = int(total_num * test_ratio)
    val_num = int(total_num * val_ratio)

    test_split_index = test_num + look_back + horizon - 1
    val_split_index = test_split_index + val_num

    train_df = df.iloc[:-(test_num + val_num)]
    val_df = df.iloc[-val_split_index: -test_num]
    test_df = df.iloc[-test_split_index:]

    return train_df, val_df, test_df

@@ -608,27 +608,8 @@
  pd%0A
-import numpy as np%0A
 %0A%0Ade
@@ -1396,267 +1396,195 @@
     
-id_list = list(np.unique(df%5Bid_col%5D))%0A    train_df_list = %5B%5D%0A    valid_df_list = %5B%5D%0A    test_df_list = %5B%5D%0A    for id_name in id_list:%0A        train_df_id, valid_df_id, test_df_id = %5C%0A            split_single_timeseries_dataframe(df=df%5Bdf%5Bid_col%5D == id_name%5D,%0A
+split_result = df.groupby(id_col).apply(lambda df:%0A                                            split_single_timeseries_dataframe(%0A                                                df=df,%0A  
     
@@ -1684,32 +1684,34 @@
                 
+  
 test_ratio=test_
@@ -1755,32 +1755,34 @@
                 
+ 
+ 
 look_back=look_b
@@ -1786,16 +1786,18 @@
 k_back,%0A
+  
         
@@ -1863,240 +1863,268 @@
     
-train_df_list.append(train_df_id)%0A        valid_df_list.append(valid_df_id)%0A        test_df_list.append(test_df_id)%0A    train_df = pd.concat(train_df_list)%0A    valid_df = pd.concat(valid_df_list)%0A    test_df = pd.concat(test_df_list
+                                    )%0A    train_df = pd.concat(%5Bsplit_result%5Bi%5D%5B0%5D for i in split_result.index%5D)%0A    valid_df = pd.concat(%5Bsplit_result%5Bi%5D%5B1%5D for i in split_result.index%5D)%0A    test_df = pd.concat(%5Bsplit_result%5Bi%5D%5B2%5D for i in split_result.index%5D
 )%0A  

enhance split (#4455)