while True:
        with open(out_file + "_%d" % i, "wb") as f:
          i += 1
          for _ in range(100):
            for j in sess.run(iterator.get_next()):
              tf.logging.info(j)
              f.write(bytearray(j))

"""TensorFlow NMT model implementation."""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import argparse
import os
import random
import sys

import numpy as np
import six
from six.moves import range
import tensorflow.compat.v1 as tf

# copybara:strip_begin
from REDACTED import REDACTED
from REDACTED.tensorflow.contrib import training as contrib_training
# copybara:strip_end
from REDACTED.tensorflow.contrib.training.python.training import evaluation
from REDACTED.tensorflow.python.ops import control_flow_util
from REDACTED.mlp_log import mlp_log
from REDACTED.nmt import estimator
from REDACTED.nmt.utils import iterator_utils
from REDACTED.nmt.utils import misc_utils as utils
from REDACTED.nmt.utils import vocab_utils

utils.check_tensorflow_version()

FLAGS = None

def add_arguments(parser):
  """Build ArgumentParser."""
  parser.register("type", "bool", lambda v: v.lower() == "true")

  # network
  parser.add_argument(
      "--num_units", type=int, default=1024, help="Network size.")
  parser.add_argument(
      "--num_layers", type=int, default=4, help="Network depth.")
  parser.add_argument("--num_encoder_layers", type=int, default=None,
                      help="Encoder depth, equal to num_layers if None.")
  parser.add_argument("--num_decoder_layers", type=int, default=None,
                      help="Decoder depth, equal to num_layers if None.")
  parser.add_argument("--num_embeddings_partitions", type=int, default=0,
                      help="Number of partitions for embedding vars.")

  # optimizer
  parser.add_argument(
      "--optimizer", type=str, default="adam", help="sgd | adam")
  parser.add_argument(
      "--learning_rate",
      type=float,
      default=0.001,
      help="Learning rate. Adam: 0.001 | 0.0001")
  parser.add_argument(
      "--warmup_steps",
      type=int,
      default=200,
      help="How many steps we inverse-decay learning.")
  parser.add_argument("--warmup_scheme", type=str, default="t2t", help="""\
      How to warmup learning rates. Options include:
        t2t: Tensor2Tensor's way, start with lr 100 times smaller, then
             exponentiate until the specified lr.\
      """)
  parser.add_argument(
      "--decay_start", type=int, default=3000, help="step to start decay")
  parser.add_argument(
      "--decay_interval",
      type=int,
      default=400,
      help="interval steps between 2 decays")
  parser.add_argument(
      "--decay_steps", type=int, default=5, help="number of decays")
  parser.add_argument(
      "--decay_factor", type=float, default=0.66, help="decay rate")

  parser.add_argument(
      "--max_train_epochs", type=int, default=8,
      help="Maximum number of training epochs.")
  parser.add_argument("--num_examples_per_epoch", type=int, default=3442299,
                      help="Number of examples in one epoch")
  parser.add_argument("--label_smoothing", type=float, default=0.1,
                      help=("If nonzero, smooth the labels towards "
                            "1/num_classes."))

  parser.add_argument("--init_op", type=str, default="uniform",
                      help="uniform | glorot_normal | glorot_uniform")
  parser.add_argument("--init_weight", type=float, default=0.1,
                      help=("for uniform init_op, initialize weights "
                            "between [-this, this]."))

  # data
  parser.add_argument(
      "--src", type=str, default="en", help="Source suffix, e.g., en.")
  parser.add_argument(
      "--tgt", type=str, default="de", help="Target suffix, e.g., de.")
  parser.add_argument(
      "--data_dir", type=str, default="", help="Training/eval data directory.")

  parser.add_argument(
      "--train_prefix",
      type=str,
      default="train.tok.clean.bpe.32000",
      help="Train prefix, expect files with src/tgt suffixes.")
  parser.add_argument(
      "--test_prefix",
      type=str,
      default="newstest2014",
      help="Test prefix, expect files with src/tgt suffixes.")
  parser.add_argument(
      "--use_preprocessed_data",
      type="bool",
      default=True,
      help="Whether to use preprocessed training data.")

  parser.add_argument(
      "--out_dir", type=str, default=None, help="Store log/model files.")

  # Vocab
  parser.add_argument(
      "--vocab_prefix",
      type=str,
      default="vocab.bpe.32000",
      help="""\
      Vocab prefix, expect files with src/tgt suffixes.\
      """)

  parser.add_argument("--check_special_token", type="bool", default=True,
                      help="""\
                      Whether check special sos, eos, unk tokens exist in the
                      vocab files.\
                      """)

  # Sequence lengths
  parser.add_argument(
      "--src_max_len",
      type=int,
      default=48,
      help="Max length of src sequences during training.")
  parser.add_argument(
      "--tgt_max_len",
      type=int,
      default=48,
      help="Max length of tgt sequences during training.")
  parser.add_argument(
      "--src_max_len_infer",
      type=int,
      default=160,
      help="Max length of src sequences during inference.")
  parser.add_argument(
      "--tgt_max_len_infer",
      type=int,
      default=160,
      help="""\
      Max length of tgt sequences during inference.  Also use to restrict the
      maximum decoding length.\
      """)

  # Default settings works well (rarely need to change)
  parser.add_argument("--forget_bias", type=float, default=0.0,
                      help="Forget bias for BasicLSTMCell.")
  parser.add_argument("--dropout", type=float, default=0.2,
                      help="Dropout rate (not keep_prob)")
  parser.add_argument("--max_gradient_norm", type=float, default=5.0,
                      help="Clip gradients to this norm.")
  parser.add_argument("--batch_size", type=int, default=512, help="Batch size.")

  parser.add_argument("--steps_per_stats", type=int, default=5,
                      help=("How many training steps to do per stats logging."
                            "Save checkpoint every 10x steps_per_stats"))
  parser.add_argument(
      "--num_buckets",
      type=int,
      default=5,
      help="Put data into similar-length buckets.")
  parser.add_argument(
      "--choose_buckets",
      type=int,
      default=1,
      help="Choose from this number of length buckets per training step.")

  # SPM
  parser.add_argument("--subword_option", type=str, default="bpe",
                      choices=["", "bpe", "spm"],
                      help="""\
                      Set to bpe or spm to activate subword desegmentation.\
                      """)

  # Misc
  parser.add_argument(
      "--num_shards", type=int,
      default=8, help="Number of shards (TPU cores).")
  parser.add_argument(
      "--num_shards_per_host", type=int,
      default=8, help="Number of shards (TPU cores) per host.")
  parser.add_argument(
      "--num_gpus", type=int, default=4, help="Number of gpus in each worker.")
  parser.add_argument(
      "--num_infeed_workers",
      type=int,
      default=1,
      help="Number of TPU workers used for input generation.")
  parser.add_argument(
      "--num_tpu_workers",
      type=int,
      default=1,
      help="Number of TPU workers; if set, uses the distributed-sync pipeline.")
  parser.add_argument("--hparams_path", type=str, default=None,
                      help=("Path to standard hparams json file that overrides"
                            "hparams values from FLAGS."))
  parser.add_argument(
      "--random_seed",
      type=int,
      default=None,
      help="Random seed (>0, set a specific seed).")

  # Inference
  parser.add_argument("--ckpt", type=str, default="",
                      help="Checkpoint file to load a model for inference.")
  parser.add_argument(
      "--infer_batch_size",
      type=int,
      default=512,
      help="Batch size for inference mode.")
  parser.add_argument(
      "--examples_to_infer",
      type=int,
      default=3003,
      help="Number of examples to infer.")
  parser.add_argument("--detokenizer_file", type=str,
                      default="mosesdecoder/scripts/tokenizer/detokenizer.perl",
                      help=("""Detokenizer script file."""))
  parser.add_argument("--use_REDACTED", type=bool, default=False)
  parser.add_argument(
      "--target_bleu", type=float, default=24.0, help="Target accuracy.")

  # Advanced inference arguments
  parser.add_argument("--infer_mode", type=str, default="beam_search",
                      choices=["greedy", "sample", "beam_search"],
                      help="Which type of decoder to use during inference.")
  parser.add_argument("--beam_width", type=int, default=5,
                      help=("""\
      beam width when using beam search decoder. If 0 (default), use standard
      decoder with greedy helper.\
      """))
  parser.add_argument(
      "--length_penalty_weight",
      type=float,
      default=0.6,
      help="Length penalty for beam search.")
  parser.add_argument(
      "--coverage_penalty_weight",
      type=float,
      default=0.1,
      help="Coverage penalty for beam search.")

  # Job info
  parser.add_argument("--jobid", type=int, default=0,
                      help="Task id of the worker.")

  # TPU
  parser.add_argument("--use_tpu", type=bool, default=True)
  parser.add_argument("--master", type=str, default="",
                      help=("Address of the master. Either --master or "
                            "--tpu_name must be specified."))
  parser.add_argument("--tpu_name", type=str, default=None,
                      help=("Name of the TPU for Cluster Resolvers. Either "
                            "--tpu_name or --master must be specified."))
  parser.add_argument("--use_dynamic_rnn", type=bool, default=False)
  parser.add_argument("--use_synthetic_data", type=bool, default=False)
  parser.add_argument(
      "--mode",
      type=str,
      default="train_and_eval",
      choices=["train", "train_and_eval", "infer", "preprocess"])
  parser.add_argument(
      "--activation_dtype",
      type=str,
      default="bfloat16",
      choices=["float32", "bfloat16"])
  parser.add_argument("--tpu_job_name", type=str, default=None)

  # copybara:strip_begin
  # Vizier
  parser.add_argument("--client_handle", type=str, default="",
                      help=("Client_handle for the tuner."))
  parser.add_argument("--study_name", type=str, default=None,
                      help=("Name of Vizier hparams tuning study."))
  parser.add_argument("--REDACTED", type=int,
                      default=REDACTED.StudyConfig.RANDOM_SEARCH,
                      help=("Vizier search algorithm to use."))
  # copybara:strip_end

