import json
import joblib
import pandas as pd
from sklearn.preprocessing import StandardScaler
from pickle import load
import numpy as np


class RTUPipeline:
    scaler = None

    def __init__(self, rtus=[1, 2], scaler_path=None):

        outputs = [
            "sa_temp",
            "oadmpr_pct",
            "ra_temp",
            "oa_temp",
            "ma_temp",
            "sf_vfd_spd_fbk_tn",
            "rf_vfd_spd_fbk_tn",
        ]
        self.output_col_names = [
            "hp_hws_temp",
        ]

        for rtu in rtus:
            for output in outputs:
                self.output_col_names.append(f"rtu_00{rtu}_{output}")

        self.input_col_names = [
            "air_temp_set_1",
            "air_temp_set_2",
            "dew_point_temperature_set_1d",
            "relative_humidity_set_1",
            "solar_radiation_set_1",
        ]
        self.num_inputs = len(self.input_col_names)
        self.num_outputs = len(self.output_col_names)
        self.column_names = self.output_col_names + self.input_col_names

        if scaler_path:
            self.scaler = self.get_scaler(scaler_path)
        self.df = pd.DataFrame(columns=self.column_names)

    def get_scaler(self, scaler_path):
        return joblib.load(scaler_path)

    def get_window(self, df):
        len_df = len(df)
        if len_df > 30:
            return df[len_df - 31 : len_df].astype("float32")
        else:
            return None

    def transform_window(self, df_window):
        return self.scaler.transform(df_window)

    def prepare_input(self, df_trans):
        return df_trans[:30, :].reshape((1, 30, len(self.column_names)))

    def extract_data_from_message(self, message):
        payload = json.loads(message.payload.decode())

        len_df = len(self.df)

        k = {}
        for col in self.column_names:
            k[col] = payload[col]
        self.df.loc[len_df] = k
        return self.df

    def fit(self, message):
        df = self.extract_data_from_message(message)
        df_window = self.get_window(df)
        if df_window is not None:
            df_trans = self.transform_window(df_window)
            df_new = self.prepare_input(df_trans)
        else:
            df_new = None
            df_trans = None
        return df_new, df_trans