# Training to a set of multiple objects (e.g. ShapeNet or DTU)
# tensorboard logs available in logs/<expname>

import sys
import os

sys.path.insert(
    0, os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "src"))
)

import warnings
import trainlib
from model import make_model, loss
from render import NeRFRenderer
from data import get_split_dataset
import util
import numpy as np
import torch.nn.functional as F
import torch
from dotmap import DotMap


def extra_args(parser):
    parser.add_argument(
        "--batch_size", "-B", type=int, default=4, help="Object batch size ('SB')"
    )
    parser.add_argument(
        "--nviews",
        "-V",
        type=str,
        default="1",
        help="Number of source views (multiview); put multiple (space delim) to pick randomly per batch ('NV')",
    )
    parser.add_argument(
        "--freeze_enc",
        action="store_true",
        default=None,
        help="Freeze encoder weights and only train MLP",
    )

    parser.add_argument(
        "--no_bbox_step",
        type=int,
        default=100000,
        help="Step to stop using bbox sampling",
    )
    parser.add_argument(
        "--fixed_test",
        action="store_true",
        default=None,
        help="Freeze encoder weights and only train MLP",
    )
    return parser


args, conf = util.args.parse_args(extra_args, training=True, default_ray_batch_size=128)
device = util.get_cuda(args.gpu_id[0])

dset, val_dset, _ = get_split_dataset(args.dataset_format, args.datadir)
print(
    "dset z_near {}, z_far {}, lindisp {}".format(dset.z_near, dset.z_far, dset.lindisp)
)

net = make_model(conf["model"]).to(device=device)
net.stop_encoder_grad = args.freeze_enc
if args.freeze_enc:
    print("Encoder frozen")
    net.encoder.eval()

renderer = NeRFRenderer.from_conf(conf["renderer"], lindisp=dset.lindisp,).to(
    device=device
)

# Parallize
render_par = renderer.bind_parallel(net, args.gpu_id).eval()

nviews = list(map(int, args.nviews.split()))


