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,
    predict_swap_time,
    update_cycle_SoC,
)
from Utilities.DataVis import format_dataframe
import matplotlib.pyplot as pl
import pandas as pd
from concurrent.futures import ThreadPoolExecutor

import threading  ### <<< CONTROL ADDED >>>
import time  ### <<< CONTROL ADDED >>>

# 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
c["sim_start_time"] = get_start_time()
dt = c["sim_time"]["time_step_minutes"] * 60
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)
c["sim_time"]["batch_process_seconds"] = c["sim_time"]["batch_process_hours"] * 60 * 60
load_profiles_since_start = None
status_df = None

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


### <<< CONTROL ADDED >>> Initialize simulation state globals
sim_i = 0
running = False
is_running_in_async = False
sim_lock = threading.Lock()


# initialise BESS
def _init_state():
    global bess_data, bess_soc_since_start, bess_soc_for_cycle, cumulative_load_profiles, load_profiles_since_start

    bd = initialise_SoC(bess_data.copy())
    bd = initial_site_assignment(c, bd)
    bess_data = bd

    bess_soc_since_start = pd.DataFrame(
        columns=[unit["name"] for unit in bess_data["units"]]
    )
    init_df = pd.DataFrame(columns=["Timestamp", "SoC"])
    bess_soc_for_cycle = {unit["name"]: init_df for unit in bess_data["units"]}

    cumulative_load_profiles = get_load_profiles(
        c, dt, c["sim_start_time"], c["sim_time"]["batch_process_seconds"]
    )


# do initial setup
_init_state()


def simulation_loop():
    """Runs the loop, stepping through timestamps until stopped or finished."""
    global sim_i, running, is_running_in_async, cumulative_load_profiles, bess_data, bess_soc_for_cycle, load_profiles_since_start, status_df
    with ThreadPoolExecutor() as executor:
        while True:
            with sim_lock:
                if not running or sim_i >= len(timestamps):
                    break

                i = sim_i
                sim_i += 1

            # pre-fetch next batch if needed
            if len(cumulative_load_profiles) <= len(timestamps):
                if not is_running_in_async:
                    future = executor.submit(
                        get_load_profiles,
                        c,
                        dt,
                        c["sim_start_time"],
                        c["sim_time"]["batch_process_seconds"],
                    )
                    is_running_in_async = True
            else:
                is_running_in_async = False

            # discharge BESS for each site
            for site in c["site_info"]["sites"]:
                name = site["name"]
                p = cumulative_load_profiles[name].iloc[i]
                bess_data = discharge_bess(bess_data, name, dt, p)

            # record SoC
            temp_soc = [u["SoC"] for u in bess_data["units"]]
            bess_soc_since_start.loc[timestamps[i]] = temp_soc

            # update cycle SoC and predict swaps
            bess_soc_for_cycle = update_cycle_SoC(
                bess_data, bess_soc_for_cycle, timestamps[i]
            )
            swap_times = predict_swap_time(bess_soc_for_cycle)

            # integrate newly fetched profiles
            if is_running_in_async and future.done():
                load_profiles = future.result()
                cumulative_load_profiles = pd.concat(
                    [cumulative_load_profiles, load_profiles], axis=0
                )
                print(len(cumulative_load_profiles), "profiles generated")
                is_running_in_async = False

            load_profiles_since_start = cumulative_load_profiles.iloc[: i + 1]

            # format data for display
            status_df = format_dataframe(
                bess_soc_for_cycle, bess_data, load_profiles_since_start, swap_times
            )

            # small sleep to allow dashboard to refresh / release GIL
            time.sleep(0.01)


### <<< CONTROL ADDED >>> Control functions
def start_sim():
    """Starts the simulation in a background thread."""
    global running, sim_thread
    if not running:
        running = True
        sim_thread = threading.Thread(target=simulation_loop, daemon=True)
        sim_thread.start()


def stop_sim():
    """Stops the simulation loop."""
    global running
    running = False


def reset_sim():
    """Stops and re-initializes the simulation state."""
    global running, sim_i
    running = False
    sim_i = 0
    _init_state()