import os

import fire

from pytorch_lightning import Trainer

import numpy as np

import pdb

from util import init_exp_folder, Args

from lightning import (get_task,

                       load_task,

                       get_ckpt_callback, 

                       get_early_stop_callback,

                       get_logger)

def train(dataset_folder="./data_files",

          save_dir="./sandbox",

          exp_name="DemoExperiment",

          model="ResNet18",

          task='classification',

          gpus=1,

          pretrained=True,

          num_classes=1,

          accelerator=None,

          logger_type='test_tube',

          gradient_clip_val=0.5,

          max_epochs=1,

          patience=10,

          stochastic_weight_avg=True,

          limit_train_batches=1.0,

          tb_path="./sandbox/tb",

          loss_fn="BCE",

          weights_summary=None,

          augmentation = 'none',

          num_workers=0,

          auto_lr_find= True,

          lr = 0.001,

          batch_size = 2,

 #         pretraining = False,

 #         aux_task = None

          ):

    """

    Run the training experiment.

    Args:

        save_dir: Path to save the checkpoints and logs

        exp_name: Name of the experiment

        model: Model name

        gpus: int. (ie: 2 gpus)

             OR list to specify which GPUs [0, 1] OR '0,1'

             OR '-1' / -1 to use all available gpus

        pretrained: Whether or not to use the pretrained model

        num_classes: Number of classes

        accelerator: Distributed computing mode

        logger_type: 'wandb' or 'test_tube'

        gradient_clip_val:  Clip value of gradient norm

        limit_train_batches: Proportion of training data to use

        max_epochs: Max number of epochs

        patience: number of epochs with no improvement after

                  which training will be stopped.

        stochastic_weight_avg: Whether to use stochastic weight averaging.

        tb_path: Path to global tb folder

        loss_fn: Loss function to use

        weights_summary: Prints a summary of the weights when training begins.

    Returns: None

    """

    args = Args(locals()) #Allows you to access stuff in the dictionary as args.exp_name

    init_exp_folder(args) #Sets up experiment directory 

    task = get_task(args) #Have to define this pytorch lightning module, for implementation, where constructor in segmentation.py is called

    #Then you instantiate trainer and start training

    trainer = Trainer(gpus=gpus, 

                      accelerator=accelerator,

                      logger=get_logger(logger_type, save_dir, exp_name), #Logging tool

                      callbacks=[get_early_stop_callback(patience), #Set number of epochs without improvement before stopping

                                 get_ckpt_callback(save_dir, exp_name)], #Save model checkpoints to folder, defined by certain metrics

                      weights_save_path=os.path.join(save_dir, exp_name), 

                      gradient_clip_val=None,

                      limit_train_batches=limit_train_batches, #When debugging, limit number of batches to run (percentage of data)

                      weights_summary=weights_summary,

                      stochastic_weight_avg=stochastic_weight_avg,

                      max_epochs=max_epochs,

                      auto_lr_find=True,

                      reload_dataloaders_every_n_epochs=1,

                      log_every_n_steps=1) #Handles functionality of training

    trainer.fit(task)

def test(ckpt_path,

         gpus=0,

         **kwargs):

    """

    Run the testing experiment.

    Args:

        ckpt_path: Path for the experiment to load

        gpus: int. (ie: 2 gpus)

             OR list to specify which GPUs [0, 1] OR '0,1'

             OR '-1' / -1 to use all available gpus

    Returns: None

    """

    task = load_task(ckpt_path, **kwargs)

    trainer = Trainer(gpus=gpus)

    trainer.test(task)

def predict(ckpt_path, gpus=1, prediction_path="predictions.pt", **kwargs):

    # couldn't figure out how to pass in a specific dataset as an argument

    # by default, this makes predictions over the test dataset

    # can change the prediction dataset in predict_dataloader() function in segmentation.py

    task = load_task(ckpt_path, **kwargs)

    trainer = Trainer(gpus=gpus)

    trainer.predict(task)

    preds_tensor = task.evaluator.preds

    preds = preds_tensor.cpu().detach()

    torch.save(prediction_path, predictions)

if __name__ == "__main__":

    fire.Fire() #Allows you to run functions and supply arguments directly in command line