import itertools
import pandas as pd
from copy import deepcopy
from anytree import NodeMixin, RenderTree
from inspect import getfullargspec
from dcptree.helper_functions import *
from dcptree.analysis import to_group_data
from dcptree.data import check_data, has_intercept
from dcptree.group_helper import check_groups
from dcptree.classification_models import ClassificationModel
from dcptree.decoupled_set import build_model_assignment_map, DecoupledClassifierSet, exact_mcn_test

#### Scoring and Selection ####

def preference_score(data, partition, pooled_model):
    """

    scores a partition based on the probability of at least 1 fair_decoupling violation

    :param data:
    :param partition:
    :return:

    """
    assert callable(pooled_model)
    assert isinstance(partition, (list, DecoupledLeafSet))
    n_parts = len(partition)
    assert n_parts >= 1
    if isinstance(partition, DecoupledLeafSet):
        assert partition.is_partition
        assert partition.has_models
        predictors, indices = partition.predictors, partition.indices
    else:
        predictors, indices = zip(*partition)
        assert np.all(np.sum(np.stack(indices), axis = 0) == 1)

    X, Y = data['X'], data['Y']

    error_matrix = np.empty(shape = (n_parts, n_parts))
    envy_free_scores = np.empty(shape = (n_parts, n_parts))
    rationality_scores = np.empty(shape = (n_parts))

    for i, idx in enumerate(indices):
        n_samples = np.sum(idx)
        Yg, Xg = Y[idx], X[idx, :]
        for j, h in enumerate(predictors):
            mistakes = np.not_equal(Yg, h(Xg))
            error_matrix[i, j] = np.mean(mistakes)

        # compute errors for group i
        self_error = error_matrix[i, i]
        pooled_error = np.mean(np.not_equal(Yg, pooled_model(Xg)))

        # compute score components for group i
        rationality_scores[i] = 0.5 * n_samples * np.square(self_error - pooled_error)
        envy_free_scores[i, :] = 0.5 * n_samples * np.power(self_error - error_matrix[i, :], 2.0)

    rationality_scores = 4.0 * np.exp(-rationality_scores)
    envy_free_scores = 4.0 * np.exp(-envy_free_scores)
    np.fill_diagonal(envy_free_scores, 0.0)
    score = np.sum(rationality_scores) + np.sum(envy_free_scores)
    return score

def minimize_score_over_partitions(partitions):
    """
    :param partitions: list of dicts containing 'score'
    :return: index of partition that minimizes the score
    """
    assert isinstance(partitions, list)
    min_score = float('inf')
    min_idx = -1
    for i, p in enumerate(partitions):
        if p['score'] < min_score:
            min_score = p['score']
            min_idx = i
    return min_idx

#### Classes ####

class DecoupledTrainer(object):

    INFEASIBILITY_DICT = {
        'unspecified': 'unspecified',
        'n_samples': 'not enough samples',
        'n_samples_fold': 'not enough samples for a fold',
        'envyfree_violation': 'label prefers the model of another label',
        'decoupling_violation': 'label prefers not to decouple',
        'group_violation': 'label belongs to a group with a different infeasible label',
        }

    INFEASIBILITY_CODES = set(INFEASIBILITY_DICT.keys())

    #### initialization ####

    def __init__(self, groups, node_idx, pos_idx, neg_idx):

        assert check_groups(groups)
        self._node_idx = np.array(node_idx, dtype = bool).flatten()
        self._pos_idx = np.array(pos_idx, dtype = bool).flatten()
        self._neg_idx = np.array(neg_idx, dtype = bool).flatten()

        self._n = np.sum(self._node_idx)
        self._n_pos = np.sum(self._pos_idx)
        self._n_neg = np.sum(self._neg_idx)

        self._min_samples = 2
        self._min_samples_pos = 1
        self._min_samples_neg = 1

        self._names = list(groups.keys())
        self._labels = {g: list(groups[g]['labels']) for g in groups}
        self._indices = {g: np.array(groups[g]['indices'], dtype = int) for g in groups}

        self._ids = self._new_split_dictionary()
        for name, labels in self._ids.items():
            for l in labels:
                self._ids[name][l] = self._labels[name].index(l)

        self._models = self._new_split_dictionary()
        self._training_order = self.splits
        self._infeasible_splits = {}
        self._check_rep()
        self.set_sample_infeasibility()

    def _new_split_dictionary(self):
        d = dict.fromkeys(self._names)
        for name in d.keys():
            d[name] = dict.fromkeys(self._labels[name])
        return d

    #### validation ####

    def _check_rep(self):
        names = set(self.names)
        n_groups = len(names)
        assert len(self._training_order) >= 2 * n_groups
        assert names == set(self._ids.keys())
        assert names == set(self._indices.keys())
        assert names == set(self._models.keys())
        assert np.all(self._node_idx == np.logical_or(self._pos_idx, self._neg_idx))

    def _check_name_and_label(self, name, label):
        assert name in self._names
        assert label in self._labels[name]
        return True

    def _check_infeasibility_code(self, code):
        assert code in self.INFEASIBILITY_CODES

    #### properties

    @property
    def names(self):
        return list(self._names)

    @property
    def splits(self):
        splits = []
        for name, labels in self._labels.items():
            for l in labels:
                splits.append((name, l))
        return splits

    def labels(self, name):
        return list(self._labels[name])

    def split_indices(self, name, label):
        return np.logical_and(self._node_idx, self._ids[name][label] == self._indices[name])

    def split_indices_pos(self, name, label):
        return np.logical_and(self._pos_idx, self.split_indices(name, label))

    def split_indices_neg(self, name, label):
        return np.logical_and(self._neg_idx, self.split_indices(name, label))

    @property
    def feasible_names(self):
        return [n for n in self.names if n not in self._infeasible_splits]

    @property
    def infeasible_names(self):
        return [n for n in self.names if n in self._infeasible_splits]

    @property
    def feasible_splits(self):
        return list(filter(lambda v: not self.is_infeasible(v[0], v[1]), self.splits))

    @property
    def infeasible_splits(self):
        return list(filter(lambda v: self.is_infeasible(v[0], v[1]), self.splits))

    @property
    def min_samples(self):
        return int(self._min_samples)

    @property
    def min_samples_pos(self):
        return int(self._min_samples_pos)

    @property
    def min_samples_neg(self):
        return int(self._min_samples_neg)

    @min_samples.setter
    def min_samples(self, n):
        assert is_integer(n), 'min samples must be integer'
        assert n >= 2, 'min samples must be >= 2'
        self._min_samples = int(n)

    @min_samples_pos.setter
    def min_samples_pos(self, n):
        assert is_integer(n), 'min + samples must be integer'
        assert n >= 1, 'min + samples must be >= 2'
        n = int(n)
        self._min_samples_pos = min(n, self._n_pos)
        self._min_samples = max(self._min_samples_neg + self._min_samples_pos, self._min_samples)

    @min_samples_neg.setter
    def min_samples_neg(self, n):
        assert is_integer(n), 'min - samples must be integer'
        assert n >= 1, 'min - samples must be >= 2'
        n = int(n)
        self._min_samples_neg = min(n, self._n_neg)
        self._min_samples = max(self._min_samples_neg + self._min_samples_pos, self._min_samples)

    #### model management ####

    def has_model(self, name, label):
        return self._models[name][label] is not None

    def get_model(self, name, label):
        if self.has_model(name, label):
            return deepcopy(self._models[name][label])
        else:
            raise ValueError('no model exists for %s = %s' % (name, label))

    def set_model(self, name, label, model):
        self._check_name_and_label(name, label)
        self._models[name][label] = deepcopy(model)

    def needs_model(self, name, label):
        if self.is_infeasible(name, label):
            return False
        else:
            return not self.has_model(name, label)

    def get_models_for_name(self, name, include_infeasible = True):

        assert name in self.names
        models = {}

        if include_infeasible:
            for label in self.labels(name):
                model_item = {label: deepcopy(self._models[name][label])}
                models.update(model_item)
        else:
            assert name in self.feasible_names
            for label in self.labels(name):
                model_item = {label: deepcopy(self._models[name][label])}
                models.update(model_item)

        return models

    def get_feasible_models(self, propagate_infeasibility = True):

        feasible_names = self.feasible_names
        if propagate_infeasibility:
            feasible_names = filter(lambda n: not self.has_infeasible_split(n), feasible_names)

        models = {}
        for name in feasible_names:
            models[name] = {}
            for label in self.labels(name):
                model = self.get_model(name, label)
                models[name].update({label: model})

        return models

    #### infeasibility ####

    def is_infeasible(self, name, label):
        return name in self._infeasible_splits and label in self._infeasible_splits[name]

    def has_infeasible_split(self, name):
        return name in self._infeasible_splits

    def mark_as_infeasible(self, name, label, code = 'unspecified'):

        self._check_name_and_label(name, label)

        if name in self._infeasible_splits:
            self._infeasible_splits[name].update({label: code})
        else:
            self._infeasible_splits[name] = {label: code}

        self._check_rep()

    def mark_as_feasible(self, name, label):

        if self.is_infeasible(name, label):
            self._infeasible_splits[name].pop(label)
            if len(self._infeasible_splits[name]) == 0:
                self._infeasible_splits.pop(name)

            self._check_rep()

    def infeasibility_code(self, name, label):
        if self.is_infeasible(name, label):
            return str(self._infeasible_splits[name][label])
        else:
            return None

    def infeasibility_reason(self, name, label):
        if self.is_infeasible(name, label):
            code = self._infeasible_splits[name][label]
            return str(self.INFEASIBILITY_DICT[code])
        else:
            return None

    #### infeasibility setting ####

    def set_sample_infeasibility(self, propagate_to_group = True):

        for name, label in self.feasible_splits:
            n_samples = np.sum(self.split_indices(name, label))
            n_pos = np.sum(self.split_indices_pos(name, label))
            n_neg = np.sum(self.split_indices_neg(name, label))
            if (n_samples < self.min_samples) or (n_pos < self.min_samples_pos) or (n_neg < self.min_samples_neg):
                self.mark_as_infeasible(name, label, 'n_samples')

        if propagate_to_group:
            self.propagate_infeasibility()

        self._check_rep()

    def propagate_infeasibility(self):

        for name in self.infeasible_names:

            all_labels = set(self.labels(name))
            infeasible_labels = set(map(lambda s: s[1], filter(lambda s: s[0] == name, self.infeasible_splits)))
            feasible_labels = all_labels.difference(infeasible_labels)

            for label in feasible_labels:
                self.mark_as_infeasible(name, label, 'group_violation')

        self._check_rep()

    #### training ####

    def next(self):
        name, label, idx = None, None, None
        for n, l in self._training_order:
            if self.needs_model(n, l):
                name = n
                label = l
                idx = self.split_indices(name, label)
                break
        return name, label, idx

    def finished(self):
        return not any(map(lambda v: self.needs_model(v[0], v[1]), self._training_order))