def create_hparams(flags):
  """Create training hparams."""
  return contrib_training.HParams(
      # Data
      src=flags.src,
      tgt=flags.tgt,
      train_prefix=flags.data_dir + flags.train_prefix,
      test_prefix=flags.data_dir + flags.test_prefix,
      vocab_prefix=flags.data_dir + flags.vocab_prefix,
      out_dir=flags.out_dir,

      # Networks
      num_units=flags.num_units,
      num_encoder_layers=(flags.num_encoder_layers or flags.num_layers),
      num_decoder_layers=(flags.num_decoder_layers or flags.num_layers),
      dropout=flags.dropout,
      num_embeddings_partitions=flags.num_embeddings_partitions,

      # Train
      optimizer=flags.optimizer,
      max_train_epochs=flags.max_train_epochs,
      num_examples_per_epoch=flags.num_examples_per_epoch,
      batch_size=flags.batch_size,
      num_train_steps=int(flags.num_examples_per_epoch / flags.batch_size *
                          flags.max_train_epochs),
      init_op=flags.init_op,
      init_weight=flags.init_weight,
      max_gradient_norm=flags.max_gradient_norm,
      learning_rate=flags.learning_rate,
      label_smoothing=flags.label_smoothing,
      warmup_steps=flags.warmup_steps,
      warmup_scheme=flags.warmup_scheme,
      decay_start=flags.decay_start,
      decay_interval=flags.decay_interval,
      decay_steps=flags.decay_steps,
      decay_factor=flags.decay_factor,

      # Data constraints
      num_buckets=flags.num_buckets,
      choose_buckets=flags.choose_buckets,
      src_max_len=flags.src_max_len,
      tgt_max_len=flags.tgt_max_len,
      use_preprocessed_data=flags.use_preprocessed_data,

      # Inference
      src_max_len_infer=flags.src_max_len_infer,
      tgt_max_len_infer=flags.tgt_max_len_infer,
      infer_batch_size=flags.infer_batch_size,
      examples_to_infer=flags.examples_to_infer,
      detokenizer_file=flags.data_dir + flags.detokenizer_file,
      use_REDACTED=flags.use_REDACTED,
      target_bleu=flags.target_bleu,

      # Advanced inference arguments
      infer_mode=flags.infer_mode,
      beam_width=flags.beam_width,
      length_penalty_weight=flags.length_penalty_weight,
      coverage_penalty_weight=flags.coverage_penalty_weight,

      # Vocab
      sos=vocab_utils.SOS,
      eos=vocab_utils.EOS,
      subword_option=flags.subword_option,
      check_special_token=flags.check_special_token,

      # Misc
      forget_bias=flags.forget_bias,
      num_shards=flags.num_shards,
      num_shards_per_host=flags.num_shards_per_host,
      num_gpus=flags.num_gpus,
      num_infeed_workers=flags.num_infeed_workers,
      epoch_step=0,  # record where we were within an epoch.
      steps_per_stats=flags.steps_per_stats,
      random_seed=flags.random_seed,

      # TPU
      use_tpu=flags.use_tpu,
      master=flags.master,
      tpu_name=flags.tpu_name,
      use_dynamic_rnn=flags.use_dynamic_rnn,
      use_synthetic_data=flags.use_synthetic_data,
      mode=flags.mode,
      activation_dtype=flags.activation_dtype,
      tpu_job_name=flags.tpu_job_name)

def _add_argument(hparams, key, value, update=True):
  """Add an argument to hparams; if exists, change the value if update==True."""
  if hasattr(hparams, key):
    if update:
      setattr(hparams, key, value)
  else:
    hparams.add_hparam(key, value)

def extend_hparams(hparams):
  """Add new arguments to hparams."""
  # Sanity checks
  if hparams.subword_option and hparams.subword_option not in ["spm", "bpe"]:
    raise ValueError("subword option must be either spm, or bpe")
  if hparams.infer_mode == "beam_search" and hparams.beam_width <= 0:
    raise ValueError("beam_width must greater than 0 when using beam_search"
                     "decoder.")

  # Different number of encoder / decoder layers
  assert hparams.num_encoder_layers == hparams.num_decoder_layers

  # The first unidirectional layer (after the bi-directional layer) in
  # the GNMT encoder can't have residual connection due to the input is
  # the concatenation of fw_cell and bw_cell's outputs.
  num_encoder_residual_layers = hparams.num_encoder_layers - 2
  num_decoder_residual_layers = num_encoder_residual_layers
  _add_argument(hparams, "num_encoder_residual_layers",
                num_encoder_residual_layers)
  _add_argument(hparams, "num_decoder_residual_layers",
                num_decoder_residual_layers)

  # Get vocab file names first
  if hparams.vocab_prefix:
    src_vocab_file = six.ensure_str(
        hparams.vocab_prefix) + "." + six.ensure_str(hparams.src)
    tgt_vocab_file = six.ensure_str(
        hparams.vocab_prefix) + "." + six.ensure_str(hparams.tgt)
  else:
    raise ValueError("hparams.vocab_prefix must be provided.")

  # Source vocab
  src_vocab_size, src_vocab_file = vocab_utils.check_vocab(
      src_vocab_file,
      hparams.out_dir,
      check_special_token=hparams.check_special_token,
      sos=hparams.sos,
      eos=hparams.eos,
      unk=vocab_utils.UNK)

  # Target vocab
  utils.print_out("  using source vocab for target")
  tgt_vocab_file = src_vocab_file
  tgt_vocab_size = src_vocab_size
  _add_argument(hparams, "src_vocab_size", src_vocab_size)
  _add_argument(hparams, "tgt_vocab_size", tgt_vocab_size)
  _add_argument(hparams, "src_vocab_file", src_vocab_file)
  _add_argument(hparams, "tgt_vocab_file", tgt_vocab_file)

  # Num embedding partitions
  _add_argument(
      hparams, "num_enc_emb_partitions", hparams.num_embeddings_partitions)
  _add_argument(
      hparams, "num_dec_emb_partitions", hparams.num_embeddings_partitions)

  # Pretrained Embeddings
  _add_argument(hparams, "src_embed_file", "")
  _add_argument(hparams, "tgt_embed_file", "")

  return hparams

def create_or_load_hparams(default_hparams, hparams_path):
  """Create hparams or load hparams from out_dir."""
  hparams = utils.maybe_parse_standard_hparams(default_hparams, hparams_path)
  hparams = extend_hparams(hparams)
  # Print HParams
  utils.print_hparams(hparams)
  return hparams

def prepare_dataset(flags):
  """Generate the preprocessed dataset."""
  src_file = "%s.%s" % (flags.data_dir + flags.train_prefix, flags.src)
  tgt_file = "%s.%s" % (flags.data_dir + flags.train_prefix, flags.tgt)
  vocab_file = flags.data_dir + flags.vocab_prefix
  _, vocab_file = vocab_utils.check_vocab(vocab_file, flags.out_dir)
  out_file = six.ensure_str(flags.out_dir) + "preprocessed_dataset"
  src_vocab_table, tgt_vocab_table = vocab_utils.create_vocab_tables(vocab_file)
  src_dataset = tf.data.TextLineDataset(src_file)
  tgt_dataset = tf.data.TextLineDataset(tgt_file)
  iterator = iterator_utils.get_iterator(
      src_dataset,
      tgt_dataset,
      src_vocab_table,
      tgt_vocab_table,
      batch_size=1,
      global_batch_size=1,
      sos=vocab_utils.SOS,
      eos=vocab_utils.EOS,
      random_seed=1,
      num_buckets=flags.num_buckets,
      src_max_len=flags.src_max_len,
      tgt_max_len=flags.tgt_max_len,
      filter_oversized_sequences=True,
      return_raw=True).make_initializable_iterator()

  with tf.Session() as sess:
    sess.run(tf.tables_initializer())
    sess.run(iterator.initializer)
    try:
      i = 0

    except tf.errors.OutOfRangeError:
      pass

def run_main(flags, default_hparams, estimator_fn):
  """Run main."""
  # Job
  jobid = flags.jobid
  utils.print_out("# Job id %d" % jobid)

  # Random
  random_seed = flags.random_seed
  if random_seed is not None and random_seed > 0:
    utils.print_out("# Set random seed to %d" % random_seed)
    random.seed(random_seed + jobid)
    np.random.seed(random_seed + jobid)
    tf.set_random_seed(random_seed)

  mlp_log.mlperf_print("cache_clear", True)
  mlp_log.mlperf_print("init_start", None)
  mlp_log.mlperf_print("submission_benchmark", "resnet")
  mlp_log.mlperf_print("submission_division", "closed")
  mlp_log.mlperf_print("submission_org", "google")
  mlp_log.mlperf_print("submission_platform", "tpu-v3-%d" % FLAGS.num_shards)
  mlp_log.mlperf_print("submission_status", "research")

  mlp_log.mlperf_print("global_batch_size", FLAGS.batch_size)
  mlp_log.mlperf_print("opt_learning_rate_alt_decay_func", "True")
  mlp_log.mlperf_print("opt_base_learning_rate", FLAGS.learning_rate)
  mlp_log.mlperf_print("opt_learning_rate_decay_interval", FLAGS.decay_interval)
  mlp_log.mlperf_print("opt_learning_rate_decay_factor", FLAGS.decay_factor)
  mlp_log.mlperf_print("opt_learning_rate_decay_steps", FLAGS.decay_steps)
  mlp_log.mlperf_print("opt_learning_rate_remain_steps", FLAGS.decay_start)
  mlp_log.mlperf_print("opt_learning_rate_alt_warmup_func", FLAGS.warmup_scheme)
  mlp_log.mlperf_print("opt_learning_rate_warmup_steps", FLAGS.warmup_steps)
  mlp_log.mlperf_print(
      "max_sequence_length", FLAGS.src_max_len, metadata={"method": "discard"})
  mlp_log.mlperf_print("train_samples", FLAGS.num_examples_per_epoch)
  mlp_log.mlperf_print("eval_samples", FLAGS.examples_to_infer)

  # Model output directory
  out_dir = flags.out_dir
  if out_dir and not tf.gfile.Exists(out_dir):
    utils.print_out("# Creating output directory %s ..." % out_dir)
    tf.gfile.MakeDirs(out_dir)

  # Load hparams.
  hparams = create_or_load_hparams(default_hparams, flags.hparams_path)

  # Train or Evaluation
  return estimator_fn(hparams)

def main(unused_argv):
  control_flow_util.ENABLE_CONTROL_FLOW_V2 = True

  if FLAGS.mode == "preprocess":
    prepare_dataset(FLAGS)
  elif FLAGS.mode == "train":
    print("Running training mode.")
    default_hparams = create_hparams(FLAGS)
    run_main(FLAGS, default_hparams, estimator.train_fn)
  elif FLAGS.mode == "train_and_eval":
    print("Running training and evaluation mode.")
    default_hparams = create_hparams(FLAGS)
    run_main(FLAGS, default_hparams,
             estimator.train_and_eval_with_low_level_api)
  else:
    print("Running inference mode.")
    default_hparams = create_hparams(FLAGS)
    current_epoch = 0
    last_step = 0
    # Run evaluation when there's a new checkpoint
    for ckpt in evaluation.checkpoints_iterator(FLAGS.out_dir):
      # Terminate eval job once target score is reached
      current_step = int(six.ensure_str(os.path.basename(ckpt)).split("-")[1])
      if current_step <= last_step:
        continue
      last_step = current_step
      tf.logging.info("Starting to evaluate...%s", ckpt)
      try:
        score = run_main(FLAGS, default_hparams, estimator.eval_fn)
        current_epoch += 1
        if score > FLAGS.target_bleu:
          tf.logging.info(
              "Evaluation finished after training step %d" % current_step)
          break
        # Terminate eval job when final checkpoint is reached
        max_steps = default_hparams.num_train_steps
        if current_step >= max_steps:
          tf.logging.info(
              "Evaluation finished but failed to reach target score.")
          break

      except tf.errors.NotFoundError:
        tf.logging.info(
            "Checkpoint %s no longer exists, skipping checkpoint" % ckpt)

