import numpy as np
import torch
from torch import nn

from utils.dataset_creation import inverse_scale_feature


def plot_prediction_windows(index: int,
                            training_configuration: dict,
                            item_x: np.ndarray,
                            item_y: np.ndarray,
                            preds: np.ndarray,
                            scalers: dict,
                            features_to_plot: list,
                            output_feature_names: list,
                            fig_widget):
    window_features = item_x[index]
    indices = np.arange(window_features.shape[0])
    actual = item_y[index].ravel()
    predicted = preds[index]

    # clear previous traces
    fig_widget.data = []

    for feature in features_to_plot:
        feature_index = feature["index"]
        feature_name = feature["name"]
        fig_widget.add_scatter(
            x=indices,
            y=window_features[:, feature_index],
            mode="lines",
            name=feature_name,
        )

    num_outputs = predicted.shape[-1]
    scaled_preds = list()
    for i in range(num_outputs):
        if isinstance(training_configuration["loss_functions"][i], nn.BCEWithLogitsLoss):
            output = torch.sigmoid(torch.tensor(predicted[i]))
        else:
            output = predicted[i]

        output = float(output)

        scaled_output = inverse_scale_feature(output,
                                              output_feature_names[i],
                                              scalers)

        scaled_preds.append(scaled_output)

    scaled_actuals = list()
    for i in range(num_outputs):
        output = float(actual[i].numpy())
        scaled_output = inverse_scale_feature(output,
                                              output_feature_names[i],
                                              scalers)

        scaled_actuals.append(scaled_output)

    # add actual and predicted values
    print(f"actual: {scaled_actuals}, predicted: {scaled_preds}")