import logging
from pathlib import Path

import numpy as np
import pytorch_lightning as pl
import torch
from pytorch_lightning.callbacks import EarlyStopping, ModelCheckpoint
from pytorch_lightning.utilities import rank_zero_only
from utils_rag import save_json






def     __UpperCAmelCase					(			lowercase		):







  """simple docstring"""
  _UpperCAmelCase						 =      filter(lambda lowercase		: p.requires_grad   ,model.parameters()		)
  _UpperCAmelCase						 =      sum([np.prod(p.size()		) for p in model_parameters]		)
  return params


UpperCAmelCase__  		=						logging.getLogger(__name__)






def     __UpperCAmelCase					(			lowercase   ,lowercase		):







  """simple docstring"""






  if metric == "rouge2":
    _UpperCAmelCase						 =      """{val_avg_rouge2:.4f}-{step_count}"""
  elif metric == "bleu":
    _UpperCAmelCase						 =      """{val_avg_bleu:.4f}-{step_count}"""
  elif metric == "em":
    _UpperCAmelCase						 =      """{val_avg_em:.4f}-{step_count}"""
  elif metric == "loss":
    _UpperCAmelCase						 =      """{val_avg_loss:.4f}-{step_count}"""
  else:
    raise NotImplementedError(
        f'''seq2seq callbacks only support rouge2 and bleu, got {metric}, You can make your own by adding to this'''
        """ function."""		)

  _UpperCAmelCase						 =      ModelCheckpoint(
      dirpath=a__   ,filename=a__   ,monitor=f'''val_{metric}'''   ,mode="""max"""   ,save_top_k=1   ,every_n_epochs=1   ,)

  return checkpoint_callback






def     __UpperCAmelCase					(			lowercase   ,lowercase		):







  """simple docstring"""
  return EarlyStopping(
      monitor=f'''val_{metric}'''   ,mode="""min""" if """loss""" in metric else """max"""   ,patience=a__   ,verbose=a__   ,)


class        a				(		pl.Callback   ):


    def      lowerCAmelCase_  (					self	:  List[str]   ,				__lowerCAmelCase	:  Dict   ,				__lowerCAmelCase	:  Optional[int]					):
      _UpperCAmelCase						 =      {f'''lr_group_{i}''': param["""lr"""] for i, param in enumerate(pl_module.trainer.optimizers[0].param_groups					)}
      pl_module.logger.log_metrics(__A					)


    @rank_zero_only
    def      lowerCAmelCase_  (					self	:  Union[str, Any]   ,				__lowerCAmelCase	:  Optional[int]   ,				__lowerCAmelCase	:  List[str]   ,				__lowerCAmelCase	:  int   ,				__lowerCAmelCase	:  List[str]=True					):
      logger.info(f'''***** {type_path} results at step {trainer.global_step:05d} *****'''					)
      _UpperCAmelCase						 =      trainer.callback_metrics
      trainer.logger.log_metrics({k: v for k, v in metrics.items() if k not in ["""log""", """progress_bar""", """preds"""]}					)
      # Log results
      _UpperCAmelCase						 =      Path(pl_module.hparams.output_dir					)
      if type_path == "test":
        _UpperCAmelCase						 =      od / """test_results.txt"""
        _UpperCAmelCase						 =      od / """test_generations.txt"""
      else:
        # this never gets hit. I prefer not to save intermediate generations, and results are in metrics.json
        # If people want this it will be easy enough to add back.
        _UpperCAmelCase						 =      od / f'''{type_path}_results/{trainer.global_step:05d}.txt'''
        _UpperCAmelCase						 =      od / f'''{type_path}_generations/{trainer.global_step:05d}.txt'''
        results_file.parent.mkdir(exist_ok=__A					)
        generations_file.parent.mkdir(exist_ok=__A					)
      with open(__A   ,				"""a+"""					) as writer:
        for key in sorted(__A					):
          if key in ["log", "progress_bar", "preds"]:
            continue
          _UpperCAmelCase						 =      metrics[key]
          if isinstance(__A   ,				torch.Tensor					):
            _UpperCAmelCase						 =      val.item()
          _UpperCAmelCase						 =      f'''{key}: {val:.6f}\n'''
          writer.write(__A					)

      if not save_generations:
        return

      if "preds" in metrics:
        _UpperCAmelCase						 =      """\n""".join(metrics["""preds"""]					)
        generations_file.open("""w+"""					).write(__A					)


    @rank_zero_only
    def      lowerCAmelCase_  (					self	:  Optional[int]   ,				__lowerCAmelCase	:  List[Any]   ,				__lowerCAmelCase	:  str					):
      try:
        _UpperCAmelCase						 =      pl_module.model.model.num_parameters()
      except AttributeError:
        _UpperCAmelCase						 =      pl_module.model.num_parameters()

      _UpperCAmelCase						 =      count_trainable_parameters(__A					)
      # mp stands for million parameters
      trainer.logger.log_metrics({"""n_params""": npars, """mp""": npars / 1e6, """grad_mp""": n_trainable_pars / 1e6}					)


    @rank_zero_only
    def      lowerCAmelCase_  (					self	:  Dict   ,				__lowerCAmelCase	:  Union[str, Any]   ,				__lowerCAmelCase	:  Optional[int]					):
      save_json(pl_module.metrics   ,				pl_module.metrics_save_path					)
      return self._write_logs(__A   ,				__A   ,				"""test"""					)


    @rank_zero_only
    def      lowerCAmelCase_  (					self	:  Any   ,				__lowerCAmelCase	:  Dict   ,				__lowerCAmelCase	:  Any					):
      save_json(pl_module.metrics   ,				pl_module.metrics_save_path					)
      # Uncommenting this will save val generations
      # return self._write_logs(trainer, pl_module, "valid")

