from typing import Tuple, Union
import random
from collections import defaultdict

import numpy as np
from sklearn.tree import DecisionTreeRegressor, DecisionTreeClassifier
from sklearn.tree._tree import TREE_LEAF
from sklearn.preprocessing import StandardScaler
from sklearn.ensemble import RandomForestRegressor, RandomForestClassifier


def class_sampling(
        X: np.ndarray,
        proba: np.ndarray,
        sample_weights: np.ndarray=None,
        max_duplication: int=2) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
    """ Duplicate each feature vector for each class, while adjusting
    sample_weights to reflect the given probabilities.

    Parameters
    ----------
    X : numpy.ndarray
        The features.
    proba : numpy.ndarray
        Class probabilities for each feature vector. Shape: (n_samples, n_classes)
    sample_weights: numpy.ndarray
        Sample weights before duplication. Shape: (n_samples, )
    max_duplication: int
        Maximum number of times a row is duplicated.
        If there are two classes, each row will be duplicated once (duplication=1).
    Returns
    -------
    X : numpy.ndarray
        Feature rows purposefully duplicated. Shape: (n_samples, dim)
    sample_weight : numpy.ndarray
        Sample weights for training a new random forest. Shape: (n_samples, )
    """
    if sample_weights is None:
        sample_weights = np.ones(X.shape[0])

    dup = min(proba.shape[1], max_duplication + 1)
    argtop = (-proba).argsort(axis=1)[:, :dup]
    row_indices = np.arange(len(proba))[:, np.newaxis]
    rough_proba = proba[row_indices, argtop]
    rough_proba /= rough_proba.sum(axis=1, keepdims=True)

    rows = []
    y = []
    weights = []
    for k in range(dup):
        rows.append(X)
        y.append(argtop[:, k])
        weights.append(sample_weights * rough_proba[:, k])

    return np.vstack(rows), np.concatenate(y), np.concatenate(weights)


class CountingTree:
    """ Wrapper to count the number of samples in each branch of the tree.
        There is no reason for the user to directly use this class. """
    def __init__(self, base_tree: Union[DecisionTreeRegressor, DecisionTreeClassifier]) -> None:
        self._tree = base_tree
        self._count = np.zeros(self._tree.tree_.node_count)

    def count(self, X: np.ndarray, weight: float) -> None:
        # indicator is a sparse matrix of shape (n_samples, n_nodes)
        indicator = self._tree.decision_path(X)

        # iterate over the indicator matrix to count the samples in each branch
        rows, cols = indicator.nonzero()
        np.add.at(self._count, cols, weight)
        # equivalent to:
        # for j in cols:
        #     self._count[j] += weight

    def left_probability(self, node_index: int) -> float:
        left = self._count[self._tree.tree_.children_left[node_index]]
        right = self._count[self._tree.tree_.children_right[node_index]]
        return left / (left + right)


