from __future__ import division, print_function, absolute_import

import os
import pickle
import json
import copy
import logging
import importlib
import collections
import operator
import functools
import diskcache

from satsim.math.random import gen_sample

logger = logging.getLogger(__name__)
caches = {}
_config = {}


def parse_function(param):
    """Parse a function. The input argument, `param`, should contain
    the following keys and values which describe the function to run.
    For example::

        param = {
            'module': 'my_module',
            'function': 'my_function',
            'kwargs': { 'param_1': 10.0, 'param_2': 'astring' }
        }

    Which is comparable to::

        import my_module
        return lambda: my_function(param_1 = 10.0, param_2 = 'astring')

    Args:
        param: `dict`, a dictionary that contain keys: `module`, `function`,
            `kwargs`

    Returns:
        A `function`, a lambda function encapsulating the parsed function and
            arguments
    """
    module = importlib.import_module(param['module'])

    if '$function' in param:
        function = getattr(module, param['$function'])
    else:
        function = getattr(module, param['function'])

    if 'kwargs' in param:
        kwargs2 = param['kwargs']

        def f(*args, **kwargs):
            kwargs2.update(**kwargs)
            return function(*args, **kwargs2)

        return f
    else:
        return function


def parse_function_pipeline(param):
    """Parse a pipeline of functions. See `parse_function`.

    Args:
        param: `dict`, a dictionary that contain keys: `module`, `function`,
            `kwargs`

    Returns:
        An `func`, a function encapsulating the pipeline
    """

    ff = list(map(parse_function, param))

    def func(x, t=None, **kwargs):
        for f in ff:
            try:  # TODO refactor, this is messy
                from inspect import signature
                sig = signature(f)
                if 't' in sig.parameters:
                    x = f(x, t, **kwargs)
                else:
                    x = f(x, **kwargs)
            except Exception:
                x = f(x, t, **kwargs)
        return x

    return func


def parse_cache(param):
    """Parse a cache parameter. The input argument, `param`, may contain
    the following keys and values which describe the function to run.
    For example::

        param = {
            '$cache': './path/to/cache'
        }

    Cache is not used if `$cache` key is not present.

    Args:
        param: `dict`, a dictionary that contain keys: `$cache`

    Returns:
        A `tuple`,
            a dictionary, loaded from cache if available
            a boolean, True if cache loaded, False otherwise
    """
    if _has_rkey(param, '$cache'):
        cache = _create_or_get_diskcache(param)
        key = _param_to_cache_key(param)
        if key in cache:
            return cache[key], True

    return param, False


def parse_import(param):
    """Parse a import key. The input argument, `param`, should contain the
    following keys and values which describe the function to run.
    For example::

        param = {
            '$import': './base.json',
            'override': { 'param_1': 10.0, 'param_2': 'override string' }
        }

    Args:
        param: `dict`, a dictionary that contain keys: `import`, and optionally
            `override`

    Returns:
        A `dict`, a dictionary generated by the generator function
    """
    with open(param[_rkey(param, '$import')]) as json_file:
        data = json.load(json_file)

        if 'key' in param:
            keys = param['key'].split('.')
            data = functools.reduce(operator.getitem, keys, data)

        if 'override' in param:
            dict_merge(data, param['override'])

        return data


def parse_generator(param):
    """Parse and run a generator function. See `parse_function`.

    Args:
        param: `dict`, a dictionary that contain keys: `module`, `function`,
            `kwargs`

    Returns:
        A `dict`, a dictionary generated by the generator function
    """
    f = parse_function(param)

    return f()


