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

preprocess_sets.py
~~~~~~~~~~~~~~~~~~

Sets & sub-sets defined to reduce size of decision variables & constraints.

The markers preceded with `###` are used to auto-include the set description
in the documentation

###PART TO INCLUDE IN DOCUMENTATION STARTS HERE###

Main sets
=========

* timesteps
* techs
* techs_non_transmission
* techs_transmission
* locs
* costs
* resources
* carriers
* carrier_tiers
* techlists

Location-technology subsets
===========================

* loc_techs

Technology groups

* loc_techs_storage
* loc_techs_transmission
* loc_techs_demand
* loc_techs_supply
* loc_techs_supply_plus
* loc_techs_conversion
* loc_techs_conversion_plus

Subsets based on active constraints

* loc_techs_area
* loc_techs_store
* loc_techs_finite_resource
* loc_techs_finite_resource_supply
* loc_techs_finite_resource_demand
* loc_techs_finite_resource_supply_plus
* loc_techs_ramping
* loc_techs_export
* loc_techs_purchase
* loc_techs_milp
* loc_techs_cost
* loc_techs_investment_cost
* loc_techs_om_cost
* loc_techs_out_2
* loc_techs_out_3
* loc_techs_in_2
* loc_techs_in_3


Subsets that include carrier

* loc_tech_carriers_prod
* loc_tech_carriers_con
* loc_tech_carriers_conversion_plus
* loc_tech_carriers_export
* loc_carriers

###PART TO INCLUDE IN DOCUMENTATION ENDS HERE###