class DecouplingTree(NodeMixin):

    ROOT_GROUP_NAME = 'root'
    ROOT_GROUP_LABEL = 'all'
    ROOT_GROUP_ID = 0
    TRAINING_HANDLE_ARGS = ['data']

    def __init__(self, data = None, groups = None, parent = None, **kwargs):

        self._halted = False
        self._split_attempt = False
        self._partition = None

        if parent is None:

            # check inputs
            assert check_data(data, ready_for_training = True)
            assert check_groups(groups, data)

            # initialize data
            self._data = {
                'X': np.array(data['X']),
                'Y': np.array(data['Y']),
                'variable_names' : list(data['variable_names'])
                }

            self._groups = deepcopy(groups)
            self._sample_indices_pos = self.data['Y'] == 1
            self._sample_indices_neg = self.data['Y'] == -1

            # set handles
            assert 'training_handle' in kwargs
            self.training_handle = kwargs.get('training_handle')
            self.scoring_handle = kwargs.get('scoring_handle', preference_score)
            self.selection_handle = kwargs.get('selection_handle', minimize_score_over_partitions)

            self._group_name = str(DecouplingTree.ROOT_GROUP_NAME)
            self._group_label = str(DecouplingTree.ROOT_GROUP_LABEL)
            self._node_indices = np.ones(self.data['X'].shape[0], dtype = bool)
            self._max_depth = len(self.groups)

            self._branch_log = []
            self._branch_iteration = 0
            self._log_branching = False

            # min sample sizes
            d = data['X'].shape[1] - has_intercept(data)
            min_samples = kwargs.get('min_samples', 2 * d)
            self._min_samples = min_samples

            min_samples_pos, min_samples_neg = kwargs.get('min_samples_pos', d), kwargs.get('min_samples_neg', d)
            assert isinstance(min_samples_pos, int) and min_samples_pos > 0
            assert isinstance(min_samples_neg, int) and min_samples_neg > 0
            self._min_samples_pos = min_samples_pos
            self._min_samples_neg = min_samples_neg

        else:

            assert all(k in kwargs for k in ['model', 'group_name', 'group_label'])
            group_name, group_label = kwargs.get('group_name'), kwargs.get('group_label')
            assert group_name in parent.root.groups
            assert group_label in parent.root.groups[group_name]['labels']

            self.parent = parent
            self._group_name = str(group_name)
            self._group_label = str(group_label)
            group_id = int(list(self.groups[group_name]['labels']).index(group_label))
            group_indices = self.groups[group_name]['indices'] == group_id
            self._node_indices = np.logical_and(self.parent.node_indices, group_indices)

        # model
        self._model = kwargs.get('model')

        # sample sizes
        self._n = np.sum(self.node_indices)
        self._n_pos = np.sum(self.node_indices_pos)
        self._n_neg = np.sum(self.node_indices_neg)

        # initialize split helper object
        self.split_helper = DecoupledTrainer(self.potential_split_groups, self.node_indices, self.node_indices_pos, self.node_indices_neg)
        self.split_helper.min_samples = self.min_samples
        self.split_helper.min_samples_pos = self.min_samples_pos
        self.split_helper.min_samples_neg = self.min_samples_neg
        self.check_rep()

    def check_rep(self):

        assert callable(self.training_handle)
        assert callable(self.scoring_handle)
        assert callable(self.selection_handle)

        if self.has_model:
            assert isinstance(self.model, ClassificationModel)

        assert self.min_samples_pos >= 1
        assert self.min_samples_neg >= 1
        assert self.min_samples >= self.min_samples_pos + self.min_samples_neg
        assert self.n >= self.min_samples
        assert self.n == self.n_pos + self.n_neg

        assert self.root.height <= self.max_depth
        assert self.depth <= self.max_depth
        if self.depth == self.max_depth:
            assert self.terminal

        if self.split_attempt:
            assert not self.terminal
            assert self.has_model
            assert self.halted or len(self.children) >= 2

        if self.halted:
            assert not self.terminal
            assert self.split_attempt
            assert self.has_model

    def check_tree(self):
        """
        checks the integrity of the entire tree
        :return:
        """

        # check tree must be run from the root
        if not self.is_root:
            return self.root.check_tree()

        for d in self.descendants:

            assert d.has_model, \
                'child nodes must have a model'

            assert d.group_name in self.groups, \
                "group_name must belong to groups"

            assert d.group_name not in [a.group_name for a in d.ancestors], \
                'group was previously branched on'

            assert d.group_label in self.groups[d.group_name]['labels'], \
                "group_label must belong to groups[group_name]['labels']"

            assert d.group_name not in [a.group_name for a in d.descendants], \
                'group is branched on later'

            if len(d.children) > 0:
                assert len(d.children) >= 2
                assert d.child_indices.shape == (len(d.node_indices), len(d.leaves))
                assert np.all(np.sum(d.child_indices, axis = 1) == d.node_indices)

            assert len(d.siblings) >= 1

            assert all(map(lambda s: d.group_name == s.group_name, d.siblings)), \
                'group names for siblings are not identical'

            assert all(map(lambda s: d.group_label != s.group_label, d.siblings)), \
                'group labels for siblings are not distinct'

        return True

    #### printing ####

    def __str__(self):
        return self.info

    def __repr__(self):
        return self.info

    @property
    def name(self):
        if self.is_root:
            s = str(DecouplingTree.ROOT_GROUP_NAME)
        else:
            s = []
            for name, label in self.partition_splits:
                s.append('%s:%s' % (name, label))
            s = ' & '.join(s)
        return s

    @property
    def info(self):

        info = [
            'node: %s' % self.name,
            'n (+/-): %d (%d,%d)' % (self.n, self.n_pos, self.n_neg),
            'children: %d' % len(self.children),
            'terminal: %s' % ('True' if self.terminal else 'False'),
            'has_model: %s' % ('True' if self.has_model else 'False'),
            ]

        if self.split_attempt and self.halted:
            info.extend(['split: halted'])
        elif self.split_attempt:
            info.extend(['split: attempted'])
        else:
            info.extend(['split: no attempt'])

        if self.has_model:

            mistakes = self.mistakes()
            error_rate = np.mean(mistakes)
            error_count = int(np.sum(mistakes))
            model_info = [
                '-' * 20,
                'train error: %1.1f%% (%d/%d)' % (100.0 * error_rate, error_count, self.n)
                ]

            info.extend(model_info)

            if len(self.children) > 0:
                X, Y = self.data['X'], self.data['Y']
                for c in self.children:
                    idx = c.node_indices
                    mistakes = self.mistakes(X[idx, :], Y[idx])
                    error_rate = np.mean(mistakes)
                    error_count = int(np.sum(mistakes))
                    child_info = ['-' * 5,
                                  'child group: %s' % c.group_label,
                                  'training error: %1.1f%% (%d/%d)' % (100.0 * error_rate, error_count, c.n)
                                  ]
                    info.extend(child_info)
        info = [' ', '=' * 20] + info + ['=' * 20]
        return '\n'.join(info)

    def render(self):
        if self.is_root:
            print(RenderTree(self).by_attr("_render_lines"))

    @property
    def _render_lines(self):

        group_string = '%s:%s' % (self._group_name, self._group_label)

        partition = self.partition
        if partition is not None:
            if self in partition:
                group_string = '[%s]' % group_string
            elif any([d in partition for d in self.descendants]):
                group_string = '%s' % group_string
            else:
                group_string = '#### %s' % group_string

        lines = ['%s (n = %d)' % (group_string, self.n)]

        if self.has_model:
            mistakes = self.mistakes()
            model_string = '-error: %1.1f%% (%d)' % (100.0 * np.mean(mistakes), int(np.sum(mistakes)))
            lines.append(model_string)

        if not self.is_root:
            parents = list(self.ancestors)
            parents.reverse()
            for p in parents:
                mistakes = p.mistakes(self.X, self.y)
                parent_string = '-error@%s: %1.1f%% (%d)' % (p.name, 100.0 * np.mean(mistakes), int(np.sum(mistakes)))
                lines.append(parent_string)

        lines.append('\n')
        return lines

    #### properties defined at ROOT for entire TREE ####

    @property
    def data(self):
        if self.is_root:
            return self._data
        else:
            return self.root.data

    @property
    def groups(self):
        if self.is_root:
            return self._groups
        else:
            return self.root.groups

    @property
    def sample_indices_pos(self):
        if self.is_root:
            return self._sample_indices_pos
        else:
            return self.root.sample_indices_pos

    @property
    def sample_indices_neg(self):
        if self.is_root:
            return self._sample_indices_neg
        else:
            return self.root.sample_indices_neg

    @property
    def min_samples_pos(self):
        if self.is_root:
            return int(self._min_samples_pos)
        else:
            return self.root.min_samples_pos

    @min_samples_pos.setter
    def min_samples_pos(self, n):
        if self.is_root:
            assert is_integer(n), 'min samples must be integer'
            assert n >= 2, 'min samples must be at least 2'
            self._min_samples_pos = min(int(n), self.n_pos)
        else:
            self.root.min_samples_pos(self, n)

    @property
    def min_samples_neg(self):
        if self.is_root:
            return int(self._min_samples_neg)
        else:
            return self.root.min_samples_neg

    @min_samples_neg.setter
    def min_samples_neg(self, n):
        if self.is_root:
            assert is_integer(n), 'min samples must be integer'
            assert n >= 2, 'min samples must be at least 2'
            self._min_samples_neg = min(int(n), self.n_neg)
        else:
            self.root.min_samples_neg(self, n)

    @property
    def min_samples(self):
        if self.is_root:
            return int(self._min_samples)
        else:
            return self.root.min_samples

    @min_samples.setter
    def min_samples(self, n):
        if self.is_root:
            assert is_integer(n), 'min samples must be integer'
            assert n >= 2, 'min samples must be at least 2'
            self._min_samples = min(int(n), self.n)
            self._min_samples = max(self._min_samples, self._min_samples_pos + self._min_samples_neg)
        else:
            self.root.min_samples(self, n)

    @property
    def max_depth(self):
        """
        :return: maximum depth of the tree
        """
        if self.is_root:
            return int(self._max_depth)
        else:
            return self.root.max_depth

    @max_depth.setter
    def max_depth(self, d):
        if self.is_root:
            assert is_integer(d), 'max depth must be integer'
            assert d >= 0, 'max depth must be non-negative'
            self._max_depth = min(int(d), len(self.groups))
        else:
            self.root.max_depth(self, d)

    @property
    def training_handle(self):
        """
        :return: function handle used for training
        """
        if self.is_root:
            return self._training_handle
        else:
            return self.root.training_handle

    @training_handle.setter
    def training_handle(self, handle):
        if not self.is_root:
            self.root.training_handle = handle
        else:
            if self.split_attempt:
                raise AttributeError('cannot change training handle after tree has already been branched')
            assert callable(handle)
            spec = getfullargspec(handle)
            for arg in self.TRAINING_HANDLE_ARGS:
                assert arg in spec.args
            self._training_handle = handle

    @property
    def scoring_handle(self):
        """
        :return: function handle used to score different partitions of the data
        """
        if self.is_root:
            return self._scoring_handle
        else:
            return self.root._scoring_handle

    @scoring_handle.setter
    def scoring_handle(self, handle):
        if not self.is_root:
            self.root.scoring_handle = handle
        else:
            if self.split_attempt:
                raise AttributeError('cannot change scoring handle after tree has already been branched')
            assert callable(handle)
            # spec = getfullargspec(handle)
            # for arg in self.SCORING_HANDLE_ARGS:
            #     assert arg in spec.args
            self._scoring_handle = handle

    @property
    def selection_handle(self):
        """
        :return: function handle used to branch_node by choosing partition among multiple feasible partitions
        """
        if self.is_root:
            return self._selection_handle
        else:
            return self.root.selection_handle

    @selection_handle.setter
    def selection_handle(self, handle):
        if not self.is_root:
            self.root.selection_handle = handle
        else:
            if self.split_attempt:
                raise AttributeError('cannot change selection handle after tree has already been branched')
            assert callable(handle)
            # spec = getfullargspec(handle)
            # for arg in self.SELECTION_HANDLE_ARGS:
            #     assert arg in spec.args
            self._selection_handle = handle

    #### properties defined for NODE at INITIALIZATION ####

    @property
    def n(self):
        return int(self._n)

    @property
    def n_pos(self):
        return int(self._n_pos)

    @property
    def n_neg(self):
        return int(self._n_neg)

    @property
    def X(self):
        return np.array(self.data['X'][self._node_indices, :])

    @property
    def y(self):
        return np.array(self.data['Y'][self._node_indices])

    @property
    def group_name(self):
        return str(self._group_name)

    @property
    def group_label(self):
        return str(self._group_label)

    @property
    def node_indices(self):
        return self._node_indices

    @property
    def node_indices_pos(self):
        return np.logical_and(self.node_indices, self.root.sample_indices_pos)

    @property
    def node_indices_neg(self):
        return np.logical_and(self.node_indices, self.root.sample_indices_neg)

    @property
    def child_indices(self):
        if self.is_leaf:
            return self.node_indices
        else:
            return np.transpose(list(map(lambda l: l.node_indices, self.leaves)))

    @property
    def path_nodes(self):
        node_list = []
        node = self
        while True:
            node_list.append(node)
            if node.is_root:
                break
            else:
                node = node.parent
        return node_list

    #### TREE METHODS ####

    @property
    def leaves(self):
        if self.is_leaf:
            return [self]
        else:
            return list(filter(lambda n: n.is_leaf, self.descendants))

    @property
    def other_leaves(self):
        if self.is_leaf:
            return list(filter(lambda n: n != self, self.root.leaves))
        else:
            return self.leaves

    def subtrees(self):

        """
        returns list of all pointers to all possible subtrees
        :return:
        """

        assert self.is_root
        initial_tree = DecoupledLeafSet(leaf_set = self.leaves)
        subtree_list = [initial_tree]
        stack = [initial_tree]
        while stack:
            t = stack.pop()
            for s in t.collapsed_trees:
                if s not in subtree_list:
                    subtree_list.append(s)
                    stack.append(s)

        assert all(map(lambda s: s.has_models, subtree_list))
        assert all(map(lambda s: s.is_partition, subtree_list))

        return subtree_list

    #### paths ####
    @property
    def path_splits(self):
        return list(map(lambda n: (n.group_name, n.group_label), self.path))

    @property
    def path_labels(self):
        return list(map(lambda n: n.group_label, self.path))

    @property
    def path_names(self):
        return list(map(lambda n: n.group_name, self.path))

    @property
    def partition_splits(self):
        return list(map(lambda n: (n.group_name, n.group_label), self.path[1:]))

    #### model training at the node ####
    def filter_data(self, indices = None):

        if indices is None:
            return self.data
        else:
            return {
                'X': np.array(self.data['X'][indices, :]),
                'Y': np.array(self.data['Y'][indices]),
                'variable_names': list(self.data['variable_names'])
                }

    def contains(self, group_names, group_values):
        """
        return indices for samples at this node that match a given group
        :param group_names: names of the group attributes (e.g., AgeGroup)
        :param group_values: values of the group attributes (e.g., Old)
        :return: indices of sample points for individuals in this group (for this node)
        """
        if self.is_root:
            idx = np.ones(group_values.shape[0], dtype = bool)
        else:
            names, labels = zip(*self.partition_splits)
            common_name_idx = [group_names.index(n) for n in names]
            idx = np.all(group_values[:, common_name_idx] == labels, axis = 1)
        return idx

    @property
    def has_model(self):
        return self._model is not None

    def train_model(self, overwrite = False):
        if (not self.has_model) or overwrite:
            self._model = self.training_handle(data = self.filter_data(indices = self.node_indices))
            assert self.has_model

    @property
    def model(self):
        """
        :return: model at this node
        """
        if self.has_model:
            return deepcopy(self._model)
        else:
            raise ValueError("no model at node")

    def predict(self, X):
        """
        predict using the model at this node
        :param X:
        :return:
        """
        if self.has_model:
            return self.model.predict(X)
        else:
            raise ValueError("no model at node")

    def mistakes(self, X = None, y = None):

        if not self.has_model:
            raise ValueError("no model at node")

        if X is None and y is None:
            mistakes = np.not_equal(self.y, self.model.predict(self.X))
        else:
            assert X is not None and y is not None
            assert isinstance(X, np.ndarray)
            assert isinstance(y, np.ndarray)
            assert len(y) == X.shape[0]
            mistakes = np.not_equal(y.flatten(), self.model.predict(X))

        return mistakes

    def error(self, X = None, y = None):
        return np.mean(self.mistakes(X, y))

    #### groups that can be split on ####

    @property
    def potential_split_names(self):
        """
        returns names of groups that can still be branched on at this node
        """
        return set(self.root.groups.keys()).difference(self.path_names)

    @property
    def potential_split_groups(self):
        return {d: self.root.groups[d] for d in self.potential_split_names}

    @property
    def has_potential_splits(self):
        """
        :return: True if there exists any group that can be split on at this node
        """
        return len(self.potential_split_names) > 0

    @property
    def feasible_split_names(self):
        self.split_helper.propagate_infeasibility()
        return self.split_helper.feasible_names

    @property
    def feasible_split_groups(self):
        return {d: self.root.groups[d] for d in self.feasible_split_names}

    @property
    def has_feasible_splits(self):
        return len(self.split_helper.feasible_names) > 0

    #### branching state ####
    @property
    def split_attempt(self):
        return bool(self._split_attempt)

    @split_attempt.setter
    def split_attempt(self, flag):
        assert isinstance(flag, bool)
        if flag:
            assert self.has_model
        self._split_attempt = bool(flag)

    @property
    def halted(self):
        return self._halted

    @halted.setter
    def halted(self, flag):
        assert isinstance(flag, bool)
        if flag:
            assert self.has_model
            assert self.split_attempt
        self._halted = bool(flag)

    @property
    def is_fully_grown(self):
        return not any(map(lambda n: n.splittable, self.leaves))

    @property
    def splittable(self):
        return self.is_leaf and (not self.terminal) and (not self.halted)

    @property
    def terminal(self):
        return (self.depth == self.max_depth) or (not self.has_potential_splits)

    ##### branch logging ####

    @property
    def branch_iteration(self):
        return int(self._branch_iteration)

    @branch_iteration.setter
    def branch_iteration(self, k):
        assert isinstance(k, int)
        assert k > self._branch_iteration
        self._branch_iteration = k

    @property
    def log_branching(self):
        if self.is_root:
            return self._log_branching
        else:
            return self.root.log_branching

    @log_branching.setter
    def log_branching(self, flag):
        if self.is_root:
            assert isinstance(flag, bool)
            self._log_branching = flag
        else:
            self.root.log_branching = flag

    @property
    def branch_log(self):
        if self.is_root:
            return self._branch_log
        else:
            return self.root.branch_log

    ##### branching  ####

    def train_split_models(self):

        sh = self.split_helper
        sh.set_sample_infeasibility(propagate_to_group = True)

        while not sh.finished():
            name, label, idx = sh.next()
            model = self.training_handle(data = self.filter_data(indices = idx))
            self.split_helper.set_model(name, label, model)

    def branch(self):

        """ split one node after checking all possible splits in the current tree """
        assert self.is_root

        # for each leaf, and each split, train all models
        splittable_leaves = [l for l in self.leaves if l.splittable]
        for l in splittable_leaves:
            l.train_split_models()
            l.split_attempt = True
            l.halted = len(l.feasible_split_groups) == 0

        if self.is_fully_grown:
            return

        # recompute splittable since some leaves may have been halted
        splittable_leaves = [l for l in self.leaves if l.splittable]

        # score each split based on envy-freeness of the partition
        potential_splits = []
        pooled_model = self.root.predict
        for leaf in splittable_leaves:

            other_leaves = [l for l in self.leaves if l != leaf]
            for group_name, group_info in leaf.feasible_split_groups.items():

                split_leaves = []
                for i, label in zip(group_info['ids'], group_info['labels']):

                    split_leaf = CandidateNode(model = leaf.split_helper.get_model(group_name, label),
                                               indices = np.logical_and(leaf.node_indices, i == group_info['indices']),
                                               group_names = leaf.partition_names + [group_name],
                                               group_labels = leaf.partition_labels + [label],
                                               parent = self)

                    split_leaves.append(split_leaf)

                partition = DecoupledLeafSet(leaf_set = split_leaves + other_leaves, node_type = NodeMixin)
                partition.score = self.scoring_handle(data = self.data, partition = partition, pooled_model = pooled_model)

                split_info = {
                    'leaf': leaf,
                    'group_name': group_name,
                    'partition': partition,
                    'score': partition.score,
                    }

                potential_splits.append(split_info)

        # choose best tree
        best_idx = self.selection_handle(potential_splits)
        split_leaf = potential_splits[best_idx]['leaf']
        split_name = potential_splits[best_idx]['group_name']

        # add child nodes to tree
        for split_label in self.groups[split_name]['labels']:
            DecouplingTree(parent = split_leaf,
                           group_name = split_name,
                           group_label = split_label,
                           model = split_leaf.split_helper.get_model(split_name, split_label))

        # update history
        self.branch_iteration = self.branch_iteration + 1
        if self.log_branching:
            potential_splits[best_idx]['partition'].selected = True
            for s in potential_splits:
                p = s['partition']
                p.order = self.branch_iteration
                self._branch_log += [{'iteration': self.branch_iteration, 'score': s['partition'].score, 'partition': s['partition']}]

    def grow(self, print_flag = True, log_flag = True):
        """
        branch_node recursively until fully trained
        :return:
        """
        self.log_branching = log_flag
        self.train_model()

        while not self.is_fully_grown:
            self.branch()
            if print_flag:
                self.render()
            self.check_tree()

    ### pruning methods ####

    def prune(self, alpha = 0.1, test_correction = True, atomic_groups = True, objectives = ['group_gain_min', 'group_gain_max', 'group_gain_med'], **data_args):

        assert 0.0 <= alpha <= 1.0
        stats = []
        subtrees = [t for t in self.subtrees() if len(t) >= 2]

        for p in subtrees:

            if atomic_groups:
                training_sizes = p.partition_sample_size_stats(parent_type = 'self', metric_type = 'n', **data_args)
                group_sizes = p.group_sample_size_stats(metric_type = 'n', **data_args)
                gain = p.group_decoupling_stats(metric_type = 'error_gap', parent_type = 'root', **data_args)
                pvals_gain = p.group_decoupling_stats(metric_type = 'pvalue', parent_type = 'root', **data_args)
                pvals_switch = p.group_switch_stats(metric_type = 'pvalue', **data_args)
            else:
                training_sizes = p.partition_sample_size_stats(parent_type = 'self', metric_type = 'n', **data_args)
                group_sizes = training_sizes
                gain = p.partition_decoupling_stats(metric_type = 'error_gap', parent_type = 'root', **data_args)
                pvals_gain = p.partition_decoupling_stats(metric_type = 'pvalue', parent_type = 'root', **data_args)
                pvals_switch = p.partition_stat_matrix(metric_type = 'pvalue', **data_args)

            # in case some groups aren't observed
            group_weights = group_sizes / np.sum(group_sizes)
            keep_idx = np.greater(group_sizes, 0)
            gain = gain[keep_idx]
            group_weights = group_weights[keep_idx]

            # compute correct alpha
            n_tests = np.isfinite(pvals_switch).sum() + np.isfinite(pvals_gain).sum()
            alpha_test = alpha / n_tests if test_correction else alpha

            stats.append({
                'name': p.name,
                'n_leaves': len(p),
                'alpha_raw': alpha,
                'n_tests': n_tests,
                'alpha': alpha / n_tests,
                #
                'p_violation_min': np.minimum(np.nanmin(pvals_gain), np.nanmin(pvals_switch)),
                'violations_switch': np.less_equal(pvals_switch, alpha_test, where = np.isfinite(pvals_switch)).sum(),
                'violations_decouple': np.less_equal(pvals_gain, alpha_test).sum(),
                #
                'weighted_gain': group_weights.dot(gain),
                #
                'group_gain_min': np.min(gain),
                'group_gain_med': np.median(gain),
                'group_gain_max': np.max(gain),
                #
                'sample_size_min': np.min(training_sizes),
                'sample_size_med': np.median(training_sizes),
                'sample_size_max': np.max(training_sizes),
                })

        # prune
        df = pd.DataFrame(stats)
        df['pruned'] = df['alpha'] >= df['p_violation_min']
        pruning_df = df.query('p_violation_min >= alpha').sort_values(objectives, ascending = False)
        self.partition = subtrees[pruning_df.index[0]]

        # return alternative trees
        candidate_trees = []
        pooled_model = self.root.model
        data = data_args.get('data', self.root.data)
        groups = data_args.get('groups', self.root.groups)
        for t in subtrees:
            dcs = DecoupledClassifierSet(data = data, groups = groups, pooled_model = pooled_model, decoupled_models = [l.model for l in t.leaves], groups_to_models = build_model_assignment_map(t))
            candidate_trees.append(dcs)

        return df, candidate_trees

    @property
    def partition(self):
        if self.is_root:
            return self._partition
        else:
            return self.root.partition

    @partition.setter
    def partition(self, partition):
        if self.is_root:
            assert isinstance(partition, DecoupledLeafSet)
            assert partition.has_models
            assert partition.is_partition
            assert partition.node_type == DecouplingTree

            # check that we have a valid partition
            for l in partition.leaves:
                assert l in self.root.descendants

            self._partition = DecoupledLeafSet(partition.leaves)
        else:
            self.root.partition = partition

    ### todo remove

    def predict_at_partition(self, X, group_names = None, group_values = None):
        """
        predict using the leaves from this node

        :param X:
        :param group_names:
        :param group_values:
        :return:
        """
        if self.partition is None:
            leaves = self.leaves
        else:
            leaves = self._partition.leaves

        assert group_names is not None and group_values is not None
        assert all(map(lambda l: map(lambda n: n in group_names, l.partition_names), leaves))

        yhat = np.repeat(np.nan, X.shape[0])
        for l in leaves:
            idx = l.contains(group_names, group_values)
            yhat[idx] = l.predict(X[idx, ])
        assert np.all(np.isfinite(yhat))
        return yhat

    @property
    def partition_names(self):
        return list(map(lambda n: n.group_name, self.path[1:]))

    @property
    def partition_labels(self):
        return list(map(lambda n: n.group_label, self.path[1:]))

