# Copyright (c) Microsoft Corporation.
# Licensed under the MIT licen

import os
from dataclasses import asdict, dataclass, field
from typing import Any, Dict, Optional, Tuple

import numpy as np
import torch

from archai.common.distributed_utils import (
    get_world_size,
    init_distributed,
    sync_workers,
)
from archai.common.file_utils import get_full_path


def _get_amlt_dirs() -> Tuple[str, str]:
    data_dir = os.environ.get("AMLT_DATA_DIR", "")
    output_dir = os.environ.get("AMLT_OUTPUT_DIR", "")

    return data_dir, output_dir


def _get_default_dataroot() -> str:
    is_amlt_available = os.environ.get("AMLT_OUTPUT_DIR", None)

    return "/var/tmp/dataroot" if is_amlt_available else "~/dataroot"


def _create_dirs(
    dataroot: str,
    dataset_name: str,
    experiment_name: Optional[str] = "tmp",
    output_dir: Optional[str] = "~/logdir",
    pretrained_path: Optional[str] = "",
    cache_dir: Optional[str] = "",
) -> Tuple[str, str, str, str]:
    def _get_dataset_dir_name(dataset_name: str) -> str:
        if dataset_name == "wt2":
            return "wikitext-2"
        if dataset_name == "wt103":
            return "wikitext-103"
        if dataset_name == "lm1b":
            return "one-billion-words"
        if dataset_name.startswith("olx_"):
            return dataset_name

        raise RuntimeError(f"Dataset: {dataset_name} is not supported yet.")

    pt_data_dir, pt_output_dir = _get_amlt_dirs()
    if pt_output_dir:
        pt_output_dir = os.path.join(pt_output_dir, experiment_name)

    dataroot = dataroot or pt_data_dir or _get_default_dataroot()
    dataroot = get_full_path(dataroot)

    dataset_dir = get_full_path(os.path.join(dataroot, "textpred", _get_dataset_dir_name(dataset_name)))
    output_dir = get_full_path(pt_output_dir or os.path.join(output_dir, experiment_name), create_folder=True)

    if not os.path.isabs(cache_dir):
        cache_dir = os.path.join(dataset_dir, cache_dir)
    cache_dir = get_full_path(cache_dir, create_folder=True)

    if not os.path.isabs(pretrained_path) and pretrained_path:
        pretrained_path = os.path.join(os.path.dirname(output_dir), pretrained_path)

    return dataset_dir, output_dir, pretrained_path, cache_dir


