#!/usr/bin/env python
import argparse
import os.path as osp
import logging
import time
import sys
sys.path.insert(0, osp.dirname(__file__) + '/..')

import torch
import torch.nn as nn

from fastmvsnet.config import load_cfg_from_file
from fastmvsnet.utils.io import mkdir
from fastmvsnet.utils.logger import setup_logger
from fastmvsnet.utils.torch_utils import set_random_seed
from fastmvsnet.model1 import build_pointmvsnet as build_model
from fastmvsnet.solver import build_optimizer, build_scheduler
from fastmvsnet.utils.checkpoint import Checkpointer
from fastmvsnet.dataset1 import build_data_loader
from fastmvsnet.utils.tensorboard_logger import TensorboardLogger
from fastmvsnet.utils.metric_logger import MetricLogger
from fastmvsnet.utils.file_logger import file_logger


def parse_args():
    parser = argparse.ArgumentParser(description="PyTorch Fast-MVSNet Training")
    parser.add_argument(
        "--cfg",
        dest="config_file",
        default="",
        metavar="FILE",
        help="path to config file",
        type=str,
    )
    parser.add_argument(
        "opts",
        help="Modify config options using the command-line",
        default=None,
        nargs=argparse.REMAINDER,
    )

    args = parser.parse_args()
    return args


def train_model(model,
                loss_fn,
                metric_fn,
                image_scales,
                inter_scales,
                isFlow,
                data_loader,
                optimizer,
                curr_epoch,
                tensorboard_logger,
                log_period=1,
                output_dir="",
                ):
    logger = logging.getLogger("fastmvsnet.train")
    meters = MetricLogger(delimiter="  ")
    model.train()
    end = time.time()
    total_iteration = data_loader.__len__()
    path_list = []

    for iteration, data_batch in enumerate(data_loader):

        data_time = time.time() - end
        curr_ref_img_path = data_batch["ref_img_path"]
        path_list.extend(curr_ref_img_path)
        data_batch = {k: v.cuda(non_blocking=True) for k, v in data_batch.items() if isinstance(v, torch.Tensor)}

        preds = model(data_batch, image_scales, inter_scales, isFlow)
        optimizer.zero_grad()

        loss_dict = loss_fn(preds, data_batch, isFlow)
        metric_dict = metric_fn(preds, data_batch, isFlow)
        losses = sum(loss_dict.values())
        #print("LOSS DICT", loss_dict['coarse_loss'])
        #print("LOSSES", loss_dict.values())
        meters.update(loss=losses, **loss_dict, **metric_dict)

        losses.backward()
       # print(poop)
        optimizer.step()

        batch_time = time.time() - end
        end = time.time()
        meters.update(time=batch_time, data=data_time)

        if iteration % log_period == 0:
            logger.info(
                meters.delimiter.join(
                    [
                        "EPOCH: {epoch:2d}",
                        "iter: {iter:4d}",
                        "{meters}",
                        "lr: {lr:.2e}",
                        "max mem: {memory:.0f}",
                    ]
                ).format(
                    epoch=curr_epoch,
                    iter=iteration,
                    meters=str(meters),
                    lr=optimizer.param_groups[0]["lr"],
                    memory=torch.cuda.max_memory_allocated() / (1024.0 ** 2),
                )
            )
            tensorboard_logger.add_scalars(loss_dict, curr_epoch * total_iteration + iteration, prefix="train")
            tensorboard_logger.add_scalars(metric_dict, curr_epoch * total_iteration + iteration, prefix="train")

        if iteration % (100 * log_period) == 0:
            file_logger(data_batch, preds, curr_epoch * total_iteration + iteration, output_dir, prefix="train")

    return meters