if __name__ == "__main__":
  tf.logging.set_verbosity(tf.logging.INFO)
  nmt_parser = argparse.ArgumentParser()
  add_arguments(nmt_parser)
  FLAGS, unparsed = nmt_parser.parse_known_args()
  tf.app.run(main=main, argv=[sys.argv[0]] + unparsed)


        mod = Modulator(
                name=self.name,
                start=r_start,
                profile=( (r_end-r_start,), (1,) ) ,
                generator=state,
                )
        self.mod_list.append(mod)

from __future__ import print_function, division, absolute_import
from .mod_util import Attr

class Modulator(object):
    """ A Helper class that make generator of pydot modulator string easier.

    Attributes
    -----------------
    name : str, 'modulator'
        redundant attribute, can be only 'modulator'
    start : str
        start time
    profile : tuple of tuple
        See Anomaly.get_profile_with_ano
    """
    def __init__(self, **desc):
        self.desc = desc

    def __getitem__(self, name):
        return self.desc[name]

    def __str__(self):
        return str( Attr(**self.desc) )

    @property
    def profile(self): return self.desc['profile']
    @property
    def start(self):
        return self.desc['start']

from .Behaviour import DTMCBehaviour, CTMCBehaviour
class MarkovModulator(Modulator):
    """ Markov Modulator

    Parameters
    -------------------
    name : str
        name of modular. now it is useless
    start : str
        a expression that can generate start time (using eval)
    profile : tuple of tuple
        profile
    generator_states : a list of generator names
        representing the states
    node_para : dict
        * type : {'DTMC', 'CTMC'}
        * P : nxn matrix
            Pij is transition probability from i to j when type=='DTMC',
            Pij is rate of possion process that transition i to j will
            happen when type=='CTMC'
        * interval : float
            mean of interval time. useful for only 'DTMC'

    """
    def __init__(self, name, start, profile, generator_states, node_para):
        # start is an expression that could generate a start time (randomly). We
        # want to frozen start time here.
        self.frozen_start_time = eval(start)
        Modulator.__init__(self, name=name, start=str(self.frozen_start_time),
                           profile=profile)
        self.name = name
        self.states = generator_states
        # self.start = start
        # self.profile = profile
        self.mod_list = []
        type_ = node_para['type']
        P = node_para['P']
        if type_ == 'DTMC':
            it = node_para['interval']
            self.behaviour = DTMCBehaviour(P, generator_states, it)
        elif type_ == 'CTMC':
            self.behaviour = CTMCBehaviour(P, generator_states)
        else:
            raise Exception("unknown Behaviour Type")

        self.sync()

    def sync(self):
        self.behaviour.behave_with_profile(self.frozen_start_time, self.profile, self.stage)

    def __str__(self):
        return ' '.join([str(mod) for mod in self.mod_list])

    def stage(self, r_start, r_end, state):
        """ One markov stage. A new modulator will be added for each stage """

#     Modulator defines the behaviour of generator within a range of time.
#     implement the stage function.
#     joint_dist should be numpy array


while stack:
            node, visited = stack.pop()
            if not visited:
                if node.right:
                    stack.append([node.right, False])
                stack.append([node, True])
                if node.left:
                    stack.append([node.left, False])
            else:
                if pre == -1:
                    pre = node.val
                else:
                    ret = min(ret, abs(node.val - pre))
                    pre = node.val

# Definition for a binary tree node.
#         self.val = x

class Solution:
    def getMinimumDifference(self, root: TreeNode) -> int:
        ret = float('inf')
        pre = -1
        stack = [[root, False]]


    try:
        yield
    finally:
        try:
            delattr(obj, attr)
        except AttributeError:
            pass
        obj._accessors.discard(attr)

import contextlib

import pytest

import pandas as pd
import pandas._testing as tm

@contextlib.contextmanager
def ensure_removed(obj, attr):
    """Ensure that an attribute added to 'obj' during the test is
    removed when we're done"""

class MyAccessor:
    def __init__(self, obj):
        self.obj = obj
        self.item = "item"

    @property
    def prop(self):
        return self.item

    def method(self):
        return self.item

@pytest.mark.parametrize(
    "obj, registrar",
    [
        (pd.Series, pd.api.extensions.register_series_accessor),
        (pd.DataFrame, pd.api.extensions.register_dataframe_accessor),
        (pd.Index, pd.api.extensions.register_index_accessor),
    ],
)
def test_register(obj, registrar):
    with ensure_removed(obj, "mine"):
        before = set(dir(obj))
        registrar("mine")(MyAccessor)
        o = obj([]) if obj is not pd.Series else obj([], dtype=object)
        assert o.mine.prop == "item"
        after = set(dir(obj))
        assert (before ^ after) == {"mine"}
        assert "mine" in obj._accessors

def test_accessor_works():
    with ensure_removed(pd.Series, "mine"):
        pd.api.extensions.register_series_accessor("mine")(MyAccessor)

        s = pd.Series([1, 2])
        assert s.mine.obj is s

        assert s.mine.prop == "item"
        assert s.mine.method() == "item"

def test_overwrite_warns():
    # Need to restore mean
    mean = pd.Series.mean
    try:
        with tm.assert_produces_warning(UserWarning) as w:
            pd.api.extensions.register_series_accessor("mean")(MyAccessor)
            s = pd.Series([1, 2])
            assert s.mean.prop == "item"
        msg = str(w[0].message)
        assert "mean" in msg
        assert "MyAccessor" in msg
        assert "Series" in msg
    finally:
        pd.Series.mean = mean

def test_raises_attribute_error():

    with ensure_removed(pd.Series, "bad"):

        @pd.api.extensions.register_series_accessor("bad")
        class Bad:
            def __init__(self, data):
                raise AttributeError("whoops")

        with pytest.raises(AttributeError, match="whoops"):
            pd.Series([], dtype=object).bad


        tstr = self.make_sql(**kwargs)
        all_data = None
        try:
            csr = self._db.cursor()
            csr.execute(tstr)
            # print tstr
            all_data = csr.fetchall()
            csr.close()
        except:
            traceback.print_exc()
            self._db.rollback()

        return self.pack(all_data)

import traceback
from base64 import b64encode, b64decode

import MySQLdb

from pwnlib.logdata import logdata
from .log import getLogger

'''
A class for log rotate. this class will get data from sql database and print data or pack as a josn
'''

class logrotate(object):
    _sql = 'SELECT * FROM connections '
    _sql_id = 'con_id = %d '
    _sql_hash = 'con_hash = "%s" '
    _sql_token = 'token = "%s" '
    _sql_host = 'host = "%s" '
    _sql_ip = 'ip = "%s" '
    _sql_dport = 'dport = %d '
    _sql_sport = 'sport = %d'
    _sql_con_time = 'con_time > %d '
    _sql_fin_time = 'fin_time < %d '
    _sql_target = 'target = "%s" '

    _sql_flow = 'SELECT * FROM flow WHERE con_hash = "%s" '

    def __init__(self, sqluser, sqlpwd, host='localhost', database='pwnlog'):
        """
        The user name and password to the mysql database. the default database if pwnlog@localhost.
        """
        self._db = MySQLdb.connect(host, sqluser, sqlpwd, database)

    def find(self, **kwargs):
        """
        Find data.The key of kwargs is also the field of connections.
        use con_time and fin_time select the log between the times
        It's return a list of logdata class.Every logdata class is a pack of one pwn attach
        """

    def make_sql(self, **kwargs):
        """
        get the enter limited and return a sql cmd.
        """

        dataList = []
        dataList.append((self._sql_id, kwargs.get('con_id', None)))
        dataList.append((self._sql_hash, kwargs.get('con_hash', None)))
        token = kwargs.get('token', None)
        if token != None:
            token = b64encode(token)
        dataList.append((self._sql_token, token))
        dataList.append((self._sql_host, kwargs.get('host', None)))
        dataList.append((self._sql_ip, kwargs.get('ip', None)))
        dataList.append((self._sql_dport, kwargs.get('dport', None)))
        dataList.append((self._sql_sport, kwargs.get('sport', None)))
        dataList.append((self._sql_con_time, kwargs.get('con_time', None)))
        dataList.append((self._sql_fin_time, kwargs.get('fin_time', None)))
        dataList.append((self._sql_target, kwargs.get('target', None)))

        # make sql just add string together
        flag = False
        tstr = self._sql
        for data in dataList:
            if data[1] != None:
                if flag == False:
                    tstr += 'WHERE '
                    flag = True
                else:
                    tstr += 'AND '
                tstr += data[0] % data[1]

        return tstr

    def pack(self, allData):
        '''
        split the data to a list.and change the sql select data to a dict
        '''
        dataList = []
        for row in allData:
            IO_data = self.get_IO_data(row[1])  # use con_hash to find send and recv data in flow
            temp_dict = {}
            temp_dict['token'] = b64decode(row[2])
            temp_dict['host'] = row[3]
            temp_dict['ip'] = row[4]
            temp_dict['dport'] = row[5]
            temp_dict['sport'] = row[6]
            temp_dict['con_time'] = row[7]
            temp_dict['fin_time'] = row[8]
            temp_dict['target'] = row[9]
            temp_dict['data'] = IO_data
            dataList.append(logdata(temp_dict))
        return dataList

    def get_IO_data(self, hash):
        """
        It's will use the hash two find the send and recv dates in flow table.And the data will sort by time.
        """
        temp_data = ()
        data = []
        try:
            csr = self._db.cursor()
            csr.execute(self._sql_flow % hash)
            temp_data = csr.fetchall()
            csr.close()
        except:
            traceback.print_exc()
            self._db.rollback()

        for row in temp_data:
            data.append((row[2], row[3], b64decode(row[4])))

        data.sort()

        return data


        ans = self.extensions(all=all, key=key)
        if len(ans) == 1:
            # There is a single service, so return it.
            return ans[0]
        elif not ans:
            return None
        else:
            raise PluginError("The ExtensionPoint does not have a unique "
                              "service!  %d services are defined for interface"
                              " '%s' (key=%s)." %
                              (len(ans), self._interface.__name__, str(key)))

#  ___________________________________________________________________________
#  Pyomo: Python Optimization Modeling Objects
#  Engineering Solutions of Sandia, LLC, the U.S. Government retains certain
#  rights in this software.
#  ___________________________________________________________________________
#  This module was originally developed as part of the PyUtilib project
#  Under the terms of Contract DE-AC04-94AL85000 with Sandia Corporation,
#  the U.S. Government retains certain rights in this software.
#  ___________________________________________________________________________