"""

from itertools import product

import numpy as np

from calliope.core.attrdict import AttrDict
from calliope.core.preprocess.util import \
    get_all_carriers, split_loc_techs_transmission, \
    concat_iterable, flatten_list


def generate_simple_sets(model_run):
    """
    Generate basic sets for a given pre-processed ``model_run``.

    Parameters
    ----------
    model_run : AttrDict

    """
    sets = AttrDict()

    flat_techs = model_run.techs.as_dict(flat=True)
    flat_locations = model_run.locations.as_dict(flat=True)

    sets.resources = set(flatten_list(
        v for k, v in flat_techs.items()
        if '.carrier' in k
    ))

    sets.carriers = sets.resources - set(['resource'])

    sets.carrier_tiers = set(
        key.split('.carrier_')[1]
        for key in flat_techs.keys()
        if '.carrier_' in key
    )

    sets.costs = set(
        k.split('costs.')[-1].split('.')[0]
        for k in flat_locations.keys()
        if '.costs.' in k)

    sets.locs = set(model_run.locations.keys())

    sets.techs_non_transmission = set(
        k for k, v in model_run.techs.items()
        if v.inheritance[-1] != 'transmission')

    sets.techs_transmission_names = set(
        k for k, v in model_run.techs.items()
        if v.inheritance[-1] == 'transmission')

    # This builds the "tech:loc" expansion of transmission technologies
    techs_transmission = set()
    for loc_name, loc_config in model_run.locations.items():
        for link_name, link_config in loc_config.get('links', {}).items():
            for tech_name in link_config.techs:
                techs_transmission.add('{}:{}'.format(tech_name, link_name))
    sets.techs_transmission = techs_transmission

    sets.techs = sets.techs_non_transmission | sets.techs_transmission_names

    # this extracts location coordinate information
    coordinates = set(
        k.split('.')[-1] for k in flat_locations.keys()
        if '.coordinates.' in k)

    if coordinates:
        sets.coordinates = coordinates

    # `timesteps` set is built from the first key in timeseries_data,
    # since we already ensure in process_timeseries_data() that all
    # timeseries have the exact same index
    first_key = list(model_run.timeseries_data.keys())[0]
    sets.timesteps = list(model_run.timeseries_data[first_key].index.astype(str))

    # `techlists` are strings with comma-separated techs used for grouping in
    # some model-wide constraints
    sets.techlists = set()
    for k in model_run.model.get_key('group_share', {}).keys():
        sets.techlists.add(k)

    return sets


def generate_loc_tech_sets(model_run, simple_sets):
    """
    Generate loc-tech sets for a given pre-processed ``model_run``

    Parameters
    ----------
    model_run : AttrDict
    simple_sets : AttrDict
        Simple sets returned by ``generate_simple_sets(model_run)``.

    """
    sets = AttrDict()

    ##
    # First deal with transmission techs, which can show up only in
    # loc_techs_transmission, loc_techs_milp, and loc_techs_purchase
    ##

    # All `tech:loc` expanded transmission technologies
    sets.loc_techs_transmission = set(concat_iterable([
        (i, u, j) for i, j, u in product(  # (loc, loc, tech) product
            simple_sets.locs,
            simple_sets.locs,
            simple_sets.techs_transmission_names)
        if model_run.get_key(
            'locations.{}.links.{}.techs.{}'.format(i, j, u), None
        )
    ], ['::', ':']))

    # A dict of transmission tech config objects
    # to make parsing for set membership easier
    loc_techs_transmission_config = {
        k: model_run.get_key(
            'locations.{loc_from}.links.{loc_to}.techs.{tech}'
            .format(**split_loc_techs_transmission(k))
        )
        for k in sets.loc_techs_transmission
    }

    ##
    # Now deal with the rest of the techs and other sets
    ##

    # Only loc-tech combinations that actually exist
    sets.loc_techs_non_transmission = set(concat_iterable([
        (l, t) for l, t in product(
            simple_sets.locs,
            simple_sets.techs_non_transmission)
        if model_run.get_key('locations.{}.techs.{}'.format(l, t), None)
    ], ['::']))

    sets.loc_techs = sets.loc_techs_non_transmission | sets.loc_techs_transmission

    # A dict of non-transmission tech config objects
    # to make parsing for set membership easier
    loc_techs_config = {
        k: model_run.get_key(
            'locations.{}.techs.{}'.format(*k.split('::'))
        )
        for k in sets.loc_techs_non_transmission
    }

    loc_techs_all_config = {**loc_techs_config, **loc_techs_transmission_config}

    ##
    # Sets based on membership in abstract base technology groups
    ##

    for group in [
            'storage', 'demand', 'supply', 'supply_plus',
            'conversion', 'conversion_plus']:
        tech_set = set(
            k for k in sets.loc_techs_non_transmission
            if model_run.techs[k.split('::')[1]].inheritance[-1] == group
        )
        sets['loc_techs_{}'.format(group)] = tech_set

    sets.loc_techs_non_conversion = set(
        k for k in sets.loc_techs_non_transmission
        if k not in sets.loc_techs_conversion and
        k not in sets.loc_techs_conversion_plus
    ) | sets.loc_techs_transmission

    # Techs that introduce energy into the system
    sets.loc_techs_supply_all = (
        sets.loc_techs_supply |
        sets.loc_techs_supply_plus |
        sets.loc_techs_conversion |
        sets.loc_techs_conversion_plus
    )

    ##
    # Sets based on specific constraints being active
    ##

    # Technologies that specify resource_area constraints
    sets.loc_techs_area = set(
        k for k in sets.loc_techs_non_transmission
        if any('resource_area' in i
               for i in loc_techs_config[k].constraints.keys_nested())
    )

    # Technologies that define storage, which can include `supply_plus`
    # and `storage` groups.
    sets.loc_techs_store = set(
        k for k in sets.loc_techs_supply_plus
        if any('storage_' in i
               for i in loc_techs_config[k].constraints.keys_nested())
    ) | sets.loc_techs_storage

    # technologies that specify a finite resource
    sets.loc_techs_finite_resource = set(
        k for k in sets.loc_techs_non_transmission
        if loc_techs_config[k].constraints.get('resource') and
        not (loc_techs_config[k].constraints.get('resource') in ['inf', np.inf])
    )

    # `supply` technologies that specify a finite resource
    sets.loc_techs_finite_resource_supply = (
        sets.loc_techs_finite_resource.intersection(sets.loc_techs_supply)
    )

    # `demand` technologies that specify a finite resource
    sets.loc_techs_finite_resource_demand = (
        sets.loc_techs_finite_resource.intersection(sets.loc_techs_demand)
    )

    # `supply_plus` technologies that specify a finite resource
    sets.loc_techs_finite_resource_supply_plus = (
        sets.loc_techs_finite_resource.intersection(sets.loc_techs_supply_plus)
    )

    # Technologies that define ramping constraints
    sets.loc_techs_ramping = set(
        k for k in sets.loc_techs_non_transmission
        if 'energy_ramping' in loc_techs_config[k].constraints
    )

    # Technologies that allow export
    sets.loc_techs_export = set(
        k for k in sets.loc_techs_non_transmission
        if 'export_carrier' in loc_techs_config[k].constraints
    )

    # Technologies that allow purchasing discrete units
    # NB: includes transmission techs!
    loc_techs_purchase = set(
        k for k in sets.loc_techs_non_transmission
        if any('.purchase' in i
               for i in loc_techs_config[k].get(
                   'costs', AttrDict()).keys_nested()) and not
        any('units_' in i
            for i in loc_techs_config[k].get(
                'constraints', AttrDict()).keys_nested())
    )

    transmission_purchase = set(
        k for k in sets.loc_techs_transmission
        if any('.purchase' in i
               for i in loc_techs_transmission_config[k].get(
                   'costs', AttrDict()).keys_nested()) and not
        any('units_' in i
            for i in loc_techs_transmission_config[k].get(
                'constraints', AttrDict()).keys_nested())
    )

    sets.loc_techs_purchase = loc_techs_purchase | transmission_purchase

    # Technologies with MILP constraints
    loc_techs_milp = set(
        k for k in sets.loc_techs_non_transmission
        if any('units_' in i
               for i in loc_techs_config[k].constraints.keys_nested())
    )

    transmission_milp = set(
        k for k in sets.loc_techs_transmission
        if any('units_' in i
               for i in loc_techs_transmission_config[k].constraints.keys_nested())
    )

    sets.loc_techs_milp = loc_techs_milp | transmission_milp

    ##
    # Sets based on specific costs being active
    # NB includes transmission techs
    ##

    loc_techs_costs = set(
        k for k in sets.loc_techs_non_transmission
        if any('costs' in i
               for i in loc_techs_config[k].keys())
    )

    loc_techs_transmission_costs = set(
        k for k in sets.loc_techs_transmission
        if any('costs' in i
               for i in loc_techs_transmission_config[k].keys())
    )

    # Any capacity or fixed annual costs
    loc_techs_investment_costs = set(
        k for k in loc_techs_costs
        if any('_cap' in i or '.purchase' in i or '_area' in i
               for i in loc_techs_config[k].costs.keys_nested())
    )
    loc_techs_transmission_investment_costs = set(
        k for k in loc_techs_transmission_costs
        if any('_cap' in i or '.purchase' in i or '_area' in i
               for i in loc_techs_transmission_config[k].costs.keys_nested())
    )

    # Any operation and maintenance
    loc_techs_om_costs = set(
        k for k in loc_techs_costs
        if any('om_' in i or 'export' in i
               for i in loc_techs_config[k].costs.keys_nested())
    )
    loc_techs_transmission_om_costs = set(
        k for k in loc_techs_transmission_costs
        if any('om_' in i
               for i in loc_techs_transmission_config[k].costs.keys_nested())
    )

    # Any export costs
    sets.loc_techs_costs_export = set(
        k for k in loc_techs_costs
        if any('export' in i
               for i in loc_techs_config[k].costs.keys_nested())
    )

    sets.loc_techs_cost = loc_techs_costs | loc_techs_transmission_costs
    sets.loc_techs_investment_cost = (
        loc_techs_investment_costs |
        loc_techs_transmission_investment_costs)
    sets.loc_techs_om_cost = loc_techs_om_costs | loc_techs_transmission_om_costs

    ##
    # Subsets of costs for different abstract base technologies
    ##

    sets.loc_techs_om_cost_conversion = loc_techs_om_costs.intersection(sets.loc_techs_conversion)
    sets.loc_techs_om_cost_conversion_plus = loc_techs_om_costs.intersection(sets.loc_techs_conversion_plus)
    sets.loc_techs_om_cost_supply = loc_techs_om_costs.intersection(sets.loc_techs_supply)
    sets.loc_techs_om_cost_supply_plus = loc_techs_om_costs.intersection(sets.loc_techs_supply_plus)

    ##
    # Subsets of `conversion_plus` technologies
    ##

    # `conversion_plus` technologies with secondary carrier(s) out
    sets.loc_techs_out_2 = set(
        k for k in sets.loc_techs_conversion_plus if 'carrier_out_2' in
        model_run.techs[k.split('::')[1].split(':')[0]].essentials
    )

    # `conversion_plus` technologies with tertiary carrier(s) out
    sets.loc_techs_out_3 = set(
        k for k in sets.loc_techs_conversion_plus if 'carrier_out_3' in
        model_run.techs[k.split('::')[1].split(':')[0]].essentials
    )

    # `conversion_plus` technologies with secondary carrier(s) in
    sets.loc_techs_in_2 = set(
        k for k in sets.loc_techs_conversion_plus if 'carrier_in_2' in
        model_run.techs[k.split('::')[1].split(':')[0]].essentials
    )

    # `conversion_plus` technologies with tertiary carrier(s) in
    sets.loc_techs_in_3 = set(
        k for k in sets.loc_techs_conversion_plus if 'carrier_in_3' in
        model_run.techs[k.split('::')[1].split(':')[0]].essentials
    )

    ##
    # `loc_tech_carrier` sets
    ##

    # loc_tech_carriers for all technologies that have energy_prod=True
    sets.loc_tech_carriers_prod = set(
        '{}::{}'.format(k, carrier)
        for k in sets.loc_techs
        if loc_techs_all_config[k].constraints.get_key('energy_prod', False)
        for carrier in get_all_carriers(model_run.techs[k.split('::')[1].split(':')[0]].essentials, direction='out')
    )

    # loc_tech_carriers for all technologies that have energy_con=True
    sets.loc_tech_carriers_con = set(
        '{}::{}'.format(k, carrier)
        for k in sets.loc_techs
        if loc_techs_all_config[k].constraints.get_key('energy_con', False)
        for carrier in get_all_carriers(model_run.techs[k.split('::')[1].split(':')[0]].essentials, direction='in')
    )

    # loc_tech_carriers for all supply technologies
    sets.loc_tech_carriers_supply_all = set(
        '{}::{}'.format(k, carrier)
        for k in sets.loc_techs_supply_all
        for carrier in get_all_carriers(model_run.techs[k.split('::')[1].split(':')[0]].essentials, direction='out')
    )

    # loc_tech_carriers for all demand technologies
    sets.loc_tech_carriers_demand = set(
        '{}::{}'.format(k, carrier)
        for k in sets.loc_techs_demand
        for carrier in get_all_carriers(model_run.techs[k.split('::')[1].split(':')[0]].essentials, direction='in')
    )

    # loc_tech_carriers for all technologies that have export
    sets.loc_tech_carriers_export = set(
        '{}::{}'.format(k, loc_techs_all_config[k].constraints.export_carrier)
        for k in sets.loc_techs
        if loc_techs_all_config[k].constraints.get_key('export_carrier', False)
    )

    # loc_tech_carriers for `conversion_plus` technologies
    sets.loc_tech_carriers_conversion_plus = set(
        k for k in sets.loc_tech_carriers_con | sets.loc_tech_carriers_prod
        if k.rsplit('::', 1)[0] in sets.loc_techs_conversion_plus
    )

    # loc_carrier combinations that exist with either a con or prod tech
    sets.loc_carriers = set(
        '{0}::{2}'.format(*k.split('::'))
        for k in sets.loc_tech_carriers_prod | sets.loc_tech_carriers_con
    )

    return sets
