PY2 = sys.version_info[0] == 2
def with_metaclass(meta, *bases):
        def __new__(cls, name, this_bases, d):
def csv_open(filename, mode='r'):
    class metaclass(meta): class Cache(dict):
    def invalidate(self, period, entity_name, variable=None):
        """
        Invalidates all keys matching period, entity_name and possibly variable

        if variable is None, it matches all keys for that period and entity
        """
        # print("invalidate", variable.name, period, entity_name)
        for i, key in enumerate(self.keys()):
            # print(key)
            c_expr, c_period, c_entity_name, c_filter_expr = key
            # XXX: do we also need to invalidate when name not in expr but
            # name in filter_expr?
            expr_match = variable is None or variable in c_expr
            if (c_period == period and c_entity_name == entity_name and
                    expr_match):
                # print("matches", key, " => invalidating")
                del self[key] class EntityContext(object):
    def __init__(self, eval_ctx, entity, extra=None):
        self.eval_ctx = eval_ctx
        self.entity = entity
        if extra is None:
            extra = {}
        self.extra = extra

    def __getitem__(self, key):
        if key == 'period':
            return self.eval_ctx.period

        try:
            return self.extra[key]
        except KeyError:
            period = self.eval_ctx.period
            array_period = self.entity.array_period
            if period == array_period:
                try:
                    return self.entity.temp_variables[key]
                except KeyError:
                    try:
                        return self.entity.array[key]
                    except ValueError:
                        raise KeyError(key)
            else:
                # FIXME: lags will break if used from a context subset (eg in
                # new() or groupby(): all individuals will be returned instead
                # of only the "filtered" ones.
                if (self.entity.array_lag is not None and
                    array_period is not None and
                    period == array_period - 1 and
                        key in self.entity.array_lag.dtype.fields):
                    return self.entity.array_lag[key]

                bounds = self.entity.output_rows.get(period)
                if bounds is not None:
                    startrow, stoprow = bounds
                else:
                    startrow, stoprow = 0, 0
                return self.entity.table.read(start=startrow, stop=stoprow,
                                              field=key)

    # is the current array period the same as the context period?
    @property
    def is_array_period(self):
        return self.entity.array_period == self.eval_ctx.period

    def __setitem__(self, key, value):
        self.extra[key] = value

    def __delitem__(self, key):
        del self.extra[key]

    def __contains__(self, key):
        entity = self.entity
        period = self.eval_ctx.period
        array_period = entity.array_period
        # entity.array can be None! (eg. with "explore")
        keyinarray = (self.is_array_period and
                      (key in entity.temp_variables or
                       key in entity.array.dtype.fields))
        # we need to check explicitly whether the key is in array_lag because
        # with output=None it can contain more fields than table.
        keyinlagarray = (entity.array_lag is not None and
                         array_period is not None and
                         period == array_period - 1 and
                         key in entity.array_lag.dtype.fields)
        keyintable = (entity.table is not None and
                      key in entity.table.dtype.fields)
        return key in self.extra or keyinarray or keyinlagarray or keyintable

    def keys(self, extra=True):
        res = list(self.entity.array.dtype.names)
        res.extend(sorted(self.entity.temp_variables.keys()))
        if extra:
            res.extend(sorted(self.extra.keys()))
        # in theory, this should not be needed because we present defining a function argument with the same
        # name than an entity field
        assert list(unique(res)) == res
        return res

    def get(self, key, elsevalue=None):
        try:
            return self[key]
        except KeyError:
            return elsevalue

    def copy(self):
        return EntityContext(self.eval_ctx, self.entity, self.extra.copy())

    def clone(self, **kwargs):
        res = self.copy()
        allowed_kwargs = {'eval_ctx', 'entity', 'extra'}
        for k, v in kwargs.items():
            assert k in allowed_kwargs, "%s is not a valid kwarg" % k
            setattr(res, k, v)
        return res

    def length(self):
        if self.is_array_period:
            return len(self.entity.array)
        else:
            period = self.eval_ctx.period
            bounds = self.entity.output_rows.get(period)
            if bounds is not None:
                startrow, stoprow = bounds
                return stoprow - startrow
            else:
                return 0

    def __len__(self):
        return self.length()

    def list_periods(self):
        return sorted(self.entity.output_index.keys())

    @property
    def id_to_rownum(self):
        period = self.eval_ctx.period
        if self.is_array_period:
            return self.entity.id_to_rownum
        elif period in self.entity.output_index:
            return self.entity.output_index[period]
        else:
            # FIXME: yes, it's true, that if period is not in output_index, it
            # probably means that we are before start_period and in that case,
            # input_index == output_index, but it would be cleaner to simply
            # initialise output_index correctly
            return self.entity.input_index[period] class EvaluationContext(object):
    def __init__(self, simulation=None, entities=None, global_tables=None,
                 period=None, entity_name=None, filter_expr=None,
                 entities_data=None):
        """
        :param simulation: Simulation
        :param entities: dict of entities {name: entity}
        :param global_tables: dict of ndarrays (structured or not)
        :param period: int of the current period
        :param entity_name: name (str) of the current entity
        :param filter_expr: contextual filter expression (Expr)
        :param entities_data: dict of data for entities (dict of
                              EntityContext or dict of dict)
        :return:
        """
        self.simulation = simulation
        self.entities = entities
        self.global_tables = global_tables
        self.period = period
        self.entity_name = entity_name
        self.filter_expr = filter_expr
        if entities_data is None:
            entities_data = {name: EntityContext(self, entity)
                             for name, entity in entities.items()}
        self.entities_data = entities_data

    def copy(self, fresh_data=False):
        # FIXME: when fresh_data is False, entities_data should clone each
        # EntityContext and set their eval_ctx attribute to the newly created
        # EvaluationContext (res), otherwise each EntityContext still points to
        # the old EvaluationContext (which obviously can have some different
        # attributes, eg period).
        # The problem is that the fix currently breaks many models because in
        # some (all?) cases EvaluableExpression are stored in a context (extra)
        # that is not the same than the context where it is used (and thus the
        # temporary field is not found if extra is not exactly the same object).
        # entities_data = None if fresh_data else {}
        entities_data = None if fresh_data else self.entities_data.copy()
        # FIXME: when switching entities, filter should not come along, or maybe
        # filter should be a per-entity dict?
        res = EvaluationContext(self.simulation, self.entities,
                                self.global_tables, self.period,
                                self.entity_name, self.filter_expr,
                                entities_data)
        # res.entities_data = {name: ent_ctx.clone(eval_ctx=res)
        #                      for name, ent_ctx
        #                      in self.entities_data.items()}
        return res

    def clone(self, fresh_data=False, **kwargs):
        res = self.copy(fresh_data=fresh_data)
        allowed_kwargs = {'simulation', 'entities', 'global_tables',
                          'period', 'entity_name', 'filter_expr',
                          'entities_data', 'entity_data'}
        for k, v in kwargs.items():
            assert k in allowed_kwargs, "%s is not a valid kwarg" % k
            setattr(res, k, v)
        return res

    @property
    def entity(self):
        return self.entities.get(self.entity_name)

    @entity.setter
    def entity(self, value):
        assert value.name in self.entities, '%s not in %s' % (value,
                                                              self.entities)
        self.entity_name = value.name

    @property
    def id_to_rownum(self):
        entity_context = self.entity_data
        if hasattr(entity_context, 'id_to_rownum'):
            return entity_context.id_to_rownum
        else:
            # fall back on the entity itself
            return self.entity.id_to_rownum

    @property
    def entity_data(self):
        return self.entities_data[self.entity_name]

    @entity_data.setter
    def entity_data(self, value):
        self.entities_data[self.entity_name] = value

    def __getitem__(self, key):
        if key == 'period':
            return self.period
        else:
            return self.entity_data[key]

    def get(self, key, elsevalue=None):
        try:
            return self[key]
        except KeyError:
            return elsevalue

    def __setitem__(self, key, value):
        # XXX: how do we set a new global?
        self.entity_data[key] = value

    def __contains__(self, key):
        from liam2.expr import Variable
        if isinstance(key, Variable):
            entity, name = key.entity, key.name
            if entity is None:
                # FIXME: this is wrong (but currently needed because some
                # Variable are created without entity)
                return True
            else:
                return name in self.entities_data[entity.name]
        else:
            return key in self.entity_data

    def length(self):
        return context_length(self.entity_data)

    def keys(self, extra=True):
        entity_data = self.entity_data
        if isinstance(entity_data, EntityContext):
            return entity_data.keys(extra)
        else:
            assert isinstance(entity_data, dict)
            return list(entity_data.keys())

    def update(self, other, **kwargs):
        self.entity_data.update(other, **kwargs)

    def __len__(self):
        return self.length()

    def subset(self, index=None, keys=None, filter_expr=None):
        """
        returns a copy of the context with only a subset of the current entity.
        The main use case is to take a subset of rows. Since this is a
        costly operation, the user can also provide keys so that only the
        columns he needs are filtered.
        :param index: indices to take (list or ndarray)
        :param keys: list of column names or None (take all)
        :param filter_expr: expression used to compute the index. This is
        only used to compute the cache key
        :return:
        """
        data_subset = context_subset(self.entity_data, index, keys)
        return self.clone(entity_data=data_subset, filter_expr=filter_expr)

    def empty(self, length=None):
        """
        returns a copy of the context with the same length but no data.
        """
        return self.clone(entity_data=empty_context(length)) class ExplainTypeError(type):
    def __call__(cls, *args, **kwargs):
        try:
            # noinspection PyArgumentList
            return type.__call__(cls, *args, **kwargs)
        except TypeError as e:
            funcname = cls.funcname
            msg = e.args[0].replace('__init__()', funcname)

            def repl_py2(matchobj):
                needed, given = int(matchobj.group(1)), int(matchobj.group(3))
                word = matchobj.group(2)
                return '%d %s (%d given)' % (needed - 1, word, given - 1)

            def repl_py3_toomany_from_to(matchobj):
                nfrom, nto, given = [int(matchobj.group(n)) for n in (1, 2, 4)]
                word = matchobj.group(3)
                verb = 'were' if given != 2 else 'was'
                return 'from %d to %d positional %s but %d %s given' \
                       % (nfrom - 1, nto - 1, word, given - 1, verb)

            def repl_py3_toomany(matchobj):
                needed, given = int(matchobj.group(1)), int(matchobj.group(3))
                word = matchobj.group(2)
                verb = 'were' if given != 2 else 'was'
                return 'takes %d positional %s but %d %s given' \
                       % (needed - 1, word, given - 1, verb)

            def repl_py3_missing(matchobj):
                missing = int(matchobj.group(1))
                return 'missing %d positional argument' % (missing - 1)

            # Python2 style
            msg = re.sub('(\d+) (arguments?) \((\d+) given\)', repl_py2, msg)
            # Python3 style for too many args in the presence of default values
            msg = re.sub('from (\d+) to (\d+) positional (arguments?) but '
                         '(\d+) were given',
                         repl_py3_toomany_from_to, msg)
            # Python3 style for too many args with no default values
            # "takes" is included to not match from_to again
            msg = re.sub('takes (\d+) positional (arguments?) but '
                         '(\d+) were given',
                         repl_py3_toomany, msg)
            # Python3 style for missing
            msg = re.sub('missing (\d+) positional argument',
                         repl_py3_missing, msg)
            raise TypeError(msg) def safe_take(a, indices, missing_value):
    """
    like np.take but out-of-bounds indices return the missing value
    """
    indexed = a.take(indices, mode='clip')
    return ne.evaluate('where((idx < 0) | (idx >= maxidx), missing, indexed)',
                       {'idx': indices, 'maxidx': len(a),
                        'missing': missing_value, 'indexed': indexed}) class IrregularNDArray(object):
    """
    A wrapper for collections of arrays (eg list of arrays or arrays of
    arrays) to make them act somewhat like a 2D (numpy) array. This makes it
    possible to have irregular lengths in the second dimension.
    """
    def __init__(self, data):
        self.data = data

    prod = _make_aggregate(np.prod)
    sum = _make_aggregate(np.sum)
    min = _make_aggregate(np.min)
    max = _make_aggregate(np.max)

    def __getattr__(self, key):
        return getattr(self.data, key)

    def __getitem__(self, key):
        return self.data[key] class NiceArgSpec(FullArgSpec):
    def __str__(self):
        if self.defaults:
            ndef = len(self.defaults)
            l = list(self.args[:-ndef])
            l.extend([k + '=' + repr(v)
                      for k, v in zip(self.args[-ndef:], self.defaults)])
        else:
            l = list(self.args)
        if self.varargs:
            l.append('*' + self.varargs)
        if self.varkw:
            l.append('**' + self.varkw)
        if self.kwonlyargs:
            l.extend([k + '=' + repr(self.kwonlydefaults[k])
                      for k in self.kwonlyargs])
        return ', '.join(l) def englishenum(iterable):
    """
    Returns an "english enumeration" of the strings in the iterable.
    >>> englishenum(['a', 'b', 'c'])
    'a, b, and c'
    >>> englishenum('abc')
    'a, b, and c'
    >>> englishenum('ab')
    'a and b'
    >>> englishenum('a')
    'a'
    """
    l = list(iterable)
    if len(l) == 2:
        return '%s and %s' % tuple(l)
    elif len(l) > 2:
        l[-1] = 'and ' + l[-1]
    return ', '.join(l) def make_hashable(obj):
    if isinstance(obj, (list, tuple)):
        return tuple(make_hashable(e) for e in obj)
    elif isinstance(obj, slice):
        return (slice, obj.start, obj.stop, obj.step)
    else:
        return obj def add_context(exception, s):
    if isinstance(exception, SyntaxError) and exception.offset is not None:
        # most SyntaxError are clearer if left unmodified since they already
        # contain the faulty string but some do not (eg non-keyword arg after
        # keyword arg).
        # SyntaxError instances have 'filename', 'lineno', 'offset' and 'text'
        # attributes.
        return
    exception.liam2context = s def array_nan_equal(a, b):
    # np.array_equal is not implemented on strings in numpy < 1.9
    if np.issubdtype(a.dtype, np.str) and np.issubdtype(b.dtype, np.str):
        try:
            return (a == b).all()
        except ValueError:
            return False
    else:
        return np.all((a == b) | (np.isnan(a) & np.isnan(b)))
