#
# Copyright 2016 The BigDL 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 pytest
import numpy as np
import pandas as pd
import random
import os

from unittest import TestCase
from bigdl.chronos.data import TSDataset
from bigdl.chronos.data.experimental import XShardsTSDataset
from bigdl.orca.data.pandas import read_csv
from bigdl.orca.common import init_orca_context, stop_orca_context, OrcaContext

from pandas.testing import assert_frame_equal
from numpy.testing import assert_array_almost_equal


def generate_spark_df():
    init_orca_context(cores=8)
    sc = OrcaContext.get_spark_context()
    rdd = sc.range(0, 100)
    from pyspark.ml.linalg import DenseVector
    df = rdd.map(lambda x: (DenseVector(np.random.randn(1, ).astype(np.float)),
                            int(np.random.randint(0, 2, size=())),
                            int(x))).toDF(["feature", "id", "date"])
    return df

def get_ugly_ts_df():
    data = np.random.random_sample((100, 5))
    mask = np.random.random_sample((100, 5))
    newmask = mask.copy()
    mask[newmask >= 0.4] = 2
    mask[newmask < 0.4] = 1
    mask[newmask < 0.2] = 0
    data[mask == 0] = None
    data[mask == 1] = np.nan
    df = pd.DataFrame(data, columns=['a', 'b', 'c', 'd', 'e'])
    df['a'][0] = np.nan  # make sure column 'a' has a N/A
    df["datetime"] = pd.date_range('1/1/2019', periods=100)
    df.loc[50:100, "datetime"] = pd.date_range('1/1/2019', periods=50)
    df["id"] = np.array(['00']*50 + ['01']*50)
    return df

