#
# 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

[(34, 'bigdl.orca.common.init_orca_context', 'init_orca_context', False, 'from bigdl.orca.common import init_orca_context, stop_orca_context, OrcaContext\n'), (35, 'bigdl.orca.common.OrcaContext.get_spark_context', 'OrcaContext.get_spark_context', False, 'from bigdl.orca.common import init_orca_context, stop_orca_context, OrcaContext\n'), (44, 'numpy.random.random_sample', 'np.random.random_sample', True, 'import numpy as np\n'), (45, 'numpy.random.random_sample', 'np.random.random_sample', True, 'import numpy as np\n'), (52, 'pandas.DataFrame', 'pd.DataFrame', True, 'import pandas as pd\n'), (54, 'pandas.date_range', 'pd.date_range', True, 'import pandas as pd\n'), (55, 'pandas.date_range', 'pd.date_range', True, 'import pandas as pd\n'), (56, 'numpy.array', 'np.array', True, 'import numpy as np\n'), (82, 'bigdl.chronos.data.experimental.XShardsTSDataset.from_xshards', 'XShardsTSDataset.from_xshards', False, 'from bigdl.chronos.data.experimental import XShardsTSDataset\n'), (90, 'bigdl.chronos.data.experimental.XShardsTSDataset.from_xshards', 'XShardsTSDataset.from_xshards', False, 'from bigdl.chronos.data.experimental import XShardsTSDataset\n'), (99, 'bigdl.chronos.data.experimental.XShardsTSDataset.from_xshards', 'XShardsTSDataset.from_xshards', False, 'from bigdl.chronos.data.experimental import XShardsTSDataset\n'), (111, 'bigdl.chronos.data.experimental.XShardsTSDataset.from_xshards', 'XShardsTSDataset.from_xshards', False, 'from bigdl.chronos.data.experimental import XShardsTSDataset\n'), (120, 'bigdl.chronos.data.experimental.XShardsTSDataset.from_xshards', 'XShardsTSDataset.from_xshards', False, 'from bigdl.chronos.data.experimental import XShardsTSDataset\n'), (129, 'bigdl.chronos.data.experimental.XShardsTSDataset.from_xshards', 'XShardsTSDataset.from_xshards', False, 'from bigdl.chronos.data.experimental import XShardsTSDataset\n'), (142, 'bigdl.chronos.data.experimental.XShardsTSDataset.from_xshards', 'XShardsTSDataset.from_xshards', False, 'from bigdl.chronos.data.experimental import XShardsTSDataset\n'), (147, 'bigdl.chronos.data.experimental.XShardsTSDataset.from_xshards', 'XShardsTSDataset.from_xshards', False, 'from bigdl.chronos.data.experimental import XShardsTSDataset\n'), (164, 'random.randint', 'random.randint', False, 'import random\n'), (165, 'random.randint', 'random.randint', False, 'import random\n'), (167, 'bigdl.chronos.data.experimental.XShardsTSDataset.from_xshards', 'XShardsTSDataset.from_xshards', False, 'from bigdl.chronos.data.experimental import XShardsTSDataset\n'), (203, 'random.randint', 'random.randint', False, 'import random\n'), (234, 'bigdl.chronos.data.experimental.XShardsTSDataset.from_sparkdf', 'XShardsTSDataset.from_sparkdf', False, 'from bigdl.chronos.data.experimental import XShardsTSDataset\n'), (246, 'bigdl.chronos.data.experimental.XShardsTSDataset.from_sparkdf', 'XShardsTSDataset.from_sparkdf', False, 'from bigdl.chronos.data.experimental import XShardsTSDataset\n'), (74, 'pyspark.SparkContext.getOrCreate', 'SparkContext.getOrCreate', False, 'from pyspark import SparkContext\n'), (81, 'os.path.join', 'os.path.join', False, 'import os\n'), (109, 'os.path.join', 'os.path.join', False, 'import os\n'), (139, 'os.path.join', 'os.path.join', False, 'import os\n'), (163, 'os.path.join', 'os.path.join', False, 'import os\n'), (171, 'pytest.raises', 'pytest.raises', False, 'import pytest\n'), (214, 'tempfile.TemporaryDirectory', 'TemporaryDirectory', False, 'from tempfile import TemporaryDirectory\n'), (215, 'os.path.join', 'os.path.join', False, 'import os\n'), (217, 'bigdl.orca.data.pandas.read_csv', 'read_csv', False, 'from bigdl.orca.data.pandas import read_csv\n'), (62, 'os.path.split', 'os.path.split', False, 'import os\n'), (77, 'bigdl.dllib.nncontext.stop_spark_standalone', 'stop_spark_standalone', False, 'from bigdl.dllib.nncontext import stop_spark_standalone\n'), (220, 'bigdl.chronos.data.experimental.XShardsTSDataset.from_xshards', 'XShardsTSDataset.from_xshards', False, 'from bigdl.chronos.data.experimental import XShardsTSDataset\n'), (225, 'pandas.concat', 'pd.concat', True, 'import pandas as pd\n'), (39, 'numpy.random.randint', 'np.random.randint', True, 'import numpy as np\n'), (38, 'numpy.random.randn', 'np.random.randn', True, 'import numpy as np\n')]