import collections
import inspect
import sys
from weakref import ref as weakref_ref

from pyomo.common.errors import PyomoException
from pyomo.common.deprecation import deprecated, deprecation_warning

if sys.version_info[:2] >= (3,7):
    _deterministic_dict = dict
else:
    from pyomo.common.collections import OrderedDict
    _deterministic_dict = OrderedDict

class PluginGlobals(object):
    @staticmethod
    @deprecated("The PluginGlobals environment manager is deprecated: "
                "Pyomo only supports a single global environment",
                version='6.0')
    def add_env(name):
        pass

    @staticmethod
    @deprecated("The PluginGlobals environment manager is deprecated: "
                "Pyomo only supports a single global environment",
                version='6.0')
    def pop_env():
        pass

    @staticmethod
    @deprecated("The PluginGlobals environment manager is deprecated: "
                "Pyomo only supports a single global environment",
                version='6.0')
    def clear():
        pass

class PluginError(PyomoException):
    pass

def alias(name, doc=None, subclass=None):
    if subclass is not None:
        deprecation_warning(
            "The Pyomo plugin infrastructure alias() function does "
            "not support the subclass flag.", version='6.0')
    calling_frame = inspect.currentframe().f_back
    locals_ = calling_frame.f_locals
    # Some sanity checks
    assert locals_ is not calling_frame.f_globals and '__module__' in locals_, \
        'implements() can only be used in a class definition'
    locals_.setdefault('__plugin_aliases__', []).append((name, doc))

def implements(interface, inherit=None, namespace=None, service=False):
    if namespace is not None:
        deprecation_warning(
            "The Pyomo plugin infrastructure only supports a "
            "single global namespace.", version='6.0')
    calling_frame = inspect.currentframe().f_back
    locals_ = calling_frame.f_locals
    # Some sanity checks
    assert locals_ is not calling_frame.f_globals and '__module__' in locals_, \
        'implements() can only be used in a class definition'
    assert issubclass(interface, Interface)
    locals_.setdefault('__implements__', []).append(
        (interface, inherit, service)
    )

class InterfaceMeta(type):
    def __new__(cls, name, bases, classdict, *args, **kwargs):
        # Ensure that all interfaces have their own _plugins & _aliases
        # dictionaries
        classdict.setdefault('_next_id', 0)
        classdict.setdefault('_plugins', {})
        classdict.setdefault('_aliases', {})
        return super().__new__(cls, name, bases, classdict, *args, **kwargs)

class Interface(metaclass=InterfaceMeta):
    pass

class _deprecated_plugin_dict(dict):
    def __init__(self, name, classdict):
        super().__init__()
        msg = classdict.pop('__deprecated_message__', None)
        if not msg:
            msg = 'The %s interface has been deprecated' % (name,)
        version = classdict.pop('__deprecated_version__', None)
        remove_in = classdict.pop('__deprecated_remove_in__', None)
        self._deprecation_info = {
            'msg': msg, 'version': version, 'remove_in': remove_in
        }

    def __setitem__(self, key, val):
        deprecation_warning(**self._deprecation_info)
        super().__setitem__(key, val)

    def items(self):
        deprecation_warning(**self._deprecation_info)
        return super().items()

class DeprecatedInterfaceMeta(InterfaceMeta):
    def __new__(cls, name, bases, classdict, *args, **kwargs):
        classdict.setdefault(
            '_plugins', _deprecated_plugin_dict(name, classdict)
        )
        return super().__new__(cls, name, bases, classdict, *args, **kwargs)

class DeprecatedInterface(Interface, metaclass=DeprecatedInterfaceMeta):
    pass

class PluginMeta(type):

    def __new__(cls, name, bases, classdict, *args, **kwargs):
        # This plugin is a singleton plugin based on the __singleton__
        # singleton plugin
        _singleton = classdict.pop(
            '__singleton__',
            any(getattr(base, '__singleton__', None) is not None
                for base in bases)
        )
        # This prevents base class __singleton__, __plugin_aliases__,
        # and __implements__ from implicitly bleeding through and being
        # accidentally shared across subclasses.
        classdict['__singleton__'] = None
        aliases = classdict.setdefault('__plugin_aliases__', [])
        implements = classdict.setdefault('__implements__', [])
        # the same interface, use standard Python rules to determine
        # which implements() should govern (i.e. classdict supersedes
        interfaces = set(impl[0] for impl in implements)
        for base in bases:
            implements.extend(
                ep for ep in getattr(base, '__implements__', [])
                if ep[0] not in interfaces
            )
            interfaces.update(impl[0] for impl in implements)
        for interface, inherit, service in implements:
            if not inherit:
                continue
            if not any(issubclass(base, interface) for base in bases):
                bases = bases + (interface,)
                # Python requires that a class' metaclass be a
                # (nonstrict) subclass of the metaclasses of all its
                # base classes.  Check, and declare a new metaclass if
                if not issubclass(cls, type(interface)):
                    class tmp_meta(cls, type(interface)):
                        def __new__(cls, name, bases, classdict,
                                    *args, **kwargs):
                            # This is a plugin and not an Interface.  Do
                            # not set up dicts for the interface
                            classdict.setdefault('_plugins', None)
                            classdict.setdefault('_aliases', None)
                            return super().__new__(
                                cls, name, bases, classdict, *args, **kwargs)
                    cls = tmp_meta

        new_class = super().__new__(
            cls, name, bases, classdict, *args, **kwargs)

        # Register the new class with the interfaces
        for interface, inherit, service in implements:
            interface._plugins[new_class] = _deterministic_dict()
            interface._aliases.update(
                {name: (new_class, doc) for name, doc in aliases}
            )

        if _singleton:
            new_class.__singleton__ = new_class()

        return new_class

class Plugin(object, metaclass=PluginMeta):
    def __new__(cls):
        if cls.__singleton__ is not None:
            raise RuntimeError(
                "Cannot create multiple singleton plugin instances of type %s"
                % (cls,))
        obj = super().__new__(cls)
        obj._plugin_ids = {}
        # Record this instance (service) with all Interfaces
        for interface, inherit, service in cls.__implements__:
            _id = interface._next_id
            interface._next_id += 1
            obj._plugin_ids[interface] = _id
            interface._plugins[cls][_id] = (weakref_ref(obj), service)
        return obj

    def activate(self):
        cls = self.__class__
        for interface, inherit, service in cls.__implements__:
            _id = self._plugin_ids[interface]
            obj, service = interface._plugins[cls][_id]
            if not service:
                interface._plugins[cls][_id] = obj, True
    enable = activate

    def deactivate(self):
        cls = self.__class__
        for interface, inherit, service in cls.__implements__:
            _id = self._plugin_ids[interface]
            obj, service = interface._plugins[cls][_id]
            if service:
                interface._plugins[cls][_id] = obj, False
    disable = deactivate

    def enabled(self):
        cls = self.__class__
        return any(interface._plugins[cls][self._plugin_ids[interface]][1]
                   for interface, inherit, service in cls.__implements__)

class SingletonPlugin(Plugin):
    __singleton__ = True

class ExtensionPoint(object):
    def __init__(self, interface):
        assert issubclass(interface, Interface)
        self._interface = interface

    def __iter__(self, key=None, all=False):
        for cls, plugins in self._interface._plugins.items():
            remove = []
            for i, (obj, service) in plugins.items():
                if not obj():
                    remove.append(i)
                elif ((all or service) and
                      (key is None or key is cls or key == cls.__name__)):
                    yield obj()
            for i in remove:
                del plugins[i]

    def __len__(self):
        return len(list(self.__iter__()))

    def extensions(self, all=False, key=None):
        return list(self.__iter__(key=key, all=all))

    def __call__(self, key=None, all=False):
        return self.extensions(all=all, key=key)

    def service(self, key=None, all=False):
        """Return the unique service that matches the interface of this
        extension point.  An exception occurs if no service matches the
        specified key, or if multiple services match.
        """

class PluginFactory(object):

    def __init__(self, interface):
        self.interface = interface

    def __call__(self, name, *args, **kwds):
        name = str(name)
        if name not in self.interface._aliases:
            return None
        else:
            return self.interface._aliases[name][0](*args, **kwds)

    def services(self):
        return list(self.interface._aliases)

    def get_class(self, name):
        return self.interface._aliases.get(name, [None])[0]

    def doc(self, name):
        name = str(name)
        if name not in self.interface._aliases:
            return ""
        else:
            return self.interface._aliases[name][1]

    def deactivate(self, name):
        if isinstance(name, str):
            cls = self.get_class(name)
        if cls is None:
            return
        for service in ExtensionPoint(self.interface)(key=cls):
            service.deactivate()

    def activate(self, name):
        if isinstance(name, str):
            cls = self.get_class(name)
        if cls is None:
            return
        for service in ExtensionPoint(self.interface)(all=True, key=cls):
            service.activate()

# Old name for creating plugin factories
CreatePluginFactory = PluginFactory


	processor = OCRProcessor(kwargs)
	processor.traverse(data_path)

import click
import numpy as np
import logging
import functools

from pathlib import Path

from calamari_ocr.ocr import Predictor, MultiPredictor
from calamari_ocr.ocr.voting.confidence_voter import ConfidenceVoter

from origami.batch.core.processor import Processor
from origami.batch.core.io import Artifact, Stage, Input, Output
from origami.batch.core.lines import LineExtractor
from origami.batch.core.utils import RegionsFilter

