from util import pos, clamp
import parameters
from time import time, sleep
import zmq
import syslab
from syslab.comm.LogUtils import setup_event_logger
logging = setup_event_logger()


### Parameters
# If soc is below this limit, divert all power to the battery.
base_soc_lower_limit = 0.2
# If soc is above this limit, divert all power to the mobile load.
base_soc_upper_limit = 0.8
# This much renewable production shifts our soc_lower_limit down by 0.1 and soc_upper_limit up by 0.1.
# I.e. if there is a lot of renewable production, the range over which we split between battery and
# load widens.
base_renewable_shift = 10.0


# # Variables
# Controller variables
x_battery = 0.5  # Default splitting factor
x_load = 0.5  # Default splitting factor
# Info on battery state
battery_soc = 0.5  # Default SOC (= state of charge)

soc_lower_limit = 0.2
soc_upper_limit = 0.8

# # Communication
# Handle the sockets we need
# Make a context that we use to set up sockets
context = zmq.Context()

# Set up a socket we can use to broadcast splitting factors on
splitting_out_socket = context.socket(zmq.PUB)
splitting_out_socket.bind(f"tcp://*:{parameters.SUPERVISOR_PORT}")

# Set up a socket to subscribe to the battery's soc
soc_in_socket = context.socket(zmq.SUB)
soc_in_socket.connect(f"tcp://{parameters.BATTERY_SOC_IP}:{parameters.BATTERY_SOC_PORT}")

# Ensure we only see message on the battery's soc
soc_in_socket.subscribe(parameters.TOPIC_BATTERY_SOC)

### Unit connections
# TODO step 1.3: Set up connection to the units

try:
    while True:
        try:
            # Try to connect the the battery.
            # If we have a connection to the battery, get its current soc.
            # If none have come in, continue with previous soc.
            # We put this in a while-loop to ensure we empty the queue each time,
            # so we always have the latest value.
            while True:
                # Receive the latest splitting factor
                incoming_str = soc_in_socket.recv_string(flags=zmq.NOBLOCK)
                # The incoming string will look like "batt_split;0.781",
                # so we split it up and take the last bit forward.
                battery_soc = float(incoming_str.split(" ")[-1])
                logging.info(f"New battery soc: {battery_soc}")
        except zmq.Again as e:
            # No (more) new messages, move along
            pass

        # Poll the grid connection to get the current production of renewables.
        wind_p = 4.0  # TODO Task 1.3: Change into syslab connection. Remember that the measurements need to be collected directly from the units since vswitchboard is not working
        solar_p = 7.0
        logging.info(f"Current renewable production: {wind_p + solar_p} kW.")

        soc_lower_limit = base_soc_lower_limit - (wind_p + solar_p) / base_renewable_shift / 10.0
        soc_upper_limit = base_soc_upper_limit + (wind_p + solar_p) / base_renewable_shift / 10.0

        # Calculate a new splitting factor for the battery.
        x_battery = clamp(0, (soc_upper_limit - battery_soc)/(soc_upper_limit - soc_lower_limit), 1)

        # If anything is not taken by the battery, put it in the load.
        x_load = 1 - x_battery

        # Publish the new splitting factors.
        splitting_out_socket.send_string(f"{parameters.TOPIC_BATTERY_SPLITTING} {x_battery:.06f}")
        splitting_out_socket.send_string(f"{parameters.TOPIC_LOAD_SPLITTING} {x_load:.06f}")
        logging.info(f"Battery split: {x_battery:.03f}; Load split: {x_load:.03f}")

        # Loop once more in a second
        sleep(1)
finally:
    # Clean up by closing our sockets.
    soc_in_socket.close()
    splitting_out_socket.close()