def validate_model(model,
                   loss_fn,
                   metric_fn,
                   image_scales,
                   inter_scales,
                   isFlow,
                   data_loader,
                   curr_epoch,
                   tensorboard_logger,
                   log_period=1,
                   output_dir="",
                   ):
    logger = logging.getLogger("fastmvsnet.validate")
    meters = MetricLogger(delimiter="  ")
    model.train()
    end = time.time()
    total_iteration = data_loader.__len__()
    with torch.no_grad():
        for iteration, data_batch in enumerate(data_loader):
            data_time = time.time() - end
            curr_ref_img_path = data_batch["ref_img_path"]

            data_batch = {k: v.cuda(non_blocking=True) for k, v in data_batch.items() if isinstance(v, torch.Tensor)}

            preds = model(data_batch, image_scales, inter_scales, isFlow)
            loss_dict = loss_fn(preds, data_batch, isFlow)
            metric_dict = metric_fn(preds, data_batch, isFlow)
            losses = sum(loss_dict.values())
            meters.update(loss=losses, **loss_dict, **metric_dict)
            batch_time = time.time() - end
            end = time.time()
            meters.update(time=batch_time, data=data_time)

            if iteration % log_period == 0:
                logger.info(
                    meters.delimiter.join(
                        [
                            "EPOCH: {epoch:2d}",
                            "iter: {iter:4d}",
                            "{meters}",
                        ]
                    ).format(
                        epoch=curr_epoch,
                        iter=iteration,
                        meters=str(meters),
                    )
                )
                tensorboard_logger.add_scalars(meters.meters, curr_epoch * total_iteration + iteration, prefix="valid")

            if iteration % (100 * log_period) == 0:
                file_logger(data_batch, preds, curr_epoch * total_iteration + iteration, output_dir, prefix="valid")

    return meters


def train(cfg, output_dir=""):
    logger = logging.getLogger("fastmvsnet.trainer")

    # build model
    set_random_seed(cfg.RNG_SEED)
    model, loss_fn, metric_fn = build_model(cfg)
    logger.info("Build model:\n{}".format(str(model)))
    model = nn.DataParallel(model).cuda()

    # build optimizer
    optimizer = build_optimizer(cfg, model)

    # build lr scheduler
    scheduler = build_scheduler(cfg, optimizer)

    # build checkpointer
    checkpointer = Checkpointer(model,
                                optimizer=optimizer,
                                scheduler=scheduler,
                                save_dir=output_dir,
                                logger=logger)

    checkpoint_data = checkpointer.load(cfg.MODEL.WEIGHT, resume=cfg.AUTO_RESUME)
    ckpt_period = cfg.TRAIN.CHECKPOINT_PERIOD

    # build data loader
    train_data_loader = build_data_loader(cfg, mode="train")
    val_period = cfg.TRAIN.VAL_PERIOD
    val_data_loader = build_data_loader(cfg, mode="val") if val_period > 0 else None

    # build tensorboard logger (optionally by comment)
    tensorboard_logger = TensorboardLogger(output_dir)

    # train
    max_epoch = cfg.SCHEDULER.MAX_EPOCH
    start_epoch = checkpoint_data.get("epoch", 0)
    best_metric_name = "best_{}".format(cfg.TRAIN.VAL_METRIC)
    best_metric = checkpoint_data.get(best_metric_name, None)
    logger.info("Start training from epoch {}".format(start_epoch))
    for epoch in range(start_epoch, max_epoch):
        cur_epoch = epoch + 1
        scheduler.step()
        start_time = time.time()
        train_meters = train_model(model,
                                   loss_fn,
                                   metric_fn,
                                   image_scales=cfg.MODEL.TRAIN.IMG_SCALES,
                                   inter_scales=cfg.MODEL.TRAIN.INTER_SCALES,
                                   isFlow=(cur_epoch > cfg.SCHEDULER.INIT_EPOCH),
                                   data_loader=train_data_loader,
                                   optimizer=optimizer,
                                   curr_epoch=epoch,
                                   tensorboard_logger=tensorboard_logger,
                                   log_period=cfg.TRAIN.LOG_PERIOD,
                                   output_dir=output_dir,
                                   )
        epoch_time = time.time() - start_time
        logger.info("Epoch[{}]-Train {}  total_time: {:.2f}s".format(
            cur_epoch, train_meters.summary_str, epoch_time))

        # checkpoint
        if cur_epoch % ckpt_period == 0 or cur_epoch == max_epoch:
            checkpoint_data["epoch"] = cur_epoch
            checkpoint_data[best_metric_name] = best_metric
            checkpointer.save("model_{:03d}".format(cur_epoch), **checkpoint_data)

        # validate
        if val_period < 1:
            continue
        if cur_epoch % val_period == 0 or cur_epoch == max_epoch:
            val_meters = validate_model(model,
                                        loss_fn,
                                        metric_fn,
                                        image_scales=cfg.MODEL.VAL.IMG_SCALES,
                                        inter_scales=cfg.MODEL.VAL.INTER_SCALES,
                                        isFlow=(cur_epoch > cfg.SCHEDULER.INIT_EPOCH),
                                        data_loader=val_data_loader,
                                        curr_epoch=epoch,
                                        tensorboard_logger=tensorboard_logger,
                                        log_period=cfg.TEST.LOG_PERIOD,
                                        output_dir=output_dir,
                                        )
            logger.info("Epoch[{}]-Val {}".format(cur_epoch, val_meters.summary_str))

            # best validation
            cur_metric = val_meters.meters[cfg.TRAIN.VAL_METRIC].global_avg
            if best_metric is None or cur_metric > best_metric:
                best_metric = cur_metric
                checkpoint_data["epoch"] = cur_epoch
                checkpoint_data[best_metric_name] = best_metric
                checkpointer.save("model_best", **checkpoint_data)

    logger.info("Best val-{} = {}".format(cfg.TRAIN.VAL_METRIC, best_metric))

    return model


