# title: Fuel distribution
#
# description: |
#     Here, we add in the ability to track the distribution of commodities in the system that do not travel along distinct networks.
#     This allows the commodities to imported/exported to/from nodes without tracking exactly where they have come from / where they are going.
#     The advantage is a simpler model definition; the disadvantage is not being able to track the source of a commodity that has been imported.
#     We refer to "fuels" instead of "commodities" in the below math, but this could refer equally to other commodities that have a corresponding carrier (waste, water, ...).
#
#     New indexed parameters:
#
#     - `fuel_import_max`
#     - `fuel_export_max`
#     - `cost_fuel_distribution`
#     - `allow_fuel_distribution` <- lookup array with a value of `True` for each carrier where you want to track its distribution
#
#     Helper functions used:
#
#     - `any` (where)
#     - `sum` (expression)
#
# ---

parameters:
  fuel_import_max:
    description: The maximum amount of fuel that can be imported
    default: .inf
    unit: energy
  fuel_export_max:
    description: The maximum amount of fuel that can be exported
    default: .inf
    unit: energy
  cost_fuel_distribution:
    description: The cost to import / revenue to export fuel
    default: 0
    unit: energy

lookups:
  allow_fuel_distribution:
    description: Lookup table of nodes / carriers that are allowed to participate in fuel distribution.
    dtype: bool

variables:
  fuel_distributor:
    description: >
      Fuel distributor, allowing copperplate transfer of specific carriers
      between model nodes. Positive values indicate carrier imports at a node,
      negative values for carrier exports.
    foreach: [nodes, carriers, timesteps]
    unit: energy
    default: 0
    # change [fuel1, fuel2] to match the carriers you want to allow to
    # be distributed within the system.
    where: "allow_fuel_distribution"
    bounds:
      min: -.inf
      max: .inf

constraints:
  # Add variable to existing system balance constraint
  system_balance:
    equations:
      - expression: >
          sum(flow_out, over=techs) - sum(flow_in, over=techs) - $flow_export
          + $unmet_demand_and_unused_supply + $fuel_distributor == 0
    sub_expressions:
      fuel_distributor:
        - where: fuel_distributor
          expression: fuel_distributor
        - where: NOT fuel_distributor
          expression: "0"

  restrict_total_imports_and_exports:
    description: >
      Ensure all fuel distribution in the system balances to zero.
      I.e., all regional fuel exports must equal regional fuel imports.
      Setting this constraint to an inequality would allow for net
      imports/exports from/to the system.
    foreach: [carriers, timesteps]
    where: fuel_distributor
    equations:
      - expression: sum(fuel_distributor, over=nodes) == 0

  # fuel_import_max and fuel_export_max could be defined per node and carrier and
  # could be collapsed into one constraining limit e.g., `fuel_distribution_max`.
  # If using one constraining limit, two constraints will still be required.
  restrict_nodal_imports:
    description: >
      Ensure all fuel distribution in the system balances to zero.
      I.e., all regional fuel exports must equal regional fuel imports.
      Setting this constraint to an inequality would allow for net
      imports/exports from/to the system.
    foreach: [nodes, carriers, timesteps]
    where: fuel_distributor AND fuel_import_max
    equations:
      - expression: fuel_distributor <= fuel_import_max

  restrict_nodal_exports:
    description: >
      Ensure all fuel distribution in the system balances to zero.
      I.e., all regional fuel exports must equal regional fuel imports.
      Setting this constraint to an inequality would allow for net
      imports/exports from/to the system.
    foreach: [nodes, carriers, timesteps]
    where: fuel_distributor AND fuel_export_max
    equations:
      - expression: -1 * fuel_distributor <= fuel_export_max

global_expressions:
  # To apply a different cost for imports vs exports, you will need to separate
  # out the fuel_distributor decision variable into two positive decision
  # variables representing imports and exports.
  # You will then need to propagate that change throughout all
  # constraints/global expressions/objective function given in this file.
  # This change will increase model complexity and if different costs do exist
  # you risk having "unrealistic" simultaneous imports/exports at a node to
  # accrue revenue.
  cost_var_fuel_distribution:
    description: >
      Cost of importing / exporting fuel. Exporting fuel will
      provide a node with revenue, while importing it will incur a cost.
    unit: cost
    default: 0
    foreach: [nodes, carriers, costs, timesteps]
    where: fuel_distributor AND cost_fuel_distribution
    equations:
      - expression: timestep_weights * fuel_distributor * cost_fuel_distribution

objectives:
  # Update objective to include fuel distribution costs
  # NOTE: these additional costs will have no impact on the objective function
  # value _unless_ the cost of fuel distribution is different per node OR a
  # systemwide imbalance in fuel distribution (inequality in
  # `restrict_total_imports_and_exports`) is enabled.
  min_cost_optimisation:
    equations:
      - expression: $cost_sum + $unmet_demand + $cost_fuel_distribution_sum
    sub_expressions:
      cost_sum:
        - where: "any(cost, over=[nodes, techs, costs])"
          expression: >
            sum(sum(cost, over=[nodes, techs]) * objective_cost_weights, over=costs)
        - where: "NOT any(cost, over=[nodes, techs, costs])"
          expression: "0"
      unmet_demand:
        - where: "config.ensure_feasibility==True"
          expression: >
            sum(sum(unmet_demand - unused_supply, over=[carriers, nodes])
            * timestep_weights, over=timesteps) * bigM
        - where: "NOT config.ensure_feasibility==True"
          expression: "0"
      cost_fuel_distribution_sum:
        - where: "cost_var_fuel_distribution"
          expression: >
            sum(sum(cost_var_fuel_distribution, over=[nodes, carriers, timesteps])
            * objective_cost_weights, over=costs)
        - where: >
            NOT any(cost_var_fuel_distribution,
            over=[nodes, carriers, costs, timesteps])
          expression: "0"