@dataclass
class NvidiaTrainingArguments:
    """Define arguments used in the NVIDIA training pipeline.

    Args:
        experiment_name: Name of the experiment.
        checkpoint_file_path: Path to the checkpoint file.
        output_dir: Output folder.
        seed: Random seed.
        no_cuda: Whether CUDA should not be used.
        logging_steps: Number of steps between logs.
        do_eval: Whether to enable evaluation.
        eval_steps: Number of steps between evaluations.
        save_all_checkpoints: Whether to save all checkpoints from `eval_steps` steps.
        dataset_name: Name of the dataset.
        dataset_dir: Dataset folder.
        dataset_cache_dir: Dataset cache folder.
        dataset_refresh_cache: Whether cache should be refreshed.
        vocab_type: Name of the vocabulary/tokenizer.
        vocab_size: Size of the vocabulary.
        iterator_roll: Whether iterator should be rolled.
        global_batch_size: Global batch size.
        per_device_global_batch_size: Individual GPU batch size.
        seq_len: Sequence length.
        strategy: Distributed training strategy.
        local_rank: Local rank of process.
        find_unused_parameters: Whether unused parameters should be found.
        max_steps: Maximum number of training steps.
        gradient_accumulation_steps: Number of gradient accumulation steps.
        fp16: Whether FP16 precision should be used.
        optim: Name of the optimizer.
        learning_rate: Optimizer learning rate.
        weight_decay: Optimizer weight decay.
        momentum: Optimizer momentum.
        max_grad_norm: Optimizer gradients clipping value.
        lr_scheduler_type: Name of the scheduler.
        lr_qat_scheduler_type: Name of the QAT-based scheduler.
        lr_scheduler_max_steps: Maximum number of scheduler steps.
        lr_scheduler_warmup_steps: Number of warmup steps for the scheduler.
        lr_scheduler_patience: Scheduler patience.
        lr_scheduler_min_lr: Scheduler minimum learning rate.
        lr_scheduler_decay_rate: Scheduler decay rate.
        qat: Whether QAT should be used during training.
        mixed_qat: Whether MixedQAT should be used during training.

    """

    experiment_name: str = field(metadata={"help": "Name of the experiment."})

    checkpoint_file_path: str = field(default="", metadata={"help": "Path to the checkpoint file."})

    output_dir: str = field(default="~/logdir", metadata={"help": "Output folder."})

    seed: int = field(default=42, metadata={"help": "Random seed."})

    no_cuda: bool = field(default=False, metadata={"help": "Whether CUDA should not be used."})

    logging_steps: int = field(default=10, metadata={"help": "Number of steps between logs."})

    do_eval: bool = field(default=True, metadata={"help": "Whether to enable evaluation."})

    eval_steps: int = field(default=100, metadata={"help": "Number of steps between evaluations."})

    save_all_checkpoints: bool = field(
        default=False, metadata={"help": "Whether to save all checkpoints from `eval_steps` steps."}
    )

    dataset_name: str = field(default="wt103", metadata={"help": "Name of the dataset."})

    dataset_dir: str = field(default="", metadata={"help": "Dataset folder."})

    dataset_cache_dir: str = field(default="cache", metadata={"help": "Dataset cache folder."})

    dataset_refresh_cache: bool = field(default=False, metadata={"help": "Whether cache should be refreshed."})

    vocab_type: str = field(default="gpt2", metadata={"help": "Type of the tokenizer."})

    vocab_size: int = field(default=10000, metadata={"help": "Size of the vocabulary"})

    iterator_roll: bool = field(default=True, metadata={"help": "Whether iterator should be rolled."})

    global_batch_size: int = field(default=256, metadata={"help": "Global batch size."})

    per_device_global_batch_size: int = field(default=None, metadata={"help": "Individual GPU batch size."})

    seq_len: int = field(default=192, metadata={"help": "Sequence length."})

    strategy: str = field(default="ddp", metadata={"help": "Multi-GPU strategy."})

    local_rank: int = field(default=os.getenv("LOCAL_RANK", 0), metadata={"help": "Local rank of process."})

    find_unused_parameters: bool = field(default=False, metadata={"help": "Whether unused parameters should be found."})

    max_steps: int = field(default=40000, metadata={"help": "Maximum number of training steps."})

    gradient_accumulation_steps: int = field(default=1, metadata={"help": "Number of gradient accumulation steps."})

    fp16: bool = field(default=False, metadata={"help": "Whether FP16 precision should be used."})

    optim: str = field(default="jitlamb", metadata={"help": "Name of the optimizer."})

    learning_rate: float = field(default=0.01, metadata={"help": "Optimizer learning rate."})

    weight_decay: float = field(default=0.0, metadata={"help": "Optimizer weight decay."})

    momentum: float = field(default=0.0, metadata={"help": "Optimizer momentum."})

    max_grad_norm: float = field(default=0.25, metadata={"help": "Optimizer gradients clipping value."})

    lr_scheduler_type: str = field(default="cosine", metadata={"help": "Name of the scheduler."})

    lr_qat_scheduler_type: str = field(default="cosine", metadata={"help": "Name of the QAT-based scheduler."})

    lr_scheduler_max_steps: int = field(default=None, metadata={"help": "Maximum number of scheduler steps."})

    lr_scheduler_warmup_steps: int = field(default=1000, metadata={"help": "Number of scheduler warmup steps."})

    lr_scheduler_patience: float = field(default=0, metadata={"help": "Scheduler patience."})

    lr_scheduler_min_lr: float = field(default=0.001, metadata={"help": "Scheduler minimum learning rate."})

    lr_scheduler_decay_rate: float = field(default=0.5, metadata={"help": "Scheduler decay rate."})

    qat: bool = field(default=False, metadata={"help": "Whether QAT should be used during training."})

    mixed_qat: bool = field(default=False, metadata={"help": "Whether MixedQAT should be used during training."})

    @property
    def device(self) -> torch.device:
        """Return a PyTorch device instance."""

        return torch.device("cuda" if not self.no_cuda else "cpu")

    def __post_init__(self) -> None:
        """Override post-initialization with custom instructions.

        Ensure that `qat` and `mixed_qat` are not used together, set the random seed,
        initialize distributed training, create necessary directories,
        and set the global batch size.

        """

        assert not (self.qat and self.mixed_qat), "`qat` and `mixed_qat` should not be used together."

        np.random.seed(self.seed)
        torch.manual_seed(self.seed)

        self.local_rank = int(self.local_rank)
        if not self.no_cuda:
            torch.cuda.set_device(self.local_rank)
            init_distributed(True)

        (self.dataset_dir, self.output_dir, self.checkpoint_file_path, self.dataset_cache_dir,) = _create_dirs(
            self.dataset_dir,
            self.dataset_name,
            self.experiment_name,
            self.output_dir,
            self.checkpoint_file_path,
            self.dataset_cache_dir,
        )

        with sync_workers() as rank:
            if rank == 0:
                os.makedirs(self.output_dir, exist_ok=True)

        if self.per_device_global_batch_size is not None:
            world_size = get_world_size()
            self.global_batch_size = world_size * self.per_device_global_batch_size

    def to_dict(self) -> Dict[str, Any]:
        """Convert attributes into a dictionary representation.

        Returns:
            Attributes encoded as a dictionary.

        """

        return asdict(self)
