# coding=utf-8
# Copyright 2023 The Google Research Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License 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.

"""Evaluation script for mipNeRF360."""

import functools
import gc
from os import path
import sys
import time

from absl import app
from absl import logging
import chex
import flax
from flax.metrics import tensorboard
from flax.training import checkpoints
import gin
from internal import alignment
from internal import camera_utils
from internal import configs
from internal import datasets
from internal import image_io
from internal import image_utils
from internal import models
from internal import ref_utils
from internal import train_utils
from internal import utils
from internal import vis
import jax
from jax import random
import jax.numpy as jnp
import jaxcam
import numpy as np


configs.define_common_flags()
jax.config.parse_flags_with_absl()


def plot_camera_metrics(
    *,
    summary_writer,
    camera_params,
    train_cameras,
    train_cameras_gt,
    config,
    step,
    tag,
):
  """Plots camera statistics to TensorBoard."""
  camera_delta = config.camera_delta_cls()
  optimized_cameras: jaxcam.Camera = camera_delta.apply(
      camera_params, train_cameras
  )
  diffs = camera_utils.compute_camera_metrics(
      train_cameras_gt, optimized_cameras
  )
  reduce_fns = {
      'mean': np.mean,
      'max': np.max,
      'std': np.std,
  }
  for reduce_name, reduce_fn in reduce_fns.items():
    for stat_name, stat in diffs.items():
      summary_writer.scalar(
          f'eval_train_camera_{tag}_{reduce_name}/{stat_name}',
          reduce_fn(np.array(stat)),
          step=step,
      )

  # pylint: disable=cell-var-from-loop
  return [
      jax.tree_util.tree_map(lambda x: float(x[i]), diffs)
      for i in range(len(train_cameras))
  ]
  # pylint: enable=cell-var-from-loop


def main(unused_argv):
  config = configs.load_config(save_config=False)

  train_dataset = datasets.load_dataset('train', config.data_dir, config)
  test_dataset = datasets.load_dataset('test', config.data_dir, config)

  key = random.PRNGKey(20200823)
  model, state, render_eval_pfn, _, _ = train_utils.setup_model(
      config, key, dataset=train_dataset
  )
  if config.rawnerf_mode:
    postprocess_fn = test_dataset.metadata['postprocess_fn']
  else:
    postprocess_fn = lambda z: z

  metric_harness = image_utils.MetricHarness(
      **config.metric_harness_eval_config
  )

  last_step = 0
  out_dir = path.join(
      config.checkpoint_dir,
      'path_renders' if config.render_path else 'test_preds',
  )
  path_fn = lambda x: path.join(out_dir, x)

  if not config.eval_only_once:
    summary_writer = tensorboard.SummaryWriter(
        path.join(config.checkpoint_dir, 'eval')
    )

  jnp_cameras = None
  if config.cast_rays_in_eval_step:
    np_to_jax = lambda x: jnp.array(x) if isinstance(x, np.ndarray) else x
    jnp_cameras = jax.tree_util.tree_map(np_to_jax, test_dataset.cameras)

  jnp_cameras_replicated = flax.jax_utils.replicate(jnp_cameras)

  last_eval_time = time.time()
  while True:
    state = checkpoints.restore_checkpoint(config.checkpoint_dir, state)
    step = int(state.step)
    state_params_replicated = flax.jax_utils.replicate(state.params)

    if step  config.eval_checkpoint_wait_timeout_sec:
        raise RuntimeError(
            'Waited for a new checkpoint for'
            f' {config.eval_checkpoint_wait_timeout_sec} seconds, got no new'
            ' checkpoint. This likely means that the training script has died.'
            ' Exiting. If this is expected, increase'
            ' config.eval_checkpoint_wait_timeout_sec.'
        )
      logging.info(
          'Checkpoint step %d