class OCRProcessor(Processor):
	def __init__(self, options):
		super().__init__(options)
		self._options = options
		self._ocr = self._options["ocr"]

		if self._ocr == "FAKE":
			self._model_path = None
			self._models = []
			self._line_height = 48
			self._chunk_size = 1
		else:
			if not options["model"]:
				raise click.BadParameter(
					"Please specify a model path", param="model")
			self._model_path = Path(options["model"])

			models = list(self._model_path.glob("*.json"))
			if not options["legacy_model"]:
				models = [m for m in models if m.with_suffix(".h5").exists()]
			if len(models) < 1:
				raise FileNotFoundError(
					"no Calamari models found at %s" % self._model_path)
			self._models = models

			self._line_height = None
			self._chunk_size = None

		self._predictor = None
		self._voter = None

		self._ignored = RegionsFilter(options["ignore"])

		if self._ocr != "FULL":
			logging.getLogger().setLevel(logging.INFO)

	@staticmethod
	def options(f):
		options = [
			click.option(
				'--legacy-model',
				is_flag=True,
				default=False,
				help='support Calamari legacy models (pre 1.0)'),
			click.option(
				'-b', '--batch-size',
				type=int,
				default=-1,
				required=False),
			click.option(
				'--ignore',
				type=str,
				default="regions/ILLUSTRATION"),
			click.option(
				'--ocr',
				type=click.Choice(['FULL', 'DRY', 'FAKE'], case_sensitive=False),
				default="FULL")
		]

		return functools.reduce(lambda x, opt: opt(x), options, f)

	@property
	def processor_name(self):
		return __loader__.name

	def _load_models(self):
		if self._predictor is not None:
			return

		if self._ocr == "FAKE":
			return

		batch_size = self._options["batch_size"]
		if batch_size > 0:
			batch_size_kwargs = dict(batch_size=batch_size)
		else:
			batch_size_kwargs = dict()
		self._chunk_size = batch_size

		if len(self._models) == 1:
			self._predictor = Predictor(
				str(self._models[0]), **batch_size_kwargs)
			self._predict_kwargs = batch_size_kwargs
			self._voter = None
			self._line_height = int(self._predictor.model_params.line_height)
		else:
			logging.info("using Calamari voting with %d models." % len(self._models))
			self._predictor = MultiPredictor(
				checkpoints=[str(p) for p in self._models],
				**batch_size_kwargs)
			self._predict_kwargs = dict()
			self._voter = ConfidenceVoter()
			self._line_height = int(self._predictor.predictors[0].model_params.line_height)

	def artifacts(self):
		return [
			("reliable", Input(
				Artifact.LINES, Artifact.TABLES,
				stage=Stage.RELIABLE)),
			("output", Output(Artifact.OCR)),
		]

	def process(self, page_path: Path, reliable, output):
		self._load_models()

		lines = reliable.lines.by_path

		extractor = LineExtractor(
			reliable.tables,
			self._line_height,
			self._options,
			min_confidence=reliable.lines.min_confidence)

		min_width = 6
		min_height = 6

		names = []
		empty_names = []
		images = []
		for stem, im in extractor(lines, ignored=self._ignored):
			if im.width >= min_width and im.height >= min_height:
				names.append("/".join(stem))
				images.append(np.array(im))
			else:
				empty_names.append("/".join(stem))

		if self._ocr == "DRY":
			logging.info("will ocr the following lines:\n%s" % "\n".join(sorted(names)))
			return

		chunk_size = self._chunk_size
		if chunk_size <= 0:
			chunk_size = len(images)

		texts = []

		if self._ocr == "FAKE":
			for name in names:
				texts.append("text for %s." % name)
		else:
			for i in range(0, len(images), chunk_size):
				for prediction in self._predictor.predict_raw(
					images[i:i + chunk_size], progress_bar=False, **self._predict_kwargs):

					if self._voter is not None:
						prediction = self._voter.vote_prediction_result(prediction)
					texts.append(prediction.sentence)

		with output.ocr() as zf:
			for name, text in zip(names, texts):
				zf.writestr("%s.txt" % name, text)
			for name in empty_names:
				zf.writestr("%s.txt" % name, "")

@click.command()
@Processor.options
@LineExtractor.options
@OCRProcessor.options
@click.option(
	'-m', '--model',
	required=False,
	type=str,
	help='path that contains Calamari model(s)')
def make_processor(**kwargs):
	return OCRProcessor(kwargs)

@click.command()
@click.argument(
	'data_path',
	type=click.Path(exists=True),
	required=True)
@click.option(
	'-m', '--model',
	required=False,
	type=click.Path(exists=True),
	help='path that contains Calamari model(s)')
@Processor.options
@LineExtractor.options
@OCRProcessor.options
def run_ocr(data_path, **kwargs):
	""" Perform OCR on all recognized lines in DATA_PATH. """


    stack = inspect.stack()
    try:
        for frame_info in stack[skip_frames + 1:]:
            frame = None

            try:
                frame = frame_info.frame

                if name in frame.f_locals or (global_ok and name in frame.f_globals):
                    return Scope(globals_=frame.f_globals, locals_=frame.f_locals)
            finally:
                del frame
    finally:
        del stack

    return None

"""
jishaku.repl.scope
~~~~~~~~~~~~~~~~~~

The Scope class and functions relating to it.

"""

import inspect
import typing

class Scope:
    """
    Class that represents a global and local scope for both scope inspection and creation.

    Many REPL functions expect or return this class.

    .. code:: python3

        scope = Scope()  # an empty Scope

        scope = Scope(globals(), locals())  # a Scope imitating the current, real scope.

        scope = Scope({'a': 3})  # a Scope with a pre-existing global scope key, and an empty local scope.
    """

    __slots__ = ('globals', 'locals')

    def __init__(self, globals_: dict = None, locals_: dict = None):
        self.globals: dict = globals_ or {}
        self.locals: dict = locals_ or {}

    def clear_intersection(self, other_dict):
        """
        Clears out locals and globals from this scope where the key-value pair matches
        with other_dict.

        This allows cleanup of temporary variables that may have washed up into this
        Scope.

        Parameters
        -----------
        other_dict: :class:`dict`
            The dictionary to be used to determine scope clearance.

            If a key from this dict matches an entry in the globals or locals of this scope,
            and the value is identical, it is removed from the scope.

        Returns
        -------
        Scope
            The updated scope (self).
        """

        for key, value in other_dict.items():
            if key in self.globals and self.globals[key] is value:
                del self.globals[key]
            if key in self.locals and self.locals[key] is value:
                del self.locals[key]

        return self

    def update(self, other):
        """
        Updates this scope with the content of another scope.

        Parameters
        ---------
        other: :class:`Scope`
            The scope to overlay onto this one.

        Returns
        -------
        Scope
            The updated scope (self).
        """

        self.globals.update(other.globals)
        self.locals.update(other.locals)
        return self

    def update_globals(self, other: dict):
        """
        Updates this scope's globals with a dict.

        Parameters
        -----------
        other: :class:`dict`
            The dictionary to be merged into this scope.

        Returns
        -------
        Scope
            The updated scope (self).
        """

        self.globals.update(other)
        return self

    def update_locals(self, other: dict):
        """
        Updates this scope's locals with a dict.

        Parameters
        -----------
        other: :class:`dict`
            The dictionary to be merged into this scope.

        Returns
        -------
        Scope
            The updated scope (self).
        """

        self.locals.update(other)
        return self

def get_parent_scope_from_var(name, global_ok=False, skip_frames=0) -> typing.Optional[Scope]:
    """
    Iterates up the frame stack looking for a frame-scope containing the given variable name.

    Returns
    --------
    Optional[Scope]
        The relevant :class:`Scope` or None
    """

def get_parent_var(name, global_ok=False, default=None, skip_frames=0):
    """
    Directly gets a variable from a parent frame-scope.

    Returns
    --------
    Any
        The content of the variable found by the given name, or None.
    """

    scope = get_parent_scope_from_var(name, global_ok=global_ok, skip_frames=skip_frames + 1)

    if not scope:
        return default

    if name in scope.locals:
        return scope.locals.get(name, default)

    return scope.globals.get(name, default)


    if itr < start_itrs:
        return mul_lr
    else:
        if itr < start_itrs + warmup_itrs:
            return (itr - start_itrs) / warmup_itrs * mul_lr
        else:
            return mul_lr

import shutil
import traceback
from functools import partial
import pprint
import logging
import argparse
import os
import numpy as np
import math
from tqdm import tqdm
from datetime import datetime
import copy
from collections import deque
from einops import repeat, rearrange

import torch
import torch.distributed as dist
import torch.multiprocessing as mp
import torch.nn as nn
import torch.nn.functional as F
from torch.nn.parallel import DistributedDataParallel as DDP
from torchvision.utils import save_image, make_grid

from torch_ema import ExponentialMovingAverage

from tl2.launch.launch_utils import update_parser_defaults_from_yaml, global_cfg
from tl2.proj.fvcore import set_global_cfg
from tl2.modelarts import modelarts_utils
from tl2.proj.fvcore import build_model
from tl2.proj.stylegan2_ada import ada_utils
from tl2.proj.logger.textlogger import summary_dict2txtfig, global_textlogger
from tl2 import tl2_utils
from tl2.proj.pytorch import torch_utils

from generators import generators
from discriminators import discriminators
from siren import siren
from exp.dev.nerf_inr import fid_evaluation
from exp.dev.nerf_inr import curriculums
from exp.pigan import datasets
from exp.comm import comm_model_utils
from exp.comm import diff_aug

def setup(rank, world_size, port):
    os.environ['MASTER_ADDR'] = 'localhost'
    os.environ['MASTER_PORT'] = port

    # initialize the process group
    dist.init_process_group("nccl", rank=rank, world_size=world_size)

def cleanup():
    dist.destroy_process_group()

def load_images(images, curriculum, device):
    return_images = []
    head = 0
    for stage in curriculum['stages']:
        stage_images = images[head:head + stage['batch_size']]
        stage_images = F.interpolate(stage_images, size=stage['img_size'],  mode='bilinear', align_corners=True)
        return_images.append(stage_images)
        head += stage['batch_size']
    return return_images

def z_sampler(shape, device, dist):
    if dist == 'gaussian':
        z = torch.randn(shape, device=device)
    elif dist == 'uniform':
        z = torch.rand(shape, device=device) * 2 - 1
    return z

def saved_models(G_ema,
                 generator_ddp,
                 discriminator_ddp,
                 optimizer_G,
                 optimizer_D,
                 scaler_G,
                 scaler_D,
                 state_dict,
                 step,
                 epoch,
                 saved_dir=None,
                 metadata=None,
                 ):
    if saved_dir is None:
        ckpt_max2keep = tl2_utils.MaxToKeep.get_named_max_to_keep(
            name='ckpt', max_to_keep=4, use_circle_number=True)
        saved_dir = ckpt_max2keep.step_and_ret_circle_dir(
            global_cfg.tl_ckptdir, info_msg=step)
    os.makedirs(saved_dir, exist_ok=True)

    if metadata is not None:
        tl2_utils.json_dump(metadata, f"{saved_dir}/metadata.json")
    global_cfg.dump_to_file_with_command(
        f"{saved_dir}/config_command.yaml", global_cfg.tl_command)

    state_dict['epoch'] = epoch
    state_dict['step'] = step
    model_dict = {
        'generator': generator_ddp.module,
        'G_ema': G_ema,
        'discriminator': discriminator_ddp.module,
        'optimizer_G': optimizer_G,
        'optimizer_D': optimizer_D,
        'scaler_G': scaler_G,
        'scaler_D': scaler_D,
        'state_dict': state_dict,
    }

    torch_utils.save_models(save_dir=saved_dir,
                            model_dict=model_dict,
                            info_msg=f"epoch: {epoch}\n"
                                     f"step: {step}")

    torch.cuda.empty_cache()

    pass