class CandidateNode(NodeMixin):

    def __init__(self, model, indices, group_names, group_labels, parent = None):

        self._group_names = group_names
        self._group_labels = group_labels
        self._model = model
        self._indices = indices

        assert parent is None or isinstance(parent, DecouplingTree)
        self._parent_node = parent

    def contains(self, group_names, group_values):
        common_names = [group_names.index(n) for n in self._group_names]
        return np.all(group_values[:, common_names] == self._group_labels, axis = 1)

    @property
    def node_indices(self):
        return np.array(self._indices)

    @property
    def path_nodes(self):
        node_list = [self] + self._parent_node.path_nodes
        return node_list

    def predict(self, X):
        return self._model.predict(X)

    @property
    def model(self):
        return self._model

    @property
    def has_model(self):
        return True

    @property
    def partition_splits(self):
        return [(n, l) for n,l in zip(self._group_names, self._group_labels)]

    @property
    def partition_names(self):
        return self._group_names

    @property
    def partition_labels(self):
        return self._group_labels

    @property
    def name(self):
        s = []
        for name, label in zip(self._group_names, self._group_labels):
            s.append('%s:%s' % (name, label))
        s = ' & '.join(s)
        return s

class DecoupledLeafSet(object):

    _CONTAINMENT_CODE_MISSING = -1
    _CONTAINMENT_CODE_SUBSET = 0
    _CONTAINMENT_CODE_EXACT = 1
    _COMPARISON_METRIC_TYPES = {'error', 'error_gap', 'mistakes', 'mistakes_gap', 'pvalue'}
    _PARENT_TYPES = ['root', 'next', 'best', 'self']

    def __init__(self, leaf_set, node_type = DecouplingTree):

        """
        Wrapper class to represent and manipulate subtrees from a tree
        Must be initialized with a list of leaves for a tree (leaf_set), or a valid tree object (tree)

        :param leaf_set: list containing the leaves of a tree (must be distinct)
        :param node_type: class of each node in the tree (either NodeMixin, or EnvyFreeTree)
        """
        assert isinstance(node_type, object)
        assert isinstance(leaf_set, list)
        self._score = float('nan')
        self._order = -1
        self._selected = False
        self._node_type = node_type
        self._leaf_set = list(leaf_set)
        assert self._check_rep()

    def __len__(self):
        return len(self._leaf_set)

    def __str__(self):
        s = ['EnvyFreePartition (%d parts)' % len(self)]
        s += ['%s' % l.name for l in self._leaf_set]
        return '\n'.join(s)

    def __repr__(self):
        return '{%s}' % ','.join('%s' % l.name for l in self._leaf_set)

    def __eq__(self, other):
        return set(self._leaf_set) == set(other.leaves)

    def __contains__(self, node):
        return any([node == l for l in self._leaf_set])

    def _check_rep(self):

        leaf_set = self._leaf_set
        assert isinstance(leaf_set, list)
        n = len(leaf_set)
        assert n > 0, 'leaf set must be non-empty'
        assert n == len(list(set(leaf_set))), 'leaf set is not unique'

        if self.node_type == DecouplingTree:

            for l in leaf_set:
                assert isinstance(l, self.node_type)

            if n == 1:
                assert leaf_set[0].is_root

        elif self.node_type == CandidateNode:

            for l in leaf_set:
                assert l.is_root
                assert isinstance(l, self.node_type)

            assert not self.attached

        elif self.node_type == NodeMixin:

            if n >= 2:
                for l in leaf_set:
                    if isinstance(l, CandidateNode):
                        assert l.is_root
                    elif isinstance(l, DecouplingTree):
                        assert not l.is_root

            assert not self.attached

        return True

    @property
    def name(self):
        return ','.join('%s' % l.name for l in self._leaf_set)

    @property
    def split_names(self):
        return [l.name for l in self._leaf_set]

    @property
    def score(self):
        return self._score

    @score.setter
    def score(self, s):
        assert isinstance(s, (float, int))
        assert np.isfinite(s)
        self._score = s

    @property
    def order(self):
        return self._order

    @order.setter
    def order(self, k):
        assert isinstance(k, int)
        assert k >= 0
        self._order = int(k)

    @property
    def node_type(self):
        return self._node_type

    @property
    def selected(self):
        return self._selected

    @selected.setter
    def selected(self, flag):
        assert isinstance(flag, bool)
        self._selected = True

    @property
    def attached(self):
        check_type = all([isinstance(l, DecouplingTree) for l in self._leaf_set])
        root_name = DecouplingTree.ROOT_GROUP_NAME
        if len(self) == 1:
            check_root = self._leaf_set[0].name == root_name
        else:
            check_root = all([l.root.name == root_name for l in self._leaf_set])
        return check_type and check_root

    @property
    def leaves(self):
        """
        :return: list containing pointers to the leaf nodes of the tree
        """
        return list(self._leaf_set)

    @property
    def indices(self):
        return [l.node_indices for l in self._leaf_set]

    @property
    def is_partition(self):
        """
        :return: True iff the current set of leaves represents a valid partition
        """
        return np.sum(list(map(lambda l: l.node_indices, self._leaf_set)), axis = 0).all()

    @property
    def has_models(self):
        """
        :return: True iff the current set of leaves contains fully trained models
        """
        return all(map(lambda l: l.has_model, self._leaf_set))

    @property
    def parent_nodes(self):
        return list({l.parent for l in self._leaf_set})

    @property
    def collapsable_nodes(self):
        nodes = []
        for p in self.parent_nodes:
            if set(p.children).issubset(self._leaf_set):
                nodes.append(p)
        return nodes

    def collapse(self, parent_node):
        """
        :param parent_node:
        :return:
        """
        assert isinstance(parent_node, self.node_type)
        assert parent_node in self.collapsable_nodes
        collapsed_leaf_set = [l for l in self._leaf_set if l.parent != parent_node]
        collapsed_leaf_set.append(parent_node)
        return DecoupledLeafSet(collapsed_leaf_set)

    @property
    def collapsed_trees(self):
        if len(self) >= 2:
            return [self.collapse(p) for p in self.collapsable_nodes]
        else:
            return self.leaves

    @property
    def predictors(self):
        return [l.predict for l in self._leaf_set]

    def predict(self, X, group_names, group_values):
        """
        predict using all of the leaves in this partition
        :param X:
        :param group_names:
        :param group_values:
        :return:
        """
        assert all(map(lambda l: map(lambda n: n in group_names, l.partition_names), self._leaf_set))
        yhat = np.repeat(np.nan, X.shape[0])
        for l in self._leaf_set:
            idx = l.contains(group_names, group_values)
            yhat[idx] = l.predict(X[idx, ])
        assert np.all(np.isfinite(yhat))
        return yhat

    #### MATCHING AND ANALYSIS

    @property
    def data(self):
        assert self.node_type == DecouplingTree
        return self._leaf_set[0].root.data

    @property
    def groups(self):
        assert self.node_type == DecouplingTree
        return self._leaf_set[0].root.groups

    def splits(self, groups = None):
        """
        :param groups: group data structure
        :return: list of splits, where each split = [(group_name, group_label)] for all possible subgroups
        """
        if groups is None:
            groups = self.groups
        splits = [[(k, l) for l in v['labels']] for k, v in groups.items()]
        return list(itertools.product(*splits))

    def _parse_data_args(self, **kwargs):
        """
        helper function to parse X, y, group_names, and group_values
        from keyword inputs to different methods

        :param kwargs:
        :return:
        """
        if len(kwargs) == 0:
            return to_group_data(data = self.data, groups = self.groups, stat_field = 'train')

        elif set(['X', 'y', 'group_names', 'group_values']).issubset(kwargs):
            return kwargs['X'], kwargs['y'], kwargs['group_names'], kwargs['group_values']

        elif set(['data', 'groups']).issubset(kwargs):
            return to_group_data(data = kwargs['data'], groups = kwargs['groups'], stat_field = kwargs.get('stat_field') or 'train')

        else:
            raise ValueError('unsupport input arguments')

    def _parse_split(self, *args, **kwargs):
        """
        helper function to parse a split from arguments or keyword arguments to different methods
        split(G, v) == (G[1] = v[1], G[2] = v[2]... ) <- set identifying a specific groups
        :param args:
        :param kwargs:
        :return: split as a set
        """

        assert len(args) == 0 or len(kwargs) == 0
        split = None

        if len(args) == 1:
            split = args[0]
        elif len(args) == 2:
            assert len(args[0]) == len(args[1])
            split = zip(args[0], args[1])
        elif 'split' in kwargs:
            split = kwargs['split']
        elif ['group_names', 'group_labels']:
            split = zip(kwargs['group_names'], kwargs['group_labels'])

        if split is None:
            raise ValueError('could not parse split from args or kwargs')

        split = set(split)

        return split

    def _drop_missing_splits(self, splits, group_values):
        new = []
        for s in splits:
            group_labels = [t[1] for t in s]
            if np.all(group_values == group_labels, axis = 1).any():
                new.append(s)
        return new

    def match(self, *args, **kwargs):
        """
        V[i] = True <-> leaf i contains samples from split
        :param split
        :param group_names:
        :param group_labels:
        :return: a matching vector V such that

        -1 if group is not contained in a leaf
         0 if group is contained along with other groups in a leaf
         1 if group has own leaf
         k if group is contained in k leaves
        """
        split = self._parse_split(*args, *kwargs)
        return [split.issuperset(l.partition_splits) for l in self._leaf_set]

    def matched_leaves(self, *args, **kwargs):
        """
        return True if this node contains any samples for a group with group_labels
        :param group_names: names of the group attributes (e.g., AgeGroup)
        :param group_values: matrix of group attributes (e.g., Old)
        :param group_labels: labels that must match
        :return: indices of sample points for individuals in this group (for this node)
        """
        return list(itertools.compress(self._leaf_set, self.match(*args, **kwargs)))

    def match_code(self, *args, **kwargs):
        """
        :param group_names:
        :param group_labels:
        :return:

        -1 if group is not contained in a leaf
         0 if group is contained along with other groups in a leaf
         1 if group has own leaf
         k if group is contained in k leaves
        """
        split = self._parse_split(*args, **kwargs)
        match_exact = [split == set(l.partition_splits) for l in self._leaf_set]
        if any(match_exact):
            assert np.sum(match_exact) == 1
            return self._CONTAINMENT_CODE_EXACT

        match_subset = [split.issuperset(l.partition_splits) for l in self._leaf_set]
        match_subset = np.logical_and(match_subset, np.logical_not(match_exact))
        n_matches = np.sum(match_subset)
        if n_matches == 0:
            return self._CONTAINMENT_CODE_MISSING
        elif n_matches == 1:
            return self._CONTAINMENT_CODE_SUBSET
        else:
            return n_matches

    def partition_stat_matrix(self, metric_type = 'error', **kwargs):

        """
        return a matrix computing

        :param metric_type: 'error', 'error_gap', 'mistakes', 'mistake_gap', 'pvalue'

        :param kwargs: can be either:
            - data, groups
            - data, groups, stat_field
            - X, y, group_names, group_values

        :return: n_parts x n_parts matrix M of statistics where:

        - n_parts = # of leaves in partition
        - M[i, j] = value of metric on data in leaf i using model from leaf j
        """

        assert metric_type in DecoupledLeafSet._COMPARISON_METRIC_TYPES
        X, y, group_names, group_values = self._parse_data_args(**kwargs)
        n_parts = len(self)
        M = np.empty(shape = (n_parts, n_parts))

        # fast return for comparison based metrics
        is_root = len(self) == 1 and (self._leaf_set[0].name == DecouplingTree.ROOT_GROUP_NAME)
        if is_root and metric_type not in ['error', 'mistakes']:
            return np.array([np.nan])

        if metric_type in ['error', 'error_gap']:

            for i, leaf in enumerate(self._leaf_set):
                idx = leaf.contains(group_names, group_values)
                ys, Xs = y[idx], X[idx, :]
                for j, other in enumerate(self._leaf_set):
                    M[i, j] = np.not_equal(ys, other.predict(Xs)).mean()

        elif metric_type in ['mistakes', 'mistakes_gap']:

            for i, leaf in enumerate(self._leaf_set):
                idx = leaf.contains(group_names, group_values)
                ys, Xs = y[idx], X[idx, :]
                for j, other in enumerate(self._leaf_set):
                    M[i, j] = np.not_equal(ys, other.predict(Xs)).sum()

        elif metric_type == 'pvalue':

            for i, leaf in enumerate(self._leaf_set):
                idx = leaf.contains(group_names, group_values)
                ys, Xs = y[idx], X[idx, :]
                hs = leaf.predict(X[idx, :])
                for j, other in enumerate(self._leaf_set):
                    if i != j:
                        M[i, j], _ = exact_mcn_test(y = ys, yhat1 = hs, yhat2 = other.predict(Xs))

            np.fill_diagonal(M, np.nan)

        if '_gap' in metric_type:

            M = M - np.diag(M)[:, None]
            np.fill_diagonal(M, np.nan)

        return M

    def partition_sample_size_stats(self, metric_type = 'n', **kwargs):

        _, y, group_names, group_values = self._parse_data_args(**kwargs)
        S = np.zeros(len(self))

        for i, l in enumerate(self._leaf_set):
            idx = l.contains(group_names, group_values)
            if np.any(idx):
                if metric_type in ('n', 'p'):
                    S[i] = np.sum(idx)
                elif metric_type in ('n_pos', 'p_pos'):
                    S[i] = np.sum(y[idx] == 1)
                elif metric_type in ('n_neg', 'p_neg'):
                    S[i] = np.sum(y[idx] == -1)

        assert len(y) == np.sum(S)

        if metric_type in ('p', 'p_pos', 'p_neg'):
            S = S / len(y)

        return S

    def partition_decoupling_stats(self, metric_type = 'error_gap', parent_type = 'root', **kwargs):

        assert parent_type in DecoupledLeafSet._PARENT_TYPES
        X, y, group_names, group_values = self._parse_data_args(**kwargs)
        S = np.repeat(np.nan, len(self))

        # fast return for comparison based metrics
        no_parents = parent_type == 'self'
        is_root = len(self) == 1 and (self._leaf_set[0].name == DecouplingTree.ROOT_GROUP_NAME)
        no_comparison_node = is_root or no_parents
        if no_comparison_node and metric_type not in ['error', 'mistakes']:
            return S

        for i, l in enumerate(self._leaf_set):

            idx = l.contains(group_names, group_values)
            ys, Xs = y[idx], X[idx, :]
            predictions = [n.predict(Xs) for n in l.path_nodes]

            if metric_type == 'error':

                errors = [np.not_equal(ys, yhat).mean() for yhat in predictions]
                if parent_type == 'self':
                    metric_value = errors.pop(0)
                elif parent_type == 'root':
                    metric_value = errors.pop(-1)
                elif parent_type == 'next':
                    metric_value = errors.pop(1)
                elif parent_type == 'best':
                    metric_value = np.min(errors)

            elif metric_type in ['error_gap', 'error_relgap']:

                errors = [np.not_equal(ys, yhat).mean() for yhat in predictions]
                base_value = errors.pop(0)
                if parent_type == 'root':
                    comp_value = errors.pop(-1)
                elif parent_type == 'next':
                    comp_value = errors.pop(0)
                elif parent_type == 'best':
                    comp_value = np.min(errors)

                metric_value = comp_value - base_value
                if '_relgap' in metric_type:
                    metric_value = (comp_value - base_value) / base_value

            elif metric_type == 'pvalue':

                yhat1 = predictions.pop(0)
                pvals = [exact_mcn_test(ys, yhat1, yhat)[0] for yhat in predictions]
                if parent_type == 'root':
                    metric_value = pvals[-1]
                elif parent_type == 'next':
                    metric_value = pvals[0]
                elif parent_type == 'best':
                    metric_value = np.min(pvals)

            S[i] = metric_value

        return S

    def group_sample_size_stats(self, metric_type = 'n', drop_missing = False, **kwargs):
        _, y, group_names, group_values = self._parse_data_args(**kwargs)
        splits = self.splits(kwargs.get('groups'))
        if drop_missing:
            splits = self._drop_missing_splits(splits, group_values)

        S = np.repeat(np.nan, len(splits))
        for i, s in enumerate(splits):
            group_labels = [t[1] for t in s]
            idx = np.all(group_values == group_labels, axis = 1)
            if metric_type in ('n', 'p'):
                S[i] = np.sum(idx)
            elif metric_type in ('n_pos', 'p_pos'):
                S[i] = np.sum(y[idx] == 1)
            elif metric_type in ('n_neg', 'p_neg'):
                S[i] = np.sum(y[idx] == -1)
        if metric_type in ('p', 'p_pos', 'p_neg'):
            S = S / len(y)

        return S

    def group_switch_stats(self, metric_type = 'error_gap', drop_missing = False, **kwargs):

        X, y, group_names, group_values = self._parse_data_args(**kwargs)
        splits = self.splits(kwargs.get('groups'))
        if drop_missing:
            splits = self._drop_missing_splits(splits, group_values)

        n_groups, n_models = len(splits), len(self)
        M = np.tile(np.nan, [n_groups, len(self)])
        for i, s in enumerate(splits):
            group_labels = [t[1] for t in s]
            idx = np.all(group_values == group_labels, axis = 1)
            if any(idx):

                ys, Xs = y[idx], X[idx, :]
                assign_idx = np.array(self.match(s))
                match_idx = int(np.flatnonzero(assign_idx))
                other_idx = np.flatnonzero(np.logical_not(assign_idx))

                h = self._leaf_set[match_idx]
                ys_hat = h.predict(Xs)

                for k in other_idx:
                    yk_hat = self._leaf_set[int(k)].predict(Xs)

                    if metric_type == 'pvalue':
                        M[i, k] = exact_mcn_test(ys, ys_hat, yk_hat)[0]

                    elif metric_type == 'error':
                        M[i, k] = np.not_equal(ys, yk_hat).mean()

                    elif metric_type == 'error_gap':
                        M[i, k] = np.not_equal(ys, yk_hat).mean() - np.not_equal(ys, ys_hat).mean()

                    elif metric_type == 'error_relgap':
                        M[i, k] = (np.not_equal(ys, yk_hat).mean() - np.not_equal(ys, ys_hat).mean()) / np.not_equal(ys, ys_hat).mean()

                if metric_type == 'error':
                    M[i, match_idx] = np.not_equal(ys, ys_hat).mean()

                elif metric_type == 'error_relgap':
                    M[i, match_idx] = 0.0

        return M

    def group_decoupling_stats(self, metric_type = 'error_gap', parent_type = 'root',  drop_missing = False, **kwargs):

        assert parent_type in DecoupledLeafSet._PARENT_TYPES
        X, y, group_names, group_values = self._parse_data_args(**kwargs)
        splits = self.splits(kwargs.get('groups'))
        if drop_missing:
            splits = self._drop_missing_splits(splits, group_values)

        leaves = self._leaf_set
        S = np.repeat(np.nan, len(splits))

        if parent_type == 'self':

            if metric_type == 'error':
                for i, s in enumerate(splits):
                    group_labels = [t[1] for t in s]
                    idx = np.all(group_values == group_labels, axis = 1)
                    ys, Xs = y[idx], X[idx, :]
                    yhat = self.predict(Xs, group_names, group_values[idx, :])
                    S[i] = np.not_equal(ys, yhat).mean()

        elif parent_type == 'root':

            for i, s in enumerate(splits):

                group_labels = [t[1] for t in s]
                idx = np.all(group_values == group_labels, axis = 1)
                ys, Xs = y[idx], X[idx, :]

                yhat = self.predict(Xs, group_names, group_values[idx, :])
                yhat_root = leaves[0].root.predict(Xs)

                if metric_type == 'pvalue':
                    S[i], _ = exact_mcn_test(ys, yhat, yhat_root)

                elif metric_type in ['error']:
                    S[i] = np.not_equal(ys, yhat_root).mean()

                elif metric_type in ['error_gap', 'error_relgap']:
                    base_value = np.not_equal(ys, yhat).mean()
                    comp_value = np.not_equal(ys, yhat_root).mean()
                    metric_value = comp_value - base_value
                    if '_relgap' in metric_type:
                        metric_value = metric_value / base_value
                    S[i] = metric_value

        return S


