# # Copyright 2019 The ASReview Authors. All Rights Reserved.
# #
# # Licensed under the Apache License, Version 2.0 (the "License");
# # you may not use this file except in compliance with the License.
# # You may obtain a copy of the License at
# #
# #     http://www.apache.org/licenses/LICENSE-2.0
# #
# # Unless required by applicable law or agreed to in writing, software
# # distributed under the License is distributed on an "AS IS" BASIS,
# # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# # See the License for the specific language governing permissions and
# # limitations under the License.
# from typing import Optional

# import logging

# import numpy as np # type: ignore

# from .base import BaseBalancer
# from .double import _one_weight
# from .double import DoubleBalancer
# from .double import _zero_weight
# from .double import fill_training
# from .double import random_round
# from ..utils import get_random_generator 

# from ..activelearning.catalog import ALCatalog

# LOGGER = logging.getLogger(__name__)

# class TripleBalancer(BaseBalancer):
#     """
#     Class to get the three way rebalancing function and arguments.
#     It divides the data into three groups: 1's, 0's from random sampling,
#     and 0's from max sampling. Thus it only makes sense to use this class in
#     combination with the rand_max query strategy.
#     """
#     name = "triple"

#     def __init__(self,
#                  a: float=2.155,
#                  alpha: float=0.94,
#                  b: float=0.789,
#                  beta: float=1.0,
#                  c: float=0.835,
#                  gamma: float=2.0,
#                  shuffle: bool=True,
#                  random_state: Optional[np.random.Generator] = None): # type: ignore
#         """Initialize the triple balance strategy.

#         Arguments
#         ---------
#         a: float
#             Governs the weight of the 1's. Higher values mean linearly more 1's
#             in your training sample.
#         alpha: float
#             Governs the scaling the weight of the 1's, as a function of the
#             ratio of ones to zeros. A positive value means that the lower the
#             ratio of zeros to ones, the higher the weight of the ones.
#         b: float
#             Governs how strongly we want to sample depending on the total
#             number of samples. A value of 1 means no dependence on the total
#             number of samples, while lower values mean increasingly stronger
#             dependence on the number of samples.
#         beta: float
#             Governs the scaling of the weight of the zeros depending on the
#             number of samples. Higher values means that larger samples are more
#             strongly penalizing zeros.
#         c: float
#             Value between one and zero that governs the weight of samples done
#             with maximal sampling. Higher values mean higher weight.
#         gamma: float
#             Governs the scaling of the weight of the max samples as a function
#             of the % of papers read. Higher values mean stronger scaling.
#         """
#         super(TripleBalancer, self).__init__()
#         self.a = a
#         self.alpha = alpha
#         self.b = b
#         self.beta = beta
#         self.c = c
#         self.gamma = gamma
#         self.shuffle = shuffle
#         self.fallback_model = DoubleBalancer(a=a, alpha=alpha, b=b, beta=beta,
#                                             rng=random_state)
#         self._random_state = get_random_generator(random_state)
    
#     def resample(self, X, y):
#         max_idx = self._environment.logger.get_instances_by_method(ALCatalog.QueryType.LABELMAXIMIZER)
#         max_idx = np.array(shared["query_src"].get("max", []), dtype=np.int)
#         rand_idx = np.array([], dtype=np.int)
#         for qtype in shared["query_src"]:
#             if qtype != "max":
#                 rand_idx = np.append(rand_idx, shared["query_src"][qtype])

#         rand_idx = rand_idx.astype(int)
#         # Write them back for next round.
#         if self.shuffle:
#             self._random_state.shuffle(rand_idx)
#             self._random_state.shuffle(max_idx)

#         if len(rand_idx) == 0 or len(max_idx) == 0:
#             LOGGER.debug("Warning: trying to use triple balance, but unable"
#                           f" to, because we have {len(max_idx)} max samples "
#                           f"and {len(rand_idx)} random samples.")
#             return self.fallback_model.resample(X, y)

#         # Split the idx into three groups: 1's, random 0's, max 0's.
#         one_idx = train_idx[np.where(y[train_idx] == 1)]
#         zero_max_idx = max_idx[np.where(y[max_idx] == 0)]
#         zero_rand_idx = rand_idx[np.where(y[rand_idx] == 0)]

