text: # Copyright 2018 Amazon.com, Inc. or its affiliates. All Rights Reserved.
  #
  # Licensed under the Apache License, Version 2.0 (the "License").
  # You may not use this file except in compliance with the License.
  # A copy of the License is located at
  #
  #     http://www.apache.org/licenses/LICENSE-2.0
  #
  # or in the "license" file accompanying this file. This file is distributed
  # on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
  # express or implied. See the License for the specific language governing
  # permissions and limitations under the License.
  
  import torch
  import argparse
  from pathlib import Path
  import numpy as np
  from .model_utils import build_model
  from .datasets import (
      GTSUnivariateDataset,
      BouncingBallDataset,
      ThreeModeSystemDataset,
      BeeDataset,
  )
  from .evaluation import evaluate_segmentation, evaluate_gts_dataset
  from .torch_utils import torch2numpy
  
  
  def get_test_dataset(config, test_path=Path("./data/")):
      if config["experiment"] == "gts_univariate":
          test_dataset = GTSUnivariateDataset(
              config["dataset"], batch_size=50, mode="test"
          )
      else:
          if config["dataset"] == "3modesystem":
              test_dataset = ThreeModeSystemDataset(
                  path=str(test_path / f"{config['dataset']}_test.npz")
              )
          elif config["dataset"] == "bee":
              test_dataset = BeeDataset(path="./data/bee_test.npz")
          elif config["dataset"] == "bouncing_ball":
              test_dataset = BouncingBallDataset(
                  path="./data/bouncing_ball_test.npz"
              )
          else:
              raise ValueError(f"There is no {config['dataset']} dataset.")
  
      return test_dataset
  
  
  if __name__ == "__main__":
      parser = argparse.ArgumentParser()
      parser.add_argument("--ckpt", type=str, required=True)
      parser.add_argument("--dev", type=str, default="cuda:0")
      args = parser.parse_args()
  
      ckpt_file = args.ckpt
      try:
          ckpt = torch.load(ckpt_file, map_location="cpu")
      except FileNotFoundError:
          print(f"Can't find {ckpt_file}!")
  
      config = ckpt["config"]
      device = torch.device(args.dev)
      model = build_model(config=config)
      model.load_state_dict(ckpt["model"])
      model = model.to(device)
  
      test_dataset = get_test_dataset(config)
  
      if config["experiment"] == "gts_univariate":
          test_loader = torch.utils.data.DataLoader(
              test_dataset, batch_size=None
          )
  
          ctx_len = config["context_length"]
          pred_len = config["prediction_length"]
  
          agg_metrics = evaluate_gts_dataset(
              test_dataset,
              model,
              max_len=np.inf,
              device=device,
              deterministic_z=False,
              deterministic_x=False,
              deterministic_y=False,
              batch_size=100,
          )
  
          CRPS = agg_metrics["mean_wQuantileLoss"]
          print(config["dataset"], CRPS)
  
      else:
          test_loader = torch.utils.data.DataLoader(test_dataset, batch_size=80)
          extra_args = dict()
          if config["model"] == "REDSDS":
              extra_args = {"dur_temperature": 1.0}
          true_segs = []
          pred_segs = []
          true_tss = []
          recons_tss = []
          for test_batch, test_label in test_loader:
              test_batch = test_batch.to(device)
              test_result = model(
                  test_batch,
                  switch_temperature=1.0,
                  num_samples=1,
                  deterministic_inference=True,
                  **extra_args,
              )
              pred_seg = torch2numpy(
                  torch.argmax(test_result["log_gamma"][0], dim=-1)
              )
              true_seg = torch2numpy(test_label[:, : config["context_length"]])
              true_ts = torch2numpy(test_result["inputs"])
              recons_ts = torch2numpy(test_result["reconstructions"])
              true_tss.append(true_ts)
              recons_tss.append(recons_ts)
              true_segs.append(true_seg)
              pred_segs.append(pred_seg)
          true_tss = np.concatenate(true_tss, 0)
          recons_tss = np.concatenate(recons_tss, 0)
          true_segs = np.concatenate(true_segs, 0)
          pred_segs = np.concatenate(pred_segs, 0)
          seg_metrics = evaluate_segmentation(
              true_segs, pred_segs, K=config["num_categories"]
          )
          print(config["dataset"], seg_metrics)
  