from gevent import sleep
from collections import deque
from mrkt.framework.role.task import Task
from mrkt.framework.role import Cluster
from mrkt.common.consts import *

class Pool(Cluster):
    def __init__(self, *args, **kwargs):
        super(Pool, self).__init__(*args, **kwargs)
        self.task_queue = deque()

    def schedule(self):
        while True:
            for worker in self.workers:
                if self.need_sync(worker):
                    self.sync_worker(worker)
                else:
                    while worker.utilization() < 1:
                        if self.task_queue:
                            task = self.task_queue.popleft()
                            self.processing_tasks.append(task)
                            task.assign_to(worker)
                        else:
                            break
            sleep(CLUSTER_SCHEDULE_INTERVAL)

    def submit(self, func, *args, **kwargs):
        task = Task(func, args, kwargs)
        self.task_queue.append(task)
        return task

    def map(self, func, *iterables):
        tasks = [self.submit(func, *args) for args in zip(*iterables)]
        for task in tasks:
            task.join()
            yield task.ret


# -*- coding: utf-8 -*-
from __future__ import unicode_literals

from django.db import models, migrations

class Migration(migrations.Migration):

    dependencies = [
        ('jtx_video', '0004_auto_20150805_2247'),
    ]

    operations = [
        migrations.CreateModel(
            name='BaseFile',
            fields=[
                ('id', models.AutoField(verbose_name='ID', auto_created=True, serialize=False, primary_key=True)),
                ('filename', models.CharField(max_length=254)),
                ('extension', models.CharField(blank=True, max_length=10)),
            ],
        ),
        migrations.CreateModel(
            name='Folder',
            fields=[
                ('basefile_ptr', models.OneToOneField(primary_key=True, parent_link=True, serialize=False, auto_created=True, to='jtx_video.BaseFile')),
            ],
            bases=('jtx_video.basefile',),
        ),
        migrations.CreateModel(
            name='SubtitleFile',
            fields=[
                ('basefile_ptr', models.OneToOneField(primary_key=True, parent_link=True, serialize=False, auto_created=True, to='jtx_video.BaseFile')),
                ('video', models.OneToOneField(blank=True, to='jtx_video.Video', related_name='subtitles')),
            ],
            bases=('jtx_video.basefile',),
        ),
        migrations.CreateModel(
            name='VideoFile',
            fields=[
                ('basefile_ptr', models.OneToOneField(primary_key=True, parent_link=True, serialize=False, auto_created=True, to='jtx_video.BaseFile')),
                ('bitrate', models.FloatField(blank=True)),
                ('width', models.IntegerField(blank=True)),
                ('height', models.IntegerField(blank=True)),
                ('video', models.ForeignKey(blank=True, to='jtx_video.Video', related_name='files')),
            ],
            bases=('jtx_video.basefile',),
        ),
        migrations.AddField(
            model_name='basefile',
            name='parent',
            field=models.ForeignKey(blank=True, to='jtx_video.BaseFile', null=True),
        ),
    ]


