import pandas as pd
from oemof import solph

date_time_index = pd.date_range("1/1/2020", periods=4, freq="H")


energysystem = solph.EnergySystem(
    timeindex=date_time_index,
    infer_last_interval=False,
)

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=[1, 0, 0],
            investment=solph.Investment(ep_costs=1, lifetime=1),
        )
    },
)

load = solph.components.Sink(
    label="load",
    inputs={el_bus: solph.flows.Flow(fix=[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)

print(results["storage->None"])