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

model_data.py
~~~~~~~~~~~~~~~~~~

Functionality to build the model-internal data array and process
time-varying parameters.

"""

import collections

import ruamel.yaml
import xarray as xr
import numpy as np

from calliope.core.attrdict import AttrDict
from calliope._version import __version__
from calliope.core.preprocess import checks
from calliope.core.preprocess.util import split_loc_techs_transmission, concat_iterable
from calliope.core.preprocess.time import add_time_dimension
from calliope.core.preprocess.lookup import add_lookup_arrays


def build_model_data(model_run, debug=False):
    """
    Take a Calliope model_run and convert it into an xarray Dataset, ready for
    constraint generation. Timeseries data is also extracted from file at this
    point, and the time dimension added to the data

    Parameters
    ----------
    model_run : AttrDict
        preprocessed model_run dictionary, as produced by
        Calliope.core.preprocess_model
    debug : bool, default = False
        Used to debug steps within build_model_data, particularly before/after
        time dimension addition. If True, more information is returned

    Returns
    -------
    data : xarray Dataset
        Dataset with optimisation parameters as variables, optimisation sets as
        coordinates, and other information in attributes.
    data_dict : dict, only returned if debug = True
        dictionary of parameters, prior to time dimension addition. Used here to
        populate the Dataset (using `from_dict()`)
    data_pre_time : xarray Dataset, only returned if debug = True
        Dataset, prior to time dimension addition, with optimisation parameters
        as variables, optimisation sets as coordinates, and other information
        in attributes.
    """
    # We build up a dictionary of the data, then convert it to an xarray Dataset
    # before applying time dimensions
    data = xr.Dataset(
        coords=add_sets(model_run),
        attrs=add_attributes(model_run)
    )

    data_dict = dict()
    data_dict.update(constraints_to_dataset(model_run))
    data_dict.update(costs_to_dataset(model_run))
    data_dict.update(location_specific_to_dataset(model_run))
    data_dict.update(tech_specific_to_dataset(model_run))
    data_dict.update(carrier_specific_to_dataset(model_run))

    data.merge(xr.Dataset.from_dict(data_dict), inplace=True)

    add_lookup_arrays(data, model_run)

    if debug:
        data_pre_time = data.copy(deep=True)

    add_time_dimension(data, model_run)

    # Carrier information uses DataArray indexing in the function, so we merge
    # these directly into the main xarray Dataset

    if debug:
        return data, data_dict, data_pre_time
    else:
        return data


def add_sets(model_run):
    coords = dict()
    for key, value in model_run.sets.items():
        if value:
            coords[key] = value
    for key, value in model_run.constraint_sets.items():
        if value:
            coords[key] = value
    return coords


def constraints_to_dataset(model_run):
    """
    Extract all constraints from the processed dictionary (model.model_run) and
    return an xarray Dataset with all the constraints as DataArray variables and
    model sets as Dataset dimensions.

    Parameters
    ----------
    model_run : AttrDict
        processed Calliope model_run dict

    Returns
    -------
    data_dict : dict conforming to xarray conventions

    """
    data_dict = dict()

    # FIXME: hardcoding == bad
    def _get_set(constraint):
        """
        return the set of loc_techs over which the given constraint should be
        built
        """
        if '_area' in constraint:
            return 'loc_techs_area'
        elif any(i in constraint for i in ['resource_cap',  'parasitic', 'resource_min_use']):
            return 'loc_techs_supply_plus'
        elif 'resource' in constraint: # i.e. everything with 'resource' in the name that isn't resource_cap
            return 'loc_techs_finite_resource'
        elif 'storage' in constraint or 'charge_rate' in constraint:
            return 'loc_techs_store'
        elif 'purchase' in constraint:
            return 'loc_techs_purchase'
        elif 'units_' in constraint:
            return 'loc_techs_milp'
        elif 'export' in constraint:
            return 'loc_techs_export'
        else:
            return 'loc_techs'

    # find all constraints which are actually defined in the yaml file
    relevant_constraints = set(i.split('.constraints.')[1]
                               for i in model_run.locations.as_dict_flat().keys()
                               if '.constraints.' in i and
                               '.carrier_ratios.'not in i)
    for constraint in relevant_constraints:
        data_dict[constraint] = dict(dims=_get_set(constraint), data=[])
        for loc_tech in model_run.sets[_get_set(constraint)]:
            loc, tech = loc_tech.split('::', 1)
            # for transmission technologies, we also need to go into link nesting
            if ':' in tech:  # i.e. transmission technologies
                tech, link = tech.split(':')
                loc_tech_dict = model_run.locations[loc].links[link].techs[tech]
            else:  # all other technologies
                loc_tech_dict = model_run.locations[loc].techs[tech]
            constraint_value = loc_tech_dict.constraints.get(constraint, np.nan)
            # inf is assumed to be string on import, so we need to np.inf it
            if constraint_value == 'inf':
                constraint_value = np.inf
            # add the value for the particular location & technology combination to the list
            data_dict[constraint]['data'].append(constraint_value)
        # once we've looped through all technology & location combinations, add the array to the dataset

    # Additional system-wide constraints from model_run.model
    # FIXME: hardcoding == bad
    data_dict['reserve_margin'] = {
        'data': [model_run.model.get('reserve_margin', {}).get(c, np.nan)
                 for c in model_run.sets['carriers']],
        'dims': 'carriers'
    }

    group_share_data = {}
    group_constraints = ['energy_cap_min', 'energy_cap_max', 'energy_cap_equals']
    group_constraints_carrier = ['carrier_prod_min', 'carrier_prod_max', 'carrier_prod_equals']

    for constraint in [  # Only process constraints that are defined
            c for c in group_constraints
            if c in ''.join(model_run.model.get_key('group_share', AttrDict()).keys_nested())]:
        group_share_data[constraint] = [
            model_run.model.get_key('group_share.{}.{}'.format(techlist, constraint), np.nan)
            for techlist in model_run.sets['techlists']
        ]

    for constraint in [  # Only process constraints that are defined
            c for c in group_constraints_carrier
            if c in ''.join(model_run.model.get_key('group_share', AttrDict()).keys_nested())]:
        group_share_data[constraint] = [
            [
                model_run.model.get_key('group_share.{}.{}.{}'.format(techlist, constraint, carrier), np.nan)
                for techlist in model_run.sets['techlists']
            ]
            for carrier in model_run.sets['carriers']
        ]

    # Add to data_dict and set dims correctly
    for k in group_share_data:
        data_dict['group_share_' + k] = {
            'data': group_share_data[k],
            'dims': 'techlists' if k in group_constraints else ('carriers', 'techlists')
        }

    return data_dict


def costs_to_dataset(model_run):
    """
    Extract all costs from the processed dictionary (model.model_run) and
    return an xarray Dataset with all the costs as DataArray variables. Variable
    names will be prepended with `cost_` to differentiate from other constraints

    Parameters