import json
from rtu.RTUAnomalizer import RTUAnomalizer
from rtu.RTUPipeline import RTUPipeline
from vav.VAVPipeline import VAVPipeline
from vav.VAVAnomalizer import VAVAnomalizer
import paho.mqtt.client as mqtt


def main():
    rtu_data_pipeline = RTUPipeline(scaler_path="src/rtu/models/scaler_1.pkl")
    rtu_anomalizer = RTUAnomalizer(
        prediction_model_path="src/rtu/models/lstm_4rtu_smooth_02.keras",
        clustering_model_paths=[
            "src/rtu/models/kmeans_model1.pkl",
            "src/rtu/models/kmeans_model2.pkl",
            "src/rtu/models/kmeans_model3.pkl",
            "src/rtu/models/kmeans_model4.pkl",
        ],
        num_inputs=rtu_data_pipeline.num_inputs,
        num_outputs=rtu_data_pipeline.num_outputs,
    )

    vav_pipeline = VAVPipeline(rtu_id=1, scaler_path="src/vav/models/scaler_vav_1.pkl")

    vav_anomalizer = VAVAnomalizer(prediction_model_path="src/vav/models/lstm__vav_01")
    # print(vav_pipeline.input_col_names)

    # print(len(vav_pipeline.output_col_names))

    def on_message(client, userdata, message):
        # print(json.loads(message.payload.decode()))
        df_new, df_trans = rtu_data_pipeline.fit(message)
        if not df_new is None and not df_trans is None:
            out = rtu_anomalizer.pipeline(df_new, df_trans, rtu_data_pipeline.scaler)

    broker_address = "localhost"
    broker_port = 1883
    topic = "sensor_data"
    client = mqtt.Client(mqtt.CallbackAPIVersion.VERSION1)
    client.on_message = on_message
    client.connect(broker_address, broker_port)
    client.subscribe(topic)
    client.loop_forever()


if __name__ == "__main__":
    main()