class FeatureGenerator:
    def __init__(self) -> None:
        # we need means and stds to initialize the generated data
        self._scaler = StandardScaler()
        self._total_samples = 0

    def generate(self, approx_n: int) -> Tuple[np.ndarray, float]:
        """ Generate data near the decision boundaries.
        Parameters
        ----------
        x : int
            The approximate number of samples to generate.

        Returns
        -------
        data : numpy.ndarray
            Generated features. Shape: (n_samples, dim)
        sample_weight : float
            Weight of each sample
        """
        # number of samples per tree
        n_per_tree = approx_n // self._rf.n_estimators
        n = self._rf.n_estimators * n_per_tree  # actual number of samples

        # default values (some features won't be set by the below algorithm)
        stds = np.sqrt(self._scaler.var_)
        X = np.random.normal(size=(n, self._dim)) * stds + self._scaler.mean_

        # generate n_per_tree samples from each tree
        for i_tree, estimator in enumerate(self._rf.estimators_):
            tree = estimator.tree_
            for i in range(n_per_tree):
                row_index = i_tree * n_per_tree + i
                node_index = 0
                right_bound = np.ones(self._dim) * np.inf
                left_bound = -right_bound

                # randomly pick one path in the tree
                while node_index != TREE_LEAF and \
                    tree.children_left[node_index] != tree.children_right[node_index]:
                    threshold = tree.threshold[node_index]
                    feature_i = tree.feature[node_index]

                    # probability of branching left or right
                    left_prob = self._counting_trees[i_tree].left_probability(node_index)

                    # we pick a value close to the threshold...
                    shift = 0.05 * np.abs(np.random.normal()) * stds[feature_i]
                    if random.random() <= left_prob:
                        value = threshold - shift
                    else:
                        value = threshold + shift
                    # ... but still within the known bounds
                    value = min(right_bound[feature_i], max(left_bound[feature_i], value))
                    # alternatively, we could keep the value already set, but I believe
                    # the chosen method restricts the value to be even closer to the
                    # decision boundary
                    X[row_index, feature_i] = value

                    # branching
                    if value <= threshold:
                        node_index = tree.children_left[node_index]
                        right_bound[feature_i] = min(right_bound[feature_i], threshold)
                    else:
                        node_index = tree.children_right[node_index]
                        left_bound[feature_i] = max(left_bound[feature_i], threshold)

        return X, self._total_samples / X.shape[0]

    def register(self, rf: Union[RandomForestRegressor, RandomForestClassifier]) -> "FeatureGenerator":
        """ Call this function after training your random forest, whether or not
        the forest was trained from generated data.

        Parameters
        ----------
        rf : RandomForestRegressor or RandomForestClassifier
             A newly trained random forest.
        """
        self._rf = rf
        self._dim = rf.n_features_
        self._counting_trees = [CountingTree(t) for t in rf.estimators_]

        return self

    def update_moments(self, X: np.ndarray) -> "FeatureGenerator":
        """ Updates mean and variance of the data distribution

        Parameters
        ----------
        X : numpy.ndarray
            Features the random forest is run on. Shape: (nrows, dim)
            Typically, X is an input prediction batch.
        """
        self._scaler.partial_fit(X)
        self._total_samples += X.shape[0]

        return self

    def reinforce(self, X: np.ndarray, weight: float=1) -> "FeatureGenerator":
        """ updates the sample counts at the tree node level.
        Always call this method after training

        Parameters
        ----------
        X : numpy.ndarray
            Features the random forest is run on. Shape: (nrows, dim)
            Typically the training data.
        weight: float
            X's sample weights
        """
        for wrapper in self._counting_trees:
            wrapper.count(X, weight)

        return self

    def predict(self, X: np.ndarray) -> np.ndarray:
        """ This is just for convenience.
        You can directly call scikit RandomForest's predict() instead."""
        return self._rf.predict(X)

    def predict_proba(self, X: np.ndarray) -> np.ndarray:
        """ This is just for convenience.
        You can directly call scikit RandomForest's predict() instead.
        It only works if the model registered to the generator was a
        RandomForestClassifier.
        """
        return self._rf.predict_proba(X)

    def predict_log_proba(self, X: np.ndarray) -> np.ndarray:
        """ This is just for convenience.
        You can directly call scikit RandomForest's predict() instead.
        It only works if the model registered to the generator was a
        RandomForestClassifier.
        """
        return self._rf.predict_log_proba(X)

from typing import Tuple, Union
import random
from collections import defaultdict

import numpy as np
from sklearn.tree import DecisionTreeRegressor, DecisionTreeClassifier
from sklearn.tree._tree import TREE_LEAF
from sklearn.preprocessing import StandardScaler
from sklearn.ensemble import RandomForestRegressor, RandomForestClassifier


def class_sampling(
        X: np.ndarray,
        proba: np.ndarray,
        sample_weights: np.ndarray=None,
        max_duplication: int=2) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
    """ Duplicate each feature vector for each class, while adjusting
    sample_weights to reflect the given probabilities.

    Parameters
    ----------
    X : numpy.ndarray
        The features.
    proba : numpy.ndarray
        Class probabilities for each feature vector. Shape: (n_samples, n_classes)
    sample_weights: numpy.ndarray
        Sample weights before duplication. Shape: (n_samples, )
    max_duplication: int
        Maximum number of times a row is duplicated.
        If there are two classes, each row will be duplicated once (duplication=1).
    Returns
    -------
    X : numpy.ndarray
        Feature rows purposefully duplicated. Shape: (n_samples, dim)
    sample_weight : numpy.ndarray
        Sample weights for training a new random forest. Shape: (n_samples, )
    """
    if sample_weights is None:
        sample_weights = np.ones(X.shape[0])

    dup = min(proba.shape[1], max_duplication + 1)
    argtop = (-proba).argsort(axis=1)[:, :dup]
    row_indices = np.arange(len(proba))[:, np.newaxis]
    rough_proba = proba[row_indices, argtop]
    rough_proba /= rough_proba.sum(axis=1, keepdims=True)

    rows = []
    y = []
    weights = []
    for k in range(dup):
        rows.append(X)
        y.append(argtop[:, k])
        weights.append(sample_weights * rough_proba[:, k])

    return np.vstack(rows), np.concatenate(y), np.concatenate(weights)