from sklearn.cross_validation import train_test_split
from sklearn.tree import DecisionTreeRegressor
from sklearn.metrics import r2_score

def function_def(X, y):
    """
    Shuffles and splits data into 75% training and 25% testing subsets,
    then returns the training and testing subsets.
    :param X: numpy array. Features
    :param y: numpy array. Labels
    """
    X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.25, random_state=0)
    return (X_train, y_train, X_test, y_test)
l_score = []
for s_col in data.columns:
    new_data = data.copy()
    y = new_data[s_col]
    X = new_data.drop([s_col], axis=1)
    X_train, y_train, X_test, y_test = function_def(X, y)
    regressor = DecisionTreeRegressor(random_state=0)
    regressor.fit(X_train, y_train)
    score = r2_score(regressor.predict(X_test), y_test)
    l_score.append(score)
df_r2 = pd.DataFrame(l_score, index=data.columns, columns=['Score'])
df_r2.index.names = ['PREDICTED']
df_r2


from django.urls import include, path

from . import views

urlpatterns = [
	path("", views.index),
	path("yazar/<slug:slug>/", views.author, name="author"),
	path("y/<slug:slug>/", views.article, name="article"),
	path("kat/<slug:slug>/", views.cat),
	path("kat/<slug:slug>/<int:page>", views.cat, name="cat"),
	path("arsiv/", views.archive),
	path("arsiv/<int:page>", views.archive, name="archive"),
	path("hakkinda/", views.about, name="about"),
	path("kunye/", views.team, name="team"),
	path("ara/", views.search, name="search"),
]


