from energy_prediction.EnergyPredictionModel import EnergyPredictionModel
from energy_prediction.EnergyPredictionPipeline import EnergyPredictionPipeline
import paho.mqtt.client as mqtt
import json

broker_address = "localhost"
broker_port = 1883
topic = "sensor_data"
client = mqtt.Client(mqtt.CallbackAPIVersion.VERSION2)

def main():
    prediction_data_pipeline_north = EnergyPredictionPipeline(scaler_path="src\energy_prediction\models\scalerNorth.pkl", wing='north')
    prediction_data_pipeline_south = EnergyPredictionPipeline(scaler_path="src\energy_prediction\models\scalerSouth.pkl", wing='south')

    # Energy Prediction North wing
    energy_prediction_north = EnergyPredictionModel(
        model_path="src/energy_prediction/models/lstm_energy_north_01.keras"
    )
    # Energy Prediction South wing
    energy_prediction_south = EnergyPredictionModel(
        model_path="src/energy_prediction/models/lstm_energy_south_01.keras"
    )

    def on_message(client, userdata, message):
        dfN = prediction_data_pipeline_north.fit(message)
        dfS = prediction_data_pipeline_south.fit(message)
        
        if not(dfN is None and dfS is None):
            outN = energy_prediction_north.pipeline(dfN, prediction_data_pipeline_north.scaler)
            outS = energy_prediction_south.pipeline(dfS, prediction_data_pipeline_south.scaler)
            return outN, outS
        else:
            return None
        
    print("Connecting to broker")
    client.on_message = on_message
    client.connect(broker_address, broker_port)
    client.subscribe(topic)
    client.loop_forever()

if __name__ == "__main__":
    main()