from pathlib import Path
import multiprocessing as mp

mp.set_start_method("spawn", force=True)

import hydra
from omegaconf import OmegaConf
from accelerate.state import PartialState

from accel_hydra import Trainer
from accel_hydra.models import CountParamsBase
from accel_hydra.utils import (
    register_omegaconf_resolvers,
    setup_resume_cfg,
    init_dataloader_from_config,
)
from accel_hydra.utils.lr_scheduler import (
    get_warmup_steps,
    get_dataloader_one_pass_outside_steps,
    get_total_training_steps,
    get_steps_inside_accelerator_from_outside_steps,
    get_dataloader_one_pass_steps_inside_accelerator,
    lr_scheduler_param_adapter,
)

register_omegaconf_resolvers()


def main():

    configs = []

    @hydra.main(version_base=None, config_path="configs", config_name="train")
    def parse_config_from_command_line(config):
        config = OmegaConf.to_container(config, resolve=True)
        configs.append(config)

    parse_config_from_command_line()
    config = configs[0]

    # helper state for accessing information about the current training environment
    state = PartialState()

    if config.get("dump_config", None) is not None:
        if state.is_main_process:
            with open(config["dump_config"], "w") as f:
                OmegaConf.save(config, f)
                print(f'config.yaml saved to {f.name}')
        return

    config = setup_resume_cfg(config, do_print=state.is_main_process)

    model: CountParamsBase = hydra.utils.instantiate(
        config["model"], _convert_="all"
    )
    train_dataloader = init_dataloader_from_config(config["train_dataloader"])
    if "val_dataloader" in config and config["val_dataloader"] is not None:
        val_dataloader = init_dataloader_from_config(config["val_dataloader"])
    else:
        val_dataloader = None
    optimizer = hydra.utils.instantiate(
        config["optimizer"], params=model.parameters(), _convert_="all"
    )

    # `accelerator.prepare` is very confusing for multi-gpu, gradient accumulation scenario:
    # For more information: see https://github.com/huggingface/diffusers/issues/4387,
    # https://github.com/huggingface/diffusers/issues/9633, and
    # https://github.com/huggingface/diffusers/issues/3954
    dataloader_one_pass_outside_steps = get_dataloader_one_pass_outside_steps(
        train_dataloader, state.num_processes
    )
    total_training_steps = get_total_training_steps(
        train_dataloader, config["epochs"], state.num_processes,
        config["epoch_length"]
    )
    dataloader_one_pass_steps_inside_accelerator = (
        get_dataloader_one_pass_steps_inside_accelerator(
            dataloader_one_pass_outside_steps,
            config["gradient_accumulation_steps"], state.num_processes
        )
    )
    num_training_updates = get_steps_inside_accelerator_from_outside_steps(
        total_training_steps, dataloader_one_pass_outside_steps,
        dataloader_one_pass_steps_inside_accelerator,
        config["gradient_accumulation_steps"], state.num_processes
    )

    if "warmup_params" in config:
        num_warmup_steps = get_warmup_steps(
            **config["warmup_params"],
            dataloader_one_pass_outside_steps=dataloader_one_pass_outside_steps
        )
        num_warmup_updates = get_steps_inside_accelerator_from_outside_steps(
            num_warmup_steps, dataloader_one_pass_outside_steps,
            dataloader_one_pass_steps_inside_accelerator,
            config["gradient_accumulation_steps"], state.num_processes
        )
    else:
        num_warmup_updates = None

    lr_scheduler_config = lr_scheduler_param_adapter(
        config_dict=config["lr_scheduler"],
        num_training_steps=num_training_updates,
        num_warmup_steps=num_warmup_updates,
    )

    lr_scheduler = hydra.utils.instantiate(
        lr_scheduler_config, optimizer=optimizer, _convert_="all"
    )
    loss_fn = hydra.utils.instantiate(config["loss_fn"], _convert_="all")
    trainer: Trainer = hydra.utils.instantiate(
        config["trainer"],
        train_dataloader=train_dataloader,
        val_dataloader=val_dataloader,
        model=model,
        optimizer=optimizer,
        lr_scheduler=lr_scheduler,
        loss_fn=loss_fn,
        _convert_="all"
    )
    trainer.config_dict = config  # assign here, don't instantiate it
    trainer.train(seed=config["seed"])


if __name__ == "__main__":
    main()