class TestXShardsTSDataset(TestCase):

    def setUp(self):
        self.resource_path = os.path.join(os.path.split(__file__)[0], "../../resources/")

    def tearDown(self):
        pass
    
    @classmethod
    def tearDownClass(cls):
        # stop possible active_spark_context
        from pyspark import SparkContext
        from bigdl.orca.ray import OrcaRayContext
        if SparkContext._active_spark_context is not None:
            print("Stopping spark_orca context")
            sc = SparkContext.getOrCreate()
            if sc.getConf().get("spark.master").startswith("spark://"):
                from bigdl.dllib.nncontext import stop_spark_standalone
                stop_spark_standalone()
            sc.stop()

    def test_xshardstsdataset_initialization(self):
        shards_single = read_csv(os.path.join(self.resource_path, "single.csv"))
        tsdata = XShardsTSDataset.from_xshards(shards_single, dt_col="datetime", target_col="value",
                                               extra_feature_col=["extra feature"], id_col="id")
        assert tsdata._id_list == [0]
        assert tsdata.feature_col == ["extra feature"]
        assert tsdata.target_col == ["value"]
        assert tsdata.dt_col == "datetime"
        assert tsdata.shards.num_partitions() == 1

        tsdata = XShardsTSDataset.from_xshards(shards_single, dt_col="datetime",
                                               target_col=["value"],
                                               extra_feature_col="extra feature", id_col="id")
        assert tsdata._id_list == [0]
        assert tsdata.feature_col == ["extra feature"]
        assert tsdata.target_col == ["value"]
        assert tsdata.dt_col == "datetime"
        assert tsdata.shards.num_partitions() == 1

        tsdata = XShardsTSDataset.from_xshards(shards_single, dt_col="datetime",
                                               target_col=["value"],
                                               extra_feature_col="extra feature")
        assert tsdata._id_list == ["0"]
        assert tsdata.feature_col == ["extra feature"]
        assert tsdata.target_col == ["value"]
        assert tsdata.dt_col == "datetime"
        assert tsdata.shards.num_partitions() == 1

    def test_xshardstsdataset_initialization_multiple(self):
        shards_multiple = read_csv(os.path.join(self.resource_path, "multiple.csv"))
        # legal input
        tsdata = XShardsTSDataset.from_xshards(shards_multiple, dt_col="datetime",
                                               target_col="value",
                                               extra_feature_col=["extra feature"], id_col="id")
        assert tsdata._id_list == [0, 1]
        assert tsdata.feature_col == ["extra feature"]
        assert tsdata.target_col == ["value"]
        assert tsdata.dt_col == "datetime"
        assert tsdata.shards.num_partitions() == 2

        tsdata = XShardsTSDataset.from_xshards(shards_multiple, dt_col="datetime",
                                               target_col=["value"],
                                               extra_feature_col="extra feature", id_col="id")
        assert tsdata._id_list == [0, 1]
        assert tsdata.feature_col == ["extra feature"]
        assert tsdata.target_col == ["value"]
        assert tsdata.dt_col == "datetime"
        assert tsdata.shards.num_partitions() == 2

        tsdata = XShardsTSDataset.from_xshards(shards_multiple, dt_col="datetime",
                                               target_col=["value"],
                                               extra_feature_col="extra feature")
        assert tsdata._id_list == ['0']
        assert tsdata.feature_col == ["extra feature"]
        assert tsdata.target_col == ["value"]
        assert tsdata.dt_col == "datetime"
        assert tsdata.shards.num_partitions() == 1

    def test_xshardstsdataset_split(self):
        shards_multiple = read_csv(os.path.join(self.resource_path, "multiple.csv"))
        # only train and test
        tsdata_train, tsdata_valid, tsdata_test =\
            XShardsTSDataset.from_xshards(shards_multiple, dt_col="datetime", target_col="value",
                                          extra_feature_col=["extra feature"], id_col="id",
                                          with_split=True, val_ratio=0, test_ratio=0.1)
        # standard split with all three sets
        tsdata_train, tsdata_valid, tsdata_test =\
            XShardsTSDataset.from_xshards(shards_multiple, dt_col="datetime", target_col="value",
                                          extra_feature_col=["extra feature"], id_col="id",
                                          with_split=True, val_ratio=0.1, test_ratio=0.1,
                                          largest_look_back=5, largest_horizon=2)

        tsdata_train.feature_col.append("new extra feature")
        assert len(tsdata_train.feature_col) == 2
        assert len(tsdata_valid.feature_col) == 1
        assert len(tsdata_test.feature_col) == 1

        tsdata_train.target_col[0] = "new value"
        assert tsdata_train.target_col[0] == "new value"
        assert tsdata_valid.target_col[0] != "new value"
        assert tsdata_test.target_col[0] != "new value"

    def test_xshardstsdataset_roll_multiple_id(self):
        shards_multiple = read_csv(os.path.join(self.resource_path, "multiple.csv"))
        horizon = random.randint(1, 10)
        lookback = random.randint(1, 20)

        tsdata = XShardsTSDataset.from_xshards(shards_multiple, dt_col="datetime",
                                               target_col="value",
                                               extra_feature_col=["extra feature"], id_col="id")

        with pytest.raises(RuntimeError):
            tsdata.to_xshards()

        # roll train
        tsdata.roll(lookback=lookback, horizon=horizon)
        shards_numpy = tsdata.to_xshards()
        collected_numpy = shards_numpy.collect()  # collect and valid
        x = np.concatenate([collected_numpy[i]['x'] for i in range(len(collected_numpy))], axis=0)
        y = np.concatenate([collected_numpy[i]['y'] for i in range(len(collected_numpy))], axis=0)
        assert x.shape == ((50-lookback-horizon+1)*2, lookback, 2)
        assert y.shape == ((50-lookback-horizon+1)*2, horizon, 1)

        tsdata.roll(lookback=lookback, horizon=horizon,
                    feature_col=["extra feature"], target_col="value")
        shards_numpy = tsdata.to_xshards()
        collected_numpy = shards_numpy.collect()  # collect and valid
        x = np.concatenate([collected_numpy[i]['x'] for i in range(len(collected_numpy))], axis=0)
        y = np.concatenate([collected_numpy[i]['y'] for i in range(len(collected_numpy))], axis=0)
        assert x.shape == ((50-lookback-horizon+1)*2, lookback, 2)
        assert y.shape == ((50-lookback-horizon+1)*2, horizon, 1)

        tsdata.roll(lookback=lookback, horizon=horizon,
                    feature_col=[], target_col="value")
        shards_numpy = tsdata.to_xshards()
        collected_numpy = shards_numpy.collect()  # collect and valid
        x = np.concatenate([collected_numpy[i]['x'] for i in range(len(collected_numpy))], axis=0)
        y = np.concatenate([collected_numpy[i]['y'] for i in range(len(collected_numpy))], axis=0)
        assert x.shape == ((50-lookback-horizon+1)*2, lookback, 1)
        assert y.shape == ((50-lookback-horizon+1)*2, horizon, 1)

        # roll test
        horizon = 0
        lookback = random.randint(1, 20)

        tsdata.roll(lookback=lookback, horizon=horizon)
        shards_numpy = tsdata.to_xshards()
        collected_numpy = shards_numpy.collect()  # collect and valid
        x = np.concatenate([collected_numpy[i]['x'] for i in range(len(collected_numpy))], axis=0)
        assert x.shape == ((50-lookback-horizon+1)*2, lookback, 2)

    def test_xshardstsdataset_impute(self):
        from tempfile import TemporaryDirectory
        tmp_df = get_ugly_ts_df()
        with TemporaryDirectory() as tmpdir:
            file_name = os.path.join(tmpdir, 'impute.csv')
            tmp_df.to_csv(file_name, index=False)
            shards_tmp = read_csv(file_name)
        
            for val in ["last", "const", "linear"]:
                tsdata = XShardsTSDataset.from_xshards(shards_tmp, 
                                                    dt_col="datetime", target_col="e",
                                            extra_feature_col=["a", "b", "c", "d"], id_col="id")
                tsdata.impute(mode=val)
                collected_df = tsdata.shards.collect()
                collected_df = pd.concat(collected_df, axis=0)
                
                assert collected_df.isna().sum().sum() == 0
                assert len(collected_df) == 100

    def test_xshardstsdataset_sparkdf(self):
        df = generate_spark_df()

        # with id
        tsdata = XShardsTSDataset.from_sparkdf(df, dt_col="date",
                                               target_col="feature",
                                               id_col="id")
        tsdata.roll(lookback=4, horizon=2)
        data = tsdata.to_xshards().collect()
        assert data[0]['x'].shape[1] == 4
        assert data[0]['x'].shape[2] == 1
        assert data[0]['y'].shape[1] == 2
        assert data[0]['y'].shape[2] == 1
        assert tsdata.shards.num_partitions() == 2

        # with only 1 id
        tsdata = XShardsTSDataset.from_sparkdf(df, dt_col="date",
                                               target_col="feature")
        tsdata.roll(lookback=4, horizon=2)
        data = tsdata.to_xshards().collect()
        assert data[0]['x'].shape[1] == 4
        assert data[0]['x'].shape[2] == 1
        assert data[0]['y'].shape[1] == 2
        assert data[0]['y'].shape[2] == 1
        assert tsdata.shards.num_partitions() == 1

