# Path: difflogic/tqdm_utils.py
# def tqdm_for(total, func):
#   """wrapper of the for function with message showing on the progress bar."""
#   # Not support break cases for now.
#   with tqdm_pbar(total=total) as pbar:
#     for i in range(total):
#       message = func(i)
#       pbar.set_description(message)
#       pbar.update()
# 
# Path: difflogic/dataset/utils.py
# class RandomlyIterDataset(object):
#   """Collect data and iterate the dataset in random order."""
# 
#   def __init__(self):
#     super().__init__()
#     self.data = []
#     self.ind = 0
# 
#   def append(self, data):
#     self.data.append(data)
# 
#   @property
#   def size(self):
#     return len(self.data)
# 
#   def reset(self):
#     self.ind = 0
# 
#   def get(self):
#     """iterate the dataset with random order."""
#     # Shuffle before the iteration starts.
#     # The iteration should better be separated with collection.
#     if self.ind == 0:
#       random.shuffle(self.data)
#     ret = self.data[self.ind]
#     self.ind += 1
#     if self.ind == self.size:
#       self.ind = 0
#     return copy.deepcopy(ret)
# 
# Path: difflogic/thutils.py
# def monitor_gradrms(model):
#   """Monitor the rms of the gradients of the parameters."""
#   monitors = {}
#   for name, p in model.named_parameters():
#     if p.grad is not None:
#       monitors['gradrms/' + name] = rms(p.grad) / max(rms(p), 1e-8)
#   return monitors


import argparse
import collections
import copy
import jacinle.random as random
from ..tqdm_utils import tqdm_for
from difflogic.dataset.utils import RandomlyIterDataset
from difflogic.thutils import monitor_gradrms
from jacinle.logging import get_logger
from jacinle.utils.meter import GroupMeters
from jacinle.utils.tqdm import tqdm_pbar
from jactorch.train.env import TrainerEnv
from jactorch.utils.meta import as_cuda

    Returns:
      succ: The result of the episode, success or not, to classify as pos/neg.
      number: The number of objects/blocks
      backup: The clone of the environment, for interacting multiple times.
    """
    raise NotImplementedError()

  def _get_train_data(self, index, meters):
    return self._get_result(index, meters, mode='train')

  def _inherit_neg_data(self, neg_data, old_neg_data, meters,
                        mining_dataset_size):
    """To avoid wasting already collect negative data, re-exam them."""
    if not self.inherit_neg_data or \
            (old_neg_data is None or old_neg_data.size == 0):
      return
    original_size = neg_data.size
    old_neg_data.reset()
    maximum_inherit_size = min(old_neg_data.size, mining_dataset_size)

    def inherit_func(index):
      number, player = old_neg_data.get()
      message, result = self._get_result_given_player(
          index, meters, number, player, mode='inherit')
      positive, number, backup = self._extract_info(result)
      if not positive:
        neg_data.append((number, backup))
      return message


# Path: difflogic/tqdm_utils.py
# def tqdm_for(total, func):
#   """wrapper of the for function with message showing on the progress bar."""
#   # Not support break cases for now.
#   with tqdm_pbar(total=total) as pbar:
#     for i in range(total):
#       message = func(i)
#       pbar.set_description(message)
#       pbar.update()
# 
# Path: difflogic/dataset/utils.py
# class RandomlyIterDataset(object):
#   """Collect data and iterate the dataset in random order."""
# 
#   def __init__(self):
#     super().__init__()
#     self.data = []
#     self.ind = 0
# 
#   def append(self, data):
#     self.data.append(data)
# 
#   @property
#   def size(self):
#     return len(self.data)
# 
#   def reset(self):
#     self.ind = 0
# 
#   def get(self):
#     """iterate the dataset with random order."""
#     # Shuffle before the iteration starts.
#     # The iteration should better be separated with collection.
#     if self.ind == 0:
#       random.shuffle(self.data)
#     ret = self.data[self.ind]
#     self.ind += 1
#     if self.ind == self.size:
#       self.ind = 0
#     return copy.deepcopy(ret)
# 
# Path: difflogic/thutils.py
# def monitor_gradrms(model):
#   """Monitor the rms of the gradients of the parameters."""
#   monitors = {}
#   for name, p in model.named_parameters():
#     if p.grad is not None:
#       monitors['gradrms/' + name] = rms(p.grad) / max(rms(p), 1e-8)
#   return monitors


# Path: difflogic/train/train.py
import argparse
import collections
import copy
import jacinle.random as random
from ..tqdm_utils import tqdm_for
from difflogic.dataset.utils import RandomlyIterDataset
from difflogic.thutils import monitor_gradrms
from jacinle.logging import get_logger
from jacinle.utils.meter import GroupMeters
from jacinle.utils.tqdm import tqdm_pbar
from jactorch.train.env import TrainerEnv
from jactorch.utils.meta import as_cuda


    Returns:
      succ: The result of the episode, success or not, to classify as pos/neg.
      number: The number of objects/blocks
      backup: The clone of the environment, for interacting multiple times.
    """
    raise NotImplementedError()

  def _get_train_data(self, index, meters):
    return self._get_result(index, meters, mode='train')

  def _inherit_neg_data(self, neg_data, old_neg_data, meters,
                        mining_dataset_size):
    """To avoid wasting already collect negative data, re-exam them."""
    if not self.inherit_neg_data or \
            (old_neg_data is None or old_neg_data.size == 0):
      return
    original_size = neg_data.size
    old_neg_data.reset()
    maximum_inherit_size = min(old_neg_data.size, mining_dataset_size)

    def inherit_func(index):
      number, player = old_neg_data.get()
      message, result = self._get_result_given_player(
          index, meters, number, player, mode='inherit')
      positive, number, backup = self._extract_info(result)
      if not positive:
        neg_data.append((number, backup))
      return message


    tqdm_for(maximum_inherit_size, inherit_func)