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

# # Optimization without pyomo
# 
# In this example we demonstrate the behaviour of the Linear Optimal Power Flow (LOPF) calculation without using pyomo. This requires to set `pyomo` to `False` in the `lopf` function. Then, the communication with the solvers happens via in-house functions which leads to a much faster solving process. 

# In[ ]:


import pypsa
import pandas as pd
import os


# In[ ]:


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


# Modify the network a bit: We set gas generators to non-extendable

# In[ ]:


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


# Add ramp limit

# In[ ]:


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


# Add additional storage units (cyclic and non-cyclic) and fix one state_of_charge

# In[ ]:


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


# Add an additional store

# In[ ]:


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,
);


# ## Extra functionalities:

# In[ ]:


from pypsa.linopt import get_var, linexpr, join_exprs, define_constraints


# One of the most important functions is linexpr which take one or more tuples of coefficient and variable pairs which should go into the left hand side (lhs) of the constraint. 
# 
# 1. Add mimimum for state_of_charge

# 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")


# 2. Fix the ratio between ingoing and outgoing capacity of the Store

# 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["battery_power"]), (-eff, vars_link["battery_discharge"])
    )
    define_constraints(n, lhs, "=", 0, "battery_discharge", attr="fixratio")


# 3. Every bus must in total produce the 20% of the total demand
# 
# This requires the function `pypsa.linopt.join_exprs` which sums up arrays of linear expressions

# 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"
    )


# Combine them ...

# In[ ]:


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


# ...and run the lopf with `pyomo=False`

# In[ ]:


n.lopf(
    pyomo=False,
    extra_functionality=extra_functionalities,
    keep_shadowprices=["Bus", "battery_discharge", "StorageUnit"],
)


# The `keep_shadowprices` argument in the lopf now decides which shadow prices (SP) should be retrieved. It can either be set to `True`, then all SP are kept. It also can be a list of names of the constraints. Therefore the `name` argument in `define_constraints` is necessary, in our case 'battery_discharge', 'StorageUnit' and 'Bus'. 

# ## Analysing the constraints

# Let's see if the system got our own constraints. We look at `n.constraints` which combines summarises constraints going into the linear problem

# In[ ]:


n.constraints


# The last three entries show our constraints. As 'soc_lower_bound' is time-dependent, the `pnl` value is set to `True`. 
# 
# Let's check whether out two custom constraint are fulfilled:

# In[ ]:


n.links.loc[["battery_power", "battery_discharge"], ["p_nom_opt"]]


# In[ ]:


n.storage_units_t.state_of_charge


# In[ ]:


n.generators_t.p.groupby(n.generators.bus, axis=1).sum().sum() / n.loads_t.p.sum().sum()


# Looks good! Now, let's see which dual values were parsed. Therefore we have a look into `n.dualvalues` 
# 

# In[ ]:


n.dualvalues


# Again we see the last two entries reflect our constraints (the values in the columns play only a minor role). Having a look what the values are:

# In[ ]:


from pypsa.linopt import get_dual


# In[ ]:


get_dual(n, "StorageUnit", "soc_lower_bound")


# In[ ]:


get_dual(n, "battery_discharge", "fixratio")


# In[ ]:


get_dual(n, "Bus", "production_share")


# ### Side note
# Some of the predefined constraints are stored in components itself like `n.lines_t.mu_upper` or `n.buses_t.marginal_price`, this is the case if their are designated columns are spots for those. All other dual are under the hook stored in `n.duals`