import os 
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
%matplotlib inline

raw_df = pd.read_csv("data/data.csv")

"""
Explanation: Network Traffic Forecasting with AutoTSEstimator
In telco, accurate forecast of KPIs (e.g. network traffic, utilizations, user experience, etc.) for communication networks ( 2G/3G/4G/5G/wired) can help predict network failures, allocate resource, or save energy. 
In this notebook, we demostrate a reference use case where we use the network traffic KPI(s) in the past to predict traffic KPI(s) in the future. We demostrate how to use AutoTS in project Chronos to do time series forecasting in an automated and distributed way.
For demonstration, we use the publicly available network traffic data repository maintained by the WIDE project and in particular, the network traffic traces aggregated every 2 hours (i.e. AverageRate in Mbps/Gbps and Total Bytes) in year 2018 and 2019 at the transit link of WIDE to the upstream ISP (dataset link). 
Download raw dataset and load into dataframe
Now we download the dataset and load it into a pandas dataframe. Steps are as below. 


First, run the script get_data.sh to download the raw data. It will download the monthly aggregated traffic data in year 2018 and 2019 into data folder. The raw data contains aggregated network traffic (average MBPs and total bytes) as well as other metrics. 


Second, run extract_data.sh to extract relavant traffic KPI's from raw data, i.e. AvgRate for average use rate, and total for total bytes. The script will extract the KPI's with timestamps into data/data.csv.


Finally, use pandas to load data/data.csv into a dataframe as shown below
End of explanation
"""


raw_df.head()

"""
Explanation: Below are some example records of the data
End of explanation
"""


df = pd.DataFrame(pd.to_datetime(raw_df.StartTime))

# we can find 'AvgRate' is of two scales: 'Mbps' and 'Gbps'
raw_df.AvgRate.str[-4:].unique() 

# Unify AvgRate value
df['AvgRate'] = raw_df.AvgRate.apply(lambda x:float(x[:-4]) if x.endswith("Mbps") else float(x[:-4])*1000)

df["total"] = raw_df["total"]

df.head()

df.describe()

"""
Explanation: Data pre-processing
Now we need to do data cleaning and preprocessing on the raw data. Note that this part could vary for different dataset. 
For the network traffic data we're using, the processing contains 2 parts:
1. Convert string datetime to TimeStamp
2. Unify the measurement scale for AvgRate value - some uses Mbps, some uses Gbps
End of explanation
"""


ax = df.plot(y='AvgRate',figsize=(12,5), title="AvgRate of network traffic data")

"""
Explanation: Plot the data to see how the KPI's look like
End of explanation
"""


from zoo.orca import init_orca_context
init_orca_context(cores=10, init_ray_on_spark=True)

"""
Explanation: Time series forecasting with AutoTS
AutoTS provides AutoML support for building end-to-end time series analysis pipelines (including automatic feature generation, model selection and hyperparameter tuning).
The general workflow using automated training contains below two steps. 
   1. create a AutoTSEstimator to train a TSPipeline, save it to file to use later or elsewhere if you wish.
   2. use TSPipeline to do prediction, evaluation, and incremental fitting as well. 
Chronos uses Orca to enable distributed training and AutoML capabilities. Init orca as below. View Orca Context for more details. Note that argument init_ray_on_spark must be True for Chronos.
End of explanation
"""


from zoo.chronos.autots import AutoTSEstimator, TSPipeline
import torch
import zoo.orca.automl.hp as hp
auto_estimator = AutoTSEstimator(model='lstm',
                                 search_space="normal",
                                 past_seq_len=hp.randint(50, 100),
                                 future_seq_len=1,
                                 metric="mse",
                                 cpus_per_trial=2)

"""
Explanation: Then we initialize a AutoTSEstimator.
End of explanation
"""


from zoo.chronos.data import TSDataset
from sklearn.preprocessing import StandardScaler

tsdata_train, tsdata_val, tsdata_test = TSDataset.from_pandas(df, 
                                                              dt_col="StartTime", 
                                                              target_col="AvgRate", 
                                                              with_split=True, 
                                                              val_ratio=0.1, 
                                                              test_ratio=0.1)