"""simple docstring"""


from unittest.mock import patch

import pyspark

from datasets.packaged_modules.spark.spark import (
    Spark,
    SparkExamplesIterable,
    _generate_iterable_examples,
)

from ..utils import (
    require_dill_gt_0_3_2,
    require_not_windows,
)






def     __UpperCAmelCase					(			lowercase   ,lowercase		):







  """simple docstring"""
  _UpperCAmelCase						 =      []
  for part_id in partition_order:
    _UpperCAmelCase						 =      df.where(f'''SPARK_PARTITION_ID() = {part_id}'''		).collect()
    for row_idx, row in enumerate(lowercase		):
      expected_row_ids_and_row_dicts.append((f'''{part_id}_{row_idx}''', row.asDict())		)
  return expected_row_ids_and_row_dicts






@require_not_windows
@require_dill_gt_0_3_2
def     __UpperCAmelCase					(			):







  """simple docstring"""
  _UpperCAmelCase						 =      pyspark.sql.SparkSession.builder.master("""local[*]"""		).appName("""pyspark"""		).getOrCreate()
  _UpperCAmelCase						 =      spark.range(1_00		).repartition(1		)
  _UpperCAmelCase						 =      Spark(lowercase		)
  # The id ints will be converted to Pyarrow int64s, so each row will be 8 bytes. Setting a max_shard_size of 16 means
  # that each partition can hold 2 rows.
  spark_builder._repartition_df_if_needed(max_shard_size=16		)
  # Given that the dataframe has 100 rows and each partition has 2 rows, we expect 50 partitions.
  assert spark_builder.df.rdd.getNumPartitions() == 50






@require_not_windows
@require_dill_gt_0_3_2
def     __UpperCAmelCase					(			):







  """simple docstring"""
  _UpperCAmelCase						 =      pyspark.sql.SparkSession.builder.master("""local[*]"""		).appName("""pyspark"""		).getOrCreate()
  _UpperCAmelCase						 =      spark.range(10		).repartition(2		)
  _UpperCAmelCase						 =      [1, 0]
  _UpperCAmelCase						 =      _generate_iterable_examples(lowercase   ,lowercase		)  # Reverse the partitions.
  _UpperCAmelCase						 =      _get_expected_row_ids_and_row_dicts_for_partition_order(lowercase   ,lowercase		)

  for i, (row_id, row_dict) in enumerate(generate_fn()		):
    _UpperCAmelCase ,				_UpperCAmelCase						 =      expected_row_ids_and_row_dicts[i]
    assert row_id == expected_row_id
    assert row_dict == expected_row_dict






