# """My demo train script."""

import argparse
import logging
import os
import random
import time
import numpy as np
import torch

from torch import nn, optim, Tensor
from torch.utils.data import DataLoader, Dataset


def parse_args() -> argparse.Namespace:
    """Parse arguments."""
    parser = argparse.ArgumentParser(description="Training")
    parser.add_argument("--seed", type=int, help="Fix random seed", default=123)
    parser.add_argument(
        "--log_file", type=str, help="Log file", default="test_train.log"
    )
    parser.add_argument(
        "--log_path", type=str, help="Model path", default="./training_log/"
    )
    parser.add_argument(
        "--train_epochs", type=int, help="Epochs of training", default=5
    )
    parser.add_argument("--batch_size", type=int, help="Batch size", default=32)
    parser.add_argument(
        "--learning_rate",
        type=float,
        help="Learning rate",
        default=1e-3,
    )
    parser.add_argument("--device", type=str, help="Run on which device", default="cpu")
    parser.add_argument(
        "--cuda_visible_devices", type=str, help="Cuda visible devices", default="0"
    )
    return parser.parse_args()


def init_logging(log_file: str, level: str = "INFO") -> None:
    """Initialize logging."""
    logging.basicConfig(
        filename=log_file,
        filemode="w",
        format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
        level=level,
    )
    logging.getLogger().addHandler(logging.StreamHandler())


def set_seed(seed: int) -> None:
    """Set seed for reproducibility."""
    os.environ["PYTHONHASHSEED"] = str(seed)
    os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":4096:8"
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)
    torch.backends.cudnn.benchmark = False
    torch.backends.cudnn.deterministic = True
    torch.use_deterministic_algorithms(True)


def seed_worker(work_id: int) -> None:
    """Set seed for worker."""
    np.random.seed(work_id)
    random.seed(work_id)


class DatasetClass(Dataset):
    """My demo dataset class."""

    def __init__(self):
        self.input = np.random.rand(1000000, 2).astype(np.float32)
        # self.input[:, 1] = 0.0
        self.target = np.zeros([1000000, 1])
        self.target[:, 0] = self.input[:, 0] + 1.0

    def __len__(self):
        return len(self.input)

    def __getitem__(self, idx: int) -> tuple:
        return self.input[idx], self.target[idx]


class ModelClass(torch.nn.Module):
    """My demo model class."""

    def __init__(self):
        super().__init__()
        self.my_layer = nn.Linear(2, 1)

    def forward(self, inputs: Tensor) -> Tensor:
        """My demo forward function."""
        outputs = self.my_layer(inputs)
        return outputs


def get_loss(model_output: Tensor, target: Tensor) -> Tensor:
    """My demo loss function."""
    loss = torch.norm(model_output - target, dim=-1).sum()
    return loss


def training() -> None:
    """My demo training function."""
    train_set = DatasetClass()
    g = torch.Generator()
    g.manual_seed(args.seed)
    train_loader = DataLoader(
        dataset=train_set,
        batch_size=args.batch_size,
        shuffle=True,
        num_workers=os.cpu_count(),
        pin_memory=True,
        worker_init_fn=seed_worker,
        generator=g,
    )
    model = ModelClass()
    if args.device == "cuda":
        model = nn.DataParallel(model)
    model.to(args.device)
    optimizer = optim.AdamW(model.parameters(), lr=args.learning_rate)
    for epoch in range(args.train_epochs):
        model.train()
        for batch_index, (features, labels) in enumerate(train_loader):
            features = features.to(args.device)
            labels = labels.to(args.device)
            model_outputs = model(features)
            optimizer.zero_grad(set_to_none=True)
            loss = get_loss(model_outputs, labels)
            loss.backward()
            optimizer.step()
            if batch_index % 1000 == 0:
                logging.info(
                    "Epoch: %s, Batch index: %s, Loss: %s",
                    epoch,
                    batch_index,
                    loss.item(),
                )
    torch.save(model.state_dict(), f"{args.log_path}/trained_model.pth")


def testing() -> None:
    """My demo testing function."""
    test_set = DatasetClass()
    g = torch.Generator()
    g.manual_seed(args.seed)
    test_loader = DataLoader(
        dataset=test_set,
        batch_size=args.batch_size,
        shuffle=True,
        num_workers=os.cpu_count(),
        pin_memory=True,
        worker_init_fn=seed_worker,
        generator=g,
    )
    model = ModelClass()
    if args.device == "cuda":
        model = nn.DataParallel(model)
    model.load_state_dict(torch.load(f"{args.log_path}/trained_model.pth"))
    model.to(args.device)
    model.eval()
    with torch.no_grad():
        for batch_index, (features, labels) in enumerate(test_loader):
            features = features.to(args.device)
            labels = labels.to(args.device)
            model_outputs = model(features)
            loss = get_loss(model_outputs, labels)
            if batch_index % 1000 == 0:
                logging.info(
                    "Batch index: %s, Loss: %s",
                    batch_index,
                    loss.item() / args.batch_size,
                )


if __name__ == "__main__":
    args = parse_args()
    set_seed(args.seed)
    init_logging(args.log_file)
    os.environ["CUDA_VISIBLE_DEVICES"] = args.cuda_visible_devices
    main_start_time = time.time()
    training()
    main_end_time = time.time()
    logging.info("Main time: %s", main_end_time - main_start_time)
    testing()