import numpy as np
import yaml
from Utilities.Time import get_start_time
from Utilities.LoadProfile import get_load_profiles
from Utilities.BESS import initialise_SoC, initial_site_assignment, discharge_bess
import matplotlib.pyplot as pl
import pandas as pd
from concurrent.futures import ThreadPoolExecutor

# read config file
c = yaml.safe_load(open("YAMLs/config.yml"))

# read BESS data
bess_data = yaml.safe_load(open(c["paths"]["bess"]))

## simulation time setup
# get current time
c["sim_start_time"] = get_start_time()
# get time step in minutes, then convert to seconds
dt = c["sim_time"]["time_step_minutes"] * 60
# compute end time based on duration in days
duration = c["sim_time"]["duration_days"] * 24 * 60 * 60
c["sim_end_time"] = c["sim_start_time"] + duration
timestamps = np.arange(c["sim_start_time"], c["sim_end_time"] + 1, dt)

# batch process hours in seconds
c["sim_time"]["batch_process_seconds"] = c["sim_time"]["batch_process_hours"] * 60 * 60

# load site info
c["site_info"] = yaml.safe_load(open(c["paths"]["site_info"]))


def generate_and_cache_profiles(c, dt):
    """Generates load profiles for all sites and caches them."""
    return get_load_profiles(
        c, dt, c["sim_start_time"], c["sim_time"]["batch_process_seconds"]
    )


# initialise BESS
bess_data = initialise_SoC(bess_data)
bess_data = initial_site_assignment(c, bess_data)
# bess SoC dataframe
bess_soc = pd.DataFrame(columns=[unit["name"] for unit in bess_data["units"]])

# get initial load profiles
cumulative_load_profiles = get_load_profiles(
    c, dt, c["sim_start_time"], c["sim_time"]["batch_process_seconds"]
)

# async function is running
is_running_in_async = False

# loop through
with ThreadPoolExecutor() as executor:
    for i in range(0, len(timestamps)):
        # start generating load profiles 200 seconds before data required
        if len(cumulative_load_profiles) <= len(timestamps):
            if is_running_in_async is False:
                # generate load profiles
                future = executor.submit(generate_and_cache_profiles, c, dt)
                is_running_in_async = True
        else:
            is_running_in_async = False

        # discharge BESS for each site
        for site in c["site_info"]["sites"]:
            site_name = site["name"]
            discharge_power = cumulative_load_profiles[site_name].iloc[i]
            bess_data = discharge_bess(bess_data, site_name, dt, discharge_power)
        temp_soc = [unit["SoC"] for unit in bess_data["units"]]

        # append SoC to dataframe
        bess_soc = pd.concat(
            [
                bess_soc,
                pd.DataFrame(
                    [temp_soc], columns=bess_soc.columns, index=[timestamps[i]]
                ),
            ],
            axis=0,
        )

        # add to cumulative load profiles
        # check if future exists and is done
        if is_running_in_async:
            if future.done():
                load_profiles = future.result()
                cumulative_load_profiles = pd.concat(
                    [
                        cumulative_load_profiles,
                        load_profiles,
                    ],
                    axis=0,
                )
                print(len(cumulative_load_profiles), "load profiles generated")
                is_running_in_async = False

pl.plot(bess_soc.index, bess_soc.values, label="BESS SoC", alpha=0.5)
pl.show()
pl.xlabel("Time (s since epoch)")