# imports
import pdb
import warnings
from tqdm import tqdm
from copy import deepcopy
from functools import partial
from itertools import permutations
from IPython.display import display

import numpy as np
import pandas as pd

import seaborn as sns
from matplotlib import pyplot as plt

from statsmodels.tsa.api import VAR
from statsmodels.tsa.stattools import grangercausalitytests

from sklearn.metrics import (
    mean_squared_error,
    mean_absolute_error,
    mean_absolute_percentage_error,
)
from sklearn.preprocessing import StandardScaler
from sklearn.tree import DecisionTreeRegressor
from sklearn.multioutput import RegressorChain
from sklearn.ensemble import RandomForestRegressor
from sklearn.linear_model import LinearRegression, TheilSenRegressor

from sklearn.utils._testing import ignore_warnings

# notebook display config
sns.set()
tqdm.pandas()

# set which columns are metadata and which ones are smart attribuetes
meta_cols = ["date", "serial_number", "model", "capacity_bytes", "failure"]

# set which columns to use in analysis
# NOTE: this determined based on backblaze research, ibm paper, and SMART wiki
smart_stats_to_keep = [
    5,
    10,
    187,
    188,
    197,
    198,
]
smart_cols = [f"smart_{i}_raw" for i in smart_stats_to_keep]
smart_cols += [f"smart_{i}_normalized" for i in smart_stats_to_keep]

# training setup config
X_COLS = deepcopy(smart_cols)
Y_COLS = deepcopy(smart_cols)

NDAYS_DATA = 6
NDAYS_TO_PREDICT = 6

MODELS = ["baseline", "var", "linreg", "wlinreg", "ridge", "theilsen", "dt", "rf"]

# read in cleaned data from previous notebook
# NOTE: we're only using failed drive timeseries because working drive timeseries dont have a lot of "fluctuations"
# or trends. Even the baseline methods achieve <0.25 MSE for working drives, so including that data would not
# contribute much, and rather worsen the class imbalance
failed_ts_df = pd.read_parquet(
    "../../data/interim/failed_ts_df_1612199672.650661.parquet"
)

# minor preprocessing - failed data

# keep only timeseries with enough data
ts_lengths = failed_ts_df.groupby("serial_number").size()
keep_serials = ts_lengths[ts_lengths > NDAYS_DATA]
failed_ts_df = failed_ts_df[failed_ts_df["serial_number"].isin(keep_serials.index)]

# convert index type to datetime instead of object
failed_ts_df.index = pd.DatetimeIndex(failed_ts_df.index)

failed_ts_df = failed_ts_df.groupby("serial_number").apply(
    # NOTE: ffill HAS TO BE CALLED after interpolate coz serial number and model
    # are strings and canot be interpolated so they stay nan after interpolate
    # if left nan, then during groupby-apply-trainmodel-calcualteloss that row
    # wont appear and therefore there wont be any point in resampling @ D
    lambda g: g.resample("D")
    .interpolate()
    .ffill()
)

# convert back serial number from index to external column
failed_ts_df.reset_index(level=0, drop=True, inplace=True)

# remove serial number from index, as it is already a colunm in df
if failed_ts_df.index.nlevels > 1:
    failed_ts_df.reset_index(level=0, drop=True, inplace=True)

# granger test for each permuatation of smart metrics
feat_perms = [i for i in permutations(Y_COLS, 2)]


def apply_grangercausalitytests(drive_ts, perms, lags=6):
    """
    This bit of code is mostly taken from this blog post:
    https://rishi-a.github.io/2020/05/25/granger-causality.html
    """
    # too little data
    if len(drive_ts) < NDAYS_DATA:
        return np.nan

    resultsdf = pd.Series(index=perms, dtype=np.float)
    for i, c in enumerate(perms):
        # timeseries for the two metrics in question
        df = drive_ts.loc[:, c]

        try:
            # apply test and get minimum p-value for
            res = grangercausalitytests(df, maxlag=lags, verbose=False)
            p_values = [res[i + 1][0]["ssr_chi2test"][1] for i in range(lags)]
            resultsdf.iloc[i] = min(p_values)
        except ValueError:
            resultsdf.iloc[i] = np.nan

    return resultsdf


with warnings.catch_warnings():
    warnings.filterwarnings("ignore")
    ret = failed_ts_df.groupby("serial_number").progress_apply(
        partial(apply_grangercausalitytests, perms=feat_perms)
    )

# lets look at, on average what is the pvalue
mean_pvals = ret.dropna(axis=1, how="all").dropna(axis=0, how="all").mean()
mean_pvals.sort_values().head(25)