def save_images_horizontal_flip(saved_imgdir,
                                metadata,
                                G_ema,
                                fixed_z,
                                device,
                                step):
    sub_fixed_z = {}
    for name, z_ in fixed_z.items():
        sub_fixed_z[name] = z_[:20]
    fixed_z = sub_fixed_z

    with torch.no_grad():
        with torch.cuda.amp.autocast(global_cfg.use_amp_G):
            copied_metadata = copy.deepcopy(metadata)
            copied_metadata['num_steps'] = max(12, copied_metadata['num_steps'])
            copied_metadata['h_stddev'] = copied_metadata['v_stddev'] = 0
            copied_metadata['h_mean'] = 1.44
            if copied_metadata['img_size'] > 128:
                sub_fixed_z = {}
                for name, z_ in fixed_z.items():
                    sub_fixed_z[name] = z_[:10]
                fixed_z = sub_fixed_z
            gen_imgs1 = G_ema(fixed_z,
                              return_aux_img=True,
                              forward_points=128 ** 2,
                              **copied_metadata)[0]

            copied_metadata = copy.deepcopy(metadata)
            copied_metadata['num_steps'] = max(12, copied_metadata['num_steps'])
            copied_metadata['h_stddev'] = copied_metadata['v_stddev'] = 0
            copied_metadata['h_mean'] = 1.70
            gen_imgs2 = G_ema(fixed_z,
                              return_aux_img=True,
                              forward_points=128 ** 2,
                              **copied_metadata)[0]

    gen_imgs = torch.cat([gen_imgs1, gen_imgs2])
    save_image(gen_imgs,
               os.path.join(saved_imgdir, f"gen_flip_ema.jpg"),
               nrow=len(list(fixed_z.values())[0]) // 2,
               normalize=True)
    pass

def save_images(generator_ddp,
                G_ema,
                discriminator,
                metadata,
                fixed_z,
                device,
                ):
    generator_ddp.eval()
    fixed_z_ori = fixed_z
    bs = len(list(fixed_z.values())[0])

    img_max2keep = tl2_utils.MaxToKeep.get_named_max_to_keep('img')
    saved_imgdir = img_max2keep.step_and_ret_circle_dir(root_dir=global_cfg.tl_imgdir,
                                                        info_msg=discriminator.step)
    with torch.no_grad():
        with torch.cuda.amp.autocast(global_cfg.use_amp_G):
            copied_metadata = copy.deepcopy(metadata)
            copied_metadata['num_steps'] = max(12, copied_metadata['num_steps'])
            copied_metadata['h_stddev'] = copied_metadata['v_stddev'] = 0
            if copied_metadata['img_size'] > 128:
                sub_fixed_z = {}
                for name, z_ in fixed_z_ori.items():
                    sub_fixed_z[name] = z_[0:9]
                fixed_z = sub_fixed_z
                bs = 9
            gen_imgs = generator_ddp.module(fixed_z,
                                            return_aux_img=True,
                                            forward_points=128**2,
                                            **copied_metadata)[0]
    save_image(gen_imgs,
               os.path.join(saved_imgdir, f"gen_fixed.jpg"),
               nrow=int(math.sqrt(bs)),
               normalize=True)

    with torch.no_grad():
        with torch.cuda.amp.autocast(global_cfg.use_amp_G):
            copied_metadata = copy.deepcopy(metadata)
            copied_metadata['num_steps'] = max(12, copied_metadata['num_steps'])
            copied_metadata['h_stddev'] = copied_metadata['v_stddev'] = 0
            copied_metadata['h_mean'] += 0.5
            gen_imgs = generator_ddp.module(fixed_z,
                                            return_aux_img=True,
                                            forward_points=128 ** 2,
                                            **copied_metadata)[0]
    save_image(gen_imgs,
               os.path.join(saved_imgdir, f"gen_tilted.jpg"),
               nrow=int(math.sqrt(bs)),
               normalize=True)

    G_ema.eval()

    with torch.no_grad():
        with torch.cuda.amp.autocast(global_cfg.use_amp_G):
            copied_metadata = copy.deepcopy(metadata)
            copied_metadata['num_steps'] = max(12, copied_metadata['num_steps'])
            copied_metadata['h_stddev'] = copied_metadata['v_stddev'] = 0
            gen_imgs = G_ema(fixed_z,
                             return_aux_img=True,
                             forward_points=128 ** 2,
                             **copied_metadata)[0]
    save_image(gen_imgs,
               os.path.join(saved_imgdir, f"gen_fixed_ema.jpg"),
               nrow=int(math.sqrt(bs)),
               normalize=True)

    with torch.no_grad():
        with torch.cuda.amp.autocast(global_cfg.use_amp_G):
            copied_metadata = copy.deepcopy(metadata)
            copied_metadata['num_steps'] = max(12, copied_metadata['num_steps'])
            copied_metadata['h_stddev'] = copied_metadata['v_stddev'] = 0
            copied_metadata['h_mean'] += 0.5
            gen_imgs = G_ema(fixed_z,
                             return_aux_img=True,
                             forward_points=128 ** 2,
                             **copied_metadata)[0]
    save_image(gen_imgs,
               os.path.join(saved_imgdir, f"gen_tilted_ema.jpg"),
               nrow=int(math.sqrt(bs)),
               normalize=True)

    save_images_horizontal_flip(saved_imgdir=saved_imgdir,
                                metadata=metadata,
                                G_ema=G_ema,
                                fixed_z=fixed_z_ori,
                                device=device,
                                step=discriminator.step)

    with torch.no_grad():
        with torch.cuda.amp.autocast(global_cfg.use_amp_G):
            copied_metadata = copy.deepcopy(metadata)
            copied_metadata['num_steps'] = max(12, copied_metadata['num_steps'])
            copied_metadata['h_stddev'] = copied_metadata['v_stddev'] = 0
            copied_metadata['psi'] = 0.7
            gen_imgs = G_ema(fixed_z,
                             return_aux_img=True,
                             forward_points=128 ** 2,
                             **copied_metadata)[0]
    save_image(gen_imgs,
               os.path.join(saved_imgdir, f"gen_random.jpg"),
               nrow=int(math.sqrt(bs)),
               normalize=True)

    pass

def get_curriculum(curriculum_name):
    curriculum = getattr(curriculums, curriculum_name)

    # update curriculum_step
    for curriculum_step in curriculum.keys():
        if type(curriculum_step) == int:
            # stage
            curriculum_step_str = str(curriculum_step)
            if curriculum_step_str in global_cfg.curriculum:
                curriculum[curriculum_step].update(global_cfg.curriculum.get(curriculum_step_str).to_dict())
        else:
            # update attrs
            if curriculum_step in global_cfg.curriculum:
                curriculum[curriculum_step] = global_cfg.curriculum[curriculum_step]
    for new_attr, value in global_cfg.curriculum.get('new_attrs', {}).items():
        assert new_attr not in curriculum
        curriculum[new_attr] = value
    return curriculum

def output_images(generator,
                  input_metadata,
                  rank,
                  world_size,
                  fake_dir,
                  num_imgs=2048,
                  img_size=128,):
    os.makedirs(fake_dir, exist_ok=True)
    metadata = copy.deepcopy(input_metadata)
    metadata['num_steps'] = max(12, metadata['num_steps'])
    metadata['img_size'] = img_size
    metadata['batch_size'] = 1

    metadata['h_stddev'] = metadata.get('h_stddev_eval', metadata['h_stddev'])
    metadata['v_stddev'] = metadata.get('v_stddev_eval', metadata['v_stddev'])
    metadata['sample_dist'] = metadata.get('sample_dist_eval', metadata['sample_dist'])
    metadata['psi'] = 1

    generator.eval()

    img_counter = rank
    if rank == 0: pbar = tqdm(desc=f"generating images at {img_size}x{img_size}", total = num_imgs)
    with torch.no_grad():
        while img_counter < num_imgs:
            zs = generator.get_zs(metadata['batch_size'])
            generated_imgs = generator(zs,
                                       forward_points=256 ** 2,
                                       **metadata)[0]

            for img in generated_imgs:
                save_image(img,
                           os.path.join(fake_dir, f'{img_counter:0>5}.jpg'),
                           normalize=True,
                           value_range=(-1, 1))
                img_counter += world_size
                if rank == 0: pbar.update(world_size)
    if rank == 0: pbar.close()
    pass

def lambda_foo(itr,
               start_itrs,
               warmup_itrs,
               mul_lr=1):
    """

    :param itr: itr >= start_itrs
    :param start_itrs:
    :param warmup_itrs:
    :param mul_lr:
    :return:
    """

