from typing import List, Optional
from omegaconf import DictConfig
from .utils import utils
from lightning.pytorch import (Callback,
                               LightningDataModule,
                               LightningModule,
                               Trainer,
                               seed_everything)
from lightning.pytorch.loggers import Logger
import hydra
import os

os.environ["TOKENIZERS_PARALLELISM"] = "false"

log = utils.get_logger(__name__)


@hydra.main(config_path="configs/", config_name="config.yaml", version_base='1.1')
def run(config: DictConfig) -> Optional[float]:
    """Contains training pipeline.
    Instantiates all PyTorch Lightning objects from config.

    Args:
        config (DictConfig): Configuration composed by Hydra.

    Returns:
        Optional[float]: Metric score for hyperparameter optimization.
    """
    
    utils.extras(config)
    if config.get("print_config"):
        utils.print_config(config, resolve=True)
        
    if config.get("use_hf_mirror"):
        os.environ["HF_ENDPOINT"] = "https://hf-mirror.com"
        log.info("Using hf-mirror to download datasets and models!")
        
    # Set seed for random number generators in pytorch, numpy and python.random
    if config.get("seed"):
        seed_everything(config.seed, workers=True)

    # Init lightning datamodule
    if len(config.datamodule) > 0:
        log.info(f"Instantiating datamodule <{config.datamodule._target_}>")
        datamodule: LightningDataModule = hydra.utils.instantiate(config.datamodule)
        datamodule.prepare_data()
        datamodule.setup(stage='fit')
    else : datamodule = None

    # Init lightning model
    if datamodule is not None:
        log.info(f"Instantiating model <{config.model._target_}>")
        model: LightningModule = hydra.utils.instantiate(config.model, **datamodule.hparams)
    else: 
        log.info(f"Instantiating model <{config.model._target_}>")
        model = hydra.utils.instantiate(config.model)

    # Init lightning callbacks
    callbacks: List[Callback] = []
    if "callbacks" in config:
        for _, cb_conf in config.callbacks.items():
            if "_target_" in cb_conf:
                log.info(f"Instantiating callback <{cb_conf._target_}>")
                callbacks.append(hydra.utils.instantiate(cb_conf))

    # Init lightning loggers
    logger: List[Logger] = []
    if "logger" in config:
            log.info(f"Instantiating logger <{config.logger._target_}>")
            logger.append(hydra.utils.instantiate(config.logger))

    # Init lightning trainer
    log.info(f"Instantiating trainer <{config.trainer._target_}>")
    trainer: Trainer = hydra.utils.instantiate(config.trainer, 
                                               callbacks=callbacks, 
                                               logger=logger)

    # Send some parameters from config to all lightning loggers
    log.info("Logging hyperparameters!")
    utils.log_hyperparameters(config=config,
                              model=model,
                              datamodule=datamodule,
                              trainer=trainer,
                              callbacks=callbacks,
                              logger=logger)


    # Train the model
    log.info("Starting training!")
    trainer.fit(model=model, datamodule=datamodule)

    # Get metric score for hyperparameter optimization
    optimized_metric = config.get("optimized_metric")
    if optimized_metric and optimized_metric not in trainer.callback_metrics:
        raise Exception(
            "Metric for hyperparameter optimization not found! "
            "Make sure the `optimized_metric` in `hparams_search` config is correct!"
        )
    score = trainer.callback_metrics.get(optimized_metric)

    # Test the model
    if config.get("test_after_training") and not config.trainer.get("fast_dev_run"):
        log.info("Starting testing!")
        datamodule.setup(stage='test')
        trainer.test(model=model, datamodule=datamodule, ckpt_path="best")

    # Make sure everything closed properly
    log.info("Finalizing!")
    utils.finish(config=config,
                 model=model,
                 datamodule=datamodule,
                 trainer=trainer,
                 callbacks=callbacks,
                 logger=logger)

    # Print path to best checkpoint
    if not config.trainer.get("fast_dev_run"):
        log.info(f"Best model ckpt at {trainer.checkpoint_callback.best_model_path}")

    # Return metric score for hyperparameter optimization
    return score


if __name__ == "__main__":
    run()