# remove smart 10
Y_COLS.remove("smart_10_raw")
Y_COLS.remove("smart_10_normalized")
X_COLS.remove("smart_10_raw")
X_COLS.remove("smart_10_normalized")

def mean_absolute_scaled_error(y_true, y_pred, y_train):
    """
    Slightly modified version of implementation at
    https://github.com/scikit-learn/scikit-learn/issues/18685
    """
    e_j = mean_absolute_error(y_true, y_pred, multioutput="raw_values")
    scale = mean_absolute_error(y_train[1:], y_train[:-1], multioutput="raw_values")
    #     return np.divide(e_j, scale, out=np.zeros_like(e_j), where=e_j!=0)
    return e_j / (1 + scale)

def mean_absolute_error_by_mean(y_true, y_pred, y_train):
    """
    Similar idea as MASE, but scaling by mean only instead
    of mean naive error
    """
    mae = mean_absolute_error(y_true, y_pred, multioutput="raw_values")
    #     return np.divide(mae, y_train.mean(), out=np.zeros_like(mae), where=mae!=0)
    return mae / (1 + y_train.mean())

def plot_losses_metricwise(results_df, ts_df, to_plot=15, loss_name="MAPE"):
    global Y_COLS

    # stds of each device time series
    device_ts_stds = (
        ts_df[Y_COLS + ["serial_number"]].groupby("serial_number").std(ddof=0)
    )

    for col in Y_COLS:
        # plot losses for devices where the timeseries had the most fluctuations
        sers = device_ts_stds[col].sort_values(ascending=False).index[:to_plot]
        results_df.loc[sers, col].unstack(level=0).plot(figsize=(15, 5))

        # plot details
        plt.title(f"{loss_name} over time for {col}")
        plt.ylabel(loss_name)
        plt.xticks(rotation=90)
        plt.show()

def calculate_predictionwindow_loss(y_true, y_pred, y_train=None, loss_name="mse"):
    if loss_name == "mase":
        return mean_absolute_scaled_error(
            y_true=y_true,
            y_pred=y_pred,
            y_train=y_train,
        )
    elif loss_name == "mape":
        return mean_absolute_percentage_error(
            y_true=y_true,
            y_pred=y_pred,
            multioutput="raw_values",
        )
    elif loss_name == "maebm":
        return mean_absolute_error_by_mean(
            y_true=y_true,
            y_pred=y_pred,
            y_train=y_train,
        )
    elif loss_name == "mse":
        try:
            ret = mean_squared_error(
                y_true=y_true,
                y_pred=y_pred,
                multioutput="raw_values",
            )
        except ValueError as ve:
            print(ve)
            print("Error too high. Returning infs")
            ret = [np.inf] * y_pred.shape[1]
        return ret
    else:
        raise NotImplementedError

def get_baseline_windowwise_losses(drive_ts, loss="mse"):
    """
    Predict the next N days the value will be same as today.
    Calculate MSE at the end of each training window
    """
    losses = pd.DataFrame(
        columns=Y_COLS, index=deepcopy(drive_ts.index), dtype=np.float
    )

    # if train window slides beyond len(df)-1 then we wont
    # have any data to calculate loss
    for end_idx in range(NDAYS_DATA, len(drive_ts) - NDAYS_TO_PREDICT + 1):

        # smart metrics values on the last day
        last_day_data = drive_ts[Y_COLS].iloc[end_idx]

        # smart metrics values for next min(NDAYS_TO_PRED, df) days
        future_data = drive_ts[Y_COLS].iloc[end_idx : end_idx + NDAYS_TO_PREDICT]

        # calculate loss
        losses.iloc[end_idx - 1] = calculate_predictionwindow_loss(
            loss_name=loss,
            y_true=future_data,
            y_train=drive_ts[Y_COLS].iloc[end_idx - NDAYS_DATA : end_idx],
            y_pred=last_day_data.values.reshape(1, -1).repeat(
                axis=0,
                repeats=len(future_data),
            ),
        )

    return losses

# calculate pred=same_as_last_day baseline
failedts_baseline_losses = failed_ts_df.groupby("serial_number").progress_apply(
    get_baseline_windowwise_losses
)

# drop rows where there are no results (coz not enough historical data)
failedts_baseline_losses.dropna(axis=0, how="all", inplace=True)

# average mse across days
failedts_baseline_losses.mean(level=0)

# visualize performance for time series with highest stds
plot_losses_metricwise(
    failedts_baseline_losses,
    failed_ts_df,
    loss_name="MSE",
    to_plot=5,
)