def train(rank,
          world_size,
          opt,
          cfg):
    setup(rank, world_size, opt.port)
    torch.cuda.set_device(rank)
    device = torch.device(rank)

    update_parser_defaults_from_yaml(parser=None, is_main_process=(rank == 0))
    if rank == 0:
        modelarts_utils.setup_tl_outdir_obs(global_cfg, unzip_code=False)
        modelarts_utils.modelarts_sync_results_dir(global_cfg, join=True)
        modelarts_utils.prepare_dataset(global_cfg.get('modelarts_download', {}), global_cfg=global_cfg)

    logger = logging.getLogger('tl')
    torch_utils.init_seeds(seed=global_cfg.seed, rank=rank)

    curriculum = get_curriculum(curriculum_name=opt.curriculum)

    metadata = curriculums.extract_metadata(curriculum, 0)

    SIREN = getattr(siren, metadata['model'])
    CHANNELS = 3

    scaler_G = torch.cuda.amp.GradScaler(enabled=global_cfg.use_amp_G)
    scaler_D = torch.cuda.amp.GradScaler(enabled=global_cfg.use_amp_D)
    state_dict = {
        'best_fid': np.inf,
    }

    if 'G_cfg' in global_cfg:
        generator = build_model(cfg=global_cfg.G_cfg).to(device)
    else:
        generator = getattr(generators, metadata['generator'])(SIREN, metadata['latent_dim']).to(device)
    if 'D_cfg' in global_cfg:
        discriminator = build_model(cfg=global_cfg.D_cfg,
                                    kwargs_priority=True).to(device)
    else:
        discriminator = getattr(discriminators, metadata['discriminator'])().to(device)

    G_ema = copy.deepcopy(generator)
    ema_model = comm_model_utils.EMA(source=generator, target=G_ema, decay=0.999, start_itr=1000)

    if global_cfg.tl_resume:
        resume_dir = f"{global_cfg.tl_resumedir}/ckptdir/resume"
        model_dict = {
            'generator': generator,
            'discriminator': discriminator,
            'state_dict': state_dict,
        }
        if global_cfg.load_G_ema:
            model_dict['G_ema'] = G_ema
        torch_utils.load_models(save_dir=resume_dir, model_dict=model_dict, strict=False, rank=rank)
        if global_cfg.load_G_ema:
            ema_model.update_target_dict(G_ema.state_dict())
        else:
            ema_model.update_target_dict(generator.state_dict())

        logger.info(pprint.pformat(state_dict))
        discriminator.step = state_dict['step']
        generator.step = state_dict['step']
        discriminator.epoch = state_dict['epoch']
        generator.epoch = state_dict['epoch']

    generator_ddp = DDP(generator, device_ids=[rank], find_unused_parameters=True, broadcast_buffers=False)
    discriminator_ddp = DDP(discriminator, device_ids=[rank], find_unused_parameters=True, broadcast_buffers=False)
    generator = generator_ddp.module
    discriminator = discriminator_ddp.module

    if metadata.get('unique_lr', False):
        mapping_network_param_names = [name for name, _ in generator_ddp.module.siren.mapping_network.named_parameters()]
        mapping_network_parameters = [p for n, p in generator_ddp.named_parameters() if n in mapping_network_param_names]
        generator_parameters = [p for n, p in generator_ddp.named_parameters() if n not in mapping_network_param_names]
        optimizer_G = torch.optim.Adam(
            [
                {'params': generator_parameters,
                 'name': 'generator'},
                {'params': mapping_network_parameters,
                 'name': 'mapping_network',
                 'lr': metadata['gen_lr']*5e-2}
            ],
            lr=metadata['gen_lr'],
            betas=metadata['betas'],
            weight_decay=metadata['weight_decay'])
    else:
        if hasattr(generator, 'get_optimizer'):
            optimizer_G = generator.get_optimizer(**global_cfg.G_cfg.optim)
        else:
            optimizer_G = torch.optim.Adam(
                [{'params': generator_ddp.parameters(),
                  'initial_lr': metadata['gen_lr']}],
                lr=metadata['gen_lr'],
                betas=metadata['betas'],
                weight_decay=metadata['weight_decay'])

    optimizer_D = torch.optim.Adam(
        [{'params': discriminator_ddp.parameters(),
          'initial_lr': metadata['disc_lr']}],
        lr=metadata['disc_lr'],
        betas=metadata['betas'],
        weight_decay=metadata['weight_decay'])

    if global_cfg.tl_resume and global_cfg.load_optimizers:
        model_dict = {
            'optimizer_G': optimizer_G,
            'optimizer_D': optimizer_D,
            'scaler_G': scaler_G,
            'scaler_D': scaler_D,
        }
        torch_utils.load_models(save_dir=resume_dir, model_dict=model_dict, strict=False, rank=rank)

    if not hasattr(generator, 'get_optimizer'):
        optimizer_G.param_groups[0]['initial_lr'] = metadata['gen_lr']
        optimizer_G.param_groups[0]['betas'] = metadata['betas']
    optimizer_D.param_groups[0]['initial_lr'] = metadata['disc_lr']
    optimizer_D.param_groups[0]['betas'] = metadata['betas']
    optimizer_D.param_groups[0]['weight_decay'] = global_cfg.D_wd

    start_itr = discriminator.step
    lambda_lr_G = partial(lambda_foo,
                          start_itrs=discriminator.step,
                          warmup_itrs=global_cfg.warmup_itrs,
                          mul_lr=global_cfg.mul_lr)
    scheduler_G = torch.optim.lr_scheduler.LambdaLR(
        optimizer_G, lr_lambda=lambda_lr_G, last_epoch=discriminator.step)
    lambda_lr_D = partial(lambda_foo,
                          start_itrs=discriminator.step,
                          warmup_itrs=global_cfg.warmup_itrs,
                          mul_lr=global_cfg.mul_lr)
    scheduler_D = torch.optim.lr_scheduler.LambdaLR(
        optimizer_D, lr_lambda=lambda_lr_D, last_epoch=discriminator.step)

    if opt.set_step != None:
        generator.step = opt.set_step
        discriminator.step = opt.set_step

    generator.set_device(device)

    fixed_z = generator.get_zs(global_cfg.get('fixed_z_bs', 25))
    use_diffaug = global_cfg.get('use_diffaug', False)
    #  Training

    dummy_tensor = torch.tensor([0], device=device)

    with open(os.path.join(opt.output_dir, 'options.txt'), 'w') as f:
        f.write(str(opt))
        f.write('\n\n')
        f.write(str(generator))
        f.write('\n\n')
        f.write(str(discriminator))
        f.write('\n\n')
        f.write(str(curriculum))

    torch.manual_seed(rank)
    dataloader = None
    total_progress_bar = tqdm(total = opt.n_epochs, desc = "Epoch", dynamic_ncols=True)
    total_progress_bar.update(discriminator.epoch)
    interior_step_bar = tqdm(dynamic_ncols=True)
    for _ in range (discriminator.epoch, opt.n_epochs):
        total_progress_bar.update(1)

        metadata = curriculums.extract_metadata(curriculum, discriminator.step)

        # Set learning rates

        if not dataloader or dataloader.batch_size != metadata['batch_size']:
            dataloader, CHANNELS = datasets.get_dataset_distributed(
                name=metadata['dataset'], world_size=world_size, rank=rank, num_workers=global_cfg.num_workers,
                **metadata)

            step_next_upsample = curriculums.next_upsample_step(curriculum, discriminator.step)
            step_last_upsample = curriculums.last_upsample_step(curriculum, discriminator.step)

            interior_step_bar.reset(total=(step_next_upsample - step_last_upsample))
            interior_step_bar.set_description(f"Stage {step_last_upsample}-{step_next_upsample}")
            interior_step_bar.update((discriminator.step - step_last_upsample))

        for i, (imgs, _) in enumerate(dataloader):
            metadata = curriculums.extract_metadata(curriculum, discriminator.step)

            scheduler_G.step()
            scheduler_D.step()

            if dataloader.batch_size != metadata['batch_size']: break

            generator_ddp.train()
            discriminator_ddp.train()

            max_alpha = 1.
            if global_cfg.D_first_layer_warmup:
                alpha = min(max_alpha, (discriminator.step - start_itr) / (metadata['fade_steps']))
            else:
                alpha = max_alpha

            real_imgs = imgs.to(device, non_blocking=True)

            metadata['nerf_noise'] = max(0, 1. - discriminator.step/5000.)

            # TRAIN DISCRIMINATOR
            torch_utils.requires_grad(generator_ddp, False)
            torch_utils.requires_grad(discriminator_ddp, True)

            aux_reg = global_cfg.train_aux_img and (i % global_cfg.update_aux_every == 0)

            with torch.cuda.amp.autocast(global_cfg.use_amp_D):
                # Generate images for discriminator training
                with torch.no_grad():
                    zs = generator.get_zs(real_imgs.shape[0])
                    split_batch_size = real_imgs.shape[0] // metadata['batch_split']
                    gen_imgs = []
                    gen_positions = []
                    if metadata['img_size'] >= 256 and global_cfg.forward_points is not None:
                        forward_points = global_cfg.forward_points ** 2
                    else:
                        forward_points = None
                    for split in range(metadata['batch_split']):
                        subset_z = {}
                        for name, z_ in zs.items():
                            subset_z[name] = z_[split * split_batch_size:(split+1) * split_batch_size]
                        g_imgs, g_pos = generator_ddp(subset_z,
                                                      return_aux_img=aux_reg,
                                                      forward_points=forward_points,
                                                      grad_points=None,
                                                      **metadata)

                        gen_imgs.append(g_imgs)
                        gen_positions.append(g_pos)

                    gen_imgs = torch.cat(gen_imgs, axis=0)
                    if use_diffaug:
                        gen_imgs = diff_aug.DiffAugment(gen_imgs)
                    gen_positions = torch.cat(gen_positions, axis=0)

                if aux_reg:
                    real_imgs = torch.cat([real_imgs, real_imgs], dim=0)
                if use_diffaug:
                    real_imgs = diff_aug.DiffAugment(real_imgs)
                real_imgs.requires_grad_()
                r_preds, _, _ = discriminator_ddp(real_imgs, alpha, use_aux_disc=aux_reg, **metadata)

            d_regularize = i % global_cfg.d_reg_every == 0

            if metadata['r1_lambda'] > 0 and d_regularize:
                # Gradient penalty
                grad_real = torch.autograd.grad(
                    outputs=scaler_D.scale(r_preds.sum()),
                    inputs=real_imgs,
                    create_graph=True)
                inv_scale = 1./scaler_D.get_scale()
                grad_real = [p * inv_scale for p in grad_real][0]
            with torch.cuda.amp.autocast(global_cfg.use_amp_D):
                if metadata['r1_lambda'] > 0 and d_regularize:
                    grad_penalty = (grad_real.view(grad_real.size(0), -1).norm(2, dim=1) ** 2).mean()
                    grad_penalty = 0.5 * metadata['r1_lambda'] * global_cfg.d_reg_every * grad_penalty + 0 * r_preds[0]
                else:
                    grad_penalty = dummy_tensor

                g_preds, g_pred_latent, g_pred_position = discriminator_ddp(
                    gen_imgs, alpha, use_aux_disc=aux_reg, **metadata)
                if metadata['z_lambda'] > 0 or metadata['pos_lambda'] > 0:
                    if metadata['z_lambda'] > 0:
                        latent_penalty = torch.nn.functional.mse_loss(
                            g_pred_latent, zs) * metadata['z_lambda']
                    else:
                        latent_penalty = 0
                    position_penalty = torch.nn.functional.mse_loss(
                        g_pred_position, gen_positions) * metadata['pos_lambda']
                    identity_penalty = latent_penalty + position_penalty
                else:
                    identity_penalty=0

                d_loss = torch.nn.functional.softplus(g_preds).mean() + \
                         torch.nn.functional.softplus(-r_preds).mean() + \
                         grad_penalty + \
                         identity_penalty

            optimizer_D.zero_grad()
            scaler_D.scale(d_loss).backward()
            scaler_D.unscale_(optimizer_D)
            try:
                D_total_norm = torch.nn.utils.clip_grad_norm_(discriminator_ddp.parameters(),
                                                              metadata['grad_clip'],
                                                              )
            except:
                logger.info(traceback.format_exc())
                saved_models(G_ema=G_ema,
                             generator_ddp=generator_ddp,
                             discriminator_ddp=discriminator_ddp,
                             optimizer_G=optimizer_G,
                             optimizer_D=optimizer_D,
                             scaler_G=scaler_G,
                             scaler_D=scaler_D,
                             state_dict=state_dict,
                             step=discriminator.step,
                             epoch=discriminator.epoch,
                             saved_dir=f"{global_cfg.tl_ckptdir}/D_crupted")
                optimizer_D.zero_grad()

            scaler_D.step(optimizer_D)
            scaler_D.update()

            # TRAIN GENERATOR
            if metadata['gen_lr'] > 0.:
                torch_utils.requires_grad(generator_ddp, True)
                torch_utils.requires_grad(discriminator_ddp, False)

                zs = generator.get_zs(imgs.shape[0])

                split_batch_size = imgs.shape[0] // metadata['batch_split']

                if global_cfg.grad_points is not None:
                    grad_points = global_cfg.grad_points ** 2
                else:
                    grad_points = None
                for split in range(metadata['batch_split']):
                    with torch.cuda.amp.autocast(global_cfg.use_amp_G):
                        subset_z = {}
                        for name, z_ in zs.items():
                            subset_z[name] = z_[split * split_batch_size:(split + 1) * split_batch_size]
                        gen_imgs, gen_positions = generator_ddp(subset_z,
                                                                return_aux_img=aux_reg,
                                                                grad_points=grad_points,
                                                                forward_points=None,
                                                                **metadata)
                        if use_diffaug:
                            gen_imgs = diff_aug.DiffAugment(gen_imgs)
                        with torch.cuda.amp.autocast(global_cfg.use_amp_D):
                          g_preds, g_pred_latent, g_pred_position = discriminator_ddp(
                              gen_imgs.to(torch.float32), alpha, use_aux_disc=aux_reg, **metadata)

                        if g_preds.shape[0] == imgs.shape[0] * 2:
                            aux_mask = torch.tensor(global_cfg.main_aux_mask, device=g_preds.device)
                            aux_mask = repeat(aux_mask, "b -> b r", r=imgs.shape[0])
                            aux_mask = rearrange(aux_mask, "b r -> (b r) 1")
                            g_preds = g_preds * aux_mask

                        if metadata['topk_v'] > 0:
                            topk_percentage = \
                                max(
                                    0.99 ** (discriminator.step/metadata['topk_interval']),
                                    metadata['topk_v']
                                ) if 'topk_interval' in metadata and 'topk_v' in metadata \
                                    else 1
                            topk_num = math.ceil(topk_percentage * g_preds.shape[0])
                            g_preds = torch.topk(g_preds, topk_num, dim=0).values
                        else:
                            topk_num = g_preds.shape[0]

                        if metadata['z_lambda'] > 0 or metadata['pos_lambda'] > 0:
                            if metadata['z_lambda'] > 0:
                                latent_penalty = torch.nn.functional.mse_loss(
                                    g_pred_latent, subset_z) * metadata['z_lambda']
                            else:
                                latent_penalty = 0
                            position_penalty = torch.nn.functional.mse_loss(
                                g_pred_position, gen_positions) * metadata['pos_lambda']
                            identity_penalty = latent_penalty + position_penalty
                        else:
                            identity_penalty = 0

                        g_loss = torch.nn.functional.softplus(-g_preds).mean() + identity_penalty

                    scaler_G.scale(g_loss).backward()

                scaler_G.unscale_(optimizer_G)
                try:
                    G_total_norm = torch.nn.utils.clip_grad_norm_(generator_ddp.parameters(),
                                                                  metadata.get('grad_clip', 0.3),
                                                                  )
                except:
                    logger.info(traceback.format_exc())
                    saved_models(G_ema=G_ema,
                                 generator_ddp=generator_ddp,
                                 discriminator_ddp=discriminator_ddp,
                                 optimizer_G=optimizer_G,
                                 optimizer_D=optimizer_D,
                                 scaler_G=scaler_G,
                                 scaler_D=scaler_D,
                                 state_dict=state_dict,
                                 step=discriminator.step,
                                 epoch=discriminator.epoch,
                                 saved_dir=f"{global_cfg.tl_ckptdir}/G_crupted",
                                 metadata=metadata)
                    optimizer_G.zero_grad()
                scaler_G.step(optimizer_G)
                scaler_G.update()
                optimizer_G.zero_grad()
                ema_model.update(itr=discriminator.step, source_dict=generator_ddp.module.state_dict())
            else:
                grad_points = None
                G_total_norm = dummy_tensor
                g_loss = dummy_tensor
                topk_num = 0

            if rank == 0:
                interior_step_bar.update(1)
                if i % 10 == 0:
                    tqdm.write(f"[Experiment: {opt.output_dir}] "
                               f"[GPU: {os.environ['CUDA_VISIBLE_DEVICES']}] "
                               f"[Epoch: {discriminator.epoch}/{opt.n_epochs}] "
                               f"[D loss: {d_loss.item():.3f}] "
                               f"[G loss: {g_loss.item():.3f}] "
                               f"[grad_penalty: {grad_penalty.item():.3f}] "
                               f"[Step: {discriminator.step}] "
                               f"[G_lr: {optimizer_G.param_groups[0]['lr']:.8f}] "
                               f"[D_lr: {optimizer_D.param_groups[0]['lr']:.8f}] "
                               f"[Img Size: {metadata['img_size']}] "
                               f"[grad_points: {grad_points}]"
                               f"[Batch Size: {metadata['batch_size']}] "
                               f"[TopK: {topk_num}] "
                               f"[scaler_G: {scaler_G.get_scale()}] "
                               f"[scaler_D: {scaler_D.get_scale()}] "
                               f"[r1_lambda: {metadata['r1_lambda']}] "
                               f"[alpha({metadata['fade_steps']}): {alpha}] ")

                if global_cfg.tl_debug or \
                      (discriminator.step % global_cfg.get('log_every', 100) == 0) and discriminator.step > 1000:
                    summary_dict = {
                        'd_loss': d_loss.item(),
                        'g_loss': g_loss.item(),
                        'grad_penalty': grad_penalty.item(),
                        'G_lr': optimizer_G.param_groups[0]['lr'],
                        'D_lr': optimizer_D.param_groups[0]['lr'],
                        'img_size': metadata['img_size'],
                        'grad_points': grad_points if grad_points is not None else metadata['img_size'] ** 2,
                        'batch_size': metadata['batch_size'],
                        'topk_num': topk_num,
                        'scaler_G': scaler_G.get_scale(),
                        'scaler_D': scaler_D.get_scale(),
                        'r1_lambda': metadata['r1_lambda'],
                        'alpha': alpha,
                        'grad_clip': metadata['grad_clip'],
                        'D_total_norm': D_total_norm.item(),
                        'G_total_norm': G_total_norm.item(),
                    }
                    summary_dict2txtfig(summary_dict, prefix="train", step=discriminator.step,
                                        textlogger=global_textlogger)

                if discriminator.step % global_cfg.get('log_img_every', 500) == 0:
                    save_images(generator_ddp=generator_ddp,
                                G_ema=G_ema,
                                discriminator=discriminator,
                                metadata=metadata,
                                fixed_z=fixed_z,
                                device=device)

            if global_cfg.tl_debug or (discriminator.step) % global_cfg.get('eval_every', 5000) == 0:
                if rank == 0:
                    saved_models(G_ema=G_ema,
                                 generator_ddp=generator_ddp,
                                 discriminator_ddp=discriminator_ddp,
                                 optimizer_G=optimizer_G,
                                 optimizer_D=optimizer_D,
                                 scaler_G=scaler_G,
                                 scaler_D=scaler_D,
                                 state_dict=state_dict,
                                 step=discriminator.step,
                                 epoch=discriminator.epoch,
                                 metadata=metadata)
                    saved_models(G_ema=G_ema,
                                 generator_ddp=generator_ddp,
                                 discriminator_ddp=discriminator_ddp,
                                 optimizer_G=optimizer_G,
                                 optimizer_D=optimizer_D,
                                 scaler_G=scaler_G,
                                 scaler_D=scaler_D,
                                 state_dict=state_dict,
                                 step=discriminator.step,
                                 epoch=discriminator.epoch,
                                 saved_dir=f"{global_cfg.tl_ckptdir}/resume",
                                 metadata=metadata)
                    # output real images
                    if global_cfg.get('del_fid_real_images', False):
                        if os.path.exists(f"{opt.output_dir}/fid/real"):
                            shutil.rmtree(f"{opt.output_dir}/fid/real")
                    fid_evaluation.setup_evaluation(dataset_name=metadata['dataset'],
                                                    image_list_file=metadata['image_list_file'],
                                                    real_dir=f"{opt.output_dir}/fid/real",
                                                    target_size=min(metadata['img_size'], 128),
                                                    num_imgs=global_cfg.get('num_images_real_eval', 8000),
                                                    horizontal_flip=False)
                dist.barrier()
                generator_ddp.eval()
                # output fake images
                output_images(generator=G_ema,
                              input_metadata=metadata,
                              rank=rank,
                              world_size=world_size,
                              fake_dir=f"{opt.output_dir}/fid/fake",
                              num_imgs=global_cfg.get('num_images_gen_eval', 2048),
                              img_size=min(metadata['img_size'], 128))
                dist.barrier()
                if rank == 0:
                    fid = fid_evaluation.calculate_fid(real_dir=f"{opt.output_dir}/fid/real",
                                                       fake_dir=f"{opt.output_dir}/fid/fake",
                                                       batch_size=128)
                    logger.info(f"\nstep: {discriminator.step}, fid: {fid}\n")
                    summary_dict = {
                        'fid': fid
                    }
                    summary_dict2txtfig(summary_dict, prefix='eval', step=discriminator.step,
                                        textlogger=global_textlogger)

                    best_dir = f"{global_cfg.tl_ckptdir}/best_fid"
                    if not os.path.exists(best_dir):
                        state_dict = {'best_fid': np.inf}
                    if state_dict['best_fid'] > fid:
                        state_dict['best_fid'] = fid
                        saved_models(G_ema=G_ema,
                                     generator_ddp=generator_ddp,
                                     discriminator_ddp=discriminator_ddp,
                                     optimizer_G=optimizer_G,
                                     optimizer_D=optimizer_D,
                                     scaler_G=scaler_G,
                                     scaler_D=scaler_D,
                                     state_dict=state_dict,
                                     step=discriminator.step,
                                     epoch=discriminator.epoch,
                                     saved_dir=best_dir,
                                     metadata=metadata)

                    modelarts_utils.modelarts_sync_results_dir(cfg=global_cfg, join=False)
                dist.barrier()
                torch.cuda.empty_cache()

            discriminator.step += 1
            generator.step += 1
        discriminator.epoch += 1
        generator.epoch += 1

    cleanup()

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument("--n_epochs", type=int, default=60000, help="number of epochs of training")
    parser.add_argument("--sample_interval", type=int, default=200, help="interval between image sampling")
    parser.add_argument('--output_dir', type=str, default='debug')
    parser.add_argument('--load_dir', type=str, default='')
    parser.add_argument('--curriculum', type=str, required=True)
    parser.add_argument('--eval_freq', type=int, default=5000)
    parser.add_argument('--port', type=str, default='12355')
    parser.add_argument('--set_step', type=int, default=None)
    parser.add_argument('--model_save_interval', type=int, default=5000)

    update_parser_defaults_from_yaml(parser)

    modelarts_utils.setup_tl_outdir_obs(global_cfg, unzip_code=False)

    opt = parser.parse_args()
    print(opt)
    os.makedirs(opt.output_dir, exist_ok=True)
    num_gpus = len(os.environ['CUDA_VISIBLE_DEVICES'].split(','))
    mp.spawn(train, args=(num_gpus, opt, global_cfg), nprocs=num_gpus, join=True)

    modelarts_utils.prepare_dataset(global_cfg.get('modelarts_upload', {}), global_cfg=global_cfg, download=False)
    modelarts_utils.modelarts_sync_results_dir(global_cfg, join=True)