#!/usr/bin/env python
# coding: utf-8

# # Optimization with Linopy - Migrate Extra Functionalities
# 
# The extra funcionalities for the native pypsa optimization code are mostly using the function of the `pypsa.linopt` model. Here we show how you can recycle large parts of your code by using the compatibility functions from the `pypsa.optimization.compat` module.
# 
# These are 
# 
# * `define_variables`
# * `define_constraints`
# * `get_var`
# * `linexpr`
# 
# You might want to use them if you have `extra_functionalities` written for the native optimization code. However, expect some hick-ups, as some operations might behave differently. 
# 
# Let's import pypsa and the compat functions:

# In[ ]:


import pypsa
from pypsa.optimization.compat import get_var, define_constraints, linexpr


# We load the same network from the previous section into memory:

# In[ ]:


n = pypsa.examples.ac_dc_meshed(from_master=True)

n.generators.loc[n.generators.carrier == "gas", "p_nom_extendable"] = False
n.generators.loc[n.generators.carrier == "gas", "ramp_limit_down"] = 0.2
n.generators.loc[n.generators.carrier == "gas", "ramp_limit_up"] = 0.2

n.add(
    "StorageUnit",
    "su",
    bus="Manchester",
    marginal_cost=10,
    inflow=50,
    p_nom_extendable=True,
    capital_cost=10,
    p_nom=2000,
    efficiency_dispatch=0.5,
    cyclic_state_of_charge=True,
    state_of_charge_initial=1000,
)

n.add(
    "StorageUnit",
    "su2",
    bus="Manchester",
    marginal_cost=10,
    p_nom_extendable=True,
    capital_cost=50,
    p_nom=2000,
    efficiency_dispatch=0.5,
    carrier="gas",
    cyclic_state_of_charge=False,
    state_of_charge_initial=1000,
)

n.storage_units_t.state_of_charge_set.loc[n.snapshots[7], "su"] = 100

n.add("Bus", "storebus", carrier="hydro", x=-5, y=55)
n.madd(
    "Link",
    ["battery_power", "battery_discharge"],
    "",
    bus0=["Manchester", "storebus"],
    bus1=["storebus", "Manchester"],
    p_nom=100,
    efficiency=0.9,
    p_nom_extendable=True,
    p_nom_max=1000,
)
n.madd(
    "Store",
    ["store"],
    bus="storebus",
    e_nom=2000,
    e_nom_extendable=True,
    marginal_cost=10,
    capital_cost=10,
    e_nom_max=5000,
    e_initial=100,
    e_cyclic=True,
);


# And define the extra functionalities as we defined them for the native code in [here](https://pypsa.readthedocs.io/en/latest/examples/lopf_with_pyomo_False.html)

# In[ ]:


def minimal_state_of_charge(n, snapshots):
    vars_soc = get_var(n, "StorageUnit", "state_of_charge")
    lhs = linexpr((1, vars_soc))
    define_constraints(n, lhs, ">", 50, "StorageUnit", "soc_lower_bound")


# With the compat functions, this will work as expected. Let's go on to the next one.

# In[ ]:


def fix_link_cap_ratio(n, snapshots):
    vars_link = get_var(n, "Link", "p_nom")
    eff = n.links.at["battery_power", "efficiency"]
    lhs = linexpr(
        (1, vars_link.loc["battery_power"]), (-eff, vars_link.loc["battery_discharge"])
    )
    define_constraints(n, lhs, "=", 0, "battery_discharge", attr="fixratio")


# This function as well should not make any problems. Let's go on.

# In[ ]:


def fix_bus_production(n, snapshots):
    total_demand = n.loads_t.p_set.sum().sum()
    prod_per_bus = (
        linexpr((1, get_var(n, "Generator", "p")))
        .groupby(n.generators.bus, axis=1)
        .apply(join_exprs)
    )
    define_constraints(
        n, prod_per_bus, ">=", total_demand / 5, "Bus", "production_share"
    )


# Here, we come into difficult terrain. The groupby function won't work since the `linexpr` functions returns some sort of a xarray object (a `LinearExpression` object, derived from `xarray.Dataset`). 
# 
# Instead, we have to rewrite parts: 
# 
# * remove the `axis` argument
# * use a xaray DataArray as argument
# * explicitly sum over all snapshots afterwards. This has nothing to do with groupby but with the fact that we want to limit to total production.
# 

# In[ ]:


def fix_bus_production(n, snapshots):
    total_demand = n.loads_t.p_set.sum().sum()
    prod_per_bus = (
        linexpr((1, get_var(n, "Generator", "p")))
        .groupby(n.generators.bus.to_xarray())
        .sum()
        .sum()
    )
    define_constraints(
        n, prod_per_bus, ">=", total_demand / 5, "Bus", "production_share"
    )


# In[ ]:


def extra_functionalities(n, snapshots):
    minimal_state_of_charge(n, snapshots)
    fix_link_cap_ratio(n, snapshots)
    fix_bus_production(n, snapshots)


# In[ ]:


n.optimize(
    extra_functionality=extra_functionalities,
)