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


class RTUPipeline:
    scaler1 = None  # RTU 1,2
    scaler2 = None  # RTU 3,4

    def __init__(self, rtus=[1, 2, 3, 4], scaler1_path=None, scaler2_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 = []

        for rtu in rtus:
            self.input_col_names.append(f"rtu_00{rtu}_sat_sp_tn")

        self.input_col_names = 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) - 2
        self.num_outputs = len(self.output_col_names) - 14
        self.column_names = self.output_col_names + self.input_col_names

        if scaler1_path:
            self.scaler1 = self.get_scaler(scaler1_path)
        if scaler2_path:
            self.scaler2 = self.get_scaler(scaler2_path)

        self.df = pd.DataFrame(columns=self.column_names)

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

    def transform_window(self, df_window):
        columns_scaler1 = [0] + list(range(1, 15)) + [29, 30] + list(range(33, 38))
        columns_scaler2 = [0] + list(range(15, 29)) + [31, 32] + list(range(33, 38))
        return self.scaler1.transform(
            df_window.iloc[:, columns_scaler1].values
        ), self.scaler2.transform(df_window.iloc[:, columns_scaler2].values)

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

    def extract_data_from_message(self, df: pd.DataFrame):
        df = df[self.column_names]
        len_df = len(self.df)
        if len(self.df) != 0:
            self.df = pd.concat([self.df, df], axis=0)
        else:
            self.df = df

        if len_df > 31:
            self.df = self.df.iloc[len_df - 31 : len_df]
            self.df.loc[len_df] = self.df.mean()
            return self.df
        else:
            return None

    def fit(self, df: pd.DataFrame):

        df_window = self.extract_data_from_message(df)
        if df_window is not None:
            df_trans1, df_trans2 = self.transform_window(df_window)
            df_new1 = self.prepare_input(df_trans1)
            df_new2 = self.prepare_input(df_trans2)
        else:
            df_new1 = None
            df_trans1 = None
            df_new2 = None
            df_trans2 = None
        return df_new1, df_trans1, df_new2, df_trans2