class PixelNeRFTrainer(trainlib.Trainer):
    def __init__(self):
        super().__init__(net, dset, val_dset, args, conf["train"], device=device)
        self.renderer_state_path = "%s/%s/_renderer" % (
            self.args.checkpoints_path,
            self.args.name,
        )

        self.lambda_coarse = conf.get_float("loss.lambda_coarse")
        self.lambda_fine = conf.get_float("loss.lambda_fine", 1.0)
        print(
            "lambda coarse {} and fine {}".format(self.lambda_coarse, self.lambda_fine)
        )
        self.rgb_coarse_crit = loss.get_rgb_loss(conf["loss.rgb"], True)
        fine_loss_conf = conf["loss.rgb"]
        if "rgb_fine" in conf["loss"]:
            print("using fine loss")
            fine_loss_conf = conf["loss.rgb_fine"]
        self.rgb_fine_crit = loss.get_rgb_loss(fine_loss_conf, False)

        if args.resume:
            if os.path.exists(self.renderer_state_path):
                renderer.load_state_dict(
                    torch.load(self.renderer_state_path, map_location=device)
                )

        self.z_near = dset.z_near
        self.z_far = dset.z_far

        self.use_bbox = args.no_bbox_step > 0

    def post_batch(self, epoch, batch):
        renderer.sched_step(args.batch_size)

    def extra_save_state(self):
        torch.save(renderer.state_dict(), self.renderer_state_path)

    def calc_losses(self, data, is_train=True, global_step=0):
        if "images" not in data:
            return {}
        all_images = data["images"].to(device=device)  # (SB, NV, 3, H, W)

        SB, NV, _, H, W = all_images.shape
        all_poses = data["poses"].to(device=device)  # (SB, NV, 4, 4)
        all_bboxes = data.get("bbox")  # (SB, NV, 4)  cmin rmin cmax rmax
        all_focals = data["focal"]  # (SB)
        all_c = data.get("c")  # (SB)

        if self.use_bbox and global_step >= args.no_bbox_step:
            self.use_bbox = False
            print(">>> Stopped using bbox sampling @ iter", global_step)

        if not is_train or not self.use_bbox:
            all_bboxes = None

        all_rgb_gt = []
        all_rays = []

        curr_nviews = nviews[torch.randint(0, len(nviews), ()).item()]
        if curr_nviews == 1:
            image_ord = torch.randint(0, NV, (SB, 1))
        else:
            image_ord = torch.empty((SB, curr_nviews), dtype=torch.long)
        for obj_idx in range(SB):
            if all_bboxes is not None:
                bboxes = all_bboxes[obj_idx]
            images = all_images[obj_idx]  # (NV, 3, H, W)
            poses = all_poses[obj_idx]  # (NV, 4, 4)
            focal = all_focals[obj_idx]
            c = None
            if "c" in data:
                c = data["c"][obj_idx]
            if curr_nviews > 1:
                # Somewhat inefficient, don't know better way
                image_ord[obj_idx] = torch.from_numpy(
                    np.random.choice(NV, curr_nviews, replace=False)
                )
            images_0to1 = images * 0.5 + 0.5

            cam_rays = util.gen_rays(
                poses, W, H, focal, self.z_near, self.z_far, c=c
            )  # (NV, H, W, 8)
            rgb_gt_all = images_0to1
            rgb_gt_all = (
                rgb_gt_all.permute(0, 2, 3, 1).contiguous().reshape(-1, 3)
            )  # (NV, H, W, 3)

            if all_bboxes is not None:
                pix = util.bbox_sample(bboxes, args.ray_batch_size)
                pix_inds = pix[..., 0] * H * W + pix[..., 1] * W + pix[..., 2]
            else:
                pix_inds = torch.randint(0, NV * H * W, (args.ray_batch_size,))

            rgb_gt = rgb_gt_all[pix_inds]  # (ray_batch_size, 3)
            rays = cam_rays.view(-1, cam_rays.shape[-1])[pix_inds].to(
                device=device
            )  # (ray_batch_size, 8)

            all_rgb_gt.append(rgb_gt)
            all_rays.append(rays)

        all_rgb_gt = torch.stack(all_rgb_gt)  # (SB, ray_batch_size, 3)
        all_rays = torch.stack(all_rays)  # (SB, ray_batch_size, 8)

        image_ord = image_ord.to(device)
        src_images = util.batched_index_select_nd(
            all_images, image_ord
        )  # (SB, NS, 3, H, W)
        src_poses = util.batched_index_select_nd(all_poses, image_ord)  # (SB, NS, 4, 4)

        all_bboxes = all_poses = all_images = None

        net.encode(
            src_images,
            src_poses,
            all_focals.to(device=device),
            c=all_c.to(device=device) if all_c is not None else None,
        )

        render_dict = DotMap(render_par(all_rays, want_weights=True,))
        coarse = render_dict.coarse
        fine = render_dict.fine
        using_fine = len(fine) > 0

        loss_dict = {}

        rgb_loss = self.rgb_coarse_crit(coarse.rgb, all_rgb_gt)
        if rgb_loss.isnan().any().item()==True:
          raise()
        loss_dict["rc"] = rgb_loss.item() * self.lambda_coarse
        if using_fine:
            fine_loss = self.rgb_fine_crit(fine.rgb, all_rgb_gt)
            rgb_loss = rgb_loss * self.lambda_coarse + fine_loss * self.lambda_fine
            loss_dict["rf"] = fine_loss.item() * self.lambda_fine

        loss = rgb_loss
        if is_train:
            loss.backward()
        loss_dict["t"] = loss.item()

        return loss_dict

    def train_step(self, data, global_step):
        return self.calc_losses(data, is_train=True, global_step=global_step)

    def eval_step(self, data, global_step):
        renderer.eval()
        losses = self.calc_losses(data, is_train=False, global_step=global_step)
        renderer.train()
        return losses

    def vis_step(self, data, global_step, idx=None):
        if "images" not in data:
            return {}
        if idx is None:
            batch_idx = np.random.randint(0, data["images"].shape[0])
        else:
            print(idx)
            batch_idx = idx
        images = data["images"][batch_idx].to(device=device)  # (NV, 3, H, W)
        poses = data["poses"][batch_idx].to(device=device)  # (NV, 4, 4)
        focal = data["focal"][batch_idx : batch_idx + 1]  # (1)
        c = data.get("c")
        if c is not None:
            c = c[batch_idx : batch_idx + 1]  # (1)
        NV, _, H, W = images.shape
        cam_rays = util.gen_rays(
            poses, W, H, focal, self.z_near, self.z_far, c=c
        )  # (NV, H, W, 8)
        images_0to1 = images * 0.5 + 0.5  # (NV, 3, H, W)

        curr_nviews = nviews[torch.randint(0, len(nviews), (1,)).item()]
        views_src = np.sort(np.random.choice(NV, curr_nviews, replace=False))
        view_dest = np.random.randint(0, NV - curr_nviews)
        for vs in range(curr_nviews):
            view_dest += view_dest >= views_src[vs]
        views_src = torch.from_numpy(views_src)

        # set renderer net to eval mode
        renderer.eval()
        source_views = (
            images_0to1[views_src]
            .permute(0, 2, 3, 1)
            .cpu()
            .numpy()
            .reshape(-1, H, W, 3)
        )

        gt = images_0to1[view_dest].permute(1, 2, 0).cpu().numpy().reshape(H, W, 3)
        with torch.no_grad():
            test_rays = cam_rays[view_dest]  # (H, W, 8)
            test_images = images[views_src]  # (NS, 3, H, W)
            net.encode(
                test_images.unsqueeze(0),
                poses[views_src].unsqueeze(0),
                focal.to(device=device),
                c=c.to(device=device) if c is not None else None,
            )
            test_rays = test_rays.reshape(1, H * W, -1)
            render_dict = DotMap(render_par(test_rays, want_weights=True))
            coarse = render_dict.coarse
            fine = render_dict.fine

            using_fine = len(fine) > 0

            alpha_coarse_np = coarse.weights[0].sum(dim=-1).cpu().numpy().reshape(H, W)
            rgb_coarse_np = coarse.rgb[0].cpu().numpy().reshape(H, W, 3)
            depth_coarse_np = coarse.depth[0].cpu().numpy().reshape(H, W)

            if using_fine:
                alpha_fine_np = fine.weights[0].sum(dim=1).cpu().numpy().reshape(H, W)
                depth_fine_np = fine.depth[0].cpu().numpy().reshape(H, W)
                rgb_fine_np = fine.rgb[0].cpu().numpy().reshape(H, W, 3)

        print("c rgb min {} max {}".format(rgb_coarse_np.min(), rgb_coarse_np.max()))
        print(
            "c alpha min {}, max {}".format(
                alpha_coarse_np.min(), alpha_coarse_np.max()
            )
        )
        alpha_coarse_cmap = util.cmap(alpha_coarse_np) / 255
        depth_coarse_cmap = util.cmap(depth_coarse_np) / 255
        vis_list = [
            *source_views,
            gt,
            depth_coarse_cmap,
            rgb_coarse_np,
            alpha_coarse_cmap,
        ]

        vis_coarse = np.hstack(vis_list)
        vis = vis_coarse

        if using_fine:
            print("f rgb min {} max {}".format(rgb_fine_np.min(), rgb_fine_np.max()))
            print(
                "f alpha min {}, max {}".format(
                    alpha_fine_np.min(), alpha_fine_np.max()
                )
            )
            depth_fine_cmap = util.cmap(depth_fine_np) / 255
            alpha_fine_cmap = util.cmap(alpha_fine_np) / 255
            vis_list = [
                *source_views,
                gt,
                depth_fine_cmap,
                rgb_fine_np,
                alpha_fine_cmap,
            ]

            vis_fine = np.hstack(vis_list)
            vis = np.vstack((vis_coarse, vis_fine))
            rgb_psnr = rgb_fine_np
        else:
            rgb_psnr = rgb_coarse_np

        psnr = util.psnr(rgb_psnr, gt)
        vals = {"psnr": psnr}
        print("psnr", psnr)

        # set the renderer network back to train mode
        renderer.train()
        return vis, vals


