# Imports
# =============================================================================
import os
from distutils.util import strtobool
from pathlib import Path
from argparse import ArgumentParser

import warnings

import torch
from pytorch_lightning import Trainer
from pytorch_lightning.callbacks import ModelCheckpoint

from lib.modules.utils import LightningBaseModule
from lib.utils.config import Config
from lib.utils.logging import Logger

warnings.filterwarnings('ignore', category=FutureWarning)
warnings.filterwarnings('ignore', category=UserWarning)

_ROOT = Path(__file__).parent

# Paramter Configuration
# =============================================================================
# Argument Parser
main_arg_parser = ArgumentParser(description="parser for fast-neural-style")

# Main Parameters
main_arg_parser.add_argument("--main_debug", type=strtobool, default=False, help="")
main_arg_parser.add_argument("--main_eval", type=strtobool, default=False, help="")
main_arg_parser.add_argument("--main_seed", type=int, default=69, help="")

# Data Parameters
main_arg_parser.add_argument("--data_worker", type=int, default=0, help="")
main_arg_parser.add_argument("--data_batchsize", type=int, default=100, help="")
main_arg_parser.add_argument("--data_root", type=str, default='/data/', help="")
main_arg_parser.add_argument("--data_map_root", type=str, default='/res/maps', help="")

# Transformations
main_arg_parser.add_argument("--transformations_to_tensor", type=strtobool, default=False, help="")

# Transformations
main_arg_parser.add_argument("--train_outpath", type=str, default="output", help="")
main_arg_parser.add_argument("--train_version", type=strtobool, required=False, help="")
main_arg_parser.add_argument("--train_epochs", type=int, default=10, help="")
main_arg_parser.add_argument("--train_batch_size", type=int, default=512, help="")
main_arg_parser.add_argument("--train_lr", type=float, default=0.002, help="")

# Model
main_arg_parser.add_argument("--model_type", type=str, default="classifier_cnn", help="")
main_arg_parser.add_argument("--model_activation", type=str, default="relu", help="")
main_arg_parser.add_argument("--model_filters", type=str, default="[32, 16, 4]", help="")
main_arg_parser.add_argument("--model_classes", type=int, default=2, help="")
main_arg_parser.add_argument("--model_use_bias", type=strtobool, default=True, help="")
main_arg_parser.add_argument("--model_use_norm", type=strtobool, default=True, help="")
main_arg_parser.add_argument("--model_dropout", type=float, default=0.00, help="")

# Project
main_arg_parser.add_argument("--project_name", type=str, default='traj-gen', help="")
main_arg_parser.add_argument("--project_owner", type=str, default='si11ium', help="")
main_arg_parser.add_argument("--project_neptune_key", type=str, default=os.getenv('NEPTUNE_KEY'), help="")

# Parse it
args = main_arg_parser.parse_args()
config = Config.read_namespace(args)

# Logger
# =============================================================================
logger = Logger(config)

# Checkpoint Callback
# =============================================================================
checkpoint_callback = ModelCheckpoint(
            filepath=str(logger.log_dir / 'ckpt_weights'),
            verbose=True,
            period=1
)


if __name__ == "__main__":

    # Model
    # =============================================================================
    # Init
    model: LightningBaseModule = config.model_class(config.model_paramters)
    model.init_weights()

    # Trainer
    # =============================================================================
    trainer = Trainer(max_nb_epochs=config.train.epochs,
                      show_progress_bar=True,
                      weights_save_path=logger.log_dir,
                      gpus=[0] if torch.cuda.is_available() else None,
                      row_log_interval=model.data_len // 40,  # TODO: Better Value / Setting
                      log_save_interval=model.data_len // 10,  # TODO: Better Value / Setting
                      checkpoint_callback=checkpoint_callback,
                      logger=logger,
                      fast_dev_run=config.main.debug,
                      early_stop_callback=None
                      )

    # Train it
    trainer.fit(model)

    # Save the last state & all parameters
    config.exp_path.mkdir(parents=True, exist_ok=True)  # Todo: do i need this?
    trainer.save_checkpoint(logger.log_dir / 'weights.ckpt')
    model.save_to_disk(logger.log_dir)

    # TODO: Eval here!