[((1051, 1077), 'bigdl.orca.common.init_orca_context', 'init_orca_context', ([], {'cores': '(8)'}), '(cores=8)\n', (1068, 1077), False, 'from bigdl.orca.common import init_orca_context, stop_orca_context, OrcaContext\n'), ((1087, 1118), 'bigdl.orca.common.OrcaContext.get_spark_context', 'OrcaContext.get_spark_context', ([], {}), '()\n', (1116, 1118), False, 'from bigdl.orca.common import init_orca_context, stop_orca_context, OrcaContext\n'), ((1456, 1489), 'numpy.random.random_sample', 'np.random.random_sample', (['(100, 5)'], {}), '((100, 5))\n', (1479, 1489), True, 'import numpy as np\n'), ((1501, 1534), 'numpy.random.random_sample', 'np.random.random_sample', (['(100, 5)'], {}), '((100, 5))\n', (1524, 1534), True, 'import numpy as np\n'), ((1711, 1764), 'pandas.DataFrame', 'pd.DataFrame', (['data'], {'columns': "['a', 'b', 'c', 'd', 'e']"}), "(data, columns=['a', 'b', 'c', 'd', 'e'])\n", (1723, 1764), True, 'import pandas as pd\n'), ((1844, 1882), 'pandas.date_range', 'pd.date_range', (['"""1/1/2019"""'], {'periods': '(100)'}), "('1/1/2019', periods=100)\n", (1857, 1882), True, 'import pandas as pd\n'), ((1916, 1953), 'pandas.date_range', 'pd.date_range', (['"""1/1/2019"""'], {'periods': '(50)'}), "('1/1/2019', periods=50)\n", (1929, 1953), True, 'import pandas as pd\n'), ((1969, 2004), 'numpy.array', 'np.array', (["(['00'] * 50 + ['01'] * 50)"], {}), "(['00'] * 50 + ['01'] * 50)\n", (1977, 2004), True, 'import numpy as np\n'), ((2899, 3037), 'bigdl.chronos.data.experimental.XShardsTSDataset.from_xshards', 'XShardsTSDataset.from_xshards', (['shards_single'], {'dt_col': '"""datetime"""', 'target_col': '"""value"""', 'extra_feature_col': "['extra feature']", 'id_col': '"""id"""'}), "(shards_single, dt_col='datetime', target_col=\n    'value', extra_feature_col=['extra feature'], id_col='id')\n", (2928, 3037), False, 'from bigdl.chronos.data.experimental import XShardsTSDataset\n'), ((3331, 3469), 'bigdl.chronos.data.experimental.XShardsTSDataset.from_xshards', 'XShardsTSDataset.from_xshards', (['shards_single'], {'dt_col': '"""datetime"""', 'target_col': "['value']", 'extra_feature_col': '"""extra feature"""', 'id_col': '"""id"""'}), "(shards_single, dt_col='datetime', target_col=\n    ['value'], extra_feature_col='extra feature', id_col='id')\n", (3360, 3469), False, 'from bigdl.chronos.data.experimental import XShardsTSDataset\n'), ((3810, 3935), 'bigdl.chronos.data.experimental.XShardsTSDataset.from_xshards', 'XShardsTSDataset.from_xshards', (['shards_single'], {'dt_col': '"""datetime"""', 'target_col': "['value']", 'extra_feature_col': '"""extra feature"""'}), "(shards_single, dt_col='datetime', target_col=\n    ['value'], extra_feature_col='extra feature')\n", (3839, 3935), False, 'from bigdl.chronos.data.experimental import XShardsTSDataset\n'), ((4446, 4585), 'bigdl.chronos.data.experimental.XShardsTSDataset.from_xshards', 'XShardsTSDataset.from_xshards', (['shards_multiple'], {'dt_col': '"""datetime"""', 'target_col': '"""value"""', 'extra_feature_col': "['extra feature']", 'id_col': '"""id"""'}), "(shards_multiple, dt_col='datetime',\n    target_col='value', extra_feature_col=['extra feature'], id_col='id')\n", (4475, 4585), False, 'from bigdl.chronos.data.experimental import XShardsTSDataset\n'), ((4930, 5069), 'bigdl.chronos.data.experimental.XShardsTSDataset.from_xshards', 'XShardsTSDataset.from_xshards', (['shards_multiple'], {'dt_col': '"""datetime"""', 'target_col': "['value']", 'extra_feature_col': '"""extra feature"""', 'id_col': '"""id"""'}), "(shards_multiple, dt_col='datetime',\n    target_col=['value'], extra_feature_col='extra feature', id_col='id')\n", (4959, 5069), False, 'from bigdl.chronos.data.experimental import XShardsTSDataset\n'), ((5414, 5540), 'bigdl.chronos.data.experimental.XShardsTSDataset.from_xshards', 'XShardsTSDataset.from_xshards', (['shards_multiple'], {'dt_col': '"""datetime"""', 'target_col': "['value']", 'extra_feature_col': '"""extra feature"""'}), "(shards_multiple, dt_col='datetime',\n    target_col=['value'], extra_feature_col='extra feature')\n", (5443, 5540), False, 'from bigdl.chronos.data.experimental import XShardsTSDataset\n'), ((6088, 6277), 'bigdl.chronos.data.experimental.XShardsTSDataset.from_xshards', 'XShardsTSDataset.from_xshards', (['shards_multiple'], {'dt_col': '"""datetime"""', 'target_col': '"""value"""', 'extra_feature_col': "['extra feature']", 'id_col': '"""id"""', 'with_split': '(True)', 'val_ratio': '(0)', 'test_ratio': '(0.1)'}), "(shards_multiple, dt_col='datetime',\n    target_col='value', extra_feature_col=['extra feature'], id_col='id',\n    with_split=True, val_ratio=0, test_ratio=0.1)\n", (6117, 6277), False, 'from bigdl.chronos.data.experimental import XShardsTSDataset\n'), ((6462, 6697), 'bigdl.chronos.data.experimental.XShardsTSDataset.from_xshards', 'XShardsTSDataset.from_xshards', (['shards_multiple'], {'dt_col': '"""datetime"""', 'target_col': '"""value"""', 'extra_feature_col': "['extra feature']", 'id_col': '"""id"""', 'with_split': '(True)', 'val_ratio': '(0.1)', 'test_ratio': '(0.1)', 'largest_look_back': '(5)', 'largest_horizon': '(2)'}), "(shards_multiple, dt_col='datetime',\n    target_col='value', extra_feature_col=['extra feature'], id_col='id',\n    with_split=True, val_ratio=0.1, test_ratio=0.1, largest_look_back=5,\n    largest_horizon=2)\n", (6491, 6697), False, 'from bigdl.chronos.data.experimental import XShardsTSDataset\n'), ((7401, 7422), 'random.randint', 'random.randint', (['(1)', '(10)'], {}), '(1, 10)\n', (7415, 7422), False, 'import random\n'), ((7442, 7463), 'random.randint', 'random.randint', (['(1)', '(20)'], {}), '(1, 20)\n', (7456, 7463), False, 'import random\n'), ((7482, 7621), 'bigdl.chronos.data.experimental.XShardsTSDataset.from_xshards', 'XShardsTSDataset.from_xshards', (['shards_multiple'], {'dt_col': '"""datetime"""', 'target_col': '"""value"""', 'extra_feature_col': "['extra feature']", 'id_col': '"""id"""'}), "(shards_multiple, dt_col='datetime',\n    target_col='value', extra_feature_col=['extra feature'], id_col='id')\n", (7511, 7621), False, 'from bigdl.chronos.data.experimental import XShardsTSDataset\n'), ((9498, 9519), 'random.randint', 'random.randint', (['(1)', '(20)'], {}), '(1, 20)\n', (9512, 9519), False, 'import random\n'), ((10893, 10980), 'bigdl.chronos.data.experimental.XShardsTSDataset.from_sparkdf', 'XShardsTSDataset.from_sparkdf', (['df'], {'dt_col': '"""date"""', 'target_col': '"""feature"""', 'id_col': '"""id"""'}), "(df, dt_col='date', target_col='feature',\n    id_col='id')\n", (10922, 10980), False, 'from bigdl.chronos.data.experimental import XShardsTSDataset\n'), ((11421, 11491), 'bigdl.chronos.data.experimental.XShardsTSDataset.from_sparkdf', 'XShardsTSDataset.from_sparkdf', (['df'], {'dt_col': '"""date"""', 'target_col': '"""feature"""'}), "(df, dt_col='date', target_col='feature')\n", (11450, 11491), False, 'from bigdl.chronos.data.experimental import XShardsTSDataset\n'), ((2515, 2541), 'pyspark.SparkContext.getOrCreate', 'SparkContext.getOrCreate', ([], {}), '()\n', (2539, 2541), False, 'from pyspark import SparkContext\n'), ((2834, 2880), 'os.path.join', 'os.path.join', (['self.resource_path', '"""single.csv"""'], {}), "(self.resource_path, 'single.csv')\n", (2846, 2880), False, 'import os\n'), ((4357, 4405), 'os.path.join', 'os.path.join', (['self.resource_path', '"""multiple.csv"""'], {}), "(self.resource_path, 'multiple.csv')\n", (4369, 4405), False, 'import os\n'), ((5945, 5993), 'os.path.join', 'os.path.join', (['self.resource_path', '"""multiple.csv"""'], {}), "(self.resource_path, 'multiple.csv')\n", (5957, 5993), False, 'import os\n'), ((7333, 7381), 'os.path.join', 'os.path.join', (['self.resource_path', '"""multiple.csv"""'], {}), "(self.resource_path, 'multiple.csv')\n", (7345, 7381), False, 'import os\n'), ((7726, 7753), 'pytest.raises', 'pytest.raises', (['RuntimeError'], {}), '(RuntimeError)\n', (7739, 7753), False, 'import pytest\n'), ((9996, 10016), 'tempfile.TemporaryDirectory', 'TemporaryDirectory', ([], {}), '()\n', (10014, 10016), False, 'from tempfile import TemporaryDirectory\n'), ((10052, 10086), 'os.path.join', 'os.path.join', (['tmpdir', '"""impute.csv"""'], {}), "(tmpdir, 'impute.csv')\n", (10064, 10086), False, 'import os\n'), ((10162, 10181), 'bigdl.orca.data.pandas.read_csv', 'read_csv', (['file_name'], {}), '(file_name)\n', (10170, 10181), False, 'from bigdl.orca.data.pandas import read_csv\n'), ((2118, 2141), 'os.path.split', 'os.path.split', (['__file__'], {}), '(__file__)\n', (2131, 2141), False, 'import os\n'), ((2702, 2725), 'bigdl.dllib.nncontext.stop_spark_standalone', 'stop_spark_standalone', ([], {}), '()\n', (2723, 2725), False, 'from bigdl.dllib.nncontext import stop_spark_standalone\n'), ((10268, 10401), 'bigdl.chronos.data.experimental.XShardsTSDataset.from_xshards', 'XShardsTSDataset.from_xshards', (['shards_tmp'], {'dt_col': '"""datetime"""', 'target_col': '"""e"""', 'extra_feature_col': "['a', 'b', 'c', 'd']", 'id_col': '"""id"""'}), "(shards_tmp, dt_col='datetime', target_col='e',\n    extra_feature_col=['a', 'b', 'c', 'd'], id_col='id')\n", (10297, 10401), False, 'from bigdl.chronos.data.experimental import XShardsTSDataset\n'), ((10621, 10652), 'pandas.concat', 'pd.concat', (['collected_df'], {'axis': '(0)'}), '(collected_df, axis=0)\n', (10630, 10652), True, 'import pandas as pd\n'), ((1304, 1336), 'numpy.random.randint', 'np.random.randint', (['(0)', '(2)'], {'size': '()'}), '(0, 2, size=())\n', (1321, 1336), True, 'import numpy as np\n'), ((1232, 1250), 'numpy.random.randn', 'np.random.randn', (['(1)'], {}), '(1)\n', (1247, 1250), True, 'import numpy as np\n')]