def parse_random_sample(param):
    """Parses a random sample parameter and returns a sample. The value of
    the key `$sample` should be the format `random.DISTRIBUTION`, where
    `DISTRIBUTION` is a NumPy random distribution. For example::

        p = { "$sample": "random.uniform", "low": 5.0, "high": 22.0 }
        x = parse_random_sample(p)
        # x is a value between 5 and 22

    If the key `$sample` is equal to `random.choice`, one option will be chosen
    from the key `choices`. For example::

        p = { "$sample": "random.choice", "choices": [1, 2, 4, 8] }
        x = parse_random_sample(p)
        # x will be 1, 2, 4, or 8

    If the key `$sample` is equal to `random.list`, a list of random samples
    will be generated. For example::

        p = {
            "$sample": "random.list",
            "length": 5,
            "value": { "$sample": "random.randint", "low": 0, "high": 10 }
        }
        x = parse_random_sample(p)
        # x will a list of length 5, i.e. x == [7,5,8,5,1]

    Args:
        param: `dict`, random sample parameter

    Returns:
        A `number`, random sample
    """
    compat_name = _rkey(param, '$sample')
    stype, rtype = param[compat_name].split('.')

    if stype == 'random':

        # handle special types first
        if rtype == 'choice':
            c = param['choices']
            return c[gen_sample('randint', low=0, high=len(c))]
        # elif rtype == 'bins': #TODO
        #     return param
        elif rtype == 'list':
            samples = []
            if 'list' in param:  # backward compat
                logger.warning('Deprecated list replacement. Use $sample keyword inline instead.')
                for i in range(parse_param(param['length'])):
                    samples.append(_transform(copy.deepcopy(param['list'])))
                param['list'] = samples
                del param[compat_name]
                del param['length']
            else:
                for i in range(parse_param(param['length'])):
                    samples.append(_transform(copy.deepcopy(param['value'])))
                del param[compat_name]
                param = samples

            return param

        # everything else is numpy random
        else:
            del param[compat_name]
            param = _ref(param, _config)
            param = _transform(param)
            return gen_sample(rtype, **param)


def parse_param(param, dirname=None, run_generator=False, eval_python=False, run_compound=False):
    """Parses a parameter and recursively parses any children parameters.

    * if the parameter contains the key `sample`, result from
      `parse_random_sample` will be returned.
    * if the parameter contains the key `file`, the contents of the loaded
      pickle file will be returned.
    * if the parameter contains the key `generator`, result from
      `parse_generator` will be returned.
    * else if no special keys are found, `param` is returned.

    Args:
        param: `dict`, root parameter
        dirname: `str`, root directory for files referenced in `param`
        run_generator: `boolean`, evaluate generator functions during this pass

    Returns:
        A `any`, a random sample, contents of pickle file, or original value
    """
    # traverse dict
    if isinstance(param, dict):
        if _has_rkey(param, '$sample'):
            return parse_random_sample(param)
        elif _has_rkey(param, '$file'):
            with open(os.path.join(dirname, param[_rkey(param, '$file')]), 'rb') as f:
                return pickle.load(f)
        elif run_generator and '$function' in param:
            val = parse_function(param)()
            save_cache(param, val)
            return val
        elif run_generator and _has_rkey(param, '$generator'):
            return parse_generator(param[_rkey(param, '$generator')])
        elif eval_python and _has_rkey(param, '$pipeline'):
            return parse_function_pipeline(param[_rkey(param, '$pipeline')])
        elif run_compound and _has_rkey(param, '$compound'):
            val = functools.reduce(lambda x, y: x + y, param['$compound'])
            save_cache(param, val)
            return val
        else:
            return param

    # traverse arrays
    elif isinstance(param, list):
        param_out = []
        for i in range(len(param)):
            # TODO this needs to be cleaned up
            if isinstance(param[i], list):
                p = parse_param(param[i], dirname, run_generator, eval_python)
                param_out.append(p)
            else:
                p = parse_param(param[i], dirname, run_generator, eval_python)
                if isinstance(p, list):
                    param_out.extend(p)
                else:
                    param_out.append(p)
        return param_out

    # static variable
    else:
        return param