# multilevel index, columns df to store all results (raw / unaggregated)

# for easy indexing into multidim df
midx = pd.IndexSlice

# init df
all_losses = pd.DataFrame(
    index=failedts_baseline_losses.index,
    columns=pd.MultiIndex.from_product([Y_COLS, MODELS]),
)

# add baseline loss values
all_losses.loc[:, midx[:, "baseline"]] = failedts_baseline_losses.values
all_losses

# mean and median loss for each train-predict window for each metric

# average value of "average loss foreach disk"
print("Average Value of Mean Loss in each Disk")
display(all_losses.mean(level=0).mean())

# average value of "median loss foreach disk"
print("Average Value of Median Loss in each Disk")
display(all_losses.median(level=0).mean())

def get_var_windowwise_losses(drive_ts, loss="mse"):
    """
    Fit VAR model and calculate loss
    for each train_ts-predict window the given drive
    """
    losses = pd.DataFrame(
        columns=Y_COLS, index=deepcopy(drive_ts.index), dtype=np.float
    )

    for end_idx in range(NDAYS_DATA, len(drive_ts) - NDAYS_TO_PREDICT + 1):

        # historical NDAYS_DATA data
        train_ts = drive_ts[Y_COLS].iloc[end_idx - NDAYS_DATA : end_idx]

        # de-mean the data
        means = train_ts.mean()
        train_ts -= means

        # EXPLICIT YEO JOHNSON TRANSFORM WITH LAMBDA=0
        train_ts_transf = np.where(
            train_ts >= 0,
            np.log(train_ts + 1),
            -((1 - train_ts) ** 2 - 1) / 2,
        )

        # log transform
        train_ts = pd.DataFrame(
            train_ts_transf,
            index=train_ts.index,
            columns=train_ts.columns,
        )

        # 1st order difference
        train_ts = train_ts.diff(1)

        # remove nans coz not supported by statsmodels
        train_ts.dropna(inplace=True)

        # next NDAYS_TO_PREDICT data
        test_ts = drive_ts[Y_COLS].iloc[end_idx : end_idx + NDAYS_TO_PREDICT]

        # init model
        try:
            model = VAR(train_ts, freq="D")
        except IndexError as e:
            print(e)
            pdb.set_trace()
            print("fix this")

        # when col has var=0, dont add constant
        # otherwise add constant, linear, quadratic trends
        if train_ts.std().min() == 0:
            model = model.fit(trend="n")
        else:
            model = model.fit(trend="ctt")

        # model preds
        pred_ts = model.forecast(
            train_ts.values,
            steps=len(test_ts),
        )

        # un difference preds
        pred_ts = pred_ts.cumsum(axis=0) + train_ts.iloc[-1:].values

        # inverse scale results
        #         pred_ts = scaler.inverse_transform(pred_ts)
        pred_ts = np.where(
            pred_ts >= 0, np.exp(pred_ts) - 1, 1 - (-2 * pred_ts + 1) ** (1 / 2)
        )

        # re-mean the data
        pred_ts += means

        # calculate loss
        losses.iloc[end_idx - 1] = calculate_predictionwindow_loss(
            y_train=train_ts,
            y_true=test_ts,
            y_pred=pred_ts,
        )

    return losses

# calculate VAR model performance
# with np.errstate(divide='ignore'):
with warnings.catch_warnings():
    warnings.filterwarnings("ignore")
    failedts_var_losses = failed_ts_df.groupby("serial_number").progress_apply(
        get_var_windowwise_losses
    )

# drop rows where there are no results (coz not enough historical data)
failedts_var_losses.dropna(axis=0, how="all", inplace=True)

# add to results
all_losses.loc[:, midx[:, "var"]] = failedts_var_losses.values

# average loss across days foreach disk
failedts_var_losses.mean(level=0)

# visualize performance for time series with highest stds
plot_losses_metricwise(
    failedts_var_losses.clip(upper=1e9),
    failed_ts_df,
    loss_name="MSE",
    to_plot=5,
)

# improvement in mse in each train-predict window for each metric for each disk
diff = failedts_baseline_losses - failedts_var_losses
diff = diff[diff != -np.inf]

# average improvement for each disk
mean_diff = diff.mean(level=0)

# average value of average improvement (averaged across disks)
print('Average Value of "Mean Improvement Over Baseline in each Disk"')
display(mean_diff.mean())

# median improvement for each disk (mean more sensitive to outliers)
median_diff = diff.median(level=0)