trainer = PixelNeRFTrainer()
trainer.start()

[((56, 13, 56, 88), 'util.args.parse_args', 'util.args.parse_args', (), '', False, 'import util\n'), ((57, 9, 57, 38), 'util.get_cuda', 'util.get_cuda', ({(57, 23, 57, 37): 'args.gpu_id[0]'}, {}), '(args.gpu_id[0])', False, 'import util\n'), ((59, 20, 59, 72), 'data.get_split_dataset', 'get_split_dataset', ({(59, 38, 59, 57): 'args.dataset_format', (59, 59, 59, 71): 'args.datadir'}, {}), '(args.dataset_format, args.datadir)', False, 'from data import get_split_dataset\n'), ((64, 6, 64, 31), 'model.make_model', 'make_model', ({(64, 17, 64, 30): "conf['model']"}, {}), "(conf['model'])", False, 'from model import make_model, loss\n'), ((70, 11, 70, 74), 'render.NeRFRenderer.from_conf', 'NeRFRenderer.from_conf', (), '', False, 'from render import NeRFRenderer\n'), ((93, 31, 93, 72), 'model.loss.get_rgb_loss', 'loss.get_rgb_loss', ({(93, 49, 93, 65): "conf['loss.rgb']", (93, 67, 93, 71): 'True'}, {}), "(conf['loss.rgb'], True)", False, 'from model import make_model, loss\n'), ((98, 29, 98, 69), 'model.loss.get_rgb_loss', 'loss.get_rgb_loss', ({(98, 47, 98, 61): 'fine_loss_conf', (98, 63, 98, 68): 'False'}, {}), '(fine_loss_conf, False)', False, 'from model import make_model, loss\n'), ((181, 21, 181, 44), 'torch.stack', 'torch.stack', ({(181, 33, 181, 43): 'all_rgb_gt'}, {}), '(all_rgb_gt)', False, 'import torch\n'), ((182, 19, 182, 40), 'torch.stack', 'torch.stack', ({(182, 31, 182, 39): 'all_rays'}, {}), '(all_rays)', False, 'import torch\n'), ((185, 21, 187, 9), 'util.batched_index_select_nd', 'util.batched_index_select_nd', ({(186, 12, 186, 22): 'all_images', (186, 24, 186, 33): 'image_ord'}, {}), '(all_images, image_ord)', False, 'import util\n'), ((188, 20, 188, 70), 'util.batched_index_select_nd', 'util.batched_index_select_nd', ({(188, 49, 188, 58): 'all_poses', (188, 60, 188, 69): 'image_ord'}, {}), '(all_poses, image_ord)', False, 'import util\n'), ((218, 25, 218, 36), 'model.loss.item', 'loss.item', ({}, {}), '()', False, 'from model import make_model, loss\n'), ((246, 19, 248, 9), 'util.gen_rays', 'util.gen_rays', (), '', False, 'import util\n'), ((253, 20, 253, 58), 'numpy.random.randint', 'np.random.randint', ({(253, 38, 253, 39): '0', (253, 41, 253, 57): 'NV - curr_nviews'}, {}), '(0, NV - curr_nviews)', True, 'import numpy as np\n'), ((256, 20, 256, 47), 'torch.from_numpy', 'torch.from_numpy', ({(256, 37, 256, 46): 'views_src'}, {}), '(views_src)', False, 'import torch\n'), ((310, 21, 310, 40), 'numpy.hstack', 'np.hstack', ({(310, 31, 310, 39): 'vis_list'}, {}), '(vis_list)', True, 'import numpy as np\n'), ((336, 15, 336, 38), 'util.psnr', 'util.psnr', ({(336, 25, 336, 33): 'rgb_psnr', (336, 35, 336, 37): 'gt'}, {}), '(rgb_psnr, gt)', False, 'import util\n'), ((8, 36, 8, 61), 'os.path.dirname', 'os.path.dirname', ({(8, 52, 8, 60): '__file__'}, {}), '(__file__)', False, 'import os\n'), ((101, 15, 101, 55), 'os.path.exists', 'os.path.exists', ({(101, 30, 101, 54): 'self.renderer_state_path'}, {}), '(self.renderer_state_path)', False, 'import os\n'), ((140, 24, 140, 53), 'torch.randint', 'torch.randint', ({(140, 38, 140, 39): '0', (140, 41, 140, 43): 'NV', (140, 45, 140, 52): '(SB, 1)'}, {}), '(0, NV, (SB, 1))', False, 'import torch\n'), ((142, 24, 142, 72), 'torch.empty', 'torch.empty', (), '', False, 'import torch\n'), ((159, 23, 161, 13), 'util.gen_rays', 'util.gen_rays', (), '', False, 'import util\n'), ((217, 12, 217, 27), 'model.loss.backward', 'loss.backward', ({}, {}), '()', False, 'from model import make_model, loss\n'), ((235, 24, 235, 69), 'numpy.random.randint', 'np.random.randint', ({(235, 42, 235, 43): '0', (235, 45, 235, 68): "data['images'].shape[0]"}, {}), "(0, data['images'].shape[0])", True, 'import numpy as np\n'), ((252, 28, 252, 76), 'numpy.random.choice', 'np.random.choice', (), '', True, 'import numpy as np\n'), ((269, 13, 269, 28), 'torch.no_grad', 'torch.no_grad', ({}, {}), '()', False, 'import torch\n'), ((300, 28, 300, 54), 'util.cmap', 'util.cmap', ({(300, 38, 300, 53): 'alpha_coarse_np'}, {}), '(alpha_coarse_np)', False, 'import util\n'), ((301, 28, 301, 54), 'util.cmap', 'util.cmap', ({(301, 38, 301, 53): 'depth_coarse_np'}, {}), '(depth_coarse_np)', False, 'import util\n'), ((330, 23, 330, 42), 'numpy.hstack', 'np.hstack', ({(330, 33, 330, 41): 'vis_list'}, {}), '(vis_list)', True, 'import numpy as np\n'), ((331, 18, 331, 51), 'numpy.vstack', 'np.vstack', ({(331, 28, 331, 50): '(vis_coarse, vis_fine)'}, {}), '((vis_coarse, vis_fine))', True, 'import numpy as np\n'), ((168, 22, 168, 67), 'util.bbox_sample', 'util.bbox_sample', ({(168, 39, 168, 45): 'bboxes', (168, 47, 168, 66): 'args.ray_batch_size'}, {}), '(bboxes, args.ray_batch_size)', False, 'import util\n'), ((171, 27, 171, 79), 'torch.randint', 'torch.randint', ({(171, 41, 171, 42): '0', (171, 44, 171, 54): 'NV * H * W', (171, 56, 171, 78): '(args.ray_batch_size,)'}, {}), '(0, NV * H * W, (args.ray_batch_size,))', False, 'import torch\n'), ((320, 30, 320, 54), 'util.cmap', 'util.cmap', ({(320, 40, 320, 53): 'depth_fine_np'}, {}), '(depth_fine_np)', False, 'import util\n'), ((321, 30, 321, 54), 'util.cmap', 'util.cmap', ({(321, 40, 321, 53): 'alpha_fine_np'}, {}), '(alpha_fine_np)', False, 'import util\n'), ((103, 20, 103, 77), 'torch.load', 'torch.load', (), '', False, 'import torch\n'), ((155, 20, 155, 68), 'numpy.random.choice', 'np.random.choice', (), '', True, 'import numpy as np\n')]