class CountingTree:
    """ Wrapper to count the number of samples in each branch of the tree.
        There is no reason for the user to directly use this class. """
    def __init__(self, base_tree: Union[DecisionTreeRegressor, DecisionTreeClassifier]) -> None:
        self._tree = base_tree
        self._count = np.zeros(self._tree.tree_.node_count)

    def count(self, X: np.ndarray, weight: float) -> None:
        # indicator is a sparse matrix of shape (n_samples, n_nodes)
        indicator = self._tree.decision_path(X)

        # iterate over the indicator matrix to count the samples in each branch
        rows, cols = indicator.nonzero()
        np.add.at(self._count, cols, weight)
        # equivalent to:
        # for j in cols:
        #     self._count[j] += weight

    def left_probability(self, node_index: int) -> float:
        left = self._count[self._tree.tree_.children_left[node_index]]
        right = self._count[self._tree.tree_.children_right[node_index]]
        return left / (left + right)


class FeatureGenerator:
    def __init__(self) -> None:
        # we need means and stds to initialize the generated data
        self._scaler = StandardScaler()
        self._total_samples = 0

    def generate(self, approx_n: int) -> Tuple[np.ndarray, float]:
        """ Generate data near the decision boundaries.
        Parameters
        ----------
        x : int
            The approximate number of samples to generate.

        Returns
        -------
        data : numpy.ndarray
            Generated features. Shape: (n_samples, dim)
        sample_weight : float
            Weight of each sample
        """
        # number of samples per tree
        n_per_tree = approx_n // self._rf.n_estimators
        n = self._rf.n_estimators * n_per_tree  # actual number of samples

        # default values (some features won't be set by the below algorithm)
        stds = np.sqrt(self._scaler.var_)
        X = np.random.normal(size=(n, self._dim)) * stds + self._scaler.mean_

        # generate n_per_tree samples from each tree
        for i_tree, estimator in enumerate(self._rf.estimators_):
            tree = estimator.tree_
            for i in range(n_per_tree):
                row_index = i_tree * n_per_tree + i
                node_index = 0
                right_bound = np.ones(self._dim) * np.inf
                left_bound = -right_bound

                # randomly pick one path in the tree
                while node_index != TREE_LEAF and \
                    tree.children_left[node_index] != tree.children_right[node_index]:
                    threshold = tree.threshold[node_index]
                    feature_i = tree.feature[node_index]

                    # probability of branching left or right
                    left_prob = self._counting_trees[i_tree].left_probability(node_index)

                    # we pick a value close to the threshold...
                    shift = 0.05 * np.abs(np.random.normal()) * stds[feature_i]
                    if random.random() <= left_prob:
                        value = threshold - shift
                    else:
                        value = threshold + shift
                    # ... but still within the known bounds
                    value = min(right_bound[feature_i], max(left_bound[feature_i], value))
                    # alternatively, we could keep the value already set, but I believe
                    # the chosen method restricts the value to be even closer to the
                    # decision boundary
                    X[row_index, feature_i] = value

                    # branching
                    if value <= threshold:
                        node_index = tree.children_left[node_index]
                        right_bound[feature_i] = min(right_bound[feature_i], threshold)
                    else:
                        node_index = tree.children_right[node_index]
                        left_bound[feature_i] = max(left_bound[feature_i], threshold)

        return X, self._total_samples / X.shape[0]

    def register(self, rf: Union[RandomForestRegressor, RandomForestClassifier]) -> "FeatureGenerator":
        """ Call this function after training your random forest, whether or not
        the forest was trained from generated data.

        Parameters
        ----------
        rf : RandomForestRegressor or RandomForestClassifier
             A newly trained random forest.
        """
        self._rf = rf
        self._dim = rf.n_features_
        self._counting_trees = [CountingTree(t) for t in rf.estimators_]

        return self

    def update_moments(self, X: np.ndarray) -> "FeatureGenerator":
        """ Updates mean and variance of the data distribution

        Parameters
        ----------
        X : numpy.ndarray
            Features the random forest is run on. Shape: (nrows, dim)
            Typically, X is an input prediction batch.
        """
        self._scaler.partial_fit(X)
        self._total_samples += X.shape[0]

        return self

    def reinforce(self, X: np.ndarray, weight: float=1) -> "FeatureGenerator":
        """ updates the sample counts at the tree node level.
        Always call this method after training

        Parameters
        ----------
        X : numpy.ndarray
            Features the random forest is run on. Shape: (nrows, dim)
            Typically the training data.
        weight: float
            X's sample weights
        """
        for wrapper in self._counting_trees:
            wrapper.count(X, weight)

        return self

    def predict(self, X: np.ndarray) -> np.ndarray:
        """ This is just for convenience.
        You can directly call scikit RandomForest's predict() instead."""
        return self._rf.predict(X)

    def predict_proba(self, X: np.ndarray) -> np.ndarray:
        """ This is just for convenience.
        You can directly call scikit RandomForest's predict() instead.
        It only works if the model registered to the generator was a
        RandomForestClassifier.
        """
        return self._rf.predict_proba(X)

    def predict_log_proba(self, X: np.ndarray) -> np.ndarray:
        """ This is just for convenience.
        You can directly call scikit RandomForest's predict() instead.
        It only works if the model registered to the generator was a
        RandomForestClassifier.
        """
        return self._rf.predict_log_proba(X)

