# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
# 
# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.

#!/usr/bin/env python
# coding: utf-8

import os
import sys

import torch as th
import numpy as np
from omegaconf import OmegaConf

from torch.nn.parallel import DistributedDataParallel as DDP
import torch.distributed as dist

from torch.utils.data import DataLoader

from dva.ray_marcher import RayMarcher, generate_colored_boxes
from dva.zju_dataset import ZJUDataset
from dva.io import load_static_assets, load_from_config
from dva.losses import process_losses
from dva.utils import to_device

from torchvision.utils import make_grid, save_image

import logging

device = th.device("cuda")

logger = logging.getLogger("train.py")


def render_mvp_boxes(rm, batch, preds):
    with th.no_grad():
        boxes_rgba = generate_colored_boxes(
            preds["prim_rgba"],
            preds["prim_rot"],
        )
        preds_boxes = rm(
            prim_rgba=boxes_rgba,
            prim_pos=preds["prim_pos"],
            prim_scale=preds["prim_scale"],
            prim_rot=preds["prim_rot"],
            RT=batch["Rt"],
            K=batch["K"],
        )

    return preds_boxes["rgba_image"][:, :3].permute(0, 2, 3, 1)


def save_image_summary(path, batch, preds):
    # TODO: process_fn here?
    rgb = preds["rgb"].detach().permute(0, 3, 1, 2)
    rgb_gt = batch["image"]
    rgb_boxes = preds["rgb_boxes"].detach().permute(0, 3, 1, 2)
    img = make_grid(th.cat([rgb, rgb_gt, rgb_boxes], dim=2) / 255.0).clip(0.0, 1.0)
    save_image(img, path)


def main(config):
    dist.init_process_group("nccl")

    logging.basicConfig(level=logging.INFO)

    local_rank = int(os.environ["LOCAL_RANK"])
    device = th.device(f"cuda:{local_rank}")
    th.cuda.set_device(device)

    static_assets = load_static_assets(config)
    os.makedirs(f"{config.output_dir}/checkpoints", exist_ok=True)
    OmegaConf.save(config, f"{config.output_dir}/config.yml")
    logger.info(f"saving results to {config.output_dir}")

    logger.info(f"starting training with the config: {OmegaConf.to_yaml(config)}")

    model = load_from_config(
        config.model,
        assets=static_assets,
    ).to(device)

    # computing values for the given viewpoints
    rm = RayMarcher(
        config.image_height,
        config.image_width,
        **config.rm,
    ).to(device)

    loss_fn = load_from_config(config.loss).to(device)
    optimizer = load_from_config(config.optimizer, params=model.parameters())

    model_ddp = th.nn.parallel.DistributedDataParallel(
        model, device_ids=[device], find_unused_parameters=True
    )

    dataset = ZJUDataset(
        **config.data,
        cameras=config.cameras_train,
        cond_cameras=config.cameras_cond,
    )

    train_sampler = th.utils.data.distributed.DistributedSampler(dataset)

    loader = DataLoader(
        dataset,
        batch_size=config.train.get("batch_size", 4),
        pin_memory=False,
        sampler=train_sampler,
        num_workers=config.train.get("n_workers", 1),
        drop_last=False,
        worker_init_fn=lambda _: np.random.seed(),
    )

    iteration = 0

    for epoch in range(config.train.n_epochs):
        for b, batch in enumerate(loader):
            batch = to_device(batch, device)

            if local_rank == 0 and batch is None:
                logger.info(f"batch {b} is None, skipping")
                continue

            if local_rank == 0 and iteration >= config.train.n_max_iters:
                logger.info(f"stopping after {config.train.n_max_iters}")
                break

            preds = model_ddp(**batch, train_iter=iteration)

            # rendering and raymarching
            rm_preds = rm(
                prim_rgba=preds["prim_rgba"],
                prim_pos=preds["prim_pos"],
                prim_scale=preds["prim_scale"],
                prim_rot=preds["prim_rot"],
                RT=batch["Rt"],
                K=batch["K"],
            )

            rgba = rm_preds["rgba_image"].permute(0, 2, 3, 1)

            preds.update(
                alpha=rgba[..., -1].contiguous(), rgb=rgba[..., :3].contiguous()
            )

            loss, loss_dict = loss_fn(batch, preds, iteration)
            _loss_dict = process_losses(loss_dict)

            if th.isnan(loss):
                loss_str = " ".join([f"{k}={v:.4f}" for k, v in _loss_dict.items()])
                logger.warning(f"some of the losses is NaN, skipping: {loss_str}")
                continue

            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

            if local_rank == 0 and iteration % config.train.log_every_n_steps == 0:
                loss_str = " ".join([f"{k}={v:.4f}" for k, v in _loss_dict.items()])
                logger.info(f"epoch={epoch}, iter={iteration}: {loss_str}")

            if (
                local_rank == 0
                # and iteration
                and iteration % config.train.summary_every_n_steps == 0
            ):
                logger.info(
                    f"saving summary to {config.output_dir} after {iteration} steps"
                )
                with th.no_grad():
                    preds["rgb_boxes"] = render_mvp_boxes(rm, batch, preds)
                save_image_summary(f"{config.output_dir}/train.png", batch, preds)

            if (
                local_rank == 0
                and iteration
                and iteration % config.train.ckpt_every_n_steps == 0
            ):
                logger.info(f"saving checkpoint after {iteration} steps")

                params = {
                    "model_state_dict": model.state_dict(),
                    "optimizer": optimizer.state_dict(),
                }
                th.save(params, f"{config.output_dir}/checkpoints/{iteration:06d}.pt")

            iteration += 1

        pass


if __name__ == "__main__":
    th.backends.cudnn.benchmark = True
    # set config
    config = OmegaConf.load(str(sys.argv[1]))
    config_cli = OmegaConf.from_cli(args_list=sys.argv[2:])
    if config_cli:
        logger.info("overriding with following values from args:")
        logger.info(OmegaConf.to_yaml(config_cli))
        config = OmegaConf.merge(config, config_cli)

    main(config)