standard_scaler = StandardScaler()
for tsdata in [tsdata_train, tsdata_val, tsdata_test]:
    tsdata.gen_dt_feature(one_hot_features=["HOUR", "WEEKDAY"])\
          .impute(mode="last")\
          .scale(standard_scaler, fit=(tsdata is tsdata_train))

"""
Explanation: We need to split the data frame into train, validation and test data frame before training.
Then we impute the data to handle missing data and scale the data.
You can use TSDataset as an easy way to finish it.
End of explanation
"""


%%time
ts_pipeline = auto_estimator.fit(data=tsdata_train,
                   epochs=20,
                   batch_size=128,
                   validation_data=tsdata_val,
                   n_sampling=4)

"""
Explanation: Then we fit on train data and evaluate on validation data.
End of explanation
"""


best_config = auto_estimator.get_best_config()
best_config

"""
Explanation: We get a TSPipeline after training. Let's print the hyper paramters selected.
Note that past_seq_len is the lookback value that is automatically chosen
End of explanation
"""


y_pred = ts_pipeline.predict(tsdata_test)

"""
Explanation: We use tspipeline to predict and evaluate.
End of explanation
"""


# plot the predicted values and actual values
lookback = best_config['past_seq_len']
plt.figure(figsize=(16,6))
test_df = tsdata_test.unscale().to_pandas()
tsdata_test.scale(standard_scaler, fit=False)
plt.plot(test_df.StartTime[lookback - 1:], y_pred[:,0,0], color='red', label='predicted values')
plt.plot(test_df.StartTime[lookback - 1:], test_df.AvgRate[lookback - 1:], color='blue', label='actual values')
plt.title('the predicted values and actual values (for the test data)')
plt.xlabel('datetime')
plt.legend(loc='upper left')
plt.show()

"""
Explanation: plot actual and prediction values for AvgRate KPI
End of explanation
"""


mse, smape = ts_pipeline.evaluate(tsdata_test, metrics=["mse", "smape"])
print("Evaluate: the mean square error is", mse)
print("Evaluate: the smape value is", smape)

"""
Explanation: Calculate mean square error and the symetric mean absolute percentage error.
End of explanation
"""


# save pipeline file
my_ppl_file_path = "/tmp/saved_pipeline"
ts_pipeline.save(my_ppl_file_path)

"""
Explanation: You can save the pipeline to file and reload it to do incremental fitting or others.
End of explanation
"""


from zoo.orca import stop_orca_context
stop_orca_context()

"""
Explanation: You can stop the orca context after auto training.
End of explanation
"""


new_ts_pipeline = TSPipeline.load(my_ppl_file_path)

"""
Explanation: Next, we demonstrate how to do incremental fitting with your saved pipeline file.
First load saved pipeline file.
End of explanation
"""


new_ts_pipeline.fit(tsdata_val)

"""
Explanation: Then do incremental fitting with TSPipeline.fit(). We use validation data frame as additional data for demonstration. You can use your new data frame.
End of explanation
"""


# predict results of test_df
y_pred = new_ts_pipeline.predict(tsdata_test)
lookback = best_config['past_seq_len']
plt.figure(figsize=(16,6))
test_df = tsdata_test.unscale().to_pandas()
tsdata_test.scale(standard_scaler, fit=False)
plt.plot(test_df.StartTime[lookback - 1:], y_pred[:,0,0], color='red', label='predicted values')
plt.plot(test_df.StartTime[lookback - 1:], test_df.AvgRate[lookback - 1:], color='blue', label='actual values')
plt.title('the predicted values and actual values (for the test data)')
plt.xlabel('datetime')
plt.legend(loc='upper left')
plt.show()

"""
Explanation: predict and plot the result after incremental fitting.
End of explanation
"""


mse, smape = new_ts_pipeline.evaluate(tsdata_test, metrics=["mse", "smape"])
print("Evaluate: the mean square error is", mse)
print("Evaluate: the smape value is", smape)

"""
Explanation: Calculate mean square error and the symetric mean absolute percentage error.
End of explanation
"""


