# Copyright 2022 The Plenoptix Team. 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.
# 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.

"""
Code to train model.
"""
from __future__ import annotations

import dataclasses
import functools
import logging
import os
import time
from pathlib import Path
from typing import Any, Dict, List, Tuple

import torch
from rich import console
from torch.cuda.amp.grad_scaler import GradScaler

from nerfactory.configs import base as cfg
from nerfactory.optimizers.optimizers import Optimizers, setup_optimizers
from nerfactory.pipelines.base import Pipeline
from nerfactory.utils import profiler, writer
from nerfactory.utils.callbacks import (
    TrainingCallback,
    TrainingCallbackAttributes,
    TrainingCallbackLocation,
)
from nerfactory.utils.decorators import check_main_thread, check_viewer_enabled
from nerfactory.utils.misc import step_check
from nerfactory.utils.writer import EventName, TimeWriter
from nerfactory.viewer.server import viewer_utils

logging.getLogger("PIL").setLevel(logging.WARNING)
CONSOLE = console.Console()


def train_loop(local_rank: int, world_size: int, config: cfg.Config) -> Any:
    """Main training function that sets up and runs the trainer per process

    Args:
        local_rank: current rank of process
        world_size: total number of gpus available
        config: config file specifying training regimen
    """
    trainer = Trainer(config, local_rank, world_size)
    trainer.setup()
    trainer.train()
    return 0