from pytube import YouTube

link = "https://youtu.be/Jxd63PGOlYE?list=PLauivoElc3girg6gOwzd1f8BQlnfjOM6P"
youtube_1 = YouTube(link)

videos = youtube_1.streams.all()

for i, video in enumerate(videos):
    print(f"{i}: {video}")

stream_number = int(input("Enter the stream number to download: "))
selected_video = videos[stream_number]
selected_video.download()
print("Video downloaded successfully!")

# from pytube import YouTube
#
# link = "https://youtu.be/Jxd63PGOlYE?list=PLauivoElc3girg6gOwzd1f8BQlnfjOM6P"
# youtube_1 = YouTube(link)
#
# #youtube function
# # print(youtube_1.title)
# # print(youtube_1.thumbnail_url)
#
# videos = youtube_1.streams.all()
# #all streams qualitya valible i get
# #videos= youtube_1.streams.filter(only_audio=True)
#
# #through enumerate i can get index
# vid = list(enumerate(videos))
# for i in vid:
#     print(i)
# print()
# strm = int(input("enter:"))
# videos[strm].download()
# print('Successfully')

#Playlist function
# from pytube import Playlist
# py = Playlist("https://www.youtube.com/playlist?list=PLauivoElc3girg6gOwzd1f8BQlnfjOM6P")
# print(f'Downloading:{py.title}')
# for video in py.videos:
#     video.streams.first().download()
# print('Successfully')


from datetime import datetime
from flask import render_template, flash, redirect, request, url_for, current_app, abort
from flask_login import current_user, login_required
from sqlalchemy import func

from app.blueprints.main import bp
from app import db, images
from app.models import User, Post, Vote, Image, Chat, Sub, Comment
from app.blueprints.main.forms import EditProfileForm, NewPostForm, NewCommentForm, NewChatForm, EditChatForm, SearchForm

def paginate(items, per_page=None):
    if per_page is None:
        per_page = current_app.config['DEFAULT_PER_PAGE']

    page = request.args.get('page', 1, type=int)
    items = items.paginate(page, per_page, False)

    next_url = url_for(request.endpoint, page=items.next_num) if items.has_next else None
    prev_url = url_for(request.endpoint, page=items.prev_num) if items.has_prev else None
    return items, next_url, prev_url

@bp.before_app_request
def before_request():
    if current_user.is_authenticated:
        current_user.last_seen = datetime.utcnow()
        db.session.commit()

# ------------------------------ FRONT PAGES ------------------------------

@bp.route('/')
@bp.route('/index')
def index():
    return render_template('main/home.html', title='Home')

@bp.route('/feed')
@login_required
def feed():
    posts = current_user.subscribed_posts().order_by(Post.created_at.desc())

    posts, next_url, prev_url = paginate(posts, current_app.config['POSTS_PER_PAGE'])

    return render_template('main/feed.html', title='Your Feed', posts=posts.items, next_url=next_url, prev_url=prev_url)

@bp.route('/popular', methods=['GET', 'POST'])
def popular():
    search_form = SearchForm()
    posts = db.session.query(Post)

    if search_form.validate_on_submit():
        posts = posts.filter(Post.title.like('%' + search_form.search.data + '%') |
                             Post.body.like('%' + search_form.search.data + '%'))

    posts = posts.join(Vote).group_by(Post.id).order_by(func.count().desc())
    posts, next_url, prev_url = paginate(posts, current_app.config['POSTS_PER_PAGE'])

    return render_template('main/popular.html', title='Popular',
                           posts=posts.items, next_url=next_url, prev_url=prev_url,
                           form=search_form, include_chat=True)

@bp.route('/explore_chats', methods=['GET', 'POST'])
def explore_chats():
    form = SearchForm()
    chats = db.session.query(Chat)

    if form.validate_on_submit():
        chats = chats.filter(Chat.name.like('%' + form.search.data + '%') |
                             Chat.about.like('%' + form.search.data + '%'))

    chats = chats.join(Sub).group_by(Chat.id).order_by(func.count().desc())
    chats, next_url, prev_url = paginate(chats, current_app.config['CHATS_PER_PAGE'])

    return render_template('main/explore_chats.html', title='Explore',
                           chats=chats.items, next_url=next_url, prev_url=prev_url,
                           form=form)

@bp.route('/leaderboard', methods=['GET', 'POST'])
def leaderboard():
    search_form = SearchForm()
    users = db.session.query(User)

    if search_form.validate_on_submit():
        users = users.filter(User.username.like('%' + search_form.search.data + '%'))

    users = users.join(Post).join(Vote).group_by(User.id).order_by(func.count().desc())
    users, next_url, prev_url = paginate(users, current_app.config['USERS_PER_PAGE'])

    return render_template('main/leaderboard.html', title='Leaderboard',
                           users=users.items, next_url=next_url, prev_url=prev_url,
                           form=search_form)