def transform(config, dirname=None, max_stages=5, with_debug=False, max_imports=20, max_ref=5):
    """Transforms a SatSim configuration, evaluating any random sample
    parameters, references and/or generator functions. The transformation will
    happen in two passes. The first pass evaluates `sample` parameters; the
    second pass evaluates `generator` parameters. This is repeated for
    `max_stages` or until no more `generator` or `sample` keys are found.

    Args:
        config: `dict`, configuration as a dictionary.
        dirname: `str`, root directory for files referenced in `config`.
        max_stages: `int`, maximum number of times to evaluate (one stage is
            equivalent to running pass 1 and pass 2 once each).
        with_debug: `boolean`, return config after each stage.
        max_imports: `int`, maximum number of times to evaluate imports.
        max_refs: `int`, maximum number of times to evaluate refs.

    Returns:
        A `dict` or `tuple`,
            if `with_debug` is `False`:
                `dict`, in place transformation of `config`
            if `with_debug` is `True`:
                config: `dict`, in place transformation of `config`
                stages: `array`, array of dictionaries after each transformation stage
    """
    stages = []
    global _config  # TODO needed to evaluate refs, make local
    _config = config

    # save version and sim parameters
    sim_param = config['sim']
    ver_param = config['version']

    def eval_sample(config):
        for i in range(max_stages):

            # evaluate random samples
            config = _transform(config, dirname, False)
            stages.append(copy.deepcopy(config))

            if _has_rkey_deep(config, '$sample') is False:
                break

        return config

    config = eval_sample(config)

    for i in range(max_imports):

        # evaluate all imports
        config = _import({'root': config})['root']

        if(_has_rkey_deep(config, '$import') is False):
            break

    # restore sim parameters
    config['sim'] = sim_param
    config['version'] = ver_param
    stages.append(copy.deepcopy(config))

    config = eval_sample(config)

    for i in range(max_ref):

        config = _ref(config, config)

        # evaluate all refs
        if _has_rkey_deep(config, '$ref') is False:
            break

    stages.append(copy.deepcopy(config))
    config = eval_sample(config)

    # eval all caches (done before function and generators)
    config = _cache(config, config)

    for i in range(max_stages):

        # pass 1: transform generator functions
        config = _transform(config, dirname, True)
        stages.append(copy.deepcopy(config))

        # pass 2: transform random samples
        config = _transform(config, dirname, False)
        stages.append(copy.deepcopy(config))

        if(_has_rkey_deep(config, '$sample') is False and _has_rkey_deep(config, '$generator') is False and has_key_deep(config, '$function') is False):
            break

    for i in range(max_stages):

        if(_has_rkey_deep(config, '$compound')) is False:
            break

        # evaluate all compounds
        config = _transform(config, dirname, run_compound=True)
        stages.append(copy.deepcopy(config))

    if with_debug:
        return config, stages
    else:
        return config


def save_debug(configs, output_dir):
    """Writes an array of configs to files for debugging.

    Args:
        configs: `array`, configuration stages output by `transform`
        output_dir: `str`, directory to output files
    """
    if output_dir is not None:
        for config, stage in zip(configs, range(len(configs))):
            with open(os.path.join(output_dir, 'config_pass_{}.json'.format(stage + 1)), 'w') as json_file:
                json.dump(config, json_file, indent=4, default=lambda o: "n/a")


def _cache(config, original):
    """Evaluate caches.

    Args:
        config: `dict`, configuration as a dictionary
        original: `dict`, original non-modified config

    Returns:
        A `dict`, in place transformation of `config`
    """
    for k, v in config.items():
        if isinstance(v, dict):
            if _has_rkey(v, '$cache'):
                config[k], is_cache = parse_cache(v)
                logger.debug('Cache found for {}: {}'.format(k, is_cache))
            else:
                config[k] = _cache(v, original)

    return config


def _ref(config, original):
    """Evaluate references.

    Args:
        config: `dict`, configuration as a dictionary
        original: `dict`, original non-modified config

    Returns:
        A `dict`, in place transformation of `config`
    """
    for k, v in config.items():
        if isinstance(v, dict):
            if _has_rkey(v, '$ref'):
                name = _rkey(v, '$ref')
                keys = v[name].split('.')
                config[k] = functools.reduce(operator.getitem, keys, original)
            else:
                config[k] = _ref(v, original)
        elif isinstance(v, list):  # TODO make this recursive
            for k2 in v:
                if isinstance(k2, dict):
                    _ref(k2, original)

    return config


def _import(config):
    """Imports a partial SatSim configuration file.

    Args:
        config: `dict`, configuration as a dictionary

    Returns:
        A `dict`, in place transformation of `config`
    """
    for k, v in config.items():
        if isinstance(v, dict):
            if _has_rkey(v, '$import'):
                config[k] = parse_import(v)
            else:
                config[k] = _import(v)
        elif isinstance(v, list):  # TODO make this recursive
            for k2 in v:
                if isinstance(k2, dict):
                    _import(k2)

    return config


def _transform(config, dirname=None, run_generator=False, eval_python=False, run_compound=False):
    """Transforms a SatSim configuration, evaluating any random sample
    parameters or references.

    Args:
        config: `dict`, configuration as a dictionary
        dirname: `str`, root directory for files referenced in `config`
        run_generator: `boolean`, evaluate generator functions during this pass

    Returns:
        A `dict`, in place transformation of `config`
    """
    if isinstance(config, dict):
        items = config.items()
    elif isinstance(config, list):
        items = enumerate(config)
    else:
        return config

    for k, v in items:
        tv = parse_param(v, dirname, run_generator, eval_python, run_compound)
        config[k] = _transform(tv, dirname, run_generator, eval_python, run_compound)

    return config