# average value of median improvement (averaged across disks)
print('Average Value of "Median Improvement Over Baseline in each Disk"')
display(median_diff.mean())

def get_model_windowwise_losses(
    drive_ts, model, sample_weight=None, ensemble=False, loss="mse"
):
    """
    Fit provided model and calculate loss
    for each ytrain-predict window the given drive
    """
    losses = pd.DataFrame(
        columns=Y_COLS,
        index=deepcopy(drive_ts.index),
        dtype=np.float,
    )

    # withstd=False because constant cols cause 0 std
    # which causes divide by zero errors
    scaler = StandardScaler(with_std=False)

    # dummy independent variable
    xtrain = np.arange(start=0, stop=NDAYS_DATA).reshape(-1, 1)

    for end_idx in range(NDAYS_DATA, len(drive_ts) - NDAYS_TO_PREDICT + 1):

        # historical NDAYS_DATA data
        ytrain = drive_ts[Y_COLS].iloc[end_idx - NDAYS_DATA : end_idx]

        ytrain = pd.DataFrame(
            scaler.fit_transform(ytrain.values),
            index=ytrain.index,
            columns=ytrain.columns,
        )

        # next NDAYS_TO_PREDICT data
        ytest = drive_ts[Y_COLS].iloc[end_idx : end_idx + NDAYS_TO_PREDICT]

        # dummy independent variables for forecasts
        xtest = np.arange(
            start=NDAYS_DATA,
            stop=NDAYS_DATA + len(ytest),
        ).reshape(-1, 1)

        # fit model
        if ensemble:
            model.fit(
                X=xtrain,
                Y=ytrain,
            )
        else:
            model.fit(
                X=xtrain,
                y=ytrain,
                sample_weight=sample_weight,
            )

        # forecast
        ypred = scaler.inverse_transform(model.predict(xtest))

        # calculate loss
        losses.iloc[end_idx - 1] = calculate_predictionwindow_loss(
            y_true=ytest,
            y_train=ytrain,
            y_pred=ypred,
        )

    return losses

# calculate linreg model performance
failedts_linreg_losses = failed_ts_df.groupby("serial_number").progress_apply(
    partial(
        get_model_windowwise_losses,
        model=LinearRegression(normalize=False),
    )
)

# drop rows where there are no results (coz not enough historical data)
failedts_linreg_losses.dropna(axis=0, how="all", inplace=True)

# add to results
all_losses.loc[:, midx[:, "linreg"]] = failedts_linreg_losses.dropna(how="all").values

# average mse across days
failedts_linreg_losses.mean(level=0)

# visualize performance for time series with highest stds
plot_losses_metricwise(
    failedts_linreg_losses,  # .clip(upper=1e9),
    failed_ts_df,
    loss_name="MSE",
    to_plot=5,
)

# improvement in mse in each train-predict window for each metric for each disk
diff = failedts_baseline_losses - failedts_linreg_losses

# average improvement for each disk
mean_diff = diff.mean(level=0)

# average value of average improvement (averaged across disks)
print("Average Values of Mean Improvement in each Disk")
display(mean_diff.mean())
# display(mean_diff.describe())

# median improvement for each disk (mean more sensitive to outliers)
median_diff = diff.median(level=0)

# average value of median improvement (averaged across disks)
print("Average Values of Median Improvement in each Disk")
display(median_diff.mean())
# display(mean_diff.describe())

# calculate weightedlinreg model performance

# twice the vanilla ols weighting then fall off exponentially
alpha = 2 * (1 / NDAYS_DATA)
exp_weights = [alpha * (1 - alpha) ** i for i in range(NDAYS_DATA - 1)] + [
    (1 - alpha) ** (NDAYS_DATA - 1)
]

failedts_weightedlinreg_losses = failed_ts_df.groupby("serial_number").progress_apply(
    partial(
        get_model_windowwise_losses,
        sample_weight=exp_weights,
        model=LinearRegression(normalize=False),
    )
)

# drop rows where there are no results (coz not enough historical data)
failedts_weightedlinreg_losses.dropna(axis=0, how="all", inplace=True)

# add to results
all_losses.loc[:, midx[:, "wlinreg"]] = failedts_weightedlinreg_losses.values

# average mse across days
failedts_weightedlinreg_losses.mean(level=0)

# visualize performance for time series with highest stds
plot_losses_metricwise(
    failedts_weightedlinreg_losses,  # .clip(upper=1e9),
    failed_ts_df,
    loss_name="MSE",
    to_plot=5,
)