# ------------------------------ USERS ------------------------------

@bp.route('/user/<username>')
def show_user(username):
    user = User.get_by_username(username)
    if user is None:
        return abort(404)

    posts = Post.query.filter_by(author=user).order_by(Post.created_at.desc())
    posts, next_url, prev_url = paginate(posts, current_app.config['POSTS_PER_PAGE'])

    return render_template('main/user.html', user=user,
                           posts=posts.items, next_url=next_url, prev_url=prev_url,
                           include_chat=True)

@bp.route('/edit_profile', methods=['GET', 'POST'])
@login_required
def edit_profile():
    form = EditProfileForm()
    if form.validate_on_submit():
        current_user.about_me = form.about_me.data
        db.session.commit()
        flash('Your changes have been saved.')
        return redirect(url_for('main.show_user', username=current_user.username))
    elif request.method == 'GET':
        form.about_me.data = current_user.about_me
    return render_template('main/edit_profile.html', title='Edit Profile',
                           form=form)

# ------------------------------ CHATS ------------------------------

@bp.route('/chat/<name>', methods=['GET', 'POST'])
def show_chat(name):
    chat = Chat.get_by_name(name)
    if chat is None:
        return abort(404)

    posts = Post.query.filter_by(chat=chat).order_by(Post.created_at.desc())
    posts, next_url, prev_url = paginate(posts, current_app.config['POSTS_PER_PAGE'])

    return render_template('main/chat.html', title=chat.name, chat=chat,
                           posts=posts.items, next_url=next_url, prev_url=prev_url)

@bp.route('/create_chat', methods=['GET', 'POST'])
@login_required
def create_chat():
    form = NewChatForm()

    if form.validate_on_submit():
        new_chat = Chat(name=form.name.data, about=form.about.data, creator=current_user)
        current_user.subscribe(new_chat)
        db.session.add(new_chat)
        db.session.commit()
        flash('Your new chat is now live!')
        return redirect(url_for('main.show_chat', name=new_chat.name))

    return render_template('main/create_chat.html', title="New Chat", form=form)

@bp.route('/edit_chat/<name>', methods=['GET', 'POST'])
@login_required
def edit_chat(name):
    chat = Chat.get_by_name(name)
    if chat is None:
        return abort(404)
    if chat not in current_user.chats:
        flash('You cannot edit this chat.')
        return redirect(url_for('show_chat', name=name))
    form = EditChatForm()
    if form.validate_on_submit():
        chat.about = form.about.data
        db.session.commit()
        flash('Your changes have been saved.')
        return redirect(url_for('main.show_chat', name=chat.name))
    elif request.method == 'GET':
        form.about.data = chat.about
    return render_template('main/edit_chat.html', title='Edit chat/{}'.format(chat.name), chat=chat,
                           form=form)

@bp.route('/subscribe/<name>')
@login_required
def subscribe(name):
    chat = Chat.get_by_name(name)
    if chat is None:
        flash('Chat {} not found.'.format(name))
        return redirect(url_for('main.index'))
    current_user.subscribe(chat)
    db.session.commit()
    flash('You subscribed to chat/{}!'.format(name))
    return redirect(request.referrer or url_for('main.show_chat', name=name))

@bp.route('/unsubscribe/<name>')
@login_required
def unsubscribe(name):
    chat = Chat.get_by_name(name)
    if chat is None:
        flash('Chat {} not found.'.format(name))
        return redirect(url_for('main.index'))
    current_user.unsubscribe(chat)
    db.session.commit()
    flash('You unsubscribed from chat/{}.'.format(name))
    return redirect(request.referrer or url_for('main.show_chat', name=name))

# ------------------------------ POSTS ------------------------------

@bp.route('/post/<post_id>', methods=['GET', 'POST'])
def show_post(post_id):
    post = Post.query.filter_by(id=post_id).first_or_404()

    form = NewCommentForm()
    if form.validate_on_submit():
        comment = Comment(body=form.body.data,
                          author=None if current_user.is_anonymous else current_user)
        comment.post = post
        db.session.add(comment)
        db.session.commit()
        flash('Your comment is now live!')
        return redirect(url_for('main.show_post', post_id=post_id))

    comments = Comment.query.filter_by(post_id=post_id).order_by(Comment.created_at.desc())
    comments, next_url, prev_url = paginate(comments, current_app.config['COMMENTS_PER_PAGE'])

    return render_template('main/post.html', title=post.title, post=post,
                           comments=comments.items, next_url=next_url, prev_url=prev_url,
                           form=form)

@bp.route('/make_post', methods=['GET', 'POST'], defaults={'chat_name': None})
@bp.route('/make_post/<chat_name>', methods=['GET', 'POST'])
@login_required
def make_post(chat_name):
    form = NewPostForm(chat_name=chat_name)
    if form.validate_on_submit():
        chat_name = form.chat_name.data

        new_post = Post(title=form.title.data, body=form.body.data, author=current_user)
        new_post.chat = Chat.query.filter_by(name=chat_name).first()

        if 'image' in request.files and request.files['image']:
            filename = images.save(form.image.data)
            url = images.url(filename)
            image = Image(filename=filename, url=url)
            new_post.attachment = image

        db.session.add(new_post)
        db.session.commit()
        flash('Your new post is now live!')
        return redirect(url_for('main.show_chat', name=chat_name))
    return render_template('main/make_post.html', title="New Post", form=form)

@bp.route('/delete_post/<post_id>', methods=['POST'])
@login_required
def delete_post(post_id):
    post = Post.query.filter_by(id=post_id).first()
    if post is None:
        flash('The post you are trying to delete does not exist.')
        return redirect(request.referrer or url_for('main.index'))
    chat_name = post.chat.name

    db.session.delete(post)
    db.session.commit()
    flash('You deleted the post.')
    return redirect(request.referrer or url_for('main.show_chat', name=chat_name))

@bp.route('/upvote/<post_id>')
@login_required
def upvote(post_id):
    post = Post.query.filter_by(id=post_id).first()
    if post is None:
        flash('Post not found.')
        return redirect(url_for('main.index'))
    vote = current_user.upvote(post)
    db.session.merge(vote)
    db.session.commit()
    flash('You upvoted the post!')
    return redirect(request.referrer or url_for('main.show_post', post_id=post_id))

@bp.route('/downvote/<post_id>')
@login_required
def downvote(post_id):
    post = Post.query.filter_by(id=post_id).first()
    if post is None:
        flash('Post not found.')
        return redirect(url_for('main.index'))
    vote = current_user.downvote(post)
    db.session.merge(vote)
    db.session.commit()
    flash('You downvoted the post!')
    return redirect(request.referrer or url_for('main.show_post', post_id=post_id))


__author__ = "Jonas Lauener"
__copyright__ = "2020, Jonas Lauener & Wingtra AG"
__licence__ = "Mozilla Public License, v. 2.0, https://mozilla.org/MPL/2.0/"

import argparse

import matplotlib.pyplot as plt
from nptdms import TdmsFile
from pathlib import Path

def main(tdms_file_path):
    tdms_file = TdmsFile.read(tdms_file_path)
    data = tdms_file["data"]

    time = data['Time (s)']
    channel1 = data['Channel 1 (V)']
    channel2 = data['Channel 2 (V)']

    plt.plot(time, channel1)
    plt.plot(time, channel2)
    plt.show()

def existing_tdms_file_path(string):
    path = Path(string)
    if path.exists() and path.suffix == ".tdms":
        return path
    else:
        raise FileNotFoundError(string)

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description='Plot TDMS file.')
    parser.add_argument("tdms_file", help="Path to TDMS file.", type=existing_tdms_file_path)

    args = parser.parse_args()
    main(args.tdms_file)


import functools
import numpy as np
import gymnasium
from gymnasium.spaces import Discrete,Box

from pettingzoo import ParallelEnv
from pettingzoo.utils import parallel_to_aec, wrappers
from ..render.renderer import Renderer
from epsim.core import World,WorldObj,Slot,Crane,CraneAction,SHARE
from epsim.core.componets import Color
from epsim.core import SHARE

import logging
logger = logging.getLogger(__name__)
def env(render_mode=None,args: "DictConfig"  = None):
    """
    The env function often wraps the environment in wrappers by default.
    You can find full documentation for these methods
    elsewhere in the developer documentation.
    """
    env = raw_env(render_mode=render_mode,args=args)
    # this wrapper helps error handling for discrete action spaces
    env = wrappers.AssertOutOfBoundsWrapper(env)
    # Provides a wide vareity of helpful user errors
    # Strongly recommended
    env = wrappers.OrderEnforcingWrapper(env)
    return env

def raw_env(render_mode=None,args: "DictConfig"  = None):
    """
    To support the AEC API, the raw_env() function just uses the from_parallel
    function to convert from a ParallelEnv to an AEC env
    """
    env = parallel_env(render_mode=render_mode,args=args)
    env = parallel_to_aec(env)
    return env