#         if len(zero_rand_idx) == 0 or len(zero_max_idx) == 0:
#             LOGGER.debug("Warning: trying to use triple balance, but unable "
#                           f"to, because we have {len(zero_max_idx)} zero max"
#                           f"samples and {len(zero_rand_idx)} random samples.")
#             return self.fallback_model.sample(X, y, train_idx, shared)

#         n_one = len(one_idx)
#         n_zero_rand = len(zero_rand_idx)
#         n_zero_max = len(zero_max_idx)
#         n_samples = len(y)
#         n_train = len(train_idx)

#         # Get the distribution of 1's, and random 0's and max 0's.
#         n_one_train, n_zero_rand_train, n_zero_max_train = _get_triple_dist(
#             n_one, n_zero_rand, n_zero_max, n_samples, n_train, self.a,
#             self.alpha, self.b, self.beta, self.c, self.gamma,
#             self._random_state)
#         logging.debug(f"(1, 0_rand, 0_max) = ({n_one_train}, "
#                       f"{n_zero_rand_train}, {n_zero_max_train})")

#         one_train_idx = fill_training(one_idx, n_one_train, self._random_state)
#         zero_rand_train_idx = fill_training(zero_rand_idx, n_zero_rand_train,
#                                             self._random_state)
#         zero_max_train_idx = fill_training(zero_max_idx, n_zero_max_train,
#                                            self._random_state)

#         all_idx = np.concatenate(
#             [one_train_idx, zero_rand_train_idx, zero_max_train_idx])
#         self._random_state.shuffle(all_idx)

#         return X[all_idx], y[all_idx]

#     def full_hyper_space(self):
#         from hyperopt import hp
#         parameter_space = {
#             "bal_a": hp.lognormal("bal_a", 0, 1),
#             "bal_alpha": hp.uniform("bal_alpha", 0, 2),
#             "bal_b": hp.uniform("bal_b", 0, 1),
#             # "bal_zero_beta": hp.uniform("bal_zero_beta", 0, 2),
#             "bal_c": hp.uniform("bal_c", 0, 1),
#             # "bal_zero_max_gamma": hp.uniform("bal_zero_max_gamma", 0.01, 2)
#         }
#         return parameter_space, {}


# def _zero_max_weight(fraction_read, c, gamma):
#     """
#     Get the weight ratio between ones and zeros.

#     Parameters
#     ----------
#     beta:
#         Exponent governing decay of 1's to 0's.
#     delta:
#         Asymptotic ratio for infinite number of samples.

#     Returns
#     -------
#     float:
#         Weight ratio between 1's and 0's
#     """
#     weight = 1 - (1 - c) * (1 - fraction_read)**gamma
#     return weight


# def _get_triple_dist(n_one, n_zero_rand, n_zero_max, n_samples, n_train, one_a,
#                      one_alpha, zero_b, zero_beta, zero_max_c, zero_max_gamma,
#                      random_state):
#     " Get the number of 1's, random 0's and max 0's in each mini epoch. "
#     n_zero = n_zero_rand + n_zero_max
#     n_read = n_one + n_zero
#     one_weight = _one_weight(n_one, n_zero, one_a, one_alpha)
#     zero_weight = _zero_weight(n_read, zero_b, zero_beta)
#     zero_max_weight = _zero_max_weight(n_read / n_samples, zero_max_c,
#                                        zero_max_gamma)

#     tot_zo_weight = one_weight * n_one + zero_weight * n_zero

#     n_one_train = random_round(one_weight * n_one * n_train / tot_zo_weight,
#                                random_state)
#     n_one_train = max(1, min(n_train - 2, n_one_train))
#     n_zero_train = n_train - n_one_train

#     tot_rm_weight = 1 * n_zero_rand + zero_max_weight * n_zero_max
#     n_zero_rand_train = random_round(
#         n_zero_train * 1 * n_zero_rand / tot_rm_weight, random_state)
#     n_zero_rand_train = max(1, min(n_zero_rand - 1, n_zero_rand_train))
#     n_zero_max_train = n_zero_train - n_zero_rand_train

#     return n_one_train, n_zero_rand_train, n_zero_max_train