# improvement in mse in each train-predict window for each metric for each disk
diff = failedts_baseline_losses - failedts_weightedlinreg_losses

# average improvement for each disk
mean_diff = diff.mean(level=0)

# average value of average improvement (averaged across disks)
print("Average Value of Mean Improvement in each Disk")
display(mean_diff.mean())

# median improvement for each disk (mean more sensitive to outliers)
median_diff = diff.median(level=0)

# average value of median improvement (averaged across disks)
print("Average Value of Median Improvement in each Disk")
display(median_diff.mean())

# calculate thielsenreg model performance
with ignore_warnings("ConvergenceWarning"):
    failedts_thielsenreg_losses = failed_ts_df.groupby("serial_number").progress_apply(
        partial(
            get_model_windowwise_losses,
            ensemble=True,
            model=RegressorChain(TheilSenRegressor(random_state=0)),
        )
    )

# drop rows where there are no results (coz not enough historical data)
failedts_thielsenreg_losses.dropna(axis=0, how="all", inplace=True)

# add to results
all_losses.loc[:, midx[:, "theilsen"]] = failedts_thielsenreg_losses.values

# average mse across days
failedts_thielsenreg_losses.mean(level=0)

# visualize performance for time series with highest stds
plot_losses_metricwise(
    failedts_thielsenreg_losses,  # .clip(upper=1e9),
    failed_ts_df,
    loss_name="MSE",
    to_plot=5,
)

# improvement in mse in each train-predict window for each metric for each disk
diff = failedts_baseline_losses - failedts_thielsenreg_losses

# average improvement for each disk
mean_diff = diff.mean(level=0)

# average value of average improvement (averaged across disks)
print("Average Value of Mean Improvement in each Disk")
display(mean_diff.mean())

# median improvement for each disk (mean more sensitive to outliers)
median_diff = diff.median(level=0)

# average value of median improvement (averaged across disks)
print("Average Value of Median Improvement in each Disk")
display(median_diff.mean())

# calculate dt model performance
failedts_dt_losses = failed_ts_df.groupby("serial_number").progress_apply(
    partial(
        get_model_windowwise_losses,
        model=DecisionTreeRegressor(random_state=0),
    )
)

# drop rows where there are no results (coz not enough historical data)
failedts_dt_losses.dropna(axis=0, how="all", inplace=True)

# add to results
all_losses.loc[:, midx[:, "dt"]] = failedts_dt_losses.values

# average mse across days
failedts_dt_losses.mean(level=0)

# visualize performance for time series with highest stds
plot_losses_metricwise(
    failedts_dt_losses,  # .clip(upper=1e9),
    failed_ts_df,
    loss_name="MSE",
    to_plot=5,
)

# improvement in mse in each train-predict window for each metric for each disk
diff = failedts_baseline_losses - failedts_dt_losses

# average improvement for each disk
mean_diff = diff.mean(level=0)

# average value of average improvement (averaged across disks)
print("Average Value of Mean Improvement in each Disk")
display(mean_diff.mean())

# median improvement for each disk (mean more sensitive to outliers)
median_diff = diff.median(level=0)

# average value of median improvement (averaged across disks)
print("Average Value of Median Improvement in each Disk")
display(median_diff.mean())

# calculate rf model performance
failedts_rf_losses = failed_ts_df.groupby("serial_number").progress_apply(
    partial(
        get_model_windowwise_losses,
        ensemble=True,
        model=RegressorChain(
            RandomForestRegressor(
                n_estimators=5,
                random_state=0,
            )
        ),
    )
)

# drop rows where there are no results (coz not enough historical data)
failedts_rf_losses.dropna(axis=0, how="all", inplace=True)

# add to results
all_losses.loc[:, midx[:, "rf"]] = failedts_rf_losses.values

# average mse across days
failedts_rf_losses.mean(level=0)

# visualize performance for time series with highest stds
plot_losses_metricwise(
    failedts_rf_losses,  # .clip(upper=1e9),
    failed_ts_df,
    loss_name="MSE",
    to_plot=5,
)

# improvement in mse in each train-predict window for each metric for each disk
diff = failedts_baseline_losses - failedts_rf_losses

# average improvement for each disk
mean_diff = diff.mean(level=0)

# average value of average improvement (averaged across disks)
print("Average Value of Mean Improvement in each Disk")
display(mean_diff.mean())

# median improvement for each disk (mean more sensitive to outliers)
median_diff = diff.median(level=0)

# average value of median improvement (averaged across disks)
print("Average Value of Median Improvement in each Disk")
display(median_diff.mean())