class Trainer:
    """Trainer class

    Args:
        config: The configuration object.
        local_rank: Local rank of the process.
        world_size: World size of the process.

    Attributes:
        config: The configuration object.
        local_rank: Local rank of the process.
        world_size: World size of the process.
        device: The device to run the training on.
        pipeline: The pipeline object.
        optimizers: The optimizers object.
        callbacks: The callbacks object.
    """

    pipeline: Pipeline
    optimizers: Optimizers
    callbacks: List[TrainingCallback]

    def __init__(self, config: cfg.Config, local_rank: int = 0, world_size: int = 1):
        self.config = config
        self.local_rank = local_rank
        self.world_size = world_size
        self.device = "cpu" if world_size == 0 else f"cuda:{local_rank}"
        self.mixed_precision = self.config.trainer.mixed_precision
        if self.device == "cpu":
            self.mixed_precision = False
            logging.warning("Mixed precision is disabled for CPU training.")
        self._start_step = 0
        # optimizers
        self.grad_scaler = GradScaler(enabled=self.mixed_precision)

        self.base_dir = config.get_base_dir()
        # directory to save checkpoints
        self.checkpoint_dir = Path(self.base_dir / self.config.trainer.relative_model_dir)
        logging.info("Saving checkpoints to: %s", self.checkpoint_dir)
        # set up viewer if enabled
        viewer_log_path = self.base_dir / config.viewer.relative_log_filename
        self.viewer_state, banner_messages = viewer_utils.setup_viewer(config.viewer, log_filename=viewer_log_path)
        self._check_viewer_warnings()
        # set up writers/profilers if enabled
        writer_log_path = self.base_dir / config.logging.relative_log_dir
        writer.setup_event_writer(config.logging, log_dir=writer_log_path)
        writer.setup_local_writer(
            config.logging, max_iter=config.trainer.max_num_iterations, banner_messages=banner_messages
        )
        writer.put_config(name="config", config_dict=dataclasses.asdict(config), step=0)
        profiler.setup_profiler(config.logging)

    def setup(self, test_mode=False):
        """Setup the Trainer by calling other setup functions.

        Args:
            test_mode: Whether to setup for testing. Defaults to False.
        """
        self.pipeline: Pipeline = self.config.pipeline.setup(
            device=self.device, test_mode=test_mode, world_size=self.world_size, local_rank=self.local_rank
        )
        self.optimizers = setup_optimizers(self.config.optimizers, self.pipeline.get_param_groups())

        self._load_checkpoint()

        # TODO(ethan): do this for pipeline, not pipeline.model
        self.callbacks = self.pipeline.get_training_callbacks(
            TrainingCallbackAttributes(
                optimizers=self.optimizers,  # type: ignore
                grad_scaler=self.grad_scaler,  # type: ignore
                pipeline=self.pipeline,  # type: ignore
            )
        )

    def train(self) -> None:
        """Train the model."""
        assert self.pipeline.datamanager.train_input_dataset is not None, "Missing DatsetInputs"

        self._init_viewer_scene()
        with TimeWriter(writer, EventName.TOTAL_TRAIN_TIME):
            num_iterations = self.config.trainer.max_num_iterations
            for step in range(self._start_step, self._start_step + num_iterations):
                # if the viewer used, the rendering of the viewer will be included in the iteration train time
                with TimeWriter(writer, EventName.ITER_TRAIN_TIME, step=step) as train_t:

                    self.pipeline.train()

                    # training callbacks before the training iteration
                    for callback in self.callbacks:
                        callback.run_callback_at_location(
                            step, location=TrainingCallbackLocation.BEFORE_TRAIN_ITERATION
                        )

                    # time the forward pass
                    loss, loss_dict, metrics_dict = self.train_iteration(step)

                    # training callbacks after the training iteration
                    for callback in self.callbacks:
                        callback.run_callback_at_location(step, location=TrainingCallbackLocation.AFTER_TRAIN_ITERATION)

                writer.put_time(
                    name=EventName.TRAIN_RAYS_PER_SEC,
                    duration=self.config.pipeline.datamanager.train_num_rays_per_batch / train_t.duration,
                    step=step,
                    avg_over_steps=True,
                )

                self._update_viewer_state(step)

                # a batch of train rays
                if step_check(step, self.config.logging.steps_per_log, run_at_zero=True):
                    writer.put_scalar(name="Train Loss", scalar=loss, step=step)
                    writer.put_dict(name="Train Loss Dict", scalar_dict=loss_dict, step=step)
                    writer.put_dict(name="Train Metrics Dict", scalar_dict=metrics_dict, step=step)

                # a batch of eval rays
                if step_check(step, self.config.trainer.steps_per_eval_batch, run_at_zero=True):
                    _, eval_loss_dict, eval_metrics_dict = self.pipeline.get_eval_loss_dict(step=step)
                    eval_loss = functools.reduce(torch.add, eval_loss_dict.values())
                    writer.put_scalar(name="Eval Loss", scalar=eval_loss, step=step)
                    writer.put_dict(name="Eval Loss Dict", scalar_dict=eval_loss_dict, step=step)
                    writer.put_dict(name="Eval Metrics Dict", scalar_dict=eval_metrics_dict, step=step)

                # one eval image
                if step_check(step, self.config.trainer.steps_per_eval_image):
                    with TimeWriter(writer, EventName.TEST_RAYS_PER_SEC, write=False) as test_t:
                        metrics_dict, images_dict = self.pipeline.get_eval_image_metrics_and_images(step=step)
                    writer.put_time(
                        name=EventName.TEST_RAYS_PER_SEC,
                        duration=metrics_dict["num_rays"] / test_t.duration,
                        step=step,
                        avg_over_steps=True,
                    )
                    writer.put_dict(name="Eval Images Metrics", scalar_dict=metrics_dict, step=step)
                    group = "Eval Images"
                    for image_name, image in images_dict.items():
                        writer.put_image(name=group + "/" + image_name, image=image, step=step)

                # all eval images
                if step_check(step, self.config.trainer.steps_per_eval_all_images):
                    metrics_dict = self.pipeline.get_average_eval_image_metrics(step=step)
                    writer.put_dict(name="Eval Images Metrics Dict (all images)", scalar_dict=metrics_dict, step=step)

                if step != 0 and self.config.trainer.steps_per_save and step % self.config.trainer.steps_per_save == 0:
                    self._save_checkpoint(step)

                self._write_out_storage(step)

        self._write_out_storage(num_iterations)

    def _check_viewer_warnings(self) -> None:
        """Helper to print out any warnings regarding the way the viewer/loggers are enabled"""
        if self.config.viewer.enable:
            if self.config.logging.event_writer != "none":
                string = (
                    "[WARNING]: Tensorboard or Wandb enabled with Viewer will slow down Viewer. "
                    "Please set `--logging.event_writer none` for faster rendering"
                )
                CONSOLE.print(f"[bold red]{string}")
            self.config.trainer.steps_per_eval_batch = self.config.trainer.max_num_iterations
            self.config.trainer.steps_per_eval_image = self.config.trainer.max_num_iterations
            self.config.trainer.steps_per_eval_all_images = self.config.trainer.max_num_iterations
            string = "[WARNING] Disabling eval iterations since viewer is enabled."
            CONSOLE.print(f"[bold red]{string}")

    @check_viewer_enabled
    def _init_viewer_scene(self) -> None:
        """Initializes viewer scene with given train dataset"""
        assert self.viewer_state and self.pipeline.datamanager.train_input_dataset
        self.viewer_state.init_scene(
            dataset=self.pipeline.datamanager.train_input_dataset,
            start_train=self.config.viewer.start_train,
        )

    @check_viewer_enabled
    def _update_viewer_state(self, step: int):
        """Updates the viewer state by rendering out scene with current pipeline
        Returns the time taken to render scene.

        Args:
            step: current train step
        """
        assert self.viewer_state is not None
        with TimeWriter(writer, EventName.ITER_VIS_TIME, step=step) as _:
            num_rays_per_batch = self.config.pipeline.datamanager.train_num_rays_per_batch
            try:
                self.viewer_state.update_scene(step, self.pipeline.model, num_rays_per_batch)
            except RuntimeError:
                time.sleep(0.03)  # sleep to allow buffer to reset
                assert self.viewer_state.vis is not None
                self.viewer_state.vis["renderingState/log_errors"].write(
                    "Error: GPU out of memory. Reduce resolution to prevent viewer from crashing."
                )

    @check_viewer_enabled
    def _update_viewer_rays_per_sec(self, train_t: TimeWriter, vis_t: TimeWriter, step: int):
        """Performs update on rays/sec calclation for training

        Args:
            train_t: timer object carrying time to execute total training iteration
            vis_t: timer object carrying time to execute visualization step
            step: current step
        """
        train_num_rays_per_batch = self.config.pipeline.datamanager.train_num_rays_per_batch
        writer.put_time(
            name=EventName.TRAIN_RAYS_PER_SEC,
            duration=train_num_rays_per_batch / (train_t.duration - vis_t.duration),
            step=step,
            avg_over_steps=True,
        )

    def _write_out_storage(self, step: int) -> None:
        """Perform writes only during appropriate time steps

        Args:
            step: Current training step.
        """
        if (
            step % self.config.logging.steps_per_log == 0
            or step % self.config.trainer.steps_per_eval_batch == 0
            or step % self.config.trainer.steps_per_eval_image == 0
            or step % self.config.trainer.steps_per_eval_all_images == 0
        ):
            writer.write_out_storage()

    def _load_checkpoint(self) -> None:
        """Helper function to load pipeline and optimizer from prespecified checkpoint"""
        load_dir = self.config.trainer.load_dir
        if load_dir is not None:
            load_step = self.config.trainer.load_step
            if load_step is None:
                print("Loading latest checkpoint from load_dir")
                # NOTE: this is specific to the checkpoint name format
                load_step = sorted(int(x[x.find("-") + 1 : x.find(".")]) for x in os.listdir(load_dir))[-1]
            load_path = load_dir / f"step-{load_step:09d}.ckpt"
            assert load_path.exists(), f"Checkpoint {load_path} does not exist"
            loaded_state = torch.load(load_path, map_location="cpu")
            self._start_step = loaded_state["step"] + 1
            # load the checkpoints for pipeline, optimizers, and gradient scalar
            self.pipeline.load_pipeline(loaded_state["pipeline"])
            self.optimizers.load_optimizers(loaded_state["optimizers"])
            self.grad_scaler.load_state_dict(loaded_state["scalers"])
            logging.info("done loading checkpoint from %s", load_path)
        else:
            logging.info("No checkpoints to load, training from scratch")

    @check_main_thread
    def _save_checkpoint(self, step: int) -> None:
        """Save the model and optimizers

        Args:
            step: number of steps in training for given checkpoint
        """
        if not self.checkpoint_dir.exists():
            self.checkpoint_dir.mkdir(parents=True, exist_ok=True)
        ckpt_path = self.checkpoint_dir / f"step-{step:09d}.ckpt"
        if hasattr(self.pipeline, "module"):
            pipeline = self.pipeline.module.state_dict()  # type: ignore
        else:
            pipeline = self.pipeline.state_dict()
        torch.save(
            {
                "step": step,
                "pipeline": pipeline,
                "optimizers": {k: v.state_dict() for (k, v) in self.optimizers.optimizers.items()},
                "scalers": self.grad_scaler.state_dict(),
            },
            ckpt_path,
        )

    @profiler.time_function
    def train_iteration(self, step: int) -> Tuple[torch.Tensor, Dict[str, torch.Tensor], Dict[str, torch.Tensor]]:
        """Run one iteration with a batch of inputs. Returns dictionary of model losses.

        Args:
            step: Current training step.
        """
        self.optimizers.zero_grad_all()
        cpu_or_cuda_str = self.device.split(":")[0]
        with torch.autocast(device_type=cpu_or_cuda_str, enabled=self.mixed_precision):
            _, loss_dict, metrics_dict = self.pipeline.get_train_loss_dict(step=step)
            loss = functools.reduce(torch.add, loss_dict.values())
        self.grad_scaler.scale(loss).backward()  # type: ignore
        self.optimizers.optimizer_scaler_step_all(self.grad_scaler)
        self.grad_scaler.update()
        self.optimizers.scheduler_step_all(step)

        # Merging loss and metrics dict into a single output.
        return loss, loss_dict, metrics_dict