import inspect
import types
import numpy as np
import larray as la
    import numexpr
from collections import Counter
from liam2.compat import basestring, PY2, zip, with_metaclass
from liam2.cache import Cache
from liam2.config import debug
from liam2.context import EntityContext, EvaluationContext
from liam2.utils import (ExplainTypeError, safe_take, IrregularNDArray, NiceArgSpec, englishenum, make_hashable,
                         add_context, array_nan_equal)
        from liam2.process import Assignment, Function
            cls.argspec = NiceArgSpec._make(spec + extra)

    def get_compute_func(cls):
        raise NotImplementedError()


class AbstractFunction(with_metaclass(FillFuncNameMeta, Expr)):
    __children__ = ('args', 'kwargs')

    funcname = None
    argspec = None

    def __init__(self, *args, **kwargs):
        # The behavior/error messages match Python 3.4 (and probably other 3.x)
        argnames = self.argspec.args
        maxargs = len(argnames)
        defaults = self.argspec.defaults
        nreqargs = maxargs - (len(defaults) if defaults is not None else 0)
        reqargnames = argnames[:nreqargs]
        allowed_kwargs = set(argnames) | set(self.argspec.kwonlyargs)
        funcname = self.funcname
        assert funcname is not None

        nargs = len(args)
        availposargnames = set(argnames[:nargs])
        availkwargnames = set(kwargs.keys())
        dupeargnames = availposargnames & availkwargnames
        if dupeargnames:
            raise TypeError("%s() got multiple values for argument '%s'"
                            % (funcname, dupeargnames.pop()))

        # Check that we do not have invalid kwargs
        extra_kwargs = availkwargnames - allowed_kwargs
        # def f(**kwargs) => argspec.varkw = 'kwargs'
        if extra_kwargs and self.argspec.varkw is None:
            raise TypeError("%s() got an unexpected keyword argument '%s'"
                            % (funcname, extra_kwargs.pop()))

        # Check that we do not have too many args
        if self.argspec.varargs is None and nargs > maxargs:
            # f() takes 3 positional arguments but 4 were given
            # f() takes from 1 to 3 positional arguments but 4 were given
            # + 1 to be consistent with Python (to account for self) but
            # those will be modified again (-1) in ExplainTypeError
            posargs = str(nreqargs + 1) if nreqargs == maxargs \
                else "from %d to %d" % (nreqargs + 1, maxargs + 1)

            msg = "%s() takes %s positional argument%s but %d were given"
            raise TypeError(msg % (funcname, posargs,
                                   's' if maxargs > 1 else '', nargs + 1))

        # Check that we have all required args (passed either as args or kwargs)
        missing = [name for name in reqargnames
                   if name not in (availposargnames | availkwargnames)]
        if missing:
            nmissing = len(missing)
            # f() missing 1 required positional argument: 'a'
            # f() missing 2 required positional arguments: 'a' and 'b'
            # f() missing 3 required positional arguments: 'a', 'b', and 'c'
            # + 1 to be consistent with Python (to account for self) but
            # those will be modified again (-1) in ExplainTypeError
            raise TypeError("%s() missing %d positional argument%s: %s"
                            % (funcname,
                               nmissing + 1,
                               's' if nmissing > 1 else '',
                               englishenum(repr(a) for a in missing)))

        # save original arguments before we mess with them
        self._original_args = args, sorted(kwargs.items())

        # move all "non-kwonly" kwargs to args
        # def func(a, b, c, d, e=1, f=1):
        #     pass
        # nreqargs = 4, maxargs = 6
        # >>> func(1, 2, c=3, d=4, f=5)
        # nargs = 2
        # >>> func(1, 2, 3, 4, 5)
        # nargs = 5
        # 1) required arguments (without a default value) passed as kwargs
        #    pop() should not raise otherwise the "if missing" test above would
        #    have triggered an exception)
        extra_args = [kwargs.pop(name) for name in argnames[nargs:nreqargs]]

        # 2) optional args (with a default value) not passed as positional args
        if defaults is not None:
            # number of optional args passed as positional args
            nposopt = max(nargs - nreqargs, 0)
            extra_args.extend([kwargs.pop(argname, default)
                               for argname, default
                               in zip(argnames[nreqargs + nposopt:],
                                      defaults[nposopt:])])

        args = args + tuple(extra_args)
        kwargs = tuple(sorted(kwargs.items()))
        self.args = args
        self.kwargs = kwargs
        self._variables = None

    @staticmethod
    def format_args_str(args, kwargs):
        """
        :param args: list of strings
        :param kwargs: list of (k, v) where both k and v are strings
        :return: a single string
        """
        return ', '.join(list(args) + ['%s=%s' % (k, v) for k, v in kwargs])

    @staticmethod
    def format(funcname, args, kwargs):
        args = [repr(a) for a in args]
        kwargs = [(str(k), repr(v)) for k, v in kwargs]
        return '%s(%s)' % (funcname,
                           AbstractFunction.format_args_str(args, kwargs))

    def __repr__(self):
        return self.format(self.funcname, *self._original_args)


# this needs to stay in the expr module because of ExprAttribute, which uses
# DynamicFunctionCall -> GenericFunctionCall -> FunctionExpr

class FunctionExpr(with_metaclass(FillArgSpecMeta, EvaluableExpression, AbstractFunction)):