@require_not_windows
@require_dill_gt_0_3_2
def     __UpperCAmelCase					(			):







  """simple docstring"""
  _UpperCAmelCase						 =      pyspark.sql.SparkSession.builder.master("""local[*]"""		).appName("""pyspark"""		).getOrCreate()
  _UpperCAmelCase						 =      spark.range(10		).repartition(1		)
  _UpperCAmelCase						 =      SparkExamplesIterable(lowercase		)
  assert it.n_shards == 1
  for i, (row_id, row_dict) in enumerate(lowercase		):
    assert row_id == f'''0_{i}'''
    assert row_dict == {"id": i}






@require_not_windows
@require_dill_gt_0_3_2
def     __UpperCAmelCase					(			):







  """simple docstring"""
  _UpperCAmelCase						 =      pyspark.sql.SparkSession.builder.master("""local[*]"""		).appName("""pyspark"""		).getOrCreate()
  _UpperCAmelCase						 =      spark.range(30		).repartition(3		)
  # Mock the generator so that shuffle reverses the partition indices.
  with patch("""numpy.random.Generator"""		) as generator_mock:
    _UpperCAmelCase						 =      lambda lowercase		: x.reverse()
    _UpperCAmelCase						 =      _get_expected_row_ids_and_row_dicts_for_partition_order(lowercase   ,[2, 1, 0]		)

    _UpperCAmelCase						 =      SparkExamplesIterable(lowercase		).shuffle_data_sources(lowercase		)
    assert shuffled_it.n_shards == 3
    for i, (row_id, row_dict) in enumerate(lowercase		):
      _UpperCAmelCase ,				_UpperCAmelCase						 =      expected_row_ids_and_row_dicts[i]
      assert row_id == expected_row_id
      assert row_dict == expected_row_dict






@require_not_windows
@require_dill_gt_0_3_2
def     __UpperCAmelCase					(			):







  """simple docstring"""
  _UpperCAmelCase						 =      pyspark.sql.SparkSession.builder.master("""local[*]"""		).appName("""pyspark"""		).getOrCreate()
  _UpperCAmelCase						 =      spark.range(20		).repartition(4		)

  # Partitions 0 and 2
  _UpperCAmelCase						 =      SparkExamplesIterable(lowercase		).shard_data_sources(worker_id=0   ,num_workers=2		)
  assert shard_it_a.n_shards == 2
  _UpperCAmelCase						 =      _get_expected_row_ids_and_row_dicts_for_partition_order(lowercase   ,[0, 2]		)
  for i, (row_id, row_dict) in enumerate(lowercase		):
    _UpperCAmelCase ,				_UpperCAmelCase						 =      expected_row_ids_and_row_dicts_a[i]
    assert row_id == expected_row_id
    assert row_dict == expected_row_dict

  # Partitions 1 and 3
  _UpperCAmelCase						 =      SparkExamplesIterable(lowercase		).shard_data_sources(worker_id=1   ,num_workers=2		)
  assert shard_it_a.n_shards == 2
  _UpperCAmelCase						 =      _get_expected_row_ids_and_row_dicts_for_partition_order(lowercase   ,[1, 3]		)
  for i, (row_id, row_dict) in enumerate(lowercase		):
    _UpperCAmelCase ,				_UpperCAmelCase						 =      expected_row_ids_and_row_dicts_a[i]
    assert row_id == expected_row_id
    assert row_dict == expected_row_dict






@require_not_windows
@require_dill_gt_0_3_2
def     __UpperCAmelCase					(			):







  """simple docstring"""
  _UpperCAmelCase						 =      pyspark.sql.SparkSession.builder.master("""local[*]"""		).appName("""pyspark"""		).getOrCreate()
  _UpperCAmelCase						 =      spark.range(1_00		).repartition(1		)
  _UpperCAmelCase						 =      Spark(lowercase		)
  # Choose a small max_shard_size for maximum partitioning.
  spark_builder._repartition_df_if_needed(max_shard_size=1		)
  # The new number of partitions should not be greater than the number of rows.
  assert spark_builder.df.rdd.getNumPartitions() == 1_00
