# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# 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.

import argparse
import importlib
import os
import sys
from datetime import datetime

import torch
import torch.distributed
from loguru import logger as logging
from rich import print

from imaginaire.config import Config, pretty_print_overrides
from imaginaire.lazy_config import instantiate
from imaginaire.lazy_config.lazy import LazyConfig
from imaginaire.utils import distributed, log, misc
from imaginaire.utils.config_helper import get_config_module, override

####################################################
# NOTE (Dian): patch for SageMaker parsing hydra args, because SM always passes in args
# like --config1 <val1> --config2 <val2>, which is not compatible with hydra
# NOTE (bvh): SageMaker sorts hyperparameters alphabetically, so we must extract
# --config/--dryrun from anywhere in argv (not assume they come first).
####################################################

if "--" not in sys.argv:
    # Extract --config and --dryrun; everything else becomes hydra overrides after "--"
    pre_separator = [sys.argv[0]]
    post_separator = []
    i = 1
    while i < len(sys.argv):
        arg = sys.argv[i]
        if arg.startswith("--config"):
            if "=" in arg:
                pre_separator.append(arg)
                i += 1
            else:
                pre_separator.extend([arg, sys.argv[i + 1]])
                i += 2
        elif arg.startswith("--dryrun"):
            pre_separator.append(arg)
            i += 1
        else:
            post_separator.append(arg)
            i += 1
    sys.argv = pre_separator + ["--"] + post_separator

    # for overrides after "--", strip the "--" added by SageMaker and add "=" between hydra args
    # e.g. --config cosmos_predict2/configs/base/config.py -- --experiment <exp> becomes
    # --config cosmos_predict2/configs/base/config.py -- experiment=<exp>
    # NOTE: only reformat if args look like SageMaker style (start with --)
    hydra_idx = sys.argv.index("--") + 1
    remaining = sys.argv[hydra_idx:]
    has_sagemaker_style = any(arg.startswith("--") for arg in remaining)
    if has_sagemaker_style:
        new_argv = sys.argv[:hydra_idx]
        for i, arg in enumerate(remaining):
            if arg.startswith("--"):
                new_argv.append(arg.lstrip("-"))
            else:
                new_argv[-1] = new_argv[-1] + "=" + arg
        sys.argv = new_argv


@logging.catch(reraise=True)
def launch(config: Config, args: argparse.Namespace) -> None:
    # Need to initialize the distributed environment before calling config.validate() because it tries to synchronize
    # a buffer across ranks. If you don't do this, then you end up allocating a bunch of buffers on rank 0, and also that
    # check doesn't actually do anything.
    print(f'[italic][magenta]{datetime.now()} | scripts/train.py:launch() begin')

    # NOTE(bvh): For debugging, enable through TORCH_ANOMALY flag:
    if os.environ.get('TORCH_ANOMALY', False):
        torch.autograd.set_detect_anomaly(True)
        logging.warning(f'WARNING: TORCH_ANOMALY = 1 => torch.autograd.set_detect_anomaly(True)')
        print(f'[bold][orange3]{datetime.now()} | WARNING: TORCH_ANOMALY = 1 => torch.autograd.set_detect_anomaly(True)')

    distributed.init()

    # Check that the config is valid
    config.validate()

    # NOTE(bvh): checkpoint.resume restores the full model state, which takes precedence over the
    # dit_path warm-start; unset dit_path to skip the redundant download + load + surgery.
    if config.checkpoint.resume and config.model.config.model_manager_config.dit_path:
        log.warning(f'checkpoint.resume is set; ignoring dit_path '
                    f'({config.model.config.model_manager_config.dit_path})')
        config.model.config.model_manager_config.dit_path = ''

    # Freeze the config so developers don't change it during training.
    config.freeze()  # type: ignore
    trainer = config.trainer.type(config)
    
    # Create the model
    with misc.timer('instantiate model'):
        model = instantiate(config.model)

    # Push tar-retry knobs from a4d_config into env vars so the webdataset
    # readers in DataLoader worker subprocesses pick them up via lazy lookup.
    # Must happen BEFORE dataloader instantiation so forked/spawned workers
    # inherit the env. Set tar_retry_attempts=1 to disable retry,
    # tar_thread_join_timeout<=0 to disable the join() safety net.
    cfg = getattr(model, 'config', None)
    if cfg is not None:
        attempts = getattr(cfg, 'tar_retry_attempts', None)
        if attempts is not None:
            os.environ['ANYDATA_TAR_RETRY_ATTEMPTS'] = str(int(attempts))
        join_to = getattr(cfg, 'tar_thread_join_timeout', None)
        if join_to is not None:
            os.environ['ANYDATA_TAR_THREAD_JOIN_TIMEOUT'] = str(float(join_to))

    # Create the dataloaders.
    with misc.timer('instantiate dataloader_train'):
        dataloader_train = instantiate(config.dataloader_train)
    with misc.timer('instantiate dataloader_val'):
        dataloader_val = instantiate(config.dataloader_val)

    # Use same name for experiment and run
    # if model.recipe is not None:
    #     model.recipe.model.config.experiment_tag = config.job.name
    
    # Start training
    trainer.train(
        model,
        dataloader_train,
        dataloader_val,
    )

    # Clean shutdown to avoid NCCL warnings.
    if torch.distributed.is_initialized():
        torch.distributed.destroy_process_group()


if __name__ == "__main__":
    # Usage: torchrun --nproc_per_node=1 -m scripts.train --config=cosmos_predict2/configs/base/config.py -- experiments=predict2_video2world_training_2b_cosmos_nemo_assets

    print(f'[italic][magenta]{datetime.now()} | scripts/train.py:__main__ begin')
    
    # Get the config file from the input arguments.
    parser = argparse.ArgumentParser(description="Training")
    parser.add_argument("--config", help="Path to the config file",
                        default='cosmos_predict2/configs/base/config.py')
    parser.add_argument(
        "opts",
        help="""
Modify config options at the end of the command. For Yacs configs, use
space-separated "PATH.KEY VALUE" pairs.
For python-based LazyConfig, use "path.key=value".
        """.strip(),
        default=None,
        nargs=argparse.REMAINDER,
    )
    parser.add_argument(
        "--dryrun",
        action="store_true",
        help="Do a dry run without training. Useful for debugging the config.",
    )
    args = parser.parse_args()
    config_module = get_config_module(args.config)
    config = importlib.import_module(config_module).make_config()
    config = override(config, args.opts)
    
    if args.dryrun:
        logging.info(
            "Config:\n" + config.pretty_print(use_color=True) + "\n" + pretty_print_overrides(args.opts, use_color=True)
        )
        os.makedirs(config.job.path_local, exist_ok=True)
        LazyConfig.save_yaml(config, f"{config.job.path_local}/config.yaml")
        print(f"{config.job.path_local}/config.yaml")
    
    else:
        # Launch the training job.
        launch(config, args)