class parallel_env(ParallelEnv):
    metadata = {"render_modes":  ["human", "rgb_array"],"name": "electroplating_v1"}

    def __init__(self, render_mode=None,args: "DictConfig"  = None):
        self.args=args
        Renderer.LANG=args.language
        SHARE.TILE_SIZE=args.tile_size
        SHARE.SHORT_ALARM_TIME=args.alarm.short_time
        SHARE.LONG_ALARM_TIME=args.alarm.long_time
        SHARE.OBSERVATION_IMAGE=args.observation_image
        self.world=World(args.data_directory,args.max_steps,args.cool_down_time,args.auto_put_starts,args.auto_dispatch_crane)

        ncols=args.screen_columns
        max_x=max(list(self.world.pos_slots.keys()))
        SHARE.MAX_X=max_x

        nrows=max_x//ncols+2
        draw_text=render_mode in ['human', 'rgb_array']

        self.renderer=Renderer(self.world,args.fps,nrows,ncols,draw_text)
        self.possible_agents=[SHARE.DISPATCH_CODE]

        self.possible_agents.extend([crane.cfg.name for crane in self.world.all_cranes])

        # optional: a mapping between agent name and ID
        self.agent_name_mapping = dict(
            zip(self.possible_agents, list(range(len(self.possible_agents))))
        )
        self.render_mode = render_mode

    @property
    def one_observation_size(self):
        return SHARE.OBJ_TYPE_SIZE+SHARE.OP_TYPE1_SIZE+SHARE.OP_TYPE2_SIZE+SHARE.PRODUCT_TYPE_SIZE+4

    @functools.lru_cache(maxsize=None)
    def observation_space(self, agent):#todo  dispatch see all obs
        rt=None
        tsize=SHARE.TILE_SIZE
        if agent==SHARE.DISPATCH_CODE:
            if SHARE.OBSERVATION_IMAGE:
                rt=Box(0,255,(3*tsize,SHARE.MAX_STATE_LIST_LEN*tsize,3),dtype=np.uint8)
            else:
                rt=Box(-1,1,(SHARE.MAX_STATE_LIST_LEN*self.one_observation_size,),dtype=np.float32)
        else:
            obs_len=SHARE.MAX_OBS_LIST_LEN
            if SHARE.OBSERVATION_IMAGE:
                rt=Box(0,255,(3*tsize,obs_len*tsize,3),dtype=np.uint8)
            else:
                rt=Box(-1,1,(obs_len*self.one_observation_size,),dtype=np.float32)
        return rt

    # Action space should be defined here.
    # If your spaces change over time, remove this line (disable caching).
    @functools.lru_cache(maxsize=None)
    def action_space(self, agent):
        return Discrete(3) if  agent==SHARE.DISPATCH_CODE else  Discrete(5)

    def render(self):
        if self.render_mode is None:
            gymnasium.logger.warn(
                "You are calling render method without specifying any render mode."
            )
            return

        self.screen_img=self.renderer.render(self.render_mode)
        return self.screen_img

    def close(self):
        self.renderer.close()

    # def observe(self, agent: str):
    #     print(f'observe for {agent}')
    #     return self.observations[agent]

    def _make_jobs(self):
        ps=[]
        for p in self.args.products[self.args.data_directory]:
            ps.extend([p.code]*p.num)
        #print(ps)
        self.world.add_jobs(ps)
        for agv in self.world.all_cranes:
            agv.color=Color(255,255,255)
        self.world.cur_crane.color=Color(255,0,0)

    def reset(self, seed=None, options=None):
        """
        Reset needs to initialize the `agents` attribute and must set up the
        environment so that render(), and step() can be called without issues.
        Here it initializes the `num_moves` variable which counts the number of
        hands that are played.
        Returns the observations for each agent
        """
        nrows=self.renderer.nrows
        ncols=self.renderer.ncols
        self.world.reset()
        self.agents = self.possible_agents[:]
        if self.render_mode == "human" or self.render_mode == "rgb_array":
            screen_img=self.render()
            self.world.get_state_img(screen_img,nrows,ncols)

        observations, infos = self._make_info()

        self._make_jobs()
        return observations, infos

    def _make_info(self):
        observations = {}#agent: NONE for agent in self.agents}

        observations[SHARE.DISPATCH_CODE]=self.world.get_state()
        if SHARE.OBSERVATION_IMAGE:
            observations[SHARE.DISPATCH_CODE]=self.world.get_state_img()

        self.world.get_dispatch_masks()

        for idx,agv in enumerate(self.world.all_cranes):
            if SHARE.OBSERVATION_IMAGE:
                observations[agv.cfg.name]=self.world.get_observation_img(agv)
            else:
                observations[agv.cfg.name]=self.world.get_observation(agv)
            self.world.get_masks(agv)

            #infos[agv.cfg.name]={"action_masks":mask}
        self.observations = observations
        #self.infos = infos
        # for name,obs,mask in zip(infos.keys(),observations.values(),infos.values()):
        #     self.observations[name]={'observation':obs,'action_masks':mask}
        return observations,self.world.masks

    def step(self, actions:dict):
        act=0
        acts={}#[0]*len(self.world.all_cranes)
        for key,action in actions.items():
            if key==SHARE.DISPATCH_CODE:
                act=action
            else:
                #idx=self.agent_name_mapping[key]-1
                acts[key]=action

        self.world.do_dispatch(act)
        self.world.set_crane_commands(acts)

        self.world.update()
        nrows=self.renderer.nrows
        ncols=self.renderer.ncols
        if self.render_mode == "human" or self.render_mode == "rgb_array":
            screen_img=self.render()
            self.world.get_state_img(screen_img,nrows,ncols)

        # rewards for all agents are placed in the rewards dictionary to be returned
        rewards = self.world._rewards

        terminations = {agent: self.world.is_over for agent in self.agents}
        truncations = {agent: self.world.is_timeout for agent in self.agents}
        observations, infos = self._make_info()

        # if self.world.is_over:
        #     self.agents = []

        return observations, rewards, terminations, truncations, infos


"""
使用颜色映射 和 自动保存图表
颜色映射 （ colormap ）是一系列颜色，它们从起始颜色渐变到结束颜色。
"""
import matplotlib.pyplot as plt

x_values = list(range(1, 1001))
y_values = [x**2 for x in x_values]

# 但绘制很多点时，黑色轮廓可能会粘连在一起。要删除数据点的轮廓，可在调用 scatter() 时传递实参 edgecolor='none'
# 可传递参数 c ，并将其设置为一个元组，其中包含三个 0~1 之间的小数值，它们分别表示红色、绿色和蓝色分量。
# 我们将参数 c 设置成了一个 y  值列表，并使用参数 cmap 告诉 pyplot 使用哪个颜色映射。这些代码将 y  值较小的点显示为浅蓝色，并将 y  值较大的点显示为深蓝色
plt.scatter(x_values, y_values,c=y_values, cmap=plt.cm.Blues, edgecolor='none', s=40)

#  设置图表标题并给坐标轴加上标签
plt.title("Square Numbers", fontsize=24)
plt.xlabel("Value", fontsize=14)
plt.ylabel("Square of Value", fontsize=14)

#  设置刻度标记的大小
plt.tick_params(axis='both', which='major', labelsize=14)

#  设置每个坐标轴的取值范围
plt.axis([0, 1100, 0, 1100000])

# 要让程序自动将图表保存到文件中，可将对 plt.show() 的调用替换为对 plt.savefig() 的调用
plt.savefig('squares_plot.png', bbox_inches='tight')


from django.conf.urls import url
from django.urls import path, re_path
from . import views

app_name = 'polls'
urlpatterns = [
    path('', views.IndexView, name='index'),
    path('signup/', views.signup_view, name='signup'),
    path('login/', views.login_view, name='login'),
    path('todolists/', views.TodoListView, name='todolists'),
    path('createtodolist/', views.CreateTodoList, name='createtodolist'),
    path('entries/', views.EntriesView, name='entries'),
    path('createentry/', views.CreateEntry, name='createentry'),
    path('target/', views.TargetViewFunction, name='target'),
    path('addmember/', views.AddMember, name='addmember'),
    path('mytodolists/', views.ViewMyTodoLists, name='mytodolists'),
    path('sharedtodolists/', views.ViewSharedTodoLists, name='sharedtodolists'),
    path('profile/', views.ViewProfile, name='profile'),
    path('base/', views.logout_view, name='logout'),
    path('renamelist/', views.RenameListView, name='renamelist'),
    path('editprofile/', views.EditProfile, name='editprofile'),
    path('ajax/entrystatus/', views.EntryStatus, name='entrystatus'),
    path('editentry/', views.EditEntry, name='editentry'),
    path('ajax/activatebuttons/', views.ActivateButtons, name='activatebuttons'),
    path('ajax/change_member_accessibility/', views.ChangeMemberAccessibility, name='changememberaccessibility'),

]


import sqlite3
import pandas as pd

def count(file,user_id):
    sql_con = sqlite3.connect(file)
    df= pd.read_sql("SELECT SUM(minutes) FROM data WHERE user_id ="+ user_id,sql_con)
    if df.loc[0][0] == None:
        sum_minutes=0
    else:
        sum_minutes=int(df.loc[0][0])
#    print(sum_minutes)
    return sum_minutes

count("users_data.sqlite","8490")
#count("users_data.sqlite","8901111")


import json
import os
import time

import pandas as pd

class Logger:
    MAIN_FOLDER = "logs"

    def __init__(self, folder_name):
        self.folder_name = folder_name
        self.features_path = None
        self.portfolio_path = None

    def initialize(self):
        timestamp = int(time.time())
        os.makedirs(f'{Logger.MAIN_FOLDER}', exist_ok=True)
        subfolder_path = f'{Logger.MAIN_FOLDER}/{self.folder_name}_{timestamp}'
        os.makedirs(subfolder_path, exist_ok=True)

        self.features_path = f"{subfolder_path}/features.log"
        self.portfolio_path = f"{subfolder_path}/portfolio.log"

    def log(self, feature_service, portfolio_manager, price_data):
        with open(self.features_path, "a+") as f_features:
            f_features.write(json.dumps(feature_service.status) + "\n")
        with open(self.portfolio_path, "a+") as f_portfolio:
            f_portfolio.write(json.dumps(portfolio_manager.log_wallets(price_data)) + "\n")

        f_features.close()
        f_portfolio.close()

    def load_owner_wallet(self, owner_name):
        with open(self.portfolio_path, "r") as f:
            history = pd.read_json(f, lines=True)
        # print(history)
        for column in history.columns:
            if column not in ["timestamp", "BTC_price"]:
                history[column] = history[column].apply(lambda x: x[owner_name])
        return history


from pdf2image import convert_from_path
import sys
import os
import re
import itertools
from collections import defaultdict
from os import listdir
from os.path import isfile, join
from math import inf

def main(argv):
    if len(argv) >= 6:
        num_vis = int(argv[4])
        if len(argv) == num_vis*2 + 5:
            filename = argv[1]
            dir = argv[2]
            last_page = int(argv[3])
            ini = []
            fin = []
            it = 5
            for i in range(0, num_vis):
                ini.append(int(argv[it]))
                fin.append(int(argv[it+1]))
                it+=2
            ini.append(inf)
            fin.append(inf)
        else:
            print ('usage:\n    python txt_2md.py <txt> <diretorio do pdf> <pagina final documento> <numero de visualizaçoes> <pagina inicial vis 1> <pagina final vis1> ... <pagina inicial vis N> <pagina final vis N>')
            return
    else:
        print ('usage:\n    python txt_2md.py <txt> <diretorio do pdf> <pagina final documento> <numero de visualizaçoes> <pagina inicial vis 1> <pagina final vis1> ... <pagina inicial vis N> <pagina final vis N>')
        return

    cpp_dir = dir+'cpp/'
    images_dir = 'images/'

    only_numbers = re.compile(r'^[\s\d]+$', re.U)
    not_only_numbers = re.compile(r'[^\W\d_]', re.U)
    code_flag = re.compile(r'{{{{{{{{{{{{{{{{{{{{{}}}}}}}}}}}}}}}}}}}}}')
    codes = [join(cpp_dir, f) for f in listdir(cpp_dir) if isfile(join(cpp_dir, f))]
    it = 0
    lines = open(filename).readlines()
    flag = False
    vis_it = 0
    d = defaultdict(lambda: 0)

    for i in range(1, len(lines)):
        if only_numbers.search(lines[i-1]) and not_only_numbers.search(lines[i]):
            current_page = int(lines[i-1].strip(' ').strip('\n'))
            if(current_page < ini[vis_it] or current_page > fin[vis_it]):
                flag = True
            else:
                flag = False
            if(current_page+1 == fin[vis_it]):
                print('## ', lines[i].strip(' ').strip('\n'))
                log = '![]({0}movie.gif)\nPara acessar álbum de imagens [clique aqui]({0}).\n'.format(images_dir+'vis-'+str(vis_it)+'/')
                print(log)
                vis_it+=1
                d[lines[i]] = 1
            if(current_page >= last_page):
                flag = False

        if flag == True and not_only_numbers.search(lines[i]) and not d[lines[i]]:
            line = lines[i].strip(' ').strip('\n').replace('•', '*')
            if only_numbers.search(lines[i-1]):
                line = "## "+line
            print(line)
            d[lines[i]] = 1

        if code_flag.search(lines[i]):
            code = open(codes[it], 'r')
            print('```C++')
            print(code.read())
            print('```')
            it+=1
            code.close()

    return
if __name__ == '__main__':
    main(sys.argv)