def main():
    args = parse_args()
    num_gpus = torch.cuda.device_count()

    cfg = load_cfg_from_file(args.config_file)
    cfg.merge_from_list(args.opts)
    cfg.freeze()

    output_dir = cfg.OUTPUT_DIR
    if output_dir:
        config_path = osp.splitext(args.config_file)[0]
        config_path = config_path.replace("configs", "outputs1")
        output_dir = output_dir.replace('@', config_path)
        mkdir(output_dir)

    logger = setup_logger("fastmvsnet", output_dir, prefix="train")
    logger.info("Using {} GPUs".format(num_gpus))
    logger.info(args)

    logger.info("Loaded configuration file {}".format(args.config_file))
    logger.info("Running with config:\n{}".format(cfg))

    train(cfg, output_dir)


if __name__ == "__main__":
    main()

[((26, 13, 26, 80), 'argparse.ArgumentParser', 'argparse.ArgumentParser', (), '', False, 'import argparse\n'), ((59, 13, 59, 50), 'logging.getLogger', 'logging.getLogger', ({(59, 31, 59, 49): '"""fastmvsnet.train"""'}, {}), "('fastmvsnet.train')", False, 'import logging\n'), ((60, 13, 60, 41), 'fastmvsnet.utils.metric_logger.MetricLogger', 'MetricLogger', (), '', False, 'from fastmvsnet.utils.metric_logger import MetricLogger\n'), ((62, 10, 62, 21), 'time.time', 'time.time', ({}, {}), '()', False, 'import time\n'), ((130, 13, 130, 53), 'logging.getLogger', 'logging.getLogger', ({(130, 31, 130, 52): '"""fastmvsnet.validate"""'}, {}), "('fastmvsnet.validate')", False, 'import logging\n'), ((131, 13, 131, 41), 'fastmvsnet.utils.metric_logger.MetricLogger', 'MetricLogger', (), '', False, 'from fastmvsnet.utils.metric_logger import MetricLogger\n'), ((133, 10, 133, 21), 'time.time', 'time.time', ({}, {}), '()', False, 'import time\n'), ((174, 13, 174, 52), 'logging.getLogger', 'logging.getLogger', ({(174, 31, 174, 51): '"""fastmvsnet.trainer"""'}, {}), "('fastmvsnet.trainer')", False, 'import logging\n'), ((177, 4, 177, 33), 'fastmvsnet.utils.torch_utils.set_random_seed', 'set_random_seed', ({(177, 20, 177, 32): 'cfg.RNG_SEED'}, {}), '(cfg.RNG_SEED)', False, 'from fastmvsnet.utils.torch_utils import set_random_seed\n'), ((178, 32, 178, 48), 'fastmvsnet.model1.build_pointmvsnet', 'build_model', ({(178, 44, 178, 47): 'cfg'}, {}), '(cfg)', True, 'from fastmvsnet.model1 import build_pointmvsnet as build_model\n'), ((183, 16, 183, 43), 'fastmvsnet.solver.build_optimizer', 'build_optimizer', ({(183, 32, 183, 35): 'cfg', (183, 37, 183, 42): 'model'}, {}), '(cfg, model)', False, 'from fastmvsnet.solver import build_optimizer, build_scheduler\n'), ((186, 16, 186, 47), 'fastmvsnet.solver.build_scheduler', 'build_scheduler', ({(186, 32, 186, 35): 'cfg', (186, 37, 186, 46): 'optimizer'}, {}), '(cfg, optimizer)', False, 'from fastmvsnet.solver import build_optimizer, build_scheduler\n'), ((189, 19, 193, 46), 'fastmvsnet.utils.checkpoint.Checkpointer', 'Checkpointer', (), '', False, 'from fastmvsnet.utils.checkpoint import Checkpointer\n'), ((199, 24, 199, 60), 'fastmvsnet.dataset1.build_data_loader', 'build_data_loader', (), '', False, 'from fastmvsnet.dataset1 import build_data_loader\n'), ((204, 25, 204, 54), 'fastmvsnet.utils.tensorboard_logger.TensorboardLogger', 'TensorboardLogger', ({(204, 43, 204, 53): 'output_dir'}, {}), '(output_dir)', False, 'from fastmvsnet.utils.tensorboard_logger import TensorboardLogger\n'), ((272, 15, 272, 40), 'torch.cuda.device_count', 'torch.cuda.device_count', ({}, {}), '()', False, 'import torch\n'), ((274, 10, 274, 46), 'fastmvsnet.config.load_cfg_from_file', 'load_cfg_from_file', ({(274, 29, 274, 45): 'args.config_file'}, {}), '(args.config_file)', False, 'from fastmvsnet.config import load_cfg_from_file\n'), ((285, 13, 285, 67), 'fastmvsnet.utils.logger.setup_logger', 'setup_logger', (), '', False, 'from fastmvsnet.utils.logger import setup_logger\n'), ((7, 19, 7, 40), 'os.path.dirname', 'osp.dirname', ({(7, 31, 7, 39): '__file__'}, {}), '(__file__)', True, 'import os.path as osp\n'), ((88, 14, 88, 25), 'time.time', 'time.time', ({}, {}), '()', False, 'import time\n'), ((135, 9, 135, 24), 'torch.no_grad', 'torch.no_grad', ({}, {}), '()', False, 'import torch\n'), ((201, 22, 201, 56), 'fastmvsnet.dataset1.build_data_loader', 'build_data_loader', (), '', False, 'from fastmvsnet.dataset1 import build_data_loader\n'), ((215, 21, 215, 32), 'time.time', 'time.time', ({}, {}), '()', False, 'import time\n'), ((283, 8, 283, 25), 'fastmvsnet.utils.io.mkdir', 'mkdir', ({(283, 14, 283, 24): 'output_dir'}, {}), '(output_dir)', False, 'from fastmvsnet.utils.io import mkdir\n'), ((68, 20, 68, 31), 'time.time', 'time.time', ({}, {}), '()', False, 'import time\n'), ((87, 21, 87, 32), 'time.time', 'time.time', ({}, {}), '()', False, 'import time\n'), ((113, 12, 113, 112), 'fastmvsnet.utils.file_logger.file_logger', 'file_logger', (), '', False, 'from fastmvsnet.utils.file_logger import file_logger\n'), ((148, 18, 148, 29), 'time.time', 'time.time', ({}, {}), '()', False, 'import time\n'), ((180, 12, 180, 34), 'torch.nn.DataParallel', 'nn.DataParallel', ({(180, 28, 180, 33): 'model'}, {}), '(model)', True, 'import torch.nn as nn\n'), ((229, 21, 229, 32), 'time.time', 'time.time', ({}, {}), '()', False, 'import time\n'), ((280, 22, 280, 52), 'os.path.splitext', 'osp.splitext', ({(280, 35, 280, 51): 'args.config_file'}, {}), '(args.config_file)', True, 'import os.path as osp\n'), ((137, 24, 137, 35), 'time.time', 'time.time', ({}, {}), '()', False, 'import time\n'), ((147, 25, 147, 36), 'time.time', 'time.time', ({}, {}), '()', False, 'import time\n'), ((168, 16, 168, 116), 'fastmvsnet.utils.file_logger.file_logger', 'file_logger', (), '', False, 'from fastmvsnet.utils.file_logger import file_logger\n'), ((106, 27, 106, 60), 'torch.cuda.max_memory_allocated', 'torch.cuda.max_memory_allocated', ({}, {}), '()', False, 'import torch\n')]