import pandas as pd
from oemof import solph


def multi_period(wind_ts):
    t_idx_1 = pd.date_range("1/1/2020", periods=3, freq="H")
    t_idx_2 = pd.date_range("1/1/2030", periods=3, freq="H")
    t_idx_1_series = pd.Series(index=t_idx_1, dtype="float64")
    t_idx_2_series = pd.Series(index=t_idx_2, dtype="float64")
    date_time_index = pd.concat([t_idx_1_series, t_idx_2_series]).index
    periods = [t_idx_1, t_idx_2]

    energysystem = solph.EnergySystem(
        timeindex=date_time_index,
        infer_last_interval=False,
        timeincrement=[1] * len(date_time_index),
        periods=periods,
    )

    el_bus = solph.buses.Bus(label="el-bus")
    energysystem.add(el_bus)

    volatile = solph.components.Source(
        label="wind",
        outputs={
            el_bus: solph.flows.Flow(
                fix=len(periods) * wind_ts,
                investment=solph.Investment(ep_costs=1, lifetime=1),
            )
        },
    )

    load = solph.components.Sink(
        label="load",
        inputs={
            el_bus: solph.flows.Flow(fix=len(periods) * [0, 1, 1], nominal_value=100)
        },
    )

    excess = solph.components.Sink(
        label="excess",
        inputs={el_bus: solph.flows.Flow(variable_costs=1e12)},
    )

    shortage = solph.components.Source(
        label="shortage", outputs={el_bus: solph.flows.Flow(variable_costs=1e12)}
    )

    storage = solph.components.GenericStorage(
        label="storage",
        inputs={el_bus: solph.Flow(variable_costs=1)},
        outputs={el_bus: solph.Flow()},
        loss_rate=0.00,
        # initial_storage_level=0,
        invest_relation_input_capacity=1 / 6,
        invest_relation_output_capacity=1 / 6,
        inflow_conversion_factor=1,
        outflow_conversion_factor=0.8,
        investment=solph.Investment(ep_costs=1, lifetime=1),
        lifetime_inflow=1,
        lifetime_outflow=1,
    )

    energysystem.add(
        volatile,
        load,
        excess,
        shortage,
        storage,
    )

    om = solph.Model(energysystem)

    om.solve(solver="cbc", solve_kwargs={"tee": True})

    results = solph.processing.results(om)

    rename_mapping = {
        oemof_tuple: f"{oemof_tuple[0]}->{oemof_tuple[1]}"
        for oemof_tuple in results.keys()
    }
    for old_key, new_key in rename_mapping.items():
        results[new_key] = results.pop(old_key)

    return results


if __name__ == "__main__":
    # wind production in first time step
    results_1 = multi_period(wind_ts=[1, 0, 0])
    # wind production in second time step
    results_2 = multi_period(wind_ts=[0, 1, 0])

    print("wind time series [1, 0, 0]")
    print("wind")
    print(results_1["wind->el-bus"]["period_scalars"])
    print("\n storage")
    print(results_1["storage->None"]["period_scalars"])
    print("\n shortage total usage: ")
    print(results_1["shortage->el-bus"]["sequences"].flow.sum())

    print("\n \n wind time series [0, 1, 0]")
    print("wind")
    print(results_2["wind->el-bus"]["period_scalars"])
    print("\n storage")
    print(results_2["storage->None"]["period_scalars"])
    print("\n shortage total usage: ")
    print(results_2["shortage->el-bus"]["sequences"].flow.sum())