import pypsa
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns

sns.set(rc={"figure.figsize": (9, 5)})

network = pypsa.Network()

"""
Explanation: LOPF with coupling to heating sector
In this example three locations are optimised, each with an electric bus and a heating bus and corresponding loads. At each location the electric and heating buses are connected with heat pumps; heat can also be supplied to the heat bus with a boiler. The electric buses are connected with transmission lines and there are electrical generators at two of the nodes.
End of explanation
"""


for i in range(3):
    network.add("Bus", "electric bus {}".format(i), v_nom=20.0)
    network.add("Bus", "heat bus {}".format(i), carrier="heat")
network.buses

network.buses["carrier"].value_counts()

"""
Explanation: Add three buses of AC and heat carrier each
End of explanation
"""


for i in range(3):
    network.add(
        "Line",
        "line {}".format(i),
        bus0="electric bus {}".format(i),
        bus1="electric bus {}".format((i + 1) % 3),
        x=0.1,
        s_nom=1000,
    )
network.lines

"""
Explanation: Add three lines in a ring
End of explanation
"""


for i in range(3):
    network.add(
        "Link",
        "heat pump {}".format(i),
        bus0="electric bus {}".format(i),
        bus1="heat bus {}".format(i),
        p_nom=100,
        efficiency=3.0,
    )
network.links

"""
Explanation: Connect the electric to the heat buses with heat pumps with COP 3
End of explanation
"""


network.add("Carrier", "gas", co2_emissions=0.27)
network.add("Carrier", "biomass", co2_emissions=0.0)
network.carriers

"""
Explanation: Add carriers
End of explanation
"""


network.add(
    "Generator",
    "gas generator",
    bus="electric bus 0",
    p_nom=100,
    marginal_cost=50,
    carrier="gas",
    efficiency=0.3,
)

network.add(
    "Generator",
    "biomass generator",
    bus="electric bus 1",
    p_nom=100,
    marginal_cost=100,
    efficiency=0.3,
    carrier="biomass",
)

for i in range(3):
    network.add(
        "Generator",
        "boiler {}".format(i),
        bus="heat bus {}".format(i),
        p_nom=1000,
        efficiency=0.9,
        marginal_cost=20.0,
        carrier="gas",
    )

network.generators

"""
Explanation: Add a gas generator at bus 0, a biomass generator at bus 1 and a boiler at all heat buses
End of explanation
"""


for i in range(3):
    network.add(
        "Load",
        "electric load {}".format(i),
        bus="electric bus {}".format(i),
        p_set=i * 10,
    )

for i in range(3):
    network.add(
        "Load",
        "heat load {}".format(i),
        bus="heat bus {}".format(i),
        p_set=(3 - i) * 10,
    )

network.loads

"""
Explanation: Add electric loads and heat loads.
End of explanation
"""


def run_lopf():
    network.lopf()
    df = pd.concat(
        [
            network.generators_t.p.loc["now"],
            network.links_t.p0.loc["now"],
            network.loads_t.p.loc["now"],
        ],
        keys=["Generators", "Links", "Line"],
        names=["Component", "index"],
    ).reset_index(name="Production")

    sns.barplot(data=df, x="index", y="Production", hue="Component")
    plt.title(f"Objective: {network.objective}")
    plt.xticks(rotation=90)
    plt.tight_layout()

run_lopf()

"""
Explanation: We define a function for the LOPF
End of explanation
"""


network.links.marginal_cost = 10
run_lopf()

"""
Explanation: Now, rerun with marginal costs for the heat pump operation.
End of explanation
"""


network.add("GlobalConstraint", "co2_limit", sense="<=", constant=0.0)

run_lopf()

"""
Explanation: Finally, rerun with no CO2 emissions.
End of explanation
"""


