"""
The pycity_scheduling framework


Copyright (C) 2025,
Institute for Automation of Complex Power Systems (ACS),
E.ON Energy Research Center (E.ON ERC),
RWTH Aachen University

Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated
documentation files (the "Software"), to deal in the Software without restriction, including without limitation the
rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit
persons to whom the Software is furnished to do so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in all copies or substantial portions of the
Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE
WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR
COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR
OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
"""


import numpy as np
import pyomo.environ as pyomo
import pycity_base.classes.city_district as cd

from pycity_scheduling.classes.electrical_entity import ElectricalEntity
from pycity_scheduling.classes.electrical_node import ElectricalNode
from pycity_scheduling.classes.electrical_grid import ElectricalGrid


class CityDistrict(ElectricalEntity, cd.CityDistrict):
    """
    Extension of pyCity_base class CityDistrict for scheduling purposes. Also represents the district operator.

    Parameters
    ----------
    environment : Environment
    objective : str, optional
        Objective for the district operator. Default is 'price'.

        - 'price' : Optimize for the minimum total cost given by `prices.da_prices`.
        - 'peak-shaving' : Try to 'flatten' the schedule as much as possible.
        - 'co2' : Optimize for the minimum total co2 emissions given by `prices.co2_prices`.
        - 'valley-filling' : Try to fill the 'valleys' given by `self.valley_profile`.
        - 'max-consumption' : Try to minimize the maximum power subscription.
        - 'self-consumption' : Try to maximize the self-consumption of the local power generation.
        - 'least-squares': Try to `follow` a given reference power profile.
        - 'flexibility-quantification' : To be used to quantify the flexibility potential of the city district only.
        - 'none' : No objective.
    valley_profile : numpy.ndarray, optional
        Profile to be filled by applying the valley filling objective.
    least_squares_profile : numpy.ndarray, optional
        Profile to be used by the least-squares objective.
    electrical_grid: ElectricalGrid, optional
        Assign an ElectricalGrid object to the city district representing the local power grid

    Notes
    -----
    - The constraints generated by a CD are the same as the ones created by an EntityContainer.
    """

    def __init__(self, environment, objective='price', valley_profile=None, least_squares_profile=None,
                 electrical_grid=None):
        super().__init__(environment)
        self._long_id = "CD_" + self._id_string
        self.objective = objective
        self.valley_profile = valley_profile
        self.least_squares_profile = least_squares_profile
        self.electrical_grid = electrical_grid

    def populate_model(self, model, mode="convex"):
        """
        Add city district block to pyomo ConcreteModel.

        Call parent's `populate_model` methods and set variables lower
        bounds to `None`.

        Parameters
        ----------
        model : pyomo.ConcreteModel
        mode : str, optional
            Specifies which set of constraints to use.

            - `convex`  : Use linear constraints
            - `integer`  : Use same constraints as convex mode
        """
        super().populate_model(model, mode)
        return

    def get_objective(self, coeff=1):
        if self.objective == 'valley-filling':
            obj = 0
            valley = self.valley_profile[self.op_slice]
            for t in range(self.op_horizon):
                obj += coeff * self.model.p_el_vars[t] * self.model.p_el_vars[t]
                obj += 2 * coeff * valley[t] * self.model.p_el_vars[t]
            return obj
        elif self.objective == 'price':
            prices = self.environment.prices.da_prices[self.op_slice]
            s = sum(abs(prices))
            if s > 0:
                obj = 0
                prices = prices * self.op_horizon / s
                for t in range(self.op_horizon):
                    obj += prices[t] * self.model.p_el_vars[t]
                return obj
            else:
                return 0
        else:
            return super().get_objective(coeff)

    def get_decision_var(self):
        if self.objective in ['price', 'valley-filling']:
            return self.model.p_el_vars
        else:
            return super().get_decision_var()

    def get_lower_entities(self):
        for node in self.nodes.values():
            yield node['entity']

    def account_imbalance(self):
        """
        Changes the current schedule to account imbalances.

        The imbalances are determined from the current schedules of the contained entities.
        """

        total = np.zeros_like(self.schedule["p_el"])
        for e in self.get_lower_entities():
            total += e.schedule["p_el"]

        self.schedule["p_el"] = total
        return

    def addElectricalGrid(self, environment, max_p_slack, max_q_slack, slack_V=None, ref_V = 400, V_min=0.95,
                          V_max=1.05):
        """
        Add electrical grid to the city district.

        Create an `ElectricalGrid` object and link it to the city district through the `electrical_grid` attribute.
        Add an electrical node the electrical grid for each entity contained in the city district.
        
        Parameters
        ----------
        environment : Environment
        max_p_slack : float
            Maximum active power that can be exchanged through slack node of the grid in kW.
        max_q_slack : float
            Maximum reactive power that can be exchanged through slack node of the grid in kVar.
        slack_V : float
            If not `None`, voltage at which the slack node should be operated in V. Default is None.
        ref_V : float
            Nominal voltage in the electrical grid in V. Default is 400V.
        V_min : float
            Minimal per unit voltage allowed in the electrical grid. Default is 0.95.
        V_max : float
            Maximal per unit voltage allowed in the electrical grid. Default is 1.05.
        """
        self.electrical_grid = ElectricalGrid(environment, self, max_p_slack, max_q_slack, slack_V, ref_V, V_min, V_max)
        for entity in self.get_lower_entities():
            related_node = ElectricalNode(environment, entity, None)
            self.electrical_grid.add_node(related_node)