Duplicate each feature vector for each class, while adjusting     sample_weights to reflect the given probabilities.      Parameters     ----------     X : numpy.ndarray         The features.     proba : numpy.ndarray         Class probabilities for each feature vector. Shape: (n_samples, n_classes)     sample_weights: numpy.ndarray         Sample weights before duplication. Shape: (n_samples, )     max_duplication: int         Maximum number of times a row is duplicated.         If there are two classes, each row will be duplicated once (duplication=1).     Returns     -------     X : numpy.ndarray         Feature rows purposefully duplicated. Shape: (n_samples, dim)     sample_weight : numpy.ndarray         Sample weights for training a new random forest. Shape: (n_samples, ) Wrapper to count the number of samples in each branch of the tree.         There is no reason for the user to directly use this class. # indicator is a sparse matrix of shape (n_samples, n_nodes) # iterate over the indicator matrix to count the samples in each branch # equivalent to: # for j in cols: #     self._count[j] += weight # we need means and stds to initialize the generated data Generate data near the decision boundaries.         Parameters         ----------         x : int             The approximate number of samples to generate.          Returns         -------         data : numpy.ndarray             Generated features. Shape: (n_samples, dim)         sample_weight : float             Weight of each sample # number of samples per tree # actual number of samples # default values (some features won't be set by the below algorithm) # generate n_per_tree samples from each tree # randomly pick one path in the tree # probability of branching left or right # we pick a value close to the threshold... # ... but still within the known bounds # alternatively, we could keep the value already set, but I believe # the chosen method restricts the value to be even closer to the # decision boundary # branching Call this function after training your random forest, whether or not         the forest was trained from generated data.          Parameters         ----------         rf : RandomForestRegressor or RandomForestClassifier              A newly trained random forest. Updates mean and variance of the data distribution          Parameters         ----------         X : numpy.ndarray             Features the random forest is run on. Shape: (nrows, dim)             Typically, X is an input prediction batch. updates the sample counts at the tree node level.         Always call this method after training          Parameters         ----------         X : numpy.ndarray             Features the random forest is run on. Shape: (nrows, dim)             Typically the training data.         weight: float             X's sample weights This is just for convenience.         You can directly call scikit RandomForest's predict() instead. This is just for convenience.         You can directly call scikit RandomForest's predict() instead.         It only works if the model registered to the generator was a         RandomForestClassifier. This is just for convenience.         You can directly call scikit RandomForest's predict() instead.         It only works if the model registered to the generator was a         RandomForestClassifier.