import pandas as pd


def initialise_SoC(bess):
    """Initialise the state of charge (SoC) for the BESS."""
    for i in range(0, len(bess["units"])):  # initially fully charged
        bess["units"][i]["SoC"] = 1
    return bess


def initial_site_assignment(c, bess):
    """Initialise the site assignment for each BESS."""

    for k in range(0, len(bess["units"])):
        # assign each BESS unit to a site
        if k < len(c["site_info"]["sites"]):
            bess["units"][k]["site"] = c["site_info"]["sites"][k]["name"]
        else:
            bess["units"][k]["site"] = "Unassigned"

    return bess


def discharge_bess(bess, site_name, dt, discharge_power):
    # convert discharge power to discharge energy (kW to kWh)
    discharge_energy = discharge_power * dt / 3600

    """Discharge the BESS for a specific site."""
    for index, unit in enumerate(bess["units"]):
        if unit["site"] == "Unassigned":
            continue

        if unit["site"] == site_name:
            new_soc = unit["SoC"] - (dt * discharge_energy) / unit["capacity_kWh"]
            new_soc = 0 if new_soc < 0 else new_soc
        else:
            # maintain SoC if not assigned to the site
            new_soc = unit["SoC"]

        # update SoC
        bess["units"][index]["SoC"] = new_soc
    return bess


def predict_swap_time(bess_soc_for_cycle):
    """Predict the swap time for each BESS unit based on its SoC history."""
    swap_times = {}
    min2sec = 60
    threshold = 2 * min2sec  # 2 minutes in seconds

    for unit_name, df in bess_soc_for_cycle.items():
        # need to be at least 1 min of operation to start estimation
        if len(df) < threshold:
            swap_times[unit_name] = None
            continue

        # linear extrapolation to estimate swap time
        # calculate the slope of the SoC over time
        m = (df["SoC"].iloc[-1] - df["SoC"].iloc[0]) / (
            df["Timestamp"].iloc[-1] - df["Timestamp"].iloc[0]
        )

        if m == 0:
            swap_times[unit_name] = None
            continue

        # solve for the time when SoC reaches 0
        swap_time = (0 - df["SoC"].iloc[0]) / m + df["Timestamp"].iloc[0]
        # assign to swap_times
        swap_times[unit_name] = swap_time

    return swap_times


def update_cycle_SoC(bess_data, bess_soc_for_cycle, timestamps):
    init_df = pd.DataFrame(columns=["Timestamp", "SoC"])
    # assign SoC for cycle
    for unit in bess_data["units"]:
        unit_name = unit["name"]
        # reset df if SoC is 0. Start a new cycle
        if unit["SoC"] == 0:
            bess_soc_for_cycle[unit_name] = init_df

        bess_soc_for_cycle[unit_name] = pd.concat(
            [
                bess_soc_for_cycle[unit_name],
                pd.DataFrame(
                    [[timestamps[i], unit["SoC"]]],
                    columns=["Timestamp", "SoC"],
                ),
            ],
            axis=0,
        )

def arrange_swap(bess_data, c):
    for unit in bess_data["units"]:
        if unit["SoC"] < c["bess"]["buffer"]:   
            # find for unassigned BESS unit with SOC at 100%