def _param_to_cache_key(param):
    """ Helper function to convert parameter to key or cache. """

    param_copy = copy.deepcopy(param)
    return json.dumps(param_copy)


def _create_or_get_diskcache(param):
    """ Helper function to get or create disk cache. """

    cache_path = param['$cache']
    if cache_path not in caches:
        caches[cache_path] = diskcache.Cache(os.path.join(cache_path, '.satsim_cache'), eviction_policy="none")
    cache = caches[cache_path]

    return cache


def save_cache(param, value):
    """ Saves a parameter to the disk cache. Parameter will be loaded from cache
    on any subsequent calls.

    Args:
        param: `dict`, root parameter to be used as the cache key
        value: `any`, the value to save to cache
    """

    if _has_rkey(param, '$cache'):
        cache = _create_or_get_diskcache(param)
        key = _param_to_cache_key(param)
        if key not in cache:
            logger.debug('Saving cache: {}'.format(key))
            cache[key] = value


def save_json(filename, config, save_pickle=False):
    """Saves a configuration to a JSON file. Any items that are not JSON
    serializable are saved to a binary pickle file and saved in the same
    directory as the file `filename`.

    Args:
        filename: `str`, path to file
        config: `dict`, configuration as a dictionary
    """
    def serialize(o):

        if save_pickle:
            picklename = '{:05d}.pickle'.format(serialize.i)
            with open(os.path.join(os.path.dirname(filename), picklename), 'wb') as picklefile:
                pickle.dump(o, picklefile)
        else:
            picklename = 'n/a'

        serialize.i = serialize.i + 1

        return {
            '$file': picklename
        }

    serialize.i = 0

    with open(filename, 'w') as outfile:
        json.dump(config, outfile, indent=4, default=serialize)


def load_json(filename):
    """Opens a JSON file and returns the result as a dictionary.

    Args:
        filename: `str`, path to file

    Returns:
        A `dict`, the result.
    """
    with open(filename, 'r') as f:
        ssp = json.load(f)
        return ssp


def load_yaml(filename):
    """Opens a YAML file and returns the result as a dictionary.

    Args:
        filename: `str`, path to file

    Returns:
        A `dict`, the result.
    """

    import yaml

    with open(filename, 'r') as f:
        ssp = yaml.load(f, Loader=yaml.SafeLoader)
        return ssp


def has_key_deep(d, key):
    """ Nested check to see if a dictionary has a key.

    Args:
        d: `dict` or `list`
        key: `str`, key name to check for
    """
    if isinstance(d, dict):
        items = d.items()
    elif isinstance(d, list):
        items = enumerate(d)
    else:
        return False

    for k, v in items:
        if isinstance(v, dict) or isinstance(v, list):
            if has_key_deep(v, key):
                return True

        if k == key:
            return True

    return False


def dict_merge(dct, merge_dct):
    """ Recursive dict merge. Inspired by :meth:`dict.update()`, instead of
    updating only top-level keys, dict_merge recurses down into dicts nested
    to an arbitrary depth, updating keys. The `merge_dct` is merged into
    `dct`.

    Param:
        dct: `dict`, onto which the merge is executed
        merge_dct: `dict`, dct merged into dct
    """
    for k, v in merge_dct.items():
        if (k in dct and isinstance(dct[k], dict) and isinstance(merge_dct[k], collections.abc.Mapping)):
            if '$sample' in dct[k]:  # TODO order dependency issue, need graph traversal to fix properly
                dct[k] = merge_dct[k]
            else:
                dict_merge(dct[k], merge_dct[k])
        else:
            dct[k] = merge_dct[k]


def _rkey(d, key):
    """ As of SatSim v0.8.2 SatSim replacement keywords should be prefixed with
    $. This function provides backward compatibility.
    """
    if key in d:
        return key
    elif key[1:] in d:
        logger.warning('Deprecated replacement keyword. Prefix keyword with $.')
        return key[1:]
    else:
        logger.error('Replacement keyword not found.')
        return None


def _has_rkey(d, key):
    """ As of SatSim v0.8.2 SatSim replacement keywords should be prefixed with
    $. This function provides backward compatibility.
    """
    if key in d:
        return True
    elif key[1:] in d:
        return True
    else:
        return False


def _has_rkey_deep(d, key):
    """ As of SatSim v0.8.2 SatSim replacement keywords should be prefixed with
    $. This function provides backward compatibility.
    """
    if has_key_deep(d, key):
        return True
    if has_key_deep(d, key[1:]):
        return True
    return False
