"""
Copyright (C) 2013-2019 Calliope contributors listed in AUTHORS.
Licensed under the Apache 2.0 License (see LICENSE file).

conversion.py
~~~~~~~~~~~~~~~~~

Conversion technology constraints.

"""

import pyomo.core as po  # pylint: disable=import-error

from calliope.backend.pyomo.util import \
    get_param

ORDER = 20  # order in which to invoke constraints relative to other constraint files


def load_constraints(backend_model):
    sets = backend_model.__calliope_model_data['sets']

    if 'loc_techs_balance_conversion_constraint' in sets:
        backend_model.balance_conversion_constraint = po.Constraint(
            backend_model.loc_techs_balance_conversion_constraint, backend_model.timesteps,
            rule=balance_conversion_constraint_rule
        )

    if 'loc_techs_cost_var_conversion_constraint' in sets:
        backend_model.cost_var_conversion_constraint = po.Constraint(
            backend_model.costs, backend_model.loc_techs_cost_var_conversion_constraint,
            backend_model.timesteps,
            rule=cost_var_conversion_constraint_rule
        )


def balance_conversion_constraint_rule(backend_model, loc_tech, timestep):
    """
    Balance energy carrier consumption and production

    .. container:: scrolling-wrapper

        .. math::

            -1 * \\boldsymbol{carrier_{con}}(loc::tech::carrier, timestep)
            \\times \\eta_{energy}(loc::tech, timestep)
            = \\boldsymbol{carrier_{prod}}(loc::tech::carrier, timestep)
            \\times \\eta_{energy}(loc::tech, timestep)
            \\quad \\forall loc::tech \\in locs::techs_{conversion},
            \\forall timestep \\in timesteps
    """
    model_data_dict = backend_model.__calliope_model_data['data']

    loc_tech_carrier_out = model_data_dict['lookup_loc_techs_conversion'][('out', loc_tech)]
    loc_tech_carrier_in = model_data_dict['lookup_loc_techs_conversion'][('in', loc_tech)]

    energy_eff = get_param(backend_model, 'energy_eff', (loc_tech, timestep))

    return (
        backend_model.carrier_prod[loc_tech_carrier_out, timestep] == -1 *
        backend_model.carrier_con[loc_tech_carrier_in, timestep] * energy_eff
    )


def cost_var_conversion_constraint_rule(backend_model, cost, loc_tech, timestep):
    """
    Add time-varying conversion technology costs

    .. container:: scrolling-wrapper

        .. math::

            \\boldsymbol{cost_{var}}(loc::tech, cost, timestep) =
            \\boldsymbol{carrier_{prod}}(loc::tech::carrier, timestep)
            \\times timestep_{weight}(timestep) \\times cost_{om, prod}(loc::tech, cost, timestep)
            +
            \\boldsymbol{carrier_{con}}(loc::tech::carrier, timestep)
            \\times timestep_{weight}(timestep) \\times cost_{om, con}(loc::tech, cost, timestep)
            \\quad \\forall loc::tech \\in loc::techs_{cost_{var}, conversion}
    """
    model_data_dict = backend_model.__calliope_model_data
    weight = backend_model.timestep_weights[timestep]

    loc_tech_carrier_in = (
        model_data_dict['data']['lookup_loc_techs_conversion'][('in', loc_tech)]
    )

    loc_tech_carrier_out = (
        model_data_dict['data']['lookup_loc_techs_conversion'][('out', loc_tech)]
    )

    cost_om_prod = get_param(backend_model, 'cost_om_prod',
                                (cost, loc_tech, timestep))
    cost_om_con = get_param(backend_model, 'cost_om_con',
                            (cost, loc_tech, timestep))
    if po.value(cost_om_prod):
        cost_prod = (cost_om_prod * weight *
            backend_model.carrier_prod[loc_tech_carrier_out, timestep])
    else:
        cost_prod = 0

    if po.value(cost_om_con):
        cost_con = (cost_om_con * weight * -1 *
            backend_model.carrier_con[loc_tech_carrier_in, timestep])
    else:
        cost_con = 0

    backend_model.cost_var_rhs[cost, loc_tech, timestep] = cost_prod + cost_con

    return (backend_model.cost_var[cost, loc_tech, timestep] ==
            backend_model.cost_var_rhs[cost, loc_tech, timestep])
