text: # Copyright 2017-2022 Lawrence Livermore National Security, LLC and other
  # Hatchet Project Developers. See the top-level LICENSE file for details.
  #
  # SPDX-License-Identifier: MIT
  
  import subprocess
  import numpy as np
  
  import pytest
  import sys
  
  from hatchet import GraphFrame
  from hatchet.readers.caliper_reader import CaliperReader
  from hatchet.util.executable import which
  from hatchet.external.console import ConsoleRenderer
  
  caliperreader_avail = True
  try:
      import caliperreader
  except ImportError:
      caliperreader_avail = False
  
  annotations = [
      "main",
      "LagrangeLeapFrog",
      "LagrangeElements",
      "ApplyMaterialPropertiesForElems",
      "EvalEOSForElems",
      "CalcEnergyForElems",
      "CalcPressureForElems",
      "CalcSoundSpeedForElems",
      "UpdateVolumesForElems",
      "CalcTimeConstraintsForElems",
      "CalcCourantConstraintForElems",
      "CalcHydroConstraintForElems",
      "TimeIncrement",
      "LagrangeNodal",
      "CalcForceForNodes",
      "CalcVolumeForceForElems",
      "IntegrateStressForElems",
      "CalcHourglassControlForElems",
      "CalcFBHourglassForceForElems",
      "CalcLagrangeElements",
      "CalcKinematicsForElems",
      "CalcQForElems",
      "CalcMonotonicQGradientsForElems",
      "CalcMonotonicQRegionForElems",
  ]
  
  
  def test_graphframe(lulesh_caliper_json):
      """Sanity test a GraphFrame object with known data."""
      gf = GraphFrame.from_caliper(str(lulesh_caliper_json))
  
      assert len(gf.dataframe.groupby("name")) == 24
  
      for col in gf.dataframe.columns:
          if col in ("time (inc)", "time"):
              assert gf.dataframe[col].dtype == np.float64
          elif col in ("nid", "rank"):
              assert gf.dataframe[col].dtype == np.int64
          elif col in ("name", "node"):
              assert gf.dataframe[col].dtype == object
  
      # TODO: add tests to confirm values in dataframe
  
  
  def test_read_lulesh_json(lulesh_caliper_json):
      """Sanity check the Caliper reader by examining a known input."""
      reader = CaliperReader(str(lulesh_caliper_json))
      reader.read_json_sections()
  
      assert len(reader.json_data) == 192
      assert len(reader.json_cols) == 4
      assert len(reader.json_cols_mdata) == 4
      assert len(reader.json_nodes) == 24
  
      reader.create_graph()
      assert all(an in reader.idx_to_label.values() for an in annotations)
  
  
  def test_calc_pi_json(calc_pi_caliper_json):
      """Sanity test a GraphFrame object with known data."""
      gf = GraphFrame.from_caliper(str(calc_pi_caliper_json))
  
      assert len(gf.dataframe.groupby("name")) == 100
  
  
  @pytest.mark.skipif(not which("cali-query"), reason="needs cali-query to be in path")
  def test_lulesh_cali(lulesh_caliper_cali):
      """Sanity check the Caliper reader ingesting a .cali file."""
      grouping_attribute = "function"
      default_metric = "sum(sum#time.duration),inclusive_sum(sum#time.duration)"
      query = "select function,%s group by %s format json-split" % (
          default_metric,
          grouping_attribute,
      )
  
      gf = GraphFrame.from_caliper(str(lulesh_caliper_cali), query)
  
      assert len(gf.dataframe.groupby("name")) == 18
  
  
  @pytest.mark.skipif(not which("cali-query"), reason="needs cali-query to be in path")
  def test_lulesh_json_stream(lulesh_caliper_cali):
      """Sanity check the Caliper reader ingesting a JSON string literal."""
      cali_query = which("cali-query")
      grouping_attribute = "function"
      default_metric = "sum(sum#time.duration),inclusive_sum(sum#time.duration)"
      query = "select function,%s group by %s format json-split" % (
          default_metric,
          grouping_attribute,
      )
  
      cali_json = subprocess.Popen(
          [cali_query, "-q", query, lulesh_caliper_cali], stdout=subprocess.PIPE
      )
  
      gf = GraphFrame.from_caliper(cali_json.stdout)
  
      assert len(gf.dataframe.groupby("name")) == 18
  
  
  @pytest.mark.skipif(sys.version_info > (3, 8), reason="Temporarily allow this to fail.")
  def test_filter_squash_unify_caliper_data(lulesh_caliper_json):
      """Sanity test a GraphFrame object with known data."""
      gf1 = GraphFrame.from_caliper(str(lulesh_caliper_json))
      gf2 = GraphFrame.from_caliper(str(lulesh_caliper_json))
  
      assert gf1.graph is not gf2.graph
  
      gf1_index_names = gf1.dataframe.index.names
      gf2_index_names = gf2.dataframe.index.names
  
      gf1.dataframe.reset_index(inplace=True)
      gf2.dataframe.reset_index(inplace=True)
  
      # indexes are the same since we are reading in the same dataset
      assert all(gf1.dataframe["node"] == gf2.dataframe["node"])
  
      gf1.dataframe.set_index(gf1_index_names, inplace=True)
      gf2.dataframe.set_index(gf2_index_names, inplace=True)
  
      squash_gf1 = gf1.filter(lambda x: x["name"].startswith("Calc"))
      squash_gf2 = gf2.filter(lambda x: x["name"].startswith("Calc"))
  
      squash_gf1.unify(squash_gf2)
  
      assert squash_gf1.graph is squash_gf2.graph
  
      squash_gf1.dataframe.reset_index(inplace=True)
      squash_gf2.dataframe.reset_index(inplace=True)
  
      # Indexes should still be the same after unify. Sort indexes before comparing.
      assert all(squash_gf1.dataframe["node"] == squash_gf2.dataframe["node"])
  
      squash_gf1.dataframe.set_index(gf1_index_names, inplace=True)
      squash_gf2.dataframe.set_index(gf2_index_names, inplace=True)
  
  
  def test_tree(lulesh_caliper_json):
      """Sanity test a GraphFrame object with known data."""
      gf = GraphFrame.from_caliper(str(lulesh_caliper_json))
  
      output = ConsoleRenderer(unicode=True, color=False).render(
          gf.graph.roots,
          gf.dataframe,
          metric_column="time",
          precision=3,
          name_column="name",
          expand_name=False,
          context_column="file",
          rank=0,
          thread=0,
          depth=10000,
          highlight_name=False,
          colormap="RdYlGn",
          invert_colormap=False,
      )
      assert "121489.000 main" in output
      assert "663.000 LagrangeElements" in output
      assert "21493.000 CalcTimeConstraintsForElems" in output
  
      output = ConsoleRenderer(unicode=True, color=False).render(
          gf.graph.roots,
          gf.dataframe,
          metric_column="time (inc)",
          precision=3,
          name_column="name",
          expand_name=False,
          context_column="file",
          rank=0,
          thread=0,
          depth=10000,
          highlight_name=False,
          colormap="RdYlGn",
          invert_colormap=False,
      )
      assert "662712.000 EvalEOSForElems" in output
      assert "2895319.000 LagrangeNodal" in output
  
  
  def test_graphframe_to_literal(lulesh_caliper_json):
      """Sanity test a GraphFrame object with known data."""
      gf = GraphFrame.from_caliper(str(lulesh_caliper_json))
      graph_literal = gf.to_literal()
  
      gf2 = GraphFrame.from_literal(graph_literal)
  
      assert len(gf.graph) == len(gf2.graph)
  
  
  def test_graphframe_native_lulesh_from_file(lulesh_caliper_cali):
      """Sanity check the native Caliper reader by examining a known input."""
  
      gf = GraphFrame.from_caliperreader(str(lulesh_caliper_cali))
  
      assert len(gf.dataframe.groupby("name")) == 19
      assert "cali.caliper.version" in gf.metadata.keys()
  
      for col in gf.dataframe.columns:
          if col in ("time (inc)", "time"):
              assert gf.dataframe[col].dtype == np.float64
          elif col in ("nid", "rank"):
              assert gf.dataframe[col].dtype == np.int64
          elif col in ("name", "node"):
              assert gf.dataframe[col].dtype == object
  
  
  @pytest.mark.skipif(
      not caliperreader_avail, reason="needs caliper-reader package to be loaded"
  )
  def test_graphframe_native_lulesh_from_caliperreader(lulesh_caliper_cali):
      """Sanity check the native Caliper reader by examining a known input."""
      r = caliperreader.CaliperReader()
      r.read(lulesh_caliper_cali)
  
      gf = GraphFrame.from_caliperreader(r)
  
      assert len(gf.dataframe.groupby("name")) == 19
      assert "cali.caliper.version" in gf.metadata.keys()
  
      for col in gf.dataframe.columns:
          if col in ("time (inc)", "time"):
              assert gf.dataframe[col].dtype == np.float64
          elif col in ("nid", "rank"):
              assert gf.dataframe[col].dtype == np.int64
          elif col in ("name", "node"):
              assert gf.dataframe[col].dtype == object
  
  
  def test_inclusive_time_calculation(lulesh_caliper_json):
      """Validate update_inclusive_columns() on known dataset containing per-rank data."""
      gf = GraphFrame.from_caliper(str(lulesh_caliper_json))
  
      # save original time (inc) column for correctness check
      gf.dataframe["orig_inc_time"] = gf.dataframe["time (inc)"]
  
      # remove original time (inc) column since it will be generated by update_inclusive_columns()
      del gf.dataframe["time (inc)"]
  
      gf.update_inclusive_columns()
      assert all(
          gf.dataframe["time (inc)"].values == gf.dataframe["orig_inc_time"].values
      )
  
  
  def test_sw4_cuda_from_caliperreader(sw4_caliper_cuda_activity_profile_cali):
      gf = GraphFrame.from_caliperreader(sw4_caliper_cuda_activity_profile_cali)
  
      assert len(gf.graph) == 549
      assert all(
          metric in gf.dataframe.columns for metric in gf.exc_metrics + gf.inc_metrics
      )
  
      for col in gf.dataframe.columns:
          if col in ("#scale#cupti.activity.duration", "#scale#sum#cupti.host.duration"):
              assert gf.dataframe[col].dtype == np.float64
          elif col in "rank":
              assert gf.dataframe[col].dtype == np.int64
          elif col in "name":
              assert gf.dataframe[col].dtype == object
  
      for col in gf.exc_metrics + gf.inc_metrics:
          assert col in gf.dataframe.columns
  
  
  def test_sw4_cuda_summary_from_caliperreader(
      sw4_caliper_cuda_activity_profile_summary_cali,
  ):
      gf = GraphFrame.from_caliperreader(sw4_caliper_cuda_activity_profile_summary_cali)
  
      assert len(gf.graph) == 393
      assert all(
          metric in gf.dataframe.columns for metric in gf.exc_metrics + gf.inc_metrics
      )
  
      for col in gf.exc_metrics + gf.inc_metrics:
          assert col in gf.dataframe.columns
  